From 55e9e0590ff160873d149a3feb7ac06c29c9d133 Mon Sep 17 00:00:00 2001 From: Matt Harris Date: Thu, 30 Apr 2026 17:19:25 -0400 Subject: [PATCH] fix: review findings from initial scaffold MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two real bugs caught by extended smoke + four parallel reviewers: Bugs (would have shipped broken in 0.0.1): - _enrich: client.task_run.execute(output=dict) raises TypeError; switched to client.beta.task_run.create(task_spec={"output_schema": {"type":"json", "json_schema": ...}}) + task_run.result(api_timeout=...) for the create+poll path. The dict envelope is required by the API contract. - extract: full_content=True was a top-level kwarg the SDK rejects. Now nested under advanced_settings={"full_content": True}. Code quality: - Tool names now public (web_search, web_fetch, extract, deep_research, enrich) instead of leading-underscore — LLMs handle clean names better. - Submodules renamed to _search, _extract, _task to avoid attribute/submodule shadowing in the package namespace. - ParallelTracingPlugin keys per-call timing on ToolContext.function_call_id (concurrency-safe vs the prior id(tool_args) approach which leaked on errors) and adds on_tool_error_callback to record failures. - __version__ now read from importlib.metadata (was hardcoded, drifted). - enrich passes api_timeout to result() instead of timeout to create(). - Dropped enrich_with_model helper (was untested + unused public API). Packaging / CI: - pyproject: parallel-web>=0.5.1 (was >=0.3, broken on older), google-adk>=1.31. - pyproject: Documentation URL points at GitHub README until docs page exists. - auto-release.yml: read commit message via env var, not direct interpolation (closes shell-injection vector via crafted commit messages). - release.sh: refuses to run off main, refuses if release tag already exists. - README + integration-docs: fixed StreamableHTTPServerParams (doesn't exist) to StreamableHTTPConnectionParams. Tests: - New test_search.py, test_extract.py, test_task.py with monkeypatched SDK client (string-target setattr to handle the namespace shadow). - test_tracing.py: added concurrent-call test and on_tool_error_callback test. - test_imports.py: parametrized tool name check; version asserted against importlib.metadata. - Moved smoke_live*.py to tests/integration/ with addopts ignore so default pytest no longer needs --ignore flags. - Live extended smoke (tests/integration/smoke_live_full.py) verifies _enrich end-to-end against api.parallel.ai + LLM picking web_fetch and extract. --- .github/workflows/auto-release.yml | 6 +- .github/workflows/ci.yml | 4 +- README.md | 8 +- pyproject.toml | 16 +- scripts/capture_screenshot.py | 4 +- scripts/release.sh | 28 ++- src/parallel_google_adk/__init__.py | 24 ++- .../{extract.py => _extract.py} | 17 +- .../{search.py => _search.py} | 24 +-- src/parallel_google_adk/_task.py | 122 +++++++++++ src/parallel_google_adk/task.py | 156 -------------- src/parallel_google_adk/tracing.py | 121 +++++++---- tests/integration/smoke_live_full.py | 201 ++++++++++++++++++ tests/smoke_live.py | 125 ----------- tests/test_extract.py | 63 ++++++ tests/test_imports.py | 34 ++- tests/test_search.py | 103 +++++++++ tests/test_task.py | 61 ++++++ tests/test_tracing.py | 68 ++++-- 19 files changed, 799 insertions(+), 386 deletions(-) rename src/parallel_google_adk/{extract.py => _extract.py} (63%) rename src/parallel_google_adk/{search.py => _search.py} (76%) create mode 100644 src/parallel_google_adk/_task.py delete mode 100644 src/parallel_google_adk/task.py create mode 100644 tests/integration/smoke_live_full.py delete mode 100644 tests/smoke_live.py create mode 100644 tests/test_extract.py create mode 100644 tests/test_search.py create mode 100644 tests/test_task.py diff --git a/.github/workflows/auto-release.yml b/.github/workflows/auto-release.yml index 980b1bd..7533378 100644 --- a/.github/workflows/auto-release.yml +++ b/.github/workflows/auto-release.yml @@ -22,8 +22,12 @@ jobs: steps: - name: Check if this is a version bump commit id: check + # Read commit message via env var, NOT direct ${{ }} interpolation in + # the run script — interpolating user-controllable text into a shell + # heredoc enables command injection via crafted commit messages. + env: + COMMIT_MSG: ${{ github.event.head_commit.message }} run: | - COMMIT_MSG="${{ github.event.head_commit.message }}" if [[ "$COMMIT_MSG" =~ ^chore:\ bump\ version\ to\ ([0-9]+\.[0-9]+\.[0-9]+(rc[0-9]+)?) ]]; then VERSION="${BASH_REMATCH[1]}" echo "should_release=true" >> "$GITHUB_OUTPUT" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c7a1e55..31be7a2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,7 +26,9 @@ jobs: pip install -e ".[dev]" - name: Run tests - run: pytest -v --ignore=tests/smoke_live.py + # Live-API integration tests under tests/integration/ are skipped via + # pytest addopts in pyproject.toml. + run: pytest -v lint: runs-on: ubuntu-latest diff --git a/README.md b/README.md index aad6cea..2e529ea 100644 --- a/README.md +++ b/README.md @@ -56,11 +56,12 @@ runner = Runner(agent=agent, app_name="my-app", plugins=[ParallelTracingPlugin() If you only want Search and don't want a Python dependency, the [Parallel Search MCP](https://docs.parallel.ai/integrations/mcp/search-mcp) works with ADK's `MCPToolset`: ```python +import os from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset -from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPServerParams +from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams MCPToolset( - connection_params=StreamableHTTPServerParams( + connection_params=StreamableHTTPConnectionParams( url="https://search.parallel.ai/mcp", headers={"Authorization": f"Bearer {os.environ['PARALLEL_API_KEY']}"}, ), @@ -77,7 +78,8 @@ See [`examples/research_agent.py`](examples/research_agent.py) for a runnable de ```bash pip install -e ".[dev]" -pytest +pytest # unit tests (no network) +pytest tests/integration/ # live API smoke (needs PARALLEL_API_KEY + GOOGLE_API_KEY) ``` ## License diff --git a/pyproject.toml b/pyproject.toml index fc09ce8..d241561 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "parallel-google-adk" -version = "0.1.0" +version = "0.0.1" description = "Google ADK FunctionTools and Plugin for Parallel APIs (search, extract, deep research)." readme = "README.md" requires-python = ">=3.10" @@ -23,8 +23,13 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "google-adk>=1.7", - "parallel-web>=0.3", + # google-adk 1.31.x is the version we test against; older minor versions + # have different plugin/MCPToolset APIs that this package doesn't support. + "google-adk>=1.31", + # parallel-web 0.5.x is required: client.search() expects max_results + # inside advanced_settings, and task_run.execute() / .result() shapes + # are used heavily. + "parallel-web>=0.5.1", "pydantic>=2.0", ] @@ -38,7 +43,7 @@ dev = [ [project.urls] Homepage = "https://parallel.ai" -Documentation = "https://docs.parallel.ai/integrations/google-adk" +Documentation = "https://github.com/parallel-web/parallel-google-adk#readme" Repository = "https://github.com/parallel-web/parallel-google-adk" Issues = "https://github.com/parallel-web/parallel-google-adk/issues" @@ -48,3 +53,6 @@ packages = ["src/parallel_google_adk"] [tool.pytest.ini_options] asyncio_mode = "auto" testpaths = ["tests"] +# Skip the live-API integration scripts on default `pytest` runs. Run them +# manually with: pytest tests/integration/ or `python tests/integration/`. +addopts = "--ignore=tests/integration" diff --git a/scripts/capture_screenshot.py b/scripts/capture_screenshot.py index c7bda59..d322d09 100644 --- a/scripts/capture_screenshot.py +++ b/scripts/capture_screenshot.py @@ -63,9 +63,7 @@ async def main() -> int: plugins=[ParallelTracingPlugin()], ) - prompt = ( - "What does Parallel's Search API do? Cite parallel.ai in one sentence." - ) + prompt = "What does Parallel's Search API do? Cite parallel.ai in one sentence." buf.write("$ python research_agent.py\n") buf.write(f"\n>>> {prompt}\n\n") diff --git a/scripts/release.sh b/scripts/release.sh index 1f05dbf..41ca15c 100755 --- a/scripts/release.sh +++ b/scripts/release.sh @@ -74,14 +74,40 @@ main() { die "working tree is not clean. Commit or stash changes first." fi - local current new branch + # Must run from main so the release branch forks from a clean upstream tip. + local current_branch + current_branch="$(git rev-parse --abbrev-ref HEAD)" + if [[ "$current_branch" != "main" ]]; then + die "must be on 'main' (currently on '$current_branch'). Run: git checkout main && git pull" + fi + + # Make sure we're up to date with origin/main so the release branch is + # forking from the latest tip — otherwise the version bump can land on a + # stale base. + git fetch origin main --quiet + if [[ "$(git rev-parse HEAD)" != "$(git rev-parse origin/main)" ]]; then + die "local main is not in sync with origin/main. Run: git pull --ff-only" + fi + + local current new branch tag current="$(get_current_version)" new="$(calculate_next_version "$current" "$bump")" branch="release/$new" + tag="v$new" + + # Refuse if the tag already exists locally OR on the remote — auto-release + # would otherwise double-tag and fail mid-flight. + if git rev-parse "$tag" >/dev/null 2>&1; then + die "tag $tag already exists locally" + fi + if git ls-remote --tags origin "$tag" | grep -q "refs/tags/$tag$"; then + die "tag $tag already exists on origin" + fi echo "current: $current" echo "new: $new" echo "branch: $branch" + echo "tag: $tag (will be created by auto-release.yml after merge)" read -rp "proceed? [y/N] " confirm [[ "$confirm" =~ ^[Yy]$ ]] || die "aborted" diff --git a/src/parallel_google_adk/__init__.py b/src/parallel_google_adk/__init__.py index 9fd6cc3..9295d1b 100644 --- a/src/parallel_google_adk/__init__.py +++ b/src/parallel_google_adk/__init__.py @@ -14,21 +14,27 @@ model="gemini-flash-latest", name="research_agent", tools=[web_search, web_fetch, extract, deep_research, enrich], - plugins=[ParallelTracingPlugin()], ) + runner = Runner(agent=agent, app_name="my-app", plugins=[ParallelTracingPlugin()]) -For Search alone, the Search MCP at https://search.parallel.ai/mcp is the simpler -default — see the integration page for the MCPToolset wiring. This package is the -typed-FunctionTool path for users who want fine-grained schemas, hidden polling, -and tracing. +For Search alone, the Search MCP at https://search.parallel.ai/mcp is a simpler +default — see the integration docs for the MCPToolset wiring. This package is +the typed-FunctionTool path for users who want fine-grained schemas, hidden +polling, structured enrichment, and tracing. """ -from .extract import extract -from .search import web_fetch, web_search -from .task import deep_research, enrich +from importlib.metadata import PackageNotFoundError, version + +from ._extract import extract +from ._search import web_fetch, web_search +from ._task import deep_research, enrich from .tracing import ParallelTracingPlugin -__version__ = "0.1.0" +try: + __version__ = version("parallel-google-adk") +except PackageNotFoundError: + # Editable install / source checkout without metadata. + __version__ = "0.0.0+unknown" __all__ = [ "ParallelTracingPlugin", diff --git a/src/parallel_google_adk/extract.py b/src/parallel_google_adk/_extract.py similarity index 63% rename from src/parallel_google_adk/extract.py rename to src/parallel_google_adk/_extract.py index b56478c..9273313 100644 --- a/src/parallel_google_adk/extract.py +++ b/src/parallel_google_adk/_extract.py @@ -9,21 +9,18 @@ from ._client import get_client -def _extract( +def extract( urls: Annotated[ list[str], - "URLs to extract clean content from. 1–20 per call. Fully-qualified " - "http(s) URLs only.", + "URLs to extract content from. 1–20 per call. Fully-qualified http(s) URLs only.", ], objective: Annotated[ str | None, - "Optional focus: what to keep relevant in the extraction. Improves " - "signal-to-noise on long pages. Omit for full extraction.", + "Optional focus: what to keep relevant. Improves signal-to-noise on long pages.", ] = None, full_content: Annotated[ bool, - "If True, returns the full cleaned page content. If False (default), " - "returns LLM-optimized excerpts only.", + "If True, returns full cleaned page content. If False (default), returns excerpts only.", ] = False, ) -> dict[str, Any]: """Extract clean content from web pages and PDFs in batch. @@ -36,10 +33,12 @@ def _extract( kwargs: dict[str, Any] = {"urls": list(urls)} if objective is not None: kwargs["objective"] = objective + # full_content goes inside advanced_settings, not at the top level — + # the SDK rejects it as an unexpected kwarg otherwise. if full_content: - kwargs["full_content"] = True + kwargs["advanced_settings"] = {"full_content": True} response = client.extract(**kwargs) return response.model_dump() -extract = FunctionTool(_extract) +extract = FunctionTool(extract) diff --git a/src/parallel_google_adk/search.py b/src/parallel_google_adk/_search.py similarity index 76% rename from src/parallel_google_adk/search.py rename to src/parallel_google_adk/_search.py index 22c5a14..37fa22e 100644 --- a/src/parallel_google_adk/search.py +++ b/src/parallel_google_adk/_search.py @@ -15,16 +15,14 @@ from ._client import get_client -def _web_search( +def web_search( objective: Annotated[ str, - "Natural-language description of what you want to find. Be specific — " - "the model uses this to score and select the most relevant excerpts.", + "Natural-language description of what you want to find.", ], queries: Annotated[ list[str], - "1–5 keyword search queries (3–6 words each). Each query is run " - "independently; results are merged and re-ranked.", + "1–5 keyword search queries (3–6 words each).", ], max_results: Annotated[ int, @@ -32,15 +30,14 @@ def _web_search( ] = 5, max_chars_total: Annotated[ int | None, - "Optional ceiling on total characters across all excerpts. Useful for " - "controlling token budget on long-running agent loops.", + "Optional ceiling on total characters across all excerpts.", ] = None, ) -> dict[str, Any]: """Search the web for current, citation-aware information. Returns LLM-optimized excerpts with source URLs. Use this when you need - grounded facts the model wouldn't know, or when you don't already have - URLs to fetch. Prefer web_fetch when you have specific URLs in hand. + grounded facts or don't have specific URLs in hand. Prefer web_fetch when + you already know which URL to read. """ client = get_client() kwargs: dict[str, Any] = { @@ -54,12 +51,11 @@ def _web_search( return response.model_dump() -def _web_fetch( +def web_fetch( url: Annotated[str, "The URL to fetch. Must be a fully-qualified http(s) URL."], objective: Annotated[ str | None, - "Optional focus for the extraction — what to keep relevant. If omitted, " - "returns the full cleaned content.", + "Optional focus for the extraction — what to keep relevant. Omit for full extraction.", ] = None, ) -> dict[str, Any]: """Fetch a single URL and return clean, LLM-ready content. @@ -82,5 +78,5 @@ def _web_fetch( ) -web_search = FunctionTool(_web_search) -web_fetch = FunctionTool(_web_fetch) +web_search = FunctionTool(web_search) +web_fetch = FunctionTool(web_fetch) diff --git a/src/parallel_google_adk/_task.py b/src/parallel_google_adk/_task.py new file mode 100644 index 0000000..646032b --- /dev/null +++ b/src/parallel_google_adk/_task.py @@ -0,0 +1,122 @@ +"""Parallel Task API as ADK FunctionTools. + +Two tools: + - deep_research: long-running multi-hop research with cited results + - enrich: structured-output extraction against a JSON Schema + +``deep_research`` uses ``client.task_run.execute(...)`` which handles create + +poll internally and accepts no output schema (returns a markdown report). + +``enrich`` needs to pass a JSON Schema dict, which ``execute()`` rejects (it +only accepts Pydantic classes / strings); so it routes via +``client.beta.task_run.create(task_spec={"output_schema": {"type": "json", "json_schema": ...}})`` +followed by ``client.task_run.result(run_id, api_timeout=...)`` for polling. + +Default processor is ``pro-fast`` for deep research (good speed/quality +tradeoff); override for the slower ``pro`` (2–10 min) or ``ultra`` (5–25 min). +""" + +from __future__ import annotations + +import json +from typing import Annotated, Any + +from google.adk.tools import FunctionTool + +from ._client import get_client + +DEFAULT_DEEP_RESEARCH_PROCESSOR = "pro-fast" +DEFAULT_ENRICH_PROCESSOR = "core" +DEFAULT_TIMEOUT_S = 1200.0 # 20 minutes + + +def _format_result(result: Any) -> dict[str, Any]: + """Convert SDK TaskRunResult / ParsedTaskRunResult to a plain dict.""" + payload = result.model_dump() if hasattr(result, "model_dump") else dict(result) + parsed = getattr(result, "parsed", None) + if parsed is not None: + payload["parsed"] = ( + parsed.model_dump() if hasattr(parsed, "model_dump") else parsed + ) + return payload + + +def deep_research( + objective: Annotated[ + str, + "Natural-language description of what to research. Be specific.", + ], + processor: Annotated[ + str, + "'pro-fast' (default, ~1–3 min), 'pro' (2–10 min), or 'ultra' (5–25 min).", + ] = DEFAULT_DEEP_RESEARCH_PROCESSOR, + timeout_seconds: Annotated[ + float, + "Maximum time to wait. Defaults to 20 minutes.", + ] = DEFAULT_TIMEOUT_S, +) -> dict[str, Any]: + """Run a multi-hop deep research task with cited results. + + Returns a comprehensive report with per-claim citations. Use for thorough + investigations — not quick fact lookups (use web_search for those). + """ + client = get_client() + result = client.task_run.execute( + input=objective, + processor=processor, + timeout=timeout_seconds, + ) + return _format_result(result) + + +def enrich( + objective: Annotated[ + str, + "What to extract about each input. e.g. 'Find the company's CEO and HQ city.'", + ], + inputs_json: Annotated[ + str, + 'JSON-encoded list of entity dicts. Example: \'[{"name":"Acme","url":"acme.com"}]\'.', + ], + output_schema_json: Annotated[ + str, + 'JSON-encoded JSON Schema for the per-input output. Example: \'{"type":"object","properties":{"ceo":{"type":"string"}}}\'.', + ], + processor: Annotated[ + str, + "'core' (default), 'pro' for harder questions, or 'lite' for cheap/fast runs.", + ] = DEFAULT_ENRICH_PROCESSOR, + timeout_seconds: Annotated[float, "Max wait time."] = DEFAULT_TIMEOUT_S, +) -> dict[str, Any]: + """Enrich a list of entities with structured, cited fields. + + Use this when you want typed output (a schema-conforming dict) rather than + a markdown report. Both inputs_json and output_schema_json are JSON strings. + """ + try: + inputs = json.loads(inputs_json) + except json.JSONDecodeError as exc: + raise ValueError(f"inputs_json must be valid JSON: {exc}") from exc + try: + schema = json.loads(output_schema_json) + except json.JSONDecodeError as exc: + raise ValueError(f"output_schema_json must be valid JSON: {exc}") from exc + + client = get_client() + payload = {"objective": objective, "inputs": inputs} + # task_run.execute() rejects raw JSON Schema dicts (it accepts only Pydantic + # classes or strings), so we use the create() + result() path. The schema + # must be wrapped in the discriminated-union envelope: + # {"type": "json", "json_schema": }. + run = client.beta.task_run.create( + input=payload, + processor=processor, + task_spec={"output_schema": {"type": "json", "json_schema": schema}}, + ) + # Poll for the result. api_timeout controls the SDK long-poll window. + result = client.task_run.result(run.run_id, api_timeout=int(timeout_seconds)) + return _format_result(result) + + +deep_research = FunctionTool(deep_research) +enrich = FunctionTool(enrich) diff --git a/src/parallel_google_adk/task.py b/src/parallel_google_adk/task.py deleted file mode 100644 index a50aade..0000000 --- a/src/parallel_google_adk/task.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Parallel Task API as ADK FunctionTools. - -Two tools: - - deep_research: long-running multi-hop research with cited results - - enrich: structured-output extraction against a JSON Schema or Pydantic model - -Both use the SDK's ``client.task_run.execute(...)`` which handles create + poll -internally, so the FunctionTool returns a finished result in one call. Default -processor is ``pro-fast`` for deep research (good speed/quality tradeoff); -override for the slower-but-more-thorough ``pro`` (2–10 min) or ``ultra`` -(5–25 min) processors. -""" - -from __future__ import annotations - -import json -from typing import Annotated, Any - -from google.adk.tools import FunctionTool -from pydantic import BaseModel - -from ._client import get_client - -DEFAULT_DEEP_RESEARCH_PROCESSOR = "pro-fast" -DEFAULT_ENRICH_PROCESSOR = "core" -DEFAULT_TIMEOUT_S = 1200.0 # 20 minutes - - -def _format_result(result: Any) -> dict[str, Any]: - """Convert SDK TaskRunResult / ParsedTaskRunResult to a plain dict. - - Surfaces run id, output, basis (citations + reasoning), and lifts the - parsed pydantic instance to a top-level ``parsed`` field if present. - """ - payload = result.model_dump() if hasattr(result, "model_dump") else dict(result) - parsed = getattr(result, "parsed", None) - if parsed is not None: - payload["parsed"] = ( - parsed.model_dump() if hasattr(parsed, "model_dump") else parsed - ) - return payload - - -def _deep_research( - objective: Annotated[ - str, - "Natural-language description of what to research. Be specific about " - "the question, the depth wanted, and any constraints (recency, sources, " - "geography).", - ], - processor: Annotated[ - str, - "Task processor. 'pro-fast' (default, ~1–3 min) for fast research; " - "'pro' (2–10 min) for thorough; 'ultra' (5–25 min) for the most " - "comprehensive investigative reports.", - ] = DEFAULT_DEEP_RESEARCH_PROCESSOR, - timeout_seconds: Annotated[ - float, - "Maximum time to wait for the task to complete. Defaults to 20 minutes.", - ] = DEFAULT_TIMEOUT_S, -) -> dict[str, Any]: - """Run a multi-hop deep research task with cited results. - - Returns a comprehensive report with per-claim citations. Use this when the - user wants a thorough investigation — not just a quick fact lookup. For - fact lookups use web_search; for structured extraction use enrich. - """ - client = get_client() - result = client.task_run.execute( - input=objective, - processor=processor, - timeout=timeout_seconds, - ) - return _format_result(result) - - -def _enrich( - objective: Annotated[ - str, - "What to extract about each input. e.g. 'Find the company's CEO, " - "headquarters city, and last funding round.'", - ], - inputs_json: Annotated[ - str, - "JSON-encoded list of entities to enrich. Each entry is an object with " - 'keys like name/url/domain. Example: \'[{"name":"Acme","url":"acme.com"}]\'.', - ], - output_schema_json: Annotated[ - str, - "JSON-encoded JSON Schema describing the desired per-input output shape. " - 'Example: \'{"type":"object","properties":{"ceo":{"type":"string"}},"required":["ceo"]}\'.', - ], - processor: Annotated[ - str, - "Task processor. 'core' (default) is the standard enrichment processor; " - "'pro' for harder questions; 'lite' for cheaper/faster runs.", - ] = DEFAULT_ENRICH_PROCESSOR, - timeout_seconds: Annotated[float, "Max wait time."] = DEFAULT_TIMEOUT_S, -) -> dict[str, Any]: - """Enrich a list of entities with structured, cited fields. - - Use this when you want typed output (a schema-conforming dict) rather than - a markdown report. Common use: pulling firmographic fields about companies, - biographic fields about people, or product specs. - - Both inputs_json and output_schema_json are passed as JSON strings — the - tool decodes them before calling the Task API. - """ - try: - inputs = json.loads(inputs_json) - except json.JSONDecodeError as exc: - raise ValueError(f"inputs_json must be valid JSON: {exc}") from exc - try: - schema = json.loads(output_schema_json) - except json.JSONDecodeError as exc: - raise ValueError(f"output_schema_json must be valid JSON: {exc}") from exc - - client = get_client() - payload = { - "objective": objective, - "inputs": inputs, - } - result = client.task_run.execute( - input=payload, - processor=processor, - output=schema, - timeout=timeout_seconds, - ) - return _format_result(result) - - -def enrich_with_model( - objective: str, - inputs: list[dict[str, Any]], - output_model: type[BaseModel], - processor: str = DEFAULT_ENRICH_PROCESSOR, - timeout_seconds: float = DEFAULT_TIMEOUT_S, -) -> dict[str, Any]: - """Pythonic enrich helper that accepts a Pydantic model class. - - Not exposed as a FunctionTool (LLMs can't pass class references), but useful - for direct programmatic calls from agent harness code. - """ - client = get_client() - payload = {"objective": objective, "inputs": inputs} - result = client.task_run.execute( - input=payload, - processor=processor, - output=output_model, - timeout=timeout_seconds, - ) - return _format_result(result) - - -deep_research = FunctionTool(_deep_research) -enrich = FunctionTool(_enrich) diff --git a/src/parallel_google_adk/tracing.py b/src/parallel_google_adk/tracing.py index a8b984a..b7a7ed8 100644 --- a/src/parallel_google_adk/tracing.py +++ b/src/parallel_google_adk/tracing.py @@ -1,12 +1,20 @@ """Optional ADK Plugin that traces Parallel tool calls. Captures, per session: tool name, latency, citation count, and (when present -in the tool result) cost. Stored on the invocation context state so downstream -observability/Plugin layers can read it. +in the tool result) cost. Stored on the session state so downstream +observability layers can read it. Keep narrow — no policy, no quota, no side-effecting export. If you want -those, layer your own Plugin on top or pipe ctx.state["_parallel_calls"] to +those, layer your own Plugin on top or pipe ``state["_parallel_calls"]`` to your existing observability stack. + +Usage: + runner = Runner(agent=agent, app_name="my-app", plugins=[ParallelTracingPlugin()]) + +After a run:: + final_session = await session_service.get_session(...) + calls = final_session.state.get("_parallel_calls", []) + # [{"tool": "web_search", "latency_s": 1.16, "citation_count": 5}, ...] """ from __future__ import annotations @@ -15,27 +23,26 @@ from typing import Any from google.adk.plugins.base_plugin import BasePlugin +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.tool_context import ToolContext -PARALLEL_TOOL_NAMES = { - "_web_search", - "_web_fetch", - "_extract", - "_deep_research", - "_enrich", - # Be permissive about the public names too in case ADK uses them. - "web_search", - "web_fetch", - "extract", - "deep_research", - "enrich", -} +# Public tool names registered by parallel_google_adk. Tracing applies to these +# only — other tools the agent calls are passed through untouched. +PARALLEL_TOOL_NAMES = frozenset( + {"web_search", "web_fetch", "extract", "deep_research", "enrich"} +) class ParallelTracingPlugin(BasePlugin): """Traces calls to Parallel FunctionTools. - Reads/writes ``ctx.state["_parallel_calls"]`` — a list of dicts with - keys: tool, latency_s, citation_count, cost_usd (optional), error (optional). + Records latency, citation count, and (when reported) cost in + ``tool_context.state["_parallel_calls"]`` — a list of dicts with keys: + ``tool``, ``latency_s``, ``citation_count``, ``cost_usd`` (optional), + ``error`` (optional). + + Per-call timing is keyed on ``ToolContext.function_call_id`` so concurrent + sessions sharing one Plugin instance don't collide. """ def __init__( @@ -45,28 +52,31 @@ def __init__( ) -> None: super().__init__(name=name) self._state_key = state_key - self._starts: dict[int, float] = {} + self._starts: dict[str, float] = {} async def before_tool_callback( - self, *, tool: Any, tool_args: dict[str, Any], **_: Any + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, ) -> None: - if getattr(tool, "name", None) not in PARALLEL_TOOL_NAMES: + if tool.name not in PARALLEL_TOOL_NAMES: return - self._starts[id(tool_args)] = time.monotonic() + key = self._call_key(tool_context, tool) + self._starts[key] = time.monotonic() async def after_tool_callback( self, *, - tool: Any, + tool: BaseTool, tool_args: dict[str, Any], - result: Any, - tool_context: Any = None, - **_: Any, + tool_context: ToolContext, + result: dict, ) -> None: - if getattr(tool, "name", None) not in PARALLEL_TOOL_NAMES: + if tool.name not in PARALLEL_TOOL_NAMES: return - started_at = self._starts.pop(id(tool_args), None) - latency_s = time.monotonic() - started_at if started_at is not None else None + latency_s = self._consume_latency(tool_context, tool) entry: dict[str, Any] = { "tool": tool.name, @@ -76,10 +86,46 @@ async def after_tool_callback( cost = _extract_cost(result) if cost is not None: entry["cost_usd"] = cost + self._append_entry(tool_context, entry) - state = ( - getattr(tool_context, "state", None) if tool_context is not None else None + async def on_tool_error_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + error: Exception, + ) -> None: + if tool.name not in PARALLEL_TOOL_NAMES: + return + latency_s = self._consume_latency(tool_context, tool) + self._append_entry( + tool_context, + { + "tool": tool.name, + "latency_s": latency_s, + "error": repr(error), + }, ) + + # ----- internal helpers ----- + + @staticmethod + def _call_key(tool_context: ToolContext, tool: BaseTool) -> str: + # function_call_id is unique per LLM-emitted tool call. Falls back to + # a less-unique key if ADK ever omits it. + call_id = getattr(tool_context, "function_call_id", None) + return f"{call_id}:{tool.name}" if call_id else f"_:{tool.name}" + + def _consume_latency( + self, tool_context: ToolContext, tool: BaseTool + ) -> float | None: + key = self._call_key(tool_context, tool) + started_at = self._starts.pop(key, None) + return time.monotonic() - started_at if started_at is not None else None + + def _append_entry(self, tool_context: ToolContext, entry: dict[str, Any]) -> None: + state = getattr(tool_context, "state", None) if state is None: return bucket = state.get(self._state_key) or [] @@ -90,17 +136,14 @@ async def after_tool_callback( def _count_citations(result: Any) -> int: """Best-effort citation count across the API response shapes we wrap. - - Search API: each ``results[]`` entry carries a ``url`` — count those. - - Extract API: same — ``results[]`` with ``url``. - - Task API (deep_research / enrich): citations live under ``output.basis`` - (list of basis blocks) or each entry's ``citations`` field. + - Search / Extract API: each ``results[]`` entry carries a ``url`` — count those. + - web_fetch (single-URL): one envelope dict with ``url`` at top — count as 1. + - Task API (deep_research / enrich): citations live under ``output.basis``. """ if not isinstance(result, dict): return 0 - # Direct top-level ``citations`` (rare, but handle it). if isinstance(result.get("citations"), list): return len(result["citations"]) - # Search / Extract: count URL-bearing results. results_list = result.get("results") if isinstance(results_list, list): url_count = sum(1 for r in results_list if isinstance(r, dict) and r.get("url")) @@ -110,7 +153,8 @@ def _count_citations(result: Any) -> int: len(r.get("citations", [])) if isinstance(r, dict) else 0 for r in results_list ) - # Task API: ``output.basis`` is a list of basis entries with citations. + if isinstance(result.get("url"), str): + return 1 output = result.get("output") if isinstance(output, dict): basis = output.get("basis") @@ -121,7 +165,6 @@ def _count_citations(result: Any) -> int: total += len(entry["citations"]) if total: return total - # Legacy / older envelopes with top-level basis. basis = result.get("basis") if isinstance(basis, dict) and isinstance(basis.get("citations"), list): return len(basis["citations"]) diff --git a/tests/integration/smoke_live_full.py b/tests/integration/smoke_live_full.py new file mode 100644 index 0000000..5360f65 --- /dev/null +++ b/tests/integration/smoke_live_full.py @@ -0,0 +1,201 @@ +"""End-to-end smoke test exercising every tool + the tracing plugin shapes. + +Coverage gaps this fills (vs the lighter tests/smoke_live.py): + - _enrich() against the live Task API (was never called) + - LLM agent picking web_fetch and extract (only web_search was exercised before) + - Tracing plugin citation_count heuristic on Task API output shape + +Skips deep_research (already verified, slow + expensive). Set RUN_DEEP=1 to include. + +Run: + PARALLEL_API_KEY=... .venv/bin/python tests/smoke_live_full.py +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +import traceback + +from google.adk.agents import LlmAgent +from google.adk.runners import Runner +from google.adk.sessions import InMemorySessionService +from google.genai import types + +from parallel_google_adk import ( + ParallelTracingPlugin, + deep_research, + enrich, + extract, + web_fetch, + web_search, +) +from parallel_google_adk.tracing import _count_citations + +# The public re-exports are FunctionTool wrappers; .func gives the inner +# Python function so we can call it directly without going through ADK. +_enrich = enrich.func + + +def _line(label: str) -> None: + print(f"\n--- {label} ---") + + +async def _run_agent_prompt(prompt: str, expected_tool: str) -> tuple[str, list]: + """Run a single prompt through Gemini Flash and return final text + trace.""" + agent = LlmAgent( + model="gemini-flash-latest", + name="research_agent", + instruction=( + "You are a research assistant. Use the most appropriate tool. " + "Always cite sources with their URL." + ), + tools=[web_search, web_fetch, extract, deep_research, enrich], + ) + session_service = InMemorySessionService() + session = await session_service.create_session(app_name="smoke", user_id="u") + runner = Runner( + agent=agent, + app_name="smoke", + session_service=session_service, + plugins=[ParallelTracingPlugin()], + ) + user_message = types.Content(role="user", parts=[types.Part(text=prompt)]) + final_text = "" + async for event in runner.run_async( + user_id="u", session_id=session.id, new_message=user_message + ): + if event.is_final_response() and event.content and event.content.parts: + for part in event.content.parts: + if part.text: + final_text += part.text + final_session = await session_service.get_session( + app_name="smoke", user_id="u", session_id=session.id + ) + calls = final_session.state.get("_parallel_calls", []) if final_session else [] + print(f"prompt: {prompt[:80]}...") + print(f"expected tool: {expected_tool}") + print(f"trace: {calls}") + return final_text, calls + + +async def main() -> int: + if "PARALLEL_API_KEY" not in os.environ: + print("PARALLEL_API_KEY not set; aborting.", file=sys.stderr) + return 1 + + failures: list[str] = [] + + # 1. Direct _enrich() — first time hitting the live Task API enrich path. + _line("_enrich (direct, lite processor)") + try: + result = _enrich( + objective="Find the company's official website URL.", + inputs_json=json.dumps( + [{"name": "Parallel Web Systems"}, {"name": "Anthropic"}] + ), + output_schema_json=json.dumps( + { + "type": "object", + "properties": { + "website": {"type": "string", "description": "Official URL"} + }, + "required": ["website"], + } + ), + processor="lite", + timeout_seconds=300, + ) + # Result shape varies; just assert we got something useful back. + run_id = result.get("run_id") or (result.get("run") or {}).get("run_id") + output = result.get("output") or result.get("parsed") + print(f"run_id: {run_id}") + if output: + preview = ( + json.dumps(output)[:300] + if not isinstance(output, str) + else output[:300] + ) + print(f"output[0:300]: {preview}") + if not run_id and not output: + failures.append("_enrich: empty result") + else: + # Also verify citation_count heuristic on Task API shape. + citations = _count_citations(result) + print(f"_count_citations on Task result: {citations}") + except Exception as e: + failures.append(f"_enrich: {e}") + traceback.print_exc() + + web_tools = {"web_search", "web_fetch", "extract"} + + # 2. LLM agent — web_fetch + _line("agent: should pick web_fetch") + try: + text, trace = await _run_agent_prompt( + "Fetch the page at https://parallel.ai/ and summarize the headline in one sentence.", + expected_tool="web_fetch", + ) + tools_called = [c.get("tool") for c in trace] + if not any(c.get("tool") in web_tools for c in trace): + failures.append( + f"agent web_fetch: no relevant tool called (got {tools_called})" + ) + for call in trace: + if call.get("citation_count", 0) == 0 and call.get("tool") in web_tools: + failures.append( + f"agent web_fetch: tracing citation_count=0 for {call['tool']}" + ) + except Exception as e: + failures.append(f"agent web_fetch: {e}") + traceback.print_exc() + + # 3. LLM agent — extract (multi-URL hint) + _line("agent: should pick extract") + try: + text, trace = await _run_agent_prompt( + "Extract the main content from these two URLs and tell me what they have in common: " + "https://parallel.ai/ and https://docs.parallel.ai/", + expected_tool="extract", + ) + tools_called = [c.get("tool") for c in trace] + if not any(c.get("tool") in web_tools for c in trace): + failures.append( + f"agent extract: no relevant tool called (got {tools_called})" + ) + except Exception as e: + failures.append(f"agent extract: {e}") + traceback.print_exc() + + # 4. Optional — deep_research via agent (slow + expensive). + if os.environ.get("RUN_DEEP") == "1": + _line("agent: should pick deep_research") + try: + text, trace = await _run_agent_prompt( + "Do deep research: what are the three core APIs Parallel offers, with a " + "one-sentence description of each? Use the deep_research tool.", + expected_tool="deep_research", + ) + if not any(c.get("tool") == "deep_research" for c in trace): + tools_called = [c.get("tool") for c in trace] + failures.append(f"agent deep_research: not called (got {tools_called})") + except Exception as e: + failures.append(f"agent deep_research: {e}") + traceback.print_exc() + else: + print("\n(skipping agent deep_research; set RUN_DEEP=1 to enable)") + + print("\n=========================") + if failures: + print(f"FAILED: {len(failures)} check(s)") + for f in failures: + print(f" - {f}") + return 2 + print("All extended smoke tests passed.") + return 0 + + +if __name__ == "__main__": + sys.exit(asyncio.run(main())) diff --git a/tests/smoke_live.py b/tests/smoke_live.py deleted file mode 100644 index 6eb463e..0000000 --- a/tests/smoke_live.py +++ /dev/null @@ -1,125 +0,0 @@ -"""Live smoke test against api.parallel.ai. - -Run manually: - PARALLEL_API_KEY=... .venv/bin/python tests/smoke_live.py -""" - -from __future__ import annotations - -import json -import os -import sys -import traceback - - -def _line(label: str) -> None: - print(f"\n--- {label} ---") - - -def _summary(name: str, result: dict) -> None: - if ( - isinstance(result, dict) - and "results" in result - and isinstance(result["results"], list) - ): - print(f"{name}: ok ({len(result['results'])} results)") - elif isinstance(result, dict) and "url" in result: - print(f"{name}: ok (url={result['url']!r})") - elif isinstance(result, dict): - keys = list(result.keys())[:5] - print(f"{name}: ok (keys={keys})") - else: - print(f"{name}: ok ({type(result).__name__})") - - -def main() -> int: - if "PARALLEL_API_KEY" not in os.environ: - print("PARALLEL_API_KEY not set; skipping live test.", file=sys.stderr) - return 1 - - failures: list[str] = [] - - from parallel_google_adk.search import _web_fetch, _web_search - from parallel_google_adk.extract import _extract - from parallel_google_adk.task import _deep_research - - _line("web_search") - try: - result = _web_search( - objective="What is the Parallel AI web research API?", - queries=["parallel.ai search API", "parallel AI web research"], - max_results=3, - ) - _summary("web_search", result) - # Print one excerpt for sanity. - results_list = result.get("results") or [] - if results_list: - first = results_list[0] - print(f" first.url: {first.get('url')}") - excerpts = first.get("excerpts") or [] - if excerpts: - preview = excerpts[0][:120].replace("\n", " ") - print(f" first.excerpt[0:120]: {preview}...") - except Exception as e: - failures.append(f"web_search: {e}") - traceback.print_exc() - - _line("web_fetch") - try: - result = _web_fetch(url="https://parallel.ai/") - _summary("web_fetch", result) - print(f" url: {result.get('url')}") - print(f" title: {result.get('title')}") - except Exception as e: - failures.append(f"web_fetch: {e}") - traceback.print_exc() - - _line("extract (multi-URL)") - try: - result = _extract( - urls=["https://parallel.ai/", "https://docs.parallel.ai/"], - objective="What products does Parallel offer?", - ) - _summary("extract", result) - for r in (result.get("results") or [])[:2]: - print(f" - {r.get('url')} (excerpts={len((r.get('excerpts') or []))})") - except Exception as e: - failures.append(f"extract: {e}") - traceback.print_exc() - - # Deep research is slow + costs $$. Only run if explicitly opted in. - if os.environ.get("RUN_DEEP_RESEARCH") == "1": - _line("deep_research (pro-fast)") - try: - result = _deep_research( - objective="In one paragraph: what does Parallel AI's Search API do? Cite parallel.ai.", - processor="pro-fast", - timeout_seconds=300, - ) - _summary("deep_research", result) - output = result.get("output") or result.get("parsed") - if output: - preview = ( - json.dumps(output)[:300] - if not isinstance(output, str) - else output[:300] - ) - print(f" output[0:300]: {preview}...") - except Exception as e: - failures.append(f"deep_research: {e}") - traceback.print_exc() - else: - print("\n(skipping deep_research; set RUN_DEEP_RESEARCH=1 to enable)") - - print("\n=========================") - if failures: - print(f"FAILED: {len(failures)} tool(s)") - for f in failures: - print(f" - {f}") - return 2 - print("All smoke tests passed.") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/test_extract.py b/tests/test_extract.py new file mode 100644 index 0000000..d497504 --- /dev/null +++ b/tests/test_extract.py @@ -0,0 +1,63 @@ +"""Unit tests for extract.py — verify SDK call shape including full_content nesting.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from parallel_google_adk import extract as extract_tool + + +class _FakeResponse: + def __init__(self, payload: dict[str, Any]) -> None: + self._payload = payload + + def model_dump(self) -> dict[str, Any]: + return self._payload + + +def _patch_client(monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]: + captured: dict[str, Any] = {} + + def fake_extract(**kwargs: Any) -> _FakeResponse: + captured.update(kwargs) + return _FakeResponse({"results": [{"url": u} for u in kwargs["urls"]]}) + + fake_client = SimpleNamespace(extract=fake_extract) + # String target resolves via sys.modules — bypasses the package's + # re-exported `extract` FunctionTool that would otherwise shadow the + # submodule under attribute access. + monkeypatch.setattr("parallel_google_adk._extract.get_client", lambda: fake_client) + return captured + + +def test_extract_default_call_shape(monkeypatch: pytest.MonkeyPatch) -> None: + captured = _patch_client(monkeypatch) + extract_tool.func(urls=["https://a", "https://b"]) + + assert captured["urls"] == ["https://a", "https://b"] + assert "objective" not in captured + # full_content defaults to False -> advanced_settings must NOT be added. + assert "advanced_settings" not in captured + + +def test_extract_passes_objective(monkeypatch: pytest.MonkeyPatch) -> None: + captured = _patch_client(monkeypatch) + extract_tool.func(urls=["https://a"], objective="What products do they sell?") + assert captured["objective"] == "What products do they sell?" + + +def test_extract_full_content_nests_under_advanced_settings( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression: full_content must be inside advanced_settings. + + The SDK rejects ``client.extract(full_content=True)`` as an unexpected + kwarg — the value lives in the advanced_settings envelope instead. + """ + captured = _patch_client(monkeypatch) + extract_tool.func(urls=["https://a"], full_content=True) + assert "full_content" not in captured # NOT a top-level kwarg + assert captured["advanced_settings"] == {"full_content": True} diff --git a/tests/test_imports.py b/tests/test_imports.py index 459772c..c4c77d0 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -2,6 +2,10 @@ from __future__ import annotations +from importlib.metadata import version as pkg_version + +import pytest + def test_public_imports() -> None: from parallel_google_adk import ( @@ -13,15 +17,15 @@ def test_public_imports() -> None: web_search, ) - # Tools should be FunctionTool instances with the right name attribute. + # Tools should be FunctionTool instances with public, LLM-friendly names. from google.adk.tools import FunctionTool for tool, expected_name in [ - (web_search, "_web_search"), - (web_fetch, "_web_fetch"), - (extract, "_extract"), - (deep_research, "_deep_research"), - (enrich, "_enrich"), + (web_search, "web_search"), + (web_fetch, "web_fetch"), + (extract, "extract"), + (deep_research, "deep_research"), + (enrich, "enrich"), ]: assert isinstance(tool, FunctionTool) assert tool.name == expected_name @@ -30,7 +34,21 @@ def test_public_imports() -> None: assert plugin.name == "parallel-tracing" -def test_version() -> None: +@pytest.mark.parametrize( + "tool_name", + ["web_search", "web_fetch", "extract", "deep_research", "enrich"], +) +def test_tool_name_in_tracing_set(tool_name: str) -> None: + """Every public tool must be recognized by the tracing plugin.""" + from parallel_google_adk.tracing import PARALLEL_TOOL_NAMES + + assert tool_name in PARALLEL_TOOL_NAMES + + +def test_version_matches_metadata() -> None: + """__version__ should reflect the installed package metadata.""" import parallel_google_adk - assert parallel_google_adk.__version__ == "0.1.0" + # When installed editable, importlib.metadata returns pyproject's version. + expected = pkg_version("parallel-google-adk") + assert parallel_google_adk.__version__ == expected diff --git a/tests/test_search.py b/tests/test_search.py new file mode 100644 index 0000000..79571ed --- /dev/null +++ b/tests/test_search.py @@ -0,0 +1,103 @@ +"""Unit tests for search.py — verify SDK call shape, not HTTP transport.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from parallel_google_adk import web_fetch, web_search + + +class _FakeResponse: + def __init__(self, payload: dict[str, Any]) -> None: + self._payload = payload + + def model_dump(self) -> dict[str, Any]: + return self._payload + + +def _patch_client( + monkeypatch: pytest.MonkeyPatch, *, response_payload: dict[str, Any] +) -> dict[str, Any]: + """Replace get_client() with a fake; return a dict where the fake records the kwargs it was called with.""" + captured: dict[str, Any] = {} + + def fake_search(**kwargs: Any) -> _FakeResponse: + captured["search"] = kwargs + return _FakeResponse(response_payload) + + def fake_extract(**kwargs: Any) -> _FakeResponse: + captured["extract"] = kwargs + return _FakeResponse(response_payload) + + fake_client = SimpleNamespace(search=fake_search, extract=fake_extract) + monkeypatch.setattr("parallel_google_adk._search.get_client", lambda: fake_client) + return captured + + +def test_web_search_passes_advanced_settings_max_results( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = _patch_client( + monkeypatch, response_payload={"results": [{"url": "https://example.com"}]} + ) + result = web_search.func( + objective="grounded objective", queries=["q1", "q2"], max_results=7 + ) + + sent = captured["search"] + assert sent["objective"] == "grounded objective" + assert sent["search_queries"] == ["q1", "q2"] + # max_results MUST be inside advanced_settings — top-level is rejected by the SDK. + assert sent["advanced_settings"] == {"max_results": 7} + assert "max_chars_total" not in sent # omitted when None + assert result == {"results": [{"url": "https://example.com"}]} + + +def test_web_search_passes_max_chars_total_when_set( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = _patch_client(monkeypatch, response_payload={"results": []}) + web_search.func(objective="x", queries=["q"], max_chars_total=10_000) + assert captured["search"]["max_chars_total"] == 10_000 + + +def test_web_fetch_returns_first_result(monkeypatch: pytest.MonkeyPatch) -> None: + captured = _patch_client( + monkeypatch, + response_payload={ + "results": [ + {"url": "https://x", "title": "X", "excerpts": ["…"]}, + ], + "warnings": [], + }, + ) + result = web_fetch.func(url="https://x") + + assert captured["extract"] == {"urls": ["https://x"]} + assert result["url"] == "https://x" + assert result["title"] == "X" + + +def test_web_fetch_passes_objective(monkeypatch: pytest.MonkeyPatch) -> None: + captured = _patch_client( + monkeypatch, response_payload={"results": [{"url": "https://x"}]} + ) + web_fetch.func(url="https://x", objective="be brief") + assert captured["extract"]["objective"] == "be brief" + + +def test_web_fetch_empty_results_returns_envelope( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_client( + monkeypatch, response_payload={"results": [], "warnings": ["unreachable"]} + ) + result = web_fetch.func(url="https://nope") + assert result == { + "url": "https://nope", + "results": [], + "warnings": ["unreachable"], + } diff --git a/tests/test_task.py b/tests/test_task.py new file mode 100644 index 0000000..e3f78ae --- /dev/null +++ b/tests/test_task.py @@ -0,0 +1,61 @@ +"""Unit tests for the task module — JSON parse error paths and _format_result.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from parallel_google_adk import enrich +from parallel_google_adk._task import _format_result + + +def test_enrich_rejects_invalid_inputs_json() -> None: + enrich_fn = enrich.func # unwrap the FunctionTool + with pytest.raises(ValueError, match="inputs_json"): + enrich_fn( + objective="x", + inputs_json="not json {", + output_schema_json='{"type":"object"}', + ) + + +def test_enrich_rejects_invalid_output_schema_json() -> None: + enrich_fn = enrich.func + with pytest.raises(ValueError, match="output_schema_json"): + enrich_fn( + objective="x", + inputs_json="[]", + output_schema_json="also not json", + ) + + +def test_format_result_with_model_dump() -> None: + obj = SimpleNamespace(model_dump=lambda: {"output": {"x": 1}, "run": {"id": "r1"}}) + out = _format_result(obj) + assert out == {"output": {"x": 1}, "run": {"id": "r1"}} + + +def test_format_result_lifts_parsed_pydantic_instance() -> None: + parsed = SimpleNamespace(model_dump=lambda: {"website": "parallel.ai"}) + obj = SimpleNamespace(model_dump=lambda: {"output": {"basis": []}}, parsed=parsed) + out = _format_result(obj) + assert out["parsed"] == {"website": "parallel.ai"} + assert out["output"]["basis"] == [] + + +def test_format_result_dict_fallback() -> None: + # An object without model_dump should fall through to dict(...) — we use a + # mapping-compatible object so dict() succeeds. + out = _format_result({"output": {"x": 1}}) + assert out == {"output": {"x": 1}} + + +def test_format_result_with_parsed_plain_dict() -> None: + """If parsed isn't a Pydantic model, surface it verbatim.""" + obj = SimpleNamespace( + model_dump=lambda: {"output": {}}, parsed={"already": "a dict"} + ) + out: dict[str, Any] = _format_result(obj) + assert out["parsed"] == {"already": "a dict"} diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 33f5e61..a749698 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -19,56 +19,99 @@ def __init__(self, name: str) -> None: class _FakeContext: - def __init__(self) -> None: + def __init__(self, function_call_id: str = "fc-1") -> None: + self.function_call_id = function_call_id self.state: dict[str, Any] = {} @pytest.mark.asyncio async def test_plugin_records_parallel_call() -> None: plugin = ParallelTracingPlugin() - tool = _FakeTool("_web_search") + tool = _FakeTool("web_search") ctx = _FakeContext() args: dict[str, Any] = {"objective": "x", "queries": ["x"]} - await plugin.before_tool_callback(tool=tool, tool_args=args) + await plugin.before_tool_callback(tool=tool, tool_args=args, tool_context=ctx) await plugin.after_tool_callback( tool=tool, tool_args=args, - result={"results": [{"citations": ["a", "b"]}, {"citations": ["c"]}]}, tool_context=ctx, + result={"results": [{"url": "https://a"}, {"url": "https://b"}]}, ) calls = ctx.state["_parallel_calls"] assert len(calls) == 1 entry = calls[0] - assert entry["tool"] == "_web_search" - assert entry["citation_count"] == 3 + assert entry["tool"] == "web_search" + assert entry["citation_count"] == 2 assert entry["latency_s"] is not None and entry["latency_s"] >= 0 +@pytest.mark.asyncio +async def test_plugin_records_error() -> None: + plugin = ParallelTracingPlugin() + tool = _FakeTool("deep_research") + ctx = _FakeContext(function_call_id="fc-err") + args: dict[str, Any] = {"objective": "x"} + await plugin.before_tool_callback(tool=tool, tool_args=args, tool_context=ctx) + await plugin.on_tool_error_callback( + tool=tool, + tool_args=args, + tool_context=ctx, + error=RuntimeError("boom"), + ) + calls = ctx.state["_parallel_calls"] + assert len(calls) == 1 + entry = calls[0] + assert entry["tool"] == "deep_research" + assert "boom" in entry["error"] + assert entry["latency_s"] is not None + + @pytest.mark.asyncio async def test_plugin_skips_non_parallel_tool() -> None: plugin = ParallelTracingPlugin() tool = _FakeTool("some_other_tool") ctx = _FakeContext() args: dict[str, Any] = {} - await plugin.before_tool_callback(tool=tool, tool_args=args) + await plugin.before_tool_callback(tool=tool, tool_args=args, tool_context=ctx) await plugin.after_tool_callback( - tool=tool, tool_args=args, result={}, tool_context=ctx + tool=tool, tool_args=args, tool_context=ctx, result={} ) assert ctx.state.get("_parallel_calls") is None +@pytest.mark.asyncio +async def test_concurrent_calls_use_distinct_keys() -> None: + """Two overlapping tool calls (different function_call_ids) must not race.""" + plugin = ParallelTracingPlugin() + tool = _FakeTool("web_search") + ctx_a = _FakeContext(function_call_id="fc-a") + ctx_b = _FakeContext(function_call_id="fc-b") + args_a: dict[str, Any] = {} + args_b: dict[str, Any] = {} + + await plugin.before_tool_callback(tool=tool, tool_args=args_a, tool_context=ctx_a) + await plugin.before_tool_callback(tool=tool, tool_args=args_b, tool_context=ctx_b) + await plugin.after_tool_callback( + tool=tool, tool_args=args_b, tool_context=ctx_b, result={"results": []} + ) + await plugin.after_tool_callback( + tool=tool, tool_args=args_a, tool_context=ctx_a, result={"results": []} + ) + assert len(ctx_a.state["_parallel_calls"]) == 1 + assert len(ctx_b.state["_parallel_calls"]) == 1 + assert ctx_a.state["_parallel_calls"][0]["latency_s"] is not None + assert ctx_b.state["_parallel_calls"][0]["latency_s"] is not None + + def test_count_citations_handles_shapes() -> None: - # Top-level citations list. assert _count_citations({"citations": [1, 2, 3]}) == 3 - # Search/Extract: URL-bearing results. assert ( _count_citations({"results": [{"url": "https://a"}, {"url": "https://b"}]}) == 2 ) - # Older shape: per-result citations lists. + assert _count_citations({"url": "https://a", "title": "x", "excerpts": ["y"]}) == 1 assert ( _count_citations({"results": [{"citations": [1]}, {"citations": [2, 3]}]}) == 3 ) - # Task API: output.basis list with per-entry citations. assert ( _count_citations( { @@ -82,7 +125,6 @@ def test_count_citations_handles_shapes() -> None: ) == 3 ) - # Legacy basis dict shape. assert _count_citations({"basis": {"citations": [1, 2]}}) == 2 assert _count_citations({}) == 0 assert _count_citations(None) == 0