deepresearch-flow 0.1.1__py3-none-any.whl

deepresearch_flow/__init__.py

@@ -0,0 +1,5 @@
+"""DeepResearch Flow package."""
+
+__all__ = ["__version__"]
+
+__version__ = "0.1.0"

deepresearch_flow/cli.py

@@ -0,0 +1,23 @@
+"""CLI entrypoint for deepresearch-flow."""
+
+import click
+
+from deepresearch_flow.paper.cli import paper
+from deepresearch_flow.recognize.cli import recognize
+
+
+@click.group()
+def cli() -> None:
+    """DeepResearch Flow command line interface."""
+
+
+cli.add_command(paper)
+cli.add_command(recognize)
+
+
+def main() -> None:
+    cli()
+
+
+if __name__ == "__main__":
+    main()

deepresearch_flow/paper/__init__.py

@@ -0,0 +1 @@
+"""Paper extraction and database tools."""

deepresearch_flow/paper/cli.py

@@ -0,0 +1,286 @@
+"""CLI commands for paper workflows."""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+import click
+
+from deepresearch_flow.paper.config import load_config, resolve_api_keys
+from deepresearch_flow.paper.extract import extract_documents, parse_model_ref, configure_logging
+from deepresearch_flow.paper.db import register_db_commands
+from deepresearch_flow.paper.schema import load_schema, validate_schema, SchemaError
+from deepresearch_flow.paper.template_registry import list_template_names, load_schema_for_template
+
+
+@click.group()
+def paper() -> None:
+    """Paper extraction and database commands."""
+
+
+@paper.command()
+@click.option("-c", "--config", "config_path", default="config.toml", help="Path to config.toml")
+@click.option(
+    "-i",
+    "--input",
+    "inputs",
+    multiple=True,
+    required=True,
+    help="Input markdown file or directory (repeatable)",
+)
+@click.option("-g", "--glob", "glob_pattern", default=None, help="Glob filter when input is a directory")
+@click.option(
+    "-s",
+    "--schema-json",
+    "--schema",
+    "schema_path",
+    default=None,
+    help="Path to JSON schema",
+)
+@click.option("--prompt-system", "prompt_system", default=None, help="Custom system prompt template path")
+@click.option("--prompt-user", "prompt_user", default=None, help="Custom user prompt template path")
+@click.option(
+    "--template-dir",
+    "template_dir",
+    default=None,
+    help="Directory containing system.j2, user.j2, schema.json, render.j2",
+)
+@click.option(
+    "--prompt-template",
+    "prompt_template",
+    default="simple",
+    type=click.Choice(list_template_names()),
+    show_default=True,
+    help="Built-in prompt template",
+)
+@click.option(
+    "--language",
+    "output_language",
+    default="en",
+    show_default=True,
+    help="Output language hint for prompts",
+)
+@click.option("-m", "--model", "model_ref", required=True, help="provider/model")
+@click.option("-o", "--output", "output_path", default=None, help="Aggregated JSON output path")
+@click.option("-e", "--errors", "errors_path", default=None, help="Error JSON output path")
+@click.option("--split", is_flag=True, help="Write per-document JSON outputs")
+@click.option("--split-dir", "split_dir", default=None, help="Directory for split outputs")
+@click.option("--force", is_flag=True, help="Force re-extraction")
+@click.option("--retry-failed", is_flag=True, help="Retry only failed documents")
+@click.option("--dry-run", is_flag=True, help="Discover inputs without calling providers")
+@click.option("--max-concurrency", "max_concurrency", type=int, default=None, help="Override max concurrency")
+@click.option("--sleep-every", "sleep_every", type=int, default=None, help="Sleep after every N requests")
+@click.option("--sleep-time", "sleep_time", type=float, default=None, help="Sleep duration in seconds")
+@click.option("--render-md", "render_md", is_flag=True, help="Render markdown outputs after extraction")
+@click.option(
+    "--render-output-dir",
+    "render_output_dir",
+    default=None,
+    help="Output directory for rendered markdown (defaults to --output parent when provided)",
+)
+@click.option(
+    "--render-markdown-template",
+    "--render-template",
+    "render_template_path",
+    default=None,
+    help="Jinja2 template path for extract-time rendering",
+)
+@click.option(
+    "--render-template-name",
+    "render_template_name",
+    default=None,
+    type=click.Choice(list_template_names()),
+    help="Built-in render template name",
+)
+@click.option(
+    "--render-template-dir",
+    "render_template_dir",
+    default=None,
+    help="Directory containing render.j2 for extract-time rendering",
+)
+@click.option("-v", "--verbose", is_flag=True, help="Enable verbose logging")
+def extract(
+    config_path: str,
+    inputs: tuple[str, ...],
+    glob_pattern: str | None,
+    schema_path: str | None,
+    prompt_template: str,
+    output_language: str,
+    prompt_system: str | None,
+    prompt_user: str | None,
+    template_dir: str | None,
+    model_ref: str,
+    output_path: str | None,
+    errors_path: str | None,
+    split: bool,
+    split_dir: str | None,
+    force: bool,
+    retry_failed: bool,
+    dry_run: bool,
+    max_concurrency: int | None,
+    sleep_every: int | None,
+    sleep_time: float | None,
+    render_md: bool,
+    render_output_dir: str,
+    render_template_path: str | None,
+    render_template_name: str | None,
+    render_template_dir: str | None,
+    verbose: bool,
+) -> None:
+    """Extract structured information from markdown documents."""
+    config = load_config(config_path)
+    provider, model_name = parse_model_ref(model_ref, config.providers)
+
+    if provider.structured_mode not in {"json_schema", "json_object", "none"}:
+        raise click.ClickException("structured_mode must be json_schema, json_object, or none")
+
+    if config.extract.truncate_strategy not in {"head", "head_tail"}:
+        raise click.ClickException("truncate_strategy must be head or head_tail")
+
+    if config.extract.max_concurrency <= 0:
+        raise click.ClickException("max_concurrency must be positive")
+    if config.extract.max_retries <= 0:
+        raise click.ClickException("max_retries must be positive")
+    if max_concurrency is not None and max_concurrency <= 0:
+        raise click.ClickException("--max-concurrency must be positive")
+    if sleep_every is not None and sleep_every <= 0:
+        raise click.ClickException("--sleep-every must be positive")
+    if sleep_time is not None and sleep_time <= 0:
+        raise click.ClickException("--sleep-time must be positive")
+    if (sleep_every is None) != (sleep_time is None):
+        raise click.ClickException("Both --sleep-every and --sleep-time are required")
+
+    if provider.type in {
+        "openai_compatible",
+        "dashscope",
+        "gemini_ai_studio",
+        "azure_openai",
+        "claude",
+    }:
+        resolved = resolve_api_keys(provider.api_keys)
+        if not resolved:
+            raise click.ClickException(f"{provider.type} providers require api_keys")
+
+    if template_dir and (prompt_system or prompt_user or schema_path):
+        raise click.ClickException("template-dir cannot be combined with custom prompt or schema flags")
+
+    if (prompt_system and not prompt_user) or (prompt_user and not prompt_system):
+        raise click.ClickException("Both --prompt-system and --prompt-user are required")
+
+    custom_prompt = bool(prompt_system or prompt_user or template_dir)
+    if custom_prompt and prompt_template != "simple":
+        raise click.ClickException("Custom prompts cannot be combined with built-in prompt templates")
+
+    schema_override = schema_path or None
+    prompt_system_path = Path(prompt_system) if prompt_system else None
+    prompt_user_path = Path(prompt_user) if prompt_user else None
+    template_dir_path = Path(template_dir) if template_dir else None
+    if template_dir_path:
+        prompt_system_path = template_dir_path / "system.j2"
+        prompt_user_path = template_dir_path / "user.j2"
+        schema_override = str(template_dir_path / "schema.json")
+
+    for prompt_path in (prompt_system_path, prompt_user_path):
+        if prompt_path and not prompt_path.exists():
+            raise click.ClickException(f"Prompt template not found: {prompt_path}")
+
+    if not render_md and any(
+        item is not None
+        for item in (render_template_path, render_template_name, render_template_dir)
+    ):
+        raise click.ClickException("Render template options require --render-md")
+    if not render_md and render_output_dir is not None:
+        raise click.ClickException("--render-output-dir requires --render-md")
+    if render_md and sum(
+        bool(item) for item in (render_template_path, render_template_name, render_template_dir)
+    ) > 1:
+        raise click.ClickException(
+            "Use only one of --render-markdown-template/--render-template, --render-template-name, or --render-template-dir"
+        )
+    render_template_path_effective = render_template_path
+    render_template_name_effective = render_template_name
+    render_template_dir_effective = render_template_dir
+    render_output_dir_effective: Path | None = None
+
+    if render_md and not any(
+        item is not None
+        for item in (render_template_path, render_template_name, render_template_dir)
+    ):
+        if template_dir:
+            render_template_dir_effective = template_dir
+        elif not custom_prompt:
+            render_template_name_effective = prompt_template
+    if render_md:
+        if render_output_dir is not None:
+            render_output_dir_effective = Path(render_output_dir)
+        elif output_path is not None:
+            render_output_dir_effective = Path(output_path).parent
+        else:
+            render_output_dir_effective = Path("rendered_md")
+
+    if render_template_path_effective and not Path(render_template_path_effective).exists():
+        raise click.ClickException(f"Render template not found: {render_template_path_effective}")
+    if render_template_dir_effective:
+        render_template_dir_path = Path(render_template_dir_effective)
+        render_template_file = render_template_dir_path / "render.j2"
+        if not render_template_file.exists():
+            raise click.ClickException(f"Render template not found: {render_template_file}")
+
+    try:
+        if schema_override:
+            schema = load_schema(schema_override)
+        elif prompt_template:
+            schema = load_schema_for_template(prompt_template)
+        else:
+            schema = load_schema(config.extract.schema_path)
+        validator = validate_schema(schema)
+    except SchemaError as exc:
+        raise click.ClickException(str(exc)) from exc
+
+    output = Path(output_path or config.extract.output)
+    errors = Path(errors_path or config.extract.errors)
+    split_out = Path(split_dir) if split_dir else None
+
+    configure_logging(verbose)
+
+    asyncio.run(
+        extract_documents(
+            inputs=inputs,
+            glob_pattern=glob_pattern,
+            provider=provider,
+            model=model_name,
+            schema=schema,
+            validator=validator,
+            config=config,
+            output_path=output,
+            errors_path=errors,
+            split=split,
+            split_dir=split_out,
+            force=force,
+            retry_failed=retry_failed,
+            dry_run=dry_run,
+            max_concurrency_override=max_concurrency,
+            prompt_template=prompt_template,
+            output_language=output_language,
+            custom_prompt=custom_prompt,
+            prompt_system_path=prompt_system_path,
+            prompt_user_path=prompt_user_path,
+            render_md=render_md,
+            render_output_dir=render_output_dir_effective,
+            render_template_path=render_template_path_effective,
+            render_template_name=render_template_name_effective,
+            render_template_dir=render_template_dir_effective,
+            sleep_every=sleep_every,
+            sleep_time=sleep_time,
+            verbose=verbose,
+        )
+    )
+
+
+@paper.group()
+def db() -> None:
+    """Database management commands."""
+
+
+register_db_commands(db)

deepresearch_flow/paper/config.py

@@ -0,0 +1,249 @@
+"""Configuration loading and validation for paper tools."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+import os
+import tomllib
+
+
+@dataclass(frozen=True)
+class ExtractConfig:
+    output: str
+    errors: str
+    max_concurrency: int
+    max_retries: int
+    backoff_base_seconds: float
+    backoff_max_seconds: float
+    truncate_strategy: str
+    truncate_max_chars: int
+    cost_estimate: bool
+    schema_path: str | None
+
+
+@dataclass(frozen=True)
+class RenderConfig:
+    template_path: str | None
+
+
+@dataclass(frozen=True)
+class ProviderConfig:
+    name: str
+    type: str
+    base_url: str
+    api_keys: list[str]
+    api_version: str | None
+    deployment: str | None
+    project_id: str | None
+    location: str | None
+    credentials_path: str | None
+    anthropic_version: str | None
+    structured_mode: str
+    extra_headers: dict[str, str]
+    system_prompt: str | None
+    user_prompt: str | None
+    model_list: list[str]
+
+
+@dataclass(frozen=True)
+class PaperConfig:
+    extract: ExtractConfig
+    render: RenderConfig
+    providers: list[ProviderConfig]
+
+
+DEFAULT_EXTRACT = ExtractConfig(
+    output="paper_infos.json",
+    errors="paper_errors.json",
+    max_concurrency=6,
+    max_retries=3,
+    backoff_base_seconds=1.0,
+    backoff_max_seconds=20.0,
+    truncate_strategy="head_tail",
+    truncate_max_chars=20000,
+    cost_estimate=True,
+    schema_path=None,
+)
+
+DEFAULT_RENDER = RenderConfig(template_path=None)
+
+
+def _as_list(value: Any) -> list[str]:
+    if value is None:
+        return []
+    if isinstance(value, list):
+        return [str(item) for item in value]
+    return [str(value)]
+
+
+def _as_bool(value: Any, default: bool) -> bool:
+    if value is None:
+        return default
+    return bool(value)
+
+
+def _as_int(value: Any, default: int) -> int:
+    if value is None:
+        return default
+    return int(value)
+
+
+def _as_float(value: Any, default: float) -> float:
+    if value is None:
+        return default
+    return float(value)
+
+
+def _as_str(value: Any, default: str | None = None) -> str | None:
+    if value is None:
+        return default
+    return str(value)
+
+
+def load_config(path: str) -> PaperConfig:
+    config_path = Path(path)
+    if not config_path.exists():
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+
+    data = tomllib.loads(config_path.read_text(encoding="utf-8"))
+
+    extract_data = data.get("extract", {})
+    extract = ExtractConfig(
+        output=_as_str(extract_data.get("output"), DEFAULT_EXTRACT.output) or DEFAULT_EXTRACT.output,
+        errors=_as_str(extract_data.get("errors"), DEFAULT_EXTRACT.errors) or DEFAULT_EXTRACT.errors,
+        max_concurrency=_as_int(extract_data.get("max_concurrency"), DEFAULT_EXTRACT.max_concurrency),
+        max_retries=_as_int(extract_data.get("max_retries"), DEFAULT_EXTRACT.max_retries),
+        backoff_base_seconds=_as_float(
+            extract_data.get("backoff_base_seconds"), DEFAULT_EXTRACT.backoff_base_seconds
+        ),
+        backoff_max_seconds=_as_float(
+            extract_data.get("backoff_max_seconds"), DEFAULT_EXTRACT.backoff_max_seconds
+        ),
+        truncate_strategy=_as_str(
+            extract_data.get("truncate_strategy"), DEFAULT_EXTRACT.truncate_strategy
+        )
+        or DEFAULT_EXTRACT.truncate_strategy,
+        truncate_max_chars=_as_int(
+            extract_data.get("truncate_max_chars"), DEFAULT_EXTRACT.truncate_max_chars
+        ),
+        cost_estimate=_as_bool(extract_data.get("cost_estimate"), DEFAULT_EXTRACT.cost_estimate),
+        schema_path=_as_str(extract_data.get("schema_path"), DEFAULT_EXTRACT.schema_path),
+    )
+
+    render_data = data.get("render", {})
+    render = RenderConfig(template_path=_as_str(render_data.get("template_path"), DEFAULT_RENDER.template_path))
+
+    providers_data = data.get("providers", [])
+    providers: list[ProviderConfig] = []
+    for provider in providers_data:
+        name = _as_str(provider.get("name"))
+        provider_type = _as_str(provider.get("type"))
+        if not name or not provider_type:
+            raise ValueError("Each provider must include name and type")
+
+        base_url = _as_str(provider.get("base_url"))
+        endpoint = _as_str(provider.get("endpoint"))
+        if not base_url:
+            if provider_type == "ollama":
+                base_url = "http://localhost:11434"
+            elif provider_type == "openai_compatible":
+                base_url = "https://api.openai.com/v1"
+            elif provider_type == "azure_openai" and endpoint:
+                base_url = endpoint
+            elif provider_type in {"dashscope", "gemini_ai_studio", "gemini_vertex", "claude"}:
+                base_url = ""
+            else:
+                raise ValueError(f"Provider '{name}' requires base_url")
+        elif provider_type == "azure_openai" and endpoint:
+            base_url = endpoint
+
+        api_keys = _as_list(provider.get("api_keys"))
+        if not api_keys:
+            api_key_single = provider.get("api_key")
+            api_keys = _as_list(api_key_single)
+
+        structured_mode = _as_str(provider.get("structured_mode"), None)
+        if structured_mode is None:
+            if provider_type == "ollama":
+                structured_mode = "json_object"
+            elif provider_type in {"dashscope", "gemini_ai_studio", "gemini_vertex", "claude"}:
+                structured_mode = "none"
+            else:
+                structured_mode = "json_schema"
+
+        extra_headers: dict[str, str] = {}
+        headers = provider.get("extra_headers")
+        if isinstance(headers, dict):
+            extra_headers = {str(k): str(v) for k, v in headers.items()}
+
+        model_list = _as_list(provider.get("model_list"))
+        if not model_list:
+            raise ValueError(f"Provider '{name}' must include model_list")
+
+        api_version = _as_str(provider.get("api_version"), None)
+        deployment = _as_str(provider.get("deployment"), None)
+        project_id = _as_str(provider.get("project_id"), None)
+        location = _as_str(provider.get("location"), None)
+        credentials_path = _as_str(provider.get("credentials_path"), None)
+        anthropic_version = _as_str(provider.get("anthropic_version"), None)
+
+        if provider_type == "azure_openai":
+            if not base_url:
+                raise ValueError(f"Provider '{name}' requires endpoint")
+            if not api_version:
+                raise ValueError(f"Provider '{name}' requires api_version")
+            if not deployment:
+                raise ValueError(f"Provider '{name}' requires deployment")
+        if provider_type == "gemini_ai_studio" and not api_keys:
+            raise ValueError(f"Provider '{name}' requires api_keys")
+        if provider_type == "gemini_vertex":
+            if not project_id:
+                raise ValueError(f"Provider '{name}' requires project_id")
+            if not location:
+                raise ValueError(f"Provider '{name}' requires location")
+        if provider_type == "claude":
+            if not api_keys:
+                raise ValueError(f"Provider '{name}' requires api_keys")
+            if not anthropic_version:
+                raise ValueError(f"Provider '{name}' requires anthropic_version")
+
+        providers.append(
+            ProviderConfig(
+                name=name,
+                type=provider_type,
+                base_url=base_url,
+                api_keys=api_keys,
+                api_version=api_version,
+                deployment=deployment,
+                project_id=project_id,
+                location=location,
+                credentials_path=credentials_path,
+                anthropic_version=anthropic_version,
+                structured_mode=structured_mode,
+                extra_headers=extra_headers,
+                system_prompt=_as_str(provider.get("system_prompt"), None),
+                user_prompt=_as_str(provider.get("user_prompt"), None),
+                model_list=model_list,
+            )
+        )
+
+    if not providers:
+        raise ValueError("Config must include at least one [[providers]] entry")
+
+    return PaperConfig(extract=extract, render=render, providers=providers)
+
+
+def resolve_api_keys(entries: list[str]) -> list[str]:
+    resolved: list[str] = []
+    for entry in entries:
+        entry = str(entry)
+        if entry.startswith("env:"):
+            env_name = entry.split(":", 1)[1]
+            value = os.environ.get(env_name)
+            if value:
+                resolved.append(value)
+        else:
+            resolved.append(entry)
+    return resolved