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

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_factory.py

@@ -11,6 +11,7 @@ from kodit.enrichment.enrichment_service import (
     EnrichmentService,
     LLMEnrichmentService,
 )
+from kodit.log import log_event
 
 
 def _get_endpoint_configuration(app_context: AppContext) -> Endpoint | None:
@@ -24,6 +25,7 @@ def enrichment_factory(app_context: AppContext) -> EnrichmentService:
     endpoint = app_context.enrichment_endpoint or app_context.default_endpoint or None
 
     if endpoint and endpoint.type == "openai":
+        log_event("kodit.enrichment", {"provider": "openai"})
         from openai import AsyncOpenAI
 
         enrichment_provider = OpenAIEnrichmentProvider(
@@ -34,6 +36,7 @@ def enrichment_factory(app_context: AppContext) -> EnrichmentService:
             model_name=endpoint.model or "gpt-4o-mini",
         )
     else:
+        log_event("kodit.enrichment", {"provider": "local"})
         enrichment_provider = LocalEnrichmentProvider()
 
     return LLMEnrichmentService(enrichment_provider=enrichment_provider)

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,10 +22,12 @@ 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
 from kodit.indexing.indexing_repository import IndexRepository
+from kodit.log import log_event
 from kodit.snippets.snippets import SnippetService
 from kodit.source.source_service import SourceService
 from kodit.util.spinner import Spinner
@@ -45,7 +47,7 @@ class IndexView(pydantic.BaseModel):
     created_at: datetime
     updated_at: datetime | None = None
     source: str | None = None
-    num_snippets: int | None = None
+    num_snippets: int
 
 
 class SearchRequest(pydantic.BaseModel):
@@ -119,6 +121,8 @@ class IndexService:
             ValueError: If the source doesn't exist or already has an index.
 
         """
+        log_event("kodit.index.create")
+
         # Check if the source exists
         source = await self.source_service.get(source_id)
 
@@ -129,6 +133,8 @@ class IndexService:
         return IndexView(
             id=index.id,
             created_at=index.created_at,
+            num_snippets=await self.repository.num_snippets_for_index(index.id),
+            source=source.uri,
         )
 
     async def list_indexes(self) -> list[IndexView]:
@@ -142,19 +148,33 @@ class IndexService:
         indexes = await self.repository.list_indexes()
 
         # Transform database results into DTOs
-        return [
+        indexes = [
             IndexView(
                 id=index.id,
                 created_at=index.created_at,
                 updated_at=index.updated_at,
-                num_snippets=await self.repository.num_snippets_for_index(index.id),
+                num_snippets=await self.repository.num_snippets_for_index(index.id)
+                or 0,
                 source=source.uri,
             )
             for index, source in indexes
         ]
 
+        # Help Kodit by measuring how much people are using indexes
+        log_event(
+            "kodit.index.list",
+            {
+                "num_indexes": len(indexes),
+                "num_snippets": sum([index.num_snippets for index in indexes]),
+            },
+        )
+
+        return indexes
+
     async def run(self, index_id: int) -> None:
         """Run the indexing process for a specific index."""
+        log_event("kodit.index.run")
+
         # Get and validate index
         index = await self.repository.get_by_id(index_id)
         if not index:
@@ -181,43 +201,50 @@ 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)
 
     async def search(self, request: SearchRequest) -> list[SearchResult]:
         """Search for relevant data."""
+        log_event("kodit.index.search")
+
         fusion_list: list[list[FusionRequest]] = []
         if request.keywords:
             # Gather results for each keyword