Update app.py

app.py

@@ -1,24 +1,23 @@
 import streamlit as st
 import os
-import openai
 import json
 from dataclasses import dataclass, field
 from typing import List
-from tavily import TavilyClient
+from groq import Groq
 from json.decoder import JSONDecodeError
 from pydantic_settings import BaseSettings
+from tavily import TavilyClient
 
 # Configuration class for API keys and settings
 class Config(BaseSettings):
-    SAMBANOVA_API_KEY: str
-    SAMBANOVA_BASE_URL: str = "https://api.sambanova.ai/v1"
-    LLM_REASONING: str = "DeepSeek-R1-Distill-Llama-70B"
-    LLM_REGULAR: str = "Meta-Llama-3.3-70B-Instruct"
+    GROQ_API_KEY: str
+    LLM_REASONING: str = "llama3-70b-8192"  # Example Groq model for reasoning tasks
+    LLM_REGULAR: str = "llama3-8b-8192"     # Example Groq model for regular tasks
     TAVILY_API_KEY: str
 
 # Load configuration from environment variables
 config = Config(
-    SAMBANOVA_API_KEY=os.environ["SAMBANOVA_API_KEY"],
+    GROQ_API_KEY=os.environ["GROQ_API_KEY"],
     TAVILY_API_KEY=os.environ["TAVILY_API_KEY"]
 )
 
@@ -47,19 +46,24 @@ class State:
 
 # Helper functions
 def remove_reasoning_from_output(output):
+    """Remove reasoning tags from output if present."""
     return output.split("</think>")[-1].strip()
 
 def clean_json_tags(text):
+    """Remove JSON markdown tags from text."""
     return text.replace("```json\n", "").replace("\n```", "")
 
 def clean_markdown_tags(text):
+    """Remove markdown tags from text."""
     return text.replace("```markdown\n", "").replace("\n```", "")
 
 def tavily_search(query, include_raw_content=True, max_results=3):
+    """Perform a web search using Tavily."""
     tavily_client = TavilyClient(api_key=config.TAVILY_API_KEY)
     return tavily_client.search(query, include_raw_content=include_raw_content, max_results=max_results)
 
 def update_state_with_search_results(search_results, idx_paragraph, state):
+    """Update the state with search results."""
     for search_result in search_results["results"]:
         search = Search(url=search_result["url"], content=search_result["raw_content"])
         state.paragraphs[idx_paragraph].research.search_history.append(search)
@@ -94,14 +98,11 @@ Only return the json object, no explanation or additional text.
 
 class ReportStructureAgent:
     def __init__(self, query: str):
-        self.openai_client = openai.OpenAI(
-            api_key=config.SAMBANOVA_API_KEY,
-            base_url=config.SAMBANOVA_BASE_URL
-        )
+        self.groq_client = Groq(api_key=config.GROQ_API_KEY)
         self.query = query
 
     def run(self) -> str:
-        response = self.openai_client.chat.completions.create(
+        response = self.groq_client.chat.completions.create(
             model=config.LLM_REASONING,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT_REPORT_STRUCTURE},
@@ -157,13 +158,10 @@ Only return the json object, no explanation or additional text.
 
 class FirstSearchAgent:
     def __init__(self):
-        self.openai_client = openai.OpenAI(
-            api_key=config.SAMBANOVA_API_KEY,
-            base_url=config.SAMBANOVA_BASE_URL
-        )
+        self.groq_client = Groq(api_key=config.GROQ_API_KEY)
 
     def run(self, message) -> str:
-        response = self.openai_client.chat.completions.create(
+        response = self.groq_client.chat.completions.create(
             model=config.LLM_REGULAR,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT_FIRST_SEARCH},
@@ -216,13 +214,10 @@ Only return the json object, no explanation or additional text.
 
 class FirstSummaryAgent:
     def __init__(self):
-        self.openai_client = openai.OpenAI(
-            api_key=config.SAMBANOVA_API_KEY,
-            base_url=config.SAMBANOVA_BASE_URL
-        )
+        self.groq_client = Groq(api_key=config.GROQ_API_KEY)
 
     def run(self, message) -> str:
-        response = self.openai_client.chat.completions.create(
+        response = self.groq_client.chat.completions.create(
             model=config.LLM_REGULAR,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT_FIRST_SUMMARY},
@@ -281,13 +276,10 @@ Only return the json object, no explanation or additional text.
 
 class ReflectionAgent:
     def __init__(self):
-        self.openai_client = openai.OpenAI(
-            api_key=config.SAMBANOVA_API_KEY,
-            base_url=config.SAMBANOVA_BASE_URL
-        )
+        self.groq_client = Groq(api_key=config.GROQ_API_KEY)
 
     def run(self, message) -> str:
-        response = self.openai_client.chat.completions.create(
+        response = self.groq_client.chat.completions.create(
             model=config.LLM_REGULAR,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT_REFLECTION},
@@ -346,13 +338,10 @@ Only return the json object, no explanation or additional text.
 
 class ReflectionSummaryAgent:
     def __init__(self):
-        self.openai_client = openai.OpenAI(
-            api_key=config.SAMBANOVA_API_KEY,
-            base_url=config.SAMBANOVA_BASE_URL
-        )
+        self.groq_client = Groq(api_key=config.GROQ_API_KEY)
 
     def run(self, message) -> str:
-        response = self.openai_client.chat.completions.create(
+        response = self.groq_client.chat.completions.create(
             model=config.LLM_REGULAR,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT_REFLECTION_SUMMARY},
@@ -399,13 +388,10 @@ Use titles of the paragraphs to create a title for the report.
 
 class ReportFormattingAgent:
     def __init__(self):
-        self.openai_client = openai.OpenAI(
-            api_key=config.SAMBANOVA_API_KEY,
-            base_url=config.SAMBANOVA_BASE_URL
-        )
+        self.groq_client = Groq(api_key=config.GROQ_API_KEY)
 
     def run(self, message) -> str:
-        response = self.openai_client.chat.completions.create(
+        response = self.groq_client.chat.completions.create(
             model=config.LLM_REASONING,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT_REPORT_FORMATTING},
@@ -418,7 +404,7 @@ class ReportFormattingAgent:
         return summary
 
 # Streamlit Web Interface
-st.title("Research Agent")
+st.title("Research Agent with Groq")
 
 query = st.text_input("Enter your research query:", "")