fix: Update milvus.py (#19602)

* Update milvus.py

* Update milvus.py

* Update milvus.py

* Update milvus.py

* Update milvus.py

---------

Co-authored-by: Tim Baek <tim@openwebui.com>
This commit is contained in:
Classic298 2025-12-02 21:30:31 +01:00 committed by GitHub
parent 192c2af7ba
commit 12f237ff80
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -200,23 +200,24 @@ class MilvusClient(VectorDBBase):
def query(self, collection_name: str, filter: dict, limit: int = -1): def query(self, collection_name: str, filter: dict, limit: int = -1):
connections.connect(uri=MILVUS_URI, token=MILVUS_TOKEN, db_name=MILVUS_DB) connections.connect(uri=MILVUS_URI, token=MILVUS_TOKEN, db_name=MILVUS_DB)
# Construct the filter string for querying
collection_name = collection_name.replace("-", "_") collection_name = collection_name.replace("-", "_")
if not self.has_collection(collection_name): if not self.has_collection(collection_name):
log.warning( log.warning(
f"Query attempted on non-existent collection: {self.collection_prefix}_{collection_name}" f"Query attempted on non-existent collection: {self.collection_prefix}_{collection_name}"
) )
return None return None
filter_string = " && ".join(
[ filter_expressions = []
f'metadata["{key}"] == {json.dumps(value)}' for key, value in filter.items():
for key, value in filter.items() if isinstance(value, str):
] filter_expressions.append(f'metadata["{key}"] == "{value}"')
) else:
filter_expressions.append(f'metadata["{key}"] == {value}')
filter_string = " && ".join(filter_expressions)
collection = Collection(f"{self.collection_prefix}_{collection_name}") collection = Collection(f"{self.collection_prefix}_{collection_name}")
collection.load() collection.load()
all_results = []
try: try:
log.info( log.info(
@ -224,24 +225,25 @@ class MilvusClient(VectorDBBase):
) )
iterator = collection.query_iterator( iterator = collection.query_iterator(
filter=filter_string, expr=filter_string,
output_fields=[ output_fields=[
"id", "id",
"data", "data",
"metadata", "metadata",
], ],
limit=limit, # Pass the limit directly; -1 means no limit. limit=limit if limit > 0 else -1,
) )
all_results = []
while True: while True:
result = iterator.next() batch = iterator.next()
if not result: if not batch:
iterator.close() iterator.close()
break break
all_results += result all_results.extend(batch)
log.info(f"Total results from query: {len(all_results)}") log.debug(f"Total results from query: {len(all_results)}")
return self._result_to_get_result([all_results]) return self._result_to_get_result([all_results] if all_results else [[]])
except Exception as e: except Exception as e:
log.exception( log.exception(