Update agent.py

agent.py

@@ -1,129 +1,137 @@
-import os
-from dotenv import load_dotenv
-from langgraph.graph import START, END, StateGraph, MessagesState
-from langgraph.prebuilt import ToolNode
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
-from langchain_core.tools import tool
-from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
-from sympy import sympify
-
-load_dotenv()
-
-#############################
-# TOOLS
-#############################
-
-@tool
-def evaluate_math_expression(expr: str) -> str:
-    """Evaluate a math expression using symbolic math."""
-    try:
-        result = sympify(expr)
-        return str(result.evalf())
-    except Exception as e:
-        return f"Math evaluation error: {e}"
-
-@tool
-def search_arxiv(query: str) -> str:
-    """Search Arxiv for scientific papers related to a query."""
-    docs = ArxivLoader(query=query, load_max_docs=2).load()
-    return "\n---\n".join(doc.page_content[:1000] for doc in docs)
-
-@tool
-def search_wikipedia(query: str) -> str:
-    """Search Wikipedia and return summarized results."""
-    docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    return "\n---\n".join(doc.page_content[:500] for doc in docs)
-
-@tool
-def web_search(query: str) -> str:
-    """Web search using Tavily (3 results)."""
-    return TavilySearchResults(max_results=3).invoke(query=query)
-
-#############################
-# INITIALIZE TOOLS
-#############################
-
-tools = [
-    evaluate_math_expression,
-    search_arxiv,
-    search_wikipedia,
-    web_search
-]
-
-#############################
-# SYSTEM MESSAGE
-#############################
-
-sys_msg = SystemMessage(content="""
-You are a research assistant agent. Given a question, attempt to:
-1. Identify whether it requires mathematical computation or factual research.
-2. Use the appropriate tool or set of tools.
-3. Answer concisely and correctly.
-""")
-
-#############################
-# NODES
-#############################
-
-def classifier_node(state: MessagesState):
-    question = state['messages'][-1].content.lower()
-    if any(k in question for k in ["calculate", "what is", "solve", "+", "-", "*", "/"]):
-        return "math"
-    return "search"
-
-def math_node(state: MessagesState):
-    expr = state['messages'][-1].content
-    response = evaluate_math_expression.invoke(expr)
-    return {"messages": state['messages'] + [AIMessage(content=response)]}
-
-def search_node(state: MessagesState):
-    query = state['messages'][-1].content
-    wiki = search_wikipedia.invoke(query)
-    return {"messages": state['messages'] + [HumanMessage(content=f"Here is some info from Wikipedia:\n{wiki}")]}
-
-def assistant_node(state: MessagesState):
-    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
-    return {"messages": state['messages'] + [llm.invoke(state['messages'])]}
-
-def answer_validator_node(state: MessagesState):
-    last_response = state['messages'][-1].content
-    if "error" in last_response.lower():
-        return {"messages": state['messages'] + [AIMessage(content="I'm sorry, something went wrong.")]}
-    return {"messages": state['messages']}
-
-#############################
-# BUILD GRAPH
-#############################
-
-def build_graph():
-    builder = StateGraph(MessagesState)
-
-    builder.add_node("classifier", classifier_node)
-    builder.add_node("math", math_node)
-    builder.add_node("search", search_node)
-    builder.add_node("assistant", assistant_node)
-    builder.add_node("validate", answer_validator_node)
-
-    builder.set_entry_point("classifier")
-    builder.add_edge("math", "validate")
-    builder.add_edge("search", "assistant")
-    builder.add_edge("assistant", "validate")
-    builder.add_edge("validate", END)
-
-    builder.add_conditional_edges("classifier", lambda state: classifier_node(state))
-
-    return builder.compile()
-
-#############################
-# TEST MAIN
-#############################
-
-if __name__ == "__main__":
-    question = "What is the result of (3^2 + 4^2)?"
-    graph = build_graph()
-    messages = [HumanMessage(content=question)]
-    output = graph.invoke({"messages": messages})
-    for m in output["messages"]:
-        m.pretty_print()
+import os
+from dotenv import load_dotenv
+from langgraph.graph import START, END, StateGraph, MessagesState
+from langgraph.prebuilt import ToolNode
+from langchain_google_genai import ChatGoogleGenerativeAI
+llm = ChatGoogleGenerativeAI(
+    model="gemini-2.0-flash",  # Specify your model
+    google_api_key=
+)
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+from langchain_core.tools import tool
+from langchain_community.tools.tavily_search import TavilySearchResults
+search = TavilySearchResults(
+    max_results=5,  # Adjust as needed
+    api_key=
+)
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from sympy import sympify
+
+load_dotenv()
+
+#############################
+# TOOLS
+#############################
+
+@tool
+def evaluate_math_expression(expr: str) -> str:
+    """Evaluate a math expression using symbolic math."""
+    try:
+        result = sympify(expr)
+        return str(result.evalf())
+    except Exception as e:
+        return f"Math evaluation error: {e}"
+
+@tool
+def search_arxiv(query: str) -> str:
+    """Search Arxiv for scientific papers related to a query."""
+    docs = ArxivLoader(query=query, load_max_docs=2).load()
+    return "\n---\n".join(doc.page_content[:1000] for doc in docs)
+
+@tool
+def search_wikipedia(query: str) -> str:
+    """Search Wikipedia and return summarized results."""
+    docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    return "\n---\n".join(doc.page_content[:500] for doc in docs)
+
+@tool
+def web_search(query: str) -> str:
+    """Web search using Tavily (3 results)."""
+    return TavilySearchResults(max_results=3).invoke(query=query)
+
+#############################
+# INITIALIZE TOOLS
+#############################
+
+tools = [
+    evaluate_math_expression,
+    search_arxiv,
+    search_wikipedia,
+    web_search
+]
+
+#############################
+# SYSTEM MESSAGE
+#############################
+
+sys_msg = SystemMessage(content="""
+You are a research assistant agent. Given a question, attempt to:
+1. Identify whether it requires mathematical computation or factual research.
+2. Use the appropriate tool or set of tools.
+3. Answer concisely and correctly.
+""")
+
+#############################
+# NODES
+#############################
+
+def classifier_node(state: MessagesState):
+    question = state['messages'][-1].content.lower()
+    if any(k in question for k in ["calculate", "what is", "solve", "+", "-", "*", "/"]):
+        return "math"
+    return "search"
+
+def math_node(state: MessagesState):
+    expr = state['messages'][-1].content
+    response = evaluate_math_expression.invoke(expr)
+    return {"messages": state['messages'] + [AIMessage(content=response)]}
+
+def search_node(state: MessagesState):
+    query = state['messages'][-1].content
+    wiki = search_wikipedia.invoke(query)
+    return {"messages": state['messages'] + [HumanMessage(content=f"Here is some info from Wikipedia:\n{wiki}")]}
+
+def assistant_node(state: MessagesState):
+    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
+    return {"messages": state['messages'] + [llm.invoke(state['messages'])]}
+
+def answer_validator_node(state: MessagesState):
+    last_response = state['messages'][-1].content
+    if "error" in last_response.lower():
+        return {"messages": state['messages'] + [AIMessage(content="I'm sorry, something went wrong.")]}
+    return {"messages": state['messages']}
+
+#############################
+# BUILD GRAPH
+#############################
+
+def build_graph():
+    builder = StateGraph(MessagesState)
+
+    builder.add_node("classifier", classifier_node)
+    builder.add_node("math", math_node)
+    builder.add_node("search", search_node)
+    builder.add_node("assistant", assistant_node)
+    builder.add_node("validate", answer_validator_node)
+
+    builder.set_entry_point("classifier")
+    builder.add_edge("math", "validate")
+    builder.add_edge("search", "assistant")
+    builder.add_edge("assistant", "validate")
+    builder.add_edge("validate", END)
+
+    builder.add_conditional_edges("classifier", lambda state: classifier_node(state))
+
+    return builder.compile()
+
+#############################
+# TEST MAIN
+#############################
+
+if __name__ == "__main__":
+    question = "What is the result of (3^2 + 4^2)?"
+    graph = build_graph()
+    messages = [HumanMessage(content=question)]
+    output = graph.invoke({"messages": messages})
+    for m in output["messages"]:
+        m.pretty_print()