to 22

backend/src/infrastructure/repositories/qdrant/__init__.py

@@ -0,0 +1,4 @@
+"""
+Qdrant repositories
+"""
+

backend/src/infrastructure/repositories/qdrant/vector_repository.py

@@ -0,0 +1,84 @@
+"""
+Qdrant реализация векторного хранилища
+"""
+from typing import Sequence
+from uuid import UUID
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
+from src.domain.entities.chunk import DocumentChunk
+from src.domain.repositories.vector_repository import IVectorRepository
+
+
+class QdrantVectorRepository(IVectorRepository):
+    def __init__(
+        self,
+        client: QdrantClient,
+        collection_name: str = "documents",
+        vector_size: int = 768,
+    ):
+        self.client = client
+        self.collection_name = collection_name
+        self.vector_size = vector_size
+        self._ensure_collection()
+
+    def _ensure_collection(self) -> None:
+        """Создает коллекцию при отсутствии"""
+        if self.collection_name in [c.name for c in self.client.get_collections().collections]:
+            return
+        self.client.create_collection(
+            collection_name=self.collection_name,
+            vectors_config=VectorParams(size=self.vector_size, distance=Distance.COSINE),
+        )
+
+    async def upsert_chunks(
+        self,
+        chunks: Sequence[DocumentChunk],
+        embeddings: Sequence[list[float]],
+        model_version: str,
+    ) -> None:
+        points = []
+        for chunk, vector in zip(chunks, embeddings):
+            points.append(
+                PointStruct(
+                    id=str(chunk.chunk_id),
+                    vector=vector,
+                    payload={
+                        "document_id": str(chunk.document_id),
+                        "collection_id": str(chunk.collection_id),
+                        "content": chunk.content,
+                        "order": chunk.order,
+                        "model_version": model_version,
+                        "title": chunk.metadata.get("title", ""),
+                    },
+                )
+            )
+        self.client.upsert(collection_name=self.collection_name, points=points)
+
+    async def search(
+        self,
+        query_embedding: list[float],
+        collection_id: UUID,
+        limit: int = 20,
+    ) -> list[tuple[DocumentChunk, float]]:
+        res = self.client.search(
+            collection_name=self.collection_name,
+            query_vector=query_embedding,
+            query_filter=Filter(
+                must=[FieldCondition(key="collection_id", match=MatchValue(value=str(collection_id)))]
+            ),
+            limit=limit,
+        )
+        results: list[tuple[DocumentChunk, float]] = []
+        for hit in res:
+            payload = hit.payload or {}
+            chunk = DocumentChunk(
+                document_id=UUID(payload["document_id"]),
+                collection_id=UUID(payload["collection_id"]),
+                content=payload.get("content", ""),
+                chunk_id=UUID(hit.id),
+                order=payload.get("order", 0),
+                metadata={"title": payload.get("title", ""), "model_version": payload.get("model_version", "")},
+            )
+            results.append((chunk, hit.score))
+        return results
+