Update app.py

app.py

@@ -14,124 +14,76 @@ GROQ_KEY = os.environ['GROQ_KEY']
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
-class BasicAgent:
+
+class LLaMaAgent:
     def __init__(self):
-        print("BasicAgent initialized.")
-        self.client = Groq(api_key=GROQ_KEY)
-        
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        try:
-            model = LiteLLMModel(
+        self.model = model = LiteLLMModel(
                 "llama-3.3-70b-versatile",
                 api_base="https://api.groq.com/openai/v1",
                 api_key=GROQ_KEY,
             )
-            model.flatten_messages_as_text = True
-    
-            agent = CodeAgent(
-                tools=[DuckDuckGoSearchTool(), FinalAnswerTool(), VisitWebpageTool(), PythonInterpreterTool()],
-                model=model,
-            )
-            response = agent.run(question)
-            
-        except Exception:
-            try:
-                message = [  
-                {
-                    "role": "user",
-                    "content": question
-                }]
-                completion = self.client.chat.completions.create(
-                    messages=message,
-                    model="compound-beta",
-                )
-                answer = completion.choices[0].message.content
-                print(f"First answer: {answer}")
-                message=[
-                    {
-                        "role": "system",
-                        "content": """
-                        You are an expert in identifying and extracting definitive answers. Your sole task is to analyze the provided text, which is an agent's response, and extract only the conclusive final answer to the original user's query.
-        
-                        Output only this core answer. Do not include any explanations, pleasantries, introductory phrases, or any surrounding text.
-        
-                        Here The Examples:
-                        
-                        Input: ... Final answer: 12 ...
-                        You should output: 12
-        
-                        Input: $\\boxed{b,c,e}$
-                        Output: b, c, e
-        
-                        Input: Jan
-                        Output: Jan
-        
-                        Input: The Yankee with the most walks in the 1977 regular season was Reggie Jackson, with 58 walks. He had 357 at bats that season.
-                        Output: 357
-        
-                        Input: broccoli, bell pepper, celery, fresh basil, green beans, lettuce, sweet potatoes, zucchini
-                        Output: broccoli, bell pepper, celery, fresh basil, green beans, lettuce, sweet potatoes, zucchini
-                        """
-                    },
-                    {
-                        "role": "user",
-                        "content": answer
-                    }
-                ]
-                completion = self.client.chat.completions.create(
-                    messages=message,
-                    model="llama-3.1-8b-instant",
-                )
-                response = completion.choices[0].message.content
-            except Exception:
-                message = [  
-                {
-                    "role": "user",
-                    "content": question
-                }]
-                completion = self.client.chat.completions.create(
-                    messages=message,
-                    model="compound-beta-mini",
-                )
-                answer = completion.choices[0].message.content
-                print(f"First answer: {answer}")
-                message=[
-                    {
-                        "role": "system",
-                        "content": """
-                        You are an expert in identifying and extracting definitive answers. Your sole task is to analyze the provided text, which is an agent's response, and extract only the conclusive final answer to the original user's query.
-        
-                        Output only this core answer. Do not include any explanations, pleasantries, introductory phrases, or any surrounding text.
-        
-                        Here The Examples:
-                        
-                        Input: ... Final answer: 12 ...
-                        You should output: 12
-        
-                        Input: $\\boxed{b,c,e}$
-                        Output: b, c, e
-        
-                        Input: Jan
-                        Output: Jan
+        self.model.flatten_messages_as_text = True
+
+        self.agent = CodeAgent(
+            tools=[DuckDuckGoSearchTool(), FinalAnswerTool(), VisitWebpageTool(), PythonInterpreterTool()],
+            model=model,
+        )
         
-                        Input: The Yankee with the most walks in the 1977 regular season was Reggie Jackson, with 58 walks. He had 357 at bats that season.
-                        Output: 357
+     def __call__(self, question: str) -> str:
+         response = self.agent.run(question)
+         return response
+
+class CompoundAgent:
+    def __init__(self):
+        self.client = Groq(api_key=GROQ_KEY)
         
-                        Input: broccoli, bell pepper, celery, fresh basil, green beans, lettuce, sweet potatoes, zucchini
-                        Output: broccoli, bell pepper, celery, fresh basil, green beans, lettuce, sweet potatoes, zucchini
-                        """
-                    },
-                    {
-                        "role": "user",
-                        "content": answer
-                    }
-                ]
-                completion = self.client.chat.completions.create(
-                    messages=message,
-                    model="llama-3.1-8b-instant",
-                )
-                response = completion.choices[0].message.content
+    def __call__(self, question: str) -> str:
+        message = [  
+        {
+            "role": "user",
+            "content": question
+        }]
+        completion = self.client.chat.completions.create(
+            messages=message,
+            model="compound-beta",
+        )
+        answer = completion.choices[0].message.content
+        message=[
+            {
+                "role": "system",
+                "content": """
+                You are an expert in identifying and extracting definitive answers. Your sole task is to analyze the provided text, which is an agent's response, and extract only the conclusive final answer to the original user's query.
+
+                Output only this core answer. Do not include any explanations, pleasantries, introductory phrases, or any surrounding text.
+
+                Here The Examples:
+                
+                Input: ... Final answer: 12 ...
+                You should output: 12
+
+                Input: $\\boxed{b,c,e}$
+                Output: b, c, e
+
+                Input: Jan
+                Output: Jan
+
+                Input: The Yankee with the most walks in the 1977 regular season was Reggie Jackson, with 58 walks. He had 357 at bats that season.
+                Output: 357
+
+                Input: broccoli, bell pepper, celery, fresh basil, green beans, lettuce, sweet potatoes, zucchini
+                Output: broccoli, bell pepper, celery, fresh basil, green beans, lettuce, sweet potatoes, zucchini
+                """
+            },
+            {
+                "role": "user",
+                "content": answer
+            }
+        ]
+        completion = self.client.chat.completions.create(
+            messages=message,
+            model="llama-3.1-8b-instant",
+        )
+        response = completion.choices[0].message.content
         return response
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
@@ -155,7 +107,8 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
 
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
-        agent = BasicAgent()
+        llama = LLaMaAgent()
+        compound = CompoundAgent()
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None
@@ -191,12 +144,14 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
     for item in questions_data:
         task_id = item.get("task_id")
         question_text = item.get("question")
-        time.sleep(45)
+        time.sleep(30)
         if not task_id or question_text is None:
             print(f"Skipping item with missing task_id or question: {item}")
             continue
         try:
-            submitted_answer = agent(question_text)
+            submitted_answer = llama(question_text)
+            if submitted_answer.contains('token'):
+                submitted_answer = compound(question_text)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         except Exception as e: