kodit 0.1.9__tar.gz → 0.1.11__tar.gz

{kodit-0.1.9 → kodit-0.1.11}/.gitignore

@@ -174,3 +174,6 @@ cython_debug/
 .pypirc
 .kodit/
 .DS_Store
+.kodit.db
+benchmark.db
+profile.prof

{kodit-0.1.9 → kodit-0.1.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.1.9
+Version: 0.1.11
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/
@@ -29,11 +29,13 @@ Requires-Dist: dotenv>=0.9.9
 Requires-Dist: fastapi[standard]>=0.115.12
 Requires-Dist: fastmcp>=2.3.3
 Requires-Dist: gitpython>=3.1.44
+Requires-Dist: hf-xet>=1.1.2
 Requires-Dist: httpx-retries>=0.3.2
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: posthog>=4.0.1
 Requires-Dist: pydantic-settings>=2.9.1
 Requires-Dist: pytable-formatter>=0.1.1
+Requires-Dist: sentence-transformers>=4.1.0
 Requires-Dist: sqlalchemy[asyncio]>=2.0.40
 Requires-Dist: structlog>=25.3.0
 Requires-Dist: tdqm>=0.0.1

{kodit-0.1.9 → kodit-0.1.11}/alembic.ini

@@ -3,7 +3,7 @@
 [alembic]
 # path to migration scripts
 # Use forward slashes (/) also on windows to provide an os agnostic path
-script_location = src/kodit/alembic
+script_location = src/kodit/migrations
 
 # template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
 # Uncomment the line below if you want the files to be prepended with date and time
@@ -63,7 +63,7 @@ version_path_separator = os
 # are written from script.py.mako
 # output_encoding = utf-8
 
-sqlalchemy.url = sqlite+aiosqlite:///%(here)s/.kodit/kodit.db
+sqlalchemy.url = sqlite+aiosqlite:///%(here)s/.kodit.db
 
 
 [post_write_hooks]

{kodit-0.1.9 → kodit-0.1.11}/docs/_index.md

@@ -57,13 +57,16 @@ pip install kodit
 Kodit has two key parts. A configuration CLI to manage what gets indexed and an MCP
 server to expose your code to an AI coding assistant.
 
-1. Index a local path: `kodit index /path/to/your/code`
-2. Or index a public git repository: `kodit index https://github.com/pydantic/pydantic-ai`
-3. Test retrieval on your index: `kodit retrieve "test"`
-4. Start an MCP server: `kodit serve`
-
-Now browse to your AI coding assistant and add the MCP server. You will also need to
-tell your assistant to use this server in coding tasks, otherwise it won't get called!
+1. Index a source:
+    1. a local path: `kodit index /path/to/your/code`
+    2. or index a public git repository: `kodit index https://github.com/pydantic/pydantic-ai`
+2. Manually search your index:
+    1. with a keyword: `kodit search keyword "test"`
+    2. or with code: `kodit search code "def main()"`
+    3. or via hybrid search:  `kodit search code hybrid --keywords "main" --code "def main()"`
+3. Start an MCP server: `kodit serve`
+
+Now add the Kodit MCP server to your AI coding assistant.
 
 ### Integration with Cursor
 

{kodit-0.1.9 → kodit-0.1.11}/docs/developer/index.md

@@ -13,8 +13,9 @@ All database operations are handled by SQLAlchemy and Alembic.
 
 1. Make changes to your models
 2. Ensure the model is referenced in [alembic's env.py](src/kodit/alembic/env.py)
-3. Run `alembic revision --autogenerate -m "your message"`
-4. The new migration will be applied when you next run a kodit command
+3. Run `alembic upgrade head` to create a temporary DB to compute the upgrade
+4. Run `alembic revision --autogenerate -m "your message"`
+5. The new migration will be applied when you next run a kodit command
 
 ## Releasing
 

{kodit-0.1.9 → kodit-0.1.11}/pyproject.toml

@@ -44,6 +44,8 @@ dependencies = [
     "pydantic-settings>=2.9.1",
     "bm25s[core]>=0.2.12",
     "gitpython>=3.1.44",
+    "sentence-transformers>=4.1.0",
+    "hf-xet>=1.1.2",
 ]
 
 [dependency-groups]
@@ -54,6 +56,7 @@ dev = [
     "pytest>=8.3.5",
     "pytest-cov>=6.1.1",
     "ruff>=0.11.8",
+    "snakeviz>=2.2.2",
 ]
 
 [project.urls]

{kodit-0.1.9 → kodit-0.1.11}/src/kodit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.9'
-__version_tuple__ = version_tuple = (0, 1, 9)
+__version__ = version = '0.1.11'
+__version_tuple__ = version_tuple = (0, 1, 11)

{kodit-0.1.9 → kodit-0.1.11}/src/kodit/bm25/bm25.py

@@ -38,7 +38,7 @@ class BM25Service:
         self.log.debug("Indexing corpus")
         vocab = self._tokenize(corpus)
         self.retriever = bm25s.BM25()
-        self.retriever.index(vocab)
+        self.retriever.index(vocab, show_progress=False)
         self.retriever.save(self.index_path)
 
     def retrieve(

{kodit-0.1.9 → kodit-0.1.11}/src/kodit/cli.py

@@ -15,6 +15,7 @@ from kodit.config import (
     DEFAULT_BASE_DIR,
     DEFAULT_DB_URL,
     DEFAULT_DISABLE_TELEMETRY,
+    DEFAULT_EMBEDDING_MODEL_NAME,
     DEFAULT_LOG_FORMAT,
     DEFAULT_LOG_LEVEL,
     AppContext,
@@ -23,7 +24,7 @@ from kodit.config import (
 )
 from kodit.indexing.repository import IndexRepository
 from kodit.indexing.service import IndexService
-from kodit.logging import configure_logging, configure_telemetry, log_event
+from kodit.log import configure_logging, configure_telemetry, log_event
 from kodit.retreival.repository import RetrievalRepository
 from kodit.retreival.service import RetrievalRequest, RetrievalService
 from kodit.sources.repository import SourceRepository
@@ -97,7 +98,12 @@ async def index(
     source_repository = SourceRepository(session)
     source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
-    service = IndexService(repository, source_service, app_context.get_data_dir())
+    service = IndexService(
+        repository,
+        source_service,
+        app_context.get_data_dir(),
+        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
+    )
 
     if not sources:
         # No source specified, list all indexes
@@ -133,20 +139,106 @@ async def index(
         await service.run(index.id)
 
 
-@cli.command()
+@cli.group()
+def search() -> None:
+    """Search for snippets in the database."""
+
+
+@search.command()
 @click.argument("query")
 @click.option("--top-k", default=10, help="Number of snippets to retrieve")
 @with_app_context
 @with_session
-async def retrieve(
-    session: AsyncSession, app_context: AppContext, query: str, top_k: int
+async def code(
+    session: AsyncSession,
+    app_context: AppContext,
+    query: str,
+    top_k: int,
+) -> None:
+    """Search for snippets using semantic code search.
+
+    This works best if your query is code.
+    """
+    repository = RetrievalRepository(session)
+    service = RetrievalService(
+        repository,
+        app_context.get_data_dir(),
+        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
+    )
+
+    snippets = await service.retrieve(RetrievalRequest(code_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(snippet.content)
+        click.echo("-" * 80)
+        click.echo()
+
+
+@search.command()
+@click.argument("keywords", nargs=-1)
+@click.option("--top-k", default=10, help="Number of snippets to retrieve")
+@with_app_context
+@with_session
+async def keyword(
+    session: AsyncSession,
+    app_context: AppContext,
+    keywords: list[str],
+    top_k: int,
 ) -> None:
-    """Retrieve snippets from the database."""
+    """Search for snippets using keyword search."""
     repository = RetrievalRepository(session)
-    service = RetrievalService(repository, app_context.get_data_dir())
-    # Temporary request while we don't have all search capabilities
+    service = RetrievalService(
+        repository,
+        app_context.get_data_dir(),
+        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
+    )
+
+    snippets = await service.retrieve(RetrievalRequest(keywords=keywords, 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(snippet.content)
+        click.echo("-" * 80)
+        click.echo()
+
+
+@search.command()
+@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")
+@with_app_context
+@with_session
+async def hybrid(
+    session: AsyncSession,
+    app_context: AppContext,
+    top_k: int,
+    keywords: str,
+    code: str,
+) -> None:
+    """Search for snippets using hybrid search."""
+    repository = RetrievalRepository(session)
+    service = RetrievalService(
+        repository,
+        app_context.get_data_dir(),
+        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
+    )
+
+    # Parse keywords into a list of strings
+    keywords_list = [k.strip().lower() for k in keywords.split(",")]
+
     snippets = await service.retrieve(
-        RetrievalRequest(keywords=query.split(","), top_k=top_k)
+        RetrievalRequest(keywords=keywords_list, code_query=code, top_k=top_k)
     )
 
     if len(snippets) == 0:

{kodit-0.1.9 → kodit-0.1.11}/src/kodit/config.py

@@ -11,12 +11,14 @@ from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from kodit.database import Database
+from kodit.embedding.embedding import TINY
 
 DEFAULT_BASE_DIR = Path.home() / ".kodit"
 DEFAULT_DB_URL = f"sqlite+aiosqlite:///{DEFAULT_BASE_DIR}/kodit.db"
 DEFAULT_LOG_LEVEL = "INFO"
 DEFAULT_LOG_FORMAT = "pretty"
 DEFAULT_DISABLE_TELEMETRY = False
+DEFAULT_EMBEDDING_MODEL_NAME = TINY
 T = TypeVar("T")
 
 

{kodit-0.1.9 → kodit-0.1.11}/src/kodit/database.py

@@ -15,7 +15,7 @@ from sqlalchemy.ext.asyncio import (
 )
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 
-from kodit import alembic
+from kodit import migrations
 
 
 class Base(AsyncAttrs, DeclarativeBase):
@@ -57,7 +57,7 @@ class Database:
         # Create Alembic configuration and run migrations
         alembic_cfg = AlembicConfig()
         alembic_cfg.set_main_option(
-            "script_location", str(Path(alembic.__file__).parent)
+            "script_location", str(Path(migrations.__file__).parent)
         )
         alembic_cfg.set_main_option("sqlalchemy.url", db_url)
         self.log.debug("Running migrations", db_url=db_url)

kodit-0.1.11/src/kodit/embedding/__init__.py

@@ -0,0 +1 @@
+"""Embedding module."""

kodit-0.1.11/src/kodit/embedding/embedding.py

@@ -0,0 +1,52 @@
+"""Embedding service."""
+
+import os
+from collections.abc import Generator
+
+import structlog
+from sentence_transformers import SentenceTransformer
+
+TINY = "tiny"
+CODE = "code"
+TEST = "test"
+
+COMMON_EMBEDDING_MODELS = {
+    TINY: "ibm-granite/granite-embedding-30m-english",
+    CODE: "flax-sentence-embeddings/st-codesearch-distilroberta-base",
+    TEST: "minishlab/potion-base-4M",
+}
+
+
+class EmbeddingService:
+    """Service for embeddings."""
+
+    def __init__(self, model_name: str) -> None:
+        """Initialize the embedding service."""
+        self.log = structlog.get_logger(__name__)
+        self.model_name = COMMON_EMBEDDING_MODELS.get(model_name, model_name)
+        self.embedding_model = None
+
+    def _model(self) -> SentenceTransformer:
+        """Get the embedding model."""
+        if self.embedding_model is None:
+            os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
+            self.embedding_model = SentenceTransformer(
+                self.model_name,
+                trust_remote_code=True,
+                device="cpu",  # Force CPU so we don't have to install accelerate, etc.
+            )
+        return self.embedding_model
+
+    def embed(self, snippets: list[str]) -> Generator[list[float], None, None]:
+        """Embed a list of documents."""
+        model = self._model()
+        embeddings = model.encode(snippets, show_progress_bar=False, batch_size=4)
+        for embedding in embeddings:
+            yield [float(x) for x in embedding]
+
+    def query(self, query: list[str]) -> Generator[list[float], None, None]:
+        """Query the embedding model."""
+        model = self._model()
+        embeddings = model.encode(query, show_progress_bar=False, batch_size=4)
+        for embedding in embeddings:
+            yield [float(x) for x in embedding]

kodit-0.1.11/src/kodit/embedding/models.py

@@ -0,0 +1,28 @@
+"""Embedding models."""
+
+from enum import Enum
+
+from sqlalchemy import JSON, ForeignKey
+from sqlalchemy import Enum as SQLAlchemyEnum
+from sqlalchemy.orm import Mapped, mapped_column
+
+from kodit.database import Base, CommonMixin
+
+
+class EmbeddingType(Enum):
+    """Embedding type."""
+
+    CODE = 1
+    TEXT = 2
+
+
+class Embedding(Base, CommonMixin):
+    """Embedding model."""
+
+    __tablename__ = "embeddings"
+
+    snippet_id: Mapped[int] = mapped_column(ForeignKey("snippets.id"), index=True)
+    type: Mapped[EmbeddingType] = mapped_column(
+        SQLAlchemyEnum(EmbeddingType), index=True
+    )
+    embedding: Mapped[list[float]] = mapped_column(JSON)

{kodit-0.1.9 → kodit-0.1.11}/src/kodit/indexing/repository.py

@@ -11,6 +11,7 @@ from typing import TypeVar
 from sqlalchemy import delete, func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from kodit.embedding.models import Embedding
 from kodit.indexing.models import Index, Snippet
 from kodit.sources.models import File, Source
 
@@ -165,3 +166,13 @@ class IndexRepository:
         query = select(Snippet).order_by(Snippet.id)
         result = await self.session.execute(query)
         return list(result.scalars())
+
+    async def add_embedding(self, embedding: Embedding) -> None:
+        """Add a new embedding to the database.
+
+        Args:
+            embedding: The Embedding instance to add.
+
+        """
+        self.session.add(embedding)
+        await self.session.commit()

{kodit-0.1.9 → kodit-0.1.11}/src/kodit/indexing/service.py

@@ -14,6 +14,8 @@ import structlog
 from tqdm.asyncio import tqdm
 
 from kodit.bm25.bm25 import BM25Service
+from kodit.embedding.embedding import EmbeddingService
+from kodit.embedding.models import Embedding, EmbeddingType
 from kodit.indexing.models import Snippet
 from kodit.indexing.repository import IndexRepository
 from kodit.snippets.snippets import SnippetService
@@ -50,6 +52,7 @@ class IndexService:
         repository: IndexRepository,
         source_service: SourceService,
         data_dir: Path,
+        embedding_model_name: str,
     ) -> None:
         """Initialize the index service.
 
@@ -63,6 +66,7 @@ class IndexService:
         self.snippet_service = SnippetService()
         self.log = structlog.get_logger(__name__)
         self.bm25 = BM25Service(data_dir)
+        self.code_embedding_service = EmbeddingService(model_name=embedding_model_name)
 
     async def create(self, source_id: int) -> IndexView:
         """Create a new index for a source.
@@ -128,9 +132,26 @@ class IndexService:
         # Create snippets for supported file types
         await self._create_snippets(index_id)
 
-        # Update BM25 index
         snippets = await self.repository.get_all_snippets()
-        self.bm25.index([snippet.content for snippet in snippets])
+
+        self.log.info("Creating keyword index")
+        self.bm25.index(
+            [
+                snippet.content
+                for snippet in tqdm(snippets, total=len(snippets), leave=False)
+            ]
+        )
+
+        self.log.info("Creating semantic code index")
+        for snippet in tqdm(snippets, total=len(snippets), leave=False):
+            embedding = next(self.code_embedding_service.embed([snippet.content]))
+            await self.repository.add_embedding(
+                Embedding(
+                    snippet_id=snippet.id,
+                    embedding=embedding,
+                    type=EmbeddingType.CODE,
+                )
+            )
 
         # Update index timestamp
         await self.repository.update_index_timestamp(index)
@@ -148,7 +169,7 @@ class IndexService:
 
         """
         files = await self.repository.files_for_index(index_id)
-        for file in tqdm(files, total=len(files)):
+        for file in tqdm(files, total=len(files), leave=False):
             # Skip unsupported file types
             if file.mime_type in MIME_BLACKLIST:
                 self.log.debug("Skipping mime type", mime_type=file.mime_type)

kodit-0.1.9/src/kodit/logging.py → kodit-0.1.11/src/kodit/log.py

@@ -87,7 +87,13 @@ def configure_logging(app_context: AppContext) -> None:
     # Configure uvicorn loggers to use our structlog setup
     # Uvicorn spits out loads of exception logs when sse server doesn't shut down
     # gracefully, so we hide them unless in DEBUG mode
-    for _log in ["uvicorn", "uvicorn.error", "uvicorn.access"]:
+    for _log in [
+        "uvicorn",
+        "uvicorn.error",
+        "uvicorn.access",
+        "bm25s",
+        "sentence_transformers.SentenceTransformer",
+    ]:
         if root_logger.getEffectiveLevel() == logging.DEBUG:
             logging.getLogger(_log).handlers.clear()
             logging.getLogger(_log).propagate = True

{kodit-0.1.9 → kodit-0.1.11}/src/kodit/mcp.py

@@ -12,7 +12,7 @@ from pydantic import Field
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit._version import version
-from kodit.config import AppContext
+from kodit.config import DEFAULT_EMBEDDING_MODEL_NAME, AppContext
 from kodit.database import Database
 from kodit.retreival.repository import RetrievalRepository, RetrievalResult
 from kodit.retreival.service import RetrievalRequest, RetrievalService
@@ -115,18 +115,12 @@ async def retrieve_relevant_snippets(
     retrieval_service = RetrievalService(
         repository=retrieval_repository,
         data_dir=mcp_context.data_dir,
+        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
     )
 
-    log.debug("Fusing input")
-    input_query = input_fusion(
-        user_intent=user_intent,
-        related_file_paths=related_file_paths,
-        related_file_contents=related_file_contents,
-        keywords=keywords,
-    )
-    log.debug("Input", input_query=input_query)
     retrieval_request = RetrievalRequest(
         keywords=keywords,
+        code_query="\n".join(related_file_contents),
     )
     log.debug("Retrieving snippets")
     snippets = await retrieval_service.retrieve(request=retrieval_request)

{kodit-0.1.9/src/kodit/alembic → kodit-0.1.11/src/kodit/migrations}/env.py

@@ -8,6 +8,7 @@ from sqlalchemy import pool
 from sqlalchemy.engine import Connection
 from sqlalchemy.ext.asyncio import async_engine_from_config
 
+import kodit.embedding.models
 import kodit.indexing.models
 import kodit.sources.models
 from kodit.database import Base

kodit-0.1.11/src/kodit/migrations/versions/7c3bbc2ab32b_add_embeddings_table.py

@@ -0,0 +1,47 @@
+# ruff: noqa
+"""add embeddings table
+
+Revision ID: 7c3bbc2ab32b
+Revises: 85155663351e
+Create Date: 2025-05-23 17:23:09.924980
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '7c3bbc2ab32b'
+down_revision: Union[str, None] = '85155663351e'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('embeddings',
+    sa.Column('snippet_id', sa.Integer(), nullable=False),
+    sa.Column('type', sa.Enum('CODE', 'TEXT', name='embeddingtype'), nullable=False),
+    sa.Column('embedding', sa.JSON(), nullable=False),
+    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
+    sa.Column('created_at', sa.DateTime(), nullable=False),
+    sa.Column('updated_at', sa.DateTime(), nullable=False),
+    sa.ForeignKeyConstraint(['snippet_id'], ['snippets.id'], ),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index(op.f('ix_embeddings_snippet_id'), 'embeddings', ['snippet_id'], unique=False)
+    op.create_index(op.f('ix_embeddings_type'), 'embeddings', ['type'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_embeddings_type'), table_name='embeddings')
+    op.drop_index(op.f('ix_embeddings_snippet_id'), table_name='embeddings')
+    op.drop_table('embeddings')
+    # ### end Alembic commands ###