Update main.py

main.py

@@ -88,23 +88,46 @@ USER QUERY:{user_input}
 TOPICS TO BE EXCLUDED:{excluded_topics}
 YOUR OUTPUT SHOULD CONSIST ONLY A VALID PYTHON LIST, DO NOT ADD ADDITIONAL COMMENTS
 """
+### ------LLM CONFIG-------- ###
 
-def together_response(message, model = "meta-llama/Llama-3-8b-chat-hf", SysPrompt = SysPromptDefault, temperature=0.2):
-    client = OpenAI(
+together_client = OpenAI(
         api_key=TOGETHER_API_KEY, 
         base_url="https://together.hconeai.com/v1", 
         default_headers={ "Helicone-Auth": f"Bearer {HELICON_API_KEY}"})
+
+groq_client = OpenAI(
+        api_key=GROQ_API_KEY, 
+        base_url="https://groq.hconeai.com/openai/v1", 
+        default_headers={ "Helicone-Auth": f"Bearer {HELICON_API_KEY}"})
+
+# Groq model names
+llm_default_small = "llama3-8b-8192"
+llm_default_medium = "llama3-70b-8192"
+
+# Together Model names (fallback)
+llm_fallback_small = "meta-llama/Llama-3-8b-chat-hf"
+llm_fallback_medium = "meta-llama/Llama-3-70b-chat-hf"
+
+### ------END OF LLM CONFIG-------- ###
+
+def together_response(message, model = llm_default_small, SysPrompt = SysPromptDefault, temperature=0.2, frequency_penalty =0.1, max_tokens= 2000):
     
     messages=[{"role": "system", "content": SysPrompt},{"role": "user", "content": message}]
+    params = {
+      "model": model,
+      "messages": messages,
+      "temperature": temperature,
+      "frequency_penalty": frequency_penalty,
+      "max_tokens": max_tokens
+    }
+    try:
+      response = groq_client.chat.completions.create(**params)
+      return response.choices[0].message.content
     
-    response = client.chat.completions.create(
-        model=model,
-        messages=messages,
-        temperature=temperature,
-        frequency_penalty=0.2
-        )
-    return response.choices[0].message.content
-
+    except Exception as e:
+      params["model"] = llm_fallback_small if model == llm_default_small else llm_fallback_medium 
+      response = together_client.chat.completions.create(**params)
+      return response.choices[0].message.content
 
 
 
@@ -127,20 +150,20 @@ def json_from_text(text):
     
 def generate_topics(user_input, num_topics, previous_queries):
     prompt = prompt_topics.format(user_input=user_input, num_topics=num_topics)
-    response_topics = together_response(prompt, model="meta-llama/Llama-3-70b-chat-hf", SysPrompt=SysPromptList, temperature=1)
+    response_topics = together_response(prompt, model=llm_default_medium, SysPrompt=SysPromptList, temperature=1)
     subtopics = json_from_text(response_topics)
     return subtopics
 
 def generate_subtopics(main_task,user_input,num_topics,excluded_topics):
     excluded_topics = ",".join(excluded_topics)
     prompt = prompt_subtopics.format(main_task = main_task,user_input=user_input, num_topics=num_topics, excluded_topics=excluded_topics)
-    response_topics = together_response(prompt, model="meta-llama/Llama-3-70b-chat-hf", SysPrompt=SysPromptList, temperature=1)
+    response_topics = together_response(prompt, model=llm_default_medium, SysPrompt=SysPromptList, temperature=1)
     subtopics = json_from_text(response_topics)
     return subtopics
 
 def generate_report(topic, description):
     prompt = f"""create a detailed report on: {topic} by following the instructions: {description}"""
-    md_report = together_response(prompt, model = "meta-llama/Llama-3-70b-chat-hf", SysPrompt = SysPromptMdOffline)
+    md_report = together_response(prompt, model = llm_default_medium, SysPrompt = SysPromptMdOffline)
     return md_to_html(md_report)
 
 def get_images(query, num_results):
@@ -218,7 +241,7 @@ async def generate_recommendations(input: RecommendationInput):
 
     response_topics = json_from_text(
             together_response(
-                prompt, model="meta-llama/Llama-3-8b-chat-hf", SysPrompt=SysPromptList,temperature=1
+                prompt, model=llm_default_small, SysPrompt=SysPromptList,temperature=1
             )
         )
     return {"recommendations": response_topics}