118 lines
3.5 KiB
Python
118 lines
3.5 KiB
Python
import hashlib
|
|
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.http.models import (
|
|
PointStruct,
|
|
Filter,
|
|
FieldCondition,
|
|
MatchValue,
|
|
VectorParams,
|
|
Distance,
|
|
)
|
|
from qdrant_client.http.exceptions import UnexpectedResponse
|
|
|
|
from app.config import settings
|
|
|
|
client = QdrantClient(url=settings.qdrant_url)
|
|
COLLECTION_NAME = "tickets"
|
|
_vector_size: int | None = None
|
|
|
|
|
|
def get_vector_size() -> int:
|
|
global _vector_size
|
|
if _vector_size is not None:
|
|
return _vector_size
|
|
|
|
collections = client.get_collections().collections
|
|
for c in collections:
|
|
if c.name == COLLECTION_NAME:
|
|
info = client.get_collection(COLLECTION_NAME)
|
|
_vector_size = info.config.params.vectors.size
|
|
return _vector_size
|
|
|
|
return 0
|
|
|
|
|
|
async def ensure_collection(size: int):
|
|
global _vector_size
|
|
collections = client.get_collections().collections
|
|
exists = any(c.name == COLLECTION_NAME for c in collections)
|
|
|
|
if not exists:
|
|
client.create_collection(
|
|
collection_name=COLLECTION_NAME,
|
|
vectors_config=VectorParams(size=size, distance=Distance.COSINE),
|
|
)
|
|
_vector_size = size
|
|
else:
|
|
get_vector_size()
|
|
|
|
|
|
def _point_id(tenant_id: str, ticket_id: str) -> int:
|
|
return int(hashlib.md5(f"{tenant_id}:{ticket_id}".encode()).hexdigest()[:16], 16)
|
|
|
|
|
|
def get_existing_ticket_ids(tenant_id: str, collection_id: str) -> set[str]:
|
|
try:
|
|
scroll_filter = Filter(must=[
|
|
FieldCondition(key="tenant_id", match=MatchValue(value=tenant_id)),
|
|
FieldCondition(key="collection_id", match=MatchValue(value=collection_id)),
|
|
])
|
|
points, _ = client.scroll(
|
|
collection_name=COLLECTION_NAME,
|
|
limit=10000,
|
|
scroll_filter=scroll_filter,
|
|
with_payload=True,
|
|
with_vectors=False,
|
|
)
|
|
return {p.payload["ticket_id"] for p in points if "ticket_id" in (p.payload or {})}
|
|
except UnexpectedResponse as e:
|
|
if b"doesn't exist" in e.content:
|
|
return set()
|
|
raise
|
|
|
|
|
|
async def upsert_chunks(tenant_id: str, chunks: list[dict], vectors: list[list[float]]):
|
|
size = len(vectors[0])
|
|
await ensure_collection(size)
|
|
|
|
points = []
|
|
for chunk, vector in zip(chunks, vectors):
|
|
points.append(
|
|
PointStruct(
|
|
id=_point_id(tenant_id, chunk["ticket_id"]),
|
|
vector=vector,
|
|
payload={
|
|
"tenant_id": tenant_id,
|
|
"ticket_id": chunk["ticket_id"],
|
|
"client": chunk.get("client", ""),
|
|
"category": chunk.get("category", ""),
|
|
"collection_id": chunk.get("collection_id", ""),
|
|
"search_text": chunk.get("search_text", ""),
|
|
"full_text": chunk.get("full_text", ""),
|
|
},
|
|
)
|
|
)
|
|
|
|
client.upsert(collection_name=COLLECTION_NAME, points=points)
|
|
|
|
|
|
async def search(
|
|
tenant_id: str,
|
|
query_vector: list[float],
|
|
collection_id: str | None = None,
|
|
k: int = 20,
|
|
):
|
|
must_conditions = [FieldCondition(key="tenant_id", match=MatchValue(value=tenant_id))]
|
|
if collection_id:
|
|
must_conditions.append(FieldCondition(key="collection_id", match=MatchValue(value=collection_id)))
|
|
|
|
results = client.query_points(
|
|
collection_name=COLLECTION_NAME,
|
|
query=query_vector,
|
|
query_filter=Filter(must=must_conditions),
|
|
limit=k,
|
|
with_payload=True,
|
|
)
|
|
|
|
return results.points
|