deepresearch-flow 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

deepresearch_flow/paper/snapshot/mcp_server.py

@@ -0,0 +1,686 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import json
+import os
+from pathlib import Path
+import re
+from typing import Any
+
+import httpx
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.routing import Mount
+
+from fastmcp import FastMCP
+
+from deepresearch_flow.paper.snapshot.common import ApiLimits, _open_ro_conn
+from deepresearch_flow.paper.snapshot.text import merge_adjacent_markers, remove_cjk_spaces, rewrite_search_query
+
+_SUPPORTED_PROTOCOL_VERSIONS = {"2025-03-26", "2025-06-18"}
+_DEFAULT_MAX_CHARS = 50_000
+_DEFAULT_TIMEOUT = 10.0
+_PAPER_ID_PATTERN = re.compile(r'^[a-zA-Z0-9_-]+$')
+
+
+class McpToolError(Exception):
+    """MCP tool exception for standardized error handling.
+    
+    FastMCP will catch this exception and convert it to a proper
+    JSON-RPC error response that the client can understand.
+    """
+    
+    def __init__(self, code: str, message: str, **details):
+        self.code = code
+        self.message = message
+        self.details = details
+        super().__init__(message)
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to error dictionary format."""
+        return {"error": self.code, "message": self.message, **self.details}
+
+
+@dataclass(frozen=True)
+class McpSnapshotConfig:
+    snapshot_db: Path
+    static_base_url: str
+    static_export_dir: Path | None
+    limits: ApiLimits
+    origin_allowlist: list[str]
+    max_chars_default: int = _DEFAULT_MAX_CHARS
+    http_timeout: float = _DEFAULT_TIMEOUT
+    max_paper_id_length: int = 64
+    # HTTP client stored in object __dict__ to avoid dataclass frozen restriction
+    _http_client: httpx.Client | None = field(default=None, repr=False, compare=False)
+    
+    def get_http_client(self) -> httpx.Client:
+        """Get or create a shared HTTP client with connection pooling."""
+        if self._http_client is None:
+            object.__setattr__(
+                self,
+                '_http_client',
+                httpx.Client(
+                    timeout=self.http_timeout,
+                    follow_redirects=True,
+                    limits=httpx.Limits(
+                        max_keepalive_connections=10,
+                        max_connections=20
+                    )
+                )
+            )
+        return self._http_client
+
+
+class McpRequestGuardMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app, *, origin_allowlist: list[str]) -> None:
+        super().__init__(app)
+        self._allowlist = [origin.lower() for origin in origin_allowlist]
+
+    async def dispatch(self, request: Request, call_next):  # type: ignore[override]
+        if request.method == "GET":
+            return Response("Method Not Allowed", status_code=405)
+        if request.method not in {"POST", "OPTIONS"}:
+            return Response("Method Not Allowed", status_code=405)
+        origin = request.headers.get("origin")
+        if origin and not self._is_allowed_origin(origin):
+            return Response("Forbidden", status_code=403)
+        protocol = request.headers.get("mcp-protocol-version")
+        if protocol and protocol not in _SUPPORTED_PROTOCOL_VERSIONS:
+            return Response("Bad Request", status_code=400)
+        return await call_next(request)
+
+    def _is_allowed_origin(self, origin: str) -> bool:
+        if not self._allowlist or "*" in self._allowlist:
+            return True
+        return origin.lower() in self._allowlist
+
+
+_CONFIG: McpSnapshotConfig | None = None
+mcp = FastMCP("Paper DB MCP")
+
+
+def configure(config: McpSnapshotConfig) -> None:
+    global _CONFIG
+    _CONFIG = config
+
+
+def create_mcp_app(config: McpSnapshotConfig) -> tuple[Starlette, Any]:
+    """Create MCP app with middleware and return it along with lifespan.
+    
+    Returns:
+        Tuple of (wrapped_app, lifespan_context) for use by parent Starlette.
+    """
+    configure(config)
+    mcp_app = mcp.http_app(path="/", stateless_http=True)
+    wrapped = Starlette(
+        routes=[Mount("/", app=mcp_app)],
+        middleware=[
+            Middleware(McpRequestGuardMiddleware, origin_allowlist=config.origin_allowlist),
+        ],
+    )
+    return wrapped, mcp_app.lifespan
+
+
+def _get_config() -> McpSnapshotConfig:
+    if _CONFIG is None:
+        raise RuntimeError("MCP server not configured")
+    return _CONFIG
+
+
+def _validate_query(query: str, cfg: McpSnapshotConfig) -> str:
+    """Validate search query string.
+    
+    Raises:
+        McpToolError: If query is invalid or too long.
+    """
+    if not query or not query.strip():
+        raise McpToolError("invalid_query", "Query cannot be empty")
+    if len(query) > cfg.limits.max_query_length:
+        raise McpToolError(
+            "query_too_long",
+            f"Query exceeds maximum length of {cfg.limits.max_query_length}",
+            length=len(query),
+            max_length=cfg.limits.max_query_length
+        )
+    return query.strip()
+
+
+def _validate_paper_id(paper_id: str, cfg: McpSnapshotConfig) -> str:
+    """Validate paper ID format.
+    
+    Raises:
+        McpToolError: If paper_id is invalid.
+    """
+    if not paper_id:
+        raise McpToolError("invalid_paper_id", "Paper ID cannot be empty")
+    if len(paper_id) > cfg.max_paper_id_length:
+        raise McpToolError(
+            "paper_id_too_long",
+            f"Paper ID exceeds maximum length of {cfg.max_paper_id_length}",
+            length=len(paper_id),
+            max_length=cfg.max_paper_id_length
+        )
+    if not _PAPER_ID_PATTERN.match(paper_id):
+        raise McpToolError(
+            "invalid_paper_id_format",
+            "Paper ID must contain only alphanumeric characters, hyphens, and underscores",
+            paper_id=paper_id
+        )
+    return paper_id
+
+
+def _truncate(text: str, max_chars: int | None) -> str:
+    """Truncate text with marker."""
+    if max_chars is None or max_chars <= 0:
+        return text
+    if len(text) <= max_chars:
+        return text
+    remaining = len(text) - max_chars
+    return f"{text[:max_chars]}\n[truncated: {remaining} more chars]"
+
+
+def _read_static_text(rel_path: str) -> str | None:
+    """Read static text from local export directory if available."""
+    cfg = _get_config()
+    if cfg.static_export_dir:
+        path = cfg.static_export_dir / rel_path
+        if path.exists():
+            return path.read_text(encoding="utf-8")
+    return None
+
+
+def _fetch_static_text(rel_path: str) -> str:
+    """Fetch static text from HTTP remote."""
+    cfg = _get_config()
+    if cfg.static_base_url:
+        base = cfg.static_base_url.rstrip("/")
+        url = f"{base}/{rel_path.lstrip('/')}"
+        client = cfg.get_http_client()
+        response = client.get(url)
+        response.raise_for_status()
+        return response.text
+    raise FileNotFoundError("static_base_url not configured")
+
+
+def _load_static_text(rel_path: str) -> str:
+    """Load static text with fallback: local first, then HTTP."""
+    try:
+        text = _read_static_text(rel_path)
+        if text is not None:
+            return text
+        return _fetch_static_text(rel_path)
+    except httpx.HTTPStatusError as exc:
+        raise RuntimeError(f"asset_fetch_failed:{exc.response.status_code}") from exc
+    except httpx.RequestError as exc:
+        raise RuntimeError("asset_fetch_failed:request_error") from exc
+    except FileNotFoundError as exc:
+        raise RuntimeError("asset_fetch_failed:not_configured") from exc
+
+
+def _load_summary_json(paper_id: str, template: str | None) -> tuple[str | None, list[str] | None]:
+    """Load summary JSON content and return available templates list."""
+    cfg = _get_config()
+    conn = _open_ro_conn(cfg.snapshot_db)
+    try:
+        row = conn.execute(
+            "SELECT preferred_summary_template, summary_asset_paths_json FROM paper WHERE paper_id = ?",
+            (paper_id,),
+        ).fetchone()
+        if not row:
+            return None, None
+        preferred = row["preferred_summary_template"]
+        asset_paths = json.loads(row["summary_asset_paths_json"] or "{}")
+        available = sorted(asset_paths.keys())
+        selected = template if template else preferred
+        if not selected or selected not in asset_paths:
+            return None, available
+        rel_path = asset_paths[selected]
+        return _load_static_text(rel_path), available
+    finally:
+        conn.close()
+
+
+def _load_source_markdown(paper_id: str) -> str | None:
+    """Load source markdown for paper."""
+    cfg = _get_config()
+    conn = _open_ro_conn(cfg.snapshot_db)
+    try:
+        row = conn.execute(
+            "SELECT source_md_content_hash FROM paper WHERE paper_id = ?",
+            (paper_id,),
+        ).fetchone()
+        if not row or not row["source_md_content_hash"]:
+            return None
+        rel_path = f"md/{row['source_md_content_hash']}.md"
+        return _load_static_text(rel_path)
+    finally:
+        conn.close()
+
+
+def _load_translation_markdown(paper_id: str, lang: str) -> str | None:
+    """Load translation markdown for paper and language."""
+    cfg = _get_config()
+    conn = _open_ro_conn(cfg.snapshot_db)
+    try:
+        row = conn.execute(
+            "SELECT translations_json FROM paper WHERE paper_id = ?",
+            (paper_id,),
+        ).fetchone()
+        if not row or not row["translations_json"]:
+            return None
+        translations = json.loads(row["translations_json"])
+        rel_path = translations.get(lang)
+        if not rel_path:
+            return None
+        return _load_static_text(rel_path)
+    finally:
+        conn.close()
+
+
+# ==================== MCP Tools ====================
+
+@mcp.tool()
+def search_papers(query: str, limit: int = 10) -> list[dict[str, Any]]:
+    """Full-text search for papers (relevance-ranked).
+    
+    Use when you only have topic keywords.
+    Returns paper_id, title, year, venue, snippet_markdown.
+    """
+    cfg = _get_config()
+    query = _validate_query(query, cfg)
+    limit = min(max(1, int(limit)), cfg.limits.max_page_size)
+    
+    conn = _open_ro_conn(cfg.snapshot_db)
+    try:
+        cur = conn.execute(
+            """
+            SELECT paper_id, title, year, venue, abstract
+            FROM paper_search
+            WHERE paper_search MATCH ?
+            ORDER BY rank
+            LIMIT ?
+            """,
+            (rewrite_search_query(query), limit),
+        )
+        rows = cur.fetchall()
+        results: list[dict[str, Any]] = []
+        for row in rows:
+            snippet = str(row["abstract"] or "")
+            snippet = remove_cjk_spaces(snippet)
+            snippet, markers = merge_adjacent_markers(snippet)
+            results.append({
+                "paper_id": str(row["paper_id"]),
+                "title": str(row["title"]),
+                "year": str(row["year"]),
+                "venue": str(row["venue"]),
+                "snippet_markdown": snippet,
+            })
+        return results
+    finally:
+        conn.close()
+
+
+@mcp.tool()
+def search_papers_by_keyword(keyword: str, limit: int = 10) -> list[dict[str, Any]]:
+    """Search papers by keyword/tag (exact match).
+    
+    Use when you know specific keywords or tags.
+    """
+    cfg = _get_config()
+    limit = min(max(1, int(limit)), cfg.limits.max_page_size)
+    
+    conn = _open_ro_conn(cfg.snapshot_db)
+    try:
+        rows = conn.execute(
+            """
+            SELECT DISTINCT p.paper_id, p.title, p.year, p.venue, p.abstract
+            FROM paper p
+            JOIN paper_keyword pk ON pk.paper_id = p.paper_id
+            JOIN keyword k ON k.keyword_id = pk.keyword_id
+            WHERE k.value LIKE ?
+            ORDER BY p.year DESC, p.title ASC
+            LIMIT ?
+            """,
+            (f"%{keyword}%", limit),
+        ).fetchall()
+        results: list[dict[str, Any]] = []
+        for row in rows:
+            snippet = str(row["abstract"] or "")
+            snippet = remove_cjk_spaces(snippet)
+            snippet, markers = merge_adjacent_markers(snippet)
+            results.append({
+                "paper_id": str(row["paper_id"]),
+                "title": str(row["title"]),
+                "year": str(row["year"]),
+                "venue": str(row["venue"]),
+                "snippet_markdown": snippet,
+            })
+        return results
+    finally:
+        conn.close()
+
+
+@mcp.tool()
+def get_paper_metadata(paper_id: str) -> dict[str, Any]:
+    """Get paper metadata and available summary templates.
+    
+    Call this first before requesting a summary to discover available templates.
+    """
+    cfg = _get_config()
+    paper_id = _validate_paper_id(paper_id, cfg)
+    
+    conn = _open_ro_conn(cfg.snapshot_db)
+    try:
+        row = conn.execute(
+            """
+            SELECT paper_id, title, year, venue, doi, arxiv_id, openreview_id, paper_pw_url,
+                   preferred_summary_template, summary_asset_paths_json
+            FROM paper WHERE paper_id = ?
+            """,
+            (paper_id,),
+        ).fetchone()
+        if not row:
+            raise McpToolError("not_found", "paper not found", paper_id=paper_id)
+        
+        asset_paths = json.loads(row["summary_asset_paths_json"] or "{}")
+        available = sorted(asset_paths.keys())
+        return {
+            "paper_id": str(row["paper_id"]),
+            "title": str(row["title"]),
+            "year": str(row["year"]),
+            "venue": str(row["venue"]),
+            "doi": row["doi"],
+            "arxiv_id": row["arxiv_id"],
+            "openreview_id": row["openreview_id"],
+            "paper_pw_url": row["paper_pw_url"],
+            "preferred_summary_template": row["preferred_summary_template"],
+            "available_summary_templates": available,
+        }
+    finally:
+        conn.close()
+
+
+@mcp.tool()
+def get_paper_summary(paper_id: str, template: str | None = None, max_chars: int | None = None) -> str:
+    """Get summary JSON as raw string.
+    
+    Uses preferred template if template is not specified.
+    Returns the full JSON content (not a URL).
+    """
+    cfg = _get_config()
+    paper_id = _validate_paper_id(paper_id, cfg)
+    max_chars = max_chars if max_chars is not None else cfg.max_chars_default
+    
+    try:
+        payload, available = _load_summary_json(paper_id, template)
+    except RuntimeError as exc:
+        raise McpToolError(
+            "asset_fetch_failed",
+            "Failed to fetch summary asset",
+            paper_id=paper_id,
+            template=template,
+            detail=str(exc),
+        ) from exc
+    
+    if payload is None:
+        raise McpToolError(
+            "template_not_available",
+            "Template not available",
+            paper_id=paper_id,
+            template=template,
+            available_summary_templates=available,
+        )
+    
+    return _truncate(payload, max_chars)
+
+
+@mcp.tool()
+def get_paper_source(paper_id: str, max_chars: int | None = None) -> str:
+    """Get source markdown text.
+    
+    Content may be large; use max_chars to limit size.
+    """
+    cfg = _get_config()
+    paper_id = _validate_paper_id(paper_id, cfg)
+    max_chars = max_chars if max_chars is not None else cfg.max_chars_default
+    
+    try:
+        content = _load_source_markdown(paper_id)
+    except RuntimeError as exc:
+        raise McpToolError(
+            "asset_fetch_failed",
+            "Failed to fetch source asset",
+            paper_id=paper_id,
+            detail=str(exc),
+        ) from exc
+    
+    if content is None:
+        raise McpToolError(
+            "source_not_available",
+            "Source markdown not available",
+            paper_id=paper_id
+        )
+    
+    return _truncate(content, max_chars)
+
+
+@mcp.tool()
+def get_database_stats() -> dict[str, Any]:
+    """Get database statistics.
+    
+    Returns totals, year/month distributions, and top facets
+    (authors, venues, keywords, institutions, tags).
+    """
+    cfg = _get_config()
+    conn = _open_ro_conn(cfg.snapshot_db)
+    try:
+        total_row = conn.execute("SELECT COUNT(*) AS c FROM paper").fetchone()
+        total = int(total_row["c"]) if total_row else 0
+        
+        def top(table: str, limit: int = 20) -> list[dict[str, Any]]:
+            rows = conn.execute(
+                f"SELECT value, paper_count FROM {table} ORDER BY paper_count DESC, value ASC LIMIT ?",
+                (limit,),
+            ).fetchall()
+            return [{"value": str(r["value"]), "paper_count": int(r["paper_count"])} for r in rows]
+        
+        years = conn.execute(
+            """
+            SELECT year AS value, paper_count
+            FROM year_count
+            ORDER BY CASE WHEN year GLOB '[0-9][0-9][0-9][0-9]' THEN 0 ELSE 1 END,
+                     CAST(year AS INT) DESC, year ASC
+            LIMIT 50
+            """,
+        ).fetchall()
+        months = conn.execute(
+            """
+            SELECT month AS value, paper_count
+            FROM month_count
+            ORDER BY CASE WHEN month GLOB '[0-1][0-9]' THEN 0 ELSE 1 END,
+                     CAST(month AS INT) ASC, month ASC
+            """,
+        ).fetchall()
+        
+        return {
+            "total": total,
+            "years": [{"value": str(r["value"]), "paper_count": int(r["paper_count"])} for r in years],
+            "months": [{"value": str(r["value"]), "paper_count": int(r["paper_count"])} for r in months],
+            "authors": top("author"),
+            "venues": top("venue"),
+            "institutions": top("institution"),
+            "keywords": top("keyword"),
+            "tags": top("tag"),
+        }
+    finally:
+        conn.close()
+
+
+@mcp.tool()
+def list_top_facets(category: str, limit: int = 20) -> list[dict[str, Any]]:
+    """List top facet values.
+    
+    Category: author | venue | keyword | institution | tag
+    """
+    table_map = {
+        "author": "author",
+        "venue": "venue",
+        "keyword": "keyword",
+        "institution": "institution",
+        "tag": "tag",
+    }
+    table = table_map.get((category or "").strip().lower())
+    if not table:
+        raise McpToolError(
+            "invalid_category",
+            f"Invalid category: {category}. Must be one of: {', '.join(table_map.keys())}",
+            category=category
+        )
+    
+    limit = max(1, int(limit))
+    cfg = _get_config()
+    conn = _open_ro_conn(cfg.snapshot_db)
+    try:
+        rows = conn.execute(
+            f"SELECT value, paper_count FROM {table} ORDER BY paper_count DESC, value ASC LIMIT ?",
+            (limit,),
+        ).fetchall()
+        return [{"value": str(r["value"]), "paper_count": int(r["paper_count"])} for r in rows]
+    finally:
+        conn.close()
+
+
+@mcp.tool()
+def filter_papers(
+    author: str | None = None,
+    venue: str | None = None,
+    year: str | None = None,
+    keyword: str | None = None,
+    tag: str | None = None,
+    limit: int = 10,
+) -> list[dict[str, Any]]:
+    """Filter papers by structured fields.
+    
+    Use for precise filtering by author, venue, year, keyword, or tag.
+    """
+    cfg = _get_config()
+    limit = min(max(1, int(limit)), cfg.limits.max_page_size)
+    
+    query = "SELECT DISTINCT p.paper_id, p.title, p.year, p.venue FROM paper p"
+    joins: list[str] = []
+    conditions: list[str] = []
+    params: list[Any] = []
+    
+    if author:
+        joins.append("JOIN paper_author pa ON pa.paper_id = p.paper_id")
+        joins.append("JOIN author a ON a.author_id = pa.author_id")
+        conditions.append("a.value LIKE ?")
+        params.append(f"%{author}%")
+    if keyword:
+        joins.append("JOIN paper_keyword pk ON pk.paper_id = p.paper_id")
+        joins.append("JOIN keyword k ON k.keyword_id = pk.keyword_id")
+        conditions.append("k.value LIKE ?")
+        params.append(f"%{keyword}%")
+    if tag:
+        joins.append("JOIN paper_tag pt ON pt.paper_id = p.paper_id")
+        joins.append("JOIN tag t ON t.tag_id = pt.tag_id")
+        conditions.append("t.value LIKE ?")
+        params.append(f"%{tag}%")
+    if venue:
+        conditions.append("p.venue LIKE ?")
+        params.append(f"%{venue}%")
+    if year:
+        conditions.append("p.year = ?")
+        params.append(str(year))
+    
+    if joins:
+        query += " " + " ".join(joins)
+    if conditions:
+        query += " WHERE " + " AND ".join(conditions)
+    query += " ORDER BY p.year DESC, p.title ASC LIMIT ?"
+    params.append(limit)
+    
+    conn = _open_ro_conn(cfg.snapshot_db)
+    try:
+        rows = conn.execute(query, tuple(params)).fetchall()
+        return [
+            {
+                "paper_id": str(row["paper_id"]),
+                "title": str(row["title"]),
+                "year": str(row["year"]),
+                "venue": str(row["venue"]),
+            }
+            for row in rows
+        ]
+    finally:
+        conn.close()
+
+
+# ==================== MCP Resources ====================
+
+@mcp.resource("paper://{paper_id}/metadata")
+def resource_metadata(paper_id: str) -> str:
+    """Resource: metadata as JSON string."""
+    payload = get_paper_metadata(paper_id)
+    return json.dumps(payload, ensure_ascii=False)
+
+
+@mcp.resource("paper://{paper_id}/summary")
+def resource_summary_default(paper_id: str) -> str:
+    """Resource: preferred summary JSON string."""
+    payload = get_paper_summary(paper_id)
+    return payload  # Already a JSON string
+
+
+@mcp.resource("paper://{paper_id}/summary/{template}")
+def resource_summary_template(paper_id: str, template: str) -> str:
+    """Resource: summary JSON string for a specific template."""
+    payload = get_paper_summary(paper_id, template=template)
+    return payload  # Already a JSON string
+
+
+@mcp.resource("paper://{paper_id}/source")
+def resource_source(paper_id: str) -> str:
+    """Resource: source markdown text."""
+    payload = get_paper_source(paper_id)
+    return payload
+
+
+@mcp.resource("paper://{paper_id}/translation/{lang}")
+def resource_translation(paper_id: str, lang: str) -> str:
+    """Resource: translated markdown text."""
+    cfg = _get_config()
+    paper_id = _validate_paper_id(paper_id, cfg)
+    
+    try:
+        content = _load_translation_markdown(paper_id, lang.lower())
+    except RuntimeError as exc:
+        raise McpToolError(
+            "asset_fetch_failed",
+            "Failed to fetch translation asset",
+            paper_id=paper_id,
+            lang=lang,
+            detail=str(exc),
+        ) from exc
+    
+    if content is None:
+        raise McpToolError(
+            "translation_not_available",
+            "Translation not available",
+            paper_id=paper_id,
+            lang=lang,
+        )
+    
+    return _truncate(content, cfg.max_chars_default)
+
+
+def resolve_static_export_dir() -> Path | None:
+    """Resolve static export directory from environment variable."""
+    value = os.getenv("PAPER_DB_STATIC_EXPORT_DIR")
+    if not value:
+        return None
+    return Path(value)