replace query_vector args in search method of mongodb vector_stores (#3379)
This commit is contained in:
@@ -111,15 +111,13 @@ class MongoDB(VectorStoreBase):
|
||||
except PyMongoError as e:
|
||||
logger.error(f"Error inserting data: {e}")
|
||||
|
||||
def search(
|
||||
self, query: str, query_vector: List[float], limit=5, filters: Optional[Dict] = None
|
||||
) -> List[OutputData]:
|
||||
def search(self, query: str, vectors: List[float], limit=5, filters: Optional[Dict] = None) -> List[OutputData]:
|
||||
"""
|
||||
Search for similar vectors using the vector search index.
|
||||
|
||||
Args:
|
||||
query (str): Query string
|
||||
query_vector (List[float]): Query vector.
|
||||
vectors (List[float]): Query vector.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search.
|
||||
|
||||
@@ -141,24 +139,24 @@ class MongoDB(VectorStoreBase):
|
||||
"index": self.index_name,
|
||||
"limit": limit,
|
||||
"numCandidates": limit,
|
||||
"queryVector": query_vector,
|
||||
"queryVector": vectors,
|
||||
"path": "embedding",
|
||||
}
|
||||
},
|
||||
{"$set": {"score": {"$meta": "vectorSearchScore"}}},
|
||||
{"$project": {"embedding": 0}},
|
||||
]
|
||||
|
||||
|
||||
# Add filter stage if filters are provided
|
||||
if filters:
|
||||
filter_conditions = []
|
||||
for key, value in filters.items():
|
||||
filter_conditions.append({"payload." + key: value})
|
||||
|
||||
|
||||
if filter_conditions:
|
||||
# Add a $match stage after vector search to apply filters
|
||||
pipeline.insert(1, {"$match": {"$and": filter_conditions}})
|
||||
|
||||
|
||||
results = list(collection.aggregate(pipeline))
|
||||
logger.info(f"Vector search completed. Found {len(results)} documents.")
|
||||
except Exception as e:
|
||||
@@ -290,7 +288,7 @@ class MongoDB(VectorStoreBase):
|
||||
filter_conditions.append({"payload." + key: value})
|
||||
if filter_conditions:
|
||||
query = {"$and": filter_conditions}
|
||||
|
||||
|
||||
cursor = self.collection.find(query).limit(limit)
|
||||
results = [OutputData(id=str(doc["_id"]), score=None, payload=doc.get("payload")) for doc in cursor]
|
||||
logger.info(f"Retrieved {len(results)} documents from collection '{self.collection_name}'.")
|
||||
|
||||
Reference in New Issue
Block a user