kodit 0.1.4__py3-none-any.whl → 0.1.6__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.4'
-__version_tuple__ = version_tuple = (0, 1, 4)
+__version__ = version = '0.1.6'
+__version_tuple__ = version_tuple = (0, 1, 6)

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/app.py

@@ -4,15 +4,11 @@ from asgi_correlation_id import CorrelationIdMiddleware
 from fastapi import FastAPI
 
 from kodit.mcp import mcp
-from kodit.middleware import logging_middleware
-from kodit.sse import create_sse_server
+from kodit.middleware import ASGICancelledErrorMiddleware, logging_middleware
 
-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,12 @@ 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)
+
+# Wrap the entire app with ASGI middleware after all routes are added to suppress
+# CancelledError at the ASGI level
+app = ASGICancelledErrorMiddleware(app)

kodit/bm25/__init__.py

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

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/cli.py

@@ -1,41 +1,87 @@
 """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,
+    AppContext,
+    with_app_context,
+    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]",
+    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
-    configure_logging(log_level, log_format)
+    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:
-        disable_posthog()
-    configure_database()
+        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()
@@ -44,11 +90,12 @@ def sources() -> None:
 
 
 @sources.command(name="list")
+@with_app_context
 @with_session
-async def list_sources(session: AsyncSession) -> None:
+async def list_sources(session: AsyncSession, app_context: AppContext) -> None:
     """List all code sources."""
     repository = SourceRepository(session)
-    service = SourceService(repository)
+    service = SourceService(app_context.get_clone_dir(), repository)
     sources = await service.list_sources()
 
     # Define headers and data
@@ -62,11 +109,14 @@ async def list_sources(session: AsyncSession) -> None:
 
 @sources.command(name="create")
 @click.argument("uri")
+@with_app_context
 @with_session
-async def create_source(session: AsyncSession, uri: str) -> None:
+async def create_source(
+    session: AsyncSession, app_context: AppContext, uri: str
+) -> None:
     """Add a new code source."""
     repository = SourceRepository(session)
-    service = SourceService(repository)
+    service = SourceService(app_context.get_clone_dir(), repository)
     source = await service.create(uri)
     click.echo(f"Source created: {source.id}")
 
@@ -78,25 +128,29 @@ def indexes() -> None:
 
 @indexes.command(name="create")
 @click.argument("source_id")
+@with_app_context
 @with_session
-async def create_index(session: AsyncSession, source_id: int) -> None:
+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(source_repository)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
-    service = IndexService(repository, source_service)
+    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) -> None:
+async def list_indexes(session: AsyncSession, app_context: AppContext) -> None:
     """List all indexes."""
     source_repository = SourceRepository(session)
-    source_service = SourceService(source_repository)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
-    service = IndexService(repository, source_service)
+    service = IndexService(repository, source_service, app_context.get_data_dir())
     indexes = await service.list_indexes()
 
     # Define headers and data
@@ -123,52 +177,84 @@ async def list_indexes(session: AsyncSession) -> None:
 
 @indexes.command(name="run")
 @click.argument("index_id")
+@with_app_context
 @with_session
-async def run_index(session: AsyncSession, index_id: int) -> None:
+async def run_index(
+    session: AsyncSession, app_context: AppContext, index_id: int
+) -> None:
     """Run an index."""
     source_repository = SourceRepository(session)
-    source_service = SourceService(source_repository)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
-    service = IndexService(repository, source_service)
+    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, query: str) -> None:
+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)
-    snippets = await service.retrieve(RetrievalRequest(query=query))
+    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")
-@click.option("--reload", is_flag=True, help="Enable auto-reload for development")
+@with_app_context
 def serve(
+    app_context: AppContext,
     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(
+
+    # 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=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/config.py

@@ -1,5 +1,97 @@
-"""Configuration for the kodit project."""
+"""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
 
-DATA_DIR = Path.home() / ".kodit"
+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

kodit/database.py

@@ -1,15 +1,11 @@
 """Database configuration for kodit."""
 
-import asyncio
-from collections.abc import AsyncGenerator, Callable
-from contextlib import asynccontextmanager
 from datetime import UTC, datetime
-from functools import wraps
 from pathlib import Path
-from typing import Any, TypeVar
 
+import structlog
 from alembic import command
-from alembic.config import Config
+from alembic.config import Config as AlembicConfig
 from sqlalchemy import DateTime
 from sqlalchemy.ext.asyncio import (
     AsyncAttrs,
@@ -20,23 +16,6 @@ from sqlalchemy.ext.asyncio import (
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 
 from kodit import alembic
-from kodit.config import DATA_DIR
-
-# Constants
-DB_URL = f"sqlite+aiosqlite:///{DATA_DIR}/kodit.db"
-
-# Create data directory if it doesn't exist
-DATA_DIR.mkdir(exist_ok=True)
-
-# Create async engine with file-based SQLite
-engine = create_async_engine(DB_URL, echo=False)
-
-# Create async session factory
-async_session_factory = async_sessionmaker(
-    engine,
-    class_=AsyncSession,
-    expire_on_commit=False,
-)
 
 
 class Base(AsyncAttrs, DeclarativeBase):
@@ -55,37 +34,42 @@ class CommonMixin:
     )
 
 
-@asynccontextmanager
-async def get_session() -> AsyncGenerator[AsyncSession, None]:
-    """Get a database session."""
-    async with async_session_factory() as session:
-        try:
-            yield session
-        finally:
-            await session.close()
-
-
-T = TypeVar("T")
-
-
-def with_session(func: Callable[..., T]) -> Callable[..., T]:
-    """Provide an async session to CLI commands."""
-
-    @wraps(func)
-    def wrapper(*args: Any, **kwargs: Any) -> T:
-        async def _run() -> T:
-            async with async_session_factory() as session:
-                return await func(session, *args, **kwargs)
-
-        return asyncio.run(_run())
-
-    return wrapper
-
-
-def configure_database() -> None:
-    """Configure the database by initializing it and running any pending migrations."""
-    # Create Alembic configuration and run migrations
-    alembic_cfg = Config()
-    alembic_cfg.set_main_option("script_location", str(Path(alembic.__file__).parent))
-    alembic_cfg.set_main_option("sqlalchemy.url", DB_URL)
-    command.upgrade(alembic_cfg, "head")
+class Database:
+    """Database class for kodit."""
+
+    def __init__(self, db_url: str) -> None:
+        """Initialize the database."""
+        self.log = structlog.get_logger(__name__)
+        self.db_engine = create_async_engine(db_url, echo=False)
+        self.db_session_factory = async_sessionmaker(
+            self.db_engine,
+            class_=AsyncSession,
+            expire_on_commit=False,
+        )
+
+    @property
+    def session_factory(self) -> async_sessionmaker[AsyncSession]:
+        """Get the session factory."""
+        return self.db_session_factory
+
+    async 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)
+
+        async with self.db_engine.begin() as conn:
+            await conn.run_sync(self.run_upgrade, alembic_cfg)
+
+    def run_upgrade(self, connection, cfg) -> None:  # noqa: ANN001
+        """Make sure the database is up to date."""
+        cfg.attributes["connection"] = connection
+        command.upgrade(cfg, "head")
+
+    async def close(self) -> None:
+        """Close the database."""
+        await self.db_engine.dispose()

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())