kodit 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

kodit/infrastructure/indexing/index_repository.py

@@ -0,0 +1,296 @@
+"""Infrastructure implementation of the index repository."""
+
+from datetime import UTC, datetime
+from typing import TypeVar
+
+from sqlalchemy import delete, func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.entities import Embedding, File, Index, Snippet, Source
+from kodit.domain.services.indexing_service import IndexRepository
+from kodit.domain.value_objects import IndexView
+
+T = TypeVar("T")
+
+
+class SQLAlchemyIndexRepository(IndexRepository):
+    """SQLAlchemy implementation of the index repository."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        """Initialize the index repository.
+
+        Args:
+            session: The SQLAlchemy async session to use for database operations.
+
+        """
+        self.session = session
+
+    async def create_index(self, source_id: int) -> IndexView:
+        """Create a new index for a source.
+
+        Args:
+            source_id: The ID of the source to create an index for.
+
+        Returns:
+            The created index view.
+
+        """
+        # Check if index already exists
+        existing_index = await self.get_index_by_source_id(source_id)
+        if existing_index:
+            return existing_index
+
+        index = Index(source_id=source_id)
+        self.session.add(index)
+        await self.session.commit()
+
+        # Get source for the view
+        source_query = select(Source).where(Source.id == source_id)
+        source_result = await self.session.execute(source_query)
+        source = source_result.scalar_one()
+
+        return IndexView(
+            id=index.id,
+            created_at=index.created_at,
+            updated_at=index.updated_at,
+            source=source.uri,
+            num_snippets=0,
+        )
+
+    async def _get_index_view(self, index: Index, source: Source) -> IndexView:
+        """Create an IndexView from Index and Source entities.
+
+        Args:
+            index: The index entity
+            source: The source entity
+
+        Returns:
+            The index view
+
+        """
+        num_snippets = await self.num_snippets_for_index(index.id)
+        return IndexView(
+            id=index.id,
+            created_at=index.created_at,
+            updated_at=index.updated_at,
+            source=source.uri,
+            num_snippets=num_snippets,
+        )
+
+    async def get_index_by_id(self, index_id: int) -> IndexView | None:
+        """Get an index by its ID.
+
+        Args:
+            index_id: The ID of the index to retrieve.
+
+        Returns:
+            The index view if found, None otherwise.
+
+        """
+        query = (
+            select(Index, Source)
+            .join(Source, Index.source_id == Source.id)
+            .where(Index.id == index_id)
+        )
+        result = await self.session.execute(query)
+        row = result.first()
+
+        if not row:
+            return None
+
+        index, source = row
+        return await self._get_index_view(index, source)
+
+    async def get_index_by_source_id(self, source_id: int) -> IndexView | None:
+        """Get an index by its source ID.
+
+        Args:
+            source_id: The ID of the source to retrieve an index for.
+
+        Returns:
+            The index view if found, None otherwise.
+
+        """
+        query = (
+            select(Index, Source)
+            .join(Source, Index.source_id == Source.id)
+            .where(Index.source_id == source_id)
+        )
+        result = await self.session.execute(query)
+        row = result.first()
+
+        if not row:
+            return None
+
+        index, source = row
+        return await self._get_index_view(index, source)
+
+    async def list_indexes(self) -> list[IndexView]:
+        """List all indexes.
+
+        Returns:
+            A list of index views.
+
+        """
+        query = select(Index, Source).join(
+            Source, Index.source_id == Source.id, full=True
+        )
+        result = await self.session.execute(query)
+        rows = result.tuples()
+
+        indexes = []
+        for index, source in rows:
+            index_view = await self._get_index_view(index, source)
+            indexes.append(index_view)
+
+        return indexes
+
+    async def update_index_timestamp(self, index_id: int) -> None:
+        """Update the timestamp of an index.
+
+        Args:
+            index_id: The ID of the index to update.
+
+        """
+        query = select(Index).where(Index.id == index_id)
+        result = await self.session.execute(query)
+        index = result.scalar_one_or_none()
+
+        if index:
+            index.updated_at = datetime.now(UTC)
+            await self.session.commit()
+
+    async def delete_all_snippets(self, index_id: int) -> None:
+        """Delete all snippets for an index.
+
+        Args:
+            index_id: The ID of the index to delete snippets for.
+
+        """
+        # First get all snippets for this index
+        snippets = await self.get_snippets_for_index(index_id)
+
+        # Delete all embeddings for these snippets, if there are any
+        for snippet in snippets:
+            query = delete(Embedding).where(Embedding.snippet_id == snippet.id)
+            await self.session.execute(query)
+
+        # Now delete the snippets
+        query = delete(Snippet).where(Snippet.index_id == index_id)
+        await self.session.execute(query)
+        await self.session.commit()
+
+    async def get_snippets_for_index(self, index_id: int) -> list[Snippet]:
+        """Get all snippets for an index.
+
+        Args:
+            index_id: The ID of the index to get snippets for.
+
+        Returns:
+            A list of Snippet entities.
+
+        """
+        query = select(Snippet).where(Snippet.index_id == index_id)
+        result = await self.session.execute(query)
+        return list(result.scalars())
+
+    async def add_snippet(self, snippet: dict) -> None:
+        """Add a snippet to the database.
+
+        Args:
+            snippet: The snippet to add.
+
+        """
+        db_snippet = Snippet(
+            file_id=snippet["file_id"],
+            index_id=snippet["index_id"],
+            content=snippet["content"],
+        )
+        self.session.add(db_snippet)
+        await self.session.commit()
+
+    async def update_snippet_content(self, snippet_id: int, content: str) -> None:
+        """Update the content of an existing snippet.
+
+        Args:
+            snippet_id: The ID of the snippet to update.
+            content: The new content for the snippet.
+
+        """
+        query = select(Snippet).where(Snippet.id == snippet_id)
+        result = await self.session.execute(query)
+        snippet = result.scalar_one_or_none()
+
+        if snippet:
+            snippet.content = content
+            # SQLAlchemy will automatically track this change
+            await self.session.commit()
+
+    async def list_snippets_by_ids(self, ids: list[int]) -> list[tuple[dict, dict]]:
+        """List snippets by IDs.
+
+        Args:
+            ids: List of snippet IDs to retrieve.
+
+        Returns:
+            List of (file, snippet) tuples.
+
+        """
+        query = (
+            select(Snippet, File)
+            .where(Snippet.id.in_(ids))
+            .join(File, Snippet.file_id == File.id)
+        )
+        rows = await self.session.execute(query)
+
+        # Create a dictionary for O(1) lookup of results by ID
+        id_to_result = {}
+        for snippet, file in rows.all():
+            id_to_result[snippet.id] = (
+                {
+                    "id": file.id,
+                    "source_id": file.source_id,
+                    "mime_type": file.mime_type,
+                    "uri": file.uri,
+                    "cloned_path": file.cloned_path,
+                    "sha256": file.sha256,
+                    "size_bytes": file.size_bytes,
+                    "extension": file.extension,
+                    "created_at": file.created_at,
+                    "updated_at": file.updated_at,
+                },
+                {
+                    "id": snippet.id,
+                    "file_id": snippet.file_id,
+                    "index_id": snippet.index_id,
+                    "content": snippet.content,
+                    "created_at": snippet.created_at,
+                    "updated_at": snippet.updated_at,
+                },
+            )
+
+        # Check that all IDs are present
+        if len(id_to_result) != len(ids):
+            # Create a list of missing IDs
+            missing_ids = [
+                snippet_id for snippet_id in ids if snippet_id not in id_to_result
+            ]
+            msg = f"Some IDs are not present: {missing_ids}"
+            raise ValueError(msg)
+
+        # Rebuild the list in the same order that it was passed in
+        return [id_to_result[i] for i in ids]
+
+    async def num_snippets_for_index(self, index_id: int) -> int:
+        """Get the number of snippets for an index.
+
+        Args:
+            index_id: The ID of the index.
+
+        Returns:
+            The number of snippets.
+
+        """
+        query = select(func.count()).where(Snippet.index_id == index_id)
+        result = await self.session.execute(query)
+        return result.scalar_one()

kodit/infrastructure/indexing/indexing_factory.py

@@ -0,0 +1,111 @@
+"""Factory for creating indexing services."""
+
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.application.services.indexing_application_service import (
+    IndexingApplicationService,
+)
+from kodit.application.services.snippet_application_service import (
+    SnippetApplicationService,
+)
+from kodit.domain.services.bm25_service import BM25DomainService
+from kodit.domain.services.indexing_service import IndexingDomainService
+from kodit.domain.services.source_service import SourceService
+from kodit.infrastructure.bm25.bm25_factory import bm25_repository_factory
+from kodit.infrastructure.embedding.embedding_factory import (
+    embedding_domain_service_factory,
+)
+from kodit.infrastructure.enrichment.enrichment_factory import (
+    create_enrichment_domain_service,
+)
+from kodit.infrastructure.indexing.fusion_service import ReciprocalRankFusionService
+from kodit.infrastructure.indexing.index_repository import SQLAlchemyIndexRepository
+from kodit.infrastructure.snippet_extraction.snippet_extraction_factory import (
+    create_snippet_extraction_domain_service,
+)
+from kodit.infrastructure.sqlalchemy.file_repository import (
+    SqlAlchemyFileRepository,
+)
+from kodit.infrastructure.sqlalchemy.snippet_repository import (
+    SqlAlchemySnippetRepository,
+)
+
+
+def create_snippet_application_service(
+    session: AsyncSession,
+) -> SnippetApplicationService:
+    """Create a snippet application service with all dependencies."""
+    # Create domain service
+    snippet_extraction_service = create_snippet_extraction_domain_service()
+
+    # Create repositories
+    snippet_repository = SqlAlchemySnippetRepository(session)
+    file_repository = SqlAlchemyFileRepository(session)
+
+    # Create application service
+    from kodit.application.services.snippet_application_service import (
+        SnippetApplicationService,
+    )
+
+    return SnippetApplicationService(
+        snippet_extraction_service=snippet_extraction_service,
+        snippet_repository=snippet_repository,
+        file_repository=file_repository,
+    )
+
+
+def create_indexing_domain_service(session: AsyncSession) -> IndexingDomainService:
+    """Create an indexing domain service.
+
+    Args:
+        session: The database session.
+
+    Returns:
+        An indexing domain service instance.
+
+    """
+    index_repository = SQLAlchemyIndexRepository(session)
+    fusion_service = ReciprocalRankFusionService()
+
+    return IndexingDomainService(
+        index_repository=index_repository,
+        fusion_service=fusion_service,
+    )
+
+
+def create_indexing_application_service(
+    app_context: Any,
+    session: AsyncSession,
+    source_service: SourceService,
+    snippet_application_service: SnippetApplicationService,
+) -> IndexingApplicationService:
+    """Create an indexing application service.
+
+    Args:
+        app_context: The application context.
+        session: The database session.
+        source_service: The source service.
+        snippet_application_service: The snippet application service.
+
+    Returns:
+        An indexing application service instance.
+
+    """
+    # Create domain services
+    indexing_domain_service = create_indexing_domain_service(session)
+    bm25_service = BM25DomainService(bm25_repository_factory(app_context, session))
+    code_search_service = embedding_domain_service_factory("code", app_context, session)
+    text_search_service = embedding_domain_service_factory("text", app_context, session)
+    enrichment_service = create_enrichment_domain_service(app_context)
+
+    return IndexingApplicationService(
+        indexing_domain_service=indexing_domain_service,
+        source_service=source_service,
+        bm25_service=bm25_service,
+        code_search_service=code_search_service,
+        text_search_service=text_search_service,
+        enrichment_service=enrichment_service,
+        snippet_application_service=snippet_application_service,
+    )

kodit/infrastructure/snippet_extraction/__init__.py

@@ -0,0 +1 @@
+"""Infrastructure services for snippet extraction."""

kodit/infrastructure/snippet_extraction/language_detection_service.py

@@ -0,0 +1,39 @@
+"""Infrastructure implementation for language detection."""
+
+from pathlib import Path
+
+from kodit.domain.services.snippet_extraction_service import LanguageDetectionService
+
+
+class FileSystemLanguageDetectionService(LanguageDetectionService):
+    """Infrastructure implementation for language detection."""
+
+    def __init__(self, language_map: dict[str, str]) -> None:
+        """Initialize the language detection service.
+
+        Args:
+            language_map: Mapping of file extensions to programming languages
+
+        """
+        self.language_map = language_map
+
+    async def detect_language(self, file_path: Path) -> str:
+        """Detect language based on file extension.
+
+        Args:
+            file_path: Path to the file to detect language for
+
+        Returns:
+            The detected programming language
+
+        Raises:
+            ValueError: If the language is not supported
+
+        """
+        suffix = file_path.suffix.removeprefix(".").lower()
+        language = self.language_map.get(suffix)
+
+        if language is None:
+            raise ValueError(f"Unsupported language for file suffix: {suffix}")
+
+        return language

kodit/infrastructure/snippet_extraction/snippet_extraction_factory.py

@@ -0,0 +1,95 @@
+"""Factory for creating snippet extraction services."""
+
+from pathlib import Path
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.enums import SnippetExtractionStrategy
+from kodit.domain.repositories import FileRepository, SnippetRepository
+from kodit.domain.services.snippet_extraction_service import (
+    SnippetExtractionDomainService,
+)
+from kodit.infrastructure.snippet_extraction.language_detection_service import (
+    FileSystemLanguageDetectionService,
+)
+from kodit.infrastructure.snippet_extraction.snippet_query_provider import (
+    FileSystemSnippetQueryProvider,
+)
+from kodit.infrastructure.snippet_extraction.tree_sitter_snippet_extractor import (
+    TreeSitterSnippetExtractor,
+)
+from kodit.infrastructure.sqlalchemy.file_repository import SqlAlchemyFileRepository
+from kodit.infrastructure.sqlalchemy.snippet_repository import (
+    SqlAlchemySnippetRepository,
+)
+
+
+def create_snippet_extraction_domain_service() -> SnippetExtractionDomainService:
+    """Create a snippet extraction domain service with all dependencies.
+
+    Returns:
+        Configured snippet extraction domain service
+
+    """
+    # Language mapping from the existing languages module
+    language_map = {
+        # JavaScript/TypeScript
+        "js": "javascript",
+        "jsx": "javascript",
+        "ts": "typescript",
+        "tsx": "typescript",
+        # Python
+        "py": "python",
+        # Rust
+        "rs": "rust",
+        # Go
+        "go": "go",
+        # C/C++
+        "cpp": "cpp",
+        "hpp": "cpp",
+        "c": "c",
+        "h": "c",
+        # C#
+        "cs": "csharp",
+        # Ruby
+        "rb": "ruby",
+        # Java
+        "java": "java",
+        # PHP
+        "php": "php",
+        # Swift
+        "swift": "swift",
+        # Kotlin
+        "kt": "kotlin",
+    }
+
+    # Create infrastructure services
+    language_detector = FileSystemLanguageDetectionService(language_map)
+    query_provider = FileSystemSnippetQueryProvider(Path(__file__).parent / "languages")
+
+    # Create snippet extractors
+    method_extractor = TreeSitterSnippetExtractor(query_provider)
+
+    snippet_extractors = {
+        SnippetExtractionStrategy.METHOD_BASED: method_extractor,
+    }
+
+    # Create domain service
+    return SnippetExtractionDomainService(language_detector, snippet_extractors)
+
+
+def create_snippet_repositories(
+    session: AsyncSession,
+) -> tuple[SnippetRepository, FileRepository]:
+    """Create snippet and file repositories.
+
+    Args:
+        session: SQLAlchemy session
+
+    Returns:
+        Tuple of (snippet_repository, file_repository)
+
+    """
+    snippet_repository = SqlAlchemySnippetRepository(session)
+    file_repository = SqlAlchemyFileRepository(session)
+    return snippet_repository, file_repository

kodit/infrastructure/snippet_extraction/snippet_query_provider.py

@@ -0,0 +1,45 @@
+"""Infrastructure implementation for loading snippet queries from files."""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class SnippetQueryProvider(ABC):
+    """Abstract interface for providing snippet queries."""
+
+    @abstractmethod
+    async def get_query(self, language: str) -> str:
+        """Get the query for a specific language."""
+
+
+class FileSystemSnippetQueryProvider(SnippetQueryProvider):
+    """Infrastructure implementation for loading snippet queries from files."""
+
+    def __init__(self, query_directory: Path) -> None:
+        """Initialize the query provider.
+
+        Args:
+            query_directory: Directory containing query files
+
+        """
+        self.query_directory = query_directory
+
+    async def get_query(self, language: str) -> str:
+        """Load query from file system.
+
+        Args:
+            language: The programming language to get the query for
+
+        Returns:
+            The query string for the language
+
+        Raises:
+            FileNotFoundError: If the query file doesn't exist
+
+        """
+        query_path = self.query_directory / f"{language}.scm"
+
+        if not query_path.exists():
+            raise FileNotFoundError(f"Query file not found: {query_path}")
+
+        return query_path.read_text()