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

kodit/{snippets/method_snippets.py → infrastructure/snippet_extraction/tree_sitter_snippet_extractor.py}

@@ -1,17 +1,132 @@
-"""Extract method snippets from source code."""
+"""Infrastructure implementation using tree-sitter for method extraction."""
+
+from pathlib import Path
+from typing import cast
 
 from tree_sitter import Node, Query
 from tree_sitter_language_pack import SupportedLanguage, get_language, get_parser
 
+from kodit.domain.services.snippet_extraction_service import SnippetExtractor
+from kodit.infrastructure.snippet_extraction.snippet_query_provider import (
+    SnippetQueryProvider,
+)
+
+
+class TreeSitterSnippetExtractor(SnippetExtractor):
+    """Infrastructure implementation using tree-sitter for method extraction."""
+
+    def __init__(self, query_provider: SnippetQueryProvider) -> None:
+        """Initialize the tree-sitter snippet extractor.
+
+        Args:
+            query_provider: Provider for snippet queries
+
+        """
+        self.query_provider = query_provider
+
+    async def extract(self, file_path: Path, language: str) -> list[str]:
+        """Extract snippets using tree-sitter parsing.
+
+        Args:
+            file_path: Path to the file to extract snippets from
+            language: The programming language of the file
+
+        Returns:
+            List of extracted code snippets
+
+        Raises:
+            ValueError: If the file cannot be read or language is not supported
+
+        """
+        try:
+            # Get the query for the language
+            query = await self.query_provider.get_query(language)
+        except FileNotFoundError as e:
+            raise ValueError(f"Unsupported language: {file_path}") from e
+
+        # Get parser and language for tree-sitter
+        try:
+            tree_sitter_language = get_language(cast("SupportedLanguage", language))
+            parser = get_parser(cast("SupportedLanguage", language))
+        except Exception as e:
+            raise ValueError(f"Unsupported language: {file_path}") from e
+
+        # Create query object
+        query_obj = Query(tree_sitter_language, query)
+
+        # Read file content
+        try:
+            file_bytes = file_path.read_bytes()
+        except Exception as e:
+            raise ValueError(f"Failed to read file: {file_path}") from e
 
-class MethodSnippets:
-    """Extract method snippets from source code."""
+        # Parse and extract snippets
+        tree = parser.parse(file_bytes)
+        captures_by_name = query_obj.captures(tree.root_node)
+        lines = file_bytes.decode().splitlines()
 
-    def __init__(self, language: SupportedLanguage, query: str) -> None:
-        """Initialize the MethodSnippets class."""
-        self.language = get_language(language)
-        self.parser = get_parser(language)
-        self.query = Query(self.language, query)
+        # Extract snippets using the existing logic
+        snippets = self._extract_snippets_from_captures(captures_by_name, lines)
+
+        # If there are no results, return the entire file
+        if not snippets:
+            return [file_bytes.decode()]
+
+        return snippets
+
+    def _extract_snippets_from_captures(
+        self, captures_by_name: dict[str, list[Node]], lines: list[str]
+    ) -> list[str]:
+        """Extract snippets from tree-sitter captures.
+
+        Args:
+            captures_by_name: Captures organized by name
+            lines: Lines of the source file
+
+        Returns:
+            List of extracted code snippets
+
+        """
+        # Find all leaf functions
+        leaf_functions = self._get_leaf_functions(captures_by_name)
+
+        # Find all imports
+        imports = self._get_imports(captures_by_name)
+
+        results = []
+
+        # For each leaf function, find all lines this function is dependent on
+        for func_node in leaf_functions:
+            all_lines_to_keep = set()
+
+            ancestors = self._get_ancestors(captures_by_name, func_node)
+
+            # Add self to keep
+            all_lines_to_keep.update(
+                range(func_node.start_point[0], func_node.end_point[0] + 1)
+            )
+
+            # Add imports to keep
+            for import_node in imports:
+                all_lines_to_keep.update(
+                    range(import_node.start_point[0], import_node.end_point[0] + 1)
+                )
+
+            # Add ancestors to keep
+            for node in ancestors:
+                # Get the first line of the node for now
+                start = node.start_point[0]
+                end = node.start_point[0]
+                all_lines_to_keep.update(range(start, end + 1))
+
+            pseudo_code = []
+            for i, line in enumerate(lines):
+                if i in all_lines_to_keep:
+                    pseudo_code.append(line)
+
+            results.append("\n".join(pseudo_code))
+
+        return results
 
     def _get_leaf_functions(
         self, captures_by_name: dict[str, list[Node]]
@@ -65,56 +180,3 @@ class MethodSnippets:
                 ancestors.append(parent)
             parent = parent.parent
         return ancestors
-
-    def extract(self, source_code: bytes) -> list[str]:
-        """Extract method snippets from source code."""
-        tree = self.parser.parse(source_code)
-
-        captures_by_name = self.query.captures(tree.root_node)
-
-        lines = source_code.decode().splitlines()
-
-        # Find all leaf functions
-        leaf_functions = self._get_leaf_functions(captures_by_name)
-
-        # Find all imports
-        imports = self._get_imports(captures_by_name)
-
-        results = []
-
-        # For each leaf function, find all lines this function is dependent on
-        for func_node in leaf_functions:
-            all_lines_to_keep = set()
-
-            ancestors = self._get_ancestors(captures_by_name, func_node)
-
-            # Add self to keep
-            all_lines_to_keep.update(
-                range(func_node.start_point[0], func_node.end_point[0] + 1)
-            )
-
-            # Add imports to keep
-            for import_node in imports:
-                all_lines_to_keep.update(
-                    range(import_node.start_point[0], import_node.end_point[0] + 1)
-                )
-
-            # Add ancestors to keep
-            for node in ancestors:
-                # Get the first line of the node for now
-                start = node.start_point[0]
-                end = node.start_point[0]
-                all_lines_to_keep.update(range(start, end + 1))
-
-            pseudo_code = []
-            for i, line in enumerate(lines):
-                if i in all_lines_to_keep:
-                    pseudo_code.append(line)
-
-            results.append("\n".join(pseudo_code))
-
-        # If there are no results, then return the entire file
-        if not results:
-            return [source_code.decode()]
-
-        return results

kodit/infrastructure/sqlalchemy/__init__.py

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

kodit/{embedding → infrastructure/sqlalchemy}/embedding_repository.py

@@ -1,35 +1,32 @@
-"""Repository for managing embeddings."""
+"""SQLAlchemy implementation of embedding repository."""
 
 import numpy as np
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from kodit.embedding.embedding_models import Embedding, EmbeddingType
+from kodit.domain.entities import Embedding, EmbeddingType
 
 
-class EmbeddingRepository:
-    """Repository for managing embeddings.
+class SqlAlchemyEmbeddingRepository:
+    """SQLAlchemy implementation of embedding repository."""
 
-    This class provides methods for creating and retrieving embeddings from the
-    database. It handles the low-level database operations and transaction management.
-
-    Args:
-        session: The SQLAlchemy async session to use for database operations.
+    def __init__(self, session: AsyncSession) -> None:
+        """Initialize the SQLAlchemy embedding repository.
 
-    """
+        Args:
+            session: The SQLAlchemy async session to use for database operations
 
-    def __init__(self, session: AsyncSession) -> None:
-        """Initialize the embedding repository."""
+        """
         self.session = session
 
     async def create_embedding(self, embedding: Embedding) -> Embedding:
         """Create a new embedding record in the database.
 
         Args:
-            embedding: The Embedding instance to create.
+            embedding: The Embedding instance to create
 
         Returns:
-            The created Embedding instance.
+            The created Embedding instance
 
         """
         self.session.add(embedding)
@@ -42,11 +39,11 @@ class EmbeddingRepository:
         """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.
+            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.
+            The Embedding instance if found, None otherwise
 
         """
         query = select(Embedding).where(
@@ -62,10 +59,10 @@ class EmbeddingRepository:
         """List all embeddings of a given type.
 
         Args:
-            embedding_type: The type of embeddings to list.
+            embedding_type: The type of embeddings to list
 
         Returns:
-            A list of Embedding instances.
+            A list of Embedding instances
 
         """
         query = select(Embedding).where(Embedding.type == embedding_type)
@@ -76,7 +73,7 @@ class EmbeddingRepository:
         """Delete all embeddings for a snippet.
 
         Args:
-            snippet_id: The ID of the snippet to delete embeddings for.
+            snippet_id: The ID of the snippet to delete embeddings for
 
         """
         query = select(Embedding).where(Embedding.snippet_id == snippet_id)
@@ -181,7 +178,27 @@ class EmbeddingRepository:
         """
         stored_norms = np.linalg.norm(stored_vecs, axis=1)
         query_norm = np.linalg.norm(query_vec)
-        return np.dot(stored_vecs, query_vec) / (stored_norms * query_norm)
+
+        # Handle zero vectors to avoid division by zero
+        if query_norm == 0:
+            # If query vector is zero, return zeros for all similarities
+            return np.zeros(len(stored_vecs))
+
+        # Handle stored vectors with zero norms
+        zero_stored_mask = stored_norms == 0
+        similarities = np.zeros(len(stored_vecs))
+
+        # Only compute similarities for non-zero stored vectors
+        non_zero_mask = ~zero_stored_mask
+        if np.any(non_zero_mask):
+            non_zero_stored_vecs = stored_vecs[non_zero_mask]
+            non_zero_stored_norms = stored_norms[non_zero_mask]
+            non_zero_similarities = np.dot(non_zero_stored_vecs, query_vec) / (
+                non_zero_stored_norms * query_norm
+            )
+            similarities[non_zero_mask] = non_zero_similarities
+
+        return similarities
 
     def _get_top_k_results(
         self,
@@ -200,7 +217,6 @@ class EmbeddingRepository:
             List of (snippet_id, similarity_score) tuples
 
         """
+        # Get indices of top-k similarities
         top_indices = np.argsort(similarities)[::-1][:top_k]
-        return [
-            (embeddings[i][0], float(similarities[i])) for i in top_indices
-        ]  # Use index 0 to get snippet_id
+        return [(embeddings[i][0], float(similarities[i])) for i in top_indices]

kodit/infrastructure/sqlalchemy/file_repository.py

@@ -0,0 +1,73 @@
+"""SQLAlchemy implementation of file repository."""
+
+from collections.abc import Sequence
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.entities import File, Index
+from kodit.domain.repositories import FileRepository
+
+
+class SqlAlchemyFileRepository(FileRepository):
+    """SQLAlchemy implementation of file repository."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        """Initialize the SQLAlchemy file repository.
+
+        Args:
+            session: The SQLAlchemy async session to use for database operations
+
+        """
+        self.session = session
+
+    async def get_files_for_index(self, index_id: int) -> Sequence[File]:
+        """Get all files for an index.
+
+        Args:
+            index_id: The ID of the index to get files for
+
+        Returns:
+            A list of File instances
+
+        """
+        # Get the index first to find its source_id
+        index_query = select(Index).where(Index.id == index_id)
+        index_result = await self.session.execute(index_query)
+        index = index_result.scalar_one_or_none()
+
+        if not index:
+            return []
+
+        # Get all files for the source
+        query = select(File).where(File.source_id == index.source_id)
+        result = await self.session.execute(query)
+        return list(result.scalars())
+
+    async def get_by_id(self, file_id: int) -> File | None:
+        """Get a file by ID.
+
+        Args:
+            file_id: The ID of the file to retrieve
+
+        Returns:
+            The File instance if found, None otherwise
+
+        """
+        query = select(File).where(File.id == file_id)
+        result = await self.session.execute(query)
+        return result.scalar_one_or_none()
+
+    async def save(self, file: File) -> File:
+        """Save file using SQLAlchemy.
+
+        Args:
+            file: The file to save
+
+        Returns:
+            The saved file
+
+        """
+        self.session.add(file)
+        await self.session.commit()
+        return file

kodit/infrastructure/sqlalchemy/repository.py

@@ -0,0 +1,121 @@
+"""SQLAlchemy repository."""
+
+from collections.abc import Sequence
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.entities import Author, AuthorFileMapping, File, Source, SourceType
+from kodit.domain.repositories import AuthorRepository, SourceRepository
+
+
+class SqlAlchemySourceRepository(SourceRepository):
+    """SQLAlchemy source repository."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        """Initialize the repository."""
+        self._session = session
+
+    async def get(self, source_id: int) -> Source | None:
+        """Get a source by ID."""
+        return await self._session.get(Source, source_id)
+
+    async def get_by_uri(self, uri: str) -> Source | None:
+        """Get a source by URI."""
+        stmt = select(Source).where(Source.uri == uri)
+        return await self._session.scalar(stmt)  # None if no row
+
+    async def list(self, *, source_type: SourceType | None = None) -> Sequence[Source]:
+        """List sources."""
+        stmt = select(Source)
+        if source_type is not None:
+            stmt = stmt.where(Source.type == source_type)
+        return (await self._session.scalars(stmt)).all()
+
+    async def add(self, source: Source) -> None:
+        """Add a source."""
+        self._session.add(source)  # INSERT on flush
+        await self._session.flush()  # Flush to get the ID
+
+    async def create_source(self, source: Source) -> Source:
+        """Create a source and commit it."""
+        self._session.add(source)
+        await self._session.commit()
+        return source
+
+    async def remove(self, source: Source) -> None:
+        """Remove a source."""
+        await self._session.delete(source)  # DELETE on flush
+
+    async def create_file(self, file: File) -> File:
+        """Create a new file record."""
+        self._session.add(file)
+        await self._session.commit()
+        return file
+
+    async def upsert_author(self, author: Author) -> Author:
+        """Create a new author or return existing one if email already exists."""
+        # First check if author already exists with same name and email
+        stmt = select(Author).where(
+            Author.name == author.name, Author.email == author.email
+        )
+        existing_author = await self._session.scalar(stmt)
+
+        if existing_author:
+            return existing_author
+
+        # Author doesn't exist, create new one
+        self._session.add(author)
+        await self._session.commit()
+        return author
+
+    async def upsert_author_file_mapping(
+        self, mapping: AuthorFileMapping
+    ) -> AuthorFileMapping:
+        """Create a new author file mapping or return existing one if already exists."""
+        # First check if mapping already exists with same author_id and file_id
+        stmt = select(AuthorFileMapping).where(
+            AuthorFileMapping.author_id == mapping.author_id,
+            AuthorFileMapping.file_id == mapping.file_id,
+        )
+        existing_mapping = await self._session.scalar(stmt)
+
+        if existing_mapping:
+            return existing_mapping
+
+        # Mapping doesn't exist, create new one
+        self._session.add(mapping)
+        await self._session.commit()
+        return mapping
+
+
+class SqlAlchemyAuthorRepository(AuthorRepository):
+    """SQLAlchemy author repository."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        """Initialize the repository."""
+        self._session = session
+
+    async def get(self, author_id: int) -> Author | None:
+        """Get an author by ID."""
+        return await self._session.get(Author, author_id)
+
+    async def get_by_name(self, name: str) -> Author | None:
+        """Get an author by name."""
+        return await self._session.scalar(select(Author).where(Author.name == name))
+
+    async def get_by_email(self, email: str) -> Author | None:
+        """Get an author by email."""
+        return await self._session.scalar(select(Author).where(Author.email == email))
+
+    async def list(self) -> Sequence[Author]:
+        """List authors."""
+        return (await self._session.scalars(select(Author))).all()
+
+    async def add(self, author: Author) -> None:
+        """Add an author."""
+        self._session.add(author)
+
+    async def remove(self, author: Author) -> None:
+        """Remove an author."""
+        await self._session.delete(author)

kodit/infrastructure/sqlalchemy/snippet_repository.py

@@ -0,0 +1,75 @@
+"""SQLAlchemy implementation of snippet repository."""
+
+from collections.abc import Sequence
+
+from sqlalchemy import delete, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.entities import Snippet
+from kodit.domain.repositories import SnippetRepository
+
+
+class SqlAlchemySnippetRepository(SnippetRepository):
+    """SQLAlchemy implementation of snippet repository."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        """Initialize the SQLAlchemy snippet repository.
+
+        Args:
+            session: The SQLAlchemy async session to use for database operations
+
+        """
+        self.session = session
+
+    async def save(self, snippet: Snippet) -> Snippet:
+        """Save snippet using SQLAlchemy.
+
+        Args:
+            snippet: The snippet to save
+
+        Returns:
+            The saved snippet
+
+        """
+        self.session.add(snippet)
+        await self.session.commit()
+        return snippet
+
+    async def get_by_id(self, snippet_id: int) -> Snippet | None:
+        """Get a snippet by ID.
+
+        Args:
+            snippet_id: The ID of the snippet to retrieve
+
+        Returns:
+            The Snippet instance if found, None otherwise
+
+        """
+        query = select(Snippet).where(Snippet.id == snippet_id)
+        result = await self.session.execute(query)
+        return result.scalar_one_or_none()
+
+    async def get_by_index(self, index_id: int) -> Sequence[Snippet]:
+        """Get all snippets for an index.
+
+        Args:
+            index_id: The ID of the index to get snippets for
+
+        Returns:
+            A list of Snippet instances
+
+        """
+        query = select(Snippet).where(Snippet.index_id == index_id)
+        result = await self.session.execute(query)
+        return list(result.scalars())
+
+    async def delete_by_index(self, index_id: int) -> None:
+        """Delete all snippets for an index.
+
+        Args:
+            index_id: The ID of the index to delete snippets for
+
+        """
+        query = delete(Snippet).where(Snippet.index_id == index_id)
+        await self.session.execute(query)
+        await self.session.commit()

kodit/infrastructure/ui/__init__.py

@@ -0,0 +1 @@
+"""UI infrastructure module."""