Update agent.py

agent.py

@@ -1,141 +1,90 @@
 import os
-import re
-from openai import OpenAI as OpenAIClient
-from duckduckgo_search import DDGS
+import requests
+import base64
+from langchain_openai import ChatOpenAI
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain.agents import initialize_agent, Tool
+from langchain.agents.agent_types import AgentType
+from langchain.memory import ConversationBufferMemory
+from langchain_core.messages import HumanMessage
 
-def duckduckgo_search(query: str) -> str:
-    try:
-        with DDGS() as ddg:
-            results = ddg.text(query=query, region="wt-wt", max_results=5)
-            bodies = [r.get('body', '') for r in results if r.get('body')]
-            return "\n".join(bodies[:3])
-    except Exception as e:
-        return f"ERROR: {e}"
+DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-def eval_python_code(code: str) -> str:
-    try:
-        return str(eval(code, {"__builtins__": {}}))
-    except Exception as e:
-        return f"ERROR: {e}"
+class BasicAgent:
+    def __init__(self):
+        print("BasicAgent initialized.")
+        # Use DuckDuckGo Search only
+        tools = [
+            Tool(
+                name="DuckDuckGo Search",
+                func=DuckDuckGoSearchRun().run,
+                description="Use this tool to find factual information or recent events."
+            ),
+            Tool(
+                name="Image Analyzer",
+                func=self.describe_image,
+                description="Analyzes and describes what's in an image. Input is an image path."
+            )
+        ]
 
-def format_gaia_answer(answer: str, question: str = "") -> str:
-    """Strict GAIA output, eliminate apologies, extract only answer value."""
-    if not answer:
-        return ""
-    # Remove apologies and anything after
-    answer = re.sub(
-        r'(?i)(I[\' ]?m sorry.*|Unfortunately.*|I cannot.*|I am unable.*|error:.*|no file.*|but.*|however.*|unable to.*|not available.*|if you have access.*|I can\'t.*)',
-        '', answer).strip()
-    # Remove everything after the first period if it's not a list
-    if not ("list" in question or "ingredient" in question or "vegetable" in question):
-        answer = answer.split('\n')[0].split('.')[0]
-    # Remove quotes/brackets
-    answer = answer.strip(' "\'[](),;:')
-    # Only numbers for count questions
-    if re.search(r'how many|number of|albums|at bats|total sales|output', question, re.I):
-        match = re.search(r'(\d+)', answer)
-        if match:
-            return match.group(1)
-    # Only last word for "surname", first for "first name"
-    if "surname" in question:
-        return answer.split()[-1]
-    if "first name" in question:
-        return answer.split()[0]
-    # For code outputs, numbers only
-    if "output" in question and "python" in question:
-        num = re.search(r'(\d+)', answer)
-        return num.group(1) if num else answer
-    # Only country code (3+ uppercase letters or digits)
-    if re.search(r'IOC country code|award number|NASA', question, re.I):
-        code = re.search(r'[A-Z0-9]{3,}', answer)
-        if code:
-            return code.group(0)
-    # For lists: comma-separated, alpha, deduped, merged phrases
-    if "list" in question or "ingredient" in question or "vegetable" in question:
-        items = [x.strip(' "\'') for x in re.split(r'[,\n]', answer) if x.strip()]
-        merged = []
-        skip = False
-        for i, item in enumerate(items):
-            if skip:
-                skip = False
-                continue
-            if i + 1 < len(items) and item in ['sweet', 'green', 'lemon', 'ripe', 'whole', 'fresh', 'bell']:
-                merged.append(f"{item} {items[i+1]}")
-                skip = True
-            else:
-                merged.append(item)
-        merged = [x.lower() for x in merged]
-        merged = sorted(set(merged))
-        return ', '.join(merged)
-    # For chess: algebraic move
-    if "algebraic notation" in question or "chess" in question:
-        move = re.findall(r'[KQRBN]?[a-h]?[1-8]?x?[a-h][1-8][+#]?', answer)
-        if move:
-            return move[-1]
-    return answer.strip(' "\'[](),;:')
+        memory = ConversationBufferMemory(memory_key="chat_history")
+        self.model = ChatOpenAI(model="gpt-4.1-mini", temperature=0)
 
-class GaiaAgent:
-    def __init__(self):
-        self.llm = OpenAIClient(api_key=os.getenv("OPENAI_API_KEY"))
+        self.agent = initialize_agent(
+            tools=tools,
+            llm=self.model,
+            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+            verbose=True,
+            memory=memory
+        )
+        
+    def describe_image(self, img_path: str) -> str:
+        all_text = ""
+        try:
+            r = requests.get(img_path, timeout=10)
+            image_bytes = r.content
+            image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+            message = [
+                HumanMessage(
+                    content=[
+                        {
+                            "type": "text",
+                            "text": (
+                                "You're a chess assistant. Answer only with the best move in algebraic notation (e.g., Qd1#)."
+                            ),
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{image_base64}"
+                            },
+                        },
+                    ]
+                )
+            ]
+            response = self.model.invoke(message)
+            all_text += response.content + "\n\n"
+            return all_text.strip()
+        except Exception as e:
+            error_msg = f"Error extracting text: {str(e)}"
+            print(error_msg)
+            return ""
 
-    def __call__(self, question: str, task_id: str = None) -> str:
-        search_keywords = [
-            "who", "when", "what", "which", "how many", "number", "name", "albums", "surname", "at bats",
-            "nasa", "city", "winner", "code", "vegetable", "ingredient", "magda m.", "featured article"
-        ]
-        needs_search = any(kw in question.lower() for kw in search_keywords)
-        if needs_search:
-            web_result = duckduckgo_search(question)
-            llm_answer = self.llm.chat.completions.create(
-                model="gpt-4o",
-                messages=[
-                    {"role": "system", "content": "You are a research assistant. Based on the web search results and question, answer strictly and concisely for the GAIA benchmark. Only the answer, no explanations or apologies."},
-                    {"role": "user", "content": f"Web search results:\n{web_result}\n\nQuestion: {question}"}
-                ],
-                temperature=0.0,
-                max_tokens=256,
-            ).choices[0].message.content.strip()
-            formatted = format_gaia_answer(llm_answer, question)
-            # Retry if apology/empty/incorrect
-            if not formatted or "sorry" in formatted.lower() or "unable" in formatted.lower():
-                llm_answer2 = self.llm.chat.completions.create(
-                    model="gpt-4o",
-                    messages=[
-                        {"role": "system", "content": "Only answer with the value. No explanation. Do not apologize. Do not begin with 'I'm sorry', 'Unfortunately', or similar."},
-                        {"role": "user", "content": f"Web search results:\n{web_result}\n\nQuestion: {question}"}
-                    ],
-                    temperature=0.0,
-                    max_tokens=128,
-                ).choices[0].message.content.strip()
-                formatted = format_gaia_answer(llm_answer2, question)
-            return formatted
-        # For code/math output
-        if "output" in question.lower() and "python" in question.lower():
-            code_match = re.search(r'```python(.*?)```', question, re.DOTALL)
-            code = code_match.group(1) if code_match else ""
-            result = eval_python_code(code)
-            return format_gaia_answer(result, question)
-        # For lists/ingredients, always web search and format
-        if "list" in question.lower() or "ingredient" in question.lower() or "vegetable" in question.lower():
-            web_result = duckduckgo_search(question)
-            llm_answer = self.llm.chat.completions.create(
-                model="gpt-4o",
-                messages=[
-                    {"role": "system", "content": "You are a research assistant. Based on the web search results and question, answer strictly and concisely for the GAIA benchmark. Only the answer, no explanations or apologies."},
-                    {"role": "user", "content": f"Web search results:\n{web_result}\n\nQuestion: {question}"}
-                ],
-                temperature=0.0,
-                max_tokens=256,
-            ).choices[0].message.content.strip()
-            return format_gaia_answer(llm_answer, question)
-        # Fallback: strict LLM answer, formatted
-        llm_answer = self.llm.chat.completions.create(
-            model="gpt-4o",
-            messages=[
-                {"role": "system", "content": "You are a research assistant. Answer strictly and concisely for the GAIA benchmark. Only the answer, no explanations or apologies."},
-                {"role": "user", "content": question}
-            ],
-            temperature=0.0,
-            max_tokens=128,
-        ).choices[0].message.content.strip()
-        return format_gaia_answer(llm_answer, question)
+    def fetch_file(self, task_id):
+        try:
+            url = f"{DEFAULT_API_URL}/files/{task_id}"
+            r = requests.get(url, timeout=10)
+            r.raise_for_status()
+            return url, r.content, r.headers.get("Content-Type", "")
+        except:
+            return None, None, None
+        
+    def __call__(self, question: str, task_id: str) -> str:
+        print(f"Agent received question (first 50 chars) {task_id}: {question[:50]}...")
+        file_url, file_content, file_type = self.fetch_file(task_id)
+        print(f"Fetched file {file_type}")
+        if file_url is not None:
+            question = f"{question} This task has assigned file with URL: {file_url}"
+        fixed_answer = self.agent.run(question)
+        print(f"Agent returning fixed answer: {fixed_answer}")
+        return fixed_answer