kodit 0.1.15__py3-none-any.whl → 0.1.17__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.17'
+__version_tuple__ = version_tuple = (0, 1, 17)

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)
@@ -134,8 +142,11 @@ class VectorChordVectorSearchService(VectorSearchService):
     async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
         """Query the embedding model."""
         embedding = await self.embedding_provider.embed([query])
+        if len(embedding) == 0 or len(embedding[0]) == 0:
+            return []
         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)