From 8ad3651edd12658614f24ea636805e954fa72dae Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 13:24:37 -0500 Subject: [PATCH 01/10] fix: lazy import integrations to avoid pyspark bundling error The CLI binary was failing with "No such file or directory: error-conditions.json" because pyspark was being eagerly imported when any integration was accessed. Changes: - Make integrations/__init__.py use lazy imports via __getattr__ - Update PyInstaller spec to use correct integration paths - Remove library-only integrations (polars, duckdb, spark) from bundle - Only bundle integrations with deploy capability (bigquery, snowflake) --- parallel-web-tools.spec | 23 ++++----- parallel_web_tools/integrations/__init__.py | 52 +++++++++++++++++---- 2 files changed, 52 insertions(+), 23 deletions(-) diff --git a/parallel-web-tools.spec b/parallel-web-tools.spec index ae0fc52..55950c3 100644 --- a/parallel-web-tools.spec +++ b/parallel-web-tools.spec @@ -36,24 +36,18 @@ a = Analysis( 'parallel_web_tools.cli', 'parallel_web_tools.cli.commands', 'parallel_web_tools.cli.planner', - # Processors + # Processors (for local file/db enrichment) 'parallel_web_tools.processors', 'parallel_web_tools.processors.csv', 'parallel_web_tools.processors.duckdb', 'parallel_web_tools.processors.bigquery', - # Integrations - 'parallel_web_tools.polars', - 'parallel_web_tools.polars.enrich', - 'parallel_web_tools.duckdb', - 'parallel_web_tools.duckdb.batch', - 'parallel_web_tools.duckdb.udf', - 'parallel_web_tools.bigquery', - 'parallel_web_tools.bigquery.deploy', - 'parallel_web_tools.snowflake', - 'parallel_web_tools.snowflake.deploy', - 'parallel_web_tools.spark', - 'parallel_web_tools.spark.udf', - 'parallel_web_tools.spark.streaming', + # Integrations with deploy capability (bundled in CLI) + 'parallel_web_tools.integrations', + 'parallel_web_tools.integrations.bigquery', + 'parallel_web_tools.integrations.bigquery.deploy', + 'parallel_web_tools.integrations.snowflake', + 'parallel_web_tools.integrations.snowflake.deploy', + # Note: polars, duckdb, spark integrations are library-only (no deploy step) # Dependencies that might not be auto-detected 'click', 'questionary', @@ -65,7 +59,6 @@ a = Analysis( 'parallel', 'duckdb', 'sqlalchemy', - 'polars', 'pandas', 'certifi', ], diff --git a/parallel_web_tools/integrations/__init__.py b/parallel_web_tools/integrations/__init__.py index 6aae97c..364cf77 100644 --- a/parallel_web_tools/integrations/__init__.py +++ b/parallel_web_tools/integrations/__init__.py @@ -1,12 +1,31 @@ -"""Database and data platform integrations for parallel-web-tools.""" +"""Database and data platform integrations for parallel-web-tools. -from parallel_web_tools.integrations import ( - bigquery, - duckdb, - polars, - snowflake, - spark, -) +Submodules are lazily imported to avoid loading heavy dependencies (like pyspark) +when they're not needed. Import the specific integration you need: + + from parallel_web_tools.integrations import bigquery + from parallel_web_tools.integrations import spark # Only loads pyspark here +""" + +import importlib +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from parallel_web_tools.integrations import ( + bigquery as bigquery, + ) + from parallel_web_tools.integrations import ( + duckdb as duckdb, + ) + from parallel_web_tools.integrations import ( + polars as polars, + ) + from parallel_web_tools.integrations import ( + snowflake as snowflake, + ) + from parallel_web_tools.integrations import ( + spark as spark, + ) __all__ = [ "bigquery", @@ -15,3 +34,20 @@ "snowflake", "spark", ] + +_SUBMODULES = { + "bigquery": "parallel_web_tools.integrations.bigquery", + "duckdb": "parallel_web_tools.integrations.duckdb", + "polars": "parallel_web_tools.integrations.polars", + "snowflake": "parallel_web_tools.integrations.snowflake", + "spark": "parallel_web_tools.integrations.spark", +} + + +def __getattr__(name: str): + """Lazily import submodules on first access.""" + if name in _SUBMODULES: + module = importlib.import_module(_SUBMODULES[name]) + globals()[name] = module + return module + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From 6cb6f7ea6b5fa43ab6bdfe6534c4d508141a347b Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 13:32:16 -0500 Subject: [PATCH 02/10] feat: add --output option to search and extract commands --- parallel_web_tools/cli/commands.py | 78 ++++++++++++++++++------------ 1 file changed, 47 insertions(+), 31 deletions(-) diff --git a/parallel_web_tools/cli/commands.py b/parallel_web_tools/cli/commands.py index 7ea30e3..6da4662 100644 --- a/parallel_web_tools/cli/commands.py +++ b/parallel_web_tools/cli/commands.py @@ -232,6 +232,7 @@ def logout_cmd(): @click.option("--include-domains", multiple=True, help="Only search these domains") @click.option("--exclude-domains", multiple=True, help="Exclude these domains") @click.option("--after-date", help="Only results after this date (YYYY-MM-DD)") +@click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to file (JSON)") @click.option("--json", "output_json", is_flag=True, help="Output as JSON") def search( objective: str | None, @@ -241,6 +242,7 @@ def search( include_domains: tuple[str, ...], exclude_domains: tuple[str, ...], after_date: str | None, + output_file: str | None, output_json: bool, ): """Search the web using Parallel's AI-powered search.""" @@ -275,16 +277,22 @@ def search( result = client.beta.search(**search_kwargs) + output_data = { + "search_id": result.search_id, + "results": [ + {"url": r.url, "title": r.title, "publish_date": r.publish_date, "excerpts": r.excerpts} + for r in result.results + ], + "warnings": result.warnings if hasattr(result, "warnings") else [], + } + + if output_file: + with open(output_file, "w") as f: + json.dump(output_data, f, indent=2) + console.print(f"[dim]Results saved to {output_file}[/dim]\n") + if output_json: - output = { - "search_id": result.search_id, - "results": [ - {"url": r.url, "title": r.title, "publish_date": r.publish_date, "excerpts": r.excerpts} - for r in result.results - ], - "warnings": result.warnings if hasattr(result, "warnings") else [], - } - print(json.dumps(output, indent=2)) + print(json.dumps(output_data, indent=2)) else: console.print(f"[bold green]Found {len(result.results)} results[/bold green]\n") for i, r in enumerate(result.results, 1): @@ -313,6 +321,7 @@ def search( @click.option("-q", "--query", multiple=True, help="Keywords to prioritize (can be repeated)") @click.option("--full-content", is_flag=True, help="Include complete page content") @click.option("--no-excerpts", is_flag=True, help="Exclude excerpts from output") +@click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to file (JSON)") @click.option("--json", "output_json", is_flag=True, help="Output as JSON") def extract( urls: tuple[str, ...], @@ -320,6 +329,7 @@ def extract( query: tuple[str, ...], full_content: bool, no_excerpts: bool, + output_file: str | None, output_json: bool, ): """Extract content from URLs as clean markdown.""" @@ -346,29 +356,35 @@ def extract( result = client.beta.extract(**extract_kwargs) + results_list = [] + for r in result.results: + result_dict: dict[str, Any] = {"url": r.url, "title": r.title} + if hasattr(r, "excerpts") and r.excerpts: + result_dict["excerpts"] = r.excerpts + if hasattr(r, "full_content") and r.full_content: + result_dict["full_content"] = r.full_content + results_list.append(result_dict) + + errors_list = [] + if hasattr(result, "errors") and result.errors: + for e in result.errors: + errors_list.append( + { + "url": getattr(e, "url", None), + "error": str(getattr(e, "error", "")), + "status_code": getattr(e, "status_code", None), + } + ) + + output_data = {"extract_id": result.extract_id, "results": results_list, "errors": errors_list} + + if output_file: + with open(output_file, "w") as f: + json.dump(output_data, f, indent=2) + console.print(f"[dim]Results saved to {output_file}[/dim]\n") + if output_json: - results_list = [] - for r in result.results: - result_dict: dict[str, Any] = {"url": r.url, "title": r.title} - if hasattr(r, "excerpts") and r.excerpts: - result_dict["excerpts"] = r.excerpts - if hasattr(r, "full_content") and r.full_content: - result_dict["full_content"] = r.full_content - results_list.append(result_dict) - - errors_list = [] - if hasattr(result, "errors") and result.errors: - for e in result.errors: - errors_list.append( - { - "url": getattr(e, "url", None), - "error": str(getattr(e, "error", "")), - "status_code": getattr(e, "status_code", None), - } - ) - - output = {"extract_id": result.extract_id, "results": results_list, "errors": errors_list} - print(json.dumps(output, indent=2)) + print(json.dumps(output_data, indent=2)) else: if result.errors: console.print(f"[yellow]Warning: {len(result.errors)} URL(s) failed[/yellow]\n") From ead1eeec30c1ed308c607136fef64618cf6bd86f Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 13:46:44 -0500 Subject: [PATCH 03/10] feat: add deep research command to CLI Add `parallel-cli research` command group for running deep research tasks: - `research run` - Create and run research tasks with polling - `research status` - Check task status - `research poll` - Resume polling existing tasks - `research processors` - List available processor tiers Also includes: - New core/research.py module with research API functions - Comprehensive tests with mocking (36 tests) - Refactored shared polling logic to reduce duplication - Added write_json_output() helper for consistent output handling --- parallel_web_tools/cli/commands.py | 301 ++++++++++++++- parallel_web_tools/core/__init__.py | 15 + parallel_web_tools/core/research.py | 303 +++++++++++++++ tests/test_research.py | 553 ++++++++++++++++++++++++++++ 4 files changed, 1158 insertions(+), 14 deletions(-) create mode 100644 parallel_web_tools/core/research.py create mode 100644 tests/test_research.py diff --git a/parallel_web_tools/cli/commands.py b/parallel_web_tools/cli/commands.py index 6da4662..016e9a1 100644 --- a/parallel_web_tools/cli/commands.py +++ b/parallel_web_tools/cli/commands.py @@ -15,11 +15,16 @@ from parallel_web_tools.core import ( AVAILABLE_PROCESSORS, JSON_SCHEMA_TYPE_MAP, + RESEARCH_PROCESSORS, + create_research_task, get_api_key, get_auth_status, + get_research_status, logout, + poll_research, run_enrichment, run_enrichment_from_dict, + run_research, ) logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO) @@ -29,6 +34,28 @@ load_dotenv(".env.local") +# ============================================================================= +# Output Helpers +# ============================================================================= + + +def write_json_output(data: dict[str, Any], output_file: str | None, output_json: bool) -> None: + """Write output data to file and/or stdout as JSON. + + Args: + data: The data dictionary to output. + output_file: Optional file path to save JSON to. + output_json: If True, print JSON to stdout. + """ + if output_file: + with open(output_file, "w") as f: + json.dump(data, f, indent=2) + console.print(f"[dim]Results saved to {output_file}[/dim]\n") + + if output_json: + print(json.dumps(data, indent=2)) + + def parse_columns(columns_json: str | None) -> list[dict[str, str]] | None: """Parse columns from JSON string.""" if not columns_json: @@ -286,14 +313,9 @@ def search( "warnings": result.warnings if hasattr(result, "warnings") else [], } - if output_file: - with open(output_file, "w") as f: - json.dump(output_data, f, indent=2) - console.print(f"[dim]Results saved to {output_file}[/dim]\n") + write_json_output(output_data, output_file, output_json) - if output_json: - print(json.dumps(output_data, indent=2)) - else: + if not output_json: console.print(f"[bold green]Found {len(result.results)} results[/bold green]\n") for i, r in enumerate(result.results, 1): console.print(f"[bold cyan]{i}. {r.title}[/bold cyan]") @@ -378,14 +400,9 @@ def extract( output_data = {"extract_id": result.extract_id, "results": results_list, "errors": errors_list} - if output_file: - with open(output_file, "w") as f: - json.dump(output_data, f, indent=2) - console.print(f"[dim]Results saved to {output_file}[/dim]\n") + write_json_output(output_data, output_file, output_json) - if output_json: - print(json.dumps(output_data, indent=2)) - else: + if not output_json: if result.errors: console.print(f"[yellow]Warning: {len(result.errors)} URL(s) failed[/yellow]\n") @@ -638,5 +655,261 @@ def enrich_deploy(system: str, project: str | None, region: str, api_key: str | raise click.Abort() from None +# ============================================================================= +# Research Command Group +# ============================================================================= + + +@main.group() +def research(): + """Deep research commands for open-ended questions.""" + pass + + +@research.command(name="run") +@click.argument("query", required=False) +@click.option("--input-file", "-f", type=click.Path(exists=True), help="Read query from file") +@click.option( + "--processor", + "-p", + type=click.Choice(list(RESEARCH_PROCESSORS.keys())), + default="pro-fast", + show_default=True, + help="Processor tier (higher = more thorough but slower)", +) +@click.option("--timeout", type=int, default=3600, show_default=True, help="Max wait time in seconds") +@click.option("--poll-interval", type=int, default=45, show_default=True, help="Seconds between status checks") +@click.option("--no-wait", is_flag=True, help="Return immediately after creating task (don't poll)") +@click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to file (markdown)") +@click.option("--json", "output_json", is_flag=True, help="Output as JSON") +@click.option("--no-basis", is_flag=True, help="Exclude citations/sources from output") +def research_run( + query: str | None, + input_file: str | None, + processor: str, + timeout: int, + poll_interval: int, + no_wait: bool, + output_file: str | None, + output_json: bool, + no_basis: bool, +): + """Run deep research on a question or topic. + + QUERY is the research question (max 15,000 chars). Alternatively, use --input-file. + + Examples: + + parallel-cli research run "What are the latest developments in quantum computing?" + + parallel-cli research run -f question.txt --processor ultra -o report.md + """ + # Get query from argument or file + if input_file: + with open(input_file) as f: + query = f.read().strip() + elif not query: + console.print("[bold red]Error: Provide a query or use --input-file[/bold red]") + raise click.Abort() + + if len(query) > 15000: + console.print(f"[yellow]Warning: Query truncated from {len(query)} to 15,000 characters[/yellow]") + query = query[:15000] + + try: + if no_wait: + # Create task and return immediately + console.print(f"[dim]Creating research task with processor: {processor}...[/dim]") + result = create_research_task(query, processor=processor) + + console.print(f"\n[bold green]Task created: {result['run_id']}[/bold green]") + console.print(f"Track progress: {result['result_url']}") + console.print("\n[dim]Use 'parallel-cli research status ' to check status[/dim]") + console.print("[dim]Use 'parallel-cli research poll ' to wait for results[/dim]") + + if output_json: + print(json.dumps(result, indent=2)) + else: + # Run and wait for results + console.print(f"[bold cyan]Starting deep research with processor: {processor}[/bold cyan]") + console.print(f"[dim]This may take {RESEARCH_PROCESSORS[processor]}[/dim]\n") + + def on_status(status: str, run_id: str): + if status == "created": + console.print(f"[green]Task created: {run_id}[/green]") + console.print(f"[dim]Track progress: https://platform.parallel.ai/tasks/{run_id}[/dim]\n") + else: + console.print(f"[dim]Status: {status}[/dim]") + + result = run_research( + query, + processor=processor, + timeout=timeout, + poll_interval=poll_interval, + include_basis=not no_basis, + on_status=on_status, + ) + + _output_research_result(result, output_file, output_json, no_basis) + + except TimeoutError as e: + console.print(f"[bold yellow]Timeout: {e}[/bold yellow]") + console.print("[dim]The task is still running. Use 'parallel-cli research poll ' to resume.[/dim]") + raise click.Abort() from None + except RuntimeError as e: + console.print(f"[bold red]Error: {e}[/bold red]") + raise click.Abort() from None + except Exception as e: + console.print(f"[bold red]Error: {e}[/bold red]") + raise click.Abort() from None + + +@research.command(name="status") +@click.argument("run_id") +@click.option("--json", "output_json", is_flag=True, help="Output as JSON") +def research_status(run_id: str, output_json: bool): + """Check the status of a research task. + + RUN_ID is the task identifier (e.g., trun_xxx). + """ + try: + result = get_research_status(run_id) + + if output_json: + print(json.dumps(result, indent=2)) + else: + status = result["status"] + status_color = { + "completed": "green", + "running": "cyan", + "pending": "yellow", + "failed": "red", + "cancelled": "red", + }.get(status, "white") + + console.print(f"[bold]Task:[/bold] {run_id}") + console.print(f"[bold]Status:[/bold] [{status_color}]{status}[/{status_color}]") + console.print(f"[bold]URL:[/bold] {result['result_url']}") + + if status == "completed": + console.print("\n[dim]Use 'parallel-cli research poll ' to retrieve results[/dim]") + + except Exception as e: + console.print(f"[bold red]Error: {e}[/bold red]") + raise click.Abort() from None + + +@research.command(name="poll") +@click.argument("run_id") +@click.option("--timeout", type=int, default=3600, show_default=True, help="Max wait time in seconds") +@click.option("--poll-interval", type=int, default=45, show_default=True, help="Seconds between status checks") +@click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to file (markdown)") +@click.option("--json", "output_json", is_flag=True, help="Output as JSON") +@click.option("--no-basis", is_flag=True, help="Exclude citations/sources from output") +def research_poll( + run_id: str, + timeout: int, + poll_interval: int, + output_file: str | None, + output_json: bool, + no_basis: bool, +): + """Poll an existing research task until completion. + + RUN_ID is the task identifier (e.g., trun_xxx). + """ + try: + console.print(f"[bold cyan]Polling task: {run_id}[/bold cyan]") + console.print(f"[dim]Track progress: https://platform.parallel.ai/tasks/{run_id}[/dim]\n") + + def on_status(status: str, run_id: str): + console.print(f"[dim]Status: {status}[/dim]") + + result = poll_research( + run_id, + timeout=timeout, + poll_interval=poll_interval, + include_basis=not no_basis, + on_status=on_status, + ) + + _output_research_result(result, output_file, output_json, no_basis) + + except TimeoutError as e: + console.print(f"[bold yellow]Timeout: {e}[/bold yellow]") + raise click.Abort() from None + except RuntimeError as e: + console.print(f"[bold red]Error: {e}[/bold red]") + raise click.Abort() from None + except Exception as e: + console.print(f"[bold red]Error: {e}[/bold red]") + raise click.Abort() from None + + +@research.command(name="processors") +def research_processors(): + """List available research processors and their characteristics.""" + console.print("[bold]Available Research Processors:[/bold]\n") + for proc, desc in RESEARCH_PROCESSORS.items(): + console.print(f" [cyan]{proc:15}[/cyan] {desc}") + console.print("\n[dim]Use --processor/-p to select a processor[/dim]") + + +def _output_research_result( + result: dict, + output_file: str | None, + output_json: bool, + no_basis: bool, +): + """Output research result to console and/or file.""" + content = result.get("content", "") + basis = result.get("basis", []) + + # Build JSON output + output_data = { + "run_id": result.get("run_id"), + "result_url": result.get("result_url"), + "status": result.get("status"), + "content": content, + } + if not no_basis and basis: + output_data["basis"] = basis + + # Save to file if requested + if output_file: + # Write markdown content + with open(output_file, "w") as f: + f.write(content) + + # Write JSON metadata + json_file = output_file.rsplit(".", 1)[0] + ".json" if "." in output_file else output_file + ".json" + with open(json_file, "w") as f: + json.dump(output_data, f, indent=2) + + console.print("\n[dim]Results saved to:[/dim]") + console.print(f" [green]Markdown:[/green] {output_file}") + console.print(f" [green]JSON:[/green] {json_file}") + + # Output to console + if output_json: + print(json.dumps(output_data, indent=2)) + else: + console.print("\n[bold green]Research Complete![/bold green]") + console.print(f"[dim]Task: {result.get('run_id')}[/dim]\n") + + # Print content (truncate for console if very long) + if len(content) > 5000 and not output_file: + console.print(content[:5000]) + console.print(f"\n[yellow]... truncated ({len(content)} chars total)[/yellow]") + console.print("[dim]Use --output to save full content to a file[/dim]") + else: + console.print(content) + + # Show citation summary + if basis and not no_basis: + total_citations = sum(len(b.get("citations", [])) for b in basis if isinstance(b, dict)) + console.print(f"\n[dim]Sources: {total_citations} citations from {len(basis)} fields[/dim]") + + if __name__ == "__main__": main() diff --git a/parallel_web_tools/core/__init__.py b/parallel_web_tools/core/__init__.py index f558cc1..7bb846b 100644 --- a/parallel_web_tools/core/__init__.py +++ b/parallel_web_tools/core/__init__.py @@ -15,6 +15,14 @@ extract_basis, run_tasks, ) +from parallel_web_tools.core.research import ( + RESEARCH_PROCESSORS, + create_research_task, + get_research_result, + get_research_status, + poll_research, + run_research, +) from parallel_web_tools.core.result import EnrichmentResult from parallel_web_tools.core.runner import ( run_enrichment, @@ -65,6 +73,13 @@ # Runner "run_enrichment", "run_enrichment_from_dict", + # Research + "RESEARCH_PROCESSORS", + "create_research_task", + "get_research_result", + "get_research_status", + "poll_research", + "run_research", # Result "EnrichmentResult", ] diff --git a/parallel_web_tools/core/research.py b/parallel_web_tools/core/research.py new file mode 100644 index 0000000..16055d6 --- /dev/null +++ b/parallel_web_tools/core/research.py @@ -0,0 +1,303 @@ +"""Deep Research using the Parallel Task API. + +Deep research is designed for open-ended research questions that require +comprehensive multi-step web exploration. Unlike batch enrichment which +processes structured data, deep research takes a natural language query +and returns analyst-grade intelligence reports. +""" + +from __future__ import annotations + +import json +import time +from collections.abc import Callable +from typing import Any + +from parallel_web_tools.core.auth import resolve_api_key +from parallel_web_tools.core.batch import extract_basis + +# Processor tiers for deep research with expected latency +RESEARCH_PROCESSORS = { + "pro-fast": "1-5 min - exploratory research (default)", + "pro": "2-10 min - exploratory research, fresher data", + "ultra-fast": "2-12 min - multi-source deep research", + "ultra": "5-25 min - advanced deep research, fresher data", + "ultra2x-fast": "2-25 min - difficult deep research", + "ultra2x": "5-50 min - difficult deep research, fresher data", + "ultra4x-fast": "2-45 min - very difficult research", + "ultra4x": "5-90 min - very difficult research, fresher data", + "ultra8x-fast": "2-60 min - most challenging research", + "ultra8x": "5min-2hr - most challenging research, fresher data", +} + +TERMINAL_STATUSES = ("completed", "failed", "cancelled") + + +def create_research_task( + query: str, + processor: str = "pro-fast", + api_key: str | None = None, +) -> dict[str, Any]: + """Create a deep research task without waiting for results. + + Args: + query: Research question or topic (max 15,000 chars). + processor: Processor tier (see RESEARCH_PROCESSORS). + api_key: Optional API key. + + Returns: + Dict with run_id, result_url, and other task metadata. + """ + from parallel import Parallel + + client = Parallel(api_key=resolve_api_key(api_key)) + + task = client.task_run.create( + input=query[:15000], + processor=processor, + ) + + return { + "run_id": task.run_id, + "result_url": getattr(task, "result_url", f"https://platform.parallel.ai/tasks/{task.run_id}"), + "processor": processor, + "status": getattr(task, "status", "pending"), + } + + +def get_research_status( + run_id: str, + api_key: str | None = None, +) -> dict[str, Any]: + """Get the current status of a research task. + + Args: + run_id: The task run ID. + api_key: Optional API key. + + Returns: + Dict with status and other task info. + """ + from parallel import Parallel + + client = Parallel(api_key=resolve_api_key(api_key)) + status = client.task_run.retrieve(run_id=run_id) + + return { + "run_id": run_id, + "status": status.status, + "result_url": f"https://platform.parallel.ai/tasks/{run_id}", + } + + +def get_research_result( + run_id: str, + api_key: str | None = None, + include_basis: bool = True, +) -> dict[str, Any]: + """Get the result of a completed research task. + + Args: + run_id: The task run ID. + api_key: Optional API key. + include_basis: Whether to include citations/sources. + + Returns: + Dict with content, basis (if included), and metadata. + """ + from parallel import Parallel + + client = Parallel(api_key=resolve_api_key(api_key)) + result = client.task_run.result(run_id=run_id) + + output = result.output if hasattr(result, "output") else {} + content = _extract_content(output) + + response: dict[str, Any] = { + "run_id": run_id, + "status": "completed", + "content": content, + } + + if include_basis and hasattr(output, "basis"): + response["basis"] = extract_basis(output) + + return response + + +def _poll_until_complete( + client, + run_id: str, + result_url: str, + timeout: int, + poll_interval: int, + include_basis: bool, + on_status: Callable[[str, str], None] | None, +) -> dict[str, Any]: + """Poll a research task until completion and return the result. + + This is the shared polling logic used by both run_research and poll_research. + + Args: + client: Parallel client instance. + run_id: The task run ID to poll. + result_url: URL to view results. + timeout: Maximum wait time in seconds. + poll_interval: Seconds between status checks. + include_basis: Whether to include citations/sources. + on_status: Optional callback called with (status, run_id) on each poll. + + Returns: + Dict with content, basis (if included), and metadata. + + Raises: + TimeoutError: If the task doesn't complete within timeout. + RuntimeError: If the task fails or is cancelled. + """ + deadline = time.time() + timeout + + while time.time() < deadline: + status = client.task_run.retrieve(run_id=run_id) + current_status = status.status + + if on_status: + on_status(current_status, run_id) + + if current_status in TERMINAL_STATUSES: + if current_status == "completed": + result = client.task_run.result(run_id=run_id) + output = result.output if hasattr(result, "output") else {} + content = _extract_content(output) + + response: dict[str, Any] = { + "run_id": run_id, + "result_url": result_url, + "status": "completed", + "content": content, + } + + if include_basis and hasattr(output, "basis"): + response["basis"] = extract_basis(output) + + return response + + error = getattr(status, "error", None) or f"Task {current_status}" + raise RuntimeError(f"Research {current_status}: {error}") + + time.sleep(poll_interval) + + raise TimeoutError(f"Research task {run_id} timed out after {timeout} seconds") + + +def run_research( + query: str, + processor: str = "pro-fast", + api_key: str | None = None, + timeout: int = 3600, + poll_interval: int = 45, + include_basis: bool = True, + on_status: Callable[[str, str], None] | None = None, +) -> dict[str, Any]: + """Run deep research and wait for results. + + This is the main entry point for running research. It creates a task, + polls for completion, and returns the result. + + Args: + query: Research question or topic (max 15,000 chars). + processor: Processor tier (see RESEARCH_PROCESSORS). + api_key: Optional API key. + timeout: Maximum wait time in seconds (default: 3600 = 1 hour). + poll_interval: Seconds between status checks (default: 45). + include_basis: Whether to include citations/sources. + on_status: Optional callback called with (status, run_id) on each poll. + + Returns: + Dict with content, basis (if included), and metadata. + + Raises: + TimeoutError: If the task doesn't complete within timeout. + RuntimeError: If the task fails or is cancelled. + """ + from parallel import Parallel + + client = Parallel(api_key=resolve_api_key(api_key)) + + task = client.task_run.create( + input=query[:15000], + processor=processor, + ) + run_id = task.run_id + result_url = getattr(task, "result_url", f"https://platform.parallel.ai/tasks/{run_id}") + + if on_status: + on_status("created", run_id) + + return _poll_until_complete(client, run_id, result_url, timeout, poll_interval, include_basis, on_status) + + +def poll_research( + run_id: str, + api_key: str | None = None, + timeout: int = 3600, + poll_interval: int = 45, + include_basis: bool = True, + on_status: Callable[[str, str], None] | None = None, +) -> dict[str, Any]: + """Resume polling an existing research task. + + Use this to reconnect to a task that was created earlier. + + Args: + run_id: The task run ID to poll. + api_key: Optional API key. + timeout: Maximum wait time in seconds. + poll_interval: Seconds between status checks. + include_basis: Whether to include citations/sources. + on_status: Optional callback called with (status, run_id) on each poll. + + Returns: + Dict with content, basis (if included), and metadata. + """ + from parallel import Parallel + + client = Parallel(api_key=resolve_api_key(api_key)) + result_url = f"https://platform.parallel.ai/tasks/{run_id}" + + if on_status: + on_status("polling", run_id) + + return _poll_until_complete(client, run_id, result_url, timeout, poll_interval, include_basis, on_status) + + +def _extract_content(output: Any) -> str: + """Extract the content string from various output formats.""" + if output is None: + return "" + + if isinstance(output, str): + return output + + if isinstance(output, dict): + # Priority: content > markdown > text > JSON dump + for key in ("content", "markdown", "text"): + if key in output: + return str(output[key]) + return json.dumps(output, indent=2, default=str) + + # Handle SDK response objects + if hasattr(output, "content"): + content = output.content + if isinstance(content, str): + return content + if isinstance(content, dict): + return json.dumps(content, indent=2, default=str) + return str(content) + + if hasattr(output, "markdown"): + return str(output.markdown) + + if hasattr(output, "text"): + return str(output.text) + + return str(output) diff --git a/tests/test_research.py b/tests/test_research.py new file mode 100644 index 0000000..6bff5ac --- /dev/null +++ b/tests/test_research.py @@ -0,0 +1,553 @@ +"""Tests for the deep research functionality.""" + +import json +from unittest import mock + +import pytest +from click.testing import CliRunner + +from parallel_web_tools.cli.commands import main +from parallel_web_tools.core.research import ( + RESEARCH_PROCESSORS, + _extract_content, + create_research_task, + get_research_result, + get_research_status, + poll_research, + run_research, +) + + +@pytest.fixture +def runner(): + """Create a CLI test runner.""" + return CliRunner() + + +@pytest.fixture +def mock_parallel_client(): + """Create a mock Parallel client.""" + with mock.patch("parallel.Parallel") as mock_cls: + mock_client = mock.MagicMock() + mock_cls.return_value = mock_client + yield mock_client + + +@pytest.fixture +def mock_api_key(): + """Mock the API key resolution.""" + with mock.patch("parallel_web_tools.core.research.resolve_api_key", return_value="test-key"): + yield + + +# ============================================================================= +# Core Research Function Tests +# ============================================================================= + + +class TestExtractContent: + """Tests for _extract_content helper function.""" + + def test_extract_none(self): + """Should return empty string for None.""" + assert _extract_content(None) == "" + + def test_extract_string(self): + """Should return string directly.""" + assert _extract_content("hello world") == "hello world" + + def test_extract_dict_content(self): + """Should extract content key from dict.""" + output = {"content": "research results"} + assert _extract_content(output) == "research results" + + def test_extract_dict_markdown(self): + """Should extract markdown key from dict.""" + output = {"markdown": "# Report"} + assert _extract_content(output) == "# Report" + + def test_extract_dict_text(self): + """Should extract text key from dict.""" + output = {"text": "plain text"} + assert _extract_content(output) == "plain text" + + def test_extract_dict_priority(self): + """Content should have priority over markdown.""" + output = {"content": "content", "markdown": "markdown"} + assert _extract_content(output) == "content" + + def test_extract_dict_json_fallback(self): + """Should JSON dump dict without known keys.""" + output = {"custom_key": "value"} + result = _extract_content(output) + assert "custom_key" in result + assert "value" in result + + def test_extract_object_content(self): + """Should extract content attribute from object.""" + obj = mock.MagicMock() + obj.content = "object content" + assert _extract_content(obj) == "object content" + + def test_extract_object_markdown(self): + """Should extract markdown attribute from object.""" + obj = mock.MagicMock(spec=["markdown"]) + obj.markdown = "# Heading" + assert _extract_content(obj) == "# Heading" + + +class TestCreateResearchTask: + """Tests for create_research_task function.""" + + def test_create_task_basic(self, mock_parallel_client, mock_api_key): + """Should create a task and return metadata.""" + mock_task = mock.MagicMock() + mock_task.run_id = "trun_123" + mock_task.result_url = "https://platform.parallel.ai/tasks/trun_123" + mock_task.status = "pending" + mock_parallel_client.task_run.create.return_value = mock_task + + result = create_research_task("What is AI?", processor="pro-fast") + + assert result["run_id"] == "trun_123" + assert result["processor"] == "pro-fast" + assert "result_url" in result + mock_parallel_client.task_run.create.assert_called_once() + + def test_create_task_truncates_query(self, mock_parallel_client, mock_api_key): + """Should truncate query to 15000 chars.""" + mock_task = mock.MagicMock() + mock_task.run_id = "trun_123" + mock_parallel_client.task_run.create.return_value = mock_task + + long_query = "x" * 20000 + create_research_task(long_query) + + call_args = mock_parallel_client.task_run.create.call_args + assert len(call_args.kwargs["input"]) == 15000 + + +class TestGetResearchStatus: + """Tests for get_research_status function.""" + + def test_get_status(self, mock_parallel_client, mock_api_key): + """Should retrieve task status.""" + mock_status = mock.MagicMock() + mock_status.status = "running" + mock_parallel_client.task_run.retrieve.return_value = mock_status + + result = get_research_status("trun_123") + + assert result["run_id"] == "trun_123" + assert result["status"] == "running" + mock_parallel_client.task_run.retrieve.assert_called_once_with(run_id="trun_123") + + +class TestGetResearchResult: + """Tests for get_research_result function.""" + + def test_get_result_basic(self, mock_parallel_client, mock_api_key): + """Should retrieve completed task result.""" + mock_output = mock.MagicMock() + mock_output.content = "Research findings" + mock_output.basis = [] + + mock_result = mock.MagicMock() + mock_result.output = mock_output + mock_parallel_client.task_run.result.return_value = mock_result + + result = get_research_result("trun_123") + + assert result["run_id"] == "trun_123" + assert result["status"] == "completed" + assert result["content"] == "Research findings" + + def test_get_result_with_basis(self, mock_parallel_client, mock_api_key): + """Should include basis when requested.""" + mock_basis = mock.MagicMock() + mock_basis.field = "summary" + mock_basis.citations = [] + mock_basis.reasoning = "Based on sources" + mock_basis.confidence = "HIGH" + + mock_output = mock.MagicMock() + mock_output.content = "Findings" + mock_output.basis = [mock_basis] + + mock_result = mock.MagicMock() + mock_result.output = mock_output + mock_parallel_client.task_run.result.return_value = mock_result + + result = get_research_result("trun_123", include_basis=True) + + assert "basis" in result + + +class TestRunResearch: + """Tests for run_research function.""" + + def test_run_research_success(self, mock_parallel_client, mock_api_key): + """Should create task and poll until completion.""" + # Mock task creation + mock_task = mock.MagicMock() + mock_task.run_id = "trun_123" + mock_parallel_client.task_run.create.return_value = mock_task + + # Mock status polling - first running, then completed + mock_status_running = mock.MagicMock() + mock_status_running.status = "running" + + mock_status_completed = mock.MagicMock() + mock_status_completed.status = "completed" + + mock_parallel_client.task_run.retrieve.side_effect = [ + mock_status_running, + mock_status_completed, + ] + + # Mock result retrieval + mock_output = mock.MagicMock() + mock_output.content = "Research complete" + mock_output.basis = [] + + mock_result = mock.MagicMock() + mock_result.output = mock_output + mock_parallel_client.task_run.result.return_value = mock_result + + with mock.patch("parallel_web_tools.core.research.time.sleep"): + result = run_research("What is AI?", poll_interval=1, timeout=10) + + assert result["status"] == "completed" + assert result["content"] == "Research complete" + + def test_run_research_timeout(self, mock_parallel_client, mock_api_key): + """Should raise TimeoutError when task doesn't complete.""" + mock_task = mock.MagicMock() + mock_task.run_id = "trun_123" + mock_parallel_client.task_run.create.return_value = mock_task + + mock_status = mock.MagicMock() + mock_status.status = "running" + mock_parallel_client.task_run.retrieve.return_value = mock_status + + with mock.patch("parallel_web_tools.core.research.time.sleep"): + with mock.patch("parallel_web_tools.core.research.time.time") as mock_time: + # Simulate timeout by returning increasing time values + mock_time.side_effect = [0, 0, 5, 10, 15] + + with pytest.raises(TimeoutError): + run_research("What is AI?", timeout=10, poll_interval=1) + + def test_run_research_failed(self, mock_parallel_client, mock_api_key): + """Should raise RuntimeError when task fails.""" + mock_task = mock.MagicMock() + mock_task.run_id = "trun_123" + mock_parallel_client.task_run.create.return_value = mock_task + + mock_status = mock.MagicMock() + mock_status.status = "failed" + mock_status.error = "Processing error" + mock_parallel_client.task_run.retrieve.return_value = mock_status + + with mock.patch("parallel_web_tools.core.research.time.sleep"): + with pytest.raises(RuntimeError, match="failed"): + run_research("What is AI?", poll_interval=1) + + def test_run_research_on_status_callback(self, mock_parallel_client, mock_api_key): + """Should call on_status callback during polling.""" + mock_task = mock.MagicMock() + mock_task.run_id = "trun_123" + mock_parallel_client.task_run.create.return_value = mock_task + + mock_status = mock.MagicMock() + mock_status.status = "completed" + mock_parallel_client.task_run.retrieve.return_value = mock_status + + mock_output = mock.MagicMock() + mock_output.content = "Done" + mock_output.basis = [] + mock_result = mock.MagicMock() + mock_result.output = mock_output + mock_parallel_client.task_run.result.return_value = mock_result + + statuses = [] + + def on_status(status, run_id): + statuses.append((status, run_id)) + + with mock.patch("parallel_web_tools.core.research.time.sleep"): + run_research("What is AI?", on_status=on_status, poll_interval=1) + + assert ("created", "trun_123") in statuses + assert ("completed", "trun_123") in statuses + + +class TestPollResearch: + """Tests for poll_research function.""" + + def test_poll_existing_task(self, mock_parallel_client, mock_api_key): + """Should poll existing task until completion.""" + mock_status = mock.MagicMock() + mock_status.status = "completed" + mock_parallel_client.task_run.retrieve.return_value = mock_status + + mock_output = mock.MagicMock() + mock_output.content = "Results" + mock_output.basis = [] + mock_result = mock.MagicMock() + mock_result.output = mock_output + mock_parallel_client.task_run.result.return_value = mock_result + + with mock.patch("parallel_web_tools.core.research.time.sleep"): + result = poll_research("trun_123", poll_interval=1) + + assert result["status"] == "completed" + assert result["run_id"] == "trun_123" + + +class TestResearchProcessors: + """Tests for RESEARCH_PROCESSORS constant.""" + + def test_processors_defined(self): + """Should have expected processors.""" + assert "pro-fast" in RESEARCH_PROCESSORS + assert "ultra" in RESEARCH_PROCESSORS + assert "ultra8x" in RESEARCH_PROCESSORS + + def test_processors_have_descriptions(self): + """All processors should have descriptions.""" + for _proc, desc in RESEARCH_PROCESSORS.items(): + assert isinstance(desc, str) + assert len(desc) > 0 + + +# ============================================================================= +# CLI Research Command Tests +# ============================================================================= + + +class TestResearchGroup: + """Tests for the research command group.""" + + def test_research_help(self, runner): + """Should show research subcommands.""" + result = runner.invoke(main, ["research", "--help"]) + assert result.exit_code == 0 + assert "run" in result.output + assert "status" in result.output + assert "poll" in result.output + assert "processors" in result.output + + def test_research_in_main_help(self, runner): + """Research should appear in main CLI help.""" + result = runner.invoke(main, ["--help"]) + assert result.exit_code == 0 + assert "research" in result.output + + +class TestResearchRunCommand: + """Tests for the research run command.""" + + def test_research_run_help(self, runner): + """Should show research run help.""" + result = runner.invoke(main, ["research", "run", "--help"]) + assert result.exit_code == 0 + assert "--processor" in result.output + assert "--timeout" in result.output + assert "--no-wait" in result.output + assert "--output" in result.output + + def test_research_run_no_query(self, runner): + """Should error without query or input file.""" + result = runner.invoke(main, ["research", "run"]) + assert result.exit_code != 0 + assert "query" in result.output.lower() or "input" in result.output.lower() + + def test_research_run_with_input_file(self, runner, tmp_path): + """Should read query from file.""" + query_file = tmp_path / "query.txt" + query_file.write_text("What is quantum computing?") + + with mock.patch("parallel_web_tools.cli.commands.create_research_task") as mock_create: + mock_create.return_value = { + "run_id": "trun_123", + "result_url": "https://platform.parallel.ai/tasks/trun_123", + "status": "pending", + } + + result = runner.invoke(main, ["research", "run", "--input-file", str(query_file), "--no-wait"]) + + assert result.exit_code == 0 + mock_create.assert_called_once() + call_args = mock_create.call_args + assert "quantum computing" in call_args[0][0] + + def test_research_run_no_wait(self, runner): + """Should return immediately with --no-wait.""" + with mock.patch("parallel_web_tools.cli.commands.create_research_task") as mock_create: + mock_create.return_value = { + "run_id": "trun_123", + "result_url": "https://platform.parallel.ai/tasks/trun_123", + "status": "pending", + } + + result = runner.invoke(main, ["research", "run", "What is AI?", "--no-wait"]) + + assert result.exit_code == 0 + assert "trun_123" in result.output + mock_create.assert_called_once() + + def test_research_run_json_output(self, runner): + """Should output JSON with --json flag.""" + with mock.patch("parallel_web_tools.cli.commands.create_research_task") as mock_create: + mock_create.return_value = { + "run_id": "trun_123", + "result_url": "https://platform.parallel.ai/tasks/trun_123", + "status": "pending", + } + + result = runner.invoke(main, ["research", "run", "What is AI?", "--no-wait", "--json"]) + + assert result.exit_code == 0 + # Find the JSON in the output (it starts with { and ends with }) + lines = result.output.strip().split("\n") + json_lines = [] + in_json = False + for line in lines: + if line.strip().startswith("{"): + in_json = True + if in_json: + json_lines.append(line) + if in_json and line.strip().startswith("}"): + break + output = json.loads("\n".join(json_lines)) + assert output["run_id"] == "trun_123" + + def test_research_run_with_wait(self, runner): + """Should poll and return results without --no-wait.""" + with mock.patch("parallel_web_tools.cli.commands.run_research") as mock_run: + mock_run.return_value = { + "run_id": "trun_123", + "result_url": "https://platform.parallel.ai/tasks/trun_123", + "status": "completed", + "content": "AI research findings", + } + + result = runner.invoke(main, ["research", "run", "What is AI?", "--poll-interval", "1", "--timeout", "10"]) + + assert result.exit_code == 0 + assert "Research Complete" in result.output or "AI research findings" in result.output + mock_run.assert_called_once() + + +class TestResearchStatusCommand: + """Tests for the research status command.""" + + def test_research_status_help(self, runner): + """Should show status help.""" + result = runner.invoke(main, ["research", "status", "--help"]) + assert result.exit_code == 0 + assert "RUN_ID" in result.output + + def test_research_status(self, runner): + """Should show task status.""" + with mock.patch("parallel_web_tools.cli.commands.get_research_status") as mock_status: + mock_status.return_value = { + "run_id": "trun_123", + "status": "running", + "result_url": "https://platform.parallel.ai/tasks/trun_123", + } + + result = runner.invoke(main, ["research", "status", "trun_123"]) + + assert result.exit_code == 0 + assert "trun_123" in result.output + assert "running" in result.output.lower() + + def test_research_status_json(self, runner): + """Should output JSON with --json flag.""" + with mock.patch("parallel_web_tools.cli.commands.get_research_status") as mock_status: + mock_status.return_value = { + "run_id": "trun_123", + "status": "completed", + "result_url": "https://platform.parallel.ai/tasks/trun_123", + } + + result = runner.invoke(main, ["research", "status", "trun_123", "--json"]) + + assert result.exit_code == 0 + output = json.loads(result.output) + assert output["status"] == "completed" + + +class TestResearchPollCommand: + """Tests for the research poll command.""" + + def test_research_poll_help(self, runner): + """Should show poll help.""" + result = runner.invoke(main, ["research", "poll", "--help"]) + assert result.exit_code == 0 + assert "RUN_ID" in result.output + assert "--timeout" in result.output + + def test_research_poll(self, runner): + """Should poll and return results.""" + with mock.patch("parallel_web_tools.cli.commands.poll_research") as mock_poll: + mock_poll.return_value = { + "run_id": "trun_123", + "result_url": "https://platform.parallel.ai/tasks/trun_123", + "status": "completed", + "content": "Research results here", + } + + result = runner.invoke(main, ["research", "poll", "trun_123", "--poll-interval", "1"]) + + assert result.exit_code == 0 + assert "Research Complete" in result.output or "Research results" in result.output + + +class TestResearchProcessorsCommand: + """Tests for the research processors command.""" + + def test_research_processors(self, runner): + """Should list all processors.""" + result = runner.invoke(main, ["research", "processors"]) + assert result.exit_code == 0 + assert "pro-fast" in result.output + assert "ultra" in result.output + assert "ultra8x" in result.output + + +class TestResearchOutputFile: + """Tests for saving research results to files.""" + + def test_research_save_to_file(self, runner, tmp_path): + """Should save results to markdown and JSON files.""" + output_file = tmp_path / "report.md" + + with mock.patch("parallel_web_tools.cli.commands.run_research") as mock_run: + mock_run.return_value = { + "run_id": "trun_123", + "result_url": "https://platform.parallel.ai/tasks/trun_123", + "status": "completed", + "content": "# Research Report\n\nFindings here.", + "basis": [], + } + + result = runner.invoke( + main, + ["research", "run", "What is AI?", "-o", str(output_file), "--poll-interval", "1"], + ) + + assert result.exit_code == 0 + + # Check markdown file + assert output_file.exists() + content = output_file.read_text() + assert "Research Report" in content + + # Check JSON file + json_file = tmp_path / "report.json" + assert json_file.exists() + data = json.loads(json_file.read_text()) + assert data["run_id"] == "trun_123" From 734f4d850bdbeed36c500de3f6a9b9d250c0b975 Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 13:49:26 -0500 Subject: [PATCH 04/10] fix: separate markdown content from JSON metadata in research output --- parallel_web_tools/cli/commands.py | 39 +++++++++++++++++++++--------- tests/test_research.py | 12 ++++++--- 2 files changed, 36 insertions(+), 15 deletions(-) diff --git a/parallel_web_tools/cli/commands.py b/parallel_web_tools/cli/commands.py index 016e9a1..ae5505c 100644 --- a/parallel_web_tools/cli/commands.py +++ b/parallel_web_tools/cli/commands.py @@ -862,34 +862,49 @@ def _output_research_result( no_basis: bool, ): """Output research result to console and/or file.""" + from datetime import datetime + content = result.get("content", "") basis = result.get("basis", []) - # Build JSON output - output_data = { - "run_id": result.get("run_id"), - "result_url": result.get("result_url"), - "status": result.get("status"), - "content": content, - } - if not no_basis and basis: - output_data["basis"] = basis - # Save to file if requested if output_file: # Write markdown content with open(output_file, "w") as f: f.write(content) - # Write JSON metadata + # Write JSON metadata (without content - that's in the markdown file) json_file = output_file.rsplit(".", 1)[0] + ".json" if "." in output_file else output_file + ".json" + metadata = { + "run_id": result.get("run_id"), + "result_url": result.get("result_url"), + "status": result.get("status"), + "downloaded_at": datetime.now().isoformat(), + "files": { + "markdown": output_file, + "json": json_file, + }, + } + if not no_basis and basis: + metadata["basis"] = basis + with open(json_file, "w") as f: - json.dump(output_data, f, indent=2) + json.dump(metadata, f, indent=2) console.print("\n[dim]Results saved to:[/dim]") console.print(f" [green]Markdown:[/green] {output_file}") console.print(f" [green]JSON:[/green] {json_file}") + # Build console/stdout JSON output (includes content for piping) + output_data = { + "run_id": result.get("run_id"), + "result_url": result.get("result_url"), + "status": result.get("status"), + "content": content, + } + if not no_basis and basis: + output_data["basis"] = basis + # Output to console if output_json: print(json.dumps(output_data, indent=2)) diff --git a/tests/test_research.py b/tests/test_research.py index 6bff5ac..79fdaa8 100644 --- a/tests/test_research.py +++ b/tests/test_research.py @@ -531,7 +531,7 @@ def test_research_save_to_file(self, runner, tmp_path): "result_url": "https://platform.parallel.ai/tasks/trun_123", "status": "completed", "content": "# Research Report\n\nFindings here.", - "basis": [], + "basis": [{"field": "summary", "citations": []}], } result = runner.invoke( @@ -541,13 +541,19 @@ def test_research_save_to_file(self, runner, tmp_path): assert result.exit_code == 0 - # Check markdown file + # Check markdown file has content assert output_file.exists() content = output_file.read_text() assert "Research Report" in content - # Check JSON file + # Check JSON file has metadata (not content) json_file = tmp_path / "report.json" assert json_file.exists() data = json.loads(json_file.read_text()) assert data["run_id"] == "trun_123" + assert data["status"] == "completed" + assert "downloaded_at" in data + assert "files" in data + assert data["files"]["markdown"] == str(output_file) + assert "basis" in data # Should have basis metadata + assert "content" not in data # Content should NOT be in JSON (it's in markdown) From 09dd375ce89257fe4e1943aa69c643b51e5afaac Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 13:52:45 -0500 Subject: [PATCH 05/10] fix: update processor latencies to match documentation --- parallel_web_tools/core/research.py | 31 +++++++++++++++++++---------- tests/test_research.py | 5 +++++ 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/parallel_web_tools/core/research.py b/parallel_web_tools/core/research.py index 16055d6..4a020de 100644 --- a/parallel_web_tools/core/research.py +++ b/parallel_web_tools/core/research.py @@ -16,17 +16,28 @@ from parallel_web_tools.core.auth import resolve_api_key from parallel_web_tools.core.batch import extract_basis -# Processor tiers for deep research with expected latency +# Processor tiers for deep research with expected latency (from docs) +# Fast variants are 2-5x faster but may use slightly less fresh data RESEARCH_PROCESSORS = { - "pro-fast": "1-5 min - exploratory research (default)", - "pro": "2-10 min - exploratory research, fresher data", - "ultra-fast": "2-12 min - multi-source deep research", - "ultra": "5-25 min - advanced deep research, fresher data", - "ultra2x-fast": "2-25 min - difficult deep research", - "ultra2x": "5-50 min - difficult deep research, fresher data", - "ultra4x-fast": "2-45 min - very difficult research", - "ultra4x": "5-90 min - very difficult research, fresher data", - "ultra8x-fast": "2-60 min - most challenging research", + # Fast processors (optimized for speed) + "lite-fast": "10-20s - quick lookups", + "base-fast": "15-50s - simple questions", + "core-fast": "15s-100s - moderate research", + "core2x-fast": "15s-3min - extended research", + "pro-fast": "30s-5min - exploratory research (default)", + "ultra-fast": "1-10min - multi-source deep research", + "ultra2x-fast": "1-20min - difficult deep research", + "ultra4x-fast": "1-40min - very difficult research", + "ultra8x-fast": "1min-1hr - most challenging research", + # Standard processors (fresher data) + "lite": "10-60s - quick lookups, fresher data", + "base": "15-100s - simple questions, fresher data", + "core": "1-5min - moderate research, fresher data", + "core2x": "1-10min - extended research, fresher data", + "pro": "2-10min - exploratory research, fresher data", + "ultra": "5-25min - advanced deep research, fresher data", + "ultra2x": "5-50min - difficult deep research, fresher data", + "ultra4x": "5-90min - very difficult research, fresher data", "ultra8x": "5min-2hr - most challenging research, fresher data", } diff --git a/tests/test_research.py b/tests/test_research.py index 79fdaa8..29c8038 100644 --- a/tests/test_research.py +++ b/tests/test_research.py @@ -310,7 +310,12 @@ class TestResearchProcessors: def test_processors_defined(self): """Should have expected processors.""" + # Fast variants + assert "lite-fast" in RESEARCH_PROCESSORS assert "pro-fast" in RESEARCH_PROCESSORS + assert "ultra8x-fast" in RESEARCH_PROCESSORS + # Standard variants + assert "lite" in RESEARCH_PROCESSORS assert "ultra" in RESEARCH_PROCESSORS assert "ultra8x" in RESEARCH_PROCESSORS From 114bc16ee2f08213445c5621a21d42781c4b07ee Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 14:04:10 -0500 Subject: [PATCH 06/10] fix: use TextSchemaParam to get markdown output from deep research --- parallel_web_tools/core/research.py | 49 ++++++++++++++++++++++++++--- 1 file changed, 44 insertions(+), 5 deletions(-) diff --git a/parallel_web_tools/core/research.py b/parallel_web_tools/core/research.py index 4a020de..0976435 100644 --- a/parallel_web_tools/core/research.py +++ b/parallel_web_tools/core/research.py @@ -48,6 +48,7 @@ def create_research_task( query: str, processor: str = "pro-fast", api_key: str | None = None, + output_format: str = "text", ) -> dict[str, Any]: """Create a deep research task without waiting for results. @@ -55,17 +56,25 @@ def create_research_task( query: Research question or topic (max 15,000 chars). processor: Processor tier (see RESEARCH_PROCESSORS). api_key: Optional API key. + output_format: "text" for markdown report (default), "auto" for structured JSON. Returns: Dict with run_id, result_url, and other task metadata. """ from parallel import Parallel + from parallel.types import TaskSpecParam, TextSchemaParam client = Parallel(api_key=resolve_api_key(api_key)) + # Build task spec based on output format + task_spec = None + if output_format == "text": + task_spec = TaskSpecParam(output_schema=TextSchemaParam(type="text")) + task = client.task_run.create( input=query[:15000], processor=processor, + task_spec=task_spec, ) return { @@ -208,6 +217,7 @@ def run_research( poll_interval: int = 45, include_basis: bool = True, on_status: Callable[[str, str], None] | None = None, + output_format: str = "text", ) -> dict[str, Any]: """Run deep research and wait for results. @@ -222,6 +232,7 @@ def run_research( poll_interval: Seconds between status checks (default: 45). include_basis: Whether to include citations/sources. on_status: Optional callback called with (status, run_id) on each poll. + output_format: "text" for markdown report (default), "auto" for structured JSON. Returns: Dict with content, basis (if included), and metadata. @@ -231,12 +242,19 @@ def run_research( RuntimeError: If the task fails or is cancelled. """ from parallel import Parallel + from parallel.types import TaskSpecParam, TextSchemaParam client = Parallel(api_key=resolve_api_key(api_key)) + # Build task spec based on output format + task_spec = None + if output_format == "text": + task_spec = TaskSpecParam(output_schema=TextSchemaParam(type="text")) + task = client.task_run.create( input=query[:15000], processor=processor, + task_spec=task_spec, ) run_id = task.run_id result_url = getattr(task, "result_url", f"https://platform.parallel.ai/tasks/{run_id}") @@ -282,7 +300,16 @@ def poll_research( def _extract_content(output: Any) -> str: - """Extract the content string from various output formats.""" + """Extract the content string from various output formats. + + The Parallel API can return content in several formats: + - Direct string (markdown text) + - Dict with 'content', 'markdown', or 'text' keys + - SDK object with .content, .markdown, or .text attributes + - Nested structures where content contains another dict + + We prioritize finding actual text content over JSON dumping structured data. + """ if output is None: return "" @@ -290,20 +317,32 @@ def _extract_content(output: Any) -> str: return output if isinstance(output, dict): - # Priority: content > markdown > text > JSON dump + # Priority: content > markdown > text for key in ("content", "markdown", "text"): if key in output: - return str(output[key]) + value = output[key] + # Recursively extract if the value is also a dict/object + if isinstance(value, str): + return value + return _extract_content(value) + # Fallback to JSON dump if no text keys found return json.dumps(output, indent=2, default=str) - # Handle SDK response objects + # Handle SDK response objects with attributes if hasattr(output, "content"): content = output.content if isinstance(content, str): return content + # If content is a dict, look for text fields within it if isinstance(content, dict): + for key in ("content", "markdown", "text"): + if key in content: + value = content[key] + if isinstance(value, str): + return value + return _extract_content(value) return json.dumps(content, indent=2, default=str) - return str(content) + return _extract_content(content) if hasattr(output, "markdown"): return str(output.markdown) From 3ac04d60c859baca4fe142f599ae0358aef840f6 Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 15:49:07 -0500 Subject: [PATCH 07/10] bump version and research updates --- parallel_web_tools/__init__.py | 2 +- parallel_web_tools/cli/commands.py | 88 +++------- parallel_web_tools/core/batch.py | 5 +- parallel_web_tools/core/research.py | 155 +++++------------- .../bigquery/cloud_function/requirements.txt | 2 +- pyproject.toml | 2 +- tests/test_cli.py | 2 +- tests/test_research.py | 141 ++++------------ uv.lock | 2 +- 9 files changed, 106 insertions(+), 293 deletions(-) diff --git a/parallel_web_tools/__init__.py b/parallel_web_tools/__init__.py index ebc5689..eb41914 100644 --- a/parallel_web_tools/__init__.py +++ b/parallel_web_tools/__init__.py @@ -27,7 +27,7 @@ run_tasks, ) -__version__ = "0.0.2" +__version__ = "0.0.3" __all__ = [ # Auth diff --git a/parallel_web_tools/cli/commands.py b/parallel_web_tools/cli/commands.py index ae5505c..3fb2233 100644 --- a/parallel_web_tools/cli/commands.py +++ b/parallel_web_tools/cli/commands.py @@ -682,7 +682,6 @@ def research(): @click.option("--no-wait", is_flag=True, help="Return immediately after creating task (don't poll)") @click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to file (markdown)") @click.option("--json", "output_json", is_flag=True, help="Output as JSON") -@click.option("--no-basis", is_flag=True, help="Exclude citations/sources from output") def research_run( query: str | None, input_file: str | None, @@ -692,7 +691,6 @@ def research_run( no_wait: bool, output_file: str | None, output_json: bool, - no_basis: bool, ): """Run deep research on a question or topic. @@ -737,7 +735,9 @@ def research_run( def on_status(status: str, run_id: str): if status == "created": console.print(f"[green]Task created: {run_id}[/green]") - console.print(f"[dim]Track progress: https://platform.parallel.ai/tasks/{run_id}[/dim]\n") + console.print( + f"[dim]Track progress: https://platform.parallel.ai/play/deep-research/{run_id}[/dim]\n" + ) else: console.print(f"[dim]Status: {status}[/dim]") @@ -746,11 +746,10 @@ def on_status(status: str, run_id: str): processor=processor, timeout=timeout, poll_interval=poll_interval, - include_basis=not no_basis, on_status=on_status, ) - _output_research_result(result, output_file, output_json, no_basis) + _output_research_result(result, output_file, output_json) except TimeoutError as e: console.print(f"[bold yellow]Timeout: {e}[/bold yellow]") @@ -805,14 +804,12 @@ def research_status(run_id: str, output_json: bool): @click.option("--poll-interval", type=int, default=45, show_default=True, help="Seconds between status checks") @click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to file (markdown)") @click.option("--json", "output_json", is_flag=True, help="Output as JSON") -@click.option("--no-basis", is_flag=True, help="Exclude citations/sources from output") def research_poll( run_id: str, timeout: int, poll_interval: int, output_file: str | None, output_json: bool, - no_basis: bool, ): """Poll an existing research task until completion. @@ -820,7 +817,7 @@ def research_poll( """ try: console.print(f"[bold cyan]Polling task: {run_id}[/bold cyan]") - console.print(f"[dim]Track progress: https://platform.parallel.ai/tasks/{run_id}[/dim]\n") + console.print(f"[dim]Track progress: https://platform.parallel.ai/play/deep-research/{run_id}[/dim]\n") def on_status(status: str, run_id: str): console.print(f"[dim]Status: {status}[/dim]") @@ -829,11 +826,10 @@ def on_status(status: str, run_id: str): run_id, timeout=timeout, poll_interval=poll_interval, - include_basis=not no_basis, on_status=on_status, ) - _output_research_result(result, output_file, output_json, no_basis) + _output_research_result(result, output_file, output_json) except TimeoutError as e: console.print(f"[bold yellow]Timeout: {e}[/bold yellow]") @@ -859,71 +855,35 @@ def _output_research_result( result: dict, output_file: str | None, output_json: bool, - no_basis: bool, ): - """Output research result to console and/or file.""" - from datetime import datetime - - content = result.get("content", "") - basis = result.get("basis", []) - - # Save to file if requested - if output_file: - # Write markdown content - with open(output_file, "w") as f: - f.write(content) - - # Write JSON metadata (without content - that's in the markdown file) - json_file = output_file.rsplit(".", 1)[0] + ".json" if "." in output_file else output_file + ".json" - metadata = { - "run_id": result.get("run_id"), - "result_url": result.get("result_url"), - "status": result.get("status"), - "downloaded_at": datetime.now().isoformat(), - "files": { - "markdown": output_file, - "json": json_file, - }, - } - if not no_basis and basis: - metadata["basis"] = basis - - with open(json_file, "w") as f: - json.dump(metadata, f, indent=2) - - console.print("\n[dim]Results saved to:[/dim]") - console.print(f" [green]Markdown:[/green] {output_file}") - console.print(f" [green]JSON:[/green] {json_file}") - - # Build console/stdout JSON output (includes content for piping) + """Output research result to console and/or file as JSON.""" output_data = { "run_id": result.get("run_id"), "result_url": result.get("result_url"), "status": result.get("status"), - "content": content, + "output": result.get("output", {}), } - if not no_basis and basis: - output_data["basis"] = basis + + # Save to file if requested + if output_file: + with open(output_file, "w") as f: + json.dump(output_data, f, indent=2, default=str) + console.print(f"\n[green]Results saved to:[/green] {output_file}") # Output to console if output_json: - print(json.dumps(output_data, indent=2)) + print(json.dumps(output_data, indent=2, default=str)) else: console.print("\n[bold green]Research Complete![/bold green]") - console.print(f"[dim]Task: {result.get('run_id')}[/dim]\n") - - # Print content (truncate for console if very long) - if len(content) > 5000 and not output_file: - console.print(content[:5000]) - console.print(f"\n[yellow]... truncated ({len(content)} chars total)[/yellow]") - console.print("[dim]Use --output to save full content to a file[/dim]") - else: - console.print(content) - - # Show citation summary - if basis and not no_basis: - total_citations = sum(len(b.get("citations", [])) for b in basis if isinstance(b, dict)) - console.print(f"\n[dim]Sources: {total_citations} citations from {len(basis)} fields[/dim]") + console.print(f"[dim]Task: {result.get('run_id')}[/dim]") + console.print(f"[dim]URL: {result.get('result_url')}[/dim]\n") + + # Show summary of output + output = result.get("output", {}) + if isinstance(output, dict): + console.print(f"[dim]Output contains {len(output)} fields[/dim]") + if not output_file: + console.print("[dim]Use --output to save full JSON to a file, or --json to print to stdout[/dim]") if __name__ == "__main__": diff --git a/parallel_web_tools/core/batch.py b/parallel_web_tools/core/batch.py index c215b36..3047e27 100644 --- a/parallel_web_tools/core/batch.py +++ b/parallel_web_tools/core/batch.py @@ -13,8 +13,11 @@ def build_output_schema(output_columns: list[str]) -> dict[str, Any]: """Build a JSON schema from output column descriptions.""" properties = {} for col in output_columns: + # Extract base name before any annotations like (type), [hint], {note} base_name = col.split("(")[0].split("[")[0].split("{")[0].strip() - prop_name = base_name.lower().strip().replace(" ", "_").replace("-", "_") + + # Convert to valid property name + prop_name = base_name.lower().replace(" ", "_").replace("-", "_") prop_name = "".join(c for c in prop_name if c.isalnum() or c == "_") if prop_name and not prop_name[0].isalpha(): prop_name = "col_" + prop_name diff --git a/parallel_web_tools/core/research.py b/parallel_web_tools/core/research.py index 0976435..e8c62fb 100644 --- a/parallel_web_tools/core/research.py +++ b/parallel_web_tools/core/research.py @@ -8,13 +8,14 @@ from __future__ import annotations -import json import time from collections.abc import Callable from typing import Any from parallel_web_tools.core.auth import resolve_api_key -from parallel_web_tools.core.batch import extract_basis + +# Base URL for viewing results +PLATFORM_BASE = "https://platform.parallel.ai" # Processor tiers for deep research with expected latency (from docs) # Fast variants are 2-5x faster but may use slightly less fresh data @@ -44,11 +45,35 @@ TERMINAL_STATUSES = ("completed", "failed", "cancelled") +def _serialize_output(output: Any) -> dict[str, Any]: + """Serialize SDK output object to a dictionary. + + The Parallel SDK returns Pydantic-like objects that can be + serialized via model_dump() or to_dict(). + """ + if output is None: + return {} + + if isinstance(output, dict): + return output + + # Try common serialization methods + if hasattr(output, "model_dump"): + return output.model_dump() + + if hasattr(output, "to_dict"): + return output.to_dict() + + if hasattr(output, "__dict__"): + return output.__dict__ + + return {"raw": str(output)} + + def create_research_task( query: str, processor: str = "pro-fast", api_key: str | None = None, - output_format: str = "text", ) -> dict[str, Any]: """Create a deep research task without waiting for results. @@ -56,30 +81,22 @@ def create_research_task( query: Research question or topic (max 15,000 chars). processor: Processor tier (see RESEARCH_PROCESSORS). api_key: Optional API key. - output_format: "text" for markdown report (default), "auto" for structured JSON. Returns: Dict with run_id, result_url, and other task metadata. """ from parallel import Parallel - from parallel.types import TaskSpecParam, TextSchemaParam client = Parallel(api_key=resolve_api_key(api_key)) - # Build task spec based on output format - task_spec = None - if output_format == "text": - task_spec = TaskSpecParam(output_schema=TextSchemaParam(type="text")) - task = client.task_run.create( input=query[:15000], processor=processor, - task_spec=task_spec, ) return { "run_id": task.run_id, - "result_url": getattr(task, "result_url", f"https://platform.parallel.ai/tasks/{task.run_id}"), + "result_url": f"{PLATFORM_BASE}/play/deep-research/{task.run_id}", "processor": processor, "status": getattr(task, "status", "pending"), } @@ -106,24 +123,22 @@ def get_research_status( return { "run_id": run_id, "status": status.status, - "result_url": f"https://platform.parallel.ai/tasks/{run_id}", + "result_url": f"{PLATFORM_BASE}/play/deep-research/{run_id}", } def get_research_result( run_id: str, api_key: str | None = None, - include_basis: bool = True, ) -> dict[str, Any]: """Get the result of a completed research task. Args: run_id: The task run ID. api_key: Optional API key. - include_basis: Whether to include citations/sources. Returns: - Dict with content, basis (if included), and metadata. + Dict with output data and metadata. """ from parallel import Parallel @@ -131,19 +146,15 @@ def get_research_result( result = client.task_run.result(run_id=run_id) output = result.output if hasattr(result, "output") else {} - content = _extract_content(output) + output_data = _serialize_output(output) - response: dict[str, Any] = { + return { "run_id": run_id, + "result_url": f"{PLATFORM_BASE}/play/deep-research/{run_id}", "status": "completed", - "content": content, + "output": output_data, } - if include_basis and hasattr(output, "basis"): - response["basis"] = extract_basis(output) - - return response - def _poll_until_complete( client, @@ -151,24 +162,20 @@ def _poll_until_complete( result_url: str, timeout: int, poll_interval: int, - include_basis: bool, on_status: Callable[[str, str], None] | None, ) -> dict[str, Any]: """Poll a research task until completion and return the result. - This is the shared polling logic used by both run_research and poll_research. - Args: client: Parallel client instance. run_id: The task run ID to poll. result_url: URL to view results. timeout: Maximum wait time in seconds. poll_interval: Seconds between status checks. - include_basis: Whether to include citations/sources. on_status: Optional callback called with (status, run_id) on each poll. Returns: - Dict with content, basis (if included), and metadata. + Dict with content and metadata. Raises: TimeoutError: If the task doesn't complete within timeout. @@ -187,20 +194,15 @@ def _poll_until_complete( if current_status == "completed": result = client.task_run.result(run_id=run_id) output = result.output if hasattr(result, "output") else {} - content = _extract_content(output) + output_data = _serialize_output(output) - response: dict[str, Any] = { + return { "run_id": run_id, "result_url": result_url, "status": "completed", - "content": content, + "output": output_data, } - if include_basis and hasattr(output, "basis"): - response["basis"] = extract_basis(output) - - return response - error = getattr(status, "error", None) or f"Task {current_status}" raise RuntimeError(f"Research {current_status}: {error}") @@ -215,9 +217,7 @@ def run_research( api_key: str | None = None, timeout: int = 3600, poll_interval: int = 45, - include_basis: bool = True, on_status: Callable[[str, str], None] | None = None, - output_format: str = "text", ) -> dict[str, Any]: """Run deep research and wait for results. @@ -230,39 +230,30 @@ def run_research( api_key: Optional API key. timeout: Maximum wait time in seconds (default: 3600 = 1 hour). poll_interval: Seconds between status checks (default: 45). - include_basis: Whether to include citations/sources. on_status: Optional callback called with (status, run_id) on each poll. - output_format: "text" for markdown report (default), "auto" for structured JSON. Returns: - Dict with content, basis (if included), and metadata. + Dict with content and metadata. Raises: TimeoutError: If the task doesn't complete within timeout. RuntimeError: If the task fails or is cancelled. """ from parallel import Parallel - from parallel.types import TaskSpecParam, TextSchemaParam client = Parallel(api_key=resolve_api_key(api_key)) - # Build task spec based on output format - task_spec = None - if output_format == "text": - task_spec = TaskSpecParam(output_schema=TextSchemaParam(type="text")) - task = client.task_run.create( input=query[:15000], processor=processor, - task_spec=task_spec, ) run_id = task.run_id - result_url = getattr(task, "result_url", f"https://platform.parallel.ai/tasks/{run_id}") + result_url = f"{PLATFORM_BASE}/play/deep-research/{run_id}" if on_status: on_status("created", run_id) - return _poll_until_complete(client, run_id, result_url, timeout, poll_interval, include_basis, on_status) + return _poll_until_complete(client, run_id, result_url, timeout, poll_interval, on_status) def poll_research( @@ -270,7 +261,6 @@ def poll_research( api_key: str | None = None, timeout: int = 3600, poll_interval: int = 45, - include_basis: bool = True, on_status: Callable[[str, str], None] | None = None, ) -> dict[str, Any]: """Resume polling an existing research task. @@ -282,72 +272,17 @@ def poll_research( api_key: Optional API key. timeout: Maximum wait time in seconds. poll_interval: Seconds between status checks. - include_basis: Whether to include citations/sources. on_status: Optional callback called with (status, run_id) on each poll. Returns: - Dict with content, basis (if included), and metadata. + Dict with content and metadata. """ from parallel import Parallel client = Parallel(api_key=resolve_api_key(api_key)) - result_url = f"https://platform.parallel.ai/tasks/{run_id}" + result_url = f"{PLATFORM_BASE}/play/deep-research/{run_id}" if on_status: on_status("polling", run_id) - return _poll_until_complete(client, run_id, result_url, timeout, poll_interval, include_basis, on_status) - - -def _extract_content(output: Any) -> str: - """Extract the content string from various output formats. - - The Parallel API can return content in several formats: - - Direct string (markdown text) - - Dict with 'content', 'markdown', or 'text' keys - - SDK object with .content, .markdown, or .text attributes - - Nested structures where content contains another dict - - We prioritize finding actual text content over JSON dumping structured data. - """ - if output is None: - return "" - - if isinstance(output, str): - return output - - if isinstance(output, dict): - # Priority: content > markdown > text - for key in ("content", "markdown", "text"): - if key in output: - value = output[key] - # Recursively extract if the value is also a dict/object - if isinstance(value, str): - return value - return _extract_content(value) - # Fallback to JSON dump if no text keys found - return json.dumps(output, indent=2, default=str) - - # Handle SDK response objects with attributes - if hasattr(output, "content"): - content = output.content - if isinstance(content, str): - return content - # If content is a dict, look for text fields within it - if isinstance(content, dict): - for key in ("content", "markdown", "text"): - if key in content: - value = content[key] - if isinstance(value, str): - return value - return _extract_content(value) - return json.dumps(content, indent=2, default=str) - return _extract_content(content) - - if hasattr(output, "markdown"): - return str(output.markdown) - - if hasattr(output, "text"): - return str(output.text) - - return str(output) + return _poll_until_complete(client, run_id, result_url, timeout, poll_interval, on_status) diff --git a/parallel_web_tools/integrations/bigquery/cloud_function/requirements.txt b/parallel_web_tools/integrations/bigquery/cloud_function/requirements.txt index 11c8ccd..c82b529 100644 --- a/parallel_web_tools/integrations/bigquery/cloud_function/requirements.txt +++ b/parallel_web_tools/integrations/bigquery/cloud_function/requirements.txt @@ -1,5 +1,5 @@ # Cloud Function dependencies for BigQuery Remote Function functions-framework>=3.0.0 flask>=3.0.0 -parallel-web-tools>=0.0.2 +parallel-web-tools>=0.0.3 google-cloud-secret-manager>=2.20.0 diff --git a/pyproject.toml b/pyproject.toml index 597e9cc..3238b42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "parallel-web-tools" -version = "0.0.2" +version = "0.0.3" description = "Parallel Tools: CLI and data enrichment utilities for the Parallel API" readme = "README.md" requires-python = ">=3.12" diff --git a/tests/test_cli.py b/tests/test_cli.py index 0f44692..e118e6f 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -111,7 +111,7 @@ def test_version(self, runner): """Should show version.""" result = runner.invoke(main, ["--version"]) assert result.exit_code == 0 - assert "0.0.2" in result.output + assert "0.0.3" in result.output class TestAuthCommand: diff --git a/tests/test_research.py b/tests/test_research.py index 29c8038..9bf12a4 100644 --- a/tests/test_research.py +++ b/tests/test_research.py @@ -9,7 +9,6 @@ from parallel_web_tools.cli.commands import main from parallel_web_tools.core.research import ( RESEARCH_PROCESSORS, - _extract_content, create_research_task, get_research_result, get_research_status, @@ -45,57 +44,6 @@ def mock_api_key(): # ============================================================================= -class TestExtractContent: - """Tests for _extract_content helper function.""" - - def test_extract_none(self): - """Should return empty string for None.""" - assert _extract_content(None) == "" - - def test_extract_string(self): - """Should return string directly.""" - assert _extract_content("hello world") == "hello world" - - def test_extract_dict_content(self): - """Should extract content key from dict.""" - output = {"content": "research results"} - assert _extract_content(output) == "research results" - - def test_extract_dict_markdown(self): - """Should extract markdown key from dict.""" - output = {"markdown": "# Report"} - assert _extract_content(output) == "# Report" - - def test_extract_dict_text(self): - """Should extract text key from dict.""" - output = {"text": "plain text"} - assert _extract_content(output) == "plain text" - - def test_extract_dict_priority(self): - """Content should have priority over markdown.""" - output = {"content": "content", "markdown": "markdown"} - assert _extract_content(output) == "content" - - def test_extract_dict_json_fallback(self): - """Should JSON dump dict without known keys.""" - output = {"custom_key": "value"} - result = _extract_content(output) - assert "custom_key" in result - assert "value" in result - - def test_extract_object_content(self): - """Should extract content attribute from object.""" - obj = mock.MagicMock() - obj.content = "object content" - assert _extract_content(obj) == "object content" - - def test_extract_object_markdown(self): - """Should extract markdown attribute from object.""" - obj = mock.MagicMock(spec=["markdown"]) - obj.markdown = "# Heading" - assert _extract_content(obj) == "# Heading" - - class TestCreateResearchTask: """Tests for create_research_task function.""" @@ -103,7 +51,6 @@ def test_create_task_basic(self, mock_parallel_client, mock_api_key): """Should create a task and return metadata.""" mock_task = mock.MagicMock() mock_task.run_id = "trun_123" - mock_task.result_url = "https://platform.parallel.ai/tasks/trun_123" mock_task.status = "pending" mock_parallel_client.task_run.create.return_value = mock_task @@ -149,8 +96,7 @@ class TestGetResearchResult: def test_get_result_basic(self, mock_parallel_client, mock_api_key): """Should retrieve completed task result.""" mock_output = mock.MagicMock() - mock_output.content = "Research findings" - mock_output.basis = [] + mock_output.model_dump.return_value = {"content": {"text": "Research findings"}, "basis": []} mock_result = mock.MagicMock() mock_result.output = mock_output @@ -160,27 +106,8 @@ def test_get_result_basic(self, mock_parallel_client, mock_api_key): assert result["run_id"] == "trun_123" assert result["status"] == "completed" - assert result["content"] == "Research findings" - - def test_get_result_with_basis(self, mock_parallel_client, mock_api_key): - """Should include basis when requested.""" - mock_basis = mock.MagicMock() - mock_basis.field = "summary" - mock_basis.citations = [] - mock_basis.reasoning = "Based on sources" - mock_basis.confidence = "HIGH" - - mock_output = mock.MagicMock() - mock_output.content = "Findings" - mock_output.basis = [mock_basis] - - mock_result = mock.MagicMock() - mock_result.output = mock_output - mock_parallel_client.task_run.result.return_value = mock_result - - result = get_research_result("trun_123", include_basis=True) - - assert "basis" in result + assert "output" in result + assert result["output"]["content"]["text"] == "Research findings" class TestRunResearch: @@ -207,8 +134,7 @@ def test_run_research_success(self, mock_parallel_client, mock_api_key): # Mock result retrieval mock_output = mock.MagicMock() - mock_output.content = "Research complete" - mock_output.basis = [] + mock_output.model_dump.return_value = {"content": {"text": "Research complete"}} mock_result = mock.MagicMock() mock_result.output = mock_output @@ -218,7 +144,7 @@ def test_run_research_success(self, mock_parallel_client, mock_api_key): result = run_research("What is AI?", poll_interval=1, timeout=10) assert result["status"] == "completed" - assert result["content"] == "Research complete" + assert "output" in result def test_run_research_timeout(self, mock_parallel_client, mock_api_key): """Should raise TimeoutError when task doesn't complete.""" @@ -264,8 +190,7 @@ def test_run_research_on_status_callback(self, mock_parallel_client, mock_api_ke mock_parallel_client.task_run.retrieve.return_value = mock_status mock_output = mock.MagicMock() - mock_output.content = "Done" - mock_output.basis = [] + mock_output.model_dump.return_value = {"content": {"text": "Done"}} mock_result = mock.MagicMock() mock_result.output = mock_output mock_parallel_client.task_run.result.return_value = mock_result @@ -292,8 +217,7 @@ def test_poll_existing_task(self, mock_parallel_client, mock_api_key): mock_parallel_client.task_run.retrieve.return_value = mock_status mock_output = mock.MagicMock() - mock_output.content = "Results" - mock_output.basis = [] + mock_output.model_dump.return_value = {"content": {"text": "Results"}} mock_result = mock.MagicMock() mock_result.output = mock_output mock_parallel_client.task_run.result.return_value = mock_result @@ -303,6 +227,7 @@ def test_poll_existing_task(self, mock_parallel_client, mock_api_key): assert result["status"] == "completed" assert result["run_id"] == "trun_123" + assert "output" in result class TestResearchProcessors: @@ -376,7 +301,7 @@ def test_research_run_with_input_file(self, runner, tmp_path): with mock.patch("parallel_web_tools.cli.commands.create_research_task") as mock_create: mock_create.return_value = { "run_id": "trun_123", - "result_url": "https://platform.parallel.ai/tasks/trun_123", + "result_url": "https://platform.parallel.ai/play/deep-research/trun_123", "status": "pending", } @@ -392,7 +317,7 @@ def test_research_run_no_wait(self, runner): with mock.patch("parallel_web_tools.cli.commands.create_research_task") as mock_create: mock_create.return_value = { "run_id": "trun_123", - "result_url": "https://platform.parallel.ai/tasks/trun_123", + "result_url": "https://platform.parallel.ai/play/deep-research/trun_123", "status": "pending", } @@ -407,14 +332,14 @@ def test_research_run_json_output(self, runner): with mock.patch("parallel_web_tools.cli.commands.create_research_task") as mock_create: mock_create.return_value = { "run_id": "trun_123", - "result_url": "https://platform.parallel.ai/tasks/trun_123", + "result_url": "https://platform.parallel.ai/play/deep-research/trun_123", "status": "pending", } result = runner.invoke(main, ["research", "run", "What is AI?", "--no-wait", "--json"]) assert result.exit_code == 0 - # Find the JSON in the output (it starts with { and ends with }) + # Find the JSON in the output lines = result.output.strip().split("\n") json_lines = [] in_json = False @@ -433,15 +358,15 @@ def test_research_run_with_wait(self, runner): with mock.patch("parallel_web_tools.cli.commands.run_research") as mock_run: mock_run.return_value = { "run_id": "trun_123", - "result_url": "https://platform.parallel.ai/tasks/trun_123", + "result_url": "https://platform.parallel.ai/play/deep-research/trun_123", "status": "completed", - "content": "AI research findings", + "output": {"content": {"text": "AI research findings"}}, } result = runner.invoke(main, ["research", "run", "What is AI?", "--poll-interval", "1", "--timeout", "10"]) assert result.exit_code == 0 - assert "Research Complete" in result.output or "AI research findings" in result.output + assert "Research Complete" in result.output mock_run.assert_called_once() @@ -460,7 +385,7 @@ def test_research_status(self, runner): mock_status.return_value = { "run_id": "trun_123", "status": "running", - "result_url": "https://platform.parallel.ai/tasks/trun_123", + "result_url": "https://platform.parallel.ai/play/deep-research/trun_123", } result = runner.invoke(main, ["research", "status", "trun_123"]) @@ -475,7 +400,7 @@ def test_research_status_json(self, runner): mock_status.return_value = { "run_id": "trun_123", "status": "completed", - "result_url": "https://platform.parallel.ai/tasks/trun_123", + "result_url": "https://platform.parallel.ai/play/deep-research/trun_123", } result = runner.invoke(main, ["research", "status", "trun_123", "--json"]) @@ -500,15 +425,15 @@ def test_research_poll(self, runner): with mock.patch("parallel_web_tools.cli.commands.poll_research") as mock_poll: mock_poll.return_value = { "run_id": "trun_123", - "result_url": "https://platform.parallel.ai/tasks/trun_123", + "result_url": "https://platform.parallel.ai/play/deep-research/trun_123", "status": "completed", - "content": "Research results here", + "output": {"content": {"text": "Research results here"}}, } result = runner.invoke(main, ["research", "poll", "trun_123", "--poll-interval", "1"]) assert result.exit_code == 0 - assert "Research Complete" in result.output or "Research results" in result.output + assert "Research Complete" in result.output class TestResearchProcessorsCommand: @@ -527,16 +452,15 @@ class TestResearchOutputFile: """Tests for saving research results to files.""" def test_research_save_to_file(self, runner, tmp_path): - """Should save results to markdown and JSON files.""" - output_file = tmp_path / "report.md" + """Should save results to JSON file.""" + output_file = tmp_path / "report.json" with mock.patch("parallel_web_tools.cli.commands.run_research") as mock_run: mock_run.return_value = { "run_id": "trun_123", - "result_url": "https://platform.parallel.ai/tasks/trun_123", + "result_url": "https://platform.parallel.ai/play/deep-research/trun_123", "status": "completed", - "content": "# Research Report\n\nFindings here.", - "basis": [{"field": "summary", "citations": []}], + "output": {"content": {"text": "Research findings"}, "basis": []}, } result = runner.invoke( @@ -546,19 +470,10 @@ def test_research_save_to_file(self, runner, tmp_path): assert result.exit_code == 0 - # Check markdown file has content + # Check JSON file has output assert output_file.exists() - content = output_file.read_text() - assert "Research Report" in content - - # Check JSON file has metadata (not content) - json_file = tmp_path / "report.json" - assert json_file.exists() - data = json.loads(json_file.read_text()) + data = json.loads(output_file.read_text()) assert data["run_id"] == "trun_123" assert data["status"] == "completed" - assert "downloaded_at" in data - assert "files" in data - assert data["files"]["markdown"] == str(output_file) - assert "basis" in data # Should have basis metadata - assert "content" not in data # Content should NOT be in JSON (it's in markdown) + assert "output" in data + assert data["output"]["content"]["text"] == "Research findings" diff --git a/uv.lock b/uv.lock index 6741a64..f901fff 100644 --- a/uv.lock +++ b/uv.lock @@ -1057,7 +1057,7 @@ wheels = [ [[package]] name = "parallel-web-tools" -version = "0.0.1" +version = "0.0.3" source = { editable = "." } dependencies = [ { name = "pandas" }, From 5ad65d4bb50ab0eb8f4a6b9ed6a731520baa6272 Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 15:52:17 -0500 Subject: [PATCH 08/10] json --- parallel_web_tools/cli/commands.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/parallel_web_tools/cli/commands.py b/parallel_web_tools/cli/commands.py index 3fb2233..0ad3d61 100644 --- a/parallel_web_tools/cli/commands.py +++ b/parallel_web_tools/cli/commands.py @@ -680,8 +680,8 @@ def research(): @click.option("--timeout", type=int, default=3600, show_default=True, help="Max wait time in seconds") @click.option("--poll-interval", type=int, default=45, show_default=True, help="Seconds between status checks") @click.option("--no-wait", is_flag=True, help="Return immediately after creating task (don't poll)") -@click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to file (markdown)") -@click.option("--json", "output_json", is_flag=True, help="Output as JSON") +@click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to JSON file") +@click.option("--json", "output_json", is_flag=True, help="Output JSON to stdout") def research_run( query: str | None, input_file: str | None, @@ -700,7 +700,7 @@ def research_run( parallel-cli research run "What are the latest developments in quantum computing?" - parallel-cli research run -f question.txt --processor ultra -o report.md + parallel-cli research run -f question.txt --processor ultra -o report.json """ # Get query from argument or file if input_file: @@ -802,8 +802,8 @@ def research_status(run_id: str, output_json: bool): @click.argument("run_id") @click.option("--timeout", type=int, default=3600, show_default=True, help="Max wait time in seconds") @click.option("--poll-interval", type=int, default=45, show_default=True, help="Seconds between status checks") -@click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to file (markdown)") -@click.option("--json", "output_json", is_flag=True, help="Output as JSON") +@click.option("-o", "--output", "output_file", type=click.Path(), help="Save results to JSON file") +@click.option("--json", "output_json", is_flag=True, help="Output JSON to stdout") def research_poll( run_id: str, timeout: int, From 06107dc5a5fc05c3d700aee1de907954896a4672 Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 15:57:27 -0500 Subject: [PATCH 09/10] fix: add type assertions for pyrefly type checker Add assertions after validation to narrow optional types to non-None, fixing 10 pyrefly bad-argument-type errors in enrich_run and enrich_plan. --- parallel_web_tools/cli/commands.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/parallel_web_tools/cli/commands.py b/parallel_web_tools/cli/commands.py index 0ad3d61..042ed0b 100644 --- a/parallel_web_tools/cli/commands.py +++ b/parallel_web_tools/cli/commands.py @@ -480,7 +480,13 @@ def enrich_run( console.print(f"[bold cyan]Running enrichment from {config_file}...[/bold cyan]\n") run_enrichment(config_file) else: + # After validation, these are guaranteed non-None + assert source_type is not None + assert source is not None + assert target is not None + src_cols = parse_columns(source_columns) + assert src_cols is not None # Validated above if intent: console.print("[dim]Getting suggestions from Parallel API...[/dim]") @@ -490,6 +496,7 @@ def enrich_run( console.print(f"[green]AI suggested {len(enr_cols)} columns, processor: {final_processor}[/green]\n") else: enr_cols = parse_columns(enriched_columns) + assert enr_cols is not None # Validated above final_processor = processor or "core-fast" config = build_config_from_args( @@ -539,7 +546,12 @@ def enrich_plan( if has_cli_args: validate_enrich_args(source_type, source, target, source_columns, enriched_columns, intent) + # After validation, these are guaranteed non-None + assert source_type is not None + assert source is not None + assert target is not None src_cols = parse_columns(source_columns) + assert src_cols is not None # Validated above if intent: console.print("[dim]Getting suggestions from Parallel API...[/dim]") @@ -549,6 +561,7 @@ def enrich_plan( console.print(f"[green]AI suggested {len(enr_cols)} columns, processor: {final_processor}[/green]") else: enr_cols = parse_columns(enriched_columns) + assert enr_cols is not None # Validated above final_processor = processor or "core-fast" config = build_config_from_args( From dd9552f8504f4e47f0a087b01be7c6cd96b77937 Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Fri, 23 Jan 2026 16:13:13 -0500 Subject: [PATCH 10/10] fix: resolve pyrefly type checking errors - Add explicit urllib.error import in auth.py - Use SDK BetaRunInputParam type in batch.py and spark/streaming.py - Add explicit type annotations to fix type inference issues - Use row.asDict() for proper PySpark Row field access - Add type narrowing assertion in test_cli.py - Exclude non-core files from pyrefly (cloud_function, examples, notebooks, scripts) - Add pyrefly pre-commit hook (local, uses uv run) --- .pre-commit-config.yaml | 8 +++++++ parallel_web_tools/core/auth.py | 1 + parallel_web_tools/core/batch.py | 6 +++-- parallel_web_tools/core/schema.py | 22 +++++++++---------- .../integrations/spark/streaming.py | 11 ++++++---- pyproject.toml | 7 ++++++ tests/test_cli.py | 1 + 7 files changed, 38 insertions(+), 18 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 4f1e3f8..695274d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -14,3 +14,11 @@ repos: - id: ruff-check # Run the formatter. - id: ruff-format + - repo: local + hooks: + - id: pyrefly-check + name: pyrefly (type checking) + entry: uv run pyrefly check + language: system + pass_filenames: false + types: [python] diff --git a/parallel_web_tools/core/auth.py b/parallel_web_tools/core/auth.py index 8341ed3..c7d9da7 100644 --- a/parallel_web_tools/core/auth.py +++ b/parallel_web_tools/core/auth.py @@ -8,6 +8,7 @@ import secrets import socketserver import sys +import urllib.error import urllib.parse import urllib.request import webbrowser diff --git a/parallel_web_tools/core/batch.py b/parallel_web_tools/core/batch.py index 3047e27..0d75e3f 100644 --- a/parallel_web_tools/core/batch.py +++ b/parallel_web_tools/core/batch.py @@ -99,6 +99,7 @@ def enrich_batch( List of result dictionaries in same order as inputs. """ from parallel.types import JsonSchemaParam, TaskSpecParam + from parallel.types.beta import BetaRunInputParam if not inputs: return [] @@ -114,8 +115,8 @@ def enrich_batch( task_group = client.beta.task_group.create() taskgroup_id = task_group.task_group_id - # Add runs - run_inputs = [{"input": inp, "processor": processor} for inp in inputs] + # Add runs - use SDK type for proper typing + run_inputs: list[BetaRunInputParam] = [{"input": inp, "processor": processor} for inp in inputs] response = client.beta.task_group.add_runs( taskgroup_id, default_task_spec=task_spec, @@ -149,6 +150,7 @@ def enrich_batch( run_id = event.run.run_id if event.output and hasattr(event.output, "content"): content = event.output.content + result: dict[str, Any] if isinstance(content, dict): result = dict(content) elif isinstance(content, str): diff --git a/parallel_web_tools/core/schema.py b/parallel_web_tools/core/schema.py index f34f357..fbad505 100644 --- a/parallel_web_tools/core/schema.py +++ b/parallel_web_tools/core/schema.py @@ -153,17 +153,15 @@ def parse_input_and_output_models( schema: InputSchema, ) -> tuple[type[BaseModel], type[BaseModel]]: """Create Pydantic models from schema.""" - InputModel = create_model( - "InputModel", - **{col.name: (str, Field(description=col.description)) for col in schema.source_columns}, - ) - - OutputModel = create_model( - "OutputModel", - **{ - col.name: (TYPE_MAP.get(col.type, str), Field(description=col.description)) - for col in schema.enriched_columns - }, - ) + # Build field definitions with proper typing for create_model + input_fields: dict[str, Any] = { + col.name: (str, Field(description=col.description)) for col in schema.source_columns + } + output_fields: dict[str, Any] = { + col.name: (TYPE_MAP.get(col.type, str), Field(description=col.description)) for col in schema.enriched_columns + } + + InputModel = create_model("InputModel", **input_fields) + OutputModel = create_model("OutputModel", **output_fields) return InputModel, OutputModel diff --git a/parallel_web_tools/integrations/spark/streaming.py b/parallel_web_tools/integrations/spark/streaming.py index c384372..64df1ac 100644 --- a/parallel_web_tools/integrations/spark/streaming.py +++ b/parallel_web_tools/integrations/spark/streaming.py @@ -165,6 +165,7 @@ def enrich_streaming_batch( >>> query = stream_df.writeStream.foreachBatch(process_batch).start() """ from parallel.types import JsonSchemaParam, TaskSpecParam + from parallel.types.beta import BetaRunInputParam # Collect batch data (this is safe in micro-batches, which are already small) rows = batch_df.collect() @@ -193,16 +194,18 @@ def enrich_streaming_batch( taskgroup_id = task_group.task_group_id # Build inputs from batch rows - run_inputs = [] + run_inputs: list[BetaRunInputParam] = [] for row in rows: - input_data = {} + # Convert Row to dict for reliable field access + row_dict = row.asDict() + input_data: dict[str, object] = {} for parallel_name, col_name in input_columns.items(): # Get value from row, handle None gracefully - value = row[col_name] + value = row_dict.get(col_name) if value is not None: input_data[parallel_name] = str(value) - run_input = { + run_input: BetaRunInputParam = { "input": input_data, "processor": processor, } diff --git a/pyproject.toml b/pyproject.toml index 3238b42..f48848d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,6 +94,13 @@ project-includes = [ "**/*.py*", "**/*.ipynb", ] +# Exclude files that have external dependencies not in the main project +project-excludes = [ + "**/integrations/bigquery/cloud_function/**", # Separate deployment with flask, functions_framework + "scripts/runtime_hook_ssl.py", # PyInstaller runtime hook with sys._MEIPASS + "notebooks/**", # Jupyter notebooks with Databricks display() + "examples/**", # Example scripts +] [tool.ruff] line-length = 120 diff --git a/tests/test_cli.py b/tests/test_cli.py index e118e6f..7ceaf9d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -34,6 +34,7 @@ def test_parse_multiple_columns(self): """Should parse multiple columns.""" json_str = '[{"name": "a", "description": "A"}, {"name": "b", "description": "B"}]' result = parse_columns(json_str) + assert result is not None assert len(result) == 2 assert result[0]["name"] == "a" assert result[1]["name"] == "b"