Add ollama support (#7)

src/components/ModelPicker.tsx

@@ -1,17 +1,25 @@
 import type { LargeLanguageModel, ModelProvider } from "@/lib/schemas";
 import { Button } from "@/components/ui/button";
-import {
-  Popover,
-  PopoverContent,
-  PopoverTrigger,
-} from "@/components/ui/popover";
 import {
   Tooltip,
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
-import { useState } from "react";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuLabel,
+  DropdownMenuSeparator,
+  DropdownMenuTrigger,
+  DropdownMenuSub,
+  DropdownMenuSubTrigger,
+  DropdownMenuSubContent,
+} from "@/components/ui/dropdown-menu";
+import { useEffect, useState } from "react";
 import { MODEL_OPTIONS } from "@/constants/models";
+import { useLocalModels } from "@/hooks/useLocalModels";
+import { ChevronDown } from "lucide-react";
 
 interface ModelPickerProps {
   selectedModel: LargeLanguageModel;
@@ -23,9 +31,37 @@ export function ModelPicker({
   onModelSelect,
 }: ModelPickerProps) {
   const [open, setOpen] = useState(false);
-  const modelDisplayName = MODEL_OPTIONS[selectedModel.provider].find(
-    (model) => model.name === selectedModel.name
-  )?.displayName;
+  const {
+    models: localModels,
+    loading: localModelsLoading,
+    error: localModelsError,
+    loadModels,
+  } = useLocalModels();
+
+  // Load local models when the component mounts or the dropdown opens
+  useEffect(() => {
+    if (open) {
+      loadModels();
+    }
+  }, [open, loadModels]);
+
+  // Get display name for the selected model
+  const getModelDisplayName = () => {
+    if (selectedModel.provider === "ollama") {
+      return (
+        localModels.find((model) => model.modelName === selectedModel.name)
+          ?.displayName || selectedModel.name
+      );
+    }
+
+    return (
+      MODEL_OPTIONS[selectedModel.provider]?.find(
+        (model) => model.name === selectedModel.name
+      )?.displayName || selectedModel.name
+    );
+  };
+
+  const modelDisplayName = getModelDisplayName();
 
   // Flatten the model options into a single array with provider information
   const allModels = Object.entries(MODEL_OPTIONS).flatMap(
@@ -36,9 +72,13 @@ export function ModelPicker({
       }))
   );
 
+  // Determine if we have local models available
+  const hasLocalModels =
+    !localModelsLoading && !localModelsError && localModels.length > 0;
+
   return (
-    <Popover open={open} onOpenChange={setOpen}>
-      <PopoverTrigger asChild>
+    <DropdownMenu open={open} onOpenChange={setOpen}>
+      <DropdownMenuTrigger asChild>
         <Button
           variant="outline"
           size="sm"
@@ -48,46 +88,127 @@ export function ModelPicker({
             <span className="text-xs text-muted-foreground">Model:</span>{" "}
             {modelDisplayName}
           </span>
+          <ChevronDown className="h-4 w-4" />
         </Button>
-      </PopoverTrigger>
-      <PopoverContent className="w-56 p-1" align="start">
-        <div className="grid gap-2">
-          {allModels.map((model) => (
-            <Tooltip key={model.name}>
-              <TooltipTrigger asChild>
-                <Button
-                  variant={
-                    selectedModel.name === model.name ? "secondary" : "ghost"
+      </DropdownMenuTrigger>
+      <DropdownMenuContent className="w-56" align="start">
+        <DropdownMenuLabel>Cloud Models</DropdownMenuLabel>
+        <DropdownMenuSeparator />
+
+        {/* Cloud models */}
+        {allModels.map((model) => (
+          <Tooltip key={`${model.provider}-${model.name}`}>
+            <TooltipTrigger asChild>
+              <DropdownMenuItem
+                className={
+                  selectedModel.provider === model.provider &&
+                  selectedModel.name === model.name
+                    ? "bg-secondary"
+                    : ""
+                }
+                onClick={() => {
+                  onModelSelect({
+                    name: model.name,
+                    provider: model.provider,
+                  });
+                  setOpen(false);
+                }}
+              >
+                <div className="flex justify-between items-start w-full">
+                  <span className="flex flex-col items-start">
+                    <span>{model.displayName}</span>
+                    <span className="text-xs text-muted-foreground">
+                      {model.provider}
+                    </span>
+                  </span>
+                  {model.tag && (
+                    <span className="text-[10px] bg-primary/10 text-primary px-1.5 py-0.5 rounded-full font-medium">
+                      {model.tag}
+                    </span>
+                  )}
+                </div>
+              </DropdownMenuItem>
+            </TooltipTrigger>
+            <TooltipContent side="right">{model.description}</TooltipContent>
+          </Tooltip>
+        ))}
+
+        <DropdownMenuSeparator />
+
+        {/* Ollama Models Dropdown */}
+        <DropdownMenuSub>
+          <DropdownMenuSubTrigger
+            disabled={localModelsLoading || !hasLocalModels}
+            className="w-full font-normal"
+          >
+            <div className="flex flex-col items-start">
+              <span>Local models (Ollama)</span>
+              {localModelsLoading ? (
+                <span className="text-xs text-muted-foreground">
+                  Loading...
+                </span>
+              ) : !hasLocalModels ? (
+                <span className="text-xs text-muted-foreground">
+                  None available
+                </span>
+              ) : (
+                <span className="text-xs text-muted-foreground">
+                  {localModels.length} models
+                </span>
+              )}
+            </div>
+          </DropdownMenuSubTrigger>
+          <DropdownMenuSubContent className="w-56">
+            <DropdownMenuLabel>Ollama Models</DropdownMenuLabel>
+            <DropdownMenuSeparator />
+
+            {localModelsLoading ? (
+              <div className="text-xs text-center py-2 text-muted-foreground">
+                Loading models...
+              </div>
+            ) : localModelsError ? (
+              <div className="text-xs text-center py-2 text-muted-foreground">
+                Error loading models
+              </div>
+            ) : localModels.length === 0 ? (
+              <div className="px-2 py-1.5 text-sm">
+                <div className="flex flex-col">
+                  <span>No local models available</span>
+                  <span className="text-xs text-muted-foreground">
+                    Start Ollama to use local models
+                  </span>
+                </div>
+              </div>
+            ) : (
+              localModels.map((model) => (
+                <DropdownMenuItem
+                  key={`local-${model.modelName}`}
+                  className={
+                    selectedModel.provider === "ollama" &&
+                    selectedModel.name === model.modelName
+                      ? "bg-secondary"
+                      : ""
                   }
-                  className="w-full justify-start font-normal"
                   onClick={() => {
                     onModelSelect({
-                      name: model.name,
-                      provider: model.provider,
+                      name: model.modelName,
+                      provider: "ollama",
                     });
                     setOpen(false);
                   }}
                 >
-                  <div className="flex justify-between items-start w-full">
-                    <span className="flex flex-col items-start">
-                      <span>{model.displayName}</span>
-                      <span className="text-xs text-muted-foreground">
-                        {model.provider}
-                      </span>
+                  <div className="flex flex-col">
+                    <span>{model.displayName}</span>
+                    <span className="text-xs text-muted-foreground truncate">
+                      {model.modelName}
                     </span>
-                    {model.tag && (
-                      <span className="text-[10px] bg-primary/10 text-primary px-1.5 py-0.5 rounded-full font-medium">
-                        {model.tag}
-                      </span>
-                    )}
                   </div>
-                </Button>
-              </TooltipTrigger>
-              <TooltipContent side="right">{model.description}</TooltipContent>
-            </Tooltip>
-          ))}
-        </div>
-      </PopoverContent>
-    </Popover>
+                </DropdownMenuItem>
+              ))
+            )}
+          </DropdownMenuSubContent>
+        </DropdownMenuSub>
+      </DropdownMenuContent>
+    </DropdownMenu>
   );
 }

src/components/ui/dropdown-menu.tsx

@@ -1,13 +1,13 @@
-import * as React from "react"
-import * as DropdownMenuPrimitive from "@radix-ui/react-dropdown-menu"
-import { CheckIcon, ChevronRightIcon, CircleIcon } from "lucide-react"
+import * as React from "react";
+import * as DropdownMenuPrimitive from "@radix-ui/react-dropdown-menu";
+import { CheckIcon, ChevronRightIcon, CircleIcon } from "lucide-react";
 
-import { cn } from "@/lib/utils"
+import { cn } from "@/lib/utils";
 
 function DropdownMenu({
   ...props
 }: React.ComponentProps<typeof DropdownMenuPrimitive.Root>) {
-  return <DropdownMenuPrimitive.Root data-slot="dropdown-menu" {...props} />
+  return <DropdownMenuPrimitive.Root data-slot="dropdown-menu" {...props} />;
 }
 
 function DropdownMenuPortal({
@@ -15,7 +15,7 @@ function DropdownMenuPortal({
 }: React.ComponentProps<typeof DropdownMenuPrimitive.Portal>) {
   return (
     <DropdownMenuPrimitive.Portal data-slot="dropdown-menu-portal" {...props} />
-  )
+  );
 }
 
 function DropdownMenuTrigger({
@@ -26,7 +26,7 @@ function DropdownMenuTrigger({
       data-slot="dropdown-menu-trigger"
       {...props}
     />
-  )
+  );
 }
 
 function DropdownMenuContent({
@@ -46,7 +46,7 @@ function DropdownMenuContent({
         {...props}
       />
     </DropdownMenuPrimitive.Portal>
-  )
+  );
 }
 
 function DropdownMenuGroup({
@@ -54,7 +54,7 @@ function DropdownMenuGroup({
 }: React.ComponentProps<typeof DropdownMenuPrimitive.Group>) {
   return (
     <DropdownMenuPrimitive.Group data-slot="dropdown-menu-group" {...props} />
-  )
+  );
 }
 
 function DropdownMenuItem({
@@ -63,8 +63,8 @@ function DropdownMenuItem({
   variant = "default",
   ...props
 }: React.ComponentProps<typeof DropdownMenuPrimitive.Item> & {
-  inset?: boolean
-  variant?: "default" | "destructive"
+  inset?: boolean;
+  variant?: "default" | "destructive";
 }) {
   return (
     <DropdownMenuPrimitive.Item
@@ -72,12 +72,12 @@ function DropdownMenuItem({
       data-inset={inset}
       data-variant={variant}
       className={cn(
-        "focus:bg-accent focus:text-accent-foreground data-[variant=destructive]:text-destructive data-[variant=destructive]:focus:bg-destructive/10 dark:data-[variant=destructive]:focus:bg-destructive/20 data-[variant=destructive]:focus:text-destructive data-[variant=destructive]:*:[svg]:!text-destructive [&_svg:not([class*='text-'])]:text-muted-foreground relative flex cursor-default items-center gap-2 rounded-sm px-2 py-1.5 text-sm outline-hidden select-none data-[disabled]:pointer-events-none data-[disabled]:opacity-50 data-[inset]:pl-8 [&_svg]:pointer-events-none [&_svg]:shrink-0 [&_svg:not([class*='size-'])]:size-4",
+        "focus:bg-accent focus:text-accent-foreground data-[variant=destructive]:text-destructive data-[variant=destructive]:focus:bg-destructive/10 dark:data-[variant=destructive]:focus:bg-destructive/20 data-[variant=destructive]:focus:text-destructive data-[variant=destructive]:*:[svg]:!text-destructive [&_svg:not([class*='text-'])]:text-muted-foreground relative flex cursor-default items-center gap-1 rounded-sm px-2 py-1 text-sm outline-hidden select-none data-[disabled]:pointer-events-none data-[disabled]:opacity-50 data-[inset]:pl-8 [&_svg]:pointer-events-none [&_svg]:shrink-0 [&_svg:not([class*='size-'])]:size-4",
         className
       )}
       {...props}
     />
-  )
+  );
 }
 
 function DropdownMenuCheckboxItem({
@@ -103,7 +103,7 @@ function DropdownMenuCheckboxItem({
       </span>
       {children}
     </DropdownMenuPrimitive.CheckboxItem>
-  )
+  );
 }
 
 function DropdownMenuRadioGroup({
@@ -114,7 +114,7 @@ function DropdownMenuRadioGroup({
       data-slot="dropdown-menu-radio-group"
       {...props}
     />
-  )
+  );
 }
 
 function DropdownMenuRadioItem({
@@ -138,7 +138,7 @@ function DropdownMenuRadioItem({
       </span>
       {children}
     </DropdownMenuPrimitive.RadioItem>
-  )
+  );
 }
 
 function DropdownMenuLabel({
@@ -146,7 +146,7 @@ function DropdownMenuLabel({
   inset,
   ...props
 }: React.ComponentProps<typeof DropdownMenuPrimitive.Label> & {
-  inset?: boolean
+  inset?: boolean;
 }) {
   return (
     <DropdownMenuPrimitive.Label
@@ -158,7 +158,7 @@ function DropdownMenuLabel({
       )}
       {...props}
     />
-  )
+  );
 }
 
 function DropdownMenuSeparator({
@@ -171,7 +171,7 @@ function DropdownMenuSeparator({
       className={cn("bg-border -mx-1 my-1 h-px", className)}
       {...props}
     />
-  )
+  );
 }
 
 function DropdownMenuShortcut({
@@ -187,13 +187,13 @@ function DropdownMenuShortcut({
       )}
       {...props}
     />
-  )
+  );
 }
 
 function DropdownMenuSub({
   ...props
 }: React.ComponentProps<typeof DropdownMenuPrimitive.Sub>) {
-  return <DropdownMenuPrimitive.Sub data-slot="dropdown-menu-sub" {...props} />
+  return <DropdownMenuPrimitive.Sub data-slot="dropdown-menu-sub" {...props} />;
 }
 
 function DropdownMenuSubTrigger({
@@ -202,7 +202,7 @@ function DropdownMenuSubTrigger({
   children,
   ...props
 }: React.ComponentProps<typeof DropdownMenuPrimitive.SubTrigger> & {
-  inset?: boolean
+  inset?: boolean;
 }) {
   return (
     <DropdownMenuPrimitive.SubTrigger
@@ -217,7 +217,7 @@ function DropdownMenuSubTrigger({
       {children}
       <ChevronRightIcon className="ml-auto size-4" />
     </DropdownMenuPrimitive.SubTrigger>
-  )
+  );
 }
 
 function DropdownMenuSubContent({
@@ -233,7 +233,7 @@ function DropdownMenuSubContent({
       )}
       {...props}
     />
-  )
+  );
 }
 
 export {
@@ -252,4 +252,4 @@ export {
   DropdownMenuSub,
   DropdownMenuSubTrigger,
   DropdownMenuSubContent,
-}
+};