mirror of
https://github.com/pydantic/pydantic-ai-harness.git
synced 2026-07-21 02:45:34 +00:00
feat(experimental): OverflowingToolOutput capability for oversized tool returns (#290)
* feat(experimental): OverflowingToolOutput capability for oversized tool returns An oversized tool return persists in history and is re-sent on every later model request, paying its token cost for the rest of the run. This capability intercepts a return in `after_tool_execute`, reduces it once, and lets the reduced form persist instead of recomputing it per request. Three freely-combinable modes selected by an ordered list of size bands: truncate (cheap clamp), spill (lossless persist + `read_tool_result` handle for on-demand slice/grep/tail), and summarize (size-gated LLM summary inheriting the run's model). The default band is `Spill(then=Truncate())`: lossless when the store accepts the write, a bounded truncation otherwise, never a silent drop. Spilled payloads go through a narrow `OverflowStore` protocol (local-file default) so a durable backend or core #4352's queryable-file primitive can slot in behind the same handle. Mirrors the experimental/compaction package shape, export pattern, and model-inherit path; truncation / ANSI / binary helpers are harvested from PR #185, whose one-way spill this supersedes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(overflow): address review -- store hardening, content overflow, bounds, no-Any, TTL cleanup S1: harden LocalFileStore while preserving cross-agent sharing. Keep a stable, configurable root (so a later run can read an earlier spill), but create it with 0700 perms and verify on read that the resolved target stays within the root, rejecting symlink/`..`/absolute escapes. Sharing is the goal; isolation is not the security mechanism. C1: large ToolReturn.content bypassed reduction -- core renders it as a separate model-visible part that also persists in history. Measure and reduce content with the same band logic (distinct spill handle); non-text content that overflows is left unreduced with a warning. C2: bound read_tool_result -- require offset >= 0 and limit >= 1, clamp limit to a line cap, cap the joined output, and make `pattern` a literal substring instead of a regex so a model-supplied value cannot trigger catastrophic backtracking. A1: drop avoidable Any -- typed tokenizer, narrowed payload/metadata handling, and removed the reportUnknownArgumentType ignores via bool guards. M1: summarize tests now assert the real claims (ctx.model inheritance, explicit model override, usage=ctx.usage threading) with a FunctionModel instead of a wholesale Agent mock. Cleanup feature: opt-in LocalFileStore(cleanup_after=timedelta) prunes files older than the TTL in a background daemon thread; keep-forever stays the default. Pruning is non-blocking and swallows all errors into warnings, never failing a run. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor(overflow): convert remaining avoidable Any to precise types (A1) `_is_mapping`/`_is_text_sequence` are now `TypeGuard`s, so `json_sketch` takes `object` and the sketch helpers take `Mapping[object, object]` / `Sequence[object]` -- no `Any`, no `Unknown`. In the capability, `original` is `object`, `metadata` / `existing` are `object`, `_copy_mapping` / `_with_handles` return `dict[str, object]`, and the assembled `ToolReturn` locals and `_assemble` return type drop `Any` for `ToolReturn[object]` / `object`. The only `Any` left is the core-inherited `after_tool_execute` hook signature (`args`/`result`/return). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: David SF <david.sanchez@pydantic.dev> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
David SF
co-authored by
David SF
Claude Opus 4.8
David SF
Claude Opus 4.8
parent
c56a4b42f8
commit
3ac6034f7e
@@ -0,0 +1,201 @@
|
||||
# Overflow capability
|
||||
|
||||
> [!WARNING]
|
||||
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
|
||||
> change or be removed in any release, without a deprecation period. Import it from the
|
||||
> experimental path -- there is no top-level export:
|
||||
>
|
||||
> ```python
|
||||
> from pydantic_ai_harness.experimental.overflow import OverflowingToolOutput
|
||||
> ```
|
||||
>
|
||||
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
|
||||
> harness experimental warnings with a single filter (no per-capability lines needed):
|
||||
>
|
||||
> ```python
|
||||
> import warnings
|
||||
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
|
||||
>
|
||||
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
|
||||
> ```
|
||||
|
||||
A tool can return a payload large enough to dominate the context window. Tool returns
|
||||
persist in history as `ToolReturnPart`s, so an oversized one is re-sent on every later
|
||||
model request -- paying its token cost for the rest of the run. `OverflowingToolOutput`
|
||||
intercepts a return in the `after_tool_execute` hook, reduces it once, and lets the reduced
|
||||
form persist. The reduction is not recomputed per request.
|
||||
|
||||
This is the overflow-to-file follow-up the `compaction` README names as out of scope: it
|
||||
moves large tool outputs *out* of the window at production time, rather than compressing or
|
||||
dropping context already inside it.
|
||||
|
||||
## The three modes
|
||||
|
||||
| Mode | Cost | Lossy? | What the model gets |
|
||||
|---|---|---|---|
|
||||
| `Truncate` | zero-LLM | yes | A head / tail / head+tail clamp of the text |
|
||||
| `Spill` | zero-LLM | no | A handle + preview + shape sketch; full payload read back on demand |
|
||||
| `Summarize` | one LLM call | yes | A size-gated summary (inherits the run's model by default) |
|
||||
|
||||
`Spill` is lossless: the full payload is persisted and the model reads slices of it through
|
||||
the registered `read_tool_result(handle, offset, limit, from_end, pattern)` tool (the Claude
|
||||
Code pattern, the core [#4352](https://github.com/pydantic/pydantic-ai/issues/4352) design).
|
||||
That tool is bounded: `offset >= 0`, `limit` clamped to a built-in line cap, the joined output
|
||||
capped, and `pattern` is a literal substring (not a regex), so a model-supplied value cannot
|
||||
hang the host with catastrophic backtracking.
|
||||
|
||||
### Both `return_value` and `content` are reduced
|
||||
|
||||
A `ToolReturn` carries a `return_value` and an optional `content` that core renders as a
|
||||
separate, model-visible part which also persists in history. This capability measures and
|
||||
reduces both with the same band logic (they spill to distinct handles). Text `content` is
|
||||
reduced in place; non-text `content` (multimodal parts) that overflows is left unreduced with
|
||||
a `warnings.warn`, since it cannot be safely truncated.
|
||||
|
||||
## Bands: combine the modes
|
||||
|
||||
Configure an ordered list of size `bands`. Each band is a `(over, action)` pair: when a
|
||||
return's measured size reaches `over`, its action runs. The band with the largest threshold
|
||||
that fits wins; anything below the smallest threshold passes through.
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.overflow import (
|
||||
Band,
|
||||
OverflowingToolOutput,
|
||||
Spill,
|
||||
Summarize,
|
||||
Truncate,
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
'openai:gpt-4o',
|
||||
capabilities=[
|
||||
OverflowingToolOutput(
|
||||
bands=[
|
||||
Band(over=100_000, action=Spill()), # huge: keep losslessly, read back on demand
|
||||
Band(over=20_000, action=Summarize()), # large: compress with the run's model
|
||||
Band(over=5_000, action=Truncate()), # medium: cheap clamp
|
||||
],
|
||||
# below 5,000: passthrough
|
||||
)
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
The default band, when you pass no `bands`, is `Spill(then=Truncate())`: lossless when a
|
||||
store accepts the write, a bounded truncation otherwise -- zero LLM cost and no silent drop.
|
||||
|
||||
### Fallbacks with `then`
|
||||
|
||||
Every action takes an optional `then`, applied when the action cannot run: a `Spill` whose
|
||||
store errors, a `Truncate` / `Summarize` on a binary payload, a `Summarize` whose model call
|
||||
raises. `then` chains, so `Summarize(then=Spill(then=Truncate()))` degrades summarize ->
|
||||
spill -> truncate.
|
||||
|
||||
### Per-tool overrides and filtering
|
||||
|
||||
`per_tool` replaces the global band list for named tools (file reads to `head`, logs to
|
||||
`tail`); `tool_filter` (a `ToolSelector`) scopes which tools the capability touches at all.
|
||||
|
||||
```python
|
||||
OverflowingToolOutput(
|
||||
per_tool={
|
||||
'read_file': [Band(over=8_000, action=Truncate(strategy=TruncationStrategy.head))],
|
||||
'run_shell': [Band(over=8_000, action=Truncate(strategy=TruncationStrategy.tail))],
|
||||
},
|
||||
tool_filter=['read_file', 'run_shell', 'search'],
|
||||
)
|
||||
```
|
||||
|
||||
## Size unit
|
||||
|
||||
Thresholds are measured in characters by default. Set `over_tokens=True` to measure in
|
||||
estimated tokens (the same ~4-chars-per-token heuristic as `compaction`); pass a `tokenizer`
|
||||
callable for accuracy. `Truncate.max_chars` is always characters -- truncation is a
|
||||
character operation regardless of the threshold unit.
|
||||
|
||||
## Spill store
|
||||
|
||||
Spilled payloads go through the narrow `OverflowStore` protocol. The default `LocalFileStore`
|
||||
writes one file per `(run_id, tool_call_id, retry)` under a stable root directory and keeps it
|
||||
after the run, so a later `read_tool_result` -- in this run or a subsequent agent/run -- can
|
||||
still reach it. The handle is backend-addressable (a relative key), not an absolute local
|
||||
path, so a durable backend (Temporal, a blob store, or the core `ExecutionEnvironment`
|
||||
workspace once #4352 lands) can resolve the same handle in another process. Supply your own
|
||||
backend with `store=...`.
|
||||
|
||||
```python
|
||||
class OverflowStore(Protocol):
|
||||
async def write(self, key: str, data: bytes) -> str: ... # returns a handle
|
||||
async def read(self, handle: str) -> bytes: ...
|
||||
```
|
||||
|
||||
### Security model (shared root, not isolation)
|
||||
|
||||
The store root is stable and shareable on purpose -- spilled files must be readable by a later
|
||||
agent or run -- so security does not come from per-instance isolation. It comes from two
|
||||
mechanisms: the root is created with `0700` (owner-only) permissions, and `read` resolves the
|
||||
target (following symlinks) and rejects anything that escapes the root via symlink, `..`, or
|
||||
an absolute path. Handle segments are also sanitized so a crafted handle cannot traverse out.
|
||||
|
||||
### Cleanup: keep-forever by default, opt-in TTL pruning
|
||||
|
||||
By default the store keeps spilled files forever -- deleting on run end would break a later
|
||||
agent that still wants to read a spill. To bound disk use, opt into age-based pruning:
|
||||
|
||||
```python
|
||||
from datetime import timedelta
|
||||
|
||||
store = LocalFileStore(cleanup_after=timedelta(hours=6)) # default: None = keep forever
|
||||
```
|
||||
|
||||
When set, a `write` schedules a background prune (a daemon thread, off the hot path) that
|
||||
deletes files whose modification time (`st_mtime`) is older than `cleanup_after`. Pruning is
|
||||
non-blocking and non-erroring: any failure is caught and surfaced via `warnings.warn`, never
|
||||
propagated into the agent run, so cleanup can never fail a run or block the hot path.
|
||||
Last-read time (`st_atime`) is unreliable on `noatime`/`relatime` mounts and is not used.
|
||||
|
||||
Prefer external cleanup (cron, a sweeper) over the in-process TTL? Point it at the store root
|
||||
and delete by mtime:
|
||||
|
||||
```python
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
root = Path('/tmp/pyai_harness_overflow') # or your configured base_dir
|
||||
cutoff = time.time() - 6 * 3600
|
||||
for path in root.rglob('*'):
|
||||
if path.is_file() and path.stat().st_mtime < cutoff:
|
||||
path.unlink(missing_ok=True)
|
||||
```
|
||||
|
||||
## Usage accounting
|
||||
|
||||
A `Summarize` call is a real request to the model, so its full usage -- tokens and the
|
||||
request itself -- folds into the run's `ctx.usage`, exactly like `SummarizingCompaction`. No
|
||||
token caps are imposed on the summary call. A `UsageLimits` request limit will see it.
|
||||
|
||||
## Edge cases
|
||||
|
||||
- Binary returns spill verbatim and are never stringify-truncated; `Truncate` / `Summarize`
|
||||
on binary fall through to `then`.
|
||||
- Structured / nested returns spill (or summarize) by preference -- truncating JSON produces
|
||||
invalid JSON. `Spill` includes a one-line shape sketch of the top level.
|
||||
- `ModelRetry` and tool errors never reach this hook (they are raised, not returned), so the
|
||||
model always gets the full error it needs to recover.
|
||||
- A large `ToolReturn.content` is reduced with the same bands as `return_value`; non-text
|
||||
content that overflows is left unreduced with a warning.
|
||||
- Multiple oversized returns in one step get distinct handles (keyed per `tool_call_id`);
|
||||
retries get distinct handles too (keyed per `retry`), so a retried call never clobbers the
|
||||
earlier attempt's spill.
|
||||
|
||||
## Relationship to other capabilities
|
||||
|
||||
- Supersedes the spill scope of PR #185 `ToolOutputManagement` (one-way truncate / spill with
|
||||
no read-back); this capability's truncation and ANSI / binary handling are harvested from it.
|
||||
- Consumes core [#4352](https://github.com/pydantic/pydantic-ai/issues/4352) (the canonical
|
||||
queryable-file primitive) through the `OverflowStore` seam once it lands.
|
||||
- Distinct from `compaction`, which compresses or drops context already inside the window, and
|
||||
from `ClampOversizedMessages` (PR #286), which clamps runaway model responses, not tool
|
||||
returns.
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Overflow capability: reduce oversized tool returns at production time.
|
||||
|
||||
`OverflowingToolOutput` intercepts a tool return when it is produced and reduces it --
|
||||
truncating, spilling to a queryable file, or summarizing -- so an oversized payload does
|
||||
not persist in history and get re-sent on every later model request. Combine the three
|
||||
modes through an ordered list of size `bands`.
|
||||
|
||||
Spilled payloads are read back on demand through the registered `read_tool_result` tool;
|
||||
the `OverflowStore` protocol is the seam for a durable backend (the local-file default
|
||||
ships for single-process runs).
|
||||
"""
|
||||
|
||||
from pydantic_ai_harness.experimental._warn import warn_experimental
|
||||
from pydantic_ai_harness.experimental.overflow._bands import (
|
||||
Action,
|
||||
Band,
|
||||
Passthrough,
|
||||
Spill,
|
||||
Summarize,
|
||||
SummarizeFunc,
|
||||
Truncate,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.overflow._capability import READ_TOOL_NAME, OverflowingToolOutput
|
||||
from pydantic_ai_harness.experimental.overflow._payload import TruncationStrategy
|
||||
from pydantic_ai_harness.experimental.overflow._store import LocalFileStore, OverflowStore
|
||||
|
||||
warn_experimental('overflow')
|
||||
|
||||
__all__ = [
|
||||
'READ_TOOL_NAME',
|
||||
'Action',
|
||||
'Band',
|
||||
'LocalFileStore',
|
||||
'OverflowStore',
|
||||
'OverflowingToolOutput',
|
||||
'Passthrough',
|
||||
'Spill',
|
||||
'Summarize',
|
||||
'SummarizeFunc',
|
||||
'Truncate',
|
||||
'TruncationStrategy',
|
||||
]
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Size bands and the actions they trigger.
|
||||
|
||||
A band is a `(over, action)` pair: when a tool return's measured size is at least `over`,
|
||||
its `action` runs. `OverflowingToolOutput` holds an ordered band list and picks the first
|
||||
match (largest threshold that fits), passing through anything below the smallest threshold.
|
||||
|
||||
Every action carries an optional `then` fallback, applied when the action cannot run --
|
||||
`Spill` whose store errors, `Truncate`/`Summarize` on a binary payload, a `Summarize`
|
||||
whose model call raises. `Spill(then=Truncate())` is the default: lossless when the store
|
||||
works, a bounded truncation otherwise, never a silent drop.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic_ai_harness.experimental.overflow._payload import TruncationStrategy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic_ai.models import Model
|
||||
|
||||
_DEFAULT_TRUNCATE_CHARS = 4_000
|
||||
_DEFAULT_PREVIEW_CHARS = 1_000
|
||||
|
||||
# A summarizer callable: `(tool_name, full_text) -> summary`, sync or async.
|
||||
SummarizeFunc = Callable[[str, str], str | Awaitable[str]]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Passthrough:
|
||||
"""Leave the tool return untouched. Useful as an explicit no-op band."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Truncate:
|
||||
"""Clamp the stringified return to `max_chars`. Lossy, zero-cost, no read-back.
|
||||
|
||||
`max_chars` is always characters, independent of the capability's `over_tokens` size
|
||||
unit (truncation is a character operation). Falls back to `then` for binary payloads,
|
||||
which cannot be stringify-truncated.
|
||||
"""
|
||||
|
||||
strategy: TruncationStrategy = TruncationStrategy.head_tail
|
||||
max_chars: int = _DEFAULT_TRUNCATE_CHARS
|
||||
then: Action | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Spill:
|
||||
"""Persist the full return and replace it with a handle, preview, and shape sketch.
|
||||
|
||||
Lossless: the model gets a `read_tool_result` handle to slice / grep / tail the full
|
||||
payload on demand. Falls back to `then` when no store accepts the write.
|
||||
"""
|
||||
|
||||
preview_chars: int = _DEFAULT_PREVIEW_CHARS
|
||||
then: Action | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Summarize:
|
||||
"""Replace the return with a size-gated LLM summary. The expensive band.
|
||||
|
||||
`model=None` inherits the running agent's model (`ctx.model`), mirroring
|
||||
`SummarizingCompaction`. Pass a model id / instance to override, or a `summarize`
|
||||
callable to bypass the built-in prompt entirely. Summary usage folds into `ctx.usage`;
|
||||
no token caps are imposed. Falls back to `then` on a binary payload or a failed call.
|
||||
"""
|
||||
|
||||
model: str | Model | None = None
|
||||
summarize: SummarizeFunc | None = None
|
||||
then: Action | None = None
|
||||
|
||||
|
||||
Action = Passthrough | Truncate | Spill | Summarize
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Band:
|
||||
"""Trigger `action` once a return's measured size reaches `over` (chars or tokens)."""
|
||||
|
||||
over: int
|
||||
action: Action
|
||||
@@ -0,0 +1,552 @@
|
||||
"""`OverflowingToolOutput` -- reduce oversized tool returns at production time."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, TypeGuard
|
||||
|
||||
from pydantic_ai import FunctionToolset
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.messages import ToolCallPart, ToolReturn, ToolReturnContent, UserContent
|
||||
from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition, ToolSelector, matches_tool_selector
|
||||
from pydantic_ai.toolsets import AgentToolset
|
||||
|
||||
from pydantic_ai_harness.experimental.overflow._bands import (
|
||||
Action,
|
||||
Band,
|
||||
Passthrough,
|
||||
Spill,
|
||||
Summarize,
|
||||
Truncate,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.overflow._payload import (
|
||||
is_binary,
|
||||
json_sketch,
|
||||
measure,
|
||||
strip_ansi,
|
||||
to_bytes,
|
||||
to_text,
|
||||
truncate_text,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.overflow._store import LocalFileStore, OverflowStore
|
||||
|
||||
READ_TOOL_NAME = 'read_tool_result'
|
||||
"""Name of the registered read-back tool. Its own returns are exempt from reduction."""
|
||||
|
||||
_DEFAULT_THRESHOLD = 10_000
|
||||
"""Default band threshold (characters) -- below this, returns pass through untouched."""
|
||||
|
||||
_DEFAULT_SUMMARY_PROMPT = """\
|
||||
The following output from the `{tool_name}` tool is too large to keep in full. Summarize it \
|
||||
so the summary carries everything needed to keep working: concrete values, identifiers, \
|
||||
errors, and structure. Respond ONLY with the summary, no preamble.
|
||||
|
||||
<output>
|
||||
{output}
|
||||
</output>\
|
||||
"""
|
||||
|
||||
|
||||
def _default_bands() -> list[Band]:
|
||||
"""Lossless spill with a bounded truncation fallback: zero LLM cost, no silent drop."""
|
||||
return [Band(over=_DEFAULT_THRESHOLD, action=Spill(then=Truncate()))]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Unit:
|
||||
"""One reducible piece of a tool return: its `return_value` or its `content`.
|
||||
|
||||
`suffix` distinguishes the two so they spill to distinct handles for the same call.
|
||||
"""
|
||||
|
||||
binary: bool
|
||||
text: str | None
|
||||
data: bytes
|
||||
value: ToolReturnContent
|
||||
suffix: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class OverflowingToolOutput(AbstractCapability[AgentDepsT]):
|
||||
"""Reduce oversized tool returns when they are produced, persisting the reduction.
|
||||
|
||||
A tool can return a payload large enough to dominate the context window. Tool returns
|
||||
persist in history, so an oversized one is re-sent on every later request. This
|
||||
capability intercepts a return in `after_tool_execute`, reduces it once, and lets the
|
||||
reduced form persist -- it is not recomputed per request.
|
||||
|
||||
Three reduction modes, freely combined through an ordered list of size `bands`:
|
||||
|
||||
- `Truncate`: clamp to a character budget. Lossy, zero-cost.
|
||||
- `Spill`: persist the full payload, hand the model a `read_tool_result` handle plus a
|
||||
preview. Lossless.
|
||||
- `Summarize`: size-gated LLM summary. Inherits the run's model by default.
|
||||
|
||||
The first band whose `over` threshold the measured size meets wins; smaller returns pass
|
||||
through. `per_tool` replaces the band list for named tools; `tool_filter` scopes which
|
||||
tools are touched at all. The default is `Spill(then=Truncate())`: lossless when a store
|
||||
accepts the write, a bounded truncation otherwise.
|
||||
|
||||
`ModelRetry` and other errors never reach this hook (they are raised, not returned), so
|
||||
error payloads the model needs to recover are never spilled or summarized.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.overflow import (
|
||||
Band,
|
||||
OverflowingToolOutput,
|
||||
Spill,
|
||||
Summarize,
|
||||
Truncate,
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
'openai:gpt-4o',
|
||||
capabilities=[
|
||||
OverflowingToolOutput(
|
||||
bands=[
|
||||
Band(over=100_000, action=Spill()),
|
||||
Band(over=20_000, action=Summarize()),
|
||||
Band(over=5_000, action=Truncate()),
|
||||
],
|
||||
)
|
||||
],
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
bands: Sequence[Band] = field(default_factory=_default_bands)
|
||||
"""Ordered size bands. The first band whose `over` threshold is met wins."""
|
||||
|
||||
per_tool: Mapping[str, Sequence[Band]] = field(default_factory=dict[str, Sequence[Band]])
|
||||
"""Per-tool band lists that replace `bands` for the named tools."""
|
||||
|
||||
tool_filter: ToolSelector[AgentDepsT] = 'all'
|
||||
"""Which tools this capability touches. Non-matching tools always pass through."""
|
||||
|
||||
over_tokens: bool = False
|
||||
"""Measure band thresholds in estimated tokens instead of characters."""
|
||||
|
||||
tokenizer: Callable[[str], int] | None = None
|
||||
"""Optional `(str) -> int` tokenizer for `over_tokens`. Defaults to a ~4-char heuristic."""
|
||||
|
||||
store: OverflowStore | None = None
|
||||
"""Backend for spilled payloads. Defaults to a `LocalFileStore`."""
|
||||
|
||||
strip_ansi: bool = False
|
||||
"""Strip ANSI escape sequences from text returns before measuring and reducing."""
|
||||
|
||||
summary_prompt: str = _DEFAULT_SUMMARY_PROMPT
|
||||
"""Prompt template for `Summarize`. Must contain `{tool_name}` and `{output}`."""
|
||||
|
||||
_store: OverflowStore = field(init=False, repr=False)
|
||||
_bands: list[Band] = field(init=False, repr=False)
|
||||
_per_tool: dict[str, list[Band]] = field(init=False, repr=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._store = self.store if self.store is not None else LocalFileStore()
|
||||
self._bands = self._prepare_bands(self.bands)
|
||||
self._per_tool = {name: self._prepare_bands(bands) for name, bands in self.per_tool.items()}
|
||||
|
||||
@staticmethod
|
||||
def _prepare_bands(bands: Sequence[Band]) -> list[Band]:
|
||||
"""Validate thresholds and order bands largest-first so first-match means largest-fit."""
|
||||
for band in bands:
|
||||
if band.over < 0:
|
||||
raise ValueError('Band.over must be non-negative.')
|
||||
return sorted(bands, key=lambda b: b.over, reverse=True)
|
||||
|
||||
# --- toolset ---
|
||||
|
||||
def get_toolset(self) -> AgentToolset[AgentDepsT] | None:
|
||||
"""Register the `read_tool_result` tool for reading spilled payloads on demand."""
|
||||
store = self._store
|
||||
|
||||
async def read_tool_result(
|
||||
ctx: RunContext[AgentDepsT],
|
||||
handle: str,
|
||||
offset: int = 0,
|
||||
limit: int = 200,
|
||||
from_end: bool = False,
|
||||
pattern: str | None = None,
|
||||
) -> str:
|
||||
"""Read a slice of a spilled tool result.
|
||||
|
||||
Args:
|
||||
ctx: The run context (supplied by the agent).
|
||||
handle: The handle from the overflowed tool return.
|
||||
offset: Number of matching lines to skip from the start (or end). Must be >= 0.
|
||||
limit: Maximum number of lines to return (>= 1; clamped to a built-in cap).
|
||||
from_end: Count `offset`/`limit` from the end of the result.
|
||||
pattern: Optional literal substring; only lines containing it are returned.
|
||||
"""
|
||||
return await _read_slice(store, handle, offset, limit, from_end, pattern)
|
||||
|
||||
return FunctionToolset([read_tool_result])
|
||||
|
||||
# --- reduction ---
|
||||
|
||||
async def after_tool_execute(
|
||||
self,
|
||||
ctx: RunContext[AgentDepsT],
|
||||
*,
|
||||
call: ToolCallPart,
|
||||
tool_def: ToolDefinition,
|
||||
args: dict[str, Any],
|
||||
result: Any,
|
||||
) -> Any:
|
||||
"""Reduce the tool result -- both `return_value` and model-visible `content`."""
|
||||
original: object = result
|
||||
if call.tool_name == READ_TOOL_NAME:
|
||||
return original
|
||||
if not await matches_tool_selector(self.tool_filter, ctx, tool_def):
|
||||
return original
|
||||
|
||||
metadata: object
|
||||
if isinstance(result, ToolReturn):
|
||||
return_value: ToolReturnContent = result.return_value
|
||||
content = result.content
|
||||
metadata = result.metadata
|
||||
wrapped = True
|
||||
else:
|
||||
return_value = result
|
||||
content = None
|
||||
metadata = None
|
||||
wrapped = False
|
||||
|
||||
if isinstance(return_value, BaseException):
|
||||
return original
|
||||
|
||||
bands = self._per_tool.get(call.tool_name, self._bands)
|
||||
value_unit = self._make_unit(return_value, suffix='')
|
||||
value_text, value_handle = await self._reduce(ctx, call, bands, value_unit)
|
||||
content_text, content_handle = await self._reduce_content(ctx, call, bands, content)
|
||||
|
||||
if value_text is None and content_text is None:
|
||||
return original
|
||||
|
||||
return self._assemble(
|
||||
wrapped=wrapped,
|
||||
return_value=return_value,
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
value_unit=value_unit,
|
||||
value_text=value_text,
|
||||
value_handle=value_handle,
|
||||
content_text=content_text,
|
||||
content_handle=content_handle,
|
||||
)
|
||||
|
||||
def _assemble(
|
||||
self,
|
||||
*,
|
||||
wrapped: bool,
|
||||
return_value: ToolReturnContent,
|
||||
content: str | Sequence[UserContent] | None,
|
||||
metadata: object,
|
||||
value_unit: _Unit,
|
||||
value_text: str | None,
|
||||
value_handle: str | None,
|
||||
content_text: str | None,
|
||||
content_handle: str | None,
|
||||
) -> object:
|
||||
"""Rebuild the tool result from the reduced parts, preserving the envelope."""
|
||||
if wrapped:
|
||||
new_metadata = metadata
|
||||
if value_handle is not None or content_handle is not None:
|
||||
new_metadata = _with_handles(metadata, value_handle, len(value_unit.data), content_handle)
|
||||
wrapped_out: ToolReturn[object] = ToolReturn(
|
||||
return_value=value_text if value_text is not None else return_value,
|
||||
content=content_text if content_text is not None else content,
|
||||
metadata=new_metadata,
|
||||
)
|
||||
return wrapped_out
|
||||
|
||||
# A plain (non-`ToolReturn`) result has no separate content part.
|
||||
if value_handle is not None:
|
||||
spilled_out: ToolReturn[object] = ToolReturn(
|
||||
return_value=value_text, metadata=_with_handles(None, value_handle, len(value_unit.data))
|
||||
)
|
||||
return spilled_out
|
||||
return value_text
|
||||
|
||||
def _make_unit(self, value: ToolReturnContent, *, suffix: str) -> _Unit:
|
||||
"""Pre-render a value into the text / bytes the reduction pipeline needs."""
|
||||
if is_binary(value):
|
||||
return _Unit(binary=True, text=None, data=to_bytes(value), value=value, suffix=suffix)
|
||||
text = to_text(value)
|
||||
if self.strip_ansi:
|
||||
text = strip_ansi(text)
|
||||
return _Unit(binary=False, text=text, data=text.encode('utf-8'), value=value, suffix=suffix)
|
||||
|
||||
async def _reduce(
|
||||
self,
|
||||
ctx: RunContext[AgentDepsT],
|
||||
call: ToolCallPart,
|
||||
bands: Sequence[Band],
|
||||
unit: _Unit,
|
||||
) -> tuple[str | None, str | None]:
|
||||
"""Select a band for `unit` and apply it. Returns `(replacement, handle)`.
|
||||
|
||||
`replacement` is None when the unit passes through unchanged; `handle` is set only
|
||||
when the unit was spilled.
|
||||
"""
|
||||
size = (
|
||||
len(unit.data)
|
||||
if unit.binary
|
||||
else measure(unit.text or '', over_tokens=self.over_tokens, tokenizer=self.tokenizer)
|
||||
)
|
||||
action = _select_action(bands, size)
|
||||
if action is None:
|
||||
return None, None
|
||||
return await self._apply(ctx, call, action, unit)
|
||||
|
||||
async def _reduce_content(
|
||||
self,
|
||||
ctx: RunContext[AgentDepsT],
|
||||
call: ToolCallPart,
|
||||
bands: Sequence[Band],
|
||||
content: str | Sequence[UserContent] | None,
|
||||
) -> tuple[str | None, str | None]:
|
||||
"""Reduce model-visible `content`. Text content is reduced; other content warns."""
|
||||
if content is None:
|
||||
return None, None
|
||||
if isinstance(content, str):
|
||||
return await self._reduce(ctx, call, bands, self._make_unit(content, suffix='.content'))
|
||||
|
||||
text = ''.join(part for part in content if isinstance(part, str))
|
||||
size = measure(text, over_tokens=self.over_tokens, tokenizer=self.tokenizer)
|
||||
action = _select_action(bands, size)
|
||||
if action is not None and not isinstance(action, Passthrough):
|
||||
warnings.warn(
|
||||
f'OverflowingToolOutput: tool {call.tool_name!r} returned large non-text '
|
||||
f'content ({len(content)} parts); leaving it unreduced.',
|
||||
stacklevel=2,
|
||||
)
|
||||
return None, None
|
||||
|
||||
async def _apply(
|
||||
self,
|
||||
ctx: RunContext[AgentDepsT],
|
||||
call: ToolCallPart,
|
||||
action: Action,
|
||||
unit: _Unit,
|
||||
) -> tuple[str | None, str | None]:
|
||||
"""Apply one action to a unit, falling back to its `then` when it cannot run."""
|
||||
if isinstance(action, Passthrough):
|
||||
return None, None
|
||||
|
||||
if isinstance(action, Truncate):
|
||||
if unit.binary:
|
||||
return await self._fallback(ctx, call, action.then, unit)
|
||||
assert unit.text is not None
|
||||
return truncate_text(unit.text, action.max_chars, action.strategy), None
|
||||
|
||||
if isinstance(action, Spill):
|
||||
return await self._spill(ctx, call, action, unit)
|
||||
|
||||
return await self._summarize_action(ctx, call, action, unit)
|
||||
|
||||
async def _fallback(
|
||||
self,
|
||||
ctx: RunContext[AgentDepsT],
|
||||
call: ToolCallPart,
|
||||
then: Action | None,
|
||||
unit: _Unit,
|
||||
) -> tuple[str | None, str | None]:
|
||||
"""Run the fallback action, or keep the unit unchanged when there is none."""
|
||||
if then is None:
|
||||
return None, None
|
||||
return await self._apply(ctx, call, then, unit)
|
||||
|
||||
async def _spill(
|
||||
self,
|
||||
ctx: RunContext[AgentDepsT],
|
||||
call: ToolCallPart,
|
||||
action: Spill,
|
||||
unit: _Unit,
|
||||
) -> tuple[str | None, str | None]:
|
||||
key = _handle_key(ctx, call, unit.suffix)
|
||||
try:
|
||||
handle = await self._store.write(key, unit.data)
|
||||
except Exception:
|
||||
return await self._fallback(ctx, call, action.then, unit)
|
||||
|
||||
preview = _build_spill_preview(handle, unit, action.preview_chars, over_tokens=self.over_tokens)
|
||||
return preview, handle
|
||||
|
||||
async def _summarize_action(
|
||||
self,
|
||||
ctx: RunContext[AgentDepsT],
|
||||
call: ToolCallPart,
|
||||
action: Summarize,
|
||||
unit: _Unit,
|
||||
) -> tuple[str | None, str | None]:
|
||||
if unit.binary:
|
||||
return await self._fallback(ctx, call, action.then, unit)
|
||||
assert unit.text is not None
|
||||
try:
|
||||
summary = await self._summarize(ctx, call, action, unit.text)
|
||||
except Exception:
|
||||
return await self._fallback(ctx, call, action.then, unit)
|
||||
return summary, None
|
||||
|
||||
async def _summarize(
|
||||
self,
|
||||
ctx: RunContext[AgentDepsT],
|
||||
call: ToolCallPart,
|
||||
action: Summarize,
|
||||
text: str,
|
||||
) -> str:
|
||||
"""Generate the summary via a custom callable or the inherited-model agent."""
|
||||
if action.summarize is not None:
|
||||
outcome = action.summarize(call.tool_name, text)
|
||||
if isinstance(outcome, Awaitable):
|
||||
return await outcome
|
||||
return outcome
|
||||
|
||||
from pydantic_ai import Agent
|
||||
|
||||
model = action.model if action.model is not None else ctx.model
|
||||
prompt = self.summary_prompt.format(tool_name=call.tool_name, output=text)
|
||||
agent: Agent[None, str] = Agent(model, instructions='You summarize oversized tool output.')
|
||||
run = await agent.run(prompt, usage=ctx.usage)
|
||||
return run.output.strip()
|
||||
|
||||
|
||||
def _select_action(bands: Sequence[Band], size: int) -> Action | None:
|
||||
"""Return the first (largest-threshold) band action whose threshold `size` meets."""
|
||||
for band in bands:
|
||||
if size >= band.over:
|
||||
return band.action
|
||||
return None
|
||||
|
||||
|
||||
def _handle_key(ctx: RunContext[AgentDepsT], call: ToolCallPart, suffix: str = '') -> str:
|
||||
"""Build a per-run, per-call, per-retry key so concurrent and retried calls never clash.
|
||||
|
||||
`suffix` keeps a return's `return_value` and `content` spills on distinct handles.
|
||||
"""
|
||||
run_id = ctx.run_id or 'run'
|
||||
call_id = call.tool_call_id or 'call'
|
||||
return f'{run_id}/{call_id}.{ctx.retry}{suffix}'
|
||||
|
||||
|
||||
def _is_mapping(value: object) -> TypeGuard[Mapping[object, object]]:
|
||||
"""`TypeGuard` so a mapping narrows to a known element type, not `Unknown`."""
|
||||
return isinstance(value, Mapping)
|
||||
|
||||
|
||||
def _with_handles(
|
||||
existing: object,
|
||||
value_handle: str | None,
|
||||
value_bytes: int,
|
||||
content_handle: str | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""Stash spill handle(s) in `ToolReturn.metadata` (app-only, costs no model tokens)."""
|
||||
base: dict[str, object] = {}
|
||||
if _is_mapping(existing):
|
||||
base.update(_copy_mapping(existing))
|
||||
if value_handle is not None:
|
||||
base['overflow_handle'] = value_handle
|
||||
base['overflow_bytes'] = value_bytes
|
||||
if content_handle is not None:
|
||||
base['overflow_content_handle'] = content_handle
|
||||
return base
|
||||
|
||||
|
||||
def _copy_mapping(source: Mapping[object, object]) -> dict[str, object]:
|
||||
"""Copy an arbitrary mapping with stringified keys (tool metadata is app-defined)."""
|
||||
return {str(key): source[key] for key in source}
|
||||
|
||||
|
||||
def _build_spill_preview(handle: str, unit: _Unit, preview_chars: int, *, over_tokens: bool) -> str:
|
||||
"""Compose the model-visible spill stand-in: marker, sketch, and a head/tail preview."""
|
||||
if unit.binary:
|
||||
size_desc = f'{len(unit.data):,} bytes (binary)'
|
||||
body = f'<{len(unit.data):,} bytes of binary data>'
|
||||
sketch = ''
|
||||
else:
|
||||
text = unit.text or ''
|
||||
size_unit = 'tokens' if over_tokens else 'chars'
|
||||
amount = measure(text, over_tokens=over_tokens, tokenizer=None) if over_tokens else len(text)
|
||||
size_desc = f'{amount:,} {size_unit}'
|
||||
body = _head_tail_preview(text, preview_chars)
|
||||
sketch = json_sketch(unit.value)
|
||||
|
||||
header = (
|
||||
f'[Tool output too large ({size_desc}); stored to handle {handle!r}. '
|
||||
f'Read it with read_tool_result(handle={handle!r}, offset=0, limit=200, '
|
||||
f'from_end=False, pattern=None).]'
|
||||
)
|
||||
parts = [header]
|
||||
if sketch:
|
||||
parts.append(f'shape: {sketch}')
|
||||
parts.append(body)
|
||||
return '\n'.join(parts)
|
||||
|
||||
|
||||
def _head_tail_preview(text: str, preview_chars: int) -> str:
|
||||
"""Return a head+tail slice of `text` with a middle-elision marker."""
|
||||
if len(text) <= preview_chars:
|
||||
return text
|
||||
head_chars = preview_chars // 2
|
||||
tail_chars = preview_chars - head_chars
|
||||
omitted = len(text) - head_chars - tail_chars
|
||||
return f'{text[:head_chars]}\n...[{omitted:,} chars omitted]...\n{text[-tail_chars:]}'
|
||||
|
||||
|
||||
_MAX_READ_LINES = 1_000
|
||||
"""Hard cap on lines returned by one `read_tool_result` call."""
|
||||
|
||||
_MAX_READ_CHARS = 50_000
|
||||
"""Hard cap on characters returned by one `read_tool_result` call."""
|
||||
|
||||
|
||||
async def _read_slice(
|
||||
store: OverflowStore,
|
||||
handle: str,
|
||||
offset: int,
|
||||
limit: int,
|
||||
from_end: bool,
|
||||
pattern: str | None,
|
||||
) -> str:
|
||||
"""Filter and slice a spilled payload for `read_tool_result`, bounded in both axes.
|
||||
|
||||
`pattern` is a literal substring (not a regex), so a model-supplied value cannot hang
|
||||
the host with catastrophic backtracking. `limit` is clamped and the joined output is
|
||||
capped, so one call can never return an unbounded amount of text.
|
||||
"""
|
||||
if offset < 0:
|
||||
raise ModelRetry('`offset` must be >= 0.')
|
||||
if limit < 1:
|
||||
raise ModelRetry('`limit` must be >= 1.')
|
||||
limit = min(limit, _MAX_READ_LINES)
|
||||
|
||||
try:
|
||||
data = await store.read(handle)
|
||||
except OSError as exc:
|
||||
raise ModelRetry(f'No stored tool result for handle {handle!r}: {exc}.') from exc
|
||||
|
||||
lines = data.decode('utf-8', errors='replace').splitlines()
|
||||
if pattern is not None:
|
||||
lines = [line for line in lines if pattern in line]
|
||||
|
||||
total = len(lines)
|
||||
if from_end:
|
||||
end = max(0, total - offset)
|
||||
window = lines[max(0, end - limit) : end]
|
||||
else:
|
||||
window = lines[offset : offset + limit]
|
||||
|
||||
body = '\n'.join(window)
|
||||
capped = ''
|
||||
if len(body) > _MAX_READ_CHARS:
|
||||
body = body[:_MAX_READ_CHARS]
|
||||
capped = ', output capped'
|
||||
header = f'[handle {handle!r}: {total:,} matching line(s); showing {len(window)}{capped}]'
|
||||
return f'{header}\n{body}' if body else header
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Size measurement, stringification, truncation, and binary detection.
|
||||
|
||||
Harvested from PR #185 (`ToolOutputManagement`) and adapted: character-based truncation
|
||||
strategies, ANSI stripping, and binary detection. Token measurement reuses the compaction
|
||||
heuristic via `estimate_token_count` so the two capabilities stay aligned.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from enum import Enum
|
||||
from typing import TypeGuard
|
||||
|
||||
from pydantic_ai.messages import ModelMessage, ModelRequest, SystemPromptPart
|
||||
from pydantic_core import to_json
|
||||
|
||||
from pydantic_ai_harness.experimental.compaction._shared import estimate_token_count
|
||||
|
||||
|
||||
class TruncationStrategy(str, Enum):
|
||||
"""Which end(s) of an oversized text to keep when truncating."""
|
||||
|
||||
head = 'head'
|
||||
"""Keep the first characters (good for headers / schemas)."""
|
||||
|
||||
tail = 'tail'
|
||||
"""Keep the last characters (good for build / test output, where errors land last)."""
|
||||
|
||||
head_tail = 'head_tail'
|
||||
"""Keep the first and last characters, eliding the middle."""
|
||||
|
||||
|
||||
# CSI sequences, OSC sequences, and simple escapes. Terminal tool output is full of color
|
||||
# codes that waste tokens and can confuse models.
|
||||
_ANSI_ESCAPE_RE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07|\x1b[^[\]()]')
|
||||
|
||||
|
||||
def strip_ansi(text: str) -> str:
|
||||
"""Remove ANSI escape sequences from `text`."""
|
||||
return _ANSI_ESCAPE_RE.sub('', text)
|
||||
|
||||
|
||||
def is_binary(value: object) -> bool:
|
||||
"""Return True for raw byte payloads, which must never be stringify-truncated."""
|
||||
return isinstance(value, (bytes, bytearray, memoryview))
|
||||
|
||||
|
||||
def to_bytes(value: object) -> bytes:
|
||||
"""Serialize any tool return value to the bytes that get spilled.
|
||||
|
||||
Strings spill as UTF-8 text; byte payloads spill verbatim; everything else spills as
|
||||
JSON so the stored payload stays valid and grep-able.
|
||||
"""
|
||||
if isinstance(value, str):
|
||||
return value.encode('utf-8')
|
||||
if isinstance(value, memoryview):
|
||||
return value.tobytes()
|
||||
if isinstance(value, (bytes, bytearray)):
|
||||
return bytes(value)
|
||||
return to_json(value)
|
||||
|
||||
|
||||
def to_text(value: object) -> str:
|
||||
"""Render a non-binary tool return value as the text used for measuring and truncating.
|
||||
|
||||
Strings pass through; structured values become JSON (truncating JSON is lossy, so prefer
|
||||
spill or summarize for them -- see the README).
|
||||
"""
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
return to_json(value).decode('utf-8', errors='replace')
|
||||
|
||||
|
||||
def measure(text: str, *, over_tokens: bool, tokenizer: Callable[[str], int] | None) -> int:
|
||||
"""Measure `text` in characters (default) or estimated tokens (`over_tokens=True`)."""
|
||||
if not over_tokens:
|
||||
return len(text)
|
||||
message: ModelMessage = ModelRequest(parts=[SystemPromptPart(content=text)])
|
||||
return estimate_token_count([message], tokenizer)
|
||||
|
||||
|
||||
def json_sketch(value: object) -> str:
|
||||
"""Build a one-line shape hint for a structured value, or '' for anything else.
|
||||
|
||||
The `_is_*` guards are `TypeGuard`s, so a `Mapping`/`Sequence` value narrows to a known
|
||||
element type (`object`) the strict type checker accepts -- no `Any` and no `Unknown`.
|
||||
"""
|
||||
if _is_mapping(value):
|
||||
return _sketch_mapping(value)
|
||||
if _is_text_sequence(value):
|
||||
return _sketch_sequence(value)
|
||||
return ''
|
||||
|
||||
|
||||
def _is_mapping(value: object) -> TypeGuard[Mapping[object, object]]:
|
||||
return isinstance(value, Mapping)
|
||||
|
||||
|
||||
def _is_text_sequence(value: object) -> TypeGuard[Sequence[object]]:
|
||||
return isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray))
|
||||
|
||||
|
||||
def _sketch_mapping(mapping: Mapping[object, object]) -> str:
|
||||
keys = list(mapping)
|
||||
shown = ', '.join(f'{key!r}: {type(mapping[key]).__name__}' for key in keys[:10])
|
||||
more = '' if len(keys) <= 10 else f', ... ({len(keys)} keys)'
|
||||
return f'{{{shown}{more}}}'
|
||||
|
||||
|
||||
def _sketch_sequence(items: Sequence[object]) -> str:
|
||||
elem = type(items[0]).__name__ if items else 'empty'
|
||||
return f'[{len(items)} items of {elem}]'
|
||||
|
||||
|
||||
def truncate_text(text: str, max_chars: int, strategy: TruncationStrategy) -> str:
|
||||
"""Cut `text` down to roughly `max_chars`, annotating what was removed.
|
||||
|
||||
Returns `text` unchanged when it already fits.
|
||||
"""
|
||||
total = len(text)
|
||||
if total <= max_chars:
|
||||
return text
|
||||
|
||||
if strategy is TruncationStrategy.head:
|
||||
return f'{text[:max_chars]}\n\n[truncated: showing first {max_chars:,} of {total:,} chars]'
|
||||
if strategy is TruncationStrategy.tail:
|
||||
return f'[truncated: showing last {max_chars:,} of {total:,} chars]\n\n{text[-max_chars:]}'
|
||||
|
||||
head_chars = max_chars * 2 // 5
|
||||
tail_chars = max_chars - head_chars
|
||||
omitted = total - head_chars - tail_chars
|
||||
return (
|
||||
f'{text[:head_chars]}\n\n'
|
||||
f'[truncated: {omitted:,} chars omitted from the middle; '
|
||||
f'showing first {head_chars:,} + last {tail_chars:,} of {total:,} chars]\n\n'
|
||||
f'{text[-tail_chars:]}'
|
||||
)
|
||||
@@ -0,0 +1,153 @@
|
||||
"""Storage backend for spilled tool outputs.
|
||||
|
||||
`OverflowStore` is a narrow protocol: persist a payload under a key, read it back by
|
||||
handle. `LocalFileStore` is the dependency-free default -- it writes each payload to a
|
||||
file under a stable root directory. The handle is backend-addressable (a relative key),
|
||||
not an absolute local path, so a durable backend (Temporal, a blob store) can resolve the
|
||||
same handle in another process. This is the seam for consuming the core queryable-file
|
||||
primitive (pydantic-ai #4352 / `ExecutionEnvironment`) once it lands.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import warnings
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class OverflowStore(Protocol):
|
||||
"""Persist and retrieve spilled tool-output payloads.
|
||||
|
||||
`write` takes a caller-chosen `key` and returns a `handle`. The handle is the only
|
||||
thing a later `read` needs, so it must be self-contained (a backend can encode the
|
||||
run, the call, and the retry into it). Implementations may return the key unchanged.
|
||||
"""
|
||||
|
||||
async def write(self, key: str, data: bytes) -> str:
|
||||
"""Persist `data` under `key` and return a handle that `read` accepts."""
|
||||
... # pragma: no cover
|
||||
|
||||
async def read(self, handle: str) -> bytes:
|
||||
"""Return the payload previously stored for `handle`.
|
||||
|
||||
Raise `FileNotFoundError` (or another `OSError`) when the handle is unknown.
|
||||
"""
|
||||
... # pragma: no cover
|
||||
|
||||
|
||||
_UNSAFE_SEGMENT = re.compile(r'[^A-Za-z0-9._-]+')
|
||||
|
||||
|
||||
def _safe_segment(segment: str) -> str:
|
||||
"""Make one path segment filesystem-safe without collapsing distinct keys.
|
||||
|
||||
Empty or dot-only segments are replaced so a handle can never escape the root via
|
||||
`.`/`..`; the resolve-within-root check in `read` is the second line of defense.
|
||||
"""
|
||||
cleaned = _UNSAFE_SEGMENT.sub('_', segment)
|
||||
if cleaned in ('', '.', '..'):
|
||||
return '_'
|
||||
return cleaned
|
||||
|
||||
|
||||
@dataclass
|
||||
class LocalFileStore:
|
||||
"""Dependency-free `OverflowStore` that writes each payload to a local file.
|
||||
|
||||
The handle equals the key: a relative `run_id/tool_call_id.retry` path under
|
||||
`base_dir`. The root is stable and shareable on purpose -- a later agent or run can
|
||||
read a spill a previous run produced, so the store is not isolated per instance.
|
||||
|
||||
Security comes from two mechanisms, not isolation: the root is created with `0700`
|
||||
perms (owner-only), and `read` resolves the target (following symlinks) and rejects
|
||||
anything that escapes the root via symlink, `..`, or an absolute path. Handle segments
|
||||
are also sanitized by `_safe_segment`.
|
||||
|
||||
Files are kept after the run by default (a later `read_tool_result` may need them).
|
||||
Set `cleanup_after` to opt into age-based pruning; see that field.
|
||||
"""
|
||||
|
||||
base_dir: Path | None = None
|
||||
"""Root directory for spilled files. Defaults to a stable temp subdirectory."""
|
||||
|
||||
cleanup_after: timedelta | None = None
|
||||
"""Opt-in TTL for spilled files. `None` (default) keeps files forever.
|
||||
|
||||
When set, a `write` schedules a background prune (a daemon thread, off the hot path)
|
||||
that deletes files whose modification time is older than `cleanup_after`. Pruning is
|
||||
best-effort: any failure is caught and surfaced via `warnings.warn`, never propagated
|
||||
into the agent run. Modification time (`st_mtime`) is the age signal; last-read time
|
||||
(`st_atime`) is unreliable on `noatime`/`relatime` mounts and is not used.
|
||||
"""
|
||||
|
||||
_root: Path = field(init=False, repr=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._root = (
|
||||
self.base_dir if self.base_dir is not None else Path(tempfile.gettempdir()) / 'pyai_harness_overflow'
|
||||
)
|
||||
|
||||
def _path(self, key: str) -> Path:
|
||||
segments = [_safe_segment(part) for part in key.split('/') if part]
|
||||
if not segments:
|
||||
segments = ['_']
|
||||
return self._root.joinpath(*segments)
|
||||
|
||||
def _ensure_root(self) -> None:
|
||||
"""Create the root directory owned by the current user with `0700` perms."""
|
||||
self._root.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
self._root.chmod(0o700)
|
||||
except OSError: # pragma: no cover - best effort on a root we do not own
|
||||
pass
|
||||
|
||||
async def write(self, key: str, data: bytes) -> str:
|
||||
self._ensure_root()
|
||||
path = self._path(key)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_bytes(data)
|
||||
self._schedule_cleanup()
|
||||
return key
|
||||
|
||||
async def read(self, handle: str) -> bytes:
|
||||
target = self._path(handle).resolve()
|
||||
root = self._root.resolve()
|
||||
if not target.is_relative_to(root):
|
||||
raise PermissionError(f'Handle {handle!r} resolves outside the store root.')
|
||||
return target.read_bytes()
|
||||
|
||||
# --- opt-in TTL pruning (non-blocking, non-erroring) ---
|
||||
|
||||
def _schedule_cleanup(self) -> threading.Thread | None:
|
||||
"""Fire a background prune when `cleanup_after` is set. Never blocks `write`."""
|
||||
if self.cleanup_after is None:
|
||||
return None
|
||||
thread = threading.Thread(target=self._run_prune, name='overflow-prune', daemon=True)
|
||||
thread.start()
|
||||
return thread
|
||||
|
||||
def _run_prune(self) -> None:
|
||||
try:
|
||||
self._prune_sync()
|
||||
except Exception as exc: # never let cleanup fail a run or block the hot path
|
||||
warnings.warn(f'LocalFileStore cleanup failed: {exc}', stacklevel=2)
|
||||
|
||||
def _prune_sync(self) -> None:
|
||||
"""Delete files older than `cleanup_after` (by `st_mtime`)."""
|
||||
assert self.cleanup_after is not None
|
||||
cutoff = time.time() - self.cleanup_after.total_seconds()
|
||||
for path in self._root.rglob('*'):
|
||||
if not path.is_file():
|
||||
continue
|
||||
try:
|
||||
if path.stat().st_mtime < cutoff:
|
||||
path.unlink()
|
||||
except OSError: # pragma: no cover - file vanished mid-prune
|
||||
continue
|
||||
@@ -0,0 +1,754 @@
|
||||
"""Tests for pydantic_ai_harness.experimental.overflow."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import os
|
||||
import time
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.messages import (
|
||||
BinaryContent,
|
||||
ModelMessage,
|
||||
ModelResponse,
|
||||
TextPart,
|
||||
ToolCallPart,
|
||||
ToolReturn,
|
||||
ToolReturnPart,
|
||||
)
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
from pydantic_ai.models.test import TestModel
|
||||
from pydantic_ai.tools import ToolDefinition
|
||||
from pydantic_ai.usage import RunUsage
|
||||
|
||||
from pydantic_ai_harness.experimental.overflow import (
|
||||
Band,
|
||||
LocalFileStore,
|
||||
OverflowingToolOutput,
|
||||
Passthrough,
|
||||
Spill,
|
||||
Summarize,
|
||||
Truncate,
|
||||
TruncationStrategy,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.overflow._capability import (
|
||||
READ_TOOL_NAME,
|
||||
_build_spill_preview,
|
||||
_handle_key,
|
||||
_head_tail_preview,
|
||||
_read_slice,
|
||||
_select_action,
|
||||
_Unit,
|
||||
_with_handles,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.overflow._payload import (
|
||||
is_binary,
|
||||
json_sketch,
|
||||
measure,
|
||||
strip_ansi,
|
||||
to_bytes,
|
||||
to_text,
|
||||
truncate_text,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.overflow._store import _safe_segment
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_ctx(*, run_id: str | None = 'run-1', retry: int = 0, model: Any = None) -> Any:
|
||||
"""Build a minimal RunContext-like object for testing the hook directly."""
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _FakeModel:
|
||||
model_id: str = 'test-model'
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _FakeCtx:
|
||||
usage: RunUsage
|
||||
run_id: str | None
|
||||
retry: int
|
||||
tool_call_id: str | None = 'call-1'
|
||||
model: Any = dataclasses.field(default_factory=_FakeModel)
|
||||
deps: None = None
|
||||
|
||||
ctx = _FakeCtx(usage=RunUsage(), run_id=run_id, retry=retry)
|
||||
if model is not None:
|
||||
ctx.model = model
|
||||
return ctx
|
||||
|
||||
|
||||
def _fixed_model(text: str) -> FunctionModel:
|
||||
"""A `FunctionModel` whose single text response is `text` (no tool calls)."""
|
||||
|
||||
def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
return ModelResponse(parts=[TextPart(content=text)])
|
||||
|
||||
return FunctionModel(respond)
|
||||
|
||||
|
||||
def _call(tool_name: str = 'big_tool', tool_call_id: str = 'call-1') -> ToolCallPart:
|
||||
return ToolCallPart(tool_name=tool_name, args='{}', tool_call_id=tool_call_id)
|
||||
|
||||
|
||||
def _tool_def(name: str = 'big_tool') -> ToolDefinition:
|
||||
return ToolDefinition(name=name)
|
||||
|
||||
|
||||
async def _run(cap: OverflowingToolOutput[None], result: Any, *, ctx: Any = None, tool_name: str = 'big_tool') -> Any:
|
||||
return await cap.after_tool_execute(
|
||||
ctx if ctx is not None else _make_ctx(),
|
||||
call=_call(tool_name),
|
||||
tool_def=_tool_def(tool_name),
|
||||
args={},
|
||||
result=result,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend() -> str:
|
||||
return 'asyncio'
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _payload helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPayloadHelpers:
|
||||
def test_strip_ansi(self):
|
||||
assert strip_ansi('\x1b[31mred\x1b[0m') == 'red'
|
||||
|
||||
def test_is_binary(self):
|
||||
assert is_binary(b'x') is True
|
||||
assert is_binary(bytearray(b'x')) is True
|
||||
assert is_binary('x') is False
|
||||
|
||||
def test_to_bytes_variants(self):
|
||||
assert to_bytes('hi') == b'hi'
|
||||
assert to_bytes(memoryview(b'mv')) == b'mv'
|
||||
assert to_bytes(bytearray(b'ba')) == b'ba'
|
||||
assert to_bytes({'a': 1}) == b'{"a":1}'
|
||||
|
||||
def test_to_text_variants(self):
|
||||
assert to_text('hi') == 'hi'
|
||||
assert to_text({'a': 1}) == '{"a":1}'
|
||||
|
||||
def test_measure_chars_and_tokens(self):
|
||||
assert measure('x' * 100, over_tokens=False, tokenizer=None) == 100
|
||||
assert measure('x' * 100, over_tokens=True, tokenizer=None) == 25
|
||||
assert measure('abcd', over_tokens=True, tokenizer=lambda s: len(s)) == 4
|
||||
|
||||
def test_json_sketch_mapping(self):
|
||||
assert json_sketch({'a': 1, 'b': 'x'}) == "{'a': int, 'b': str}"
|
||||
|
||||
def test_json_sketch_mapping_truncated(self):
|
||||
big = {f'k{i}': i for i in range(12)}
|
||||
assert json_sketch(big).endswith('... (12 keys)}')
|
||||
|
||||
def test_json_sketch_sequence(self):
|
||||
assert json_sketch([1, 2, 3]) == '[3 items of int]'
|
||||
|
||||
def test_json_sketch_empty_sequence(self):
|
||||
assert json_sketch([]) == '[0 items of empty]'
|
||||
|
||||
def test_json_sketch_scalar(self):
|
||||
assert json_sketch(42) == ''
|
||||
assert json_sketch('plain') == ''
|
||||
|
||||
def test_truncate_under_limit(self):
|
||||
assert truncate_text('short', 100, TruncationStrategy.head_tail) == 'short'
|
||||
|
||||
def test_truncate_head(self):
|
||||
out = truncate_text('a' * 100, 10, TruncationStrategy.head)
|
||||
assert out.startswith('aaaaaaaaaa')
|
||||
assert 'showing first 10' in out
|
||||
|
||||
def test_truncate_tail(self):
|
||||
out = truncate_text('a' * 100, 10, TruncationStrategy.tail)
|
||||
assert out.endswith('aaaaaaaaaa')
|
||||
assert 'showing last 10' in out
|
||||
|
||||
def test_truncate_head_tail(self):
|
||||
out = truncate_text('a' * 100, 10, TruncationStrategy.head_tail)
|
||||
assert 'omitted from the middle' in out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Store: write/read, S1 hardening
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStore:
|
||||
def test_safe_segment(self):
|
||||
assert _safe_segment('a b!@#') == 'a_b_'
|
||||
assert _safe_segment('') == '_'
|
||||
assert _safe_segment('..') == '_'
|
||||
assert _safe_segment('.') == '_'
|
||||
assert _safe_segment('ok-1.2') == 'ok-1.2'
|
||||
|
||||
def test_default_root(self):
|
||||
store = LocalFileStore()
|
||||
assert store._root.name == 'pyai_harness_overflow'
|
||||
|
||||
async def test_write_read_roundtrip(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path / 'store')
|
||||
handle = await store.write('run-1/call-1.0', b'payload')
|
||||
assert handle == 'run-1/call-1.0'
|
||||
assert await store.read(handle) == b'payload'
|
||||
|
||||
async def test_root_created_0700(self, tmp_path: Path):
|
||||
root = tmp_path / 'store'
|
||||
store = LocalFileStore(base_dir=root)
|
||||
await store.write('run/c.0', b'x')
|
||||
assert oct(root.stat().st_mode & 0o777) == '0o700'
|
||||
|
||||
async def test_empty_key(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path / 'store')
|
||||
handle = await store.write('', b'data')
|
||||
assert await store.read(handle) == b'data'
|
||||
|
||||
async def test_read_missing_raises(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path / 'store')
|
||||
with pytest.raises(OSError):
|
||||
await store.read('nope/x.0')
|
||||
|
||||
async def test_dotdot_handle_stays_in_root(self, tmp_path: Path):
|
||||
# `_safe_segment` neutralizes `..`, so the read resolves inside the root and 404s
|
||||
# rather than escaping.
|
||||
store = LocalFileStore(base_dir=tmp_path / 'store')
|
||||
await store.write('run/c.0', b'inside')
|
||||
with pytest.raises(OSError):
|
||||
await store.read('../c.0')
|
||||
|
||||
async def test_symlink_escape_rejected(self, tmp_path: Path):
|
||||
secret = tmp_path / 'secret.txt'
|
||||
secret.write_bytes(b'top secret')
|
||||
root = tmp_path / 'store'
|
||||
store = LocalFileStore(base_dir=root)
|
||||
await store.write('run/c.0', b'inside') # creates the root
|
||||
(root / 'evil').symlink_to(secret)
|
||||
with pytest.raises(PermissionError, match='outside the store root'):
|
||||
await store.read('evil')
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Store: opt-in TTL cleanup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCleanup:
|
||||
def test_prune_removes_old_keeps_new(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path, cleanup_after=timedelta(seconds=1))
|
||||
old = tmp_path / 'old.bin'
|
||||
old.write_bytes(b'x')
|
||||
new = tmp_path / 'new.bin'
|
||||
new.write_bytes(b'y')
|
||||
(tmp_path / 'sub').mkdir() # a directory rglob yields -- must be skipped
|
||||
past = time.time() - 100
|
||||
os.utime(old, (past, past))
|
||||
|
||||
store._prune_sync()
|
||||
|
||||
assert not old.exists()
|
||||
assert new.exists()
|
||||
|
||||
def test_run_prune_swallows_errors(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
store = LocalFileStore(base_dir=tmp_path, cleanup_after=timedelta(seconds=1))
|
||||
|
||||
def boom() -> None:
|
||||
raise OSError('disk gone')
|
||||
|
||||
monkeypatch.setattr(store, '_prune_sync', boom)
|
||||
with pytest.warns(UserWarning, match='cleanup failed'):
|
||||
store._run_prune()
|
||||
|
||||
def test_schedule_none_when_disabled(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
assert store._schedule_cleanup() is None
|
||||
|
||||
def test_schedule_starts_thread(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path, cleanup_after=timedelta(seconds=1))
|
||||
(tmp_path / 'f.bin').write_bytes(b'z')
|
||||
thread = store._schedule_cleanup()
|
||||
assert thread is not None
|
||||
thread.join(timeout=5)
|
||||
assert not thread.is_alive()
|
||||
|
||||
async def test_write_schedules_cleanup(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
store = LocalFileStore(base_dir=tmp_path / 'store', cleanup_after=timedelta(seconds=1))
|
||||
scheduled: list[int] = []
|
||||
monkeypatch.setattr(store, '_schedule_cleanup', lambda: scheduled.append(1))
|
||||
await store.write('run/c.0', b'data')
|
||||
assert scheduled == [1]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Capability construction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestConstruction:
|
||||
def test_default_band_is_spill_then_truncate(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput()
|
||||
assert len(cap._bands) == 1
|
||||
action = cap._bands[0].action
|
||||
assert isinstance(action, Spill)
|
||||
assert isinstance(action.then, Truncate)
|
||||
|
||||
def test_bands_sorted_descending(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=10, action=Truncate()), Band(over=100, action=Spill())]
|
||||
)
|
||||
assert [b.over for b in cap._bands] == [100, 10]
|
||||
|
||||
def test_negative_threshold_rejected(self):
|
||||
with pytest.raises(ValueError, match='non-negative'):
|
||||
OverflowingToolOutput(bands=[Band(over=-1, action=Passthrough())])
|
||||
|
||||
def test_provided_store_used(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(store=store)
|
||||
assert cap._store is store
|
||||
|
||||
def test_per_tool_prepared(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
per_tool={'read_file': [Band(over=5, action=Truncate())]}
|
||||
)
|
||||
assert 'read_file' in cap._per_tool
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Passthrough / filtering / guards
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPassthrough:
|
||||
async def test_read_tool_exempt(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Truncate(max_chars=2))])
|
||||
out = await _run(cap, 'x' * 100, tool_name=READ_TOOL_NAME)
|
||||
assert out == 'x' * 100
|
||||
|
||||
async def test_tool_filter_skips_unmatched(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=1, action=Truncate(max_chars=2))], tool_filter=['other']
|
||||
)
|
||||
out = await _run(cap, 'x' * 100)
|
||||
assert out == 'x' * 100
|
||||
|
||||
async def test_callable_filter(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=1, action=Truncate(max_chars=2))],
|
||||
tool_filter=lambda ctx, td: td.name == 'big_tool',
|
||||
)
|
||||
out = await _run(cap, 'x' * 100)
|
||||
assert isinstance(out, str) and 'truncated' in out
|
||||
|
||||
async def test_below_threshold_passthrough(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1000, action=Truncate())])
|
||||
out = await _run(cap, 'small')
|
||||
assert out == 'small'
|
||||
|
||||
async def test_exception_result_passthrough(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Truncate(max_chars=2))])
|
||||
err = ValueError('boom')
|
||||
assert await _run(cap, err) is err
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Truncate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTruncate:
|
||||
async def test_truncates_text(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=10, action=Truncate(max_chars=20, strategy=TruncationStrategy.head))]
|
||||
)
|
||||
out = await _run(cap, 'a' * 100)
|
||||
assert isinstance(out, str) and out.startswith('a' * 20)
|
||||
|
||||
async def test_strip_ansi_applied(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=5, action=Truncate(max_chars=1000))], strip_ansi=True
|
||||
)
|
||||
out = await _run(cap, '\x1b[31m' + 'red text ' * 10 + '\x1b[0m')
|
||||
assert isinstance(out, str) and '\x1b[' not in out
|
||||
|
||||
async def test_binary_truncate_falls_back_to_passthrough(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Truncate())])
|
||||
data = b'\x00\x01' * 100
|
||||
assert await _run(cap, data) == data
|
||||
|
||||
async def test_tool_return_envelope_preserved(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=10, action=Truncate(max_chars=20))])
|
||||
out = await _run(cap, ToolReturn(return_value='a' * 100, content='note', metadata={'k': 1}))
|
||||
assert isinstance(out, ToolReturn)
|
||||
assert out.content == 'note'
|
||||
assert out.metadata == {'k': 1}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Spill
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSpill:
|
||||
async def test_spill_roundtrip(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=10, action=Spill(preview_chars=20))], store=store
|
||||
)
|
||||
text = 'line\n' * 1000
|
||||
out = await _run(cap, text)
|
||||
assert isinstance(out, ToolReturn)
|
||||
assert isinstance(out.return_value, str) and 'too large' in out.return_value
|
||||
handle = out.metadata['overflow_handle']
|
||||
assert handle == 'run-1/call-1.0'
|
||||
assert await store.read(handle) == text.encode('utf-8')
|
||||
|
||||
async def test_spill_binary_verbatim(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Spill())], store=store)
|
||||
data = b'\x00\xff' * 100
|
||||
out = await _run(cap, data)
|
||||
assert isinstance(out, ToolReturn)
|
||||
assert 'binary' in out.return_value # type: ignore[operator]
|
||||
assert await store.read(out.metadata['overflow_handle']) == data
|
||||
|
||||
async def test_spill_structured_includes_sketch(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Spill())], store=store)
|
||||
out = await _run(cap, {'rows': list(range(1000)), 'ok': True})
|
||||
assert isinstance(out, ToolReturn)
|
||||
assert 'shape:' in out.return_value # type: ignore[operator]
|
||||
|
||||
async def test_spill_failure_falls_back_to_truncate(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=10, action=Spill(then=Truncate(max_chars=15)))], store=_BrokenStore()
|
||||
)
|
||||
out = await _run(cap, 'a' * 100)
|
||||
assert isinstance(out, str) and 'truncated' in out
|
||||
|
||||
async def test_spill_failure_no_fallback_returns_original(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=10, action=Spill())], store=_BrokenStore()
|
||||
)
|
||||
out = await _run(cap, 'a' * 100)
|
||||
assert out == 'a' * 100
|
||||
|
||||
async def test_handle_distinct_per_retry(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Spill())], store=store)
|
||||
out0 = await _run(cap, 'a' * 100, ctx=_make_ctx(retry=0))
|
||||
out1 = await _run(cap, 'b' * 100, ctx=_make_ctx(retry=1))
|
||||
assert out0.metadata['overflow_handle'] != out1.metadata['overflow_handle'] # type: ignore[union-attr]
|
||||
|
||||
async def test_spill_merges_existing_metadata(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Spill())], store=store)
|
||||
out = await _run(cap, ToolReturn(return_value='a' * 100, metadata={'orig': True}))
|
||||
assert isinstance(out, ToolReturn)
|
||||
assert out.metadata['orig'] is True
|
||||
assert 'overflow_handle' in out.metadata
|
||||
|
||||
|
||||
class _BrokenStore:
|
||||
"""An `OverflowStore` whose writes always fail (for fallback tests)."""
|
||||
|
||||
async def write(self, key: str, data: bytes) -> str:
|
||||
raise OSError('disk full')
|
||||
|
||||
async def read(self, handle: str) -> bytes: # pragma: no cover - never reached
|
||||
raise FileNotFoundError(handle)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# C1: model-visible ToolReturn.content is reduced too
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestContentReduction:
|
||||
async def test_large_content_spilled(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=100, action=Spill(preview_chars=20))], store=store
|
||||
)
|
||||
out = await _run(cap, ToolReturn(return_value='small', content='C' * 5000))
|
||||
assert isinstance(out, ToolReturn)
|
||||
assert out.return_value == 'small' # small return_value untouched
|
||||
assert isinstance(out.content, str) and 'too large' in out.content
|
||||
handle = out.metadata['overflow_content_handle']
|
||||
assert await store.read(handle) == ('C' * 5000).encode('utf-8')
|
||||
|
||||
async def test_large_content_truncated(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=10, action=Truncate(max_chars=20))])
|
||||
out = await _run(cap, ToolReturn(return_value='small', content='C' * 200))
|
||||
assert isinstance(out, ToolReturn)
|
||||
assert isinstance(out.content, str) and 'truncated' in out.content
|
||||
|
||||
async def test_both_value_and_content_reduced(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=50, action=Spill())], store=store)
|
||||
out = await _run(cap, ToolReturn(return_value='v' * 500, content='c' * 500))
|
||||
assert isinstance(out, ToolReturn)
|
||||
assert out.metadata['overflow_handle'] != out.metadata['overflow_content_handle']
|
||||
|
||||
async def test_nontext_content_warns_and_passes_through(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=10, action=Truncate())])
|
||||
content = ['x' * 5000, BinaryContent(data=b'\x00', media_type='application/octet-stream')]
|
||||
original = ToolReturn(return_value='small', content=content)
|
||||
with pytest.warns(UserWarning, match='non-text content'):
|
||||
out = await _run(cap, original)
|
||||
assert out is original
|
||||
|
||||
async def test_nontext_content_passthrough_action_no_warn(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Passthrough())])
|
||||
content = ['x' * 5000, BinaryContent(data=b'\x00', media_type='application/octet-stream')]
|
||||
original = ToolReturn(return_value='small', content=content)
|
||||
out = await _run(cap, original) # Passthrough action -> no warning, returned unchanged
|
||||
assert out is original
|
||||
|
||||
async def test_small_nontext_content_no_warn(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=10_000, action=Truncate())])
|
||||
content = ['tiny', BinaryContent(data=b'\x00', media_type='application/octet-stream')]
|
||||
original = ToolReturn(return_value='small', content=content)
|
||||
out = await _run(cap, original)
|
||||
assert out is original
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Summarize (M1: assert model + usage, not just a wholesale mock)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSummarize:
|
||||
async def test_custom_sync_summarizer(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=5, action=Summarize(summarize=lambda name, text: f'{name}:{len(text)}'))]
|
||||
)
|
||||
out = await _run(cap, 'x' * 100)
|
||||
assert out == 'big_tool:100'
|
||||
|
||||
async def test_custom_async_summarizer(self):
|
||||
async def summ(name: str, text: str) -> str:
|
||||
return f'async:{len(text)}'
|
||||
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Summarize(summarize=summ))])
|
||||
out = await _run(cap, 'x' * 100)
|
||||
assert out == 'async:100'
|
||||
|
||||
async def test_inherited_model_and_usage(self):
|
||||
# model=None resolves to ctx.model, and the call threads usage=ctx.usage.
|
||||
ctx = _make_ctx(model=_fixed_model('THE SUMMARY'))
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Summarize())])
|
||||
out = await _run(cap, 'x' * 100, ctx=ctx)
|
||||
assert out == 'THE SUMMARY'
|
||||
assert ctx.usage.requests == 1
|
||||
|
||||
async def test_explicit_model_overrides_ctx(self):
|
||||
ctx = _make_ctx(model=_fixed_model('FROM CTX MODEL'))
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=5, action=Summarize(model=_fixed_model('FROM EXPLICIT MODEL')))]
|
||||
)
|
||||
out = await _run(cap, 'x' * 100, ctx=ctx)
|
||||
assert out == 'FROM EXPLICIT MODEL'
|
||||
assert ctx.usage.requests == 1
|
||||
|
||||
async def test_binary_summarize_falls_back(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=1, action=Summarize(then=Passthrough()))]
|
||||
)
|
||||
data = b'\x00' * 100
|
||||
assert await _run(cap, data) == data
|
||||
|
||||
async def test_summarize_failure_falls_back(self):
|
||||
def boom(name: str, text: str) -> str:
|
||||
raise RuntimeError('model down')
|
||||
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=5, action=Summarize(summarize=boom, then=Truncate(max_chars=10)))]
|
||||
)
|
||||
out = await _run(cap, 'a' * 100)
|
||||
assert isinstance(out, str) and 'truncated' in out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Passthrough action + per-tool + band selection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestActionsAndSelection:
|
||||
async def test_passthrough_action(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Passthrough())])
|
||||
assert await _run(cap, 'x' * 100) == 'x' * 100
|
||||
|
||||
async def test_per_tool_replaces_bands(self):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=1, action=Truncate(max_chars=5))],
|
||||
per_tool={'big_tool': [Band(over=100_000, action=Truncate())]},
|
||||
)
|
||||
# global band would truncate, but per_tool threshold is huge -> passthrough
|
||||
assert await _run(cap, 'x' * 100) == 'x' * 100
|
||||
|
||||
def test_select_action_no_match(self):
|
||||
assert _select_action([Band(over=100, action=Passthrough())], 50) is None
|
||||
|
||||
def test_select_action_first_match(self):
|
||||
bands = [Band(over=100, action=Spill()), Band(over=10, action=Truncate())]
|
||||
assert isinstance(_select_action(bands, 50), Truncate)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestInternals:
|
||||
def test_handle_key_defaults(self):
|
||||
ctx = _make_ctx(run_id=None, retry=2)
|
||||
ctx.tool_call_id = None
|
||||
key = _handle_key(ctx, ToolCallPart(tool_name='t', args='{}', tool_call_id=''))
|
||||
assert key == 'run/call.2'
|
||||
|
||||
def test_handle_key_suffix(self):
|
||||
key = _handle_key(_make_ctx(), ToolCallPart(tool_name='t', args='{}', tool_call_id='c'), '.content')
|
||||
assert key.endswith('.content')
|
||||
|
||||
def test_with_handles_non_mapping(self):
|
||||
meta = _with_handles('not-a-mapping', 'h/1.0', 42)
|
||||
assert meta == {'overflow_handle': 'h/1.0', 'overflow_bytes': 42}
|
||||
|
||||
def test_with_handles_content_only(self):
|
||||
meta = _with_handles({'orig': 1}, None, 0, 'h/1.0.content')
|
||||
assert meta == {'orig': 1, 'overflow_content_handle': 'h/1.0.content'}
|
||||
|
||||
def test_head_tail_preview_under(self):
|
||||
assert _head_tail_preview('short', 1000) == 'short'
|
||||
|
||||
def test_head_tail_preview_over(self):
|
||||
assert 'omitted' in _head_tail_preview('a' * 100, 10)
|
||||
|
||||
def test_build_spill_preview_tokens_unit(self):
|
||||
unit = _Unit(binary=False, text='x' * 100, data=b'x' * 100, value='x' * 100, suffix='')
|
||||
assert 'tokens' in _build_spill_preview('h/1.0', unit, 20, over_tokens=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# read_tool_result / _read_slice (C2 bounds + literal pattern)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestReadBack:
|
||||
async def test_read_slice_basic(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
await store.write('h/1.0', '\n'.join(f'line {i}' for i in range(50)).encode('utf-8'))
|
||||
out = await _read_slice(store, 'h/1.0', offset=0, limit=3, from_end=False, pattern=None)
|
||||
assert 'line 0' in out and 'line 2' in out and 'line 3' not in out
|
||||
|
||||
async def test_read_slice_from_end(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
await store.write('h/1.0', '\n'.join(f'line {i}' for i in range(50)).encode('utf-8'))
|
||||
out = await _read_slice(store, 'h/1.0', offset=0, limit=2, from_end=True, pattern=None)
|
||||
assert 'line 49' in out and 'line 48' in out
|
||||
|
||||
async def test_read_slice_literal_pattern(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
await store.write('h/1.0', b'apple\nbanana\navocado\ncherry')
|
||||
out = await _read_slice(store, 'h/1.0', offset=0, limit=200, from_end=False, pattern='av')
|
||||
assert 'avocado' in out and 'apple' not in out and 'banana' not in out
|
||||
|
||||
async def test_read_slice_pattern_is_literal_not_regex(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
await store.write('h/1.0', b'plain line\n^anchored')
|
||||
# A regex metacharacter is matched literally, so it cannot trigger backtracking.
|
||||
out = await _read_slice(store, 'h/1.0', offset=0, limit=200, from_end=False, pattern='^a')
|
||||
assert 'anchored' in out and 'plain line' not in out
|
||||
|
||||
async def test_read_slice_offset_negative(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
await store.write('h/1.0', b'data')
|
||||
with pytest.raises(ModelRetry, match='offset'):
|
||||
await _read_slice(store, 'h/1.0', offset=-1, limit=10, from_end=False, pattern=None)
|
||||
|
||||
async def test_read_slice_limit_too_small(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
await store.write('h/1.0', b'data')
|
||||
with pytest.raises(ModelRetry, match='limit'):
|
||||
await _read_slice(store, 'h/1.0', offset=0, limit=0, from_end=False, pattern=None)
|
||||
|
||||
async def test_read_slice_limit_clamped(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
await store.write('h/1.0', '\n'.join(f'l{i}' for i in range(2000)).encode('utf-8'))
|
||||
out = await _read_slice(store, 'h/1.0', offset=0, limit=10_000, from_end=False, pattern=None)
|
||||
assert out.count('\n') <= 1_000 # clamped to the line cap
|
||||
|
||||
async def test_read_slice_output_capped(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
await store.write('h/1.0', ('x' * 60_000).encode('utf-8'))
|
||||
out = await _read_slice(store, 'h/1.0', offset=0, limit=10, from_end=False, pattern=None)
|
||||
assert 'output capped' in out
|
||||
assert len(out) < 60_000
|
||||
|
||||
async def test_read_slice_missing_handle(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
with pytest.raises(ModelRetry, match='No stored tool result'):
|
||||
await _read_slice(store, 'missing/1.0', offset=0, limit=10, from_end=False, pattern=None)
|
||||
|
||||
async def test_get_toolset_registers_read_tool(self, tmp_path: Path):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
await store.write('h/1.0', b'hello\nworld')
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(store=store)
|
||||
toolset = cap.get_toolset()
|
||||
assert toolset is not None
|
||||
tool = toolset.tools[READ_TOOL_NAME] # type: ignore[union-attr]
|
||||
out = await tool.function(_make_ctx(), 'h/1.0') # type: ignore[attr-defined]
|
||||
assert 'hello' in out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Agent-path integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAgentIntegration:
|
||||
async def test_spill_persists_in_history(self, tmp_path: Path, anyio_backend: str):
|
||||
store = LocalFileStore(base_dir=tmp_path)
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=100, action=Spill(preview_chars=50))], store=store
|
||||
)
|
||||
agent = Agent(TestModel(call_tools=['big_tool']), capabilities=[cap])
|
||||
|
||||
@agent.tool_plain
|
||||
def big_tool() -> str:
|
||||
return 'data line\n' * 500
|
||||
|
||||
result = await agent.run('go')
|
||||
returns = [p for m in result.all_messages() for p in m.parts if isinstance(p, ToolReturnPart)]
|
||||
spilled = [p for p in returns if p.tool_name == 'big_tool']
|
||||
assert spilled
|
||||
part = spilled[0]
|
||||
assert isinstance(part.content, str) and 'too large' in part.content
|
||||
assert part.metadata is not None and 'overflow_handle' in part.metadata
|
||||
assert await store.read(part.metadata['overflow_handle']) == ('data line\n' * 500).encode('utf-8')
|
||||
|
||||
async def test_small_output_untouched(self, tmp_path: Path, anyio_backend: str):
|
||||
cap: OverflowingToolOutput[None] = OverflowingToolOutput(
|
||||
bands=[Band(over=10_000, action=Spill())], store=LocalFileStore(base_dir=tmp_path)
|
||||
)
|
||||
agent = Agent(TestModel(call_tools=['small_tool']), capabilities=[cap])
|
||||
|
||||
@agent.tool_plain
|
||||
def small_tool() -> str:
|
||||
return 'tiny'
|
||||
|
||||
result = await agent.run('go')
|
||||
returns = [p for m in result.all_messages() for p in m.parts if isinstance(p, ToolReturnPart)]
|
||||
small = [p for p in returns if p.tool_name == 'small_tool']
|
||||
assert small and small[0].content == 'tiny'
|
||||
Reference in New Issue
Block a user