Increase embeding priority for query request
This commit is contained in:
@@ -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()
|
||||
|
@@ -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
|
||||
|
@@ -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,
|
||||
|
@@ -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()
|
||||
|
@@ -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()
|
||||
|
@@ -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)
|
||||
|
@@ -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],
|
||||
|
@@ -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())) + "]"
|
||||
|
Reference in New Issue
Block a user