diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 655612e4f2c..0052716eaa1 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -171,6 +171,25 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID Shared serialization normalizes stored data; codecs and connector overrides remain trusted Python code - **`SupportsVectorUpsert`** / **`SupportsVectorSearch`** - Structural protocols for vector store capabilities +### Compaction (`_compaction.py`) + +- **`AgentTriggeredCompactionStrategy`** - Wraps one existing compaction strategy and invokes it only after its + explicitly registered trigger tool completes in the latest message group. Trigger consumption is stored on the + completed trigger call occurrence rather than on the strategy instance so shared agents remain safe across + concurrent runs and composed trigger tools do not suppress each other. The trigger group + is temporarily excluded while the inner strategy runs, so it does not consume retention counts, then restored by + default; set `exclude_trigger=True` to remove a trigger-only group from the next model call. Mixed groups that + contain unrelated tool calls are always restored atomically. The agent-triggered compaction APIs are experimental + under the shared `AGENT_TRIGGERED_COMPACTION` feature ID and require client-managed history (`store=False` with no + service-managed conversation ID). Strategy wrappers and compositions must propagate the optional + `requires_client_managed_history` capability; `TokenBudgetComposedStrategy` does this for its nested strategies. +- **`UnloadSkillsCompactionStrategy`** - Intended as an inner agent-triggered strategy. It supplies an + `unload_skill` trigger-tool factory that `AgentTriggeredCompactionStrategy` adopts automatically, then removes + complete `load_skill` and `read_skill_resource` groups for one named skill, or all skills when no name is + supplied. Loaded skill identity comes from successful result provenance recorded by `SkillsProvider` after + resolving the effective invocation, not the original model arguments. Groups with missing/invalid provenance, + unresolved calls, mixed tools, or `run_skill_script` results remain intact. + ### Middleware (`_middleware.py`) - **`AgentMiddleware`** - Intercepts agent `run()` calls diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 8132f659c9e..e695e8d1e69 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -73,6 +73,7 @@ "SUMMARIZED_BY_SUMMARY_ID_KEY", "SUMMARY_OF_GROUP_IDS_KEY", "SUMMARY_OF_MESSAGE_IDS_KEY", + "AgentTriggeredCompactionStrategy", "CharacterEstimatorTokenizer", "CompactionProvider", "CompactionStrategy", @@ -84,6 +85,7 @@ "TokenizerProtocol", "ToolResultCompactionStrategy", "TruncationStrategy", + "UnloadSkillsCompactionStrategy", "annotate_message_groups", "apply_compaction", "included_messages", @@ -456,6 +458,7 @@ "AgentResponseUpdate", "AgentRunInputs", "AgentSession", + "AgentTriggeredCompactionStrategy", "AggregatingSkillsSource", "Annotation", "BackgroundAgentsProvider", @@ -649,6 +652,7 @@ "ToolTypes", "TruncationStrategy", "TypeCompatibilityError", + "UnloadSkillsCompactionStrategy", "UpdateT", "UsageDetails", "UserInputRequiredException", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 9a7e0fdfade..a46b7dcdc94 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -35,6 +35,7 @@ from ._compaction import ( SUMMARIZED_BY_SUMMARY_ID_KEY, SUMMARY_OF_GROUP_IDS_KEY, SUMMARY_OF_MESSAGE_IDS_KEY, + AgentTriggeredCompactionStrategy, CharacterEstimatorTokenizer, CompactionProvider, CompactionStrategy, @@ -46,6 +47,7 @@ from ._compaction import ( TokenizerProtocol, ToolResultCompactionStrategy, TruncationStrategy, + UnloadSkillsCompactionStrategy, annotate_message_groups, apply_compaction, included_messages, @@ -411,6 +413,7 @@ __all__ = [ "AgentResponseUpdate", "AgentRunInputs", "AgentSession", + "AgentTriggeredCompactionStrategy", "AggregatingSkillsSource", "Annotation", "BackgroundAgentsProvider", @@ -604,6 +607,7 @@ __all__ = [ "ToolTypes", "TruncationStrategy", "TypeCompatibilityError", + "UnloadSkillsCompactionStrategy", "UpdateT", "UsageDetails", "UserInputRequiredException", diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 447cb730b0b..d353b0a5521 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -40,6 +40,7 @@ ResponseStream, validate_chat_options, ) +from .exceptions import ChatClientInvalidRequestException if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover @@ -295,6 +296,8 @@ async def _stream(): When ``True``, the agent skips auto-injecting ``InMemoryHistoryProvider`` unless the user explicitly sets ``store=False``. """ + SERVICE_CONTINUATION_OPTION_NAMES: ClassVar[tuple[str, ...]] = ("conversation_id",) + """Option names that activate provider-managed conversation continuation.""" # OTEL_PROVIDER_NAME is used for OTel setup, should be overridden in subclasses def __init__( @@ -406,8 +409,37 @@ async def _prepare_messages_for_model_call( *, compaction_strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None, + options: Mapping[str, Any] | None = None, ) -> list[Message]: prepared_messages = list(messages) + if compaction_strategy is not None: + from ._compaction import _requires_client_managed_history # pyright: ignore[reportPrivateUsage] + from ._sessions import is_local_history_conversation_id + + if _requires_client_managed_history(compaction_strategy): + store = options.get("store") if options is not None else None + service_stores_history = ( + getattr(self, "STORES_BY_DEFAULT", False) if store is None else store is not False + ) + has_service_conversation = False + if options is not None: + for option_name in self.SERVICE_CONTINUATION_OPTION_NAMES: + value = options.get(option_name) + if value is None or value == "": + continue + if ( + option_name == "conversation_id" + and isinstance(value, str) + and is_local_history_conversation_id(value) + ): + continue + has_service_conversation = True + break + if service_stores_history or has_service_conversation: + raise ChatClientInvalidRequestException( + "Agent-triggered compaction requires client-managed history. " + "Set store=False and do not use a service-managed continuation option." + ) if compaction_strategy is None: if tokenizer is None: return prepared_messages @@ -564,6 +596,7 @@ async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]] prepared_messages = await self._prepare_messages_for_model_call( messages, **compaction_overrides, + options=options, ) stream_response = self._inner_get_response( messages=prepared_messages, @@ -584,6 +617,7 @@ async def _get_response() -> ChatResponse[Any]: prepared_messages = await self._prepare_messages_for_model_call( messages, **compaction_overrides, + options=options, ) return await self._inner_get_response( messages=prepared_messages, diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 6f4078eca35..28ec1a7f017 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -4,6 +4,7 @@ import json import logging +from collections import deque from collections.abc import Iterable, Mapping, Sequence from typing import ( TYPE_CHECKING, @@ -16,12 +17,14 @@ runtime_checkable, ) +from ._feature_stage import ExperimentalFeature, experimental from ._sessions import ContextProvider from ._telemetry import FeatureIndex, mark_feature_used -from ._types import ChatResponse, Content, Message +from ._types import APPROVAL_REJECTION_RESULT_KEY, SKILL_CONTEXT_NAME_KEY, ChatResponse, Content, Message if TYPE_CHECKING: from ._clients import SupportsChatGetResponse + from ._tools import ApprovalMode, FunctionTool GroupKind: TypeAlias = Literal["system", "user", "assistant_text", "tool_call"] GROUP_ANNOTATION_KEY = "_group" @@ -41,6 +44,9 @@ SUMMARY_OF_MESSAGE_IDS_KEY = "_summary_of_message_ids" SUMMARY_OF_GROUP_IDS_KEY = "_summary_of_group_ids" SUMMARIZED_BY_SUMMARY_ID_KEY = "_summarized_by_summary_id" +_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY = "_agent_compaction_trigger_consumed" +_SKILL_UNLOAD_SELECTION_KEY = "_skill_unload_selection" +_SKILL_CONTEXT_TOOL_NAMES: Final[frozenset[str]] = frozenset({"load_skill", "read_skill_resource"}) logger = logging.getLogger("agent_framework") @@ -52,6 +58,7 @@ "code_interpreter_tool_call", "shell_tool_call", "image_generation_tool_call", + "search_tool_call", } @@ -104,7 +111,12 @@ def count_tokens(self, text: str) -> int: @runtime_checkable class CompactionStrategy(Protocol): - """Protocol for in-place message compaction strategies.""" + """Protocol for in-place message compaction strategies. + + Strategies that require local/client-managed message history should expose + ``requires_client_managed_history = True``. Wrappers and compositions must + propagate that optional capability from nested strategies. + """ async def __call__(self, messages: list[Message]) -> bool: """Mutate message annotations and/or list contents in place. @@ -1180,6 +1192,532 @@ def _count_included_tokens(messages: list[Message]) -> int: return included_token_count(messages) +def _latest_group(messages: list[Message], *, include_excluded: bool) -> list[Message]: + grouped = _group_messages_by_id(messages) + for message in reversed(messages): + if not include_excluded and message.additional_properties.get(EXCLUDED_KEY, False): + continue + group_id = _group_id(message) + if group_id is not None: + return grouped.get(group_id, []) + return [] + + +def _latest_included_group(messages: list[Message]) -> list[Message]: + return _latest_group(messages, include_excluded=False) + + +def _function_call_result_pairs( + messages: Sequence[Message], + *, + include_excluded: bool = False, +) -> list[tuple[Content, Content]]: + candidate_messages = ( + list(messages) + if include_excluded + else [message for message in messages if not message.additional_properties.get(EXCLUDED_KEY, False)] + ) + open_calls_by_call_id: dict[str, deque[Content]] = {} + open_calls_by_occurrence_id: dict[tuple[str, str], deque[Content]] = {} + closed_call_identities: set[int] = set() + pairs: list[tuple[Content, Content]] = [] + + for message in candidate_messages: + for content in message.contents: + if content.type == "function_call" and content.call_id: + open_calls_by_call_id.setdefault(content.call_id, deque()).append(content) + if content.id: + open_calls_by_occurrence_id.setdefault((content.id, content.call_id), deque()).append(content) + continue + if content.type != "function_result" or not content.call_id: + continue + + matched_call: Content | None = None + if content.id is not None: + occurrence_queue = open_calls_by_occurrence_id.get((content.id, content.call_id)) + while occurrence_queue and id(occurrence_queue[0]) in closed_call_identities: + occurrence_queue.popleft() + if occurrence_queue: + matched_call = occurrence_queue.popleft() + else: + call_queue = open_calls_by_call_id.get(content.call_id) + while call_queue and id(call_queue[0]) in closed_call_identities: + call_queue.popleft() + if call_queue: + matched_call = call_queue.popleft() + if matched_call is None: + continue + + closed_call_identities.add(id(matched_call)) + pairs.append((matched_call, content)) + + return pairs + + +def _completed_function_call_pairs( + messages: Sequence[Message], + *, + tool_name: str, + include_consumed: bool = False, + include_excluded: bool = False, +) -> list[tuple[Content, Content]]: + return [ + (call, result) + for call, result in _function_call_result_pairs(messages, include_excluded=include_excluded) + if ( + call.name == tool_name + and result.exception is None + and result.additional_properties.get(APPROVAL_REJECTION_RESULT_KEY) is not True + and (include_consumed or call.additional_properties.get(_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY) is not True) + ) + ] + + +def _completed_function_calls( + messages: Sequence[Message], + *, + tool_name: str, + include_consumed: bool = False, + include_excluded: bool = False, +) -> list[Content]: + return [ + call + for call, _ in _completed_function_call_pairs( + messages, + tool_name=tool_name, + include_consumed=include_consumed, + include_excluded=include_excluded, + ) + ] + + +def _requires_client_managed_history(strategy: object | None) -> bool: + return getattr(strategy, "requires_client_managed_history", False) is True + + +@experimental(feature_id=ExperimentalFeature.AGENT_TRIGGERED_COMPACTION) +class AgentTriggeredCompactionStrategy: + """Run an existing compaction strategy after an agent invokes a trigger tool. + + The trigger tool must be registered explicitly on the agent alongside this + strategy. A completed trigger call in the latest included message group is + consumed once, then the wrapped strategy runs before the next model call. + Wrapped strategies may provide ``create_trigger_tool()`` to replace the + default ``compact_context`` tool with one specialized for their inputs. + + The consumed marker is stored on the completed trigger call occurrence + rather than on this strategy instance, so one strategy can be shared safely + by concurrent runs and composed trigger tools do not suppress each other. + + This strategy requires client-managed history (for example ``store=False``). + Chat clients reject it when service-side storage or a service-managed + conversation ID is active because local compaction cannot remove remote + conversation context. + """ + + DEFAULT_TOOL_NAME: Final[str] = "compact_context" + requires_client_managed_history = True + + def __init__( + self, + strategy: CompactionStrategy, + *, + tool_name: str = DEFAULT_TOOL_NAME, + approval_mode: ApprovalMode = "never_require", + exclude_trigger: bool = False, + ) -> None: + """Create an agent-triggered wrapper around one compaction strategy. + + Args: + strategy: Existing strategy to invoke after a completed trigger call. + + Keyword Args: + tool_name: Name for the default trigger tool when the wrapped + strategy does not provide ``create_trigger_tool()``. + approval_mode: Approval mode for the generated default trigger tool. + exclude_trigger: Whether to exclude the completed trigger call/result + group before the next model call. Defaults to False so the model + can observe that compaction completed. A trigger group containing + unrelated tool calls is always retained atomically. + """ + if not callable(strategy): + raise TypeError("strategy must be callable.") + if not tool_name.strip(): + raise ValueError("tool_name must not be empty.") + + trigger_tool_factory = getattr(strategy, "create_trigger_tool", None) + tool = ( + trigger_tool_factory() + if callable(trigger_tool_factory) + else self._create_default_tool(tool_name, approval_mode) + ) + from ._tools import FunctionTool + + if not isinstance(tool, FunctionTool): + raise TypeError("create_trigger_tool() must return a FunctionTool.") + if not tool.name.strip(): + raise ValueError("trigger tool name must not be empty.") + + self.strategy = strategy + self.tool = tool + self.exclude_trigger = exclude_trigger + + @property + def tokenizer(self) -> TokenizerProtocol | None: + """Return the wrapped strategy's tokenizer, when it exposes one.""" + tokenizer = getattr(self.strategy, "tokenizer", None) + return tokenizer if isinstance(tokenizer, TokenizerProtocol) else None + + @staticmethod + def _create_default_tool(tool_name: str, approval_mode: ApprovalMode) -> FunctionTool: + from ._tools import FunctionTool + + def request_compaction(reason: str | None = None) -> str: + """Request configured context compaction before the next model call.""" + message = "Compaction requested. The configured strategy will run before the next model call." + normalized_reason = reason.strip() if isinstance(reason, str) else "" + return f"{message} Reason: {normalized_reason}" if normalized_reason else message + + return FunctionTool( + name=tool_name, + description=( + "Request that the configured context compaction strategy run before the next model call. " + "Call this tool by itself, after any other tool calls whose results should be eligible for compaction; " + "do not call it in parallel with other tools." + ), + func=request_compaction, + approval_mode=approval_mode, + input_model={ + "type": "object", + "properties": { + "reason": { + "type": "string", + "description": "Why the agent is requesting compaction before the next model call.", + }, + }, + "additionalProperties": False, + }, + ) + + async def __call__(self, messages: list[Message]) -> bool: + """Run the wrapped strategy once for the latest completed trigger call.""" + original_messages = list(messages) + annotate_message_groups(messages) + grouped = _group_messages_by_id(messages) + trigger_group = _latest_included_group(messages) + trigger_calls = _completed_function_calls( + trigger_group, + tool_name=self.tool.name, + ) + if not trigger_calls: + return False + + groups_to_hide: list[list[Message]] = [trigger_group] + trigger_group_ids = {id(message) for message in trigger_group} + for group in grouped.values(): + if any(id(message) in trigger_group_ids for message in group): + continue + if not any( + content.type == "function_call" + and content.name == self.tool.name + and content.additional_properties.get(_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY) is True + for message in group + if not message.additional_properties.get(EXCLUDED_KEY, False) + for content in message.contents + ): + continue + if self._is_trigger_only_group(group): + groups_to_hide.append(group) + + hidden_messages = [message for group in groups_to_hide for message in group] + trigger_state = [ + ( + message, + EXCLUDED_KEY in message.additional_properties, + message.additional_properties.get(EXCLUDED_KEY), + EXCLUDE_REASON_KEY in message.additional_properties, + message.additional_properties.get(EXCLUDE_REASON_KEY), + ) + for message in hidden_messages + ] + for message in hidden_messages: + set_excluded(message, excluded=True, reason="agent_compaction_trigger") + + try: + changed = await self.strategy(messages) + except BaseException: + self._restore_hidden_messages(messages, original_messages, trigger_state) + raise + + for trigger_call in trigger_calls: + trigger_call.additional_properties[_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY] = True + + has_unrelated_tool_calls = not self._is_trigger_only_group(trigger_group, include_excluded=True) + trigger_message_ids = {id(message) for message in trigger_group} + has_other_included_context = any( + id(message) not in trigger_message_ids and not message.additional_properties.get(EXCLUDED_KEY, False) + for message in messages + ) + trigger_changed = self.exclude_trigger and not has_unrelated_tool_calls and has_other_included_context + prior_trigger_state = [state for state in trigger_state if id(state[0]) not in trigger_message_ids] + self._restore_hidden_messages(messages, original_messages, prior_trigger_state) + if not trigger_changed: + current_trigger_state = [state for state in trigger_state if id(state[0]) in trigger_message_ids] + self._restore_hidden_messages(messages, original_messages, current_trigger_state) + return changed or trigger_changed + + def _is_trigger_only_group( + self, + group: Sequence[Message], + *, + include_excluded: bool = False, + ) -> bool: + tool_calls = [ + content + for message in group + if include_excluded or not message.additional_properties.get(EXCLUDED_KEY, False) + for content in message.contents + if content.type in _TOOL_CALL_CONTENT_TYPES + ] + return bool(tool_calls) and all( + content.type == "function_call" and content.name == self.tool.name for content in tool_calls + ) + + @staticmethod + def _restore_trigger_group( + trigger_state: Sequence[tuple[Message, bool, Any, bool, Any]], + ) -> None: + for message, had_excluded, excluded, had_reason, reason in trigger_state: + if had_excluded: + message.additional_properties[EXCLUDED_KEY] = excluded + else: + message.additional_properties.pop(EXCLUDED_KEY, None) + if had_reason: + message.additional_properties[EXCLUDE_REASON_KEY] = reason + else: + message.additional_properties.pop(EXCLUDE_REASON_KEY, None) + + @classmethod + def _restore_hidden_messages( + cls, + messages: list[Message], + original_messages: Sequence[Message], + trigger_state: Sequence[tuple[Message, bool, Any, bool, Any]], + ) -> None: + restore_ids = {id(state[0]) for state in trigger_state} + present_ids = {id(message) for message in messages} + original_index_by_id = {id(message): index for index, message in enumerate(original_messages)} + + for original_message in original_messages: + message_id = id(original_message) + if message_id not in restore_ids or message_id in present_ids: + continue + original_index = original_index_by_id[message_id] + next_anchor_id = next( + ( + id(candidate) + for candidate in original_messages[original_index + 1 :] + if id(candidate) in present_ids + ), + None, + ) + if next_anchor_id is not None: + insert_index = next( + index for index, candidate in enumerate(messages) if id(candidate) == next_anchor_id + ) + else: + previous_anchor_id = next( + ( + id(candidate) + for candidate in reversed(original_messages[:original_index]) + if id(candidate) in present_ids + ), + None, + ) + if previous_anchor_id is None: + insert_index = len(messages) + else: + insert_index = ( + next(index for index, candidate in enumerate(messages) if id(candidate) == previous_anchor_id) + + 1 + ) + messages.insert(insert_index, original_message) + present_ids.add(message_id) + + cls._restore_trigger_group(trigger_state) + + +@experimental(feature_id=ExperimentalFeature.AGENT_TRIGGERED_COMPACTION) +class UnloadSkillsCompactionStrategy: + """Exclude loaded skill bodies and resources selected by an unload tool. + + This strategy targets complete tool-call groups containing only + ``load_skill`` and ``read_skill_resource`` calls. The latest completed + ``unload_skill`` request selects one skill by name, or all skills when its + ``skill_name`` argument is omitted or blank. Skill identity comes from + successful load/resource result provenance, not the original model arguments. + Groups without that provenance are retained. + + Use this as the wrapped strategy of :class:`AgentTriggeredCompactionStrategy`. + The wrapper obtains the specialized ``unload_skill`` tool through + :meth:`create_trigger_tool` and exposes that single tool to the caller. + """ + + DEFAULT_TOOL_NAME: Final[str] = "unload_skill" + + def __init__( + self, + *, + tool_name: str = DEFAULT_TOOL_NAME, + approval_mode: ApprovalMode = "never_require", + ) -> None: + """Create a skill-unloading strategy and its trigger tool. + + Keyword Args: + tool_name: Name of the tool that requests skill unloading. + approval_mode: Approval mode for the generated unload tool. + """ + if not tool_name.strip(): + raise ValueError("tool_name must not be empty.") + self.tool_name: str = tool_name + self.approval_mode: ApprovalMode = approval_mode + + def create_trigger_tool(self) -> FunctionTool: + """Create the ``unload_skill`` trigger tool used by the wrapper.""" + from ._tools import FunctionTool + + def request_skill_unload(skill_name: str | Sequence[str] | None = None) -> list[Content]: + """Unload the full instructions and resources for selected skills.""" + if skill_name is None: + requested_names: Sequence[str] = [] + elif isinstance(skill_name, str): + requested_names = [skill_name] + elif isinstance(skill_name, Sequence) and not isinstance(skill_name, (bytes, bytearray)): + if any(not isinstance(name, str) for name in skill_name): + raise TypeError("skill_name must contain only strings.") + requested_names = skill_name + else: + raise TypeError("skill_name must be a string, a sequence of strings, or None.") + normalized_names = [name.strip() for name in requested_names if name.strip()] + if normalized_names: + formatted_names = ", ".join(f"'{name}'" for name in normalized_names) + message = f"Skill(s) {formatted_names} will be unloaded before the next model call." + else: + message = "All loaded skills will be unloaded before the next model call." + return [ + Content.from_text( + message, + additional_properties={_SKILL_UNLOAD_SELECTION_KEY: normalized_names}, + ) + ] + + return FunctionTool( + name=self.tool_name, + description=( + "Unloads the full instructions and resources for one or more loaded skills. " + "Omit skill_name to unload all loaded skills. Call this tool by itself after the relevant skill work " + "is complete; do not call it in parallel with other tools." + ), + func=request_skill_unload, + approval_mode=self.approval_mode, + input_model={ + "type": "object", + "properties": { + "skill_name": { + "oneOf": [ + {"type": "string"}, + {"type": "array", "items": {"type": "string"}}, + ], + "description": ( + "The name of a skill to unload, a list of skill names to unload, " + "or omit to unload all loaded skills." + ), + }, + }, + "additionalProperties": False, + }, + ) + + def _requested_skill_names(self, messages: list[Message]) -> set[str] | None: + requests = _completed_function_call_pairs( + _latest_group(messages, include_excluded=True), + tool_name=self.tool_name, + include_consumed=True, + include_excluded=True, + ) + if not requests: + return set() + + requested_names: set[str] = set() + for _, result in requests: + selections = [ + item.additional_properties[_SKILL_UNLOAD_SELECTION_KEY] + for item in (result.items or []) + if item.type == "text" and _SKILL_UNLOAD_SELECTION_KEY in item.additional_properties + ] + if len(selections) != 1 or not isinstance(selections[0], list): + return set() + skill_names = cast("list[Any]", selections[0]) + if any(not isinstance(skill_name, str) for skill_name in skill_names): + return set() + normalized_names = {skill_name.strip().casefold() for skill_name in skill_names if skill_name.strip()} + if not normalized_names: + return None + requested_names.update(normalized_names) + return requested_names + + @staticmethod + def _skill_names_for_removable_group(group: Sequence[Message]) -> set[str] | None: + included_messages = [message for message in group if not message.additional_properties.get(EXCLUDED_KEY, False)] + tool_calls = [ + content + for message in included_messages + for content in message.contents + if content.type in _TOOL_CALL_CONTENT_TYPES + ] + if not tool_calls or any( + call.type != "function_call" or call.name not in _SKILL_CONTEXT_TOOL_NAMES for call in tool_calls + ): + return None + calls = tool_calls + + results_by_call_identity = {id(call): result for call, result in _function_call_result_pairs(included_messages)} + skill_names: set[str] = set() + for call in calls: + result = results_by_call_identity.get(id(call)) + if ( + result is None + or result.exception is not None + or result.additional_properties.get(APPROVAL_REJECTION_RESULT_KEY) is True + ): + return None + skill_name = result.additional_properties.get(SKILL_CONTEXT_NAME_KEY) + if not isinstance(skill_name, str) or not skill_name.strip(): + return None + skill_names.add(skill_name.strip().casefold()) + return skill_names + + async def __call__(self, messages: list[Message]) -> bool: + """Exclude matching loaded skill context groups.""" + ordered_group_ids = annotate_message_groups(messages) + requested_skill_names = self._requested_skill_names(messages) + if requested_skill_names == set(): + return False + + grouped = _group_messages_by_id(messages) + changed = False + for group_id in ordered_group_ids: + group = grouped.get(group_id, []) + skill_names = self._skill_names_for_removable_group(group) + if skill_names is None: + continue + if requested_skill_names is not None and not skill_names.issubset(requested_skill_names): + continue + for message in group: + changed = set_excluded(message, excluded=True, reason="skill_unload") or changed + return changed + + class TruncationStrategy: """Oldest-first compaction using a single metric threshold. @@ -1429,7 +1967,11 @@ async def __call__(self, messages: list[Message]) -> bool: for content in msg.contents: if content.type == "function_call" and content.call_id and content.name: call_id_to_name[content.call_id] = content.name - elif content.type == "mcp_server_tool_call" and content.call_id and content.tool_name: + elif ( + content.type in ("mcp_server_tool_call", "search_tool_call") + and content.call_id + and content.tool_name + ): call_id_to_name[content.call_id] = content.tool_name # Collect tool results with the tool name for context. @@ -1448,6 +1990,11 @@ async def __call__(self, messages: list[Message]) -> bool: tool_name = call_id_to_name.get(content.call_id or "", "") label = f"{tool_name}: {result_text}" if tool_name else result_text tool_results.append(label.strip()) + elif content.type == "search_tool_result": + result_text = _tool_result_text(content.result if content.result is not None else content.items) + tool_name = call_id_to_name.get(content.call_id or "", "") + label = f"{tool_name}: {result_text}" if tool_name else result_text + tool_results.append(label.strip()) summary_text = self._summary_text(tool_results) summary_id = f"tool_summary_{group_id}" @@ -1963,6 +2510,11 @@ def __init__( self.strategies = list(strategies) self.early_stop = early_stop + @property + def requires_client_managed_history(self) -> bool: + """Return whether any nested strategy requires client-managed history.""" + return any(_requires_client_managed_history(strategy) for strategy in self.strategies) + async def __call__(self, messages: list[Message]) -> bool: annotate_message_groups(messages) annotate_token_counts(messages, tokenizer=self.tokenizer) @@ -2369,6 +2921,7 @@ async def __call__(self, messages: list[Message]) -> bool: "SUMMARIZED_BY_SUMMARY_ID_KEY", "SUMMARY_OF_GROUP_IDS_KEY", "SUMMARY_OF_MESSAGE_IDS_KEY", + "AgentTriggeredCompactionStrategy", "CharacterEstimatorTokenizer", "CompactionProvider", "CompactionStrategy", @@ -2381,6 +2934,7 @@ async def __call__(self, messages: list[Message]) -> bool: "TokenizerProtocol", "ToolResultCompactionStrategy", "TruncationStrategy", + "UnloadSkillsCompactionStrategy", "annotate_message_groups", "annotate_token_counts", "append_compaction_message", diff --git a/python/packages/core/agent_framework/_feature_stage.py b/python/packages/core/agent_framework/_feature_stage.py index 9421c12cc27..b4988db3c19 100644 --- a/python/packages/core/agent_framework/_feature_stage.py +++ b/python/packages/core/agent_framework/_feature_stage.py @@ -50,6 +50,7 @@ class ExperimentalFeature(str, Enum): on enum membership or attribute presence over time. """ + AGENT_TRIGGERED_COMPACTION = "AGENT_TRIGGERED_COMPACTION" AGENT_HOOKS = "AGENT_HOOKS" COMPUTER_USE = "COMPUTER_USE" DECLARATIVE_AGENTS = "DECLARATIVE_AGENTS" diff --git a/python/packages/core/agent_framework/_skills.py b/python/packages/core/agent_framework/_skills.py index 8e323145f0b..464d79ac7e2 100644 --- a/python/packages/core/agent_framework/_skills.py +++ b/python/packages/core/agent_framework/_skills.py @@ -92,7 +92,13 @@ from ._middleware import FunctionInvocationContext from ._sessions import ContextProvider from ._telemetry import FeatureIndex, mark_feature_used -from ._tools import ApprovalMode, FunctionTool +from ._tools import ( + _FUNCTION_RESULT_CARRIER_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] + ApprovalMode, + FunctionTool, + _FunctionResultCarrier, # pyright: ignore[reportPrivateUsage] +) +from ._types import SKILL_CONTEXT_NAME_KEY if TYPE_CHECKING: from mcp import Client, ClientSession @@ -2717,8 +2723,8 @@ def _create_tools( A list of :class:`FunctionTool` instances. """ - async def _load(skill_name: str) -> str: - return await self._load_skill(skills, skill_name) + async def _load(ctx: FunctionInvocationContext, skill_name: str) -> str: + return await self._load_skill(skills, skill_name, context=ctx) # The runtime kwargs these handlers forward come from the host's # ``agent.run(function_invocation_kwargs=...)``, never from the model. Two @@ -2752,6 +2758,7 @@ async def _read_resource( skill_name, resource_name, runtime_kwargs=ctx.kwargs, + context=ctx, ) async def _run_script( @@ -2863,7 +2870,26 @@ def _find_skill(skills: Sequence[Skill], name: str) -> Skill | None: name_lower = name.lower() return next((s for s in skills if s.frontmatter.name.lower() == name_lower), None) - async def _load_skill(self, skills: Sequence[Skill], skill_name: str) -> str: + @staticmethod + def _record_skill_context(context: FunctionInvocationContext | None, skill: Skill) -> None: + """Attach the resolved skill identity without changing model-facing output.""" + if context is not None: + carrier = context.metadata.get(_FUNCTION_RESULT_CARRIER_CONTEXT_KEY) + if isinstance(carrier, _FunctionResultCarrier): + carrier.additional_properties[SKILL_CONTEXT_NAME_KEY] = skill.frontmatter.name + else: + context.metadata[_FUNCTION_RESULT_CARRIER_CONTEXT_KEY] = _FunctionResultCarrier( + additional_properties={SKILL_CONTEXT_NAME_KEY: skill.frontmatter.name}, + item_additional_properties={}, + ) + + async def _load_skill( + self, + skills: Sequence[Skill], + skill_name: str, + *, + context: FunctionInvocationContext | None = None, + ) -> str: """Return the full content for the named skill. Delegates to the skill's :meth:`~Skill.get_content` method, which @@ -2873,6 +2899,9 @@ async def _load_skill(self, skills: Sequence[Skill], skill_name: str) -> str: skills: The skills to look up the skill from. skill_name: The name of the skill to load. + Keyword Args: + context: Invocation context for recording successful load provenance. + Returns: The skill content text, or a user-facing error message if *skill_name* is empty or not found. @@ -2886,7 +2915,9 @@ async def _load_skill(self, skills: Sequence[Skill], skill_name: str) -> str: logger.info("Loading skill: %s", skill_name) - return await skill.get_content() + content = await skill.get_content() + self._record_skill_context(context, skill) + return content async def _run_skill_script( self, @@ -2954,6 +2985,7 @@ async def _read_skill_resource( resource_name: str, *, runtime_kwargs: Mapping[str, Any] | None = None, + context: FunctionInvocationContext | None = None, ) -> Any: """Read a named resource from a skill. @@ -2971,6 +3003,7 @@ async def _read_skill_resource( ``agent.run(..., function_invocation_kwargs={"user_id": "123"})`` and delivered through :class:`FunctionInvocationContext`; they are never sourced from model-supplied tool arguments. + context: Invocation context for recording successful resource provenance. Returns: The resource content (any type). Returns a user-facing error @@ -2999,11 +3032,14 @@ async def _read_skill_resource( return f"Error: Resource '{resource_name}' not found in skill '{skill_name}'." try: - return await resource.read(**(runtime_kwargs or {})) + result = await resource.read(**(runtime_kwargs or {})) except Exception: logger.exception("Failed to read resource '%s' from skill '%s'", resource_name, skill_name) raise + self._record_skill_context(context, skill) + return result + # endregion diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 08f00fefc11..53b64bce9d3 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -416,9 +416,7 @@ def _parse_inputs( # pyright: ignore[reportUnusedFunction] if inputs is None: return [] - from ._types import ( - Content, - ) + from ._types import Content parsed_inputs: list[Content] = [] if not isinstance(inputs, list): @@ -4449,9 +4447,7 @@ def _replace_approval_contents_with_results( Returns: The terminal contents produced while resolving the approval responses, in response order. """ - from ._types import ( - Content, - ) + from ._types import APPROVAL_REJECTION_RESULT_KEY, Content result_groups_by_call_id: dict[str, deque[list[Content]]] = {} for result_group in approved_function_result_groups: @@ -4578,7 +4574,10 @@ def find_approval_occurrence(approval_id: str) -> _ApprovalCallOccurrence | None call_id=call_id, id=occurrence.function_call.id if occurrence is not None else content.function_call.id, result="Error: Tool call invocation was rejected by user.", - additional_properties=content.function_call.additional_properties, + additional_properties={ + **content.function_call.additional_properties, + APPROVAL_REJECTION_RESULT_KEY: True, + }, ) ] if not replacements: diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 7eca7a0fe28..b7febaf4256 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -459,6 +459,8 @@ def _restore_compaction_annotation_in_additional_properties( "code_interpreter_tool_call", "code_interpreter_tool_result", }) +APPROVAL_REJECTION_RESULT_KEY = "_approval_rejection_result" +SKILL_CONTEXT_NAME_KEY = "_skill_context_name" _CONTENT_TYPES: Final[frozenset[str]] = frozenset(get_args(ContentType)) diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index b28da937cce..8d6ec0033a4 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -23,6 +23,7 @@ AgentResponse, AgentResponseUpdate, AgentSession, + AgentTriggeredCompactionStrategy, ChatContext, ChatOptions, ChatResponse, @@ -30,9 +31,12 @@ Content, ContextProvider, Embedding, + FunctionMiddleware, FunctionTool, GeneratedEmbeddings, HistoryProvider, + InlineSkill, + InlineSkillResource, InMemoryHistoryProvider, InMemoryStore, Message, @@ -41,12 +45,15 @@ ResponseStream, ServiceSessionId, SessionContext, + SkillFrontmatter, + SkillsProvider, SlidingWindowStrategy, SupportsAgentRun, SupportsChatGetResponse, ToolApprovalMiddleware, ToolResultCompactionStrategy, TruncationStrategy, + UnloadSkillsCompactionStrategy, VectorStoreHistoryProvider, chat_middleware, enqueue_messages, @@ -56,6 +63,7 @@ from agent_framework._mcp import MCPTool, _build_prefixed_mcp_name, _normalize_mcp_name from agent_framework._middleware import FunctionInvocationContext from agent_framework._tools import _PARENT_SERVICE_SESSION_STATE_KEYS_CONTEXT_KEY +from agent_framework._types import SKILL_CONTEXT_NAME_KEY from agent_framework.exceptions import ( AgentInvalidRequestException, ChatClientInvalidResponseException, @@ -3493,6 +3501,275 @@ async def capturing_inner( assert captured_token_counts == [[9, 9]] +@pytest.mark.parametrize("stream", [False, True]) +async def test_agent_triggered_compaction_applies_before_next_model_call( + chat_client_base: MockBaseChatClient, + stream: bool, +) -> None: + agent_compaction = AgentTriggeredCompactionStrategy( + TruncationStrategy(max_n=2, compact_to=2, preserve_first_user_group=True), + ) + function_call = Content.from_function_call( + call_id="compact-1", + name=agent_compaction.tool.name, + arguments={"reason": "planning complete"}, + ) + if stream: + chat_client_base.streaming_responses = [ + [ChatResponseUpdate(contents=[function_call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[Content.from_text("done")], role="assistant", finish_reason="stop")], + ] + else: + chat_client_base.run_responses = [ + ChatResponse(messages=Message(role="assistant", contents=[function_call])), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + captured_roles: list[list[str]] = [] + captured_texts: list[list[str | None]] = [] + original_inner = chat_client_base._inner_get_response + + def capture_inner( + *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any + ) -> Any: + captured_roles.append([message.role for message in messages]) + captured_texts.append([message.text for message in messages]) + return original_inner(messages=messages, stream=stream, options=options, **kwargs) + + agent = Agent( + client=chat_client_base, + tools=[agent_compaction.tool], + compaction_strategy=agent_compaction, + ) + + with patch.object(chat_client_base, "_inner_get_response", side_effect=capture_inner): + run_messages = [ + Message(role="user", contents=["original task"]), + Message(role="assistant", contents=["obsolete planning detail"]), + Message(role="user", contents=["plan"]), + ] + if stream: + result_text = "".join([update.text or "" async for update in agent.run(run_messages, stream=True)]) + else: + result_text = (await agent.run(run_messages)).text + + assert result_text == "done" + assert captured_roles == [ + ["user", "assistant", "user"], + ["user", "user", "assistant", "tool"], + ] + assert "obsolete planning detail" in captured_texts[0] + assert "obsolete planning detail" not in captured_texts[1] + + +async def test_unload_skills_uses_middleware_rewritten_selection( + chat_client_base: MockBaseChatClient, +) -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + + class RewriteUnloadSelectionMiddleware(FunctionMiddleware): + async def process(self, context: FunctionInvocationContext, call_next: Any) -> None: + if context.function.name == strategy.tool.name: + context.arguments = {"skill_name": "excel"} + await call_next() + if context.function.name == strategy.tool.name: + assert isinstance(context.result, list) + context.result.append(Content.from_text("middleware note")) + + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="unload-1", + name=strategy.tool.name, + arguments={"skill_name": "pdf"}, + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + captured_calls: list[list[tuple[str | None, Any]]] = [] + original_inner = chat_client_base._inner_get_response + + def capture_inner( + *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any + ) -> Any: + captured_calls.append([ + (content.name, content.parse_arguments()) + for message in messages + for content in message.contents + if content.type == "function_call" + ]) + return original_inner(messages=messages, stream=stream, options=options, **kwargs) + + agent = Agent( + client=chat_client_base, + tools=[strategy.tool], + compaction_strategy=strategy, + middleware=[RewriteUnloadSelectionMiddleware()], + ) + run_messages = [ + Message(role="user", contents=["work"]), + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-pdf", + name="load_skill", + arguments={"skill_name": "pdf"}, + ) + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="load-pdf", + result="pdf instructions", + additional_properties={SKILL_CONTEXT_NAME_KEY: "pdf"}, + ) + ], + ), + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-excel", + name="load_skill", + arguments={"skill_name": "excel"}, + ) + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="load-excel", + result="excel instructions", + additional_properties={SKILL_CONTEXT_NAME_KEY: "excel"}, + ) + ], + ), + Message(role="user", contents=["unload the selected skill"]), + ] + + with patch.object(chat_client_base, "_inner_get_response", side_effect=capture_inner): + result = await agent.run(run_messages) + + assert result.text == "done" + assert ("load_skill", {"skill_name": "pdf"}) in captured_calls[1] + assert ("load_skill", {"skill_name": "excel"}) not in captured_calls[1] + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool_name", ["load_skill", "read_skill_resource"]) +@pytest.mark.parametrize("unload_name", ["pdf", "excel"]) +@pytest.mark.parametrize("rewrite_load", [False, True]) +async def test_unload_skills_uses_effective_loaded_skill_identity( + chat_client_base: MockBaseChatClient, + stream: bool, + tool_name: str, + unload_name: str, + rewrite_load: bool, +) -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + provider = SkillsProvider( + [ + InlineSkill( + frontmatter=SkillFrontmatter(name=skill_name, description=f"{skill_name} skill"), + instructions=f"{skill_name} instructions", + resources=[ + InlineSkillResource(name=f"{skill_name}-reference", content=f"{skill_name} resource content") + ], + ) + for skill_name in ("pdf", "excel") + ], + disable_load_skill_approval=True, + disable_read_skill_resource_approval=True, + ) + + class RewriteLoadedSkillMiddleware(FunctionMiddleware): + async def process(self, context: FunctionInvocationContext, call_next: Any) -> None: + if context.function.name == tool_name: + assert isinstance(context.arguments, dict) + context.arguments = {**context.arguments, "skill_name": "EXCEL"} + if tool_name == "read_skill_resource": + context.arguments["resource_name"] = "excel-reference" + await call_next() + + original_arguments = {"skill_name": "pdf"} + if tool_name == "read_skill_resource": + original_arguments["resource_name"] = "pdf-reference" + load_call = Content.from_function_call(call_id="load-1", name=tool_name, arguments=original_arguments) + unload_call = Content.from_function_call( + call_id="unload-1", + name=strategy.tool.name, + arguments={"skill_name": unload_name}, + ) + chat_client_base.function_invocation_configuration["max_iterations"] = 4 + if stream: + chat_client_base.streaming_responses = [ + [ChatResponseUpdate(role="assistant", contents=[load_call], finish_reason="tool_calls")], + [ChatResponseUpdate(role="assistant", contents=[unload_call], finish_reason="tool_calls")], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")], finish_reason="stop")], + ] + else: + chat_client_base.run_responses = [ + ChatResponse(messages=Message(role="assistant", contents=[load_call])), + ChatResponse(messages=Message(role="assistant", contents=[unload_call])), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + model_inputs: list[list[Message]] = [] + original_inner = chat_client_base._inner_get_response + + def capture_inner( + *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any + ) -> Any: + model_inputs.append([Message.from_dict(message.to_dict()) for message in messages]) + return original_inner(messages=messages, stream=stream, options=options, **kwargs) + + agent = Agent( + client=chat_client_base, + tools=[strategy.tool], + context_providers=[provider], + middleware=[RewriteLoadedSkillMiddleware()] if rewrite_load else [], + compaction_strategy=strategy, + default_options=ChatOptions(store=False), + ) + with patch.object(chat_client_base, "_inner_get_response", side_effect=capture_inner): + if stream: + response_stream = agent.run("Read a skill, then unload the selected skill.", stream=True) + async for _ in response_stream: + pass + result = await response_stream.get_final_response() + else: + result = await agent.run("Read a skill, then unload the selected skill.") + + assert result.messages[-1].text == "done" + assert len(model_inputs) == 3 + loaded_contents = [content for message in model_inputs[1] for content in message.contents] + captured_load = next(content for content in loaded_contents if content.type == "function_call") + captured_result = next(content for content in loaded_contents if content.type == "function_result") + assert captured_load.parse_arguments() == original_arguments + assert captured_result.exception is None + effective_name = "excel" if rewrite_load else "pdf" + assert captured_result.additional_properties[SKILL_CONTEXT_NAME_KEY] == effective_name + assert ( + f"{effective_name} instructions" in captured_result.result + or f"{effective_name} resource content" in captured_result.result + ) + final_contents = [content for message in model_inputs[2] for content in message.contents] + retained_load = any(content.call_id == "load-1" for content in final_contents) + assert retained_load is (unload_name != effective_name) + assert any(content.type == "function_call" and content.call_id == "unload-1" for content in final_contents) + assert any(content.type == "function_result" and content.call_id == "unload-1" for content in final_contents) + + async def test_chat_agent_uses_client_compaction_defaults_when_agent_unset(chat_client_base: Any) -> None: captured_roles: list[list[str]] = [] original_inner = chat_client_base._inner_get_response diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index 55a4c470450..1bdc4ce5c82 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -16,8 +16,10 @@ GROUP_TOKEN_COUNT_KEY, SUMMARIZED_BY_SUMMARY_ID_KEY, SUMMARY_OF_MESSAGE_IDS_KEY, + AgentTriggeredCompactionStrategy, BaseChatClient, ChatMiddleware, + ChatOptions, ChatResponse, ChatResponseUpdate, Content, @@ -25,6 +27,7 @@ SlidingWindowStrategy, SummarizationStrategy, SupportsChatGetResponse, + TokenBudgetComposedStrategy, ToolResultCompactionStrategy, TruncationStrategy, annotate_message_groups, @@ -32,6 +35,8 @@ included_messages, tool, ) +from agent_framework._sessions import LOCAL_HISTORY_CONVERSATION_ID +from agent_framework.exceptions import ChatClientInvalidRequestException class _NoOpChatMiddleware(ChatMiddleware): @@ -329,6 +334,124 @@ async def _capture( assert captured_roles == [["assistant"]] +@pytest.mark.parametrize("stream", [False, True]) +async def test_agent_triggered_compaction_rejects_default_service_storage( + chat_client_base: SupportsChatGetResponse, + stream: bool, +) -> None: + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.STORES_BY_DEFAULT = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + + with pytest.raises(ChatClientInvalidRequestException, match="store=False"): + if stream: + response_stream = chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + stream=True, + compaction_strategy=strategy, + ) + async for _ in response_stream: + pass + else: + response = chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + stream=False, + compaction_strategy=strategy, + ) + _ = await response + + +async def test_agent_triggered_compaction_rejects_explicit_store_true( + chat_client_base: SupportsChatGetResponse, +) -> None: + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + + with pytest.raises(ChatClientInvalidRequestException, match="store=False"): + await chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + options=ChatOptions(store=True), + compaction_strategy=strategy, + ) + + +async def test_nested_agent_triggered_compaction_rejects_service_storage( + chat_client_base: SupportsChatGetResponse, +) -> None: + class StrategyWrapper: + def __init__(self, strategy: Any) -> None: + self.inner = strategy + self.requires_client_managed_history = strategy.requires_client_managed_history + + async def __call__(self, messages: list[Message]) -> bool: + return await self.inner(messages) + + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.STORES_BY_DEFAULT = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = StrategyWrapper(AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1))) + + with pytest.raises(ChatClientInvalidRequestException, match="store=False"): + await chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + compaction_strategy=strategy, + ) + + +async def test_composed_agent_triggered_compaction_rejects_service_storage( + chat_client_base: SupportsChatGetResponse, +) -> None: + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.STORES_BY_DEFAULT = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + tokenizer = _FixedTokenizer(1) + strategy = TokenBudgetComposedStrategy( + token_budget=1, + tokenizer=tokenizer, + strategies=[AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1))], + ) + + assert strategy.requires_client_managed_history is True + with pytest.raises(ChatClientInvalidRequestException, match="store=False"): + await chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + compaction_strategy=strategy, + ) + + +async def test_agent_triggered_compaction_rejects_service_conversation_id( + chat_client_base: SupportsChatGetResponse, +) -> None: + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + + with pytest.raises(ChatClientInvalidRequestException, match="service-managed continuation"): + await chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + options=ChatOptions(store=False, conversation_id="service-conversation"), + compaction_strategy=strategy, + ) + + +@pytest.mark.parametrize("conversation_id", [None, LOCAL_HISTORY_CONVERSATION_ID]) +async def test_agent_triggered_compaction_allows_client_managed_history( + chat_client_base: SupportsChatGetResponse, + conversation_id: str | None, +) -> None: + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.STORES_BY_DEFAULT = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + options = ChatOptions(store=False) + if conversation_id is not None: + options["conversation_id"] = conversation_id + + response = await chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + options=options, + compaction_strategy=strategy, + ) + + assert response.text == "test response - Hello" + + async def test_base_client_per_call_tokenizer_override_annotates_messages( chat_client_base: SupportsChatGetResponse, ) -> None: diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 7c5ed25de54..8e4f6f54eab 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -19,6 +19,7 @@ SUMMARIZED_BY_SUMMARY_ID_KEY, SUMMARY_OF_GROUP_IDS_KEY, SUMMARY_OF_MESSAGE_IDS_KEY, + AgentTriggeredCompactionStrategy, CharacterEstimatorTokenizer, ChatResponse, CompactionProvider, @@ -31,18 +32,24 @@ TokenBudgetComposedStrategy, ToolResultCompactionStrategy, TruncationStrategy, + UnloadSkillsCompactionStrategy, annotate_message_groups, apply_compaction, included_messages, included_token_count, ) from agent_framework._compaction import ( + _SKILL_UNLOAD_SELECTION_KEY, + _completed_function_calls, _format_summary_message, + _function_call_result_pairs, _select_summary_input_groups, _serialize_message, append_compaction_message, extend_compaction_messages, ) +from agent_framework._feature_stage import _WARNED_FEATURES, ExperimentalFeature, ExperimentalWarning +from agent_framework._types import APPROVAL_REJECTION_RESULT_KEY, SKILL_CONTEXT_NAME_KEY def _assistant_function_call(call_id: str) -> Message: @@ -52,6 +59,17 @@ def _assistant_function_call(call_id: str) -> Message: ) +def _assistant_named_function_call( + call_id: str, + name: str, + arguments: dict[str, Any] | None = None, +) -> Message: + return Message( + role="assistant", + contents=[Content.from_function_call(call_id=call_id, name=name, arguments=arguments or {})], + ) + + def _assistant_function_call_with_result(call_id: str, result: str) -> Message: return Message( role="assistant", @@ -112,6 +130,36 @@ def _tool_result(call_id: str, result: str, message_id: str | None = None) -> Me ) +def _skill_unload_result(call_id: str, skill_names: list[str]) -> Message: + return Message( + role="tool", + contents=[ + Content.from_function_result( + call_id=call_id, + result=[ + Content.from_text( + "Skills will be unloaded before the next model call.", + additional_properties={_SKILL_UNLOAD_SELECTION_KEY: skill_names}, + ) + ], + ) + ], + ) + + +def _skill_context_result(call_id: str, skill_name: str, result: str) -> Message: + return Message( + role="tool", + contents=[ + Content.from_function_result( + call_id=call_id, + result=result, + additional_properties={SKILL_CONTEXT_NAME_KEY: skill_name}, + ) + ], + ) + + def _group_id(message: Message) -> str | None: annotation = message.additional_properties.get(GROUP_ANNOTATION_KEY) if not isinstance(annotation, dict): @@ -848,6 +896,720 @@ def test_truncation_strategy_validates_token_targets() -> None: raise AssertionError("Expected ValueError when compact_to is greater than max_n.") +class _RecordingCompactionStrategy: + def __init__(self, *, changed: bool = True) -> None: + self.call_count = 0 + self.changed = changed + + async def __call__(self, messages: list[Message]) -> bool: + self.call_count += 1 + return self.changed + + +def test_agent_triggered_compaction_apis_are_experimental() -> None: + feature_id = ExperimentalFeature.AGENT_TRIGGERED_COMPACTION.value + for api in (AgentTriggeredCompactionStrategy, UnloadSkillsCompactionStrategy): + assert getattr(api, "__feature_stage__", None) == "experimental" + assert getattr(api, "__feature_id__", None) == feature_id + assert ".. warning:: Experimental" in (api.__doc__ or "") + + _WARNED_FEATURES.discard((ExperimentalWarning, feature_id)) + with pytest.warns(ExperimentalWarning, match=feature_id): + AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + + +async def test_agent_triggered_compaction_runs_only_after_completed_trigger() -> None: + inner = _RecordingCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(inner) + messages = [ + Message(role="user", contents=["plan"]), + _assistant_named_function_call("compact-1", strategy.tool.name, {"reason": "planning complete"}), + ] + + assert await strategy(messages) is False + assert inner.call_count == 0 + + messages.append(_tool_result("compact-1", "Compaction requested.")) + + assert await strategy(messages) is True + assert inner.call_count == 1 + assert messages[1].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is not True + assert await strategy(messages) is False + assert inner.call_count == 1 + + +async def test_agent_triggered_compaction_consumes_request_when_inner_strategy_is_noop() -> None: + inner = _RecordingCompactionStrategy(changed=False) + strategy = AgentTriggeredCompactionStrategy(inner) + messages = [ + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is False + assert await strategy(messages) is False + assert inner.call_count == 1 + + +async def test_agent_triggered_compaction_can_exclude_trigger_group() -> None: + inner = _RecordingCompactionStrategy(changed=False) + strategy = AgentTriggeredCompactionStrategy(inner, exclude_trigger=True) + messages = [ + Message(role="user", contents=["keep working"]), + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is True + assert messages[0].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + + +async def test_agent_trigger_does_not_count_toward_tool_group_retention() -> None: + strategy = AgentTriggeredCompactionStrategy(SelectiveToolCallCompactionStrategy(keep_last_tool_call_groups=1)) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("work-1", "work_tool"), + _tool_result("work-1", "work result"), + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + +@pytest.mark.parametrize("trigger_position", [0, 1, 2]) +@pytest.mark.parametrize("exclude_trigger", [False, True]) +async def test_agent_trigger_in_mixed_tool_batch_is_position_independent_and_retained( + trigger_position: int, + exclude_trigger: bool, +) -> None: + strategy = AgentTriggeredCompactionStrategy( + SelectiveToolCallCompactionStrategy(keep_last_tool_call_groups=0), + exclude_trigger=exclude_trigger, + ) + batch_calls = [ + Content.from_function_call(call_id="work-a", name="work_a", arguments={}), + Content.from_function_call(call_id="work-b", name="work_b", arguments={}), + ] + batch_calls.insert( + trigger_position, + Content.from_function_call(call_id="compact-1", name=strategy.tool.name, arguments={}), + ) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("old-work", "old_work"), + _tool_result("old-work", "old result"), + Message(role="assistant", contents=batch_calls), + Message( + role="tool", + contents=[ + Content.from_function_result(call_id=call.call_id or "", result=f"{call.name} result") + for call in batch_calls + ], + ), + ] + + changed = await strategy(messages) + + assert changed is True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + assert messages[3].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[4].additional_properties.get(EXCLUDED_KEY) is not True + + +async def test_agent_trigger_in_mixed_search_batch_is_retained() -> None: + strategy = AgentTriggeredCompactionStrategy( + SelectiveToolCallCompactionStrategy(keep_last_tool_call_groups=0), + exclude_trigger=True, + ) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("old-work", "old_work"), + _tool_result("old-work", "old result"), + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="compact-1", name=strategy.tool.name, arguments={}), + Content.from_search_tool_call( + call_id="search-1", + tool_name="web_search", + arguments={"query": "latest"}, + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result(call_id="compact-1", result="Compaction requested."), + Content.from_search_tool_result( + call_id="search-1", + tool_name="web_search", + result="search result", + ), + ], + ), + ] + + changed = await strategy(messages) + + assert changed is True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + assert messages[3].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[4].additional_properties.get(EXCLUDED_KEY) is not True + + +async def test_list_mutating_strategy_restores_mixed_trigger_batch_only() -> None: + class ProjectIncludedMessagesStrategy: + async def __call__(self, messages: list[Message]) -> bool: + messages[:] = included_messages(messages) + return True + + strategy = AgentTriggeredCompactionStrategy( + ProjectIncludedMessagesStrategy(), + exclude_trigger=True, + ) + removed_by_inner = Message( + role="assistant", + contents=["already excluded"], + additional_properties={EXCLUDED_KEY: True}, + ) + mixed_call_message = Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="work-1", name="work_tool", arguments={}), + Content.from_function_call(call_id="compact-1", name=strategy.tool.name, arguments={}), + ], + ) + mixed_result_message = Message( + role="tool", + contents=[ + Content.from_function_result(call_id="work-1", result="work result"), + Content.from_function_result(call_id="compact-1", result="Compaction requested."), + ], + ) + user_message = Message(role="user", contents=["work"]) + messages = [ + removed_by_inner, + user_message, + mixed_call_message, + mixed_result_message, + ] + + changed = await strategy(messages) + + assert changed is True + assert messages == [user_message, mixed_call_message, mixed_result_message] + assert removed_by_inner not in messages + assert mixed_call_message.additional_properties.get(EXCLUDED_KEY) is not True + assert mixed_result_message.additional_properties.get(EXCLUDED_KEY) is not True + + +async def test_successive_triggers_do_not_consume_tool_group_retention() -> None: + strategy = AgentTriggeredCompactionStrategy(SelectiveToolCallCompactionStrategy(keep_last_tool_call_groups=2)) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("work-1", "work_tool"), + _tool_result("work-1", "work result 1"), + _assistant_named_function_call("work-2", "work_tool"), + _tool_result("work-2", "work result 2"), + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is False + + messages.extend([ + _assistant_named_function_call("work-3", "work_tool"), + _tool_result("work-3", "work result 3"), + _assistant_named_function_call("compact-2", strategy.tool.name), + _tool_result("compact-2", "Compaction requested."), + ]) + + assert await strategy(messages) is True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + assert messages[3].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[4].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[5].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[6].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[7].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[8].additional_properties.get(EXCLUDED_KEY) is not True + + +async def test_agent_triggered_compaction_restores_request_when_inner_strategy_fails() -> None: + class _FailingStrategy: + async def __call__(self, messages: list[Message]) -> bool: + raise RuntimeError("compaction failed") + + strategy = AgentTriggeredCompactionStrategy(_FailingStrategy()) + trigger_message = _assistant_named_function_call("compact-1", strategy.tool.name) + messages = [trigger_message, _tool_result("compact-1", "Compaction requested.")] + + with pytest.raises(RuntimeError, match="compaction failed"): + await strategy(messages) + + assert "_agent_compaction_trigger_consumed" not in trigger_message.contents[0].additional_properties + + +@pytest.mark.parametrize("reverse_order", [False, True]) +async def test_composed_trigger_tools_each_execute_once(reverse_order: bool) -> None: + first_inner = _RecordingCompactionStrategy(changed=False) + second_inner = _RecordingCompactionStrategy(changed=False) + first = AgentTriggeredCompactionStrategy(first_inner, tool_name="compact_first") + second = AgentTriggeredCompactionStrategy(second_inner, tool_name="compact_second") + messages = [ + Message(role="user", contents=["work"]), + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="first", name=first.tool.name, arguments={}), + Content.from_function_call(call_id="second", name=second.tool.name, arguments={}), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result(call_id="first", result="First compaction requested."), + Content.from_function_result(call_id="second", result="Second compaction requested."), + ], + ), + ] + strategies = [second, first] if reverse_order else [first, second] + + assert await strategies[0](messages) is False + assert await strategies[1](messages) is False + assert first_inner.call_count == 1 + assert second_inner.call_count == 1 + assert await first(messages) is False + assert await second(messages) is False + assert first_inner.call_count == 1 + assert second_inner.call_count == 1 + + +async def test_agent_triggered_compaction_ignores_trigger_outside_latest_group() -> None: + inner = _RecordingCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(inner) + messages = [ + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + Message(role="assistant", contents=["continued"]), + ] + + assert await strategy(messages) is False + assert inner.call_count == 0 + + +async def test_agent_triggered_compaction_uses_terminal_result_for_nonadjacent_group_recency() -> None: + inner = _RecordingCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(inner) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("compact-1", strategy.tool.name), + Message(role="assistant", contents=["intervening context"]), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is True + assert inner.call_count == 1 + + +async def test_agent_triggered_compaction_ignores_rejected_approval_result() -> None: + inner = _RecordingCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(inner) + rejection_result = Content.from_function_result( + call_id="compact-1", + result="Error: Tool call invocation was rejected by user.", + additional_properties={APPROVAL_REJECTION_RESULT_KEY: True}, + ) + messages = [ + _assistant_named_function_call("compact-1", strategy.tool.name), + Message(role="tool", contents=[rejection_result]), + ] + + assert await strategy(messages) is False + assert inner.call_count == 0 + + +def test_completed_function_calls_correlates_reused_call_id_by_occurrence() -> None: + rejected_call = Content.from_function_call( + id="occ-rejected", + call_id="shared", + name="compact_context", + arguments={"reason": "rejected"}, + ) + approved_call = Content.from_function_call( + id="occ-approved", + call_id="shared", + name="compact_context", + arguments={"reason": "approved"}, + ) + unresolved_call = Content.from_function_call( + id="occ-unresolved", + call_id="shared", + name="compact_context", + arguments={"reason": "unresolved"}, + ) + messages = [ + Message(role="assistant", contents=[rejected_call, approved_call, unresolved_call]), + Message( + role="tool", + contents=[ + Content.from_function_result( + id="occ-rejected", + call_id="shared", + result="Error: Tool call invocation was rejected by user.", + additional_properties={APPROVAL_REJECTION_RESULT_KEY: True}, + ), + Content.from_function_result( + id="occ-approved", + call_id="shared", + result="Compaction requested.", + ), + ], + ), + ] + + completed = _completed_function_calls(messages, tool_name="compact_context") + + assert completed == [approved_call] + + +def test_completed_function_calls_uses_fifo_for_idless_reused_call_id() -> None: + first_call = Content.from_function_call(call_id="shared", name="compact_context", arguments={}) + second_call = Content.from_function_call(call_id="shared", name="compact_context", arguments={}) + messages = [ + Message(role="assistant", contents=[first_call, second_call]), + Message( + role="tool", + contents=[Content.from_function_result(call_id="shared", result="Compaction requested.")], + ), + ] + + completed = _completed_function_calls(messages, tool_name="compact_context") + + assert completed == [first_call] + + +def test_function_call_result_pairs_scales_linearly_for_idless_reused_call_id() -> None: + calls = [ + Content.from_function_call(call_id="shared", name="compact_context", arguments={"index": index}) + for index in range(500) + ] + results = [Content.from_function_result(call_id="shared", result=f"result-{index}") for index in range(500)] + messages = [ + Message(role="assistant", contents=calls), + Message(role="tool", contents=results), + ] + + pairs = _function_call_result_pairs(messages) + + assert pairs == list(zip(calls, results)) + + +def test_agent_triggered_compaction_tool_describes_reason_parameter() -> None: + strategy = AgentTriggeredCompactionStrategy(_RecordingCompactionStrategy()) + + reason_schema = strategy.tool.parameters()["properties"]["reason"] + + assert reason_schema["description"] == "Why the agent is requesting compaction before the next model call." + assert "Call this tool by itself" in strategy.tool.description + assert "do not call it in parallel" in strategy.tool.description + + +async def test_unload_skills_compaction_removes_named_skill_only() -> None: + skill_strategy = UnloadSkillsCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(skill_strategy) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("load-pdf", "load_skill", {"skill_name": "pdf"}), + _skill_context_result("load-pdf", "pdf", "pdf instructions"), + _assistant_named_function_call("load-excel", "load_skill", {"skill_name": "excel"}), + _skill_context_result("load-excel", "excel", "excel instructions"), + _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), + _skill_unload_result("unload-pdf", ["pdf"]), + ] + + changed = await strategy(messages) + + assert changed is True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + assert messages[3].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[4].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[5].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[6].additional_properties.get(EXCLUDED_KEY) is not True + + +async def test_unload_skills_compaction_removes_multiple_named_skills() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + messages = [ + _assistant_named_function_call("load-pdf", "load_skill", {"skill_name": "pdf"}), + _skill_context_result("load-pdf", "pdf", "pdf instructions"), + _assistant_named_function_call("load-excel", "load_skill", {"skill_name": "excel"}), + _skill_context_result("load-excel", "excel", "excel instructions"), + _assistant_named_function_call("load-word", "load_skill", {"skill_name": "word"}), + _skill_context_result("load-word", "word", "word instructions"), + _assistant_named_function_call( + "unload-selected", + strategy.tool.name, + {"skill_name": ["pdf", "excel"]}, + ), + _skill_unload_result("unload-selected", ["pdf", "excel"]), + ] + + changed = await strategy(messages) + + assert changed is True + assert all(message.additional_properties.get(EXCLUDED_KEY) is True for message in messages[:4]) + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages[4:6]) + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages[6:]) + + +def test_unload_skills_tool_describes_skill_name_parameter() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + + skill_name_schema = strategy.tool.parameters()["properties"]["skill_name"] + + assert skill_name_schema["description"] == ( + "The name of a skill to unload, a list of skill names to unload, or omit to unload all loaded skills." + ) + assert "Call this tool by itself" in strategy.tool.description + assert "do not call it in parallel" in strategy.tool.description + + +@pytest.mark.parametrize( + "skill_name", + [123, [123], ["pdf", 123], b"pdf"], +) +async def test_unload_skills_tool_rejects_malformed_skill_names(skill_name: Any) -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + + with pytest.raises(TypeError, match="skill_name"): + await strategy.tool.invoke(arguments={"skill_name": skill_name}) + + +async def test_unload_skills_compaction_empty_name_removes_all_skill_context() -> None: + skill_strategy = UnloadSkillsCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(skill_strategy) + messages = [ + _assistant_named_function_call("load-pdf", "load_skill", {"skill_name": "pdf"}), + _skill_context_result("load-pdf", "pdf", "pdf instructions"), + _assistant_named_function_call( + "resource-excel", + "read_skill_resource", + {"skill_name": "excel", "resource_name": "reference"}, + ), + _skill_context_result("resource-excel", "excel", "excel reference"), + _assistant_named_function_call("unload-all", strategy.tool.name, {"skill_name": " "}), + _skill_unload_result("unload-all", []), + ] + + changed = await strategy(messages) + + assert changed is True + assert all(message.additional_properties.get(EXCLUDED_KEY) is True for message in messages[:4]) + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages[4:]) + + +async def test_unload_skills_compaction_preserves_script_results_and_mixed_groups() -> None: + skill_strategy = UnloadSkillsCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(skill_strategy) + mixed_call = Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-pdf", + name="load_skill", + arguments={"skill_name": "pdf"}, + ), + Content.from_function_call( + call_id="search", + name="search", + arguments={"query": "pdf"}, + ), + ], + ) + messages = [ + mixed_call, + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="load-pdf", + result="pdf instructions", + additional_properties={SKILL_CONTEXT_NAME_KEY: "pdf"}, + ), + Content.from_function_result(call_id="search", result="search result"), + ], + ), + _assistant_named_function_call( + "run-pdf", + "run_skill_script", + {"skill_name": "pdf", "script_name": "convert"}, + ), + _tool_result("run-pdf", "converted.pdf"), + _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), + _skill_unload_result("unload-pdf", ["pdf"]), + ] + + changed = await strategy(messages) + + assert changed is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + +async def test_unload_skills_compaction_preserves_group_with_hosted_tool_call() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-pdf", + name="load_skill", + arguments={"skill_name": "pdf"}, + ), + Content.from_shell_tool_call(call_id="shell-1", commands=["echo hi"]), + ], + ), + _skill_context_result("load-pdf", "pdf", "pdf instructions"), + _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), + _skill_unload_result("unload-pdf", ["pdf"]), + ] + + changed = await strategy(messages) + + assert changed is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + +async def test_unload_skills_compaction_preserves_group_with_search_tool_call() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-pdf", + name="load_skill", + arguments={"skill_name": "pdf"}, + ), + Content.from_search_tool_call( + call_id="search-1", + tool_name="web_search", + arguments={"query": "pdf"}, + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="load-pdf", + result="pdf instructions", + additional_properties={SKILL_CONTEXT_NAME_KEY: "pdf"}, + ), + Content.from_search_tool_result( + call_id="search-1", + tool_name="web_search", + result="search result", + ), + ], + ), + _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), + _skill_unload_result("unload-pdf", ["pdf"]), + ] + + changed = await strategy(messages) + + assert changed is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + +async def test_unload_skills_compaction_preserves_reused_id_group_with_unresolved_call() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call( + id="load-pdf", + call_id="shared", + name="load_skill", + arguments={"skill_name": "pdf"}, + ), + Content.from_function_call( + id="load-excel", + call_id="shared", + name="load_skill", + arguments={"skill_name": "excel"}, + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + id="load-pdf", + call_id="shared", + result="pdf instructions", + additional_properties={SKILL_CONTEXT_NAME_KEY: "pdf"}, + ) + ], + ), + _assistant_named_function_call("unload-all", strategy.tool.name), + _skill_unload_result("unload-all", []), + ] + + changed = await strategy(messages) + + assert changed is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + +@pytest.mark.parametrize( + ("properties", "exception"), + [ + ({}, None), + ({SKILL_CONTEXT_NAME_KEY: ""}, None), + ({SKILL_CONTEXT_NAME_KEY: 123}, None), + ({SKILL_CONTEXT_NAME_KEY: "pdf"}, "Load failed."), + ({SKILL_CONTEXT_NAME_KEY: "pdf", APPROVAL_REJECTION_RESULT_KEY: True}, None), + ], +) +async def test_unload_skills_preserves_groups_without_successful_provenance( + properties: dict[str, Any], + exception: str | None, +) -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("load-pdf", "load_skill", {"skill_name": "pdf"}), + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="load-pdf", + result="pdf instructions", + exception=exception, + additional_properties=properties, + ) + ], + ), + _assistant_named_function_call("unload-all", strategy.tool.name), + _skill_unload_result("unload-all", []), + ] + + assert await strategy(messages) is False + assert included_messages(messages) == messages + + async def test_selective_tool_call_strategy_excludes_older_tool_groups() -> None: messages = [ Message(role="user", contents=["u"]), @@ -1673,6 +2435,42 @@ async def test_apply_compaction_logs_changed_context_without_content(caplog: Any # --- ToolResultCompactionStrategy tests --- +async def test_tool_result_compaction_preserves_hosted_search_result() -> None: + messages = [ + Message(role="user", contents=["find docs"]), + Message( + role="assistant", + contents=[ + Content.from_search_tool_call( + call_id="search-1", + tool_name="web_search", + arguments={"query": "agent framework"}, + ) + ], + ), + Message( + role="tool", + contents=[ + Content.from_search_tool_result( + call_id="search-1", + tool_name="web_search", + result="Found the Agent Framework documentation.", + ) + ], + ), + Message(role="assistant", contents=["done"]), + ] + annotate_message_groups(messages) + + changed = await ToolResultCompactionStrategy(keep_last_tool_call_groups=0)(messages) + + assert changed is True + summaries = [message.text or "" for message in messages if (message.text or "").startswith("[Tool results:")] + assert len(summaries) == 1 + assert "web_search: Found the Agent Framework documentation." in summaries[0] + assert "no results" not in summaries[0] + + async def test_tool_result_compaction_collapses_old_groups_into_summary() -> None: """Old tool-call groups are collapsed into summary messages, newest kept.""" messages = [ diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index f185504d2e3..4b8a34dc829 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -7956,6 +7956,7 @@ def test_replace_approval_contents_with_results_correlates_reused_call_id_pendin def test_replace_approval_contents_with_results_resolves_rejected_request() -> None: """A rejected pending call produces one terminal result.""" from agent_framework._tools import _collect_approval_responses, _replace_approval_contents_with_results + from agent_framework._types import APPROVAL_REJECTION_RESULT_KEY function_call, request, _ = _build_approved_tool_roundtrip( call_id="call_rejected", @@ -7977,6 +7978,7 @@ def test_replace_approval_contents_with_results_resolves_rejected_request() -> N results = [content for message in messages for content in message.contents if content.type == "function_result"] assert len(results) == 1 assert results[0].result == "Error: Tool call invocation was rejected by user." + assert results[0].additional_properties[APPROVAL_REJECTION_RESULT_KEY] is True assert resolved_contents == results diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index fcf1e4f3cd1..604cc77654d 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -623,6 +623,11 @@ class RawOpenAIChatClient( INJECTABLE: ClassVar[set[str]] = {"client"} STORES_BY_DEFAULT: ClassVar[bool] = True + SERVICE_CONTINUATION_OPTION_NAMES: ClassVar[tuple[str, ...]] = ( + "conversation_id", + "previous_response_id", + "conversation", + ) SUPPORTS_RICH_FUNCTION_OUTPUT: ClassVar[bool] = True _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.OPENAI diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index 5cb15482393..1644d5c74ee 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -21,6 +21,7 @@ import pytest from agent_framework import ( Agent, + AgentTriggeredCompactionStrategy, Annotation, ChatOptions, ChatResponse, @@ -36,6 +37,7 @@ SupportsMCPTool, SupportsWebSearchTool, TextSpanRegion, + TruncationStrategy, WorkflowBuilder, tool, ) @@ -104,6 +106,43 @@ ) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + ("option_name", "option_value"), + [ + ("previous_response_id", "resp_existing"), + ("conversation", "conv_existing"), + ], +) +async def test_agent_triggered_compaction_rejects_native_service_continuation_aliases( + stream: bool, + option_name: str, + option_value: str, +) -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + options = cast(OpenAIChatOptions, {"store": False, option_name: option_value}) + + with pytest.raises(ChatClientInvalidRequestException, match="service-managed continuation"): + if stream: + response_stream = client.get_response( + [Message(role="user", contents=["Hello"])], + stream=True, + options=options, + compaction_strategy=strategy, + ) + async for _ in response_stream: + pass + else: + response = client.get_response( + [Message(role="user", contents=["Hello"])], + stream=False, + options=options, + compaction_strategy=strategy, + ) + _ = await response + + class OutputStruct(BaseModel): """A structured output for testing purposes.""" diff --git a/python/samples/02-agents/compaction/README.md b/python/samples/02-agents/compaction/README.md index 779893eefeb..420e3791e2c 100644 --- a/python/samples/02-agents/compaction/README.md +++ b/python/samples/02-agents/compaction/README.md @@ -11,6 +11,8 @@ This folder demonstrates context compaction patterns introduced by ADR-0019. - `custom.py` — defines a custom strategy implementing the `CompactionStrategy` protocol. - `tiktoken_tokenizer.py` — shows a `TokenizerProtocol` implementation backed by `tiktoken`. - `compaction_provider.py` — uses `CompactionProvider` with an agent and `InMemoryHistoryProvider`. +- `agent_triggered_compaction.py` — lets an agent trigger in-run compaction to unload one skill's context. +- `agent_triggered_compaction_basic.py` — demonstrates the generic wrapper with an existing tool-call strategy. Run samples with: @@ -22,6 +24,8 @@ uv run samples/02-agents/compaction/agent_client_overrides.py uv run samples/02-agents/compaction/custom.py uv run samples/02-agents/compaction/tiktoken_tokenizer.py uv run samples/02-agents/compaction/compaction_provider.py # requires OPENAI_API_KEY +uv run samples/02-agents/compaction/agent_triggered_compaction.py # requires Foundry configuration +uv run samples/02-agents/compaction/agent_triggered_compaction_basic.py # requires Foundry configuration ``` ## Security Considerations diff --git a/python/samples/02-agents/compaction/agent_triggered_compaction.py b/python/samples/02-agents/compaction/agent_triggered_compaction.py new file mode 100644 index 00000000000..6c3c79d0836 --- /dev/null +++ b/python/samples/02-agents/compaction/agent_triggered_compaction.py @@ -0,0 +1,168 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os +from itertools import count +from textwrap import dedent + +from agent_framework import ( + Agent, + AgentTriggeredCompactionStrategy, + CompactionStrategy, + InlineSkill, + Message, + SkillFrontmatter, + SkillsProvider, + UnloadSkillsCompactionStrategy, + included_messages, +) +from agent_framework.foundry import FoundryChatClient +from azure.identity import AzureCliCredential +from dotenv import load_dotenv + +load_dotenv() + +"""This sample demonstrates agent-triggered in-run compaction. + +The agent follows three steps in one function-calling run: + +1. Call ``load_skill`` to load a planning skill. +2. Use the skill, then call ``unload_skill``. +3. Continue without the loaded skill body in model context. + +``AgentTriggeredCompactionStrategy`` controls when compaction runs. +``UnloadSkillsCompactionStrategy`` controls what context is removed. +The trigger tool is registered explicitly alongside the strategy. + +A small inspection wrapper prints every model input after compaction so the +skill body is visibly present after ``load_skill`` and absent after +``unload_skill``. + +Environment variables: + FOUNDRY_PROJECT_ENDPOINT — Microsoft Foundry project endpoint. + FOUNDRY_MODEL — Model deployment name. + +Run with: + uv run samples/02-agents/compaction/agent_triggered_compaction.py +""" + + +MODEL_CALLS = count(1) + + +class InspectingCompactionStrategy: + """Print the projected messages after applying another strategy.""" + + def __init__(self, strategy: CompactionStrategy) -> None: + self.strategy = strategy + self.requires_client_managed_history = getattr(strategy, "requires_client_managed_history", False) + + async def __call__(self, messages: list[Message]) -> bool: + changed = await self.strategy(messages) + visible_messages = included_messages(messages) + + print(f"\n--- Model call {next(MODEL_CALLS)} ---") + print(f"Visible messages: {len(visible_messages)}") + for message in visible_messages: + parts: list[str] = [] + for content in message.contents: + if content.type == "text": + text = (content.text or "").replace("\n", " ") + parts.append(text[:100]) + elif content.type == "function_call": + parts.append(f"{content.name}({content.arguments})") + elif content.type == "function_result": + result = str(content.result).replace("\n", " ") + parts.append(f"result={result[:100]}") + else: + parts.append(content.type) + print(f"[{message.role}] {' | '.join(parts)}") + return changed + + +planning_skill = InlineSkill( + frontmatter=SkillFrontmatter( + name="planning", + description="Create a concise implementation plan before execution.", + ), + instructions=dedent("""\ + # Planning Skill + + Before implementation: + 1. Identify the required behavior. + 2. List the smallest code changes. + 3. Define the targeted validation. + """), +) + + +async def main() -> None: + # 1. Wrap the skill strategy. The wrapper automatically adopts its + # unload_skill trigger and exposes that single tool. + triggered_compaction = AgentTriggeredCompactionStrategy( + UnloadSkillsCompactionStrategy(), + ) + + # 2. Use the standard Foundry chat client. Run `az login` before starting. + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["FOUNDRY_MODEL"], + credential=AzureCliCredential(), + ) + + # 3. The trigger tool is registered explicitly. SkillsProvider contributes + # load_skill, while disabling approval keeps this sample automatic. + agent = Agent( + client=client, + name="PlanningAgent", + instructions=dedent("""\ + If the current context contains a completed unload_skill result, the planning + skill is already unloaded: do not call any tools, return the plan, and say + that execution can begin. + + Otherwise, follow this sequence exactly: + 1. Call load_skill with skill_name="planning". + 2. Use the loaded instructions to create a concise plan. + 3. Call unload_skill with skill_name="planning". + + Do not call load_skill and unload_skill in parallel. + """), + tools=[triggered_compaction.tool], + context_providers=[ + SkillsProvider( + planning_skill, + disable_load_skill_approval=True, + ) + ], + compaction_strategy=InspectingCompactionStrategy(triggered_compaction), + default_options={"store": False}, + ) + + # 4. The final model call printed by the middleware no longer contains the + # load_skill call or its large planning-skill result. + result = await agent.run("Plan a small implementation, unload the planning skill, then move to execution.") + print(f"\nAgent: {result.messages[-1].text}") + + +if __name__ == "__main__": + asyncio.run(main()) + +""" +Example output (model wording varies): + +--- Model call 1 --- +[user] Plan a small implementation, unload the planning skill, then move to execution. + +--- Model call 2 --- +[user] Plan a small implementation, unload the planning skill, then move to execution. +[assistant] load_skill({"skill_name":"planning"}) +[tool] result=planning Create a concise implementation plan... + +--- Model call 3 --- +Visible messages: 3 +[user] Plan a small implementation, unload the planning skill, then move to execution. +[assistant] unload_skill({"skill_name":"planning"}) +[tool] result=Skill(s) 'planning' will be unloaded before the next model call. + +Agent: [The generated plan, followed by confirmation that execution can begin.] +""" diff --git a/python/samples/02-agents/compaction/agent_triggered_compaction_basic.py b/python/samples/02-agents/compaction/agent_triggered_compaction_basic.py new file mode 100644 index 00000000000..8b7eaecc439 --- /dev/null +++ b/python/samples/02-agents/compaction/agent_triggered_compaction_basic.py @@ -0,0 +1,141 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os +from itertools import count + +from agent_framework import ( + Agent, + AgentTriggeredCompactionStrategy, + CompactionStrategy, + Message, + ToolResultCompactionStrategy, + included_messages, + tool, +) +from agent_framework.foundry import FoundryChatClient +from azure.identity import AzureCliCredential +from dotenv import load_dotenv + +load_dotenv() + +"""This sample demonstrates generic agent-triggered compaction. + +The agent first calls ``collect_context`` to produce a deliberately large tool +result. It then calls ``compact_context``. Before the next model call, +``ToolResultCompactionStrategy`` replaces the older context-producing tool +group with a compact summary, and the wrapper removes the completed compaction +trigger group. + +Environment variables: + FOUNDRY_PROJECT_ENDPOINT — Microsoft Foundry project endpoint. + FOUNDRY_MODEL — Model deployment name. + +Run with: + uv run samples/02-agents/compaction/agent_triggered_compaction_basic.py +""" + + +MODEL_CALLS = count(1) + + +@tool(approval_mode="never_require") +def collect_context(topic: str) -> str: + """Collect intentionally verbose context for a topic.""" + return "\n".join( + f"Exploration note {index}: {topic} should remain simple, testable, and observable." for index in range(1, 21) + ) + + +class InspectingCompactionStrategy: + """Print the projected messages after applying another strategy.""" + + def __init__(self, strategy: CompactionStrategy) -> None: + self.strategy = strategy + self.requires_client_managed_history = getattr(strategy, "requires_client_managed_history", False) + + async def __call__(self, messages: list[Message]) -> bool: + changed = await self.strategy(messages) + visible_messages = included_messages(messages) + + print(f"\n--- Model call {next(MODEL_CALLS)} ---") + print(f"Visible messages: {len(visible_messages)}") + for message in visible_messages: + parts: list[str] = [] + for content in message.contents: + if content.type == "text": + text = (content.text or "").replace("\n", " ") + parts.append(text[:100]) + elif content.type == "function_call": + parts.append(f"{content.name}({content.arguments})") + elif content.type == "function_result": + result = str(content.result).replace("\n", " ") + parts.append(f"result={result[:100]}") + else: + parts.append(content.type) + print(f"[{message.role}] {' | '.join(parts)}") + return changed + + +async def main() -> None: + # 1. Wrap an existing strategy. The trigger does not count toward the + # keep-last setting, so 0 compacts every workload tool-call group. + # A summary remains, allowing this sample to remove the trigger too. + triggered_compaction = AgentTriggeredCompactionStrategy( + ToolResultCompactionStrategy(keep_last_tool_call_groups=0), + exclude_trigger=True, + ) + + # 2. Use Foundry with store=False so Agent Framework owns the full in-run + # message list. Run `az login` before starting. + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["FOUNDRY_MODEL"], + credential=AzureCliCredential(), + ) + + # 3. Register both the ordinary context-producing tool and the explicit + # compaction trigger tool. + agent = Agent( + client=client, + name="ResearchAgent", + instructions=( + "If the current context contains an assistant message starting with '[Tool results:', " + "compaction is complete: do not call any tools and answer the user in one sentence. " + "Otherwise, call collect_context once and then call compact_context. " + "Do not call the tools in parallel." + ), + tools=[collect_context, triggered_compaction.tool], + compaction_strategy=InspectingCompactionStrategy(triggered_compaction), + default_options={"store": False}, + ) + + # 4. The third model input contains a compact tool-result summary, but + # neither of the original tool-call groups. + result = await agent.run("Identify the main implementation priorities.") + print(f"\nAgent: {result.messages[-1].text}") + + +if __name__ == "__main__": + asyncio.run(main()) + +""" +Example output (model wording varies): + +--- Model call 1 --- +Visible messages: 1 +[user] Identify the main implementation priorities. + +--- Model call 2 --- +Visible messages: 3 +[user] Identify the main implementation priorities. +[assistant] collect_context({"topic":"implementation priorities"}) +[tool] result=Exploration note 1: implementation priorities should remain simple, testable... + +--- Model call 3 --- +Visible messages: 2 +[user] Identify the main implementation priorities. +[assistant] [Tool results: collect_context: Exploration note 1: implementation priorities...] + +Agent: Keep the implementation simple, testable, and observable. +"""