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

kodit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.15'
-__version_tuple__ = version_tuple = (0, 1, 15)
+__version__ = version = '0.1.16'
+__version_tuple__ = version_tuple = (0, 1, 16)

kodit/cli.py

@@ -17,11 +17,10 @@ from kodit.config import (
     with_session,
 )
 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
+from kodit.indexing.indexing_service import IndexService, SearchRequest
 from kodit.log import configure_logging, configure_telemetry, log_event
-from kodit.search.search_repository import SearchRepository
-from kodit.search.search_service import SearchRequest, SearchService
 from kodit.source.source_repository import SourceRepository
 from kodit.source.source_service import SourceService
 
@@ -72,9 +71,13 @@ async def index(
         repository=repository,
         source_service=source_service,
         keyword_search_provider=keyword_search_factory(app_context, session),
-        vector_search_service=embedding_factory(
-            app_context=app_context, session=session
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
         ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
     )
 
     if not sources:
@@ -131,11 +134,20 @@ async def code(
 
     This works best if your query is code.
     """
-    repository = SearchRepository(session)
-    service = SearchService(
-        repository,
+    source_repository = SourceRepository(session)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
+    repository = IndexRepository(session)
+    service = IndexService(
+        repository=repository,
+        source_service=source_service,
         keyword_search_provider=keyword_search_factory(app_context, session),
-        embedding_service=embedding_factory(app_context=app_context, session=session),
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
+        ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
     )
 
     snippets = await service.search(SearchRequest(code_query=query, top_k=top_k))
@@ -147,6 +159,7 @@ async def code(
     for snippet in snippets:
         click.echo("-" * 80)
         click.echo(f"{snippet.uri}")
+        click.echo(f"Original scores: {snippet.original_scores}")
         click.echo(snippet.content)
         click.echo("-" * 80)
         click.echo()
@@ -164,11 +177,20 @@ async def keyword(
     top_k: int,
 ) -> None:
     """Search for snippets using keyword search."""
-    repository = SearchRepository(session)
-    service = SearchService(
-        repository,
+    source_repository = SourceRepository(session)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
+    repository = IndexRepository(session)
+    service = IndexService(
+        repository=repository,
+        source_service=source_service,
         keyword_search_provider=keyword_search_factory(app_context, session),
-        embedding_service=embedding_factory(app_context=app_context, session=session),
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
+        ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
     )
 
     snippets = await service.search(SearchRequest(keywords=keywords, top_k=top_k))
@@ -180,6 +202,53 @@ async def keyword(
     for snippet in snippets:
         click.echo("-" * 80)
         click.echo(f"{snippet.uri}")
+        click.echo(f"Original scores: {snippet.original_scores}")
+        click.echo(snippet.content)
+        click.echo("-" * 80)
+        click.echo()
+
+
+@search.command()
+@click.argument("query")
+@click.option("--top-k", default=10, help="Number of snippets to retrieve")
+@with_app_context
+@with_session
+async def text(
+    session: AsyncSession,
+    app_context: AppContext,
+    query: str,
+    top_k: int,
+) -> None:
+    """Search for snippets using semantic text search.
+
+    This works best if your query is text.
+    """
+    source_repository = SourceRepository(session)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
+    repository = IndexRepository(session)
+    service = IndexService(
+        repository=repository,
+        source_service=source_service,
+        keyword_search_provider=keyword_search_factory(app_context, session),
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
+        ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
+    )
+
+    snippets = await service.search(SearchRequest(text_query=query, top_k=top_k))
+
+    if len(snippets) == 0:
+        click.echo("No snippets found")
+        return
+
+    for snippet in snippets:
+        click.echo("-" * 80)
+        click.echo(f"{snippet.uri}")
+        click.echo(f"Original scores: {snippet.original_scores}")
         click.echo(snippet.content)
         click.echo("-" * 80)
         click.echo()
@@ -189,28 +258,44 @@ async def keyword(
 @click.option("--top-k", default=10, help="Number of snippets to retrieve")
 @click.option("--keywords", required=True, help="Comma separated list of keywords")
 @click.option("--code", required=True, help="Semantic code search query")
+@click.option("--text", required=True, help="Semantic text search query")
 @with_app_context
 @with_session
-async def hybrid(
+async def hybrid(  # noqa: PLR0913
     session: AsyncSession,
     app_context: AppContext,
     top_k: int,
     keywords: str,
     code: str,
+    text: str,
 ) -> None:
     """Search for snippets using hybrid search."""
-    repository = SearchRepository(session)
-    service = SearchService(
-        repository,
+    source_repository = SourceRepository(session)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
+    repository = IndexRepository(session)
+    service = IndexService(
+        repository=repository,
+        source_service=source_service,
         keyword_search_provider=keyword_search_factory(app_context, session),
-        embedding_service=embedding_factory(app_context=app_context, session=session),
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
+        ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
     )
 
     # Parse keywords into a list of strings
     keywords_list = [k.strip().lower() for k in keywords.split(",")]
 
     snippets = await service.search(
-        SearchRequest(keywords=keywords_list, code_query=code, top_k=top_k)
+        SearchRequest(
+            text_query=text,
+            keywords=keywords_list,
+            code_query=code,
+            top_k=top_k,
+        )
     )
 
     if len(snippets) == 0:
@@ -220,6 +305,7 @@ async def hybrid(
     for snippet in snippets:
         click.echo("-" * 80)
         click.echo(f"{snippet.uri}")
+        click.echo(f"Original scores: {snippet.original_scores}")
         click.echo(snippet.content)
         click.echo("-" * 80)
         click.echo()

kodit/embedding/embedding_factory.py

@@ -21,7 +21,7 @@ from kodit.embedding.vectorchord_vector_search_service import (
 
 
 def embedding_factory(
-    app_context: AppContext, session: AsyncSession
+    task_name: str, app_context: AppContext, session: AsyncSession
 ) -> VectorSearchService:
     """Create an embedding service."""
     embedding_repository = EmbeddingRepository(session=session)
@@ -33,7 +33,7 @@ def embedding_factory(
         embedding_provider = LocalEmbeddingProvider(CODE)
 
     if app_context.default_search.provider == "vectorchord":
-        return VectorChordVectorSearchService(session, embedding_provider)
+        return VectorChordVectorSearchService(task_name, session, embedding_provider)
     if app_context.default_search.provider == "sqlite":
         return LocalVectorSearchService(
             embedding_repository=embedding_repository,

kodit/embedding/embedding_provider/embedding_provider.py

@@ -38,8 +38,15 @@ def split_sub_batches(encoding: tiktoken.Encoding, data: list[str]) -> list[list
             item_tokens = len(encoding.encode(next_item))
 
             if item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
-                log.warning("Skipping too long snippet", snippet=data_to_process.pop(0))
-                continue
+                # Loop around trying to truncate the snippet until it fits in the max
+                # embedding size
+                while item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
+                    next_item = next_item[:-1]
+                    item_tokens = len(encoding.encode(next_item))
+
+                data_to_process[0] = next_item
+
+                log.warning("Truncated snippet", snippet=next_item)
 
             if current_tokens + item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
                 break

kodit/embedding/embedding_provider/openai_embedding_provider.py

@@ -38,26 +38,38 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
         # Process batches in parallel with a semaphore to limit concurrent requests
         sem = asyncio.Semaphore(OPENAI_NUM_PARALLEL_TASKS)
 
-        async def process_batch(batch: list[str]) -> list[Vector]:
+        # Create a list of tuples with a temporary id for each batch
+        # We need to do this so that we can return the results in the same order as the
+        # input data
+        input_data = [(i, batch) for i, batch in enumerate(batched_data)]
+
+        async def process_batch(
+            data: tuple[int, list[str]],
+        ) -> tuple[int, list[Vector]]:
+            batch_id, batch = data
             async with sem:
                 try:
                     response = await self.openai_client.embeddings.create(
                         model=self.model_name,
                         input=batch,
                     )
-                    return [
+                    return batch_id, [
                         [float(x) for x in embedding.embedding]
                         for embedding in response.data
                     ]
                 except Exception as e:
                     self.log.exception("Error embedding batch", error=str(e))
-                    return []
+                    return batch_id, []
 
         # Create tasks for all batches
-        tasks = [process_batch(batch) for batch in batched_data]
+        tasks = [process_batch(batch) for batch in input_data]
 
         # Process all batches and yield results as they complete
-        results: list[Vector] = []
+        results: list[tuple[int, list[Vector]]] = []
         for task in asyncio.as_completed(tasks):
-            results.extend(await task)
-        return results
+            result = await task
+            results.append(result)
+
+        # Output in the same order as the input data
+        ordered_results = [result for _, result in sorted(results, key=lambda x: x[0])]
+        return [item for sublist in ordered_results for item in sublist]

kodit/embedding/vectorchord_vector_search_service.py

@@ -12,23 +12,20 @@ from kodit.embedding.vector_search_service import (
     VectorSearchService,
 )
 
-TABLE_NAME = "vectorchord_embeddings"
-INDEX_NAME = f"{TABLE_NAME}_idx"
-
 # SQL Queries
 CREATE_VCHORD_EXTENSION = """
 CREATE EXTENSION IF NOT EXISTS vchord CASCADE;
 """
 
-CHECK_VCHORD_EMBEDDING_DIMENSION = f"""
+CHECK_VCHORD_EMBEDDING_DIMENSION = """
 SELECT a.atttypmod as dimension
 FROM pg_attribute a
 JOIN pg_class c ON a.attrelid = c.oid
 WHERE c.relname = '{TABLE_NAME}'
 AND a.attname = 'embedding';
-"""  # noqa: S608
+"""
 
-CREATE_VCHORD_INDEX = f"""
+CREATE_VCHORD_INDEX = """
 CREATE INDEX IF NOT EXISTS {INDEX_NAME}
 ON {TABLE_NAME}
 USING vchordrq (embedding vector_l2_ops) WITH (options = $$
@@ -38,21 +35,21 @@ lists = []
 $$);
 """
 
-INSERT_QUERY = f"""
+INSERT_QUERY = """
 INSERT INTO {TABLE_NAME} (snippet_id, embedding)
 VALUES (:snippet_id, :embedding)
 ON CONFLICT (snippet_id) DO UPDATE
 SET embedding = EXCLUDED.embedding
-"""  # noqa: S608
+"""
 
 # Note that <=> in vectorchord is cosine distance
 # So scores go from 0 (similar) to 2 (opposite)
-SEARCH_QUERY = f"""
+SEARCH_QUERY = """
 SELECT snippet_id, embedding <=> :query as score
 FROM {TABLE_NAME}
 ORDER BY score ASC
 LIMIT :top_k;
-"""  # noqa: S608
+"""
 
 
 class VectorChordVectorSearchService(VectorSearchService):
@@ -60,6 +57,7 @@ class VectorChordVectorSearchService(VectorSearchService):
 
     def __init__(
         self,
+        task_name: str,
         session: AsyncSession,
         embedding_provider: EmbeddingProvider,
     ) -> None:
@@ -67,6 +65,8 @@ class VectorChordVectorSearchService(VectorSearchService):
         self.embedding_provider = embedding_provider
         self._session = session
         self._initialized = False
+        self.table_name = f"vectorchord_{task_name}_embeddings"
+        self.index_name = f"{self.table_name}_idx"
 
     async def _initialize(self) -> None:
         """Initialize the VectorChord environment."""
@@ -88,15 +88,23 @@ class VectorChordVectorSearchService(VectorSearchService):
         vector_dim = (await self.embedding_provider.embed(["dimension"]))[0]
         await self._session.execute(
             text(
-                f"""CREATE TABLE IF NOT EXISTS {TABLE_NAME} (
+                f"""CREATE TABLE IF NOT EXISTS {self.table_name} (
                     id SERIAL PRIMARY KEY,
                     snippet_id INT NOT NULL UNIQUE,
                     embedding VECTOR({len(vector_dim)}) NOT NULL
                 );"""
             )
         )
-        await self._session.execute(text(CREATE_VCHORD_INDEX))
-        result = await self._session.execute(text(CHECK_VCHORD_EMBEDDING_DIMENSION))
+        await self._session.execute(
+            text(
+                CREATE_VCHORD_INDEX.format(
+                    TABLE_NAME=self.table_name, INDEX_NAME=self.index_name
+                )
+            )
+        )
+        result = await self._session.execute(
+            text(CHECK_VCHORD_EMBEDDING_DIMENSION.format(TABLE_NAME=self.table_name))
+        )
         vector_dim_from_db = result.scalar_one()
         if vector_dim_from_db != len(vector_dim):
             msg = (
@@ -123,7 +131,7 @@ class VectorChordVectorSearchService(VectorSearchService):
         embeddings = await self.embedding_provider.embed([doc.text for doc in data])
         # Execute inserts
         await self._execute(
-            text(INSERT_QUERY),
+            text(INSERT_QUERY.format(TABLE_NAME=self.table_name)),
             [
                 {"snippet_id": doc.snippet_id, "embedding": str(embedding)}
                 for doc, embedding in zip(data, embeddings, strict=True)
@@ -135,7 +143,8 @@ class VectorChordVectorSearchService(VectorSearchService):
         """Query the embedding model."""
         embedding = await self.embedding_provider.embed([query])
         result = await self._execute(
-            text(SEARCH_QUERY), {"query": str(embedding[0]), "top_k": top_k}
+            text(SEARCH_QUERY.format(TABLE_NAME=self.table_name)),
+            {"query": str(embedding[0]), "top_k": top_k},
         )
         rows = result.mappings().all()
 

kodit/enrichment/__init__.py

@@ -0,0 +1 @@
+"""Enrichment."""

kodit/enrichment/enrichment_factory.py

@@ -0,0 +1,23 @@
+"""Embedding service."""
+
+from kodit.config import AppContext
+from kodit.enrichment.enrichment_provider.local_enrichment_provider import (
+    LocalEnrichmentProvider,
+)
+from kodit.enrichment.enrichment_provider.openai_enrichment_provider import (
+    OpenAIEnrichmentProvider,
+)
+from kodit.enrichment.enrichment_service import (
+    EnrichmentService,
+    LLMEnrichmentService,
+)
+
+
+def enrichment_factory(app_context: AppContext) -> EnrichmentService:
+    """Create an embedding service."""
+    openai_client = app_context.get_default_openai_client()
+    if openai_client is not None:
+        enrichment_provider = OpenAIEnrichmentProvider(openai_client=openai_client)
+        return LLMEnrichmentService(enrichment_provider)
+
+    return LLMEnrichmentService(LocalEnrichmentProvider())

kodit/enrichment/enrichment_provider/__init__.py

@@ -0,0 +1 @@
+"""Enrichment provider."""

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

@@ -196,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,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.
@@ -152,9 +184,93 @@ class IndexService:
                 ]
             )
 
+        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,

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/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-0.1.15.dist-info → kodit-0.1.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.1.15
+Version: 0.1.16
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/
@@ -15,6 +15,7 @@ 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
@@ -42,6 +43,7 @@ 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

{kodit-0.1.15.dist-info → kodit-0.1.16.dist-info}/RECORD

@@ -1,12 +1,12 @@
 kodit/.gitignore,sha256=ztkjgRwL9Uud1OEi36hGQeDGk3OLK1NfDEO8YqGYy8o,11
 kodit/__init__.py,sha256=aEKHYninUq1yh6jaNfvJBYg-6fenpN132nJt1UU6Jxs,59
-kodit/_version.py,sha256=OX-WIjJlMaFvqRmCfLtOYEOYoiov9NdOA089N36rG-g,513
+kodit/_version.py,sha256=VYJNWHISWEW-KD_clKUYcTY_Z30r993Sjws4URJIL0g,513
 kodit/app.py,sha256=Mr5BFHOHx5zppwjC4XPWVvHjwgl1yrKbUjTWXKubJQM,891
-kodit/cli.py,sha256=wAaMZQs-h6hyashWB3DBR2GIf496vfHmepcXhpa7-eM,8085
+kodit/cli.py,sha256=i7eEt0FdIQGEfXKFte-8fBcZZGE8BPXBp40aGwJDQGI,11323
 kodit/config.py,sha256=2W2u5J8j-Mbt-C4xzOuK-PeuDCx0S_rnCXPhBwvfLT4,4353
 kodit/database.py,sha256=WB1KpVxUYPgiJGU0gJa2hqytYB8wJEJ5z3WayhWzNMU,2403
 kodit/log.py,sha256=HU1OmuxO4FcVw61k4WW7Y4WM7BrDaeplw1PcBHhuIZY,5434
-kodit/mcp.py,sha256=HA3R7YG0Al1A6MjSCSIi0hEGXG3WP7tix-N5AROasCM,5278
+kodit/mcp.py,sha256=QruyPskWB0_x59pkfj5BBeXuR13GMny5TAZEa2j4U9s,5752
 kodit/middleware.py,sha256=I6FOkqG9-8RH5kR1-0ZoQWfE4qLCB8lZYv8H_OCH29o,2714
 kodit/bm25/__init__.py,sha256=j8zyriNWhbwE5Lbybzg1hQAhANlU9mKHWw4beeUR6og,19
 kodit/bm25/keyword_search_factory.py,sha256=rp-wx3DJsc2KlELK1V337EyeYvmwnMQwUqOo1WVPSmg,631
@@ -14,21 +14,29 @@ kodit/bm25/keyword_search_service.py,sha256=aBbWQKgQmi2re3EIHdXFS00n7Wj3b2D0pZsL
 kodit/bm25/local_bm25.py,sha256=AAbFhbQDqyL3d7jsPL7W4HsLxdoYctaDsREUXOLy6jM,3260
 kodit/bm25/vectorchord_bm25.py,sha256=_nGrkUReYLLV-L8RIuIVLwjuhSYZl9T532n5OVf0kWs,6393
 kodit/embedding/__init__.py,sha256=h9NXzDA1r-K23nvBajBV-RJzHJN0p3UJ7UQsmdnOoRw,24
-kodit/embedding/embedding_factory.py,sha256=qzoxBS3scR-ABd-u9215uGES7c6clYy2DiKcSDQivnA,1603
+kodit/embedding/embedding_factory.py,sha256=UGnFRyyQXazSUOwyW4Hg7Vq2-kfAoDj9lD4CTLu8x04,1630
 kodit/embedding/embedding_models.py,sha256=rN90vSs86dYiqoawcp8E9jtwY31JoJXYfaDlsJK7uqc,656
 kodit/embedding/embedding_repository.py,sha256=-ux3scpBzel8c0pMH9fNOEsSXFIzl-IfgaWrkTb1szo,6907
 kodit/embedding/local_vector_search_service.py,sha256=hkF0qlfzjyGt400qIX9Mr6B7b7i8WvYIYWN2Z2C_pcs,1907
 kodit/embedding/vector_search_service.py,sha256=pQJ129QjGrAWOXzqkywmgtDRpy8_gtzYgkivyqF9Vrs,1009
-kodit/embedding/vectorchord_vector_search_service.py,sha256=OsVeM3gpoT8Ihzh-kEIzBm3xh_a4D-sErPvsQSKCME8,4732
+kodit/embedding/vectorchord_vector_search_service.py,sha256=KSs0IMFHHIllwq2d3A0LGqGGZDqO1Ht6K-BCfBBWW0Y,5051
 kodit/embedding/embedding_provider/__init__.py,sha256=h9NXzDA1r-K23nvBajBV-RJzHJN0p3UJ7UQsmdnOoRw,24
-kodit/embedding/embedding_provider/embedding_provider.py,sha256=NKs4nriup47R8xRciP07NE1-eZE9RPHklS7VH910UZ4,1537
+kodit/embedding/embedding_provider/embedding_provider.py,sha256=Tf3bwUsUMzAgoyLFM5qBtOLqPp1qr03TzrwGczkDvy0,1835
 kodit/embedding/embedding_provider/hash_embedding_provider.py,sha256=nAhlhh8j8PqqCCbhVl26Y8ntFBm2vJBCtB4X04g5Wwg,2638
 kodit/embedding/embedding_provider/local_embedding_provider.py,sha256=4ER-UPq506Y0TWU6qcs0nUqw6bSKQkSrdog-DhNQWM8,1906
-kodit/embedding/embedding_provider/openai_embedding_provider.py,sha256=bmUpegDgaF5Qj9uWcj1az4ADA2cKHUjraaMjGGPr83U,2076
+kodit/embedding/embedding_provider/openai_embedding_provider.py,sha256=V_jdUXiaGdslplwxMlfgFc4_hAVS2eaJXMTs2C7RiLI,2666
+kodit/enrichment/__init__.py,sha256=vBEolHpKaHUhfINX0dSGyAPlvgpLNAer9YzFtdvCB24,18
+kodit/enrichment/enrichment_factory.py,sha256=vKjkUTdhj74IW2S4GENDWdWMJx6BwUSZjJGDC0i7DSk,787
+kodit/enrichment/enrichment_service.py,sha256=87Sd3gGbEMJYb_wVrHG8L1yGIZmQNR7foUS4_y94azI,977
+kodit/enrichment/enrichment_provider/__init__.py,sha256=klf8iuLVWX4iRz-DZQauFFNAoJC5CByczh48TBZPW-o,27
+kodit/enrichment/enrichment_provider/enrichment_provider.py,sha256=E0H5rq3OENM0yYbA8K_3nSnj5lUHCpoIOqpWLo-2MVU,413
+kodit/enrichment/enrichment_provider/local_enrichment_provider.py,sha256=bR6HR1gH7wtZdMLOwaKdASjvllRo1FlNW9GyZC11zAM,2164
+kodit/enrichment/enrichment_provider/openai_enrichment_provider.py,sha256=gYuFTAeIVdQNlCUvNSPgRoiRwCvRD0C8419h8ubyABA,2725
 kodit/indexing/__init__.py,sha256=cPyi2Iej3G1JFWlWr7X80_UrsMaTu5W5rBwgif1B3xo,75
+kodit/indexing/fusion.py,sha256=TZb4fPAedXdEUXzwzOofW98QIOymdbclBOP1KOijuEk,1674
 kodit/indexing/indexing_models.py,sha256=6NX9HVcj6Pu9ePwHC7n-PWSyAgukpJq0nCNmUIigtbo,1282
-kodit/indexing/indexing_repository.py,sha256=4RJ3zY8p6QxHrYW7dDjru_w94Eu19v2gQ4mdlTgcXvY,6331
-kodit/indexing/indexing_service.py,sha256=T_dxOzNW_0OCpR4Fha1hHuNkmtLcDMZwL6t5xeu5VXQ,6613
+kodit/indexing/indexing_repository.py,sha256=GYHoACUWYKQdVTwP7tfik_TMUD1WUK76nywH88eCSwg,7006
+kodit/indexing/indexing_service.py,sha256=tKcZpi0pzsmF6OpqnqF0Q5HfSXxi5iLTysrVSou4JiQ,10579
 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
@@ -36,14 +44,12 @@ kodit/migrations/script.py.mako,sha256=zWziKtiwYKEWuwPV_HBNHwa9LCT45_bi01-uSNFaO
 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=6q0k7JMTM_7hPK2TSA30CykGbc5N16kCL7HTjlbai0w,1563
-kodit/search/search_service.py,sha256=-XlbP_50e1dKFJ5jBvex5FjLnffW43LcwQV_SeYNFB0,3944
 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/go.scm,sha256=SEX9mTOrhP2KiQW7oflDKkd21u5dK56QbJ4LvTDxY8A,533
 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
@@ -51,8 +57,8 @@ kodit/source/source_repository.py,sha256=0EksMpoLzdkfe8S4eeCm4Sf7TuxsOzOzaF4BBsM
 kodit/source/source_service.py,sha256=u_GaH07ewakThQJRfT8O_yZ54A52qLtJuM1bF3xUT2A,9633
 kodit/util/__init__.py,sha256=bPu6CtqDWCRGU7VgW2_aiQrCBi8G89FS6k1PjvDajJ0,37
 kodit/util/spinner.py,sha256=R9bzrHtBiIH6IfLbmsIVHL53s8vg-tqW4lwGGALu4dw,1932
-kodit-0.1.15.dist-info/METADATA,sha256=8E-bw8L-Df5Hdt16R5IWkyw7uUAr13CwYfcEyFExaPw,2380
-kodit-0.1.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-kodit-0.1.15.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
-kodit-0.1.15.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-kodit-0.1.15.dist-info/RECORD,,
+kodit-0.1.16.dist-info/METADATA,sha256=1lR4ZSTiRBzUv9Gj8FPspv4GU2vWGQU6HSiffWgU2Do,2467
+kodit-0.1.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+kodit-0.1.16.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
+kodit-0.1.16.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+kodit-0.1.16.dist-info/RECORD,,

kodit/search/__init__.py

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

kodit/search/search_repository.py

@@ -1,57 +0,0 @@
-"""Repository for searching for relevant snippets."""
-
-from typing import TypeVar
-
-from sqlalchemy import (
-    select,
-)
-from sqlalchemy.ext.asyncio import AsyncSession
-
-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]

kodit/search/search_service.py

@@ -1,135 +0,0 @@
-"""Search service."""
-
-import pydantic
-import structlog
-
-from kodit.bm25.keyword_search_service import BM25Result, KeywordSearchProvider
-from kodit.embedding.vector_search_service import VectorSearchService
-from kodit.search.search_repository import SearchRepository
-
-
-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."""
-
-    content: str
-    file_path: str
-
-
-class SearchService:
-    """Service for searching for relevant data."""
-
-    def __init__(
-        self,
-        repository: SearchRepository,
-        keyword_search_provider: KeywordSearchProvider,
-        embedding_service: VectorSearchService,
-    ) -> None:
-        """Initialize the search service."""
-        self.repository = repository
-        self.log = structlog.get_logger(__name__)
-        self.keyword_search_provider = keyword_search_provider
-        self.code_embedding_service = embedding_service
-
-    async def search(self, request: SearchRequest) -> list[SearchResult]:
-        """Search for relevant data."""
-        fusion_list = []
-        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)
-
-            # Sort results by score
-            result_ids.sort(key=lambda x: x[1], reverse=True)
-
-            self.log.debug("Search results (BM25)", results=result_ids)
-
-            bm25_results = [x[0] for x in result_ids]
-            fusion_list.append(bm25_results)
-
-        # Compute embedding for semantic query
-        semantic_results = []
-        if request.code_query:
-            query_embedding = await self.code_embedding_service.retrieve(
-                request.code_query, top_k=request.top_k
-            )
-            semantic_results = [x.snippet_id for x in query_embedding]
-            fusion_list.append(semantic_results)
-
-        if len(fusion_list) == 0:
-            return []
-
-        # Combine all results together with RFF if required
-        final_results = reciprocal_rank_fusion(fusion_list, k=60)
-
-        # Extract ids from final results
-        final_ids = [x[0] for x in final_results]
-
-        # Get snippets from database (up to 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(
-    rankings: list[list[int]], k: float = 60
-) -> list[tuple[int, float]]:
-    """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] = float(0)
-
-    for ranker in rankings:
-        for i, rank in enumerate(ranker):
-            scores[rank] += 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)
-
-    return results