Update agent.py

agent.py

@@ -15,6 +15,8 @@ from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain.tools.retriever import create_retriever_tool
 from supabase.client import Client, create_client
+from langchain_core.documents import Document
+from langchain_community.vectorstores import Chroma
 
 load_dotenv()
 
@@ -126,11 +128,17 @@ embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-b
 supabase: Client = create_client(
     os.environ.get("SUPABASE_URL"), 
     os.environ.get("SUPABASE_SERVICE_KEY"))
-vector_store = SupabaseVectorStore(
-    client=supabase,
-    embedding= embeddings,
-    table_name="documents",
-    query_name="match_documents_langchain",
+# vector_store = SupabaseVectorStore(
+#     client=supabase,
+#     embedding= embeddings,
+#     table_name="documents",
+#     query_name="match_documents_langchain",
+# )
+
+vector_store = Chroma(
+    collection_name = 'documents',
+    embedding_function=embeddings,
+    persist_directory="./vector_db"  # Omit for in-memory only
 )
 create_retriever_tool = create_retriever_tool(
     retriever=vector_store.as_retriever(),
@@ -182,6 +190,7 @@ def build_graph(provider: str = "groq"):
     def retriever(state: MessagesState):
         """Retriever node"""
         similar_question = vector_store.similarity_search(state["messages"][0].content)
+        # results = vector_store.similarity_search(query=state["messages"][0].content,k=3)
         example_msg = HumanMessage(
             content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
         )