kodit 0.2.3__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.3'
-__version_tuple__ = version_tuple = (0, 2, 3)
+__version__ = version = '0.2.4'
+__version_tuple__ = version_tuple = (0, 2, 4)

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/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, disallowed_special=()))
+            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, disallowed_special=()))
-
-                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."""

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/enrichment/enrichment_provider/enrichment_provider.py

@@ -1,6 +1,8 @@
 """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.
@@ -8,9 +10,27 @@ 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
-    async def enrich(self, data: list[str]) -> list[str]:
+    def enrich(
+        self, data: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
         """Enrich a list of strings."""

kodit/enrichment/enrichment_provider/local_enrichment_provider.py

@@ -1,15 +1,19 @@
 """Local embedding service."""
 
 import os
+from collections.abc import AsyncGenerator
 
 import structlog
 import tiktoken
-from tqdm import tqdm
 
-from kodit.embedding.embedding_provider.embedding_provider import split_sub_batches
+from kodit.embedding.embedding_provider.embedding_provider import (
+    EmbeddingRequest,
+)
 from kodit.enrichment.enrichment_provider.enrichment_provider import (
     ENRICHMENT_SYSTEM_PROMPT,
     EnrichmentProvider,
+    EnrichmentRequest,
+    EnrichmentResponse,
 )
 
 DEFAULT_ENRICHMENT_MODEL = "Qwen/Qwen3-0.6B"
@@ -32,11 +36,16 @@ class LocalEnrichmentProvider(EnrichmentProvider):
         self.tokenizer = None
         self.encoding = tiktoken.encoding_for_model("text-embedding-3-small")
 
-    async def enrich(self, data: list[str]) -> list[str]:
+    async def enrich(
+        self, data: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
         """Enrich a list of strings."""
+        # Remove empty snippets
+        data = [snippet for snippet in data if snippet.text]
+
         if not data or len(data) == 0:
             self.log.warning("Data is empty, skipping enrichment")
-            return []
+            return
 
         from transformers.models.auto.modeling_auto import (
             AutoModelForCausalLM,
@@ -57,36 +66,38 @@ class LocalEnrichmentProvider(EnrichmentProvider):
             )
 
         # Prepare prompts
-        prompts = [
-            self.tokenizer.apply_chat_template(
-                [
-                    {"role": "system", "content": ENRICHMENT_SYSTEM_PROMPT},
-                    {"role": "user", "content": snippet},
-                ],
-                tokenize=False,
-                add_generation_prompt=True,
-                enable_thinking=False,
+        prompts: list[EmbeddingRequest] = [
+            EmbeddingRequest(
+                id=snippet.snippet_id,
+                text=self.tokenizer.apply_chat_template(
+                    [
+                        {"role": "system", "content": ENRICHMENT_SYSTEM_PROMPT},
+                        {"role": "user", "content": snippet.text},
+                    ],
+                    tokenize=False,
+                    add_generation_prompt=True,
+                    enable_thinking=False,
+                ),
             )
             for snippet in data
         ]
 
-        # Batch prompts using split_sub_batches
-        batched_prompts = split_sub_batches(
-            self.encoding, prompts, max_context_window=self.context_window
-        )
-        results = []
-        for batch in tqdm(batched_prompts, leave=False, total=len(batched_prompts)):
+        for prompt in prompts:
             model_inputs = self.tokenizer(
-                batch, return_tensors="pt", padding=True, truncation=True
+                prompt.text,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
             ).to(self.model.device)
             generated_ids = self.model.generate(
                 **model_inputs, max_new_tokens=self.context_window
             )
-            # For each prompt in the batch, decode only the generated part
-            for i, input_ids in enumerate(model_inputs["input_ids"]):
-                output_ids = generated_ids[i][len(input_ids) :].tolist()
-                content = self.tokenizer.decode(
-                    output_ids, skip_special_tokens=True
-                ).strip("\n")
-                results.append(content)
-        return results
+            input_ids = model_inputs["input_ids"][0]
+            output_ids = generated_ids[0][len(input_ids) :].tolist()
+            content = self.tokenizer.decode(output_ids, skip_special_tokens=True).strip(
+                "\n"
+            )
+            yield EnrichmentResponse(
+                snippet_id=prompt.id,
+                text=content,
+            )

kodit/enrichment/enrichment_provider/openai_enrichment_provider.py

@@ -1,15 +1,17 @@
 """OpenAI embedding service."""
 
 import asyncio
+from collections.abc import AsyncGenerator
 
 import structlog
 import tiktoken
 from openai import AsyncOpenAI
-from tqdm import tqdm
 
 from kodit.enrichment.enrichment_provider.enrichment_provider import (
     ENRICHMENT_SYSTEM_PROMPT,
     EnrichmentProvider,
+    EnrichmentRequest,
+    EnrichmentResponse,
 )
 
 OPENAI_NUM_PARALLEL_TASKS = 10
@@ -29,25 +31,24 @@ class OpenAIEnrichmentProvider(EnrichmentProvider):
         self.model_name = model_name
         self.encoding = tiktoken.encoding_for_model("gpt-4o-mini")  # Approximation
 
-    async def enrich(self, data: list[str]) -> list[str]:
+    async def enrich(
+        self, data: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
         """Enrich a list of documents."""
         if not data or len(data) == 0:
             self.log.warning("Data is empty, skipping enrichment")
-            return []
+            return
 
         # 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 snippet
-        # We need to do this so that we can return the results in the same order as the
-        # input data
-        input_data = [(i, snippet) for i, snippet in enumerate(data)]
-
-        async def process_data(data: tuple[int, str]) -> tuple[int, str]:
-            snippet_id, snippet = data
-            if not snippet:
-                return snippet_id, ""
+        async def process_data(data: EnrichmentRequest) -> EnrichmentResponse:
             async with sem:
+                if not data.text:
+                    return EnrichmentResponse(
+                        snippet_id=data.snippet_id,
+                        text="",
+                    )
                 try:
                     response = await self.openai_client.chat.completions.create(
                         model=self.model_name,
@@ -56,26 +57,23 @@ class OpenAIEnrichmentProvider(EnrichmentProvider):
                                 "role": "system",
                                 "content": ENRICHMENT_SYSTEM_PROMPT,
                             },
-                            {"role": "user", "content": snippet},
+                            {"role": "user", "content": data.text},
                         ],
                     )
-                    return snippet_id, response.choices[0].message.content or ""
+                    return EnrichmentResponse(
+                        snippet_id=data.snippet_id,
+                        text=response.choices[0].message.content or "",
+                    )
                 except Exception as e:
                     self.log.exception("Error enriching data", error=str(e))
-                    return snippet_id, ""
+                    return EnrichmentResponse(
+                        snippet_id=data.snippet_id,
+                        text="",
+                    )
 
         # Create tasks for all data
-        tasks = [process_data(snippet) for snippet in input_data]
+        tasks = [process_data(snippet) for snippet in data]
 
         # Process all data and yield results as they complete
-        results: list[tuple[int, str]] = []
-        for task in tqdm(
-            asyncio.as_completed(tasks),
-            total=len(tasks),
-            leave=False,
-        ):
-            result = await task
-            results.append(result)
-
-        # Output in the same order as the input data
-        return [result for _, result in sorted(results, key=lambda x: x[0])]
+        for task in asyncio.as_completed(tasks):
+            yield await task

kodit/enrichment/enrichment_service.py

@@ -1,24 +1,34 @@
 """Enrichment service."""
 
 from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
 
-from kodit.enrichment.enrichment_provider.enrichment_provider import EnrichmentProvider
+from kodit.enrichment.enrichment_provider.enrichment_provider import (
+    EnrichmentProvider,
+    EnrichmentRequest,
+    EnrichmentResponse,
+)
 
 
 class EnrichmentService(ABC):
     """Enrichment service."""
 
     @abstractmethod
-    async def enrich(self, data: list[str]) -> list[str]:
+    def enrich(
+        self, data: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
         """Enrich a list of strings."""
 
 
 class NullEnrichmentService(EnrichmentService):
     """Null enrichment service."""
 
-    async def enrich(self, data: list[str]) -> list[str]:
+    async def enrich(
+        self, data: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
         """Enrich a list of strings."""
-        return [""] * len(data)
+        for request in data:
+            yield EnrichmentResponse(snippet_id=request.snippet_id, text="")
 
 
 class LLMEnrichmentService(EnrichmentService):
@@ -28,6 +38,8 @@ class LLMEnrichmentService(EnrichmentService):
         """Initialize the enrichment service."""
         self.enrichment_provider = enrichment_provider
 
-    async def enrich(self, data: list[str]) -> list[str]:
-        """Enrich a list of strings."""
-        return await self.enrichment_provider.enrich(data)
+    def enrich(
+        self, data: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
+        """Enrich a list of snippets."""
+        return self.enrichment_provider.enrich(data)

kodit/indexing/indexing_service.py

@@ -22,6 +22,7 @@ from kodit.embedding.vector_search_service import (
     VectorSearchRequest,
     VectorSearchService,
 )
+from kodit.enrichment.enrichment_provider.enrichment_provider import EnrichmentRequest
 from kodit.enrichment.enrichment_service import EnrichmentService
 from kodit.indexing.fusion import FusionRequest, reciprocal_rank_fusion
 from kodit.indexing.indexing_models import Snippet
@@ -200,37 +201,42 @@ class IndexService:
             )
 
         self.log.info("Creating semantic code index")
-        with Spinner():
-            await self.code_search_service.index(
+        with tqdm(total=len(snippets), leave=False) as pbar:
+            async for result in self.code_search_service.index(
                 [
                     VectorSearchRequest(snippet.id, snippet.content)
                     for snippet in snippets
                 ]
-            )
+            ):
+                pbar.update(len(result))
 
         self.log.info("Enriching snippets", num_snippets=len(snippets))
-        enriched_contents = await self.enrichment_service.enrich(
-            [snippet.content for snippet in snippets]
-        )
+        enriched_contents = []
+        with tqdm(total=len(snippets), leave=False) as pbar:
+            async for result in self.enrichment_service.enrich(
+                [
+                    EnrichmentRequest(snippet_id=snippet.id, text=snippet.content)
+                    for snippet in snippets
+                ]
+            ):
+                snippet = next(s for s in snippets if s.id == result.snippet_id)
+                if snippet:
+                    snippet.content = (
+                        result.text + "\n\n```\n" + snippet.content + "\n```"
+                    )
+                    await self.repository.add_snippet(snippet)
+                    enriched_contents.append(result)
+                pbar.update(1)
 
         self.log.info("Creating semantic text index")
-        with Spinner():
-            await self.text_search_service.index(
+        with tqdm(total=len(snippets), leave=False) as pbar:
+            async for result in self.text_search_service.index(
                 [
-                    VectorSearchRequest(snippet.id, enriched_content)
-                    for snippet, enriched_content in zip(
-                        snippets, enriched_contents, strict=True
-                    )
+                    VectorSearchRequest(snippet.id, snippet.content)
+                    for snippet in snippets
                 ]
-            )
-            # Add the enriched text back to the snippets and write to the database
-            for snippet, enriched_content in zip(
-                snippets, enriched_contents, strict=True
             ):
-                snippet.content = (
-                    enriched_content + "\n\n```\n" + snippet.content + "\n```"
-                )
-                await self.repository.add_snippet(snippet)
+                pbar.update(len(result))
 
         # Update index timestamp
         await self.repository.update_index_timestamp(index)

{kodit-0.2.3.dist-info → kodit-0.2.4.dist-info}/METADATA

@@ -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.dist-info → kodit-0.2.4.dist-info}/RECORD

@@ -1,6 +1,6 @@
 kodit/.gitignore,sha256=ztkjgRwL9Uud1OEi36hGQeDGk3OLK1NfDEO8YqGYy8o,11
 kodit/__init__.py,sha256=aEKHYninUq1yh6jaNfvJBYg-6fenpN132nJt1UU6Jxs,59
-kodit/_version.py,sha256=wD8hnA5gV5UmPkQnpT3xR6V2csgj9K5NEADogbLK79M,511
+kodit/_version.py,sha256=1LUN_sRKOiFInoB6AlW6TYoQMCh1Z4KutwcHNvHcfB0,511
 kodit/app.py,sha256=qKBWJ0VNSY_M6G3VFfAQ0133q5bnS99cUFD0p396taw,1032
 kodit/cli.py,sha256=wKFXGUMX-fDLooaK-3po2TBpNNRBwgSD7BRbUddg-_M,11562
 kodit/config.py,sha256=3yh7hfLSILjZK_qJMhcExwRcrWJ0b5Eb1JjjOvMPJZo,4146
@@ -14,29 +14,29 @@ kodit/bm25/keyword_search_service.py,sha256=aBbWQKgQmi2re3EIHdXFS00n7Wj3b2D0pZsL
 kodit/bm25/local_bm25.py,sha256=nokrd_xAeqXi3m68X5P1R5KBhRRB1E2L_J6Zgm26PCg,3869
 kodit/bm25/vectorchord_bm25.py,sha256=0p_FgliaoevB8GLSmzWnV3zUjdcWgCgOKIpLURr7Qfo,6549
 kodit/embedding/__init__.py,sha256=h9NXzDA1r-K23nvBajBV-RJzHJN0p3UJ7UQsmdnOoRw,24
-kodit/embedding/embedding_factory.py,sha256=-WuXNleQ_mqdw1E4TczjtOawNeaXKAiDPFqN_XX7Mmg,2419
+kodit/embedding/embedding_factory.py,sha256=lFcgqsDxw8L5mygq-TppQ2wtoIA2p2OL7XmtOyX8Omw,2683
 kodit/embedding/embedding_models.py,sha256=rN90vSs86dYiqoawcp8E9jtwY31JoJXYfaDlsJK7uqc,656
 kodit/embedding/embedding_repository.py,sha256=-ux3scpBzel8c0pMH9fNOEsSXFIzl-IfgaWrkTb1szo,6907
-kodit/embedding/local_vector_search_service.py,sha256=dgMi8hQNUbYEgHnEYmLIpon4yLduoNUpu7k7VP6sOHI,2042
-kodit/embedding/vector_search_service.py,sha256=pQJ129QjGrAWOXzqkywmgtDRpy8_gtzYgkivyqF9Vrs,1009
-kodit/embedding/vectorchord_vector_search_service.py,sha256=TKNR3HgWHwwWtJ1SsvSaj_BXLJ_uw6Bdr_tpaePMeAA,5383
+kodit/embedding/local_vector_search_service.py,sha256=yZm0ahQQKhfYZ943yxKHp04cairmzgGBUNi5PB_GDbo,3002
+kodit/embedding/vector_search_service.py,sha256=frN9baAlqFmsY3xiv1ZeSgsfhK9FzKPkVR55MEvMV4I,1416
+kodit/embedding/vectorchord_vector_search_service.py,sha256=JQeIl9mtR4E_izOoFD_4ZRfENHNfwoKr16pQkkGoK3o,6884
 kodit/embedding/embedding_provider/__init__.py,sha256=h9NXzDA1r-K23nvBajBV-RJzHJN0p3UJ7UQsmdnOoRw,24
-kodit/embedding/embedding_provider/embedding_provider.py,sha256=T6AKMWwzEJ1vDe-cEIg-qxdjhUEZ0PKs9YqQMWaLaQ0,1928
-kodit/embedding/embedding_provider/hash_embedding_provider.py,sha256=nAhlhh8j8PqqCCbhVl26Y8ntFBm2vJBCtB4X04g5Wwg,2638
-kodit/embedding/embedding_provider/local_embedding_provider.py,sha256=WP8lw6XG7v1_5Mw4_rhIOETooYRsxhkwmFaXCqCouQU,1977
-kodit/embedding/embedding_provider/openai_embedding_provider.py,sha256=-phz5FKYM_tI3Q4_3SPzjzIOK3k92Uk52TAOTmoVoWI,2722
+kodit/embedding/embedding_provider/embedding_provider.py,sha256=WDHifrsQOnpXwIDzSfau32Eq8z8BF3XNeVYd6X989uc,2841
+kodit/embedding/embedding_provider/hash_embedding_provider.py,sha256=AhGize94EoScyQMhCjo26zlO0eP_m3F_1qvrVmB6MTE,2941
+kodit/embedding/embedding_provider/local_embedding_provider.py,sha256=kqbGd7TW6BUsOq_f_IzPCsD7z8LsFieTOZ7saY11I8o,2877
+kodit/embedding/embedding_provider/openai_embedding_provider.py,sha256=2FTIL34yVstf0NTJNSi-sjk38OJd4Aa66TH5FMPJul0,2425
 kodit/enrichment/__init__.py,sha256=vBEolHpKaHUhfINX0dSGyAPlvgpLNAer9YzFtdvCB24,18
 kodit/enrichment/enrichment_factory.py,sha256=AAzvxgjo-FQU5aAm9Zla4DAwUMKGrcw8mQwJsMhIsHY,1566
-kodit/enrichment/enrichment_service.py,sha256=87Sd3gGbEMJYb_wVrHG8L1yGIZmQNR7foUS4_y94azI,977
+kodit/enrichment/enrichment_service.py,sha256=z7VrrQ-Jhb-oO26rQCaqlpmkGRlDQGAu7qVsI0cwHak,1310
 kodit/enrichment/enrichment_provider/__init__.py,sha256=klf8iuLVWX4iRz-DZQauFFNAoJC5CByczh48TBZPW-o,27
-kodit/enrichment/enrichment_provider/enrichment_provider.py,sha256=E0H5rq3OENM0yYbA8K_3nSnj5lUHCpoIOqpWLo-2MVU,413
-kodit/enrichment/enrichment_provider/local_enrichment_provider.py,sha256=RqwUD0BnwRQ8zlkFNkaKq8d58r33k2jIdnSdf6zla1w,3325
-kodit/enrichment/enrichment_provider/openai_enrichment_provider.py,sha256=0Yw7h9RXptoI4bKuqJSKIRQXPUUhNV7eACavgoy_T8s,2874
+kodit/enrichment/enrichment_provider/enrichment_provider.py,sha256=kiDgg2G8G85K4KqwCQKHE_ANybANURPO6NbASf4yAr0,751
+kodit/enrichment/enrichment_provider/local_enrichment_provider.py,sha256=JipvgZwfNvHaECqLJRmQo3W10yb9rOVSrV2U0Jpp4d8,3456
+kodit/enrichment/enrichment_provider/openai_enrichment_provider.py,sha256=xAzbGHJHqGxZxa3yPvHAcPgjOMzQ05qLes0XW6OIdYc,2758
 kodit/indexing/__init__.py,sha256=cPyi2Iej3G1JFWlWr7X80_UrsMaTu5W5rBwgif1B3xo,75
 kodit/indexing/fusion.py,sha256=TZb4fPAedXdEUXzwzOofW98QIOymdbclBOP1KOijuEk,1674
 kodit/indexing/indexing_models.py,sha256=6NX9HVcj6Pu9ePwHC7n-PWSyAgukpJq0nCNmUIigtbo,1282
 kodit/indexing/indexing_repository.py,sha256=dqOS0pxKM6bUjMXWqYukAK8XdiD36OnskFASgZRXRQM,6955
-kodit/indexing/indexing_service.py,sha256=3hW7vbFyabLEkLU-PRoGR49yVLewANdOKlye4GhR-tw,11467
+kodit/indexing/indexing_service.py,sha256=UD7RKQRkAlpmepl20vcdEgQapwEA2kDJQBmn4_kGWwU,11841
 kodit/migrations/README,sha256=ISVtAOvqvKk_5ThM5ioJE-lMkvf9IbknFUFVU_vPma4,58
 kodit/migrations/__init__.py,sha256=lP5MuwlyWRMO6UcDWnQcQ3G-GYHcFb6rl9gYPHJ1sjo,40
 kodit/migrations/env.py,sha256=w1M7OZh-ZeR2dPHS0ByXAUxQjfZQ8xIzMseWuzLDTWw,2469
@@ -64,8 +64,8 @@ kodit/source/source_repository.py,sha256=eme0C3pRqwFZ1ZSbqq4Z6SV9CC6AvRmiOjy3eHQ
 kodit/source/source_service.py,sha256=E1KPG7TrorqdreJVHxZPx8CVLncOxGEvZ5uDQ6yZugo,5050
 kodit/util/__init__.py,sha256=bPu6CtqDWCRGU7VgW2_aiQrCBi8G89FS6k1PjvDajJ0,37
 kodit/util/spinner.py,sha256=R9bzrHtBiIH6IfLbmsIVHL53s8vg-tqW4lwGGALu4dw,1932
-kodit-0.2.3.dist-info/METADATA,sha256=ccZ0bl5PCGFNC30XSr-4ljL-JrkTgTfYlqZUlUPYba8,5867
-kodit-0.2.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-kodit-0.2.3.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
-kodit-0.2.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-kodit-0.2.3.dist-info/RECORD,,
+kodit-0.2.4.dist-info/METADATA,sha256=PLQQVNKVnMyyliP9TEapeXUuog_N1bTFlup6F89B7NU,5867
+kodit-0.2.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+kodit-0.2.4.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
+kodit-0.2.4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+kodit-0.2.4.dist-info/RECORD,,