diff --git a/.konflux/requirements-build.txt b/.konflux/requirements-build.txt index 2ee97da11..361fef270 100644 --- a/.konflux/requirements-build.txt +++ b/.konflux/requirements-build.txt @@ -42,6 +42,7 @@ hatchling==1.33.0 # logfire # logfire-sdk # pydantic-ai + # pydantic-ai-harness # pydantic-ai-skills # pydantic-ai-slim # pydantic-evals @@ -95,6 +96,7 @@ uv-dynamic-versioning==0.14.1 # via # genai-prices # pydantic-ai + # pydantic-ai-harness # pydantic-ai-slim # pydantic-evals # pydantic-graph diff --git a/.konflux/requirements.hashes.source.txt b/.konflux/requirements.hashes.source.txt index 0f27ae2de..495b1de59 100644 --- a/.konflux/requirements.hashes.source.txt +++ b/.konflux/requirements.hashes.source.txt @@ -41,6 +41,9 @@ presidio-anonymizer==2.2.366 \ pydantic-ai==2.38.0 \ --hash=sha256:2dc423fefdaf04998889440404ec84ea9060f7bc9dfadbc94be1f5c8cbcc0324 \ --hash=sha256:90fff4dd3b1db9994861bd102c786d471b81e07a3165c751f403b414bd02e3f9 +pydantic-ai-harness==0.30.0 \ + --hash=sha256:226d35fc3602852e92adf26b5a29717721d66da0f99d5cbff6043bf49ecb0fee \ + --hash=sha256:f9eaa6c19abf47e7af9a84c86124e96c59b432afe8b9f7deae3d9e98fa4677ef pydantic-ai-skills==1.4.0 \ --hash=sha256:7303e0738a837218415f8b2e7bd5b5cbc0ceeb74850b1902c2d079550e0b9b83 \ --hash=sha256:bce8731a042f50c45965acc520dc883163e511c34b208f720eb62067ef5be686 diff --git a/docs/devel_doc/openapi.json b/docs/devel_doc/openapi.json index c582d4272..9ef479985 100644 --- a/docs/devel_doc/openapi.json +++ b/docs/devel_doc/openapi.json @@ -13417,12 +13417,24 @@ "title": "Buffer max ratio", "description": "Maximum fraction of context window the buffer zone can occupy, regardless of buffer_turns.", "default": 0.3 + }, + "summary_prompt": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Summary prompt", + "description": "Prompt template for summarizing older turns in the Pydantic AI compaction. The conversation transcript replaces the {messages} placeholder, which is required; no other replacement field is allowed, so double any literal braces. When unset, the default prompt is used." } }, "additionalProperties": false, "type": "object", "title": "CompactionConfiguration", - "description": "Configuration for conversation history compaction.\n\nCompaction summarizes older conversation turns when their estimated\ntoken count approaches the context window limit, keeping the\nconversation usable instead of failing with HTTP 413. The\nconfiguration here controls when compaction triggers and how much\nrecent context is preserved verbatim.\n\nAttributes:\n enabled: Master switch. When False, compaction never triggers\n and other fields are inert.\n threshold_ratio: Trigger compaction when estimated input tokens\n exceed this fraction of the model's context window\n (clamped to 0.0..1.0).\n token_floor: Minimum estimated token count before compaction\n can trigger, regardless of threshold_ratio. Prevents\n triggering on very small context windows.\n buffer_turns: Initial number of recent turns to keep verbatim.\n The runtime applies a degrading guard \u2014 if these turns\n exceed the available budget, it reduces buffer_turns by\n one repeatedly until the budget fits, down to zero.\n buffer_max_ratio: Hard cap on the fraction of the context\n window the buffer zone may occupy, regardless of\n buffer_turns." + "description": "Configuration for conversation history compaction.\n\nCompaction summarizes older conversation turns when their estimated\ntoken count approaches the context window limit, keeping the\nconversation usable instead of failing with HTTP 413. The\nconfiguration here controls when compaction triggers and how much\nrecent context is preserved verbatim.\n\nAttributes:\n enabled: Master switch. When False, compaction never triggers\n and other fields are inert.\n threshold_ratio: Trigger compaction when estimated input tokens\n exceed this fraction of the model's context window\n (clamped to 0.0..1.0).\n token_floor: Minimum estimated token count before compaction\n can trigger, regardless of threshold_ratio. Prevents\n triggering on very small context windows.\n buffer_turns: Initial number of recent turns to keep verbatim.\n The runtime applies a degrading guard \u2014 if these turns\n exceed the available budget, it reduces buffer_turns by\n one repeatedly until the budget fits, down to zero.\n buffer_max_ratio: Hard cap on the fraction of the context\n window the buffer zone may occupy, regardless of\n buffer_turns.\n summary_prompt: Optional prompt template for summarizing older\n turns in the Pydantic AI compaction. It must contain the\n ``{messages}`` placeholder and no other replacement field;\n when unset, the default prompt is used." }, "Configuration": { "properties": { diff --git a/docs/devel_doc/openapi.md b/docs/devel_doc/openapi.md index a5f8d5dee..e9be0e35a 100644 --- a/docs/devel_doc/openapi.md +++ b/docs/devel_doc/openapi.md @@ -6078,6 +6078,10 @@ Attributes: buffer_max_ratio: Hard cap on the fraction of the context window the buffer zone may occupy, regardless of buffer_turns. + summary_prompt: Optional prompt template for summarizing older + turns in the Pydantic AI compaction. It must contain the + ``{messages}`` placeholder and no other replacement field; + when unset, the default prompt is used. | Field | Type | Description | @@ -6087,6 +6091,7 @@ Attributes: | token_floor | integer | Minimum token count before compaction can trigger. Prevents triggering on very small context windows. | | buffer_turns | integer | Number of recent turns to keep verbatim. | | buffer_max_ratio | number | Maximum fraction of context window the buffer zone can occupy, regardless of buffer_turns. | +| summary_prompt | | Prompt template for summarizing older turns in the Pydantic AI compaction. The conversation transcript replaces the {messages} placeholder, which is required; no other replacement field is allowed, so double any literal braces. When unset, the default prompt is used. | ## Configuration diff --git a/docs/models/successful_responses.json b/docs/models/successful_responses.json index 00978ba70..127de4c1c 100644 --- a/docs/models/successful_responses.json +++ b/docs/models/successful_responses.json @@ -868,7 +868,7 @@ }, "CompactionConfiguration": { "additionalProperties": false, - "description": "Configuration for conversation history compaction.\n\nCompaction summarizes older conversation turns when their estimated\ntoken count approaches the context window limit, keeping the\nconversation usable instead of failing with HTTP 413. The\nconfiguration here controls when compaction triggers and how much\nrecent context is preserved verbatim.\n\nAttributes:\n enabled: Master switch. When False, compaction never triggers\n and other fields are inert.\n threshold_ratio: Trigger compaction when estimated input tokens\n exceed this fraction of the model's context window\n (clamped to 0.0..1.0).\n token_floor: Minimum estimated token count before compaction\n can trigger, regardless of threshold_ratio. Prevents\n triggering on very small context windows.\n buffer_turns: Initial number of recent turns to keep verbatim.\n The runtime applies a degrading guard \u2014 if these turns\n exceed the available budget, it reduces buffer_turns by\n one repeatedly until the budget fits, down to zero.\n buffer_max_ratio: Hard cap on the fraction of the context\n window the buffer zone may occupy, regardless of\n buffer_turns.", + "description": "Configuration for conversation history compaction.\n\nCompaction summarizes older conversation turns when their estimated\ntoken count approaches the context window limit, keeping the\nconversation usable instead of failing with HTTP 413. The\nconfiguration here controls when compaction triggers and how much\nrecent context is preserved verbatim.\n\nAttributes:\n enabled: Master switch. When False, compaction never triggers\n and other fields are inert.\n threshold_ratio: Trigger compaction when estimated input tokens\n exceed this fraction of the model's context window\n (clamped to 0.0..1.0).\n token_floor: Minimum estimated token count before compaction\n can trigger, regardless of threshold_ratio. Prevents\n triggering on very small context windows.\n buffer_turns: Initial number of recent turns to keep verbatim.\n The runtime applies a degrading guard \u2014 if these turns\n exceed the available budget, it reduces buffer_turns by\n one repeatedly until the budget fits, down to zero.\n buffer_max_ratio: Hard cap on the fraction of the context\n window the buffer zone may occupy, regardless of\n buffer_turns.\n summary_prompt: Optional prompt template for summarizing older\n turns in the Pydantic AI compaction. It must contain the\n ``{messages}`` placeholder and no other replacement field;\n when unset, the default prompt is used.", "properties": { "enabled": { "default": false, @@ -901,6 +901,13 @@ "description": "Maximum fraction of context window the buffer zone can occupy, regardless of buffer_turns.", "title": "Buffer max ratio", "type": "number" + }, + "summary_prompt": { + "type": "string", + "nullable": true, + "default": null, + "description": "Prompt template for summarizing older turns in the Pydantic AI compaction. The conversation transcript replaces the {messages} placeholder, which is required; no other replacement field is allowed, so double any literal braces. When unset, the default prompt is used.", + "title": "Summary prompt" } }, "title": "CompactionConfiguration", diff --git a/docs/models/successful_responses.md b/docs/models/successful_responses.md index 7c2dbbbe6..3f246fcc0 100644 --- a/docs/models/successful_responses.md +++ b/docs/models/successful_responses.md @@ -418,6 +418,10 @@ Attributes: buffer_max_ratio: Hard cap on the fraction of the context window the buffer zone may occupy, regardless of buffer_turns. + summary_prompt: Optional prompt template for summarizing older + turns in the Pydantic AI compaction. It must contain the + ``{messages}`` placeholder and no other replacement field; + when unset, the default prompt is used. | Field | Type | Description | @@ -427,6 +431,7 @@ Attributes: | token_floor | integer | Minimum token count before compaction can trigger. Prevents triggering on very small context windows. | | buffer_turns | integer | Number of recent turns to keep verbatim. | | buffer_max_ratio | number | Maximum fraction of context window the buffer zone can occupy, regardless of buffer_turns. | +| summary_prompt | string | Prompt template for summarizing older turns in the Pydantic AI compaction. The conversation transcript replaces the {messages} placeholder, which is required; no other replacement field is allowed, so double any literal braces. When unset, the default prompt is used. | ## Configuration diff --git a/docs/user_doc/config.html b/docs/user_doc/config.html index c26ede4cd..f4447f581 100644 --- a/docs/user_doc/config.html +++ b/docs/user_doc/config.html @@ -555,7 +555,10 @@

CompactionConfiguration

— if these turns exceed the available budget, it reduces buffer_turns by one repeatedly until the budget fits, down to zero. buffer_max_ratio: Hard cap on the fraction of the context window the buffer zone may -occupy, regardless of buffer_turns.

+occupy, regardless of buffer_turns. summary_prompt: Optional prompt +template for summarizing older turns in the Pydantic AI compaction. It +must contain the {messages} placeholder and no other +replacement field; when unset, the default prompt is used.

@@ -599,6 +602,15 @@

CompactionConfiguration

+ + + + +
Maximum fraction of context window the buffer zone can occupy, regardless of buffer_turns.
summary_promptstringPrompt template for summarizing older turns in the Pydantic AI +compaction. The conversation transcript replaces the {messages} +placeholder, which is required; no other replacement field is allowed, +so double any literal braces. When unset, the default prompt is +used.

Configuration

diff --git a/docs/user_doc/config.json b/docs/user_doc/config.json index 92a58a49b..2b5da8e7a 100644 --- a/docs/user_doc/config.json +++ b/docs/user_doc/config.json @@ -389,7 +389,7 @@ }, "CompactionConfiguration": { "additionalProperties": false, - "description": "Configuration for conversation history compaction.\n\nCompaction summarizes older conversation turns when their estimated\ntoken count approaches the context window limit, keeping the\nconversation usable instead of failing with HTTP 413. The\nconfiguration here controls when compaction triggers and how much\nrecent context is preserved verbatim.\n\nAttributes:\n enabled: Master switch. When False, compaction never triggers\n and other fields are inert.\n threshold_ratio: Trigger compaction when estimated input tokens\n exceed this fraction of the model's context window\n (clamped to 0.0..1.0).\n token_floor: Minimum estimated token count before compaction\n can trigger, regardless of threshold_ratio. Prevents\n triggering on very small context windows.\n buffer_turns: Initial number of recent turns to keep verbatim.\n The runtime applies a degrading guard — if these turns\n exceed the available budget, it reduces buffer_turns by\n one repeatedly until the budget fits, down to zero.\n buffer_max_ratio: Hard cap on the fraction of the context\n window the buffer zone may occupy, regardless of\n buffer_turns.", + "description": "Configuration for conversation history compaction.\n\nCompaction summarizes older conversation turns when their estimated\ntoken count approaches the context window limit, keeping the\nconversation usable instead of failing with HTTP 413. The\nconfiguration here controls when compaction triggers and how much\nrecent context is preserved verbatim.\n\nAttributes:\n enabled: Master switch. When False, compaction never triggers\n and other fields are inert.\n threshold_ratio: Trigger compaction when estimated input tokens\n exceed this fraction of the model's context window\n (clamped to 0.0..1.0).\n token_floor: Minimum estimated token count before compaction\n can trigger, regardless of threshold_ratio. Prevents\n triggering on very small context windows.\n buffer_turns: Initial number of recent turns to keep verbatim.\n The runtime applies a degrading guard — if these turns\n exceed the available budget, it reduces buffer_turns by\n one repeatedly until the budget fits, down to zero.\n buffer_max_ratio: Hard cap on the fraction of the context\n window the buffer zone may occupy, regardless of\n buffer_turns.\n summary_prompt: Optional prompt template for summarizing older\n turns in the Pydantic AI compaction. It must contain the\n ``{messages}`` placeholder and no other replacement field;\n when unset, the default prompt is used.", "properties": { "enabled": { "default": false, @@ -422,6 +422,13 @@ "description": "Maximum fraction of context window the buffer zone can occupy, regardless of buffer_turns.", "title": "Buffer max ratio", "type": "number" + }, + "summary_prompt": { + "type": "string", + "nullable": true, + "default": null, + "description": "Prompt template for summarizing older turns in the Pydantic AI compaction. The conversation transcript replaces the {messages} placeholder, which is required; no other replacement field is allowed, so double any literal braces. When unset, the default prompt is used.", + "title": "Summary prompt" } }, "title": "CompactionConfiguration", diff --git a/docs/user_doc/config.md b/docs/user_doc/config.md index 95d5a96d3..bb192924c 100644 --- a/docs/user_doc/config.md +++ b/docs/user_doc/config.md @@ -209,15 +209,20 @@ Attributes: buffer_max_ratio: Hard cap on the fraction of the context window the buffer zone may occupy, regardless of buffer_turns. - - -| Field | Type | Description | -|------------------|---------|-------------- | -| enabled | boolean | When true, older conversation turns are summarized when estimated tokens approach the context window limit. | -| threshold_ratio | number | Trigger compaction when estimated tokens exceed this fraction of the model's context window (0.0-1.0). | -| token_floor | integer | Minimum token count before compaction can trigger. Prevents triggering on very small context windows. | -| buffer_turns | integer | Number of recent turns to keep verbatim. | -| buffer_max_ratio | number | Maximum fraction of context window the buffer zone can occupy, regardless of buffer_turns. | + summary_prompt: Optional prompt template for summarizing older + turns in the Pydantic AI compaction. It must contain the + ``{messages}`` placeholder and no other replacement field; + when unset, the default prompt is used. + + +| Field | Type | Description | +|------------------|---------|-------------- | +| enabled | boolean | When true, older conversation turns are summarized when estimated tokens approach the context window limit. | +| threshold_ratio | number | Trigger compaction when estimated tokens exceed this fraction of the model's context window (0.0-1.0). | +| token_floor | integer | Minimum token count before compaction can trigger. Prevents triggering on very small context windows. | +| buffer_turns | integer | Number of recent turns to keep verbatim. | +| buffer_max_ratio | number | Maximum fraction of context window the buffer zone can occupy, regardless of buffer_turns. | +| summary_prompt | string | Prompt template for summarizing older turns in the Pydantic AI compaction. The conversation transcript replaces the {messages} placeholder, which is required; no other replacement field is allowed, so double any literal braces. When unset, the default prompt is used. | ## Configuration diff --git a/pyproject.toml b/pyproject.toml index 01a53ea07..f89127c51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,6 +85,9 @@ dependencies = [ "genai-prices>=0.1.12", # Used for Pydantic AI "pydantic-ai==2.38.0", + # Used for conversation compaction (TieredCompaction, LCORE-4692); 0.30.0 is the + # newest release that supports pydantic-ai 2.38.x (0.31+ needs pydantic-ai-slim>=2.40) + "pydantic-ai-harness==0.30.0", "pydantic-ai-skills>=0.11.0,<2.0.0", # Used for OpenTelemetry instrumentation "opentelemetry-distro>=0.49b0", diff --git a/src/models/config.py b/src/models/config.py index 63218416a..4c5bf38e2 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -8,6 +8,7 @@ from functools import cached_property from pathlib import Path from re import Pattern +from string import Formatter from typing import Annotated, Any, Literal, Optional, Self import jsonpath_ng @@ -2164,6 +2165,10 @@ class CompactionConfiguration(ConfigurationBase): buffer_max_ratio: Hard cap on the fraction of the context window the buffer zone may occupy, regardless of buffer_turns. + summary_prompt: Optional prompt template for summarizing older + turns in the Pydantic AI compaction. It must contain the + ``{messages}`` placeholder and no other replacement field; + when unset, the default prompt is used. """ @@ -2196,6 +2201,14 @@ class CompactionConfiguration(ConfigurationBase): description="Maximum fraction of context window the buffer zone " "can occupy, regardless of buffer_turns.", ) + summary_prompt: Optional[str] = Field( + default=None, + title="Summary prompt", + description="Prompt template for summarizing older turns in the Pydantic " + "AI compaction. The conversation transcript replaces the {messages} " + "placeholder, which is required; no other replacement field is allowed, " + "so double any literal braces. When unset, the default prompt is used.", + ) @field_validator("threshold_ratio") @classmethod @@ -2213,6 +2226,44 @@ def _validate_buffer_max_ratio(cls, value: float) -> float: raise ValueError("buffer_max_ratio must be between 0.0 and 1.0 (inclusive)") return value + @field_validator("summary_prompt") + @classmethod + def _validate_summary_prompt(cls, value: Optional[str]) -> Optional[str]: + """Check that a summary prompt formats the transcript and nothing else. + + The prompt is filled with ``str.format(messages=...)``. Without the + placeholder the transcript would be silently dropped, and any other + replacement field would fail only when a conversation is summarized, + so both are rejected when the configuration is loaded. + + Parameters: + value: Prompt template, or None for the default prompt. + + Returns: + The prompt, unchanged. + + Raises: + ValueError: If the prompt does not contain ``{messages}``, contains + any other replacement field, or is not a valid format string. + + """ + if value is None: + return value + try: + fields = [ + (name, spec, conversion) + for _, name, spec, conversion in Formatter().parse(value) + if name is not None + ] + except ValueError as e: + raise ValueError(f"summary_prompt is not a valid format string: {e}") from e + if not fields or any(field != ("messages", "", None) for field in fields): + raise ValueError( + "summary_prompt must contain the {messages} placeholder and no other " + "replacement field; double any literal braces" + ) + return value + class ConversationHistoryConfiguration(ConfigurationBase): """Conversation history configuration.""" diff --git a/src/utils/README.md b/src/utils/README.md index 85f3c3ed3..5a4e62cd5 100644 --- a/src/utils/README.md +++ b/src/utils/README.md @@ -124,6 +124,10 @@ SSE formatting helpers for streaming query responses. Session ID utility functions. +## [tiered_compaction.py](tiered_compaction.py) + +Build the harness ``TieredCompaction`` capability from the compaction config. + ## [token_counter.py](token_counter.py) Helper classes to count tokens sent and received by the LLM. diff --git a/src/utils/tiered_compaction.py b/src/utils/tiered_compaction.py new file mode 100644 index 000000000..2e9713f17 --- /dev/null +++ b/src/utils/tiered_compaction.py @@ -0,0 +1,102 @@ +"""Build the harness ``TieredCompaction`` capability from the compaction config. + +Conversation compaction after the Pydantic AI cutover runs as one agent +capability: ``TieredCompaction`` from ``pydantic-ai-harness``. Before every +model request it estimates the history size and, when the history exceeds +``threshold_ratio`` of the context window, applies three tiers in a fixed +order until the history fits: clamp oversized messages, clear old tool +results, then summarize the older history with the chat model. + +The builder is not wired into ``build_agent`` or any request handler yet; +the cutover attaches it. With ``keep_user_messages=True``, harness 0.30.0 +gives retained user turns the ``keep_messages`` tail slots and drops the +unsummarized messages that no longer fit. At the default configuration this +usually includes the request being answered, so this must be resolved +before the cutover attaches the capability. ``token_floor`` has no harness +equivalent and is ignored, and ``buffer_turns`` becomes a count of messages, +not turn pairs. +""" + +from dataclasses import replace +from typing import Final, Optional + +from pydantic_ai.models import Model +from pydantic_ai_harness.compaction import ( + ClampOversizedMessages, + ClearToolResults, + SummarizingCompaction, + TieredCompaction, +) + +from constants import DEFAULT_CONTEXT_WINDOW +from models.config import CompactionConfiguration +from utils.token_estimator import estimate_tokens + +# The harness requires a clamp threshold and the configuration has no knob for +# it; 50_000 is the value the harness documentation uses. +CLAMP_MAX_PART_TOKENS: Final[int] = 50_000 +"""Size in tokens above which an assistant text or tool-call part is clamped.""" + +# TieredCompaction decides when to compact and calls each tier directly, so a +# tier's own trigger is never read; the harness still requires one to be set. +_UNUSED_TIER_TRIGGER: Final[int] = 1 + + +def build_tiered_compaction( + config: CompactionConfiguration, + model: Model | str, + context_window: Optional[int] = None, +) -> Optional[TieredCompaction[object]]: + """Build the tiered compaction capability for an agent. + + Parameters: + config: Compaction configuration. Its ``summary_prompt``, already + validated when the configuration was loaded, replaces the harness + default prompt of the summarizing tier when set. + model: Chat model that answers the user query; the summarizing tier + uses it too. + context_window: Context window of the model in tokens, when known. + When omitted, the harness resolves it from the model name and + falls back to ``DEFAULT_CONTEXT_WINDOW``. + + Returns: + The capability, or None when compaction is disabled. + + Raises: + ValueError: If the harness rejects the configuration (for example + ``threshold_ratio`` of 0). + + """ + if not config.enabled: + return None + summarizer = SummarizingCompaction[object]( + model=model, + max_messages=_UNUSED_TIER_TRIGGER, + keep_messages=config.buffer_turns, + keep_tokens=( + None + if context_window is None + else int(context_window * config.buffer_max_ratio) + ), + tokenizer=estimate_tokens, + incremental=True, + preserve_first_user_message=True, + keep_user_messages=True, + ) + if config.summary_prompt is not None: + summarizer = replace(summarizer, summary_prompt=config.summary_prompt) + return TieredCompaction[object]( + tiers=[ + ClampOversizedMessages[object]( + max_part_tokens=CLAMP_MAX_PART_TOKENS, tokenizer=estimate_tokens + ), + ClearToolResults[object]( + max_tokens=_UNUSED_TIER_TRIGGER, tokenizer=estimate_tokens + ), + summarizer, + ], + target_fraction=config.threshold_ratio, + context_window=context_window, + fallback_context_window=DEFAULT_CONTEXT_WINDOW, + tokenizer=estimate_tokens, + ) diff --git a/tests/unit/models/config/test_compaction_configuration.py b/tests/unit/models/config/test_compaction_configuration.py index 1d63910c2..54bf54e5f 100644 --- a/tests/unit/models/config/test_compaction_configuration.py +++ b/tests/unit/models/config/test_compaction_configuration.py @@ -1,6 +1,7 @@ """Unit tests for CompactionConfiguration and its placement on Configuration.""" import pytest +from pydantic import ValidationError from models.config import CompactionConfiguration, Configuration @@ -17,6 +18,7 @@ def test_default_values() -> None: assert config.token_floor == 4096 assert config.buffer_turns == 4 assert config.buffer_max_ratio == 0.3 + assert config.summary_prompt is None def test_enabled_can_be_turned_on() -> None: @@ -87,6 +89,35 @@ def test_buffer_turns_zero_is_allowed() -> None: assert CompactionConfiguration(buffer_turns=0).buffer_turns == 0 +@pytest.mark.parametrize( + "summary_prompt", + ["Summarize for support:\n{messages}", "Keep {{literal}} braces.\n{messages}"], +) +def test_summary_prompt_accepts_messages_placeholder(summary_prompt: str) -> None: + """A prompt whose only replacement field is {messages} is kept unchanged.""" + config = CompactionConfiguration.model_validate({"summary_prompt": summary_prompt}) + assert config.summary_prompt == summary_prompt + + +@pytest.mark.parametrize( + ("summary_prompt", "message"), + [ + ("Summarize the conversation.", "must contain the {messages} placeholder"), + ("", "must contain the {messages} placeholder"), + ("{messages} about {product}", "no other replacement field"), + ("{messages:.100}", "no other replacement field"), + ("{messages!r}", "no other replacement field"), + ("{messages", "not a valid format string"), + ], +) +def test_summary_prompt_rejects_invalid_template( + summary_prompt: str, message: str +) -> None: + """A prompt without {messages}, or with any other field, fails to load.""" + with pytest.raises(ValidationError, match=message): + CompactionConfiguration.model_validate({"summary_prompt": summary_prompt}) + + def test_rejects_unknown_field() -> None: """Unknown fields are forbidden via ConfigurationBase's extra='forbid'.""" with pytest.raises(ValueError): diff --git a/tests/unit/models/config/test_dump_configuration.py b/tests/unit/models/config/test_dump_configuration.py index 3182ec105..1313dc16e 100644 --- a/tests/unit/models/config/test_dump_configuration.py +++ b/tests/unit/models/config/test_dump_configuration.py @@ -224,6 +224,7 @@ def test_dump_configuration_minimal_cfg(tmp_path: Path) -> None: "token_floor": 4096, "buffer_turns": 4, "buffer_max_ratio": 0.3, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { @@ -465,6 +466,7 @@ def test_dump_configuration_valid_values(tmp_path: Path) -> None: "token_floor": 4096, "buffer_turns": 4, "buffer_max_ratio": 0.3, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { @@ -845,6 +847,7 @@ def test_dump_configuration_with_quota_limiters(tmp_path: Path) -> None: "token_floor": 4096, "buffer_turns": 4, "buffer_max_ratio": 0.3, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { @@ -1123,6 +1126,7 @@ def test_dump_configuration_with_quota_limiters_different_values( "token_floor": 4096, "buffer_turns": 4, "buffer_max_ratio": 0.3, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { @@ -1438,6 +1442,7 @@ def test_dump_configuration_byok(tmp_path: Path) -> None: "token_floor": 4096, "buffer_turns": 4, "buffer_max_ratio": 0.3, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { @@ -1694,6 +1699,7 @@ def test_dump_configuration_pg_namespace(tmp_path: Path) -> None: "token_floor": 4096, "buffer_turns": 4, "buffer_max_ratio": 0.3, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { @@ -2093,6 +2099,7 @@ def test_dump_configuration_allow_degraded_mode(tmp_path: Path) -> None: "token_floor": 4096, "buffer_turns": 4, "buffer_max_ratio": 0.3, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { @@ -2338,6 +2345,7 @@ def test_dump_configuration_max_retries_settings(tmp_path: Path) -> None: "token_floor": 4096, "buffer_turns": 4, "buffer_max_ratio": 0.3, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { @@ -2583,6 +2591,7 @@ def test_dump_configuration_retry_count_settings(tmp_path: Path) -> None: "token_floor": 4096, "buffer_turns": 4, "buffer_max_ratio": 0.3, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { @@ -2835,6 +2844,7 @@ def test_dump_configuration_specific_compaction_values(tmp_path: Path) -> None: "token_floor": 1024, "buffer_turns": 8, "buffer_max_ratio": 0.5, + "summary_prompt": None, }, "approvals": _DEFAULT_APPROVALS_DUMP, "vector_store": { diff --git a/tests/unit/utils/README.md b/tests/unit/utils/README.md index fa88bea01..893b6bbd1 100644 --- a/tests/unit/utils/README.md +++ b/tests/unit/utils/README.md @@ -112,6 +112,10 @@ Unit tests for utils/streaming_sse.py. Unit tests for functions defined in utils.suid module. +## [test_tiered_compaction.py](test_tiered_compaction.py) + +Unit tests for utils/tiered_compaction. + ## [test_token_estimator.py](test_token_estimator.py) Unit tests for utils/token_estimator. diff --git a/tests/unit/utils/test_tiered_compaction.py b/tests/unit/utils/test_tiered_compaction.py new file mode 100644 index 000000000..e8b745513 --- /dev/null +++ b/tests/unit/utils/test_tiered_compaction.py @@ -0,0 +1,169 @@ +"""Unit tests for utils/tiered_compaction.""" + +from typing import Optional + +import pytest +from pydantic_ai import Agent +from pydantic_ai.messages import ( + ModelMessage, + ModelRequest, + ModelResponse, + TextPart, + UserPromptPart, +) +from pydantic_ai.models.function import AgentInfo, FunctionModel +from pydantic_ai.models.test import TestModel +from pydantic_ai_harness.compaction import ( + ClampOversizedMessages, + ClearToolResults, + SummarizingCompaction, + TieredCompaction, +) + +from constants import DEFAULT_CONTEXT_WINDOW +from models.config import CompactionConfiguration +from utils.tiered_compaction import CLAMP_MAX_PART_TOKENS, build_tiered_compaction +from utils.token_estimator import estimate_tokens + +ENABLED = CompactionConfiguration( + enabled=True, threshold_ratio=0.5, buffer_turns=6, buffer_max_ratio=0.25 +) + + +def _build( + config: CompactionConfiguration = ENABLED, + model: Optional[TestModel] = None, + context_window: Optional[int] = None, +) -> TieredCompaction[object]: + """Build the capability for an enabled configuration and assert it exists.""" + compaction = build_tiered_compaction(config, model or TestModel(), context_window) + assert compaction is not None + return compaction + + +def _summarizer(compaction: TieredCompaction[object]) -> SummarizingCompaction[object]: + """Return the summarizing tier, which the builder puts last.""" + summarizer = compaction.tiers[-1] + assert isinstance(summarizer, SummarizingCompaction) + return summarizer + + +def test_disabled_configuration_builds_nothing() -> None: + """A disabled configuration yields no capability, whatever else is passed.""" + config = CompactionConfiguration(enabled=False, summary_prompt="{messages}") + assert build_tiered_compaction(config, TestModel(), 128_000) is None + + +def test_tiers_come_in_fixed_order() -> None: + """Clamp, then clear tool results, then summarize.""" + tiers = _build().tiers + assert [type(tier) for tier in tiers] == [ + ClampOversizedMessages, + ClearToolResults, + SummarizingCompaction, + ] + assert isinstance(tiers[0], ClampOversizedMessages) + assert tiers[0].max_part_tokens == CLAMP_MAX_PART_TOKENS + + +@pytest.mark.parametrize("context_window", [None, 32_000]) +def test_threshold_and_context_window_mapping(context_window: Optional[int]) -> None: + """threshold_ratio becomes target_fraction; the window passes through.""" + compaction = _build(context_window=context_window) + assert compaction.target_fraction == 0.5 + assert compaction.target_tokens is None + assert compaction.context_window == context_window + assert compaction.fallback_context_window == DEFAULT_CONTEXT_WINDOW + + +@pytest.mark.parametrize( + ("context_window", "expected_keep_tokens"), + [(None, None), (10_000, 2_500), (999, 249)], +) +def test_summarizing_tier_mapping( + context_window: Optional[int], expected_keep_tokens: Optional[int] +) -> None: + """buffer_turns and buffer_max_ratio map onto the summarizing tier.""" + model = TestModel() + summarizer = _summarizer(_build(model=model, context_window=context_window)) + assert summarizer.model is model + assert summarizer.keep_messages == 6 + assert summarizer.keep_tokens == expected_keep_tokens + assert summarizer.incremental + assert summarizer.preserve_first_user_message + assert summarizer.keep_user_messages + + +def test_estimate_tokens_is_the_tokenizer_everywhere() -> None: + """The orchestrator and every tier count tokens with estimate_tokens.""" + compaction = _build(context_window=32_000) + assert compaction.tokenizer is estimate_tokens + for tier in compaction.tiers: + assert isinstance( + tier, (ClampOversizedMessages, ClearToolResults, SummarizingCompaction) + ) + assert tier.tokenizer is estimate_tokens + + +def test_token_floor_is_ignored() -> None: + """Builds that differ only in token_floor are identical.""" + model = TestModel() + low = ENABLED.model_copy(update={"token_floor": 0}) + high = ENABLED.model_copy(update={"token_floor": 100_000}) + assert _build(low, model, 32_000) == _build(high, model, 32_000) + + +def test_zero_threshold_ratio_is_rejected() -> None: + """The harness cannot compact at a zero fraction of the window.""" + config = ENABLED.model_copy(update={"threshold_ratio": 0.0}) + with pytest.raises(ValueError, match="target_fraction"): + build_tiered_compaction(config, TestModel()) + + +def test_default_summary_prompt_is_kept() -> None: + """Without a configured prompt the harness default stays in place.""" + default = SummarizingCompaction[object](max_messages=1).summary_prompt + assert _summarizer(_build()).summary_prompt == default + + +@pytest.mark.parametrize( + "summary_prompt", + ["Summarize for support:\n{messages}", "Keep {{literal}} braces.\n{messages}"], +) +def test_configured_summary_prompt_is_used(summary_prompt: str) -> None: + """The configured prompt replaces the harness default.""" + config = CompactionConfiguration.model_validate( + ENABLED.model_dump() | {"summary_prompt": summary_prompt} + ) + summarizer = _summarizer(_build(config)) + assert summarizer.summary_prompt == summary_prompt + + +@pytest.mark.asyncio +async def test_capability_summarizes_history_inside_an_agent() -> None: + """An agent with the capability sends the summary, not the old turns.""" + sent: list[list[ModelMessage]] = [] + + def chat(messages: list[ModelMessage], _: AgentInfo) -> ModelResponse: + """Record what the chat model receives and answer.""" + sent.append(messages) + return ModelResponse(parts=[TextPart(content="answer")]) + + # User turns outgrow keep_tokens (250): none is kept, so tail loss spares the query. + history: list[ModelMessage] = [] + for turn in range(3): + history.append( + ModelRequest(parts=[UserPromptPart(content=f"q{turn} " + "x " * 300)]) + ) + history.append( + ModelResponse(parts=[TextPart(content=f"a{turn} " + "y " * 300)]) + ) + compaction = _build( + model=TestModel(custom_output_text="SUMMARY"), context_window=1_000 + ) + await Agent(FunctionModel(chat), capabilities=[compaction]).run( + "current question", message_history=history + ) + assert "SUMMARY" in repr(sent[-1]) + assert "current question" in repr(sent[-1]) + assert "q0 " not in repr(sent[-1]) diff --git a/uv.lock b/uv.lock index 4573cff11..f5c08af22 100644 --- a/uv.lock +++ b/uv.lock @@ -2436,6 +2436,7 @@ dependencies = [ { name = "psycopg2-binary" }, { name = "pyasn1" }, { name = "pydantic-ai" }, + { name = "pydantic-ai-harness" }, { name = "pydantic-ai-skills" }, { name = "pyjwt" }, { name = "python-dotenv" }, @@ -2559,6 +2560,7 @@ requires-dist = [ { name = "psycopg2-binary", specifier = ">=2.9.10" }, { name = "pyasn1", specifier = ">=0.6.3" }, { name = "pydantic-ai", specifier = "==2.38.0" }, + { name = "pydantic-ai-harness", specifier = "==0.30.0" }, { name = "pydantic-ai-skills", specifier = ">=0.11.0,<2.0.0" }, { name = "pyjwt", specifier = ">=2.15.1" }, { name = "python-dotenv", specifier = ">=1.2.2" }, @@ -4386,6 +4388,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/69/3a/b552b097d57a45ecb2a0a972f830b77428ae900b41c86880303c66f59369/pydantic_ai-2.38.0-py3-none-any.whl", hash = "sha256:2dc423fefdaf04998889440404ec84ea9060f7bc9dfadbc94be1f5c8cbcc0324", size = 9454, upload-time = "2026-09-03T07:57:18.519Z" }, ] +[[package]] +name = "pydantic-ai-harness" +version = "0.30.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "genai-prices" }, + { name = "httpx" }, + { name = "pydantic-ai-slim" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/eb/18/2f1950a8e3421b7a1616ca3e4fafccbe5c2b707f2e6cc72e1ed42db086a7/pydantic_ai_harness-0.30.0.tar.gz", hash = "sha256:226d35fc3602852e92adf26b5a29717721d66da0f99d5cbff6043bf49ecb0fee", size = 2560749, upload-time = "2026-09-09T03:49:00.64Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/07/8ff61abd60e794f739ee7588fbedb50d0d61af1e47e78b5a0c7372bd806c/pydantic_ai_harness-0.30.0-py3-none-any.whl", hash = "sha256:f9eaa6c19abf47e7af9a84c86124e96c59b432afe8b9f7deae3d9e98fa4677ef", size = 804266, upload-time = "2026-09-09T03:48:58.759Z" }, +] + [[package]] name = "pydantic-ai-skills" version = "1.4.0"