kodit 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

kodit/application/services/code_indexing_application_service.py

@@ -4,10 +4,12 @@ from dataclasses import replace
 from datetime import UTC, datetime
 
 import structlog
-from sqlalchemy.ext.asyncio import AsyncSession
 
+from kodit.application.services.reporting import (
+    ProgressTracker,
+    TaskOperation,
+)
 from kodit.domain.entities import Index, Snippet
-from kodit.domain.interfaces import ProgressCallback
 from kodit.domain.protocols import IndexRepository
 from kodit.domain.services.bm25_service import BM25DomainService
 from kodit.domain.services.embedding_service import EmbeddingDomainService
@@ -23,9 +25,9 @@ from kodit.domain.value_objects import (
     SearchRequest,
     SearchResult,
     SnippetSearchFilters,
+    TrackableType,
 )
 from kodit.log import log_event
-from kodit.reporting import Reporter
 
 
 class CodeIndexingApplicationService:
@@ -40,7 +42,7 @@ class CodeIndexingApplicationService:
         code_search_service: EmbeddingDomainService,
         text_search_service: EmbeddingDomainService,
         enrichment_service: EnrichmentDomainService,
-        session: AsyncSession,
+        operation: ProgressTracker,
     ) -> None:
         """Initialize the code indexing application service."""
         self.index_domain_service = indexing_domain_service
@@ -50,7 +52,7 @@ class CodeIndexingApplicationService:
         self.code_search_service = code_search_service
         self.text_search_service = text_search_service
         self.enrichment_service = enrichment_service
-        self.session = session
+        self.operation = operation
         self.log = structlog.get_logger(__name__)
 
     async def does_index_exist(self, uri: str) -> bool:
@@ -60,107 +62,131 @@ class CodeIndexingApplicationService:
         existing_index = await self.index_repository.get_by_uri(sanitized_uri)
         return existing_index is not None
 
-    async def create_index_from_uri(
-        self, uri: str, progress_callback: ProgressCallback | None = None
-    ) -> Index:
+    async def create_index_from_uri(self, uri: str) -> Index:
         """Create a new index for a source."""
         log_event("kodit.index.create")
+        async with self.operation.create_child(TaskOperation.CREATE_INDEX) as operation:
+            # Check if index already exists
+            sanitized_uri, _ = self.index_domain_service.sanitize_uri(uri)
+            self.log.info("Creating index from URI", uri=str(sanitized_uri))
+            existing_index = await self.index_repository.get_by_uri(sanitized_uri)
+            if existing_index:
+                self.log.debug(
+                    "Index already exists",
+                    uri=str(sanitized_uri),
+                    index_id=existing_index.id,
+                )
+                return existing_index
 
-        # Check if index already exists
-        sanitized_uri, _ = self.index_domain_service.sanitize_uri(uri)
-        existing_index = await self.index_repository.get_by_uri(sanitized_uri)
-        if existing_index:
-            self.log.debug(
-                "Index already exists",
-                uri=str(sanitized_uri),
-                index_id=existing_index.id,
-            )
-            return existing_index
-
-        # Only prepare working copy if we need to create a new index
-        working_copy = await self.index_domain_service.prepare_index(
-            uri, progress_callback
-        )
+            # Only prepare working copy if we need to create a new index
+            self.log.info("Preparing working copy", uri=str(sanitized_uri))
+            working_copy = await self.index_domain_service.prepare_index(uri, operation)
 
-        # Create new index
-        index = await self.index_repository.create(sanitized_uri, working_copy)
-        await self.session.commit()
-        return index
+            # Create new index
+            self.log.info("Creating index", uri=str(sanitized_uri))
+            return await self.index_repository.create(sanitized_uri, working_copy)
 
-    async def run_index(
-        self, index: Index, progress_callback: ProgressCallback | None = None
-    ) -> None:
+    async def run_index(self, index: Index) -> None:
         """Run the complete indexing process for a specific index."""
-        log_event("kodit.index.run")
-
-        if not index or not index.id:
-            msg = f"Index has no ID: {index}"
-            raise ValueError(msg)
-
-        # Refresh working copy
-        index.source.working_copy = (
-            await self.index_domain_service.refresh_working_copy(
-                index.source.working_copy
-            )
-        )
-        if len(index.source.working_copy.changed_files()) == 0:
-            self.log.info("No new changes to index", index_id=index.id)
-            return
-
-        # Delete the old snippets from the files that have changed
-        await self.index_repository.delete_snippets_by_file_ids(
-            [file.id for file in index.source.working_copy.changed_files() if file.id]
-        )
-
-        # Extract and create snippets (domain service handles progress)
-        self.log.info("Creating snippets for files", index_id=index.id)
-        index = await self.index_domain_service.extract_snippets_from_index(
-            index=index, progress_callback=progress_callback
-        )
-
-        await self.index_repository.update(index)
-        await self.session.flush()
-
-        # Refresh index to get snippets with IDs, required as a ref for subsequent steps
-        flushed_index = await self.index_repository.get(index.id)
-        if not flushed_index:
-            msg = f"Index {index.id} not found after snippet extraction"
-            raise ValueError(msg)
-        index = flushed_index
-        if len(index.snippets) == 0:
-            self.log.info("No snippets to index after extraction", index_id=index.id)
-            return
-
-        # Create BM25 index
-        self.log.info("Creating keyword index")
-        await self._create_bm25_index(index.snippets, progress_callback)
-
-        # Create code embeddings
-        self.log.info("Creating semantic code index")
-        await self._create_code_embeddings(index.snippets, progress_callback)
-
-        # Enrich snippets
-        self.log.info("Enriching snippets", num_snippets=len(index.snippets))
-        enriched_snippets = await self.index_domain_service.enrich_snippets_in_index(
-            snippets=index.snippets, progress_callback=progress_callback
-        )
-        # Update snippets in repository
-        await self.index_repository.update_snippets(index.id, enriched_snippets)
-
-        # Create text embeddings (on enriched content)
-        self.log.info("Creating semantic text index")
-        await self._create_text_embeddings(enriched_snippets, progress_callback)
-
-        # Update index timestamp
-        await self.index_repository.update_index_timestamp(index.id)
-
-        # Now that all file dependencies have been captured, enact the file processing
-        # statuses
-        index.source.working_copy.clear_file_processing_statuses()
-        await self.index_repository.update(index)
+        # Create a new operation
+        async with self.operation.create_child(
+            TaskOperation.RUN_INDEX,
+            trackable_type=TrackableType.INDEX,
+            trackable_id=index.id,
+        ) as operation:
+            # TODO(philwinder): Move this into a reporter # noqa: TD003, FIX002
+            log_event("kodit.index.run")
+
+            if not index or not index.id:
+                msg = f"Index has no ID: {index}"
+                raise ValueError(msg)
+
+            # Refresh working copy
+            async with operation.create_child(
+                TaskOperation.REFRESH_WORKING_COPY
+            ) as step:
+                index.source.working_copy = (
+                    await self.index_domain_service.refresh_working_copy(
+                        index.source.working_copy, step
+                    )
+                )
+                if len(index.source.working_copy.changed_files()) == 0:
+                    self.log.info("No new changes to index", index_id=index.id)
+                    await step.skip("No new changes to index")
+                    return
+
+            # Delete the old snippets from the files that have changed
+            async with operation.create_child(
+                TaskOperation.DELETE_OLD_SNIPPETS
+            ) as step:
+                await self.index_repository.delete_snippets_by_file_ids(
+                    [
+                        file.id
+                        for file in index.source.working_copy.changed_files()
+                        if file.id
+                    ]
+                )
 
-        # Single transaction commit for the entire operation
-        await self.session.commit()
+            # Extract and create snippets (domain service handles progress)
+            async with operation.create_child(TaskOperation.EXTRACT_SNIPPETS) as step:
+                index = await self.index_domain_service.extract_snippets_from_index(
+                    index=index, step=step
+                )
+                await self.index_repository.update(index)
+
+                # Refresh index to get snippets with IDs, required for subsequent steps
+                flushed_index = await self.index_repository.get(index.id)
+                if not flushed_index:
+                    msg = f"Index {index.id} not found after snippet extraction"
+                    raise ValueError(msg)
+                index = flushed_index
+                if len(index.snippets) == 0:
+                    self.log.info(
+                        "No snippets to index after extraction", index_id=index.id
+                    )
+                    await step.skip("No snippets to index after extraction")
+                    return
+
+            # Create BM25 index
+            self.log.info("Creating keyword index")
+            async with operation.create_child(TaskOperation.CREATE_BM25_INDEX) as step:
+                await self._create_bm25_index(index.snippets)
+
+            # Create code embeddings
+            async with operation.create_child(
+                TaskOperation.CREATE_CODE_EMBEDDINGS
+            ) as step:
+                await self._create_code_embeddings(index.snippets, step)
+
+            # Enrich snippets
+            async with operation.create_child(TaskOperation.ENRICH_SNIPPETS) as step:
+                enriched_snippets = (
+                    await self.index_domain_service.enrich_snippets_in_index(
+                        snippets=index.snippets,
+                        reporting_step=step,
+                    )
+                )
+                # Update snippets in repository
+                await self.index_repository.update_snippets(index.id, enriched_snippets)
+
+            # Create text embeddings (on enriched content)
+            async with operation.create_child(
+                TaskOperation.CREATE_TEXT_EMBEDDINGS
+            ) as step:
+                await self._create_text_embeddings(enriched_snippets, step)
+
+            # Update index timestamp
+            async with operation.create_child(
+                TaskOperation.UPDATE_INDEX_TIMESTAMP
+            ) as step:
+                await self.index_repository.update_index_timestamp(index.id)
+
+            # After indexing, clear the file processing statuses
+            async with operation.create_child(
+                TaskOperation.CLEAR_FILE_PROCESSING_STATUSES
+            ) as step:
+                index.source.working_copy.clear_file_processing_statuses()
+                await self.index_repository.update(index)
 
     async def search(self, request: MultiSearchRequest) -> list[MultiSearchResult]:
         """Search for relevant snippets across all indexes."""
@@ -312,15 +338,7 @@ class CodeIndexingApplicationService:
         ]
 
     # FUTURE: BM25 index enriched content too
-    async def _create_bm25_index(
-        self, snippets: list[Snippet], progress_callback: ProgressCallback | None = None
-    ) -> None:
-        reporter = Reporter(self.log, progress_callback)
-        await reporter.start("bm25_index", len(snippets), "Creating keyword index...")
-
-        for _snippet in snippets:
-            pass
-
+    async def _create_bm25_index(self, snippets: list[Snippet]) -> None:
         await self.bm25_service.index_documents(
             IndexRequest(
                 documents=[
@@ -331,16 +349,10 @@ class CodeIndexingApplicationService:
             )
         )
 
-        await reporter.done("bm25_index", "Keyword index created")
-
     async def _create_code_embeddings(
-        self, snippets: list[Snippet], progress_callback: ProgressCallback | None = None
+        self, snippets: list[Snippet], reporting_step: ProgressTracker
     ) -> None:
-        reporter = Reporter(self.log, progress_callback)
-        await reporter.start(
-            "code_embeddings", len(snippets), "Creating code embeddings..."
-        )
-
+        await reporting_step.set_total(len(snippets))
         processed = 0
         async for result in self.code_search_service.index_documents(
             IndexRequest(
@@ -352,23 +364,13 @@ class CodeIndexingApplicationService:
             )
         ):
             processed += len(result)
-            await reporter.step(
-                "code_embeddings",
-                processed,
-                len(snippets),
-                "Creating code embeddings...",
+            await reporting_step.set_current(
+                processed, f"Creating code embeddings for {processed} snippets"
             )
 
-        await reporter.done("code_embeddings")
-
     async def _create_text_embeddings(
-        self, snippets: list[Snippet], progress_callback: ProgressCallback | None = None
+        self, snippets: list[Snippet], reporting_step: ProgressTracker
     ) -> None:
-        reporter = Reporter(self.log, progress_callback)
-        await reporter.start(
-            "text_embeddings", len(snippets), "Creating text embeddings..."
-        )
-
         # Only create text embeddings for snippets that have summary content
         documents_with_summaries = []
         for snippet in snippets:
@@ -384,23 +386,21 @@ class CodeIndexingApplicationService:
                     continue
 
         if not documents_with_summaries:
-            await reporter.done("text_embeddings", "No summaries to index")
+            await reporting_step.skip(
+                "No snippets with summaries to create text embeddings"
+            )
             return
 
+        await reporting_step.set_total(len(documents_with_summaries))
         processed = 0
         async for result in self.text_search_service.index_documents(
             IndexRequest(documents=documents_with_summaries)
         ):
             processed += len(result)
-            await reporter.step(
-                "text_embeddings",
-                processed,
-                len(snippets),
-                "Creating text embeddings...",
+            await reporting_step.set_current(
+                processed, f"Creating text embeddings for {processed} snippets"
             )
 
-        await reporter.done("text_embeddings")
-
     async def delete_index(self, index: Index) -> None:
         """Delete an index."""
         # Delete the index from the domain
@@ -408,4 +408,3 @@ class CodeIndexingApplicationService:
 
         # Delete index from the database
         await self.index_repository.delete(index)
-        await self.session.commit()

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:
@@ -140,15 +136,15 @@ class IndexingWorkerService:
         # Create a fresh database connection for this thread's event loop
         db = await self.app_context.new_db(run_migrations=True)
         try:
-            async with db.session_factory() as session:
-                service = create_code_indexing_application_service(
-                    app_context=self.app_context,
-                    session=session,
-                )
-                index = await service.index_repository.get(index_id)
-                if not index:
-                    raise ValueError(f"Index not found: {index_id}")
-
-                await service.run_index(index)
+            service = create_code_indexing_application_service(
+                app_context=self.app_context,
+                session_factory=self.session_factory,
+                operation=operation,
+            )
+            index = await service.index_repository.get(index_id)
+            if not index:
+                raise ValueError(f"Index not found: {index_id}")
+
+            await service.run_index(index)
         finally:
             await db.close()

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,14 +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."""
-        repo = SqlAlchemyTaskRepository(self.session)
-        return await repo.get(task_id)
+        return await self.task_repository.get(task_id)

kodit/application/services/reporting.py

@@ -0,0 +1,104 @@
+"""Reporting."""
+
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING
+
+import structlog
+
+from kodit.domain.entities import TaskStatus
+from kodit.domain.value_objects import TaskOperation, TrackableType
+
+if TYPE_CHECKING:
+    from kodit.domain.protocols import ReportingModule
+
+
+class ProgressTracker:
+    """Progress tracker.
+
+    Provides a reactive wrapper around TaskStatus domain entities that automatically
+    propagates state changes to the database and reporting modules. This pattern was
+    chosen over a traditional service-repository approach because:
+    - State changes must trigger immediate side effects (database writes, notifications)
+    - Multiple consumers need real-time updates without polling
+    - The wrapper pattern allows transparent interception of all state mutations
+
+    The tracker monitors all modifications to the underlying TaskStatus and ensures
+    consistency across all downstream systems.
+    """
+
+    def __init__(
+        self,
+        task_status: TaskStatus,
+    ) -> None:
+        """Initialize the progress tracker."""
+        self.task_status = task_status
+        self._log = structlog.get_logger(__name__)
+        self._subscribers: list[ReportingModule] = []
+
+    @staticmethod
+    def create(
+        operation: TaskOperation,
+        parent: "TaskStatus | None" = None,
+        trackable_type: TrackableType | None = None,
+        trackable_id: int | None = None,
+    ) -> "ProgressTracker":
+        """Create a progress tracker."""
+        return ProgressTracker(
+            TaskStatus.create(
+                operation=operation,
+                trackable_type=trackable_type,
+                trackable_id=trackable_id,
+                parent=parent,
+            )
+        )
+
+    @asynccontextmanager
+    async def create_child(
+        self,
+        operation: TaskOperation,
+        trackable_type: TrackableType | None = None,
+        trackable_id: int | None = None,
+    ) -> AsyncGenerator["ProgressTracker", None]:
+        """Create a child step."""
+        c = ProgressTracker.create(
+            operation=operation,
+            parent=self.task_status,
+            trackable_type=trackable_type or self.task_status.trackable_type,
+            trackable_id=trackable_id or self.task_status.trackable_id,
+        )
+        try:
+            for subscriber in self._subscribers:
+                c.subscribe(subscriber)
+
+            await c.notify_subscribers()
+            yield c
+        except Exception as e:  # noqa: BLE001
+            c.task_status.fail(str(e))
+        finally:
+            c.task_status.complete()
+            await c.notify_subscribers()
+
+    async def skip(self, reason: str) -> None:
+        """Skip the step."""
+        self.task_status.skip(reason)
+        await self.notify_subscribers()
+
+    def subscribe(self, subscriber: "ReportingModule") -> None:
+        """Subscribe to the step."""
+        self._subscribers.append(subscriber)
+
+    async def set_total(self, total: int) -> None:
+        """Set the total for the step."""
+        self.task_status.set_total(total)
+        await self.notify_subscribers()
+
+    async def set_current(self, current: int, message: str | None = None) -> None:
+        """Progress the step."""
+        self.task_status.set_current(current, message)
+        await self.notify_subscribers()
+
+    async def notify_subscribers(self) -> None:
+        """Notify the subscribers only if progress has changed."""
+        for subscriber in self._subscribers:
+            await subscriber.on_change(self.task_status)