forked from innovacion/searchbox
41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
from collections.abc import Sequence
|
|
from typing import Any, final
|
|
|
|
from qdrant_client import AsyncQdrantClient, models
|
|
|
|
from .config import Settings
|
|
from .models import SearchRow
|
|
|
|
|
|
@final
|
|
class QdrantEngine:
|
|
def __init__(self) -> None:
|
|
self.settings = Settings() # type: ignore[reportCallIssue]
|
|
self.client = AsyncQdrantClient(
|
|
url=self.settings.url, api_key=self.settings.api_key
|
|
)
|
|
|
|
async def semantic_search(
|
|
self,
|
|
embedding: Sequence[float] | models.NamedVector,
|
|
collection: str,
|
|
limit: int = 10,
|
|
conditions: Any | None = None,
|
|
threshold: float | None = None,
|
|
) -> list[SearchRow]:
|
|
points = await self.client.search(
|
|
collection_name=collection,
|
|
query_vector=embedding,
|
|
query_filter=conditions,
|
|
limit=limit,
|
|
with_payload=True,
|
|
with_vectors=False,
|
|
score_threshold=threshold,
|
|
)
|
|
|
|
return [
|
|
SearchRow(chunk_id=str(point.id), score=point.score, payload=point.payload)
|
|
for point in points
|
|
if point.payload is not None
|
|
]
|