kodit 0.1.14__py3-none-any.whl → 0.1.15__py3-none-any.whl

kodit/embedding/embedding_provider/local_embedding_provider.py

@@ -0,0 +1,58 @@
+"""Local embedding service."""
+
+import os
+
+import structlog
+import tiktoken
+from sentence_transformers import SentenceTransformer
+from tqdm import tqdm
+
+from kodit.embedding.embedding_provider.embedding_provider import (
+    EmbeddingProvider,
+    Vector,
+    split_sub_batches,
+)
+
+TINY = "tiny"
+CODE = "code"
+TEST = "test"
+
+COMMON_EMBEDDING_MODELS = {
+    TINY: "ibm-granite/granite-embedding-30m-english",
+    CODE: "flax-sentence-embeddings/st-codesearch-distilroberta-base",
+    TEST: "minishlab/potion-base-4M",
+}
+
+
+class LocalEmbeddingProvider(EmbeddingProvider):
+    """Local embedder."""
+
+    def __init__(self, model_name: str) -> None:
+        """Initialize the local embedder."""
+        self.log = structlog.get_logger(__name__)
+        self.model_name = COMMON_EMBEDDING_MODELS.get(model_name, model_name)
+        self.embedding_model = None
+        self.encoding = tiktoken.encoding_for_model("text-embedding-3-small")
+
+    def _model(self) -> SentenceTransformer:
+        """Get the embedding model."""
+        if self.embedding_model is None:
+            os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
+            self.embedding_model = SentenceTransformer(
+                self.model_name,
+                trust_remote_code=True,
+                device="cpu",  # Force CPU so we don't have to install accelerate, etc.
+            )
+        return self.embedding_model
+
+    async def embed(self, data: list[str]) -> list[Vector]:
+        """Embed a list of strings."""
+        model = self._model()
+
+        batched_data = split_sub_batches(self.encoding, data)
+
+        results: list[Vector] = []
+        for batch in tqdm(batched_data, total=len(batched_data), leave=False):
+            embeddings = model.encode(batch, show_progress_bar=False, batch_size=4)
+            results.extend([[float(x) for x in embedding] for embedding in embeddings])
+        return results

kodit/embedding/embedding_provider/openai_embedding_provider.py

@@ -0,0 +1,63 @@
+"""OpenAI embedding service."""
+
+import asyncio
+
+import structlog
+import tiktoken
+from openai import AsyncOpenAI
+
+from kodit.embedding.embedding_provider.embedding_provider import (
+    EmbeddingProvider,
+    Vector,
+    split_sub_batches,
+)
+
+OPENAI_NUM_PARALLEL_TASKS = 10
+
+
+class OpenAIEmbeddingProvider(EmbeddingProvider):
+    """OpenAI embedder."""
+
+    def __init__(
+        self,
+        openai_client: AsyncOpenAI,
+        model_name: str = "text-embedding-3-small",
+    ) -> None:
+        """Initialize the OpenAI embedder."""
+        self.log = structlog.get_logger(__name__)
+        self.openai_client = openai_client
+        self.model_name = model_name
+        self.encoding = tiktoken.encoding_for_model(model_name)
+
+    async def embed(self, data: list[str]) -> list[Vector]:
+        """Embed a list of documents."""
+        # First split the list into a list of list where each sublist has fewer than
+        # max tokens.
+        batched_data = split_sub_batches(self.encoding, data)
+
+        # Process batches in parallel with a semaphore to limit concurrent requests
+        sem = asyncio.Semaphore(OPENAI_NUM_PARALLEL_TASKS)
+
+        async def process_batch(batch: list[str]) -> list[Vector]:
+            async with sem:
+                try:
+                    response = await self.openai_client.embeddings.create(
+                        model=self.model_name,
+                        input=batch,
+                    )
+                    return [
+                        [float(x) for x in embedding.embedding]
+                        for embedding in response.data
+                    ]
+                except Exception as e:
+                    self.log.exception("Error embedding batch", error=str(e))
+                    return []
+
+        # Create tasks for all batches
+        tasks = [process_batch(batch) for batch in batched_data]
+
+        # Process all batches and yield results as they complete
+        results: list[Vector] = []
+        for task in asyncio.as_completed(tasks):
+            results.extend(await task)
+        return results

kodit/embedding/embedding_repository.py

@@ -0,0 +1,206 @@
+"""Repository for managing embeddings."""
+
+import numpy as np
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.embedding.embedding_models import Embedding, EmbeddingType
+
+
+class EmbeddingRepository:
+    """Repository for managing embeddings.
+
+    This class provides methods for creating and retrieving embeddings from the
+    database. It handles the low-level database operations and transaction management.
+
+    Args:
+        session: The SQLAlchemy async session to use for database operations.
+
+    """
+
+    def __init__(self, session: AsyncSession) -> None:
+        """Initialize the embedding repository."""
+        self.session = session
+
+    async def create_embedding(self, embedding: Embedding) -> Embedding:
+        """Create a new embedding record in the database.
+
+        Args:
+            embedding: The Embedding instance to create.
+
+        Returns:
+            The created Embedding instance.
+
+        """
+        self.session.add(embedding)
+        await self.session.commit()
+        return embedding
+
+    async def get_embedding_by_snippet_id_and_type(
+        self, snippet_id: int, embedding_type: EmbeddingType
+    ) -> Embedding | None:
+        """Get an embedding by its snippet ID and type.
+
+        Args:
+            snippet_id: The ID of the snippet to get the embedding for.
+            embedding_type: The type of embedding to get.
+
+        Returns:
+            The Embedding instance if found, None otherwise.
+
+        """
+        query = select(Embedding).where(
+            Embedding.snippet_id == snippet_id,
+            Embedding.type == embedding_type,
+        )
+        result = await self.session.execute(query)
+        return result.scalar_one_or_none()
+
+    async def list_embeddings_by_type(
+        self, embedding_type: EmbeddingType
+    ) -> list[Embedding]:
+        """List all embeddings of a given type.
+
+        Args:
+            embedding_type: The type of embeddings to list.
+
+        Returns:
+            A list of Embedding instances.
+
+        """
+        query = select(Embedding).where(Embedding.type == embedding_type)
+        result = await self.session.execute(query)
+        return list(result.scalars())
+
+    async def delete_embeddings_by_snippet_id(self, snippet_id: int) -> None:
+        """Delete all embeddings for a snippet.
+
+        Args:
+            snippet_id: The ID of the snippet to delete embeddings for.
+
+        """
+        query = select(Embedding).where(Embedding.snippet_id == snippet_id)
+        result = await self.session.execute(query)
+        embeddings = result.scalars().all()
+        for embedding in embeddings:
+            await self.session.delete(embedding)
+        await self.session.commit()
+
+    async def list_semantic_results(
+        self, embedding_type: EmbeddingType, embedding: list[float], top_k: int = 10
+    ) -> list[tuple[int, float]]:
+        """List semantic results using cosine similarity.
+
+        This implementation fetches all embeddings of the given type and computes
+        cosine similarity in Python using NumPy for better performance.
+
+        Args:
+            embedding_type: The type of embeddings to search
+            embedding: The query embedding vector
+            top_k: Number of results to return
+
+        Returns:
+            List of (snippet_id, similarity_score) tuples, sorted by similarity
+
+        """
+        # Step 1: Fetch embeddings from database
+        embeddings = await self._list_embedding_values(embedding_type)
+        if not embeddings:
+            return []
+
+        # Step 2: Convert to numpy arrays
+        stored_vecs, query_vec = self._prepare_vectors(embeddings, embedding)
+
+        # Step 3: Compute similarities
+        similarities = self._compute_similarities(stored_vecs, query_vec)
+
+        # Step 4: Get top-k results
+        return self._get_top_k_results(similarities, embeddings, top_k)
+
+    async def _list_embedding_values(
+        self, embedding_type: EmbeddingType
+    ) -> list[tuple[int, list[float]]]:
+        """List all embeddings of a given type from the database.
+
+        Args:
+            embedding_type: The type of embeddings to fetch
+
+        Returns:
+            List of (snippet_id, embedding) tuples
+
+        """
+        # Only select the fields we need and use a more efficient query
+        query = select(Embedding.snippet_id, Embedding.embedding).where(
+            Embedding.type == embedding_type
+        )
+        rows = await self.session.execute(query)
+        return [tuple(row) for row in rows.all()]  # Convert Row objects to tuples
+
+    def _prepare_vectors(
+        self, embeddings: list[tuple[int, list[float]]], query_embedding: list[float]
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Convert embeddings to numpy arrays.
+
+        Args:
+            embeddings: List of (snippet_id, embedding) tuples
+            query_embedding: Query embedding vector
+
+        Returns:
+            Tuple of (stored_vectors, query_vector) as numpy arrays
+
+        """
+        try:
+            stored_vecs = np.array(
+                [emb[1] for emb in embeddings]
+            )  # Use index 1 to get embedding
+        except ValueError as e:
+            if "inhomogeneous" in str(e):
+                msg = (
+                    "The database has returned embeddings of different sizes. If you"
+                    "have recently updated the embedding model, you will need to"
+                    "delete your database and re-index your snippets."
+                )
+                raise ValueError(msg) from e
+            raise
+
+        query_vec = np.array(query_embedding)
+        return stored_vecs, query_vec
+
+    def _compute_similarities(
+        self, stored_vecs: np.ndarray, query_vec: np.ndarray
+    ) -> np.ndarray:
+        """Compute cosine similarities between stored vectors and query vector.
+
+        Args:
+            stored_vecs: Array of stored embedding vectors
+            query_vec: Query embedding vector
+
+        Returns:
+            Array of similarity scores
+
+        """
+        stored_norms = np.linalg.norm(stored_vecs, axis=1)
+        query_norm = np.linalg.norm(query_vec)
+        return np.dot(stored_vecs, query_vec) / (stored_norms * query_norm)
+
+    def _get_top_k_results(
+        self,
+        similarities: np.ndarray,
+        embeddings: list[tuple[int, list[float]]],
+        top_k: int,
+    ) -> list[tuple[int, float]]:
+        """Get top-k results by similarity score.
+
+        Args:
+            similarities: Array of similarity scores
+            embeddings: List of (snippet_id, embedding) tuples
+            top_k: Number of results to return
+
+        Returns:
+            List of (snippet_id, similarity_score) tuples
+
+        """
+        top_indices = np.argsort(similarities)[::-1][:top_k]
+        return [
+            (embeddings[i][0], float(similarities[i])) for i in top_indices
+        ]  # Use index 0 to get snippet_id

kodit/embedding/local_vector_search_service.py

@@ -0,0 +1,50 @@
+"""Local vector search."""
+
+import structlog
+import tiktoken
+
+from kodit.embedding.embedding_models import Embedding, EmbeddingType
+from kodit.embedding.embedding_provider.embedding_provider import EmbeddingProvider
+from kodit.embedding.embedding_repository import EmbeddingRepository
+from kodit.embedding.vector_search_service import (
+    VectorSearchRequest,
+    VectorSearchResponse,
+    VectorSearchService,
+)
+
+
+class LocalVectorSearchService(VectorSearchService):
+    """Local vector search."""
+
+    def __init__(
+        self,
+        embedding_repository: EmbeddingRepository,
+        embedding_provider: EmbeddingProvider,
+    ) -> None:
+        """Initialize the local embedder."""
+        self.log = structlog.get_logger(__name__)
+        self.embedding_repository = embedding_repository
+        self.embedding_provider = embedding_provider
+        self.encoding = tiktoken.encoding_for_model("text-embedding-3-small")
+
+    async def index(self, data: list[VectorSearchRequest]) -> None:
+        """Embed a list of documents."""
+        embeddings = await self.embedding_provider.embed([i.text for i in data])
+        for i, x in zip(data, embeddings, strict=False):
+            await self.embedding_repository.create_embedding(
+                Embedding(
+                    snippet_id=i.snippet_id,
+                    embedding=[float(y) for y in x],
+                    type=EmbeddingType.CODE,
+                )
+            )
+
+    async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
+        """Query the embedding model."""
+        embedding = (await self.embedding_provider.embed([query]))[0]
+        results = await self.embedding_repository.list_semantic_results(
+            EmbeddingType.CODE, [float(x) for x in embedding], top_k
+        )
+        return [
+            VectorSearchResponse(snippet_id, score) for snippet_id, score in results
+        ]

kodit/embedding/vector_search_service.py

@@ -0,0 +1,38 @@
+"""Embedding service."""
+
+from abc import ABC, abstractmethod
+from typing import NamedTuple
+
+
+class VectorSearchResponse(NamedTuple):
+    """Embedding result."""
+
+    snippet_id: int
+    score: float
+
+
+class VectorSearchRequest(NamedTuple):
+    """Input for embedding."""
+
+    snippet_id: int
+    text: str
+
+
+class VectorSearchService(ABC):
+    """Semantic search service interface."""
+
+    @abstractmethod
+    async def index(self, data: list[VectorSearchRequest]) -> None:
+        """Embed a list of documents.
+
+        The embedding service accepts a massive list of id,strings to embed. Behind the
+        scenes it batches up requests and parallelizes them for performance according to
+        the specifics of the embedding service.
+
+        The id reference is required because the parallelization may return results out
+        of order.
+        """
+
+    @abstractmethod
+    async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
+        """Query the embedding model."""

kodit/embedding/vectorchord_vector_search_service.py

@@ -0,0 +1,145 @@
+"""Vectorchord vector search."""
+
+from typing import Any
+
+from sqlalchemy import Result, TextClause, text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.embedding.embedding_provider.embedding_provider import EmbeddingProvider
+from kodit.embedding.vector_search_service import (
+    VectorSearchRequest,
+    VectorSearchResponse,
+    VectorSearchService,
+)
+
+TABLE_NAME = "vectorchord_embeddings"
+INDEX_NAME = f"{TABLE_NAME}_idx"
+
+# SQL Queries
+CREATE_VCHORD_EXTENSION = """
+CREATE EXTENSION IF NOT EXISTS vchord CASCADE;
+"""
+
+CHECK_VCHORD_EMBEDDING_DIMENSION = f"""
+SELECT a.atttypmod as dimension
+FROM pg_attribute a
+JOIN pg_class c ON a.attrelid = c.oid
+WHERE c.relname = '{TABLE_NAME}'
+AND a.attname = 'embedding';
+"""  # noqa: S608
+
+CREATE_VCHORD_INDEX = f"""
+CREATE INDEX IF NOT EXISTS {INDEX_NAME}
+ON {TABLE_NAME}
+USING vchordrq (embedding vector_l2_ops) WITH (options = $$
+residual_quantization = true
+[build.internal]
+lists = []
+$$);
+"""
+
+INSERT_QUERY = f"""
+INSERT INTO {TABLE_NAME} (snippet_id, embedding)
+VALUES (:snippet_id, :embedding)
+ON CONFLICT (snippet_id) DO UPDATE
+SET embedding = EXCLUDED.embedding
+"""  # noqa: S608
+
+# Note that <=> in vectorchord is cosine distance
+# So scores go from 0 (similar) to 2 (opposite)
+SEARCH_QUERY = f"""
+SELECT snippet_id, embedding <=> :query as score
+FROM {TABLE_NAME}
+ORDER BY score ASC
+LIMIT :top_k;
+"""  # noqa: S608
+
+
+class VectorChordVectorSearchService(VectorSearchService):
+    """VectorChord vector search."""
+
+    def __init__(
+        self,
+        session: AsyncSession,
+        embedding_provider: EmbeddingProvider,
+    ) -> None:
+        """Initialize the VectorChord BM25."""
+        self.embedding_provider = embedding_provider
+        self._session = session
+        self._initialized = False
+
+    async def _initialize(self) -> None:
+        """Initialize the VectorChord environment."""
+        try:
+            await self._create_extensions()
+            await self._create_tables()
+            self._initialized = True
+        except Exception as e:
+            msg = f"Failed to initialize VectorChord repository: {e}"
+            raise RuntimeError(msg) from e
+
+    async def _create_extensions(self) -> None:
+        """Create the necessary extensions."""
+        await self._session.execute(text(CREATE_VCHORD_EXTENSION))
+        await self._commit()
+
+    async def _create_tables(self) -> None:
+        """Create the necessary tables."""
+        vector_dim = (await self.embedding_provider.embed(["dimension"]))[0]
+        await self._session.execute(
+            text(
+                f"""CREATE TABLE IF NOT EXISTS {TABLE_NAME} (
+                    id SERIAL PRIMARY KEY,
+                    snippet_id INT NOT NULL UNIQUE,
+                    embedding VECTOR({len(vector_dim)}) NOT NULL
+                );"""
+            )
+        )
+        await self._session.execute(text(CREATE_VCHORD_INDEX))
+        result = await self._session.execute(text(CHECK_VCHORD_EMBEDDING_DIMENSION))
+        vector_dim_from_db = result.scalar_one()
+        if vector_dim_from_db != len(vector_dim):
+            msg = (
+                f"Embedding vector dimension does not match database, "
+                f"please delete your index: {vector_dim_from_db} != {len(vector_dim)}"
+            )
+            raise ValueError(msg)
+        await self._commit()
+
+    async def _execute(
+        self, query: TextClause, param_list: list[Any] | dict[str, Any] | None = None
+    ) -> Result:
+        """Execute a query."""
+        if not self._initialized:
+            await self._initialize()
+        return await self._session.execute(query, param_list)
+
+    async def _commit(self) -> None:
+        """Commit the session."""
+        await self._session.commit()
+
+    async def index(self, data: list[VectorSearchRequest]) -> None:
+        """Embed a list of documents."""
+        embeddings = await self.embedding_provider.embed([doc.text for doc in data])
+        # Execute inserts
+        await self._execute(
+            text(INSERT_QUERY),
+            [
+                {"snippet_id": doc.snippet_id, "embedding": str(embedding)}
+                for doc, embedding in zip(data, embeddings, strict=True)
+            ],
+        )
+        await self._commit()
+
+    async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
+        """Query the embedding model."""
+        embedding = await self.embedding_provider.embed([query])
+        result = await self._execute(
+            text(SEARCH_QUERY), {"query": str(embedding[0]), "top_k": top_k}
+        )
+        rows = result.mappings().all()
+
+        return [
+            VectorSearchResponse(snippet_id=row["snippet_id"], score=row["score"])
+            for row in rows
+        ]

kodit/indexing/indexing_repository.py

@@ -10,6 +10,7 @@ from typing import TypeVar
 
 from sqlalchemy import delete, func, select
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm.exc import MultipleResultsFound
 
 from kodit.embedding.embedding_models import Embedding
 from kodit.indexing.indexing_models import Index, Snippet
@@ -124,15 +125,34 @@ class IndexRepository:
         index.updated_at = datetime.now(UTC)
         await self.session.commit()
 
-    async def add_snippet(self, snippet: Snippet) -> None:
-        """Add a new snippet to the database.
+    async def add_snippet_or_update_content(self, snippet: Snippet) -> None:
+        """Add a new snippet to the database if it doesn't exist, otherwise update it.
 
         Args:
             snippet: The Snippet instance to add.
 
         """
-        self.session.add(snippet)
-        await self.session.commit()
+        query = select(Snippet).where(
+            Snippet.file_id == snippet.file_id,
+            Snippet.index_id == snippet.index_id,
+        )
+        result = await self.session.execute(query)
+        try:
+            existing_snippet = result.scalar_one_or_none()
+
+            if existing_snippet:
+                existing_snippet.content = snippet.content
+            else:
+                self.session.add(snippet)
+
+            await self.session.commit()
+        except MultipleResultsFound as e:
+            msg = (
+                f"Multiple snippets found for file_id {snippet.file_id}, this "
+                "shouldn't happen. "
+                "Please report this as a bug then delete your index and start again."
+            )
+            raise ValueError(msg) from e
 
     async def delete_all_snippets(self, index_id: int) -> None:
         """Delete all snippets for an index.