From 9a288fd82e9d50b88e620c2c40e3aa29eca71d2f Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Thu, 15 May 2025 16:46:29 -0700
Subject: [PATCH] Fix local models (#176)

Fixes #167
---
 src/components/ProviderSettings.tsx      | 116 ++++++++++++-----------
 src/ipc/shared/language_model_helpers.ts |  40 +++++++-
 2 files changed, 94 insertions(+), 62 deletions(-)

diff --git a/src/components/ProviderSettings.tsx b/src/components/ProviderSettings.tsx
index 1ba1f5b..fe80833 100644
--- a/src/components/ProviderSettings.tsx
+++ b/src/components/ProviderSettings.tsx
@@ -102,66 +102,68 @@ export function ProviderSettingsGrid() {
     <div className="p-6">
       <h2 className="text-2xl font-bold mb-6">AI Providers</h2>
       <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-4">
-        {providers?.map((provider: LanguageModelProvider) => {
-          const isCustom = provider.type === "custom";
+        {providers
+          ?.filter((p) => p.type !== "local")
+          .map((provider: LanguageModelProvider) => {
+            const isCustom = provider.type === "custom";
 
-          return (
-            <Card
-              key={provider.id}
-              className="relative transition-all hover:shadow-md border-border"
-            >
-              <CardHeader
-                className="p-4 cursor-pointer"
-                onClick={() => handleProviderClick(provider.id)}
+            return (
+              <Card
+                key={provider.id}
+                className="relative transition-all hover:shadow-md border-border"
               >
-                <CardTitle className="text-xl flex items-center justify-between">
-                  {provider.name}
-                  {isProviderSetup(provider.id) ? (
-                    <span className="ml-3 text-sm font-medium text-green-500 bg-green-50 dark:bg-green-900/30 border border-green-500/50 dark:border-green-500/50 px-2 py-1 rounded-full">
-                      Ready
-                    </span>
-                  ) : (
-                    <span className="text-sm text-gray-500 bg-gray-50 dark:bg-gray-900 dark:text-gray-300 px-2 py-1 rounded-full">
-                      Needs Setup
-                    </span>
-                  )}
-                </CardTitle>
-                <CardDescription>
-                  {provider.hasFreeTier && (
-                    <span className="text-blue-600 mt-2 dark:text-blue-400 text-sm font-medium bg-blue-100 dark:bg-blue-900/30 px-2 py-1 rounded-full inline-flex items-center">
-                      <GiftIcon className="w-4 h-4 mr-1" />
-                      Free tier available
-                    </span>
-                  )}
-                </CardDescription>
-              </CardHeader>
-
-              {isCustom && (
-                <div
-                  className="absolute top-2 right-2"
-                  onClick={(e) => e.stopPropagation()}
+                <CardHeader
+                  className="p-4 cursor-pointer"
+                  onClick={() => handleProviderClick(provider.id)}
                 >
-                  <DropdownMenu>
-                    <DropdownMenuTrigger className="focus:outline-none">
-                      <div className="p-1 hover:bg-muted rounded-full">
-                        <MoreVertical className="h-4 w-4 text-muted-foreground" />
-                      </div>
-                    </DropdownMenuTrigger>
-                    <DropdownMenuContent align="end">
-                      <DropdownMenuItem
-                        className="text-destructive focus:text-destructive"
-                        onClick={() => setProviderToDelete(provider.id)}
-                      >
-                        <Trash2 className="h-4 w-4 mr-2" />
-                        Delete Provider
-                      </DropdownMenuItem>
-                    </DropdownMenuContent>
-                  </DropdownMenu>
-                </div>
-              )}
-            </Card>
-          );
-        })}
+                  <CardTitle className="text-xl flex items-center justify-between">
+                    {provider.name}
+                    {isProviderSetup(provider.id) ? (
+                      <span className="ml-3 text-sm font-medium text-green-500 bg-green-50 dark:bg-green-900/30 border border-green-500/50 dark:border-green-500/50 px-2 py-1 rounded-full">
+                        Ready
+                      </span>
+                    ) : (
+                      <span className="text-sm text-gray-500 bg-gray-50 dark:bg-gray-900 dark:text-gray-300 px-2 py-1 rounded-full">
+                        Needs Setup
+                      </span>
+                    )}
+                  </CardTitle>
+                  <CardDescription>
+                    {provider.hasFreeTier && (
+                      <span className="text-blue-600 mt-2 dark:text-blue-400 text-sm font-medium bg-blue-100 dark:bg-blue-900/30 px-2 py-1 rounded-full inline-flex items-center">
+                        <GiftIcon className="w-4 h-4 mr-1" />
+                        Free tier available
+                      </span>
+                    )}
+                  </CardDescription>
+                </CardHeader>
+
+                {isCustom && (
+                  <div
+                    className="absolute top-2 right-2"
+                    onClick={(e) => e.stopPropagation()}
+                  >
+                    <DropdownMenu>
+                      <DropdownMenuTrigger className="focus:outline-none">
+                        <div className="p-1 hover:bg-muted rounded-full">
+                          <MoreVertical className="h-4 w-4 text-muted-foreground" />
+                        </div>
+                      </DropdownMenuTrigger>
+                      <DropdownMenuContent align="end">
+                        <DropdownMenuItem
+                          className="text-destructive focus:text-destructive"
+                          onClick={() => setProviderToDelete(provider.id)}
+                        >
+                          <Trash2 className="h-4 w-4 mr-2" />
+                          Delete Provider
+                        </DropdownMenuItem>
+                      </DropdownMenuContent>
+                    </DropdownMenu>
+                  </div>
+                )}
+              </Card>
+            );
+          })}
 
         {/* Add custom provider button */}
         <Card
diff --git a/src/ipc/shared/language_model_helpers.ts b/src/ipc/shared/language_model_helpers.ts
index 9f98d8a..8480d00 100644
--- a/src/ipc/shared/language_model_helpers.ts
+++ b/src/ipc/shared/language_model_helpers.ts
@@ -113,7 +113,7 @@ export const PROVIDER_TO_ENV_VAR: Record<string, string> = {
   openrouter: "OPENROUTER_API_KEY",
 };
 
-export const PROVIDERS: Record<
+export const CLOUD_PROVIDERS: Record<
   string,
   {
     displayName: string;
@@ -153,6 +153,23 @@ export const PROVIDERS: Record<
   },
 };
 
+const LOCAL_PROVIDERS: Record<
+  string,
+  {
+    displayName: string;
+    hasFreeTier: boolean;
+  }
+> = {
+  ollama: {
+    displayName: "Ollama",
+    hasFreeTier: true,
+  },
+  lmstudio: {
+    displayName: "LM Studio",
+    hasFreeTier: true,
+  },
+};
+
 /**
  * Fetches language model providers from both the database (custom) and hardcoded constants (cloud),
  * merging them with custom providers taking precedence.
@@ -181,11 +198,11 @@ export async function getLanguageModelProviders(): Promise<
 
   // Get hardcoded cloud providers
   const hardcodedProviders: LanguageModelProvider[] = [];
-  for (const providerKey in PROVIDERS) {
-    if (Object.prototype.hasOwnProperty.call(PROVIDERS, providerKey)) {
+  for (const providerKey in CLOUD_PROVIDERS) {
+    if (Object.prototype.hasOwnProperty.call(CLOUD_PROVIDERS, providerKey)) {
       // Ensure providerKey is a key of PROVIDERS
-      const key = providerKey as keyof typeof PROVIDERS;
-      const providerDetails = PROVIDERS[key];
+      const key = providerKey as keyof typeof CLOUD_PROVIDERS;
+      const providerDetails = CLOUD_PROVIDERS[key];
       if (providerDetails) {
         // Ensure providerDetails is not undefined
         hardcodedProviders.push({
@@ -202,6 +219,19 @@ export async function getLanguageModelProviders(): Promise<
     }
   }
 
+  for (const providerKey in LOCAL_PROVIDERS) {
+    if (Object.prototype.hasOwnProperty.call(LOCAL_PROVIDERS, providerKey)) {
+      const key = providerKey as keyof typeof LOCAL_PROVIDERS;
+      const providerDetails = LOCAL_PROVIDERS[key];
+      hardcodedProviders.push({
+        id: key,
+        name: providerDetails.displayName,
+        hasFreeTier: providerDetails.hasFreeTier,
+        type: "local",
+      });
+    }
+  }
+
   return [...hardcodedProviders, ...customProvidersMap.values()];
 }