Update agent.py

agent.py

@@ -354,51 +354,44 @@ def wiki_search(query : str) -> str:
 tools = [weather_tool, wiki_search, web_search,
             add, subtract, multiply, divide, square, cube, power, factorial, mean, standard_deviation]
 
+# === LLM with Tools ===
+
 llm = ChatGroq(
     temperature=0,
-    model_name="qwen-qwq-32b",  # Updated to working model
+    model_name="qwen-qwq-32b",
     groq_api_key=os.getenv("GROQ_API_KEY")
 )
 
-llm_with_tools = llm.bind_tools(tools)
+tools = [weather_tool, wiki_search, web_search,
+         add, subtract, multiply, divide, square, cube,
+         power, factorial, mean, standard_deviation]
 
-# === NODES ===
+llm_with_tools = llm.bind_tools(tools)
 
-def retriever(state: MessagesState):
-    """Retrieve similar context and inject"""
-    query = state["messages"][0].content
-    similar_docs = vector_store.similarity_search(query)
+# === LangGraph State ===
 
-    if similar_docs:
-        ref_msg = HumanMessage(
-            content=f"Here is a similar question and answer for reference:\n\n{similar_docs[0].page_content}"
-        )
-        return {"messages": [sys_msg] + state["messages"] + [ref_msg]}
-    else:
-        return {"messages": [sys_msg] + state["messages"]}
+class ToolAgentState(TypedDict):
+    messages: Annotated[List[HumanMessage], "Messages in the conversation"]
 
-def assistant(state: MessagesState):
-    """Invoke LLM with tools"""
-    return {"messages": [llm_with_tools.invoke(state["messages"])]}
+def assistant(state: ToolAgentState):
+    return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
 
-# === GRAPH BUILD ===
+# === Build Graph ===
 
 def build_graph():
-    builder = StateGraph(MessagesState)
-    builder.add_node("retriever", retriever)
+    builder = StateGraph(ToolAgentState)
     builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode([retriever_tool] + tools))
+    builder.add_node("tools", ToolNode(tools))
 
-    builder.set_entry_point("retriever")
-    builder.add_edge("retriever", "assistant")
+    builder.set_entry_point("assistant")
     builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")
-
     return builder.compile()
 
-# === TEST ===
+# === Run ===
+
 if __name__ == "__main__":
-    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
+    question = "When did India won a world cup in cricket before 2000?"
     graph = build_graph()
     messages = [HumanMessage(content=question)]
     result = graph.invoke({"messages": messages})