fix local llm chat

different-ai · Dec 20, 2024 · 283c364 · 283c364
1 parent b01b2aa
commit 283c364
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 41 deletions.
diff --git a/packages/plugin/settings.ts b/packages/plugin/settings.ts
@@ -36,6 +36,7 @@ export class FileOrganizerSettings {
   showLocalLLMInChat = false;
   customFolderInstructions = "";
   selectedModel: "gpt-4o" | "llama3.2" = "gpt-4o";
+  customModelName = "llama3.2";
   tagScoreThreshold = 70;
   formatBehavior: "override" | "newFile" = "override";
   useInbox = false;

diff --git a/packages/plugin/views/ai-chat/chat.tsx b/packages/plugin/views/ai-chat/chat.tsx
@@ -107,29 +107,25 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
         return fetch(url, options);
       }
 
-      if (selectedModel === "llama3.2") {
-        const { messages, newUnifiedContext, currentDatetime, enableScreenpipe } =
-          JSON.parse(options.body as string);
-        logger.debug("llama3.2 context", {
-          contextLength: newUnifiedContext.length,
-          contextPreview: newUnifiedContext.slice(0, 200),
-          messageCount: messages.length,
-        });
-        const result = await streamText({
-          model: ollama("llama3.2"),
-          system: getChatSystemPrompt(
-            newUnifiedContext,
-            enableScreenpipe,
-            currentDatetime
-          ),
-          messages: convertToCoreMessages(messages),
-        });
-
-        return result.toDataStreamResponse();
-      }
-
-      // Default fetch behavior for remote API
-      return fetch(url, options);
+      // Handle local models (llama3.2 or custom)
+      const { messages, newUnifiedContext, currentDatetime, enableScreenpipe } =
+        JSON.parse(options.body as string);
+      logger.debug("local model context", {
+        model: selectedModel,
+        contextLength: newUnifiedContext.length,
+        contextPreview: newUnifiedContext.slice(0, 200),
+        messageCount: messages.length,
+      });
+      const result = await streamText({
+        model: ollama(selectedModel),
+        system: `
+          ${newUnifiedContext},
+          currentDatetime: ${currentDatetime},
+          `,
+        messages: convertToCoreMessages(messages),
+      });
+
+      return result.toDataStreamResponse();
     },
     onToolCall({ toolCall }) {
       logMessage("toolCall", toolCall);
@@ -187,7 +183,14 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
   const [maxContextSize] = useState(80 * 1000); // Keep this one
 
   // Update state to default to gpt-4
-  const [selectedModel, setSelectedModel] = useState<ModelType>("gpt-4o");
+  const [selectedModel, setSelectedModel] = useState<ModelType>(
+    plugin.settings.selectedModel
+  );
+
+  useEffect(() => {
+    // Update selectedModel when plugin settings change
+    setSelectedModel(plugin.settings.selectedModel);
+  }, [plugin.settings.selectedModel]);
 
   const handleTranscriptionComplete = (text: string) => {
     handleInputChange({
@@ -222,19 +225,15 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
 
       <div className="border-t border-[--background-modifier-border] p-4">
         <div className="flex items-center space-x-2 mb-4">
-
           <ContextItems />
 
           <ClearAllButton />
         </div>
 
-        <form
-          onSubmit={handleSendMessage}
-          className="flex items-end"
-        >
+        <form onSubmit={handleSendMessage} className="flex items-end">
           <div className="flex-grow overflow-y-auto relative" ref={inputRef}>
             <Tiptap
-            value={input}
+              value={input}
               onChange={handleTiptapChange}
               onKeyDown={handleKeyDown}
             />

diff --git a/packages/plugin/views/ai-chat/model-selector.tsx b/packages/plugin/views/ai-chat/model-selector.tsx
@@ -11,13 +11,30 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
   selectedModel,
   onModelSelect,
 }) => {
-  const [isModelSelectorOpen, setIsModelSelectorOpen] = React.useState(false);
   const plugin = usePlugin();
+  const [isModelSelectorOpen, setIsModelSelectorOpen] = React.useState(false);
+  const [isCustomizing, setIsCustomizing] = React.useState(false);
+  const [customModel, setCustomModel] = React.useState(plugin.settings.customModelName || "llama3.2");
 
   const handleModelSelect = async (model: ModelType) => {
+    if (model === "custom") {
+      setIsCustomizing(true);
+      return;
+    }
     onModelSelect(model);
-    plugin.settings.selectedModel = model;
+    if (model === "gpt-4o" || model === "llama3.2") {
+      plugin.settings.selectedModel = model;
+    }
+    await plugin.saveSettings();
+    setIsModelSelectorOpen(false);
+  };
+
+  const handleCustomModelSave = async () => {
+    plugin.settings.customModelName = customModel;
+    plugin.settings.selectedModel = customModel as "gpt-4o" | "llama3.2";
     await plugin.saveSettings();
+    onModelSelect(customModel);
+    setIsCustomizing(false);
     setIsModelSelectorOpen(false);
   };
 
@@ -28,7 +45,7 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
           onClick={() => plugin.settings.showLocalLLMInChat && setIsModelSelectorOpen(!isModelSelectorOpen)}
           className={`flex items-center space-x-1 text-[--text-muted] hover:text-[--text-normal] text-sm bg-[--background-primary-alt] ${plugin.settings.showLocalLLMInChat ? 'hover:bg-[--background-modifier-hover] cursor-pointer' : ''} px-2 py-1 rounded`}
         >
-          <span>{selectedModel === "gpt-4o" ? "gpt-4o" : "llama 3.2"}</span>
+          <span>{selectedModel}</span>
           {plugin.settings.showLocalLLMInChat && (
             <svg
               xmlns="http://www.w3.org/2000/svg"
@@ -46,7 +63,6 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
             </svg>
           )}
         </div>
-
         {isModelSelectorOpen && plugin.settings.showLocalLLMInChat && (
           <div className="absolute bottom-full right-0 mb-1 bg-[--background-primary] border border-[--background-modifier-border] rounded shadow-lg">
             <div className="py-1">
@@ -56,12 +72,38 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
               >
                 gpt-4o
               </div>
-              <div
-                onClick={() => handleModelSelect("llama3.2")}
-                className="cursor-pointer block w-full text-left px-4 py-2 text-sm text-[--text-normal] hover:bg-[--background-modifier-hover]"
-              >
-                llama 3.2
-              </div>
+              {isCustomizing ? (
+                <div className="px-4 py-2">
+                  <input
+                    type="text"
+                    value={customModel}
+                    onChange={(e) => setCustomModel(e.target.value)}
+                    className="w-full px-2 py-1 text-sm border rounded bg-[--background-primary] text-[--text-normal] border-[--background-modifier-border]"
+                    placeholder="Enter model name..."
+                  />
+                  <div className="flex justify-end mt-2 space-x-2">
+                    <button
+                      onClick={() => setIsCustomizing(false)}
+                      className="px-2 py-1 text-xs text-[--text-muted] hover:text-[--text-normal]"
+                    >
+                      Cancel
+                    </button>
+                    <button
+                      onClick={handleCustomModelSave}
+                      className="px-2 py-1 text-xs text-[--text-accent] hover:text-[--text-accent-hover]"
+                    >
+                      Save
+                    </button>
+                  </div>
+                </div>
+              ) : (
+                <div
+                  onClick={() => handleModelSelect("custom")}
+                  className="cursor-pointer block w-full text-left px-4 py-2 text-sm text-[--text-normal] hover:bg-[--background-modifier-hover]"
+                >
+                  Custom Model
+                </div>
+              )}
             </div>
           </div>
         )}

diff --git a/packages/plugin/views/ai-chat/types.ts b/packages/plugin/views/ai-chat/types.ts
@@ -1 +1 @@
-export type ModelType = "gpt-4o" | "llama3.2";
+export type ModelType = "gpt-4o" | "llama3.2" | string;
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		export type ModelType = "gpt-4o" \| "llama3.2";
		export type ModelType = "gpt-4o" \| "llama3.2" \| string;