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

kodit/indexing/fusion.py

@@ -1,67 +0,0 @@
-"""Fusion functions for combining search results."""
-
-from collections import defaultdict
-from dataclasses import dataclass
-
-
-@dataclass
-class FusionResult:
-    """Result of a fusion operation."""
-
-    id: int
-    score: float
-    original_scores: list[float]
-
-
-@dataclass
-class FusionRequest:
-    """Result of a RRF operation."""
-
-    id: int
-    score: float
-
-
-def reciprocal_rank_fusion(
-    rankings: list[list[FusionRequest]], k: float = 60
-) -> list[FusionResult]:
-    """RRF prioritises results that are present in all results.
-
-    Args:
-        rankings: List of rankers, each containing a list of document ids. Top of the
-        list is considered to be the best result.
-        k: Parameter for RRF.
-
-    Returns:
-        Dictionary of ids and their scores.
-
-    """
-    scores = {}
-    for ranker in rankings:
-        for rank in ranker:
-            scores[rank.id] = float(0)
-
-    for ranker in rankings:
-        for i, rank in enumerate(ranker):
-            scores[rank.id] += 1.0 / (k + i)
-
-    # Create a list of tuples of ids and their scores
-    results = [(rank, scores[rank]) for rank in scores]
-
-    # Sort results by score
-    results.sort(key=lambda x: x[1], reverse=True)
-
-    # Create a map of original scores to ids
-    original_scores_to_ids = defaultdict(list)
-    for ranker in rankings:
-        for rank in ranker:
-            original_scores_to_ids[rank.id].append(rank.score)
-
-    # Rebuild a list of final results with their original scores
-    return [
-        FusionResult(
-            id=result[0],
-            score=result[1],
-            original_scores=original_scores_to_ids[result[0]],
-        )
-        for result in results
-    ]

kodit/indexing/indexing_models.py

@@ -1,43 +0,0 @@
-"""Index models for managing code indexes.
-
-This module defines the SQLAlchemy models used for storing and managing code indexes,
-including files and snippets. It provides the data structures for tracking indexed
-files and their content.
-"""
-
-from sqlalchemy import ForeignKey, UnicodeText
-from sqlalchemy.orm import Mapped, mapped_column
-
-from kodit.database import Base, CommonMixin
-
-
-class Index(Base, CommonMixin):
-    """Index model."""
-
-    __tablename__ = "indexes"
-
-    source_id: Mapped[int] = mapped_column(
-        ForeignKey("sources.id"), unique=True, index=True
-    )
-
-    def __init__(self, source_id: int) -> None:
-        """Initialize the index."""
-        super().__init__()
-        self.source_id = source_id
-
-
-class Snippet(Base, CommonMixin):
-    """Snippet model."""
-
-    __tablename__ = "snippets"
-
-    file_id: Mapped[int] = mapped_column(ForeignKey("files.id"), index=True)
-    index_id: Mapped[int] = mapped_column(ForeignKey("indexes.id"), index=True)
-    content: Mapped[str] = mapped_column(UnicodeText, default="")
-
-    def __init__(self, file_id: int, index_id: int, content: str) -> None:
-        """Initialize the snippet."""
-        super().__init__()
-        self.file_id = file_id
-        self.index_id = index_id
-        self.content = content

kodit/indexing/indexing_repository.py

@@ -1,216 +0,0 @@
-"""Repository for managing code indexes and their associated files and snippets.
-
-This module provides the IndexRepository class which handles all database operations
-related to code indexes, including creating indexes, managing files and snippets,
-and retrieving index information with their associated metadata.
-"""
-
-from datetime import UTC, datetime
-from typing import TypeVar
-
-from sqlalchemy import delete, func, select
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from kodit.embedding.embedding_models import Embedding
-from kodit.indexing.indexing_models import Index, Snippet
-from kodit.source.source_models import File, Source
-
-T = TypeVar("T")
-
-
-class IndexRepository:
-    """Repository for managing code indexes and their associated data.
-
-    This class provides methods for creating and managing code indexes, including
-    their associated files and snippets. It handles all database operations related
-    to indexing code sources.
-    """
-
-    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(self, source_id: int) -> Index:
-        """Create a new index for a source.
-
-        Args:
-            source_id: The ID of the source to create an index for.
-
-        Returns:
-            The newly created Index instance.
-
-        """
-        index = Index(source_id=source_id)
-        self.session.add(index)
-        await self.session.commit()
-        return index
-
-    async def get_by_id(self, index_id: int) -> Index | None:
-        """Get an index by its ID.
-
-        Args:
-            index_id: The ID of the index to retrieve.
-
-        Returns:
-            The Index instance if found, None otherwise.
-
-        """
-        query = select(Index).where(Index.id == index_id)
-        result = await self.session.execute(query)
-        return result.scalar_one_or_none()
-
-    async def get_by_source_id(self, source_id: int) -> Index | None:
-        """Get an index by its source ID.
-
-        Args:
-            source_id: The ID of the source to retrieve an index for.
-
-        """
-        query = select(Index).where(Index.source_id == source_id)
-        result = await self.session.execute(query)
-        return result.scalar_one_or_none()
-
-    async def files_for_index(self, index_id: int) -> list[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.
-
-        """
-        query = (
-            select(File)
-            .join(Source, File.source_id == Source.id)
-            .join(Index, Index.source_id == Source.id)
-            .where(Index.id == index_id)
-        )
-        result = await self.session.execute(query)
-        return list(result.scalars())
-
-    async def list_indexes(self) -> list[tuple[Index, Source]]:
-        """List all indexes.
-
-        Returns:
-            A list of tuples containing index information, source details,
-            and counts of files and snippets.
-
-        """
-        query = select(Index, Source).join(
-            Source, Index.source_id == Source.id, full=True
-        )
-        result = await self.session.execute(query)
-        return list(result.tuples())
-
-    async def num_snippets_for_index(self, index_id: int) -> int:
-        """Get the number of snippets for an index."""
-        query = select(func.count()).where(Snippet.index_id == index_id)
-        result = await self.session.execute(query)
-        return result.scalar_one()
-
-    async def update_index_timestamp(self, index: Index) -> None:
-        """Update the updated_at timestamp of an index.
-
-        Args:
-            index: The Index instance to update.
-
-        """
-        index.updated_at = datetime.now(UTC)
-        await self.session.commit()
-
-    async def add_snippet(self, snippet: Snippet) -> None:
-        """Add a new snippet to the database if it doesn't exist, otherwise update it.
-
-        Args:
-            snippet: The Snippet instance to add.
-
-        """
-        self.session.add(snippet)
-        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.
-
-        """
-        query = select(Snippet).where(Snippet.index_id == index_id)
-        result = await self.session.execute(query)
-        return list(result.scalars())
-
-    async def get_all_snippets(self, index_id: int) -> list[Snippet]:
-        """Get all snippets.
-
-        Returns:
-            A list of all snippets.
-
-        """
-        query = select(Snippet).where(Snippet.index_id == index_id).order_by(Snippet.id)
-        result = await self.session.execute(query)
-        return list(result.scalars())
-
-    async def add_embedding(self, embedding: Embedding) -> None:
-        """Add a new embedding to the database.
-
-        Args:
-            embedding: The Embedding instance to add.
-
-        """
-        self.session.add(embedding)
-        await self.session.commit()
-
-    async def list_snippets_by_ids(self, ids: list[int]) -> list[tuple[File, Snippet]]:
-        """List snippets by IDs.
-
-        Returns:
-            A list of snippets in the same order as the input IDs.
-
-        """
-        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 = {snippet.id: (file, snippet) for snippet, file in rows.all()}
-
-        # 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]

kodit/indexing/indexing_service.py

@@ -1,344 +0,0 @@
-"""Index service for managing code indexes.
-
-This module provides the IndexService class which handles the business logic for
-creating, listing, and running code indexes. It orchestrates the interaction between the
-file system, database operations (via IndexRepository), and provides a clean API for
-index management.
-"""
-
-from datetime import datetime
-from pathlib import Path
-
-import pydantic
-import structlog
-from tqdm.asyncio import tqdm
-
-from kodit.bm25.keyword_search_service import (
-    BM25Document,
-    BM25Result,
-    KeywordSearchProvider,
-)
-from kodit.embedding.vector_search_service import (
-    VectorSearchRequest,
-    VectorSearchService,
-)
-from kodit.enrichment.enrichment_provider.enrichment_provider import EnrichmentRequest
-from kodit.enrichment.enrichment_service import EnrichmentService
-from kodit.indexing.fusion import FusionRequest, reciprocal_rank_fusion
-from kodit.indexing.indexing_models import Snippet
-from kodit.indexing.indexing_repository import IndexRepository
-from kodit.log import log_event
-from kodit.snippets.snippets import SnippetService
-from kodit.source.source_service import SourceService
-from kodit.util.spinner import Spinner
-
-# List of MIME types that are blacklisted from being indexed
-MIME_BLACKLIST = ["unknown/unknown"]
-
-
-class IndexView(pydantic.BaseModel):
-    """Data transfer object for index information.
-
-    This model represents the public interface for index data, providing a clean
-    view of index information without exposing internal implementation details.
-    """
-
-    id: int
-    created_at: datetime
-    updated_at: datetime | None = None
-    source: str | None = None
-    num_snippets: int
-
-
-class SearchRequest(pydantic.BaseModel):
-    """Request for a search."""
-
-    text_query: str | None = None
-    code_query: str | None = None
-    keywords: list[str] | None = None
-    top_k: int = 10
-
-
-class SearchResult(pydantic.BaseModel):
-    """Data transfer object for search results.
-
-    This model represents a single search result, containing both the file path
-    and the matching snippet content.
-    """
-
-    id: int
-    uri: str
-    content: str
-    original_scores: list[float]
-
-
-class IndexService:
-    """Service for managing code indexes.
-
-    This service handles the business logic for creating, listing, and running code
-    indexes. It coordinates between file system operations, database operations (via
-    IndexRepository), and provides a clean API for index management.
-    """
-
-    def __init__(  # noqa: PLR0913
-        self,
-        repository: IndexRepository,
-        source_service: SourceService,
-        keyword_search_provider: KeywordSearchProvider,
-        code_search_service: VectorSearchService,
-        text_search_service: VectorSearchService,
-        enrichment_service: EnrichmentService,
-    ) -> None:
-        """Initialize the index service.
-
-        Args:
-            repository: The repository instance to use for database operations.
-            source_service: The source service instance to use for source validation.
-
-        """
-        self.repository = repository
-        self.source_service = source_service
-        self.snippet_service = SnippetService()
-        self.log = structlog.get_logger(__name__)
-        self.keyword_search_provider = keyword_search_provider
-        self.code_search_service = code_search_service
-        self.text_search_service = text_search_service
-        self.enrichment_service = enrichment_service
-
-    async def create(self, source_id: int) -> IndexView:
-        """Create a new index for a source.
-
-        This method creates a new index for the specified source, after validating
-        that the source exists and doesn't already have an index.
-
-        Args:
-            source_id: The ID of the source to create an index for.
-
-        Returns:
-            An Index object representing the newly created index.
-
-        Raises:
-            ValueError: If the source doesn't exist or already has an index.
-
-        """
-        log_event("kodit.index.create")
-
-        # Check if the source exists
-        source = await self.source_service.get(source_id)
-
-        # Check if the index already exists
-        index = await self.repository.get_by_source_id(source.id)
-        if not index:
-            index = await self.repository.create(source.id)
-        return IndexView(
-            id=index.id,
-            created_at=index.created_at,
-            num_snippets=await self.repository.num_snippets_for_index(index.id),
-            source=source.uri,
-        )
-
-    async def list_indexes(self) -> list[IndexView]:
-        """List all available indexes with their details.
-
-        Returns:
-            A list of Index objects containing information about each index,
-            including file and snippet counts.
-
-        """
-        indexes = await self.repository.list_indexes()
-
-        # Transform database results into DTOs
-        indexes = [
-            IndexView(
-                id=index.id,
-                created_at=index.created_at,
-                updated_at=index.updated_at,
-                num_snippets=await self.repository.num_snippets_for_index(index.id)
-                or 0,
-                source=source.uri,
-            )
-            for index, source in indexes
-        ]
-
-        # Help Kodit by measuring how much people are using indexes
-        log_event(
-            "kodit.index.list",
-            {
-                "num_indexes": len(indexes),
-                "num_snippets": sum([index.num_snippets for index in indexes]),
-            },
-        )
-
-        return indexes
-
-    async def run(self, index_id: int) -> None:
-        """Run the indexing process for a specific index."""
-        log_event("kodit.index.run")
-
-        # Get and validate index
-        index = await self.repository.get_by_id(index_id)
-        if not index:
-            msg = f"Index not found: {index_id}"
-            raise ValueError(msg)
-
-        # Delete old snippets so we don't duplicate. In the future should probably check
-        # which files have changed and only change those.
-        await self.repository.delete_all_snippets(index.id)
-
-        # Create snippets for supported file types
-        self.log.info("Creating snippets for files", index_id=index.id)
-        await self._create_snippets(index.id)
-
-        snippets = await self.repository.get_all_snippets(index.id)
-
-        self.log.info("Creating keyword index")
-        with Spinner():
-            await self.keyword_search_provider.index(
-                [
-                    BM25Document(snippet_id=snippet.id, text=snippet.content)
-                    for snippet in snippets
-                ]
-            )
-
-        self.log.info("Creating semantic code index")
-        with tqdm(total=len(snippets), leave=False) as pbar:
-            async for result in self.code_search_service.index(
-                [
-                    VectorSearchRequest(snippet.id, snippet.content)
-                    for snippet in snippets
-                ]
-            ):
-                pbar.update(len(result))
-
-        self.log.info("Enriching snippets", num_snippets=len(snippets))
-        enriched_contents = []
-        with tqdm(total=len(snippets), leave=False) as pbar:
-            async for result in self.enrichment_service.enrich(
-                [
-                    EnrichmentRequest(snippet_id=snippet.id, text=snippet.content)
-                    for snippet in snippets
-                ]
-            ):
-                snippet = next(s for s in snippets if s.id == result.snippet_id)
-                if snippet:
-                    snippet.content = (
-                        result.text + "\n\n```\n" + snippet.content + "\n```"
-                    )
-                    await self.repository.add_snippet(snippet)
-                    enriched_contents.append(result)
-                pbar.update(1)
-
-        self.log.info("Creating semantic text index")
-        with tqdm(total=len(snippets), leave=False) as pbar:
-            async for result in self.text_search_service.index(
-                [
-                    VectorSearchRequest(snippet.id, snippet.content)
-                    for snippet in snippets
-                ]
-            ):
-                pbar.update(len(result))
-
-        # Update index timestamp
-        await self.repository.update_index_timestamp(index)
-
-    async def search(self, request: SearchRequest) -> list[SearchResult]:
-        """Search for relevant data."""
-        log_event("kodit.index.search")
-
-        fusion_list: list[list[FusionRequest]] = []
-        if request.keywords:
-            # Gather results for each keyword
-            result_ids: list[BM25Result] = []
-            for keyword in request.keywords:
-                results = await self.keyword_search_provider.retrieve(
-                    keyword, request.top_k
-                )
-                result_ids.extend(results)
-
-            fusion_list.append(
-                [FusionRequest(id=x.snippet_id, score=x.score) for x in result_ids]
-            )
-
-        # Compute embedding for semantic query
-        if request.code_query:
-            query_embedding = await self.code_search_service.retrieve(
-                request.code_query, top_k=request.top_k
-            )
-            fusion_list.append(
-                [FusionRequest(id=x.snippet_id, score=x.score) for x in query_embedding]
-            )
-
-        if request.text_query:
-            query_embedding = await self.text_search_service.retrieve(
-                request.text_query, top_k=request.top_k
-            )
-            fusion_list.append(
-                [FusionRequest(id=x.snippet_id, score=x.score) for x in query_embedding]
-            )
-
-        if len(fusion_list) == 0:
-            return []
-
-        # Combine all results together with RFF if required
-        final_results = reciprocal_rank_fusion(
-            rankings=fusion_list,
-            k=60,
-        )
-
-        # Only keep top_k results
-        final_results = final_results[: request.top_k]
-
-        # Get snippets from database (up to top_k)
-        search_results = await self.repository.list_snippets_by_ids(
-            [x.id for x in final_results]
-        )
-
-        return [
-            SearchResult(
-                id=snippet.id,
-                uri=file.uri,
-                content=snippet.content,
-                original_scores=fr.original_scores,
-            )
-            for (file, snippet), fr in zip(search_results, final_results, strict=True)
-        ]
-
-    async def _create_snippets(
-        self,
-        index_id: int,
-    ) -> None:
-        """Create snippets for supported files.
-
-        Args:
-            index: The index to create snippets for.
-            file_list: List of files to create snippets from.
-            existing_snippets_set: Set of file IDs that already have snippets.
-
-        """
-        files = await self.repository.files_for_index(index_id)
-        if not files:
-            self.log.warning("No files to create snippets for")
-            return
-
-        for file in tqdm(files, total=len(files), leave=False):
-            # Skip unsupported file types
-            if file.mime_type in MIME_BLACKLIST:
-                self.log.debug("Skipping mime type", mime_type=file.mime_type)
-                continue
-
-            # Create snippet from file content
-            try:
-                snippets = self.snippet_service.snippets_for_file(
-                    Path(file.cloned_path)
-                )
-            except ValueError as e:
-                self.log.debug("Skipping file", file=file.cloned_path, error=e)
-                continue
-
-            for snippet in snippets:
-                s = Snippet(
-                    index_id=index_id,
-                    file_id=file.id,
-                    content=snippet.text,
-                )
-                await self.repository.add_snippet(s)

kodit/snippets/__init__.py

@@ -1 +0,0 @@
-"""Extract method snippets from source code."""