From ca633865469297ccf5e2d1631197c96b5565fcbc Mon Sep 17 00:00:00 2001 From: yangdx Date: Mon, 28 Apr 2025 20:10:39 +0800 Subject: [PATCH] Increase embeding priority for query request --- lightrag/kg/chroma_impl.py | 4 +++- lightrag/kg/faiss_impl.py | 4 +++- lightrag/kg/milvus_impl.py | 4 +++- lightrag/kg/mongo_impl.py | 4 +++- lightrag/kg/nano_vector_db_impl.py | 6 ++++-- lightrag/kg/postgres_impl.py | 4 +++- lightrag/kg/qdrant_impl.py | 4 +++- lightrag/kg/tidb_impl.py | 4 +++- 8 files changed, 25 insertions(+), 9 deletions(-) diff --git a/lightrag/kg/chroma_impl.py b/lightrag/kg/chroma_impl.py index 020e358f..627cf480 100644 --- a/lightrag/kg/chroma_impl.py +++ b/lightrag/kg/chroma_impl.py @@ -161,7 +161,9 @@ class ChromaVectorDBStorage(BaseVectorStorage): self, query: str, top_k: int, ids: list[str] | None = None ) -> list[dict[str, Any]]: try: - embedding = await self.embedding_func([query]) + embedding = await self.embedding_func( + [query], _priority=5 + ) # higher priority for query results = self._collection.query( query_embeddings=embedding.tolist() diff --git a/lightrag/kg/faiss_impl.py b/lightrag/kg/faiss_impl.py index 3022e6f6..d4e6ad01 100644 --- a/lightrag/kg/faiss_impl.py +++ b/lightrag/kg/faiss_impl.py @@ -175,7 +175,9 @@ class FaissVectorDBStorage(BaseVectorStorage): """ Search by a textual query; returns top_k results with their metadata + similarity distance. """ - embedding = await self.embedding_func([query]) + embedding = await self.embedding_func( + [query], _priority=5 + ) # higher priority for query # embedding is shape (1, dim) embedding = np.array(embedding, dtype=np.float32) faiss.normalize_L2(embedding) # we do in-place normalization diff --git a/lightrag/kg/milvus_impl.py b/lightrag/kg/milvus_impl.py index 2cff0079..74e0f039 100644 --- a/lightrag/kg/milvus_impl.py +++ b/lightrag/kg/milvus_impl.py @@ -104,7 +104,9 @@ class MilvusVectorDBStorage(BaseVectorStorage): async def query( self, query: str, top_k: int, ids: list[str] | None = None ) -> list[dict[str, Any]]: - embedding = await self.embedding_func([query]) + embedding = await self.embedding_func( + [query], _priority=5 + ) # higher priority for query results = self._client.search( collection_name=self.namespace, data=embedding, diff --git a/lightrag/kg/mongo_impl.py b/lightrag/kg/mongo_impl.py index bc4bb240..f4b15d60 100644 --- a/lightrag/kg/mongo_impl.py +++ b/lightrag/kg/mongo_impl.py @@ -1032,7 +1032,9 @@ class MongoVectorDBStorage(BaseVectorStorage): ) -> list[dict[str, Any]]: """Queries the vector database using Atlas Vector Search.""" # Generate the embedding - embedding = await self.embedding_func([query]) + embedding = await self.embedding_func( + [query], _priority=5 + ) # higher priority for query # Convert numpy array to a list to ensure compatibility with MongoDB query_vector = embedding[0].tolist() diff --git a/lightrag/kg/nano_vector_db_impl.py b/lightrag/kg/nano_vector_db_impl.py index 3003eea0..589dd9ee 100644 --- a/lightrag/kg/nano_vector_db_impl.py +++ b/lightrag/kg/nano_vector_db_impl.py @@ -124,8 +124,10 @@ class NanoVectorDBStorage(BaseVectorStorage): async def query( self, query: str, top_k: int, ids: list[str] | None = None ) -> list[dict[str, Any]]: - # Execute embedding outside of lock to avoid long lock times - embedding = await self.embedding_func([query]) + # Execute embedding outside of lock to avoid improve cocurrent + embedding = await self.embedding_func( + [query], _priority=5 + ) # higher priority for query embedding = embedding[0] client = await self._get_client() diff --git a/lightrag/kg/postgres_impl.py b/lightrag/kg/postgres_impl.py index 07b30d08..cb302e8c 100644 --- a/lightrag/kg/postgres_impl.py +++ b/lightrag/kg/postgres_impl.py @@ -644,7 +644,9 @@ class PGVectorStorage(BaseVectorStorage): async def query( self, query: str, top_k: int, ids: list[str] | None = None ) -> list[dict[str, Any]]: - embeddings = await self.embedding_func([query]) + embeddings = await self.embedding_func( + [query], _priority=5 + ) # higher priority for query embedding = embeddings[0] embedding_string = ",".join(map(str, embedding)) # Use parameterized document IDs (None means search across all documents) diff --git a/lightrag/kg/qdrant_impl.py b/lightrag/kg/qdrant_impl.py index d758ca5c..e0d0c9c3 100644 --- a/lightrag/kg/qdrant_impl.py +++ b/lightrag/kg/qdrant_impl.py @@ -124,7 +124,9 @@ class QdrantVectorDBStorage(BaseVectorStorage): async def query( self, query: str, top_k: int, ids: list[str] | None = None ) -> list[dict[str, Any]]: - embedding = await self.embedding_func([query]) + embedding = await self.embedding_func( + [query], _priority=5 + ) # higher priority for query results = self._client.search( collection_name=self.namespace, query_vector=embedding[0], diff --git a/lightrag/kg/tidb_impl.py b/lightrag/kg/tidb_impl.py index 33ad1567..fb52b744 100644 --- a/lightrag/kg/tidb_impl.py +++ b/lightrag/kg/tidb_impl.py @@ -390,7 +390,9 @@ class TiDBVectorDBStorage(BaseVectorStorage): self, query: str, top_k: int, ids: list[str] | None = None ) -> list[dict[str, Any]]: """Search from tidb vector""" - embeddings = await self.embedding_func([query]) + embeddings = await self.embedding_func( + [query], _priority=5 + ) # higher priority for query embedding = embeddings[0] embedding_string = "[" + ", ".join(map(str, embedding.tolist())) + "]"