Skip to content

Commit 3e03a30

Browse files
eavanvalkenburgCopilot
eavanvalkenburg
and
Copilot
authored
Python: Implement annotation-based context compaction (#4469)
* Implement annotation-based context compaction Co-authored-by: Copilot <[email protected]> * Handle missing compaction attributes in BaseChatClient Co-authored-by: Copilot <[email protected]> * Fix CI typing and bandit issues Co-authored-by: Copilot <[email protected]> * Optimize incremental compaction annotation pass Co-authored-by: Copilot <[email protected]> * refinement * Python: add ToolResultCompactionStrategy and CompactionProvider Add ToolResultCompactionStrategy that collapses older tool-call groups into short summary messages (e.g. [Tool calls: get_weather]) while keeping the most recent groups verbatim. This mirrors the .NET ToolResultCompactionStrategy from PR #4533. Add CompactionProvider as a context-provider that auto-applies compaction before each agent turn and stores compacted history in session state after each turn. Includes tests and samples for both features. Co-authored-by: Copilot <[email protected]> * refinement and alignment with dotnet PR * updated tool result compaction * updated tool result compaction * Python: add ToolResultCompactionStrategy, CompactionProvider, and skip_excluded - ToolResultCompactionStrategy collapses older tool-call groups into [Tool results: func_name: result] summaries with bidirectional tracing (same pattern as SummarizationStrategy). - CompactionProvider as BaseContextProvider with separate before_strategy and after_strategy parameters. before_strategy compacts loaded context; after_strategy compacts stored history via history_source_id. - InMemoryHistoryProvider gains skip_excluded flag to filter out messages marked as excluded by compaction strategies. - Tests, samples, and exports updated. Co-authored-by: Copilot <[email protected]> * fixed checks * fix mypy * Fix: ensure summary messages from both strategies get full compaction annotations SummarizationStrategy was not calling annotate_message_groups after inserting its summary message, so the summary lacked core group annotations (id, kind, index, has_reasoning, _excluded). Added the missing call. ToolResultCompactionStrategy already had it. Added tests verifying both strategies produce fully annotated summaries. Co-authored-by: Copilot <[email protected]> * updated propagation * fix mypy --------- Co-authored-by: Copilot <[email protected]>
1 parent 565c0b1 commit 3e03a30

29 files changed

Lines changed: 4401 additions & 209 deletions

‎docs/decisions/0019-python-context-compaction-strategy.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1240,3 +1240,10 @@ class AttributionAwareStrategy(CompactionStrategy):
12401240

12411241
- [ADR-0016: Unifying Context Management with ContextPlugin](0016-python-context-middleware.md) — Parent ADR that established `ContextProvider`, `HistoryProvider`, and `AgentSession` architecture.
12421242
- [Context Compaction Limitations Analysis](https://gist.github.com/victordibia/ec3f3baf97345f7e47da025cf55b999f) — Detailed analysis of why current architecture cannot support in-run compaction, with attempted solutions and their failure modes. Option 4 in this ADR corresponds to "Option A: Middleware Access to Mutable Message Source" from that analysis; Options 1-3 correspond to "Option B: Tool Loop Hook", adapted here to a `BaseChatClient` hook instead of `FunctionInvocationConfiguration`.
1243+
1244+
### Implementation Rollout Note
1245+
1246+
Implementation is split into two phases:
1247+
1248+
1. **Phase 1 (PR 1):** runtime compaction foundation in `agent_framework/_compaction.py`, in-run integration, and extensive core tests, plus in-run compaction samples (`basics`, `advanced`, `custom`).
1249+
2. **Phase 2 (PR 2):** history/storage compaction (`upsert`-based full replacement), provider support, storage tests, and storage-focused sample (`storage`).

‎python/packages/azure-ai-search/tests/test_aisearch_context_provider.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,13 @@
1616
# -- Helpers -------------------------------------------------------------------
1717

1818

19+
@pytest.fixture(autouse=True)
20+
def clear_azure_search_environment(monkeypatch: pytest.MonkeyPatch) -> None:
21+
for key in tuple(os.environ):
22+
if key.startswith("AZURE_SEARCH_"):
23+
monkeypatch.delenv(key, raising=False)
24+
25+
1926
class MockSearchResults:
2027
"""Async-iterable mock for Azure SearchClient.search() results."""
2128

‎python/packages/core/agent_framework/__init__.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,34 @@
2929
SupportsMCPTool,
3030
SupportsWebSearchTool,
3131
)
32+
from ._compaction import (
33+
COMPACTION_STATE_KEY,
34+
EXCLUDE_REASON_KEY,
35+
EXCLUDED_KEY,
36+
GROUP_ANNOTATION_KEY,
37+
GROUP_HAS_REASONING_KEY,
38+
GROUP_ID_KEY,
39+
GROUP_INDEX_KEY,
40+
GROUP_KIND_KEY,
41+
GROUP_TOKEN_COUNT_KEY,
42+
SUMMARIZED_BY_SUMMARY_ID_KEY,
43+
SUMMARY_OF_GROUP_IDS_KEY,
44+
SUMMARY_OF_MESSAGE_IDS_KEY,
45+
CharacterEstimatorTokenizer,
46+
CompactionProvider,
47+
CompactionStrategy,
48+
SelectiveToolCallCompactionStrategy,
49+
SlidingWindowStrategy,
50+
SummarizationStrategy,
51+
TokenBudgetComposedStrategy,
52+
TokenizerProtocol,
53+
ToolResultCompactionStrategy,
54+
TruncationStrategy,
55+
annotate_message_groups,
56+
apply_compaction,
57+
included_messages,
58+
included_token_count,
59+
)
3260
from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPWebsocketTool
3361
from ._middleware import (
3462
AgentContext,
@@ -196,7 +224,19 @@
196224
__all__ = [
197225
"AGENT_FRAMEWORK_USER_AGENT",
198226
"APP_INFO",
227+
"COMPACTION_STATE_KEY",
199228
"DEFAULT_MAX_ITERATIONS",
229+
"EXCLUDED_KEY",
230+
"EXCLUDE_REASON_KEY",
231+
"GROUP_ANNOTATION_KEY",
232+
"GROUP_HAS_REASONING_KEY",
233+
"GROUP_ID_KEY",
234+
"GROUP_INDEX_KEY",
235+
"GROUP_KIND_KEY",
236+
"GROUP_TOKEN_COUNT_KEY",
237+
"SUMMARIZED_BY_SUMMARY_ID_KEY",
238+
"SUMMARY_OF_GROUP_IDS_KEY",
239+
"SUMMARY_OF_MESSAGE_IDS_KEY",
200240
"USER_AGENT_KEY",
201241
"USER_AGENT_TELEMETRY_DISABLED_ENV_VAR",
202242
"Agent",
@@ -218,6 +258,7 @@
218258
"BaseEmbeddingClient",
219259
"BaseHistoryProvider",
220260
"Case",
261+
"CharacterEstimatorTokenizer",
221262
"ChatAndFunctionMiddlewareTypes",
222263
"ChatContext",
223264
"ChatMiddleware",
@@ -227,6 +268,8 @@
227268
"ChatResponse",
228269
"ChatResponseUpdate",
229270
"CheckpointStorage",
271+
"CompactionProvider",
272+
"CompactionStrategy",
230273
"Content",
231274
"ContinuationToken",
232275
"Default",
@@ -273,15 +316,18 @@
273316
"Runner",
274317
"RunnerContext",
275318
"SecretString",
319+
"SelectiveToolCallCompactionStrategy",
276320
"SessionContext",
277321
"SingleEdgeGroup",
278322
"Skill",
279323
"SkillResource",
280324
"SkillScript",
281325
"SkillScriptRunner",
282326
"SkillsProvider",
327+
"SlidingWindowStrategy",
283328
"SubWorkflowRequestMessage",
284329
"SubWorkflowResponseMessage",
330+
"SummarizationStrategy",
285331
"SupportsAgentRun",
286332
"SupportsChatGetResponse",
287333
"SupportsCodeInterpreterTool",
@@ -294,8 +340,12 @@
294340
"SwitchCaseEdgeGroupCase",
295341
"SwitchCaseEdgeGroupDefault",
296342
"TextSpanRegion",
343+
"TokenBudgetComposedStrategy",
344+
"TokenizerProtocol",
297345
"ToolMode",
346+
"ToolResultCompactionStrategy",
298347
"ToolTypes",
348+
"TruncationStrategy",
299349
"TypeCompatibilityError",
300350
"UpdateT",
301351
"UsageDetails",
@@ -322,12 +372,16 @@
322372
"__version__",
323373
"add_usage_details",
324374
"agent_middleware",
375+
"annotate_message_groups",
376+
"apply_compaction",
325377
"chat_middleware",
326378
"create_edge_runner",
327379
"detect_media_type_from_base64",
328380
"executor",
329381
"function_middleware",
330382
"handler",
383+
"included_messages",
384+
"included_token_count",
331385
"load_settings",
332386
"map_chat_to_agent_update",
333387
"merge_chat_options",

‎python/packages/core/agent_framework/_agents.py‎

Lines changed: 45 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@
7474
from typing_extensions import Self, TypedDict # pragma: no cover
7575

7676
if TYPE_CHECKING:
77+
from ._compaction import CompactionStrategy, TokenizerProtocol
7778
from ._types import ChatOptions
7879

7980
logger = logging.getLogger("agent_framework")
@@ -177,6 +178,8 @@ class _RunContext(TypedDict):
177178
session_messages: Sequence[Message]
178179
agent_name: str
179180
chat_options: MutableMapping[str, Any]
181+
compaction_strategy: CompactionStrategy | None
182+
tokenizer: TokenizerProtocol | None
180183
filtered_kwargs: Mapping[str, Any]
181184
finalize_kwargs: Mapping[str, Any]
182185

@@ -665,6 +668,8 @@ def __init__(
665668
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
666669
default_options: OptionsCoT | None = None,
667670
context_providers: Sequence[BaseContextProvider] | None = None,
671+
compaction_strategy: CompactionStrategy | None = None,
672+
tokenizer: TokenizerProtocol | None = None,
668673
**kwargs: Any,
669674
) -> None:
670675
"""Initialize a Agent instance.
@@ -688,6 +693,10 @@ def __init__(
688693
Note: response_format typing does not flow into run outputs when set via default_options.
689694
These can be overridden at runtime via the ``options`` parameter of ``run()``.
690695
tools: The tools to use for the request.
696+
compaction_strategy: Optional agent-level in-run compaction.
697+
If both this and a compaction_strategy on the underlying client are set, this one is used.
698+
tokenizer: Optional agent-level tokenizer.
699+
If both this and a tokenizer on the underlying client are set, this one is used.
691700
kwargs: Any additional keyword arguments. Will be stored as ``additional_properties``.
692701
"""
693702
opts = dict(default_options) if default_options else {}
@@ -705,6 +714,8 @@ def __init__(
705714
**kwargs,
706715
)
707716
self.client = client
717+
self.compaction_strategy = compaction_strategy
718+
self.tokenizer = tokenizer
708719

709720
# Get tools from options or named parameter (named param takes precedence)
710721
tools_ = tools if tools is not None else opts.pop("tools", None)
@@ -799,6 +810,8 @@ def run(
799810
session: AgentSession | None = None,
800811
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
801812
options: ChatOptions[ResponseModelBoundT],
813+
compaction_strategy: CompactionStrategy | None = None,
814+
tokenizer: TokenizerProtocol | None = None,
802815
**kwargs: Any,
803816
) -> Awaitable[AgentResponse[ResponseModelBoundT]]: ...
804817

@@ -811,6 +824,8 @@ def run(
811824
session: AgentSession | None = None,
812825
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
813826
options: OptionsCoT | ChatOptions[None] | None = None,
827+
compaction_strategy: CompactionStrategy | None = None,
828+
tokenizer: TokenizerProtocol | None = None,
814829
**kwargs: Any,
815830
) -> Awaitable[AgentResponse[Any]]: ...
816831

@@ -823,6 +838,8 @@ def run(
823838
session: AgentSession | None = None,
824839
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
825840
options: OptionsCoT | ChatOptions[Any] | None = None,
841+
compaction_strategy: CompactionStrategy | None = None,
842+
tokenizer: TokenizerProtocol | None = None,
826843
**kwargs: Any,
827844
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
828845

@@ -834,6 +851,8 @@ def run(
834851
session: AgentSession | None = None,
835852
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
836853
options: OptionsCoT | ChatOptions[Any] | None = None,
854+
compaction_strategy: CompactionStrategy | None = None,
855+
tokenizer: TokenizerProtocol | None = None,
837856
**kwargs: Any,
838857
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
839858
"""Run the agent with the given messages and options.
@@ -857,8 +876,14 @@ def run(
857876
``Agent[OpenAIChatOptions]``, this enables IDE autocomplete for
858877
provider-specific options including temperature, max_tokens, model_id,
859878
tool_choice, and provider-specific options like reasoning_effort.
860-
kwargs: Additional keyword arguments for the agent.
861-
Will only be passed to functions that are called.
879+
compaction_strategy: Optional per-run compaction override passed to
880+
``client.get_response()``. When omitted, the agent-level override
881+
is used, falling back to the client default.
882+
tokenizer: Optional per-run tokenizer override passed to
883+
``client.get_response()``. When omitted, the agent-level override
884+
is used, falling back to the client default.
885+
kwargs: Additional keyword arguments for the agent. These are only
886+
passed to functions that are called.
862887
863888
Returns:
864889
When stream=False: An Awaitable[AgentResponse] containing the agent's response.
@@ -873,6 +898,8 @@ async def _run_non_streaming() -> AgentResponse[Any]:
873898
session=session,
874899
tools=tools,
875900
options=options,
901+
compaction_strategy=compaction_strategy,
902+
tokenizer=tokenizer,
876903
kwargs=kwargs,
877904
)
878905
response = cast(
@@ -881,6 +908,8 @@ async def _run_non_streaming() -> AgentResponse[Any]:
881908
messages=ctx["session_messages"],
882909
stream=False,
883910
options=ctx["chat_options"], # type: ignore[reportArgumentType]
911+
compaction_strategy=ctx["compaction_strategy"],
912+
tokenizer=ctx["tokenizer"],
884913
**ctx["filtered_kwargs"],
885914
),
886915
)
@@ -954,13 +983,17 @@ async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]
954983
session=session,
955984
tools=tools,
956985
options=options,
986+
compaction_strategy=compaction_strategy,
987+
tokenizer=tokenizer,
957988
kwargs=kwargs,
958989
)
959990
ctx: _RunContext = ctx_holder["ctx"] # type: ignore[assignment] # Safe: we just assigned it
960991
return self.client.get_response( # type: ignore[call-overload, no-any-return]
961992
messages=ctx["session_messages"],
962993
stream=True,
963994
options=ctx["chat_options"], # type: ignore[reportArgumentType]
995+
compaction_strategy=ctx["compaction_strategy"],
996+
tokenizer=ctx["tokenizer"],
964997
**ctx["filtered_kwargs"],
965998
)
966999

@@ -1047,6 +1080,8 @@ async def _prepare_run_context(
10471080
session: AgentSession | None,
10481081
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None,
10491082
options: Mapping[str, Any] | None,
1083+
compaction_strategy: CompactionStrategy | None,
1084+
tokenizer: TokenizerProtocol | None,
10501085
kwargs: dict[str, Any],
10511086
) -> _RunContext:
10521087
opts = dict(options) if options else {}
@@ -1081,9 +1116,10 @@ async def _prepare_run_context(
10811116
options=opts,
10821117
)
10831118

1119+
agent_name = self._get_agent_name()
1120+
10841121
# Normalize tools
10851122
normalized_tools = normalize_tools(tools_)
1086-
agent_name = self._get_agent_name()
10871123

10881124
# Resolve final tool list (runtime provided tools + local MCP server tools)
10891125
final_tools: list[FunctionTool | Callable[..., Any] | dict[str, Any] | Any] = []
@@ -1153,6 +1189,8 @@ async def _prepare_run_context(
11531189
"session_messages": session_messages,
11541190
"agent_name": agent_name,
11551191
"chat_options": co,
1192+
"compaction_strategy": compaction_strategy or self.compaction_strategy,
1193+
"tokenizer": tokenizer or self.tokenizer,
11561194
"filtered_kwargs": filtered_kwargs,
11571195
"finalize_kwargs": finalize_kwargs,
11581196
}
@@ -1408,6 +1446,8 @@ def __init__(
14081446
default_options: OptionsCoT | None = None,
14091447
context_providers: Sequence[BaseContextProvider] | None = None,
14101448
middleware: Sequence[MiddlewareTypes] | None = None,
1449+
compaction_strategy: CompactionStrategy | None = None,
1450+
tokenizer: TokenizerProtocol | None = None,
14111451
**kwargs: Any,
14121452
) -> None:
14131453
"""Initialize a Agent instance."""
@@ -1421,5 +1461,7 @@ def __init__(
14211461
default_options=default_options,
14221462
context_providers=context_providers,
14231463
middleware=middleware,
1464+
compaction_strategy=compaction_strategy,
1465+
tokenizer=tokenizer,
14241466
**kwargs,
14251467
)

0 commit comments

Comments
 (0)