kodit 0.1.4__tar.gz → 0.1.6__tar.gz

{kodit-0.1.4 → kodit-0.1.6}/.github/workflows/pypi.yaml

@@ -32,7 +32,6 @@ jobs:
 
           REPO_NAME=${{ github.event.repository.name }}
           REPO_TAG=${{ github.event.release.tag_name }}
-          REPO_TAG=0.1.3
 
           # Get the first letter of the repo name
           REPO_NAME_FIRST_LETTER=${REPO_NAME:0:1}

{kodit-0.1.4 → kodit-0.1.6}/.github/workflows/test.yaml

@@ -14,6 +14,7 @@ permissions:
 jobs:
   test:
     runs-on: ubuntu-latest
+    timeout-minutes: 10
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -44,6 +45,7 @@ jobs:
 
   build-package:
     runs-on: ubuntu-latest
+    timeout-minutes: 10
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -67,6 +69,7 @@ jobs:
   test-package:
     needs: build-package
     runs-on: ubuntu-latest
+    timeout-minutes: 10
     steps:
       - uses: actions/checkout@v4
         with:
@@ -97,5 +100,8 @@ jobs:
       - name: Run simple version command test
         run: kodit version
 
+      - name: Delete kodit data_dir
+        run: rm -rf ${HOME}/.kodit
+
       - name: Run smoke test
         run: ./tests/smoke.sh

kodit-0.1.6/.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: Module",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "src.kodit.cli",
+            "args": [
+                "retrieve",
+                "hello",
+            ]
+        }
+    ]
+}

{kodit-0.1.4 → kodit-0.1.6}/.vscode/settings.json

@@ -24,5 +24,5 @@
     "python.analysis.autoImportCompletions": true,
     "python.analysis.completeFunctionParens": true,
     "python.analysis.inlayHints.functionReturnTypes": true,
-    "python.analysis.typeCheckingMode": "strict",
+    "python.analysis.typeCheckingMode": "standard",
 }

{kodit-0.1.4 → kodit-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.1.4
+Version: 0.1.6
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/
@@ -22,18 +22,22 @@ Requires-Dist: aiosqlite>=0.20.0
 Requires-Dist: alembic>=1.15.2
 Requires-Dist: asgi-correlation-id>=4.3.4
 Requires-Dist: better-exceptions>=0.3.3
+Requires-Dist: bm25s[core]>=0.2.12
 Requires-Dist: click>=8.1.8
 Requires-Dist: colorama>=0.4.6
 Requires-Dist: dotenv>=0.9.9
 Requires-Dist: fastapi[standard]>=0.115.12
+Requires-Dist: fastmcp>=2.3.3
 Requires-Dist: httpx-retries>=0.3.2
 Requires-Dist: httpx>=0.28.1
-Requires-Dist: mcp>=1.6.0
 Requires-Dist: posthog>=4.0.1
+Requires-Dist: pydantic-settings>=2.9.1
 Requires-Dist: pytable-formatter>=0.1.1
 Requires-Dist: sqlalchemy[asyncio]>=2.0.40
 Requires-Dist: structlog>=25.3.0
 Requires-Dist: tdqm>=0.0.1
+Requires-Dist: tree-sitter-language-pack>=0.7.3
+Requires-Dist: tree-sitter>=0.24.0
 Requires-Dist: uritools>=5.0.0
 Description-Content-Type: text/markdown
 

{kodit-0.1.4 → kodit-0.1.6}/docs/_index.md

@@ -87,6 +87,7 @@ description:
 globs: 
 alwaysApply: true
 ---
-You MUST use the code-search MCP tool and always include any file context the user has provided in the chat.```
+You MUST use the code-search MCP tool and always include any file context the user has provided in the chat.
+```
 
 Alternatively, you can browse to the cursor settings and set this prompt globally.

{kodit-0.1.4 → kodit-0.1.6}/pyproject.toml

@@ -29,7 +29,6 @@ dependencies = [
     "fastapi[standard]>=0.115.12",
     "httpx-retries>=0.3.2",
     "httpx>=0.28.1",
-    "mcp>=1.6.0",
     "structlog>=25.3.0",
     "posthog>=4.0.1",
     "sqlalchemy[asyncio]>=2.0.40",
@@ -39,6 +38,11 @@ dependencies = [
     "aiofiles>=24.1.0",
     "tdqm>=0.0.1",
     "uritools>=5.0.0",
+    "tree-sitter-language-pack>=0.7.3",
+    "tree-sitter>=0.24.0",
+    "fastmcp>=2.3.3",
+    "pydantic-settings>=2.9.1",
+    "bm25s[core]>=0.2.12",
 ]
 
 [dependency-groups]
@@ -104,7 +108,7 @@ ignore = [
   "PGH004", # If I've disabled all, I mean disable all
 ]
 select = ["ALL"]
-exclude = []
+exclude = ["./tests/*"]
 
 [[tool.uv.index]]
 name = "pypi"

{kodit-0.1.4 → kodit-0.1.6}/src/kodit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.4'
-__version_tuple__ = version_tuple = (0, 1, 4)
+__version__ = version = '0.1.6'
+__version_tuple__ = version_tuple = (0, 1, 6)

{kodit-0.1.4 → kodit-0.1.6}/src/kodit/alembic/env.py

@@ -3,7 +3,6 @@
 
 import asyncio
 
-import structlog
 from alembic import context
 from sqlalchemy import pool
 from sqlalchemy.engine import Connection
@@ -66,8 +65,6 @@ async def run_async_migrations() -> None:
         prefix="sqlalchemy.",
         poolclass=pool.NullPool,
     )
-    log = structlog.get_logger(__name__)
-    log.debug("Running migrations on %s", connectable.url)
 
     async with connectable.connect() as connection:
         await connection.run_sync(do_run_migrations)
@@ -77,7 +74,11 @@ async def run_async_migrations() -> None:
 
 def run_migrations_online() -> None:
     """Run migrations in 'online' mode."""
-    asyncio.run(run_async_migrations())
+    connectable = config.attributes.get("connection", None)
+    if connectable is None:
+        asyncio.run(run_async_migrations())
+    else:
+        do_run_migrations(connectable)
 
 
 if context.is_offline_mode():

kodit-0.1.6/src/kodit/app.py

@@ -0,0 +1,29 @@
+"""FastAPI application for kodit API."""
+
+from asgi_correlation_id import CorrelationIdMiddleware
+from fastapi import FastAPI
+
+from kodit.mcp import mcp
+from kodit.middleware import ASGICancelledErrorMiddleware, logging_middleware
+
+# See https://gofastmcp.com/deployment/asgi#fastapi-integration
+mcp_app = mcp.sse_app()
+app = FastAPI(title="kodit API", lifespan=mcp_app.router.lifespan_context)
+
+# Add middleware
+app.middleware("http")(logging_middleware)
+app.add_middleware(CorrelationIdMiddleware)
+
+
+@app.get("/")
+async def root() -> dict[str, str]:
+    """Return a welcome message for the kodit API."""
+    return {"message": "Hello, World!"}
+
+
+# Add mcp routes last, otherwise previous routes aren't added
+app.mount("", mcp_app)
+
+# Wrap the entire app with ASGI middleware after all routes are added to suppress
+# CancelledError at the ASGI level
+app = ASGICancelledErrorMiddleware(app)

kodit-0.1.6/src/kodit/bm25/__init__.py

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

kodit-0.1.6/src/kodit/bm25/bm25.py

@@ -0,0 +1,71 @@
+"""BM25 service."""
+
+from pathlib import Path
+
+import bm25s
+import Stemmer
+import structlog
+from bm25s.tokenization import Tokenized
+
+
+class BM25Service:
+    """Service for BM25."""
+
+    def __init__(self, data_dir: Path) -> None:
+        """Initialize the BM25 service."""
+        self.log = structlog.get_logger(__name__)
+        self.index_path = data_dir / "bm25s_index"
+        try:
+            self.log.debug("Loading BM25 index")
+            self.retriever = bm25s.BM25.load(self.index_path, mmap=True)
+        except FileNotFoundError:
+            self.log.debug("BM25 index not found, creating new index")
+            self.retriever = bm25s.BM25()
+
+        self.stemmer = Stemmer.Stemmer("english")
+
+    def _tokenize(self, corpus: list[str]) -> list[list[str]] | Tokenized:
+        return bm25s.tokenize(
+            corpus,
+            stopwords="en",
+            stemmer=self.stemmer,
+            return_ids=False,
+            show_progress=True,
+        )
+
+    def index(self, corpus: list[str]) -> None:
+        """Index a new corpus."""
+        self.log.debug("Indexing corpus")
+        vocab = self._tokenize(corpus)
+        self.retriever = bm25s.BM25()
+        self.retriever.index(vocab)
+        self.retriever.save(self.index_path)
+
+    def retrieve(
+        self, doc_ids: list[int], query: str, top_k: int = 2
+    ) -> list[tuple[int, float]]:
+        """Retrieve from the index."""
+        if top_k == 0:
+            self.log.warning("Top k is 0, returning empty list")
+            return []
+        if len(doc_ids) == 0:
+            self.log.warning("No documents to retrieve from, returning empty list")
+            return []
+
+        top_k = min(top_k, len(doc_ids))
+        self.log.debug(
+            "Retrieving from index", query=query, top_k=top_k, num_docs=len(doc_ids)
+        )
+
+        query_tokens = self._tokenize([query])
+
+        self.log.debug("Query tokens", query_tokens=query_tokens)
+
+        results, scores = self.retriever.retrieve(
+            query_tokens=query_tokens, corpus=doc_ids, k=top_k
+        )
+        self.log.debug("Raw results", results=results, scores=scores)
+        return [
+            (int(result), float(score))
+            for result, score in zip(results[0], scores[0], strict=False)
+        ]

kodit-0.1.6/src/kodit/cli.py

@@ -0,0 +1,272 @@
+"""Command line interface for kodit."""
+
+import os
+import signal
+from pathlib import Path
+from typing import Any
+
+import click
+import structlog
+import uvicorn
+from pytable_formatter import Table
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.config import (
+    DEFAULT_BASE_DIR,
+    DEFAULT_DB_URL,
+    DEFAULT_DISABLE_TELEMETRY,
+    DEFAULT_LOG_FORMAT,
+    DEFAULT_LOG_LEVEL,
+    AppContext,
+    with_app_context,
+    with_session,
+)
+from kodit.indexing.repository import IndexRepository
+from kodit.indexing.service import IndexService
+from kodit.logging 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
+from kodit.sources.service import SourceService
+
+
+@click.group(context_settings={"max_content_width": 100})
+@click.option("--log-level", help=f"Log level [default: {DEFAULT_LOG_LEVEL}]")
+@click.option("--log-format", help=f"Log format [default: {DEFAULT_LOG_FORMAT}]")
+@click.option(
+    "--disable-telemetry",
+    is_flag=True,
+    help=f"Disable telemetry [default: {DEFAULT_DISABLE_TELEMETRY}]",
+)
+@click.option("--db-url", help=f"Database URL [default: {DEFAULT_DB_URL}]")
+@click.option("--data-dir", help=f"Data directory [default: {DEFAULT_BASE_DIR}]")
+@click.option(
+    "--env-file",
+    help="Path to a .env file [default: .env]",
+    type=click.Path(
+        exists=True,
+        dir_okay=False,
+        resolve_path=True,
+        path_type=Path,
+    ),
+)
+@click.pass_context
+def cli(  # noqa: PLR0913
+    ctx: click.Context,
+    log_level: str | None,
+    log_format: str | None,
+    disable_telemetry: bool | None,
+    db_url: str | None,
+    data_dir: str | None,
+    env_file: Path | None,
+) -> None:
+    """kodit CLI - Code indexing for better AI code generation."""  # noqa: D403
+    config = AppContext()
+    # First check if env-file is set and reload config if it is
+    if env_file:
+        config = AppContext(_env_file=env_file)  # type: ignore[reportCallIssue]
+
+    # Now override with CLI arguments, if set
+    if data_dir:
+        config.data_dir = Path(data_dir)
+    if db_url:
+        config.db_url = db_url
+    if log_level:
+        config.log_level = log_level
+    if log_format:
+        config.log_format = log_format
+    if disable_telemetry:
+        config.disable_telemetry = disable_telemetry
+    configure_logging(config)
+    configure_telemetry(config)
+
+    # Set the app context in the click context for downstream cli
+    ctx.obj = config
+
+
+@cli.group()
+def sources() -> None:
+    """Manage code sources."""
+
+
+@sources.command(name="list")
+@with_app_context
+@with_session
+async def list_sources(session: AsyncSession, app_context: AppContext) -> None:
+    """List all code sources."""
+    repository = SourceRepository(session)
+    service = SourceService(app_context.get_clone_dir(), repository)
+    sources = await service.list_sources()
+
+    # Define headers and data
+    headers = ["ID", "Created At", "URI"]
+    data = [[source.id, source.created_at, source.uri] for source in sources]
+
+    # Create and display the table
+    table = Table(headers=headers, data=data)
+    click.echo(table)
+
+
+@sources.command(name="create")
+@click.argument("uri")
+@with_app_context
+@with_session
+async def create_source(
+    session: AsyncSession, app_context: AppContext, uri: str
+) -> None:
+    """Add a new code source."""
+    repository = SourceRepository(session)
+    service = SourceService(app_context.get_clone_dir(), repository)
+    source = await service.create(uri)
+    click.echo(f"Source created: {source.id}")
+
+
+@cli.group()
+def indexes() -> None:
+    """Manage indexes."""
+
+
+@indexes.command(name="create")
+@click.argument("source_id")
+@with_app_context
+@with_session
+async def create_index(
+    session: AsyncSession, app_context: AppContext, source_id: int
+) -> None:
+    """Create an index for a source."""
+    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())
+    index = await service.create(source_id)
+    click.echo(f"Index created: {index.id}")
+
+
+@indexes.command(name="list")
+@with_app_context
+@with_session
+async def list_indexes(session: AsyncSession, app_context: AppContext) -> None:
+    """List all indexes."""
+    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())
+    indexes = await service.list_indexes()
+
+    # Define headers and data
+    headers = [
+        "ID",
+        "Created At",
+        "Updated At",
+        "Num Snippets",
+    ]
+    data = [
+        [
+            index.id,
+            index.created_at,
+            index.updated_at,
+            index.num_snippets,
+        ]
+        for index in indexes
+    ]
+
+    # Create and display the table
+    table = Table(headers=headers, data=data)
+    click.echo(table)
+
+
+@indexes.command(name="run")
+@click.argument("index_id")
+@with_app_context
+@with_session
+async def run_index(
+    session: AsyncSession, app_context: AppContext, index_id: int
+) -> None:
+    """Run an 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())
+    await service.run(index_id)
+
+
+@cli.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
+) -> None:
+    """Retrieve snippets from the database."""
+    repository = RetrievalRepository(session)
+    service = RetrievalService(repository, app_context.get_data_dir())
+    # Temporary request while we don't have all search capabilities
+    snippets = await service.retrieve(
+        RetrievalRequest(keywords=query.split(","), 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()
+
+
+@cli.command()
+@click.option("--host", default="127.0.0.1", help="Host to bind the server to")
+@click.option("--port", default=8080, help="Port to bind the server to")
+@with_app_context
+def serve(
+    app_context: AppContext,
+    host: str,
+    port: int,
+) -> None:
+    """Start the kodit server, which hosts the MCP server and the kodit API."""
+    log = structlog.get_logger(__name__)
+    log.info("Starting kodit server", host=host, port=port)
+    log_event("kodit_server_started")
+
+    # Dump AppContext to a dictionary of strings, and set the env vars
+    app_context_dict = {k: str(v) for k, v in app_context.model_dump().items()}
+    os.environ.update(app_context_dict)
+
+    # Configure uvicorn with graceful shutdown
+    config = uvicorn.Config(
+        "kodit.app:app",
+        host=host,
+        port=port,
+        reload=False,
+        log_config=None,  # Setting to None forces uvicorn to use our structlog setup
+        access_log=False,  # Using own middleware for access logging
+        timeout_graceful_shutdown=0,  # The mcp server does not shutdown cleanly, force
+    )
+    server = uvicorn.Server(config)
+
+    def handle_sigint(signum: int, frame: Any) -> None:
+        """Handle SIGINT (Ctrl+C)."""
+        log.info("Received shutdown signal, force killing MCP connections")
+        server.handle_exit(signum, frame)
+
+    signal.signal(signal.SIGINT, handle_sigint)
+    server.run()
+
+
+@cli.command()
+def version() -> None:
+    """Show the version of kodit."""
+    try:
+        from kodit import _version
+    except ImportError:
+        print("unknown, try running `uv build`, which is what happens in ci")  # noqa: T201
+    else:
+        print(_version.version)  # noqa: T201
+
+
+if __name__ == "__main__":
+    cli()

kodit-0.1.6/src/kodit/config.py

@@ -0,0 +1,97 @@
+"""Global configuration for the kodit project."""
+
+import asyncio
+from collections.abc import Callable, Coroutine
+from functools import wraps
+from pathlib import Path
+from typing import Any, TypeVar
+
+import click
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from kodit.database import Database
+
+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
+T = TypeVar("T")
+
+
+class AppContext(BaseSettings):
+    """Global context for the kodit project. Provides a shared state for the app."""
+
+    model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
+
+    data_dir: Path = Field(default=DEFAULT_BASE_DIR)
+    db_url: str = Field(default=DEFAULT_DB_URL)
+    log_level: str = Field(default=DEFAULT_LOG_LEVEL)
+    log_format: str = Field(default=DEFAULT_LOG_FORMAT)
+    disable_telemetry: bool = Field(default=DEFAULT_DISABLE_TELEMETRY)
+    _db: Database | None = None
+
+    def model_post_init(self, _: Any) -> None:
+        """Post-initialization hook."""
+        # Call this to ensure the data dir exists for the default db location
+        self.get_data_dir()
+
+    def get_data_dir(self) -> Path:
+        """Get the data directory."""
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+        return self.data_dir
+
+    def get_clone_dir(self) -> Path:
+        """Get the clone directory."""
+        clone_dir = self.get_data_dir() / "clones"
+        clone_dir.mkdir(parents=True, exist_ok=True)
+        return clone_dir
+
+    async def get_db(self, *, run_migrations: bool = True) -> Database:
+        """Get the database."""
+        if self._db is None:
+            self._db = Database(self.db_url)
+        if run_migrations:
+            await self._db.run_migrations(self.db_url)
+        return self._db
+
+
+with_app_context = click.make_pass_decorator(AppContext)
+
+T = TypeVar("T")
+
+
+def wrap_async(f: Callable[..., Coroutine[Any, Any, T]]) -> Callable[..., T]:
+    """Decorate async Click commands.
+
+    This decorator wraps an async function to run it with asyncio.run().
+    It should be used after the Click command decorator.
+
+    Example:
+        @cli.command()
+        @wrap_async
+        async def my_command():
+            ...
+
+    """
+
+    @wraps(f)
+    def wrapper(*args: Any, **kwargs: Any) -> T:
+        return asyncio.run(f(*args, **kwargs))
+
+    return wrapper
+
+
+def with_session(f: Callable[..., Coroutine[Any, Any, T]]) -> Callable[..., T]:
+    """Provide a database session to CLI commands."""
+
+    @wraps(f)
+    @with_app_context
+    @wrap_async
+    async def wrapper(app_context: AppContext, *args: Any, **kwargs: Any) -> T:
+        db = await app_context.get_db()
+        async with db.session_factory() as session:
+            return await f(session, *args, **kwargs)
+
+    return wrapper