From 75da34c16d3a6c9f6fc5c636741af9b378ff0528 Mon Sep 17 00:00:00 2001 From: Sean Brar Date: Sat, 11 Jul 2026 19:34:56 -0700 Subject: [PATCH] feat(core): harden v2 interaction contracts --- docs/agent-loop.md | 38 +- docs/conversations-and-agents.md | 64 +-- docs/migrating-to-v2.md | 16 +- docs/reference/api.md | 2 + docs/reference/provider-capabilities.md | 3 + src/pollux/__init__.py | 33 +- src/pollux/_lifecycle.py | 43 ++ src/pollux/interaction/__init__.py | 5 +- src/pollux/interaction/continuation.py | 439 ++++++++++++------ src/pollux/interaction/environment.py | 100 ++-- src/pollux/interaction/execute.py | 7 +- src/pollux/interaction/validate.py | 8 +- src/pollux/providers/_compile.py | 17 +- src/pollux/providers/anthropic.py | 5 + src/pollux/providers/gemini.py | 5 + src/pollux/providers/openai.py | 5 + src/pollux/source.py | 16 + tests/helpers.py | 35 +- tests/interaction/test_continuation.py | 239 +++++----- tests/interaction/test_continuation_compat.py | 19 +- tests/interaction/test_environment.py | 136 ++++-- tests/interaction/test_execute_interaction.py | 3 +- tests/interaction/test_input.py | 7 +- tests/interaction/test_interact_frontdoor.py | 75 +++ tests/interaction/test_stream.py | 105 +++++ tests/providers/test_anthropic_contract.py | 7 +- tests/providers/test_anthropic_stream.py | 9 +- tests/providers/test_gemini_stream.py | 9 +- tests/providers/test_local_contract.py | 15 +- tests/providers/test_local_stream.py | 7 +- tests/providers/test_openai_contract.py | 7 +- tests/providers/test_openai_stream.py | 9 +- tests/providers/test_openrouter_contract.py | 16 +- tests/providers/test_openrouter_stream.py | 7 +- tests/providers/test_tool_history.py | 19 +- tests/test_api.py | 4 +- 36 files changed, 1064 insertions(+), 470 deletions(-) create mode 100644 src/pollux/_lifecycle.py diff --git a/docs/agent-loop.md b/docs/agent-loop.md index f73d37c2..4b89aeb3 100644 --- a/docs/agent-loop.md +++ b/docs/agent-loop.md @@ -183,6 +183,25 @@ Do not call `asyncio.run(...)` from inside an already-running event loop runtime). In those environments, keep your Pollux adapter async and `await` `interact()`, `run()`, or `stream()` directly. +### Closing streams early + +Cancellation of a task awaiting `interact()` or a stream iteration propagates +through Pollux and releases the active provider request. If a consumer may stop +streaming before the terminal `done` event, explicitly close the iterator; a +bare `break` does not guarantee synchronous async-generator cleanup: + +```python +from contextlib import aclosing + +async with aclosing(stream(env, input, config=config)) as events: + async for event in events: + if should_stop(event): + break +``` + +Closing a `Session.stream()` releases only that request. The session remains +open and reusable until its own context manager exits or `aclose()` is called. + ## Variations These are all small modifications to the same loop structure. @@ -234,27 +253,22 @@ out = await interact( ``` If your application owns an OpenAI Chat Completions-style transcript for resume, -compaction, or audit logs, keep that transcript as the durable record and import -it into a fresh Pollux continuation for each turn: +compaction, or audit logs, convert its portable messages into typed history: ```python -from pollux import Continuation, Input +from pollux import Input, Message -continuation = Continuation.from_openai_messages(messages, provider="local") +history = [Message.from_openai(message) for message in messages] out = await interact( env, - Input(content="Continue.", continuation=continuation), + Input(content="Continue.", history=history), config=config, ) ``` -`ToolCall.to_openai()`, `Message.to_openai()`, and -`Continuation.to_openai_messages()` provide the reverse mapping for harnesses -that still dispatch OpenAI-shaped tool calls. For supported text and tool -message shapes, importing with `Continuation.from_openai_messages(...)` and -exporting with `to_openai_messages()` preserves `role`, `content`, -`tool_calls`, and `tool_call_id`. Display reasoning is output data and is not -replayed as transcript history. +`ToolCall.to_openai()` and `Message.to_openai()` provide the reverse mapping for +harnesses that dispatch OpenAI-shaped tool calls. System instructions belong on +`Environment`, and display reasoning is not replayed as transcript history. ### Guiding tool use with system instructions diff --git a/docs/conversations-and-agents.md b/docs/conversations-and-agents.md index 51f6ad6e..27f67eb8 100644 --- a/docs/conversations-and-agents.md +++ b/docs/conversations-and-agents.md @@ -32,7 +32,11 @@ next turn. ## Continuing a Conversation with `Continuation` -Pass a prior result's `continuation` back into the next `Input(continuation=...)` to automatically resume a conversation. Pollux unpacks the initial prompt, the assistant's previous response, and any tool calls directly into the context payload. +Pass a prior result's `continuation` back into the next +`Input(continuation=...)` to resume the provider-correct conversation. Treat the +value as opaque: serialize it with `to_jsonable()`, restore it with +`Continuation.from_jsonable()`, and otherwise pass it back unchanged. Do not +inspect, edit, merge, or summarize its serialized provider state. To get a `continuation` for subsequent turns in plain conversational calls, the first turn must opt into conversation tracking by passing `history=[]` (or an empty list/tuple). Without it, Pollux treats the call as stateless and does not build continuation state. @@ -74,8 +78,7 @@ If you need to inject mid-conversation context, groom old context out to save tokens, or resume a chat from a database, a prior `continuation` alone is not enough. -Instead, pass an explicit `history` list of dictionaries containing `role` -and `content`: +Instead, pass explicit typed `Message` history: ```python import asyncio @@ -104,7 +107,9 @@ asyncio.run(manual_history_injection()) ``` Pollux treats the `history` block chronologically *before* the prompt you -provide to `interact()`. +provide to `interact()`. This deliberately gives up response IDs and opaque +provider replay state. A successful interaction creates a fresh continuation +for the active provider. ## Persisting Agent Transcripts @@ -123,42 +128,43 @@ product transcript: If your application is the source of truth for history because it supports resume, compaction, truncation, or audit logs, keep that transcript in your own -store and rebuild a `Continuation` for each Pollux turn. In this pattern, -`Continuation` is the provider replay object for the next call, not the durable -application transcript. - -`Continuation.from_openai_messages(...)` is a compatibility bridge for text -Chat Completions-style transcripts and tool turns. It extracts text from -OpenAI text parts, preserves tool calls, and deliberately does not turn -provider-shaped media attachments into Pollux `Source` objects. That keeps -media explicit at the Pollux boundary: +store and rebuild typed `Message` history for each Pollux turn: ```python -from pollux import Continuation, Environment, Input, Source, interact +from pollux import Environment, Input, Message, Source, interact -continuation = Continuation.from_openai_messages(text_messages, provider="local") +history = [Message.from_openai(message) for message in text_messages] env = Environment(sources=[Source.from_file("current-report.pdf")]) result = await interact( env, - Input(content="Continue the analysis.", continuation=continuation), + Input(content="Continue the analysis.", history=history), config=config, ) ``` -For supported text and tool message shapes, -`Continuation.from_openai_messages(messages).to_openai_messages()` is a -lossless round trip for `role`, `content`, `tool_calls`, and `tool_call_id`. -Display reasoning is not part of that replay contract: `Output.reasoning` is -diagnostic/output data and should not be copied into future transcript messages. -Provider-specific opaque reasoning blocks are replayed through -`provider_state` only when a provider requires them. +`Message.from_openai()` and `Message.to_openai()` bridge individual portable +text/tool messages. Move OpenAI `system` messages to `Environment.instructions`. +Display reasoning is diagnostic output and should not be copied into transcript +messages. If you need to resume an older session that included files or images, load those application records and rebuild `Source.from_file(...)`, `Source.from_uri(...)`, or another source constructor explicitly. Pollux's replay messages are for conversation turns, not hidden media transport. +### Portable History Shapes + +Portable history consists of non-empty user text, assistant text and/or +normalized `ToolCall` values, and tool messages with the matching +`tool_call_id`. Keep parallel tool results in call order. System messages, +media, reasoning, response IDs, and provider-native state are not portable. + +Anthropic extended-thinking tool turns are a special boundary: signed thinking +blocks live only in the untouched continuation. Return pending tool results with +that continuation, then compact at a completed interaction boundary. Groomed +history cannot recreate the signatures Anthropic requires. + ## Handling Tool Messages in History If your conversation includes tool execution (the model asked for data, you @@ -256,14 +262,20 @@ next_result = await interact( - **Reasoning is output data unless a provider requires opaque replay.** Pollux surfaces reasoning text on `Output.reasoning` for display and debugging. Provider-specific signed thinking or reasoning blocks are kept - inside continuation `provider_state` only when the provider requires them - for valid follow-up turns; do not copy display reasoning into future user or - assistant messages yourself. + inside the opaque continuation only when the provider requires them for valid + follow-up turns; do not copy display reasoning into future messages. - **Provider differences exist.** Gemini, OpenAI, and Anthropic support tool calling and tool messages in history. OpenRouter supports them on models that advertise tool support. See [Provider Capabilities](reference/provider-capabilities.md) for details. +## Durable Environment Identity + +Use `environment.fingerprint(provider=config.provider)` to bind saved work to +the model-facing instructions, sources, tools, and provider. Compose +`config.model` plus application policy/schema identifiers separately. Cache +preferences and environment metadata intentionally do not affect this identity. + --- Now that you understand the conversation mechanics, see diff --git a/docs/migrating-to-v2.md b/docs/migrating-to-v2.md index 22946a02..24acf653 100644 --- a/docs/migrating-to-v2.md +++ b/docs/migrating-to-v2.md @@ -104,11 +104,17 @@ produced them, and `from_jsonable()` rejects an incompatible version (and, when pass `expected_provider=`, a mismatched provider) with an actionable error instead of misreading it. -A continuation is bound to its provider: its `provider_state` (response ids, -provider-specific replay blocks) is not portable, so reusing one under a different -provider is rejected before dispatch. Across the 1.x → 2.0 boundary, plan to re-run -work rather than reusing old serialized blobs; persist enough application state to -rebuild the request when that is the right recovery path. +A continuation is opaque and bound to its provider. Applications serialize, +restore, and pass it back unchanged; provider response IDs and replay blocks are +not editable or portable. The final v2 RC uses continuation schema version 2 and +rejects schema-v1 RC artifacts. `Continuation.from_openai_messages()` and +`to_openai_messages()` were removed; use typed `Message` history after grooming +or importing an application transcript. + +For durable run identity, use +`environment.fingerprint(provider=config.provider)` and compose `config.model` +plus application policy/schema versions separately. `EnvironmentSnapshot` is +now internal and is no longer a supported identity route. ## What To Do In 1.x diff --git a/docs/reference/api.md b/docs/reference/api.md index feb9fec3..cd046e30 100644 --- a/docs/reference/api.md +++ b/docs/reference/api.md @@ -88,6 +88,8 @@ and `ResultEnvelope` types are no longer part of the public API. ::: pollux.Continuation +::: pollux.Message + ::: pollux.ToolDeclaration ::: pollux.ToolCall diff --git a/docs/reference/provider-capabilities.md b/docs/reference/provider-capabilities.md index 4970c28f..752efe04 100644 --- a/docs/reference/provider-capabilities.md +++ b/docs/reference/provider-capabilities.md @@ -144,6 +144,9 @@ jobs, see [Building With Deferred Delivery](../building-with-deferred-delivery.m holds for `stream()` too: signed thinking blocks are reassembled from the stream, so a streamed extended-thinking + tool turn continues identically to the non-streaming path. +- Do not compact between an extended-thinking tool call and its tool results. + The signed blocks exist only in the opaque continuation; complete the tool + exchange first, then switch to application-authored `Message` history. - `max_tokens`: limits the output length. Default is `16384` for Anthropic, which leaves room for thinking output at all effort levels. Other providers currently ignore this option. diff --git a/src/pollux/__init__.py b/src/pollux/__init__.py index 13c94c49..a1bfd38f 100644 --- a/src/pollux/__init__.py +++ b/src/pollux/__init__.py @@ -20,6 +20,7 @@ import logging from typing import TYPE_CHECKING, Any, cast +from pollux._lifecycle import close_async_iterator from pollux.config import ( _API_KEY_ENV_VARS, _LOCAL_BASE_URL_ENV_VAR, @@ -53,10 +54,10 @@ CacheSetting, Continuation, Environment, - EnvironmentSnapshot, Event, Input, Message, + MessageRole, Output, OutputCollection, OutputRequirements, @@ -66,6 +67,7 @@ ToolDeclaration, ToolResult, ) +from pollux.interaction.environment import EnvironmentSnapshot as _EnvironmentSnapshot from pollux.interaction.execute import ( execute_interaction, execute_interactions, @@ -81,7 +83,7 @@ from pollux.source import Source if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable, Sequence + from collections.abc import AsyncGenerator, Callable, Sequence from pollux.interaction.schema import ResponseSchemaInput from pollux.providers.base import Provider @@ -287,7 +289,7 @@ async def stream( reasoning_budget_tokens: int | None = None, tool_choice: ToolChoice | None = None, provider_options: dict[str, dict[str, Any]] | None = None, -) -> AsyncIterator[Event]: +) -> AsyncGenerator[Event, None]: """Stream one explicit v2 interaction as :class:`Event` objects. The streaming sibling of :func:`interact`: same environment/input/config and @@ -326,7 +328,7 @@ async def stream( result = event.output """ async with Session(config) as session: - async for event in session.stream( + events = session.stream( environment, input, output=output, @@ -338,8 +340,12 @@ async def stream( reasoning_budget_tokens=reasoning_budget_tokens, tool_choice=tool_choice, provider_options=provider_options, - ): - yield event + ) + try: + async for event in events: + yield event + finally: + await close_async_iterator(events) class Session: @@ -421,7 +427,7 @@ async def stream( reasoning_budget_tokens: int | None = None, tool_choice: ToolChoice | None = None, provider_options: dict[str, dict[str, Any]] | None = None, - ) -> AsyncIterator[Event]: + ) -> AsyncGenerator[Event, None]: """Stream one interaction using the session's provider instance.""" self._ensure_open() requirements = _build_requirements( @@ -435,10 +441,14 @@ async def stream( tool_choice=tool_choice, provider_options=provider_options, ) - async for event in stream_interaction( + events = stream_interaction( environment, input, requirements, self.config, self._provider - ): - yield event + ) + try: + async for event in events: + yield event + finally: + await close_async_iterator(events) async def run_many( self, @@ -733,7 +743,7 @@ async def prepare_environment( if isinstance(cache, CachePolicy): provider = _get_provider(config) try: - snapshot = EnvironmentSnapshot.from_environment( + snapshot = _EnvironmentSnapshot.from_environment( environment, provider=config.provider ) await resolve_persistent_cache(snapshot, config, provider) @@ -965,6 +975,7 @@ def _resolve_deferred_provider(handle: DeferredHandle) -> Provider: "Input", "InternalError", "Message", + "MessageRole", "Output", "OutputCollection", "OutputRequirements", diff --git a/src/pollux/_lifecycle.py b/src/pollux/_lifecycle.py new file mode 100644 index 00000000..579bc898 --- /dev/null +++ b/src/pollux/_lifecycle.py @@ -0,0 +1,43 @@ +"""Shared asynchronous resource-lifecycle helpers.""" + +from __future__ import annotations + +import asyncio +import inspect +import logging +from typing import Any +from weakref import WeakSet + +logger = logging.getLogger(__name__) +_closed_iterators: WeakSet[Any] = WeakSet() + + +async def close_async_iterator(iterator: Any) -> None: + """Close an async iterator without masking an interaction's primary error.""" + close = getattr(iterator, "aclose", None) + if not callable(close): + close = getattr(iterator, "close", None) + if not callable(close): + return + try: + if iterator in _closed_iterators: + return + _closed_iterators.add(iterator) + except TypeError: + # Some third-party stream wrappers cannot be weak-referenced. Prefer a + # private marker when they allow attributes; otherwise rely on the SDK's + # idempotent close contract. + try: + if getattr(iterator, "_pollux_closed", False): + return + iterator._pollux_closed = True + except (AttributeError, TypeError): + pass + try: + result = close() + if inspect.isawaitable(result): + await result + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning("Async iterator cleanup failed: %s", exc) diff --git a/src/pollux/interaction/__init__.py b/src/pollux/interaction/__init__.py index 49aa7c3e..dce84636 100644 --- a/src/pollux/interaction/__init__.py +++ b/src/pollux/interaction/__init__.py @@ -13,12 +13,11 @@ from __future__ import annotations from pollux.interaction.collection import CollectionStatus, OutputCollection -from pollux.interaction.continuation import Continuation, Message +from pollux.interaction.continuation import Continuation, Message, MessageRole from pollux.interaction.environment import ( CachePolicy, CacheSetting, Environment, - EnvironmentSnapshot, ) from pollux.interaction.event import Event, EventType from pollux.interaction.input import Input @@ -47,12 +46,12 @@ "Continuation", "Diagnostics", "Environment", - "EnvironmentSnapshot", "Event", "EventType", "Input", "JSONValue", "Message", + "MessageRole", "Metrics", "Output", "OutputCollection", diff --git a/src/pollux/interaction/continuation.py b/src/pollux/interaction/continuation.py index 5f113847..c6c6dbb6 100644 --- a/src/pollux/interaction/continuation.py +++ b/src/pollux/interaction/continuation.py @@ -1,149 +1,268 @@ -"""The v2 ``Continuation`` primitive and its typed replay messages. - -``Continuation`` is the public, serializable state Pollux needs to continue a -provider-correct interaction. It replaces v1.x's private ``_conversation_state`` -dict. It is not memory: Pollux does not summarize, rank, or compact it. -""" +"""Portable transcript messages and opaque provider replay continuations.""" from __future__ import annotations +from collections.abc import Mapping, Sequence +from copy import deepcopy from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast -from pollux.errors import PolluxError +from pollux.errors import ConfigurationError, PolluxError from pollux.interaction.tools import ToolCall if TYPE_CHECKING: - from collections.abc import Mapping - from pollux.interaction.input import Input from pollux.providers.models import ProviderResponse -#: Bump when the serialized shape changes incompatibly. -SCHEMA_VERSION = 1 + +#: Bump when the serialized continuation shape or replay semantics change. +SCHEMA_VERSION = 2 + +MessageRole = Literal["user", "assistant", "tool"] +_MESSAGE_ROLES = {"user", "assistant", "tool"} @dataclass(frozen=True, slots=True) class Message: - """A typed conversational turn preserved for provider-correct replay.""" + """A portable, application-authored text or tool transcript message.""" - role: str + role: MessageRole content: str = "" - tool_calls: tuple[ToolCall, ...] = () + tool_calls: Sequence[ToolCall] = () tool_call_id: str | None = None - provider_state: dict[str, Any] | None = None + + def __post_init__(self) -> None: + """Freeze tool calls and reject shapes adapters cannot portably replay.""" + if self.role not in _MESSAGE_ROLES: + raise ConfigurationError( + f"Unsupported history message role: {self.role!r}", + hint="Use 'user', 'assistant', or 'tool'. Put system instructions " + "on Environment.instructions.", + ) + if not isinstance(self.content, str): + raise ConfigurationError( + "Message content must be text", + hint="Keep media in Source values or the current Input.", + ) + calls = tuple(self.tool_calls) + if not all(isinstance(call, ToolCall) for call in calls): + raise ConfigurationError( + "Message tool_calls must contain ToolCall values", + hint="Normalize provider tool calls with ToolCall.from_text(...).", + ) + if any(call.provider_state is not None for call in calls): + raise ConfigurationError( + "Portable history tool calls cannot contain provider state", + hint="Rebuild transcript calls with ToolCall.from_text(...).", + ) + object.__setattr__(self, "tool_calls", calls) + + if self.role == "user": + if not self.content.strip(): + raise ConfigurationError("User history messages require non-empty text") + if calls or self.tool_call_id is not None: + raise ConfigurationError( + "User history messages cannot contain tool-call fields" + ) + elif self.role == "assistant": + if not self.content and not calls: + raise ConfigurationError( + "Assistant history messages require text or tool calls" + ) + if self.tool_call_id is not None: + raise ConfigurationError( + "Assistant history messages cannot have tool_call_id" + ) + else: + if not isinstance(self.tool_call_id, str) or not self.tool_call_id.strip(): + raise ConfigurationError( + "Tool history messages require a non-empty tool_call_id" + ) + if calls: + raise ConfigurationError( + "Tool history messages cannot contain nested tool calls" + ) def to_jsonable(self) -> dict[str, Any]: - """Serialize to a compact JSON-compatible dict (optional facets omitted).""" + """Serialize this portable transcript message.""" payload: dict[str, Any] = {"role": self.role, "content": self.content} if self.tool_calls: - payload["tool_calls"] = [tc.to_jsonable() for tc in self.tool_calls] + payload["tool_calls"] = [call.to_jsonable() for call in self.tool_calls] if self.tool_call_id is not None: payload["tool_call_id"] = self.tool_call_id - if self.provider_state is not None: - payload["provider_state"] = self.provider_state return payload @classmethod def from_jsonable(cls, data: Mapping[str, Any]) -> Message: - """Parse a serialized message, type-guarding each facet.""" - raw_tool_calls = data.get("tool_calls") - tool_calls: tuple[ToolCall, ...] = () - if isinstance(raw_tool_calls, list): - tool_calls = tuple( - ToolCall.from_text( - id=str(tc.get("id", "")), - name=str(tc.get("name", "")), - arguments_text=str(tc.get("arguments_text", "")), - index=tc.get("index") if isinstance(tc.get("index"), int) else None, - provider_state=tc.get("provider_state") - if isinstance(tc.get("provider_state"), dict) - else None, - ) - for tc in raw_tool_calls - if isinstance(tc, dict) + """Parse and validate a portable transcript message.""" + role = data.get("role") + if role not in _MESSAGE_ROLES: + raise ConfigurationError( + f"Unsupported history message role: {role!r}", + hint="Use 'user', 'assistant', or 'tool'. Put system instructions " + "on Environment.instructions.", ) + raw_calls = data.get("tool_calls") + calls = _tool_calls_from_jsonable(raw_calls) content = data.get("content", "") + if not isinstance(content, str): + raise ConfigurationError("Message content must be text") tool_call_id = data.get("tool_call_id") - provider_state = data.get("provider_state") return cls( - role=str(data.get("role", "user")), - content=content if isinstance(content, str) else str(content), - tool_calls=tool_calls, + role=cast("MessageRole", role), + content=content, + tool_calls=calls, tool_call_id=tool_call_id if isinstance(tool_call_id, str) else None, - provider_state=provider_state if isinstance(provider_state, dict) else None, ) @classmethod def from_openai(cls, data: Mapping[str, Any]) -> Message: - """Build a replay message from an OpenAI Chat Completions message dict. + """Build a portable message from one OpenAI Chat Completions message. - This importer is for text transcript replay. Media attachments belong in - Pollux ``Source`` values or the current turn's ``Input`` content, not in - provider-shaped history. + System messages are intentionally rejected: stable system context belongs + on :attr:`Environment.instructions`, not in portable turn history. """ - raw_tool_calls = data.get("tool_calls") - tool_calls: tuple[ToolCall, ...] = () - if isinstance(raw_tool_calls, list): - tool_calls = tuple( - ToolCall.from_openai(tc) - for tc in raw_tool_calls - if isinstance(tc, dict) + role = data.get("role") + if role not in _MESSAGE_ROLES: + raise ConfigurationError( + f"Unsupported OpenAI history role: {role!r}; move system " + "messages to Environment.instructions", + hint="Move system messages to Environment.instructions.", + ) + raw_calls = data.get("tool_calls") + imported_calls = ( + tuple( + ToolCall.from_openai(call) + for call in raw_calls + if isinstance(call, dict) + ) + if isinstance(raw_calls, list) + else () + ) + calls = tuple( + ToolCall.from_text( + id=call.id, + name=call.name, + arguments_text=call.arguments_text, + index=call.index, ) + for call in imported_calls + ) tool_call_id = data.get("tool_call_id") return cls( - role=str(data.get("role", "user")), + role=cast("MessageRole", role), content=_openai_text_content(data.get("content")), - tool_calls=tool_calls, + tool_calls=calls, tool_call_id=tool_call_id if isinstance(tool_call_id, str) else None, - provider_state={"openai": dict(data)}, ) def to_openai(self) -> dict[str, Any]: - """Serialize as an OpenAI Chat Completions message dict.""" + """Serialize as one OpenAI Chat Completions transcript message.""" payload: dict[str, Any] = {"role": self.role, "content": self.content} if self.tool_calls: - payload["tool_calls"] = [tc.to_openai() for tc in self.tool_calls] + payload["tool_calls"] = [call.to_openai() for call in self.tool_calls] if self.tool_call_id is not None: payload["tool_call_id"] = self.tool_call_id return payload @dataclass(frozen=True, slots=True) -class Continuation: - """Serializable state for continuing a provider-correct interaction. +class _ReplayMessage: + """One internal continuation message, including opaque provider state.""" + + role: MessageRole + content: str = "" + tool_calls: tuple[ToolCall, ...] = () + tool_call_id: str | None = None + provider_state: dict[str, Any] | None = None + + @classmethod + def from_message(cls, message: Message) -> _ReplayMessage: + return cls( + role=message.role, + content=message.content, + tool_calls=tuple(message.tool_calls), + tool_call_id=message.tool_call_id, + ) + + @classmethod + def from_jsonable(cls, data: Mapping[str, Any]) -> _ReplayMessage: + role = data.get("role") + if role not in _MESSAGE_ROLES: + raise PolluxError( + f"Incompatible continuation message role: {role!r}", + hint="Start a new interaction instead of editing continuation state.", + ) + content = data.get("content", "") + if not isinstance(content, str): + raise PolluxError("Incompatible continuation message content") + tool_call_id = data.get("tool_call_id") + provider_state = data.get("provider_state") + return cls( + role=cast("MessageRole", role), + content=content, + tool_calls=_tool_calls_from_jsonable(data.get("tool_calls")), + tool_call_id=tool_call_id if isinstance(tool_call_id, str) else None, + provider_state=deepcopy(provider_state) + if isinstance(provider_state, dict) + else None, + ) + + def to_jsonable(self) -> dict[str, Any]: + payload: dict[str, Any] = {"role": self.role, "content": self.content} + if self.tool_calls: + payload["tool_calls"] = [ + deepcopy(call.to_jsonable()) for call in self.tool_calls + ] + if self.tool_call_id is not None: + payload["tool_call_id"] = self.tool_call_id + if self.provider_state is not None: + payload["provider_state"] = deepcopy(self.provider_state) + return payload - Read it from ``output.continuation`` and pass it back as - ``Input(continuation=...)`` to take the next turn. Persist it across processes - with :meth:`to_jsonable` / :meth:`from_jsonable`, which stamp and verify a - schema version (and, optionally, the producing provider). - A continuation is bound to the provider that produced it — its - ``provider_state`` (response ids, provider-specific replay blocks) is not - portable. Reusing one under a different provider is rejected before dispatch. - It is not memory: Pollux does not summarize, rank, or compact it. +@dataclass(frozen=True, slots=True) +class _ContinuationState: + messages: tuple[_ReplayMessage, ...] + response_id: str | None + provider: str + provider_state: dict[str, Any] | None + + +@dataclass(frozen=True, slots=True, init=False, repr=False) +class Continuation: + """Opaque, serializable state for provider-correct replay. + + Applications receive this value from :attr:`Output.continuation`, persist it + through :meth:`to_jsonable` / :meth:`from_jsonable`, and pass it back through + ``Input(continuation=...)``. Replay fields are intentionally not public. """ SCHEMA_VERSION: ClassVar[int] = SCHEMA_VERSION + __state: _ContinuationState - messages: tuple[Message, ...] = () - response_id: str | None = None - provider: str | None = None - provider_state: dict[str, Any] | None = None - version: int = SCHEMA_VERSION + def __init__(self) -> None: + raise TypeError( + "Continuation values are created by Pollux outputs or " + "Continuation.from_jsonable()" + ) + + def __repr__(self) -> str: + """Return a representation that does not reveal replay state.""" + return f"Continuation(version={SCHEMA_VERSION})" def to_jsonable(self) -> dict[str, Any]: - """Serialize to a JSON-compatible dict with version/provider markers.""" + """Return a defensive JSON-compatible serialization of this handle.""" + state = self.__state payload: dict[str, Any] = { - "version": self.version, - "messages": [m.to_jsonable() for m in self.messages], + "version": SCHEMA_VERSION, + "provider": state.provider, + "messages": [message.to_jsonable() for message in state.messages], } - if self.provider is not None: - payload["provider"] = self.provider - if self.response_id is not None: - payload["response_id"] = self.response_id - if self.provider_state is not None: - payload["provider_state"] = self.provider_state + if state.response_id is not None: + payload["response_id"] = state.response_id + if state.provider_state is not None: + payload["provider_state"] = deepcopy(state.provider_state) return payload @classmethod @@ -153,14 +272,8 @@ def from_jsonable( *, expected_provider: str | None = None, ) -> Continuation: - """Parse a serialized continuation, rejecting incompatible artifacts. - - A continuation written by an incompatible schema version is refused with - a clear error rather than misread. When *expected_provider* is given, a - continuation produced by a different provider is also refused. - """ - raw_version = data.get("version") - version = raw_version if isinstance(raw_version, int) else None + """Restore a versioned artifact and optionally verify its provider.""" + version = data.get("version") if version != SCHEMA_VERSION: raise PolluxError( f"Incompatible continuation: expected schema version " @@ -169,7 +282,11 @@ def from_jsonable( "version. Start a new interaction instead of reusing it.", ) provider = data.get("provider") - provider = provider if isinstance(provider, str) else None + if not isinstance(provider, str) or not provider: + raise PolluxError( + "Incompatible continuation: missing provider identity", + hint="Start a new interaction instead of editing continuation state.", + ) if expected_provider is not None and provider != expected_provider: raise PolluxError( f"Continuation provider {provider!r} does not match the active " @@ -177,123 +294,145 @@ def from_jsonable( hint="Reuse a continuation only with the provider that produced it.", ) raw_messages = data.get("messages") - messages: tuple[Message, ...] = () - if isinstance(raw_messages, list): - messages = tuple( - Message.from_jsonable(m) for m in raw_messages if isinstance(m, dict) + if not isinstance(raw_messages, list): + raise PolluxError("Incompatible continuation: messages must be a list") + if not all(isinstance(message, Mapping) for message in raw_messages): + raise PolluxError( + "Incompatible continuation: every message must be an object" ) response_id = data.get("response_id") provider_state = data.get("provider_state") - return cls( - messages=messages, + return _new_continuation( + messages=tuple( + _ReplayMessage.from_jsonable(message) for message in raw_messages + ), response_id=response_id if isinstance(response_id, str) else None, provider=provider, - provider_state=provider_state if isinstance(provider_state, dict) else None, - version=version, + provider_state=deepcopy(provider_state) + if isinstance(provider_state, dict) + else None, ) - @classmethod - def from_openai_messages( - cls, - messages: list[Mapping[str, Any]] | tuple[Mapping[str, Any], ...], - *, - provider: str | None = None, - response_id: str | None = None, - ) -> Continuation: - """Build a continuation from OpenAI Chat Completions replay messages.""" - return cls( - messages=tuple(Message.from_openai(message) for message in messages), + +def _new_continuation( + *, + messages: tuple[_ReplayMessage, ...], + response_id: str | None, + provider: str, + provider_state: dict[str, Any] | None, +) -> Continuation: + continuation = object.__new__(Continuation) + object.__setattr__( + continuation, + "_Continuation__state", + _ContinuationState( + messages=messages, response_id=response_id, provider=provider, - ) + provider_state=deepcopy(provider_state), + ), + ) + return continuation + - def to_openai_messages(self) -> list[dict[str, Any]]: - """Serialize continuation messages as OpenAI Chat Completions messages.""" - return [message.to_openai() for message in self.messages] +def _continuation_state(continuation: Continuation) -> _ContinuationState: + """Return opaque replay state for Pollux internals.""" + return cast( + "_ContinuationState", + object.__getattribute__(continuation, "_Continuation__state"), + ) + + +def _tool_calls_from_jsonable(raw: Any) -> tuple[ToolCall, ...]: + if not isinstance(raw, list): + return () + return tuple( + ToolCall.from_text( + id=str(call.get("id", "")), + name=str(call.get("name", "")), + arguments_text=str(call.get("arguments_text", "")), + index=call.get("index") if isinstance(call.get("index"), int) else None, + provider_state=deepcopy(call.get("provider_state")) + if isinstance(call.get("provider_state"), dict) + else None, + ) + for call in raw + if isinstance(call, dict) + ) def _openai_text_content(content: Any) -> str: - """Extract text from OpenAI string or text-part message content.""" if content is None: return "" if isinstance(content, str): return content if isinstance(content, list): - text_parts: list[str] = [] - for part in content: - if not isinstance(part, dict): - continue - text = part.get("text") - if isinstance(text, str): - text_parts.append(text) - return "\n\n".join(text_parts) + return "\n\n".join( + part["text"] + for part in content + if isinstance(part, dict) and isinstance(part.get("text"), str) + ) return str(content) -def _prior_messages(input: Input) -> tuple[Message, ...]: # noqa: A002 - """Replay messages preceding this turn: prior state plus returned tool results.""" +def _prior_messages(input: Input) -> tuple[_ReplayMessage, ...]: # noqa: A002 if input.continuation is not None: - prior = input.continuation.messages + prior = _continuation_state(input.continuation).messages elif input.history is not None: - prior = tuple(input.history) + prior = tuple(_ReplayMessage.from_message(message) for message in input.history) else: prior = () tool_messages = tuple( - Message(role="tool", content=tr.content, tool_call_id=tr.call_id) - for tr in input.tool_results + _ReplayMessage(role="tool", content=result.content, tool_call_id=result.call_id) + for result in input.tool_results ) return prior + tool_messages def build_continuation( - input: Input, # noqa: A002 - "input" is the canonical v2 primitive name + input: Input, # noqa: A002 response: ProviderResponse, *, user_content: str | None, provider: str | None, ) -> Continuation | None: - """Assemble the next-turn continuation for one interaction, or ``None``. - - A continuation is produced when the caller opted into conversation continuity - (the input carried prior ``continuation`` or ``history``) or when the response - carries tool calls the caller may need to return results for. It appends this - turn's user message (when present) and the assistant reply to the prior replay - messages; the assistant message and the continuation both carry the response's - opaque ``provider_state`` for correct replay. - """ + """Assemble the next opaque continuation after a successful interaction.""" wants_conversation = input.continuation is not None or input.history is not None response_tool_calls = response.tool_calls or () if not (wants_conversation or response_tool_calls): return None + if provider is None: + raise PolluxError("Cannot create a continuation without provider identity") messages = list(_prior_messages(input)) turn_user_content = user_content or input.content if turn_user_content is not None: - messages.append(Message(role="user", content=turn_user_content)) + messages.append(_ReplayMessage(role="user", content=turn_user_content)) provider_state = ( - response.provider_state if isinstance(response.provider_state, dict) else None + deepcopy(response.provider_state) + if isinstance(response.provider_state, dict) + else None ) messages.append( - Message( + _ReplayMessage( role="assistant", content=response.text, tool_calls=tuple( - ToolCall.from_text(id=tc.id, name=tc.name, arguments_text=tc.arguments) - for tc in response_tool_calls + ToolCall.from_text( + id=call.id, + name=call.name, + arguments_text=call.arguments, + ) + for call in response_tool_calls ), provider_state=provider_state, ) ) - response_id = ( - response.response_id - if isinstance(response.response_id, str) - else input.continuation.response_id - if input.continuation is not None - else None + response.response_id if isinstance(response.response_id, str) else None ) - return Continuation( + return _new_continuation( messages=tuple(messages), response_id=response_id, provider=provider, diff --git a/src/pollux/interaction/environment.py b/src/pollux/interaction/environment.py index fe1cbddf..21cf5736 100644 --- a/src/pollux/interaction/environment.py +++ b/src/pollux/interaction/environment.py @@ -3,8 +3,8 @@ An :class:`Environment` is the stable context around one or more interactions: instructions, sources, tool declarations, and a cache preference. It does not contain conversation history or application memory. An -:class:`EnvironmentSnapshot` is the planned, immutable provider-facing form whose -``fingerprint`` backs continuation and cache compatibility checks. +:class:`EnvironmentSnapshot` is the internal, planned, immutable provider-facing +form used by the execution path. """ from __future__ import annotations @@ -21,6 +21,9 @@ from pollux.source import Source +_FINGERPRINT_VERSION = 1 + + @dataclass(frozen=True, slots=True) class CachePolicy: """An explicit persistent-cache preference for an environment.""" @@ -52,15 +55,65 @@ def __post_init__(self) -> None: object.__setattr__(self, "sources", tuple(self.sources)) object.__setattr__(self, "tools", tuple(self.tools)) + def fingerprint(self, *, provider: str) -> str: + """Return stable identity for the provider-visible environment. + + The fingerprint covers instructions, ordered sources, ordered tools, + and the active provider. Model identity belongs to :class:`Config` and + must be composed separately by durable runtimes. Cache preferences and + metadata are deliberately excluded because they do not change the + model-visible environment. + + Fingerprints remain stable across compatible Pollux releases. A future + semantic change will increment the embedded fingerprint version. + """ + if not isinstance(provider, str) or not provider: + from pollux.errors import ConfigurationError + + raise ConfigurationError( + "Environment fingerprint requires a non-empty provider", + hint="Pass config.provider when creating durable identity.", + ) + payload = { + "version": _FINGERPRINT_VERSION, + "provider": provider, + "instructions": self.instructions, + "sources": [ + source._environment_identity(provider=provider) + for source in self.sources + ], + "tools": [ + { + "name": tool.name, + "description": tool.description, + "parameters": tool.parameters, + "strict": tool.strict, + } + for tool in self.tools + ], + } + try: + encoded = json.dumps( + payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True + ).encode("utf-8") + except (TypeError, ValueError) as exc: + from pollux.errors import ConfigurationError + + raise ConfigurationError( + "Environment identity is not JSON serializable", + hint="Use JSON-compatible tool schemas and provider hints.", + ) from exc + return hashlib.sha256(encoded).hexdigest() + @dataclass(frozen=True, slots=True) class EnvironmentSnapshot: """The planned, immutable provider-facing environment for one interaction. ``instructions``/``sources``/``tools``/``cache``/``provider`` describe the - environment's identity (and back :meth:`fingerprint`). The remaining fields - are core-populated transport state, frozen onto the snapshot by the execution - path just before ``Provider.generate`` so adapters compile from primitives: + planned request. The remaining fields are core-populated transport state, + frozen onto the snapshot by the execution path just before + ``Provider.generate`` so adapters compile from primitives: - ``prepared_parts``: the environment's shared source parts with local files already uploaded (single-flight, once per fan-out); empty when a persistent @@ -68,7 +121,8 @@ class EnvironmentSnapshot: - ``cache_name``: the resolved provider persistent-cache name, if any. - ``implicit_caching``: whether provider-managed implicit caching is enabled. - These derived fields are intentionally excluded from :meth:`fingerprint`. + Applications use :meth:`Environment.fingerprint` for durable identity; the + snapshot is an internal planning type. """ instructions: str | None = None @@ -80,33 +134,6 @@ class EnvironmentSnapshot: cache_name: str | None = None implicit_caching: bool = False - def fingerprint(self) -> str: - """Return a stable hash of the provider-facing environment identity. - - Used to reject reuse of a continuation or cache handle against an - environment whose instructions, sources, or tools have changed. - """ - payload = { - "instructions": self.instructions, - "sources": [ - source.cache_identity_hash(provider=self.provider) - for source in self.sources - ], - "tools": [ - { - "name": tool.name, - "description": tool.description, - "parameters": tool.parameters, - } - for tool in self.tools - ], - "cache": _cache_fingerprint(self.cache), - } - encoded = json.dumps( - payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True - ).encode("utf-8") - return hashlib.sha256(encoded).hexdigest() - @classmethod def from_environment( cls, environment: Environment, *, provider: str | None = None @@ -119,10 +146,3 @@ def from_environment( cache=environment.cache, provider=provider, ) - - -def _cache_fingerprint(cache: CacheSetting) -> Any: - """Reduce a cache setting to a JSON-stable fingerprint component.""" - if isinstance(cache, CachePolicy): - return {"ttl_seconds": cache.ttl_seconds} - return cache diff --git a/src/pollux/interaction/execute.py b/src/pollux/interaction/execute.py index 9b424909..0272d082 100644 --- a/src/pollux/interaction/execute.py +++ b/src/pollux/interaction/execute.py @@ -16,6 +16,7 @@ import time from typing import TYPE_CHECKING +from pollux._lifecycle import close_async_iterator from pollux.cache import create_cache_impl from pollux.errors import APIError, ConfigurationError, InternalError, PolluxError from pollux.interaction._uploads import cleanup_uploads, substitute_upload_parts @@ -387,11 +388,10 @@ async def stream_interaction( finish_reason: str | None = None response_id: str | None = None + provider_stream = provider.stream_generate(snapshot, input, requirements, config) try: yield Event(type="start") - async for chunk in provider.stream_generate( - snapshot, input, requirements, config - ): + async for chunk in provider_stream: if chunk.text: text_parts.append(chunk.text) yield Event(type="text_delta", text=chunk.text) @@ -449,4 +449,5 @@ async def stream_interaction( ) yield Event(type="done", output=output) finally: + await close_async_iterator(provider_stream) await cleanup_uploads(upload_cache, provider) diff --git a/src/pollux/interaction/validate.py b/src/pollux/interaction/validate.py index 8b87ec72..fdb5ccd2 100644 --- a/src/pollux/interaction/validate.py +++ b/src/pollux/interaction/validate.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING from pollux.errors import ConfigurationError +from pollux.interaction.continuation import _continuation_state if TYPE_CHECKING: from collections.abc import Sequence @@ -46,11 +47,12 @@ def _reject_incompatible_continuations( return for inp in inputs: continuation = inp.continuation - if continuation is None or continuation.provider is None: + if continuation is None: continue - if continuation.provider != active: + producing_provider = _continuation_state(continuation).provider + if producing_provider != active: raise ConfigurationError( - f"Continuation was produced by provider {continuation.provider!r}, " + f"Continuation was produced by provider {producing_provider!r}, " f"but the active provider is {active!r}", hint="Reuse a continuation only with the provider that produced " "it, or start a new interaction.", diff --git a/src/pollux/providers/_compile.py b/src/pollux/providers/_compile.py index 6d4f92f5..24f173df 100644 --- a/src/pollux/providers/_compile.py +++ b/src/pollux/providers/_compile.py @@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any +from pollux.interaction.continuation import _continuation_state from pollux.providers.models import ( Message as ProviderMessage, ) @@ -102,21 +103,25 @@ def prior_turns( state is folded under a ``"history"`` key so the transport can replay opaque blocks (e.g. reasoning) during continuation. """ - prior: tuple[Message, ...] = () + prior: tuple[Any, ...] = () previous_response_id: str | None = None provider_state: dict[str, object] | None = None if input.continuation is not None: - prior = input.continuation.messages - previous_response_id = input.continuation.response_id - if input.continuation.provider_state is not None: - provider_state = dict(input.continuation.provider_state) + state = _continuation_state(input.continuation) + prior = state.messages + previous_response_id = state.response_id + if state.provider_state is not None: + provider_state = dict(state.provider_state) elif input.history is not None: prior = tuple(input.history) messages = [_provider_message(m) for m in prior] item_states: list[dict[str, object] | None] = [ - dict(m.provider_state) if m.provider_state is not None else None for m in prior + dict(m.provider_state) + if getattr(m, "provider_state", None) is not None + else None + for m in prior ] messages.extend(_tool_result_message(tr) for tr in input.tool_results) diff --git a/src/pollux/providers/anthropic.py b/src/pollux/providers/anthropic.py index a5f65576..252f38f1 100644 --- a/src/pollux/providers/anthropic.py +++ b/src/pollux/providers/anthropic.py @@ -9,6 +9,7 @@ import logging from typing import TYPE_CHECKING, Any, cast +from pollux._lifecycle import close_async_iterator from pollux.errors import APIError, ConfigurationError from pollux.interaction.tools import ToolCallDelta from pollux.parts import build_shared_parts @@ -620,6 +621,7 @@ async def stream_generate( create_kwargs["stream"] = True assembler = _AnthropicStreamAssembler() + stream: Any = None try: stream = await client.messages.create(**create_kwargs) async for event in stream: @@ -641,6 +643,9 @@ async def stream_generate( allow_network_errors=True, message="Anthropic stream failed", ) from e + finally: + if stream is not None: + await close_async_iterator(stream) async def _resolve_deferred_parts( self, diff --git a/src/pollux/providers/gemini.py b/src/pollux/providers/gemini.py index ab1ecfd1..a9e12c11 100644 --- a/src/pollux/providers/gemini.py +++ b/src/pollux/providers/gemini.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any import uuid +from pollux._lifecycle import close_async_iterator from pollux.errors import APIError, ConfigurationError from pollux.interaction.tools import ToolCallDelta from pollux.parts import build_shared_parts @@ -669,6 +670,7 @@ async def stream_generate( contents = self._build_contents(parts, history or None) tool_call_index = 0 + stream: Any = None try: stream = await client.aio.models.generate_content_stream( model=config.model, @@ -690,6 +692,9 @@ async def stream_generate( allow_network_errors=True, message="Gemini stream failed", ) from e + finally: + if stream is not None: + await close_async_iterator(stream) def _stream_response_to_chunks( self, response: Any, tool_call_index: int diff --git a/src/pollux/providers/openai.py b/src/pollux/providers/openai.py index f12ae18e..f6d46b3b 100644 --- a/src/pollux/providers/openai.py +++ b/src/pollux/providers/openai.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any from urllib.parse import urlparse +from pollux._lifecycle import close_async_iterator from pollux.errors import APIError, ConfigurationError from pollux.interaction.tools import ToolCallDelta from pollux.parts import build_shared_parts @@ -549,6 +550,7 @@ async def stream_generate( ) create_kwargs["stream"] = True + stream: Any = None try: stream = await client.responses.create(**create_kwargs) async for event in stream: @@ -567,6 +569,9 @@ async def stream_generate( allow_network_errors=True, message="OpenAI stream failed", ) from e + finally: + if stream is not None: + await close_async_iterator(stream) def _stream_event_to_chunk(self, event: Any) -> ProviderStreamChunk | None: """Map one Responses API stream event to a normalized chunk.""" diff --git a/src/pollux/source.py b/src/pollux/source.py index 6cfc8c31..3ac2d6cd 100644 --- a/src/pollux/source.py +++ b/src/pollux/source.py @@ -248,6 +248,22 @@ def cache_identity_hash(self, *, provider: str | None = None) -> str: ) return hashlib.sha256(combined.encode("utf-8")).hexdigest() + def _environment_identity(self, *, provider: str) -> dict[str, Any]: + """Return provider-visible source identity for environment hashing.""" + identity: dict[str, Any] = { + "source_type": self.source_type, + "mime_type": self.mime_type, + "content_hash": self._content_hash(), + } + if self.source_type == "file": + identity["identifier"] = Path(self.identifier).name + elif self.source_type not in {"text", "json"}: + identity["identifier"] = self.identifier + provider_hints = self.provider_hints_for(provider) + if provider_hints is not None: + identity["provider_hints"] = provider_hints + return identity + def with_gemini_video_settings( self, *, diff --git a/tests/helpers.py b/tests/helpers.py index e36c34d0..65b316e3 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -12,6 +12,7 @@ from pollux.config import Config from pollux.errors import APIError, ConfigurationError +from pollux.interaction.continuation import SCHEMA_VERSION, Continuation, Message from pollux.interaction.environment import EnvironmentSnapshot from pollux.interaction.input import Input from pollux.interaction.requirements import OutputRequirements @@ -29,7 +30,39 @@ from collections.abc import Sequence from pollux.config import ProviderName - from pollux.interaction.continuation import Continuation, Message + + +def make_continuation( + *, + messages: Sequence[Message] = (), + serialized_messages: Sequence[dict[str, Any]] | None = None, + response_id: str | None = None, + provider: str = "anthropic", + provider_state: dict[str, Any] | None = None, + message_provider_states: Sequence[dict[str, Any] | None] | None = None, +) -> Continuation: + """Restore an opaque continuation fixture through its public artifact API.""" + message_payloads = ( + [dict(message) for message in serialized_messages] + if serialized_messages is not None + else [message.to_jsonable() for message in messages] + ) + if message_provider_states is not None: + for message, state in zip( + message_payloads, message_provider_states, strict=False + ): + if state is not None: + message["provider_state"] = state + payload: dict[str, Any] = { + "version": SCHEMA_VERSION, + "provider": provider, + "messages": message_payloads, + } + if response_id is not None: + payload["response_id"] = response_id + if provider_state is not None: + payload["provider_state"] = provider_state + return Continuation.from_jsonable(payload) def make_interaction( diff --git a/tests/interaction/test_continuation.py b/tests/interaction/test_continuation.py index d52b1334..141ad38e 100644 --- a/tests/interaction/test_continuation.py +++ b/tests/interaction/test_continuation.py @@ -1,165 +1,158 @@ -"""Unit tests for the v2 ``Continuation`` primitive.""" +"""Unit tests for portable messages and opaque continuations.""" from __future__ import annotations -from typing import TYPE_CHECKING, Any +from copy import deepcopy import pytest -from pollux.errors import PolluxError -from pollux.interaction.continuation import SCHEMA_VERSION, Continuation, Message +from pollux.errors import ConfigurationError, PolluxError +from pollux.interaction.continuation import ( + SCHEMA_VERSION, + Continuation, + Message, + build_continuation, +) +from pollux.interaction.input import Input from pollux.interaction.tools import ToolCall - -if TYPE_CHECKING: - from collections.abc import Mapping +from pollux.providers.models import ProviderResponse +from tests.helpers import make_continuation pytestmark = pytest.mark.unit -def test_continuation_roundtrips_through_jsonable(): - cont = Continuation( +def test_continuation_roundtrips_and_defensively_serializes() -> None: + original = make_continuation( messages=( Message(role="user", content="hi"), Message( role="assistant", - content="", tool_calls=( - ToolCall.from_text(id="c1", name="f", arguments_text='{"a": 1}'), + ToolCall.from_text(id="c1", name="f", arguments_text="{}"), ), ), ), response_id="r1", provider="anthropic", + provider_state={"secret": {"value": 1}}, ) - restored = Continuation.from_jsonable(cont.to_jsonable()) - assert restored.response_id == "r1" - assert restored.provider == "anthropic" - assert restored.messages[0].content == "hi" - assert restored.messages[1].tool_calls[0].name == "f" - assert restored.messages[1].tool_calls[0].arguments == {"a": 1} - + blob = original.to_jsonable() + restored = Continuation.from_jsonable(blob, expected_provider="anthropic") + assert restored.to_jsonable() == blob + assert blob["version"] == SCHEMA_VERSION == 2 -def test_continuation_stamps_current_schema_version(): - assert Continuation().to_jsonable()["version"] == SCHEMA_VERSION + mutated = deepcopy(blob) + mutated["provider_state"]["secret"]["value"] = 99 + assert original.to_jsonable()["provider_state"]["secret"]["value"] == 1 -def test_continuation_rejects_incompatible_version(): - blob = Continuation(provider="mock").to_jsonable() - blob["version"] = SCHEMA_VERSION + 1 - with pytest.raises(PolluxError, match="Incompatible continuation"): - Continuation.from_jsonable(blob) +def test_continuation_has_no_public_constructor_or_replay_fields() -> None: + with pytest.raises(TypeError, match="created by Pollux"): + Continuation() + continuation = make_continuation(provider="openai") + for name in ("messages", "response_id", "provider", "provider_state"): + assert not hasattr(continuation, name) + assert not hasattr(Continuation, "from_openai_messages") + assert not hasattr(continuation, "to_openai_messages") -def test_continuation_rejects_missing_version(): +@pytest.mark.parametrize("version", [1, SCHEMA_VERSION + 1, None]) +def test_continuation_rejects_incompatible_versions(version: int | None) -> None: with pytest.raises(PolluxError, match="Incompatible continuation"): - Continuation.from_jsonable({"messages": []}) + Continuation.from_jsonable( + {"version": version, "provider": "openai", "messages": []} + ) -def test_continuation_rejects_provider_mismatch(): - blob = Continuation(provider="anthropic").to_jsonable() +def test_continuation_rejects_missing_or_mismatched_provider() -> None: + with pytest.raises(PolluxError, match="missing provider"): + Continuation.from_jsonable({"version": SCHEMA_VERSION, "messages": []}) + blob = make_continuation(provider="anthropic").to_jsonable() with pytest.raises(PolluxError, match="does not match"): Continuation.from_jsonable(blob, expected_provider="openai") -def test_continuation_accepts_matching_provider(): - blob = Continuation(provider="openai").to_jsonable() - restored = Continuation.from_jsonable(blob, expected_provider="openai") - assert restored.provider == "openai" - - -def test_openai_messages_import_tool_calls(): - continuation = Continuation.from_openai_messages( - [ - {"role": "user", "content": "What is the weather?"}, - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city":"Paris"}', - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_1", - "content": '{"temp_c": 21}', - }, - ], - provider="local", - ) - - assert continuation.provider == "local" - assert continuation.messages[1].tool_calls[0].name == "get_weather" - assert continuation.messages[1].tool_calls[0].arguments_dict() == {"city": "Paris"} - assert continuation.messages[2].tool_call_id == "call_1" - - -def test_openai_messages_round_trip_supported_transcript_shapes(): - messages: list[Mapping[str, Any]] = [ - {"role": "system", "content": "Be concise."}, - {"role": "user", "content": "What is the weather?"}, - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city":"Paris"}', - }, - "index": 0, - } - ], - }, +@pytest.mark.parametrize( + "message", + [ + Message(role="user", content="question"), + Message(role="assistant", content="answer"), + Message( + role="assistant", + tool_calls=(ToolCall.from_text(id="c1", name="lookup"),), + ), + Message(role="tool", tool_call_id="c1", content=""), + ], +) +def test_portable_message_shapes_roundtrip(message: Message) -> None: + assert Message.from_jsonable(message.to_jsonable()) == message + + +@pytest.mark.parametrize( + "kwargs", + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": ""}, + {"role": "user", "content": "x", "tool_call_id": "c1"}, + {"role": "assistant", "content": ""}, + {"role": "assistant", "content": "x", "tool_call_id": "c1"}, + {"role": "tool", "content": "result"}, { "role": "tool", - "tool_call_id": "call_1", - "content": '{"temp_c": 21}', + "content": "result", + "tool_call_id": "c1", + "tool_calls": [ToolCall.from_text(id="nested", name="bad")], }, - {"role": "assistant", "content": "It is 21 C in Paris."}, - ] + ], +) +def test_rejects_nonportable_message_shapes(kwargs: dict[str, object]) -> None: + with pytest.raises(ConfigurationError): + Message(**kwargs) # type: ignore[arg-type] + + +def test_message_openai_conversion_preserves_text_and_tools() -> None: + raw = { + "role": "assistant", + "content": "checking", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q":"x"}'}, + } + ], + } + assert Message.from_openai(raw).to_openai() == raw - continuation = Continuation.from_openai_messages(messages, provider="local") - assert continuation.to_openai_messages() == messages +def test_message_openai_conversion_rejects_system_role() -> None: + with pytest.raises(ConfigurationError, match=r"Environment\.instructions"): + Message.from_openai({"role": "system", "content": "Be concise"}) -def test_openai_messages_round_trip_tool_call_arguments_text(): - continuation = Continuation( - messages=( - Message( - role="assistant", - tool_calls=( - ToolCall.from_text( - id="call_1", - name="run", - arguments_text='{"cmd":"pwd"}', - ), - ), - ), - ) +def test_successful_manual_history_creates_fresh_provider_continuation() -> None: + continuation = build_continuation( + Input(content="next", history=[Message(role="user", content="summary")]), + ProviderResponse(text="answer", response_id="resp_new"), + user_content="next", + provider="openai", ) - - messages = continuation.to_openai_messages() - - assert messages == [ - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "run", "arguments": '{"cmd":"pwd"}'}, - } - ], - } - ] + assert continuation is not None + serialized = continuation.to_jsonable() + assert serialized["provider"] == "openai" + assert serialized["response_id"] == "resp_new" + + +def test_manual_history_does_not_create_opaque_anthropic_thinking_state() -> None: + continuation = build_continuation( + Input( + tool_results=[], + content="continue", + history=[Message(role="user", content="compacted summary")], + ), + ProviderResponse(text="answer"), + user_content="continue", + provider="anthropic", + ) + assert continuation is not None + assert "anthropic_thinking_blocks" not in str(continuation.to_jsonable()) diff --git a/tests/interaction/test_continuation_compat.py b/tests/interaction/test_continuation_compat.py index 16fa061a..5b35381e 100644 --- a/tests/interaction/test_continuation_compat.py +++ b/tests/interaction/test_continuation_compat.py @@ -18,6 +18,7 @@ from pollux.interaction.requirements import OutputRequirements from pollux.providers.base import ProviderCapabilities from tests.conftest import ANTHROPIC_MODEL, FakeProvider +from tests.helpers import make_continuation pytestmark = pytest.mark.integration @@ -34,8 +35,8 @@ def _cfg() -> Config: return Config(provider="anthropic", model=ANTHROPIC_MODEL, use_mock=True) -def _continuation(provider: str | None) -> Continuation: - return Continuation( +def _continuation(provider: str) -> Continuation: + return make_continuation( messages=(Message(role="user", content="earlier"),), provider=provider, ) @@ -63,17 +64,3 @@ async def test_accepts_continuation_from_the_matching_provider() -> None: _conversational_provider(), ) assert out.text == "ok:next" - - -@pytest.mark.asyncio -async def test_accepts_continuation_without_a_provider_marker() -> None: - # Hand-built or history-derived continuations carry no provider marker and - # are left alone by the compatibility check. - out = await execute_interaction( - Environment(), - Input(content="next", continuation=_continuation(None)), - OutputRequirements(), - _cfg(), - _conversational_provider(), - ) - assert out.text == "ok:next" diff --git a/tests/interaction/test_environment.py b/tests/interaction/test_environment.py index 82a094f7..5430517d 100644 --- a/tests/interaction/test_environment.py +++ b/tests/interaction/test_environment.py @@ -1,62 +1,132 @@ -"""Unit tests for ``Environment`` and ``EnvironmentSnapshot``.""" +"""Unit tests for public ``Environment`` identity.""" from __future__ import annotations +from typing import TYPE_CHECKING + import pytest -from pollux.interaction.environment import ( - CachePolicy, - Environment, - EnvironmentSnapshot, -) +import pollux +import pollux.interaction +from pollux.interaction.environment import CachePolicy, Environment from pollux.interaction.tools import ToolDeclaration from pollux.source import Source +if TYPE_CHECKING: + from pathlib import Path + pytestmark = pytest.mark.unit -def test_environment_freezes_sequences_to_tuples(): +def test_environment_freezes_sequences_to_tuples() -> None: env = Environment( sources=[Source.from_text("a")], tools=[ToolDeclaration(name="f", description="d")], ) assert isinstance(env.sources, tuple) assert isinstance(env.tools, tuple) - assert env.tools[0].name == "f" -def test_environment_accepts_declaration_objects(): - decl = ToolDeclaration(name="f", description="d") - env = Environment(tools=[decl]) - assert env.tools[0] is decl +def test_fingerprint_is_stable_and_normalizes_mapping_order() -> None: + first = Environment( + instructions="sys", + tools=[ + ToolDeclaration( + name="f", + parameters={"type": "object", "properties": {"a": {}, "b": {}}}, + ) + ], + ) + second = Environment( + instructions="sys", + tools=[ + ToolDeclaration( + name="f", + parameters={"properties": {"b": {}, "a": {}}, "type": "object"}, + ) + ], + ) + assert first.fingerprint(provider="openai") == second.fingerprint(provider="openai") + + +@pytest.mark.parametrize( + ("first", "second"), + [ + (Environment(instructions="a"), Environment(instructions="b")), + ( + Environment(sources=[Source.from_text("a"), Source.from_text("b")]), + Environment(sources=[Source.from_text("b"), Source.from_text("a")]), + ), + ( + Environment(tools=[ToolDeclaration(name="f", strict=True)]), + Environment(tools=[ToolDeclaration(name="f", strict=False)]), + ), + ( + Environment( + tools=[ToolDeclaration(name="f", parameters={"type": "object"})] + ), + Environment( + tools=[ToolDeclaration(name="f", parameters={"type": "string"})] + ), + ), + ], +) +def test_fingerprint_changes_with_model_visible_environment( + first: Environment, second: Environment +) -> None: + assert first.fingerprint(provider="openai") != second.fingerprint(provider="openai") -def test_snapshot_fingerprint_is_stable(): - env = Environment(instructions="sys", sources=(Source.from_text("doc"),)) - snap_a = EnvironmentSnapshot.from_environment(env) - snap_b = EnvironmentSnapshot.from_environment(env) - assert snap_a.fingerprint() == snap_b.fingerprint() +def test_fingerprint_changes_with_provider_and_provider_hints() -> None: + source = Source.from_youtube("https://youtu.be/example").with_gemini_video_settings( + fps=1 + ) + env = Environment(sources=[source]) + assert env.fingerprint(provider="gemini") != env.fingerprint(provider="openai") + + +def test_fingerprint_tracks_local_file_content_mime_and_identifier( + tmp_path: Path, +) -> None: + first_path = tmp_path / "first.bin" + second_path = tmp_path / "second.bin" + first_path.write_bytes(b"one") + second_path.write_bytes(b"one") + + original = Environment(sources=[Source.from_file(first_path)]) + different_name = Environment(sources=[Source.from_file(second_path)]) + different_mime = Environment( + sources=[Source.from_file(first_path, mime_type="text/plain")] + ) + original_hash = original.fingerprint(provider="anthropic") + assert original_hash != different_name.fingerprint(provider="anthropic") + assert original_hash != different_mime.fingerprint(provider="anthropic") -def test_snapshot_fingerprint_changes_with_instructions(): - base = EnvironmentSnapshot.from_environment(Environment(instructions="a")) - other = EnvironmentSnapshot.from_environment(Environment(instructions="b")) - assert base.fingerprint() != other.fingerprint() + first_path.write_bytes(b"two") + assert original_hash != original.fingerprint(provider="anthropic") -def test_snapshot_fingerprint_changes_with_sources(): - one = EnvironmentSnapshot.from_environment( - Environment(sources=(Source.from_text("x"),)) - ) - two = EnvironmentSnapshot.from_environment( - Environment(sources=(Source.from_text("y"),)) +def test_uri_fingerprint_is_network_free(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "urllib.request.urlopen", + lambda *_args, **_kwargs: pytest.fail("fingerprinting performed network I/O"), ) - assert one.fingerprint() != two.fingerprint() + env = Environment(sources=[Source.from_uri("https://example.com/report.pdf")]) + assert env.fingerprint(provider="gemini") -def test_snapshot_fingerprint_changes_with_cache_policy(): - none_cache = EnvironmentSnapshot.from_environment(Environment()) - ttl_cache = EnvironmentSnapshot.from_environment( - Environment(cache=CachePolicy(ttl_seconds=3600)) +def test_cache_and_metadata_do_not_change_fingerprint() -> None: + base = Environment(instructions="sys") + policy = Environment( + instructions="sys", + cache=CachePolicy(ttl_seconds=3600), + metadata={"runtime": "different"}, ) - assert none_cache.fingerprint() != ttl_cache.fingerprint() + assert base.fingerprint(provider="gemini") == policy.fingerprint(provider="gemini") + + +def test_environment_snapshot_is_not_publicly_exported() -> None: + assert "EnvironmentSnapshot" not in pollux.__all__ + assert "EnvironmentSnapshot" not in pollux.interaction.__all__ + assert not hasattr(pollux, "EnvironmentSnapshot") diff --git a/tests/interaction/test_execute_interaction.py b/tests/interaction/test_execute_interaction.py index c38fd45a..94f1ff62 100644 --- a/tests/interaction/test_execute_interaction.py +++ b/tests/interaction/test_execute_interaction.py @@ -92,7 +92,8 @@ async def test_tool_calls_and_continuation(): assert out.tool_calls[0].arguments == {"city": "P"} assert out.metrics.completion_status == "clean" assert out.continuation is not None - assert any(message.tool_calls for message in out.continuation.messages) + serialized = out.continuation.to_jsonable() + assert any(message.get("tool_calls") for message in serialized["messages"]) @pytest.mark.asyncio diff --git a/tests/interaction/test_input.py b/tests/interaction/test_input.py index c1ae5daa..5c1bf82e 100644 --- a/tests/interaction/test_input.py +++ b/tests/interaction/test_input.py @@ -5,9 +5,10 @@ import pytest from pollux.errors import ConfigurationError -from pollux.interaction.continuation import Continuation, Message +from pollux.interaction.continuation import Message from pollux.interaction.input import Input from pollux.interaction.tools import ToolResult +from tests.helpers import make_continuation pytestmark = pytest.mark.unit @@ -39,13 +40,13 @@ def test_rejects_history_and_continuation_together(): Input( content="x", history=(Message(role="user", content="prior"),), - continuation=Continuation(provider="mock"), + continuation=make_continuation(provider="mock"), ) def test_continuation_only_with_tool_results_is_valid(): inp = Input( - continuation=Continuation(provider="mock"), + continuation=make_continuation(provider="mock"), tool_results=[ToolResult(call_id="c1", content="ok")], ) assert inp.continuation is not None diff --git a/tests/interaction/test_interact_frontdoor.py b/tests/interaction/test_interact_frontdoor.py index bde1404d..08b63705 100644 --- a/tests/interaction/test_interact_frontdoor.py +++ b/tests/interaction/test_interact_frontdoor.py @@ -2,6 +2,9 @@ from __future__ import annotations +import asyncio +from typing import Any + from pydantic import BaseModel import pytest @@ -169,3 +172,75 @@ def test_local_reasoning_returns_scoped_provider_options() -> None: assert pollux.local_reasoning(enabled=False) == { "local": {"chat_template_kwargs": {"enable_thinking": False}} } + + +@pytest.mark.asyncio +async def test_cancelled_one_shot_interact_releases_request_and_provider( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class BlockingProvider(ScriptedProvider): + started = asyncio.Event() + request_released = False + provider_closed = False + + async def generate(self, *_args: Any) -> ProviderResponse: + self.started.set() + try: + await asyncio.Event().wait() + raise AssertionError("blocking request unexpectedly resumed") + finally: + self.request_released = True + + async def aclose(self) -> None: + self.provider_closed = True + + provider = BlockingProvider() + monkeypatch.setattr(pollux, "_get_provider", lambda _config: provider) + task = asyncio.create_task(interact(Environment(), Input("wait"), config=_cfg())) + await provider.started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + assert provider.request_released is True + assert provider.provider_closed is True + + +@pytest.mark.asyncio +async def test_cancelled_session_interact_keeps_session_reusable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class ReusableProvider(ScriptedProvider): + started = asyncio.Event() + request_released = False + provider_closed = False + calls = 0 + + async def generate(self, *_args: Any) -> ProviderResponse: + self.calls += 1 + if self.calls > 1: + return ProviderResponse(text="reused", usage={"total_tokens": 1}) + self.started.set() + try: + await asyncio.Event().wait() + raise AssertionError("blocking request unexpectedly resumed") + finally: + self.request_released = True + + async def aclose(self) -> None: + self.provider_closed = True + + provider = ReusableProvider() + monkeypatch.setattr(pollux, "_get_provider", lambda _config: provider) + + async with pollux.Session(_cfg()) as session: + task = asyncio.create_task(session.interact(Environment(), Input("wait"))) + await provider.started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert provider.request_released is True + assert provider.provider_closed is False + assert (await session.interact(Environment(), Input("again"))).text == "reused" + + assert provider.provider_closed is True diff --git a/tests/interaction/test_stream.py b/tests/interaction/test_stream.py index 9ad0c2c9..d587bc7d 100644 --- a/tests/interaction/test_stream.py +++ b/tests/interaction/test_stream.py @@ -7,6 +7,7 @@ from __future__ import annotations +import asyncio from dataclasses import dataclass, field from typing import Any @@ -14,6 +15,7 @@ import pollux from pollux import Environment, Event, Input +from pollux._lifecycle import close_async_iterator from pollux.config import Config from pollux.errors import APIError, ConfigurationError from pollux.interaction.execute import execute_interaction, stream_interaction @@ -229,3 +231,106 @@ async def test_stream_frontdoor_yields_events(monkeypatch: pytest.MonkeyPatch) - ] assert types == ["start", "text_delta", "usage", "finish", "done"] + + +@dataclass +class _CloseAwareStreamProvider: + """Block after one chunk and record request-stream/provider cleanup.""" + + stream_closed: bool = False + provider_closed: bool = False + waiting: asyncio.Event = field(default_factory=asyncio.Event) + + @property + def capabilities(self) -> ProviderCapabilities: + return ProviderCapabilities( + persistent_cache=False, + uploads=False, + conversation=True, + ) + + async def generate(self, *_args: Any) -> ProviderResponse: + return ProviderResponse(text="reused", usage={"total_tokens": 1}) + + async def stream_generate(self, *_args: Any) -> Any: + try: + yield ProviderStreamChunk(text="partial") + self.waiting.set() + await asyncio.Event().wait() + finally: + self.stream_closed = True + + async def aclose(self) -> None: + self.provider_closed = True + + +@pytest.mark.asyncio +async def test_public_stream_aclose_propagates_and_closes_owned_provider( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = _CloseAwareStreamProvider() + monkeypatch.setattr(pollux, "_get_provider", lambda _config: provider) + events = pollux.stream(Environment(), Input("hi"), config=_cfg()) + assert (await anext(events)).type == "start" + assert (await anext(events)).type == "text_delta" + + await events.aclose() + + assert provider.stream_closed is True + assert provider.provider_closed is True + + +@pytest.mark.asyncio +async def test_session_stream_aclose_releases_stream_but_keeps_session_reusable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = _CloseAwareStreamProvider() + monkeypatch.setattr(pollux, "_get_provider", lambda _config: provider) + + async with pollux.Session(_cfg()) as session: + events = session.stream(Environment(), Input("hi")) + assert (await anext(events)).type == "start" + assert (await anext(events)).type == "text_delta" + await events.aclose() + assert provider.stream_closed is True + assert provider.provider_closed is False + assert (await session.interact(Environment(), Input("again"))).text == "reused" + + assert provider.provider_closed is True + + +@pytest.mark.asyncio +async def test_cancelling_blocked_stream_iteration_closes_active_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = _CloseAwareStreamProvider() + monkeypatch.setattr(pollux, "_get_provider", lambda _config: provider) + events = pollux.stream(Environment(), Input("hi"), config=_cfg()) + await anext(events) + await anext(events) + + async def next_event() -> Event: + return await anext(events) + + pending: asyncio.Task[Event] = asyncio.create_task(next_event()) + await provider.waiting.wait() + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + assert provider.stream_closed is True + assert provider.provider_closed is True + + +@pytest.mark.asyncio +async def test_async_iterator_close_helper_is_idempotent() -> None: + class CloseCounter: + calls = 0 + + async def aclose(self) -> None: + self.calls += 1 + + iterator = CloseCounter() + await close_async_iterator(iterator) + await close_async_iterator(iterator) + assert iterator.calls == 1 diff --git a/tests/providers/test_anthropic_contract.py b/tests/providers/test_anthropic_contract.py index 6a00ed29..1b405a9d 100644 --- a/tests/providers/test_anthropic_contract.py +++ b/tests/providers/test_anthropic_contract.py @@ -8,7 +8,7 @@ import pytest from pollux.errors import APIError, ConfigurationError -from pollux.interaction.continuation import Continuation, Message +from pollux.interaction.continuation import Message from pollux.interaction.tools import ToolCall, ToolResult from pollux.providers.anthropic import AnthropicProvider from pollux.providers.models import ( @@ -17,7 +17,7 @@ from tests.conftest import ( ANTHROPIC_MODEL, ) -from tests.helpers import make_interaction +from tests.helpers import make_continuation, make_interaction def _anthropic(**kwargs: Any) -> tuple[Any, Any, Any, Any]: @@ -711,7 +711,8 @@ async def test_anthropic_generate_history_replays_preserved_thinking_blocks() -> await provider.generate( *_anthropic( - continuation=Continuation( + continuation=make_continuation( + provider="anthropic", messages=( Message( role="assistant", diff --git a/tests/providers/test_anthropic_stream.py b/tests/providers/test_anthropic_stream.py index f36ed033..52629fea 100644 --- a/tests/providers/test_anthropic_stream.py +++ b/tests/providers/test_anthropic_stream.py @@ -103,6 +103,7 @@ def __init__( ): self._events = events self._raise = raise_exc + self.closed = False def __aiter__(self) -> _FakeStream: return self @@ -114,6 +115,9 @@ async def __anext__(self) -> SimpleNamespace: raise StopAsyncIteration return self._events.pop(0) + async def close(self) -> None: + self.closed = True + class _FakeMessages: def __init__( @@ -122,10 +126,12 @@ def __init__( self._events = events self._raise = raise_exc self.last_kwargs: dict[str, Any] | None = None + self.stream: _FakeStream | None = None async def create(self, **kwargs: Any) -> _FakeStream: self.last_kwargs = kwargs - return _FakeStream(self._events, raise_exc=self._raise) + self.stream = _FakeStream(self._events, raise_exc=self._raise) + return self.stream def _provider_with_stream( @@ -148,6 +154,7 @@ async def test_anthropic_stream_generate_reassembles_signed_thinking_blocks() -> assert messages.last_kwargs is not None assert messages.last_kwargs["stream"] is True + assert messages.stream is not None and messages.stream.closed is True state_chunks = [c for c in chunks if c.provider_state is not None] assert len(state_chunks) == 1 diff --git a/tests/providers/test_gemini_stream.py b/tests/providers/test_gemini_stream.py index bf9e8079..84b1c22e 100644 --- a/tests/providers/test_gemini_stream.py +++ b/tests/providers/test_gemini_stream.py @@ -75,6 +75,7 @@ def __init__( ): self._chunks = chunks self._raise = raise_exc + self.closed = False def __aiter__(self) -> _AsyncChunks: return self @@ -86,6 +87,9 @@ async def __anext__(self) -> SimpleNamespace: raise StopAsyncIteration return self._chunks.pop(0) + async def aclose(self) -> None: + self.closed = True + def _provider_with( chunks: list[SimpleNamespace], raise_exc: Exception | None = None @@ -96,7 +100,9 @@ async def fake_stream(*, model: str, contents: Any, config: Any) -> _AsyncChunks captured["model"] = model captured["contents"] = contents captured["config"] = config - return _AsyncChunks(chunks, raise_exc=raise_exc) + stream = _AsyncChunks(chunks, raise_exc=raise_exc) + captured["stream"] = stream + return stream provider = GeminiProvider("test-key") provider._client = SimpleNamespace( @@ -131,6 +137,7 @@ async def test_gemini_stream_through_interaction_assembles_output() -> None: ] assert captured["model"] == GEMINI_MODEL + assert captured["stream"].closed is True types = [e.type for e in events] assert types[0] == "start" diff --git a/tests/providers/test_local_contract.py b/tests/providers/test_local_contract.py index b8c7279e..714744af 100644 --- a/tests/providers/test_local_contract.py +++ b/tests/providers/test_local_contract.py @@ -10,13 +10,13 @@ import pytest from pollux.errors import APIError, ConfigurationError, ToolCallParseError -from pollux.interaction.continuation import Continuation, Message, build_continuation +from pollux.interaction.continuation import Message, build_continuation from pollux.interaction.tools import ToolCall, ToolResult from pollux.providers.local import LocalProvider from tests.conftest import ( LOCAL_MODEL, ) -from tests.helpers import make_interaction +from tests.helpers import make_continuation, make_interaction pytestmark = pytest.mark.contract @@ -313,10 +313,10 @@ async def test_local_reasoning_is_display_only_and_not_replayed() -> None: input_, response, user_content="2+2?", provider="local" ) assert continuation is not None - assistant = continuation.messages[-1] - assert assistant.role == "assistant" - assert assistant.content == "4" - assert assistant.provider_state is None + assistant = continuation.to_jsonable()["messages"][-1] + assert assistant["role"] == "assistant" + assert assistant["content"] == "4" + assert "provider_state" not in assistant assert reasoning_text not in json.dumps(continuation.to_jsonable()) # Replaying the continuation must not send reasoning back to the server. @@ -476,7 +476,8 @@ async def test_local_generate_replays_tool_history() -> None: await provider.generate( *_local( content="", - continuation=Continuation( + continuation=make_continuation( + provider="local", messages=( Message(role="user", content="What's the weather in NYC?"), Message( diff --git a/tests/providers/test_local_stream.py b/tests/providers/test_local_stream.py index f0029308..01234921 100644 --- a/tests/providers/test_local_stream.py +++ b/tests/providers/test_local_stream.py @@ -126,11 +126,13 @@ def json(self) -> Any: class _FakeStreamCM: def __init__(self, response: _FakeStreamResponse) -> None: self._response = response + self.exited = False async def __aenter__(self) -> _FakeStreamResponse: return self._response async def __aexit__(self, *_exc: object) -> bool: + self.exited = True return False @@ -147,13 +149,15 @@ def __init__( self._status_code = status_code self._error_body = error_body self.closed = False + self.last_stream: _FakeStreamCM | None = None def stream(self, method: str, path: str, *, json: dict[str, Any]) -> _FakeStreamCM: del method, path self.last_json = json - return _FakeStreamCM( + self.last_stream = _FakeStreamCM( _FakeStreamResponse(self._lines, self._status_code, self._error_body) ) + return self.last_stream async def aclose(self) -> None: self.closed = True @@ -188,6 +192,7 @@ async def test_local_stream_generate_parses_sse_chunks() -> None: assert any(c.finish_reason == "stop" for c in chunks) assert any(c.usage and c.usage.get("total_tokens") == 5 for c in chunks) assert any(c.response_id == "c1" for c in chunks) + assert fake.last_stream is not None and fake.last_stream.exited is True @pytest.mark.asyncio diff --git a/tests/providers/test_openai_contract.py b/tests/providers/test_openai_contract.py index 5b176d26..669fc307 100644 --- a/tests/providers/test_openai_contract.py +++ b/tests/providers/test_openai_contract.py @@ -7,7 +7,7 @@ import pytest from pollux.errors import APIError, ConfigurationError -from pollux.interaction.continuation import Continuation, Message +from pollux.interaction.continuation import Message from pollux.providers._utils import to_strict_schema from pollux.providers.models import ( ProviderFileAsset, @@ -16,7 +16,7 @@ from tests.conftest import ( OPENAI_MODEL, ) -from tests.helpers import make_interaction +from tests.helpers import make_continuation, make_interaction from tests.providers.helpers import FakeResponses, async_return @@ -252,7 +252,8 @@ async def test_openai_generate_forwards_conversation_and_instructions() -> None: await provider.generate( *_openai( content="And now?", - continuation=Continuation( + continuation=make_continuation( + provider="openai", response_id="resp_123", messages=(Message(role="user", content="This should be skipped."),), ), diff --git a/tests/providers/test_openai_stream.py b/tests/providers/test_openai_stream.py index 00e08661..4f8cdfd0 100644 --- a/tests/providers/test_openai_stream.py +++ b/tests/providers/test_openai_stream.py @@ -76,6 +76,7 @@ def __init__( ): self._events = events self._raise = raise_exc + self.closed = False def __aiter__(self) -> _AsyncEvents: return self @@ -87,6 +88,9 @@ async def __anext__(self) -> SimpleNamespace: raise StopAsyncIteration return self._events.pop(0) + async def close(self) -> None: + self.closed = True + class _FakeResponses: def __init__( @@ -95,10 +99,12 @@ def __init__( self._events = events self._raise = raise_exc self.last_kwargs: dict[str, Any] | None = None + self.stream: _AsyncEvents | None = None async def create(self, **kwargs: Any) -> _AsyncEvents: self.last_kwargs = kwargs - return _AsyncEvents(self._events, raise_exc=self._raise) + self.stream = _AsyncEvents(self._events, raise_exc=self._raise) + return self.stream def _provider_with( @@ -137,6 +143,7 @@ async def test_openai_stream_through_interaction_assembles_output() -> None: assert responses.last_kwargs is not None assert responses.last_kwargs["stream"] is True + assert responses.stream is not None and responses.stream.closed is True types = [e.type for e in events] assert types[0] == "start" diff --git a/tests/providers/test_openrouter_contract.py b/tests/providers/test_openrouter_contract.py index a86aedb3..8818c442 100644 --- a/tests/providers/test_openrouter_contract.py +++ b/tests/providers/test_openrouter_contract.py @@ -9,7 +9,7 @@ import pytest from pollux.errors import APIError, ConfigurationError -from pollux.interaction.continuation import Continuation, Message +from pollux.interaction.continuation import Message from pollux.interaction.tools import ToolCall, ToolResult from pollux.providers.openrouter import ( OpenRouterProvider, @@ -19,7 +19,7 @@ from tests.conftest import ( OPENROUTER_MODEL, ) -from tests.helpers import make_interaction +from tests.helpers import make_continuation, make_interaction def _openrouter(**kwargs: Any) -> tuple[Any, Any, Any, Any]: @@ -264,7 +264,8 @@ async def test_openrouter_generate_characterizes_tool_history_and_schema_shape() await provider.generate( *_openrouter( model="openai/gpt-4.1-mini", - continuation=Continuation( + continuation=make_continuation( + provider="openrouter", messages=( Message(role="user", content="Need orbit code."), Message( @@ -474,10 +475,11 @@ async def test_openrouter_generate_replays_reasoning_only_history() -> None: *_openrouter( model="openai/gpt-4.1-mini", content="Continue.", - continuation=Continuation( - messages=( - Message(role="user", content="Think first."), - Message(role="assistant", content=""), + continuation=make_continuation( + provider="openrouter", + serialized_messages=( + {"role": "user", "content": "Think first."}, + {"role": "assistant", "content": ""}, ), provider_state={ "history": [ diff --git a/tests/providers/test_openrouter_stream.py b/tests/providers/test_openrouter_stream.py index 4dfbb8aa..8467b052 100644 --- a/tests/providers/test_openrouter_stream.py +++ b/tests/providers/test_openrouter_stream.py @@ -84,11 +84,13 @@ def json(self) -> Any: class _FakeStreamCM: def __init__(self, response: _FakeStreamResponse) -> None: self._response = response + self.exited = False async def __aenter__(self) -> _FakeStreamResponse: return self._response async def __aexit__(self, *_exc: object) -> bool: + self.exited = True return False @@ -104,6 +106,7 @@ def __init__( self._lines = lines if lines is not None else _TEXT_STREAM self._status_code = status_code self._error_body = error_body + self.last_stream: _FakeStreamCM | None = None async def get(self, path: str) -> Any: request = httpx.Request("GET", f"{_BASE_URL}{path}") @@ -112,9 +115,10 @@ async def get(self, path: str) -> Any: def stream(self, method: str, path: str, *, json: dict[str, Any]) -> _FakeStreamCM: del method, path self.last_json = json - return _FakeStreamCM( + self.last_stream = _FakeStreamCM( _FakeStreamResponse(self._lines, self._status_code, self._error_body) ) + return self.last_stream async def aclose(self) -> None: return None @@ -144,6 +148,7 @@ async def test_openrouter_stream_generate_parses_sse_and_reasoning() -> None: assert any(c.reasoning == "thinking" for c in chunks) assert any(c.finish_reason == "stop" for c in chunks) assert any(c.usage and c.usage.get("total_tokens") == 5 for c in chunks) + assert fake.last_stream is not None and fake.last_stream.exited is True @pytest.mark.asyncio diff --git a/tests/providers/test_tool_history.py b/tests/providers/test_tool_history.py index c469b619..baa3f936 100644 --- a/tests/providers/test_tool_history.py +++ b/tests/providers/test_tool_history.py @@ -7,7 +7,7 @@ import pytest -from pollux.interaction.continuation import Continuation, Message +from pollux.interaction.continuation import Message from pollux.interaction.tools import ToolCall, ToolResult from pollux.providers.gemini import GeminiProvider from pollux.providers.openai import OpenAIProvider @@ -15,7 +15,7 @@ GEMINI_MODEL, OPENAI_MODEL, ) -from tests.helpers import make_interaction +from tests.helpers import make_continuation, make_interaction from tests.providers.helpers import FakeResponses pytestmark = pytest.mark.contract @@ -48,7 +48,8 @@ async def test_openai_maps_tool_history_to_responses_api_format() -> None: await provider.generate( *_openai( content="Continue the conversation", - continuation=Continuation( + continuation=make_continuation( + provider="openai", messages=( Message(role="user", content="What's the weather?"), Message( @@ -104,7 +105,8 @@ async def test_openai_preserves_assistant_text_with_tool_calls() -> None: await provider.generate( *_openai( content="Continue", - continuation=Continuation( + continuation=make_continuation( + provider="openai", messages=( Message( role="assistant", @@ -142,7 +144,8 @@ async def test_openai_keeps_tool_outputs_when_previous_response_id_is_set() -> N await provider.generate( *_openai( content="Continue", - continuation=Continuation( + continuation=make_continuation( + provider="openai", response_id="resp_prev", messages=( Message(role="user", content="What's the weather?"), @@ -202,7 +205,8 @@ async def fake_generate_content( await provider.generate( *_gemini( content="Continue the conversation", - continuation=Continuation( + continuation=make_continuation( + provider="gemini", messages=( Message(role="user", content="What's the weather?"), Message( @@ -273,7 +277,8 @@ async def fake_generate_content( await provider.generate( *_gemini( content="Proceed.", - continuation=Continuation( + continuation=make_continuation( + provider="gemini", messages=( Message(role="user", content="What's the weather?"), Message( diff --git a/tests/test_api.py b/tests/test_api.py index 5c035d6c..01fc54ef 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -287,12 +287,12 @@ async def test_live_tool_calls_conversation_and_reasoning_roundtrip( assert second.structured.get("secret_code") == "K9-ORBIT" assert second.continuation is not None - messages = second.continuation.messages + messages = second.continuation.to_jsonable()["messages"] assert messages tool_indexes = [ i for i, msg in enumerate(messages) - if msg.role == "tool" and msg.tool_call_id == tool_ref + if msg["role"] == "tool" and msg.get("tool_call_id") == tool_ref ] assert tool_indexes assert tool_indexes[0] < len(messages) - 1