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

kodit/infrastructure/cloning/git/working_copy.py

@@ -1,5 +1,6 @@
 """Working copy provider for git-based sources."""
 
+import asyncio
 import hashlib
 import shutil
 from pathlib import Path
@@ -39,7 +40,7 @@ class GitWorkingCopyProvider:
         clone_path.mkdir(parents=True, exist_ok=True)
 
         step_record = []
-        step.set_total(12)
+        await step.set_total(12)
 
         def _clone_progress_callback(
             a: int, _: str | float | None, __: str | float | None, _d: str
@@ -49,7 +50,13 @@ class GitWorkingCopyProvider:
 
             # Git reports a really weird format. This is a quick hack to get some
             # progress.
-            step.set_current(len(step_record))
+            # Normally this would fail because the loop is already running,
+            # but in this case, this callback is called by some git sub-thread.
+            asyncio.run(
+                step.set_current(
+                    len(step_record), f"Cloning repository ({step_record[-1]})"
+                )
+            )
 
         try:
             self.log.info(

kodit/infrastructure/enrichment/local_enrichment_provider.py

@@ -1,7 +1,9 @@
 """Local enrichment provider implementation."""
 
+import asyncio
 import os
 from collections.abc import AsyncGenerator
+from typing import Any
 
 import structlog
 import tiktoken
@@ -60,23 +62,26 @@ class LocalEnrichmentProvider(EnrichmentProvider):
             self.log.warning("No valid requests for enrichment")
             return
 
-        from transformers.models.auto.modeling_auto import (
-            AutoModelForCausalLM,
-        )
-        from transformers.models.auto.tokenization_auto import AutoTokenizer
-
-        if self.tokenizer is None:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_name, padding_side="left"
-            )
-        if self.model is None:
-            os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                torch_dtype="auto",
-                trust_remote_code=True,
-                device_map="auto",
+        def _init_model() -> None:
+            from transformers.models.auto.modeling_auto import (
+                AutoModelForCausalLM,
             )
+            from transformers.models.auto.tokenization_auto import AutoTokenizer
+
+            if self.tokenizer is None:
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    self.model_name, padding_side="left"
+                )
+            if self.model is None:
+                os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_name,
+                    torch_dtype="auto",
+                    trust_remote_code=True,
+                    device_map="auto",
+                )
+
+        await asyncio.to_thread(_init_model)
 
         # Prepare prompts
         prompts = [
@@ -96,20 +101,26 @@ class LocalEnrichmentProvider(EnrichmentProvider):
         ]
 
         for prompt in prompts:
-            model_inputs = self.tokenizer(  # type: ignore[misc]
-                prompt["text"],
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-            ).to(self.model.device)  # type: ignore[attr-defined]
-            generated_ids = self.model.generate(  # type: ignore[attr-defined]
-                **model_inputs, max_new_tokens=self.context_window
-            )
-            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(  # type: ignore[attr-defined]
-                "\n"
-            )
+
+            def process_prompt(prompt: dict[str, Any]) -> str:
+                model_inputs = self.tokenizer(  # type: ignore[misc]
+                    prompt["text"],
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                ).to(self.model.device)  # type: ignore[attr-defined]
+                generated_ids = self.model.generate(  # type: ignore[attr-defined]
+                    **model_inputs, max_new_tokens=self.context_window
+                )
+                input_ids = model_inputs["input_ids"][0]
+                output_ids = generated_ids[0][len(input_ids) :].tolist()
+                return self.tokenizer.decode(  # type: ignore[attr-defined]
+                    output_ids, skip_special_tokens=True
+                ).strip(  # type: ignore[attr-defined]
+                    "\n"
+                )
+
+            content = await asyncio.to_thread(process_prompt, prompt)
             # Remove thinking tags from the response
             cleaned_content = clean_thinking_tags(content)
             yield EnrichmentResponse(

kodit/infrastructure/mappers/task_status_mapper.py

@@ -0,0 +1,85 @@
+"""Task status mapper."""
+
+from kodit.domain import entities as domain_entities
+from kodit.domain.value_objects import ReportingState, TaskOperation, TrackableType
+from kodit.infrastructure.sqlalchemy import entities as db_entities
+
+
+class TaskStatusMapper:
+    """Mapper for converting between domain TaskStatus and database entities."""
+
+    @staticmethod
+    def from_domain_task_status(
+        task_status: domain_entities.TaskStatus,
+    ) -> db_entities.TaskStatus:
+        """Convert domain TaskStatus to database TaskStatus."""
+        return db_entities.TaskStatus(
+            id=task_status.id,
+            operation=task_status.operation,
+            created_at=task_status.created_at,
+            updated_at=task_status.updated_at,
+            trackable_id=task_status.trackable_id,
+            trackable_type=(
+                task_status.trackable_type.value if task_status.trackable_type else None
+            ),
+            parent=task_status.parent.id if task_status.parent else None,
+            state=(
+                task_status.state.value
+                if isinstance(task_status.state, ReportingState)
+                else task_status.state
+            ),
+            error=task_status.error,
+            total=task_status.total,
+            current=task_status.current,
+            message=task_status.message,
+        )
+
+    @staticmethod
+    def to_domain_task_status(
+        db_status: db_entities.TaskStatus,
+    ) -> domain_entities.TaskStatus:
+        """Convert database TaskStatus to domain TaskStatus."""
+        return domain_entities.TaskStatus(
+            id=db_status.id,
+            operation=TaskOperation(db_status.operation),
+            state=ReportingState(db_status.state),
+            created_at=db_status.created_at,
+            updated_at=db_status.updated_at,
+            trackable_id=db_status.trackable_id,
+            trackable_type=(
+                TrackableType(db_status.trackable_type)
+                if db_status.trackable_type
+                else None
+            ),
+            parent=None,  # Parent relationships need to be reconstructed separately
+            error=db_status.error if db_status.error else None,
+            total=db_status.total,
+            current=db_status.current,
+            message=db_status.message,
+        )
+
+    @staticmethod
+    def to_domain_task_status_with_hierarchy(
+        db_statuses: list[db_entities.TaskStatus],
+    ) -> list[domain_entities.TaskStatus]:
+        """Convert database TaskStatus list to domain with parent-child hierarchy.
+
+        This method performs a two-pass conversion:
+        1. First pass: Convert all DB entities to domain entities
+        2. Second pass: Reconstruct parent-child relationships using ID mapping
+        """
+        # First pass: Convert all database entities to domain entities
+        domain_statuses = [
+            TaskStatusMapper.to_domain_task_status(db_status)
+            for db_status in db_statuses
+        ]
+
+        # Create ID-to-entity mapping for efficient parent lookup
+        id_to_entity = {status.id: status for status in domain_statuses}
+
+        # Second pass: Reconstruct parent-child relationships
+        for db_status, domain_status in zip(db_statuses, domain_statuses, strict=True):
+            if db_status.parent and db_status.parent in id_to_entity:
+                domain_status.parent = id_to_entity[db_status.parent]
+
+        return domain_statuses

kodit/infrastructure/reporting/db_progress.py

@@ -0,0 +1,23 @@
+"""Log progress using structlog."""
+
+import structlog
+
+from kodit.config import ReportingConfig
+from kodit.domain.entities import TaskStatus
+from kodit.domain.protocols import ReportingModule, TaskStatusRepository
+
+
+class DBProgressReportingModule(ReportingModule):
+    """Database progress reporting module."""
+
+    def __init__(
+        self, task_status_repository: TaskStatusRepository, config: ReportingConfig
+    ) -> None:
+        """Initialize the database progress reporting module."""
+        self.task_status_repository = task_status_repository
+        self.config = config
+        self._log = structlog.get_logger(__name__)
+
+    async def on_change(self, progress: TaskStatus) -> None:
+        """On step changed - update task status in database."""
+        await self.task_status_repository.save(progress)

kodit/infrastructure/reporting/log_progress.py

@@ -1,13 +1,13 @@
 """Log progress using structlog."""
 
-import time
 from datetime import UTC, datetime
 
 import structlog
 
 from kodit.config import ReportingConfig
+from kodit.domain.entities import TaskStatus
 from kodit.domain.protocols import ReportingModule
-from kodit.domain.value_objects import Progress, ProgressState, ReportingState
+from kodit.domain.value_objects import ReportingState
 
 
 class LoggingReportingModule(ReportingModule):
@@ -19,47 +19,19 @@ class LoggingReportingModule(ReportingModule):
         self._log = structlog.get_logger(__name__)
         self._last_log_time: datetime = datetime.now(UTC)
 
-    def on_change(self, step: Progress) -> None:
+    async def on_change(self, progress: TaskStatus) -> None:
         """On step changed."""
         current_time = datetime.now(UTC)
         time_since_last_log = current_time - self._last_log_time
+        step = progress
 
         if (
             step.state != ReportingState.IN_PROGRESS
             or time_since_last_log >= self.config.log_time_interval
         ):
             self._log.info(
-                step.name,
+                step.operation,
                 state=step.state,
-                message=step.message,
                 completion_percent=step.completion_percent,
             )
             self._last_log_time = current_time
-
-
-class LogProgress(Progress):
-    """Log progress using structlog with time-based throttling."""
-
-    def __init__(self, config: ReportingConfig | None = None) -> None:
-        """Initialize the log progress."""
-        self.log = structlog.get_logger()
-        self.config = config or ReportingConfig()
-        self.last_log_time: float = 0
-
-    def on_update(self, state: ProgressState) -> None:
-        """Log the progress with time-based throttling."""
-        current_time = time.time()
-        time_since_last_log = current_time - self.last_log_time
-
-        if time_since_last_log >= self.config.log_time_interval.total_seconds():
-            self.log.info(
-                "Progress...",
-                operation=state.operation,
-                percentage=state.percentage,
-                message=state.message,
-            )
-            self.last_log_time = current_time
-
-    def on_complete(self) -> None:
-        """Log the completion."""
-        self.log.info("Completed")

kodit/infrastructure/reporting/tdqm_progress.py

@@ -3,8 +3,9 @@
 from tqdm import tqdm
 
 from kodit.config import ReportingConfig
+from kodit.domain.entities import TaskStatus
 from kodit.domain.protocols import ReportingModule
-from kodit.domain.value_objects import Progress, ProgressState, ReportingState
+from kodit.domain.value_objects import ReportingState
 
 
 class TQDMReportingModule(ReportingModule):
@@ -15,59 +16,23 @@ class TQDMReportingModule(ReportingModule):
         self.config = config
         self.pbar = tqdm()
 
-    def on_change(self, step: Progress) -> None:
+    async def on_change(self, progress: TaskStatus) -> None:
         """On step changed."""
+        step = progress
         if step.state == ReportingState.COMPLETED:
             self.pbar.close()
             return
 
-        self.pbar.set_description(step.message)
+        self.pbar.set_description(step.operation)
         self.pbar.refresh()
         # Update description if message is provided
-        if step.message:
+        if step.error:
             # Fix the event message to a specific size so it's not jumping around
             # If it's too small, add spaces
             # If it's too large, truncate
-            if len(step.message) < 30:
-                self.pbar.set_description(step.message + " " * (30 - len(step.message)))
+            if len(step.error) < 30:
+                self.pbar.set_description(step.error + " " * (30 - len(step.error)))
             else:
-                self.pbar.set_description(step.message[-30:])
+                self.pbar.set_description(step.error[-30:])
         else:
-            self.pbar.set_description(step.name)
-
-
-class TQDMProgress(Progress):
-    """TQDM-based progress callback implementation."""
-
-    def __init__(self, config: ReportingConfig | None = None) -> None:
-        """Initialize with a TQDM progress bar."""
-        self.config = config or ReportingConfig()
-        self.pbar = tqdm()
-
-    def on_update(self, state: ProgressState) -> None:
-        """Update the TQDM progress bar."""
-        # Update total if it changes
-        if state.total != self.pbar.total:
-            self.pbar.total = state.total
-
-        # Update the progress bar
-        self.pbar.n = state.current
-        self.pbar.refresh()
-
-        # Update description if message is provided
-        if state.message:
-            # Fix the event message to a specific size so it's not jumping around
-            # If it's too small, add spaces
-            # If it's too large, truncate
-            if len(state.message) < 30:
-                self.pbar.set_description(
-                    state.message + " " * (30 - len(state.message))
-                )
-            else:
-                self.pbar.set_description(state.message[-30:])
-        else:
-            self.pbar.set_description(state.operation)
-
-    def on_complete(self) -> None:
-        """Complete the progress bar."""
-        self.pbar.close()
+            self.pbar.set_description(step.operation)

kodit/infrastructure/sqlalchemy/entities.py

@@ -262,3 +262,64 @@ class Task(Base, CommonMixin):
         self.type = type
         self.payload = payload
         self.priority = priority
+
+
+class TaskStatus(Base):
+    """Task status model."""
+
+    __tablename__ = "task_status"
+    id: Mapped[str] = mapped_column(
+        String(255), primary_key=True, index=True, nullable=False
+    )
+    created_at: Mapped[datetime] = mapped_column(
+        TZDateTime, nullable=False, default=lambda: datetime.now(UTC)
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        TZDateTime,
+        nullable=False,
+        default=lambda: datetime.now(UTC),
+        onupdate=lambda: datetime.now(UTC),
+    )
+    operation: Mapped[str] = mapped_column(String(255), index=True, nullable=False)
+    trackable_id: Mapped[int | None] = mapped_column(Integer, index=True, nullable=True)
+    trackable_type: Mapped[str | None] = mapped_column(
+        String(255), index=True, nullable=True
+    )
+    parent: Mapped[str | None] = mapped_column(
+        ForeignKey("task_status.id"), index=True, nullable=True
+    )
+    message: Mapped[str] = mapped_column(UnicodeText, default="")
+    state: Mapped[str] = mapped_column(String(255), default="")
+    error: Mapped[str] = mapped_column(UnicodeText, default="")
+    total: Mapped[int] = mapped_column(Integer, default=0)
+    current: Mapped[int] = mapped_column(Integer, default=0)
+
+    def __init__(  # noqa: PLR0913
+        self,
+        id: str,  # noqa: A002
+        operation: str,
+        created_at: datetime,
+        updated_at: datetime,
+        trackable_id: int | None,
+        trackable_type: str | None,
+        parent: str | None,
+        state: str,
+        error: str | None,
+        total: int,
+        current: int,
+        message: str,
+    ) -> None:
+        """Initialize the task status."""
+        super().__init__()
+        self.id = id
+        self.operation = operation
+        self.created_at = created_at
+        self.updated_at = updated_at
+        self.trackable_id = trackable_id
+        self.trackable_type = trackable_type
+        self.parent = parent
+        self.state = state
+        self.error = error or ""
+        self.total = total
+        self.current = current
+        self.message = message or ""

kodit/infrastructure/sqlalchemy/task_status_repository.py

@@ -0,0 +1,79 @@
+"""Task repository for the task queue."""
+
+from collections.abc import Callable
+
+import structlog
+from sqlalchemy import delete, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain import entities as domain_entities
+from kodit.domain.protocols import TaskStatusRepository
+from kodit.infrastructure.mappers.task_status_mapper import TaskStatusMapper
+from kodit.infrastructure.sqlalchemy import entities as db_entities
+from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
+
+
+def create_task_status_repository(
+    session_factory: Callable[[], AsyncSession],
+) -> TaskStatusRepository:
+    """Create an index repository."""
+    uow = SqlAlchemyUnitOfWork(session_factory=session_factory)
+    return SqlAlchemyTaskStatusRepository(uow)
+
+
+class SqlAlchemyTaskStatusRepository(TaskStatusRepository):
+    """Repository for persisting TaskStatus entities."""
+
+    def __init__(self, uow: SqlAlchemyUnitOfWork) -> None:
+        """Initialize the repository."""
+        self.uow = uow
+        self.log = structlog.get_logger(__name__)
+        self.mapper = TaskStatusMapper()
+
+    async def save(self, status: domain_entities.TaskStatus) -> None:
+        """Save a TaskStatus to database."""
+        async with self.uow:
+            # Convert domain entity to database entity
+            db_status = self.mapper.from_domain_task_status(status)
+            stmt = select(db_entities.TaskStatus).where(
+                db_entities.TaskStatus.id == db_status.id,
+            )
+            result = await self.uow.session.execute(stmt)
+            existing = result.scalar_one_or_none()
+
+            if not existing:
+                self.uow.session.add(db_status)
+            else:
+                # Update existing record with new values
+                existing.operation = db_status.operation
+                existing.state = db_status.state
+                existing.error = db_status.error
+                existing.total = db_status.total
+                existing.current = db_status.current
+                existing.updated_at = db_status.updated_at
+                existing.parent = db_status.parent
+                existing.trackable_id = db_status.trackable_id
+                existing.trackable_type = db_status.trackable_type
+
+    async def load_with_hierarchy(
+        self, trackable_type: str, trackable_id: int
+    ) -> list[domain_entities.TaskStatus]:
+        """Load TaskStatus entities with hierarchy from database."""
+        async with self.uow:
+            stmt = select(db_entities.TaskStatus).where(
+                db_entities.TaskStatus.trackable_id == trackable_id,
+                db_entities.TaskStatus.trackable_type == trackable_type,
+            )
+            result = await self.uow.session.execute(stmt)
+            db_statuses = list(result.scalars().all())
+
+            # Use mapper to convert and reconstruct hierarchy
+            return self.mapper.to_domain_task_status_with_hierarchy(db_statuses)
+
+    async def delete(self, status: domain_entities.TaskStatus) -> None:
+        """Delete a TaskStatus."""
+        async with self.uow:
+            stmt = delete(db_entities.TaskStatus).where(
+                db_entities.TaskStatus.id == status.id,
+            )
+            await self.uow.session.execute(stmt)

kodit/mcp.py

@@ -23,6 +23,9 @@ from kodit.domain.value_objects import (
     MultiSearchResult,
     SnippetSearchFilters,
 )
+from kodit.infrastructure.sqlalchemy.task_status_repository import (
+    create_task_status_repository,
+)
 
 # Global database connection for MCP server
 _mcp_db: Database | None = None
@@ -179,9 +182,10 @@ def register_mcp_tools(mcp_server: FastMCP) -> None:
         # Use the unified application service
         service = create_code_indexing_application_service(
             app_context=mcp_context.app_context,
-            session=mcp_context.session,
             session_factory=mcp_context.session_factory,
-            operation=create_server_operation(),
+            operation=create_server_operation(
+                create_task_status_repository(mcp_context.session_factory)
+            ),
         )
 
         log.debug("Searching for snippets")

kodit/migrations/env.py

@@ -1,4 +1,3 @@
-# ruff: noqa: F401
 """Alembic environment file for kodit."""
 
 import asyncio

kodit/migrations/versions/b9cd1c3fd762_add_task_status.py

@@ -0,0 +1,77 @@
+# ruff: noqa
+"""add task status
+
+Revision ID: b9cd1c3fd762
+Revises: 9cf0e87de578
+Create Date: 2025-09-05 13:41:29.645898
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "b9cd1c3fd762"
+down_revision: Union[str, None] = "9cf0e87de578"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "task_status",
+        sa.Column("id", sa.String(length=255), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("operation", sa.String(length=255), nullable=False),
+        sa.Column("trackable_id", sa.Integer(), nullable=True),
+        sa.Column("trackable_type", sa.String(length=255), nullable=True),
+        sa.Column("parent", sa.String(length=255), nullable=True),
+        sa.Column("message", sa.UnicodeText(), nullable=False),
+        sa.Column("state", sa.String(length=255), nullable=False),
+        sa.Column("error", sa.UnicodeText(), nullable=False),
+        sa.Column("total", sa.Integer(), nullable=False),
+        sa.Column("current", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["parent"],
+            ["task_status.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(op.f("ix_task_status_id"), "task_status", ["id"], unique=False)
+    op.create_index(
+        op.f("ix_task_status_operation"), "task_status", ["operation"], unique=False
+    )
+    op.create_index(
+        op.f("ix_task_status_parent"), "task_status", ["parent"], unique=False
+    )
+    op.create_index(
+        op.f("ix_task_status_trackable_id"),
+        "task_status",
+        ["trackable_id"],
+        unique=False,
+    )
+    op.create_index(
+        op.f("ix_task_status_trackable_type"),
+        "task_status",
+        ["trackable_type"],
+        unique=False,
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_task_status_trackable_type"), table_name="task_status")
+    op.drop_index(op.f("ix_task_status_trackable_id"), table_name="task_status")
+    op.drop_index(op.f("ix_task_status_parent"), table_name="task_status")
+    op.drop_index(op.f("ix_task_status_operation"), table_name="task_status")
+    op.drop_index(op.f("ix_task_status_id"), table_name="task_status")
+    op.drop_table("task_status")
+    # ### end Alembic commands ###

{kodit-0.4.2.dist-info → kodit-0.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.4.2
+Version: 0.4.3
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/