Update agent.py

agent.py

@@ -70,14 +70,21 @@ 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,
+    embedding=embeddings,
     table_name="documents",
     query_name="match_documents_langchain",
+    text_key="content",
+    embedding_key="embedding"
 )
-create_retriever_tool = create_retriever_tool(
-    retriever=vector_store.as_retriever(),
+
+retriever_tool = create_retriever_tool(
+    retriever=vector_store.as_retriever(
+        search_type="similarity",
+        search_kwargs={"k": 5}
+    ),
     name="Question Search",
     description="A tool to retrieve similar questions from a vector store.",
 )
@@ -86,6 +93,7 @@ tools = [
     wiki_search,
     web_search,
     arvix_search,
+    retriever_tool,
 ]
 
 def build_graph():
@@ -109,7 +117,7 @@ def build_graph():
         return {"messages": [sys_msg] + state["messages"] + [example_msg]}
 
     builder = StateGraph(MessagesState)
-    #builder.add_node("retriever", retriever)
+    builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(tools))
     builder.add_edge(START, "retriever")