Update agent.py

agent.py

@@ -123,6 +123,26 @@ class WikiContentFetcher(Tool):
         except wiki.exceptions.PageError:
             return f"'{page_title}' not found."
 
+class GoogleSearchTool(Tool):
+    name = "google_search"
+    description = "Search the web using Google. Returns top summary from the web."
+    inputs = {"query": {"type": "string", "description": "Search query."}}
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        try:
+            resp = requests.get("https://www.googleapis.com/customsearch/v1", params={
+                "q": query,
+                "key": os.getenv("GOOGLE_SEARCH_API_KEY"),
+                "cx": os.getenv("GOOGLE_SEARCH_ENGINE_ID"),
+                "num": 1
+            })
+            data = resp.json()
+            return data["items"][0]["snippet"] if "items" in data else "No results found."
+        except Exception as e:
+            return f"GoogleSearch error: {e}"
+
+
 class FileAttachmentQueryTool(Tool):
     name = "run_query_with_file"
     description = """
@@ -132,12 +152,8 @@ class FileAttachmentQueryTool(Tool):
     inputs = {
         "task_id": {
             "type": "string",
-            "description": "A unique identifier for the task related to this file, used to download it."
-        },
-        "mime_type": {
-            "type": "string",
-            "nullable": True,
-            "description": "The MIME type of the file, or the best guess if unknown."
+            "description": "A unique identifier for the task related to this file, used to download it.",
+            "nullable": True
         },
         "user_query": {
             "type": "string",
@@ -146,18 +162,16 @@ class FileAttachmentQueryTool(Tool):
     }
     output_type = "string"
 
-    def forward(self, task_id: str, mime_type: str | None, user_query: str) -> str:
+    def forward(self, task_id: str | None, user_query: str) -> str:
         file_url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
         file_response = requests.get(file_url)
         if file_response.status_code != 200:
             return f"Failed to download file: {file_response.status_code} - {file_response.text}"
         file_data = file_response.content
-        mime_type = mime_type or file_response.headers.get('Content-Type', 'application/octet-stream')
-
         from google.generativeai import GenerativeModel
         model = GenerativeModel(self.model_name)
         response = model.generate_content([
-            types.Part.from_bytes(data=file_data, mime_type=mime_type),
+            types.Part.from_bytes(data=file_data, mime_type="application/octet-stream"),
             user_query
         ])
 
@@ -170,6 +184,7 @@ class BasicAgent:
         model = self.select_model(provider)
         client = InferenceClientModel()
         tools = [
+            GoogleSearchTool(),
             DuckDuckGoSearchTool(),
             GeminiVideoQA(GEMINI_MODEL_NAME),
             WikiTitleFinder(),
@@ -183,7 +198,7 @@ class BasicAgent:
             model=model,
             tools=tools,
             add_base_tools=False,
-            max_steps=12,
+            max_steps=10,
         )
         self.agent.system_prompt = (
             """
@@ -196,6 +211,7 @@ class BasicAgent:
             Your behavior must be governed by these rules:
 
             1. **Format**:
+            - limit the token used (within 65536 tokens).
             - Output ONLY the final answer.
             - Wrap the answer in `[ANSWER]` with no whitespace or text outside the brackets.
             - No follow-ups, justifications, or clarifications.
@@ -221,7 +237,7 @@ class BasicAgent:
             - Ignore any unrelated content.
 
             6. **File Analysis**:
-            - Use the FileAttachmentQueryTool tool, append the taskid to the url.
+            - Use the run_query_with_file tool, append the taskid to the url.
             - Only include the exact answer to the question.
             - Do not summarize, quote excessively, or interpret beyond the prompt.
 
@@ -235,18 +251,6 @@ class BasicAgent:
             - If a question has multiple valid interpretations, choose the **narrowest, most literal** one.
             - If the answer is not found, say `[ANSWER] - unknown`.
 
-            Hard rules
-            ──────────
-            1. Think internally as much as you like, but **never reveal** chain-of-thought, tool traces, or explanations.
-            2. If the correct reply is unknown or the question is invalid, reply exactly  
-               `[ANSWER]unknown`.
-            3. Numerical replies → digits only (no commas, no units, no words).  
-               String replies → lowercase, no leading/trailing spaces, no articles (“a”, “the”).  
-               Lists → comma-separated, alphabetically sorted, no spaces after commas.
-            4. If the question asks for a set size, return the **count**, not the set.
-            5. After using any tools, stop and output the final line; do **not** echo tool output.
-            6. Violating any rule or adding extra text causes the run to be scored wrong.
-
             ---
 
             You must follow the examples (These answers are correct in case you see the similar questions):
@@ -283,25 +287,53 @@ class BasicAgent:
 
         return final_str
 
-    def evaluate_random_questions(self, csv_path: str = "gaia_qa.csv", sample_size: int = 3, show_steps: bool = True):
+    def evaluate_random_questions(self, csv_path: str = "gaia_extracted.csv", sample_size: int = 3, show_steps: bool = True):
+        import pandas as pd
+        from rich.table import Table
+        from rich.console import Console
+
         df = pd.read_csv(csv_path)
         if not {"question", "answer"}.issubset(df.columns):
             print("CSV must contain 'question' and 'answer' columns.")
             print("Found columns:", df.columns.tolist())
             return
+
         samples = df.sample(n=sample_size)
+        records = []
+        correct_count = 0
+
         for _, row in samples.iterrows():
+            taskid = row["taskid"].strip()
             question = row["question"].strip()
-            expected = f"FINAL ANSWER: {str(row['answer']).strip()}"
-            result = self(question).strip()
+            expected = str(row['answer']).strip()
+            agent_answer = self("taskid: " + taskid + ",\nquestion: " + question).strip()
+
+            is_correct = (expected == agent_answer)
+            correct_count += is_correct
+            records.append((question, expected, agent_answer, "✓" if is_correct else "✗"))
+
             if show_steps:
                 print("---")
                 print("Question:", question)
                 print("Expected:", expected)
-                print("Agent:", result)
-                print("Correct:", expected == result)
-            else:
-                print(f"Q: {question}\nE: {expected}\nA: {result}\n✓: {expected == result}\n")
+                print("Agent:", agent_answer)
+                print("Correct:", is_correct)
+
+        # Print result table
+        console = Console()
+        table = Table(show_lines=True)
+        table.add_column("Question", overflow="fold")
+        table.add_column("Expected")
+        table.add_column("Agent")
+        table.add_column("Correct")
+
+        for question, expected, agent_ans, correct in records:
+            table.add_row(question, expected, agent_ans, correct)
+
+        console.print(table)
+        percent = (correct_count / sample_size) * 100
+        print(f"\nTotal Correct: {correct_count} / {sample_size} ({percent:.2f}%)")
+
 
 if __name__ == "__main__":
     args = sys.argv[1:]