Update agent.py

agent.py

@@ -19,6 +19,8 @@ from langchain_openai import ChatOpenAI
 
 from langchain.tools import Tool
 from code_interpreter import CodeInterpreter
+from langchain_chroma import Chroma
+from langchain_core.messages import AIMessage
 
 interpreter_instance = CodeInterpreter()
 
@@ -228,16 +230,13 @@ execute_code_multilang_tool = Tool(
 
 ### ======================================== RETRIEVER TOOLS ======================================== ###
 
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") #  dim=768
-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",
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+vector_store = Chroma(
+    collection_name="gaia_dataset",
+    embedding_function=embeddings,
+    persist_directory="./data/chroma_langchain_db",  # Where to save data locally, remove if not necessary
 )
+# It's not going to be used later
 create_retriever_tool = create_retriever_tool(
     retriever=vector_store.as_retriever(),
     name="Question Search",
@@ -312,8 +311,6 @@ def build_graph(provider: str = "huggingface"):
         """Assistant node"""
         return {"messages": [chat_with_tools.invoke(state["messages"])]}
 
-    from langchain_core.messages import AIMessage
-
     def retriever(state: MessagesState):
         query = state["messages"][-1].content
         similar_doc = vector_store.similarity_search(query, k=1)[0]
@@ -324,7 +321,7 @@ def build_graph(provider: str = "huggingface"):
         else:
             answer = content.strip()
 
-        return {"messages": [AIMessage(content=answer)]}
+        return {"messages": [AIMessage(content=answer)]}    
 
     """
     Graph with retriever and tools