kodit 0.1.14__py3-none-any.whl → 0.1.16__py3-none-any.whl

kodit/enrichment/enrichment_provider/enrichment_provider.py

@@ -0,0 +1,16 @@
+"""Enrichment provider."""
+
+from abc import ABC, abstractmethod
+
+ENRICHMENT_SYSTEM_PROMPT = """
+You are a professional software developer. You will be given a snippet of code.
+Please provide a concise explanation of the code.
+"""
+
+
+class EnrichmentProvider(ABC):
+    """Enrichment provider."""
+
+    @abstractmethod
+    async def enrich(self, data: list[str]) -> list[str]:
+        """Enrich a list of strings."""

kodit/enrichment/enrichment_provider/local_enrichment_provider.py

@@ -0,0 +1,63 @@
+"""Local embedding service."""
+
+import os
+
+import structlog
+from transformers.models.auto.modeling_auto import AutoModelForCausalLM
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+
+from kodit.enrichment.enrichment_provider.enrichment_provider import (
+    ENRICHMENT_SYSTEM_PROMPT,
+    EnrichmentProvider,
+)
+
+
+class LocalEnrichmentProvider(EnrichmentProvider):
+    """Local embedder."""
+
+    def __init__(self, model_name: str = "Qwen/Qwen3-0.6B") -> None:
+        """Initialize the local enrichment provider."""
+        self.log = structlog.get_logger(__name__)
+        self.model_name = model_name
+        self.model = None
+        self.tokenizer = None
+
+    async def enrich(self, data: list[str]) -> list[str]:
+        """Enrich a list of strings."""
+        if self.tokenizer is None:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        if self.model is None:
+            os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                torch_dtype="auto",
+                trust_remote_code=True,
+            )
+
+        results = []
+        for snippet in data:
+            # prepare the model input
+            messages = [
+                {"role": "system", "content": ENRICHMENT_SYSTEM_PROMPT},
+                {"role": "user", "content": snippet},
+            ]
+            text = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=False,
+            )
+            model_inputs = self.tokenizer([text], return_tensors="pt").to(
+                self.model.device
+            )
+
+            # conduct text completion
+            generated_ids = self.model.generate(**model_inputs, max_new_tokens=32768)
+            output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()
+            content = self.tokenizer.decode(output_ids, skip_special_tokens=True).strip(
+                "\n"
+            )
+
+            results.append(content)
+
+        return results

kodit/enrichment/enrichment_provider/openai_enrichment_provider.py

@@ -0,0 +1,77 @@
+"""OpenAI embedding service."""
+
+import asyncio
+
+import structlog
+import tiktoken
+from openai import AsyncOpenAI
+from tqdm import tqdm
+
+from kodit.enrichment.enrichment_provider.enrichment_provider import (
+    ENRICHMENT_SYSTEM_PROMPT,
+    EnrichmentProvider,
+)
+
+OPENAI_NUM_PARALLEL_TASKS = 10
+
+
+class OpenAIEnrichmentProvider(EnrichmentProvider):
+    """OpenAI enrichment provider."""
+
+    def __init__(
+        self,
+        openai_client: AsyncOpenAI,
+        model_name: str = "gpt-4o-mini",
+    ) -> None:
+        """Initialize the OpenAI enrichment provider."""
+        self.log = structlog.get_logger(__name__)
+        self.openai_client = openai_client
+        self.model_name = model_name
+        self.encoding = tiktoken.encoding_for_model(model_name)
+
+    async def enrich(self, data: list[str]) -> list[str]:
+        """Enrich a list of documents."""
+        # Process batches in parallel with a semaphore to limit concurrent requests
+        sem = asyncio.Semaphore(OPENAI_NUM_PARALLEL_TASKS)
+
+        # Create a list of tuples with a temporary id for each snippet
+        # We need to do this so that we can return the results in the same order as the
+        # input data
+        input_data = [(i, snippet) for i, snippet in enumerate(data)]
+
+        async def process_data(data: tuple[int, str]) -> tuple[int, str]:
+            snippet_id, snippet = data
+            if not snippet:
+                return snippet_id, ""
+            async with sem:
+                try:
+                    response = await self.openai_client.chat.completions.create(
+                        model=self.model_name,
+                        messages=[
+                            {
+                                "role": "system",
+                                "content": ENRICHMENT_SYSTEM_PROMPT,
+                            },
+                            {"role": "user", "content": snippet},
+                        ],
+                    )
+                    return snippet_id, response.choices[0].message.content or ""
+                except Exception as e:
+                    self.log.exception("Error enriching data", error=str(e))
+                    return snippet_id, ""
+
+        # Create tasks for all data
+        tasks = [process_data(snippet) for snippet in input_data]
+
+        # Process all data and yield results as they complete
+        results: list[tuple[int, str]] = []
+        for task in tqdm(
+            asyncio.as_completed(tasks),
+            total=len(tasks),
+            leave=False,
+        ):
+            result = await task
+            results.append(result)
+
+        # Output in the same order as the input data
+        return [result for _, result in sorted(results, key=lambda x: x[0])]

kodit/enrichment/enrichment_service.py

@@ -0,0 +1,33 @@
+"""Enrichment service."""
+
+from abc import ABC, abstractmethod
+
+from kodit.enrichment.enrichment_provider.enrichment_provider import EnrichmentProvider
+
+
+class EnrichmentService(ABC):
+    """Enrichment service."""
+
+    @abstractmethod
+    async def enrich(self, data: list[str]) -> list[str]:
+        """Enrich a list of strings."""
+
+
+class NullEnrichmentService(EnrichmentService):
+    """Null enrichment service."""
+
+    async def enrich(self, data: list[str]) -> list[str]:
+        """Enrich a list of strings."""
+        return [""] * len(data)
+
+
+class LLMEnrichmentService(EnrichmentService):
+    """Enrichment service using an LLM."""
+
+    def __init__(self, enrichment_provider: EnrichmentProvider) -> None:
+        """Initialize the enrichment service."""
+        self.enrichment_provider = enrichment_provider
+
+    async def enrich(self, data: list[str]) -> list[str]:
+        """Enrich a list of strings."""
+        return await self.enrichment_provider.enrich(data)

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,6 +10,7 @@ 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
@@ -124,15 +125,34 @@ class IndexRepository:
         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.
+    async def add_snippet_or_update_content(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()
+        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
 
     async def delete_all_snippets(self, index_id: int) -> None:
         """Delete all snippets for an index.
@@ -176,3 +196,23 @@ 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()}
+
+        # Return results in the same order as input IDs
+        return [id_to_result[i] for i in ids]

kodit/indexing/indexing_service.py

@@ -13,13 +13,22 @@ import pydantic
 import structlog
 from tqdm.asyncio import tqdm
 
-from kodit.bm25.bm25 import BM25Service
-from kodit.embedding.embedding import Embedder, EmbeddingInput
-from kodit.embedding.embedding_models import Embedding, EmbeddingType
+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
 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"]
@@ -39,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.
 
@@ -47,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,
-        data_dir: Path,
-        embedding_service: Embedder,
+        keyword_search_provider: KeywordSearchProvider,
+        code_search_service: VectorSearchService,
+        text_search_service: VectorSearchService,
+        enrichment_service: EnrichmentService,
     ) -> None:
         """Initialize the index service.
 
@@ -65,8 +98,10 @@ class IndexService:
         self.source_service = source_service
         self.snippet_service = SnippetService()
         self.log = structlog.get_logger(__name__)
-        self.bm25 = BM25Service(data_dir)
-        self.code_embedding_service = embedding_service
+        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.
@@ -126,41 +161,116 @@ class IndexService:
             msg = f"Index not found: {index_id}"
             raise ValueError(msg)
 
-        # First delete all old snippets, if they exist
-        await self.repository.delete_all_snippets(index_id)
-
         # Create snippets for supported file types
         await self._create_snippets(index_id)
 
         snippets = await self.repository.get_all_snippets(index_id)
 
         self.log.info("Creating keyword index")
-        self.bm25.index(
-            [
-                snippet.content
-                for snippet in tqdm(snippets, total=len(snippets), leave=False)
-            ]
-        )
+        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")
-        async for e in tqdm(
-            self.code_embedding_service.embed(
-                [EmbeddingInput(snippet.id, snippet.content) for snippet in snippets]
-            ),
-            total=len(snippets),
-            leave=False,
-        ):
-            await self.repository.add_embedding(
-                Embedding(
-                    snippet_id=e.id,
-                    embedding=e.embedding,
-                    type=EmbeddingType.CODE,
-                )
+        with Spinner():
+            await self.code_search_service.index(
+                [
+                    VectorSearchRequest(snippet.id, snippet.content)
+                    for snippet in snippets
+                ]
             )
 
+        self.log.info("Enriching 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_or_update_content(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,
@@ -174,6 +284,7 @@ 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:
@@ -195,4 +306,4 @@ class IndexService:
                     file_id=file.id,
                     content=snippet.text,
                 )
-                await self.repository.add_snippet(s)
+                await self.repository.add_snippet_or_update_content(s)

kodit/mcp.py

@@ -12,11 +12,15 @@ from pydantic import Field
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit._version import version
+from kodit.bm25.keyword_search_factory import keyword_search_factory
 from kodit.config import AppContext
 from kodit.database import Database
-from kodit.embedding.embedding import embedding_factory
-from kodit.search.search_repository import SearchRepository
-from kodit.search.search_service import SearchRequest, SearchResult, SearchService
+from kodit.embedding.embedding_factory import embedding_factory
+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
@@ -122,29 +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(
-        mcp_context.app_context.get_default_openai_client()
-    )
-
-    log.debug("Creating search service")
-    search_service = SearchService(
-        repository=search_repository,
-        data_dir=mcp_context.app_context.get_data_dir(),
-        embedding_service=embedding_service,
+    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,
+        ),
+        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/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