diff --git a/mem0/vector_stores/mongodb.py b/mem0/vector_stores/mongodb.py index 381bfde4e..01cb17bb8 100644 --- a/mem0/vector_stores/mongodb.py +++ b/mem0/vector_stores/mongodb.py @@ -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}'.")