Update agent.py

agent.py

@@ -1,9 +1,11 @@
 import os
 
 from dotenv import load_dotenv
+from langchain_community.vectorstores import Chroma
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.tools import tool
-from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
+from langchain_huggingface import (ChatHuggingFace, HuggingFaceEmbeddings,
+                                   HuggingFaceEndpoint)
 from langgraph.graph import START, MessagesState, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
 
@@ -41,6 +43,16 @@ YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma sepa
 Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
 """
 
+# System message
+sys_msg = SystemMessage(content=system_prompt)
+
+# Embeddings + Chroma Vector Store
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+vector_store = Chroma(
+    collection_name="langgraph-documents",
+    embedding_function=embeddings,
+    persist_directory="chroma_db"  # Use a persistent directory
+)
 
 def build_graph():
     """Build the graph""" 
@@ -71,17 +83,22 @@ def build_graph():
         """Assistant node"""
         return {"messages": [llm_with_tools.invoke([system_prompt] + state["messages"])]}
     
-    builder = StateGraph(MessagesState)
+    def retriever(state: MessagesState):
+        similar = vector_store.similarity_search(state["messages"][0].content)
+        if similar:
+            example_msg = HumanMessage(content=f"Here is a similar question:\n\n{similar[0].page_content}")
+            return {"messages": [sys_msg] + state["messages"] + [example_msg]}
+        return {"messages": [sys_msg] + state["messages"]}
 
-    # Nodes
+    builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(tools))
-
-    # Edges
-    builder.add_edge(START, "assistant")
+    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()