kodit 0.1.15__py3-none-any.whl → 0.1.17__py3-none-any.whl

kodit/indexing/fusion.py

@@ -0,0 +1,67 @@
+"""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_repository.py

@@ -10,7 +10,6 @@ from typing import TypeVar
 
 from sqlalchemy import delete, func, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm.exc import MultipleResultsFound
 
 from kodit.embedding.embedding_models import Embedding
 from kodit.indexing.indexing_models import Index, Snippet
@@ -125,34 +124,15 @@ class IndexRepository:
         index.updated_at = datetime.now(UTC)
         await self.session.commit()
 
-    async def add_snippet_or_update_content(self, snippet: Snippet) -> None:
+    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.
 
         """
-        query = select(Snippet).where(
-            Snippet.file_id == snippet.file_id,
-            Snippet.index_id == snippet.index_id,
-        )
-        result = await self.session.execute(query)
-        try:
-            existing_snippet = result.scalar_one_or_none()
-
-            if existing_snippet:
-                existing_snippet.content = snippet.content
-            else:
-                self.session.add(snippet)
-
-            await self.session.commit()
-        except MultipleResultsFound as e:
-            msg = (
-                f"Multiple snippets found for file_id {snippet.file_id}, this "
-                "shouldn't happen. "
-                "Please report this as a bug then delete your index and start again."
-            )
-            raise ValueError(msg) from e
+        self.session.add(snippet)
+        await self.session.commit()
 
     async def delete_all_snippets(self, index_id: int) -> None:
         """Delete all snippets for an index.
@@ -161,6 +141,15 @@ class IndexRepository:
             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()
@@ -196,3 +185,32 @@ class IndexRepository:
         """
         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

@@ -13,11 +13,17 @@ import pydantic
 import structlog
 from tqdm.asyncio import tqdm
 
-from kodit.bm25.keyword_search_service import BM25Document, KeywordSearchProvider
+from kodit.bm25.keyword_search_service import (
+    BM25Document,
+    BM25Result,
+    KeywordSearchProvider,
+)
 from kodit.embedding.vector_search_service import (
     VectorSearchRequest,
     VectorSearchService,
 )
+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.snippets.snippets import SnippetService
@@ -42,6 +48,28 @@ class IndexView(pydantic.BaseModel):
     num_snippets: int | None = None
 
 
+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.
 
@@ -50,12 +78,14 @@ class IndexService:
     IndexRepository), and provides a clean API for index management.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         repository: IndexRepository,
         source_service: SourceService,
         keyword_search_provider: KeywordSearchProvider,
-        vector_search_service: VectorSearchService,
+        code_search_service: VectorSearchService,
+        text_search_service: VectorSearchService,
+        enrichment_service: EnrichmentService,
     ) -> None:
         """Initialize the index service.
 
@@ -69,7 +99,9 @@ class IndexService:
         self.snippet_service = SnippetService()
         self.log = structlog.get_logger(__name__)
         self.keyword_search_provider = keyword_search_provider
-        self.code_search_service = vector_search_service
+        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.
@@ -129,10 +161,15 @@ class IndexService:
             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
-        await self._create_snippets(index_id)
+        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)
+        snippets = await self.repository.get_all_snippets(index.id)
 
         self.log.info("Creating keyword index")
         with Spinner():
@@ -152,9 +189,93 @@ class IndexService:
                 ]
             )
 
+        self.log.info("Enriching snippets", num_snippets=len(snippets))
+        enriched_contents = await self.enrichment_service.enrich(
+            [snippet.content for snippet in snippets]
+        )
+
+        self.log.info("Creating semantic text index")
+        with Spinner():
+            await self.text_search_service.index(
+                [
+                    VectorSearchRequest(snippet.id, enriched_content)
+                    for snippet, enriched_content in zip(
+                        snippets, enriched_contents, strict=True
+                    )
+                ]
+            )
+            # Add the enriched text back to the snippets and write to the database
+            for snippet, enriched_content in zip(
+                snippets, enriched_contents, strict=True
+            ):
+                snippet.content = (
+                    enriched_content + "\n\n```\n" + snippet.content + "\n```"
+                )
+                await self.repository.add_snippet(snippet)
+
         # Update index timestamp
         await self.repository.update_index_timestamp(index)
 
+    async def search(self, request: SearchRequest) -> list[SearchResult]:
+        """Search for relevant data."""
+        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,
@@ -168,7 +289,6 @@ class IndexService:
 
         """
         files = await self.repository.files_for_index(index_id)
-        self.log.info("Creating snippets for files", index_id=index_id)
         for file in tqdm(files, total=len(files), leave=False):
             # Skip unsupported file types
             if file.mime_type in MIME_BLACKLIST:
@@ -190,4 +310,4 @@ class IndexService:
                     file_id=file.id,
                     content=snippet.text,
                 )
-                await self.repository.add_snippet_or_update_content(s)
+                await self.repository.add_snippet(s)

kodit/mcp.py

@@ -16,8 +16,11 @@ from kodit.bm25.keyword_search_factory import keyword_search_factory
 from kodit.config import AppContext
 from kodit.database import Database
 from kodit.embedding.embedding_factory import embedding_factory
-from kodit.search.search_repository import SearchRepository
-from kodit.search.search_service import SearchRequest, SearchResult, SearchService
+from kodit.enrichment.enrichment_factory import enrichment_factory
+from kodit.indexing.indexing_repository import IndexRepository
+from kodit.indexing.indexing_service import IndexService, SearchRequest, SearchResult
+from kodit.source.source_repository import SourceRepository
+from kodit.source.source_service import SourceService
 
 
 @dataclass
@@ -123,32 +126,38 @@ async def search(
 
     mcp_context: MCPContext = ctx.request_context.lifespan_context
 
-    log.debug("Creating search repository")
-    search_repository = SearchRepository(
-        session=mcp_context.session,
+    source_repository = SourceRepository(mcp_context.session)
+    source_service = SourceService(
+        mcp_context.app_context.get_clone_dir(), source_repository
     )
-
-    log.debug("Creating embedding service")
-    embedding_service = embedding_factory(
-        app_context=mcp_context.app_context, session=mcp_context.session
-    )
-
-    log.debug("Creating search service")
-    search_service = SearchService(
-        repository=search_repository,
+    repository = IndexRepository(mcp_context.session)
+    service = IndexService(
+        repository=repository,
+        source_service=source_service,
         keyword_search_provider=keyword_search_factory(
+            mcp_context.app_context, mcp_context.session
+        ),
+        code_search_service=embedding_factory(
+            task_name="code",
             app_context=mcp_context.app_context,
             session=mcp_context.session,
         ),
-        embedding_service=embedding_service,
+        text_search_service=embedding_factory(
+            task_name="text",
+            app_context=mcp_context.app_context,
+            session=mcp_context.session,
+        ),
+        enrichment_service=enrichment_factory(mcp_context.app_context),
     )
 
     search_request = SearchRequest(
         keywords=keywords,
         code_query="\n".join(related_file_contents),
+        text_query=user_intent,
     )
+
     log.debug("Searching for snippets")
-    snippets = await search_service.search(request=search_request)
+    snippets = await service.search(request=search_request)
 
     log.debug("Fusing output")
     output = output_fusion(snippets=snippets)

kodit/migrations/versions/c3f5137d30f5_index_all_the_things.py

@@ -0,0 +1,44 @@
+# ruff: noqa
+"""index all the things
+
+Revision ID: c3f5137d30f5
+Revises: 7c3bbc2ab32b
+Create Date: 2025-06-05 17:17:32.440740
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'c3f5137d30f5'
+down_revision: Union[str, None] = '7c3bbc2ab32b'
+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_index(op.f('ix_files_cloned_path'), 'files', ['cloned_path'], unique=False)
+    op.create_index(op.f('ix_files_mime_type'), 'files', ['mime_type'], unique=False)
+    op.create_index(op.f('ix_files_uri'), 'files', ['uri'], unique=False)
+    op.create_index(op.f('ix_snippets_file_id'), 'snippets', ['file_id'], unique=False)
+    op.create_index(op.f('ix_snippets_index_id'), 'snippets', ['index_id'], unique=False)
+    op.create_index(op.f('ix_sources_cloned_path'), 'sources', ['cloned_path'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_sources_cloned_path'), table_name='sources')
+    op.drop_index(op.f('ix_snippets_index_id'), table_name='snippets')
+    op.drop_index(op.f('ix_snippets_file_id'), table_name='snippets')
+    op.drop_index(op.f('ix_files_uri'), table_name='files')
+    op.drop_index(op.f('ix_files_mime_type'), table_name='files')
+    op.drop_index(op.f('ix_files_cloned_path'), table_name='files')
+    # ### end Alembic commands ###

kodit/snippets/languages/go.scm

@@ -0,0 +1,26 @@
+(function_declaration
+  name: (identifier) @function.name
+  body: (block) @function.body
+) @function.def
+
+(method_declaration
+  name: (field_identifier) @method.name
+  body: (block) @method.body
+) @method.def
+
+(import_declaration
+  (import_spec
+    path: (interpreted_string_literal) @import.name
+  )
+) @import.statement
+
+(identifier) @ident
+
+(parameter_declaration
+  name: (identifier) @param.name
+)
+
+(package_clause "package" (package_identifier) @name.definition.module)
+
+;; Exclude comments from being captured
+(comment) @comment

kodit/source/source_models.py

@@ -31,7 +31,7 @@ class Source(Base, CommonMixin):
 
     __tablename__ = "sources"
     uri: Mapped[str] = mapped_column(String(1024), index=True, unique=True)
-    cloned_path: Mapped[str] = mapped_column(String(1024))
+    cloned_path: Mapped[str] = mapped_column(String(1024), index=True)
 
     def __init__(self, uri: str, cloned_path: str) -> None:
         """Initialize a new Source instance for typing purposes."""
@@ -46,9 +46,9 @@ class File(Base, CommonMixin):
     __tablename__ = "files"
 
     source_id: Mapped[int] = mapped_column(ForeignKey("sources.id"))
-    mime_type: Mapped[str] = mapped_column(String(255), default="")
-    uri: Mapped[str] = mapped_column(String(1024), default="")
-    cloned_path: Mapped[str] = mapped_column(String(1024))
+    mime_type: Mapped[str] = mapped_column(String(255), default="", index=True)
+    uri: Mapped[str] = mapped_column(String(1024), default="", index=True)
+    cloned_path: Mapped[str] = mapped_column(String(1024), index=True)
     sha256: Mapped[str] = mapped_column(String(64), default="", index=True)
     size_bytes: Mapped[int] = mapped_column(Integer, default=0)
 

kodit-0.1.17.dist-info/METADATA

@@ -0,0 +1,152 @@
+Metadata-Version: 2.4
+Name: kodit
+Version: 0.1.17
+Summary: Code indexing for better AI code generation
+Project-URL: Homepage, https://docs.helixml.tech/kodit/
+Project-URL: Documentation, https://docs.helixml.tech/kodit/
+Project-URL: Repository, https://github.com/helixml/kodit.git
+Project-URL: Issues, https://github.com/helixml/kodit/issues
+Project-URL: Changelog, https://github.com/helixml/kodit/releases
+Author-email: "Helix.ML" <founders@helix.ml>
+Maintainer-email: "Helix.ML" <founders@helix.ml>
+License-Expression: Apache-2.0
+License-File: LICENSE
+Keywords: ai,indexing,mcp,rag
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Software Development :: Code Generators
+Requires-Python: >=3.12
+Requires-Dist: aiofiles>=24.1.0
+Requires-Dist: aiosqlite>=0.20.0
+Requires-Dist: alembic>=1.15.2
+Requires-Dist: asgi-correlation-id>=4.3.4
+Requires-Dist: asyncpg>=0.30.0
+Requires-Dist: better-exceptions>=0.3.3
+Requires-Dist: bm25s[core]>=0.2.12
+Requires-Dist: click>=8.1.8
+Requires-Dist: colorama>=0.4.6
+Requires-Dist: dotenv>=0.9.9
+Requires-Dist: fastapi[standard]>=0.115.12
+Requires-Dist: fastmcp>=2.3.3
+Requires-Dist: gitpython>=3.1.44
+Requires-Dist: hf-xet>=1.1.2
+Requires-Dist: httpx-retries>=0.3.2
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: openai>=1.82.0
+Requires-Dist: posthog>=4.0.1
+Requires-Dist: pydantic-settings>=2.9.1
+Requires-Dist: pytable-formatter>=0.1.1
+Requires-Dist: sentence-transformers>=4.1.0
+Requires-Dist: sqlalchemy[asyncio]>=2.0.40
+Requires-Dist: structlog>=25.3.0
+Requires-Dist: tdqm>=0.0.1
+Requires-Dist: tiktoken>=0.9.0
+Requires-Dist: transformers>=4.51.3
+Requires-Dist: tree-sitter-language-pack>=0.7.3
+Requires-Dist: tree-sitter>=0.24.0
+Requires-Dist: uritools>=5.0.0
+Description-Content-Type: text/markdown
+
+<p align="center">
+    <a href="https://docs.helix.ml/kodit/"><img src="https://docs.helix.ml/images/helix-kodit-logo.png" alt="Helix Kodit Logo" width="300"></a>
+</p>
+
+<h1 align="center">
+Kodit: A Code Indexing MCP Server
+</h1>
+
+<p align="center">
+Kodit connects your AI coding assistant to external codebases to provide accurate and up-to-date snippets of code.
+</p>
+
+<div align="center">
+
+[![Documentation](https://img.shields.io/badge/Documentation-6B46C1?style=for-the-badge&logo=readthedocs&logoColor=white)](https://docs.helix.ml/kodit/)
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg?style=for-the-badge)](./LICENSE)
+[![Discussions](https://img.shields.io/badge/Discussions-181717?style=for-the-badge&logo=github&logoColor=white)](https://github.com/helixml/kodit/discussions)
+
+</div>
+
+**Helix Kodit** is an **MCP server** that connects your AI coding assistant to external codebases. It can:
+
+- Improve your AI-assisted code by providing canonical examples direct from the source
+- Index local and public codebases
+- Integrates with any AI coding assistant via MCP
+- Search using keyword and semantic search
+- Integrate with any OpenAI-compatible or custom API/model
+
+If you're an engineer working with AI-powered coding assistants, Kodit helps by
+providing relevant and up-to-date examples of your task so that LLMs make less mistakes
+and produce fewer hallucinations.
+
+## ✨ Features
+
+### Codebase Indexing
+
+Kodit connects to a variety of local and remote codebases to build an index of your
+code. This index is used to build a snippet library, ready for ingestion into an LLM.
+
+- Index local directories and public Git repositories
+- Build comprehensive snippet libraries for LLM ingestion
+- Support for multiple codebase types and languages
+- Efficient indexing and search capabilities
+
+### MCP Server
+
+Relevant snippets are exposed to an AI coding assistant via an MCP server. This allows
+the assistant to request relevant snippets by providing keywords, code, and semantic
+intent. Kodit has been tested to work well with:
+
+- Seamless integration with popular AI coding assistants
+- Tested and verified with:
+  - [Cursor](https://docs.helix.ml/kodit/#integration-with-cursor)
+  - [Cline](https://docs.helix.ml/kodit/#integration-with-cline)
+- Please contribute more instructions! ... any other assistant is likely to work ...
+
+### Enterprise Ready
+
+Out of the box, Kodit works with a local SQLite database and very small, local models.
+But enterprises can scale out with performant databases and dedicated models. Everything
+can even run securely, privately, with on-premise LLM platforms like
+[Helix](https://helix.ml).
+
+Supported databases:
+
+- SQLite
+- [Vectorchord](https://github.com/tensorchord/VectorChord)
+
+Supported providers:
+
+- Local (which uses tiny CPU-only open-source models)
+- OpenAI
+- Secure, private LLM enclave with [Helix](https://helix.ml).
+- Any other OpenAI compatible API
+
+## 🚀 Quick Start
+
+1. [Install Kodit](https://docs.helix.ml/kodit/#installation)
+2. [Index codebases](https://docs.helix.ml/kodit/#quick-start)
+3. [Integrate with your coding assistant](https://docs.helix.ml/kodit/#integrating-kodit-with-coding-assistants)
+
+### Documentation
+
+- [Installation Guide](https://docs.helix.ml/kodit/#installation)
+- [Usage Guide](https://docs.helix.ml/kodit/#quick-start)
+- [Connecting to Kodit](https://docs.helix.ml/kodit/#integrating-kodit-with-coding-assistants)
+- [Configuration Options](https://docs.helix.ml/kodit/#configuring-kodit)
+- [Contribution Guidelines](.github/CONTRIBUTING.md)
+
+## Roadmap
+
+The roadmap is currently maintained as a [Github Project](https://github.com/orgs/helixml/projects/4).
+
+## 💬 Support
+
+For commercial support, please contact [Helix.ML](founders@helix.ml). To ask a question,
+please [open a discussion](https://github.com/helixml/kodit/discussions).
+
+## License
+
+[Apache 2.0 © 2025 HelixML, Inc.](./LICENSE)