replace query_vector args in search method of mongodb vector_stores (#3379)

This commit is contained in:
Rupam Jana
2025-08-29 19:05:25 +05:30
committed by GitHub
parent 84af5ad265
commit 458d7ab8a3
+7 -9
View File
@@ -111,15 +111,13 @@ class MongoDB(VectorStoreBase):
except PyMongoError as e: except PyMongoError as e:
logger.error(f"Error inserting data: {e}") logger.error(f"Error inserting data: {e}")
def search( def search(self, query: str, vectors: List[float], limit=5, filters: Optional[Dict] = None) -> List[OutputData]:
self, query: str, query_vector: List[float], limit=5, filters: Optional[Dict] = None
) -> List[OutputData]:
""" """
Search for similar vectors using the vector search index. Search for similar vectors using the vector search index.
Args: Args:
query (str): Query string 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. limit (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. filters (Dict, optional): Filters to apply to the search.
@@ -141,24 +139,24 @@ class MongoDB(VectorStoreBase):
"index": self.index_name, "index": self.index_name,
"limit": limit, "limit": limit,
"numCandidates": limit, "numCandidates": limit,
"queryVector": query_vector, "queryVector": vectors,
"path": "embedding", "path": "embedding",
} }
}, },
{"$set": {"score": {"$meta": "vectorSearchScore"}}}, {"$set": {"score": {"$meta": "vectorSearchScore"}}},
{"$project": {"embedding": 0}}, {"$project": {"embedding": 0}},
] ]
# Add filter stage if filters are provided # Add filter stage if filters are provided
if filters: if filters:
filter_conditions = [] filter_conditions = []
for key, value in filters.items(): for key, value in filters.items():
filter_conditions.append({"payload." + key: value}) filter_conditions.append({"payload." + key: value})
if filter_conditions: if filter_conditions:
# Add a $match stage after vector search to apply filters # Add a $match stage after vector search to apply filters
pipeline.insert(1, {"$match": {"$and": filter_conditions}}) pipeline.insert(1, {"$match": {"$and": filter_conditions}})
results = list(collection.aggregate(pipeline)) results = list(collection.aggregate(pipeline))
logger.info(f"Vector search completed. Found {len(results)} documents.") logger.info(f"Vector search completed. Found {len(results)} documents.")
except Exception as e: except Exception as e:
@@ -290,7 +288,7 @@ class MongoDB(VectorStoreBase):
filter_conditions.append({"payload." + key: value}) filter_conditions.append({"payload." + key: value})
if filter_conditions: if filter_conditions:
query = {"$and": filter_conditions} query = {"$and": filter_conditions}
cursor = self.collection.find(query).limit(limit) cursor = self.collection.find(query).limit(limit)
results = [OutputData(id=str(doc["_id"]), score=None, payload=doc.get("payload")) for doc in cursor] 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}'.") logger.info(f"Retrieved {len(results)} documents from collection '{self.collection_name}'.")