kodit 0.2.3__tar.gz → 0.2.4__tar.gz

{kodit-0.2.3 → kodit-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.2.3
+Version: 0.2.4
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/

{kodit-0.2.3 → kodit-0.2.4}/src/kodit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.3'
-__version_tuple__ = version_tuple = (0, 2, 3)
+__version__ = version = '0.2.4'
+__version_tuple__ = version_tuple = (0, 2, 4)

{kodit-0.2.3 → kodit-0.2.4}/src/kodit/embedding/embedding_factory.py

@@ -3,6 +3,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.config import AppContext, Endpoint
+from kodit.embedding.embedding_models import EmbeddingType
 from kodit.embedding.embedding_provider.local_embedding_provider import (
     CODE,
     LocalEmbeddingProvider,
@@ -54,9 +55,14 @@ def embedding_factory(
         return VectorChordVectorSearchService(task_name, session, embedding_provider)
     if app_context.default_search.provider == "sqlite":
         log_event("kodit.database", {"provider": "sqlite"})
+        if task_name == "code":
+            embedding_type = EmbeddingType.CODE
+        elif task_name == "text":
+            embedding_type = EmbeddingType.TEXT
         return LocalVectorSearchService(
             embedding_repository=embedding_repository,
             embedding_provider=embedding_provider,
+            embedding_type=embedding_type,
         )
 
     msg = f"Invalid semantic search provider: {app_context.default_search.provider}"

kodit-0.2.4/src/kodit/embedding/embedding_provider/embedding_provider.py

@@ -0,0 +1,92 @@
+"""Embedding provider."""
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
+
+import structlog
+import tiktoken
+
+OPENAI_MAX_EMBEDDING_SIZE = 8192
+
+Vector = list[float]
+
+
+@dataclass
+class EmbeddingRequest:
+    """Embedding request."""
+
+    id: int
+    text: str
+
+
+@dataclass
+class EmbeddingResponse:
+    """Embedding response."""
+
+    id: int
+    embedding: Vector
+
+
+class EmbeddingProvider(ABC):
+    """Embedding provider."""
+
+    @abstractmethod
+    def embed(
+        self, data: list[EmbeddingRequest]
+    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
+        """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[EmbeddingRequest],
+    max_context_window: int = OPENAI_MAX_EMBEDDING_SIZE,
+) -> list[list[EmbeddingRequest]]:
+    """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.text.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.text, disallowed_special=()))
+
+            if item_tokens > max_context_window:
+                # Optimise truncation by operating on tokens directly instead of
+                # removing one character at a time and repeatedly re-encoding.
+                tokens = encoding.encode(next_item.text, disallowed_special=())
+                if len(tokens) > max_context_window:
+                    # Keep only the first *max_context_window* tokens.
+                    tokens = tokens[:max_context_window]
+                    # Convert back to text. This requires only one decode call and
+                    # guarantees that the resulting string fits the token budget.
+                    next_item.text = encoding.decode(tokens)
+                    item_tokens = max_context_window  # We know the exact size now
+
+                    data_to_process[0] = next_item
+
+                    log.warning(
+                        "Truncated snippet because it was too long to embed",
+                        snippet=next_item.text[:100] + "...",
+                    )
+
+            if current_tokens + item_tokens > max_context_window:
+                break
+
+            next_batch.append(data_to_process.pop(0))
+            current_tokens += item_tokens
+
+        if next_batch:
+            result.append(next_batch)
+
+    return result

{kodit-0.2.3 → kodit-0.2.4}/src/kodit/embedding/embedding_provider/hash_embedding_provider.py

@@ -3,10 +3,12 @@
 import asyncio
 import hashlib
 import math
-from collections.abc import Generator, Sequence
+from collections.abc import AsyncGenerator, Generator, Sequence
 
 from kodit.embedding.embedding_provider.embedding_provider import (
     EmbeddingProvider,
+    EmbeddingRequest,
+    EmbeddingResponse,
     Vector,
 )
 
@@ -31,27 +33,34 @@ class HashEmbeddingProvider(EmbeddingProvider):
         self.dim = dim
         self.batch_size = batch_size
 
-    async def embed(self, data: list[str]) -> list[Vector]:
+    async def embed(
+        self, data: list[EmbeddingRequest]
+    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
         """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 []
+            yield []
 
         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)
+            for chunk in self._chunked([i.text for i in data], self.batch_size)
         ]
 
-        vectors: list[Vector] = []
         for task in tasks:
-            vectors.extend(await task)
-        return vectors
+            result = await task
+            yield [
+                EmbeddingResponse(
+                    id=item.id,
+                    embedding=embedding,
+                )
+                for item, embedding in zip(data, result, strict=True)
+            ]
 
     @staticmethod
     def _chunked(seq: Sequence[str], size: int) -> Generator[Sequence[str], None, None]:

kodit-0.2.4/src/kodit/embedding/embedding_provider/local_embedding_provider.py

@@ -0,0 +1,96 @@
+"""Local embedding service."""
+
+from __future__ import annotations
+
+import os
+from time import time
+from typing import TYPE_CHECKING
+
+import structlog
+
+from kodit.embedding.embedding_provider.embedding_provider import (
+    EmbeddingProvider,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    split_sub_batches,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+    from sentence_transformers import SentenceTransformer
+    from tiktoken import Encoding
+
+
+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.encoding_name = "text-embedding-3-small"
+        self.embedding_model = None
+        self.encoding = None
+
+    def _encoding(self) -> Encoding:
+        if self.encoding is None:
+            from tiktoken import encoding_for_model
+
+            start_time = time()
+            self.encoding = encoding_for_model(self.encoding_name)
+            self.log.debug(
+                "Encoding loaded",
+                model_name=self.encoding_name,
+                duration=time() - start_time,
+            )
+        return self.encoding
+
+    def _model(self) -> SentenceTransformer:
+        """Get the embedding model."""
+        if self.embedding_model is None:
+            os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
+            from sentence_transformers import SentenceTransformer
+
+            start_time = time()
+            self.embedding_model = SentenceTransformer(
+                self.model_name,
+                trust_remote_code=True,
+            )
+            self.log.debug(
+                "Model loaded",
+                model_name=self.model_name,
+                duration=time() - start_time,
+            )
+        return self.embedding_model
+
+    async def embed(
+        self, data: list[EmbeddingRequest]
+    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
+        """Embed a list of strings."""
+        model = self._model()
+
+        batched_data = split_sub_batches(self._encoding(), data)
+
+        for batch in batched_data:
+            embeddings = model.encode(
+                [i.text for i in batch], show_progress_bar=False, batch_size=4
+            )
+            yield [
+                EmbeddingResponse(
+                    id=item.id,
+                    embedding=[float(x) for x in embedding],
+                )
+                for item, embedding in zip(batch, embeddings, strict=True)
+            ]

{kodit-0.2.3 → kodit-0.2.4}/src/kodit/embedding/embedding_provider/openai_embedding_provider.py

@@ -1,6 +1,7 @@
 """OpenAI embedding service."""
 
 import asyncio
+from collections.abc import AsyncGenerator
 
 import structlog
 import tiktoken
@@ -8,7 +9,8 @@ from openai import AsyncOpenAI
 
 from kodit.embedding.embedding_provider.embedding_provider import (
     EmbeddingProvider,
-    Vector,
+    EmbeddingRequest,
+    EmbeddingResponse,
     split_sub_batches,
 )
 
@@ -31,7 +33,9 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
             "text-embedding-3-small"
         )  # Sensible default
 
-    async def embed(self, data: list[str]) -> list[Vector]:
+    async def embed(
+        self, data: list[EmbeddingRequest]
+    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
         """Embed a list of documents."""
         # First split the list into a list of list where each sublist has fewer than
         # max tokens.
@@ -40,38 +44,30 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
         # 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
+            data: list[EmbeddingRequest],
+        ) -> list[EmbeddingResponse]:
             async with sem:
                 try:
                     response = await self.openai_client.embeddings.create(
                         model=self.model_name,
-                        input=batch,
+                        input=[i.text for i in data],
                     )
-                    return batch_id, [
-                        [float(x) for x in embedding.embedding]
-                        for embedding in response.data
+                    return [
+                        EmbeddingResponse(
+                            id=item.id,
+                            embedding=embedding.embedding,
+                        )
+                        for item, embedding in zip(data, response.data, strict=True)
                     ]
                 except Exception as e:
                     self.log.exception("Error embedding batch", error=str(e))
-                    return batch_id, []
+                    return []
 
         # Create tasks for all batches
-        tasks = [process_batch(batch) for batch in input_data]
+        tasks = [process_batch(batch) for batch in batched_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]
+            yield result

kodit-0.2.4/src/kodit/embedding/local_vector_search_service.py

@@ -0,0 +1,87 @@
+"""Local vector search."""
+
+from collections.abc import AsyncGenerator
+
+import structlog
+import tiktoken
+
+from kodit.embedding.embedding_models import Embedding, EmbeddingType
+from kodit.embedding.embedding_provider.embedding_provider import (
+    EmbeddingProvider,
+    EmbeddingRequest,
+)
+from kodit.embedding.embedding_repository import EmbeddingRepository
+from kodit.embedding.vector_search_service import (
+    IndexResult,
+    VectorSearchRequest,
+    VectorSearchResponse,
+    VectorSearchService,
+)
+
+
+class LocalVectorSearchService(VectorSearchService):
+    """Local vector search."""
+
+    def __init__(
+        self,
+        embedding_repository: EmbeddingRepository,
+        embedding_provider: EmbeddingProvider,
+        embedding_type: EmbeddingType = EmbeddingType.CODE,
+    ) -> 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")
+        self.embedding_type = embedding_type
+
+    async def index(
+        self, data: list[VectorSearchRequest]
+    ) -> AsyncGenerator[list[IndexResult], None]:
+        """Embed a list of documents."""
+        if not data or len(data) == 0:
+            return
+
+        requests = [EmbeddingRequest(id=doc.snippet_id, text=doc.text) for doc in data]
+
+        async for batch in self.embedding_provider.embed(requests):
+            for result in batch:
+                await self.embedding_repository.create_embedding(
+                    Embedding(
+                        snippet_id=result.id,
+                        embedding=result.embedding,
+                        type=self.embedding_type,
+                    )
+                )
+                yield [IndexResult(snippet_id=result.id)]
+
+    async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
+        """Query the embedding model."""
+        # Build a single-item request and collect its embedding.
+        req = EmbeddingRequest(id=0, text=query)
+        embedding_vec: list[float] | None = None
+        async for batch in self.embedding_provider.embed([req]):
+            if batch:
+                embedding_vec = [float(v) for v in batch[0].embedding]
+                break
+
+        if not embedding_vec:
+            return []
+
+        results = await self.embedding_repository.list_semantic_results(
+            self.embedding_type, embedding_vec, top_k
+        )
+        return [
+            VectorSearchResponse(snippet_id, score) for snippet_id, score in results
+        ]
+
+    async def has_embedding(
+        self, snippet_id: int, embedding_type: EmbeddingType
+    ) -> bool:
+        """Check if a snippet has an embedding."""
+        return (
+            await self.embedding_repository.get_embedding_by_snippet_id_and_type(
+                snippet_id, embedding_type
+            )
+            is not None
+        )

{kodit-0.2.3 → kodit-0.2.4}/src/kodit/embedding/vector_search_service.py

@@ -1,8 +1,11 @@
 """Embedding service."""
 
 from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
 from typing import NamedTuple
 
+from kodit.embedding.embedding_models import EmbeddingType
+
 
 class VectorSearchResponse(NamedTuple):
     """Embedding result."""
@@ -18,11 +21,19 @@ class VectorSearchRequest(NamedTuple):
     text: str
 
 
+class IndexResult(NamedTuple):
+    """Result of indexing."""
+
+    snippet_id: int
+
+
 class VectorSearchService(ABC):
     """Semantic search service interface."""
 
     @abstractmethod
-    async def index(self, data: list[VectorSearchRequest]) -> None:
+    def index(
+        self, data: list[VectorSearchRequest]
+    ) -> AsyncGenerator[list[IndexResult], None]:
         """Embed a list of documents.
 
         The embedding service accepts a massive list of id,strings to embed. Behind the
@@ -36,3 +47,9 @@ class VectorSearchService(ABC):
     @abstractmethod
     async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
         """Query the embedding model."""
+
+    @abstractmethod
+    async def has_embedding(
+        self, snippet_id: int, embedding_type: EmbeddingType
+    ) -> bool:
+        """Check if a snippet has an embedding."""

{kodit-0.2.3 → kodit-0.2.4}/src/kodit/embedding/vectorchord_vector_search_service.py

@@ -1,13 +1,19 @@
 """Vectorchord vector search."""
 
+from collections.abc import AsyncGenerator
 from typing import Any, Literal
 
 import structlog
 from sqlalchemy import Result, TextClause, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from kodit.embedding.embedding_provider.embedding_provider import EmbeddingProvider
+from kodit.embedding.embedding_models import EmbeddingType
+from kodit.embedding.embedding_provider.embedding_provider import (
+    EmbeddingProvider,
+    EmbeddingRequest,
+)
 from kodit.embedding.vector_search_service import (
+    IndexResult,
     VectorSearchRequest,
     VectorSearchResponse,
     VectorSearchService,
@@ -52,6 +58,10 @@ ORDER BY score ASC
 LIMIT :top_k;
 """
 
+CHECK_VCHORD_EMBEDDING_EXISTS = """
+SELECT EXISTS(SELECT 1 FROM {TABLE_NAME} WHERE snippet_id = :snippet_id)
+"""
+
 TaskName = Literal["code", "text"]
 
 
@@ -89,7 +99,15 @@ class VectorChordVectorSearchService(VectorSearchService):
 
     async def _create_tables(self) -> None:
         """Create the necessary tables."""
-        vector_dim = (await self.embedding_provider.embed(["dimension"]))[0]
+        req = EmbeddingRequest(id=0, text="dimension")
+        vector_dim: list[float] | None = None
+        async for batch in self.embedding_provider.embed([req]):
+            if batch:
+                vector_dim = batch[0].embedding
+                break
+        if vector_dim is None:
+            msg = "Failed to obtain embedding dimension from provider"
+            raise RuntimeError(msg)
         await self._session.execute(
             text(
                 f"""CREATE TABLE IF NOT EXISTS {self.table_name} (
@@ -130,31 +148,48 @@ class VectorChordVectorSearchService(VectorSearchService):
         """Commit the session."""
         await self._session.commit()
 
-    async def index(self, data: list[VectorSearchRequest]) -> None:
+    async def index(
+        self, data: list[VectorSearchRequest]
+    ) -> AsyncGenerator[list[IndexResult], None]:
         """Embed a list of documents."""
         if not data or len(data) == 0:
             self.log.warning("Embedding data is empty, skipping embedding")
             return
 
-        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()
+        requests = [EmbeddingRequest(id=doc.snippet_id, text=doc.text) for doc in data]
+
+        async for batch in self.embedding_provider.embed(requests):
+            await self._execute(
+                text(INSERT_QUERY.format(TABLE_NAME=self.table_name)),
+                [
+                    {
+                        "snippet_id": result.id,
+                        "embedding": str(result.embedding),
+                    }
+                    for result in batch
+                ],
+            )
+            await self._commit()
+            yield [IndexResult(snippet_id=result.id) for result in batch]
 
     async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
         """Query the embedding model."""
-        embedding = await self.embedding_provider.embed([query])
-        if len(embedding) == 0 or len(embedding[0]) == 0:
+        from kodit.embedding.embedding_provider.embedding_provider import (
+            EmbeddingRequest,
+        )
+
+        req = EmbeddingRequest(id=0, text=query)
+        embedding_vec: list[float] | None = None
+        async for batch in self.embedding_provider.embed([req]):
+            if batch:
+                embedding_vec = batch[0].embedding
+                break
+
+        if not embedding_vec:
             return []
         result = await self._execute(
             text(SEARCH_QUERY.format(TABLE_NAME=self.table_name)),
-            {"query": str(embedding[0]), "top_k": top_k},
+            {"query": str(embedding_vec), "top_k": top_k},
         )
         rows = result.mappings().all()
 
@@ -162,3 +197,15 @@ class VectorChordVectorSearchService(VectorSearchService):
             VectorSearchResponse(snippet_id=row["snippet_id"], score=row["score"])
             for row in rows
         ]
+
+    async def has_embedding(
+        self,
+        snippet_id: int,
+        embedding_type: EmbeddingType,  # noqa: ARG002
+    ) -> bool:
+        """Check if a snippet has an embedding."""
+        result = await self._execute(
+            text(CHECK_VCHORD_EMBEDDING_EXISTS.format(TABLE_NAME=self.table_name)),
+            {"snippet_id": snippet_id},
+        )
+        return result.scalar_one()

kodit-0.2.4/src/kodit/enrichment/enrichment_provider/enrichment_provider.py

@@ -0,0 +1,36 @@
+"""Enrichment provider."""
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
+
+ENRICHMENT_SYSTEM_PROMPT = """
+You are a professional software developer. You will be given a snippet of code.
+Please provide a concise explanation of the code.
+"""
+
+
+@dataclass
+class EnrichmentRequest:
+    """Enrichment request."""
+
+    snippet_id: int
+    text: str
+
+
+@dataclass
+class EnrichmentResponse:
+    """Enrichment response."""
+
+    snippet_id: int
+    text: str
+
+
+class EnrichmentProvider(ABC):
+    """Enrichment provider."""
+
+    @abstractmethod
+    def enrich(
+        self, data: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
+        """Enrich a list of strings."""