kodit 0.1.3__tar.gz → 0.1.5__tar.gz

kodit-0.1.5/.cursor/rules/kodit.mdc

@@ -0,0 +1,6 @@
+---
+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.

{kodit-0.1.3 → kodit-0.1.5}/.github/workflows/pypi.yaml

@@ -28,10 +28,26 @@ jobs:
     steps:
       - name: Wait for PyPI release to be available
         run: |
+          set -x
+
+          REPO_NAME=${{ github.event.repository.name }}
+          REPO_TAG=${{ github.event.release.tag_name }}
+
+          # Get the first letter of the repo name
+          REPO_NAME_FIRST_LETTER=${REPO_NAME:0:1}
+
+          # Initialize counter for 60 second timeout
+          count=0
           while true; do
-            if curl -s https://pypi.io/packages/source/k/${{ github.event.repository.name }}/${{ github.event.repository.name }}-${{ github.event.release.tag_name }}.tar.gz; then
-              break
-            fi
+              if curl -sfL https://pypi.io/packages/source/${REPO_NAME_FIRST_LETTER}/${REPO_NAME}/${REPO_NAME}-${REPO_TAG}.tar.gz > /dev/null; then
+                  break
+              fi
+              sleep 1
+              ((count++))
+              if [ $count -ge 60 ]; then
+                  echo "Timeout reached after 60 seconds"
+                  exit 1
+              fi
           done
       - uses: mislav/bump-homebrew-formula-action@v3
         with:

{kodit-0.1.3 → kodit-0.1.5}/.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:

kodit-0.1.5/.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.3 → kodit-0.1.5}/.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.3 → kodit-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.1.3
+Version: 0.1.5
 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.5/docs/_index.md

@@ -0,0 +1,93 @@
+---
+title: "kodit: Code Indexing MCP Server"
+linkTitle: kodit Docs
+cascade:
+  type: docs
+menu:
+  main:
+    name: kodit Docs
+    weight: 3
+# next: /helix/getting-started
+weight: 1
+aliases:
+- /coda
+---
+
+## Installation
+
+Please choose your preferred installation method. They all ultimately install the kodit
+cli, which contains the kodit MCP server and other tools to manage your data sources.
+
+### Docker
+
+```sh
+docker run -it --rm registry.helix.ml/helix/kodit:latest
+```
+
+Always replace latest with a specific version.
+
+### pipx
+
+```sh
+pipx install kodit
+```
+
+### homebrew
+
+```sh
+brew install helixml/kodit/kodit
+```
+
+### uv
+
+```sh
+uv tool install kodit
+```
+
+### pip
+
+Use this if you want to use kodit as a python library:
+
+```sh
+pip install kodit
+```
+
+## Usage
+
+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. Add a source: `kodit sources create /path/to/your/code`
+2. Create an index on that source: `kodit indexes create 1`
+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!
+
+### Integration with Cursor
+
+Add the following to `.cursor/mcp.json`:
+
+```json
+{
+  "mcpServers": {
+    "code-search": {
+      "url": "http://localhost:8080/sse"
+    }
+  }
+}
+```
+
+Now add the following prompt to `.cursor/rules/kodit.mdc` in your project directory:
+
+```markdown
+---
+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.
+```
+
+Alternatively, you can browse to the cursor settings and set this prompt globally.

kodit-0.1.5/docs/developer/index.md

@@ -0,0 +1,30 @@
+---
+title: "Kodit Developer Documentation"
+linkTitle: Developer Docs
+# next: /helix/getting-started
+weight: 99
+---
+
+## Database
+
+All database operations are handled by SQLAlchemy and Alembic.
+
+### Creating a Database Migration
+
+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
+
+## Releasing
+
+Performing a release is designed to be fully automated. If you spot opportunities to
+improve the CI to help performing an automated release, please do so.
+
+1. Create a new release in GitHub.
+2. Set the version number. Use patch versions for bugfixes or minor small improvements.
+   Use minor versions when adding significant new functionality. Use major versions for
+   overhauls.
+3. Generate the release notes. <- this could be improved, because we use a strict
+   pr/commit naming structure.
+4. Wait for all jobs to succeed, then you should be able to brew install, pipx install, etc.

{kodit-0.1.3 → kodit-0.1.5}/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.3 → kodit-0.1.5}/src/kodit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.3'
-__version_tuple__ = version_tuple = (0, 1, 3)
+__version__ = version = '0.1.5'
+__version_tuple__ = version_tuple = (0, 1, 5)

{kodit-0.1.3 → kodit-0.1.5}/src/kodit/alembic/env.py

@@ -66,8 +66,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)

{kodit-0.1.3 → kodit-0.1.5}/src/kodit/app.py

@@ -5,14 +5,10 @@ from fastapi import FastAPI
 
 from kodit.mcp import mcp
 from kodit.middleware import logging_middleware
-from kodit.sse import create_sse_server
 
-app = FastAPI(title="kodit API")
-
-# Get the SSE routes from the Starlette app hosting the MCP server
-sse_app = create_sse_server(mcp)
-for route in sse_app.routes:
-    app.router.routes.append(route)
+# 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)
@@ -22,4 +18,8 @@ app.add_middleware(CorrelationIdMiddleware)
 @app.get("/")
 async def root() -> dict[str, str]:
     """Return a welcome message for the kodit API."""
-    return {"message": "Welcome to kodit API"}
+    return {"message": "Hello, World!"}
+
+
+# Add mcp routes last, otherwise previous routes aren't added
+app.mount("", mcp_app)

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

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

kodit-0.1.5/src/kodit/bm25/bm25.py

@@ -0,0 +1,71 @@
+"""BM25 service."""
+
+import bm25s
+import Stemmer
+import structlog
+from bm25s.tokenization import Tokenized
+
+from kodit.config import Config
+
+
+class BM25Service:
+    """Service for BM25."""
+
+    def __init__(self, config: Config) -> None:
+        """Initialize the BM25 service."""
+        self.log = structlog.get_logger(__name__)
+        self.index_path = config.get_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.3 → kodit-0.1.5}/src/kodit/cli.py

@@ -1,41 +1,74 @@
 """Command line interface for kodit."""
 
 import os
+import signal
+from pathlib import Path
+from typing import Any
 
 import click
 import structlog
 import uvicorn
-from dotenv import dotenv_values
 from pytable_formatter import Table
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from kodit.database import configure_database, with_session
+from kodit.config import (
+    DEFAULT_BASE_DIR,
+    DEFAULT_DB_URL,
+    DEFAULT_DISABLE_TELEMETRY,
+    DEFAULT_LOG_FORMAT,
+    DEFAULT_LOG_LEVEL,
+    get_config,
+    reset_config,
+    with_session,
+)
 from kodit.indexing.repository import IndexRepository
 from kodit.indexing.service import IndexService
-from kodit.logging import LogFormat, configure_logging, disable_posthog, log_event
+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
 
-env_vars = dict(dotenv_values())
-os.environ.update(env_vars)
 
-
-@click.group(context_settings={"auto_envvar_prefix": "KODIT", "show_default": True})
-@click.option("--log-level", default="INFO", help="Log level")
-@click.option("--log-format", default=LogFormat.PRETTY, help="Log format")
-@click.option("--disable-telemetry", is_flag=True, help="Disable telemetry")
-def cli(
-    log_level: str,
-    log_format: LogFormat,
-    disable_telemetry: bool,  # noqa: FBT001
+@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]")
+def cli(  # noqa: PLR0913
+    log_level: str | None,
+    log_format: str | None,
+    disable_telemetry: bool | None,
+    db_url: str | None,
+    data_dir: str | None,
+    env_file: str | None,
 ) -> None:
     """kodit CLI - Code indexing for better AI code generation."""  # noqa: D403
-    configure_logging(log_level, log_format)
+    # First check if env-file is set and reload config if it is
+    if env_file:
+        reset_config()
+        get_config(env_file)
+
+    # Override global config with cli args, if set
+    config = get_config()
+    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:
-        disable_posthog()
-    configure_database()
+        config.disable_telemetry = disable_telemetry
+    configure_logging(config)
+    configure_telemetry(config)
 
 
 @cli.group()
@@ -48,7 +81,7 @@ def sources() -> None:
 async def list_sources(session: AsyncSession) -> None:
     """List all code sources."""
     repository = SourceRepository(session)
-    service = SourceService(repository)
+    service = SourceService(get_config().get_clone_dir(), repository)
     sources = await service.list_sources()
 
     # Define headers and data
@@ -66,7 +99,7 @@ async def list_sources(session: AsyncSession) -> None:
 async def create_source(session: AsyncSession, uri: str) -> None:
     """Add a new code source."""
     repository = SourceRepository(session)
-    service = SourceService(repository)
+    service = SourceService(get_config().get_clone_dir(), repository)
     source = await service.create(uri)
     click.echo(f"Source created: {source.id}")
 
@@ -82,9 +115,9 @@ def indexes() -> None:
 async def create_index(session: AsyncSession, source_id: int) -> None:
     """Create an index for a source."""
     source_repository = SourceRepository(session)
-    source_service = SourceService(source_repository)
+    source_service = SourceService(get_config().get_clone_dir(), source_repository)
     repository = IndexRepository(session)
-    service = IndexService(repository, source_service)
+    service = IndexService(get_config(), repository, source_service)
     index = await service.create(source_id)
     click.echo(f"Index created: {index.id}")
 
@@ -94,9 +127,9 @@ async def create_index(session: AsyncSession, source_id: int) -> None:
 async def list_indexes(session: AsyncSession) -> None:
     """List all indexes."""
     source_repository = SourceRepository(session)
-    source_service = SourceService(source_repository)
+    source_service = SourceService(get_config().get_clone_dir(), source_repository)
     repository = IndexRepository(session)
-    service = IndexService(repository, source_service)
+    service = IndexService(get_config(), repository, source_service)
     indexes = await service.list_indexes()
 
     # Define headers and data
@@ -104,7 +137,6 @@ async def list_indexes(session: AsyncSession) -> None:
         "ID",
         "Created At",
         "Updated At",
-        "Source URI",
         "Num Snippets",
     ]
     data = [
@@ -112,7 +144,6 @@ async def list_indexes(session: AsyncSession) -> None:
             index.id,
             index.created_at,
             index.updated_at,
-            index.source_uri,
             index.num_snippets,
         ]
         for index in indexes
@@ -129,48 +160,69 @@ async def list_indexes(session: AsyncSession) -> None:
 async def run_index(session: AsyncSession, index_id: int) -> None:
     """Run an index."""
     source_repository = SourceRepository(session)
-    source_service = SourceService(source_repository)
+    source_service = SourceService(get_config().get_clone_dir(), source_repository)
     repository = IndexRepository(session)
-    service = IndexService(repository, source_service)
+    service = IndexService(get_config(), repository, source_service)
     await service.run(index_id)
 
 
 @cli.command()
 @click.argument("query")
+@click.option("--top-k", default=10, help="Number of snippets to retrieve")
 @with_session
-async def retrieve(session: AsyncSession, query: str) -> None:
+async def retrieve(session: AsyncSession, query: str, top_k: int) -> None:
     """Retrieve snippets from the database."""
     repository = RetrievalRepository(session)
-    service = RetrievalService(repository)
-    snippets = await service.retrieve(RetrievalRequest(query=query))
+    service = RetrievalService(get_config(), repository)
+    # 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")
-@click.option("--reload", is_flag=True, help="Enable auto-reload for development")
 def serve(
     host: str,
     port: int,
-    reload: bool,  # noqa: FBT001
 ) -> 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, reload=reload)
+    log.info("Starting kodit server", host=host, port=port)
     log_event("kodit_server_started")
-    uvicorn.run(
+    os.environ["HELLO"] = "WORLD"
+
+    # Configure uvicorn with graceful shutdown
+    config = uvicorn.Config(
         "kodit.app:app",
         host=host,
         port=port,
-        reload=reload,
+        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()