Update agent.py

agent.py

@@ -1087,41 +1087,21 @@ def tools_condition(state: dict) -> str:
             
 
 from langgraph.graph import StateGraph
+from langchain_core.messages import SystemMessage
+from langchain_core.runnables import RunnableLambda
 
-# Build graph using AgentState as the shared schema
-def build_graph() -> StateGraph:
-    builder = StateGraph(AgentState)
-
-    from langchain_core.messages import SystemMessage, HumanMessage
-    #from langgraph.graph import StateGraph, ToolNode
-    from langchain_core.runnables import RunnableLambda
-    from some_module import vector_store  # Make sure this is defined/imported
-
-
+def build_graph(vector_store, provider: str, model_config: dict) -> StateGraph:
     llm = get_llm(provider, model_config)
 
-    # Step 2: Define tools
-    tools = [
-        wiki_search, 
-        calculator, 
-        web_search, 
-        arxiv_search,
-        get_youtube_transcript, 
-        extract_video_id, 
-        analyze_attachment, 
-        wikidata_query
-    ]
+    tools = [wiki_search, calculator, web_search, arxiv_search,
+             get_youtube_transcript, extract_video_id, analyze_attachment, wikidata_query]
 
     global tool_map
     tool_map = {t.name: t for t in tools}
 
-    # Step 3: Bind tools to LLM
     llm_with_tools = llm.bind_tools(tools)
-
-    # Step 4: Build stateful graph logic
     sys_msg = SystemMessage(content="You are a helpful assistant.")
 
-    # --- Define nodes ---
     retriever = RunnableLambda(lambda state: {
         **state,
         "retrieved_docs": vector_store.similarity_search(state["input"])
@@ -1129,31 +1109,27 @@ def build_graph() -> StateGraph:
 
     assistant = RunnableLambda(lambda state: {
         **state,
-        "messages": state["messages"] + [SystemMessage(content="You are a helpful assistant.")]
+        "messages": [sys_msg] + state["messages"]
     })
 
-    call_llm = llm.bind_tools(tools)
+    call_llm = llm_with_tools
 
-    # --- Build 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.set_entry_point("retriever")
-
     builder.add_edge("retriever", "assistant")
     builder.add_edge("assistant", "call_llm")
-
     builder.add_conditional_edges("call_llm", should_call_tool, {
         "call_tool": "call_tool",
         "end": None
     })
-
     builder.add_edge("call_tool", "call_llm")
 
-    graph = builder.compile()
-    return graph
+    return builder.compile()