deepresearch-flow 0.1.1__py3-none-any.whl

deepresearch_flow/paper/providers/ollama.py

@@ -0,0 +1,46 @@
+"""Ollama provider implementation."""
+
+from __future__ import annotations
+
+from typing import Any
+import httpx
+
+from deepresearch_flow.paper.providers.base import ProviderError
+
+
+async def chat(
+    client: httpx.AsyncClient,
+    base_url: str,
+    model: str,
+    messages: list[dict[str, str]],
+    structured_mode: str,
+    headers: dict[str, str] | None,
+    timeout: float,
+) -> str:
+    url = base_url.rstrip("/") + "/api/chat"
+    payload: dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+        "stream": False,
+    }
+    if structured_mode in ("json_schema", "json_object"):
+        payload["format"] = "json"
+
+    try:
+        response = await client.post(url, json=payload, headers=headers, timeout=timeout)
+    except httpx.RequestError as exc:
+        raise ProviderError(str(exc), retryable=True) from exc
+
+    if response.status_code == 429:
+        raise ProviderError(response.text, status_code=429, retryable=True)
+    if response.status_code >= 500:
+        raise ProviderError(response.text, status_code=response.status_code, retryable=True)
+    if response.status_code >= 400:
+        raise ProviderError(response.text, status_code=response.status_code)
+
+    data = response.json()
+    message = data.get("message", {})
+    content = message.get("content")
+    if not content:
+        raise ProviderError("Ollama response missing content", status_code=response.status_code)
+    return content

deepresearch_flow/paper/providers/openai_compatible.py

@@ -0,0 +1,60 @@
+"""OpenAI-compatible provider implementation."""
+
+from __future__ import annotations
+
+from typing import Any
+import httpx
+
+from deepresearch_flow.paper.providers.base import ProviderError
+
+
+async def chat(
+    client: httpx.AsyncClient,
+    base_url: str,
+    model: str,
+    messages: list[dict[str, str]],
+    structured_mode: str,
+    headers: dict[str, str] | None,
+    timeout: float,
+    schema: dict[str, Any] | None = None,
+) -> str:
+    url = base_url.rstrip("/") + "/chat/completions"
+    payload: dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+        "temperature": 0,
+    }
+
+    if structured_mode == "json_schema" and schema is not None:
+        payload["response_format"] = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": "paper_extract",
+                "schema": schema,
+            },
+        }
+    elif structured_mode == "json_object":
+        payload["response_format"] = {"type": "json_object"}
+
+    try:
+        response = await client.post(url, json=payload, headers=headers, timeout=timeout)
+    except httpx.RequestError as exc:
+        raise ProviderError(str(exc), retryable=True) from exc
+
+    if response.status_code == 429:
+        raise ProviderError(response.text, status_code=429, retryable=True)
+    if response.status_code >= 500:
+        raise ProviderError(response.text, status_code=response.status_code, retryable=True)
+    if response.status_code >= 400:
+        structured_error = structured_mode in ("json_schema", "json_object")
+        raise ProviderError(response.text, status_code=response.status_code, structured_error=structured_error)
+
+    data = response.json()
+    choices = data.get("choices") or []
+    if not choices:
+        raise ProviderError("OpenAI-compatible response missing choices", status_code=response.status_code)
+    message = choices[0].get("message", {})
+    content = message.get("content")
+    if not content:
+        raise ProviderError("OpenAI-compatible response missing content", status_code=response.status_code)
+    return content

deepresearch_flow/paper/render.py

@@ -0,0 +1,64 @@
+"""Markdown rendering helpers for paper outputs."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Iterable
+import importlib.resources as resources
+
+from jinja2 import Environment, FileSystemLoader, Template
+
+from deepresearch_flow.paper.template_registry import load_render_template
+from deepresearch_flow.paper.utils import split_output_name, unique_split_name
+
+
+def load_default_template() -> Template:
+    template_path = resources.files("deepresearch_flow.paper.templates").joinpath(
+        "default_paper.md.j2"
+    )
+    with template_path.open("r", encoding="utf-8") as handle:
+        return Environment().from_string(handle.read())
+
+
+def resolve_render_template(
+    template_path: str | None,
+    template_name: str | None,
+    template_dir: str | None,
+) -> Template:
+    if sum(bool(item) for item in (template_path, template_name, template_dir)) > 1:
+        raise ValueError(
+            "Use only one of template path, template name, or template directory"
+        )
+    if template_dir:
+        template_path = str(Path(template_dir) / "render.j2")
+    if template_path:
+        return Environment(loader=FileSystemLoader(Path(template_path).parent)).get_template(
+            Path(template_path).name
+        )
+    if template_name:
+        return load_render_template(template_name)
+    return load_default_template()
+
+
+def render_papers(
+    papers: Iterable[dict[str, Any]],
+    output_dir: Path,
+    template: Template,
+    fallback_language: str,
+) -> int:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    used: set[str] = set()
+    count = 0
+    for paper in papers:
+        source_path = paper.get("source_path")
+        base = paper.get("paper_title") or "paper"
+        if source_path:
+            base = split_output_name(Path(source_path))
+        filename = unique_split_name(base, used, source_path or base)
+        context = dict(paper)
+        if not context.get("output_language"):
+            context["output_language"] = fallback_language
+        markdown = template.render(**context)
+        (output_dir / f"{filename}.md").write_text(markdown, encoding="utf-8")
+        count += 1
+    return count

deepresearch_flow/paper/schema.py

@@ -0,0 +1,58 @@
+"""Schema loading and validation helpers."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+import importlib.resources as resources
+
+from jsonschema import Draft7Validator
+
+
+class SchemaError(ValueError):
+    pass
+
+
+def load_schema(schema_path: str | None) -> dict[str, Any]:
+    if schema_path:
+        path = Path(schema_path)
+        if not path.exists():
+            raise SchemaError(f"Schema file not found: {path}")
+        return json.loads(path.read_text(encoding="utf-8"))
+
+    with resources.files("deepresearch_flow.paper.schemas").joinpath(
+        "default_paper_schema.json"
+    ).open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def validate_schema(schema: dict[str, Any]) -> Draft7Validator:
+    if schema.get("type") != "object":
+        raise SchemaError("Schema must describe a JSON object at the top level")
+
+    required = schema.get("required", [])
+    if "paper_title" not in required:
+        raise SchemaError("Schema must require 'paper_title'")
+    if "paper_authors" not in required:
+        raise SchemaError("Schema must require 'paper_authors'")
+
+    properties = schema.get("properties", {})
+    authors_schema = properties.get("paper_authors", {})
+    if authors_schema.get("type") != "array":
+        raise SchemaError("'paper_authors' must be an array in the schema")
+
+    items = authors_schema.get("items", {})
+    if items.get("type") != "string":
+        raise SchemaError("'paper_authors' items must be strings")
+
+    schema_for_validation = dict(schema)
+    if schema_for_validation.get("additionalProperties", True) is False:
+        schema_for_validation = dict(schema_for_validation)
+        schema_for_validation["additionalProperties"] = True
+
+    return Draft7Validator(schema_for_validation)
+
+
+def schema_to_prompt(schema: dict[str, Any]) -> str:
+    return json.dumps(schema, ensure_ascii=False, indent=2)

deepresearch_flow/paper/schemas/__init__.py

@@ -0,0 +1 @@
+"""Built-in JSON schemas."""

deepresearch_flow/paper/schemas/deep_read_schema.json

@@ -0,0 +1,46 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "type": "object",
+  "additionalProperties": false,
+  "required": [
+    "paper_title",
+    "paper_authors",
+    "publication_date",
+    "publication_venue",
+    "module_a",
+    "module_b",
+    "module_c1",
+    "module_c2",
+    "module_c3",
+    "module_c4",
+    "module_c5",
+    "module_c6",
+    "module_c7",
+    "module_c8",
+    "module_d",
+    "module_e",
+    "module_h"
+  ],
+  "properties": {
+    "paper_title": {"type": "string", "minLength": 1},
+    "paper_authors": {"type": "array", "items": {"type": "string"}},
+    "publication_date": {"type": "string"},
+    "publication_venue": {"type": "string"},
+    "module_a": {"type": "string"},
+    "module_b": {"type": "string"},
+    "module_c1": {"type": "string"},
+    "module_c2": {"type": "string"},
+    "module_c3": {"type": "string"},
+    "module_c4": {"type": "string"},
+    "module_c5": {"type": "string"},
+    "module_c6": {"type": "string"},
+    "module_c7": {"type": "string"},
+    "module_c8": {"type": "string"},
+    "module_d": {"type": "string"},
+    "module_e": {"type": "string"},
+    "module_f": {"type": "string"},
+    "module_g": {"type": "string"},
+    "module_h": {"type": "string"},
+    "output_language": {"type": "string"}
+  }
+}

deepresearch_flow/paper/schemas/default_paper_schema.json

@@ -0,0 +1,47 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "type": "object",
+  "additionalProperties": false,
+  "required": [
+    "paper_title",
+    "paper_authors",
+    "publication_date",
+    "publication_venue",
+    "abstract",
+    "keywords",
+    "summary"
+  ],
+  "properties": {
+    "paper_title": {
+      "type": "string",
+      "minLength": 1
+    },
+    "paper_authors": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      }
+    },
+    "abstract": {
+      "type": "string"
+    },
+    "publication_date": {
+      "type": "string"
+    },
+    "publication_venue": {
+      "type": "string"
+    },
+    "keywords": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      }
+    },
+    "summary": {
+      "type": "string"
+    },
+    "output_language": {
+      "type": "string"
+    }
+  }
+}

deepresearch_flow/paper/schemas/eight_questions_schema.json

@@ -0,0 +1,34 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "type": "object",
+  "additionalProperties": false,
+  "required": [
+    "paper_title",
+    "paper_authors",
+    "publication_date",
+    "publication_venue",
+    "question1",
+    "question2",
+    "question3",
+    "question4",
+    "question5",
+    "question6",
+    "question7",
+    "question8"
+  ],
+  "properties": {
+    "paper_title": {"type": "string", "minLength": 1},
+    "paper_authors": {"type": "array", "items": {"type": "string"}},
+    "publication_date": {"type": "string"},
+    "publication_venue": {"type": "string"},
+    "question1": {"type": "string"},
+    "question2": {"type": "string"},
+    "question3": {"type": "string"},
+    "question4": {"type": "string"},
+    "question5": {"type": "string"},
+    "question6": {"type": "string"},
+    "question7": {"type": "string"},
+    "question8": {"type": "string"},
+    "output_language": {"type": "string"}
+  }
+}

deepresearch_flow/paper/schemas/three_pass_schema.json

@@ -0,0 +1,24 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "type": "object",
+  "additionalProperties": false,
+  "required": [
+    "paper_title",
+    "paper_authors",
+    "publication_date",
+    "publication_venue",
+    "step1_summary",
+    "step2_analysis",
+    "step3_analysis"
+  ],
+  "properties": {
+    "paper_title": {"type": "string", "minLength": 1},
+    "paper_authors": {"type": "array", "items": {"type": "string"}},
+    "publication_date": {"type": "string"},
+    "publication_venue": {"type": "string"},
+    "step1_summary": {"type": "string"},
+    "step2_analysis": {"type": "string"},
+    "step3_analysis": {"type": "string"},
+    "output_language": {"type": "string"}
+  }
+}

deepresearch_flow/paper/template_registry.py

@@ -0,0 +1,189 @@
+"""Template registry for extract prompts and render output."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+from typing import Any
+import importlib.resources as resources
+
+from jinja2 import Environment
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class TemplateBundle:
+    name: str
+    prompt_system: str
+    prompt_user: str
+    schema_file: str
+    render_template: str
+
+
+@dataclass(frozen=True)
+class StageDefinition:
+    name: str
+    fields: list[str]
+
+
+_TEMPLATES: dict[str, TemplateBundle] = {
+    "simple": TemplateBundle(
+        name="simple",
+        prompt_system="simple_system.j2",
+        prompt_user="simple_user.j2",
+        schema_file="default_paper_schema.json",
+        render_template="default_paper.md.j2",
+    ),
+    "deep_read": TemplateBundle(
+        name="deep_read",
+        prompt_system="deep_read_system.j2",
+        prompt_user="deep_read_user.j2",
+        schema_file="deep_read_schema.json",
+        render_template="deep_read.md.j2",
+    ),
+    "eight_questions": TemplateBundle(
+        name="eight_questions",
+        prompt_system="eight_questions_system.j2",
+        prompt_user="eight_questions_user.j2",
+        schema_file="eight_questions_schema.json",
+        render_template="eight_questions.md.j2",
+    ),
+    "three_pass": TemplateBundle(
+        name="three_pass",
+        prompt_system="three_pass_system.j2",
+        prompt_user="three_pass_user.j2",
+        schema_file="three_pass_schema.json",
+        render_template="three_pass.md.j2",
+    ),
+}
+
+_STAGES: dict[str, list[StageDefinition]] = {
+    "deep_read": [
+        StageDefinition("module_a", ["module_a"]),
+        StageDefinition("module_b", ["module_b"]),
+        StageDefinition("module_c1", ["module_c1"]),
+        StageDefinition("module_c2", ["module_c2"]),
+        StageDefinition("module_c3", ["module_c3"]),
+        StageDefinition("module_c4", ["module_c4"]),
+        StageDefinition("module_c5", ["module_c5"]),
+        StageDefinition("module_c6", ["module_c6"]),
+        StageDefinition("module_c7", ["module_c7"]),
+        StageDefinition("module_c8", ["module_c8"]),
+        StageDefinition("module_d", ["module_d"]),
+        StageDefinition("module_e", ["module_e"]),
+        StageDefinition("module_f", ["module_f"]),
+        StageDefinition("module_g", ["module_g"]),
+        StageDefinition("module_h", ["module_h"]),
+    ],
+    "eight_questions": [
+        StageDefinition(
+            "questions_1to4",
+            ["question1", "question2", "question3", "question4"],
+        ),
+        StageDefinition(
+            "questions_5to8",
+            ["question5", "question6", "question7", "question8"],
+        ),
+    ],
+    "three_pass": [
+        StageDefinition("step1_summary", ["step1_summary"]),
+        StageDefinition("step2_analysis", ["step2_analysis"]),
+        StageDefinition("step3_analysis", ["step3_analysis"]),
+    ],
+}
+
+
+def list_template_names() -> list[str]:
+    return sorted(_TEMPLATES.keys())
+
+
+def list_template_names_in_registry_order() -> list[str]:
+    return list(_TEMPLATES.keys())
+
+
+def get_stage_definitions(template_name: str) -> list[StageDefinition]:
+    return _STAGES.get(template_name, [])
+
+
+def get_template_bundle(name: str) -> TemplateBundle:
+    try:
+        return _TEMPLATES[name]
+    except KeyError as exc:
+        available = ", ".join(list_template_names())
+        raise ValueError(f"Unknown template '{name}'. Available: {available}") from exc
+
+
+def load_schema_for_template(name: str) -> dict[str, Any]:
+    bundle = get_template_bundle(name)
+    schema_path = resources.files("deepresearch_flow.paper.schemas").joinpath(bundle.schema_file)
+    with schema_path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def load_prompt_templates(
+    name: str,
+    *,
+    content: str,
+    schema: str,
+    output_language: str,
+    stage_name: str | None = None,
+    stage_fields: list[str] | None = None,
+    previous_outputs: str | None = None,
+) -> tuple[str, str]:
+    bundle = get_template_bundle(name)
+    env = Environment()
+    stage_fields = stage_fields or []
+    previous_outputs = previous_outputs or ""
+    system_text = _render_prompt_template(
+        "deepresearch_flow.paper.prompt_templates",
+        bundle.prompt_system,
+        env,
+        {
+            "output_language": output_language,
+            "stage_name": stage_name,
+            "stage_fields": stage_fields,
+            "previous_outputs": previous_outputs,
+        },
+    )
+    user_text = _render_prompt_template(
+        "deepresearch_flow.paper.prompt_templates",
+        bundle.prompt_user,
+        env,
+        {
+            "content": content,
+            "schema": schema,
+            "output_language": output_language,
+            "stage_name": stage_name,
+            "stage_fields": stage_fields,
+            "previous_outputs": previous_outputs,
+        },
+    )
+    return system_text, user_text
+
+
+def load_render_template(name: str):
+    bundle = get_template_bundle(name)
+    template_path = resources.files("deepresearch_flow.paper.templates").joinpath(
+        bundle.render_template
+    )
+    with template_path.open("r", encoding="utf-8") as handle:
+        return Environment().from_string(handle.read())
+
+
+def load_custom_prompt_templates(
+    system_path: Path,
+    user_path: Path,
+    context: dict[str, Any],
+) -> tuple[str, str]:
+    env = Environment()
+    system_text = env.from_string(system_path.read_text(encoding="utf-8")).render(**context)
+    user_text = env.from_string(user_path.read_text(encoding="utf-8")).render(**context)
+    return system_text, user_text
+
+
+def _render_prompt_template(
+    package: str, filename: str, env: Environment, context: dict[str, Any]
+) -> str:
+    template_path = resources.files(package).joinpath(filename)
+    with template_path.open("r", encoding="utf-8") as handle:
+        return env.from_string(handle.read()).render(**context)

deepresearch_flow/paper/templates/__init__.py

@@ -0,0 +1 @@
+"""Built-in Markdown templates."""

deepresearch_flow/paper/templates/deep_read.md.j2

@@ -0,0 +1,79 @@
+# {{ paper_title }}
+
+{% set is_zh = output_language == "zh" %}
+**{{ "作者 / Authors" if is_zh else "Authors" }}:** {{ paper_authors | join(", ") }}
+
+{% if output_language %}
+**{{ "输出语言 / Output Language" if is_zh else "Output Language" }}:** {{ output_language }}
+{% endif %}
+
+{% if publication_date %}
+**{{ "发表日期 / Publication Date" if is_zh else "Publication Date" }}:** {{ publication_date }}
+{% endif %}
+
+{% if publication_venue %}
+**{{ "期刊/会议 / Publication Venue" if is_zh else "Publication Venue" }}:** {{ publication_venue }}
+{% endif %}
+
+## {{ "模块 A：阅读对齐与输入校验 / Module A: Reading Alignment and Input Check" if is_zh else "Module A: Reading Alignment and Input Check" }}
+
+{{ module_a }}
+
+## {{ "模块 B：第一遍鸟瞰扫描（类别/上下文/正确性/贡献/清晰度） / Module B: First Pass Overview (category, context, correctness, contribution, clarity)" if is_zh else "Module B: First Pass Overview (category, context, correctness, contribution, clarity)" }}
+
+{{ module_b }}
+
+## {{ "模块 C1：论文试图解决什么问题？ / Module C1: What problem does the paper try to solve?" if is_zh else "Module C1: What problem does the paper try to solve?" }}
+
+{{ module_c1 }}
+
+## {{ "模块 C2：有哪些相关研究？（最强基线/最相近思路/正交方向） / Module C2: What related research exists? (strongest baseline, closest approach, orthogonal directions)" if is_zh else "Module C2: What related research exists? (strongest baseline, closest approach, orthogonal directions)" }}
+
+{{ module_c2 }}
+
+## {{ "模块 C3：论文如何解决这个问题？（流程图文字版） / Module C3: How does the paper solve the problem? (include a text flowchart)" if is_zh else "Module C3: How does the paper solve the problem? (include a text flowchart)" }}
+
+{{ module_c3 }}
+
+## {{ "模块 C4：实验数据集与实现设置（数据集/场景、为何选择、验证哪些关键特性） / Module C4: Datasets and implementation settings (what datasets/scenarios, why chosen, what they validate)" if is_zh else "Module C4: Datasets and implementation settings (what datasets/scenarios, why chosen, what they validate)" }}
+
+{{ module_c4 }}
+
+## {{ "模块 C5：指标与结果（指标、结果、各指标对应的创新点） / Module C5: Metrics and results (metrics, results, which innovations each metric reflects)" if is_zh else "Module C5: Metrics and results (metrics, results, which innovations each metric reflects)" }}
+
+{{ module_c5 }}
+
+## {{ "模块 C6：有什么可以进一步探索的点？（理论/工程/数据/评测） / Module C6: What can be further explored? (theory/engineering/data/evaluation)" if is_zh else "Module C6: What can be further explored? (theory/engineering/data/evaluation)" }}
+
+{{ module_c6 }}
+
+## {{ "模块 C7：现实场景如何体现/实现？（部署约束、算力、延迟、鲁棒性） / Module C7: How would this work in real-world scenarios? (deployment constraints, compute, latency, robustness)" if is_zh else "Module C7: How would this work in real-world scenarios? (deployment constraints, compute, latency, robustness)" }}
+
+{{ module_c7 }}
+
+## {{ "模块 C8：总结主要内容（一段摘要 + 5 条要点 + 1 个核心公式/机制） / Module C8: Summarize the main content (one paragraph + 5 bullets + one core formula/mechanism)." if is_zh else "Module C8: Summarize the main content (one paragraph + 5 bullets + one core formula/mechanism)." }}
+
+{{ module_c8 }}
+
+## {{ "模块 D：第二遍内容理解（方法结构/关键图表/实验） / Module D: Second Pass (method structure, key figures, experiments)" if is_zh else "Module D: Second Pass (method structure, key figures, experiments)" }}
+
+{{ module_d }}
+
+## {{ "模块 E：第三遍深度审视（重建/假设/失败模式/审稿意见） / Module E: Third Pass (reconstruction, assumptions, failure modes, review)" if is_zh else "Module E: Third Pass (reconstruction, assumptions, failure modes, review)" }}
+
+{{ module_e }}
+
+{% if module_f %}## {{ "模块 F：术语表与概念依赖图 / Module F: Glossary and Concept Map" if is_zh else "Module F: Glossary and Concept Map" }}
+
+{{ module_f }}
+
+{% endif %}
+
+{% if module_g %}## {{ "模块 G：联网核验与引用校对 / Module G: Verification and Citation Check" if is_zh else "Module G: Verification and Citation Check" }}
+
+{{ module_g }}
+{% endif %}
+
+## {{ "模块 H：图表与表格抽取（索引/打分/深读） / Module H: Figure/Table Mining (index, scoring, top deep dive)" if is_zh else "Module H: Figure/Table Mining (index, scoring, top deep dive)" }}
+
+{{ module_h }}