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

kodit/infrastructure/git/git_utils.py

@@ -3,6 +3,7 @@
 import tempfile
 
 import git
+import git.cmd
 import structlog
 
 
@@ -19,10 +20,10 @@ def is_valid_clone_target(target: str) -> bool:
     """
     with tempfile.TemporaryDirectory() as temp_dir:
         try:
-            git.Repo.clone_from(target, temp_dir)
+            git.cmd.Git(temp_dir).ls_remote(target)
         except git.GitCommandError as e:
             structlog.get_logger(__name__).warning(
-                "Failed to clone git repository",
+                "Failed to list git repository",
                 target=target,
                 error=e,
             )

kodit/infrastructure/mappers/index_mapper.py

@@ -15,6 +15,7 @@ from kodit.domain.value_objects import (
 from kodit.infrastructure.sqlalchemy import entities as db_entities
 
 
+# TODO(Phil): Make this a pure mapper without any DB access # noqa: TD003, FIX002
 class IndexMapper:
     """Mapper for converting between domain Index aggregate and database entities."""
 

kodit/infrastructure/reporting/__init__.py

@@ -0,0 +1 @@
+"""Reporting infrastructure."""

kodit/infrastructure/reporting/log_progress.py

@@ -0,0 +1,65 @@
+"""Log progress using structlog."""
+
+import time
+from datetime import UTC, datetime
+
+import structlog
+
+from kodit.config import ReportingConfig
+from kodit.domain.protocols import ReportingModule
+from kodit.domain.value_objects import Progress, ProgressState, ReportingState
+
+
+class LoggingReportingModule(ReportingModule):
+    """Logging reporting module."""
+
+    def __init__(self, config: ReportingConfig) -> None:
+        """Initialize the logging reporting module."""
+        self.config = config
+        self._log = structlog.get_logger(__name__)
+        self._last_log_time: datetime = datetime.now(UTC)
+
+    def on_change(self, step: Progress) -> None:
+        """On step changed."""
+        current_time = datetime.now(UTC)
+        time_since_last_log = current_time - self._last_log_time
+
+        if (
+            step.state != ReportingState.IN_PROGRESS
+            or time_since_last_log >= self.config.log_time_interval
+        ):
+            self._log.info(
+                step.name,
+                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

@@ -0,0 +1,73 @@
+"""TQDM progress."""
+
+from tqdm import tqdm
+
+from kodit.config import ReportingConfig
+from kodit.domain.protocols import ReportingModule
+from kodit.domain.value_objects import Progress, ProgressState, ReportingState
+
+
+class TQDMReportingModule(ReportingModule):
+    """TQDM reporting module."""
+
+    def __init__(self, config: ReportingConfig) -> None:
+        """Initialize the TQDM reporting module."""
+        self.config = config
+        self.pbar = tqdm()
+
+    def on_change(self, step: Progress) -> None:
+        """On step changed."""
+        if step.state == ReportingState.COMPLETED:
+            self.pbar.close()
+            return
+
+        self.pbar.set_description(step.message)
+        self.pbar.refresh()
+        # Update description if message is provided
+        if step.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(step.message) < 30:
+                self.pbar.set_description(step.message + " " * (30 - len(step.message)))
+            else:
+                self.pbar.set_description(step.message[-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()

kodit/infrastructure/sqlalchemy/embedding_repository.py

@@ -1,85 +1,64 @@
 """SQLAlchemy implementation of embedding repository."""
 
+from collections.abc import Callable
+
 import numpy as np
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.infrastructure.sqlalchemy.entities import Embedding, EmbeddingType
+from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
 
 
-class SqlAlchemyEmbeddingRepository:
-    """SQLAlchemy implementation of embedding repository."""
-
-    def __init__(self, session: AsyncSession) -> None:
-        """Initialize the SQLAlchemy embedding repository.
-
-        Args:
-            session: The SQLAlchemy async session to use for database operations
-
-        """
-        self.session = session
+def create_embedding_repository(
+    session_factory: Callable[[], AsyncSession],
+) -> "SqlAlchemyEmbeddingRepository":
+    """Create an embedding repository."""
+    uow = SqlAlchemyUnitOfWork(session_factory=session_factory)
+    return SqlAlchemyEmbeddingRepository(uow)
 
-    async def create_embedding(self, embedding: Embedding) -> Embedding:
-        """Create a new embedding record in the database.
 
-        Args:
-            embedding: The Embedding instance to create
+class SqlAlchemyEmbeddingRepository:
+    """SQLAlchemy implementation of embedding repository."""
 
-        Returns:
-            The created Embedding instance
+    def __init__(self, uow: SqlAlchemyUnitOfWork) -> None:
+        """Initialize the SQLAlchemy embedding repository."""
+        self.uow = uow
 
-        """
-        self.session.add(embedding)
-        return embedding
+    async def create_embedding(self, embedding: Embedding) -> None:
+        """Create a new embedding record in the database."""
+        async with self.uow:
+            self.uow.session.add(embedding)
 
     async def get_embedding_by_snippet_id_and_type(
         self, snippet_id: int, embedding_type: EmbeddingType
     ) -> Embedding | None:
-        """Get an embedding by its snippet ID and type.
-
-        Args:
-            snippet_id: The ID of the snippet to get the embedding for
-            embedding_type: The type of embedding to get
-
-        Returns:
-            The Embedding instance if found, None otherwise
-
-        """
-        query = select(Embedding).where(
-            Embedding.snippet_id == snippet_id,
-            Embedding.type == embedding_type,
-        )
-        result = await self.session.execute(query)
-        return result.scalar_one_or_none()
+        """Get an embedding by its snippet ID and type."""
+        async with self.uow:
+            query = select(Embedding).where(
+                Embedding.snippet_id == snippet_id,
+                Embedding.type == embedding_type,
+            )
+            result = await self.uow.session.execute(query)
+            return result.scalar_one_or_none()
 
     async def list_embeddings_by_type(
         self, embedding_type: EmbeddingType
     ) -> list[Embedding]:
-        """List all embeddings of a given type.
-
-        Args:
-            embedding_type: The type of embeddings to list
-
-        Returns:
-            A list of Embedding instances
-
-        """
-        query = select(Embedding).where(Embedding.type == embedding_type)
-        result = await self.session.execute(query)
-        return list(result.scalars())
+        """List all embeddings of a given type."""
+        async with self.uow:
+            query = select(Embedding).where(Embedding.type == embedding_type)
+            result = await self.uow.session.execute(query)
+            return list(result.scalars())
 
     async def delete_embeddings_by_snippet_id(self, snippet_id: int) -> None:
-        """Delete all embeddings for a snippet.
-
-        Args:
-            snippet_id: The ID of the snippet to delete embeddings for
-
-        """
-        query = select(Embedding).where(Embedding.snippet_id == snippet_id)
-        result = await self.session.execute(query)
-        embeddings = result.scalars().all()
-        for embedding in embeddings:
-            await self.session.delete(embedding)
+        """Delete all embeddings for a snippet."""
+        async with self.uow:
+            query = select(Embedding).where(Embedding.snippet_id == snippet_id)
+            result = await self.uow.session.execute(query)
+            embeddings = result.scalars().all()
+            for embedding in embeddings:
+                await self.uow.session.delete(embedding)
 
     async def list_semantic_results(
         self,
@@ -130,17 +109,17 @@ class SqlAlchemyEmbeddingRepository:
             List of (snippet_id, embedding) tuples
 
         """
-        # Only select the fields we need and use a more efficient query
-        query = select(Embedding.snippet_id, Embedding.embedding).where(
-            Embedding.type == embedding_type
-        )
+        async with self.uow:
+            query = select(Embedding.snippet_id, Embedding.embedding).where(
+                Embedding.type == embedding_type
+            )
 
-        # Add snippet_ids filter if provided
-        if snippet_ids is not None:
-            query = query.where(Embedding.snippet_id.in_(snippet_ids))
+            # Add snippet_ids filter if provided
+            if snippet_ids is not None:
+                query = query.where(Embedding.snippet_id.in_(snippet_ids))
 
-        rows = await self.session.execute(query)
-        return [tuple(row) for row in rows.all()]  # Convert Row objects to tuples
+            rows = await self.uow.session.execute(query)
+            return [tuple(row) for row in rows.all()]  # Convert Row objects to tuples
 
     def _prepare_vectors(
         self, embeddings: list[tuple[int, list[float]]], query_embedding: list[float]

kodit/infrastructure/sqlalchemy/entities.py

@@ -2,6 +2,7 @@
 
 from datetime import UTC, datetime
 from enum import Enum
+from typing import Any
 
 from git import Actor
 from sqlalchemy import (
@@ -9,6 +10,7 @@ from sqlalchemy import (
     ForeignKey,
     Integer,
     String,
+    TypeDecorator,
     UnicodeText,
     UniqueConstraint,
 )
@@ -18,6 +20,29 @@ from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 from sqlalchemy.types import JSON
 
 
+# See <https://docs.sqlalchemy.org/en/20/core/custom_types.html#store-timezone-aware-timestamps-as-timezone-naive-utc>
+# And [this issue](https://github.com/sqlalchemy/sqlalchemy/issues/1985)
+class TZDateTime(TypeDecorator):
+    """Timezone-aware datetime type."""
+
+    impl = DateTime
+    cache_ok = True
+
+    def process_bind_param(self, value: Any, dialect: Any) -> Any:  # noqa: ARG002
+        """Process bind param."""
+        if value is not None:
+            if not value.tzinfo or value.tzinfo.utcoffset(value) is None:
+                raise TypeError("tzinfo is required")
+            value = value.astimezone(UTC).replace(tzinfo=None)
+        return value
+
+    def process_result_value(self, value: Any, dialect: Any) -> Any:  # noqa: ARG002
+        """Process result value."""
+        if value is not None:
+            value = value.replace(tzinfo=UTC)
+        return value
+
+
 class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models."""
 
@@ -27,10 +52,11 @@ class CommonMixin:
 
     id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
     created_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), default=lambda: datetime.now(UTC)
+        TZDateTime, nullable=False, default=lambda: datetime.now(UTC)
     )
     updated_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True),
+        TZDateTime,
+        nullable=False,
         default=lambda: datetime.now(UTC),
         onupdate=lambda: datetime.now(UTC),
     )