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

deepresearch_flow/paper/db.py

@@ -31,7 +31,8 @@ from deepresearch_flow.paper.template_registry import (
 from deepresearch_flow.paper.render import resolve_render_template, render_papers
 
 try:
-    from pybtex.database import parse_file
+    from pybtex.database import BibliographyData, parse_file
+    from pybtex.database.output.bibtex import Writer
     PYBTEX_AVAILABLE = True
 except ImportError:
     PYBTEX_AVAILABLE = False
@@ -423,6 +424,87 @@ def register_db_commands(db_group: click.Group) -> None:
         click.echo(f"Wrote snapshot DB: {opts.output_db}")
         click.echo(f"Wrote static export: {opts.static_export_dir}")
 
+    @snapshot_group.group("unpack")
+    def snapshot_unpack_group() -> None:
+        """Unpack snapshot artifacts."""
+
+    @snapshot_unpack_group.command("md")
+    @click.option("--snapshot-db", "snapshot_db", required=True, help="Path to snapshot database")
+    @click.option(
+        "--static-export-dir",
+        "static_export_dir",
+        required=True,
+        help="Path to static export directory",
+    )
+    @click.option(
+        "--pdf-root",
+        "pdf_roots",
+        multiple=True,
+        required=True,
+        help="PDF root directories for name alignment (repeatable)",
+    )
+    @click.option("--md-output-dir", "md_output_dir", required=True, help="Output directory for Markdown")
+    @click.option(
+        "--md-translated-output-dir",
+        "md_translated_output_dir",
+        required=True,
+        help="Output directory for translated Markdown",
+    )
+    def snapshot_unpack_md(
+        snapshot_db: str,
+        static_export_dir: str,
+        pdf_roots: tuple[str, ...],
+        md_output_dir: str,
+        md_translated_output_dir: str,
+    ) -> None:
+        """Unpack source/translated markdown and align filenames to PDFs."""
+        from deepresearch_flow.paper.snapshot.unpacker import SnapshotUnpackMdOptions, unpack_md
+
+        opts = SnapshotUnpackMdOptions(
+            snapshot_db=Path(snapshot_db),
+            static_export_dir=Path(static_export_dir),
+            pdf_roots=[Path(path) for path in pdf_roots],
+            md_output_dir=Path(md_output_dir),
+            md_translated_output_dir=Path(md_translated_output_dir),
+        )
+        unpack_md(opts)
+
+    @snapshot_unpack_group.command("info")
+    @click.option("--snapshot-db", "snapshot_db", required=True, help="Path to snapshot database")
+    @click.option(
+        "--static-export-dir",
+        "static_export_dir",
+        required=True,
+        help="Path to static export directory",
+    )
+    @click.option(
+        "--pdf-root",
+        "pdf_roots",
+        multiple=True,
+        required=True,
+        help="PDF root directories for name alignment (repeatable)",
+    )
+    @click.option("--template", "template", required=True, help="Summary template tag")
+    @click.option("--output-json", "output_json", required=True, help="Output JSON file path")
+    def snapshot_unpack_info(
+        snapshot_db: str,
+        static_export_dir: str,
+        pdf_roots: tuple[str, ...],
+        template: str,
+        output_json: str,
+    ) -> None:
+        """Unpack aggregated paper_infos.json from snapshot summaries."""
+        from deepresearch_flow.paper.snapshot.unpacker import SnapshotUnpackInfoOptions, unpack_info
+
+        opts = SnapshotUnpackInfoOptions(
+            snapshot_db=Path(snapshot_db),
+            static_export_dir=Path(static_export_dir),
+            pdf_roots=[Path(path) for path in pdf_roots],
+            template=template,
+            output_json=Path(output_json),
+        )
+        unpack_info(opts)
+
     @db_group.group("api")
     def api_group() -> None:
         """Read-only JSON API server backed by a snapshot DB."""
@@ -1015,11 +1097,18 @@ def register_db_commands(db_group: click.Group) -> None:
     def merge_group() -> None:
         """Merge paper JSON inputs."""
 
-    def _summarize_merge(output_path: Path, merged: list[dict[str, Any]], *, input_count: int) -> None:
+    def _summarize_merge(output_path: Path, merged: Any, *, input_count: int) -> None:
+        items: list[dict[str, Any]] = []
+        if isinstance(merged, dict):
+            raw_items = merged.get("papers")
+            if isinstance(raw_items, list):
+                items = [item for item in raw_items if isinstance(item, dict)]
+        elif isinstance(merged, list):
+            items = [item for item in merged if isinstance(item, dict)]
+
         field_set: set[str] = set()
-        for item in merged:
-            if isinstance(item, dict):
-                field_set.update(item.keys())
+        for item in items:
+            field_set.update(item.keys())
         field_list = sorted(field_set)
 
         console = Console()
@@ -1027,7 +1116,7 @@ def register_db_commands(db_group: click.Group) -> None:
         summary.add_column("Metric", style="bold")
         summary.add_column("Value")
         summary.add_row("Inputs", str(input_count))
-        summary.add_row("Items", str(len(merged)))
+        summary.add_row("Items", str(len(items)))
         summary.add_row("Fields", str(len(field_list)))
         summary.add_row("Output", str(output_path))
         console.print(summary)
@@ -1039,17 +1128,65 @@ def register_db_commands(db_group: click.Group) -> None:
                 field_table.add_row(name)
             console.print(field_table)
 
+    def _bibtex_entry_score(entry: Any) -> int:
+        fields = getattr(entry, "fields", {}) or {}
+        persons = getattr(entry, "persons", {}) or {}
+        person_count = sum(len(people) for people in persons.values())
+        return len(fields) + len(persons) + person_count
+
+    def _summarize_bibtex_merge(output_path: Path, *, input_count: int, entry_count: int, duplicate_count: int) -> None:
+        summary = Table(title="BibTeX Merge Summary")
+        summary.add_column("Metric", style="bold")
+        summary.add_column("Value")
+        summary.add_row("Inputs", str(input_count))
+        summary.add_row("Entries", str(entry_count))
+        summary.add_row("Duplicates", str(duplicate_count))
+        summary.add_row("Output", str(output_path))
+        Console().print(summary)
+
     @merge_group.command("library")
     @click.option("-i", "--inputs", "input_paths", multiple=True, required=True, help="Input JSON files")
+    @click.option("--template-tag", "template_tag", default=None, help="Template tag for merged output")
     @click.option("-o", "--output", "output_path", required=True, help="Output JSON file path")
-    def merge_library(input_paths: Iterable[str], output_path: str) -> None:
+    def merge_library(input_paths: Iterable[str], template_tag: str | None, output_path: str) -> None:
         paths = [Path(path) for path in input_paths]
         merged: list[dict[str, Any]] = []
+        tag_candidates: list[str] = []
         for path in paths:
-            merged.extend(load_json(path))
+            payload = load_json(path)
+            if isinstance(payload, dict):
+                tag = str(payload.get("template_tag") or "")
+                if tag:
+                    tag_candidates.append(tag)
+                papers = payload.get("papers")
+                if isinstance(papers, list):
+                    merged.extend(papers)
+                else:
+                    raise click.ClickException("Input JSON must be a list or {template_tag, papers}")
+            elif isinstance(payload, list):
+                merged.extend(payload)
+            else:
+                raise click.ClickException("Input JSON must be a list or {template_tag, papers}")
+        if not template_tag:
+            inferred = ""
+            for paper in merged:
+                if not isinstance(paper, dict):
+                    continue
+                inferred = str(paper.get("prompt_template") or paper.get("template_tag") or "")
+                if inferred:
+                    break
+            if inferred:
+                template_tag = inferred
+        if tag_candidates and not template_tag:
+            template_tag = tag_candidates[0]
+        if not template_tag:
+            template_tag = "unknown"
+        if tag_candidates and any(tag != template_tag for tag in tag_candidates):
+            click.echo("Warning: multiple template_tag values detected in inputs; using first")
         output = Path(output_path)
-        write_json(output, merged)
-        _summarize_merge(output, merged, input_count=len(paths))
+        bundle = {"template_tag": template_tag, "papers": merged}
+        write_json(output, bundle)
+        _summarize_merge(output, bundle, input_count=len(paths))
 
     @merge_group.command("templates")
     @click.option("-i", "--inputs", "input_paths", multiple=True, required=True, help="Input JSON files")
@@ -1201,6 +1338,62 @@ def register_db_commands(db_group: click.Group) -> None:
                 sample_table.add_row(*row)
             Console().print(sample_table)
 
+    @merge_group.command("bibtex")
+    @click.option("-i", "--input", "input_paths", multiple=True, required=True, help="Input BibTeX file paths")
+    @click.option("-o", "--output", "output_path", required=True, help="Output BibTeX file path")
+    def merge_bibtex(input_paths: Iterable[str], output_path: str) -> None:
+        if not PYBTEX_AVAILABLE:
+            raise click.ClickException("pybtex is required for merge bibtex")
+
+        paths = [Path(path) for path in input_paths]
+        if not paths:
+            raise click.ClickException("No BibTeX inputs provided")
+
+        for path in paths:
+            if not path.is_file():
+                raise click.ClickException(f"BibTeX file not found: {path}")
+
+        merged_entries: dict[str, tuple[Any, int]] = {}
+        duplicate_keys: list[str] = []
+        duplicate_seen: set[str] = set()
+
+        for path in paths:
+            bib_data = parse_file(str(path))
+            for key, entry in bib_data.entries.items():
+                score = _bibtex_entry_score(entry)
+                if key not in merged_entries:
+                    merged_entries[key] = (entry, score)
+                    continue
+                if key not in duplicate_seen:
+                    duplicate_seen.add(key)
+                    duplicate_keys.append(key)
+                _, existing_score = merged_entries[key]
+                if score > existing_score:
+                    merged_entries[key] = (entry, score)
+
+        output = Path(output_path)
+        output.parent.mkdir(parents=True, exist_ok=True)
+        out_data = BibliographyData()
+        for key, (entry, _) in merged_entries.items():
+            out_data.entries[key] = entry
+        with output.open("w", encoding="utf-8") as handle:
+            Writer().write_stream(out_data, handle)
+
+        _summarize_bibtex_merge(
+            output,
+            input_count=len(paths),
+            entry_count=len(merged_entries),
+            duplicate_count=len(duplicate_keys),
+        )
+
+        if duplicate_keys:
+            preview_limit = 20
+            preview = ", ".join(duplicate_keys[:preview_limit])
+            if len(duplicate_keys) > preview_limit:
+                preview = f"{preview}, ... (+{len(duplicate_keys) - preview_limit} more)"
+            note = "Kept entry with most fields; ties keep first input order."
+            Console().print(Panel(f"{note}\n{preview}", title=f"Duplicate keys ({len(duplicate_keys)})", style="yellow"))
+
     @db_group.command("render-md")
     @click.option("-i", "--input", "input_path", required=True, help="Input JSON file path")
     @click.option("-d", "--output-dir", "output_dir", default="rendered_md", help="Output directory")

deepresearch_flow/paper/snapshot/api.py

@@ -11,8 +11,9 @@ from starlette.applications import Starlette
 from starlette.middleware.cors import CORSMiddleware
 from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
-from starlette.routing import Route
+from starlette.routing import Mount, Route
 
+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
 
 _WHITESPACE_RE = re.compile(r"\s+")
@@ -87,13 +88,6 @@ _FACET_TYPE_TO_KEY = {
 }
 
 
-@dataclass(frozen=True)
-class ApiLimits:
-    max_query_length: int = 500
-    max_page_size: int = 100
-    max_pagination_offset: int = 10_000  # page * page_size
-
-
 @dataclass(frozen=True)
 class SnapshotApiConfig:
     snapshot_db: Path
@@ -110,12 +104,6 @@ def _json_error(status_code: int, *, error: str, detail: str) -> JSONResponse:
     return JSONResponse({"error": error, "detail": detail}, status_code=status_code)
 
 
-def _open_ro_conn(db_path: Path) -> sqlite3.Connection:
-    uri = f"file:{db_path.as_posix()}?mode=ro"
-    conn = sqlite3.connect(uri, uri=True)
-    conn.row_factory = sqlite3.Row
-    conn.execute("PRAGMA query_only=ON;")
-    return conn
 
 
 def _snapshot_build_id(conn: sqlite3.Connection) -> str:
@@ -917,6 +905,22 @@ def create_app(
         limits=limits or ApiLimits(),
     )
 
+    # Lazy import to avoid circular dependency
+    from deepresearch_flow.paper.snapshot.mcp_server import (
+        McpSnapshotConfig,
+        create_mcp_app,
+        resolve_static_export_dir,
+    )
+
+    mcp_config = McpSnapshotConfig(
+        snapshot_db=snapshot_db,
+        static_base_url=_normalize_base_url(static_base_url),
+        static_export_dir=resolve_static_export_dir(),
+        limits=limits or ApiLimits(),
+        origin_allowlist=cors_allowed_origins or ["*"],
+    )
+    mcp_app, mcp_lifespan = create_mcp_app(mcp_config)
+
     routes = [
         Route("/api/v1/config", _api_config, methods=["GET"]),
         Route("/api/v1/search", _api_search, methods=["GET"]),
@@ -927,9 +931,15 @@ def create_app(
         Route("/api/v1/facets/{facet:str}/{facet_id:str}/stats", _api_facet_stats, methods=["GET"]),
         Route("/api/v1/facets/{facet:str}/by-value/{value:str}/papers", _api_facet_by_value_papers, methods=["GET"]),
         Route("/api/v1/facets/{facet:str}/by-value/{value:str}/stats", _api_facet_by_value_stats, methods=["GET"]),
+        Mount("/mcp", app=mcp_app),
     ]
 
-    app = Starlette(routes=routes)
+    # Pass MCP lifespan to ensure session manager initializes properly
+    # https://gofastmcp.com/deployment/http#mounting-in-starlette
+    app = Starlette(
+        routes=routes,
+        lifespan=mcp_lifespan,
+    )
     if cfg.cors_allowed_origins:
         app.add_middleware(
             CORSMiddleware,

deepresearch_flow/paper/snapshot/common.py

@@ -0,0 +1,34 @@
+"""Shared utilities for snapshot API and MCP server.
+
+This module contains common types, configuration, and utilities used by both
+the snapshot REST API and the MCP server to avoid circular imports.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+import sqlite3
+
+
+@dataclass(frozen=True)
+class ApiLimits:
+    """API rate and size limits."""
+
+    max_query_length: int = 500
+    max_page_size: int = 100
+    max_pagination_offset: int = 10_000
+
+
+def _open_ro_conn(db_path: Path) -> sqlite3.Connection:
+    """Open a read-only SQLite connection with Row factory.
+
+    Args:
+        db_path: Path to the SQLite database file.
+
+    Returns:
+        sqlite3.Connection: A read-only connection with row_factory set to Row.
+    """
+    conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True, check_same_thread=False)
+    conn.row_factory = sqlite3.Row
+    return conn