CyberAssassin commited on
Commit
5339f28
·
verified ·
1 Parent(s): b11b3e6

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +19 -8
agent.py CHANGED
@@ -126,21 +126,32 @@ embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-b
126
  supabase: Client = create_client(
127
  os.environ.get("SUPABASE_URL"),
128
  os.environ.get("SUPABASE_SERVICE_KEY"))
 
129
  vector_store = SupabaseVectorStore(
130
  client=supabase,
131
  embedding=embeddings,
132
  table_name="documents",
133
- query_name="match_documents", # Your unified function name
 
 
 
 
 
134
  )
 
 
 
 
 
 
 
 
 
 
135
  retriever_tool = create_retriever_tool(
136
- retriever=vector_store.as_retriever(
137
- search_kwargs={
138
- "k": 3,
139
- "filter": {}
140
- }
141
- ),
142
  name="document_retriever",
143
- description="Searches for similar documents",
144
  )
145
 
146
 
 
126
  supabase: Client = create_client(
127
  os.environ.get("SUPABASE_URL"),
128
  os.environ.get("SUPABASE_SERVICE_KEY"))
129
+ # Initialize vector store with proper column mappings
130
  vector_store = SupabaseVectorStore(
131
  client=supabase,
132
  embedding=embeddings,
133
  table_name="documents",
134
+ query_name="match_documents",
135
+ columns={
136
+ "id": "document_id", # Maps to function output
137
+ "content": "document_content", # Maps to function output
138
+ "metadata": "document_metadata" # Maps to function output
139
+ }
140
  )
141
+
142
+ # Create retriever with search parameters
143
+ retriever = vector_store.as_retriever(
144
+ search_kwargs={
145
+ "k": 3,
146
+ "filter": {}
147
+ }
148
+ )
149
+
150
+ # Create tool with the configured retriever
151
  retriever_tool = create_retriever_tool(
152
+ retriever=retriever,
 
 
 
 
 
153
  name="document_retriever",
154
+ description="Searches for similar documents in knowledge base",
155
  )
156
 
157