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

kodit/embedding/embedding_provider/embedding_provider.py

@@ -0,0 +1,60 @@
+"""Embedding provider."""
+
+from abc import ABC, abstractmethod
+
+import structlog
+import tiktoken
+
+OPENAI_MAX_EMBEDDING_SIZE = 8192
+
+Vector = list[float]
+
+
+class EmbeddingProvider(ABC):
+    """Embedding provider."""
+
+    @abstractmethod
+    async def embed(self, data: list[str]) -> list[Vector]:
+        """Embed a list of strings.
+
+        The embedding provider is responsible for embedding a list of strings into a
+        list of vectors. The embedding provider is responsible for splitting the list of
+        strings into smaller sub-batches and embedding them in parallel.
+        """
+
+
+def split_sub_batches(encoding: tiktoken.Encoding, data: list[str]) -> list[list[str]]:
+    """Split a list of strings into smaller sub-batches."""
+    log = structlog.get_logger(__name__)
+    result = []
+    data_to_process = [s for s in data if s.strip()]  # Filter out empty strings
+
+    while data_to_process:
+        next_batch = []
+        current_tokens = 0
+
+        while data_to_process:
+            next_item = data_to_process[0]
+            item_tokens = len(encoding.encode(next_item))
+
+            if item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
+                # Loop around trying to truncate the snippet until it fits in the max
+                # embedding size
+                while item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
+                    next_item = next_item[:-1]
+                    item_tokens = len(encoding.encode(next_item))
+
+                data_to_process[0] = next_item
+
+                log.warning("Truncated snippet", snippet=next_item)
+
+            if current_tokens + item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
+                break
+
+            next_batch.append(data_to_process.pop(0))
+            current_tokens += item_tokens
+
+        if next_batch:
+            result.append(next_batch)
+
+    return result

kodit/embedding/embedding_provider/hash_embedding_provider.py

@@ -0,0 +1,77 @@
+"""Hash embedding provider, for use in tests only."""
+
+import asyncio
+import hashlib
+import math
+from collections.abc import Generator, Sequence
+
+from kodit.embedding.embedding_provider.embedding_provider import (
+    EmbeddingProvider,
+    Vector,
+)
+
+
+class HashEmbeddingProvider(EmbeddingProvider):
+    """A minimal test-time embedding provider.
+
+    • Zero third-party dependencies (uses only std-lib)
+    • Distinguishes strings by hashing with SHA-256
+    • Maps the digest to a fixed-size float vector, then ℓ₂-normalises
+    • Splits work into small asynchronous chunks for speed in event loops
+    """
+
+    def __init__(self, dim: int = 16, batch_size: int = 64) -> None:
+        """Initialize the hash embedding provider."""
+        if dim <= 0:
+            msg = f"dim must be > 0, got {dim}"
+            raise ValueError(msg)
+        if batch_size <= 0:
+            msg = f"batch_size must be > 0, got {batch_size}"
+            raise ValueError(msg)
+        self.dim = dim
+        self.batch_size = batch_size
+
+    async def embed(self, data: list[str]) -> list[Vector]:
+        """Embed every string in *data*, preserving order.
+
+        Work is sliced into *batch_size* chunks and scheduled concurrently
+        (still CPU-bound, but enough to cooperate with an asyncio loop).
+        """
+        if not data:
+            return []
+
+        async def _embed_chunk(chunk: Sequence[str]) -> list[Vector]:
+            return [self._string_to_vector(text) for text in chunk]
+
+        tasks = [
+            asyncio.create_task(_embed_chunk(chunk))
+            for chunk in self._chunked(data, self.batch_size)
+        ]
+
+        vectors: list[Vector] = []
+        for task in tasks:
+            vectors.extend(await task)
+        return vectors
+
+    @staticmethod
+    def _chunked(seq: Sequence[str], size: int) -> Generator[Sequence[str], None, None]:
+        """Yield successive *size*-sized slices from *seq*."""
+        for i in range(0, len(seq), size):
+            yield seq[i : i + size]
+
+    def _string_to_vector(self, text: str) -> Vector:
+        """Deterministically convert *text* to a normalised float vector."""
+        digest = hashlib.sha256(text.encode("utf-8")).digest()
+
+        # Build the vector from 4-byte windows of the digest.
+        vec = [
+            int.from_bytes(
+                digest[(i * 4) % len(digest) : (i * 4) % len(digest) + 4], "big"
+            )
+            / 0xFFFFFFFF
+            for i in range(self.dim)
+        ]
+
+        # ℓ₂-normalise so magnitudes are comparable.
+        norm = math.sqrt(sum(x * x for x in vec)) or 1.0
+        return [x / norm for x in vec]

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,75 @@
+"""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)
+
+        # Create a list of tuples with a temporary id for each batch
+        # We need to do this so that we can return the results in the same order as the
+        # input data
+        input_data = [(i, batch) for i, batch in enumerate(batched_data)]
+
+        async def process_batch(
+            data: tuple[int, list[str]],
+        ) -> tuple[int, list[Vector]]:
+            batch_id, batch = data
+            async with sem:
+                try:
+                    response = await self.openai_client.embeddings.create(
+                        model=self.model_name,
+                        input=batch,
+                    )
+                    return batch_id, [
+                        [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 batch_id, []
+
+        # Create tasks for all batches
+        tasks = [process_batch(batch) for batch in input_data]
+
+        # Process all batches and yield results as they complete
+        results: list[tuple[int, list[Vector]]] = []
+        for task in asyncio.as_completed(tasks):
+            result = await task
+            results.append(result)
+
+        # Output in the same order as the input data
+        ordered_results = [result for _, result in sorted(results, key=lambda x: x[0])]
+        return [item for sublist in ordered_results for item in sublist]

kodit/{search/search_repository.py → embedding/embedding_repository.py}

@@ -1,62 +1,90 @@
-"""Repository for searching for relevant snippets."""
-
-from typing import TypeVar
+"""Repository for managing embeddings."""
 
 import numpy as np
-from sqlalchemy import (
-    select,
-)
+from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.embedding.embedding_models import Embedding, EmbeddingType
-from kodit.indexing.indexing_models import Snippet
-from kodit.source.source_models import File
 
-T = TypeVar("T")
 
+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.
 
-class SearchRepository:
-    """Repository for searching for relevant snippets."""
+    """
 
     def __init__(self, session: AsyncSession) -> None:
-        """Initialize the search repository.
+        """Initialize the embedding repository."""
+        self.session = session
+
+    async def create_embedding(self, embedding: Embedding) -> Embedding:
+        """Create a new embedding record in the database.
 
         Args:
-            session: The SQLAlchemy async session to use for database operations.
+            embedding: The Embedding instance to create.
+
+        Returns:
+            The created Embedding instance.
 
         """
-        self.session = session
+        self.session.add(embedding)
+        await self.session.commit()
+        return embedding
 
-    async def list_snippet_ids(self) -> list[int]:
-        """List all snippet IDs.
+    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:
-            A list of all snippets.
+            The Embedding instance if found, None otherwise.
 
         """
-        query = select(Snippet.id)
-        rows = await self.session.execute(query)
-        return list(rows.scalars().all())
+        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_snippets_by_ids(self, ids: list[int]) -> list[tuple[File, Snippet]]:
-        """List snippets by IDs.
+    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 snippets in the same order as the input IDs.
+            A list of Embedding instances.
 
         """
-        query = (
-            select(Snippet, File)
-            .where(Snippet.id.in_(ids))
-            .join(File, Snippet.file_id == File.id)
-        )
-        rows = await self.session.execute(query)
+        query = select(Embedding).where(Embedding.type == embedding_type)
+        result = await self.session.execute(query)
+        return list(result.scalars())
 
-        # Create a dictionary for O(1) lookup of results by ID
-        id_to_result = {snippet.id: (file, snippet) for snippet, file in rows.all()}
+    async def delete_embeddings_by_snippet_id(self, snippet_id: int) -> None:
+        """Delete all embeddings for a snippet.
 
-        # Return results in the same order as input IDs
-        return [id_to_result[i] for i in ids]
+        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

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,154 @@
+"""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,
+)
+
+# SQL Queries
+CREATE_VCHORD_EXTENSION = """
+CREATE EXTENSION IF NOT EXISTS vchord CASCADE;
+"""
+
+CHECK_VCHORD_EMBEDDING_DIMENSION = """
+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';
+"""
+
+CREATE_VCHORD_INDEX = """
+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 = """
+INSERT INTO {TABLE_NAME} (snippet_id, embedding)
+VALUES (:snippet_id, :embedding)
+ON CONFLICT (snippet_id) DO UPDATE
+SET embedding = EXCLUDED.embedding
+"""
+
+# Note that <=> in vectorchord is cosine distance
+# So scores go from 0 (similar) to 2 (opposite)
+SEARCH_QUERY = """
+SELECT snippet_id, embedding <=> :query as score
+FROM {TABLE_NAME}
+ORDER BY score ASC
+LIMIT :top_k;
+"""
+
+
+class VectorChordVectorSearchService(VectorSearchService):
+    """VectorChord vector search."""
+
+    def __init__(
+        self,
+        task_name: str,
+        session: AsyncSession,
+        embedding_provider: EmbeddingProvider,
+    ) -> None:
+        """Initialize the VectorChord BM25."""
+        self.embedding_provider = embedding_provider
+        self._session = session
+        self._initialized = False
+        self.table_name = f"vectorchord_{task_name}_embeddings"
+        self.index_name = f"{self.table_name}_idx"
+
+    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 {self.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.format(
+                    TABLE_NAME=self.table_name, INDEX_NAME=self.index_name
+                )
+            )
+        )
+        result = await self._session.execute(
+            text(CHECK_VCHORD_EMBEDDING_DIMENSION.format(TABLE_NAME=self.table_name))
+        )
+        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.format(TABLE_NAME=self.table_name)),
+            [
+                {"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.format(TABLE_NAME=self.table_name)),
+            {"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/enrichment/__init__.py

@@ -0,0 +1 @@
+"""Enrichment."""

kodit/enrichment/enrichment_factory.py

@@ -0,0 +1,23 @@
+"""Embedding service."""
+
+from kodit.config import AppContext
+from kodit.enrichment.enrichment_provider.local_enrichment_provider import (
+    LocalEnrichmentProvider,
+)
+from kodit.enrichment.enrichment_provider.openai_enrichment_provider import (
+    OpenAIEnrichmentProvider,
+)
+from kodit.enrichment.enrichment_service import (
+    EnrichmentService,
+    LLMEnrichmentService,
+)
+
+
+def enrichment_factory(app_context: AppContext) -> EnrichmentService:
+    """Create an embedding service."""
+    openai_client = app_context.get_default_openai_client()
+    if openai_client is not None:
+        enrichment_provider = OpenAIEnrichmentProvider(openai_client=openai_client)
+        return LLMEnrichmentService(enrichment_provider)
+
+    return LLMEnrichmentService(LocalEnrichmentProvider())

kodit/enrichment/enrichment_provider/__init__.py

@@ -0,0 +1 @@
+"""Enrichment provider."""