Update agent.py

agent.py

@@ -4,7 +4,7 @@ import os
 from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import tools_condition
-from langgraph.prebuilt import ToolNode
+#from langgraph.prebuilt import ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
@@ -839,44 +839,74 @@ def process_question(question: str):
 
 
 
+
 def call_llm(state):
     messages = state["messages"]
     response = llm.invoke(messages)
     return {"messages": messages + [response]}
 
-tool_node = ToolNode(tools)
+builder.set_entry_point("call_llm")
+builder.add_conditional_edges("call_llm", should_call_tool, {
+    "call_tool": "call_tool",
+    "end": None
+})
 
+# After tool runs, go back to the LLM
+builder.add_edge("call_tool", "call_llm")
 
-from langgraph.graph import StateGraph
-from typing import TypedDict
 
-# Define the state schema
-class AgentState(TypedDict):
-    input: str
-    result: str
 
-# Create the state graph builder with state_schema
-builder = StateGraph(AgentState)
+from langchain.schema import AIMessage
 
+def tool_dispatcher(state: AgentState) -> AgentState:
+    last_msg = state["messages"][-1]
+
+    # Make sure it's an AI message with tool_calls
+    if isinstance(last_msg, AIMessage) and last_msg.tool_calls:
+        tool_call = last_msg.tool_calls[0]
+        tool_name = tool_call["name"]
+        tool_input = tool_call["args"]  # Adjust based on your actual schema
+
+        tool_func = tool_map.get(tool_name, default_tool)
+
+        # If args is a dict and your tool expects unpacked values:
+        if isinstance(tool_input, dict):
+            result = tool_func.invoke(tool_input) if hasattr(tool_func, "invoke") else tool_func(**tool_input)
+        else:
+            result = tool_func.invoke(tool_input) if hasattr(tool_func, "invoke") else tool_func(tool_input)
+
+        # You can choose to append this to messages, or just save result
+        return {
+            **state,
+            "result": result,
+            # Optionally add: "messages": state["messages"] + [ToolMessage(...)]
+        }
+
+    # No tool call detected, return state unchanged
+    return state
+
+    
 
-builder.add_node("call_llm", call_llm)
-builder.add_node("call_tool", tool_node)
 
 # Decide what to do next: if tool call → call_tool, else → end
-def should_call_tool(state):
+def call_tool(state):
     last_msg = state["messages"][-1]
     if isinstance(last_msg, AIMessage) and last_msg.tool_calls:
         return "call_tool"
     return "end"
 
-builder.set_entry_point("call_llm")
-builder.add_conditional_edges("call_llm", should_call_tool, {
-    "call_tool": "call_tool",
-    "end": None
-})
+from typing import TypedDict, List, Optional, Union
+from langchain.schema import BaseMessage
+
+class AgentState(TypedDict):
+    messages: List[BaseMessage]   # chat history
+    input: str                    # original input
+    intent: str                   # derived or predicted intent
+    result: Optional[str]         # tool output, if any
+
+builder.add_node("call_tool", tool_dispatcher)
+
 
-# After tool runs, go back to the LLM
-builder.add_edge("call_tool", "call_llm")
 
 
 
@@ -1009,7 +1039,7 @@ model_config = {
 
 def build_graph(provider, model_config):
     from langchain_core.messages import SystemMessage, HumanMessage
-    from langgraph.graph import StateGraph, ToolNode
+    #from langgraph.graph import StateGraph, ToolNode
     from langchain_core.runnables import RunnableLambda
     from some_module import vector_store  # Make sure this is defined/imported
 
@@ -1078,20 +1108,39 @@ def build_graph(provider, model_config):
         else:
             return "END"
 
-    # Step 5: Define LangGraph StateGraph
-    builder = StateGraph(dict)  # Using dict as state type here
 
-    builder.add_node("retriever", retriever)
-    builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
 
-    builder.set_entry_point("retriever")
-    builder.add_edge("retriever", "assistant")
-    builder.add_conditional_edges("assistant", tools_condition)
-    builder.add_edge("tools", "assistant")
+from langgraph.graph import StateGraph
+
+# Build graph using AgentState as the shared schema
+builder = StateGraph(AgentState)
+
+# Add nodes
+builder.add_node("retriever", retriever)
+builder.add_node("assistant", assistant)
+builder.add_node("call_llm", call_llm)
+builder.add_node("call_tool", tool_dispatcher)  # one name is enough
+
+# Entry point
+builder.set_entry_point("retriever")
+
+# Define the flow
+builder.add_edge("retriever", "assistant")
+builder.add_edge("assistant", "call_llm")
+
+# Conditional edge from LLM to tool or end
+builder.add_conditional_edges("call_llm", should_call_tool, {
+    "call_tool": "call_tool",
+    "end": None
+})
+
+# Loop back after tool execution
+builder.add_edge("call_tool", "call_llm")
+
+# Compile
+graph = builder.compile()
+return graph
 
-    graph = builder.compile()
-    return graph
 
 
 # call build_graph AFTER it’s defined