Create agent.py

agent.py

@@ -0,0 +1,69 @@
+import os
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import tool
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+
+# Load system prompt from system_prompt.txt
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    SYSTEM_PROMPT = f.read()
+
+# Define tools
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return 2 results."""
+    docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    return "\n\n---\n\n".join([doc.page_content for doc in docs])
+
+@tool
+def web_search(query: str) -> str:
+    """Search the web using Tavily API and return 3 results."""
+    docs = TavilySearchResults(max_results=3).invoke(query)
+    return "\n\n---\n\n".join([doc.page_content for doc in docs])
+
+tools = [wiki_search, web_search]
+
+# Build LangGraph agent
+def build_graph():
+    llm = ChatOpenAI(
+        model="gpt-4o",
+        temperature=0,
+        api_key=os.getenv("OPENAI_API_KEY")
+    )
+    llm_with_tools = llm.bind_tools(tools)
+
+    def system_node(state: MessagesState):
+        return {"messages": [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]}
+
+    def assistant_node(state: MessagesState):
+        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+
+    graph = StateGraph(MessagesState)
+    graph.add_node("system", system_node)
+    graph.add_node("assistant", assistant_node)
+    graph.add_node("tools", ToolNode(tools))
+
+    graph.add_edge(START, "system")
+    graph.add_edge("system", "assistant")
+    graph.add_conditional_edges("assistant", tools_condition)
+    graph.add_edge("tools", "assistant")
+
+    return graph.compile()
+
+# Final class for submission agent
+class BasicAgent:
+    def __init__(self):
+        print("Initializing LangGraph GAIA Agent...")
+        self.graph = build_graph()
+
+    def __call__(self, question: str) -> str:
+        try:
+            messages = [HumanMessage(content=question)]
+            result = self.graph.invoke({"messages": messages})
+            final_msg = result["messages"][-1]
+            return final_msg.content.strip()
+        except Exception as e:
+            return f"Agent error: {str(e)}"