kodit 0.5.2__tar.gz → 0.5.4__tar.gz

{kodit-0.5.2 → kodit-0.5.4}/.github/workflows/test.yaml

@@ -39,6 +39,7 @@ jobs:
         run: make no-database-changes-check
 
       - name: API Doc check
+        if: "!startsWith(github.ref, 'refs/tags/')"
         run: make docs-check
 
       - name: Run tests

{kodit-0.5.2 → kodit-0.5.4}/Makefile

@@ -21,7 +21,7 @@ type:
 lint:
 	uv run ruff check --fix --unsafe-fixes
 
-test: lint type docs-check
+test: lint type
 	uv run pytest -s --cov=src --cov-report=xml tests/kodit
 
 no-database-changes-check:

{kodit-0.5.2 → kodit-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.5.2
+Version: 0.5.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.5.2 → kodit-0.5.4}/docs/reference/api/index.md

@@ -12,7 +12,7 @@ look at the [hosted version](https://kodit.helix.ml/docs).
 This is the REST API for the Kodit server. Please refer to the
 [Kodit documentation](https://docs.helix.ml/kodit/) for more information.
     
-Current version: 0.5.1
+Current version: 0.5.3
 
 ## Authentication
 

{kodit-0.5.2 → kodit-0.5.4}/docs/reference/api/openapi.json

@@ -3,7 +3,7 @@
   "info": {
     "title": "kodit API",
     "description": "\nThis is the REST API for the Kodit server. Please refer to the\n[Kodit documentation](https://docs.helix.ml/kodit/) for more information.\n    ",
-    "version": "0.5.1"
+    "version": "0.5.3"
   },
   "paths": {
     "/healthz": {

{kodit-0.5.2 → kodit-0.5.4}/docs/reference/configuration/index.md

@@ -24,6 +24,9 @@ Global context for the kodit project. Provides a shared state for the app.
 | `EMBEDDING_ENDPOINT_NUM_PARALLEL_TASKS` | int | `10` | Number of parallel tasks to use for the endpoint |
 | `EMBEDDING_ENDPOINT_SOCKET_PATH` | `str | None` | `None` | Unix socket path for local communication (e.g., /tmp/openai.sock) |
 | `EMBEDDING_ENDPOINT_TIMEOUT` | float | `60` | Request timeout in seconds |
+| `EMBEDDING_ENDPOINT_MAX_RETRIES` | int | `5` | Maximum number of retries for the endpoint |
+| `EMBEDDING_ENDPOINT_INITIAL_DELAY` | float | `2.0` | Initial delay in seconds for the endpoint |
+| `EMBEDDING_ENDPOINT_BACKOFF_FACTOR` | float | `2.0` | Backoff factor for the endpoint |
 | `EMBEDDING_ENDPOINT_EXTRA_PARAMS` | `dict | None` | `None` | Extra provider-specific non-secret parameters for LiteLLM |
 | `EMBEDDING_ENDPOINT_MAX_TOKENS` | int | `8000` | Conservative token limit for the embedding model |
 | `ENRICHMENT_ENDPOINT` | `Endpoint | None` | `None` | Endpoint to use for enrichment. |
@@ -33,6 +36,9 @@ Global context for the kodit project. Provides a shared state for the app.
 | `ENRICHMENT_ENDPOINT_NUM_PARALLEL_TASKS` | int | `10` | Number of parallel tasks to use for the endpoint |
 | `ENRICHMENT_ENDPOINT_SOCKET_PATH` | `str | None` | `None` | Unix socket path for local communication (e.g., /tmp/openai.sock) |
 | `ENRICHMENT_ENDPOINT_TIMEOUT` | float | `60` | Request timeout in seconds |
+| `ENRICHMENT_ENDPOINT_MAX_RETRIES` | int | `5` | Maximum number of retries for the endpoint |
+| `ENRICHMENT_ENDPOINT_INITIAL_DELAY` | float | `2.0` | Initial delay in seconds for the endpoint |
+| `ENRICHMENT_ENDPOINT_BACKOFF_FACTOR` | float | `2.0` | Backoff factor for the endpoint |
 | `ENRICHMENT_ENDPOINT_EXTRA_PARAMS` | `dict | None` | `None` | Extra provider-specific non-secret parameters for LiteLLM |
 | `ENRICHMENT_ENDPOINT_MAX_TOKENS` | int | `8000` | Conservative token limit for the embedding model |
 | `DEFAULT_SEARCH` | Search | `provider='sqlite'` |  |

{kodit-0.5.2 → kodit-0.5.4}/src/kodit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.5.2'
-__version_tuple__ = version_tuple = (0, 5, 2)
+__version__ = version = '0.5.4'
+__version_tuple__ = version_tuple = (0, 5, 4)
 
 __commit_id__ = commit_id = None

{kodit-0.5.2 → kodit-0.5.4}/src/kodit/application/factories/server_factory.py

@@ -1,7 +1,6 @@
 """Create a big object that contains all the application services."""
 
 from collections.abc import Callable
-from typing import TYPE_CHECKING
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -86,9 +85,6 @@ from kodit.infrastructure.sqlalchemy.task_status_repository import (
 )
 from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
 
-if TYPE_CHECKING:
-    from kodit.domain.services.enrichment_service import EnrichmentDomainService
-
 
 class ServerFactory:
     """Factory for creating server application services."""
@@ -109,7 +105,6 @@ class ServerFactory:
         self._commit_indexing_application_service: (
             CommitIndexingApplicationService | None
         ) = None
-        self._enrichment_service: EnrichmentDomainService | None = None
         self._enricher_service: Enricher | None = None
         self._task_status_repository: TaskStatusRepository | None = None
         self._operation: ProgressTracker | None = None

{kodit-0.5.2 → kodit-0.5.4}/src/kodit/config.py

@@ -70,6 +70,18 @@ class Endpoint(BaseModel):
         default=60,
         description="Request timeout in seconds",
     )
+    max_retries: int = Field(
+        default=5,
+        description="Maximum number of retries for the endpoint",
+    )
+    initial_delay: float = Field(
+        default=2.0,
+        description="Initial delay in seconds for the endpoint",
+    )
+    backoff_factor: float = Field(
+        default=2.0,
+        description="Backoff factor for the endpoint",
+    )
     extra_params: dict[str, Any] | None = Field(
         default=None,
         description="Extra provider-specific non-secret parameters for LiteLLM",

kodit-0.5.4/src/kodit/infrastructure/embedding/embedding_providers/litellm_embedding_provider.py

@@ -0,0 +1,107 @@
+"""LiteLLM embedding provider implementation."""
+
+from collections.abc import AsyncGenerator
+
+import structlog
+import tiktoken
+
+from kodit.config import Endpoint
+from kodit.domain.services.embedding_service import EmbeddingProvider
+from kodit.domain.value_objects import EmbeddingRequest, EmbeddingResponse
+from kodit.infrastructure.embedding.embedding_providers.batching import (
+    split_sub_batches,
+)
+from kodit.infrastructure.providers.async_batch_processor import (
+    process_items_concurrently,
+)
+from kodit.infrastructure.providers.litellm_provider import LiteLLMProvider
+
+
+class LiteLLMEmbeddingProvider(EmbeddingProvider):
+    """LiteLLM embedding provider that supports 100+ providers."""
+
+    def __init__(
+        self,
+        endpoint: Endpoint,
+    ) -> None:
+        """Initialize the LiteLLM embedding provider.
+
+        Args:
+            endpoint: The endpoint configuration containing all settings.
+
+        """
+        self.endpoint = endpoint
+        self.log = structlog.get_logger(__name__)
+        self._encoding: tiktoken.Encoding | None = None
+        self.provider: LiteLLMProvider = LiteLLMProvider(self.endpoint)
+
+    def _split_sub_batches(
+        self, encoding: tiktoken.Encoding, data: list[EmbeddingRequest]
+    ) -> list[list[EmbeddingRequest]]:
+        """Proxy to the shared batching utility (kept for backward-compat)."""
+        return split_sub_batches(
+            encoding,
+            data,
+            max_tokens=self.endpoint.max_tokens,
+            batch_size=self.endpoint.num_parallel_tasks,
+        )
+
+    async def embed(
+        self, data: list[EmbeddingRequest]
+    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
+        """Embed a list of strings using LiteLLM."""
+        if not data:
+            yield []
+            return
+
+        # Split into batches
+        encoding = self._get_encoding()
+        batched_data = self._split_sub_batches(encoding, data)
+
+        async def _process_batch(
+            batch: list[EmbeddingRequest],
+        ) -> list[EmbeddingResponse]:
+            texts = [item.text for item in batch]
+            response = await self.provider.embedding(texts)
+            embeddings_data = response.get("data", [])
+
+            # Handle mismatch between batch size and response size
+            if len(embeddings_data) != len(batch):
+                preview_response = embeddings_data[:3] if embeddings_data else None
+                self.log.error(
+                    "Embedding response size mismatch",
+                    batch_size=len(batch),
+                    response_size=len(embeddings_data),
+                    texts_preview=[t[:50] for t in texts[:3]],
+                    response_preview=preview_response,
+                )
+                raise ValueError(
+                    f"Expected {len(batch)} embeddings, got {len(embeddings_data)}"
+                )
+
+            return [
+                EmbeddingResponse(
+                    snippet_id=item.snippet_id,
+                    embedding=emb_data.get("embedding", []),
+                )
+                for item, emb_data in zip(batch, embeddings_data, strict=True)
+            ]
+
+        async for result in process_items_concurrently(
+            batched_data,
+            _process_batch,
+            self.endpoint.num_parallel_tasks,
+        ):
+            yield result
+
+    async def close(self) -> None:
+        """Close the provider."""
+        await self.provider.close()
+
+    def _get_encoding(self) -> tiktoken.Encoding:
+        """Return (and cache) the tiktoken encoding for the chosen model."""
+        if self._encoding is None:
+            self._encoding = tiktoken.get_encoding(
+                "o200k_base"
+            )  # Reasonable default for most models, but might not be perfect.
+        return self._encoding

kodit-0.5.4/src/kodit/infrastructure/enricher/litellm_enricher.py

@@ -0,0 +1,83 @@
+"""LiteLLM enricher implementation."""
+
+from collections.abc import AsyncGenerator
+
+import structlog
+
+from kodit.config import Endpoint
+from kodit.domain.enrichments.enricher import Enricher
+from kodit.domain.enrichments.request import EnrichmentRequest
+from kodit.domain.enrichments.response import EnrichmentResponse
+from kodit.infrastructure.enricher.utils import clean_thinking_tags
+from kodit.infrastructure.providers.async_batch_processor import (
+    process_items_concurrently,
+)
+from kodit.infrastructure.providers.litellm_provider import LiteLLMProvider
+
+
+class LiteLLMEnricher(Enricher):
+    """LiteLLM enricher that supports 100+ providers."""
+
+    def __init__(
+        self,
+        endpoint: Endpoint,
+    ) -> None:
+        """Initialize the LiteLLM enricher.
+
+        Args:
+            endpoint: The endpoint configuration containing all settings.
+
+        """
+        self.log = structlog.get_logger(__name__)
+        self.provider: LiteLLMProvider = LiteLLMProvider(endpoint)
+        self.endpoint = endpoint
+
+    async def enrich(
+        self, requests: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
+        """Enrich a list of requests using LiteLLM.
+
+        Args:
+            requests: List of generic enrichment requests.
+
+        Yields:
+            Generic enrichment responses as they are processed.
+
+        """
+        if not requests:
+            self.log.warning("No requests for enrichment")
+            return
+
+        async def process_request(
+            request: EnrichmentRequest,
+        ) -> EnrichmentResponse:
+            if not request.text:
+                return EnrichmentResponse(
+                    id=request.id,
+                    text="",
+                )
+            messages = [
+                {
+                    "role": "system",
+                    "content": request.system_prompt,
+                },
+                {"role": "user", "content": request.text},
+            ]
+            response = await self.provider.chat_completion(messages)
+            content = (
+                response.get("choices", [{}])[0].get("message", {}).get("content", "")
+            )
+            cleaned_content = clean_thinking_tags(content or "")
+            return EnrichmentResponse(
+                id=request.id,
+                text=cleaned_content,
+            )
+
+        async for result in process_items_concurrently(
+            requests, process_request, self.endpoint.num_parallel_tasks
+        ):
+            yield result
+
+    async def close(self) -> None:
+        """Close the enricher and cleanup HTTPX client if using Unix sockets."""
+        await self.provider.close()

kodit-0.5.4/src/kodit/infrastructure/providers/__init__.py

@@ -0,0 +1 @@
+"""Provider utilities for LiteLLM and async batch processing."""

kodit-0.5.4/src/kodit/infrastructure/providers/async_batch_processor.py

@@ -0,0 +1,51 @@
+"""Generic async batch processor with semaphore-controlled concurrency."""
+
+import asyncio
+from collections.abc import AsyncGenerator, Awaitable, Callable
+from typing import TypeVar
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+async def process_items_concurrently(
+    items: list[T],
+    process_fn: Callable[[T], Awaitable[R]],
+    max_parallel_tasks: int,
+) -> AsyncGenerator[R, None]:
+    """Process items concurrently with semaphore-controlled concurrency.
+
+    Args:
+        items: List of items to process.
+        process_fn: Async function to process each item.
+        max_parallel_tasks: Maximum number of concurrent tasks.
+
+    Yields:
+        Results as they are completed (not necessarily in order).
+
+    """
+    if not items:
+        return
+
+    sem = asyncio.Semaphore(max_parallel_tasks)
+
+    async def _process_with_semaphore(item: T) -> R:
+        async with sem:
+            return await process_fn(item)
+
+    tasks: list[asyncio.Task[R]] = [
+        asyncio.create_task(_process_with_semaphore(item)) for item in items
+    ]
+
+    try:
+        for task in asyncio.as_completed(tasks):
+            yield await task
+    finally:
+        # Cancel any remaining tasks when generator exits
+        # (due to exception, Ctrl+C, or early consumer termination)
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+
+        # Wait for all tasks to finish cancelling
+        await asyncio.gather(*tasks, return_exceptions=True)

kodit-0.5.4/src/kodit/infrastructure/providers/litellm_provider.py

@@ -0,0 +1,154 @@
+"""LiteLLM provider implementation."""
+
+import asyncio
+import functools
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+import httpx
+import litellm
+import structlog
+from litellm import acompletion, aembedding
+
+from kodit.config import Endpoint
+
+ProviderMaxRetriesError = Exception("LiteLLM API error: Max retries exceeded")
+
+RETRYABLE_ERRORS = (
+    litellm.exceptions.Timeout,
+    litellm.exceptions.RateLimitError,
+    litellm.exceptions.InternalServerError,
+    litellm.exceptions.ServiceUnavailableError,
+    litellm.exceptions.APIConnectionError,
+    litellm.exceptions.MidStreamFallbackError,
+)
+
+
+def litellm_retry(
+    func: Callable[..., Coroutine[Any, Any, Any]],
+) -> Callable[..., Coroutine[Any, Any, Any]]:
+    """Retry decorator for LiteLLM API calls with exponential backoff.
+
+    Extracts retry configuration from the endpoint attribute of the first
+    argument (self) if it's a LiteLLMProvider instance.
+    """
+
+    @functools.wraps(func)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Extract endpoint configuration from self if available
+        endpoint = None
+        if args and hasattr(args[0], "endpoint"):
+            endpoint = args[0].endpoint
+
+        # Use endpoint configuration or fall back to defaults
+        max_retries = endpoint.max_retries if endpoint else 5
+        initial_delay = endpoint.initial_delay if endpoint else 2.0
+        backoff_factor = endpoint.backoff_factor if endpoint else 2.0
+
+        retries = max_retries
+        delay = initial_delay
+        log: structlog.stdlib.BoundLogger = structlog.get_logger(__name__)
+
+        while True:
+            try:
+                return await func(*args, **kwargs)
+            except (asyncio.CancelledError, KeyboardInterrupt):
+                raise
+            except Exception as e:
+                if isinstance(e, RETRYABLE_ERRORS) and retries > 0:
+                    log.warning(
+                        "LiteLLM API error: Retrying",
+                        error=e,
+                        retries=retries,
+                        backoff=delay,
+                    )
+                    try:
+                        await asyncio.sleep(delay)
+                    except (asyncio.CancelledError, KeyboardInterrupt):
+                        # Cancellation during sleep should stop retries immediately
+                        log.info("Retry cancelled during backoff")
+                        raise
+                    retries -= 1
+                    delay *= backoff_factor
+                    continue
+
+                exception_info = {
+                    attr: getattr(e, attr)
+                    for attr in dir(e)
+                    if not attr.startswith("_")
+                }
+                log.exception(
+                    "LiteLLM API error, check provider logs for details",
+                    error=e,
+                    exception_info=exception_info,
+                    retries=retries,
+                    backoff=delay,
+                )
+                raise
+
+    return wrapper
+
+
+class LiteLLMProvider:
+    """LiteLLM provider that supports 100+ providers."""
+
+    def __init__(self, endpoint: Endpoint) -> None:
+        """Initialize the LiteLLM provider."""
+        self.endpoint = endpoint
+        self._setup_litellm_client()
+
+    def _setup_litellm_client(self) -> None:
+        """Set up LiteLLM with custom HTTPX client for Unix socket support."""
+        if self.endpoint.socket_path:
+            # Create HTTPX client with Unix socket transport
+            transport = httpx.AsyncHTTPTransport(uds=self.endpoint.socket_path)
+            unix_client = httpx.AsyncClient(
+                transport=transport,
+                base_url="http://localhost",  # Base URL for Unix socket
+                timeout=self.endpoint.timeout,
+            )
+            # Set as LiteLLM's async client session
+            litellm.aclient_session = unix_client
+
+    def _populate_base_kwargs(self) -> dict[str, Any]:
+        """Populate base kwargs common to all API calls."""
+        kwargs = {
+            "model": self.endpoint.model,
+            "timeout": self.endpoint.timeout,
+        }
+
+        if self.endpoint.api_key:
+            kwargs["api_key"] = self.endpoint.api_key
+
+        if self.endpoint.base_url:
+            kwargs["api_base"] = self.endpoint.base_url
+
+        kwargs.update(self.endpoint.extra_params or {})
+
+        return kwargs
+
+    @litellm_retry
+    async def chat_completion(self, messages: list[dict[str, str]]) -> Any:
+        """Call the chat completion API using LiteLLM."""
+        kwargs = self._populate_base_kwargs()
+        kwargs["messages"] = messages
+        response = await acompletion(max_retries=0, **kwargs)
+        return response.model_dump()
+
+    @litellm_retry
+    async def embedding(self, texts: list[str]) -> Any:
+        """Call the embedding API using LiteLLM."""
+        kwargs = self._populate_base_kwargs()
+        kwargs["input"] = texts
+        response = await aembedding(max_retries=0, **kwargs)
+        return response.model_dump()
+
+    async def close(self) -> None:
+        """Close the provider and cleanup HTTPX client if using Unix sockets."""
+        if (
+            self.endpoint.socket_path
+            and hasattr(litellm, "aclient_session")
+            and litellm.aclient_session
+        ):
+            await litellm.aclient_session.aclose()
+            litellm.aclient_session = None

{kodit-0.5.2 → kodit-0.5.4}/src/kodit/log.py

@@ -1,6 +1,7 @@
 """Logging configuration for kodit."""
 
 import logging
+import os
 import platform
 import re
 import shutil
@@ -11,6 +12,9 @@ from functools import lru_cache
 from pathlib import Path
 from typing import Any
 
+# Set litellm logging level BEFORE import to prevent broken logging objects
+os.environ["LITELLM_LOG"] = "ERROR"
+
 import litellm
 import rudderstack.analytics as rudder_analytics  # type: ignore[import-untyped]
 import structlog
@@ -108,9 +112,14 @@ def configure_logging(app_context: AppContext) -> None:
         else:
             logging.getLogger(_log).disabled = True
 
-    # More litellm logging cruft
+    # Disable litellm's internal debug logging
     litellm.suppress_debug_info = True
 
+    # Monkey-patch litellm's Logging class to add missing debug method
+    # This prevents AttributeError when litellm tries to call logging_obj.debug()
+    if not hasattr(litellm.Logging, "debug"):
+        litellm.Logging.debug = lambda _self, *_args, **_kwargs: None  # type: ignore[attr-defined]
+
     # Configure SQLAlchemy loggers to use our structlog setup
     for _log in ["sqlalchemy.engine", "alembic"]:
         engine_logger = logging.getLogger(_log)