kodit 0.1.4__tar.gz → 0.1.5__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.1.4
+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.4 → kodit-0.1.5}/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.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.4 → 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.4'
-__version_tuple__ = version_tuple = (0, 1, 4)
+__version__ = version = '0.1.5'
+__version_tuple__ = version_tuple = (0, 1, 5)

{kodit-0.1.4 → 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.4 → 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.4 → 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
@@ -127,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()

kodit-0.1.5/src/kodit/config.py

@@ -0,0 +1,89 @@
+"""Global configuration for the kodit project."""
+
+import asyncio
+from collections.abc import Callable
+from functools import wraps
+from pathlib import Path
+from typing import Any, TypeVar
+
+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 Config(BaseSettings):
+    """Global configuration for the kodit project."""
+
+    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
+
+    def get_db(self, *, run_migrations: bool = True) -> Database:
+        """Get the database."""
+        if self._db is None:
+            self._db = Database(self.db_url, run_migrations=run_migrations)
+        return self._db
+
+
+# Global config instance for mcp Apps
+config = None
+
+
+def get_config(env_file: str | None = None) -> Config:
+    """Get the global config instance."""
+    global config  # noqa: PLW0603
+    if config is None:
+        config = Config(_env_file=env_file)
+    return config
+
+
+def reset_config() -> None:
+    """Reset the global config instance."""
+    global config  # noqa: PLW0603
+    config = None
+
+
+def with_session(func: Callable[..., T]) -> Callable[..., T]:
+    """Provide an async session to CLI commands."""
+
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> T:
+        # Create DB connection before starting event loop
+        db = get_config().get_db()
+
+        async def _run() -> T:
+            async with db.get_session() as session:
+                return await func(session, *args, **kwargs)
+
+        return asyncio.run(_run())
+
+    return wrapper

kodit-0.1.5/src/kodit/database.py

@@ -0,0 +1,72 @@
+"""Database configuration for kodit."""
+
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from datetime import UTC, datetime
+from pathlib import Path
+
+import structlog
+from alembic import command
+from alembic.config import Config as AlembicConfig
+from sqlalchemy import DateTime
+from sqlalchemy.ext.asyncio import (
+    AsyncAttrs,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+from kodit import alembic
+
+
+class Base(AsyncAttrs, DeclarativeBase):
+    """Base class for all models."""
+
+
+class CommonMixin:
+    """Common mixin for all models."""
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime, default=lambda: datetime.now(UTC)
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)
+    )
+
+
+class Database:
+    """Database class for kodit."""
+
+    def __init__(self, db_url: str, *, run_migrations: bool = True) -> None:
+        """Initialize the database."""
+        self.log = structlog.get_logger(__name__)
+        if run_migrations:
+            self._run_migrations(db_url)
+        db_engine = create_async_engine(db_url, echo=False)
+        self.db_session_factory = async_sessionmaker(
+            db_engine,
+            class_=AsyncSession,
+            expire_on_commit=False,
+        )
+
+    @asynccontextmanager
+    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
+        """Get a database session."""
+        async with self.db_session_factory() as session:
+            try:
+                yield session
+            finally:
+                await session.close()
+
+    def _run_migrations(self, db_url: str) -> None:
+        """Run any pending migrations."""
+        # Create Alembic configuration and run migrations
+        alembic_cfg = AlembicConfig()
+        alembic_cfg.set_main_option(
+            "script_location", str(Path(alembic.__file__).parent)
+        )
+        alembic_cfg.set_main_option("sqlalchemy.url", db_url)
+        self.log.debug("Running migrations", db_url=db_url)
+        command.upgrade(alembic_cfg, "head")

{kodit-0.1.4 → kodit-0.1.5}/src/kodit/indexing/repository.py

@@ -130,3 +130,14 @@ class IndexRepository:
         query = select(Snippet).where(Snippet.index_id == index_id)
         result = await self.session.execute(query)
         return list(result.scalars())
+
+    async def get_all_snippets(self) -> list[Snippet]:
+        """Get all snippets.
+
+        Returns:
+            A list of all snippets.
+
+        """
+        query = select(Snippet).order_by(Snippet.id)
+        result = await self.session.execute(query)
+        return list(result.scalars())