diff --git a/backend/open_webui/retrieval/vector/dbs/chroma.py b/backend/open_webui/retrieval/vector/dbs/chroma.py index 67268e9dc..b7ea5244b 100755 --- a/backend/open_webui/retrieval/vector/dbs/chroma.py +++ b/backend/open_webui/retrieval/vector/dbs/chroma.py @@ -69,7 +69,11 @@ class ChromaClient(VectorDBBase): return self.client.delete_collection(name=collection_name) def search( - self, collection_name: str, vectors: list[list[float | int]], filter: Optional[dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: list[list[float | int]], + filter: Optional[dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: # Search for the nearest neighbor items based on the vectors and return 'limit' number of results. try: diff --git a/backend/open_webui/retrieval/vector/dbs/elasticsearch.py b/backend/open_webui/retrieval/vector/dbs/elasticsearch.py index 46cd1ad92..e209453f5 100644 --- a/backend/open_webui/retrieval/vector/dbs/elasticsearch.py +++ b/backend/open_webui/retrieval/vector/dbs/elasticsearch.py @@ -153,7 +153,11 @@ class ElasticsearchClient(VectorDBBase): # Status: works def search( - self, collection_name: str, vectors: list[list[float]], filter: Optional[dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: list[list[float]], + filter: Optional[dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: query = { "size": limit, diff --git a/backend/open_webui/retrieval/vector/dbs/milvus.py b/backend/open_webui/retrieval/vector/dbs/milvus.py index 69ca3988d..35cf6b382 100644 --- a/backend/open_webui/retrieval/vector/dbs/milvus.py +++ b/backend/open_webui/retrieval/vector/dbs/milvus.py @@ -179,7 +179,11 @@ class MilvusClient(VectorDBBase): ) def search( - self, collection_name: str, vectors: list[list[float | int]], filter: Optional[dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: list[list[float | int]], + filter: Optional[dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: # Search for the nearest neighbor items based on the vectors and return 'limit' number of results. collection_name = collection_name.replace("-", "_") diff --git a/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py b/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py index 5dfad33de..c58189b2a 100644 --- a/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py +++ b/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py @@ -157,7 +157,11 @@ class MilvusClient(VectorDBBase): collection.insert(entities) def search( - self, collection_name: str, vectors: List[List[float]], filter: Optional[Dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: List[List[float]], + filter: Optional[Dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: if not vectors: return None diff --git a/backend/open_webui/retrieval/vector/dbs/opensearch.py b/backend/open_webui/retrieval/vector/dbs/opensearch.py index e6bd99b97..dc9c35805 100644 --- a/backend/open_webui/retrieval/vector/dbs/opensearch.py +++ b/backend/open_webui/retrieval/vector/dbs/opensearch.py @@ -113,7 +113,11 @@ class OpenSearchClient(VectorDBBase): self.client.indices.delete(index=self._get_index_name(collection_name)) def search( - self, collection_name: str, vectors: list[list[float | int]], filter: Optional[dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: list[list[float | int]], + filter: Optional[dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: try: if not self.has_collection(collection_name): diff --git a/backend/open_webui/retrieval/vector/dbs/oracle23ai.py b/backend/open_webui/retrieval/vector/dbs/oracle23ai.py index db7b86943..9f16f82bc 100644 --- a/backend/open_webui/retrieval/vector/dbs/oracle23ai.py +++ b/backend/open_webui/retrieval/vector/dbs/oracle23ai.py @@ -521,7 +521,11 @@ class Oracle23aiClient(VectorDBBase): raise def search( - self, collection_name: str, vectors: List[List[Union[float, int]]], filter: Optional[dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: List[List[Union[float, int]]], + filter: Optional[dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: """ Search for similar vectors in the database. diff --git a/backend/open_webui/retrieval/vector/dbs/pgvector.py b/backend/open_webui/retrieval/vector/dbs/pgvector.py index 6ae09cf05..15430db11 100644 --- a/backend/open_webui/retrieval/vector/dbs/pgvector.py +++ b/backend/open_webui/retrieval/vector/dbs/pgvector.py @@ -477,7 +477,7 @@ class PgvectorClient(VectorDBBase): # Build the lateral subquery for each query vector where_clauses = [DocumentChunk.collection_name == collection_name] - + # Apply metadata filter if provided if filter: for key, value in filter.items(): @@ -487,20 +487,27 @@ class PgvectorClient(VectorDBBase): if PGVECTOR_PGCRYPTO: where_clauses.append( pgcrypto_decrypt( - DocumentChunk.vmetadata, PGVECTOR_PGCRYPTO_KEY, JSONB + DocumentChunk.vmetadata, + PGVECTOR_PGCRYPTO_KEY, + JSONB, )[key].astext.in_([str(v) for v in in_values]) ) else: where_clauses.append( - DocumentChunk.vmetadata[key].astext.in_([str(v) for v in in_values]) + DocumentChunk.vmetadata[key].astext.in_( + [str(v) for v in in_values] + ) ) else: # Handle simple equality: {"field": "value"} if PGVECTOR_PGCRYPTO: where_clauses.append( pgcrypto_decrypt( - DocumentChunk.vmetadata, PGVECTOR_PGCRYPTO_KEY, JSONB - )[key].astext == str(value) + DocumentChunk.vmetadata, + PGVECTOR_PGCRYPTO_KEY, + JSONB, + )[key].astext + == str(value) ) else: where_clauses.append( diff --git a/backend/open_webui/retrieval/vector/dbs/pinecone.py b/backend/open_webui/retrieval/vector/dbs/pinecone.py index 22b9cb98b..fc3c98f8c 100644 --- a/backend/open_webui/retrieval/vector/dbs/pinecone.py +++ b/backend/open_webui/retrieval/vector/dbs/pinecone.py @@ -391,7 +391,11 @@ class PineconeClient(VectorDBBase): ) def search( - self, collection_name: str, vectors: List[List[Union[float, int]]], filter: Optional[dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: List[List[Union[float, int]]], + filter: Optional[dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: """Search for similar vectors in a collection.""" if not vectors or not vectors[0]: diff --git a/backend/open_webui/retrieval/vector/dbs/qdrant.py b/backend/open_webui/retrieval/vector/dbs/qdrant.py index efa33a681..d42984e1d 100644 --- a/backend/open_webui/retrieval/vector/dbs/qdrant.py +++ b/backend/open_webui/retrieval/vector/dbs/qdrant.py @@ -145,7 +145,11 @@ class QdrantClient(VectorDBBase): ) def search( - self, collection_name: str, vectors: list[list[float | int]], filter: Optional[dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: list[list[float | int]], + filter: Optional[dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: # Search for the nearest neighbor items based on the vectors and return 'limit' number of results. if limit is None: diff --git a/backend/open_webui/retrieval/vector/dbs/qdrant_multitenancy.py b/backend/open_webui/retrieval/vector/dbs/qdrant_multitenancy.py index 70ec6d606..f87f85a23 100644 --- a/backend/open_webui/retrieval/vector/dbs/qdrant_multitenancy.py +++ b/backend/open_webui/retrieval/vector/dbs/qdrant_multitenancy.py @@ -254,7 +254,11 @@ class QdrantClient(VectorDBBase): ) def search( - self, collection_name: str, vectors: List[List[float | int]], filter: Optional[Dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: List[List[float | int]], + filter: Optional[Dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: """ Search for the nearest neighbor items based on the vectors with tenant isolation. diff --git a/backend/open_webui/retrieval/vector/dbs/s3vector.py b/backend/open_webui/retrieval/vector/dbs/s3vector.py index d6784f9f2..96e487f11 100644 --- a/backend/open_webui/retrieval/vector/dbs/s3vector.py +++ b/backend/open_webui/retrieval/vector/dbs/s3vector.py @@ -295,7 +295,11 @@ class S3VectorClient(VectorDBBase): raise def search( - self, collection_name: str, vectors: List[List[Union[float, int]]], filter: Optional[dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: List[List[Union[float, int]]], + filter: Optional[dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: """ Search for similar vectors in a collection using multiple query vectors. diff --git a/backend/open_webui/retrieval/vector/dbs/weaviate.py b/backend/open_webui/retrieval/vector/dbs/weaviate.py index 680a6c973..d204e8293 100644 --- a/backend/open_webui/retrieval/vector/dbs/weaviate.py +++ b/backend/open_webui/retrieval/vector/dbs/weaviate.py @@ -159,7 +159,11 @@ class WeaviateClient(VectorDBBase): ) def search( - self, collection_name: str, vectors: List[List[Union[float, int]]], filter: Optional[dict] = None, limit: int = 10 + self, + collection_name: str, + vectors: List[List[Union[float, int]]], + filter: Optional[dict] = None, + limit: int = 10, ) -> Optional[SearchResult]: sane_collection_name = self._sanitize_collection_name(collection_name) if not self.client.collections.exists(sane_collection_name): diff --git a/backend/open_webui/tools/builtin.py b/backend/open_webui/tools/builtin.py index 7c691f2e2..eb3b7cfc9 100644 --- a/backend/open_webui/tools/builtin.py +++ b/backend/open_webui/tools/builtin.py @@ -1624,7 +1624,11 @@ async def query_knowledge_bases( if search_results and search_results.ids and search_results.ids[0]: result_ids = search_results.ids[0] - result_distances = search_results.distances[0] if search_results.distances else [0] * len(result_ids) + result_distances = ( + search_results.distances[0] + if search_results.distances + else [0] * len(result_ids) + ) for knowledge_base_id, distance in zip(result_ids, result_distances): if knowledge_base_id in seen_ids: @@ -1634,7 +1638,9 @@ async def query_knowledge_bases( if len(top_results_heap) < count: heapq.heappush(top_results_heap, (distance, knowledge_base_id)) elif distance > top_results_heap[0][0]: - heapq.heapreplace(top_results_heap, (distance, knowledge_base_id)) + heapq.heapreplace( + top_results_heap, (distance, knowledge_base_id) + ) page_offset += page_size if len(accessible_knowledge_bases.items) < page_size: @@ -1649,16 +1655,17 @@ async def query_knowledge_bases( for distance, knowledge_base_id in sorted_results: knowledge_base = Knowledges.get_knowledge_by_id(knowledge_base_id) if knowledge_base: - matching_knowledge_bases.append({ - "id": knowledge_base.id, - "name": knowledge_base.name, - "description": knowledge_base.description or "", - "similarity": round(distance, 4), - }) + matching_knowledge_bases.append( + { + "id": knowledge_base.id, + "name": knowledge_base.name, + "description": knowledge_base.description or "", + "similarity": round(distance, 4), + } + ) return json.dumps(matching_knowledge_bases, ensure_ascii=False) except Exception as e: log.exception(f"query_knowledge_bases error: {e}") return json.dumps({"error": str(e)}) -