mirror of
https://github.com/pydantic/pydantic-ai-harness.git
synced 2026-07-21 02:45:34 +00:00
Emit a span when compaction happens (#312)
* Emit compaction telemetry for observability Compaction silently rewrites a run's message history, so an operator could not see when or how much it reclaimed. Each compacting strategy now emits a `compact_messages` span (with before/after message and token counts) when it actually runs, so compaction is visible alongside the rest of a run's traces. The span is emitted in `before_model_request` rather than `compact` so that `TieredCompaction` produces one span for the escalation instead of one per tier. It rides on the run's tracer, which is a no-op without instrumentation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015kRjwKq2AmcuWZcDfJ4WHs * Suppress compaction span on no-op and uninstrumented runs A trigger can fire while compact() returns the history unchanged (most notably DeduplicateFileReads, which has no threshold), so wrapping compact() inside the span emitted a compact_messages span on ordinary turns with identical before/after counts. Run compact() first and only start the span when the returned history actually changed. Building the token attributes also ran estimate_token_count (and any user tokenizer) over the full before/after histories on every compaction, even when ctx.tracer is a no-op tracer and the span is discarded. Gate the attribute computation on span.is_recording() so uninstrumented runs pay nothing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015kRjwKq2AmcuWZcDfJ4WHs * Make compact_with_span config params keyword-only Forcing strategy/messages/compact/tokenizer to be passed by keyword prevents argument-order mistakes at the six strategy call sites, where several same-typed callables/lists are passed positionally. * making compact_with_span * only after ctx * Type compaction test tracer field with opentelemetry Tracer The fake RunContext tracer field was typed as Any. opentelemetry provides the proper type and core's RunContext uses it, so type the field as opentelemetry.trace.Tracer to match. * Type compaction test fake context model with pydantic_ai Model The fake RunContext model field was typed as Any and backed by a hand-rolled stub carrying only model_id. Back it with TestModel, a real Model subclass, and annotate the field as pydantic_ai.models.Model so the fake conforms to the type the production code expects from ctx.model. * Return original messages on no-op compaction The no-op early-return branch of compact_with_span only runs when the history is unchanged, so returning the original messages instead of the compacted list makes the no-op explicit at the call site without changing behavior. * redundant list length check * history covers would_clamp * Set gen_ai.conversation.compacted on the compaction span The OpenTelemetry GenAI semantic convention models context compaction as the `gen_ai.conversation.compacted` boolean, set true-only. The span emits solely when the history actually changed, so the flag is only ever true. The `compaction.*` attributes have no convention equivalent and stay harness-specific. --------- Co-authored-by: Aditya Vardhan <adtyavrdhn@gmail.com>
This commit is contained in:
co-authored by
Aditya Vardhan
parent
4e21f7bcbf
commit
326c3d648d
@@ -176,6 +176,32 @@ def my_file_key(call: ToolCallPart) -> str | None:
|
||||
return args.get('path') if isinstance(args, dict) else None
|
||||
```
|
||||
|
||||
## Tracing
|
||||
|
||||
When core instrumentation is active (the `Instrumentation` capability, `agent.instrument`, or
|
||||
`Agent.instrument_all()`), each strategy emits a `compact_messages` span on the run's tracer the
|
||||
moment it actually compacts -- that is, in `before_model_request`, once the strategy's threshold is
|
||||
exceeded (`ClampOversizedMessages` emits only when a part is actually clamped). `TieredCompaction`
|
||||
emits a single span for the whole escalation rather than one per tier, because it drives each tier's
|
||||
`compact` directly. Without instrumentation the tracer is a no-op, so the span adds no overhead.
|
||||
|
||||
The span name is the static `compact_messages`; the strategy is an attribute, not part of the name,
|
||||
to keep span cardinality low. Attributes:
|
||||
|
||||
| Attribute | Type | Meaning |
|
||||
|---|---|---|
|
||||
| `gen_ai.conversation.compacted` | bool | Always `true`; the OpenTelemetry GenAI convention's flag for a compacted context |
|
||||
| `compaction.strategy` | str | Strategy class name (e.g. `SlidingWindow`, `SummarizingCompaction`) |
|
||||
| `compaction.messages_before` | int | Message count before compaction |
|
||||
| `compaction.messages_after` | int | Message count after compaction |
|
||||
| `compaction.tokens_before` | int | Estimated token count before compaction |
|
||||
| `compaction.tokens_after` | int | Estimated token count after compaction |
|
||||
|
||||
`gen_ai.conversation.compacted` is the GenAI semantic convention's flag; the rest is
|
||||
harness-specific. Token counts use the strategy's `tokenizer` when set, otherwise the
|
||||
~4-chars-per-token heuristic.
|
||||
Raw message content is not recorded.
|
||||
|
||||
## Out of scope
|
||||
|
||||
These strategies compress or drop context *inside* the window. Moving large tool outputs *out* of the
|
||||
|
||||
@@ -17,7 +17,7 @@ from pydantic_ai.messages import (
|
||||
)
|
||||
from pydantic_ai.tools import RunContext
|
||||
|
||||
from pydantic_ai_harness.experimental.compaction._shared import estimate_text_tokens
|
||||
from pydantic_ai_harness.experimental.compaction._shared import compact_with_span, estimate_text_tokens
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic_ai.models import ModelRequestContext
|
||||
@@ -170,5 +170,12 @@ class ClampOversizedMessages(AbstractCapability[AgentDepsT]):
|
||||
request_context: ModelRequestContext,
|
||||
) -> ModelRequestContext:
|
||||
"""Clamp any oversized response part before the request is sent."""
|
||||
request_context.messages = await self.compact(list(request_context.messages), ctx)
|
||||
messages: list[ModelMessage] = list(request_context.messages)
|
||||
request_context.messages = await compact_with_span(
|
||||
ctx,
|
||||
strategy='ClampOversizedMessages',
|
||||
messages=messages,
|
||||
compact=lambda: self.compact(messages, ctx),
|
||||
tokenizer=self.tokenizer,
|
||||
)
|
||||
return request_context
|
||||
|
||||
@@ -12,6 +12,7 @@ from pydantic_ai.messages import ModelMessage
|
||||
from pydantic_ai.tools import RunContext
|
||||
|
||||
from pydantic_ai_harness.experimental.compaction._shared import (
|
||||
compact_with_span,
|
||||
estimate_token_count,
|
||||
exceeds,
|
||||
iter_tool_pairs,
|
||||
@@ -131,5 +132,11 @@ class ClearToolResults(AbstractCapability[AgentDepsT]):
|
||||
messages: list[ModelMessage] = list(request_context.messages)
|
||||
if not exceeds(messages, self.max_messages, self.max_tokens, self.tokenizer):
|
||||
return request_context
|
||||
request_context.messages = await self.compact(messages, ctx)
|
||||
request_context.messages = await compact_with_span(
|
||||
ctx,
|
||||
strategy='ClearToolResults',
|
||||
messages=messages,
|
||||
compact=lambda: self.compact(messages, ctx),
|
||||
tokenizer=self.tokenizer,
|
||||
)
|
||||
return request_context
|
||||
|
||||
@@ -11,7 +11,12 @@ from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.messages import ModelMessage, ToolCallPart
|
||||
from pydantic_ai.tools import RunContext
|
||||
|
||||
from pydantic_ai_harness.experimental.compaction._shared import exceeds, iter_tool_pairs, rebuild_with_cleared
|
||||
from pydantic_ai_harness.experimental.compaction._shared import (
|
||||
compact_with_span,
|
||||
exceeds,
|
||||
iter_tool_pairs,
|
||||
rebuild_with_cleared,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic_ai.models import ModelRequestContext
|
||||
@@ -107,5 +112,11 @@ class DeduplicateFileReads(AbstractCapability[AgentDepsT]):
|
||||
if self.max_messages is not None or self.max_tokens is not None:
|
||||
if not exceeds(messages, self.max_messages, self.max_tokens, self.tokenizer):
|
||||
return request_context
|
||||
request_context.messages = await self.compact(messages, ctx)
|
||||
request_context.messages = await compact_with_span(
|
||||
ctx,
|
||||
strategy='DeduplicateFileReads',
|
||||
messages=messages,
|
||||
compact=lambda: self.compact(messages, ctx),
|
||||
tokenizer=self.tokenizer,
|
||||
)
|
||||
return request_context
|
||||
|
||||
@@ -6,7 +6,7 @@ preservation, and in-place tool-result clearing -- anything used by more than on
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Sequence
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Protocol
|
||||
|
||||
@@ -110,6 +110,68 @@ def exceeds(
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tracing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_SPAN_NAME = 'compact_messages'
|
||||
"""Static, low-cardinality span name emitted whenever a strategy compacts. The strategy name
|
||||
goes in the `compaction.strategy` attribute rather than the span name to keep cardinality low."""
|
||||
|
||||
|
||||
def _history_changed(before: list[ModelMessage], after: list[ModelMessage]) -> bool:
|
||||
"""Return True if *after* differs from *before*.
|
||||
|
||||
The same list object, or an equal-length list that compares equal element-wise, counts as
|
||||
unchanged; anything else is a change.
|
||||
"""
|
||||
# `!=` short-circuits on identity element-wise, so this also covers `before is after`; an
|
||||
# unequal length already implies an unequal list, so a separate length check is redundant.
|
||||
return before != after
|
||||
|
||||
|
||||
async def compact_with_span(
|
||||
ctx: RunContext[AgentDepsT],
|
||||
*,
|
||||
strategy: str,
|
||||
messages: list[ModelMessage],
|
||||
compact: Callable[[], Awaitable[list[ModelMessage]]],
|
||||
tokenizer: Callable[[str], int] | None = None,
|
||||
) -> list[ModelMessage]:
|
||||
"""Run *compact* and emit a `compact_messages` span when it changes the history.
|
||||
|
||||
*compact* runs before the span so a no-op compaction (a trigger fired but the history is
|
||||
returned unchanged) emits nothing. The span is started on `ctx.tracer`, which is a no-op
|
||||
tracer unless core's instrumentation is active, so this adds no overhead to a
|
||||
non-instrumented run; the before/after attributes are only computed when the span records.
|
||||
|
||||
Args:
|
||||
ctx: Run context whose `tracer` the span is started on.
|
||||
strategy: Strategy name recorded in the `compaction.strategy` attribute.
|
||||
messages: The pre-compaction messages, measured for the `*_before` attributes.
|
||||
compact: Zero-argument async callable returning the compacted message list.
|
||||
tokenizer: Optional tokenizer for the `compaction.tokens_*` estimates. When `None`,
|
||||
uses the same ~4 characters-per-token heuristic as `estimate_token_count`.
|
||||
"""
|
||||
compacted = await compact()
|
||||
if not _history_changed(messages, compacted):
|
||||
return messages
|
||||
with ctx.tracer.start_as_current_span(_SPAN_NAME) as span:
|
||||
if span.is_recording():
|
||||
span.set_attributes(
|
||||
{
|
||||
# GenAI semconv flag; the convention says set `true` only, never `false`.
|
||||
'gen_ai.conversation.compacted': True,
|
||||
'compaction.strategy': strategy,
|
||||
'compaction.messages_before': len(messages),
|
||||
'compaction.messages_after': len(compacted),
|
||||
'compaction.tokens_before': estimate_token_count(messages, tokenizer),
|
||||
'compaction.tokens_after': estimate_token_count(compacted, tokenizer),
|
||||
}
|
||||
)
|
||||
return compacted
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Compaction strategy protocol
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -12,6 +12,7 @@ from pydantic_ai.messages import ModelMessage
|
||||
from pydantic_ai.tools import RunContext
|
||||
|
||||
from pydantic_ai_harness.experimental.compaction._shared import (
|
||||
compact_with_span,
|
||||
exceeds,
|
||||
find_safe_cutoff,
|
||||
find_token_cutoff,
|
||||
@@ -112,5 +113,11 @@ class SlidingWindow(AbstractCapability[AgentDepsT]):
|
||||
messages: list[ModelMessage] = list(request_context.messages)
|
||||
if not exceeds(messages, self.max_messages, self.max_tokens, self.tokenizer):
|
||||
return request_context
|
||||
request_context.messages = await self.compact(messages, ctx)
|
||||
request_context.messages = await compact_with_span(
|
||||
ctx,
|
||||
strategy='SlidingWindow',
|
||||
messages=messages,
|
||||
compact=lambda: self.compact(messages, ctx),
|
||||
tokenizer=self.tokenizer,
|
||||
)
|
||||
return request_context
|
||||
|
||||
@@ -21,6 +21,7 @@ from pydantic_ai.messages import (
|
||||
from pydantic_ai.tools import RunContext
|
||||
|
||||
from pydantic_ai_harness.experimental.compaction._shared import (
|
||||
compact_with_span,
|
||||
exceeds,
|
||||
find_first_user_message,
|
||||
find_safe_cutoff,
|
||||
@@ -259,7 +260,13 @@ class SummarizingCompaction(AbstractCapability[AgentDepsT]):
|
||||
messages: list[ModelMessage] = list(request_context.messages)
|
||||
if not exceeds(messages, self.max_messages, self.max_tokens, self.tokenizer):
|
||||
return request_context
|
||||
request_context.messages = await self.compact(messages, ctx)
|
||||
request_context.messages = await compact_with_span(
|
||||
ctx,
|
||||
strategy='SummarizingCompaction',
|
||||
messages=messages,
|
||||
compact=lambda: self.compact(messages, ctx),
|
||||
tokenizer=self.tokenizer,
|
||||
)
|
||||
return request_context
|
||||
|
||||
async def _summarize(
|
||||
|
||||
@@ -11,7 +11,11 @@ from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.messages import ModelMessage
|
||||
from pydantic_ai.tools import RunContext
|
||||
|
||||
from pydantic_ai_harness.experimental.compaction._shared import CompactionStrategy, estimate_token_count
|
||||
from pydantic_ai_harness.experimental.compaction._shared import (
|
||||
CompactionStrategy,
|
||||
compact_with_span,
|
||||
estimate_token_count,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic_ai.models import ModelRequestContext
|
||||
@@ -91,5 +95,11 @@ class TieredCompaction(AbstractCapability[AgentDepsT]):
|
||||
messages: list[ModelMessage] = list(request_context.messages)
|
||||
if estimate_token_count(messages, self.tokenizer) <= self.target_tokens:
|
||||
return request_context
|
||||
request_context.messages = await self.compact(messages, ctx)
|
||||
request_context.messages = await compact_with_span(
|
||||
ctx,
|
||||
strategy='TieredCompaction',
|
||||
messages=messages,
|
||||
compact=lambda: self.compact(messages, ctx),
|
||||
tokenizer=self.tokenizer,
|
||||
)
|
||||
return request_context
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing import Any
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from opentelemetry.trace import NoOpTracer, Tracer
|
||||
from pydantic_ai.messages import (
|
||||
ModelMessage,
|
||||
ModelRequest,
|
||||
@@ -17,7 +18,8 @@ from pydantic_ai.messages import (
|
||||
ToolReturnPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
from pydantic_ai.models import ModelRequestContext, ModelRequestParameters
|
||||
from pydantic_ai.models import Model, ModelRequestContext, ModelRequestParameters
|
||||
from pydantic_ai.models.test import TestModel
|
||||
from pydantic_ai.usage import RunUsage
|
||||
|
||||
from pydantic_ai_harness.experimental.compaction import (
|
||||
@@ -35,7 +37,9 @@ from pydantic_ai_harness.experimental.compaction._clamp_oversized_messages impor
|
||||
_CLAMP_MARKER,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.compaction._shared import (
|
||||
_history_changed,
|
||||
_is_safe_cutoff,
|
||||
compact_with_span,
|
||||
find_first_user_message,
|
||||
find_safe_cutoff,
|
||||
find_token_cutoff,
|
||||
@@ -49,6 +53,13 @@ from pydantic_ai_harness.experimental.compaction._summarizing_compaction import
|
||||
_format_messages,
|
||||
)
|
||||
|
||||
try:
|
||||
from logfire.testing import CaptureLogfire
|
||||
|
||||
logfire_installed = True
|
||||
except ImportError: # pragma: no cover
|
||||
logfire_installed = False
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -62,17 +73,14 @@ def _make_ctx(
|
||||
) -> Any:
|
||||
"""Build a minimal RunContext-like object for testing hooks."""
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _FakeModel:
|
||||
model_id: str = 'test-model'
|
||||
|
||||
usage = RunUsage(requests=requests, input_tokens=input_tokens, output_tokens=output_tokens)
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _FakeCtx:
|
||||
usage: RunUsage
|
||||
model: Any = dataclasses.field(default_factory=_FakeModel)
|
||||
model: Model = dataclasses.field(default_factory=TestModel)
|
||||
deps: None = None
|
||||
tracer: Tracer = dataclasses.field(default_factory=NoOpTracer)
|
||||
|
||||
return _FakeCtx(usage=usage)
|
||||
|
||||
@@ -2075,3 +2083,310 @@ class TestSummarizingCompactionPreserveBranches:
|
||||
1 for m in result.messages if isinstance(m, ModelRequest) for p in m.parts if isinstance(p, UserPromptPart)
|
||||
)
|
||||
assert user_count == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OTel / Logfire instrumentation: the `compact_messages` span
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _compact_spans(capfire: CaptureLogfire) -> list[dict[str, Any]]:
|
||||
"""Return only the `compact_messages` spans, which this package controls.
|
||||
|
||||
Core's own instrumentation span and attribute names have changed across pydantic-ai
|
||||
versions, so the assertions here stay on the spans and attributes this package emits.
|
||||
"""
|
||||
return [s for s in capfire.exporter.exported_spans_as_dict() if s['name'] == 'compact_messages']
|
||||
|
||||
|
||||
def _make_ctx_with_tracer() -> Any:
|
||||
"""A fake RunContext whose `tracer` exports to the active `CaptureLogfire` provider.
|
||||
|
||||
The `capfire` fixture configures the global OTel provider, so a tracer fetched from it
|
||||
captures the `compact_messages` span without needing a full instrumented `Agent` run.
|
||||
"""
|
||||
from opentelemetry.trace import get_tracer
|
||||
|
||||
ctx = _make_ctx()
|
||||
ctx.tracer = get_tracer('test')
|
||||
return ctx
|
||||
|
||||
|
||||
@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
|
||||
class TestCompactionSpan:
|
||||
@pytest.fixture
|
||||
def anyio_backend(self) -> str:
|
||||
# A full agent.run only needs the asyncio backend; trio hits a TestModel
|
||||
# event-loop quirk in core unrelated to compaction.
|
||||
return 'asyncio'
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_span_emitted_when_threshold_exceeded(self, capfire: CaptureLogfire) -> None:
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.instrumented import InstrumentationSettings
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
agent: Agent[None, str] = Agent(
|
||||
TestModel(),
|
||||
capabilities=[SlidingWindow(max_tokens=1, keep_messages=1)],
|
||||
)
|
||||
agent.instrument = InstrumentationSettings()
|
||||
history: list[ModelMessage] = [_user('first'), _assistant('a'), _user('second'), _assistant('b')]
|
||||
await agent.run('a reasonably long prompt that exceeds one token', message_history=history)
|
||||
|
||||
spans = _compact_spans(capfire)
|
||||
# Exactly one compaction runs (TestModel makes a single request); `== 1` also guards against
|
||||
# double-emission, and the `>` checks assert the window actually shrank rather than just
|
||||
# reporting integers.
|
||||
assert len(spans) == 1
|
||||
attrs = spans[0]['attributes']
|
||||
assert attrs['gen_ai.conversation.compacted'] is True
|
||||
assert attrs['compaction.strategy'] == 'SlidingWindow'
|
||||
assert attrs['compaction.messages_before'] > attrs['compaction.messages_after']
|
||||
assert attrs['compaction.tokens_before'] > attrs['compaction.tokens_after']
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_no_span_when_threshold_not_exceeded(self, capfire: CaptureLogfire) -> None:
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.instrumented import InstrumentationSettings
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
agent: Agent[None, str] = Agent(
|
||||
TestModel(),
|
||||
capabilities=[SlidingWindow(max_tokens=1_000_000, keep_messages=1)],
|
||||
)
|
||||
agent.instrument = InstrumentationSettings()
|
||||
await agent.run('short prompt')
|
||||
|
||||
assert _compact_spans(capfire) == []
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_summarizing_compaction_emits_span(self, capfire: CaptureLogfire) -> None:
|
||||
comp = SummarizingCompaction(model='test:m', max_messages=2, keep_messages=1, incremental=False)
|
||||
messages: list[ModelMessage] = [_user('first'), _assistant('a'), _user('b'), _assistant('c')]
|
||||
rc = _make_request_context(messages)
|
||||
ctx = _make_ctx_with_tracer()
|
||||
|
||||
mock_result = AsyncMock()
|
||||
mock_result.output = 'Summary.'
|
||||
with patch('pydantic_ai.Agent') as MockAgent:
|
||||
mock_agent_instance = AsyncMock()
|
||||
mock_agent_instance.run.return_value = mock_result
|
||||
MockAgent.return_value = mock_agent_instance
|
||||
await comp.before_model_request(ctx, rc)
|
||||
|
||||
spans = _compact_spans(capfire)
|
||||
assert len(spans) == 1
|
||||
assert spans[0]['attributes']['compaction.strategy'] == 'SummarizingCompaction'
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_clamp_emits_span_only_when_a_part_is_clamped(self, capfire: CaptureLogfire) -> None:
|
||||
comp = ClampOversizedMessages(max_part_chars=4, keep_head_chars=1, keep_tail_chars=1)
|
||||
|
||||
not_oversized: list[ModelMessage] = [_assistant('ab')]
|
||||
await comp.before_model_request(_make_ctx_with_tracer(), _make_request_context(not_oversized))
|
||||
assert _compact_spans(capfire) == []
|
||||
|
||||
oversized: list[ModelMessage] = [_assistant('a' * 50)]
|
||||
await comp.before_model_request(_make_ctx_with_tracer(), _make_request_context(oversized))
|
||||
spans = _compact_spans(capfire)
|
||||
assert len(spans) == 1
|
||||
assert spans[0]['attributes']['compaction.strategy'] == 'ClampOversizedMessages'
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_clamp_emits_span_for_oversized_tool_call_args(self, capfire: CaptureLogfire) -> None:
|
||||
comp = ClampOversizedMessages(max_part_chars=4, keep_head_chars=1, keep_tail_chars=1, clamp_tool_call_args=True)
|
||||
messages: list[ModelMessage] = [
|
||||
ModelResponse(parts=[ToolCallPart(tool_name='fn', args={'q': 'x' * 50}, tool_call_id='tc1')])
|
||||
]
|
||||
await comp.before_model_request(_make_ctx_with_tracer(), _make_request_context(messages))
|
||||
|
||||
spans = _compact_spans(capfire)
|
||||
assert len(spans) == 1
|
||||
assert spans[0]['attributes']['compaction.strategy'] == 'ClampOversizedMessages'
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_clamp_no_span_for_non_oversized_or_skipped_parts(self, capfire: CaptureLogfire) -> None:
|
||||
from pydantic_ai.messages import ThinkingPart
|
||||
|
||||
comp = ClampOversizedMessages(max_part_chars=1_000, clamp_tool_call_args=True)
|
||||
messages: list[ModelMessage] = [
|
||||
ModelResponse(
|
||||
parts=[
|
||||
ToolCallPart(tool_name='fn', args={'q': 'x'}, tool_call_id='tc1'),
|
||||
ThinkingPart(content='thinking'),
|
||||
]
|
||||
)
|
||||
]
|
||||
await comp.before_model_request(_make_ctx_with_tracer(), _make_request_context(messages))
|
||||
|
||||
assert _compact_spans(capfire) == []
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_tiered_emits_single_span_not_one_per_tier(self, capfire: CaptureLogfire) -> None:
|
||||
comp: TieredCompaction[None] = TieredCompaction(
|
||||
tiers=[
|
||||
ClearToolResults(max_tokens=1, keep_pairs=0),
|
||||
SlidingWindow(max_tokens=1, keep_messages=1),
|
||||
],
|
||||
target_tokens=1,
|
||||
)
|
||||
messages: list[ModelMessage] = [
|
||||
_user('first'),
|
||||
_tool_call('fn', 'tc1'),
|
||||
_tool_return('fn', 'tc1', 'a long tool result that takes up space'),
|
||||
_assistant('done'),
|
||||
]
|
||||
await comp.before_model_request(_make_ctx_with_tracer(), _make_request_context(messages))
|
||||
|
||||
spans = _compact_spans(capfire)
|
||||
# The orchestrator drives each tier's `compact` directly, so only one span is emitted.
|
||||
assert len(spans) == 1
|
||||
assert spans[0]['attributes']['compaction.strategy'] == 'TieredCompaction'
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_no_span_when_compaction_is_noop(self, capfire: CaptureLogfire) -> None:
|
||||
# DeduplicateFileReads has no threshold, so its trigger always fires, but with no
|
||||
# superseded reads `compact` returns the history unchanged and no span should be emitted.
|
||||
comp = DeduplicateFileReads(file_key=_file_key)
|
||||
messages: list[ModelMessage] = [
|
||||
_read_call('tc1', 'a.py'),
|
||||
_read_return('tc1', 'a body'),
|
||||
_read_call('tc2', 'b.py'),
|
||||
_read_return('tc2', 'b body'),
|
||||
]
|
||||
await comp.before_model_request(_make_ctx_with_tracer(), _make_request_context(messages))
|
||||
|
||||
assert _compact_spans(capfire) == []
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_span_emitted_when_dedup_changes_history(self, capfire: CaptureLogfire) -> None:
|
||||
comp = DeduplicateFileReads(file_key=_file_key)
|
||||
messages: list[ModelMessage] = [
|
||||
_read_call('tc1', 'a.py'),
|
||||
_read_return('tc1', 'first'),
|
||||
_read_call('tc2', 'a.py'),
|
||||
_read_return('tc2', 'second'),
|
||||
]
|
||||
await comp.before_model_request(_make_ctx_with_tracer(), _make_request_context(messages))
|
||||
|
||||
spans = _compact_spans(capfire)
|
||||
assert len(spans) == 1
|
||||
assert spans[0]['attributes']['compaction.strategy'] == 'DeduplicateFileReads'
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_clear_tool_results_emits_span(self, capfire: CaptureLogfire) -> None:
|
||||
# ClearToolResults is otherwise only exercised inside TieredCompaction, which reports the
|
||||
# orchestrator's name -- so this is the only check on its own `strategy` literal.
|
||||
comp = ClearToolResults(max_tokens=1, keep_pairs=0)
|
||||
messages: list[ModelMessage] = [
|
||||
_user('first'),
|
||||
_tool_call('fn', 'tc1'),
|
||||
_tool_return('fn', 'tc1', 'a long tool result that takes up space'),
|
||||
_assistant('done'),
|
||||
]
|
||||
await comp.before_model_request(_make_ctx_with_tracer(), _make_request_context(messages))
|
||||
|
||||
spans = _compact_spans(capfire)
|
||||
assert len(spans) == 1
|
||||
assert spans[0]['attributes']['compaction.strategy'] == 'ClearToolResults'
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# compact_with_span helper internals
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestHistoryChanged:
|
||||
def test_same_object_is_unchanged(self):
|
||||
msgs: list[ModelMessage] = [_user('a')]
|
||||
assert _history_changed(msgs, msgs) is False
|
||||
|
||||
def test_different_length_is_changed(self):
|
||||
before: list[ModelMessage] = [_user('a')]
|
||||
after: list[ModelMessage] = [_user('a'), _user('b')]
|
||||
assert _history_changed(before, after) is True
|
||||
|
||||
def test_same_length_equal_is_unchanged(self):
|
||||
# Distinct list objects holding equal elements compares unchanged. A shared message
|
||||
# object avoids the per-message timestamp that would otherwise break equality.
|
||||
shared = _user('a')
|
||||
before: list[ModelMessage] = [shared]
|
||||
after: list[ModelMessage] = [shared]
|
||||
assert before is not after
|
||||
assert _history_changed(before, after) is False
|
||||
|
||||
def test_same_length_unequal_is_changed(self):
|
||||
before: list[ModelMessage] = [_user('a')]
|
||||
after: list[ModelMessage] = [_user('b')]
|
||||
assert _history_changed(before, after) is True
|
||||
|
||||
|
||||
class TestCompactWithSpan:
|
||||
@pytest.mark.anyio
|
||||
async def test_no_op_returns_without_starting_span(self):
|
||||
# A NoOpTracer would never record anyway; this guards that `compact` runs and the
|
||||
# unchanged result is returned without attempting to start a span.
|
||||
messages: list[ModelMessage] = [_user('a')]
|
||||
|
||||
async def _compact() -> list[ModelMessage]:
|
||||
return messages
|
||||
|
||||
result = await compact_with_span(_make_ctx(), strategy='Strat', messages=messages, compact=_compact)
|
||||
assert result is messages
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
|
||||
async def test_recording_span_sets_attributes(self, capfire: CaptureLogfire) -> None:
|
||||
# Distinct text lengths plus a character-counting tokenizer pin the exact attribute values.
|
||||
# A before/after swap, computing both token counts from one list, or ignoring the strategy
|
||||
# tokenizer would each change a number this test checks.
|
||||
before: list[ModelMessage] = [_user('aaaa'), _user('bb')]
|
||||
after: list[ModelMessage] = [_user('aaaa')]
|
||||
seen: list[str] = []
|
||||
|
||||
def _tokenizer(text: str) -> int:
|
||||
seen.append(text)
|
||||
return len(text)
|
||||
|
||||
async def _compact() -> list[ModelMessage]:
|
||||
return after
|
||||
|
||||
result = await compact_with_span(
|
||||
_make_ctx_with_tracer(), strategy='Strat', messages=before, compact=_compact, tokenizer=_tokenizer
|
||||
)
|
||||
assert result is after
|
||||
|
||||
spans = _compact_spans(capfire)
|
||||
assert len(spans) == 1
|
||||
attrs = spans[0]['attributes']
|
||||
assert attrs['gen_ai.conversation.compacted'] is True
|
||||
assert attrs['compaction.strategy'] == 'Strat'
|
||||
assert attrs['compaction.messages_before'] == 2
|
||||
assert attrs['compaction.messages_after'] == 1
|
||||
# tokenizer counts characters: before = len('aaaa') + len('bb') = 6, after = len('aaaa') = 4.
|
||||
assert attrs['compaction.tokens_before'] == 6
|
||||
assert attrs['compaction.tokens_after'] == 4
|
||||
assert attrs['compaction.tokens_before'] > attrs['compaction.tokens_after']
|
||||
assert seen # the strategy tokenizer reached the span attributes, not the default heuristic
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_non_recording_tracer_skips_attributes(self):
|
||||
# A no-op tracer returns a non-recording span, so attribute computation is skipped.
|
||||
before: list[ModelMessage] = [_user('a'), _user('b')]
|
||||
after: list[ModelMessage] = [_user('a')]
|
||||
called = False
|
||||
|
||||
def _tokenizer(_text: str) -> int: # pragma: no cover - asserted never called
|
||||
nonlocal called
|
||||
called = True
|
||||
return 1
|
||||
|
||||
async def _compact() -> list[ModelMessage]:
|
||||
return after
|
||||
|
||||
result = await compact_with_span(
|
||||
_make_ctx(), strategy='Strat', messages=before, compact=_compact, tokenizer=_tokenizer
|
||||
)
|
||||
assert result is after
|
||||
assert called is False
|
||||
|
||||
Reference in New Issue
Block a user