Refactor main_v2.py to update task formatting for dual answer requests, enhancing response structure. Implement error handling for JSON parsing in agent results, ensuring robust output. Add unit tests in test_questions.py to validate succinct answer accuracy against expected values. Remove unused extract_final_answer utility from utils.py, streamlining the codebase.

.vscode/tasks.json

@@ -4,7 +4,7 @@
     "version": "2.0.0",
     "tasks": [
         {
-            "label": "serve",
+            "label": "prerun",
             "type": "shell",
             "command": "uv run python -m phoenix.server.main serve"
         },
@@ -12,6 +12,11 @@
             "label": "run",
             "type": "shell",
             "command": "uv run python -m main_v2"
+        },
+        {
+            "label": "test",
+            "type": "shell",
+            "command": "uv run python -m unittest test_questions.py"
         }
     ]
-}
+}

app.py

@@ -30,7 +30,7 @@ class BasicAgent:
         # question_counter += 1
 
         # if question_counter > 1:
-            # return "This is a default answer."
+        # return "This is a default answer."
 
         # fixed_answer = "This is a default answer."
         # print(f"Agent returning fixed answer: {fixed_answer}")

main.py

@@ -14,14 +14,14 @@ from langgraph.graph import END, START, StateGraph
 from openinference.instrumentation.smolagents import SmolagentsInstrumentor
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from phoenix.otel import register
-from prompts import MANAGER_SYSTEM_PROMPT
 from smolagents import CodeAgent, LiteLLMModel
 from smolagents.memory import ActionStep, FinalAnswerStep
 from smolagents.monitoring import LogLevel
+from utils import extract_final_answer
 
 from agents import create_data_analysis_agent, create_media_agent, create_web_agent
+from prompts import MANAGER_SYSTEM_PROMPT
 from tools import perform_calculation, web_search
-from utils import extract_final_answer
 
 litellm._turn_on_debug()
 

main_v2.py

@@ -13,7 +13,6 @@ from smolagents import CodeAgent, LiteLLMModel
 from smolagents.monitoring import LogLevel
 
 from tools.smart_search.tool import SmartSearchTool
-from utils import extract_final_answer
 
 _disable_debugging()
 
@@ -78,20 +77,32 @@ agent.visualize()
 
 
 def main(task: str):
-    # Format the task with GAIA-style instructions
-#     gaia_task = f"""Instructions:
-# 1. Your response must contain ONLY the answer to the question, nothing else
-# 2. Do not repeat the question or any part of it
-# 3. Do not include any explanations, reasoning, or context
-# 4. Do not include source attribution or references
-# 5. Do not use phrases like "The answer is" or "I found that"
-# 6. Do not include any formatting, bullet points, or line breaks
-# 7. If the answer is a number, return only the number
-# 8. If the answer requires multiple items, separate them with commas
-# 9. If the answer requires ordering, maintain the specified order
-# 10. Use the most direct and succinct form possible
-
-# {task}"""
+    # Format the task to request both succinct and verbose answers
+    formatted_task = f"""Please provide two answers to the following question:
+
+1. A succinct answer that follows these rules:
+   - Contains ONLY the answer, nothing else
+   - Does not repeat the question
+   - Does not include explanations, reasoning, or context
+   - Does not include source attribution or references
+   - Does not use phrases like "The answer is" or "I found that"
+   - Does not include formatting, bullet points, or line breaks
+   - If the answer is a number, return only the number
+   - If the answer requires multiple items, separate them with commas
+   - If the answer requires ordering, maintain the specified order
+   - Uses the most direct and succinct form possible
+
+2. A verbose answer that includes:
+   - The complete answer with all relevant details
+   - Explanations and reasoning
+   - Context and background information
+   - Source attribution where appropriate
+
+Question: {task}
+
+Please format your response as a JSON object with two keys:
+- "succinct_answer": The concise answer following the rules above
+- "verbose_answer": The detailed explanation with context"""
 
     result = agent.run(
         additional_args=None,
@@ -99,13 +110,23 @@ def main(task: str):
         max_steps=3,
         reset=True,
         stream=False,
-        task=task,
-        # task=gaia_task,
+        task=formatted_task,
     )
 
-    logger.info(f"Result: {result}")
+    # Parse the result into a dictionary
+    try:
+        import json
 
-    return extract_final_answer(result)
+        # Find the JSON object in the response
+        json_str = result[result.find("{") : result.rfind("}") + 1]
+        parsed_result = json.loads(json_str)
+    except (ValueError, AttributeError) as e:
+        logger.error(f"Error parsing result: {e}")
+        # If parsing fails, return the raw result
+        return result
+
+    logger.info(f"Result: {parsed_result}")
+    return parsed_result["succinct_answer"]
 
 
 if __name__ == "__main__":
@@ -119,7 +140,7 @@ if __name__ == "__main__":
     response.raise_for_status()
     questions_data = response.json()
 
-    for question_data in questions_data[:1]:
+    for question_data in questions_data[:2]:
         file_name = question_data["file_name"]
         level = question_data["Level"]
         question = question_data["question"]

test_questions.py

@@ -0,0 +1,35 @@
+import unittest
+import requests
+from main_v2 import main
+
+class TestQuestions(unittest.TestCase):
+    def setUp(self):
+        self.api_url = "https://agents-course-unit4-scoring.hf.space"
+        self.questions_url = f"{self.api_url}/questions"
+        
+        # Get questions from the API
+        response = requests.get(self.questions_url, timeout=15)
+        response.raise_for_status()
+        self.questions = response.json()
+        
+        # Expected answers for each question
+        self.expected_answers = {
+            "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia.": "3",
+            # Add more expected answers as you verify them
+        }
+
+    def test_questions(self):
+        """Test each question and verify the succinct answer matches the expected value."""
+        for question_data in self.questions:
+            question = question_data["question"]
+            if question in self.expected_answers:
+                expected_answer = self.expected_answers[question]
+                actual_answer = main(question)
+                self.assertEqual(
+                    actual_answer,
+                    expected_answer,
+                    f"Question: {question}\nExpected: {expected_answer}\nGot: {actual_answer}"
+                )
+
+if __name__ == "__main__":
+    unittest.main()

tools/smart_search/tool.py

@@ -1,10 +1,10 @@
 import logging
 import re
-from typing import List, Dict, Optional
+from typing import Optional
+
+import requests
 from smolagents import Tool
 from smolagents.default_tools import DuckDuckGoSearchTool
-import requests
-from bs4 import BeautifulSoup
 
 logger = logging.getLogger(__name__)
 
@@ -12,7 +12,12 @@ logger = logging.getLogger(__name__)
 class SmartSearchTool(Tool):
     name = "smart_search"
     description = """A smart search tool that searches Wikipedia for information."""
-    inputs = {"query": {"type": "string", "description": "The search query to find information"}}
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The search query to find information",
+        }
+    }
     output_type = "string"
 
     def __init__(self):
@@ -20,30 +25,30 @@ class SmartSearchTool(Tool):
         self.web_search_tool = DuckDuckGoSearchTool(max_results=1)
         self.api_url = "https://en.wikipedia.org/w/api.php"
         self.headers = {
-            'User-Agent': 'SmartSearchTool/1.0 (https://github.com/yourusername/yourproject; [email protected])'
+            "User-Agent": "SmartSearchTool/1.0 (https://github.com/yourusername/yourproject; [email protected])"
         }
 
     def get_wikipedia_page(self, title: str) -> Optional[str]:
         """Get the raw wiki markup of a Wikipedia page."""
         try:
             params = {
-                'action': 'query',
-                'prop': 'revisions',
-                'rvprop': 'content',
-                'rvslots': 'main',
-                'format': 'json',
-                'titles': title,
-                'redirects': 1
+                "action": "query",
+                "prop": "revisions",
+                "rvprop": "content",
+                "rvslots": "main",
+                "format": "json",
+                "titles": title,
+                "redirects": 1,
             }
             response = requests.get(self.api_url, params=params, headers=self.headers)
             response.raise_for_status()
             data = response.json()
-            
+
             # Extract page content
-            pages = data.get('query', {}).get('pages', {})
+            pages = data.get("query", {}).get("pages", {})
             for page_id, page_data in pages.items():
-                if 'revisions' in page_data:
-                    return page_data['revisions'][0]['slots']['main']['*']
+                if "revisions" in page_data:
+                    return page_data["revisions"][0]["slots"]["main"]["*"]
             return None
         except Exception as e:
             logger.error(f"Error getting Wikipedia page: {e}")
@@ -52,117 +57,126 @@ class SmartSearchTool(Tool):
     def clean_wiki_content(self, content: str) -> str:
         """Clean Wikipedia content by removing markup and formatting."""
         # Remove citations
-        content = re.sub(r'\[\d+\]', '', content)
+        content = re.sub(r"\[\d+\]", "", content)
         # Remove edit links
-        content = re.sub(r'\[edit\]', '', content)
+        content = re.sub(r"\[edit\]", "", content)
         # Remove file links
-        content = re.sub(r'\[\[File:.*?\]\]', '', content)
+        content = re.sub(r"\[\[File:.*?\]\]", "", content)
         # Convert links to just text
-        content = re.sub(r'\[\[(?:[^|\]]*\|)?([^\]]+)\]\]', r'\1', content)
+        content = re.sub(r"\[\[(?:[^|\]]*\|)?([^\]]+)\]\]", r"\1", content)
         # Remove HTML comments
-        content = re.sub(r'<!--.*?-->', '', content, flags=re.DOTALL)
+        content = re.sub(r"<!--.*?-->", "", content, flags=re.DOTALL)
         # Remove templates
-        content = re.sub(r'\{\{.*?\}\}', '', content)
+        content = re.sub(r"\{\{.*?\}\}", "", content)
         # Remove small tags
-        content = re.sub(r'<small>.*?</small>', '', content)
+        content = re.sub(r"<small>.*?</small>", "", content)
         # Normalize whitespace
-        content = re.sub(r'\n\s*\n', '\n\n', content)
+        content = re.sub(r"\n\s*\n", "\n\n", content)
         return content.strip()
 
     def format_wiki_table(self, table_content: str) -> str:
         """Format a Wikipedia table into readable text."""
         # Split into rows
-        rows = table_content.strip().split('\n')
+        rows = table_content.strip().split("\n")
         formatted_rows = []
         current_row = []
-        
+
         for row in rows:
             # Skip empty rows and table structure markers
-            if not row.strip() or row.startswith('|-') or row.startswith('|+'):
+            if not row.strip() or row.startswith("|-") or row.startswith("|+"):
                 if current_row:
-                    formatted_rows.append('\t'.join(current_row))
+                    formatted_rows.append("\t".join(current_row))
                     current_row = []
                 continue
-                
+
             # Extract cells
             cells = []
             # Split the row into cells using | or ! as separators
-            cell_parts = re.split(r'[|!]', row)
+            cell_parts = re.split(r"[|!]", row)
             for cell in cell_parts[1:]:  # Skip the first empty part
                 # Clean up the cell content
                 cell = cell.strip()
                 # Remove any remaining markup
-                cell = re.sub(r'<.*?>', '', cell)  # Remove HTML tags
-                cell = re.sub(r'\[\[.*?\|(.*?)\]\]', r'\1', cell)  # Convert links
-                cell = re.sub(r'\[\[(.*?)\]\]', r'\1', cell)  # Convert simple links
-                cell = re.sub(r'\{\{.*?\}\}', '', cell)  # Remove templates
-                cell = re.sub(r'<small>.*?</small>', '', cell)  # Remove small tags
-                cell = re.sub(r'rowspan="\d+"', '', cell)  # Remove rowspan
-                cell = re.sub(r'colspan="\d+"', '', cell)  # Remove colspan
-                cell = re.sub(r'class=".*?"', '', cell)  # Remove class attributes
-                cell = re.sub(r'style=".*?"', '', cell)  # Remove style attributes
-                cell = re.sub(r'align=".*?"', '', cell)  # Remove align attributes
-                cell = re.sub(r'width=".*?"', '', cell)  # Remove width attributes
-                cell = re.sub(r'bgcolor=".*?"', '', cell)  # Remove bgcolor attributes
-                cell = re.sub(r'valign=".*?"', '', cell)  # Remove valign attributes
-                cell = re.sub(r'border=".*?"', '', cell)  # Remove border attributes
-                cell = re.sub(r'cellpadding=".*?"', '', cell)  # Remove cellpadding attributes
-                cell = re.sub(r'cellspacing=".*?"', '', cell)  # Remove cellspacing attributes
-                cell = re.sub(r'<ref.*?</ref>', '', cell)  # Remove references
-                cell = re.sub(r'<ref.*?/>', '', cell)  # Remove empty references
-                cell = re.sub(r'<br\s*/?>', ' ', cell)  # Replace line breaks with spaces
-                cell = re.sub(r'\s+', ' ', cell)  # Normalize whitespace
+                cell = re.sub(r"<.*?>", "", cell)  # Remove HTML tags
+                cell = re.sub(r"\[\[.*?\|(.*?)\]\]", r"\1", cell)  # Convert links
+                cell = re.sub(r"\[\[(.*?)\]\]", r"\1", cell)  # Convert simple links
+                cell = re.sub(r"\{\{.*?\}\}", "", cell)  # Remove templates
+                cell = re.sub(r"<small>.*?</small>", "", cell)  # Remove small tags
+                cell = re.sub(r'rowspan="\d+"', "", cell)  # Remove rowspan
+                cell = re.sub(r'colspan="\d+"', "", cell)  # Remove colspan
+                cell = re.sub(r'class=".*?"', "", cell)  # Remove class attributes
+                cell = re.sub(r'style=".*?"', "", cell)  # Remove style attributes
+                cell = re.sub(r'align=".*?"', "", cell)  # Remove align attributes
+                cell = re.sub(r'width=".*?"', "", cell)  # Remove width attributes
+                cell = re.sub(r'bgcolor=".*?"', "", cell)  # Remove bgcolor attributes
+                cell = re.sub(r'valign=".*?"', "", cell)  # Remove valign attributes
+                cell = re.sub(r'border=".*?"', "", cell)  # Remove border attributes
+                cell = re.sub(
+                    r'cellpadding=".*?"', "", cell
+                )  # Remove cellpadding attributes
+                cell = re.sub(
+                    r'cellspacing=".*?"', "", cell
+                )  # Remove cellspacing attributes
+                cell = re.sub(r"<ref.*?</ref>", "", cell)  # Remove references
+                cell = re.sub(r"<ref.*?/>", "", cell)  # Remove empty references
+                cell = re.sub(
+                    r"<br\s*/?>", " ", cell
+                )  # Replace line breaks with spaces
+                cell = re.sub(r"\s+", " ", cell)  # Normalize whitespace
                 cells.append(cell)
-            
+
             if cells:
                 current_row.extend(cells)
-        
+
         if current_row:
-            formatted_rows.append('\t'.join(current_row))
-        
+            formatted_rows.append("\t".join(current_row))
+
         if formatted_rows:
-            return '\n'.join(formatted_rows)
-        return ''
+            return "\n".join(formatted_rows)
+        return ""
 
     def extract_wikipedia_title(self, search_result: str) -> Optional[str]:
         """Extract Wikipedia page title from search result."""
         # Look for Wikipedia links in the format [Title - Wikipedia](url)
-        wiki_match = re.search(r'\[([^\]]+)\s*-\s*Wikipedia\]\(https://en\.wikipedia\.org/wiki/[^)]+\)', search_result)
+        wiki_match = re.search(
+            r"\[([^\]]+)\s*-\s*Wikipedia\]\(https://en\.wikipedia\.org/wiki/[^)]+\)",
+            search_result,
+        )
         if wiki_match:
             return wiki_match.group(1).strip()
         return None
 
     def forward(self, query: str) -> str:
         logger.info(f"Starting smart search for query: {query}")
-        
+
         # First do a web search to find the Wikipedia page
         search_result = self.web_search_tool.forward(query)
         logger.info(f"Web search results: {search_result[:100]}...")
-        
+
         # Extract Wikipedia page title from search results
         wiki_title = self.extract_wikipedia_title(search_result)
         if not wiki_title:
             return f"Could not find Wikipedia page in search results for '{query}'."
-        
+
         # Get Wikipedia page content
         page_content = self.get_wikipedia_page(wiki_title)
         if not page_content:
             return f"Could not find Wikipedia page for '{wiki_title}'."
-        
+
         # Format tables and content
         formatted_content = []
         current_section = []
         in_table = False
         table_content = []
-        
-        for line in page_content.split('\n'):
-            if line.startswith('{|'):
+
+        for line in page_content.split("\n"):
+            if line.startswith("{|"):
                 in_table = True
                 table_content = [line]
-            elif line.startswith('|}'):
+            elif line.startswith("|}"):
                 in_table = False
                 table_content.append(line)
-                formatted_table = self.format_wiki_table('\n'.join(table_content))
+                formatted_table = self.format_wiki_table("\n".join(table_content))
                 if formatted_table:
                     current_section.append(formatted_table)
             elif in_table:
@@ -171,49 +185,49 @@ class SmartSearchTool(Tool):
                 if line.strip():
                     current_section.append(line)
                 elif current_section:
-                    formatted_content.append('\n'.join(current_section))
+                    formatted_content.append("\n".join(current_section))
                     current_section = []
-        
+
         if current_section:
-            formatted_content.append('\n'.join(current_section))
-        
+            formatted_content.append("\n".join(current_section))
+
         # Clean and return the formatted content
-        cleaned_content = self.clean_wiki_content('\n\n'.join(formatted_content))
+        cleaned_content = self.clean_wiki_content("\n\n".join(formatted_content))
         return f"Wikipedia content for '{wiki_title}':\n\n{cleaned_content}"
 
 
 def main(query: str) -> str:
     """
     Test function to run the SmartSearchTool directly.
-    
+
     Args:
         query: The search query to test
-        
+
     Returns:
         The search results
     """
     # Configure logging
     logging.basicConfig(
         level=logging.INFO,
-        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
     )
-    
+
     # Create and run the tool
     tool = SmartSearchTool()
     result = tool.forward(query)
-    
+
     # Print the result
     print("\nSearch Results:")
     print("-" * 80)
     print(result)
     print("-" * 80)
-    
+
     return result
 
 
 if __name__ == "__main__":
     import sys
-    
+
     if len(sys.argv) > 1:
         query = " ".join(sys.argv[1:])
         main(query)

utils.py

@@ -1,70 +0,0 @@
-import re
-from typing import Union
-
-
-def extract_final_answer(result: Union[str, dict]) -> str:
-    """
-    Extract the final answer from the agent's result, removing explanations.
-    GAIA requires concise, properly formatted answers.
-
-    Args:
-        result: The full result from the agent, either a string or a dictionary
-
-    Returns:
-        Extracted final answer
-    """
-    # Handle dictionary input
-    if isinstance(result, dict):
-        # Check for final_answer_text first (from agent output)
-        if "final_answer_text" in result:
-            return str(result["final_answer_text"])
-        # Fall back to final_answer key
-        if "final_answer" in result:
-            return str(result["final_answer"])
-        return "No final answer found in result"
-
-    # Handle string input (original logic)
-    # First check if there's a specific final_answer marker
-    if "final_answer(" in result:
-        # Try to extract the answer from final_answer call
-        pattern = r"final_answer\(['\"](.*?)['\"]\)"
-        matches = re.findall(pattern, result)
-        if matches:
-            return matches[-1]  # Return the last final_answer if multiple exist
-
-    # If no final_answer marker, look for lines that might contain the answer
-    lines = result.strip().split("\n")
-
-    # Check for typical patterns indicating a final answer
-    for line in reversed(lines):  # Start from the end
-        line = line.strip()
-
-        # Skip empty lines
-        if not line:
-            continue
-
-        # Look for patterns like "Answer:", "Final answer:", etc.
-        if re.match(r"^(answer|final answer|result):?\s+", line.lower()):
-            return line.split(":", 1)[1].strip()
-
-        # Check for answers that are comma-separated lists (common in GAIA)
-        if (
-            "," in line
-            and len(line.split(",")) > 1
-            and not line.startswith("#")
-            and not line.startswith("print(")
-        ):
-            # It might be a comma-separated list answer
-            return line
-
-    # If no clear answer pattern is found, return the last non-empty line
-    # (often the answer is simply the last output)
-    for line in reversed(lines):
-        if (
-            line.strip()
-            and not line.strip().startswith("#")
-            and not line.strip().startswith("print(")
-        ):
-            return line.strip()
-
-    return "No answer found"