Update app.py

app.py

@@ -4,7 +4,17 @@ import requests
 import inspect
 import pandas as pd
 from langchain_core.messages import HumanMessage
-from agent import BasicAgent
+from typing import TypedDict, Annotated, Sequence, Dict, Any, List
+from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_core.tools import tool
+from langchain_openai import ChatOpenAI
+from langgraph.graph import END, StateGraph
+from langgraph.prebuilt import ToolNode
+from langchain_community.tools import DuckDuckGoSearchResults
+from langchain_community.utilities import WikipediaAPIWrapper
+from langchain.agents import create_tool_calling_agent, AgentExecutor
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+import operator
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -12,6 +22,29 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
+
+
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+    sender: str
+
+@tool
+def wikipedia_search(query: str) -> str:
+    """Search Wikipedia for information."""
+    return WikipediaAPIWrapper().run(query)
+
+@tool
+def web_search(query: str, num_results: int = 3) -> list:
+    """Search the web for current information."""
+    return DuckDuckGoSearchResults(num_results=num_results).run(query)
+
+@tool
+def calculate(expression: str) -> str:
+    """Evaluate mathematical expressions."""
+    from langchain_experimental.utilities import PythonREPL
+    python_repl = PythonREPL()
+    return python_repl.run(expression)
+
 class BasicAgent:
     """A langgraph agent."""
     def __init__(self):
@@ -26,6 +59,67 @@ class BasicAgent:
         answer = messages['messages'][-1].content
         return answer[14:]
 
+    
+    def _build_workflow(self) -> StateGraph:
+        """Build and return the compiled workflow"""
+        workflow = StateGraph(AgentState)
+        
+        workflow.add_node("agent", self._run_agent)
+        workflow.add_node("tools", ToolNode(self.tools))
+        
+        workflow.set_entry_point("agent")
+        workflow.add_conditional_edges(
+            "agent",
+            self._should_continue,
+            {"continue": "tools", "end": END}
+        )
+        workflow.add_edge("tools", "agent")
+        
+        return workflow.compile()
+    
+    def _run_agent(self, state: AgentState) -> Dict[str, Any]:
+        """Execute the agent"""
+        response = self.agent_executor.invoke({"messages": state["messages"]})
+        return {"messages": [response["output"]]}
+    
+    def _should_continue(self, state: AgentState) -> str:
+        """Determine if the workflow should continue"""
+        last_message = state["messages"][-1]
+        return "continue" if last_message.additional_kwargs.get("tool_calls") else "end"
+    
+    def __call__(self, query: str) -> Dict[str, Any]:
+        """Process a user query"""
+        state = AgentState(messages=[HumanMessage(content=query)], sender="user")
+        
+        for output in self.workflow.stream(state):
+            for key, value in output.items():
+                if key == "messages":
+                    for message in value:
+                        if isinstance(message, BaseMessage):
+                            return {
+                                "response": message.content,
+                                "sources": self._extract_sources(state["messages"]),
+                                "steps": self._extract_steps(state["messages"])
+                            }
+        return {"response": "No response generated", "sources": [], "steps": []}
+    
+    def _extract_sources(self, messages: Sequence[BaseMessage]) -> List[str]:
+        """Extract sources from tool messages"""
+        return [
+            f"{msg.additional_kwargs.get('name', 'unknown')}: {msg.content}"
+            for msg in messages
+            if hasattr(msg, 'additional_kwargs') and 'name' in msg.additional_kwargs
+        ]
+    
+    def _extract_steps(self, messages: Sequence[BaseMessage]) -> List[str]:
+        """Extract reasoning steps"""
+        steps = []
+        for msg in messages:
+            if hasattr(msg, 'additional_kwargs') and 'tool_calls' in msg.additional_kwargs:
+                for call in msg.additional_kwargs['tool_calls']:
+                    steps.append(f"Used {call['function']['name']}: {call['function']['arguments']}")
+        return steps
+        
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,