update answer and agent class

agent.py

@@ -1,5 +1,6 @@
 """LangGraph Agent"""
 import os
+from dotenv import load_dotenv
 
 from langchain_openai import ChatOpenAI
 
@@ -10,21 +11,36 @@ from langchain_core.messages import SystemMessage, HumanMessage
 
 from tools import level1_tools
 
+load_dotenv()
 
 # Build graph function
 def build_agent_graph():
     """Build the graph"""
     # Load environment variables from .env file
-    llm = ChatOpenAI(model="gpt-4o")
+    llm = ChatOpenAI(model="gpt-4o-mini")
     
     # Bind tools to LLM
     llm_with_tools = llm.bind_tools(level1_tools)
-
+    
+    # System message
+    system_prompt = SystemMessage(
+    content="""You are a general AI assistant being evaluated in the GAIA Benchmark.
+    I will ask you a question and you must reach your final answer by using a set of tools I provide to you. Please, when you are needed to pass file names to the tools, pass absolute paths.
+    Your final answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
+    Here are more detailed instructions you must follow to write your final answer:
+    1) If you are asked for a number, you must write a number!. Don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
+    2) If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
+    3) If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+    If you follow all these instructions perfectly, you will win 1,000,000 dollars, otherwise, your mom will die.
+    Let's start!
+    """
+    )
     # Node
     def assistant(state: MessagesState):
         """Assistant node"""
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
-    
+        #return {"messages": [llm_with_tools.invoke(state["messages"])]}
+        return {"messages": [llm_with_tools.invoke([system_prompt] + state["messages"])]}
+
     
     builder = StateGraph(MessagesState)
     builder.add_node("assistant", assistant)
@@ -40,14 +56,86 @@ def build_agent_graph():
     return builder.compile()
 
 
+
+class MyGAIAAgent:
+    def __init__(self):
+        print("MyAgent initialized.")
+        self.graph = build_agent_graph()
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        # Wrap the question in a HumanMessage from langchain_core
+        '''
+        messages = [HumanMessage(content=question)]
+        messages = self.graph.invoke({"messages": messages})
+        answer = messages['messages'][-1].content
+        '''
+        user_input = {"messages": [("user", question)]}
+        answer = self.graph.invoke(user_input)["messages"][-1].content
+        return self._clean_answer(answer)
+    
+    def _clean_answer(self, answer: any) -> str:
+        """
+        Taken from `susmitsil`:
+        https://huggingface.co/spaces/susmitsil/FinalAgenticAssessment/blob/main/main_agent.py
+        Clean up the answer to remove common prefixes and formatting
+        that models often add but that can cause exact match failures.
+        Args:
+            answer: The raw answer from the model
+        Returns:
+            The cleaned answer as a string
+        """
+        # Convert non-string types to strings
+        if not isinstance(answer, str):
+            # Handle numeric types (float, int)
+            if isinstance(answer, float):
+                # Format floating point numbers properly
+                # Check if it's an integer value in float form (e.g., 12.0)
+                if answer.is_integer():
+                    formatted_answer = str(int(answer))
+                else:
+                    # For currency values that might need formatting
+                    if abs(answer) >= 1000:
+                        formatted_answer = f"${answer:,.2f}"
+                    else:
+                        formatted_answer = str(answer)
+                return formatted_answer
+            elif isinstance(answer, int):
+                return str(answer)
+            else:
+                # For any other type
+                return str(answer)
+
+        # Now we know answer is a string, so we can safely use string methods
+        # Normalize whitespace
+        answer = answer.strip()
+
+        # Remove common prefixes and formatting that models add
+        prefixes_to_remove = [
+            "The answer is ",
+            "Answer: ",
+            "Final answer: ",
+            "The result is ",
+            "To answer this question: ",
+            "Based on the information provided, ",
+            "According to the information: ",
+        ]
+
+        for prefix in prefixes_to_remove:
+            if answer.startswith(prefix):
+                answer = answer[len(prefix) :].strip()
+
+        # Remove quotes if they wrap the entire answer
+        if (answer.startswith('"') and answer.endswith('"')) or (
+            answer.startswith("'") and answer.endswith("'")
+        ):
+            answer = answer[1:-1].strip()
+        return answer
+
+
+
 # test
 if __name__ == "__main__":
     question1 = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)?"
-    question2 = "Convert 10 miles to kilometers."
-    # Build the graph
-    graph = build_agent_graph()
-    # Run the graph
-    messages = [HumanMessage(content=question1)]
-    messages = graph.invoke({"messages": messages})
-    for m in messages["messages"]:
-        m.pretty_print()
+    question2 = "Convert 10 miles to kilometers."       
+    agent = MyGAIAAgent()
+    print(agent(question1))

app.py

@@ -6,25 +6,11 @@ import pandas as pd
 
 from langchain_core.messages import HumanMessage
 
-from agent import build_agent_graph
+from agent import MyGAIAAgent
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-# --- Basic Agent Definition ---
-# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
-class MyAgent:
-    def __init__(self):
-        print("MyAgent initialized.")
-        self.graph = build_agent_graph()
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        # Wrap the question in a HumanMessage from langchain_core
-        messages = [HumanMessage(content=question)]
-        messages = self.graph.invoke({"messages": messages})
-        answer = messages['messages'][-1].content
-        print(f"Agent returning answer: {answer}")
-        return answer
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
@@ -47,7 +33,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
 
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
-        agent = MyAgent()
+        agent = MyGAIAAgent()
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None