From 074e2f3016706a87ea72b7f58c720bf9f4f4f1d8 Mon Sep 17 00:00:00 2001
From: Tommy <tommyh86@gmail.com>
Date: Wed, 6 Nov 2024 11:10:08 +0100
Subject: [PATCH] Moved provider and setProvider variables to the higher level
 component so that it can be accessed in sendMessage. Added provider to
 message queue in sendMessage. Changed streamText to extract both model and
 provider.

---
 app/components/chat/BaseChat.tsx    |  9 ++++--
 app/components/chat/Chat.client.tsx |  9 ++++--
 app/lib/.server/llm/stream-text.ts  | 47 +++++++++++++++++------------
 3 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/app/components/chat/BaseChat.tsx b/app/components/chat/BaseChat.tsx
index c1175f700..a47b6adfd 100644
--- a/app/components/chat/BaseChat.tsx
+++ b/app/components/chat/BaseChat.tsx
@@ -24,8 +24,7 @@ const EXAMPLE_PROMPTS = [
 
 const providerList = [...new Set(MODEL_LIST.map((model) => model.provider))]
 
-const ModelSelector = ({ model, setModel, modelList, providerList }) => {
-  const [provider, setProvider] = useState(DEFAULT_PROVIDER);
+const ModelSelector = ({ model, setModel, provider, setProvider, modelList, providerList }) => {
   return (
     <div className="mb-2">
       <select
@@ -79,6 +78,8 @@ interface BaseChatProps {
   input?: string;
   model: string;
   setModel: (model: string) => void;
+  provider: string;
+  setProvider: (provider: string) => void;
   handleStop?: () => void;
   sendMessage?: (event: React.UIEvent, messageInput?: string) => void;
   handleInputChange?: (event: React.ChangeEvent<HTMLTextAreaElement>) => void;
@@ -100,6 +101,8 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
       input = '',
       model,
       setModel,
+      provider,
+      setProvider,
       sendMessage,
       handleInputChange,
       enhancePrompt,
@@ -157,6 +160,8 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
                   model={model}
                   setModel={setModel}
                   modelList={MODEL_LIST}
+                  provider={provider}
+                  setProvider={setProvider}
                   providerList={providerList}
                 />
                 <div
diff --git a/app/components/chat/Chat.client.tsx b/app/components/chat/Chat.client.tsx
index 458bd8364..506d82dc2 100644
--- a/app/components/chat/Chat.client.tsx
+++ b/app/components/chat/Chat.client.tsx
@@ -11,7 +11,7 @@ import { useChatHistory } from '~/lib/persistence';
 import { chatStore } from '~/lib/stores/chat';
 import { workbenchStore } from '~/lib/stores/workbench';
 import { fileModificationsToHTML } from '~/utils/diff';
-import { DEFAULT_MODEL } from '~/utils/constants';
+import { DEFAULT_MODEL, DEFAULT_PROVIDER } from '~/utils/constants';
 import { cubicEasingFn } from '~/utils/easings';
 import { createScopedLogger, renderLogger } from '~/utils/logger';
 import { BaseChat } from './BaseChat';
@@ -74,6 +74,7 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
 
   const [chatStarted, setChatStarted] = useState(initialMessages.length > 0);
   const [model, setModel] = useState(DEFAULT_MODEL);
+  const [provider, setProvider] = useState(DEFAULT_PROVIDER);
 
   const { showChat } = useStore(chatStore);
 
@@ -182,7 +183,7 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
        * manually reset the input and we'd have to manually pass in file attachments. However, those
        * aren't relevant here.
        */
-      append({ role: 'user', content: `[Model: ${model}]\n\n${diff}\n\n${_input}` });
+      append({ role: 'user', content: `[Model: ${model}]\n\n[Provider: ${provider}]\n\n${diff}\n\n${_input}` });
 
       /**
        * After sending a new message we reset all modifications since the model
@@ -190,7 +191,7 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
        */
       workbenchStore.resetAllFileModifications();
     } else {
-      append({ role: 'user', content: `[Model: ${model}]\n\n${_input}` });
+      append({ role: 'user', content: `[Model: ${model}]\n\n[Provider: ${provider}]\n\n${_input}` });
     }
 
     setInput('');
@@ -215,6 +216,8 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
       sendMessage={sendMessage}
       model={model}
       setModel={setModel}
+      provider={provider}
+      setProvider={setProvider}
       messageRef={messageRef}
       scrollRef={scrollRef}
       handleInputChange={handleInputChange}
diff --git a/app/lib/.server/llm/stream-text.ts b/app/lib/.server/llm/stream-text.ts
index de3d5bfa8..70bb3a917 100644
--- a/app/lib/.server/llm/stream-text.ts
+++ b/app/lib/.server/llm/stream-text.ts
@@ -24,42 +24,51 @@ export type Messages = Message[];
 
 export type StreamingOptions = Omit<Parameters<typeof _streamText>[0], 'model'>;
 
-function extractModelFromMessage(message: Message): { model: string; content: string } {
+function extractPropertiesFromMessage(message: Message): { model: string; provider: string; content: string } {
   const modelRegex = /^\[Model: (.*?)\]\n\n/;
-  const match = message.content.match(modelRegex);
+  const providerRegex = /\[Provider: (.*?)\]\n\n/;
 
-  if (match) {
-    const model = match[1];
-    const content = message.content.replace(modelRegex, '');
-    return { model, content };
-  }
+  // Extract model
+  const modelMatch = message.content.match(modelRegex);
+  const model = modelMatch ? modelMatch[1] : DEFAULT_MODEL;
 
-  // Default model if not specified
-  return { model: DEFAULT_MODEL, content: message.content };
+  // Extract provider
+  const providerMatch = message.content.match(providerRegex);
+  const provider = providerMatch ? providerMatch[1] : DEFAULT_PROVIDER;
+
+  // Remove model and provider lines from content
+  const cleanedContent = message.content
+    .replace(modelRegex, '')
+    .replace(providerRegex, '')
+    .trim();
+
+  return { model, provider, content: cleanedContent };
 }
 
 export function streamText(messages: Messages, env: Env, options?: StreamingOptions) {
   let currentModel = DEFAULT_MODEL;
+  let currentProvider = DEFAULT_PROVIDER;
+
   const processedMessages = messages.map((message) => {
     if (message.role === 'user') {
-      const { model, content } = extractModelFromMessage(message);
-      if (model && MODEL_LIST.find((m) => m.name === model)) {
-        currentModel = model; // Update the current model
+      const { model, provider, content } = extractPropertiesFromMessage(message);
+
+      if (MODEL_LIST.find((m) => m.name === model)) {
+        currentModel = model;
       }
+
+      currentProvider = provider;
+
       return { ...message, content };
     }
-    return message;
-  });
 
-  const provider = MODEL_LIST.find((model) => model.name === currentModel)?.provider || DEFAULT_PROVIDER;
+    return message; // No changes for non-user messages
+  });
 
   return _streamText({
-    model: getModel(provider, currentModel, env),
+    model: getModel(currentProvider, currentModel, env),
     system: getSystemPrompt(),
     maxTokens: MAX_TOKENS,
-    // headers: {
-    //   'anthropic-beta': 'max-tokens-3-5-sonnet-2024-07-15',
-    // },
     messages: convertToCoreMessages(processedMessages),
     ...options,
   });