Update agent.py

agent.py

@@ -155,18 +155,28 @@ def build_graph(provider: str = "openai"):
 
     from langchain_core.messages import AIMessage
 
-def retriever(state: MessagesState):
-    query = state["messages"][-1].content
-    results = vector_store.similarity_search(query, k=1)
-    if not results:
-        return {"messages": [AIMessage(content="FINAL ANSWER: No relevant answer found.")]}
-    similar_doc = results[0]
-    content = similar_doc.page_content
-    if "FINAL ANSWER:" in content:
-        answer = content.split("FINAL ANSWER:")[-1].strip()
-        return {"messages": [AIMessage(content=f"FINAL ANSWER: {answer}")]}
-    else:
-        return {"messages": [AIMessage(content=content.strip())]}
+    def retriever(state: MessagesState):
+            query = state["messages"][-1].content
+            results = vector_store.similarity_search(query, k=1)
+            if not results:
+                return {"messages": [AIMessage(content="FINAL ANSWER: No relevant answer found.")]}
+            similar_doc = results[0]
+            content = similar_doc.page_content
+            if "FINAL ANSWER:" in content:
+                answer = content.split("FINAL ANSWER:")[-1].strip()
+                return {"messages": [AIMessage(content=f"FINAL ANSWER: {answer}")]}
+            else:
+                return {"messages": [AIMessage(content=content.strip())]}
+        
+        
+            builder = StateGraph(MessagesState)
+            builder.add_node("retriever", retriever)
+            builder.add_node("assistant", assistant)
+            builder.add_edge(START, "retriever")
+            builder.add_edge("retriever", "assistant")
+            builder.set_finish_point("assistant")
+        
+            return builder.compile()
 
 
     builder = StateGraph(MessagesState)