Update excel_chat.py

excel_chat.py

@@ -21,7 +21,7 @@ def ask_llm(query, input, client_index):
                     "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}."   
+    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": "user",
@@ -40,12 +40,34 @@ def ask_llm(query, input, client_index):
             messages=messages,
             model='mixtral-8x7b-32768',
         )
-    elif client_index == "Mistral":
+    elif client_index == "Mistral Small":
         client = MistralClient(api_key=os.environ['MISTRAL_API_KEY'])
         chat_completion = client.chat(
             messages=messages,
             model='mistral-small-latest',
         )
+    elif client_index == "Mistral Tiny":
+        client = MistralClient(api_key=os.environ['MISTRAL_API_KEY'])
+        chat_completion = client.chat(
+            messages=messages,
+            model='mistral-tiny',
+        )
+    elif client_index == "Mistral Medium":
+        client = MistralClient(api_key=os.environ['MISTRAL_API_KEY'])
+        chat_completion = client.chat(
+            messages=messages,
+            model='mistral-medium',
+        )
+    elif client_index == "Claude Opus":
+        client = anthropic.Anthropic(api_key=os.environ['CLAUDE_API_KEY'])
+        chat_completion = client.messages.create(
+            model="claude-3-opus-20240229",
+            max_tokens=350,
+            temperature=0,
+            system=systemC,
+            messages=messageC
+        ).content[0].text
+        return chat_completion
     else:
         client = anthropic.Anthropic(api_key=os.environ['CLAUDE_API_KEY'])
         chat_completion = client.messages.create(