Increase embeding priority for query request

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()

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

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,

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()

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
+        # Execute embedding outside of lock to avoid improve cocurrent
-        embedding = await self.embedding_func([query])
+        embedding = await self.embedding_func(
+            [query], _priority=5
+        )  # higher priority for query
         embedding = embedding[0]
 
         client = await self._get_client()

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)

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],

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())) + "]"