graph-dependency-analyzer 0.2.0__py3-none-any.whl

src/infrastructure/persistence/qdrant_vector_repository.py

@@ -0,0 +1,357 @@
+"""
+Qdrant Vector Repository Implementation (Task 5.2)
+
+Implements persistence layer for code chunk embeddings in Qdrant.
+Supports batch upsert with 1000 points per batch for optimal performance.
+"""
+from dataclasses import dataclass
+from uuid import uuid4
+from typing import Any
+import structlog
+
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue, MatchAny
+from qdrant_client.http.exceptions import UnexpectedResponse
+
+logger = structlog.get_logger(__name__)
+
+
+@dataclass
+class VectorStorageResult:
+    """Result of storing embeddings in Qdrant"""
+    points_upserted: int
+
+
+class QdrantVectorRepository:
+    """
+    Qdrant-based repository for storing code chunk embeddings.
+
+    Implements:
+    - Collection management (create if not exists)
+    - Batch upsert with 1000 points per batch
+    - 1536-dimensional embeddings (text-embedding-3-small)
+    - Complete metadata payload
+    - Unique repository listing
+
+    Requirements: 9.1-9.6
+    """
+
+    def __init__(
+        self,
+        client: QdrantClient,
+        collection_name: str = "code_chunks",
+        batch_size: int = 1000,
+        vector_size: int = 1536,
+    ):
+        """
+        Initialize repository with Qdrant client.
+
+        Args:
+            client: QdrantClient instance
+            collection_name: Name of Qdrant collection
+            batch_size: Batch size for upsert operations (default 1000)
+            vector_size: Embedding dimension — must match the provider model
+                         (1536 for OpenAI/Flow, 768 for nomic-embed-text, etc.)
+        """
+        self.client = client
+        self.collection_name = collection_name
+        self.batch_size = batch_size
+        self.vector_size = vector_size
+
+    async def upsert_batch(
+        self,
+        chunks: list,
+        embeddings: list[list[float]],
+        batch_size: int | None = None
+    ) -> VectorStorageResult:
+        """
+        Upsert code chunks with embeddings to Qdrant.
+
+        Uses batch processing for performance. Each point contains:
+        - UUID v4 id (from chunk.id)
+        - 1536-dimensional embedding vector
+        - Complete metadata payload
+
+        Args:
+            chunks: List of CodeChunk entities
+            embeddings: List of embedding vectors (1536-dim each)
+            batch_size: Override default batch size (optional)
+
+        Returns:
+            VectorStorageResult with count of points upserted
+
+        Raises:
+            ValueError: If validation fails (dimensions, length mismatch)
+        """
+        # Validation
+        self._validate_inputs(chunks, embeddings)
+
+        # Use provided batch_size or default
+        effective_batch_size = batch_size or self.batch_size
+
+        # Ensure collection exists
+        await self.ensure_collection_exists()
+
+        # Process in batches
+        total_upserted = 0
+        for i in range(0, len(chunks), effective_batch_size):
+            batch_chunks = chunks[i:i + effective_batch_size]
+            batch_embeddings = embeddings[i:i + effective_batch_size]
+
+            # Create points for this batch
+            points = self._create_points(batch_chunks, batch_embeddings)
+
+            # Upsert to Qdrant
+            self.client.upsert(
+                collection_name=self.collection_name,
+                points=points
+            )
+
+            total_upserted += len(points)
+
+        logger.info(
+            "embeddings_upserted",
+            collection=self.collection_name,
+            points_upserted=total_upserted,
+            batch_count=(len(chunks) + effective_batch_size - 1) // effective_batch_size
+        )
+
+        return VectorStorageResult(points_upserted=total_upserted)
+
+    def _validate_inputs(
+        self,
+        chunks: list,
+        embeddings: list[list[float]]
+    ) -> None:
+        """
+        Validate chunks and embeddings inputs.
+
+        Args:
+            chunks: List of CodeChunk
+            embeddings: List of embedding vectors
+
+        Raises:
+            ValueError: If validation fails
+        """
+        # Validate lengths match
+        if len(chunks) != len(embeddings):
+            raise ValueError(
+                f"Chunks and embeddings must have same length. "
+                f"Got {len(chunks)} chunks and {len(embeddings)} embeddings."
+            )
+
+        # Validate embedding dimensions — uses self.vector_size set at init
+        for i, embedding in enumerate(embeddings):
+            if len(embedding) != self.vector_size:
+                raise ValueError(
+                    f"Embedding at index {i} must have {self.vector_size} dimensions, "
+                    f"got {len(embedding)}. "
+                    f"Check QDRANT_VECTOR_SIZE in .env matches your embedding model."
+                )
+
+    def _create_points(
+        self,
+        chunks: list,
+        embeddings: list[list[float]]
+    ) -> list[PointStruct]:
+        """
+        Create Qdrant PointStruct from chunks and embeddings.
+
+        Args:
+            chunks: List of CodeChunk
+            embeddings: List of embedding vectors
+
+        Returns:
+            List of PointStruct ready for upsert
+        """
+        points = []
+
+        for chunk, embedding in zip(chunks, embeddings):
+            # Build payload from chunk metadata
+            payload = {
+                "repository": chunk.metadata.get("repository", ""),
+                "file_path": chunk.metadata.get("file_path", ""),
+                "language": chunk.metadata.get("language", "java"),
+                "type": chunk.chunk_type,
+                "name": chunk.metadata.get("name", ""),
+                "qualified_name": chunk.metadata.get("qualified_name", ""),
+                "content": chunk.content,
+                "line_start": chunk.metadata.get("line_start", 0),
+                "line_end": chunk.metadata.get("line_end", 0),
+            }
+
+            # Add optional metadata fields
+            if "parameters" in chunk.metadata:
+                payload["parameters"] = chunk.metadata["parameters"]
+            if "return_type" in chunk.metadata:
+                payload["return_type"] = chunk.metadata["return_type"]
+            if "parent_class" in chunk.metadata:
+                payload["parent_class"] = chunk.metadata["parent_class"]
+
+            # Create point with UUID from chunk
+            point = PointStruct(
+                id=str(chunk.id),
+                vector=embedding,
+                payload=payload
+            )
+
+            points.append(point)
+
+        return points
+
+    async def search_similar(
+        self,
+        query_vector: list[float],
+        limit: int = 10,
+        repositories: list[str] | None = None,
+        score_threshold: float = 0.35,
+    ) -> list[dict]:
+        """
+        Search for code chunks semantically similar to the query vector.
+
+        Args:
+            query_vector: 1536-dimensional embedding of the search query
+            limit: Maximum number of results to return
+            repositories: Optional list of repo names to filter by
+            score_threshold: Minimum cosine similarity score (0-1)
+
+        Returns:
+            List of dicts with keys: score, repository, file_path,
+            qualified_name, name, type, content, line_start, line_end
+        """
+        # Build optional filter for repository
+        query_filter = None
+        if repositories:
+            if len(repositories) == 1:
+                query_filter = Filter(
+                    must=[FieldCondition(
+                        key="repository",
+                        match=MatchValue(value=repositories[0])
+                    )]
+                )
+            else:
+                # MatchAny matches any of the given values — equivalent to OR
+                query_filter = Filter(
+                    must=[FieldCondition(
+                        key="repository",
+                        match=MatchAny(any=repositories)
+                    )]
+                )
+
+        try:
+            # qdrant-client >= 1.7 uses query_points; older versions use search.
+            # Try query_points first, fall back to search.
+            if hasattr(self.client, "query_points"):
+                response = self.client.query_points(
+                    collection_name=self.collection_name,
+                    query=query_vector,
+                    limit=limit,
+                    query_filter=query_filter,
+                    score_threshold=score_threshold,
+                    with_payload=True,
+                )
+                results = response.points
+            else:
+                results = self.client.search(
+                    collection_name=self.collection_name,
+                    query_vector=query_vector,
+                    limit=limit,
+                    query_filter=query_filter,
+                    score_threshold=score_threshold,
+                    with_payload=True,
+                    with_vectors=False,
+                )
+        except Exception:
+            return []
+
+        hits = []
+        for point in results:
+            payload = point.payload or {}
+            hits.append({
+                "score":          round(point.score, 4),
+                "repository":     payload.get("repository", ""),
+                "file_path":      payload.get("file_path", ""),
+                "qualified_name": payload.get("qualified_name", ""),
+                "name":           payload.get("name", ""),
+                "type":           payload.get("type", ""),
+                "language":       payload.get("language", ""),
+                "line_start":     payload.get("line_start", 0),
+                "line_end":       payload.get("line_end", 0),
+                "content":        (payload.get("content", "")[:500]
+                                   if payload.get("content") else ""),
+            })
+
+        logger.info(
+            "semantic_search_completed",
+            results=len(hits),
+            repositories=repositories,
+            score_threshold=score_threshold,
+        )
+        return hits
+
+    async def get_unique_repositories(self) -> list[str]:
+        """
+        Get list of unique repository names indexed in Qdrant.
+
+        Returns:
+            List of unique repository names
+
+        Raises:
+            Exception: If Qdrant query fails
+        """
+        # Scroll through all points to collect unique repositories
+        # Using scroll instead of search since we want all data
+        records, next_page = self.client.scroll(
+            collection_name=self.collection_name,
+            limit=10000,  # Large limit to get all in one call
+            with_payload=["repository"],
+            with_vectors=False  # Don't need vectors for this
+        )
+
+        # Extract unique repository names
+        repositories = set()
+        for record in records:
+            if "repository" in record.payload:
+                repositories.add(record.payload["repository"])
+
+        logger.debug(
+            "unique_repositories_retrieved",
+            collection=self.collection_name,
+            count=len(repositories)
+        )
+
+        return sorted(list(repositories))
+
+    async def ensure_collection_exists(self) -> None:
+        """
+        Ensure Qdrant collection exists, create if not.
+
+        Creates collection with:
+        - vector_size: 1536 (text-embedding-3-small)
+        - distance: Cosine
+
+        Raises:
+            Exception: If collection creation fails
+        """
+        try:
+            # Try to get collection info (QdrantClient is synchronous)
+            self.client.get_collection(self.collection_name)
+            logger.debug("collection_exists", collection=self.collection_name)
+
+        except UnexpectedResponse as e:
+            if e.status_code == 404:
+                # Collection doesn't exist, create it
+                logger.info("creating_collection", collection=self.collection_name)
+
+                self.client.create_collection(
+                    collection_name=self.collection_name,
+                    vectors_config=VectorParams(
+                        size=self.vector_size,
+                        distance=Distance.COSINE
+                    )
+                )
+
+                logger.info("collection_created", collection=self.collection_name)
+            else:
+                # Other error, re-raise
+                raise