kodit 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

kodit/logging.py

@@ -11,6 +11,8 @@ import structlog
 from posthog import Posthog
 from structlog.types import EventDict
 
+from kodit.config import Config
+
 log = structlog.get_logger(__name__)
 
 
@@ -27,14 +29,8 @@ class LogFormat(Enum):
     JSON = "json"
 
 
-def configure_logging(log_level: str, log_format: LogFormat) -> None:
-    """Configure logging for the application.
-
-    Args:
-        json_logs: Whether to use JSON format for logs
-        log_level: The minimum log level to display
-
-    """
+def configure_logging(config: Config) -> None:
+    """Configure logging for the application."""
     timestamper = structlog.processors.TimeStamper(fmt="iso")
 
     shared_processors: list[structlog.types.Processor] = [
@@ -48,7 +44,7 @@ def configure_logging(log_level: str, log_format: LogFormat) -> None:
         structlog.processors.StackInfoRenderer(),
     ]
 
-    if log_format == LogFormat.JSON:
+    if config.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)
@@ -64,7 +60,7 @@ def configure_logging(log_level: str, log_format: LogFormat) -> None:
     )
 
     log_renderer: structlog.types.Processor
-    if log_format == LogFormat.JSON:
+    if config.log_format == LogFormat.JSON:
         log_renderer = structlog.processors.JSONRenderer()
     else:
         log_renderer = structlog.dev.ConsoleRenderer()
@@ -86,18 +82,23 @@ def configure_logging(log_level: str, log_format: LogFormat) -> None:
     handler.setFormatter(formatter)
     root_logger = logging.getLogger()
     root_logger.addHandler(handler)
-    root_logger.setLevel(log_level.upper())
+    root_logger.setLevel(config.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
+    # gracefully, so we hide them unless in DEBUG mode
     for _log in ["uvicorn", "uvicorn.error", "uvicorn.access"]:
-        logging.getLogger(_log).handlers.clear()
-        logging.getLogger(_log).propagate = True
+        if root_logger.getEffectiveLevel() == logging.DEBUG:
+            logging.getLogger(_log).handlers.clear()
+            logging.getLogger(_log).propagate = True
+        else:
+            logging.getLogger(_log).disabled = True
 
     # Configure SQLAlchemy loggers to use our structlog setup
     for _log in ["sqlalchemy.engine", "alembic"]:
         engine_logger = logging.getLogger(_log)
         engine_logger.setLevel(logging.WARNING)  # Hide INFO logs by default
-        if log_level.upper() == "DEBUG":
+        if config.log_level.upper() == "DEBUG":
             engine_logger.setLevel(
                 logging.DEBUG
             )  # Only show all logs when in DEBUG mode
@@ -142,10 +143,11 @@ def get_mac_address() -> str:
     return f"{mac:012x}" if mac != uuid.getnode() else str(uuid.uuid4())
 
 
-def disable_posthog() -> None:
-    """Disable telemetry for the application."""
-    structlog.stdlib.get_logger(__name__).info("Telemetry has been disabled")
-    posthog.disabled = True
+def configure_telemetry(config: Config) -> None:
+    """Configure telemetry for the application."""
+    if config.disable_telemetry:
+        structlog.stdlib.get_logger(__name__).info("Telemetry has been disabled")
+        posthog.disabled = True
 
 
 def log_event(event: str, properties: dict[str, Any] | None = None) -> None:

kodit/mcp.py

@@ -4,10 +4,11 @@ from pathlib import Path
 from typing import Annotated
 
 import structlog
-from mcp.server.fastmcp import FastMCP
+from fastmcp import FastMCP
 from pydantic import Field
 
-from kodit.database import get_session
+from kodit._version import version
+from kodit.config import get_config
 from kodit.retreival.repository import RetrievalRepository, RetrievalResult
 from kodit.retreival.service import RetrievalRequest, RetrievalService
 
@@ -62,7 +63,11 @@ async def retrieve_relevant_snippets(
         file_contents=related_file_contents,
     )
 
-    async with get_session() as session:
+    # 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,
@@ -70,6 +75,7 @@ async def retrieve_relevant_snippets(
 
         log.debug("Creating retrieval service")
         retrieval_service = RetrievalService(
+            config=config,
             repository=retrieval_repository,
         )
 
@@ -82,7 +88,7 @@ async def retrieve_relevant_snippets(
         )
         log.debug("Input", input_query=input_query)
         retrieval_request = RetrievalRequest(
-            query=input_query,
+            keywords=keywords,
         )
         log.debug("Retrieving snippets")
         snippets = await retrieval_service.retrieve(request=retrieval_request)
@@ -108,3 +114,9 @@ def input_fusion(
 def output_fusion(snippets: list[RetrievalResult]) -> str:
     """Fuse the snippets into a single output."""
     return "\n\n".join(f"{snippet.uri}\n{snippet.content}" for snippet in snippets)
+
+
+@mcp.tool()
+async def get_version() -> str:
+    """Get the version of the kodit project."""
+    return version

kodit/retreival/repository.py

@@ -74,3 +74,35 @@ class RetrievalRepository:
             )
             for snippet, file in results
         ]
+
+    async def list_snippet_ids(self) -> list[int]:
+        """List all snippet IDs.
+
+        Returns:
+            A list of all snippets.
+
+        """
+        query = select(Snippet.id)
+        rows = await self.session.execute(query)
+        return list(rows.scalars().all())
+
+    async def list_snippets_by_ids(self, ids: list[int]) -> list[RetrievalResult]:
+        """List snippets by IDs.
+
+        Returns:
+            A list of snippets.
+
+        """
+        query = (
+            select(Snippet, File)
+            .where(Snippet.id.in_(ids))
+            .join(File, Snippet.file_id == File.id)
+        )
+        rows = await self.session.execute(query)
+        return [
+            RetrievalResult(
+                uri=file.uri,
+                content=snippet.content,
+            )
+            for snippet, file in rows.all()
+        ]

kodit/retreival/service.py

@@ -1,14 +1,18 @@
 """Retrieval service."""
 
 import pydantic
+import structlog
 
+from kodit.bm25.bm25 import BM25Service
+from kodit.config import Config
 from kodit.retreival.repository import RetrievalRepository, RetrievalResult
 
 
 class RetrievalRequest(pydantic.BaseModel):
     """Request for a retrieval."""
 
-    query: str
+    keywords: list[str]
+    top_k: int = 10
 
 
 class Snippet(pydantic.BaseModel):
@@ -21,10 +25,44 @@ class Snippet(pydantic.BaseModel):
 class RetrievalService:
     """Service for retrieving relevant data."""
 
-    def __init__(self, repository: RetrievalRepository) -> None:
+    def __init__(self, config: Config, repository: RetrievalRepository) -> None:
         """Initialize the retrieval service."""
         self.repository = repository
+        self.log = structlog.get_logger(__name__)
+        self.bm25 = BM25Service(config)
+
+    async def _load_bm25_index(self) -> None:
+        """Load the BM25 index."""
 
     async def retrieve(self, request: RetrievalRequest) -> list[RetrievalResult]:
         """Retrieve relevant data."""
-        return await self.repository.string_search(request.query)
+        snippet_ids = await self.repository.list_snippet_ids()
+
+        # Gather results for each keyword
+        result_ids: list[tuple[int, float]] = []
+        for keyword in request.keywords:
+            results = self.bm25.retrieve(snippet_ids, keyword, request.top_k)
+            result_ids.extend(results)
+
+        if len(result_ids) == 0:
+            return []
+
+        # Sort results by score
+        result_ids.sort(key=lambda x: x[1], reverse=True)
+
+        self.log.debug(
+            "Retrieval results",
+            total_results=len(result_ids),
+            max_score=result_ids[0][1],
+            min_score=result_ids[-1][1],
+            median_score=result_ids[len(result_ids) // 2][1],
+        )
+
+        # Don't return zero score results
+        result_ids = [x for x in result_ids if x[1] > 0]
+
+        # Build final list of doc ids up to top_k
+        final_doc_ids = [x[0] for x in result_ids[: request.top_k]]
+
+        # Get snippets from database
+        return await self.repository.list_snippets_by_ids(final_doc_ids)

kodit/snippets/__init__.py

@@ -0,0 +1 @@
+"""Extract method snippets from source code."""

kodit/snippets/languages/__init__.py

@@ -0,0 +1,53 @@
+"""Detect the language of a file."""
+
+from pathlib import Path
+from typing import cast
+
+from tree_sitter_language_pack import SupportedLanguage
+
+# Mapping of file extensions to programming languages
+LANGUAGE_MAP: dict[str, str] = {
+    # JavaScript/TypeScript
+    "js": "javascript",
+    "jsx": "javascript",
+    "ts": "typescript",
+    "tsx": "typescript",
+    # Python
+    "py": "python",
+    # Rust
+    "rs": "rust",
+    # Go
+    "go": "go",
+    # C/C++
+    "cpp": "cpp",
+    "hpp": "cpp",
+    "c": "c",
+    "h": "c",
+    # C#
+    "cs": "csharp",
+    # Ruby
+    "rb": "ruby",
+    # Java
+    "java": "java",
+    # PHP
+    "php": "php",
+    # Swift
+    "swift": "swift",
+    # Kotlin
+    "kt": "kotlin",
+}
+
+
+def detect_language(file_path: Path) -> SupportedLanguage:
+    """Detect the language of a file."""
+    suffix = file_path.suffix.removeprefix(".").lower()
+    msg = f"Unsupported language for file suffix: {suffix}"
+    lang = LANGUAGE_MAP.get(suffix)
+    if lang is None:
+        raise ValueError(msg)
+
+    # Try to cast the language to a SupportedLanguage
+    try:
+        return cast("SupportedLanguage", lang)
+    except Exception as e:
+        raise ValueError(msg) from e

kodit/snippets/languages/csharp.scm

@@ -0,0 +1,12 @@
+(method_declaration
+  name: (identifier) @function.name
+  body: (block) @function.body
+) @function.def
+
+(class_declaration
+  name: (identifier) @class.name
+) @class.def
+
+(using_directive) @import.name
+
+(identifier) @ident

kodit/snippets/languages/python.scm

@@ -0,0 +1,22 @@
+(function_definition
+  name: (identifier) @function.name
+  body: (block) @function.body
+) @function.def
+
+(class_definition
+  name: (identifier) @class.name
+) @class.def
+
+(import_statement
+  name: (dotted_name (identifier) @import.name))
+
+(import_from_statement
+  module_name: (dotted_name (identifier) @import.from))
+
+(identifier) @ident
+
+(assignment
+  left: (identifier) @assignment.lhs)
+
+(parameters
+  (identifier) @param.name)

kodit/snippets/method_snippets.py

@@ -0,0 +1,120 @@
+"""Extract method snippets from source code."""
+
+from tree_sitter import Node, Query
+from tree_sitter_language_pack import SupportedLanguage, get_language, get_parser
+
+
+class MethodSnippets:
+    """Extract method snippets from source code."""
+
+    def __init__(self, language: SupportedLanguage, query: str) -> None:
+        """Initialize the MethodSnippets class."""
+        self.language = get_language(language)
+        self.parser = get_parser(language)
+        self.query = Query(self.language, query)
+
+    def _get_leaf_functions(
+        self, captures_by_name: dict[str, list[Node]]
+    ) -> list[Node]:
+        """Return all leaf functions in the AST."""
+        return [
+            node
+            for node in captures_by_name.get("function.body", [])
+            if self._is_leaf_function(captures_by_name, node)
+        ]
+
+    def _is_leaf_function(
+        self, captures_by_name: dict[str, list[Node]], node: Node
+    ) -> bool:
+        """Return True if the node is a leaf function."""
+        for other in captures_by_name.get("function.body", []):
+            if other == node:  # Skip self
+                continue
+            # if other is inside node, it's not a leaf function
+            if other.start_byte >= node.start_byte and other.end_byte <= node.end_byte:
+                return False
+        return True
+
+    def _get_imports(self, captures_by_name: dict[str, list[Node]]) -> list[Node]:
+        """Return all imports in the AST."""
+        return captures_by_name.get("import.name", []) + captures_by_name.get(
+            "import.from", []
+        )
+
+    def _classes_and_functions(
+        self, captures_by_name: dict[str, list[Node]]
+    ) -> list[int]:
+        """Return all classes and functions in the AST."""
+        return [
+            node.id
+            for node in {
+                *captures_by_name.get("function.def", []),
+                *captures_by_name.get("class.def", []),
+            }
+        ]
+
+    def _get_ancestors(
+        self, captures_by_name: dict[str, list[Node]], node: Node
+    ) -> list[Node]:
+        """Return all ancestors of the node."""
+        valid_ancestors = self._classes_and_functions(captures_by_name)
+        ancestors = []
+        parent = node.parent
+        while parent:
+            if parent.id in valid_ancestors:
+                ancestors.append(parent)
+            parent = parent.parent
+        return ancestors
+
+    def extract(self, source_code: bytes) -> list[str]:
+        """Extract method snippets from source code."""
+        tree = self.parser.parse(source_code)
+
+        captures_by_name = self.query.captures(tree.root_node)
+
+        lines = source_code.decode().splitlines()
+
+        # Find all leaf functions
+        leaf_functions = self._get_leaf_functions(captures_by_name)
+
+        # Find all imports
+        imports = self._get_imports(captures_by_name)
+
+        results = []
+
+        # For each leaf function, find all lines this function is dependent on
+        for func_node in leaf_functions:
+            all_lines_to_keep = set()
+
+            ancestors = self._get_ancestors(captures_by_name, func_node)
+
+            # Add self to keep
+            all_lines_to_keep.update(
+                range(func_node.start_point[0], func_node.end_point[0] + 1)
+            )
+
+            # Add imports to keep
+            for import_node in imports:
+                all_lines_to_keep.update(
+                    range(import_node.start_point[0], import_node.end_point[0] + 1)
+                )
+
+            # Add ancestors to keep
+            for node in ancestors:
+                # Get the first line of the node for now
+                start = node.start_point[0]
+                end = node.start_point[0]
+                all_lines_to_keep.update(range(start, end + 1))
+
+            pseudo_code = []
+            for i, line in enumerate(lines):
+                if i in all_lines_to_keep:
+                    pseudo_code.append(line)
+
+            results.append("\n".join(pseudo_code))
+
+        # If there are no results, then return the entire file
+        if not results:
+            return [source_code.decode()]
+
+        return results

kodit/snippets/snippets.py

@@ -0,0 +1,48 @@
+"""Generate snippets from a file."""
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from kodit.snippets.languages import detect_language
+from kodit.snippets.method_snippets import MethodSnippets
+
+
+@dataclass
+class Snippet:
+    """A snippet of code."""
+
+    text: str
+
+
+class SnippetService:
+    """Factory for generating snippets from a file.
+
+    This is required because there's going to be multiple ways to generate snippets.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the snippet factory."""
+        self.language_dir = Path(__file__).parent / "languages"
+
+    def snippets_for_file(self, file_path: Path) -> list[Snippet]:
+        """Generate snippets from a file."""
+        language = detect_language(file_path)
+
+        try:
+            query_path = self.language_dir / f"{language}.scm"
+            with query_path.open() as f:
+                query = f.read()
+        except Exception as e:
+            msg = f"Unsupported language: {file_path}"
+            raise ValueError(msg) from e
+
+        method_analser = MethodSnippets(language, query)
+
+        try:
+            file_bytes = file_path.read_bytes()
+        except Exception as e:
+            msg = f"Failed to read file: {file_path}"
+            raise ValueError(msg) from e
+
+        method_snippets = method_analser.extract(file_bytes)
+        return [Snippet(text=snippet) for snippet in method_snippets]

kodit/sources/service.py

@@ -18,12 +18,9 @@ import structlog
 from tqdm import tqdm
 from uritools import isuri, urisplit
 
-from kodit.config import DATA_DIR
 from kodit.sources.models import File, Source
 from kodit.sources.repository import SourceRepository
 
-CLONE_DIR = DATA_DIR / "clones"
-
 
 class SourceView(pydantic.BaseModel):
     """View model for displaying source information.
@@ -53,13 +50,14 @@ class SourceService:
     SourceRepository), and provides a clean API for source management.
     """
 
-    def __init__(self, repository: SourceRepository) -> None:
+    def __init__(self, clone_dir: Path, repository: SourceRepository) -> None:
         """Initialize the source service.
 
         Args:
             repository: The repository instance to use for database operations.
 
         """
+        self.clone_dir = clone_dir
         self.repository = repository
         self.log = structlog.get_logger(__name__)
 
@@ -129,7 +127,7 @@ class SourceService:
             raise ValueError(msg)
 
         # Clone into a local directory
-        clone_path = CLONE_DIR / directory.as_posix().replace("/", "_")
+        clone_path = self.clone_dir / directory.as_posix().replace("/", "_")
         clone_path.mkdir(parents=True, exist_ok=True)
 
         # Copy all files recursively, preserving directory structure, ignoring hidden

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

@@ -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.5.dist-info/RECORD

@@ -0,0 +1,40 @@
+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/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/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/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/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/snippets/__init__.py,sha256=-2coNoCRjTixU9KcP6alpmt7zqf37tCRWH3D7FPJ8dg,48
+kodit/snippets/method_snippets.py,sha256=EVHhSNWahAC5nSXv9fWVFJY2yq25goHdCSCuENC07F8,4145
+kodit/snippets/snippets.py,sha256=QumvhltWoxXw41SyKb-RbSvAr3m6V3lUy9n0AI8jcto,1409
+kodit/snippets/languages/__init__.py,sha256=Bj5KKZSls2MQ8ZY1S_nHg447MgGZW-2WZM-oq6vjwwA,1187
+kodit/snippets/languages/csharp.scm,sha256=gbBN4RiV1FBuTJF6orSnDFi8H9JwTw-d4piLJYsWUsc,222
+kodit/snippets/languages/python.scm,sha256=ee85R9PBzwye3IMTE7-iVoKWd_ViU3EJISTyrFGrVeo,429
+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/sse.py

@@ -1,61 +0,0 @@
-"""Server-Sent Events (SSE) implementation for kodit."""
-
-from collections.abc import Coroutine
-from typing import Any
-
-from fastapi import Request
-from mcp.server.fastmcp import FastMCP
-from mcp.server.session import ServerSession
-from mcp.server.sse import SseServerTransport
-from starlette.applications import Starlette
-from starlette.routing import Mount, Route
-
-####################################################################################
-# Temporary monkeypatch which avoids crashing when a POST message is received
-# before a connection has been initialized, e.g: after a deployment.
-old__received_request = ServerSession._received_request  # noqa: SLF001
-
-
-async def _received_request(self: ServerSession, *args: Any, **kwargs: Any) -> None:
-    """Handle a received request, catching RuntimeError to avoid crashes.
-
-    This is a temporary monkeypatch to avoid crashing when a POST message is
-    received before a connection has been initialized, e.g: after a deployment.
-    """
-    try:
-        return await old__received_request(self, *args, **kwargs)
-    except RuntimeError:
-        pass
-
-
-# pylint: disable-next=protected-access
-ServerSession._received_request = _received_request  # noqa: SLF001
-####################################################################################
-
-
-def create_sse_server(mcp: FastMCP) -> Starlette:
-    """Create a Starlette app that handles SSE connections and message handling."""
-    transport = SseServerTransport("/messages/")
-
-    # Define handler functions
-    async def handle_sse(request: Request) -> Coroutine[Any, Any, None]:
-        """Handle SSE connections."""
-        async with transport.connect_sse(
-            request.scope,
-            request.receive,
-            request._send,  # noqa: SLF001
-        ) as streams:
-            await mcp._mcp_server.run(  # noqa: SLF001
-                streams[0],
-                streams[1],
-                mcp._mcp_server.create_initialization_options(),  # noqa: SLF001
-            )
-
-    # Create Starlette routes for SSE and message handling
-    routes = [
-        Route("/sse/", endpoint=handle_sse),
-        Mount("/messages/", app=transport.handle_post_message),
-    ]
-
-    # Create a Starlette app
-    return Starlette(routes=routes)