chore: import upstream snapshot with attribution
CodeQL / Analyze (csharp) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
dotnet-build-and-test / dotnet-test-functions (push) Has been cancelled
dotnet-build-and-test / paths-filter (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Debug, windows-latest, net9.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net8.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-foundry-hosted-it (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test-check (push) Has been cancelled
dotnet-build-and-test / Integration Test Report (push) Has been cancelled
CodeQL / Analyze (csharp) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
dotnet-build-and-test / dotnet-test-functions (push) Has been cancelled
dotnet-build-and-test / paths-filter (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Debug, windows-latest, net9.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net8.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-foundry-hosted-it (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test-check (push) Has been cancelled
dotnet-build-and-test / Integration Test Report (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
+15
@@ -0,0 +1,15 @@
|
||||
{
|
||||
// Use IntelliSense to learn about possible attributes.
|
||||
// Hover to view descriptions of existing attributes.
|
||||
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
|
||||
"version": "0.2.0",
|
||||
"configurations": [
|
||||
{
|
||||
"name": "Python Debugger: Current File",
|
||||
"type": "debugpy",
|
||||
"request": "launch",
|
||||
"program": "${file}",
|
||||
"console": "integratedTerminal"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,233 @@
|
||||
# Core Package (agent-framework-core)
|
||||
|
||||
The foundation package containing all core abstractions, types, and built-in OpenAI/Azure OpenAI support.
|
||||
|
||||
## Module Structure
|
||||
|
||||
```
|
||||
agent_framework/
|
||||
├── __init__.py # Lazy runtime public API exports
|
||||
├── __init__.pyi # Public API typing surface for lazy root exports
|
||||
├── security.py # Public security primitives, middleware, and tools
|
||||
├── _agents.py # Agent implementations
|
||||
├── _clients.py # Chat client base classes and protocols
|
||||
├── _types.py # Core types (Message, ChatResponse, Content, etc.)
|
||||
├── _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
|
||||
├── azure/ # Lazy-loading entry point for Azure integrations
|
||||
└── <provider>/ # Other lazy-loading provider folders
|
||||
```
|
||||
|
||||
## Core Classes
|
||||
|
||||
### Root Public API (`__init__.py` / `__init__.pyi`)
|
||||
|
||||
- `agent_framework.__init__` uses lazy module-level `__getattr__` for most public exports to keep cold
|
||||
`import agent_framework` lightweight.
|
||||
- Keep `_LAZY_MODULE_EXPORTS`, `_LAZY_EXPORTS`, the explicit runtime `__all__`, and `__init__.pyi` synchronized
|
||||
whenever adding, removing, or moving a root public export.
|
||||
- Runtime `__all__` is still required for `from agent_framework import *`; the `.pyi` file is for type checkers
|
||||
and editors and does not replace runtime exports.
|
||||
- Public deprecation behavior for a lazy export belongs in the owning module. The root package should delegate via
|
||||
the normal lazy export map instead of carrying one-off branches.
|
||||
|
||||
### Agents (`_agents.py`)
|
||||
|
||||
- **`SupportsAgentRun`** - Protocol defining the agent interface
|
||||
- **`BaseAgent`** - Abstract base class for agents
|
||||
- **`Agent`** - Main agent class wrapping a chat client with tools, instructions, and middleware
|
||||
|
||||
### Chat Clients (`_clients.py`)
|
||||
|
||||
- **`SupportsChatGetResponse`** - Protocol for chat client implementations
|
||||
- **`BaseChatClient`** - Abstract base class with middleware support; subclasses implement `_inner_get_response()` and `_inner_get_streaming_response()`
|
||||
|
||||
### Types (`_types.py`)
|
||||
|
||||
- **`Message`** - Represents a chat message with role, content, and metadata
|
||||
- **`ChatResponse`** - Response from a chat client containing messages and usage
|
||||
- **`ChatResponseUpdate`** - Streaming response update
|
||||
- **`AgentResponse`** / **`AgentResponseUpdate`** - Agent-level response wrappers
|
||||
- **`Content`** - Base class for message content (text, function calls, images, etc.)
|
||||
- **`ChatOptions`** - TypedDict for chat request options
|
||||
|
||||
### Tools (`_tools.py`)
|
||||
|
||||
- **`ToolProtocol`** - Protocol for tool definitions
|
||||
- **`FunctionTool`** - Wraps Python functions as tools with JSON schema generation
|
||||
- **`@tool`** decorator - Converts functions to tools
|
||||
- **`use_function_invocation()`** - Decorator to add automatic function calling to chat clients
|
||||
|
||||
### Middleware (`_middleware.py`)
|
||||
|
||||
- **`AgentMiddleware`** - Intercepts agent `run()` calls
|
||||
- **`ChatMiddleware`** - Intercepts chat client `get_response()` calls
|
||||
- **`FunctionMiddleware`** - Intercepts function/tool invocations
|
||||
- **`AgentContext`** / **`ChatContext`** / **`FunctionInvocationContext`** - Context objects passed through middleware. A tool can declare a `FunctionInvocationContext` parameter to receive it; `context.tools` is the live, mutable tools list for the run, and `context.add_tools(...)` / `context.remove_tools(...)` enable progressive tool exposure (changes apply on the next function-calling iteration).
|
||||
- **`MessageInjectionMiddleware`** - Session-scoped chat middleware that lets tools or other code enqueue messages for the next model call in the current `AgentSession`; it drains queued messages into the next call and loops only when no function calls need to be handled by the function invocation layer.
|
||||
|
||||
### Sessions (`_sessions.py`)
|
||||
|
||||
- **`AgentSession`** - Manages conversation state and session metadata
|
||||
- **`ServiceSessionId`** - Mapping alias for structured service-owned continuation handles used in `AgentSession.service_session_id`
|
||||
- **`SessionContext`** - Context object for session-scoped data during agent runs
|
||||
- **`ContextProvider`** - Base class for context providers (RAG, memory systems)
|
||||
- **`HistoryProvider`** - Base class for conversation history storage
|
||||
- **`InMemoryHistoryProvider`** - Built-in session-state history provider for local runs
|
||||
- **`FileHistoryProvider`** - JSON Lines file-backed history provider storing one file per session with one message record per line
|
||||
|
||||
### Skills (`_skills.py`)
|
||||
|
||||
- **`Skill`** - Abstract base for a skill definition bundling instructions (`content`) with frontmatter metadata, resources, and scripts. Concrete subclasses (`InlineSkill`, `FileSkill`, `ClassSkill`) accept a `frontmatter=SkillFrontmatter(...)` argument carrying the spec fields. Adding new spec fields is done in one place — on `SkillFrontmatter` — keeping the subclass constructors stable.
|
||||
- **`SkillFrontmatter`** - L1 discovery metadata for a skill (`name`, `description`, `license`, `compatibility`, `allowed_tools`, `metadata`). All fields are mutable plain attributes; the constructor validates `name`, `description`, and `compatibility` against the spec but post-construction assignments are not re-validated. Spec fields are reachable on every skill via `skill.frontmatter`.
|
||||
- **`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.
|
||||
- **`SkillScriptArgumentParser`** - Public type alias for an optional callable `(raw args: dict | list[str] | str | None) -> dict | None` that converts the raw `args` value before an `InlineSkillScript` runs (applied before the inline list-args guard). It is an opt-in customization hook (port of .NET PR #6498) that lets callers support backends sending tool-call arguments in a non-conforming shape (e.g. vLLM JSON strings). The output is constrained to a `dict` (named keyword arguments) or `None`, because inline scripts bind arguments by keyword name. Supply it via the `argument_parser=` constructor arg on `InlineSkillScript`, `InlineSkill` (default for scripts added via `@skill.script`), or `ClassSkill` (default for scripts discovered via `@ClassSkill.script`). When `None` (the default), the raw value is used unchanged. File-based scripts are unaffected (their runner owns arg handling).
|
||||
- **`SkillsProvider`** - Context provider (extends `ContextProvider`) 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. By default all three tools it exposes (`load_skill`, `read_skill_resource`, `run_skill_script`) are registered with `approval_mode="always_require"`, so every skill operation needs approval. To run unattended, pass one of the static auto-approval rules to `ToolApprovalMiddleware` (via `auto_approval_rules`): `SkillsProvider.read_only_tools_auto_approval_rule` approves only the read-only tools (`load_skill`, `read_skill_resource`) while still prompting for `run_skill_script`, and `SkillsProvider.all_tools_auto_approval_rule` approves every skill tool including script execution. Both rules reject any call carrying a `server_label` so they stay scoped to this provider's local tools and never auto-approve a same-named hosted tool. Alternatively, for trusted skills, the constructor / `from_paths` kwargs `disable_load_skill_approval`, `disable_read_skill_resource_approval`, and `disable_run_skill_script_approval` (all default `False`) opt individual tools out of approval entirely by registering them with `approval_mode="never_require"` (the auto-approval rules only apply to tools that still require approval). The tool names are also exposed as class constants (`LOAD_SKILL_TOOL_NAME`, `READ_SKILL_RESOURCE_TOOL_NAME`, `RUN_SKILL_SCRIPT_TOOL_NAME`).
|
||||
- **`SkillsSource` decorators** - Skill sources are composable: `SkillsSource` is the abstract base, with concrete sources (`InMemorySkillsSource`, `FileSkillsSource`, `MCPSkillsSource`) and decorators that wrap an inner source — `AggregatingSkillsSource` (concatenate several sources), `FilteringSkillsSource` (predicate filter), `DeduplicatingSkillsSource` (first-wins by name), and `CachingSkillsSource` (cache the inner source's skills list). `DelegatingSkillsSource` is the abstract base for decorators. **`get_skills` takes a `SkillsSourceContext`**: every source/decorator implements `async def get_skills(self, context: SkillsSourceContext) -> list[Skill]` and forwards `context` to inner sources. `SkillsSourceContext` (frozen) carries the invoking `agent` (`SupportsAgentRun`) and optional `session` (`AgentSession | None`); `SkillsProvider` builds it from `before_run`'s `agent`/`session` and passes it into the pipeline. `FilteringSkillsSource`'s predicate is context-aware: `Callable[[Skill, SkillsSourceContext], bool]` (port of .NET #6797). **Default caching is applied only to the built-in, context-independent leaf sources**: for the `Skill` / sequence-of-skills / `from_paths` constructors, `SkillsProvider` builds `DeduplicatingSkillsSource(CachingSkillsSource(<file|in-memory leaf>))` so expensive filesystem/network discovery runs once. A **caller-supplied `SkillsSource` is used as-is — never auto-wrapped in caching or deduplication** — because auto-caching a context-aware caller source in a single shared bucket would replay the first invocation's skills for later `SkillsSourceContext`s and leak skills across agents/tenants (matches .NET, whose custom-source constructor also adds no caching/dedup). Callers who want caching on a custom pipeline compose `CachingSkillsSource(inner, cache_isolation_key_selector=...)` themselves. `disable_caching=True` only affects the built-in leaf caching (it has no effect on a caller-supplied source, which is never cached). `CachingSkillsSource` shares a single in-flight fetch across concurrent callers (per cache key) and does not update its cache on a failed fetch, so the next call retries (an initial failure leaves the cache empty; a refresh failure keeps the previously cached list). By default all callers share one cache bucket; pass `cache_isolation_key_selector=Callable[[SkillsSourceContext], str | None]` to cache separately per key (e.g. per agent name) for context-aware inner sources — the key should be low-cardinality and stable, and returning `None` (or leaving the selector `None`) uses the shared bucket. By default a cached list never expires; pass `refresh_interval=timedelta(...)` (port of .NET `CachingAgentSkillsSourceOptions.RefreshInterval`) to treat a cached list as stale once it is older than the interval so the next call re-queries the inner source (useful when an inner source such as `MCPSkillsSource` changes over the process lifetime; a zero/negative interval makes every result immediately stale, and a failed refresh keeps the prior list and retries). Freshness is measured with a monotonic clock (`time.monotonic()`). `SkillsProvider.__init__` / `from_paths` expose a `cache_refresh_interval` kwarg that is threaded into the built-in `CachingSkillsSource` (it has no effect on a caller-supplied source or when `disable_caching=True`).
|
||||
|
||||
### Model Context Protocol (`_mcp.py`)
|
||||
|
||||
- **`MCPTool`** - Base wrapper that owns the MCP `ClientSession` and exposes the remote server's tools as `FunctionTool`s.
|
||||
- **`MCPStdioTool`** / **`MCPStreamableHTTPTool`** / **`MCPWebsocketTool`** - Transport-specific subclasses.
|
||||
- **Argument allowlist (`_prepare_call_kwargs`)** - Before each `tools/call`, kwargs are filtered to an **allowlist** built from the tool's declared parameters (`inputSchema.properties`) plus any user-configured extras. Framework runtime kwargs injected through the function-invocation pipeline (e.g. `thread`, `conversation_id`, `chat_options`, `options`, `response_format`) are stripped by default rather than forwarded. A tool that declares no usable `properties` (including schemas with `additionalProperties: true`) forwards only the configured extras. The `_MCP_FRAMEWORK_DENYLIST` is a safety net for framework-named params a server *declares* in its schema (those are dropped); names explicitly opted in via `additional_tool_argument_names` always win. The reserved `_meta` key is never forwarded as an argument; trusted caller/runtime `_meta` is validated as MCP request metadata, model-supplied `_meta` is discarded in generated MCP functions, and metadata precedence is caller/runtime < OpenTelemetry < tools/list metadata.
|
||||
- **`allowed_tools`** (constructor arg on all `MCPTool` subclasses) - Restricts exposed MCP tools by raw remote MCP tool identity. Prefixed local names remain accepted only when the raw remote name already matches its normalized form; normalized/local aliases do not authorize a different raw remote name. If multiple raw remote tool names map to the same local function name, tool loading raises `ToolExecutionException` instead of first-one-wins shadowing.
|
||||
- **Progressive MCP disclosure** (`use_progressive_disclosure`, `always_load`) - When enabled on any `MCPTool` subclass, the initial model-facing surface is loader tools (`list_mcp_tools` / `load_tool` / `unload_tool`, prefixed by `tool_name_prefix` when configured) plus allowed tools selected by `always_load` and tools loaded earlier on the same `MCPTool` instance. `list_mcp_tools` only reports tools that pass `allowed_tools`; filtered tools are not listed or loadable. Loader tool names are reserved in progressive mode: remote MCP tools whose local generated name collides with a loader name are omitted from the initial/listed surface, and explicit `load_tool` calls return a model-visible message pointing callers to `tool_name_prefix` or excluding the colliding tool. `load_tool` accepts one tool name or a list of tool names and uses `FunctionInvocationContext.add_tools(...)` so the selected generated MCP `FunctionTool`s become available on the next function-calling iteration while keeping existing approval mode, argument filtering, header-provider runtime kwargs, result parsing, OTel, and task behavior. `unload_tool` accepts one dynamically loaded tool name or a list of names and removes them from the live tool list and persisted progressive surface, but it does not remove tools configured in `always_load`. Invalid `always_load` entries are ignored like unmatched `allowed_tools` entries.
|
||||
- **`additional_tool_argument_names`** (constructor arg on all `MCPTool` subclasses) - Opt extra argument names back into the allowlist. Accepts a `Sequence[str]` (applied to every tool) or a `Mapping[str, Sequence[str]]` keyed by **remote tool name**, where the reserved key `"*"` denotes global extras. It is configured only in user code at construction; there is **no per-call/runtime override**, so a model-issued tool call cannot change which names pass through. To use a server that accepts `additionalProperties: true`, list the extra names here and then either (1) manually extend that tool's `inputSchema` (via the `.functions` list after connecting) so the model is prompted to supply them, or (2) supply the values yourself via `function_invocation_kwargs`. If a normal forwarded argument name is supplied by both the model and `function_invocation_kwargs`, the model-supplied value wins; `_meta` is the exception and only trusted runtime/caller metadata is used.
|
||||
- **Sampling guardrails** (`sampling_callback`) - Passing `client=` advertises `SamplingCapability` so the server can send `sampling/createMessage`. Because remote servers are untrusted (confused-deputy risk), the default `sampling_callback` is **deny-by-default** and applies, in order: a per-session rate limit (`sampling_max_requests`, default `_DEFAULT_SAMPLING_MAX_REQUESTS`), an approval gate (`sampling_approval_callback`), and a `maxTokens` cap (`sampling_max_tokens`, default `_DEFAULT_SAMPLING_MAX_TOKENS`). The approval callback (constructor arg on all subclasses; exported type alias `SamplingApprovalCallback`) receives the raw `CreateMessageRequestParams`, may be sync or async, and must return truthy to approve. When it is `None` (the default) every sampling request is denied; pass `lambda params: True` to restore legacy auto-approve as an explicit opt-in. Requests and denials are logged at WARNING (content is not logged). The per-session counter resets in `_reset_session_state`.
|
||||
- **`MCPTaskOptions`** (experimental, `MCP_LONG_RUNNING_TASKS` feature, **frozen**) - Per-tool-instance options controlling the SEP-2663 long-running task lifecycle. When the server advertises a tool with `execution.taskSupport == "required"`, `MCPTool.call_tool` transparently routes through `call_tool_as_task`, which sends an augmented `tools/call`, polls `tasks/get` until terminal, and reinterprets `tasks/result` as a normal `CallToolResult`. Instances are immutable; replace via `MCPTool.task_options = MCPTaskOptions(...)`. Fields:
|
||||
- `default_ttl: timedelta | None` — forwarded to the server as `params.task.ttl` (milliseconds). When `None`, the server's default applies.
|
||||
- `cancel_remote_task_on_local_cancellation: bool = True` — only gates the `CancelledError` path. Abandonment paths (see below) always cancel.
|
||||
- `max_task_wait: timedelta | None` — client-side deadline for the whole post-create lifecycle (poll + result fetch). When exceeded, raises `ToolExecutionException` and fires a best-effort `tasks/cancel`. `None` (default) means no client-side bound. Bounds sleeps, sends, AND reconnects via `asyncio.wait_for`.
|
||||
- **Permissive fallback**: servers that ignore the augmentation (return `CallToolResult` directly) or reject the unknown `task` field with `METHOD_NOT_FOUND` / `INVALID_PARAMS` fall back to the plain `session.call_tool(...)` path so legacy servers keep working. An unparseable success response (server accepted the augmented call but returned a payload that is neither `CreateTaskResult` nor `CallToolResult`) **does not** fall back — it raises `ToolExecutionException` to avoid double-executing a side-effecting tool.
|
||||
- **Submit-vs-track reconnect policy**: a dropped connection before a `task_id` is known raises `ToolExecutionException("connection lost; task state unknown")` without re-issuing the augmented `tools/call`, so a server that accepted the request but lost the response cannot be made to start the same operation twice; once a `task_id` exists, `tasks/get` / `tasks/result` reconnect once and retry against the same id (a shared `_send_with_one_reconnect` helper).
|
||||
- **Cancel-on-abandonment vs terminal failure**: any path where the remote task may still be running (max-wait exceeded, hard `McpError` in poll, malformed `tasks/get`, second connection loss in poll/fetch, reconnect failure) fires best-effort `tasks/cancel` before raising. Terminal failures (`failed`/`cancelled`/`input_required` server-side, `completed+isError`, malformed `tasks/result` after server completed) do **not** cancel — the server is already done. `_MCPTaskAbandoned` is the private marker distinguishing the two.
|
||||
- **Transient poll retry**: a slow `tasks/get` that surfaces as `McpError(code=408 REQUEST_TIMEOUT)` is retried (bounded by `max_task_wait`). All other non-connection `McpError`s during poll are treated as abandonment. `tasks/result` does not get transient retry — the server has already completed, so a slow payload fetch is anomalous.
|
||||
|
||||
### File Access Harness (`_harness/_file_access.py`)
|
||||
|
||||
- **`AgentFileStore`** - Abstract async store backing the file-access harness. Implementations expose `write`, `read`, `delete`, `list_children`, `file_exists`, `search`, and `create_directory` over forward-slash relative paths. `list_children` returns the direct children (files and subdirectories, subdirectories first) as `FileStoreEntry` instances; `search` accepts a keyword-only `recursive` flag (default `False`) and, when `recursive=True`, walks all descendants and returns `file_name` values relative to the search directory.
|
||||
- **`InMemoryAgentFileStore`** - Dict-backed store suitable for tests and lightweight scenarios.
|
||||
- **`FileSystemAgentFileStore`** - Disk-backed store rooted under a configurable directory. Enforces relative-path normalization, root containment, and rejects symlink/reparse-point segments to prevent escape.
|
||||
- **`FileSearchResult`** / **`FileSearchMatch`** - `SerializationMixin` DTOs returned by `search`, carrying the matching file name, a context snippet, and the matching lines with 1-based line numbers.
|
||||
- **`FileStoreEntry`** - `SerializationMixin` DTO returned by `list_children`, carrying an entry `name` and `type` (`"file"` or `"directory"`).
|
||||
- **`FileAccessProvider`** - `ContextProvider` that adds shared file-access tools (`file_access_write`, `file_access_read`, `file_access_delete`, `file_access_ls`, `file_access_grep`, `file_access_replace`, `file_access_replace_lines`) plus default usage instructions to each invocation. `file_access_ls` enumerates direct children (both files and subdirectories) as `{name, type}` entries with an optional `glob_pattern`, so the agent can walk the tree level by level; `file_access_grep` searches recursively from an optional base `directory` and returns relative `file_name` paths, scoped via an `fnmatch` `glob_pattern` (where `*` crosses `/`, e.g. `*.md`, `reports/*`). `file_access_replace` substitutes `old_string` with `new_string` (failing if not found, or if multiple matches and `replace_all` is false); `file_access_replace_lines` replaces whole 1-based lines with literal text (each `new_line` includes its own trailing newline; an empty `new_line` deletes the line, including its line break). All tools are registered with `approval_mode="always_require"` by default, so every file operation needs host approval. Pass `disable_write_tools=True` to advertise only the read-only tools. To run unattended you can disable approval at the source with `disable_readonly_tool_approval=True` (read, ls, grep) and/or `disable_write_tool_approval=True` (write, delete, replace, replace_lines), which register the affected tools with `approval_mode="never_require"`; alternatively, keep approval on and pass one of the static auto-approval rules to `ToolApprovalMiddleware` (via `auto_approval_rules`): `FileAccessProvider.read_only_tools_auto_approval_rule` approves only the read-only tools (read, ls, grep), while `FileAccessProvider.all_tools_auto_approval_rule` approves every file-access tool including the write tools. Both rules reject any call carrying a `server_label` so they stay scoped to this provider's local tools and never auto-approve a same-named hosted tool. The tool names are also exposed as class constants (`WRITE_TOOL_NAME`, `READ_TOOL_NAME`, `DELETE_TOOL_NAME`, `LS_TOOL_NAME`, `GREP_TOOL_NAME`, `REPLACE_TOOL_NAME`, `REPLACE_LINES_TOOL_NAME`). Unlike `MemoryContextProvider`, the store is intentionally shared across sessions and agents.
|
||||
|
||||
### File Memory Harness (`_harness/_file_memory.py`)
|
||||
|
||||
- **`FileMemoryProvider`** - `ContextProvider` that gives an agent a session-scoped, file-based memory backed by the same `AgentFileStore` abstraction. Adds tools (`file_memory_write`, `file_memory_read`, `file_memory_delete`, `file_memory_ls`, `file_memory_grep`, `file_memory_replace`, `file_memory_replace_lines`) plus default usage instructions. Port of the .NET `FileMemoryProvider`.
|
||||
- **Scoping** - Memories are isolated per session by default: each session writes under a working folder derived from `context.session_id`. Pass an explicit `scope` (e.g. a user id) to group memories across sessions, mirroring `FoundryMemoryProvider`'s `scope` arg.
|
||||
- **Descriptions & index** - `file_memory_write` accepts an optional `description`, stored in a companion `<stem>_description.md` sidecar. After each write/delete the provider rebuilds a capped (50-entry) `memories.md` index, and `before_run` injects that index as a `user` context message so the model knows what memories exist. Sidecars and the index are internal files hidden from `file_memory_ls`/`file_memory_grep` and rejected as write targets.
|
||||
- **`DEFAULT_FILE_MEMORY_SOURCE_ID`** / **`DEFAULT_FILE_MEMORY_INSTRUCTIONS`** - Public defaults for the provider's source id and instruction banner.
|
||||
- **Harness wiring** - `create_harness_agent` includes both `FileMemoryProvider` and `FileAccessProvider` by default. Disable via `disable_file_memory` / `disable_file_access`; override the backing store via `file_memory_store` / `file_access_store`. When no store is supplied, defaults are `FileSystemAgentFileStore` rooted at `{cwd}/agent-file-memory` (memory) and `{cwd}/working` (access), mirroring the .NET `HarnessAgent`. `create_harness_agent` also wires in `MessageInjectionMiddleware` by default (mirroring the .NET harness's `UseMessageInjection`); it is always on with no opt-out because it is a no-op when no messages are queued for the session.
|
||||
|
||||
### Tool Approval Harness (`_harness/_tool_approval.py`)
|
||||
|
||||
- **`ToolApprovalMiddleware`** - Experimental opt-in agent middleware that coordinates session-backed approval
|
||||
rules, heuristic `auto_approval_rules`, queued approval requests, collected approval responses, and
|
||||
streaming/non-streaming approval prompts. Heuristic callbacks receive the underlying `function_call` content.
|
||||
- **`ToolApprovalRule`** / **`ToolApprovalState`** - Serializable state models for standing approvals and queued
|
||||
approval flow. `ToolApprovalRule.arguments is None` means a tool-wide rule; an empty dict `{}` means an exact
|
||||
no-argument call for `create_always_approve_tool_with_arguments_response`.
|
||||
- **`create_always_approve_tool_response`** / **`create_always_approve_tool_with_arguments_response`** - Helpers
|
||||
that return normal `function_approval_response` content with `additional_properties` metadata consumed by
|
||||
`ToolApprovalMiddleware`. Standing rules for hosted tools include the `server_label` boundary, so same-named tools
|
||||
on different hosted servers do not share approvals.
|
||||
- Mixed tool-call batches use a default .NET-style bypass in the function invocation loop: when a session is
|
||||
available, approval requests for known non-approval-required tools are treated as already approved, hidden, stored
|
||||
in session state keyed to the visible approval request ids from that batch, and reinjected only when that visible
|
||||
approval flow resumes.
|
||||
### Agent Loop (`_harness/_loop.py`)
|
||||
|
||||
- **`AgentLoopMiddleware`** - `AgentMiddleware` that re-runs an agent in a loop by calling `call_next()` repeatedly (the pipeline re-reads `context.messages` each time). One configurable class covers two patterns: a required user `should_continue` predicate (sync or async, the first positional/keyword arg), and a chat-client judge built via the `.with_judge(...)` factory (a second chat client decides whether the original request was answered; loops while it is *not*, using a `JudgeVerdict` structured-output response — internally just an async `should_continue` predicate). The constructor covers the predicate pattern directly; only the judge has a convenience classmethod factory (`.with_judge(judge_client, ...)`) that forwards to `__init__`. Supports both streaming and non-streaming runs. By default a non-streaming run returns an aggregated `AgentResponse` containing every iteration's messages plus the injected `next_message` "nudge" messages (as `user` messages); set `return_final_only=True` to return only the last iteration's response. Streaming runs always yield each iteration's updates and emit the injected nudge messages as `user` updates between iterations (the `return_final_only` flag has no effect on streaming, and the final response reflects the last iteration; `MiddlewareTermination` is handled cleanly). `should_continue` is required; other constructor args are optional: `max_iterations` (safety cap; defaults to `DEFAULT_MAX_ITERATIONS`=10, explicit `None`→unbounded, positive int caps; `.with_judge` uses `DEFAULT_JUDGE_MAX_ITERATIONS`=5 as its default), `next_message` (defaults to a short "continue" nudge), `return_final_only`, and `additional_instructions` (an extra `system` message injected ahead of the input before the agent runs — becomes part of the original messages so it survives `fresh_context` resets and persists via a session). The judge is configured only through `.with_judge` (`judge_client`/`instructions`/`criteria`), not the constructor, and its `reasoning` is fed back to the agent as the next iteration's input; the judge forwards the original request messages and the agent's latest response messages verbatim so multi-modal content is preserved. `criteria` (a `list[str]`) is both injected as the agent's `additional_instructions` and rendered into the judge instructions wherever the `{{criteria}}` placeholder (`CRITERIA_PLACEHOLDER`) appears (`DEFAULT_JUDGE_INSTRUCTIONS` ends with it; custom `instructions` may include it, and it is stripped when no criteria are given). The `should_continue`/`next_message` callables are invoked with keyword args (`iteration`, `last_result`, `messages`, `original_messages`, `session`, `agent`, `progress`, `feedback`) and may be sync or async; declare only what you need plus `**kwargs`. `should_continue` may return a plain `bool` or a `(bool, str | None)` tuple whose second item is feedback surfaced to `next_message`/`record_feedback` via the `feedback` kwarg (the judge uses this to relay its `reasoning`). Stop precedence per iteration is `max_iterations` → `should_continue`, evaluated before `record_feedback` so the feedback is available to it.
|
||||
- **Feedback tracking** - `record_feedback` captures a per-iteration progress entry (called with the loop kwargs; if it returns a truthy string the entry is appended, otherwise the agent's response text is used as the fallback entry). The accumulated log is exposed to every callback via the `progress` keyword (a per-iteration copy of prior entries) and, when `inject_progress=True` (default), injected into the next iteration's input as a `user` message (the full log without a session, only the latest entry with a session to avoid duplicating history). `fresh_context=True` restarts each iteration from the original task plus the progress log; when a session is attached it is snapshotted (`to_dict()`) before the loop and restored (`from_dict` + field copy) between iterations so the local transcript and any service-side conversation id reset too (in-loop working-state is discarded, pre-loop state preserved, continuity carried only by the progress log).
|
||||
- **`todos_remaining(*, looping_modes=None)`** / **`todos_remaining_message`** - Helper factories for todo-driven loops (the Python counterpart of .NET's `TodoCompletionLoopEvaluator`), designed for `create_harness_agent` but usable with any agent that registers a `TodoProvider` via `context_providers`. They resolve the `TodoProvider`/`AgentModeProvider` from the *running agent* (`agent.context_providers`, via `_resolve_context_provider`) rather than taking the provider as an argument, so they can be wired directly into `loop_should_continue`/`loop_next_message`. `todos_remaining` returns a `should_continue` predicate that loops while any todo is open; pass `looping_modes=[...]` to gate looping to specific operating modes (case-insensitive; honors the `AgentModeProvider`'s `source_id`/`available_modes`), `looping_modes=None` (default) applies in every mode, and an empty sequence raises `ValueError`. `todos_remaining_message` is a `next_message` callable that lists the still-open todo titles and tells the agent to finish them, returning `None` when the session/agent/provider is unavailable or nothing is open (in which case the middleware's default `None` handling applies: reuse the previous iteration's messages verbatim under the default `fresh_context=False`, or `DEFAULT_NEXT_MESSAGE` only when `fresh_context=True`).
|
||||
- **`background_tasks_running()`** / **`background_tasks_running_message`** - Helper factories for background-agent-driven loops, mirroring the `todos_remaining` pair. They resolve the `BackgroundAgentsProvider` from the *running agent* (`agent.context_providers`, via `_resolve_context_provider`) rather than taking the provider as an argument, so they can be wired directly into `create_harness_agent`'s `loop_should_continue`/`loop_next_message`. `background_tasks_running` returns a `should_continue` predicate that loops while the provider's persisted state shows any task with `status == RUNNING` (pair it with `max_iterations` so the loop is bounded even if a task's persisted status is never refreshed). `background_tasks_running_message` is a `next_message` callable that lists the still-running tasks (`#<id> (<agent_name>): <description>`) and tells the agent to wait for them to finish and retrieve their results, returning `None` when the session/agent/provider is unavailable or no task is running.
|
||||
- **Approval escape hatch** - `_has_pending_approval_request(result)` checks whether an iteration's response carries a pending tool-approval request (any content with `type == "function_approval_request"`). Both the streaming and non-streaming loops stop and return that response to the caller *before* evaluating `should_continue`/`max_iterations` or injecting `next_message`, so the loop is HITL-safe even when wrapped outermost around a `ToolApprovalMiddleware` (mirrors the C# `LoopAgent`'s `HasPendingApprovalRequests`).
|
||||
- **Harness integration** - `create_harness_agent` enables the loop when a `loop_should_continue` callable is passed; it prepends `AgentLoopMiddleware(loop_should_continue, max_iterations=loop_max_iterations, next_message=loop_next_message)` ahead of `ToolApprovalMiddleware` so the loop is the outermost middleware (each iteration is a full agent run including tool approval, and the escape hatch hands pending approvals back to the caller). `loop_next_message` and `loop_max_iterations` only take effect together with `loop_should_continue` (with no `loop_should_continue` there is no loop, so they are ignored); `loop_max_iterations` defaults to the loop's default cap (`None` → unbounded).
|
||||
|
||||
### Workflows (`_workflows/`)
|
||||
|
||||
- **`Workflow`** - Graph-based workflow definition
|
||||
- **`WorkflowBuilder`** - Fluent API for building workflows, including explicit
|
||||
`output_from` / `intermediate_output_from` selection for caller-facing emissions. `output_from`
|
||||
is an allow-list for **Workflow Output**; unselected executor payloads are hidden unless
|
||||
`intermediate_output_from` selects them as **Intermediate Output**. Use `output_from="all"` for
|
||||
explicit all-output behavior and `intermediate_output_from="all_other"` for visible progress from
|
||||
every output-capable executor not selected by `output_from`.
|
||||
- **`WorkflowRunResult`** - Non-streaming workflow result with Workflow Output `get_outputs()`
|
||||
and Intermediate Output `get_intermediate_outputs()` accessors
|
||||
- **Orchestrators**: `SequentialOrchestrator`, `ConcurrentOrchestrator`, `GroupChatOrchestrator`, `MagenticOrchestrator`, `HandoffOrchestrator`
|
||||
|
||||
## Built-in Providers
|
||||
|
||||
### OpenAI (`openai/`)
|
||||
|
||||
- **`OpenAIChatClient`** - Chat client for the OpenAI Responses API
|
||||
- **`OpenAIChatCompletionClient`** - Chat client for the OpenAI Chat Completions API
|
||||
|
||||
### Foundry (`foundry/`)
|
||||
|
||||
- **`FoundryChatClient`** - Chat client for Azure AI Foundry project endpoints
|
||||
|
||||
## Key Patterns
|
||||
|
||||
### Creating an Agent
|
||||
|
||||
```python
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
agent = Agent(
|
||||
client=OpenAIChatClient(),
|
||||
instructions="You are helpful.",
|
||||
tools=[my_function],
|
||||
)
|
||||
response = await agent.run("Hello")
|
||||
```
|
||||
|
||||
### Using `as_agent()` Shorthand
|
||||
|
||||
```python
|
||||
agent = OpenAIChatClient().as_agent(
|
||||
name="Assistant",
|
||||
instructions="You are helpful.",
|
||||
)
|
||||
```
|
||||
|
||||
### Middleware Pipeline
|
||||
|
||||
```python
|
||||
from agent_framework import Agent, AgentMiddleware, AgentContext
|
||||
|
||||
|
||||
class LoggingMiddleware(AgentMiddleware):
|
||||
async def process(self, context: AgentContext, call_next) -> None:
|
||||
print(f"Input: {context.messages}")
|
||||
await call_next()
|
||||
print(f"Output: {context.result}")
|
||||
|
||||
|
||||
agent = Agent(..., middleware=[LoggingMiddleware()])
|
||||
```
|
||||
|
||||
### Custom Chat Client
|
||||
|
||||
```python
|
||||
from agent_framework import BaseChatClient, ChatResponse, Message
|
||||
|
||||
|
||||
class MyClient(BaseChatClient):
|
||||
async def _inner_get_response(self, *, messages, options, **kwargs) -> ChatResponse:
|
||||
# Call your LLM here
|
||||
return ChatResponse(messages=[Message(role="assistant", contents=["Hi!"])])
|
||||
|
||||
async def _inner_get_streaming_response(self, *, messages, options, **kwargs):
|
||||
yield ChatResponseUpdate(...)
|
||||
```
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
@@ -0,0 +1,225 @@
|
||||
# Get Started with Microsoft Agent Framework
|
||||
|
||||
Highlights
|
||||
|
||||
- Flexible Agent Framework: build, orchestrate, and deploy AI agents and multi-agent systems
|
||||
- Multi-Agent Orchestration: Group chat, sequential, concurrent, and handoff patterns
|
||||
- Plugin Ecosystem: Extend with native functions, OpenAPI, Model Context Protocol (MCP), and more
|
||||
- LLM Support: OpenAI, Foundry, Anthropic, and more
|
||||
- Runtime Support: In-process and distributed agent execution
|
||||
- Multimodal: Text, vision, and function calling
|
||||
- Cross-Platform: .NET and Python implementations
|
||||
|
||||
## Quick Install
|
||||
|
||||
```bash
|
||||
pip install agent-framework-core
|
||||
# Optional: Add Azure AI Foundry integration
|
||||
pip install agent-framework-foundry
|
||||
# Optional: Add OpenAI integration
|
||||
pip install agent-framework-openai
|
||||
```
|
||||
|
||||
Supported Platforms:
|
||||
|
||||
- Python: 3.10+
|
||||
- OS: Windows, macOS, Linux
|
||||
|
||||
## 1. Setup API Keys
|
||||
|
||||
Depending on the client you want to use, there are various environment variables you can set to configure the chat clients. This can be done in the environment itself, or with a `.env` file in your project root, some examples of environment variables include:
|
||||
|
||||
```bash
|
||||
FOUNDRY_PROJECT_ENDPOINT=...
|
||||
FOUNDRY_MODEL=...
|
||||
...
|
||||
OPENAI_API_KEY=sk-...
|
||||
OPENAI_CHAT_COMPLETION_MODEL=...
|
||||
OPENAI_CHAT_MODEL=...
|
||||
...
|
||||
AZURE_OPENAI_API_KEY=...
|
||||
AZURE_OPENAI_ENDPOINT=...
|
||||
AZURE_OPENAI_MODEL=...
|
||||
```
|
||||
|
||||
You can also override environment variables by explicitly passing configuration parameters to the chat client constructor:
|
||||
|
||||
```python
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
client = OpenAIChatClient(
|
||||
api_key="",
|
||||
model="",
|
||||
)
|
||||
```
|
||||
|
||||
See the following [getting started samples](https://github.com/microsoft/agent-framework/tree/main/python/samples/01-get-started) for more information.
|
||||
|
||||
## 2. Create a Simple Agent
|
||||
|
||||
Create agents and invoke them directly:
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
agent = Agent(
|
||||
client=OpenAIChatClient(),
|
||||
instructions="""
|
||||
1) A robot may not injure a human being...
|
||||
2) A robot must obey orders given it by human beings...
|
||||
3) A robot must protect its own existence...
|
||||
|
||||
Give me the TLDR in exactly 5 words.
|
||||
"""
|
||||
)
|
||||
|
||||
result = asyncio.run(agent.run("Summarize the Three Laws of Robotics"))
|
||||
print(result)
|
||||
# Output: Protect humans, obey, self-preserve, prioritized.
|
||||
```
|
||||
|
||||
## 3. Directly Use Chat Clients (No Agent Required)
|
||||
|
||||
You can use the chat client classes directly for advanced workflows:
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
from agent_framework import Message, Role
|
||||
|
||||
async def main():
|
||||
client = OpenAIChatClient()
|
||||
|
||||
response = await client.get_response([
|
||||
Message("system", ["You are a helpful assistant."]),
|
||||
Message("user", ["Write a haiku about Agent Framework."])
|
||||
])
|
||||
print(response.messages[0].text)
|
||||
|
||||
"""
|
||||
Output:
|
||||
|
||||
Agents work in sync,
|
||||
Framework threads through each task—
|
||||
Code sparks collaboration.
|
||||
"""
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
## 4. Build an Agent with Tools and Functions
|
||||
|
||||
Enhance your agent with custom tools and function calling:
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from typing import Annotated
|
||||
from random import randint
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
|
||||
def get_weather(
|
||||
location: Annotated[str, "The location to get the weather for."],
|
||||
) -> str:
|
||||
"""Get the weather for a given location."""
|
||||
conditions = ["sunny", "cloudy", "rainy", "stormy"]
|
||||
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
|
||||
|
||||
|
||||
def get_menu_specials() -> str:
|
||||
"""Get today's menu specials."""
|
||||
return """
|
||||
Special Soup: Clam Chowder
|
||||
Special Salad: Cobb Salad
|
||||
Special Drink: Chai Tea
|
||||
"""
|
||||
|
||||
|
||||
async def main():
|
||||
agent = Agent(
|
||||
client=OpenAIChatClient(),
|
||||
instructions="You are a helpful assistant that can provide weather and restaurant information.",
|
||||
tools=[get_weather, get_menu_specials]
|
||||
)
|
||||
|
||||
response = await agent.run("What's the weather in Amsterdam and what are today's specials?")
|
||||
print(response)
|
||||
|
||||
# Output:
|
||||
# The weather in Amsterdam is sunny with a high of 22°C. Today's specials include
|
||||
# Clam Chowder soup, Cobb Salad, and Chai Tea as the special drink.
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
You can explore additional agent samples [here](https://github.com/microsoft/agent-framework/tree/main/python/samples/02-agents).
|
||||
|
||||
## 5. Multi-Agent Orchestration
|
||||
|
||||
Coordinate multiple agents to collaborate on complex tasks using orchestration patterns:
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
|
||||
async def main():
|
||||
# Create specialized agents
|
||||
writer = Agent(
|
||||
client=OpenAIChatClient(),
|
||||
name="Writer",
|
||||
instructions="You are a creative content writer. Generate and refine slogans based on feedback."
|
||||
)
|
||||
|
||||
reviewer = Agent(
|
||||
client=OpenAIChatClient(),
|
||||
name="Reviewer",
|
||||
instructions="You are a critical reviewer. Provide detailed feedback on proposed slogans."
|
||||
)
|
||||
|
||||
# Sequential workflow: Writer creates, Reviewer provides feedback
|
||||
task = "Create a slogan for a new electric SUV that is affordable and fun to drive."
|
||||
|
||||
# Step 1: Writer creates initial slogan
|
||||
initial_result = await writer.run(task)
|
||||
print(f"Writer: {initial_result}")
|
||||
|
||||
# Step 2: Reviewer provides feedback
|
||||
feedback_request = f"Please review this slogan: {initial_result}"
|
||||
feedback = await reviewer.run(feedback_request)
|
||||
print(f"Reviewer: {feedback}")
|
||||
|
||||
# Step 3: Writer refines based on feedback
|
||||
refinement_request = f"Please refine this slogan based on the feedback: {initial_result}\nFeedback: {feedback}"
|
||||
final_result = await writer.run(refinement_request)
|
||||
print(f"Final Slogan: {final_result}")
|
||||
|
||||
# Example Output:
|
||||
# Writer: "Charge Forward: Affordable Adventure Awaits!"
|
||||
# Reviewer: "Good energy, but 'Charge Forward' is overused in EV marketing..."
|
||||
# Final Slogan: "Power Up Your Adventure: Premium Feel, Smart Price!"
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
**Note**: Sequential, Concurrent, Group Chat, Handoff, and Magentic orchestrations are available. See examples in [orchestration samples](https://github.com/microsoft/agent-framework/tree/main/python/samples/03-workflows/orchestrations).
|
||||
|
||||
## More Examples & Samples
|
||||
|
||||
- [Getting Started with Agents](https://github.com/microsoft/agent-framework/tree/main/python/samples/02-agents): Basic agent creation and tool usage
|
||||
- [Chat Client Examples](https://github.com/microsoft/agent-framework/tree/main/python/samples/02-agents/chat_client): Direct chat client usage patterns
|
||||
- [Foundry Integration](https://github.com/microsoft/agent-framework/tree/main/python/packages/foundry): Foundry integration
|
||||
- [Workflows Samples](https://github.com/microsoft/agent-framework/tree/main/python/samples/03-workflows): Advanced multi-agent patterns
|
||||
|
||||
## Agent Framework Documentation
|
||||
|
||||
- [Agent Framework Repository](https://github.com/microsoft/agent-framework)
|
||||
- [Python Package Documentation](https://github.com/microsoft/agent-framework/tree/main/python)
|
||||
- [.NET Package Documentation](https://github.com/microsoft/agent-framework/tree/main/dotnet)
|
||||
- [Design Documents](https://github.com/microsoft/agent-framework/tree/main/docs/design)
|
||||
- [Learn Documentation](https://learn.microsoft.com/agent-framework/)
|
||||
@@ -0,0 +1,652 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Public API surface for Agent Framework core.
|
||||
|
||||
This module exposes the primary abstractions for agents, chat clients, tools, sessions,
|
||||
middleware, observability, and workflows. Most public exports are resolved lazily to keep
|
||||
``import agent_framework`` lightweight; importing a specific symbol still loads the module
|
||||
that owns that symbol.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# pyright: reportUnsupportedDunderAll=false
|
||||
# ruff: noqa: F822
|
||||
import importlib
|
||||
import importlib.metadata
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
try:
|
||||
_version = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
_version = "0.0.0" # Fallback for development mode
|
||||
__version__: Final[str] = _version
|
||||
|
||||
from ._telemetry import (
|
||||
AGENT_FRAMEWORK_USER_AGENT,
|
||||
APP_INFO,
|
||||
USER_AGENT_KEY,
|
||||
USER_AGENT_TELEMETRY_DISABLED_ENV_VAR,
|
||||
prepend_agent_framework_to_user_agent,
|
||||
)
|
||||
from .exceptions import (
|
||||
AgentFrameworkException,
|
||||
MiddlewareException,
|
||||
UserInputRequiredException,
|
||||
WorkflowCheckpointException,
|
||||
WorkflowConvergenceException,
|
||||
WorkflowException,
|
||||
WorkflowRunnerException,
|
||||
)
|
||||
|
||||
_LAZY_MODULE_EXPORTS: Final[Mapping[str, tuple[str, ...]]] = {
|
||||
"._agents": ("Agent", "BaseAgent", "RawAgent", "SupportsAgentRun"),
|
||||
"._clients": (
|
||||
"BaseChatClient",
|
||||
"BaseEmbeddingClient",
|
||||
"SupportsChatGetResponse",
|
||||
"SupportsCodeInterpreterTool",
|
||||
"SupportsFileSearchTool",
|
||||
"SupportsGetEmbeddings",
|
||||
"SupportsImageGenerationTool",
|
||||
"SupportsMCPTool",
|
||||
"SupportsShellTool",
|
||||
"SupportsWebSearchTool",
|
||||
),
|
||||
"._compaction": (
|
||||
"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",
|
||||
"ContextWindowCompactionStrategy",
|
||||
"SelectiveToolCallCompactionStrategy",
|
||||
"SlidingWindowStrategy",
|
||||
"SummarizationStrategy",
|
||||
"TokenBudgetComposedStrategy",
|
||||
"TokenizerProtocol",
|
||||
"ToolResultCompactionStrategy",
|
||||
"TruncationStrategy",
|
||||
"annotate_message_groups",
|
||||
"apply_compaction",
|
||||
"included_messages",
|
||||
"included_token_count",
|
||||
),
|
||||
"._evaluation": (
|
||||
"AgentEvalConverter",
|
||||
"CheckResult",
|
||||
"ConversationSplit",
|
||||
"ConversationSplitter",
|
||||
"EvalItem",
|
||||
"EvalItemResult",
|
||||
"EvalNotPassedError",
|
||||
"EvalResults",
|
||||
"EvalScoreResult",
|
||||
"Evaluator",
|
||||
"ExpectedToolCall",
|
||||
"LocalEvaluator",
|
||||
"RubricScore",
|
||||
"evaluate_agent",
|
||||
"evaluate_workflow",
|
||||
"evaluator",
|
||||
"keyword_check",
|
||||
"tool_call_args_match",
|
||||
"tool_called_check",
|
||||
"tool_calls_present",
|
||||
),
|
||||
"._feature_stage": ("ExperimentalFeature", "ReleaseCandidateFeature"),
|
||||
"._harness._agent": ("DEFAULT_HARNESS_INSTRUCTIONS", "create_harness_agent"),
|
||||
"._harness._background_agents": (
|
||||
"DEFAULT_BACKGROUND_AGENTS_SOURCE_ID",
|
||||
"BackgroundAgentsProvider",
|
||||
"BackgroundTaskInfo",
|
||||
"BackgroundTaskStatus",
|
||||
),
|
||||
"._harness._file_access": (
|
||||
"DEFAULT_FILE_ACCESS_INSTRUCTIONS",
|
||||
"DEFAULT_FILE_ACCESS_SOURCE_ID",
|
||||
"AgentFileStore",
|
||||
"FileAccessProvider",
|
||||
"FileSearchMatch",
|
||||
"FileSearchResult",
|
||||
"FileStoreEntry",
|
||||
"FileSystemAgentFileStore",
|
||||
"InMemoryAgentFileStore",
|
||||
),
|
||||
"._harness._file_memory": (
|
||||
"DEFAULT_FILE_MEMORY_INSTRUCTIONS",
|
||||
"DEFAULT_FILE_MEMORY_SOURCE_ID",
|
||||
"FileMemoryProvider",
|
||||
),
|
||||
"._harness._loop": (
|
||||
"AgentLoopMiddleware",
|
||||
"JudgeVerdict",
|
||||
"background_tasks_running",
|
||||
"background_tasks_running_message",
|
||||
"todos_remaining",
|
||||
"todos_remaining_message",
|
||||
),
|
||||
"._harness._memory": (
|
||||
"DEFAULT_MEMORY_SOURCE_ID",
|
||||
"MemoryContextProvider",
|
||||
"MemoryFileStore",
|
||||
"MemoryIndexEntry",
|
||||
"MemoryStore",
|
||||
"MemoryTopicRecord",
|
||||
),
|
||||
"._harness._mode": ("DEFAULT_MODE_SOURCE_ID", "AgentModeProvider", "get_agent_mode", "set_agent_mode"),
|
||||
"._harness._todo": (
|
||||
"DEFAULT_TODO_SOURCE_ID",
|
||||
"TodoFileStore",
|
||||
"TodoInput",
|
||||
"TodoItem",
|
||||
"TodoProvider",
|
||||
"TodoSessionStore",
|
||||
"TodoStore",
|
||||
),
|
||||
"._harness._tool_approval": (
|
||||
"DEFAULT_TOOL_APPROVAL_SOURCE_ID",
|
||||
"ToolApprovalMiddleware",
|
||||
"ToolApprovalRule",
|
||||
"ToolApprovalRuleCallback",
|
||||
"ToolApprovalState",
|
||||
"create_always_approve_tool_response",
|
||||
"create_always_approve_tool_with_arguments_response",
|
||||
),
|
||||
"._mcp": (
|
||||
"MCPStdioTool",
|
||||
"MCPStreamableHTTPTool",
|
||||
"MCPTaskOptions",
|
||||
"MCPWebsocketTool",
|
||||
"SamplingApprovalCallback",
|
||||
),
|
||||
"._middleware": (
|
||||
"AgentContext",
|
||||
"AgentMiddleware",
|
||||
"AgentMiddlewareLayer",
|
||||
"AgentMiddlewareTypes",
|
||||
"ChatAndFunctionMiddlewareTypes",
|
||||
"ChatContext",
|
||||
"ChatMiddleware",
|
||||
"ChatMiddlewareLayer",
|
||||
"ChatMiddlewareTypes",
|
||||
"FunctionInvocationContext",
|
||||
"FunctionMiddleware",
|
||||
"FunctionMiddlewareTypes",
|
||||
"MiddlewareTermination",
|
||||
"MiddlewareType",
|
||||
"MiddlewareTypes",
|
||||
"agent_middleware",
|
||||
"chat_middleware",
|
||||
"function_middleware",
|
||||
),
|
||||
"._sessions": (
|
||||
"AgentSession",
|
||||
"ContextProvider",
|
||||
"FileHistoryProvider",
|
||||
"HistoryProvider",
|
||||
"InMemoryHistoryProvider",
|
||||
"MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY",
|
||||
"MessageInjectionMiddleware",
|
||||
"ServiceSessionId",
|
||||
"SessionContext",
|
||||
"enqueue_messages",
|
||||
"register_state_type",
|
||||
),
|
||||
"._settings": ("SecretString", "load_settings"),
|
||||
"._skills": (
|
||||
"AggregatingSkillsSource",
|
||||
"CachingSkillsSource",
|
||||
"ClassSkill",
|
||||
"DeduplicatingSkillsSource",
|
||||
"DelegatingSkillsSource",
|
||||
"FileSkill",
|
||||
"FileSkillScript",
|
||||
"FileSkillsSource",
|
||||
"FilteringSkillsSource",
|
||||
"InlineSkill",
|
||||
"InlineSkillResource",
|
||||
"InlineSkillScript",
|
||||
"InMemorySkillsSource",
|
||||
"MCPSkill",
|
||||
"MCPSkillResource",
|
||||
"MCPSkillsSource",
|
||||
"Skill",
|
||||
"SkillFrontmatter",
|
||||
"SkillResource",
|
||||
"SkillScript",
|
||||
"SkillScriptArgumentParser",
|
||||
"SkillScriptRunner",
|
||||
"SkillsProvider",
|
||||
"SkillsSource",
|
||||
"SkillsSourceContext",
|
||||
),
|
||||
"._tools": (
|
||||
"SKIP_PARSING",
|
||||
"FunctionInvocationConfiguration",
|
||||
"FunctionInvocationLayer",
|
||||
"FunctionTool",
|
||||
"ToolTypes",
|
||||
"normalize_function_invocation_configuration",
|
||||
"tool",
|
||||
),
|
||||
"._types": (
|
||||
"AgentResponse",
|
||||
"AgentResponseUpdate",
|
||||
"AgentRunInputs",
|
||||
"Annotation",
|
||||
"ChatOptions",
|
||||
"ChatResponse",
|
||||
"ChatResponseUpdate",
|
||||
"Content",
|
||||
"ContinuationToken",
|
||||
"Embedding",
|
||||
"EmbeddingGenerationOptions",
|
||||
"EmbeddingInputT",
|
||||
"EmbeddingT",
|
||||
"FinalT",
|
||||
"FinishReason",
|
||||
"FinishReasonLiteral",
|
||||
"GeneratedEmbeddings",
|
||||
"Message",
|
||||
"OuterFinalT",
|
||||
"OuterUpdateT",
|
||||
"ResponseStream",
|
||||
"Role",
|
||||
"RoleLiteral",
|
||||
"TextSpanRegion",
|
||||
"ToolMode",
|
||||
"UpdateT",
|
||||
"UsageDetails",
|
||||
"add_usage_details",
|
||||
"detect_media_type_from_base64",
|
||||
"map_chat_to_agent_update",
|
||||
"merge_chat_options",
|
||||
"normalize_messages",
|
||||
"normalize_tools",
|
||||
"prepend_instructions_to_messages",
|
||||
"validate_chat_options",
|
||||
"validate_tool_mode",
|
||||
"validate_tools",
|
||||
),
|
||||
"._workflows._agent": ("WorkflowAgent",),
|
||||
"._workflows._agent_executor": ("AgentExecutor", "AgentExecutorRequest", "AgentExecutorResponse"),
|
||||
"._workflows._agent_utils": ("resolve_agent_id",),
|
||||
"._workflows._checkpoint": (
|
||||
"CheckpointID",
|
||||
"CheckpointStorage",
|
||||
"FileCheckpointStorage",
|
||||
"InMemoryCheckpointStorage",
|
||||
"WorkflowCheckpoint",
|
||||
),
|
||||
"._workflows._const": ("DEFAULT_MAX_ITERATIONS",),
|
||||
"._workflows._edge": (
|
||||
"Case",
|
||||
"Default",
|
||||
"Edge",
|
||||
"EdgeCondition",
|
||||
"FanInEdgeGroup",
|
||||
"FanOutEdgeGroup",
|
||||
"SingleEdgeGroup",
|
||||
"SwitchCaseEdgeGroup",
|
||||
"SwitchCaseEdgeGroupCase",
|
||||
"SwitchCaseEdgeGroupDefault",
|
||||
),
|
||||
"._workflows._edge_runner": ("create_edge_runner",),
|
||||
"._workflows._events": (
|
||||
"WorkflowErrorDetails",
|
||||
"WorkflowEvent",
|
||||
"WorkflowEventSource",
|
||||
"WorkflowEventType",
|
||||
"WorkflowRunState",
|
||||
),
|
||||
"._workflows._executor": ("Executor", "handler"),
|
||||
"._workflows._function_executor": ("FunctionExecutor", "executor"),
|
||||
"._workflows._functional": (
|
||||
"FunctionalWorkflow",
|
||||
"FunctionalWorkflowAgent",
|
||||
"RunContext",
|
||||
"StepWrapper",
|
||||
"get_run_context",
|
||||
"step",
|
||||
"workflow",
|
||||
),
|
||||
"._workflows._request_info_mixin": ("response_handler",),
|
||||
"._workflows._runner": ("Runner",),
|
||||
"._workflows._runner_context": ("InProcRunnerContext", "RunnerContext", "WorkflowMessage"),
|
||||
"._workflows._validation": (
|
||||
"EdgeDuplicationError",
|
||||
"GraphConnectivityError",
|
||||
"TypeCompatibilityError",
|
||||
"ValidationTypeEnum",
|
||||
"WorkflowValidationError",
|
||||
"validate_workflow_graph",
|
||||
),
|
||||
"._workflows._viz": ("WorkflowViz",),
|
||||
"._workflows._workflow": ("Workflow", "WorkflowRunResult"),
|
||||
"._workflows._workflow_builder": ("WorkflowBuilder",),
|
||||
"._workflows._workflow_context": ("WorkflowContext",),
|
||||
"._workflows._workflow_executor": (
|
||||
"SubWorkflowRequestMessage",
|
||||
"SubWorkflowResponseMessage",
|
||||
"WorkflowExecutor",
|
||||
),
|
||||
}
|
||||
_LAZY_EXPORTS: Final[dict[str, str]] = {
|
||||
name: module_name for module_name, names in _LAZY_MODULE_EXPORTS.items() for name in names
|
||||
}
|
||||
|
||||
__all__ = [
|
||||
"AGENT_FRAMEWORK_USER_AGENT",
|
||||
"APP_INFO",
|
||||
"COMPACTION_STATE_KEY",
|
||||
"DEFAULT_BACKGROUND_AGENTS_SOURCE_ID",
|
||||
"DEFAULT_FILE_ACCESS_INSTRUCTIONS",
|
||||
"DEFAULT_FILE_ACCESS_SOURCE_ID",
|
||||
"DEFAULT_FILE_MEMORY_INSTRUCTIONS",
|
||||
"DEFAULT_FILE_MEMORY_SOURCE_ID",
|
||||
"DEFAULT_HARNESS_INSTRUCTIONS",
|
||||
"DEFAULT_MAX_ITERATIONS",
|
||||
"DEFAULT_MEMORY_SOURCE_ID",
|
||||
"DEFAULT_MODE_SOURCE_ID",
|
||||
"DEFAULT_TODO_SOURCE_ID",
|
||||
"DEFAULT_TOOL_APPROVAL_SOURCE_ID",
|
||||
"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",
|
||||
"MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY",
|
||||
"SKIP_PARSING",
|
||||
"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",
|
||||
"AgentContext",
|
||||
"AgentEvalConverter",
|
||||
"AgentExecutor",
|
||||
"AgentExecutorRequest",
|
||||
"AgentExecutorResponse",
|
||||
"AgentFileStore",
|
||||
"AgentFrameworkException",
|
||||
"AgentLoopMiddleware",
|
||||
"AgentMiddleware",
|
||||
"AgentMiddlewareLayer",
|
||||
"AgentMiddlewareTypes",
|
||||
"AgentModeProvider",
|
||||
"AgentResponse",
|
||||
"AgentResponseUpdate",
|
||||
"AgentRunInputs",
|
||||
"AgentSession",
|
||||
"AggregatingSkillsSource",
|
||||
"Annotation",
|
||||
"BackgroundAgentsProvider",
|
||||
"BackgroundTaskInfo",
|
||||
"BackgroundTaskStatus",
|
||||
"BaseAgent",
|
||||
"BaseChatClient",
|
||||
"BaseEmbeddingClient",
|
||||
"CachingSkillsSource",
|
||||
"Case",
|
||||
"CharacterEstimatorTokenizer",
|
||||
"ChatAndFunctionMiddlewareTypes",
|
||||
"ChatContext",
|
||||
"ChatMiddleware",
|
||||
"ChatMiddlewareLayer",
|
||||
"ChatMiddlewareTypes",
|
||||
"ChatOptions",
|
||||
"ChatResponse",
|
||||
"ChatResponseUpdate",
|
||||
"CheckResult",
|
||||
"CheckpointID",
|
||||
"CheckpointStorage",
|
||||
"ClassSkill",
|
||||
"CompactionProvider",
|
||||
"CompactionStrategy",
|
||||
"Content",
|
||||
"ContextProvider",
|
||||
"ContextWindowCompactionStrategy",
|
||||
"ContinuationToken",
|
||||
"ConversationSplit",
|
||||
"ConversationSplitter",
|
||||
"DeduplicatingSkillsSource",
|
||||
"Default",
|
||||
"DelegatingSkillsSource",
|
||||
"Edge",
|
||||
"EdgeCondition",
|
||||
"EdgeDuplicationError",
|
||||
"Embedding",
|
||||
"EmbeddingGenerationOptions",
|
||||
"EmbeddingInputT",
|
||||
"EmbeddingT",
|
||||
"EvalItem",
|
||||
"EvalItemResult",
|
||||
"EvalNotPassedError",
|
||||
"EvalResults",
|
||||
"EvalScoreResult",
|
||||
"Evaluator",
|
||||
"Executor",
|
||||
"ExpectedToolCall",
|
||||
"ExperimentalFeature",
|
||||
"FanInEdgeGroup",
|
||||
"FanOutEdgeGroup",
|
||||
"FileAccessProvider",
|
||||
"FileCheckpointStorage",
|
||||
"FileHistoryProvider",
|
||||
"FileMemoryProvider",
|
||||
"FileSearchMatch",
|
||||
"FileSearchResult",
|
||||
"FileSkill",
|
||||
"FileSkillScript",
|
||||
"FileSkillsSource",
|
||||
"FileStoreEntry",
|
||||
"FileSystemAgentFileStore",
|
||||
"FilteringSkillsSource",
|
||||
"FinalT",
|
||||
"FinishReason",
|
||||
"FinishReasonLiteral",
|
||||
"FunctionExecutor",
|
||||
"FunctionInvocationConfiguration",
|
||||
"FunctionInvocationContext",
|
||||
"FunctionInvocationLayer",
|
||||
"FunctionMiddleware",
|
||||
"FunctionMiddlewareTypes",
|
||||
"FunctionTool",
|
||||
"FunctionalWorkflow",
|
||||
"FunctionalWorkflowAgent",
|
||||
"GeneratedEmbeddings",
|
||||
"GraphConnectivityError",
|
||||
"HistoryProvider",
|
||||
"InMemoryAgentFileStore",
|
||||
"InMemoryCheckpointStorage",
|
||||
"InMemoryHistoryProvider",
|
||||
"InMemorySkillsSource",
|
||||
"InProcRunnerContext",
|
||||
"InlineSkill",
|
||||
"InlineSkillResource",
|
||||
"InlineSkillScript",
|
||||
"JudgeVerdict",
|
||||
"LocalEvaluator",
|
||||
"MCPSkill",
|
||||
"MCPSkillResource",
|
||||
"MCPSkillsSource",
|
||||
"MCPStdioTool",
|
||||
"MCPStreamableHTTPTool",
|
||||
"MCPTaskOptions",
|
||||
"MCPWebsocketTool",
|
||||
"MemoryContextProvider",
|
||||
"MemoryFileStore",
|
||||
"MemoryIndexEntry",
|
||||
"MemoryStore",
|
||||
"MemoryTopicRecord",
|
||||
"Message",
|
||||
"MessageInjectionMiddleware",
|
||||
"MiddlewareException",
|
||||
"MiddlewareTermination",
|
||||
"MiddlewareType",
|
||||
"MiddlewareTypes",
|
||||
"OuterFinalT",
|
||||
"OuterUpdateT",
|
||||
"RawAgent",
|
||||
"ReleaseCandidateFeature",
|
||||
"ResponseStream",
|
||||
"Role",
|
||||
"RoleLiteral",
|
||||
"RubricScore",
|
||||
"RunContext",
|
||||
"Runner",
|
||||
"RunnerContext",
|
||||
"SamplingApprovalCallback",
|
||||
"SecretString",
|
||||
"SelectiveToolCallCompactionStrategy",
|
||||
"ServiceSessionId",
|
||||
"SessionContext",
|
||||
"SingleEdgeGroup",
|
||||
"Skill",
|
||||
"SkillFrontmatter",
|
||||
"SkillResource",
|
||||
"SkillScript",
|
||||
"SkillScriptArgumentParser",
|
||||
"SkillScriptRunner",
|
||||
"SkillsProvider",
|
||||
"SkillsSource",
|
||||
"SkillsSourceContext",
|
||||
"SlidingWindowStrategy",
|
||||
"StepWrapper",
|
||||
"SubWorkflowRequestMessage",
|
||||
"SubWorkflowResponseMessage",
|
||||
"SummarizationStrategy",
|
||||
"SupportsAgentRun",
|
||||
"SupportsChatGetResponse",
|
||||
"SupportsCodeInterpreterTool",
|
||||
"SupportsFileSearchTool",
|
||||
"SupportsGetEmbeddings",
|
||||
"SupportsImageGenerationTool",
|
||||
"SupportsMCPTool",
|
||||
"SupportsShellTool",
|
||||
"SupportsWebSearchTool",
|
||||
"SwitchCaseEdgeGroup",
|
||||
"SwitchCaseEdgeGroupCase",
|
||||
"SwitchCaseEdgeGroupDefault",
|
||||
"TextSpanRegion",
|
||||
"TodoFileStore",
|
||||
"TodoInput",
|
||||
"TodoItem",
|
||||
"TodoProvider",
|
||||
"TodoSessionStore",
|
||||
"TodoStore",
|
||||
"TokenBudgetComposedStrategy",
|
||||
"TokenizerProtocol",
|
||||
"ToolApprovalMiddleware",
|
||||
"ToolApprovalRule",
|
||||
"ToolApprovalRuleCallback",
|
||||
"ToolApprovalState",
|
||||
"ToolMode",
|
||||
"ToolResultCompactionStrategy",
|
||||
"ToolTypes",
|
||||
"TruncationStrategy",
|
||||
"TypeCompatibilityError",
|
||||
"UpdateT",
|
||||
"UsageDetails",
|
||||
"UserInputRequiredException",
|
||||
"ValidationTypeEnum",
|
||||
"Workflow",
|
||||
"WorkflowAgent",
|
||||
"WorkflowBuilder",
|
||||
"WorkflowCheckpoint",
|
||||
"WorkflowCheckpointException",
|
||||
"WorkflowContext",
|
||||
"WorkflowConvergenceException",
|
||||
"WorkflowErrorDetails",
|
||||
"WorkflowEvent",
|
||||
"WorkflowEventSource",
|
||||
"WorkflowEventType",
|
||||
"WorkflowException",
|
||||
"WorkflowExecutor",
|
||||
"WorkflowMessage",
|
||||
"WorkflowRunResult",
|
||||
"WorkflowRunState",
|
||||
"WorkflowRunnerException",
|
||||
"WorkflowValidationError",
|
||||
"WorkflowViz",
|
||||
"__version__",
|
||||
"add_usage_details",
|
||||
"agent_middleware",
|
||||
"annotate_message_groups",
|
||||
"apply_compaction",
|
||||
"background_tasks_running",
|
||||
"background_tasks_running_message",
|
||||
"chat_middleware",
|
||||
"create_always_approve_tool_response",
|
||||
"create_always_approve_tool_with_arguments_response",
|
||||
"create_edge_runner",
|
||||
"create_harness_agent",
|
||||
"detect_media_type_from_base64",
|
||||
"enqueue_messages",
|
||||
"evaluate_agent",
|
||||
"evaluate_workflow",
|
||||
"evaluator",
|
||||
"executor",
|
||||
"function_middleware",
|
||||
"get_agent_mode",
|
||||
"get_run_context",
|
||||
"handler",
|
||||
"included_messages",
|
||||
"included_token_count",
|
||||
"keyword_check",
|
||||
"load_settings",
|
||||
"map_chat_to_agent_update",
|
||||
"merge_chat_options",
|
||||
"normalize_function_invocation_configuration",
|
||||
"normalize_messages",
|
||||
"normalize_tools",
|
||||
"prepend_agent_framework_to_user_agent",
|
||||
"prepend_instructions_to_messages",
|
||||
"register_state_type",
|
||||
"resolve_agent_id",
|
||||
"response_handler",
|
||||
"set_agent_mode",
|
||||
"step",
|
||||
"todos_remaining",
|
||||
"todos_remaining_message",
|
||||
"tool",
|
||||
"tool_call_args_match",
|
||||
"tool_called_check",
|
||||
"tool_calls_present",
|
||||
"validate_chat_options",
|
||||
"validate_tool_mode",
|
||||
"validate_tools",
|
||||
"validate_workflow_graph",
|
||||
"workflow",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
"""Lazily resolve public names exported from ``agent_framework``."""
|
||||
if module_name := _LAZY_EXPORTS.get(name):
|
||||
value = getattr(importlib.import_module(module_name, __name__), name)
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
"""Return public names for interactive discovery."""
|
||||
return sorted(set(globals()) | set(__all__))
|
||||
@@ -0,0 +1,604 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import Final
|
||||
|
||||
__version__: Final[str]
|
||||
|
||||
from ._agents import Agent, BaseAgent, RawAgent, SupportsAgentRun
|
||||
from ._clients import (
|
||||
BaseChatClient,
|
||||
BaseEmbeddingClient,
|
||||
SupportsChatGetResponse,
|
||||
SupportsCodeInterpreterTool,
|
||||
SupportsFileSearchTool,
|
||||
SupportsGetEmbeddings,
|
||||
SupportsImageGenerationTool,
|
||||
SupportsMCPTool,
|
||||
SupportsShellTool,
|
||||
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,
|
||||
ContextWindowCompactionStrategy,
|
||||
SelectiveToolCallCompactionStrategy,
|
||||
SlidingWindowStrategy,
|
||||
SummarizationStrategy,
|
||||
TokenBudgetComposedStrategy,
|
||||
TokenizerProtocol,
|
||||
ToolResultCompactionStrategy,
|
||||
TruncationStrategy,
|
||||
annotate_message_groups,
|
||||
apply_compaction,
|
||||
included_messages,
|
||||
included_token_count,
|
||||
)
|
||||
from ._evaluation import (
|
||||
AgentEvalConverter,
|
||||
CheckResult,
|
||||
ConversationSplit,
|
||||
ConversationSplitter,
|
||||
EvalItem,
|
||||
EvalItemResult,
|
||||
EvalNotPassedError,
|
||||
EvalResults,
|
||||
EvalScoreResult,
|
||||
Evaluator,
|
||||
ExpectedToolCall,
|
||||
LocalEvaluator,
|
||||
RubricScore,
|
||||
evaluate_agent,
|
||||
evaluate_workflow,
|
||||
evaluator,
|
||||
keyword_check,
|
||||
tool_call_args_match,
|
||||
tool_called_check,
|
||||
tool_calls_present,
|
||||
)
|
||||
from ._feature_stage import ExperimentalFeature, ReleaseCandidateFeature
|
||||
from ._harness._agent import DEFAULT_HARNESS_INSTRUCTIONS, create_harness_agent
|
||||
from ._harness._background_agents import (
|
||||
DEFAULT_BACKGROUND_AGENTS_SOURCE_ID,
|
||||
BackgroundAgentsProvider,
|
||||
BackgroundTaskInfo,
|
||||
BackgroundTaskStatus,
|
||||
)
|
||||
from ._harness._file_access import (
|
||||
DEFAULT_FILE_ACCESS_INSTRUCTIONS,
|
||||
DEFAULT_FILE_ACCESS_SOURCE_ID,
|
||||
AgentFileStore,
|
||||
FileAccessProvider,
|
||||
FileSearchMatch,
|
||||
FileSearchResult,
|
||||
FileStoreEntry,
|
||||
FileSystemAgentFileStore,
|
||||
InMemoryAgentFileStore,
|
||||
)
|
||||
from ._harness._file_memory import DEFAULT_FILE_MEMORY_INSTRUCTIONS, DEFAULT_FILE_MEMORY_SOURCE_ID, FileMemoryProvider
|
||||
from ._harness._loop import (
|
||||
AgentLoopMiddleware,
|
||||
JudgeVerdict,
|
||||
background_tasks_running,
|
||||
background_tasks_running_message,
|
||||
todos_remaining,
|
||||
todos_remaining_message,
|
||||
)
|
||||
from ._harness._memory import (
|
||||
DEFAULT_MEMORY_SOURCE_ID,
|
||||
MemoryContextProvider,
|
||||
MemoryFileStore,
|
||||
MemoryIndexEntry,
|
||||
MemoryStore,
|
||||
MemoryTopicRecord,
|
||||
)
|
||||
from ._harness._mode import DEFAULT_MODE_SOURCE_ID, AgentModeProvider, get_agent_mode, set_agent_mode
|
||||
from ._harness._todo import (
|
||||
DEFAULT_TODO_SOURCE_ID,
|
||||
TodoFileStore,
|
||||
TodoInput,
|
||||
TodoItem,
|
||||
TodoProvider,
|
||||
TodoSessionStore,
|
||||
TodoStore,
|
||||
)
|
||||
from ._harness._tool_approval import (
|
||||
DEFAULT_TOOL_APPROVAL_SOURCE_ID,
|
||||
ToolApprovalMiddleware,
|
||||
ToolApprovalRule,
|
||||
ToolApprovalRuleCallback,
|
||||
ToolApprovalState,
|
||||
create_always_approve_tool_response,
|
||||
create_always_approve_tool_with_arguments_response,
|
||||
)
|
||||
from ._mcp import (
|
||||
MCPStdioTool,
|
||||
MCPStreamableHTTPTool,
|
||||
MCPTaskOptions,
|
||||
MCPWebsocketTool,
|
||||
SamplingApprovalCallback,
|
||||
)
|
||||
from ._middleware import (
|
||||
AgentContext,
|
||||
AgentMiddleware,
|
||||
AgentMiddlewareLayer,
|
||||
AgentMiddlewareTypes,
|
||||
ChatAndFunctionMiddlewareTypes,
|
||||
ChatContext,
|
||||
ChatMiddleware,
|
||||
ChatMiddlewareLayer,
|
||||
ChatMiddlewareTypes,
|
||||
FunctionInvocationContext,
|
||||
FunctionMiddleware,
|
||||
FunctionMiddlewareTypes,
|
||||
MiddlewareTermination,
|
||||
MiddlewareType,
|
||||
MiddlewareTypes,
|
||||
agent_middleware,
|
||||
chat_middleware,
|
||||
function_middleware,
|
||||
)
|
||||
from ._sessions import (
|
||||
MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY,
|
||||
AgentSession,
|
||||
ContextProvider,
|
||||
FileHistoryProvider,
|
||||
HistoryProvider,
|
||||
InMemoryHistoryProvider,
|
||||
MessageInjectionMiddleware,
|
||||
ServiceSessionId,
|
||||
SessionContext,
|
||||
enqueue_messages,
|
||||
register_state_type,
|
||||
)
|
||||
from ._settings import SecretString, load_settings
|
||||
from ._skills import (
|
||||
AggregatingSkillsSource,
|
||||
CachingSkillsSource,
|
||||
ClassSkill,
|
||||
DeduplicatingSkillsSource,
|
||||
DelegatingSkillsSource,
|
||||
FileSkill,
|
||||
FileSkillScript,
|
||||
FileSkillsSource,
|
||||
FilteringSkillsSource,
|
||||
InlineSkill,
|
||||
InlineSkillResource,
|
||||
InlineSkillScript,
|
||||
InMemorySkillsSource,
|
||||
MCPSkill,
|
||||
MCPSkillResource,
|
||||
MCPSkillsSource,
|
||||
Skill,
|
||||
SkillFrontmatter,
|
||||
SkillResource,
|
||||
SkillScript,
|
||||
SkillScriptArgumentParser,
|
||||
SkillScriptRunner,
|
||||
SkillsProvider,
|
||||
SkillsSource,
|
||||
SkillsSourceContext,
|
||||
)
|
||||
from ._telemetry import (
|
||||
AGENT_FRAMEWORK_USER_AGENT,
|
||||
APP_INFO,
|
||||
USER_AGENT_KEY,
|
||||
USER_AGENT_TELEMETRY_DISABLED_ENV_VAR,
|
||||
prepend_agent_framework_to_user_agent,
|
||||
)
|
||||
from ._tools import (
|
||||
SKIP_PARSING,
|
||||
FunctionInvocationConfiguration,
|
||||
FunctionInvocationLayer,
|
||||
FunctionTool,
|
||||
ToolTypes,
|
||||
normalize_function_invocation_configuration,
|
||||
tool,
|
||||
)
|
||||
from ._types import (
|
||||
AgentResponse,
|
||||
AgentResponseUpdate,
|
||||
AgentRunInputs,
|
||||
Annotation,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Content,
|
||||
ContinuationToken,
|
||||
Embedding,
|
||||
EmbeddingGenerationOptions,
|
||||
EmbeddingInputT,
|
||||
EmbeddingT,
|
||||
FinalT,
|
||||
FinishReason,
|
||||
FinishReasonLiteral,
|
||||
GeneratedEmbeddings,
|
||||
Message,
|
||||
OuterFinalT,
|
||||
OuterUpdateT,
|
||||
ResponseStream,
|
||||
Role,
|
||||
RoleLiteral,
|
||||
TextSpanRegion,
|
||||
ToolMode,
|
||||
UpdateT,
|
||||
UsageDetails,
|
||||
add_usage_details,
|
||||
detect_media_type_from_base64,
|
||||
map_chat_to_agent_update,
|
||||
merge_chat_options,
|
||||
normalize_messages,
|
||||
normalize_tools,
|
||||
prepend_instructions_to_messages,
|
||||
validate_chat_options,
|
||||
validate_tool_mode,
|
||||
validate_tools,
|
||||
)
|
||||
from ._workflows._agent import WorkflowAgent
|
||||
from ._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
|
||||
from ._workflows._agent_utils import resolve_agent_id
|
||||
from ._workflows._checkpoint import (
|
||||
CheckpointID,
|
||||
CheckpointStorage,
|
||||
FileCheckpointStorage,
|
||||
InMemoryCheckpointStorage,
|
||||
WorkflowCheckpoint,
|
||||
)
|
||||
from ._workflows._const import DEFAULT_MAX_ITERATIONS
|
||||
from ._workflows._edge import (
|
||||
Case,
|
||||
Default,
|
||||
Edge,
|
||||
EdgeCondition,
|
||||
FanInEdgeGroup,
|
||||
FanOutEdgeGroup,
|
||||
SingleEdgeGroup,
|
||||
SwitchCaseEdgeGroup,
|
||||
SwitchCaseEdgeGroupCase,
|
||||
SwitchCaseEdgeGroupDefault,
|
||||
)
|
||||
from ._workflows._edge_runner import create_edge_runner
|
||||
from ._workflows._events import (
|
||||
WorkflowErrorDetails,
|
||||
WorkflowEvent,
|
||||
WorkflowEventSource,
|
||||
WorkflowEventType,
|
||||
WorkflowRunState,
|
||||
)
|
||||
from ._workflows._executor import Executor, handler
|
||||
from ._workflows._function_executor import FunctionExecutor, executor
|
||||
from ._workflows._functional import (
|
||||
FunctionalWorkflow,
|
||||
FunctionalWorkflowAgent,
|
||||
RunContext,
|
||||
StepWrapper,
|
||||
get_run_context,
|
||||
step,
|
||||
workflow,
|
||||
)
|
||||
from ._workflows._request_info_mixin import response_handler
|
||||
from ._workflows._runner import Runner
|
||||
from ._workflows._runner_context import InProcRunnerContext, RunnerContext, WorkflowMessage
|
||||
from ._workflows._validation import (
|
||||
EdgeDuplicationError,
|
||||
GraphConnectivityError,
|
||||
TypeCompatibilityError,
|
||||
ValidationTypeEnum,
|
||||
WorkflowValidationError,
|
||||
validate_workflow_graph,
|
||||
)
|
||||
from ._workflows._viz import WorkflowViz
|
||||
from ._workflows._workflow import Workflow, WorkflowRunResult
|
||||
from ._workflows._workflow_builder import WorkflowBuilder
|
||||
from ._workflows._workflow_context import WorkflowContext
|
||||
from ._workflows._workflow_executor import SubWorkflowRequestMessage, SubWorkflowResponseMessage, WorkflowExecutor
|
||||
from .exceptions import (
|
||||
AgentFrameworkException,
|
||||
MiddlewareException,
|
||||
UserInputRequiredException,
|
||||
WorkflowCheckpointException,
|
||||
WorkflowConvergenceException,
|
||||
WorkflowException,
|
||||
WorkflowRunnerException,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AGENT_FRAMEWORK_USER_AGENT",
|
||||
"APP_INFO",
|
||||
"COMPACTION_STATE_KEY",
|
||||
"DEFAULT_BACKGROUND_AGENTS_SOURCE_ID",
|
||||
"DEFAULT_FILE_ACCESS_INSTRUCTIONS",
|
||||
"DEFAULT_FILE_ACCESS_SOURCE_ID",
|
||||
"DEFAULT_FILE_MEMORY_INSTRUCTIONS",
|
||||
"DEFAULT_FILE_MEMORY_SOURCE_ID",
|
||||
"DEFAULT_HARNESS_INSTRUCTIONS",
|
||||
"DEFAULT_MAX_ITERATIONS",
|
||||
"DEFAULT_MEMORY_SOURCE_ID",
|
||||
"DEFAULT_MODE_SOURCE_ID",
|
||||
"DEFAULT_TODO_SOURCE_ID",
|
||||
"DEFAULT_TOOL_APPROVAL_SOURCE_ID",
|
||||
"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",
|
||||
"MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY",
|
||||
"SKIP_PARSING",
|
||||
"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",
|
||||
"AgentContext",
|
||||
"AgentEvalConverter",
|
||||
"AgentExecutor",
|
||||
"AgentExecutorRequest",
|
||||
"AgentExecutorResponse",
|
||||
"AgentFileStore",
|
||||
"AgentFrameworkException",
|
||||
"AgentLoopMiddleware",
|
||||
"AgentMiddleware",
|
||||
"AgentMiddlewareLayer",
|
||||
"AgentMiddlewareTypes",
|
||||
"AgentModeProvider",
|
||||
"AgentResponse",
|
||||
"AgentResponseUpdate",
|
||||
"AgentRunInputs",
|
||||
"AgentSession",
|
||||
"AggregatingSkillsSource",
|
||||
"Annotation",
|
||||
"BackgroundAgentsProvider",
|
||||
"BackgroundTaskInfo",
|
||||
"BackgroundTaskStatus",
|
||||
"BaseAgent",
|
||||
"BaseChatClient",
|
||||
"BaseEmbeddingClient",
|
||||
"CachingSkillsSource",
|
||||
"Case",
|
||||
"CharacterEstimatorTokenizer",
|
||||
"ChatAndFunctionMiddlewareTypes",
|
||||
"ChatContext",
|
||||
"ChatMiddleware",
|
||||
"ChatMiddlewareLayer",
|
||||
"ChatMiddlewareTypes",
|
||||
"ChatOptions",
|
||||
"ChatResponse",
|
||||
"ChatResponseUpdate",
|
||||
"CheckResult",
|
||||
"CheckpointID",
|
||||
"CheckpointStorage",
|
||||
"ClassSkill",
|
||||
"CompactionProvider",
|
||||
"CompactionStrategy",
|
||||
"Content",
|
||||
"ContextProvider",
|
||||
"ContextWindowCompactionStrategy",
|
||||
"ContinuationToken",
|
||||
"ConversationSplit",
|
||||
"ConversationSplitter",
|
||||
"DeduplicatingSkillsSource",
|
||||
"Default",
|
||||
"DelegatingSkillsSource",
|
||||
"Edge",
|
||||
"EdgeCondition",
|
||||
"EdgeDuplicationError",
|
||||
"Embedding",
|
||||
"EmbeddingGenerationOptions",
|
||||
"EmbeddingInputT",
|
||||
"EmbeddingT",
|
||||
"EvalItem",
|
||||
"EvalItemResult",
|
||||
"EvalNotPassedError",
|
||||
"EvalResults",
|
||||
"EvalScoreResult",
|
||||
"Evaluator",
|
||||
"Executor",
|
||||
"ExpectedToolCall",
|
||||
"ExperimentalFeature",
|
||||
"FanInEdgeGroup",
|
||||
"FanOutEdgeGroup",
|
||||
"FileAccessProvider",
|
||||
"FileCheckpointStorage",
|
||||
"FileHistoryProvider",
|
||||
"FileMemoryProvider",
|
||||
"FileSearchMatch",
|
||||
"FileSearchResult",
|
||||
"FileSkill",
|
||||
"FileSkillScript",
|
||||
"FileSkillsSource",
|
||||
"FileStoreEntry",
|
||||
"FileSystemAgentFileStore",
|
||||
"FilteringSkillsSource",
|
||||
"FinalT",
|
||||
"FinishReason",
|
||||
"FinishReasonLiteral",
|
||||
"FunctionExecutor",
|
||||
"FunctionInvocationConfiguration",
|
||||
"FunctionInvocationContext",
|
||||
"FunctionInvocationLayer",
|
||||
"FunctionMiddleware",
|
||||
"FunctionMiddlewareTypes",
|
||||
"FunctionTool",
|
||||
"FunctionalWorkflow",
|
||||
"FunctionalWorkflowAgent",
|
||||
"GeneratedEmbeddings",
|
||||
"GraphConnectivityError",
|
||||
"HistoryProvider",
|
||||
"InMemoryAgentFileStore",
|
||||
"InMemoryCheckpointStorage",
|
||||
"InMemoryHistoryProvider",
|
||||
"InMemorySkillsSource",
|
||||
"InProcRunnerContext",
|
||||
"InlineSkill",
|
||||
"InlineSkillResource",
|
||||
"InlineSkillScript",
|
||||
"JudgeVerdict",
|
||||
"LocalEvaluator",
|
||||
"MCPSkill",
|
||||
"MCPSkillResource",
|
||||
"MCPSkillsSource",
|
||||
"MCPStdioTool",
|
||||
"MCPStreamableHTTPTool",
|
||||
"MCPTaskOptions",
|
||||
"MCPWebsocketTool",
|
||||
"MemoryContextProvider",
|
||||
"MemoryFileStore",
|
||||
"MemoryIndexEntry",
|
||||
"MemoryStore",
|
||||
"MemoryTopicRecord",
|
||||
"Message",
|
||||
"MessageInjectionMiddleware",
|
||||
"MiddlewareException",
|
||||
"MiddlewareTermination",
|
||||
"MiddlewareType",
|
||||
"MiddlewareTypes",
|
||||
"OuterFinalT",
|
||||
"OuterUpdateT",
|
||||
"RawAgent",
|
||||
"ReleaseCandidateFeature",
|
||||
"ResponseStream",
|
||||
"Role",
|
||||
"RoleLiteral",
|
||||
"RubricScore",
|
||||
"RunContext",
|
||||
"Runner",
|
||||
"RunnerContext",
|
||||
"SamplingApprovalCallback",
|
||||
"SecretString",
|
||||
"SelectiveToolCallCompactionStrategy",
|
||||
"ServiceSessionId",
|
||||
"SessionContext",
|
||||
"SingleEdgeGroup",
|
||||
"Skill",
|
||||
"SkillFrontmatter",
|
||||
"SkillResource",
|
||||
"SkillScript",
|
||||
"SkillScriptArgumentParser",
|
||||
"SkillScriptRunner",
|
||||
"SkillsProvider",
|
||||
"SkillsSource",
|
||||
"SkillsSourceContext",
|
||||
"SlidingWindowStrategy",
|
||||
"StepWrapper",
|
||||
"SubWorkflowRequestMessage",
|
||||
"SubWorkflowResponseMessage",
|
||||
"SummarizationStrategy",
|
||||
"SupportsAgentRun",
|
||||
"SupportsChatGetResponse",
|
||||
"SupportsCodeInterpreterTool",
|
||||
"SupportsFileSearchTool",
|
||||
"SupportsGetEmbeddings",
|
||||
"SupportsImageGenerationTool",
|
||||
"SupportsMCPTool",
|
||||
"SupportsShellTool",
|
||||
"SupportsWebSearchTool",
|
||||
"SwitchCaseEdgeGroup",
|
||||
"SwitchCaseEdgeGroupCase",
|
||||
"SwitchCaseEdgeGroupDefault",
|
||||
"TextSpanRegion",
|
||||
"TodoFileStore",
|
||||
"TodoInput",
|
||||
"TodoItem",
|
||||
"TodoProvider",
|
||||
"TodoSessionStore",
|
||||
"TodoStore",
|
||||
"TokenBudgetComposedStrategy",
|
||||
"TokenizerProtocol",
|
||||
"ToolApprovalMiddleware",
|
||||
"ToolApprovalRule",
|
||||
"ToolApprovalRuleCallback",
|
||||
"ToolApprovalState",
|
||||
"ToolMode",
|
||||
"ToolResultCompactionStrategy",
|
||||
"ToolTypes",
|
||||
"TruncationStrategy",
|
||||
"TypeCompatibilityError",
|
||||
"UpdateT",
|
||||
"UsageDetails",
|
||||
"UserInputRequiredException",
|
||||
"ValidationTypeEnum",
|
||||
"Workflow",
|
||||
"WorkflowAgent",
|
||||
"WorkflowBuilder",
|
||||
"WorkflowCheckpoint",
|
||||
"WorkflowCheckpointException",
|
||||
"WorkflowContext",
|
||||
"WorkflowConvergenceException",
|
||||
"WorkflowErrorDetails",
|
||||
"WorkflowEvent",
|
||||
"WorkflowEventSource",
|
||||
"WorkflowEventType",
|
||||
"WorkflowException",
|
||||
"WorkflowExecutor",
|
||||
"WorkflowMessage",
|
||||
"WorkflowRunResult",
|
||||
"WorkflowRunState",
|
||||
"WorkflowRunnerException",
|
||||
"WorkflowValidationError",
|
||||
"WorkflowViz",
|
||||
"__version__",
|
||||
"add_usage_details",
|
||||
"agent_middleware",
|
||||
"annotate_message_groups",
|
||||
"apply_compaction",
|
||||
"background_tasks_running",
|
||||
"background_tasks_running_message",
|
||||
"chat_middleware",
|
||||
"create_always_approve_tool_response",
|
||||
"create_always_approve_tool_with_arguments_response",
|
||||
"create_edge_runner",
|
||||
"create_harness_agent",
|
||||
"detect_media_type_from_base64",
|
||||
"enqueue_messages",
|
||||
"evaluate_agent",
|
||||
"evaluate_workflow",
|
||||
"evaluator",
|
||||
"executor",
|
||||
"function_middleware",
|
||||
"get_agent_mode",
|
||||
"get_run_context",
|
||||
"handler",
|
||||
"included_messages",
|
||||
"included_token_count",
|
||||
"keyword_check",
|
||||
"load_settings",
|
||||
"map_chat_to_agent_update",
|
||||
"merge_chat_options",
|
||||
"normalize_function_invocation_configuration",
|
||||
"normalize_messages",
|
||||
"normalize_tools",
|
||||
"prepend_agent_framework_to_user_agent",
|
||||
"prepend_instructions_to_messages",
|
||||
"register_state_type",
|
||||
"resolve_agent_id",
|
||||
"response_handler",
|
||||
"set_agent_mode",
|
||||
"step",
|
||||
"todos_remaining",
|
||||
"todos_remaining_message",
|
||||
"tool",
|
||||
"tool_call_args_match",
|
||||
"tool_called_check",
|
||||
"tool_calls_present",
|
||||
"validate_chat_options",
|
||||
"validate_tool_mode",
|
||||
"validate_tools",
|
||||
"validate_workflow_graph",
|
||||
"workflow",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,110 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import textwrap
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Any
|
||||
|
||||
_GOOGLE_SECTION_HEADERS = (
|
||||
"Args:",
|
||||
"Keyword Args:",
|
||||
"Attributes:",
|
||||
"Returns:",
|
||||
"Raises:",
|
||||
"Examples:",
|
||||
"Note:",
|
||||
"Notes:",
|
||||
"Warning:",
|
||||
"Warnings:",
|
||||
)
|
||||
|
||||
|
||||
def _find_section_index(lines: list[str], header: str) -> int | None:
|
||||
for index, line in enumerate(lines):
|
||||
if line == header:
|
||||
return index
|
||||
return None
|
||||
|
||||
|
||||
def _find_next_section_index(lines: list[str], start: int) -> int:
|
||||
for index in range(start, len(lines)):
|
||||
if lines[index] in _GOOGLE_SECTION_HEADERS:
|
||||
return index
|
||||
return len(lines)
|
||||
|
||||
|
||||
def _format_keyword_arg_lines(extra_keyword_args: Mapping[str, str]) -> list[str]:
|
||||
formatted_lines: list[str] = []
|
||||
for name, description in extra_keyword_args.items():
|
||||
description_lines = inspect.cleandoc(description).splitlines()
|
||||
if not description_lines:
|
||||
formatted_lines.append(f" {name}:")
|
||||
continue
|
||||
formatted_lines.append(f" {name}: {description_lines[0]}")
|
||||
formatted_lines.extend(f" {line}" for line in description_lines[1:])
|
||||
return formatted_lines
|
||||
|
||||
|
||||
def insert_docstring_block(docstring: str | None, *, block: str) -> str | None:
|
||||
"""Insert a preformatted block before the first Google-style section."""
|
||||
cleaned_block = textwrap.dedent(block).strip()
|
||||
if not cleaned_block:
|
||||
return docstring
|
||||
if not docstring:
|
||||
return cleaned_block
|
||||
|
||||
lines = inspect.cleandoc(docstring).splitlines()
|
||||
block_lines = cleaned_block.splitlines()
|
||||
insert_index = _find_next_section_index(lines, 0)
|
||||
|
||||
insertion: list[str] = []
|
||||
if insert_index > 0 and lines[insert_index - 1] != "":
|
||||
insertion.append("")
|
||||
insertion.extend(block_lines)
|
||||
if insert_index < len(lines) and insertion[-1] != "":
|
||||
insertion.append("")
|
||||
|
||||
lines[insert_index:insert_index] = insertion
|
||||
return "\n".join(lines).rstrip()
|
||||
|
||||
|
||||
def build_layered_docstring(
|
||||
source: Callable[..., Any],
|
||||
*,
|
||||
extra_keyword_args: Mapping[str, str] | None = None,
|
||||
) -> str | None:
|
||||
"""Build a Google-style docstring from a lower-layer implementation."""
|
||||
docstring = inspect.getdoc(source)
|
||||
if not docstring:
|
||||
return None
|
||||
if not extra_keyword_args:
|
||||
return docstring
|
||||
|
||||
lines = docstring.splitlines()
|
||||
formatted_keyword_arg_lines = _format_keyword_arg_lines(extra_keyword_args)
|
||||
keyword_args_index = _find_section_index(lines, "Keyword Args:")
|
||||
|
||||
if keyword_args_index is None:
|
||||
args_index = _find_section_index(lines, "Args:")
|
||||
if args_index is not None:
|
||||
insert_index = _find_next_section_index(lines, args_index + 1)
|
||||
else:
|
||||
insert_index = _find_next_section_index(lines, 0)
|
||||
lines[insert_index:insert_index] = ["", "Keyword Args:", *formatted_keyword_arg_lines]
|
||||
return "\n".join(lines).rstrip()
|
||||
|
||||
insert_index = _find_next_section_index(lines, keyword_args_index + 1)
|
||||
lines[insert_index:insert_index] = formatted_keyword_arg_lines
|
||||
return "\n".join(lines).rstrip()
|
||||
|
||||
|
||||
def apply_layered_docstring(
|
||||
target: Callable[..., Any],
|
||||
source: Callable[..., Any],
|
||||
*,
|
||||
extra_keyword_args: Mapping[str, str] | None = None,
|
||||
) -> None:
|
||||
"""Copy a lower-layer docstring onto a wrapper and extend it when needed."""
|
||||
target.__doc__ = build_layered_docstring(source, extra_keyword_args=extra_keyword_args)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,403 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import asyncio.coroutines
|
||||
import contextlib
|
||||
import functools
|
||||
import inspect
|
||||
import os
|
||||
import sys
|
||||
import typing
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from enum import Enum
|
||||
from types import MethodType
|
||||
from typing import Any, Literal, TypeVar, cast
|
||||
|
||||
from ._docstrings import insert_docstring_block
|
||||
|
||||
FeatureStageT = TypeVar("FeatureStageT", bound=Callable[..., Any])
|
||||
|
||||
FeatureStageName = Literal["experimental", "release_candidate"]
|
||||
|
||||
# Optional feature-stage metadata for warnings and best-effort introspection.
|
||||
_FEATURE_ID_ATTR = "__feature_id__"
|
||||
_FEATURE_STAGE_ATTR = "__feature_stage__"
|
||||
_WARNED_FEATURES: set[tuple[type[Warning], str]] = set()
|
||||
_EXPERIMENTAL_DOCSTRING = """\
|
||||
.. warning:: Experimental
|
||||
|
||||
This API is experimental and subject to change or removal
|
||||
in future versions without notice.
|
||||
"""
|
||||
_RELEASE_CANDIDATE_DOCSTRING = """\
|
||||
.. note:: Release candidate
|
||||
|
||||
This API is in release-candidate stage and may receive
|
||||
minor refinements before it is considered generally available.
|
||||
"""
|
||||
|
||||
|
||||
class ExperimentalFeature(str, Enum):
|
||||
"""Current experimental feature IDs.
|
||||
|
||||
This enum is a stage-scoped inventory, not a stable introspection surface.
|
||||
Members may move or be removed as features advance. The `__feature_id__`
|
||||
attribute is also optional stage metadata and may disappear when a feature
|
||||
is released, so consumer code should use `getattr(...)` rather than relying
|
||||
on enum membership or attribute presence over time.
|
||||
"""
|
||||
|
||||
DECLARATIVE_AGENTS = "DECLARATIVE_AGENTS"
|
||||
EVALS = "EVALS"
|
||||
FILE_HISTORY = "FILE_HISTORY"
|
||||
FIDES = "FIDES"
|
||||
FOUNDRY_TOOLS = "FOUNDRY_TOOLS"
|
||||
FOUNDRY_PREVIEW_TOOLS = "FOUNDRY_PREVIEW_TOOLS"
|
||||
FUNCTIONAL_WORKFLOWS = "FUNCTIONAL_WORKFLOWS"
|
||||
HARNESS = "HARNESS"
|
||||
MCP_LONG_RUNNING_TASKS = "MCP_LONG_RUNNING_TASKS"
|
||||
MCP_SKILLS = "MCP_SKILLS"
|
||||
PROGRESSIVE_TOOLS = "PROGRESSIVE_TOOLS"
|
||||
TO_PROMPT_AGENT = "TO_PROMPT_AGENT"
|
||||
|
||||
|
||||
class ReleaseCandidateFeature(str, Enum):
|
||||
"""Current release-candidate feature IDs.
|
||||
|
||||
This enum is a stage-scoped inventory, not a stable introspection surface.
|
||||
Members may move or be removed as features advance. The `__feature_id__`
|
||||
attribute is also optional stage metadata and may disappear when a feature
|
||||
is released, so consumer code should use `getattr(...)` rather than relying
|
||||
on enum membership or attribute presence over time.
|
||||
"""
|
||||
|
||||
|
||||
class FeatureStageWarning(FutureWarning):
|
||||
"""Base warning category for staged APIs."""
|
||||
|
||||
|
||||
class ExperimentalWarning(FeatureStageWarning):
|
||||
"""Warning emitted when an experimental API is used."""
|
||||
|
||||
|
||||
# Sentinel attribute used to detect (and reuse) a formatter we've already
|
||||
# installed. This lets the install be idempotent across re-imports / reloads
|
||||
# and keeps a stable reference to the previous formatter for testing or
|
||||
# external restoration via ``warnings.formatwarning = original``.
|
||||
_FEATURE_STAGE_FORMATTER_MARKER = "__feature_stage_formatter__"
|
||||
|
||||
|
||||
def _install_feature_stage_formatter() -> None:
|
||||
"""Install a single-line formatter for FeatureStageWarning categories.
|
||||
|
||||
The stdlib default formatter emits two lines (header + source snippet)
|
||||
which is noisy for our warnings — the offending class/function name is
|
||||
already in the message, so a one-line ``file:lineno: Category: message``
|
||||
is enough. Other warning categories are delegated to the previous
|
||||
formatter so we never change behaviour for unrelated warnings.
|
||||
|
||||
The install is idempotent: if a formatter installed by this module is
|
||||
already in place, we leave it alone so re-imports (and any third-party
|
||||
formatter wrapped on top of ours) don't get wrapped multiple times.
|
||||
"""
|
||||
current = warnings.formatwarning
|
||||
if getattr(current, _FEATURE_STAGE_FORMATTER_MARKER, False):
|
||||
return
|
||||
|
||||
def _formatwarning(
|
||||
message: Warning | str,
|
||||
category: type[Warning],
|
||||
filename: str,
|
||||
lineno: int,
|
||||
line: str | None = None,
|
||||
) -> str:
|
||||
if issubclass(category, FeatureStageWarning):
|
||||
return f"{filename}:{lineno}: {category.__name__}: {message}\n"
|
||||
return current(message, category, filename, lineno, line)
|
||||
|
||||
setattr(_formatwarning, _FEATURE_STAGE_FORMATTER_MARKER, True)
|
||||
# Keep a reference to the wrapped formatter so callers (tests, embedders)
|
||||
# can restore the previous behaviour if they need to.
|
||||
_formatwarning.__wrapped__ = current # type: ignore[attr-defined]
|
||||
warnings.formatwarning = _formatwarning
|
||||
|
||||
|
||||
_install_feature_stage_formatter()
|
||||
|
||||
|
||||
def _normalize_feature_id(feature_id: str | Enum) -> str:
|
||||
return str(feature_id.value if isinstance(feature_id, Enum) else feature_id)
|
||||
|
||||
|
||||
def _get_object_name(obj: Any) -> str:
|
||||
return str(getattr(obj, "__qualname__", getattr(obj, "__name__", type(obj).__name__)))
|
||||
|
||||
|
||||
def _get_descriptor_callable(obj: Any) -> Callable[..., Any]:
|
||||
return cast(Callable[..., Any], obj.__func__)
|
||||
|
||||
|
||||
def _is_protocol_class(obj: Any) -> bool:
|
||||
return isinstance(obj, type) and bool(getattr(obj, "_is_protocol", False))
|
||||
|
||||
|
||||
def _build_stage_warning_message(*, stage: FeatureStageName, feature_id: str, object_name: str) -> str:
|
||||
if stage == "experimental":
|
||||
return (
|
||||
f"[{feature_id}] {object_name} is experimental and may change or be removed in future versions "
|
||||
"without notice."
|
||||
)
|
||||
|
||||
return (
|
||||
f"[{feature_id}] {object_name} is in release-candidate stage and may receive minor refinements before it is "
|
||||
"considered generally available."
|
||||
)
|
||||
|
||||
|
||||
def _set_feature_stage_metadata(obj: Any, *, stage: FeatureStageName, feature_id: str) -> None:
|
||||
setattr(obj, _FEATURE_STAGE_ATTR, stage)
|
||||
setattr(obj, _FEATURE_ID_ATTR, feature_id)
|
||||
|
||||
|
||||
_INTERNAL_FRAME_FILE = os.path.normcase(__file__)
|
||||
# Module names whose frames we never want to surface as the caller. ``abc`` is
|
||||
# the big one (its ``__new__`` shows up as ``<frozen abc>:106`` for ABC-driven
|
||||
# subclass creation on modern CPython, so we cannot rely on filename matching).
|
||||
# ``functools``/``typing``/``contextlib`` are added because they often wrap our
|
||||
# decorators or appear in the metaclass call path.
|
||||
_INTERNAL_FRAME_MODULES: frozenset[str] = frozenset({
|
||||
abc.__name__,
|
||||
functools.__name__,
|
||||
typing.__name__,
|
||||
contextlib.__name__,
|
||||
})
|
||||
|
||||
|
||||
def _is_internal_frame(frame: Any) -> bool:
|
||||
if os.path.normcase(frame.f_code.co_filename) == _INTERNAL_FRAME_FILE:
|
||||
return True
|
||||
module_name = frame.f_globals.get("__name__", "")
|
||||
if module_name in _INTERNAL_FRAME_MODULES:
|
||||
return True
|
||||
# Submodules of the skipped stdlib packages (``typing.ext``, ``functools``
|
||||
# wrappers under ``concurrent.futures._base``, etc.) are also wrappers we
|
||||
# don't want to surface.
|
||||
return any(module_name.startswith(prefix + ".") for prefix in _INTERNAL_FRAME_MODULES)
|
||||
|
||||
|
||||
def _resolve_user_frame() -> tuple[str, int, str] | None:
|
||||
"""Resolve the user frame that triggered an experimental warning.
|
||||
|
||||
Walk the stack and return ``(filename, lineno, module_name)`` for the first
|
||||
frame outside this module and the wrapping/metaclass machinery.
|
||||
|
||||
Returns ``None`` if no such frame is found; callers fall back to plain
|
||||
``warnings.warn`` with a fixed stacklevel.
|
||||
"""
|
||||
# Frame objects participate in reference cycles (``frame -> f_locals ->
|
||||
# frame``) and can delay GC if held implicitly. Capture the user frame's
|
||||
# data into plain values inside the try, and explicitly delete the frame
|
||||
# references in finally so we never leak frames across this call. This
|
||||
# follows CPython's own guidance for code that uses ``inspect.currentframe``.
|
||||
frame = inspect.currentframe()
|
||||
candidate: Any = None
|
||||
try:
|
||||
if frame is None:
|
||||
return None
|
||||
# Skip _resolve_user_frame itself + the warn helper that called it.
|
||||
candidate = frame.f_back.f_back if frame.f_back and frame.f_back.f_back else None
|
||||
while candidate is not None:
|
||||
if not _is_internal_frame(candidate):
|
||||
return (
|
||||
candidate.f_code.co_filename,
|
||||
candidate.f_lineno,
|
||||
candidate.f_globals.get("__name__", "<unknown>"),
|
||||
)
|
||||
candidate = candidate.f_back
|
||||
return None
|
||||
finally:
|
||||
del frame, candidate
|
||||
|
||||
|
||||
def _warn_on_feature_use(
|
||||
*,
|
||||
stage: FeatureStageName,
|
||||
feature_id: str | Enum,
|
||||
object_name: str,
|
||||
category: type[Warning],
|
||||
) -> None:
|
||||
normalized_feature_id = _normalize_feature_id(feature_id)
|
||||
warning_key = (category, normalized_feature_id)
|
||||
if warning_key in _WARNED_FEATURES:
|
||||
return
|
||||
|
||||
message = _build_stage_warning_message(stage=stage, feature_id=normalized_feature_id, object_name=object_name)
|
||||
user_frame = _resolve_user_frame()
|
||||
if user_frame is None:
|
||||
# Last-resort fallback: emit at the immediate caller of this helper.
|
||||
warnings.warn(message, category=category, stacklevel=2)
|
||||
else:
|
||||
filename, lineno, module = user_frame
|
||||
warnings.warn_explicit(
|
||||
message,
|
||||
category=category,
|
||||
filename=filename,
|
||||
lineno=lineno,
|
||||
module=module,
|
||||
)
|
||||
_WARNED_FEATURES.add(warning_key)
|
||||
|
||||
|
||||
def _add_runtime_warning(
|
||||
obj: FeatureStageT,
|
||||
*,
|
||||
stage: FeatureStageName,
|
||||
feature_id: str,
|
||||
category: type[Warning],
|
||||
) -> FeatureStageT:
|
||||
object_name = _get_object_name(obj)
|
||||
|
||||
if isinstance(obj, type):
|
||||
experimental_class = cast(type[Any], obj)
|
||||
original_new: Any = experimental_class.__new__
|
||||
|
||||
@functools.wraps(original_new)
|
||||
def __new__(cls: type[Any], /, *args: Any, **kwargs: Any) -> Any:
|
||||
if cls is experimental_class:
|
||||
_warn_on_feature_use(
|
||||
stage=stage,
|
||||
feature_id=feature_id,
|
||||
object_name=object_name,
|
||||
category=category,
|
||||
)
|
||||
if original_new is not object.__new__:
|
||||
return original_new(cls, *args, **kwargs)
|
||||
if cls.__init__ is object.__init__ and (args or kwargs):
|
||||
raise TypeError(f"{cls.__name__}() takes no arguments")
|
||||
return original_new(cls)
|
||||
|
||||
experimental_class.__new__ = staticmethod(__new__)
|
||||
|
||||
original_init_subclass: Any = experimental_class.__init_subclass__
|
||||
if isinstance(original_init_subclass, MethodType):
|
||||
original_init_subclass_func = original_init_subclass.__func__
|
||||
|
||||
@functools.wraps(original_init_subclass_func)
|
||||
def bound_init_subclass_wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
_warn_on_feature_use(
|
||||
stage=stage,
|
||||
feature_id=feature_id,
|
||||
object_name=object_name,
|
||||
category=category,
|
||||
)
|
||||
return original_init_subclass_func(*args, **kwargs)
|
||||
|
||||
experimental_class.__init_subclass__ = classmethod(bound_init_subclass_wrapper) # type: ignore[assignment]
|
||||
else:
|
||||
|
||||
@functools.wraps(original_init_subclass)
|
||||
def init_subclass_wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
_warn_on_feature_use(
|
||||
stage=stage,
|
||||
feature_id=feature_id,
|
||||
object_name=object_name,
|
||||
category=category,
|
||||
)
|
||||
return original_init_subclass(*args, **kwargs)
|
||||
|
||||
experimental_class.__init_subclass__ = init_subclass_wrapper
|
||||
|
||||
return cast(FeatureStageT, experimental_class)
|
||||
|
||||
@functools.wraps(obj)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
_warn_on_feature_use(
|
||||
stage=stage,
|
||||
feature_id=feature_id,
|
||||
object_name=object_name,
|
||||
category=category,
|
||||
)
|
||||
return obj(*args, **kwargs)
|
||||
|
||||
if inspect.iscoroutinefunction(obj):
|
||||
if sys.version_info >= (3, 12):
|
||||
wrapper = inspect.markcoroutinefunction(wrapper)
|
||||
else:
|
||||
wrapper._is_coroutine = asyncio.coroutines._is_coroutine # type: ignore[attr-defined]
|
||||
|
||||
return cast(FeatureStageT, wrapper)
|
||||
|
||||
|
||||
def _feature_stage(
|
||||
*,
|
||||
stage: FeatureStageName,
|
||||
feature_id: str | Enum,
|
||||
docstring_block: str,
|
||||
warning_category: type[Warning] | None,
|
||||
) -> Callable[[FeatureStageT], FeatureStageT]:
|
||||
normalized_feature_id = _normalize_feature_id(feature_id)
|
||||
|
||||
def decorator(obj: FeatureStageT) -> FeatureStageT:
|
||||
descriptor_wrapper: Callable[[Any], Any] | None = None
|
||||
target: Any = obj
|
||||
|
||||
if isinstance(obj, staticmethod):
|
||||
descriptor_wrapper = staticmethod
|
||||
target = _get_descriptor_callable(obj)
|
||||
elif isinstance(obj, classmethod):
|
||||
descriptor_wrapper = classmethod
|
||||
target = _get_descriptor_callable(obj)
|
||||
|
||||
if not callable(target):
|
||||
raise TypeError(f"{stage} decorator can only be applied to classes and callables, not {obj!r}.")
|
||||
|
||||
is_protocol_class = _is_protocol_class(target)
|
||||
decorated: Any = target
|
||||
if warning_category is not None and not is_protocol_class:
|
||||
decorated = _add_runtime_warning(
|
||||
target,
|
||||
stage=stage,
|
||||
feature_id=normalized_feature_id,
|
||||
category=warning_category,
|
||||
)
|
||||
|
||||
updated_docstring = insert_docstring_block(decorated.__doc__, block=docstring_block)
|
||||
if updated_docstring is not None:
|
||||
decorated.__doc__ = updated_docstring
|
||||
|
||||
# runtime_checkable Protocol classes treat added class attributes as protocol members
|
||||
# on older Python versions, which breaks isinstance/issubclass checks.
|
||||
if not is_protocol_class:
|
||||
_set_feature_stage_metadata(decorated, stage=stage, feature_id=normalized_feature_id)
|
||||
if descriptor_wrapper is not None:
|
||||
return cast(FeatureStageT, descriptor_wrapper(decorated))
|
||||
|
||||
return cast(FeatureStageT, decorated)
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def experimental(*, feature_id: ExperimentalFeature) -> Callable[[FeatureStageT], FeatureStageT]:
|
||||
"""Mark a class or callable as experimental."""
|
||||
return _feature_stage(
|
||||
stage="experimental",
|
||||
feature_id=feature_id,
|
||||
docstring_block=_EXPERIMENTAL_DOCSTRING,
|
||||
warning_category=ExperimentalWarning,
|
||||
)
|
||||
|
||||
|
||||
def release_candidate(
|
||||
*,
|
||||
feature_id: ReleaseCandidateFeature,
|
||||
) -> Callable[[FeatureStageT], FeatureStageT]:
|
||||
"""Mark a class or callable as release-candidate."""
|
||||
return _feature_stage(
|
||||
stage="release_candidate",
|
||||
feature_id=feature_id,
|
||||
docstring_block=_RELEASE_CANDIDATE_DOCSTRING,
|
||||
warning_category=None,
|
||||
)
|
||||
@@ -0,0 +1,609 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Harness agent factory: a pre-configured bundled agent with batteries included.
|
||||
|
||||
This module provides :func:`create_harness_agent`, a factory function that assembles
|
||||
the full agent pipeline from a chat client, wiring up function invocation,
|
||||
per-service-call history persistence, compaction, and a rich set of default
|
||||
context providers (todo, mode, memory, skills).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Callable, Sequence
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, TypedDict
|
||||
|
||||
from .._agents import Agent, SupportsAgentRun
|
||||
from .._clients import SupportsShellTool, SupportsWebSearchTool
|
||||
from .._compaction import CompactionProvider, ContextWindowCompactionStrategy, ToolResultCompactionStrategy
|
||||
from .._feature_stage import ExperimentalFeature, experimental
|
||||
from .._sessions import ContextProvider, HistoryProvider, InMemoryHistoryProvider, MessageInjectionMiddleware
|
||||
from .._skills import SkillsProvider
|
||||
from .._types import ChatOptions
|
||||
from ._background_agents import BackgroundAgentsProvider
|
||||
from ._file_access import AgentFileStore, FileAccessProvider, FileSystemAgentFileStore
|
||||
from ._file_memory import FileMemoryProvider
|
||||
from ._loop import DEFAULT_MAX_ITERATIONS, AgentLoopMiddleware
|
||||
from ._mode import AgentModeProvider
|
||||
from ._todo import TodoProvider
|
||||
from ._tool_approval import ToolApprovalMiddleware
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # pragma: no cover
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Mapping
|
||||
|
||||
from agent_framework_tools.shell import ShellEnvironmentProviderOptions, ShellExecutor
|
||||
|
||||
from .._clients import SupportsChatGetResponse
|
||||
from .._compaction import CompactionStrategy, TokenizerProtocol
|
||||
from .._middleware import MiddlewareTypes
|
||||
from .._tools import ToolTypes
|
||||
from ._loop import NextMessageCallable, ShouldContinueCallable
|
||||
from ._tool_approval import ToolApprovalRuleCallback
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_HARNESS_INSTRUCTIONS = """\
|
||||
You are a helpful AI assistant that uses tools to complete tasks.
|
||||
|
||||
## General guidelines
|
||||
|
||||
- Think through the task before acting. Break complex work into clear steps.
|
||||
- Use the tools available to you to gather information, perform actions, and verify results.
|
||||
- Explain your reasoning and thought process as you work through tasks.
|
||||
- Explain what you learned and what you are going to do next between tool calls, \
|
||||
so the user can follow along with your thought process.
|
||||
- Avoid making more than 4 tool calls in a row without explaining what you are doing.
|
||||
- If a tool call fails or returns unexpected results, adapt your approach rather than \
|
||||
repeating the same call.
|
||||
- When you have completed the task, present a clear and concise summary of what you did \
|
||||
and what you found.
|
||||
"""
|
||||
|
||||
|
||||
def _assemble_instructions(
|
||||
harness_instructions: str | None,
|
||||
agent_instructions: str | None,
|
||||
) -> str | None:
|
||||
"""Assemble final instructions from harness + agent instructions."""
|
||||
harness = harness_instructions if harness_instructions is not None else DEFAULT_HARNESS_INSTRUCTIONS
|
||||
|
||||
return f"{harness}\n\n{agent_instructions or ''}".strip() or None
|
||||
|
||||
|
||||
def _assemble_compaction_provider(
|
||||
*,
|
||||
disable_compaction: bool,
|
||||
max_context_window_tokens: int | None,
|
||||
max_output_tokens: int | None,
|
||||
history_source_id: str,
|
||||
before_compaction_strategy: CompactionStrategy | None,
|
||||
after_compaction_strategy: CompactionStrategy | None,
|
||||
tokenizer: TokenizerProtocol | None,
|
||||
) -> CompactionProvider | None:
|
||||
"""Build the compaction provider from parameters or defaults.
|
||||
|
||||
The token-budget defaults (``ContextWindowCompactionStrategy`` for the before phase and
|
||||
``ToolResultCompactionStrategy`` for the after phase) are only applied when the token
|
||||
params are provided. Caller-supplied strategies are always honored. Either phase may end
|
||||
up ``None``, which ``CompactionProvider`` interprets as "skip that phase".
|
||||
|
||||
Returns None when compaction is explicitly disabled, or when neither phase has a strategy
|
||||
(no custom strategies and no token budget to build the defaults).
|
||||
"""
|
||||
if disable_compaction:
|
||||
return None
|
||||
|
||||
# Resolve the before-strategy: custom strategy wins; otherwise fall back to the
|
||||
# token-budget-aware default when token params are available.
|
||||
before_strategy = before_compaction_strategy
|
||||
if before_strategy is None and max_context_window_tokens is not None and max_output_tokens is not None:
|
||||
before_strategy = ContextWindowCompactionStrategy(
|
||||
max_context_window_tokens=max_context_window_tokens,
|
||||
max_output_tokens=max_output_tokens,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
|
||||
# Resolve the after-strategy: custom strategy wins; otherwise fall back to the default
|
||||
# when token params are available.
|
||||
after_strategy = after_compaction_strategy
|
||||
if after_strategy is None and max_context_window_tokens is not None and max_output_tokens is not None:
|
||||
after_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=2)
|
||||
|
||||
# Nothing to compact in either phase: skip the provider entirely.
|
||||
if before_strategy is None and after_strategy is None:
|
||||
return None
|
||||
|
||||
return CompactionProvider(
|
||||
before_strategy=before_strategy,
|
||||
after_strategy=after_strategy,
|
||||
tokenizer=tokenizer,
|
||||
history_source_id=history_source_id,
|
||||
)
|
||||
|
||||
|
||||
def _assemble_context_providers(
|
||||
*,
|
||||
history_provider: HistoryProvider,
|
||||
compaction_provider: CompactionProvider | None,
|
||||
disable_todo: bool,
|
||||
todo_provider: TodoProvider | None,
|
||||
disable_mode: bool,
|
||||
mode_provider: AgentModeProvider | None,
|
||||
disable_file_memory: bool,
|
||||
file_memory_store: AgentFileStore | None,
|
||||
disable_file_access: bool,
|
||||
file_access_store: AgentFileStore | None,
|
||||
file_access_disable_write_tools: bool,
|
||||
file_access_disable_readonly_tool_approval: bool,
|
||||
file_access_disable_write_tool_approval: bool,
|
||||
skills_provider: SkillsProvider | None,
|
||||
skills_paths: str | Path | Sequence[str | Path] | None,
|
||||
background_agents: Sequence[SupportsAgentRun] | None,
|
||||
background_agents_instructions: str | None,
|
||||
shell_context_provider: ContextProvider | None,
|
||||
extra_context_providers: Sequence[ContextProvider] | None,
|
||||
) -> list[ContextProvider]:
|
||||
"""Assemble the ordered list of context providers."""
|
||||
providers: list[ContextProvider] = []
|
||||
|
||||
# History first so other providers can access loaded messages.
|
||||
providers.append(history_provider)
|
||||
|
||||
# Compaction runs after history loads messages.
|
||||
if compaction_provider is not None:
|
||||
providers.append(compaction_provider)
|
||||
|
||||
if not disable_todo:
|
||||
providers.append(todo_provider or TodoProvider())
|
||||
|
||||
if not disable_mode:
|
||||
providers.append(mode_provider or AgentModeProvider())
|
||||
|
||||
# File-based session memory (on by default). Default store is rooted at
|
||||
# ``{cwd}/agent-file-memory``; the provider isolates memories per session
|
||||
# via its default ``scope=session_id``.
|
||||
if not disable_file_memory:
|
||||
memory_store = file_memory_store or FileSystemAgentFileStore(Path.cwd() / "agent-file-memory")
|
||||
providers.append(FileMemoryProvider(memory_store))
|
||||
|
||||
# Shared file access (on by default). Default store is rooted at ``{cwd}/working``.
|
||||
if not disable_file_access:
|
||||
access_store = file_access_store or FileSystemAgentFileStore(Path.cwd() / "working")
|
||||
providers.append(
|
||||
FileAccessProvider(
|
||||
access_store,
|
||||
disable_write_tools=file_access_disable_write_tools,
|
||||
disable_readonly_tool_approval=file_access_disable_readonly_tool_approval,
|
||||
disable_write_tool_approval=file_access_disable_write_tool_approval,
|
||||
)
|
||||
)
|
||||
|
||||
# Skills are opt-in: only added when skills_provider or skills_paths is provided.
|
||||
if skills_provider:
|
||||
providers.append(skills_provider)
|
||||
if skills_paths:
|
||||
providers.append(SkillsProvider.from_paths(skills_paths))
|
||||
|
||||
# Background agents are opt-in: only added when agents are provided.
|
||||
if background_agents:
|
||||
providers.append(BackgroundAgentsProvider(background_agents, instructions=background_agents_instructions))
|
||||
|
||||
# Shell environment provider is opt-in: only added when a shell tool was wired.
|
||||
if shell_context_provider is not None:
|
||||
providers.append(shell_context_provider)
|
||||
|
||||
# Append any user-supplied additional providers.
|
||||
if extra_context_providers:
|
||||
providers.extend(extra_context_providers)
|
||||
|
||||
return providers
|
||||
|
||||
|
||||
def _assemble_shell(
|
||||
client: SupportsChatGetResponse[Any],
|
||||
shell_executor: ShellExecutor | None,
|
||||
shell_environment_provider_options: ShellEnvironmentProviderOptions | None,
|
||||
) -> tuple[ToolTypes | None, ContextProvider | None]:
|
||||
"""Build the shell tool and environment provider when a shell executor is supplied.
|
||||
|
||||
Returns a ``(tool, provider)`` tuple. Both are ``None`` when no shell executor is
|
||||
provided, or when the client does not support shell tools (a warning is logged in the
|
||||
latter case, since the environment provider is not useful without an execution path).
|
||||
|
||||
Raises:
|
||||
TypeError: If ``shell_executor`` does not expose a callable ``as_function()`` method.
|
||||
"""
|
||||
if shell_executor is None:
|
||||
return None, None
|
||||
|
||||
# ShellExecutor is a protocol without ``as_function()``, so the
|
||||
# contract is validated at runtime: a shell tool such as LocalShellTool/DockerShellTool exposes it.
|
||||
as_function = getattr(shell_executor, "as_function", None)
|
||||
if not callable(as_function):
|
||||
raise TypeError(
|
||||
f"shell_executor must expose a callable 'as_function()' method "
|
||||
f"(e.g. a LocalShellTool or DockerShellTool from agent-framework-tools), "
|
||||
f"but got {type(shell_executor).__name__}."
|
||||
)
|
||||
|
||||
if not isinstance(client, SupportsShellTool):
|
||||
logger.warning(
|
||||
"Shell tool not available: client %r does not implement SupportsShellTool. "
|
||||
"Skipping the shell tool and environment provider.",
|
||||
type(client).__name__,
|
||||
)
|
||||
return None, None
|
||||
|
||||
# Imported lazily: the shell types live in the separate agent-framework-tools package,
|
||||
# which depends on core, so core cannot import them at module load time.
|
||||
from agent_framework_tools.shell import ShellEnvironmentProvider
|
||||
|
||||
shell_tool = client.get_shell_tool(func=as_function())
|
||||
shell_provider = ShellEnvironmentProvider(shell_executor, shell_environment_provider_options)
|
||||
return shell_tool, shell_provider
|
||||
|
||||
|
||||
HARNESS_AGENT_PROVIDER_NAME = "microsoft.agent_framework.harness"
|
||||
|
||||
OptionsCoT = TypeVar(
|
||||
"OptionsCoT",
|
||||
bound=TypedDict, # type: ignore[valid-type]
|
||||
default="ChatOptions[None]",
|
||||
)
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
def create_harness_agent(
|
||||
client: SupportsChatGetResponse[OptionsCoT],
|
||||
*,
|
||||
id: str | None = None,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
harness_instructions: str | None = None,
|
||||
agent_instructions: str | None = None,
|
||||
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
|
||||
max_context_window_tokens: int | None = None,
|
||||
max_output_tokens: int | None = None,
|
||||
history_provider: HistoryProvider | None = None,
|
||||
disable_compaction: bool = False,
|
||||
before_compaction_strategy: CompactionStrategy | None = None,
|
||||
after_compaction_strategy: CompactionStrategy | None = None,
|
||||
tokenizer: TokenizerProtocol | None = None,
|
||||
disable_todo: bool = False,
|
||||
todo_provider: TodoProvider | None = None,
|
||||
disable_mode: bool = False,
|
||||
mode_provider: AgentModeProvider | None = None,
|
||||
disable_file_memory: bool = False,
|
||||
file_memory_store: AgentFileStore | None = None,
|
||||
disable_file_access: bool = False,
|
||||
file_access_store: AgentFileStore | None = None,
|
||||
file_access_disable_write_tools: bool = False,
|
||||
file_access_disable_readonly_tool_approval: bool = False,
|
||||
file_access_disable_write_tool_approval: bool = False,
|
||||
skills_provider: SkillsProvider | None = None,
|
||||
skills_paths: str | Path | Sequence[str | Path] | None = None,
|
||||
background_agents: Sequence[SupportsAgentRun] | None = None,
|
||||
background_agents_instructions: str | None = None,
|
||||
shell_executor: ShellExecutor | None = None,
|
||||
shell_environment_provider_options: ShellEnvironmentProviderOptions | None = None,
|
||||
disable_web_search: bool = False,
|
||||
disable_tool_auto_approval: bool = False,
|
||||
auto_approval_rules: Sequence[ToolApprovalRuleCallback] | None = None,
|
||||
loop_should_continue: ShouldContinueCallable | None = None,
|
||||
loop_next_message: NextMessageCallable | None = None,
|
||||
loop_max_iterations: int | None = DEFAULT_MAX_ITERATIONS,
|
||||
otel_provider_name: str | None = None,
|
||||
context_providers: Sequence[ContextProvider] | None = None,
|
||||
middleware: Sequence[MiddlewareTypes] | None = None,
|
||||
default_options: Mapping[str, Any] | None = None,
|
||||
) -> Agent[OptionsCoT]:
|
||||
"""Create a pre-configured agent with batteries included.
|
||||
|
||||
Assembles an :class:`~agent_framework.Agent` from a chat client, automatically wiring:
|
||||
|
||||
- **Function invocation** — automatic tool calling loop
|
||||
- **Per-service-call history persistence** — persists history after every model call
|
||||
- **Compaction** — context-window compaction before/after each run
|
||||
- **TodoProvider** — todo list management
|
||||
- **AgentModeProvider** — plan/execute mode tracking
|
||||
- **FileMemoryProvider** — file-based session memory (on by default)
|
||||
- **FileAccessProvider** — shared file read/write tools (on by default)
|
||||
- **SkillsProvider** — skill discovery and progressive loading
|
||||
- **BackgroundAgentsProvider** — delegate work to background sub-agents
|
||||
- **Tool approval** — "don't ask again" standing approval rules plus heuristic
|
||||
auto-approval callbacks
|
||||
- **Looping** — re-run the agent until a ``should_continue`` predicate is satisfied
|
||||
- **OpenTelemetry** — observability via ``AgentTelemetryLayer``
|
||||
|
||||
Each feature can be disabled or customized via keyword arguments.
|
||||
|
||||
Examples:
|
||||
Basic usage:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import create_harness_agent
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
agent = create_harness_agent(
|
||||
OpenAIChatClient(model="gpt-4o"),
|
||||
)
|
||||
session = agent.create_session()
|
||||
response = await agent.run("Plan a weekend trip to Seattle", session=session)
|
||||
|
||||
With customization:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
agent = create_harness_agent(
|
||||
client=client,
|
||||
max_context_window_tokens=200_000,
|
||||
max_output_tokens=32_000,
|
||||
name="research-agent",
|
||||
agent_instructions="Focus on academic sources.",
|
||||
disable_todo=True,
|
||||
skills_paths=["./skills", "./custom-skills"],
|
||||
)
|
||||
|
||||
Args:
|
||||
client: The chat client providing access to the underlying AI model.
|
||||
|
||||
Keyword Args:
|
||||
id: Optional agent ID (auto-generated UUID if omitted).
|
||||
name: Optional agent name.
|
||||
description: Optional agent description.
|
||||
harness_instructions: Override the default harness-level system instructions that
|
||||
govern agent behavior (how to use tools, report progress, structure responses).
|
||||
These provide general "operating guidelines" independent of any specific task.
|
||||
When None, ``DEFAULT_HARNESS_INSTRUCTIONS`` is used. Set to empty string ``""``
|
||||
to omit harness instructions entirely.
|
||||
agent_instructions: Domain or task-specific instructions appended after harness
|
||||
instructions. Use this for the agent's purpose, persona, or specialization
|
||||
(e.g., "You are a research assistant focused on academic sources.").
|
||||
tools: Additional tools to include in the agent's toolset.
|
||||
max_context_window_tokens: Maximum tokens the model's context window supports.
|
||||
Used to construct the default token-budget-aware compaction strategies. When None
|
||||
(default) and no custom ``before_compaction_strategy`` / ``after_compaction_strategy``
|
||||
is provided, compaction is automatically disabled.
|
||||
max_output_tokens: Maximum output tokens per response.
|
||||
Used to construct the default compaction strategies and sets a default max_tokens
|
||||
chat option. When None (default), no default max_tokens option is set, and unless a
|
||||
custom compaction strategy is provided, compaction is automatically disabled.
|
||||
history_provider: Custom history provider. When None, an InMemoryHistoryProvider is used.
|
||||
disable_compaction: When True, skip compaction provider setup.
|
||||
before_compaction_strategy: Custom before-run compaction strategy. When provided,
|
||||
compaction runs even if token params are omitted. Defaults to
|
||||
ContextWindowCompactionStrategy (token-budget aware) when token params are provided.
|
||||
after_compaction_strategy: Custom after-run compaction strategy. When provided,
|
||||
compaction runs even if token params are omitted. Defaults to
|
||||
ToolResultCompactionStrategy when token params are provided.
|
||||
tokenizer: Custom tokenizer for compaction strategies.
|
||||
disable_todo: When True, skip the TodoProvider.
|
||||
todo_provider: Custom TodoProvider instance. Ignored when disable_todo is True.
|
||||
disable_mode: When True, skip the AgentModeProvider.
|
||||
mode_provider: Custom AgentModeProvider instance. Ignored when disable_mode is True.
|
||||
disable_file_memory: When True, skip the FileMemoryProvider. When False (default),
|
||||
a FileMemoryProvider is added, giving the agent session-scoped, file-based memory.
|
||||
file_memory_store: Custom AgentFileStore backing the FileMemoryProvider. When None
|
||||
(and disable_file_memory is False), a FileSystemAgentFileStore rooted at
|
||||
``{cwd}/agent-file-memory`` is created. Ignored when disable_file_memory is True.
|
||||
disable_file_access: When True, skip the FileAccessProvider. When False (default),
|
||||
a FileAccessProvider is added, giving the agent shared read/write file tools.
|
||||
file_access_store: Custom AgentFileStore backing the FileAccessProvider. When None
|
||||
(and disable_file_access is False), a FileSystemAgentFileStore rooted at
|
||||
``{cwd}/working`` is created. Ignored when disable_file_access is True.
|
||||
file_access_disable_write_tools: When True, the FileAccessProvider advertises only its
|
||||
read-only tools (read, ls, grep); the write tools (write, delete, replace,
|
||||
replace_lines) are hidden. When False (default), all tools are advertised. Ignored
|
||||
when disable_file_access is True.
|
||||
file_access_disable_readonly_tool_approval: When True, the FileAccessProvider's read-only
|
||||
tools (read, ls, grep) are registered with ``approval_mode="never_require"`` so they
|
||||
run without host approval. When False (default), they require approval. Ignored when
|
||||
disable_file_access is True.
|
||||
file_access_disable_write_tool_approval: When True, the FileAccessProvider's write tools
|
||||
(write, delete, replace, replace_lines) are registered with
|
||||
``approval_mode="never_require"`` so they run without host approval. When False
|
||||
(default), they require approval. Ignored when disable_file_access is True.
|
||||
skills_provider: Custom SkillsProvider instance for code-defined skills.
|
||||
Can be combined with ``skills_paths`` to aggregate file and code-based skills.
|
||||
**Security:** if the provider is configured with an external skill source (e.g.
|
||||
:class:`~agent_framework.MCPSkillsSource`), the skill content it loads is untrusted input
|
||||
— only enable sources you trust; see :class:`~agent_framework.SkillsSource`.
|
||||
skills_paths: Paths for file-based skill discovery (looks for SKILL.md files).
|
||||
Accepts a single ``str`` or :class:`~pathlib.Path`, or a sequence of
|
||||
``str | Path``. Can be combined with ``skills_provider``. When neither
|
||||
``skills_provider`` nor ``skills_paths`` is provided, no SkillsProvider
|
||||
is added.
|
||||
background_agents: Collection of agents available for background task delegation.
|
||||
When provided, a ``BackgroundAgentsProvider`` is automatically included,
|
||||
enabling the agent to start, monitor, and retrieve results from background tasks.
|
||||
Each agent must have a non-empty, unique name (case-insensitive).
|
||||
**Security:** supplied agents receive text input from this agent and their output is fed
|
||||
back into its context, so only supply agents you have vetted and trust — see
|
||||
:class:`~agent_framework.BackgroundAgentsProvider` for the exfiltration and
|
||||
prompt-injection risks of untrusted agents.
|
||||
background_agents_instructions: Optional instruction override for the
|
||||
``BackgroundAgentsProvider``. May include ``{background_agents}`` placeholder
|
||||
which will be replaced with the agent listing.
|
||||
shell_executor: Optional shell tool that enables shell command execution. When
|
||||
provided, the shell tool and a ``ShellEnvironmentProvider`` are automatically
|
||||
added (provided the client supports shell tools; otherwise a warning is logged
|
||||
and both are skipped). The object must expose ``as_function()`` and satisfy the
|
||||
``ShellExecutor`` protocol -- e.g. a ``LocalShellTool`` or ``DockerShellTool`` from
|
||||
the ``agent-framework-tools`` package. The caller owns the executor's lifecycle.
|
||||
shell_environment_provider_options: Optional ``ShellEnvironmentProviderOptions``
|
||||
(from ``agent-framework-tools``) used to customize the ``ShellEnvironmentProvider``
|
||||
environment probing and instructions. Only used when ``shell_executor`` is provided.
|
||||
disable_web_search: When True, skip automatic web search tool inclusion.
|
||||
When False (default), the web search tool is automatically added if the
|
||||
client implements SupportsWebSearchTool. A warning is logged if the client
|
||||
does not support web search.
|
||||
disable_tool_auto_approval: When True, do not wire the tool auto-approval middleware.
|
||||
When False (default), a :class:`~agent_framework.ToolApprovalMiddleware` is added
|
||||
(outermost) to coordinate "don't ask again" standing approval rules and queued
|
||||
approval prompts; callers must pass an :class:`~agent_framework.AgentSession` to
|
||||
:meth:`~agent_framework.Agent.run` when enabled.
|
||||
auto_approval_rules: Optional heuristic callbacks that can auto-approve a function call
|
||||
that would otherwise require approval. Each callback receives the ``function_call``
|
||||
content and returns ``True`` to approve it. Rules are evaluated after standing rules
|
||||
(derived from prior user approvals) but before prompting the user. Only used when
|
||||
``disable_tool_auto_approval`` is False.
|
||||
loop_should_continue: Optional predicate that enables the looping middleware. When provided, the
|
||||
agent is re-run in a loop (via :class:`~agent_framework.AgentLoopMiddleware`, wired as
|
||||
the outermost middleware so each iteration is a full agent run including tool approval)
|
||||
for as long as the predicate returns ``True``, up to ``loop_max_iterations``. If an
|
||||
iteration returns a pending tool-approval request, the loop stops and returns it so the
|
||||
caller can approve before continuing. When None (default), no loop is added.
|
||||
loop_next_message: Optional callable controlling the input for the next loop iteration.
|
||||
Only takes effect when ``loop_should_continue`` is set (otherwise no loop is added and
|
||||
this is ignored).
|
||||
loop_max_iterations: Safety cap on the number of loop iterations. ``None`` means unbounded;
|
||||
a positive integer caps the loop (defaults to the loop middleware's default cap). Only
|
||||
takes effect when ``loop_should_continue`` is set (otherwise no loop is added and this
|
||||
is ignored).
|
||||
otel_provider_name: Custom OpenTelemetry provider/source name for telemetry.
|
||||
context_providers: Additional context providers to include after the built-in ones.
|
||||
middleware: Additional middleware to include.
|
||||
default_options: Provider-specific chat options (temperature, max_tokens, etc.).
|
||||
|
||||
Returns:
|
||||
A fully configured :class:`~agent_framework.Agent` instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If max_context_window_tokens is provided and <= 0, or
|
||||
max_output_tokens is provided and <= 0, or max_output_tokens >=
|
||||
max_context_window_tokens when both are provided.
|
||||
"""
|
||||
if max_context_window_tokens is not None and max_context_window_tokens <= 0:
|
||||
raise ValueError("max_context_window_tokens must be positive.")
|
||||
if max_output_tokens is not None and max_output_tokens <= 0:
|
||||
raise ValueError("max_output_tokens must be positive.")
|
||||
if (
|
||||
max_context_window_tokens is not None
|
||||
and max_output_tokens is not None
|
||||
and max_output_tokens >= max_context_window_tokens
|
||||
):
|
||||
raise ValueError("max_output_tokens must be less than max_context_window_tokens.")
|
||||
|
||||
# Build history provider.
|
||||
resolved_history = history_provider or InMemoryHistoryProvider()
|
||||
|
||||
# Build compaction provider.
|
||||
compaction_provider = _assemble_compaction_provider(
|
||||
disable_compaction=disable_compaction,
|
||||
max_context_window_tokens=max_context_window_tokens,
|
||||
max_output_tokens=max_output_tokens,
|
||||
history_source_id=resolved_history.source_id,
|
||||
before_compaction_strategy=before_compaction_strategy,
|
||||
after_compaction_strategy=after_compaction_strategy,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
|
||||
# Build the shell tool and environment provider (opt-in via shell_executor).
|
||||
shell_tool, shell_provider = _assemble_shell(
|
||||
client,
|
||||
shell_executor,
|
||||
shell_environment_provider_options,
|
||||
)
|
||||
|
||||
# Build context providers.
|
||||
assembled_providers = _assemble_context_providers(
|
||||
history_provider=resolved_history,
|
||||
compaction_provider=compaction_provider,
|
||||
disable_todo=disable_todo,
|
||||
todo_provider=todo_provider,
|
||||
disable_mode=disable_mode,
|
||||
mode_provider=mode_provider,
|
||||
disable_file_memory=disable_file_memory,
|
||||
file_memory_store=file_memory_store,
|
||||
disable_file_access=disable_file_access,
|
||||
file_access_store=file_access_store,
|
||||
file_access_disable_write_tools=file_access_disable_write_tools,
|
||||
file_access_disable_readonly_tool_approval=file_access_disable_readonly_tool_approval,
|
||||
file_access_disable_write_tool_approval=file_access_disable_write_tool_approval,
|
||||
skills_provider=skills_provider,
|
||||
skills_paths=skills_paths,
|
||||
background_agents=background_agents,
|
||||
background_agents_instructions=background_agents_instructions,
|
||||
shell_context_provider=shell_provider,
|
||||
extra_context_providers=context_providers,
|
||||
)
|
||||
|
||||
# Build instructions.
|
||||
instructions = _assemble_instructions(harness_instructions, agent_instructions)
|
||||
|
||||
# Assemble tools, auto-adding web search if supported.
|
||||
assembled_tools: list[ToolTypes | Callable[..., Any]] = []
|
||||
if not disable_web_search:
|
||||
if isinstance(client, SupportsWebSearchTool):
|
||||
assembled_tools.append(client.get_web_search_tool())
|
||||
else:
|
||||
logger.warning(
|
||||
"Web search tool not available: client %r does not implement SupportsWebSearchTool. "
|
||||
"Set disable_web_search=True to suppress this warning.",
|
||||
type(client).__name__,
|
||||
)
|
||||
if shell_tool is not None:
|
||||
assembled_tools.append(shell_tool)
|
||||
if tools is not None:
|
||||
if isinstance(tools, Sequence):
|
||||
assembled_tools.extend(tools) # pyright: ignore[reportUnknownArgumentType]
|
||||
else:
|
||||
assembled_tools.append(tools)
|
||||
final_tools: list[ToolTypes | Callable[..., Any]] | None = assembled_tools or None
|
||||
|
||||
# Build default options dict.
|
||||
default_opts: dict[str, Any] = dict(default_options) if default_options else {}
|
||||
if max_output_tokens is not None:
|
||||
default_opts.setdefault("max_tokens", max_output_tokens)
|
||||
|
||||
# Assemble middleware. Tool approval is enabled by default (like the .NET harness) and is
|
||||
# placed first so it sits outermost: it intercepts inbound "always approve" responses and
|
||||
# outbound approval requests at the caller boundary, and its re-invocation loop re-runs any
|
||||
# user-supplied middleware. ToolApprovalMiddleware requires an AgentSession at run time.
|
||||
# When should_continue is supplied, the loop is prepended ahead of tool approval so it sits
|
||||
# outermost of all: each loop iteration is a full agent run (including tool approval), and the
|
||||
# loop's approval escape hatch returns any pending approval request to the caller.
|
||||
assembled_middleware: list[MiddlewareTypes] = []
|
||||
if not disable_tool_auto_approval:
|
||||
assembled_middleware.append(ToolApprovalMiddleware(auto_approval_rules=auto_approval_rules))
|
||||
if loop_should_continue is not None:
|
||||
assembled_middleware.insert(
|
||||
0,
|
||||
AgentLoopMiddleware(
|
||||
loop_should_continue,
|
||||
max_iterations=loop_max_iterations,
|
||||
next_message=loop_next_message,
|
||||
),
|
||||
)
|
||||
# Message injection is always on. It is a no-op when no messages are queued for the session,
|
||||
# so there is no opt-out.
|
||||
assembled_middleware.append(MessageInjectionMiddleware())
|
||||
if middleware:
|
||||
assembled_middleware.extend(middleware)
|
||||
|
||||
agent = Agent(
|
||||
client,
|
||||
instructions,
|
||||
id=id,
|
||||
name=name,
|
||||
description=description,
|
||||
tools=final_tools,
|
||||
default_options=default_opts, # type: ignore[arg-type]
|
||||
context_providers=assembled_providers,
|
||||
middleware=assembled_middleware or None,
|
||||
require_per_service_call_history_persistence=True,
|
||||
)
|
||||
|
||||
# Set the telemetry provider name after construction.
|
||||
agent.otel_provider_name = otel_provider_name or HARNESS_AGENT_PROVIDER_NAME
|
||||
|
||||
return agent
|
||||
@@ -0,0 +1,92 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol
|
||||
|
||||
from typing_extensions import TypedDict, TypeVar
|
||||
|
||||
from .._agents import Agent, SupportsAgentRun
|
||||
from .._clients import SupportsChatGetResponse
|
||||
from .._compaction import CompactionStrategy, TokenizerProtocol
|
||||
from .._middleware import MiddlewareTypes
|
||||
from .._sessions import ContextProvider, HistoryProvider
|
||||
from .._skills import SkillsProvider
|
||||
from .._tools import ToolTypes
|
||||
from .._types import ChatOptions
|
||||
from ._file_access import AgentFileStore
|
||||
from ._loop import DEFAULT_MAX_ITERATIONS, NextMessageCallable, ShouldContinueCallable
|
||||
from ._mode import AgentModeProvider
|
||||
from ._todo import TodoProvider
|
||||
from ._tool_approval import ToolApprovalRuleCallback
|
||||
|
||||
DEFAULT_HARNESS_INSTRUCTIONS: str
|
||||
HARNESS_AGENT_PROVIDER_NAME: str
|
||||
|
||||
OptionsCoT = TypeVar(
|
||||
"OptionsCoT",
|
||||
bound=TypedDict, # type: ignore[valid-type]
|
||||
default=ChatOptions[None],
|
||||
)
|
||||
|
||||
class _ShellExecutorLike(Protocol):
|
||||
def as_function(self, *args: Any, **kwargs: Any) -> Any: ...
|
||||
|
||||
class _ShellEnvironmentProviderOptionsLike(Protocol):
|
||||
@property
|
||||
def probe_tools(self) -> Sequence[str]: ...
|
||||
@property
|
||||
def override_family(self) -> Any | None: ...
|
||||
@property
|
||||
def probe_timeout(self) -> float: ...
|
||||
@property
|
||||
def instructions_formatter(self) -> Callable[[Any], str] | None: ...
|
||||
|
||||
def _assemble_instructions(
|
||||
harness_instructions: str | None,
|
||||
agent_instructions: str | None,
|
||||
) -> str | None: ...
|
||||
def create_harness_agent(
|
||||
client: SupportsChatGetResponse[OptionsCoT],
|
||||
*,
|
||||
id: str | None = None,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
harness_instructions: str | None = None,
|
||||
agent_instructions: str | None = None,
|
||||
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
|
||||
max_context_window_tokens: int | None = None,
|
||||
max_output_tokens: int | None = None,
|
||||
history_provider: HistoryProvider | None = None,
|
||||
disable_compaction: bool = False,
|
||||
before_compaction_strategy: CompactionStrategy | None = None,
|
||||
after_compaction_strategy: CompactionStrategy | None = None,
|
||||
tokenizer: TokenizerProtocol | None = None,
|
||||
disable_todo: bool = False,
|
||||
todo_provider: TodoProvider | None = None,
|
||||
disable_mode: bool = False,
|
||||
mode_provider: AgentModeProvider | None = None,
|
||||
disable_file_memory: bool = False,
|
||||
file_memory_store: AgentFileStore | None = None,
|
||||
disable_file_access: bool = False,
|
||||
file_access_store: AgentFileStore | None = None,
|
||||
file_access_disable_write_tools: bool = False,
|
||||
file_access_disable_readonly_tool_approval: bool = False,
|
||||
file_access_disable_write_tool_approval: bool = False,
|
||||
skills_provider: SkillsProvider | None = None,
|
||||
skills_paths: str | Path | Sequence[str | Path] | None = None,
|
||||
background_agents: Sequence[SupportsAgentRun] | None = None,
|
||||
background_agents_instructions: str | None = None,
|
||||
shell_executor: _ShellExecutorLike | None = None,
|
||||
shell_environment_provider_options: _ShellEnvironmentProviderOptionsLike | None = None,
|
||||
disable_web_search: bool = False,
|
||||
disable_tool_auto_approval: bool = False,
|
||||
auto_approval_rules: Sequence[ToolApprovalRuleCallback] | None = None,
|
||||
loop_should_continue: ShouldContinueCallable | None = None,
|
||||
loop_next_message: NextMessageCallable | None = None,
|
||||
loop_max_iterations: int | None = DEFAULT_MAX_ITERATIONS,
|
||||
otel_provider_name: str | None = None,
|
||||
context_providers: Sequence[ContextProvider] | None = None,
|
||||
middleware: Sequence[MiddlewareTypes] | None = None,
|
||||
default_options: Mapping[str, Any] | None = None,
|
||||
) -> Agent[OptionsCoT]: ...
|
||||
@@ -0,0 +1,534 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""BackgroundAgentsProvider: enables an agent to delegate work to background sub-agents asynchronously.
|
||||
|
||||
This module provides :class:`BackgroundAgentsProvider`, a context provider that allows
|
||||
a parent agent to start background tasks on child agents, wait for their completion,
|
||||
and retrieve results. Each background task runs in its own session concurrently.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, MutableMapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, ClassVar, cast
|
||||
|
||||
from .._agents import SupportsAgentRun
|
||||
from .._feature_stage import ExperimentalFeature, experimental
|
||||
from .._serialization import SerializationMixin
|
||||
from .._sessions import AgentSession, ContextProvider, SessionContext
|
||||
from .._tools import tool
|
||||
from .._types import AgentResponse, Message
|
||||
|
||||
DEFAULT_BACKGROUND_AGENTS_SOURCE_ID = "background_agents"
|
||||
|
||||
DEFAULT_BACKGROUND_AGENTS_INSTRUCTIONS = """\
|
||||
## Background Agents
|
||||
|
||||
You have access to background agents that can perform work on your behalf.
|
||||
|
||||
- Use the `background_agents_*` tools to start tasks on background agents and check their results.
|
||||
- Creating a background task does not block, and background tasks run concurrently.
|
||||
- Important: Always wait for outstanding tasks to finish before you finish processing.
|
||||
- Important: After retrieving results from a completed task, clear it with \
|
||||
background_agents_clear_completed_task to free memory, unless you plan to continue it with \
|
||||
background_agents_continue_task.
|
||||
|
||||
{background_agents}"""
|
||||
|
||||
|
||||
class BackgroundTaskStatus(str, Enum):
|
||||
"""Status of a background task."""
|
||||
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
LOST = "lost"
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class BackgroundTaskInfo(SerializationMixin):
|
||||
"""Metadata for a single background task."""
|
||||
|
||||
DEFAULT_EXCLUDE: ClassVar[set[str]] = set()
|
||||
|
||||
id: int
|
||||
agent_name: str
|
||||
description: str
|
||||
status: BackgroundTaskStatus
|
||||
result_text: str | None
|
||||
error_text: str | None
|
||||
__slots__ = ("agent_name", "description", "error_text", "id", "result_text", "status")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
id: int,
|
||||
agent_name: str,
|
||||
description: str,
|
||||
status: BackgroundTaskStatus = BackgroundTaskStatus.RUNNING,
|
||||
result_text: str | None = None,
|
||||
error_text: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize a background task info entry."""
|
||||
self.id = id
|
||||
self.agent_name = agent_name
|
||||
self.description = description
|
||||
self.status = status
|
||||
self.result_text = result_text
|
||||
self.error_text = error_text
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
|
||||
"""Serialize for session state persistence."""
|
||||
del exclude
|
||||
data: dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"agent_name": self.agent_name,
|
||||
"description": self.description,
|
||||
"status": self.status.value,
|
||||
}
|
||||
if not exclude_none or self.result_text is not None:
|
||||
data["result_text"] = self.result_text
|
||||
if not exclude_none or self.error_text is not None:
|
||||
data["error_text"] = self.error_text
|
||||
return data
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: MutableMapping[str, Any], **kwargs: Any) -> BackgroundTaskInfo:
|
||||
"""Deserialize from session state."""
|
||||
return cls(
|
||||
id=data["id"],
|
||||
agent_name=data["agent_name"],
|
||||
description=data["description"],
|
||||
status=BackgroundTaskStatus(data["status"]),
|
||||
result_text=data.get("result_text"),
|
||||
error_text=data.get("error_text"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RuntimeState:
|
||||
"""Non-serializable per-session runtime state for background tasks."""
|
||||
|
||||
in_flight_tasks: dict[int, asyncio.Task[AgentResponse[Any]]] = field(default_factory=lambda: {})
|
||||
background_sessions: dict[int, AgentSession] = field(default_factory=lambda: {})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module-level helper functions (following ModeProvider pattern)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _run_agent(awaitable: Awaitable[AgentResponse[Any]]) -> AgentResponse[Any]:
|
||||
"""Wrap an Awaitable in a proper coroutine for use with asyncio.create_task."""
|
||||
return await awaitable
|
||||
|
||||
|
||||
def _validate_and_build_agent_dict(agents: Sequence[SupportsAgentRun]) -> dict[str, SupportsAgentRun]:
|
||||
"""Validate agents and build a case-insensitive lookup dict.
|
||||
|
||||
Raises:
|
||||
ValueError: If agents is empty, an agent has no name, or names are not unique.
|
||||
"""
|
||||
if not agents:
|
||||
raise ValueError("At least one background agent must be provided.")
|
||||
|
||||
agent_dict: dict[str, SupportsAgentRun] = {}
|
||||
for agent in agents:
|
||||
name = agent.name
|
||||
if not name or not name.strip():
|
||||
raise ValueError("All background agents must have a non-empty name.")
|
||||
key = name.lower()
|
||||
if key in agent_dict:
|
||||
raise ValueError(
|
||||
f"Duplicate background agent name: '{name}'. Agent names must be unique (case-insensitive)."
|
||||
)
|
||||
agent_dict[key] = agent
|
||||
return agent_dict
|
||||
|
||||
|
||||
def _build_agent_list_text(agents: dict[str, SupportsAgentRun]) -> str:
|
||||
"""Build text listing available background agents."""
|
||||
lines = ["Available background agents:"]
|
||||
for agent in agents.values():
|
||||
line = f"- {agent.name}"
|
||||
if agent.description:
|
||||
line += f": {agent.description}"
|
||||
lines.append(line)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _get_provider_state(session: AgentSession, *, source_id: str) -> dict[str, Any]:
|
||||
"""Load or initialize serializable provider state from session."""
|
||||
state = session.state.get(source_id)
|
||||
if state is None:
|
||||
initial: dict[str, Any] = {"next_task_id": 1, "tasks": []}
|
||||
session.state[source_id] = initial
|
||||
return initial
|
||||
return cast(dict[str, Any], state)
|
||||
|
||||
|
||||
def _save_provider_state(session: AgentSession, state: dict[str, Any], *, source_id: str) -> None:
|
||||
"""Persist serializable state to session."""
|
||||
session.state[source_id] = state
|
||||
|
||||
|
||||
def _get_tasks(state: dict[str, Any]) -> list[BackgroundTaskInfo]:
|
||||
"""Parse task list from state dict."""
|
||||
return [BackgroundTaskInfo.from_dict(t) for t in state.get("tasks", [])]
|
||||
|
||||
|
||||
def _save_tasks(state: dict[str, Any], tasks: list[BackgroundTaskInfo]) -> None:
|
||||
"""Serialize task list back to state dict."""
|
||||
state["tasks"] = [t.to_dict() for t in tasks]
|
||||
|
||||
|
||||
def _finalize_task(
|
||||
task_info: BackgroundTaskInfo,
|
||||
completed_task: asyncio.Task[AgentResponse[Any]],
|
||||
runtime: _RuntimeState,
|
||||
) -> None:
|
||||
"""Extract results from a completed asyncio task and update task info."""
|
||||
if completed_task.cancelled():
|
||||
task_info.status = BackgroundTaskStatus.FAILED
|
||||
task_info.error_text = "Task was canceled."
|
||||
else:
|
||||
exception = completed_task.exception()
|
||||
if exception is not None:
|
||||
task_info.status = BackgroundTaskStatus.FAILED
|
||||
task_info.error_text = str(exception)
|
||||
else:
|
||||
task_info.status = BackgroundTaskStatus.COMPLETED
|
||||
task_info.result_text = completed_task.result().text
|
||||
runtime.in_flight_tasks.pop(task_info.id, None)
|
||||
|
||||
|
||||
def _refresh_task_state(
|
||||
session: AgentSession, state: dict[str, Any], runtime: _RuntimeState, *, source_id: str
|
||||
) -> list[BackgroundTaskInfo]:
|
||||
"""Refresh status of in-flight tasks and return updated task list."""
|
||||
tasks = _get_tasks(state)
|
||||
changed = False
|
||||
|
||||
for task_info in tasks:
|
||||
if task_info.status != BackgroundTaskStatus.RUNNING:
|
||||
continue
|
||||
|
||||
in_flight = runtime.in_flight_tasks.get(task_info.id)
|
||||
if in_flight is None:
|
||||
task_info.status = BackgroundTaskStatus.LOST
|
||||
changed = True
|
||||
continue
|
||||
|
||||
if in_flight.done():
|
||||
_finalize_task(task_info, in_flight, runtime)
|
||||
changed = True
|
||||
|
||||
if changed:
|
||||
_save_tasks(state, tasks)
|
||||
_save_provider_state(session, state, source_id=source_id)
|
||||
|
||||
return tasks
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Provider class
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class BackgroundAgentsProvider(ContextProvider):
|
||||
"""Context provider that enables an agent to delegate work to background sub-agents.
|
||||
|
||||
The ``BackgroundAgentsProvider`` allows a parent agent to start background tasks on child agents,
|
||||
wait for their completion, and retrieve results. Each background task runs in its own session and
|
||||
executes concurrently.
|
||||
|
||||
This provider exposes the following tools to the agent:
|
||||
|
||||
- ``background_agents_start_task`` — Start a background task on a named agent with text input.
|
||||
- ``background_agents_wait_for_first_completion`` — Block until the first of the specified tasks completes.
|
||||
- ``background_agents_get_task_results`` — Retrieve the text output of a completed background task.
|
||||
- ``background_agents_get_all_tasks`` — List all background tasks with their IDs, statuses, and descriptions.
|
||||
- ``background_agents_continue_task`` — Send follow-up input to a completed task's session to resume work.
|
||||
- ``background_agents_clear_completed_task`` — Remove a completed task and release its session.
|
||||
|
||||
Security considerations:
|
||||
The agents passed to the constructor are delegated arbitrary work by the parent agent — the
|
||||
parent sends them text input (which may include content derived from the parent's own
|
||||
untrusted context) and receives back whatever text they produce. A compromised or malicious
|
||||
supplied agent (for example, one with a compromised system prompt, tools, or upstream model)
|
||||
could exfiltrate that input to an external system, or return adversarial output designed to
|
||||
influence the parent agent via indirect prompt injection once its result is retrieved. Only
|
||||
supply background agents you have vetted and trust with the data the parent may pass to them.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agents: Sequence[SupportsAgentRun],
|
||||
*,
|
||||
source_id: str = DEFAULT_BACKGROUND_AGENTS_SOURCE_ID,
|
||||
instructions: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize the background agents provider.
|
||||
|
||||
Args:
|
||||
agents: Collection of background agents available for delegation.
|
||||
Each agent must have a non-empty, unique name (case-insensitive).
|
||||
**Security:** each supplied agent should be vetted and trusted, since it will receive
|
||||
text input from the parent agent and its output is fed back into the parent's
|
||||
context — see the class-level security considerations for the exfiltration and
|
||||
prompt-injection risks of untrusted agents.
|
||||
|
||||
Keyword Args:
|
||||
source_id: Unique source ID for serializable task state in session.
|
||||
instructions: Optional instruction override. May include ``{background_agents}``
|
||||
placeholder which will be replaced with the agent listing.
|
||||
|
||||
Raises:
|
||||
ValueError: If agents is empty, an agent has no name, or names are not unique.
|
||||
"""
|
||||
super().__init__(source_id)
|
||||
|
||||
self._agents = _validate_and_build_agent_dict(agents)
|
||||
|
||||
# Build instructions with agent listing.
|
||||
base_instructions = instructions if instructions is not None else DEFAULT_BACKGROUND_AGENTS_INSTRUCTIONS
|
||||
agent_list_text = _build_agent_list_text(self._agents)
|
||||
self._instructions = base_instructions.replace("{background_agents}", agent_list_text)
|
||||
|
||||
# Per-session runtime state (non-serializable), keyed by session_id.
|
||||
# Note: Runtime state (in-flight asyncio.Task objects, child AgentSession handles)
|
||||
# is inherently non-serializable and cannot survive process restarts. If the provider
|
||||
# instance is lost, _refresh_task_state() marks orphaned tasks as LOST.
|
||||
self._runtime: dict[str, _RuntimeState] = {}
|
||||
|
||||
def _get_runtime(self, session: AgentSession) -> _RuntimeState:
|
||||
"""Get or create runtime state for a session."""
|
||||
session_id = session.session_id
|
||||
if session_id not in self._runtime:
|
||||
self._runtime[session_id] = _RuntimeState()
|
||||
return self._runtime[session_id]
|
||||
|
||||
async def before_run(
|
||||
self,
|
||||
*,
|
||||
agent: Any,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Inject background agent tools and instructions before the model runs."""
|
||||
del agent, state
|
||||
|
||||
provider_state = _get_provider_state(session, source_id=self.source_id)
|
||||
runtime = self._get_runtime(session)
|
||||
source_id = self.source_id
|
||||
|
||||
@tool(name="background_agents_start_task", approval_mode="never_require")
|
||||
def background_agents_start_task(agent_name: str, input: str, description: str) -> str:
|
||||
"""Start a background task on a named agent. Returns a confirmation with the task ID."""
|
||||
key = agent_name.lower()
|
||||
if key not in self._agents:
|
||||
available = ", ".join(a.name or "" for a in self._agents.values())
|
||||
return f"Error: No background agent found with name '{agent_name}'. Available agents: {available}"
|
||||
|
||||
bg_agent = self._agents[key]
|
||||
task_id = provider_state.get("next_task_id", 1)
|
||||
provider_state["next_task_id"] = task_id + 1
|
||||
|
||||
task_info = BackgroundTaskInfo(
|
||||
id=task_id,
|
||||
agent_name=agent_name,
|
||||
description=description,
|
||||
)
|
||||
tasks = _get_tasks(provider_state)
|
||||
tasks.append(task_info)
|
||||
_save_tasks(provider_state, tasks)
|
||||
|
||||
# Create a dedicated session for this background task.
|
||||
sub_session = bg_agent.create_session()
|
||||
|
||||
# Start the task concurrently.
|
||||
async_task = asyncio.create_task(_run_agent(bg_agent.run(input, session=sub_session)))
|
||||
runtime.in_flight_tasks[task_id] = async_task
|
||||
runtime.background_sessions[task_id] = sub_session
|
||||
|
||||
_save_provider_state(session, provider_state, source_id=source_id)
|
||||
return f"Background task {task_id} started on agent '{agent_name}'."
|
||||
|
||||
background_agents_start_task._invoke_sync_on_event_loop = True # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
@tool(name="background_agents_wait_for_first_completion", approval_mode="never_require")
|
||||
async def background_agents_wait_for_first_completion(task_ids: list[int]) -> str:
|
||||
"""Block until the first of the specified background tasks completes. Returns the completed task's ID."""
|
||||
if not task_ids:
|
||||
return "Error: No task IDs provided."
|
||||
|
||||
# Collect in-flight tasks matching the requested IDs.
|
||||
waitable: list[tuple[int, asyncio.Task[AgentResponse[Any]]]] = []
|
||||
for tid in task_ids:
|
||||
in_flight = runtime.in_flight_tasks.get(tid)
|
||||
if in_flight is not None:
|
||||
waitable.append((tid, in_flight))
|
||||
|
||||
if not waitable:
|
||||
# Refresh state to catch any that completed.
|
||||
tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
|
||||
already_complete = next(
|
||||
(t for t in tasks if t.id in task_ids and t.status != BackgroundTaskStatus.RUNNING), None
|
||||
)
|
||||
if already_complete is not None:
|
||||
return (
|
||||
f"Task {already_complete.id} is not running; current status: {already_complete.status.value}."
|
||||
)
|
||||
return "Error: None of the specified task IDs correspond to running tasks."
|
||||
|
||||
# Wait for the first one to complete.
|
||||
done, _ = await asyncio.wait(
|
||||
[t for _, t in waitable],
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
|
||||
# Find which ID completed.
|
||||
completed_id: int | None = None
|
||||
for tid, task in waitable:
|
||||
if task in done:
|
||||
completed_id = tid
|
||||
break
|
||||
|
||||
# Finalize the completed task.
|
||||
tasks = _get_tasks(provider_state)
|
||||
task_info = next((t for t in tasks if t.id == completed_id), None)
|
||||
if task_info is not None and completed_id is not None:
|
||||
completed_task = runtime.in_flight_tasks.get(completed_id)
|
||||
if completed_task is not None:
|
||||
_finalize_task(task_info, completed_task, runtime)
|
||||
_save_tasks(provider_state, tasks)
|
||||
_save_provider_state(session, provider_state, source_id=source_id)
|
||||
|
||||
status_str = task_info.status.value if task_info else "Unknown"
|
||||
return f"Task {completed_id} finished with status: {status_str}."
|
||||
|
||||
@tool(name="background_agents_get_task_results", approval_mode="never_require")
|
||||
def background_agents_get_task_results(task_id: int) -> str:
|
||||
"""Get the text output of a background task by its ID."""
|
||||
tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
|
||||
task_info = next((t for t in tasks if t.id == task_id), None)
|
||||
|
||||
if task_info is None:
|
||||
return f"Error: No task found with ID {task_id}."
|
||||
|
||||
if task_info.status == BackgroundTaskStatus.COMPLETED:
|
||||
return task_info.result_text or "(no output)"
|
||||
if task_info.status == BackgroundTaskStatus.FAILED:
|
||||
return f"Task failed: {task_info.error_text or 'Unknown error'}"
|
||||
if task_info.status == BackgroundTaskStatus.LOST:
|
||||
return "Task state was lost (reference unavailable)."
|
||||
if task_info.status == BackgroundTaskStatus.RUNNING:
|
||||
return f"Task {task_id} is still running."
|
||||
return f"Task {task_id} has status: {task_info.status.value}."
|
||||
|
||||
@tool(name="background_agents_get_all_tasks", approval_mode="never_require")
|
||||
def background_agents_get_all_tasks() -> str:
|
||||
"""List all background tasks with their IDs, statuses, agent names, and descriptions."""
|
||||
tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
|
||||
|
||||
if not tasks:
|
||||
return "No tasks."
|
||||
|
||||
lines = ["Tasks:"]
|
||||
for t in tasks:
|
||||
lines.append(f"- Task {t.id} [{t.status.value}] ({t.agent_name}): {t.description}")
|
||||
return "\n".join(lines)
|
||||
|
||||
@tool(name="background_agents_continue_task", approval_mode="never_require")
|
||||
def background_agents_continue_task(task_id: int, text: str) -> str:
|
||||
"""Send follow-up input to a completed or failed task to resume its work."""
|
||||
tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
|
||||
task_info = next((t for t in tasks if t.id == task_id), None)
|
||||
|
||||
if task_info is None:
|
||||
return f"Error: No task found with ID {task_id}."
|
||||
|
||||
if task_info.status == BackgroundTaskStatus.LOST:
|
||||
return (
|
||||
f"Error: Task {task_id} cannot be continued because its session was lost. Start a new task instead."
|
||||
)
|
||||
|
||||
if task_info.status == BackgroundTaskStatus.RUNNING:
|
||||
return f"Error: Task {task_id} is still running. Wait for it to complete before continuing."
|
||||
|
||||
key = task_info.agent_name.lower()
|
||||
if key not in self._agents:
|
||||
return f"Error: Agent '{task_info.agent_name}' is no longer available."
|
||||
|
||||
sub_session = runtime.background_sessions.get(task_id)
|
||||
if sub_session is None:
|
||||
return f"Error: Session for task {task_id} is no longer available."
|
||||
|
||||
bg_agent = self._agents[key]
|
||||
|
||||
# Reset task state and start a new run on the existing session.
|
||||
task_info.status = BackgroundTaskStatus.RUNNING
|
||||
task_info.result_text = None
|
||||
task_info.error_text = None
|
||||
_save_tasks(provider_state, tasks)
|
||||
|
||||
async_task = asyncio.create_task(_run_agent(bg_agent.run(text, session=sub_session)))
|
||||
runtime.in_flight_tasks[task_id] = async_task
|
||||
|
||||
_save_provider_state(session, provider_state, source_id=source_id)
|
||||
return f"Task {task_id} continued with new input."
|
||||
|
||||
background_agents_continue_task._invoke_sync_on_event_loop = True # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
@tool(name="background_agents_clear_completed_task", approval_mode="never_require")
|
||||
def background_agents_clear_completed_task(task_id: int) -> str:
|
||||
"""Remove a completed or failed task and release its session to free memory."""
|
||||
tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
|
||||
task_info = next((t for t in tasks if t.id == task_id), None)
|
||||
|
||||
if task_info is None:
|
||||
return f"Error: No task found with ID {task_id}."
|
||||
|
||||
if task_info.status == BackgroundTaskStatus.RUNNING:
|
||||
return f"Error: Task {task_id} is still running. Wait for it to complete before clearing."
|
||||
|
||||
# Remove the task from state.
|
||||
tasks = [t for t in tasks if t.id != task_id]
|
||||
_save_tasks(provider_state, tasks)
|
||||
|
||||
# Clean up runtime references.
|
||||
runtime.in_flight_tasks.pop(task_id, None)
|
||||
runtime.background_sessions.pop(task_id, None)
|
||||
|
||||
_save_provider_state(session, provider_state, source_id=source_id)
|
||||
return f"Task {task_id} cleared."
|
||||
|
||||
# Inject instructions and current task status.
|
||||
context.extend_instructions(self.source_id, [self._instructions])
|
||||
context.extend_tools(
|
||||
self.source_id,
|
||||
[
|
||||
background_agents_start_task,
|
||||
background_agents_wait_for_first_completion,
|
||||
background_agents_get_task_results,
|
||||
background_agents_get_all_tasks,
|
||||
background_agents_continue_task,
|
||||
background_agents_clear_completed_task,
|
||||
],
|
||||
)
|
||||
|
||||
# Include current task status as context message if there are tasks.
|
||||
# Refresh first to get accurate statuses for any tasks that completed between turns.
|
||||
tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
|
||||
if tasks:
|
||||
status_lines = ["### Current background tasks"]
|
||||
for t in tasks:
|
||||
status_lines.append(f"- Task {t.id} [{t.status.value}] ({t.agent_name}): {t.description}")
|
||||
context.extend_messages(
|
||||
self.source_id,
|
||||
[Message(role="user", contents=["\n".join(status_lines)])],
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,533 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""File-based memory harness provider backed by an ``AgentFileStore``.
|
||||
|
||||
:class:`FileMemoryProvider` gives an agent a session-scoped, file-based memory
|
||||
system. Each memory is stored as an individual file with a meaningful name, and
|
||||
large files can carry a companion description file (suffixed with
|
||||
``_description.md``) that provides a short summary used for discovery. A
|
||||
``memories.md`` index file is maintained automatically and injected into the
|
||||
agent's context so the model knows what memories already exist.
|
||||
|
||||
File access is mediated through the :class:`~agent_framework.AgentFileStore`
|
||||
abstraction (shared with :class:`~agent_framework.FileAccessProvider`), so the
|
||||
same in-memory, local-disk, or remote-blob backends can be reused here.
|
||||
|
||||
Unlike :class:`~agent_framework.FileAccessProvider`, which exposes a *shared*
|
||||
store visible across sessions and agents, :class:`FileMemoryProvider` isolates
|
||||
memories per session by default: every session writes under its own working
|
||||
folder (derived from the session id). Pass an explicit ``scope`` to group
|
||||
memories differently, for example by user id.
|
||||
|
||||
The provider exposes the following tools to the agent (registered on the
|
||||
per-invocation :class:`~agent_framework.SessionContext` in
|
||||
:meth:`FileMemoryProvider.before_run`):
|
||||
|
||||
* ``file_memory_write`` — Write a memory file (with an optional description).
|
||||
* ``file_memory_read`` — Read the content of a memory file by name.
|
||||
* ``file_memory_delete`` — Delete a memory file (and its description).
|
||||
* ``file_memory_ls`` — List memory files with their descriptions.
|
||||
* ``file_memory_grep`` — Search memory file contents with a regex.
|
||||
* ``file_memory_replace`` — Replace a substring within a memory file.
|
||||
* ``file_memory_replace_lines`` — Replace whole lines within a memory file.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Annotated, Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .._feature_stage import ExperimentalFeature, experimental
|
||||
from .._sessions import AgentSession, ContextProvider, SessionContext
|
||||
from .._tools import tool
|
||||
from .._types import Message
|
||||
from ._file_access import (
|
||||
AgentFileStore,
|
||||
FileStoreEntry,
|
||||
_apply_replace, # pyright: ignore[reportPrivateUsage]
|
||||
_apply_replace_lines, # pyright: ignore[reportPrivateUsage]
|
||||
_line_edits, # pyright: ignore[reportPrivateUsage]
|
||||
_matches_glob, # pyright: ignore[reportPrivateUsage]
|
||||
_normalize_relative_path, # pyright: ignore[reportPrivateUsage]
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_FILE_MEMORY_SOURCE_ID = "file_memory"
|
||||
|
||||
DEFAULT_FILE_MEMORY_INSTRUCTIONS = (
|
||||
"## File Based Memory\n"
|
||||
"You have access to a session-scoped, file-based memory system via the `file_memory_*` tools "
|
||||
"for storing and retrieving information across interactions. "
|
||||
"These files act as your working memory for the current session and are isolated from other sessions. "
|
||||
"Use these tools to store plans, memories, processing results, or downloaded data.\n\n"
|
||||
'- Use descriptive file names (e.g., "projectarchitecture.md", "userpreferences.md").\n'
|
||||
"- Include a description when writing a file to help with future discovery.\n"
|
||||
"- Before starting new tasks, use file_memory_ls and file_memory_grep to check for "
|
||||
"relevant existing memories to avoid duplicate work.\n"
|
||||
"- Keep memories up-to-date by overwriting files when information changes.\n"
|
||||
"- When you receive large amounts of data (e.g., downloaded web pages, API responses, research results), "
|
||||
"write them to files if they will be required later, so that they are not lost when older context is "
|
||||
"compacted or truncated. This ensures important data remains accessible across long-running sessions."
|
||||
)
|
||||
|
||||
_DESCRIPTION_SUFFIX = "_description.md"
|
||||
_MEMORY_INDEX_FILE_NAME = "memories.md"
|
||||
_MAX_INDEX_ENTRIES = 50
|
||||
|
||||
|
||||
def _description_file_name(file_name: str) -> str:
|
||||
"""Return the companion description file name for ``file_name``.
|
||||
|
||||
The suffix replaces the original extension when present (so ``notes.md``
|
||||
becomes ``notes_description.md``); otherwise it is appended.
|
||||
"""
|
||||
dot_index = file_name.rfind(".")
|
||||
if dot_index > 0:
|
||||
return f"{file_name[:dot_index]}{_DESCRIPTION_SUFFIX}"
|
||||
return f"{file_name}{_DESCRIPTION_SUFFIX}"
|
||||
|
||||
|
||||
def _is_internal_file(file_name: str) -> bool:
|
||||
"""Return whether ``file_name`` is an internal file hidden from the agent.
|
||||
|
||||
Internal files are the description sidecars and the ``memories.md`` index.
|
||||
"""
|
||||
lowered = file_name.lower()
|
||||
return lowered.endswith(_DESCRIPTION_SUFFIX) or lowered == _MEMORY_INDEX_FILE_NAME
|
||||
|
||||
|
||||
def _combine_paths(base_path: str, relative_path: str) -> str:
|
||||
"""Join a working-folder path with a relative path using forward slashes."""
|
||||
if not base_path:
|
||||
return relative_path
|
||||
if not relative_path:
|
||||
return base_path
|
||||
return f"{base_path.rstrip('/')}/{relative_path.lstrip('/')}"
|
||||
|
||||
|
||||
def _is_nested_path(normalized_file_name: str) -> bool:
|
||||
"""Return whether a normalized file name points into a subdirectory.
|
||||
|
||||
File memory is a flat, session-scoped space: every discovery surface
|
||||
(the ``memories.md`` index, ``file_memory_ls``, and non-recursive
|
||||
``file_memory_grep``) only enumerates direct children of the
|
||||
working folder. A nested name such as ``"notes/plan.md"`` would therefore
|
||||
be written but never surface again, so such names are rejected up front.
|
||||
``_normalize_relative_path`` already converts backslashes to forward
|
||||
slashes, so checking for ``/`` covers both separators.
|
||||
"""
|
||||
return "/" in normalized_file_name
|
||||
|
||||
|
||||
class _WriteFileInput(BaseModel):
|
||||
"""Input schema for ``file_memory_write``."""
|
||||
|
||||
file_name: Annotated[str, Field(description="Flat file name to write under; must not contain path separators.")]
|
||||
content: Annotated[str, Field(description="Full text content to write to the file.")]
|
||||
description: Annotated[
|
||||
str | None,
|
||||
Field(
|
||||
default=None,
|
||||
description="Optional summary used to aid future discovery; recommended for large files.",
|
||||
),
|
||||
] = None
|
||||
|
||||
|
||||
class _ReadFileInput(BaseModel):
|
||||
"""Input schema for ``file_memory_read``."""
|
||||
|
||||
file_name: Annotated[str, Field(description="Name of the memory file to read.")]
|
||||
|
||||
|
||||
class _DeleteFileInput(BaseModel):
|
||||
"""Input schema for ``file_memory_delete``."""
|
||||
|
||||
file_name: Annotated[str, Field(description="Name of the memory file to delete.")]
|
||||
|
||||
|
||||
class _ListInput(BaseModel):
|
||||
"""Input schema for ``file_memory_ls``."""
|
||||
|
||||
glob_pattern: Annotated[
|
||||
str | None,
|
||||
Field(
|
||||
default=None,
|
||||
description='Optional glob (e.g. "*.md") matched against file names to filter the listing.',
|
||||
),
|
||||
] = None
|
||||
|
||||
|
||||
class _ReplaceInput(BaseModel):
|
||||
"""Input schema for ``file_memory_replace``."""
|
||||
|
||||
file_name: Annotated[str, Field(description="Name of the memory file to modify.")]
|
||||
old_string: Annotated[str, Field(description="Substring to find and replace.")]
|
||||
new_string: Annotated[str, Field(description="Replacement text.")]
|
||||
replace_all: Annotated[
|
||||
bool,
|
||||
Field(
|
||||
default=False,
|
||||
description="When true, replace every occurrence; when false, fail unless exactly one occurrence exists.",
|
||||
),
|
||||
] = False
|
||||
|
||||
|
||||
class _LineEdit(BaseModel):
|
||||
"""A single literal line replacement for ``file_memory_replace_lines``."""
|
||||
|
||||
line_number: Annotated[int, Field(description="1-based line number to replace.")]
|
||||
new_line: Annotated[
|
||||
str,
|
||||
Field(
|
||||
description=(
|
||||
"Literal replacement text for the line, including any trailing newline you want to keep "
|
||||
"(the editor does not add one). Set to an empty string to delete the line entirely, "
|
||||
"including its line break."
|
||||
)
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class _ReplaceLinesInput(BaseModel):
|
||||
"""Input schema for ``file_memory_replace_lines``."""
|
||||
|
||||
file_name: Annotated[str, Field(description="Name of the memory file to modify.")]
|
||||
edits: Annotated[
|
||||
list[_LineEdit],
|
||||
Field(description="List of 1-based line numbers and their literal replacement text."),
|
||||
]
|
||||
|
||||
|
||||
class _SearchFilesInput(BaseModel):
|
||||
"""Input schema for ``file_memory_grep``."""
|
||||
|
||||
regex_pattern: Annotated[
|
||||
str,
|
||||
Field(description="Case-insensitive regex matched against file contents; 256 characters or fewer."),
|
||||
]
|
||||
glob_pattern: Annotated[
|
||||
str | None,
|
||||
Field(
|
||||
default=None,
|
||||
description='Optional glob to filter which files are searched (e.g. "*.md", "research*").',
|
||||
),
|
||||
] = None
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class FileMemoryProvider(ContextProvider):
|
||||
"""Context provider that gives an agent session-scoped, file-based memory.
|
||||
|
||||
The provider exposes the following tools to the agent via the per-invocation
|
||||
:class:`~agent_framework.SessionContext`:
|
||||
|
||||
- ``file_memory_write`` — Write a memory file with an optional description.
|
||||
- ``file_memory_read`` — Read the content of a memory file by name.
|
||||
- ``file_memory_delete`` — Delete a memory file and its description.
|
||||
- ``file_memory_ls`` — List memory files with their descriptions.
|
||||
- ``file_memory_grep`` — Search memory file contents with a regex.
|
||||
- ``file_memory_replace`` — Replace a substring within a memory file.
|
||||
- ``file_memory_replace_lines`` — Replace whole lines within a memory file.
|
||||
|
||||
Memories are isolated per session: each session reads and writes under a
|
||||
working folder derived from its session id. Pass an explicit ``scope`` to
|
||||
group memories differently (for example, per user id) across sessions.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: AgentFileStore,
|
||||
*,
|
||||
source_id: str = DEFAULT_FILE_MEMORY_SOURCE_ID,
|
||||
scope: str | None = None,
|
||||
instructions: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize the file memory provider.
|
||||
|
||||
Args:
|
||||
store: The file store implementation used for storage operations.
|
||||
|
||||
Keyword Args:
|
||||
source_id: Unique source ID for the provider.
|
||||
scope: The namespace that logically groups and isolates memories
|
||||
(for example, a user ID). Used as the working folder within the
|
||||
store. When ``None`` (the default), the active session's
|
||||
``session_id`` is used, isolating memories per session.
|
||||
instructions: Optional instruction override. When ``None`` the
|
||||
default file-memory instructions are used.
|
||||
"""
|
||||
super().__init__(source_id)
|
||||
self.store = store
|
||||
self.scope = scope
|
||||
self.instructions = instructions or DEFAULT_FILE_MEMORY_INSTRUCTIONS
|
||||
# Serializes write/delete operations (and their index rebuilds) so the
|
||||
# ``memories.md`` index stays consistent. A single per-instance lock is
|
||||
# sufficient for v1; concurrent writes across scopes are rare in practice.
|
||||
self._write_lock = asyncio.Lock()
|
||||
|
||||
def _resolve_working_folder(self, context: SessionContext) -> str:
|
||||
"""Resolve the working folder for the current invocation.
|
||||
|
||||
Uses the configured ``scope`` when set, otherwise the session id. The
|
||||
result is normalized as a relative directory path so it cannot escape
|
||||
the store root.
|
||||
"""
|
||||
raw_scope = self.scope or context.session_id or ""
|
||||
return _normalize_relative_path(raw_scope, is_directory=True)
|
||||
|
||||
async def _rebuild_index(self, working_folder: str) -> None:
|
||||
"""Rebuild the ``memories.md`` index for ``working_folder``.
|
||||
|
||||
Lists the non-internal files, sorts them deterministically, reads any
|
||||
companion descriptions, and writes a capped markdown summary.
|
||||
"""
|
||||
entries = await self.store.list_children(working_folder)
|
||||
file_names = [entry.name for entry in entries if entry.type == FileStoreEntry.FILE]
|
||||
sorted_files = sorted((name for name in file_names if not _is_internal_file(name)), key=str.lower)
|
||||
|
||||
lines = ["# Memory Index", ""]
|
||||
for file_name in sorted_files[:_MAX_INDEX_ENTRIES]:
|
||||
description = await self.store.read(_combine_paths(working_folder, _description_file_name(file_name)))
|
||||
if description and description.strip():
|
||||
lines.append(f"- **{file_name}**: {description.strip()}")
|
||||
else:
|
||||
lines.append(f"- **{file_name}**")
|
||||
|
||||
index_path = _combine_paths(working_folder, _MEMORY_INDEX_FILE_NAME)
|
||||
await self.store.write(index_path, "\n".join(lines) + "\n")
|
||||
|
||||
async def before_run(
|
||||
self,
|
||||
*,
|
||||
agent: Any,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Inject file-memory tools, instructions, and the memory index."""
|
||||
working_folder = self._resolve_working_folder(context)
|
||||
|
||||
if working_folder:
|
||||
await self.store.create_directory(working_folder)
|
||||
|
||||
@tool(name="file_memory_write", schema=_WriteFileInput, approval_mode="never_require")
|
||||
async def file_memory_write(file_name: str, content: str, description: str | None = None) -> str:
|
||||
"""Write a memory file with the given name and content. Overwrites the file if it already exists. Include a description for large files to provide a summary that helps with future discovery.""" # noqa: E501
|
||||
try:
|
||||
normalized = _normalize_relative_path(file_name)
|
||||
except ValueError as exc:
|
||||
return f"Could not write file '{file_name}': {exc}"
|
||||
if _is_nested_path(normalized):
|
||||
return (
|
||||
f"Could not write file '{file_name}': memory files must not be written into a "
|
||||
"subdirectory. Please choose a flat file name without path separators."
|
||||
)
|
||||
if _is_internal_file(normalized):
|
||||
return (
|
||||
f"Could not write file '{file_name}': the file name is reserved for internal use. "
|
||||
"Please choose a different file name."
|
||||
)
|
||||
|
||||
path = _combine_paths(working_folder, normalized)
|
||||
desc_path = _combine_paths(working_folder, _description_file_name(normalized))
|
||||
async with self._write_lock:
|
||||
try:
|
||||
await self.store.write(path, content)
|
||||
if description and description.strip():
|
||||
await self.store.write(desc_path, description)
|
||||
else:
|
||||
await self.store.delete(desc_path)
|
||||
await self._rebuild_index(working_folder)
|
||||
except ValueError as exc:
|
||||
return f"Could not write file '{file_name}': {exc}"
|
||||
except OSError as exc:
|
||||
return f"Could not write file '{file_name}': {exc.strerror or exc}"
|
||||
if description and description.strip():
|
||||
return f"File '{file_name}' written with description."
|
||||
return f"File '{file_name}' written."
|
||||
|
||||
@tool(name="file_memory_read", schema=_ReadFileInput, approval_mode="never_require")
|
||||
async def file_memory_read(file_name: str) -> str:
|
||||
"""Read the content of a memory file by name. Returns the file content or a message indicating the file was not found.""" # noqa: E501
|
||||
try:
|
||||
normalized = _normalize_relative_path(file_name)
|
||||
except ValueError as exc:
|
||||
return f"Could not read file '{file_name}': {exc}"
|
||||
if _is_nested_path(normalized):
|
||||
return f"File '{file_name}' not found."
|
||||
try:
|
||||
content = await self.store.read(_combine_paths(working_folder, normalized))
|
||||
except ValueError as exc:
|
||||
return f"Could not read file '{file_name}': {exc}"
|
||||
except OSError as exc:
|
||||
return f"Could not read file '{file_name}': {exc.strerror or exc}"
|
||||
return content if content is not None else f"File '{file_name}' not found."
|
||||
|
||||
@tool(name="file_memory_delete", schema=_DeleteFileInput, approval_mode="never_require")
|
||||
async def file_memory_delete(file_name: str) -> str:
|
||||
"""Delete a memory file by name. Also removes its companion description file if one exists."""
|
||||
try:
|
||||
normalized = _normalize_relative_path(file_name)
|
||||
except ValueError as exc:
|
||||
return f"Could not delete file '{file_name}': {exc}"
|
||||
if _is_nested_path(normalized):
|
||||
return f"File '{file_name}' not found."
|
||||
|
||||
path = _combine_paths(working_folder, normalized)
|
||||
desc_path = _combine_paths(working_folder, _description_file_name(normalized))
|
||||
async with self._write_lock:
|
||||
try:
|
||||
deleted = await self.store.delete(path)
|
||||
await self.store.delete(desc_path)
|
||||
await self._rebuild_index(working_folder)
|
||||
except ValueError as exc:
|
||||
return f"Could not delete file '{file_name}': {exc}"
|
||||
except OSError as exc:
|
||||
return f"Could not delete file '{file_name}': {exc.strerror or exc}"
|
||||
return f"File '{file_name}' deleted." if deleted else f"File '{file_name}' not found."
|
||||
|
||||
@tool(name="file_memory_ls", schema=_ListInput, approval_mode="never_require")
|
||||
async def file_memory_ls(glob_pattern: str | None = None) -> list[dict[str, Any]] | str:
|
||||
"""List all memory files with their descriptions (if available). Optionally filter file names with a glob_pattern (e.g. "*.md"). Internal files (description sidecars and the memory index) are not shown. Each entry is {"name": <name>, "type": "file", "description": <desc-or-null>}.""" # noqa: E501
|
||||
try:
|
||||
entries = await self.store.list_children(working_folder)
|
||||
except OSError as exc:
|
||||
return f"Could not list memory files: {exc.strerror or exc}"
|
||||
|
||||
file_names = [entry.name for entry in entries if entry.type == FileStoreEntry.FILE]
|
||||
available = set(file_names)
|
||||
results: list[dict[str, Any]] = []
|
||||
for file_name in file_names:
|
||||
if _is_internal_file(file_name):
|
||||
continue
|
||||
if not _matches_glob(file_name, glob_pattern):
|
||||
continue
|
||||
description: str | None = None
|
||||
desc_file_name = _description_file_name(file_name)
|
||||
if desc_file_name in available:
|
||||
description = await self.store.read(_combine_paths(working_folder, desc_file_name))
|
||||
results.append({"name": file_name, "type": "file", "description": description})
|
||||
return results
|
||||
|
||||
@tool(name="file_memory_replace", schema=_ReplaceInput, approval_mode="never_require")
|
||||
async def file_memory_replace(
|
||||
file_name: str, old_string: str, new_string: str, replace_all: bool = False
|
||||
) -> str:
|
||||
"""Replace occurrences of old_string with new_string in a memory file. Fails if old_string is not found, or if it occurs more than once and replace_all is false. Returns the number of occurrences replaced.""" # noqa: E501
|
||||
try:
|
||||
normalized = _normalize_relative_path(file_name)
|
||||
except ValueError as exc:
|
||||
return f"Could not replace in file '{file_name}': {exc}"
|
||||
if _is_nested_path(normalized):
|
||||
return f"File '{file_name}' not found."
|
||||
if _is_internal_file(normalized):
|
||||
return (
|
||||
f"Could not replace in file '{file_name}': the file name is reserved for internal use. "
|
||||
"Please choose a different file name."
|
||||
)
|
||||
path = _combine_paths(working_folder, normalized)
|
||||
async with self._write_lock:
|
||||
try:
|
||||
content = await self.store.read(path)
|
||||
if content is None:
|
||||
return f"File '{file_name}' not found."
|
||||
new_content, count = _apply_replace(content, old_string, new_string, replace_all)
|
||||
await self.store.write(path, new_content)
|
||||
except ValueError as exc:
|
||||
return f"Could not replace in file '{file_name}': {exc}"
|
||||
except OSError as exc:
|
||||
return f"Could not replace in file '{file_name}': {exc.strerror or exc}"
|
||||
return f"Replaced {count} occurrence(s) in '{file_name}'."
|
||||
|
||||
@tool(name="file_memory_replace_lines", schema=_ReplaceLinesInput, approval_mode="never_require")
|
||||
async def file_memory_replace_lines(file_name: str, edits: list[_LineEdit]) -> str:
|
||||
"""Replace lines in a memory file. Provide a list of edits, each with a 1-based line_number and a literal new_line (include your own trailing newline); an empty new_line deletes the line, including its line break. Fails on out-of-range or duplicate line numbers.""" # noqa: E501
|
||||
try:
|
||||
normalized = _normalize_relative_path(file_name)
|
||||
except ValueError as exc:
|
||||
return f"Could not edit file '{file_name}': {exc}"
|
||||
if _is_nested_path(normalized):
|
||||
return f"File '{file_name}' not found."
|
||||
if _is_internal_file(normalized):
|
||||
return (
|
||||
f"Could not edit file '{file_name}': the file name is reserved for internal use. "
|
||||
"Please choose a different file name."
|
||||
)
|
||||
path = _combine_paths(working_folder, normalized)
|
||||
async with self._write_lock:
|
||||
try:
|
||||
content = await self.store.read(path)
|
||||
if content is None:
|
||||
return f"File '{file_name}' not found."
|
||||
new_content = _apply_replace_lines(content, _line_edits(edits))
|
||||
await self.store.write(path, new_content)
|
||||
except ValueError as exc:
|
||||
return f"Could not edit file '{file_name}': {exc}"
|
||||
except OSError as exc:
|
||||
return f"Could not edit file '{file_name}': {exc.strerror or exc}"
|
||||
return f"Replaced {len(edits)} line(s) in '{file_name}'."
|
||||
|
||||
@tool(name="file_memory_grep", schema=_SearchFilesInput, approval_mode="never_require")
|
||||
async def file_memory_grep(
|
||||
regex_pattern: str,
|
||||
glob_pattern: str | None = None,
|
||||
) -> list[dict[str, Any]] | str:
|
||||
"""Search memory file contents using a case-insensitive regular expression. Optionally filter which files to search using a glob pattern (e.g., "*.md", "research*"). Returns matching file names, content snippets, and matching lines with line numbers. The regex_pattern must be 256 characters or fewer.""" # noqa: E501
|
||||
glob_filter = glob_pattern if glob_pattern and glob_pattern.strip() else None
|
||||
try:
|
||||
results = await self.store.search(working_folder, regex_pattern, glob_filter, recursive=False)
|
||||
except ValueError as exc:
|
||||
return f"Could not search memory files: {exc}"
|
||||
except OSError as exc:
|
||||
return f"Could not search memory files: {exc.strerror or exc}"
|
||||
return [result.to_dict() for result in results if not _is_internal_file(result.file_name)]
|
||||
|
||||
context.extend_instructions(self.source_id, [self.instructions])
|
||||
context.extend_tools(
|
||||
self.source_id,
|
||||
[
|
||||
file_memory_write,
|
||||
file_memory_read,
|
||||
file_memory_delete,
|
||||
file_memory_ls,
|
||||
file_memory_grep,
|
||||
file_memory_replace,
|
||||
file_memory_replace_lines,
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
index_content = await self.store.read(_combine_paths(working_folder, _MEMORY_INDEX_FILE_NAME))
|
||||
except (OSError, ValueError) as exc:
|
||||
# A corrupt/unavailable index (e.g. non-UTF8 bytes on disk or a store
|
||||
# error) must not block the run. Skip index injection for this run; it
|
||||
# self-heals on the next successful write/delete that rebuilds the index.
|
||||
logger.warning("Could not read memory index; skipping index injection: %s", exc)
|
||||
index_content = None
|
||||
if index_content and index_content.strip():
|
||||
context.extend_messages(
|
||||
self.source_id,
|
||||
[
|
||||
Message(
|
||||
role="user",
|
||||
contents=[
|
||||
(
|
||||
"The following is your memory index — a list of files you have previously written. "
|
||||
"You can read any of these files using the file_memory_read tool.\n\n"
|
||||
f"{index_content}"
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_FILE_MEMORY_INSTRUCTIONS",
|
||||
"DEFAULT_FILE_MEMORY_SOURCE_ID",
|
||||
"FileMemoryProvider",
|
||||
]
|
||||
@@ -0,0 +1,972 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AgentLoopMiddleware: re-run an agent in a loop until a criterion is met.
|
||||
|
||||
This module provides :class:`AgentLoopMiddleware`, an :class:`~agent_framework.AgentMiddleware`
|
||||
that repeatedly re-invokes the wrapped agent while a ``should_continue`` predicate says to keep
|
||||
going. It serves two common patterns through a single configurable class:
|
||||
|
||||
1. A user-supplied ``should_continue`` predicate - for example, keep looping while a response does
|
||||
not yet contain a completion marker, while a :class:`~agent_framework.TodoProvider` still has
|
||||
open items, or while a :class:`~agent_framework.BackgroundAgentsProvider` still has running
|
||||
tasks (see the :func:`todos_remaining` and :func:`background_tasks_running` helpers, which resolve
|
||||
their provider from the running agent). The loop
|
||||
can track a **feedback log** across iterations (``record_feedback``): each pass contributes an
|
||||
entry that is exposed to every callback via the ``progress`` keyword and (by default) injected
|
||||
into the next iteration's input. Set ``fresh_context=True`` to restart each pass from the
|
||||
original task plus the progress log (with a session attached, the session is also snapshotted
|
||||
before the loop and restored between iterations so no accumulated history leaks back in).
|
||||
``max_iterations`` bounds the loop as a safety cap.
|
||||
2. A chat-client judge (via :meth:`AgentLoopMiddleware.with_judge`) - a second chat client decides
|
||||
whether the user's original request has been answered (via a :class:`JudgeVerdict` structured
|
||||
output); the loop continues while the answer is "no". This is a convenience wrapper that builds an
|
||||
async ``should_continue`` predicate, so it is a special case of (1).
|
||||
|
||||
In every case, the input for the next iteration is controlled by the ``next_message`` callable.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from typing import TYPE_CHECKING, Any, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import Self
|
||||
|
||||
from .._feature_stage import ExperimentalFeature, experimental
|
||||
from .._middleware import AgentContext, AgentMiddleware, MiddlewareTermination
|
||||
from .._types import (
|
||||
AgentResponse,
|
||||
AgentResponseUpdate,
|
||||
AgentRunInputs,
|
||||
Message,
|
||||
ResponseStream,
|
||||
UsageDetails,
|
||||
add_usage_details,
|
||||
normalize_messages,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .._clients import SupportsChatGetResponse
|
||||
|
||||
__all__ = [
|
||||
"AgentLoopMiddleware",
|
||||
"JudgeVerdict",
|
||||
"background_tasks_running",
|
||||
"background_tasks_running_message",
|
||||
"todos_remaining",
|
||||
"todos_remaining_message",
|
||||
]
|
||||
|
||||
DEFAULT_NEXT_MESSAGE = "Continue working on the task. If it is complete, say so."
|
||||
|
||||
# Placeholder substituted with the rendered ``criteria`` block in judge instructions (see
|
||||
# :meth:`AgentLoopMiddleware.with_judge`). User-supplied instructions may include it to control
|
||||
# where the criteria are inserted; if absent, the criteria are not added to the judge instructions.
|
||||
CRITERIA_PLACEHOLDER = "{{criteria}}"
|
||||
|
||||
# Verdict markers the judge is asked to emit for clients that do not honor structured output. They
|
||||
# are deliberately non-overlapping: neither marker is a substring of the other, nor of the JSON
|
||||
# field name ``answered``, so the text fallback in :func:`_build_judge_condition` cannot misclassify
|
||||
# a negative verdict (e.g. ``{"answered": false}``) as a positive one.
|
||||
JUDGE_VERDICT_DONE = "VERDICT: DONE"
|
||||
JUDGE_VERDICT_MORE = "VERDICT: MORE"
|
||||
|
||||
DEFAULT_JUDGE_INSTRUCTIONS = (
|
||||
"You are an evaluator. You are given a user's original request and an agent's latest response. "
|
||||
"Decide whether the agent has fully addressed the original request. "
|
||||
"Set 'answered' to true if the request has been fully addressed, or false if more work is still "
|
||||
"required, and use 'reasoning' to briefly justify your decision. "
|
||||
f"If you cannot return structured output, end your reply with a line reading exactly "
|
||||
f"'{JUDGE_VERDICT_DONE}' when the request has been fully addressed or '{JUDGE_VERDICT_MORE}' "
|
||||
f"when more work is still required."
|
||||
"{{criteria}}"
|
||||
)
|
||||
|
||||
|
||||
def _render_criteria_block(criteria: Sequence[str] | None) -> str:
|
||||
"""Render a list of criteria into a bullet block for the judge instructions (``""`` if none)."""
|
||||
if not criteria:
|
||||
return ""
|
||||
bullets = "\n".join(f"- {item}" for item in criteria)
|
||||
return f"\n\nThe response must satisfy all of the following criteria:\n{bullets}"
|
||||
|
||||
|
||||
def _criteria_agent_instruction(criteria: Sequence[str]) -> str:
|
||||
"""Render the criteria into an extra instruction injected for the agent before each run."""
|
||||
bullets = "\n".join(f"- {item}" for item in criteria)
|
||||
return f"Your response must satisfy all of the following criteria:\n{bullets}"
|
||||
|
||||
|
||||
class JudgeVerdict(BaseModel):
|
||||
"""Structured verdict returned by the judge chat client."""
|
||||
|
||||
answered: bool = Field(
|
||||
description=(
|
||||
"True if the agent has fully addressed the original request and it adheres to the other "
|
||||
"judging standards, otherwise False."
|
||||
),
|
||||
)
|
||||
reasoning: str = Field(
|
||||
default="",
|
||||
description="Brief justification for the verdict.",
|
||||
)
|
||||
|
||||
|
||||
# Default iteration cap applied when ``max_iterations`` is not provided. Loops are bounded by
|
||||
# default to guard against runaway re-invocation; pass ``max_iterations=None`` explicitly to opt
|
||||
# into an unbounded loop.
|
||||
DEFAULT_MAX_ITERATIONS = 10
|
||||
|
||||
# Default iteration cap for judge-driven loops. LLM-judged loops are costly and probabilistic, so
|
||||
# they are bounded by a smaller default. Pass ``max_iterations=None`` explicitly to opt into an
|
||||
# unbounded judge loop.
|
||||
DEFAULT_JUDGE_MAX_ITERATIONS = 5
|
||||
|
||||
|
||||
# A callable invoked between iterations. It always receives the loop keyword arguments
|
||||
# (``iteration``, ``last_result``, ``messages``, ``original_messages``, ``session``, ``agent``,
|
||||
# ``progress``, ``feedback``). Callers declare only the keywords they need plus ``**kwargs`` to
|
||||
# ignore the rest. ``should_continue`` may return a plain ``bool`` (continue/stop) or a
|
||||
# ``(bool, str | None)`` tuple whose second item is feedback surfaced to the ``next_message`` and
|
||||
# ``record_feedback`` callables via the ``feedback`` keyword argument.
|
||||
ShouldContinueResult: TypeAlias = "bool | tuple[bool, str | None]"
|
||||
ShouldContinueCallable = Callable[..., "ShouldContinueResult | Awaitable[ShouldContinueResult]"]
|
||||
NextMessageCallable = Callable[..., "AgentRunInputs | Awaitable[AgentRunInputs | None] | None"]
|
||||
|
||||
# A callable invoked once per work iteration to capture a progress-log entry from that iteration. It
|
||||
# receives the loop keyword arguments and returns a string entry (appended to the log) or ``None``
|
||||
# (record nothing for that iteration).
|
||||
FeedbackCallable = Callable[..., "str | Awaitable[str | None] | None"]
|
||||
|
||||
|
||||
async def _maybe_await(value: Any) -> Any:
|
||||
"""Await ``value`` if it is awaitable, otherwise return it as-is."""
|
||||
if inspect.isawaitable(value):
|
||||
return await value
|
||||
return value
|
||||
|
||||
|
||||
def _build_judge_condition(
|
||||
judge_client: SupportsChatGetResponse,
|
||||
instructions: str,
|
||||
) -> tuple[ShouldContinueCallable, NextMessageCallable]:
|
||||
"""Build the ``should_continue`` predicate and ``next_message`` callable for a judge loop.
|
||||
|
||||
The judge is called directly (no agent tools, session, or middleware) with fresh messages, so
|
||||
the loop's evaluation cannot recurse back through the agent pipeline. The original input messages
|
||||
are forwarded verbatim (rather than collapsed to text) so multi-modal requests are preserved. The
|
||||
judge is asked for a :class:`JudgeVerdict` structured output; if the client does not honor
|
||||
structured output the verdict falls back to the explicit, non-overlapping ``VERDICT: DONE`` /
|
||||
``VERDICT: MORE`` markers (``MORE`` wins, keeping the loop running, when the marker is ambiguous
|
||||
or absent).
|
||||
|
||||
The predicate returns a ``(continue, reasoning)`` tuple; the loop surfaces that ``reasoning`` to
|
||||
the next-message callable as the ``feedback`` keyword argument, which feeds it back to the agent
|
||||
so it knows *why* its previous answer was judged incomplete.
|
||||
"""
|
||||
|
||||
async def _judge(
|
||||
*, last_result: AgentResponse, original_messages: list[Message], **kwargs: Any
|
||||
) -> tuple[bool, str | None]:
|
||||
judge_messages = [
|
||||
Message(role="system", contents=[instructions]),
|
||||
Message(
|
||||
role="user",
|
||||
contents=["Evaluate the agent's work. The user's original request follows:"],
|
||||
),
|
||||
*original_messages,
|
||||
Message(role="user", contents=["The agent's latest response was:"]),
|
||||
*last_result.messages,
|
||||
Message(role="user", contents=["Has the original request been fully addressed?"]),
|
||||
]
|
||||
response = await judge_client.get_response(judge_messages, options={"response_format": JudgeVerdict})
|
||||
verdict = response.value
|
||||
if isinstance(verdict, JudgeVerdict):
|
||||
answered = verdict.answered
|
||||
reasoning = verdict.reasoning
|
||||
else:
|
||||
# Fallback for clients that do not honor structured output: look for the explicit,
|
||||
# non-overlapping verdict markers. ``FAIL`` (more work needed) takes precedence so an
|
||||
# ambiguous or marker-less reply keeps looping rather than stopping on an incomplete
|
||||
# answer.
|
||||
text = response.text.upper()
|
||||
# ``MORE`` (more work needed) takes precedence so an ambiguous reply keeps looping.
|
||||
answered = False if JUDGE_VERDICT_MORE in text else JUDGE_VERDICT_DONE in text
|
||||
reasoning = response.text.strip()
|
||||
# Continue looping while the request is not yet answered, surfacing the reasoning as feedback.
|
||||
return (not answered), (reasoning or None)
|
||||
|
||||
def _next_message(*, feedback: str | None = None, **kwargs: Any) -> AgentRunInputs:
|
||||
# Feed the judge's reasoning back to the agent so the next iteration addresses the gap.
|
||||
if feedback:
|
||||
return (
|
||||
"An evaluator reviewed your previous response and judged that it does not yet fully "
|
||||
f"address the original request.\n\nEvaluator feedback: {feedback}\n\n"
|
||||
"Revise and continue so the original request is fully addressed."
|
||||
)
|
||||
return DEFAULT_NEXT_MESSAGE
|
||||
|
||||
return _judge, _next_message
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class AgentLoopMiddleware(AgentMiddleware):
|
||||
"""Re-run an agent in a loop until a criterion is met (or never).
|
||||
|
||||
This middleware repeatedly invokes the wrapped agent. After each run it decides whether to run
|
||||
again based on ``should_continue`` and ``max_iterations``, and uses ``next_message`` to build
|
||||
the input for the next iteration. Use :meth:`with_judge` to drive the loop with a chat-client
|
||||
judge instead of a hand-written predicate.
|
||||
|
||||
By default a non-streaming run returns an aggregated :class:`~agent_framework.AgentResponse`
|
||||
containing every iteration's messages plus the injected ``next_message`` "nudge" messages (set
|
||||
``return_final_only=True`` to return only the last iteration's response). Streaming runs always
|
||||
yield each iteration's updates and emit the injected nudge messages as ``user`` updates between
|
||||
iterations.
|
||||
|
||||
The ``should_continue`` and ``next_message`` callables are invoked with keyword arguments, so a
|
||||
caller only needs to declare the ones it uses plus ``**kwargs``. The keywords are:
|
||||
|
||||
- ``iteration`` (int): the number of completed runs so far (1-based after the first run).
|
||||
- ``last_result`` (AgentResponse): the result of the iteration that just completed.
|
||||
- ``messages`` (list[Message]): the messages used for the iteration that just completed.
|
||||
- ``original_messages`` (list[Message]): the input used for the first iteration.
|
||||
- ``session`` (AgentSession | None): the active session, used by the provider helpers.
|
||||
- ``agent``: the agent being looped.
|
||||
- ``progress`` (list[str]): the feedback log accumulated so far (see ``record_feedback``).
|
||||
- ``feedback`` (str | None): the feedback string returned by ``should_continue`` for this
|
||||
iteration (``None`` when it returned a plain bool). ``should_continue`` may return either a
|
||||
``bool`` or a ``(bool, str | None)`` tuple; the string is surfaced here so ``next_message``
|
||||
and ``record_feedback`` can reference it.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Agent, AgentResponse
|
||||
from agent_framework._harness._loop import AgentLoopMiddleware
|
||||
|
||||
|
||||
async def should_continue(*, iteration: int, last_result: AgentResponse, **kwargs) -> bool:
|
||||
return iteration < 3 and "DONE" not in last_result.text
|
||||
|
||||
|
||||
agent = Agent(client=client, middleware=[AgentLoopMiddleware(should_continue)])
|
||||
|
||||
Note:
|
||||
``max_iterations`` acts as a safety cap and defaults to ``DEFAULT_MAX_ITERATIONS`` (10). Pass
|
||||
an explicit ``None`` to make the loop unbounded, in which case it relies entirely on
|
||||
``should_continue`` to stop, so make sure the predicate can eventually return ``False``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
should_continue: ShouldContinueCallable,
|
||||
*,
|
||||
max_iterations: int | None = DEFAULT_MAX_ITERATIONS,
|
||||
next_message: NextMessageCallable | None = None,
|
||||
record_feedback: FeedbackCallable | None = None,
|
||||
inject_progress: bool = True,
|
||||
fresh_context: bool = False,
|
||||
return_final_only: bool = False,
|
||||
additional_instructions: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize the agent loop middleware.
|
||||
|
||||
Args:
|
||||
should_continue: Predicate that decides whether to run the agent again. May be sync or
|
||||
async and is called with the loop keyword arguments (``iteration``, ``last_result``,
|
||||
``messages``, ``original_messages``, ``session``, ``agent``, ``progress``, and
|
||||
``feedback`` -- see the class docstring for what each one carries; declare only the
|
||||
ones you need plus ``**kwargs``). Return ``True``/``False`` to
|
||||
continue/stop, or a ``(bool, str | None)`` tuple to also provide feedback; the
|
||||
feedback string is surfaced to the ``next_message`` and ``record_feedback`` callables
|
||||
via the ``feedback`` keyword argument. To loop on a chat-client judge instead, build
|
||||
the middleware via :meth:`with_judge`.
|
||||
|
||||
Keyword Args:
|
||||
max_iterations: Maximum number of agent runs, used as a safety cap. Defaults to
|
||||
``DEFAULT_MAX_ITERATIONS`` (10); pass an explicit ``None`` for an unbounded loop, or
|
||||
a positive integer to set a custom cap. (The :meth:`with_judge` factory uses
|
||||
``DEFAULT_JUDGE_MAX_ITERATIONS`` (5) as its default instead.)
|
||||
next_message: Callable that produces the input for the next iteration, called with the
|
||||
loop keyword arguments. Defaults to a short "continue" nudge. Returning ``None``
|
||||
reuses the previous iteration's messages verbatim (in which case the progress log is
|
||||
*not* injected; see ``inject_progress``).
|
||||
record_feedback: Optional callable invoked once per work iteration to capture a feedback
|
||||
entry. Called as ``record_feedback(**loop_kwargs)`` and returns a
|
||||
string entry appended to the progress log, or ``None`` to record nothing for that
|
||||
iteration. When not provided, the iteration's response text (``last_result.text``) is
|
||||
recorded instead. The accumulated log is exposed to every callback via the
|
||||
``progress`` loop keyword argument. For production loops prefer a ``record_feedback``
|
||||
that returns a terse summary rather than relying on the full response text.
|
||||
inject_progress: When ``True`` (default), the accumulated progress log is injected into
|
||||
the next iteration's input as a single ``user`` message ("Progress so far: ..."). To
|
||||
avoid duplication, only the most recent entry is injected when a session is attached
|
||||
(the session already retains earlier turns); the full log is injected when there is
|
||||
no session or ``fresh_context`` is set. When ``False`` the log is only exposed via the
|
||||
``progress`` loop keyword argument and never injected automatically.
|
||||
fresh_context: When ``True``, each iteration starts from a clean context: ``context``
|
||||
messages are reset to the original input messages (plus the injected progress log)
|
||||
instead of accumulating the prior conversation. When a session is attached, the
|
||||
session is snapshotted once before the loop and restored to that pre-loop baseline
|
||||
before each subsequent iteration, so the local transcript and any service-side
|
||||
conversation id are reset too and the agent does not re-read the accumulated history.
|
||||
In-loop working-state mutations are discarded; pre-loop state is preserved; continuity
|
||||
is carried only by the progress log.
|
||||
return_final_only: Controls what a non-streaming run returns. When ``False`` (default),
|
||||
the returned :class:`~agent_framework.AgentResponse` aggregates every iteration: each
|
||||
iteration's response messages plus the injected ``next_message`` "nudge" messages
|
||||
(as ``user`` messages), so the caller sees the full back-and-forth. When ``True``,
|
||||
only the final iteration's :class:`~agent_framework.AgentResponse` is returned. This
|
||||
flag has no effect on streaming runs (the stream cannot know in advance which
|
||||
iteration is last); streaming always yields each iteration's updates and injects the
|
||||
``next_message`` messages as ``user`` updates between iterations.
|
||||
additional_instructions: Optional extra instruction injected as a ``system`` message
|
||||
ahead of the input messages before the agent runs. It becomes part of the original
|
||||
messages, so it is preserved across ``fresh_context`` resets and (with a session)
|
||||
persists server-side across iterations. Used by :meth:`with_judge` to tell the agent
|
||||
about the criteria its response must satisfy, but available to any loop.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``max_iterations`` is not ``None`` and is less than 1.
|
||||
"""
|
||||
if max_iterations is not None and max_iterations < 1:
|
||||
raise ValueError("max_iterations must be None or a positive integer (>= 1).")
|
||||
|
||||
self.max_iterations: int | None = max_iterations
|
||||
self.should_continue: ShouldContinueCallable = should_continue
|
||||
self.next_message = next_message
|
||||
self.record_feedback = record_feedback
|
||||
self.inject_progress = inject_progress
|
||||
self.fresh_context = fresh_context
|
||||
self.return_final_only = return_final_only
|
||||
self.additional_instructions = additional_instructions
|
||||
|
||||
@classmethod
|
||||
def with_judge(
|
||||
cls,
|
||||
judge_client: SupportsChatGetResponse,
|
||||
*,
|
||||
criteria: Sequence[str] | None = None,
|
||||
instructions: str | None = None,
|
||||
max_iterations: int | None = DEFAULT_JUDGE_MAX_ITERATIONS,
|
||||
next_message: NextMessageCallable | None = None,
|
||||
fresh_context: bool = False,
|
||||
) -> Self:
|
||||
"""Create a loop that continues until a judge chat client decides the request was answered.
|
||||
|
||||
Convenience factory for the judge pattern: ``judge_client`` is queried with a
|
||||
:class:`JudgeVerdict` structured-output response after each iteration and the loop continues
|
||||
while the request is *not* answered. The judge's ``reasoning`` is fed back to the agent as
|
||||
the next iteration's input (unless a custom ``next_message`` is provided), so the agent knows
|
||||
why its previous answer was judged incomplete. See :meth:`__init__` for the full meaning of
|
||||
each argument.
|
||||
|
||||
Security considerations:
|
||||
Using a judge is an explicit opt-in — the caller must supply a ``judge_client`` — and
|
||||
introduces a second external LLM boundary in addition to the agent's own model. On every
|
||||
iteration the judge is sent the original request and the agent's latest response, both of
|
||||
which may contain sensitive or untrusted content. A compromised or malicious judge
|
||||
endpoint could exfiltrate that data, or return a manipulated :class:`JudgeVerdict` / gap
|
||||
analysis that is fed back into the loop as feedback, potentially steering the agent via
|
||||
indirect prompt injection. Only configure a ``judge_client`` that points at a service you
|
||||
trust as much as the primary model.
|
||||
|
||||
Args:
|
||||
judge_client: Chat client used to judge whether the original request was answered.
|
||||
**Security:** this client is sent the original request and the agent's latest
|
||||
response on every iteration, so only point it at a service you trust as much as the
|
||||
primary model — see the security considerations above for the exfiltration and
|
||||
prompt-injection risks of an untrusted judge.
|
||||
|
||||
Keyword Args:
|
||||
criteria: Optional list of criteria the response must satisfy. When provided, they are
|
||||
(1) injected as an extra ``system`` instruction for the agent before it runs (via
|
||||
``additional_instructions``) and (2) rendered into the judge instructions wherever
|
||||
the ``{{criteria}}`` placeholder appears (``CRITERIA_PLACEHOLDER``).
|
||||
instructions: Optional system instructions for the judge. Defaults to
|
||||
``DEFAULT_JUDGE_INSTRUCTIONS``. May contain the ``{{criteria}}`` placeholder, which
|
||||
is replaced with the rendered ``criteria`` (or removed when no criteria are given).
|
||||
max_iterations: Maximum number of agent runs. Defaults to
|
||||
``DEFAULT_JUDGE_MAX_ITERATIONS`` (5); pass ``None`` for unbounded, or a positive
|
||||
integer to set a custom cap.
|
||||
next_message: Callable that produces the next iteration's input. Defaults to one that
|
||||
relays the judge's ``reasoning`` back to the agent.
|
||||
fresh_context: When ``True``, each iteration restarts from the original input messages
|
||||
(plus the injected progress log and judge feedback) instead of accumulating the prior
|
||||
conversation; an attached session is snapshotted before the loop and restored to that
|
||||
baseline between iterations. See :meth:`__init__` for the full semantics. Defaults to
|
||||
``False``.
|
||||
"""
|
||||
judge_instructions = (instructions or DEFAULT_JUDGE_INSTRUCTIONS).replace(
|
||||
CRITERIA_PLACEHOLDER, _render_criteria_block(criteria)
|
||||
)
|
||||
should_continue, judge_next_message = _build_judge_condition(judge_client, judge_instructions)
|
||||
return cls(
|
||||
should_continue=should_continue,
|
||||
max_iterations=max_iterations,
|
||||
next_message=next_message or judge_next_message,
|
||||
fresh_context=fresh_context,
|
||||
additional_instructions=_criteria_agent_instruction(criteria) if criteria else None,
|
||||
)
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: AgentContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Run the wrapped agent in a loop."""
|
||||
if self.additional_instructions is not None:
|
||||
# Inject the extra instruction as a system message ahead of the input so it is present
|
||||
# on every iteration and preserved across fresh_context resets (which restart from
|
||||
# ``original_messages``).
|
||||
context.messages = [
|
||||
Message(role="system", contents=[self.additional_instructions]),
|
||||
*context.messages,
|
||||
]
|
||||
original_messages = list(context.messages)
|
||||
# For a truly fresh context per iteration the session must also be reset, otherwise the
|
||||
# next run reloads the local transcript or re-threads the service-side conversation and the
|
||||
# model still sees the accumulated history. Snapshot the session once here (the pre-loop
|
||||
# baseline) and restore it before each subsequent iteration so every pass starts clean.
|
||||
snapshot = context.session.to_dict() if self.fresh_context and context.session is not None else None
|
||||
if context.stream:
|
||||
self._process_streaming(context, call_next, original_messages, snapshot)
|
||||
else:
|
||||
await self._process_non_streaming(context, call_next, original_messages, snapshot)
|
||||
|
||||
@staticmethod
|
||||
def _has_pending_approval_request(result: AgentResponse | None) -> bool:
|
||||
"""Return ``True`` if ``result`` carries a pending tool-approval request.
|
||||
|
||||
When the loop sits outermost (e.g. around a tool-approval middleware), an iteration may
|
||||
return a response that asks the caller to approve a tool call rather than a completed turn.
|
||||
In that case the loop must stop and hand the response back so a human can approve, instead
|
||||
of continuing or injecting the next message. This mirrors the C# ``LoopAgent`` escape hatch
|
||||
(``HasPendingApprovalRequests``). A pending request is any content whose ``type`` is
|
||||
``"function_approval_request"``.
|
||||
"""
|
||||
if result is None:
|
||||
return False
|
||||
return any(
|
||||
getattr(content, "type", None) == "function_approval_request"
|
||||
for message in result.messages
|
||||
for content in message.contents
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _restore_session(session: Any, snapshot: dict[str, Any]) -> None:
|
||||
"""Restore a session in place to a previously captured ``to_dict()`` snapshot.
|
||||
|
||||
Re-hydrates the snapshot via :meth:`AgentSession.from_dict` and copies the mutable fields
|
||||
(``service_session_id`` and ``state``) back onto the live ``session`` instance, so any
|
||||
reference held by the agent/context observes the reset. ``session_id`` is preserved (the
|
||||
snapshot carries the same id). A fresh ``from_dict`` is built on every call so repeated
|
||||
restores from one snapshot do not alias the same state dict.
|
||||
"""
|
||||
from .._sessions import AgentSession
|
||||
|
||||
restored = AgentSession.from_dict(snapshot)
|
||||
session.service_session_id = restored.service_session_id
|
||||
session.state = restored.state
|
||||
|
||||
async def _process_non_streaming(
|
||||
self,
|
||||
context: AgentContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
original_messages: list[Message],
|
||||
snapshot: dict[str, Any] | None,
|
||||
) -> None:
|
||||
iteration = 0
|
||||
work_iterations = 0
|
||||
progress: list[str] = []
|
||||
# Aggregated transcript across iterations: each iteration's response messages plus the
|
||||
# injected "nudge" messages, used to build the combined response when return_final_only=False.
|
||||
aggregated: list[Message] = []
|
||||
aggregated_usage: UsageDetails | None = None
|
||||
final_result: AgentResponse | None = None
|
||||
while True:
|
||||
await call_next()
|
||||
iteration += 1
|
||||
|
||||
result = context.result
|
||||
if not isinstance(result, AgentResponse):
|
||||
raise TypeError(
|
||||
"AgentLoopMiddleware expected an AgentResponse from a non-streaming run, "
|
||||
f"got {type(result).__name__}."
|
||||
)
|
||||
|
||||
final_result = result
|
||||
aggregated.extend(result.messages)
|
||||
if result.usage_details is not None:
|
||||
aggregated_usage = add_usage_details(aggregated_usage, result.usage_details)
|
||||
|
||||
# Escape hatch: if this iteration is asking for tool approval, stop and return the
|
||||
# response so the caller can approve, instead of continuing or injecting next_message.
|
||||
if self._has_pending_approval_request(result):
|
||||
break
|
||||
|
||||
messages_used = context.messages
|
||||
loop_kwargs = self._build_loop_kwargs(
|
||||
context=context,
|
||||
iteration=iteration,
|
||||
last_result=result,
|
||||
messages_used=messages_used,
|
||||
original_messages=original_messages,
|
||||
progress=progress,
|
||||
)
|
||||
|
||||
work_iterations += 1
|
||||
# Decide whether to stop and capture any feedback from should_continue first, so the
|
||||
# feedback is available to both the progress and next-message callables this iteration.
|
||||
stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations)
|
||||
loop_kwargs = self._build_loop_kwargs(
|
||||
context=context,
|
||||
iteration=iteration,
|
||||
last_result=result,
|
||||
messages_used=messages_used,
|
||||
original_messages=original_messages,
|
||||
progress=progress,
|
||||
feedback=feedback,
|
||||
)
|
||||
# Capture this iteration's progress entry, then refresh loop_kwargs so the next-message
|
||||
# resolution sees the latest entry.
|
||||
if await self._record_progress(result, loop_kwargs, progress):
|
||||
loop_kwargs = self._build_loop_kwargs(
|
||||
context=context,
|
||||
iteration=iteration,
|
||||
last_result=result,
|
||||
messages_used=messages_used,
|
||||
original_messages=original_messages,
|
||||
progress=progress,
|
||||
feedback=feedback,
|
||||
)
|
||||
if stop:
|
||||
break
|
||||
if snapshot is not None and context.session is not None:
|
||||
# Reset the session to the pre-loop baseline so the next run starts fresh; only the
|
||||
# progress log (injected by _resolve_next_message) carries continuity forward.
|
||||
self._restore_session(context.session, snapshot)
|
||||
next_messages = await self._resolve_next_message(loop_kwargs, messages_used, original_messages)
|
||||
context.messages = next_messages
|
||||
aggregated.extend(next_messages)
|
||||
|
||||
if not self.return_final_only:
|
||||
context.result = self._aggregate_response(final_result, aggregated, aggregated_usage)
|
||||
|
||||
def _process_streaming(
|
||||
self,
|
||||
context: AgentContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
original_messages: list[Message],
|
||||
snapshot: dict[str, Any] | None,
|
||||
) -> None:
|
||||
# Holds the last iteration's final response so the outer stream's finalizer can return it
|
||||
# rather than an aggregate of every iteration.
|
||||
holder: dict[str, AgentResponse | None] = {"final": None}
|
||||
|
||||
async def _generator() -> Any:
|
||||
iteration = 0
|
||||
work_iterations = 0
|
||||
progress: list[str] = []
|
||||
while True:
|
||||
try:
|
||||
await call_next()
|
||||
inner = context.result
|
||||
if not isinstance(inner, ResponseStream):
|
||||
raise TypeError(
|
||||
"AgentLoopMiddleware expected a ResponseStream from a streaming run, "
|
||||
f"got {type(inner).__name__}."
|
||||
)
|
||||
|
||||
async for update in inner:
|
||||
yield update
|
||||
|
||||
holder["final"] = await inner.get_final_response()
|
||||
except MiddlewareTermination:
|
||||
# The pipeline's MiddlewareTermination suppression is no longer active once
|
||||
# process() has returned (the stream is consumed lazily), so a termination
|
||||
# raised by a downstream middleware or during stream consumption surfaces here.
|
||||
# Stop cleanly and keep whatever final response we have from a prior iteration.
|
||||
return
|
||||
|
||||
iteration += 1
|
||||
|
||||
messages_used = context.messages
|
||||
final = holder["final"]
|
||||
# Escape hatch: if this iteration is asking for tool approval, stop the loop and
|
||||
# let the caller approve, instead of continuing or injecting next_message.
|
||||
if self._has_pending_approval_request(final):
|
||||
return
|
||||
loop_kwargs = self._build_loop_kwargs(
|
||||
context=context,
|
||||
iteration=iteration,
|
||||
last_result=final,
|
||||
messages_used=messages_used,
|
||||
original_messages=original_messages,
|
||||
progress=progress,
|
||||
)
|
||||
|
||||
work_iterations += 1
|
||||
# Decide whether to stop and capture any feedback from should_continue first, so the
|
||||
# feedback is available to both the progress and next-message callables this iteration.
|
||||
stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations)
|
||||
loop_kwargs = self._build_loop_kwargs(
|
||||
context=context,
|
||||
iteration=iteration,
|
||||
last_result=final,
|
||||
messages_used=messages_used,
|
||||
original_messages=original_messages,
|
||||
progress=progress,
|
||||
feedback=feedback,
|
||||
)
|
||||
if await self._record_progress(final, loop_kwargs, progress):
|
||||
loop_kwargs = self._build_loop_kwargs(
|
||||
context=context,
|
||||
iteration=iteration,
|
||||
last_result=final,
|
||||
messages_used=messages_used,
|
||||
original_messages=original_messages,
|
||||
progress=progress,
|
||||
feedback=feedback,
|
||||
)
|
||||
if stop:
|
||||
return
|
||||
if snapshot is not None and context.session is not None:
|
||||
# Reset the session to the pre-loop baseline before the next run. The final
|
||||
# response was already awaited above, so the service-side conversation id has
|
||||
# been propagated and is safe to discard here.
|
||||
self._restore_session(context.session, snapshot)
|
||||
next_messages = await self._resolve_next_message(loop_kwargs, messages_used, original_messages)
|
||||
context.messages = next_messages
|
||||
# Surface the injected "nudge" messages in the stream so consumers see the user
|
||||
# turns that drive each subsequent iteration (the equivalent of the aggregated
|
||||
# transcript that non-streaming runs return).
|
||||
for message in next_messages:
|
||||
yield self._message_to_update(message)
|
||||
|
||||
def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse:
|
||||
if holder["final"] is not None:
|
||||
return holder["final"]
|
||||
return AgentResponse.from_updates(updates)
|
||||
|
||||
context.result = ResponseStream(_generator(), finalizer=_finalize)
|
||||
|
||||
def _build_loop_kwargs(
|
||||
self,
|
||||
*,
|
||||
context: AgentContext,
|
||||
iteration: int,
|
||||
last_result: AgentResponse | None,
|
||||
messages_used: list[Message],
|
||||
original_messages: list[Message],
|
||||
progress: list[str],
|
||||
feedback: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"iteration": iteration,
|
||||
"last_result": last_result,
|
||||
"messages": messages_used,
|
||||
"original_messages": original_messages,
|
||||
"session": context.session,
|
||||
"agent": context.agent,
|
||||
# A copy so user callbacks cannot mutate the loop's internal progress log.
|
||||
"progress": list(progress),
|
||||
# Feedback returned by ``should_continue`` for this iteration (``None`` if it returned a
|
||||
# plain bool, or the stop was decided by ``max_iterations``).
|
||||
"feedback": feedback,
|
||||
}
|
||||
|
||||
async def _record_progress(
|
||||
self,
|
||||
last_result: AgentResponse | None,
|
||||
loop_kwargs: dict[str, Any],
|
||||
progress: list[str],
|
||||
) -> bool:
|
||||
"""Capture this iteration's feedback into ``progress``. Returns ``True`` if an entry was added."""
|
||||
if self.record_feedback is not None:
|
||||
entry = await _maybe_await(self.record_feedback(**loop_kwargs))
|
||||
else:
|
||||
entry = last_result.text.strip() if last_result is not None else None
|
||||
if entry:
|
||||
progress.append(entry)
|
||||
return True
|
||||
return False
|
||||
|
||||
async def _evaluate_stop(self, loop_kwargs: dict[str, Any], work_iterations: int) -> tuple[bool, str | None]:
|
||||
"""Decide whether the loop should stop, returning ``(stop, feedback)``.
|
||||
|
||||
``max_iterations`` is a safety cap that short-circuits before ``should_continue`` is
|
||||
evaluated (so an expensive predicate/judge is not called once the cap has fired). Any
|
||||
feedback returned by ``should_continue`` is propagated so the progress and next-message
|
||||
callables can reference it.
|
||||
"""
|
||||
if self.max_iterations is not None and work_iterations >= self.max_iterations:
|
||||
return True, None
|
||||
keep_going, feedback = await self._should_continue(loop_kwargs)
|
||||
return (not keep_going), feedback
|
||||
|
||||
async def _should_continue(self, loop_kwargs: dict[str, Any]) -> tuple[bool, str | None]:
|
||||
"""Evaluate the predicate, normalizing its result to ``(continue, feedback)``."""
|
||||
result = await _maybe_await(self.should_continue(**loop_kwargs))
|
||||
return (bool(result[0]), result[1]) if isinstance(result, tuple) else (bool(result), None) # type: ignore
|
||||
|
||||
@staticmethod
|
||||
def _message_to_update(message: Message) -> AgentResponseUpdate:
|
||||
"""Wrap an injected loop message as a streaming update so consumers see it inline."""
|
||||
return AgentResponseUpdate(
|
||||
contents=message.contents,
|
||||
role=message.role,
|
||||
author_name=message.author_name,
|
||||
message_id=message.message_id,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _aggregate_response(
|
||||
final: AgentResponse,
|
||||
messages: list[Message],
|
||||
usage: UsageDetails | None,
|
||||
) -> AgentResponse:
|
||||
"""Build a combined response carrying every iteration's messages and summed usage.
|
||||
|
||||
Metadata (``response_id``, structured ``value``, etc.) is taken from the final iteration; the
|
||||
structured value is passed through pre-parsed so it is not re-derived from the aggregated text.
|
||||
"""
|
||||
return AgentResponse(
|
||||
messages=messages,
|
||||
response_id=final.response_id,
|
||||
agent_id=final.agent_id,
|
||||
created_at=final.created_at,
|
||||
finish_reason=final.finish_reason, # pyright: ignore[reportArgumentType]
|
||||
usage_details=usage,
|
||||
value=final.value,
|
||||
additional_properties=dict(final.additional_properties) if final.additional_properties else None,
|
||||
raw_representation=final.raw_representation,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _render_progress(entries: list[str]) -> Message:
|
||||
"""Format progress-log entries into a single ``user`` message."""
|
||||
body = "\n".join(f"- {entry}" for entry in entries)
|
||||
return Message(role="user", contents=[f"Progress so far:\n{body}"])
|
||||
|
||||
async def _resolve_next_message(
|
||||
self,
|
||||
loop_kwargs: dict[str, Any],
|
||||
messages_used: list[Message],
|
||||
original_messages: list[Message],
|
||||
) -> list[Message]:
|
||||
# Compute the base next input. A ``next_message`` callable returning None requests a verbatim
|
||||
# reuse of the previous messages (no progress injection); in fresh-context mode that escape
|
||||
# hatch does not apply, so fall back to the default nudge instead.
|
||||
if self.next_message is None:
|
||||
next_msgs = normalize_messages(DEFAULT_NEXT_MESSAGE)
|
||||
else:
|
||||
next_input = await _maybe_await(self.next_message(**loop_kwargs))
|
||||
if next_input is None:
|
||||
if not self.fresh_context:
|
||||
return list(messages_used)
|
||||
next_msgs = normalize_messages(DEFAULT_NEXT_MESSAGE)
|
||||
else:
|
||||
next_msgs = normalize_messages(next_input)
|
||||
|
||||
progress: list[str] = loop_kwargs.get("progress") or []
|
||||
session = loop_kwargs.get("session")
|
||||
progress_msg: Message | None = None
|
||||
if self.inject_progress and progress:
|
||||
# With a session the earlier entries are already retained in the conversation, so only
|
||||
# the latest entry is injected to avoid duplication. Otherwise inject the full log.
|
||||
entries = progress if (session is None or self.fresh_context) else progress[-1:]
|
||||
progress_msg = self._render_progress(entries)
|
||||
|
||||
if self.fresh_context:
|
||||
result = list(original_messages)
|
||||
if progress_msg is not None:
|
||||
result.append(progress_msg)
|
||||
result.extend(next_msgs)
|
||||
return result
|
||||
|
||||
if progress_msg is not None:
|
||||
return [progress_msg, *next_msgs]
|
||||
return list(next_msgs)
|
||||
|
||||
|
||||
def _running_background_tasks(session: Any, agent: Any) -> list[Any]:
|
||||
"""Return the still-running ``BackgroundTaskInfo`` entries for the agent's provider.
|
||||
|
||||
Resolves the :class:`~agent_framework.BackgroundAgentsProvider` from the running agent
|
||||
(``agent.context_providers``) and reads its persisted task state. Returns an empty list when the
|
||||
session/agent/provider is unavailable or no task is currently running.
|
||||
"""
|
||||
from ._background_agents import BackgroundAgentsProvider, BackgroundTaskInfo, BackgroundTaskStatus
|
||||
|
||||
if session is None or agent is None:
|
||||
return []
|
||||
provider = _resolve_context_provider(agent, BackgroundAgentsProvider)
|
||||
if provider is None:
|
||||
return []
|
||||
state = session.state.get(provider.source_id)
|
||||
if not state:
|
||||
return []
|
||||
tasks = [BackgroundTaskInfo.from_dict(task) for task in state.get("tasks", [])]
|
||||
return [task for task in tasks if task.status == BackgroundTaskStatus.RUNNING]
|
||||
|
||||
|
||||
def background_tasks_running() -> ShouldContinueCallable:
|
||||
"""Build a ``should_continue`` predicate that loops while the agent's background tasks are busy.
|
||||
|
||||
This resolves the :class:`~agent_framework.BackgroundAgentsProvider` from the running agent
|
||||
(``agent.context_providers``).
|
||||
|
||||
The predicate inspects the provider's persisted task state and continues while any task is still
|
||||
marked as running. Pair it with ``max_iterations`` so the loop is guaranteed to stop even if a
|
||||
task's persisted status is never refreshed.
|
||||
|
||||
Returns:
|
||||
A predicate suitable for :class:`AgentLoopMiddleware`'s ``should_continue`` argument (and for
|
||||
``create_harness_agent``'s ``loop_should_continue``).
|
||||
"""
|
||||
|
||||
def _should_continue(*, session: Any = None, agent: Any = None, **kwargs: Any) -> bool:
|
||||
return bool(_running_background_tasks(session, agent))
|
||||
|
||||
return _should_continue
|
||||
|
||||
|
||||
def background_tasks_running_message(*, session: Any = None, agent: Any = None, **kwargs: Any) -> str | None:
|
||||
"""``next_message`` callable that reminds the agent which background tasks are still running.
|
||||
|
||||
Designed to pair with :func:`background_tasks_running` as a loop's ``next_message`` (e.g.
|
||||
``create_harness_agent``'s ``loop_next_message``): between iterations it resolves the
|
||||
:class:`~agent_framework.BackgroundAgentsProvider` from the agent, lists the still-running tasks,
|
||||
and instructs the agent to wait for them to finish (and retrieve their results) before finishing.
|
||||
|
||||
Returns ``None`` when the session/agent/provider is unavailable or no task is running. In that
|
||||
case the loop's default ``next_message`` handling applies. In normal looping a ``None`` here is
|
||||
rare, since "no running tasks" also makes :func:`background_tasks_running` stop the loop before
|
||||
the next message is consulted.
|
||||
"""
|
||||
running = _running_background_tasks(session, agent)
|
||||
if not running:
|
||||
return None
|
||||
task_lines = "\n".join(f"- #{task.id} ({task.agent_name}): {task.description}" for task in running)
|
||||
return (
|
||||
f"You still have {len(running)} background task(s) running that must finish before you can "
|
||||
f"complete the work:\n{task_lines}\n\n"
|
||||
"Wait for these tasks to complete, retrieve their results, and incorporate them. Only stop "
|
||||
"once every background task has finished."
|
||||
)
|
||||
|
||||
|
||||
def _resolve_context_provider(agent: Any, provider_type: type) -> Any:
|
||||
"""Return the first ``provider_type`` instance on ``agent.context_providers`` (or ``None``).
|
||||
|
||||
The harness exposes its built-in context providers (``TodoProvider``, ``AgentModeProvider``,
|
||||
...) on ``agent.context_providers``, so loop callbacks can reuse the same instances that
|
||||
:func:`~agent_framework.create_harness_agent` wired up instead of constructing their own.
|
||||
"""
|
||||
return next(
|
||||
(provider for provider in getattr(agent, "context_providers", []) if isinstance(provider, provider_type)),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def todos_remaining(*, looping_modes: Sequence[str] | None = None) -> ShouldContinueCallable:
|
||||
"""Build a ``should_continue`` predicate that loops while the Agent's ``TodoProvider`` has open items.
|
||||
|
||||
This resolves the :class:`~agent_framework.TodoProvider` from the running agent
|
||||
(``agent.context_providers``) rather than taking it as an argument, so it can be used directly
|
||||
with :func:`~agent_framework.create_harness_agent` (whose providers are built internally) as well
|
||||
as with any agent that registers a ``TodoProvider`` via ``context_providers``. It is the Python
|
||||
counterpart of the .NET ``TodoCompletionLoopEvaluator``.
|
||||
|
||||
Args:
|
||||
looping_modes: When provided, the loop only continues while the agent's current operating
|
||||
mode (read from its :class:`~agent_framework.AgentModeProvider`) is one of these modes;
|
||||
in any other mode the predicate returns ``False`` so the agent stays interactive. Mode
|
||||
matching is case-insensitive. When ``None`` (default), the loop applies in every mode. An
|
||||
empty sequence is rejected (there would be no mode in which the loop could ever run).
|
||||
Restricting looping to certain modes is useful when, for example, the agent has a planning
|
||||
and execution mode, and you only want to loop on the execution mode until all todos are
|
||||
complete. Looping until completion in planning is usually undesirable since the agent is
|
||||
still building the list of todos to complete.
|
||||
|
||||
Returns:
|
||||
A predicate suitable for :class:`AgentLoopMiddleware`'s ``should_continue`` argument (and for
|
||||
``create_harness_agent``'s ``loop_should_continue``).
|
||||
|
||||
Raises:
|
||||
ValueError: ``looping_modes`` is an empty sequence.
|
||||
"""
|
||||
if looping_modes is not None:
|
||||
allowed_modes: set[str] | None = {mode.strip().lower() for mode in looping_modes}
|
||||
if not allowed_modes:
|
||||
raise ValueError("looping_modes must be None or a non-empty sequence of mode names.")
|
||||
else:
|
||||
allowed_modes = None
|
||||
|
||||
async def _should_continue(*, session: Any = None, agent: Any = None, **kwargs: Any) -> bool:
|
||||
from ._mode import AgentModeProvider, get_agent_mode
|
||||
from ._todo import TodoProvider
|
||||
|
||||
if session is None or agent is None:
|
||||
return False
|
||||
|
||||
if allowed_modes is not None:
|
||||
mode_provider = _resolve_context_provider(agent, AgentModeProvider)
|
||||
if mode_provider is not None:
|
||||
current_mode = get_agent_mode(
|
||||
session,
|
||||
source_id=mode_provider.source_id,
|
||||
default_mode=mode_provider.default_mode,
|
||||
available_modes=mode_provider.available_modes,
|
||||
)
|
||||
else:
|
||||
current_mode = get_agent_mode(session)
|
||||
if current_mode.strip().lower() not in allowed_modes:
|
||||
return False
|
||||
|
||||
todo_provider = _resolve_context_provider(agent, TodoProvider)
|
||||
if todo_provider is None:
|
||||
return False
|
||||
items = await todo_provider.store.load_items(session, source_id=todo_provider.source_id)
|
||||
return any(not item.is_complete for item in items)
|
||||
|
||||
return _should_continue
|
||||
|
||||
|
||||
async def todos_remaining_message(*, session: Any = None, agent: Any = None, **kwargs: Any) -> str | None:
|
||||
"""``next_message`` callable that reminds the agent which todos are still open.
|
||||
|
||||
Designed to pair with :func:`todos_remaining` as a loop's ``next_message`` (e.g.
|
||||
``create_harness_agent``'s ``loop_next_message``): between iterations it resolves the harness
|
||||
:class:`~agent_framework.TodoProvider` from the agent, lists the still-open todo items, and
|
||||
instructs the agent to complete them all before finishing.
|
||||
|
||||
Returns ``None`` when the session/agent/provider is unavailable or no todos are open. In that
|
||||
case the loop's default ``next_message`` handling applies: with ``fresh_context=False`` (the
|
||||
default, used by ``create_harness_agent``) it reuses the previous iteration's messages verbatim
|
||||
(skipping progress injection); only with ``fresh_context=True`` does it fall back to
|
||||
``DEFAULT_NEXT_MESSAGE``. In normal looping a ``None`` here is rare, since "no open todos" also
|
||||
makes :func:`todos_remaining` stop the loop before the next message is consulted.
|
||||
"""
|
||||
from ._todo import TodoProvider
|
||||
|
||||
if session is None or agent is None:
|
||||
return None
|
||||
todo_provider = _resolve_context_provider(agent, TodoProvider)
|
||||
if todo_provider is None:
|
||||
return None
|
||||
items = await todo_provider.store.load_items(session, source_id=todo_provider.source_id)
|
||||
open_items = [item for item in items if not item.is_complete]
|
||||
if not open_items:
|
||||
return None
|
||||
todo_lines = "\n".join(f"- {item.title}" for item in open_items)
|
||||
return (
|
||||
f"You still have {len(open_items)} open todo item(s) that must be addressed before you can "
|
||||
f"finish:\n{todo_lines}\n\n"
|
||||
"Continue working through them now. Mark each todo complete as you finish it, and only stop "
|
||||
"once every todo item is complete."
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,325 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, cast
|
||||
|
||||
from .._feature_stage import ExperimentalFeature, experimental
|
||||
from .._sessions import AgentSession, ContextProvider, SessionContext
|
||||
from .._tools import tool
|
||||
from .._types import Message
|
||||
|
||||
DEFAULT_MODE_SOURCE_ID = "agent_mode"
|
||||
DEFAULT_MODE_INSTRUCTIONS = (
|
||||
"## Agent Mode\n\n"
|
||||
"- You can operate in different modes. Depending on the mode you are in, "
|
||||
"you will be required to follow different processes.\n\n"
|
||||
"Use the mode_get tool to check your current operating mode.\n"
|
||||
"Use the mode_set tool to switch between modes as your work progresses. "
|
||||
"Only use mode_set if the user explicitly instructs/allows you to change modes.\n\n"
|
||||
"You are currently operating in the {current_mode} mode.\n\n"
|
||||
"### Mandatory Mode based Workflow\n\n"
|
||||
"For every new substantive user request, including short factual questions, "
|
||||
"your behavior is determined by the mode you are in.\n\n"
|
||||
"{available_modes}\n"
|
||||
)
|
||||
DEFAULT_MODE_CHANGE_NOTIFICATION = (
|
||||
'[Mode changed: The operating mode has been switched from "{previous_mode}" to "{current_mode}". '
|
||||
'You must now adjust your behavior to match the "{current_mode}" mode.]'
|
||||
)
|
||||
DEFAULT_MODE_DESCRIPTIONS: dict[str, str] = {
|
||||
"plan": (
|
||||
"Use this mode when analyzing requirements, breaking down tasks, and creating plans. "
|
||||
"This is the interactive mode — ask clarifying questions, discuss options, and get user approval before "
|
||||
"proceeding.\n\n"
|
||||
"Process to follow when in plan mode:\n"
|
||||
"1. Analyze the request with the purpose of building a research plan.\n"
|
||||
"2. Create a list of todo items.\n"
|
||||
"3. If needed, use the provided tools to do some exploratory checks to help build a plan and determine "
|
||||
"what clarifying questions you may need from the user.\n"
|
||||
"4. Ask for clarifications from the user where needed.\n"
|
||||
" 1. Ask each clarification one by one.\n"
|
||||
" 2. When asking for clarification and you have specific options in mind, present them to the user, "
|
||||
"so they can choose the option instead of having to retype the entire response.\n"
|
||||
" 3. Do not proceed until you have received all the needed clarifications.\n"
|
||||
" 4. Do short exploratory research if it helps with being able to ask sensible clarifications from "
|
||||
"the user.\n"
|
||||
"5. Write the plan to a memory file, so that it is retained even if compaction happens. "
|
||||
"Make sure to update the plan file if the user requests changes.\n"
|
||||
"6. Present the plan to the user and ask for approval to switch to execute mode and process the plan.\n"
|
||||
"7. When approval is granted, always switch to execute mode (using the `mode_set` tool), "
|
||||
"and follow the steps for *Execute mode*."
|
||||
),
|
||||
"execute": (
|
||||
"Determine the type of ask:\n"
|
||||
"1. Simple question that doesn't require any further work to answer.\n"
|
||||
"2. Any other work, including complex user request that requires a multi-step process to satisfy.\n\n"
|
||||
"If 1. just answer the question directly.\n"
|
||||
"If 2. Work autonomously using your best judgment — do not ask the user questions or wait for feedback "
|
||||
"and follow the following process:\n"
|
||||
"1. If you don't have a plan or tasks yet, analyze the user request and create tasks and a plan. "
|
||||
"(**Skip this step if you came from plan mode**)\n"
|
||||
"2. Work autonomously — use your best judgment to make decisions and keep progressing without asking "
|
||||
"the user questions. The goal is to have a complete, useful result ready when the user returns.\n"
|
||||
"3. If you encounter ambiguity or an unexpected situation during execution, choose the most reasonable "
|
||||
"option, note your choice, and keep going.\n"
|
||||
"4. Mark tasks as completed as you finish them.\n"
|
||||
"5. Continue working, thinking and calling tools until you have the research result for the user."
|
||||
),
|
||||
}
|
||||
|
||||
_PREVIOUS_MODE_STATE_KEY = "previous_mode_for_notification"
|
||||
|
||||
|
||||
def _get_mode_state(session: AgentSession, *, source_id: str) -> dict[str, Any]:
|
||||
"""Return the mutable session state used by the mode provider."""
|
||||
provider_state = session.state.get(source_id)
|
||||
if isinstance(provider_state, dict):
|
||||
return cast(dict[str, Any], provider_state)
|
||||
if provider_state is not None:
|
||||
raise TypeError(
|
||||
f"Session state for source_id {source_id!r} must be a dict, got {type(provider_state).__name__}."
|
||||
)
|
||||
state: dict[str, Any] = {}
|
||||
session.state[source_id] = state
|
||||
return state
|
||||
|
||||
|
||||
def _normalize_available_modes(available_modes: Sequence[str]) -> dict[str, str]:
|
||||
"""Return normalized mode names mapped to display names."""
|
||||
normalized_modes: dict[str, str] = {}
|
||||
for mode in available_modes:
|
||||
display_mode = mode.strip()
|
||||
normalized_mode = display_mode.lower()
|
||||
if normalized_mode in normalized_modes:
|
||||
raise ValueError(f"Duplicate mode configured: {mode}.")
|
||||
normalized_modes[normalized_mode] = display_mode
|
||||
return normalized_modes
|
||||
|
||||
|
||||
def _normalize_mode(mode: str, *, available_modes: Mapping[str, str]) -> str:
|
||||
"""Validate and normalize a mode string."""
|
||||
normalized = mode.strip().lower()
|
||||
if normalized not in available_modes:
|
||||
supported_modes = ", ".join(repr(item) for item in available_modes.values())
|
||||
raise ValueError(f"Invalid mode: {mode}. Supported modes are {supported_modes}.")
|
||||
return normalized
|
||||
|
||||
|
||||
def _resolve_default_mode(default_mode: str | None, *, available_modes: Mapping[str, str]) -> str:
|
||||
"""Resolve the default mode, falling back to the first configured mode when omitted."""
|
||||
if default_mode is None:
|
||||
return next(iter(available_modes))
|
||||
return _normalize_mode(default_mode, available_modes=available_modes)
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
def get_agent_mode(
|
||||
session: AgentSession,
|
||||
*,
|
||||
source_id: str = DEFAULT_MODE_SOURCE_ID,
|
||||
default_mode: str | None = None,
|
||||
available_modes: Sequence[str] | None = None,
|
||||
) -> str:
|
||||
"""Get the current operating mode from session state.
|
||||
|
||||
Args:
|
||||
session: The agent session to read the mode from.
|
||||
|
||||
Keyword Args:
|
||||
source_id: Unique source ID for the provider state.
|
||||
default_mode: Initial mode used when no mode is stored yet. When omitted, the first entry of
|
||||
``available_modes`` is used.
|
||||
available_modes: Supported modes to validate against. Defaults to the built-in modes.
|
||||
|
||||
Returns:
|
||||
The current mode string.
|
||||
"""
|
||||
normalized_modes = _normalize_available_modes(tuple(available_modes or DEFAULT_MODE_DESCRIPTIONS))
|
||||
normalized_default_mode = _resolve_default_mode(default_mode, available_modes=normalized_modes)
|
||||
provider_state = _get_mode_state(session, source_id=source_id)
|
||||
current_mode = provider_state.get("current_mode")
|
||||
if isinstance(current_mode, str):
|
||||
try:
|
||||
return _normalize_mode(current_mode, available_modes=normalized_modes)
|
||||
except ValueError:
|
||||
# Stored mode is no longer in the configured set (e.g. available_modes was reconfigured).
|
||||
# Fall through and reset to the default mode.
|
||||
pass
|
||||
provider_state["current_mode"] = normalized_default_mode
|
||||
return normalized_default_mode
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
def set_agent_mode(
|
||||
session: AgentSession,
|
||||
mode: str,
|
||||
*,
|
||||
source_id: str = DEFAULT_MODE_SOURCE_ID,
|
||||
available_modes: Sequence[str] | None = None,
|
||||
) -> str:
|
||||
"""Set the current operating mode in session state.
|
||||
|
||||
External callers (e.g., a slash-command handler) should use this helper rather than mutating
|
||||
session state directly. When the mode actually changes, the prior mode is recorded so that the
|
||||
next :meth:`AgentModeProvider.before_run` invocation injects a user message announcing the
|
||||
switch — system instructions alone are insufficient to redirect a model that has already seen
|
||||
its own ``set_mode`` tool call earlier in the chat history.
|
||||
|
||||
Args:
|
||||
session: The agent session to update the mode in.
|
||||
mode: The new mode to set.
|
||||
|
||||
Keyword Args:
|
||||
source_id: Unique source ID for the provider state.
|
||||
available_modes: Supported modes to validate against. Defaults to the built-in modes.
|
||||
|
||||
Returns:
|
||||
The normalized mode string that was stored.
|
||||
|
||||
Raises:
|
||||
ValueError: The requested mode is not configured.
|
||||
"""
|
||||
normalized_modes = _normalize_available_modes(tuple(available_modes or DEFAULT_MODE_DESCRIPTIONS))
|
||||
normalized_mode = _normalize_mode(mode, available_modes=normalized_modes)
|
||||
provider_state = _get_mode_state(session, source_id=source_id)
|
||||
previous_mode = provider_state.get("current_mode")
|
||||
provider_state["current_mode"] = normalized_mode
|
||||
# When the mode is changed externally (i.e. not via the agent's own ``set_mode`` tool), record the
|
||||
# prior mode so the next ``before_run`` can inject a user message announcing the switch. Without
|
||||
# that injection, the model often anchors on the earlier ``set_mode`` tool call in the chat history
|
||||
# and keeps behaving as if it were still in that mode — system instructions alone are insufficient.
|
||||
if isinstance(previous_mode, str) and previous_mode != normalized_mode:
|
||||
provider_state[_PREVIOUS_MODE_STATE_KEY] = previous_mode
|
||||
return normalized_mode
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class AgentModeProvider(ContextProvider):
|
||||
"""Track the agent's operating mode in session state and provide mode tools.
|
||||
|
||||
The ``AgentModeProvider`` enables agents to operate in distinct modes during long-running complex tasks.
|
||||
The current mode is persisted in the ``AgentSession`` state and is included in the instructions provided to the
|
||||
agent on each invocation.
|
||||
|
||||
The set of available modes is configurable with ``mode_descriptions``. By default, two modes are provided:
|
||||
``"plan"`` (interactive planning) and ``"execute"`` (autonomous execution).
|
||||
|
||||
This provider exposes the following tools to the agent:
|
||||
- ``mode_set``: Switch the agent's operating mode.
|
||||
- ``mode_get``: Retrieve the agent's current operating mode.
|
||||
|
||||
Public helper functions ``get_agent_mode`` and ``set_agent_mode`` allow external code to programmatically read
|
||||
and change the mode.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str = DEFAULT_MODE_SOURCE_ID,
|
||||
*,
|
||||
default_mode: str | None = None,
|
||||
mode_descriptions: Mapping[str, str] | None = None,
|
||||
instructions: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize a new agent mode provider.
|
||||
|
||||
Args:
|
||||
source_id: Unique source ID for the provider.
|
||||
|
||||
Keyword Args:
|
||||
default_mode: Initial mode used when no mode is stored yet. When omitted, the first entry of
|
||||
``mode_descriptions`` is used.
|
||||
mode_descriptions: Mapping of supported modes to descriptions of when and how to use each mode.
|
||||
instructions: Custom instructions for using the mode tools. The instructions can contain an
|
||||
``{available_modes}`` placeholder for the configured list of modes and a ``{current_mode}`` placeholder
|
||||
for the currently active mode. When omitted, the provider uses a default set of instructions.
|
||||
|
||||
Raises:
|
||||
ValueError: No modes are configured, or the default mode is not configured.
|
||||
"""
|
||||
super().__init__(source_id)
|
||||
mode_descriptions = dict(DEFAULT_MODE_DESCRIPTIONS if mode_descriptions is None else mode_descriptions)
|
||||
self._mode_display_names = _normalize_available_modes(tuple(mode_descriptions))
|
||||
if not self._mode_display_names:
|
||||
raise ValueError("mode_descriptions must contain at least one mode.")
|
||||
self.mode_descriptions = {mode.strip().lower(): description for mode, description in mode_descriptions.items()}
|
||||
self.available_modes = tuple(self._mode_display_names)
|
||||
self.default_mode = _resolve_default_mode(default_mode, available_modes=self._mode_display_names)
|
||||
self.instructions = instructions
|
||||
|
||||
def _build_instructions(self, current_mode: str) -> str:
|
||||
"""Build the mode guidance injected for the current session."""
|
||||
mode_lines = "".join(
|
||||
f"#### {self._mode_display_names[mode]}\n\n{description}\n\n"
|
||||
for mode, description in self.mode_descriptions.items()
|
||||
)
|
||||
instructions = self.instructions or DEFAULT_MODE_INSTRUCTIONS
|
||||
return instructions.replace("{available_modes}", mode_lines).replace("{current_mode}", current_mode)
|
||||
|
||||
async def before_run(
|
||||
self,
|
||||
*,
|
||||
agent: Any,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Inject mode tools and instructions before the model runs.
|
||||
|
||||
Args:
|
||||
agent: The agent being invoked.
|
||||
session: The agent session whose state stores the current mode.
|
||||
context: The session context to receive instructions and tools.
|
||||
state: Per-provider invocation state.
|
||||
"""
|
||||
del agent, state
|
||||
current_mode = get_agent_mode(
|
||||
session,
|
||||
source_id=self.source_id,
|
||||
default_mode=self.default_mode,
|
||||
available_modes=self.available_modes,
|
||||
)
|
||||
# Pop the external-mode-change marker (set by ``set_agent_mode``) before injecting tools so
|
||||
# the agent only sees the notification once.
|
||||
provider_state = _get_mode_state(session, source_id=self.source_id)
|
||||
previous_mode = provider_state.pop(_PREVIOUS_MODE_STATE_KEY, None)
|
||||
|
||||
@tool(name="mode_set", approval_mode="never_require")
|
||||
def mode_set(mode: str) -> str:
|
||||
"""Switch the agent's operating mode."""
|
||||
# The agent invoked the tool itself, so it knows the mode just changed — bypass
|
||||
# ``set_agent_mode`` to avoid triggering a notification message on the next turn.
|
||||
normalized_mode = _normalize_mode(mode, available_modes=self._mode_display_names)
|
||||
tool_state = _get_mode_state(session, source_id=self.source_id)
|
||||
tool_state["current_mode"] = normalized_mode
|
||||
return json.dumps({"mode": normalized_mode, "message": f"Mode changed to '{normalized_mode}'."})
|
||||
|
||||
@tool(name="mode_get", approval_mode="never_require")
|
||||
def mode_get() -> str:
|
||||
"""Get the agent's current operating mode."""
|
||||
current_mode_value = get_agent_mode(
|
||||
session,
|
||||
source_id=self.source_id,
|
||||
default_mode=self.default_mode,
|
||||
available_modes=self.available_modes,
|
||||
)
|
||||
return json.dumps({"mode": current_mode_value})
|
||||
|
||||
context.extend_instructions(
|
||||
self.source_id,
|
||||
[self._build_instructions(current_mode)],
|
||||
)
|
||||
context.extend_tools(self.source_id, [mode_set, mode_get])
|
||||
if isinstance(previous_mode, str) and previous_mode != current_mode:
|
||||
# Inject a user-role message announcing the external mode change. System instructions
|
||||
# always render first in the chat history, so the agent can otherwise stay anchored to
|
||||
# the most recent ``mode_set`` tool call rather than the new mode.
|
||||
previous_display = self._mode_display_names.get(previous_mode, previous_mode)
|
||||
current_display = self._mode_display_names.get(current_mode, current_mode)
|
||||
notification = DEFAULT_MODE_CHANGE_NOTIFICATION.format(
|
||||
previous_mode=previous_display,
|
||||
current_mode=current_display,
|
||||
)
|
||||
context.extend_messages(self, [Message(role="user", contents=[notification])])
|
||||
@@ -0,0 +1,619 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import weakref
|
||||
from abc import ABC, abstractmethod
|
||||
from base64 import urlsafe_b64encode
|
||||
from collections.abc import Mapping, MutableMapping
|
||||
from pathlib import Path
|
||||
from typing import Any, ClassVar, cast
|
||||
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
from .._feature_stage import ExperimentalFeature, experimental
|
||||
from .._serialization import SerializationMixin
|
||||
from .._sessions import AgentSession, ContextProvider, SessionContext
|
||||
from .._tools import tool
|
||||
from .._types import Message
|
||||
|
||||
DEFAULT_TODO_SOURCE_ID = "todo"
|
||||
DEFAULT_TODO_INSTRUCTIONS = (
|
||||
"## Todo Items\n\n"
|
||||
"You have access to a todo list for tracking work items.\n"
|
||||
"When a user asks you to perform a task, follow these steps to manage your work:\n"
|
||||
"1. Determine whether the ask requires multiple steps to complete (complex) or can be completed "
|
||||
"using a single step (simple).\n"
|
||||
"2. If complex, turn the task into manageable todo items and add them to the list.\n"
|
||||
"3. If simple, don't add a todo item, but rather just complete the task directly.\n\n"
|
||||
"### General TODO Guidelines\n"
|
||||
"Ask questions from the user where clarification is needed to create effective todos.\n"
|
||||
"If the user provides feedback on your plan, adjust your todos accordingly by adding new items "
|
||||
"or removing irrelevant ones.\n"
|
||||
"During execution, use the todo list to keep track of what needs to be done, "
|
||||
"mark items as complete when finished, and remove any items that are no longer needed.\n"
|
||||
"When a user changes the topic, changes their mind or switches to a new request, ensure that you update "
|
||||
"the todo list accordingly by removing irrelevant/old items, clearing the list, or adding new ones as needed.\n\n"
|
||||
"Use these tools to manage your tasks:\n"
|
||||
"- Use todos_add to break down complex work into trackable items (supports adding one or many at once).\n"
|
||||
"- Use todos_complete to mark items as done when finished (supports one or many at once). "
|
||||
"Include a reason describing how the items were completed.\n"
|
||||
"- Use todos_get_remaining to check what work is still pending.\n"
|
||||
"- Use todos_get_all to review the full list including completed items.\n"
|
||||
"- Use todos_remove to remove items that are no longer needed (supports one or many at once)."
|
||||
)
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class TodoItem(SerializationMixin):
|
||||
"""Represent one todo item tracked for the current session."""
|
||||
|
||||
id: int
|
||||
title: str
|
||||
description: str | None
|
||||
is_complete: bool
|
||||
|
||||
def __init__(self, id: int, title: str, description: str | None = None, is_complete: bool = False) -> None:
|
||||
"""Initialize one todo item."""
|
||||
self.id = id
|
||||
self.title = title
|
||||
self.description = description
|
||||
self.is_complete = is_complete
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
|
||||
"""Serialize the todo item for persistence."""
|
||||
del exclude
|
||||
payload = {
|
||||
"id": self.id,
|
||||
"title": self.title,
|
||||
"description": self.description,
|
||||
"is_complete": self.is_complete,
|
||||
}
|
||||
return {key: value for key, value in payload.items() if value is not None or not exclude_none}
|
||||
|
||||
@classmethod
|
||||
def from_dict(
|
||||
cls, raw_item: MutableMapping[str, Any], /, *, dependencies: MutableMapping[str, Any] | None = None
|
||||
) -> TodoItem:
|
||||
"""Parse one todo item loaded from storage."""
|
||||
del dependencies
|
||||
item_id = raw_item.get("id")
|
||||
title = raw_item.get("title")
|
||||
description = raw_item.get("description")
|
||||
is_complete = raw_item.get("is_complete", False)
|
||||
if not isinstance(item_id, int):
|
||||
raise ValueError("Todo item id must be an integer.")
|
||||
if not isinstance(title, str) or not title.strip():
|
||||
raise ValueError("Todo item title must be a non-empty string.")
|
||||
if description is not None and not isinstance(description, str):
|
||||
raise ValueError("Todo item description must be a string or null.")
|
||||
if not isinstance(is_complete, bool):
|
||||
raise ValueError("Todo item is_complete must be a boolean.")
|
||||
return cls(id=item_id, title=title, description=description, is_complete=is_complete)
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
"""Return whether two todo items have the same values."""
|
||||
return isinstance(other, TodoItem) and self.to_dict() == other.to_dict()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Return a helpful debug representation."""
|
||||
return (
|
||||
"TodoItem("
|
||||
f"id={self.id!r}, title={self.title!r}, description={self.description!r}, is_complete={self.is_complete!r})"
|
||||
)
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class TodoInput(SerializationMixin):
|
||||
"""Describe one todo item to create."""
|
||||
|
||||
title: str
|
||||
description: str | None
|
||||
|
||||
def __init__(self, title: str, description: str | None = None) -> None:
|
||||
"""Initialize one todo input."""
|
||||
normalized_title = title.strip()
|
||||
if not normalized_title:
|
||||
raise ValueError("Todo input title must be a non-empty string.")
|
||||
if description is not None and not isinstance(description, str):
|
||||
raise ValueError("Todo input description must be a string or null.")
|
||||
self.title = normalized_title
|
||||
self.description = description
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
|
||||
"""Serialize the todo input."""
|
||||
del exclude
|
||||
payload = {"title": self.title, "description": self.description}
|
||||
return {key: value for key, value in payload.items() if value is not None or not exclude_none}
|
||||
|
||||
@classmethod
|
||||
def from_dict(
|
||||
cls, raw_todo: MutableMapping[str, Any], /, *, dependencies: MutableMapping[str, Any] | None = None
|
||||
) -> TodoInput:
|
||||
"""Parse one todo input loaded from tool arguments."""
|
||||
del dependencies
|
||||
title = raw_todo.get("title")
|
||||
description = raw_todo.get("description")
|
||||
if not isinstance(title, str):
|
||||
raise ValueError("Todo input title must be a string.")
|
||||
return cls(title=title, description=description)
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class TodoCompleteInput(SerializationMixin):
|
||||
"""Describe one todo item to mark as complete."""
|
||||
|
||||
id: int
|
||||
reason: str
|
||||
|
||||
def __init__(self, id: int, reason: str) -> None:
|
||||
"""Initialize one todo complete input."""
|
||||
if not isinstance(id, int):
|
||||
raise ValueError("Todo complete input id must be an integer.")
|
||||
if not isinstance(reason, str) or not reason.strip():
|
||||
raise ValueError("Todo complete input reason must be a non-empty string.")
|
||||
self.id = id
|
||||
self.reason = reason.strip()
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
|
||||
"""Serialize the todo complete input."""
|
||||
del exclude, exclude_none
|
||||
return {"id": self.id, "reason": self.reason}
|
||||
|
||||
@classmethod
|
||||
def from_dict(
|
||||
cls, raw_item: MutableMapping[str, Any], /, *, dependencies: MutableMapping[str, Any] | None = None
|
||||
) -> TodoCompleteInput:
|
||||
"""Parse one todo complete input from tool arguments."""
|
||||
del dependencies
|
||||
item_id = raw_item.get("id")
|
||||
reason = raw_item.get("reason")
|
||||
if not isinstance(item_id, int):
|
||||
raise ValueError("Todo complete input id must be an integer.")
|
||||
if not isinstance(reason, str):
|
||||
raise ValueError("Todo complete input reason must be a string.")
|
||||
return cls(id=item_id, reason=reason)
|
||||
|
||||
|
||||
class _TodoAddItemSchema(TypedDict):
|
||||
"""Schema for a single todo item in the todos_add tool."""
|
||||
|
||||
title: str
|
||||
description: NotRequired[str]
|
||||
|
||||
|
||||
class _TodoCompleteItemSchema(TypedDict):
|
||||
"""Schema for a single item in the todos_complete tool."""
|
||||
|
||||
id: int
|
||||
reason: str
|
||||
|
||||
|
||||
def _parse_todo_items(items_payload: list[Any], *, source_description: str) -> list[TodoItem]:
|
||||
"""Parse persisted todo item payloads with clear corruption errors."""
|
||||
items: list[TodoItem] = []
|
||||
for index, item in enumerate(items_payload):
|
||||
if not isinstance(item, Mapping):
|
||||
raise ValueError(
|
||||
f"Todo item at index {index} in {source_description} must be a mapping; got {type(item).__name__}."
|
||||
)
|
||||
items.append(TodoItem.from_dict(dict(cast(Mapping[str, Any], item))))
|
||||
return items
|
||||
|
||||
|
||||
def _coerce_todo_input(todo: TodoInput | dict[str, Any] | Any) -> TodoInput:
|
||||
"""Normalize tool-provided todo input into a TodoInput model."""
|
||||
if isinstance(todo, TodoInput):
|
||||
return todo
|
||||
if isinstance(todo, MutableMapping):
|
||||
return TodoInput.from_dict(cast(MutableMapping[str, Any], todo))
|
||||
raise ValueError("Todo input must be a TodoInput instance or JSON object.")
|
||||
|
||||
|
||||
def _coerce_todo_complete_input(item: TodoCompleteInput | dict[str, Any] | Any) -> TodoCompleteInput:
|
||||
"""Normalize tool-provided complete input into a TodoCompleteInput model."""
|
||||
if isinstance(item, TodoCompleteInput):
|
||||
return item
|
||||
if isinstance(item, MutableMapping):
|
||||
return TodoCompleteInput.from_dict(cast(MutableMapping[str, Any], item))
|
||||
raise ValueError("Todo complete input must be a TodoCompleteInput instance or JSON object.")
|
||||
|
||||
|
||||
def _safe_next_id(items: list[TodoItem], next_id: int) -> int:
|
||||
"""Clamp ``next_id`` so it cannot collide with any persisted item id."""
|
||||
return max(next_id, max((item.id for item in items), default=0) + 1)
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class TodoStore(ABC):
|
||||
"""Abstract backing store for session todo items."""
|
||||
|
||||
@abstractmethod
|
||||
async def load_state(self, session: AgentSession, *, source_id: str) -> tuple[list[TodoItem], int]:
|
||||
"""Load persisted todo items and the next available ID."""
|
||||
|
||||
@abstractmethod
|
||||
async def save_state(self, session: AgentSession, items: list[TodoItem], *, next_id: int, source_id: str) -> None:
|
||||
"""Persist todo items and the next available ID."""
|
||||
|
||||
async def load_items(self, session: AgentSession, *, source_id: str) -> list[TodoItem]:
|
||||
"""Load todo items for one session."""
|
||||
items, _ = await self.load_state(session, source_id=source_id)
|
||||
return items
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class TodoSessionStore(TodoStore):
|
||||
"""Store todo state inside ``AgentSession.state``."""
|
||||
|
||||
async def load_state(self, session: AgentSession, *, source_id: str) -> tuple[list[TodoItem], int]:
|
||||
"""Load todo state from session state."""
|
||||
provider_state_value = session.state.get(source_id)
|
||||
if provider_state_value is None:
|
||||
provider_state: dict[str, Any] = {}
|
||||
session.state[source_id] = provider_state
|
||||
elif isinstance(provider_state_value, dict):
|
||||
provider_state = cast(dict[str, Any], provider_state_value)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Session state for source_id {source_id!r} must be a dict; got {type(provider_state_value).__name__}."
|
||||
)
|
||||
|
||||
raw_items = provider_state.get("items", [])
|
||||
if not isinstance(raw_items, list):
|
||||
raise ValueError(
|
||||
f"Session state for source_id {source_id!r} has a non-list 'items' field; "
|
||||
f"got {type(raw_items).__name__}."
|
||||
)
|
||||
raw_next_id = provider_state.get("next_id", 1)
|
||||
if not isinstance(raw_next_id, int):
|
||||
raise ValueError(
|
||||
f"Session state for source_id {source_id!r} has a non-integer 'next_id' field; "
|
||||
f"got {type(raw_next_id).__name__}."
|
||||
)
|
||||
items_payload: list[Any] = cast(Any, raw_items)
|
||||
items = _parse_todo_items(items_payload, source_description="session todo state")
|
||||
return items, _safe_next_id(items, raw_next_id)
|
||||
|
||||
async def save_state(self, session: AgentSession, items: list[TodoItem], *, next_id: int, source_id: str) -> None:
|
||||
"""Persist todo state back into session state."""
|
||||
provider_state_value = session.state.get(source_id)
|
||||
provider_state = cast(dict[str, Any], provider_state_value) if isinstance(provider_state_value, dict) else {}
|
||||
if not isinstance(provider_state_value, dict):
|
||||
session.state[source_id] = provider_state
|
||||
provider_state["items"] = [item.to_dict(exclude_none=False) for item in items]
|
||||
provider_state["next_id"] = _safe_next_id(items, next_id)
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class TodoFileStore(TodoStore):
|
||||
"""Store todo state in one JSON file per session and source ID."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_path: str | Path,
|
||||
*,
|
||||
kind: str = "todos",
|
||||
owner_prefix: str = "",
|
||||
owner_state_key: str | None = None,
|
||||
state_filename: str = "todos.json",
|
||||
) -> None:
|
||||
"""Initialize the file-backed todo store.
|
||||
|
||||
Args:
|
||||
base_path: Root storage directory.
|
||||
|
||||
Keyword Args:
|
||||
kind: Storage bucket name under each owner directory.
|
||||
owner_prefix: Optional prefix applied to the resolved owner ID.
|
||||
owner_state_key: Session-state key holding the logical owner ID.
|
||||
state_filename: File name used for the persisted todo state.
|
||||
"""
|
||||
self.base_path = Path(base_path)
|
||||
self.kind = kind
|
||||
self.owner_prefix = owner_prefix
|
||||
self.owner_state_key = owner_state_key
|
||||
self.state_filename = state_filename
|
||||
self._base_root = self.base_path.resolve()
|
||||
|
||||
_ENCODED_SEGMENT_PREFIX: ClassVar[str] = "~todo-"
|
||||
_WINDOWS_RESERVED_FILE_STEMS: ClassVar[frozenset[str]] = frozenset({
|
||||
"CON",
|
||||
"PRN",
|
||||
"AUX",
|
||||
"NUL",
|
||||
"COM1",
|
||||
"COM2",
|
||||
"COM3",
|
||||
"COM4",
|
||||
"COM5",
|
||||
"COM6",
|
||||
"COM7",
|
||||
"COM8",
|
||||
"COM9",
|
||||
"LPT1",
|
||||
"LPT2",
|
||||
"LPT3",
|
||||
"LPT4",
|
||||
"LPT5",
|
||||
"LPT6",
|
||||
"LPT7",
|
||||
"LPT8",
|
||||
"LPT9",
|
||||
})
|
||||
|
||||
def _get_state_path(self, session: AgentSession, *, source_id: str) -> Path:
|
||||
"""Return the JSON file path for one session and source ID."""
|
||||
session_directory = self.base_path
|
||||
if self.owner_state_key is not None:
|
||||
owner_value = session.state.get(self.owner_state_key)
|
||||
if owner_value is None:
|
||||
raise RuntimeError(
|
||||
f"TodoFileStore requires session.state[{self.owner_state_key!r}] to be set for file-backed storage."
|
||||
)
|
||||
owner_segment = self._path_segment(owner_value, label="owner")
|
||||
session_directory = session_directory / f"{self.owner_prefix}{owner_segment}" / self.kind
|
||||
session_directory = session_directory / self._path_segment(
|
||||
session.session_id, label="session_id", reject_path_separators=True
|
||||
)
|
||||
state_path = (session_directory / self._state_filename(source_id)).resolve()
|
||||
if not state_path.is_relative_to(self._base_root):
|
||||
raise ValueError(f"Todo file path escaped base directory for session_id {session.session_id!r}.")
|
||||
return state_path
|
||||
|
||||
@classmethod
|
||||
def _path_segment(cls, value: object, *, label: str, reject_path_separators: bool = False) -> str:
|
||||
"""Return a filesystem-safe path segment for user-controlled state values."""
|
||||
raw_value = str(value)
|
||||
if reject_path_separators and ("/" in raw_value or "\\" in raw_value):
|
||||
raise ValueError(f"TodoFileStore {label} must not contain path separators: {raw_value!r}")
|
||||
if cls._is_literal_path_segment_safe(raw_value):
|
||||
return raw_value
|
||||
encoded_value = urlsafe_b64encode(raw_value.encode("utf-8")).decode("ascii").rstrip("=")
|
||||
return f"{cls._ENCODED_SEGMENT_PREFIX}{encoded_value or label}"
|
||||
|
||||
@classmethod
|
||||
def _is_literal_path_segment_safe(cls, value: str) -> bool:
|
||||
"""Return whether a value can be used directly as one path segment."""
|
||||
if (
|
||||
not value
|
||||
or value.startswith(".")
|
||||
or value.endswith((" ", "."))
|
||||
or value.upper() in cls._WINDOWS_RESERVED_FILE_STEMS
|
||||
):
|
||||
return False
|
||||
if any(ord(character) < 32 for character in value):
|
||||
return False
|
||||
return all(character.isalnum() or character in "._-" for character in value)
|
||||
|
||||
def _state_filename(self, source_id: str) -> str:
|
||||
"""Return a source-specific JSON state filename."""
|
||||
state_path = Path(self.state_filename)
|
||||
source_segment = self._path_segment(source_id, label="source_id")
|
||||
if state_path.suffix:
|
||||
return f"{state_path.stem}.{source_segment}{state_path.suffix}"
|
||||
return f"{state_path.name}.{source_segment}.json"
|
||||
|
||||
async def load_state(self, session: AgentSession, *, source_id: str) -> tuple[list[TodoItem], int]:
|
||||
"""Load todo state from disk."""
|
||||
state_path = self._get_state_path(session, source_id=source_id)
|
||||
return await asyncio.to_thread(self._load_state_sync, state_path)
|
||||
|
||||
@staticmethod
|
||||
def _load_state_sync(state_path: Path) -> tuple[list[TodoItem], int]:
|
||||
"""Synchronous helper that performs the disk I/O for ``load_state``."""
|
||||
if not state_path.exists():
|
||||
return [], 1
|
||||
payload = cast(dict[str, Any], json.loads(state_path.read_text(encoding="utf-8")))
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError(f"Todo file {state_path} must contain a JSON object.")
|
||||
raw_items = payload.get("items", [])
|
||||
raw_next_id = payload.get("next_id", 1)
|
||||
if not isinstance(raw_items, list):
|
||||
raise ValueError(f"Todo file {state_path} has a non-list 'items' field.")
|
||||
if not isinstance(raw_next_id, int):
|
||||
raise ValueError(f"Todo file {state_path} has a non-integer 'next_id' field.")
|
||||
items_payload: list[Any] = cast(Any, raw_items)
|
||||
items = _parse_todo_items(items_payload, source_description=f"todo file {state_path}")
|
||||
return items, _safe_next_id(items, raw_next_id)
|
||||
|
||||
async def save_state(self, session: AgentSession, items: list[TodoItem], *, next_id: int, source_id: str) -> None:
|
||||
"""Persist todo state to disk."""
|
||||
state_path = self._get_state_path(session, source_id=source_id)
|
||||
payload = (
|
||||
json.dumps({
|
||||
"items": [item.to_dict(exclude_none=False) for item in items],
|
||||
"next_id": _safe_next_id(items, next_id),
|
||||
})
|
||||
+ "\n"
|
||||
)
|
||||
await asyncio.to_thread(self._save_state_sync, state_path, payload)
|
||||
|
||||
@staticmethod
|
||||
def _save_state_sync(state_path: Path, payload: str) -> None:
|
||||
"""Synchronous helper that atomically writes the JSON state file."""
|
||||
state_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Write to a sibling temp file then atomically replace, so a crash mid-write cannot leave
|
||||
# a truncated state file that breaks every subsequent tool call.
|
||||
temp_path = state_path.with_name(f"{state_path.name}.tmp.{os.getpid()}")
|
||||
try:
|
||||
temp_path.write_text(payload, encoding="utf-8")
|
||||
os.replace(temp_path, state_path)
|
||||
finally:
|
||||
if temp_path.exists():
|
||||
temp_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class TodoProvider(ContextProvider):
|
||||
"""Provide todo management tools and instructions to an agent.
|
||||
|
||||
The ``TodoProvider`` enables agents to create, complete, remove, and query todo items as part of their planning
|
||||
and execution workflow. Todo state is stored in the configured ``TodoStore`` and persists across agent invocations
|
||||
within the same session. By default, state is stored in ``AgentSession.state`` with ``TodoSessionStore``; callers
|
||||
can provide ``TodoFileStore`` or another store implementation for file-backed or custom persistence.
|
||||
|
||||
This provider exposes the following tools to the agent:
|
||||
- ``todos_add``: Add one or more todo items, each with a title and optional description.
|
||||
- ``todos_complete``: Mark one or more todo items as complete by their IDs and reasons.
|
||||
- ``todos_remove``: Remove one or more todo items by their IDs.
|
||||
- ``todos_get_remaining``: Retrieve only incomplete todo items.
|
||||
- ``todos_get_all``: Retrieve all todo items, complete and incomplete.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str = DEFAULT_TODO_SOURCE_ID,
|
||||
*,
|
||||
instructions: str | None = None,
|
||||
store: TodoStore | None = None,
|
||||
) -> None:
|
||||
"""Initialize the todo provider.
|
||||
|
||||
Args:
|
||||
source_id: Unique source ID for the provider.
|
||||
|
||||
Keyword Args:
|
||||
instructions: Optional instruction override.
|
||||
store: Optional todo store override.
|
||||
"""
|
||||
super().__init__(source_id)
|
||||
self.instructions = instructions or DEFAULT_TODO_INSTRUCTIONS
|
||||
self.store = store or TodoSessionStore()
|
||||
# WeakKeyDictionary so per-session locks are evicted automatically when the session is GC'd
|
||||
# rather than accumulating in long-running services that create many sessions.
|
||||
self._mutation_locks: weakref.WeakKeyDictionary[AgentSession, asyncio.Lock] = weakref.WeakKeyDictionary()
|
||||
|
||||
def _mutation_lock(self, session: AgentSession) -> asyncio.Lock:
|
||||
"""Return the per-session lock for read-modify-write todo operations."""
|
||||
lock = self._mutation_locks.get(session)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
self._mutation_locks[session] = lock
|
||||
return lock
|
||||
|
||||
async def before_run(
|
||||
self,
|
||||
*,
|
||||
agent: Any,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Inject todo tools and instructions before the model runs."""
|
||||
del agent, state
|
||||
|
||||
@tool(name="todos_add", approval_mode="never_require")
|
||||
async def todos_add(todos: list[_TodoAddItemSchema]) -> str:
|
||||
"""Add one or more todo items for the current session."""
|
||||
if not todos:
|
||||
raise ValueError("todos must contain at least one item.")
|
||||
|
||||
async with self._mutation_lock(session):
|
||||
existing_items, next_id = await self.store.load_state(session, source_id=self.source_id)
|
||||
created_items: list[TodoItem] = []
|
||||
for raw_todo in todos:
|
||||
todo = _coerce_todo_input(raw_todo)
|
||||
created_item = TodoItem(
|
||||
id=next_id,
|
||||
title=todo.title,
|
||||
description=todo.description.strip() if todo.description is not None else None,
|
||||
)
|
||||
existing_items.append(created_item)
|
||||
created_items.append(created_item)
|
||||
next_id += 1
|
||||
|
||||
await self.store.save_state(session, existing_items, next_id=next_id, source_id=self.source_id)
|
||||
return json.dumps([item.to_dict(exclude_none=False) for item in created_items])
|
||||
|
||||
@tool(name="todos_complete", approval_mode="never_require")
|
||||
async def todos_complete(items: list[_TodoCompleteItemSchema]) -> str:
|
||||
"""Mark one or more todo items as complete.
|
||||
|
||||
Each entry has an id (int) and a reason (string) describing how/why the item was completed.
|
||||
"""
|
||||
if not items:
|
||||
raise ValueError("items must contain at least one entry.")
|
||||
|
||||
parsed = [_coerce_todo_complete_input(entry) for entry in items]
|
||||
ids = [entry.id for entry in parsed]
|
||||
|
||||
async with self._mutation_lock(session):
|
||||
existing_items, next_id = await self.store.load_state(session, source_id=self.source_id)
|
||||
id_set = set(ids)
|
||||
completed_count = 0
|
||||
updated_items: list[TodoItem] = []
|
||||
for item in existing_items:
|
||||
if not item.is_complete and item.id in id_set:
|
||||
updated_items.append(
|
||||
TodoItem(
|
||||
id=item.id,
|
||||
title=item.title,
|
||||
description=item.description,
|
||||
is_complete=True,
|
||||
)
|
||||
)
|
||||
completed_count += 1
|
||||
else:
|
||||
updated_items.append(item)
|
||||
|
||||
if completed_count:
|
||||
await self.store.save_state(session, updated_items, next_id=next_id, source_id=self.source_id)
|
||||
return json.dumps({"completed": completed_count})
|
||||
|
||||
@tool(name="todos_remove", approval_mode="never_require")
|
||||
async def todos_remove(ids: list[int]) -> str:
|
||||
"""Remove one or more todo items by ID."""
|
||||
if not ids:
|
||||
raise ValueError("ids must contain at least one todo ID.")
|
||||
|
||||
async with self._mutation_lock(session):
|
||||
items, next_id = await self.store.load_state(session, source_id=self.source_id)
|
||||
remaining_items = [item for item in items if item.id not in set(ids)]
|
||||
removed_count = len(items) - len(remaining_items)
|
||||
if removed_count:
|
||||
await self.store.save_state(session, remaining_items, next_id=next_id, source_id=self.source_id)
|
||||
return json.dumps({"removed": removed_count})
|
||||
|
||||
@tool(name="todos_get_remaining", approval_mode="never_require")
|
||||
async def todos_get_remaining() -> str:
|
||||
"""Retrieve only incomplete todo items for the current session."""
|
||||
items = [
|
||||
item for item in await self.store.load_items(session, source_id=self.source_id) if not item.is_complete
|
||||
]
|
||||
return json.dumps([item.to_dict(exclude_none=False) for item in items])
|
||||
|
||||
@tool(name="todos_get_all", approval_mode="never_require")
|
||||
async def todos_get_all() -> str:
|
||||
"""Retrieve all todo items for the current session."""
|
||||
items = await self.store.load_items(session, source_id=self.source_id)
|
||||
return json.dumps([item.to_dict(exclude_none=False) for item in items])
|
||||
|
||||
context.extend_instructions(self.source_id, [self.instructions])
|
||||
context.extend_tools(
|
||||
self.source_id,
|
||||
[todos_add, todos_complete, todos_remove, todos_get_remaining, todos_get_all],
|
||||
)
|
||||
current_items = await self.store.load_items(session, source_id=self.source_id)
|
||||
context.extend_messages(
|
||||
self.source_id,
|
||||
[
|
||||
Message(
|
||||
role="user",
|
||||
contents=[
|
||||
"### Current todo list\n"
|
||||
+ (
|
||||
"\n".join(
|
||||
f"- {item.id} [{'done' if item.is_complete else 'open'}] {item.title}"
|
||||
+ (f": {item.description}" if item.description else "")
|
||||
for item in current_items
|
||||
)
|
||||
or "- none yet"
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,632 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import inspect
|
||||
import json
|
||||
from asyncio import sleep
|
||||
from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
from .._feature_stage import ExperimentalFeature, experimental
|
||||
from .._middleware import AgentContext, AgentMiddleware
|
||||
from .._serialization import SerializationMixin
|
||||
from .._sessions import AgentSession
|
||||
from .._types import (
|
||||
AgentResponse,
|
||||
AgentResponseUpdate,
|
||||
Content,
|
||||
FinishReason,
|
||||
FinishReasonLiteral,
|
||||
Message,
|
||||
ResponseStream,
|
||||
)
|
||||
|
||||
DEFAULT_TOOL_APPROVAL_SOURCE_ID = "tool_approval"
|
||||
_FUNCTION_INVOCATION_BUDGET_STATE_KEY = "_function_invocation_budget_state"
|
||||
ALWAYS_APPROVE_PROPERTY = "tool_approval"
|
||||
ALWAYS_APPROVE_SCOPE_PROPERTY = "always_approve"
|
||||
ALWAYS_APPROVE_TOOL: Literal["tool"] = "tool"
|
||||
ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS: Literal["tool_with_arguments"] = "tool_with_arguments"
|
||||
|
||||
_RULES_KEY = "rules"
|
||||
_QUEUED_APPROVAL_REQUESTS_KEY = "queued_approval_requests"
|
||||
_COLLECTED_APPROVAL_RESPONSES_KEY = "collected_approval_responses"
|
||||
|
||||
ToolApprovalScope = Literal["tool", "tool_with_arguments"]
|
||||
ToolApprovalRuleCallback = Callable[[Content], bool | Awaitable[bool]]
|
||||
|
||||
|
||||
def _parse_function_arguments(function_call: Content) -> dict[str, Any]:
|
||||
arguments = function_call.parse_arguments()
|
||||
return dict(arguments or {})
|
||||
|
||||
|
||||
def _serialize_argument_value(value: Any) -> str:
|
||||
return json.dumps(value, sort_keys=True, separators=(",", ":"), default=str)
|
||||
|
||||
|
||||
def _serialize_arguments(function_call: Content) -> dict[str, str]:
|
||||
"""Serialize arguments for exact matching.
|
||||
|
||||
``None`` is reserved on :class:`ToolApprovalRule` for tool-wide rules.
|
||||
An argument-scoped rule for a no-argument call stores ``{}``, so it only
|
||||
matches future no-argument calls and never becomes a wildcard.
|
||||
"""
|
||||
arguments = _parse_function_arguments(function_call)
|
||||
return {key: _serialize_argument_value(value) for key, value in arguments.items()}
|
||||
|
||||
|
||||
def _server_label(function_call: Content) -> str | None:
|
||||
"""Return the hosted-tool server boundary for a function call, if present."""
|
||||
value = function_call.additional_properties.get("server_label")
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
def _content_from_state(value: Any) -> Content:
|
||||
if isinstance(value, Content):
|
||||
return value
|
||||
if isinstance(value, Mapping):
|
||||
return Content.from_dict(cast(Mapping[str, Any], value))
|
||||
raise TypeError(f"Expected Content or mapping state item, got {type(value).__name__}.")
|
||||
|
||||
|
||||
def _contents_from_state(values: Any) -> list[Content]:
|
||||
if not isinstance(values, list):
|
||||
return []
|
||||
state_items = list(cast(Iterable[Any], values))
|
||||
return [_content_from_state(value) for value in state_items]
|
||||
|
||||
|
||||
def _content_to_state(content: Content) -> dict[str, Any]:
|
||||
return content.to_dict()
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class ToolApprovalRule(SerializationMixin):
|
||||
"""A standing rule for approving future matching tool calls."""
|
||||
|
||||
tool_name: str
|
||||
arguments: dict[str, str] | None
|
||||
server_label: str | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: Mapping[str, str] | None = None,
|
||||
*,
|
||||
server_label: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize a tool approval rule.
|
||||
|
||||
Args:
|
||||
tool_name: The function tool name this rule applies to.
|
||||
arguments: Optional canonicalized argument values. When omitted, the
|
||||
rule applies to every call to the tool. Use an empty mapping to
|
||||
match only no-argument calls.
|
||||
|
||||
Keyword Args:
|
||||
server_label: Optional hosted-tool server boundary. Hosted approvals
|
||||
only match future approvals from the same server label.
|
||||
"""
|
||||
normalized_name = tool_name.strip()
|
||||
if not normalized_name:
|
||||
raise ValueError("Tool approval rule tool_name must be a non-empty string.")
|
||||
self.tool_name = normalized_name
|
||||
self.arguments = dict(arguments) if arguments is not None else None
|
||||
self.server_label = server_label
|
||||
|
||||
@classmethod
|
||||
def from_dict(
|
||||
cls,
|
||||
value: MutableMapping[str, Any],
|
||||
/,
|
||||
*,
|
||||
dependencies: MutableMapping[str, Any] | None = None,
|
||||
) -> ToolApprovalRule:
|
||||
"""Create a rule from serialized state."""
|
||||
del dependencies
|
||||
tool_name = value.get("tool_name")
|
||||
if not isinstance(tool_name, str):
|
||||
raise ValueError("Tool approval rule tool_name must be a string.")
|
||||
raw_arguments = value.get("arguments")
|
||||
if raw_arguments is not None and not isinstance(raw_arguments, Mapping):
|
||||
raise ValueError("Tool approval rule arguments must be a mapping or None.")
|
||||
server_label = value.get("server_label")
|
||||
if server_label is not None and not isinstance(server_label, str):
|
||||
raise ValueError("Tool approval rule server_label must be a string or None.")
|
||||
arguments = (
|
||||
{str(key): str(argument_value) for key, argument_value in cast(Mapping[str, Any], raw_arguments).items()}
|
||||
if isinstance(raw_arguments, Mapping)
|
||||
else None
|
||||
)
|
||||
return cls(tool_name=tool_name, arguments=arguments, server_label=server_label)
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
|
||||
"""Serialize the rule."""
|
||||
exclude = exclude or set()
|
||||
payload: dict[str, Any] = {"tool_name": self.tool_name}
|
||||
if "type" not in exclude:
|
||||
payload["type"] = self._get_type_identifier()
|
||||
if self.arguments is not None or not exclude_none:
|
||||
payload["arguments"] = self.arguments
|
||||
if self.server_label is not None or not exclude_none:
|
||||
payload["server_label"] = self.server_label
|
||||
return payload
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class ToolApprovalState(SerializationMixin):
|
||||
"""Session-backed state used by :class:`ToolApprovalMiddleware`."""
|
||||
|
||||
rules: list[ToolApprovalRule]
|
||||
queued_approval_requests: list[Content]
|
||||
collected_approval_responses: list[Content]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
rules: Sequence[ToolApprovalRule | Mapping[str, Any]] | None = None,
|
||||
queued_approval_requests: Sequence[Content | Mapping[str, Any]] | None = None,
|
||||
collected_approval_responses: Sequence[Content | Mapping[str, Any]] | None = None,
|
||||
) -> None:
|
||||
"""Initialize approval state."""
|
||||
self.rules = [
|
||||
rule if isinstance(rule, ToolApprovalRule) else ToolApprovalRule.from_dict(dict(rule))
|
||||
for rule in (rules or [])
|
||||
]
|
||||
self.queued_approval_requests = [
|
||||
item if isinstance(item, Content) else Content.from_dict(item) for item in (queued_approval_requests or [])
|
||||
]
|
||||
self.collected_approval_responses = [
|
||||
item if isinstance(item, Content) else Content.from_dict(item)
|
||||
for item in (collected_approval_responses or [])
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def from_dict(
|
||||
cls,
|
||||
value: MutableMapping[str, Any],
|
||||
/,
|
||||
*,
|
||||
dependencies: MutableMapping[str, Any] | None = None,
|
||||
) -> ToolApprovalState:
|
||||
"""Create state from serialized state."""
|
||||
del dependencies
|
||||
return cls(
|
||||
rules=cast(Sequence[Mapping[str, Any]], value.get("rules", [])),
|
||||
queued_approval_requests=cast(Sequence[Mapping[str, Any]], value.get("queued_approval_requests", [])),
|
||||
collected_approval_responses=cast(
|
||||
Sequence[Mapping[str, Any]],
|
||||
value.get("collected_approval_responses", []),
|
||||
),
|
||||
)
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
|
||||
"""Serialize state."""
|
||||
del exclude_none
|
||||
exclude = exclude or set()
|
||||
payload: dict[str, Any] = {
|
||||
"rules": [rule.to_dict() for rule in self.rules],
|
||||
"queued_approval_requests": [_content_to_state(item) for item in self.queued_approval_requests],
|
||||
"collected_approval_responses": [_content_to_state(item) for item in self.collected_approval_responses],
|
||||
}
|
||||
if "type" not in exclude:
|
||||
payload["type"] = self._get_type_identifier()
|
||||
return payload
|
||||
|
||||
|
||||
def create_always_approve_tool_response(request: Content, *, reason: str | None = None) -> Content:
|
||||
"""Create an approval response that records a standing rule for the whole tool.
|
||||
|
||||
Args:
|
||||
request: The ``function_approval_request`` content to approve.
|
||||
|
||||
Keyword Args:
|
||||
reason: Optional approval reason stored in ``additional_properties``.
|
||||
|
||||
Returns:
|
||||
A ``function_approval_response`` with metadata consumed by
|
||||
:class:`ToolApprovalMiddleware`.
|
||||
"""
|
||||
return _create_always_approve_response(request, ALWAYS_APPROVE_TOOL, reason=reason)
|
||||
|
||||
|
||||
def create_always_approve_tool_with_arguments_response(request: Content, *, reason: str | None = None) -> Content:
|
||||
"""Create an approval response that records a standing rule for the tool and exact arguments."""
|
||||
return _create_always_approve_response(request, ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS, reason=reason)
|
||||
|
||||
|
||||
def _create_always_approve_response(request: Content, scope: ToolApprovalScope, *, reason: str | None) -> Content:
|
||||
response = request.to_function_approval_response(approved=True)
|
||||
metadata: dict[str, Any] = {ALWAYS_APPROVE_SCOPE_PROPERTY: scope}
|
||||
if reason is not None:
|
||||
metadata["reason"] = reason
|
||||
response.additional_properties[ALWAYS_APPROVE_PROPERTY] = metadata
|
||||
return response
|
||||
|
||||
|
||||
def _get_state(session: AgentSession, *, source_id: str) -> ToolApprovalState:
|
||||
raw_state = session.state.get(source_id)
|
||||
if isinstance(raw_state, ToolApprovalState):
|
||||
return raw_state
|
||||
if isinstance(raw_state, MutableMapping):
|
||||
raw_state_mapping = cast(MutableMapping[str, Any], raw_state)
|
||||
return ToolApprovalState(
|
||||
rules=cast(Sequence[Mapping[str, Any]], raw_state_mapping.get(_RULES_KEY, [])),
|
||||
queued_approval_requests=_contents_from_state(raw_state_mapping.get(_QUEUED_APPROVAL_REQUESTS_KEY, [])),
|
||||
collected_approval_responses=_contents_from_state(
|
||||
raw_state_mapping.get(_COLLECTED_APPROVAL_RESPONSES_KEY, []),
|
||||
),
|
||||
)
|
||||
if raw_state is not None:
|
||||
raise TypeError(f"Session state for {source_id!r} must be a mapping, got {type(raw_state).__name__}.")
|
||||
state = ToolApprovalState()
|
||||
session.state[source_id] = state.to_dict(exclude={"type"})
|
||||
return state
|
||||
|
||||
|
||||
def _save_state(session: AgentSession, state: ToolApprovalState, *, source_id: str) -> None:
|
||||
serialized = state.to_dict(exclude={"type"})
|
||||
existing = session.state.get(source_id)
|
||||
if isinstance(existing, MutableMapping):
|
||||
for key, value in cast(MutableMapping[str, Any], existing).items():
|
||||
if key not in serialized and key != "type":
|
||||
serialized[key] = value
|
||||
session.state[source_id] = serialized
|
||||
|
||||
|
||||
def _rule_exists(rules: Sequence[ToolApprovalRule], new_rule: ToolApprovalRule) -> bool:
|
||||
for rule in rules:
|
||||
if rule.tool_name != new_rule.tool_name:
|
||||
continue
|
||||
if rule.server_label != new_rule.server_label:
|
||||
continue
|
||||
if rule.arguments == new_rule.arguments:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _add_rule_if_missing(state: ToolApprovalState, rule: ToolApprovalRule) -> None:
|
||||
if not _rule_exists(state.rules, rule):
|
||||
state.rules.append(rule)
|
||||
|
||||
|
||||
def _function_call_from_request(request: Content) -> Content | None:
|
||||
function_call = request.function_call
|
||||
if function_call is None or function_call.type != "function_call" or function_call.name is None:
|
||||
return None
|
||||
return function_call
|
||||
|
||||
|
||||
def _arguments_match(rule_arguments: Mapping[str, str], function_call: Content) -> bool:
|
||||
call_arguments = _serialize_arguments(function_call) or {}
|
||||
if len(rule_arguments) != len(call_arguments):
|
||||
return False
|
||||
return all(call_arguments.get(key) == value for key, value in rule_arguments.items())
|
||||
|
||||
|
||||
def _matches_rule(request: Content, rules: Sequence[ToolApprovalRule]) -> bool:
|
||||
function_call = _function_call_from_request(request)
|
||||
if function_call is None:
|
||||
return False
|
||||
for rule in rules:
|
||||
if rule.tool_name != function_call.name:
|
||||
continue
|
||||
if rule.server_label != _server_label(function_call):
|
||||
continue
|
||||
if rule.arguments is None:
|
||||
return True
|
||||
if _arguments_match(rule.arguments, function_call):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _get_always_approve_scope(response: Content) -> ToolApprovalScope | None:
|
||||
metadata = response.additional_properties.get(ALWAYS_APPROVE_PROPERTY)
|
||||
if not isinstance(metadata, Mapping):
|
||||
return None
|
||||
metadata_mapping = cast(Mapping[str, Any], metadata)
|
||||
scope = metadata_mapping.get(ALWAYS_APPROVE_SCOPE_PROPERTY)
|
||||
if scope == ALWAYS_APPROVE_TOOL:
|
||||
return ALWAYS_APPROVE_TOOL
|
||||
if scope == ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS:
|
||||
return ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS
|
||||
return None
|
||||
|
||||
|
||||
def _clone_without_always_approve_metadata(response: Content) -> Content:
|
||||
cloned = copy.deepcopy(response)
|
||||
cloned.additional_properties.pop(ALWAYS_APPROVE_PROPERTY, None)
|
||||
return cloned
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.HARNESS)
|
||||
class ToolApprovalMiddleware(AgentMiddleware):
|
||||
"""Coordinate standing tool approvals and queued approval prompts for an agent.
|
||||
|
||||
This middleware is opt-in and requires callers to run the agent with an
|
||||
:class:`AgentSession`, because approval rules and queued requests are stored
|
||||
in session state.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
source_id: str = DEFAULT_TOOL_APPROVAL_SOURCE_ID,
|
||||
auto_approval_rules: Sequence[ToolApprovalRuleCallback] | None = None,
|
||||
) -> None:
|
||||
"""Initialize the middleware.
|
||||
|
||||
Keyword Args:
|
||||
source_id: Session-state key used by this middleware.
|
||||
auto_approval_rules: Optional callbacks that can auto-approve a
|
||||
``function_call``. Each callback receives the function-call
|
||||
content and returns ``True`` to approve it.
|
||||
"""
|
||||
self.source_id = source_id
|
||||
self.auto_approval_rules = tuple(auto_approval_rules or ())
|
||||
|
||||
async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
"""Process one agent invocation."""
|
||||
if context.session is None:
|
||||
raise RuntimeError("ToolApprovalMiddleware requires an AgentSession.")
|
||||
|
||||
state = _get_state(context.session, source_id=self.source_id)
|
||||
context.client_kwargs.setdefault(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, {})
|
||||
context.messages = self._prepare_inbound_messages(context.messages, state)
|
||||
await self._drain_auto_approvable_queue(state)
|
||||
if next_queued := self._pop_next_queued_request(state):
|
||||
_save_state(context.session, state, source_id=self.source_id)
|
||||
context.result = self._response_for_queued_request(next_queued, stream=context.stream)
|
||||
return
|
||||
if context.stream:
|
||||
context.result = self._process_stream(context, call_next, state)
|
||||
return
|
||||
|
||||
while True:
|
||||
context.messages = self._inject_collected_responses(context.messages, state)
|
||||
state_changed = bool(state.collected_approval_responses)
|
||||
state.collected_approval_responses.clear()
|
||||
if state_changed:
|
||||
_save_state(context.session, state, source_id=self.source_id)
|
||||
|
||||
await call_next()
|
||||
if isinstance(context.result, ResponseStream):
|
||||
return
|
||||
if context.result is None:
|
||||
_save_state(context.session, state, source_id=self.source_id)
|
||||
return
|
||||
|
||||
all_auto_approved = await self._process_outbound_messages(context.result.messages, state)
|
||||
_save_state(context.session, state, source_id=self.source_id)
|
||||
if not all_auto_approved:
|
||||
return
|
||||
context.messages = []
|
||||
context.result = None
|
||||
|
||||
def _response_for_queued_request(
|
||||
self,
|
||||
request: Content,
|
||||
*,
|
||||
stream: bool,
|
||||
) -> AgentResponse | ResponseStream[AgentResponseUpdate, AgentResponse]:
|
||||
if not stream:
|
||||
return AgentResponse(messages=[Message(role="assistant", contents=[request])])
|
||||
|
||||
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
|
||||
await sleep(0)
|
||||
yield AgentResponseUpdate(role="assistant", contents=[request])
|
||||
|
||||
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
|
||||
|
||||
def _process_stream(
|
||||
self,
|
||||
context: AgentContext,
|
||||
call_next: Callable[[], Awaitable[None]],
|
||||
state: ToolApprovalState,
|
||||
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
|
||||
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
|
||||
if context.session is None:
|
||||
raise RuntimeError("ToolApprovalMiddleware requires an AgentSession.")
|
||||
while True:
|
||||
context.messages = self._inject_collected_responses(context.messages, state)
|
||||
state_changed = bool(state.collected_approval_responses)
|
||||
state.collected_approval_responses.clear()
|
||||
if state_changed:
|
||||
_save_state(context.session, state, source_id=self.source_id)
|
||||
|
||||
await call_next()
|
||||
if not isinstance(context.result, ResponseStream):
|
||||
raise ValueError("Streaming ToolApprovalMiddleware requires a ResponseStream result.")
|
||||
|
||||
approval_requests: list[Content] = []
|
||||
async for update in context.result:
|
||||
approval_contents = [
|
||||
content for content in update.contents if content.type == "function_approval_request"
|
||||
]
|
||||
if not approval_contents:
|
||||
yield update
|
||||
continue
|
||||
approval_requests.extend(approval_contents)
|
||||
remaining_contents = [
|
||||
content for content in update.contents if content.type != "function_approval_request"
|
||||
]
|
||||
if remaining_contents:
|
||||
raw_finish_reason = update.finish_reason
|
||||
finish_reason: FinishReasonLiteral | FinishReason | None
|
||||
if isinstance(raw_finish_reason, str):
|
||||
finish_reason = FinishReason(raw_finish_reason)
|
||||
else:
|
||||
finish_reason = cast(FinishReasonLiteral | FinishReason | None, raw_finish_reason)
|
||||
yield AgentResponseUpdate(
|
||||
contents=remaining_contents,
|
||||
role=update.role,
|
||||
author_name=update.author_name,
|
||||
agent_id=update.agent_id,
|
||||
response_id=update.response_id,
|
||||
message_id=update.message_id,
|
||||
created_at=update.created_at,
|
||||
finish_reason=finish_reason,
|
||||
continuation_token=update.continuation_token,
|
||||
additional_properties=update.additional_properties,
|
||||
raw_representation=update.raw_representation,
|
||||
)
|
||||
await context.result.get_final_response()
|
||||
if not approval_requests:
|
||||
return
|
||||
|
||||
response_messages = [Message(role="assistant", contents=approval_requests)]
|
||||
all_auto_approved = await self._process_outbound_messages(response_messages, state)
|
||||
_save_state(context.session, state, source_id=self.source_id)
|
||||
if not all_auto_approved:
|
||||
for message in response_messages:
|
||||
if message.contents:
|
||||
yield AgentResponseUpdate(role=message.role, contents=message.contents)
|
||||
return
|
||||
context.messages = []
|
||||
context.result = None
|
||||
|
||||
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
|
||||
|
||||
def _prepare_inbound_messages(self, messages: Sequence[Message], state: ToolApprovalState) -> list[Message]:
|
||||
prepared: list[Message] = []
|
||||
for message in messages:
|
||||
replacement_contents: list[Content] = []
|
||||
changed = False
|
||||
for content in message.contents:
|
||||
if content.type == "function_approval_response":
|
||||
replacement = self._handle_inbound_approval_response(content, state)
|
||||
state.collected_approval_responses.append(replacement)
|
||||
changed = True
|
||||
continue
|
||||
replacement_contents.append(content)
|
||||
|
||||
if not changed:
|
||||
prepared.append(message)
|
||||
continue
|
||||
if replacement_contents:
|
||||
cloned = copy.copy(message)
|
||||
cloned.contents = replacement_contents
|
||||
prepared.append(cloned)
|
||||
return prepared
|
||||
|
||||
def _handle_inbound_approval_response(self, response: Content, state: ToolApprovalState) -> Content:
|
||||
scope = _get_always_approve_scope(response)
|
||||
if scope is None or not response.approved:
|
||||
return response
|
||||
|
||||
function_call = response.function_call
|
||||
if function_call is not None and function_call.type == "function_call" and function_call.name is not None:
|
||||
if scope == ALWAYS_APPROVE_TOOL:
|
||||
_add_rule_if_missing(
|
||||
state,
|
||||
ToolApprovalRule(
|
||||
tool_name=function_call.name,
|
||||
server_label=_server_label(function_call),
|
||||
),
|
||||
)
|
||||
else:
|
||||
_add_rule_if_missing(
|
||||
state,
|
||||
ToolApprovalRule(
|
||||
tool_name=function_call.name,
|
||||
arguments=_serialize_arguments(function_call),
|
||||
server_label=_server_label(function_call),
|
||||
),
|
||||
)
|
||||
return _clone_without_always_approve_metadata(response)
|
||||
|
||||
def _inject_collected_responses(self, messages: Sequence[Message], state: ToolApprovalState) -> list[Message]:
|
||||
if not state.collected_approval_responses:
|
||||
return list(messages)
|
||||
return [Message(role="user", contents=list(state.collected_approval_responses)), *messages]
|
||||
|
||||
async def _drain_auto_approvable_queue(self, state: ToolApprovalState) -> None:
|
||||
remaining: list[Content] = []
|
||||
for request in state.queued_approval_requests:
|
||||
if _matches_rule(request, state.rules) or await self._matches_auto_rule(request):
|
||||
state.collected_approval_responses.append(request.to_function_approval_response(approved=True))
|
||||
continue
|
||||
remaining.append(request)
|
||||
state.queued_approval_requests = remaining
|
||||
|
||||
def _pop_next_queued_request(self, state: ToolApprovalState) -> Content | None:
|
||||
if not state.queued_approval_requests:
|
||||
return None
|
||||
return state.queued_approval_requests.pop(0)
|
||||
|
||||
async def _process_outbound_messages(self, messages: list[Message], state: ToolApprovalState) -> bool:
|
||||
approval_requests = [
|
||||
content
|
||||
for message in messages
|
||||
for content in message.contents
|
||||
if content.type == "function_approval_request"
|
||||
]
|
||||
if not approval_requests:
|
||||
return False
|
||||
|
||||
auto_approved: set[int] = set()
|
||||
unresolved: list[Content] = []
|
||||
for request in approval_requests:
|
||||
if _matches_rule(request, state.rules) or await self._matches_auto_rule(request):
|
||||
state.collected_approval_responses.append(request.to_function_approval_response(approved=True))
|
||||
auto_approved.add(id(request))
|
||||
else:
|
||||
unresolved.append(request)
|
||||
|
||||
if not auto_approved and len(unresolved) <= 1:
|
||||
return False
|
||||
|
||||
queued_ids: set[int] = set()
|
||||
for request in unresolved[1:]:
|
||||
queued_ids.add(id(request))
|
||||
state.queued_approval_requests.append(request)
|
||||
|
||||
remove_ids = auto_approved | queued_ids
|
||||
self._remove_approval_requests(messages, remove_ids)
|
||||
return not unresolved
|
||||
|
||||
@staticmethod
|
||||
def _remove_approval_requests(messages: list[Message], remove_ids: set[int]) -> None:
|
||||
for message_index in range(len(messages) - 1, -1, -1):
|
||||
message = messages[message_index]
|
||||
filtered = [
|
||||
content
|
||||
for content in message.contents
|
||||
if content.type != "function_approval_request" or id(content) not in remove_ids
|
||||
]
|
||||
if len(filtered) == len(message.contents):
|
||||
continue
|
||||
if filtered:
|
||||
message.contents = filtered
|
||||
else:
|
||||
messages.pop(message_index)
|
||||
|
||||
async def _matches_auto_rule(self, request: Content) -> bool:
|
||||
function_call = _function_call_from_request(request)
|
||||
if function_call is None:
|
||||
return False
|
||||
for rule in self.auto_approval_rules:
|
||||
result = rule(function_call)
|
||||
if inspect.isawaitable(result):
|
||||
result = await result
|
||||
if result:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ALWAYS_APPROVE_PROPERTY",
|
||||
"ALWAYS_APPROVE_SCOPE_PROPERTY",
|
||||
"ALWAYS_APPROVE_TOOL",
|
||||
"ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS",
|
||||
"DEFAULT_TOOL_APPROVAL_SOURCE_ID",
|
||||
"ToolApprovalMiddleware",
|
||||
"ToolApprovalRule",
|
||||
"ToolApprovalRuleCallback",
|
||||
"ToolApprovalState",
|
||||
"create_always_approve_tool_response",
|
||||
"create_always_approve_tool_with_arguments_response",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,661 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Mapping, MutableMapping
|
||||
from dataclasses import asdict, is_dataclass
|
||||
from datetime import date, datetime
|
||||
from typing import Any, ClassVar, Protocol, TypeVar, runtime_checkable
|
||||
|
||||
logger = logging.getLogger("agent_framework")
|
||||
|
||||
ClassT = TypeVar("ClassT", bound="SerializationMixin")
|
||||
ProtocolT = TypeVar("ProtocolT", bound="SerializationProtocol")
|
||||
|
||||
# Regex pattern for converting CamelCase to snake_case
|
||||
_CAMEL_TO_SNAKE_PATTERN = re.compile(r"(?<!^)(?=[A-Z])")
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SerializationProtocol(Protocol):
|
||||
"""Protocol for objects that support serialization and deserialization.
|
||||
|
||||
This protocol defines the interface that classes must implement to be compatible
|
||||
with the agent framework's serialization system. Any class implementing both
|
||||
``to_dict()`` and ``from_dict()`` methods will automatically satisfy this protocol
|
||||
and can be used seamlessly with other serializable components.
|
||||
|
||||
The protocol enables type safety and duck typing for serializable objects,
|
||||
ensuring consistent behavior across the framework.
|
||||
|
||||
Examples:
|
||||
The framework's ``Message`` class demonstrates the protocol in action:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Message
|
||||
from agent_framework._serialization import SerializationProtocol
|
||||
|
||||
|
||||
# Message implements SerializationProtocol via SerializationMixin
|
||||
user_msg = Message(role="user", contents=["What's the weather like today?"])
|
||||
|
||||
# Serialize to dictionary - automatic type identification and nested serialization
|
||||
msg_dict = user_msg.to_dict()
|
||||
# Result: {
|
||||
# "type": "chat_message",
|
||||
# "role": {"type": "role", "value": "user"},
|
||||
# "contents": [{"type": "text_content", "text": "What's the weather like today?"}],
|
||||
# "message_id": "...",
|
||||
# "additional_properties": {}
|
||||
# }
|
||||
|
||||
# Deserialize back to Message instance - automatic type reconstruction
|
||||
restored_msg = Message.from_dict(msg_dict)
|
||||
print(restored_msg.text) # "What's the weather like today?"
|
||||
print(restored_msg.role) # "user"
|
||||
|
||||
# Verify protocol compliance (useful for type checking and validation)
|
||||
assert isinstance(user_msg, SerializationProtocol)
|
||||
assert isinstance(restored_msg, SerializationProtocol)
|
||||
|
||||
The protocol is also implemented by simpler classes like ``UsageDetails``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import UsageDetails
|
||||
|
||||
# Create usage tracking instance
|
||||
usage = UsageDetails(input_token_count=150, output_token_count=75, total_token_count=225)
|
||||
|
||||
# Seamless serialization with type preservation
|
||||
usage_dict = usage.to_dict()
|
||||
restored_usage = UsageDetails.from_dict(usage_dict)
|
||||
|
||||
# Both satisfy the SerializationProtocol
|
||||
assert isinstance(usage, SerializationProtocol)
|
||||
assert restored_usage.total_token_count == 225
|
||||
|
||||
The protocol ensures consistent serialization behavior across all framework components,
|
||||
enabling reliable data persistence, API communication, and object reconstruction
|
||||
throughout the agent framework ecosystem.
|
||||
"""
|
||||
|
||||
def to_dict(self, **kwargs: Any) -> dict[str, Any]:
|
||||
"""Convert the instance to a dictionary.
|
||||
|
||||
Keyword Args:
|
||||
kwargs: Additional keyword arguments for serialization.
|
||||
|
||||
Returns:
|
||||
Dictionary representation of the instance.
|
||||
"""
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls: type[ProtocolT], value: MutableMapping[str, Any], /, **kwargs: Any) -> ProtocolT:
|
||||
"""Create an instance from a dictionary.
|
||||
|
||||
Args:
|
||||
value: Dictionary containing the instance data (positional-only).
|
||||
|
||||
Keyword Args:
|
||||
kwargs: Additional keyword arguments for deserialization.
|
||||
|
||||
Returns:
|
||||
New instance of the class.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
def is_serializable(value: Any) -> bool:
|
||||
"""Check if a value is JSON serializable.
|
||||
|
||||
This function tests whether a value can be directly serialized to JSON
|
||||
without custom encoding. It checks for basic Python types that have
|
||||
direct JSON equivalents.
|
||||
|
||||
Args:
|
||||
value: The value to check for JSON serializability.
|
||||
|
||||
Returns:
|
||||
True if the value is one of the basic JSON-serializable types
|
||||
(str, int, float, bool, None, list, dict), False otherwise.
|
||||
|
||||
Note:
|
||||
This function only checks for direct JSON compatibility. Complex objects
|
||||
that implement ``SerializationProtocol`` require conversion via ``to_dict()``
|
||||
before JSON serialization.
|
||||
"""
|
||||
return isinstance(value, (str, int, float, bool, type(None), list, dict))
|
||||
|
||||
|
||||
class SerializationMixin:
|
||||
"""Mixin class providing comprehensive serialization and deserialization capabilities.
|
||||
|
||||
.. note::
|
||||
SerializationMixin is in active development. The API may change in future versions
|
||||
as we continue to improve and extend its functionality.
|
||||
|
||||
This mixin enables classes to automatically handle complex serialization scenarios
|
||||
including nested objects, dependency injection, and type conversion. It provides
|
||||
robust support for converting objects to/from dictionaries and JSON strings while
|
||||
maintaining object relationships and handling external dependencies.
|
||||
|
||||
**Key Features:**
|
||||
|
||||
- Automatic serialization of nested SerializationProtocol objects
|
||||
- Support for lists and dictionaries containing serializable objects
|
||||
- Dependency injection system for non-serializable external dependencies
|
||||
- Flexible exclusion of fields from serialization
|
||||
- Type-safe deserialization with automatic type conversion
|
||||
|
||||
**Constructor Pattern for Nested Objects:**
|
||||
|
||||
Classes using this mixin should handle ``MutableMapping`` inputs in their ``__init__`` method
|
||||
for any parameters that expect ``SerializationMixin`` or ``SerializationProtocol`` instances.
|
||||
This enables automatic conversion of dictionaries to proper object instances during deserialization.
|
||||
|
||||
**Dependency Injection System:**
|
||||
|
||||
The mixin supports injecting external dependencies (like database connections, API clients,
|
||||
or configuration objects) that shouldn't be serialized but are needed at runtime.
|
||||
Fields marked in ``INJECTABLE`` are excluded during serialization and can be provided
|
||||
during deserialization via the ``dependencies`` parameter.
|
||||
|
||||
Examples:
|
||||
**Nested object serialization:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Message
|
||||
from agent_framework._sessions import AgentSession
|
||||
|
||||
|
||||
# AgentSession uses SerializationMixin for state serialization
|
||||
session = AgentSession(session_id="test")
|
||||
|
||||
# Serialization produces a clean dict representation
|
||||
session_dict = session.to_dict()
|
||||
|
||||
# Reconstruction from dictionaries
|
||||
restored = AgentSession.from_dict(session_dict)
|
||||
|
||||
**Framework tools with exclusion patterns:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework._tools import BaseTool
|
||||
|
||||
|
||||
class WeatherTool(BaseTool):
|
||||
\"\"\"Example tool that extends BaseTool with additional properties exclusion.\"\"\"
|
||||
|
||||
# Inherits DEFAULT_EXCLUDE = {"additional_properties"} from BaseTool
|
||||
|
||||
def __init__(self, name: str, api_key: str, **kwargs):
|
||||
super().__init__(name=name, description="Get weather information", **kwargs)
|
||||
self.api_key = api_key # Will be serialized
|
||||
|
||||
# Additional properties are excluded from serialization
|
||||
self.additional_properties = {"version": "1.0", "internal_config": {...}}
|
||||
|
||||
|
||||
weather_tool = WeatherTool(name="get_weather", api_key="secret-key")
|
||||
|
||||
# Serialization excludes additional_properties but includes other fields
|
||||
tool_dict = weather_tool.to_dict()
|
||||
# Result: {
|
||||
# "type": "weather_tool",
|
||||
# "name": "get_weather",
|
||||
# "description": "Get weather information",
|
||||
# "api_key": "secret-key"
|
||||
# # additional_properties excluded due to DEFAULT_EXCLUDE
|
||||
# }
|
||||
|
||||
**Agent framework with injectable dependencies:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import BaseAgent
|
||||
|
||||
|
||||
class CustomAgent(BaseAgent):
|
||||
\"\"\"Custom agent extending BaseAgent with additional functionality.\"\"\"
|
||||
|
||||
# Inherits DEFAULT_EXCLUDE = {"additional_properties"} from BaseAgent
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(name="custom-agent", description="A custom agent", **kwargs)
|
||||
|
||||
# additional_properties stores runtime configuration but isn't serialized
|
||||
self.additional_properties.update({
|
||||
"runtime_context": {...},
|
||||
"session_data": {...}
|
||||
})
|
||||
|
||||
|
||||
agent = CustomAgent(
|
||||
context_provider=[...],
|
||||
middleware=[...]
|
||||
)
|
||||
|
||||
# Serialization captures agent configuration but excludes runtime data
|
||||
agent_dict = agent.to_dict()
|
||||
# Result: {
|
||||
# "type": "custom_agent",
|
||||
# "id": "...",
|
||||
# "name": "custom-agent",
|
||||
# "description": "A custom agent",
|
||||
# "context_provider": [...],
|
||||
# "middleware": [...]
|
||||
# # additional_properties excluded
|
||||
# }
|
||||
|
||||
# Agent can be reconstructed with the same configuration
|
||||
restored_agent = CustomAgent.from_dict(agent_dict)
|
||||
|
||||
This approach enables the agent framework to maintain clean separation between
|
||||
persistent configuration and transient runtime state, allowing agents and tools
|
||||
to be serialized for storage or transmission while preserving their functionality.
|
||||
"""
|
||||
|
||||
DEFAULT_EXCLUDE: ClassVar[set[str]] = set()
|
||||
INJECTABLE: ClassVar[set[str]] = set()
|
||||
_SHALLOW_COPY_FIELDS: ClassVar[set[str]] = {"raw_representation"}
|
||||
|
||||
def __deepcopy__(self, memo: dict[int, Any]) -> SerializationMixin:
|
||||
"""Create a deep copy, preserving ``_SHALLOW_COPY_FIELDS`` by reference.
|
||||
|
||||
Fields listed in ``_SHALLOW_COPY_FIELDS`` may contain LLM SDK objects
|
||||
(e.g., proto/gRPC responses) that are not safe to deep-copy. They are
|
||||
kept as shallow references in the copy; all other attributes are
|
||||
deep-copied normally.
|
||||
"""
|
||||
cls = type(self)
|
||||
result = cls.__new__(cls)
|
||||
memo[id(self)] = result
|
||||
for k, v in self.__dict__.items():
|
||||
if k in cls._SHALLOW_COPY_FIELDS:
|
||||
object.__setattr__(result, k, v)
|
||||
else:
|
||||
object.__setattr__(result, k, copy.deepcopy(v, memo))
|
||||
return result
|
||||
|
||||
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
|
||||
"""Convert the instance and any nested objects to a dictionary.
|
||||
|
||||
This method performs deep serialization, automatically converting nested
|
||||
``SerializationProtocol`` objects, lists, and dictionaries containing
|
||||
serializable objects. Non-serializable objects are skipped with debug logging.
|
||||
|
||||
Fields marked in ``DEFAULT_EXCLUDE`` and ``INJECTABLE`` are automatically
|
||||
excluded from the output, as are any private attributes (starting with '_').
|
||||
|
||||
Keyword Args:
|
||||
exclude: Additional field names to exclude from serialization beyond
|
||||
the default exclusions (``DEFAULT_EXCLUDE`` and ``INJECTABLE``).
|
||||
exclude_none: Whether to exclude None values from the output. When True,
|
||||
None values are omitted from the dictionary. Defaults to True.
|
||||
|
||||
Returns:
|
||||
Dictionary representation of the instance including a 'type' field
|
||||
for type identification during deserialization (unless 'type' is excluded).
|
||||
"""
|
||||
# Combine exclude sets
|
||||
combined_exclude = set(self.DEFAULT_EXCLUDE)
|
||||
if exclude:
|
||||
combined_exclude.update(exclude)
|
||||
combined_exclude.update(self.INJECTABLE)
|
||||
|
||||
# Get all instance attributes
|
||||
result: dict[str, Any] = {} if "type" in combined_exclude else {"type": self._get_type_identifier()}
|
||||
for key, value in self.__dict__.items():
|
||||
if key not in combined_exclude and not key.startswith("_"):
|
||||
if exclude_none and value is None:
|
||||
continue
|
||||
# Recursively serialize SerializationProtocol objects
|
||||
if isinstance(value, SerializationProtocol):
|
||||
result[key] = value.to_dict(exclude=exclude, exclude_none=exclude_none)
|
||||
continue
|
||||
# Handle lists containing SerializationProtocol objects
|
||||
if isinstance(value, list):
|
||||
value_as_list: list[Any] = []
|
||||
for item in value: # pyright: ignore[reportUnknownVariableType]
|
||||
if isinstance(item, SerializationProtocol):
|
||||
value_as_list.append(item.to_dict(exclude=exclude, exclude_none=exclude_none))
|
||||
continue
|
||||
if is_serializable(item):
|
||||
value_as_list.append(item)
|
||||
continue
|
||||
logger.debug(
|
||||
f"Skipping non-serializable item in list attribute '{key}' of type {type(item).__name__}" # pyright: ignore[reportUnknownArgumentType]
|
||||
)
|
||||
result[key] = value_as_list
|
||||
continue
|
||||
# Handle dicts containing SerializationProtocol values
|
||||
if isinstance(value, dict):
|
||||
from datetime import date, datetime, time
|
||||
|
||||
serialized_dict: dict[str, Any] = {}
|
||||
for raw_key, v in value.items(): # pyright: ignore[reportUnknownVariableType]
|
||||
dict_key = str(raw_key) # pyright: ignore[reportUnknownArgumentType]
|
||||
if isinstance(v, SerializationProtocol):
|
||||
serialized_dict[dict_key] = v.to_dict(exclude=exclude, exclude_none=exclude_none)
|
||||
continue
|
||||
# Convert datetime objects to strings
|
||||
if isinstance(v, (datetime, date, time)):
|
||||
serialized_dict[dict_key] = str(v)
|
||||
continue
|
||||
# Check if the value is JSON serializable
|
||||
if is_serializable(v):
|
||||
serialized_dict[dict_key] = v
|
||||
continue
|
||||
logger.debug(
|
||||
f"Skipping non-serializable value for key '{dict_key}' in dict attribute '{key}' "
|
||||
f"of type {type(v).__name__}" # pyright: ignore[reportUnknownArgumentType]
|
||||
)
|
||||
result[key] = serialized_dict
|
||||
continue
|
||||
# Directly include JSON serializable values
|
||||
if is_serializable(value):
|
||||
result[key] = value
|
||||
continue
|
||||
logger.debug(f"Skipping non-serializable attribute '{key}' of type {type(value).__name__}")
|
||||
|
||||
return result
|
||||
|
||||
def to_json(self, *, exclude: set[str] | None = None, exclude_none: bool = True, **kwargs: Any) -> str:
|
||||
"""Convert the instance to a JSON string.
|
||||
|
||||
This is a convenience method that calls ``to_dict()`` and then serializes
|
||||
the result using ``json.dumps()``. All the same serialization rules apply
|
||||
as in ``to_dict()``, including automatic exclusion of injectable dependencies
|
||||
and deep serialization of nested objects.
|
||||
|
||||
Keyword Args:
|
||||
exclude: Additional field names to exclude from serialization.
|
||||
exclude_none: Whether to exclude None values from the output. Defaults to True.
|
||||
**kwargs: Additional keyword arguments passed through to ``json.dumps()``.
|
||||
Common options include ``indent`` for pretty-printing and
|
||||
``ensure_ascii`` for Unicode handling.
|
||||
|
||||
Returns:
|
||||
JSON string representation of the instance.
|
||||
"""
|
||||
return json.dumps(self.to_dict(exclude=exclude, exclude_none=exclude_none), **kwargs)
|
||||
|
||||
@classmethod
|
||||
def from_dict(
|
||||
cls: type[ClassT], value: MutableMapping[str, Any], /, *, dependencies: MutableMapping[str, Any] | None = None
|
||||
) -> ClassT:
|
||||
"""Create an instance from a dictionary with optional dependency injection.
|
||||
|
||||
This method reconstructs an object from its dictionary representation, automatically
|
||||
handling type conversion and dependency injection. It supports three patterns of
|
||||
dependency injection to handle different scenarios where external dependencies
|
||||
need to be provided at deserialization time.
|
||||
|
||||
Args:
|
||||
value: The dictionary containing the instance data (positional-only).
|
||||
Must include a 'type' field matching the class type identifier.
|
||||
|
||||
Keyword Args:
|
||||
dependencies: A nested dictionary mapping type identifiers to their injectable dependencies.
|
||||
The structure varies based on injection pattern:
|
||||
|
||||
- **Simple injection**: ``{"<type>": {"<parameter>": value}}``
|
||||
- **Dict parameter injection**: ``{"<type>": {"<dict-parameter>": {"<key>": value}}}``
|
||||
- **Instance-specific injection**: ``{"<type>": {"<field>:<value>": {"<parameter>": value}}}``
|
||||
|
||||
Returns:
|
||||
New instance of the class with injected dependencies.
|
||||
|
||||
Raises:
|
||||
ValueError: If the 'type' field in the data doesn't match the class type identifier.
|
||||
|
||||
Examples:
|
||||
**Simple Client Injection** - OpenAI client dependency injection:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
|
||||
# OpenAI chat client requires an AsyncOpenAI client instance.
|
||||
# The client dependency is excluded from serialization.
|
||||
|
||||
# Serialized data contains only the model configuration
|
||||
client_data = {
|
||||
"type": "open_ai_chat_client",
|
||||
"model": "gpt-4o-mini",
|
||||
# client is excluded from serialization
|
||||
}
|
||||
|
||||
# Provide the OpenAI client during deserialization
|
||||
openai_client = AsyncOpenAI(api_key="your-api-key")
|
||||
dependencies = {"open_ai_chat_client": {"client": openai_client}}
|
||||
|
||||
# The chat client is reconstructed with the OpenAI client injected
|
||||
client = OpenAIChatClient.from_dict(client_data, dependencies=dependencies)
|
||||
# Now ready to make API calls with the injected client
|
||||
|
||||
**Function Injection for Tools** - FunctionTool runtime dependency:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import FunctionTool
|
||||
from typing import Annotated
|
||||
|
||||
|
||||
# Define a function to be wrapped
|
||||
async def get_current_weather(location: Annotated[str, "The city name"]) -> str:
|
||||
# In real implementation, this would call a weather API
|
||||
return f"Current weather in {location}: 72°F and sunny"
|
||||
|
||||
|
||||
# FunctionTool has INJECTABLE = {"func"}
|
||||
function_data = {
|
||||
"type": "function_tool",
|
||||
"name": "get_weather",
|
||||
"description": "Get current weather for a location",
|
||||
# func is excluded from serialization
|
||||
}
|
||||
|
||||
# Inject the actual function implementation during deserialization
|
||||
dependencies = {"function_tool": {"func": get_current_weather}}
|
||||
|
||||
# Reconstruct the FunctionTool with the callable injected
|
||||
weather_func = FunctionTool.from_dict(function_data, dependencies=dependencies)
|
||||
# The function is now callable and ready for agent use
|
||||
|
||||
**MiddlewareTypes Context Injection** - Agent execution context:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework._middleware import AgentContext
|
||||
from agent_framework import BaseAgent
|
||||
|
||||
# AgentContext has INJECTABLE = {"agent", "result"}
|
||||
context_data = {
|
||||
"type": "agent_context",
|
||||
"messages": [{"role": "user", "text": "Hello"}],
|
||||
"stream": False,
|
||||
"metadata": {"session_id": "abc123"},
|
||||
# agent and result are excluded from serialization
|
||||
}
|
||||
|
||||
# Inject agent and result during middleware processing
|
||||
my_agent = BaseAgent(name="test-agent")
|
||||
dependencies = {
|
||||
"agent_context": {
|
||||
"agent": my_agent,
|
||||
"result": None, # Will be populated during execution
|
||||
}
|
||||
}
|
||||
|
||||
# Reconstruct context with agent dependency for middleware chain
|
||||
context = AgentContext.from_dict(context_data, dependencies=dependencies)
|
||||
# MiddlewareTypes can now access context.agent and process the execution
|
||||
|
||||
This injection system allows the agent framework to maintain clean separation
|
||||
between serializable configuration and runtime dependencies like API clients,
|
||||
functions, and execution contexts that cannot or should not be persisted.
|
||||
"""
|
||||
if dependencies is None:
|
||||
dependencies = {}
|
||||
|
||||
# Get the type identifier
|
||||
type_id = cls._get_type_identifier(value)
|
||||
|
||||
if (supplied_type := value.get("type")) and supplied_type != type_id:
|
||||
raise ValueError(f"Type mismatch: expected '{type_id}', got '{supplied_type}'")
|
||||
|
||||
# Create a copy of the value dict to work with, filtering out the 'type' key
|
||||
kwargs = {k: v for k, v in value.items() if k != "type"}
|
||||
|
||||
# Process dependencies using dict-based structure
|
||||
type_deps = dependencies.get(type_id, {})
|
||||
for dep_key, dep_value in type_deps.items():
|
||||
# Check if this is an instance-specific dependency (field:name format)
|
||||
if ":" in dep_key:
|
||||
field, name = dep_key.split(":", 1)
|
||||
# Only apply if the instance matches
|
||||
if kwargs.get(field) == name and isinstance(dep_value, dict):
|
||||
# Apply instance-specific dependencies
|
||||
for raw_param_name, param_value in dep_value.items(): # pyright: ignore[reportUnknownVariableType]
|
||||
param_name = str(raw_param_name) # pyright: ignore[reportUnknownArgumentType]
|
||||
if param_name not in cls.INJECTABLE:
|
||||
logger.debug(
|
||||
f"Dependency '{param_name}' for type '{type_id}' is not in INJECTABLE set. "
|
||||
f"Available injectable parameters: {cls.INJECTABLE}"
|
||||
)
|
||||
# Handle nested dict parameters
|
||||
if (
|
||||
isinstance(param_value, dict)
|
||||
and param_name in kwargs
|
||||
and isinstance(kwargs[param_name], dict)
|
||||
):
|
||||
kwargs[param_name].update(param_value)
|
||||
else:
|
||||
kwargs[param_name] = param_value
|
||||
else:
|
||||
# Regular parameter dependency
|
||||
if dep_key not in cls.INJECTABLE:
|
||||
logger.debug(
|
||||
f"Dependency '{dep_key}' for type '{type_id}' is not in INJECTABLE set. "
|
||||
f"Available injectable parameters: {cls.INJECTABLE}"
|
||||
)
|
||||
# Handle dict parameters - merge if both are dicts
|
||||
if isinstance(dep_value, dict) and dep_key in kwargs and isinstance(kwargs[dep_key], dict):
|
||||
kwargs[dep_key].update(dep_value)
|
||||
else:
|
||||
kwargs[dep_key] = dep_value
|
||||
|
||||
return cls(**kwargs)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls: type[ClassT], value: str, /, *, dependencies: MutableMapping[str, Any] | None = None) -> ClassT:
|
||||
"""Create an instance from a JSON string.
|
||||
|
||||
This is a convenience method that parses the JSON string using ``json.loads()``
|
||||
and then calls ``from_dict()`` to reconstruct the object. All dependency injection
|
||||
capabilities are available through the ``dependencies`` parameter.
|
||||
|
||||
Args:
|
||||
value: The JSON string containing the instance data (positional-only).
|
||||
Must be valid JSON that deserializes to a dictionary with a 'type' field.
|
||||
|
||||
Keyword Args:
|
||||
dependencies: A nested dictionary mapping type identifiers to their injectable dependencies.
|
||||
See :meth:`from_dict` for detailed structure and examples of the three
|
||||
injection patterns (simple, dict parameter, and instance-specific).
|
||||
|
||||
Returns:
|
||||
New instance of the class with any specified dependencies injected.
|
||||
|
||||
Raises:
|
||||
json.JSONDecodeError: If the JSON string is malformed.
|
||||
ValueError: If the parsed data doesn't contain a valid 'type' field.
|
||||
"""
|
||||
data = json.loads(value)
|
||||
return cls.from_dict(data, dependencies=dependencies)
|
||||
|
||||
@classmethod
|
||||
def _get_type_identifier(cls, value: Mapping[str, Any] | None = None) -> str:
|
||||
"""Get the type identifier for this class.
|
||||
|
||||
The type identifier is used in serialized data to enable proper deserialization.
|
||||
It follows a priority order to determine the identifier:
|
||||
|
||||
1. If ``value`` contains a 'type' field, return that value (for ``from_dict``)
|
||||
2. If the class has a ``type`` attribute, use that value (instance-level)
|
||||
3. If the class has a ``TYPE`` attribute, use that value (class-level constant)
|
||||
4. Otherwise, convert the class name to snake_case as fallback
|
||||
|
||||
Args:
|
||||
value: Optional mapping containing serialized data that may have a 'type' field.
|
||||
|
||||
Returns:
|
||||
Type identifier string used for serialization and dependency injection mapping.
|
||||
"""
|
||||
# for from_dict
|
||||
if value and (type_ := value.get("type")) and isinstance(type_, str):
|
||||
return type_
|
||||
# for todict when defined per instance
|
||||
if (type_ := getattr(cls, "type", None)) and isinstance(type_, str):
|
||||
return type_
|
||||
# for both when defined on class.
|
||||
if (type_ := getattr(cls, "TYPE", None)) and isinstance(type_, str):
|
||||
return type_
|
||||
# Fallback and default
|
||||
# Convert class name to snake_case
|
||||
return _CAMEL_TO_SNAKE_PATTERN.sub("_", cls.__name__).lower()
|
||||
|
||||
|
||||
def make_json_safe(obj: Any) -> Any:
|
||||
"""Recursively convert an object to a JSON-serializable form.
|
||||
|
||||
Handles dataclasses, Pydantic models, objects with ``to_dict``/``dict``/``__dict__``,
|
||||
datetimes, lists, dicts, and primitives. Falls back to ``str()`` for any remaining
|
||||
non-serializable value so that ``json.dumps`` never raises a ``TypeError``.
|
||||
|
||||
Args:
|
||||
obj: Object to make JSON safe.
|
||||
|
||||
Returns:
|
||||
A JSON-serializable version of the object.
|
||||
"""
|
||||
if obj is None or isinstance(obj, (str, int, float, bool)):
|
||||
return obj
|
||||
if isinstance(obj, (datetime, date)):
|
||||
return obj.isoformat()
|
||||
if is_dataclass(obj) and not isinstance(obj, type):
|
||||
return make_json_safe(asdict(obj))
|
||||
if callable(getattr(obj, "model_dump", None)):
|
||||
try:
|
||||
return make_json_safe(obj.model_dump()) # type: ignore[no-any-return]
|
||||
except TypeError:
|
||||
pass
|
||||
if callable(getattr(obj, "to_dict", None)):
|
||||
try:
|
||||
return make_json_safe(obj.to_dict()) # type: ignore[no-any-return]
|
||||
except TypeError:
|
||||
pass
|
||||
if callable(getattr(obj, "dict", None)):
|
||||
try:
|
||||
return make_json_safe(obj.dict()) # type: ignore[no-any-return]
|
||||
except TypeError:
|
||||
pass
|
||||
if isinstance(obj, dict):
|
||||
return {str(key): make_json_safe(value) for key, value in obj.items()} # type: ignore[misc]
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [make_json_safe(item) for item in obj] # type: ignore[misc]
|
||||
if hasattr(obj, "__dict__"):
|
||||
return {key: make_json_safe(value) for key, value in vars(obj).items()}
|
||||
return str(obj)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,293 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Generic settings loader with environment variable resolution.
|
||||
|
||||
This module provides a ``load_settings()`` function that populates a ``TypedDict``
|
||||
from environment variables, ``.env`` files, and explicit overrides. It replaces
|
||||
the previous pydantic-settings-based ``AFBaseSettings`` with a lighter-weight,
|
||||
function-based approach that has no pydantic-settings dependency.
|
||||
|
||||
Usage::
|
||||
|
||||
class MySettings(TypedDict, total=False):
|
||||
api_key: str | None # optional — resolves to None if not set
|
||||
model: str | None # optional by default
|
||||
source_a: str | None
|
||||
source_b: str | None
|
||||
|
||||
|
||||
# Make model required; require exactly one of source_a / source_b:
|
||||
settings = load_settings(
|
||||
MySettings,
|
||||
env_prefix="MY_APP_",
|
||||
required_fields=["model", ("source_a", "source_b")],
|
||||
model="gpt-4",
|
||||
source_a="value",
|
||||
)
|
||||
settings["api_key"] # type-checked dict access
|
||||
settings["model"] # str | None per type, but guaranteed not None at runtime
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Callable, Sequence
|
||||
from contextlib import suppress
|
||||
from typing import Any, Union, get_args, get_origin, get_type_hints
|
||||
|
||||
from dotenv import dotenv_values
|
||||
|
||||
from .exceptions import SettingNotFoundError
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # pragma: no cover
|
||||
|
||||
|
||||
SettingsT = TypeVar("SettingsT", default=dict[str, Any])
|
||||
|
||||
|
||||
class SecretString(str):
|
||||
"""A string subclass that masks its value in repr() to prevent accidental exposure.
|
||||
|
||||
SecretString behaves exactly like a regular string in all operations,
|
||||
but its repr() shows '**********' instead of the actual value.
|
||||
This helps prevent secrets from being accidentally logged or displayed.
|
||||
|
||||
It also provides a ``get_secret_value()`` method for backward compatibility
|
||||
with code that previously used ``pydantic.SecretStr``.
|
||||
|
||||
Example:
|
||||
```python
|
||||
api_key = SecretString("sk-secret-key")
|
||||
print(api_key) # sk-secret-key (normal string behavior)
|
||||
print(repr(api_key)) # SecretString('**********')
|
||||
print(f"Key: {api_key}") # Key: sk-secret-key
|
||||
print(api_key.get_secret_value()) # sk-secret-key
|
||||
```
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Return a masked representation to prevent secret exposure."""
|
||||
return "SecretString('**********')"
|
||||
|
||||
def get_secret_value(self) -> str:
|
||||
"""Return the underlying string value.
|
||||
|
||||
Provided for backward compatibility with ``pydantic.SecretStr``.
|
||||
Since SecretString *is* a str, this simply returns ``str(self)``.
|
||||
"""
|
||||
return str(self)
|
||||
|
||||
|
||||
def _coerce_value(value: str, target_type: type) -> Any:
|
||||
"""Coerce a string value to the target type."""
|
||||
origin = get_origin(target_type)
|
||||
args = get_args(target_type)
|
||||
|
||||
# Handle Union types (e.g., str | None) — try each non-None arm
|
||||
if origin is type(None):
|
||||
return None
|
||||
|
||||
if args and type(None) in args:
|
||||
for arg in args:
|
||||
if arg is not type(None):
|
||||
with suppress(ValueError, TypeError):
|
||||
return _coerce_value(value, arg)
|
||||
return value
|
||||
|
||||
# Handle SecretString
|
||||
if target_type is SecretString or (isinstance(target_type, type) and issubclass(target_type, SecretString)):
|
||||
return SecretString(value)
|
||||
|
||||
# Handle basic types
|
||||
if target_type is str:
|
||||
return value
|
||||
if target_type is int:
|
||||
return int(value)
|
||||
if target_type is float:
|
||||
return float(value)
|
||||
if target_type is bool:
|
||||
return value.lower() in ("true", "1", "yes", "on")
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def _check_override_type(value: Any, field_type: type, field_name: str) -> None:
|
||||
"""Validate that *value* is compatible with *field_type*.
|
||||
|
||||
Raises ``ValueError`` when the override is clearly
|
||||
incompatible (e.g. a ``dict`` passed where ``str`` is expected).
|
||||
Callable values and ``None`` are always accepted.
|
||||
"""
|
||||
if value is None:
|
||||
return
|
||||
|
||||
# Callables are always allowed (e.g. lazy token providers)
|
||||
if callable(value) and not isinstance(value, (str, bytes)):
|
||||
return
|
||||
|
||||
# Collect the concrete types that *field_type* allows
|
||||
origin = get_origin(field_type)
|
||||
args = get_args(field_type)
|
||||
|
||||
allowed: tuple[type, ...]
|
||||
if origin is Union or origin is type(int | str):
|
||||
allowed = tuple(a for a in args if isinstance(a, type) and a is not type(None))
|
||||
# If any arm is a Callable, allow anything callable
|
||||
if any(get_origin(a) is Callable or a is Callable for a in args):
|
||||
return
|
||||
elif isinstance(field_type, type):
|
||||
allowed = (field_type,)
|
||||
else:
|
||||
return # complex / unknown annotation — skip check
|
||||
|
||||
if not allowed:
|
||||
return
|
||||
|
||||
if not isinstance(value, allowed):
|
||||
# Allow str for SecretString fields (will be coerced)
|
||||
if isinstance(value, str) and any(isinstance(a, type) and issubclass(a, str) for a in allowed):
|
||||
return
|
||||
# Allow int for float fields (standard numeric promotion)
|
||||
if isinstance(value, int) and float in allowed:
|
||||
return
|
||||
|
||||
allowed_names = ", ".join(t.__name__ for t in allowed)
|
||||
raise ValueError(
|
||||
f"Invalid type for setting '{field_name}': expected {allowed_names}, got {type(value).__name__}."
|
||||
)
|
||||
|
||||
|
||||
def load_settings(
|
||||
settings_type: type[SettingsT],
|
||||
*,
|
||||
env_prefix: str = "",
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
required_fields: Sequence[str | tuple[str, ...]] | None = None,
|
||||
**overrides: Any,
|
||||
) -> SettingsT:
|
||||
"""Load settings from explicit overrides, an optional ``.env`` file, and environment variables.
|
||||
|
||||
The *settings_type* must be a ``TypedDict`` subclass. Values are resolved in
|
||||
this order (highest priority first):
|
||||
|
||||
1. Explicit keyword *overrides* (``None`` values are filtered out).
|
||||
2. A ``.env`` file (when *env_file_path* is explicitly provided).
|
||||
3. Environment variables (``<env_prefix><FIELD_NAME>``).
|
||||
4. Default values — fields with class-level defaults on the TypedDict, or
|
||||
``None`` for optional fields.
|
||||
|
||||
Entries in *required_fields* are validated after resolution:
|
||||
|
||||
- A **string** entry means the field must resolve to a non-``None`` value.
|
||||
- A **tuple** entry means exactly one field in the group must be non-``None``
|
||||
(mutually exclusive).
|
||||
|
||||
Args:
|
||||
settings_type: A ``TypedDict`` class describing the settings schema.
|
||||
env_prefix: Prefix for environment variable lookup (e.g. ``"OPENAI_"``).
|
||||
env_file_path: Path to ``.env`` file. When provided, the file is required
|
||||
and values are resolved before process environment variables.
|
||||
env_file_encoding: Encoding for reading the ``.env`` file. Defaults to ``"utf-8"``.
|
||||
required_fields: Field names (``str``) that must resolve to a non-``None``
|
||||
value, or tuples of field names where exactly one must be set.
|
||||
**overrides: Field values. ``None`` values are ignored so that callers can
|
||||
forward optional parameters without masking env-var / default resolution.
|
||||
|
||||
Returns:
|
||||
A populated dict matching *settings_type*.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If *env_file_path* was provided but the file does not exist.
|
||||
SettingNotFoundError: If a required field could not be resolved from any
|
||||
source, or if a mutually exclusive constraint is violated.
|
||||
ValueError: If an override value has an incompatible type.
|
||||
"""
|
||||
encoding = env_file_encoding or "utf-8"
|
||||
|
||||
loaded_dotenv_values: dict[str, str] = {}
|
||||
if env_file_path is not None:
|
||||
if not os.path.exists(env_file_path):
|
||||
raise FileNotFoundError(env_file_path)
|
||||
|
||||
raw_dotenv_values = dotenv_values(dotenv_path=env_file_path, encoding=encoding)
|
||||
loaded_dotenv_values = {key: value for key, value in raw_dotenv_values.items() if value is not None}
|
||||
|
||||
# Filter out None overrides so defaults / env vars are preserved
|
||||
overrides = {k: v for k, v in overrides.items() if v is not None}
|
||||
|
||||
# Get field type hints from the TypedDict
|
||||
hints = get_type_hints(settings_type)
|
||||
|
||||
result: dict[str, Any] = {}
|
||||
for field_name, field_type in hints.items():
|
||||
# 1. Explicit override wins
|
||||
if field_name in overrides:
|
||||
override_value = overrides[field_name]
|
||||
_check_override_type(override_value, field_type, field_name)
|
||||
# Coerce plain str → SecretString if the annotation expects it
|
||||
if isinstance(override_value, str) and not isinstance(override_value, SecretString):
|
||||
with suppress(ValueError, TypeError):
|
||||
coerced = _coerce_value(override_value, field_type)
|
||||
if isinstance(coerced, SecretString):
|
||||
override_value = coerced
|
||||
result[field_name] = override_value
|
||||
continue
|
||||
|
||||
env_var_name = f"{env_prefix}{field_name.upper()}"
|
||||
|
||||
# 2. Optional .env value (only when env_file_path is explicitly provided)
|
||||
if loaded_dotenv_values:
|
||||
dotenv_value = loaded_dotenv_values.get(env_var_name)
|
||||
if dotenv_value is not None:
|
||||
try:
|
||||
result[field_name] = _coerce_value(dotenv_value, field_type)
|
||||
except (ValueError, TypeError):
|
||||
result[field_name] = dotenv_value
|
||||
continue
|
||||
|
||||
# 3. Environment variable
|
||||
env_value = os.getenv(env_var_name)
|
||||
if env_value is not None:
|
||||
try:
|
||||
result[field_name] = _coerce_value(env_value, field_type)
|
||||
except (ValueError, TypeError):
|
||||
result[field_name] = env_value
|
||||
continue
|
||||
|
||||
# 4. Default from TypedDict class-level defaults, or None for optional fields
|
||||
if hasattr(settings_type, field_name):
|
||||
result[field_name] = getattr(settings_type, field_name)
|
||||
else:
|
||||
result[field_name] = None
|
||||
|
||||
# Validate required fields after all resolution
|
||||
if required_fields:
|
||||
for entry in required_fields:
|
||||
if isinstance(entry, str):
|
||||
# Single required field
|
||||
if result.get(entry) is None:
|
||||
env_var_name = f"{env_prefix}{entry.upper()}"
|
||||
raise SettingNotFoundError(
|
||||
f"Required setting '{entry}' was not provided. "
|
||||
f"Set it via the '{entry}' parameter or the "
|
||||
f"'{env_var_name}' environment variable."
|
||||
)
|
||||
else:
|
||||
# Mutually exclusive group — exactly one must be set
|
||||
set_fields = [f for f in entry if result.get(f) is not None]
|
||||
if len(set_fields) == 0:
|
||||
names = ", ".join(f"'{f}'" for f in entry)
|
||||
raise SettingNotFoundError(f"Exactly one of {names} must be provided, but none was set.")
|
||||
if len(set_fields) > 1:
|
||||
all_names = ", ".join(f"'{f}'" for f in entry)
|
||||
set_names = ", ".join(f"'{f}'" for f in set_fields)
|
||||
raise SettingNotFoundError(
|
||||
f"Only one of {all_names} may be provided, but multiple were set: {set_names}."
|
||||
)
|
||||
|
||||
return result # type: ignore[return-value]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,134 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Final
|
||||
|
||||
from . import __version__ as version_info
|
||||
|
||||
logger = logging.getLogger("agent_framework")
|
||||
|
||||
|
||||
# Note that if this environment variable does not exist, user agent telemetry is enabled.
|
||||
USER_AGENT_TELEMETRY_DISABLED_ENV_VAR = "AGENT_FRAMEWORK_USER_AGENT_DISABLED"
|
||||
IS_TELEMETRY_ENABLED = os.environ.get(USER_AGENT_TELEMETRY_DISABLED_ENV_VAR, "false").lower() not in ["true", "1"]
|
||||
|
||||
APP_INFO = (
|
||||
{
|
||||
"agent-framework-version": f"python/{version_info}",
|
||||
}
|
||||
if IS_TELEMETRY_ENABLED
|
||||
else None
|
||||
)
|
||||
USER_AGENT_KEY: Final[str] = "User-Agent"
|
||||
HTTP_USER_AGENT: Final[str] = "agent-framework-python"
|
||||
AGENT_FRAMEWORK_USER_AGENT = f"{HTTP_USER_AGENT}/{version_info}"
|
||||
|
||||
# This environment variable is reserved by the Foundry hosting environment to
|
||||
# indicate that the agent is running in a hosted environment.
|
||||
_FOUNDRY_HOSTING_ENV_VAR = "FOUNDRY_HOSTING_ENVIRONMENT"
|
||||
# This prefix is added to the user agent string when the agent is running in a hosted environment.
|
||||
_HOSTED_USER_AGENT_PREFIX = "foundry-hosting"
|
||||
|
||||
_user_agent_prefixes: set[str] = set()
|
||||
_hosted_env_detected: bool = False
|
||||
|
||||
|
||||
def _add_user_agent_prefix(prefix: str) -> None:
|
||||
"""Permanently add a prefix to the user agent string.
|
||||
|
||||
This is used by hosting layers to identify themselves in telemetry.
|
||||
Once added, the prefix applies to all subsequent user agent strings.
|
||||
|
||||
Args:
|
||||
prefix: The prefix to add (e.g. "foundry-hosting").
|
||||
"""
|
||||
if prefix:
|
||||
_user_agent_prefixes.add(prefix)
|
||||
|
||||
|
||||
def _detect_hosted_environment() -> None:
|
||||
"""Detect if running in a hosted environment and add the user agent prefix.
|
||||
|
||||
Checks the ``FOUNDRY_HOSTING_ENVIRONMENT`` env var first, then falls back
|
||||
to checking whether the agent server SDK is installed (via
|
||||
``importlib.util.find_spec``) before importing it, to avoid unnecessary
|
||||
import overhead for non-hosted scenarios.
|
||||
"""
|
||||
global _hosted_env_detected
|
||||
if _hosted_env_detected:
|
||||
return
|
||||
|
||||
if (env_value := os.environ.get(_FOUNDRY_HOSTING_ENV_VAR)) is not None:
|
||||
# Env var exists — trust its value and skip the fallback.
|
||||
if env_value:
|
||||
_add_user_agent_prefix(_HOSTED_USER_AGENT_PREFIX)
|
||||
_hosted_env_detected = True
|
||||
return
|
||||
|
||||
# Env var not set — fall back to AgentConfig as a second layer of defense.
|
||||
# Use find_spec to avoid the cost of a full import when the SDK is not installed.
|
||||
import importlib.util
|
||||
|
||||
try:
|
||||
if importlib.util.find_spec("azure.ai.agentserver.core") is None:
|
||||
return
|
||||
except (ImportError, ValueError):
|
||||
return
|
||||
with contextlib.suppress(ImportError, AttributeError):
|
||||
from azure.ai.agentserver.core import (
|
||||
AgentConfig,
|
||||
)
|
||||
|
||||
if AgentConfig.from_env().is_hosted:
|
||||
_add_user_agent_prefix(_HOSTED_USER_AGENT_PREFIX)
|
||||
_hosted_env_detected = True
|
||||
|
||||
|
||||
def get_user_agent() -> str:
|
||||
"""Return the full user agent string including any registered prefixes."""
|
||||
_detect_hosted_environment()
|
||||
if not _user_agent_prefixes:
|
||||
return AGENT_FRAMEWORK_USER_AGENT
|
||||
return f"{'/'.join(sorted(_user_agent_prefixes))}/{AGENT_FRAMEWORK_USER_AGENT}"
|
||||
|
||||
|
||||
def prepend_agent_framework_to_user_agent(headers: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
"""Prepend "agent-framework" to the User-Agent in the headers.
|
||||
|
||||
When user agent telemetry is disabled through the ``AGENT_FRAMEWORK_USER_AGENT_DISABLED``
|
||||
environment variable, the User-Agent header will not include the agent-framework information.
|
||||
It will be sent back as is, or as an empty dict when None is passed.
|
||||
|
||||
Args:
|
||||
headers: The existing headers dictionary.
|
||||
|
||||
Returns:
|
||||
A new dict with "User-Agent" set to "agent-framework-python/{version}" if headers is None.
|
||||
The modified headers dictionary with "agent-framework-python/{version}" prepended to the User-Agent.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import prepend_agent_framework_to_user_agent
|
||||
|
||||
# Add agent-framework to new headers
|
||||
headers = prepend_agent_framework_to_user_agent()
|
||||
print(headers["User-Agent"]) # "agent-framework-python/0.1.0"
|
||||
|
||||
# Prepend to existing headers
|
||||
existing = {"User-Agent": "my-app/1.0"}
|
||||
headers = prepend_agent_framework_to_user_agent(existing)
|
||||
print(headers["User-Agent"]) # "agent-framework-python/0.1.0 my-app/1.0"
|
||||
"""
|
||||
if not IS_TELEMETRY_ENABLED:
|
||||
return headers or {}
|
||||
user_agent = get_user_agent()
|
||||
if not headers:
|
||||
return {USER_AGENT_KEY: user_agent}
|
||||
headers[USER_AGENT_KEY] = f"{user_agent} {headers[USER_AGENT_KEY]}" if USER_AGENT_KEY in headers else user_agent
|
||||
|
||||
return headers
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
@@ -0,0 +1,922 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
import uuid
|
||||
from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload
|
||||
|
||||
from .._agents import BaseAgent
|
||||
from .._sessions import (
|
||||
AgentSession,
|
||||
ContextProvider,
|
||||
HistoryProvider,
|
||||
InMemoryHistoryProvider,
|
||||
SessionContext,
|
||||
)
|
||||
from .._types import (
|
||||
AgentResponse,
|
||||
AgentResponseUpdate,
|
||||
AgentRunInputs,
|
||||
Content,
|
||||
Message,
|
||||
ResponseStream,
|
||||
UsageDetails,
|
||||
add_usage_details,
|
||||
)
|
||||
from ..exceptions import AgentException, AgentInvalidRequestException, AgentInvalidResponseException
|
||||
from ._checkpoint import CheckpointStorage
|
||||
from ._events import (
|
||||
AGENT_FORWARDED_EVENT_TYPES,
|
||||
WorkflowEvent,
|
||||
WorkflowRunState,
|
||||
)
|
||||
from ._message_utils import normalize_messages_input
|
||||
from ._typing_utils import is_instance_of, is_type_compatible
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypedDict # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypedDict # pragma: no cover
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._workflow import Workflow
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WorkflowAgent(BaseAgent):
|
||||
"""An `Agent` subclass that wraps a workflow and exposes it as an agent."""
|
||||
|
||||
# Class variable for the request info function name
|
||||
REQUEST_INFO_FUNCTION_NAME: ClassVar[str] = "request_info"
|
||||
|
||||
@dataclass
|
||||
class RequestInfoFunctionArgs:
|
||||
request_id: str
|
||||
request_event: WorkflowEvent
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"request_id": self.request_id, "request_event": self.request_event.to_dict()}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, payload: dict[str, Any]) -> WorkflowAgent.RequestInfoFunctionArgs:
|
||||
if "request_id" not in payload or "request_event" not in payload:
|
||||
raise ValueError(
|
||||
"Invalid payload for RequestInfoFunctionArgs. 'request_id' and 'request_event' are required."
|
||||
)
|
||||
if not payload["request_id"]:
|
||||
raise ValueError("request_id cannot be empty.")
|
||||
|
||||
return cls(
|
||||
request_id=payload.get("request_id", ""),
|
||||
request_event=WorkflowEvent.from_dict(payload.get("request_event", {})),
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workflow: Workflow,
|
||||
*,
|
||||
id: str | None = None,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
context_providers: Sequence[ContextProvider] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize the WorkflowAgent.
|
||||
|
||||
Args:
|
||||
workflow: The workflow to wrap as an agent.
|
||||
|
||||
Keyword Args:
|
||||
id: Unique identifier for the agent. If None, will be generated.
|
||||
name: Optional name for the agent.
|
||||
description: Optional description of the agent.
|
||||
context_providers: Optional sequence of context providers for the agent.
|
||||
**kwargs: Additional keyword arguments passed to BaseAgent.
|
||||
|
||||
Note:
|
||||
Only output events (type='output') and request_info events (type='request_info') from
|
||||
the workflow are considered and converted to agent responses of the WorkflowAgent.
|
||||
Other workflow events are ignored. Use `output_from` in WorkflowBuilder to control
|
||||
which executors' outputs are surfaced as agent responses.
|
||||
"""
|
||||
if id is None:
|
||||
id = f"WorkflowAgent_{uuid.uuid4().hex[:8]}"
|
||||
# Initialize with standard BaseAgent parameters first
|
||||
# Validate the workflow's start executor can handle agent-facing message inputs
|
||||
try:
|
||||
start_executor = workflow.get_start_executor()
|
||||
except KeyError as exc: # Defensive: workflow lacks a configured entry point
|
||||
raise ValueError("Workflow's start executor is not defined.") from exc
|
||||
|
||||
if not any(is_type_compatible(list[Message], input_type) for input_type in start_executor.input_types):
|
||||
raise ValueError("Workflow's start executor cannot handle list[Message]")
|
||||
|
||||
super().__init__(
|
||||
id=id,
|
||||
name=name,
|
||||
description=description,
|
||||
context_providers=context_providers,
|
||||
**kwargs,
|
||||
)
|
||||
self._workflow: Workflow = workflow
|
||||
|
||||
@property
|
||||
def workflow(self) -> Workflow:
|
||||
return self._workflow
|
||||
|
||||
# region Run Methods
|
||||
|
||||
@overload
|
||||
def run(
|
||||
self,
|
||||
messages: AgentRunInputs | None = None,
|
||||
*,
|
||||
stream: Literal[True],
|
||||
session: AgentSession | None = None,
|
||||
checkpoint_id: str | None = None,
|
||||
checkpoint_storage: CheckpointStorage | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ...
|
||||
|
||||
@overload
|
||||
async def run(
|
||||
self,
|
||||
messages: AgentRunInputs | None = None,
|
||||
*,
|
||||
stream: Literal[False] = ...,
|
||||
session: AgentSession | None = None,
|
||||
checkpoint_id: str | None = None,
|
||||
checkpoint_storage: CheckpointStorage | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
) -> AgentResponse: ...
|
||||
|
||||
def run(
|
||||
self,
|
||||
messages: AgentRunInputs | None = None,
|
||||
*,
|
||||
stream: bool = False,
|
||||
session: AgentSession | None = None,
|
||||
checkpoint_id: str | None = None,
|
||||
checkpoint_storage: CheckpointStorage | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
) -> ResponseStream[AgentResponseUpdate, AgentResponse] | Awaitable[AgentResponse]:
|
||||
"""Get a response from the workflow agent.
|
||||
|
||||
Args:
|
||||
messages: The message(s) to send to the workflow. Required for new runs,
|
||||
could be None if only restoring the underlying workflow from a checkpoint.
|
||||
|
||||
Keyword Args:
|
||||
stream: If True, returns an async iterable of updates. If False (default),
|
||||
returns an awaitable AgentResponse.
|
||||
session: The agent session for conversation context.
|
||||
checkpoint_id: ID of checkpoint to restore from. If provided, the workflow
|
||||
resumes from this checkpoint instead of starting fresh.
|
||||
checkpoint_storage: Runtime checkpoint storage. When provided with checkpoint_id,
|
||||
used to load and restore the checkpoint. When provided without checkpoint_id,
|
||||
enables checkpointing for this run.
|
||||
function_invocation_kwargs: Keyword arguments forwarded to tool invocations in
|
||||
subagents. Either a mapping of agent name/executor id to kwargs, or a flat
|
||||
mapping of kwargs for all tool invocations.
|
||||
client_kwargs: Keyword arguments forwarded to chat client calls in
|
||||
subagents. Either a mapping of agent name/executor id to kwargs, or a flat
|
||||
mapping of kwargs for all chat client calls.
|
||||
|
||||
Returns:
|
||||
When stream=True: An AsyncIterable[AgentResponseUpdate] for streaming updates.
|
||||
When stream=False: An Awaitable[AgentResponse] with the complete response.
|
||||
|
||||
Output events (type='output') from the workflow will be converted to ChatMessages
|
||||
or AgentResponseUpdate objects. Request info events (type='request_info') will be
|
||||
converted to function call and approval request contents.
|
||||
"""
|
||||
if messages is None:
|
||||
messages = []
|
||||
response_id = str(uuid.uuid4())
|
||||
if stream:
|
||||
return ResponseStream(
|
||||
self._run_stream_impl(
|
||||
messages,
|
||||
response_id,
|
||||
session,
|
||||
checkpoint_id,
|
||||
checkpoint_storage,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
),
|
||||
finalizer=AgentResponse.from_updates,
|
||||
)
|
||||
return self._run_impl(
|
||||
messages,
|
||||
response_id,
|
||||
session,
|
||||
checkpoint_id,
|
||||
checkpoint_storage,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
)
|
||||
|
||||
async def _run_impl(
|
||||
self,
|
||||
messages: AgentRunInputs,
|
||||
response_id: str,
|
||||
session: AgentSession | None,
|
||||
checkpoint_id: str | None = None,
|
||||
checkpoint_storage: CheckpointStorage | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
) -> AgentResponse:
|
||||
"""Internal implementation of non-streaming execution.
|
||||
|
||||
Args:
|
||||
messages: Normalized input messages to process.
|
||||
response_id: The unique response ID for this workflow execution.
|
||||
session: The agent session for conversation context.
|
||||
checkpoint_id: ID of checkpoint to restore from.
|
||||
checkpoint_storage: Runtime checkpoint storage.
|
||||
function_invocation_kwargs: Optional kwargs for tool invocations.
|
||||
client_kwargs: Optional kwargs for chat client calls.
|
||||
|
||||
Returns:
|
||||
An AgentResponse representing the workflow execution results.
|
||||
"""
|
||||
input_messages = normalize_messages_input(messages)
|
||||
|
||||
if (
|
||||
not any(
|
||||
provider.load_messages for provider in self.context_providers if isinstance(provider, HistoryProvider)
|
||||
)
|
||||
and session is not None
|
||||
):
|
||||
self.context_providers.append(InMemoryHistoryProvider())
|
||||
|
||||
provider_session = session
|
||||
if provider_session is None and self.context_providers:
|
||||
provider_session = AgentSession()
|
||||
|
||||
# run the context providers with the session
|
||||
session_context = SessionContext(
|
||||
session_id=provider_session.session_id if provider_session else None,
|
||||
service_session_id=provider_session.service_session_id if provider_session else None,
|
||||
input_messages=input_messages or [],
|
||||
options={},
|
||||
)
|
||||
for provider in self.context_providers:
|
||||
if isinstance(provider, HistoryProvider) and not provider.load_messages:
|
||||
continue
|
||||
if provider_session is None:
|
||||
raise RuntimeError("Provider session must be available when context providers are configured.")
|
||||
await provider.before_run(
|
||||
agent=self,
|
||||
session=provider_session,
|
||||
context=session_context,
|
||||
state=provider_session.state.setdefault(provider.source_id, {}),
|
||||
)
|
||||
# combine the messages
|
||||
session_messages: list[Message] = session_context.get_messages(include_input=True)
|
||||
|
||||
output_events: list[WorkflowEvent[Any]] = []
|
||||
async for event in self._run_core(
|
||||
session_messages,
|
||||
checkpoint_id,
|
||||
checkpoint_storage,
|
||||
streaming=False,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
):
|
||||
if event.type in AGENT_FORWARDED_EVENT_TYPES:
|
||||
output_events.append(event)
|
||||
|
||||
result = self._convert_workflow_events_to_agent_response(response_id, output_events)
|
||||
|
||||
# Set the response on the context so after_run providers (e.g. InMemoryHistoryProvider)
|
||||
# can persist the response messages alongside input messages.
|
||||
session_context._response = result # type: ignore[assignment]
|
||||
|
||||
await self._run_after_providers(session=provider_session, context=session_context)
|
||||
return result
|
||||
|
||||
async def _run_stream_impl(
|
||||
self,
|
||||
messages: AgentRunInputs,
|
||||
response_id: str,
|
||||
session: AgentSession | None,
|
||||
checkpoint_id: str | None = None,
|
||||
checkpoint_storage: CheckpointStorage | None = None,
|
||||
function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
) -> AsyncIterable[AgentResponseUpdate]:
|
||||
"""Internal implementation of streaming execution.
|
||||
|
||||
Args:
|
||||
messages: Input messages to process.
|
||||
response_id: The unique response ID for this workflow execution.
|
||||
session: The agent session for conversation context.
|
||||
checkpoint_id: ID of checkpoint to restore from.
|
||||
checkpoint_storage: Runtime checkpoint storage.
|
||||
function_invocation_kwargs: Optional kwargs for tool invocations.
|
||||
client_kwargs: Optional kwargs for chat client calls.
|
||||
|
||||
Yields:
|
||||
AgentResponseUpdate objects representing the workflow execution progress.
|
||||
"""
|
||||
input_messages = normalize_messages_input(messages)
|
||||
|
||||
if (
|
||||
not any(
|
||||
provider.load_messages for provider in self.context_providers if isinstance(provider, HistoryProvider)
|
||||
)
|
||||
and session is not None
|
||||
):
|
||||
self.context_providers.append(InMemoryHistoryProvider())
|
||||
|
||||
provider_session = session
|
||||
if provider_session is None and self.context_providers:
|
||||
provider_session = AgentSession()
|
||||
|
||||
# run the context providers with the session
|
||||
session_context = SessionContext(
|
||||
session_id=provider_session.session_id if provider_session else None,
|
||||
service_session_id=provider_session.service_session_id if provider_session else None,
|
||||
input_messages=input_messages or [],
|
||||
options={},
|
||||
)
|
||||
for provider in self.context_providers:
|
||||
if isinstance(provider, HistoryProvider) and not provider.load_messages:
|
||||
continue
|
||||
if provider_session is None:
|
||||
raise RuntimeError("Provider session must be available when context providers are configured.")
|
||||
await provider.before_run(
|
||||
agent=self,
|
||||
session=provider_session,
|
||||
context=session_context,
|
||||
state=provider_session.state.setdefault(provider.source_id, {}),
|
||||
)
|
||||
# combine the messages
|
||||
|
||||
session_messages: list[Message] = session_context.get_messages(include_input=True)
|
||||
all_updates: list[AgentResponseUpdate] = []
|
||||
async for event in self._run_core(
|
||||
session_messages,
|
||||
checkpoint_id,
|
||||
checkpoint_storage,
|
||||
streaming=True,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
):
|
||||
updates = self._convert_workflow_event_to_agent_response_updates(response_id, event)
|
||||
for update in updates:
|
||||
all_updates.append(update)
|
||||
yield update
|
||||
|
||||
# Build the final response from collected updates so after_run providers
|
||||
# (e.g. InMemoryHistoryProvider) can persist the response messages.
|
||||
if all_updates:
|
||||
session_context._response = AgentResponse.from_updates(all_updates) # type: ignore[assignment]
|
||||
|
||||
await self._run_after_providers(session=provider_session, context=session_context)
|
||||
|
||||
async def _run_core(
|
||||
self,
|
||||
input_messages: Sequence[Message],
|
||||
checkpoint_id: str | None,
|
||||
checkpoint_storage: CheckpointStorage | None,
|
||||
streaming: bool,
|
||||
function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None,
|
||||
) -> AsyncIterable[WorkflowEvent]:
|
||||
"""Core implementation that yields workflow events for both streaming and non-streaming modes.
|
||||
|
||||
Args:
|
||||
input_messages: Normalized input messages to process.
|
||||
checkpoint_id: ID of checkpoint to restore from.
|
||||
checkpoint_storage: Runtime checkpoint storage.
|
||||
streaming: Whether to use streaming workflow methods.
|
||||
function_invocation_kwargs: Optional kwargs for tool invocations.
|
||||
client_kwargs: Optional kwargs for chat client calls.
|
||||
|
||||
Yields:
|
||||
WorkflowEvent objects from the workflow execution.
|
||||
"""
|
||||
# Restore the workflow state if a checkpoint is provided
|
||||
if checkpoint_id is not None:
|
||||
if checkpoint_storage is None:
|
||||
raise AgentInvalidRequestException("checkpoint_storage must be provided when checkpoint_id is provided")
|
||||
logger.debug(f"Restoring workflow from checkpoint {checkpoint_id}")
|
||||
# Restore the workflow from checkpoint
|
||||
if streaming:
|
||||
async for _ in self.workflow.run(
|
||||
stream=True,
|
||||
checkpoint_id=checkpoint_id,
|
||||
checkpoint_storage=checkpoint_storage,
|
||||
):
|
||||
pass
|
||||
else:
|
||||
_ = await self.workflow.run(
|
||||
checkpoint_id=checkpoint_id,
|
||||
checkpoint_storage=checkpoint_storage,
|
||||
)
|
||||
if not input_messages:
|
||||
logger.info("No input messages provided; the workflow has been restored to the checkpoint state.")
|
||||
return
|
||||
|
||||
final_state = self._workflow.status
|
||||
logger.debug(f"Workflow state: {final_state}")
|
||||
|
||||
if final_state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
|
||||
# Extract function responses from input messages, and ensure that
|
||||
# only function responses are present in messages if there is any
|
||||
# pending request.
|
||||
# NOTE: It is possible that some pending requests are not fulfilled,
|
||||
# and we will let the workflow to handle this -- the agent does not
|
||||
# have an opinion on this.
|
||||
function_responses = self._extract_function_responses(input_messages)
|
||||
if streaming:
|
||||
async for event in self.workflow.run(
|
||||
responses=function_responses,
|
||||
stream=True,
|
||||
checkpoint_storage=checkpoint_storage,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
):
|
||||
yield event
|
||||
else:
|
||||
for event in await self.workflow.run(
|
||||
responses=function_responses,
|
||||
checkpoint_storage=checkpoint_storage,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
):
|
||||
yield event
|
||||
elif final_state == WorkflowRunState.IDLE:
|
||||
if streaming:
|
||||
async for event in self.workflow.run(
|
||||
message=input_messages,
|
||||
stream=True,
|
||||
checkpoint_storage=checkpoint_storage,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
):
|
||||
yield event
|
||||
else:
|
||||
for event in await self.workflow.run(
|
||||
message=input_messages,
|
||||
checkpoint_storage=checkpoint_storage,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
):
|
||||
yield event
|
||||
else:
|
||||
raise AgentException(f"The underlying workflow is in an invalid state to restart: {final_state}.")
|
||||
|
||||
# endregion Run Methods
|
||||
|
||||
def _convert_workflow_events_to_agent_response(
|
||||
self,
|
||||
response_id: str,
|
||||
output_events: list[WorkflowEvent[Any]],
|
||||
) -> AgentResponse:
|
||||
"""Convert a list of workflow events to an AgentResponse.
|
||||
|
||||
Caller-facing workflow events are forwarded as agent messages. Terminal and
|
||||
intermediate event payloads keep their original content types.
|
||||
"""
|
||||
messages: list[Message] = []
|
||||
raw_representations: list[object] = []
|
||||
merged_usage: UsageDetails | None = None
|
||||
latest_created_at: str | None = None
|
||||
|
||||
for output_event in output_events:
|
||||
if output_event.type == "request_info":
|
||||
request_content = self._process_request_info_event(output_event)
|
||||
messages.append(
|
||||
Message(
|
||||
contents=[request_content],
|
||||
role="assistant",
|
||||
author_name=output_event.source_executor_id,
|
||||
message_id=str(uuid.uuid4()),
|
||||
raw_representation=output_event,
|
||||
)
|
||||
)
|
||||
raw_representations.append(output_event)
|
||||
else:
|
||||
data = output_event.data
|
||||
# Anything that isn't `output` is intermediate — this branch only sees
|
||||
# events that already passed the lifecycle filter and weren't request_info.
|
||||
is_intermediate = output_event.type != "output"
|
||||
|
||||
if isinstance(data, AgentResponseUpdate):
|
||||
# AgentResponseUpdate is a streaming-only payload. Accepting it
|
||||
# in non-streaming runs would make message ordering depend on
|
||||
# partial chunks for both terminal and intermediate events.
|
||||
event_label = "Intermediate" if is_intermediate else "Output"
|
||||
raise AgentInvalidRequestException(
|
||||
f"{event_label} event with AgentResponseUpdate data cannot be emitted "
|
||||
"in non-streaming mode. Please ensure executors emit AgentResponse "
|
||||
"for non-streaming workflows."
|
||||
)
|
||||
|
||||
if isinstance(data, AgentResponse):
|
||||
messages.extend(data.messages)
|
||||
raw_representations.append(data.raw_representation)
|
||||
merged_usage = add_usage_details(merged_usage, data.usage_details)
|
||||
latest_created_at = (
|
||||
data.created_at
|
||||
if not latest_created_at
|
||||
else max(latest_created_at, data.created_at)
|
||||
if data.created_at
|
||||
else latest_created_at
|
||||
)
|
||||
elif isinstance(data, Message):
|
||||
messages.append(data)
|
||||
raw_representations.append(data.raw_representation)
|
||||
elif is_instance_of(data, list[Message]):
|
||||
chat_messages = cast(list[Message], data)
|
||||
messages.extend(chat_messages)
|
||||
raw_representations.append(data)
|
||||
else:
|
||||
contents = self._extract_contents(data)
|
||||
if not contents:
|
||||
continue
|
||||
|
||||
messages.append(
|
||||
Message(
|
||||
contents=contents,
|
||||
role="assistant",
|
||||
author_name=output_event.executor_id,
|
||||
message_id=str(uuid.uuid4()),
|
||||
raw_representation=data,
|
||||
)
|
||||
)
|
||||
raw_representations.append(data)
|
||||
|
||||
return AgentResponse(
|
||||
messages=messages,
|
||||
response_id=response_id,
|
||||
created_at=latest_created_at,
|
||||
usage_details=merged_usage,
|
||||
raw_representation=raw_representations,
|
||||
)
|
||||
|
||||
def _convert_workflow_event_to_agent_response_updates(
|
||||
self,
|
||||
response_id: str,
|
||||
event: WorkflowEvent[Any],
|
||||
) -> list[AgentResponseUpdate]:
|
||||
"""Convert a workflow event to a list of AgentResponseUpdate objects.
|
||||
|
||||
Forwarding rule:
|
||||
|
||||
- ``type='output'`` — terminal user-facing emission. Forwarded as-is.
|
||||
- ``type='intermediate'`` (and the deprecated ``type='data'``) — forwarded
|
||||
as-is.
|
||||
- ``type='request_info'`` — request-info translation (unchanged).
|
||||
- Everything else (lifecycle, diagnostics, executor bookkeeping,
|
||||
orchestration-internal events like ``group_chat``/``handoff_sent``/
|
||||
``magentic_orchestrator``) is dropped.
|
||||
"""
|
||||
# TODO(evmattso): https://github.com/microsoft/agent-framework/issues/5885
|
||||
if event.type not in AGENT_FORWARDED_EVENT_TYPES:
|
||||
return []
|
||||
|
||||
if event.type != "request_info":
|
||||
data = event.data
|
||||
executor_id = event.executor_id
|
||||
|
||||
if isinstance(data, AgentResponseUpdate):
|
||||
# Construct a fresh AgentResponseUpdate so we don't mutate a payload
|
||||
# that AgentExecutor still holds a reference to in its `updates` list.
|
||||
return [
|
||||
AgentResponseUpdate(
|
||||
contents=list(data.contents),
|
||||
role=data.role,
|
||||
author_name=data.author_name or executor_id,
|
||||
response_id=data.response_id,
|
||||
message_id=data.message_id,
|
||||
created_at=data.created_at,
|
||||
raw_representation=data.raw_representation,
|
||||
)
|
||||
]
|
||||
if isinstance(data, AgentResponse):
|
||||
# Convert each message in AgentResponse to an AgentResponseUpdate
|
||||
updates: list[AgentResponseUpdate] = []
|
||||
for msg in data.messages:
|
||||
updates.append(
|
||||
AgentResponseUpdate(
|
||||
contents=list(msg.contents),
|
||||
role=msg.role,
|
||||
author_name=msg.author_name or executor_id,
|
||||
response_id=data.response_id or response_id,
|
||||
message_id=msg.message_id or str(uuid.uuid4()),
|
||||
created_at=data.created_at
|
||||
or datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
|
||||
raw_representation=msg,
|
||||
)
|
||||
)
|
||||
return updates
|
||||
if isinstance(data, Message):
|
||||
return [
|
||||
AgentResponseUpdate(
|
||||
contents=list(data.contents),
|
||||
role=data.role,
|
||||
author_name=data.author_name or executor_id,
|
||||
response_id=response_id,
|
||||
message_id=str(uuid.uuid4()),
|
||||
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
|
||||
raw_representation=data,
|
||||
)
|
||||
]
|
||||
if is_instance_of(data, list[Message]):
|
||||
# Convert each Message to an AgentResponseUpdate
|
||||
chat_messages = cast(list[Message], data)
|
||||
updates = []
|
||||
for msg in chat_messages:
|
||||
updates.append(
|
||||
AgentResponseUpdate(
|
||||
contents=list(msg.contents),
|
||||
role=msg.role,
|
||||
author_name=msg.author_name or executor_id,
|
||||
response_id=response_id,
|
||||
message_id=msg.message_id or str(uuid.uuid4()),
|
||||
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
|
||||
raw_representation=msg,
|
||||
)
|
||||
)
|
||||
return updates
|
||||
contents = self._extract_contents(data)
|
||||
if not contents:
|
||||
return []
|
||||
return [
|
||||
AgentResponseUpdate(
|
||||
contents=contents,
|
||||
role="assistant",
|
||||
author_name=executor_id,
|
||||
response_id=response_id,
|
||||
message_id=str(uuid.uuid4()),
|
||||
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
|
||||
raw_representation=data,
|
||||
)
|
||||
]
|
||||
|
||||
if event.type == "request_info":
|
||||
request_content = self._process_request_info_event(event)
|
||||
return [
|
||||
AgentResponseUpdate(
|
||||
contents=[request_content],
|
||||
role="assistant",
|
||||
author_name=self.name,
|
||||
response_id=response_id,
|
||||
message_id=str(uuid.uuid4()),
|
||||
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
|
||||
raw_representation=event,
|
||||
)
|
||||
]
|
||||
|
||||
# Ignore workflow-internal events
|
||||
return []
|
||||
|
||||
def _process_request_info_event(
|
||||
self,
|
||||
event: WorkflowEvent[Any],
|
||||
) -> Content:
|
||||
"""Convert a request_info event to FunctionApprovalRequestContent.
|
||||
|
||||
Args:
|
||||
event: A WorkflowEvent with type='request_info'.
|
||||
|
||||
Returns:
|
||||
A content object representing the request info. The content can be a `function_approval_request`
|
||||
or a `function_call` depending on the structure of the event data.
|
||||
|
||||
Note:
|
||||
If the event data is already a FunctionApprovalRequestContent, it will be returned as-is.
|
||||
"""
|
||||
if isinstance(event.data, Content) and event.data.user_input_request:
|
||||
# Return the event data as-is if it's already a properly formed FunctionApprovalRequestContent
|
||||
return event.data
|
||||
|
||||
request_id = event.request_id
|
||||
args = self.RequestInfoFunctionArgs(request_id=request_id, request_event=event).to_dict()
|
||||
|
||||
return Content.from_function_call(
|
||||
call_id=request_id,
|
||||
name=self.REQUEST_INFO_FUNCTION_NAME,
|
||||
arguments=args,
|
||||
)
|
||||
|
||||
def _extract_function_responses(self, input_messages: Sequence[Message]) -> dict[str, Any]:
|
||||
"""Extract function responses from input messages.
|
||||
|
||||
The responses are for pending requests that the workflow is waiting on, and
|
||||
will be passed to the workflow. The pending requests are processed to either
|
||||
`function_approval_request` or `function_call` content by `_process_request_info_event`.
|
||||
"""
|
||||
function_responses: dict[str, Any] = {}
|
||||
for message in input_messages:
|
||||
for content in message.contents:
|
||||
if content.type == "function_approval_response":
|
||||
request_id: str = content.id # type: ignore[assignment]
|
||||
function_responses[request_id] = content
|
||||
elif content.type == "function_result":
|
||||
response_data = content.result if hasattr(content, "result") else str(content)
|
||||
function_responses[content.call_id] = response_data # type: ignore
|
||||
else:
|
||||
raise AgentInvalidResponseException(
|
||||
"Unexpected content type while awaiting request info responses."
|
||||
)
|
||||
|
||||
return function_responses
|
||||
|
||||
def _extract_contents(self, data: Any) -> list[Content]:
|
||||
"""Recursively extract Content from workflow output data."""
|
||||
if isinstance(data, list):
|
||||
return [c for item in data for c in self._extract_contents(item)] # type: ignore
|
||||
if isinstance(data, Content):
|
||||
return [data]
|
||||
if isinstance(data, str):
|
||||
return [Content.from_text(text=data)]
|
||||
return [Content.from_text(text=str(data))]
|
||||
|
||||
class _ResponseState(TypedDict):
|
||||
"""State for grouping response updates by message_id."""
|
||||
|
||||
by_msg: dict[str, list[AgentResponseUpdate]]
|
||||
dangling: list[AgentResponseUpdate]
|
||||
|
||||
@staticmethod
|
||||
def merge_updates(updates: list[AgentResponseUpdate], response_id: str) -> AgentResponse:
|
||||
"""Merge streaming updates into a single AgentResponse.
|
||||
|
||||
Behavior:
|
||||
- Group updates by response_id; within each response_id, group by message_id and keep a dangling bucket for
|
||||
updates without message_id.
|
||||
- Convert each group (per message and dangling) into an intermediate AgentResponse via
|
||||
AgentResponse.from_updates, then sort by created_at and merge.
|
||||
- Append messages from updates without any response_id at the end (global dangling), while aggregating metadata.
|
||||
|
||||
Args:
|
||||
updates: The list of AgentResponseUpdate objects to merge.
|
||||
response_id: The response identifier to set on the returned AgentResponse.
|
||||
|
||||
Returns:
|
||||
An AgentResponse with messages in processing order and aggregated metadata.
|
||||
"""
|
||||
# PHASE 1: GROUP UPDATES BY RESPONSE_ID AND MESSAGE_ID
|
||||
# First pass: build call_id -> response_id map from FunctionCallContent updates
|
||||
call_id_to_response_id: dict[str, str] = {}
|
||||
for u in updates:
|
||||
if u.response_id:
|
||||
for content in u.contents:
|
||||
if content.type == "function_call" and content.call_id:
|
||||
call_id_to_response_id[content.call_id] = u.response_id
|
||||
|
||||
# Second pass: group updates, associating FunctionResultContent with their calls
|
||||
states: dict[str, WorkflowAgent._ResponseState] = {}
|
||||
global_dangling: list[AgentResponseUpdate] = []
|
||||
|
||||
for u in updates:
|
||||
effective_response_id = u.response_id
|
||||
# If no response_id, check if this is a FunctionResultContent that matches a call
|
||||
if not effective_response_id:
|
||||
for content in u.contents:
|
||||
if content.type == "function_result" and content.call_id:
|
||||
effective_response_id = call_id_to_response_id.get(content.call_id)
|
||||
if effective_response_id:
|
||||
break
|
||||
|
||||
if effective_response_id:
|
||||
state = states.setdefault(effective_response_id, {"by_msg": {}, "dangling": []})
|
||||
by_msg = state["by_msg"]
|
||||
dangling = state["dangling"]
|
||||
if u.message_id:
|
||||
by_msg.setdefault(u.message_id, []).append(u)
|
||||
else:
|
||||
dangling.append(u)
|
||||
else:
|
||||
global_dangling.append(u)
|
||||
|
||||
# HELPER FUNCTIONS
|
||||
def _parse_dt(value: str | None) -> tuple[int, datetime | str | None]:
|
||||
if not value:
|
||||
return (1, None)
|
||||
v = value
|
||||
if v.endswith("Z"):
|
||||
v = v[:-1] + "+00:00"
|
||||
try:
|
||||
return (0, datetime.fromisoformat(v))
|
||||
except Exception:
|
||||
return (0, v)
|
||||
|
||||
def _merge_responses(current: AgentResponse | None, incoming: AgentResponse) -> AgentResponse:
|
||||
if current is None:
|
||||
return incoming
|
||||
raw_list: list[object] = []
|
||||
|
||||
def _add_raw(value: object) -> None:
|
||||
if isinstance(value, list):
|
||||
raw_list.extend(cast(list[object], value))
|
||||
else:
|
||||
raw_list.append(value)
|
||||
|
||||
if current.raw_representation is not None:
|
||||
_add_raw(current.raw_representation)
|
||||
if incoming.raw_representation is not None:
|
||||
_add_raw(incoming.raw_representation)
|
||||
return AgentResponse(
|
||||
messages=(current.messages or []) + (incoming.messages or []),
|
||||
response_id=current.response_id or incoming.response_id,
|
||||
created_at=incoming.created_at or current.created_at,
|
||||
usage_details=add_usage_details(current.usage_details, incoming.usage_details),
|
||||
raw_representation=raw_list if raw_list else None,
|
||||
additional_properties=incoming.additional_properties or current.additional_properties,
|
||||
)
|
||||
|
||||
# PHASE 2: CONVERT GROUPED UPDATES TO RESPONSES AND MERGE
|
||||
final_messages: list[Message] = []
|
||||
merged_usage: UsageDetails | None = None
|
||||
latest_created_at: str | None = None
|
||||
merged_additional_properties: dict[str, Any] | None = None
|
||||
raw_representations: list[object] = []
|
||||
|
||||
for grouped_response_id in states:
|
||||
state = states[grouped_response_id]
|
||||
by_msg = state["by_msg"]
|
||||
dangling = state["dangling"]
|
||||
|
||||
per_message_responses: list[AgentResponse] = []
|
||||
for _, msg_updates in by_msg.items():
|
||||
if msg_updates:
|
||||
per_message_responses.append(AgentResponse.from_updates(msg_updates))
|
||||
if dangling:
|
||||
per_message_responses.append(AgentResponse.from_updates(dangling))
|
||||
|
||||
per_message_responses.sort(key=lambda r: _parse_dt(r.created_at))
|
||||
|
||||
aggregated: AgentResponse | None = None
|
||||
for resp in per_message_responses:
|
||||
if resp.response_id and grouped_response_id and resp.response_id != grouped_response_id:
|
||||
resp.response_id = grouped_response_id
|
||||
aggregated = _merge_responses(aggregated, resp)
|
||||
|
||||
if aggregated:
|
||||
final_messages.extend(aggregated.messages)
|
||||
if aggregated.usage_details:
|
||||
merged_usage = add_usage_details(merged_usage, aggregated.usage_details)
|
||||
if aggregated.created_at and (
|
||||
not latest_created_at or _parse_dt(aggregated.created_at) > _parse_dt(latest_created_at)
|
||||
):
|
||||
latest_created_at = aggregated.created_at
|
||||
if aggregated.additional_properties:
|
||||
if merged_additional_properties is None:
|
||||
merged_additional_properties = {}
|
||||
merged_additional_properties.update(aggregated.additional_properties)
|
||||
raw_value = aggregated.raw_representation
|
||||
if raw_value:
|
||||
cast_value = cast(object | list[object], raw_value)
|
||||
if isinstance(cast_value, list):
|
||||
raw_representations.extend(cast(list[object], cast_value))
|
||||
else:
|
||||
raw_representations.append(cast_value)
|
||||
|
||||
# PHASE 3: HANDLE GLOBAL DANGLING UPDATES (NO RESPONSE_ID)
|
||||
# These are updates that couldn't be associated with any response_id
|
||||
# (e.g., orphan FunctionResultContent with no matching FunctionCallContent)
|
||||
if global_dangling:
|
||||
flattened = AgentResponse.from_updates(global_dangling)
|
||||
final_messages.extend(flattened.messages)
|
||||
if flattened.usage_details:
|
||||
merged_usage = add_usage_details(merged_usage, flattened.usage_details)
|
||||
if flattened.created_at and (
|
||||
not latest_created_at or _parse_dt(flattened.created_at) > _parse_dt(latest_created_at)
|
||||
):
|
||||
latest_created_at = flattened.created_at
|
||||
if flattened.additional_properties:
|
||||
if merged_additional_properties is None:
|
||||
merged_additional_properties = {}
|
||||
merged_additional_properties.update(flattened.additional_properties)
|
||||
flat_raw = flattened.raw_representation
|
||||
if flat_raw:
|
||||
cast_flat = cast(object | list[object], flat_raw)
|
||||
if isinstance(cast_flat, list):
|
||||
raw_representations.extend(cast(list[object], cast_flat))
|
||||
else:
|
||||
raw_representations.append(cast_flat)
|
||||
|
||||
# PHASE 4: CONSTRUCT FINAL RESPONSE WITH INPUT RESPONSE_ID
|
||||
return AgentResponse(
|
||||
messages=final_messages,
|
||||
response_id=response_id,
|
||||
created_at=latest_created_at,
|
||||
usage_details=merged_usage,
|
||||
raw_representation=raw_representations if raw_representations else None,
|
||||
additional_properties=merged_additional_properties,
|
||||
)
|
||||
@@ -0,0 +1,602 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
from typing_extensions import Never
|
||||
|
||||
from agent_framework import Content
|
||||
|
||||
from .._agents import SupportsAgentRun
|
||||
from .._sessions import AgentSession
|
||||
from .._types import AgentResponse, AgentResponseUpdate, Message, ResponseStream
|
||||
from ._agent_utils import resolve_agent_id
|
||||
from ._const import GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY
|
||||
from ._executor import Executor, handler
|
||||
from ._message_utils import normalize_messages_input
|
||||
from ._request_info_mixin import response_handler
|
||||
from ._typing_utils import is_chat_agent
|
||||
from ._workflow_context import WorkflowContext
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # pragma: no cover
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentExecutorRequest:
|
||||
"""A request to an agent executor.
|
||||
|
||||
Attributes:
|
||||
messages: A list of chat messages to be processed by the agent.
|
||||
should_respond: A flag indicating whether the agent should respond to the messages.
|
||||
If False, the messages will be saved to the executor's cache but not sent to the agent.
|
||||
"""
|
||||
|
||||
messages: list[Message]
|
||||
should_respond: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentExecutorResponse:
|
||||
"""A response from an agent executor.
|
||||
|
||||
Attributes:
|
||||
executor_id: The ID of the executor that generated the response.
|
||||
agent_response: The underlying agent run response (unaltered from client).
|
||||
full_conversation: The full conversation context (prior inputs + all assistant/tool outputs) that
|
||||
should be used when chaining to another AgentExecutor. This prevents downstream agents losing
|
||||
user prompts.
|
||||
"""
|
||||
|
||||
executor_id: str
|
||||
agent_response: AgentResponse
|
||||
full_conversation: list[Message]
|
||||
|
||||
def with_text(self, text: str) -> "AgentExecutorResponse":
|
||||
"""Create a new AgentExecutorResponse with replaced text, preserving the conversation history.
|
||||
|
||||
Use this in custom executors that transform agent output text (e.g. upper-casing, summarising)
|
||||
when you need downstream AgentExecutors to still have access to the full prior conversation.
|
||||
|
||||
Without this helper, sending a plain ``str`` from a custom executor breaks the context chain:
|
||||
the downstream ``AgentExecutor.from_str`` handler only adds that one string to its cache and
|
||||
loses all prior messages. By using ``with_text`` the response type stays
|
||||
``AgentExecutorResponse``, so ``AgentExecutor.from_response`` is invoked instead and the full
|
||||
conversation is preserved.
|
||||
|
||||
Args:
|
||||
text: The replacement assistant message text.
|
||||
|
||||
Returns:
|
||||
A new ``AgentExecutorResponse`` whose ``agent_response`` contains a single assistant
|
||||
message with ``text``, and whose ``full_conversation`` is the prior conversation
|
||||
(everything before the original agent turn) followed by the new assistant message.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import AgentExecutorResponse, WorkflowContext, executor
|
||||
|
||||
|
||||
@executor(
|
||||
id="upper_case_executor",
|
||||
input=AgentExecutorResponse,
|
||||
output=AgentExecutorResponse,
|
||||
workflow_output=str,
|
||||
)
|
||||
async def upper_case(
|
||||
response: AgentExecutorResponse,
|
||||
ctx: WorkflowContext[AgentExecutorResponse, str],
|
||||
) -> None:
|
||||
upper_text = response.agent_response.text.upper()
|
||||
await ctx.send_message(response.with_text(upper_text))
|
||||
await ctx.yield_output(upper_text)
|
||||
"""
|
||||
new_message = Message("assistant", [text])
|
||||
new_agent_response = AgentResponse(messages=[new_message])
|
||||
|
||||
# Strip off the original agent turn and replace with the new text.
|
||||
n_agent_messages = len(self.agent_response.messages)
|
||||
prior_messages = (
|
||||
self.full_conversation[:-n_agent_messages] if n_agent_messages else list(self.full_conversation)
|
||||
)
|
||||
new_full_conversation = [*prior_messages, new_message]
|
||||
|
||||
return AgentExecutorResponse(
|
||||
executor_id=self.executor_id,
|
||||
agent_response=new_agent_response,
|
||||
full_conversation=new_full_conversation,
|
||||
)
|
||||
|
||||
|
||||
class AgentExecutor(Executor):
|
||||
"""built-in executor that wraps an agent for handling messages.
|
||||
|
||||
AgentExecutor adapts its behavior based on the workflow execution mode:
|
||||
- run(stream=True): Emits incremental output events (type='output') as the agent produces tokens
|
||||
- run(): Emits a single output event (type='output') containing the complete response
|
||||
|
||||
Use `output_from` in WorkflowBuilder to control whether the AgentResponse
|
||||
or AgentResponseUpdate objects are yielded as workflow outputs.
|
||||
|
||||
Messages sent to downstream executors will always be the complete AgentResponse. In
|
||||
streaming mode, incremental AgentResponseUpdates will be concatenated to form the full
|
||||
response to be sent downstream.
|
||||
|
||||
The executor automatically detects the mode via WorkflowContext.is_streaming().
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: SupportsAgentRun,
|
||||
*,
|
||||
session: AgentSession | None = None,
|
||||
id: str | None = None,
|
||||
context_mode: Literal["full", "last_agent", "custom"] | None = None,
|
||||
context_filter: Callable[[list[Message]], list[Message]] | None = None,
|
||||
):
|
||||
"""Initialize the executor with a unique identifier.
|
||||
|
||||
Args:
|
||||
agent: The agent to be wrapped by this executor.
|
||||
session: The session to use for running the agent. If None, a new session will be created.
|
||||
id: A unique identifier for the executor. If None, the agent's name will be used if available.
|
||||
context_mode: Configuration for how the executor should manage conversation context upon
|
||||
receiving an AgentExecutorResponse as input. Options:
|
||||
- "full": append the full conversation (all prior messages + latest agent response) to the
|
||||
cache for the agent run. This is the default mode.
|
||||
- "last_agent": provide only the messages from the latest agent response as context for
|
||||
the agent run.
|
||||
- "custom": use the provided context_filter function to determine which messages to include
|
||||
as context for the agent run.
|
||||
context_filter: A function that takes the full conversation (list of Messages) as input and returns
|
||||
a filtered list of Messages to be used as context for the agent run. This is required
|
||||
if context_mode is set to "custom".
|
||||
"""
|
||||
# Prefer provided id; else use agent.name if present; else generate deterministic prefix
|
||||
exec_id = id or resolve_agent_id(agent)
|
||||
if not exec_id:
|
||||
raise ValueError("Agent must have a non-empty name or id or an explicit id must be provided.")
|
||||
super().__init__(exec_id)
|
||||
self._agent = agent
|
||||
self._session = session or self._agent.create_session()
|
||||
|
||||
self._pending_agent_requests: dict[str, Content] = {}
|
||||
self._pending_responses_to_agent: list[Content] = []
|
||||
|
||||
# AgentExecutor maintains an internal cache of messages in between runs
|
||||
self._cache: list[Message] = []
|
||||
# This tracks the full conversation after each run
|
||||
self._full_conversation: list[Message] = []
|
||||
|
||||
# Context mode validation
|
||||
self._context_mode = context_mode or "full"
|
||||
self._context_filter = context_filter
|
||||
if self._context_mode not in {"full", "last_agent", "custom"}:
|
||||
raise ValueError("context_mode must be one of 'full', 'last_agent', or 'custom'.")
|
||||
if self._context_mode == "custom" and not self._context_filter:
|
||||
raise ValueError("context_filter must be provided when context_mode is set to 'custom'.")
|
||||
|
||||
@property
|
||||
def agent(self) -> SupportsAgentRun:
|
||||
"""Get the underlying agent wrapped by this executor."""
|
||||
return self._agent
|
||||
|
||||
@property
|
||||
def description(self) -> str | None:
|
||||
"""Get the description of the underlying agent."""
|
||||
return self._agent.description
|
||||
|
||||
@handler
|
||||
async def run(
|
||||
self,
|
||||
request: AgentExecutorRequest,
|
||||
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
|
||||
) -> None:
|
||||
"""Handle an AgentExecutorRequest (canonical input).
|
||||
|
||||
This is the standard path: extend cache with provided messages; if should_respond
|
||||
run the agent and emit an AgentExecutorResponse downstream.
|
||||
"""
|
||||
self._cache.extend(request.messages)
|
||||
|
||||
if request.should_respond:
|
||||
await self._run_agent_and_emit(ctx)
|
||||
|
||||
@handler
|
||||
async def from_response(
|
||||
self,
|
||||
prior: AgentExecutorResponse,
|
||||
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
|
||||
) -> None:
|
||||
"""Enable seamless chaining: accept a prior AgentExecutorResponse as input.
|
||||
|
||||
Strategy: treat the prior response's messages as the conversation state and
|
||||
immediately run the agent to produce a new response.
|
||||
"""
|
||||
if self._context_mode == "full":
|
||||
self._cache.extend(prior.full_conversation)
|
||||
elif self._context_mode == "last_agent":
|
||||
self._cache.extend(prior.agent_response.messages)
|
||||
else:
|
||||
if not self._context_filter:
|
||||
# This should never happen due to validation in __init__, but mypy doesn't track that well
|
||||
raise ValueError("context_filter function must be provided for 'custom' context_mode.")
|
||||
self._cache.extend(self._context_filter(prior.full_conversation))
|
||||
|
||||
await self._run_agent_and_emit(ctx)
|
||||
|
||||
@handler
|
||||
async def from_str(
|
||||
self, text: str, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate]
|
||||
) -> None:
|
||||
"""Accept a raw user prompt string and run the agent.
|
||||
|
||||
The new string input will be added to the cache which is used as the conversation context for the agent run.
|
||||
|
||||
Warning:
|
||||
If the upstream executor received an ``AgentExecutorResponse`` but emits a plain
|
||||
``str``, this handler will be invoked instead of ``from_response``. This resets
|
||||
the conversation context because only the new string is added to the cache and
|
||||
all prior messages from the upstream agent are lost.
|
||||
|
||||
To preserve the full conversation when transforming agent output in a custom
|
||||
executor, use ``AgentExecutorResponse.with_text(...)`` so that the message type
|
||||
stays ``AgentExecutorResponse`` and ``from_response`` is called instead.
|
||||
"""
|
||||
if not self._cache and ctx.source_executor_ids != ["Workflow"]:
|
||||
logger.warning(
|
||||
"AgentExecutor '%s': from_str handler invoked with an empty cache. "
|
||||
"If you are chaining from an AgentExecutor, the upstream custom executor may be "
|
||||
"emitting a plain str instead of using AgentExecutorResponse.with_text(...), "
|
||||
"which causes the full conversation context to be lost.",
|
||||
self.id,
|
||||
)
|
||||
self._cache.extend(normalize_messages_input(text))
|
||||
await self._run_agent_and_emit(ctx)
|
||||
|
||||
@handler
|
||||
async def from_message(
|
||||
self,
|
||||
message: Message,
|
||||
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
|
||||
) -> None:
|
||||
"""Accept a single Message as input.
|
||||
|
||||
The new message will be added to the cache which is used as the conversation context for the agent run.
|
||||
"""
|
||||
self._cache.extend(normalize_messages_input(message))
|
||||
await self._run_agent_and_emit(ctx)
|
||||
|
||||
@handler
|
||||
async def from_messages(
|
||||
self,
|
||||
messages: list[str | Message],
|
||||
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
|
||||
) -> None:
|
||||
"""Accept a list of chat inputs (strings or Message) as conversation context.
|
||||
|
||||
The new messages will be added to the cache which is used as the conversation context for the agent run.
|
||||
"""
|
||||
self._cache.extend(normalize_messages_input(messages))
|
||||
await self._run_agent_and_emit(ctx)
|
||||
|
||||
@response_handler
|
||||
async def handle_user_input_response(
|
||||
self,
|
||||
original_request: Content,
|
||||
response: Content,
|
||||
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
|
||||
) -> None:
|
||||
"""Handle user input responses for function approvals during agent execution.
|
||||
|
||||
This will hold the executor's execution until all pending user input requests are resolved.
|
||||
|
||||
Args:
|
||||
original_request: The original function approval request sent by the agent.
|
||||
response: The user's response to the function approval request.
|
||||
ctx: The workflow context for emitting events and outputs.
|
||||
"""
|
||||
self._pending_responses_to_agent.append(response)
|
||||
self._pending_agent_requests.pop(original_request.id, None) # type: ignore[arg-type]
|
||||
|
||||
if not self._pending_agent_requests:
|
||||
# All pending requests have been resolved; resume agent execution.
|
||||
# Use role="tool" for function_result responses (from declaration-only tools)
|
||||
# so the LLM receives proper tool results instead of orphaned tool_calls.
|
||||
role = "tool" if all(r.type == "function_result" for r in self._pending_responses_to_agent) else "user"
|
||||
self._cache = normalize_messages_input(Message(role=role, contents=self._pending_responses_to_agent))
|
||||
self._pending_responses_to_agent.clear()
|
||||
await self._run_agent_and_emit(ctx)
|
||||
|
||||
@override
|
||||
async def on_checkpoint_save(self) -> dict[str, Any]:
|
||||
"""Capture current executor state for checkpointing.
|
||||
|
||||
NOTE: if the session uses service-side storage, the full session state
|
||||
may not be serialized locally.
|
||||
|
||||
Returns:
|
||||
Dict containing serialized cache and session state
|
||||
"""
|
||||
serialized_session = self._session.to_dict()
|
||||
|
||||
return {
|
||||
"cache": self._cache,
|
||||
"full_conversation": self._full_conversation,
|
||||
"agent_session": serialized_session,
|
||||
"pending_agent_requests": self._pending_agent_requests,
|
||||
"pending_responses_to_agent": self._pending_responses_to_agent,
|
||||
}
|
||||
|
||||
@override
|
||||
async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
|
||||
"""Restore executor state from checkpoint.
|
||||
|
||||
Args:
|
||||
state: Checkpoint data dict
|
||||
"""
|
||||
cache_payload = state.get("cache")
|
||||
self._cache = cache_payload or []
|
||||
|
||||
full_conversation_payload = state.get("full_conversation")
|
||||
self._full_conversation = full_conversation_payload or []
|
||||
|
||||
session_payload = state.get("agent_session")
|
||||
if session_payload:
|
||||
try:
|
||||
self._session = AgentSession.from_dict(session_payload)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to restore agent session: %s", exc)
|
||||
self._session = self._agent.create_session()
|
||||
else:
|
||||
self._session = self._agent.create_session()
|
||||
|
||||
pending_requests_payload = state.get("pending_agent_requests")
|
||||
self._pending_agent_requests = pending_requests_payload or {}
|
||||
|
||||
pending_responses_payload = state.get("pending_responses_to_agent")
|
||||
self._pending_responses_to_agent = pending_responses_payload or []
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the internal cache of the executor."""
|
||||
logger.debug("AgentExecutor %s: Resetting cache", self.id)
|
||||
self._cache.clear()
|
||||
|
||||
async def _run_agent_and_emit(
|
||||
self,
|
||||
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
|
||||
) -> None:
|
||||
"""Execute the underlying agent, emit events, and enqueue response.
|
||||
|
||||
Checks ctx.is_streaming() to determine whether to emit output events (type='output')
|
||||
containing incremental updates (streaming mode) or a single output event (type='output')
|
||||
containing the complete response (non-streaming mode).
|
||||
"""
|
||||
if ctx.is_streaming():
|
||||
# Streaming mode: emit incremental updates
|
||||
response = await self._run_agent_streaming(cast(WorkflowContext[Never, AgentResponseUpdate], ctx))
|
||||
else:
|
||||
# Non-streaming mode: use run() and emit single event
|
||||
response = await self._run_agent(cast(WorkflowContext[Never, AgentResponse], ctx))
|
||||
|
||||
# Snapshot current conversation as cache + latest agent outputs.
|
||||
# Do not append to prior snapshots: callers may provide full-history messages
|
||||
# in request.messages, and extending would duplicate prior turns.
|
||||
self._full_conversation = [*self._cache, *(list(response.messages) if response else [])]
|
||||
|
||||
if response is None:
|
||||
# Agent did not complete (e.g., waiting for user input); do not emit response
|
||||
logger.info("AgentExecutor %s: Agent did not complete, awaiting user input", self.id)
|
||||
return
|
||||
|
||||
agent_response = AgentExecutorResponse(self.id, response, full_conversation=self._full_conversation)
|
||||
await ctx.send_message(agent_response)
|
||||
self._cache.clear()
|
||||
|
||||
async def _run_agent(self, ctx: WorkflowContext[Never, AgentResponse]) -> AgentResponse | None:
|
||||
"""Execute the underlying agent in non-streaming mode.
|
||||
|
||||
Args:
|
||||
ctx: The workflow context for emitting events.
|
||||
|
||||
Returns:
|
||||
The complete AgentResponse, or None if waiting for user input.
|
||||
"""
|
||||
function_invocation_kwargs, client_kwargs = self._prepare_agent_run_args(
|
||||
ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {})
|
||||
)
|
||||
|
||||
if not self._cache:
|
||||
logger.warning(
|
||||
"AgentExecutor %s: Running agent with empty message cache. "
|
||||
"This could lead to service error for some LLM providers.",
|
||||
self.id,
|
||||
)
|
||||
|
||||
run_agent = cast(Callable[..., Awaitable[AgentResponse[Any]]], self._agent.run)
|
||||
response = await run_agent(
|
||||
self._cache,
|
||||
stream=False,
|
||||
session=self._session,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
)
|
||||
|
||||
# Handle any user input requests
|
||||
if response.user_input_requests:
|
||||
user_input_request_count = len(response.user_input_requests)
|
||||
total_message_content_count = sum(len(msg.contents) for msg in response.messages)
|
||||
if user_input_request_count != total_message_content_count:
|
||||
logger.warning(
|
||||
"Response %s contains %d user input requests but total message contents are %d. "
|
||||
"This indicates the response contains both user input requests and message contents. "
|
||||
"Double check if this is the intended behavior, as non user input request contents in "
|
||||
"this response will not be emitted.",
|
||||
response.response_id,
|
||||
user_input_request_count,
|
||||
total_message_content_count,
|
||||
)
|
||||
for user_input_request in response.user_input_requests:
|
||||
self._pending_agent_requests[user_input_request.id] = user_input_request # type: ignore[index]
|
||||
await ctx.request_info(user_input_request, Content, request_id=user_input_request.id)
|
||||
return None
|
||||
|
||||
# Only yield output if the response is complete and not waiting for user input.
|
||||
# This is to avoid emitting two events of different types ('output' and 'request_info')
|
||||
# that carry the same payload.
|
||||
await ctx.yield_output(response)
|
||||
return response
|
||||
|
||||
async def _run_agent_streaming(self, ctx: WorkflowContext[Never, AgentResponseUpdate]) -> AgentResponse | None:
|
||||
"""Execute the underlying agent in streaming mode and collect the full response.
|
||||
|
||||
Args:
|
||||
ctx: The workflow context for emitting events.
|
||||
|
||||
Returns:
|
||||
The complete AgentResponse, or None if waiting for user input.
|
||||
"""
|
||||
function_invocation_kwargs, client_kwargs = self._prepare_agent_run_args(
|
||||
ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {})
|
||||
)
|
||||
|
||||
if not self._cache:
|
||||
logger.warning(
|
||||
"AgentExecutor %s: Running agent with empty message cache. "
|
||||
"This could lead to service error for some LLM providers.",
|
||||
self.id,
|
||||
)
|
||||
|
||||
updates: list[AgentResponseUpdate] = []
|
||||
streamed_user_input_requests: list[Content] = []
|
||||
run_agent_stream = cast(Callable[..., ResponseStream[AgentResponseUpdate, AgentResponse[Any]]], self._agent.run)
|
||||
stream = run_agent_stream(
|
||||
self._cache,
|
||||
stream=True,
|
||||
session=self._session,
|
||||
function_invocation_kwargs=function_invocation_kwargs,
|
||||
client_kwargs=client_kwargs,
|
||||
)
|
||||
async for update in stream:
|
||||
updates.append(update)
|
||||
if update.user_input_requests:
|
||||
user_input_request_count = len(update.user_input_requests)
|
||||
total_message_content_count = len(update.contents)
|
||||
if user_input_request_count != total_message_content_count:
|
||||
logger.warning(
|
||||
"Response update %s contains %d user input requests but total message contents are %d. "
|
||||
"This indicates the response update contains both user input requests and message contents. "
|
||||
"Double check if this is the intended behavior, as non user input request contents will "
|
||||
"not be emitted.",
|
||||
update.response_id,
|
||||
user_input_request_count,
|
||||
total_message_content_count,
|
||||
)
|
||||
streamed_user_input_requests.extend(update.user_input_requests)
|
||||
else:
|
||||
# Only yield output events for updates that do not contain user input requests.
|
||||
# This is to avoid emitting two events of different types ('output' and 'request_info')
|
||||
# that carry the same payload.
|
||||
await ctx.yield_output(update)
|
||||
|
||||
# Prefer stream finalization when available so result hooks run
|
||||
# (e.g., thread conversation updates). Fall back to reconstructing from updates
|
||||
# for compatibility/custom agents that return a plain async iterable.
|
||||
# TODO(evmattso): Integrate workflow agent run handling around ResponseStream so
|
||||
# AgentExecutor does not need this conditional stream-finalization branch.
|
||||
maybe_get_final_response = getattr(stream, "get_final_response", None)
|
||||
get_final_response = maybe_get_final_response if callable(maybe_get_final_response) else None
|
||||
response: AgentResponse[Any]
|
||||
if get_final_response is not None:
|
||||
response = await cast(Callable[[], Awaitable[AgentResponse[Any]]], get_final_response)()
|
||||
elif is_chat_agent(self._agent):
|
||||
response_format = self._agent.default_options.get("response_format")
|
||||
response = AgentResponse.from_updates(
|
||||
updates,
|
||||
output_format_type=response_format,
|
||||
)
|
||||
else:
|
||||
response = AgentResponse.from_updates(updates)
|
||||
|
||||
# Handle any user input requests after the streaming completes
|
||||
user_input_requests: list[Content] = []
|
||||
seen_request_ids: set[str] = set()
|
||||
for user_input_request in [*streamed_user_input_requests, *response.user_input_requests]:
|
||||
request_id = getattr(user_input_request, "id", None)
|
||||
if isinstance(request_id, str) and request_id:
|
||||
if request_id in seen_request_ids:
|
||||
continue
|
||||
seen_request_ids.add(request_id)
|
||||
user_input_requests.append(user_input_request)
|
||||
|
||||
if user_input_requests:
|
||||
for user_input_request in user_input_requests:
|
||||
self._pending_agent_requests[user_input_request.id] = user_input_request # type: ignore[index]
|
||||
await ctx.request_info(user_input_request, Content, request_id=user_input_request.id)
|
||||
return None
|
||||
|
||||
return response
|
||||
|
||||
def _prepare_agent_run_args(
|
||||
self,
|
||||
raw_run_kwargs: dict[str, Any],
|
||||
) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
|
||||
"""Prepare function_invocation_kwargs and client_kwargs for agent.run().
|
||||
|
||||
Extracts ``function_invocation_kwargs`` and ``client_kwargs`` from the
|
||||
workflow state dict, resolving per-executor entries using ``self.id``. The
|
||||
``__global__`` sentinel key (set by ``Workflow._resolve_invocation_kwargs``) denotes
|
||||
global kwargs that apply to all executors. Per-executor dicts use executor IDs as
|
||||
keys; this executor extracts only its own entry.
|
||||
|
||||
Returns:
|
||||
A 2-tuple of (function_invocation_kwargs, client_kwargs).
|
||||
"""
|
||||
fi_resolved = raw_run_kwargs.get("function_invocation_kwargs")
|
||||
ci_resolved = raw_run_kwargs.get("client_kwargs")
|
||||
|
||||
function_invocation_kwargs = self._resolve_executor_kwargs(fi_resolved)
|
||||
client_kwargs = self._resolve_executor_kwargs(ci_resolved)
|
||||
|
||||
return function_invocation_kwargs, client_kwargs
|
||||
|
||||
def _resolve_executor_kwargs(self, resolved: dict[str, Any] | None) -> dict[str, Any] | None:
|
||||
"""Extract this executor's kwargs from a resolved invocation kwargs dict.
|
||||
|
||||
Args:
|
||||
resolved: The resolved dict produced by ``Workflow._resolve_invocation_kwargs``,
|
||||
containing either a ``__global__`` key (global kwargs) or executor-ID keys
|
||||
(per-executor kwargs). May also be ``None``.
|
||||
|
||||
Returns:
|
||||
The kwargs for this executor, or ``None`` if not applicable.
|
||||
"""
|
||||
if not isinstance(resolved, dict):
|
||||
return None
|
||||
# Use explicit key-presence checks so that an empty per-executor dict is
|
||||
# honoured (e.g. to clear kwargs) instead of falling through to global.
|
||||
if self.id in resolved:
|
||||
executor_kwargs = resolved[self.id]
|
||||
elif GLOBAL_KWARGS_KEY in resolved:
|
||||
executor_kwargs = resolved[GLOBAL_KWARGS_KEY]
|
||||
else:
|
||||
return None
|
||||
|
||||
if not isinstance(executor_kwargs, dict):
|
||||
logger.warning(
|
||||
"Executor %s expected a dict for its kwargs, but got %s. Ignoring.",
|
||||
self.id,
|
||||
type(executor_kwargs), # type: ignore
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
return executor_kwargs # type: ignore
|
||||
@@ -0,0 +1,17 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .._agents import SupportsAgentRun
|
||||
|
||||
|
||||
def resolve_agent_id(agent: SupportsAgentRun) -> str:
|
||||
"""Resolve the unique identifier for an agent.
|
||||
|
||||
Prefers the `.name` attribute if set; otherwise falls back to `.id`.
|
||||
|
||||
Args:
|
||||
agent: The agent whose identifier is to be resolved.
|
||||
|
||||
Returns:
|
||||
The resolved unique identifier for the agent.
|
||||
"""
|
||||
return agent.name if agent.name else agent.id
|
||||
@@ -0,0 +1,450 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field, fields
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Protocol, TypeAlias
|
||||
|
||||
from ..exceptions import WorkflowCheckpointException
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._events import WorkflowEvent
|
||||
from ._runner_context import WorkflowMessage
|
||||
|
||||
# Type alias for checkpoint IDs in case we want to change the
|
||||
# underlying type in the future (e.g., to UUID or a custom class)
|
||||
CheckpointID: TypeAlias = str
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class WorkflowCheckpoint:
|
||||
"""Represents a complete checkpoint of workflow state.
|
||||
|
||||
Checkpoints capture the full execution state of a workflow at a specific point,
|
||||
enabling workflows to be paused and resumed.
|
||||
|
||||
Note that a checkpoint is not tied to a specific workflow instance, but rather to
|
||||
a workflow definition (identified by workflow_name and graph_signature_hash). Thus,
|
||||
the ID of the workflow instance that created the checkpoint is not included in the
|
||||
checkpoint data. This allows checkpoints to be shared and restored across different
|
||||
workflow instances of the same workflow definition.
|
||||
|
||||
Attributes:
|
||||
workflow_name: Name of the workflow this checkpoint belongs to. This acts as a
|
||||
logical grouping for checkpoints and can be used to filter checkpoints by
|
||||
workflow. Workflows with the same name are expected to have compatible graph
|
||||
structures for checkpointing.
|
||||
graph_signature_hash: Hash of the workflow graph topology to validate checkpoint
|
||||
compatibility during restore
|
||||
checkpoint_id: Unique identifier for this checkpoint
|
||||
previous_checkpoint_id: ID of the previous checkpoint in the chain, if any. This
|
||||
allows chaining checkpoints together to form a history of workflow states.
|
||||
timestamp: ISO 8601 timestamp when checkpoint was created
|
||||
messages: Messages exchanged between executors
|
||||
state: Committed workflow state including user data and executor states.
|
||||
This contains only committed state; pending state changes are not
|
||||
included in checkpoints. Executor states are stored under the
|
||||
reserved key '_executor_state'.
|
||||
pending_request_info_events: Any pending request info events that have not
|
||||
yet been processed at the time of checkpointing. This allows the workflow
|
||||
to resume with the correct pending events after a restore.
|
||||
iteration_count: Current iteration number when checkpoint was created
|
||||
metadata: Additional metadata (e.g., superstep info, graph signature)
|
||||
version: Checkpoint format version
|
||||
|
||||
Note:
|
||||
The state dict may contain reserved keys managed by the framework.
|
||||
See State class documentation for details on reserved keys.
|
||||
"""
|
||||
|
||||
workflow_name: str
|
||||
graph_signature_hash: str
|
||||
|
||||
checkpoint_id: CheckpointID = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
previous_checkpoint_id: CheckpointID | None = None
|
||||
timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
|
||||
|
||||
# Core workflow state
|
||||
messages: dict[str, list[WorkflowMessage]] = field(default_factory=dict) # type: ignore[misc]
|
||||
state: dict[str, Any] = field(default_factory=dict) # type: ignore[misc]
|
||||
pending_request_info_events: dict[str, WorkflowEvent[Any]] = field(default_factory=dict) # type: ignore[misc]
|
||||
|
||||
# Runtime state
|
||||
iteration_count: int = 0
|
||||
|
||||
# Metadata
|
||||
metadata: dict[str, Any] = field(default_factory=dict) # type: ignore[misc]
|
||||
version: str = "1.0"
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert the WorkflowCheckpoint to a dictionary.
|
||||
|
||||
Notes:
|
||||
1. This method does not recursively convert nested dataclasses to dicts.
|
||||
2. This is a shallow conversion. The resulting dict will contain the same
|
||||
references to nested objects as the original dataclass.
|
||||
"""
|
||||
return {f.name: getattr(self, f.name) for f in fields(self)}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Mapping[str, Any]) -> WorkflowCheckpoint:
|
||||
"""Create a WorkflowCheckpoint from a dictionary.
|
||||
|
||||
Args:
|
||||
data: Dictionary containing checkpoint fields.
|
||||
|
||||
Returns:
|
||||
A new WorkflowCheckpoint instance.
|
||||
|
||||
Raises:
|
||||
WorkflowCheckpointException: If required fields are missing.
|
||||
"""
|
||||
try:
|
||||
return cls(**data)
|
||||
except Exception as ex:
|
||||
raise WorkflowCheckpointException(f"Failed to create WorkflowCheckpoint from dict: {ex}") from ex
|
||||
|
||||
|
||||
class CheckpointStorage(Protocol):
|
||||
"""Protocol for checkpoint storage backends."""
|
||||
|
||||
async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
|
||||
"""Save a checkpoint and return its ID.
|
||||
|
||||
Args:
|
||||
checkpoint: The WorkflowCheckpoint object to save.
|
||||
|
||||
Returns:
|
||||
The unique ID of the saved checkpoint.
|
||||
"""
|
||||
...
|
||||
|
||||
async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
|
||||
"""Load a checkpoint by ID.
|
||||
|
||||
Args:
|
||||
checkpoint_id: The unique ID of the checkpoint to load.
|
||||
|
||||
Returns:
|
||||
The WorkflowCheckpoint object corresponding to the given ID.
|
||||
|
||||
Raises:
|
||||
WorkflowCheckpointException: If no checkpoint with the given ID exists.
|
||||
"""
|
||||
...
|
||||
|
||||
async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]:
|
||||
"""List checkpoint objects for a given workflow name.
|
||||
|
||||
Args:
|
||||
workflow_name: The name of the workflow to list checkpoints for.
|
||||
|
||||
Returns:
|
||||
A list of WorkflowCheckpoint objects for the specified workflow name.
|
||||
"""
|
||||
...
|
||||
|
||||
async def delete(self, checkpoint_id: CheckpointID) -> bool:
|
||||
"""Delete a checkpoint by ID.
|
||||
|
||||
Args:
|
||||
checkpoint_id: The unique ID of the checkpoint to delete.
|
||||
|
||||
Returns:
|
||||
True if the checkpoint was successfully deleted, False if no checkpoint with the given ID exists.
|
||||
"""
|
||||
...
|
||||
|
||||
async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None:
|
||||
"""Get the latest checkpoint for a given workflow name.
|
||||
|
||||
Args:
|
||||
workflow_name: The name of the workflow to get the latest checkpoint for.
|
||||
|
||||
Returns:
|
||||
The latest WorkflowCheckpoint object for the specified workflow name, or None if no checkpoints exist.
|
||||
"""
|
||||
...
|
||||
|
||||
async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]:
|
||||
"""List checkpoint IDs for a given workflow name.
|
||||
|
||||
Args:
|
||||
workflow_name: The name of the workflow to list checkpoint IDs for.
|
||||
|
||||
Returns:
|
||||
A list of checkpoint IDs for the specified workflow name.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class InMemoryCheckpointStorage:
|
||||
"""In-memory checkpoint storage for testing and development."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize the memory storage."""
|
||||
self._checkpoints: dict[CheckpointID, WorkflowCheckpoint] = {}
|
||||
|
||||
async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
|
||||
"""Save a checkpoint and return its ID."""
|
||||
self._checkpoints[checkpoint.checkpoint_id] = copy.deepcopy(checkpoint)
|
||||
logger.debug(f"Saved checkpoint {checkpoint.checkpoint_id} to memory")
|
||||
return checkpoint.checkpoint_id
|
||||
|
||||
async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
|
||||
"""Load a checkpoint by ID."""
|
||||
checkpoint = self._checkpoints.get(checkpoint_id)
|
||||
if checkpoint:
|
||||
logger.debug(f"Loaded checkpoint {checkpoint_id} from memory")
|
||||
return checkpoint
|
||||
raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}")
|
||||
|
||||
async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]:
|
||||
"""List checkpoint objects for a given workflow name."""
|
||||
return [cp for cp in self._checkpoints.values() if cp.workflow_name == workflow_name]
|
||||
|
||||
async def delete(self, checkpoint_id: CheckpointID) -> bool:
|
||||
"""Delete a checkpoint by ID."""
|
||||
if checkpoint_id in self._checkpoints:
|
||||
del self._checkpoints[checkpoint_id]
|
||||
logger.debug(f"Deleted checkpoint {checkpoint_id} from memory")
|
||||
return True
|
||||
return False
|
||||
|
||||
async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None:
|
||||
"""Get the latest checkpoint for a given workflow name."""
|
||||
checkpoints = [cp for cp in self._checkpoints.values() if cp.workflow_name == workflow_name]
|
||||
if not checkpoints:
|
||||
return None
|
||||
latest_checkpoint = max(checkpoints, key=lambda cp: datetime.fromisoformat(cp.timestamp))
|
||||
logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint.checkpoint_id}")
|
||||
return latest_checkpoint
|
||||
|
||||
async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]:
|
||||
"""List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
|
||||
return [cp.checkpoint_id for cp in self._checkpoints.values() if cp.workflow_name == workflow_name]
|
||||
|
||||
|
||||
class FileCheckpointStorage:
|
||||
"""File-based checkpoint storage for persistence.
|
||||
|
||||
This storage implements a hybrid approach where the checkpoint metadata and structure are
|
||||
stored in JSON format, while the actual state data (which may contain complex Python objects)
|
||||
is serialized using pickle and embedded as base64-encoded strings within the JSON. This allows
|
||||
for human-readable checkpoint files while preserving the ability to store complex Python objects.
|
||||
|
||||
By default, checkpoint deserialization is restricted to a built-in set of safe Python types
|
||||
(primitives, datetime, uuid, ...), all ``agent_framework`` internal types, and OpenAI SDK types
|
||||
(``openai.types``). To allow additional application-specific types, pass them via the
|
||||
``allowed_checkpoint_types`` parameter using ``"module:qualname"`` format.
|
||||
|
||||
Example::
|
||||
|
||||
storage = FileCheckpointStorage(
|
||||
"/tmp/checkpoints",
|
||||
allowed_checkpoint_types=[
|
||||
"my_app.models:MyState",
|
||||
],
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
storage_path: str | Path,
|
||||
*,
|
||||
allowed_checkpoint_types: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Initialize the file storage.
|
||||
|
||||
Args:
|
||||
storage_path: Directory path where checkpoint files will be stored.
|
||||
allowed_checkpoint_types: Additional types (beyond the built-in safe set
|
||||
and framework types) that are permitted during checkpoint
|
||||
deserialization. Each entry should be a ``"module:qualname"``
|
||||
string (e.g., ``"my_app.models:MyState"``).
|
||||
"""
|
||||
self.storage_path = Path(storage_path)
|
||||
self.storage_path.mkdir(parents=True, exist_ok=True)
|
||||
self._allowed_types: frozenset[str] = frozenset(allowed_checkpoint_types or [])
|
||||
logger.info(f"Initialized file checkpoint storage at {self.storage_path}")
|
||||
|
||||
def _validate_file_path(self, checkpoint_id: CheckpointID) -> Path:
|
||||
"""Validate that a checkpoint ID resolves to a path within the storage directory.
|
||||
|
||||
This can prevent someone from crafting a checkpoint ID that points to an arbitrary
|
||||
file on the filesystem.
|
||||
|
||||
Args:
|
||||
checkpoint_id: The checkpoint ID to validate.
|
||||
|
||||
Returns:
|
||||
The validated file path.
|
||||
|
||||
Raises:
|
||||
WorkflowCheckpointException: If the checkpoint ID would resolve outside the storage directory.
|
||||
"""
|
||||
file_path = (self.storage_path / f"{checkpoint_id}.json").resolve()
|
||||
if not file_path.is_relative_to(self.storage_path.resolve()):
|
||||
raise WorkflowCheckpointException(f"Invalid checkpoint ID: {checkpoint_id}")
|
||||
return file_path
|
||||
|
||||
async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
|
||||
"""Save a checkpoint and return its ID.
|
||||
|
||||
Args:
|
||||
checkpoint: The WorkflowCheckpoint object to save.
|
||||
|
||||
Returns:
|
||||
The unique ID of the saved checkpoint.
|
||||
"""
|
||||
from ._checkpoint_encoding import encode_checkpoint_value
|
||||
|
||||
file_path = self._validate_file_path(checkpoint.checkpoint_id)
|
||||
checkpoint_dict = checkpoint.to_dict()
|
||||
encoded_checkpoint = encode_checkpoint_value(checkpoint_dict)
|
||||
|
||||
def _write_atomic() -> None:
|
||||
tmp_path = file_path.with_suffix(".json.tmp")
|
||||
with open(tmp_path, "w") as f:
|
||||
json.dump(encoded_checkpoint, f, indent=2, ensure_ascii=False)
|
||||
os.replace(tmp_path, file_path)
|
||||
|
||||
await asyncio.to_thread(_write_atomic)
|
||||
|
||||
logger.info(f"Saved checkpoint {checkpoint.checkpoint_id} to {file_path}")
|
||||
return checkpoint.checkpoint_id
|
||||
|
||||
async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
|
||||
"""Load a checkpoint by ID.
|
||||
|
||||
Args:
|
||||
checkpoint_id: The unique ID of the checkpoint to load.
|
||||
|
||||
Returns:
|
||||
The WorkflowCheckpoint object corresponding to the given ID.
|
||||
|
||||
Raises:
|
||||
WorkflowCheckpointException: If no checkpoint with the given ID exists,
|
||||
or if checkpoint decoding fails.
|
||||
"""
|
||||
file_path = self._validate_file_path(checkpoint_id)
|
||||
|
||||
if not file_path.exists():
|
||||
raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}")
|
||||
|
||||
def _read() -> dict[str, Any]:
|
||||
with open(file_path) as f:
|
||||
return json.load(f)
|
||||
|
||||
encoded_checkpoint = await asyncio.to_thread(_read)
|
||||
|
||||
from ._checkpoint_encoding import decode_checkpoint_value
|
||||
|
||||
try:
|
||||
decoded_checkpoint_dict = decode_checkpoint_value(encoded_checkpoint, allowed_types=self._allowed_types)
|
||||
except WorkflowCheckpointException:
|
||||
raise
|
||||
checkpoint = WorkflowCheckpoint.from_dict(decoded_checkpoint_dict)
|
||||
logger.info(f"Loaded checkpoint {checkpoint_id} from {file_path}")
|
||||
return checkpoint
|
||||
|
||||
async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]:
|
||||
"""List checkpoint objects for a given workflow name.
|
||||
|
||||
Args:
|
||||
workflow_name: The name of the workflow to list checkpoints for.
|
||||
|
||||
Returns:
|
||||
A list of WorkflowCheckpoint objects for the specified workflow name.
|
||||
"""
|
||||
|
||||
def _list_checkpoints() -> list[WorkflowCheckpoint]:
|
||||
checkpoints: list[WorkflowCheckpoint] = []
|
||||
for file_path in self.storage_path.glob("*.json"):
|
||||
try:
|
||||
with open(file_path) as f:
|
||||
encoded_checkpoint = json.load(f)
|
||||
from ._checkpoint_encoding import decode_checkpoint_value
|
||||
|
||||
decoded_checkpoint_dict = decode_checkpoint_value(
|
||||
encoded_checkpoint, allowed_types=self._allowed_types
|
||||
)
|
||||
checkpoint = WorkflowCheckpoint.from_dict(decoded_checkpoint_dict)
|
||||
if checkpoint.workflow_name == workflow_name:
|
||||
checkpoints.append(checkpoint)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to read checkpoint file {file_path}: {e}")
|
||||
return checkpoints
|
||||
|
||||
return await asyncio.to_thread(_list_checkpoints)
|
||||
|
||||
async def delete(self, checkpoint_id: CheckpointID) -> bool:
|
||||
"""Delete a checkpoint by ID.
|
||||
|
||||
Args:
|
||||
checkpoint_id: The unique ID of the checkpoint to delete.
|
||||
|
||||
Returns:
|
||||
True if the checkpoint was successfully deleted, False if no checkpoint with the given ID exists.
|
||||
"""
|
||||
file_path = self._validate_file_path(checkpoint_id)
|
||||
|
||||
def _delete() -> bool:
|
||||
if file_path.exists():
|
||||
file_path.unlink()
|
||||
logger.info(f"Deleted checkpoint {checkpoint_id} from {file_path}")
|
||||
return True
|
||||
return False
|
||||
|
||||
return await asyncio.to_thread(_delete)
|
||||
|
||||
async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None:
|
||||
"""Get the latest checkpoint for a given workflow name.
|
||||
|
||||
Args:
|
||||
workflow_name: The name of the workflow to get the latest checkpoint for.
|
||||
|
||||
Returns:
|
||||
The latest WorkflowCheckpoint object for the specified workflow name, or None if no checkpoints exist.
|
||||
"""
|
||||
checkpoints = await self.list_checkpoints(workflow_name=workflow_name)
|
||||
if not checkpoints:
|
||||
return None
|
||||
latest_checkpoint = max(checkpoints, key=lambda cp: datetime.fromisoformat(cp.timestamp))
|
||||
logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint.checkpoint_id}")
|
||||
return latest_checkpoint
|
||||
|
||||
async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]:
|
||||
"""List checkpoint IDs for a given workflow name.
|
||||
|
||||
Args:
|
||||
workflow_name: The name of the workflow to list checkpoint IDs for.
|
||||
|
||||
Returns:
|
||||
A list of checkpoint IDs for the specified workflow name.
|
||||
"""
|
||||
|
||||
def _list_ids() -> list[CheckpointID]:
|
||||
checkpoint_ids: list[CheckpointID] = []
|
||||
for file_path in self.storage_path.glob("*.json"):
|
||||
try:
|
||||
with open(file_path) as f:
|
||||
data = json.load(f)
|
||||
if data.get("workflow_name") == workflow_name:
|
||||
checkpoint_ids.append(data.get("checkpoint_id", file_path.stem))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to read checkpoint file {file_path}: {e}")
|
||||
return checkpoint_ids
|
||||
|
||||
return await asyncio.to_thread(_list_ids)
|
||||
@@ -0,0 +1,311 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Checkpoint encoding using JSON structure with pickle+base64 for arbitrary data.
|
||||
|
||||
This hybrid approach provides:
|
||||
- Human-readable JSON structure for debugging and inspection of primitives and collections
|
||||
- Full Python object fidelity via pickle for data values (non-JSON-native types)
|
||||
- Base64 encoding to embed binary pickle data in JSON strings
|
||||
|
||||
When ``allowed_types`` is supplied to :func:`decode_checkpoint_value`, a
|
||||
``RestrictedUnpickler`` is used that limits which classes may be instantiated
|
||||
during deserialization. The default built-in safe set covers common Python
|
||||
value types (primitives, datetime, uuid, ...), all ``agent_framework`` internal
|
||||
types, and all ``openai.types`` types. Callers can extend the set by passing
|
||||
additional ``"module:qualname"`` strings.
|
||||
|
||||
Security Model
|
||||
--------------
|
||||
Checkpoint storage is treated as a **trusted data source**. The serialization
|
||||
format uses Python's ``pickle`` module which can execute arbitrary code during
|
||||
deserialization. The ``RestrictedUnpickler`` provides a defense-in-depth
|
||||
allowlist that limits instantiable classes, but it is **not** a security
|
||||
boundary — certain allowlisted builtins (e.g. ``getattr``) are required for
|
||||
legitimate object reconstruction (enums, named tuples) and cannot be removed
|
||||
without breaking compatibility.
|
||||
|
||||
Developers **must** ensure that:
|
||||
|
||||
1. The checkpoint storage backend (file system, Cosmos DB, Azure Blob, Durable
|
||||
Functions storage) is access-controlled and not writable by untrusted
|
||||
parties.
|
||||
2. Data flowing into ``decode_checkpoint_value`` originates exclusively from
|
||||
the application's own checkpoint storage — never from user-supplied HTTP
|
||||
requests, message payloads, or other untrusted sources.
|
||||
3. The ``allowed_types`` parameter is specified whenever possible to restrict
|
||||
the set of reconstructible types to the minimum required by the application.
|
||||
4. Never pass untrusted external input to ``decode_checkpoint_value``. If you
|
||||
must accept external JSON that might contain checkpoint markers, sanitize it
|
||||
first (for example, :func:`agent_framework_durabletask._workflows.serialization.strip_pickle_markers`).
|
||||
|
||||
The allowlist is a mitigation that reduces attack surface but does not
|
||||
eliminate the inherent risks of deserializing untrusted pickle data. Treat
|
||||
your checkpoint storage with the same access controls you would apply to
|
||||
application secrets or database credentials.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
import pickle # nosec # noqa: S403
|
||||
from typing import Any
|
||||
|
||||
from ..exceptions import WorkflowCheckpointException
|
||||
|
||||
logger = logging.getLogger("agent_framework")
|
||||
|
||||
# Marker to identify pickled values in serialized JSON
|
||||
_PICKLE_MARKER = "__pickled__"
|
||||
_TYPE_MARKER = "__type__"
|
||||
|
||||
# Types that are natively JSON-serializable and don't need pickling
|
||||
_JSON_NATIVE_TYPES = (str, int, float, bool, type(None))
|
||||
|
||||
# Module prefix for framework-internal types that are always allowed
|
||||
_FRAMEWORK_MODULE_PREFIX = "agent_framework."
|
||||
|
||||
# Module prefix for OpenAI SDK types that are always allowed
|
||||
_OPENAI_MODULE_PREFIX = "openai.types."
|
||||
|
||||
# Built-in types considered safe for checkpoint deserialization.
|
||||
# Each entry is a ``module:qualname`` string matching the format produced by
|
||||
# :func:`_type_to_key`. These are the classes for which pickle's
|
||||
# ``find_class`` will be called when unpickling common Python value types.
|
||||
_BUILTIN_ALLOWED_TYPE_KEYS: frozenset[str] = frozenset({
|
||||
# builtins
|
||||
"builtins:object",
|
||||
"builtins:complex",
|
||||
"builtins:range",
|
||||
"builtins:slice",
|
||||
"builtins:int",
|
||||
"builtins:float",
|
||||
"builtins:str",
|
||||
"builtins:bytes",
|
||||
"builtins:bytearray",
|
||||
"builtins:bool",
|
||||
"builtins:set",
|
||||
"builtins:frozenset",
|
||||
"builtins:list",
|
||||
"builtins:dict",
|
||||
"builtins:tuple",
|
||||
"builtins:type",
|
||||
# getattr is used by pickle to reconstruct enum members
|
||||
"builtins:getattr",
|
||||
# copyreg helpers used by pickle for object reconstruction
|
||||
"copyreg:_reconstructor",
|
||||
# datetime
|
||||
"datetime:datetime",
|
||||
"datetime:date",
|
||||
"datetime:time",
|
||||
"datetime:timedelta",
|
||||
"datetime:timezone",
|
||||
# uuid
|
||||
"uuid:UUID",
|
||||
# decimal
|
||||
"decimal:Decimal",
|
||||
# collections
|
||||
"collections:OrderedDict",
|
||||
"collections:defaultdict",
|
||||
"collections:deque",
|
||||
})
|
||||
|
||||
|
||||
class _RestrictedUnpickler(pickle.Unpickler): # noqa: S301
|
||||
"""Unpickler that restricts which classes may be instantiated.
|
||||
|
||||
Only classes whose ``module:qualname`` key appears in the combined allow
|
||||
set (built-in safe types + framework types + OpenAI SDK types +
|
||||
caller-specified extras) are permitted. All other classes raise
|
||||
:class:`pickle.UnpicklingError`.
|
||||
"""
|
||||
|
||||
def __init__(self, data: bytes, allowed_types: frozenset[str]) -> None:
|
||||
super().__init__(io.BytesIO(data))
|
||||
self._allowed_types = allowed_types
|
||||
|
||||
def find_class(self, module: str, name: str) -> type:
|
||||
type_key = f"{module}:{name}"
|
||||
|
||||
if (
|
||||
type_key in _BUILTIN_ALLOWED_TYPE_KEYS
|
||||
or type_key in self._allowed_types
|
||||
or module.startswith(_FRAMEWORK_MODULE_PREFIX)
|
||||
or module.startswith(_OPENAI_MODULE_PREFIX)
|
||||
):
|
||||
return super().find_class(module, name)
|
||||
|
||||
raise pickle.UnpicklingError(
|
||||
f"Checkpoint deserialization blocked for type '{type_key}'. "
|
||||
f"To allow this type, either include its 'module:qualname' key in the "
|
||||
f"'allowed_types' set passed to 'decode_checkpoint_value', or add it to "
|
||||
f"'allowed_checkpoint_types' on your checkpoint storage "
|
||||
f"(for example, 'FileCheckpointStorage.allowed_checkpoint_types')."
|
||||
)
|
||||
|
||||
|
||||
def encode_checkpoint_value(value: Any) -> Any:
|
||||
"""Encode a Python value for checkpoint storage.
|
||||
|
||||
JSON-native types (str, int, float, bool, None) pass through unchanged.
|
||||
Collections (dict, list) are recursed with their values encoded.
|
||||
All other types (dataclasses, custom objects, datetime, etc.) are pickled
|
||||
and stored as base64-encoded strings.
|
||||
|
||||
Args:
|
||||
value: Any Python value to encode.
|
||||
|
||||
Returns:
|
||||
A JSON-serializable representation of the value.
|
||||
"""
|
||||
return _encode(value)
|
||||
|
||||
|
||||
def decode_checkpoint_value(value: Any, *, allowed_types: frozenset[str] | None = None) -> Any:
|
||||
"""Decode a value from checkpoint storage.
|
||||
|
||||
Reverses the encoding performed by encode_checkpoint_value.
|
||||
Pickled values (identified by _PICKLE_MARKER) are decoded and unpickled.
|
||||
|
||||
Args:
|
||||
value: A JSON-deserialized value from checkpoint storage.
|
||||
allowed_types: If not ``None``, restrict pickle deserialization to the
|
||||
built-in safe set, framework types, and the types listed here.
|
||||
Each entry should use ``"module:qualname"`` format — that is, the
|
||||
dotted module path followed by a colon and the class
|
||||
``__qualname__``. For example, given a user-defined class::
|
||||
|
||||
# my_app/models.py
|
||||
class MyState: ...
|
||||
|
||||
the corresponding entry would be ``"my_app.models:MyState"``::
|
||||
|
||||
decode_checkpoint_value(
|
||||
data,
|
||||
allowed_types=frozenset({"my_app.models:MyState"}),
|
||||
)
|
||||
|
||||
When using :class:`FileCheckpointStorage`, pass the same strings
|
||||
via ``allowed_checkpoint_types``::
|
||||
|
||||
storage = FileCheckpointStorage(
|
||||
"/tmp/checkpoints",
|
||||
allowed_checkpoint_types=["my_app.models:MyState"],
|
||||
)
|
||||
|
||||
If ``None``, no restriction is applied (backward-compatible
|
||||
behavior).
|
||||
|
||||
Returns:
|
||||
The original Python value.
|
||||
|
||||
Raises:
|
||||
WorkflowCheckpointException: If the unpickled object's type doesn't match
|
||||
the recorded type, indicating corruption, if the base64/pickle
|
||||
data is malformed, or if a disallowed type is encountered during
|
||||
restricted deserialization.
|
||||
"""
|
||||
return _decode(value, allowed_types=allowed_types)
|
||||
|
||||
|
||||
def _encode(value: Any) -> Any:
|
||||
"""Recursively encode a value for JSON storage."""
|
||||
# JSON-native types pass through
|
||||
if isinstance(value, _JSON_NATIVE_TYPES):
|
||||
return value
|
||||
|
||||
# Recursively encode dict values (keys become strings)
|
||||
if isinstance(value, dict):
|
||||
return {str(k): _encode(v) for k, v in value.items()} # type: ignore
|
||||
|
||||
# Recursively encode list items (lists are JSON-native collections)
|
||||
if isinstance(value, list):
|
||||
return [_encode(item) for item in value] # type: ignore
|
||||
|
||||
# Everything else (tuples, sets, dataclasses, custom objects, etc.): pickle and base64 encode
|
||||
return {
|
||||
_PICKLE_MARKER: _pickle_to_base64(value),
|
||||
_TYPE_MARKER: _type_to_key(type(value)), # type: ignore
|
||||
}
|
||||
|
||||
|
||||
def _decode(value: Any, *, allowed_types: frozenset[str] | None = None) -> Any:
|
||||
"""Recursively decode a value from JSON storage."""
|
||||
# JSON-native types pass through
|
||||
if isinstance(value, _JSON_NATIVE_TYPES):
|
||||
return value
|
||||
|
||||
# Handle encoded dicts
|
||||
if isinstance(value, dict):
|
||||
# Pickled value: decode, unpickle, and verify type
|
||||
if _PICKLE_MARKER in value and _TYPE_MARKER in value:
|
||||
obj = _base64_to_unpickle(value[_PICKLE_MARKER], allowed_types=allowed_types) # type: ignore
|
||||
_verify_type(obj, value.get(_TYPE_MARKER)) # type: ignore
|
||||
return obj
|
||||
|
||||
# Regular dict: decode values recursively
|
||||
return {k: _decode(v, allowed_types=allowed_types) for k, v in value.items()} # type: ignore
|
||||
|
||||
# Handle encoded lists
|
||||
if isinstance(value, list):
|
||||
return [_decode(item, allowed_types=allowed_types) for item in value] # type: ignore
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def _verify_type(obj: Any, expected_type_key: str) -> None:
|
||||
"""Verify that an unpickled object matches its recorded type.
|
||||
|
||||
This is a post-deserialization integrity check that detects accidental
|
||||
corruption or type mismatches. It does not prevent arbitrary code execution
|
||||
from malicious pickle payloads, since ``pickle.loads()`` has already
|
||||
executed by the time this function is called.
|
||||
|
||||
Args:
|
||||
obj: The unpickled object.
|
||||
expected_type_key: The recorded type key (module:qualname format).
|
||||
|
||||
Raises:
|
||||
WorkflowCheckpointException: If the types don't match.
|
||||
"""
|
||||
actual_type_key = _type_to_key(type(obj)) # type: ignore
|
||||
if actual_type_key != expected_type_key:
|
||||
raise WorkflowCheckpointException(
|
||||
f"Type mismatch during checkpoint decoding: "
|
||||
f"expected '{expected_type_key}', got '{actual_type_key}'. "
|
||||
f"The checkpoint may be corrupted or tampered with."
|
||||
)
|
||||
|
||||
|
||||
def _pickle_to_base64(value: Any) -> str:
|
||||
"""Pickle a value and encode as base64 string."""
|
||||
pickled = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
return base64.b64encode(pickled).decode("ascii")
|
||||
|
||||
|
||||
def _base64_to_unpickle(encoded: str, *, allowed_types: frozenset[str] | None = None) -> Any:
|
||||
"""Decode base64 string and unpickle.
|
||||
|
||||
Args:
|
||||
encoded: Base64-encoded pickle data.
|
||||
allowed_types: If not ``None``, use restricted unpickling that only
|
||||
permits built-in safe types, framework types, and the specified
|
||||
extra types.
|
||||
|
||||
Raises:
|
||||
WorkflowCheckpointException: If the base64 data is corrupted, the pickle
|
||||
format is incompatible, or a disallowed type is encountered.
|
||||
"""
|
||||
try:
|
||||
pickled = base64.b64decode(encoded.encode("ascii"))
|
||||
if allowed_types is not None:
|
||||
return _RestrictedUnpickler(pickled, allowed_types).load()
|
||||
return pickle.loads(pickled) # nosec # noqa: S301
|
||||
except Exception as exc:
|
||||
raise WorkflowCheckpointException(f"Failed to decode pickled checkpoint data: {exc}") from exc
|
||||
|
||||
|
||||
def _type_to_key(t: type) -> str:
|
||||
"""Convert a type to a module:qualname string."""
|
||||
return f"{t.__module__}:{t.__qualname__}"
|
||||
@@ -0,0 +1,24 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
# Default maximum iterations for workflow execution.
|
||||
DEFAULT_MAX_ITERATIONS = 100
|
||||
|
||||
# Key used to store executor state in state.
|
||||
EXECUTOR_STATE_KEY = "_executor_state"
|
||||
|
||||
# Source identifier for internal workflow messages.
|
||||
INTERNAL_SOURCE_PREFIX = "internal"
|
||||
|
||||
# State key for storing run kwargs that should be passed to agent invocations.
|
||||
# Used by all orchestration patterns (Sequential, Concurrent, GroupChat, Handoff, Magentic)
|
||||
# to pass kwargs from workflow.run() through to agent.run() and @tool functions.
|
||||
WORKFLOW_RUN_KWARGS_KEY = "_workflow_run_kwargs"
|
||||
|
||||
# Sentinel key used in resolved invocation kwargs dicts to denote global kwargs
|
||||
# that apply to all executors (as opposed to per-executor keyed entries).
|
||||
GLOBAL_KWARGS_KEY = "__global__"
|
||||
|
||||
|
||||
def INTERNAL_SOURCE_ID(executor_id: str) -> str:
|
||||
"""Generate an internal source ID for a given executor."""
|
||||
return f"{INTERNAL_SOURCE_PREFIX}:{executor_id}"
|
||||
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from .._types import Message
|
||||
|
||||
"""Helpers for managing chat conversation history.
|
||||
|
||||
These utilities operate on standard `list[Message]` collections and simple
|
||||
dictionary snapshots so orchestrators can share logic without new mixins.
|
||||
"""
|
||||
|
||||
|
||||
def latest_user_message(conversation: Sequence[Message]) -> Message:
|
||||
"""Return the most recent user-authored message from `conversation`."""
|
||||
for message in reversed(conversation):
|
||||
role_value = getattr(message.role, "value", message.role)
|
||||
if str(role_value).lower() == "user":
|
||||
return message
|
||||
raise ValueError("No user message in conversation")
|
||||
|
||||
|
||||
def ensure_author(message: Message, fallback: str) -> Message:
|
||||
"""Attach `fallback` author if message is missing `author_name`."""
|
||||
message.author_name = message.author_name or fallback
|
||||
return message
|
||||
@@ -0,0 +1,943 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, ClassVar, TypeAlias, TypeVar
|
||||
|
||||
from .._agents import SupportsAgentRun
|
||||
from ._const import INTERNAL_SOURCE_ID
|
||||
from ._executor import Executor
|
||||
from ._model_utils import DictConvertible, encode_value
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Type alias for edge condition functions.
|
||||
# Conditions receive the message data and return bool (sync or async).
|
||||
EdgeCondition: TypeAlias = Callable[[Any], bool | Awaitable[bool]]
|
||||
|
||||
# TypeVar for EdgeGroup subclasses used in class methods
|
||||
EdgeGroupT = TypeVar("EdgeGroupT", bound="EdgeGroup")
|
||||
|
||||
|
||||
def _extract_function_name(func: Callable[..., Any]) -> str:
|
||||
"""Map a Python callable to a concise, human-focused identifier.
|
||||
|
||||
The workflow graph persists references to callables by recording only an
|
||||
identifier. This helper inspects standard callable metadata and picks a
|
||||
stable value so that serialized representations remain intelligible when
|
||||
they are later rendered in logs or reconstructed during deserialization.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
def threshold(value: float) -> bool:
|
||||
return value > 0.5
|
||||
|
||||
|
||||
assert _extract_function_name(threshold) == "threshold"
|
||||
"""
|
||||
if hasattr(func, "__name__"):
|
||||
name = func.__name__
|
||||
return name if name != "<lambda>" else "<lambda>"
|
||||
return "<callable>"
|
||||
|
||||
|
||||
def _missing_callable(name: str) -> Callable[..., Any]:
|
||||
"""Create a defensive placeholder for callables that cannot be restored.
|
||||
|
||||
When a workflow is deserialized in an environment that lacks the original
|
||||
Python callable, we install a proxy that fails loudly. Surfacing the error
|
||||
at invocation time preserves a clean separation between I/O concerns and
|
||||
runtime execution, while making it obvious which callable needs to be
|
||||
re-registered.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
guard = _missing_callable("transform_price")
|
||||
try:
|
||||
guard()
|
||||
except RuntimeError as exc:
|
||||
assert "transform_price" in str(exc)
|
||||
"""
|
||||
|
||||
def _raise(*_: Any, **__: Any) -> Any:
|
||||
raise RuntimeError(f"Callable '{name}' is unavailable after serialization")
|
||||
|
||||
return _raise
|
||||
|
||||
|
||||
@dataclass(init=False)
|
||||
class Edge(DictConvertible):
|
||||
"""Model a directed, optionally-conditional hand-off between two executors.
|
||||
|
||||
Each `Edge` captures the minimal metadata required to move a message from
|
||||
one executor to another inside the workflow graph. It optionally embeds a
|
||||
boolean predicate that decides if the edge should be taken at runtime. By
|
||||
serialising the edge down to primitives we can reconstruct the topology of
|
||||
a workflow irrespective of the original Python process.
|
||||
|
||||
Edge conditions receive the message data and return a boolean (sync or async).
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
edge = Edge(source_id="ingest", target_id="score", condition=lambda data: data["ready"])
|
||||
assert await edge.should_route({"ready": True}) is True
|
||||
"""
|
||||
|
||||
ID_SEPARATOR: ClassVar[str] = "->"
|
||||
|
||||
source_id: str
|
||||
target_id: str
|
||||
condition_name: str | None
|
||||
_condition: EdgeCondition | None = field(default=None, repr=False, compare=False)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str,
|
||||
target_id: str,
|
||||
condition: EdgeCondition | None = None,
|
||||
*,
|
||||
condition_name: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize a fully-specified edge between two workflow executors.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source_id:
|
||||
Canonical identifier of the upstream executor instance.
|
||||
target_id:
|
||||
Canonical identifier of the downstream executor instance.
|
||||
condition:
|
||||
Optional predicate that receives the message data and returns
|
||||
`True` when the edge should be traversed. Can be sync or async.
|
||||
When omitted, the edge is unconditionally active.
|
||||
condition_name:
|
||||
Optional override that pins a human-friendly name for the condition
|
||||
when the callable cannot be introspected (for example after
|
||||
deserialization).
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
edge = Edge("fetch", "parse", condition=lambda data: data.is_valid)
|
||||
assert edge.source_id == "fetch"
|
||||
assert edge.target_id == "parse"
|
||||
"""
|
||||
if not source_id:
|
||||
raise ValueError("Edge source_id must be a non-empty string")
|
||||
if not target_id:
|
||||
raise ValueError("Edge target_id must be a non-empty string")
|
||||
self.source_id = source_id
|
||||
self.target_id = target_id
|
||||
self._condition = condition
|
||||
self.condition_name = (
|
||||
_extract_function_name(condition) if condition is not None and condition_name is None else condition_name
|
||||
)
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
"""Return the stable identifier used to reference this edge.
|
||||
|
||||
The identifier combines the source and target executor identifiers with
|
||||
a deterministic separator. This allows other graph structures such as
|
||||
adjacency lists or visualisations to refer to an edge without carrying
|
||||
the full object.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
edge = Edge("reader", "writer")
|
||||
assert edge.id == "reader->writer"
|
||||
"""
|
||||
return f"{self.source_id}{self.ID_SEPARATOR}{self.target_id}"
|
||||
|
||||
@property
|
||||
def has_condition(self) -> bool:
|
||||
"""Check if this edge has a condition.
|
||||
|
||||
Returns True if the edge was configured with a condition function.
|
||||
"""
|
||||
return self._condition is not None
|
||||
|
||||
async def should_route(self, data: Any) -> bool:
|
||||
"""Evaluate the edge predicate against payload.
|
||||
|
||||
When the edge was defined without an explicit predicate the method
|
||||
returns `True`, signalling an unconditional routing rule. Otherwise the
|
||||
user-supplied callable decides whether the message should proceed along
|
||||
this edge. Any exception raised by the callable is deliberately allowed
|
||||
to surface to the caller to avoid masking logic bugs.
|
||||
|
||||
The condition receives the message data and may be sync or async.
|
||||
|
||||
Args:
|
||||
data: The message payload
|
||||
|
||||
Returns:
|
||||
True if the edge should be traversed, False otherwise.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
edge = Edge("stage1", "stage2", condition=lambda data: data["score"] > 0.8)
|
||||
assert await edge.should_route({"score": 0.9}) is True
|
||||
assert await edge.should_route({"score": 0.4}) is False
|
||||
"""
|
||||
if self._condition is None:
|
||||
return True
|
||||
result = self._condition(data)
|
||||
if inspect.isawaitable(result):
|
||||
return bool(await result)
|
||||
return bool(result)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Produce a JSON-serialisable view of the edge metadata.
|
||||
|
||||
The representation includes the source and target executor identifiers
|
||||
plus the condition name when it is known. Serialisation intentionally
|
||||
omits the live callable to keep payloads transport-friendly.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
edge = Edge("reader", "writer", condition=lambda payload: payload["ok"])
|
||||
snapshot = edge.to_dict()
|
||||
assert snapshot == {"source_id": "reader", "target_id": "writer", "condition_name": "<lambda>"}
|
||||
"""
|
||||
payload = {"source_id": self.source_id, "target_id": self.target_id}
|
||||
if self.condition_name is not None:
|
||||
payload["condition_name"] = self.condition_name
|
||||
return payload
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> Edge:
|
||||
"""Reconstruct an `Edge` from its serialised dictionary form.
|
||||
|
||||
The deserialised edge will lack the executable predicate because we do
|
||||
not attempt to hydrate Python callables from storage. Instead, the
|
||||
stored `condition_name` is preserved so that downstream consumers can
|
||||
detect missing callables and re-register them where appropriate.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
payload = {"source_id": "reader", "target_id": "writer", "condition_name": "is_ready"}
|
||||
edge = Edge.from_dict(payload)
|
||||
assert edge.source_id == "reader"
|
||||
assert edge.condition_name == "is_ready"
|
||||
"""
|
||||
return cls(
|
||||
source_id=data["source_id"],
|
||||
target_id=data["target_id"],
|
||||
condition=None,
|
||||
condition_name=data.get("condition_name"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Case:
|
||||
"""Runtime wrapper combining a switch-case predicate with its target.
|
||||
|
||||
Each `Case` couples a boolean predicate with the executor that should
|
||||
handle the message when the predicate evaluates to `True`. The runtime
|
||||
keeps this lightweight container separate from the serialisable
|
||||
`SwitchCaseEdgeGroupCase` so that execution can operate with live callables
|
||||
without polluting persisted state.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
class JsonExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="json", defer_discovery=True)
|
||||
|
||||
|
||||
processor = JsonExecutor()
|
||||
case = Case(condition=lambda payload: payload["kind"] == "json", target=processor)
|
||||
assert case.target.id == "json"
|
||||
"""
|
||||
|
||||
condition: Callable[[Any], bool]
|
||||
target: Executor | SupportsAgentRun
|
||||
|
||||
|
||||
@dataclass
|
||||
class Default:
|
||||
"""Runtime representation of the default branch in a switch-case group.
|
||||
|
||||
The default branch is invoked only when no other case predicates match. In
|
||||
practice it is guaranteed to exist so that routing never produces an empty
|
||||
target.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
class DeadLetterExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="dead_letter", defer_discovery=True)
|
||||
|
||||
|
||||
fallback = Default(target=DeadLetterExecutor())
|
||||
assert fallback.target.id == "dead_letter"
|
||||
"""
|
||||
|
||||
target: Executor | SupportsAgentRun
|
||||
|
||||
|
||||
@dataclass(init=False)
|
||||
class EdgeGroup(DictConvertible):
|
||||
"""Bundle edges that share a common routing semantics under a single id.
|
||||
|
||||
The workflow runtime manipulates `EdgeGroup` instances rather than raw
|
||||
edges so it can reason about higher-order routing behaviours such as
|
||||
fan-out, fan-in, switch-case, and other graph patterns. The base class stores the
|
||||
identifying information and handles serialisation duties so specialised
|
||||
groups need only maintain their additional state.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = EdgeGroup([Edge("source", "sink")])
|
||||
assert group.source_executor_ids == ["source"]
|
||||
"""
|
||||
|
||||
id: str
|
||||
type: str
|
||||
edges: list[Edge]
|
||||
|
||||
from builtins import type as builtin_type
|
||||
|
||||
_TYPE_REGISTRY: ClassVar[dict[str, builtin_type[EdgeGroup]]] = {}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
edges: Sequence[Edge] | None = None,
|
||||
*,
|
||||
id: str | None = None,
|
||||
type: str | None = None,
|
||||
) -> None:
|
||||
"""Construct an edge group shell around a set of `Edge` instances.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
edges:
|
||||
Sequence of edges that participate in this group. When omitted we
|
||||
start from an empty list so subclasses can append later.
|
||||
id:
|
||||
Stable identifier for the group. Defaults to a random UUID so
|
||||
serialised graphs remain uniquely addressable.
|
||||
type:
|
||||
Logical discriminator used to recover the appropriate subclass when
|
||||
de-serialising.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
edges = [Edge("validate", "persist")]
|
||||
group = EdgeGroup(edges, id="stage", type="Custom")
|
||||
assert group.to_dict()["type"] == "Custom"
|
||||
"""
|
||||
self.id = id or f"{self.__class__.__name__}/{uuid.uuid4()}"
|
||||
self.type = type or self.__class__.__name__
|
||||
self.edges = list(edges) if edges is not None else []
|
||||
|
||||
@property
|
||||
def source_executor_ids(self) -> list[str]:
|
||||
"""Return the deduplicated list of upstream executor ids.
|
||||
|
||||
The property preserves order-of-first-appearance so the caller can rely
|
||||
on deterministic iteration when reconstructing graph topology.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = EdgeGroup([Edge("read", "write"), Edge("read", "archive")])
|
||||
assert group.source_executor_ids == ["read"]
|
||||
"""
|
||||
return list(dict.fromkeys(edge.source_id for edge in self.edges))
|
||||
|
||||
@property
|
||||
def target_executor_ids(self) -> list[str]:
|
||||
"""Return the ordered, deduplicated list of downstream executor ids.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = EdgeGroup([Edge("read", "write"), Edge("read", "archive")])
|
||||
assert group.target_executor_ids == ["write", "archive"]
|
||||
"""
|
||||
return list(dict.fromkeys(edge.target_id for edge in self.edges))
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialise the group metadata and contained edges into primitives.
|
||||
|
||||
The payload captures each edge through its own `to_dict` call, enabling
|
||||
round-tripping through formats such as JSON without leaking Python
|
||||
objects.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = EdgeGroup([Edge("read", "write")])
|
||||
snapshot = group.to_dict()
|
||||
assert snapshot["edges"][0]["source_id"] == "read"
|
||||
"""
|
||||
return {
|
||||
"id": self.id,
|
||||
"type": self.type,
|
||||
"edges": [edge.to_dict() for edge in self.edges],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def register(cls, subclass: builtin_type[EdgeGroupT]) -> builtin_type[EdgeGroupT]:
|
||||
"""Register a subclass so deserialisation can recover the right type.
|
||||
|
||||
Registration is typically performed via the decorator syntax applied to
|
||||
each concrete edge group. The registry stores classes by their
|
||||
`__name__`, which must therefore remain stable across versions when
|
||||
persisted workflows are in circulation.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
@EdgeGroup.register
|
||||
class CustomGroup(EdgeGroup):
|
||||
pass
|
||||
|
||||
|
||||
assert EdgeGroup._TYPE_REGISTRY["CustomGroup"] is CustomGroup
|
||||
"""
|
||||
cls._TYPE_REGISTRY[subclass.__name__] = subclass
|
||||
return subclass
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> EdgeGroup:
|
||||
"""Hydrate the correct `EdgeGroup` subclass from serialised state.
|
||||
|
||||
The method inspects the `type` field, allocates the corresponding class
|
||||
without executing subclass `__init__`, and then manually restores any
|
||||
subtype-specific attributes. This keeps deserialisation deterministic
|
||||
even for complex group types that configure additional runtime
|
||||
callables.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
payload = {"type": "EdgeGroup", "edges": [{"source_id": "a", "target_id": "b"}]}
|
||||
group = EdgeGroup.from_dict(payload)
|
||||
assert isinstance(group, EdgeGroup)
|
||||
"""
|
||||
group_type = data.get("type", "EdgeGroup")
|
||||
target_cls = cls._TYPE_REGISTRY.get(group_type, EdgeGroup)
|
||||
edges = [Edge.from_dict(entry) for entry in data.get("edges", [])]
|
||||
|
||||
obj = target_cls.__new__(target_cls)
|
||||
EdgeGroup.__init__(obj, edges=edges, id=data.get("id"), type=group_type)
|
||||
|
||||
# Handle FanOutEdgeGroup-specific attributes
|
||||
if isinstance(obj, FanOutEdgeGroup):
|
||||
obj.selection_func_name = data.get("selection_func_name")
|
||||
obj._selection_func = ( # type: ignore[attr-defined]
|
||||
None if obj.selection_func_name is None else _missing_callable(obj.selection_func_name)
|
||||
)
|
||||
obj._target_ids = [edge.target_id for edge in obj.edges] # type: ignore[attr-defined]
|
||||
|
||||
# Handle SwitchCaseEdgeGroup-specific attributes
|
||||
if isinstance(obj, SwitchCaseEdgeGroup):
|
||||
cases_payload = data.get("cases", [])
|
||||
restored_cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault] = []
|
||||
for case_data in cases_payload:
|
||||
case_type = case_data.get("type")
|
||||
if case_type == "Default":
|
||||
restored_cases.append(SwitchCaseEdgeGroupDefault.from_dict(case_data))
|
||||
else:
|
||||
restored_cases.append(SwitchCaseEdgeGroupCase.from_dict(case_data))
|
||||
obj.cases = restored_cases
|
||||
obj._selection_func = _missing_callable("switch_case_selection") # type: ignore[attr-defined]
|
||||
|
||||
return obj
|
||||
|
||||
|
||||
@EdgeGroup.register
|
||||
@dataclass(init=False)
|
||||
class SingleEdgeGroup(EdgeGroup):
|
||||
"""Convenience wrapper for a solitary edge, keeping the group API uniform."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str,
|
||||
target_id: str,
|
||||
condition: EdgeCondition | None = None,
|
||||
*,
|
||||
id: str | None = None,
|
||||
) -> None:
|
||||
"""Create a one-to-one edge group between two executors.
|
||||
|
||||
Args:
|
||||
source_id: The source executor ID.
|
||||
target_id: The target executor ID.
|
||||
condition: Optional condition function `(data) -> bool | Awaitable[bool]`.
|
||||
id: Optional explicit ID for the edge group.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = SingleEdgeGroup("ingest", "validate")
|
||||
assert group.edges[0].source_id == "ingest"
|
||||
"""
|
||||
edge = Edge(source_id=source_id, target_id=target_id, condition=condition)
|
||||
super().__init__([edge], id=id, type=self.__class__.__name__)
|
||||
|
||||
|
||||
@EdgeGroup.register
|
||||
@dataclass(init=False)
|
||||
class FanOutEdgeGroup(EdgeGroup):
|
||||
"""Represent a broadcast-style edge group with optional selection logic.
|
||||
|
||||
A fan-out forwards a message produced by a single source executor to one
|
||||
or more downstream executors. At runtime we may further narrow the targets
|
||||
by executing a `selection_func` that inspects the payload and returns the
|
||||
subset of ids that should receive the message.
|
||||
"""
|
||||
|
||||
selection_func_name: str | None
|
||||
_selection_func: Callable[[Any, list[str]], list[str]] | None
|
||||
_target_ids: list[str]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str,
|
||||
target_ids: Sequence[str],
|
||||
selection_func: Callable[[Any, list[str]], list[str]] | None = None,
|
||||
*,
|
||||
selection_func_name: str | None = None,
|
||||
id: str | None = None,
|
||||
) -> None:
|
||||
"""Create a fan-out mapping from a single source to many targets.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source_id:
|
||||
Identifier of the upstream executor broadcasting the message.
|
||||
target_ids:
|
||||
Ordered set of downstream executor identifiers that may receive the
|
||||
message. At least two targets are required to preserve the fan-out
|
||||
semantics.
|
||||
selection_func:
|
||||
Optional callable that returns the subset of `target_ids` that
|
||||
should be active for a given payload. The callable receives the
|
||||
original message plus a copy of all configured target ids.
|
||||
selection_func_name:
|
||||
Static identifier used when persisting the fan-out. Needed when the
|
||||
callable cannot be introspected or is unavailable during
|
||||
deserialisation.
|
||||
id:
|
||||
Stable identifier for the group; defaults to an autogenerated UUID.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
def choose_targets(message: dict[str, Any], available: list[str]) -> list[str]:
|
||||
return [target for target in available if message.get(target)]
|
||||
|
||||
|
||||
group = FanOutEdgeGroup("sensor", ["db", "cache"], selection_func=choose_targets)
|
||||
assert group.selection_func is choose_targets
|
||||
"""
|
||||
if len(target_ids) <= 1:
|
||||
raise ValueError("FanOutEdgeGroup must contain at least two targets.")
|
||||
|
||||
edges = [Edge(source_id=source_id, target_id=target) for target in target_ids]
|
||||
super().__init__(edges, id=id, type=self.__class__.__name__)
|
||||
|
||||
self._target_ids = list(target_ids)
|
||||
self._selection_func = selection_func
|
||||
self.selection_func_name = (
|
||||
_extract_function_name(selection_func) if selection_func is not None else selection_func_name
|
||||
)
|
||||
|
||||
@property
|
||||
def target_ids(self) -> list[str]:
|
||||
"""Return a shallow copy of the configured downstream executor ids.
|
||||
|
||||
The list is defensively copied to prevent callers from mutating the
|
||||
internal state while still providing deterministic ordering.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = FanOutEdgeGroup("node", ["alpha", "beta"])
|
||||
assert group.target_ids == ["alpha", "beta"]
|
||||
"""
|
||||
return list(self._target_ids)
|
||||
|
||||
@property
|
||||
def selection_func(self) -> Callable[[Any, list[str]], list[str]] | None:
|
||||
"""Expose the runtime callable used to select active fan-out targets.
|
||||
|
||||
When no selection function was supplied the property returns `None`,
|
||||
signalling that all targets must receive the payload.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = FanOutEdgeGroup("source", ["x", "y"], selection_func=None)
|
||||
assert group.selection_func is None
|
||||
"""
|
||||
return self._selection_func
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialise the fan-out group while preserving selection metadata.
|
||||
|
||||
In addition to the base `EdgeGroup` payload we embed the human-friendly
|
||||
name of the selection function. The callable itself is not persisted.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = FanOutEdgeGroup("source", ["a", "b"], selection_func=lambda *_: ["a"])
|
||||
snapshot = group.to_dict()
|
||||
assert snapshot["selection_func_name"] == "<lambda>"
|
||||
"""
|
||||
payload = super().to_dict()
|
||||
payload["selection_func_name"] = self.selection_func_name
|
||||
return payload
|
||||
|
||||
|
||||
@EdgeGroup.register
|
||||
@dataclass(init=False)
|
||||
class FanInEdgeGroup(EdgeGroup):
|
||||
"""Represent a converging set of edges that feed a single downstream executor.
|
||||
|
||||
Fan-in groups are typically used when multiple upstream stages independently
|
||||
produce messages that should all arrive at the same downstream processor.
|
||||
"""
|
||||
|
||||
def __init__(self, source_ids: Sequence[str], target_id: str, *, id: str | None = None) -> None:
|
||||
"""Build a fan-in mapping that merges several sources into one target.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source_ids:
|
||||
Sequence of upstream executor identifiers contributing messages.
|
||||
target_id:
|
||||
Downstream executor that receives every message emitted by the
|
||||
sources.
|
||||
id:
|
||||
Optional explicit identifier for the edge group.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = FanInEdgeGroup(["parser", "enricher"], target_id="writer")
|
||||
assert group.to_dict()["edges"][0]["target_id"] == "writer"
|
||||
"""
|
||||
if len(source_ids) <= 1:
|
||||
raise ValueError("FanInEdgeGroup must contain at least two sources.")
|
||||
|
||||
edges = [Edge(source_id=source, target_id=target_id) for source in source_ids]
|
||||
super().__init__(edges, id=id, type=self.__class__.__name__)
|
||||
|
||||
|
||||
@dataclass(init=False)
|
||||
class SwitchCaseEdgeGroupCase(DictConvertible):
|
||||
"""Persistable description of a single conditional branch in a switch-case.
|
||||
|
||||
Unlike the runtime `Case` object this serialisable variant stores only the
|
||||
target identifier and a descriptive name for the predicate. When the
|
||||
underlying callable is unavailable during deserialisation we substitute a
|
||||
proxy placeholder that fails loudly, ensuring the missing dependency is
|
||||
immediately visible.
|
||||
"""
|
||||
|
||||
target_id: str
|
||||
condition_name: str | None
|
||||
type: str
|
||||
_condition: Callable[[Any], bool] = field(repr=False, compare=False)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
condition: Callable[[Any], bool] | None,
|
||||
target_id: str,
|
||||
*,
|
||||
condition_name: str | None = None,
|
||||
) -> None:
|
||||
"""Record the routing metadata for a conditional case branch.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
condition:
|
||||
Optional live predicate. When omitted we fall back to a placeholder
|
||||
that raises at runtime to highlight missing registrations.
|
||||
target_id:
|
||||
Identifier of the executor that should handle messages when the
|
||||
predicate succeeds.
|
||||
condition_name:
|
||||
Human-friendly label for the predicate used for diagnostics and
|
||||
on-disk persistence.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
case = SwitchCaseEdgeGroupCase(lambda payload: payload["type"] == "csv", target_id="csv_handler")
|
||||
assert case.condition_name == "<lambda>"
|
||||
"""
|
||||
if not target_id:
|
||||
raise ValueError("SwitchCaseEdgeGroupCase requires a target_id")
|
||||
self.target_id = target_id
|
||||
self.type = "Case"
|
||||
if condition is not None:
|
||||
self._condition = condition
|
||||
self.condition_name = _extract_function_name(condition)
|
||||
else:
|
||||
safe_name = condition_name or "<missing_condition>"
|
||||
self._condition = _missing_callable(safe_name)
|
||||
self.condition_name = condition_name
|
||||
|
||||
@property
|
||||
def condition(self) -> Callable[[Any], bool]:
|
||||
"""Return the predicate associated with this case.
|
||||
|
||||
The placeholder installed during deserialisation raises a
|
||||
`RuntimeError` when invoked so that workflow authors are forced to
|
||||
provide the missing callable explicitly.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
case = SwitchCaseEdgeGroupCase(None, target_id="missing", condition_name="needs_registration")
|
||||
guard = case.condition
|
||||
try:
|
||||
guard({})
|
||||
except RuntimeError:
|
||||
pass
|
||||
"""
|
||||
return self._condition
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialise the case metadata without the executable predicate.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
case = SwitchCaseEdgeGroupCase(lambda _: True, target_id="handler")
|
||||
assert case.to_dict()["target_id"] == "handler"
|
||||
"""
|
||||
payload = {"target_id": self.target_id, "type": self.type}
|
||||
if self.condition_name is not None:
|
||||
payload["condition_name"] = self.condition_name
|
||||
return payload
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> SwitchCaseEdgeGroupCase:
|
||||
"""Instantiate a case from its serialised dictionary payload.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
payload = {"target_id": "handler", "condition_name": "is_ready"}
|
||||
case = SwitchCaseEdgeGroupCase.from_dict(payload)
|
||||
assert case.target_id == "handler"
|
||||
"""
|
||||
return cls(
|
||||
condition=None,
|
||||
target_id=data["target_id"],
|
||||
condition_name=data.get("condition_name"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(init=False)
|
||||
class SwitchCaseEdgeGroupDefault(DictConvertible):
|
||||
"""Persistable descriptor for the fallback branch of a switch-case group.
|
||||
|
||||
The default branch is guaranteed to exist and is invoked when every other
|
||||
case predicate fails to match the payload.
|
||||
"""
|
||||
|
||||
target_id: str
|
||||
type: str
|
||||
|
||||
def __init__(self, target_id: str) -> None:
|
||||
"""Point the default branch toward the given executor identifier.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
fallback = SwitchCaseEdgeGroupDefault(target_id="dead_letter")
|
||||
assert fallback.target_id == "dead_letter"
|
||||
"""
|
||||
if not target_id:
|
||||
raise ValueError("SwitchCaseEdgeGroupDefault requires a target_id")
|
||||
self.target_id = target_id
|
||||
self.type = "Default"
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialise the default branch metadata for persistence or logging.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
fallback = SwitchCaseEdgeGroupDefault("dead_letter")
|
||||
assert fallback.to_dict()["type"] == "Default"
|
||||
"""
|
||||
return {"target_id": self.target_id, "type": self.type}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> SwitchCaseEdgeGroupDefault:
|
||||
"""Recreate the default branch from its persisted form.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
payload = {"target_id": "dead_letter", "type": "Default"}
|
||||
fallback = SwitchCaseEdgeGroupDefault.from_dict(payload)
|
||||
assert fallback.target_id == "dead_letter"
|
||||
"""
|
||||
return cls(target_id=data["target_id"])
|
||||
|
||||
|
||||
@EdgeGroup.register
|
||||
@dataclass(init=False)
|
||||
class SwitchCaseEdgeGroup(FanOutEdgeGroup):
|
||||
"""Fan-out variant that mimics a traditional switch/case control flow.
|
||||
|
||||
Each case inspects the message payload and decides whether it should handle
|
||||
the message. Exactly one case-or the default branch-returns a target at
|
||||
runtime, preserving single-dispatch semantics.
|
||||
"""
|
||||
|
||||
cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str,
|
||||
cases: Sequence[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault],
|
||||
*,
|
||||
id: str | None = None,
|
||||
) -> None:
|
||||
"""Configure a switch/case routing structure for a single source executor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source_id:
|
||||
Identifier of the executor producing the message to be routed.
|
||||
cases:
|
||||
Ordered sequence of case descriptors concluding with a
|
||||
`SwitchCaseEdgeGroupDefault`. Ordering matters because the runtime
|
||||
evaluates each branch sequentially until one matches.
|
||||
id:
|
||||
Optional explicit identifier for the edge group.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
cases = [
|
||||
SwitchCaseEdgeGroupCase(lambda payload: payload["kind"] == "csv", target_id="process_csv"),
|
||||
SwitchCaseEdgeGroupDefault(target_id="process_default"),
|
||||
]
|
||||
group = SwitchCaseEdgeGroup("router", cases)
|
||||
encoded = group.to_dict()
|
||||
assert encoded["cases"][0]["type"] == "Case"
|
||||
"""
|
||||
if len(cases) < 2:
|
||||
raise ValueError("SwitchCaseEdgeGroup must contain at least two cases (including the default case).")
|
||||
|
||||
default_cases = [case for case in cases if isinstance(case, SwitchCaseEdgeGroupDefault)]
|
||||
if len(default_cases) != 1:
|
||||
raise ValueError("SwitchCaseEdgeGroup must contain exactly one default case.")
|
||||
|
||||
if not isinstance(cases[-1], SwitchCaseEdgeGroupDefault):
|
||||
logger.warning(
|
||||
"Default case in the switch-case edge group is not the last case. "
|
||||
"This may result in unexpected behavior."
|
||||
)
|
||||
|
||||
def selection_func(message: Any, targets: list[str]) -> list[str]:
|
||||
for case in cases:
|
||||
if isinstance(case, SwitchCaseEdgeGroupDefault):
|
||||
return [case.target_id]
|
||||
try:
|
||||
if case.condition(message):
|
||||
return [case.target_id]
|
||||
except Exception as exc: # pragma: no cover - defensive logging
|
||||
logger.warning("Error evaluating condition for case %s: %s", case.target_id, exc)
|
||||
raise RuntimeError("No matching case found in SwitchCaseEdgeGroup")
|
||||
|
||||
target_ids = [case.target_id for case in cases]
|
||||
# Call FanOutEdgeGroup constructor directly to avoid type checking issues
|
||||
edges = [Edge(source_id=source_id, target_id=target) for target in target_ids]
|
||||
EdgeGroup.__init__(self, edges, id=id, type=self.__class__.__name__)
|
||||
|
||||
# Initialize FanOutEdgeGroup-specific attributes
|
||||
self._target_ids = list(target_ids)
|
||||
self._selection_func = selection_func
|
||||
self.selection_func_name = None
|
||||
self.cases = list(cases)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialise the switch-case group, capturing all case descriptors.
|
||||
|
||||
Each case is converted using `encode_value` to respect dataclass
|
||||
semantics as well as any nested serialisable structures.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
group = SwitchCaseEdgeGroup(
|
||||
"router",
|
||||
[
|
||||
SwitchCaseEdgeGroupCase(lambda _: True, target_id="handler"),
|
||||
SwitchCaseEdgeGroupDefault(target_id="fallback"),
|
||||
],
|
||||
)
|
||||
snapshot = group.to_dict()
|
||||
assert len(snapshot["cases"]) == 2
|
||||
"""
|
||||
payload = super().to_dict()
|
||||
payload["cases"] = [encode_value(case) for case in self.cases]
|
||||
return payload
|
||||
|
||||
|
||||
@EdgeGroup.register
|
||||
@dataclass(init=False)
|
||||
class InternalEdgeGroup(EdgeGroup):
|
||||
"""Special edge group used to route internal messages to executors.
|
||||
|
||||
This group is created automatically when a new executor is added to the workflow
|
||||
builder. It contains a single edge that routes messages from the internal source
|
||||
to the executor itself. Internal source represent messages that are generated by
|
||||
the system rather than by another executor. This includes request and response
|
||||
handling.
|
||||
|
||||
This edge group only contains one edge from the internal source to the executor.
|
||||
And it does not support any conditions or complex routing logic.
|
||||
|
||||
During workflow serialization and deserialization, the internal edge group is
|
||||
preserved and visible to systems consuming the workflow definition.
|
||||
|
||||
Messages sent along this edge will also be captured by monitoring and logging systems,
|
||||
allowing for observability into internal message flows (when tracing is enabled).
|
||||
"""
|
||||
|
||||
def __init__(self, executor_id: str) -> None:
|
||||
"""Create an internal edge group from the given edges.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
executor_id:
|
||||
Identifier of the internal executor that should receive messages.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
edge_group = InternalEdgeGroup("executor_a")
|
||||
"""
|
||||
edge = Edge(source_id=INTERNAL_SOURCE_ID(executor_id), target_id=executor_id)
|
||||
super().__init__([edge])
|
||||
@@ -0,0 +1,429 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
|
||||
from ..observability import EdgeGroupDeliveryStatus, OtelAttr, create_edge_group_processing_span
|
||||
from ._edge import (
|
||||
Edge,
|
||||
EdgeGroup,
|
||||
FanInEdgeGroup,
|
||||
FanOutEdgeGroup,
|
||||
InternalEdgeGroup,
|
||||
SingleEdgeGroup,
|
||||
SwitchCaseEdgeGroup,
|
||||
)
|
||||
from ._executor import Executor
|
||||
from ._runner_context import RunnerContext, WorkflowMessage
|
||||
from ._state import State
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EdgeRunner(ABC):
|
||||
"""Abstract base class for edge runners that handle message delivery."""
|
||||
|
||||
def __init__(self, edge_group: EdgeGroup, executors: dict[str, Executor]) -> None:
|
||||
"""Initialize the edge runner with an edge group and executor map.
|
||||
|
||||
Args:
|
||||
edge_group: The edge group to run.
|
||||
executors: Map of executor IDs to executor instances.
|
||||
"""
|
||||
self._edge_group = edge_group
|
||||
self._executors = executors
|
||||
|
||||
@abstractmethod
|
||||
async def send_message(
|
||||
self,
|
||||
message: WorkflowMessage,
|
||||
state: State,
|
||||
ctx: RunnerContext,
|
||||
) -> bool:
|
||||
"""Send a message through the edge group.
|
||||
|
||||
Args:
|
||||
message: The message to send.
|
||||
state: The workflow state.
|
||||
ctx: The context for the runner.
|
||||
|
||||
Returns:
|
||||
bool: True if the message was processed successfully,
|
||||
False if the target executor cannot handle the message.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def _can_handle(self, executor_id: str, message: WorkflowMessage) -> bool:
|
||||
"""Check if an executor can handle the given message data."""
|
||||
if executor_id not in self._executors:
|
||||
return False
|
||||
return self._executors[executor_id].can_handle(message)
|
||||
|
||||
async def _execute_on_target(
|
||||
self,
|
||||
target_id: str,
|
||||
source_ids: list[str],
|
||||
message: WorkflowMessage,
|
||||
state: State,
|
||||
ctx: RunnerContext,
|
||||
) -> None:
|
||||
"""Execute a message on a target executor with trace context."""
|
||||
if target_id not in self._executors:
|
||||
raise RuntimeError(f"Target executor {target_id} not found.")
|
||||
|
||||
target_executor = self._executors[target_id]
|
||||
|
||||
# Execute with trace context parameters
|
||||
await target_executor.execute(
|
||||
message,
|
||||
source_ids, # source_executor_ids
|
||||
state, # state
|
||||
ctx, # runner_context
|
||||
trace_contexts=message.trace_contexts, # Pass trace contexts
|
||||
source_span_ids=message.source_span_ids, # Pass source span IDs for linking
|
||||
)
|
||||
|
||||
|
||||
class SingleEdgeRunner(EdgeRunner):
|
||||
"""Runner for single edge groups."""
|
||||
|
||||
def __init__(self, edge_group: SingleEdgeGroup | InternalEdgeGroup, executors: dict[str, Executor]) -> None:
|
||||
super().__init__(edge_group, executors)
|
||||
self._edge = edge_group.edges[0]
|
||||
|
||||
async def send_message(
|
||||
self,
|
||||
message: WorkflowMessage,
|
||||
state: State,
|
||||
ctx: RunnerContext,
|
||||
) -> bool:
|
||||
"""Send a message through the single edge."""
|
||||
should_execute = False
|
||||
target_id: str | None = None
|
||||
source_id: str | None = None
|
||||
with create_edge_group_processing_span(
|
||||
self._edge_group.__class__.__name__,
|
||||
edge_group_id=self._edge_group.id,
|
||||
message_source_id=message.source_id,
|
||||
message_target_id=message.target_id,
|
||||
source_trace_contexts=message.trace_contexts,
|
||||
source_span_ids=message.source_span_ids,
|
||||
) as span:
|
||||
try:
|
||||
if message.target_id and message.target_id != self._edge.target_id:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value,
|
||||
})
|
||||
return False
|
||||
|
||||
if self._can_handle(self._edge.target_id, message):
|
||||
route_result = await self._edge.should_route(message.data)
|
||||
|
||||
if route_result:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: True,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
|
||||
})
|
||||
should_execute = True
|
||||
target_id = self._edge.target_id
|
||||
source_id = self._edge.source_id
|
||||
else:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value,
|
||||
})
|
||||
# Return True here because message was processed, just condition failed
|
||||
return True
|
||||
else:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,
|
||||
})
|
||||
return False
|
||||
except Exception as e:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
|
||||
})
|
||||
raise e
|
||||
|
||||
# Execute outside the span
|
||||
if should_execute and target_id and source_id:
|
||||
await self._execute_on_target(target_id, [source_id], message, state, ctx)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
class FanOutEdgeRunner(EdgeRunner):
|
||||
"""Runner for fan-out edge groups."""
|
||||
|
||||
def __init__(self, edge_group: FanOutEdgeGroup, executors: dict[str, Executor]) -> None:
|
||||
super().__init__(edge_group, executors)
|
||||
self._edges = edge_group.edges
|
||||
self._target_ids = edge_group.target_executor_ids
|
||||
self._target_map = {edge.target_id: edge for edge in self._edges}
|
||||
self._selection_func = cast(
|
||||
Callable[[Any, list[str]], list[str]] | None, getattr(edge_group, "selection_func", None)
|
||||
)
|
||||
|
||||
async def send_message(
|
||||
self,
|
||||
message: WorkflowMessage,
|
||||
state: State,
|
||||
ctx: RunnerContext,
|
||||
) -> bool:
|
||||
"""Send a message through all edges in the fan-out edge group."""
|
||||
deliverable_edges: list[Edge] = []
|
||||
single_target_edge: Edge | None = None
|
||||
# Process routing logic within span
|
||||
with create_edge_group_processing_span(
|
||||
self._edge_group.__class__.__name__,
|
||||
edge_group_id=self._edge_group.id,
|
||||
message_source_id=message.source_id,
|
||||
message_target_id=message.target_id,
|
||||
source_trace_contexts=message.trace_contexts,
|
||||
source_span_ids=message.source_span_ids,
|
||||
) as span:
|
||||
try:
|
||||
selection_results = (
|
||||
self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids
|
||||
)
|
||||
if not self._validate_selection_result(selection_results):
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
|
||||
})
|
||||
raise RuntimeError(
|
||||
f"Invalid selection result: {selection_results}. "
|
||||
f"Expected selections to be a subset of valid target executor IDs: {self._target_ids}."
|
||||
)
|
||||
|
||||
if message.target_id:
|
||||
# If the target ID is specified and the selection result contains it, send the message to that edge
|
||||
if message.target_id in selection_results:
|
||||
edge = self._target_map.get(message.target_id)
|
||||
if edge and self._can_handle(edge.target_id, message):
|
||||
route_result = await edge.should_route(message.data)
|
||||
|
||||
if route_result:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: True,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
|
||||
})
|
||||
single_target_edge = edge
|
||||
else:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value, # noqa: E501
|
||||
})
|
||||
# For targeted messages with condition failure, return True (message was processed)
|
||||
return True
|
||||
else:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value, # noqa: E501
|
||||
})
|
||||
# For targeted messages that can't be handled, return False
|
||||
return False
|
||||
else:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value,
|
||||
})
|
||||
# For targeted messages not in selection, return False
|
||||
return False
|
||||
else:
|
||||
# If no target ID, send the message to the selected targets
|
||||
for target_id in selection_results:
|
||||
edge = self._target_map[target_id]
|
||||
if self._can_handle(edge.target_id, message):
|
||||
route_result = await edge.should_route(message.data)
|
||||
if route_result:
|
||||
deliverable_edges.append(edge)
|
||||
|
||||
if len(deliverable_edges) > 0:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: True,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
|
||||
})
|
||||
else:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
|
||||
})
|
||||
raise e
|
||||
|
||||
# Execute outside the span
|
||||
if single_target_edge:
|
||||
await self._execute_on_target(
|
||||
single_target_edge.target_id,
|
||||
[single_target_edge.source_id],
|
||||
message,
|
||||
state,
|
||||
ctx,
|
||||
)
|
||||
return True
|
||||
|
||||
if deliverable_edges:
|
||||
|
||||
async def send_to_edge(edge: Edge) -> bool:
|
||||
await self._execute_on_target(edge.target_id, [edge.source_id], message, state, ctx)
|
||||
return True
|
||||
|
||||
tasks = [send_to_edge(edge) for edge in deliverable_edges]
|
||||
results = await asyncio.gather(*tasks)
|
||||
return any(results)
|
||||
|
||||
# If we get here, it's a broadcast message with no deliverable edges
|
||||
return False
|
||||
|
||||
def _validate_selection_result(self, selection_results: list[str]) -> bool:
|
||||
"""Validate the selection results to ensure all IDs are valid target executor IDs."""
|
||||
return all(result in self._target_ids for result in selection_results)
|
||||
|
||||
|
||||
class FanInEdgeRunner(EdgeRunner):
|
||||
"""Runner for fan-in edge groups."""
|
||||
|
||||
def __init__(self, edge_group: FanInEdgeGroup, executors: dict[str, Executor]) -> None:
|
||||
super().__init__(edge_group, executors)
|
||||
self._edges = edge_group.edges
|
||||
# Buffer to hold messages before sending them to the target executor
|
||||
# Key is the source executor ID, value is a list of messages
|
||||
self._buffer: dict[str, list[WorkflowMessage]] = defaultdict(list)
|
||||
|
||||
async def send_message(
|
||||
self,
|
||||
message: WorkflowMessage,
|
||||
state: State,
|
||||
ctx: RunnerContext,
|
||||
) -> bool:
|
||||
"""Send a message through all edges in the fan-in edge group."""
|
||||
execution_data: dict[str, Any] | None = None
|
||||
with create_edge_group_processing_span(
|
||||
self._edge_group.__class__.__name__,
|
||||
edge_group_id=self._edge_group.id,
|
||||
message_source_id=message.source_id,
|
||||
message_target_id=message.target_id,
|
||||
source_trace_contexts=message.trace_contexts,
|
||||
source_span_ids=message.source_span_ids,
|
||||
) as span:
|
||||
try:
|
||||
if message.target_id and message.target_id != self._edges[0].target_id:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value,
|
||||
})
|
||||
return False
|
||||
|
||||
# Check if target can handle list of message data (fan-in aggregates multiple messages)
|
||||
if self._can_handle(
|
||||
self._edges[0].target_id, WorkflowMessage(data=[message.data], source_id=message.source_id)
|
||||
):
|
||||
# If the edge can handle the data, buffer the message
|
||||
self._buffer[message.source_id].append(message)
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: True,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.BUFFERED.value,
|
||||
})
|
||||
else:
|
||||
# If the edge cannot handle the data, return False
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,
|
||||
})
|
||||
return False
|
||||
|
||||
if self._is_ready_to_send():
|
||||
# If all edges in the group have data, prepare for execution
|
||||
messages_to_send = [msg for edge in self._edges for msg in self._buffer[edge.source_id]]
|
||||
self._buffer.clear()
|
||||
# Send aggregated data to target
|
||||
aggregated_data = [msg.data for msg in messages_to_send]
|
||||
|
||||
# Collect all trace contexts and source span IDs for fan-in linking
|
||||
trace_contexts = [msg.trace_context for msg in messages_to_send if msg.trace_context]
|
||||
source_span_ids = [msg.source_span_id for msg in messages_to_send if msg.source_span_id]
|
||||
|
||||
# Create a new Message object for the aggregated data
|
||||
aggregated_message = WorkflowMessage(
|
||||
data=aggregated_data,
|
||||
source_id=self._edge_group.__class__.__name__, # This won't be used in self._execute_on_target.
|
||||
trace_contexts=trace_contexts,
|
||||
source_span_ids=source_span_ids,
|
||||
)
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: True,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
|
||||
})
|
||||
|
||||
# Store execution data for later
|
||||
execution_data = {
|
||||
"target_id": self._edges[0].target_id,
|
||||
"source_ids": [edge.source_id for edge in self._edges],
|
||||
"message": aggregated_message,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
span.set_attributes({
|
||||
OtelAttr.EDGE_GROUP_DELIVERED: False,
|
||||
OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
|
||||
})
|
||||
raise e
|
||||
|
||||
# Execute outside the span if needed
|
||||
if execution_data:
|
||||
await self._execute_on_target(
|
||||
execution_data["target_id"],
|
||||
execution_data["source_ids"],
|
||||
execution_data["message"],
|
||||
state,
|
||||
ctx,
|
||||
)
|
||||
return True
|
||||
|
||||
return True # Return True for buffered messages (waiting for more)
|
||||
|
||||
def _is_ready_to_send(self) -> bool:
|
||||
"""Check if all edges in the group have data to send."""
|
||||
return all(self._buffer[edge.source_id] for edge in self._edges)
|
||||
|
||||
|
||||
class SwitchCaseEdgeRunner(FanOutEdgeRunner):
|
||||
"""Runner for switch-case edge groups (inherits from FanOutEdgeRunner)."""
|
||||
|
||||
def __init__(self, edge_group: SwitchCaseEdgeGroup, executors: dict[str, Executor]) -> None:
|
||||
super().__init__(edge_group, executors)
|
||||
|
||||
|
||||
def create_edge_runner(edge_group: EdgeGroup, executors: dict[str, Executor]) -> EdgeRunner:
|
||||
"""Factory function to create the appropriate edge runner for an edge group.
|
||||
|
||||
Args:
|
||||
edge_group: The edge group to create a runner for.
|
||||
executors: Map of executor IDs to executor instances.
|
||||
|
||||
Returns:
|
||||
The appropriate EdgeRunner instance.
|
||||
"""
|
||||
if isinstance(edge_group, (SingleEdgeGroup, InternalEdgeGroup)):
|
||||
return SingleEdgeRunner(edge_group, executors)
|
||||
if isinstance(edge_group, SwitchCaseEdgeGroup):
|
||||
return SwitchCaseEdgeRunner(edge_group, executors)
|
||||
if isinstance(edge_group, FanOutEdgeGroup):
|
||||
return FanOutEdgeRunner(edge_group, executors)
|
||||
if isinstance(edge_group, FanInEdgeGroup):
|
||||
return FanInEdgeRunner(edge_group, executors)
|
||||
raise ValueError(f"Unsupported edge group type: {type(edge_group)}")
|
||||
@@ -0,0 +1,448 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import sys
|
||||
import traceback as _traceback
|
||||
import warnings
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any, Generic, Literal, cast
|
||||
|
||||
from ._typing_utils import deserialize_type, serialize_type
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # pragma: no cover
|
||||
|
||||
DataT = TypeVar("DataT", default=Any)
|
||||
|
||||
|
||||
class WorkflowEventSource(str, Enum):
|
||||
"""Identifies whether a workflow event came from the framework or an executor.
|
||||
|
||||
Use `FRAMEWORK` for events emitted by built-in orchestration paths—even when the
|
||||
code that raises them lives in runner-related modules—and `EXECUTOR` for events
|
||||
surfaced by developer-provided executor implementations.
|
||||
"""
|
||||
|
||||
FRAMEWORK = "FRAMEWORK" # Framework-owned orchestration, regardless of module location
|
||||
EXECUTOR = "EXECUTOR" # User-supplied executor code and callbacks
|
||||
|
||||
|
||||
_event_origin_context: ContextVar[WorkflowEventSource] = ContextVar(
|
||||
"workflow_event_origin", default=WorkflowEventSource.EXECUTOR
|
||||
)
|
||||
|
||||
|
||||
def _current_event_origin() -> WorkflowEventSource:
|
||||
"""Return the origin to associate with newly created workflow events."""
|
||||
return _event_origin_context.get()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _framework_event_origin() -> Generator[None]: # pyright: ignore[reportUnusedFunction]
|
||||
"""Temporarily mark subsequently created events as originating from the framework (internal)."""
|
||||
token = _event_origin_context.set(WorkflowEventSource.FRAMEWORK)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_event_origin_context.reset(token)
|
||||
|
||||
|
||||
class WorkflowRunState(str, Enum):
|
||||
"""Run-level state of a workflow execution."""
|
||||
|
||||
STARTED = "STARTED"
|
||||
IN_PROGRESS = "IN_PROGRESS"
|
||||
IN_PROGRESS_PENDING_REQUESTS = "IN_PROGRESS_PENDING_REQUESTS"
|
||||
IDLE = "IDLE"
|
||||
IDLE_WITH_PENDING_REQUESTS = "IDLE_WITH_PENDING_REQUESTS"
|
||||
FAILED = "FAILED"
|
||||
CANCELLED = "CANCELLED"
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkflowErrorDetails:
|
||||
"""Structured error information to surface in error events/results."""
|
||||
|
||||
error_type: str
|
||||
message: str
|
||||
traceback: str | None = None
|
||||
executor_id: str | None = None
|
||||
extra: dict[str, Any] | None = None
|
||||
|
||||
@classmethod
|
||||
def from_exception(
|
||||
cls,
|
||||
exc: BaseException,
|
||||
*,
|
||||
executor_id: str | None = None,
|
||||
extra: dict[str, Any] | None = None,
|
||||
) -> WorkflowErrorDetails:
|
||||
tb = None
|
||||
try:
|
||||
tb = "".join(_traceback.format_exception(type(exc), exc, exc.__traceback__))
|
||||
except Exception:
|
||||
tb = None
|
||||
return cls(
|
||||
error_type=exc.__class__.__name__,
|
||||
message=str(exc),
|
||||
traceback=tb,
|
||||
executor_id=executor_id,
|
||||
extra=extra,
|
||||
)
|
||||
|
||||
|
||||
# Type discriminator for workflow events.
|
||||
# Includes both framework lifecycle types and well-known orchestration types.
|
||||
WorkflowEventType = Literal[
|
||||
# Lifecycle events (workflow-level)
|
||||
"started", # Workflow run began
|
||||
"status", # Workflow state changed (use .state)
|
||||
"failed", # Workflow terminated with error (use .details)
|
||||
# Data events
|
||||
"output", # Executor yielded final terminal output (use .executor_id, .data)
|
||||
"intermediate", # Executor emitted intermediate (non-terminal) output (use .executor_id, .data)
|
||||
"data", # DEPRECATED — compatibility alias for intermediate emissions; use type='intermediate' instead.
|
||||
# Request events (human-in-the-loop)
|
||||
"request_info", # Executor requests external info (use .request_id, .source_executor_id)
|
||||
# Diagnostic events (warnings/errors from user code)
|
||||
"warning", # Warning from user code (use .data as str)
|
||||
"error", # Error from user code, non-fatal (use .data as Exception)
|
||||
# Iteration events (supersteps)
|
||||
"superstep_started", # Superstep began (use .iteration)
|
||||
"superstep_completed", # Superstep ended (use .iteration)
|
||||
# Executor lifecycle events
|
||||
"executor_invoked", # Executor handler was called (use .executor_id, .data)
|
||||
"executor_completed", # Executor handler completed (use .executor_id, .data)
|
||||
"executor_failed", # Executor handler raised error (use .executor_id, .details)
|
||||
"executor_bypassed", # Executor skipped via cache hit during replay (use .executor_id, .data)
|
||||
# Orchestration event types (use .data for typed payload)
|
||||
"group_chat", # Group chat orchestrator events (use .data as GroupChatRequestSentEvent | GroupChatResponseReceivedEvent) # noqa: E501
|
||||
"handoff_sent", # Handoff routing events (use .data as HandoffSentEvent)
|
||||
"magentic_orchestrator", # Magentic orchestrator events (use .data as MagenticOrchestratorEvent)
|
||||
]
|
||||
|
||||
|
||||
# Event types forwarded across the ``workflow.as_agent()`` boundary. Anything not
|
||||
# in this set — lifecycle events, diagnostics, executor bookkeeping, and
|
||||
# orchestration-internal events (``group_chat``, ``handoff_sent``,
|
||||
# ``magentic_orchestrator``) — stays inside the workflow and is not surfaced to
|
||||
# agent callers. Internal to the ``_workflows`` package.
|
||||
AGENT_FORWARDED_EVENT_TYPES: frozenset[str] = frozenset({
|
||||
"output",
|
||||
"intermediate",
|
||||
"data", # deprecated alias for intermediate; retained for backward compat
|
||||
"request_info",
|
||||
})
|
||||
|
||||
|
||||
class WorkflowEvent(Generic[DataT]):
|
||||
"""Unified event for all workflow emissions.
|
||||
|
||||
This single generic class handles all workflow events through a `type` discriminator,
|
||||
following the same pattern as the `Content` class.
|
||||
|
||||
Use factory methods for convenient construction of lifecycle, diagnostic, request,
|
||||
and executor bookkeeping events. Workflow ``output`` and ``intermediate`` events
|
||||
are emitted by ``ctx.yield_output(...)`` based on workflow output selection.
|
||||
|
||||
- `WorkflowEvent.started()` - workflow run began
|
||||
- `WorkflowEvent.status(state)` - workflow state changed
|
||||
- `WorkflowEvent.failed(details)` - workflow terminated with error
|
||||
- `WorkflowEvent.warning(message)` - warning from user code
|
||||
- `WorkflowEvent.error(exception)` - error from user code
|
||||
- `WorkflowEvent.request_info(...)` - executor requests external info
|
||||
- `WorkflowEvent.superstep_started(iteration)` - superstep began
|
||||
- `WorkflowEvent.superstep_completed(iteration)` - superstep ended
|
||||
- `WorkflowEvent.executor_invoked(executor_id)` - executor handler called
|
||||
- `WorkflowEvent.executor_completed(executor_id)` - executor handler completed
|
||||
- `WorkflowEvent.executor_failed(executor_id, details)` - executor handler failed
|
||||
- `WorkflowEvent.executor_bypassed(executor_id)` - executor skipped via cache hit
|
||||
|
||||
The generic parameter DataT represents the type of the event's data payload:
|
||||
- Lifecycle events: `WorkflowEvent[None]` (data is None)
|
||||
- Data events: `WorkflowEvent[DataT]` where DataT is the payload type (e.g., AgentResponse)
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
# Create lifecycle events via factory methods
|
||||
started = WorkflowEvent.started()
|
||||
status = WorkflowEvent.status(WorkflowRunState.IN_PROGRESS)
|
||||
|
||||
# Type-safe access to event data
|
||||
event: WorkflowEvent[AgentResponse] = WorkflowEvent("data", executor_id="agent1", data=response)
|
||||
data: AgentResponse = event.data
|
||||
|
||||
# Check event type
|
||||
if event.type == "status":
|
||||
print(f"State: {event.state}")
|
||||
elif event.type == "output":
|
||||
print(f"Output from {event.executor_id}: {event.data}")
|
||||
elif event.type == "data":
|
||||
if isinstance(event.data, AgentResponse):
|
||||
print(f"Agent response: {event.data.text}")
|
||||
"""
|
||||
|
||||
type: WorkflowEventType
|
||||
data: DataT
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
type: WorkflowEventType,
|
||||
data: DataT | None = None,
|
||||
*,
|
||||
# Event context fields
|
||||
origin: WorkflowEventSource | None = None,
|
||||
# STATUS event fields
|
||||
state: WorkflowRunState | None = None,
|
||||
# FAILED event fields
|
||||
details: WorkflowErrorDetails | None = None,
|
||||
# OUTPUT/DATA event fields
|
||||
executor_id: str | None = None,
|
||||
# REQUEST_INFO event fields
|
||||
request_id: str | None = None,
|
||||
source_executor_id: str | None = None,
|
||||
request_type: builtins.type[Any] | None = None,
|
||||
response_type: builtins.type[Any] | None = None,
|
||||
# SUPERSTEP event fields
|
||||
iteration: int | None = None,
|
||||
) -> None:
|
||||
"""Initialize the workflow event.
|
||||
|
||||
Prefer using factory methods like `WorkflowEvent.started()` instead of __init__ directly.
|
||||
"""
|
||||
self.type = type
|
||||
self.data = data # type: ignore[assignment]
|
||||
self.origin = origin if origin is not None else _current_event_origin()
|
||||
|
||||
# Event-specific fields
|
||||
self.state = state
|
||||
self.details = details
|
||||
self.executor_id = executor_id
|
||||
self._request_id = request_id
|
||||
self._source_executor_id = source_executor_id
|
||||
self._request_type = request_type
|
||||
self._response_type = response_type
|
||||
self.iteration = iteration
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Return a string representation of the workflow event."""
|
||||
parts = [f"type={self.type!r}"]
|
||||
if self.state is not None:
|
||||
parts.append(f"state={self.state.value}")
|
||||
if self.executor_id is not None:
|
||||
parts.append(f"executor_id={self.executor_id!r}")
|
||||
if self.iteration is not None:
|
||||
parts.append(f"iteration={self.iteration}")
|
||||
if self._request_id is not None:
|
||||
parts.append(f"request_id={self._request_id!r}")
|
||||
if self.data is not None:
|
||||
parts.append(f"data={self.data!r}")
|
||||
return f"WorkflowEvent({', '.join(parts)})" # pragma: no cover
|
||||
|
||||
# ==========================================================================
|
||||
# Factory methods
|
||||
# ==========================================================================
|
||||
|
||||
@classmethod
|
||||
def started(cls, data: DataT | None = None) -> WorkflowEvent[DataT]:
|
||||
"""Create a 'started' event when a workflow run begins."""
|
||||
return cls("started", data=data)
|
||||
|
||||
@classmethod
|
||||
def status(cls, state: WorkflowRunState, data: DataT | None = None) -> WorkflowEvent[DataT]:
|
||||
"""Create a 'status' event for workflow state transitions."""
|
||||
return cls("status", data=data, state=state)
|
||||
|
||||
@classmethod
|
||||
def failed(cls, details: WorkflowErrorDetails, data: DataT | None = None) -> WorkflowEvent[DataT]:
|
||||
"""Create a 'failed' event when a workflow terminates with error."""
|
||||
return cls("failed", data=data, details=details)
|
||||
|
||||
@classmethod
|
||||
def warning(cls, message: str) -> WorkflowEvent[str]:
|
||||
"""Create a 'warning' event from user code."""
|
||||
return WorkflowEvent("warning", data=message)
|
||||
|
||||
@classmethod
|
||||
def error(cls, exception: Exception) -> WorkflowEvent[Exception]:
|
||||
"""Create an 'error' event from user code."""
|
||||
return WorkflowEvent("error", data=exception)
|
||||
|
||||
@classmethod
|
||||
def emit(cls, executor_id: str, data: DataT) -> WorkflowEvent[DataT]:
|
||||
"""Create a 'data' event (deprecated alias for intermediate emissions).
|
||||
|
||||
.. deprecated::
|
||||
Use ``ctx.yield_output(...)`` and configure ``intermediate_output_from`` instead.
|
||||
Will be removed in a future major release along with the ``type='data'`` event variant.
|
||||
"""
|
||||
warnings.warn(
|
||||
"WorkflowEvent.emit() / type='data' are deprecated; use ctx.yield_output() from an "
|
||||
"intermediate-designated executor. Will be removed in a future major release.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
return cls("data", executor_id=executor_id, data=data)
|
||||
|
||||
@classmethod
|
||||
def request_info(
|
||||
cls,
|
||||
request_id: str,
|
||||
source_executor_id: str,
|
||||
request_data: DataT,
|
||||
response_type: builtins.type[Any],
|
||||
) -> WorkflowEvent[DataT]:
|
||||
"""Create a 'request_info' event when an executor requests external information."""
|
||||
return cls(
|
||||
"request_info",
|
||||
data=request_data,
|
||||
request_id=request_id,
|
||||
source_executor_id=source_executor_id,
|
||||
request_type=type(request_data),
|
||||
response_type=response_type,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def superstep_started(cls, iteration: int, data: DataT | None = None) -> WorkflowEvent[DataT]:
|
||||
"""Create a 'superstep_started' event when a superstep begins."""
|
||||
return cls("superstep_started", iteration=iteration, data=data)
|
||||
|
||||
@classmethod
|
||||
def superstep_completed(cls, iteration: int, data: DataT | None = None) -> WorkflowEvent[DataT]:
|
||||
"""Create a 'superstep_completed' event when a superstep ends."""
|
||||
return cls("superstep_completed", iteration=iteration, data=data)
|
||||
|
||||
@classmethod
|
||||
def executor_invoked(cls, executor_id: str, data: DataT | None = None) -> WorkflowEvent[DataT]:
|
||||
"""Create an 'executor_invoked' event when an executor handler is called."""
|
||||
return cls("executor_invoked", executor_id=executor_id, data=data)
|
||||
|
||||
@classmethod
|
||||
def executor_completed(cls, executor_id: str, data: DataT | None = None) -> WorkflowEvent[DataT]:
|
||||
"""Create an 'executor_completed' event when an executor handler completes."""
|
||||
return cls("executor_completed", executor_id=executor_id, data=data)
|
||||
|
||||
@classmethod
|
||||
def executor_failed(cls, executor_id: str, details: WorkflowErrorDetails) -> WorkflowEvent[WorkflowErrorDetails]:
|
||||
"""Create an 'executor_failed' event when an executor handler raises an error."""
|
||||
return WorkflowEvent("executor_failed", executor_id=executor_id, data=details, details=details)
|
||||
|
||||
@classmethod
|
||||
def executor_bypassed(cls, executor_id: str, data: DataT | None = None) -> WorkflowEvent[DataT]:
|
||||
"""Create an 'executor_bypassed' event when a step is skipped via cache hit during replay."""
|
||||
return cls("executor_bypassed", executor_id=executor_id, data=data)
|
||||
|
||||
# ==========================================================================
|
||||
# Property for type-safe access
|
||||
# ==========================================================================
|
||||
|
||||
@property
|
||||
def request_id(self) -> str:
|
||||
"""Get request_id for request_info events.
|
||||
|
||||
Returns:
|
||||
The request ID as a non-None string.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called on an event that is not a request_info event,
|
||||
or if the event is malformed (request_info without request_id).
|
||||
"""
|
||||
if self.type != "request_info" or self._request_id is None:
|
||||
raise RuntimeError(f"request_id is only available for request_info events, got type={self.type!r}")
|
||||
return self._request_id
|
||||
|
||||
@property
|
||||
def source_executor_id(self) -> str:
|
||||
"""Get source_executor_id for request_info events.
|
||||
|
||||
Returns:
|
||||
The source executor ID as a non-None string.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called on an event that is not a request_info event,
|
||||
or if the event is malformed (request_info without source_executor_id).
|
||||
"""
|
||||
if self.type != "request_info" or self._source_executor_id is None:
|
||||
raise RuntimeError(f"source_executor_id is only available for request_info events, got type={self.type!r}")
|
||||
return self._source_executor_id
|
||||
|
||||
@property
|
||||
def request_type(self) -> builtins.type[Any]:
|
||||
"""Get request_type for request_info events.
|
||||
|
||||
Returns:
|
||||
The request data type as a non-None type object.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called on an event that is not a request_info event,
|
||||
or if the event is malformed (request_info without request_type).
|
||||
"""
|
||||
if self.type != "request_info" or self._request_type is None:
|
||||
raise RuntimeError(f"request_type is only available for request_info events, got type={self.type!r}")
|
||||
return self._request_type
|
||||
|
||||
@property
|
||||
def response_type(self) -> builtins.type[Any]:
|
||||
"""Get response_type for request_info events.
|
||||
|
||||
Returns:
|
||||
The response data type as a non-None type object.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called on an event that is not a request_info event,
|
||||
or if the event is malformed (request_info without response_type).
|
||||
"""
|
||||
if self.type != "request_info" or self._response_type is None:
|
||||
raise RuntimeError(f"response_type is only available for request_info events, got type={self.type!r}")
|
||||
return self._response_type
|
||||
|
||||
# ==========================================================================
|
||||
# Serialization methods (primarily for REQUEST_INFO events)
|
||||
# ==========================================================================
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dictionary for serialization.
|
||||
|
||||
Currently only implemented for 'request_info' events for checkpoint storage.
|
||||
"""
|
||||
if self.type != "request_info":
|
||||
raise ValueError(f"to_dict() only supported for 'request_info' events, got '{self.type}'")
|
||||
return {
|
||||
"type": self.type,
|
||||
"data": self.data,
|
||||
"request_id": self._request_id,
|
||||
"source_executor_id": self._source_executor_id,
|
||||
"request_type": serialize_type(self._request_type) if self._request_type else None,
|
||||
"response_type": serialize_type(self._response_type) if self._response_type else None,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> WorkflowEvent[Any]:
|
||||
"""Create a REQUEST_INFO event from a dictionary."""
|
||||
for prop in ["data", "request_id", "source_executor_id", "request_type", "response_type"]:
|
||||
if prop not in data:
|
||||
raise KeyError(f"Missing '{prop}' field in WorkflowEvent dictionary.")
|
||||
|
||||
request_data = data["data"]
|
||||
request_type = deserialize_type(data["request_type"])
|
||||
|
||||
if request_type is not type(request_data):
|
||||
raise TypeError(
|
||||
"Mismatch between deserialized request_data type and request_type field in WorkflowEvent dictionary."
|
||||
)
|
||||
|
||||
return cls.request_info(
|
||||
request_id=data["request_id"],
|
||||
source_executor_id=data["source_executor_id"],
|
||||
request_data=cast(Any, request_data), # type: ignore
|
||||
response_type=deserialize_type(data["response_type"]),
|
||||
)
|
||||
@@ -0,0 +1,814 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import copy
|
||||
import functools
|
||||
import inspect
|
||||
import logging
|
||||
import types
|
||||
import typing
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, TypeVar, overload
|
||||
|
||||
from ..observability import create_processing_span
|
||||
from ._events import (
|
||||
WorkflowErrorDetails,
|
||||
WorkflowEvent,
|
||||
_framework_event_origin, # type: ignore[reportPrivateUsage]
|
||||
)
|
||||
from ._model_utils import DictConvertible
|
||||
from ._request_info_mixin import RequestInfoMixin
|
||||
from ._runner_context import MessageType, RunnerContext, WorkflowMessage
|
||||
from ._state import State
|
||||
from ._typing_utils import contains_typevar, is_instance_of, normalize_type_to_list, resolve_type_annotation
|
||||
from ._workflow_context import WorkflowContext, validate_workflow_context_annotation
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# region Executor
|
||||
class Executor(RequestInfoMixin, DictConvertible):
|
||||
"""Base class for all workflow executors that process messages and perform computations.
|
||||
|
||||
## Overview
|
||||
Executors are the fundamental building blocks of workflows, representing individual processing
|
||||
units that receive messages, perform operations, and produce outputs. Each executor is uniquely
|
||||
identified and can handle specific message types through decorated handler methods.
|
||||
|
||||
## Type System
|
||||
Executors have a rich type system that defines their capabilities:
|
||||
|
||||
### Input Types
|
||||
The types of messages an executor can process, discovered from handler method signatures:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class MyExecutor(Executor):
|
||||
@handler
|
||||
async def handle_string(self, message: str, ctx: WorkflowContext) -> None:
|
||||
# This executor can handle 'str' input types
|
||||
Access via the `input_types` property.
|
||||
|
||||
### Output Types
|
||||
The types of messages an executor can send to other executors via `ctx.send_message()`:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class MyExecutor(Executor):
|
||||
@handler
|
||||
async def handle_data(self, message: str, ctx: WorkflowContext[int | bool]) -> None:
|
||||
# This executor can send 'int' or 'bool' messages
|
||||
Access via the `output_types` property.
|
||||
|
||||
### Workflow Output Types
|
||||
The types of data an executor can emit as workflow-level outputs via `ctx.yield_output()`:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class MyExecutor(Executor):
|
||||
@handler
|
||||
async def process(self, message: str, ctx: WorkflowContext[int, str]) -> None:
|
||||
# Can send 'int' messages AND yield 'str' workflow outputs
|
||||
Access via the `workflow_output_types` property.
|
||||
|
||||
## Handler Discovery
|
||||
Executors discover their capabilities through decorated methods:
|
||||
|
||||
### @handler Decorator
|
||||
Marks methods that process incoming messages:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class MyExecutor(Executor):
|
||||
@handler
|
||||
async def handle_text(self, message: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(message.upper())
|
||||
|
||||
### Sub-workflow Request Interception
|
||||
Use @handler methods to intercept sub-workflow requests:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class ParentExecutor(Executor):
|
||||
@handler
|
||||
async def handle_subworkflow_request(
|
||||
self,
|
||||
request: SubWorkflowRequestMessage,
|
||||
ctx: WorkflowContext[SubWorkflowResponseMessage],
|
||||
) -> None:
|
||||
if self.is_allowed(request.domain):
|
||||
response = request.create_response(data=True)
|
||||
await ctx.send_message(response, target_id=request.executor_id)
|
||||
else:
|
||||
await ctx.request_info(request.source_event, response_type=request.source_event.response_type)
|
||||
|
||||
## Context Types
|
||||
Handler methods receive different WorkflowContext variants based on their type annotations:
|
||||
|
||||
### WorkflowContext (no type parameters)
|
||||
For handlers that only perform side effects without sending messages or yielding outputs:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class LoggingExecutor(Executor):
|
||||
@handler
|
||||
async def log_message(self, msg: str, ctx: WorkflowContext) -> None:
|
||||
print(f"Received: {msg}") # Only logging, no outputs
|
||||
|
||||
### WorkflowContext[OutT]
|
||||
Enables sending messages of type OutT via `ctx.send_message()`:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class ProcessorExecutor(Executor):
|
||||
@handler
|
||||
async def handler(self, msg: str, ctx: WorkflowContext[int]) -> None:
|
||||
await ctx.send_message(42) # Can send int messages
|
||||
|
||||
### WorkflowContext[OutT, W_OutT]
|
||||
Enables both sending messages (OutT) and yielding workflow outputs (W_OutT):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class DualOutputExecutor(Executor):
|
||||
@handler
|
||||
async def handler(self, msg: str, ctx: WorkflowContext[int, str]) -> None:
|
||||
await ctx.send_message(42) # Send int message
|
||||
await ctx.yield_output("done") # Yield str workflow output
|
||||
|
||||
## Function Executors
|
||||
Simple functions can be converted to executors using the `@executor` decorator:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@executor
|
||||
async def process_text(text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(text.upper())
|
||||
|
||||
|
||||
# Or with custom ID:
|
||||
@executor(id="text_processor")
|
||||
def sync_process(text: str, ctx: WorkflowContext[str]) -> None:
|
||||
ctx.send_message(text.lower()) # Sync functions run in thread pool
|
||||
|
||||
## Sub-workflow Composition
|
||||
Executors can contain sub-workflows using WorkflowExecutor. Sub-workflows can make requests
|
||||
that parent workflows can intercept. See WorkflowExecutor documentation for details on
|
||||
workflow composition patterns and request/response handling.
|
||||
|
||||
## State Management
|
||||
Executors can contain states that persist across workflow runs and checkpoints. Override the
|
||||
`on_checkpoint_save` and `on_checkpoint_restore` methods to implement custom state
|
||||
serialization and restoration logic.
|
||||
|
||||
## Implementation Notes
|
||||
- Do not call `execute()` directly - it's invoked by the workflow engine
|
||||
- Do not override `execute()` - define handlers using decorators instead
|
||||
- Each executor must have at least one `@handler` method
|
||||
- Handler method signatures are validated at initialization time
|
||||
"""
|
||||
|
||||
# Provide a default so static analyzers (e.g., pyright) don't require passing `id`.
|
||||
# Runtime still sets a concrete value in __init__.
|
||||
def __init__(
|
||||
self,
|
||||
id: str,
|
||||
*,
|
||||
type: str | None = None,
|
||||
type_: str | None = None,
|
||||
defer_discovery: bool = False,
|
||||
**_: Any,
|
||||
) -> None:
|
||||
"""Initialize the executor with a unique identifier.
|
||||
|
||||
Args:
|
||||
id: A unique identifier for the executor.
|
||||
|
||||
Keyword Args:
|
||||
type: The executor type name. If not provided, uses class name.
|
||||
type_: Alternative parameter name for executor type.
|
||||
defer_discovery: If True, defer handler method discovery until later.
|
||||
**_: Additional keyword arguments. Unused in this implementation.
|
||||
"""
|
||||
if not id:
|
||||
raise ValueError("Executor ID must be a non-empty string.")
|
||||
|
||||
resolved_type = type or type_ or self.__class__.__name__
|
||||
self.id = id
|
||||
self.type = resolved_type
|
||||
self.type_ = resolved_type
|
||||
|
||||
# Serialize per-executor message processing. Within a superstep the runner may
|
||||
# dispatch deliveries to the same target executor from multiple sources
|
||||
# concurrently; this lock guarantees the executor processes them one at a time
|
||||
# (and, per source, in the order they were sent).
|
||||
#
|
||||
# The lock must be created lazily under the running loop (see ``_get_execution_lock``)
|
||||
# in order to support executors that are constructed outside an event loop and reused
|
||||
# across multiple async loops (e.g., successive ``asyncio.run`` calls on the same workflow).
|
||||
# Binding a single lock to one loop would raise "bound to a different event loop" on reuse.
|
||||
self._execution_lock: asyncio.Lock | None = None
|
||||
self._execution_lock_loop: asyncio.AbstractEventLoop | None = None
|
||||
|
||||
from builtins import type as builtin_type
|
||||
|
||||
self._handlers: dict[
|
||||
builtin_type[Any] | types.UnionType, Callable[[Any, WorkflowContext[Any, Any]], Awaitable[None]]
|
||||
] = {}
|
||||
self._handler_specs: list[dict[str, Any]] = []
|
||||
if not defer_discovery:
|
||||
self._discover_handlers()
|
||||
|
||||
if not self._handlers:
|
||||
raise ValueError(
|
||||
f"Executor {self.__class__.__name__} has no handlers defined. "
|
||||
"Please define at least one handler using the @handler decorator."
|
||||
)
|
||||
|
||||
# Initialize RequestInfoMixin to discover response handlers
|
||||
self._discover_response_handlers()
|
||||
|
||||
def _get_execution_lock(self) -> asyncio.Lock:
|
||||
"""Return this executor's serialization lock, bound to the running event loop.
|
||||
|
||||
The lock is created lazily and re-created if the running loop has changed since it
|
||||
was last used (for example, the executor is reused across successive
|
||||
``asyncio.run`` calls), avoiding ``asyncio.Lock`` "bound to a different event loop"
|
||||
errors. Must be called from within a running loop.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
if self._execution_lock is None or self._execution_lock_loop is not loop:
|
||||
self._execution_lock = asyncio.Lock()
|
||||
self._execution_lock_loop = loop
|
||||
return self._execution_lock
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
message: Any,
|
||||
source_executor_ids: list[str],
|
||||
state: State,
|
||||
runner_context: RunnerContext,
|
||||
trace_contexts: list[dict[str, str]] | None = None,
|
||||
source_span_ids: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Execute the executor with a given message and context parameters.
|
||||
|
||||
- Do not call this method directly - it is invoked by the workflow engine.
|
||||
- Do not override this method. Instead, define handlers using @handler decorator.
|
||||
|
||||
Args:
|
||||
message: The message to be processed by the executor.
|
||||
source_executor_ids: The IDs of the source executors that sent messages to this executor.
|
||||
state: The state for the workflow.
|
||||
runner_context: The runner context that provides methods to send messages and events.
|
||||
trace_contexts: Optional trace contexts from multiple sources for OpenTelemetry propagation.
|
||||
source_span_ids: Optional source span IDs from multiple sources for linking.
|
||||
|
||||
Returns:
|
||||
An awaitable that resolves to the result of the execution.
|
||||
"""
|
||||
async with self._get_execution_lock():
|
||||
# Create processing span for tracing (gracefully handles disabled tracing)
|
||||
with create_processing_span(
|
||||
self.id,
|
||||
self.__class__.__name__,
|
||||
str(MessageType.STANDARD if not isinstance(message, WorkflowMessage) else message.type),
|
||||
type(message).__name__,
|
||||
source_trace_contexts=trace_contexts,
|
||||
source_span_ids=source_span_ids,
|
||||
):
|
||||
# Find the handler and handler spec that matches the message type.
|
||||
handler = self._find_handler(message)
|
||||
|
||||
original_message = message
|
||||
if isinstance(message, WorkflowMessage):
|
||||
# Unwrap raw data for handler call
|
||||
message = message.data
|
||||
|
||||
# Create the appropriate WorkflowContext based on handler specs
|
||||
context = self._create_context_for_handler(
|
||||
source_executor_ids=source_executor_ids,
|
||||
state=state,
|
||||
runner_context=runner_context,
|
||||
trace_contexts=trace_contexts,
|
||||
source_span_ids=source_span_ids,
|
||||
request_id=original_message.original_request_info_event.request_id
|
||||
if isinstance(original_message, WorkflowMessage) and original_message.original_request_info_event
|
||||
else None,
|
||||
)
|
||||
|
||||
# Invoke the handler with the message and context
|
||||
# Use deepcopy to capture original input state before handler can mutate it
|
||||
with _framework_event_origin():
|
||||
invoke_event = WorkflowEvent.executor_invoked(self.id, copy.deepcopy(message))
|
||||
await context.add_event(invoke_event)
|
||||
try:
|
||||
await handler(message, context)
|
||||
except Exception as exc:
|
||||
# Surface structured executor failure before propagating
|
||||
with _framework_event_origin():
|
||||
failure_event = WorkflowEvent.executor_failed(self.id, WorkflowErrorDetails.from_exception(exc))
|
||||
await context.add_event(failure_event)
|
||||
raise
|
||||
with _framework_event_origin():
|
||||
# Include sent messages and yielded outputs as the completion data
|
||||
sent_messages = context.get_sent_messages()
|
||||
yielded_outputs = context.get_yielded_outputs()
|
||||
completion_data = sent_messages + yielded_outputs
|
||||
completed_event = WorkflowEvent.executor_completed(
|
||||
self.id, completion_data if completion_data else None
|
||||
)
|
||||
await context.add_event(completed_event)
|
||||
|
||||
def _create_context_for_handler(
|
||||
self,
|
||||
source_executor_ids: list[str],
|
||||
state: State,
|
||||
runner_context: RunnerContext,
|
||||
trace_contexts: list[dict[str, str]] | None = None,
|
||||
source_span_ids: list[str] | None = None,
|
||||
request_id: str | None = None,
|
||||
) -> WorkflowContext[Any]:
|
||||
"""Create the appropriate WorkflowContext based on the handler's context annotation.
|
||||
|
||||
Args:
|
||||
source_executor_ids: The IDs of the source executors that sent messages to this executor.
|
||||
state: The state for the workflow.
|
||||
runner_context: The runner context that provides methods to send messages and events.
|
||||
trace_contexts: Optional trace contexts from multiple sources for OpenTelemetry propagation.
|
||||
source_span_ids: Optional source span IDs from multiple sources for linking.
|
||||
request_id: Optional request ID if this context is for a `handle_response` handler.
|
||||
|
||||
Returns:
|
||||
WorkflowContext[Any] based on the handler's context annotation.
|
||||
"""
|
||||
# Create WorkflowContext
|
||||
return WorkflowContext(
|
||||
executor=self,
|
||||
source_executor_ids=source_executor_ids,
|
||||
state=state,
|
||||
runner_context=runner_context,
|
||||
trace_contexts=trace_contexts,
|
||||
source_span_ids=source_span_ids,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
def _discover_handlers(self) -> None:
|
||||
"""Discover message handlers in the executor class."""
|
||||
# Use __class__.__dict__ to avoid accessing pydantic's dynamic attributes
|
||||
for attr_name in dir(self.__class__):
|
||||
try:
|
||||
attr = getattr(self.__class__, attr_name)
|
||||
except AttributeError:
|
||||
# Skip attributes that may not be accessible (e.g., dynamic descriptors)
|
||||
continue
|
||||
|
||||
# Discover @handler methods
|
||||
if callable(attr) and hasattr(attr, "_handler_spec"):
|
||||
handler_spec = attr._handler_spec # type: ignore
|
||||
message_type = handler_spec["message_type"]
|
||||
|
||||
# Keep full generic types for handler registration to avoid conflicts
|
||||
if self._handlers.get(message_type) is not None:
|
||||
raise ValueError(f"Duplicate handler for type {message_type} in {self.__class__.__name__}")
|
||||
|
||||
# Get the bound method
|
||||
bound_method = getattr(self, attr_name)
|
||||
self._handlers[message_type] = bound_method
|
||||
|
||||
# Add to unified handler specs list
|
||||
self._handler_specs.append({**handler_spec})
|
||||
|
||||
def can_handle(self, message: WorkflowMessage) -> bool:
|
||||
"""Check if the executor can handle a given message type.
|
||||
|
||||
Args:
|
||||
message: The message to check.
|
||||
|
||||
Returns:
|
||||
True if the executor can handle the message type, False otherwise.
|
||||
"""
|
||||
if message.type == MessageType.RESPONSE:
|
||||
if message.original_request_info_event is None:
|
||||
logger.warning(
|
||||
f"Executor {self.__class__.__name__} received a response message without an original request event."
|
||||
)
|
||||
return False
|
||||
|
||||
return any(
|
||||
is_instance_of(message.original_request_info_event.data, message_type[0])
|
||||
and is_instance_of(message.data, message_type[1])
|
||||
for message_type in self._response_handlers
|
||||
)
|
||||
|
||||
return any(is_instance_of(message.data, message_type) for message_type in self._handlers)
|
||||
|
||||
def _register_instance_handler(
|
||||
self,
|
||||
name: str,
|
||||
func: Callable[[Any, WorkflowContext[Any]], Awaitable[Any]],
|
||||
message_type: type | types.UnionType,
|
||||
ctx_annotation: Any,
|
||||
output_types: list[type[Any] | types.UnionType],
|
||||
workflow_output_types: list[type[Any] | types.UnionType],
|
||||
) -> None:
|
||||
"""Register a handler at instance level.
|
||||
|
||||
Args:
|
||||
name: Name of the handler function for error reporting
|
||||
func: The async handler function to register
|
||||
message_type: Type of message this handler processes
|
||||
ctx_annotation: The WorkflowContext[T] annotation from the function
|
||||
output_types: List of output types for send_message()
|
||||
workflow_output_types: List of workflow output types for yield_output()
|
||||
"""
|
||||
if message_type in self._handlers:
|
||||
raise ValueError(f"Handler for type {message_type} already registered in {self.__class__.__name__}")
|
||||
|
||||
self._handlers[message_type] = func
|
||||
self._handler_specs.append({
|
||||
"name": name,
|
||||
"message_type": message_type,
|
||||
"ctx_annotation": ctx_annotation,
|
||||
"output_types": output_types,
|
||||
"workflow_output_types": workflow_output_types,
|
||||
})
|
||||
|
||||
@property
|
||||
def input_types(self) -> list[type[Any] | types.UnionType]:
|
||||
"""Get the list of input types that this executor can handle.
|
||||
|
||||
Returns:
|
||||
A list of the message types that this executor's handlers can process.
|
||||
"""
|
||||
return list(self._handlers.keys())
|
||||
|
||||
@property
|
||||
def output_types(self) -> list[type[Any] | types.UnionType]:
|
||||
"""Get the list of output types that this executor can produce via send_message().
|
||||
|
||||
Returns:
|
||||
A list of the output types inferred from the handlers' WorkflowContext[T] annotations.
|
||||
"""
|
||||
output_types: set[type[Any] | types.UnionType] = set()
|
||||
|
||||
# Collect output types from all handlers
|
||||
for handler_spec in self._handler_specs + self._response_handler_specs:
|
||||
handler_output_types = handler_spec.get("output_types", [])
|
||||
output_types.update(handler_output_types)
|
||||
|
||||
return list(output_types)
|
||||
|
||||
@property
|
||||
def workflow_output_types(self) -> list[type[Any] | types.UnionType]:
|
||||
"""Get the list of workflow output types that this executor can produce via yield_output().
|
||||
|
||||
Returns:
|
||||
A list of the workflow output types inferred from handlers' WorkflowContext[T, U] annotations.
|
||||
"""
|
||||
output_types: set[type[Any] | types.UnionType] = set()
|
||||
|
||||
# Collect workflow output types from all handlers
|
||||
for handler_spec in self._handler_specs + self._response_handler_specs:
|
||||
handler_workflow_output_types = handler_spec.get("workflow_output_types", [])
|
||||
output_types.update(handler_workflow_output_types)
|
||||
|
||||
return list(output_types)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialize executor definition for workflow topology export."""
|
||||
return {"id": self.id, "type": self.type}
|
||||
|
||||
def _find_handler(self, message: Any) -> Callable[[Any, WorkflowContext[Any, Any]], Awaitable[None]]:
|
||||
"""Find the handler for a given message.
|
||||
|
||||
Args:
|
||||
message: The message to find the handler for.
|
||||
|
||||
Returns:
|
||||
The handler function if found, None otherwise
|
||||
"""
|
||||
if isinstance(message, WorkflowMessage):
|
||||
# Case where Message wrapper is passed instead of raw data
|
||||
# Handler can be a standard handler or a response handler
|
||||
if message.type == MessageType.STANDARD:
|
||||
for message_type in self._handlers:
|
||||
if is_instance_of(message.data, message_type):
|
||||
return self._handlers[message_type]
|
||||
raise RuntimeError(
|
||||
f"Executor {self.__class__.__name__} cannot handle message of type {type(message.data)}."
|
||||
)
|
||||
# Response message case - find response handler based on original request and response types
|
||||
if message.original_request_info_event is None:
|
||||
raise RuntimeError(
|
||||
f"Executor {self.__class__.__name__} received a response message without an original request event."
|
||||
)
|
||||
handler = self._find_response_handler(message.original_request_info_event.data, message.data)
|
||||
if not handler:
|
||||
raise RuntimeError(
|
||||
f"Executor {self.__class__.__name__} cannot handle request of type "
|
||||
f"{type(message.original_request_info_event.data)} and response of type {type(message.data)}."
|
||||
)
|
||||
return handler
|
||||
|
||||
# Standard raw message data case - only standard handlers apply
|
||||
for message_type in self._handlers:
|
||||
if is_instance_of(message, message_type):
|
||||
return self._handlers[message_type]
|
||||
raise RuntimeError(f"Executor {self.__class__.__name__} cannot handle message of type {type(message)}.")
|
||||
|
||||
async def on_checkpoint_save(self) -> dict[str, Any]:
|
||||
"""Hook called when the workflow is being saved to a checkpoint.
|
||||
|
||||
Override this method in subclasses to implement custom logic that should
|
||||
return state to be saved in the checkpoint.
|
||||
|
||||
The returned state dictionary will be passed to `on_checkpoint_restore`
|
||||
when the workflow is restored from the checkpoint. The dictionary should
|
||||
only contain JSON-serializable data.
|
||||
|
||||
Returns:
|
||||
A state dictionary to be saved during checkpointing.
|
||||
"""
|
||||
return {}
|
||||
|
||||
async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
|
||||
"""Hook called when the workflow is restored from a checkpoint.
|
||||
|
||||
Override this method in subclasses to implement custom logic that should
|
||||
run when the workflow is restored from a checkpoint.
|
||||
|
||||
Args:
|
||||
state: The state dictionary that was saved during checkpointing.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
# endregion: Executor
|
||||
|
||||
# region Handler Decorator
|
||||
|
||||
|
||||
ExecutorT = TypeVar("ExecutorT", bound="Executor")
|
||||
ContextT = TypeVar("ContextT", bound="WorkflowContext[Any, Any]")
|
||||
|
||||
|
||||
@overload
|
||||
def handler(
|
||||
func: Callable[[ExecutorT, Any, ContextT], Awaitable[Any]],
|
||||
) -> Callable[[ExecutorT, Any, ContextT], Awaitable[Any]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def handler(
|
||||
*,
|
||||
input: type | types.UnionType | str | None = None,
|
||||
output: type | types.UnionType | str | None = None,
|
||||
workflow_output: type | types.UnionType | str | None = None,
|
||||
) -> Callable[
|
||||
[Callable[..., Awaitable[Any]]],
|
||||
Callable[..., Awaitable[Any]],
|
||||
]: ...
|
||||
|
||||
|
||||
def handler(
|
||||
func: Callable[[ExecutorT, Any, ContextT], Awaitable[Any]] | None = None,
|
||||
*,
|
||||
input: type | types.UnionType | str | None = None,
|
||||
output: type | types.UnionType | str | None = None,
|
||||
workflow_output: type | types.UnionType | str | None = None,
|
||||
) -> (
|
||||
Callable[[ExecutorT, Any, ContextT], Awaitable[Any]]
|
||||
| Callable[
|
||||
[Callable[[ExecutorT, Any, ContextT], Awaitable[Any]]],
|
||||
Callable[[ExecutorT, Any, ContextT], Awaitable[Any]],
|
||||
]
|
||||
):
|
||||
"""Decorator to register a handler for an executor.
|
||||
|
||||
Type information can be provided in two mutually exclusive ways:
|
||||
|
||||
1. **Introspection** (default): Types are inferred from function signature annotations.
|
||||
Use type annotations on the message parameter and WorkflowContext generic parameters.
|
||||
|
||||
2. **Explicit parameters**: Types are specified via decorator parameters (input, output,
|
||||
workflow_output). When ANY explicit parameter is provided, ALL types must come from
|
||||
explicit parameters - introspection is completely disabled. The ``input`` parameter
|
||||
is required; ``output`` and ``workflow_output`` are optional (default to no outputs).
|
||||
|
||||
Args:
|
||||
func: The function to decorate. Can be None when used with parameters.
|
||||
input: Explicit input type(s) for this handler. Required when using explicit mode.
|
||||
Supports union types (e.g., ``str | int``) and string forward references.
|
||||
output: Explicit output type(s) that can be sent via ``ctx.send_message()``.
|
||||
Optional; defaults to no outputs if not specified.
|
||||
workflow_output: Explicit output type(s) that can be yielded via ``ctx.yield_output()``.
|
||||
Optional; defaults to no outputs if not specified.
|
||||
|
||||
Returns:
|
||||
The decorated function with handler metadata.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
# Mode 1: Introspection - types from annotations
|
||||
@handler
|
||||
async def handle_string(self, message: str, ctx: WorkflowContext[str]) -> None: ...
|
||||
|
||||
|
||||
# Mode 2: Explicit types - ALL types from decorator params
|
||||
# Note: No type annotations on function parameters when using explicit types
|
||||
@handler(input=str | int, output=bool)
|
||||
async def handle_data(self, message, ctx): ...
|
||||
|
||||
|
||||
# Explicit with string forward references
|
||||
@handler(input="MyCustomType | int", output="ResponseType")
|
||||
async def handle_custom(self, message, ctx): ...
|
||||
|
||||
|
||||
# Explicit with all three type parameters
|
||||
@handler(input=str, output=int, workflow_output=bool)
|
||||
async def handle_full(self, message, ctx):
|
||||
await ctx.send_message(42) # int - matches output
|
||||
await ctx.yield_output(True) # bool - matches workflow_output
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[[ExecutorT, Any, ContextT], Awaitable[Any]],
|
||||
) -> Callable[[ExecutorT, Any, ContextT], Awaitable[Any]]:
|
||||
# Check if ANY explicit type parameter was provided - if so, use ONLY explicit params.
|
||||
# This is "all or nothing" - no mixing of explicit params with introspection.
|
||||
use_explicit_types = input is not None or output is not None or workflow_output is not None
|
||||
|
||||
if use_explicit_types:
|
||||
# Resolve string forward references using the function's globals
|
||||
resolved_input_type = resolve_type_annotation(input, func.__globals__) if input is not None else None
|
||||
resolved_output_type = resolve_type_annotation(output, func.__globals__) if output is not None else None
|
||||
resolved_workflow_output_type = (
|
||||
resolve_type_annotation(workflow_output, func.__globals__) if workflow_output is not None else None
|
||||
)
|
||||
|
||||
# Check for unresolved TypeVars in explicit type parameters
|
||||
for param_name, param_type in [
|
||||
("input", resolved_input_type),
|
||||
("output", resolved_output_type),
|
||||
("workflow_output", resolved_workflow_output_type),
|
||||
]:
|
||||
if param_type is not None and contains_typevar(param_type):
|
||||
raise ValueError(
|
||||
f"Handler '{func.__name__}' has an unresolved TypeVar '{param_type}' "
|
||||
f"as its {param_name} type. "
|
||||
f"Use @handler(input=ConcreteType, output=ConcreteType) with concrete types "
|
||||
f"for parameterized executors."
|
||||
)
|
||||
|
||||
# Validate signature structure (correct number of params, ctx is WorkflowContext)
|
||||
# but skip type extraction since we're using explicit types
|
||||
_validate_handler_signature(func, skip_message_annotation=True)
|
||||
|
||||
# Use explicit types only - missing params default to empty
|
||||
message_type = resolved_input_type
|
||||
if message_type is None:
|
||||
raise ValueError(f"Handler {func.__name__} with explicit type parameters must specify 'input' type")
|
||||
|
||||
final_output_types = normalize_type_to_list(resolved_output_type) if resolved_output_type else []
|
||||
final_workflow_output_types = (
|
||||
normalize_type_to_list(resolved_workflow_output_type) if resolved_workflow_output_type else []
|
||||
)
|
||||
# Get ctx_annotation for consistency (even though types come from explicit params)
|
||||
ctx_annotation = (
|
||||
inspect.signature(func).parameters[list(inspect.signature(func).parameters.keys())[2]].annotation
|
||||
)
|
||||
else:
|
||||
# Use introspection for ALL types - no explicit params provided
|
||||
introspected_message_type, ctx_annotation, inferred_output_types, inferred_workflow_output_types = (
|
||||
_validate_handler_signature(func, skip_message_annotation=False)
|
||||
)
|
||||
|
||||
message_type = introspected_message_type
|
||||
if message_type is None:
|
||||
raise ValueError(
|
||||
f"Handler {func.__name__} requires either a message parameter type annotation "
|
||||
"or explicit type parameters (input, output, workflow_output)"
|
||||
)
|
||||
|
||||
# Check for unresolved TypeVar in introspected message type
|
||||
if contains_typevar(message_type):
|
||||
raise ValueError(
|
||||
f"Handler '{func.__name__}' has an unresolved TypeVar '{message_type}' "
|
||||
f"as its message type. "
|
||||
f"Use @handler(input=ConcreteType, output=ConcreteType) with concrete types "
|
||||
f"for parameterized executors."
|
||||
)
|
||||
|
||||
final_output_types = inferred_output_types
|
||||
final_workflow_output_types = inferred_workflow_output_types
|
||||
|
||||
# Get signature for preservation
|
||||
sig = inspect.signature(func)
|
||||
|
||||
@functools.wraps(func)
|
||||
async def wrapper(self: ExecutorT, message: Any, ctx: ContextT) -> Any:
|
||||
"""Wrapper function to call the handler."""
|
||||
return await func(self, message, ctx)
|
||||
|
||||
# Preserve the original function signature for introspection during validation
|
||||
with contextlib.suppress(AttributeError, TypeError):
|
||||
wrapper.__signature__ = sig # type: ignore[attr-defined]
|
||||
|
||||
wrapper._handler_spec = { # type: ignore
|
||||
"name": func.__name__,
|
||||
"message_type": message_type,
|
||||
# Keep output_types and workflow_output_types in spec for validators
|
||||
"output_types": final_output_types,
|
||||
"workflow_output_types": final_workflow_output_types,
|
||||
"ctx_annotation": ctx_annotation,
|
||||
}
|
||||
|
||||
return wrapper
|
||||
|
||||
# Handle both @handler and @handler(...) usage patterns
|
||||
if func is not None:
|
||||
# Called as @handler without parentheses
|
||||
return decorator(func)
|
||||
# Called as @handler(...) with parentheses
|
||||
return decorator
|
||||
|
||||
|
||||
# endregion: Handler Decorator
|
||||
|
||||
# region Handler Validation
|
||||
|
||||
|
||||
def _validate_handler_signature(
|
||||
func: Callable[..., Any],
|
||||
*,
|
||||
skip_message_annotation: bool = False,
|
||||
) -> tuple[type | None, Any, list[type[Any] | types.UnionType], list[type[Any] | types.UnionType]]:
|
||||
"""Validate function signature for executor functions.
|
||||
|
||||
Args:
|
||||
func: The function to validate
|
||||
skip_message_annotation: If True, skip validation that message parameter has a type
|
||||
annotation. Used when input_type is explicitly provided to the @handler decorator.
|
||||
|
||||
Returns:
|
||||
Tuple of (message_type, ctx_annotation, output_types, workflow_output_types).
|
||||
message_type may be None if skip_message_annotation is True and no annotation exists.
|
||||
|
||||
Raises:
|
||||
ValueError: If the function signature is invalid
|
||||
"""
|
||||
signature = inspect.signature(func)
|
||||
params = list(signature.parameters.values())
|
||||
|
||||
expected_counts = 3 # self, message, ctx
|
||||
param_description = "(self, message: T, ctx: WorkflowContext[U, V])"
|
||||
if len(params) != expected_counts:
|
||||
raise ValueError(f"Handler {func.__name__} must have {param_description}. Got {len(params)} parameters.")
|
||||
|
||||
# Check message parameter has type annotation (unless skipped)
|
||||
message_param = params[1]
|
||||
if not skip_message_annotation and message_param.annotation == inspect.Parameter.empty:
|
||||
raise ValueError(f"Handler {func.__name__} must have a type annotation for the message parameter")
|
||||
|
||||
# Resolve string annotations from `from __future__ import annotations`.
|
||||
# Fall back to raw annotations if resolution fails (e.g. unresolvable forward refs,
|
||||
# AttributeError, or RecursionError), so registration failures are easier to diagnose.
|
||||
try:
|
||||
type_hints = typing.get_type_hints(func)
|
||||
except (NameError, AttributeError, RecursionError):
|
||||
type_hints = {p.name: p.annotation for p in params}
|
||||
|
||||
message_type = type_hints.get(message_param.name, message_param.annotation)
|
||||
if message_type == inspect.Parameter.empty:
|
||||
message_type = None
|
||||
|
||||
# Reject unresolved TypeVar in message annotation -- these are not supported
|
||||
# for workflow type validation and must be replaced with concrete types.
|
||||
if not skip_message_annotation and contains_typevar(message_type):
|
||||
raise ValueError(
|
||||
f"Handler {func.__name__} has an unresolved TypeVar '{message_type}' as its message type annotation. "
|
||||
"Generic TypeVar annotations are not supported for workflow type validation. "
|
||||
"Use @handler(input=<concrete_type>, output=<concrete_type>) to specify explicit types."
|
||||
)
|
||||
|
||||
# Validate ctx parameter is WorkflowContext and extract type args
|
||||
ctx_param = params[2]
|
||||
ctx_annotation = type_hints.get(ctx_param.name, ctx_param.annotation)
|
||||
if skip_message_annotation and ctx_annotation == inspect.Parameter.empty:
|
||||
# When explicit types are provided via @handler(input=..., output=...),
|
||||
# the ctx parameter doesn't need a type annotation - types come from the decorator.
|
||||
output_types: list[type[Any] | types.UnionType] = []
|
||||
workflow_output_types: list[type[Any] | types.UnionType] = []
|
||||
else:
|
||||
output_types, workflow_output_types = validate_workflow_context_annotation(
|
||||
ctx_annotation, f"parameter '{ctx_param.name}'", "Handler"
|
||||
)
|
||||
|
||||
return message_type, ctx_annotation, output_types, workflow_output_types
|
||||
|
||||
|
||||
# endregion: Handler Validation
|
||||
@@ -0,0 +1,397 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Function-based Executor and decorator utilities.
|
||||
|
||||
This module provides:
|
||||
- FunctionExecutor: an Executor subclass that wraps a standalone user-defined function
|
||||
with signature (message) or (message, ctx: WorkflowContext[T]). Both sync and async functions are supported.
|
||||
Synchronous functions are executed in a thread pool using asyncio.to_thread() to avoid blocking the event loop.
|
||||
- executor decorator: converts a standalone module-level function into a ready-to-use Executor instance
|
||||
with proper type validation and handler registration.
|
||||
|
||||
Design Pattern:
|
||||
- Use @executor for standalone module-level or local functions
|
||||
- Use Executor subclass with @handler for class-based executors with state/dependencies
|
||||
- Do NOT use @executor with @staticmethod or @classmethod
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import sys
|
||||
import types
|
||||
import typing
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from ._executor import Executor
|
||||
from ._typing_utils import contains_typevar, normalize_type_to_list, resolve_type_annotation
|
||||
from ._workflow_context import WorkflowContext, validate_workflow_context_annotation
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import overload # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import overload # pragma: no cover
|
||||
|
||||
|
||||
class FunctionExecutor(Executor):
|
||||
"""Executor that wraps a user-defined function.
|
||||
|
||||
This executor allows users to define simple functions (both sync and async) and use them
|
||||
as workflow executors without needing to create full executor classes.
|
||||
|
||||
Synchronous functions are executed in a thread pool using asyncio.to_thread() to avoid
|
||||
blocking the event loop.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
func: Callable[..., Any],
|
||||
id: str | None = None,
|
||||
*,
|
||||
input: type | types.UnionType | str | None = None,
|
||||
output: type | types.UnionType | str | None = None,
|
||||
workflow_output: type | types.UnionType | str | None = None,
|
||||
):
|
||||
"""Initialize the FunctionExecutor with a user-defined function.
|
||||
|
||||
Args:
|
||||
func: The function to wrap as an executor (can be sync or async)
|
||||
id: Optional executor ID. If None, uses the function name.
|
||||
input: Optional explicit input type(s) for this executor. Supports union types
|
||||
(e.g., ``str | int``) and string forward references (e.g., ``"MyType | int"``).
|
||||
When provided, takes precedence over introspection from the function's message
|
||||
parameter annotation.
|
||||
output: Optional explicit output type(s) that can be sent via ``ctx.send_message()``.
|
||||
Supports union types (e.g., ``str | int``) and string forward references.
|
||||
When provided, takes precedence over introspection from the ``WorkflowContext``
|
||||
first generic parameter (OutT).
|
||||
workflow_output: Optional explicit output type(s) that can be yielded via
|
||||
``ctx.yield_output()``. Supports union types (e.g., ``str | int``) and string
|
||||
forward references. When provided, takes precedence over introspection from the
|
||||
``WorkflowContext`` second generic parameter (W_OutT).
|
||||
|
||||
Raises:
|
||||
ValueError: If func is a staticmethod or classmethod (use @handler on instance methods instead)
|
||||
"""
|
||||
# Detect misuse of @executor with staticmethod/classmethod
|
||||
if isinstance(func, (staticmethod, classmethod)):
|
||||
descriptor_type = "staticmethod" if isinstance(func, staticmethod) else "classmethod"
|
||||
raise ValueError(
|
||||
f"The @executor decorator cannot be used with @{descriptor_type}. "
|
||||
f"Use the @executor decorator on standalone module-level functions, "
|
||||
f"or create an Executor subclass and use @handler on instance methods instead."
|
||||
)
|
||||
|
||||
# Resolve string forward references using the function's globals
|
||||
resolved_input_type = resolve_type_annotation(input, func.__globals__) if input is not None else None
|
||||
resolved_output_type = resolve_type_annotation(output, func.__globals__) if output is not None else None
|
||||
resolved_workflow_output_type = (
|
||||
resolve_type_annotation(workflow_output, func.__globals__) if workflow_output is not None else None
|
||||
)
|
||||
|
||||
# Validate function signature and extract types
|
||||
introspected_message_type, ctx_annotation, inferred_output_types, inferred_workflow_output_types = (
|
||||
_validate_function_signature(func, skip_message_annotation=resolved_input_type is not None)
|
||||
)
|
||||
|
||||
# Check for unresolved TypeVars in explicit type parameters
|
||||
for param_name, param_type in [
|
||||
("input", resolved_input_type),
|
||||
("output", resolved_output_type),
|
||||
("workflow_output", resolved_workflow_output_type),
|
||||
]:
|
||||
if param_type is not None and contains_typevar(param_type):
|
||||
raise ValueError(
|
||||
f"Executor '{func.__name__}' has an unresolved TypeVar '{param_type}' "
|
||||
f"as its {param_name} type. "
|
||||
f"Use @executor(input=ConcreteType, output=ConcreteType) with concrete types."
|
||||
)
|
||||
|
||||
# Use explicit types if provided, otherwise fall back to introspection
|
||||
message_type = resolved_input_type if resolved_input_type is not None else introspected_message_type
|
||||
output_types: list[type[Any] | types.UnionType] = (
|
||||
normalize_type_to_list(resolved_output_type)
|
||||
if resolved_output_type is not None
|
||||
else list(inferred_output_types)
|
||||
)
|
||||
final_workflow_output_types: list[type[Any] | types.UnionType] = (
|
||||
normalize_type_to_list(resolved_workflow_output_type)
|
||||
if resolved_workflow_output_type is not None
|
||||
else list(inferred_workflow_output_types)
|
||||
)
|
||||
|
||||
# Validate that we have a message type - provides a clear error if type information is missing
|
||||
if message_type is None:
|
||||
raise ValueError(
|
||||
f"Function {func.__name__} requires either a message parameter type annotation "
|
||||
"or an explicit input_type parameter"
|
||||
)
|
||||
|
||||
# Check for unresolved TypeVar in introspected message type
|
||||
if contains_typevar(message_type):
|
||||
raise ValueError(
|
||||
f"Executor '{func.__name__}' has an unresolved TypeVar '{message_type}' "
|
||||
f"as its message type. "
|
||||
f"Use @executor(input=ConcreteType, output=ConcreteType) with concrete types."
|
||||
)
|
||||
|
||||
# Store the original function
|
||||
self._original_func = func
|
||||
# Determine if function has WorkflowContext parameter
|
||||
self._has_context = ctx_annotation is not None
|
||||
# Determine if the function is an async function
|
||||
self._is_async = inspect.iscoroutinefunction(func)
|
||||
|
||||
# Initialize parent WITHOUT calling _discover_handlers yet
|
||||
# We'll manually set up the attributes first
|
||||
executor_id = str(id or getattr(func, "__name__", "FunctionExecutor"))
|
||||
kwargs = {"type": "FunctionExecutor"}
|
||||
|
||||
super().__init__(id=executor_id, defer_discovery=True, **kwargs)
|
||||
|
||||
# Create a wrapper function that always accepts both message and context
|
||||
if self._has_context and self._is_async:
|
||||
# Async function with context - already has the right signature
|
||||
wrapped_func: Callable[[Any, WorkflowContext[Any]], Awaitable[Any]] = func # type: ignore
|
||||
elif self._has_context and not self._is_async:
|
||||
# Sync function with context - wrap to make async using thread pool
|
||||
async def wrapped_func(message: Any, ctx: WorkflowContext[Any]) -> Any:
|
||||
# Call the sync function with both parameters in a thread
|
||||
return await asyncio.to_thread(func, message, ctx)
|
||||
|
||||
elif not self._has_context and self._is_async:
|
||||
# Async function without context - wrap to ignore context
|
||||
async def wrapped_func(message: Any, ctx: WorkflowContext[Any]) -> Any:
|
||||
# Call the async function with just the message
|
||||
return await func(message)
|
||||
|
||||
else:
|
||||
# Sync function without context - wrap to make async and ignore context using thread pool
|
||||
async def wrapped_func(message: Any, ctx: WorkflowContext[Any]) -> Any:
|
||||
# Call the sync function with just the message in a thread
|
||||
return await asyncio.to_thread(func, message)
|
||||
|
||||
# Now register our instance handler
|
||||
self._register_instance_handler(
|
||||
name=func.__name__,
|
||||
func=wrapped_func,
|
||||
message_type=message_type,
|
||||
ctx_annotation=ctx_annotation,
|
||||
output_types=output_types,
|
||||
workflow_output_types=final_workflow_output_types,
|
||||
)
|
||||
|
||||
# Now we can safely call _discover_handlers (it won't find any class-level handlers)
|
||||
self._discover_handlers()
|
||||
self._discover_response_handlers()
|
||||
|
||||
if not self._handlers:
|
||||
raise ValueError(
|
||||
f"FunctionExecutor {self.__class__.__name__} failed to register handler for {func.__name__}"
|
||||
)
|
||||
|
||||
|
||||
# region Decorator
|
||||
|
||||
|
||||
@overload
|
||||
def executor(func: Callable[..., Any]) -> FunctionExecutor: ...
|
||||
|
||||
|
||||
@overload
|
||||
def executor(
|
||||
*,
|
||||
id: str | None = None,
|
||||
input: type | types.UnionType | str | None = None,
|
||||
output: type | types.UnionType | str | None = None,
|
||||
workflow_output: type | types.UnionType | str | None = None,
|
||||
) -> Callable[[Callable[..., Any]], FunctionExecutor]: ...
|
||||
|
||||
|
||||
def executor(
|
||||
func: Callable[..., Any] | None = None,
|
||||
*,
|
||||
id: str | None = None,
|
||||
input: type | types.UnionType | str | None = None,
|
||||
output: type | types.UnionType | str | None = None,
|
||||
workflow_output: type | types.UnionType | str | None = None,
|
||||
) -> Callable[[Callable[..., Any]], FunctionExecutor] | FunctionExecutor:
|
||||
"""Decorator that converts a standalone function into a FunctionExecutor instance.
|
||||
|
||||
The @executor decorator is designed for **standalone module-level functions only**.
|
||||
For class-based executors, use the Executor base class with @handler on instance methods.
|
||||
|
||||
Supports both synchronous and asynchronous functions. Synchronous functions
|
||||
are executed in a thread pool to avoid blocking the event loop.
|
||||
|
||||
Important:
|
||||
- Use @executor for standalone functions (module-level or local functions)
|
||||
- Do NOT use @executor with @staticmethod or @classmethod
|
||||
- For class-based executors, subclass Executor and use @handler on instance methods
|
||||
|
||||
Usage:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Standalone async function (RECOMMENDED):
|
||||
@executor(id="upper_case")
|
||||
async def to_upper(text: str, ctx: WorkflowContext[str]):
|
||||
await ctx.send_message(text.upper())
|
||||
|
||||
|
||||
# Standalone sync function (runs in thread pool):
|
||||
@executor
|
||||
def process_data(data: str):
|
||||
return data.upper()
|
||||
|
||||
|
||||
# Using explicit types (takes precedence over introspection):
|
||||
# Note: No type annotations on function parameters when using explicit types
|
||||
@executor(id="my_executor", input=str | int, output=bool)
|
||||
async def process(message, ctx):
|
||||
await ctx.send_message(True)
|
||||
|
||||
|
||||
# Using string forward references:
|
||||
@executor(input="MyCustomType | int", output="ResponseType")
|
||||
async def process(message, ctx): ...
|
||||
|
||||
|
||||
# Specifying both output types (send_message and yield_output):
|
||||
@executor(input=str, output=int, workflow_output=bool)
|
||||
async def process(message, ctx):
|
||||
await ctx.send_message(42) # int - matches output
|
||||
await ctx.yield_output(True) # bool - matches workflow_output
|
||||
|
||||
|
||||
# For class-based executors, use @handler instead:
|
||||
class MyExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="my_executor")
|
||||
|
||||
@handler
|
||||
async def process(self, data: str, ctx: WorkflowContext[str]):
|
||||
await ctx.send_message(data.upper())
|
||||
|
||||
Args:
|
||||
func: The function to decorate (when used without parentheses)
|
||||
id: Optional custom ID for the executor. If None, uses the function name.
|
||||
input: Optional explicit input type(s) for this executor. Supports union types
|
||||
(e.g., ``str | int``) and string forward references (e.g., ``"MyType | int"``).
|
||||
When provided, takes precedence over introspection from the function's message
|
||||
parameter annotation.
|
||||
output: Optional explicit output type(s) that can be sent via ``ctx.send_message()``.
|
||||
Supports union types (e.g., ``str | int``) and string forward references.
|
||||
When provided, takes precedence over introspection from the ``WorkflowContext``
|
||||
first generic parameter (OutT).
|
||||
workflow_output: Optional explicit output type(s) that can be yielded via
|
||||
``ctx.yield_output()``. Supports union types (e.g., ``str | int``) and string
|
||||
forward references. When provided, takes precedence over introspection from the
|
||||
``WorkflowContext`` second generic parameter (W_OutT).
|
||||
|
||||
Warning:
|
||||
When placing a custom ``@executor`` **between** two ``AgentExecutor`` nodes, be
|
||||
careful about the output type. If the custom executor receives an
|
||||
``AgentExecutorResponse`` but emits a plain ``str``, the downstream
|
||||
``AgentExecutor.from_str`` handler is invoked instead of ``from_response``.
|
||||
This resets the conversation context because only the new string is added to
|
||||
the cache and all prior messages from the upstream agent are lost.
|
||||
|
||||
To preserve the full conversation, use
|
||||
``AgentExecutorResponse.with_text(new_text)`` to create a new response that
|
||||
keeps the prior history, and set ``output=AgentExecutorResponse`` on the
|
||||
decorator.
|
||||
|
||||
Returns:
|
||||
A FunctionExecutor instance that can be wired into a Workflow.
|
||||
|
||||
Raises:
|
||||
ValueError: If used with @staticmethod or @classmethod (unsupported pattern)
|
||||
"""
|
||||
|
||||
def wrapper(func: Callable[..., Any]) -> FunctionExecutor:
|
||||
return FunctionExecutor(func, id=id, input=input, output=output, workflow_output=workflow_output)
|
||||
|
||||
# If func is provided, this means @executor was used without parentheses
|
||||
if func is not None:
|
||||
return wrapper(func)
|
||||
|
||||
# Otherwise, return the wrapper for @executor() or @executor(id="...")
|
||||
return wrapper
|
||||
|
||||
|
||||
# endregion: Decorator
|
||||
|
||||
# region Function Validation
|
||||
|
||||
|
||||
def _validate_function_signature(
|
||||
func: Callable[..., Any],
|
||||
*,
|
||||
skip_message_annotation: bool = False,
|
||||
) -> tuple[type | None, Any, list[type[Any] | types.UnionType], list[type[Any] | types.UnionType]]:
|
||||
"""Validate function signature for executor functions.
|
||||
|
||||
Args:
|
||||
func: The function to validate
|
||||
skip_message_annotation: If True, skip validation that message parameter has a type
|
||||
annotation. Used when input is explicitly provided to the @executor decorator.
|
||||
|
||||
Returns:
|
||||
Tuple of (message_type, ctx_annotation, output_types, workflow_output_types).
|
||||
message_type may be None if skip_message_annotation is True and no annotation exists.
|
||||
|
||||
Raises:
|
||||
ValueError: If the function signature is invalid
|
||||
"""
|
||||
signature = inspect.signature(func)
|
||||
params = list(signature.parameters.values())
|
||||
|
||||
expected_counts = (1, 2) # Function executor: (message) or (message, ctx)
|
||||
param_description = "(message: T) or (message: T, ctx: WorkflowContext[U])"
|
||||
if len(params) not in expected_counts:
|
||||
raise ValueError(
|
||||
f"Function instance {func.__name__} must have {param_description}. Got {len(params)} parameters."
|
||||
)
|
||||
|
||||
# Check message parameter has type annotation (unless skipped)
|
||||
message_param = params[0]
|
||||
if not skip_message_annotation and message_param.annotation == inspect.Parameter.empty:
|
||||
raise ValueError(f"Function instance {func.__name__} must have a type annotation for the message parameter")
|
||||
|
||||
# Resolve string annotations from `from __future__ import annotations`.
|
||||
# Fall back to raw annotations if resolution fails (e.g. unresolvable forward refs,
|
||||
# AttributeError, or RecursionError), so registration failures are easier to diagnose.
|
||||
try:
|
||||
type_hints = typing.get_type_hints(func)
|
||||
except (NameError, AttributeError, RecursionError):
|
||||
type_hints = {p.name: p.annotation for p in params}
|
||||
message_type = type_hints.get(message_param.name, message_param.annotation)
|
||||
if message_type == inspect.Parameter.empty:
|
||||
message_type = None
|
||||
|
||||
# Reject unresolved TypeVar in message annotation -- these are not supported
|
||||
# for workflow type validation and must be replaced with concrete types.
|
||||
if not skip_message_annotation and contains_typevar(message_type):
|
||||
raise ValueError(
|
||||
f"Function instance {func.__name__} has an unresolved TypeVar '{message_type}' as its message type "
|
||||
"annotation. Generic TypeVar annotations are not supported for workflow type validation. "
|
||||
"Use @executor(input=<concrete_type>, output=<concrete_type>) to specify explicit types."
|
||||
)
|
||||
|
||||
# Check if there's a context parameter
|
||||
if len(params) == 2:
|
||||
ctx_param = params[1]
|
||||
ctx_annotation = type_hints.get(ctx_param.name, ctx_param.annotation)
|
||||
output_types, workflow_output_types = validate_workflow_context_annotation(
|
||||
ctx_annotation, f"parameter '{ctx_param.name}'", "Function instance"
|
||||
)
|
||||
else:
|
||||
# No context parameter (only valid for function executors)
|
||||
output_types, workflow_output_types = [], []
|
||||
ctx_annotation = None
|
||||
|
||||
return message_type, ctx_annotation, output_types, workflow_output_types
|
||||
|
||||
|
||||
# endregion: Function Validation
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,42 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Shared helpers for normalizing workflow message inputs."""
|
||||
|
||||
from agent_framework import Content, Message
|
||||
from agent_framework._types import AgentRunInputs
|
||||
|
||||
|
||||
def normalize_messages_input(
|
||||
messages: AgentRunInputs | None = None,
|
||||
) -> list[Message]:
|
||||
"""Normalize heterogeneous message inputs to a list of Message objects.
|
||||
|
||||
Args:
|
||||
messages: String, Content, Message, or sequence of those values. None yields empty list.
|
||||
|
||||
Returns:
|
||||
List of Message instances suitable for workflow consumption.
|
||||
"""
|
||||
if messages is None:
|
||||
return []
|
||||
|
||||
if isinstance(messages, str):
|
||||
return [Message(role="user", contents=[messages])]
|
||||
|
||||
if isinstance(messages, Content):
|
||||
return [Message(role="user", contents=[messages])]
|
||||
|
||||
if isinstance(messages, Message):
|
||||
return [messages]
|
||||
|
||||
normalized: list[Message] = []
|
||||
for item in messages:
|
||||
if isinstance(item, (str, Content)):
|
||||
normalized.append(Message(role="user", contents=[item]))
|
||||
elif isinstance(item, Message):
|
||||
normalized.append(item)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Messages sequence must contain only str, Content, or Message instances; found {type(item).__name__}."
|
||||
)
|
||||
return normalized
|
||||
@@ -0,0 +1,51 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import copy
|
||||
import sys
|
||||
from typing import Any, TypeVar, cast
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import Self # pragma: no cover
|
||||
|
||||
ModelT = TypeVar("ModelT", bound="DictConvertible")
|
||||
|
||||
|
||||
class DictConvertible:
|
||||
"""Mixin providing conversion helpers for plain Python models."""
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls: type[ModelT], data: dict[str, Any]) -> ModelT:
|
||||
return cls(**data)
|
||||
|
||||
def clone(self, *, deep: bool = True) -> Self:
|
||||
return copy.deepcopy(self) if deep else copy.copy(self)
|
||||
|
||||
def to_json(self) -> str:
|
||||
import json
|
||||
|
||||
return json.dumps(self.to_dict())
|
||||
|
||||
@classmethod
|
||||
def from_json(cls: type[ModelT], raw: str) -> ModelT:
|
||||
import json
|
||||
|
||||
data = json.loads(raw)
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("JSON payload must decode to a mapping")
|
||||
return cls.from_dict(cast(dict[str, Any], data))
|
||||
|
||||
|
||||
def encode_value(value: Any) -> Any:
|
||||
"""Recursively encode values for JSON-friendly serialization."""
|
||||
if isinstance(value, DictConvertible):
|
||||
return value.to_dict()
|
||||
if isinstance(value, dict):
|
||||
return {k: encode_value(v) for k, v in value.items()} # type: ignore[misc]
|
||||
if isinstance(value, (list, tuple, set)):
|
||||
return [encode_value(v) for v in value] # type: ignore[misc]
|
||||
return value
|
||||
@@ -0,0 +1,369 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import inspect
|
||||
import logging
|
||||
import sys
|
||||
import types
|
||||
from builtins import type as builtin_type
|
||||
from collections.abc import Awaitable, Callable
|
||||
from types import UnionType
|
||||
from typing import TYPE_CHECKING, Any, TypeVar, cast
|
||||
|
||||
from ._typing_utils import is_instance_of, is_type_compatible, normalize_type_to_list, resolve_type_annotation
|
||||
from ._workflow_context import WorkflowContext, validate_workflow_context_annotation
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import overload # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import overload # pragma: no cover
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._executor import Executor
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RequestInfoMixin:
|
||||
"""Mixin providing common functionality for request info handling."""
|
||||
|
||||
def is_request_supported(self, request_type: builtin_type[Any], response_type: builtin_type[Any]) -> bool:
|
||||
"""Check if the executor supports request of the given type and handling a response of the given type.
|
||||
|
||||
Args:
|
||||
request_type: The type of the request message
|
||||
response_type: The type of the expected response message
|
||||
Returns:
|
||||
True if a response handler is registered for the given request and response types, False otherwise
|
||||
"""
|
||||
if not hasattr(self, "_response_handlers"):
|
||||
return False
|
||||
|
||||
for request_type_key, response_type_key in self._response_handlers:
|
||||
if is_type_compatible(request_type, request_type_key) and is_type_compatible(
|
||||
response_type, response_type_key
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _find_response_handler(self, request: Any, response: Any) -> Callable[..., Awaitable[None]] | None:
|
||||
"""Find a registered response handler for the given request and response types.
|
||||
|
||||
Args:
|
||||
request: The original request
|
||||
response: The response message
|
||||
Returns:
|
||||
The response handler function with the request bound as the first argument, or None if not found
|
||||
"""
|
||||
if not hasattr(self, "_response_handlers"):
|
||||
return None
|
||||
|
||||
for (request_type, response_type), handler in self._response_handlers.items():
|
||||
if is_instance_of(request, request_type) and is_instance_of(response, response_type):
|
||||
return functools.partial(handler, request)
|
||||
|
||||
return None
|
||||
|
||||
def _discover_response_handlers(self) -> None:
|
||||
"""Discover and register response handlers defined in the class."""
|
||||
# Initialize handler storage if not already present
|
||||
if not hasattr(self, "_response_handlers"):
|
||||
self._response_handlers: dict[
|
||||
tuple[builtin_type[Any], builtin_type[Any]], # key
|
||||
Callable[[Any, Any, WorkflowContext[Any, Any]], Awaitable[None]], # value
|
||||
] = {}
|
||||
if not hasattr(self, "_response_handler_specs"):
|
||||
self._response_handler_specs: list[dict[str, Any]] = []
|
||||
|
||||
for attr_name in dir(self.__class__):
|
||||
try:
|
||||
attr = getattr(self.__class__, attr_name)
|
||||
if callable(attr) and hasattr(attr, "_response_handler_spec"):
|
||||
handler_spec = attr._response_handler_spec # type: ignore
|
||||
|
||||
request_type = handler_spec["request_type"]
|
||||
response_type = handler_spec["response_type"]
|
||||
|
||||
if self._response_handlers.get((request_type, response_type)):
|
||||
raise ValueError(
|
||||
f"Duplicate response handler for request type {request_type} "
|
||||
f"and response type {response_type} in {self.__class__.__name__}"
|
||||
)
|
||||
|
||||
self._response_handlers[request_type, response_type] = getattr(self, attr_name)
|
||||
self._response_handler_specs.append({**handler_spec, "source": "class_method"})
|
||||
except AttributeError:
|
||||
continue # Skip non-callable attributes or those without handler spec
|
||||
|
||||
# A request sent via `request_info` must be handled by a response handler inside the same executor.
|
||||
# It is safe to assume that an executor is request-response capable if it has at least one response
|
||||
# handler, and that the executor could send a request.
|
||||
self.is_request_response_capable = bool(self._response_handlers)
|
||||
|
||||
|
||||
ExecutorT = TypeVar("ExecutorT", bound="Executor")
|
||||
ContextT = TypeVar("ContextT", bound="WorkflowContext[Any, Any]")
|
||||
|
||||
# region Handler Decorator
|
||||
|
||||
|
||||
@overload
|
||||
def response_handler(
|
||||
func: Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]],
|
||||
) -> Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def response_handler(
|
||||
func: None = None,
|
||||
*,
|
||||
request: type | types.UnionType | str | None = None,
|
||||
response: type | types.UnionType | str | None = None,
|
||||
output: type | types.UnionType | str | None = None,
|
||||
workflow_output: type | types.UnionType | str | None = None,
|
||||
) -> Callable[
|
||||
[Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]]],
|
||||
Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]],
|
||||
]: ...
|
||||
|
||||
|
||||
def response_handler(
|
||||
func: Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]] | None = None,
|
||||
*,
|
||||
request: type | types.UnionType | str | None = None,
|
||||
response: type | types.UnionType | str | None = None,
|
||||
output: type | types.UnionType | str | None = None,
|
||||
workflow_output: type | types.UnionType | str | None = None,
|
||||
) -> (
|
||||
Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]]
|
||||
| Callable[
|
||||
[Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]]],
|
||||
Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]],
|
||||
]
|
||||
):
|
||||
"""Decorator to register a handler to handle responses for a request.
|
||||
|
||||
Type information can be provided in two mutually exclusive ways:
|
||||
|
||||
1. **Introspection** (default): Types are inferred from function signature annotations.
|
||||
Use type annotations on the original_request, response parameters and WorkflowContext
|
||||
generic parameters.
|
||||
|
||||
2. **Explicit parameters**: Types are specified via decorator parameters (request, response,
|
||||
output, workflow_output). When ANY explicit parameter is provided, ALL types must come
|
||||
from explicit parameters - introspection is completely disabled. The ``request`` and
|
||||
``response`` parameters are required; ``output`` and ``workflow_output`` are optional
|
||||
(default to no outputs).
|
||||
|
||||
Args:
|
||||
func: The function to decorate. Can be None when used with parameters.
|
||||
request: Explicit request type for this handler (the original_request parameter type).
|
||||
Required when using explicit mode. Supports union types and string forward references.
|
||||
response: Explicit response type for this handler (the response parameter type).
|
||||
Required when using explicit mode. Supports union types and string forward references.
|
||||
output: Explicit output type(s) that can be sent via ``ctx.send_message()``.
|
||||
Optional; defaults to no outputs if not specified.
|
||||
workflow_output: Explicit output type(s) that can be yielded via ``ctx.yield_output()``.
|
||||
Optional; defaults to no outputs if not specified.
|
||||
|
||||
Returns:
|
||||
The decorated function with handler metadata.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
# Mode 1: Introspection - types from annotations
|
||||
@handler
|
||||
async def run(self, message: int, context: WorkflowContext[str]) -> None:
|
||||
# Example of a handler that sends a request
|
||||
...
|
||||
# Send a request with a `CustomRequest` payload and expect a `str` response.
|
||||
await context.request_info(CustomRequest(...), str)
|
||||
|
||||
|
||||
@response_handler
|
||||
async def handle_response(
|
||||
self,
|
||||
original_request: CustomRequest,
|
||||
response: str,
|
||||
context: WorkflowContext[str],
|
||||
) -> None:
|
||||
# Example of a response handler for the above request
|
||||
...
|
||||
|
||||
|
||||
# Mode 2: Explicit types - ALL types from decorator params
|
||||
# Note: No type annotations on function parameters when using explicit types
|
||||
@response_handler(request=CustomRequest, response=dict, output=int)
|
||||
async def handle_response(self, original_request, response, context):
|
||||
# Example of a response handler with explicit types
|
||||
await context.send_message(42)
|
||||
|
||||
|
||||
# Explicit with string forward references
|
||||
@response_handler(request="MyRequest", response="MyResponse")
|
||||
async def handle_response(self, original_request, response, context): ...
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]],
|
||||
) -> Callable[[ExecutorT, Any, Any, ContextT], Awaitable[None]]:
|
||||
# Check if ANY explicit type parameter was provided - if so, use ONLY explicit params.
|
||||
# This is "all or nothing" - no mixing of explicit params with introspection.
|
||||
use_explicit_types = (
|
||||
request is not None or response is not None or output is not None or workflow_output is not None
|
||||
)
|
||||
|
||||
if use_explicit_types:
|
||||
# Resolve string forward references using the function's globals
|
||||
resolved_request_type = resolve_type_annotation(request, func.__globals__) if request is not None else None
|
||||
resolved_response_type = (
|
||||
resolve_type_annotation(response, func.__globals__) if response is not None else None
|
||||
)
|
||||
resolved_output_type = resolve_type_annotation(output, func.__globals__) if output is not None else None
|
||||
resolved_workflow_output_type = (
|
||||
resolve_type_annotation(workflow_output, func.__globals__) if workflow_output is not None else None
|
||||
)
|
||||
|
||||
# Validate signature structure but skip type extraction
|
||||
_validate_response_handler_signature(func, skip_annotations=True)
|
||||
|
||||
# Validate required parameters
|
||||
if resolved_request_type is None:
|
||||
raise ValueError(
|
||||
f"Response handler {func.__name__} with explicit type parameters must specify 'request' type"
|
||||
)
|
||||
if resolved_response_type is None:
|
||||
raise ValueError(
|
||||
f"Response handler {func.__name__} with explicit type parameters must specify 'response' type"
|
||||
)
|
||||
|
||||
final_request_type = resolved_request_type
|
||||
final_response_type = resolved_response_type
|
||||
final_output_types = normalize_type_to_list(resolved_output_type) if resolved_output_type else []
|
||||
final_workflow_output_types = (
|
||||
normalize_type_to_list(resolved_workflow_output_type) if resolved_workflow_output_type else []
|
||||
)
|
||||
# Get ctx_annotation for consistency
|
||||
ctx_annotation = (
|
||||
inspect.signature(func).parameters[list(inspect.signature(func).parameters.keys())[3]].annotation
|
||||
)
|
||||
if ctx_annotation == inspect.Parameter.empty:
|
||||
ctx_annotation = None
|
||||
else:
|
||||
# Use introspection - all types from annotations
|
||||
(
|
||||
inferred_request_type,
|
||||
inferred_response_type,
|
||||
ctx_annotation,
|
||||
final_output_types,
|
||||
final_workflow_output_types,
|
||||
) = _validate_response_handler_signature(func)
|
||||
# In introspection mode, validation ensures these are not None (raises ValueError if missing)
|
||||
final_request_type = cast(type, inferred_request_type)
|
||||
final_response_type = cast(type, inferred_response_type)
|
||||
|
||||
# Get signature for preservation
|
||||
sig = inspect.signature(func)
|
||||
|
||||
@functools.wraps(func)
|
||||
async def wrapper(self: ExecutorT, original_request: Any, response_msg: Any, ctx: ContextT) -> Any:
|
||||
"""Wrapper function to call the handler."""
|
||||
return await func(self, original_request, response_msg, ctx)
|
||||
|
||||
# Preserve the original function signature for introspection during validation
|
||||
with contextlib.suppress(AttributeError, TypeError):
|
||||
wrapper.__signature__ = sig # type: ignore[attr-defined]
|
||||
|
||||
wrapper._response_handler_spec = { # type: ignore
|
||||
"name": func.__name__,
|
||||
"request_type": final_request_type,
|
||||
"response_type": final_response_type,
|
||||
# Keep output_types and workflow_output_types in spec for validators
|
||||
"output_types": final_output_types,
|
||||
"workflow_output_types": final_workflow_output_types,
|
||||
"ctx_annotation": ctx_annotation,
|
||||
}
|
||||
|
||||
return wrapper
|
||||
|
||||
# If func is provided, this means @response_handler was used without parentheses
|
||||
if func is not None:
|
||||
return decorator(func)
|
||||
|
||||
# Otherwise, return the wrapper for @response_handler(...) with parameters
|
||||
return decorator
|
||||
|
||||
|
||||
# endregion: Handler Decorator
|
||||
|
||||
# region Response Handler Validation
|
||||
|
||||
|
||||
def _validate_response_handler_signature(
|
||||
func: Callable[..., Any],
|
||||
*,
|
||||
skip_annotations: bool = False,
|
||||
) -> tuple[type | None, type | None, Any, list[type[Any] | UnionType], list[type[Any] | UnionType]]:
|
||||
"""Validate function signature for response handler functions.
|
||||
|
||||
Args:
|
||||
func: The function to validate
|
||||
skip_annotations: If True, skip validation that request/response parameters have type
|
||||
annotations. Used when types are explicitly provided to the @response_handler decorator.
|
||||
|
||||
Returns:
|
||||
Tuple of (request_type, response_type, ctx_annotation, output_types, workflow_output_types).
|
||||
request_type and response_type may be None if skip_annotations is True and no annotations exist.
|
||||
|
||||
Raises:
|
||||
ValueError: If the function signature is invalid
|
||||
"""
|
||||
signature = inspect.signature(func)
|
||||
params = list(signature.parameters.values())
|
||||
|
||||
# Note that the original_request parameter must be the second parameter
|
||||
# such that we can wrap the handler with functools.partial to bind it
|
||||
# to the original request when registering the handler, while maintaining
|
||||
# the order of parameters as if the response handler is a normal handler.
|
||||
expected_counts = 4 # self, original_request, message, ctx
|
||||
param_description = "(self, original_request, response, ctx)"
|
||||
if len(params) != expected_counts:
|
||||
raise ValueError(
|
||||
f"Response handler {func.__name__} must have {param_description}. Got {len(params)} parameters."
|
||||
)
|
||||
|
||||
# Check original_request parameter exists and has annotation (unless skipped)
|
||||
original_request_param = params[1]
|
||||
if not skip_annotations and original_request_param.annotation == inspect.Parameter.empty:
|
||||
raise ValueError(
|
||||
f"Response handler {func.__name__} must have a type annotation for the original_request parameter"
|
||||
)
|
||||
|
||||
# Check response parameter has type annotation (unless skipped)
|
||||
response_param = params[2]
|
||||
if not skip_annotations and response_param.annotation == inspect.Parameter.empty:
|
||||
raise ValueError(f"Response handler {func.__name__} must have a type annotation for the response parameter")
|
||||
|
||||
# Validate ctx parameter is WorkflowContext and extract type args (if annotated)
|
||||
ctx_param = params[3]
|
||||
if ctx_param.annotation != inspect.Parameter.empty:
|
||||
output_types, workflow_output_types = validate_workflow_context_annotation(
|
||||
ctx_param.annotation, f"parameter '{ctx_param.name}'", "Response handler"
|
||||
)
|
||||
else:
|
||||
output_types, workflow_output_types = [], []
|
||||
|
||||
request_type = (
|
||||
original_request_param.annotation if original_request_param.annotation != inspect.Parameter.empty else None
|
||||
)
|
||||
response_type = response_param.annotation if response_param.annotation != inspect.Parameter.empty else None
|
||||
ctx_annotation = ctx_param.annotation if ctx_param.annotation != inspect.Parameter.empty else None
|
||||
|
||||
return request_type, response_type, ctx_annotation, output_types, workflow_output_types
|
||||
|
||||
|
||||
# endregion: Response Handler Validation
|
||||
@@ -0,0 +1,433 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import warnings
|
||||
from collections import defaultdict
|
||||
from collections.abc import AsyncGenerator, Sequence
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..exceptions import (
|
||||
WorkflowCheckpointException,
|
||||
WorkflowConvergenceException,
|
||||
)
|
||||
from ._checkpoint import CheckpointID, CheckpointStorage, WorkflowCheckpoint
|
||||
from ._const import EXECUTOR_STATE_KEY
|
||||
from ._edge import EdgeGroup
|
||||
from ._edge_runner import EdgeRunner, create_edge_runner
|
||||
from ._events import WorkflowEvent
|
||||
from ._executor import Executor
|
||||
from ._runner_context import (
|
||||
RunnerContext,
|
||||
WorkflowMessage,
|
||||
)
|
||||
from ._state import State
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def warn_runner_deprecated() -> None:
|
||||
"""Emit a deprecation warning when ``Runner`` is accessed from the public API.
|
||||
|
||||
``Runner`` remains importable from ``agent_framework`` for backward
|
||||
compatibility, but it is intended for internal use only and will be removed
|
||||
from the public API in a future version.
|
||||
"""
|
||||
warnings.warn(
|
||||
"`Runner` is deprecated and will be removed from the public API in a future version. "
|
||||
"It is intended for internal use only.",
|
||||
DeprecationWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
|
||||
|
||||
class RunnerImpl:
|
||||
"""A class to run a workflow in Pregel supersteps."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
edge_groups: Sequence[EdgeGroup],
|
||||
executors: dict[str, Executor],
|
||||
state: State,
|
||||
ctx: RunnerContext,
|
||||
workflow_name: str,
|
||||
graph_signature_hash: str,
|
||||
max_iterations: int = 100,
|
||||
) -> None:
|
||||
"""Initialize the runner with edges, state, and context.
|
||||
|
||||
Args:
|
||||
edge_groups: The edge groups of the workflow.
|
||||
executors: Map of executor IDs to executor instances.
|
||||
state: The state for the workflow.
|
||||
ctx: The runner context for the workflow.
|
||||
workflow_name: The name of the workflow, used for checkpoint labeling.
|
||||
graph_signature_hash: A hash representing the workflow graph topology for checkpoint validation.
|
||||
max_iterations: The maximum number of iterations to run.
|
||||
"""
|
||||
# Workflow instance related attributes
|
||||
self._executors = executors
|
||||
self._edge_runners = [create_edge_runner(group, executors) for group in edge_groups]
|
||||
self._edge_runner_map = self._parse_edge_runners(self._edge_runners)
|
||||
self._ctx = ctx
|
||||
self._workflow_name = workflow_name
|
||||
self._graph_signature_hash = graph_signature_hash
|
||||
|
||||
# Runner state related attributes
|
||||
self._iteration = 0
|
||||
self._max_iterations = max_iterations
|
||||
self._state = state
|
||||
|
||||
# Checkpointing related attributes
|
||||
self._resumed_from_checkpoint = False
|
||||
self._previous_checkpoint_id: CheckpointID | None = None
|
||||
|
||||
@property
|
||||
def context(self) -> RunnerContext:
|
||||
"""Get the runner context for message, event, and checkpoint handling."""
|
||||
return self._ctx
|
||||
|
||||
@property
|
||||
def state(self) -> State:
|
||||
"""Get the shared state for the workflow."""
|
||||
return self._state
|
||||
|
||||
def reset_iteration_count(self) -> None:
|
||||
"""Reset the iteration count to zero.
|
||||
|
||||
This is useful when the workflow resumes from a new set of messages.
|
||||
|
||||
Note:
|
||||
When a workflow is resumed from a response (for a request_info_event)
|
||||
or a checkpoint, the iteration count is normally NOT reset.
|
||||
"""
|
||||
self._iteration = 0
|
||||
|
||||
async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
|
||||
"""Run the workflow until no more messages are sent."""
|
||||
try:
|
||||
# Emit any events already produced prior to entering loop
|
||||
if await self._ctx.has_events():
|
||||
logger.info("Yielding pre-loop events")
|
||||
for event in await self._ctx.drain_events():
|
||||
yield event
|
||||
|
||||
# Create a checkpoint before a run starts. Checkpoints are usually considered to be created at the
|
||||
# end of an iteration, we can think of this checkpoint as being created at the end of "superstep 0"
|
||||
# which captures the states after which the start executor has run. Note that we execute the start
|
||||
# executor outside of the main iteration loop.
|
||||
if await self._ctx.has_messages() and self._iteration == 0 and not self._resumed_from_checkpoint:
|
||||
await self.create_checkpoint_if_enabled()
|
||||
|
||||
while self._iteration < self._max_iterations:
|
||||
logger.info(f"Starting superstep {self._iteration + 1}")
|
||||
yield WorkflowEvent.superstep_started(iteration=self._iteration + 1)
|
||||
|
||||
# Run iteration concurrently with live event streaming: we poll
|
||||
# for new events while the iteration coroutine progresses.
|
||||
iteration_task = asyncio.create_task(self._run_iteration())
|
||||
try:
|
||||
while not iteration_task.done():
|
||||
try:
|
||||
# Wait briefly for any new event; timeout allows progress checks
|
||||
event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
|
||||
yield event
|
||||
except asyncio.TimeoutError:
|
||||
# Periodically continue to let iteration advance
|
||||
continue
|
||||
except asyncio.CancelledError:
|
||||
# Propagate cancellation to the iteration task to avoid orphaned work
|
||||
iteration_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await iteration_task
|
||||
raise
|
||||
|
||||
# Propagate errors from iteration, but first surface any pending events
|
||||
try:
|
||||
await iteration_task
|
||||
except Exception:
|
||||
# Make sure failure-related events (like ExecutorFailedEvent) are surfaced
|
||||
if await self._ctx.has_events():
|
||||
for event in await self._ctx.drain_events():
|
||||
yield event
|
||||
raise
|
||||
self._iteration += 1
|
||||
|
||||
# Drain any straggler events emitted at tail end
|
||||
if await self._ctx.has_events():
|
||||
for event in await self._ctx.drain_events():
|
||||
yield event
|
||||
|
||||
logger.info(f"Completed superstep {self._iteration}")
|
||||
|
||||
# Commit pending state changes at superstep boundary
|
||||
self._state.commit()
|
||||
|
||||
# Create checkpoint after each superstep iteration
|
||||
await self.create_checkpoint_if_enabled()
|
||||
|
||||
yield WorkflowEvent.superstep_completed(iteration=self._iteration)
|
||||
|
||||
# Check for convergence: no more messages to process
|
||||
if not await self._ctx.has_messages():
|
||||
break
|
||||
|
||||
logger.info(f"Workflow completed after {self._iteration} supersteps")
|
||||
|
||||
if self._iteration >= self._max_iterations and await self._ctx.has_messages():
|
||||
raise WorkflowConvergenceException(f"Runner did not converge after {self._max_iterations} iterations.")
|
||||
finally:
|
||||
# Reset the resume flag so stale resume state never leaks into the next run on this
|
||||
# instance - even if convergence raised before completing (e.g. an executor failure
|
||||
# during a resumed run).
|
||||
self._resumed_from_checkpoint = False
|
||||
|
||||
async def _run_iteration(self) -> None:
|
||||
"""Run a single iteration of the workflow.
|
||||
|
||||
Messages are delivered through edge runners. A source executor may have multiple outgoing edge
|
||||
runners. All edge runners run concurrently, but messages sent through the same edge runner are
|
||||
delivered in the order they were sent to preserve message ordering guarantees per edge.
|
||||
|
||||
What this means in practice:
|
||||
- A message from a source to multiple target is delivered to all targets concurrently.
|
||||
- Multiple messages from a source to the same target are delivered in the order they were sent.
|
||||
- Multiple messages from different sources to the same target can be delivered to the target one
|
||||
at a time in any order, because true parallelism is not realized in Python.
|
||||
- Multiple message from different sources to different targets are delivered concurrently to all
|
||||
targets, assuming each message is targeting a unique target, or it falls back to the previous
|
||||
rules if there are multiple messages targeting the same target.
|
||||
- Special case: if using a fan-out edge runner (or derived edge runner that replicates messages
|
||||
to multiple targets such as multi-selection or switch-case) to send messages to targets from
|
||||
a source by specifying the target, the messages will be delivered to the specified targets
|
||||
in the order they were sent. This is because all messages go through the same edge runner instance
|
||||
which preserves message order.
|
||||
"""
|
||||
|
||||
async def _deliver_messages(source_executor_id: str, source_messages: list[WorkflowMessage]) -> None:
|
||||
"""Outer loop to concurrently deliver messages from all sources to their targets."""
|
||||
|
||||
async def _deliver_message_inner(edge_runner: EdgeRunner, message: WorkflowMessage) -> bool:
|
||||
"""Inner loop to deliver a single message through an edge runner."""
|
||||
return await edge_runner.send_message(message, self._state, self._ctx)
|
||||
|
||||
async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None:
|
||||
# Preserve message order per edge runner (and therefore per routed target path)
|
||||
# while still allowing parallelism across different edge runners.
|
||||
for message in source_messages:
|
||||
await _deliver_message_inner(edge_runner, message)
|
||||
|
||||
# Route all messages through normal workflow edges
|
||||
associated_edge_runners = self._edge_runner_map.get(source_executor_id, [])
|
||||
if not associated_edge_runners:
|
||||
# This is expected for terminal nodes (e.g., EndWorkflow, last action in workflow)
|
||||
logger.debug(f"No outgoing edges found for executor {source_executor_id}; dropping messages.")
|
||||
return
|
||||
|
||||
tasks = [_deliver_messages_for_edge_runner(edge_runner) for edge_runner in associated_edge_runners]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
message_batches = await self._ctx.drain_messages()
|
||||
tasks = [
|
||||
_deliver_messages(source_executor_id, source_messages)
|
||||
for source_executor_id, source_messages in message_batches.items()
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
async def _prepare_checkpoint_state(self) -> None:
|
||||
"""Persist executor snapshots into committed shared state.
|
||||
|
||||
This is used by checkpoint capture paths that need a complete, restorable
|
||||
state payload without necessarily writing to a checkpoint storage backend.
|
||||
"""
|
||||
await self._save_executor_states()
|
||||
self._state.commit()
|
||||
|
||||
async def create_checkpoint_if_enabled(self) -> None:
|
||||
"""Create a checkpoint if checkpointing is enabled and attach a label and metadata."""
|
||||
if not self._ctx.has_checkpointing():
|
||||
return
|
||||
|
||||
try:
|
||||
# Save executor states into committed state before creating the checkpoint.
|
||||
await self._prepare_checkpoint_state()
|
||||
|
||||
checkpoint_id = await self._ctx.create_checkpoint(
|
||||
self._workflow_name,
|
||||
self._graph_signature_hash,
|
||||
self._state,
|
||||
self._previous_checkpoint_id,
|
||||
self._iteration,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Created checkpoint: %s with parent checkpoint at iteration %d: %s",
|
||||
checkpoint_id,
|
||||
self._iteration,
|
||||
self._previous_checkpoint_id,
|
||||
)
|
||||
self._previous_checkpoint_id = checkpoint_id
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to create checkpoint at iteration %d: %s. "
|
||||
"Note that this does not fail the workflow run. "
|
||||
"The next successfully-created checkpoint will be parented to the last successful checkpoint: %s",
|
||||
self._iteration,
|
||||
e,
|
||||
self._previous_checkpoint_id,
|
||||
)
|
||||
|
||||
async def restore_from_checkpoint(
|
||||
self,
|
||||
checkpoint_id: CheckpointID,
|
||||
checkpoint_storage: CheckpointStorage | None = None,
|
||||
) -> None:
|
||||
"""Restore the runner from a checkpoint.
|
||||
|
||||
Args:
|
||||
checkpoint_id: The ID of the checkpoint to restore from
|
||||
checkpoint_storage: Optional storage to load checkpoints from when the
|
||||
runner context itself is not configured with checkpointing.
|
||||
|
||||
Returns:
|
||||
None on success.
|
||||
|
||||
Raises:
|
||||
WorkflowCheckpointException on failure.
|
||||
"""
|
||||
try:
|
||||
# Load the checkpoint
|
||||
checkpoint: WorkflowCheckpoint | None
|
||||
if self._ctx.has_checkpointing():
|
||||
checkpoint = await self._ctx.load_checkpoint(checkpoint_id)
|
||||
elif checkpoint_storage is not None:
|
||||
checkpoint = await checkpoint_storage.load(checkpoint_id)
|
||||
else:
|
||||
raise WorkflowCheckpointException(
|
||||
"Cannot load checkpoint: no checkpointing configured in context or external storage provided."
|
||||
)
|
||||
|
||||
if not checkpoint:
|
||||
logger.error(f"Checkpoint {checkpoint_id} not found")
|
||||
raise WorkflowCheckpointException(f"Checkpoint {checkpoint_id} not found")
|
||||
|
||||
# Validate the loaded checkpoint against the workflow
|
||||
if self._graph_signature_hash != checkpoint.graph_signature_hash:
|
||||
raise WorkflowCheckpointException(
|
||||
"Workflow graph has changed since the checkpoint was created. "
|
||||
"Please rebuild the original workflow before resuming."
|
||||
)
|
||||
|
||||
# Restore state. Clear first so import_state (which merges) does
|
||||
# not leak stale keys from a prior run on this Workflow instance.
|
||||
# This matters more now that Workflow.run() no longer wipes state
|
||||
# per call - the only reset point for shared state on a reused
|
||||
# instance is at restore time.
|
||||
self._state.clear()
|
||||
self._state.import_state(checkpoint.state)
|
||||
# Restore executor states using the restored state
|
||||
await self._restore_executor_states()
|
||||
# Apply the checkpoint to the context
|
||||
await self._ctx.apply_checkpoint(checkpoint)
|
||||
# Mark the runner as resumed
|
||||
self._mark_resumed(checkpoint)
|
||||
|
||||
logger.info(f"Successfully restored workflow from checkpoint: {checkpoint_id}")
|
||||
except WorkflowCheckpointException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}")
|
||||
raise WorkflowCheckpointException(f"Failed to restore from checkpoint {checkpoint_id}") from e
|
||||
|
||||
async def _save_executor_states(self) -> None:
|
||||
"""Populate executor state by calling checkpoint hooks on executors."""
|
||||
for exec_id, executor in self._executors.items():
|
||||
# Try the updated behavior only if backward compatibility did not yield state
|
||||
try:
|
||||
state_dict = await executor.on_checkpoint_save()
|
||||
await self._set_executor_state(exec_id, state_dict)
|
||||
except WorkflowCheckpointException:
|
||||
raise
|
||||
except Exception as ex: # pragma: no cover
|
||||
raise WorkflowCheckpointException(f"Executor {exec_id} on_checkpoint_save failed") from ex
|
||||
|
||||
async def _restore_executor_states(self) -> None:
|
||||
"""Restore executor state by calling restore hooks on executors."""
|
||||
has_executor_states = self._state.has(EXECUTOR_STATE_KEY)
|
||||
if not has_executor_states:
|
||||
return
|
||||
|
||||
executor_states = self._state.get(EXECUTOR_STATE_KEY)
|
||||
if not isinstance(executor_states, dict):
|
||||
raise WorkflowCheckpointException("Executor states in shared state is not a dictionary. Unable to restore.")
|
||||
|
||||
for executor_id, state in executor_states.items(): # pyright: ignore[reportUnknownVariableType]
|
||||
if not isinstance(executor_id, str):
|
||||
raise WorkflowCheckpointException("Executor ID in executor states is not a string. Unable to restore.")
|
||||
if not isinstance(state, dict) or not all(isinstance(k, str) for k in state): # pyright: ignore[reportUnknownVariableType]
|
||||
raise WorkflowCheckpointException(
|
||||
f"Executor state for {executor_id} is not a dict[str, Any]. Unable to restore."
|
||||
)
|
||||
|
||||
executor = self._executors.get(executor_id)
|
||||
if not executor:
|
||||
raise WorkflowCheckpointException(f"Executor {executor_id} not found during state restoration.")
|
||||
|
||||
# Try the updated behavior only if backward compatibility did not restore
|
||||
try:
|
||||
await executor.on_checkpoint_restore(state) # pyright: ignore[reportUnknownArgumentType]
|
||||
except Exception as ex: # pragma: no cover - defensive
|
||||
raise WorkflowCheckpointException(f"Executor {executor_id} on_checkpoint_restore failed") from ex
|
||||
|
||||
def _parse_edge_runners(self, edge_runners: list[EdgeRunner]) -> dict[str, list[EdgeRunner]]:
|
||||
"""Parse the edge runners of the workflow into a mapping where each source executor ID maps to its edge runners.
|
||||
|
||||
Args:
|
||||
edge_runners: A list of edge runners in the workflow.
|
||||
|
||||
Returns:
|
||||
A dictionary mapping each source executor ID to a list of edge runners.
|
||||
"""
|
||||
parsed: defaultdict[str, list[EdgeRunner]] = defaultdict(list)
|
||||
for runner in edge_runners:
|
||||
# Accessing protected attribute (_edge_group) intentionally for internal wiring.
|
||||
for source_executor_id in runner._edge_group.source_executor_ids: # type: ignore[attr-defined]
|
||||
parsed[source_executor_id].append(runner)
|
||||
|
||||
return parsed
|
||||
|
||||
def _mark_resumed(self, checkpoint: WorkflowCheckpoint) -> None:
|
||||
"""Mark the runner as having resumed from a checkpoint.
|
||||
|
||||
Optionally set the current iteration and max iterations.
|
||||
"""
|
||||
self._resumed_from_checkpoint = True
|
||||
self._iteration = checkpoint.iteration_count
|
||||
self._previous_checkpoint_id = checkpoint.checkpoint_id
|
||||
|
||||
async def _set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:
|
||||
"""Store executor state in state under a reserved key.
|
||||
|
||||
Executors call this with a JSON-serializable dict capturing the minimal
|
||||
state needed to resume. It replaces any previously stored state.
|
||||
"""
|
||||
existing_states = self._state.get(EXECUTOR_STATE_KEY, {})
|
||||
|
||||
if not isinstance(existing_states, dict):
|
||||
raise WorkflowCheckpointException("Existing executor states in state is not a dictionary.")
|
||||
|
||||
existing_states[executor_id] = state
|
||||
self._state.set(EXECUTOR_STATE_KEY, existing_states)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
Runner = RunnerImpl
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
"""Lazily expose deprecated module-level public names."""
|
||||
if name == "Runner":
|
||||
warn_runner_deprecated()
|
||||
return RunnerImpl
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,518 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from copy import copy
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any, Literal, Protocol, TypeVar, runtime_checkable
|
||||
|
||||
from ._checkpoint import CheckpointID, CheckpointStorage, WorkflowCheckpoint
|
||||
from ._const import INTERNAL_SOURCE_ID
|
||||
from ._events import WorkflowEvent
|
||||
from ._state import State
|
||||
from ._typing_utils import is_instance_of
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
YieldOutputEventType = Literal["output", "intermediate"]
|
||||
YieldOutputClassifier = Callable[[str], YieldOutputEventType | None]
|
||||
|
||||
|
||||
class MessageType(Enum):
|
||||
"""Enumeration of WorkflowMessage types in the workflow."""
|
||||
|
||||
STANDARD = "standard"
|
||||
"""A standard WorkflowMessage between executors."""
|
||||
|
||||
RESPONSE = "response"
|
||||
"""A response WorkflowMessage to a pending request."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkflowMessage:
|
||||
"""A class representing a WorkflowMessage in the workflow."""
|
||||
|
||||
data: Any
|
||||
source_id: str
|
||||
target_id: str | None = None
|
||||
type: MessageType = MessageType.STANDARD
|
||||
|
||||
# OpenTelemetry trace context fields for WorkflowMessage propagation
|
||||
# These are plural to support fan-in scenarios where multiple messages are aggregated
|
||||
trace_contexts: list[dict[str, str]] | None = None # W3C Trace Context headers from multiple sources
|
||||
source_span_ids: list[str] | None = None # Publishing span IDs for linking from multiple sources
|
||||
|
||||
# For response messages, the original request data
|
||||
original_request_info_event: WorkflowEvent[Any] | None = None
|
||||
|
||||
# Backward compatibility properties
|
||||
@property
|
||||
def trace_context(self) -> dict[str, str] | None:
|
||||
"""Get the first trace context for backward compatibility."""
|
||||
return self.trace_contexts[0] if self.trace_contexts else None
|
||||
|
||||
@property
|
||||
def source_span_id(self) -> str | None:
|
||||
"""Get the first source span ID for backward compatibility."""
|
||||
return self.source_span_ids[0] if self.source_span_ids else None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert the WorkflowMessage to a dictionary for serialization."""
|
||||
return {
|
||||
"data": self.data,
|
||||
"source_id": self.source_id,
|
||||
"target_id": self.target_id,
|
||||
"type": self.type.value,
|
||||
"trace_contexts": self.trace_contexts,
|
||||
"source_span_ids": self.source_span_ids,
|
||||
"original_request_info_event": self.original_request_info_event,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def from_dict(data: dict[str, Any]) -> WorkflowMessage:
|
||||
"""Create a WorkflowMessage from a dictionary."""
|
||||
# Validation
|
||||
if "data" not in data:
|
||||
raise KeyError("Missing 'data' field in WorkflowMessage dictionary.")
|
||||
|
||||
if "source_id" not in data:
|
||||
raise KeyError("Missing 'source_id' field in WorkflowMessage dictionary.")
|
||||
|
||||
return WorkflowMessage(
|
||||
data=data["data"],
|
||||
source_id=data["source_id"],
|
||||
target_id=data.get("target_id"),
|
||||
type=MessageType(data.get("type", "standard")),
|
||||
trace_contexts=data.get("trace_contexts"),
|
||||
source_span_ids=data.get("source_span_ids"),
|
||||
original_request_info_event=data.get("original_request_info_event"),
|
||||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RunnerContext(Protocol):
|
||||
"""Protocol for the execution context used by the runner.
|
||||
|
||||
A single context that supports messaging, events, and optional checkpointing.
|
||||
If checkpoint storage is not configured, checkpoint methods may raise.
|
||||
"""
|
||||
|
||||
async def send_message(self, message: WorkflowMessage) -> None:
|
||||
"""Send a WorkflowMessage from the executor to the context.
|
||||
|
||||
Args:
|
||||
message: The WorkflowMessage to be sent.
|
||||
"""
|
||||
...
|
||||
|
||||
async def drain_messages(self) -> dict[str, list[WorkflowMessage]]:
|
||||
"""Drain all messages from the context.
|
||||
|
||||
Returns:
|
||||
A dictionary mapping executor IDs to lists of messages.
|
||||
"""
|
||||
...
|
||||
|
||||
async def has_messages(self) -> bool:
|
||||
"""Check if there are any messages in the context.
|
||||
|
||||
Returns:
|
||||
True if there are messages, False otherwise.
|
||||
"""
|
||||
...
|
||||
|
||||
async def add_event(self, event: WorkflowEvent) -> None:
|
||||
"""Add an event to the execution context.
|
||||
|
||||
Args:
|
||||
event: The event to be added.
|
||||
"""
|
||||
...
|
||||
|
||||
async def drain_events(self) -> list[WorkflowEvent]:
|
||||
"""Drain all events from the context.
|
||||
|
||||
Returns:
|
||||
A list of events that were added to the context.
|
||||
"""
|
||||
...
|
||||
|
||||
async def has_events(self) -> bool:
|
||||
"""Check if there are any events in the context.
|
||||
|
||||
Returns:
|
||||
True if there are events, False otherwise.
|
||||
"""
|
||||
...
|
||||
|
||||
async def next_event(self) -> WorkflowEvent: # pragma: no cover - interface only
|
||||
"""Wait for and return the next event emitted by the workflow run."""
|
||||
...
|
||||
|
||||
# Checkpointing capability
|
||||
def has_checkpointing(self) -> bool:
|
||||
"""Check if the context supports checkpointing.
|
||||
|
||||
Returns:
|
||||
True if checkpointing is supported, False otherwise.
|
||||
"""
|
||||
...
|
||||
|
||||
def set_runtime_checkpoint_storage(self, storage: CheckpointStorage) -> None:
|
||||
"""Set runtime checkpoint storage to override build-time configuration.
|
||||
|
||||
Args:
|
||||
storage: The checkpoint storage to use for this run.
|
||||
"""
|
||||
...
|
||||
|
||||
def clear_runtime_checkpoint_storage(self) -> None:
|
||||
"""Clear runtime checkpoint storage override."""
|
||||
...
|
||||
|
||||
def reset_for_new_run(self) -> None:
|
||||
"""Reset the context for a new workflow run."""
|
||||
...
|
||||
|
||||
def set_streaming(self, streaming: bool) -> None:
|
||||
"""Set whether agents should stream incremental updates.
|
||||
|
||||
Args:
|
||||
streaming: True for streaming mode (stream=True), False for non-streaming (stream=False).
|
||||
"""
|
||||
...
|
||||
|
||||
def is_streaming(self) -> bool:
|
||||
"""Check if the workflow is in streaming mode.
|
||||
|
||||
Returns:
|
||||
True if streaming mode is enabled, False otherwise.
|
||||
"""
|
||||
...
|
||||
|
||||
async def create_checkpoint(
|
||||
self,
|
||||
workflow_name: str,
|
||||
graph_signature_hash: str,
|
||||
state: State,
|
||||
previous_checkpoint_id: CheckpointID | None,
|
||||
iteration_count: int,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> CheckpointID:
|
||||
"""Create a checkpoint of the current workflow state.
|
||||
|
||||
Args:
|
||||
workflow_name: The name of the workflow for which the checkpoint is being created.
|
||||
graph_signature_hash: Hash of the workflow graph topology to
|
||||
validate checkpoint compatibility during restore.
|
||||
state: The state to include in the checkpoint.
|
||||
This is needed to capture the full state of the workflow.
|
||||
The state is not managed by the context itself.
|
||||
previous_checkpoint_id: The ID of the previous checkpoint, if any, to form a checkpoint chain.
|
||||
iteration_count: The current iteration count of the workflow.
|
||||
metadata: Optional metadata to associate with the checkpoint.
|
||||
|
||||
Returns:
|
||||
The ID of the created checkpoint.
|
||||
"""
|
||||
...
|
||||
|
||||
async def load_checkpoint(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint | None:
|
||||
"""Load a checkpoint without mutating the current context state.
|
||||
|
||||
Args:
|
||||
checkpoint_id: The ID of the checkpoint to load.
|
||||
|
||||
Returns:
|
||||
The loaded checkpoint, or None if it does not exist.
|
||||
"""
|
||||
...
|
||||
|
||||
async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
|
||||
"""Apply a checkpoint to the current context, mutating its state.
|
||||
|
||||
Args:
|
||||
checkpoint: The checkpoint whose state is to be applied.
|
||||
"""
|
||||
...
|
||||
|
||||
async def add_request_info_event(self, event: WorkflowEvent[Any]) -> None:
|
||||
"""Add a request_info event to the context and track it for correlation.
|
||||
|
||||
Args:
|
||||
event: The WorkflowEvent with type='request_info' to be added.
|
||||
"""
|
||||
...
|
||||
|
||||
async def send_request_info_response(self, request_id: str, response: Any) -> None:
|
||||
"""Send a response correlated to a pending request.
|
||||
|
||||
Args:
|
||||
request_id: The ID of the original request.
|
||||
response: The response data to be sent.
|
||||
"""
|
||||
...
|
||||
|
||||
async def get_pending_request_info_events(self) -> dict[str, WorkflowEvent[Any]]:
|
||||
"""Get the mapping of request IDs to their corresponding request_info events.
|
||||
|
||||
Returns:
|
||||
A dictionary mapping request IDs to their corresponding WorkflowEvent (type='request_info').
|
||||
"""
|
||||
...
|
||||
|
||||
def set_yield_output_classifier(self, classifier: YieldOutputClassifier) -> None:
|
||||
"""Set the classifier used by WorkflowContext.yield_output()."""
|
||||
...
|
||||
|
||||
def classify_yielded_output(self, executor_id: str) -> YieldOutputEventType | None:
|
||||
"""Classify an executor's yield_output payload as output, intermediate, or hidden."""
|
||||
...
|
||||
|
||||
|
||||
class InProcRunnerContext:
|
||||
"""In-process execution context for local execution and optional checkpointing."""
|
||||
|
||||
def __init__(self, checkpoint_storage: CheckpointStorage | None = None):
|
||||
"""Initialize the in-process execution context.
|
||||
|
||||
Args:
|
||||
checkpoint_storage: Optional storage to enable checkpointing.
|
||||
"""
|
||||
self._messages: dict[str, list[WorkflowMessage]] = {}
|
||||
|
||||
# The queue must be created lazily under the running loop (see ``_get_event_queue``)
|
||||
# in order to support contexts that are constructed outside an event loop and reused
|
||||
# across multiple async loops (e.g., successive ``asyncio.run`` calls on the same workflow).
|
||||
# Binding a single queue to one loop would raise "bound to a different event loop" on reuse.
|
||||
self._event_queue: asyncio.Queue[WorkflowEvent] | None = None
|
||||
self._event_queue_loop: asyncio.AbstractEventLoop | None = None
|
||||
|
||||
# An additional storage for pending request info events
|
||||
self._pending_request_info_events: dict[str, WorkflowEvent[Any]] = {}
|
||||
|
||||
# Checkpointing configuration/state
|
||||
self._checkpoint_storage = checkpoint_storage
|
||||
self._runtime_checkpoint_storage: CheckpointStorage | None = None
|
||||
|
||||
# Streaming flag - set by workflow's run(..., stream=True) vs run(..., stream=False)
|
||||
self._streaming: bool = False
|
||||
self._yield_output_classifier: YieldOutputClassifier = lambda _executor_id: "output"
|
||||
|
||||
# region Messaging and Events
|
||||
async def send_message(self, message: WorkflowMessage) -> None:
|
||||
self._messages.setdefault(message.source_id, [])
|
||||
self._messages[message.source_id].append(message)
|
||||
|
||||
async def drain_messages(self) -> dict[str, list[WorkflowMessage]]:
|
||||
messages = copy(self._messages)
|
||||
self._messages.clear()
|
||||
return messages
|
||||
|
||||
async def has_messages(self) -> bool:
|
||||
return bool(self._messages)
|
||||
|
||||
def _get_event_queue(self) -> asyncio.Queue[WorkflowEvent]:
|
||||
"""Return the event queue bound to the running loop, re-creating it on loop change."""
|
||||
loop = asyncio.get_running_loop()
|
||||
if self._event_queue is None or self._event_queue_loop is not loop:
|
||||
self._event_queue = asyncio.Queue()
|
||||
self._event_queue_loop = loop
|
||||
return self._event_queue
|
||||
|
||||
async def add_event(self, event: WorkflowEvent) -> None:
|
||||
"""Add an event to the context immediately.
|
||||
|
||||
Events are enqueued so runners can stream them in real time instead of
|
||||
waiting for superstep boundaries.
|
||||
"""
|
||||
await self._get_event_queue().put(event)
|
||||
|
||||
async def drain_events(self) -> list[WorkflowEvent]:
|
||||
"""Drain all currently queued events without blocking for new ones."""
|
||||
events: list[WorkflowEvent] = []
|
||||
queue = self._get_event_queue()
|
||||
while True:
|
||||
try:
|
||||
events.append(queue.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
return events
|
||||
|
||||
async def has_events(self) -> bool:
|
||||
return not self._get_event_queue().empty()
|
||||
|
||||
async def next_event(self) -> WorkflowEvent:
|
||||
"""Wait for and return the next event.
|
||||
|
||||
Used by the runner to interleave event emission with ongoing iteration work.
|
||||
"""
|
||||
return await self._get_event_queue().get()
|
||||
|
||||
# endregion Messaging and Events
|
||||
|
||||
# region Checkpointing
|
||||
|
||||
def _get_effective_checkpoint_storage(self) -> CheckpointStorage | None:
|
||||
"""Get the effective checkpoint storage (runtime override or build-time)."""
|
||||
return self._runtime_checkpoint_storage or self._checkpoint_storage
|
||||
|
||||
def set_runtime_checkpoint_storage(self, storage: CheckpointStorage) -> None:
|
||||
"""Set runtime checkpoint storage to override build-time configuration.
|
||||
|
||||
Args:
|
||||
storage: The checkpoint storage to use for this run.
|
||||
"""
|
||||
self._runtime_checkpoint_storage = storage
|
||||
|
||||
def clear_runtime_checkpoint_storage(self) -> None:
|
||||
"""Clear runtime checkpoint storage override.
|
||||
|
||||
This is called automatically by workflow execution methods after a run completes,
|
||||
ensuring runtime storage doesn't leak across runs.
|
||||
"""
|
||||
self._runtime_checkpoint_storage = None
|
||||
|
||||
def has_checkpointing(self) -> bool:
|
||||
return self._get_effective_checkpoint_storage() is not None
|
||||
|
||||
async def create_checkpoint(
|
||||
self,
|
||||
workflow_name: str,
|
||||
graph_signature_hash: str,
|
||||
state: State,
|
||||
previous_checkpoint_id: CheckpointID | None,
|
||||
iteration_count: int,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> CheckpointID:
|
||||
storage = self._get_effective_checkpoint_storage()
|
||||
if not storage:
|
||||
raise ValueError("Checkpoint storage not configured")
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
workflow_name=workflow_name,
|
||||
graph_signature_hash=graph_signature_hash,
|
||||
previous_checkpoint_id=previous_checkpoint_id,
|
||||
messages=dict(self._messages),
|
||||
state=state.export_state(),
|
||||
pending_request_info_events=dict(self._pending_request_info_events),
|
||||
iteration_count=iteration_count,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
checkpoint_id = await storage.save(checkpoint)
|
||||
logger.debug(f"Created checkpoint {checkpoint_id}")
|
||||
return checkpoint_id
|
||||
|
||||
async def load_checkpoint(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
|
||||
storage = self._get_effective_checkpoint_storage()
|
||||
if not storage:
|
||||
raise ValueError("Checkpoint storage not configured")
|
||||
return await storage.load(checkpoint_id)
|
||||
|
||||
def reset_for_new_run(self) -> None:
|
||||
"""Reset the context for a new workflow run.
|
||||
|
||||
This clears messages, events, and resets streaming flag.
|
||||
Runtime checkpoint storage is NOT cleared here as it's managed at the workflow level.
|
||||
"""
|
||||
self._messages.clear()
|
||||
# Drop any pending events. The queue and its loop marker are cleared so the queue
|
||||
# rebinds lazily under the running loop on next use.
|
||||
self._event_queue = None
|
||||
self._event_queue_loop = None
|
||||
self._streaming = False # Reset streaming flag
|
||||
|
||||
async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
|
||||
"""Apply a checkpoint to the current context, mutating its state."""
|
||||
# Restore messages
|
||||
self._messages.clear()
|
||||
messages_data = checkpoint.messages
|
||||
for source_id, message_list in messages_data.items():
|
||||
self._messages[source_id] = list(message_list)
|
||||
|
||||
# Restore pending request info events
|
||||
self._pending_request_info_events.clear()
|
||||
for request_id, request_info_event in checkpoint.pending_request_info_events.items():
|
||||
self._pending_request_info_events[request_id] = request_info_event
|
||||
await self.add_event(request_info_event)
|
||||
|
||||
# endregion Checkpointing
|
||||
|
||||
def set_streaming(self, streaming: bool) -> None:
|
||||
"""Set whether agents should stream incremental updates.
|
||||
|
||||
Args:
|
||||
streaming: True for streaming mode (run(stream=True)), False for non-streaming.
|
||||
"""
|
||||
self._streaming = streaming
|
||||
|
||||
def is_streaming(self) -> bool:
|
||||
"""Check if the workflow is in streaming mode.
|
||||
|
||||
Returns:
|
||||
True if streaming mode is enabled, False otherwise.
|
||||
"""
|
||||
return self._streaming
|
||||
|
||||
async def add_request_info_event(self, event: WorkflowEvent[Any]) -> None:
|
||||
"""Add a request_info event to the context and track it for correlation.
|
||||
|
||||
Args:
|
||||
event: The WorkflowEvent with type='request_info' to be added.
|
||||
"""
|
||||
if event.type != "request_info":
|
||||
raise ValueError("Event type must be 'request_info'")
|
||||
self._pending_request_info_events[event.request_id] = event
|
||||
await self.add_event(event)
|
||||
|
||||
async def send_request_info_response(self, request_id: str, response: Any) -> None:
|
||||
"""Send a response correlated to a pending request.
|
||||
|
||||
Args:
|
||||
request_id: The ID of the original request.
|
||||
response: The response data to be sent.
|
||||
"""
|
||||
event = self._pending_request_info_events.pop(request_id, None)
|
||||
if not event:
|
||||
raise ValueError(f"No pending request found for request_id: {request_id}")
|
||||
|
||||
# Validate response type if specified
|
||||
if event.response_type and not is_instance_of(response, event.response_type):
|
||||
raise TypeError(
|
||||
f"Response type mismatch for request_id {request_id}: "
|
||||
f"expected {event.response_type.__name__}, got {type(response).__name__}"
|
||||
)
|
||||
|
||||
source_executor_id = event.source_executor_id
|
||||
|
||||
# Create ResponseMessage instance
|
||||
response_msg = WorkflowMessage(
|
||||
data=response,
|
||||
source_id=INTERNAL_SOURCE_ID(source_executor_id),
|
||||
target_id=source_executor_id,
|
||||
type=MessageType.RESPONSE,
|
||||
original_request_info_event=event,
|
||||
)
|
||||
|
||||
await self.send_message(response_msg)
|
||||
|
||||
async def get_pending_request_info_events(self) -> dict[str, WorkflowEvent[Any]]:
|
||||
"""Get the mapping of request IDs to their corresponding request_info events.
|
||||
|
||||
Returns:
|
||||
A dictionary mapping request IDs to their corresponding WorkflowEvent (type='request_info').
|
||||
"""
|
||||
return dict(self._pending_request_info_events)
|
||||
|
||||
def set_yield_output_classifier(self, classifier: YieldOutputClassifier) -> None:
|
||||
"""Set the classifier used by WorkflowContext.yield_output()."""
|
||||
self._yield_output_classifier = classifier
|
||||
|
||||
def classify_yielded_output(self, executor_id: str) -> YieldOutputEventType | None:
|
||||
"""Classify an executor's yield_output payload as output, intermediate, or hidden."""
|
||||
return self._yield_output_classifier(executor_id)
|
||||
@@ -0,0 +1,127 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
class State:
|
||||
"""Manages shared state across executors within a workflow.
|
||||
|
||||
State provides access to workflow state data that is shared across executors
|
||||
during workflow execution. It implements superstep caching semantics where
|
||||
writes are staged in a pending buffer and only committed to the actual state
|
||||
at superstep boundaries.
|
||||
|
||||
Superstep Semantics:
|
||||
- `set()` writes to a pending buffer, not directly to committed state
|
||||
- `get()` checks pending buffer first, then committed state
|
||||
- `commit()` moves all pending changes to committed state (called by Runner at superstep boundary)
|
||||
- `discard()` clears pending changes without committing
|
||||
|
||||
Reserved Keys:
|
||||
Keys starting with underscore (_) are reserved for internal framework use.
|
||||
Do not use these in user code.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize the state."""
|
||||
self._committed: dict[str, Any] = {}
|
||||
self._pending: dict[str, Any] = {}
|
||||
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
"""Set a value in the pending state buffer.
|
||||
|
||||
The value will be visible to subsequent `get()` calls but won't be
|
||||
committed to the actual state until `commit()` is called.
|
||||
|
||||
Note:
|
||||
When multiple executors run concurrently within the same superstep,
|
||||
each executor's writes go to the same pending buffer. The last write
|
||||
for a given key wins when commit() is called. This is consistent with
|
||||
the .NET behavior and the superstep execution model where all executors
|
||||
in a superstep see the same committed state at the start.
|
||||
"""
|
||||
self._pending[key] = value
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
"""Get a value from state, checking pending first then committed.
|
||||
|
||||
Args:
|
||||
key: The key to retrieve.
|
||||
default: Value to return if key is not found. Defaults to None.
|
||||
|
||||
Returns:
|
||||
The value if found, otherwise the default value.
|
||||
"""
|
||||
if key in self._pending:
|
||||
value = self._pending[key]
|
||||
if value is _DeleteSentinel:
|
||||
return default
|
||||
return value
|
||||
return self._committed.get(key, default)
|
||||
|
||||
def has(self, key: str) -> bool:
|
||||
"""Check if a key exists in pending or committed state."""
|
||||
if key in self._pending:
|
||||
return self._pending[key] is not _DeleteSentinel
|
||||
return key in self._committed
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
"""Mark a key for deletion.
|
||||
|
||||
If the key exists in committed state, a sentinel is stored in pending
|
||||
to indicate deletion at commit time. If it only exists in pending,
|
||||
it is removed from pending.
|
||||
"""
|
||||
if key not in self._pending and key not in self._committed:
|
||||
raise KeyError(f"Key '{key}' not found in state.")
|
||||
|
||||
if key in self._committed:
|
||||
# Mark for deletion from committed state at commit time
|
||||
self._pending[key] = _DeleteSentinel
|
||||
elif key in self._pending:
|
||||
# Only exists in pending, safe to just remove
|
||||
del self._pending[key]
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear both committed and pending state."""
|
||||
self._committed.clear()
|
||||
self._pending.clear()
|
||||
|
||||
def commit(self) -> None:
|
||||
"""Commit pending changes to the committed state.
|
||||
|
||||
Called by the Runner at superstep boundaries after successful execution.
|
||||
"""
|
||||
for key, value in self._pending.items():
|
||||
if value is _DeleteSentinel:
|
||||
self._committed.pop(key, None)
|
||||
else:
|
||||
self._committed[key] = value
|
||||
self._pending.clear()
|
||||
|
||||
def discard(self) -> None:
|
||||
"""Discard all pending changes without committing."""
|
||||
self._pending.clear()
|
||||
|
||||
def export_state(self) -> dict[str, Any]:
|
||||
"""Export a serialized copy of the committed state.
|
||||
|
||||
Note: Does not include pending changes.
|
||||
"""
|
||||
return dict(self._committed)
|
||||
|
||||
def import_state(self, state: dict[str, Any]) -> None:
|
||||
"""Import state from a serialized dictionary.
|
||||
|
||||
Merges into committed state. Does not affect pending changes.
|
||||
"""
|
||||
self._committed.update(state)
|
||||
|
||||
|
||||
class _DeleteSentinelType:
|
||||
"""Sentinel type to mark keys for deletion in pending state."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
_DeleteSentinel = _DeleteSentinelType()
|
||||
@@ -0,0 +1,414 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import typing
|
||||
from types import UnionType
|
||||
from typing import Any, TypeGuard, Union, cast, get_args, get_origin
|
||||
|
||||
import typing_extensions
|
||||
|
||||
from .._agents import Agent
|
||||
|
||||
# Pre-compute the TypeVar types for runtime-safe detection.
|
||||
# isinstance(x, TypeVar) can fail if TypeVar is a factory/callable
|
||||
# on some Python versions, so we compare against the actual runtime type.
|
||||
_TYPEVAR_TYPES: tuple[type, ...] = (type(typing.TypeVar("_T")), type(typing_extensions.TypeVar("_T"))) # pyright: ignore[reportUnknownVariableType]
|
||||
|
||||
|
||||
def is_typevar(x: Any) -> bool:
|
||||
"""Check if x is an unresolved TypeVar instance (from typing or typing_extensions).
|
||||
|
||||
Args:
|
||||
x: The value to check.
|
||||
|
||||
Returns:
|
||||
True if x is a TypeVar instance, False otherwise.
|
||||
"""
|
||||
return isinstance(x, _TYPEVAR_TYPES)
|
||||
|
||||
|
||||
def contains_typevar(annotation: Any) -> bool:
|
||||
"""Check if an annotation contains an unresolved TypeVar at any nesting level.
|
||||
|
||||
Args:
|
||||
annotation: The annotation to inspect.
|
||||
|
||||
Returns:
|
||||
True if the annotation or any nested type argument is a TypeVar, False otherwise.
|
||||
"""
|
||||
if is_typevar(annotation):
|
||||
return True
|
||||
|
||||
return any(contains_typevar(arg) for arg in get_args(annotation))
|
||||
|
||||
|
||||
def is_chat_agent(agent: Any) -> TypeGuard[Agent]:
|
||||
"""Check if the given agent is a Agent.
|
||||
|
||||
Args:
|
||||
agent (Any): The agent to check.
|
||||
|
||||
Returns:
|
||||
TypeGuard[Agent]: True if the agent is a Agent, False otherwise.
|
||||
"""
|
||||
return isinstance(agent, Agent)
|
||||
|
||||
|
||||
def resolve_type_annotation(
|
||||
type_annotation: type[Any] | UnionType | str | None,
|
||||
globalns: dict[str, Any] | None = None,
|
||||
localns: dict[str, Any] | None = None,
|
||||
) -> type[Any] | UnionType | None:
|
||||
"""Resolve a type annotation, including string forward references.
|
||||
|
||||
Args:
|
||||
type_annotation: A type, union type, string forward reference, or None
|
||||
globalns: Global namespace for resolving forward references (typically func.__globals__)
|
||||
localns: Local namespace for resolving forward references
|
||||
|
||||
Returns:
|
||||
The resolved type annotation. For string annotations, evaluates them in the
|
||||
provided namespace. Returns None if type_annotation is None.
|
||||
|
||||
Raises:
|
||||
NameError: If a forward reference cannot be resolved in the provided namespaces
|
||||
SyntaxError: If a string annotation contains invalid Python syntax
|
||||
|
||||
Note:
|
||||
This function uses eval() to resolve string type annotations. This is the same
|
||||
approach used by Python's typing.get_type_hints() and typing.ForwardRef internally.
|
||||
Security is managed by: (1) strings come from decorator parameters in source code,
|
||||
not runtime user input, and (2) the eval namespace is restricted to the function's
|
||||
module globals plus Union/Optional from typing.
|
||||
|
||||
Examples:
|
||||
- resolve_type_annotation(str) -> str
|
||||
- resolve_type_annotation("str | int", {"str": str, "int": int}) -> str | int
|
||||
- resolve_type_annotation("MyClass", {"MyClass": MyClass}) -> MyClass
|
||||
"""
|
||||
if type_annotation is None:
|
||||
return None
|
||||
|
||||
if isinstance(type_annotation, str):
|
||||
# Resolve string forward reference by evaluating it.
|
||||
# This uses eval() which is the same approach as Python's typing.get_type_hints()
|
||||
# and typing.ForwardRef._evaluate(). The namespace is restricted to the function's
|
||||
# globals plus typing constructs, and input comes from developer source code.
|
||||
eval_globalns = globalns.copy() if globalns else {}
|
||||
eval_globalns.setdefault("Union", Union)
|
||||
eval_globalns.setdefault("Optional", __import__("typing").Optional)
|
||||
|
||||
try:
|
||||
return cast(
|
||||
"type[Any] | UnionType",
|
||||
eval(type_annotation, eval_globalns, localns), # noqa: S307 # nosec B307
|
||||
)
|
||||
except NameError as e:
|
||||
raise NameError(
|
||||
f"Could not resolve type annotation '{type_annotation}'. "
|
||||
f"Make sure the type is defined or imported. Original error: {e}"
|
||||
) from e
|
||||
|
||||
return type_annotation
|
||||
|
||||
|
||||
def normalize_type_to_list(type_annotation: type[Any] | UnionType | None) -> list[type[Any] | UnionType]:
|
||||
"""Normalize a type annotation (possibly a union) to a list of concrete types.
|
||||
|
||||
Args:
|
||||
type_annotation: A type, union type (using | or Union[]), or None
|
||||
|
||||
Returns:
|
||||
A list of types. For union types, returns all members.
|
||||
For None, returns an empty list.
|
||||
For Optional[T] (Union[T, None]), returns [T, type(None)].
|
||||
|
||||
Examples:
|
||||
- normalize_type_to_list(str) -> [str]
|
||||
- normalize_type_to_list(str | int) -> [str, int]
|
||||
- normalize_type_to_list(Union[str, int]) -> [str, int]
|
||||
- normalize_type_to_list(None) -> []
|
||||
"""
|
||||
if type_annotation is None:
|
||||
return []
|
||||
|
||||
origin = get_origin(type_annotation)
|
||||
|
||||
# Handle Union types (str | int or Union[str, int])
|
||||
if origin is Union or origin is UnionType:
|
||||
return list(get_args(type_annotation))
|
||||
|
||||
# Single type
|
||||
return [type_annotation]
|
||||
|
||||
|
||||
def is_instance_of(data: Any, target_type: type | UnionType | Any) -> bool:
|
||||
"""Check if the data is an instance of the target type.
|
||||
|
||||
Args:
|
||||
data (Any): The data to check.
|
||||
target_type (type): The type to check against.
|
||||
|
||||
Returns:
|
||||
bool: True if data is an instance of target_type, False otherwise.
|
||||
"""
|
||||
# Case 0: target_type is Any - always return True
|
||||
if target_type is Any:
|
||||
return True
|
||||
|
||||
origin = get_origin(target_type)
|
||||
args = get_args(target_type)
|
||||
|
||||
# Case 1: origin is None, meaning target_type is not a generic type
|
||||
if origin is None:
|
||||
return isinstance(data, target_type)
|
||||
|
||||
# Case 2: target_type is Optional[T] or Union[T1, T2, ...]
|
||||
# Optional[T] is really just as Union[T, None]
|
||||
if origin is UnionType:
|
||||
return any(is_instance_of(data, arg) for arg in args)
|
||||
|
||||
# Case 2b: Handle typing.Union (legacy Union syntax)
|
||||
if origin is Union:
|
||||
return any(is_instance_of(data, arg) for arg in args)
|
||||
|
||||
# Case 3: target_type is a generic type
|
||||
if origin in [list, set]:
|
||||
return isinstance(data, origin) and (
|
||||
not args or all(any(is_instance_of(item, arg) for arg in args) for item in data) # type: ignore[misc]
|
||||
)
|
||||
|
||||
# Case 4: target_type is a tuple
|
||||
if origin is tuple:
|
||||
if len(args) == 2 and args[1] is Ellipsis: # Tuple[T, ...] case
|
||||
element_type = args[0]
|
||||
return isinstance(data, tuple) and all(is_instance_of(item, element_type) for item in data) # type: ignore[misc]
|
||||
if len(args) == 1 and args[0] is Ellipsis: # Tuple[...] case
|
||||
return isinstance(data, tuple)
|
||||
if len(args) == 0:
|
||||
return isinstance(data, tuple)
|
||||
return (
|
||||
isinstance(data, tuple)
|
||||
and len(data) == len(args) # type: ignore
|
||||
and all(is_instance_of(item, arg) for item, arg in zip(data, args, strict=False)) # type: ignore
|
||||
)
|
||||
|
||||
# Case 5: target_type is a dict
|
||||
if origin is dict:
|
||||
return isinstance(data, dict) and (
|
||||
not args
|
||||
or all(
|
||||
is_instance_of(key, args[0]) and is_instance_of(value, args[1])
|
||||
for key, value in data.items() # type: ignore
|
||||
)
|
||||
)
|
||||
|
||||
# Case 6: Other custom generic classes - check origin type only
|
||||
# For generic classes, we check if data is an instance of the origin type
|
||||
# We don't validate the generic parameters at runtime since that's handled by type system
|
||||
if origin and hasattr(origin, "__name__"):
|
||||
return isinstance(data, origin)
|
||||
|
||||
# Fallback: if we reach here, we assume data is an instance of the target_type
|
||||
return isinstance(data, target_type)
|
||||
|
||||
|
||||
def try_coerce_to_type(data: Any, target_type: type | UnionType | Any) -> Any:
|
||||
"""Try to coerce data to the target type.
|
||||
|
||||
Attempts lightweight type coercion for common cases where raw data
|
||||
(e.g., from JSON deserialization) needs to be converted to the expected type.
|
||||
|
||||
Returns the coerced value if successful, or the original value if coercion
|
||||
is not needed or not possible.
|
||||
|
||||
Args:
|
||||
data: The data to coerce.
|
||||
target_type: The type to coerce to.
|
||||
|
||||
Returns:
|
||||
The coerced value, or the original value if coercion fails.
|
||||
"""
|
||||
original_data = data
|
||||
|
||||
# If already the right type, return as-is
|
||||
if is_instance_of(data, target_type):
|
||||
return data
|
||||
|
||||
# Can't coerce to non-concrete targets (Union, generic, etc.)
|
||||
if not isinstance(target_type, type):
|
||||
return original_data
|
||||
|
||||
target_cls: type[Any] = target_type
|
||||
|
||||
# int -> float (JSON integers for float fields)
|
||||
if isinstance(data, int) and target_cls is float:
|
||||
return float(data)
|
||||
|
||||
# dict -> dataclass or pydantic model
|
||||
if isinstance(data, dict):
|
||||
from dataclasses import is_dataclass
|
||||
|
||||
if is_dataclass(target_cls):
|
||||
try:
|
||||
return target_cls(**data)
|
||||
except (TypeError, ValueError):
|
||||
return original_data
|
||||
|
||||
model_validate = getattr(target_cls, "model_validate", None)
|
||||
if callable(model_validate):
|
||||
try:
|
||||
return model_validate(data)
|
||||
except Exception:
|
||||
return original_data
|
||||
|
||||
return original_data
|
||||
|
||||
|
||||
def serialize_type(t: type) -> str:
|
||||
"""Serialize a type to a string.
|
||||
|
||||
For example,
|
||||
|
||||
serialize_type(int) => "builtins.int"
|
||||
"""
|
||||
return f"{t.__module__}.{t.__qualname__}"
|
||||
|
||||
|
||||
def deserialize_type(serialized_type_string: str) -> type:
|
||||
"""Deserialize a serialized type string.
|
||||
|
||||
For example,
|
||||
|
||||
deserialize_type("builtins.int") => int
|
||||
"""
|
||||
import importlib
|
||||
|
||||
module_name, _, type_name = serialized_type_string.rpartition(".")
|
||||
module = importlib.import_module(module_name)
|
||||
|
||||
return cast(type, getattr(module, type_name))
|
||||
|
||||
|
||||
def is_type_compatible(source_type: type | UnionType | Any, target_type: type | UnionType | Any) -> bool:
|
||||
"""Check if source_type is compatible with target_type.
|
||||
|
||||
A type is compatible if values of source_type can be assigned to variables of target_type.
|
||||
For example:
|
||||
- list[Message] is compatible with list[str | Message]
|
||||
- str is compatible with str | int
|
||||
- int is compatible with Any
|
||||
|
||||
Args:
|
||||
source_type: The type being assigned from
|
||||
target_type: The type being assigned to
|
||||
|
||||
Returns:
|
||||
bool: True if source_type is compatible with target_type, False otherwise
|
||||
"""
|
||||
# Case 0: target_type is Any - always compatible
|
||||
if target_type is Any:
|
||||
return True
|
||||
|
||||
# Case 1: exact type match
|
||||
if source_type == target_type:
|
||||
return True
|
||||
|
||||
source_origin = get_origin(source_type)
|
||||
source_args = get_args(source_type)
|
||||
target_origin = get_origin(target_type)
|
||||
target_args = get_args(target_type)
|
||||
|
||||
# Case 2: target is Union/Optional - source is compatible if it matches any target member
|
||||
if target_origin is Union or target_origin is UnionType:
|
||||
# Special case: if source is also a Union, check that each source member
|
||||
# is compatible with at least one target member
|
||||
if source_origin is Union or source_origin is UnionType:
|
||||
return all(
|
||||
any(is_type_compatible(source_arg, target_arg) for target_arg in target_args)
|
||||
for source_arg in source_args
|
||||
)
|
||||
# If source is not a Union, check if it's compatible with any target member
|
||||
return any(is_type_compatible(source_type, arg) for arg in target_args)
|
||||
|
||||
# Case 3: source is Union (and target is not Union) - each source member must be compatible with target
|
||||
if source_origin is Union or source_origin is UnionType:
|
||||
return all(is_type_compatible(arg, target_type) for arg in source_args)
|
||||
|
||||
# Case 4: both are non-generic types
|
||||
if source_origin is None and target_origin is None:
|
||||
# Only call issubclass if both are actual types, not UnionType or Any
|
||||
if isinstance(source_type, type) and isinstance(target_type, type):
|
||||
try:
|
||||
return issubclass(source_type, target_type)
|
||||
except TypeError:
|
||||
# Handle cases where issubclass doesn't work (e.g., with special forms)
|
||||
return False
|
||||
return source_type == target_type
|
||||
|
||||
# Case 5: different container types are not compatible
|
||||
if source_origin != target_origin:
|
||||
return False
|
||||
|
||||
# Case 6: same container type - check generic arguments
|
||||
if source_origin in [list, set]:
|
||||
if not source_args and not target_args:
|
||||
return True # Both are untyped
|
||||
if not source_args or not target_args:
|
||||
return True # One is untyped - assume compatible
|
||||
# For collections, source element type must be compatible with target element type
|
||||
return is_type_compatible(source_args[0], target_args[0])
|
||||
|
||||
# Case 7: tuple compatibility
|
||||
if source_origin is tuple:
|
||||
if not source_args and not target_args:
|
||||
return True # Both are untyped tuples
|
||||
if not source_args or not target_args:
|
||||
return True # One is untyped - assume compatible
|
||||
|
||||
# Handle Tuple[T, ...] (variable length)
|
||||
if len(source_args) == 2 and source_args[1] is Ellipsis:
|
||||
if len(target_args) == 2 and target_args[1] is Ellipsis:
|
||||
return is_type_compatible(source_args[0], target_args[0])
|
||||
return False # Variable length can't be assigned to fixed length
|
||||
|
||||
if len(target_args) == 2 and target_args[1] is Ellipsis:
|
||||
# Fixed length can be assigned to variable length if element types are compatible
|
||||
return all(is_type_compatible(source_arg, target_args[0]) for source_arg in source_args)
|
||||
|
||||
# Fixed length tuples must have same length and compatible element types
|
||||
if len(source_args) != len(target_args):
|
||||
return False
|
||||
return all(is_type_compatible(s_arg, t_arg) for s_arg, t_arg in zip(source_args, target_args, strict=False))
|
||||
|
||||
# Case 8: dict compatibility
|
||||
if source_origin is dict:
|
||||
if not source_args and not target_args:
|
||||
return True # Both are untyped dicts
|
||||
if not source_args or not target_args:
|
||||
return True # One is untyped - assume compatible
|
||||
if len(source_args) != 2 or len(target_args) != 2:
|
||||
return False # Malformed dict types
|
||||
# Both key and value types must be compatible
|
||||
return is_type_compatible(source_args[0], target_args[0]) and is_type_compatible(source_args[1], target_args[1])
|
||||
|
||||
# Case 9: custom generic classes - check if origins are the same and args are compatible
|
||||
if source_origin and target_origin and source_origin == target_origin:
|
||||
if not source_args and not target_args:
|
||||
return True # Both are untyped generics
|
||||
if not source_args or not target_args:
|
||||
return True # One is untyped - assume compatible
|
||||
if len(source_args) != len(target_args):
|
||||
return False # Different number of type parameters
|
||||
return all(is_type_compatible(s_arg, t_arg) for s_arg, t_arg in zip(source_args, target_args, strict=False))
|
||||
|
||||
# Case 10: fallback - check if source is subclass of target (for non-generic types)
|
||||
if source_origin is None and target_origin is None:
|
||||
try:
|
||||
# Only call issubclass if both are actual types, not UnionType or Any
|
||||
if isinstance(source_type, type) and isinstance(target_type, type):
|
||||
return issubclass(source_type, target_type)
|
||||
return source_type == target_type
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,462 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
import types
|
||||
from collections import defaultdict
|
||||
from collections.abc import Sequence
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from ..exceptions import WorkflowException
|
||||
from ._edge import Edge, EdgeGroup, FanInEdgeGroup, InternalEdgeGroup
|
||||
from ._executor import Executor
|
||||
from ._typing_utils import is_type_compatible
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# region Enums and Base Classes
|
||||
class ValidationTypeEnum(Enum):
|
||||
"""Enumeration of workflow validation types."""
|
||||
|
||||
EDGE_DUPLICATION = "EDGE_DUPLICATION"
|
||||
EXECUTOR_DUPLICATION = "EXECUTOR_DUPLICATION"
|
||||
TYPE_COMPATIBILITY = "TYPE_COMPATIBILITY"
|
||||
GRAPH_CONNECTIVITY = "GRAPH_CONNECTIVITY"
|
||||
HANDLER_OUTPUT_ANNOTATION = "HANDLER_OUTPUT_ANNOTATION"
|
||||
OUTPUT_VALIDATION = "OUTPUT_VALIDATION"
|
||||
|
||||
|
||||
class WorkflowValidationError(WorkflowException):
|
||||
"""Base exception for workflow validation errors."""
|
||||
|
||||
def __init__(self, message: str, validation_type: ValidationTypeEnum):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.validation_type = validation_type
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"[{self.validation_type.value}] {self.message}"
|
||||
|
||||
|
||||
class EdgeDuplicationError(WorkflowValidationError):
|
||||
"""Exception raised when duplicate edges are detected in the workflow."""
|
||||
|
||||
def __init__(self, edge_id: str):
|
||||
super().__init__(
|
||||
message=f"Duplicate edge detected: {edge_id}. Each edge in the workflow must be unique.",
|
||||
validation_type=ValidationTypeEnum.EDGE_DUPLICATION,
|
||||
)
|
||||
self.edge_id = edge_id
|
||||
|
||||
|
||||
class TypeCompatibilityError(WorkflowValidationError):
|
||||
"""Exception raised when type incompatibility is detected between connected executors."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_executor_id: str,
|
||||
target_executor_id: str,
|
||||
source_types: list[type[Any] | types.UnionType],
|
||||
target_types: list[type[Any] | types.UnionType],
|
||||
):
|
||||
# Use a placeholder for incompatible types - will be computed in WorkflowGraphValidator
|
||||
super().__init__(
|
||||
message=f"Type incompatibility between executors '{source_executor_id}' -> '{target_executor_id}'. "
|
||||
f"Source executor outputs types {[str(t) for t in source_types]} but target executor "
|
||||
f"can only handle types {[str(t) for t in target_types]}.",
|
||||
validation_type=ValidationTypeEnum.TYPE_COMPATIBILITY,
|
||||
)
|
||||
self.source_executor_id = source_executor_id
|
||||
self.target_executor_id = target_executor_id
|
||||
self.source_types = source_types
|
||||
self.target_types = target_types
|
||||
|
||||
|
||||
class GraphConnectivityError(WorkflowValidationError):
|
||||
"""Exception raised when graph connectivity issues are detected."""
|
||||
|
||||
def __init__(self, message: str):
|
||||
super().__init__(message, validation_type=ValidationTypeEnum.GRAPH_CONNECTIVITY)
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region Workflow Graph Validator
|
||||
class WorkflowGraphValidator:
|
||||
"""Validator for workflow graphs.
|
||||
|
||||
This validator performs multiple validation checks:
|
||||
1. Edge duplication validation
|
||||
2. Type compatibility validation between connected executors
|
||||
3. Graph connectivity validation
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._edges: list[Edge] = []
|
||||
self._executors: dict[str, Executor] = {}
|
||||
|
||||
# region Core Validation Methods
|
||||
def validate_workflow(
|
||||
self,
|
||||
edge_groups: Sequence[EdgeGroup],
|
||||
executors: dict[str, Executor],
|
||||
start_executor: Executor,
|
||||
output_executors: list[str],
|
||||
intermediate_executors: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Validate the entire workflow graph.
|
||||
|
||||
Args:
|
||||
edge_groups: list of edge groups in the workflow
|
||||
executors: Map of executor IDs to executor instances
|
||||
start_executor: The starting executor
|
||||
output_executors: List of output executor IDs
|
||||
intermediate_executors: List of intermediate executor IDs
|
||||
|
||||
Raises:
|
||||
WorkflowValidationError: If any validation fails
|
||||
"""
|
||||
self._executors = executors
|
||||
self._edges = [edge for group in edge_groups for edge in group.edges]
|
||||
self._edge_groups = edge_groups
|
||||
|
||||
# If only the start executor exists, add it to the executor map
|
||||
# Handle the special case where the workflow consists of only a single executor and no edges.
|
||||
# In this scenario, the executor map will be empty because there are no edge groups to reference executors.
|
||||
# Adding the start executor to the map ensures that single-executor workflows (without any edges) are supported,
|
||||
# allowing validation and execution to proceed for workflows that do not require inter-executor communication.
|
||||
if not self._executors:
|
||||
self._executors[start_executor.id] = start_executor
|
||||
|
||||
# Validate that start_executor exists in the graph
|
||||
# It should because we check for it in the WorkflowBuilder
|
||||
# but we do it here for completeness.
|
||||
if start_executor.id not in self._executors:
|
||||
raise GraphConnectivityError(f"Start executor '{start_executor.id}' is not present in the workflow graph")
|
||||
|
||||
# Additional presence verification:
|
||||
# A start executor that is only injected via the builder (present in the executors map)
|
||||
# but not referenced by any edge while other executors ARE referenced indicates a
|
||||
# configuration error: the chosen start node is effectively disconnected / unknown to the
|
||||
# defined graph topology. For single-node workflows (no edges) we allow the start executor
|
||||
# to stand alone (handled above when we inject it into the map). We perform this refined
|
||||
# check only when there is at least one edge group defined.
|
||||
if self._edges: # Only evaluate when the workflow defines edges
|
||||
edge_executor_ids: set[str] = set()
|
||||
for e in self._edges:
|
||||
edge_executor_ids.add(e.source_id)
|
||||
edge_executor_ids.add(e.target_id)
|
||||
if start_executor.id not in edge_executor_ids:
|
||||
raise GraphConnectivityError(
|
||||
f"Start executor '{start_executor.id}' is not present in the workflow graph"
|
||||
)
|
||||
|
||||
# Run all checks
|
||||
self._validate_edge_duplication()
|
||||
self._validate_handler_output_annotations()
|
||||
self._validate_type_compatibility()
|
||||
self._validate_graph_connectivity(start_executor.id)
|
||||
self._validate_self_loops()
|
||||
self._validate_dead_ends()
|
||||
self._output_validation(output_executors, intermediate_executors or [])
|
||||
|
||||
def _validate_handler_output_annotations(self) -> None:
|
||||
"""Validate that each handler's ctx parameter is annotated with WorkflowContext[T].
|
||||
|
||||
Note: This validation is now primarily handled at handler registration time
|
||||
via the unified validation functions in _workflow_context.py when the @handler
|
||||
decorator is applied. This method is kept minimal for any edge cases.
|
||||
"""
|
||||
# The comprehensive validation is already done during handler registration:
|
||||
# 1. @handler and @response_handler decorators already have validation logic
|
||||
# 2. FunctionExecutor constructor also has validation logic
|
||||
# 3. Both use validate_workflow_context_annotation() for WorkflowContext validation
|
||||
#
|
||||
# All executors in the workflow must have gone through one of these paths,
|
||||
# so redundant validation here is unnecessary and has been removed.
|
||||
pass
|
||||
|
||||
# endregion
|
||||
|
||||
# region Edge and Type Validation
|
||||
def _validate_edge_duplication(self) -> None:
|
||||
"""Validate that there are no duplicate edges in the workflow.
|
||||
|
||||
Raises:
|
||||
EdgeDuplicationError: If duplicate edges are found
|
||||
"""
|
||||
seen_edge_ids: set[str] = set()
|
||||
|
||||
for edge in self._edges:
|
||||
edge_id = edge.id
|
||||
if edge_id in seen_edge_ids:
|
||||
raise EdgeDuplicationError(edge_id)
|
||||
seen_edge_ids.add(edge_id)
|
||||
|
||||
def _validate_type_compatibility(self) -> None:
|
||||
"""Validate type compatibility between connected executors.
|
||||
|
||||
This checks that the output types of source executors are compatible
|
||||
with the input types expected by target executors.
|
||||
|
||||
Raises:
|
||||
TypeCompatibilityError: If type incompatibility is detected
|
||||
"""
|
||||
for edge_group in self._edge_groups:
|
||||
for edge in edge_group.edges:
|
||||
self._validate_edge_type_compatibility(edge, edge_group)
|
||||
|
||||
def _validate_edge_type_compatibility(self, edge: Edge, edge_group: EdgeGroup) -> None:
|
||||
"""Validate type compatibility for a specific edge.
|
||||
|
||||
This checks that the output types of the source executor are compatible
|
||||
with the input types expected by the target executor.
|
||||
|
||||
Args:
|
||||
edge: The edge to validate
|
||||
edge_group: The edge group containing this edge
|
||||
|
||||
Raises:
|
||||
TypeCompatibilityError: If type incompatibility is detected
|
||||
"""
|
||||
if isinstance(edge_group, InternalEdgeGroup):
|
||||
# Skip type compatibility validation for internal edges
|
||||
return
|
||||
|
||||
source_executor = self._executors[edge.source_id]
|
||||
target_executor = self._executors[edge.target_id]
|
||||
|
||||
# Get output types from source executor
|
||||
source_output_types = list(source_executor.output_types)
|
||||
|
||||
# Get input types from target executor
|
||||
target_input_types = target_executor.input_types
|
||||
|
||||
# If either executor has no type information, log warning and skip validation
|
||||
# This allows for dynamic typing scenarios but warns about reduced validation coverage
|
||||
if not source_output_types or not target_input_types:
|
||||
if not source_output_types:
|
||||
logger.warning(
|
||||
f"Executor '{source_executor.id}' has no output type annotations. "
|
||||
f"Type compatibility validation will be skipped for edges from this executor. "
|
||||
f"Consider adding WorkflowContext[T] generics in handlers for better validation."
|
||||
)
|
||||
if not target_input_types:
|
||||
logger.warning(
|
||||
f"Executor '{target_executor.id}' has no input type annotations. "
|
||||
f"Type compatibility validation will be skipped for edges to this executor. "
|
||||
f"Consider adding type annotations to message handler parameters for better validation."
|
||||
)
|
||||
return
|
||||
|
||||
# Check if any source output type is compatible with any target input type
|
||||
compatible = False
|
||||
compatible_pairs: list[tuple[type[Any] | types.UnionType, type[Any] | types.UnionType]] = []
|
||||
|
||||
for source_type in source_output_types:
|
||||
for target_type in target_input_types:
|
||||
if isinstance(edge_group, FanInEdgeGroup):
|
||||
# If the edge is part of an edge group, the target expects a list of data types
|
||||
if is_type_compatible(list[source_type], target_type):
|
||||
compatible = True
|
||||
compatible_pairs.append((list[source_type], target_type))
|
||||
else:
|
||||
if is_type_compatible(source_type, target_type):
|
||||
compatible = True
|
||||
compatible_pairs.append((source_type, target_type))
|
||||
|
||||
# Log successful type compatibility for debugging
|
||||
if compatible:
|
||||
logger.debug(
|
||||
f"Type compatibility validated for edge '{source_executor.id}' -> '{target_executor.id}'. "
|
||||
f"Compatible type pairs: {[(str(s), str(t)) for s, t in compatible_pairs]}"
|
||||
)
|
||||
|
||||
if not compatible:
|
||||
# Enhanced error with more detailed information
|
||||
raise TypeCompatibilityError(
|
||||
source_executor.id,
|
||||
target_executor.id,
|
||||
source_output_types,
|
||||
target_input_types,
|
||||
)
|
||||
|
||||
# endregion
|
||||
|
||||
# region Graph Connectivity Validation
|
||||
def _validate_graph_connectivity(self, start_executor_id: str) -> None:
|
||||
"""Validate graph connectivity and detect potential issues.
|
||||
|
||||
This performs several checks:
|
||||
- Detects unreachable executors from the start node
|
||||
- Detects isolated executors (no incoming or outgoing edges)
|
||||
- Warns about potential infinite loops
|
||||
|
||||
Args:
|
||||
start_executor_id: The ID of the starting executor
|
||||
|
||||
Raises:
|
||||
GraphConnectivityError: If connectivity issues are detected
|
||||
"""
|
||||
# Build adjacency list for the graph
|
||||
graph: dict[str, list[str]] = defaultdict(list)
|
||||
all_executors = set(self._executors.keys())
|
||||
|
||||
for edge in self._edges:
|
||||
graph[edge.source_id].append(edge.target_id)
|
||||
|
||||
# Find reachable nodes from start
|
||||
reachable = self._find_reachable_nodes(graph, start_executor_id)
|
||||
|
||||
# Check for unreachable executors
|
||||
unreachable = all_executors - reachable
|
||||
if unreachable:
|
||||
raise GraphConnectivityError(
|
||||
f"The following executors are unreachable from the start executor '{start_executor_id}': "
|
||||
f"{sorted(unreachable)}. This may indicate a disconnected workflow graph."
|
||||
)
|
||||
|
||||
# Check for isolated executors (no edges)
|
||||
isolated_executors: list[str] = []
|
||||
for executor_id in all_executors:
|
||||
has_incoming = any(edge.target_id == executor_id for edge in self._edges)
|
||||
has_outgoing = any(edge.source_id == executor_id for edge in self._edges)
|
||||
|
||||
if not has_incoming and not has_outgoing and executor_id != start_executor_id:
|
||||
isolated_executors.append(executor_id)
|
||||
|
||||
if isolated_executors:
|
||||
raise GraphConnectivityError(
|
||||
f"The following executors are isolated (no incoming or outgoing edges): "
|
||||
f"{sorted(isolated_executors)}. Isolated executors will never be executed."
|
||||
)
|
||||
|
||||
def _find_reachable_nodes(self, graph: dict[str, list[str]], start: str) -> set[str]:
|
||||
"""Find all nodes reachable from the start node using DFS.
|
||||
|
||||
Args:
|
||||
graph: Adjacency list representation of the graph
|
||||
start: Starting node ID
|
||||
|
||||
Returns:
|
||||
Set of reachable node IDs
|
||||
"""
|
||||
visited: set[str] = set()
|
||||
stack = [start]
|
||||
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
if node not in visited:
|
||||
visited.add(node)
|
||||
stack.extend(graph[node])
|
||||
|
||||
return visited
|
||||
|
||||
# endregion
|
||||
|
||||
# region Output Validation
|
||||
|
||||
def _output_validation(self, output_executors: list[str], intermediate_executors: list[str]) -> None:
|
||||
"""Validate that designated executors exist and have workflow output annotations."""
|
||||
overlap = sorted(set(output_executors).intersection(intermediate_executors))
|
||||
if overlap:
|
||||
raise WorkflowValidationError(
|
||||
f"Executors cannot be both output and intermediate designated: {overlap}",
|
||||
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
|
||||
)
|
||||
|
||||
for output_id in output_executors:
|
||||
if output_id not in self._executors:
|
||||
raise WorkflowValidationError(
|
||||
f"Output executor '{output_id}' is not present in the workflow graph",
|
||||
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
|
||||
)
|
||||
|
||||
output_executor = self._executors[output_id]
|
||||
if not output_executor.workflow_output_types:
|
||||
raise WorkflowValidationError(
|
||||
f"Output executor '{output_id}' must have output type annotations defined.",
|
||||
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
|
||||
)
|
||||
|
||||
for intermediate_id in intermediate_executors:
|
||||
if intermediate_id not in self._executors:
|
||||
raise WorkflowValidationError(
|
||||
f"Intermediate executor '{intermediate_id}' is not present in the workflow graph",
|
||||
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
|
||||
)
|
||||
|
||||
intermediate_executor = self._executors[intermediate_id]
|
||||
if not intermediate_executor.workflow_output_types:
|
||||
raise WorkflowValidationError(
|
||||
f"Intermediate executor '{intermediate_id}' must have output type annotations defined.",
|
||||
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
|
||||
)
|
||||
|
||||
# endregion
|
||||
|
||||
# region Additional Validation Scenarios
|
||||
def _validate_self_loops(self) -> None:
|
||||
"""Detect and log self-loops (edges from executor to itself).
|
||||
|
||||
Self-loops might indicate recursive processing which could be intentional
|
||||
but should be highlighted for review.
|
||||
"""
|
||||
self_loops = [edge for edge in self._edges if edge.source_id == edge.target_id]
|
||||
|
||||
for edge in self_loops:
|
||||
logger.warning(
|
||||
f"Self-loop detected: Executor '{edge.source_id}' connects to itself. "
|
||||
f"This may cause infinite recursion if not properly handled with conditions."
|
||||
)
|
||||
|
||||
def _validate_dead_ends(self) -> None:
|
||||
"""Identify executors that have no outgoing edges (potential dead ends).
|
||||
|
||||
These might be intentional final nodes or could indicate missing connections.
|
||||
"""
|
||||
executors_with_outgoing = {edge.source_id for edge in self._edges}
|
||||
all_executor_ids = set(self._executors.keys())
|
||||
dead_ends = all_executor_ids - executors_with_outgoing
|
||||
|
||||
if dead_ends:
|
||||
logger.info(
|
||||
f"Dead-end executors detected (no outgoing edges): {sorted(dead_ends)}. "
|
||||
f"Verify these are intended as final nodes in the workflow."
|
||||
)
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
def validate_workflow_graph(
|
||||
edge_groups: Sequence[EdgeGroup],
|
||||
executors: dict[str, Executor],
|
||||
start_executor: Executor,
|
||||
output_executors: list[str],
|
||||
intermediate_executors: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Convenience function to validate a workflow graph.
|
||||
|
||||
Args:
|
||||
edge_groups: list of edge groups in the workflow
|
||||
executors: Map of executor IDs to executor instances
|
||||
start_executor: The starting executor instance
|
||||
output_executors: List of output executor IDs
|
||||
intermediate_executors: List of intermediate executor IDs
|
||||
|
||||
Raises:
|
||||
WorkflowValidationError: If any validation fails
|
||||
"""
|
||||
validator = WorkflowGraphValidator()
|
||||
validator.validate_workflow(
|
||||
edge_groups,
|
||||
executors,
|
||||
start_executor,
|
||||
output_executors,
|
||||
intermediate_executors,
|
||||
)
|
||||
@@ -0,0 +1,442 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import hashlib
|
||||
import re
|
||||
import tempfile
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from ._edge import FanInEdgeGroup, InternalEdgeGroup
|
||||
from ._workflow import Workflow
|
||||
|
||||
# Import of WorkflowExecutor is performed lazily inside methods to avoid cycles
|
||||
|
||||
"""Workflow visualization module using graphviz and Mermaid."""
|
||||
|
||||
|
||||
class WorkflowViz:
|
||||
"""A class for visualizing workflows using graphviz and Mermaid."""
|
||||
|
||||
def __init__(self, workflow: Workflow):
|
||||
"""Initialize the WorkflowViz with a workflow.
|
||||
|
||||
Args:
|
||||
workflow: The workflow to visualize.
|
||||
"""
|
||||
self._workflow = workflow
|
||||
|
||||
def to_digraph(self, include_internal_executors: bool = False) -> str:
|
||||
"""Export the workflow as a DOT format digraph string.
|
||||
|
||||
Args:
|
||||
include_internal_executors (bool): Whether to include internal executors in the visualization.
|
||||
Default is False.
|
||||
|
||||
Returns:
|
||||
A string representation of the workflow in DOT format.
|
||||
"""
|
||||
lines = ["digraph Workflow {"]
|
||||
lines.append(" rankdir=TD;") # Top to bottom layout
|
||||
lines.append(" node [shape=box, style=filled, fillcolor=lightblue];")
|
||||
lines.append(" edge [color=black, arrowhead=vee];")
|
||||
lines.append("")
|
||||
|
||||
# Emit the top-level workflow nodes/edges
|
||||
self._emit_workflow_digraph(
|
||||
self._workflow,
|
||||
lines,
|
||||
indent=" ",
|
||||
include_internal_executors=include_internal_executors,
|
||||
)
|
||||
|
||||
# Emit sub-workflows hosted by WorkflowExecutor as nested clusters
|
||||
self._emit_sub_workflows_digraph(
|
||||
self._workflow,
|
||||
lines,
|
||||
indent=" ",
|
||||
include_internal_executors=include_internal_executors,
|
||||
)
|
||||
|
||||
lines.append("}")
|
||||
return "\n".join(lines)
|
||||
|
||||
def export(
|
||||
self,
|
||||
format: Literal["svg", "png", "pdf", "dot"] = "svg",
|
||||
filename: str | None = None,
|
||||
include_internal_executors: bool = False,
|
||||
) -> str:
|
||||
"""Export the workflow visualization to a file or return the file path.
|
||||
|
||||
Args:
|
||||
format: The output format. Supported formats: 'svg', 'png', 'pdf', 'dot'.
|
||||
filename: Optional filename to save the output. If None, creates a temporary file.
|
||||
include_internal_executors (bool): Whether to include internal executors in the visualization.
|
||||
Default is False.
|
||||
|
||||
Returns:
|
||||
The path to the saved file.
|
||||
|
||||
Raises:
|
||||
ImportError: If graphviz is not installed.
|
||||
ValueError: If an unsupported format is specified.
|
||||
"""
|
||||
# Validate format first
|
||||
if format not in ["svg", "png", "pdf", "dot"]:
|
||||
raise ValueError(f"Unsupported format: {format}. Supported formats: svg, png, pdf, dot")
|
||||
|
||||
if format == "dot":
|
||||
content = self.to_digraph(include_internal_executors=include_internal_executors)
|
||||
if filename:
|
||||
with open(filename, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
return filename
|
||||
# Create temporary file for dot format
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".dot", delete=False, encoding="utf-8") as temp_file:
|
||||
temp_file.write(content)
|
||||
return temp_file.name
|
||||
|
||||
try:
|
||||
import graphviz
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"viz extra is required for export. Install it with: pip install graphviz>=0.20.0 "
|
||||
"The version needs to be at least 0.20.0. "
|
||||
"You also need to install graphviz separately. E.g., sudo apt-get install graphviz on Debian/Ubuntu "
|
||||
"or brew install graphviz on macOS. See https://graphviz.org/download/ for details."
|
||||
) from e
|
||||
|
||||
# Create a temporary graphviz Source object
|
||||
dot_content = self.to_digraph(include_internal_executors=include_internal_executors)
|
||||
source = graphviz.Source(dot_content)
|
||||
|
||||
try:
|
||||
if filename:
|
||||
# Save to specified file
|
||||
output_path = Path(filename)
|
||||
if output_path.suffix and output_path.suffix[1:] != format:
|
||||
raise ValueError(f"File extension {output_path.suffix} doesn't match format {format}")
|
||||
|
||||
# Remove extension if present since graphviz.render() adds it
|
||||
base_name = str(output_path.with_suffix(""))
|
||||
source.render(base_name, format=format, cleanup=True) # type: ignore
|
||||
|
||||
# Return the actual filename with extension
|
||||
return f"{base_name}.{format}"
|
||||
# Create temporary file
|
||||
with tempfile.NamedTemporaryFile(suffix=f".{format}", delete=False) as temp_file:
|
||||
temp_path = Path(temp_file.name)
|
||||
base_name = str(temp_path.with_suffix(""))
|
||||
|
||||
source.render(base_name, format=format, cleanup=True) # type: ignore
|
||||
return f"{base_name}.{format}"
|
||||
except graphviz.backend.execute.ExecutableNotFound as e:
|
||||
raise ImportError(
|
||||
"The graphviz executables are not found. The graphviz Python package is installed, but the "
|
||||
"graphviz executables (dot, neato, etc.) are not available on your system's PATH. "
|
||||
"Install graphviz executables: sudo apt-get install graphviz on Debian/Ubuntu, "
|
||||
"brew install graphviz on macOS, or download from https://graphviz.org/download/ for other platforms."
|
||||
) from e
|
||||
|
||||
def save_svg(self, filename: str, include_internal_executors: bool = False) -> str:
|
||||
"""Convenience method to save as SVG.
|
||||
|
||||
Args:
|
||||
filename: The filename to save the SVG file.
|
||||
include_internal_executors (bool): Whether to include internal executors in the visualization.
|
||||
Default is False.
|
||||
|
||||
Returns:
|
||||
The path to the saved SVG file.
|
||||
"""
|
||||
return self.export(format="svg", filename=filename, include_internal_executors=include_internal_executors)
|
||||
|
||||
def save_png(self, filename: str, include_internal_executors: bool = False) -> str:
|
||||
"""Convenience method to save as PNG.
|
||||
|
||||
Args:
|
||||
filename: The filename to save the PNG file.
|
||||
include_internal_executors (bool): Whether to include internal executors in the visualization.
|
||||
Default is False.
|
||||
|
||||
Returns:
|
||||
The path to the saved PNG file.
|
||||
"""
|
||||
return self.export(format="png", filename=filename, include_internal_executors=include_internal_executors)
|
||||
|
||||
def save_pdf(self, filename: str, include_internal_executors: bool = False) -> str:
|
||||
"""Convenience method to save as PDF.
|
||||
|
||||
Args:
|
||||
filename: The filename to save the PDF file.
|
||||
include_internal_executors (bool): Whether to include internal executors in the visualization.
|
||||
Default is False.
|
||||
|
||||
Returns:
|
||||
The path to the saved PDF file.
|
||||
"""
|
||||
return self.export(format="pdf", filename=filename, include_internal_executors=include_internal_executors)
|
||||
|
||||
def to_mermaid(self, include_internal_executors: bool = False) -> str:
|
||||
"""Export the workflow as a Mermaid flowchart string.
|
||||
|
||||
Args:
|
||||
include_internal_executors (bool): Whether to include internal executors in the visualization.
|
||||
Default is False.
|
||||
|
||||
Returns:
|
||||
A string representation of the workflow in Mermaid flowchart syntax.
|
||||
"""
|
||||
lines: list[str] = ["flowchart TD"]
|
||||
|
||||
# Emit top-level workflow
|
||||
self._emit_workflow_mermaid(
|
||||
self._workflow,
|
||||
lines,
|
||||
indent=" ",
|
||||
include_internal_executors=include_internal_executors,
|
||||
)
|
||||
|
||||
# Emit sub-workflows as Mermaid subgraphs
|
||||
self._emit_sub_workflows_mermaid(
|
||||
self._workflow,
|
||||
lines,
|
||||
indent=" ",
|
||||
include_internal_executors=include_internal_executors,
|
||||
)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
# region Private helpers
|
||||
|
||||
def _fan_in_digest(self, target: str, sources: list[str]) -> str:
|
||||
sources_sorted = sorted(sources)
|
||||
return hashlib.sha256((target + "|" + "|".join(sources_sorted)).encode("utf-8")).hexdigest()[:8]
|
||||
|
||||
def _compute_fan_in_descriptors(self, workflow: Workflow | None = None) -> list[tuple[str, list[str], str]]:
|
||||
"""Return list of (node_id, sources, target) for fan-in groups.
|
||||
|
||||
node_id is DOT-oriented: fan_in::target::digest
|
||||
"""
|
||||
result: list[tuple[str, list[str], str]] = []
|
||||
workflow = workflow or self._workflow
|
||||
for group in workflow.edge_groups:
|
||||
if isinstance(group, FanInEdgeGroup):
|
||||
target = group.target_executor_ids[0]
|
||||
sources = list(group.source_executor_ids)
|
||||
digest = self._fan_in_digest(target, sources)
|
||||
node_id = f"fan_in::{target}::{digest}"
|
||||
result.append((node_id, sorted(sources), target))
|
||||
return result
|
||||
|
||||
def _compute_normal_edges(
|
||||
self,
|
||||
workflow: Workflow | None = None,
|
||||
include_internal_executors: bool = False,
|
||||
) -> list[tuple[str, str, bool]]:
|
||||
"""Return list of (source_id, target_id, is_conditional) for non-fan-in groups."""
|
||||
edges: list[tuple[str, str, bool]] = []
|
||||
workflow = workflow or self._workflow
|
||||
for group in workflow.edge_groups:
|
||||
if isinstance(group, FanInEdgeGroup):
|
||||
continue
|
||||
if isinstance(group, InternalEdgeGroup) and not include_internal_executors:
|
||||
continue
|
||||
for edge in group.edges:
|
||||
is_cond = getattr(edge, "_condition", None) is not None
|
||||
edges.append((edge.source_id, edge.target_id, is_cond))
|
||||
return edges
|
||||
|
||||
# endregion
|
||||
|
||||
# region Internal emitters (DOT)
|
||||
|
||||
def _emit_workflow_digraph(
|
||||
self,
|
||||
workflow: Workflow,
|
||||
lines: list[str],
|
||||
indent: str,
|
||||
ns: str | None = None,
|
||||
include_internal_executors: bool = False,
|
||||
) -> None:
|
||||
"""Emit DOT nodes/edges for the given workflow.
|
||||
|
||||
If ns (namespace) is provided, node ids are prefixed with f"{ns}/" for uniqueness,
|
||||
but labels remain the original executor ids.
|
||||
"""
|
||||
|
||||
def map_id(x: str) -> str:
|
||||
return f"{ns}/{x}" if ns else x
|
||||
|
||||
# Nodes
|
||||
start_executor_id = workflow.start_executor_id
|
||||
lines.append(
|
||||
f'{indent}"{map_id(start_executor_id)}" [fillcolor=lightgreen, label="{start_executor_id}\\n(Start)"];'
|
||||
)
|
||||
for executor_id in workflow.executors:
|
||||
if executor_id != start_executor_id:
|
||||
lines.append(f'{indent}"{map_id(executor_id)}" [label="{executor_id}"];')
|
||||
|
||||
# Fan-in nodes
|
||||
fan_in_nodes = self._compute_fan_in_descriptors(workflow)
|
||||
if fan_in_nodes:
|
||||
lines.append("")
|
||||
for node_id, _, _ in fan_in_nodes:
|
||||
lines.append(f'{indent}"{map_id(node_id)}" [shape=ellipse, fillcolor=lightgoldenrod, label="fan-in"];')
|
||||
|
||||
# Fan-in edges
|
||||
for node_id, sources, target in fan_in_nodes:
|
||||
for src in sources:
|
||||
lines.append(f'{indent}"{map_id(src)}" -> "{map_id(node_id)}";')
|
||||
lines.append(f'{indent}"{map_id(node_id)}" -> "{map_id(target)}";')
|
||||
|
||||
# Normal edges
|
||||
for src, tgt, is_cond in self._compute_normal_edges(
|
||||
workflow, include_internal_executors=include_internal_executors
|
||||
):
|
||||
edge_attr = ' [style=dashed, label="conditional"]' if is_cond else ""
|
||||
lines.append(f'{indent}"{map_id(src)}" -> "{map_id(tgt)}"{edge_attr};')
|
||||
|
||||
def _emit_sub_workflows_digraph(
|
||||
self,
|
||||
workflow: Workflow,
|
||||
lines: list[str],
|
||||
indent: str,
|
||||
include_internal_executors: bool = False,
|
||||
) -> None:
|
||||
"""Emit DOT subgraphs for any WorkflowExecutor instances found in the workflow."""
|
||||
# Lazy import to avoid any potential import cycles
|
||||
try:
|
||||
from ._workflow_executor import WorkflowExecutor
|
||||
except ImportError: # pragma: no cover - best-effort; if unavailable, skip subgraphs
|
||||
return
|
||||
|
||||
for exec_id, exec_obj in workflow.executors.items():
|
||||
if isinstance(exec_obj, WorkflowExecutor) and hasattr(exec_obj, "workflow") and exec_obj.workflow:
|
||||
subgraph_id = f"cluster_{uuid.uuid5(uuid.NAMESPACE_OID, exec_id).hex[:8]}"
|
||||
lines.append(f"{indent}subgraph {subgraph_id} {{")
|
||||
lines.append(f'{indent} label="sub-workflow: {exec_id}";')
|
||||
lines.append(f"{indent} style=dashed;")
|
||||
|
||||
# Emit the nested workflow inside this cluster using a namespace
|
||||
ns = exec_id
|
||||
self._emit_workflow_digraph(
|
||||
exec_obj.workflow,
|
||||
lines,
|
||||
indent=f"{indent} ",
|
||||
ns=ns,
|
||||
include_internal_executors=include_internal_executors,
|
||||
)
|
||||
|
||||
# Recurse into deeper nested sub-workflows
|
||||
self._emit_sub_workflows_digraph(
|
||||
exec_obj.workflow,
|
||||
lines,
|
||||
indent=f"{indent} ",
|
||||
include_internal_executors=include_internal_executors,
|
||||
)
|
||||
|
||||
lines.append(f"{indent}}}")
|
||||
|
||||
# endregion
|
||||
|
||||
# region Internal emitters (Mermaid)
|
||||
|
||||
def _emit_workflow_mermaid(
|
||||
self,
|
||||
workflow: Workflow,
|
||||
lines: list[str],
|
||||
indent: str,
|
||||
ns: str | None = None,
|
||||
include_internal_executors: bool = False,
|
||||
) -> None:
|
||||
def _san(s: str) -> str:
|
||||
s2 = re.sub(r"[^0-9A-Za-z_]", "_", s)
|
||||
if not s2 or not s2[0].isalpha():
|
||||
s2 = f"n_{s2}"
|
||||
return s2
|
||||
|
||||
def map_id(x: str) -> str:
|
||||
if ns:
|
||||
return f"{_san(ns)}__{_san(x)}"
|
||||
return _san(x)
|
||||
|
||||
# Nodes
|
||||
start_executor_id = workflow.start_executor_id
|
||||
lines.append(f'{indent}{map_id(start_executor_id)}["{start_executor_id} (Start)"];')
|
||||
for executor_id in workflow.executors:
|
||||
if executor_id == start_executor_id:
|
||||
continue
|
||||
lines.append(f'{indent}{map_id(executor_id)}["{executor_id}"];')
|
||||
|
||||
# Fan-in nodes
|
||||
fan_in_nodes_dot = self._compute_fan_in_descriptors(workflow)
|
||||
fan_in_nodes: list[tuple[str, list[str], str]] = []
|
||||
for dot_node_id, sources, target in fan_in_nodes_dot:
|
||||
digest = dot_node_id.split("::")[-1]
|
||||
base = f"{target}__{digest}"
|
||||
fan_node_id = f"fan_in__{_san(ns) + '__' if ns else ''}{_san(base)}"
|
||||
fan_in_nodes.append((fan_node_id, sources, target))
|
||||
|
||||
for fan_node_id, _, _ in fan_in_nodes:
|
||||
# Keep this line without trailing semicolon to match existing tests
|
||||
lines.append(f"{indent}{fan_node_id}((fan-in))")
|
||||
|
||||
# Fan-in edges
|
||||
for fan_node_id, sources, target in fan_in_nodes:
|
||||
for s in sources:
|
||||
lines.append(f"{indent}{map_id(s)} --> {fan_node_id};")
|
||||
lines.append(f"{indent}{fan_node_id} --> {map_id(target)};")
|
||||
|
||||
# Normal edges
|
||||
for src, tgt, is_cond in self._compute_normal_edges(
|
||||
workflow, include_internal_executors=include_internal_executors
|
||||
):
|
||||
s = map_id(src)
|
||||
t = map_id(tgt)
|
||||
if is_cond:
|
||||
lines.append(f"{indent}{s} -. conditional .-> {t};")
|
||||
else:
|
||||
lines.append(f"{indent}{s} --> {t};")
|
||||
|
||||
def _emit_sub_workflows_mermaid(
|
||||
self,
|
||||
workflow: Workflow,
|
||||
lines: list[str],
|
||||
indent: str,
|
||||
include_internal_executors: bool = False,
|
||||
) -> None:
|
||||
try:
|
||||
from ._workflow_executor import WorkflowExecutor
|
||||
except ImportError: # pragma: no cover
|
||||
return
|
||||
|
||||
def _san(s: str) -> str:
|
||||
s2 = re.sub(r"[^0-9A-Za-z_]", "_", s)
|
||||
if not s2 or not s2[0].isalpha():
|
||||
s2 = f"n_{s2}"
|
||||
return s2
|
||||
|
||||
for exec_id, exec_obj in workflow.executors.items():
|
||||
if isinstance(exec_obj, WorkflowExecutor) and hasattr(exec_obj, "workflow") and exec_obj.workflow:
|
||||
sg_id = _san(exec_id)
|
||||
lines.append(f"{indent}subgraph {sg_id}")
|
||||
# Render nested workflow within this subgraph using namespacing
|
||||
self._emit_workflow_mermaid(
|
||||
exec_obj.workflow,
|
||||
lines,
|
||||
indent=f"{indent} ",
|
||||
ns=exec_id,
|
||||
include_internal_executors=include_internal_executors,
|
||||
)
|
||||
# Recurse into deeper sub-workflows
|
||||
self._emit_sub_workflows_mermaid(
|
||||
exec_obj.workflow,
|
||||
lines,
|
||||
indent=f"{indent} ",
|
||||
include_internal_executors=include_internal_executors,
|
||||
)
|
||||
lines.append(f"{indent}end")
|
||||
|
||||
# endregion
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,896 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
import sys
|
||||
import uuid
|
||||
import warnings
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Literal
|
||||
|
||||
from .._agents import SupportsAgentRun
|
||||
from ..observability import OtelAttr, capture_exception, create_workflow_span
|
||||
from ._agent_executor import AgentExecutor
|
||||
from ._agent_utils import resolve_agent_id
|
||||
from ._checkpoint import CheckpointStorage
|
||||
from ._const import DEFAULT_MAX_ITERATIONS
|
||||
from ._edge import (
|
||||
Case,
|
||||
Default,
|
||||
EdgeCondition,
|
||||
EdgeGroup,
|
||||
FanInEdgeGroup,
|
||||
FanOutEdgeGroup,
|
||||
InternalEdgeGroup,
|
||||
SingleEdgeGroup,
|
||||
SwitchCaseEdgeGroup,
|
||||
SwitchCaseEdgeGroupCase,
|
||||
SwitchCaseEdgeGroupDefault,
|
||||
)
|
||||
from ._executor import Executor
|
||||
from ._runner_context import InProcRunnerContext
|
||||
from ._validation import ValidationTypeEnum, WorkflowValidationError, validate_workflow_graph
|
||||
from ._workflow import (
|
||||
_MISSING, # pyright: ignore[reportPrivateUsage]
|
||||
Workflow,
|
||||
_coalesce_output_from_kwarg, # pyright: ignore[reportPrivateUsage]
|
||||
)
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import Self # pragma: no cover
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_ALL_OUTPUTS: Literal["all"] = "all"
|
||||
_ALL_OTHER_OUTPUTS: Literal["all_other"] = "all_other"
|
||||
_OutputSelection = list[Executor | SupportsAgentRun] | Literal["all"] | None
|
||||
_IntermediateOutputSelection = list[Executor | SupportsAgentRun] | Literal["all", "all_other"] | None
|
||||
_AnyOutputSelection = _OutputSelection | _IntermediateOutputSelection
|
||||
|
||||
|
||||
class WorkflowBuilder:
|
||||
"""A builder class for constructing workflows.
|
||||
|
||||
This class provides a fluent API for defining workflow graphs by connecting executors
|
||||
with edges and configuring execution parameters. Call :meth:`build` to create an
|
||||
immutable :class:`Workflow` instance.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from typing_extensions import Never
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
|
||||
class UpperCaseExecutor(Executor):
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(text.upper())
|
||||
|
||||
|
||||
class ReverseExecutor(Executor):
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[Never, str]) -> None:
|
||||
await ctx.yield_output(text[::-1])
|
||||
|
||||
|
||||
upper = UpperCaseExecutor(id="upper")
|
||||
reverse = ReverseExecutor(id="reverse")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=upper).add_edge(upper, reverse).build()
|
||||
|
||||
# Run the workflow
|
||||
events = await workflow.run("hello")
|
||||
print(events.get_outputs()) # ['OLLEH']
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_iterations: int = DEFAULT_MAX_ITERATIONS,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
*,
|
||||
start_executor: Executor | SupportsAgentRun,
|
||||
checkpoint_storage: CheckpointStorage | None = None,
|
||||
output_from: list[Executor | SupportsAgentRun] | Literal["all"] | None = _MISSING,
|
||||
intermediate_output_from: _IntermediateOutputSelection = _MISSING,
|
||||
output_executors: list[Executor | SupportsAgentRun] | None = _MISSING,
|
||||
):
|
||||
"""Initialize the WorkflowBuilder.
|
||||
|
||||
Args:
|
||||
max_iterations: Maximum number of iterations for workflow convergence. Default is 100.
|
||||
name: A human-readable name for the workflow builder. This name will be the identifier
|
||||
for all workflow instances created from this builder. If not provided, a unique name
|
||||
will be generated. This will be useful for versioning, monitoring, checkpointing, and
|
||||
debugging workflows. Keeping this name unique across versions of your workflow definitions
|
||||
is recommended for better observability and management.
|
||||
description: Optional description of what the workflow does.
|
||||
start_executor: The starting executor for the workflow. Can be an Executor instance
|
||||
or SupportsAgentRun instance.
|
||||
checkpoint_storage: Optional checkpoint storage for enabling workflow state persistence.
|
||||
output_from: Designates which executors emit workflow output
|
||||
(``type='output'`` workflow events). Pass ``"all"`` to explicitly select every
|
||||
executor with declared workflow output types.
|
||||
intermediate_output_from: Designates which executors emit intermediate output
|
||||
(``type='intermediate'`` workflow events). Pass ``"all"`` to select every executor
|
||||
with declared workflow output types as intermediate (no executor emits ``output``).
|
||||
Pass ``"all_other"`` to select every executor with declared workflow output types
|
||||
that is not selected by ``output_from``.
|
||||
If neither ``output_from`` nor ``intermediate_output_from`` is provided,
|
||||
omitted-selection compatibility behavior applies and every ``yield_output`` produces
|
||||
``type='output'``. If either is provided, explicit mode applies: listed
|
||||
workflow-output executors emit ``output``, listed intermediate executors emit
|
||||
``intermediate``, and unlisted executor yields are hidden.
|
||||
|
||||
Output selection behavior:
|
||||
- Omit both selections: every ``yield_output`` emits ``output`` for compatibility,
|
||||
with a deprecation warning.
|
||||
- ``output_from="all"``: every output-capable executor emits ``output``.
|
||||
- ``output_from=[A]``: only A emits ``output``; other executor payloads are hidden.
|
||||
- ``output_from=[A], intermediate_output_from="all_other"``: A emits ``output``;
|
||||
all other output-capable executors emit ``intermediate``.
|
||||
- ``intermediate_output_from="all_other"``: no executor emits ``output``; every
|
||||
output-capable executor emits ``intermediate``.
|
||||
- ``output_from=[], intermediate_output_from="all_other"``: no executor emits
|
||||
``output``; every output-capable executor emits ``intermediate``.
|
||||
- ``output_from=[A], intermediate_output_from=[B, C]``: A emits ``output``; B and C
|
||||
emit ``intermediate``; other executor payloads are hidden.
|
||||
output_executors: **Deprecated** alias for ``output_from``. Will be removed in a
|
||||
future version.
|
||||
"""
|
||||
output_from = _coalesce_output_from_kwarg(output_from, output_executors)
|
||||
if intermediate_output_from is _MISSING:
|
||||
intermediate_output_from = None
|
||||
self._edge_groups: list[EdgeGroup] = []
|
||||
self._executors: dict[str, Executor] = {}
|
||||
self._start_executor: Executor | None = None
|
||||
self._checkpoint_storage: CheckpointStorage | None = checkpoint_storage
|
||||
self._max_iterations: int = max_iterations
|
||||
self._name: str = name or f"WorkflowBuilder-{uuid.uuid4()!s}"
|
||||
self._description: str | None = description
|
||||
# Maps underlying SupportsAgentRun object id -> wrapped Executor so we reuse the same wrapper
|
||||
# across start_executor / add_edge calls. This avoids multiple AgentExecutor instances
|
||||
# being created for the same agent.
|
||||
self._agent_wrappers: dict[str, Executor] = {}
|
||||
|
||||
# ``None`` for both means omitted-selection compatibility behavior
|
||||
# (every yield_output produces type='output').
|
||||
# If either is provided, explicit mode applies and unlisted executor yields are hidden.
|
||||
self._output_from: _OutputSelection = self._coerce_output_from(output_from)
|
||||
self._intermediate_output_from: _IntermediateOutputSelection = self._coerce_intermediate_output_from(
|
||||
intermediate_output_from
|
||||
)
|
||||
|
||||
# Set the start executor
|
||||
self._set_start_executor(start_executor)
|
||||
|
||||
# Agents auto-wrapped by builder now always stream incremental updates.
|
||||
|
||||
def _add_executor(self, executor: Executor) -> str:
|
||||
"""Add an executor to the map and return its ID."""
|
||||
existing = self._executors.get(executor.id)
|
||||
if existing is not None:
|
||||
if existing is executor:
|
||||
# Already added
|
||||
return executor.id
|
||||
# ID conflict
|
||||
raise ValueError(f"Duplicate executor ID '{executor.id}' detected in workflow.")
|
||||
|
||||
# New executor
|
||||
self._executors[executor.id] = executor
|
||||
# Add an internal edge group for each unique executor
|
||||
self._edge_groups.append(InternalEdgeGroup(executor.id))
|
||||
|
||||
return executor.id
|
||||
|
||||
def _maybe_wrap_agent(self, candidate: Executor | SupportsAgentRun) -> Executor:
|
||||
"""If the provided object implements SupportsAgentRun, wrap it in an AgentExecutor.
|
||||
|
||||
This allows fluent builder APIs to directly accept agents instead of
|
||||
requiring callers to manually instantiate AgentExecutor.
|
||||
|
||||
Args:
|
||||
candidate: The executor or agent to wrap.
|
||||
|
||||
Returns:
|
||||
An Executor instance, wrapping the agent if necessary.
|
||||
"""
|
||||
try: # Local import to avoid hard dependency at import time
|
||||
from agent_framework import SupportsAgentRun
|
||||
except Exception: # pragma: no cover - defensive
|
||||
SupportsAgentRun = object
|
||||
|
||||
if isinstance(candidate, Executor): # Already an executor
|
||||
return candidate
|
||||
if isinstance(candidate, SupportsAgentRun):
|
||||
# Reuse existing wrapper for the same agent instance if present
|
||||
agent_instance_id = str(id(candidate))
|
||||
existing = self._agent_wrappers.get(agent_instance_id)
|
||||
if existing is not None:
|
||||
return existing
|
||||
executor_id = resolve_agent_id(candidate)
|
||||
if executor_id in self._executors:
|
||||
raise ValueError(
|
||||
f"Duplicate executor ID '{executor_id}' from agent. "
|
||||
"Agent IDs or names must be unique within a workflow."
|
||||
)
|
||||
wrapper = AgentExecutor(candidate, id=executor_id)
|
||||
self._agent_wrappers[agent_instance_id] = wrapper
|
||||
return wrapper
|
||||
|
||||
raise TypeError(
|
||||
f"WorkflowBuilder expected an Executor or SupportsAgentRun instance; got {type(candidate).__name__}."
|
||||
)
|
||||
|
||||
def add_edge(
|
||||
self,
|
||||
source: Executor | SupportsAgentRun,
|
||||
target: Executor | SupportsAgentRun,
|
||||
condition: EdgeCondition | None = None,
|
||||
) -> Self:
|
||||
"""Add a directed edge between two executors.
|
||||
|
||||
The output types of the source and the input types of the target must be compatible.
|
||||
Messages sent by the source executor will be routed to the target executor.
|
||||
|
||||
Args:
|
||||
source: The source executor or agent for the edge.
|
||||
target: The target executor or agent for the edge.
|
||||
condition: An optional condition function `(data) -> bool | Awaitable[bool]`
|
||||
that determines whether the edge should be traversed.
|
||||
Example: `lambda data: data["ready"]`.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from typing_extensions import Never
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
|
||||
class ProcessorA(Executor):
|
||||
@handler
|
||||
async def process(self, data: str, ctx: WorkflowContext[int]) -> None:
|
||||
await ctx.send_message(len(data))
|
||||
|
||||
|
||||
class ProcessorB(Executor):
|
||||
@handler
|
||||
async def process(self, count: int, ctx: WorkflowContext[Never, str]) -> None:
|
||||
await ctx.yield_output(f"Processed {count} characters")
|
||||
|
||||
|
||||
a = ProcessorA(id="a")
|
||||
b = ProcessorB(id="b")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=a).add_edge(a, b).build()
|
||||
"""
|
||||
source_exec = self._maybe_wrap_agent(source)
|
||||
target_exec = self._maybe_wrap_agent(target)
|
||||
source_id = self._add_executor(source_exec)
|
||||
target_id = self._add_executor(target_exec)
|
||||
self._edge_groups.append(SingleEdgeGroup(source_id, target_id, condition))
|
||||
return self
|
||||
|
||||
def add_fan_out_edges(
|
||||
self,
|
||||
source: Executor | SupportsAgentRun,
|
||||
targets: Sequence[Executor | SupportsAgentRun],
|
||||
) -> Self:
|
||||
"""Add multiple edges to the workflow where messages from the source will be sent to all targets.
|
||||
|
||||
The output types of the source and the input types of the targets must be compatible.
|
||||
Messages from the source will be broadcast to all target executors concurrently.
|
||||
|
||||
Args:
|
||||
source: The source executor or agent for the edges.
|
||||
targets: A list of target executors or agents for the edges.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
|
||||
class DataSource(Executor):
|
||||
@handler
|
||||
async def generate(self, count: int, ctx: WorkflowContext[str]) -> None:
|
||||
for i in range(count):
|
||||
await ctx.send_message(f"data_{i}")
|
||||
|
||||
|
||||
class ValidatorA(Executor):
|
||||
@handler
|
||||
async def validate(self, data: str, ctx: WorkflowContext) -> None:
|
||||
print(f"ValidatorA: {data}")
|
||||
|
||||
|
||||
class ValidatorB(Executor):
|
||||
@handler
|
||||
async def validate(self, data: str, ctx: WorkflowContext) -> None:
|
||||
print(f"ValidatorB: {data}")
|
||||
|
||||
|
||||
source = DataSource(id="source")
|
||||
val_a = ValidatorA(id="val_a")
|
||||
val_b = ValidatorB(id="val_b")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=source).add_fan_out_edges(source, [val_a, val_b]).build()
|
||||
"""
|
||||
source_exec = self._maybe_wrap_agent(source)
|
||||
target_execs = [self._maybe_wrap_agent(t) for t in targets]
|
||||
source_id = self._add_executor(source_exec)
|
||||
target_ids = [self._add_executor(t) for t in target_execs]
|
||||
self._edge_groups.append(FanOutEdgeGroup(source_id, target_ids))
|
||||
|
||||
return self
|
||||
|
||||
def add_switch_case_edge_group(
|
||||
self,
|
||||
source: Executor | SupportsAgentRun,
|
||||
cases: Sequence[Case | Default],
|
||||
) -> Self:
|
||||
"""Add an edge group that represents a switch-case statement.
|
||||
|
||||
The output types of the source and the input types of the targets must be compatible.
|
||||
Messages from the source executor will be sent to one of the target executors based on
|
||||
the provided conditions.
|
||||
|
||||
Think of this as a switch statement where each target executor corresponds to a case.
|
||||
Each condition function will be evaluated in order, and the first one that returns True
|
||||
will determine which target executor receives the message.
|
||||
|
||||
The default case (if provided) will receive messages that fall through all conditions
|
||||
(i.e., no condition matched).
|
||||
|
||||
Args:
|
||||
source: The source executor or agent for the edge group.
|
||||
cases: A list of case objects that determine the target executor for each message.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler, Case, Default
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class Result:
|
||||
score: int
|
||||
|
||||
|
||||
class Evaluator(Executor):
|
||||
@handler
|
||||
async def evaluate(self, text: str, ctx: WorkflowContext[Result]) -> None:
|
||||
await ctx.send_message(Result(score=len(text)))
|
||||
|
||||
|
||||
class HighScoreHandler(Executor):
|
||||
@handler
|
||||
async def handle(self, result: Result, ctx: WorkflowContext) -> None:
|
||||
print(f"High score: {result.score}")
|
||||
|
||||
|
||||
class LowScoreHandler(Executor):
|
||||
@handler
|
||||
async def handle(self, result: Result, ctx: WorkflowContext) -> None:
|
||||
print(f"Low score: {result.score}")
|
||||
|
||||
|
||||
evaluator = Evaluator(id="eval")
|
||||
high = HighScoreHandler(id="high")
|
||||
low = LowScoreHandler(id="low")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=evaluator)
|
||||
.add_switch_case_edge_group(
|
||||
evaluator,
|
||||
[
|
||||
Case(condition=lambda r: r.score > 10, target=high),
|
||||
Default(target=low),
|
||||
],
|
||||
)
|
||||
.build()
|
||||
)
|
||||
"""
|
||||
source_exec = self._maybe_wrap_agent(source)
|
||||
source_id = self._add_executor(source_exec)
|
||||
# Convert case data types to internal types that only uses target_id.
|
||||
internal_cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault] = []
|
||||
for case in cases:
|
||||
# Allow case targets to be agents
|
||||
case.target = self._maybe_wrap_agent(case.target)
|
||||
self._add_executor(case.target)
|
||||
if isinstance(case, Default):
|
||||
internal_cases.append(SwitchCaseEdgeGroupDefault(target_id=case.target.id))
|
||||
else:
|
||||
internal_cases.append(SwitchCaseEdgeGroupCase(condition=case.condition, target_id=case.target.id))
|
||||
self._edge_groups.append(SwitchCaseEdgeGroup(source_id, internal_cases))
|
||||
|
||||
return self
|
||||
|
||||
def add_multi_selection_edge_group(
|
||||
self,
|
||||
source: Executor | SupportsAgentRun,
|
||||
targets: Sequence[Executor | SupportsAgentRun],
|
||||
selection_func: Callable[[Any, list[str]], list[str]],
|
||||
) -> Self:
|
||||
"""Add an edge group that represents a multi-selection execution model.
|
||||
|
||||
The output types of the source and the input types of the targets must be compatible.
|
||||
Messages from the source executor will be sent to multiple target executors based on
|
||||
the provided selection function.
|
||||
|
||||
The selection function should take a message and a list of target executor IDs,
|
||||
and return a list of executor IDs indicating which target executors should receive the message.
|
||||
|
||||
Args:
|
||||
source: The source executor or agent for the edge group.
|
||||
targets: A list of target executors or agents for the edges.
|
||||
selection_func: A function that selects target executors for messages.
|
||||
Takes (message, list[executor_id]) and returns list[executor_id].
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class Task:
|
||||
priority: str
|
||||
data: str
|
||||
|
||||
|
||||
class TaskDispatcher(Executor):
|
||||
@handler
|
||||
async def dispatch(self, text: str, ctx: WorkflowContext[Task]) -> None:
|
||||
priority = "high" if len(text) > 10 else "low"
|
||||
await ctx.send_message(Task(priority=priority, data=text))
|
||||
|
||||
|
||||
class WorkerA(Executor):
|
||||
@handler
|
||||
async def process(self, task: Task, ctx: WorkflowContext) -> None:
|
||||
print(f"WorkerA processing: {task.data}")
|
||||
|
||||
|
||||
class WorkerB(Executor):
|
||||
@handler
|
||||
async def process(self, task: Task, ctx: WorkflowContext) -> None:
|
||||
print(f"WorkerB processing: {task.data}")
|
||||
|
||||
|
||||
dispatcher = TaskDispatcher(id="dispatcher")
|
||||
worker_a = WorkerA(id="worker_a")
|
||||
worker_b = WorkerB(id="worker_b")
|
||||
|
||||
|
||||
# Select workers based on task priority
|
||||
def select_workers(task: Task, available: list[str]) -> list[str]:
|
||||
if task.priority == "high":
|
||||
return available # Send to all workers
|
||||
return [available[0]] # Send to first worker only
|
||||
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=dispatcher)
|
||||
.add_multi_selection_edge_group(
|
||||
dispatcher,
|
||||
[worker_a, worker_b],
|
||||
selection_func=select_workers,
|
||||
)
|
||||
.build()
|
||||
)
|
||||
"""
|
||||
source_exec = self._maybe_wrap_agent(source)
|
||||
target_execs = [self._maybe_wrap_agent(t) for t in targets]
|
||||
source_id = self._add_executor(source_exec)
|
||||
target_ids = [self._add_executor(t) for t in target_execs]
|
||||
self._edge_groups.append(FanOutEdgeGroup(source_id, target_ids, selection_func))
|
||||
|
||||
return self
|
||||
|
||||
def add_fan_in_edges(
|
||||
self,
|
||||
sources: Sequence[Executor | SupportsAgentRun],
|
||||
target: Executor | SupportsAgentRun,
|
||||
) -> Self:
|
||||
"""Add multiple edges from sources to a single target executor.
|
||||
|
||||
The edges will be grouped together for synchronized processing, meaning
|
||||
the target executor will only be executed once all source executors have completed.
|
||||
|
||||
The target executor will receive a list of messages aggregated from all source executors.
|
||||
Thus the input types of the target executor must be compatible with a list of the output
|
||||
types of the source executors.
|
||||
|
||||
Args:
|
||||
sources: A list of source executors or agents for the edges.
|
||||
target: The target executor or agent for the edges.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from typing_extensions import Never
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
|
||||
class Producer(Executor):
|
||||
@handler
|
||||
async def produce(self, seed: int, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(f"result_{seed}")
|
||||
|
||||
|
||||
class Aggregator(Executor):
|
||||
@handler
|
||||
async def aggregate(self, results: list[str], ctx: WorkflowContext[Never, str]) -> None:
|
||||
combined = ", ".join(results)
|
||||
await ctx.yield_output(f"Combined: {combined}")
|
||||
|
||||
|
||||
prod_1 = Producer(id="prod_1")
|
||||
prod_2 = Producer(id="prod_2")
|
||||
agg = Aggregator(id="agg")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=prod_1).add_fan_in_edges([prod_1, prod_2], agg).build()
|
||||
"""
|
||||
source_execs = [self._maybe_wrap_agent(s) for s in sources]
|
||||
target_exec = self._maybe_wrap_agent(target)
|
||||
source_ids = [self._add_executor(s) for s in source_execs]
|
||||
target_id = self._add_executor(target_exec)
|
||||
self._edge_groups.append(FanInEdgeGroup(source_ids, target_id))
|
||||
|
||||
return self
|
||||
|
||||
def add_chain(self, executors: Sequence[Executor | SupportsAgentRun]) -> Self:
|
||||
"""Add a chain of executors to the workflow.
|
||||
|
||||
The output of each executor in the chain will be sent to the next executor in the chain.
|
||||
The input types of each executor must be compatible with the output types of the previous executor.
|
||||
|
||||
Cycles in the chain are not allowed, meaning an executor cannot appear more than once in the chain.
|
||||
|
||||
Args:
|
||||
executors: A list of executors or agents to chain together.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from typing_extensions import Never
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
|
||||
class Step1(Executor):
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(text.upper())
|
||||
|
||||
|
||||
class Step2(Executor):
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(text[::-1])
|
||||
|
||||
|
||||
class Step3(Executor):
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[Never, str]) -> None:
|
||||
await ctx.yield_output(f"Final: {text}")
|
||||
|
||||
|
||||
step1 = Step1(id="step1")
|
||||
step2 = Step2(id="step2")
|
||||
step3 = Step3(id="step3")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=step1).add_chain([step1, step2, step3]).build()
|
||||
"""
|
||||
if len(executors) < 2:
|
||||
raise ValueError("At least two executors are required to form a chain.")
|
||||
|
||||
# Wrap each candidate first to ensure stable IDs before adding edges
|
||||
wrapped: list[Executor] = [self._maybe_wrap_agent(e) for e in executors]
|
||||
for i in range(len(wrapped) - 1):
|
||||
self.add_edge(wrapped[i], wrapped[i + 1])
|
||||
return self
|
||||
|
||||
def _set_start_executor(self, executor: Executor | SupportsAgentRun) -> None:
|
||||
"""Set the starting executor for the workflow (internal method).
|
||||
|
||||
Args:
|
||||
executor: The starting executor, which can be an Executor instance or SupportsAgentRun instance.
|
||||
"""
|
||||
if self._start_executor is not None:
|
||||
logger.warning(f"Overwriting existing start executor: {self._start_executor.id} for the workflow.")
|
||||
|
||||
wrapped = self._maybe_wrap_agent(executor)
|
||||
self._start_executor = wrapped
|
||||
# Ensure the start executor is present in the executor map so validation succeeds
|
||||
# even if no edges are added yet, or before edges wrap the same agent again.
|
||||
existing = self._executors.get(wrapped.id)
|
||||
if existing is not wrapped:
|
||||
self._add_executor(wrapped)
|
||||
|
||||
def _coerce_output_from(self, output_from: Any) -> _OutputSelection:
|
||||
"""Coerce workflow-output selection while preserving the explicit ``"all"`` literal."""
|
||||
if output_from is None:
|
||||
return None
|
||||
if output_from == _ALL_OUTPUTS:
|
||||
return _ALL_OUTPUTS
|
||||
if isinstance(output_from, str):
|
||||
raise ValueError(f"Unsupported output_from literal {output_from!r}; use 'all' or a list of executors.")
|
||||
return list(output_from)
|
||||
|
||||
def _coerce_intermediate_output_from(self, intermediate_output_from: Any) -> _IntermediateOutputSelection:
|
||||
"""Coerce intermediate-output selection and reject output-only literals."""
|
||||
if intermediate_output_from is None:
|
||||
return None
|
||||
if isinstance(intermediate_output_from, str):
|
||||
if intermediate_output_from == _ALL_OUTPUTS:
|
||||
return _ALL_OUTPUTS
|
||||
if intermediate_output_from == _ALL_OTHER_OUTPUTS:
|
||||
return _ALL_OTHER_OUTPUTS
|
||||
raise ValueError(
|
||||
f"Unsupported intermediate_output_from literal {intermediate_output_from!r}; "
|
||||
"use 'all', 'all_other', or a list of executors."
|
||||
)
|
||||
return list(intermediate_output_from)
|
||||
|
||||
def _resolve_designated_executor_ids(
|
||||
self,
|
||||
designated: _AnyOutputSelection,
|
||||
) -> list[str] | None:
|
||||
"""Resolve an optional designation list into executor IDs without mutating the graph."""
|
||||
if designated is None:
|
||||
return None
|
||||
if designated == _ALL_OUTPUTS:
|
||||
return [executor_id for executor_id, executor in self._executors.items() if executor.workflow_output_types]
|
||||
if designated == _ALL_OTHER_OUTPUTS:
|
||||
raise ValueError("intermediate_output_from='all_other' must be expanded relative to output_from.")
|
||||
ids: list[str] = []
|
||||
for item in designated:
|
||||
if isinstance(item, Executor):
|
||||
ids.append(item.id)
|
||||
elif isinstance(item, SupportsAgentRun):
|
||||
ids.append(resolve_agent_id(item))
|
||||
else:
|
||||
raise TypeError(
|
||||
"WorkflowBuilder expected designation entries to be Executor or SupportsAgentRun instances; "
|
||||
f"got {type(item).__name__}."
|
||||
)
|
||||
return ids
|
||||
|
||||
def _validate_designation_lists(
|
||||
self,
|
||||
output_executor_ids: list[str] | None,
|
||||
intermediate_executor_ids: list[str] | None,
|
||||
) -> None:
|
||||
"""Validate builder-level designation rules that need omitted-vs-explicit context."""
|
||||
explicit_mode = output_executor_ids is not None or intermediate_executor_ids is not None
|
||||
if not explicit_mode:
|
||||
return
|
||||
|
||||
output_ids = output_executor_ids or []
|
||||
intermediate_ids = intermediate_executor_ids or []
|
||||
if not output_ids and not intermediate_ids:
|
||||
raise WorkflowValidationError(
|
||||
"Explicit workflow output designation must include at least one output or intermediate executor.",
|
||||
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
|
||||
)
|
||||
|
||||
duplicate_outputs = sorted({executor_id for executor_id in output_ids if output_ids.count(executor_id) > 1})
|
||||
if duplicate_outputs:
|
||||
raise WorkflowValidationError(
|
||||
f"Duplicate output executor designation(s): {duplicate_outputs}",
|
||||
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
|
||||
)
|
||||
|
||||
duplicate_intermediates = sorted({
|
||||
executor_id for executor_id in intermediate_ids if intermediate_ids.count(executor_id) > 1
|
||||
})
|
||||
if duplicate_intermediates:
|
||||
raise WorkflowValidationError(
|
||||
f"Duplicate intermediate executor designation(s): {duplicate_intermediates}",
|
||||
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
|
||||
)
|
||||
|
||||
overlap = sorted(set(output_ids).intersection(intermediate_ids))
|
||||
if overlap:
|
||||
raise WorkflowValidationError(
|
||||
f"Executors cannot be both output and intermediate designated: {overlap}",
|
||||
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
|
||||
)
|
||||
|
||||
def build(self) -> Workflow:
|
||||
"""Build and return the constructed workflow.
|
||||
|
||||
This method performs validation before building the workflow to ensure:
|
||||
- A starting executor has been set
|
||||
- All edges connect valid executors
|
||||
- The graph is properly connected
|
||||
- Type compatibility between connected executors
|
||||
|
||||
Returns:
|
||||
Workflow: An immutable Workflow instance ready for execution.
|
||||
|
||||
Raises:
|
||||
ValueError: If starting executor is not set.
|
||||
WorkflowValidationError: If workflow validation fails (includes EdgeDuplicationError,
|
||||
TypeCompatibilityError, and GraphConnectivityError subclasses).
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from typing_extensions import Never
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
|
||||
class MyExecutor(Executor):
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[Never, str]) -> None:
|
||||
await ctx.yield_output(text.upper())
|
||||
|
||||
|
||||
executor = MyExecutor(id="executor")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=executor).build()
|
||||
|
||||
# The workflow is now immutable and ready to run
|
||||
events = await workflow.run("hello")
|
||||
print(events.get_outputs()) # ['HELLO']
|
||||
|
||||
# Workflows can be reused multiple times
|
||||
events2 = await workflow.run("world")
|
||||
print(events2.get_outputs()) # ['WORLD']
|
||||
|
||||
# Select one executor as Workflow Output.
|
||||
workflow = WorkflowBuilder(start_executor=executor, output_from=[executor]).build()
|
||||
events = await workflow.run("hello")
|
||||
print(events.get_outputs()) # ['HELLO']
|
||||
print(events.get_intermediate_outputs()) # []
|
||||
|
||||
# Make one executor Workflow Output and every other output-capable executor Intermediate Output.
|
||||
workflow = (
|
||||
WorkflowBuilder(
|
||||
start_executor=planner,
|
||||
output_from=[answerer],
|
||||
intermediate_output_from="all_other",
|
||||
)
|
||||
.add_edge(planner, answerer)
|
||||
.build()
|
||||
)
|
||||
events = await workflow.run("hello")
|
||||
print(events.get_outputs()) # outputs from answerer
|
||||
print(events.get_intermediate_outputs()) # outputs from planner
|
||||
|
||||
# Build a progress-only workflow: no Workflow Output, all output-capable executors are intermediate.
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=planner, intermediate_output_from="all_other")
|
||||
.add_edge(planner, answerer)
|
||||
.build()
|
||||
)
|
||||
events = await workflow.run("hello")
|
||||
print(events.get_outputs()) # []
|
||||
print(events.get_intermediate_outputs()) # outputs from planner and answerer
|
||||
|
||||
# Explicitly preserve all-output behavior without relying on omitted-selection compatibility.
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=planner, output_from="all").add_edge(planner, answerer).build()
|
||||
)
|
||||
events = await workflow.run("hello")
|
||||
print(events.get_outputs()) # outputs from planner and answerer
|
||||
"""
|
||||
# Create workflow build span that includes validation and workflow creation
|
||||
with create_workflow_span(OtelAttr.WORKFLOW_BUILD_SPAN) as span:
|
||||
try:
|
||||
# Add workflow build started event
|
||||
span.add_event(OtelAttr.BUILD_STARTED)
|
||||
|
||||
if not self._start_executor:
|
||||
raise ValueError(
|
||||
"Starting executor must be set via the start_executor constructor parameter before building."
|
||||
)
|
||||
|
||||
if self._output_from is None and self._intermediate_output_from is None:
|
||||
warnings.warn(
|
||||
"WorkflowBuilder built without explicit output_from or intermediate_output_from; "
|
||||
"every yield_output produces type='output' for compatibility. Pass output_from='all', "
|
||||
"output_from=[...], or intermediate_output_from=[...] to opt into explicit designation - "
|
||||
"explicit designation will be required in a future version.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
start_executor = self._start_executor
|
||||
executors = self._executors
|
||||
edge_groups = self._edge_groups
|
||||
output_ids = self._resolve_designated_executor_ids(self._output_from)
|
||||
intermediate_output_ids: list[str] | None
|
||||
if self._intermediate_output_from == _ALL_OTHER_OUTPUTS:
|
||||
output_ids_for_all_other = output_ids or []
|
||||
intermediate_output_ids = [
|
||||
executor_id
|
||||
for executor_id, executor in self._executors.items()
|
||||
if executor.workflow_output_types and executor_id not in output_ids_for_all_other
|
||||
]
|
||||
else:
|
||||
intermediate_output_ids = self._resolve_designated_executor_ids(self._intermediate_output_from)
|
||||
self._validate_designation_lists(output_ids, intermediate_output_ids)
|
||||
|
||||
explicit_mode = output_ids is not None or intermediate_output_ids is not None
|
||||
output_for_workflow: list[str] | None = output_ids if explicit_mode else None
|
||||
if explicit_mode and output_for_workflow is None:
|
||||
output_for_workflow = []
|
||||
intermediate_output_for_workflow: list[str] | None = intermediate_output_ids if explicit_mode else None
|
||||
if explicit_mode and intermediate_output_for_workflow is None:
|
||||
intermediate_output_for_workflow = []
|
||||
|
||||
# Perform validation before creating the workflow
|
||||
validate_workflow_graph(
|
||||
edge_groups,
|
||||
executors,
|
||||
start_executor,
|
||||
output_for_workflow or [],
|
||||
intermediate_output_for_workflow or [],
|
||||
)
|
||||
|
||||
# Add validation completed event
|
||||
span.add_event(OtelAttr.BUILD_VALIDATION_COMPLETED)
|
||||
|
||||
context = InProcRunnerContext(self._checkpoint_storage)
|
||||
|
||||
# Create workflow instance after validation
|
||||
workflow = Workflow(
|
||||
edge_groups,
|
||||
executors,
|
||||
start_executor,
|
||||
context,
|
||||
self._name,
|
||||
description=self._description,
|
||||
max_iterations=self._max_iterations,
|
||||
output_from=output_for_workflow,
|
||||
intermediate_output_from=intermediate_output_for_workflow,
|
||||
)
|
||||
build_attributes: dict[str, Any] = {
|
||||
OtelAttr.WORKFLOW_BUILDER_NAME: self._name,
|
||||
OtelAttr.WORKFLOW_ID: workflow.id,
|
||||
OtelAttr.WORKFLOW_DEFINITION: workflow.to_json(),
|
||||
}
|
||||
if self._description:
|
||||
build_attributes[OtelAttr.WORKFLOW_BUILDER_DESCRIPTION] = self._description
|
||||
span.set_attributes(build_attributes)
|
||||
|
||||
# Add workflow build completed event
|
||||
span.add_event(OtelAttr.BUILD_COMPLETED)
|
||||
|
||||
return workflow
|
||||
|
||||
except Exception as exc:
|
||||
attributes = {
|
||||
OtelAttr.BUILD_ERROR_MESSAGE: str(exc),
|
||||
OtelAttr.BUILD_ERROR_TYPE: type(exc).__name__,
|
||||
}
|
||||
span.add_event(OtelAttr.BUILD_ERROR, attributes) # type: ignore[reportArgumentType, arg-type]
|
||||
capture_exception(span, exc)
|
||||
raise
|
||||
@@ -0,0 +1,489 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import inspect
|
||||
import logging
|
||||
import uuid
|
||||
from types import UnionType
|
||||
from typing import TYPE_CHECKING, Any, Generic, Union, cast, get_args, get_origin
|
||||
|
||||
from opentelemetry.propagate import inject
|
||||
from opentelemetry.trace import SpanKind
|
||||
from typing_extensions import Never, TypeVar
|
||||
|
||||
from ..observability import OtelAttr, create_workflow_span
|
||||
from ._events import (
|
||||
WorkflowEvent,
|
||||
WorkflowEventSource,
|
||||
_framework_event_origin, # type: ignore
|
||||
)
|
||||
from ._runner_context import RunnerContext, WorkflowMessage
|
||||
from ._state import State
|
||||
from ._typing_utils import contains_typevar
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._executor import Executor
|
||||
|
||||
OutT = TypeVar("OutT", default=Never)
|
||||
W_OutT = TypeVar("W_OutT", default=Never)
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def infer_output_types_from_ctx_annotation(
|
||||
ctx_annotation: Any,
|
||||
) -> tuple[list[type[Any] | UnionType], list[type[Any] | UnionType]]:
|
||||
"""Infer message types and workflow output types from the WorkflowContext generic parameters.
|
||||
|
||||
Examples:
|
||||
- WorkflowContext -> ([], [])
|
||||
- WorkflowContext[str] -> ([str], [])
|
||||
- WorkflowContext[str, int] -> ([str], [int])
|
||||
- WorkflowContext[str | int, bool | int] -> ([str, int], [bool, int])
|
||||
- WorkflowContext[Union[str, int], Union[bool, int]] -> ([str, int], [bool, int])
|
||||
- WorkflowContext[Any] -> ([Any], [])
|
||||
- WorkflowContext[Any, Any] -> ([Any], [Any])
|
||||
- WorkflowContext[Never, Never] -> ([], [])
|
||||
- WorkflowContext[Never, int] -> ([], [int])
|
||||
|
||||
Returns:
|
||||
Tuple of (message_types, workflow_output_types)
|
||||
"""
|
||||
# If no annotation or not parameterized, return empty lists
|
||||
try:
|
||||
origin = get_origin(ctx_annotation)
|
||||
except Exception:
|
||||
origin = None
|
||||
|
||||
# If annotation is unsubscripted WorkflowContext, nothing to infer
|
||||
if origin is None:
|
||||
return [], []
|
||||
|
||||
# Expecting WorkflowContext[OutT, W_OutT]
|
||||
if origin is not WorkflowContext:
|
||||
return [], []
|
||||
|
||||
args = list(get_args(ctx_annotation))
|
||||
if not args:
|
||||
return [], []
|
||||
|
||||
# WorkflowContext[OutT] -> message_types from OutT, no workflow output types
|
||||
if len(args) == 1:
|
||||
t = args[0]
|
||||
t_origin = get_origin(t)
|
||||
if t is Any:
|
||||
return [cast(type[Any], Any)], []
|
||||
|
||||
if t_origin in (Union, UnionType):
|
||||
msg_types: list[type[Any] | UnionType] = [arg for arg in get_args(t) if arg is not Any and arg is not Never]
|
||||
return msg_types, []
|
||||
|
||||
if t is Never:
|
||||
return [], []
|
||||
return [t], []
|
||||
|
||||
# WorkflowContext[OutT, W_OutT] -> message_types from OutT, workflow_output_types from W_OutT
|
||||
t_out, t_w_out = args[:2] # Take first two args in case there are more
|
||||
|
||||
# Process OutT for message_types
|
||||
message_types: list[type[Any] | UnionType] = []
|
||||
t_out_origin = get_origin(t_out)
|
||||
if t_out is Any:
|
||||
message_types = [cast(type[Any], Any)]
|
||||
elif t_out is not Never:
|
||||
if t_out_origin in (Union, UnionType):
|
||||
message_types = [arg for arg in get_args(t_out) if arg is not Any and arg is not Never]
|
||||
else:
|
||||
message_types = [t_out]
|
||||
|
||||
# Process W_OutT for workflow_output_types
|
||||
workflow_output_types: list[type[Any] | UnionType] = []
|
||||
t_w_out_origin = get_origin(t_w_out)
|
||||
if t_w_out is Any:
|
||||
workflow_output_types = [cast(type[Any], Any)]
|
||||
elif t_w_out is not Never:
|
||||
if t_w_out_origin in (Union, UnionType):
|
||||
workflow_output_types = [arg for arg in get_args(t_w_out) if arg is not Any and arg is not Never]
|
||||
else:
|
||||
workflow_output_types = [t_w_out]
|
||||
|
||||
return message_types, workflow_output_types
|
||||
|
||||
|
||||
def _is_workflow_context_type(annotation: Any) -> bool:
|
||||
"""Check if an annotation represents WorkflowContext, WorkflowContext[T], or WorkflowContext[T, U]."""
|
||||
origin = get_origin(annotation)
|
||||
if origin is WorkflowContext:
|
||||
return True
|
||||
# Also handle the case where the raw class is used
|
||||
return annotation is WorkflowContext
|
||||
|
||||
|
||||
def validate_workflow_context_annotation(
|
||||
annotation: Any,
|
||||
parameter_name: str,
|
||||
context_description: str,
|
||||
) -> tuple[list[type[Any] | UnionType], list[type[Any] | UnionType]]:
|
||||
"""Validate a WorkflowContext annotation and return inferred types.
|
||||
|
||||
Args:
|
||||
annotation: The type annotation to validate
|
||||
parameter_name: Name of the parameter (for error messages)
|
||||
context_description: Description of the context (e.g., "Function func1", "Handler method")
|
||||
|
||||
Returns:
|
||||
Tuple of (output_types, workflow_output_types)
|
||||
|
||||
Raises:
|
||||
ValueError: If the annotation is invalid
|
||||
"""
|
||||
if annotation == inspect.Parameter.empty:
|
||||
raise ValueError(
|
||||
f"{context_description} {parameter_name} must have a WorkflowContext, "
|
||||
f"WorkflowContext[T] or WorkflowContext[T, U] type annotation, "
|
||||
f"where T is output message type and U is workflow output type"
|
||||
)
|
||||
|
||||
if not _is_workflow_context_type(annotation):
|
||||
raise ValueError(
|
||||
f"{context_description} {parameter_name} must be annotated as "
|
||||
f"WorkflowContext, WorkflowContext[T], or WorkflowContext[T, U], "
|
||||
f"got {annotation}"
|
||||
)
|
||||
|
||||
# Validate type arguments for WorkflowContext[T] or WorkflowContext[T, U]
|
||||
type_args = get_args(annotation)
|
||||
|
||||
if len(type_args) > 2:
|
||||
raise ValueError(
|
||||
f"{context_description} {parameter_name} must have at most 2 type arguments, "
|
||||
"WorkflowContext, WorkflowContext[T], or WorkflowContext[T, U], "
|
||||
f"got {len(type_args)} arguments"
|
||||
)
|
||||
|
||||
if type_args:
|
||||
# Helper function to check if a value is a valid type annotation
|
||||
def _is_type_like(x: Any) -> bool:
|
||||
"""Check if a value is a type-like entity (class, type, or typing construct)."""
|
||||
return isinstance(x, type) or get_origin(x) is not None or x is Never
|
||||
|
||||
for i, type_arg in enumerate(type_args):
|
||||
param_description = "OutT" if i == 0 else "W_OutT"
|
||||
|
||||
# Allow Any explicitly
|
||||
if type_arg is Any:
|
||||
continue
|
||||
|
||||
# Check for unresolved TypeVar early with an actionable error message
|
||||
if contains_typevar(type_arg):
|
||||
raise ValueError(
|
||||
f"{context_description} {parameter_name} {param_description} "
|
||||
f"contains an unresolved TypeVar in '{type_arg}'. "
|
||||
f"Use @handler(input=ConcreteType, output=ConcreteType) with concrete types "
|
||||
f"for parameterized executors."
|
||||
)
|
||||
|
||||
# Check if it's a union type and validate each member
|
||||
union_origin = get_origin(type_arg)
|
||||
if union_origin in (Union, UnionType):
|
||||
union_members = get_args(type_arg)
|
||||
invalid_members = [m for m in union_members if not _is_type_like(m) and m is not Any]
|
||||
if invalid_members:
|
||||
raise ValueError(
|
||||
f"{context_description} {parameter_name} {param_description} "
|
||||
f"contains invalid type entries: {invalid_members}. "
|
||||
f"Use proper types or typing generics"
|
||||
)
|
||||
else:
|
||||
# Check if it's a valid type
|
||||
if not _is_type_like(type_arg):
|
||||
raise ValueError(
|
||||
f"{context_description} {parameter_name} {param_description} "
|
||||
f"contains invalid type entry: {type_arg}. "
|
||||
f"Use proper types or typing generics"
|
||||
)
|
||||
|
||||
return infer_output_types_from_ctx_annotation(annotation)
|
||||
|
||||
|
||||
# Event types reserved for framework lifecycle (not allowed from user code)
|
||||
_FRAMEWORK_LIFECYCLE_EVENT_TYPES: frozenset[str] = frozenset({"started", "status", "failed"})
|
||||
_OUTPUT_SELECTION_EVENT_TYPES: frozenset[str] = frozenset({"output", "intermediate"})
|
||||
|
||||
|
||||
class WorkflowContext(Generic[OutT, W_OutT]):
|
||||
"""Execution context that enables executors to interact with workflows and other executors.
|
||||
|
||||
## Overview
|
||||
WorkflowContext provides a controlled interface for executors to send messages, yield outputs,
|
||||
manage state, and interact with the broader workflow ecosystem. It enforces type safety through
|
||||
generic parameters while preventing direct access to internal runtime components.
|
||||
|
||||
## Type Parameters
|
||||
The context is parameterized to enforce type safety for different operations:
|
||||
|
||||
### WorkflowContext (no parameters)
|
||||
For executors that only perform side effects without sending messages or yielding outputs:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
async def log_handler(message: str, ctx: WorkflowContext) -> None:
|
||||
print(f"Received: {message}") # Only side effects
|
||||
|
||||
### WorkflowContext[OutT]
|
||||
Enables sending messages of type OutT to other executors:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
async def processor(message: str, ctx: WorkflowContext[int]) -> None:
|
||||
result = len(message)
|
||||
await ctx.send_message(result) # Send int to downstream executors
|
||||
|
||||
### WorkflowContext[OutT, W_OutT]
|
||||
Enables both sending messages (OutT) and yielding workflow outputs (W_OutT):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
async def dual_output(message: str, ctx: WorkflowContext[int, str]) -> None:
|
||||
await ctx.send_message(42) # Send int message
|
||||
await ctx.yield_output("complete") # Yield str workflow output
|
||||
|
||||
### Union Types
|
||||
Multiple types can be specified using union notation:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
async def flexible(message: str, ctx: WorkflowContext[int | str, bool | dict]) -> None:
|
||||
await ctx.send_message("text") # or send 42
|
||||
await ctx.yield_output(True) # or yield {"status": "done"}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
executor: Executor,
|
||||
source_executor_ids: list[str],
|
||||
state: State,
|
||||
runner_context: RunnerContext,
|
||||
trace_contexts: list[dict[str, str]] | None = None,
|
||||
source_span_ids: list[str] | None = None,
|
||||
request_id: str | None = None,
|
||||
):
|
||||
"""Initialize the executor context with the given workflow context.
|
||||
|
||||
Args:
|
||||
executor: The executor instance that this context belongs to.
|
||||
source_executor_ids: The IDs of the source executors that sent messages to this executor.
|
||||
This is a list to support fan_in scenarios where multiple sources send aggregated
|
||||
messages to the same executor.
|
||||
state: The workflow state.
|
||||
runner_context: The runner context that provides methods to send messages and events.
|
||||
trace_contexts: Optional trace contexts from multiple sources for OpenTelemetry propagation.
|
||||
source_span_ids: Optional source span IDs from multiple sources for linking (not for nesting).
|
||||
request_id: Optional request ID if this context is for a `handle_response` handler.
|
||||
"""
|
||||
self._executor = executor
|
||||
self._executor_id = executor.id
|
||||
self._source_executor_ids = source_executor_ids
|
||||
self._runner_context = runner_context
|
||||
self._state = state
|
||||
|
||||
# Track messages sent via send_message() for executor_completed event (type='executor_completed')
|
||||
self._sent_messages: list[Any] = []
|
||||
|
||||
# Track outputs yielded via yield_output() for executor_completed event (type='executor_completed')
|
||||
self._yielded_outputs: list[Any] = []
|
||||
|
||||
# Store trace contexts and source span IDs for linking (supporting multiple sources)
|
||||
self._trace_contexts = trace_contexts or []
|
||||
self._source_span_ids = source_span_ids or []
|
||||
|
||||
# request info related
|
||||
self._request_id: str | None = request_id
|
||||
|
||||
if not self._source_executor_ids:
|
||||
raise ValueError("source_executor_ids cannot be empty. At least one source executor ID is required.")
|
||||
|
||||
@property
|
||||
def request_id(self) -> str | None:
|
||||
"""Get the request ID if this context is for a `handle_response` handler.
|
||||
|
||||
Returns:
|
||||
The request ID string or None if not applicable.
|
||||
"""
|
||||
return self._request_id
|
||||
|
||||
async def send_message(self, message: OutT, target_id: str | None = None) -> None:
|
||||
"""Send a message to the workflow context.
|
||||
|
||||
Args:
|
||||
message: The message to send. This must conform to the output type(s) declared on this context.
|
||||
target_id: The ID of the target executor to send the message to.
|
||||
If None, the message will be sent to all target executors.
|
||||
"""
|
||||
global OBSERVABILITY_SETTINGS
|
||||
from ..observability import OBSERVABILITY_SETTINGS
|
||||
|
||||
# Create publishing span (inherits current trace context automatically)
|
||||
attributes: dict[str, str] = {OtelAttr.MESSAGE_TYPE: type(message).__name__}
|
||||
if target_id:
|
||||
attributes[OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID] = target_id
|
||||
with create_workflow_span(OtelAttr.MESSAGE_SEND_SPAN, attributes, kind=SpanKind.PRODUCER) as span:
|
||||
# Create Message wrapper
|
||||
msg = WorkflowMessage(data=message, source_id=self._executor_id, target_id=target_id)
|
||||
|
||||
# Track sent message for executor_completed event (type='executor_completed')
|
||||
self._sent_messages.append(message)
|
||||
|
||||
# Inject current trace context if tracing enabled
|
||||
if OBSERVABILITY_SETTINGS.ENABLED and span and span.is_recording():
|
||||
trace_context: dict[str, str] = {}
|
||||
inject(trace_context) # Inject current trace context for message propagation
|
||||
|
||||
msg.trace_contexts = [trace_context]
|
||||
msg.source_span_ids = [format(span.get_span_context().span_id, "016x")]
|
||||
|
||||
await self._runner_context.send_message(msg)
|
||||
|
||||
async def yield_output(self, output: W_OutT) -> None:
|
||||
"""Yield an output from this executor.
|
||||
|
||||
The framework labels the resulting workflow event based on the workflow's explicit
|
||||
output designation:
|
||||
|
||||
- Omitted-selection compatibility behavior: every yield produces ``type='output'``.
|
||||
- Explicit mode: output-designated executors produce ``type='output'``,
|
||||
intermediate-designated executors produce ``type='intermediate'``, and
|
||||
unlisted executor yields are hidden from caller-facing events.
|
||||
|
||||
Whether a given executor produces ``output`` or ``intermediate`` events is fixed at
|
||||
workflow-build time via ``output_from`` / ``intermediate_output_from`` on
|
||||
:class:`WorkflowBuilder`; an executor cannot vary the label per yield. To change an
|
||||
executor's role, list it under a different designation when building the workflow.
|
||||
|
||||
Args:
|
||||
output: The output to yield. This must conform to the workflow output type(s)
|
||||
declared on this context.
|
||||
"""
|
||||
# Track yielded output for executor_completed event (type='executor_completed')
|
||||
# (deepcopy to capture state at yield time)
|
||||
self._yielded_outputs.append(copy.deepcopy(output))
|
||||
|
||||
event_type = self._runner_context.classify_yielded_output(self._executor_id)
|
||||
if event_type is None:
|
||||
return
|
||||
|
||||
with _framework_event_origin():
|
||||
event = WorkflowEvent(event_type, executor_id=self._executor_id, data=output)
|
||||
await self._runner_context.add_event(event)
|
||||
|
||||
async def add_event(self, event: WorkflowEvent[Any]) -> None:
|
||||
"""Add an event to the workflow context."""
|
||||
if event.origin == WorkflowEventSource.EXECUTOR and event.type in _OUTPUT_SELECTION_EVENT_TYPES:
|
||||
warning_msg = (
|
||||
f"Executor '{self._executor_id}' attempted to emit a '{event.type}' event directly, "
|
||||
"which is reserved for ctx.yield_output(). The event was ignored."
|
||||
)
|
||||
logger.warning(warning_msg)
|
||||
await self._runner_context.add_event(WorkflowEvent.warning(warning_msg))
|
||||
return
|
||||
if event.origin == WorkflowEventSource.EXECUTOR and event.type in _FRAMEWORK_LIFECYCLE_EVENT_TYPES:
|
||||
warning_msg = (
|
||||
f"Executor '{self._executor_id}' attempted to emit a '{event.type}' event, "
|
||||
"which is reserved for framework lifecycle notifications. The "
|
||||
"event was ignored."
|
||||
)
|
||||
logger.warning(warning_msg)
|
||||
await self._runner_context.add_event(WorkflowEvent.warning(warning_msg))
|
||||
return
|
||||
await self._runner_context.add_event(event)
|
||||
|
||||
async def request_info(self, request_data: object, response_type: type, *, request_id: str | None = None) -> None:
|
||||
"""Request information from outside of the workflow.
|
||||
|
||||
Calling this method will cause the workflow to emit a request_info event (type='request_info'), carrying the
|
||||
provided request_data and request_type. External systems listening for such events
|
||||
can then process the request and respond accordingly.
|
||||
|
||||
Executors must have the corresponding response handlers defined using the
|
||||
@response_handler decorator to handle the incoming responses.
|
||||
|
||||
Args:
|
||||
request_data: The data associated with the information request.
|
||||
response_type: The expected type of the response, used for validation.
|
||||
request_id: Optional unique identifier for the request. If not provided,
|
||||
a new UUID will be generated. This allows executors to track requests and responses.
|
||||
"""
|
||||
request_type: type = type(request_data)
|
||||
if not self._executor.is_request_supported(request_type, response_type):
|
||||
logger.warning(
|
||||
f"Executor '{self._executor_id}' requested info of type {request_type.__name__} "
|
||||
f"with expected response type {response_type.__name__}, but no matching "
|
||||
"response handler is defined. The request will not be ignored but responses will "
|
||||
"not be processed. Please define a response handler using the @response_handler decorator."
|
||||
)
|
||||
|
||||
request_info_event = WorkflowEvent.request_info(
|
||||
request_id=request_id or str(uuid.uuid4()),
|
||||
source_executor_id=self._executor_id,
|
||||
request_data=request_data,
|
||||
response_type=response_type,
|
||||
)
|
||||
await self._runner_context.add_request_info_event(request_info_event)
|
||||
|
||||
def get_state(self, key: str, default: Any = None) -> Any:
|
||||
"""Get a value from the workflow state."""
|
||||
return self._state.get(key, default)
|
||||
|
||||
def set_state(self, key: str, value: Any) -> None:
|
||||
"""Set a value in the workflow state."""
|
||||
self._state.set(key, value)
|
||||
|
||||
def get_source_executor_id(self) -> str:
|
||||
"""Get the ID of the source executor that sent the message to this executor.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If there are multiple source executors, this method raises an error.
|
||||
"""
|
||||
if len(self._source_executor_ids) > 1:
|
||||
raise RuntimeError(
|
||||
"Cannot get source executor ID when there are multiple source executors. "
|
||||
"Access the full list via the source_executor_ids property instead."
|
||||
)
|
||||
return self._source_executor_ids[0]
|
||||
|
||||
@property
|
||||
def source_executor_ids(self) -> list[str]:
|
||||
"""Get the IDs of the source executors that sent messages to this executor."""
|
||||
return self._source_executor_ids
|
||||
|
||||
@property
|
||||
def state(self) -> State:
|
||||
"""Get the workflow state."""
|
||||
return self._state
|
||||
|
||||
def get_sent_messages(self) -> list[Any]:
|
||||
"""Get all messages sent via send_message() during this handler execution.
|
||||
|
||||
Returns:
|
||||
A list of messages that were sent to downstream executors.
|
||||
"""
|
||||
return self._sent_messages.copy()
|
||||
|
||||
def get_yielded_outputs(self) -> list[Any]:
|
||||
"""Get all outputs yielded via yield_output() during this handler execution.
|
||||
|
||||
Returns:
|
||||
A list of outputs that were yielded as workflow outputs.
|
||||
"""
|
||||
return self._yielded_outputs.copy()
|
||||
|
||||
def is_streaming(self) -> bool:
|
||||
"""Check if the workflow is running in streaming mode.
|
||||
|
||||
Returns:
|
||||
True if the workflow was started with stream=True, False otherwise.
|
||||
"""
|
||||
return self._runner_context.is_streaming()
|
||||
@@ -0,0 +1,697 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import sys
|
||||
import types
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._workflow import Workflow
|
||||
|
||||
from ._checkpoint_encoding import decode_checkpoint_value
|
||||
from ._const import GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY
|
||||
from ._events import (
|
||||
WorkflowEvent,
|
||||
WorkflowRunState,
|
||||
_framework_event_origin, # type: ignore[reportPrivateUsage]
|
||||
)
|
||||
from ._executor import Executor, handler
|
||||
from ._request_info_mixin import response_handler
|
||||
from ._runner_context import WorkflowMessage
|
||||
from ._typing_utils import is_instance_of
|
||||
from ._workflow import WorkflowRunResult
|
||||
from ._workflow_context import WorkflowContext
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # pragma: no cover
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExecutionContext:
|
||||
"""Context for tracking a single sub-workflow execution."""
|
||||
|
||||
# The ID of the execution context
|
||||
execution_id: str
|
||||
|
||||
# Responses that have been collected so far for requests that
|
||||
# were sent out in the previous iteration
|
||||
collected_responses: dict[str, Any] # request_id -> response_data
|
||||
|
||||
# Number of responses to be expected. If the WorkflowExecutor has
|
||||
# not received all responses, it won't run the sub workflow.
|
||||
expected_response_count: int
|
||||
|
||||
# Pending requests to be fulfilled. This will get updated as the
|
||||
# WorkflowExecutor receives responses.
|
||||
pending_requests: dict[str, WorkflowEvent] # request_id -> request_info_event
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubWorkflowResponseMessage:
|
||||
"""Message sent from a parent workflow to a sub-workflow via WorkflowExecutor to provide requested information.
|
||||
|
||||
This message wraps the response data along with the original WorkflowEvent emitted by the sub-workflow executor.
|
||||
|
||||
Attributes:
|
||||
data: The response data to the original request.
|
||||
source_event: The original WorkflowEvent emitted by the sub-workflow executor.
|
||||
"""
|
||||
|
||||
data: Any
|
||||
source_event: WorkflowEvent
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubWorkflowRequestMessage:
|
||||
"""Message sent from a sub-workflow to an executor in the parent workflow to request information.
|
||||
|
||||
This message wraps a WorkflowEvent emitted by the executor in the sub-workflow.
|
||||
|
||||
Attributes:
|
||||
source_event: The original WorkflowEvent emitted by the sub-workflow executor.
|
||||
executor_id: The ID of the WorkflowExecutor in the parent workflow that is
|
||||
responsible for this sub-workflow. This can be used to ensure that the response
|
||||
is sent back to the correct sub-workflow instance.
|
||||
"""
|
||||
|
||||
source_event: WorkflowEvent
|
||||
executor_id: str
|
||||
|
||||
def create_response(self, data: Any) -> SubWorkflowResponseMessage:
|
||||
"""Validate and wrap response data into a SubWorkflowResponseMessage.
|
||||
|
||||
Validation ensures the response data type matches the expected type from the original request.
|
||||
"""
|
||||
expected_data_type = self.source_event.response_type
|
||||
if not is_instance_of(data, expected_data_type):
|
||||
raise TypeError(
|
||||
f"Response data type {type(data)} does not match expected type {expected_data_type} "
|
||||
f"for request_id {self.source_event.request_id}"
|
||||
)
|
||||
|
||||
return SubWorkflowResponseMessage(data=data, source_event=self.source_event)
|
||||
|
||||
|
||||
class WorkflowExecutor(Executor):
|
||||
"""An executor that wraps a workflow to enable hierarchical workflow composition.
|
||||
|
||||
## Overview
|
||||
WorkflowExecutor makes a workflow behave as a single executor within a parent workflow,
|
||||
enabling nested workflow architectures. It handles the complete lifecycle of sub-workflow
|
||||
execution including event processing, output forwarding, and request/response coordination
|
||||
between parent and child workflows.
|
||||
|
||||
## Execution Model
|
||||
When invoked, WorkflowExecutor:
|
||||
1. Starts the wrapped workflow with the input message
|
||||
2. Runs the sub-workflow to completion or until it needs external input
|
||||
3. Processes the sub-workflow's complete event stream after execution
|
||||
4. Forwards outputs to the parent workflow as messages
|
||||
5. Handles external requests by routing them to the parent workflow
|
||||
6. Accumulates responses and resumes sub-workflow execution
|
||||
|
||||
## Event Stream Processing
|
||||
WorkflowExecutor processes events after sub-workflow completion:
|
||||
|
||||
### Output Forwarding
|
||||
All outputs from the sub-workflow are automatically forwarded to the parent:
|
||||
|
||||
#### When `allow_direct_output` is False (default):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# An executor in the sub-workflow yields outputs
|
||||
await ctx.yield_output("sub-workflow result")
|
||||
|
||||
# WorkflowExecutor forwards to parent via ctx.send_message()
|
||||
# Parent receives the output as a regular message
|
||||
|
||||
#### When `allow_direct_output` is True:
|
||||
|
||||
.. code-block:: python
|
||||
# An executor in the sub-workflow yields outputs
|
||||
await ctx.yield_output("sub-workflow result")
|
||||
|
||||
# WorkflowExecutor yields output directly to parent workflow's event stream
|
||||
# The output of the sub-workflow is considered the output of the parent workflow
|
||||
# Caller of the parent workflow receives the output directly
|
||||
|
||||
### Request/Response Coordination
|
||||
When sub-workflows need external information:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# An executor in the sub-workflow makes request
|
||||
request = MyDataRequest(query="user info")
|
||||
|
||||
# WorkflowExecutor captures WorkflowEvent and wraps it in a SubWorkflowRequestMessage
|
||||
# then send it to the receiving executor in parent workflow. The executor in parent workflow
|
||||
# can handle the request locally or forward it to an external source.
|
||||
# The WorkflowExecutor tracks the pending request, and implements a response handler.
|
||||
# When the response is received, it executes the response handler to accumulate responses
|
||||
# and resume the sub-workflow when all expected responses are received.
|
||||
# The response handler expects a SubWorkflowResponseMessage wrapping the response data.
|
||||
|
||||
### State Management
|
||||
WorkflowExecutor maintains execution state across request/response cycles:
|
||||
- Tracks pending requests by request_id
|
||||
- Accumulates responses until all expected responses are received
|
||||
- Resumes sub-workflow execution with complete response batch
|
||||
- Handles concurrent executions and multiple pending requests
|
||||
|
||||
## Type System Integration
|
||||
WorkflowExecutor inherits its type signature from the wrapped workflow:
|
||||
|
||||
### Input Types
|
||||
Matches the wrapped workflow's start executor input types:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# If sub-workflow accepts str, WorkflowExecutor accepts str
|
||||
workflow_executor = WorkflowExecutor(my_workflow, id="wrapper")
|
||||
assert workflow_executor.input_types == my_workflow.input_types
|
||||
|
||||
### Output Types
|
||||
Combines sub-workflow outputs with request coordination types:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Includes all sub-workflow output types
|
||||
# Plus SubWorkflowRequestMessage if sub-workflow can make requests
|
||||
output_types = workflow.output_types + [SubWorkflowRequestMessage] # if applicable
|
||||
|
||||
## Error Handling
|
||||
WorkflowExecutor propagates sub-workflow failures:
|
||||
- Captures failed event (type='failed') from sub-workflow
|
||||
- Converts to error event in parent context
|
||||
- Provides detailed error information including sub-workflow ID
|
||||
|
||||
## Concurrent Execution Support
|
||||
WorkflowExecutor fully supports multiple concurrent sub-workflow executions:
|
||||
|
||||
### Per-Execution State Isolation
|
||||
Each sub-workflow invocation creates an isolated ExecutionContext:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Multiple concurrent invocations are supported
|
||||
workflow_executor = WorkflowExecutor(my_workflow, id="concurrent_executor")
|
||||
|
||||
# Each invocation gets its own execution context
|
||||
# Execution 1: processes input_1 independently
|
||||
# Execution 2: processes input_2 independently
|
||||
# No state interference between executions
|
||||
|
||||
### Request/Response Coordination
|
||||
Responses are correctly routed to the originating execution:
|
||||
- Each execution tracks its own pending requests and expected responses
|
||||
- Request-to-execution mapping ensures responses reach the correct sub-workflow
|
||||
- Response accumulation is isolated per execution
|
||||
- Automatic cleanup when execution completes
|
||||
|
||||
### Memory Management
|
||||
- Unlimited concurrent executions supported
|
||||
- Each execution has unique UUID-based identification
|
||||
- Cleanup of completed execution contexts
|
||||
- Thread-safe state management for concurrent access
|
||||
|
||||
### Important Considerations
|
||||
**Shared Workflow Instance**: All concurrent executions use the same underlying workflow instance.
|
||||
For proper isolation, ensure that the wrapped workflow and its executors are stateless.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Avoid: Stateful executor with instance variables
|
||||
class StatefulExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="stateful")
|
||||
self.data = [] # This will be shared across concurrent executions!
|
||||
|
||||
## Integration with Parent Workflows
|
||||
Parent workflows can intercept sub-workflow requests:
|
||||
|
||||
.. code-block:: python
|
||||
class ParentExecutor(Executor):
|
||||
@handler
|
||||
async def handle_subworkflow_request(
|
||||
self,
|
||||
request: SubWorkflowRequestMessage,
|
||||
ctx: WorkflowContext[SubWorkflowResponseMessage],
|
||||
) -> None:
|
||||
# Handle request locally or forward to external source
|
||||
if self.can_handle_locally(request):
|
||||
# Send response back to sub-workflow
|
||||
response = request.create_response(data="local response data")
|
||||
await ctx.send_message(response, target_id=request.source_executor_id)
|
||||
else:
|
||||
# Forward to external handler
|
||||
await ctx.request_info(request.source_event, response_type=request.source_event.response_type)
|
||||
|
||||
## Implementation Notes
|
||||
- Sub-workflows run to completion before processing their results
|
||||
- Event processing is atomic - all outputs are forwarded before requests
|
||||
- Response accumulation ensures sub-workflows receive complete response batches
|
||||
- Execution state is maintained for proper resumption after external requests
|
||||
- Concurrent executions are fully isolated and do not interfere with each other
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workflow: "Workflow",
|
||||
id: str,
|
||||
allow_direct_output: bool = False,
|
||||
propagate_request: bool = False,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize the WorkflowExecutor.
|
||||
|
||||
Args:
|
||||
workflow: The workflow to execute as a sub-workflow.
|
||||
id: Unique identifier for this executor.
|
||||
allow_direct_output: Whether to allow direct output from the sub-workflow.
|
||||
By default, outputs from the sub-workflow are sent to
|
||||
other executors in the parent workflow as messages.
|
||||
When this is set to true, the outputs are yielded
|
||||
directly from the WorkflowExecutor to the parent
|
||||
workflow's event stream.
|
||||
propagate_request: Whether to propagate requests from the sub-workflow to the
|
||||
parent workflow. If set to true, requests from the sub-workflow
|
||||
will be propagated as the original WorkflowEvent to the parent
|
||||
workflow. Otherwise, they will be wrapped in a SubWorkflowRequestMessage,
|
||||
which should be handled by an executor in the parent workflow.
|
||||
|
||||
Keyword Args:
|
||||
**kwargs: Additional keyword arguments passed to the parent constructor.
|
||||
"""
|
||||
super().__init__(id, **kwargs)
|
||||
self.workflow = workflow
|
||||
self.allow_direct_output = allow_direct_output
|
||||
|
||||
# Track execution contexts for concurrent sub-workflow executions
|
||||
self._execution_contexts: dict[str, ExecutionContext] = {} # execution_id -> ExecutionContext
|
||||
# Map request_id to execution_id for response routing
|
||||
self._request_to_execution: dict[str, str] = {} # request_id -> execution_id
|
||||
self._propagate_request = propagate_request
|
||||
|
||||
@property
|
||||
def input_types(self) -> list[type[Any] | types.UnionType]:
|
||||
"""Get the input types based on the underlying workflow's input types plus WorkflowExecutor-specific types.
|
||||
|
||||
Returns:
|
||||
A list of input types that the WorkflowExecutor can accept.
|
||||
"""
|
||||
input_types: list[type[Any] | types.UnionType] = list(self.workflow.input_types)
|
||||
|
||||
# WorkflowExecutor can also handle SubWorkflowResponseMessage for sub-workflow responses
|
||||
if SubWorkflowResponseMessage not in input_types:
|
||||
input_types.append(SubWorkflowResponseMessage)
|
||||
|
||||
return input_types
|
||||
|
||||
@property
|
||||
def output_types(self) -> list[type[Any] | types.UnionType]:
|
||||
"""Get the output types based on the underlying workflow's output types.
|
||||
|
||||
Returns:
|
||||
A list of output types that the underlying workflow can produce.
|
||||
Includes the SubWorkflowRequestMessage type if any executor in the
|
||||
sub-workflow is request-response capable.
|
||||
"""
|
||||
output_types: list[type[Any] | types.UnionType] = list(self.workflow.output_types)
|
||||
|
||||
is_request_response_capable = any(
|
||||
executor.is_request_response_capable for executor in self.workflow.executors.values()
|
||||
)
|
||||
|
||||
if is_request_response_capable:
|
||||
output_types.append(SubWorkflowRequestMessage)
|
||||
|
||||
return output_types
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
data = super().to_dict()
|
||||
data["workflow"] = self.workflow.to_dict()
|
||||
return data
|
||||
|
||||
def can_handle(self, message: WorkflowMessage) -> bool:
|
||||
"""Override can_handle to only accept messages that the wrapped workflow can handle.
|
||||
|
||||
This prevents the WorkflowExecutor from accepting messages that should go to other
|
||||
executors because the handler `process_workflow` has no type restrictions.
|
||||
"""
|
||||
if isinstance(message.data, SubWorkflowResponseMessage):
|
||||
# Always handle SubWorkflowResponseMessage
|
||||
return True
|
||||
|
||||
if (
|
||||
message.original_request_info_event is not None
|
||||
and message.original_request_info_event.request_id in self._request_to_execution
|
||||
):
|
||||
# Handle propagated responses for known requests
|
||||
return True
|
||||
|
||||
# For other messages, only handle if the wrapped workflow can accept them as input
|
||||
return any(is_instance_of(message.data, input_type) for input_type in self.workflow.input_types)
|
||||
|
||||
@handler
|
||||
async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, Any]) -> None:
|
||||
"""Execute the sub-workflow with raw input data.
|
||||
|
||||
This handler starts a new sub-workflow execution. When the sub-workflow
|
||||
needs external information, it pauses and sends a request to the parent.
|
||||
|
||||
Args:
|
||||
input_data: The input data to send to the sub-workflow.
|
||||
ctx: The workflow context from the parent.
|
||||
"""
|
||||
# Create execution context for this sub-workflow run
|
||||
execution_id = str(uuid.uuid4())
|
||||
execution_context = ExecutionContext(
|
||||
execution_id=execution_id,
|
||||
collected_responses={},
|
||||
expected_response_count=0,
|
||||
pending_requests={},
|
||||
)
|
||||
self._execution_contexts[execution_id] = execution_context
|
||||
|
||||
logger.debug(f"WorkflowExecutor {self.id} starting sub-workflow {self.workflow.id} execution {execution_id}")
|
||||
|
||||
try:
|
||||
# Get kwargs from parent workflow's State to propagate to subworkflow
|
||||
parent_kwargs: dict[str, Any] = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {})
|
||||
|
||||
# Extract invocation kwargs recognised by Workflow.run()
|
||||
# The state stores resolved format (with __global__ wrapper for global kwargs).
|
||||
# Unwrap __global__ before passing to the subworkflow so it gets re-resolved
|
||||
# against the subworkflow's own executor IDs.
|
||||
fi_kwargs: dict[str, Any] | None = None
|
||||
ci_kwargs: dict[str, Any] | None = None
|
||||
for key in ("function_invocation_kwargs", "client_kwargs"):
|
||||
resolved = parent_kwargs.get(key)
|
||||
if isinstance(resolved, dict):
|
||||
# Unwrap global sentinel; pass per-executor dicts as-is
|
||||
unwrapped: dict[str, Any] = resolved.get(GLOBAL_KWARGS_KEY, resolved) # type: ignore
|
||||
if key == "function_invocation_kwargs":
|
||||
fi_kwargs = unwrapped # type: ignore
|
||||
else:
|
||||
ci_kwargs = unwrapped # type: ignore
|
||||
|
||||
# Run the sub-workflow and collect all events, passing parent kwargs
|
||||
result = await self.workflow.run(
|
||||
input_data,
|
||||
function_invocation_kwargs=fi_kwargs, # type: ignore
|
||||
client_kwargs=ci_kwargs, # type: ignore
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"WorkflowExecutor {self.id} sub-workflow {self.workflow.id} "
|
||||
f"execution {execution_id} completed with {len(result)} events"
|
||||
)
|
||||
|
||||
# Process the workflow result using shared logic
|
||||
await self._process_workflow_result(result, execution_context, ctx)
|
||||
finally:
|
||||
# Clean up execution context if it's completed (no pending requests)
|
||||
if execution_id in self._execution_contexts:
|
||||
exec_ctx = self._execution_contexts[execution_id]
|
||||
if not exec_ctx.pending_requests:
|
||||
del self._execution_contexts[execution_id]
|
||||
|
||||
@handler
|
||||
async def handle_message_wrapped_request_response(
|
||||
self,
|
||||
response: SubWorkflowResponseMessage,
|
||||
ctx: WorkflowContext[Any, Any],
|
||||
) -> None:
|
||||
"""Handle response from parent for a forwarded request.
|
||||
|
||||
This handler accumulates responses and only resumes the sub-workflow
|
||||
when all expected responses have been received for that execution.
|
||||
|
||||
Args:
|
||||
response: The response to a previous request.
|
||||
ctx: The workflow context.
|
||||
"""
|
||||
request_id = response.source_event.request_id
|
||||
await self._handle_response(
|
||||
request_id=request_id,
|
||||
response=response.data,
|
||||
ctx=ctx,
|
||||
)
|
||||
|
||||
@response_handler
|
||||
async def handle_propagated_request_response(
|
||||
self,
|
||||
original_request: Any,
|
||||
response: object,
|
||||
ctx: WorkflowContext[Any],
|
||||
) -> None:
|
||||
"""Handle response for a request that was propagated to the parent workflow.
|
||||
|
||||
Args:
|
||||
original_request: The original WorkflowEvent.
|
||||
response: The response data.
|
||||
ctx: The workflow context.
|
||||
"""
|
||||
if ctx.request_id is None:
|
||||
raise RuntimeError("WorkflowExecutor received a propagated response without a request ID in the context.")
|
||||
|
||||
await self._handle_response(
|
||||
request_id=ctx.request_id,
|
||||
response=response,
|
||||
ctx=ctx,
|
||||
)
|
||||
|
||||
@override
|
||||
async def on_checkpoint_save(self) -> dict[str, Any]:
|
||||
"""Get the current state of the WorkflowExecutor for checkpointing purposes."""
|
||||
return {
|
||||
"execution_contexts": {
|
||||
execution_id: execution_context for execution_id, execution_context in self._execution_contexts.items()
|
||||
},
|
||||
"request_to_execution": dict(self._request_to_execution),
|
||||
}
|
||||
|
||||
@override
|
||||
async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
|
||||
"""Restore the WorkflowExecutor state from a checkpoint snapshot."""
|
||||
# Validate the state contains the right keys
|
||||
if "execution_contexts" not in state:
|
||||
raise KeyError("Missing 'execution_contexts' in WorkflowExecutor state.")
|
||||
if "request_to_execution" not in state:
|
||||
raise KeyError("Missing 'request_to_execution' in WorkflowExecutor state.")
|
||||
|
||||
# Validate the execution contexts stored in the state have the right keys and values
|
||||
execution_contexts: dict[str, ExecutionContext] | None = None
|
||||
try:
|
||||
execution_contexts = {
|
||||
key: decode_checkpoint_value(value) for key, value in state["execution_contexts"].items()
|
||||
}
|
||||
except Exception as ex:
|
||||
raise RuntimeError("Failed to deserialize execution context.") from ex
|
||||
|
||||
if not all(
|
||||
isinstance(key, str) and isinstance(value, ExecutionContext) for key, value in execution_contexts.items()
|
||||
):
|
||||
raise ValueError("Execution contexts must have 'str' as key and 'ExecutionContext' as value.")
|
||||
if not all(key == value.execution_id for key, value in execution_contexts.items()):
|
||||
raise ValueError("Execution contexts must have matching keys and IDs.")
|
||||
|
||||
# Validate the request_to_execution map contain the right data
|
||||
request_to_execution = state["request_to_execution"]
|
||||
if not all(isinstance(key, str) and isinstance(value, str) for key, value in request_to_execution.items()):
|
||||
raise ValueError("Request to execution map must have 'str' as key and 'str' as value.")
|
||||
if not all(value in execution_contexts for value in request_to_execution.values()):
|
||||
raise ValueError(
|
||||
"'request_to_execution` contains unknown execution ID that is not part of the execution contexts."
|
||||
)
|
||||
|
||||
self._execution_contexts = execution_contexts
|
||||
self._request_to_execution = request_to_execution
|
||||
|
||||
# Add the `request_info_event`s back to the sub workflow.
|
||||
# This is only a temporary solution to rehydrate the sub workflow with the requests.
|
||||
# The proper way would be to rehydrate the workflow from a checkpoint on a Workflow
|
||||
# API instead of the '_runner_context' object that should be hidden. And the sub workflow
|
||||
# should be rehydrated from a checkpoint object instead of from a subset of the state.
|
||||
# TODO(@taochen): Issue #1614 - how to handle the case when the parent workflow has checkpointing
|
||||
# set up but not the sub workflow?
|
||||
request_info_events = [
|
||||
request_info_event
|
||||
for execution_context in self._execution_contexts.values()
|
||||
for request_info_event in execution_context.pending_requests.values()
|
||||
]
|
||||
await asyncio.gather(*[
|
||||
self.workflow._runner_context.add_request_info_event(event) # pyright: ignore[reportPrivateUsage]
|
||||
for event in request_info_events
|
||||
])
|
||||
|
||||
async def _process_workflow_result(
|
||||
self,
|
||||
result: WorkflowRunResult,
|
||||
execution_context: ExecutionContext,
|
||||
ctx: WorkflowContext[Any],
|
||||
) -> None:
|
||||
"""Process the result from a workflow execution.
|
||||
|
||||
This method handles the common logic for processing outputs, request info events,
|
||||
and final states that is shared between process_workflow and handle_response.
|
||||
|
||||
Args:
|
||||
result: The workflow execution result.
|
||||
execution_context: The execution context for this sub-workflow run.
|
||||
ctx: The workflow context.
|
||||
"""
|
||||
# Collect all events from the workflow
|
||||
request_info_events = result.get_request_info_events()
|
||||
outputs = result.get_outputs()
|
||||
intermediate_outputs = result.get_intermediate_outputs()
|
||||
workflow_run_state = result.get_final_state()
|
||||
logger.debug(
|
||||
f"WorkflowExecutor {self.id} processing workflow result with "
|
||||
f"{len(outputs)} outputs, {len(intermediate_outputs)} intermediate outputs, "
|
||||
f"and {len(request_info_events)} request info events. "
|
||||
f"Workflow run state: {workflow_run_state}"
|
||||
)
|
||||
|
||||
# Process outputs
|
||||
if self.allow_direct_output:
|
||||
# Note that the executor is allowed to continue its own execution after yielding outputs.
|
||||
await asyncio.gather(*[ctx.yield_output(output) for output in outputs])
|
||||
else:
|
||||
await asyncio.gather(*[ctx.send_message(output) for output in outputs])
|
||||
|
||||
# Pipe sub-workflow intermediate emissions up through the parent's event stream.
|
||||
# Bypasses the parent's yield-output classifier so the 'intermediate' label is preserved
|
||||
# across the encapsulation boundary; uses this WorkflowExecutor's id as the source
|
||||
# so outer callers don't need to know the sub-workflow's internal executor layout.
|
||||
if intermediate_outputs:
|
||||
|
||||
async def _forward_intermediate_output(output: Any) -> None:
|
||||
with _framework_event_origin():
|
||||
event = WorkflowEvent("intermediate", executor_id=self.id, data=output)
|
||||
await ctx.add_event(event)
|
||||
|
||||
await asyncio.gather(*[_forward_intermediate_output(output) for output in intermediate_outputs])
|
||||
|
||||
# Process request info events
|
||||
for event in request_info_events:
|
||||
request_id = event.request_id
|
||||
response_type = event.response_type
|
||||
# Track the pending request in execution context
|
||||
execution_context.pending_requests[request_id] = event
|
||||
# Map request to execution for response routing
|
||||
self._request_to_execution[request_id] = execution_context.execution_id
|
||||
if self._propagate_request:
|
||||
# In a workflow where the parent workflow does not handle the request, the request
|
||||
# should be propagated via the `request_info` mechanism to an external source. And
|
||||
# a @response_handler would be required in the WorkflowExecutor to handle the response.
|
||||
await ctx.request_info(event.data, response_type, request_id=request_id)
|
||||
else:
|
||||
# In a workflow where the parent workflow has an executor that may intercept the
|
||||
# request and handle it directly, a message should be sent.
|
||||
await ctx.send_message(SubWorkflowRequestMessage(source_event=event, executor_id=self.id))
|
||||
|
||||
# Update expected response count for this execution
|
||||
execution_context.expected_response_count = len(request_info_events)
|
||||
|
||||
# Handle final state
|
||||
if workflow_run_state == WorkflowRunState.FAILED:
|
||||
# Find the failed event (type='failed').
|
||||
failed_events = [e for e in result if isinstance(e, WorkflowEvent) and e.type == "failed"]
|
||||
if failed_events:
|
||||
failed_event = failed_events[0]
|
||||
if failed_event.details is not None:
|
||||
error_type = failed_event.details.error_type
|
||||
error_message = failed_event.details.message
|
||||
exception = Exception(
|
||||
f"Sub-workflow {self.workflow.id} failed with error: {error_type} - {error_message}"
|
||||
)
|
||||
else:
|
||||
exception = Exception(f"Sub-workflow {self.workflow.id} failed with unknown error")
|
||||
error_event = WorkflowEvent.error(exception)
|
||||
await ctx.add_event(error_event)
|
||||
elif workflow_run_state == WorkflowRunState.IDLE:
|
||||
# Sub-workflow is idle - nothing more to do now
|
||||
logger.debug(
|
||||
f"Sub-workflow {self.workflow.id} is idle with {len(self._execution_contexts)} active executions"
|
||||
)
|
||||
elif workflow_run_state == WorkflowRunState.CANCELLED:
|
||||
# Sub-workflow was cancelled - treat as completion
|
||||
logger.debug(
|
||||
f"Sub-workflow {self.workflow.id} was cancelled with {len(self._execution_contexts)} active executions"
|
||||
)
|
||||
elif workflow_run_state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS:
|
||||
# Sub-workflow is still running with pending requests
|
||||
logger.debug(
|
||||
f"Sub-workflow {self.workflow.id} is still in progress with {len(request_info_events)} "
|
||||
f"pending requests with {len(self._execution_contexts)} active executions"
|
||||
)
|
||||
elif workflow_run_state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
|
||||
# Sub-workflow is idle but has pending requests
|
||||
logger.debug(
|
||||
f"Sub-workflow {self.workflow.id} is idle with pending requests: "
|
||||
f"{len(request_info_events)} with {len(self._execution_contexts)} active executions"
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(f"Unexpected workflow run state: {workflow_run_state}")
|
||||
|
||||
async def _handle_response(
|
||||
self,
|
||||
request_id: str,
|
||||
response: Any,
|
||||
ctx: WorkflowContext[Any],
|
||||
) -> None:
|
||||
execution_id = self._request_to_execution.get(request_id)
|
||||
if not execution_id or execution_id not in self._execution_contexts:
|
||||
logger.warning(
|
||||
f"WorkflowExecutor {self.id} received response for unknown request_id: {request_id}. "
|
||||
"This response will be ignored."
|
||||
)
|
||||
return
|
||||
|
||||
execution_context = self._execution_contexts[execution_id]
|
||||
|
||||
# Check if we have this pending request in the execution context
|
||||
if request_id not in execution_context.pending_requests:
|
||||
logger.warning(
|
||||
f"WorkflowExecutor {self.id} received response for unknown request_id: "
|
||||
f"{request_id} in execution {execution_id}, ignoring"
|
||||
)
|
||||
return
|
||||
|
||||
# Remove the request from pending list and request mapping
|
||||
execution_context.pending_requests.pop(request_id, None)
|
||||
self._request_to_execution.pop(request_id, None)
|
||||
|
||||
# Accumulate the response in this execution's context
|
||||
execution_context.collected_responses[request_id] = response
|
||||
# Check if we have all expected responses for this execution
|
||||
if len(execution_context.collected_responses) < execution_context.expected_response_count:
|
||||
logger.debug(
|
||||
f"WorkflowExecutor {self.id} execution {execution_id} waiting for more responses: "
|
||||
f"{len(execution_context.collected_responses)}/{execution_context.expected_response_count} received"
|
||||
)
|
||||
return # Wait for more responses
|
||||
|
||||
# Send all collected responses to the sub-workflow
|
||||
responses_to_send = dict(execution_context.collected_responses)
|
||||
execution_context.collected_responses.clear() # Clear for next batch
|
||||
|
||||
try:
|
||||
# Resume the sub-workflow with all collected responses
|
||||
result = await self.workflow.run(responses=responses_to_send)
|
||||
# Process the workflow result using shared logic
|
||||
await self._process_workflow_result(result, execution_context, ctx)
|
||||
finally:
|
||||
# Clean up execution context if it's completed (no pending requests)
|
||||
if not execution_context.pending_requests:
|
||||
del self._execution_contexts[execution_id]
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""A2A integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-a2a``
|
||||
|
||||
Supported classes:
|
||||
- A2AAgent
|
||||
- A2AAgentSession
|
||||
- A2AServiceSessionId
|
||||
- A2AExecutor
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
IMPORT_PATH = "agent_framework_a2a"
|
||||
PACKAGE_NAME = "agent-framework-a2a"
|
||||
_IMPORTS = ["A2AAgent", "A2AAgentSession", "A2AServiceSessionId", "A2AExecutor"]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
try:
|
||||
return getattr(importlib.import_module(IMPORT_PATH), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return _IMPORTS
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_a2a import A2AAgent, A2AAgentSession, A2AExecutor, A2AServiceSessionId
|
||||
|
||||
__all__ = ["A2AAgent", "A2AAgentSession", "A2AExecutor", "A2AServiceSessionId"]
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AG-UI integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-ag-ui``
|
||||
|
||||
Supported classes and functions:
|
||||
- AgentFrameworkAgent
|
||||
- AgentFrameworkWorkflow
|
||||
- AGUIChatClient
|
||||
- AGUIEventConverter
|
||||
- AGUIHttpService
|
||||
- AGUIThreadSnapshot
|
||||
- AGUIThreadSnapshotStore
|
||||
- InMemoryAGUIThreadSnapshotStore
|
||||
- SnapshotScopeResolver
|
||||
- add_agent_framework_fastapi_endpoint
|
||||
- state_update
|
||||
- __version__
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
IMPORT_PATH = "agent_framework_ag_ui"
|
||||
PACKAGE_NAME = "agent-framework-ag-ui"
|
||||
_IMPORTS = [
|
||||
"AgentFrameworkAgent",
|
||||
"AgentFrameworkWorkflow",
|
||||
"add_agent_framework_fastapi_endpoint",
|
||||
"AGUIChatClient",
|
||||
"AGUIEventConverter",
|
||||
"AGUIHttpService",
|
||||
"AGUIThreadSnapshot",
|
||||
"AGUIThreadSnapshotStore",
|
||||
"InMemoryAGUIThreadSnapshotStore",
|
||||
"SnapshotScopeResolver",
|
||||
"state_update",
|
||||
"__version__",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
try:
|
||||
return getattr(importlib.import_module(IMPORT_PATH), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return _IMPORTS
|
||||
@@ -0,0 +1,31 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_ag_ui import (
|
||||
AgentFrameworkAgent,
|
||||
AgentFrameworkWorkflow,
|
||||
AGUIChatClient,
|
||||
AGUIEventConverter,
|
||||
AGUIHttpService,
|
||||
AGUIThreadSnapshot,
|
||||
AGUIThreadSnapshotStore,
|
||||
InMemoryAGUIThreadSnapshotStore,
|
||||
SnapshotScopeResolver,
|
||||
__version__,
|
||||
add_agent_framework_fastapi_endpoint,
|
||||
state_update,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AGUIChatClient",
|
||||
"AGUIEventConverter",
|
||||
"AGUIHttpService",
|
||||
"AGUIThreadSnapshot",
|
||||
"AGUIThreadSnapshotStore",
|
||||
"AgentFrameworkAgent",
|
||||
"AgentFrameworkWorkflow",
|
||||
"InMemoryAGUIThreadSnapshotStore",
|
||||
"SnapshotScopeResolver",
|
||||
"__version__",
|
||||
"add_agent_framework_fastapi_endpoint",
|
||||
"state_update",
|
||||
]
|
||||
@@ -0,0 +1,50 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Amazon Bedrock integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-anthropic``
|
||||
- ``agent-framework-bedrock``
|
||||
|
||||
Supported classes:
|
||||
- AnthropicBedrockClient
|
||||
- BedrockChatClient
|
||||
- BedrockChatOptions
|
||||
- BedrockEmbeddingClient
|
||||
- BedrockEmbeddingOptions
|
||||
- BedrockEmbeddingSettings
|
||||
- BedrockGuardrailConfig
|
||||
- BedrockSettings
|
||||
- RawAnthropicBedrockClient
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"AnthropicBedrockClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"BedrockChatClient": ("agent_framework_bedrock", "agent-framework-bedrock"),
|
||||
"BedrockChatOptions": ("agent_framework_bedrock", "agent-framework-bedrock"),
|
||||
"BedrockEmbeddingClient": ("agent_framework_bedrock", "agent-framework-bedrock"),
|
||||
"BedrockEmbeddingOptions": ("agent_framework_bedrock", "agent-framework-bedrock"),
|
||||
"BedrockEmbeddingSettings": ("agent_framework_bedrock", "agent-framework-bedrock"),
|
||||
"BedrockGuardrailConfig": ("agent_framework_bedrock", "agent-framework-bedrock"),
|
||||
"BedrockSettings": ("agent_framework_bedrock", "agent-framework-bedrock"),
|
||||
"RawAnthropicBedrockClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
import_path, package_name = _IMPORTS[name]
|
||||
try:
|
||||
return getattr(importlib.import_module(import_path), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{package_name}' package is not installed, please do `pip install {package_name}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module `amazon` has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_IMPORTS.keys())
|
||||
@@ -0,0 +1,24 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_anthropic import AnthropicBedrockClient, RawAnthropicBedrockClient
|
||||
from agent_framework_bedrock import (
|
||||
BedrockChatClient,
|
||||
BedrockChatOptions,
|
||||
BedrockEmbeddingClient,
|
||||
BedrockEmbeddingOptions,
|
||||
BedrockEmbeddingSettings,
|
||||
BedrockGuardrailConfig,
|
||||
BedrockSettings,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AnthropicBedrockClient",
|
||||
"BedrockChatClient",
|
||||
"BedrockChatOptions",
|
||||
"BedrockEmbeddingClient",
|
||||
"BedrockEmbeddingOptions",
|
||||
"BedrockEmbeddingSettings",
|
||||
"BedrockGuardrailConfig",
|
||||
"BedrockSettings",
|
||||
"RawAnthropicBedrockClient",
|
||||
]
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Anthropic integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-anthropic``
|
||||
- ``agent-framework-claude``
|
||||
|
||||
Supported classes:
|
||||
- AnthropicBedrockClient
|
||||
- AnthropicClient
|
||||
- AnthropicChatOptions
|
||||
- AnthropicFoundryClient
|
||||
- AnthropicVertexClient
|
||||
- ClaudeAgent
|
||||
- ClaudeAgentOptions
|
||||
- RawAnthropicBedrockClient
|
||||
- RawAnthropicClient
|
||||
- RawAnthropicFoundryClient
|
||||
- RawClaudeAgent
|
||||
- RawAnthropicVertexClient
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"AnthropicBedrockClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"AnthropicClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"AnthropicChatOptions": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"AnthropicFoundryClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"AnthropicVertexClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"ClaudeAgent": ("agent_framework_claude", "agent-framework-claude"),
|
||||
"ClaudeAgentOptions": ("agent_framework_claude", "agent-framework-claude"),
|
||||
"RawAnthropicBedrockClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"RawAnthropicClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"RawAnthropicFoundryClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"RawClaudeAgent": ("agent_framework_claude", "agent-framework-claude"),
|
||||
"RawAnthropicVertexClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
import_path, package_name = _IMPORTS[name]
|
||||
try:
|
||||
return getattr(importlib.import_module(import_path), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{package_name}' package is not installed, please do `pip install {package_name}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module `anthropic` has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_IMPORTS.keys())
|
||||
@@ -0,0 +1,28 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_anthropic import (
|
||||
AnthropicBedrockClient,
|
||||
AnthropicChatOptions,
|
||||
AnthropicClient,
|
||||
AnthropicFoundryClient,
|
||||
AnthropicVertexClient,
|
||||
RawAnthropicBedrockClient,
|
||||
RawAnthropicClient,
|
||||
RawAnthropicFoundryClient,
|
||||
RawAnthropicVertexClient,
|
||||
)
|
||||
from agent_framework_claude import ClaudeAgent, ClaudeAgentOptions
|
||||
|
||||
__all__ = [
|
||||
"AnthropicBedrockClient",
|
||||
"AnthropicChatOptions",
|
||||
"AnthropicClient",
|
||||
"AnthropicFoundryClient",
|
||||
"AnthropicVertexClient",
|
||||
"ClaudeAgent",
|
||||
"ClaudeAgentOptions",
|
||||
"RawAnthropicBedrockClient",
|
||||
"RawAnthropicClient",
|
||||
"RawAnthropicFoundryClient",
|
||||
"RawAnthropicVertexClient",
|
||||
]
|
||||
@@ -0,0 +1,40 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Azure integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from optional Azure connector packages.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"AgentCallbackContext": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
"AgentFunctionApp": ("agent_framework_azurefunctions", "agent-framework-azurefunctions"),
|
||||
"AgentResponseCallbackProtocol": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
"AzureAISearchContextProvider": ("agent_framework_azure_ai_search", "agent-framework-azure-ai-search"),
|
||||
"AzureAISearchSettings": ("agent_framework_azure_ai_search", "agent-framework-azure-ai-search"),
|
||||
"CosmosHistoryProvider": ("agent_framework_azure_cosmos", "agent-framework-azure-cosmos"),
|
||||
"DurableAIAgent": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
"DurableAIAgentClient": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
"DurableAIAgentOrchestrationContext": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
"DurableAIAgentWorker": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
"DurableWorkflowClient": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
import_path, package_name = _IMPORTS[name]
|
||||
try:
|
||||
return getattr(importlib.import_module(import_path), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The package {package_name} is required to use `{name}`. "
|
||||
f"Please use `pip install {package_name}`, or update your requirements.txt or pyproject.toml file."
|
||||
) from exc
|
||||
raise AttributeError(f"Module `azure` has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_IMPORTS.keys())
|
||||
@@ -0,0 +1,34 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
# Type stubs for the agent_framework.azure lazy-loading namespace.
|
||||
# Install the relevant packages for full type support.
|
||||
|
||||
from agent_framework_azure_ai_search import (
|
||||
AzureAISearchContextProvider,
|
||||
AzureAISearchSettings,
|
||||
)
|
||||
from agent_framework_azure_cosmos import CosmosHistoryProvider
|
||||
from agent_framework_azurefunctions import AgentFunctionApp
|
||||
from agent_framework_durabletask import (
|
||||
AgentCallbackContext,
|
||||
AgentResponseCallbackProtocol,
|
||||
DurableAIAgent,
|
||||
DurableAIAgentClient,
|
||||
DurableAIAgentOrchestrationContext,
|
||||
DurableAIAgentWorker,
|
||||
DurableWorkflowClient,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AgentCallbackContext",
|
||||
"AgentFunctionApp",
|
||||
"AgentResponseCallbackProtocol",
|
||||
"AzureAISearchContextProvider",
|
||||
"AzureAISearchSettings",
|
||||
"CosmosHistoryProvider",
|
||||
"DurableAIAgent",
|
||||
"DurableAIAgentClient",
|
||||
"DurableAIAgentOrchestrationContext",
|
||||
"DurableAIAgentWorker",
|
||||
"DurableWorkflowClient",
|
||||
]
|
||||
@@ -0,0 +1,34 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""ChatKit integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-chatkit``
|
||||
|
||||
Supported classes and functions:
|
||||
- ThreadItemConverter
|
||||
- simple_to_agent_input
|
||||
- stream_agent_response
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
IMPORT_PATH = "agent_framework_chatkit"
|
||||
PACKAGE_NAME = "agent-framework-chatkit"
|
||||
_IMPORTS = ["ThreadItemConverter", "simple_to_agent_input", "stream_agent_response"]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
try:
|
||||
return getattr(importlib.import_module(IMPORT_PATH), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return _IMPORTS
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_chatkit import (
|
||||
ThreadItemConverter,
|
||||
simple_to_agent_input,
|
||||
stream_agent_response,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ThreadItemConverter",
|
||||
"simple_to_agent_input",
|
||||
"stream_agent_response",
|
||||
]
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Declarative integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-declarative``
|
||||
|
||||
Supported classes include:
|
||||
- AgentFactory
|
||||
- WorkflowFactory
|
||||
- ExternalInputRequest
|
||||
- ExternalInputResponse
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
IMPORT_PATH = "agent_framework_declarative"
|
||||
PACKAGE_NAME = "agent-framework-declarative"
|
||||
_IMPORTS = [
|
||||
"AgentFactory",
|
||||
"AgentExternalInputRequest",
|
||||
"AgentExternalInputResponse",
|
||||
"DeclarativeActionError",
|
||||
"DeclarativeLoaderError",
|
||||
"DeclarativeWorkflowError",
|
||||
"DefaultHttpRequestHandler",
|
||||
"DefaultMCPToolHandler",
|
||||
"ExternalInputRequest",
|
||||
"ExternalInputResponse",
|
||||
"HttpRequestHandler",
|
||||
"HttpRequestInfo",
|
||||
"HttpRequestResult",
|
||||
"MCPToolApprovalRequest",
|
||||
"MCPToolHandler",
|
||||
"MCPToolInvocation",
|
||||
"MCPToolResult",
|
||||
"ProviderLookupError",
|
||||
"ProviderTypeMapping",
|
||||
"ToolApprovalRequest",
|
||||
"ToolApprovalResponse",
|
||||
"WorkflowFactory",
|
||||
"WorkflowState",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
try:
|
||||
return getattr(importlib.import_module(IMPORT_PATH), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return _IMPORTS
|
||||
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_declarative import (
|
||||
AgentExternalInputRequest,
|
||||
AgentExternalInputResponse,
|
||||
AgentFactory,
|
||||
DeclarativeActionError,
|
||||
DeclarativeLoaderError,
|
||||
DeclarativeWorkflowError,
|
||||
DefaultHttpRequestHandler,
|
||||
DefaultMCPToolHandler,
|
||||
ExternalInputRequest,
|
||||
ExternalInputResponse,
|
||||
HttpRequestHandler,
|
||||
HttpRequestInfo,
|
||||
HttpRequestResult,
|
||||
MCPToolApprovalRequest,
|
||||
MCPToolHandler,
|
||||
MCPToolInvocation,
|
||||
MCPToolResult,
|
||||
ProviderLookupError,
|
||||
ProviderTypeMapping,
|
||||
ToolApprovalRequest,
|
||||
ToolApprovalResponse,
|
||||
WorkflowFactory,
|
||||
WorkflowState,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AgentExternalInputRequest",
|
||||
"AgentExternalInputResponse",
|
||||
"AgentFactory",
|
||||
"DeclarativeActionError",
|
||||
"DeclarativeLoaderError",
|
||||
"DeclarativeWorkflowError",
|
||||
"DefaultHttpRequestHandler",
|
||||
"DefaultMCPToolHandler",
|
||||
"ExternalInputRequest",
|
||||
"ExternalInputResponse",
|
||||
"HttpRequestHandler",
|
||||
"HttpRequestInfo",
|
||||
"HttpRequestResult",
|
||||
"MCPToolApprovalRequest",
|
||||
"MCPToolHandler",
|
||||
"MCPToolInvocation",
|
||||
"MCPToolResult",
|
||||
"ProviderLookupError",
|
||||
"ProviderTypeMapping",
|
||||
"ToolApprovalRequest",
|
||||
"ToolApprovalResponse",
|
||||
"WorkflowFactory",
|
||||
"WorkflowState",
|
||||
]
|
||||
@@ -0,0 +1,48 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""DevUI integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-devui``
|
||||
|
||||
Supported classes and functions include:
|
||||
- DevServer
|
||||
- AgentFrameworkRequest
|
||||
- DiscoveryResponse
|
||||
- ResponseStreamEvent
|
||||
- serve
|
||||
- main
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
IMPORT_PATH = "agent_framework_devui"
|
||||
PACKAGE_NAME = "agent-framework-devui"
|
||||
_IMPORTS = [
|
||||
"AgentFrameworkRequest",
|
||||
"DevServer",
|
||||
"DiscoveryResponse",
|
||||
"EntityInfo",
|
||||
"OpenAIError",
|
||||
"OpenAIResponse",
|
||||
"ResponseStreamEvent",
|
||||
"main",
|
||||
"register_cleanup",
|
||||
"serve",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
try:
|
||||
return getattr(importlib.import_module(IMPORT_PATH), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return _IMPORTS
|
||||
@@ -0,0 +1,27 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_devui import (
|
||||
AgentFrameworkRequest,
|
||||
DevServer,
|
||||
DiscoveryResponse,
|
||||
EntityInfo,
|
||||
OpenAIError,
|
||||
OpenAIResponse,
|
||||
ResponseStreamEvent,
|
||||
main,
|
||||
register_cleanup,
|
||||
serve,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AgentFrameworkRequest",
|
||||
"DevServer",
|
||||
"DiscoveryResponse",
|
||||
"EntityInfo",
|
||||
"OpenAIError",
|
||||
"OpenAIResponse",
|
||||
"ResponseStreamEvent",
|
||||
"main",
|
||||
"register_cleanup",
|
||||
"serve",
|
||||
]
|
||||
@@ -0,0 +1,263 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Exception hierarchy used across Agent Framework core and connectors.
|
||||
|
||||
See python/CODING_STANDARD.md § Exception Hierarchy for design rationale
|
||||
and guidance on choosing the correct exception class.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Literal
|
||||
|
||||
logger = logging.getLogger("agent_framework")
|
||||
|
||||
|
||||
class AgentFrameworkException(Exception):
|
||||
"""Base exception for the Agent Framework.
|
||||
|
||||
Automatically logs the message as debug.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
inner_exception: Exception | None = None,
|
||||
log_level: Literal[0] | Literal[10] | Literal[20] | Literal[30] | Literal[40] | Literal[50] | None = 10,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Create an AgentFrameworkException.
|
||||
|
||||
This emits a debug log (by default), with the inner_exception if provided.
|
||||
"""
|
||||
if log_level is not None:
|
||||
logger.log(log_level, message, exc_info=inner_exception)
|
||||
if inner_exception:
|
||||
super().__init__(message, inner_exception, *args)
|
||||
else:
|
||||
super().__init__(message, *args)
|
||||
|
||||
|
||||
# region Agent Exceptions
|
||||
|
||||
|
||||
class AgentException(AgentFrameworkException):
|
||||
"""Base class for all agent exceptions."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AgentInvalidAuthException(AgentException):
|
||||
"""An authentication error occurred in an agent."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AgentInvalidRequestException(AgentException):
|
||||
"""An invalid request was made to an agent."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AgentInvalidResponseException(AgentException):
|
||||
"""An invalid or unexpected response was received from an agent."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AgentContentFilterException(AgentException):
|
||||
"""A content filter was triggered by an agent."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Chat Client Exceptions
|
||||
|
||||
|
||||
class ChatClientException(AgentFrameworkException):
|
||||
"""Base class for all chat client exceptions."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ChatClientInvalidAuthException(ChatClientException):
|
||||
"""An authentication error occurred in a chat client."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ChatClientInvalidRequestException(ChatClientException):
|
||||
"""An invalid request was made to a chat client."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ChatClientInvalidResponseException(ChatClientException):
|
||||
"""An invalid or unexpected response was received from a chat client."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ChatClientContentFilterException(ChatClientException):
|
||||
"""A content filter was triggered by a chat client."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Integration Exceptions
|
||||
|
||||
|
||||
class IntegrationException(AgentFrameworkException):
|
||||
"""Base class for all external service/dependency integration exceptions."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class IntegrationInitializationError(IntegrationException):
|
||||
"""A wrapped dependency/service lifecycle failure occurred during setup."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class IntegrationInvalidAuthException(IntegrationException):
|
||||
"""An authentication error occurred in an external integration."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class IntegrationInvalidRequestException(IntegrationException):
|
||||
"""An invalid request was made to an external integration."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class IntegrationInvalidResponseException(IntegrationException):
|
||||
"""An invalid or unexpected response was received from an external integration."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class IntegrationContentFilterException(IntegrationException):
|
||||
"""A content filter was triggered by an external integration."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Content Exceptions
|
||||
|
||||
|
||||
class ContentError(AgentFrameworkException):
|
||||
"""An error occurred while processing content."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AdditionItemMismatch(ContentError):
|
||||
"""A type mismatch occurred while merging content items."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Tool Exceptions
|
||||
|
||||
|
||||
class ToolException(AgentFrameworkException):
|
||||
"""Base class for all tool-related exceptions."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ToolExecutionException(ToolException):
|
||||
"""A tool or prompt call failed at runtime."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class UserInputRequiredException(ToolException):
|
||||
"""Raised when a tool wrapping a sub-agent requires user input to proceed.
|
||||
|
||||
This exception carries the ``user_input_request`` Content items emitted by
|
||||
the sub-agent (e.g., ``oauth_consent_request``, ``function_approval_request``)
|
||||
so the tool invocation layer can propagate them to the parent agent's response
|
||||
instead of swallowing them as a generic tool error.
|
||||
|
||||
Args:
|
||||
contents: The user-input-request Content items from the sub-agent response.
|
||||
message: Human-readable description of why user input is needed.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
contents: list[Any],
|
||||
message: str = "Tool requires user input to proceed.",
|
||||
) -> None:
|
||||
"""Create a UserInputRequiredException.
|
||||
|
||||
Args:
|
||||
contents: The user-input-request Content items from the sub-agent response.
|
||||
message: Human-readable description of why user input is needed.
|
||||
"""
|
||||
super().__init__(message, log_level=None)
|
||||
self.contents = contents
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Middleware Exceptions
|
||||
|
||||
|
||||
class MiddlewareException(AgentFrameworkException):
|
||||
"""An error occurred during middleware execution."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Settings Exceptions
|
||||
|
||||
|
||||
class SettingNotFoundError(AgentFrameworkException):
|
||||
"""A required setting could not be resolved from any source."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Workflow Exceptions
|
||||
|
||||
|
||||
class WorkflowException(AgentFrameworkException):
|
||||
"""Base exception for workflow errors."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WorkflowRunnerException(WorkflowException):
|
||||
"""Base exception for workflow runner errors."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WorkflowConvergenceException(WorkflowRunnerException):
|
||||
"""Exception raised when a workflow runner fails to converge within the maximum iterations."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WorkflowCheckpointException(WorkflowRunnerException):
|
||||
"""Exception raised for errors related to workflow checkpoints."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
# endregion
|
||||
@@ -0,0 +1,63 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Foundry integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-anthropic``
|
||||
- ``agent-framework-azure-contentunderstanding``
|
||||
- ``agent-framework-foundry``
|
||||
- ``agent-framework-foundry-local``
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"AnalysisSection": ("agent_framework_azure_contentunderstanding", "agent-framework-azure-contentunderstanding"),
|
||||
"AnthropicFoundryClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"ContentUnderstandingContextProvider": (
|
||||
"agent_framework_azure_contentunderstanding",
|
||||
"agent-framework-azure-contentunderstanding",
|
||||
),
|
||||
"DocumentStatus": ("agent_framework_azure_contentunderstanding", "agent-framework-azure-contentunderstanding"),
|
||||
"FileSearchBackend": ("agent_framework_azure_contentunderstanding", "agent-framework-azure-contentunderstanding"),
|
||||
"FileSearchConfig": ("agent_framework_azure_contentunderstanding", "agent-framework-azure-contentunderstanding"),
|
||||
"FoundryAgent": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"FoundryAgentOptions": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"FoundryChatClient": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"FoundryChatOptions": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"FoundryEmbeddingClient": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"FoundryEmbeddingOptions": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"FoundryEmbeddingSettings": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"FoundryEvals": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"FoundryMemoryProvider": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"FoundryLocalChatOptions": ("agent_framework_foundry_local", "agent-framework-foundry-local"),
|
||||
"FoundryLocalClient": ("agent_framework_foundry_local", "agent-framework-foundry-local"),
|
||||
"FoundryLocalSettings": ("agent_framework_foundry_local", "agent-framework-foundry-local"),
|
||||
"GeneratedEvaluatorRef": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"RawAnthropicFoundryClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"RawFoundryAgent": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"RawFoundryAgentChatClient": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"RawFoundryChatClient": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"RawFoundryEmbeddingClient": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"evaluate_foundry_target": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"evaluate_traces": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
"to_prompt_agent": ("agent_framework_foundry", "agent-framework-foundry"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
import_path, package_name = _IMPORTS[name]
|
||||
try:
|
||||
return getattr(importlib.import_module(import_path), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The package {package_name} is required to use `{name}`. "
|
||||
f"Please use `pip install {package_name}`, or update your requirements.txt or pyproject.toml file."
|
||||
) from exc
|
||||
raise AttributeError(f"Module `foundry` has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_IMPORTS.keys())
|
||||
@@ -0,0 +1,65 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
# Type stubs for the agent_framework.foundry lazy-loading namespace.
|
||||
# Install the relevant packages for full type support.
|
||||
|
||||
from agent_framework_anthropic import AnthropicFoundryClient, RawAnthropicFoundryClient
|
||||
from agent_framework_azure_contentunderstanding import (
|
||||
AnalysisSection,
|
||||
ContentUnderstandingContextProvider,
|
||||
DocumentStatus,
|
||||
FileSearchBackend,
|
||||
FileSearchConfig,
|
||||
)
|
||||
from agent_framework_foundry import (
|
||||
FoundryAgent,
|
||||
FoundryChatClient,
|
||||
FoundryChatOptions,
|
||||
FoundryEmbeddingClient,
|
||||
FoundryEmbeddingOptions,
|
||||
FoundryEmbeddingSettings,
|
||||
FoundryEvals,
|
||||
FoundryMemoryProvider,
|
||||
GeneratedEvaluatorRef,
|
||||
RawFoundryAgent,
|
||||
RawFoundryAgentChatClient,
|
||||
RawFoundryChatClient,
|
||||
RawFoundryEmbeddingClient,
|
||||
evaluate_foundry_target,
|
||||
evaluate_traces,
|
||||
to_prompt_agent,
|
||||
)
|
||||
from agent_framework_foundry_local import (
|
||||
FoundryLocalChatOptions,
|
||||
FoundryLocalClient,
|
||||
FoundryLocalSettings,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AnalysisSection",
|
||||
"AnthropicFoundryClient",
|
||||
"ContentUnderstandingContextProvider",
|
||||
"DocumentStatus",
|
||||
"FileSearchBackend",
|
||||
"FileSearchConfig",
|
||||
"FoundryAgent",
|
||||
"FoundryChatClient",
|
||||
"FoundryChatOptions",
|
||||
"FoundryEmbeddingClient",
|
||||
"FoundryEmbeddingOptions",
|
||||
"FoundryEmbeddingSettings",
|
||||
"FoundryEvals",
|
||||
"FoundryLocalChatOptions",
|
||||
"FoundryLocalClient",
|
||||
"FoundryLocalSettings",
|
||||
"FoundryMemoryProvider",
|
||||
"GeneratedEvaluatorRef",
|
||||
"RawAnthropicFoundryClient",
|
||||
"RawFoundryAgent",
|
||||
"RawFoundryAgentChatClient",
|
||||
"RawFoundryChatClient",
|
||||
"RawFoundryEmbeddingClient",
|
||||
"evaluate_foundry_target",
|
||||
"evaluate_traces",
|
||||
"to_prompt_agent",
|
||||
]
|
||||
@@ -0,0 +1,40 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""GitHub integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-github-copilot``
|
||||
|
||||
Supported classes:
|
||||
- GitHubCopilotAgent
|
||||
- GitHubCopilotOptions
|
||||
- GitHubCopilotSettings
|
||||
- RawGitHubCopilotAgent
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"GitHubCopilotAgent": ("agent_framework_github_copilot", "agent-framework-github-copilot"),
|
||||
"GitHubCopilotOptions": ("agent_framework_github_copilot", "agent-framework-github-copilot"),
|
||||
"GitHubCopilotSettings": ("agent_framework_github_copilot", "agent-framework-github-copilot"),
|
||||
"RawGitHubCopilotAgent": ("agent_framework_github_copilot", "agent-framework-github-copilot"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
import_path, package_name = _IMPORTS[name]
|
||||
try:
|
||||
return getattr(importlib.import_module(import_path), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The package {package_name} is required to use `{name}`. "
|
||||
f"Please use `pip install {package_name}`, or update your requirements.txt or pyproject.toml file."
|
||||
) from exc
|
||||
raise AttributeError(f"Module `agent_framework.github` has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_IMPORTS.keys())
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_github_copilot import (
|
||||
GitHubCopilotAgent,
|
||||
GitHubCopilotOptions,
|
||||
GitHubCopilotSettings,
|
||||
RawGitHubCopilotAgent,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"GitHubCopilotAgent",
|
||||
"GitHubCopilotOptions",
|
||||
"GitHubCopilotSettings",
|
||||
"RawGitHubCopilotAgent",
|
||||
]
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Google integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports Google-hosted Anthropic clients from:
|
||||
- ``agent-framework-anthropic``
|
||||
|
||||
Supported classes:
|
||||
- AnthropicVertexClient
|
||||
- RawAnthropicVertexClient
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"AnthropicVertexClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
"RawAnthropicVertexClient": ("agent_framework_anthropic", "agent-framework-anthropic"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
import_path, package_name = _IMPORTS[name]
|
||||
try:
|
||||
return getattr(importlib.import_module(import_path), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{package_name}' package is not installed, please do `pip install {package_name}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module `google` has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_IMPORTS.keys())
|
||||
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_anthropic import AnthropicVertexClient, RawAnthropicVertexClient
|
||||
|
||||
__all__ = [
|
||||
"AnthropicVertexClient",
|
||||
"RawAnthropicVertexClient",
|
||||
]
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Hyperlight CodeAct namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from ``agent-framework-hyperlight``.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"AllowedDomain": ("agent_framework_hyperlight", "agent-framework-hyperlight"),
|
||||
"AllowedDomainInput": ("agent_framework_hyperlight", "agent-framework-hyperlight"),
|
||||
"FileMount": ("agent_framework_hyperlight", "agent-framework-hyperlight"),
|
||||
"FileMountInput": ("agent_framework_hyperlight", "agent-framework-hyperlight"),
|
||||
"HyperlightCodeActProvider": ("agent_framework_hyperlight", "agent-framework-hyperlight"),
|
||||
"HyperlightExecuteCodeTool": ("agent_framework_hyperlight", "agent-framework-hyperlight"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
import_path, package_name = _IMPORTS[name]
|
||||
try:
|
||||
return getattr(importlib.import_module(import_path), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The package {package_name} is required to use `{name}`. "
|
||||
f"Please use `pip install {package_name}`, or update your requirements.txt or pyproject.toml file."
|
||||
) from exc
|
||||
raise AttributeError(f"Module `hyperlight` has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_IMPORTS.keys())
|
||||
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Lab namespace package for experimental Agent Framework integrations.
|
||||
|
||||
This module extends the package path so experimental lab integrations can be
|
||||
distributed in separate packages under the ``agent_framework.lab`` namespace.
|
||||
"""
|
||||
|
||||
# This makes agent_framework.lab a namespace package
|
||||
__path__ = __import__("pkgutil").extend_path(__path__, __name__)
|
||||
@@ -0,0 +1,32 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Mem0 integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-mem0``
|
||||
|
||||
Supported classes:
|
||||
- Mem0ContextProvider
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
IMPORT_PATH = "agent_framework_mem0"
|
||||
PACKAGE_NAME = "agent-framework-mem0"
|
||||
_IMPORTS = ["Mem0ContextProvider"]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
try:
|
||||
return getattr(importlib.import_module(IMPORT_PATH), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return _IMPORTS
|
||||
@@ -0,0 +1,9 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_mem0 import (
|
||||
Mem0ContextProvider,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Mem0ContextProvider",
|
||||
]
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Microsoft integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-copilotstudio``
|
||||
- ``agent-framework-purview``
|
||||
|
||||
Supported classes:
|
||||
- CopilotStudioAgent
|
||||
- PurviewPolicyMiddleware
|
||||
- PurviewChatPolicyMiddleware
|
||||
- PurviewSettings
|
||||
- PurviewAppLocation
|
||||
- PurviewLocationType
|
||||
- PurviewAuthenticationError
|
||||
- PurviewPaymentRequiredError
|
||||
- PurviewRateLimitError
|
||||
- PurviewRequestError
|
||||
- PurviewServiceError
|
||||
- CacheProvider
|
||||
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"CopilotStudioAgent": ("agent_framework_copilotstudio", "agent-framework-copilotstudio"),
|
||||
"acquire_token": ("agent_framework_copilotstudio", "agent-framework-copilotstudio"),
|
||||
"PurviewPolicyMiddleware": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"PurviewChatPolicyMiddleware": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"PurviewSettings": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"PurviewAppLocation": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"PurviewLocationType": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"PurviewAuthenticationError": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"PurviewPaymentRequiredError": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"PurviewRateLimitError": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"PurviewRequestError": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"PurviewServiceError": ("agent_framework_purview", "agent-framework-purview"),
|
||||
"CacheProvider": ("agent_framework_purview", "agent-framework-purview"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
import_path, package_name = _IMPORTS[name]
|
||||
try:
|
||||
return getattr(importlib.import_module(import_path), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The package {package_name} is required to use `{name}`. "
|
||||
f"Please use `pip install {package_name}`, or update your requirements.txt or pyproject.toml file."
|
||||
) from exc
|
||||
raise AttributeError(f"Module `microsoft` has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_IMPORTS.keys())
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_copilotstudio import (
|
||||
CopilotStudioAgent,
|
||||
acquire_token,
|
||||
)
|
||||
from agent_framework_purview import (
|
||||
CacheProvider,
|
||||
PurviewAppLocation,
|
||||
PurviewAuthenticationError,
|
||||
PurviewChatPolicyMiddleware,
|
||||
PurviewLocationType,
|
||||
PurviewPaymentRequiredError,
|
||||
PurviewPolicyMiddleware,
|
||||
PurviewRateLimitError,
|
||||
PurviewRequestError,
|
||||
PurviewServiceError,
|
||||
PurviewSettings,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CacheProvider",
|
||||
"CopilotStudioAgent",
|
||||
"PurviewAppLocation",
|
||||
"PurviewAuthenticationError",
|
||||
"PurviewChatPolicyMiddleware",
|
||||
"PurviewLocationType",
|
||||
"PurviewPaymentRequiredError",
|
||||
"PurviewPolicyMiddleware",
|
||||
"PurviewRateLimitError",
|
||||
"PurviewRequestError",
|
||||
"PurviewServiceError",
|
||||
"PurviewSettings",
|
||||
"acquire_token",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,44 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Ollama integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-ollama``
|
||||
|
||||
Supported classes:
|
||||
- OllamaChatClient
|
||||
- OllamaChatOptions
|
||||
- OllamaEmbeddingClient
|
||||
- OllamaEmbeddingOptions
|
||||
- OllamaEmbeddingSettings
|
||||
- OllamaSettings
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
IMPORT_PATH = "agent_framework_ollama"
|
||||
PACKAGE_NAME = "agent-framework-ollama"
|
||||
_IMPORTS = [
|
||||
"OllamaChatClient",
|
||||
"OllamaChatOptions",
|
||||
"OllamaEmbeddingClient",
|
||||
"OllamaEmbeddingOptions",
|
||||
"OllamaEmbeddingSettings",
|
||||
"OllamaSettings",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
try:
|
||||
return getattr(importlib.import_module(IMPORT_PATH), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return _IMPORTS
|
||||
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_ollama import (
|
||||
OllamaChatClient,
|
||||
OllamaChatOptions,
|
||||
OllamaEmbeddingClient,
|
||||
OllamaEmbeddingOptions,
|
||||
OllamaEmbeddingSettings,
|
||||
OllamaSettings,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"OllamaChatClient",
|
||||
"OllamaChatOptions",
|
||||
"OllamaEmbeddingClient",
|
||||
"OllamaEmbeddingOptions",
|
||||
"OllamaEmbeddingSettings",
|
||||
"OllamaSettings",
|
||||
]
|
||||
@@ -0,0 +1,47 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""OpenAI namespace for Agent Framework clients.
|
||||
|
||||
This module lazily re-exports objects from the ``agent-framework-openai`` package.
|
||||
Install it with: ``pip install agent-framework-openai``
|
||||
|
||||
Supported classes include:
|
||||
- OpenAIChatClient (Responses API)
|
||||
- OpenAIChatCompletionClient (Chat Completions API)
|
||||
- OpenAIEmbeddingClient
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"OpenAIChatClient": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"OpenAIChatOptions": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"OpenAIContinuationToken": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"RawOpenAIChatClient": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"OpenAIChatCompletionClient": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"OpenAIChatCompletionOptions": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"RawOpenAIChatCompletionClient": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"OpenAIEmbeddingClient": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"OpenAIEmbeddingOptions": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"OpenAISettings": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"ContentFilterResultSeverity": ("agent_framework_openai", "agent-framework-openai"),
|
||||
"OpenAIContentFilterException": ("agent_framework_openai", "agent-framework-openai"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
import_path, package_name = _IMPORTS[name]
|
||||
try:
|
||||
return getattr(importlib.import_module(import_path), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The package {package_name} is required to use `{name}`. "
|
||||
f"Please use `pip install {package_name}`, or update your requirements.txt or pyproject.toml file."
|
||||
) from exc
|
||||
raise AttributeError(f"Module `openai` has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_IMPORTS.keys())
|
||||
@@ -0,0 +1,34 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
# Type stubs for the agent_framework.openai lazy-loading namespace.
|
||||
# Install agent-framework-openai for full type support.
|
||||
|
||||
from agent_framework_openai import (
|
||||
ContentFilterResultSeverity,
|
||||
OpenAIChatClient,
|
||||
OpenAIChatCompletionClient,
|
||||
OpenAIChatCompletionOptions,
|
||||
OpenAIChatOptions,
|
||||
OpenAIContentFilterException,
|
||||
OpenAIContinuationToken,
|
||||
OpenAIEmbeddingClient,
|
||||
OpenAIEmbeddingOptions,
|
||||
OpenAISettings,
|
||||
RawOpenAIChatClient,
|
||||
RawOpenAIChatCompletionClient,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ContentFilterResultSeverity",
|
||||
"OpenAIChatClient",
|
||||
"OpenAIChatCompletionClient",
|
||||
"OpenAIChatCompletionOptions",
|
||||
"OpenAIChatOptions",
|
||||
"OpenAIContentFilterException",
|
||||
"OpenAIContinuationToken",
|
||||
"OpenAIEmbeddingClient",
|
||||
"OpenAIEmbeddingOptions",
|
||||
"OpenAISettings",
|
||||
"RawOpenAIChatClient",
|
||||
"RawOpenAIChatCompletionClient",
|
||||
]
|
||||
@@ -0,0 +1,84 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Orchestrations integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-orchestrations``
|
||||
|
||||
Supported classes include:
|
||||
- SequentialBuilder
|
||||
- ConcurrentBuilder
|
||||
- GroupChatBuilder
|
||||
- MagenticBuilder
|
||||
- HandoffBuilder
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
IMPORT_PATH = "agent_framework_orchestrations"
|
||||
PACKAGE_NAME = "agent-framework-orchestrations"
|
||||
_IMPORTS = [
|
||||
# Sequential
|
||||
"SequentialBuilder",
|
||||
# Concurrent
|
||||
"ConcurrentBuilder",
|
||||
# Handoff
|
||||
"HandoffAgentExecutor",
|
||||
"HandoffAgentUserRequest",
|
||||
"HandoffBuilder",
|
||||
"HandoffConfiguration",
|
||||
"HandoffSentEvent",
|
||||
# Base orchestrator
|
||||
"BaseGroupChatOrchestrator",
|
||||
"GroupChatRequestMessage",
|
||||
"GroupChatRequestSentEvent",
|
||||
"GroupChatResponseReceivedEvent",
|
||||
"TerminationCondition",
|
||||
# Orchestration helpers
|
||||
"AgentRequestInfoResponse",
|
||||
"OrchestrationState",
|
||||
"clean_conversation_for_handoff",
|
||||
"create_completion_message",
|
||||
# Group Chat
|
||||
"AgentBasedGroupChatOrchestrator",
|
||||
"AgentOrchestrationOutput",
|
||||
"GroupChatBuilder",
|
||||
"GroupChatOrchestrator",
|
||||
"GroupChatSelectionFunction",
|
||||
"GroupChatState",
|
||||
# Magentic
|
||||
"MAGENTIC_MANAGER_NAME",
|
||||
"ORCH_MSG_KIND_INSTRUCTION",
|
||||
"ORCH_MSG_KIND_NOTICE",
|
||||
"ORCH_MSG_KIND_TASK_LEDGER",
|
||||
"ORCH_MSG_KIND_USER_TASK",
|
||||
"MagenticAgentExecutor",
|
||||
"MagenticBuilder",
|
||||
"MagenticContext",
|
||||
"MagenticManagerBase",
|
||||
"MagenticOrchestrator",
|
||||
"MagenticOrchestratorEvent",
|
||||
"MagenticOrchestratorEventType",
|
||||
"MagenticPlanReviewRequest",
|
||||
"MagenticPlanReviewResponse",
|
||||
"MagenticProgressLedger",
|
||||
"MagenticProgressLedgerItem",
|
||||
"MagenticResetSignal",
|
||||
"StandardMagenticManager",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
try:
|
||||
return getattr(importlib.import_module(IMPORT_PATH), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return _IMPORTS
|
||||
@@ -0,0 +1,75 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_orchestrations import (
|
||||
MAGENTIC_MANAGER_NAME,
|
||||
ORCH_MSG_KIND_INSTRUCTION,
|
||||
ORCH_MSG_KIND_NOTICE,
|
||||
ORCH_MSG_KIND_TASK_LEDGER,
|
||||
ORCH_MSG_KIND_USER_TASK,
|
||||
AgentBasedGroupChatOrchestrator,
|
||||
AgentOrchestrationOutput,
|
||||
AgentRequestInfoResponse,
|
||||
ConcurrentBuilder,
|
||||
GroupChatBuilder,
|
||||
GroupChatOrchestrator,
|
||||
GroupChatRequestMessage,
|
||||
GroupChatRequestSentEvent,
|
||||
GroupChatSelectionFunction,
|
||||
GroupChatState,
|
||||
HandoffAgentExecutor,
|
||||
HandoffAgentUserRequest,
|
||||
HandoffBuilder,
|
||||
HandoffConfiguration,
|
||||
HandoffSentEvent,
|
||||
MagenticAgentExecutor,
|
||||
MagenticBuilder,
|
||||
MagenticContext,
|
||||
MagenticManagerBase,
|
||||
MagenticOrchestrator,
|
||||
MagenticOrchestratorEvent,
|
||||
MagenticOrchestratorEventType,
|
||||
MagenticPlanReviewRequest,
|
||||
MagenticPlanReviewResponse,
|
||||
MagenticProgressLedger,
|
||||
MagenticProgressLedgerItem,
|
||||
MagenticResetSignal,
|
||||
SequentialBuilder,
|
||||
StandardMagenticManager,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"MAGENTIC_MANAGER_NAME",
|
||||
"ORCH_MSG_KIND_INSTRUCTION",
|
||||
"ORCH_MSG_KIND_NOTICE",
|
||||
"ORCH_MSG_KIND_TASK_LEDGER",
|
||||
"ORCH_MSG_KIND_USER_TASK",
|
||||
"AgentBasedGroupChatOrchestrator",
|
||||
"AgentOrchestrationOutput",
|
||||
"AgentRequestInfoResponse",
|
||||
"ConcurrentBuilder",
|
||||
"GroupChatBuilder",
|
||||
"GroupChatOrchestrator",
|
||||
"GroupChatRequestMessage",
|
||||
"GroupChatRequestSentEvent",
|
||||
"GroupChatSelectionFunction",
|
||||
"GroupChatState",
|
||||
"HandoffAgentExecutor",
|
||||
"HandoffAgentUserRequest",
|
||||
"HandoffBuilder",
|
||||
"HandoffConfiguration",
|
||||
"HandoffSentEvent",
|
||||
"MagenticAgentExecutor",
|
||||
"MagenticBuilder",
|
||||
"MagenticContext",
|
||||
"MagenticManagerBase",
|
||||
"MagenticOrchestrator",
|
||||
"MagenticOrchestratorEvent",
|
||||
"MagenticOrchestratorEventType",
|
||||
"MagenticPlanReviewRequest",
|
||||
"MagenticPlanReviewResponse",
|
||||
"MagenticProgressLedger",
|
||||
"MagenticProgressLedgerItem",
|
||||
"MagenticResetSignal",
|
||||
"SequentialBuilder",
|
||||
"StandardMagenticManager",
|
||||
]
|
||||
@@ -0,0 +1,33 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Redis integration namespace for optional Agent Framework connectors.
|
||||
|
||||
This module lazily re-exports objects from:
|
||||
- ``agent-framework-redis``
|
||||
|
||||
Supported classes:
|
||||
- RedisContextProvider
|
||||
- RedisHistoryProvider
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
IMPORT_PATH = "agent_framework_redis"
|
||||
PACKAGE_NAME = "agent-framework-redis"
|
||||
_IMPORTS = ["RedisContextProvider", "RedisHistoryProvider"]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _IMPORTS:
|
||||
try:
|
||||
return getattr(importlib.import_module(IMPORT_PATH), name)
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ModuleNotFoundError(
|
||||
f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`"
|
||||
) from exc
|
||||
raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return _IMPORTS
|
||||
@@ -0,0 +1,11 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_redis import (
|
||||
RedisContextProvider,
|
||||
RedisHistoryProvider,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"RedisContextProvider",
|
||||
"RedisHistoryProvider",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,140 @@
|
||||
[project]
|
||||
name = "agent-framework-core"
|
||||
description = "Microsoft Agent Framework for building AI Agents with Python. This is the core package that has all the core abstractions and implementations."
|
||||
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
version = "1.11.0"
|
||||
license-files = ["LICENSE"]
|
||||
urls.homepage = "https://aka.ms/agent-framework"
|
||||
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
|
||||
urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
|
||||
urls.issues = "https://github.com/microsoft/agent-framework/issues"
|
||||
classifiers = [
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Development Status :: 5 - Production/Stable",
|
||||
"Intended Audience :: Developers",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Programming Language :: Python :: 3.14",
|
||||
"Typing :: Typed",
|
||||
]
|
||||
dependencies = [
|
||||
"typing-extensions>=4.15.0,<5",
|
||||
"pydantic>=2,<3",
|
||||
"python-dotenv>=1,<2",
|
||||
"opentelemetry-api>=1.39.0,<2",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
all = [
|
||||
"mcp>=1.24.0,<2",
|
||||
"agent-framework-a2a",
|
||||
"agent-framework-ag-ui",
|
||||
"agent-framework-anthropic",
|
||||
"agent-framework-azure-ai-search",
|
||||
"agent-framework-azure-cosmos",
|
||||
"agent-framework-azurefunctions",
|
||||
"agent-framework-bedrock",
|
||||
"agent-framework-chatkit",
|
||||
"agent-framework-claude",
|
||||
"agent-framework-copilotstudio",
|
||||
"agent-framework-declarative",
|
||||
"agent-framework-devui",
|
||||
"agent-framework-durabletask",
|
||||
"agent-framework-foundry",
|
||||
"agent-framework-foundry-local",
|
||||
"agent-framework-github-copilot; python_version >= '3.11'",
|
||||
"agent-framework-hyperlight; ((sys_platform == 'linux' and platform_machine == 'x86_64') or (sys_platform == 'win32' and platform_machine == 'AMD64')) and python_version < '3.14'",
|
||||
"agent-framework-lab",
|
||||
"agent-framework-mem0",
|
||||
"agent-framework-ollama",
|
||||
"agent-framework-openai",
|
||||
"agent-framework-orchestrations",
|
||||
"agent-framework-purview",
|
||||
"agent-framework-redis",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
# These optional integrations depend on agent-framework-core, so core cannot list them as
|
||||
# runtime dependencies. They are declared here so isolated source checks can resolve their APIs.
|
||||
dev = [
|
||||
"agent-framework-azure-contentunderstanding",
|
||||
"agent-framework-tools",
|
||||
"azure-ai-agentserver-core>=2.0.0b7,<3",
|
||||
]
|
||||
|
||||
[tool.uv]
|
||||
prerelease = "if-necessary-or-explicit"
|
||||
environments = [
|
||||
"sys_platform == 'darwin'",
|
||||
"sys_platform == 'linux'",
|
||||
"sys_platform == 'win32'"
|
||||
]
|
||||
|
||||
[tool.uv-dynamic-versioning]
|
||||
fallback-version = "0.0.0"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ['tests']
|
||||
addopts = "-ra -q -r fEX"
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
filterwarnings = []
|
||||
timeout = 120
|
||||
markers = [
|
||||
"integration: marks tests as integration tests that require external services",
|
||||
]
|
||||
|
||||
[tool.coverage.run]
|
||||
omit = [
|
||||
"**/__init__.py"
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
extend = "../../pyproject.toml"
|
||||
|
||||
[tool.pyright]
|
||||
extends = "../../pyproject.toml"
|
||||
include = ["agent_framework"]
|
||||
|
||||
[tool.mypy]
|
||||
plugins = ['pydantic.mypy']
|
||||
strict = true
|
||||
incremental = false
|
||||
python_version = "3.10"
|
||||
ignore_missing_imports = true
|
||||
disallow_untyped_defs = true
|
||||
no_implicit_optional = true
|
||||
check_untyped_defs = true
|
||||
warn_return_any = true
|
||||
show_error_codes = true
|
||||
warn_unused_ignores = false
|
||||
disallow_incomplete_defs = true
|
||||
disallow_untyped_decorators = true
|
||||
|
||||
[tool.bandit]
|
||||
targets = ["agent_framework"]
|
||||
exclude_dirs = ["tests"]
|
||||
|
||||
[tool.poe]
|
||||
executor.type = "uv"
|
||||
include = "../../shared_tasks.toml"
|
||||
|
||||
[tool.poe.tasks.mypy]
|
||||
help = "Run MyPy for this package."
|
||||
cmd = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework"
|
||||
|
||||
[tool.poe.tasks.test]
|
||||
help = "Run the default unit test suite for this package."
|
||||
cmd = 'pytest -m "not integration" --cov=agent_framework --cov-report=term-missing:skip-covered -n auto --dist worksteal tests'
|
||||
|
||||
[tool.flit.module]
|
||||
name = "agent_framework"
|
||||
|
||||
[build-system]
|
||||
requires = ["flit-core >= 3.11,<4.0"]
|
||||
build-backend = "flit_core.buildapi"
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user