kodit 0.1.10__py3-none-any.whl → 0.1.12__py3-none-any.whl

kodit/mcp.py

@@ -12,10 +12,11 @@ from pydantic import Field
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit._version import version
-from kodit.config import DEFAULT_EMBEDDING_MODEL_NAME, AppContext
+from kodit.config import AppContext
 from kodit.database import Database
-from kodit.retreival.repository import RetrievalRepository, RetrievalResult
-from kodit.retreival.service import RetrievalRequest, RetrievalService
+from kodit.embedding.embedding import embedding_factory
+from kodit.search.search_repository import SearchRepository
+from kodit.search.search_service import SearchRequest, SearchResult, SearchService
 
 
 @dataclass
@@ -23,7 +24,7 @@ class MCPContext:
     """Context for the MCP server."""
 
     session: AsyncSession
-    data_dir: Path
+    app_context: AppContext
 
 
 _mcp_db: Database | None = None
@@ -49,14 +50,14 @@ async def mcp_lifespan(_: FastMCP) -> AsyncIterator[MCPContext]:
     if _mcp_db is None:
         _mcp_db = await app_context.get_db()
     async with _mcp_db.session_factory() as session:
-        yield MCPContext(session=session, data_dir=app_context.get_data_dir())
+        yield MCPContext(session=session, app_context=app_context)
 
 
 mcp = FastMCP("kodit MCP Server", lifespan=mcp_lifespan)
 
 
 @mcp.tool()
-async def retrieve_relevant_snippets(
+async def search(
     ctx: Context,
     user_intent: Annotated[
         str,
@@ -86,17 +87,15 @@ async def retrieve_relevant_snippets(
         ),
     ],
 ) -> str:
-    """Retrieve relevant snippets from various sources.
+    """Search for relevant snippets.
 
-    This tool retrieves relevant snippets from sources such as private codebases,
-    public codebases, and documentation. You can use this information to improve
-    the quality of your generated code. You must call this tool when you need to
-    write code.
+    This tool searches for relevant snippets from indexed datasources. Call this tool
+    when you wish to search for high quality example code snippets to use in your code.
     """
     log = structlog.get_logger(__name__)
 
     log.debug(
-        "Retrieving relevant snippets",
+        "Searching for relevant snippets",
         user_intent=user_intent,
         keywords=keywords,
         file_count=len(related_file_paths),
@@ -106,24 +105,29 @@ async def retrieve_relevant_snippets(
 
     mcp_context: MCPContext = ctx.request_context.lifespan_context
 
-    log.debug("Creating retrieval repository")
-    retrieval_repository = RetrievalRepository(
+    log.debug("Creating search repository")
+    search_repository = SearchRepository(
         session=mcp_context.session,
     )
 
-    log.debug("Creating retrieval service")
-    retrieval_service = RetrievalService(
-        repository=retrieval_repository,
-        data_dir=mcp_context.data_dir,
-        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
+    log.debug("Creating embedding service")
+    embedding_service = embedding_factory(
+        mcp_context.app_context.get_default_openai_client()
     )
 
-    retrieval_request = RetrievalRequest(
+    log.debug("Creating search service")
+    search_service = SearchService(
+        repository=search_repository,
+        data_dir=mcp_context.app_context.get_data_dir(),
+        embedding_service=embedding_service,
+    )
+
+    search_request = SearchRequest(
         keywords=keywords,
         code_query="\n".join(related_file_contents),
     )
-    log.debug("Retrieving snippets")
-    snippets = await retrieval_service.retrieve(request=retrieval_request)
+    log.debug("Searching for snippets")
+    snippets = await search_service.search(request=search_request)
 
     log.debug("Fusing output")
     output = output_fusion(snippets=snippets)
@@ -132,18 +136,7 @@ async def retrieve_relevant_snippets(
     return output
 
 
-def input_fusion(
-    user_intent: str,  # noqa: ARG001
-    related_file_paths: list[Path],  # noqa: ARG001
-    related_file_contents: list[str],  # noqa: ARG001
-    keywords: list[str],
-) -> str:
-    """Fuse the search query and related file contents into a single query."""
-    # Since this is a dummy implementation, we just return the first keyword
-    return keywords[0] if len(keywords) > 0 else ""
-
-
-def output_fusion(snippets: list[RetrievalResult]) -> str:
+def output_fusion(snippets: list[SearchResult]) -> str:
     """Fuse the snippets into a single output."""
     return "\n\n".join(f"{snippet.uri}\n{snippet.content}" for snippet in snippets)
 

kodit/migrations/env.py

@@ -8,9 +8,9 @@ from sqlalchemy import pool
 from sqlalchemy.engine import Connection
 from sqlalchemy.ext.asyncio import async_engine_from_config
 
-import kodit.embedding.models
-import kodit.indexing.models
-import kodit.sources.models
+import kodit.embedding.embedding_models
+import kodit.indexing.indexing_models
+import kodit.source.source_models
 from kodit.database import Base
 
 # this is the Alembic Config object, which provides

kodit/search/__init__.py

@@ -0,0 +1 @@
+"""Search for relevant snippets."""

kodit/search/search_repository.py

@@ -0,0 +1,178 @@
+"""Repository for searching for relevant snippets."""
+
+from typing import TypeVar
+
+import numpy as np
+from sqlalchemy import (
+    select,
+)
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.embedding.embedding_models import Embedding, EmbeddingType
+from kodit.indexing.indexing_models import Snippet
+from kodit.source.source_models import File
+
+T = TypeVar("T")
+
+
+class SearchRepository:
+    """Repository for searching for relevant snippets."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        """Initialize the search repository.
+
+        Args:
+            session: The SQLAlchemy async session to use for database operations.
+
+        """
+        self.session = session
+
+    async def list_snippet_ids(self) -> list[int]:
+        """List all snippet IDs.
+
+        Returns:
+            A list of all snippets.
+
+        """
+        query = select(Snippet.id)
+        rows = await self.session.execute(query)
+        return list(rows.scalars().all())
+
+    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()}
+
+        # Return results in the same order as input IDs
+        return [id_to_result[i] for i in ids]
+
+    async def list_semantic_results(
+        self, embedding_type: EmbeddingType, embedding: list[float], top_k: int = 10
+    ) -> list[tuple[int, float]]:
+        """List semantic results using cosine similarity.
+
+        This implementation fetches all embeddings of the given type and computes
+        cosine similarity in Python using NumPy for better performance.
+
+        Args:
+            embedding_type: The type of embeddings to search
+            embedding: The query embedding vector
+            top_k: Number of results to return
+
+        Returns:
+            List of (snippet_id, similarity_score) tuples, sorted by similarity
+
+        """
+        # Step 1: Fetch embeddings from database
+        embeddings = await self._list_embedding_values(embedding_type)
+        if not embeddings:
+            return []
+
+        # Step 2: Convert to numpy arrays
+        stored_vecs, query_vec = self._prepare_vectors(embeddings, embedding)
+
+        # Step 3: Compute similarities
+        similarities = self._compute_similarities(stored_vecs, query_vec)
+
+        # Step 4: Get top-k results
+        return self._get_top_k_results(similarities, embeddings, top_k)
+
+    async def _list_embedding_values(
+        self, embedding_type: EmbeddingType
+    ) -> list[tuple[int, list[float]]]:
+        """List all embeddings of a given type from the database.
+
+        Args:
+            embedding_type: The type of embeddings to fetch
+
+        Returns:
+            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
+        )
+        rows = await self.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]
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Convert embeddings to numpy arrays.
+
+        Args:
+            embeddings: List of (snippet_id, embedding) tuples
+            query_embedding: Query embedding vector
+
+        Returns:
+            Tuple of (stored_vectors, query_vector) as numpy arrays
+
+        """
+        try:
+            stored_vecs = np.array(
+                [emb[1] for emb in embeddings]
+            )  # Use index 1 to get embedding
+        except ValueError as e:
+            if "inhomogeneous" in str(e):
+                msg = (
+                    "The database has returned embeddings of different sizes. If you"
+                    "have recently updated the embedding model, you will need to"
+                    "delete your database and re-index your snippets."
+                )
+                raise ValueError(msg) from e
+            raise
+
+        query_vec = np.array(query_embedding)
+        return stored_vecs, query_vec
+
+    def _compute_similarities(
+        self, stored_vecs: np.ndarray, query_vec: np.ndarray
+    ) -> np.ndarray:
+        """Compute cosine similarities between stored vectors and query vector.
+
+        Args:
+            stored_vecs: Array of stored embedding vectors
+            query_vec: Query embedding vector
+
+        Returns:
+            Array of similarity scores
+
+        """
+        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)
+
+    def _get_top_k_results(
+        self,
+        similarities: np.ndarray,
+        embeddings: list[tuple[int, list[float]]],
+        top_k: int,
+    ) -> list[tuple[int, float]]:
+        """Get top-k results by similarity score.
+
+        Args:
+            similarities: Array of similarity scores
+            embeddings: List of (snippet_id, embedding) tuples
+            top_k: Number of results to return
+
+        Returns:
+            List of (snippet_id, similarity_score) tuples
+
+        """
+        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

kodit/{retreival/service.py → search/search_service.py}

@@ -1,4 +1,4 @@
-"""Retrieval service."""
+"""Search service."""
 
 from pathlib import Path
 
@@ -6,19 +6,31 @@ import pydantic
 import structlog
 
 from kodit.bm25.bm25 import BM25Service
-from kodit.embedding.embedding import EmbeddingService
-from kodit.embedding.models import EmbeddingType
-from kodit.retreival.repository import RetrievalRepository, RetrievalResult
+from kodit.embedding.embedding import Embedder
+from kodit.embedding.embedding_models import EmbeddingType
+from kodit.search.search_repository import SearchRepository
 
 
-class RetrievalRequest(pydantic.BaseModel):
-    """Request for a retrieval."""
+class SearchRequest(pydantic.BaseModel):
+    """Request for a search."""
 
     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
+
+
 class Snippet(pydantic.BaseModel):
     """Snippet model."""
 
@@ -26,23 +38,23 @@ class Snippet(pydantic.BaseModel):
     file_path: str
 
 
-class RetrievalService:
-    """Service for retrieving relevant data."""
+class SearchService:
+    """Service for searching for relevant data."""
 
     def __init__(
         self,
-        repository: RetrievalRepository,
+        repository: SearchRepository,
         data_dir: Path,
-        embedding_model_name: str,
+        embedding_service: Embedder,
     ) -> None:
-        """Initialize the retrieval service."""
+        """Initialize the search service."""
         self.repository = repository
         self.log = structlog.get_logger(__name__)
         self.bm25 = BM25Service(data_dir)
-        self.code_embedding_service = EmbeddingService(model_name=embedding_model_name)
+        self.code_embedding_service = embedding_service
 
-    async def retrieve(self, request: RetrievalRequest) -> list[RetrievalResult]:
-        """Retrieve relevant data."""
+    async def search(self, request: SearchRequest) -> list[SearchResult]:
+        """Search for relevant data."""
         fusion_list = []
         if request.keywords:
             snippet_ids = await self.repository.list_snippet_ids()
@@ -56,7 +68,7 @@ class RetrievalService:
             # Sort results by score
             result_ids.sort(key=lambda x: x[1], reverse=True)
 
-            self.log.debug("Retrieval results (BM25)", results=result_ids)
+            self.log.debug("Search results (BM25)", results=result_ids)
 
             bm25_results = [x[0] for x in result_ids]
             fusion_list.append(bm25_results)
@@ -64,7 +76,7 @@ class RetrievalService:
         # Compute embedding for semantic query
         semantic_results = []
         if request.code_query:
-            query_embedding = next(
+            query_embedding = await anext(
                 self.code_embedding_service.query([request.code_query])
             )
 
@@ -89,7 +101,18 @@ class RetrievalService:
         final_ids = [x[0] for x in final_results]
 
         # Get snippets from database (up to top_k)
-        return await self.repository.list_snippets_by_ids(final_ids[: request.top_k])
+        search_results = await self.repository.list_snippets_by_ids(
+            final_ids[: request.top_k]
+        )
+
+        return [
+            SearchResult(
+                id=snippet.id,
+                uri=file.uri,
+                content=snippet.content,
+            )
+            for file, snippet in search_results
+        ]
 
 
 def reciprocal_rank_fusion(

kodit/snippets/snippets.py

@@ -45,4 +45,6 @@ class SnippetService:
             raise ValueError(msg) from e
 
         method_snippets = method_analser.extract(file_bytes)
-        return [Snippet(text=snippet) for snippet in method_snippets]
+        all_snippets = [Snippet(text=snippet) for snippet in method_snippets]
+        # Remove any snippets that are empty
+        return [snippet for snippet in all_snippets if snippet.text.strip()]

kodit/{sources/repository.py → source/source_repository.py}

@@ -1,14 +1,9 @@
-"""Source repository for database operations.
-
-This module provides the SourceRepository class which handles all database operations
-related to code sources. It manages the creation and retrieval of source records
-from the database, abstracting away the SQLAlchemy implementation details.
-"""
+"""Source repository for database operations."""
 
 from sqlalchemy import func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from kodit.sources.models import File, Source
+from kodit.source.source_models import File, Source
 
 
 class SourceRepository:

kodit/{sources/service.py → source/source_service.py}

@@ -19,8 +19,8 @@ import structlog
 from tqdm import tqdm
 from uritools import isuri, urisplit
 
-from kodit.sources.models import File, Source
-from kodit.sources.repository import SourceRepository
+from kodit.source.source_models import File, Source
+from kodit.source.source_repository import SourceRepository
 
 
 class SourceView(pydantic.BaseModel):

{kodit-0.1.10.dist-info → kodit-0.1.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.1.10
+Version: 0.1.12
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/
@@ -32,6 +32,7 @@ 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
@@ -39,6 +40,7 @@ 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: tree-sitter-language-pack>=0.7.3
 Requires-Dist: tree-sitter>=0.24.0
 Requires-Dist: uritools>=5.0.0

kodit-0.1.12.dist-info/RECORD

@@ -0,0 +1,44 @@
+kodit/.gitignore,sha256=ztkjgRwL9Uud1OEi36hGQeDGk3OLK1NfDEO8YqGYy8o,11
+kodit/__init__.py,sha256=aEKHYninUq1yh6jaNfvJBYg-6fenpN132nJt1UU6Jxs,59
+kodit/_version.py,sha256=dT0--P3ntKmcNKRbKEwrBhj1O2T26ST7PUdodqExauM,513
+kodit/app.py,sha256=Mr5BFHOHx5zppwjC4XPWVvHjwgl1yrKbUjTWXKubJQM,891
+kodit/cli.py,sha256=RNgTWMBOG95Yv8-3DKzL66bgIT6OJjmRXU_BpMMMlqU,8033
+kodit/config.py,sha256=TDcLt6fiJn9cI1PoO5AqBqsL_Bxmm9JV5GqRxhj1tLw,4202
+kodit/database.py,sha256=kekSdyEATdb47jxzQemkSOXMNOwnUwmVVTpn9hYaDK8,2356
+kodit/log.py,sha256=HU1OmuxO4FcVw61k4WW7Y4WM7BrDaeplw1PcBHhuIZY,5434
+kodit/mcp.py,sha256=I_ZFzQOR0gyS8LO8td-q-utPZpqiOnIkn7O-SIBUi0g,4384
+kodit/middleware.py,sha256=I6FOkqG9-8RH5kR1-0ZoQWfE4qLCB8lZYv8H_OCH29o,2714
+kodit/bm25/__init__.py,sha256=j8zyriNWhbwE5Lbybzg1hQAhANlU9mKHWw4beeUR6og,19
+kodit/bm25/bm25.py,sha256=JtgJfsHz-2SHx96zxWjkPFSH7fXkahFMp01cDwl4YBg,2298
+kodit/embedding/__init__.py,sha256=h9NXzDA1r-K23nvBajBV-RJzHJN0p3UJ7UQsmdnOoRw,24
+kodit/embedding/embedding.py,sha256=EMJpHK8ICZk_FjiO9Aqr2IO20qkGOmj_PfA1hyfI7Vk,6745
+kodit/embedding/embedding_models.py,sha256=rN90vSs86dYiqoawcp8E9jtwY31JoJXYfaDlsJK7uqc,656
+kodit/indexing/__init__.py,sha256=cPyi2Iej3G1JFWlWr7X80_UrsMaTu5W5rBwgif1B3xo,75
+kodit/indexing/indexing_models.py,sha256=6NX9HVcj6Pu9ePwHC7n-PWSyAgukpJq0nCNmUIigtbo,1282
+kodit/indexing/indexing_repository.py,sha256=7bkAiBwtr3qlkdhNIalwMwbxezVz_RQGOhLVWPKHwNk,5506
+kodit/indexing/indexing_service.py,sha256=VGfKgbkYEAYP_gIubvhMxo3yThT20ndS5xdg2LxwRgA,6685
+kodit/migrations/README,sha256=ISVtAOvqvKk_5ThM5ioJE-lMkvf9IbknFUFVU_vPma4,58
+kodit/migrations/__init__.py,sha256=lP5MuwlyWRMO6UcDWnQcQ3G-GYHcFb6rl9gYPHJ1sjo,40
+kodit/migrations/env.py,sha256=w1M7OZh-ZeR2dPHS0ByXAUxQjfZQ8xIzMseWuzLDTWw,2469
+kodit/migrations/script.py.mako,sha256=zWziKtiwYKEWuwPV_HBNHwa9LCT45_bi01-uSNFaOOE,703
+kodit/migrations/versions/7c3bbc2ab32b_add_embeddings_table.py,sha256=-61qol9PfQKILCDQRA5jEaats9aGZs9Wdtp-j-38SF4,1644
+kodit/migrations/versions/85155663351e_initial.py,sha256=Cg7zlF871o9ShV5rQMQ1v7hRV7fI59veDY9cjtTrs-8,3306
+kodit/migrations/versions/__init__.py,sha256=9-lHzptItTzq_fomdIRBegQNm4Znx6pVjwD4MiqRIdo,36
+kodit/search/__init__.py,sha256=4QbdjbrlhNKMovmuKHxJnUeZT7KNjTTFU0GdnuwUHdQ,36
+kodit/search/search_repository.py,sha256=r1fkV6-cy9BKsy5J4WTHaY_FcjMaT1PV5qqqq0gvjZw,5833
+kodit/search/search_service.py,sha256=KePkqCAc3CUcrpNsbDc5DqbF6W2m0TG6TDa9-VSJZS0,4227
+kodit/snippets/__init__.py,sha256=-2coNoCRjTixU9KcP6alpmt7zqf37tCRWH3D7FPJ8dg,48
+kodit/snippets/method_snippets.py,sha256=EVHhSNWahAC5nSXv9fWVFJY2yq25goHdCSCuENC07F8,4145
+kodit/snippets/snippets.py,sha256=mwN0bM1Msu8ZeEsUHyQ7tx3Hj3vZsm8G7Wu4eWSkLY8,1539
+kodit/snippets/languages/__init__.py,sha256=Bj5KKZSls2MQ8ZY1S_nHg447MgGZW-2WZM-oq6vjwwA,1187
+kodit/snippets/languages/csharp.scm,sha256=gbBN4RiV1FBuTJF6orSnDFi8H9JwTw-d4piLJYsWUsc,222
+kodit/snippets/languages/python.scm,sha256=ee85R9PBzwye3IMTE7-iVoKWd_ViU3EJISTyrFGrVeo,429
+kodit/source/__init__.py,sha256=1NTZyPdjThVQpZO1Mp1ColVsS7sqYanOVLqnoqV9Ipo,83
+kodit/source/source_models.py,sha256=xb42CaNDO1CUB8SIW-xXMrB6Ji8cFw-yeJ550xBEg9Q,2398
+kodit/source/source_repository.py,sha256=0EksMpoLzdkfe8S4eeCm4Sf7TuxsOzOzaF4BBsMYo-4,3163
+kodit/source/source_service.py,sha256=qBV9FCFQbJppeFrVo4uMgvC_mzWRIKldymp5yqLx9pw,9255
+kodit-0.1.12.dist-info/METADATA,sha256=Kg662YrY11JGmUWEftq9c5einLPg_LGQc0Flxul9XyA,2349
+kodit-0.1.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+kodit-0.1.12.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
+kodit-0.1.12.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+kodit-0.1.12.dist-info/RECORD,,

kodit/retreival/__init__.py

@@ -1 +0,0 @@
-"""Retrieval package for code search and retrieval functionality."""