Update excel_chat.py

excel_chat.py

@@ -19,19 +19,41 @@ 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}."   
+    messageC=[
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": f"{input}"
+                }
+            ]
+        }
+    ]
 
-    if client_index == 0:
+    if client_index == "Groq":
         client = Groq(api_key=os.environ["GROQ_API_KEY"])
         chat_completion = client.chat.completions.create(
             messages=messages,
             model='mixtral-8x7b-32768',
         )
-    else:
+    elif client_index == "Mistral":
         client = MistralClient(api_key=os.environ['MISTRAL_API_KEY'])
         chat_completion = client.chat(
             messages=messages,
             model='mistral-small-latest',
         )
+    else:
+        client = anthropic.Anthropic(api_key=os.environ['CLAUDE_API_KEY'])
+        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
 
     return chat_completion.choices[0].message.content
 
@@ -61,17 +83,20 @@ def chat_with_mistral(source_cols, dest_col, prompt, excel_file, url, search_col
 
     for index, row in filtred_df.iterrows():
         concatenated_content = "\n\n".join(f"{column_name}: {str(row[column_name])}" for column_name in source_cols)
-        llm_answer = ask_llm(prompt, concatenated_content, client)
-        print(f"QUERY:\n{prompt}\nCONTENT:\n{concatenated_content[:200]}...\n\nANSWER:\n{llm_answer}")
-        df.at[index, dest_col] = llm_answer
+        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, concatenated_content, client)
+            print(f"QUERY:\n{prompt}\nCONTENT:\n{concatenated_content[:200]}...\n\nANSWER:\n{llm_answer}")
+            df.at[index, dest_col] = llm_answer
 
     df.to_excel(file_name, index=False)
     return file_name, df.head(5)
 
-def get_columns(file):
+
+def get_columns(file,progress=gr.Progress()):
     if file is not None:
         df = pd.read_excel(file)
-        columns = list(df.columns)
-        return gr.update(choices=columns), gr.update(choices=columns), gr.update(choices=columns), gr.update(choices=columns + [""]), gr.update(choices=columns + ['[ALL]']), df.head(5)
+        return gr.update(choices=list(df.columns)), gr.update(choices=list(df.columns)), df.head(5)
     else:
-        return gr.update(choices=[]), gr.update(choices=[]), gr.update(choices=[]), gr.update(choices=[]), gr.update(choices=[]), pd.DataFrame()
+        return gr.update(choices=[]), gr.update(choices=[]), pd.DataFrame()