Update app.py

updated with langgraph

app.py

@@ -3,6 +3,11 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
+import sympy
+import re
+from duckduckgo_search import ddg
+from langgraph.graph import StateGraph, END
+from typing import TypedDict
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -14,74 +19,92 @@ from smolagent import Agent, Tool
 from duckduckgo_search import ddg
 import sympy
 import re
+# --- Tools ---
 
-class WikipediaSearchTool(Tool):
-    name = "wikipedia_search"
-    description = "Searches for answers using Wikipedia"
+def wikipedia_search_tool(input: str) -> str:
+    results = ddg(input, max_results=3)
+    if results:
+        return results[0]["body"]
+    return "No information found."
 
-    def use(self, input: str) -> str:
-        results = ddg(input, max_results=3)
-        if results:
-            return results[0]["body"]
-        return "No information found."
+def math_solver_tool(input: str) -> str:
+    try:
+        expr = sympy.sympify(input)
+        return str(expr.evalf())
+    except Exception as e:
+        return f"Math Error: {e}"
 
-class MathSolverTool(Tool):
-    name = "math_solver"
-    description = "Solves math expressions and equations"
+def code_execution_tool(input: str) -> str:
+    try:
+        local_vars = {}
+        exec(f"result = {input}", {}, local_vars)
+        return str(local_vars["result"])
+    except Exception as e:
+        return f"Code Error: {e}"
 
-    def use(self, input: str) -> str:
-        try:
-            expression = sympy.sympify(input)
-            result = expression.evalf()
-            return str(result)
-        except Exception as e:
-            return f"Math Error: {e}"
+# --- State definition ---
 
-class CodeExecutionTool(Tool):
-    name = "python_executor"
-    description = "Executes Python code snippets and returns output"
+class AgentState(TypedDict):
+    question: str
+    response: str
 
-    def use(self, input: str) -> str:
-        try:
-            local_vars = {}
-            exec(f"result = {input}", {}, local_vars)
-            return str(local_vars["result"])
-        except Exception as e:
-            return f"Code Error: {e}"
+# --- Routing logic ---
 
-def classify_question(q: str) -> str:
-    q_lower = q.lower()
-    if any(keyword in q_lower for keyword in ["solve", "evaluate", "calculate", "math"]):
+def route_question(state: AgentState) -> str:
+    q = state["question"].lower()
+    if any(k in q for k in ["solve", "calculate", "evaluate", "math", "=", "^"]):
         return "math"
-    elif any(keyword in q_lower for keyword in ["code", "python", "return", "function"]):
+    elif any(k in q for k in ["code", "python", "return", "function"]):
         return "code"
-    elif any(keyword in q_lower for keyword in ["who", "what", "where", "when", "why", "how", "capital", "invented"]):
+    elif any(k in q for k in ["what", "who", "when", "where", "why", "capital", "invented"]):
         return "search"
     else:
-        return "default"
+        return "search"
+
+# --- LangGraph graph setup ---
+
+builder = StateGraph(AgentState)
+
+@builder.node()
+def math_node(state: AgentState) -> AgentState:
+    return {"question": state["question"], "response": math_solver_tool(state["question"])}
+
+@builder.node()
+def code_node(state: AgentState) -> AgentState:
+    return {"question": state["question"], "response": code_execution_tool(state["question"])}
+
+@builder.node()
+def search_node(state: AgentState) -> AgentState:
+    return {"question": state["question"], "response": wikipedia_search_tool(state["question"])}
+
+@builder.multi_choice()
+def router(state: AgentState) -> str:
+    return route_question(state)
+
+builder.set_entry_point("router", router)
+builder.add_node("math", math_node)
+builder.add_node("code", code_node)
+builder.add_node("search", search_node)
+builder.add_conditional_edges("router", {
+    "math": "math",
+    "code": "code",
+    "search": "search"
+})
+builder.add_edge("math", END)
+builder.add_edge("code", END)
+builder.add_edge("search", END)
+
+app_graph = builder.compile()
 
 class BasicAgent:
     def __init__(self):
-        self.agent = Agent(
-            tools=[
-                WikipediaSearchTool(),
-                MathSolverTool(),
-                CodeExecutionTool(),
-            ]
-        )
+        print("LangGraph Agent initialized.")
 
     def __call__(self, question: str) -> str:
-        category = classify_question(question)
-        print(f"Routing question: {question[:50]}... -> [{category}]")
-
-        if category == "math":
-            return MathSolverTool().use(question)
-        elif category == "code":
-            return CodeExecutionTool().use(question)
-        elif category == "search":
-            return WikipediaSearchTool().use(question)
-        else:
-            return self.agent.run(question)
+        state = {"question": question, "response": ""}
+        result = app_graph.invoke(state)
+        return result["response"]
+
 
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):