From 58dd4867e2bfabaa0167d325f75f95104a4cb02b Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Mon, 20 Jul 2026 14:38:17 -0700 Subject: [PATCH 01/17] =?UTF-8?q?feat:=20Python=20SDK=202.5.0=20=E2=80=94?= =?UTF-8?q?=20Research=20Background=20Mode?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds background-mode research (async task submission, polling, streaming) to the SDK. Adds you.research(background=True), you.get_research_task(), you.stream_research_task(), and the hand-maintained youdotcom.research_helpers module. All 2.4.0 features unchanged. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- CHANGELOG.md | 28 + MIGRATION.md | 76 +- README.md | 48 +- docs/errors/getresearchtaskforbiddenerror.md | 10 + .../getresearchtaskinternalservererror.md | 10 + docs/errors/getresearchtasknotfounderror.md | 10 + .../getresearchtaskunauthorizederror.md | 10 + .../streamresearchtaskforbiddenerror.md | 10 + .../streamresearchtaskinternalservererror.md | 10 + .../errors/streamresearchtasknotfounderror.md | 10 + .../streamresearchtaskunauthorizederror.md | 10 + docs/models/event.md | 24 + docs/models/getresearchtaskrequest.md | 8 + docs/models/researchrequest.md | 1 + docs/models/researchresult.md | 19 + docs/models/researchtaskstreamevent.md | 12 + docs/models/researchtaskstreameventdata.md | 15 + docs/models/result.md | 9 + docs/models/streamresearchtaskrequest.md | 9 + docs/models/taskdetail.md | 16 + docs/models/taskdetailinput.md | 9 + docs/models/taskdetailstatus.md | 22 + docs/models/taskresponse.md | 12 + docs/models/taskresponsestatus.md | 22 + docs/sdks/you/README.md | 93 ++ examples/api-example-calls.py | 76 + overlays/python_overlay.yaml | 17 + pyproject.toml | 2 +- src/youdotcom/_version.py | 4 +- src/youdotcom/errors/__init__.py | 52 + src/youdotcom/errors/getresearchtaskop.py | 92 ++ src/youdotcom/errors/streamresearchtaskop.py | 92 ++ src/youdotcom/models/__init__.py | 69 + src/youdotcom/models/getresearchtaskop.py | 18 + src/youdotcom/models/researchop.py | 24 +- .../models/researchtaskstreamevent.py | 113 ++ src/youdotcom/models/streamresearchtaskop.py | 44 + src/youdotcom/models/taskdetail.py | 111 ++ src/youdotcom/models/taskresponse.py | 47 + src/youdotcom/research_helpers.py | 690 +++++++++ src/youdotcom/sdk.py | 498 ++++++- tests/README.md | 23 +- .../internal/handler/generated_handlers.go | 2 + .../internal/handler/pathgetv1research.go | 111 ++ .../handler/pathgetv1researchstream.go | 128 ++ .../internal/handler/pathpostv1research.go | 45 + tests/test_live.py | 183 +++ tests/test_research.py | 149 +- tests/test_research_helpers.py | 1281 +++++++++++++++++ uv.lock | 2 +- 50 files changed, 4344 insertions(+), 32 deletions(-) create mode 100644 docs/errors/getresearchtaskforbiddenerror.md create mode 100644 docs/errors/getresearchtaskinternalservererror.md create mode 100644 docs/errors/getresearchtasknotfounderror.md create mode 100644 docs/errors/getresearchtaskunauthorizederror.md create mode 100644 docs/errors/streamresearchtaskforbiddenerror.md create mode 100644 docs/errors/streamresearchtaskinternalservererror.md create mode 100644 docs/errors/streamresearchtasknotfounderror.md create mode 100644 docs/errors/streamresearchtaskunauthorizederror.md create mode 100644 docs/models/event.md create mode 100644 docs/models/getresearchtaskrequest.md create mode 100644 docs/models/researchresult.md create mode 100644 docs/models/researchtaskstreamevent.md create mode 100644 docs/models/researchtaskstreameventdata.md create mode 100644 docs/models/result.md create mode 100644 docs/models/streamresearchtaskrequest.md create mode 100644 docs/models/taskdetail.md create mode 100644 docs/models/taskdetailinput.md create mode 100644 docs/models/taskdetailstatus.md create mode 100644 docs/models/taskresponse.md create mode 100644 docs/models/taskresponsestatus.md create mode 100644 src/youdotcom/errors/getresearchtaskop.py create mode 100644 src/youdotcom/errors/streamresearchtaskop.py create mode 100644 src/youdotcom/models/getresearchtaskop.py create mode 100644 src/youdotcom/models/researchtaskstreamevent.py create mode 100644 src/youdotcom/models/streamresearchtaskop.py create mode 100644 src/youdotcom/models/taskdetail.py create mode 100644 src/youdotcom/models/taskresponse.py create mode 100644 src/youdotcom/research_helpers.py create mode 100644 tests/mockserver/internal/handler/pathgetv1research.go create mode 100644 tests/mockserver/internal/handler/pathgetv1researchstream.go create mode 100644 tests/test_research_helpers.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a9764e6..cc87f1b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,34 @@ All notable changes to the You.com Python SDK will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [2.5.0] - 2026-07-20 + +### Added + +- **Research Background Mode**: The `you.research()` method now accepts an optional `background=True` parameter. When enabled, the API queues the research task and returns a `TaskResponse` (with `task_id`, `type`, `status`, `stream_url`, `created_at`) immediately instead of waiting for the inline `ResearchResponse`. Use the new methods to poll or stream the task to completion. + +- **`you.get_research_task(task_id=...)`**: Poll the status of a background research task via `GET /v1/research/{task_id}`. Returns a `TaskDetail` with `status`, `result` (populated when completed), `error` (populated when failed), and timing fields. + +- **`you.stream_research_task(task_id=..., from_id=0)`**: Stream real-time updates for a background research task via `GET /v1/research/{task_id}/stream` (Server-Sent Events). Returns an `EventStream` of `ResearchTaskStreamEvent` objects. The connection closes automatically when the task reaches a terminal state (`response.done`, `complete`, `error`, or `cancelled`). The stream may also emit `completed`, `failed`, or `cancelled` as event names corresponding to the task's terminal status. + +- **Convenience helpers** in `youdotcom.research_helpers` (hand-maintained, regen-safe): + - `research_background(you, ...)` / `research_background_async(you, ...)` — submit and return `TaskResponse` directly (no Union narrowing needed). + - `poll_research_task(you, task_id, ...)` / `poll_research_task_async(...)` — poll until terminal status (`completed`, `failed`, `cancelled`). + - `research_and_wait(you, ...)` / `research_and_wait_async(...)` — submit with `background=True`, then stream SSE events until a terminal event arrives, and fetch the final `TaskDetail`. For polling instead of streaming, use `poll_research_task` directly. + - `stream_research(you, task_id, ...)` / `stream_research_async(...)` — tolerant SSE iterator that surfaces undocumented event types as raw dicts instead of crashing on validation. **Recommended over `you.stream_research_task()` for real research tasks**, since the server emits intermediate workflow events not in the strict `Event` enum. + +- **New models**: `TaskResponse`, `TaskResponseStatus`, `TaskDetail`, `TaskDetailStatus`, `TaskDetailInput`, `Result`, `GetResearchTaskRequest`, `StreamResearchTaskRequest`, `ResearchTaskStreamEvent`, `ResearchTaskStreamEventData`, `Event`, `ResearchResult` (Union alias). + +- **New error classes**: `GetResearchTaskUnauthorizedError`, `GetResearchTaskForbiddenError`, `GetResearchTaskNotFoundError`, `GetResearchTaskInternalServerError`, `StreamResearchTaskUnauthorizedError`, `StreamResearchTaskForbiddenError`, `StreamResearchTaskNotFoundError`, `StreamResearchTaskInternalServerError` (and their `*Data` variants). + +### Changed + +- **`you.research()` return type** is now `Union[ResearchResponse, TaskResponse]` (exposed as the `ResearchResult` alias). When `background=False` (the default), the return is `ResearchResponse` as before. When `background=True`, the return is `TaskResponse`. Use `isinstance(res, TaskResponse)` to narrow, or use the `research_helpers` convenience functions. + +### Notes + +- All 2.4.0 features (finance_research, source_control, output_schema, boost_domains, max_age, env-var precedence, user_agent hook) are unchanged. + ## [2.4.0] - 2026-07-14 ### Added diff --git a/MIGRATION.md b/MIGRATION.md index 1909d95..04775f4 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -1,6 +1,80 @@ # Migration Guide -## 2.3.0 → 2.4.0 (Latest) +## 2.4.0 → 2.5.0 (Latest) + +### New Background Mode for Research + +The `you.research()` method now accepts `background=True` to queue long-running research tasks asynchronously. The return type changes from `ResearchResponse` to `Union[ResearchResponse, TaskResponse]` (exposed as the `ResearchResult` alias). + +**Synchronous mode (default, unchanged):** + +```python +from youdotcom.models import ResearchResponse +res = you.research(input="...", research_effort=ResearchEffort.DEEP) +# res is ResearchResponse — same as 2.4.0 +assert isinstance(res, ResearchResponse) +``` + +**Background mode (new):** + +```python +from youdotcom.models import TaskResponse, TaskDetail + +# Submit and get a task handle +res = you.research(input="...", research_effort=ResearchEffort.DEEP, background=True) +assert isinstance(res, TaskResponse) +print(res.task_id, res.stream_url) + +# Poll for completion +detail = you.get_research_task(task_id=res.task_id) +if detail.status.value == "completed": + payload = detail.result.model_dump() # extra="allow" preserves the full ResearchResponse + print(payload["output"]["content"]) + +# Or stream events with the tolerant helper (recommended) +from youdotcom.research_helpers import stream_research +for event in stream_research(you, task_id=res.task_id): + print(event.event, event.data) + if event.event in ("response.done", "completed", "error", "failed", "cancelled"): + break +``` + +**Or use the convenience helpers:** + +```python +from youdotcom.research_helpers import ( + research_background, poll_research_task, research_and_wait, + stream_research, +) + +# Option 1: Submit and wait in one call (simplest) +# Streams SSE events until the task completes, then fetches the result. +detail = research_and_wait( + you, input="...", research_effort=ResearchEffort.DEEP, +) + +# Option 2: Submit, then poll manually +task = research_background(you, input="...", research_effort=ResearchEffort.DEEP) +detail = poll_research_task(you, task_id=task.task_id) + +# Option 3: Stream SSE events with a tolerant decoder +# (recommended over you.stream_research_task for real tasks, since +# the server emits intermediate event types not in the strict enum) +for event in stream_research(you, task_id=task.task_id): + print(event.event, event.data) + if event.event in ("response.done", "completed"): + break +``` + +> **Note on streaming:** The generated `you.stream_research_task()` method uses a strict pydantic decoder that validates event names against a fixed `Event` enum. The server emits intermediate workflow events (e.g. `response.created`, `response.starting`, `response.output_item.added`) that are not in this enum, which causes `ResponseValidationError` on the first intermediate event. The `stream_research()` helper uses a tolerant decoder that surfaces unknown event names as raw dicts instead of crashing. For real research tasks, prefer `stream_research()`. + +### No Breaking Changes for Existing Code + +If you do not use `background=True`, your existing `you.research()` calls are unchanged at runtime. The return is still `ResearchResponse` when `background` is omitted or `False`. + +> **Type-level note for statically-typed callers:** The return type of `you.research()` widened from `ResearchResponse` to `Union[ResearchResponse, TaskResponse]`. If your code accesses `res.output` directly (or passes the result where a `ResearchResponse` is expected), `mypy`/`pyright` will flag it because `TaskResponse` has no `output` field. Add an `isinstance(res, ResearchResponse)` narrow or use `cast(ResearchResponse, res)` to satisfy type checkers. This only affects type checking, not runtime behavior. + +## 2.3.0 → 2.4.0 This guide covers breaking changes introduced in 2.4.0. If you are upgrading from 1.x or 2.0, also read the [1.x → 2.0](#1x-to-20) section below. diff --git a/README.md b/README.md index 7d7f1ef..d9d0461 100644 --- a/README.md +++ b/README.md @@ -230,6 +230,8 @@ with You( * [search_post](docs/sdks/you/README.md#search_post) - Returns a list of unified search results from web and news sources * [research](docs/sdks/you/README.md#research) - Returns comprehensive research-grade answers with multi-step reasoning +* [get_research_task](docs/sdks/you/README.md#get_research_task) - Get the status of a background research task +* [stream_research_task](docs/sdks/you/README.md#stream_research_task) - Stream updates for a background research task * [finance_research](docs/sdks/you/README.md#finance_research) - Returns comprehensive finance-grade research answers with multi-step reasoning ### [Agents.Runs](docs/sdks/runs/README.md) @@ -445,7 +447,7 @@ with You( **Primary error:** * [`YouError`](./src/youdotcom/errors/youerror.py): The base class for HTTP error responses. -
Less common errors (23) +
Less common errors (31)
@@ -456,24 +458,32 @@ with You( **Inherit from [`YouError`](./src/youdotcom/errors/youerror.py)**: -* [`UnauthorizedResponseError`](./src/youdotcom/errors/unauthorizedresponseerror.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 2 of 6 methods.* -* [`ForbiddenResponseError`](./src/youdotcom/errors/forbiddenresponseerror.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 2 of 6 methods.* -* [`UnprocessableEntityResponseError`](./src/youdotcom/errors/unprocessableentityresponseerror.py): Unprocessable Entity. Invalid request parameter combination. Status code `422`. Applicable to 2 of 6 methods.* -* [`InternalServerErrorResponse`](./src/youdotcom/errors/internalservererrorresponse.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 2 of 6 methods.* -* [`AgentRuns400ResponseError`](./src/youdotcom/errors/agentruns400responseerror.py): The message returned by the error. Status code `400`. Applicable to 1 of 6 methods.* -* [`ResearchUnauthorizedError`](./src/youdotcom/errors/researchunauthorizederror.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 1 of 6 methods.* -* [`FinanceResearchUnauthorizedError`](./src/youdotcom/errors/financeresearchunauthorizederror.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 1 of 6 methods.* -* [`ContentsUnauthorizedError`](./src/youdotcom/errors/contentsunauthorizederror.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 1 of 6 methods.* -* [`AgentRuns401ResponseError`](./src/youdotcom/errors/agentruns401responseerror.py): The message returned by the error. Status code `401`. Applicable to 1 of 6 methods.* -* [`ResearchForbiddenError`](./src/youdotcom/errors/researchforbiddenerror.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 1 of 6 methods.* -* [`FinanceResearchForbiddenError`](./src/youdotcom/errors/financeresearchforbiddenerror.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 1 of 6 methods.* -* [`ContentsForbiddenError`](./src/youdotcom/errors/contentsforbiddenerror.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 1 of 6 methods.* -* [`ResearchUnprocessableEntityError`](./src/youdotcom/errors/researchunprocessableentityerror.py): Unprocessable Entity. Request validation failed. Status code `422`. Applicable to 1 of 6 methods.* -* [`FinanceResearchUnprocessableEntityError`](./src/youdotcom/errors/financeresearchunprocessableentityerror.py): Unprocessable Entity. Request validation failed. Status code `422`. Applicable to 1 of 6 methods.* -* [`AgentRuns422ResponseError`](./src/youdotcom/errors/agentruns422responseerror.py): Unprocessable Entity - Invalid request data. Status code `422`. Applicable to 1 of 6 methods.* -* [`ResearchInternalServerError`](./src/youdotcom/errors/researchinternalservererror.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 1 of 6 methods.* -* [`FinanceResearchInternalServerError`](./src/youdotcom/errors/financeresearchinternalservererror.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 1 of 6 methods.* -* [`ContentsInternalServerError`](./src/youdotcom/errors/contentsinternalservererror.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 1 of 6 methods.* +* [`UnauthorizedResponseError`](./src/youdotcom/errors/unauthorizedresponseerror.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 2 of 8 methods.* +* [`ForbiddenResponseError`](./src/youdotcom/errors/forbiddenresponseerror.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 2 of 8 methods.* +* [`UnprocessableEntityResponseError`](./src/youdotcom/errors/unprocessableentityresponseerror.py): Unprocessable Entity. Invalid request parameter combination. Status code `422`. Applicable to 2 of 8 methods.* +* [`InternalServerErrorResponse`](./src/youdotcom/errors/internalservererrorresponse.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 2 of 8 methods.* +* [`AgentRuns400ResponseError`](./src/youdotcom/errors/agentruns400responseerror.py): The message returned by the error. Status code `400`. Applicable to 1 of 8 methods.* +* [`ResearchUnauthorizedError`](./src/youdotcom/errors/researchunauthorizederror.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 1 of 8 methods.* +* [`FinanceResearchUnauthorizedError`](./src/youdotcom/errors/financeresearchunauthorizederror.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 1 of 8 methods.* +* [`ContentsUnauthorizedError`](./src/youdotcom/errors/contentsunauthorizederror.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 1 of 8 methods.* +* [`AgentRuns401ResponseError`](./src/youdotcom/errors/agentruns401responseerror.py): The message returned by the error. Status code `401`. Applicable to 1 of 8 methods.* +* [`ResearchForbiddenError`](./src/youdotcom/errors/researchforbiddenerror.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 1 of 8 methods.* +* [`FinanceResearchForbiddenError`](./src/youdotcom/errors/financeresearchforbiddenerror.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 1 of 8 methods.* +* [`ContentsForbiddenError`](./src/youdotcom/errors/contentsforbiddenerror.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 1 of 8 methods.* +* [`ResearchUnprocessableEntityError`](./src/youdotcom/errors/researchunprocessableentityerror.py): Unprocessable Entity. Request validation failed. Status code `422`. Applicable to 1 of 8 methods.* +* [`FinanceResearchUnprocessableEntityError`](./src/youdotcom/errors/financeresearchunprocessableentityerror.py): Unprocessable Entity. Request validation failed. Status code `422`. Applicable to 1 of 8 methods.* +* [`AgentRuns422ResponseError`](./src/youdotcom/errors/agentruns422responseerror.py): Unprocessable Entity - Invalid request data. Status code `422`. Applicable to 1 of 8 methods.* +* [`ResearchInternalServerError`](./src/youdotcom/errors/researchinternalservererror.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 1 of 8 methods.* +* [`FinanceResearchInternalServerError`](./src/youdotcom/errors/financeresearchinternalservererror.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 1 of 8 methods.* +* [`ContentsInternalServerError`](./src/youdotcom/errors/contentsinternalservererror.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 1 of 8 methods.* +* [`GetResearchTaskUnauthorizedError`](./src/youdotcom/errors/getresearchtaskop.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 1 of 8 methods.* +* [`GetResearchTaskForbiddenError`](./src/youdotcom/errors/getresearchtaskop.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 1 of 8 methods.* +* [`GetResearchTaskNotFoundError`](./src/youdotcom/errors/getresearchtaskop.py): Not Found. Task does not exist or does not belong to the caller. Status code `404`. Applicable to 1 of 8 methods.* +* [`GetResearchTaskInternalServerError`](./src/youdotcom/errors/getresearchtaskop.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 1 of 8 methods.* +* [`StreamResearchTaskUnauthorizedError`](./src/youdotcom/errors/streamresearchtaskop.py): Unauthorized. Problems with API key. Status code `401`. Applicable to 1 of 8 methods.* +* [`StreamResearchTaskForbiddenError`](./src/youdotcom/errors/streamresearchtaskop.py): Forbidden. API key lacks scope for this path. Status code `403`. Applicable to 1 of 8 methods.* +* [`StreamResearchTaskNotFoundError`](./src/youdotcom/errors/streamresearchtaskop.py): Not Found. Task does not exist or does not belong to the caller. Status code `404`. Applicable to 1 of 8 methods.* +* [`StreamResearchTaskInternalServerError`](./src/youdotcom/errors/streamresearchtaskop.py): Internal Server Error during authentication/authorization middleware. Status code `500`. Applicable to 1 of 8 methods.* * [`ResponseValidationError`](./src/youdotcom/errors/responsevalidationerror.py): Type mismatch between the response data and the expected Pydantic model. Provides access to the Pydantic validation error via the `cause` attribute.
diff --git a/docs/errors/getresearchtaskforbiddenerror.md b/docs/errors/getresearchtaskforbiddenerror.md new file mode 100644 index 0000000..cc633bb --- /dev/null +++ b/docs/errors/getresearchtaskforbiddenerror.md @@ -0,0 +1,10 @@ +# GetResearchTaskForbiddenError + +Forbidden. API key lacks scope for this path. + + +## Fields + +| Field | Type | Required | Description | +| ------------------ | ------------------ | ------------------ | ------------------ | +| `detail` | *Optional[str]* | :heavy_minus_sign: | N/A | \ No newline at end of file diff --git a/docs/errors/getresearchtaskinternalservererror.md b/docs/errors/getresearchtaskinternalservererror.md new file mode 100644 index 0000000..ebed213 --- /dev/null +++ b/docs/errors/getresearchtaskinternalservererror.md @@ -0,0 +1,10 @@ +# GetResearchTaskInternalServerError + +Internal Server Error. + + +## Fields + +| Field | Type | Required | Description | +| ------------------ | ------------------ | ------------------ | ------------------ | +| `detail` | *Optional[str]* | :heavy_minus_sign: | N/A | \ No newline at end of file diff --git a/docs/errors/getresearchtasknotfounderror.md b/docs/errors/getresearchtasknotfounderror.md new file mode 100644 index 0000000..cbb274c --- /dev/null +++ b/docs/errors/getresearchtasknotfounderror.md @@ -0,0 +1,10 @@ +# GetResearchTaskNotFoundError + +Task not found or not authorized. + + +## Fields + +| Field | Type | Required | Description | Example | +| ------------------ | ------------------ | ------------------ | ------------------ | ------------------ | +| `detail` | *Optional[str]* | :heavy_minus_sign: | N/A | Task not found | \ No newline at end of file diff --git a/docs/errors/getresearchtaskunauthorizederror.md b/docs/errors/getresearchtaskunauthorizederror.md new file mode 100644 index 0000000..5ed32fa --- /dev/null +++ b/docs/errors/getresearchtaskunauthorizederror.md @@ -0,0 +1,10 @@ +# GetResearchTaskUnauthorizedError + +Unauthorized. Problems with API key. + + +## Fields + +| Field | Type | Required | Description | +| ------------------ | ------------------ | ------------------ | ------------------ | +| `detail` | *Optional[str]* | :heavy_minus_sign: | N/A | \ No newline at end of file diff --git a/docs/errors/streamresearchtaskforbiddenerror.md b/docs/errors/streamresearchtaskforbiddenerror.md new file mode 100644 index 0000000..f13d63a --- /dev/null +++ b/docs/errors/streamresearchtaskforbiddenerror.md @@ -0,0 +1,10 @@ +# StreamResearchTaskForbiddenError + +Forbidden. API key lacks scope for this path. + + +## Fields + +| Field | Type | Required | Description | +| ------------------ | ------------------ | ------------------ | ------------------ | +| `detail` | *Optional[str]* | :heavy_minus_sign: | N/A | \ No newline at end of file diff --git a/docs/errors/streamresearchtaskinternalservererror.md b/docs/errors/streamresearchtaskinternalservererror.md new file mode 100644 index 0000000..e7d0083 --- /dev/null +++ b/docs/errors/streamresearchtaskinternalservererror.md @@ -0,0 +1,10 @@ +# StreamResearchTaskInternalServerError + +Internal Server Error. + + +## Fields + +| Field | Type | Required | Description | +| ------------------ | ------------------ | ------------------ | ------------------ | +| `detail` | *Optional[str]* | :heavy_minus_sign: | N/A | \ No newline at end of file diff --git a/docs/errors/streamresearchtasknotfounderror.md b/docs/errors/streamresearchtasknotfounderror.md new file mode 100644 index 0000000..8d8ff6b --- /dev/null +++ b/docs/errors/streamresearchtasknotfounderror.md @@ -0,0 +1,10 @@ +# StreamResearchTaskNotFoundError + +Task not found or not authorized. + + +## Fields + +| Field | Type | Required | Description | Example | +| ------------------ | ------------------ | ------------------ | ------------------ | ------------------ | +| `detail` | *Optional[str]* | :heavy_minus_sign: | N/A | Task not found | \ No newline at end of file diff --git a/docs/errors/streamresearchtaskunauthorizederror.md b/docs/errors/streamresearchtaskunauthorizederror.md new file mode 100644 index 0000000..86e9ae7 --- /dev/null +++ b/docs/errors/streamresearchtaskunauthorizederror.md @@ -0,0 +1,10 @@ +# StreamResearchTaskUnauthorizedError + +Unauthorized. Problems with API key. + + +## Fields + +| Field | Type | Required | Description | +| ------------------ | ------------------ | ------------------ | ------------------ | +| `detail` | *Optional[str]* | :heavy_minus_sign: | N/A | \ No newline at end of file diff --git a/docs/models/event.md b/docs/models/event.md new file mode 100644 index 0000000..67efe4c --- /dev/null +++ b/docs/models/event.md @@ -0,0 +1,24 @@ +# Event + +The type of the SSE event. Terminal events that close the stream are: `response.done`, `complete`, `error`, and `cancelled`. The stream may also emit `completed`, `failed`, or `cancelled` as event names corresponding to the task's terminal status. + +## Example Usage + +```python +from youdotcom.models import Event + +value = Event.CONNECTED +``` + + +## Values + +| Name | Value | +| --------------- | --------------- | +| `CONNECTED` | connected | +| `RESPONSE_DONE` | response.done | +| `COMPLETE` | complete | +| `COMPLETED` | completed | +| `ERROR` | error | +| `FAILED` | failed | +| `CANCELLED` | cancelled | \ No newline at end of file diff --git a/docs/models/getresearchtaskrequest.md b/docs/models/getresearchtaskrequest.md new file mode 100644 index 0000000..ba0a7b2 --- /dev/null +++ b/docs/models/getresearchtaskrequest.md @@ -0,0 +1,8 @@ +# GetResearchTaskRequest + + +## Fields + +| Field | Type | Required | Description | +| ------------------------------ | ------------------------------ | ------------------------------ | ------------------------------ | +| `task_id` | *str* | :heavy_check_mark: | The UUID of the research task. | \ No newline at end of file diff --git a/docs/models/researchrequest.md b/docs/models/researchrequest.md index 6a08aa1..e8ebd05 100644 --- a/docs/models/researchrequest.md +++ b/docs/models/researchrequest.md @@ -7,5 +7,6 @@ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `input` | *str* | :heavy_check_mark: | The research question or complex query requiring in-depth investigation and multi-step reasoning.

Note: The maximum length of the input is 40,000 characters. | | `research_effort` | [Optional[models.ResearchEffort]](../models/researcheffort.md) | :heavy_minus_sign: | Controls how much time and effort the Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time.

Available levels:
- `lite`: Returns answers quickly. Good for straightforward questions that just need a fast, reliable answer.
- `standard`: The default. Balances speed and depth, a good fit for most questions.
- `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed.
- `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. | +| `background` | *Optional[bool]* | :heavy_minus_sign: | When true, queue a research task and return a task handle immediately instead of waiting for the result inline. Defaults to synchronous. When enabled, the response is a TaskResponse object with a task_id and stream_url for polling progress via GET /v1/research/{task_id} or streaming via GET /v1/research/{task_id}/stream. | | `source_control` | [Optional[models.SourceControl]](../models/sourcecontrol.md) | :heavy_minus_sign: | Beta. Controls which web sources the research agent searches and visits. Use this to allow specific domains, block specific domains, boost specific domains, filter by recency, or focus web results by country.

`include_domains` and `exclude_domains` cannot be used together. Each domain list is capped at 500 entries. `exclude_domains` also blocks the research agent from visiting pages on those domains during browsing. `boost_domains` gives matching domains a relative ranking boost without filtering out other domains. It can be combined with `exclude_domains` but cannot be combined with `include_domains`. | | `output_schema` | Dict[str, *Any*] | :heavy_minus_sign: | Beta. Requests structured JSON output in output.content using a supported JSON Schema subset. Supported only with research_effort values standard, deep, and exhaustive. Sending output_schema with research_effort: "lite" returns 422.

Schema rules: Root must be a JSON object. Top-level anyOf is not allowed. Every object must define properties and set additionalProperties: false. Every property must be listed in required. Recursive schemas are not supported.

Limits: Max nesting depth 5, max total properties 100, max total enum values 500, max total schema string budget 25,000. | \ No newline at end of file diff --git a/docs/models/researchresult.md b/docs/models/researchresult.md new file mode 100644 index 0000000..2a4c454 --- /dev/null +++ b/docs/models/researchresult.md @@ -0,0 +1,19 @@ +# ResearchResult + +A JSON object containing a comprehensive answer with citations and supporting search results. When background=true, returns a task handle instead. + + +## Supported Types + +### `models.ResearchResponse` + +```python +value: models.ResearchResponse = /* values here */ +``` + +### `models.TaskResponse` + +```python +value: models.TaskResponse = /* values here */ +``` + diff --git a/docs/models/researchtaskstreamevent.md b/docs/models/researchtaskstreamevent.md new file mode 100644 index 0000000..0fe2392 --- /dev/null +++ b/docs/models/researchtaskstreamevent.md @@ -0,0 +1,12 @@ +# ResearchTaskStreamEvent + +A server-sent event for a background research task stream. + + +## Fields + +| Field | Type | Required | Description | +| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `id` | *str* | :heavy_check_mark: | Sequence number of the SSE event. | +| `event` | [models.Event](../models/event.md) | :heavy_check_mark: | The type of the SSE event. Terminal events that close the stream are: `response.done`, `complete`, `error`, and `cancelled`. The stream may also emit `completed`, `failed`, or `cancelled` as event names corresponding to the task's terminal status. | +| `data` | [models.ResearchTaskStreamEventData](../models/researchtaskstreameventdata.md) | :heavy_check_mark: | The event payload. Structure varies by event type. Common fields include type, task_id, status, data (event-specific), error, and sequence. | \ No newline at end of file diff --git a/docs/models/researchtaskstreameventdata.md b/docs/models/researchtaskstreameventdata.md new file mode 100644 index 0000000..56634db --- /dev/null +++ b/docs/models/researchtaskstreameventdata.md @@ -0,0 +1,15 @@ +# ResearchTaskStreamEventData + +The event payload. Structure varies by event type. Common fields include type, task_id, status, data (event-specific), error, and sequence. + + +## Fields + +| Field | Type | Required | Description | +| ----------------------------------------------- | ----------------------------------------------- | ----------------------------------------------- | ----------------------------------------------- | +| `type` | *Optional[str]* | :heavy_minus_sign: | The event type identifier. | +| `task_id` | *Optional[str]* | :heavy_minus_sign: | The task UUID. | +| `status` | *Optional[str]* | :heavy_minus_sign: | Current task status when the event was emitted. | +| `data` | Dict[str, *Any*] | :heavy_minus_sign: | Event-specific payload data. | +| `error` | *OptionalNullable[str]* | :heavy_minus_sign: | Error message if the event represents an error. | +| `sequence` | *Optional[int]* | :heavy_minus_sign: | Event sequence number. | \ No newline at end of file diff --git a/docs/models/result.md b/docs/models/result.md new file mode 100644 index 0000000..bca20b4 --- /dev/null +++ b/docs/models/result.md @@ -0,0 +1,9 @@ +# Result + +The task result when completed. For research tasks, this contains the ResearchResponse output. + + +## Fields + +| Field | Type | Required | Description | +| ----------- | ----------- | ----------- | ----------- | \ No newline at end of file diff --git a/docs/models/streamresearchtaskrequest.md b/docs/models/streamresearchtaskrequest.md new file mode 100644 index 0000000..1b8c0b3 --- /dev/null +++ b/docs/models/streamresearchtaskrequest.md @@ -0,0 +1,9 @@ +# StreamResearchTaskRequest + + +## Fields + +| Field | Type | Required | Description | +| ----------------------------------------------- | ----------------------------------------------- | ----------------------------------------------- | ----------------------------------------------- | +| `task_id` | *str* | :heavy_check_mark: | The UUID of the research task. | +| `from_id` | *Optional[int]* | :heavy_minus_sign: | Resume from a sequence number for reconnection. | \ No newline at end of file diff --git a/docs/models/taskdetail.md b/docs/models/taskdetail.md new file mode 100644 index 0000000..a153ed2 --- /dev/null +++ b/docs/models/taskdetail.md @@ -0,0 +1,16 @@ +# TaskDetail + + +## Fields + +| Field | Type | Required | Description | Example | +| ---------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------- | +| `id` | *str* | :heavy_check_mark: | Unique identifier for the task. | | +| `task_type` | *str* | :heavy_check_mark: | Task type. | research | +| `status` | [models.TaskDetailStatus](../models/taskdetailstatus.md) | :heavy_check_mark: | Current status of the task. | | +| `created_at` | [date](https://docs.python.org/3/library/datetime.html#date-objects) | :heavy_check_mark: | When the task was created. | | +| `updated_at` | [date](https://docs.python.org/3/library/datetime.html#date-objects) | :heavy_check_mark: | When the task was last updated. | | +| `completed_at` | [date](https://docs.python.org/3/library/datetime.html#date-objects) | :heavy_minus_sign: | When the task completed, if applicable. | | +| `error` | *OptionalNullable[str]* | :heavy_minus_sign: | Error message if the task failed. | | +| `input` | [OptionalNullable[models.TaskDetailInput]](../models/taskdetailinput.md) | :heavy_minus_sign: | The original request input for the task. | | +| `result` | [OptionalNullable[models.Result]](../models/result.md) | :heavy_minus_sign: | The task result when completed. For research tasks, this contains the ResearchResponse output. | | \ No newline at end of file diff --git a/docs/models/taskdetailinput.md b/docs/models/taskdetailinput.md new file mode 100644 index 0000000..4470b82 --- /dev/null +++ b/docs/models/taskdetailinput.md @@ -0,0 +1,9 @@ +# TaskDetailInput + +The original request input for the task. + + +## Fields + +| Field | Type | Required | Description | +| ----------- | ----------- | ----------- | ----------- | \ No newline at end of file diff --git a/docs/models/taskdetailstatus.md b/docs/models/taskdetailstatus.md new file mode 100644 index 0000000..94533f5 --- /dev/null +++ b/docs/models/taskdetailstatus.md @@ -0,0 +1,22 @@ +# TaskDetailStatus + +Current status of the task. + +## Example Usage + +```python +from youdotcom.models import TaskDetailStatus + +value = TaskDetailStatus.QUEUED +``` + + +## Values + +| Name | Value | +| ----------- | ----------- | +| `QUEUED` | queued | +| `RUNNING` | running | +| `COMPLETED` | completed | +| `FAILED` | failed | +| `CANCELLED` | cancelled | \ No newline at end of file diff --git a/docs/models/taskresponse.md b/docs/models/taskresponse.md new file mode 100644 index 0000000..759621a --- /dev/null +++ b/docs/models/taskresponse.md @@ -0,0 +1,12 @@ +# TaskResponse + + +## Fields + +| Field | Type | Required | Description | Example | +| -------------------------------------------------------------------- | -------------------------------------------------------------------- | -------------------------------------------------------------------- | -------------------------------------------------------------------- | -------------------------------------------------------------------- | +| `task_id` | *str* | :heavy_check_mark: | Unique identifier for the task. | | +| `type` | *str* | :heavy_check_mark: | Task type. | research | +| `status` | [models.TaskResponseStatus](../models/taskresponsestatus.md) | :heavy_check_mark: | Current status of the task. | queued | +| `stream_url` | *str* | :heavy_check_mark: | URL to stream task events via SSE. | /v1/research/a1b2c3d4-0000-0000-0000-000000000000/stream | +| `created_at` | [date](https://docs.python.org/3/library/datetime.html#date-objects) | :heavy_check_mark: | When the task was created. | | \ No newline at end of file diff --git a/docs/models/taskresponsestatus.md b/docs/models/taskresponsestatus.md new file mode 100644 index 0000000..b6ee96e --- /dev/null +++ b/docs/models/taskresponsestatus.md @@ -0,0 +1,22 @@ +# TaskResponseStatus + +Current status of the task. + +## Example Usage + +```python +from youdotcom.models import TaskResponseStatus + +value = TaskResponseStatus.QUEUED +``` + + +## Values + +| Name | Value | +| ----------- | ----------- | +| `QUEUED` | queued | +| `RUNNING` | running | +| `COMPLETED` | completed | +| `FAILED` | failed | +| `CANCELLED` | cancelled | \ No newline at end of file diff --git a/docs/sdks/you/README.md b/docs/sdks/you/README.md index be10607..642f0a3 100644 --- a/docs/sdks/you/README.md +++ b/docs/sdks/you/README.md @@ -18,6 +18,8 @@ Comprehensive API for You.com services: * [search_post](#search_post) - Returns a list of unified search results from web and news sources * [research](#research) - Returns comprehensive research-grade answers with multi-step reasoning +* [get_research_task](#get_research_task) - Get the status of a background research task +* [stream_research_task](#stream_research_task) - Stream updates for a background research task * [finance_research](#finance_research) - Returns comprehensive finance-grade research answers with multi-step reasoning ## search_post @@ -438,6 +440,97 @@ with You( | errors.ResearchInternalServerError | 500 | application/json | | errors.YouDefaultError | 4XX, 5XX | \*/\* | +## get_research_task + +Poll the status of a background research task created with background=true. When the task is completed, the result is included in the response. + +### Example Usage + + +```python +import os +from youdotcom import You + + +with You( + api_key_auth=os.getenv("YDC_API_KEY", ""), +) as you: + + res = you.get_research_task(task_id="586a9bc3-2c52-499c-a61d-be3cc9170c51") + + # Handle response + print(res) + +``` + +### Parameters + +| Parameter | Type | Required | Description | +| ------------------------------------------------------------------- | ------------------------------------------------------------------- | ------------------------------------------------------------------- | ------------------------------------------------------------------- | +| `task_id` | *str* | :heavy_check_mark: | The UUID of the research task. | +| `retries` | [Optional[utils.RetryConfig]](../../models/utils/retryconfig.md) | :heavy_minus_sign: | Configuration to override the default retry behavior of the client. | + +### Response + +**[models.TaskDetail](../../models/taskdetail.md)** + +### Errors + +| Error Type | Status Code | Content Type | +| ----------------------------------------- | ----------------------------------------- | ----------------------------------------- | +| errors.GetResearchTaskUnauthorizedError | 401 | application/json | +| errors.GetResearchTaskForbiddenError | 403 | application/json | +| errors.GetResearchTaskNotFoundError | 404 | application/json | +| errors.GetResearchTaskInternalServerError | 500 | application/json | +| errors.YouDefaultError | 4XX, 5XX | \*/\* | + +## stream_research_task + +Stream real-time updates for a background research task via Server-Sent Events (SSE). Supports reconnection via the from_id query parameter to replay missed events. The connection closes automatically when the task reaches a terminal state. + +### Example Usage + + +```python +import os +from youdotcom import You + + +with You( + api_key_auth=os.getenv("YDC_API_KEY", ""), +) as you: + + res = you.stream_research_task(task_id="b431835b-e51d-453e-a623-25615ac31489", from_id=0) + + with res as event_stream: + for event in event_stream: + # handle event + print(event, flush=True) + +``` + +### Parameters + +| Parameter | Type | Required | Description | +| ------------------------------------------------------------------- | ------------------------------------------------------------------- | ------------------------------------------------------------------- | ------------------------------------------------------------------- | +| `task_id` | *str* | :heavy_check_mark: | The UUID of the research task. | +| `from_id` | *Optional[int]* | :heavy_minus_sign: | Resume from a sequence number for reconnection. | +| `retries` | [Optional[utils.RetryConfig]](../../models/utils/retryconfig.md) | :heavy_minus_sign: | Configuration to override the default retry behavior of the client. | + +### Response + +**[Union[eventstreaming.EventStream[models.ResearchTaskStreamEvent], eventstreaming.EventStreamAsync[models.ResearchTaskStreamEvent]]](../../models/.md)** + +### Errors + +| Error Type | Status Code | Content Type | +| -------------------------------------------- | -------------------------------------------- | -------------------------------------------- | +| errors.StreamResearchTaskUnauthorizedError | 401 | application/json | +| errors.StreamResearchTaskForbiddenError | 403 | application/json | +| errors.StreamResearchTaskNotFoundError | 404 | application/json | +| errors.StreamResearchTaskInternalServerError | 500 | application/json | +| errors.YouDefaultError | 4XX, 5XX | \*/\* | + ## finance_research The Finance Research API is purpose-built for financial questions. Like the Research API, it runs multiple searches, reads through sources, and synthesizes everything into a thorough, well-cited answer — but its retrieval index is optimized for financial data: earnings reports, SEC filings, analyst coverage, market data, and financial news. diff --git a/examples/api-example-calls.py b/examples/api-example-calls.py index a5fa041..a846471 100755 --- a/examples/api-example-calls.py +++ b/examples/api-example-calls.py @@ -19,6 +19,7 @@ """ from typing import Optional +import time from youdotcom import You from youdotcom.models import ( ResearchTool, @@ -44,6 +45,7 @@ FinanceResearchEffort, FinanceResearchResponse, ResearchResponse, + TaskResponse, ) from youdotcom.utils import eventstreaming @@ -295,6 +297,78 @@ def research_request(): print(f" - {source.title or 'Untitled'}: {source.url}") +def research_background_request(): + """ + Research API with background mode: returns a task handle instead of the final answer. + Poll status with `you.get_research_task(task_id)` or stream events with + `you.stream_research_task(task_id)`. + """ + print("\n🚀 Running Research Background Request...\n") + + assert you is not None, "SDK client not initialized" + + res = you.research( + input="Compare the profitability of NVIDIA, AMD, and Intel over the past 5 fiscal years.", + research_effort=ResearchEffort.DEEP, + background=True, + ) + + assert isinstance(res, TaskResponse) + print(f"Queued task {res.task_id} (status: {res.status.value})") + print(f"Stream URL: {res.stream_url}") + + # Optional: poll until completion + print("\nPolling for completion...") + while True: + status_res = you.get_research_task(task_id=res.task_id) + status = status_res.status.value + print(f" status: {status}") + if status in ("completed", "failed", "cancelled"): + break + time.sleep(5) + + if status == "completed": + # The Result model uses extra="allow", so model_dump() recovers + # the full ResearchResponse payload from the task detail. + print("\nFinal answer (preview):") + payload = status_res.result.model_dump() if status_res.result else {} + content = payload.get("output", {}).get("content", "") + if isinstance(content, str): + print(content[:500] + ("..." if len(content) > 500 else "")) + + +def research_and_wait_example(): + """ + Research API with the research_and_wait helper: submit in background + mode and wait for completion in one call. Returns the final TaskDetail. + """ + from youdotcom.research_helpers import research_and_wait + from youdotcom.models import TaskDetail + + print("\n🚀 Running Research and Wait (Helper)...\n") + + assert you is not None, "SDK client not initialized" + + try: + detail = research_and_wait( + you, + timeout_s=120.0, + input="Compare the profitability of NVIDIA, AMD, and Intel over the past 5 fiscal years.", + research_effort=ResearchEffort.DEEP, + ) + except TypeError: + print(" Background mode not enabled on server. Skipping.") + return + + assert isinstance(detail, TaskDetail) + print(f"Task completed: {detail.status.value}") + if detail.result: + payload = detail.result.model_dump() + content = payload.get("output", {}).get("content", "") + if isinstance(content, str): + print(f"\nAnswer (preview):\n{content[:500]}...") + + def research_output_schema_request(): """ Research API with `output_schema` for structured JSON output. @@ -409,6 +483,8 @@ def content_request_with_max_age(): {"name": "Content Request", "fn": content_request}, {"name": "Content Request (max_age)", "fn": content_request_with_max_age}, {"name": "Research Request", "fn": research_request}, + {"name": "Research Background Mode", "fn": research_background_request}, + {"name": "Research and Wait (Helper)", "fn": research_and_wait_example}, {"name": "Research with output_schema", "fn": research_output_schema_request}, {"name": "Finance Research Request", "fn": finance_research_request}, ] diff --git a/overlays/python_overlay.yaml b/overlays/python_overlay.yaml index 678fac2..9065a61 100644 --- a/overlays/python_overlay.yaml +++ b/overlays/python_overlay.yaml @@ -45,3 +45,20 @@ actions: # the overlay makes this regen-durable until the upstream spec drops it. - target: $["components"]["schemas"]["WebResult"]["properties"]["authors"] remove: true + # Background-mode research: TaskDetail.result is an inline object schema + # that Speakeasy generates as `class Result(BaseModel): pass`, which + # round-trips as an empty dict. Inject additionalProperties: true so the + # full ResearchResponse payload (output.content, sources, etc.) survives + # deserialization. The Result model in taskdetail.py already has + # extra="allow"; this overlay ensures future regens preserve it. + - target: $["components"]["schemas"]["TaskDetail"]["properties"]["result"] + update: + additionalProperties: true + # Same fix for TaskDetail.input: the server stores the original request + # as a Dict[str, Any] (input, research_effort, source_control, etc.) but + # Speakeasy generates an empty TaskDetailInput BaseModel that drops all + # fields. Inject additionalProperties: true so the input round-trips. + # The TaskDetailInput model in taskdetail.py already has extra="allow". + - target: $["components"]["schemas"]["TaskDetail"]["properties"]["input"] + update: + additionalProperties: true diff --git a/pyproject.toml b/pyproject.toml index ee856fd..e14db01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "youdotcom" -version = "2.4.0" +version = "2.5.0" description = "The official You.com Python SDK." authors = [{ name = "You.com" },] readme = "README.md" diff --git a/src/youdotcom/_version.py b/src/youdotcom/_version.py index 5233a7d..3cd2c70 100644 --- a/src/youdotcom/_version.py +++ b/src/youdotcom/_version.py @@ -3,10 +3,10 @@ import importlib.metadata __title__: str = "youdotcom" -__version__: str = "2.4.0" +__version__: str = "2.5.0" __openapi_doc_version__: str = "1.0.0" __gen_version__: str = "2.918.1" -__user_agent__: str = "speakeasy-sdk/python 2.4.0 2.918.1 1.0.0 youdotcom" +__user_agent__: str = "speakeasy-sdk/python 2.5.0 2.918.1 1.0.0 youdotcom" try: if __package__ is not None: diff --git a/src/youdotcom/errors/__init__.py b/src/youdotcom/errors/__init__.py index fe2e187..aea2e79 100644 --- a/src/youdotcom/errors/__init__.py +++ b/src/youdotcom/errors/__init__.py @@ -40,6 +40,16 @@ ForbiddenResponseError, ForbiddenResponseErrorData, ) + from .getresearchtaskop import ( + GetResearchTaskForbiddenError, + GetResearchTaskForbiddenErrorData, + GetResearchTaskInternalServerError, + GetResearchTaskInternalServerErrorData, + GetResearchTaskNotFoundError, + GetResearchTaskNotFoundErrorData, + GetResearchTaskUnauthorizedError, + GetResearchTaskUnauthorizedErrorData, + ) from .internalservererror_response import ( InternalServerErrorResponse, InternalServerErrorResponseData, @@ -56,6 +66,16 @@ ResearchUnprocessableEntityErrorData, ) from .responsevalidationerror import ResponseValidationError + from .streamresearchtaskop import ( + StreamResearchTaskForbiddenError, + StreamResearchTaskForbiddenErrorData, + StreamResearchTaskInternalServerError, + StreamResearchTaskInternalServerErrorData, + StreamResearchTaskNotFoundError, + StreamResearchTaskNotFoundErrorData, + StreamResearchTaskUnauthorizedError, + StreamResearchTaskUnauthorizedErrorData, + ) from .unauthorized_response_error import ( UnauthorizedResponseError, UnauthorizedResponseErrorData, @@ -89,6 +109,14 @@ "FinanceResearchUnprocessableEntityErrorData", "ForbiddenResponseError", "ForbiddenResponseErrorData", + "GetResearchTaskForbiddenError", + "GetResearchTaskForbiddenErrorData", + "GetResearchTaskInternalServerError", + "GetResearchTaskInternalServerErrorData", + "GetResearchTaskNotFoundError", + "GetResearchTaskNotFoundErrorData", + "GetResearchTaskUnauthorizedError", + "GetResearchTaskUnauthorizedErrorData", "InternalServerErrorResponse", "InternalServerErrorResponseData", "NoResponseError", @@ -101,6 +129,14 @@ "ResearchUnprocessableEntityError", "ResearchUnprocessableEntityErrorData", "ResponseValidationError", + "StreamResearchTaskForbiddenError", + "StreamResearchTaskForbiddenErrorData", + "StreamResearchTaskInternalServerError", + "StreamResearchTaskInternalServerErrorData", + "StreamResearchTaskNotFoundError", + "StreamResearchTaskNotFoundErrorData", + "StreamResearchTaskUnauthorizedError", + "StreamResearchTaskUnauthorizedErrorData", "UnauthorizedResponseError", "UnauthorizedResponseErrorData", "UnprocessableEntityResponseError", @@ -132,6 +168,14 @@ "FinanceResearchUnprocessableEntityErrorData": ".finance_researchop", "ForbiddenResponseError": ".forbidden_response_error", "ForbiddenResponseErrorData": ".forbidden_response_error", + "GetResearchTaskForbiddenError": ".getresearchtaskop", + "GetResearchTaskForbiddenErrorData": ".getresearchtaskop", + "GetResearchTaskInternalServerError": ".getresearchtaskop", + "GetResearchTaskInternalServerErrorData": ".getresearchtaskop", + "GetResearchTaskNotFoundError": ".getresearchtaskop", + "GetResearchTaskNotFoundErrorData": ".getresearchtaskop", + "GetResearchTaskUnauthorizedError": ".getresearchtaskop", + "GetResearchTaskUnauthorizedErrorData": ".getresearchtaskop", "InternalServerErrorResponse": ".internalservererror_response", "InternalServerErrorResponseData": ".internalservererror_response", "NoResponseError": ".no_response_error", @@ -144,6 +188,14 @@ "ResearchUnprocessableEntityError": ".researchop", "ResearchUnprocessableEntityErrorData": ".researchop", "ResponseValidationError": ".responsevalidationerror", + "StreamResearchTaskForbiddenError": ".streamresearchtaskop", + "StreamResearchTaskForbiddenErrorData": ".streamresearchtaskop", + "StreamResearchTaskInternalServerError": ".streamresearchtaskop", + "StreamResearchTaskInternalServerErrorData": ".streamresearchtaskop", + "StreamResearchTaskNotFoundError": ".streamresearchtaskop", + "StreamResearchTaskNotFoundErrorData": ".streamresearchtaskop", + "StreamResearchTaskUnauthorizedError": ".streamresearchtaskop", + "StreamResearchTaskUnauthorizedErrorData": ".streamresearchtaskop", "UnauthorizedResponseError": ".unauthorized_response_error", "UnauthorizedResponseErrorData": ".unauthorized_response_error", "UnprocessableEntityResponseError": ".unprocessableentity_response_error", diff --git a/src/youdotcom/errors/getresearchtaskop.py b/src/youdotcom/errors/getresearchtaskop.py new file mode 100644 index 0000000..1c3cda1 --- /dev/null +++ b/src/youdotcom/errors/getresearchtaskop.py @@ -0,0 +1,92 @@ +"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" + +from __future__ import annotations +from dataclasses import dataclass, field +import httpx +from typing import Optional +from youdotcom.errors import YouError +from youdotcom.types import BaseModel + + +class GetResearchTaskInternalServerErrorData(BaseModel): + detail: Optional[str] = None + + +@dataclass(unsafe_hash=True) +class GetResearchTaskInternalServerError(YouError): + r"""Internal Server Error.""" + + data: GetResearchTaskInternalServerErrorData = field(hash=False) + + def __init__( + self, + data: GetResearchTaskInternalServerErrorData, + raw_response: httpx.Response, + body: Optional[str] = None, + ): + message = body or raw_response.text + super().__init__(message, raw_response, body) + object.__setattr__(self, "data", data) + + +class GetResearchTaskNotFoundErrorData(BaseModel): + detail: Optional[str] = None + + +@dataclass(unsafe_hash=True) +class GetResearchTaskNotFoundError(YouError): + r"""Task not found or not authorized.""" + + data: GetResearchTaskNotFoundErrorData = field(hash=False) + + def __init__( + self, + data: GetResearchTaskNotFoundErrorData, + raw_response: httpx.Response, + body: Optional[str] = None, + ): + message = body or raw_response.text + super().__init__(message, raw_response, body) + object.__setattr__(self, "data", data) + + +class GetResearchTaskForbiddenErrorData(BaseModel): + detail: Optional[str] = None + + +@dataclass(unsafe_hash=True) +class GetResearchTaskForbiddenError(YouError): + r"""Forbidden. API key lacks scope for this path.""" + + data: GetResearchTaskForbiddenErrorData = field(hash=False) + + def __init__( + self, + data: GetResearchTaskForbiddenErrorData, + raw_response: httpx.Response, + body: Optional[str] = None, + ): + message = body or raw_response.text + super().__init__(message, raw_response, body) + object.__setattr__(self, "data", data) + + +class GetResearchTaskUnauthorizedErrorData(BaseModel): + detail: Optional[str] = None + + +@dataclass(unsafe_hash=True) +class GetResearchTaskUnauthorizedError(YouError): + r"""Unauthorized. Problems with API key.""" + + data: GetResearchTaskUnauthorizedErrorData = field(hash=False) + + def __init__( + self, + data: GetResearchTaskUnauthorizedErrorData, + raw_response: httpx.Response, + body: Optional[str] = None, + ): + message = body or raw_response.text + super().__init__(message, raw_response, body) + object.__setattr__(self, "data", data) diff --git a/src/youdotcom/errors/streamresearchtaskop.py b/src/youdotcom/errors/streamresearchtaskop.py new file mode 100644 index 0000000..a2ee4a0 --- /dev/null +++ b/src/youdotcom/errors/streamresearchtaskop.py @@ -0,0 +1,92 @@ +"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" + +from __future__ import annotations +from dataclasses import dataclass, field +import httpx +from typing import Optional +from youdotcom.errors import YouError +from youdotcom.types import BaseModel + + +class StreamResearchTaskInternalServerErrorData(BaseModel): + detail: Optional[str] = None + + +@dataclass(unsafe_hash=True) +class StreamResearchTaskInternalServerError(YouError): + r"""Internal Server Error.""" + + data: StreamResearchTaskInternalServerErrorData = field(hash=False) + + def __init__( + self, + data: StreamResearchTaskInternalServerErrorData, + raw_response: httpx.Response, + body: Optional[str] = None, + ): + message = body or raw_response.text + super().__init__(message, raw_response, body) + object.__setattr__(self, "data", data) + + +class StreamResearchTaskNotFoundErrorData(BaseModel): + detail: Optional[str] = None + + +@dataclass(unsafe_hash=True) +class StreamResearchTaskNotFoundError(YouError): + r"""Task not found or not authorized.""" + + data: StreamResearchTaskNotFoundErrorData = field(hash=False) + + def __init__( + self, + data: StreamResearchTaskNotFoundErrorData, + raw_response: httpx.Response, + body: Optional[str] = None, + ): + message = body or raw_response.text + super().__init__(message, raw_response, body) + object.__setattr__(self, "data", data) + + +class StreamResearchTaskForbiddenErrorData(BaseModel): + detail: Optional[str] = None + + +@dataclass(unsafe_hash=True) +class StreamResearchTaskForbiddenError(YouError): + r"""Forbidden. API key lacks scope for this path.""" + + data: StreamResearchTaskForbiddenErrorData = field(hash=False) + + def __init__( + self, + data: StreamResearchTaskForbiddenErrorData, + raw_response: httpx.Response, + body: Optional[str] = None, + ): + message = body or raw_response.text + super().__init__(message, raw_response, body) + object.__setattr__(self, "data", data) + + +class StreamResearchTaskUnauthorizedErrorData(BaseModel): + detail: Optional[str] = None + + +@dataclass(unsafe_hash=True) +class StreamResearchTaskUnauthorizedError(YouError): + r"""Unauthorized. Problems with API key.""" + + data: StreamResearchTaskUnauthorizedErrorData = field(hash=False) + + def __init__( + self, + data: StreamResearchTaskUnauthorizedErrorData, + raw_response: httpx.Response, + body: Optional[str] = None, + ): + message = body or raw_response.text + super().__init__(message, raw_response, body) + object.__setattr__(self, "data", data) diff --git a/src/youdotcom/models/__init__.py b/src/youdotcom/models/__init__.py index 0aa0f8b..ad650a3 100644 --- a/src/youdotcom/models/__init__.py +++ b/src/youdotcom/models/__init__.py @@ -85,6 +85,10 @@ from .financeresearcheffort import FinanceResearchEffort from .freshness import Freshness from .freshnessvalue import FreshnessValue, FreshnessValueTypedDict + from .getresearchtaskop import ( + GetResearchTaskRequest, + GetResearchTaskRequestTypedDict, + ) from .language import Language from .livecrawl import LiveCrawl from .livecrawlformats import LiveCrawlFormats @@ -102,6 +106,8 @@ ResearchLocTypedDict, ResearchRequest, ResearchRequestTypedDict, + ResearchResult, + ResearchResultTypedDict, SourceControl, SourceControlTypedDict, ) @@ -116,6 +122,13 @@ Source, SourceTypedDict, ) + from .researchtaskstreamevent import ( + Event, + ResearchTaskStreamEvent, + ResearchTaskStreamEventData, + ResearchTaskStreamEventDataTypedDict, + ResearchTaskStreamEventTypedDict, + ) from .researchtool import ResearchTool, ResearchToolTypedDict from .response_created import ResponseCreated, ResponseCreatedTypedDict from .response_done import ( @@ -162,6 +175,20 @@ SearchResponseTypedDict, ) from .security import Security, SecurityTypedDict + from .streamresearchtaskop import ( + StreamResearchTaskRequest, + StreamResearchTaskRequestTypedDict, + ) + from .taskdetail import ( + Result, + ResultTypedDict, + TaskDetail, + TaskDetailInput, + TaskDetailInputTypedDict, + TaskDetailStatus, + TaskDetailTypedDict, + ) + from .taskresponse import TaskResponse, TaskResponseStatus, TaskResponseTypedDict from .verbosity import Verbosity from .webresult import WebResult, WebResultTypedDict from .websearchtool import WebSearchTool, WebSearchToolTypedDict @@ -204,6 +231,7 @@ "DataTypedDict", "Detail", "DetailTypedDict", + "Event", "ExpressAgentRunsRequest", "ExpressAgentRunsRequestTypedDict", "FinanceResearchContentType", @@ -227,6 +255,8 @@ "Freshness", "FreshnessValue", "FreshnessValueTypedDict", + "GetResearchTaskRequest", + "GetResearchTaskRequestTypedDict", "Input", "InputTypedDict", "Language", @@ -251,7 +281,13 @@ "ResearchRequest", "ResearchRequestTypedDict", "ResearchResponse", + "ResearchResult", + "ResearchResultTypedDict", "ResearchResponseTypedDict", + "ResearchTaskStreamEvent", + "ResearchTaskStreamEventData", + "ResearchTaskStreamEventDataTypedDict", + "ResearchTaskStreamEventTypedDict", "ResearchTool", "ResearchToolTypedDict", "ResponseCreated", @@ -278,6 +314,8 @@ "ResponseOutputTextDeltaTypedDict", "ResponseStarting", "ResponseStartingTypedDict", + "Result", + "ResultTypedDict", "Results", "ResultsTypedDict", "Role", @@ -299,6 +337,16 @@ "SourceControl", "SourceControlTypedDict", "SourceTypedDict", + "StreamResearchTaskRequest", + "StreamResearchTaskRequestTypedDict", + "TaskDetail", + "TaskDetailInput", + "TaskDetailInputTypedDict", + "TaskDetailStatus", + "TaskDetailTypedDict", + "TaskResponse", + "TaskResponseStatus", + "TaskResponseTypedDict", "Tool", "ToolTypedDict", "Type", @@ -379,6 +427,8 @@ "Freshness": ".freshness", "FreshnessValue": ".freshnessvalue", "FreshnessValueTypedDict": ".freshnessvalue", + "GetResearchTaskRequest": ".getresearchtaskop", + "GetResearchTaskRequestTypedDict": ".getresearchtaskop", "Language": ".language", "LiveCrawl": ".livecrawl", "LiveCrawlFormats": ".livecrawlformats", @@ -396,6 +446,8 @@ "ResearchLocTypedDict": ".researchop", "ResearchRequest": ".researchop", "ResearchRequestTypedDict": ".researchop", + "ResearchResult": ".researchop", + "ResearchResultTypedDict": ".researchop", "SourceControl": ".researchop", "SourceControlTypedDict": ".researchop", "Content": ".researchresponse", @@ -407,6 +459,11 @@ "ResearchResponseTypedDict": ".researchresponse", "Source": ".researchresponse", "SourceTypedDict": ".researchresponse", + "Event": ".researchtaskstreamevent", + "ResearchTaskStreamEvent": ".researchtaskstreamevent", + "ResearchTaskStreamEventData": ".researchtaskstreamevent", + "ResearchTaskStreamEventDataTypedDict": ".researchtaskstreamevent", + "ResearchTaskStreamEventTypedDict": ".researchtaskstreamevent", "ResearchTool": ".researchtool", "ResearchToolTypedDict": ".researchtool", "ResponseCreated": ".response_created", @@ -449,6 +506,18 @@ "SearchResponseTypedDict": ".searchresponse", "Security": ".security", "SecurityTypedDict": ".security", + "StreamResearchTaskRequest": ".streamresearchtaskop", + "StreamResearchTaskRequestTypedDict": ".streamresearchtaskop", + "TaskDetail": ".taskdetail", + "TaskDetailInput": ".taskdetail", + "TaskDetailInputTypedDict": ".taskdetail", + "TaskDetailStatus": ".taskdetail", + "TaskDetailTypedDict": ".taskdetail", + "Result": ".taskdetail", + "ResultTypedDict": ".taskdetail", + "TaskResponse": ".taskresponse", + "TaskResponseStatus": ".taskresponse", + "TaskResponseTypedDict": ".taskresponse", "Verbosity": ".verbosity", "WebResult": ".webresult", "WebResultTypedDict": ".webresult", diff --git a/src/youdotcom/models/getresearchtaskop.py b/src/youdotcom/models/getresearchtaskop.py new file mode 100644 index 0000000..0330943 --- /dev/null +++ b/src/youdotcom/models/getresearchtaskop.py @@ -0,0 +1,18 @@ +"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" + +from __future__ import annotations +from typing_extensions import Annotated, TypedDict +from youdotcom.types import BaseModel +from youdotcom.utils import FieldMetadata, PathParamMetadata + + +class GetResearchTaskRequestTypedDict(TypedDict): + task_id: str + r"""The UUID of the research task.""" + + +class GetResearchTaskRequest(BaseModel): + task_id: Annotated[ + str, FieldMetadata(path=PathParamMetadata(style="simple", explode=False)) + ] + r"""The UUID of the research task.""" diff --git a/src/youdotcom/models/researchop.py b/src/youdotcom/models/researchop.py index 15231ba..58e2443 100644 --- a/src/youdotcom/models/researchop.py +++ b/src/youdotcom/models/researchop.py @@ -2,6 +2,8 @@ from __future__ import annotations from .researcheffort import ResearchEffort +from .researchresponse import ResearchResponse, ResearchResponseTypedDict +from .taskresponse import TaskResponse, TaskResponseTypedDict from pydantic import model_serializer from typing import Any, Dict, List, Optional, Union from typing_extensions import NotRequired, TypeAliasType, TypedDict @@ -87,6 +89,8 @@ class ResearchRequestTypedDict(TypedDict): - `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. """ + background: NotRequired[bool] + r"""When true, queue a research task and return a task handle immediately instead of waiting for the result inline. Defaults to synchronous. When enabled, the response is a TaskResponse object with a task_id and stream_url for polling progress via GET /v1/research/{task_id} or streaming via GET /v1/research/{task_id}/stream.""" source_control: NotRequired[SourceControlTypedDict] r"""Beta. Controls which web sources the research agent searches and visits. Use this to allow specific domains, block specific domains, boost specific domains, filter by recency, or focus web results by country. @@ -118,6 +122,9 @@ class ResearchRequest(BaseModel): - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. """ + background: Optional[bool] = False + r"""When true, queue a research task and return a task handle immediately instead of waiting for the result inline. Defaults to synchronous. When enabled, the response is a TaskResponse object with a task_id and stream_url for polling progress via GET /v1/research/{task_id} or streaming via GET /v1/research/{task_id}/stream.""" + source_control: Optional[SourceControl] = None r"""Beta. Controls which web sources the research agent searches and visits. Use this to allow specific domains, block specific domains, boost specific domains, filter by recency, or focus web results by country. @@ -134,7 +141,9 @@ class ResearchRequest(BaseModel): @model_serializer(mode="wrap") def serialize_model(self, handler): - optional_fields = set(["research_effort", "source_control", "output_schema"]) + optional_fields = set( + ["research_effort", "background", "source_control", "output_schema"] + ) serialized = handler(self) m = {} @@ -217,3 +226,16 @@ def serialize_model(self, handler): m[k] = val return m + + +ResearchResultTypedDict = TypeAliasType( + "ResearchResultTypedDict", + Union[ResearchResponseTypedDict, TaskResponseTypedDict], +) +r"""A JSON object containing a comprehensive answer with citations and supporting search results. When background=true, returns a task handle instead.""" + + +ResearchResult = TypeAliasType( + "ResearchResult", Union[ResearchResponse, TaskResponse] +) +r"""A JSON object containing a comprehensive answer with citations and supporting search results. When background=true, returns a task handle instead.""" diff --git a/src/youdotcom/models/researchtaskstreamevent.py b/src/youdotcom/models/researchtaskstreamevent.py new file mode 100644 index 0000000..fe7f123 --- /dev/null +++ b/src/youdotcom/models/researchtaskstreamevent.py @@ -0,0 +1,113 @@ +"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" + +from __future__ import annotations +from enum import Enum +from pydantic import model_serializer +from typing import Any, Dict, Optional +from typing_extensions import NotRequired, TypedDict +from youdotcom.types import BaseModel, Nullable, OptionalNullable, UNSET, UNSET_SENTINEL + + +class Event(str, Enum): + r"""The type of the SSE event. Terminal events that close the stream are: `response.done`, `complete`, `error`, and `cancelled`. The stream may also emit `completed`, `failed`, or `cancelled` as event names corresponding to the task's terminal status. + """ + + CONNECTED = "connected" + RESPONSE_DONE = "response.done" + COMPLETE = "complete" + COMPLETED = "completed" + ERROR = "error" + FAILED = "failed" + CANCELLED = "cancelled" + + +class ResearchTaskStreamEventDataTypedDict(TypedDict): + r"""The event payload. Structure varies by event type. Common fields include type, task_id, status, data (event-specific), error, and sequence.""" + + type: NotRequired[str] + r"""The event type identifier.""" + task_id: NotRequired[str] + r"""The task UUID.""" + status: NotRequired[str] + r"""Current task status when the event was emitted.""" + data: NotRequired[Dict[str, Any]] + r"""Event-specific payload data.""" + error: NotRequired[Nullable[str]] + r"""Error message if the event represents an error.""" + sequence: NotRequired[int] + r"""Event sequence number.""" + + +class ResearchTaskStreamEventData(BaseModel): + r"""The event payload. Structure varies by event type. Common fields include type, task_id, status, data (event-specific), error, and sequence.""" + + type: Optional[str] = None + r"""The event type identifier.""" + + task_id: Optional[str] = None + r"""The task UUID.""" + + status: Optional[str] = None + r"""Current task status when the event was emitted.""" + + data: Optional[Dict[str, Any]] = None + r"""Event-specific payload data.""" + + error: OptionalNullable[str] = UNSET + r"""Error message if the event represents an error.""" + + sequence: Optional[int] = None + r"""Event sequence number.""" + + @model_serializer(mode="wrap") + def serialize_model(self, handler): + optional_fields = set( + ["type", "task_id", "status", "data", "error", "sequence"] + ) + nullable_fields = set(["error"]) + serialized = handler(self) + m = {} + + for n, f in type(self).model_fields.items(): + k = f.alias or n + val = serialized.get(k, serialized.get(n)) + is_nullable_and_explicitly_set = ( + k in nullable_fields + and (self.__pydantic_fields_set__.intersection({n})) # pylint: disable=no-member + ) + + if val != UNSET_SENTINEL: + if ( + val is not None + or k not in optional_fields + or is_nullable_and_explicitly_set + ): + m[k] = val + + return m + + +class ResearchTaskStreamEventTypedDict(TypedDict): + r"""A server-sent event for a background research task stream.""" + + id: str + r"""Sequence number of the SSE event.""" + event: Event + r"""The type of the SSE event. Terminal events that close the stream are: `response.done`, `complete`, `error`, and `cancelled`. The stream may also emit `completed`, `failed`, or `cancelled` as event names corresponding to the task's terminal status. + """ + data: ResearchTaskStreamEventDataTypedDict + r"""The event payload. Structure varies by event type. Common fields include type, task_id, status, data (event-specific), error, and sequence.""" + + +class ResearchTaskStreamEvent(BaseModel): + r"""A server-sent event for a background research task stream.""" + + id: str + r"""Sequence number of the SSE event.""" + + event: Event + r"""The type of the SSE event. Terminal events that close the stream are: `response.done`, `complete`, `error`, and `cancelled`. The stream may also emit `completed`, `failed`, or `cancelled` as event names corresponding to the task's terminal status. + """ + + data: ResearchTaskStreamEventData + r"""The event payload. Structure varies by event type. Common fields include type, task_id, status, data (event-specific), error, and sequence.""" diff --git a/src/youdotcom/models/streamresearchtaskop.py b/src/youdotcom/models/streamresearchtaskop.py new file mode 100644 index 0000000..6b57b1d --- /dev/null +++ b/src/youdotcom/models/streamresearchtaskop.py @@ -0,0 +1,44 @@ +"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" + +from __future__ import annotations +from pydantic import model_serializer +from typing import Optional +from typing_extensions import Annotated, NotRequired, TypedDict +from youdotcom.types import BaseModel, UNSET_SENTINEL +from youdotcom.utils import FieldMetadata, PathParamMetadata, QueryParamMetadata + + +class StreamResearchTaskRequestTypedDict(TypedDict): + task_id: str + r"""The UUID of the research task.""" + from_id: NotRequired[int] + r"""Resume from a sequence number for reconnection.""" + + +class StreamResearchTaskRequest(BaseModel): + task_id: Annotated[ + str, FieldMetadata(path=PathParamMetadata(style="simple", explode=False)) + ] + r"""The UUID of the research task.""" + + from_id: Annotated[ + Optional[int], + FieldMetadata(query=QueryParamMetadata(style="form", explode=True)), + ] = 0 + r"""Resume from a sequence number for reconnection.""" + + @model_serializer(mode="wrap") + def serialize_model(self, handler): + optional_fields = set(["from_id"]) + serialized = handler(self) + m = {} + + for n, f in type(self).model_fields.items(): + k = f.alias or n + val = serialized.get(k, serialized.get(n)) + + if val != UNSET_SENTINEL: + if val is not None or k not in optional_fields: + m[k] = val + + return m diff --git a/src/youdotcom/models/taskdetail.py b/src/youdotcom/models/taskdetail.py new file mode 100644 index 0000000..f039294 --- /dev/null +++ b/src/youdotcom/models/taskdetail.py @@ -0,0 +1,111 @@ +"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" + +from __future__ import annotations +from datetime import datetime +from enum import Enum +from pydantic import ConfigDict, model_serializer +from typing_extensions import NotRequired, TypedDict +from youdotcom.types import BaseModel, Nullable, OptionalNullable, UNSET, UNSET_SENTINEL + + +class TaskDetailStatus(str, Enum): + r"""Current status of the task.""" + + QUEUED = "queued" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class TaskDetailInputTypedDict(TypedDict): + r"""The original request input for the task.""" + + +class TaskDetailInput(BaseModel): + r"""The original request input for the task.""" + model_config = ConfigDict(extra="allow") + + +class ResultTypedDict(TypedDict): + r"""The task result when completed. For research tasks, this contains the ResearchResponse output.""" + + +class Result(BaseModel): + r"""The task result when completed. For research tasks, this contains the ResearchResponse output.""" + model_config = ConfigDict(extra="allow") + + +class TaskDetailTypedDict(TypedDict): + id: str + r"""Unique identifier for the task.""" + task_type: str + r"""Task type.""" + status: TaskDetailStatus + r"""Current status of the task.""" + created_at: datetime + r"""When the task was created.""" + updated_at: datetime + r"""When the task was last updated.""" + completed_at: NotRequired[Nullable[datetime]] + r"""When the task completed, if applicable.""" + error: NotRequired[Nullable[str]] + r"""Error message if the task failed.""" + input: NotRequired[Nullable[TaskDetailInputTypedDict]] + r"""The original request input for the task.""" + result: NotRequired[Nullable[ResultTypedDict]] + r"""The task result when completed. For research tasks, this contains the ResearchResponse output.""" + + +class TaskDetail(BaseModel): + id: str + r"""Unique identifier for the task.""" + + task_type: str + r"""Task type.""" + + status: TaskDetailStatus + r"""Current status of the task.""" + + created_at: datetime + r"""When the task was created.""" + + updated_at: datetime + r"""When the task was last updated.""" + + completed_at: OptionalNullable[datetime] = UNSET + r"""When the task completed, if applicable.""" + + error: OptionalNullable[str] = UNSET + r"""Error message if the task failed.""" + + input: OptionalNullable[TaskDetailInput] = UNSET + r"""The original request input for the task.""" + + result: OptionalNullable[Result] = UNSET + r"""The task result when completed. For research tasks, this contains the ResearchResponse output.""" + + @model_serializer(mode="wrap") + def serialize_model(self, handler): + optional_fields = set(["completed_at", "error", "input", "result"]) + nullable_fields = set(["completed_at", "error", "input", "result"]) + serialized = handler(self) + m = {} + + for n, f in type(self).model_fields.items(): + k = f.alias or n + val = serialized.get(k, serialized.get(n)) + is_nullable_and_explicitly_set = ( + k in nullable_fields + and (self.__pydantic_fields_set__.intersection({n})) # pylint: disable=no-member + ) + + if val != UNSET_SENTINEL: + if ( + val is not None + or k not in optional_fields + or is_nullable_and_explicitly_set + ): + m[k] = val + + return m diff --git a/src/youdotcom/models/taskresponse.py b/src/youdotcom/models/taskresponse.py new file mode 100644 index 0000000..490476d --- /dev/null +++ b/src/youdotcom/models/taskresponse.py @@ -0,0 +1,47 @@ +"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" + +from __future__ import annotations +from datetime import datetime +from enum import Enum +from typing_extensions import TypedDict +from youdotcom.types import BaseModel + + +class TaskResponseStatus(str, Enum): + r"""Current status of the task.""" + + QUEUED = "queued" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class TaskResponseTypedDict(TypedDict): + task_id: str + r"""Unique identifier for the task.""" + type: str + r"""Task type.""" + status: TaskResponseStatus + r"""Current status of the task.""" + stream_url: str + r"""URL to stream task events via SSE.""" + created_at: datetime + r"""When the task was created.""" + + +class TaskResponse(BaseModel): + task_id: str + r"""Unique identifier for the task.""" + + type: str + r"""Task type.""" + + status: TaskResponseStatus + r"""Current status of the task.""" + + stream_url: str + r"""URL to stream task events via SSE.""" + + created_at: datetime + r"""When the task was created.""" diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py new file mode 100644 index 0000000..1f62982 --- /dev/null +++ b/src/youdotcom/research_helpers.py @@ -0,0 +1,690 @@ +"""Hand-maintained research workflow helpers. + +This module is NOT regenerated by Speakeasy. It adds convenience helpers on +top of the auto-generated research endpoints: + +- ``research_background`` / ``research_background_async``: Submit a research + task with ``background=True`` and return the ``TaskResponse`` directly so + callers don't need to narrow the ``Union[ResearchResponse, TaskResponse]`` + response shape. +- ``poll_research_task`` / ``poll_research_task_async``: Poll + ``GET /v1/research/{task_id}`` until the task reaches a terminal state + (``completed``, ``failed``, ``cancelled``). +- ``research_and_wait`` / ``research_and_wait_async``: Submit a research task + with ``background=True``, then stream SSE events until a terminal event + arrives, and fetch the final ``TaskDetail``. +- ``stream_research`` / ``stream_research_async``: Iterate the SSE event + stream with a tolerant decoder that surfaces non-typed event names + (``research.searching``, etc.) as raw dicts instead of crashing on + validation. Use this when the server may emit event types outside the + documented enum. +""" + +from __future__ import annotations + +import asyncio +import json +import threading +import time +from dataclasses import dataclass +from typing import Any, AsyncIterator, Iterator, Mapping, Optional + +import youdotcom.models as _models +from youdotcom import errors as _errors +from youdotcom._hooks.types import HookContext +from youdotcom.sdk import You +from youdotcom.models import TaskDetail, TaskResponse +from youdotcom.utils import eventstreaming, get_security_from_env, match_response, match_status_codes +from youdotcom.utils.serializers import stream_to_text, stream_to_text_async +from youdotcom.utils.unmarshal_json_response import unmarshal_json_response + + +_DEFAULT_POLL_INTERVAL_S = 2.0 +_DEFAULT_POLL_TIMEOUT_S = 600.0 +_TERMINAL_TASK_STATUSES = frozenset({"completed", "failed", "cancelled"}) +_TERMINAL_STREAM_EVENTS_OK = frozenset({"response.done", "complete", "completed"}) +_TERMINAL_STREAM_EVENTS_ERR = frozenset({"error", "failed", "cancelled"}) + + +@dataclass +class RawStreamEvent: + """Tolerant SSE event representation. + + Attributes: + id: SSE event ID (``id:`` directive) when supplied. + event: Raw event name as received from the server, NOT validated + against the documented enum. + data: Parsed JSON payload when the data line is valid JSON, + otherwise the raw string. + retry: Optional ``retry:`` directive from the server, if any. + """ + + id: Optional[str] + event: Optional[str] + data: Any + retry: Optional[int] + + +def _decode_raw_event(raw_json: str) -> RawStreamEvent: + """Tolerant SSE decoder used by ``stream_research``. + + Accepts any JSON object regardless of whether ``event`` matches the + documented enum, so unknown workflow events pass through instead of + failing pydantic validation. + """ + parsed = json.loads(raw_json) + if not isinstance(parsed, dict): + return RawStreamEvent(id=None, event=None, data=parsed, retry=None) + return RawStreamEvent( + id=parsed.get("id"), + event=parsed.get("event"), + data=parsed.get("data"), + retry=parsed.get("retry"), + ) + + +def research_background(client: You, **kwargs: Any) -> TaskResponse: + """Submit a research task with ``background=True`` and return ``TaskResponse``. + + Equivalent to ``client.research(..., background=True)`` but asserts the + response is a ``TaskResponse`` so callers get a strong return type + without needing to narrow ``Union[ResearchResponse, TaskResponse]``. + """ + kw = dict(kwargs) + kw["background"] = True + res = client.research(**kw) + if not isinstance(res, TaskResponse): + raise TypeError( + "research_background requires background=True so the server " + "returns TaskResponse; got " + type(res).__name__ + ) + return res + + +async def research_background_async(client: You, **kwargs: Any) -> TaskResponse: + """Async variant of :func:`research_background`.""" + kw = dict(kwargs) + kw["background"] = True + res = await client.research_async(**kw) + if not isinstance(res, TaskResponse): + raise TypeError( + "research_background_async requires background=True so the " + "server returns TaskResponse; got " + type(res).__name__ + ) + return res + + +def _poll_detail( + client: You, + *, + task_id: str, + interval_s: float, + timeout_s: float, + deadline: float, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, + http_headers: Optional[Mapping[str, str]] = None, +) -> TaskDetail: + while True: + detail = client.get_research_task( + task_id=task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + status = detail.status.value + if status in _TERMINAL_TASK_STATUSES: + if status != "completed": + raise RuntimeError( + f"research task {task_id} ended in non-completed state: {status}" + ) + return detail + if time.monotonic() >= deadline: + raise TimeoutError( + f"research task {task_id} did not complete within {timeout_s}s" + ) + time.sleep(interval_s) + + +def poll_research_task( + client: You, + task_id: str, + *, + interval_s: float = _DEFAULT_POLL_INTERVAL_S, + timeout_s: float = _DEFAULT_POLL_TIMEOUT_S, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, + http_headers: Optional[Mapping[str, str]] = None, +) -> TaskDetail: + """Poll ``GET /v1/research/{task_id}`` until ``status == "completed"``. + + Raises ``RuntimeError`` if the task ends in a non-completed terminal state + (``failed`` / ``cancelled``) and ``TimeoutError`` if ``timeout_s`` elapses + before completion. + """ + return _poll_detail( + client, + task_id=task_id, + interval_s=interval_s, + timeout_s=timeout_s, + deadline=time.monotonic() + timeout_s, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + + +async def _poll_detail_async( + client: You, + *, + task_id: str, + interval_s: float, + timeout_s: float, + deadline: float, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, + http_headers: Optional[Mapping[str, str]] = None, +) -> TaskDetail: + while True: + detail = await client.get_research_task_async( + task_id=task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + status = detail.status.value + if status in _TERMINAL_TASK_STATUSES: + if status != "completed": + raise RuntimeError( + f"research task {task_id} ended in non-completed state: {status}" + ) + return detail + if time.monotonic() >= deadline: + raise TimeoutError( + f"research task {task_id} did not complete within {timeout_s}s" + ) + await asyncio.sleep(interval_s) + + +async def poll_research_task_async( + client: You, + task_id: str, + *, + interval_s: float = _DEFAULT_POLL_INTERVAL_S, + timeout_s: float = _DEFAULT_POLL_TIMEOUT_S, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, + http_headers: Optional[Mapping[str, str]] = None, +) -> TaskDetail: + """Async variant of :func:`poll_research_task`.""" + return await _poll_detail_async( + client, + task_id=task_id, + interval_s=interval_s, + timeout_s=timeout_s, + deadline=time.monotonic() + timeout_s, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + + +def research_and_wait( + client: You, + *, + timeout_s: float = _DEFAULT_POLL_TIMEOUT_S, + **kwargs: Any, +) -> TaskDetail: + """Submit a research task in background mode and wait for completion. + + Submits via :func:`research_background`, then opens the SSE stream + and reads events until a terminal event arrives. Fetches the final + ``TaskDetail`` via one ``get_research_task`` call. + + Parameters: + timeout_s: Maximum seconds to wait for a terminal stream event. + ``**kwargs``: Forwarded to ``client.research``, including + ``input`` and ``research_effort``. ``background=True`` is set + automatically. ``server_url``, ``timeout_ms``, and + ``http_headers`` are forwarded to all sub-requests. + + Returns: + ``TaskDetail`` with ``status == completed`` and the task's + ``result`` attribute populated. + + Raises: + TimeoutError: If no terminal event arrives within ``timeout_s`` + and the task is still non-terminal. + RuntimeError: If the task ends in a non-completed terminal state. + """ + server_url = kwargs.get("server_url") + timeout_ms = kwargs.get("timeout_ms") + http_headers = kwargs.get("http_headers") + task = research_background(client, **kwargs) + + stream = _open_raw_stream( + client, task.task_id, http_headers=http_headers, + server_url=server_url, timeout_ms=timeout_ms, + ) + result: Optional[str] = None + thread_exc: Optional[BaseException] = None + consumer: Optional[threading.Thread] = None + try: + def _consume() -> None: + nonlocal result, thread_exc + try: + for evt in stream: + if evt.event in _TERMINAL_STREAM_EVENTS_OK: + result = "ok" + return + if evt.event in _TERMINAL_STREAM_EVENTS_ERR: + result = evt.event + return + except BaseException as exc: # pylint: disable=broad-except + thread_exc = exc + + consumer = threading.Thread(target=_consume, daemon=True) + consumer.start() + consumer.join(timeout=timeout_s) + timed_out = consumer.is_alive() + finally: + stream.close() + if consumer is not None: + consumer.join(timeout=2.0) + + if timed_out: + # Stream didn't emit a terminal event within timeout_s. The task + # may have completed anyway (the SSE stream doesn't always emit + # terminal events for tasks that complete mid-stream). Do a final + # GET to check the actual status before raising. + detail = client.get_research_task( + task_id=task.task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + status = detail.status.value + if status == "completed": + return detail + if status in _TERMINAL_TASK_STATUSES: + raise RuntimeError( + f"research task {task.task_id} ended in non-completed state: {status}" + ) + raise TimeoutError( + f"research task {task.task_id} did not complete within {timeout_s}s " + f"(status: {status})" + ) + if thread_exc is not None: + raise thread_exc + if result == "ok": + detail = client.get_research_task( + task_id=task.task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + if detail.status.value != "completed": + raise RuntimeError( + f"research task {task.task_id} stream signalled completion " + f"but GET returned status={detail.status.value}" + ) + return detail + if result is not None: + raise RuntimeError( + f"research task {task.task_id} ended in non-completed state: {result}" + ) + # Stream closed without a terminal event — the task may have completed + # mid-stream. Do a final GET to check the actual status. + detail = client.get_research_task( + task_id=task.task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + status = detail.status.value + if status == "completed": + return detail + if status in _TERMINAL_TASK_STATUSES: + raise RuntimeError( + f"research task {task.task_id} ended in non-completed state: {status}" + ) + raise TimeoutError( + f"research task {task.task_id} stream closed without terminal event " + f"and task is still {status}" + ) + + +async def research_and_wait_async( + client: You, + *, + timeout_s: float = _DEFAULT_POLL_TIMEOUT_S, + **kwargs: Any, +) -> TaskDetail: + """Async variant of :func:`research_and_wait`.""" + server_url = kwargs.get("server_url") + timeout_ms = kwargs.get("timeout_ms") + http_headers = kwargs.get("http_headers") + task = await research_background_async(client, **kwargs) + + async def _consume() -> Optional[str]: + async for evt in stream_research_async( + client, task.task_id, + server_url=server_url, timeout_ms=timeout_ms, + http_headers=http_headers, + ): + if evt.event in _TERMINAL_STREAM_EVENTS_OK: + return "ok" + if evt.event in _TERMINAL_STREAM_EVENTS_ERR: + return evt.event + return None + + try: + result = await asyncio.wait_for(_consume(), timeout=timeout_s) + except asyncio.TimeoutError as exc: + # Stream didn't emit a terminal event within timeout_s. The task + # may have completed anyway — do a final GET to check. + detail = await client.get_research_task_async( + task_id=task.task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + status = detail.status.value + if status == "completed": + return detail + if status in _TERMINAL_TASK_STATUSES: + raise RuntimeError( + f"research task {task.task_id} ended in non-completed state: {status}" + ) from exc + raise TimeoutError( + f"research task {task.task_id} did not complete within {timeout_s}s " + f"(status: {status})" + ) from exc + + if result == "ok": + detail = await client.get_research_task_async( + task_id=task.task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + if detail.status.value != "completed": + raise RuntimeError( + f"research task {task.task_id} stream signalled completion " + f"but GET returned status={detail.status.value}" + ) + return detail + if result is not None: + raise RuntimeError( + f"research task {task.task_id} ended in non-completed state: {result}" + ) + # Stream closed without a terminal event — the task may have completed + # mid-stream. Do a final GET to check the actual status. + detail = await client.get_research_task_async( + task_id=task.task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + status = detail.status.value + if status == "completed": + return detail + if status in _TERMINAL_TASK_STATUSES: + raise RuntimeError( + f"research task {task.task_id} ended in non-completed state: {status}" + ) + raise TimeoutError( + f"research task {task.task_id} stream closed without terminal event " + f"and task is still {status}" + ) + + +def _raise_stream_error(http_res: Any) -> None: + """Raise the typed ``StreamResearchTask*Error`` for non-200 responses. + + Mirrors the error-handling branches in the generated + ``client.stream_research_task`` so that callers of ``stream_research`` + get the same structured exceptions (401/403/404/500) as the generated + method instead of a generic ``RuntimeError``. + """ + response_data: Any = None + if match_response(http_res, "401", "application/json"): + text = stream_to_text(http_res) + response_data = unmarshal_json_response( + _errors.StreamResearchTaskUnauthorizedErrorData, http_res, text + ) + raise _errors.StreamResearchTaskUnauthorizedError(response_data, http_res, text) + if match_response(http_res, "403", "application/json"): + text = stream_to_text(http_res) + response_data = unmarshal_json_response( + _errors.StreamResearchTaskForbiddenErrorData, http_res, text + ) + raise _errors.StreamResearchTaskForbiddenError(response_data, http_res, text) + if match_response(http_res, "404", "application/json"): + text = stream_to_text(http_res) + response_data = unmarshal_json_response( + _errors.StreamResearchTaskNotFoundErrorData, http_res, text + ) + raise _errors.StreamResearchTaskNotFoundError(response_data, http_res, text) + if match_response(http_res, "500", "application/json"): + text = stream_to_text(http_res) + response_data = unmarshal_json_response( + _errors.StreamResearchTaskInternalServerErrorData, http_res, text + ) + raise _errors.StreamResearchTaskInternalServerError(response_data, http_res, text) + if match_response(http_res, "4XX", "*") or match_response(http_res, "5XX", "*"): + text = stream_to_text(http_res) + raise _errors.YouDefaultError("API error occurred", http_res, text) + text = stream_to_text(http_res) + raise _errors.YouDefaultError("Unexpected response received", http_res, text) + + +async def _raise_stream_error_async(http_res: Any) -> None: + """Async counterpart of :func:`_raise_stream_error`.""" + response_data: Any = None + if match_response(http_res, "401", "application/json"): + text = await stream_to_text_async(http_res) + response_data = unmarshal_json_response( + _errors.StreamResearchTaskUnauthorizedErrorData, http_res, text + ) + raise _errors.StreamResearchTaskUnauthorizedError(response_data, http_res, text) + if match_response(http_res, "403", "application/json"): + text = await stream_to_text_async(http_res) + response_data = unmarshal_json_response( + _errors.StreamResearchTaskForbiddenErrorData, http_res, text + ) + raise _errors.StreamResearchTaskForbiddenError(response_data, http_res, text) + if match_response(http_res, "404", "application/json"): + text = await stream_to_text_async(http_res) + response_data = unmarshal_json_response( + _errors.StreamResearchTaskNotFoundErrorData, http_res, text + ) + raise _errors.StreamResearchTaskNotFoundError(response_data, http_res, text) + if match_response(http_res, "500", "application/json"): + text = await stream_to_text_async(http_res) + response_data = unmarshal_json_response( + _errors.StreamResearchTaskInternalServerErrorData, http_res, text + ) + raise _errors.StreamResearchTaskInternalServerError(response_data, http_res, text) + if match_response(http_res, "4XX", "*") or match_response(http_res, "5XX", "*"): + text = await stream_to_text_async(http_res) + raise _errors.YouDefaultError("API error occurred", http_res, text) + text = await stream_to_text_async(http_res) + raise _errors.YouDefaultError("Unexpected response received", http_res, text) + + +def _open_raw_stream( + client: You, + task_id: str, + *, + http_headers: Optional[Mapping[str, str]], + from_id: int = 0, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, +) -> eventstreaming.EventStream[RawStreamEvent]: + """Issue the SSE request directly so we can pass a tolerant decoder. + + Mirrors the request setup used by ``client.stream_research_task`` but + wires :func:`_decode_raw_event` into the ``EventStream`` instead of the + strict pydantic ``ResearchTaskStreamEvent`` decoder. + """ + # pylint: disable=protected-access + base_url = server_url if server_url is not None else client._get_url(None, None) + if timeout_ms is None: + timeout_ms = client.sdk_configuration.timeout_ms + req = client._build_request( + method="GET", + path="/v1/research/{task_id}/stream", + base_url=base_url, + url_variables=None, + request=_models.StreamResearchTaskRequest(task_id=task_id, from_id=from_id), + request_body_required=False, + request_has_path_params=True, + request_has_query_params=True, + user_agent_header="user-agent", + accept_header_value="text/event-stream", + http_headers=http_headers, + security=client.sdk_configuration.security, + allow_empty_value=None, + timeout_ms=timeout_ms, + ) + http_res = client.do_request( + hook_ctx=HookContext( + config=client.sdk_configuration, + base_url=base_url or "", + operation_id="streamResearchTask", + oauth2_scopes=None, + security_source=get_security_from_env( + client.sdk_configuration.security, _models.Security + ), + tags=None, + extensions=None, + ), + request=req, + is_error_status_code=lambda c: match_status_codes(["4XX", "5XX"], c), + stream=True, + retry_config=None, + ) + if not match_response(http_res, "200", "text/event-stream"): + _raise_stream_error(http_res) + return eventstreaming.EventStream( + http_res, + _decode_raw_event, + client_ref=client, + data_required=False, + ) + + +def stream_research( + client: You, + task_id: str, + *, + http_headers: Optional[Mapping[str, str]] = None, + from_id: int = 0, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, +) -> Iterator[RawStreamEvent]: + """SSE iterator that accepts unknown event names without failing. + + Each yielded item is a :class:`RawStreamEvent` carrying the raw ``event`` + string and parsed JSON ``data`` payload. Useful when the server may emit + intermediate workflow events with names outside the documented enum + (e.g. ``research.searching``, ``response.created``, + ``response.output_item.added``). + + Parameters: + from_id: Sequence number to resume from (reconnection). ``0`` starts at + the beginning of the stream. + """ + # pylint: disable-next=protected-access + with _open_raw_stream( + client, task_id, http_headers=http_headers, from_id=from_id, + server_url=server_url, timeout_ms=timeout_ms, + ) as stream: + yield from stream + + +async def _open_raw_stream_async( + client: You, + task_id: str, + *, + http_headers: Optional[Mapping[str, str]], + from_id: int = 0, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, +) -> eventstreaming.EventStreamAsync[RawStreamEvent]: + """Async counterpart of :func:`_open_raw_stream`.""" + # pylint: disable=protected-access + base_url = server_url if server_url is not None else client._get_url(None, None) + if timeout_ms is None: + timeout_ms = client.sdk_configuration.timeout_ms + req = client._build_request_async( + method="GET", + path="/v1/research/{task_id}/stream", + base_url=base_url, + url_variables=None, + request=_models.StreamResearchTaskRequest(task_id=task_id, from_id=from_id), + request_body_required=False, + request_has_path_params=True, + request_has_query_params=True, + user_agent_header="user-agent", + accept_header_value="text/event-stream", + http_headers=http_headers, + security=client.sdk_configuration.security, + allow_empty_value=None, + timeout_ms=timeout_ms, + ) + http_res = await client.do_request_async( + hook_ctx=HookContext( + config=client.sdk_configuration, + base_url=base_url or "", + operation_id="streamResearchTask", + oauth2_scopes=None, + security_source=get_security_from_env( + client.sdk_configuration.security, _models.Security + ), + tags=None, + extensions=None, + ), + request=req, + is_error_status_code=lambda c: match_status_codes(["4XX", "5XX"], c), + stream=True, + retry_config=None, + ) + if not match_response(http_res, "200", "text/event-stream"): + await _raise_stream_error_async(http_res) + return eventstreaming.EventStreamAsync( + http_res, + _decode_raw_event, + client_ref=client, + data_required=False, + ) + + +async def stream_research_async( + client: You, + task_id: str, + *, + http_headers: Optional[Mapping[str, str]] = None, + from_id: int = 0, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, +) -> AsyncIterator[RawStreamEvent]: + """Async variant of :func:`stream_research`. + + Parameters: + from_id: Sequence number to resume from (reconnection). ``0`` starts at + the beginning of the stream. + """ + # Close the underlying httpx SSE response deterministically on consumer + # break/return/throw. EventStreamAsync exposes async close() (not aclose), + # so a try/finally is used rather than contextlib.aclosing. + stream = await _open_raw_stream_async( + client, task_id, http_headers=http_headers, from_id=from_id, + server_url=server_url, timeout_ms=timeout_ms, + ) + try: + async for evt in stream: + yield evt + finally: + await stream.close() diff --git a/src/youdotcom/sdk.py b/src/youdotcom/sdk.py index 5d95cb9..153869c 100644 --- a/src/youdotcom/sdk.py +++ b/src/youdotcom/sdk.py @@ -24,7 +24,7 @@ from youdotcom import errors, models, utils from youdotcom._hooks import HookContext, SDKHooks from youdotcom.types import OptionalNullable, UNSET -from youdotcom.utils import get_security_from_env +from youdotcom.utils import eventstreaming, get_security_from_env from youdotcom.utils.unmarshal_json_response import unmarshal_json_response if TYPE_CHECKING: @@ -532,6 +532,7 @@ def research( research_effort: Optional[ models.ResearchEffort ] = models.ResearchEffort.STANDARD, + background: Optional[bool] = False, source_control: Optional[ Union[models.SourceControl, models.SourceControlTypedDict] ] = None, @@ -540,7 +541,7 @@ def research( server_url: Optional[str] = None, timeout_ms: Optional[int] = None, http_headers: Optional[Mapping[str, str]] = None, - ) -> models.ResearchResponse: + ) -> models.ResearchResult: r"""Returns comprehensive research-grade answers with multi-step reasoning Research goes beyond a single web search. In response to your question, it runs multiple searches, reads through the sources, and synthesizes everything into a thorough, well-cited answer. Use it when a question is too complex for a simple lookup, and when you need a response you can actually trust and verify. @@ -555,6 +556,7 @@ def research( - `standard`: The default. Balances speed and depth, a good fit for most questions. - `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. + :param background: When true, queue a research task and return a task handle immediately instead of waiting for the result inline. Defaults to synchronous. When enabled, the response is a TaskResponse object with a task_id and stream_url for polling progress via GET /v1/research/{task_id} or streaming via GET /v1/research/{task_id}/stream. :param source_control: Beta. Controls which web sources the research agent searches and visits. Use this to allow specific domains, block specific domains, boost specific domains, filter by recency, or focus web results by country. `include_domains` and `exclude_domains` cannot be used together. Each domain list is capped at 500 entries. `exclude_domains` also blocks the research agent from visiting pages on those domains during browsing. `boost_domains` gives matching domains a relative ranking boost without filtering out other domains. It can be combined with `exclude_domains` but cannot be combined with `include_domains`. @@ -581,6 +583,7 @@ def research( request = models.ResearchRequest( input=input, research_effort=research_effort, + background=background, source_control=utils.get_pydantic_model( source_control, Optional[models.SourceControl] ), @@ -634,7 +637,7 @@ def research( response_data: Any = None if utils.match_response(http_res, "200", "application/json"): - return unmarshal_json_response(models.ResearchResponse, http_res) + return unmarshal_json_response(models.ResearchResult, http_res) if utils.match_response(http_res, "401", "application/json"): response_data = unmarshal_json_response( errors.ResearchUnauthorizedErrorData, http_res @@ -671,6 +674,7 @@ async def research_async( research_effort: Optional[ models.ResearchEffort ] = models.ResearchEffort.STANDARD, + background: Optional[bool] = False, source_control: Optional[ Union[models.SourceControl, models.SourceControlTypedDict] ] = None, @@ -679,7 +683,7 @@ async def research_async( server_url: Optional[str] = None, timeout_ms: Optional[int] = None, http_headers: Optional[Mapping[str, str]] = None, - ) -> models.ResearchResponse: + ) -> models.ResearchResult: r"""Returns comprehensive research-grade answers with multi-step reasoning Research goes beyond a single web search. In response to your question, it runs multiple searches, reads through the sources, and synthesizes everything into a thorough, well-cited answer. Use it when a question is too complex for a simple lookup, and when you need a response you can actually trust and verify. @@ -694,6 +698,7 @@ async def research_async( - `standard`: The default. Balances speed and depth, a good fit for most questions. - `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. + :param background: When true, queue a research task and return a task handle immediately instead of waiting for the result inline. Defaults to synchronous. When enabled, the response is a TaskResponse object with a task_id and stream_url for polling progress via GET /v1/research/{task_id} or streaming via GET /v1/research/{task_id}/stream. :param source_control: Beta. Controls which web sources the research agent searches and visits. Use this to allow specific domains, block specific domains, boost specific domains, filter by recency, or focus web results by country. `include_domains` and `exclude_domains` cannot be used together. Each domain list is capped at 500 entries. `exclude_domains` also blocks the research agent from visiting pages on those domains during browsing. `boost_domains` gives matching domains a relative ranking boost without filtering out other domains. It can be combined with `exclude_domains` but cannot be combined with `include_domains`. @@ -720,6 +725,7 @@ async def research_async( request = models.ResearchRequest( input=input, research_effort=research_effort, + background=background, source_control=utils.get_pydantic_model( source_control, Optional[models.SourceControl] ), @@ -773,7 +779,7 @@ async def research_async( response_data: Any = None if utils.match_response(http_res, "200", "application/json"): - return unmarshal_json_response(models.ResearchResponse, http_res) + return unmarshal_json_response(models.ResearchResult, http_res) if utils.match_response(http_res, "401", "application/json"): response_data = unmarshal_json_response( errors.ResearchUnauthorizedErrorData, http_res @@ -803,6 +809,488 @@ async def research_async( raise errors.YouDefaultError("Unexpected response received", http_res) + def get_research_task( + self, + *, + task_id: str, + retries: OptionalNullable[utils.RetryConfig] = UNSET, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, + http_headers: Optional[Mapping[str, str]] = None, + ) -> models.TaskDetail: + r"""Get the status of a background research task + + Poll the status of a background research task created with background=true. When the task is completed, the result is included in the response. + + :param task_id: The UUID of the research task. + :param retries: Override the default retry configuration for this method + :param server_url: Override the default server URL for this method + :param timeout_ms: Override the default request timeout configuration for this method in milliseconds + :param http_headers: Additional headers to set or replace on requests. + """ + base_url = None + url_variables = None + if timeout_ms is None: + timeout_ms = self.sdk_configuration.timeout_ms + + if server_url is not None: + base_url = server_url + else: + base_url = self._get_url(base_url, url_variables) + + request = models.GetResearchTaskRequest( + task_id=task_id, + ) + + req = self._build_request( + method="GET", + path="/v1/research/{task_id}", + base_url=base_url, + url_variables=url_variables, + request=request, + request_body_required=False, + request_has_path_params=True, + request_has_query_params=True, + user_agent_header="user-agent", + accept_header_value="application/json", + http_headers=http_headers, + security=self.sdk_configuration.security, + allow_empty_value=None, + timeout_ms=timeout_ms, + ) + + if retries == UNSET: + if self.sdk_configuration.retry_config is not UNSET: + retries = self.sdk_configuration.retry_config + + retry_config = None + if isinstance(retries, utils.RetryConfig): + retry_config = (retries, ["429", "500", "502", "503", "504"]) + + http_res = self.do_request( + hook_ctx=HookContext( + config=self.sdk_configuration, + base_url=base_url or "", + operation_id="getResearchTask", + oauth2_scopes=None, + security_source=get_security_from_env( + self.sdk_configuration.security, models.Security + ), + tags=None, + extensions=None, + ), + request=req, + is_error_status_code=lambda c: utils.match_status_codes(["4XX", "5XX"], c), + retry_config=retry_config, + ) + + response_data: Any = None + if utils.match_response(http_res, "200", "application/json"): + return unmarshal_json_response(models.TaskDetail, http_res) + if utils.match_response(http_res, "401", "application/json"): + response_data = unmarshal_json_response( + errors.GetResearchTaskUnauthorizedErrorData, http_res + ) + raise errors.GetResearchTaskUnauthorizedError(response_data, http_res) + if utils.match_response(http_res, "403", "application/json"): + response_data = unmarshal_json_response( + errors.GetResearchTaskForbiddenErrorData, http_res + ) + raise errors.GetResearchTaskForbiddenError(response_data, http_res) + if utils.match_response(http_res, "404", "application/json"): + response_data = unmarshal_json_response( + errors.GetResearchTaskNotFoundErrorData, http_res + ) + raise errors.GetResearchTaskNotFoundError(response_data, http_res) + if utils.match_response(http_res, "500", "application/json"): + response_data = unmarshal_json_response( + errors.GetResearchTaskInternalServerErrorData, http_res + ) + raise errors.GetResearchTaskInternalServerError(response_data, http_res) + if utils.match_response(http_res, "4XX", "*"): + http_res_text = utils.stream_to_text(http_res) + raise errors.YouDefaultError("API error occurred", http_res, http_res_text) + if utils.match_response(http_res, "5XX", "*"): + http_res_text = utils.stream_to_text(http_res) + raise errors.YouDefaultError("API error occurred", http_res, http_res_text) + + raise errors.YouDefaultError("Unexpected response received", http_res) + + async def get_research_task_async( + self, + *, + task_id: str, + retries: OptionalNullable[utils.RetryConfig] = UNSET, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, + http_headers: Optional[Mapping[str, str]] = None, + ) -> models.TaskDetail: + r"""Get the status of a background research task + + Poll the status of a background research task created with background=true. When the task is completed, the result is included in the response. + + :param task_id: The UUID of the research task. + :param retries: Override the default retry configuration for this method + :param server_url: Override the default server URL for this method + :param timeout_ms: Override the default request timeout configuration for this method in milliseconds + :param http_headers: Additional headers to set or replace on requests. + """ + base_url = None + url_variables = None + if timeout_ms is None: + timeout_ms = self.sdk_configuration.timeout_ms + + if server_url is not None: + base_url = server_url + else: + base_url = self._get_url(base_url, url_variables) + + request = models.GetResearchTaskRequest( + task_id=task_id, + ) + + req = self._build_request_async( + method="GET", + path="/v1/research/{task_id}", + base_url=base_url, + url_variables=url_variables, + request=request, + request_body_required=False, + request_has_path_params=True, + request_has_query_params=True, + user_agent_header="user-agent", + accept_header_value="application/json", + http_headers=http_headers, + security=self.sdk_configuration.security, + allow_empty_value=None, + timeout_ms=timeout_ms, + ) + + if retries == UNSET: + if self.sdk_configuration.retry_config is not UNSET: + retries = self.sdk_configuration.retry_config + + retry_config = None + if isinstance(retries, utils.RetryConfig): + retry_config = (retries, ["429", "500", "502", "503", "504"]) + + http_res = await self.do_request_async( + hook_ctx=HookContext( + config=self.sdk_configuration, + base_url=base_url or "", + operation_id="getResearchTask", + oauth2_scopes=None, + security_source=get_security_from_env( + self.sdk_configuration.security, models.Security + ), + tags=None, + extensions=None, + ), + request=req, + is_error_status_code=lambda c: utils.match_status_codes(["4XX", "5XX"], c), + retry_config=retry_config, + ) + + response_data: Any = None + if utils.match_response(http_res, "200", "application/json"): + return unmarshal_json_response(models.TaskDetail, http_res) + if utils.match_response(http_res, "401", "application/json"): + response_data = unmarshal_json_response( + errors.GetResearchTaskUnauthorizedErrorData, http_res + ) + raise errors.GetResearchTaskUnauthorizedError(response_data, http_res) + if utils.match_response(http_res, "403", "application/json"): + response_data = unmarshal_json_response( + errors.GetResearchTaskForbiddenErrorData, http_res + ) + raise errors.GetResearchTaskForbiddenError(response_data, http_res) + if utils.match_response(http_res, "404", "application/json"): + response_data = unmarshal_json_response( + errors.GetResearchTaskNotFoundErrorData, http_res + ) + raise errors.GetResearchTaskNotFoundError(response_data, http_res) + if utils.match_response(http_res, "500", "application/json"): + response_data = unmarshal_json_response( + errors.GetResearchTaskInternalServerErrorData, http_res + ) + raise errors.GetResearchTaskInternalServerError(response_data, http_res) + if utils.match_response(http_res, "4XX", "*"): + http_res_text = await utils.stream_to_text_async(http_res) + raise errors.YouDefaultError("API error occurred", http_res, http_res_text) + if utils.match_response(http_res, "5XX", "*"): + http_res_text = await utils.stream_to_text_async(http_res) + raise errors.YouDefaultError("API error occurred", http_res, http_res_text) + + raise errors.YouDefaultError("Unexpected response received", http_res) + + def stream_research_task( + self, + *, + task_id: str, + from_id: Optional[int] = 0, + retries: OptionalNullable[utils.RetryConfig] = UNSET, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, + http_headers: Optional[Mapping[str, str]] = None, + ) -> eventstreaming.EventStream[models.ResearchTaskStreamEvent]: + r"""Stream updates for a background research task + + Stream real-time updates for a background research task via Server-Sent Events (SSE). Supports reconnection via the from_id query parameter to replay missed events. The connection closes automatically when the task reaches a terminal state. + + :param task_id: The UUID of the research task. + :param from_id: Resume from a sequence number for reconnection. + :param retries: Override the default retry configuration for this method + :param server_url: Override the default server URL for this method + :param timeout_ms: Override the default request timeout configuration for this method in milliseconds + :param http_headers: Additional headers to set or replace on requests. + """ + base_url = None + url_variables = None + if timeout_ms is None: + timeout_ms = self.sdk_configuration.timeout_ms + + if server_url is not None: + base_url = server_url + else: + base_url = self._get_url(base_url, url_variables) + + request = models.StreamResearchTaskRequest( + task_id=task_id, + from_id=from_id, + ) + + req = self._build_request( + method="GET", + path="/v1/research/{task_id}/stream", + base_url=base_url, + url_variables=url_variables, + request=request, + request_body_required=False, + request_has_path_params=True, + request_has_query_params=True, + user_agent_header="user-agent", + accept_header_value="text/event-stream", + http_headers=http_headers, + security=self.sdk_configuration.security, + allow_empty_value=None, + timeout_ms=timeout_ms, + ) + + if retries == UNSET: + if self.sdk_configuration.retry_config is not UNSET: + retries = self.sdk_configuration.retry_config + + retry_config = None + if isinstance(retries, utils.RetryConfig): + retry_config = (retries, ["429", "500", "502", "503", "504"]) + + http_res = self.do_request( + hook_ctx=HookContext( + config=self.sdk_configuration, + base_url=base_url or "", + operation_id="streamResearchTask", + oauth2_scopes=None, + security_source=get_security_from_env( + self.sdk_configuration.security, models.Security + ), + tags=None, + extensions=None, + ), + request=req, + is_error_status_code=lambda c: utils.match_status_codes(["4XX", "5XX"], c), + stream=True, + retry_config=retry_config, + ) + + response_data: Any = None + if utils.match_response(http_res, "200", "text/event-stream"): + return eventstreaming.EventStream( + http_res, + lambda raw: unmarshal_json_response( + models.ResearchTaskStreamEvent, http_res, raw + ), + client_ref=self, + ) + if utils.match_response(http_res, "401", "application/json"): + http_res_text = utils.stream_to_text(http_res) + response_data = unmarshal_json_response( + errors.StreamResearchTaskUnauthorizedErrorData, http_res, http_res_text + ) + raise errors.StreamResearchTaskUnauthorizedError( + response_data, http_res, http_res_text + ) + if utils.match_response(http_res, "403", "application/json"): + http_res_text = utils.stream_to_text(http_res) + response_data = unmarshal_json_response( + errors.StreamResearchTaskForbiddenErrorData, http_res, http_res_text + ) + raise errors.StreamResearchTaskForbiddenError( + response_data, http_res, http_res_text + ) + if utils.match_response(http_res, "404", "application/json"): + http_res_text = utils.stream_to_text(http_res) + response_data = unmarshal_json_response( + errors.StreamResearchTaskNotFoundErrorData, http_res, http_res_text + ) + raise errors.StreamResearchTaskNotFoundError( + response_data, http_res, http_res_text + ) + if utils.match_response(http_res, "500", "application/json"): + http_res_text = utils.stream_to_text(http_res) + response_data = unmarshal_json_response( + errors.StreamResearchTaskInternalServerErrorData, + http_res, + http_res_text, + ) + raise errors.StreamResearchTaskInternalServerError( + response_data, http_res, http_res_text + ) + if utils.match_response(http_res, "4XX", "*"): + http_res_text = utils.stream_to_text(http_res) + raise errors.YouDefaultError("API error occurred", http_res, http_res_text) + if utils.match_response(http_res, "5XX", "*"): + http_res_text = utils.stream_to_text(http_res) + raise errors.YouDefaultError("API error occurred", http_res, http_res_text) + + http_res_text = utils.stream_to_text(http_res) + raise errors.YouDefaultError( + "Unexpected response received", http_res, http_res_text + ) + + async def stream_research_task_async( + self, + *, + task_id: str, + from_id: Optional[int] = 0, + retries: OptionalNullable[utils.RetryConfig] = UNSET, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, + http_headers: Optional[Mapping[str, str]] = None, + ) -> eventstreaming.EventStreamAsync[models.ResearchTaskStreamEvent]: + r"""Stream updates for a background research task + + Stream real-time updates for a background research task via Server-Sent Events (SSE). Supports reconnection via the from_id query parameter to replay missed events. The connection closes automatically when the task reaches a terminal state. + + :param task_id: The UUID of the research task. + :param from_id: Resume from a sequence number for reconnection. + :param retries: Override the default retry configuration for this method + :param server_url: Override the default server URL for this method + :param timeout_ms: Override the default request timeout configuration for this method in milliseconds + :param http_headers: Additional headers to set or replace on requests. + """ + base_url = None + url_variables = None + if timeout_ms is None: + timeout_ms = self.sdk_configuration.timeout_ms + + if server_url is not None: + base_url = server_url + else: + base_url = self._get_url(base_url, url_variables) + + request = models.StreamResearchTaskRequest( + task_id=task_id, + from_id=from_id, + ) + + req = self._build_request_async( + method="GET", + path="/v1/research/{task_id}/stream", + base_url=base_url, + url_variables=url_variables, + request=request, + request_body_required=False, + request_has_path_params=True, + request_has_query_params=True, + user_agent_header="user-agent", + accept_header_value="text/event-stream", + http_headers=http_headers, + security=self.sdk_configuration.security, + allow_empty_value=None, + timeout_ms=timeout_ms, + ) + + if retries == UNSET: + if self.sdk_configuration.retry_config is not UNSET: + retries = self.sdk_configuration.retry_config + + retry_config = None + if isinstance(retries, utils.RetryConfig): + retry_config = (retries, ["429", "500", "502", "503", "504"]) + + http_res = await self.do_request_async( + hook_ctx=HookContext( + config=self.sdk_configuration, + base_url=base_url or "", + operation_id="streamResearchTask", + oauth2_scopes=None, + security_source=get_security_from_env( + self.sdk_configuration.security, models.Security + ), + tags=None, + extensions=None, + ), + request=req, + is_error_status_code=lambda c: utils.match_status_codes(["4XX", "5XX"], c), + stream=True, + retry_config=retry_config, + ) + + response_data: Any = None + if utils.match_response(http_res, "200", "text/event-stream"): + return eventstreaming.EventStreamAsync( + http_res, + lambda raw: unmarshal_json_response( + models.ResearchTaskStreamEvent, http_res, raw + ), + client_ref=self, + ) + if utils.match_response(http_res, "401", "application/json"): + http_res_text = await utils.stream_to_text_async(http_res) + response_data = unmarshal_json_response( + errors.StreamResearchTaskUnauthorizedErrorData, http_res, http_res_text + ) + raise errors.StreamResearchTaskUnauthorizedError( + response_data, http_res, http_res_text + ) + if utils.match_response(http_res, "403", "application/json"): + http_res_text = await utils.stream_to_text_async(http_res) + response_data = unmarshal_json_response( + errors.StreamResearchTaskForbiddenErrorData, http_res, http_res_text + ) + raise errors.StreamResearchTaskForbiddenError( + response_data, http_res, http_res_text + ) + if utils.match_response(http_res, "404", "application/json"): + http_res_text = await utils.stream_to_text_async(http_res) + response_data = unmarshal_json_response( + errors.StreamResearchTaskNotFoundErrorData, http_res, http_res_text + ) + raise errors.StreamResearchTaskNotFoundError( + response_data, http_res, http_res_text + ) + if utils.match_response(http_res, "500", "application/json"): + http_res_text = await utils.stream_to_text_async(http_res) + response_data = unmarshal_json_response( + errors.StreamResearchTaskInternalServerErrorData, + http_res, + http_res_text, + ) + raise errors.StreamResearchTaskInternalServerError( + response_data, http_res, http_res_text + ) + if utils.match_response(http_res, "4XX", "*"): + http_res_text = await utils.stream_to_text_async(http_res) + raise errors.YouDefaultError("API error occurred", http_res, http_res_text) + if utils.match_response(http_res, "5XX", "*"): + http_res_text = await utils.stream_to_text_async(http_res) + raise errors.YouDefaultError("API error occurred", http_res, http_res_text) + + http_res_text = await utils.stream_to_text_async(http_res) + raise errors.YouDefaultError( + "Unexpected response received", http_res, http_res_text + ) + def finance_research( self, *, diff --git a/tests/README.md b/tests/README.md index 5a8c7e3..ca38eb3 100644 --- a/tests/README.md +++ b/tests/README.md @@ -55,6 +55,10 @@ pytest tests/ -v - `test_search.py` - Tests for the Search API (`/v1/search`) - `test_contents.py` - Tests for the Contents API (`/v1/contents`) - `test_runs.py` - Tests for the Agents/Runs API (`/v1/agents/runs`) +- `test_research.py` - Tests for the Research API (`/v1/research`) including background mode, output_schema, and source_control +- `test_research_helpers.py` - Tests for the hand-maintained `research_helpers` module (background submission, polling, streaming, research_and_wait) +- `test_security_env.py` - Tests for environment variable precedence (`YDC_API_KEY` / `YOU_API_KEY_AUTH`) +- `test_user_agent_hook.py` - Tests for the `YDCUserAgentOverrideHook` custom user-agent pass-through - `test_performance.py` - Performance/instrumentation tests measuring SDK overhead - `test_live.py` - Live API tests that run against the real You.com API (requires API key) @@ -82,6 +86,21 @@ Tests are organized into logical classes using pytest: - Tool configurations and verbosity - Error handling (unauthorized, forbidden, empty input) +**Research API**: +- Basic research functionality (standard, deep, exhaustive effort) +- Background mode (task submission, get_research_task, status polling) +- Output schema (structured JSON output, content_type object) +- Source control (include/exclude/boost domains, freshness, country) +- Error handling (unauthorized, forbidden, unprocessable entity, 422 combos) +- Stream research task (SSE success path + 404/401/403 error paths) + +**Research Helpers**: +- research_background / research_background_async (TaskResponse return) +- poll_research_task / poll_research_task_async (terminal status) +- research_and_wait / research_and_wait_async (submit + wait) +- stream_research / stream_research_async (tolerant SSE) +- RawStreamEvent decoder (_decode_raw_event) + ### Running Live Tests The `test_live.py` file contains tests that run against the real You.com API. These are skipped by default unless an API key is provided: @@ -115,7 +134,9 @@ The tests use a mock server located in `tests/mockserver/`. This server contains The mock server supports: - Success responses for all endpoints -- Error responses (401 Unauthorized, 403 Forbidden) +- Error responses (401 Unauthorized, 403 Forbidden, 404 Not Found) +- Background research task endpoints (GET /v1/research/{task_id}, GET /v1/research/{task_id}/stream) +- SSE streaming for research task updates - Multiple test scenarios per endpoint See [mockserver/README.md](mockserver/README.md) for more details. diff --git a/tests/mockserver/internal/handler/generated_handlers.go b/tests/mockserver/internal/handler/generated_handlers.go index 6db3e34..e2bdae4 100644 --- a/tests/mockserver/internal/handler/generated_handlers.go +++ b/tests/mockserver/internal/handler/generated_handlers.go @@ -16,6 +16,8 @@ func GeneratedHandlers(ctx context.Context, dir *logging.HTTPFileDirectory, rt * NewGeneratedHandler(ctx, http.MethodPost, "/v1/agents/runs", pathPostV1AgentsRuns(dir, rt)), NewGeneratedHandler(ctx, http.MethodPost, "/v1/contents", pathPostV1Contents(dir, rt)), NewGeneratedHandler(ctx, http.MethodPost, "/v1/research", pathPostV1Research(dir, rt)), + NewGeneratedHandler(ctx, http.MethodGet, "/v1/research/{task_id}", pathGetV1Research(dir, rt)), + NewGeneratedHandler(ctx, http.MethodGet, "/v1/research/{task_id}/stream", pathGetV1ResearchStream(dir, rt)), NewGeneratedHandler(ctx, http.MethodPost, "/v1/finance_research", pathPostV1FinanceResearch(dir, rt)), } } diff --git a/tests/mockserver/internal/handler/pathgetv1research.go b/tests/mockserver/internal/handler/pathgetv1research.go new file mode 100644 index 0000000..781dcdf --- /dev/null +++ b/tests/mockserver/internal/handler/pathgetv1research.go @@ -0,0 +1,111 @@ +package handler + +import ( + "encoding/json" + "fmt" + "log" + "mockserver/internal/handler/assert" + "mockserver/internal/logging" + "mockserver/internal/tracking" + "net/http" +) + +func pathGetV1Research(dir *logging.HTTPFileDirectory, rt *tracking.RequestTracker) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + test := req.Header.Get("x-speakeasy-test-name") + instanceID := req.Header.Get("x-speakeasy-test-instance-id") + + count := rt.GetRequestCount(test, instanceID) + + switch fmt.Sprintf("%s[%d]", test, count) { + case "get_/v1/research/{task_id}[0]": + dir.HandlerFunc("get_/v1/research/{task_id}", testGetV1ResearchTaskSuccess)(w, req) + case "get_/v1/research/{task_id}-not-found[0]": + testGetV1ResearchTaskNotFound(w, req) + case "get_/v1/research/{task_id}-unauthorized[0]": + testGetV1ResearchTaskUnauthorized(w, req) + case "get_/v1/research/{task_id}-forbidden[0]": + testGetV1ResearchTaskForbidden(w, req) + case "get_/v1/research/{task_id}-internal-error[0]": + testGetV1ResearchTaskInternalError(w, req) + default: + dir.HandlerFunc("get_/v1/research/{task_id}", testGetV1ResearchTaskSuccess)(w, req) + } + } +} + +// testGetV1ResearchTaskSuccess returns a TaskDetail with status "completed" +// and a populated result block, mirroring the structure used by the real API +// after a background research task finishes. +func testGetV1ResearchTaskSuccess(w http.ResponseWriter, req *http.Request) { + if err := assert.SecurityHeader(req, "X-API-Key", false); err != nil { + log.Printf("assertion error: %s\n", err) + http.Error(w, err.Error(), http.StatusUnauthorized) + return + } + if err := assert.HeaderExists(req, "User-Agent"); err != nil { + log.Printf("assertion error: %s\n", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + respBody := map[string]interface{}{ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "completed", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:02:30Z", + "completed_at": "2026-07-09T00:02:30Z", + "error": nil, + "input": map[string]interface{}{ + "input": "Compare NVIDIA, AMD, and Intel revenue over 5 years", + "research_effort": "deep", + }, + "result": map[string]interface{}{ + "output": map[string]interface{}{ + "content": "# Mock Research Result\n\nMock result for completed background task.", + "content_type": "text", + "sources": []map[string]interface{}{ + { + "url": "https://example.com/research/1", + "title": "Mock Research Source 1", + "snippets": []string{"This is a relevant snippet from source 1."}, + }, + }, + }, + }, + } + + respBodyBytes, err := json.Marshal(respBody) + if err != nil { + http.Error(w, "Unable to encode response body as JSON: "+err.Error(), http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(respBodyBytes) +} + +func testGetV1ResearchTaskNotFound(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"detail":"Task not found"}`)) +} + +func testGetV1ResearchTaskUnauthorized(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"message":"Invalid or expired API key"}`)) +} + +func testGetV1ResearchTaskForbidden(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"detail":"Forbidden"}`)) +} + +func testGetV1ResearchTaskInternalError(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"message":"Internal server error"}`)) +} diff --git a/tests/mockserver/internal/handler/pathgetv1researchstream.go b/tests/mockserver/internal/handler/pathgetv1researchstream.go new file mode 100644 index 0000000..3527a83 --- /dev/null +++ b/tests/mockserver/internal/handler/pathgetv1researchstream.go @@ -0,0 +1,128 @@ +package handler + +import ( + "encoding/json" + "fmt" + "log" + "mockserver/internal/handler/assert" + "mockserver/internal/logging" + "mockserver/internal/tracking" + "net/http" +) + +func pathGetV1ResearchStream(dir *logging.HTTPFileDirectory, rt *tracking.RequestTracker) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + test := req.Header.Get("x-speakeasy-test-name") + instanceID := req.Header.Get("x-speakeasy-test-instance-id") + + count := rt.GetRequestCount(test, instanceID) + + switch fmt.Sprintf("%s[%d]", test, count) { + case "get_/v1/research/{task_id}/stream[0]": + dir.HandlerFunc("get_/v1/research/{task_id}/stream", testGetV1ResearchStreamSuccess)(w, req) + case "get_/v1/research/{task_id}/stream-not-found[0]": + testGetV1ResearchStreamNotFound(w, req) + case "get_/v1/research/{task_id}/stream-unauthorized[0]": + testGetV1ResearchStreamUnauthorized(w, req) + case "get_/v1/research/{task_id}/stream-forbidden[0]": + testGetV1ResearchStreamForbidden(w, req) + case "get_/v1/research/{task_id}/stream-internal-error[0]": + testGetV1ResearchStreamInternalError(w, req) + default: + dir.HandlerFunc("get_/v1/research/{task_id}/stream", testGetV1ResearchStreamSuccess)(w, req) + } + } +} + +// testGetV1ResearchStreamSuccess emits the SSE sequence documented in the +// Research API stream spec: an opening `connected` event followed by a +// terminal `response.done` event, then a normal close. +// +// Terminal event types: {"response.done", "complete", "completed", +// "error", "failed", "cancelled"}. +func testGetV1ResearchStreamSuccess(w http.ResponseWriter, req *http.Request) { + if err := assert.SecurityHeader(req, "X-API-Key", false); err != nil { + log.Printf("assertion error: %s\n", err) + http.Error(w, err.Error(), http.StatusUnauthorized) + return + } + if err := assert.HeaderExists(req, "User-Agent"); err != nil { + log.Printf("assertion error: %s\n", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + // SSE requires these headers; flushing after each event lets the client + // receive events incrementally. + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + w.WriteHeader(http.StatusOK) + + flusher, ok := w.(http.Flusher) + if !ok { + log.Printf("response writer does not support flushing; stream may buffer") + } + + writeSSE := func(id, event string, data map[string]interface{}) bool { + payload, err := json.Marshal(data) + if err != nil { + log.Printf("error marshalling SSE data: %s", err) + return false + } + if _, err := fmt.Fprintf(w, "id: %s\nevent: %s\ndata: %s\n\n", id, event, payload); err != nil { + return false + } + if ok { + flusher.Flush() + } + return true + } + + taskID := "00000000-0000-0000-0000-000000000001" + + // 1) Opening event — sent unconditionally when the stream is opened. + if !writeSSE("0", "connected", map[string]interface{}{ + "type": "connected", + "task_id": taskID, + "status": "running", + }) { + return + } + + // 2) Terminal event — closes the stream. `response.done` is one of the + // four TERMINAL_SSE_EVENTS values that the SDK treats as stream-end. + if !writeSSE("1", "response.done", map[string]interface{}{ + "type": "response.done", + "task_id": taskID, + "status": "completed", + "sequence": 1, + }) { + return + } +} + +func testGetV1ResearchStreamNotFound(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"detail":"Task not found"}`)) +} + +func testGetV1ResearchStreamUnauthorized(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"message":"Invalid or expired API key"}`)) +} + +func testGetV1ResearchStreamForbidden(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"detail":"Forbidden"}`)) +} + +func testGetV1ResearchStreamInternalError(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"message":"Internal server error"}`)) +} diff --git a/tests/mockserver/internal/handler/pathpostv1research.go b/tests/mockserver/internal/handler/pathpostv1research.go index fff4d02..3b8d857 100644 --- a/tests/mockserver/internal/handler/pathpostv1research.go +++ b/tests/mockserver/internal/handler/pathpostv1research.go @@ -28,6 +28,8 @@ func pathPostV1Research(dir *logging.HTTPFileDirectory, rt *tracking.RequestTrac switch fmt.Sprintf("%s[%d]", test, count) { case "post_/v1/research[0]": dir.HandlerFunc("post_/v1/research", testPostV1ResearchSuccess)(w, req) + case "post_/v1/research-background[0]": + dir.HandlerFunc("post_/v1/research-background", testPostV1ResearchBackground)(w, req) case "post_/v1/research-unauthorized[0]": testPostV1ResearchUnauthorized(w, req) case "post_/v1/research-forbidden[0]": @@ -79,6 +81,13 @@ func testPostV1ResearchSuccess(w http.ResponseWriter, req *http.Request) { effort = "standard" } + // When `background: true` is set, return a TaskResponse shape so the SDK + // can deserialize the async task handle instead of an inline ResearchResponse. + if background, _ := requestBody["background"].(bool); background { + respondTaskResponse(w, "research", "queued", "/v1/research/00000000-0000-0000-0000-000000000001/stream") + return + } + respBody := buildResearchResponse(input, effort, requestBody) respBodyBytes, err := json.Marshal(respBody) if err != nil { @@ -144,6 +153,42 @@ func testPostV1ResearchUnauthorized(w http.ResponseWriter, req *http.Request) { _, _ = w.Write([]byte(`{"message":"Invalid or expired API key"}`)) } +// respondTaskResponse writes a TaskResponse payload used by background-mode +// research. +func respondTaskResponse(w http.ResponseWriter, typeValue, statusValue, streamPathSuffix string) { + respBody := map[string]interface{}{ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": typeValue, + "status": statusValue, + "stream_url": streamPathSuffix, + "created_at": "2026-07-09T00:00:00Z", + } + respBodyBytes, err := json.Marshal(respBody) + if err != nil { + http.Error(w, "Unable to encode response body as JSON: "+err.Error(), http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(respBodyBytes) +} + +// testPostV1ResearchBackground handles the explicit `post_/v1/research-background` +// test name and returns a TaskResponse shape for background-mode submission. +func testPostV1ResearchBackground(w http.ResponseWriter, req *http.Request) { + if err := assert.SecurityHeader(req, "X-API-Key", false); err != nil { + log.Printf("assertion error: %s\n", err) + http.Error(w, err.Error(), http.StatusUnauthorized) + return + } + if err := assert.ContentType(req, "application/json", true); err != nil { + log.Printf("assertion error: %s\n", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + respondTaskResponse(w, "research", "queued", "/v1/research/00000000-0000-0000-0000-000000000001/stream") +} + func testPostV1ResearchForbidden(w http.ResponseWriter, req *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusForbidden) diff --git a/tests/test_live.py b/tests/test_live.py index e2d23c7..e621470 100644 --- a/tests/test_live.py +++ b/tests/test_live.py @@ -34,8 +34,16 @@ AgentRunsBatchResponse, ResearchEffort, ResearchResponse, + TaskResponse, + TaskDetail, FinanceResearchEffort, ) +from youdotcom.research_helpers import ( + research_background, + poll_research_task, + research_and_wait, + stream_research, +) # Skip all tests in this file if no API key is provided. @@ -447,6 +455,181 @@ def test_search_post_boost_domains_list(self, you_client): assert res.results is not None +# --------------------------------------------------------------------------- +# Background-mode research (new in 2.5.0) +# --------------------------------------------------------------------------- +# These tests exercise the live API's async task path: +# POST /v1/research?background=true -> TaskResponse +# GET /v1/research/{task_id} -> TaskDetail +# GET /v1/research/{task_id}/stream -> SSE stream +# +# All tests use LITE effort so the task finishes in 10-30s on prod. +# The suite is sequential within each test: submit, then poll or stream +# until a terminal state arrives. If background mode is not enabled on +# the server, the research() call falls back to a sync ResearchResponse +# and the test is skipped with a clear message. +# --------------------------------------------------------------------------- + +_BG_TIMEOUT_S = 120.0 # generous wall-clock for LITE background tasks + + +class TestLiveResearchBackground: + """Live tests for background-mode research (POST /v1/research?background=true).""" + + def test_background_returns_task_response(self, you_client): + """research(background=True) should return a TaskResponse, not ResearchResponse.""" + with you_client as you: + res = you.research( + input="What is the capital of France?", + research_effort=ResearchEffort.LITE, + background=True, + ) + + if isinstance(res, ResearchResponse): + pytest.skip("Background mode not enabled on server (got sync ResearchResponse)") + + assert isinstance(res, TaskResponse) + assert res.task_id is not None + assert len(res.task_id) > 0 + assert res.type == "research" + assert res.status is not None + assert res.stream_url is not None + assert res.stream_url.startswith("/v1/research/") + assert res.stream_url.endswith("/stream") + assert res.created_at is not None + + def test_get_research_task_returns_task_detail(self, you_client): + """get_research_task() should return a TaskDetail for a background task.""" + with you_client as you: + task = you.research( + input="What is the capital of France?", + research_effort=ResearchEffort.LITE, + background=True, + ) + + if isinstance(task, ResearchResponse): + pytest.skip("Background mode not enabled on server") + + assert isinstance(task, TaskResponse) + + detail = you.get_research_task(task_id=task.task_id) + assert isinstance(detail, TaskDetail) + assert detail.id == task.task_id + assert detail.task_type == "research" + assert detail.status is not None + assert detail.created_at is not None + # input should be preserved (TaskDetailInput uses extra="allow") + assert detail.input is not None + input_dump = detail.input.model_dump() + assert input_dump.get("input") == "What is the capital of France?" + + def test_poll_until_completed(self, you_client): + """Poll get_research_task() until status == completed, verify result.""" + with you_client as you: + task = you.research( + input="What is the capital of France?", + research_effort=ResearchEffort.LITE, + background=True, + ) + + if isinstance(task, ResearchResponse): + pytest.skip("Background mode not enabled on server") + + assert isinstance(task, TaskResponse) + + detail = poll_research_task( + you, + task.task_id, + interval_s=3.0, + timeout_s=_BG_TIMEOUT_S, + ) + + assert detail.status.value == "completed" + assert detail.completed_at is not None + # Result should contain the ResearchResponse payload + assert detail.result is not None + result_dump = detail.result.model_dump() + assert "output" in result_dump + output = result_dump["output"] + assert "content" in output + assert len(str(output["content"])) > 0 + + def test_research_and_wait(self, you_client): + """research_and_wait submits + streams + returns completed TaskDetail.""" + with you_client as you: + try: + detail = research_and_wait( + you, + timeout_s=_BG_TIMEOUT_S, + timeout_ms=180_000, + input="What is the capital of France?", + research_effort=ResearchEffort.LITE, + ) + except TypeError: + pytest.skip("Background mode not enabled on server") + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + assert detail.result is not None + result_dump = detail.result.model_dump() + assert "output" in result_dump + + +class TestLiveResearchBackgroundHelpers: + """Live tests for the research_helpers convenience functions.""" + + def test_research_background_helper(self, you_client): + """research_background() helper asserts and returns TaskResponse.""" + with you_client as you: + try: + task = research_background( + you, + input="What is the capital of France?", + research_effort=ResearchEffort.LITE, + ) + except TypeError: + pytest.skip("Background mode not enabled on server") + + assert isinstance(task, TaskResponse) + assert task.task_id is not None + assert task.stream_url is not None + + def test_stream_research(self, you_client): + """stream_research() yields SSE events from a live task. + + The server's SSE stream sends a 'connected' event immediately, then + pings. Terminal events may not arrive for tasks that complete + mid-stream, so we iterate with a generous timeout and verify + at least the 'connected' event was received. + """ + with you_client as you: + try: + task = research_background( + you, + input="What is the capital of France?", + research_effort=ResearchEffort.LITE, + ) + except TypeError: + pytest.skip("Background mode not enabled on server") + + assert isinstance(task, TaskResponse) + + terminal = {"response.done", "complete", "completed", "error", "failed", "cancelled"} + + events = [] + for evt in stream_research(you, task_id=task.task_id, timeout_ms=180_000): + events.append(evt) + if evt.event in terminal: + break + + assert len(events) > 0, "No SSE events received" + assert events[0].event == "connected" + # The connected event should carry task_id and status + assert events[0].data is not None + assert events[0].data.get("task_id") == task.task_id + assert events[0].data.get("status") is not None + + if __name__ == "__main__": # Run with: python -m pytest tests/test_live.py -v pytest.main([__file__, "-v"]) diff --git a/tests/test_research.py b/tests/test_research.py index de579c4..1e4768f 100644 --- a/tests/test_research.py +++ b/tests/test_research.py @@ -5,7 +5,7 @@ import httpx from tests.test_client import create_test_http_client -from youdotcom import You +from youdotcom import You, errors from youdotcom.errors import ( FinanceResearchUnauthorizedError, FinanceResearchUnprocessableEntityError, @@ -20,7 +20,9 @@ FinanceResearchEffort, ResearchEffort, ResearchResponse, + TaskResponse, ) +from youdotcom.utils import eventstreaming @pytest.fixture @@ -208,6 +210,151 @@ def test_finance_research_unauthorized(self, server_url): ) +# --------------------------------------------------------------------------- +# 2.5.0 background mode: queue task, poll status, stream events. +# --------------------------------------------------------------------------- + +class TestResearchBackground: + """Direct coverage for the auto-generated background/stream SDK methods.""" + + def test_research_background_returns_task_response(self, server_url, api_key): + client = create_test_http_client("post_/v1/research-background") + + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + res = you.research( + input="What is the capital of France?", + research_effort=ResearchEffort.STANDARD, + background=True, + server_url=server_url, + ) + + assert isinstance(res, TaskResponse) + assert res.task_id == "00000000-0000-0000-0000-000000000001" + assert res.status.value == "queued" + assert res.stream_url is not None + assert res.created_at is not None + + def test_get_research_task_returns_task_detail(self, server_url, api_key): + client = create_test_http_client("get_/v1/research/{task_id}") + + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + detail = you.get_research_task( + task_id="00000000-0000-0000-0000-000000000001", + server_url=server_url, + ) + + assert detail.id == "00000000-0000-0000-0000-000000000001" + assert detail.task_type == "research" + assert detail.status.value == "completed" + assert detail.created_at is not None + assert detail.completed_at is not None + # result is populated server-side; the Result model uses + # extra="allow" so model_dump() recovers the full payload. + assert detail.result is not None + assert detail.result.model_dump().get("output") is not None + # input is populated server-side; the TaskDetailInput model uses + # extra="allow" so model_dump() recovers the original request fields. + assert detail.input is not None + assert detail.input.model_dump().get("input") == "Compare NVIDIA, AMD, and Intel revenue over 5 years" + assert detail.input.model_dump().get("research_effort") == "deep" + + def test_research_background_false_returns_research_response(self, server_url, api_key): + """When background=False (explicit), return type is ResearchResponse, not TaskResponse.""" + client = create_test_http_client("post_/v1/research") + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + res = you.research( + input="What is the capital of France?", + research_effort=ResearchEffort.STANDARD, + background=False, + server_url=server_url, + ) + assert isinstance(res, ResearchResponse) + assert not isinstance(res, TaskResponse) + + @pytest.mark.asyncio + async def test_get_research_task_async_returns_task_detail(self, server_url, api_key): + async_client = httpx.AsyncClient( + headers={ + "x-speakeasy-test-name": "get_/v1/research/{task_id}", + "x-speakeasy-test-instance-id": str(uuid.uuid4()), + }, + follow_redirects=True, + ) + async with You(server_url=server_url, async_client=async_client, api_key_auth=api_key) as you: + detail = await you.get_research_task_async( + task_id="00000000-0000-0000-0000-000000000001", + server_url=server_url, + ) + assert detail.id == "00000000-0000-0000-0000-000000000001" + assert detail.status.value == "completed" + + +class TestResearchBackgroundErrors: + """Error path coverage for get_research_task().""" + + def test_get_research_task_not_found(self, server_url, api_key): + client = create_test_http_client("get_/v1/research/{task_id}-not-found") + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + with pytest.raises(errors.GetResearchTaskNotFoundError): + you.get_research_task(task_id="nonexistent", server_url=server_url) + + def test_get_research_task_unauthorized(self, server_url): + client = create_test_http_client("get_/v1/research/{task_id}-unauthorized") + with You(server_url=server_url, client=client, api_key_auth="bad-key") as you: + with pytest.raises(errors.GetResearchTaskUnauthorizedError): + you.get_research_task(task_id="test", server_url=server_url) + + def test_get_research_task_forbidden(self, server_url, api_key): + client = create_test_http_client("get_/v1/research/{task_id}-forbidden") + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + with pytest.raises(errors.GetResearchTaskForbiddenError): + you.get_research_task(task_id="test", server_url=server_url) + + def test_get_research_task_internal_error(self, server_url, api_key): + client = create_test_http_client("get_/v1/research/{task_id}-internal-error") + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + with pytest.raises(errors.GetResearchTaskInternalServerError): + you.get_research_task(task_id="test", server_url=server_url) + + +class TestStreamResearchTask: + """Direct SDK-level coverage for stream_research_task().""" + + def test_stream_research_task_returns_event_stream(self, server_url, api_key): + client = create_test_http_client("get_/v1/research/{task_id}/stream") + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + stream = you.stream_research_task( + task_id="00000000-0000-0000-0000-000000000001", + server_url=server_url, + ) + assert isinstance(stream, eventstreaming.EventStream) + events = [] + with stream as s: + for chunk in s: + events.append(chunk.data) + # The mockserver emits connected + response.done events + assert len(events) >= 2 + + def test_stream_research_task_not_found(self, server_url, api_key): + client = create_test_http_client("get_/v1/research/{task_id}/stream-not-found") + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + with pytest.raises(errors.StreamResearchTaskNotFoundError): + you.stream_research_task(task_id="nonexistent", server_url=server_url) + + def test_stream_research_task_unauthorized(self, server_url): + client = create_test_http_client("get_/v1/research/{task_id}/stream-unauthorized") + with You(server_url=server_url, client=client, api_key_auth="bad-key") as you: + with pytest.raises(errors.StreamResearchTaskUnauthorizedError): + you.stream_research_task(task_id="test", server_url=server_url) + + def test_stream_research_task_forbidden(self, server_url, api_key): + # Requires the forbidden case added in Phase 3.10.1 + client = create_test_http_client("get_/v1/research/{task_id}/stream-forbidden") + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + with pytest.raises(errors.StreamResearchTaskForbiddenError): + you.stream_research_task(task_id="test", server_url=server_url) + + # --------------------------------------------------------------------------- # 2.4.0 beta params: output_schema (structured output_content) and # source_control (domain constraints / freshness / country). diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py new file mode 100644 index 0000000..e9a96c1 --- /dev/null +++ b/tests/test_research_helpers.py @@ -0,0 +1,1281 @@ +"""Tests for research background-mode helpers in youdotcom.research_helpers.""" + +import os +import uuid + +import httpx +import pytest + +from tests.test_client import create_test_http_client +from youdotcom import You +from youdotcom.models import ( + ResearchEffort, + TaskDetail, + TaskResponse, +) +from youdotcom.research_helpers import ( + RawStreamEvent, + research_and_wait, + research_and_wait_async, + research_background, + research_background_async, + poll_research_task, + poll_research_task_async, + stream_research, + stream_research_async, + _decode_raw_event, +) + + +@pytest.fixture +def server_url(): + return os.getenv("TEST_SERVER_URL", "http://localhost:18080") + + +@pytest.fixture +def api_key(): + return "test-api-key" + + +# --------------------------------------------------------------------------- +# research_background[Async]: assert TaskResponse return type without +# forcing callers to narrow Union[ResearchResponse, TaskResponse]. +# --------------------------------------------------------------------------- + +class TestResearchBackground: + def test_research_background_returns_task_response(self, server_url, api_key): + client = create_test_http_client("post_/v1/research-background") + + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + res = research_background( + you, + input="Compare NVIDIA, AMD, and Intel revenue over 5 years", + research_effort=ResearchEffort.DEEP, + server_url=server_url, + ) + + assert isinstance(res, TaskResponse) + assert res.task_id == "00000000-0000-0000-0000-000000000001" + assert res.type == "research" + assert res.status.value == "queued" + + @pytest.mark.asyncio + async def test_research_background_async_returns_task_response(self, server_url, api_key): + async_client = httpx.AsyncClient( + headers={ + "x-speakeasy-test-name": "post_/v1/research-background", + "x-speakeasy-test-instance-id": str(uuid.uuid4()), + }, + follow_redirects=True, + ) + + async with You( + server_url=server_url, async_client=async_client, api_key_auth=api_key + ) as you: + res = await research_background_async( + you, + input="Compare NVIDIA, AMD, and Intel revenue over 5 years", + research_effort=ResearchEffort.DEEP, + server_url=server_url, + ) + + assert isinstance(res, TaskResponse) + assert res.task_id == "00000000-0000-0000-0000-000000000001" + + +class TestResearchBackgroundTypeError: + """Tests that research_background raises TypeError when the server + returns a ResearchResponse instead of TaskResponse (e.g. when the + server-side background-mode flag is disabled).""" + + def test_research_background_raises_type_error_on_sync_response(self): + """When the server ignores background=true and returns a + ResearchResponse, research_background must raise TypeError.""" + import json + + def handler(request): + # Server ignores background=true, returns a sync ResearchResponse + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "output": { + "content_type": "text", + "content": "The capital of France is Paris.", + "sources": [], + }, + }), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(TypeError, match="TaskResponse"): + research_background( + you, + input="What is the capital of France?", + research_effort=ResearchEffort.STANDARD, + ) + + @pytest.mark.asyncio + async def test_research_background_async_raises_type_error_on_sync_response(self): + """Async mirror: TypeError when server returns ResearchResponse.""" + import json + + def handler(request): + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "output": { + "content_type": "text", + "content": "The capital of France is Paris.", + "sources": [], + }, + }), + ) + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(TypeError, match="TaskResponse"): + await research_background_async( + you, + input="What is the capital of France?", + research_effort=ResearchEffort.STANDARD, + ) + + +# --------------------------------------------------------------------------- +# poll_research_task[Async]: poll GET /v1/research/{task_id} until terminal. +# --------------------------------------------------------------------------- + +class TestPollResearchTask: + def test_poll_returns_completed_task_detail(self, server_url, api_key): + client = create_test_http_client("get_/v1/research/{task_id}") + + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + detail = poll_research_task( + you, + "00000000-0000-0000-0000-000000000001", + interval_s=0.01, + timeout_s=2.0, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + assert detail.result is not None + + @pytest.mark.asyncio + async def test_poll_async_returns_completed_task_detail(self, server_url, api_key): + async_client = httpx.AsyncClient( + headers={ + "x-speakeasy-test-name": "get_/v1/research/{task_id}", + "x-speakeasy-test-instance-id": str(uuid.uuid4()), + }, + follow_redirects=True, + ) + + async with You( + server_url=server_url, async_client=async_client, api_key_auth=api_key + ) as you: + detail = await poll_research_task_async( + you, + "00000000-0000-0000-0000-000000000001", + interval_s=0.01, + timeout_s=2.0, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + +# --------------------------------------------------------------------------- +# research_and_wait: submit background + poll + return TaskDetail. +# The Result model uses extra="allow" so detail.result.model_dump() +# recovers the full payload; the helper returns TaskDetail. +# --------------------------------------------------------------------------- + +class TestResearchAndWait: + def test_research_and_wait_returns_completed_detail(self): + """research_and_wait submits, streams until terminal event, + then fetches the final TaskDetail.""" + import json + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + chunks = [ + b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=chunks, + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + # GET /v1/research/{task_id} — final fetch after terminal event + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "completed", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:02:30Z", + "completed_at": "2026-07-09T00:02:30Z", + "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, + }), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + detail = research_and_wait( + you, + timeout_s=5.0, + input="Compare NVIDIA, AMD, and Intel revenue over 5 years", + research_effort=ResearchEffort.DEEP, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + def test_research_and_wait_error_event_raises_runtime_error(self): + """research_and_wait raises RuntimeError when the stream emits an + error terminal event.""" + import json + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + chunks = [ + b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', + b'id: 1\nevent: error\ndata: {"type":"error","task_id":"abc","message":"internal error"}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=chunks, + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + return httpx.Response(200, content="{}") + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(RuntimeError, match="non-completed state"): + research_and_wait( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + def test_research_and_wait_timeout_raises_timeout_error(self): + """research_and_wait raises TimeoutError when the stream never sends + a terminal event within timeout_s.""" + import json + import time as _time + + class _BlockingStream(httpx.SyncByteStream): + """Yields one event then blocks so the consumer thread + is still alive when the join timeout expires.""" + def __iter__(self): + yield b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n' + _time.sleep(100) + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_BlockingStream(), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + # Final GET on timeout — task still running + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "running", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:00:10Z", + }), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(TimeoutError, match="did not complete within"): + research_and_wait( + you, + timeout_s=0.5, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + def test_research_and_wait_timeout_falls_back_to_get(self): + """When the stream times out but the task has completed, the final + GET fallback returns the completed detail.""" + import json + import time as _time + + class _BlockingStream(httpx.SyncByteStream): + def __iter__(self): + yield b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n' + _time.sleep(100) + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_BlockingStream(), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + # Final GET on timeout — task completed despite no terminal event + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "completed", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:02:30Z", + "completed_at": "2026-07-09T00:02:30Z", + "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, + }), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + detail = research_and_wait( + you, + timeout_s=0.5, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + def test_research_and_wait_stream_close_falls_back_to_get(self): + """When the stream closes without a terminal event, research_and_wait + does a final GET and returns the detail if completed.""" + import json + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + # Only a connected event, no terminal event + chunks = [ + b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=chunks, + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + # GET returns completed — the fallback should return this + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "completed", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:02:30Z", + "completed_at": "2026-07-09T00:02:30Z", + "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, + }), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + detail = research_and_wait( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + def test_research_and_wait_stream_close_task_running_raises_timeout(self): + """When the stream closes without a terminal event and the final GET + shows the task is still running, research_and_wait raises TimeoutError.""" + import json + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + chunks = [ + b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=chunks, + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + # GET returns running — task hasn't completed + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "running", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:00:10Z", + }), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(TimeoutError, match="still running"): + research_and_wait( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) +# (a) helper accepts documented event types as before; +# (b) helper accepts unknown event types without raising +# pydantic.ValidationError, surfacing them as RawStreamEvent(event="..."). +# Uses an httpx.MockTransport to inject a fake SSE server response -- this +# path doesn't depend on the Go mockserver so unknown event names can be +# emitted safely without adding new mockserver fixtures. +# --------------------------------------------------------------------------- + +class TestStreamResearchEventsTolerant: + def test_tolerant_stream_yields_all_events_with_unknown_name(self): + # Inject a fake SSE stream with one workflow-internal event type + # (research.searching) that's NOT in the documented enum. The strict + # speakeasy decoder would raise ValidationError on it; our tolerant + # helper must surface it as RawStreamEvent(event="research.searching"). + recorded_ua: dict = {} + + def record_send(request): + recorded_ua["value"] = request.headers.get("User-Agent") + chunks = [ + b"id: 0\nevent: connected\ndata: " + b'{"type":"connected","task_id":"abc","status":"running"}\n\n', + b"id: 1\nevent: research.searching\ndata: " + b'{"query":"markets","phase":"searching"}\n\n', + b"id: 2\nevent: response.done\ndata: " + b'{"type":"response.done","task_id":"abc","status":"completed","sequence":2}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=chunks, + ) + + transport = httpx.MockTransport(record_send) + sdk_client = httpx.Client( + transport=transport, + headers={ + "x-speakeasy-test-name": "get_/v1/research/{task_id}/stream", + "x-speakeasy-test-instance-id": str(uuid.uuid4()), + }, + ) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + events = list( + stream_research( + you, "00000000-0000-0000-0000-000000000001", + ) + ) + + assert [e.event for e in events] == [ + "connected", "research.searching", "response.done", + ] + # Validate that the unknown event came through without raising. + assert isinstance(events[1], RawStreamEvent) + assert events[1].data == {"query": "markets", "phase": "searching"} + # And confirm the SDK still set User-Agent on the underlying request + # (the YDCUserAgentOverrideHook ran before send). + assert recorded_ua["value"] == f"youdotcom-python-sdk/{you.sdk_configuration.sdk_version}" +# --------------------------------------------------------------------------- +# _decode_raw_event sanity: known and unknown shapes. +# --------------------------------------------------------------------------- + +class TestDecodeRawEvent: + def test_known_event(self): + import json + ev = _decode_raw_event(json.dumps({"id": "1", "event": "response.done", "data": {"status": "completed"}})) + assert isinstance(ev, RawStreamEvent) + assert ev.id == "1" + assert ev.event == "response.done" + assert ev.data == {"status": "completed"} + + def test_unknown_event(self): + import json + ev = _decode_raw_event(json.dumps({"id": "2", "event": "some.workflow.step", "data": {"k": "v"}})) + assert ev.event == "some.workflow.step" + assert ev.data == {"k": "v"} + + +# --------------------------------------------------------------------------- +# stream_research typed error mapping: non-200 responses must raise the +# same structured errors as the generated stream_research_task() method. +# --------------------------------------------------------------------------- + +class TestStreamResearchTypedErrors: + def test_404_raises_not_found_error(self): + import json + + def handler(request): + return httpx.Response( + 404, + headers={"content-type": "application/json"}, + content=json.dumps({"detail": "Task not found"}), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + from youdotcom.errors import StreamResearchTaskNotFoundError + with pytest.raises(StreamResearchTaskNotFoundError): + list(stream_research(you, "00000000-0000-0000-0000-000000000001")) + + def test_401_raises_unauthorized_error(self): + import json + + def handler(request): + return httpx.Response( + 401, + headers={"content-type": "application/json"}, + content=json.dumps({"detail": "Invalid API key"}), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + from youdotcom.errors import StreamResearchTaskUnauthorizedError + with pytest.raises(StreamResearchTaskUnauthorizedError): + list(stream_research(you, "00000000-0000-0000-0000-000000000001")) + + @pytest.mark.asyncio + async def test_async_404_raises_not_found_error(self): + import json + + def handler(request): + return httpx.Response( + 404, + headers={"content-type": "application/json"}, + content=json.dumps({"detail": "Task not found"}), + ) + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + from youdotcom.errors import StreamResearchTaskNotFoundError + with pytest.raises(StreamResearchTaskNotFoundError): + async for _ in stream_research_async(you, "00000000-0000-0000-0000-000000000001"): + pass + + +# --------------------------------------------------------------------------- +# Error path tests: poll timeout, poll failed status, and stream mode. +# --------------------------------------------------------------------------- + +class TestPollResearchTaskErrorPaths: + def test_poll_timeout_raises_timeout_error(self): + """poll_research_task must raise TimeoutError when the task never + reaches a terminal state within timeout_s.""" + import json + + def handler(request): + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "running", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:00:01Z", + }), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(TimeoutError, match="did not complete"): + poll_research_task( + you, + "00000000-0000-0000-0000-000000000001", + interval_s=0.01, + timeout_s=0.05, + ) + + def test_poll_failed_status_raises_runtime_error(self): + """poll_research_task must raise RuntimeError when the task ends + in a non-completed terminal state (failed).""" + import json + + def handler(request): + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000002", + "task_type": "research", + "status": "failed", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:00:05Z", + "error": "upstream search timeout", + }), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You( + server_url="http://mock.local", + client=sdk_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(RuntimeError, match="non-completed state: failed"): + poll_research_task( + you, + "00000000-0000-0000-0000-000000000002", + interval_s=0.01, + timeout_s=2.0, + ) + +class TestPollResearchTaskAsyncErrorPaths: + @pytest.mark.asyncio + async def test_poll_async_timeout_raises_timeout_error(self): + """Async mirror of the sync timeout test.""" + import json + + def handler(request): + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000003", + "task_type": "research", + "status": "running", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:00:01Z", + }), + ) + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(TimeoutError, match="did not complete"): + await poll_research_task_async( + you, + "00000000-0000-0000-0000-000000000003", + interval_s=0.01, + timeout_s=0.05, + ) + + @pytest.mark.asyncio + async def test_poll_async_failed_status_raises_runtime_error(self): + """Async mirror of the sync failed-status test.""" + import json + + def handler(request): + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000004", + "task_type": "research", + "status": "failed", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:00:05Z", + "error": "upstream search timeout", + }), + ) + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(RuntimeError, match="non-completed state: failed"): + await poll_research_task_async( + you, + "00000000-0000-0000-0000-000000000004", + interval_s=0.01, + timeout_s=2.0, + ) + + +# --------------------------------------------------------------------------- +# Async streaming: mirror the sync TestStreamResearchEventsTolerant and +# TestResearchAndWaitStreamMode. These also exercise the try/finally cleanup +# path (replacing the broken contextlib.aclosing that called aclose()). + +class _AsyncChunks(httpx.AsyncByteStream): + """Wrap a list of bytes chunks in an AsyncByteStream for MockTransport + + AsyncClient streaming responses (httpx MockTransport returns sync + streams by default, which AsyncClient rejects for stream=True).""" + + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + + async def __aiter__(self): + for chunk in self._chunks: + yield chunk + + +class TestStreamResearchEventsTolerantAsync: + @pytest.mark.asyncio + async def test_async_tolerant_stream_yields_all_events_with_unknown_name(self): + """Async mirror of TestStreamResearchEventsTolerant — injects a + fake SSE stream with an unknown event type and verifies it surfaces + as RawStreamEvent without raising.""" + recorded_ua: dict = {} + + def record_send(request): + recorded_ua["value"] = request.headers.get("User-Agent") + chunks = [ + b"id: 0\nevent: connected\ndata: " + b'{"type":"connected","task_id":"abc","status":"running"}\n\n', + b"id: 1\nevent: research.searching\ndata: " + b'{"query":"markets","phase":"searching"}\n\n', + b"id: 2\nevent: response.done\ndata: " + b'{"type":"response.done","task_id":"abc","status":"completed","sequence":2}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_AsyncChunks(chunks), + ) + + transport = httpx.MockTransport(record_send) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + events = [ + evt + async for evt in stream_research_async( + you, "00000000-0000-0000-0000-000000000001", + ) + ] + + assert [e.event for e in events] == [ + "connected", "research.searching", "response.done", + ] + assert isinstance(events[1], RawStreamEvent) + assert events[1].data == {"query": "markets", "phase": "searching"} + assert recorded_ua["value"] == f"youdotcom-python-sdk/{you.sdk_configuration.sdk_version}" + + +class TestResearchAndWaitAsync: + @pytest.mark.asyncio + async def test_async_research_and_wait_returns_completed_detail(self): + """Async research_and_wait: submit, stream until terminal, fetch detail.""" + import json + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + chunks = [ + b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_AsyncChunks(chunks), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "completed", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:02:30Z", + "completed_at": "2026-07-09T00:02:30Z", + "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, + }), + ) + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + detail = await research_and_wait_async( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + @pytest.mark.asyncio + async def test_async_research_and_wait_error_event_raises_runtime_error(self): + """Async research_and_wait raises RuntimeError on error terminal event.""" + import json + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + chunks = [ + b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', + b'id: 1\nevent: error\ndata: {"type":"error","task_id":"abc","message":"internal error"}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_AsyncChunks(chunks), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + return httpx.Response(200, content="{}") + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(RuntimeError, match="non-completed state"): + await research_and_wait_async( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + @pytest.mark.asyncio + async def test_async_research_and_wait_timeout_raises_timeout_error(self): + """Async research_and_wait raises TimeoutError when the stream never + sends a terminal event and the task is still running.""" + import asyncio + import json + + class _BlockingAsyncStream(httpx.AsyncByteStream): + """Yields one event then blocks so asyncio.wait_for times out.""" + async def __aiter__(self): + yield b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n' + await asyncio.sleep(100) + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_BlockingAsyncStream(), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + # Final GET on timeout — task still running + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "running", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:00:10Z", + }), + ) + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(TimeoutError, match="did not complete within"): + await research_and_wait_async( + you, + timeout_s=0.5, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + @pytest.mark.asyncio + async def test_async_research_and_wait_timeout_falls_back_to_get(self): + """When the async stream times out but the task completed, the final + GET fallback returns the completed detail.""" + import asyncio + import json + + class _BlockingAsyncStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n' + await asyncio.sleep(100) + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_BlockingAsyncStream(), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + # Final GET on timeout — task completed + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "completed", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:02:30Z", + "completed_at": "2026-07-09T00:02:30Z", + "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, + }), + ) + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + detail = await research_and_wait_async( + you, + timeout_s=0.5, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + @pytest.mark.asyncio + async def test_async_research_and_wait_stream_close_falls_back_to_get(self): + """When the stream closes without a terminal event, async + research_and_wait does a final GET and returns the detail if completed.""" + import json + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + chunks = [ + b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_AsyncChunks(chunks), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "completed", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:02:30Z", + "completed_at": "2026-07-09T00:02:30Z", + "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, + }), + ) + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + detail = await research_and_wait_async( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + @pytest.mark.asyncio + async def test_async_research_and_wait_stream_close_task_running_raises_timeout(self): + """When the stream closes without a terminal event and the final GET + shows the task is still running, async research_and_wait raises TimeoutError.""" + import json + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + chunks = [ + b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', + ] + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_AsyncChunks(chunks), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", + }), + ) + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + content=json.dumps({ + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": "running", + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:00:10Z", + }), + ) + + transport = httpx.MockTransport(handler) + sdk_async_client = httpx.AsyncClient(transport=transport) + you = You( + server_url="http://mock.local", + async_client=sdk_async_client, + api_key_auth="test-api-key", + ) + + with pytest.raises(TimeoutError, match="still running"): + await research_and_wait_async( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) diff --git a/uv.lock b/uv.lock index a3a09df..4307a77 100644 --- a/uv.lock +++ b/uv.lock @@ -533,7 +533,7 @@ wheels = [ [[package]] name = "youdotcom" -version = "2.4.0" +version = "2.5.0" source = { editable = "." } dependencies = [ { name = "httpcore" }, From 09b252a2bd91801fc01f2c5e068517eb16f3bb2e Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Tue, 21 Jul 2026 16:58:14 -0700 Subject: [PATCH 02/17] refactor: ponytail shrinks + P2 fixes for research_helpers - Remove dead RawStreamEvent.retry field (never read by helpers/tests) - Extract _resolve_from_final_get[_async] helpers (dedup 4x final-GET logic) - Table-drive _raise_stream_error[_async] via _STREAM_ERROR_BRANCHES tuple - Factor _stream_build_kwargs/_stream_hook_ctx shared by sync/async stream-open - Inline _poll_detail[_async] into their single callers - P1 fix: replace daemon-thread + consumer.join with single-threaded loop using httpx read timeout (no more leaked blocked threads on socket hangs) - P2 fix: remove dead deadline/time.monotonic() inner check (unreachable - httpx.TimeoutException already enforces the timeout) - P2 fix: remove stream_exc/except BaseException indirection (fossil of threaded design; exceptions now propagate naturally) - P3: fix research_and_wait docstring (one GET -> a GET, second in fallback) - P3: document retry_config=None divergence from generated stream_research_task - Add 7 new typed-error tests (403/500/4XX/5XX/async 401/403/500) - Fix live test TypeError masking: check 'TaskResponse' in str(e) before skip - Extract _make_wait_handler factory + shared mock stream classes in tests 158/158 tests pass, pylint 10.00/10, mypy clean. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 531 +++++++++----------- tests/test_live.py | 18 +- tests/test_research_helpers.py | 810 ++++++++++-------------------- 3 files changed, 512 insertions(+), 847 deletions(-) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index 1f62982..51b98ec 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -24,11 +24,12 @@ import asyncio import json -import threading import time from dataclasses import dataclass from typing import Any, AsyncIterator, Iterator, Mapping, Optional +import httpx + import youdotcom.models as _models from youdotcom import errors as _errors from youdotcom._hooks.types import HookContext @@ -45,6 +46,16 @@ _TERMINAL_STREAM_EVENTS_OK = frozenset({"response.done", "complete", "completed"}) _TERMINAL_STREAM_EVENTS_ERR = frozenset({"error", "failed", "cancelled"}) +# Typed-error branches for stream open: (status, content_type, DataClass, ErrorClass). +# Mirrors the generated stream_research_task error handling so callers of +# stream_research get the same structured exceptions as the generated method. +_STREAM_ERROR_BRANCHES = ( + ("401", "application/json", _errors.StreamResearchTaskUnauthorizedErrorData, _errors.StreamResearchTaskUnauthorizedError), + ("403", "application/json", _errors.StreamResearchTaskForbiddenErrorData, _errors.StreamResearchTaskForbiddenError), + ("404", "application/json", _errors.StreamResearchTaskNotFoundErrorData, _errors.StreamResearchTaskNotFoundError), + ("500", "application/json", _errors.StreamResearchTaskInternalServerErrorData, _errors.StreamResearchTaskInternalServerError), +) + @dataclass class RawStreamEvent: @@ -56,13 +67,11 @@ class RawStreamEvent: against the documented enum. data: Parsed JSON payload when the data line is valid JSON, otherwise the raw string. - retry: Optional ``retry:`` directive from the server, if any. """ id: Optional[str] event: Optional[str] data: Any - retry: Optional[int] def _decode_raw_event(raw_json: str) -> RawStreamEvent: @@ -74,12 +83,11 @@ def _decode_raw_event(raw_json: str) -> RawStreamEvent: """ parsed = json.loads(raw_json) if not isinstance(parsed, dict): - return RawStreamEvent(id=None, event=None, data=parsed, retry=None) + return RawStreamEvent(id=None, event=None, data=parsed) return RawStreamEvent( id=parsed.get("id"), event=parsed.get("event"), data=parsed.get("data"), - retry=parsed.get("retry"), ) @@ -114,17 +122,23 @@ async def research_background_async(client: You, **kwargs: Any) -> TaskResponse: return res -def _poll_detail( +def poll_research_task( client: You, - *, task_id: str, - interval_s: float, - timeout_s: float, - deadline: float, + *, + interval_s: float = _DEFAULT_POLL_INTERVAL_S, + timeout_s: float = _DEFAULT_POLL_TIMEOUT_S, server_url: Optional[str] = None, timeout_ms: Optional[int] = None, http_headers: Optional[Mapping[str, str]] = None, ) -> TaskDetail: + """Poll ``GET /v1/research/{task_id}`` until ``status == "completed"``. + + Raises ``RuntimeError`` if the task ends in a non-completed terminal state + (``failed`` / ``cancelled``) and ``TimeoutError`` if ``timeout_s`` elapses + before completion. + """ + deadline = time.monotonic() + timeout_s while True: detail = client.get_research_task( task_id=task_id, @@ -146,7 +160,7 @@ def _poll_detail( time.sleep(interval_s) -def poll_research_task( +async def poll_research_task_async( client: You, task_id: str, *, @@ -156,35 +170,8 @@ def poll_research_task( timeout_ms: Optional[int] = None, http_headers: Optional[Mapping[str, str]] = None, ) -> TaskDetail: - """Poll ``GET /v1/research/{task_id}`` until ``status == "completed"``. - - Raises ``RuntimeError`` if the task ends in a non-completed terminal state - (``failed`` / ``cancelled``) and ``TimeoutError`` if ``timeout_s`` elapses - before completion. - """ - return _poll_detail( - client, - task_id=task_id, - interval_s=interval_s, - timeout_s=timeout_s, - deadline=time.monotonic() + timeout_s, - server_url=server_url, - timeout_ms=timeout_ms, - http_headers=http_headers, - ) - - -async def _poll_detail_async( - client: You, - *, - task_id: str, - interval_s: float, - timeout_s: float, - deadline: float, - server_url: Optional[str] = None, - timeout_ms: Optional[int] = None, - http_headers: Optional[Mapping[str, str]] = None, -) -> TaskDetail: + """Async variant of :func:`poll_research_task`.""" + deadline = time.monotonic() + timeout_s while True: detail = await client.get_research_task_async( task_id=task_id, @@ -206,27 +193,120 @@ async def _poll_detail_async( await asyncio.sleep(interval_s) -async def poll_research_task_async( +def _resolve_from_final_get( client: You, task_id: str, + result: Optional[str], + timeout_s: float, *, - interval_s: float = _DEFAULT_POLL_INTERVAL_S, - timeout_s: float = _DEFAULT_POLL_TIMEOUT_S, server_url: Optional[str] = None, timeout_ms: Optional[int] = None, http_headers: Optional[Mapping[str, str]] = None, + timed_out: bool = False, ) -> TaskDetail: - """Async variant of :func:`poll_research_task`.""" - return await _poll_detail_async( - client, + """Fetch the final ``TaskDetail`` and resolve based on the stream result. + + Called after the SSE stream terminates (terminal event, timeout, or close + without terminal event). When ``result`` is ``"ok"``, verifies the task is + completed. When ``result`` is a non-None error event name, raises + ``RuntimeError``. When ``result`` is None (timeout or stream close), + performs a final GET and returns the detail if completed, raises + ``RuntimeError`` for terminal non-completed, or ``TimeoutError`` if still + running. + """ + if result == "ok": + detail = client.get_research_task( + task_id=task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + if detail.status.value != "completed": + raise RuntimeError( + f"research task {task_id} stream signalled completion " + f"but GET returned status={detail.status.value}" + ) + return detail + if result is not None: + raise RuntimeError( + f"research task {task_id} ended in non-completed state: {result}" + ) + # result is None: timeout or stream close without terminal event + detail = client.get_research_task( task_id=task_id, - interval_s=interval_s, - timeout_s=timeout_s, - deadline=time.monotonic() + timeout_s, server_url=server_url, timeout_ms=timeout_ms, http_headers=http_headers, ) + status = detail.status.value + if status == "completed": + return detail + if status in _TERMINAL_TASK_STATUSES: + raise RuntimeError( + f"research task {task_id} ended in non-completed state: {status}" + ) + if timed_out: + raise TimeoutError( + f"research task {task_id} did not complete within {timeout_s}s " + f"(status: {status})" + ) + raise TimeoutError( + f"research task {task_id} stream closed without terminal event " + f"and task is still {status}" + ) + + +async def _resolve_from_final_get_async( + client: You, + task_id: str, + result: Optional[str], + timeout_s: float, + *, + server_url: Optional[str] = None, + timeout_ms: Optional[int] = None, + http_headers: Optional[Mapping[str, str]] = None, + timed_out: bool = False, +) -> TaskDetail: + """Async counterpart of :func:`_resolve_from_final_get`.""" + if result == "ok": + detail = await client.get_research_task_async( + task_id=task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + if detail.status.value != "completed": + raise RuntimeError( + f"research task {task_id} stream signalled completion " + f"but GET returned status={detail.status.value}" + ) + return detail + if result is not None: + raise RuntimeError( + f"research task {task_id} ended in non-completed state: {result}" + ) + detail = await client.get_research_task_async( + task_id=task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, + ) + status = detail.status.value + if status == "completed": + return detail + if status in _TERMINAL_TASK_STATUSES: + raise RuntimeError( + f"research task {task_id} ended in non-completed state: {status}" + ) + if timed_out: + raise TimeoutError( + f"research task {task_id} did not complete within {timeout_s}s " + f"(status: {status})" + ) + raise TimeoutError( + f"research task {task_id} stream closed without terminal event " + f"and task is still {status}" + ) def research_and_wait( @@ -239,7 +319,12 @@ def research_and_wait( Submits via :func:`research_background`, then opens the SSE stream and reads events until a terminal event arrives. Fetches the final - ``TaskDetail`` via one ``get_research_task`` call. + ``TaskDetail`` via a ``get_research_task`` call (a second GET is + issued in timeout/stream-close fallback paths). + + The stream is opened with a read timeout of ``timeout_s`` so that a + stalled server raises ``httpx.ReadTimeout`` deterministically instead + of leaking a blocked consumer thread. Parameters: timeout_s: Maximum seconds to wait for a terminal stream event. @@ -262,95 +347,32 @@ def research_and_wait( http_headers = kwargs.get("http_headers") task = research_background(client, **kwargs) + # Use timeout_s as the stream read timeout so iter_bytes raises + # httpx.ReadTimeout deterministically if the server stops sending data. + stream_timeout_ms = int(timeout_s * 1000) stream = _open_raw_stream( client, task.task_id, http_headers=http_headers, - server_url=server_url, timeout_ms=timeout_ms, + server_url=server_url, timeout_ms=stream_timeout_ms, ) result: Optional[str] = None - thread_exc: Optional[BaseException] = None - consumer: Optional[threading.Thread] = None + timed_out = False try: - def _consume() -> None: - nonlocal result, thread_exc - try: - for evt in stream: - if evt.event in _TERMINAL_STREAM_EVENTS_OK: - result = "ok" - return - if evt.event in _TERMINAL_STREAM_EVENTS_ERR: - result = evt.event - return - except BaseException as exc: # pylint: disable=broad-except - thread_exc = exc - - consumer = threading.Thread(target=_consume, daemon=True) - consumer.start() - consumer.join(timeout=timeout_s) - timed_out = consumer.is_alive() + for evt in stream: + if evt.event in _TERMINAL_STREAM_EVENTS_OK: + result = "ok" + break + if evt.event in _TERMINAL_STREAM_EVENTS_ERR: + result = evt.event + break + except httpx.TimeoutException: + timed_out = True finally: stream.close() - if consumer is not None: - consumer.join(timeout=2.0) - if timed_out: - # Stream didn't emit a terminal event within timeout_s. The task - # may have completed anyway (the SSE stream doesn't always emit - # terminal events for tasks that complete mid-stream). Do a final - # GET to check the actual status before raising. - detail = client.get_research_task( - task_id=task.task_id, - server_url=server_url, - timeout_ms=timeout_ms, - http_headers=http_headers, - ) - status = detail.status.value - if status == "completed": - return detail - if status in _TERMINAL_TASK_STATUSES: - raise RuntimeError( - f"research task {task.task_id} ended in non-completed state: {status}" - ) - raise TimeoutError( - f"research task {task.task_id} did not complete within {timeout_s}s " - f"(status: {status})" - ) - if thread_exc is not None: - raise thread_exc - if result == "ok": - detail = client.get_research_task( - task_id=task.task_id, - server_url=server_url, - timeout_ms=timeout_ms, - http_headers=http_headers, - ) - if detail.status.value != "completed": - raise RuntimeError( - f"research task {task.task_id} stream signalled completion " - f"but GET returned status={detail.status.value}" - ) - return detail - if result is not None: - raise RuntimeError( - f"research task {task.task_id} ended in non-completed state: {result}" - ) - # Stream closed without a terminal event — the task may have completed - # mid-stream. Do a final GET to check the actual status. - detail = client.get_research_task( - task_id=task.task_id, - server_url=server_url, - timeout_ms=timeout_ms, - http_headers=http_headers, - ) - status = detail.status.value - if status == "completed": - return detail - if status in _TERMINAL_TASK_STATUSES: - raise RuntimeError( - f"research task {task.task_id} ended in non-completed state: {status}" - ) - raise TimeoutError( - f"research task {task.task_id} stream closed without terminal event " - f"and task is still {status}" + return _resolve_from_final_get( + client, task.task_id, result, timeout_s, + server_url=server_url, timeout_ms=timeout_ms, http_headers=http_headers, + timed_out=timed_out, ) @@ -378,64 +400,17 @@ async def _consume() -> Optional[str]: return evt.event return None + timed_out = False try: result = await asyncio.wait_for(_consume(), timeout=timeout_s) - except asyncio.TimeoutError as exc: - # Stream didn't emit a terminal event within timeout_s. The task - # may have completed anyway — do a final GET to check. - detail = await client.get_research_task_async( - task_id=task.task_id, - server_url=server_url, - timeout_ms=timeout_ms, - http_headers=http_headers, - ) - status = detail.status.value - if status == "completed": - return detail - if status in _TERMINAL_TASK_STATUSES: - raise RuntimeError( - f"research task {task.task_id} ended in non-completed state: {status}" - ) from exc - raise TimeoutError( - f"research task {task.task_id} did not complete within {timeout_s}s " - f"(status: {status})" - ) from exc - - if result == "ok": - detail = await client.get_research_task_async( - task_id=task.task_id, - server_url=server_url, - timeout_ms=timeout_ms, - http_headers=http_headers, - ) - if detail.status.value != "completed": - raise RuntimeError( - f"research task {task.task_id} stream signalled completion " - f"but GET returned status={detail.status.value}" - ) - return detail - if result is not None: - raise RuntimeError( - f"research task {task.task_id} ended in non-completed state: {result}" - ) - # Stream closed without a terminal event — the task may have completed - # mid-stream. Do a final GET to check the actual status. - detail = await client.get_research_task_async( - task_id=task.task_id, - server_url=server_url, - timeout_ms=timeout_ms, - http_headers=http_headers, - ) - status = detail.status.value - if status == "completed": - return detail - if status in _TERMINAL_TASK_STATUSES: - raise RuntimeError( - f"research task {task.task_id} ended in non-completed state: {status}" - ) - raise TimeoutError( - f"research task {task.task_id} stream closed without terminal event " - f"and task is still {status}" + except asyncio.TimeoutError: + timed_out = True + result = None + + return await _resolve_from_final_get_async( + client, task.task_id, result, timeout_s, + server_url=server_url, timeout_ms=timeout_ms, http_headers=http_headers, + timed_out=timed_out, ) @@ -448,30 +423,11 @@ def _raise_stream_error(http_res: Any) -> None: method instead of a generic ``RuntimeError``. """ response_data: Any = None - if match_response(http_res, "401", "application/json"): - text = stream_to_text(http_res) - response_data = unmarshal_json_response( - _errors.StreamResearchTaskUnauthorizedErrorData, http_res, text - ) - raise _errors.StreamResearchTaskUnauthorizedError(response_data, http_res, text) - if match_response(http_res, "403", "application/json"): - text = stream_to_text(http_res) - response_data = unmarshal_json_response( - _errors.StreamResearchTaskForbiddenErrorData, http_res, text - ) - raise _errors.StreamResearchTaskForbiddenError(response_data, http_res, text) - if match_response(http_res, "404", "application/json"): - text = stream_to_text(http_res) - response_data = unmarshal_json_response( - _errors.StreamResearchTaskNotFoundErrorData, http_res, text - ) - raise _errors.StreamResearchTaskNotFoundError(response_data, http_res, text) - if match_response(http_res, "500", "application/json"): - text = stream_to_text(http_res) - response_data = unmarshal_json_response( - _errors.StreamResearchTaskInternalServerErrorData, http_res, text - ) - raise _errors.StreamResearchTaskInternalServerError(response_data, http_res, text) + for status, content_type, data_cls, err_cls in _STREAM_ERROR_BRANCHES: + if match_response(http_res, status, content_type): + text = stream_to_text(http_res) + response_data = unmarshal_json_response(data_cls, http_res, text) + raise err_cls(response_data, http_res, text) if match_response(http_res, "4XX", "*") or match_response(http_res, "5XX", "*"): text = stream_to_text(http_res) raise _errors.YouDefaultError("API error occurred", http_res, text) @@ -482,30 +438,11 @@ def _raise_stream_error(http_res: Any) -> None: async def _raise_stream_error_async(http_res: Any) -> None: """Async counterpart of :func:`_raise_stream_error`.""" response_data: Any = None - if match_response(http_res, "401", "application/json"): - text = await stream_to_text_async(http_res) - response_data = unmarshal_json_response( - _errors.StreamResearchTaskUnauthorizedErrorData, http_res, text - ) - raise _errors.StreamResearchTaskUnauthorizedError(response_data, http_res, text) - if match_response(http_res, "403", "application/json"): - text = await stream_to_text_async(http_res) - response_data = unmarshal_json_response( - _errors.StreamResearchTaskForbiddenErrorData, http_res, text - ) - raise _errors.StreamResearchTaskForbiddenError(response_data, http_res, text) - if match_response(http_res, "404", "application/json"): - text = await stream_to_text_async(http_res) - response_data = unmarshal_json_response( - _errors.StreamResearchTaskNotFoundErrorData, http_res, text - ) - raise _errors.StreamResearchTaskNotFoundError(response_data, http_res, text) - if match_response(http_res, "500", "application/json"): - text = await stream_to_text_async(http_res) - response_data = unmarshal_json_response( - _errors.StreamResearchTaskInternalServerErrorData, http_res, text - ) - raise _errors.StreamResearchTaskInternalServerError(response_data, http_res, text) + for status, content_type, data_cls, err_cls in _STREAM_ERROR_BRANCHES: + if match_response(http_res, status, content_type): + text = await stream_to_text_async(http_res) + response_data = unmarshal_json_response(data_cls, http_res, text) + raise err_cls(response_data, http_res, text) if match_response(http_res, "4XX", "*") or match_response(http_res, "5XX", "*"): text = await stream_to_text_async(http_res) raise _errors.YouDefaultError("API error occurred", http_res, text) @@ -513,6 +450,53 @@ async def _raise_stream_error_async(http_res: Any) -> None: raise _errors.YouDefaultError("Unexpected response received", http_res, text) +def _stream_build_kwargs( + client: You, + task_id: str, + *, + http_headers: Optional[Mapping[str, str]], + from_id: int, + server_url: Optional[str], + timeout_ms: Optional[int], +) -> dict[str, Any]: + """Build common kwargs dict for ``_build_request[_async]``.""" + # pylint: disable=protected-access + base_url = server_url if server_url is not None else client._get_url(None, None) + if timeout_ms is None: + timeout_ms = client.sdk_configuration.timeout_ms + return { + "method": "GET", + "path": "/v1/research/{task_id}/stream", + "base_url": base_url, + "url_variables": None, + "request": _models.StreamResearchTaskRequest(task_id=task_id, from_id=from_id), + "request_body_required": False, + "request_has_path_params": True, + "request_has_query_params": True, + "user_agent_header": "user-agent", + "accept_header_value": "text/event-stream", + "http_headers": http_headers, + "security": client.sdk_configuration.security, + "allow_empty_value": None, + "timeout_ms": timeout_ms, + } + + +def _stream_hook_ctx(client: You, base_url: str) -> HookContext: + """Build the shared ``HookContext`` for stream-open requests.""" + return HookContext( + config=client.sdk_configuration, + base_url=base_url or "", + operation_id="streamResearchTask", + oauth2_scopes=None, + security_source=get_security_from_env( + client.sdk_configuration.security, _models.Security + ), + tags=None, + extensions=None, + ) + + def _open_raw_stream( client: You, task_id: str, @@ -528,41 +512,21 @@ def _open_raw_stream( wires :func:`_decode_raw_event` into the ``EventStream`` instead of the strict pydantic ``ResearchTaskStreamEvent`` decoder. """ - # pylint: disable=protected-access - base_url = server_url if server_url is not None else client._get_url(None, None) - if timeout_ms is None: - timeout_ms = client.sdk_configuration.timeout_ms - req = client._build_request( - method="GET", - path="/v1/research/{task_id}/stream", - base_url=base_url, - url_variables=None, - request=_models.StreamResearchTaskRequest(task_id=task_id, from_id=from_id), - request_body_required=False, - request_has_path_params=True, - request_has_query_params=True, - user_agent_header="user-agent", - accept_header_value="text/event-stream", - http_headers=http_headers, - security=client.sdk_configuration.security, - allow_empty_value=None, - timeout_ms=timeout_ms, + kwargs = _stream_build_kwargs( + client, task_id, http_headers=http_headers, from_id=from_id, + server_url=server_url, timeout_ms=timeout_ms, ) + base_url = kwargs["base_url"] + # pylint: disable=protected-access + req = client._build_request(**kwargs) http_res = client.do_request( - hook_ctx=HookContext( - config=client.sdk_configuration, - base_url=base_url or "", - operation_id="streamResearchTask", - oauth2_scopes=None, - security_source=get_security_from_env( - client.sdk_configuration.security, _models.Security - ), - tags=None, - extensions=None, - ), + hook_ctx=_stream_hook_ctx(client, base_url), request=req, is_error_status_code=lambda c: match_status_codes(["4XX", "5XX"], c), stream=True, + # Diverges from generated stream_research_task which retries + # [429,500,502,503,504] on stream-open. We skip retries so a + # transient 5xx doesn't silently reopen a half-consumed stream. retry_config=None, ) if not match_response(http_res, "200", "text/event-stream"): @@ -614,38 +578,15 @@ async def _open_raw_stream_async( timeout_ms: Optional[int] = None, ) -> eventstreaming.EventStreamAsync[RawStreamEvent]: """Async counterpart of :func:`_open_raw_stream`.""" - # pylint: disable=protected-access - base_url = server_url if server_url is not None else client._get_url(None, None) - if timeout_ms is None: - timeout_ms = client.sdk_configuration.timeout_ms - req = client._build_request_async( - method="GET", - path="/v1/research/{task_id}/stream", - base_url=base_url, - url_variables=None, - request=_models.StreamResearchTaskRequest(task_id=task_id, from_id=from_id), - request_body_required=False, - request_has_path_params=True, - request_has_query_params=True, - user_agent_header="user-agent", - accept_header_value="text/event-stream", - http_headers=http_headers, - security=client.sdk_configuration.security, - allow_empty_value=None, - timeout_ms=timeout_ms, + kwargs = _stream_build_kwargs( + client, task_id, http_headers=http_headers, from_id=from_id, + server_url=server_url, timeout_ms=timeout_ms, ) + base_url = kwargs["base_url"] + # pylint: disable=protected-access + req = client._build_request_async(**kwargs) http_res = await client.do_request_async( - hook_ctx=HookContext( - config=client.sdk_configuration, - base_url=base_url or "", - operation_id="streamResearchTask", - oauth2_scopes=None, - security_source=get_security_from_env( - client.sdk_configuration.security, _models.Security - ), - tags=None, - extensions=None, - ), + hook_ctx=_stream_hook_ctx(client, base_url), request=req, is_error_status_code=lambda c: match_status_codes(["4XX", "5XX"], c), stream=True, diff --git a/tests/test_live.py b/tests/test_live.py index e621470..11df41c 100644 --- a/tests/test_live.py +++ b/tests/test_live.py @@ -565,8 +565,10 @@ def test_research_and_wait(self, you_client): input="What is the capital of France?", research_effort=ResearchEffort.LITE, ) - except TypeError: - pytest.skip("Background mode not enabled on server") + except TypeError as e: + if "TaskResponse" in str(e): + pytest.skip("Background mode not enabled on server") + raise assert isinstance(detail, TaskDetail) assert detail.status.value == "completed" @@ -587,8 +589,10 @@ def test_research_background_helper(self, you_client): input="What is the capital of France?", research_effort=ResearchEffort.LITE, ) - except TypeError: - pytest.skip("Background mode not enabled on server") + except TypeError as e: + if "TaskResponse" in str(e): + pytest.skip("Background mode not enabled on server") + raise assert isinstance(task, TaskResponse) assert task.task_id is not None @@ -609,8 +613,10 @@ def test_stream_research(self, you_client): input="What is the capital of France?", research_effort=ResearchEffort.LITE, ) - except TypeError: - pytest.skip("Background mode not enabled on server") + except TypeError as e: + if "TaskResponse" in str(e): + pytest.skip("Background mode not enabled on server") + raise assert isinstance(task, TaskResponse) diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py index e9a96c1..a871f89 100644 --- a/tests/test_research_helpers.py +++ b/tests/test_research_helpers.py @@ -1,5 +1,7 @@ """Tests for research background-mode helpers in youdotcom.research_helpers.""" +import asyncio +import json import os import uuid @@ -27,6 +29,112 @@ ) +# --------------------------------------------------------------------------- +# Shared test helpers: handler factories + mock stream classes. +# --------------------------------------------------------------------------- + +_TASK_RESPONSE_JSON = json.dumps({ + "task_id": "00000000-0000-0000-0000-000000000001", + "type": "research", + "status": "queued", + "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", + "created_at": "2026-07-09T00:00:00Z", +}) + +_DEFAULT_RESULT = {"output": {"content": "done", "content_type": "text", "sources": []}} + +_CONNECTED_CHUNK = b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n' + + +def _make_task_detail_json(status: str = "completed", result: dict | None = None) -> str: + """Build a TaskDetail JSON body for GET /v1/research/{task_id} responses.""" + detail: dict = { + "id": "00000000-0000-0000-0000-000000000001", + "task_type": "research", + "status": status, + "created_at": "2026-07-09T00:00:00Z", + "updated_at": "2026-07-09T00:02:30Z", + } + if status == "completed": + detail["completed_at"] = "2026-07-09T00:02:30Z" + if result is not None: + detail["result"] = result + return json.dumps(detail) + + +class _AsyncChunks(httpx.AsyncByteStream): + """Wrap a list of bytes chunks in an AsyncByteStream for MockTransport + + AsyncClient streaming responses.""" + + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + + async def __aiter__(self): + for chunk in self._chunks: + yield chunk + + +class _BlockingStream(httpx.SyncByteStream): + """Yields one event then raises ReadTimeout to simulate a stalled server + that stops sending data within the read timeout.""" + + def __iter__(self): + yield _CONNECTED_CHUNK + raise httpx.ReadTimeout("read timeout") + + +class _BlockingAsyncStream(httpx.AsyncByteStream): + """Yields one event then blocks so asyncio.wait_for times out.""" + + async def __aiter__(self): + yield _CONNECTED_CHUNK + await asyncio.sleep(100) + + +def _make_wait_handler( + *, + stream_chunks: list[bytes] | None = None, + stream_obj: httpx.SyncByteStream | httpx.AsyncByteStream | None = None, + final_status: str = "completed", + final_result: dict | None = _DEFAULT_RESULT, + is_async: bool = False, +): + """Create a MockTransport handler for research_and_wait tests. + + Returns SSE stream, POST TaskResponse, and GET TaskDetail responses. + Pass ``stream_obj`` for custom stream behavior (e.g. ``_BlockingStream``). + Pass ``stream_chunks`` for simple list-of-bytes streams. + """ + final_json = _make_task_detail_json(status=final_status, result=final_result) + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + if stream_obj is not None: + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=stream_obj, + ) + if stream_chunks is not None: + if is_async: + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, + stream=_AsyncChunks(stream_chunks), + ) + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=stream_chunks, + ) + return httpx.Response(200, content="{}") + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, headers={"content-type": "application/json"}, content=_TASK_RESPONSE_JSON, + ) + return httpx.Response( + 200, headers={"content-type": "application/json"}, content=final_json, + ) + + return handler + + @pytest.fixture def server_url(): return os.getenv("TEST_SERVER_URL", "http://localhost:18080") @@ -210,52 +318,14 @@ class TestResearchAndWait: def test_research_and_wait_returns_completed_detail(self): """research_and_wait submits, streams until terminal event, then fetches the final TaskDetail.""" - import json - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - chunks = [ - b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', - b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', - ] - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - content=chunks, - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - # GET /v1/research/{task_id} — final fetch after terminal event - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "completed", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:02:30Z", - "completed_at": "2026-07-09T00:02:30Z", - "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, - }), - ) - + handler = _make_wait_handler(stream_chunks=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ]) transport = httpx.MockTransport(handler) - sdk_client = httpx.Client(transport=transport) you = You( server_url="http://mock.local", - client=sdk_client, + client=httpx.Client(transport=transport), api_key_auth="test-api-key", ) @@ -272,39 +342,18 @@ def handler(request): def test_research_and_wait_error_event_raises_runtime_error(self): """research_and_wait raises RuntimeError when the stream emits an error terminal event.""" - import json - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - chunks = [ - b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', - b'id: 1\nevent: error\ndata: {"type":"error","task_id":"abc","message":"internal error"}\n\n', - ] - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - content=chunks, - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - return httpx.Response(200, content="{}") - + handler = _make_wait_handler( + stream_chunks=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: error\ndata: {"type":"error","task_id":"abc","message":"internal error"}\n\n', + ], + final_status="running", + final_result=None, + ) transport = httpx.MockTransport(handler) - sdk_client = httpx.Client(transport=transport) you = You( server_url="http://mock.local", - client=sdk_client, + client=httpx.Client(transport=transport), api_key_auth="test-api-key", ) @@ -318,55 +367,17 @@ def handler(request): def test_research_and_wait_timeout_raises_timeout_error(self): """research_and_wait raises TimeoutError when the stream never sends - a terminal event within timeout_s.""" - import json - import time as _time - - class _BlockingStream(httpx.SyncByteStream): - """Yields one event then blocks so the consumer thread - is still alive when the join timeout expires.""" - def __iter__(self): - yield b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n' - _time.sleep(100) - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - stream=_BlockingStream(), - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - # Final GET on timeout — task still running - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "running", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:00:10Z", - }), - ) - + a terminal event within timeout_s. The _BlockingStream simulates a + stalled server by raising httpx.ReadTimeout after the first event.""" + handler = _make_wait_handler( + stream_obj=_BlockingStream(), + final_status="running", + final_result=None, + ) transport = httpx.MockTransport(handler) - sdk_client = httpx.Client(transport=transport) you = You( server_url="http://mock.local", - client=sdk_client, + client=httpx.Client(transport=transport), api_key_auth="test-api-key", ) @@ -379,56 +390,16 @@ def handler(request): ) def test_research_and_wait_timeout_falls_back_to_get(self): - """When the stream times out but the task has completed, the final - GET fallback returns the completed detail.""" - import json - import time as _time - - class _BlockingStream(httpx.SyncByteStream): - def __iter__(self): - yield b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n' - _time.sleep(100) - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - stream=_BlockingStream(), - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - # Final GET on timeout — task completed despite no terminal event - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "completed", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:02:30Z", - "completed_at": "2026-07-09T00:02:30Z", - "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, - }), - ) - + """When the stream times out (ReadTimeout) but the task has completed, + the final GET fallback returns the completed detail.""" + handler = _make_wait_handler( + stream_obj=_BlockingStream(), + final_status="completed", + ) transport = httpx.MockTransport(handler) - sdk_client = httpx.Client(transport=transport) you = You( server_url="http://mock.local", - client=sdk_client, + client=httpx.Client(transport=transport), api_key_auth="test-api-key", ) @@ -445,52 +416,14 @@ def handler(request): def test_research_and_wait_stream_close_falls_back_to_get(self): """When the stream closes without a terminal event, research_and_wait does a final GET and returns the detail if completed.""" - import json - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - # Only a connected event, no terminal event - chunks = [ - b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', - ] - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - content=chunks, - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - # GET returns completed — the fallback should return this - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "completed", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:02:30Z", - "completed_at": "2026-07-09T00:02:30Z", - "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, - }), - ) - + handler = _make_wait_handler( + stream_chunks=[_CONNECTED_CHUNK], + final_status="completed", + ) transport = httpx.MockTransport(handler) - sdk_client = httpx.Client(transport=transport) you = You( server_url="http://mock.local", - client=sdk_client, + client=httpx.Client(transport=transport), api_key_auth="test-api-key", ) @@ -507,49 +440,15 @@ def handler(request): def test_research_and_wait_stream_close_task_running_raises_timeout(self): """When the stream closes without a terminal event and the final GET shows the task is still running, research_and_wait raises TimeoutError.""" - import json - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - chunks = [ - b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', - ] - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - content=chunks, - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - # GET returns running — task hasn't completed - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "running", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:00:10Z", - }), - ) - + handler = _make_wait_handler( + stream_chunks=[_CONNECTED_CHUNK], + final_status="running", + final_result=None, + ) transport = httpx.MockTransport(handler) - sdk_client = httpx.Client(transport=transport) you = You( server_url="http://mock.local", - client=sdk_client, + client=httpx.Client(transport=transport), api_key_auth="test-api-key", ) @@ -560,6 +459,10 @@ def handler(request): input="test query", research_effort=ResearchEffort.STANDARD, ) + + +# --------------------------------------------------------------------------- +# Stream research tolerant decoder: verify that # (a) helper accepts documented event types as before; # (b) helper accepts unknown event types without raising # pydantic.ValidationError, surfacing them as RawStreamEvent(event="..."). @@ -647,72 +550,95 @@ def test_unknown_event(self): # --------------------------------------------------------------------------- class TestStreamResearchTypedErrors: - def test_404_raises_not_found_error(self): - import json + """Verify stream_research[_async] raise the same typed errors as the + generated stream_research_task method for each status-code branch.""" + @staticmethod + def _make_error_handler(status_code: int, body: dict | None = None): + """Create a MockTransport handler that always returns an error response.""" + content = json.dumps(body or {"detail": "error"}) def handler(request): return httpx.Response( - 404, + status_code, headers={"content-type": "application/json"}, - content=json.dumps({"detail": "Task not found"}), + content=content, ) + return handler - transport = httpx.MockTransport(handler) - sdk_client = httpx.Client(transport=transport) - you = You( + @staticmethod + def _sync_you(handler): + return You( server_url="http://mock.local", - client=sdk_client, + client=httpx.Client(transport=httpx.MockTransport(handler)), + api_key_auth="test-api-key", + ) + + @staticmethod + def _async_you(handler): + return You( + server_url="http://mock.local", + async_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), api_key_auth="test-api-key", ) + _TASK = "00000000-0000-0000-0000-000000000001" + + def test_401_raises_unauthorized_error(self): + from youdotcom.errors import StreamResearchTaskUnauthorizedError + with pytest.raises(StreamResearchTaskUnauthorizedError): + list(stream_research(self._sync_you(self._make_error_handler(401)), self._TASK)) + + def test_403_raises_forbidden_error(self): + from youdotcom.errors import StreamResearchTaskForbiddenError + with pytest.raises(StreamResearchTaskForbiddenError): + list(stream_research(self._sync_you(self._make_error_handler(403)), self._TASK)) + + def test_404_raises_not_found_error(self): from youdotcom.errors import StreamResearchTaskNotFoundError with pytest.raises(StreamResearchTaskNotFoundError): - list(stream_research(you, "00000000-0000-0000-0000-000000000001")) + list(stream_research(self._sync_you(self._make_error_handler(404)), self._TASK)) - def test_401_raises_unauthorized_error(self): - import json + def test_500_raises_internal_server_error(self): + from youdotcom.errors import StreamResearchTaskInternalServerError + with pytest.raises(StreamResearchTaskInternalServerError): + list(stream_research(self._sync_you(self._make_error_handler(500)), self._TASK)) - def handler(request): - return httpx.Response( - 401, - headers={"content-type": "application/json"}, - content=json.dumps({"detail": "Invalid API key"}), - ) + def test_4xx_fallback_raises_default_error(self): + from youdotcom.errors import YouDefaultError + with pytest.raises(YouDefaultError): + list(stream_research(self._sync_you(self._make_error_handler(400)), self._TASK)) - transport = httpx.MockTransport(handler) - sdk_client = httpx.Client(transport=transport) - you = You( - server_url="http://mock.local", - client=sdk_client, - api_key_auth="test-api-key", - ) + def test_5xx_fallback_raises_default_error(self): + from youdotcom.errors import YouDefaultError + with pytest.raises(YouDefaultError): + list(stream_research(self._sync_you(self._make_error_handler(502)), self._TASK)) + @pytest.mark.asyncio + async def test_async_401_raises_unauthorized_error(self): from youdotcom.errors import StreamResearchTaskUnauthorizedError with pytest.raises(StreamResearchTaskUnauthorizedError): - list(stream_research(you, "00000000-0000-0000-0000-000000000001")) + async for _ in stream_research_async(self._async_you(self._make_error_handler(401)), self._TASK): + pass @pytest.mark.asyncio - async def test_async_404_raises_not_found_error(self): - import json - - def handler(request): - return httpx.Response( - 404, - headers={"content-type": "application/json"}, - content=json.dumps({"detail": "Task not found"}), - ) - - transport = httpx.MockTransport(handler) - sdk_async_client = httpx.AsyncClient(transport=transport) - you = You( - server_url="http://mock.local", - async_client=sdk_async_client, - api_key_auth="test-api-key", - ) + async def test_async_403_raises_forbidden_error(self): + from youdotcom.errors import StreamResearchTaskForbiddenError + with pytest.raises(StreamResearchTaskForbiddenError): + async for _ in stream_research_async(self._async_you(self._make_error_handler(403)), self._TASK): + pass + @pytest.mark.asyncio + async def test_async_404_raises_not_found_error(self): from youdotcom.errors import StreamResearchTaskNotFoundError with pytest.raises(StreamResearchTaskNotFoundError): - async for _ in stream_research_async(you, "00000000-0000-0000-0000-000000000001"): + async for _ in stream_research_async(self._async_you(self._make_error_handler(404)), self._TASK): + pass + + @pytest.mark.asyncio + async def test_async_500_raises_internal_server_error(self): + from youdotcom.errors import StreamResearchTaskInternalServerError + with pytest.raises(StreamResearchTaskInternalServerError): + async for _ in stream_research_async(self._async_you(self._make_error_handler(500)), self._TASK): pass @@ -866,18 +792,6 @@ def handler(request): # TestResearchAndWaitStreamMode. These also exercise the try/finally cleanup # path (replacing the broken contextlib.aclosing that called aclose()). -class _AsyncChunks(httpx.AsyncByteStream): - """Wrap a list of bytes chunks in an AsyncByteStream for MockTransport - + AsyncClient streaming responses (httpx MockTransport returns sync - streams by default, which AsyncClient rejects for stream=True).""" - - def __init__(self, chunks: list[bytes]): - self._chunks = chunks - - async def __aiter__(self): - for chunk in self._chunks: - yield chunk - class TestStreamResearchEventsTolerantAsync: @pytest.mark.asyncio @@ -930,51 +844,17 @@ class TestResearchAndWaitAsync: @pytest.mark.asyncio async def test_async_research_and_wait_returns_completed_detail(self): """Async research_and_wait: submit, stream until terminal, fetch detail.""" - import json - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - chunks = [ - b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', - b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', - ] - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - stream=_AsyncChunks(chunks), - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "completed", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:02:30Z", - "completed_at": "2026-07-09T00:02:30Z", - "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, - }), - ) - + handler = _make_wait_handler( + stream_chunks=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ], + is_async=True, + ) transport = httpx.MockTransport(handler) - sdk_async_client = httpx.AsyncClient(transport=transport) you = You( server_url="http://mock.local", - async_client=sdk_async_client, + async_client=httpx.AsyncClient(transport=transport), api_key_auth="test-api-key", ) @@ -991,39 +871,19 @@ def handler(request): @pytest.mark.asyncio async def test_async_research_and_wait_error_event_raises_runtime_error(self): """Async research_and_wait raises RuntimeError on error terminal event.""" - import json - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - chunks = [ - b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', - b'id: 1\nevent: error\ndata: {"type":"error","task_id":"abc","message":"internal error"}\n\n', - ] - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - stream=_AsyncChunks(chunks), - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - return httpx.Response(200, content="{}") - + handler = _make_wait_handler( + stream_chunks=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: error\ndata: {"type":"error","task_id":"abc","message":"internal error"}\n\n', + ], + final_status="running", + final_result=None, + is_async=True, + ) transport = httpx.MockTransport(handler) - sdk_async_client = httpx.AsyncClient(transport=transport) you = You( server_url="http://mock.local", - async_client=sdk_async_client, + async_client=httpx.AsyncClient(transport=transport), api_key_auth="test-api-key", ) @@ -1038,54 +898,18 @@ def handler(request): @pytest.mark.asyncio async def test_async_research_and_wait_timeout_raises_timeout_error(self): """Async research_and_wait raises TimeoutError when the stream never - sends a terminal event and the task is still running.""" - import asyncio - import json - - class _BlockingAsyncStream(httpx.AsyncByteStream): - """Yields one event then blocks so asyncio.wait_for times out.""" - async def __aiter__(self): - yield b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n' - await asyncio.sleep(100) - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - stream=_BlockingAsyncStream(), - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - # Final GET on timeout — task still running - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "running", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:00:10Z", - }), - ) - + sends a terminal event and the task is still running. asyncio.wait_for + cancels the blocked _consume() coroutine.""" + handler = _make_wait_handler( + stream_obj=_BlockingAsyncStream(), + final_status="running", + final_result=None, + is_async=True, + ) transport = httpx.MockTransport(handler) - sdk_async_client = httpx.AsyncClient(transport=transport) you = You( server_url="http://mock.local", - async_client=sdk_async_client, + async_client=httpx.AsyncClient(transport=transport), api_key_auth="test-api-key", ) @@ -1101,54 +925,15 @@ def handler(request): async def test_async_research_and_wait_timeout_falls_back_to_get(self): """When the async stream times out but the task completed, the final GET fallback returns the completed detail.""" - import asyncio - import json - - class _BlockingAsyncStream(httpx.AsyncByteStream): - async def __aiter__(self): - yield b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n' - await asyncio.sleep(100) - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - stream=_BlockingAsyncStream(), - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - # Final GET on timeout — task completed - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "completed", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:02:30Z", - "completed_at": "2026-07-09T00:02:30Z", - "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, - }), - ) - + handler = _make_wait_handler( + stream_obj=_BlockingAsyncStream(), + final_status="completed", + is_async=True, + ) transport = httpx.MockTransport(handler) - sdk_async_client = httpx.AsyncClient(transport=transport) you = You( server_url="http://mock.local", - async_client=sdk_async_client, + async_client=httpx.AsyncClient(transport=transport), api_key_auth="test-api-key", ) @@ -1166,50 +951,15 @@ def handler(request): async def test_async_research_and_wait_stream_close_falls_back_to_get(self): """When the stream closes without a terminal event, async research_and_wait does a final GET and returns the detail if completed.""" - import json - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - chunks = [ - b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', - ] - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - stream=_AsyncChunks(chunks), - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "completed", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:02:30Z", - "completed_at": "2026-07-09T00:02:30Z", - "result": {"output": {"content": "done", "content_type": "text", "sources": []}}, - }), - ) - + handler = _make_wait_handler( + stream_chunks=[_CONNECTED_CHUNK], + final_status="completed", + is_async=True, + ) transport = httpx.MockTransport(handler) - sdk_async_client = httpx.AsyncClient(transport=transport) you = You( server_url="http://mock.local", - async_client=sdk_async_client, + async_client=httpx.AsyncClient(transport=transport), api_key_auth="test-api-key", ) @@ -1227,48 +977,16 @@ def handler(request): async def test_async_research_and_wait_stream_close_task_running_raises_timeout(self): """When the stream closes without a terminal event and the final GET shows the task is still running, async research_and_wait raises TimeoutError.""" - import json - - def handler(request): - url = str(request.url) - if url.endswith("/stream") or "/stream?" in url: - chunks = [ - b'id: 0\nevent: connected\ndata: {"type":"connected","task_id":"abc","status":"running"}\n\n', - ] - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - stream=_AsyncChunks(chunks), - ) - if request.method == "POST" and "/v1/research" in url and "/stream" not in url: - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "task_id": "00000000-0000-0000-0000-000000000001", - "type": "research", - "status": "queued", - "stream_url": "/v1/research/00000000-0000-0000-0000-000000000001/stream", - "created_at": "2026-07-09T00:00:00Z", - }), - ) - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - content=json.dumps({ - "id": "00000000-0000-0000-0000-000000000001", - "task_type": "research", - "status": "running", - "created_at": "2026-07-09T00:00:00Z", - "updated_at": "2026-07-09T00:00:10Z", - }), - ) - + handler = _make_wait_handler( + stream_chunks=[_CONNECTED_CHUNK], + final_status="running", + final_result=None, + is_async=True, + ) transport = httpx.MockTransport(handler) - sdk_async_client = httpx.AsyncClient(transport=transport) you = You( server_url="http://mock.local", - async_client=sdk_async_client, + async_client=httpx.AsyncClient(transport=transport), api_key_auth="test-api-key", ) From aa6d1544f4c8b9968a14c5198f85331dcdc77e7f Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Tue, 21 Jul 2026 17:02:44 -0700 Subject: [PATCH 03/17] docs: fix 2.5.0 doc inaccuracies (terminal events, research_and_wait fallback, TypeError masking) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - CHANGELOG: fix stream_research_task terminal event list (was missing 'completed' in ok set and 'failed' in err set; now matches the actual _TERMINAL_STREAM_EVENTS_OK/ERR frozensets in research_helpers.py) - CHANGELOG: document research_and_wait timeout/stream-close fallback GET behavior (was described as just 'fetch the final TaskDetail') - MIGRATION: fix streaming example terminal event set to include 'complete' (was only 'response.done' + 'completed') - MIGRATION: document research_and_wait fallback GET behavior - examples: fix research_and_wait_example TypeError masking — check 'TaskResponse' in str(e) before skipping, re-raise real TypeErrors (matches the fix applied to test_live.py) - PR #23 description: fix research_and_wait row (was 'No polling fallback'; now documents the final-GET fallback) and stream_research_task terminal event list Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- CHANGELOG.md | 4 ++-- MIGRATION.md | 8 ++++++-- examples/api-example-calls.py | 8 +++++--- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cc87f1b..ab9b64f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,12 +13,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **`you.get_research_task(task_id=...)`**: Poll the status of a background research task via `GET /v1/research/{task_id}`. Returns a `TaskDetail` with `status`, `result` (populated when completed), `error` (populated when failed), and timing fields. -- **`you.stream_research_task(task_id=..., from_id=0)`**: Stream real-time updates for a background research task via `GET /v1/research/{task_id}/stream` (Server-Sent Events). Returns an `EventStream` of `ResearchTaskStreamEvent` objects. The connection closes automatically when the task reaches a terminal state (`response.done`, `complete`, `error`, or `cancelled`). The stream may also emit `completed`, `failed`, or `cancelled` as event names corresponding to the task's terminal status. +- **`you.stream_research_task(task_id=..., from_id=0)`**: Stream real-time updates for a background research task via `GET /v1/research/{task_id}/stream` (Server-Sent Events). Returns an `EventStream` of `ResearchTaskStreamEvent` objects. The connection closes automatically when the task reaches a terminal state. Terminal event names: `response.done`, `complete`, `completed` (success); `error`, `failed`, `cancelled` (failure). - **Convenience helpers** in `youdotcom.research_helpers` (hand-maintained, regen-safe): - `research_background(you, ...)` / `research_background_async(you, ...)` — submit and return `TaskResponse` directly (no Union narrowing needed). - `poll_research_task(you, task_id, ...)` / `poll_research_task_async(...)` — poll until terminal status (`completed`, `failed`, `cancelled`). - - `research_and_wait(you, ...)` / `research_and_wait_async(...)` — submit with `background=True`, then stream SSE events until a terminal event arrives, and fetch the final `TaskDetail`. For polling instead of streaming, use `poll_research_task` directly. + - `research_and_wait(you, ...)` / `research_and_wait_async(...)` — submit with `background=True`, then stream SSE events until a terminal event arrives, and fetch the final `TaskDetail`. If the stream times out or closes without a terminal event, a final `get_research_task` call resolves the status (returns the detail if completed, raises `RuntimeError` for terminal non-completed, or `TimeoutError` if still running). For polling instead of streaming, use `poll_research_task` directly. - `stream_research(you, task_id, ...)` / `stream_research_async(...)` — tolerant SSE iterator that surfaces undocumented event types as raw dicts instead of crashing on validation. **Recommended over `you.stream_research_task()` for real research tasks**, since the server emits intermediate workflow events not in the strict `Event` enum. - **New models**: `TaskResponse`, `TaskResponseStatus`, `TaskDetail`, `TaskDetailStatus`, `TaskDetailInput`, `Result`, `GetResearchTaskRequest`, `StreamResearchTaskRequest`, `ResearchTaskStreamEvent`, `ResearchTaskStreamEventData`, `Event`, `ResearchResult` (Union alias). diff --git a/MIGRATION.md b/MIGRATION.md index 04775f4..5662033 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -35,7 +35,7 @@ if detail.status.value == "completed": from youdotcom.research_helpers import stream_research for event in stream_research(you, task_id=res.task_id): print(event.event, event.data) - if event.event in ("response.done", "completed", "error", "failed", "cancelled"): + if event.event in ("response.done", "complete", "completed", "error", "failed", "cancelled"): break ``` @@ -49,6 +49,10 @@ from youdotcom.research_helpers import ( # Option 1: Submit and wait in one call (simplest) # Streams SSE events until the task completes, then fetches the result. +# If the stream times out or closes without a terminal event, a final +# get_research_task call resolves the status (returns the detail if +# completed, raises RuntimeError for terminal non-completed, or +# TimeoutError if still running). detail = research_and_wait( you, input="...", research_effort=ResearchEffort.DEEP, ) @@ -62,7 +66,7 @@ detail = poll_research_task(you, task_id=task.task_id) # the server emits intermediate event types not in the strict enum) for event in stream_research(you, task_id=task.task_id): print(event.event, event.data) - if event.event in ("response.done", "completed"): + if event.event in ("response.done", "complete", "completed"): break ``` diff --git a/examples/api-example-calls.py b/examples/api-example-calls.py index a846471..49acc94 100755 --- a/examples/api-example-calls.py +++ b/examples/api-example-calls.py @@ -356,9 +356,11 @@ def research_and_wait_example(): input="Compare the profitability of NVIDIA, AMD, and Intel over the past 5 fiscal years.", research_effort=ResearchEffort.DEEP, ) - except TypeError: - print(" Background mode not enabled on server. Skipping.") - return + except TypeError as e: + if "TaskResponse" in str(e): + print(" Background mode not enabled on server. Skipping.") + return + raise assert isinstance(detail, TaskDetail) print(f"Task completed: {detail.status.value}") From cc43e45d52f7f1b03e90d6f05b9efa4c055a0180 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Tue, 21 Jul 2026 17:49:37 -0700 Subject: [PATCH 04/17] fix: sync research_and_wait now enforces total wall-clock deadline (matches async) Addresses PR feedback (claude[bot] review 2026-07-22): sync research_and_wait used timeout_s only as a per-read stall timeout (httpx read timeout), while the async variant used asyncio.wait_for as a total wall-clock budget. This meant a server streaming heartbeats forever would never time out in sync mode. Fix: add a total deadline check inside the stream loop that breaks and sets timed_out=True when time.monotonic() >= deadline. Both sync and async now enforce timeout_s as a total budget. Added test_research_and_wait_total_deadline_not_stall_timeout: streams non-terminal 'ping' events every 10ms (fast enough to not trip the per-read timeout) and verifies the total deadline fires at timeout_s=0.5s. 159/159 tests pass, pylint 10.00/10, mypy clean. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 9 ++++++-- tests/test_research_helpers.py | 34 +++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index 51b98ec..37e72df 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -347,8 +347,9 @@ def research_and_wait( http_headers = kwargs.get("http_headers") task = research_background(client, **kwargs) - # Use timeout_s as the stream read timeout so iter_bytes raises - # httpx.ReadTimeout deterministically if the server stops sending data. + # Open the stream with a per-read timeout so a stalled server raises + # httpx.ReadTimeout deterministically. We also enforce a total wall-clock + # deadline below to match the async variant's asyncio.wait_for semantics. stream_timeout_ms = int(timeout_s * 1000) stream = _open_raw_stream( client, task.task_id, http_headers=http_headers, @@ -356,6 +357,7 @@ def research_and_wait( ) result: Optional[str] = None timed_out = False + deadline = time.monotonic() + timeout_s try: for evt in stream: if evt.event in _TERMINAL_STREAM_EVENTS_OK: @@ -364,6 +366,9 @@ def research_and_wait( if evt.event in _TERMINAL_STREAM_EVENTS_ERR: result = evt.event break + if time.monotonic() >= deadline: + timed_out = True + break except httpx.TimeoutException: timed_out = True finally: diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py index a871f89..6e49379 100644 --- a/tests/test_research_helpers.py +++ b/tests/test_research_helpers.py @@ -3,6 +3,7 @@ import asyncio import json import os +import time import uuid import httpx @@ -460,6 +461,39 @@ def test_research_and_wait_stream_close_task_running_raises_timeout(self): research_effort=ResearchEffort.STANDARD, ) + def test_research_and_wait_total_deadline_not_stall_timeout(self): + """research_and_wait enforces a total wall-clock deadline, not just a + per-read stall timeout. If the server keeps sending non-terminal + events forever, the total deadline fires and raises TimeoutError. + This matches the async variant's asyncio.wait_for semantics.""" + class _NonTerminalStream(httpx.SyncByteStream): + """Yields non-terminal events fast enough to not trip the per-read + timeout, but never emits a terminal event.""" + def __iter__(self): + for _ in range(10000): + yield b'id: 0\nevent: ping\ndata: {"type":"ping"}\n\n' + time.sleep(0.01) + + handler = _make_wait_handler( + stream_obj=_NonTerminalStream(), + final_status="running", + final_result=None, + ) + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + client=httpx.Client(transport=transport), + api_key_auth="test-api-key", + ) + + with pytest.raises(TimeoutError, match="did not complete within"): + research_and_wait( + you, + timeout_s=0.5, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + # --------------------------------------------------------------------------- # Stream research tolerant decoder: verify that From e000a0420f2e1e6799643a58bbdbf4173ccd7739 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Tue, 21 Jul 2026 18:21:05 -0700 Subject: [PATCH 05/17] test: fix Go mock comment (four->six terminal events) + add ok-but-get-non-completed test Addresses PR feedback (claude[bot] review 2026-07-22T00:57): - Fix Go mockserver comment: said 'four' terminal events but SDK treats six (OK: response.done, complete, completed; ERROR: error, failed, cancelled) - Add test_research_and_wait_ok_event_but_get_non_completed_raises: covers the defensive branch where stream emits response.done but the follow-up GET returns non-completed status (RuntimeError path in _resolve_from_final_get) 160/160 tests pass, pylint 10.00/10, mypy clean, Go builds clean. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- .../handler/pathgetv1researchstream.go | 4 ++- tests/test_research_helpers.py | 27 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/tests/mockserver/internal/handler/pathgetv1researchstream.go b/tests/mockserver/internal/handler/pathgetv1researchstream.go index 3527a83..4b0c57d 100644 --- a/tests/mockserver/internal/handler/pathgetv1researchstream.go +++ b/tests/mockserver/internal/handler/pathgetv1researchstream.go @@ -92,7 +92,9 @@ func testGetV1ResearchStreamSuccess(w http.ResponseWriter, req *http.Request) { } // 2) Terminal event — closes the stream. `response.done` is one of the - // four TERMINAL_SSE_EVENTS values that the SDK treats as stream-end. + // six terminal event names the SDK treats as stream-end: + // OK: response.done, complete, completed + // ERROR: error, failed, cancelled if !writeSSE("1", "response.done", map[string]interface{}{ "type": "response.done", "task_id": taskID, diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py index 6e49379..96dadf7 100644 --- a/tests/test_research_helpers.py +++ b/tests/test_research_helpers.py @@ -366,6 +366,33 @@ def test_research_and_wait_error_event_raises_runtime_error(self): research_effort=ResearchEffort.STANDARD, ) + def test_research_and_wait_ok_event_but_get_non_completed_raises(self): + """When the stream emits a terminal OK event (response.done) but the + follow-up GET returns a non-completed status, research_and_wait raises + RuntimeError. This covers the defensive branch in _resolve_from_final_get.""" + handler = _make_wait_handler( + stream_chunks=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ], + final_status="running", + final_result=None, + ) + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + client=httpx.Client(transport=transport), + api_key_auth="test-api-key", + ) + + with pytest.raises(RuntimeError, match="stream signalled completion but GET returned status=running"): + research_and_wait( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + def test_research_and_wait_timeout_raises_timeout_error(self): """research_and_wait raises TimeoutError when the stream never sends a terminal event within timeout_s. The _BlockingStream simulates a From 6ed04e8b2a53fe479787d560d67b6d551595da92 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 09:46:39 -0700 Subject: [PATCH 06/17] feat: add frontier research effort tier + streaming example Add ResearchEffort.FRONTIER for the highest-quality, longest-running research tasks (up to 4 hours). Frontier only works with background=true; sending it without returns 422. The enum value is injected via overlay (overlays/python_overlay.yaml) so it survives Speakeasy regenerations. - Update docstrings in researcheffort.py, researchop.py, sdk.py - Update CHANGELOG.md, MIGRATION.md, docs/models/researcheffort.md, docs/sdks/you/README.md - Add tests: frontier+background happy path, frontier-without-background 422 - Add research_stream_example() to examples/api-example-calls.py - Update PR #23 description Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- CHANGELOG.md | 17 +++++++++ MIGRATION.md | 38 ++++++++++++++++++++ docs/models/researcheffort.md | 4 ++- docs/sdks/you/README.md | 2 +- examples/api-example-calls.py | 49 ++++++++++++++++++++++++++ overlays/python_overlay.yaml | 33 +++++++++++++++++ src/youdotcom/models/researcheffort.py | 5 +++ src/youdotcom/models/researchop.py | 2 ++ src/youdotcom/sdk.py | 2 ++ tests/test_research.py | 39 ++++++++++++++++++++ 10 files changed, 189 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ab9b64f..9b4ea5e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- **`frontier` research effort tier**: New `ResearchEffort.FRONTIER` enum value for the highest-quality, longest-running research tasks. Frontier runs over longer durations with improved quality and accuracy. It only works with the task-based API (`background=true`); sending `frontier` without `background=true` returns a `422`. Use it with the background-mode helpers: + +```python +from youdotcom import You +from youdotcom.models import ResearchEffort +from youdotcom.research_helpers import research_and_wait + +you = You() +detail = research_and_wait( + you, + input="Evaluate the measurable global-health impact of the Gates Foundation", + research_effort=ResearchEffort.FRONTIER, + timeout_s=14400, # frontier tasks can run up to 4 hours +) +print(detail.result.model_dump()["output"]["content"]) +``` + - **Research Background Mode**: The `you.research()` method now accepts an optional `background=True` parameter. When enabled, the API queues the research task and returns a `TaskResponse` (with `task_id`, `type`, `status`, `stream_url`, `created_at`) immediately instead of waiting for the inline `ResearchResponse`. Use the new methods to poll or stream the task to completion. - **`you.get_research_task(task_id=...)`**: Poll the status of a background research task via `GET /v1/research/{task_id}`. Returns a `TaskDetail` with `status`, `result` (populated when completed), `error` (populated when failed), and timing fields. diff --git a/MIGRATION.md b/MIGRATION.md index 5662033..e1061d1 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -2,6 +2,44 @@ ## 2.4.0 → 2.5.0 (Latest) +### New `frontier` Research Effort Tier + +A new `ResearchEffort.FRONTIER` enum value has been added for the highest-quality, longest-running research tasks. Frontier runs over longer durations (up to 4 hours) with improved quality and accuracy. + +**Key constraint:** `frontier` only works with the task-based API (`background=true`). Sending `frontier` without `background=true` returns a `422`. + +```python +from youdotcom import You +from youdotcom.models import ResearchEffort +from youdotcom.research_helpers import research_and_wait + +you = You() + +# Submit a frontier task with background mode and wait for completion +detail = research_and_wait( + you, + input="Evaluate the measurable global-health impact of the Gates Foundation", + research_effort=ResearchEffort.FRONTIER, + timeout_s=14400, # frontier tasks can run up to 4 hours +) +print(detail.result.model_dump()["output"]["content"]) +``` + +You can also use the lower-level helpers for manual polling or streaming: + +```python +from youdotcom.research_helpers import research_background, poll_research_task + +task = research_background( + you, + input="Evaluate the measurable global-health impact of the Gates Foundation", + research_effort=ResearchEffort.FRONTIER, +) +detail = poll_research_task(you, task_id=task.task_id, timeout_s=14400) +``` + +No migration is required for existing code. The `frontier` tier is purely additive; existing `LITE`, `STANDARD`, `DEEP`, and `EXHAUSTIVE` values are unchanged. + ### New Background Mode for Research The `you.research()` method now accepts `background=True` to queue long-running research tasks asynchronously. The return type changes from `ResearchResponse` to `Union[ResearchResponse, TaskResponse]` (exposed as the `ResearchResult` alias). diff --git a/docs/models/researcheffort.md b/docs/models/researcheffort.md index 27ca90d..9203639 100644 --- a/docs/models/researcheffort.md +++ b/docs/models/researcheffort.md @@ -7,6 +7,7 @@ Available levels: - `standard`: The default. Balances speed and depth, a good fit for most questions. - `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. +- `frontier`: The highest-quality tier. Runs over longer durations with improved quality and accuracy. Only works with the task-based API (`background=true`); sending `frontier` without `background=true` returns a 422. ## Example Usage @@ -24,4 +25,5 @@ value = ResearchEffort.LITE | `LITE` | lite | | `STANDARD` | standard | | `DEEP` | deep | -| `EXHAUSTIVE` | exhaustive | \ No newline at end of file +| `EXHAUSTIVE` | exhaustive | +| `FRONTIER` | frontier | \ No newline at end of file diff --git a/docs/sdks/you/README.md b/docs/sdks/you/README.md index 642f0a3..e6eb9c9 100644 --- a/docs/sdks/you/README.md +++ b/docs/sdks/you/README.md @@ -421,7 +421,7 @@ with You( | Parameter | Type | Required | Description | | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `input` | *str* | :heavy_check_mark: | The research question or complex query requiring in-depth investigation and multi-step reasoning.

Note: The maximum length of the input is 40,000 characters. | -| `research_effort` | [Optional[models.ResearchEffort]](../../models/researcheffort.md) | :heavy_minus_sign: | Controls how much time and effort the Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time.

Available levels:
- `lite`: Returns answers quickly. Good for straightforward questions that just need a fast, reliable answer.
- `standard`: The default. Balances speed and depth, a good fit for most questions.
- `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed.
- `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. | +| `research_effort` | [Optional[models.ResearchEffort]](../../models/researcheffort.md) | :heavy_minus_sign: | Controls how much time and effort the Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time.

Available levels:
- `lite`: Returns answers quickly. Good for straightforward questions that just need a fast, reliable answer.
- `standard`: The default. Balances speed and depth, a good fit for most questions.
- `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed.
- `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result.
- `frontier`: The highest-quality tier. Runs over longer durations with improved quality and accuracy. Only works with the task-based API (`background=true`); sending `frontier` without `background=true` returns a 422. | | `source_control` | [Optional[models.SourceControl]](../../models/sourcecontrol.md) | :heavy_minus_sign: | Beta. Controls which web sources the research agent searches and visits. Use this to allow specific domains, block specific domains, boost specific domains, filter by recency, or focus web results by country.

`include_domains` and `exclude_domains` cannot be used together. Each domain list is capped at 500 entries. `exclude_domains` also blocks the research agent from visiting pages on those domains during browsing. `boost_domains` gives matching domains a relative ranking boost without filtering out other domains. It can be combined with `exclude_domains` but cannot be combined with `include_domains`. | | `output_schema` | Dict[str, *Any*] | :heavy_minus_sign: | Beta. Requests structured JSON output in output.content using a supported JSON Schema subset. Supported only with research_effort values standard, deep, and exhaustive. Sending output_schema with research_effort: "lite" returns 422.

Schema rules: Root must be a JSON object. Top-level anyOf is not allowed. Every object must define properties and set additionalProperties: false. Every property must be listed in required. Recursive schemas are not supported.

Limits: Max nesting depth 5, max total properties 100, max total enum values 500, max total schema string budget 25,000. | | `retries` | [Optional[utils.RetryConfig]](../../models/utils/retryconfig.md) | :heavy_minus_sign: | Configuration to override the default retry behavior of the client. | diff --git a/examples/api-example-calls.py b/examples/api-example-calls.py index 49acc94..9136685 100755 --- a/examples/api-example-calls.py +++ b/examples/api-example-calls.py @@ -371,6 +371,54 @@ def research_and_wait_example(): print(f"\nAnswer (preview):\n{content[:500]}...") +def research_stream_example(): + """ + Research API with streaming: submit in background mode, then stream + real-time SSE events with the tolerant stream_research helper. + + The tolerant helper surfaces undocumented intermediate event types + (e.g. research.searching, response.created) as raw dicts instead of + crashing on pydantic validation. Recommended over the generated + you.stream_research_task() for real tasks. + """ + from youdotcom.research_helpers import research_background, stream_research + + print("\n🚀 Running Research Stream (Helper)...\n") + + assert you is not None, "SDK client not initialized" + + try: + task = research_background( + you, + input="Compare the profitability of NVIDIA, AMD, and Intel over the past 5 fiscal years.", + research_effort=ResearchEffort.DEEP, + ) + except TypeError as e: + if "TaskResponse" in str(e): + print(" Background mode not enabled on server. Skipping.") + return + raise + + print(f"Queued task {task.task_id}, streaming events...\n") + + for event in stream_research(you, task_id=task.task_id): + print(f" event: {event.event} data: {str(event.data)[:120]}") + if event.event in ("response.done", "complete", "completed"): + print("\n Task completed via stream.") + break + if event.event in ("error", "failed", "cancelled"): + print(f"\n Task ended with: {event.event}") + break + + # Fetch the final result via a GET + detail = you.get_research_task(task_id=task.task_id) + if detail.status.value == "completed" and detail.result: + payload = detail.result.model_dump() + content = payload.get("output", {}).get("content", "") + if isinstance(content, str): + print(f"\nFinal answer (preview):\n{content[:500]}...") + + def research_output_schema_request(): """ Research API with `output_schema` for structured JSON output. @@ -487,6 +535,7 @@ def content_request_with_max_age(): {"name": "Research Request", "fn": research_request}, {"name": "Research Background Mode", "fn": research_background_request}, {"name": "Research and Wait (Helper)", "fn": research_and_wait_example}, + {"name": "Research Stream (Helper)", "fn": research_stream_example}, {"name": "Research with output_schema", "fn": research_output_schema_request}, {"name": "Finance Research Request", "fn": finance_research_request}, ] diff --git a/overlays/python_overlay.yaml b/overlays/python_overlay.yaml index 9065a61..8aac562 100644 --- a/overlays/python_overlay.yaml +++ b/overlays/python_overlay.yaml @@ -62,3 +62,36 @@ actions: - target: $["components"]["schemas"]["TaskDetail"]["properties"]["input"] update: additionalProperties: true + # Add the `frontier` research_effort tier to the ResearchEffort enum. + # Frontier is a long-running, higher-quality tier that only works with + # the task-based (background=true) API. The upstream spec doesn't yet + # include it, so this overlay injects the enum value and updates the + # description to document the new tier. Regen-durable: future Speakeasy + # regenerations will pick up `frontier` from the overlay-applied spec. + - target: $["components"]["schemas"]["ResearchEffort"]["enum"] + update: + - lite + - standard + - deep + - exhaustive + - frontier + - target: $["components"]["schemas"]["ResearchEffort"]["description"] + update: >- + Controls how much time and effort the Research API spends on your + question. Higher effort levels run more searches and dig deeper into + sources, at the cost of a longer response time. + + Available levels: + - `lite`: Returns answers quickly. Good for straightforward questions + that just need a fast, reliable answer. + - `standard`: The default. Balances speed and depth, a good fit for + most questions. + - `deep`: Spends more time researching and cross-referencing sources. + Use this when accuracy and thoroughness matter more than speed. + - `exhaustive`: The most thorough option. Explores the topic as fully + as possible, best suited for complex research tasks where you want + the highest quality result. + - `frontier`: The highest-quality tier. Runs over longer durations + with improved quality and accuracy. Only works with the task-based + API (`background=true`); sending `frontier` without `background=true` + returns a 422. diff --git a/src/youdotcom/models/researcheffort.py b/src/youdotcom/models/researcheffort.py index 2384195..506ea33 100644 --- a/src/youdotcom/models/researcheffort.py +++ b/src/youdotcom/models/researcheffort.py @@ -1,5 +1,8 @@ """Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" +# Manual overlay: `frontier` tier added via overlays/python_overlay.yaml +# until the upstream OpenAPI spec includes it. See that file for details. + from __future__ import annotations from enum import Enum @@ -12,9 +15,11 @@ class ResearchEffort(str, Enum): - `standard`: The default. Balances speed and depth, a good fit for most questions. - `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. + - `frontier`: The highest-quality tier. Runs over longer durations with improved quality and accuracy. Only works with the task-based API (``background=true``); sending ``frontier`` without ``background=true`` returns a 422. """ LITE = "lite" STANDARD = "standard" DEEP = "deep" EXHAUSTIVE = "exhaustive" + FRONTIER = "frontier" diff --git a/src/youdotcom/models/researchop.py b/src/youdotcom/models/researchop.py index 58e2443..15349a0 100644 --- a/src/youdotcom/models/researchop.py +++ b/src/youdotcom/models/researchop.py @@ -88,6 +88,7 @@ class ResearchRequestTypedDict(TypedDict): - `standard`: The default. Balances speed and depth, a good fit for most questions. - `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. + - `frontier`: The highest-quality tier. Runs over longer durations with improved quality and accuracy. Only works with the task-based API (`background=true`); sending `frontier` without `background=true` returns a 422. """ background: NotRequired[bool] r"""When true, queue a research task and return a task handle immediately instead of waiting for the result inline. Defaults to synchronous. When enabled, the response is a TaskResponse object with a task_id and stream_url for polling progress via GET /v1/research/{task_id} or streaming via GET /v1/research/{task_id}/stream.""" @@ -120,6 +121,7 @@ class ResearchRequest(BaseModel): - `standard`: The default. Balances speed and depth, a good fit for most questions. - `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. + - `frontier`: The highest-quality tier. Runs over longer durations with improved quality and accuracy. Only works with the task-based API (`background=true`); sending `frontier` without `background=true` returns a 422. """ background: Optional[bool] = False diff --git a/src/youdotcom/sdk.py b/src/youdotcom/sdk.py index 153869c..c979b82 100644 --- a/src/youdotcom/sdk.py +++ b/src/youdotcom/sdk.py @@ -556,6 +556,7 @@ def research( - `standard`: The default. Balances speed and depth, a good fit for most questions. - `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. + - `frontier`: The highest-quality tier. Runs over longer durations with improved quality and accuracy. Only works with the task-based API (`background=true`); sending `frontier` without `background=true` returns a 422. :param background: When true, queue a research task and return a task handle immediately instead of waiting for the result inline. Defaults to synchronous. When enabled, the response is a TaskResponse object with a task_id and stream_url for polling progress via GET /v1/research/{task_id} or streaming via GET /v1/research/{task_id}/stream. :param source_control: Beta. Controls which web sources the research agent searches and visits. Use this to allow specific domains, block specific domains, boost specific domains, filter by recency, or focus web results by country. @@ -698,6 +699,7 @@ async def research_async( - `standard`: The default. Balances speed and depth, a good fit for most questions. - `deep`: Spends more time researching and cross-referencing sources. Use this when accuracy and thoroughness matter more than speed. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex research tasks where you want the highest quality result. + - `frontier`: The highest-quality tier. Runs over longer durations with improved quality and accuracy. Only works with the task-based API (`background=true`); sending `frontier` without `background=true` returns a 422. :param background: When true, queue a research task and return a task handle immediately instead of waiting for the result inline. Defaults to synchronous. When enabled, the response is a TaskResponse object with a task_id and stream_url for polling progress via GET /v1/research/{task_id} or streaming via GET /v1/research/{task_id}/stream. :param source_control: Beta. Controls which web sources the research agent searches and visits. Use this to allow specific domains, block specific domains, boost specific domains, filter by recency, or focus web results by country. diff --git a/tests/test_research.py b/tests/test_research.py index 1e4768f..085d8a5 100644 --- a/tests/test_research.py +++ b/tests/test_research.py @@ -94,6 +94,23 @@ def test_research_exhaustive_effort(self, server_url, api_key): assert res.output.content is not None assert len(res.output.content) > 0 + def test_research_frontier_effort_with_background(self, server_url, api_key): + """frontier requires background=true; the mockserver returns a + TaskResponse when background=true regardless of effort tier.""" + client = create_test_http_client("post_/v1/research-background") + + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + res = you.research( + input="Evaluate the measurable global-health impact of the Gates Foundation", + research_effort=ResearchEffort.FRONTIER, + background=True, + server_url=server_url, + ) + + assert isinstance(res, TaskResponse) + assert res.task_id is not None + assert res.status.value == "queued" + def test_research_with_sources(self, server_url, api_key): client = create_test_http_client("post_/v1/research") @@ -558,6 +575,28 @@ def handler(request): ) sdk_client.close() + def test_frontier_without_background_raises_422(self): + """frontier requires background=true; sending it without returns 422.""" + def handler(request): + body = json.loads(request.content) + assert body.get("research_effort") == "frontier" + assert not body.get("background", False) + return httpx.Response( + 422, + headers={"content-type": "application/json"}, + content=json.dumps({"error": {"message": "frontier requires background=true"}}), + ) + + transport = httpx.MockTransport(handler) + sdk_client = httpx.Client(transport=transport) + you = You(server_url="http://mock.local", client=sdk_client, api_key_auth="test") + with pytest.raises(ResearchUnprocessableEntityError): + you.research( + input="Evaluate the Gates Foundation's global-health impact", + research_effort=ResearchEffort.FRONTIER, + ) + sdk_client.close() + class TestFinanceResearchAsync: @pytest.mark.asyncio From b89039ac262cfd0b35a6f779a448b35b6bb22333 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 09:55:44 -0700 Subject: [PATCH 07/17] feat: auto-adjust research_and_wait timeout for frontier + surface poll defaults research_and_wait now auto-selects timeout_s based on research_effort when omitted: 600s (10 min) for standard/deep/exhaustive, 14400s (4 hr) for frontier. Explicit timeout_s always takes precedence. poll_research_task cannot auto-detect (receives task_id, not effort) so its defaults (interval_s=2.0, timeout_s=600.0) are now documented in docstrings and MIGRATION.md with a per-tier guidance table. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- CHANGELOG.md | 6 +-- MIGRATION.md | 31 +++++++++++ src/youdotcom/research_helpers.py | 46 +++++++++++++++-- tests/test_research_helpers.py | 85 +++++++++++++++++++++++++++++++ 4 files changed, 161 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9b4ea5e..f28709c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,7 +21,7 @@ detail = research_and_wait( you, input="Evaluate the measurable global-health impact of the Gates Foundation", research_effort=ResearchEffort.FRONTIER, - timeout_s=14400, # frontier tasks can run up to 4 hours + # timeout_s auto-adjusts to 14400 (4 hours) for frontier when omitted ) print(detail.result.model_dump()["output"]["content"]) ``` @@ -34,8 +34,8 @@ print(detail.result.model_dump()["output"]["content"]) - **Convenience helpers** in `youdotcom.research_helpers` (hand-maintained, regen-safe): - `research_background(you, ...)` / `research_background_async(you, ...)` — submit and return `TaskResponse` directly (no Union narrowing needed). - - `poll_research_task(you, task_id, ...)` / `poll_research_task_async(...)` — poll until terminal status (`completed`, `failed`, `cancelled`). - - `research_and_wait(you, ...)` / `research_and_wait_async(...)` — submit with `background=True`, then stream SSE events until a terminal event arrives, and fetch the final `TaskDetail`. If the stream times out or closes without a terminal event, a final `get_research_task` call resolves the status (returns the detail if completed, raises `RuntimeError` for terminal non-completed, or `TimeoutError` if still running). For polling instead of streaming, use `poll_research_task` directly. + - `poll_research_task(you, task_id, ...)` / `poll_research_task_async(...)` — poll until terminal status (`completed`, `failed`, `cancelled`). Defaults: `interval_s=2.0`, `timeout_s=600.0` (10 minutes). For `frontier` tasks, pass `timeout_s=14400` (4 hours) explicitly since `poll_research_task` receives a `task_id` and cannot auto-detect the effort tier. + - `research_and_wait(you, ...)` / `research_and_wait_async(...)` — submit with `background=True`, then stream SSE events until a terminal event arrives, and fetch the final `TaskDetail`. If the stream times out or closes without a terminal event, a final `get_research_task` call resolves the status (returns the detail if completed, raises `RuntimeError` for terminal non-completed, or `TimeoutError` if still running). For polling instead of streaming, use `poll_research_task` directly. **`timeout_s` auto-adjusts** based on `research_effort` when omitted: 600s (10 min) for standard/deep/exhaustive, 14400s (4 hours) for `frontier`. - `stream_research(you, task_id, ...)` / `stream_research_async(...)` — tolerant SSE iterator that surfaces undocumented event types as raw dicts instead of crashing on validation. **Recommended over `you.stream_research_task()` for real research tasks**, since the server emits intermediate workflow events not in the strict `Event` enum. - **New models**: `TaskResponse`, `TaskResponseStatus`, `TaskDetail`, `TaskDetailStatus`, `TaskDetailInput`, `Result`, `GetResearchTaskRequest`, `StreamResearchTaskRequest`, `ResearchTaskStreamEvent`, `ResearchTaskStreamEventData`, `Event`, `ResearchResult` (Union alias). diff --git a/MIGRATION.md b/MIGRATION.md index e1061d1..9fc8f55 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -110,6 +110,37 @@ for event in stream_research(you, task_id=task.task_id): > **Note on streaming:** The generated `you.stream_research_task()` method uses a strict pydantic decoder that validates event names against a fixed `Event` enum. The server emits intermediate workflow events (e.g. `response.created`, `response.starting`, `response.output_item.added`) that are not in this enum, which causes `ResponseValidationError` on the first intermediate event. The `stream_research()` helper uses a tolerant decoder that surfaces unknown event names as raw dicts instead of crashing. For real research tasks, prefer `stream_research()`. +### Polling and Timeout Guidance + +The `poll_research_task` helper polls `GET /v1/research/{task_id}` at a configurable interval until the task reaches a terminal state. The `research_and_wait` helper streams SSE events and enforces a total wall-clock timeout. Both default sensibly for most effort tiers, but `frontier` tasks can run much longer. + +**`research_and_wait` auto-adjusts `timeout_s` based on `research_effort`** when you don't pass an explicit value: + +| `research_effort` | Auto `timeout_s` | Typical latency | +|-------------------|-------------------|-----------------| +| `lite` | 600s (10 min) | seconds | +| `standard` | 600s (10 min) | 30-120s | +| `deep` | 600s (10 min) | 120-300s | +| `exhaustive` | 600s (10 min) | 300-600s | +| `frontier` | 14400s (4 hours) | 300s - 4 hours | + +**`poll_research_task` does not auto-adjust** (it receives a `task_id`, not the effort level). You must set `timeout_s` explicitly for frontier: + +```python +from youdotcom.research_helpers import research_background, poll_research_task + +task = research_background( + you, input="...", research_effort=ResearchEffort.FRONTIER, +) +# poll_research_task defaults: interval_s=2.0, timeout_s=600.0 +# Override for frontier: +detail = poll_research_task( + you, task_id=task.task_id, + interval_s=5.0, # less aggressive for long-running tasks + timeout_s=14400, # 4 hours +) +``` + ### No Breaking Changes for Existing Code If you do not use `background=True`, your existing `you.research()` calls are unchanged at runtime. The return is still `ResearchResponse` when `background` is omitted or `False`. diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index 37e72df..bb72d86 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -42,10 +42,26 @@ _DEFAULT_POLL_INTERVAL_S = 2.0 _DEFAULT_POLL_TIMEOUT_S = 600.0 +_FRONTIER_TIMEOUT_S = 14400.0 # 4 hours — frontier tasks can run this long _TERMINAL_TASK_STATUSES = frozenset({"completed", "failed", "cancelled"}) _TERMINAL_STREAM_EVENTS_OK = frozenset({"response.done", "complete", "completed"}) _TERMINAL_STREAM_EVENTS_ERR = frozenset({"error", "failed", "cancelled"}) + +def _resolve_default_timeout(kwargs: Mapping[str, Any]) -> float: + """Pick a sensible default ``timeout_s`` based on ``research_effort``. + + Frontier tasks can run up to 4 hours; all other effort tiers complete + well within the standard 10-minute default. Only called when the caller + did not explicitly pass ``timeout_s``. + """ + effort = kwargs.get("research_effort") + if effort is not None: + effort_val = effort.value if hasattr(effort, "value") else str(effort) + if effort_val == "frontier": + return _FRONTIER_TIMEOUT_S + return _DEFAULT_POLL_TIMEOUT_S + # Typed-error branches for stream open: (status, content_type, DataClass, ErrorClass). # Mirrors the generated stream_research_task error handling so callers of # stream_research get the same structured exceptions as the generated method. @@ -134,6 +150,12 @@ def poll_research_task( ) -> TaskDetail: """Poll ``GET /v1/research/{task_id}`` until ``status == "completed"``. + Parameters: + interval_s: Seconds between poll requests. Defaults to 2.0. + timeout_s: Maximum seconds to wait before raising ``TimeoutError``. + Defaults to 600.0 (10 minutes). For ``frontier`` tasks (up to 4 + hours), pass ``timeout_s=14400``. + Raises ``RuntimeError`` if the task ends in a non-completed terminal state (``failed`` / ``cancelled``) and ``TimeoutError`` if ``timeout_s`` elapses before completion. @@ -170,7 +192,11 @@ async def poll_research_task_async( timeout_ms: Optional[int] = None, http_headers: Optional[Mapping[str, str]] = None, ) -> TaskDetail: - """Async variant of :func:`poll_research_task`.""" + """Async variant of :func:`poll_research_task`. + + See :func:`poll_research_task` for parameter defaults (interval_s=2.0, + timeout_s=600.0). + """ deadline = time.monotonic() + timeout_s while True: detail = await client.get_research_task_async( @@ -312,7 +338,7 @@ async def _resolve_from_final_get_async( def research_and_wait( client: You, *, - timeout_s: float = _DEFAULT_POLL_TIMEOUT_S, + timeout_s: Optional[float] = None, **kwargs: Any, ) -> TaskDetail: """Submit a research task in background mode and wait for completion. @@ -328,6 +354,10 @@ def research_and_wait( Parameters: timeout_s: Maximum seconds to wait for a terminal stream event. + When ``None`` (the default), the timeout is auto-selected based + on ``research_effort``: 600.0 (10 minutes) for standard/deep/ + exhaustive, 14400.0 (4 hours) for ``frontier``. Pass an explicit + value to override. ``**kwargs``: Forwarded to ``client.research``, including ``input`` and ``research_effort``. ``background=True`` is set automatically. ``server_url``, ``timeout_ms``, and @@ -342,6 +372,8 @@ def research_and_wait( and the task is still non-terminal. RuntimeError: If the task ends in a non-completed terminal state. """ + if timeout_s is None: + timeout_s = _resolve_default_timeout(kwargs) server_url = kwargs.get("server_url") timeout_ms = kwargs.get("timeout_ms") http_headers = kwargs.get("http_headers") @@ -384,10 +416,16 @@ def research_and_wait( async def research_and_wait_async( client: You, *, - timeout_s: float = _DEFAULT_POLL_TIMEOUT_S, + timeout_s: Optional[float] = None, **kwargs: Any, ) -> TaskDetail: - """Async variant of :func:`research_and_wait`.""" + """Async variant of :func:`research_and_wait`. + + See :func:`research_and_wait` for parameter details and auto-timeout + behavior (600s default, 14400s for frontier). + """ + if timeout_s is None: + timeout_s = _resolve_default_timeout(kwargs) server_url = kwargs.get("server_url") timeout_ms = kwargs.get("timeout_ms") http_headers = kwargs.get("http_headers") diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py index 96dadf7..ce592ed 100644 --- a/tests/test_research_helpers.py +++ b/tests/test_research_helpers.py @@ -27,6 +27,9 @@ stream_research, stream_research_async, _decode_raw_event, + _resolve_default_timeout, + _FRONTIER_TIMEOUT_S, + _DEFAULT_POLL_TIMEOUT_S, ) @@ -585,6 +588,88 @@ def record_send(request): # And confirm the SDK still set User-Agent on the underlying request # (the YDCUserAgentOverrideHook ran before send). assert recorded_ua["value"] == f"youdotcom-python-sdk/{you.sdk_configuration.sdk_version}" +# --------------------------------------------------------------------------- +# _resolve_default_timeout: auto-adjust timeout based on research_effort. +# --------------------------------------------------------------------------- + +class TestResolveDefaultTimeout: + def test_frontier_returns_4_hour_timeout(self): + assert _resolve_default_timeout({"research_effort": ResearchEffort.FRONTIER}) == _FRONTIER_TIMEOUT_S + + def test_frontier_string_returns_4_hour_timeout(self): + assert _resolve_default_timeout({"research_effort": "frontier"}) == _FRONTIER_TIMEOUT_S + + def test_standard_returns_default_timeout(self): + assert _resolve_default_timeout({"research_effort": ResearchEffort.STANDARD}) == _DEFAULT_POLL_TIMEOUT_S + + def test_deep_returns_default_timeout(self): + assert _resolve_default_timeout({"research_effort": ResearchEffort.DEEP}) == _DEFAULT_POLL_TIMEOUT_S + + def test_exhaustive_returns_default_timeout(self): + assert _resolve_default_timeout({"research_effort": ResearchEffort.EXHAUSTIVE}) == _DEFAULT_POLL_TIMEOUT_S + + def test_no_effort_returns_default_timeout(self): + assert _resolve_default_timeout({}) == _DEFAULT_POLL_TIMEOUT_S + + +# --------------------------------------------------------------------------- +# research_and_wait frontier auto-timeout: when timeout_s is omitted and +# research_effort=frontier, the helper should use 14400s (4 hours), not the +# 600s default. We verify by checking that a short stream timeout is NOT +# applied — instead the 4-hour deadline is used, so the stream consumes +# all events without a premature TimeoutError. +# --------------------------------------------------------------------------- + +class TestResearchAndWaitFrontierAutoTimeout: + def test_frontier_auto_timeout_completes_without_premature_timeout(self): + """When research_effort=frontier and timeout_s is omitted, the + auto-adjusted 4-hour timeout should not trip on a normal event + sequence that would complete within seconds.""" + handler = _make_wait_handler(stream_chunks=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ]) + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + client=httpx.Client(transport=transport), + api_key_auth="test-api-key", + ) + + detail = research_and_wait( + you, + input="Evaluate the Gates Foundation's global-health impact", + research_effort=ResearchEffort.FRONTIER, + # timeout_s intentionally omitted — should auto-adjust to 14400 + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + def test_explicit_timeout_overrides_auto_adjust(self): + """When the user passes an explicit timeout_s, it takes precedence + over the frontier auto-adjustment.""" + handler = _make_wait_handler( + stream_obj=_BlockingStream(), + final_status="running", + final_result=None, + ) + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + client=httpx.Client(transport=transport), + api_key_auth="test-api-key", + ) + + with pytest.raises(TimeoutError, match="did not complete within 0.5"): + research_and_wait( + you, + timeout_s=0.5, + input="test query", + research_effort=ResearchEffort.FRONTIER, + ) + + # --------------------------------------------------------------------------- # _decode_raw_event sanity: known and unknown shapes. # --------------------------------------------------------------------------- From edb2aeb8e9372ebeba82d8d29e78d0482fdf70c7 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 10:25:04 -0700 Subject: [PATCH 08/17] feat: add FinanceResearchEffort.LITE + fix PR review findings Add LITE tier to FinanceResearchEffort enum (overlay-injected, same pattern as ResearchEffort.FRONTIER). Update docstrings, docs, CHANGELOG, and MIGRATION. Add test_finance_research_lite_effort. Fix three PR #24 review items: 1. _resolve_from_final_get now re-polls up to 3 times (1s interval) when the stream signals OK but GET returns non-completed, tolerating backend commit races. Updated test_research_and_wait_ok_event_but_get_non_completed_raises and added test_research_and_wait_ok_event_repoll_succeeds. 2. research_and_wait_async now catches httpx.TimeoutException in addition to asyncio.TimeoutError, matching sync behavior. Added test_async_research_and_wait_httpx_readtimeout_falls_back_to_get. 3. Removed misplaced pylint disable-next=protected-access in stream_research. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- CHANGELOG.md | 2 + MIGRATION.md | 17 ++++ docs/models/financeresearcheffort.md | 2 + docs/sdks/you/README.md | 2 +- overlays/python_overlay.yaml | 24 +++++ src/youdotcom/models/finance_researchop.py | 2 + src/youdotcom/models/financeresearcheffort.py | 5 ++ src/youdotcom/research_helpers.py | 72 ++++++++------- src/youdotcom/sdk.py | 2 + tests/test_research.py | 14 +++ tests/test_research_helpers.py | 87 ++++++++++++++++++- 11 files changed, 198 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f28709c..75363e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,8 @@ detail = research_and_wait( print(detail.result.model_dump()["output"]["content"]) ``` +- **`lite` finance research effort tier**: New `FinanceResearchEffort.LITE` enum value for the Finance Research API. Returns answers quickly for straightforward financial questions. The default remains `deep`. No migration required; existing `DEEP` and `EXHAUSTIVE` values are unchanged. + - **Research Background Mode**: The `you.research()` method now accepts an optional `background=True` parameter. When enabled, the API queues the research task and returns a `TaskResponse` (with `task_id`, `type`, `status`, `stream_url`, `created_at`) immediately instead of waiting for the inline `ResearchResponse`. Use the new methods to poll or stream the task to completion. - **`you.get_research_task(task_id=...)`**: Poll the status of a background research task via `GET /v1/research/{task_id}`. Returns a `TaskDetail` with `status`, `result` (populated when completed), `error` (populated when failed), and timing fields. diff --git a/MIGRATION.md b/MIGRATION.md index 9fc8f55..d25109e 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -40,6 +40,23 @@ detail = poll_research_task(you, task_id=task.task_id, timeout_s=14400) No migration is required for existing code. The `frontier` tier is purely additive; existing `LITE`, `STANDARD`, `DEEP`, and `EXHAUSTIVE` values are unchanged. +### New `lite` Finance Research Effort Tier + +A new `FinanceResearchEffort.LITE` enum value has been added for quick, straightforward financial questions. The default remains `deep`. + +```python +from youdotcom import You +from youdotcom.models import FinanceResearchEffort + +you = You() +res = you.finance_research( + input="What was Apple's revenue in FY2024?", + research_effort=FinanceResearchEffort.LITE, +) +``` + +No migration is required. Existing `DEEP` and `EXHAUSTIVE` values are unchanged. + ### New Background Mode for Research The `you.research()` method now accepts `background=True` to queue long-running research tasks asynchronously. The return type changes from `ResearchResponse` to `Union[ResearchResponse, TaskResponse]` (exposed as the `ResearchResult` alias). diff --git a/docs/models/financeresearcheffort.md b/docs/models/financeresearcheffort.md index 9402be1..6c7ebde 100644 --- a/docs/models/financeresearcheffort.md +++ b/docs/models/financeresearcheffort.md @@ -3,6 +3,7 @@ Controls how much time and effort the Finance Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time. Available levels: +- `lite`: Returns answers quickly. Good for straightforward financial questions that just need a fast, reliable answer. - `deep`: The default. Spends more time researching and cross-referencing sources. Good for most financial questions, including multi-company comparisons, earnings analysis, and regulatory research. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex financial research tasks where you want the highest quality result. @@ -19,5 +20,6 @@ value = FinanceResearchEffort.DEEP | Name | Value | | ------------ | ------------ | +| `LITE` | lite | | `DEEP` | deep | | `EXHAUSTIVE` | exhaustive | \ No newline at end of file diff --git a/docs/sdks/you/README.md b/docs/sdks/you/README.md index e6eb9c9..cae12e2 100644 --- a/docs/sdks/you/README.md +++ b/docs/sdks/you/README.md @@ -722,7 +722,7 @@ with You( | Parameter | Type | Required | Description | Example | | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `input` | *str* | :heavy_check_mark: | The financial research question or complex query requiring in-depth investigation and multi-step reasoning.

Note: The maximum length of the input is 40,000 characters. | What were the key drivers of NVIDIA's revenue growth in fiscal year 2025? | -| `research_effort` | [Optional[models.FinanceResearchEffort]](../../models/financeresearcheffort.md) | :heavy_minus_sign: | Controls how much time and effort the Finance Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time.

Available levels:
- `deep`: The default. Spends more time researching and cross-referencing sources. Good for most financial questions, including multi-company comparisons, earnings analysis, and regulatory research.
- `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex financial research tasks where you want the highest quality result. | deep | +| `research_effort` | [Optional[models.FinanceResearchEffort]](../../models/financeresearcheffort.md) | :heavy_minus_sign: | Controls how much time and effort the Finance Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time.

Available levels:
- `lite`: Returns answers quickly. Good for straightforward financial questions that just need a fast, reliable answer.
- `deep`: The default. Spends more time researching and cross-referencing sources. Good for most financial questions, including multi-company comparisons, earnings analysis, and regulatory research.
- `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex financial research tasks where you want the highest quality result. | deep | | `retries` | [Optional[utils.RetryConfig]](../../models/utils/retryconfig.md) | :heavy_minus_sign: | Configuration to override the default retry behavior of the client. | | ### Response diff --git a/overlays/python_overlay.yaml b/overlays/python_overlay.yaml index 8aac562..905aed5 100644 --- a/overlays/python_overlay.yaml +++ b/overlays/python_overlay.yaml @@ -95,3 +95,27 @@ actions: with improved quality and accuracy. Only works with the task-based API (`background=true`); sending `frontier` without `background=true` returns a 422. + # Add the `lite` research_effort tier to the FinanceResearchEffort enum. + # The upstream spec doesn't yet include it, so this overlay injects the + # enum value and updates the description. Regen-durable. + - target: $["components"]["schemas"]["FinanceResearchEffort"]["enum"] + update: + - lite + - deep + - exhaustive + - target: $["components"]["schemas"]["FinanceResearchEffort"]["description"] + update: >- + Controls how much time and effort the Finance Research API spends on + your question. Higher effort levels run more searches and dig deeper + into sources, at the cost of a longer response time. + + Available levels: + - `lite`: Returns answers quickly. Good for straightforward financial + questions that just need a fast, reliable answer. + - `deep`: The default. Spends more time researching and + cross-referencing sources. Good for most financial questions, + including multi-company comparisons, earnings analysis, and + regulatory research. + - `exhaustive`: The most thorough option. Explores the topic as fully + as possible, best suited for complex financial research tasks where + you want the highest quality result. diff --git a/src/youdotcom/models/finance_researchop.py b/src/youdotcom/models/finance_researchop.py index 1d886fd..47259fd 100644 --- a/src/youdotcom/models/finance_researchop.py +++ b/src/youdotcom/models/finance_researchop.py @@ -19,6 +19,7 @@ class FinanceResearchRequestTypedDict(TypedDict): r"""Controls how much time and effort the Finance Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time. Available levels: + - `lite`: Returns answers quickly. Good for straightforward financial questions that just need a fast, reliable answer. - `deep`: The default. Spends more time researching and cross-referencing sources. Good for most financial questions, including multi-company comparisons, earnings analysis, and regulatory research. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex financial research tasks where you want the highest quality result. """ @@ -35,6 +36,7 @@ class FinanceResearchRequest(BaseModel): r"""Controls how much time and effort the Finance Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time. Available levels: + - `lite`: Returns answers quickly. Good for straightforward financial questions that just need a fast, reliable answer. - `deep`: The default. Spends more time researching and cross-referencing sources. Good for most financial questions, including multi-company comparisons, earnings analysis, and regulatory research. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex financial research tasks where you want the highest quality result. """ diff --git a/src/youdotcom/models/financeresearcheffort.py b/src/youdotcom/models/financeresearcheffort.py index 417d805..909b887 100644 --- a/src/youdotcom/models/financeresearcheffort.py +++ b/src/youdotcom/models/financeresearcheffort.py @@ -1,5 +1,8 @@ """Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" +# Manual overlay: `lite` tier added via overlays/python_overlay.yaml +# until the upstream OpenAPI spec includes it. See that file for details. + from __future__ import annotations from enum import Enum @@ -8,9 +11,11 @@ class FinanceResearchEffort(str, Enum): r"""Controls how much time and effort the Finance Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time. Available levels: + - `lite`: Returns answers quickly. Good for straightforward financial questions that just need a fast, reliable answer. - `deep`: The default. Spends more time researching and cross-referencing sources. Good for most financial questions, including multi-company comparisons, earnings analysis, and regulatory research. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex financial research tasks where you want the highest quality result. """ + LITE = "lite" DEEP = "deep" EXHAUSTIVE = "exhaustive" diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index bb72d86..f5e7bef 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -43,6 +43,8 @@ _DEFAULT_POLL_INTERVAL_S = 2.0 _DEFAULT_POLL_TIMEOUT_S = 600.0 _FRONTIER_TIMEOUT_S = 14400.0 # 4 hours — frontier tasks can run this long +_REPOLL_MAX_ATTEMPTS = 3 # bounded re-poll on completion race (stream OK but GET not yet completed) +_REPOLL_INTERVAL_S = 1.0 _TERMINAL_TASK_STATUSES = frozenset({"completed", "failed", "cancelled"}) _TERMINAL_STREAM_EVENTS_OK = frozenset({"response.done", "complete", "completed"}) _TERMINAL_STREAM_EVENTS_ERR = frozenset({"error", "failed", "cancelled"}) @@ -234,25 +236,33 @@ def _resolve_from_final_get( Called after the SSE stream terminates (terminal event, timeout, or close without terminal event). When ``result`` is ``"ok"``, verifies the task is - completed. When ``result`` is a non-None error event name, raises - ``RuntimeError``. When ``result`` is None (timeout or stream close), - performs a final GET and returns the detail if completed, raises - ``RuntimeError`` for terminal non-completed, or ``TimeoutError`` if still - running. + completed with a short bounded re-poll to tolerate backend commit races. + When ``result`` is a non-None error event name, raises ``RuntimeError``. + When ``result`` is None (timeout or stream close), performs a final GET + and returns the detail if completed, raises ``RuntimeError`` for terminal + non-completed, or ``TimeoutError`` if still running. """ if result == "ok": - detail = client.get_research_task( - task_id=task_id, - server_url=server_url, - timeout_ms=timeout_ms, - http_headers=http_headers, - ) - if detail.status.value != "completed": - raise RuntimeError( - f"research task {task_id} stream signalled completion " - f"but GET returned status={detail.status.value}" + # The stream signalled completion, but the persisted task status may + # not be committed yet (distributed backend race). Re-poll a few times + # before giving up instead of failing on the first GET. + for _ in range(_REPOLL_MAX_ATTEMPTS): + detail = client.get_research_task( + task_id=task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, ) - return detail + if detail.status.value == "completed": + return detail + if detail.status.value in _TERMINAL_TASK_STATUSES: + break # terminal non-completed, fall through to error below + time.sleep(_REPOLL_INTERVAL_S) + raise RuntimeError( + f"research task {task_id} stream signalled completion " + f"but GET returned status={detail.status.value} " + f"after {_REPOLL_MAX_ATTEMPTS} attempts" + ) if result is not None: raise RuntimeError( f"research task {task_id} ended in non-completed state: {result}" @@ -295,18 +305,23 @@ async def _resolve_from_final_get_async( ) -> TaskDetail: """Async counterpart of :func:`_resolve_from_final_get`.""" if result == "ok": - detail = await client.get_research_task_async( - task_id=task_id, - server_url=server_url, - timeout_ms=timeout_ms, - http_headers=http_headers, - ) - if detail.status.value != "completed": - raise RuntimeError( - f"research task {task_id} stream signalled completion " - f"but GET returned status={detail.status.value}" + for _ in range(_REPOLL_MAX_ATTEMPTS): + detail = await client.get_research_task_async( + task_id=task_id, + server_url=server_url, + timeout_ms=timeout_ms, + http_headers=http_headers, ) - return detail + if detail.status.value == "completed": + return detail + if detail.status.value in _TERMINAL_TASK_STATUSES: + break + await asyncio.sleep(_REPOLL_INTERVAL_S) + raise RuntimeError( + f"research task {task_id} stream signalled completion " + f"but GET returned status={detail.status.value} " + f"after {_REPOLL_MAX_ATTEMPTS} attempts" + ) if result is not None: raise RuntimeError( f"research task {task_id} ended in non-completed state: {result}" @@ -446,7 +461,7 @@ async def _consume() -> Optional[str]: timed_out = False try: result = await asyncio.wait_for(_consume(), timeout=timeout_s) - except asyncio.TimeoutError: + except (asyncio.TimeoutError, httpx.TimeoutException): timed_out = True result = None @@ -603,7 +618,6 @@ def stream_research( from_id: Sequence number to resume from (reconnection). ``0`` starts at the beginning of the stream. """ - # pylint: disable-next=protected-access with _open_raw_stream( client, task_id, http_headers=http_headers, from_id=from_id, server_url=server_url, timeout_ms=timeout_ms, diff --git a/src/youdotcom/sdk.py b/src/youdotcom/sdk.py index c979b82..166fce4 100644 --- a/src/youdotcom/sdk.py +++ b/src/youdotcom/sdk.py @@ -1316,6 +1316,7 @@ def finance_research( :param research_effort: Controls how much time and effort the Finance Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time. Available levels: + - `lite`: Returns answers quickly. Good for straightforward financial questions that just need a fast, reliable answer. - `deep`: The default. Spends more time researching and cross-referencing sources. Good for most financial questions, including multi-company comparisons, earnings analysis, and regulatory research. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex financial research tasks where you want the highest quality result. :param retries: Override the default retry configuration for this method @@ -1440,6 +1441,7 @@ async def finance_research_async( :param research_effort: Controls how much time and effort the Finance Research API spends on your question. Higher effort levels run more searches and dig deeper into sources, at the cost of a longer response time. Available levels: + - `lite`: Returns answers quickly. Good for straightforward financial questions that just need a fast, reliable answer. - `deep`: The default. Spends more time researching and cross-referencing sources. Good for most financial questions, including multi-company comparisons, earnings analysis, and regulatory research. - `exhaustive`: The most thorough option. Explores the topic as fully as possible, best suited for complex financial research tasks where you want the highest quality result. :param retries: Override the default retry configuration for this method diff --git a/tests/test_research.py b/tests/test_research.py index 085d8a5..84a4919 100644 --- a/tests/test_research.py +++ b/tests/test_research.py @@ -216,6 +216,20 @@ def test_basic_finance_research(self, server_url, api_key): assert source.url is not None assert source.title is not None + def test_finance_research_lite_effort(self, server_url, api_key): + client = create_test_http_client("post_/v1/finance_research") + + with You(server_url=server_url, client=client, api_key_auth=api_key) as you: + res = you.finance_research( + input="What was Apple's revenue in FY2024?", + research_effort=FinanceResearchEffort.LITE, + server_url=server_url, + ) + + assert res.output is not None + assert res.output.content is not None + assert "effort: lite" in res.output.content + def test_finance_research_unauthorized(self, server_url): client = create_test_http_client("post_/v1/finance_research-unauthorized") diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py index ce592ed..f57e3b6 100644 --- a/tests/test_research_helpers.py +++ b/tests/test_research_helpers.py @@ -388,7 +388,7 @@ def test_research_and_wait_ok_event_but_get_non_completed_raises(self): api_key_auth="test-api-key", ) - with pytest.raises(RuntimeError, match="stream signalled completion but GET returned status=running"): + with pytest.raises(RuntimeError, match="stream signalled completion but GET returned status=running after 3 attempts"): research_and_wait( you, timeout_s=5.0, @@ -396,6 +396,56 @@ def test_research_and_wait_ok_event_but_get_non_completed_raises(self): research_effort=ResearchEffort.STANDARD, ) + def test_research_and_wait_ok_event_repoll_succeeds(self): + """When the stream emits a terminal OK event and the first GET returns + running (backend commit race), research_and_wait re-polls and returns + the completed detail once the status catches up.""" + call_count = {"get": 0} + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, + content=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ], + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, headers={"content-type": "application/json"}, content=_TASK_RESPONSE_JSON, + ) + # GET: first call returns running, second returns completed + call_count["get"] += 1 + if call_count["get"] == 1: + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_make_task_detail_json(status="running", result=None), + ) + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_make_task_detail_json(status="completed", result=_DEFAULT_RESULT), + ) + + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + client=httpx.Client(transport=transport), + api_key_auth="test-api-key", + ) + + detail = research_and_wait( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + assert call_count["get"] == 2 # first running, second completed + def test_research_and_wait_timeout_raises_timeout_error(self): """research_and_wait raises TimeoutError when the stream never sends a terminal event within timeout_s. The _BlockingStream simulates a @@ -1067,6 +1117,41 @@ async def test_async_research_and_wait_timeout_raises_timeout_error(self): research_effort=ResearchEffort.STANDARD, ) + @pytest.mark.asyncio + async def test_async_research_and_wait_httpx_readtimeout_falls_back_to_get(self): + """Async research_and_wait catches httpx.ReadTimeout (not just + asyncio.TimeoutError) and falls back to a final GET. Mirrors the + sync _BlockingStream test. Without catching httpx.TimeoutException, + the raw ReadTimeout would propagate uncaught.""" + class _ReadTimeoutAsyncStream(httpx.AsyncByteStream): + """Yields one event then raises httpx.ReadTimeout to simulate + a stalled server with an explicit read timeout.""" + async def __aiter__(self): + yield _CONNECTED_CHUNK + raise httpx.ReadTimeout("read timeout") + + handler = _make_wait_handler( + stream_obj=_ReadTimeoutAsyncStream(), + final_status="completed", + is_async=True, + ) + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + async_client=httpx.AsyncClient(transport=transport), + api_key_auth="test-api-key", + ) + + detail = await research_and_wait_async( + you, + timeout_s=0.5, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + @pytest.mark.asyncio async def test_async_research_and_wait_timeout_falls_back_to_get(self): """When the async stream times out but the task completed, the final From c03c18393ba6945e606b6652651306d3627e71a3 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 10:46:16 -0700 Subject: [PATCH 09/17] test: add live tests for frontier research + finance lite effort - test_frontier_background_completes: submits frontier task with background=true via research_and_wait, verifies TaskDetail completed. Passed live in 607s. - test_frontier_without_background_raises_422: verifies 422 when frontier is sent without background=true. Passed live. - test_finance_research_lite_effort: verifies LITE returns a quick answer. Skips gracefully when server hasn't deployed the tier yet (422). Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- tests/test_live.py | 72 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/tests/test_live.py b/tests/test_live.py index 11df41c..ebe9304 100644 --- a/tests/test_live.py +++ b/tests/test_live.py @@ -44,6 +44,11 @@ research_and_wait, stream_research, ) +from youdotcom.errors import ( + FinanceResearchUnprocessableEntityError, + ResearchUnprocessableEntityError, + YouDefaultError, +) # Skip all tests in this file if no API key is provided. @@ -415,6 +420,27 @@ def test_finance_research_basic(self, you_client): for source in res.output.sources: assert source.url is not None + def test_finance_research_lite_effort(self, you_client): + """Test finance_research with LITE effort returns a quick answer. + Skipped if the server hasn't deployed the lite tier yet.""" + with you_client as you: + try: + res = you.finance_research( + input="What was Apple's revenue in fiscal year 2024?", + research_effort=FinanceResearchEffort.LITE, + ) + except (FinanceResearchUnprocessableEntityError, YouDefaultError) as e: + if "lite" in str(e).lower() or "422" in str(e): + pytest.skip("Finance research lite tier not yet deployed on server") + raise + + assert res.output is not None + assert res.output.content is not None + assert len(res.output.content) > 0 + if res.output.sources: + for source in res.output.sources: + assert source.url is not None + class TestLiveContentsMaxAge: """Live test for Contents `max_age` parameter. @@ -636,6 +662,52 @@ def test_stream_research(self, you_client): assert events[0].data.get("status") is not None +# --------------------------------------------------------------------------- +# Frontier research effort (new in 2.5.0) +# --------------------------------------------------------------------------- +# Frontier only works with background=true and can run up to 4 hours. +# For a live test we use a simple query and a generous but bounded timeout. +# Marked slow so it can be skipped with `-m "not slow"`. +# --------------------------------------------------------------------------- +class TestLiveResearchFrontier: + """Live tests for frontier research effort (requires background=true).""" + + @pytest.mark.slow + def test_frontier_background_completes(self, you_client): + """research_and_wait with frontier effort auto-adjusts timeout to 4h. + Uses a simple query so it completes in a few minutes on prod.""" + with you_client as you: + try: + detail = research_and_wait( + you, + input="Who is Bill Gates?", + research_effort=ResearchEffort.FRONTIER, + timeout_s=600, # bounded for CI; auto would be 14400 + ) + except TypeError as e: + if "TaskResponse" in str(e): + pytest.skip("Background mode not enabled on server") + raise + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + assert detail.result is not None + payload = detail.result.model_dump() + content = payload.get("output", {}).get("content", "") + assert len(content) > 0 + + @pytest.mark.slow + def test_frontier_without_background_raises_422(self, you_client): + """frontier without background=true should return 422.""" + with you_client as you: + with pytest.raises((ResearchUnprocessableEntityError, YouDefaultError)): + you.research( + input="Who is Bill Gates?", + research_effort=ResearchEffort.FRONTIER, + background=False, + ) + + if __name__ == "__main__": # Run with: python -m pytest tests/test_live.py -v pytest.main([__file__, "-v"]) From 19464c027a8957643b6520de2c5df0cd423a6702 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 10:50:52 -0700 Subject: [PATCH 10/17] fix: bound async SSE per-read timeout to timeout_s (match sync variant) research_and_wait_async now passes stream_timeout_ms=int(timeout_s*1000) when the caller didn't set an explicit timeout_ms, matching the sync variant. Without this, a caller-set You(timeout_ms=60000) would cap SSE reads at 60s and cause premature TimeoutError on frontier tasks with auto timeout_s=14400. Also updated the async docstring to note the per-read timeout parity. Addresses PR #24 review comment 3640119913. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index f5e7bef..3dabebf 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -437,7 +437,10 @@ async def research_and_wait_async( """Async variant of :func:`research_and_wait`. See :func:`research_and_wait` for parameter details and auto-timeout - behavior (600s default, 14400s for frontier). + behavior (600s default, 14400s for frontier). The SSE per-read timeout + is bounded to ``timeout_s`` (when the caller didn't set ``timeout_ms``) + to match the sync variant, preventing a caller-set client default from + causing premature ``TimeoutError`` on long-running frontier tasks. """ if timeout_s is None: timeout_s = _resolve_default_timeout(kwargs) @@ -446,10 +449,16 @@ async def research_and_wait_async( http_headers = kwargs.get("http_headers") task = await research_background_async(client, **kwargs) + # Bound the SSE per-read timeout to timeout_s when the caller didn't + # set an explicit timeout_ms, matching the sync variant. Without this, + # a caller-set You(timeout_ms=60000) would cap reads at 60s and cause + # premature TimeoutError on frontier tasks with auto timeout_s=14400. + stream_timeout_ms = timeout_ms if timeout_ms is not None else int(timeout_s * 1000) + async def _consume() -> Optional[str]: async for evt in stream_research_async( client, task.task_id, - server_url=server_url, timeout_ms=timeout_ms, + server_url=server_url, timeout_ms=stream_timeout_ms, http_headers=http_headers, ): if evt.event in _TERMINAL_STREAM_EVENTS_OK: From 039fee1e1e3c6d87e4bdd87e2940a7d40e5dd305 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 10:55:37 -0700 Subject: [PATCH 11/17] fix: distinguish early-break from exhausted re-poll in _resolve_from_final_get When the re-poll loop hits a terminal non-completed status (failed/ cancelled) on the first GET, raise immediately with a clear message instead of falling through to the misleading after-N-attempts error. Added test_research_and_wait_ok_event_but_get_failed_raises_immediately to cover the early-break branch. Addresses PR #24 review comments 3640244682 and 3640245293. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 10 ++++++++-- tests/test_research_helpers.py | 27 +++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index 3dabebf..692c9bd 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -256,7 +256,10 @@ def _resolve_from_final_get( if detail.status.value == "completed": return detail if detail.status.value in _TERMINAL_TASK_STATUSES: - break # terminal non-completed, fall through to error below + raise RuntimeError( + f"research task {task_id} ended in non-completed state: " + f"{detail.status.value}" + ) time.sleep(_REPOLL_INTERVAL_S) raise RuntimeError( f"research task {task_id} stream signalled completion " @@ -315,7 +318,10 @@ async def _resolve_from_final_get_async( if detail.status.value == "completed": return detail if detail.status.value in _TERMINAL_TASK_STATUSES: - break + raise RuntimeError( + f"research task {task_id} ended in non-completed state: " + f"{detail.status.value}" + ) await asyncio.sleep(_REPOLL_INTERVAL_S) raise RuntimeError( f"research task {task_id} stream signalled completion " diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py index f57e3b6..bdc984f 100644 --- a/tests/test_research_helpers.py +++ b/tests/test_research_helpers.py @@ -396,6 +396,33 @@ def test_research_and_wait_ok_event_but_get_non_completed_raises(self): research_effort=ResearchEffort.STANDARD, ) + def test_research_and_wait_ok_event_but_get_failed_raises_immediately(self): + """When the stream emits a terminal OK event but the follow-up GET + returns a terminal non-completed status (failed), research_and_wait + raises RuntimeError immediately without exhausting re-poll attempts.""" + handler = _make_wait_handler( + stream_chunks=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ], + final_status="failed", + final_result=None, + ) + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + client=httpx.Client(transport=transport), + api_key_auth="test-api-key", + ) + + with pytest.raises(RuntimeError, match="ended in non-completed state: failed"): + research_and_wait( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + def test_research_and_wait_ok_event_repoll_succeeds(self): """When the stream emits a terminal OK event and the first GET returns running (backend commit race), research_and_wait re-polls and returns From 714d9b875b88726578c981611c19d0a51d847801 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 10:58:47 -0700 Subject: [PATCH 12/17] fix: align sync/async stream timeout_ms handling + add async repoll tests Sync research_and_wait now honors caller timeout_ms for the stream read when provided (matching async), falling back to timeout_s*1000 when not. Updated both docstrings to accurately describe the shared behavior. Added async mirrors for repoll-succeeds and ok-event-but-get-failed tests, closing the coverage gap flagged in review. Addresses PR #24 review comments 3640262585 and 3640262855 (the latter was already fixed in 039fee1 but commented on an older revision). Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 17 +++---- tests/test_research_helpers.py | 76 +++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+), 8 deletions(-) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index 692c9bd..02fcd2f 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -369,9 +369,10 @@ def research_and_wait( ``TaskDetail`` via a ``get_research_task`` call (a second GET is issued in timeout/stream-close fallback paths). - The stream is opened with a read timeout of ``timeout_s`` so that a - stalled server raises ``httpx.ReadTimeout`` deterministically instead - of leaking a blocked consumer thread. + The stream is opened with a per-read timeout bounded to ``timeout_s`` + (or the caller's ``timeout_ms`` when provided) so that a stalled server + raises ``httpx.ReadTimeout`` deterministically instead of leaking a + blocked consumer thread. A total wall-clock deadline is also enforced. Parameters: timeout_s: Maximum seconds to wait for a terminal stream event. @@ -400,10 +401,11 @@ def research_and_wait( http_headers = kwargs.get("http_headers") task = research_background(client, **kwargs) - # Open the stream with a per-read timeout so a stalled server raises + # Bound the SSE per-read timeout: use the caller's timeout_ms when + # provided, otherwise default to timeout_s so a stalled server raises # httpx.ReadTimeout deterministically. We also enforce a total wall-clock # deadline below to match the async variant's asyncio.wait_for semantics. - stream_timeout_ms = int(timeout_s * 1000) + stream_timeout_ms = timeout_ms if timeout_ms is not None else int(timeout_s * 1000) stream = _open_raw_stream( client, task.task_id, http_headers=http_headers, server_url=server_url, timeout_ms=stream_timeout_ms, @@ -444,9 +446,8 @@ async def research_and_wait_async( See :func:`research_and_wait` for parameter details and auto-timeout behavior (600s default, 14400s for frontier). The SSE per-read timeout - is bounded to ``timeout_s`` (when the caller didn't set ``timeout_ms``) - to match the sync variant, preventing a caller-set client default from - causing premature ``TimeoutError`` on long-running frontier tasks. + is bounded to ``timeout_s`` (or the caller's ``timeout_ms`` when + provided), matching the sync variant. """ if timeout_s is None: timeout_s = _resolve_default_timeout(kwargs) diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py index bdc984f..383c848 100644 --- a/tests/test_research_helpers.py +++ b/tests/test_research_helpers.py @@ -1118,6 +1118,82 @@ async def test_async_research_and_wait_error_event_raises_runtime_error(self): research_effort=ResearchEffort.STANDARD, ) + @pytest.mark.asyncio + async def test_async_research_and_wait_ok_event_repoll_succeeds(self): + """Async mirror: stream OK + first GET running, second GET completed.""" + call_count = {"get": 0} + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, + stream=_AsyncChunks([ + _CONNECTED_CHUNK, + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ]), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, headers={"content-type": "application/json"}, content=_TASK_RESPONSE_JSON, + ) + call_count["get"] += 1 + if call_count["get"] == 1: + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_make_task_detail_json(status="running", result=None), + ) + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_make_task_detail_json(status="completed", result=_DEFAULT_RESULT), + ) + + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + async_client=httpx.AsyncClient(transport=transport), + api_key_auth="test-api-key", + ) + + detail = await research_and_wait_async( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + assert call_count["get"] == 2 + + @pytest.mark.asyncio + async def test_async_research_and_wait_ok_event_but_get_failed_raises_immediately(self): + """Async mirror: stream OK + GET returns terminal failed raises + immediately without exhausting re-poll attempts.""" + handler = _make_wait_handler( + stream_chunks=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ], + final_status="failed", + final_result=None, + is_async=True, + ) + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + async_client=httpx.AsyncClient(transport=transport), + api_key_auth="test-api-key", + ) + + with pytest.raises(RuntimeError, match="ended in non-completed state: failed"): + await research_and_wait_async( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + @pytest.mark.asyncio async def test_async_research_and_wait_timeout_raises_timeout_error(self): """Async research_and_wait raises TimeoutError when the stream never From 7f7a0bdf10a51cefcf67a7d524dd145156b38b01 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 11:08:52 -0700 Subject: [PATCH 13/17] fix: stream-open failure falls back to polling + cap per-read timeout at timeout_s 1. research_and_wait[_async] now catches stream-open failures (transient 500, 404 on just-created task, etc.) and falls back to poll_research_task so a successfully-submitted task is never abandoned. 2. SSE per-read timeout is now capped at timeout_s*1000 in both variants, preventing a caller-set timeout_ms larger than timeout_s from causing the sync wall-clock to exceed the budget. 3. Added test_async_frontier_auto_timeout_completes_without_premature_timeout to mirror the sync frontier auto-timeout coverage. Addresses PR #24 review comments 3640308132, 3640308929, 3640309427. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 51 +++++++++++++++++++++++-------- tests/test_research_helpers.py | 29 ++++++++++++++++++ 2 files changed, 67 insertions(+), 13 deletions(-) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index 02fcd2f..edc1751 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -401,15 +401,29 @@ def research_and_wait( http_headers = kwargs.get("http_headers") task = research_background(client, **kwargs) - # Bound the SSE per-read timeout: use the caller's timeout_ms when - # provided, otherwise default to timeout_s so a stalled server raises - # httpx.ReadTimeout deterministically. We also enforce a total wall-clock - # deadline below to match the async variant's asyncio.wait_for semantics. - stream_timeout_ms = timeout_ms if timeout_ms is not None else int(timeout_s * 1000) - stream = _open_raw_stream( - client, task.task_id, http_headers=http_headers, - server_url=server_url, timeout_ms=stream_timeout_ms, + # Cap the SSE per-read timeout at timeout_s so a silent stall can never + # exceed the total wait budget, even if the caller's timeout_ms is larger. + # When timeout_ms is smaller, honor it (tighter per-read bound). + stream_timeout_ms = min( + timeout_ms if timeout_ms is not None else int(timeout_s * 1000), + int(timeout_s * 1000), ) + try: + stream = _open_raw_stream( + client, task.task_id, http_headers=http_headers, + server_url=server_url, timeout_ms=stream_timeout_ms, + ) + except Exception: + # Stream-open failure (transient 500, 404 on a just-created task, etc.): + # the task is already submitted and running. Fall back to polling so + # the caller gets a result instead of an unrecoverable exception. + return poll_research_task( + client, task.task_id, + interval_s=_DEFAULT_POLL_INTERVAL_S, + timeout_s=timeout_s, + server_url=server_url, timeout_ms=timeout_ms, + http_headers=http_headers, + ) result: Optional[str] = None timed_out = False deadline = time.monotonic() + timeout_s @@ -456,11 +470,11 @@ async def research_and_wait_async( http_headers = kwargs.get("http_headers") task = await research_background_async(client, **kwargs) - # Bound the SSE per-read timeout to timeout_s when the caller didn't - # set an explicit timeout_ms, matching the sync variant. Without this, - # a caller-set You(timeout_ms=60000) would cap reads at 60s and cause - # premature TimeoutError on frontier tasks with auto timeout_s=14400. - stream_timeout_ms = timeout_ms if timeout_ms is not None else int(timeout_s * 1000) + # Cap the SSE per-read timeout at timeout_s (same rationale as sync). + stream_timeout_ms = min( + timeout_ms if timeout_ms is not None else int(timeout_s * 1000), + int(timeout_s * 1000), + ) async def _consume() -> Optional[str]: async for evt in stream_research_async( @@ -480,6 +494,17 @@ async def _consume() -> Optional[str]: except (asyncio.TimeoutError, httpx.TimeoutException): timed_out = True result = None + except Exception: + # Stream-open failure (transient 500, 404 on a just-created task, etc.): + # the task is already submitted and running. Fall back to polling so + # the caller gets a result instead of an unrecoverable exception. + return await poll_research_task_async( + client, task.task_id, + interval_s=_DEFAULT_POLL_INTERVAL_S, + timeout_s=timeout_s, + server_url=server_url, timeout_ms=timeout_ms, + http_headers=http_headers, + ) return await _resolve_from_final_get_async( client, task.task_id, result, timeout_s, diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py index 383c848..ecd7db0 100644 --- a/tests/test_research_helpers.py +++ b/tests/test_research_helpers.py @@ -1091,6 +1091,35 @@ async def test_async_research_and_wait_returns_completed_detail(self): assert isinstance(detail, TaskDetail) assert detail.status.value == "completed" + @pytest.mark.asyncio + async def test_async_frontier_auto_timeout_completes_without_premature_timeout(self): + """Async mirror: when research_effort=frontier and timeout_s is omitted, + the auto-adjusted 4-hour timeout should not trip on a normal event + sequence that completes within seconds.""" + handler = _make_wait_handler( + stream_chunks=[ + _CONNECTED_CHUNK, + b'id: 1\nevent: response.done\ndata: {"type":"response.done","task_id":"abc","status":"completed","sequence":1}\n\n', + ], + is_async=True, + ) + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + async_client=httpx.AsyncClient(transport=transport), + api_key_auth="test-api-key", + ) + + detail = await research_and_wait_async( + you, + input="Evaluate the Gates Foundation's global-health impact", + research_effort=ResearchEffort.FRONTIER, + # timeout_s intentionally omitted — should auto-adjust to 14400 + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + @pytest.mark.asyncio async def test_async_research_and_wait_error_event_raises_runtime_error(self): """Async research_and_wait raises RuntimeError on error terminal event.""" From ee28a4c1ddbe999f9ea3cd9a7b74056d58d9e0a5 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 11:17:19 -0700 Subject: [PATCH 14/17] fix: sync research_and_wait falls back to polling on mid-stream errors The sync variant previously only caught httpx.TimeoutException during stream iteration, letting non-timeout errors (RemoteProtocolError, ConnectError on dropped connections) propagate unrecoverably. Now catches Exception and falls back to poll_research_task, matching the async variant's behavior and the stream-open failure fallback. Addresses PR #24 review comment 3640374576. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index edc1751..00dfcff 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -440,6 +440,17 @@ def research_and_wait( break except httpx.TimeoutException: timed_out = True + except Exception: + # Mid-stream error (e.g. RemoteProtocolError on a dropped connection): + # the task is already submitted and running. Fall back to polling, + # matching the async variant's behavior. + return poll_research_task( + client, task.task_id, + interval_s=_DEFAULT_POLL_INTERVAL_S, + timeout_s=timeout_s, + server_url=server_url, timeout_ms=timeout_ms, + http_headers=http_headers, + ) finally: stream.close() From 330028d177878cb8747883723e6649f6eed9690d Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 11:30:59 -0700 Subject: [PATCH 15/17] fix: narrow fallback to httpx.TransportError + decouple SSE per-read from timeout_ms - Change `except Exception` to `except httpx.TransportError` for both stream-open and mid-stream fallbacks (sync + async). Typed auth errors (401/403/404/500 from _raise_stream_error) now propagate to the caller instead of being swallowed and masked by a different error type from poll_research_task. - Decouple SSE per-read timeout from caller's timeout_ms. Always use int(timeout_s * 1000) for stream reads. Previously, a caller with timeout_ms=30000 on a frontier task (timeout_s=14400) would get a misleading "did not complete within 14400s" after only 30s. - Add 5 new tests: stream-open TransportError -> poll succeeds (sync + async), stream-open 401 -> typed error propagates (sync + async), mid-stream TransportError -> poll succeeds (sync). Addresses PR #24 review comments 3640425979 and 3640426348. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 48 ++++---- tests/test_research_helpers.py | 188 ++++++++++++++++++++++++++++++ 2 files changed, 210 insertions(+), 26 deletions(-) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index 00dfcff..356f341 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -401,22 +401,20 @@ def research_and_wait( http_headers = kwargs.get("http_headers") task = research_background(client, **kwargs) - # Cap the SSE per-read timeout at timeout_s so a silent stall can never - # exceed the total wait budget, even if the caller's timeout_ms is larger. - # When timeout_ms is smaller, honor it (tighter per-read bound). - stream_timeout_ms = min( - timeout_ms if timeout_ms is not None else int(timeout_s * 1000), - int(timeout_s * 1000), - ) + # The SSE per-read timeout is bounded to timeout_s so a silent stall + # can never exceed the total wait budget. This is decoupled from the + # caller's timeout_ms (which governs individual HTTP requests like the + # POST and GET, not SSE event gaps). + stream_timeout_ms = int(timeout_s * 1000) try: stream = _open_raw_stream( client, task.task_id, http_headers=http_headers, server_url=server_url, timeout_ms=stream_timeout_ms, ) - except Exception: - # Stream-open failure (transient 500, 404 on a just-created task, etc.): - # the task is already submitted and running. Fall back to polling so - # the caller gets a result instead of an unrecoverable exception. + except httpx.TransportError: + # Transient connection failure (can't reach server, etc.): the task + # is already submitted. Fall back to polling. Typed errors (401/403/ + # 404/500 from _raise_stream_error) propagate to the caller. return poll_research_task( client, task.task_id, interval_s=_DEFAULT_POLL_INTERVAL_S, @@ -440,10 +438,10 @@ def research_and_wait( break except httpx.TimeoutException: timed_out = True - except Exception: - # Mid-stream error (e.g. RemoteProtocolError on a dropped connection): - # the task is already submitted and running. Fall back to polling, - # matching the async variant's behavior. + except httpx.TransportError: + # Mid-stream network failure (e.g. RemoteProtocolError on a dropped + # connection): the task is already submitted and running. Fall back + # to polling. Typed SDK errors propagate to the caller. return poll_research_task( client, task.task_id, interval_s=_DEFAULT_POLL_INTERVAL_S, @@ -471,8 +469,7 @@ async def research_and_wait_async( See :func:`research_and_wait` for parameter details and auto-timeout behavior (600s default, 14400s for frontier). The SSE per-read timeout - is bounded to ``timeout_s`` (or the caller's ``timeout_ms`` when - provided), matching the sync variant. + is bounded to ``timeout_s``, matching the sync variant. """ if timeout_s is None: timeout_s = _resolve_default_timeout(kwargs) @@ -481,11 +478,10 @@ async def research_and_wait_async( http_headers = kwargs.get("http_headers") task = await research_background_async(client, **kwargs) - # Cap the SSE per-read timeout at timeout_s (same rationale as sync). - stream_timeout_ms = min( - timeout_ms if timeout_ms is not None else int(timeout_s * 1000), - int(timeout_s * 1000), - ) + # The SSE per-read timeout is bounded to timeout_s (same rationale as + # sync). Decoupled from the caller's timeout_ms, which governs + # individual HTTP requests (POST, GET), not SSE event gaps. + stream_timeout_ms = int(timeout_s * 1000) async def _consume() -> Optional[str]: async for evt in stream_research_async( @@ -505,10 +501,10 @@ async def _consume() -> Optional[str]: except (asyncio.TimeoutError, httpx.TimeoutException): timed_out = True result = None - except Exception: - # Stream-open failure (transient 500, 404 on a just-created task, etc.): - # the task is already submitted and running. Fall back to polling so - # the caller gets a result instead of an unrecoverable exception. + except httpx.TransportError: + # Transient connection failure (can't reach server, dropped + # connection, etc.): the task is already submitted. Fall back to + # polling. Typed errors (401/403/404/500) propagate to the caller. return await poll_research_task_async( client, task.task_id, interval_s=_DEFAULT_POLL_INTERVAL_S, diff --git a/tests/test_research_helpers.py b/tests/test_research_helpers.py index ecd7db0..9a9a7e1 100644 --- a/tests/test_research_helpers.py +++ b/tests/test_research_helpers.py @@ -601,6 +601,122 @@ def __iter__(self): research_effort=ResearchEffort.STANDARD, ) + def test_stream_open_transport_error_falls_back_to_poll(self): + """When _open_raw_stream raises a TransportError (can't reach server), + research_and_wait falls back to poll_research_task and returns the + completed detail.""" + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + raise httpx.ConnectError("connection refused") + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_TASK_RESPONSE_JSON, + ) + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_make_task_detail_json(status="completed", result=_DEFAULT_RESULT), + ) + + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + client=httpx.Client(transport=transport), + api_key_auth="test-api-key", + ) + + detail = research_and_wait( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + def test_stream_open_401_propagates_typed_error(self): + """When _open_raw_stream gets a 401, the typed + StreamResearchTaskUnauthorizedError propagates instead of falling + back to polling.""" + from youdotcom.errors import StreamResearchTaskUnauthorizedError + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + return httpx.Response( + 401, + headers={"content-type": "application/json"}, + content='{"error": "unauthorized"}', + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_TASK_RESPONSE_JSON, + ) + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_make_task_detail_json(status="completed", result=_DEFAULT_RESULT), + ) + + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + client=httpx.Client(transport=transport), + api_key_auth="bad-key", + ) + + with pytest.raises(StreamResearchTaskUnauthorizedError): + research_and_wait( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + def test_mid_stream_transport_error_falls_back_to_poll(self): + """When a TransportError occurs mid-stream (dropped connection), + research_and_wait falls back to poll_research_task.""" + class _DroppedStream(httpx.SyncByteStream): + def __iter__(self): + yield _CONNECTED_CHUNK + raise httpx.RemoteProtocolError("connection dropped") + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, + stream=_DroppedStream(), + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_TASK_RESPONSE_JSON, + ) + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_make_task_detail_json(status="completed", result=_DEFAULT_RESULT), + ) + + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + client=httpx.Client(transport=transport), + api_key_auth="test-api-key", + ) + + detail = research_and_wait( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + # --------------------------------------------------------------------------- # Stream research tolerant decoder: verify that @@ -1360,3 +1476,75 @@ async def test_async_research_and_wait_stream_close_task_running_raises_timeout( input="test query", research_effort=ResearchEffort.STANDARD, ) + + @pytest.mark.asyncio + async def test_async_stream_open_transport_error_falls_back_to_poll(self): + """Async mirror: TransportError on stream-open falls back to polling.""" + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + raise httpx.ConnectError("connection refused") + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_TASK_RESPONSE_JSON, + ) + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_make_task_detail_json(status="completed", result=_DEFAULT_RESULT), + ) + + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + async_client=httpx.AsyncClient(transport=transport), + api_key_auth="test-api-key", + ) + + detail = await research_and_wait_async( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) + + assert isinstance(detail, TaskDetail) + assert detail.status.value == "completed" + + @pytest.mark.asyncio + async def test_async_stream_open_401_propagates_typed_error(self): + """Async mirror: 401 on stream-open propagates typed error.""" + from youdotcom.errors import StreamResearchTaskUnauthorizedError + + def handler(request): + url = str(request.url) + if url.endswith("/stream") or "/stream?" in url: + return httpx.Response( + 401, + headers={"content-type": "application/json"}, + content='{"error": "unauthorized"}', + ) + if request.method == "POST" and "/v1/research" in url and "/stream" not in url: + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_TASK_RESPONSE_JSON, + ) + return httpx.Response( + 200, headers={"content-type": "application/json"}, + content=_make_task_detail_json(status="completed", result=_DEFAULT_RESULT), + ) + + transport = httpx.MockTransport(handler) + you = You( + server_url="http://mock.local", + async_client=httpx.AsyncClient(transport=transport), + api_key_auth="bad-key", + ) + + with pytest.raises(StreamResearchTaskUnauthorizedError): + await research_and_wait_async( + you, + timeout_s=5.0, + input="test query", + research_effort=ResearchEffort.STANDARD, + ) From c1ef2e1a81c1726bfea89b5defe1c30e5cc8257d Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 11:40:55 -0700 Subject: [PATCH 16/17] fix: close stream before fallback poll in mid-stream TransportError handler The `return poll_research_task(...)` inside the `try` block meant `finally: stream.close()` only ran after the entire poll loop completed, keeping the failed SSE connection open during polling. Now calling `stream.close()` explicitly before the return so cleanup happens first. Addresses PR #24 review comment 3640517073. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index 356f341..43416bc 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -442,6 +442,7 @@ def research_and_wait( # Mid-stream network failure (e.g. RemoteProtocolError on a dropped # connection): the task is already submitted and running. Fall back # to polling. Typed SDK errors propagate to the caller. + stream.close() return poll_research_task( client, task.task_id, interval_s=_DEFAULT_POLL_INTERVAL_S, From 9a363a1a398d4edd8ad2c20803dd8de2f54b7e91 Mon Sep 17 00:00:00 2001 From: Tyler Eastman Date: Thu, 23 Jul 2026 11:48:57 -0700 Subject: [PATCH 17/17] docs: fix sync docstring to drop stale timeout_ms parenthetical The SSE per-read timeout was decoupled from timeout_ms in a previous commit but the sync docstring still mentioned "(or the caller's timeout_ms when provided)". Now matches the actual behavior and the async variant's docstring. Addresses PR #24 review comment 3640569229. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/youdotcom/research_helpers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/youdotcom/research_helpers.py b/src/youdotcom/research_helpers.py index 43416bc..a4b0078 100644 --- a/src/youdotcom/research_helpers.py +++ b/src/youdotcom/research_helpers.py @@ -370,9 +370,9 @@ def research_and_wait( issued in timeout/stream-close fallback paths). The stream is opened with a per-read timeout bounded to ``timeout_s`` - (or the caller's ``timeout_ms`` when provided) so that a stalled server - raises ``httpx.ReadTimeout`` deterministically instead of leaking a - blocked consumer thread. A total wall-clock deadline is also enforced. + so that a stalled server raises ``httpx.ReadTimeout`` deterministically + instead of leaking a blocked consumer thread. A total wall-clock deadline + is also enforced. Parameters: timeout_s: Maximum seconds to wait for a terminal stream event.