Update LG_agent.py

LG_agent.py

@@ -7,6 +7,7 @@ from langchain_openai import ChatOpenAI
 from tools import all_tools
 import inspect
 import os
+import re
 
 # 1. Setup once
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -25,6 +26,35 @@ chat_with_tools = chat.bind_tools(all_tools)
 class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
 
+def extract_gaia_answer(text: str) -> str:
+    """
+    Extract only the final answer from the LLM output using common patterns like:
+    - 'The answer is: ...'
+    - 'Answer: ...'
+    - bullet points
+    - Or just return raw answer if short and plain
+    
+    Strips everything else to conform with GAIA expectations.
+    """
+    # Common patterns GAIA outputs follow
+    patterns = [
+        r"The answer is: (.+)",
+        r"Answer: (.+)",
+        r"^([a-z0-9,\s]+)$",  # raw answer line
+    ]
+    for pattern in patterns:
+        match = re.search(pattern, text.strip(), re.IGNORECASE | re.MULTILINE)
+        if match:
+            return match.group(1).strip()
+    
+    # Fallback: if it's a single short line, return it
+    lines = text.strip().splitlines()
+    if len(lines) == 1 and len(lines[0]) < 100:
+        return lines[0].strip()
+    
+    # Fallback: return full (stripped) content
+    return text.strip()
+
 # 3. Assistant node
 
 def assistant(state: AgentState):
@@ -114,17 +144,30 @@ def build_graph(max_steps: int = 5):
 
 # 5. BasicAgent class
 
-class BasicAgent:
+# class BasicAgent:
+#     def __init__(self, max_steps: int = 5):
+#         self.graph = build_graph(max_steps)
+
+#     def __call__(self, question: str) -> str:
+#         response = self.graph({"messages": [HumanMessage(content=question)]})
+#         if response.get("messages"):
+#             final_message = response["messages"][-1]
+#             return final_message.content if hasattr(final_message, "content") else "No final message."
+#         else:
+#             return "No response."
+
+def __call__(self, question: str) -> str:
     def __init__(self, max_steps: int = 5):
         self.graph = build_graph(max_steps)
+    
+    response = self.graph({"messages": [HumanMessage(content=question)]})
+    if response.get("messages"):
+        final_message = response["messages"][-1]
+        raw_content = final_message.content if hasattr(final_message, "content") else "No final message."
+        return extract_gaia_answer(raw_content)
+    else:
+        return "No response."
 
-    def __call__(self, question: str) -> str:
-        response = self.graph({"messages": [HumanMessage(content=question)]})
-        if response.get("messages"):
-            final_message = response["messages"][-1]
-            return final_message.content if hasattr(final_message, "content") else "No final message."
-        else:
-            return "No response."
 
 if __name__ == "__main__":
     agent = BasicAgent()