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:
|
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}'.")
|
||||||
|
|||||||
Reference in New Issue
Block a user