code-graph-builder 0.2.0__py3-none-any.whl

code_graph_builder/embeddings/vector_store.py

@@ -0,0 +1,584 @@
+"""Vector store for code embeddings.
+
+This module provides abstract base class and implementations for storing
+and searching code embeddings.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from loguru import logger
+
+if TYPE_CHECKING:
+    from ..types import PropertyDict
+
+
+def cosine_similarity(a: list[float], b: list[float]) -> float:
+    """Calculate cosine similarity between two vectors.
+
+    Args:
+        a: First vector
+        b: Second vector
+
+    Returns:
+        Cosine similarity (-1 to 1)
+
+    Raises:
+        ValueError: If vectors have different lengths or are zero vectors
+    """
+    import math
+
+    if len(a) != len(b):
+        raise ValueError(f"Vectors have different lengths: {len(a)} vs {len(b)}")
+
+    dot_product = sum(x * y for x, y in zip(a, b))
+    norm_a = math.sqrt(sum(x * x for x in a))
+    norm_b = math.sqrt(sum(x * x for x in b))
+
+    if norm_a == 0 or norm_b == 0:
+        raise ValueError("Cannot compute cosine similarity for zero vectors")
+
+    return dot_product / (norm_a * norm_b)
+
+
+@dataclass
+class VectorRecord:
+    """A record in the vector store.
+
+    Attributes:
+        node_id: Unique node identifier
+        qualified_name: Fully qualified name of the code entity
+        embedding: Embedding vector
+        metadata: Additional metadata
+    """
+
+    node_id: int
+    qualified_name: str
+    embedding: list[float]
+    metadata: dict[str, str | int | float | None] = field(default_factory=dict)
+
+
+@dataclass
+class SearchResult:
+    """Result from vector similarity search.
+
+    Attributes:
+        node_id: Node identifier
+        score: Similarity score (0-1, higher is better)
+        qualified_name: Fully qualified name
+    """
+
+    node_id: int
+    score: float
+    qualified_name: str
+
+
+class VectorStore(ABC):
+    """Abstract base class for vector stores."""
+
+    @abstractmethod
+    def store_embedding(
+        self,
+        node_id: int,
+        qualified_name: str,
+        embedding: list[float],
+        metadata: PropertyDict | None = None,
+        **kwargs,
+    ) -> None:
+        """Store an embedding vector.
+
+        Args:
+            node_id: Unique node identifier
+            qualified_name: Fully qualified name of the code entity
+            embedding: Embedding vector
+            metadata: Additional metadata
+            **kwargs: Additional keyword arguments (implementation-specific)
+        """
+        ...
+
+    @abstractmethod
+    def store_embeddings_batch(
+        self,
+        records: list[VectorRecord],
+    ) -> None:
+        """Store multiple embeddings in batch.
+
+        Args:
+            records: List of vector records to store
+        """
+        ...
+
+    @abstractmethod
+    def search_similar(
+        self,
+        query_embedding: list[float],
+        top_k: int = 5,
+        filter_metadata: PropertyDict | None = None,
+    ) -> list[SearchResult]:
+        """Search for similar embeddings.
+
+        Args:
+            query_embedding: Query embedding vector
+            top_k: Number of results to return
+            filter_metadata: Optional metadata filter
+
+        Returns:
+            List of search results
+        """
+        ...
+
+    @abstractmethod
+    def delete_by_node_id(self, node_id: int) -> bool:
+        """Delete an embedding by node ID.
+
+        Args:
+            node_id: Node identifier to delete
+
+        Returns:
+            True if deleted, False if not found
+        """
+        ...
+
+    @abstractmethod
+    def clear(self) -> None:
+        """Clear all embeddings from the store."""
+        ...
+
+    @abstractmethod
+    def get_stats(self) -> dict[str, int]:
+        """Get store statistics.
+
+        Returns:
+            Dictionary with statistics (count, dimension, etc.)
+        """
+        ...
+
+
+class MemoryVectorStore(VectorStore):
+    """In-memory vector store implementation.
+
+    Uses cosine similarity for search. Suitable for testing and
+    small datasets.
+
+    Args:
+        dimension: Expected embedding dimension
+    """
+
+    def __init__(self, dimension: int = 1024):
+        self.dimension = dimension
+        self._records: dict[int, VectorRecord] = {}
+
+    def __len__(self) -> int:
+        """Return the number of stored embeddings."""
+        return len(self._records)
+
+    def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
+        """Calculate cosine similarity between two vectors.
+
+        Args:
+            a: First vector
+            b: Second vector
+
+        Returns:
+            Cosine similarity (-1 to 1)
+        """
+        try:
+            return cosine_similarity(a, b)
+        except ValueError:
+            return 0.0
+
+    def store_embedding(
+        self,
+        node_id: int,
+        qualified_name: str,
+        embedding: list[float],
+        metadata: PropertyDict | None = None,
+        **kwargs,
+    ) -> None:
+        """Store an embedding vector in memory.
+
+        Args:
+            node_id: Unique node identifier
+            qualified_name: Fully qualified name of the code entity
+            embedding: Embedding vector
+            metadata: Additional metadata dictionary
+            **kwargs: Additional keyword arguments (stored as metadata)
+
+        Raises:
+            ValueError: If embedding dimension doesn't match or embedding is empty
+        """
+        if not embedding:
+            raise ValueError("Embedding cannot be empty")
+
+        if len(embedding) != self.dimension:
+            raise ValueError(
+                f"Embedding dimension mismatch: expected {self.dimension}, got {len(embedding)}"
+            )
+
+        meta: dict[str, str | int | float | None] = {}
+        if metadata:
+            for k, v in metadata.items():
+                if isinstance(v, (str, int, float, type(None))):
+                    meta[k] = v
+                elif isinstance(v, list):
+                    meta[k] = str(v)
+                elif isinstance(v, bool):
+                    meta[k] = int(v)
+
+        # Store additional kwargs as metadata
+        for k, v in kwargs.items():
+            if isinstance(v, (str, int, float, type(None))):
+                meta[k] = v
+            elif isinstance(v, list):
+                meta[k] = str(v)
+            elif isinstance(v, bool):
+                meta[k] = int(v)
+
+        self._records[node_id] = VectorRecord(
+            node_id=node_id,
+            qualified_name=qualified_name,
+            embedding=embedding,
+            metadata=meta,
+        )
+
+    def store_embeddings_batch(
+        self,
+        records: list[VectorRecord],
+    ) -> None:
+        """Store multiple embeddings in batch."""
+        for record in records:
+            self._records[record.node_id] = record
+
+    def search_similar(
+        self,
+        query_embedding: list[float],
+        top_k: int = 5,
+        filter_metadata: PropertyDict | None = None,
+    ) -> list[SearchResult]:
+        """Search for similar embeddings using cosine similarity."""
+        if top_k < 0:
+            raise ValueError(f"top_k must be non-negative, got {top_k}")
+
+        if not self._records:
+            return []
+
+        scores: list[tuple[int, float, str]] = []
+
+        for node_id, record in self._records.items():
+            if filter_metadata:
+                match = all(
+                    record.metadata.get(k) == v for k, v in filter_metadata.items()
+                )
+                if not match:
+                    continue
+
+            similarity = self._cosine_similarity(query_embedding, record.embedding)
+            scores.append((node_id, similarity, record.qualified_name))
+
+        scores.sort(key=lambda x: x[1], reverse=True)
+
+        return [
+            SearchResult(
+                node_id=node_id,
+                score=round(score, 4),
+                qualified_name=qn,
+            )
+            for node_id, score, qn in scores[:top_k]
+        ]
+
+    def delete_by_node_id(self, node_id: int) -> bool:
+        """Delete an embedding by node ID."""
+        if node_id in self._records:
+            del self._records[node_id]
+            return True
+        return False
+
+    # Alias for compatibility with tests
+    delete_embedding = delete_by_node_id
+
+    def get_embedding(self, node_id: int) -> "VectorRecord | None":
+        """Get an embedding record by node ID.
+
+        Args:
+            node_id: Node identifier
+
+        Returns:
+            VectorRecord if found, None otherwise
+        """
+        return self._records.get(node_id)
+
+    def clear(self) -> None:
+        """Clear all embeddings."""
+        self._records.clear()
+
+    def get_stats(self) -> dict[str, int]:
+        """Get store statistics."""
+        return {
+            "count": len(self._records),
+            "dimension": self.dimension,
+        }
+
+    def get_all_records(self) -> list[VectorRecord]:
+        """Get all records (for testing/debugging).
+
+        Returns:
+            List of all vector records
+        """
+        return list(self._records.values())
+
+
+class QdrantVectorStore(VectorStore):
+    """Qdrant-based vector store implementation.
+
+    Requires qdrant-client to be installed.
+
+    Args:
+        collection_name: Name of the Qdrant collection
+        dimension: Embedding dimension
+        db_path: Path for local Qdrant storage (optional)
+        host: Qdrant server host (if not using local)
+        port: Qdrant server port
+    """
+
+    def __init__(
+        self,
+        collection_name: str = "code_embeddings",
+        dimension: int = 1024,
+        db_path: str | Path | None = None,
+        host: str | None = None,
+        port: int = 6333,
+    ):
+        self.collection_name = collection_name
+        self.dimension = dimension
+        self.db_path = Path(db_path) if db_path else None
+        self.host = host
+        self.port = port
+
+        self._client: "QdrantClient | None" = None
+        self._initialized = False
+
+    def _lazy_init(self) -> None:
+        """Lazy initialization of Qdrant client."""
+        if self._initialized:
+            return
+
+        try:
+            from qdrant_client import QdrantClient
+            from qdrant_client.models import Distance, VectorParams
+
+            if self.db_path:
+                self._client = QdrantClient(path=str(self.db_path))
+            elif self.host:
+                self._client = QdrantClient(host=self.host, port=self.port)
+            else:
+                self._client = QdrantClient(location=":memory:")
+
+            if not self._client.collection_exists(self.collection_name):
+                self._client.create_collection(
+                    collection_name=self.collection_name,
+                    vectors_config=VectorParams(
+                        size=self.dimension,
+                        distance=Distance.COSINE,
+                    ),
+                )
+                logger.info(f"Created Qdrant collection: {self.collection_name}")
+
+            self._initialized = True
+
+        except ImportError as e:
+            logger.error(f"Failed to import qdrant-client: {e}")
+            raise RuntimeError(
+                "qdrant-client required for QdrantVectorStore. "
+                "Install with: pip install qdrant-client"
+            ) from e
+        except Exception as e:
+            logger.error(f"Failed to initialize Qdrant: {e}")
+            raise
+
+    def store_embedding(
+        self,
+        node_id: int,
+        qualified_name: str,
+        embedding: list[float],
+        metadata: PropertyDict | None = None,
+        **kwargs,
+    ) -> None:
+        """Store an embedding vector in Qdrant."""
+        self._lazy_init()
+
+        from qdrant_client.models import PointStruct
+
+        payload: dict[str, str | int | float | None] = {
+            "node_id": node_id,
+            "qualified_name": qualified_name,
+        }
+        if metadata:
+            for k, v in metadata.items():
+                if isinstance(v, (str, int, float, type(None))):
+                    payload[k] = v
+                elif isinstance(v, list):
+                    payload[k] = str(v)
+                elif isinstance(v, bool):
+                    payload[k] = int(v)
+
+        # Store additional kwargs as metadata
+        for k, v in kwargs.items():
+            if isinstance(v, (str, int, float, type(None))):
+                payload[k] = v
+            elif isinstance(v, list):
+                payload[k] = str(v)
+            elif isinstance(v, bool):
+                payload[k] = int(v)
+
+        assert self._client is not None
+        self._client.upsert(
+            collection_name=self.collection_name,
+            points=[
+                PointStruct(
+                    id=node_id,
+                    vector=embedding,
+                    payload=payload,
+                )
+            ],
+        )
+
+    def store_embeddings_batch(
+        self,
+        records: list[VectorRecord],
+    ) -> None:
+        """Store multiple embeddings in batch."""
+        self._lazy_init()
+
+        from qdrant_client.models import PointStruct
+
+        points = []
+        for record in records:
+            payload: dict[str, str | int | float | None] = {
+                "node_id": record.node_id,
+                "qualified_name": record.qualified_name,
+            }
+            payload.update(record.metadata)
+
+            points.append(
+                PointStruct(
+                    id=record.node_id,
+                    vector=record.embedding,
+                    payload=payload,
+                )
+            )
+
+        if points:
+            assert self._client is not None
+            self._client.upsert(
+                collection_name=self.collection_name,
+                points=points,
+            )
+
+    def search_similar(
+        self,
+        query_embedding: list[float],
+        top_k: int = 5,
+        filter_metadata: PropertyDict | None = None,
+    ) -> list[SearchResult]:
+        """Search for similar embeddings in Qdrant."""
+        self._lazy_init()
+
+        from qdrant_client.models import Filter, FieldCondition, MatchValue
+
+        search_filter = None
+        if filter_metadata:
+            conditions = []
+            for k, v in filter_metadata.items():
+                if isinstance(v, (str, int)):
+                    conditions.append(
+                        FieldCondition(key=k, match=MatchValue(value=v))
+                    )
+            if conditions:
+                search_filter = Filter(must=conditions)
+
+        assert self._client is not None
+        results = self._client.query_points(
+            collection_name=self.collection_name,
+            query=query_embedding,
+            limit=top_k,
+            query_filter=search_filter,
+        )
+
+        return [
+            SearchResult(
+                node_id=hit.payload["node_id"],
+                score=hit.score,
+                qualified_name=str(hit.payload.get("qualified_name", "")),
+            )
+            for hit in results.points
+            if hit.payload is not None
+        ]
+
+    def delete_by_node_id(self, node_id: int) -> bool:
+        """Delete an embedding by node ID."""
+        self._lazy_init()
+
+        assert self._client is not None
+        result = self._client.delete(
+            collection_name=self.collection_name,
+            points_selector=[node_id],
+        )
+
+        return result.operation_id is not None
+
+    def clear(self) -> None:
+        """Clear all embeddings."""
+        self._lazy_init()
+
+        assert self._client is not None
+        self._client.delete_collection(self.collection_name)
+        self._initialized = False
+        self._lazy_init()
+
+    def get_stats(self) -> dict[str, int]:
+        """Get store statistics."""
+        self._lazy_init()
+
+        assert self._client is not None
+        info = self._client.get_collection(self.collection_name)
+
+        return {
+            "count": info.points_count,
+            "dimension": self.dimension,
+        }
+
+
+def create_vector_store(
+    backend: str = "memory",
+    dimension: int = 1024,
+    **kwargs: str | int | Path | None,
+) -> VectorStore:
+    """Factory function to create vector store.
+
+    Args:
+        backend: Backend type ("memory" or "qdrant")
+        dimension: Embedding dimension
+        **kwargs: Additional arguments for specific backends
+
+    Returns:
+        VectorStore instance
+
+    Raises:
+        ValueError: If backend is unknown
+    """
+    if backend == "memory":
+        return MemoryVectorStore(dimension=dimension)
+    elif backend == "qdrant":
+        return QdrantVectorStore(
+            dimension=dimension,
+            collection_name=str(kwargs.get("collection_name", "code_embeddings")),
+            db_path=kwargs.get("db_path"),
+            host=kwargs.get("host"),
+            port=int(kwargs.get("port", 6333)),
+        )
+    else:
+        raise ValueError(f"Unknown vector store backend: {backend}")