deepresearch-flow 0.5.1__py3-none-any.whl → 0.6.1__py3-none-any.whl

deepresearch_flow/paper/cli.py

@@ -67,9 +67,40 @@ def paper() -> None:
 @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(
+    "--force-stage",
+    "force_stages",
+    multiple=True,
+    help="Force re-run specific stages (multi-stage templates only)",
+)
 @click.option("--retry-failed", is_flag=True, help="Retry only failed documents")
+@click.option(
+    "--retry-failed-stages",
+    is_flag=True,
+    help="Retry only failed stages per document (multi-stage templates only)",
+)
+@click.option(
+    "--retry-list-json",
+    "retry_list_json",
+    default=None,
+    help="Retry only documents listed in a verification report",
+)
+@click.option(
+    "--stage-dag",
+    is_flag=True,
+    help="Enable dependency-aware DAG scheduling (multi-stage templates only)",
+)
+@click.option("--start-idx", "start_idx", type=int, default=0, help="Start index for inputs")
+@click.option(
+    "--end-idx",
+    "end_idx",
+    type=int,
+    default=-1,
+    help="End index (exclusive); -1 means to the last item",
+)
 @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("--timeout", "timeout_seconds", type=float, default=None, help="Request timeout in seconds")
 @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")
@@ -116,9 +147,16 @@ def extract(
     split: bool,
     split_dir: str | None,
     force: bool,
+    force_stages: tuple[str, ...],
     retry_failed: bool,
+    retry_failed_stages: bool,
+    retry_list_json: str | None,
+    stage_dag: bool,
+    start_idx: int,
+    end_idx: int,
     dry_run: bool,
     max_concurrency: int | None,
+    timeout_seconds: float | None,
     sleep_every: int | None,
     sleep_time: float | None,
     render_md: bool,
@@ -142,14 +180,28 @@ def extract(
         raise click.ClickException("max_concurrency must be positive")
     if config.extract.max_retries <= 0:
         raise click.ClickException("max_retries must be positive")
+    if config.extract.timeout <= 0:
+        raise click.ClickException("timeout must be positive")
     if max_concurrency is not None and max_concurrency <= 0:
         raise click.ClickException("--max-concurrency must be positive")
+    if timeout_seconds is not None and timeout_seconds <= 0:
+        raise click.ClickException("--timeout 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 start_idx < 0:
+        raise click.ClickException("--start-idx must be >= 0")
+    if end_idx < -1:
+        raise click.ClickException("--end-idx must be -1 or >= 0")
+    if retry_failed and retry_failed_stages:
+        raise click.ClickException("--retry-failed and --retry-failed-stages are mutually exclusive")
+    if retry_list_json and (retry_failed or retry_failed_stages):
+        raise click.ClickException(
+            "--retry-list-json cannot be combined with --retry-failed or --retry-failed-stages"
+        )
 
     if provider.type in {
         "openai_compatible",
@@ -171,6 +223,8 @@ def extract(
     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")
+    if stage_dag and custom_prompt:
+        raise click.ClickException("--stage-dag requires a built-in multi-stage prompt template")
 
     schema_override = schema_path or None
     prompt_system_path = Path(prompt_system) if prompt_system else None
@@ -240,7 +294,9 @@ def extract(
 
     output = Path(output_path or config.extract.output)
     errors = Path(errors_path or config.extract.errors)
+    retry_list_path = Path(retry_list_json) if retry_list_json else None
     split_out = Path(split_dir) if split_dir else None
+    timeout_seconds_effective = timeout_seconds if timeout_seconds is not None else config.extract.timeout
 
     configure_logging(verbose)
 
@@ -258,9 +314,16 @@ def extract(
             split=split,
             split_dir=split_out,
             force=force,
+            force_stages=list(force_stages),
             retry_failed=retry_failed,
+            retry_failed_stages=retry_failed_stages,
+            retry_list_path=retry_list_path,
+            stage_dag=stage_dag or config.extract.stage_dag,
+            start_idx=start_idx,
+            end_idx=end_idx,
             dry_run=dry_run,
             max_concurrency_override=max_concurrency,
+            timeout_seconds=timeout_seconds_effective,
             prompt_template=prompt_template,
             output_language=output_language,
             custom_prompt=custom_prompt,

deepresearch_flow/paper/config.py

@@ -15,12 +15,15 @@ class ExtractConfig:
     errors: str
     max_concurrency: int
     max_retries: int
+    timeout: float
     backoff_base_seconds: float
     backoff_max_seconds: float
+    pause_threshold_seconds: float
     truncate_strategy: str
     truncate_max_chars: int
     cost_estimate: bool
     schema_path: str | None
+    stage_dag: bool
 
 
 @dataclass(frozen=True)
@@ -28,12 +31,20 @@ class RenderConfig:
     template_path: str | None
 
 
+@dataclass(frozen=True)
+class ApiKeyConfig:
+    key: str
+    quota_duration: int | None
+    reset_time: str | None
+    quota_error_tokens: list[str]
+
+
 @dataclass(frozen=True)
 class ProviderConfig:
     name: str
     type: str
     base_url: str
-    api_keys: list[str]
+    api_keys: list[ApiKeyConfig]
     api_version: str | None
     deployment: str | None
     project_id: str | None
@@ -60,12 +71,15 @@ DEFAULT_EXTRACT = ExtractConfig(
     errors="paper_errors.json",
     max_concurrency=6,
     max_retries=3,
+    timeout=60.0,
     backoff_base_seconds=1.0,
     backoff_max_seconds=20.0,
+    pause_threshold_seconds=10.0,
     truncate_strategy="head_tail",
     truncate_max_chars=20000,
     cost_estimate=True,
     schema_path=None,
+    stage_dag=False,
 )
 
 DEFAULT_RENDER = RenderConfig(template_path=None)
@@ -103,6 +117,51 @@ def _as_str(value: Any, default: str | None = None) -> str | None:
     return str(value)
 
 
+def _parse_api_keys(value: Any) -> list[ApiKeyConfig]:
+    if value is None:
+        return []
+    entries = value if isinstance(value, list) else [value]
+    parsed: list[ApiKeyConfig] = []
+    for entry in entries:
+        if isinstance(entry, dict):
+            key = _as_str(entry.get("key"))
+            if not key:
+                raise ValueError("api_keys object entries must include key")
+            quota_duration = entry.get("quota_duration")
+            quota_duration_value = int(quota_duration) if quota_duration is not None else None
+            if quota_duration_value is not None and quota_duration_value <= 0:
+                raise ValueError("quota_duration must be positive seconds")
+            reset_time = _as_str(entry.get("reset_time"), None)
+            tokens = entry.get("quota_error_tokens")
+            if tokens is None:
+                quota_error_tokens = []
+            elif isinstance(tokens, list):
+                quota_error_tokens = [str(token) for token in tokens]
+            else:
+                quota_error_tokens = [str(tokens)]
+            parsed.append(
+                ApiKeyConfig(
+                    key=key,
+                    quota_duration=quota_duration_value,
+                    reset_time=reset_time,
+                    quota_error_tokens=quota_error_tokens,
+                )
+            )
+        else:
+            key = _as_str(entry)
+            if not key:
+                continue
+            parsed.append(
+                ApiKeyConfig(
+                    key=key,
+                    quota_duration=None,
+                    reset_time=None,
+                    quota_error_tokens=[],
+                )
+            )
+    return parsed
+
+
 def _ensure_http_scheme(base_url: str, *, default_scheme: str = "http://") -> str:
     normalized = base_url.strip()
     if normalized.startswith(("http://", "https://")):
@@ -125,12 +184,16 @@ def load_config(path: str) -> PaperConfig:
         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),
+        timeout=_as_float(extract_data.get("timeout"), DEFAULT_EXTRACT.timeout),
         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
         ),
+        pause_threshold_seconds=_as_float(
+            extract_data.get("pause_threshold_seconds"), DEFAULT_EXTRACT.pause_threshold_seconds
+        ),
         truncate_strategy=_as_str(
             extract_data.get("truncate_strategy"), DEFAULT_EXTRACT.truncate_strategy
         )
@@ -140,6 +203,7 @@ def load_config(path: str) -> PaperConfig:
         ),
         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),
+        stage_dag=_as_bool(extract_data.get("stage_dag"), DEFAULT_EXTRACT.stage_dag),
     )
 
     render_data = data.get("render", {})
@@ -171,10 +235,10 @@ def load_config(path: str) -> PaperConfig:
         if provider_type == "ollama" and base_url:
             base_url = _ensure_http_scheme(base_url)
 
-        api_keys = _as_list(provider.get("api_keys"))
+        api_keys = _parse_api_keys(provider.get("api_keys"))
         if not api_keys:
             api_key_single = provider.get("api_key")
-            api_keys = _as_list(api_key_single)
+            api_keys = _parse_api_keys(api_key_single)
 
         structured_mode = _as_str(provider.get("structured_mode"), None)
         if structured_mode is None:
@@ -250,15 +314,26 @@ def load_config(path: str) -> PaperConfig:
     return PaperConfig(extract=extract, render=render, providers=providers)
 
 
-def resolve_api_keys(entries: list[str]) -> list[str]:
-    resolved: list[str] = []
+def resolve_api_key_configs(entries: list[ApiKeyConfig]) -> list[ApiKeyConfig]:
+    resolved: list[ApiKeyConfig] = []
     for entry in entries:
-        entry = str(entry)
-        if entry.startswith("env:"):
-            env_name = entry.split(":", 1)[1]
+        key = entry.key
+        if key.startswith("env:"):
+            env_name = key.split(":", 1)[1]
             value = os.environ.get(env_name)
-            if value:
-                resolved.append(value)
-        else:
-            resolved.append(entry)
+            if not value:
+                continue
+            key = value
+        resolved.append(
+            ApiKeyConfig(
+                key=key,
+                quota_duration=entry.quota_duration,
+                reset_time=entry.reset_time,
+                quota_error_tokens=entry.quota_error_tokens,
+            )
+        )
     return resolved
+
+
+def resolve_api_keys(entries: list[ApiKeyConfig]) -> list[str]:
+    return [entry.key for entry in resolve_api_key_configs(entries)]