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

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
@@ -75,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,7 +4,7 @@ from asgi_correlation_id import CorrelationIdMiddleware
 from fastapi import FastAPI
 
 from kodit.mcp import mcp
-from kodit.middleware import logging_middleware
+from kodit.middleware import ASGICancelledErrorMiddleware, logging_middleware
 
 # See https://gofastmcp.com/deployment/asgi#fastapi-integration
 mcp_app = mcp.sse_app()
@@ -23,3 +23,7 @@ async def root() -> dict[str, str]:
 
 # 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/bm25.py

@@ -1,20 +1,20 @@
 """BM25 service."""
 
+from pathlib import Path
+
 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:
+    def __init__(self, data_dir: Path) -> None:
         """Initialize the BM25 service."""
         self.log = structlog.get_logger(__name__)
-        self.index_path = config.get_data_dir() / "bm25s_index"
+        self.index_path = data_dir / "bm25s_index"
         try:
             self.log.debug("Loading BM25 index")
             self.retriever = bm25s.BM25.load(self.index_path, mmap=True)

kodit/cli.py

@@ -8,7 +8,7 @@ from typing import Any
 import click
 import structlog
 import uvicorn
-from pytable_formatter import Table
+from pytable_formatter import Cell, Table
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.config import (
@@ -17,8 +17,8 @@ from kodit.config import (
     DEFAULT_DISABLE_TELEMETRY,
     DEFAULT_LOG_FORMAT,
     DEFAULT_LOG_LEVEL,
-    get_config,
-    reset_config,
+    AppContext,
+    with_app_context,
     with_session,
 )
 from kodit.indexing.repository import IndexRepository
@@ -40,23 +40,33 @@ from kodit.sources.service import SourceService
 )
 @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]")
+@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: 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:
-        reset_config()
-        get_config(env_file)
+        config = AppContext(_env_file=env_file)  # type: ignore[reportCallIssue]
 
-    # Override global config with cli args, if set
-    config = get_config()
+    # Now override with CLI arguments, if set
     if data_dir:
         config.data_dir = Path(data_dir)
     if db_url:
@@ -70,110 +80,76 @@ def cli(  # noqa: PLR0913
     configure_logging(config)
     configure_telemetry(config)
 
-
-@cli.group()
-def sources() -> None:
-    """Manage code sources."""
-
-
-@sources.command(name="list")
-@with_session
-async def list_sources(session: AsyncSession) -> None:
-    """List all code sources."""
-    repository = SourceRepository(session)
-    service = SourceService(get_config().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_session
-async def create_source(session: AsyncSession, uri: str) -> None:
-    """Add a new code source."""
-    repository = SourceRepository(session)
-    service = SourceService(get_config().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_session
-async def create_index(session: AsyncSession, source_id: int) -> None:
-    """Create an index for a source."""
-    source_repository = SourceRepository(session)
-    source_service = SourceService(get_config().get_clone_dir(), source_repository)
-    repository = IndexRepository(session)
-    service = IndexService(get_config(), repository, source_service)
-    index = await service.create(source_id)
-    click.echo(f"Index created: {index.id}")
+    # Set the app context in the click context for downstream cli
+    ctx.obj = config
 
 
-@indexes.command(name="list")
+@cli.command()
+@click.argument("sources", nargs=-1)
+@with_app_context
 @with_session
-async def list_indexes(session: AsyncSession) -> None:
-    """List all indexes."""
+async def index(
+    session: AsyncSession,
+    app_context: AppContext,
+    sources: list[str],
+) -> None:
+    """List indexes, or index data sources."""
     source_repository = SourceRepository(session)
-    source_service = SourceService(get_config().get_clone_dir(), source_repository)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
-    service = IndexService(get_config(), repository, source_service)
-    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,
+    service = IndexService(repository, source_service, app_context.get_data_dir())
+
+    if not sources:
+        # No source specified, list all indexes
+        indexes = await service.list_indexes()
+        headers: list[str | Cell] = [
+            "ID",
+            "Created At",
+            "Updated At",
+            "Source",
+            "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_session
-async def run_index(session: AsyncSession, index_id: int) -> None:
-    """Run an index."""
-    source_repository = SourceRepository(session)
-    source_service = SourceService(get_config().get_clone_dir(), source_repository)
-    repository = IndexRepository(session)
-    service = IndexService(get_config(), repository, source_service)
-    await service.run(index_id)
+        data = [
+            [
+                index.id,
+                index.created_at,
+                index.updated_at,
+                index.source,
+                index.num_snippets,
+            ]
+            for index in indexes
+        ]
+        click.echo(Table(headers=headers, data=data))
+        return
+    # Handle source indexing
+    for source in sources:
+        if source.startswith("https://"):
+            msg = "Web or git indexing is not implemented yet"
+            raise click.UsageError(msg)
+        if source.startswith("git"):
+            msg = "Git indexing is not implemented yet"
+            raise click.UsageError(msg)
+        if Path(source).is_file():
+            msg = "File indexing is not implemented yet"
+            raise click.UsageError(msg)
+
+        # Index directory
+        s = await source_service.create(source)
+        index = await service.create(s.id)
+        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, top_k: int) -> 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(get_config(), repository)
+    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)
@@ -194,7 +170,9 @@ async def retrieve(session: AsyncSession, query: str, top_k: int) -> None:
 @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:
@@ -202,7 +180,10 @@ def serve(
     log = structlog.get_logger(__name__)
     log.info("Starting kodit server", host=host, port=port)
     log_event("kodit_server_started")
-    os.environ["HELLO"] = "WORLD"
+
+    # 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/config.py

@@ -1,11 +1,12 @@
 """Global configuration for the kodit project."""
 
 import asyncio
-from collections.abc import Callable
+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
 
@@ -19,8 +20,8 @@ DEFAULT_DISABLE_TELEMETRY = False
 T = TypeVar("T")
 
 
-class Config(BaseSettings):
-    """Global configuration for the kodit project."""
+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")
 
@@ -47,43 +48,50 @@ class Config(BaseSettings):
         clone_dir.mkdir(parents=True, exist_ok=True)
         return clone_dir
 
-    def get_db(self, *, run_migrations: bool = True) -> Database:
+    async 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)
+            self._db = Database(self.db_url)
+        if run_migrations:
+            await self._db.run_migrations(self.db_url)
         return self._db
 
 
-# Global config instance for mcp Apps
-config = None
+with_app_context = click.make_pass_decorator(AppContext)
 
+T = TypeVar("T")
 
-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 wrap_async(f: Callable[..., Coroutine[Any, Any, T]]) -> Callable[..., T]:
+    """Decorate async Click commands.
 
-def reset_config() -> None:
-    """Reset the global config instance."""
-    global config  # noqa: PLW0603
-    config = None
+    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():
+            ...
 
-def with_session(func: Callable[..., T]) -> Callable[..., T]:
-    """Provide an async session to CLI commands."""
+    """
 
-    @wraps(func)
+    @wraps(f)
     def wrapper(*args: Any, **kwargs: Any) -> T:
-        # Create DB connection before starting event loop
-        db = get_config().get_db()
+        return asyncio.run(f(*args, **kwargs))
+
+    return wrapper
+
 
-        async def _run() -> T:
-            async with db.get_session() as session:
-                return await func(session, *args, **kwargs)
+def with_session(f: Callable[..., Coroutine[Any, Any, T]]) -> Callable[..., T]:
+    """Provide a database session to CLI commands."""
 
-        return asyncio.run(_run())
+    @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,7 +1,5 @@
 """Database configuration for kodit."""
 
-from collections.abc import AsyncGenerator
-from contextlib import asynccontextmanager
 from datetime import UTC, datetime
 from pathlib import Path
 
@@ -39,28 +37,22 @@ class CommonMixin:
 class Database:
     """Database class for kodit."""
 
-    def __init__(self, db_url: str, *, run_migrations: bool = True) -> None:
+    def __init__(self, db_url: str) -> 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_engine = create_async_engine(db_url, echo=False)
         self.db_session_factory = async_sessionmaker(
-            db_engine,
+            self.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()
+    @property
+    def session_factory(self) -> async_sessionmaker[AsyncSession]:
+        """Get the session factory."""
+        return self.db_session_factory
 
-    def _run_migrations(self, db_url: str) -> None:
+    async def run_migrations(self, db_url: str) -> None:
         """Run any pending migrations."""
         # Create Alembic configuration and run migrations
         alembic_cfg = AlembicConfig()
@@ -69,4 +61,15 @@ class Database:
         )
         alembic_cfg.set_main_option("sqlalchemy.url", db_url)
         self.log.debug("Running migrations", db_url=db_url)
-        command.upgrade(alembic_cfg, "head")
+
+        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

@@ -82,7 +82,7 @@ class IndexRepository:
         result = await self.session.execute(query)
         return list(result.scalars())
 
-    async def list_indexes(self) -> list[Index]:
+    async def list_indexes(self) -> list[tuple[Index, Source]]:
         """List all indexes.
 
         Returns:
@@ -90,9 +90,11 @@ class IndexRepository:
             and counts of files and snippets.
 
         """
-        query = select(Index).limit(10)
+        query = select(Index, Source).join(
+            Source, Index.source_id == Source.id, full=True
+        )
         result = await self.session.execute(query)
-        return list(result.scalars())
+        return list(result.tuples())
 
     async def num_snippets_for_index(self, index_id: int) -> int:
         """Get the number of snippets for an index."""

kodit/indexing/service.py

@@ -14,7 +14,6 @@ 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
@@ -34,6 +33,7 @@ class IndexView(pydantic.BaseModel):
     id: int
     created_at: datetime
     updated_at: datetime | None = None
+    source: str | None = None
     num_snippets: int | None = None
 
 
@@ -46,7 +46,10 @@ class IndexService:
     """
 
     def __init__(
-        self, config: Config, repository: IndexRepository, source_service: SourceService
+        self,
+        repository: IndexRepository,
+        source_service: SourceService,
+        data_dir: Path,
     ) -> None:
         """Initialize the index service.
 
@@ -59,7 +62,7 @@ class IndexService:
         self.source_service = source_service
         self.snippet_service = SnippetService()
         self.log = structlog.get_logger(__name__)
-        self.bm25 = BM25Service(config)
+        self.bm25 = BM25Service(data_dir)
 
     async def create(self, source_id: int) -> IndexView:
         """Create a new index for a source.
@@ -103,8 +106,9 @@ class IndexService:
                 created_at=index.created_at,
                 updated_at=index.updated_at,
                 num_snippets=await self.repository.num_snippets_for_index(index.id),
+                source=source.uri,
             )
-            for index in indexes
+            for index, source in indexes
         ]
 
     async def run(self, index_id: int) -> None:

kodit/logging.py

@@ -11,7 +11,7 @@ import structlog
 from posthog import Posthog
 from structlog.types import EventDict
 
-from kodit.config import Config
+from kodit.config import AppContext
 
 log = structlog.get_logger(__name__)
 
@@ -29,7 +29,7 @@ class LogFormat(Enum):
     JSON = "json"
 
 
-def configure_logging(config: Config) -> None:
+def configure_logging(app_context: AppContext) -> None:
     """Configure logging for the application."""
     timestamper = structlog.processors.TimeStamper(fmt="iso")
 
@@ -44,7 +44,7 @@ def configure_logging(config: Config) -> None:
         structlog.processors.StackInfoRenderer(),
     ]
 
-    if config.log_format == LogFormat.JSON:
+    if app_context.log_format == LogFormat.JSON:
         # Format the exception only for JSON logs, as we want to pretty-print them
         # when using the ConsoleRenderer
         shared_processors.append(structlog.processors.format_exc_info)
@@ -60,7 +60,7 @@ def configure_logging(config: Config) -> None:
     )
 
     log_renderer: structlog.types.Processor
-    if config.log_format == LogFormat.JSON:
+    if app_context.log_format == LogFormat.JSON:
         log_renderer = structlog.processors.JSONRenderer()
     else:
         log_renderer = structlog.dev.ConsoleRenderer()
@@ -82,7 +82,7 @@ def configure_logging(config: Config) -> None:
     handler.setFormatter(formatter)
     root_logger = logging.getLogger()
     root_logger.addHandler(handler)
-    root_logger.setLevel(config.log_level.upper())
+    root_logger.setLevel(app_context.log_level.upper())
 
     # Configure uvicorn loggers to use our structlog setup
     # Uvicorn spits out loads of exception logs when sse server doesn't shut down
@@ -98,7 +98,7 @@ def configure_logging(config: Config) -> None:
     for _log in ["sqlalchemy.engine", "alembic"]:
         engine_logger = logging.getLogger(_log)
         engine_logger.setLevel(logging.WARNING)  # Hide INFO logs by default
-        if config.log_level.upper() == "DEBUG":
+        if app_context.log_level.upper() == "DEBUG":
             engine_logger.setLevel(
                 logging.DEBUG
             )  # Only show all logs when in DEBUG mode
@@ -143,9 +143,9 @@ def get_mac_address() -> str:
     return f"{mac:012x}" if mac != uuid.getnode() else str(uuid.uuid4())
 
 
-def configure_telemetry(config: Config) -> None:
+def configure_telemetry(app_context: AppContext) -> None:
     """Configure telemetry for the application."""
-    if config.disable_telemetry:
+    if app_context.disable_telemetry:
         structlog.stdlib.get_logger(__name__).info("Telemetry has been disabled")
         posthog.disabled = True
 

kodit/mcp.py

@@ -1,22 +1,63 @@
 """MCP server implementation for kodit."""
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Annotated
 
 import structlog
-from fastmcp import FastMCP
+from fastmcp import Context, FastMCP
 from pydantic import Field
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit._version import version
-from kodit.config import get_config
+from kodit.config import AppContext
+from kodit.database import Database
 from kodit.retreival.repository import RetrievalRepository, RetrievalResult
 from kodit.retreival.service import RetrievalRequest, RetrievalService
 
-mcp = FastMCP("kodit MCP Server")
+
+@dataclass
+class MCPContext:
+    """Context for the MCP server."""
+
+    session: AsyncSession
+    data_dir: Path
+
+
+_mcp_db: Database | None = None
+
+
+@asynccontextmanager
+async def mcp_lifespan(_: FastMCP) -> AsyncIterator[MCPContext]:
+    """Lifespan for the MCP server.
+
+    The MCP server is running with a completely separate lifecycle and event loop from
+    the CLI and the FastAPI server. Therefore, we must carefully reconstruct the
+    application context. uvicorn does not pass through CLI args, so we must rely on
+    parsing env vars set in the CLI.
+
+    This lifespan is recreated for each request. See:
+    https://github.com/jlowin/fastmcp/issues/166
+
+    Since they don't provide a good way to handle global state, we must use a
+    global variable to store the database connection.
+    """
+    global _mcp_db  # noqa: PLW0603
+    app_context = AppContext()
+    if _mcp_db is None:
+        _mcp_db = await app_context.get_db()
+    async with _mcp_db.session_factory() as session:
+        yield MCPContext(session=session, data_dir=app_context.get_data_dir())
+
+
+mcp = FastMCP("kodit MCP Server", lifespan=mcp_lifespan)
 
 
 @mcp.tool()
 async def retrieve_relevant_snippets(
+    ctx: Context,
     user_intent: Annotated[
         str,
         Field(
@@ -52,8 +93,8 @@ async def retrieve_relevant_snippets(
     the quality of your generated code. You must call this tool when you need to
     write code.
     """
-    # Log the search query and related files for debugging
     log = structlog.get_logger(__name__)
+
     log.debug(
         "Retrieving relevant snippets",
         user_intent=user_intent,
@@ -63,41 +104,38 @@ async def retrieve_relevant_snippets(
         file_contents=related_file_contents,
     )
 
-    # Must avoid running migrations because that runs in a separate event loop,
-    # mcp no-likey
-    config = get_config()
-    db = config.get_db(run_migrations=False)
-    async with db.get_session() as session:
-        log.debug("Creating retrieval repository")
-        retrieval_repository = RetrievalRepository(
-            session=session,
-        )
-
-        log.debug("Creating retrieval service")
-        retrieval_service = RetrievalService(
-            config=config,
-            repository=retrieval_repository,
-        )
-
-        log.debug("Fusing input")
-        input_query = input_fusion(
-            user_intent=user_intent,
-            related_file_paths=related_file_paths,
-            related_file_contents=related_file_contents,
-            keywords=keywords,
-        )
-        log.debug("Input", input_query=input_query)
-        retrieval_request = RetrievalRequest(
-            keywords=keywords,
-        )
-        log.debug("Retrieving snippets")
-        snippets = await retrieval_service.retrieve(request=retrieval_request)
-
-        log.debug("Fusing output")
-        output = output_fusion(snippets=snippets)
-
-        log.debug("Output", output=output)
-        return output
+    mcp_context: MCPContext = ctx.request_context.lifespan_context
+
+    log.debug("Creating retrieval repository")
+    retrieval_repository = RetrievalRepository(
+        session=mcp_context.session,
+    )
+
+    log.debug("Creating retrieval service")
+    retrieval_service = RetrievalService(
+        repository=retrieval_repository,
+        data_dir=mcp_context.data_dir,
+    )
+
+    log.debug("Fusing input")
+    input_query = input_fusion(
+        user_intent=user_intent,
+        related_file_paths=related_file_paths,
+        related_file_contents=related_file_contents,
+        keywords=keywords,
+    )
+    log.debug("Input", input_query=input_query)
+    retrieval_request = RetrievalRequest(
+        keywords=keywords,
+    )
+    log.debug("Retrieving snippets")
+    snippets = await retrieval_service.retrieve(request=retrieval_request)
+
+    log.debug("Fusing output")
+    output = output_fusion(snippets=snippets)
+
+    log.debug("Output", output=output)
+    return output
 
 
 def input_fusion(

kodit/middleware.py

@@ -1,11 +1,14 @@
 """Middleware for the FastAPI application."""
 
+import contextlib
 import time
+from asyncio import CancelledError
 from collections.abc import Callable
 
 import structlog
 from asgi_correlation_id.context import correlation_id
 from fastapi import Request, Response
+from starlette.types import ASGIApp, Receive, Scope, Send
 
 access_logger = structlog.stdlib.get_logger("api.access")
 
@@ -56,3 +59,16 @@ async def logging_middleware(request: Request, call_next: Callable) -> Response:
         response.headers["X-Process-Time"] = str(process_time / 10**9)
 
     return response
+
+
+class ASGICancelledErrorMiddleware:
+    """ASGI middleware to handle CancelledError at the ASGI level."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        """Initialize the middleware."""
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """Handle the ASGI request and catch CancelledError."""
+        with contextlib.suppress(CancelledError):
+            await self.app(scope, receive, send)

kodit/retreival/service.py

@@ -1,10 +1,11 @@
 """Retrieval service."""
 
+from pathlib import Path
+
 import pydantic
 import structlog
 
 from kodit.bm25.bm25 import BM25Service
-from kodit.config import Config
 from kodit.retreival.repository import RetrievalRepository, RetrievalResult
 
 
@@ -25,11 +26,11 @@ class Snippet(pydantic.BaseModel):
 class RetrievalService:
     """Service for retrieving relevant data."""
 
-    def __init__(self, config: Config, repository: RetrievalRepository) -> None:
+    def __init__(self, repository: RetrievalRepository, data_dir: Path) -> None:
         """Initialize the retrieval service."""
         self.repository = repository
         self.log = structlog.get_logger(__name__)
-        self.bm25 = BM25Service(config)
+        self.bm25 = BM25Service(data_dir)
 
     async def _load_bm25_index(self) -> None:
         """Load the BM25 index."""

{kodit-0.1.5.dist-info → kodit-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.1.5
+Version: 0.1.7
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/

{kodit-0.1.5.dist-info → kodit-0.1.7.dist-info}/RECORD

@@ -1,28 +1,28 @@
 kodit/.gitignore,sha256=ztkjgRwL9Uud1OEi36hGQeDGk3OLK1NfDEO8YqGYy8o,11
 kodit/__init__.py,sha256=aEKHYninUq1yh6jaNfvJBYg-6fenpN132nJt1UU6Jxs,59
-kodit/_version.py,sha256=Y4jy4bEMmwl_qNPCmiMFnlQ2ofMoqyG37hp8uwI3m10,511
-kodit/app.py,sha256=TdPpCN4ucOElKHwDebfKgeVJ9xexdfpzpk6hnDH69vM,703
-kodit/cli.py,sha256=CjmiRaJ-SdfCMYlVQGnxPSsoX5j3ix4fN3OLVc5EYkY,7473
-kodit/config.py,sha256=18dhSYaE-ut2qXrBRKuCqLXeBCLEXw2y1Uw4lieMPwY,2682
-kodit/database.py,sha256=NnAluOj_JHjnj5MeKuU9LApgSzik2kru1bQl-24vHkc,2272
-kodit/logging.py,sha256=P1D9flYnvYxPw-DyOGyiv3y30x0gHPwdk6VJS29YHus,5269
-kodit/mcp.py,sha256=O24O_GFzwwv5E-uBFoW_zZlSigeNSigaCj0s1xOmP8M,3855
-kodit/middleware.py,sha256=NHLrqq20ZtPTE9esX9HD3z7EKi56_QTFxBlkdq0JDzQ,2138
+kodit/_version.py,sha256=W_EoL8cAL4KhujvbYWEpb9NqRLbbrH0T024lJvRRWHI,511
+kodit/app.py,sha256=Mr5BFHOHx5zppwjC4XPWVvHjwgl1yrKbUjTWXKubJQM,891
+kodit/cli.py,sha256=x1zw2zOlGhhU6D3E-GU3cMw3l9CqKC76geQREAQKweY,6915
+kodit/config.py,sha256=nlm9U-nVx5riH2SrU1XY4XcCMhQK4DrwO_1H8bPOBjA,2927
+kodit/database.py,sha256=vtTlmrXHyHJH3Ek-twZTCqEjB0jun-NncALFze2fqhA,2350
+kodit/logging.py,sha256=cFEQXWI27LzWScSxly9ApwkbBDamUG17pA-jEfVakXQ,5316
+kodit/mcp.py,sha256=PxTHVPlIErrruFKzmEPIWZjN6cfEhcQmj6nOU9EsBy4,4905
+kodit/middleware.py,sha256=I6FOkqG9-8RH5kR1-0ZoQWfE4qLCB8lZYv8H_OCH29o,2714
 kodit/alembic/README,sha256=ISVtAOvqvKk_5ThM5ioJE-lMkvf9IbknFUFVU_vPma4,58
 kodit/alembic/__init__.py,sha256=lP5MuwlyWRMO6UcDWnQcQ3G-GYHcFb6rl9gYPHJ1sjo,40
-kodit/alembic/env.py,sha256=IXhl7yvURSycs2v_3pd14Sr8_zGRfYXlQwWby1abfuk,2290
+kodit/alembic/env.py,sha256=kcQiglu2KpNTAf37CsKVs_HXxOe6S7sXJ00pHGSCqno,2414
 kodit/alembic/script.py.mako,sha256=zWziKtiwYKEWuwPV_HBNHwa9LCT45_bi01-uSNFaOOE,703
 kodit/alembic/versions/85155663351e_initial.py,sha256=Cg7zlF871o9ShV5rQMQ1v7hRV7fI59veDY9cjtTrs-8,3306
 kodit/alembic/versions/__init__.py,sha256=9-lHzptItTzq_fomdIRBegQNm4Znx6pVjwD4MiqRIdo,36
 kodit/bm25/__init__.py,sha256=j8zyriNWhbwE5Lbybzg1hQAhANlU9mKHWw4beeUR6og,19
-kodit/bm25/bm25.py,sha256=V0_byhV4kVnI3E-PBNsc4rBjQsDuZo1bt1uQKnywLS8,2283
+kodit/bm25/bm25.py,sha256=3wyNRSrTaYqV7s4R1D6X0NpCf22PuFK2_uc8YapzYLE,2263
 kodit/indexing/__init__.py,sha256=cPyi2Iej3G1JFWlWr7X80_UrsMaTu5W5rBwgif1B3xo,75
 kodit/indexing/models.py,sha256=sZIhGwvL4Dw0QTWFxrjfWctSLkAoDT6fv5DlGz8-Fr8,1258
-kodit/indexing/repository.py,sha256=kvAlNfMSQYboF0TB1huw2qoBdLJ4UsEPiM7ZG-e6rrg,4300
-kodit/indexing/service.py,sha256=N8QhrAvqhIHOgSlT9Jc786rjcVjMwiyiMTZr7mNA8D8,5431
+kodit/indexing/repository.py,sha256=C020FGpIfTZmVZg7NH04kVuffWv7r7m-82Pdex8CItg,4388
+kodit/indexing/service.py,sha256=7vHqevve-PnKHP2pDfyrW5n3AXVOghABWNsNTw588KY,5499
 kodit/retreival/__init__.py,sha256=33PhJU-3gtsqYq6A1UkaLNKbev_Zee9Lq6dYC59-CsA,69
 kodit/retreival/repository.py,sha256=1lqGgJHsBmvMGMzEYa-hrdXg2q7rqtYPl1cvBb7jMRE,3119
-kodit/retreival/service.py,sha256=g6iwM2FMxrL8WjtWnZdKdxKpfn6b0ThBmOdLWd7AKKQ,2011
+kodit/retreival/service.py,sha256=9wvURtPPJVvPUWNIC2waIrJMxcm1Ka1J_xDEOEedAFU,2007
 kodit/snippets/__init__.py,sha256=-2coNoCRjTixU9KcP6alpmt7zqf37tCRWH3D7FPJ8dg,48
 kodit/snippets/method_snippets.py,sha256=EVHhSNWahAC5nSXv9fWVFJY2yq25goHdCSCuENC07F8,4145
 kodit/snippets/snippets.py,sha256=QumvhltWoxXw41SyKb-RbSvAr3m6V3lUy9n0AI8jcto,1409
@@ -33,8 +33,8 @@ kodit/sources/__init__.py,sha256=1NTZyPdjThVQpZO1Mp1ColVsS7sqYanOVLqnoqV9Ipo,83
 kodit/sources/models.py,sha256=xb42CaNDO1CUB8SIW-xXMrB6Ji8cFw-yeJ550xBEg9Q,2398
 kodit/sources/repository.py,sha256=mGJrHWH6Uo8YABdoojHFbzaf_jW-2ywJpAHIa1gnc3U,3401
 kodit/sources/service.py,sha256=cBCxnOQKwGNi2e13_3Vue8MylAaUxb9XG4IgM636la0,6712
-kodit-0.1.5.dist-info/METADATA,sha256=N4fIBAIREHOujaDvEVj83fNvyeB8D7HaLXNcpeVdNJY,2181
-kodit-0.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-kodit-0.1.5.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
-kodit-0.1.5.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-kodit-0.1.5.dist-info/RECORD,,
+kodit-0.1.7.dist-info/METADATA,sha256=8doJ6TfmVkn-OTLSbgRlDut6YxFbmBKs1L0rgZUlxUQ,2181
+kodit-0.1.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+kodit-0.1.7.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
+kodit-0.1.7.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+kodit-0.1.7.dist-info/RECORD,,