diff --git a/docs/decisions/0019-python-context-compaction-strategy.md b/docs/decisions/0019-python-context-compaction-strategy.md index 11e1c091e5..8fffb185d1 100644 --- a/docs/decisions/0019-python-context-compaction-strategy.md +++ b/docs/decisions/0019-python-context-compaction-strategy.md @@ -1240,3 +1240,10 @@ class AttributionAwareStrategy(CompactionStrategy): - [ADR-0016: Unifying Context Management with ContextPlugin](0016-python-context-middleware.md) — Parent ADR that established `ContextProvider`, `HistoryProvider`, and `AgentSession` architecture. - [Context Compaction Limitations Analysis](https://gist.github.com/victordibia/ec3f3baf97345f7e47da025cf55b999f) — Detailed analysis of why current architecture cannot support in-run compaction, with attempted solutions and their failure modes. Option 4 in this ADR corresponds to "Option A: Middleware Access to Mutable Message Source" from that analysis; Options 1-3 correspond to "Option B: Tool Loop Hook", adapted here to a `BaseChatClient` hook instead of `FunctionInvocationConfiguration`. + +### Implementation Rollout Note + +Implementation is split into two phases: + +1. **Phase 1 (PR 1):** runtime compaction foundation in `agent_framework/_compaction.py`, in-run integration, and extensive core tests, plus in-run compaction samples (`basics`, `advanced`, `custom`). +2. **Phase 2 (PR 2):** history/storage compaction (`upsert`-based full replacement), provider support, storage tests, and storage-focused sample (`storage`). diff --git a/dotnet/nuget/nuget-package.props b/dotnet/nuget/nuget-package.props index ee3b144b06..7b241e9d56 100644 --- a/dotnet/nuget/nuget-package.props +++ b/dotnet/nuget/nuget-package.props @@ -2,11 +2,11 @@ 1.0.0 - 3 + 4 $(VersionPrefix)-rc$(RCNumber) - $(VersionPrefix)-$(VersionSuffix).260304.1 - $(VersionPrefix)-preview.260304.1 - 1.0.0-rc3 + $(VersionPrefix)-$(VersionSuffix).260311.1 + $(VersionPrefix)-preview.260311.1 + 1.0.0-rc4 Debug;Release;Publish true diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 7ecab1b442..9bdb542519 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -7,9 +7,32 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [1.0.0rc4] - 2026-03-11 + +### Added + +- **agent-framework-core**: Add `propagate_session` to `as_tool()` for session sharing in agent-as-tool scenarios ([#4439](https://github.com/microsoft/agent-framework/pull/4439)) +- **agent-framework-core**: Forward runtime kwargs to skill resource functions ([#4417](https://github.com/microsoft/agent-framework/pull/4417)) +- **samples**: Add A2A server sample ([#4528](https://github.com/microsoft/agent-framework/pull/4528)) + +### Changed + +- **agent-framework-github-copilot**: [BREAKING] Update integration to use `ToolInvocation` and `ToolResult` types ([#4551](https://github.com/microsoft/agent-framework/pull/4551)) +- **agent-framework-azure-ai**: [BREAKING] Upgrade to `azure-ai-projects` 2.0+ ([#4536](https://github.com/microsoft/agent-framework/pull/4536)) + ### Fixed +- **agent-framework-core**: Propagate MCP `isError` flag through the function middleware pipeline ([#4511](https://github.com/microsoft/agent-framework/pull/4511)) +- **agent-framework-core**: Fix `as_agent()` not defaulting name/description from client properties ([#4484](https://github.com/microsoft/agent-framework/pull/4484)) +- **agent-framework-core**: Exclude `conversation_id` from chat completions API options ([#4517](https://github.com/microsoft/agent-framework/pull/4517)) +- **agent-framework-core**: Fix conversation ID propagation when `chat_options` is a dict ([#4340](https://github.com/microsoft/agent-framework/pull/4340)) +- **agent-framework-core**: Auto-finalize `ResponseStream` on iteration completion ([#4478](https://github.com/microsoft/agent-framework/pull/4478)) +- **agent-framework-core**: Prevent pickle deserialization of untrusted HITL HTTP input ([#4566](https://github.com/microsoft/agent-framework/pull/4566)) +- **agent-framework-core**: Fix `executor_completed` event handling for non-copyable `raw_representation` in mixed workflows ([#4493](https://github.com/microsoft/agent-framework/pull/4493)) +- **agent-framework-core**: Fix `store=False` not overriding client default ([#4569](https://github.com/microsoft/agent-framework/pull/4569)) - **agent-framework-redis**: Fix `RedisContextProvider` compatibility with redisvl 0.14.0 by using `AggregateHybridQuery` ([#3954](https://github.com/microsoft/agent-framework/pull/3954)) +- **samples**: Fix `chat_response_cancellation` sample to use `Message` objects ([#4532](https://github.com/microsoft/agent-framework/pull/4532)) +- **agent-framework-purview**: Fix broken link in Purview README (Microsoft 365 Dev Program URL) ([#4610](https://github.com/microsoft/agent-framework/pull/4610)) ## [1.0.0rc3] - 2026-03-04 @@ -745,7 +768,8 @@ Release candidate for **agent-framework-core** and **agent-framework-azure-ai** For more information, see the [announcement blog post](https://devblogs.microsoft.com/foundry/introducing-microsoft-agent-framework-the-open-source-engine-for-agentic-ai-apps/). -[Unreleased]: https://github.com/microsoft/agent-framework/compare/python-1.0.0rc3...HEAD +[Unreleased]: https://github.com/microsoft/agent-framework/compare/python-1.0.0rc4...HEAD +[1.0.0rc4]: https://github.com/microsoft/agent-framework/compare/python-1.0.0rc3...python-1.0.0rc4 [1.0.0rc3]: https://github.com/microsoft/agent-framework/compare/python-1.0.0rc2...python-1.0.0rc3 [1.0.0rc2]: https://github.com/microsoft/agent-framework/compare/python-1.0.0rc1...python-1.0.0rc2 [1.0.0rc1]: https://github.com/microsoft/agent-framework/compare/python-1.0.0b260212...python-1.0.0rc1 diff --git a/python/packages/a2a/pyproject.toml b/python/packages/a2a/pyproject.toml index b7bfdb9275..4d015305c7 100644 --- a/python/packages/a2a/pyproject.toml +++ b/python/packages/a2a/pyproject.toml @@ -4,7 +4,7 @@ description = "A2A integration for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260304" +version = "1.0.0b260311" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -23,7 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "a2a-sdk>=0.3.5", ] diff --git a/python/packages/ag-ui/pyproject.toml b/python/packages/ag-ui/pyproject.toml index e41176e4c0..355405142e 100644 --- a/python/packages/ag-ui/pyproject.toml +++ b/python/packages/ag-ui/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "agent-framework-ag-ui" -version = "1.0.0b260304" +version = "1.0.0b260311" description = "AG-UI protocol integration for Agent Framework" readme = "README.md" license-files = ["LICENSE"] @@ -22,7 +22,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "ag-ui-protocol>=0.1.9", "fastapi>=0.115.0", "uvicorn>=0.30.0" diff --git a/python/packages/anthropic/pyproject.toml b/python/packages/anthropic/pyproject.toml index 51631bdd30..95be433e5a 100644 --- a/python/packages/anthropic/pyproject.toml +++ b/python/packages/anthropic/pyproject.toml @@ -4,7 +4,7 @@ description = "Anthropic integration for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260304" +version = "1.0.0b260311" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -23,7 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "anthropic>=0.70.0,<1", ] diff --git a/python/packages/azure-ai-search/pyproject.toml b/python/packages/azure-ai-search/pyproject.toml index 0827c2d816..d391de0d93 100644 --- a/python/packages/azure-ai-search/pyproject.toml +++ b/python/packages/azure-ai-search/pyproject.toml @@ -4,7 +4,7 @@ description = "Azure AI Search integration for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260304" +version = "1.0.0b260311" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -23,7 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "azure-search-documents==11.7.0b2", ] diff --git a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py index 3c4fb68fe8..4c065174ea 100644 --- a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py +++ b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py @@ -16,6 +16,13 @@ from agent_framework_azure_ai_search._context_provider import AzureAISearchConte # -- Helpers ------------------------------------------------------------------- +@pytest.fixture(autouse=True) +def clear_azure_search_environment(monkeypatch: pytest.MonkeyPatch) -> None: + for key in tuple(os.environ): + if key.startswith("AZURE_SEARCH_"): + monkeypatch.delenv(key, raising=False) + + class MockSearchResults: """Async-iterable mock for Azure SearchClient.search() results.""" diff --git a/python/packages/azure-ai/pyproject.toml b/python/packages/azure-ai/pyproject.toml index 2bd51729c2..0df9533a0b 100644 --- a/python/packages/azure-ai/pyproject.toml +++ b/python/packages/azure-ai/pyproject.toml @@ -4,7 +4,7 @@ description = "Azure AI Foundry integration for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0rc3" +version = "1.0.0rc4" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -23,7 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "azure-ai-agents == 1.2.0b5", "azure-ai-inference>=1.0.0b9", "aiohttp", diff --git a/python/packages/azure-cosmos/pyproject.toml b/python/packages/azure-cosmos/pyproject.toml index cae3b3168c..24ffbf8886 100644 --- a/python/packages/azure-cosmos/pyproject.toml +++ b/python/packages/azure-cosmos/pyproject.toml @@ -4,7 +4,7 @@ description = "Azure Cosmos DB history provider integration for Microsoft Agent authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260304" +version = "1.0.0b260311" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -23,7 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "azure-cosmos>=4.9.0", ] diff --git a/python/packages/azurefunctions/pyproject.toml b/python/packages/azurefunctions/pyproject.toml index 0bb2ec9612..c9e7890ede 100644 --- a/python/packages/azurefunctions/pyproject.toml +++ b/python/packages/azurefunctions/pyproject.toml @@ -4,7 +4,7 @@ description = "Azure Functions integration for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260304" +version = "1.0.0b260311" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -22,7 +22,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "agent-framework-durabletask", "azure-functions", "azure-functions-durable", diff --git a/python/packages/bedrock/pyproject.toml b/python/packages/bedrock/pyproject.toml index b99ecb91ff..4f1db9f4f3 100644 --- a/python/packages/bedrock/pyproject.toml +++ b/python/packages/bedrock/pyproject.toml @@ -4,7 +4,7 @@ description = "Amazon Bedrock integration for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260304" +version = "1.0.0b260311" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -23,7 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "boto3>=1.35.0,<2.0.0", "botocore>=1.35.0,<2.0.0", ] diff --git a/python/packages/chatkit/pyproject.toml b/python/packages/chatkit/pyproject.toml index 74d7216da6..d6fa2bb382 100644 --- a/python/packages/chatkit/pyproject.toml +++ b/python/packages/chatkit/pyproject.toml @@ -4,7 +4,7 @@ description = "OpenAI ChatKit integration for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260304" +version = "1.0.0b260311" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -22,7 +22,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "openai-chatkit>=1.4.0,<2.0.0", ] diff --git a/python/packages/claude/pyproject.toml b/python/packages/claude/pyproject.toml index f1891586f8..2f67d8d947 100644 --- a/python/packages/claude/pyproject.toml +++ b/python/packages/claude/pyproject.toml @@ -4,7 +4,7 @@ description = "Claude Agent SDK integration for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260304" +version = "1.0.0b260311" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -23,7 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "claude-agent-sdk>=0.1.25", ] diff --git a/python/packages/copilotstudio/pyproject.toml b/python/packages/copilotstudio/pyproject.toml index c37fa71ecf..c6d382b923 100644 --- a/python/packages/copilotstudio/pyproject.toml +++ b/python/packages/copilotstudio/pyproject.toml @@ -4,7 +4,7 @@ description = "Copilot Studio integration for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260304" +version = "1.0.0b260311" license-files = ["LICENSE"] urls.homepage = "https://aka.ms/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" @@ -23,7 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "agent-framework-core>=1.0.0rc3", + "agent-framework-core>=1.0.0rc4", "microsoft-agents-copilotstudio-client>=0.3.1", ] diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index a270bc1686..859858f0ef 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -13,6 +13,7 @@ agent_framework/ ├── _tools.py # Tool definitions and function invocation ├── _middleware.py # Middleware system for request/response interception ├── _sessions.py # AgentSession and context provider abstractions +├── _skills.py # Agent Skills system (models, executors, provider) ├── _mcp.py # Model Context Protocol support ├── _workflows/ # Workflow orchestration (sequential, concurrent, handoff, etc.) ├── openai/ # Built-in OpenAI client @@ -63,6 +64,14 @@ agent_framework/ - **`BaseContextProvider`** - Base class for context providers (RAG, memory systems) - **`BaseHistoryProvider`** - Base class for conversation history storage +### Skills (`_skills.py`) + +- **`Skill`** - A skill definition bundling instructions (`content`) with metadata, resources, and scripts. Supports `@skill.resource` and `@skill.script` decorators for adding components. +- **`SkillResource`** - Named supplementary content attached to a skill; holds either static `content` or a dynamic `function` (sync or async). Exactly one must be provided. +- **`SkillScript`** - An executable script attached to a skill; holds either an inline `function` (code-defined, runs in-process) or a `path` to a file on disk (file-based, delegated to a runner). Exactly one must be provided. +- **`SkillScriptRunner`** - Protocol for file-based script execution. Any callable matching `(skill, script, args) -> Any` satisfies it. Code-defined scripts do not use a runner. +- **`SkillsProvider`** - Context provider (extends `BaseContextProvider`) that discovers file-based skills from `SKILL.md` files and/or accepts code-defined `Skill` instances. Follows progressive disclosure: advertise → load → read resources / run scripts. + ### Workflows (`_workflows/`) - **`Workflow`** - Graph-based workflow definition diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index ef03652898..95d9b97d64 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -29,6 +29,34 @@ from ._clients import ( SupportsMCPTool, SupportsWebSearchTool, ) +from ._compaction import ( + COMPACTION_STATE_KEY, + EXCLUDE_REASON_KEY, + EXCLUDED_KEY, + GROUP_ANNOTATION_KEY, + GROUP_HAS_REASONING_KEY, + GROUP_ID_KEY, + GROUP_INDEX_KEY, + GROUP_KIND_KEY, + GROUP_TOKEN_COUNT_KEY, + SUMMARIZED_BY_SUMMARY_ID_KEY, + SUMMARY_OF_GROUP_IDS_KEY, + SUMMARY_OF_MESSAGE_IDS_KEY, + CharacterEstimatorTokenizer, + CompactionProvider, + CompactionStrategy, + SelectiveToolCallCompactionStrategy, + SlidingWindowStrategy, + SummarizationStrategy, + TokenBudgetComposedStrategy, + TokenizerProtocol, + ToolResultCompactionStrategy, + TruncationStrategy, + annotate_message_groups, + apply_compaction, + included_messages, + included_token_count, +) from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPWebsocketTool from ._middleware import ( AgentContext, @@ -59,7 +87,13 @@ from ._sessions import ( register_state_type, ) from ._settings import SecretString, load_settings -from ._skills import Skill, SkillResource, SkillsProvider +from ._skills import ( + Skill, + SkillResource, + SkillScript, + SkillScriptRunner, + SkillsProvider, +) from ._telemetry import ( AGENT_FRAMEWORK_USER_AGENT, APP_INFO, @@ -190,7 +224,19 @@ from .exceptions import ( __all__ = [ "AGENT_FRAMEWORK_USER_AGENT", "APP_INFO", + "COMPACTION_STATE_KEY", "DEFAULT_MAX_ITERATIONS", + "EXCLUDED_KEY", + "EXCLUDE_REASON_KEY", + "GROUP_ANNOTATION_KEY", + "GROUP_HAS_REASONING_KEY", + "GROUP_ID_KEY", + "GROUP_INDEX_KEY", + "GROUP_KIND_KEY", + "GROUP_TOKEN_COUNT_KEY", + "SUMMARIZED_BY_SUMMARY_ID_KEY", + "SUMMARY_OF_GROUP_IDS_KEY", + "SUMMARY_OF_MESSAGE_IDS_KEY", "USER_AGENT_KEY", "USER_AGENT_TELEMETRY_DISABLED_ENV_VAR", "Agent", @@ -212,6 +258,7 @@ __all__ = [ "BaseEmbeddingClient", "BaseHistoryProvider", "Case", + "CharacterEstimatorTokenizer", "ChatAndFunctionMiddlewareTypes", "ChatContext", "ChatMiddleware", @@ -221,6 +268,8 @@ __all__ = [ "ChatResponse", "ChatResponseUpdate", "CheckpointStorage", + "CompactionProvider", + "CompactionStrategy", "Content", "ContinuationToken", "Default", @@ -267,13 +316,18 @@ __all__ = [ "Runner", "RunnerContext", "SecretString", + "SelectiveToolCallCompactionStrategy", "SessionContext", "SingleEdgeGroup", "Skill", "SkillResource", + "SkillScript", + "SkillScriptRunner", "SkillsProvider", + "SlidingWindowStrategy", "SubWorkflowRequestMessage", "SubWorkflowResponseMessage", + "SummarizationStrategy", "SupportsAgentRun", "SupportsChatGetResponse", "SupportsCodeInterpreterTool", @@ -286,8 +340,12 @@ __all__ = [ "SwitchCaseEdgeGroupCase", "SwitchCaseEdgeGroupDefault", "TextSpanRegion", + "TokenBudgetComposedStrategy", + "TokenizerProtocol", "ToolMode", + "ToolResultCompactionStrategy", "ToolTypes", + "TruncationStrategy", "TypeCompatibilityError", "UpdateT", "UsageDetails", @@ -314,12 +372,16 @@ __all__ = [ "__version__", "add_usage_details", "agent_middleware", + "annotate_message_groups", + "apply_compaction", "chat_middleware", "create_edge_runner", "detect_media_type_from_base64", "executor", "function_middleware", "handler", + "included_messages", + "included_token_count", "load_settings", "map_chat_to_agent_update", "merge_chat_options", diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 5cf7ff78a2..2b35b96e58 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -74,6 +74,7 @@ else: from typing_extensions import Self, TypedDict # pragma: no cover if TYPE_CHECKING: + from ._compaction import CompactionStrategy, TokenizerProtocol from ._types import ChatOptions logger = logging.getLogger("agent_framework") @@ -177,6 +178,8 @@ class _RunContext(TypedDict): session_messages: Sequence[Message] agent_name: str chat_options: MutableMapping[str, Any] + compaction_strategy: CompactionStrategy | None + tokenizer: TokenizerProtocol | None filtered_kwargs: Mapping[str, Any] finalize_kwargs: Mapping[str, Any] @@ -665,6 +668,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, context_providers: Sequence[BaseContextProvider] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> None: """Initialize a Agent instance. @@ -688,6 +693,10 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] Note: response_format typing does not flow into run outputs when set via default_options. These can be overridden at runtime via the ``options`` parameter of ``run()``. tools: The tools to use for the request. + compaction_strategy: Optional agent-level in-run compaction. + If both this and a compaction_strategy on the underlying client are set, this one is used. + tokenizer: Optional agent-level tokenizer. + If both this and a tokenizer on the underlying client are set, this one is used. kwargs: Any additional keyword arguments. Will be stored as ``additional_properties``. """ opts = dict(default_options) if default_options else {} @@ -705,6 +714,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] **kwargs, ) self.client = client + self.compaction_strategy = compaction_strategy + self.tokenizer = tokenizer # Get tools from options or named parameter (named param takes precedence) tools_ = tools if tools is not None else opts.pop("tools", None) @@ -799,6 +810,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] session: AgentSession | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[ResponseModelBoundT], + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse[ResponseModelBoundT]]: ... @@ -811,6 +824,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] session: AgentSession | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: OptionsCoT | ChatOptions[None] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse[Any]]: ... @@ -823,6 +838,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] session: AgentSession | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: OptionsCoT | ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ... @@ -834,6 +851,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] session: AgentSession | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: OptionsCoT | ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: """Run the agent with the given messages and options. @@ -857,8 +876,14 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] ``Agent[OpenAIChatOptions]``, this enables IDE autocomplete for provider-specific options including temperature, max_tokens, model_id, tool_choice, and provider-specific options like reasoning_effort. - kwargs: Additional keyword arguments for the agent. - Will only be passed to functions that are called. + compaction_strategy: Optional per-run compaction override passed to + ``client.get_response()``. When omitted, the agent-level override + is used, falling back to the client default. + tokenizer: Optional per-run tokenizer override passed to + ``client.get_response()``. When omitted, the agent-level override + is used, falling back to the client default. + kwargs: Additional keyword arguments for the agent. These are only + passed to functions that are called. Returns: When stream=False: An Awaitable[AgentResponse] containing the agent's response. @@ -873,6 +898,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] session=session, tools=tools, options=options, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, kwargs=kwargs, ) response = cast( @@ -881,6 +908,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] messages=ctx["session_messages"], stream=False, options=ctx["chat_options"], # type: ignore[reportArgumentType] + compaction_strategy=ctx["compaction_strategy"], + tokenizer=ctx["tokenizer"], **ctx["filtered_kwargs"], ), ) @@ -954,6 +983,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] session=session, tools=tools, options=options, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, kwargs=kwargs, ) ctx: _RunContext = ctx_holder["ctx"] # type: ignore[assignment] # Safe: we just assigned it @@ -961,6 +992,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] messages=ctx["session_messages"], stream=True, options=ctx["chat_options"], # type: ignore[reportArgumentType] + compaction_strategy=ctx["compaction_strategy"], + tokenizer=ctx["tokenizer"], **ctx["filtered_kwargs"], ) @@ -1047,6 +1080,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] session: AgentSession | None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, options: Mapping[str, Any] | None, + compaction_strategy: CompactionStrategy | None, + tokenizer: TokenizerProtocol | None, kwargs: dict[str, Any], ) -> _RunContext: opts = dict(options) if options else {} @@ -1081,9 +1116,10 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] options=opts, ) + agent_name = self._get_agent_name() + # Normalize tools normalized_tools = normalize_tools(tools_) - agent_name = self._get_agent_name() # Resolve final tool list (runtime provided tools + local MCP server tools) final_tools: list[FunctionTool | Callable[..., Any] | dict[str, Any] | Any] = [] @@ -1153,6 +1189,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] "session_messages": session_messages, "agent_name": agent_name, "chat_options": co, + "compaction_strategy": compaction_strategy or self.compaction_strategy, + "tokenizer": tokenizer or self.tokenizer, "filtered_kwargs": filtered_kwargs, "finalize_kwargs": finalize_kwargs, } @@ -1408,6 +1446,8 @@ class Agent( default_options: OptionsCoT | None = None, context_providers: Sequence[BaseContextProvider] | None = None, middleware: Sequence[MiddlewareTypes] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> None: """Initialize a Agent instance.""" @@ -1421,5 +1461,7 @@ class Agent( default_options=default_options, context_providers=context_providers, middleware=middleware, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, **kwargs, ) diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 5dd049ecd3..5f9c1bb08f 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -52,6 +52,7 @@ else: if TYPE_CHECKING: from ._agents import Agent + from ._compaction import CompactionStrategy, TokenizerProtocol from ._middleware import ( MiddlewareTypes, ) @@ -134,6 +135,8 @@ class SupportsChatGetResponse(Protocol[OptionsContraT]): *, stream: Literal[False] = ..., options: ChatOptions[ResponseModelBoundT], + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[ChatResponse[ResponseModelBoundT]]: ... @@ -144,6 +147,8 @@ class SupportsChatGetResponse(Protocol[OptionsContraT]): *, stream: Literal[False] = ..., options: OptionsContraT | ChatOptions[None] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[ChatResponse[Any]]: ... @@ -154,6 +159,8 @@ class SupportsChatGetResponse(Protocol[OptionsContraT]): *, stream: Literal[True], options: OptionsContraT | ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ... @@ -163,6 +170,8 @@ class SupportsChatGetResponse(Protocol[OptionsContraT]): *, stream: bool = False, options: OptionsContraT | ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: """Send input and return the response. @@ -171,6 +180,8 @@ class SupportsChatGetResponse(Protocol[OptionsContraT]): messages: The sequence of input messages to send. stream: Whether to stream the response. Defaults to False. options: Chat options as a TypedDict. + compaction_strategy: Optional per-call compaction override. + tokenizer: Optional per-call tokenizer override. **kwargs: Additional chat options. Returns: @@ -252,7 +263,13 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): """ OTEL_PROVIDER_NAME: ClassVar[str] = "unknown" - DEFAULT_EXCLUDE: ClassVar[set[str]] = {"additional_properties"} + compaction_strategy: CompactionStrategy | None = None + tokenizer: TokenizerProtocol | None = None + DEFAULT_EXCLUDE: ClassVar[set[str]] = { + "additional_properties", + "compaction_strategy", + "tokenizer", + } STORES_BY_DEFAULT: ClassVar[bool] = False """Whether this client stores conversation history server-side by default. @@ -267,15 +284,21 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): self, *, additional_properties: dict[str, Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> None: """Initialize a BaseChatClient instance. Keyword Args: additional_properties: Additional properties for the client. + compaction_strategy: Optional compaction strategy to apply before model calls. + tokenizer: Optional tokenizer used by token-aware compaction strategies. kwargs: Additional keyword arguments (merged into additional_properties). """ self.additional_properties = additional_properties or {} + self.compaction_strategy = compaction_strategy + self.tokenizer = tokenizer super().__init__(**kwargs) def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: @@ -337,6 +360,46 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): finalizer=lambda updates: self._finalize_response_updates(updates, response_format=response_format), ) + async def _prepare_messages_for_model_call( + self, + messages: Sequence[Message], + *, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + ) -> list[Message]: + prepared_messages = list(messages) + if compaction_strategy is None: + if tokenizer is None: + return prepared_messages + from ._compaction import annotate_message_groups + + annotate_message_groups(prepared_messages, tokenizer=tokenizer) + return prepared_messages + from ._compaction import apply_compaction + + return await apply_compaction( + prepared_messages, + strategy=compaction_strategy, + tokenizer=tokenizer, + ) + + def _resolve_compaction_overrides( + self, + *, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + ) -> dict[str, Any]: + current_compaction_strategy = getattr(self, "compaction_strategy", None) + current_tokenizer = getattr(self, "tokenizer", None) + ret: dict[str, Any] = {} + if current_compaction_strategy is not None or compaction_strategy is not None: + ret["compaction_strategy"] = ( + current_compaction_strategy if compaction_strategy is None else compaction_strategy + ) + if current_tokenizer is not None or tokenizer is not None: + ret["tokenizer"] = current_tokenizer if tokenizer is None else tokenizer + return ret + # region Internal method to be implemented by derived classes @abstractmethod @@ -374,6 +437,8 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): *, stream: Literal[False] = ..., options: ChatOptions[ResponseModelBoundT], + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[ChatResponse[ResponseModelBoundT]]: ... @@ -384,6 +449,8 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): *, stream: Literal[False] = ..., options: OptionsCoT | ChatOptions[None] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[ChatResponse[Any]]: ... @@ -394,6 +461,8 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): *, stream: Literal[True], options: OptionsCoT | ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ... @@ -403,6 +472,8 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): *, stream: bool = False, options: OptionsCoT | ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: """Get a response from a chat client. @@ -411,17 +482,62 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): messages: The message or messages to send to the model. stream: Whether to stream the response. Defaults to False. options: Chat options as a TypedDict. + compaction_strategy: Optional per-call override for in-run compaction. + When omitted, the client-level default is used. + tokenizer: Optional per-call tokenizer override. When omitted, the + client-level default is used. **kwargs: Other keyword arguments, can be used to pass function specific parameters. Returns: When streaming a response stream of ChatResponseUpdates, otherwise an Awaitable ChatResponse. """ - return self._inner_get_response( - messages=messages, - stream=stream, - options=options or {}, # type: ignore[arg-type] - **kwargs, + compaction_overrides = self._resolve_compaction_overrides( + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, ) + if not compaction_overrides: + return self._inner_get_response( + messages=messages, + stream=stream, + options=options or {}, + **kwargs, + ) + + if stream: + + async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: + prepared_messages = await self._prepare_messages_for_model_call( + messages, + **compaction_overrides, + ) + stream_response = self._inner_get_response( + messages=prepared_messages, + stream=True, + options=options or {}, + **kwargs, + ) + if isinstance(stream_response, ResponseStream): + return stream_response # type: ignore[reportUnknownVariableType] + awaited_stream_response = await stream_response + if isinstance(awaited_stream_response, ResponseStream): + return awaited_stream_response + raise ValueError("Streaming responses must return a ResponseStream.") + + return ResponseStream.from_awaitable(_get_stream()) # type: ignore[reportUnknownVariableType] + + async def _get_response() -> ChatResponse[Any]: + prepared_messages = await self._prepare_messages_for_model_call( + messages, + **compaction_overrides, + ) + return await self._inner_get_response( + messages=prepared_messages, + stream=False, + options=options or {}, + **kwargs, + ) + + return _get_response() def service_url(self) -> str: """Get the URL of the service. @@ -446,6 +562,8 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): context_providers: Sequence[Any] | None = None, middleware: Sequence[MiddlewareTypes] | None = None, function_invocation_configuration: FunctionInvocationConfiguration | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Agent[OptionsCoT]: """Create a Agent with this client. @@ -468,6 +586,10 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): context_providers: Context providers to include during agent invocation. middleware: List of middleware to intercept agent and function invocations. function_invocation_configuration: Optional function invocation configuration override. + compaction_strategy: Optional agent-level compaction override. When omitted, + client-level compaction defaults remain in effect for each call. + tokenizer: Optional agent-level tokenizer override. When omitted, + client-level tokenizer defaults remain in effect for each call. kwargs: Any additional keyword arguments. Will be stored as ``additional_properties``. Returns: @@ -504,6 +626,8 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): context_providers=context_providers, middleware=middleware, function_invocation_configuration=function_invocation_configuration, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, **kwargs, ) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py new file mode 100644 index 0000000000..07d18da695 --- /dev/null +++ b/python/packages/core/agent_framework/_compaction.py @@ -0,0 +1,1310 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import json +import logging +from collections.abc import Mapping, Sequence +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + Protocol, + TypeAlias, + runtime_checkable, +) + +from ._sessions import BaseContextProvider +from ._types import ChatResponse, Content, Message + +if TYPE_CHECKING: + from ._clients import SupportsChatGetResponse + +GroupKind: TypeAlias = Literal["system", "user", "assistant_text", "tool_call"] +GROUP_ANNOTATION_KEY = "_group" +GROUP_ID_KEY = "id" +GROUP_KIND_KEY = "kind" +GROUP_INDEX_KEY = "index" +GROUP_HAS_REASONING_KEY = "has_reasoning" +GROUP_TOKEN_COUNT_KEY = "token_count" # noqa: S105 # nosec B105 - compaction metadata key, not a credential +EXCLUDED_KEY = "_excluded" +EXCLUDE_REASON_KEY = "_exclude_reason" +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" + + +logger = logging.getLogger("agent_framework") + + +@runtime_checkable +class TokenizerProtocol(Protocol): + """Protocol for token counters used by token-aware compaction strategies.""" + + def count_tokens(self, text: str) -> int: + """Count tokens for a serialized message payload.""" + ... + + +@runtime_checkable +class CompactionStrategy(Protocol): + """Protocol for in-place message compaction strategies.""" + + async def __call__(self, messages: list[Message]) -> bool: + """Mutate message annotations and/or list contents in place. + + Assumes caller has already applied grouping annotations (and token + annotations when required by the strategy). + + Returns: + True if compaction changed message inclusion or content; otherwise False. + """ + ... + + +class CharacterEstimatorTokenizer: + """Fast heuristic tokenizer using a 4-char/token estimate.""" + + def count_tokens(self, text: str) -> int: + return max(1, len(text) // 4) + + +def _has_content_type(message: Message, content_type: str) -> bool: + return any(content.type == content_type for content in message.contents) + + +def _has_function_call(message: Message) -> bool: + return _has_content_type(message, "function_call") + + +def _has_reasoning(message: Message) -> bool: + return _has_content_type(message, "text_reasoning") + + +def _is_tool_call_assistant(message: Message) -> bool: + return message.role == "assistant" and _has_function_call(message) + + +def _is_reasoning_only_assistant(message: Message) -> bool: + if message.role != "assistant" or not message.contents: + return False + return all(content.type == "text_reasoning" for content in message.contents) + + +def _ensure_message_ids(messages: list[Message]) -> None: + for index, message in enumerate(messages): + if not message.message_id: + message.message_id = f"msg_{index}" + + +def _group_id_for(message: Message, group_index: int) -> str: + if message.message_id: + return f"group_{message.message_id}" + return f"group_index_{group_index}" + + +def group_messages(messages: list[Message]) -> list[dict[str, Any]]: + """Compute group spans and metadata for annotation. + + Returns: + Ordered list of lightweight span dicts with keys: + ``group_id``, ``kind``, ``start_index``, ``end_index``, ``has_reasoning``. + """ + _ensure_message_ids(messages) + spans: list[dict[str, Any]] = [] + i = 0 + group_index = 0 + + while i < len(messages): + current = messages[i] + + if current.role == "system": + spans.append({ + "group_id": _group_id_for(current, group_index), + "kind": "system", + "start_index": i, + "end_index": i, + "has_reasoning": _has_reasoning(current), + }) + i += 1 + group_index += 1 + continue + + if current.role == "user": + spans.append({ + "group_id": _group_id_for(current, group_index), + "kind": "user", + "start_index": i, + "end_index": i, + "has_reasoning": _has_reasoning(current), + }) + i += 1 + group_index += 1 + continue + + # Reasoning prefix before an assistant function_call joins the same tool_call group. + # This includes the OpenAI Responses shape where reasoning and function_call + # contents are co-located in the same assistant message. + if _is_reasoning_only_assistant(current): + prefix_start = i + j = i + while j < len(messages) and _is_reasoning_only_assistant(messages[j]): + j += 1 + if j < len(messages) and _is_tool_call_assistant(messages[j]): + k = j + 1 + has_reasoning = True + while k < len(messages) and _is_reasoning_only_assistant(messages[k]): + has_reasoning = True + k += 1 + while k < len(messages) and messages[k].role == "tool": + k += 1 + spans.append({ + "group_id": _group_id_for(messages[prefix_start], group_index), + "kind": "tool_call", + "start_index": prefix_start, + "end_index": k - 1, + "has_reasoning": has_reasoning or _has_reasoning(messages[j]), + }) + i = k + group_index += 1 + continue + + if _is_tool_call_assistant(current): + has_reasoning = _has_reasoning(current) + k = i + 1 + while k < len(messages) and _is_reasoning_only_assistant(messages[k]): + has_reasoning = True + k += 1 + while k < len(messages) and messages[k].role == "tool": + k += 1 + spans.append({ + "group_id": _group_id_for(current, group_index), + "kind": "tool_call", + "start_index": i, + "end_index": k - 1, + "has_reasoning": has_reasoning, + }) + i = k + group_index += 1 + continue + + if current.role == "tool": + k = i + 1 + while k < len(messages) and messages[k].role == "tool": + k += 1 + spans.append({ + "group_id": _group_id_for(current, group_index), + "kind": "tool_call", + "start_index": i, + "end_index": k - 1, + "has_reasoning": False, + }) + i = k + group_index += 1 + continue + + spans.append({ + "group_id": _group_id_for(current, group_index), + "kind": "assistant_text", + "start_index": i, + "end_index": i, + "has_reasoning": _has_reasoning(current), + }) + i += 1 + group_index += 1 + + return spans + + +def _coerce_group_kind(value: object) -> GroupKind | None: + if value == "system": + return "system" + if value == "user": + return "user" + if value == "assistant_text": + return "assistant_text" + if value == "tool_call": + return "tool_call" + return None + + +def _read_group_annotation(message: Message) -> dict[str, Any] | None: + raw_annotation = _read_group_annotation_raw(message) + if raw_annotation is None: + return None + + group_id = raw_annotation.get(GROUP_ID_KEY) + group_kind = _coerce_group_kind(raw_annotation.get(GROUP_KIND_KEY)) + group_index = raw_annotation.get(GROUP_INDEX_KEY) + has_reasoning = raw_annotation.get(GROUP_HAS_REASONING_KEY) + token_count = raw_annotation.get(GROUP_TOKEN_COUNT_KEY) + if token_count is not None and not isinstance(token_count, int): + return None + if ( + not isinstance(group_id, str) + or group_kind is None + or not isinstance(group_index, int) + or not isinstance(has_reasoning, bool) + ): + return None + + return raw_annotation + + +def _read_group_annotation_raw(message: Message) -> dict[str, Any] | None: + annotation = message.additional_properties.get(GROUP_ANNOTATION_KEY) + if isinstance(annotation, Mapping): + return annotation # type: ignore[reportUnknownVariableType, return-value] + return None + + +def _set_group_summarized_by_summary_id(message: Message, summary_id: str) -> None: + annotation = _read_group_annotation_raw(message) + if annotation is None: + annotation = {} + message.additional_properties[GROUP_ANNOTATION_KEY] = annotation + annotation[SUMMARIZED_BY_SUMMARY_ID_KEY] = summary_id + + +def _write_group_annotation( + message: Message, + *, + group_id: str, + kind: GroupKind, + index: int, + has_reasoning: bool, +) -> None: + existing_raw_annotation = _read_group_annotation_raw(message) + unknown_fields: dict[str, Any] = {} + token_count: int | None = None + if existing_raw_annotation is not None: + raw_token_count = existing_raw_annotation.get(GROUP_TOKEN_COUNT_KEY) + if isinstance(raw_token_count, int) or raw_token_count is None: + token_count = raw_token_count + unknown_fields = { + key: value + for key, value in existing_raw_annotation.items() + if key + not in { + GROUP_ID_KEY, + GROUP_KIND_KEY, + GROUP_INDEX_KEY, + GROUP_HAS_REASONING_KEY, + GROUP_TOKEN_COUNT_KEY, + } + } + + annotation = { + GROUP_ID_KEY: group_id, + GROUP_KIND_KEY: kind, + GROUP_INDEX_KEY: index, + GROUP_HAS_REASONING_KEY: has_reasoning, + GROUP_TOKEN_COUNT_KEY: token_count, + } + annotation.update(unknown_fields) + message.additional_properties[GROUP_ANNOTATION_KEY] = annotation + + +def _group_id(message: Message) -> str | None: + annotation = _read_group_annotation(message) + if annotation is None: + return None + group_id = annotation.get(GROUP_ID_KEY) + return group_id if isinstance(group_id, str) else None + + +def _group_kind(message: Message) -> GroupKind | None: + annotation = _read_group_annotation(message) + if annotation is None: + return None + return _coerce_group_kind(annotation.get(GROUP_KIND_KEY)) + + +def _group_index(message: Message) -> int | None: + annotation = _read_group_annotation(message) + if annotation is None: + return None + group_index = annotation.get(GROUP_INDEX_KEY) + return group_index if isinstance(group_index, int) else None + + +def _token_count(message: Message) -> int | None: + annotation = _read_group_annotation(message) + if annotation is None: + return None + token_count = annotation.get(GROUP_TOKEN_COUNT_KEY) + return token_count if isinstance(token_count, int) else None + + +def _write_token_count(message: Message, token_count: int) -> None: + annotation = _read_group_annotation_raw(message) + if annotation is None: + return + annotation[GROUP_TOKEN_COUNT_KEY] = token_count + message.additional_properties[GROUP_ANNOTATION_KEY] = annotation + + +def _ordered_group_ids_from_annotations(messages: Sequence[Message]) -> list[str]: + ordered_group_ids: list[str] = [] + seen: set[str] = set() + for message in messages: + group_id = _group_id(message) + if group_id is not None and group_id not in seen: + seen.add(group_id) + ordered_group_ids.append(group_id) + return ordered_group_ids + + +def _first_untokenized_index(messages: Sequence[Message]) -> int | None: + for index, message in enumerate(messages): + if _token_count(message) is None: + return index + return None + + +def _first_annotation_gaps( + messages: Sequence[Message], + *, + include_tokens: bool, +) -> tuple[int | None, int | None]: + first_unannotated: int | None = None + first_untokenized: int | None = None + for index, message in enumerate(messages): + missing_group_annotation = first_unannotated is None and _group_id(message) is None + missing_token_annotation = include_tokens and first_untokenized is None and _token_count(message) is None + + if missing_group_annotation: + first_unannotated = index + if missing_token_annotation: + first_untokenized = index + + if missing_group_annotation or missing_token_annotation: + break + return first_unannotated, first_untokenized + + +def _reannotation_start(messages: Sequence[Message], index: int) -> int: + if index <= 0: + return 0 + previous_index = index - 1 + previous_group_id = _group_id(messages[previous_index]) + if previous_group_id is None: + return previous_index + while previous_index > 0: + prior_group_id = _group_id(messages[previous_index - 1]) + if prior_group_id != previous_group_id: + break + previous_index -= 1 + return previous_index + + +def annotate_message_groups( + messages: list[Message], + *, + from_index: int | None = None, + force_reannotate: bool = False, + tokenizer: TokenizerProtocol | None = None, +) -> list[str]: + """Annotate message groups while reusing existing annotations when possible. + + By default, the function re-annotates only the suffix that contains new + messages and keeps previously annotated prefixes untouched. When a + ``tokenizer`` is provided, token-count annotations are also populated + incrementally. + """ + if not messages: + return [] + + if force_reannotate: + start_index = 0 + elif from_index is not None: + start_index = max(0, min(from_index, len(messages) - 1)) + else: + first_unannotated_index, first_untokenized_index = _first_annotation_gaps( + messages, + include_tokens=tokenizer is not None, + ) + candidate_starts = [index for index in (first_unannotated_index, first_untokenized_index) if index is not None] + if not candidate_starts: + return _ordered_group_ids_from_annotations(messages) + start_index = min(candidate_starts) + + start_index = _reannotation_start(messages, start_index) + + # Continue group indices from the preserved prefix when only re-annotating a suffix. + group_index_offset = 0 + if start_index > 0: + previous_group_index = _group_index(messages[start_index - 1]) + if previous_group_index is not None: + group_index_offset = previous_group_index + 1 + + spans = group_messages(messages[start_index:]) + for span_index, span in enumerate(spans): + group_id = str(span["group_id"]) + kind = _coerce_group_kind(span["kind"]) + if kind is None: + raise ValueError(f"Unexpected group kind in span: {span['kind']}") + local_start_index = int(span["start_index"]) + local_end_index = int(span["end_index"]) + has_reasoning = bool(span["has_reasoning"]) + for idx in range(start_index + local_start_index, start_index + local_end_index + 1): + message = messages[idx] + _write_group_annotation( + message, + group_id=group_id, + kind=kind, + index=group_index_offset + span_index, + has_reasoning=has_reasoning, + ) + message.additional_properties.setdefault(EXCLUDED_KEY, False) + if tokenizer is not None and _token_count(message) is None: + _write_token_count(message, tokenizer.count_tokens(_serialize_message(message))) + return _ordered_group_ids_from_annotations(messages) + + +def _serialize_content(content: Content) -> dict[str, Any]: + payload = content.to_dict(exclude_none=True) + payload.pop("raw_representation", None) + return payload + + +def _serialize_message(message: Message) -> str: + serialized_contents = [_serialize_content(content) for content in message.contents] + payload = { + "role": message.role, + "message_id": message.message_id, + "contents": serialized_contents, + } + return json.dumps(payload, ensure_ascii=True, sort_keys=True, default=str) + + +def annotate_token_counts( + messages: list[Message], + *, + tokenizer: TokenizerProtocol, + from_index: int | None = None, + force_retokenize: bool = False, +) -> None: + """Annotate token-count metadata, incrementally by default.""" + if not messages: + return + + # Token counts are stored inside group annotations. + annotate_message_groups(messages, from_index=from_index) + + if force_retokenize: + start_index = 0 + elif from_index is not None: + start_index = max(0, min(from_index, len(messages) - 1)) + else: + first_untokenized_index = _first_untokenized_index(messages) + if first_untokenized_index is None: + return + start_index = first_untokenized_index + + for message in messages[start_index:]: + _write_token_count(message, tokenizer.count_tokens(_serialize_message(message))) + + +def extend_compaction_messages( + messages: list[Message], + new_messages: Sequence[Message], + *, + tokenizer: TokenizerProtocol | None = None, +) -> None: + """Append a batch of messages and annotate only the appended tail.""" + if not new_messages: + return + + start_index = len(messages) + messages.extend(new_messages) + annotate_message_groups( + messages, + from_index=start_index, + tokenizer=tokenizer, + ) + + +def append_compaction_message( + messages: list[Message], + message: Message, + *, + tokenizer: TokenizerProtocol | None = None, +) -> None: + """Append a single message and incrementally annotate metadata.""" + extend_compaction_messages(messages, [message], tokenizer=tokenizer) + + +def included_messages(messages: list[Message]) -> list[Message]: + return [message for message in messages if not message.additional_properties.get(EXCLUDED_KEY, False)] + + +def included_token_count(messages: list[Message]) -> int: + total = 0 + for message in included_messages(messages): + token_count = _token_count(message) + if token_count is not None: + total += token_count + return total + + +def set_excluded(message: Message, *, excluded: bool, reason: str | None = None) -> bool: + changed = bool(message.additional_properties.get(EXCLUDED_KEY, False)) != excluded + if changed: + message.additional_properties[EXCLUDED_KEY] = excluded + if reason is not None: + message.additional_properties[EXCLUDE_REASON_KEY] = reason + return changed + + +def exclude_group_ids(messages: list[Message], group_ids: set[str], *, reason: str) -> bool: + changed = False + for message in messages: + group_id = _group_id(message) + if group_id is not None and group_id in group_ids: + changed = set_excluded(message, excluded=True, reason=reason) or changed + return changed + + +def project_included_messages(messages: list[Message]) -> list[Message]: + return included_messages(messages) + + +def _group_messages_by_id(messages: list[Message]) -> dict[str, list[Message]]: + grouped: dict[str, list[Message]] = {} + for message in messages: + group_id = _group_id(message) + if group_id is None: + continue + grouped.setdefault(group_id, []).append(message) + return grouped + + +def _group_kind_map(messages: list[Message]) -> dict[str, GroupKind]: + kinds: dict[str, GroupKind] = {} + for message in messages: + group_id = _group_id(message) + group_kind = _group_kind(message) + if group_id is not None and group_kind is not None and group_id not in kinds: + kinds[group_id] = group_kind + return kinds + + +def _group_start_indices(messages: list[Message]) -> dict[str, int]: + starts: dict[str, int] = {} + for idx, message in enumerate(messages): + group_id = _group_id(message) + if group_id is not None and group_id not in starts: + starts[group_id] = idx + return starts + + +def _included_group_ids(messages: list[Message], ordered_group_ids: list[str]) -> list[str]: + grouped = _group_messages_by_id(messages) + included_ids: list[str] = [] + for group_id in ordered_group_ids: + if any(not m.additional_properties.get(EXCLUDED_KEY, False) for m in grouped.get(group_id, [])): + included_ids.append(group_id) + return included_ids + + +def _count_included_messages(messages: list[Message]) -> int: + return len(included_messages(messages)) + + +def _count_included_tokens(messages: list[Message]) -> int: + return included_token_count(messages) + + +class TruncationStrategy: + """Oldest-first compaction using a single metric threshold. + + This strategy runs after group annotations are computed and excludes whole + groups (never partial tool-call groups). The metric is: + - token count when ``tokenizer`` is provided + - included message count when ``tokenizer`` is not provided + Compaction triggers when the metric exceeds ``max_n`` and trims to + ``compact_to``. + """ + + def __init__( + self, + *, + max_n: int, + compact_to: int, + tokenizer: TokenizerProtocol | None = None, + preserve_system: bool = True, + ) -> None: + """Create a truncation strategy. + + Keyword Args: + max_n: Trigger threshold measured in tokens when ``tokenizer`` is + provided, otherwise measured in included messages. + compact_to: Target value for the same metric used by ``max_n``. + This argument is required and must be explicitly set. + tokenizer: Optional tokenizer used for token-based truncation. + preserve_system: When True, system groups remain included and only + non-system groups are eligible for exclusion. + """ + if max_n <= 0: + raise ValueError("max_n must be greater than 0.") + if compact_to <= 0: + raise ValueError("compact_to must be greater than 0.") + if compact_to > max_n: + raise ValueError("compact_to must be less than or equal to max_n.") + self.max_n = max_n + self.compact_to = compact_to + self.tokenizer = tokenizer + self.preserve_system = preserve_system + + async def __call__(self, messages: list[Message]) -> bool: + ordered_group_ids = _ordered_group_ids_from_annotations(messages) + if self.tokenizer is not None: + over_limit = _count_included_tokens(messages) > self.max_n + else: + over_limit = _count_included_messages(messages) > self.max_n + if not over_limit: + return False + + grouped = _group_messages_by_id(messages) + kinds = _group_kind_map(messages) + protected_ids: set[str] = set() + if self.preserve_system: + protected_ids = {group_id for group_id in ordered_group_ids if kinds.get(group_id) == "system"} + + changed = False + for group_id in ordered_group_ids: + if self.tokenizer is not None: + target_met = _count_included_tokens(messages) <= self.compact_to + else: + target_met = _count_included_messages(messages) <= self.compact_to + if target_met: + break + if group_id in protected_ids: + continue + for message in grouped.get(group_id, []): + changed = set_excluded(message, excluded=True, reason="truncation") or changed + return changed + + +class SlidingWindowStrategy: + """Windowed compaction that keeps the most recent non-system groups. + + The strategy preserves recency by retaining only the last + ``keep_last_groups`` included non-system groups. System groups can be kept + as stable anchors when ``preserve_system`` is enabled. + + This can remove older user and assistant groups while keeping system + instructions, which is useful when directives must persist but conversation + history grows. Use ``SelectiveToolCallCompactionStrategy`` when only tool + groups should be reduced. + """ + + def __init__(self, *, keep_last_groups: int, preserve_system: bool = True) -> None: + """Create a sliding-window strategy. + + Args: + keep_last_groups: Number of most-recent non-system groups to keep. + preserve_system: Whether system groups should always remain included. + """ + if keep_last_groups <= 0: + raise ValueError(f"keep_last_groups must be more than 0, got {keep_last_groups}") + self.keep_last_groups = keep_last_groups + self.preserve_system = preserve_system + + async def __call__(self, messages: list[Message]) -> bool: + ordered_group_ids = _ordered_group_ids_from_annotations(messages) + grouped = _group_messages_by_id(messages) + kinds = _group_kind_map(messages) + + included_group_ids = _included_group_ids(messages, ordered_group_ids) + non_system_group_ids = [group_id for group_id in included_group_ids if kinds.get(group_id) != "system"] + keep_non_system_ids = set(non_system_group_ids[-self.keep_last_groups :]) + keep_ids = set(keep_non_system_ids) + if self.preserve_system: + keep_ids.update(group_id for group_id in ordered_group_ids if kinds.get(group_id) == "system") + + changed = False + for group_id in included_group_ids: + if group_id in keep_ids: + continue + for message in grouped.get(group_id, []): + changed = set_excluded(message, excluded=True, reason="sliding_window") or changed + return changed + + +class SelectiveToolCallCompactionStrategy: + """Compaction focused on reducing tool-call history growth. + + This strategy only targets groups annotated as ``tool_call`` and keeps the + latest ``keep_last_tool_call_groups`` included tool-call groups. It is + useful when tool chatter dominates token usage. + + It does not change non-tool-call groups, so it can be combined with other + strategies that target different aspects of the message history. + """ + + def __init__(self, *, keep_last_tool_call_groups: int = 1) -> None: + """Create a tool-call-focused compaction strategy. + + Args: + keep_last_tool_call_groups: Number of newest included tool-call + groups to retain. Set to 0 to remove all included tool-call + groups. + + Raises: + ValueError: If ``keep_last_tool_call_groups`` is negative. + """ + if keep_last_tool_call_groups < 0: + raise ValueError("keep_last_tool_call_groups must be greater than or equal to 0.") + self.keep_last_tool_call_groups = keep_last_tool_call_groups + + async def __call__(self, messages: list[Message]) -> bool: + ordered_group_ids = _ordered_group_ids_from_annotations(messages) + grouped = _group_messages_by_id(messages) + kinds = _group_kind_map(messages) + + included_tool_group_ids = [ + group_id + for group_id in _included_group_ids(messages, ordered_group_ids) + if kinds.get(group_id) == "tool_call" + ] + if len(included_tool_group_ids) <= self.keep_last_tool_call_groups: + return False + + keep_ids: set[str] = ( + set(included_tool_group_ids[-self.keep_last_tool_call_groups :]) + if self.keep_last_tool_call_groups > 0 + else set() + ) + changed = False + for group_id in included_tool_group_ids: + if group_id in keep_ids: + continue + for message in grouped.get(group_id, []): + changed = set_excluded(message, excluded=True, reason="tool_call_compaction") or changed + return changed + + +class ToolResultCompactionStrategy: + """Collapse older tool-call groups into short summary messages. + + Unlike ``SelectiveToolCallCompactionStrategy`` which fully excludes old + tool-call groups, this strategy *replaces* them with a compact summary + message containing the tool results (e.g. + ``[Tool results: get_weather: sunny, 18°C]``). This preserves a readable + trace of what tools returned while reclaiming the token overhead of the + full function-call/result message structure. + + The most recent ``keep_last_tool_call_groups`` tool-call groups are left + untouched; older ones are collapsed. + """ + + def __init__(self, *, keep_last_tool_call_groups: int = 1) -> None: + """Create a tool-result compaction strategy. + + Keyword Args: + keep_last_tool_call_groups: Number of newest included tool-call + groups to retain verbatim. Older tool-call groups are collapsed + into summary messages. Set to 0 to collapse all. + + Raises: + ValueError: If ``keep_last_tool_call_groups`` is negative. + """ + if keep_last_tool_call_groups < 0: + raise ValueError("keep_last_tool_call_groups must be greater than or equal to 0.") + self.keep_last_tool_call_groups = keep_last_tool_call_groups + + async def __call__(self, messages: list[Message]) -> bool: + ordered_group_ids = _ordered_group_ids_from_annotations(messages) + grouped = _group_messages_by_id(messages) + kinds = _group_kind_map(messages) + + included_tool_group_ids = [ + group_id + for group_id in _included_group_ids(messages, ordered_group_ids) + if kinds.get(group_id) == "tool_call" + ] + if len(included_tool_group_ids) <= self.keep_last_tool_call_groups: + return False + + keep_ids: set[str] = ( + set(included_tool_group_ids[-self.keep_last_tool_call_groups :]) + if self.keep_last_tool_call_groups > 0 + else set() + ) + starts = _group_start_indices(messages) + changed = False + for group_id in included_tool_group_ids: + if group_id in keep_ids: + continue + group_msgs = grouped.get(group_id, []) + # Build a call_id → function_name map from function_call contents. + call_id_to_name: dict[str, str] = {} + for msg in group_msgs: + 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 + # Collect tool results with the function name for context. + tool_results: list[str] = [] + for msg in group_msgs: + for content in msg.contents: + if content.type == "function_result": + result_text = content.result if isinstance(content.result, str) else str(content.result) + func_name = call_id_to_name.get(content.call_id or "", "") + label = f"{func_name}: {result_text}" if func_name else result_text + tool_results.append(label.strip()) + summary_label = "; ".join(tool_results) if tool_results else "no results" + summary_text = f"[Tool results: {summary_label}]" + + summary_id = f"tool_summary_{group_id}" + original_message_ids = [msg.message_id for msg in group_msgs if msg.message_id] + + # Mark originals as excluded with back-link to the summary. + for msg in group_msgs: + _set_group_summarized_by_summary_id(msg, summary_id) + changed = set_excluded(msg, excluded=True, reason="tool_result_compaction") or changed + + # Insert summary with forward links to the originals. + summary_annotation = { + SUMMARY_OF_MESSAGE_IDS_KEY: original_message_ids, + SUMMARY_OF_GROUP_IDS_KEY: [group_id], + } + insertion_index = starts.get(group_id, 0) + summary_message = Message( + role="assistant", + text=summary_text, + message_id=summary_id, + additional_properties={ + GROUP_ANNOTATION_KEY: summary_annotation, + }, + ) + messages.insert(insertion_index, summary_message) + annotate_message_groups(messages, from_index=insertion_index, force_reannotate=False) + starts = _group_start_indices(messages) + grouped = _group_messages_by_id(messages) + + return changed + + +def _format_messages_for_summary(messages: list[Message]) -> str: + lines: list[str] = [] + for index, message in enumerate(messages, start=1): + content_text = message.text + if not content_text: + content_text = ", ".join(content.type for content in message.contents) + lines.append(f"{index}. [{message.role}] {content_text}") + return "\n".join(lines) + + +DEFAULT_SUMMARIZATION_PROMPT: Final[ + str +] = """**Generate a clear and complete summary of the entire conversation in no more than five sentences.** + +The summary must always: +- Reflect contributions from both the user and the assistant +- Preserve context to support ongoing dialogue +- Incorporate any previously provided summary +- Emphasize the most relevant and meaningful points + +The summary must never: +- Offer critique, correction, interpretation, or speculation +- Highlight errors, misunderstandings, or judgments of accuracy +- Comment on events or ideas not present in the conversation +- Omit any details included in an earlier summary +""" + + +class SummarizationStrategy: + """Summarize older included groups and replace them with linked summary text. + + The strategy monitors included non-system message count and triggers when + that count grows beyond ``target_count + threshold``. When triggered, it + summarizes the oldest groups and retains the newest content near + ``target_count`` (subject to atomic group boundaries). It writes trace + metadata in both directions: summary -> original message/group IDs and + original -> summary ID. + """ + + def __init__( + self, + *, + client: SupportsChatGetResponse[Any], + target_count: int = 4, + threshold: int | None = 2, + prompt: str | None = None, + ) -> None: + """Create a summarization strategy. + + Keyword Args: + client: A chat client compatible with ``SupportsChatGetResponse`` + used to generate summary text. + target_count: Target number of included non-system messages to + retain after summarization. Must be greater than 0. + threshold: Extra included non-system messages allowed above + ``target_count`` before summarization triggers. Must be greater + than or equal to 0 when provided. + prompt: Optional summarization instruction. If omitted, a default + prompt that preserves goals, decisions, and unresolved items is + used. + + Raises: + ValueError: If ``target_count`` is less than 1. + ValueError: If ``threshold`` is provided and is negative. + """ + if target_count <= 0: + raise ValueError("target_count must be greater than 0.") + if threshold is not None and threshold < 0: + raise ValueError("threshold must be greater than or equal to 0.") + self.client = client + self.target_count = target_count + self.threshold = threshold if threshold is not None else 0 + self.prompt = prompt or DEFAULT_SUMMARIZATION_PROMPT + + async def __call__(self, messages: list[Message]) -> bool: + ordered_group_ids = _ordered_group_ids_from_annotations(messages) + grouped = _group_messages_by_id(messages) + kinds = _group_kind_map(messages) + starts = _group_start_indices(messages) + + included_non_system_groups: list[tuple[str, list[Message]]] = [] + included_non_system_message_count = 0 + for group_id in _included_group_ids(messages, ordered_group_ids): + if kinds.get(group_id) == "system": + continue + group_messages = [ + message + for message in grouped.get(group_id, []) + if not message.additional_properties.get(EXCLUDED_KEY, False) + ] + if not group_messages: + continue + included_non_system_groups.append((group_id, group_messages)) + included_non_system_message_count += len(group_messages) + + if included_non_system_message_count <= self.target_count + self.threshold: + return False + + keep_group_ids: list[str] = [] + retained_message_count = 0 + for group_id, group_messages in reversed(included_non_system_groups): + if retained_message_count >= self.target_count and keep_group_ids: + break + keep_group_ids.append(group_id) + retained_message_count += len(group_messages) + keep_group_id_set = set(keep_group_ids) + + group_ids_to_summarize = [ + group_id for group_id, _ in included_non_system_groups if group_id not in keep_group_id_set + ] + if not group_ids_to_summarize: + return False + + messages_to_summarize: list[Message] = [] + for group_id, group_messages in included_non_system_groups: + if group_id in keep_group_id_set: + continue + messages_to_summarize.extend(group_messages) + if not messages_to_summarize: + return False + + try: + summary_response: ChatResponse[None] = await self.client.get_response( + [ + Message(role="system", text=self.prompt), + Message( + role="user", + text=_format_messages_for_summary(messages_to_summarize), + ), + ], + stream=False, + ) + except Exception as exc: + logger.warning( + "Skipping summarization compaction: summary generation failed (%s).", + exc, + ) + return False + + summary_text = summary_response.text.strip() if summary_response.text else "" + if not summary_text: + logger.warning("Skipping summarization compaction: summarizer returned no text.") + return False + summary_id = f"summary_{len(messages)}" + original_message_ids = [message.message_id for message in messages_to_summarize if message.message_id] + summary_of_group_ids = list(group_ids_to_summarize) + summary_annotation = { + SUMMARY_OF_MESSAGE_IDS_KEY: original_message_ids, + SUMMARY_OF_GROUP_IDS_KEY: summary_of_group_ids, + } + + summary_message = Message( + role="assistant", + text=summary_text, + message_id=summary_id, + additional_properties={ + GROUP_ANNOTATION_KEY: summary_annotation, + }, + ) + + for message in messages_to_summarize: + _set_group_summarized_by_summary_id(message, summary_id) + set_excluded(message, excluded=True, reason="summarized") + + insertion_index = min(starts[group_id] for group_id in group_ids_to_summarize if group_id in starts) + messages.insert(insertion_index, summary_message) + annotate_message_groups(messages, from_index=insertion_index, force_reannotate=False) + return True + + +class TokenBudgetComposedStrategy: + """Compose multiple strategies until an included-token budget is satisfied. + + Strategies run in the provided order over shared message annotations. After + each step, token counts are refreshed. If no strategy reaches budget, a + deterministic fallback excludes oldest groups (and finally anchors when + necessary) to enforce the limit. + """ + + def __init__( + self, + *, + token_budget: int, + tokenizer: TokenizerProtocol, + strategies: Sequence[CompactionStrategy], + early_stop: bool = True, + ) -> None: + """Create a composed token-budget strategy. + + Args: + token_budget: Maximum included token count allowed after compaction. + tokenizer: Tokenizer implementation used for per-message token + annotation. + strategies: Ordered strategy sequence to execute before fallback. + early_stop: When True, stop as soon as budget is satisfied. + """ + self.token_budget = token_budget + self.tokenizer = tokenizer + self.strategies = list(strategies) + self.early_stop = early_stop + + async def __call__(self, messages: list[Message]) -> bool: + annotate_message_groups(messages) + annotate_token_counts(messages, tokenizer=self.tokenizer) + if included_token_count(messages) <= self.token_budget: + return False + + changed = False + for strategy in self.strategies: + changed = (await strategy(messages)) or changed + annotate_message_groups(messages) + annotate_token_counts(messages, tokenizer=self.tokenizer) + if self.early_stop and included_token_count(messages) <= self.token_budget: + return changed + + if included_token_count(messages) <= self.token_budget: + return changed + + ordered_group_ids = annotate_message_groups(messages) + grouped = _group_messages_by_id(messages) + kinds = _group_kind_map(messages) + for group_id in ordered_group_ids: + if kinds.get(group_id) == "system": + continue + for message in grouped.get(group_id, []): + changed = set_excluded(message, excluded=True, reason="token_budget_fallback") or changed + if included_token_count(messages) <= self.token_budget: + break + if included_token_count(messages) <= self.token_budget: + return changed + + # Strict budget enforcement fallback: if anchors alone exceed budget, exclude remaining groups. + for group_id in ordered_group_ids: + if kinds.get(group_id) != "system": + continue + for message in grouped.get(group_id, []): + changed = set_excluded(message, excluded=True, reason="token_budget_fallback_strict") or changed + if included_token_count(messages) <= self.token_budget: + break + return changed + + +async def apply_compaction( + messages: list[Message], + *, + strategy: CompactionStrategy | None, + tokenizer: TokenizerProtocol | None = None, +) -> list[Message]: + """Apply configured compaction and return projected model-input messages.""" + if strategy is None: + return messages + annotate_message_groups(messages) + if tokenizer is not None: + annotate_token_counts(messages, tokenizer=tokenizer) + await strategy(messages) + return project_included_messages(messages) + + +COMPACTION_STATE_KEY: Final[str] = "_compaction_messages" + + +class CompactionProvider(BaseContextProvider): + """Context provider that compacts messages before and after agent runs. + + This provider accepts two separate strategies: + + - ``before_strategy``: Runs in ``before_run`` on messages already in the + context (loaded by earlier providers such as a history provider). + Compacts the loaded history before it reaches the model. + - ``after_strategy``: Runs in ``after_run`` on the accumulated messages + stored by a history provider in session state. This compacts the + persisted history so the next turn starts with a smaller context. + + Either strategy may be ``None`` to skip that phase. + + Examples: + .. code-block:: python + + from agent_framework import Agent, CompactionProvider, InMemoryHistoryProvider + from agent_framework._compaction import ( + SlidingWindowStrategy, + ToolResultCompactionStrategy, + ) + + history = InMemoryHistoryProvider() + compaction = CompactionProvider( + before_strategy=SlidingWindowStrategy(keep_last_groups=20), + after_strategy=ToolResultCompactionStrategy(keep_last_tool_call_groups=1), + history_source_id=history.source_id, + ) + agent = Agent( + client=client, + name="assistant", + context_providers=[history, compaction], + ) + session = agent.create_session() + await agent.run("Hello", session=session) + """ + + def __init__( + self, + *, + before_strategy: CompactionStrategy | None = None, + after_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, + source_id: str = "compaction", + history_source_id: str = "in_memory", + ) -> None: + """Create a compaction provider. + + Keyword Args: + before_strategy: Strategy applied to loaded context messages before + the model runs. ``None`` to skip pre-run compaction. + after_strategy: Strategy applied to stored history messages after + the model runs. Requires ``history_source_id`` to locate the + messages in session state. ``None`` to skip post-run compaction. + tokenizer: Optional tokenizer for token-aware strategies. + source_id: Provider source id (default ``"compaction"``). + history_source_id: The ``source_id`` of the history provider whose + stored messages the ``after_strategy`` should compact + (default ``"in_memory"``). + """ + super().__init__(source_id) + self.before_strategy = before_strategy + self.after_strategy = after_strategy + self.tokenizer = tokenizer + self.history_source_id = history_source_id + + async def before_run( + self, + *, + agent: Any, + session: Any, + context: Any, + state: dict[str, Any], + ) -> None: + """Compact messages already present in the context from earlier providers.""" + if self.before_strategy is None: + return + + all_messages: list[Message] = context.get_messages() + if not all_messages: + return + + annotate_message_groups(all_messages) + if self.tokenizer is not None: + annotate_token_counts(all_messages, tokenizer=self.tokenizer) + await self.before_strategy(all_messages) + + projected = project_included_messages(all_messages) + projected_set = {id(m) for m in projected} + for sid in list(context.context_messages): + context.context_messages[sid] = [m for m in context.context_messages[sid] if id(m) in projected_set] + + async def after_run( + self, + *, + agent: Any, + session: Any, + context: Any, + state: dict[str, Any], + ) -> None: + """Compact stored history messages after the model runs.""" + if self.after_strategy is None: + return + + # Access the history provider's stored messages from session state. + history_state_raw = session.state.get(self.history_source_id) if session else None + if not isinstance(history_state_raw, dict): + return + history_state: dict[str, Any] = history_state_raw # type: ignore[assignment] + raw_messages = history_state.get("messages") + if not isinstance(raw_messages, list) or not raw_messages: + return + stored_messages: list[Message] = raw_messages # type: ignore[assignment] + + annotate_message_groups(stored_messages) + if self.tokenizer is not None: + annotate_token_counts(stored_messages, tokenizer=self.tokenizer) + await self.after_strategy(stored_messages) + + # Keep all messages (including excluded) in storage so annotations are + # preserved. The history provider's ``skip_excluded`` flag controls + # whether excluded messages are loaded on the next turn. + + +__all__ = [ + "COMPACTION_STATE_KEY", + "EXCLUDED_KEY", + "EXCLUDE_REASON_KEY", + "GROUP_ANNOTATION_KEY", + "GROUP_HAS_REASONING_KEY", + "GROUP_ID_KEY", + "GROUP_INDEX_KEY", + "GROUP_KIND_KEY", + "GROUP_TOKEN_COUNT_KEY", + "SUMMARIZED_BY_SUMMARY_ID_KEY", + "SUMMARY_OF_GROUP_IDS_KEY", + "SUMMARY_OF_MESSAGE_IDS_KEY", + "CharacterEstimatorTokenizer", + "CompactionProvider", + "CompactionStrategy", + "GroupKind", + "SelectiveToolCallCompactionStrategy", + "SlidingWindowStrategy", + "SummarizationStrategy", + "TokenBudgetComposedStrategy", + "TokenizerProtocol", + "ToolResultCompactionStrategy", + "TruncationStrategy", + "annotate_message_groups", + "annotate_token_counts", + "append_compaction_message", + "apply_compaction", + "extend_compaction_messages", + "group_messages", + "included_messages", + "included_token_count", + "project_included_messages", +] diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 7f3f3da13d..ba11355adc 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -37,6 +37,7 @@ if TYPE_CHECKING: from ._agents import SupportsAgentRun from ._clients import SupportsChatGetResponse + from ._compaction import CompactionStrategy, TokenizerProtocol from ._sessions import AgentSession from ._tools import FunctionTool from ._types import ChatOptions, ChatResponse, ChatResponseUpdate @@ -101,6 +102,8 @@ class AgentContext: session: The agent session for this invocation, if any. options: The options for the agent invocation as a dict. stream: Whether this is a streaming invocation. + compaction_strategy: Optional per-run compaction override. + tokenizer: Optional per-run tokenizer override. metadata: Metadata dictionary for sharing data between agent middleware. result: Agent execution result. Can be observed after calling ``call_next()`` to see the actual execution result or can be set to override the execution result. @@ -139,6 +142,8 @@ class AgentContext: session: AgentSession | None = None, options: Mapping[str, Any] | None = None, stream: bool = False, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, metadata: Mapping[str, Any] | None = None, result: AgentResponse | ResponseStream[AgentResponseUpdate, AgentResponse] | None = None, kwargs: Mapping[str, Any] | None = None, @@ -158,6 +163,8 @@ class AgentContext: session: The agent session for this invocation, if any. options: The options for the agent invocation as a dict. stream: Whether this is a streaming invocation. + compaction_strategy: Optional per-run compaction override. + tokenizer: Optional per-run tokenizer override. metadata: Metadata dictionary for sharing data between agent middleware. result: Agent execution result. kwargs: Additional keyword arguments passed to the agent run method. @@ -170,6 +177,8 @@ class AgentContext: self.session = session self.options = options self.stream = stream + self.compaction_strategy = compaction_strategy + self.tokenizer = tokenizer self.metadata: dict[str, Any] = dict(metadata) if metadata is not None else {} self.result = result self.kwargs: dict[str, Any] = dict(kwargs) if kwargs is not None else {} @@ -969,6 +978,8 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]): *, stream: Literal[False] = ..., options: ChatOptions[ResponseModelBoundT], + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[ChatResponse[ResponseModelBoundT]]: ... @@ -979,6 +990,8 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]): *, stream: Literal[False] = ..., options: OptionsCoT | ChatOptions[None] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[ChatResponse[Any]]: ... @@ -989,6 +1002,8 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]): *, stream: Literal[True], options: OptionsCoT | ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ... @@ -998,11 +1013,18 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]): *, stream: bool = False, options: OptionsCoT | ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: """Execute the chat pipeline if middleware is configured.""" super_get_response = super().get_response # type: ignore[misc] + if compaction_strategy is not None: + kwargs["compaction_strategy"] = compaction_strategy + if tokenizer is not None: + kwargs["tokenizer"] = tokenizer + call_middleware = kwargs.pop("middleware", []) middleware = categorize_middleware(call_middleware) kwargs["function_middleware"] = middleware["function"] @@ -1091,6 +1113,8 @@ class AgentMiddlewareLayer: session: AgentSession | None = None, middleware: Sequence[MiddlewareTypes] | None = None, options: ChatOptions[ResponseModelBoundT], + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse[ResponseModelBoundT]]: ... @@ -1103,6 +1127,8 @@ class AgentMiddlewareLayer: session: AgentSession | None = None, middleware: Sequence[MiddlewareTypes] | None = None, options: ChatOptions[None] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse[Any]]: ... @@ -1115,6 +1141,8 @@ class AgentMiddlewareLayer: session: AgentSession | None = None, middleware: Sequence[MiddlewareTypes] | None = None, options: ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ... @@ -1126,6 +1154,8 @@ class AgentMiddlewareLayer: session: AgentSession | None = None, middleware: Sequence[MiddlewareTypes] | None = None, options: ChatOptions[Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + tokenizer: TokenizerProtocol | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: """MiddlewareTypes-enabled unified run method.""" @@ -1150,7 +1180,15 @@ class AgentMiddlewareLayer: # Execute with middleware if available if not pipeline.has_middlewares: - return super().run(messages, stream=stream, session=session, options=options, **combined_kwargs) # type: ignore[misc, no-any-return] + return super().run( # type: ignore[misc, no-any-return] + messages, + stream=stream, + session=session, + options=options, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, + **combined_kwargs, + ) context = AgentContext( agent=self, # type: ignore[arg-type] @@ -1158,6 +1196,8 @@ class AgentMiddlewareLayer: session=session, options=options, stream=stream, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, kwargs=combined_kwargs, ) @@ -1195,6 +1235,8 @@ class AgentMiddlewareLayer: stream=context.stream, session=context.session, options=context.options, + compaction_strategy=context.compaction_strategy, + tokenizer=context.tokenizer, **context.kwargs, ) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 8c3457da26..434a8d1fd4 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -547,6 +547,7 @@ class InMemoryHistoryProvider(BaseHistoryProvider): store_context_messages: bool = False, store_context_from: set[str] | None = None, store_outputs: bool = True, + skip_excluded: bool = False, ) -> None: """Initialize the in-memory history provider. @@ -558,6 +559,11 @@ class InMemoryHistoryProvider(BaseHistoryProvider): store_context_messages: Whether to store context from other providers. store_context_from: If set, only store context from these source_ids. store_outputs: Whether to store response messages. + skip_excluded: When True, ``get_messages`` omits messages whose + ``additional_properties["_excluded"]`` is truthy. This is + useful when a ``CompactionProvider`` marks messages as excluded + in stored history and you want the loaded context to reflect + those exclusions. Defaults to False (load all messages). """ super().__init__( source_id=source_id or self.DEFAULT_SOURCE_ID, @@ -567,6 +573,7 @@ class InMemoryHistoryProvider(BaseHistoryProvider): store_context_from=store_context_from, store_outputs=store_outputs, ) + self.skip_excluded = skip_excluded async def get_messages( self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any @@ -574,7 +581,10 @@ class InMemoryHistoryProvider(BaseHistoryProvider): """Retrieve messages from session state.""" if state is None: return [] - return list(state.get("messages", [])) + messages = list(state.get("messages", [])) + if self.skip_excluded: + messages = [m for m in messages if not m.additional_properties.get("_excluded", False)] + return messages async def save_messages( self, diff --git a/python/packages/core/agent_framework/_skills.py b/python/packages/core/agent_framework/_skills.py index fc71329a5f..c95fc46aa2 100644 --- a/python/packages/core/agent_framework/_skills.py +++ b/python/packages/core/agent_framework/_skills.py @@ -26,13 +26,14 @@ Only use skills from trusted sources. from __future__ import annotations import inspect +import json import logging import os import re from collections.abc import Callable, Sequence from html import escape as xml_escape from pathlib import Path, PurePosixPath -from typing import TYPE_CHECKING, Any, ClassVar, Final +from typing import TYPE_CHECKING, Any, ClassVar, Final, Protocol, runtime_checkable from ._sessions import BaseContextProvider from ._tools import FunctionTool @@ -93,6 +94,7 @@ class SkillResource: description: Optional human-readable summary shown when advertising the resource. content: Static content string. Mutually exclusive with *function*. function: Callable (sync or async) that returns content on demand. + May return any type; the value is passed through as-is. Mutually exclusive with *content*. """ if not name or not name.strip(): @@ -115,6 +117,108 @@ class SkillResource: self._accepts_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()) +class SkillScript: + """An executable script attached to a skill. + + .. warning:: Experimental + + This API is experimental and subject to change or removal + in future versions without notice. + + A script represents executable code that an agent can run. It holds + either an inline ``function`` callable (code-defined scripts) or + a ``path`` to a script file on disk (file-based scripts). + Exactly one must be provided. + + When ``function`` is set the script is treated as **code-based** + and the function is invoked directly in-process. When ``path`` is + set the script is treated as **file-based** and delegated to the + configured :class:`SkillScriptRunner`. + + Attributes: + name: Script identifier. + description: Optional human-readable summary, or ``None``. + function: Callable that implements the script, or ``None``. + path: Relative path to the script file from the skill directory, or + ``None`` for code-defined scripts. + + Examples: + Code-defined script: + + .. code-block:: python + + SkillScript(name="analyze", function=analyze_data, description="Run analysis") + + File-based script (discovered from disk): + + .. code-block:: python + + SkillScript(name="process.py", path="scripts/process.py") + """ + + def __init__( + self, + *, + name: str, + description: str | None = None, + function: Callable[..., Any] | None = None, + path: str | None = None, + ) -> None: + """Initialize a SkillScript. + + Args: + name: Identifier for this script (e.g. ``"analyze"``, ``"process.py"``). + description: Optional human-readable summary. + function: Callable (sync or async) that implements the script. + Set for code-defined scripts; ``None`` for file-based scripts. + Mutually exclusive with *path*. + path: Relative path to the script file from the skill directory. + Set automatically for file-based scripts discovered from disk; + ``None`` for code-defined scripts. + Mutually exclusive with *function*. + """ + if not name or not name.strip(): + raise ValueError("Script name cannot be empty.") + if function is None and path is None: + raise ValueError(f"Script '{name}' must have either function or path.") + if function is not None and path is not None: + raise ValueError(f"Script '{name}' must have either function or path, not both.") + + self.name = name + self.description = description + self.function = function + self.path = path + self._parameters_schema: dict[str, Any] | None = None + self._parameters_schema_resolved: bool = False + + # Precompute whether the function accepts **kwargs to avoid + # repeated inspect.signature() calls on every invocation. + self._accepts_kwargs: bool = False + if function is not None: + sig = inspect.signature(function) + self._accepts_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()) + + @property + def parameters_schema(self) -> dict[str, Any] | None: + """JSON Schema describing the script's parameters. + + .. warning:: Experimental + + This API is experimental and subject to change or removal + in future versions without notice. + + Lazily generated from the callable's signature on first access. + Returns ``None`` for file-based scripts or functions with no + introspectable parameters. + """ + if not self._parameters_schema_resolved and self.function is not None: + tool = FunctionTool(name=self.function.__name__, func=self.function) + schema = tool.parameters() + self._parameters_schema = schema if schema and schema.get("properties") else None + self._parameters_schema_resolved = True + return self._parameters_schema + + class Skill: """A skill definition with optional resources. @@ -124,15 +228,16 @@ class Skill: in future versions without notice. A skill bundles a set of instructions (``content``) with metadata and - zero or more :class:`SkillResource` instances. Resources can be - supplied at construction time or added later via the :meth:`resource` - decorator. + zero or more :class:`SkillResource` and :class:`SkillScript` instances. + Resources and scripts can be supplied at construction time or added later + via the :meth:`resource` and :meth:`script` decorators. Attributes: name: Skill name (lowercase letters, numbers, hyphens only). description: Human-readable description of the skill. content: The skill instructions body. resources: Mutable list of :class:`SkillResource` instances. + scripts: Mutable list of :class:`SkillScript` instances. path: Absolute path to the skill directory on disk, or ``None`` for code-defined skills. @@ -171,6 +276,7 @@ class Skill: description: str, content: str, resources: list[SkillResource] | None = None, + scripts: list[SkillScript] | None = None, path: str | None = None, ) -> None: """Initialize a Skill. @@ -180,6 +286,7 @@ class Skill: description: Human-readable description of the skill (≤1024 chars). content: The skill instructions body. resources: Pre-built resources to attach to this skill. + scripts: Pre-built scripts to attach to this skill. path: Absolute path to the skill directory on disk. Set automatically for file-based skills; leave as ``None`` for code-defined skills. """ @@ -192,6 +299,7 @@ class Skill: self.description = description self.content = content self.resources: list[SkillResource] = resources if resources is not None else [] + self.scripts: list[SkillScript] = scripts if scripts is not None else [] self.path = path def resource( @@ -227,7 +335,7 @@ class Skill: .. code-block:: python @skill.resource - def get_schema() -> str: + def get_schema() -> Any: return "schema..." With arguments: @@ -235,7 +343,7 @@ class Skill: .. code-block:: python @skill.resource(name="custom-name", description="Custom desc") - async def get_data() -> str: + async def get_data() -> Any: return "data..." """ @@ -255,10 +363,116 @@ class Skill: return decorator return decorator(func) + def script( + self, + func: Callable[..., Any] | None = None, + *, + name: str | None = None, + description: str | None = None, + ) -> Any: + """Decorator that registers a callable as a script on this skill. + + Supports bare usage (``@skill.script``) and parameterized usage + (``@skill.script(name="custom", description="...")``). The + decorated function is returned unchanged; a new + :class:`SkillScript` is appended to :attr:`scripts`. + + Args: + func: The function being decorated. Populated automatically when + the decorator is applied without parentheses. + + Keyword Args: + name: Script name override. Defaults to ``func.__name__``. + description: Script description override. Defaults to the + function's docstring (via :func:`inspect.getdoc`). + + Returns: + The original function unchanged, or a secondary decorator when + called with keyword arguments. + + Examples: + Bare decorator: + + .. code-block:: python + + @skill.script + def analyze_data(query: str) -> str: + \"\"\"Run data analysis.\"\"\" + return run_analysis(query) + + With arguments: + + .. code-block:: python + + @skill.script(name="fetch", description="Fetch remote data") + async def fetch_data(url: str) -> str: + return await http_get(url) + """ + + def decorator(f: Callable[..., Any]) -> Callable[..., Any]: + script_name = name or f.__name__ + script_description = description or (inspect.getdoc(f) or None) + self.scripts.append( + SkillScript( + name=script_name, + description=script_description, + function=f, + ) + ) + return f + + if func is None: + return decorator + return decorator(func) + # endregion -# region Constants +# region Script Runners + + +@runtime_checkable +class SkillScriptRunner(Protocol): + """Protocol for skill script runners. + + .. warning:: Experimental + + This API is experimental and subject to change or removal + in future versions without notice. + + A script runner determines how **file-based** skill scripts are + run. Implementations decide the execution strategy + (e.g., local subprocess, hosted code execution environment, + user-provided callable). + + Code-defined scripts (registered via the ``@skill.script`` decorator) + are always executed **in-process** and do not use a script runner. + + Any callable (sync or async) matching the ``__call__`` signature + satisfies this protocol. + """ + + def __call__(self, skill: Skill, script: SkillScript, args: dict[str, Any] | None = None) -> Any: + """Run a skill script. + + The :class:`SkillsProvider` resolves skill and script names + before calling this method, so implementations receive fully + resolved objects. + + Args: + skill: The skill that owns the script. + script: The script to run. + args: Optional keyword arguments for the script. + + Returns: + The result. May be any type; the framework + serialises it automatically via + :meth:`~FunctionTool.parse_result`. + """ + ... + + +# endregion SKILL_FILE_NAME: Final[str] = "SKILL.md" MAX_SEARCH_DEPTH: Final[int] = 2 @@ -273,8 +487,7 @@ DEFAULT_RESOURCE_EXTENSIONS: Final[tuple[str, ...]] = ( ".xml", ".txt", ) - -# endregion +DEFAULT_SCRIPT_EXTENSIONS: Final[tuple[str, ...]] = (".py",) # region Patterns and prompt template @@ -307,13 +520,19 @@ Each skill provides specialized instructions, reference documents, and assets fo When a task aligns with a skill's domain, follow these steps in exact order: -1. Use `load_skill` to retrieve the skill's instructions. -2. Follow the provided guidance. -3. Use `read_skill_resource` to read any referenced resources, using the name exactly as listed +- Use `load_skill` to retrieve the skill's instructions. +- Follow the provided guidance. +- Use `read_skill_resource` to read any referenced resources, using the name exactly as listed (e.g. `"style-guide"` not `"style-guide.md"`, `"references/FAQ.md"` not `"FAQ.md"`). - +{runner_instructions} Only load what is needed, when it is needed.""" +SCRIPT_RUNNER_INSTRUCTIONS: Final[str] = ( + "\n- Use `run_skill_script` to run referenced scripts, using the name exactly as listed." + "\n- Pass script arguments inside `args` as a JSON object" + ' (e.g. `args: {"length": 24}`), not as top-level tool parameters.\n' +) + # endregion # region SkillsProvider @@ -381,8 +600,11 @@ class SkillsProvider(BaseContextProvider): skill_paths: str | Path | Sequence[str | Path] | None = None, *, skills: Sequence[Skill] | None = None, + script_runner: SkillScriptRunner | None = None, instruction_template: str | None = None, resource_extensions: tuple[str, ...] | None = None, + script_extensions: tuple[str, ...] | None = None, + require_script_approval: bool = False, source_id: str | None = None, ) -> None: """Initialize a SkillsProvider. @@ -395,21 +617,69 @@ class SkillsProvider(BaseContextProvider): Keyword Args: skills: Code-defined :class:`Skill` instances to register. + script_runner: Strategy for running **file-based** skill + scripts. The provider resolves skill and script names, then + calls the runner directly. This parameter only + affects scripts discovered from disk (via *skill_paths*); + code-defined scripts (registered with ``@skill.script``) are + always executed in-process and ignore this setting. + When ``None``, file-based scripts are not executable. instruction_template: Custom system-prompt template for advertising skills. Must contain a ``{skills}`` placeholder for the generated skills list. Uses a built-in template when ``None``. resource_extensions: File extensions recognized as discoverable resources. Defaults to ``DEFAULT_RESOURCE_EXTENSIONS`` (``(".md", ".json", ".yaml", ".yml", ".csv", ".xml", ".txt")``). + script_extensions: File extensions recognized as discoverable + scripts. Defaults to ``DEFAULT_SCRIPT_EXTENSIONS`` + (``(".py",)``). + require_script_approval: When ``True``, skill script execution + requires explicit user approval before running. Instead of + executing immediately, the agent pauses and returns a + ``function_approval_request`` via ``result.user_input_requests``. + The application should present the request to the user, then + call ``request.to_function_approval_response(approved=True)`` + (or ``False`` to reject) and pass the response back with + ``agent.run(approval_response, session=session)``. + Rejected scripts are not executed and the agent is informed + the user declined. Defaults to ``False``. See + ``samples/02-agents/skills/script_approval/script_approval.py`` + for the full approval loop pattern. source_id: Unique identifier for this provider instance. """ super().__init__(source_id or self.DEFAULT_SOURCE_ID) - self._skills = _load_skills(skill_paths, skills, resource_extensions or DEFAULT_RESOURCE_EXTENSIONS) + self._skills = _load_skills( + skill_paths, + skills, + resource_extensions or DEFAULT_RESOURCE_EXTENSIONS, + script_extensions or DEFAULT_SCRIPT_EXTENSIONS, + ) - self._instructions = _create_instructions(instruction_template, self._skills) + # File-based skills (skill.path set) have scripts discovered from disk + has_file_scripts = any(s.scripts for s in self._skills.values() if s.path is not None) - self._tools = self._create_tools() + # Code-defined skills (skill.path is None) have scripts with callable functions + has_code_scripts = any(s.scripts for s in self._skills.values() if s.path is None) + + if has_file_scripts and script_runner is None: + raise ValueError( + "File-based skills with scripts were provided but no 'script_runner' was provided. " + "Pass a SkillScriptRunner callable to SkillsProvider." + ) + + self._script_runner = script_runner + + self._instructions = _create_instructions( + prompt_template=instruction_template, + skills=self._skills, + include_script_runner_instructions=has_file_scripts or has_code_scripts, + ) + + self._tools = self._create_tools( + include_script_runner_tool=has_file_scripts or has_code_scripts, + require_script_approval=require_script_approval, + ) async def before_run( self, @@ -425,6 +695,11 @@ class SkillsProvider(BaseContextProvider): skill is registered, appends the skill-list system prompt and the ``load_skill`` / ``read_skill_resource`` tools to *context*. + When any registered skill defines one or more scripts (file-based or + code-based), the system prompt also includes script-runner + instructions (embedded via the ``{runner_instructions}`` placeholder), + and the ``run_skill_script`` tool is included alongside the base tools. + Args: agent: The agent instance about to run. session: The current agent session. @@ -434,17 +709,30 @@ class SkillsProvider(BaseContextProvider): if not self._skills: return - if self._instructions: - context.extend_instructions(self.source_id, self._instructions) + context.extend_instructions(self.source_id, self._instructions) # type: ignore[arg-type] context.extend_tools(self.source_id, self._tools) - def _create_tools(self) -> list[FunctionTool]: + def _create_tools( + self, + include_script_runner_tool: bool, + require_script_approval: bool = False, + ) -> list[FunctionTool]: """Create the ``load_skill`` and ``read_skill_resource`` tool definitions. + When *include_script_runner_tool* is ``True``, also creates + ``run_skill_script``. + + Args: + include_script_runner_tool: Whether to include the + ``run_skill_script`` tool in the returned list. + require_script_approval: When ``True``, the + ``run_skill_script`` tool pauses for user approval + before each invocation. + Returns: - A two-element list of :class:`FunctionTool` instances. + A list of :class:`FunctionTool` instances. """ - return [ + tools = [ FunctionTool( name="load_skill", description="Loads the full instructions for a specific skill.", @@ -475,6 +763,45 @@ class SkillsProvider(BaseContextProvider): ), ] + if include_script_runner_tool: + tools.append( + FunctionTool( + name="run_skill_script", + description="Runs a script associated with a skill.", + func=self._run_skill_script, + approval_mode="always_require" if require_script_approval else "never_require", + input_model={ + "type": "object", + "properties": { + "skill_name": {"type": "string", "description": "The name of the skill."}, + "script_name": { + "type": "string", + "description": ( + "The name of the script to run as listed in the skill, " + "preserving any directory prefix exactly as shown. " + "Do not add or remove path prefixes." + ), + }, + "args": { + "type": ["object", "null"], + "additionalProperties": True, + "default": None, + "description": ( + "Arguments to pass to the script as key-value pairs. " + "Use parameter names as keys without leading dashes " + '(e.g. {"length": 24, "uppercase": true}). ' + "How these values are mapped to the underlying script " + "is determined by the script implementation or configured runner." + ), + }, + }, + "required": ["skill_name", "script_name"], + }, + ) + ) + + return tools + def _load_skill(self, skill_name: str) -> str: """Return the full instructions for the named skill. @@ -516,9 +843,79 @@ class SkillsProvider(BaseContextProvider): resource_lines = "\n".join(_create_resource_element(r) for r in skill.resources) content += f"\n\n\n{resource_lines}\n" + if skill.scripts: + script_lines = "\n".join(_create_script_element(s) for s in skill.scripts) + content += f"\n\n\n{script_lines}\n" + return content - async def _read_skill_resource(self, skill_name: str, resource_name: str, **kwargs: Any) -> str: + async def _run_skill_script( + self, skill_name: str, script_name: str, args: dict[str, Any] | None = None, **kwargs: Any + ) -> Any: + """Run a named script from a skill. + + For code-defined scripts (those with a ``function`` and no ``path``), + the function is invoked directly in-process. For file-based scripts + the configured :class:`SkillScriptRunner` is used. + + Args: + skill_name: The name of the owning skill. + script_name: The script name to look up (case-insensitive). + args: Optional keyword arguments for the script, provided by the + agent/LLM. These are mapped to the function's declared + parameters. + **kwargs: Runtime keyword arguments forwarded only to script + functions that accept ``**kwargs`` (e.g. arguments passed via + ``agent.run(user_id="123")``). + + Returns: + The result, or a user-facing error message on + failure. + """ + if not skill_name or not skill_name.strip(): + return "Error: Skill name cannot be empty." + + if not script_name or not script_name.strip(): + return "Error: Script name cannot be empty." + + skill = self._skills.get(skill_name) + if not skill: + return f"Error: Skill '{skill_name}' not found." + + script = next((s for s in skill.scripts if s.name.lower() == script_name.lower()), None) + if not script: + return f"Error: Script '{script_name}' not found in skill '{skill_name}'." + + # Code-defined scripts: run the function directly + if script.function is not None: + try: + if script._accepts_kwargs: # pyright: ignore[reportPrivateUsage] + result = script.function(**(args or {}), **kwargs) + else: + result = script.function(**(args or {})) + if inspect.isawaitable(result): + result = await result + return result + except Exception: + logger.exception("Error running code-defined script '%s' in skill '%s'", script_name, skill_name) + return f"Error: Failed to run script '{script_name}' in skill '{skill_name}'." + + # File-based scripts: delegate to the runner + if self._script_runner is None: + return ( + f"Error: Script '{script_name}' in skill '{skill_name}' requires a runner. " + "Provide a script_runner for file-based scripts." + ) + try: + result = self._script_runner(skill, script, args) + if inspect.isawaitable(result): + result = await result + return result + except Exception: + logger.exception("Error running file-based script '%s' in skill '%s'", script_name, skill_name) + return f"Error: Failed to run script '{script_name}' in skill '{skill_name}'." + + async def _read_skill_resource(self, skill_name: str, resource_name: str, **kwargs: Any) -> Any: """Read a named resource from a skill. Resolves the resource by case-insensitive name lookup. Static @@ -533,7 +930,7 @@ class SkillsProvider(BaseContextProvider): ``agent.run(user_id="123")``). Returns: - The resource content string, or a user-facing error message on + The resource content (any type), or a user-facing error message on failure. """ if not skill_name or not skill_name.strip(): @@ -565,13 +962,10 @@ class SkillsProvider(BaseContextProvider): ) else: result = resource.function(**kwargs) if resource._accepts_kwargs else resource.function() # pyright: ignore[reportPrivateUsage] - return str(result) - except Exception as exc: + return result + except Exception: logger.exception("Failed to read resource '%s' from skill '%s'", resource_name, skill_name) - return ( - f"Error ({type(exc).__name__}): Failed to read resource" - f" '{resource_name}' from skill '{skill_name}'." - ) + return f"Error: Failed to read resource '{resource_name}' from skill '{skill_name}'." return f"Error: Resource '{resource.name}' has no content or function." @@ -707,6 +1101,60 @@ def _discover_resource_files( return resources +def _discover_script_files( + skill_dir_path: str, + extensions: tuple[str, ...] = DEFAULT_SCRIPT_EXTENSIONS, +) -> list[str]: + """Scan a skill directory for script files matching *extensions*. + + Recursively walks *skill_dir_path* and collects files whose extension + is in *extensions*. Each candidate is validated against path-traversal + and symlink-escape checks; unsafe files are skipped with a warning. + + Args: + skill_dir_path: Absolute path to the skill directory to scan. + extensions: Tuple of allowed script extensions (e.g. ``(".py",)``). + + Returns: + Relative script paths (forward-slash-separated) for every + discovered file that passes security checks. + """ + skill_dir = Path(skill_dir_path).absolute() + root_directory_path = str(skill_dir) + scripts: list[str] = [] + normalized_extensions = {e.lower() for e in extensions} + + for script_file in skill_dir.rglob("*"): + if not script_file.is_file(): + continue + + if script_file.suffix.lower() not in normalized_extensions: + continue + + script_full_path = str(Path(os.path.normpath(script_file)).absolute()) + + if not _is_path_within_directory(script_full_path, root_directory_path): + logger.warning( + "Skipping script '%s': resolves outside skill directory '%s'", + script_file, + skill_dir_path, + ) + continue + + if _has_symlink_in_path(script_full_path, root_directory_path): + logger.warning( + "Skipping script '%s': symlink detected in path under skill directory '%s'", + script_file, + skill_dir_path, + ) + continue + + rel_path = script_file.relative_to(skill_dir) + scripts.append(_normalize_resource_path(str(rel_path))) + + return scripts + + def _validate_skill_metadata( name: str | None, description: str | None, @@ -902,6 +1350,7 @@ def _read_file_skill_resource(skill: Skill, resource_name: str) -> str: def _discover_file_skills( skill_paths: str | Path | Sequence[str | Path] | None, resource_extensions: tuple[str, ...] = DEFAULT_RESOURCE_EXTENSIONS, + script_extensions: tuple[str, ...] = DEFAULT_SCRIPT_EXTENSIONS, ) -> dict[str, Skill]: """Discover, parse, and load all file-based skills from the given paths. @@ -912,6 +1361,7 @@ def _discover_file_skills( Args: skill_paths: Directory path(s) to scan, or ``None`` to skip. resource_extensions: File extensions recognized as resources. + script_extensions: File extensions recognized as scripts. Returns: A dict mapping skill name → :class:`Skill`. @@ -955,6 +1405,10 @@ def _discover_file_skills( reader = (lambda s, r: lambda: _read_file_skill_resource(s, r))(file_skill, rn) file_skill.resources.append(SkillResource(name=rn, function=reader)) + # Discover and attach file-based scripts as SkillScript instances + for sn in _discover_script_files(skill_path, script_extensions): + file_skill.scripts.append(SkillScript(name=sn, path=sn)) + skills[file_skill.name] = file_skill logger.info("Loaded skill: %s", file_skill.name) @@ -966,6 +1420,7 @@ def _load_skills( skill_paths: str | Path | Sequence[str | Path] | None, skills: Sequence[Skill] | None, resource_extensions: tuple[str, ...], + script_extensions: tuple[str, ...], ) -> dict[str, Skill]: """Discover and merge skills from file paths and code-defined skills. @@ -977,11 +1432,12 @@ def _load_skills( skill_paths: Directory path(s) to scan for ``SKILL.md`` files, or ``None``. skills: Code-defined :class:`Skill` instances, or ``None``. resource_extensions: File extensions recognized as discoverable resources. + script_extensions: File extensions recognized as discoverable scripts. Returns: A dict mapping skill name → :class:`Skill`. """ - result = _discover_file_skills(skill_paths, resource_extensions) + result = _discover_file_skills(skill_paths, resource_extensions, script_extensions) if skills: for code_skill in skills: @@ -1017,19 +1473,50 @@ def _create_resource_element(resource: SkillResource) -> str: return f" " +def _create_script_element(script: SkillScript) -> str: + """Create an XML ``" + return f"