Update agent.py

agent.py

@@ -685,28 +685,28 @@ def process_question(question: str):
 
 
 # Build graph function
-def build_graph():
-    """Build the graph based on provider"""
+def build_graph(provider: str, model_config: dict):
+    from langgraph.prebuilt.tool_node import ToolNode
+
     llm = get_llm(provider, model_config)
     llm_with_tools = llm.bind_tools(tools)
 
-    # Node
+    sys_msg = SystemMessage(content="You are a helpful assistant.")
+
     def assistant(state: MessagesState):
-        """Assistant node"""
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
-    
+
     def retriever(state: MessagesState):
         user_query = state["messages"][0].content
         similar_docs = vector_store.similarity_search(user_query)
-    
+
         if not similar_docs:
-            print("No similar docs found in FAISS. Using wiki_search.")
-            wiki_result = wiki_search.invoke(user_query)
+            wiki_result = wiki_tool.run(user_query)
             return {
                 "messages": [
                     sys_msg,
                     state["messages"][0],
-                    HumanMessage(content=f"Using Wikipedia search:\n\n{wiki_result['wiki_results']}")
+                    HumanMessage(content=f"Using Wikipedia search:\n\n{wiki_result}")
                 ]
             }
         else:
@@ -714,22 +714,23 @@ def build_graph():
                 "messages": [
                     sys_msg,
                     state["messages"][0],
-                    HumanMessage(content=f"Reference question:\n\n{similar_docs[0].page_content}")
+                    HumanMessage(content=f"Reference:\n\n{similar_docs[0].page_content}")
                 ]
             }
 
+    def tools_condition(state: MessagesState) -> str:
+        if "use tool" in state["messages"][-1].content.lower():
+            return "tools"
+        else:
+            return END
 
-    
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(tools))
-    builder.add_edge(START, "retriever")
+    builder.set_entry_point("retriever")
     builder.add_edge("retriever", "assistant")
-    builder.add_conditional_edges(
-        "assistant",
-        tools_condition,
-    )
+    builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")
 
     # Compile graph