Spaces:
Sleeping
Sleeping
| """FastMCP server exposing secure code search via Scalekit auth.""" | |
| from __future__ import annotations | |
| import logging | |
| import os | |
| from html import escape | |
| from typing import Any, Protocol | |
| from fastmcp import FastMCP | |
| from fastmcp.server.auth import TokenVerifier | |
| from fastmcp.server.auth.providers.scalekit import ScalekitProvider | |
| from Retriever import Retriever | |
| class _RetrieverProtocol(Protocol): | |
| def retrieve(self, query: str, *, top_k: int = 10, | |
| repo: str | None = None, branch: str | None = None) -> list: ... | |
| logger = logging.getLogger(__name__) | |
| def _env_truthy(value: str | None) -> bool: | |
| return str(value or "").strip().lower() in {"1", "true", "yes", "on"} | |
| def _resolve_require_auth(require_auth: bool | None) -> bool: | |
| if require_auth is not None: | |
| return require_auth | |
| if _env_truthy(os.environ.get("GITRAG_MCP_DISABLE_AUTH")): | |
| return False | |
| configured = os.environ.get("GITRAG_MCP_REQUIRE_AUTH") | |
| if configured is not None: | |
| return _env_truthy(configured) | |
| return True | |
| def build_scalekit_provider(*, token_verifier: TokenVerifier | None = None, | |
| base_url: str | None = None, ) -> ScalekitProvider: | |
| """Construct Scalekit auth provider from environment variables.""" | |
| environment_url = os.environ.get("SCALEKIT_ENVIRONMENT_URL") | |
| client_id = os.environ.get("SCALEKIT_CLIENT_ID") | |
| resource_id = os.environ.get("SCALEKIT_RESOURCE_ID") | |
| missing = [ | |
| name | |
| for name, value in { | |
| "SCALEKIT_ENVIRONMENT_URL": environment_url, | |
| "SCALEKIT_CLIENT_ID": client_id, | |
| "SCALEKIT_RESOURCE_ID": resource_id, | |
| }.items() | |
| if not value | |
| ] | |
| if missing: | |
| raise RuntimeError(f"Missing required Scalekit environment variables: {', '.join(missing)}") | |
| resolved_base_url = base_url or os.environ.get("MCP_BASE_URL") or "http://127.0.0.1:8000/mcp" | |
| return ScalekitProvider( | |
| environment_url=environment_url, | |
| client_id=client_id, | |
| resource_id=resource_id, | |
| base_url=resolved_base_url, | |
| token_verifier=token_verifier, | |
| ) | |
| def create_mcp_server(*, retriever: Retriever, token_verifier: TokenVerifier | None = None, | |
| base_url: str | None = None, require_auth: bool | None = None, ) -> FastMCP: | |
| """Create an authenticated MCP server with `search_code` tool.""" | |
| resolved_require_auth = _resolve_require_auth(require_auth) | |
| auth_provider = None | |
| if resolved_require_auth: | |
| auth_provider = build_scalekit_provider(token_verifier=token_verifier, base_url=base_url) | |
| else: | |
| logger.info("Authentication disabled for MCP server") | |
| mcp = FastMCP(name="GitRag MCP Server", auth=auth_provider) | |
| def search_code( | |
| query: str, | |
| top_k: int = 5, | |
| repo: str | None = None, | |
| branch: str | None = None, | |
| ) -> dict[str, Any]: | |
| """Search indexed code and return top snippets.""" | |
| logger.info("search_code start query=%r top_k=%s repo=%r branch=%r", query, top_k, repo, branch) | |
| chunks = retriever.retrieve(query, top_k=top_k, repo=repo, branch=branch) | |
| output: list[dict[str, Any]] = [] | |
| markdown_blocks: list[str] = [] | |
| for chunk in chunks: | |
| payload = vars(chunk).copy() | |
| if isinstance(payload.get("embeddings"), (bytes, bytearray, memoryview)): | |
| payload["embeddings"] = None | |
| escaped_path = escape(str(payload.get("path", "")), quote=True) | |
| escaped_repo = escape(str(payload.get("repo", "")), quote=True) | |
| escaped_branch = escape(str(payload.get("branch", "")), quote=True) | |
| escaped_chunk = escape(str(payload.get("chunk", ""))) | |
| payload["formatted"] = ( | |
| f'<file path="{escaped_path}" repo="{escaped_repo}" ' | |
| f'branch="{escaped_branch}">\n{escaped_chunk}\n</file>' | |
| ) | |
| markdown_blocks.append( | |
| f"### {payload.get('path')}\n\n```{payload.get('language', '')}\n{payload.get('chunk', '')}\n```" | |
| ) | |
| output.append(payload) | |
| response = {"results": output, "markdown": "\n\n".join(markdown_blocks)} | |
| logger.info("search_code end query=%r results=%s", query, len(output)) | |
| return response | |
| return mcp | |