kodit 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

kodit/_version.py

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

kodit/app.py

@@ -21,6 +21,12 @@ async def root() -> dict[str, str]:
     return {"message": "Hello, World!"}
 
 
+@app.get("/healthz")
+async def healthz() -> dict[str, str]:
+    """Return a health check for the kodit API."""
+    return {"status": "ok"}
+
+
 # Add mcp routes last, otherwise previous routes aren't added
 app.mount("", mcp_app)
 

kodit/cli.py

@@ -81,6 +81,7 @@ async def index(
     )
 
     if not sources:
+        log_event("kodit.cli.index.list")
         # No source specified, list all indexes
         indexes = await service.list_indexes()
         headers: list[str | Cell] = [
@@ -108,7 +109,8 @@ async def index(
             msg = "File indexing is not implemented yet"
             raise click.UsageError(msg)
 
-        # Index directory
+        # Index source
+        log_event("kodit.cli.index.create")
         s = await source_service.create(source)
         index = await service.create(s.id)
         await service.run(index.id)
@@ -134,6 +136,7 @@ async def code(
 
     This works best if your query is code.
     """
+    log_event("kodit.cli.search.code")
     source_repository = SourceRepository(session)
     source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
@@ -177,6 +180,7 @@ async def keyword(
     top_k: int,
 ) -> None:
     """Search for snippets using keyword search."""
+    log_event("kodit.cli.search.keyword")
     source_repository = SourceRepository(session)
     source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
@@ -223,6 +227,7 @@ async def text(
 
     This works best if your query is text.
     """
+    log_event("kodit.cli.search.text")
     source_repository = SourceRepository(session)
     source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
@@ -270,6 +275,7 @@ async def hybrid(  # noqa: PLR0913
     text: str,
 ) -> None:
     """Search for snippets using hybrid search."""
+    log_event("kodit.cli.search.hybrid")
     source_repository = SourceRepository(session)
     source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
@@ -321,7 +327,7 @@ def serve(
     """Start the kodit server, which hosts the MCP server and the kodit API."""
     log = structlog.get_logger(__name__)
     log.info("Starting kodit server", host=host, port=port)
-    log_event("kodit_server_started")
+    log_event("kodit.cli.serve")
 
     # Configure uvicorn with graceful shutdown
     config = uvicorn.Config(

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,
@@ -19,6 +20,7 @@ from kodit.embedding.vectorchord_vector_search_service import (
     TaskName,
     VectorChordVectorSearchService,
 )
+from kodit.log import log_event
 
 
 def _get_endpoint_configuration(app_context: AppContext) -> Endpoint | None:
@@ -34,6 +36,7 @@ def embedding_factory(
     endpoint = _get_endpoint_configuration(app_context)
 
     if endpoint and endpoint.type == "openai":
+        log_event("kodit.embedding", {"provider": "openai"})
         from openai import AsyncOpenAI
 
         embedding_provider = OpenAIEmbeddingProvider(
@@ -44,14 +47,22 @@ def embedding_factory(
             model_name=endpoint.model or "text-embedding-3-small",
         )
     else:
+        log_event("kodit.embedding", {"provider": "local"})
         embedding_provider = LocalEmbeddingProvider(CODE)
 
     if app_context.default_search.provider == "vectorchord":
+        log_event("kodit.database", {"provider": "vectorchord"})
         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/embedding/embedding_provider/embedding_provider.py

@@ -1,6 +1,8 @@
 """Embedding provider."""
 
 from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
 
 import structlog
 import tiktoken
@@ -10,11 +12,29 @@ 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
-    async def embed(self, data: list[str]) -> list[Vector]:
+    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
@@ -25,13 +45,13 @@ class EmbeddingProvider(ABC):
 
 def split_sub_batches(
     encoding: tiktoken.Encoding,
-    data: list[str],
+    data: list[EmbeddingRequest],
     max_context_window: int = OPENAI_MAX_EMBEDDING_SIZE,
-) -> list[list[str]]:
+) -> 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.strip()]  # Filter out empty strings
+    data_to_process = [s for s in data if s.text.strip()]  # Filter out empty strings
 
     while data_to_process:
         next_batch = []
@@ -39,18 +59,26 @@ def split_sub_batches(
 
         while data_to_process:
             next_item = data_to_process[0]
-            item_tokens = len(encoding.encode(next_item))
+            item_tokens = len(encoding.encode(next_item.text, disallowed_special=()))
 
             if item_tokens > max_context_window:
-                # Loop around trying to truncate the snippet until it fits in the max
-                # embedding size
-                while item_tokens > max_context_window:
-                    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)
+                # 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

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/embedding/embedding_provider/local_embedding_provider.py

@@ -3,20 +3,24 @@
 from __future__ import annotations
 
 import os
+from time import time
 from typing import TYPE_CHECKING
 
 import structlog
-import tiktoken
-from tqdm import tqdm
 
 from kodit.embedding.embedding_provider.embedding_provider import (
     EmbeddingProvider,
-    Vector,
+    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"
@@ -36,8 +40,22 @@ class LocalEmbeddingProvider(EmbeddingProvider):
         """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 = tiktoken.encoding_for_model("text-embedding-3-small")
+        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."""
@@ -45,20 +63,34 @@ class LocalEmbeddingProvider(EmbeddingProvider):
             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[str]) -> list[Vector]:
+    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)
+        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
+        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/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/embedding/local_vector_search_service.py

@@ -1,12 +1,18 @@
 """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
+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,
@@ -20,35 +26,62 @@ class LocalVectorSearchService(VectorSearchService):
         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]) -> 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([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,
+        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."""
-        embedding = (await self.embedding_provider.embed([query]))[0]
+        # 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(
-            EmbeddingType.CODE, [float(x) for x in embedding], top_k
+            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/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."""