Update agent.py

agent.py

@@ -1094,24 +1094,29 @@ def tools_condition(state: dict) -> str:
 from langgraph.graph import StateGraph
 from langchain_core.messages import SystemMessage
 from langchain_core.runnables import RunnableLambda
-
 def build_graph(vector_store, provider: str, model_config: dict) -> StateGraph:
+    # Get LLM
     llm = get_llm(provider, model_config)
 
-    tools = [wiki_search, calculator, web_search, arxiv_search,
-             get_youtube_transcript, extract_video_id, analyze_attachment, wikidata_query]
+    # Define available tools
+    tools = [
+        wiki_search, calculator, web_search, arxiv_search,
+        get_youtube_transcript, extract_video_id, analyze_attachment, wikidata_query
+    ]
 
+    # Tool mapping (global if needed elsewhere)
     global tool_map
     tool_map = {t.name: t for t in tools}
 
+    # Bind tools only if LLM supports it
     if hasattr(llm, "bind_tools"):
         llm_with_tools = llm.bind_tools(tools)
     else:
-        llm_with_tools = llm  # fallback: no tool binding
+        llm_with_tools = llm  # fallback for non-tool-aware models
 
-    
     sys_msg = SystemMessage(content="You are a helpful assistant.")
 
+    # Define nodes as runnables
     retriever = RunnableLambda(lambda state: {
         **state,
         "retrieved_docs": vector_store.similarity_search(state["input"])
@@ -1122,35 +1127,27 @@ def build_graph(vector_store, provider: str, model_config: dict) -> StateGraph:
         "messages": [sys_msg] + state["messages"]
     })
 
-    call_llm = llm_with_tools
+    call_llm = llm_with_tools  # already configured
 
+    # Start building the graph
     builder = StateGraph(AgentState)
     builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
     builder.add_node("call_llm", call_llm)
     builder.add_node("call_tool", tool_dispatcher)
+    builder.add_node("end", lambda state: state)  # Add explicit end node
 
+    # Define graph flow
     builder.set_entry_point("retriever")
     builder.add_edge("retriever", "assistant")
     builder.add_edge("assistant", "call_llm")
-    builder.add_node("end", lambda state: state)  # Add a terminal node
-    
+
     builder.add_conditional_edges("call_llm", should_call_tool, {
         "call_tool": "call_tool",
-        "end": "end"  
+        "end": "end"  # ✅ fixed: must point to actual "end" node
     })
 
-    builder.add_edge("call_tool", "call_llm")
+    builder.add_edge("call_tool", "call_llm")  # loop back after tool call
 
     return builder.compile()
-
-
-
-# call build_graph AFTER it’s defined
-graph = build_graph(vector_store, provider, model_config)
-agent = graph
-
-# Now you can use the agent like this:
-result = agent.invoke({"messages": [HumanMessage(content=question)]})
-
-
+