kodit 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

kodit/application/services/indexing_worker_service.py

@@ -2,7 +2,6 @@
 
 import asyncio
 from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor
 from contextlib import suppress
 from datetime import UTC, datetime
 
@@ -12,10 +11,12 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from kodit.application.factories.code_indexing_factory import (
     create_code_indexing_application_service,
 )
+from kodit.application.factories.reporting_factory import create_noop_operation
+from kodit.application.services.reporting import ProgressTracker
 from kodit.config import AppContext
 from kodit.domain.entities import Task
 from kodit.domain.value_objects import TaskType
-from kodit.infrastructure.sqlalchemy.task_repository import SqlAlchemyTaskRepository
+from kodit.infrastructure.sqlalchemy.task_repository import create_task_repository
 
 
 class IndexingWorkerService:
@@ -35,17 +36,16 @@ class IndexingWorkerService:
         self.session_factory = session_factory
         self._worker_task: asyncio.Task | None = None
         self._shutdown_event = asyncio.Event()
-        self._executor = ThreadPoolExecutor(
-            max_workers=1, thread_name_prefix="indexing-worker"
-        )
+        self.task_repository = create_task_repository(session_factory)
         self.log = structlog.get_logger(__name__)
 
-    async def start(self) -> None:
+    async def start(self, operation: ProgressTracker | None = None) -> None:
         """Start the worker to process the queue."""
+        operation = operation or create_noop_operation()
         self._running = True
 
         # Start single worker task
-        self._worker_task = asyncio.create_task(self._worker_loop())
+        self._worker_task = asyncio.create_task(self._worker_loop(operation))
 
         self.log.info(
             "Indexing worker started",
@@ -57,30 +57,24 @@ class IndexingWorkerService:
         self._shutdown_event.set()
 
         if self._worker_task and not self._worker_task.done():
-            self._worker_task.cancel()
             with suppress(asyncio.CancelledError):
+                self._worker_task.cancel()
                 await self._worker_task
 
-        # Shutdown the thread pool executor
-        self._executor.shutdown(wait=True)
-
         self.log.info("Indexing worker stopped")
 
-    async def _worker_loop(self) -> None:
+    async def _worker_loop(self, operation: ProgressTracker) -> None:
         self.log.debug("Worker loop started")
 
         while not self._shutdown_event.is_set():
             try:
                 async with self.session_factory() as session:
-                    repo = SqlAlchemyTaskRepository(session)
-                    task = await repo.take()
+                    task = await self.task_repository.take()
                     await session.commit()
 
                 # If there's a task, process it in a new thread
                 if task:
-                    await asyncio.get_event_loop().run_in_executor(
-                        self._executor, self._process_task, task
-                    )
+                    await self._process_task(task, operation)
                     continue
 
                 # If no task, sleep for a bit
@@ -96,7 +90,7 @@ class IndexingWorkerService:
 
         self.log.info("Worker loop stopped")
 
-    def _process_task(self, task: Task) -> None:
+    async def _process_task(self, task: Task, operation: ProgressTracker) -> None:
         """Process a single task."""
         self.log.info(
             "Processing task",
@@ -113,7 +107,7 @@ class IndexingWorkerService:
         try:
             # Process based on task type (currently only INDEX_UPDATE is supported)
             if task.type is TaskType.INDEX_UPDATE:
-                loop.run_until_complete(self._process_index_update(task))
+                await self._process_index_update(task, operation)
             else:
                 self.log.warning(
                     "Unknown task type",
@@ -131,7 +125,9 @@ class IndexingWorkerService:
             duration_seconds=duration,
         )
 
-    async def _process_index_update(self, task: Task) -> None:
+    async def _process_index_update(
+        self, task: Task, operation: ProgressTracker
+    ) -> None:
         """Process index update/sync task."""
         index_id = task.payload.get("index_id")
         if not index_id:
@@ -144,6 +140,8 @@ class IndexingWorkerService:
                 service = create_code_indexing_application_service(
                     app_context=self.app_context,
                     session=session,
+                    session_factory=self.session_factory,
+                    operation=operation,
                 )
                 index = await service.index_repository.get(index_id)
                 if not index:

kodit/application/services/queue_service.py

@@ -1,11 +1,15 @@
 """Queue service for managing tasks."""
 
+from collections.abc import Callable
+
 import structlog
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.domain.entities import Task
 from kodit.domain.value_objects import TaskType
-from kodit.infrastructure.sqlalchemy.task_repository import SqlAlchemyTaskRepository
+from kodit.infrastructure.sqlalchemy.task_repository import (
+    create_task_repository,
+)
 
 
 class QueueService:
@@ -17,26 +21,24 @@ class QueueService:
 
     def __init__(
         self,
-        session: AsyncSession,
+        session_factory: Callable[[], AsyncSession],
     ) -> None:
         """Initialize the queue service."""
-        self.session = session
+        self.task_repository = create_task_repository(session_factory=session_factory)
         self.log = structlog.get_logger(__name__)
 
     async def enqueue_task(self, task: Task) -> None:
         """Queue a task in the database."""
-        repo = SqlAlchemyTaskRepository(self.session)
-
         # See if task already exists
-        db_task = await repo.get(task.id)
+        db_task = await self.task_repository.get(task.id)
         if db_task:
             # Task already exists, update priority
             db_task.priority = task.priority
-            await repo.update(db_task)
+            await self.task_repository.update(db_task)
             self.log.info("Task updated", task_id=task.id, task_type=task.type)
         else:
             # Otherwise, add task
-            await repo.add(task)
+            await self.task_repository.add(task)
             self.log.info(
                 "Task queued",
                 task_id=task.id,
@@ -44,9 +46,10 @@ class QueueService:
                 payload=task.payload,
             )
 
-        await self.session.commit()
-
     async def list_tasks(self, task_type: TaskType | None = None) -> list[Task]:
         """List all tasks in the queue."""
-        repo = SqlAlchemyTaskRepository(self.session)
-        return await repo.list(task_type)
+        return await self.task_repository.list(task_type)
+
+    async def get_task(self, task_id: str) -> Task | None:
+        """Get a specific task by ID."""
+        return await self.task_repository.get(task_id)

kodit/application/services/reporting.py

@@ -0,0 +1,86 @@
+"""Reporting."""
+
+from enum import StrEnum
+from types import TracebackType
+from typing import TYPE_CHECKING
+
+import structlog
+
+from kodit.domain.value_objects import Progress, ReportingState
+
+if TYPE_CHECKING:
+    from kodit.domain.protocols import ReportingModule
+
+
+class OperationType(StrEnum):
+    """Operation type."""
+
+    ROOT = "kodit.root"
+    CREATE_INDEX = "kodit.index.create"
+    RUN_INDEX = "kodit.index.run"
+
+
+class ProgressTracker:
+    """Progress tracker."""
+
+    def __init__(self, name: str, parent: "ProgressTracker | None" = None) -> None:
+        """Initialize the progress tracker."""
+        self._parent: ProgressTracker | None = parent
+        self._children: list[ProgressTracker] = []
+        self._log = structlog.get_logger(__name__)
+        self._subscribers: list[ReportingModule] = []
+        self._snapshot: Progress = Progress(name=name, state=ReportingState.IN_PROGRESS)
+
+    def __enter__(self) -> "ProgressTracker":
+        """Enter the operation."""
+        self._notify_subscribers()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        """Exit the operation."""
+        if exc_value:
+            self._snapshot = self._snapshot.with_error(exc_value)
+            self._snapshot = self._snapshot.with_state(
+                ReportingState.FAILED, str(exc_value)
+            )
+
+        if self._snapshot.state == ReportingState.IN_PROGRESS:
+            self._snapshot = self._snapshot.with_progress(100)
+            self._snapshot = self._snapshot.with_state(ReportingState.COMPLETED)
+        self._notify_subscribers()
+
+    def create_child(self, name: str) -> "ProgressTracker":
+        """Create a child step."""
+        s = ProgressTracker(name, self)
+        self._children.append(s)
+        for subscriber in self._subscribers:
+            s.subscribe(subscriber)
+        return s
+
+    def skip(self, reason: str | None = None) -> None:
+        """Skip the step."""
+        self._snapshot = self._snapshot.with_state(ReportingState.SKIPPED, reason or "")
+
+    def subscribe(self, subscriber: "ReportingModule") -> None:
+        """Subscribe to the step."""
+        self._subscribers.append(subscriber)
+
+    def set_total(self, total: int) -> None:
+        """Set the total for the step."""
+        self._snapshot = self._snapshot.with_total(total)
+        self._notify_subscribers()
+
+    def set_current(self, current: int) -> None:
+        """Progress the step."""
+        self._snapshot = self._snapshot.with_progress(current)
+        self._notify_subscribers()
+
+    def _notify_subscribers(self) -> None:
+        """Notify the subscribers."""
+        for subscriber in self._subscribers:
+            subscriber.on_change(self._snapshot)

kodit/application/services/sync_scheduler.py

@@ -12,7 +12,7 @@ from kodit.domain.entities import Task
 from kodit.domain.services.index_query_service import IndexQueryService
 from kodit.domain.value_objects import QueuePriority
 from kodit.infrastructure.indexing.fusion_service import ReciprocalRankFusionService
-from kodit.infrastructure.sqlalchemy.index_repository import SqlAlchemyIndexRepository
+from kodit.infrastructure.sqlalchemy.index_repository import create_index_repository
 
 
 class SyncSchedulerService:
@@ -67,27 +67,28 @@ class SyncSchedulerService:
         """Perform a sync operation on all indexes."""
         self.log.info("Starting sync operation")
 
-        async with self.session_factory() as session:
-            # Create services
-            queue_service = QueueService(session=session)
-            index_query_service = IndexQueryService(
-                index_repository=SqlAlchemyIndexRepository(session=session),
-                fusion_service=ReciprocalRankFusionService(),
-            )
+        # Create services
+        queue_service = QueueService(session_factory=self.session_factory)
+        index_query_service = IndexQueryService(
+            index_repository=create_index_repository(
+                session_factory=self.session_factory
+            ),
+            fusion_service=ReciprocalRankFusionService(),
+        )
 
-            # Get all existing indexes
-            all_indexes = await index_query_service.list_indexes()
+        # Get all existing indexes
+        all_indexes = await index_query_service.list_indexes()
 
-            if not all_indexes:
-                self.log.info("No indexes found to sync")
-                return
+        if not all_indexes:
+            self.log.info("No indexes found to sync")
+            return
 
-            self.log.info("Adding sync tasks to queue", count=len(all_indexes))
+        self.log.info("Adding sync tasks to queue", count=len(all_indexes))
 
-            # Sync each index
-            for index in all_indexes:
-                await queue_service.enqueue_task(
-                    Task.create_index_update_task(index.id, QueuePriority.BACKGROUND)
-                )
+        # Sync each index
+        for index in all_indexes:
+            await queue_service.enqueue_task(
+                Task.create_index_update_task(index.id, QueuePriority.BACKGROUND)
+            )
 
-            self.log.info("Sync operation completed")
+        self.log.info("Sync operation completed")

kodit/cli.py

@@ -11,7 +11,7 @@ import uvicorn
 from pytable_formatter import Cell, Table  # type: ignore[import-untyped]
 
 from kodit.application.factories.code_indexing_factory import (
-    create_code_indexing_application_service,
+    create_cli_code_indexing_application_service,
 )
 from kodit.config import (
     AppContext,
@@ -27,11 +27,7 @@ from kodit.domain.value_objects import (
 )
 from kodit.infrastructure.api.client import IndexClient, SearchClient
 from kodit.infrastructure.indexing.fusion_service import ReciprocalRankFusionService
-from kodit.infrastructure.sqlalchemy.index_repository import SqlAlchemyIndexRepository
-from kodit.infrastructure.ui.progress import (
-    create_lazy_progress_callback,
-    create_multi_stage_progress_callback,
-)
+from kodit.infrastructure.sqlalchemy.index_repository import create_index_repository
 from kodit.log import configure_logging, configure_telemetry, log_event
 from kodit.mcp import create_stdio_mcp_server
 
@@ -119,11 +115,8 @@ async def _handle_sync(
     for index in indexes_to_sync:
         click.echo(f"Syncing: {index.source.working_copy.remote_uri}")
 
-        # Create progress callback for this sync operation
-        progress_callback = create_multi_stage_progress_callback()
-
         try:
-            await service.run_index(index, progress_callback)
+            await service.run_index(index)
             click.echo(f"✓ Sync completed: {index.source.working_copy.remote_uri}")
         except Exception as e:
             log.exception("Sync failed", index_id=index.id, error=e)
@@ -191,12 +184,15 @@ async def _index_local(
     # Get database session
     db = await app_context.get_db()
     async with db.session_factory() as session:
-        service = create_code_indexing_application_service(
+        service = create_cli_code_indexing_application_service(
             app_context=app_context,
             session=session,
+            session_factory=db.session_factory,
         )
         index_query_service = IndexQueryService(
-            index_repository=SqlAlchemyIndexRepository(session=session),
+            index_repository=create_index_repository(
+                session_factory=db.session_factory
+            ),
             fusion_service=ReciprocalRankFusionService(),
         )
 
@@ -223,13 +219,11 @@ async def _index_local(
             log_event("kodit.cli.index.create")
 
             # Create a lazy progress callback that only shows progress when needed
-            progress_callback = create_lazy_progress_callback()
-            index = await service.create_index_from_uri(source, progress_callback)
+            index = await service.create_index_from_uri(source)
 
             # Create a new progress callback for the indexing operations
-            indexing_progress_callback = create_multi_stage_progress_callback()
             try:
-                await service.run_index(index, indexing_progress_callback)
+                await service.run_index(index)
             except EmptySourceError as e:
                 log.exception("Empty source error", error=e)
                 msg = f"""{e}. This could mean:
@@ -326,9 +320,10 @@ async def _search_local(  # noqa: PLR0913
     # Get database session
     db = await app_context.get_db()
     async with db.session_factory() as session:
-        service = create_code_indexing_application_service(
+        service = create_cli_code_indexing_application_service(
             app_context=app_context,
             session=session,
+            session_factory=db.session_factory,
         )
 
         filters = _parse_filters(
@@ -791,9 +786,10 @@ async def snippets(
         log_event("kodit.cli.show.snippets")
         db = await app_context.get_db()
         async with db.session_factory() as session:
-            service = create_code_indexing_application_service(
+            service = create_cli_code_indexing_application_service(
                 app_context=app_context,
                 session=session,
+                session_factory=db.session_factory,
             )
             snippets = await service.list_snippets(
                 file_path=by_path, source_uri=by_source

kodit/config.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+from datetime import timedelta
 from enum import Enum
 from functools import wraps
 from pathlib import Path
@@ -38,17 +39,31 @@ DEFAULT_LOG_FORMAT = LogFormat.PRETTY
 DEFAULT_DISABLE_TELEMETRY = False
 T = TypeVar("T")
 
-EndpointType = Literal["openai"]
+EndpointType = Literal["openai", "litellm"]
+
+
+class ReportingConfig(BaseModel):
+    """Reporting configuration."""
+
+    log_time_interval: timedelta = Field(
+        default=timedelta(seconds=5),
+        description="Time interval to log progress in seconds",
+    )
 
 
 class Endpoint(BaseModel):
     """Endpoint provides configuration for an AI service."""
 
-    type: EndpointType | None = None
     base_url: str | None = None
-    model: str | None = None
+    model: str | None = Field(
+        default=None,
+        description="Model to use for the endpoint in litellm format (e.g. 'openai/text-embedding-3-small')",  # noqa: E501
+    )
     api_key: str | None = None
-    num_parallel_tasks: int | None = None
+    num_parallel_tasks: int = Field(
+        default=10,
+        description="Number of parallel tasks to use for the endpoint",
+    )
     socket_path: str | None = Field(
         default=None,
         description="Unix socket path for local communication (e.g., /tmp/openai.sock)",
@@ -57,6 +72,17 @@ class Endpoint(BaseModel):
         default=None,
         description="Request timeout in seconds (default: 30.0)",
     )
+    extra_params: dict[str, Any] | None = Field(
+        default=None,
+        description="Extra provider-specific non-secret parameters for LiteLLM",
+    )
+    max_tokens: int = Field(
+        default=8000,  # Reasonable default (with headroom) for most models.
+        description="Conservative token limit for the embedding model",
+    )
+
+
+DEFAULT_NUM_PARALLEL_TASKS = 10  # Semaphore limit for concurrent requests
 
 
 class Search(BaseModel):
@@ -114,15 +140,11 @@ class PeriodicSyncConfig(BaseModel):
 class RemoteConfig(BaseModel):
     """Configuration for remote server connection."""
 
-    server_url: str | None = Field(
-        default=None, description="Remote Kodit server URL"
-    )
+    server_url: str | None = Field(default=None, description="Remote Kodit server URL")
     api_key: str | None = Field(default=None, description="API key for authentication")
     timeout: float = Field(default=30.0, description="Request timeout in seconds")
     max_retries: int = Field(default=3, description="Maximum retry attempts")
-    verify_ssl: bool = Field(
-        default=True, description="Verify SSL certificates"
-    )
+    verify_ssl: bool = Field(default=True, description="Verify SSL certificates")
 
 
 class CustomAutoIndexingEnvSource(EnvSettingsSource):
@@ -198,13 +220,6 @@ class AppContext(BaseSettings):
     log_level: str = Field(default=DEFAULT_LOG_LEVEL)
     log_format: LogFormat = Field(default=DEFAULT_LOG_FORMAT)
     disable_telemetry: bool = Field(default=DEFAULT_DISABLE_TELEMETRY)
-    default_endpoint: Endpoint | None = Field(
-        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.",
@@ -229,6 +244,9 @@ class AppContext(BaseSettings):
     remote: RemoteConfig = Field(
         default_factory=RemoteConfig, description="Remote server configuration"
     )
+    reporting: ReportingConfig = Field(
+        default=ReportingConfig(), description="Reporting configuration"
+    )
 
     @field_validator("api_keys", mode="before")
     @classmethod

kodit/database.py

@@ -1,5 +1,6 @@
 """Database configuration for kodit."""
 
+from collections.abc import Callable
 from pathlib import Path
 
 import structlog
@@ -28,7 +29,7 @@ class Database:
         )
 
     @property
-    def session_factory(self) -> async_sessionmaker[AsyncSession]:
+    def session_factory(self) -> Callable[[], AsyncSession]:
         """Get the session factory."""
         return self.db_session_factory
 

kodit/domain/protocols.py

@@ -6,7 +6,7 @@ from typing import Protocol
 from pydantic import AnyUrl
 
 from kodit.domain.entities import Index, Snippet, SnippetWithContext, Task, WorkingCopy
-from kodit.domain.value_objects import MultiSearchRequest, TaskType
+from kodit.domain.value_objects import MultiSearchRequest, Progress, TaskType
 
 
 class TaskRepository(Protocol):
@@ -90,3 +90,11 @@ class IndexRepository(Protocol):
     async def get_snippets_by_ids(self, ids: list[int]) -> list[SnippetWithContext]:
         """Get snippets by their IDs."""
         ...
+
+
+class ReportingModule(Protocol):
+    """Reporting module."""
+
+    def on_change(self, step: Progress) -> None:
+        """On step changed."""
+        ...

kodit/domain/services/bm25_service.py

@@ -31,12 +31,7 @@ class BM25DomainService:
     """Domain service for BM25 operations."""
 
     def __init__(self, repository: BM25Repository) -> None:
-        """Initialize the BM25 domain service.
-
-        Args:
-            repository: The BM25 repository for persistence operations
-
-        """
+        """Initialize the BM25 domain service."""
         self.repository = repository
 
     async def index_documents(self, request: IndexRequest) -> None: