Create agent.py

agent.py

@@ -0,0 +1,204 @@
+# GAIA Agent Solution with LangGraph and OpenAI
+import os
+from typing import TypedDict, Annotated, Sequence, Union
+from langgraph.graph import StateGraph, END
+from langchain_community.tools import DuckDuckGoSearchResults
+from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_openai import ChatOpenAI
+from langchain_core.tools import tool
+from langchain_core.utils.function_calling import convert_to_openai_tool
+import json
+from openai import OpenAI  # For vision capabilities
+
+# Set your OpenAI API key
+openai_api_key = os.getenv("OPENAI_API_KEY")  # Replace with your actual key
+
+# ---------------------
+# Tool Definitions
+# ---------------------
+
+# Web Search Tool
+search_tool = DuckDuckGoSearchResults(max_results=3)
+
+# Image Description Tool (using GPT-4 Vision)
+@tool
+def describe_image(image_url: str) -> str:
+    """Generate detailed description of an image from its URL"""
+    vision_client = OpenAI()
+    response = vision_client.chat.completions.create(
+        model="gpt-4-vision-preview",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Describe this image in detail. Include text, objects, colors, and context."},
+                    {"type": "image_url", "image_url": {"url": image_url}}
+                ]
+            }
+        ],
+        max_tokens=500
+    )
+    return response.choices[0].message.content
+
+# Math Tool (example - extend with more capabilities)
+@tool
+def calculate(expression: str) -> Union[float, str]:
+    """Evaluate mathematical expressions. Input must be a valid math expression."""
+    try:
+        return eval(expression)  # For real usage, use a safe evaluator like numexpr
+    except:
+        return "Error: Invalid expression"
+
+# ---------------------
+# Agent Setup
+# ---------------------
+
+# Available tools
+tools = [search_tool, describe_image, calculate]
+tools_as_openai = [convert_to_openai_tool(t) for t in tools]
+
+# Agent State Definition
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+
+# Initialize LLM (GPT-4 Turbo for best results)
+model = ChatOpenAI(model="gpt-4-turbo", temperature=0)
+
+# ---------------------
+# Graph Nodes
+# ---------------------
+
+def run_agent(state: AgentState):
+    """Node: Run the agent's reasoning"""
+    messages = state["messages"]
+    response = model.invoke(messages, tools=tools_as_openai)
+    return {"messages": [response]}
+
+def run_tools(state: AgentState):
+    """Node: Execute tools based on agent's request"""
+    messages = state["messages"]
+    last_message = messages[-1]
+    
+    tool_messages = []
+    for tool_call in last_message.additional_kwargs.get("tool_calls", []):
+        function_name = tool_call["function"]["name"]
+        function_args = json.loads(tool_call["function"]["arguments"])
+        
+        # Find matching tool
+        tool = next((t for t in tools if t.name == function_name), None)
+        
+        if tool:
+            try:
+                # Special handling for image URLs in questions
+                if function_name == "describe_image" and "http" not in function_args["image_url"]:
+                    function_args["image_url"] = find_image_url(messages, function_args["image_url"])
+                
+                # Execute tool
+                output = tool.invoke(function_args)
+                content = f"Tool Result: {str(output)}"
+            except Exception as e:
+                content = f"Error: {str(e)}"
+        else:
+            content = f"Tool {function_name} not available"
+        
+        tool_messages.append(
+            ToolMessage(
+                content=content,
+                tool_call_id=tool_call["id"]
+            )
+        )
+    
+    return {"messages": tool_messages}
+
+# ---------------------
+# Helper Functions
+# ---------------------
+
+def find_image_url(messages: Sequence[BaseMessage], reference: str) -> str:
+    """Extract actual image URL from message context"""
+    for msg in messages:
+        if reference in msg.content:
+            # Simple extraction - improve with regex for production
+            if "http" in msg.content and ("jpg" in msg.content or "png" in msg.content):
+                start = msg.content.find("http")
+                return msg.content[start:].split()[0]
+    return reference  # Fallback to original reference
+
+# ---------------------
+# Graph Construction
+# ---------------------
+
+# Decision logic for graph flow
+def should_continue(state: AgentState):
+    last_message = state["messages"][-1]
+    if last_message.tool_calls:
+        return "run_tools"
+    return "end"
+
+# Build the graph
+graph = StateGraph(AgentState)
+graph.add_node("run_agent", run_agent)
+graph.add_node("run_tools", run_tools)
+graph.set_entry_point("run_agent")
+
+graph.add_conditional_edges(
+    "run_agent",
+    should_continue,
+    {
+        "run_tools": "run_tools",
+        "end": END
+    }
+)
+
+graph.add_edge("run_tools", "run_agent")
+agent = graph.compile()
+
+# ---------------------
+# Execution Function
+# ---------------------
+
+def solve_gaia_task(question: str) -> str:
+    """Solve GAIA tasks with our agent"""
+    system_prompt = (
+        "You are a GAIA problem-solving expert. Follow these rules:\n"
+        "1. Use tools for current information\n"
+        "2. Break complex problems into steps\n"
+        "3. Verify answers before finalizing\n"
+        "4. Format final answers EXACTLY as requested:\n"
+        "   - Lists: comma-separated values\n"
+        "   - Numbers: digits only\n"
+        "   - Dates: YYYY-MM-DD format\n"
+        "5. Never include reasoning in final answers"
+    )
+    
+    # Initialize agent state
+    state = {
+        "messages": [
+            SystemMessage(content=system_prompt),
+            HumanMessage(content=question)
+        ]
+    }
+    
+    # Run the agent
+    final_state = agent.invoke(state)
+    
+    # Extract and return final answer
+    for msg in reversed(final_state["messages"]):
+        if msg.type == "ai" and not msg.tool_calls:
+            return msg.content
+    return "No final answer found"
+
+# ---------------------
+# Example Execution
+# ---------------------
+if __name__ == "__main__":
+    # Example GAIA task
+    task = (
+        "What is the current population of the country where the 2023 "
+        "World Artificial Intelligence Conference was held? "
+        "Include only the numeric value in your answer."
+    )
+    
+    result = solve_gaia_task(task)
+    print("\n--- FINAL ANSWER ---")
+    print(result)