kodit 0.2.1__py3-none-any.whl → 0.2.3__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.1'
-__version_tuple__ = version_tuple = (0, 2, 1)
+__version__ = version = '0.2.3'
+__version_tuple__ = version_tuple = (0, 2, 3)

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/bm25/local_bm25.py

@@ -64,6 +64,10 @@ class BM25Service(KeywordSearchProvider):
     async def index(self, corpus: list[BM25Document]) -> None:
         """Index a new corpus."""
         self.log.debug("Indexing corpus")
+        if not corpus or len(corpus) == 0:
+            self.log.warning("Corpus is empty, skipping bm25 index")
+            return
+
         vocab = self._tokenize([doc.text for doc in corpus])
         self._retriever().index(vocab, show_progress=False)
         self._retriever().save(self.index_path)
@@ -77,6 +81,10 @@ class BM25Service(KeywordSearchProvider):
             self.log.warning("Top k is 0, returning empty list")
             return []
 
+        # Check that the index has data
+        if not hasattr(self._retriever(), "scores"):
+            return []
+
         # Get the number of documents in the index
         num_docs = self._retriever().scores["num_docs"]
         if num_docs == 0:

kodit/bm25/vectorchord_bm25.py

@@ -2,6 +2,7 @@
 
 from typing import Any
 
+import structlog
 from sqlalchemy import Result, TextClause, bindparam, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -93,6 +94,7 @@ class VectorChordBM25(KeywordSearchProvider):
         """Initialize the VectorChord BM25."""
         self.__session = session
         self._initialized = False
+        self.log = structlog.get_logger(__name__)
 
     async def _initialize(self) -> None:
         """Initialize the VectorChord environment."""
@@ -149,7 +151,8 @@ class VectorChordBM25(KeywordSearchProvider):
             if doc.snippet_id is not None and doc.text is not None and doc.text != ""
         ]
 
-        if not corpus:
+        if not corpus or len(corpus) == 0:
+            self.log.warning("Corpus is empty, skipping bm25 index")
             return
 
         # Execute inserts

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/config.py

@@ -14,7 +14,6 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
 if TYPE_CHECKING:
     from collections.abc import Callable, Coroutine
 
-    from openai import AsyncOpenAI
 
 from kodit.database import Database
 
@@ -25,13 +24,16 @@ DEFAULT_LOG_FORMAT = "pretty"
 DEFAULT_DISABLE_TELEMETRY = False
 T = TypeVar("T")
 
+EndpointType = Literal["openai"]
+
 
 class Endpoint(BaseModel):
     """Endpoint provides configuration for an AI service."""
 
-    type: Literal["openai"] = Field(default="openai")
-    api_key: str | None = None
+    type: EndpointType | None = None
     base_url: str | None = None
+    model: str | None = None
+    api_key: str | None = None
 
 
 class Search(BaseModel):
@@ -57,15 +59,20 @@ class AppContext(BaseSettings):
     log_format: str = Field(default=DEFAULT_LOG_FORMAT)
     disable_telemetry: bool = Field(default=DEFAULT_DISABLE_TELEMETRY)
     default_endpoint: Endpoint | None = Field(
-        default=Endpoint(
-            type="openai",
-            base_url="https://api.openai.com/v1",
-        ),
+        default=None,
         description=(
             "Default endpoint to use for all AI interactions "
             "(can be overridden by task-specific configuration)."
         ),
     )
+    embedding_endpoint: Endpoint | None = Field(
+        default=None,
+        description="Endpoint to use for embedding.",
+    )
+    enrichment_endpoint: Endpoint | None = Field(
+        default=None,
+        description="Endpoint to use for enrichment.",
+    )
     default_search: Search = Field(
         default=Search(),
     )
@@ -95,23 +102,6 @@ class AppContext(BaseSettings):
             await self._db.run_migrations(self.db_url)
         return self._db
 
-    def get_default_openai_client(self) -> AsyncOpenAI | None:
-        """Get the default OpenAI client, if it is configured."""
-        from openai import AsyncOpenAI
-
-        endpoint = self.default_endpoint
-        if not (
-            endpoint
-            and endpoint.type == "openai"
-            and endpoint.api_key
-            and endpoint.base_url
-        ):
-            return None
-        return AsyncOpenAI(
-            api_key=endpoint.api_key,
-            base_url=endpoint.base_url,
-        )
-
 
 with_app_context = click.make_pass_decorator(AppContext)
 

kodit/embedding/embedding_factory.py

@@ -2,7 +2,7 @@
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from kodit.config import AppContext
+from kodit.config import AppContext, Endpoint
 from kodit.embedding.embedding_provider.local_embedding_provider import (
     CODE,
     LocalEmbeddingProvider,
@@ -16,25 +16,44 @@ from kodit.embedding.vector_search_service import (
     VectorSearchService,
 )
 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:
+    """Get the endpoint configuration for the embedding service."""
+    return app_context.embedding_endpoint or app_context.default_endpoint or None
 
 
 def embedding_factory(
-    task_name: str, app_context: AppContext, session: AsyncSession
+    task_name: TaskName, app_context: AppContext, session: AsyncSession
 ) -> VectorSearchService:
     """Create an embedding service."""
     embedding_repository = EmbeddingRepository(session=session)
-    embedding_provider = None
-    openai_client = app_context.get_default_openai_client()
-    if openai_client is not None:
-        embedding_provider = OpenAIEmbeddingProvider(openai_client=openai_client)
+    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(
+            openai_client=AsyncOpenAI(
+                api_key=endpoint.api_key or "default",
+                base_url=endpoint.base_url or "https://api.openai.com/v1",
+            ),
+            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"})
         return LocalVectorSearchService(
             embedding_repository=embedding_repository,
             embedding_provider=embedding_provider,

kodit/embedding/embedding_provider/embedding_provider.py

@@ -39,14 +39,14 @@ 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, 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))
+                    item_tokens = len(encoding.encode(next_item, disallowed_special=()))
 
                 data_to_process[0] = next_item
 

kodit/embedding/embedding_provider/openai_embedding_provider.py

@@ -27,7 +27,9 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
         self.log = structlog.get_logger(__name__)
         self.openai_client = openai_client
         self.model_name = model_name
-        self.encoding = tiktoken.encoding_for_model(model_name)
+        self.encoding = tiktoken.encoding_for_model(
+            "text-embedding-3-small"
+        )  # Sensible default
 
     async def embed(self, data: list[str]) -> list[Vector]:
         """Embed a list of documents."""

kodit/embedding/local_vector_search_service.py

@@ -29,6 +29,10 @@ class LocalVectorSearchService(VectorSearchService):
 
     async def index(self, data: list[VectorSearchRequest]) -> 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(

kodit/embedding/vectorchord_vector_search_service.py

@@ -1,7 +1,8 @@
 """Vectorchord vector search."""
 
-from typing import Any
+from typing import Any, Literal
 
+import structlog
 from sqlalchemy import Result, TextClause, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -51,13 +52,15 @@ ORDER BY score ASC
 LIMIT :top_k;
 """
 
+TaskName = Literal["code", "text"]
+
 
 class VectorChordVectorSearchService(VectorSearchService):
     """VectorChord vector search."""
 
     def __init__(
         self,
-        task_name: str,
+        task_name: TaskName,
         session: AsyncSession,
         embedding_provider: EmbeddingProvider,
     ) -> None:
@@ -67,6 +70,7 @@ class VectorChordVectorSearchService(VectorSearchService):
         self._initialized = False
         self.table_name = f"vectorchord_{task_name}_embeddings"
         self.index_name = f"{self.table_name}_idx"
+        self.log = structlog.get_logger(__name__)
 
     async def _initialize(self) -> None:
         """Initialize the VectorChord environment."""
@@ -128,6 +132,10 @@ class VectorChordVectorSearchService(VectorSearchService):
 
     async def index(self, data: list[VectorSearchRequest]) -> 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(

kodit/enrichment/enrichment_factory.py

@@ -1,6 +1,6 @@
 """Embedding service."""
 
-from kodit.config import AppContext
+from kodit.config import AppContext, Endpoint
 from kodit.enrichment.enrichment_provider.local_enrichment_provider import (
     LocalEnrichmentProvider,
 )
@@ -11,13 +11,32 @@ from kodit.enrichment.enrichment_service import (
     EnrichmentService,
     LLMEnrichmentService,
 )
+from kodit.log import log_event
+
+
+def _get_endpoint_configuration(app_context: AppContext) -> Endpoint | None:
+    """Get the endpoint configuration for the enrichment service."""
+    return app_context.enrichment_endpoint or app_context.default_endpoint or None
 
 
 def enrichment_factory(app_context: AppContext) -> EnrichmentService:
-    """Create an embedding service."""
-    openai_client = app_context.get_default_openai_client()
-    if openai_client is not None:
-        enrichment_provider = OpenAIEnrichmentProvider(openai_client=openai_client)
-        return LLMEnrichmentService(enrichment_provider)
+    """Create an enrichment service."""
+    endpoint = _get_endpoint_configuration(app_context)
+    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(
+            openai_client=AsyncOpenAI(
+                api_key=endpoint.api_key or "default",
+                base_url=endpoint.base_url or "https://api.openai.com/v1",
+            ),
+            model_name=endpoint.model or "gpt-4o-mini",
+        )
+    else:
+        log_event("kodit.enrichment", {"provider": "local"})
+        enrichment_provider = LocalEnrichmentProvider()
 
-    return LLMEnrichmentService(LocalEnrichmentProvider())
+    return LLMEnrichmentService(enrichment_provider=enrichment_provider)

kodit/enrichment/enrichment_provider/local_enrichment_provider.py

@@ -34,6 +34,10 @@ class LocalEnrichmentProvider(EnrichmentProvider):
 
     async def enrich(self, data: list[str]) -> list[str]:
         """Enrich a list of strings."""
+        if not data or len(data) == 0:
+            self.log.warning("Data is empty, skipping enrichment")
+            return []
+
         from transformers.models.auto.modeling_auto import (
             AutoModelForCausalLM,
         )

kodit/enrichment/enrichment_provider/openai_enrichment_provider.py

@@ -27,10 +27,14 @@ class OpenAIEnrichmentProvider(EnrichmentProvider):
         self.log = structlog.get_logger(__name__)
         self.openai_client = openai_client
         self.model_name = model_name
-        self.encoding = tiktoken.encoding_for_model(model_name)
+        self.encoding = tiktoken.encoding_for_model("gpt-4o-mini")  # Approximation
 
     async def enrich(self, data: list[str]) -> list[str]:
         """Enrich a list of documents."""
+        if not data or len(data) == 0:
+            self.log.warning("Data is empty, skipping enrichment")
+            return []
+
         # Process batches in parallel with a semaphore to limit concurrent requests
         sem = asyncio.Semaphore(OPENAI_NUM_PARALLEL_TASKS)
 

kodit/indexing/indexing_service.py

@@ -26,6 +26,7 @@ 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 +46,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 +120,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 +132,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 +147,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:
@@ -218,6 +237,8 @@ class IndexService:
 
     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
@@ -289,6 +310,10 @@ class IndexService:
 
         """
         files = await self.repository.files_for_index(index_id)
+        if not files:
+            self.log.warning("No files to create snippets for")
+            return
+
         for file in tqdm(files, total=len(files), leave=False):
             # Skip unsupported file types
             if file.mime_type in MIME_BLACKLIST: