Update excel_chat.py

excel_chat.py

@@ -9,7 +9,8 @@ import anthropic
 from users_management import update_json, users
 from code_df_custom import load_excel
 import zipfile
-from openai import OpenAI
+from openai import *
+import time
 
 #users = ['maksG', 'AlmaA', 'YchK']
 
@@ -34,88 +35,153 @@ def ask_llm(query, user_input, client_index, user, keys):
         }]
     }]
 
-    if "Mistral" in client_index:
-        client = MistralClient(api_key=os.environ[user['api_keys']['mistral']])
-        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[user['api_keys']['claude']])
-        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
-        )
-        return response.content[0].text
-
-    elif "GPT 4o" in client_index:
-        client = OpenAI(api_key=os.environ["OPENAI_YCHK"])
-        response = client.chat.completions.create(
-            model="gpt-4o",
-            messages=messageC
-        )
-        return response.choices[0][message][content].text
-
-    elif "Perplexity" in client_index:
-        client = OpenAI(api_key=os.environ["PERPLEXITY_ALMAA"], base_url="https://api.perplexity.ai")
-        model_map = {
-                "Perplexity Llama3 70b": "llama-3-70b-instruct",
-                "Perplexity Llama3 8b": "llama-3-8b-instruct",
-                "Perplexity Llama3 Sonar Small": "llama-3-sonar-small-32k-chat",
-                "Perplexity Llama3 Sonar Large": "llama-3-sonar-large-32k-chat"
-            }
-        response = client.chat.completions.create(
-            model=model_map[client_index],
-            messages=messageC
-        )
-
-        responseContent = str(response.choices[0].message.content)
-        print(responseContent)        
-        return responseContent,keys
+    try: 
         
-    elif "Groq" in client_index:
-        try:
-            client = Groq(api_key= os.getenv(keys[0]))
+        if "Mistral" in client_index:
+            client = MistralClient(api_key=os.environ[user['api_keys']['mistral']])
+            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[user['api_keys']['claude']])
             model_map = {
-                "Groq Mixtral": "mixtral-8x7b-32768",
-                "Groq Llama3 70b": "llama3-70b-8192",
-                "Groq Llama3 8b": "llama3-8b-8192"
+                "Claude Sonnet": "claude-3-sonnet-20240229",
+                "Claude Opus": "claude-3-opus-20240229",
             }
-            chat_completion = client.chat.completions.create(
-                messages=messages,
+            response = client.messages.create(
                 model=model_map[client_index],
+                max_tokens=350,
+                temperature=0,
+                system=systemC,
+                messages=messageC
             )
-            response = chat_completion.choices[0].message.content
-        except Exception as e:
-            print("Change key")
-            if keys[0] == keys[1][0]:
-                keys[0] = keys[1][1]
-            elif keys[0] == keys[1][1]:
-                keys[0] = keys[1][2]
-            else:
-                keys[0] = keys[1][0]
-                
-            client = Groq(api_key= os.getenv(keys[0]))
-            chat_completion = client.chat.completions.create(
-                messages=messages,
-                model='llama3-8b-8192',
+            return response.content[0].text
+    
+        elif "GPT 4o" in client_index:
+            client = OpenAI(api_key=os.environ["OPENAI_YCHK"])
+            response = client.chat.completions.create(
+                model="gpt-4o",
+                messages=messageC
             )
-            response = chat_completion.choices[0].message.content
-    else:
-        raise ValueError("Unsupported client index provided")
+            return response.choices[0][message][content].text
+    
+        elif "Perplexity" in client_index:
+            client = OpenAI(api_key=os.environ["PERPLEXITY_ALMAA"], base_url="https://api.perplexity.ai")
+            model_map = {
+                    "Perplexity Llama3 70b": "llama-3-70b-instruct",
+                    "Perplexity Llama3 8b": "llama-3-8b-instruct",
+                    "Perplexity Llama3 Sonar Small": "llama-3-sonar-small-32k-chat",
+                    "Perplexity Llama3 Sonar Large": "llama-3-sonar-large-32k-chat"
+            }
+            
+            response = client.chat.completions.create(
+                model=model_map[client_index],
+                messages=messageC
+            )
+    
+            responseContent = str(response.choices[0].message.content)
+            print(responseContent)        
+            return responseContent,keys
+            
+        elif "Groq" in client_index:
+            try:
+                client = Groq(api_key= os.getenv(keys[0]))
+                model_map = {
+                    "Groq Mixtral": "mixtral-8x7b-32768",
+                    "Groq Llama3 70b": "llama3-70b-8192",
+                    "Groq Llama3 8b": "llama3-8b-8192"
+                }
+                chat_completion = client.chat.completions.create(
+                    messages=messages,
+                    model=model_map[client_index],
+                )
+                response = chat_completion.choices[0].message.content
+            except Exception as e:
+                print("Change key")
+                if keys[0] == keys[1][0]:
+                    keys[0] = keys[1][1]
+                elif keys[0] == keys[1][1]:
+                    keys[0] = keys[1][2]
+                else:
+                    keys[0] = keys[1][0]
+                    
+                client = Groq(api_key= os.getenv(keys[0]))
+                chat_completion = client.chat.completions.create(
+                    messages=messages,
+                    model='llama3-8b-8192',
+                )
+                response = chat_completion.choices[0].message.content
+        else:
+            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,keys if client_index != "Claude" else chat_completion
+        
+    except (BadRequestError) as e:
+        
+        model_id = "meta-llama/Meta-Llama-3-70B-Instruct"
+        access_token = os.getenv("HUGGINGFACE_SPLITFILES_API_KEY")
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id,
+            padding_side="left",
+            token = access_token
+        )
+
+        user_input_tokenized = tokenizer.encode(user_input)
+        messages = []
+
+        while len(user_input_tokenized) > max_token:
+
+            user_input_divided = tokenizer.decode(user_input_tokenized[:max_token])
+            messages.append([
+            {
+                "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": user_input_divided,
+            }])
+
+            user_input_tokenized = user_input_tokenized[max_token:]
+
+        responses = []
+
+        print(len(messages))
+        for msg in messages:
+
+            responses.append(client.chat.completions.create(
+              model=model_map["Perplexity Llama3 70b"],
+              messages=msg
+            ))
+
+        response = ""
+        for resp in responses:
+            response += " " + resp.choices[0].message.content
+
+        return response
+
+    except (RateLimitError) as e:
+
+        #if model_user in keys:
+            #Swap those keys
+        #    return f()
+
+        #else:
+            #get eepy
+        time.sleep(60)
+        return ask_llm(query, user_input, client_index, user, keys)
+
+    except Exception as e:
+        print(e)
+        return e
 
-    # Return the response, handling the structure specific to Groq and Mistral clients.
-    return chat_completion.choices[0].message.content,keys if client_index != "Claude" else chat_completion