test

agent.py

@@ -99,30 +99,36 @@ def build_graph():
     llm_with_tools = llm.bind_tools(tools)
 
     # Define nodes
-    def assistant_node(state: MessagesState) -> dict:
-        # Append system message for context
-        messages = [sys_msg] + state["messages"]
-        response = llm_with_tools.invoke(messages)
-        return {"messages": [response]}
+    def assistant(state: MessagesState):
+        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+
 
 
     # Retriever node returns AIMessage
     def retriever(state: MessagesState):
         query = state["messages"][-1].content
-        similar_doc = vector_store.similarity_search(query, k=1)[0]
+        similar_docs = vector_store.similarity_search(query, k=1)
 
-        content = similar_doc.page_content
-        if "Final answer :" in content:
-            answer = content.split("Final answer :")[-1].strip()
+        if similar_docs:
+            reference = similar_docs[0].page_content
+            context_msg = HumanMessage(content=f"Here is a similar question and answer for reference:\n\n{reference}")
         else:
-            answer = content.strip()
-        return {"messages": [AIMessage(content=answer)]}
+            context_msg = HumanMessage(content="No relevant example found.")
+
+        return {
+            "messages": [sys_msg] + state["messages"] + [context_msg]
+        }
+
 
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
 
-    builder.set_entry_point("retriever")
-    builder.set_finish_point("retriever")
+    builder.add_edge(START, "retriever")
+    builder.add_edge("retriever", "assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
 
     # Compile graph
     return builder.compile()