kodit 0.1.3__py3-none-any.whl → 0.1.5__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.3'
-__version_tuple__ = version_tuple = (0, 1, 3)
+__version__ = version = '0.1.5'
+__version_tuple__ = version_tuple = (0, 1, 5)

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/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/bm25/__init__.py

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

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

kodit/config.py

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

@@ -1,14 +1,13 @@
 """Database configuration for kodit."""
 
-import asyncio
-from collections.abc import AsyncGenerator, Callable
+from collections.abc import AsyncGenerator
+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,
@@ -19,23 +18,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):
@@ -54,36 +36,37 @@ class CommonMixin:
     )
 
 
-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, *, 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/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())

kodit/indexing/service.py

@@ -7,24 +7,21 @@ index management.
 """
 
 from datetime import datetime
+from pathlib import Path
 
-import aiofiles
 import pydantic
 import structlog
 from tqdm.asyncio import tqdm
 
+from kodit.bm25.bm25 import BM25Service
+from kodit.config import Config
 from kodit.indexing.models import Snippet
 from kodit.indexing.repository import IndexRepository
+from kodit.snippets.snippets import SnippetService
 from kodit.sources.service import SourceService
 
-# List of MIME types that are supported for indexing and snippet creation
-MIME_WHITELIST = [
-    "text/plain",
-    "text/markdown",
-    "text/x-python",
-    "text/x-shellscript",
-    "text/x-sql",
-]
+# List of MIME types that are blacklisted from being indexed
+MIME_BLACKLIST = ["unknown/unknown"]
 
 
 class IndexView(pydantic.BaseModel):
@@ -37,7 +34,6 @@ class IndexView(pydantic.BaseModel):
     id: int
     created_at: datetime
     updated_at: datetime | None = None
-    source_uri: str | None = None
     num_snippets: int | None = None
 
 
@@ -50,7 +46,7 @@ class IndexService:
     """
 
     def __init__(
-        self, repository: IndexRepository, source_service: SourceService
+        self, config: Config, repository: IndexRepository, source_service: SourceService
     ) -> None:
         """Initialize the index service.
 
@@ -61,7 +57,9 @@ class IndexService:
         """
         self.repository = repository
         self.source_service = source_service
+        self.snippet_service = SnippetService()
         self.log = structlog.get_logger(__name__)
+        self.bm25 = BM25Service(config)
 
     async def create(self, source_id: int) -> IndexView:
         """Create a new index for a source.
@@ -120,6 +118,10 @@ 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])
+
         # Update index timestamp
         await self.repository.update_index_timestamp(index)
 
@@ -138,16 +140,23 @@ class IndexService:
         files = await self.repository.files_for_index(index_id)
         for file in tqdm(files, total=len(files)):
             # Skip unsupported file types
-            if file.mime_type not in MIME_WHITELIST:
+            if file.mime_type in MIME_BLACKLIST:
                 self.log.debug("Skipping mime type", mime_type=file.mime_type)
                 continue
 
             # Create snippet from file content
-            async with aiofiles.open(file.cloned_path, "rb") as f:
-                content = await f.read()
-                snippet = Snippet(
+            try:
+                snippets = self.snippet_service.snippets_for_file(
+                    Path(file.cloned_path)
+                )
+            except ValueError as e:
+                self.log.debug("Skipping file", file=file.cloned_path, error=e)
+                continue
+
+            for snippet in snippets:
+                s = Snippet(
                     index_id=index_id,
                     file_id=file.id,
-                    content=content.decode("utf-8"),
+                    content=snippet.text,
                 )
-                await self.repository.add_snippet(snippet)
+                await self.repository.add_snippet(s)