Update excel_chat.py

excel_chat.py

@@ -8,9 +8,9 @@ from groq import Groq
 import anthropic
 from users_management import update_json, users
 
-#users = ['maksG', 'Alma', 'YchK']
+#users = ['maksG', 'AlmaA', 'YchK']
 
-def ask_llm(query, input, client_index):
+def ask_llm(query, input, client_index, user):
     messages = [
                 {
                     "role": "system",
@@ -35,31 +35,31 @@ def ask_llm(query, input, client_index):
     ]
 
     if client_index == "Groq":
-        client = Groq(api_key=os.environ["GROQ_API_KEY"])
+        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['MISTRAL_API_KEY'])
+        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['MISTRAL_API_KEY'])
+        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['MISTRAL_API_KEY'])
+        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['CLAUDE_API_KEY'])
+        client = anthropic.Anthropic(api_key=os.environ[user['api_keys']['claude'])
         chat_completion = client.messages.create(
             model="claude-3-opus-20240229",
             max_tokens=350,
@@ -69,7 +69,7 @@ def ask_llm(query, input, client_index):
         ).content[0].text
         return chat_completion
     else:
-        client = anthropic.Anthropic(api_key=os.environ['CLAUDE_API_KEY'])
+        client = anthropic.Anthropic(api_key=os.environ[user['api_keys']['claude']])
         chat_completion = client.messages.create(
             model="claude-3-sonnet-20240229",
             max_tokens=350,
@@ -111,7 +111,7 @@ def chat_with_mistral(source_cols, dest_col, prompt, excel_file, url, search_col
         print('test')
         if not concatenated_content == "\n\n".join(f"{column_name}: nan" for column_name in source_cols):
             print('c bon')
-            llm_answer = ask_llm(prompt[0], concatenated_content, client)
+            llm_answer = ask_llm(prompt[0], concatenated_content, client, user)
             print(f"QUERY:\n{prompt[0]}\nCONTENT:\n{concatenated_content[:200]}...\n\nANSWER:\n{llm_answer}")
             df.at[index, dest_col] = llm_answer