Update excel_chat.py

excel_chat.py

@@ -10,76 +10,66 @@ from users_management import update_json, users
 
 #users = ['maksG', 'AlmaA', 'YchK']
 
-def ask_llm(query, input, client_index, user):
+def ask_llm(query, input, client_index):
     messages = [
-                {
-                    "role": "system",
-                    "content": f"You are a helpful assistant. Only show your final response to the **User Query**! Do not provide any explanations or details: \n# User Query:\n{query}."
-                },
-                {
-                    "role": "user",
-                    "content": f"{input}",
-                }
-            ]
-    systemC = f"You are a helpful assistant. Only show your final response to the **User Query**! Do not provide any explanations or details: \n# User Query:\n{query}."
-    messageC=[
+        {
+            "role": "system",
+            "content": f"You are a helpful assistant. Only show your final response to the **User Query**! Do not provide any explanations or details: \n# User Query:\n{query}."
+        },
         {
             "role": "user",
-            "content": [
-                {
-                    "type": "text",
-                    "text": f"{input}"
-                }
-            ]
+            "content": input,
         }
     ]
+    
+    systemC = messages[0]["content"]
+    messageC = [{
+        "role": "user",
+        "content": [{
+            "type": "text",
+            "text": input
+        }]
+    }]
 
-    if client_index == "Groq":
-        client = Groq(api_key=os.environ[user['api_keys']['groq']])
-        chat_completion = client.chat.completions.create(
-            messages=messages,
-            model='mixtral-8x7b-32768',
-        )
-    elif client_index == "Mistral Small":
-        client = MistralClient(api_key=os.environ[user['api_keys']['mistral']])
-        chat_completion = client.chat(
-            messages=messages,
-            model='mistral-small-latest',
-        )
-    elif client_index == "Mistral Tiny":
-        client = MistralClient(api_key=os.environ[user['api_keys']['mistral']])
-        chat_completion = client.chat(
-            messages=messages,
-            model='mistral-tiny',
-        )
-    elif client_index == "Mistral Medium":
-        client = MistralClient(api_key=os.environ[user['api_keys']['mistral']])
-        chat_completion = client.chat(
-            messages=messages,
-            model='mistral-medium',
-        )
-    elif client_index == "Claude Opus":
-        client = anthropic.Anthropic(api_key=os.environ[user['api_keys']['claude']])
-        chat_completion = client.messages.create(
-            model="claude-3-opus-20240229",
+    if "Mistral" in client_index:
+        client = MistralClient(api_key=os.environ['MISTRAL_API_KEY'])
+        model_map = {
+            "Mistral Tiny": "mistral-tiny",
+            "Mistral Small": "mistral-small-latest",
+            "Mistral Medium": "mistral-medium",
+        }
+        chat_completion = client.chat(messages=messages, model=model_map[client_index])
+
+    elif "Claude" in client_index:
+        client = anthropic.Anthropic(api_key=os.environ['CLAUDE_API_KEY'])
+        model_map = {
+            "Claude Sonnet": "claude-3-sonnet-20240229",
+            "Claude Opus": "claude-3-opus-20240229",
+        }
+        response = client.messages.create(
+            model=model_map[client_index],
             max_tokens=350,
             temperature=0,
             system=systemC,
             messages=messageC
-        ).content[0].text
-        return chat_completion
+        )
+        return response.content[0].text
+
+    elif client_index == "Groq (mixtral)":
+        client = Groq(api_key=os.environ["GROQ_API_KEY"])
+        chat_completion = client.chat.completions.create(
+            messages=messages,
+            model='mixtral-8x7b-32768',
+        )
+
     else:
-        client = anthropic.Anthropic(api_key=os.environ[user['api_keys']['claude']])
-        chat_completion = client.messages.create(
-            model="claude-3-sonnet-20240229",
-            max_tokens=350,
-            temperature=0,
-            system=systemC,
-            messages=messageC
-        ).content[0].text
-        return chat_completion
+        raise ValueError("Unsupported client index provided")
+
+    # Return the response, handling the structure specific to Groq and Mistral clients.
+    return chat_completion.choices[0].message.content if client_index != "Claude" else chat_completion
+
+
 
-    return chat_completion.choices[0].message.content
 
 def filter_df(df, column_name, keywords):
     if len(keywords)>0: