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
@@ -0,0 +1,88 @@
|
||||
{
|
||||
"version": "0.2",
|
||||
"languageSettings": [
|
||||
{
|
||||
"languageId": "py",
|
||||
"allowCompoundWords": true,
|
||||
"locale": "en-US"
|
||||
}
|
||||
],
|
||||
"language": "en-US",
|
||||
"patterns": [
|
||||
{
|
||||
"name": "import",
|
||||
"pattern": "import [a-zA-Z0-9_]+"
|
||||
},
|
||||
{
|
||||
"name": "from import",
|
||||
"pattern": "from [a-zA-Z0-9_]+ import [a-zA-Z0-9_]+"
|
||||
}
|
||||
],
|
||||
"ignorePaths": [
|
||||
"samples/**",
|
||||
"notebooks/**"
|
||||
],
|
||||
"words": [
|
||||
"aeiou",
|
||||
"agentserver",
|
||||
"agui",
|
||||
"aiplatform",
|
||||
"azuredocindex",
|
||||
"azuredocs",
|
||||
"azurefunctions",
|
||||
"boto",
|
||||
"codeact",
|
||||
"contentvector",
|
||||
"contoso",
|
||||
"datamodel",
|
||||
"desync",
|
||||
"dotenv",
|
||||
"endregion",
|
||||
"entra",
|
||||
"faiss",
|
||||
"finalizer",
|
||||
"finalizers",
|
||||
"genai",
|
||||
"generativeai",
|
||||
"hnsw",
|
||||
"httpx",
|
||||
"huggingface",
|
||||
"hyperlight",
|
||||
"Instrumentor",
|
||||
"logit",
|
||||
"logprobs",
|
||||
"lowlevel",
|
||||
"Magentic",
|
||||
"mistralai",
|
||||
"mongocluster",
|
||||
"nd",
|
||||
"ndarray",
|
||||
"nopep",
|
||||
"NOSQL",
|
||||
"ollama",
|
||||
"Onnx",
|
||||
"onyourdatatest",
|
||||
"OPENAI",
|
||||
"opentelemetry",
|
||||
"OTEL",
|
||||
"otlp",
|
||||
"powerfx",
|
||||
"protos",
|
||||
"pydantic",
|
||||
"pytestmark",
|
||||
"qdrant",
|
||||
"retrywrites",
|
||||
"serde",
|
||||
"streamable",
|
||||
"superstep",
|
||||
"supersteps",
|
||||
"templating",
|
||||
"uninstrument",
|
||||
"vectordb",
|
||||
"vectorizable",
|
||||
"vectorizer",
|
||||
"vectorstoremodel",
|
||||
"vertexai",
|
||||
"Weaviate"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
# Microsoft Foundry
|
||||
FOUNDRY_PROJECT_ENDPOINT=""
|
||||
# Model used for FoundryChatClient
|
||||
FOUNDRY_MODEL=""
|
||||
# Foundry Agents (prompt or hosted agents)
|
||||
FOUNDRY_AGENT_NAME=""
|
||||
FOUNDRY_AGENT_VERSION=""
|
||||
# Microsoft Foundry Models endpoint, used by embeddings
|
||||
FOUNDRY_MODELS_ENDPOINT=""
|
||||
FOUNDRY_MODELS_API_KEY=""
|
||||
FOUNDRY_EMBEDDING_MODEL=""
|
||||
FOUNDRY_IMAGE_EMBEDDING_MODEL=""
|
||||
# Bing connection for web search (optional, used by samples with web search)
|
||||
BING_CONNECTION_ID=""
|
||||
# Azure AI Search (optional, used by AzureAISearchContextProvider samples)
|
||||
AZURE_SEARCH_ENDPOINT=""
|
||||
AZURE_SEARCH_API_KEY=""
|
||||
AZURE_SEARCH_INDEX_NAME=""
|
||||
AZURE_SEARCH_SEMANTIC_CONFIG=""
|
||||
AZURE_SEARCH_KNOWLEDGE_BASE_NAME=""
|
||||
# Note: For agentic mode Knowledge Bases, also set AZURE_OPENAI_ENDPOINT below
|
||||
# (different from AZURE_AI_PROJECT_ENDPOINT - Knowledge Base needs OpenAI endpoint for model calls)
|
||||
# OpenAI
|
||||
OPENAI_API_KEY=""
|
||||
OPENAI_CHAT_COMPLETION_MODEL=""
|
||||
OPENAI_CHAT_MODEL=""
|
||||
# Azure OpenAI
|
||||
AZURE_OPENAI_ENDPOINT=""
|
||||
AZURE_OPENAI_CHAT_COMPLETION_MODEL=""
|
||||
AZURE_OPENAI_CHAT_MODEL=""
|
||||
# Mem0
|
||||
MEM0_API_KEY=""
|
||||
# Copilot Studio
|
||||
COPILOTSTUDIOAGENT__ENVIRONMENTID=""
|
||||
COPILOTSTUDIOAGENT__SCHEMANAME=""
|
||||
COPILOTSTUDIOAGENT__TENANTID=""
|
||||
COPILOTSTUDIOAGENT__AGENTAPPID=""
|
||||
# Anthropic
|
||||
ANTHROPIC_API_KEY=""
|
||||
ANTHROPIC_MODEL=""
|
||||
# Google Gemini
|
||||
GEMINI_API_KEY=""
|
||||
GEMINI_MODEL=""
|
||||
# Ollama
|
||||
OLLAMA_ENDPOINT=""
|
||||
OLLAMA_MODEL=""
|
||||
# Mistral AI
|
||||
MISTRAL_API_KEY=""
|
||||
MISTRAL_EMBEDDING_MODEL=""
|
||||
# Observability (instrumentation is enabled by default; set "ENABLE_INSTRUMENTATION" to "false" to opt out)
|
||||
ENABLE_SENSITIVE_DATA=true
|
||||
OTEL_EXPORTER_OTLP_ENDPOINT="http://localhost:4317/"
|
||||
@@ -0,0 +1,10 @@
|
||||
---
|
||||
applyTo: 'python/**'
|
||||
---
|
||||
|
||||
See [AGENTS.md](../../AGENTS.md) for project structure and package documentation.
|
||||
Detailed conventions are in the agent skills under `.github/skills/`.
|
||||
|
||||
When changing the public root API surface (`agent_framework/__init__.py`), keep the lazy runtime export
|
||||
registry, explicit runtime `__all__`, and `agent_framework/__init__.pyi` synchronized. Runtime deprecation
|
||||
behavior for a public alias should live in the owning module, not as a special case in the root package.
|
||||
@@ -0,0 +1,361 @@
|
||||
---
|
||||
name: agent-framework-py-release
|
||||
description: Use when cutting a Python release for the microsoft/agent-framework monorepo. Triggers on "bump py versions", "cut a python release", "prepare release PR for python", "release py packages", "bump python to X.Y.Z", or similar requests to bump Python package versions and prepare a release PR. Handles all four lifecycle tiers (alpha, beta, rc, released) with CHANGELOG-driven selective bumps, floor bound checks, and post-bump validation.
|
||||
---
|
||||
|
||||
# Agent Framework Python Release
|
||||
|
||||
Cuts a Python release PR for the `microsoft/agent-framework` monorepo.
|
||||
|
||||
## Core principle: CHANGELOG drives bumps
|
||||
|
||||
**No more universal lockstep.** A package only bumps if it has a CHANGELOG entry this cycle. Draft the CHANGELOG first, then derive the bump list from it.
|
||||
|
||||
Exceptions to selective bumping (these stay coupled regardless of per-package changes):
|
||||
|
||||
- **Root `agent-framework` is coupled to `agent-framework-core` via `agent-framework-core[all]==X.Y.Z`** (exact pin in `python/pyproject.toml`). If `core` bumps, root bumps. Root may also bump independently for its own changes (docs/metadata/extras).
|
||||
- **Beta-tier cohort bumps are still allowed** when the user explicitly wants a fresh date stamp across all betas for cohort signaling — but this is a deliberate user decision per release, not a default.
|
||||
|
||||
If the user explicitly states which packages to bump (or hands over a PR list per package), treat that as authoritative.
|
||||
|
||||
## Lifecycle tiers and version rules
|
||||
|
||||
Use `python/.github/skills/python-package-management/SKILL.md` as the source of truth for lifecycle
|
||||
version patterns, date-stamp cutoffs, classifier alignment, and internal dependency updates.
|
||||
|
||||
For release work, derive the live tier map at release time from `python/PACKAGE_STATUS.md` and the actual
|
||||
`version =` lines in each `pyproject.toml`. Do not hardcode counts — packages move between tiers.
|
||||
|
||||
## Inputs to confirm before bumping
|
||||
|
||||
1. **The changeset**: explicit commits/PRs the release covers, OR derive from `git log ${LAST_RELEASED_TAG}..origin/main -- python/`.
|
||||
2. **Per-package CHANGELOG entries**: which packages will get a line in the new release section. This list IS the bump list.
|
||||
3. **Per-released-package semver bump**: for each released-tier package that has a CHANGELOG entry, decide PATCH / MINOR / MAJOR.
|
||||
4. **Date stamp** (only if any alpha/beta is being bumped): default from the `python-package-management`
|
||||
date-stamp rule.
|
||||
5. **Optional cohort bump?**: ask if betas should all get a new date stamp regardless of per-package changes. Default: no.
|
||||
|
||||
If the user states target versions or a date explicitly, use exactly what they said — do not "correct" based on historical timezones.
|
||||
|
||||
## Non-negotiable rules
|
||||
|
||||
- **CHANGELOG-driven bumps**: only packages mentioned in the new CHANGELOG section get version bumps. Exceptions: root follows core (==pin); user-opted cohort bump on betas.
|
||||
- **Follow `python-package-management` for package lifecycle and versioning rules** — do not duplicate those
|
||||
rules in this release workflow.
|
||||
- **No `Co-Authored-By` trailer** on any commit.
|
||||
- **Use `uv run`** for all Python/poe commands (`uv run poe ...`, `uv run pytest ...`).
|
||||
- **Never rename an existing CHANGELOG section header** during a new release cut. Only INSERT a new section above existing ones.
|
||||
- **Footer reference links are part of the CHANGELOG edit**, not optional.
|
||||
|
||||
## Workflow
|
||||
|
||||
### 1. Orient and create branch
|
||||
|
||||
```bash
|
||||
git fetch origin main --tags --quiet
|
||||
git fetch upstream main --tags --quiet 2>/dev/null || true
|
||||
git status
|
||||
```
|
||||
|
||||
If the user already has a `bump-py-ver-release-*` branch checked out, use it. Otherwise:
|
||||
|
||||
```bash
|
||||
git checkout -b bump-py-ver-release-YYMMDD origin/main
|
||||
```
|
||||
|
||||
### 2. Build the live tier map
|
||||
|
||||
Read `python/PACKAGE_STATUS.md` to enumerate current packages and lifecycle stages, and `grep '^version' python/pyproject.toml python/packages/*/pyproject.toml` to confirm actual versions. Cross-reference — if `PACKAGE_STATUS.md` shows a package as `beta` but its version looks like `1.0.0aYYMMDD`, surface the inconsistency before proceeding.
|
||||
|
||||
Watch for: new packages added since the last release, lifecycle transitions (alpha→beta, beta→rc, rc→released), and stale packages that haven't been touched in many cycles.
|
||||
|
||||
### 3. Identify changeset
|
||||
|
||||
Use the LATEST released tag as the compare base:
|
||||
|
||||
```bash
|
||||
LAST_RELEASED_TAG=$(
|
||||
git tag -l 'python-[0-9]*.[0-9]*.[0-9]*' \
|
||||
| uv run --directory python python -c 'import sys; tags=[t.strip() for t in sys.stdin if t.strip()]; print(max(tags, key=lambda t: tuple(map(int, t[7:].split(".")))))'
|
||||
)
|
||||
echo "Compare base: $LAST_RELEASED_TAG"
|
||||
```
|
||||
|
||||
List commits and packages touched:
|
||||
|
||||
```bash
|
||||
git log --oneline ${LAST_RELEASED_TAG}..origin/main -- python/ ':!python/CHANGELOG.md'
|
||||
|
||||
# Per-commit package footprint
|
||||
for sha in $(git log --format='%H' ${LAST_RELEASED_TAG}..origin/main -- python/); do
|
||||
echo "--- $(git show -s --format='%h %s' $sha) ---"
|
||||
git show --name-only --format='' $sha | grep '^python/packages/' | \
|
||||
sed 's|^python/packages/||' | awk -F/ '{print $1}' | sort -u
|
||||
done
|
||||
```
|
||||
|
||||
If the release ultimately tags from `upstream/main` but `origin/main` is behind, also run:
|
||||
|
||||
```bash
|
||||
git log --oneline ${LAST_RELEASED_TAG}..upstream/main -- python/ ':!python/CHANGELOG.md'
|
||||
```
|
||||
|
||||
If user provides an explicit commit/PR list, treat THAT as authoritative.
|
||||
|
||||
#### 3a. Build the authoritative touched-package set
|
||||
|
||||
Aggregate the per-commit footprint into a single union across the whole range. This is the **source of truth** for what must appear in the CHANGELOG. Save it before drafting entries.
|
||||
|
||||
```bash
|
||||
# Union of all touched package directories across the range
|
||||
git log --name-only --format='' ${LAST_RELEASED_TAG}..origin/main -- python/packages/ \
|
||||
| grep '^python/packages/' \
|
||||
| sed 's|^python/packages/||' \
|
||||
| awk -F/ '{print $1}' \
|
||||
| sort -u
|
||||
|
||||
# Root-level files (drive a root agent-framework entry if substantive)
|
||||
git log --name-only --format='' ${LAST_RELEASED_TAG}..origin/main \
|
||||
-- python/pyproject.toml python/agent_framework_meta/ python/README.md \
|
||||
2>/dev/null | grep -v '^$' | sort -u
|
||||
```
|
||||
|
||||
**What "touched" means for bump purposes** (apply judgment, document in commit message):
|
||||
|
||||
| Change scope under `python/packages/<pkg>/` | Drives a bump? |
|
||||
| --- | --- |
|
||||
| Source code under `<pkg>/agent_framework_*` or the package's library tree | **Yes** — ship-affecting |
|
||||
| `pyproject.toml` (deps, classifiers, extras, version metadata) | **Yes** — publishes new metadata |
|
||||
| `README.md` substantive changes (env vars, install instructions, usage) | **Yes** — user-facing |
|
||||
| `README.md` typo/wording-only | Judgment — usually no, but record under `### Changed` if you do bump |
|
||||
| Tests only (`tests/`) | Usually no — test changes don't ship to PyPI consumers. Bump only if the test change reflects a real behavior change in the package |
|
||||
| Package-local samples (alpha packages only, before promotion) | **Yes** for alpha; samples ship with the alpha package |
|
||||
| Repo infra touching the package dir (CI config, lint config) | No — record under `- **tests**:`, `- **samples**:`, or `- **docs**:` instead |
|
||||
|
||||
Root `agent-framework` is touched when `python/pyproject.toml`, `python/agent_framework_meta/`, or root `README.md` substantive content changed.
|
||||
|
||||
### 4. Draft CHANGELOG entries (THIS DRIVES THE BUMP LIST)
|
||||
|
||||
Locate `## [Unreleased]` and the top existing release header. INSERT a new section between them.
|
||||
|
||||
**New section structure:**
|
||||
|
||||
```markdown
|
||||
## [<release-label>] - YYYY-MM-DD
|
||||
|
||||
### Added
|
||||
- **agent-framework-<pkg>**: <subject> ([#NNNN](https://github.com/microsoft/agent-framework/pull/NNNN))
|
||||
|
||||
### Changed
|
||||
- **agent-framework-<pkg>**: ...
|
||||
|
||||
### Fixed
|
||||
- **agent-framework-<pkg>**: ...
|
||||
```
|
||||
|
||||
`<release-label>` is the released-tier target if any released package is bumping (e.g. `1.7.0`), or the date stamp (e.g. `1.0.0b260528`) if the release is prerelease-only. Reflect the user's framing.
|
||||
|
||||
**Categorization heuristics (read commit subject + PR title/body):**
|
||||
- **Added**: new public APIs, new packages, new samples, `feat(...)`, new re-exports from `agent_framework`, new capabilities surfaced from underlying SDKs
|
||||
- **Changed**: metadata/classifier corrections, dependency upgrades, test infra changes, signature adjustments on existing APIs, renames, doc updates, lifecycle transitions
|
||||
- **Fixed**: `fix(...)`, bug/crash/regression repairs, noise reduction (e.g. stop emitting spurious warnings), protocol-compliance fixes
|
||||
- **Removed**: deprecations/deletions (use `[BREAKING]` prefix if breaking)
|
||||
|
||||
**Entry format:**
|
||||
- `- **agent-framework-<pkg>**: <subject, grammar-normalized> ([#NNNN](https://github.com/microsoft/agent-framework/pull/NNNN))`
|
||||
- Multiple packages touched by one PR: comma-separated bold names at the head: `- **agent-framework-core**, **agent-framework-foundry**: ...`
|
||||
- Root package changes: use `- **agent-framework**: ...`
|
||||
- Test/sample/repo infra: `- **tests**: ...`, `- **samples**: ...`, `- **docs**: ...` (these do not drive package bumps)
|
||||
|
||||
**Once entries are drafted, the set of `**agent-framework-<pkg>**` and root `**agent-framework**` mentions IS the bump list.** Anything not mentioned does not bump.
|
||||
|
||||
#### 4a. Reconcile mentions against the touched-package set (DO NOT SKIP)
|
||||
|
||||
Before moving on, prove that every ship-affecting touched package has at least one CHANGELOG entry. A package whose code changed but is missing from CHANGELOG will NOT bump — and the fix will not ship to PyPI consumers.
|
||||
|
||||
```bash
|
||||
# 1. Touched ship-affecting packages and root package files (from step 3a)
|
||||
TOUCHED_PACKAGES=$(git log --name-only --format='' ${LAST_RELEASED_TAG}..origin/main -- python/packages/ \
|
||||
| grep '^python/packages/' \
|
||||
| sed 's|^python/packages/||' \
|
||||
| awk -F/ '{print $1}' \
|
||||
| sort -u)
|
||||
|
||||
ROOT_TOUCHED=$(git log --name-only --format='' ${LAST_RELEASED_TAG}..origin/main \
|
||||
-- python/pyproject.toml python/agent_framework_meta/ python/README.md \
|
||||
2>/dev/null | grep -v '^$' | sort -u)
|
||||
|
||||
TOUCHED=$(
|
||||
{
|
||||
if [ -n "$TOUCHED_PACKAGES" ]; then echo "$TOUCHED_PACKAGES"; fi
|
||||
if [ -n "$ROOT_TOUCHED" ]; then echo "agent-framework"; fi
|
||||
} | sort -u
|
||||
)
|
||||
|
||||
# 2. Packages mentioned in the new CHANGELOG section
|
||||
# (Extract from the section you just drafted — adjust the awk range to the new section's bounds)
|
||||
MENTIONED=$(awk '
|
||||
/^## \[<release-label>\]/ { in_section=1; next }
|
||||
in_section && /^## \[/ { exit }
|
||||
in_section { print }
|
||||
' python/CHANGELOG.md \
|
||||
| grep -oE '\*\*agent-framework(-[a-z0-9_-]+)?\*\*' \
|
||||
| sed 's/\*\*//g;s/^agent-framework-//' \
|
||||
| sort -u)
|
||||
|
||||
# 3. Diff — anything in TOUCHED but missing from MENTIONED is a gap
|
||||
comm -23 <(echo "$TOUCHED") <(echo "$MENTIONED")
|
||||
```
|
||||
|
||||
For every gap, decide explicitly with the user:
|
||||
|
||||
- **Add a CHANGELOG entry and bump** — the default. Even small package changes (a dep upgrade, a metadata fix) deserve an entry under `### Changed`.
|
||||
- **Intentionally skip** — only when the touched files are demonstrably non-shipping (e.g. comments-only, test infra). Note the skip in the commit message body so reviewers see the reasoning.
|
||||
|
||||
Package-name aliasing to watch for: directory `foundry_local` → package `agent-framework-foundry-local`; `github_copilot` → `agent-framework-github-copilot`; `azure-ai-search` → `agent-framework-azure-ai-search`. The directory name and the published PyPI name are not always identical — confirm both when reconciling.
|
||||
|
||||
**Footer reference links** (REQUIRED every release):
|
||||
|
||||
```bash
|
||||
grep -n "^\[.*\]:" python/CHANGELOG.md | head -5
|
||||
```
|
||||
|
||||
Two edits when a released-tier bump is in play:
|
||||
1. Advance `[Unreleased]` compare base from `python-${OLD_RELEASED}...HEAD` to `python-${NEW_RELEASED}...HEAD`.
|
||||
2. INSERT a new `[${NEW_RELEASED}]` line ABOVE the previous version's link: `[${NEW_RELEASED}]: https://github.com/microsoft/agent-framework/compare/python-${OLD_RELEASED}...python-${NEW_RELEASED}`
|
||||
|
||||
For prerelease-only releases (no released-tier bump this cycle), footer links don't change.
|
||||
|
||||
### 5. Apply bumps per tier
|
||||
|
||||
For each package in the bump list, choose the rule for its current tier. Use anchored `sed` (`^version = "..."$`) so only the project's own version matches. Use `sed -i.bak` plus backup cleanup for portable in-place edits across macOS/BSD sed and GNU sed.
|
||||
|
||||
**Released tier (per-package semver):**
|
||||
|
||||
```bash
|
||||
# Example: openai goes 1.6.0 -> 1.6.1 (PATCH); core stays
|
||||
sed -i.bak 's/^version = "1.6.0"$/version = "1.6.1"/' python/packages/openai/pyproject.toml
|
||||
rm python/packages/openai/pyproject.toml.bak
|
||||
```
|
||||
|
||||
**Root `agent-framework` (only when `core` bumps, OR for its own changes):**
|
||||
|
||||
```bash
|
||||
sed -i.bak "s/^version = \"${OLD_ROOT}\"$/version = \"${NEW_ROOT}\"/" python/pyproject.toml
|
||||
# AND keep the exact pin in sync with core
|
||||
sed -i.bak "s/agent-framework-core\\[all\\]==${OLD_CORE}/agent-framework-core[all]==${NEW_CORE}/" python/pyproject.toml
|
||||
rm python/pyproject.toml.bak
|
||||
```
|
||||
|
||||
**RC tier (counter increment):**
|
||||
|
||||
```bash
|
||||
# ag-ui goes 1.0.0rc3 -> 1.0.0rc4 only because it has a CHANGELOG entry
|
||||
sed -i.bak 's/^version = "1.0.0rc3"$/version = "1.0.0rc4"/' python/packages/ag-ui/pyproject.toml
|
||||
rm python/packages/ag-ui/pyproject.toml.bak
|
||||
```
|
||||
|
||||
**Alpha/Beta tier (new date stamp):**
|
||||
|
||||
```bash
|
||||
OLD_DATE=260521
|
||||
NEW_DATE=260528
|
||||
|
||||
# Only the packages with CHANGELOG entries get the new stamp.
|
||||
# Example: anthropic and bedrock had entries, others did not.
|
||||
for pkg in anthropic bedrock; do
|
||||
file="python/packages/$pkg/pyproject.toml"
|
||||
sed -i.bak "s/1\\.0\\.0b${OLD_DATE}/1.0.0b${NEW_DATE}/g" "$file"
|
||||
rm "${file}.bak"
|
||||
done
|
||||
```
|
||||
|
||||
If the user opts into a cohort-wide beta bump regardless of per-package changes, enumerate every beta package from `PACKAGE_STATUS.md` and stamp all of them. State that this is a cohort bump in the commit message.
|
||||
|
||||
Spot-check with `grep '^version' python/pyproject.toml python/packages/*/pyproject.toml | sort` before moving on.
|
||||
|
||||
### 6. Floor bound updates on `agent-framework-core`
|
||||
|
||||
Only relevant when `core` itself bumped this cycle. Two policies, pick one explicitly with the user:
|
||||
|
||||
- **Conservative (default)**: raise `agent-framework-core>=X.Y.Z` to the new core version on every non-core package that is ALSO bumping this cycle. Leaves packages-not-bumped at their existing floor.
|
||||
- **Strict per-upstream-doc**: only raise the floor on packages that actually consume a new core API introduced in the bump. This requires per-package code inspection. Use only when the user is comfortable letting `validate-dependency-bounds-test` (lower-resolution pass) catch any mistakes.
|
||||
|
||||
When raising a core floor, replace only the `>=OLD` half of the bound you intend to change:
|
||||
|
||||
```bash
|
||||
file="python/packages/<pkg>/pyproject.toml"
|
||||
sed -i.bak "s/agent-framework-core>=${OLD_CORE}/agent-framework-core>=${NEW_CORE}/" "$file"
|
||||
rm "${file}.bak"
|
||||
```
|
||||
|
||||
If `core` did not bump this cycle, do not touch floors.
|
||||
|
||||
### 7. Validate
|
||||
|
||||
```bash
|
||||
cd python && uv run poe validate-dependency-bounds-test
|
||||
```
|
||||
|
||||
Must exit 0. This is the safety net for selective bumping: the lower-resolution pass catches floors set too low for code that depends on new APIs, and the upper pass catches caps that exclude installable versions. If it fails, the output names the offending bound — fix and re-run before committing. This step also regenerates `uv.lock` to match new bounds.
|
||||
|
||||
If only prereleases changed (no `core` bump, no floor changes), this validation is still required — `uv.lock` regeneration alone justifies the run.
|
||||
|
||||
### 8. Commit (expect hook retry)
|
||||
|
||||
The project's pre-commit hook runs `poe check` / `poe typing` / `uv lock` and may re-modify `uv.lock` on first commit attempt. Expected pattern — run twice:
|
||||
|
||||
```bash
|
||||
git add -u
|
||||
git commit -m "<message>" # may fail with "files were modified by this hook"
|
||||
git add -u # re-stage hook-generated changes
|
||||
git commit -m "<message>" # succeeds
|
||||
```
|
||||
|
||||
**Commit message format** — HEREDOC, no `Co-Authored-By`:
|
||||
|
||||
```bash
|
||||
git commit -m "$(cat <<'EOF'
|
||||
Bump Python package versions for <release-label> release
|
||||
|
||||
<One paragraph explaining:
|
||||
- which packages bumped and why (CHANGELOG-driven)
|
||||
- semver rationale for any released-tier bumps
|
||||
- whether a beta cohort bump was applied
|
||||
- whether core floors were raised, and under which policy>
|
||||
EOF
|
||||
)"
|
||||
```
|
||||
|
||||
### 9. Push and report
|
||||
|
||||
```bash
|
||||
git push origin bump-py-ver-release-YYMMDD
|
||||
```
|
||||
|
||||
The push output includes a `Create a pull request for '<branch>' on GitHub by visiting: <URL>` line — surface that to the user.
|
||||
|
||||
## Pitfalls from past cycles — avoid repeating
|
||||
|
||||
- **Scoping from the wrong tag.** A release may have been tagged between bump cycles. Always compute `LAST_RELEASED_TAG` fresh — do NOT reuse the previous compare base. Symptom: dragging already-shipped PRs into the new CHANGELOG section.
|
||||
- **Bumping packages that had no changes.** This was the old lockstep default. Now: a package without a CHANGELOG entry does not bump. The two exceptions (root follows core, optional beta cohort bump) are explicit choices, not defaults.
|
||||
- **Touched package missing from CHANGELOG.** The inverse failure mode of the previous pitfall, and the dangerous one. If a PR landed code into `python/packages/foo/` but the new CHANGELOG section has no `**agent-framework-foo**:` line, then `foo` will NOT bump, and the fix sits in `main` but never reaches PyPI. Always run the touched-vs-mentioned reconciliation in step 4a before applying bumps.
|
||||
- **Directory name vs PyPI name confusion.** Some package dirs use underscores or split names that differ from the published PyPI name (e.g. `foundry_local` → `agent-framework-foundry-local`). When reconciling touched dirs against CHANGELOG mentions, normalize on the PyPI form.
|
||||
- **Forgetting root when core bumps.** Root `agent-framework` has `agent-framework-core[all]==EXACT_VERSION`. If `core` bumps and root does not, the pin is broken. Always bump them together and update the pin string.
|
||||
- **Renaming an existing section header.** When drafting a new release, do NOT rewrite an existing `## [X.Y.Z]` header to a new label — that wipes the historical section. Always INSERT a new section above.
|
||||
- **Forgetting footer reference links.** When a released-tier bump is in play, the `[Unreleased]` line and the new `[X.Y.Z]` link at the bottom MUST be updated. Heading links don't resolve without them.
|
||||
- **Wrong footer compare base.** The new `[X.Y.Z]` line compares from the PREVIOUS released tag, not from two releases ago.
|
||||
- **Timezone drift.** For alpha/beta date stamps, follow the `python-package-management` date-stamp rule;
|
||||
do not infer a local timezone from the user's current shell.
|
||||
- **`Co-Authored-By` trailer.** Never add it. Rewrite/amend if it slipped in.
|
||||
- **Stale inventory in this skill.** Always read `python/PACKAGE_STATUS.md` for the live tier map. Do not trust a hardcoded list.
|
||||
- **Divergent origin vs upstream.** If the release tags from `upstream/main` but `origin/main` is behind, check both — warn if they differ and offer to sync.
|
||||
- **`--pre` README cleanup on promotion.** When a package is promoted to `released` in this cycle, grep for `pip install agent-framework-<pkg> --pre` in READMEs and drop the `--pre` flag.
|
||||
- **RC counter inflation.** Do not increment `1.0.0rcN` without a CHANGELOG entry for that package. The counter tracks iterations, not calendar.
|
||||
|
||||
## References
|
||||
|
||||
- Package lifecycle and versioning source of truth: `python/.github/skills/python-package-management/SKILL.md`
|
||||
- Lifecycle source of truth: `python/PACKAGE_STATUS.md`
|
||||
- Validator: `python/scripts/dependencies/validate_dependency_bounds.py` (runs `lowest-direct` and `highest` resolution smoke tests; catches floors/caps that don't match the code)
|
||||
- Poe task definitions: `python/pyproject.toml` `[tool.poe.tasks]`
|
||||
+116
@@ -0,0 +1,116 @@
|
||||
---
|
||||
name: pull-requests
|
||||
description: >
|
||||
Guidance for creating pull requests and handling PR review comments in the
|
||||
Agent Framework repository. Use this when writing a PR description (filling out
|
||||
the PR template) or when responding to and resolving review comments on an
|
||||
existing PR.
|
||||
---
|
||||
|
||||
# Pull Request Workflow
|
||||
|
||||
This skill covers two tasks: (1) writing a high-quality PR description, and
|
||||
(2) handling review comments on an existing PR.
|
||||
|
||||
## 1. Writing the PR description
|
||||
|
||||
Always follow the repository PR template at
|
||||
[`.github/pull_request_template.md`](../../../../.github/pull_request_template.md). Keep its
|
||||
exact structure and headings. Fill every section:
|
||||
|
||||
### `### Motivation & Context`
|
||||
Explain *why* the change is needed: the problem it solves and the scenario it
|
||||
contributes to. Describe the net change relative to `main` — this is implied, so
|
||||
do **not** spell out "vs main" explicitly.
|
||||
|
||||
### `### Description & Review Guide`
|
||||
Describe the changes, the overall approach, and the design. Answer the three
|
||||
prompts:
|
||||
- **What are the major changes?**
|
||||
- **What is the impact of these changes?**
|
||||
- **What do you want reviewers to focus on?** — This item is for **human
|
||||
reviewers only**. Automated/AI reviewers must ignore it and review the entire
|
||||
change rather than narrowing scope to it.
|
||||
|
||||
### `### Related Issue`
|
||||
Link the issue the PR fixes using a GitHub closing keyword (`Fixes #123` /
|
||||
`Closes #123`) so it closes automatically on merge. A PR with no linked issue may
|
||||
be closed regardless of how valid the change is. Before opening, confirm there is
|
||||
no other open PR for the same issue; if there is, explain how this PR differs.
|
||||
|
||||
### `### Contribution Checklist`
|
||||
Check every item that applies. For the breaking-change item:
|
||||
- Leave **"This is not a breaking change."** checked for the common case.
|
||||
- If the change **is** breaking, add the `breaking change` label **or** put
|
||||
`[BREAKING]` in the title prefix, before or after a language prefix such as
|
||||
`Python:` or `.NET:` — workflows keep the label and the title prefix in sync
|
||||
automatically (see `.github/workflows/label-title-prefix.yml` and
|
||||
`.github/workflows/label-pr.yml`).
|
||||
|
||||
### Do not
|
||||
- Do **not** add ad-hoc sections such as "Validation" or "Tests run"; CI/CD and
|
||||
the checklist already cover validation status.
|
||||
- Do **not** remove or reorder the template's headings.
|
||||
|
||||
### Creating the PR
|
||||
Open new PRs as **drafts** until they are ready for review. Example:
|
||||
|
||||
```bash
|
||||
gh pr create --repo microsoft/agent-framework --base main \
|
||||
--head <your-fork-owner>:<branch> --draft \
|
||||
--title "<concise title>" --body "<body following the template>"
|
||||
```
|
||||
|
||||
## 2. Handling review comments
|
||||
|
||||
When a PR receives review comments, follow this sequence — **do not start editing
|
||||
code before the user has reviewed the plan**:
|
||||
|
||||
1. **Review the comments.** Read every review comment and thread on the PR,
|
||||
including inline code comments and general review summaries.
|
||||
2. **Make a plan.** Produce a concrete plan describing how each comment will be
|
||||
addressed (or why it should not be, with reasoning).
|
||||
3. **Let the user review the plan.** Present the plan and wait for the user's
|
||||
approval or adjustments before implementing anything.
|
||||
4. **Implement.** Make the agreed changes.
|
||||
5. **Reply to every comment.** Add a reply to **all** comments explaining how it
|
||||
was addressed (or the agreed outcome) — leave none unanswered.
|
||||
6. **Resolve resolved threads.** Mark a review thread as resolved only when the
|
||||
comment has actually been addressed.
|
||||
|
||||
### Useful commands
|
||||
|
||||
List review comments and threads:
|
||||
|
||||
```bash
|
||||
# Inline review comments
|
||||
gh api repos/{owner}/{repo}/pulls/{pr}/comments
|
||||
|
||||
# Review threads with resolution state (GraphQL)
|
||||
gh api graphql -f query='
|
||||
query($owner:String!,$repo:String!,$pr:Int!){
|
||||
repository(owner:$owner,name:$repo){
|
||||
pullRequest(number:$pr){
|
||||
reviewThreads(first:100){
|
||||
nodes{ id isResolved comments(first:50){ nodes{ id body author{login} } } }
|
||||
}
|
||||
}
|
||||
}
|
||||
}' -F owner={owner} -F repo={repo} -F pr={pr}
|
||||
```
|
||||
|
||||
Reply to an inline review comment:
|
||||
|
||||
```bash
|
||||
gh api repos/{owner}/{repo}/pulls/{pr}/comments/{comment_id}/replies \
|
||||
-f body="Addressed in <commit>: <explanation>"
|
||||
```
|
||||
|
||||
Resolve a review thread (needs the thread node id from the GraphQL query above):
|
||||
|
||||
```bash
|
||||
gh api graphql -f query='
|
||||
mutation($threadId:ID!){
|
||||
resolveReviewThread(input:{threadId:$threadId}){ thread{ isResolved } }
|
||||
}' -F threadId={thread_id}
|
||||
```
|
||||
+135
@@ -0,0 +1,135 @@
|
||||
---
|
||||
name: python-code-quality
|
||||
description: >
|
||||
Code quality checks, linting, formatting, and type checking commands for the
|
||||
Agent Framework Python codebase. Use this when running checks, fixing lint
|
||||
errors, or troubleshooting CI failures.
|
||||
---
|
||||
|
||||
# Python Code Quality
|
||||
|
||||
## Quick Commands
|
||||
|
||||
All commands run from the `python/` directory:
|
||||
|
||||
```bash
|
||||
# Syntax formatting + checks (parallel across packages by default)
|
||||
uv run poe syntax
|
||||
uv run poe syntax -P core
|
||||
uv run poe syntax -F # Format only
|
||||
uv run poe syntax -C # Check only
|
||||
uv run poe syntax -S # Samples only
|
||||
|
||||
# Type checking
|
||||
#
|
||||
# Division of labor (see "Type checking architecture" below):
|
||||
# - Pyright (strict) is the source-code type checker.
|
||||
# - Pyright (relaxed `basic`), mypy, pyrefly, ty, zuban all check the TESTS;
|
||||
# pyright/pyrefly/ty also check the SAMPLES (mypy/zuban skip script-style samples).
|
||||
uv run poe pyright # Pyright (strict) over SOURCE, fan-out across packages
|
||||
uv run poe pyright -P core
|
||||
uv run poe pyright -A
|
||||
uv run poe test-typing # mypy + pyrefly + ty + zuban + pyright over each package's TESTS
|
||||
uv run poe test-typing -P core
|
||||
uv run poe test-typing -S # samples (pyrefly + ty + pyright)
|
||||
uv run poe test-typing -P core --checker mypy # narrow to one checker (repeatable)
|
||||
uv run poe test-typing -P core --checker pyright # relaxed pyright over the tests
|
||||
uv run poe mypy # alias: MyPy over the tests only
|
||||
uv run poe mypy -P core
|
||||
uv run poe typing # Pyright (source) + the tests checkers
|
||||
uv run poe typing -P core
|
||||
uv run poe typing -A
|
||||
|
||||
# All package-level checks in parallel (syntax + pyright)
|
||||
uv run poe check-packages
|
||||
|
||||
# Full check (packages + samples + tests + markdown)
|
||||
uv run poe check
|
||||
uv run poe check -P core
|
||||
|
||||
# Samples only
|
||||
uv run poe check -S
|
||||
uv run poe pyright -S
|
||||
|
||||
# Markdown code blocks
|
||||
uv run poe markdown-code-lint
|
||||
```
|
||||
|
||||
## Pre-commit Hooks (prek)
|
||||
|
||||
Prek hooks run automatically on commit. They stay lightweight and only check
|
||||
changed files.
|
||||
|
||||
```bash
|
||||
# Install hooks
|
||||
uv run poe prek-install
|
||||
|
||||
# Run all hooks manually
|
||||
uv run prek run -a
|
||||
|
||||
# Run on last commit
|
||||
uv run prek run --last-commit
|
||||
```
|
||||
|
||||
They run changed-package syntax formatting/checking, markdown code lint only
|
||||
when markdown files change, and sample syntax lint/pyright only when files
|
||||
under `samples/` change.
|
||||
They intentionally do not run workspace `pyright` or `mypy` by default.
|
||||
|
||||
## Type checking architecture
|
||||
|
||||
Following the "too many type checkers" approach, type checkers are split by target:
|
||||
|
||||
| Target | Checker(s) | Mode | Config |
|
||||
|--------|-----------|------|--------|
|
||||
| Source (`agent_framework*`) | **pyright** | strict | `[tool.pyright]` in `pyproject.toml` |
|
||||
| Tests | pyright, mypy, pyrefly, ty, zuban | relaxed/basic | `pyrightconfig.tests.json`, `[tool.mypy]`, `pyrefly.toml`, `ty` rules |
|
||||
| Samples | pyright, pyrefly, ty | basic | `pyrightconfig.samples.json`, `pyrefly.samples.toml`, `ty.samples.toml` |
|
||||
|
||||
- **Pyright is the only *strict* source-code checker**, and it ALSO runs in a relaxed
|
||||
`basic` profile over the tests and samples (so the surfaces customers copy from are
|
||||
validated by every checker, including pyright). MyPy was removed from source; its
|
||||
`[tool.mypy]` block is now a *relaxed* profile used only for tests/samples.
|
||||
- The extra checkers run over tests/samples because those exercise the public API the way
|
||||
users do. The profile is intentionally relaxed (private access allowed, untyped test
|
||||
bodies allowed) so authors aren't forced into ugly over-annotation.
|
||||
- **Gating checkers** are `pyright`, `mypy`, `pyrefly`, `ty`, and `zuban` — all five run by
|
||||
default and gate CI. `zuban` is the strictest of the mypy-compatible pair, so the same
|
||||
`[tool.mypy]` config yields more findings; suppress zuban-only friction with shared
|
||||
`# type: ignore[code]`. Suppress relaxed-pyright friction with `# pyright: ignore[rule]`.
|
||||
- **Samples** add `pyright` to `pyrefly` + `ty` — mypy/zuban can't resolve script-style
|
||||
sample layouts (numeric-prefixed dirs, duplicate `main.py`), but pyright handles them.
|
||||
- The strict source-pyright (`[tool.pyright]`) enforces `reportUnnecessaryTypeIgnoreComment`
|
||||
and excludes tests/samples; the relaxed test/sample pyright configs do not flag unnecessary
|
||||
ignores.
|
||||
|
||||
## Ruff Configuration
|
||||
|
||||
- Line length: 120
|
||||
- Target: Python 3.10+
|
||||
- Auto-fix enabled
|
||||
- Rules: ASYNC, B, CPY, D, E, ERA, F, FIX, I, INP, ISC, Q, RET, RSE, RUF, SIM, T20, TD, W, T100, S
|
||||
- Scripts directory is excluded from checks
|
||||
|
||||
## Pyright Configuration
|
||||
|
||||
- **Source**: strict mode (`[tool.pyright]`), `reportUnnecessaryTypeIgnoreComment = "error"`,
|
||||
excludes tests, samples, .venv, packages/devui/frontend.
|
||||
- **Tests**: relaxed `basic` profile (`pyrightconfig.tests.json`) — private import/usage and
|
||||
not-required TypedDict access allowed; runs as the `pyright` checker in `test-typing`.
|
||||
- **Samples**: relaxed `basic` profile (`pyrightconfig.samples.json`, with a py310 variant) —
|
||||
runs as the `pyright` checker in `test-typing -S`.
|
||||
|
||||
## Parallel Execution
|
||||
|
||||
The task runner (`scripts/task_runner.py`) executes the cross-product of
|
||||
(package × task) in parallel using ThreadPoolExecutor. Single items run
|
||||
in-process with streaming output.
|
||||
|
||||
## CI Workflow
|
||||
|
||||
CI splits into 4 parallel jobs:
|
||||
1. **Pre-commit hooks** — lightweight hooks (SKIP=poe-check)
|
||||
2. **Package checks** — syntax/pyright (source) via check-packages
|
||||
3. **Samples & markdown** — `check -S` plus `markdown-code-lint`
|
||||
4. **Test Typing** — change-detected mypy/pyrefly/ty over tests (`ci-test-typing`)
|
||||
+127
@@ -0,0 +1,127 @@
|
||||
---
|
||||
name: python-development
|
||||
description: >
|
||||
Coding standards, conventions, and patterns for developing Python code in the
|
||||
Agent Framework repository. Use this when writing or modifying Python source
|
||||
files in the python/ directory.
|
||||
---
|
||||
|
||||
# Python Development Standards
|
||||
|
||||
## File Header
|
||||
|
||||
Every `.py` file must start with:
|
||||
|
||||
```python
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
```
|
||||
|
||||
## Type Annotations
|
||||
|
||||
- Always specify return types and parameter types
|
||||
- Use `Type | None` instead of `Optional[Type]`
|
||||
- Use `from __future__ import annotations` to enable postponed evaluation
|
||||
- Use suffix `T` for TypeVar names: `ChatResponseT = TypeVar("ChatResponseT", bound=ChatResponse)`
|
||||
- Use `Mapping` instead of `MutableMapping` for read-only input parameters
|
||||
- Prefer `# type: ignore[...]` over unnecessary casts, or `isinstance` checks, when these are internally called and executed methods
|
||||
But make sure the ignore is specific for both mypy and pyright so that we don't miss other mistakes
|
||||
- Internal private helpers may be used across `agent_framework*` modules when intentional; use a targeted
|
||||
`# pyright: ignore[reportPrivateUsage]` instead of making the helper public just to satisfy pyright.
|
||||
- Do not add trivial pass-through or one-line helper functions solely to appease typing. Prefer targeted ignores,
|
||||
casts, or clearer annotations over adding runtime overhead without a design benefit.
|
||||
|
||||
## Function Parameters
|
||||
|
||||
- Positional parameters: up to 3 fully expected parameters
|
||||
- Use keyword-only arguments (after `*`) for optional parameters
|
||||
- Provide string-based overrides to avoid requiring extra imports:
|
||||
|
||||
```python
|
||||
def create_agent(name: str, tool_mode: Literal['auto', 'required', 'none'] | ChatToolMode) -> Agent:
|
||||
if isinstance(tool_mode, str):
|
||||
tool_mode = ChatToolMode(tool_mode)
|
||||
```
|
||||
|
||||
- Avoid shadowing built-ins (use `next_handler` instead of `next`)
|
||||
- Avoid `**kwargs` unless needed for subclass extensibility; prefer named parameters
|
||||
|
||||
## Docstrings
|
||||
|
||||
Use Google-style docstrings for all public APIs:
|
||||
|
||||
```python
|
||||
def equal(arg1: str, arg2: str) -> bool:
|
||||
"""Compares two strings and returns True if they are the same.
|
||||
|
||||
Args:
|
||||
arg1: The first string to compare.
|
||||
arg2: The second string to compare.
|
||||
|
||||
Returns:
|
||||
True if the strings are the same, False otherwise.
|
||||
|
||||
Raises:
|
||||
ValueError: If one of the strings is empty.
|
||||
"""
|
||||
```
|
||||
|
||||
- Always document Agent Framework specific exceptions
|
||||
- Explicitly use `Keyword Args` when applicable
|
||||
- Only document standard Python exceptions when the condition is non-obvious
|
||||
|
||||
## Import Structure
|
||||
|
||||
```python
|
||||
# Core
|
||||
from agent_framework import Agent, Message, tool
|
||||
|
||||
# Components
|
||||
from agent_framework.observability import enable_sensitive_telemetry
|
||||
|
||||
# Connectors (lazy-loaded)
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
from agent_framework.foundry import FoundryChatClient
|
||||
```
|
||||
|
||||
## Public API and Exports
|
||||
|
||||
In `__init__.py` files that define package-level public APIs, use direct re-export imports plus an explicit
|
||||
`__all__`. Avoid identity aliases like `from ._agents import Agent as Agent`, and avoid
|
||||
`from module import *`.
|
||||
|
||||
Do not define `__all__` in internal non-`__init__.py` modules. Exception: modules intentionally exposed as a
|
||||
public import surface (for example, `agent_framework.observability`) should define `__all__`.
|
||||
|
||||
```python
|
||||
__all__ = ["Agent", "Message", "ChatResponse"]
|
||||
|
||||
from ._agents import Agent
|
||||
from ._types import Message, ChatResponse
|
||||
```
|
||||
|
||||
Special case: the root `agent_framework/__init__.py` uses lazy runtime exports. For root public API changes:
|
||||
- Add the symbol to `_LAZY_MODULE_EXPORTS` and keep `_LAZY_EXPORTS` derived from it.
|
||||
- Keep the explicit runtime `__all__` synchronized; it is still required for `from agent_framework import *`.
|
||||
- Add the same public symbol to `agent_framework/__init__.pyi` so pyright, mypy, and editors see the typed surface.
|
||||
- Put runtime deprecation behavior in the owning module via that module's `__getattr__`; avoid root-level
|
||||
special-case branches for individual deprecated exports.
|
||||
- Identity aliases are appropriate in `.pyi` stubs because they mark re-exported names for type checkers; avoid them
|
||||
in runtime `.py` modules unless there is a specific compatibility reason.
|
||||
|
||||
## Performance Guidelines
|
||||
|
||||
- Cache expensive computations (e.g., JSON schema generation)
|
||||
- Prefer `match/case` on `.type` attribute over `isinstance()` in hot paths
|
||||
- Avoid redundant serialization — compute once, reuse
|
||||
|
||||
## Style
|
||||
|
||||
- Line length: 120 characters
|
||||
- Format only files you changed, not the entire codebase
|
||||
- Prefer attributes over inheritance when parameters are mostly the same
|
||||
- Async by default — assume everything is asynchronous
|
||||
|
||||
## Naming Conventions for Connectors
|
||||
|
||||
- `_prepare_<object>_for_<purpose>` for methods that prepare data for external services
|
||||
- `_parse_<object>_from_<source>` for methods that process data from external services
|
||||
@@ -0,0 +1,236 @@
|
||||
---
|
||||
name: python-feature-lifecycle
|
||||
description: >
|
||||
Guidance for package and feature lifecycle in the Agent Framework Python
|
||||
codebase, including stage meanings, feature-stage decorators, feature enums,
|
||||
and how to move APIs from one stage to the next.
|
||||
---
|
||||
|
||||
# Python Feature Lifecycle
|
||||
|
||||
## Two lifecycle levels
|
||||
|
||||
Agent Framework uses lifecycle at two different levels:
|
||||
|
||||
1. **Package lifecycle** — the maturity of the package as a whole
|
||||
2. **Feature lifecycle** — the maturity of a specific API or feature inside that package
|
||||
|
||||
These are related, but they are **not the same thing**.
|
||||
|
||||
- The **package stage is the default** for everything in the package.
|
||||
- **Feature-stage decorators are only for exceptions** when a feature is behind the package's default stage.
|
||||
- Do **not** decorate every class or function just because the package is experimental or release candidate.
|
||||
|
||||
### Important default
|
||||
|
||||
If a package is still in **beta / experimental preview**, all public APIs in that package are experimental by default.
|
||||
|
||||
- Do **not** add `@experimental(...)` everywhere in that package.
|
||||
- The package stage already communicates that default.
|
||||
|
||||
Once a package moves forward, you can keep individual features behind:
|
||||
|
||||
- If a package moves to **release candidate**, a feature may remain **experimental**
|
||||
- If a package moves to **released / GA**, a feature may remain **experimental** or **release candidate**
|
||||
|
||||
That is the main use case for feature-stage decorators.
|
||||
|
||||
## The four stages
|
||||
|
||||
### 1. Experimental
|
||||
|
||||
Use for features that are still unstable and may change or be removed without notice.
|
||||
|
||||
Feature-level code pattern:
|
||||
|
||||
```python
|
||||
from ._feature_stage import ExperimentalFeature, experimental
|
||||
|
||||
|
||||
@experimental(feature_id=ExperimentalFeature.MY_FEATURE)
|
||||
class MyFeature:
|
||||
...
|
||||
```
|
||||
|
||||
Behavior:
|
||||
|
||||
- Adds an experimental warning block to the docstring
|
||||
- Records feature metadata on the decorated object
|
||||
- Emits a runtime warning the first time the feature is used (once per feature by default)
|
||||
|
||||
Enum setup:
|
||||
|
||||
- Add an all-caps member to `ExperimentalFeature`
|
||||
- Reuse the same feature ID across all APIs that belong to the same conceptual feature
|
||||
|
||||
### 2. Release candidate
|
||||
|
||||
Use for features that are nearly stable but may still receive small refinements before GA.
|
||||
|
||||
Feature-level code pattern:
|
||||
|
||||
```python
|
||||
from ._feature_stage import ReleaseCandidateFeature, release_candidate
|
||||
|
||||
|
||||
@release_candidate(feature_id=ReleaseCandidateFeature.MY_FEATURE)
|
||||
class MyFeature:
|
||||
...
|
||||
```
|
||||
|
||||
Behavior:
|
||||
|
||||
- Adds a release-candidate note to the docstring
|
||||
- Records feature metadata on the decorated object
|
||||
- Does **not** emit the experimental warning
|
||||
|
||||
Enum setup:
|
||||
|
||||
- Add an all-caps member to `ReleaseCandidateFeature`
|
||||
|
||||
### 3. Released
|
||||
|
||||
Use for stable GA APIs.
|
||||
|
||||
Code pattern:
|
||||
|
||||
- **No feature-stage decorator**
|
||||
- **No entry** in `ExperimentalFeature`
|
||||
- **No entry** in `ReleaseCandidateFeature`
|
||||
|
||||
If a feature is fully released, remove any stage-specific feature annotation.
|
||||
|
||||
### 4. Deprecated
|
||||
|
||||
Use for APIs that still exist but should not be used for new code.
|
||||
|
||||
Code pattern:
|
||||
|
||||
```python
|
||||
import sys
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from warnings import deprecated # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import deprecated # type: ignore # pragma: no cover
|
||||
|
||||
|
||||
@deprecated("MyOldFeature is deprecated. Use MyNewFeature instead.")
|
||||
class MyOldFeature:
|
||||
...
|
||||
```
|
||||
|
||||
Behavior:
|
||||
|
||||
- Uses the repository's version-conditional deprecation import pattern
|
||||
- Should describe what to use instead
|
||||
|
||||
Deprecated APIs should not also carry feature-stage decorators.
|
||||
|
||||
## Expected decorators by stage
|
||||
|
||||
| Feature stage | Expected annotation |
|
||||
| --- | --- |
|
||||
| Experimental | `@experimental(feature_id=ExperimentalFeature.X)` |
|
||||
| Release candidate | `@release_candidate(feature_id=ReleaseCandidateFeature.X)` |
|
||||
| Released | No feature-stage decorator |
|
||||
| Deprecated | `@deprecated("...")` |
|
||||
|
||||
## Feature enums
|
||||
|
||||
The feature enums are the inventory of currently staged features:
|
||||
|
||||
- `ExperimentalFeature`
|
||||
- `ReleaseCandidateFeature`
|
||||
|
||||
Guidance:
|
||||
|
||||
- Use one enum member per conceptual feature, not per class
|
||||
- Ideally, an ADR already defines the overall feature boundary and therefore the feature ID that staged APIs for that feature should reuse
|
||||
- Keep feature IDs all caps
|
||||
- Reuse the same member across related APIs for the same feature
|
||||
- Remove enum members when the feature no longer belongs to that stage
|
||||
- Treat these enums as **current-stage inventories**, not as a stable consumer introspection API
|
||||
|
||||
Minimal consumer guidance:
|
||||
|
||||
- Treat `__feature_stage__` and `__feature_id__` as optional staged metadata, not as stable contracts
|
||||
- Use `getattr(obj, "__feature_stage__", None)` and `getattr(obj, "__feature_id__", None)` rather than direct attribute access
|
||||
- Treat missing metadata as "no explicit feature-stage annotation"
|
||||
- For warning filters while a feature is staged, match the literal feature ID string
|
||||
- Do **not** rely on `ExperimentalFeature.X`, `ReleaseCandidateFeature.X`, or the continued presence of `__feature_id__` after a feature moves stages or is released
|
||||
|
||||
For consumers, the enums are also re-exported from `agent_framework`.
|
||||
|
||||
For internal implementation code inside `agent_framework`, continue to import the enums and decorators from `._feature_stage`.
|
||||
|
||||
## Package stage vs feature stage
|
||||
|
||||
Use the following rules:
|
||||
|
||||
### Package is experimental / beta
|
||||
|
||||
- All public APIs are experimental by default
|
||||
- Do **not** add feature-stage decorators just to restate that
|
||||
- Only introduce feature-level annotations later if the package advances first
|
||||
|
||||
### Package is release candidate
|
||||
|
||||
- All public APIs are RC by default
|
||||
- Do **not** decorate everything
|
||||
- Add `@experimental(...)` only for features that are intentionally still behind the package
|
||||
|
||||
### Package is released / GA
|
||||
|
||||
- All public APIs are released by default
|
||||
- Add `@experimental(...)` or `@release_candidate(...)` only for features still being held back
|
||||
|
||||
## Moving a feature from one stage to the next
|
||||
|
||||
### Experimental -> Release candidate
|
||||
|
||||
1. Move the feature ID from `ExperimentalFeature` to `ReleaseCandidateFeature`
|
||||
2. Replace `@experimental(...)` with `@release_candidate(...)`
|
||||
3. Update any tests or docs that mention the old stage
|
||||
|
||||
### Experimental -> Released
|
||||
|
||||
1. Remove `@experimental(...)`
|
||||
2. Remove the feature from `ExperimentalFeature`
|
||||
3. Do not add a replacement feature-stage decorator
|
||||
|
||||
### Release candidate -> Released
|
||||
|
||||
1. Remove `@release_candidate(...)`
|
||||
2. Remove the feature from `ReleaseCandidateFeature`
|
||||
3. Leave the API undecorated
|
||||
|
||||
### Any stage -> Deprecated
|
||||
|
||||
1. Remove any feature-stage decorator
|
||||
2. Remove the feature from the stage enum
|
||||
3. Add `@deprecated("...")`
|
||||
4. Update docs/tests to reflect the replacement path
|
||||
|
||||
## Promotion guidance
|
||||
|
||||
Features do **not** have to pass through every stage.
|
||||
|
||||
- It is usually a good idea to move features in order when that reflects reality
|
||||
- But it is completely acceptable to go **experimental -> released**
|
||||
- Do **not** force a feature through release candidate if there is no real RC period
|
||||
|
||||
Likewise, when a package advances, do not automatically move every feature with it.
|
||||
|
||||
- Promote features based on actual readiness
|
||||
- Keep lagging features explicitly marked only when they are behind the package default
|
||||
|
||||
## Practical rules of thumb
|
||||
|
||||
- **Package default first, feature exceptions second**
|
||||
- **Do not decorate everything in preview packages**
|
||||
- **Do not double-annotate members of an already-staged class**
|
||||
- **Use enums only for currently staged features**
|
||||
- **Do not treat stage enums as a compatibility contract**
|
||||
- **Treat `__feature_stage__` and `__feature_id__` as optional metadata; use `getattr`**
|
||||
- **Remove stage annotations once a feature is released or deprecated**
|
||||
@@ -0,0 +1,252 @@
|
||||
---
|
||||
name: python-package-management
|
||||
description: >
|
||||
Guide for managing packages in the Agent Framework Python monorepo, including
|
||||
creating new connector packages, versioning, and the lazy-loading pattern.
|
||||
Use this when adding, modifying, or releasing packages.
|
||||
---
|
||||
|
||||
# Python Package Management
|
||||
|
||||
## Monorepo Structure
|
||||
|
||||
```
|
||||
python/
|
||||
├── pyproject.toml # Root package (agent-framework)
|
||||
├── packages/
|
||||
│ ├── core/ # agent-framework-core (main package)
|
||||
│ ├── foundry/ # agent-framework-foundry
|
||||
│ ├── anthropic/ # agent-framework-anthropic
|
||||
│ └── ... # Other connector packages
|
||||
```
|
||||
|
||||
- `agent-framework-core` contains core abstractions and OpenAI/Azure OpenAI built-in
|
||||
- Provider packages extend core with specific integrations
|
||||
- Root `agent-framework` depends on `agent-framework-core[all]`
|
||||
|
||||
## Dependency Management
|
||||
|
||||
Uses [uv](https://github.com/astral-sh/uv) for dependency management and
|
||||
[poethepoet](https://github.com/nat-n/poethepoet) for task automation.
|
||||
|
||||
```bash
|
||||
# Full setup (venv + install + prek hooks)
|
||||
uv run poe setup
|
||||
|
||||
# Install dependencies from lockfile (frozen resolution with prerelease policy)
|
||||
uv run poe install
|
||||
|
||||
# Create venv with specific Python version
|
||||
uv run poe venv --python 3.12
|
||||
|
||||
# Intentionally upgrade a specific dependency to reduce lockfile conflicts
|
||||
uv lock --upgrade-package <dependency-name> && uv run poe install
|
||||
|
||||
# Refresh exact development dependency-group pins, lockfile, and validation in one run
|
||||
uv run poe upgrade-dev-dependencies
|
||||
|
||||
# First, run workspace-wide lower/upper compatibility gates
|
||||
uv run poe validate-dependency-bounds-test
|
||||
# Defaults to --package "*"; pass a package to scope test mode
|
||||
uv run poe validate-dependency-bounds-test --package core
|
||||
|
||||
# Then expand bounds for one dependency in the target package
|
||||
uv run poe validate-dependency-bounds-project --mode both --package core --dependency "<dependency-name>"
|
||||
|
||||
# Repo-wide automation can reuse the same task
|
||||
uv run poe validate-dependency-bounds-project --mode upper --package "*"
|
||||
|
||||
# Add a dependency to one project and run both validators for that project/dependency
|
||||
uv run poe add-dependency-and-validate-bounds --package core --dependency "<dependency-spec>"
|
||||
```
|
||||
|
||||
### Dependency Bound Notes
|
||||
|
||||
- Stable dependencies (`>=1.0`) should typically be bounded as `>=<known-good>,<next-major>`.
|
||||
- Prerelease (`dev`/`a`/`b`/`rc`) and `<1.0` dependencies should use hard bounds with an explicit upper cap (avoid open-ended ranges).
|
||||
- For `<1.0` dependencies, prefer the broadest validated range the package can really support. That may be a patch line, a minor line, or multiple minor lines when checks/tests show the broader lane is compatible.
|
||||
- Prefer supporting multiple majors when practical; if APIs diverge across supported majors, use version-conditional imports/paths.
|
||||
- For dependency changes, run workspace-wide bound gates first, then `validate-dependency-bounds-project --mode both` for the target package/dependency to keep minimum and maximum constraints current. The same task can also drive repo-wide upper-bound automation by using `--package "*"` and omitting `--dependency`.
|
||||
- Prefer targeted lock updates with `uv lock --upgrade-package <dependency-name>` to reduce `uv.lock` merge conflicts.
|
||||
- Use `add-dependency-and-validate-bounds` for package-scoped dependency additions plus bound validation in one command.
|
||||
- Keep shared tooling and source/type-check support in the root or package `dev` group. Put package-specific test
|
||||
fixtures in a `test` group, and use a feature-named group for local-only executable dependencies that cannot be
|
||||
expressed in published runtime metadata.
|
||||
- Use `upgrade-dev-dependencies` for repo-wide development dependency refreshes; it repins exact dependencies
|
||||
across development groups, refreshes `uv.lock`, and reruns `check`, `typing`, and `test`.
|
||||
|
||||
## Lazy Loading Pattern
|
||||
|
||||
### Root core API
|
||||
|
||||
The root `agent_framework` package is a lazy public API surface:
|
||||
|
||||
- Runtime exports live in `packages/core/agent_framework/__init__.py`.
|
||||
- Typing/editor exports live in `packages/core/agent_framework/__init__.pyi`.
|
||||
- Add or move root exports in `_LAZY_MODULE_EXPORTS`, keep the explicit runtime `__all__` in sync, and add the same
|
||||
symbol to the `.pyi` file.
|
||||
- Keep deprecation behavior in the owning module (for example, a module-level `__getattr__` that warns and returns
|
||||
the deprecated alias). Do not add one-off deprecated-symbol branches to root `__getattr__`.
|
||||
- Validate root API changes with `uv run poe syntax -P core`, `uv run poe pyright -P core`, and import smoke tests
|
||||
for both `from agent_framework import <symbol>` and `from agent_framework import *`.
|
||||
|
||||
### Provider namespaces
|
||||
|
||||
Provider folders in core use `__getattr__` to lazy load from connector packages:
|
||||
|
||||
```python
|
||||
# In agent_framework/foundry/__init__.py
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"FoundryChatClient": ("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"Install it with: pip install {package_name}"
|
||||
) from exc
|
||||
```
|
||||
|
||||
## Adding a New Connector Package
|
||||
|
||||
**Important:** Do not create a new package unless approved by the core team.
|
||||
|
||||
Every new package starts as `alpha`.
|
||||
|
||||
### Alpha package checklist
|
||||
|
||||
1. Create directory under `packages/` (e.g., `packages/my-connector/`)
|
||||
2. Add the package to `tool.uv.sources` in root `pyproject.toml`
|
||||
3. Set the package version to the alpha pattern: `1.0.0a<date>`
|
||||
4. Set the package classifier to `Development Status :: 3 - Alpha`
|
||||
5. Include samples inside the package (e.g., `packages/my-connector/samples/`)
|
||||
6. Do **NOT** add to `[all]` extra in `packages/core/pyproject.toml`
|
||||
7. Do **NOT** create lazy loading in core yet
|
||||
8. Add the package to `python/PACKAGE_STATUS.md` and keep that file updated when packages are added,
|
||||
removed, renamed, or promoted. If the package exposes individually staged APIs, keep the feature list
|
||||
there current too.
|
||||
|
||||
Recommended dependency workflow during connector implementation:
|
||||
|
||||
1. Add the dependency to the target package:
|
||||
`uv run poe add-dependency-to-project --package core --dependency "<dependency-spec>"`
|
||||
2. Implement connector code and tests.
|
||||
3. Validate dependency bounds for that package/dependency:
|
||||
`uv run poe validate-dependency-bounds-project --mode both --package core --dependency "<dependency-name>"`
|
||||
4. If the package has meaningful tests/checks that validate dependency compatibility, you can use the add + validation flow in one command:
|
||||
`uv run poe add-dependency-and-validate-bounds --package core --dependency "<dependency-spec>"`
|
||||
If compatibility checks are not in place yet, add the dependency first, then implement tests before running bound validation.
|
||||
|
||||
### Promotion path
|
||||
|
||||
Promotion work is not isolated to the package being promoted. If a promotion changes dependency
|
||||
metadata for downstream packages, also update the dependent packages' own versions so they publish
|
||||
new metadata alongside the promoted dependency bounds.
|
||||
Apply the internal package dependency update rules from the versioning section below during
|
||||
promotions as well as standalone version update work.
|
||||
|
||||
#### Alpha -> Beta
|
||||
|
||||
Move a package to `beta` when it is stable enough to be part of the main install surface.
|
||||
|
||||
1. Update the package version to the beta pattern: `1.0.0b<date>`
|
||||
2. Update the classifier to `Development Status :: 4 - Beta`
|
||||
3. Add the package to `[all]` in `packages/core/pyproject.toml`
|
||||
4. Move samples to the root `samples/` tree and remove package-local samples
|
||||
5. Create or update the relevant lazy-loading namespace in core when the package belongs under one
|
||||
6. Update `python/PACKAGE_STATUS.md`
|
||||
|
||||
After `alpha`, there should be no samples left inside a package folder.
|
||||
|
||||
#### Beta -> RC
|
||||
|
||||
Move a package to `rc` when its API is close to the final released shape.
|
||||
|
||||
1. Update the package version to the release-candidate pattern: `1.0.0rc<number>`
|
||||
2. Keep the classifier at `Development Status :: 4 - Beta` because PyPI does not have a separate
|
||||
release-candidate classifier
|
||||
3. Keep the package in `core[all]`
|
||||
4. Keep samples only in the root `samples/` tree
|
||||
5. Update `python/PACKAGE_STATUS.md` to show the package as `rc`
|
||||
|
||||
#### RC -> Released
|
||||
|
||||
Move a package to `released` when it no longer carries a prerelease qualifier.
|
||||
|
||||
1. Update the package version to the stable pattern: `1.0.0`
|
||||
2. Update the classifier to `Development Status :: 5 - Production/Stable`
|
||||
3. Keep the package in `core[all]`
|
||||
4. Keep samples only in the root `samples/` tree
|
||||
5. Update `python/PACKAGE_STATUS.md` to show the package as `released`
|
||||
6. Update all `README.md` files that install that package with
|
||||
`pip install agent-framework-... --pre` so they use `pip install agent-framework-...` without
|
||||
the `--pre` suffix
|
||||
|
||||
## Versioning
|
||||
|
||||
### Internal package dependency updates
|
||||
|
||||
- If package A depends on package B within this repository, only update package A's dependency
|
||||
declaration when the work on package B actually affects package A.
|
||||
- If package A does not need anything from the package B change, leave package A's dependency
|
||||
declaration unchanged.
|
||||
- If package A does need something from the package B change, update package A's dependency
|
||||
declaration to the version or versioning scheme that matches what package A now requires.
|
||||
- If package B is promoted to a different lifecycle stage, update package A's dependency
|
||||
declaration to the new versioning scheme for package B even when the only change is the stage
|
||||
transition itself.
|
||||
- Use this guidance both for ordinary version updates and for package promotion work.
|
||||
|
||||
- All non-core packages declare a lower bound on `agent-framework-core`
|
||||
- When core version bumps with breaking changes, update the lower bound in all packages
|
||||
- Non-core packages version independently; only raise core bound when using new core APIs
|
||||
- If promoting a package changes a dependent package's published dependency metadata, bump the
|
||||
dependent package's own version in the correct lifecycle pattern for its current stage
|
||||
- Lifecycle version patterns:
|
||||
- `alpha`: `1.0.0a<date>` where `<date>` is the current Pacific (US west coast) `YYMMDD`
|
||||
- `beta`: `1.0.0b<date>` where `<date>` is the current Pacific (US west coast) `YYMMDD`
|
||||
- `rc`: `1.0.0rc<number>` where `<number>` increments only when the package has changes
|
||||
- `released`: `X.Y.Z` using semver per package
|
||||
- For alpha/beta date stamps, use the current Pacific date as the cutoff, not UTC and not the user's local
|
||||
timezone. Same-Pacific-day re-cuts use a `.postN` suffix. Honor an explicit user-provided date over this
|
||||
default.
|
||||
- Keep the `Development Status` classifier in `pyproject.toml` aligned with the lifecycle stage:
|
||||
- `alpha` -> `Development Status :: 3 - Alpha`
|
||||
- `beta` -> `Development Status :: 4 - Beta`
|
||||
- `rc` -> `Development Status :: 4 - Beta`
|
||||
- `released` -> `Development Status :: 5 - Production/Stable`
|
||||
- See the PyPI classifier list for the available classifier values:
|
||||
`https://pypi.org/classifiers/`
|
||||
|
||||
## Installation Options
|
||||
|
||||
```bash
|
||||
pip install agent-framework-core # Core only
|
||||
pip install agent-framework-core[all] # Core + all connectors
|
||||
pip install agent-framework # Same as core[all]
|
||||
pip install agent-framework-foundry # Specific connector (pulls in core)
|
||||
```
|
||||
|
||||
## Maintaining Documentation
|
||||
|
||||
When changing a package, check if its `AGENTS.md` needs updates:
|
||||
- Adding/removing/renaming public classes or functions
|
||||
- Changing the package's purpose or architecture
|
||||
- Modifying import paths or usage patterns
|
||||
|
||||
Keep `python/PACKAGE_STATUS.md` updated when:
|
||||
- A package is added, removed, renamed, or promoted between lifecycle stages
|
||||
- A package starts or stops exposing individually staged experimental or release-candidate APIs
|
||||
|
||||
When a package adds, removes, or renames environment variables, update the related documentation in the same
|
||||
change:
|
||||
- The package's `README.md` for package-level configuration/env var guidance
|
||||
- `samples/README.md` if the package is included in `packages/core/pyproject.toml` `[all]` and the env var is
|
||||
part of the consolidated package env-var inventory
|
||||
- Any affected sample/package-local `.env.example`, `.env.template`, or sample README files when sample setup
|
||||
changes alongside the package
|
||||
+157
@@ -0,0 +1,157 @@
|
||||
---
|
||||
name: python-testing
|
||||
description: >
|
||||
Guidelines for writing and running tests in the Agent Framework Python
|
||||
codebase. Use this when creating, modifying, or running tests.
|
||||
---
|
||||
|
||||
# Python Testing
|
||||
|
||||
We strive for at least 85% test coverage across the codebase, with a focus on core packages and critical paths. Tests should be fast, reliable, and maintainable.
|
||||
When adding new code, check that the relevant sections of the codebase are covered by tests, and add new tests as needed. When modifying existing code, update or add tests to cover the changes.
|
||||
We run tests in two stages, for a PR each commit is tested with unit tests only (using `-m "not integration"`), and the full suite including integration tests is run when merging.
|
||||
|
||||
## Running Tests
|
||||
|
||||
```bash
|
||||
# Run tests for all packages in parallel
|
||||
uv run poe test
|
||||
|
||||
# Run tests for a specific workspace package
|
||||
uv run poe test -P core
|
||||
|
||||
# Run all selected tests in a single pytest invocation
|
||||
uv run poe test -A
|
||||
|
||||
# With coverage
|
||||
uv run poe test -A -C
|
||||
uv run poe test -P core -C
|
||||
|
||||
# Run only unit tests (exclude integration tests)
|
||||
uv run poe test -A -m "not integration"
|
||||
|
||||
# Run only integration tests
|
||||
uv run poe test -A -m integration
|
||||
```
|
||||
|
||||
Direct package execution still works when you need it:
|
||||
|
||||
```bash
|
||||
uv run --directory packages/core poe test
|
||||
```
|
||||
|
||||
## Test Configuration
|
||||
|
||||
- **Async mode**: `asyncio_mode = "auto"` is enabled — do NOT use `@pytest.mark.asyncio`, but do mark tests with `async def` and use `await` for async calls
|
||||
- **Timeout**: Default 60 seconds per test
|
||||
- **Import mode**: `importlib` for cross-package isolation
|
||||
- **Parallelization**: Large packages (core, ag-ui, orchestrations, anthropic) use `pytest-xdist` (`-n auto --dist worksteal`) in their `poe test` task. The aggregate `uv run poe test -A` sweep also uses xdist across the selected packages.
|
||||
|
||||
## Test Directory Structure
|
||||
|
||||
Test directories must NOT contain `__init__.py` files.
|
||||
|
||||
Non-core packages must place tests in a uniquely-named subdirectory:
|
||||
|
||||
```
|
||||
packages/anthropic/
|
||||
├── tests/
|
||||
│ └── anthropic/ # Unique subdirectory matching package name
|
||||
│ ├── conftest.py
|
||||
│ └── test_client.py
|
||||
```
|
||||
|
||||
Core package can use `tests/` directly with topic subdirectories:
|
||||
|
||||
```
|
||||
packages/core/
|
||||
├── tests/
|
||||
│ ├── conftest.py
|
||||
│ ├── core/
|
||||
│ │ └── test_agents.py
|
||||
│ └── openai/
|
||||
│ └── test_client.py
|
||||
```
|
||||
|
||||
## Fixture Guidelines
|
||||
|
||||
- Use `conftest.py` for shared fixtures within a test directory
|
||||
- Before adding new fixtures, check if existing ones can be reused or extended
|
||||
- Use descriptive names: `mapper`, `test_request`, `mock_client`
|
||||
|
||||
## File Naming
|
||||
|
||||
- Files starting with `test_` are test files — do not use this prefix for helpers
|
||||
- Use `conftest.py` for shared utilities
|
||||
|
||||
## Integration Tests
|
||||
|
||||
Integration tests require external services (OpenAI, Azure, etc.) and are controlled by three markers:
|
||||
|
||||
1. **`@pytest.mark.flaky`** — marks the test as potentially flaky since it depends on external services
|
||||
2. **`@pytest.mark.integration`** — used for test selection, so integration tests can be included/excluded with `-m integration` / `-m "not integration"`
|
||||
3. **`@skip_if_..._integration_tests_disabled`** decorator — skips the test when the required API keys or service endpoints are missing
|
||||
|
||||
### Adding New Integration Tests
|
||||
|
||||
All three markers must be applied to every new integration test:
|
||||
|
||||
```python
|
||||
@pytest.mark.flaky
|
||||
@pytest.mark.integration
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_chat_completion() -> None:
|
||||
...
|
||||
```
|
||||
|
||||
For test files where all tests are integration tests (e.g., Azure Functions, Durable Task), use the module-level `pytestmark` list:
|
||||
|
||||
```python
|
||||
pytestmark = [
|
||||
pytest.mark.flaky,
|
||||
pytest.mark.integration,
|
||||
pytest.mark.sample("01_single_agent"),
|
||||
pytest.mark.usefixtures("function_app_for_test"),
|
||||
]
|
||||
```
|
||||
|
||||
### CI Workflow
|
||||
|
||||
The merge CI workflow (`python-merge-tests.yml`) splits integration tests into parallel jobs by provider with change-based detection:
|
||||
|
||||
- **Unit tests** — always run all non-integration tests
|
||||
- **OpenAI integration** — runs when `packages/core/agent_framework/openai/` or core infrastructure changes
|
||||
- **Azure OpenAI integration** — runs when `packages/core/agent_framework/azure/` or core changes
|
||||
- **Misc integration** — Anthropic, Ollama, MCP tests; runs when their packages or core change
|
||||
- **Functions integration** — Azure Functions + Durable Task; runs when their packages or core change
|
||||
- **Foundry integration** — runs when `packages/foundry/` or core changes
|
||||
|
||||
Core infrastructure changes (e.g., `_agents.py`, `_types.py`) trigger all integration test jobs. Scheduled and manual runs always execute all jobs.
|
||||
|
||||
### Keeping CI Workflows in Sync
|
||||
|
||||
Two workflow files define the same set of parallel test jobs:
|
||||
|
||||
- **`python-merge-tests.yml`** — runs on PRs, merge queue, schedule, and manual dispatch. Uses path-based change detection to skip unaffected integration jobs.
|
||||
- **`python-integration-tests.yml`** — called from the manual integration test orchestrator (`integration-tests-manual.yml`). Always runs all jobs (no path filtering).
|
||||
|
||||
These workflows must be kept in sync. When you add, remove, or modify a test job, update **both** files. The job structure, pytest commands, and xdist flags should match between them. The only difference is that `python-merge-tests.yml` has path filters and conditional job execution, while `python-integration-tests.yml` does not.
|
||||
|
||||
### Updating the CI When Adding Integration Tests for a New Provider
|
||||
|
||||
When adding integration tests for a new provider package, you must update **both** `python-merge-tests.yml` and `python-integration-tests.yml`:
|
||||
|
||||
1. **Add a path filter** for the new provider in the `paths-filter` job in `python-merge-tests.yml` so the CI knows which file changes should trigger those tests.
|
||||
2. **Add the test job to both workflow files** — either add them to the existing `python-tests-misc-integration` job, or create a dedicated job if the provider:
|
||||
- Has a large number of integration tests
|
||||
- Requires special infrastructure setup (emulators, Docker containers, etc.)
|
||||
- Has long-running tests that would slow down the misc job
|
||||
|
||||
The `python-tests-misc-integration` job is intended for small integration test suites that don't need dedicated infrastructure. When a provider's integration tests grow large or gain special requirements, split them out into their own job (like `python-tests-functions` was split out for Azure Functions + Durable Task).
|
||||
|
||||
## Best Practices
|
||||
|
||||
- Run only related tests, not the entire suite
|
||||
- Review existing tests to understand coding style before creating new ones
|
||||
- Use print statements for debugging, then remove them when done
|
||||
- Resolve all errors and warnings before committing
|
||||
@@ -0,0 +1,69 @@
|
||||
fail_fast: true
|
||||
exclude: ^scripts/
|
||||
repos:
|
||||
- repo: https://github.com/pre-commit/pre-commit-hooks
|
||||
rev: v6.0.0
|
||||
hooks:
|
||||
- id: check-toml
|
||||
name: Check TOML files
|
||||
files: \.toml$
|
||||
exclude: ^packages/lab/cookiecutter-agent-framework-lab/
|
||||
- id: check-yaml
|
||||
name: Check YAML files
|
||||
files: \.yaml$
|
||||
- id: check-json
|
||||
name: Check JSON files
|
||||
files: \.json$
|
||||
exclude: ^.*\.vscode\/.*|^demos/samples/chatkit-integration/frontend/(tsconfig.*\.json|package-lock\.json)$
|
||||
- id: end-of-file-fixer
|
||||
name: Fix End of File
|
||||
files: \.py$
|
||||
exclude: ^packages/lab/cookiecutter-agent-framework-lab/
|
||||
- id: mixed-line-ending
|
||||
name: Check Mixed Line Endings
|
||||
files: \.py$
|
||||
exclude: ^packages/lab/cookiecutter-agent-framework-lab/
|
||||
- id: trailing-whitespace
|
||||
name: Trim Trailing Whitespace
|
||||
exclude: ^packages/lab/cookiecutter-agent-framework-lab/
|
||||
- id: check-merge-conflict
|
||||
name: Check Merge Conflicts
|
||||
- id: detect-private-key
|
||||
name: Detect Private Keys
|
||||
- id: check-added-large-files
|
||||
name: Check Added Large Files
|
||||
- id: no-commit-to-branch
|
||||
name: Protect main branch
|
||||
args: [--branch, main]
|
||||
- id: check-ast
|
||||
name: Check Valid Python Samples
|
||||
types: ["python"]
|
||||
exclude: ^packages/lab/cookiecutter-agent-framework-lab/
|
||||
- repo: https://github.com/asottile/pyupgrade
|
||||
rev: v3.21.2
|
||||
hooks:
|
||||
- id: pyupgrade
|
||||
name: Upgrade Python syntax
|
||||
args: [--py310-plus]
|
||||
exclude: ^packages/lab/cookiecutter-agent-framework-lab/
|
||||
- repo: local
|
||||
hooks:
|
||||
- id: poe-check
|
||||
name: Run checks through Poe
|
||||
entry: uv run python scripts/workspace_poe_tasks.py prek-check
|
||||
language: system
|
||||
- repo: https://github.com/PyCQA/bandit
|
||||
rev: 1.9.4
|
||||
hooks:
|
||||
- id: bandit
|
||||
name: Bandit Security Checks
|
||||
args: ["-c", "pyproject.toml"]
|
||||
additional_dependencies: ["bandit[toml]"]
|
||||
- repo: https://github.com/astral-sh/uv-pre-commit
|
||||
# uv version.
|
||||
rev: 0.11.6
|
||||
hooks:
|
||||
# Update the uv lockfile
|
||||
- id: uv-lock
|
||||
name: Update uv lockfile
|
||||
files: pyproject.toml
|
||||
Vendored
+34
@@ -0,0 +1,34 @@
|
||||
{
|
||||
// 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",
|
||||
"justMyCode": false
|
||||
},
|
||||
{
|
||||
"name": "AG-UI Examples Server",
|
||||
"type": "debugpy",
|
||||
"request": "launch",
|
||||
"module": "agent_framework_ag_ui_examples",
|
||||
"cwd": "${workspaceFolder}/packages/ag-ui",
|
||||
"console": "integratedTerminal",
|
||||
"justMyCode": false
|
||||
},
|
||||
{
|
||||
"name": "Python Attach",
|
||||
"type": "debugpy",
|
||||
"request": "attach",
|
||||
"connect": {
|
||||
"host": "localhost",
|
||||
"port": 5678
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
Vendored
+39
@@ -0,0 +1,39 @@
|
||||
{
|
||||
"cSpell.languageSettings": [
|
||||
{
|
||||
"languageId": "py",
|
||||
"allowCompoundWords": true,
|
||||
"locale": "en-US"
|
||||
}
|
||||
],
|
||||
"[python]": {
|
||||
"editor.codeActionsOnSave": {
|
||||
"source.organizeImports.ruff": "always",
|
||||
"source.fixAll.ruff": "always"
|
||||
},
|
||||
"editor.formatOnSave": true,
|
||||
"editor.formatOnPaste": true,
|
||||
"editor.formatOnType": true,
|
||||
"editor.defaultFormatter": "charliermarsh.ruff"
|
||||
},
|
||||
"python.analysis.autoFormatStrings": true,
|
||||
"python.analysis.importFormat": "relative",
|
||||
"python.analysis.packageIndexDepths": [
|
||||
{
|
||||
"name": "agent_framework",
|
||||
"depth": 2
|
||||
},
|
||||
{
|
||||
"name": "extensions",
|
||||
"depth": 2
|
||||
},
|
||||
{
|
||||
"name": "openai",
|
||||
"depth": 2
|
||||
},
|
||||
{
|
||||
"name": "azure",
|
||||
"depth": 2
|
||||
}
|
||||
]
|
||||
}
|
||||
Vendored
+247
@@ -0,0 +1,247 @@
|
||||
{
|
||||
// See https://go.microsoft.com/fwlink/?LinkId=733558
|
||||
// for the documentation about the tasks.json format
|
||||
"version": "2.0.0",
|
||||
"tasks": [
|
||||
{
|
||||
"label": "Run Checks",
|
||||
"type": "shell",
|
||||
"command": "uv",
|
||||
"args": [
|
||||
"run",
|
||||
"poe",
|
||||
"check"
|
||||
],
|
||||
"problemMatcher": {
|
||||
"owner": "python",
|
||||
"fileLocation": [
|
||||
"relative",
|
||||
"${workspaceFolder}"
|
||||
],
|
||||
"pattern": {
|
||||
"regexp": "^(.*):(\\d+):(\\d+):\\s+(.*)$",
|
||||
"file": 1,
|
||||
"line": 2,
|
||||
"column": 3,
|
||||
"message": 4
|
||||
}
|
||||
},
|
||||
"presentation": {
|
||||
"panel": "shared"
|
||||
}
|
||||
},
|
||||
{
|
||||
"label": "Syntax",
|
||||
"type": "shell",
|
||||
"command": "uv",
|
||||
"args": [
|
||||
"run",
|
||||
"poe",
|
||||
"syntax",
|
||||
],
|
||||
"problemMatcher": {
|
||||
"owner": "python",
|
||||
"fileLocation": [
|
||||
"relative",
|
||||
"${workspaceFolder}"
|
||||
],
|
||||
"pattern": {
|
||||
"regexp": "^(.*):(\\d+):(\\d+):\\s+(.*)$",
|
||||
"file": 1,
|
||||
"line": 2,
|
||||
"column": 3,
|
||||
"message": 4
|
||||
}
|
||||
},
|
||||
"presentation": {
|
||||
"panel": "shared"
|
||||
}
|
||||
},
|
||||
{
|
||||
"label": "Syntax (format only)",
|
||||
"type": "shell",
|
||||
"command": "uv",
|
||||
"args": [
|
||||
"run",
|
||||
"poe",
|
||||
"syntax",
|
||||
"-F",
|
||||
],
|
||||
"problemMatcher": {
|
||||
"owner": "python",
|
||||
"fileLocation": [
|
||||
"relative",
|
||||
"${workspaceFolder}"
|
||||
],
|
||||
"pattern": {
|
||||
"regexp": "^(.*):(\\d+):(\\d+):\\s+(.*)$",
|
||||
"file": 1,
|
||||
"line": 2,
|
||||
"column": 3,
|
||||
"message": 4
|
||||
}
|
||||
},
|
||||
"presentation": {
|
||||
"panel": "shared"
|
||||
}
|
||||
},
|
||||
{
|
||||
"label": "Syntax (check only)",
|
||||
"type": "shell",
|
||||
"command": "uv",
|
||||
"args": [
|
||||
"run",
|
||||
"poe",
|
||||
"syntax",
|
||||
"-C",
|
||||
],
|
||||
"problemMatcher": {
|
||||
"owner": "python",
|
||||
"fileLocation": [
|
||||
"relative",
|
||||
"${workspaceFolder}"
|
||||
],
|
||||
"pattern": {
|
||||
"regexp": "^(.*):(\\d+):(\\d+):\\s+(.*)$",
|
||||
"file": 1,
|
||||
"line": 2,
|
||||
"column": 3,
|
||||
"message": 4
|
||||
}
|
||||
},
|
||||
"presentation": {
|
||||
"panel": "shared"
|
||||
}
|
||||
},
|
||||
{
|
||||
"label": "Mypy",
|
||||
"type": "shell",
|
||||
"command": "uv",
|
||||
"args": [
|
||||
"run",
|
||||
"poe",
|
||||
"mypy",
|
||||
],
|
||||
"problemMatcher": {
|
||||
"owner": "python",
|
||||
"fileLocation": [
|
||||
"relative",
|
||||
"${workspaceFolder}"
|
||||
],
|
||||
"pattern": {
|
||||
"regexp": "^(.*):(\\d+):(\\d+):\\s+(.*)$",
|
||||
"file": 1,
|
||||
"line": 2,
|
||||
"column": 3,
|
||||
"message": 4
|
||||
}
|
||||
},
|
||||
"presentation": {
|
||||
"panel": "shared"
|
||||
}
|
||||
},
|
||||
{
|
||||
"label": "Pyright",
|
||||
"type": "shell",
|
||||
"command": "uv",
|
||||
"args": [
|
||||
"run",
|
||||
"poe",
|
||||
"pyright",
|
||||
],
|
||||
"problemMatcher": {
|
||||
"owner": "python",
|
||||
"fileLocation": [
|
||||
"relative",
|
||||
"${workspaceFolder}"
|
||||
],
|
||||
"pattern": {
|
||||
"regexp": "^(.*):(\\d+):(\\d+):\\s+(.*)$",
|
||||
"file": 1,
|
||||
"line": 2,
|
||||
"column": 3,
|
||||
"message": 4
|
||||
}
|
||||
},
|
||||
"presentation": {
|
||||
"panel": "shared"
|
||||
}
|
||||
},
|
||||
{
|
||||
"label": "Test",
|
||||
"type": "shell",
|
||||
"command": "uv",
|
||||
"args": [
|
||||
"run",
|
||||
"poe",
|
||||
"test",
|
||||
],
|
||||
"problemMatcher": {
|
||||
"owner": "python",
|
||||
"fileLocation": [
|
||||
"relative",
|
||||
"${workspaceFolder}"
|
||||
],
|
||||
"pattern": {
|
||||
"regexp": "^(.*):(\\d+):(\\d+):\\s+(.*)$",
|
||||
"file": 1,
|
||||
"line": 2,
|
||||
"column": 3,
|
||||
"message": 4
|
||||
}
|
||||
},
|
||||
"presentation": {
|
||||
"panel": "shared"
|
||||
}
|
||||
},
|
||||
{
|
||||
"label": "Create Venv",
|
||||
"type": "shell",
|
||||
"command": "uv",
|
||||
"args": [
|
||||
"run",
|
||||
"poe",
|
||||
"venv",
|
||||
"-P",
|
||||
"${input:py_version}"
|
||||
],
|
||||
"presentation": {
|
||||
"reveal": "always",
|
||||
"panel": "new"
|
||||
},
|
||||
"problemMatcher": []
|
||||
},
|
||||
{
|
||||
"label": "Install all dependencies",
|
||||
"type": "shell",
|
||||
"command": "uv",
|
||||
"args": [
|
||||
"run",
|
||||
"poe",
|
||||
"setup",
|
||||
"-P",
|
||||
"${input:py_version}"
|
||||
],
|
||||
"presentation": {
|
||||
"reveal": "always",
|
||||
"panel": "new"
|
||||
},
|
||||
"problemMatcher": []
|
||||
}
|
||||
],
|
||||
"inputs": [
|
||||
{
|
||||
"type": "pickString",
|
||||
"options": [
|
||||
"3.10",
|
||||
"3.11",
|
||||
"3.12",
|
||||
"3.13",
|
||||
"3.14"
|
||||
],
|
||||
"id": "py_version",
|
||||
"description": "Python version",
|
||||
"default": "3.13"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
# AGENTS.md
|
||||
|
||||
Instructions for AI coding agents working in the Python codebase.
|
||||
|
||||
**Key Documentation:**
|
||||
- [DEV_SETUP.md](DEV_SETUP.md) - Development environment setup and available poe tasks
|
||||
- [CODING_STANDARD.md](CODING_STANDARD.md) - Coding standards, docstring format, and performance guidelines
|
||||
- [samples/SAMPLE_GUIDELINES.md](samples/SAMPLE_GUIDELINES.md) - Sample structure and guidelines
|
||||
|
||||
**Agent Skills** (`.github/skills/`) — detailed, task-specific instructions loaded on demand:
|
||||
- `python-development` — coding standards, type annotations, docstrings, logging, performance
|
||||
- `python-testing` — test structure, fixtures, async mode, running tests
|
||||
- `python-code-quality` — linting, formatting, type checking, prek hooks, CI workflow
|
||||
- `python-feature-lifecycle` — package vs feature lifecycle stages, decorators, enums, and promotion guidance
|
||||
- `python-package-management` — monorepo structure, lazy loading, versioning, new packages
|
||||
- `pull-requests` — writing PR descriptions (template) and handling/resolving PR review comments
|
||||
- `agent-framework-py-release` — Python release PR workflow, CHANGELOG-driven package bumps, lifecycle version rules, and dependency-floor validation
|
||||
|
||||
## Maintaining Documentation
|
||||
|
||||
When making changes to a package, check if the following need updates:
|
||||
- The package's `AGENTS.md` file (adding/removing/renaming public APIs, architecture changes, import path changes)
|
||||
- The agent skills in `.github/skills/` if conventions, commands, or workflows change
|
||||
|
||||
At the end of every run, re-read `AGENTS.md` and the relevant skill files and
|
||||
update any guidance that the conversation revealed to be out of date,
|
||||
incomplete, or misleading (renamed files, changed commands, new conventions
|
||||
the user confirmed, etc.). **Before adding a new principle or rule, ask the
|
||||
user whether they want it captured as a durable principle** — do not invent
|
||||
team norms from a single conversation without explicit confirmation.
|
||||
|
||||
## Terminology
|
||||
|
||||
- **Avoid "GA" for Agent Framework code.** Reserve *GA* for hosted services
|
||||
(e.g. "the Foundry service is GA"). For Agent Framework packages, features,
|
||||
and APIs use **"released"** or **"stable"** depending on context — these
|
||||
match the feature-lifecycle stages documented in the
|
||||
`python-feature-lifecycle` skill.
|
||||
|
||||
## Pull Request Description Guidance
|
||||
|
||||
When preparing a PR description:
|
||||
- Follow the repository PR template at `.github/pull_request_template.md` and keep its structure/headings.
|
||||
- Describe the net change relative to `main` (this is implied; do not call it out explicitly as "vs main").
|
||||
- Do not add ad-hoc validation sections (for example, "Validation" or "Tests run"); CI/CD and the template checklist cover validation status.
|
||||
|
||||
## Quick Reference
|
||||
|
||||
Run `uv run poe` from the `python/` directory to see available commands. See [DEV_SETUP.md](DEV_SETUP.md) for detailed usage.
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
python/
|
||||
├── packages/
|
||||
│ ├── core/ # agent-framework-core (main package)
|
||||
│ │ ├── agent_framework/ # Public API exports
|
||||
│ │ └── tests/
|
||||
│ ├── foundry/ # agent-framework-foundry
|
||||
│ ├── anthropic/ # agent-framework-anthropic
|
||||
│ ├── ollama/ # agent-framework-ollama
|
||||
│ └── ... # Other provider packages
|
||||
├── samples/ # Sample code and examples
|
||||
├── .github/skills/ # Agent skills for Copilot
|
||||
└── tests/ # Integration tests
|
||||
```
|
||||
|
||||
### Package Relationships
|
||||
|
||||
- `agent-framework-core` contains core abstractions and OpenAI/Azure OpenAI built-in
|
||||
- Provider packages (`foundry`, `anthropic`, etc.) extend core with specific integrations
|
||||
- The root `agent_framework` public API is lazy-loaded from `packages/core/agent_framework/__init__.py` and
|
||||
described for type checkers in `packages/core/agent_framework/__init__.pyi`; keep both plus `__all__` in sync.
|
||||
- Core uses lazy loading via `__getattr__` in provider folders (e.g., `agent_framework/azure/`)
|
||||
|
||||
## Package Documentation
|
||||
|
||||
### Core
|
||||
- [core](packages/core/AGENTS.md) - Core abstractions, types, and built-in OpenAI/Azure OpenAI support
|
||||
|
||||
### LLM Providers
|
||||
- [anthropic](packages/anthropic/AGENTS.md) - Anthropic Claude API
|
||||
- [bedrock](packages/bedrock/AGENTS.md) - AWS Bedrock
|
||||
- [claude](packages/claude/AGENTS.md) - Claude Agent SDK
|
||||
- [foundry_local](packages/foundry_local/AGENTS.md) - Azure AI Foundry Local
|
||||
- [ollama](packages/ollama/AGENTS.md) - Local Ollama inference
|
||||
|
||||
### Azure Integrations
|
||||
- [foundry](packages/foundry/README.md) - Microsoft Foundry chat, agent, memory, and embedding integrations
|
||||
- [azure-contentunderstanding](packages/azure-contentunderstanding/AGENTS.md) - Azure Content Understanding context provider
|
||||
- [azure-ai-search](packages/azure-ai-search/AGENTS.md) - Azure AI Search RAG
|
||||
- [azure-cosmos](packages/azure-cosmos/AGENTS.md) - Azure Cosmos DB-backed history provider
|
||||
- [azurefunctions](packages/azurefunctions/AGENTS.md) - Azure Functions hosting
|
||||
|
||||
### Protocols & UI
|
||||
- [a2a](packages/a2a/AGENTS.md) - Agent-to-Agent protocol
|
||||
- [ag-ui](packages/ag-ui/AGENTS.md) - AG-UI protocol
|
||||
- [chatkit](packages/chatkit/AGENTS.md) - OpenAI ChatKit integration
|
||||
- [devui](packages/devui/AGENTS.md) - Developer UI for testing
|
||||
|
||||
### Storage & Memory
|
||||
- [mem0](packages/mem0/AGENTS.md) - Mem0 memory integration
|
||||
- [redis](packages/redis/AGENTS.md) - Redis storage
|
||||
|
||||
### Infrastructure
|
||||
- [copilotstudio](packages/copilotstudio/AGENTS.md) - Microsoft Copilot Studio
|
||||
- [declarative](packages/declarative/AGENTS.md) - YAML/JSON agent definitions
|
||||
- [durabletask](packages/durabletask/AGENTS.md) - Durable execution
|
||||
- [github_copilot](packages/github_copilot/AGENTS.md) - GitHub Copilot extensions
|
||||
- [purview](packages/purview/AGENTS.md) - Data governance
|
||||
|
||||
### Experimental
|
||||
- [lab](packages/lab/AGENTS.md) - Experimental features
|
||||
- [monty](packages/monty/AGENTS.md) - Monty-backed CodeAct integrations (alpha)
|
||||
+1391
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,770 @@
|
||||
# Coding Standards
|
||||
|
||||
This document describes the coding standards and conventions for the Agent Framework project.
|
||||
|
||||
## Code Style and Formatting
|
||||
|
||||
We use [ruff](https://github.com/astral-sh/ruff) for both linting and formatting with the following configuration:
|
||||
|
||||
- **Line length**: 120 characters
|
||||
- **Target Python version**: 3.10+
|
||||
- **Google-style docstrings**: All public functions, classes, and modules should have docstrings following Google conventions
|
||||
|
||||
### Module Docstrings
|
||||
|
||||
Public modules must include a module-level docstring, including `__init__.py` files.
|
||||
|
||||
- Namespace-style `__init__.py` modules (for example under `agent_framework/<provider>/`) should use a structured
|
||||
docstring that includes:
|
||||
- A one-line summary of the namespace
|
||||
- A short "This module lazily re-exports objects from:" section that lists only pip install package names
|
||||
(for example `agent-framework-a2a`)
|
||||
- A short "Supported classes:" (or "Supported classes and functions:") section
|
||||
- The main `agent_framework/__init__.py` should include a concise background-oriented docstring rather than a long
|
||||
per-symbol list.
|
||||
- Core modules with broad surface area, including `agent_framework/exceptions.py` and
|
||||
`agent_framework/observability.py`, should always have explicit module docstrings.
|
||||
|
||||
## Type Annotations
|
||||
|
||||
We use typing as a helper, it is not a goal in and of itself, so be pragmatic about where and when to strictly type, versus when to use a targetted cast or ignore.
|
||||
In general, the public interfaces of our classes, are important to get right, internally it is okay to have loosely typed code, as long as tests cover the code itself.
|
||||
This includes making a conscious choice when to program defensively, you can always do `getattr(item, 'attribute')` but that might end up causing you issues down the road
|
||||
because the type of `item` in this case, should have that attribute and if it doesn't it points to a larger issue, so if the type is expected to have that attribute, you should
|
||||
use `item.attribute` to ensure it fails at that point, rather then somewhere downstream where a value is expected but none was found.
|
||||
|
||||
### Future Annotations
|
||||
|
||||
> **Note:** This convention is being adopted. See [#3578](https://github.com/microsoft/agent-framework/issues/3578) for progress.
|
||||
|
||||
Use `from __future__ import annotations` at the top of files to enable postponed evaluation of annotations. This prevents the need for string-based type hints for forward references:
|
||||
|
||||
```python
|
||||
# ✅ Preferred - use future annotations
|
||||
from __future__ import annotations
|
||||
|
||||
class Agent:
|
||||
def create_child(self) -> Agent: # No quotes needed
|
||||
...
|
||||
|
||||
# ❌ Avoid - string-based type hints
|
||||
class Agent:
|
||||
def create_child(self) -> "Agent": # Requires quotes without future annotations
|
||||
...
|
||||
```
|
||||
|
||||
### TypeVar Naming Convention
|
||||
|
||||
> **Note:** This convention is being adopted. See [#3594](https://github.com/microsoft/agent-framework/issues/3594) for progress.
|
||||
|
||||
Use the suffix `T` for TypeVar names instead of a prefix:
|
||||
|
||||
```python
|
||||
# ✅ Preferred - suffix T
|
||||
ChatResponseT = TypeVar("ChatResponseT", bound=ChatResponse)
|
||||
AgentT = TypeVar("AgentT", bound=Agent)
|
||||
|
||||
# ❌ Avoid - prefix T
|
||||
TChatResponse = TypeVar("TChatResponse", bound=ChatResponse)
|
||||
TAgent = TypeVar("TAgent", bound=Agent)
|
||||
```
|
||||
|
||||
### Mapping Types
|
||||
|
||||
> **Note:** This convention is being adopted. See [#3577](https://github.com/microsoft/agent-framework/issues/3577) for progress.
|
||||
|
||||
Use `Mapping` instead of `MutableMapping` for input parameters when mutation is not required:
|
||||
|
||||
```python
|
||||
# ✅ Preferred - Mapping for read-only access
|
||||
def process_config(config: Mapping[str, Any]) -> None:
|
||||
...
|
||||
|
||||
# ❌ Avoid - MutableMapping when mutation isn't needed
|
||||
def process_config(config: MutableMapping[str, Any]) -> None:
|
||||
...
|
||||
```
|
||||
|
||||
### Typing Ignore and Cast Policy
|
||||
|
||||
Use typing as a helper first and suppressions as a last resort:
|
||||
|
||||
- **Prefer explicit typing before suppression**: Start with clearer type annotations, helper types, overloads,
|
||||
protocols, or refactoring dynamic code into typed helpers. Prioritize performance over completeness of typing, but make a good-faith effort to reduce uncertainty with typing before ignoring. Prefer to use a cast over a typeguard function since that does add overhead.
|
||||
- **Avoid redundant casts**: Do not add `cast(...)` if the type already matches; casts should be reserved for
|
||||
unavoidable narrowing where the runtime contract is known.
|
||||
- **Avoid multiple assignments**: Avoid assigning multiple variables just to get typing to pass, that has performance impact while typing should not have that.
|
||||
- **Source vs tests/samples**: Source code (`agent_framework*`) is checked **by pyright in strict mode** — use
|
||||
`# pyright: ignore[...]` there, never `# type: ignore` (strict pyright flags unnecessary ignores as errors). Tests
|
||||
and samples are checked by pyright (relaxed `basic`), mypy, pyrefly, ty (and zuban on tests) in a relaxed/basic
|
||||
profile; prefer real fixes (`isinstance`, `cast`, annotations, asserts for Optional access) over per-line ignores,
|
||||
and keep test/sample bodies readable rather than over-annotated. When a relaxed-pyright suppression is genuinely
|
||||
needed in tests/samples, use `# pyright: ignore[rule]`; the relaxed test/sample configs do not flag unnecessary
|
||||
ignores, so combine with a mypy/zuban `# type: ignore[code]` on the same line only where both are required.
|
||||
- **Line-level pyright ignores only**: If suppression is still required in source, use a line-level rule-specific ignore
|
||||
(`# pyright: ignore[reportGeneralTypeIssues]`), file-level is allowed if there is a compelling reason for it, that should be documented right beneath the ignore.
|
||||
Never change the global suppression flags unless the dev team okays it.
|
||||
- **Private usage boundary**: Accessing private members across `agent_framework*` packages can be acceptable for this
|
||||
codebase, but private member usage for non-Agent Framework dependencies should remain flagged. Do not make an
|
||||
internal helper public merely to satisfy pyright private-usage checks inside the package; use a targeted
|
||||
`# pyright: ignore[reportPrivateUsage]` when the internal dependency is intentional.
|
||||
- **Avoid typing-only wrappers**: Do not introduce trivial pass-through functions solely to satisfy typing or private
|
||||
usage checks. Prefer a targeted ignore, cast, or clearer annotation over an extra one-line function that adds runtime
|
||||
overhead without improving the design.
|
||||
|
||||
## Function Parameter Guidelines
|
||||
|
||||
To make the code easier to use and maintain:
|
||||
|
||||
- **Positional parameters**: Only use for up to 3 fully expected parameters (this is not a hard rule, but a guideline there are instances where this does make sense to exceed)
|
||||
- **Keyword-only parameters**: Arguments after `*` in function signatures are keyword-only; prefer these for optional parameters
|
||||
- **Avoid additional imports**: Do not require the user to import additional modules to use the function, so provide string based overrides when applicable, for instance:
|
||||
```python
|
||||
def create_agent(name: str, tool_mode: ChatToolMode) -> Agent:
|
||||
# Implementation here
|
||||
```
|
||||
Should be:
|
||||
```python
|
||||
def create_agent(name: str, tool_mode: Literal['auto', 'required', 'none'] | ChatToolMode) -> Agent:
|
||||
# Implementation here
|
||||
if isinstance(tool_mode, str):
|
||||
tool_mode = ChatToolMode(tool_mode)
|
||||
```
|
||||
- **Avoid shadowing built-ins**: Do not use parameter names that shadow Python built-ins (e.g., use `next_handler` instead of `next`). See [#3583](https://github.com/microsoft/agent-framework/issues/3583) for progress.
|
||||
|
||||
### Using `**kwargs`
|
||||
|
||||
> **Note:** This convention is being adopted. See [#3642](https://github.com/microsoft/agent-framework/issues/3642) for progress.
|
||||
|
||||
Avoid `**kwargs` unless absolutely necessary. It should only be used as an escape route, not for well-known flows of data:
|
||||
|
||||
- **Prefer named parameters**: If there are known extra arguments being passed, use explicit named parameters instead of kwargs
|
||||
- **Prefer purpose-specific buckets over generic kwargs**: If a flexible payload is still needed, use an explicit named parameter such as `additional_properties`, `function_invocation_kwargs`, or `client_kwargs` rather than a blanket `**kwargs`
|
||||
- **Subclassing support**: kwargs is acceptable in methods that are part of classes designed for subclassing, allowing subclass-defined kwargs to pass through without issues. In this case, clearly document that kwargs exists for subclass extensibility and not for passing arbitrary data
|
||||
- **Make known flows explicit first**: For abstract hooks, move known data flows into explicit parameters before leaving `**kwargs` behind for subclass extensibility (for example, prefer `state=` explicitly instead of passing it through kwargs)
|
||||
- **Prefer explicit metadata containers**: For constructors that expose metadata, prefer an explicit `additional_properties` parameter.
|
||||
- **Keep SDK passthroughs narrow and documented**: A kwargs escape hatch may be acceptable for provider helper APIs that pass through to a large or unstable external SDK surface, but it should be documented as SDK passthrough and revisited regularly
|
||||
- **Do not keep passthrough kwargs on wrappers that do not use them**: Convenience wrappers and session helpers should not accept generic kwargs merely to forward or ignore them
|
||||
- **Remove when possible**: In other cases, removing kwargs is likely better than keeping it
|
||||
- **Separate kwargs by purpose**: When combining kwargs for multiple purposes, use specific parameters like `client_kwargs: dict[str, Any]` instead of mixing everything in `**kwargs`
|
||||
- **Always document**: If kwargs must be used, always document how it's used, either by referencing external documentation or explaining its purpose
|
||||
|
||||
## Method Naming Inside Connectors
|
||||
|
||||
When naming methods inside connectors, we have a loose preference for using the following conventions:
|
||||
- Use `_prepare_<object>_for_<purpose>` as a prefix for methods that prepare data for sending to the external service.
|
||||
- Use `_parse_<object>_from_<source>` as a prefix for methods that process data received from the external service.
|
||||
|
||||
This is not a strict rule, but a guideline to help maintain consistency across the codebase.
|
||||
|
||||
## Implementation Decisions
|
||||
|
||||
### Asynchronous Programming
|
||||
|
||||
It's important to note that most of this library is written with asynchronous in mind. The
|
||||
developer should always assume everything is asynchronous. One can use the function signature
|
||||
with either `async def` or `def` to understand if something is asynchronous or not.
|
||||
|
||||
### Attributes vs Inheritance
|
||||
|
||||
Prefer attributes over inheritance when parameters are mostly the same:
|
||||
|
||||
```python
|
||||
# ✅ Preferred - using attributes
|
||||
from agent_framework import Message
|
||||
|
||||
user_msg = Message("user", ["Hello, world!"])
|
||||
asst_msg = Message("assistant", ["Hello, world!"])
|
||||
|
||||
# ❌ Not preferred - unnecessary inheritance
|
||||
class UserMessage(Message):
|
||||
pass
|
||||
|
||||
class AssistantMessage(Message):
|
||||
pass
|
||||
|
||||
user_msg = UserMessage("user", ["Hello, world!"])
|
||||
asst_msg = AssistantMessage("assistant", ["Hello, world!"])
|
||||
```
|
||||
|
||||
### Import Structure
|
||||
|
||||
The package follows a flat import structure:
|
||||
|
||||
- **Core**: Import directly from `agent_framework`
|
||||
```python
|
||||
from agent_framework import Agent, tool
|
||||
```
|
||||
|
||||
- **Components**: Import from `agent_framework.<component>`
|
||||
```python
|
||||
from agent_framework.observability import enable_sensitive_telemetry, configure_otel_providers
|
||||
```
|
||||
|
||||
- **Connectors**: Import from `agent_framework.<vendor/platform>`
|
||||
```python
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
from agent_framework.foundry import FoundryChatClient
|
||||
```
|
||||
|
||||
## Exception Hierarchy
|
||||
|
||||
The Agent Framework defines a structured exception hierarchy rooted at `AgentFrameworkException`. Every AF-specific
|
||||
exception inherits from this base, so callers can catch `AgentFrameworkException` as a broad fallback. The hierarchy
|
||||
is organized into domain-specific L1 branches, each with a consistent set of leaf exceptions where applicable.
|
||||
|
||||
### Design Principles
|
||||
|
||||
- **Domain-scoped branches**: Exceptions are grouped by the subsystem that raises them (agent, chat client,
|
||||
integration, workflow, content, tool, middleware), not by HTTP status code or generic error category.
|
||||
- **Consistent suberror pattern**: The `AgentException`, `ChatClientException`, and `IntegrationException` branches
|
||||
share a parallel set of leaf exceptions (`InvalidAuth`, `InvalidRequest`, `InvalidResponse`, `ContentFilter`) so
|
||||
that callers can handle the same failure mode uniformly across domains.
|
||||
- **Built-ins for validation**: Configuration/parameter validation errors use Python built-in exceptions
|
||||
(`ValueError`, `TypeError`, `RuntimeError`) rather than AF-specific classes. AF exceptions are reserved for
|
||||
domain-level failures that callers may want to catch and handle distinctly from programming errors.
|
||||
- **No compatibility aliases**: When exceptions are renamed or removed, the old names are not kept as aliases.
|
||||
This is a deliberate trade-off for hierarchy clarity over backward compatibility.
|
||||
- **Suffix convention**: L1 branch classes use `...Exception` (e.g., `AgentException`). Leaf classes may use
|
||||
either `...Exception` or `...Error` depending on the domain convention (e.g., `ContentError`,
|
||||
`WorkflowValidationError`). Within a branch, the suffix is consistent.
|
||||
|
||||
### Full Hierarchy
|
||||
|
||||
```
|
||||
AgentFrameworkException # Base for all AF exceptions
|
||||
├── AgentException # Agent-scoped failures
|
||||
│ ├── AgentInvalidAuthException # Agent auth failures
|
||||
│ ├── AgentInvalidRequestException # Invalid request to agent (e.g., agent not found, bad input)
|
||||
│ ├── AgentInvalidResponseException # Invalid/unexpected response from agent
|
||||
│ └── AgentContentFilterException # Agent content filter triggered
|
||||
│
|
||||
├── ChatClientException # Chat client lifecycle and communication failures
|
||||
│ ├── ChatClientInvalidAuthException # Chat client auth failures
|
||||
│ ├── ChatClientInvalidRequestException # Invalid request to chat client
|
||||
│ ├── ChatClientInvalidResponseException # Invalid/unexpected response from chat client
|
||||
│ └── ChatClientContentFilterException # Chat client content filter triggered
|
||||
│
|
||||
├── IntegrationException # External service/dependency integration failures
|
||||
│ ├── IntegrationInitializationError # Wrapped dependency lifecycle failure during setup
|
||||
│ ├── IntegrationInvalidAuthException # Integration auth failures (e.g., 401/403)
|
||||
│ ├── IntegrationInvalidRequestException # Invalid request to integration
|
||||
│ ├── IntegrationInvalidResponseException # Invalid/unexpected response from integration
|
||||
│ └── IntegrationContentFilterException # Integration content filter triggered
|
||||
│
|
||||
├── ContentError # Content processing/validation failures
|
||||
│ └── AdditionItemMismatch # Type mismatch when merging content items
|
||||
│
|
||||
├── WorkflowException # Workflow engine failures
|
||||
│ ├── WorkflowRunnerException # Runtime execution failures
|
||||
│ │ ├── WorkflowConvergenceException # Runner exceeded max iterations
|
||||
│ │ └── WorkflowCheckpointException # Checkpoint save/restore/decode failures
|
||||
│ ├── WorkflowValidationError # Graph validation errors
|
||||
│ │ ├── EdgeDuplicationError # Duplicate edge in workflow graph
|
||||
│ │ ├── TypeCompatibilityError # Type mismatch between connected executors
|
||||
│ │ └── GraphConnectivityError # Graph connectivity issues
|
||||
│ ├── WorkflowActionError # User-level error from declarative ThrowException action
|
||||
│ └── DeclarativeWorkflowError # Declarative workflow definition/YAML errors
|
||||
│
|
||||
├── ToolException # Tool-related failures
|
||||
│ └── ToolExecutionException # Failure during tool execution
|
||||
│
|
||||
├── MiddlewareException # Middleware failures
|
||||
│ └── MiddlewareTermination # Control-flow: early middleware termination
|
||||
│
|
||||
└── SettingNotFoundError # Required setting not resolved from any source
|
||||
```
|
||||
|
||||
### When to Use AF Exceptions vs Built-ins
|
||||
|
||||
| Scenario | Exception to use |
|
||||
|---|---|
|
||||
| Missing or invalid constructor argument (e.g., `api_key` is `None`) | `ValueError` or `TypeError` |
|
||||
| Object in wrong state (e.g., client not initialized) | `RuntimeError` |
|
||||
| External service returns 401/403 | `IntegrationInvalidAuthException` (or `ChatClient`/`Agent` variant) |
|
||||
| External service returns unexpected response | `IntegrationInvalidResponseException` (or variant) |
|
||||
| Content filter blocks a request | `IntegrationContentFilterException` (or variant) |
|
||||
| Request validation fails before sending to service | `IntegrationInvalidRequestException` (or variant) |
|
||||
| Agent not found in registry | `AgentInvalidRequestException` |
|
||||
| Agent returned no/bad response | `AgentInvalidResponseException` |
|
||||
| Workflow runner exceeds max iterations | `WorkflowConvergenceException` |
|
||||
| Checkpoint serialization/deserialization failure | `WorkflowCheckpointException` |
|
||||
| Workflow graph has invalid structure | `WorkflowValidationError` (or specific subclass) |
|
||||
| Declarative YAML definition error | `DeclarativeWorkflowError` |
|
||||
| Tool execution failure | `ToolExecutionException` |
|
||||
| Content merge type mismatch | `AdditionItemMismatch` |
|
||||
|
||||
### Choosing Between Agent, ChatClient, and Integration Branches
|
||||
|
||||
- **`AgentException`**: The failure is scoped to agent-level logic — agent lookup, agent response handling,
|
||||
agent content filtering. Use when the agent itself is the source of the problem.
|
||||
- **`ChatClientException`**: The failure is scoped to the chat client (the LLM provider connection) — auth with
|
||||
the LLM provider, request/response format issues specific to the chat protocol, chat-level content filtering.
|
||||
- **`IntegrationException`**: The failure is in a non-chat external dependency — search services, vector stores,
|
||||
Purview, custom APIs, or any service that is not the primary LLM chat provider.
|
||||
|
||||
When in doubt: if the code is in a chat client constructor or method, use `ChatClient*`. If it's in an agent
|
||||
method, use `Agent*`. If it's talking to an external service that isn't the chat LLM, use `Integration*`.
|
||||
|
||||
## Package Structure
|
||||
|
||||
The project uses a monorepo structure with separate packages for each connector/extension:
|
||||
|
||||
```plaintext
|
||||
python/
|
||||
├── pyproject.toml # Root package (agent-framework) depends on agent-framework-core[all]
|
||||
├── samples/ # Sample code and examples
|
||||
├── packages/
|
||||
│ ├── core/ # agent-framework-core - Core abstractions and implementations
|
||||
│ │ ├── pyproject.toml # Defines [all] extra that includes all connector packages
|
||||
│ │ ├── tests/ # Tests for core package
|
||||
│ │ └── agent_framework/
|
||||
│ │ ├── __init__.py # Lazy runtime public API exports
|
||||
│ │ ├── __init__.pyi # Public API typing surface for lazy root exports
|
||||
│ │ ├── _agents.py # Agent implementations
|
||||
│ │ ├── _clients.py # Chat client protocols and base classes
|
||||
│ │ ├── _tools.py # Tool definitions
|
||||
│ │ ├── _types.py # Type definitions
|
||||
│ │ │ # Provider folders - lazy load from connector packages
|
||||
│ │ ├── openai/ # OpenAI clients (built into core)
|
||||
│ │ ├── azure/ # Lazy loads from azure-ai, azure-ai-search, azurefunctions
|
||||
│ │ ├── anthropic/ # Lazy loads from agent-framework-anthropic
|
||||
│ │ ├── ollama/ # Lazy loads from agent-framework-ollama
|
||||
│ │ ├── a2a/ # Lazy loads from agent-framework-a2a
|
||||
│ │ ├── ag_ui/ # Lazy loads from agent-framework-ag-ui
|
||||
│ │ ├── chatkit/ # Lazy loads from agent-framework-chatkit
|
||||
│ │ ├── declarative/ # Lazy loads from agent-framework-declarative
|
||||
│ │ ├── devui/ # Lazy loads from agent-framework-devui
|
||||
│ │ ├── mem0/ # Lazy loads from agent-framework-mem0
|
||||
│ │ └── redis/ # Lazy loads from agent-framework-redis
|
||||
│ │
|
||||
│ ├── foundry/ # agent-framework-foundry
|
||||
│ │ ├── pyproject.toml
|
||||
│ │ ├── tests/
|
||||
│ │ └── agent_framework_foundry/
|
||||
│ │ ├── __init__.py # Public exports
|
||||
│ │ ├── _chat_client.py # FoundryChatClient implementation
|
||||
│ │ ├── _agent.py # FoundryAgent implementation
|
||||
│ │ ├── _embedding_client.py # FoundryEmbeddingClient implementation
|
||||
│ │ ├── _memory_provider.py # Foundry memory implementation
|
||||
│ │ └── py.typed # PEP 561 marker
|
||||
│ ├── anthropic/ # agent-framework-anthropic
|
||||
│ ├── bedrock/ # agent-framework-bedrock
|
||||
│ ├── ollama/ # agent-framework-ollama
|
||||
│ └── ... # Other connector packages
|
||||
```
|
||||
|
||||
### Lazy Loading Pattern
|
||||
|
||||
The root `agent_framework` package is a lazy public API surface. When adding, removing, or moving a root export:
|
||||
|
||||
- Add the symbol to `_LAZY_MODULE_EXPORTS` in `agent_framework/__init__.py`.
|
||||
- Keep `_LAZY_EXPORTS` derived from `_LAZY_MODULE_EXPORTS`.
|
||||
- Keep the explicit runtime `__all__` synchronized; it is required for `from agent_framework import *`.
|
||||
- Add the same public symbol to `agent_framework/__init__.pyi` so type checkers and editors see the typed surface.
|
||||
- Put runtime deprecation behavior in the owning module using that module's `__getattr__`. Do not add one-off
|
||||
deprecated-symbol branches to root `agent_framework.__getattr__`.
|
||||
- Validate root API changes with `uv run poe syntax -P core`, `uv run poe pyright -P core`, and import smoke tests
|
||||
for both `from agent_framework import <symbol>` and `from agent_framework import *`.
|
||||
|
||||
Provider folders in the core package use `__getattr__` to lazy load classes from their respective connector packages. This allows users to import from a consistent location while only loading dependencies when needed:
|
||||
|
||||
```python
|
||||
# In agent_framework/foundry/__init__.py
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"FoundryChatClient": ("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"Install it with: pip install {package_name}"
|
||||
) from exc
|
||||
```
|
||||
|
||||
### Adding a New Connector Package
|
||||
|
||||
**Important:** Do not create a new package unless there is an issue that has been reviewed and approved by the core team.
|
||||
|
||||
#### Initial Release (Preview Phase)
|
||||
|
||||
For the first release of a new connector package:
|
||||
|
||||
1. Create a new directory under `packages/` (e.g., `packages/my-connector/`)
|
||||
2. Add the package to `tool.uv.sources` in the root `pyproject.toml`
|
||||
3. Include samples inside the package itself (e.g., `packages/my-connector/samples/`)
|
||||
4. **Do NOT** add the package to the `[all]` extra in `packages/core/pyproject.toml`
|
||||
5. **Do NOT** create lazy loading in core yet
|
||||
|
||||
#### Promotion to Stable
|
||||
|
||||
After the package has been released and gained a measure of confidence:
|
||||
|
||||
1. Move samples from the package to the root `samples/` folder
|
||||
2. Add the package to the `[all]` extra in `packages/core/pyproject.toml`
|
||||
3. Create a provider folder in `agent_framework/` with lazy loading `__init__.py`
|
||||
|
||||
### Versioning and Core Dependency
|
||||
|
||||
All non-core packages declare a lower bound on `agent-framework-core` (e.g., `"agent-framework-core>=1.0.0b260130"`). Follow these rules when bumping versions:
|
||||
|
||||
- **Core version changes**: When `agent-framework-core` is updated with breaking or significant changes and its version is bumped, update the `agent-framework-core>=...` lower bound in every other package's `pyproject.toml` to match the new core version.
|
||||
- **Non-core version changes**: Non-core packages (connectors, extensions) can have their own versions incremented independently while keeping the existing core lower bound pinned. Only raise the core lower bound if the non-core package actually depends on new core APIs.
|
||||
|
||||
### External Dependency Version Bounds
|
||||
|
||||
The guiding principle for external dependencies is to make the range of allowed versions as broad as possible, even if that means we have to do some conditional imports, and other tricks to allow small changes in versions.
|
||||
So we use bounded ranges for external package dependencies in `pyproject.toml`:
|
||||
|
||||
|
||||
- For stable dependencies (`>=1.0.0`), use a lower bound at a known-good version and an explicit upper bound that reflects the maximum major version we currently support (for example: `openai>=1.99.0,<3`).
|
||||
- For prerelease (`dev`/`a`/`b`/`rc`) dependencies, use a known-good lower bound with a hard upper boundary in the same prerelease line (for example: `azure-ai-projects>=2.0.0b3,<2.0.0b4`).
|
||||
- For `<1.0.0` dependencies, use a known-good bounded range with an explicit upper cap. Prefer the broadest validated range the package can actually support: that may be a patch line, a minor line, or multiple minor lines (for example: `a2a-sdk>=0.3.5,<0.4.0`, `fastapi>=0.115.0,<0.136.0`, `uvicorn>=0.30.0,<0.39.0`).
|
||||
- For prerelease (`dev`/`a`/`b`/`rc`) dependencies, use a known-good bounded range with a hard upper cap and keep the range only as broad as the package's validation coverage justifies.
|
||||
- Prefer keeping support for multiple major versions when practical. This may mean that the upper bound spans multiple major versions when the dependency maintains backward compatibility; if APIs differ between supported majors, version-conditional imports/branches are acceptable to preserve compatibility.
|
||||
- When adding or changing an external dependency, first run `uv run poe validate-dependency-bounds-test` to validate workspace-wide lower/upper compatibility, then run `uv run poe validate-dependency-bounds-project --mode both --package <workspace-package-name> --dependency "<dependency-name>"` to expand package-scoped bounds.
|
||||
|
||||
### Installation Options
|
||||
|
||||
Connectors are distributed as separate packages and are not imported by default in the core package. Users install the specific connectors they need:
|
||||
|
||||
```bash
|
||||
# Install core only
|
||||
pip install agent-framework-core
|
||||
|
||||
# Install core with all connectors
|
||||
pip install agent-framework-core[all]
|
||||
# or (equivalently):
|
||||
pip install agent-framework
|
||||
|
||||
# Install specific connector (pulls in core as dependency)
|
||||
pip install agent-framework-foundry
|
||||
```
|
||||
|
||||
## Documentation
|
||||
|
||||
Each file should have a single first line containing: # Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
We follow the [Google Docstring](https://github.com/google/styleguide/blob/gh-pages/pyguide.md#383-functions-and-methods) style guide for functions and methods.
|
||||
They are currently not checked for private functions (functions starting with '_').
|
||||
|
||||
When a change adds, removes, or renames a sample-facing environment variable in repo-level samples or
|
||||
package-local sample docs for a package included by `agent-framework-core[all]`, update the consolidated
|
||||
inventory in `samples/README.md` in the same change.
|
||||
|
||||
They should contain:
|
||||
|
||||
- Single line explaining what the function does, ending with a period.
|
||||
- If necessary to further explain the logic a newline follows the first line and then the explanation is given.
|
||||
- The following three sections are optional, and if used should be separated by a single empty line.
|
||||
- Arguments are then specified after a header called `Args:`, with each argument being specified in the following format:
|
||||
- `arg_name`: Explanation of the argument.
|
||||
- if a longer explanation is needed for a argument, it should be placed on the next line, indented by 4 spaces.
|
||||
- Type and default values do not have to be specified, they will be pulled from the definition.
|
||||
- Returns are specified after a header called `Returns:` or `Yields:`, with the return type and explanation of the return value.
|
||||
- Keyword arguments are specified after a header called `Keyword Args:`, with each argument being specified in the same format as `Args:`.
|
||||
- A header for exceptions can be added, called `Raises:`, following these guidelines:
|
||||
- **Always document** Agent Framework specific exceptions (e.g., `AgentInvalidRequestException`, `IntegrationInvalidAuthException`)
|
||||
- **Only document** standard Python exceptions (TypeError, ValueError, KeyError, etc.) when the condition is non-obvious or provides value to API users
|
||||
- Format: `ExceptionType`: Explanation of the exception.
|
||||
- If a longer explanation is needed, it should be placed on the next line, indented by 4 spaces.
|
||||
- Code examples can be added using the `Examples:` header followed by `.. code-block:: python` directive.
|
||||
|
||||
Putting them all together, gives you at minimum this:
|
||||
|
||||
```python
|
||||
def equal(arg1: str, arg2: str) -> bool:
|
||||
"""Compares two strings and returns True if they are the same."""
|
||||
...
|
||||
```
|
||||
|
||||
Or a complete version of this:
|
||||
|
||||
```python
|
||||
def equal(arg1: str, arg2: str) -> bool:
|
||||
"""Compares two strings and returns True if they are the same.
|
||||
|
||||
Here is extra explanation of the logic involved.
|
||||
|
||||
Args:
|
||||
arg1: The first string to compare.
|
||||
arg2: The second string to compare.
|
||||
|
||||
Returns:
|
||||
True if the strings are the same, False otherwise.
|
||||
"""
|
||||
```
|
||||
|
||||
A more complete example with keyword arguments and code samples:
|
||||
|
||||
```python
|
||||
def create_client(
|
||||
model: str | None = None,
|
||||
*,
|
||||
timeout: float | None = None,
|
||||
env_file_path: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Client:
|
||||
"""Create a new client with the specified configuration.
|
||||
|
||||
Args:
|
||||
model: The model ID to use. If not provided,
|
||||
it will be loaded from settings.
|
||||
|
||||
Keyword Args:
|
||||
timeout: Optional timeout for requests.
|
||||
env_file_path: If provided, settings are read from this file.
|
||||
kwargs: Additional keyword arguments passed to the underlying client.
|
||||
|
||||
Returns:
|
||||
A configured client instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If the model is invalid.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Create a client with default settings:
|
||||
client = create_client(model="gpt-4o")
|
||||
|
||||
# Or load from environment:
|
||||
client = create_client(env_file_path=".env")
|
||||
"""
|
||||
...
|
||||
```
|
||||
|
||||
Use Google-style docstrings for all public APIs:
|
||||
|
||||
```python
|
||||
def create_agent(name: str, client: SupportsChatGetResponse) -> Agent:
|
||||
"""Create a new agent with the specified configuration.
|
||||
|
||||
Args:
|
||||
name: The name of the agent.
|
||||
client: The chat client to use for communication.
|
||||
|
||||
Returns:
|
||||
True if the strings are the same, False otherwise.
|
||||
|
||||
Raises:
|
||||
ValueError: If one of the strings is empty.
|
||||
"""
|
||||
...
|
||||
```
|
||||
|
||||
If in doubt, use the link above to read much more considerations of what to do and when, or use common sense.
|
||||
|
||||
## Public API and Exports
|
||||
|
||||
### Explicit Exports
|
||||
|
||||
**All wildcard imports (`from ... import *`) are prohibited** in production code, including both `.py` and `.pyi` files. Always use explicit import lists to maintain clarity and avoid namespace pollution.
|
||||
|
||||
Do not use ``__all__`` in internal modules. Define it in the ``__init__`` file of the level you want to expose.
|
||||
If a non-``__init__`` module is intentionally part of the public API surface (for example, ``observability.py``),
|
||||
it should define ``__all__`` as well.
|
||||
|
||||
Also avoid identity alias imports in ``__init__`` files. Use ``from ._module import Symbol`` instead of
|
||||
``from ._module import Symbol as Symbol``.
|
||||
|
||||
Exception: `.pyi` stubs that describe re-exported public APIs should use identity aliases (for example,
|
||||
`from ._agents import Agent as Agent`) so type checkers recognize the symbol as exported. This applies to the
|
||||
root `agent_framework/__init__.pyi`, which mirrors the lazy runtime exports from `agent_framework/__init__.py`.
|
||||
|
||||
```python
|
||||
# ✅ Preferred - explicit __all__ and named imports
|
||||
from ._agents import Agent
|
||||
from ._types import Message, ChatResponse
|
||||
|
||||
# ✅ For many exports, use parenthesized multi-line imports
|
||||
from ._types import (
|
||||
AgentResponse,
|
||||
ChatResponse,
|
||||
Message,
|
||||
ResponseStream,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Agent",
|
||||
"AgentResponse",
|
||||
"ChatResponse",
|
||||
"Message",
|
||||
"ResponseStream",
|
||||
]
|
||||
|
||||
# ❌ Prohibited pattern: wildcard/star imports (do not use)
|
||||
# from ._agents import *
|
||||
# from ._types import *
|
||||
|
||||
# ❌ Prohibited pattern: identity alias imports (do not use)
|
||||
# from ._agents import Agent as Agent
|
||||
```
|
||||
|
||||
**Rationale:**
|
||||
- **Clarity**: Explicit imports make it clear exactly what is being exported and used
|
||||
- **IDE Support**: Enables better autocomplete, go-to-definition, and refactoring
|
||||
- **Type Checking**: Improves static analysis and type checker accuracy
|
||||
- **Maintenance**: Makes it easier to track symbol usage and detect breaking changes
|
||||
- **Performance**: Avoids unnecessary symbol resolution during module import
|
||||
|
||||
## Performance considerations
|
||||
|
||||
### Cache Expensive Computations
|
||||
|
||||
Think about caching where appropriate. Cache the results of expensive operations that are called repeatedly with the same inputs:
|
||||
|
||||
```python
|
||||
# ✅ Preferred - cache expensive computations
|
||||
class FunctionTool:
|
||||
def __init__(self, ...):
|
||||
self._cached_parameters: dict[str, Any] | None = None
|
||||
|
||||
def parameters(self) -> dict[str, Any]:
|
||||
"""Return the JSON schema for the function's parameters.
|
||||
|
||||
The result is cached after the first call for performance.
|
||||
"""
|
||||
if self._cached_parameters is None:
|
||||
self._cached_parameters = self.input_model.model_json_schema()
|
||||
return self._cached_parameters
|
||||
|
||||
# ❌ Avoid - recalculating every time
|
||||
def parameters(self) -> dict[str, Any]:
|
||||
return self.input_model.model_json_schema()
|
||||
```
|
||||
|
||||
### Prefer Attribute Access Over isinstance()
|
||||
|
||||
When checking types in hot paths, prefer checking a `type` attribute (fast string comparison) over `isinstance()` (slower due to method resolution order traversal):
|
||||
|
||||
```python
|
||||
# ✅ Preferred - use match/case with type attribute (faster)
|
||||
match content.type:
|
||||
case "function_call":
|
||||
# handle function call
|
||||
case "usage":
|
||||
# handle usage
|
||||
case _:
|
||||
# handle other types
|
||||
|
||||
# ❌ Avoid in hot paths - isinstance() is slower
|
||||
if isinstance(content, FunctionCallContent):
|
||||
# handle function call
|
||||
elif isinstance(content, UsageContent):
|
||||
# handle usage
|
||||
```
|
||||
|
||||
For inline conditionals:
|
||||
|
||||
```python
|
||||
# ✅ Preferred - type attribute comparison
|
||||
result = value if content.type == "function_call" else other
|
||||
|
||||
# ❌ Avoid - isinstance() in hot paths
|
||||
result = value if isinstance(content, FunctionCallContent) else other
|
||||
```
|
||||
|
||||
### Avoid Redundant Serialization
|
||||
|
||||
When the same data needs to be used in multiple places, compute it once and reuse it:
|
||||
|
||||
```python
|
||||
# ✅ Preferred - reuse computed representation
|
||||
otel_message = _to_otel_message(message)
|
||||
otel_messages.append(otel_message)
|
||||
logger.info(otel_message, extra={...})
|
||||
|
||||
# ❌ Avoid - computing the same thing twice
|
||||
otel_messages.append(_to_otel_message(message)) # this already serializes
|
||||
message_data = message.to_dict(exclude_none=True) # and this does so again!
|
||||
logger.info(message_data, extra={...})
|
||||
```
|
||||
|
||||
## Test Organization
|
||||
|
||||
### Test Directory Structure
|
||||
|
||||
Test folders require specific organization to avoid pytest conflicts when running tests across packages:
|
||||
|
||||
1. **No `__init__.py` in test folders**: Test directories should NOT contain `__init__.py` files. This can cause import conflicts when pytest collects tests across multiple packages.
|
||||
|
||||
2. **File naming**: Files starting with `test_` are treated as test files by pytest. Do not use this prefix for helper modules or utilities. If you need shared test utilities, put them in `conftest.py` or a file with a different name pattern (e.g., `helpers.py`, `fixtures.py`).
|
||||
|
||||
3. **Package-specific conftest location**: The `tests/conftest.py` path is reserved for the core package (`packages/core/tests/conftest.py`). Other packages must place their tests in a uniquely-named subdirectory:
|
||||
|
||||
```plaintext
|
||||
# ✅ Correct structure for non-core packages
|
||||
packages/devui/
|
||||
├── tests/
|
||||
│ └── devui/ # Unique subdirectory matching package name
|
||||
│ ├── conftest.py # Package-specific fixtures
|
||||
│ ├── test_server.py
|
||||
│ └── test_mapper.py
|
||||
|
||||
packages/anthropic/
|
||||
├── tests/
|
||||
│ └── anthropic/ # Unique subdirectory
|
||||
│ ├── conftest.py
|
||||
│ └── test_client.py
|
||||
|
||||
# ❌ Incorrect - will conflict with core package
|
||||
packages/devui/
|
||||
├── tests/
|
||||
│ ├── conftest.py # Conflicts when running all tests
|
||||
│ ├── test_server.py
|
||||
│ └── test_helpers.py # Bad name - looks like a test file
|
||||
|
||||
# ✅ Core package can use tests/ directly
|
||||
packages/core/
|
||||
├── tests/
|
||||
│ ├── conftest.py # Core's conftest.py
|
||||
│ ├── core/
|
||||
│ │ └── test_agents.py
|
||||
│ └── openai/
|
||||
│ └── test_client.py
|
||||
```
|
||||
|
||||
4. **Keep the `tests/` folder**: Even when using a subdirectory, keep the `tests/` folder at the package root. Some test discovery commands and tooling rely on this convention.
|
||||
|
||||
### Fixture Guidelines
|
||||
|
||||
- Use `conftest.py` for shared fixtures within a test directory
|
||||
- Factory functions with parameters should be regular functions, not fixtures (fixtures can't accept arguments)
|
||||
- Import factory functions explicitly: `from conftest import create_test_request`
|
||||
- Fixtures should use simple names that describe what they provide: `mapper`, `test_request`, `mock_client`
|
||||
|
||||
### Integration Test Markers
|
||||
|
||||
New integration tests that call external services must have all three markers:
|
||||
|
||||
```python
|
||||
@pytest.mark.flaky
|
||||
@pytest.mark.integration
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_chat_completion() -> None:
|
||||
...
|
||||
```
|
||||
|
||||
- `@pytest.mark.flaky` — marks the test as potentially flaky since it depends on external services
|
||||
- `@pytest.mark.integration` — enables selecting/excluding integration tests with `-m integration` / `-m "not integration"`
|
||||
- `@skip_if_..._integration_tests_disabled` — skips the test when required API keys or service endpoints are missing
|
||||
|
||||
For test modules where all tests are integration tests, use `pytestmark`:
|
||||
|
||||
```python
|
||||
pytestmark = [
|
||||
pytest.mark.flaky,
|
||||
pytest.mark.integration,
|
||||
pytest.mark.sample("01_single_agent"),
|
||||
]
|
||||
```
|
||||
|
||||
When adding integration tests for a new provider, update the path filters and job assignments in **both** `python-merge-tests.yml` and `python-integration-tests.yml` — these workflows must be kept in sync. See the `python-testing` skill for details.
|
||||
@@ -0,0 +1,442 @@
|
||||
# Dev Setup
|
||||
|
||||
This document describes how to setup your environment with Python and uv,
|
||||
if you're working on new features or a bug fix for Agent Framework, or simply
|
||||
want to run the tests included.
|
||||
|
||||
For coding standards and conventions, see [CODING_STANDARD.md](CODING_STANDARD.md).
|
||||
|
||||
## System setup
|
||||
|
||||
We are using a tool called [poethepoet](https://github.com/nat-n/poethepoet) for task management and [uv](https://github.com/astral-sh/uv) for dependency management. At the [end of this document](#available-poe-tasks), you will find the available Poe tasks.
|
||||
|
||||
## If you're on WSL
|
||||
|
||||
Check that you've cloned the repository to `~/workspace` or a similar folder.
|
||||
Avoid `/mnt/c/` and prefer using your WSL user's home directory.
|
||||
|
||||
Ensure you have the WSL extension for VSCode installed.
|
||||
|
||||
## Using uv
|
||||
|
||||
uv allows us to use AF from the local files, without worrying about paths, as
|
||||
if you had AF pip package installed.
|
||||
|
||||
To install AF and all the required tools in your system, first, navigate to the directory containing
|
||||
this DEV_SETUP using your chosen shell.
|
||||
|
||||
### For windows (non-WSL)
|
||||
|
||||
Check the [uv documentation](https://docs.astral.sh/uv/getting-started/installation/) for the installation instructions. At the time of writing this is the command to install uv:
|
||||
|
||||
```powershell
|
||||
powershell -c "irm https://astral.sh/uv/install.ps1 | iex"
|
||||
```
|
||||
|
||||
### For WSL, Linux or MacOS
|
||||
|
||||
Check the [uv documentation](https://docs.astral.sh/uv/getting-started/installation/) for the installation instructions. At the time of writing this is the command to install uv:
|
||||
|
||||
```bash
|
||||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
```
|
||||
|
||||
### Alternative for MacOS
|
||||
|
||||
For MacOS users, Homebrew provides an easy installation of uv with the [uv Formulae](https://formulae.brew.sh/formula/uv)
|
||||
|
||||
```bash
|
||||
brew install uv
|
||||
```
|
||||
|
||||
|
||||
### After installing uv
|
||||
|
||||
You can then run the following commands manually:
|
||||
|
||||
```bash
|
||||
# Install Python 3.10, 3.11, 3.12, and 3.13
|
||||
uv python install 3.10 3.11 3.12 3.13
|
||||
# Create a virtual environment with Python 3.10 (you can change this to 3.11, 3.12 or 3.13)
|
||||
PYTHON_VERSION="3.10"
|
||||
uv venv --python $PYTHON_VERSION
|
||||
# Install AF and all dependencies
|
||||
uv sync --all-groups
|
||||
# Install all the tools and dependencies
|
||||
uv run poe install
|
||||
# Install prek hooks
|
||||
uv run poe prek-install
|
||||
```
|
||||
|
||||
Alternatively, you can reinstall the venv, pacakges, dependencies and prek hooks with a single command (but this requires poe in the current env), this is especially useful if you want to switch python versions:
|
||||
|
||||
```bash
|
||||
uv run poe setup -p 3.13
|
||||
```
|
||||
|
||||
You can then run different commands through Poe the Poet, use `uv run poe` to discover which ones.
|
||||
|
||||
## VSCode Setup
|
||||
|
||||
Install the [Python extension](https://marketplace.visualstudio.com/items?itemName=ms-python.python) for VSCode.
|
||||
|
||||
Open the `python` folder in [VSCode](https://code.visualstudio.com/docs/editor/workspaces).
|
||||
> The workspace for python should be rooted in the `./python` folder.
|
||||
|
||||
Open any of the `.py` files in the project and run the `Python: Select Interpreter`
|
||||
command from the command palette. Make sure the virtual env (default path is `.venv`) created by `uv` is selected.
|
||||
|
||||
## LLM setup
|
||||
|
||||
Make sure you have an
|
||||
[OpenAI API Key](https://platform.openai.com) or
|
||||
[Azure OpenAI service key](https://learn.microsoft.com/azure/cognitive-services/openai/quickstart?pivots=rest-api)
|
||||
|
||||
There are two methods to manage keys, secrets, and endpoints:
|
||||
|
||||
1. Store them in environment variables. AF Python leverages pydantic settings to load keys, secrets, and endpoints from the environment.
|
||||
> When you are using VSCode and have the python extension setup, it automatically loads environment variables from a `.env` file, so you don't have to manually set them in the terminal.
|
||||
> During runtime on different platforms, environment settings set as part of the deployments should be used.
|
||||
|
||||
2. Store them in a separate `.env` file, like `dev.env`, you can then pass that name into the constructor for most services, to the `env_file_path` parameter, see below.
|
||||
> Make sure to add `*.env` to your `.gitignore` file.
|
||||
|
||||
### Example for file-based setup with OpenAI Chat Completions
|
||||
To configure a `.env` file with just the keys needed for OpenAI Chat Completions, you can create a `openai.env` (this name is just as an example, a single `.env` with all required keys is more common) file in the root of the `python` folder with the following content:
|
||||
|
||||
Content of `.env` or `openai.env`:
|
||||
|
||||
```env
|
||||
OPENAI_API_KEY=""
|
||||
OPENAI_MODEL="gpt-4o-mini"
|
||||
```
|
||||
|
||||
You will then configure the ChatClient class with the keyword argument `env_file_path` (alternatively you can use `load_dotenv` in your code):
|
||||
|
||||
```python
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
client = OpenAIChatClient(env_file_path="openai.env")
|
||||
```
|
||||
|
||||
## Tests
|
||||
|
||||
All the tests are located in the `tests` folder of each package. Tests marked with `@pytest.mark.integration` and `@skip_if_..._integration_tests_disabled` are integration tests that require external services (e.g., OpenAI, Azure OpenAI). They are automatically skipped when the required API keys or service endpoints are not configured in your environment or `.env` file.
|
||||
|
||||
The root `test` command now supports both project-scoped fan-out and a single aggregate sweep:
|
||||
|
||||
```bash
|
||||
# Run package-local tests across all workspace packages
|
||||
uv run poe test
|
||||
|
||||
# Run tests for one workspace package
|
||||
uv run poe test -P core
|
||||
|
||||
# Run an aggregate pytest sweep across the selected packages
|
||||
uv run poe test -A
|
||||
|
||||
# Run only unit tests in aggregate mode
|
||||
uv run poe test -A -m "not integration"
|
||||
|
||||
# Run only integration tests in aggregate mode
|
||||
uv run poe test -A -m integration
|
||||
|
||||
# Run tests with coverage for one package or an aggregate sweep
|
||||
uv run poe test -P core -C
|
||||
uv run poe test -A -C
|
||||
```
|
||||
|
||||
Alternatively, you can run them using VSCode Tasks. Open the command palette
|
||||
(`Ctrl+Shift+P`) and type `Tasks: Run Task`. Select `Test` from the list.
|
||||
|
||||
Direct package execution still works when you need it:
|
||||
|
||||
```bash
|
||||
uv run poe --directory packages/core test
|
||||
```
|
||||
|
||||
Large packages (core, ag-ui, orchestrations, anthropic) use `pytest-xdist` for parallel test execution within the package. The aggregate `test -A` sweep also uses `pytest-xdist` across the selected packages.
|
||||
|
||||
## Code quality checks
|
||||
|
||||
To run the same checks that run during a commit and the GitHub Action `Python Code Quality`, you can use this command, from the [python](../python) folder:
|
||||
|
||||
```bash
|
||||
uv run poe check
|
||||
```
|
||||
|
||||
Ideally you should run these checks before committing any changes, when you install using the instructions above the prek hooks should be installed already.
|
||||
|
||||
## Code Coverage
|
||||
|
||||
We try to maintain a high code coverage for the project. To review coverage locally, use either a package-scoped run or the aggregate sweep:
|
||||
|
||||
```bash
|
||||
uv run poe test -P core -C
|
||||
uv run poe test -A -C
|
||||
```
|
||||
|
||||
This will show you which files are not covered by the tests, including the specific lines not covered. Make sure to consider the untested lines from the code you are working on, but feel free to add other tests as well, that is always welcome!
|
||||
|
||||
## Catching up with the latest changes
|
||||
|
||||
There are many people committing to Agent Framework, so it is important to keep your local repository up to date. To do this, you can run the following commands:
|
||||
|
||||
```bash
|
||||
git fetch upstream main
|
||||
git rebase upstream/main
|
||||
git push --force-with-lease
|
||||
```
|
||||
|
||||
or:
|
||||
|
||||
```bash
|
||||
git fetch upstream main
|
||||
git merge upstream/main
|
||||
git push
|
||||
```
|
||||
|
||||
This is assuming the upstream branch refers to the main repository. If you have a different name for the upstream branch, you can replace `upstream` with the name of your upstream branch.
|
||||
|
||||
After running the rebase command, you may need to resolve any conflicts that arise. If you are unsure how to resolve a conflict, please refer to the [GitHub's documentation on resolving conflicts](https://docs.github.com/en/get-started/using-git/resolving-merge-conflicts-after-a-git-rebase), or for [VSCode](https://code.visualstudio.com/docs/sourcecontrol/overview#_merge-conflicts).
|
||||
|
||||
# Task automation
|
||||
|
||||
## Available Poe Tasks
|
||||
This project uses [poethepoet](https://github.com/nat-n/poethepoet) for task management and [uv](https://github.com/astral-sh/uv) for dependency management.
|
||||
|
||||
### Setup and Installation
|
||||
|
||||
Once uv is installed, and you do not yet have a virtual environment setup:
|
||||
|
||||
```bash
|
||||
uv venv
|
||||
```
|
||||
|
||||
and then you can run the following tasks:
|
||||
```bash
|
||||
uv sync --all-extras --all-groups
|
||||
```
|
||||
|
||||
After this initial setup, you can use the following tasks to manage your development environment. It is advised to use the following setup command since that also installs the prek hooks.
|
||||
|
||||
#### `setup`
|
||||
Set up the development environment with a virtual environment, install dependencies and prek hooks:
|
||||
```bash
|
||||
uv run poe setup
|
||||
# or with specific Python version
|
||||
uv run poe setup -P 3.12
|
||||
```
|
||||
|
||||
#### `install`
|
||||
Install all dependencies (including extras and dependency groups) from the lockfile using frozen resolution:
|
||||
```bash
|
||||
uv run poe install
|
||||
```
|
||||
The root `dev` group contains shared tooling and source/type-check support. Package-specific test fixtures use
|
||||
`test` groups, while dependencies needed for a locally executable optional feature may use a feature-named group
|
||||
such as the lab package's `tau2` group.
|
||||
For intentional dependency upgrades, run `uv lock --upgrade-package <dependency-name>` and then run `uv run poe install`.
|
||||
|
||||
For repo-wide development dependency refreshes, run `uv run poe upgrade-dev-dependencies` to repin exact
|
||||
dependencies in development groups, refresh `uv.lock`, and rerun validation, typing, and tests.
|
||||
|
||||
#### `venv`
|
||||
Create a virtual environment with specified Python version or switch python version:
|
||||
```bash
|
||||
uv run poe venv
|
||||
# or with specific Python version
|
||||
uv run poe venv -P 3.12
|
||||
```
|
||||
|
||||
#### `prek-install`
|
||||
Install prek hooks:
|
||||
```bash
|
||||
uv run poe prek-install
|
||||
```
|
||||
|
||||
### Project-scoped command families
|
||||
|
||||
These commands default to `--package "*"`, so they run across all workspace packages unless you narrow them with `-P/--package`:
|
||||
|
||||
#### `syntax`
|
||||
Run Ruff formatting plus Ruff lint checks by default:
|
||||
```bash
|
||||
uv run poe syntax
|
||||
uv run poe syntax -P core
|
||||
uv run poe syntax -F # format only
|
||||
uv run poe syntax -C # lint/check only
|
||||
```
|
||||
|
||||
#### `build`
|
||||
Build workspace packages and the root meta package:
|
||||
```bash
|
||||
uv run poe build
|
||||
uv run poe build -P core
|
||||
```
|
||||
|
||||
#### `clean-dist`
|
||||
Clean generated dist artifacts:
|
||||
```bash
|
||||
uv run poe clean-dist
|
||||
uv run poe clean-dist -P core
|
||||
```
|
||||
|
||||
### Dual-mode validation and test commands
|
||||
|
||||
These command families share the same selector model:
|
||||
|
||||
```bash
|
||||
uv run poe <command> # project fan-out over --package "*"
|
||||
uv run poe <command> -P core # one-project fan-out
|
||||
uv run poe <command> -A # aggregate sweep where supported
|
||||
```
|
||||
|
||||
#### `pyright`
|
||||
Run Pyright type checking. Pyright is the **strict source-code type checker**, and also runs
|
||||
in a relaxed `basic` profile over the tests + samples (as one of the `test-typing` checkers):
|
||||
```bash
|
||||
uv run poe pyright
|
||||
uv run poe pyright -P core
|
||||
uv run poe pyright -A
|
||||
```
|
||||
|
||||
#### `test-typing`
|
||||
Run the **tests + samples** type checkers. Source code is owned by strict Pyright; the tests
|
||||
and samples are checked by `pyright` (relaxed), `mypy`, `pyrefly`, `ty`, and `zuban` in a
|
||||
deliberately relaxed/basic profile so real public-API type errors surface without forcing
|
||||
test/sample authors to fully annotate their code. All five gate CI:
|
||||
```bash
|
||||
uv run poe test-typing # all checkers over every package's tests
|
||||
uv run poe test-typing -P core # one package
|
||||
uv run poe test-typing -S # samples (pyright + pyrefly + ty; mypy/zuban skip script-style samples)
|
||||
uv run poe test-typing -P core --checker mypy # narrow to one checker (repeatable)
|
||||
uv run poe test-typing -P core --checker pyright # relaxed pyright over the tests
|
||||
```
|
||||
|
||||
#### `mypy`
|
||||
Convenience alias that runs MyPy over the test suite (MyPy no longer runs on source):
|
||||
```bash
|
||||
uv run poe mypy
|
||||
uv run poe mypy -P core
|
||||
```
|
||||
|
||||
#### `typing`
|
||||
Run Pyright over source **and** the tests/samples checkers:
|
||||
```bash
|
||||
uv run poe typing
|
||||
uv run poe typing -P core
|
||||
uv run poe typing -A
|
||||
```
|
||||
|
||||
#### `test`
|
||||
Run package-local tests in fan-out mode, or switch to one aggregate pytest sweep with `-A`:
|
||||
```bash
|
||||
uv run poe test
|
||||
uv run poe test -P core
|
||||
uv run poe test -P core -C
|
||||
uv run poe test -A
|
||||
uv run poe test -A -C
|
||||
```
|
||||
|
||||
### Sample-target variants
|
||||
|
||||
Use `-S/--samples` for sample-only validation instead of separate top-level commands:
|
||||
|
||||
```bash
|
||||
uv run poe syntax -S
|
||||
uv run poe syntax -S -C
|
||||
uv run poe pyright -S
|
||||
uv run poe check -S
|
||||
```
|
||||
|
||||
### Workspace validation and dependency commands
|
||||
|
||||
#### `markdown-code-lint`
|
||||
Lint markdown code blocks:
|
||||
```bash
|
||||
uv run poe markdown-code-lint
|
||||
```
|
||||
|
||||
#### `check-packages`
|
||||
Run the package-level syntax sweep (`syntax`) plus `pyright` across the selected projects:
|
||||
```bash
|
||||
uv run poe check-packages
|
||||
uv run poe check-packages -P core
|
||||
```
|
||||
|
||||
#### `check`
|
||||
Run package syntax, pyright, and tests for the selected project set. Without `-P/--package`, it also includes sample checks and markdown lint:
|
||||
```bash
|
||||
uv run poe check
|
||||
uv run poe check -P core
|
||||
uv run poe check -S
|
||||
```
|
||||
|
||||
#### `validate-dependency-bounds-test`
|
||||
Run workspace-wide dependency compatibility gates at lower and upper resolutions. This runs test + pyright across all packages and stops on first failure:
|
||||
```bash
|
||||
uv run poe validate-dependency-bounds-test
|
||||
# Defaults to --package "*"; pass a package to scope test mode
|
||||
uv run poe validate-dependency-bounds-test -P core
|
||||
```
|
||||
|
||||
#### `validate-dependency-bounds-project`
|
||||
Validate and extend dependency bounds for a single dependency in a single package. Use `--mode lower`, `--mode upper`, or the default `--mode both`:
|
||||
```bash
|
||||
uv run poe validate-dependency-bounds-project -M both -P core -D "<dependency-name>"
|
||||
```
|
||||
`--package` defaults to `*`, and `--dependency` is optional. Automation can use `--mode upper --package "*"` to run the upper-bound pass across the workspace.
|
||||
For `<1.0` dependencies, prefer the broadest validated range the package can really support. That may still be a single patch or minor line, but multi-minor ranges are fine when the package's checks/tests prove they work.
|
||||
|
||||
#### `add-dependency-and-validate-bounds`
|
||||
Add an external dependency to a workspace project and run both validators for that same project/dependency:
|
||||
```bash
|
||||
uv run poe add-dependency-and-validate-bounds -P core -D "<dependency-spec>"
|
||||
```
|
||||
|
||||
#### `upgrade-dev-dependencies`
|
||||
Refresh exact development dependency pins across the workspace, run `uv lock --upgrade`, reinstall from the frozen lockfile, then rerun validation, typing, and tests:
|
||||
```bash
|
||||
uv run poe upgrade-dev-dependencies
|
||||
```
|
||||
Use this for repo-wide development tooling and dependency-group refreshes. For targeted runtime dependency upgrades, prefer `uv lock --upgrade-package <dependency-name>` plus the package-scoped bound validation tasks above.
|
||||
|
||||
### Building and Publishing
|
||||
|
||||
#### `publish`
|
||||
Publish packages to PyPI:
|
||||
```bash
|
||||
uv run poe publish
|
||||
```
|
||||
|
||||
### Compatibility aliases
|
||||
|
||||
These legacy commands still work during the transition, but prefer the newer forms above:
|
||||
|
||||
```bash
|
||||
uv run poe fmt # prefer: uv run poe syntax -F
|
||||
uv run poe format # prefer: uv run poe syntax -F
|
||||
uv run poe lint # prefer: uv run poe syntax -C
|
||||
uv run poe all-tests # prefer: uv run poe test -A
|
||||
uv run poe all-tests-cov # prefer: uv run poe test -A -C
|
||||
uv run poe samples-lint # prefer: uv run poe syntax -S -C
|
||||
uv run poe samples-syntax # prefer: uv run poe pyright -S
|
||||
```
|
||||
|
||||
## Prek Hooks
|
||||
|
||||
Prek hooks run automatically on commit and stay intentionally lightweight:
|
||||
|
||||
- changed-package syntax formatting
|
||||
- changed-package syntax lint/check
|
||||
- markdown code lint only when markdown files change
|
||||
- sample lint + sample pyright only when files under `samples/` change
|
||||
|
||||
They do **not** run workspace `pyright` or `mypy` by default. Use `uv run poe pyright`, `uv run poe mypy`, `uv run poe typing`, `uv run poe check-packages`, or `uv run poe check` when you want deeper validation.
|
||||
|
||||
You can run the installed hooks directly with:
|
||||
|
||||
```bash
|
||||
uv run prek run -a
|
||||
```
|
||||
@@ -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,85 @@
|
||||
# Python Package Status
|
||||
|
||||
This file tracks the current lifecycle state of the Python packages in this workspace. Some packages at later stages might have features within them that are not ready yet, these have feature stage decorators on the relevant APIs, and for `experimental` features warnings are raised. See the [Feature-level staged APIs](#feature-level-staged-apis) section below for details on which features are in which stage and where to find them.
|
||||
|
||||
Status is grouped into these buckets:
|
||||
|
||||
- `alpha` - initial release and early development packages that are not yet ready for general use
|
||||
- `beta` - prerelease packages that are not currently release candidates
|
||||
- `rc` - release candidate packages, these are close to ready for release but may still have some breaking changes before the final release
|
||||
- `released` - stable packages without a prerelease suffix, these are stable packages that should not have breaking changes between versions
|
||||
- `deprecated` - removed or deprecated packages that should not be used for new work
|
||||
|
||||
## Current packages
|
||||
|
||||
| Package | Path | State |
|
||||
| --- | --- | --- |
|
||||
| `agent-framework` | `python/` | `released` |
|
||||
| `agent-framework-a2a` | `python/packages/a2a` | `beta` |
|
||||
| `agent-framework-ag-ui` | `python/packages/ag-ui` | `rc` |
|
||||
| `agent-framework-anthropic` | `python/packages/anthropic` | `beta` |
|
||||
| `agent-framework-azure-contentunderstanding` | `python/packages/azure-contentunderstanding` | `alpha` |
|
||||
| `agent-framework-azure-ai-search` | `python/packages/azure-ai-search` | `beta` |
|
||||
| `agent-framework-azure-cosmos` | `python/packages/azure-cosmos` | `beta` |
|
||||
| `agent-framework-azurefunctions` | `python/packages/azurefunctions` | `beta` |
|
||||
| `agent-framework-bedrock` | `python/packages/bedrock` | `beta` |
|
||||
| `agent-framework-chatkit` | `python/packages/chatkit` | `beta` |
|
||||
| `agent-framework-claude` | `python/packages/claude` | `beta` |
|
||||
| `agent-framework-copilotstudio` | `python/packages/copilotstudio` | `beta` |
|
||||
| `agent-framework-core` | `python/packages/core` | `released` |
|
||||
| `agent-framework-declarative` | `python/packages/declarative` | `rc` |
|
||||
| `agent-framework-devui` | `python/packages/devui` | `beta` |
|
||||
| `agent-framework-durabletask` | `python/packages/durabletask` | `beta` |
|
||||
| `agent-framework-foundry` | `python/packages/foundry` | `released` |
|
||||
| `agent-framework-foundry-local` | `python/packages/foundry_local` | `beta` |
|
||||
| `agent-framework-gemini` | `python/packages/gemini` | `alpha` |
|
||||
| `agent-framework-github-copilot` | `python/packages/github_copilot` | `rc` |
|
||||
| `agent-framework-hosting` | `python/packages/hosting` | `alpha` |
|
||||
| `agent-framework-hosting-responses` | `python/packages/hosting-responses` | `alpha` |
|
||||
| `agent-framework-hyperlight` | `python/packages/hyperlight` | `beta` |
|
||||
| `agent-framework-lab` | `python/packages/lab` | `beta` |
|
||||
| `agent-framework-mem0` | `python/packages/mem0` | `beta` |
|
||||
| `agent-framework-mistral` | `python/packages/mistral` | `alpha` |
|
||||
| `agent-framework-monty` | `python/packages/monty` | `alpha` |
|
||||
| `agent-framework-ollama` | `python/packages/ollama` | `beta` |
|
||||
| `agent-framework-openai` | `python/packages/openai` | `released` |
|
||||
| `agent-framework-orchestrations` | `python/packages/orchestrations` | `released` |
|
||||
| `agent-framework-purview` | `python/packages/purview` | `beta` |
|
||||
| `agent-framework-redis` | `python/packages/redis` | `beta` |
|
||||
|
||||
## Deprecated / removed packages
|
||||
|
||||
| Package | Previous path | State | Notes |
|
||||
| --- | --- | --- | --- |
|
||||
| `agent-framework-azure-ai` | `python/packages/azure-ai` | `deprecated` | The client classes within the `azure-ai` package were renamed, sometimes changed, and moved to `agent-framework-foundry`. |
|
||||
|
||||
## Feature-level staged APIs
|
||||
|
||||
The following feature IDs have explicit feature-stage decorators on public APIs in the packages
|
||||
listed below.
|
||||
|
||||
### Experimental features
|
||||
|
||||
#### `DECLARATIVE_AGENTS`
|
||||
|
||||
- `agent-framework-declarative`: declarative agent loading APIs from
|
||||
`agent_framework_declarative`, including `AgentFactory`,
|
||||
`DeclarativeLoaderError`, `ProviderLookupError`, and `ProviderTypeMapping`
|
||||
from `agent_framework_declarative/_loader.py`
|
||||
|
||||
#### `EVALS`
|
||||
|
||||
- `agent-framework-core`: exported evaluation APIs from `agent_framework`, including
|
||||
`LocalEvaluator`, `evaluate_agent`, `evaluate_workflow`, and the related evaluation types and
|
||||
helper checks defined in `agent_framework/_evaluation.py`
|
||||
- `agent-framework-foundry`: `FoundryEvals`, `evaluate_traces`, and `evaluate_foundry_target`
|
||||
|
||||
#### `SKILLS`
|
||||
|
||||
- `agent-framework-core`: exported skills APIs from `agent_framework`, including `Skill`,
|
||||
`SkillResource`, `SkillScript`, `SkillScriptRunner`, and `SkillsProvider` from
|
||||
`agent_framework/_skills.py`
|
||||
|
||||
### Release-candidate features
|
||||
|
||||
There are currently no feature-level `rc` APIs.
|
||||
@@ -0,0 +1,261 @@
|
||||
# Get Started with Microsoft Agent Framework for Python Developers
|
||||
|
||||
## Quick Install
|
||||
|
||||
We recommend two common installation paths depending on your use case.
|
||||
|
||||
### 1. Development mode
|
||||
|
||||
If you are exploring or developing locally, install the entire framework with all sub-packages:
|
||||
|
||||
```bash
|
||||
pip install agent-framework
|
||||
```
|
||||
|
||||
This installs the core and every integration package, making sure that all features are available without additional steps. This is the simplest way to get started.
|
||||
|
||||
### 2. Selective install
|
||||
|
||||
If you only need specific integrations, you can install at a more granular level. This keeps dependencies lighter and focuses on what you actually plan to use. Some examples:
|
||||
|
||||
```bash
|
||||
# Core only
|
||||
# includes Azure OpenAI and OpenAI support by default
|
||||
# also includes workflows and orchestrations
|
||||
pip install agent-framework-core
|
||||
|
||||
# Core + Azure AI Foundry integration
|
||||
pip install agent-framework-foundry
|
||||
|
||||
# Core + Microsoft Copilot Studio integration (preview package)
|
||||
pip install agent-framework-copilotstudio --pre
|
||||
|
||||
# Core + both Microsoft Copilot Studio and Azure AI Foundry integration
|
||||
pip install --pre agent-framework-copilotstudio agent-framework-foundry
|
||||
```
|
||||
|
||||
This selective approach is useful when you know which integrations you need, and it is the recommended way to set up lightweight environments. Released packages such as `agent-framework`, `agent-framework-core`, and `agent-framework-foundry` no longer require `--pre`, while preview connectors such as `agent-framework-copilotstudio` still do.
|
||||
|
||||
Supported Platforms:
|
||||
|
||||
- Python: 3.10+
|
||||
- OS: Windows, macOS, Linux
|
||||
|
||||
## 1. Setup API Keys
|
||||
|
||||
Set as environment variables, or create a .env file at your project root:
|
||||
|
||||
```bash
|
||||
OPENAI_API_KEY=sk-...
|
||||
OPENAI_MODEL=...
|
||||
...
|
||||
AZURE_OPENAI_API_KEY=...
|
||||
AZURE_OPENAI_ENDPOINT=...
|
||||
AZURE_OPENAI_MODEL=...
|
||||
...
|
||||
FOUNDRY_PROJECT_ENDPOINT=...
|
||||
FOUNDRY_MODEL=...
|
||||
```
|
||||
|
||||
For the generic OpenAI clients (`OpenAIChatClient` and `OpenAIChatCompletionClient`), configuration
|
||||
resolves in this order:
|
||||
|
||||
1. Explicit Azure inputs such as `credential` or `azure_endpoint`
|
||||
2. `OPENAI_API_KEY` / explicit OpenAI API-key parameters
|
||||
3. Azure environment fallback such as `AZURE_OPENAI_ENDPOINT` and `AZURE_OPENAI_API_KEY`
|
||||
|
||||
This means mixed shells default to OpenAI when `OPENAI_API_KEY` is present. To force Azure routing,
|
||||
pass an explicit Azure input such as `credential=AzureCliCredential()`.
|
||||
|
||||
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='',
|
||||
azure_endpoint='',
|
||||
model='',
|
||||
api_version='',
|
||||
)
|
||||
```
|
||||
|
||||
See the following [setup guide](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
|
||||
|
||||
async def main():
|
||||
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 = await agent.run("Summarize the Three Laws of Robotics")
|
||||
print(result)
|
||||
|
||||
asyncio.run(main())
|
||||
# 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 import Message
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
async def main():
|
||||
client = OpenAIChatClient()
|
||||
|
||||
messages = [
|
||||
Message("system", ["You are a helpful assistant."]),
|
||||
Message("user", ["Write a haiku about Agent Framework."])
|
||||
]
|
||||
|
||||
response = await client.get_response(messages)
|
||||
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 pydantic import Field
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
|
||||
def get_weather(
|
||||
location: Annotated[str, Field(description="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.
|
||||
"""
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
You can explore additional agent samples [here](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())
|
||||
```
|
||||
|
||||
For more advanced orchestration patterns including Sequential, Concurrent, Group Chat, Handoff, and Magentic orchestrations, see the [orchestration samples](samples/03-workflows/orchestrations).
|
||||
|
||||
## More Examples & Samples
|
||||
|
||||
- [Getting Started with Agents](samples/02-agents): Basic agent creation and tool usage
|
||||
- [Chat Client Examples](samples/02-agents/chat_client): Direct chat client usage patterns
|
||||
- [Foundry Integration](https://github.com/microsoft/agent-framework/tree/main/python/packages/foundry): Microsoft Foundry integration
|
||||
- [Workflow Samples](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 docs are coming soon.
|
||||
@@ -0,0 +1,31 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from importlib import metadata as _metadata
|
||||
from pathlib import Path as _Path
|
||||
from typing import Any, cast
|
||||
|
||||
try:
|
||||
import tomllib as _toml # type: ignore # Python 3.11+
|
||||
except ModuleNotFoundError: # Python 3.10
|
||||
import tomli as _toml # type: ignore
|
||||
|
||||
|
||||
def _load_pyproject() -> dict[str, Any]:
|
||||
pyproject = (_Path(__file__).resolve().parents[1] / "pyproject.toml").read_text("utf-8")
|
||||
return cast(dict[str, Any], _toml.loads(pyproject)) # type: ignore
|
||||
|
||||
|
||||
def _version() -> str:
|
||||
try:
|
||||
return _metadata.version("agent-framework")
|
||||
except _metadata.PackageNotFoundError as ex:
|
||||
data = _load_pyproject()
|
||||
project = cast(dict[str, Any], data.get("project", {}))
|
||||
version = project.get("version")
|
||||
if isinstance(version, str):
|
||||
return version
|
||||
raise RuntimeError("pyproject.toml missing project.version") from ex
|
||||
|
||||
|
||||
__version__ = _version()
|
||||
__all__ = ["__version__"]
|
||||
@@ -0,0 +1,10 @@
|
||||
uv python install 3.10 3.11 3.12 3.13
|
||||
# Create a virtual environment with Python 3.10 (you can change this to 3.11, 3.12 or 3.13)
|
||||
PYTHON_VERSION="3.13"
|
||||
uv venv --python $PYTHON_VERSION
|
||||
# Install AF and all dependencies
|
||||
uv sync --dev
|
||||
# Install all the tools and dependencies
|
||||
uv run poe install
|
||||
# Install prek hooks
|
||||
uv run poe prek-install
|
||||
@@ -0,0 +1,58 @@
|
||||
# A2A Package (agent-framework-a2a)
|
||||
|
||||
Agent-to-Agent (A2A) protocol support for inter-agent communication.
|
||||
|
||||
## Main Classes
|
||||
|
||||
- **`A2AAgent`** - Client to connect to remote A2A-compliant agents.
|
||||
- **`A2AExecutor`** - Bridge to expose Agent Framework agents via the A2A protocol.
|
||||
- **`A2AServiceSessionId`** - Typed durable A2A continuation state shape stored in `AgentSession.service_session_id`.
|
||||
- **`A2AAgentSession`** - Deprecated compatibility session wrapper; prefer `AgentSession` + `A2AServiceSessionId`.
|
||||
|
||||
## Usage
|
||||
|
||||
### A2AAgent (Client)
|
||||
|
||||
```python
|
||||
from agent_framework.a2a import A2AAgent
|
||||
|
||||
# Connect to a remote A2A agent
|
||||
a2a_agent = A2AAgent(url="http://remote-agent/a2a")
|
||||
response = await a2a_agent.run("Hello!")
|
||||
```
|
||||
|
||||
### A2AExecutor (Server/Bridge)
|
||||
|
||||
```python
|
||||
from agent_framework.a2a import A2AExecutor
|
||||
from a2a.server.request_handlers import DefaultRequestHandler
|
||||
from a2a.server.routes import create_agent_card_routes, create_jsonrpc_routes
|
||||
from a2a.server.tasks import InMemoryTaskStore
|
||||
from starlette.applications import Starlette
|
||||
|
||||
# Create an A2A executor for your agent
|
||||
executor = A2AExecutor(agent=my_agent)
|
||||
|
||||
# Set up the request handler (agent_card is required)
|
||||
request_handler = DefaultRequestHandler(
|
||||
agent_executor=executor,
|
||||
task_store=InMemoryTaskStore(),
|
||||
agent_card=my_agent_card,
|
||||
)
|
||||
|
||||
# Build a Starlette app with A2A routes
|
||||
app = Starlette(
|
||||
routes=[
|
||||
*create_agent_card_routes(my_agent_card),
|
||||
*create_jsonrpc_routes(request_handler, "/"),
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
## Import Path
|
||||
|
||||
```python
|
||||
from agent_framework.a2a import A2AAgent, A2AExecutor, A2AServiceSessionId
|
||||
# or directly:
|
||||
from agent_framework_a2a import A2AAgent, A2AExecutor, A2AServiceSessionId
|
||||
```
|
||||
@@ -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,64 @@
|
||||
# Get Started with Microsoft Agent Framework A2A
|
||||
|
||||
Please install this package via pip:
|
||||
|
||||
```bash
|
||||
pip install agent-framework-a2a --pre
|
||||
```
|
||||
|
||||
## A2A Agent Integration
|
||||
|
||||
The A2A agent integration enables communication with remote A2A-compliant agents using the standardized A2A protocol. This allows your Agent Framework applications to connect to agents running on different platforms, languages, or services.
|
||||
|
||||
### A2AAgent (Client)
|
||||
|
||||
The `A2AAgent` class is a client that wraps an A2A Client to connect the Agent Framework with external A2A-compliant agents.
|
||||
|
||||
```python
|
||||
from agent_framework.a2a import A2AAgent
|
||||
|
||||
# Connect to a remote A2A agent
|
||||
a2a_agent = A2AAgent(url="http://remote-agent/a2a")
|
||||
response = await a2a_agent.run("Hello!")
|
||||
```
|
||||
|
||||
### A2AExecutor (Hosting)
|
||||
|
||||
The `A2AExecutor` class bridges local AI agents built with the `agent_framework` library to the A2A protocol, allowing them to be hosted and accessed by other A2A-compliant clients.
|
||||
|
||||
```python
|
||||
from agent_framework.a2a import A2AExecutor
|
||||
from a2a.server.apps import A2AStarletteApplication
|
||||
from a2a.server.request_handlers import DefaultRequestHandler
|
||||
from a2a.server.tasks import InMemoryTaskStore
|
||||
|
||||
# Create an A2A executor for your agent
|
||||
executor = A2AExecutor(agent=my_agent)
|
||||
|
||||
# Set up the request handler and server application
|
||||
request_handler = DefaultRequestHandler(
|
||||
agent_executor=executor,
|
||||
task_store=InMemoryTaskStore(),
|
||||
)
|
||||
|
||||
app = A2AStarletteApplication(
|
||||
agent_card=my_agent_card,
|
||||
http_handler=request_handler,
|
||||
).build()
|
||||
```
|
||||
|
||||
### Basic Usage Example
|
||||
|
||||
See the [A2A agent examples](../../samples/04-hosting/a2a/) which demonstrate:
|
||||
|
||||
- Connecting to remote A2A agents
|
||||
- Hosting local agents via A2A protocol
|
||||
- Sending messages and receiving responses
|
||||
- Handling different content types (text, files, data)
|
||||
- Streaming responses and real-time interaction
|
||||
|
||||
## Security considerations
|
||||
|
||||
The hosting example above focuses on protocol wiring and does not add authentication or authorization by itself. Production A2A hosts should protect their HTTP or JSON-RPC entry points with the deployment's normal auth layer and verify that each caller is allowed to access the requested agent, task, or session.
|
||||
|
||||
Task, thread, context, and session identifiers used by an A2A host are routing handles, not bearer credentials. Do not rely on client-supplied identifiers alone to select or mutate persisted state; bind them to authenticated user, tenant, or workspace context first.
|
||||
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import importlib.metadata
|
||||
|
||||
from ._a2a_executor import A2AExecutor
|
||||
from ._agent import A2AAgent, A2AAgentSession, A2AContinuationToken, A2AServiceSessionId
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0" # Fallback for development mode
|
||||
|
||||
__all__ = [
|
||||
"A2AAgent",
|
||||
"A2AAgentSession",
|
||||
"A2AContinuationToken",
|
||||
"A2AExecutor",
|
||||
"A2AServiceSessionId",
|
||||
"__version__",
|
||||
]
|
||||
@@ -0,0 +1,300 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import uuid
|
||||
from asyncio import CancelledError
|
||||
from collections.abc import Mapping
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
from a2a.helpers import new_task_from_user_message
|
||||
from a2a.server.agent_execution import AgentExecutor, RequestContext
|
||||
from a2a.server.events import EventQueue
|
||||
from a2a.server.tasks import TaskUpdater
|
||||
from a2a.types import Part, TaskState
|
||||
from agent_framework import (
|
||||
AgentResponseUpdate,
|
||||
AgentSession,
|
||||
Message,
|
||||
SupportsAgentRun,
|
||||
)
|
||||
from typing_extensions import override
|
||||
|
||||
from ._utils import get_uri_data
|
||||
|
||||
logger = logging.getLogger("agent_framework.a2a")
|
||||
|
||||
|
||||
class A2AExecutor(AgentExecutor):
|
||||
"""Execute AI agents using the A2A (Agent-to-Agent) protocol.
|
||||
|
||||
The A2AExecutor bridges AI agents built with the agent_framework library and the A2A protocol,
|
||||
enabling structured agent execution with event-driven communication. It handles execution
|
||||
contexts, delegates history management to the agent's session, and converts agent
|
||||
responses into A2A protocol events.
|
||||
|
||||
The executor supports executing an Agent or WorkflowAgent. It provides comprehensive
|
||||
error handling with task status updates and supports various content types including text,
|
||||
binary data, and URI-based content.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from a2a.server.request_handlers import DefaultRequestHandler
|
||||
from a2a.server.routes import create_jsonrpc_routes, create_agent_card_routes
|
||||
from a2a.server.tasks import InMemoryTaskStore
|
||||
from a2a.types import AgentCapabilities, AgentCard, AgentInterface
|
||||
from agent_framework.a2a import A2AExecutor
|
||||
from agent_framework.openai import OpenAIResponsesClient
|
||||
from starlette.applications import Starlette
|
||||
|
||||
public_agent_card = AgentCard(
|
||||
name="Food Agent",
|
||||
description="A simple agent that provides food-related information.",
|
||||
version="1.0.0",
|
||||
default_input_modes=["text"],
|
||||
default_output_modes=["text"],
|
||||
capabilities=AgentCapabilities(streaming=True),
|
||||
supported_interfaces=[
|
||||
AgentInterface(url="http://localhost:9999/", protocol_binding="JSONRPC"),
|
||||
],
|
||||
skills=[],
|
||||
)
|
||||
|
||||
# Create an agent
|
||||
agent = OpenAIResponsesClient().as_agent(
|
||||
name="Food Agent",
|
||||
instructions="A simple agent that provides food-related information.",
|
||||
)
|
||||
|
||||
# Set up the A2A server with the A2AExecutor enabled for streaming
|
||||
# and passing custom keyword arguments to the agent's run method.
|
||||
request_handler = DefaultRequestHandler(
|
||||
agent_executor=A2AExecutor(agent, stream=True, run_kwargs={"client_kwargs": {"max_tokens": 500}}),
|
||||
task_store=InMemoryTaskStore(),
|
||||
agent_card=public_agent_card,
|
||||
)
|
||||
|
||||
app = Starlette(
|
||||
routes=[
|
||||
*create_agent_card_routes(public_agent_card),
|
||||
*create_jsonrpc_routes(request_handler, "/"),
|
||||
],
|
||||
)
|
||||
|
||||
Args:
|
||||
agent: The AI agent to execute.
|
||||
stream: Whether to stream the agent response. Defaults to False.
|
||||
run_kwargs: Additional keyword arguments to pass to the agent's run method.
|
||||
"""
|
||||
|
||||
def __init__(self, agent: SupportsAgentRun, stream: bool = False, run_kwargs: Mapping[str, Any] | None = None):
|
||||
"""Initialize the A2AExecutor with the specified agent.
|
||||
|
||||
Args:
|
||||
agent: The AI agent or workflow to execute.
|
||||
stream: Whether to stream the agent response. Defaults to False.
|
||||
run_kwargs: Additional keyword arguments to pass to the agent's run method.
|
||||
Cannot contain 'session' or 'stream' as these are managed by the executor.
|
||||
|
||||
Raises:
|
||||
ValueError: If run_kwargs contains 'session' or 'stream'.
|
||||
"""
|
||||
super().__init__()
|
||||
self._agent: SupportsAgentRun = agent
|
||||
self._stream: bool = stream
|
||||
if run_kwargs:
|
||||
if "session" in run_kwargs:
|
||||
raise ValueError("run_kwargs cannot contain 'session' as it is managed by the executor.")
|
||||
if "stream" in run_kwargs:
|
||||
raise ValueError("run_kwargs cannot contain 'stream' as it is managed by the executor.")
|
||||
self._run_kwargs: Mapping[str, Any] = run_kwargs or {}
|
||||
|
||||
@override
|
||||
async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
|
||||
"""Cancel agent execution for the given request context.
|
||||
|
||||
Uses a TaskUpdater to send a cancellation event through the provided event queue.
|
||||
|
||||
Args:
|
||||
context: The request context identifying the task to cancel.
|
||||
event_queue: The event queue to publish the cancellation event to.
|
||||
|
||||
Raises:
|
||||
ValueError: If context_id is not provided in the RequestContext.
|
||||
"""
|
||||
if context.context_id is None:
|
||||
raise ValueError("Context ID must be provided in the RequestContext")
|
||||
|
||||
updater = TaskUpdater(
|
||||
event_queue=event_queue,
|
||||
task_id=context.task_id or "",
|
||||
context_id=context.context_id,
|
||||
)
|
||||
|
||||
await updater.cancel()
|
||||
|
||||
@override
|
||||
async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
|
||||
"""Execute the agent with the given context and event queue.
|
||||
|
||||
Orchestrates the agent execution process: sets up the agent session,
|
||||
executes the agent, processes response messages, and handles errors with appropriate task status updates.
|
||||
"""
|
||||
if context.context_id is None:
|
||||
raise ValueError("Context ID must be provided in the RequestContext")
|
||||
if context.message is None:
|
||||
raise ValueError("Message must be provided in the RequestContext")
|
||||
|
||||
query = context.get_user_input()
|
||||
task = context.current_task
|
||||
|
||||
if not task:
|
||||
task = new_task_from_user_message(context.message)
|
||||
await event_queue.enqueue_event(task)
|
||||
|
||||
updater = TaskUpdater(event_queue, task.id, context.context_id)
|
||||
await updater.submit()
|
||||
|
||||
try:
|
||||
await updater.start_work()
|
||||
|
||||
session = self._agent.create_session(session_id=task.context_id)
|
||||
|
||||
if self._stream:
|
||||
await self._run_stream(query, session, updater)
|
||||
else:
|
||||
await self._run(query, session, updater)
|
||||
|
||||
# Mark as complete
|
||||
await updater.complete()
|
||||
except CancelledError:
|
||||
await updater.update_status(state=TaskState.TASK_STATE_CANCELED)
|
||||
except Exception as e:
|
||||
logger.exception("A2AExecutor encountered an error during execution.", exc_info=e)
|
||||
await updater.update_status(
|
||||
state=TaskState.TASK_STATE_FAILED,
|
||||
message=updater.new_agent_message([Part(text=str(e))]),
|
||||
)
|
||||
|
||||
async def _run_stream(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None:
|
||||
"""Run the agent in streaming mode and publish updates to the task updater."""
|
||||
response_stream = self._agent.run(query, session=session, stream=True, **self._run_kwargs)
|
||||
streamed_artifact_ids: set[str] = set()
|
||||
# Generate a stable artifact ID for the entire stream so all chunks share the same ID.
|
||||
# This ensures clients can coalesce streaming tokens into a single artifact/message
|
||||
# per the A2A spec (TaskArtifactUpdateEvent with append=True on same artifactId).
|
||||
default_artifact_id = str(uuid.uuid4())
|
||||
await (
|
||||
response_stream.with_transform_hook(
|
||||
partial(
|
||||
self.handle_events,
|
||||
updater=updater,
|
||||
streamed_artifact_ids=streamed_artifact_ids,
|
||||
default_artifact_id=default_artifact_id,
|
||||
)
|
||||
)
|
||||
).get_final_response()
|
||||
|
||||
async def _run(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None:
|
||||
"""Run the agent in non-streaming mode and publish messages to the task updater."""
|
||||
response = await self._agent.run(query, session=session, stream=False, **self._run_kwargs)
|
||||
response_messages = response.messages
|
||||
|
||||
if not isinstance(response_messages, list):
|
||||
response_messages = [response_messages]
|
||||
|
||||
for message in response_messages:
|
||||
await self.handle_events(message, updater)
|
||||
|
||||
async def handle_events(
|
||||
self,
|
||||
item: Message | AgentResponseUpdate,
|
||||
updater: TaskUpdater,
|
||||
streamed_artifact_ids: set[str] | None = None,
|
||||
default_artifact_id: str | None = None,
|
||||
) -> None:
|
||||
"""Convert agent response items (Messages or Updates) to A2A protocol events.
|
||||
|
||||
Processes Message or AgentResponseUpdate objects and converts them into A2A protocol format.
|
||||
Handles text, data, and URI content. USER role messages are skipped.
|
||||
|
||||
Users can override this method in a subclass to implement custom transformations
|
||||
from their agent's output format to A2A protocol events.
|
||||
|
||||
Args:
|
||||
item: The agent response item (Message or AgentResponseUpdate) to process.
|
||||
updater: The task updater to publish events to.
|
||||
streamed_artifact_ids: A set of artifact IDs that have already been streamed.
|
||||
Used to track which artifacts need append=True on subsequent chunks.
|
||||
default_artifact_id: A stable artifact ID to use when the item does not provide one.
|
||||
This ensures all streaming chunks for a single response share the same artifact ID,
|
||||
allowing clients to coalesce them into a single message.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
class CustomA2AExecutor(A2AExecutor):
|
||||
async def handle_events(
|
||||
self,
|
||||
item: Message | AgentResponseUpdate,
|
||||
updater: TaskUpdater,
|
||||
streamed_artifact_ids: set[str] | None = None,
|
||||
default_artifact_id: str | None = None,
|
||||
) -> None:
|
||||
# Custom logic to transform item contents
|
||||
if item.role == "assistant" and item.contents:
|
||||
parts = [Part(text=f"Custom: {item.contents[0].text}")]
|
||||
await updater.update_status(
|
||||
state=TaskState.TASK_STATE_WORKING,
|
||||
message=updater.new_agent_message(parts=parts),
|
||||
)
|
||||
else:
|
||||
await super().handle_events(item, updater)
|
||||
"""
|
||||
role = getattr(item, "role", None)
|
||||
if role == "user":
|
||||
# This is a user message, we can ignore it in the context of task updates
|
||||
return
|
||||
|
||||
parts: list[Part] = []
|
||||
metadata = getattr(item, "additional_properties", None)
|
||||
|
||||
# AgentResponseUpdate uses 'contents', Message uses 'contents'
|
||||
contents = getattr(item, "contents", [])
|
||||
|
||||
for content in contents:
|
||||
if content.type == "text" and content.text:
|
||||
parts.append(Part(text=content.text))
|
||||
elif content.type == "data" and content.uri:
|
||||
base64_str = get_uri_data(content.uri)
|
||||
parts.append(Part(raw=base64.b64decode(base64_str), media_type=content.media_type or ""))
|
||||
elif content.type == "uri" and content.uri:
|
||||
parts.append(Part(url=content.uri, media_type=content.media_type or ""))
|
||||
else:
|
||||
# Silently skip unsupported content types
|
||||
logger.warning("A2AExecutor does not yet support content type: %s. Omitted.", content.type)
|
||||
|
||||
if parts:
|
||||
if isinstance(item, AgentResponseUpdate):
|
||||
# Resolve artifact ID: use item's message_id if available, otherwise fall back
|
||||
# to the stable default_artifact_id so all streaming chunks share the same ID.
|
||||
artifact_id = item.message_id or default_artifact_id
|
||||
# For streaming updates, we send TaskArtifactUpdateEvent via add_artifact
|
||||
await updater.add_artifact(
|
||||
parts=parts,
|
||||
artifact_id=artifact_id,
|
||||
metadata=metadata,
|
||||
append=(
|
||||
True if streamed_artifact_ids is not None and artifact_id in streamed_artifact_ids else None
|
||||
),
|
||||
)
|
||||
if artifact_id and streamed_artifact_ids is not None:
|
||||
streamed_artifact_ids.add(artifact_id)
|
||||
else:
|
||||
# For final messages, we send TaskStatusUpdateEvent with 'working' state
|
||||
await updater.update_status(
|
||||
state=TaskState.TASK_STATE_WORKING,
|
||||
message=updater.new_agent_message(parts=parts, metadata=metadata),
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,24 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import re
|
||||
|
||||
URI_PATTERN = re.compile(r"^data:(?P<media_type>[^;,]+(?:;[^;,=]+=[^;,]+)*);base64,(?P<base64_data>[A-Za-z0-9+/=]+)\Z")
|
||||
|
||||
|
||||
def get_uri_data(uri: str) -> str:
|
||||
"""Extracts the base64-encoded data from a data URI.
|
||||
|
||||
Args:
|
||||
uri: The data URI to parse.
|
||||
|
||||
Returns:
|
||||
The base64-encoded data part of the URI.
|
||||
|
||||
Raises:
|
||||
ValueError: If the URI format is invalid.
|
||||
"""
|
||||
match = URI_PATTERN.match(uri)
|
||||
if not match:
|
||||
raise ValueError(f"Invalid data URI format: {uri}")
|
||||
|
||||
return match.group("base64_data")
|
||||
@@ -0,0 +1,98 @@
|
||||
[project]
|
||||
name = "agent-framework-a2a"
|
||||
description = "A2A integration for Microsoft Agent Framework."
|
||||
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
version = "1.0.0b260709"
|
||||
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 :: 4 - Beta",
|
||||
"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 = [
|
||||
"agent-framework-core>=1.11.0,<2",
|
||||
"a2a-sdk>=1.0.0,<2",
|
||||
]
|
||||
|
||||
[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 = [
|
||||
"ignore:Support for class-based `config` is deprecated:DeprecationWarning:pydantic.*"
|
||||
]
|
||||
timeout = 120
|
||||
markers = [
|
||||
"integration: marks tests as integration tests that require external services",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
extend = "../../pyproject.toml"
|
||||
|
||||
[tool.coverage.run]
|
||||
omit = [
|
||||
"**/__init__.py"
|
||||
]
|
||||
|
||||
[tool.pyright]
|
||||
extends = "../../pyproject.toml"
|
||||
include = ["agent_framework_a2a"]
|
||||
|
||||
[tool.mypy]
|
||||
plugins = ['pydantic.mypy']
|
||||
strict = true
|
||||
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_a2a"]
|
||||
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_a2a"
|
||||
|
||||
[tool.poe.tasks.test]
|
||||
help = "Run the default unit test suite for this package."
|
||||
cmd = 'pytest -m "not integration" --cov=agent_framework_a2a --cov-report=term-missing:skip-covered tests'
|
||||
|
||||
[build-system]
|
||||
requires = ["flit-core >= 3.11,<4.0"]
|
||||
build-backend = "flit_core.buildapi"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,907 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from asyncio import CancelledError
|
||||
from typing import Any, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from uuid import uuid4
|
||||
|
||||
from a2a.types import Part, Task, TaskState
|
||||
from agent_framework import (
|
||||
AgentResponseUpdate,
|
||||
Content,
|
||||
Message,
|
||||
SupportsAgentRun,
|
||||
)
|
||||
from agent_framework._types import AgentResponse
|
||||
from agent_framework.a2a import A2AExecutor
|
||||
from pytest import fixture, raises
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_agent() -> MagicMock:
|
||||
"""Fixture that provides a mock SupportsAgentRun."""
|
||||
agent = MagicMock(spec=SupportsAgentRun)
|
||||
agent.run = AsyncMock()
|
||||
return agent
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_request_context() -> MagicMock:
|
||||
"""Fixture that provides a mock RequestContext."""
|
||||
request_context = MagicMock()
|
||||
request_context.context_id = str(uuid4())
|
||||
request_context.get_user_input = MagicMock(return_value="Test query")
|
||||
request_context.current_task = None
|
||||
request_context.message = None
|
||||
return request_context
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_event_queue() -> MagicMock:
|
||||
"""Fixture that provides a mock EventQueue."""
|
||||
queue = AsyncMock()
|
||||
queue.enqueue_event = AsyncMock()
|
||||
return queue
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_task() -> Task:
|
||||
"""Fixture that provides a mock Task."""
|
||||
task = MagicMock(spec=Task)
|
||||
task.id = str(uuid4())
|
||||
task.context_id = str(uuid4())
|
||||
task.state = TaskState.TASK_STATE_COMPLETED
|
||||
return task
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_task_updater() -> MagicMock:
|
||||
"""Fixture that provides a mock TaskUpdater."""
|
||||
updater = MagicMock()
|
||||
updater.submit = AsyncMock()
|
||||
updater.start_work = AsyncMock()
|
||||
updater.complete = AsyncMock()
|
||||
updater.update_status = AsyncMock()
|
||||
updater.new_agent_message = MagicMock()
|
||||
return updater
|
||||
|
||||
|
||||
@fixture
|
||||
def executor(mock_agent: MagicMock) -> A2AExecutor:
|
||||
"""Fixture that provides an A2AExecutor."""
|
||||
return A2AExecutor(agent=mock_agent)
|
||||
|
||||
|
||||
class TestA2AExecutorInitialization:
|
||||
"""Tests for A2AExecutor initialization."""
|
||||
|
||||
def test_initialization_with_agent_only(self, mock_agent: MagicMock) -> None:
|
||||
"""Arrange: Create mock agent
|
||||
Act: Initialize A2AExecutor with only agent
|
||||
Assert: Executor is created with default values
|
||||
"""
|
||||
# Act
|
||||
executor = A2AExecutor(agent=mock_agent)
|
||||
|
||||
# Assert
|
||||
assert executor._agent is mock_agent
|
||||
assert executor._stream is False
|
||||
assert executor._run_kwargs == {}
|
||||
|
||||
def test_initialization_with_stream_and_kwargs(self, mock_agent: MagicMock) -> None:
|
||||
"""Arrange: Create mock agent
|
||||
Act: Initialize A2AExecutor with stream and run_kwargs
|
||||
Assert: Executor is created with specified values
|
||||
"""
|
||||
# Arrange
|
||||
run_kwargs = {"temperature": 0.5}
|
||||
|
||||
# Act
|
||||
executor = A2AExecutor(agent=mock_agent, stream=True, run_kwargs=run_kwargs)
|
||||
|
||||
# Assert
|
||||
assert executor._agent is mock_agent
|
||||
assert executor._stream is True
|
||||
assert executor._run_kwargs == run_kwargs
|
||||
|
||||
def test_initialization_with_invalid_run_kwargs(self, mock_agent: MagicMock) -> None:
|
||||
"""Arrange: Create mock agent
|
||||
Act: Initialize A2AExecutor with reserved keys in run_kwargs
|
||||
Assert: ValueError is raised
|
||||
"""
|
||||
# Act & Assert
|
||||
with raises(ValueError, match="run_kwargs cannot contain 'session'"):
|
||||
A2AExecutor(agent=mock_agent, run_kwargs={"session": "something"})
|
||||
|
||||
with raises(ValueError, match="run_kwargs cannot contain 'stream'"):
|
||||
A2AExecutor(agent=mock_agent, run_kwargs={"stream": True})
|
||||
|
||||
|
||||
class TestA2AExecutorCancel:
|
||||
"""Tests for the cancel method."""
|
||||
|
||||
async def test_cancel_method_completes(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
) -> None:
|
||||
"""Arrange: Create executor with dependencies
|
||||
Act: Call cancel method
|
||||
Assert: Method completes without raising error
|
||||
"""
|
||||
# Arrange
|
||||
mock_request_context.task_id = "task-123"
|
||||
|
||||
# Act & Assert (should not raise)
|
||||
await executor.cancel(mock_request_context, mock_event_queue) # type: ignore
|
||||
|
||||
async def test_cancel_handles_different_contexts(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_event_queue: MagicMock,
|
||||
) -> None:
|
||||
"""Arrange: Create executor with multiple request contexts
|
||||
Act: Call cancel with different contexts
|
||||
Assert: Each cancel completes successfully
|
||||
"""
|
||||
# Arrange
|
||||
context1 = MagicMock()
|
||||
context1.context_id = "ctx-1"
|
||||
context1.task_id = "task-1"
|
||||
context2 = MagicMock()
|
||||
context2.context_id = "ctx-2"
|
||||
context2.task_id = "task-2"
|
||||
|
||||
# Act & Assert
|
||||
await executor.cancel(context1, mock_event_queue) # type: ignore
|
||||
await executor.cancel(context2, mock_event_queue) # type: ignore
|
||||
|
||||
async def test_cancel_raises_error_when_context_id_missing(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_event_queue: MagicMock,
|
||||
) -> None:
|
||||
"""Arrange: Create context without context_id
|
||||
Act: Call cancel method
|
||||
Assert: ValueError is raised
|
||||
"""
|
||||
# Arrange
|
||||
mock_context = MagicMock()
|
||||
mock_context.context_id = None
|
||||
|
||||
# Act & Assert
|
||||
with raises(ValueError) as excinfo:
|
||||
await executor.cancel(mock_context, mock_event_queue) # type: ignore
|
||||
|
||||
# Assert
|
||||
assert "Context ID" in str(excinfo.value)
|
||||
|
||||
|
||||
class TestA2AExecutorExecute:
|
||||
"""Tests for the execute method."""
|
||||
|
||||
async def test_execute_with_existing_task_succeeds(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
mock_task: Task,
|
||||
) -> None:
|
||||
"""Arrange: Create executor with mocked dependencies and existing task
|
||||
Act: Call execute method
|
||||
Assert: Execution completes successfully
|
||||
"""
|
||||
# Arrange
|
||||
mock_request_context.get_user_input = MagicMock(return_value="Hello")
|
||||
mock_request_context.current_task = mock_task
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
mock_request_context.message = MagicMock()
|
||||
|
||||
response_message = Message(role="assistant", contents=[Content.from_text(text="Hello back")])
|
||||
response = MagicMock(spec=AgentResponse)
|
||||
response.messages = [response_message]
|
||||
cast(Any, executor._agent).run = AsyncMock(return_value=response)
|
||||
cast(Any, executor._agent).create_session = MagicMock()
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.submit = AsyncMock()
|
||||
mock_updater.start_work = AsyncMock()
|
||||
mock_updater.complete = AsyncMock()
|
||||
mock_updater.update_status = AsyncMock()
|
||||
mock_updater.new_agent_message = MagicMock(return_value="message_obj")
|
||||
mock_updater_class.return_value = mock_updater
|
||||
|
||||
# Act
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
mock_updater.submit.assert_called_once()
|
||||
mock_updater.start_work.assert_called_once()
|
||||
mock_updater.complete.assert_called_once()
|
||||
cast(Any, executor._agent.create_session).assert_called_once()
|
||||
cast(Any, executor._agent.run).assert_called_once()
|
||||
|
||||
async def test_execute_creates_task_when_not_exists(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
) -> None:
|
||||
"""Arrange: Create executor with request context without task
|
||||
Act: Call execute method
|
||||
Assert: New task is created and enqueued
|
||||
"""
|
||||
# Arrange
|
||||
mock_message = MagicMock()
|
||||
mock_request_context.get_user_input = MagicMock(return_value="Hello")
|
||||
mock_request_context.current_task = None
|
||||
mock_request_context.message = mock_message
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
|
||||
response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
|
||||
response = MagicMock(spec=AgentResponse)
|
||||
response.messages = [response_message]
|
||||
cast(Any, executor._agent).run = AsyncMock(return_value=response)
|
||||
cast(Any, executor._agent).create_session = MagicMock()
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.new_task_from_user_message") as mock_new_task:
|
||||
mock_task = MagicMock(spec=Task)
|
||||
mock_task.id = "task-new"
|
||||
mock_task.context_id = "ctx-123"
|
||||
mock_new_task.return_value = mock_task
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.submit = AsyncMock()
|
||||
mock_updater.start_work = AsyncMock()
|
||||
mock_updater.complete = AsyncMock()
|
||||
mock_updater.update_status = AsyncMock()
|
||||
mock_updater.new_agent_message = MagicMock(return_value="message_obj")
|
||||
mock_updater_class.return_value = mock_updater
|
||||
|
||||
# Act
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
mock_new_task.assert_called_once()
|
||||
mock_event_queue.enqueue_event.assert_called_once()
|
||||
|
||||
async def test_execute_raises_error_when_context_id_missing(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
) -> None:
|
||||
"""Arrange: Create context without context_id
|
||||
Act: Call execute method
|
||||
Assert: ValueError is raised
|
||||
"""
|
||||
# Arrange
|
||||
mock_request_context.context_id = None
|
||||
mock_request_context.message = MagicMock()
|
||||
|
||||
# Act & Assert
|
||||
with raises(ValueError) as excinfo:
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
assert "Context ID" in str(excinfo.value)
|
||||
|
||||
async def test_execute_raises_error_when_message_missing(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
) -> None:
|
||||
"""Arrange: Create context without message
|
||||
Act: Call execute method
|
||||
Assert: ValueError is raised
|
||||
"""
|
||||
# Arrange
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
mock_request_context.message = None
|
||||
|
||||
# Act & Assert
|
||||
with raises(ValueError) as excinfo:
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
assert "Message" in str(excinfo.value)
|
||||
|
||||
async def test_execute_handles_cancelled_error(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
mock_task: Task,
|
||||
) -> None:
|
||||
"""Arrange: Create executor that raises CancelledError
|
||||
Act: Call execute method
|
||||
Assert: Error is caught and task is marked as canceled
|
||||
"""
|
||||
# Arrange
|
||||
mock_request_context.get_user_input = MagicMock(return_value="Hello")
|
||||
mock_request_context.current_task = mock_task
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
mock_request_context.message = MagicMock()
|
||||
|
||||
cast(Any, executor._agent).run = AsyncMock(side_effect=CancelledError())
|
||||
cast(Any, executor._agent).create_session = MagicMock()
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.submit = AsyncMock()
|
||||
mock_updater.start_work = AsyncMock()
|
||||
mock_updater.update_status = AsyncMock()
|
||||
mock_updater_class.return_value = mock_updater
|
||||
|
||||
# Act
|
||||
await executor.execute(mock_request_context, mock_event_queue) # type: ignore
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_called()
|
||||
call_args_list = mock_updater.update_status.call_args_list
|
||||
assert any(call[1].get("state") == TaskState.TASK_STATE_CANCELED for call in call_args_list)
|
||||
|
||||
async def test_execute_handles_generic_exception(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
mock_task: Task,
|
||||
) -> None:
|
||||
"""Arrange: Create executor that raises generic exception
|
||||
Act: Call execute method
|
||||
Assert: Error is caught and task is marked as failed
|
||||
"""
|
||||
# Arrange
|
||||
mock_request_context.get_user_input = MagicMock(return_value="Hello")
|
||||
mock_request_context.current_task = mock_task
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
mock_request_context.message = MagicMock()
|
||||
|
||||
error_message = "Test error"
|
||||
cast(Any, executor._agent).run = AsyncMock(side_effect=ValueError(error_message))
|
||||
cast(Any, executor._agent).create_session = MagicMock()
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.submit = AsyncMock()
|
||||
mock_updater.start_work = AsyncMock()
|
||||
mock_updater.update_status = AsyncMock()
|
||||
mock_updater.new_agent_message = MagicMock(return_value="error_message_obj")
|
||||
mock_updater_class.return_value = mock_updater
|
||||
|
||||
# Act
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
mock_updater.new_agent_message.assert_called_once()
|
||||
args, _ = mock_updater.new_agent_message.call_args
|
||||
parts = args[0]
|
||||
assert len(parts) == 1
|
||||
assert isinstance(parts[0], Part)
|
||||
assert parts[0].text == error_message
|
||||
|
||||
call_args_list = mock_updater.update_status.call_args_list
|
||||
assert any(
|
||||
call[1].get("state") == TaskState.TASK_STATE_FAILED and call[1].get("message") == "error_message_obj"
|
||||
for call in call_args_list
|
||||
)
|
||||
|
||||
async def test_execute_processes_multiple_response_messages(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
mock_task: Task,
|
||||
) -> None:
|
||||
"""Arrange: Create executor that returns multiple response messages
|
||||
Act: Call execute method
|
||||
Assert: All messages are processed through handle_events
|
||||
"""
|
||||
# Arrange
|
||||
mock_request_context.get_user_input = MagicMock(return_value="Hello")
|
||||
mock_request_context.current_task = mock_task
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
mock_request_context.message = MagicMock()
|
||||
|
||||
response_message1 = Message(role="assistant", contents=[Content.from_text(text="First")])
|
||||
response_message2 = Message(role="assistant", contents=[Content.from_text(text="Second")])
|
||||
response = MagicMock(spec=AgentResponse)
|
||||
response.messages = [response_message1, response_message2]
|
||||
cast(Any, executor._agent).run = AsyncMock(return_value=response)
|
||||
cast(Any, executor._agent).create_session = MagicMock()
|
||||
|
||||
# Mock handle_events
|
||||
cast(Any, executor).handle_events = AsyncMock()
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.submit = AsyncMock()
|
||||
mock_updater.start_work = AsyncMock()
|
||||
mock_updater.complete = AsyncMock()
|
||||
mock_updater.update_status = AsyncMock()
|
||||
mock_updater_class.return_value = mock_updater
|
||||
|
||||
# Act
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
assert cast(Any, executor.handle_events).call_count == 2
|
||||
|
||||
async def test_execute_passes_query_to_run(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
mock_task: Task,
|
||||
) -> None:
|
||||
"""Arrange: Create executor with request
|
||||
Act: Call execute method
|
||||
Assert: Query text is passed to run method with default stream and kwargs
|
||||
"""
|
||||
# Arrange
|
||||
query_text = "Hello agent"
|
||||
mock_request_context.get_user_input = MagicMock(return_value=query_text)
|
||||
mock_request_context.current_task = mock_task
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
mock_request_context.message = MagicMock()
|
||||
|
||||
response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
|
||||
response = MagicMock(spec=AgentResponse)
|
||||
response.messages = [response_message]
|
||||
cast(Any, executor._agent).run = AsyncMock(return_value=response)
|
||||
cast(Any, executor._agent).create_session = MagicMock()
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.submit = AsyncMock()
|
||||
mock_updater.start_work = AsyncMock()
|
||||
mock_updater.complete = AsyncMock()
|
||||
mock_updater.update_status = AsyncMock()
|
||||
mock_updater.new_agent_message = MagicMock(return_value="message_obj")
|
||||
mock_updater_class.return_value = mock_updater
|
||||
|
||||
# Act
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
cast(Any, executor._agent.run).assert_called_once_with(
|
||||
query_text, session=executor._agent.create_session(), stream=False
|
||||
)
|
||||
|
||||
async def test_execute_with_stream_enabled(
|
||||
self,
|
||||
mock_agent: MagicMock,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
mock_task: Task,
|
||||
) -> None:
|
||||
"""Arrange: Create executor with stream=True
|
||||
Act: Call execute method
|
||||
Assert: _run_stream is called and passes stream=True to run
|
||||
"""
|
||||
# Arrange
|
||||
executor = A2AExecutor(agent=mock_agent, stream=True)
|
||||
query_text = "Hello agent"
|
||||
mock_request_context.get_user_input = MagicMock(return_value=query_text)
|
||||
mock_request_context.current_task = mock_task
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
mock_request_context.message = MagicMock()
|
||||
|
||||
mock_response_stream = MagicMock()
|
||||
mock_response_stream.with_transform_hook = MagicMock(return_value=mock_response_stream)
|
||||
mock_response_stream.get_final_response = AsyncMock()
|
||||
mock_agent.run = MagicMock(return_value=mock_response_stream)
|
||||
mock_agent.create_session = MagicMock()
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.submit = AsyncMock()
|
||||
mock_updater.start_work = AsyncMock()
|
||||
mock_updater.complete = AsyncMock()
|
||||
mock_updater.update_status = AsyncMock()
|
||||
mock_updater_class.return_value = mock_updater
|
||||
|
||||
# Act
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
mock_agent.run.assert_called_once_with(query_text, session=mock_agent.create_session(), stream=True)
|
||||
mock_response_stream.with_transform_hook.assert_called_once()
|
||||
mock_response_stream.get_final_response.assert_called_once()
|
||||
|
||||
async def test_execute_with_run_kwargs(
|
||||
self,
|
||||
mock_agent: MagicMock,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
mock_task: Task,
|
||||
) -> None:
|
||||
"""Arrange: Create executor with run_kwargs
|
||||
Act: Call execute method
|
||||
Assert: run_kwargs are passed to run method
|
||||
"""
|
||||
# Arrange
|
||||
run_kwargs = {"temperature": 0.5, "max_tokens": 100}
|
||||
executor = A2AExecutor(agent=mock_agent, run_kwargs=run_kwargs)
|
||||
query_text = "Hello agent"
|
||||
mock_request_context.get_user_input = MagicMock(return_value=query_text)
|
||||
mock_request_context.current_task = mock_task
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
mock_request_context.message = MagicMock()
|
||||
|
||||
response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
|
||||
response = MagicMock(spec=AgentResponse)
|
||||
response.messages = [response_message]
|
||||
mock_agent.run = AsyncMock(return_value=response)
|
||||
mock_agent.create_session = MagicMock()
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.submit = AsyncMock()
|
||||
mock_updater.start_work = AsyncMock()
|
||||
mock_updater.complete = AsyncMock()
|
||||
mock_updater.update_status = AsyncMock()
|
||||
mock_updater_class.return_value = mock_updater
|
||||
|
||||
# Act
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
mock_agent.run.assert_called_once_with(
|
||||
query_text, session=mock_agent.create_session(), stream=False, **run_kwargs
|
||||
)
|
||||
|
||||
|
||||
class TestA2AExecutorHandleEvents:
|
||||
"""Tests for A2AExecutor.handle_events method."""
|
||||
|
||||
async def test_run_method_with_single_message(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test the private _run method with a single message (not a list)."""
|
||||
# Arrange
|
||||
query = "test query"
|
||||
session = MagicMock()
|
||||
response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
|
||||
response = MagicMock(spec=AgentResponse)
|
||||
response.messages = response_message # Not a list
|
||||
cast(Any, executor._agent).run = AsyncMock(return_value=response)
|
||||
cast(Any, executor).handle_events = AsyncMock()
|
||||
|
||||
# Act
|
||||
await executor._run(query, session, mock_updater)
|
||||
|
||||
# Assert
|
||||
cast(Any, executor.handle_events).assert_called_once_with(response_message, mock_updater)
|
||||
|
||||
@fixture
|
||||
def mock_updater(self) -> MagicMock:
|
||||
"""Create a mock execution context."""
|
||||
updater = MagicMock()
|
||||
updater.update_status = AsyncMock()
|
||||
updater.new_agent_message = MagicMock(return_value="mock_message")
|
||||
return updater
|
||||
|
||||
async def test_ignore_user_messages(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test that messages from USER role are ignored."""
|
||||
# Arrange
|
||||
message = Message(
|
||||
contents=[Content.from_text(text="User input")],
|
||||
role="user",
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_not_called()
|
||||
|
||||
async def test_ignore_messages_with_no_contents(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test that messages with no contents are ignored."""
|
||||
# Arrange
|
||||
message = Message(
|
||||
contents=[],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_not_called()
|
||||
|
||||
async def test_handle_text_content(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test handling messages with text content."""
|
||||
# Arrange
|
||||
text = "Hello, this is a test message"
|
||||
message = Message(
|
||||
contents=[Content.from_text(text=text)],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_called_once()
|
||||
call_args = mock_updater.update_status.call_args
|
||||
assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
|
||||
assert mock_updater.new_agent_message.called
|
||||
|
||||
async def test_handle_multiple_text_contents(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test handling messages with multiple text contents."""
|
||||
# Arrange
|
||||
message = Message(
|
||||
contents=[
|
||||
Content.from_text(text="First message"),
|
||||
Content.from_text(text="Second message"),
|
||||
],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_called_once()
|
||||
assert mock_updater.new_agent_message.called
|
||||
|
||||
async def test_handle_data_content(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test handling messages with data content."""
|
||||
# Arrange
|
||||
data = b"test file data"
|
||||
message = Message(
|
||||
contents=[Content.from_data(data=data, media_type="application/octet-stream")],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_called_once()
|
||||
call_args = mock_updater.update_status.call_args
|
||||
assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
|
||||
|
||||
async def test_handle_uri_content(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test handling messages with URI content."""
|
||||
# Arrange
|
||||
uri = "https://example.com/file.pdf"
|
||||
message = Message(
|
||||
contents=[Content.from_uri(uri=uri, media_type="application/pdf")],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_called_once()
|
||||
call_args = mock_updater.update_status.call_args
|
||||
assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
|
||||
|
||||
async def test_handle_mixed_content_types(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test handling messages with mixed content types."""
|
||||
# Arrange
|
||||
data = b"file data"
|
||||
|
||||
message = Message(
|
||||
contents=[
|
||||
Content.from_text(text="Processing file..."),
|
||||
Content.from_data(data=data, media_type="application/octet-stream"),
|
||||
Content.from_uri(uri="https://example.com/reference.pdf", media_type="application/pdf"),
|
||||
],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_called_once()
|
||||
call_args = mock_updater.update_status.call_args
|
||||
assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
|
||||
|
||||
async def test_handle_with_additional_properties(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test handling messages with additional properties metadata."""
|
||||
# Arrange
|
||||
additional_props = {"custom_field": "custom_value", "priority": "high"}
|
||||
message = Message(
|
||||
contents=[Content.from_text(text="Test message")],
|
||||
role="assistant",
|
||||
additional_properties=additional_props,
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_called_once()
|
||||
mock_updater.new_agent_message.assert_called_once()
|
||||
call_args = mock_updater.new_agent_message.call_args
|
||||
assert call_args.kwargs["metadata"] == additional_props
|
||||
|
||||
async def test_handle_with_no_additional_properties(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test handling messages without additional properties."""
|
||||
# Arrange
|
||||
message = Message(
|
||||
contents=[Content.from_text(text="Test message")],
|
||||
role="assistant",
|
||||
additional_properties=None,
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.update_status.assert_called_once()
|
||||
mock_updater.new_agent_message.assert_called_once()
|
||||
call_args = mock_updater.new_agent_message.call_args
|
||||
assert call_args.kwargs["metadata"] == {}
|
||||
|
||||
async def test_parts_list_passed_to_new_agent_message(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test that parts list is correctly passed to new_agent_message."""
|
||||
# Arrange
|
||||
message = Message(
|
||||
contents=[
|
||||
Content.from_text(text="Message 1"),
|
||||
Content.from_text(text="Message 2"),
|
||||
],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.new_agent_message.assert_called_once()
|
||||
call_kwargs = mock_updater.new_agent_message.call_args.kwargs
|
||||
assert "parts" in call_kwargs
|
||||
parts_list = call_kwargs["parts"]
|
||||
assert len(parts_list) == 2
|
||||
|
||||
async def test_task_state_always_working(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test that task state is always set to working."""
|
||||
# Arrange
|
||||
message = Message(
|
||||
contents=[Content.from_text(text="Any message")],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
# Act
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
call_kwargs = mock_updater.update_status.call_args.kwargs
|
||||
assert call_kwargs["state"] == TaskState.TASK_STATE_WORKING
|
||||
|
||||
async def test_handle_agent_response_update_no_streamed_set(
|
||||
self, executor: A2AExecutor, mock_updater: MagicMock
|
||||
) -> None:
|
||||
"""Test handling AgentResponseUpdate (streaming) without a tracking set."""
|
||||
# Arrange
|
||||
update = AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Streaming chunk")],
|
||||
role="assistant",
|
||||
message_id="msg-1",
|
||||
)
|
||||
mock_updater.add_artifact = AsyncMock()
|
||||
|
||||
# Act
|
||||
await executor.handle_events(update, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_updater.add_artifact.assert_called_once()
|
||||
call_kwargs = mock_updater.add_artifact.call_args.kwargs
|
||||
assert call_kwargs["artifact_id"] == "msg-1"
|
||||
assert call_kwargs["append"] is None
|
||||
|
||||
async def test_handle_agent_response_update_first_time(
|
||||
self, executor: A2AExecutor, mock_updater: MagicMock
|
||||
) -> None:
|
||||
"""Test handling AgentResponseUpdate (streaming) for the first time with a tracking set."""
|
||||
# Arrange
|
||||
update = AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Streaming chunk")],
|
||||
role="assistant",
|
||||
message_id="msg-1",
|
||||
)
|
||||
mock_updater.add_artifact = AsyncMock()
|
||||
streamed_artifact_ids: set[str] = set()
|
||||
|
||||
# Act
|
||||
await executor.handle_events(update, mock_updater, streamed_artifact_ids=streamed_artifact_ids)
|
||||
|
||||
# Assert
|
||||
mock_updater.add_artifact.assert_called_once()
|
||||
call_kwargs = mock_updater.add_artifact.call_args.kwargs
|
||||
assert call_kwargs["append"] is None
|
||||
assert "msg-1" in streamed_artifact_ids
|
||||
|
||||
async def test_handle_agent_response_update_subsequent_time(
|
||||
self, executor: A2AExecutor, mock_updater: MagicMock
|
||||
) -> None:
|
||||
"""Test handling AgentResponseUpdate (streaming) for subsequent times with a tracking set."""
|
||||
# Arrange
|
||||
update = AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Next chunk")],
|
||||
role="assistant",
|
||||
message_id="msg-1",
|
||||
)
|
||||
mock_updater.add_artifact = AsyncMock()
|
||||
streamed_artifact_ids = {"msg-1"}
|
||||
|
||||
# Act
|
||||
await executor.handle_events(update, mock_updater, streamed_artifact_ids=streamed_artifact_ids)
|
||||
|
||||
# Assert
|
||||
mock_updater.add_artifact.assert_called_once()
|
||||
call_kwargs = mock_updater.add_artifact.call_args.kwargs
|
||||
assert call_kwargs["append"] is True
|
||||
|
||||
async def test_handle_unsupported_content_type(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
|
||||
"""Test handling messages with unsupported content types."""
|
||||
# Arrange
|
||||
message = Message(
|
||||
contents=[Content(type=cast(Any, "unknown"), text="Some text")], # type: ignore[arg-type]
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
# Act
|
||||
with patch("agent_framework_a2a._a2a_executor.logger") as mock_logger:
|
||||
await executor.handle_events(message, mock_updater)
|
||||
|
||||
# Assert
|
||||
mock_logger.warning.assert_called_once()
|
||||
mock_updater.update_status.assert_not_called()
|
||||
|
||||
|
||||
class TestA2AExecutorIntegration:
|
||||
"""Integration tests for A2AExecutor."""
|
||||
|
||||
async def test_full_execution_flow_with_responses(
|
||||
self,
|
||||
executor: A2AExecutor,
|
||||
mock_request_context: MagicMock,
|
||||
mock_event_queue: MagicMock,
|
||||
mock_task: Task,
|
||||
) -> None:
|
||||
"""Arrange: Create executor with all mocked dependencies
|
||||
Act: Execute full flow from request to completion
|
||||
Assert: All components interact correctly
|
||||
"""
|
||||
# Arrange
|
||||
mock_request_context.get_user_input = MagicMock(return_value="Hello agent")
|
||||
mock_request_context.current_task = mock_task
|
||||
mock_request_context.context_id = "ctx-123"
|
||||
mock_request_context.message = MagicMock()
|
||||
|
||||
response = MagicMock(spec=AgentResponse)
|
||||
response_message = MagicMock(spec=Message)
|
||||
response.messages = [response_message]
|
||||
response_message.contents = [Content.from_text(text="Hello user")]
|
||||
response_message.role = "assistant"
|
||||
response_message.additional_properties = None
|
||||
|
||||
cast(Any, executor._agent).run = AsyncMock(return_value=response)
|
||||
cast(Any, executor._agent).create_session = MagicMock()
|
||||
cast(Any, executor).handle_events = AsyncMock()
|
||||
|
||||
with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.submit = AsyncMock()
|
||||
mock_updater.start_work = AsyncMock()
|
||||
mock_updater.complete = AsyncMock()
|
||||
mock_updater.update_status = AsyncMock()
|
||||
mock_updater_class.return_value = mock_updater
|
||||
|
||||
# Act
|
||||
await executor.execute(mock_request_context, mock_event_queue)
|
||||
|
||||
# Assert
|
||||
mock_updater.submit.assert_called_once()
|
||||
mock_updater.start_work.assert_called_once()
|
||||
cast(Any, executor.handle_events).assert_called_once()
|
||||
mock_updater.complete.assert_called_once()
|
||||
@@ -0,0 +1,51 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework_a2a._utils import get_uri_data
|
||||
|
||||
|
||||
def test_get_uri_data_valid() -> None:
|
||||
"""Test get_uri_data with valid data URIs."""
|
||||
# Simple text/plain
|
||||
uri = "data:text/plain;base64,SGVsbG8sIFdvcmxkIQ=="
|
||||
assert get_uri_data(uri) == "SGVsbG8sIFdvcmxkIQ=="
|
||||
|
||||
# Image png
|
||||
uri = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
|
||||
assert get_uri_data(uri) == "iVBORw0KGgoAAAANSUhEUgfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
|
||||
|
||||
# Application octet-stream
|
||||
uri = "data:application/octet-stream;base64,AQIDBA=="
|
||||
assert get_uri_data(uri) == "AQIDBA=="
|
||||
|
||||
# Media type with parameters
|
||||
uri = "data:text/plain;charset=utf-8;base64,SGVsbG8sIFdvcmxkIQ=="
|
||||
assert get_uri_data(uri) == "SGVsbG8sIFdvcmxkIQ=="
|
||||
|
||||
# Media type with multiple parameters
|
||||
uri = "data:text/plain;charset=utf-8;name=hello.txt;base64,SGVsbG8sIFdvcmxkIQ=="
|
||||
assert get_uri_data(uri) == "SGVsbG8sIFdvcmxkIQ=="
|
||||
|
||||
|
||||
def test_get_uri_data_invalid_format() -> None:
|
||||
"""Test get_uri_data with invalid URI formats."""
|
||||
invalid_uris = [
|
||||
"not-a-uri",
|
||||
"http://example.com",
|
||||
"data:text/plain;SGVsbG8sIFdvcmxkIQ==", # Missing base64 marker
|
||||
"data:base64,SGVsbG8sIFdvcmxkIQ==", # Missing media type
|
||||
"data:text/plain;foo;base64,SGVsbG8sIFdvcmxkIQ==", # Parameter without value
|
||||
"data:text/plain;base64;base64,SGVsbG8sIFdvcmxkIQ==", # base64 used as a parameter name
|
||||
"data:text/plain;base64,SGVsbG8sIFdvcmxkIQ== extra",
|
||||
"data:text/plain;base64,SGVsbG8sIFdvcmxkIQ==\n",
|
||||
]
|
||||
for uri in invalid_uris:
|
||||
with pytest.raises(ValueError, match="Invalid data URI format"):
|
||||
get_uri_data(uri)
|
||||
|
||||
|
||||
def test_get_uri_data_empty() -> None:
|
||||
"""Test get_uri_data with empty string."""
|
||||
with pytest.raises(ValueError, match="Invalid data URI format"):
|
||||
get_uri_data("")
|
||||
@@ -0,0 +1,51 @@
|
||||
# AG-UI Package (agent-framework-ag-ui)
|
||||
|
||||
AG-UI protocol integration for building agent UIs with the AG-UI standard.
|
||||
|
||||
## Main Classes
|
||||
|
||||
- **`AgentFrameworkAgent`** - Wraps agents for AG-UI compatibility
|
||||
- **`AgentFrameworkWorkflow`** - Wraps native `Workflow` objects, or accepts `workflow_factory(thread_id)` for thread-scoped workflow instances without subclassing
|
||||
- **`AGUIChatClient`** - Chat client that speaks AG-UI protocol
|
||||
- **`AGUIHttpService`** - HTTP service for AG-UI endpoints
|
||||
- **`AGUIEventConverter`** - Converts between Agent Framework and AG-UI events
|
||||
- **`add_agent_framework_fastapi_endpoint()`** - Add AG-UI endpoint to FastAPI app (`SupportsAgentRun` or `Workflow`)
|
||||
- **`InMemoryAGUIThreadSnapshotStore`** - Memory-only latest AG-UI Thread Snapshot store for local development, demos, and tests
|
||||
|
||||
## Types
|
||||
|
||||
- **`AGUIRequest`** / **`AGUIChatOptions`** - Request types
|
||||
- **`AGUIThreadSnapshot`** / **`AGUIThreadSnapshotStore`** - Replayable thread snapshot model and scoped async store protocol
|
||||
- **`availableInterrupts` / `resume`** - Optional canonical AG-UI `Interrupt` and `ResumeEntry` protocol data
|
||||
- **`AgentState`** / **`RunMetadata`** - State management types
|
||||
- **`PredictStateConfig`** - Configuration for state prediction
|
||||
|
||||
## Protocol Notes
|
||||
|
||||
- Outbound custom events are emitted as AG-UI `CUSTOM`.
|
||||
- Usage metadata from `Content(type="usage")` is surfaced as `CUSTOM` events with `name="usage"`.
|
||||
- Inbound custom event aliases are accepted: `CUSTOM`, `CUSTOM_EVENT`, and `custom_event`.
|
||||
- Multimodal user inputs support both legacy (`text`, `binary`) and draft-style (`image`, `audio`, `video`, `document`) shapes.
|
||||
- Interrupted runs complete with `RUN_FINISHED.outcome.type == "interrupt"` and canonical `outcome.interrupts`; do not document or add new flows that depend on the legacy top-level `RUN_FINISHED.interrupt` field.
|
||||
- `Interrupt` and `ResumeEntry` come from the `ag-ui-protocol` package (`ag_ui.core`), not from an Agent Framework-specific interrupt model.
|
||||
- SSE keepalive is endpoint-owned transport behavior configured through
|
||||
`add_agent_framework_fastapi_endpoint(keepalive_seconds=...)`. It emits SSE comments only; do not add `PING`,
|
||||
`HEARTBEAT`, or `KEEPALIVE` AG-UI events, and do not add runner-level keepalive settings.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
from fastapi import FastAPI
|
||||
|
||||
app = FastAPI()
|
||||
add_agent_framework_fastapi_endpoint(app, agent)
|
||||
```
|
||||
|
||||
## Import Path
|
||||
|
||||
```python
|
||||
from agent_framework.ag_ui import AGUIChatClient, add_agent_framework_fastapi_endpoint
|
||||
# or directly:
|
||||
from agent_framework_ag_ui import AGUIChatClient
|
||||
```
|
||||
@@ -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,378 @@
|
||||
# Agent Framework AG-UI Integration
|
||||
|
||||
AG-UI protocol integration for Agent Framework, enabling seamless integration with AG-UI's web interface and streaming protocol.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install agent-framework-ag-ui
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Server (Host an AI Agent)
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
|
||||
# Create your agent
|
||||
agent = Agent(
|
||||
name="my_agent",
|
||||
instructions="You are a helpful assistant.",
|
||||
client=OpenAIChatCompletionClient(
|
||||
azure_endpoint="https://your-resource.openai.azure.com/",
|
||||
model="gpt-4o-mini",
|
||||
api_key="your-api-key",
|
||||
),
|
||||
)
|
||||
|
||||
# Create FastAPI app and add AG-UI endpoint
|
||||
app = FastAPI()
|
||||
add_agent_framework_fastapi_endpoint(app, agent, "/")
|
||||
|
||||
# Run with: uvicorn main:app --reload
|
||||
```
|
||||
|
||||
### Server (Host a Workflow)
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI
|
||||
from agent_framework import WorkflowBuilder, WorkflowContext, executor
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
|
||||
@executor(id="start")
|
||||
async def start(message: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.yield_output(f"Workflow received: {message}")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=start).build()
|
||||
|
||||
app = FastAPI()
|
||||
add_agent_framework_fastapi_endpoint(app, workflow, "/")
|
||||
```
|
||||
|
||||
### Server (Thread-Scoped WorkflowBuilder)
|
||||
|
||||
Use `workflow_factory` when your workflow keeps runtime state (for example pending `request_info` interrupts) and must be isolated per AG-UI thread:
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI
|
||||
from agent_framework import Workflow, WorkflowBuilder
|
||||
from agent_framework.ag_ui import AgentFrameworkWorkflow, add_agent_framework_fastapi_endpoint
|
||||
|
||||
def build_workflow_for_thread(thread_id: str) -> Workflow:
|
||||
# Build a fresh workflow instance for each thread id.
|
||||
return WorkflowBuilder(start_executor=...).build()
|
||||
|
||||
app = FastAPI()
|
||||
thread_scoped_workflow = AgentFrameworkWorkflow(
|
||||
workflow_factory=build_workflow_for_thread,
|
||||
name="my_workflow",
|
||||
)
|
||||
add_agent_framework_fastapi_endpoint(app, thread_scoped_workflow, "/")
|
||||
```
|
||||
|
||||
### Client (Connect to an AG-UI Server)
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from agent_framework.ag_ui import AGUIChatClient
|
||||
|
||||
async def main():
|
||||
async with AGUIChatClient(endpoint="http://localhost:8000/") as client:
|
||||
# Stream responses
|
||||
async for update in client.get_response("Hello!", stream=True):
|
||||
for content in update.contents:
|
||||
if content.type == "text" and content.text:
|
||||
print(content.text, end="", flush=True)
|
||||
print()
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
The `AGUIChatClient` supports:
|
||||
- Streaming and non-streaming responses
|
||||
- Hybrid tool execution (client-side + server-side tools)
|
||||
- Automatic thread management for conversation continuity
|
||||
- Integration with `Agent` for client-side history management
|
||||
- Canonical interrupt/resume passthrough (`availableInterrupts` and `resume`)
|
||||
|
||||
## Tool Return Helpers
|
||||
|
||||
Use `state_update` when a backend tool needs to send different payloads to the model, the UI, and shared state. The `text` value remains the LLM-bound tool result, `tool_result` becomes the AG-UI `ToolCallResultEvent.content` for frontend rendering, and `state` is merged into durable shared state.
|
||||
|
||||
```python
|
||||
from agent_framework import Content, tool
|
||||
from agent_framework.ag_ui import state_update
|
||||
|
||||
@tool
|
||||
async def get_weather(city: str) -> Content:
|
||||
data = await fetch_weather(city)
|
||||
return state_update(
|
||||
text=f"{city}: {data['temp']}°C and {data['conditions']}",
|
||||
tool_result={
|
||||
"component": "weather-card",
|
||||
"city": city,
|
||||
"temperature": data["temp"],
|
||||
"conditions": data["conditions"],
|
||||
"humidity": data["humidity"],
|
||||
},
|
||||
state={"weather": {"city": city, **data}},
|
||||
)
|
||||
```
|
||||
|
||||
## Documentation
|
||||
|
||||
- **[Getting Started Tutorial](getting_started/)** - Step-by-step guide to building AG-UI servers and clients
|
||||
- Server setup with FastAPI
|
||||
- Client examples using `AGUIChatClient`
|
||||
- Hybrid tool execution (client-side + server-side)
|
||||
- Thread management and conversation continuity
|
||||
- **[Examples](agent_framework_ag_ui_examples/)** - Complete examples for AG-UI features
|
||||
|
||||
## Interrupts and Resume
|
||||
|
||||
Agent Framework AG-UI uses the canonical AG-UI interrupt protocol. Paused agent approval and workflow
|
||||
`request_info` runs finish with `RUN_FINISHED.outcome.type == "interrupt"` and a non-empty
|
||||
`RUN_FINISHED.outcome.interrupts` array. Agent Framework does not define a separate interrupt model; use
|
||||
`ag_ui.core.Interrupt` and `ag_ui.core.ResumeEntry` when constructing typed request data in Python.
|
||||
|
||||
Tool approval interrupts use `reason: "tool_call"` and include `toolCallId` when the pause is bound to a tool call.
|
||||
Workflow `request_info` interrupts use `reason: "input_required"`. Framework-specific details needed for resume
|
||||
validation live in each interrupt's `metadata`, while generic clients can render the human-readable `message` and
|
||||
`responseSchema`.
|
||||
|
||||
Interrupted terminal event shape:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "RUN_FINISHED",
|
||||
"outcome": {
|
||||
"type": "interrupt",
|
||||
"interrupts": [
|
||||
{
|
||||
"id": "approval_1",
|
||||
"reason": "tool_call",
|
||||
"message": "Approve tool call get_weather?",
|
||||
"toolCallId": "tool_call_1",
|
||||
"responseSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"accepted": { "type": "boolean" },
|
||||
"arguments": { "type": "object" }
|
||||
},
|
||||
"required": ["accepted"]
|
||||
},
|
||||
"metadata": {
|
||||
"agent_framework": {
|
||||
"type": "function_approval_request",
|
||||
"function_call": {
|
||||
"call_id": "tool_call_1",
|
||||
"name": "get_weather",
|
||||
"arguments": {
|
||||
"city": "Seattle"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Resume the paused thread with a canonical `resume` array. Each entry addresses exactly one open interrupt by
|
||||
`interruptId`; `status` is `resolved` or `cancelled`; resolved entries carry the approval or workflow response payload.
|
||||
|
||||
```json
|
||||
{
|
||||
"threadId": "thread-1",
|
||||
"messages": [],
|
||||
"resume": [
|
||||
{
|
||||
"interruptId": "approval_1",
|
||||
"status": "resolved",
|
||||
"payload": {
|
||||
"approved": true
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
This is a clean release-candidate breaking change before `1.0.0`: new interrupted runs use
|
||||
`RUN_FINISHED.outcome.interrupts` and do not emit a stable top-level `RUN_FINISHED.interrupt` field. Normal
|
||||
non-interrupted runs continue to finish with valid `RUN_FINISHED` terminal events.
|
||||
|
||||
## Public API Review Notes
|
||||
|
||||
The Python package is currently in release candidate stage and is targeting the released `1.0.0` API surface. The preferred application import path is `agent_framework.ag_ui`; direct package imports from `agent_framework_ag_ui` are also supported.
|
||||
|
||||
Review focus: whether these names are the right stable contract for Python users, and whether the protocol interrupt fields below match AG-UI's expected pause/resume shape.
|
||||
|
||||
| Surface | Public exports |
|
||||
| --- | --- |
|
||||
| `agent_framework.ag_ui` facade | `AgentFrameworkAgent`, `AgentFrameworkWorkflow`, `AGUIChatClient`, `AGUIEventConverter`, `AGUIHttpService`, `AGUIThreadSnapshot`, `AGUIThreadSnapshotStore`, `InMemoryAGUIThreadSnapshotStore`, `SnapshotScopeResolver`, `add_agent_framework_fastapi_endpoint`, `state_update`, `__version__` |
|
||||
| Direct `agent_framework_ag_ui` package | Facade exports plus `AGUIChatOptions`, `AGUIRequest`, `AGUIThreadID`, `AgentState`, `DEFAULT_MAX_THREAD_SNAPSHOTS`, `DEFAULT_TAGS`, `PredictStateConfig`, `RunMetadata`, `SnapshotScope`, `WorkflowFactory` |
|
||||
| AG-UI protocol package (`ag_ui.core`) | `Interrupt`, `ResumeEntry`, `RunFinishedInterruptOutcome`, and related run outcome models |
|
||||
|
||||
Interrupt support is protocol data rather than a separate Agent Framework Python class. Requests accept canonical `availableInterrupts`/`available_interrupts` and `resume` values; `AGUIChatClient` and `AGUIHttpService.post_run(...)` forward those fields with AG-UI wire aliases; agent approval and workflow `request_info` pauses emit `RUN_FINISHED.outcome.interrupts`; `AGUIEventConverter` preserves canonical interrupt outcome metadata on the final `ChatResponseUpdate`; and thread snapshot hydration replays the canonical interrupt outcome when a scoped snapshot stores an unresolved pause.
|
||||
|
||||
## Features
|
||||
|
||||
This integration supports all 7 AG-UI features:
|
||||
|
||||
1. **Agentic Chat**: Basic streaming chat with tool calling support
|
||||
2. **Backend Tool Rendering**: Tools executed on backend with results streamed to client
|
||||
3. **Human in the Loop**: Function approval requests for user confirmation before tool execution
|
||||
4. **Agentic Generative UI**: Async tools for long-running operations with progress updates
|
||||
5. **Tool-based Generative UI**: Custom UI components rendered on frontend based on tool calls
|
||||
6. **Shared State**: Bidirectional state sync between client and server
|
||||
7. **Predictive State Updates**: Stream tool arguments as optimistic state updates during execution
|
||||
|
||||
Additional compatibility and draft support:
|
||||
- Native `Workflow` endpoint registration via `add_agent_framework_fastapi_endpoint(...)`
|
||||
- Workflow-to-AG-UI event mapping (run/step/activity/tool/custom events)
|
||||
- Custom event compatibility for inbound `CUSTOM`, `CUSTOM_EVENT`, and `custom_event`
|
||||
- Pragmatic multimodal input parsing for both legacy (`binary`) and draft media-part shapes
|
||||
- Canonical interrupt/resume handling (`availableInterrupts`, `resume`, and `RUN_FINISHED.outcome.interrupts`)
|
||||
|
||||
## Security: Authentication & Authorization
|
||||
|
||||
The AG-UI endpoint does not enforce authentication by default. **For production deployments, you should add authentication** using FastAPI's dependency injection system via the `dependencies` parameter.
|
||||
|
||||
### API Key Authentication Example
|
||||
|
||||
```python
|
||||
import os
|
||||
from fastapi import Depends, FastAPI, HTTPException, Security
|
||||
from fastapi.security import APIKeyHeader
|
||||
from agent_framework import Agent
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
|
||||
# Configure API key authentication
|
||||
API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)
|
||||
EXPECTED_API_KEY = os.environ.get("AG_UI_API_KEY")
|
||||
|
||||
async def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> None:
|
||||
"""Verify the API key provided in the request header."""
|
||||
if not api_key or api_key != EXPECTED_API_KEY:
|
||||
raise HTTPException(status_code=401, detail="Invalid or missing API key")
|
||||
|
||||
# Create agent and app
|
||||
agent = Agent(name="my_agent", instructions="...", client=...)
|
||||
app = FastAPI()
|
||||
|
||||
# Register endpoint WITH authentication
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app,
|
||||
agent,
|
||||
"/",
|
||||
dependencies=[Depends(verify_api_key)], # Authentication enforced here
|
||||
)
|
||||
```
|
||||
|
||||
### Other Authentication Options
|
||||
|
||||
The `dependencies` parameter accepts any FastAPI dependency, enabling integration with:
|
||||
|
||||
- **OAuth 2.0 / OpenID Connect** - Use `fastapi.security.OAuth2PasswordBearer`
|
||||
- **JWT Tokens** - Validate tokens with libraries like `python-jose`
|
||||
- **Azure AD / Entra ID** - Use `azure-identity` for Microsoft identity platform
|
||||
- **Rate Limiting** - Add request throttling dependencies
|
||||
- **Custom Authentication** - Implement your organization's auth requirements
|
||||
|
||||
For a complete authentication example, see [getting_started/server.py](getting_started/server.py).
|
||||
|
||||
## AG-UI Thread Snapshots
|
||||
|
||||
AG-UI Thread Snapshot persistence is opt-in and disabled by default. Existing endpoints keep their current behavior
|
||||
unless you provide a `snapshot_store`.
|
||||
|
||||
Thread snapshots let an AG-UI frontend recover replayable UI state after a refresh. When snapshot persistence is
|
||||
enabled, the endpoint stores the latest replayable snapshot for an AG-UI Thread within an application-defined
|
||||
Snapshot Scope. A Hydrate Request is an AG-UI request with a known `threadId`, `messages: []`, and no `resume`
|
||||
payload. Hydration replays the stored Shared State, message snapshot, and canonical interrupt outcome when available,
|
||||
then finishes without invoking the wrapped agent or workflow.
|
||||
|
||||
Use the built-in in-memory store for local development, demos, and tests:
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI
|
||||
|
||||
from agent_framework.ag_ui import InMemoryAGUIThreadSnapshotStore, add_agent_framework_fastapi_endpoint
|
||||
|
||||
app = FastAPI()
|
||||
agent = ...
|
||||
snapshot_store = InMemoryAGUIThreadSnapshotStore(max_snapshots=500)
|
||||
|
||||
|
||||
def resolve_snapshot_scope(request):
|
||||
# Local demo scope. Production apps should derive the scope from authenticated user or tenant context.
|
||||
del request
|
||||
return "local-demo"
|
||||
|
||||
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app,
|
||||
agent,
|
||||
"/",
|
||||
snapshot_store=snapshot_store,
|
||||
snapshot_scope_resolver=resolve_snapshot_scope,
|
||||
)
|
||||
```
|
||||
|
||||
A frontend can then hydrate the latest stored snapshot for the scoped thread:
|
||||
|
||||
```json
|
||||
{
|
||||
"threadId": "thread-1",
|
||||
"messages": []
|
||||
}
|
||||
```
|
||||
|
||||
Endpoint configuration requires `snapshot_scope_resolver` whenever a snapshot store is configured, including when
|
||||
the store is already set on a pre-wrapped `AgentFrameworkAgent` or `AgentFrameworkWorkflow`. The resolver returns
|
||||
the application-defined Snapshot Scope used with the AG-UI Thread id as the storage key.
|
||||
|
||||
AG-UI Thread ids identify AG-UI Threads; they do not authorize snapshot access. Do not treat a thread id as a bearer
|
||||
credential or tenant boundary. Production applications must authenticate and authorize every AG-UI endpoint request
|
||||
and choose a Snapshot Scope that represents the app's real access boundary, such as an authenticated user, tenant,
|
||||
or workspace. Do not rely on untrusted client-provided fields by themselves to choose that boundary.
|
||||
|
||||
Tool approval resumes are validated against server-owned Approval State. The default Approval State store is
|
||||
process-local and bounded, and stores only approval-specific state needed to validate and continue pending approvals.
|
||||
It is not an authentication, tenant authorization, or distributed durability mechanism; production applications remain
|
||||
responsible for endpoint authentication, tenant authorization, and deployment/storage architecture that matches their
|
||||
availability and worker topology requirements.
|
||||
|
||||
Stored snapshots are untrusted application data with confidentiality impact. They may contain sensitive user text,
|
||||
model output, tool results, function arguments, UI payloads, Shared State, and interrupt data. The built-in
|
||||
`InMemoryAGUIThreadSnapshotStore` is in-memory only, process-local, bounded, latest-only, and not durable production
|
||||
storage. It is cleared on process restart and is not shared across workers.
|
||||
|
||||
No file-backed AG-UI snapshot store is provided by the package. Applications that need durable persistence should
|
||||
provide an app-owned implementation of the `AGUIThreadSnapshotStore` protocol and own storage hardening, including
|
||||
encryption, access control, retention, audit, data residency, and deletion behavior.
|
||||
|
||||
## Architecture
|
||||
|
||||
The package uses a clean, orchestrator-based architecture:
|
||||
|
||||
- **AgentFrameworkAgent**: Lightweight wrapper that delegates to orchestrators
|
||||
- **Orchestrators**: Handle different execution flows (default, human-in-the-loop, etc.)
|
||||
- **Confirmation Strategies**: Domain-specific confirmation messages (extensible)
|
||||
- **AgentFrameworkEventBridge**: Converts Agent Framework events to AG-UI events
|
||||
- **Message Adapters**: Bidirectional conversion between AG-UI and Agent Framework message formats
|
||||
- **FastAPI Endpoint**: Streaming HTTP endpoint with Server-Sent Events (SSE)
|
||||
|
||||
## Next Steps
|
||||
|
||||
1. **New to AG-UI?** Start with the [Getting Started Tutorial](getting_started/)
|
||||
2. **Want to see examples?** Check out the [Examples](agent_framework_ag_ui_examples/) for AG-UI features
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AG-UI protocol integration for Agent Framework."""
|
||||
|
||||
import importlib.metadata
|
||||
|
||||
from ._agent import AgentFrameworkAgent
|
||||
from ._client import AGUIChatClient
|
||||
from ._endpoint import add_agent_framework_fastapi_endpoint
|
||||
from ._event_converters import AGUIEventConverter
|
||||
from ._http_service import AGUIHttpService
|
||||
from ._snapshots import (
|
||||
DEFAULT_MAX_THREAD_SNAPSHOTS,
|
||||
AGUIThreadID,
|
||||
AGUIThreadSnapshot,
|
||||
AGUIThreadSnapshotStore,
|
||||
InMemoryAGUIThreadSnapshotStore,
|
||||
SnapshotScope,
|
||||
SnapshotScopeResolver,
|
||||
)
|
||||
from ._state import state_update
|
||||
from ._types import AgentState, AGUIChatOptions, AGUIRequest, PredictStateConfig, RunMetadata
|
||||
from ._workflow import AgentFrameworkWorkflow, WorkflowFactory
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0"
|
||||
|
||||
# Default OpenAPI tags for AG-UI endpoints
|
||||
DEFAULT_TAGS = ["AG-UI"]
|
||||
|
||||
__all__ = [
|
||||
"AgentFrameworkAgent",
|
||||
"AgentFrameworkWorkflow",
|
||||
"WorkflowFactory",
|
||||
"add_agent_framework_fastapi_endpoint",
|
||||
"AGUIChatClient",
|
||||
"AGUIChatOptions",
|
||||
"AGUIEventConverter",
|
||||
"AGUIHttpService",
|
||||
"AGUIRequest",
|
||||
"AGUIThreadID",
|
||||
"AGUIThreadSnapshot",
|
||||
"AGUIThreadSnapshotStore",
|
||||
"AgentState",
|
||||
"InMemoryAGUIThreadSnapshotStore",
|
||||
"PredictStateConfig",
|
||||
"RunMetadata",
|
||||
"SnapshotScope",
|
||||
"SnapshotScopeResolver",
|
||||
"DEFAULT_MAX_THREAD_SNAPSHOTS",
|
||||
"DEFAULT_TAGS",
|
||||
"state_update",
|
||||
"__version__",
|
||||
]
|
||||
@@ -0,0 +1,146 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AgentFrameworkAgent wrapper for AG-UI protocol."""
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any, cast
|
||||
|
||||
from ag_ui.core import BaseEvent
|
||||
from agent_framework import SupportsAgentRun
|
||||
|
||||
from ._agent_run import PendingApprovalEntry, PendingApprovalKey, run_agent_stream
|
||||
from ._approval_state import InMemoryAGUIApprovalStateStore
|
||||
from ._snapshots import AGUIThreadSnapshotStore
|
||||
|
||||
|
||||
class AgentConfig:
|
||||
"""Configuration for agent wrapper."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
state_schema: Any | None = None,
|
||||
predict_state_config: dict[str, dict[str, str]] | None = None,
|
||||
use_service_session: bool = False,
|
||||
require_confirmation: bool = True,
|
||||
snapshot_store: AGUIThreadSnapshotStore | None = None,
|
||||
):
|
||||
"""Initialize agent configuration.
|
||||
|
||||
Args:
|
||||
state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
|
||||
predict_state_config: Configuration for predictive state updates
|
||||
use_service_session: Whether the agent session is service-managed
|
||||
require_confirmation: Whether predictive updates require user confirmation before applying
|
||||
snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless
|
||||
endpoint setup also provides an explicit Snapshot Scope resolver.
|
||||
"""
|
||||
self.state_schema = self._normalize_state_schema(state_schema)
|
||||
self.predict_state_config = predict_state_config or {}
|
||||
self.use_service_session = use_service_session
|
||||
self.require_confirmation = require_confirmation
|
||||
self.snapshot_store = snapshot_store
|
||||
|
||||
@staticmethod
|
||||
def _normalize_state_schema(state_schema: Any | None) -> dict[str, Any]:
|
||||
"""Accept dict or Pydantic model/class and return a properties dict."""
|
||||
if state_schema is None:
|
||||
return {}
|
||||
|
||||
if isinstance(state_schema, dict):
|
||||
return cast(dict[str, Any], state_schema)
|
||||
|
||||
base_model_type: type[Any] | None
|
||||
try:
|
||||
from pydantic import BaseModel as ImportedBaseModel
|
||||
|
||||
base_model_type = ImportedBaseModel
|
||||
except Exception: # pragma: no cover
|
||||
base_model_type = None
|
||||
|
||||
if base_model_type is not None and isinstance(state_schema, base_model_type):
|
||||
schema_dict = state_schema.__class__.model_json_schema()
|
||||
return schema_dict.get("properties", {}) or {}
|
||||
|
||||
if base_model_type is not None and isinstance(state_schema, type) and issubclass(state_schema, base_model_type):
|
||||
schema_dict = state_schema.model_json_schema() # type: ignore[union-attr]
|
||||
return schema_dict.get("properties", {}) or {} # type: ignore
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
class AgentFrameworkAgent:
|
||||
"""Wraps Agent Framework agents for AG-UI protocol compatibility.
|
||||
|
||||
Translates between Agent Framework's SupportsAgentRun and AG-UI's event-based
|
||||
protocol. Follows a simple linear flow: RunStarted -> content events -> RunFinished.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: SupportsAgentRun,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
state_schema: Any | None = None,
|
||||
predict_state_config: dict[str, dict[str, str]] | None = None,
|
||||
require_confirmation: bool = True,
|
||||
use_service_session: bool = False,
|
||||
snapshot_store: AGUIThreadSnapshotStore | None = None,
|
||||
):
|
||||
"""Initialize the AG-UI compatible agent wrapper.
|
||||
|
||||
Args:
|
||||
agent: The Agent Framework agent to wrap
|
||||
name: Optional name for the agent
|
||||
description: Optional description
|
||||
state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
|
||||
predict_state_config: Configuration for predictive state updates
|
||||
require_confirmation: Whether predictive updates require user confirmation before applying
|
||||
use_service_session: Whether the agent session is service-managed
|
||||
snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless
|
||||
endpoint setup also provides an explicit Snapshot Scope resolver.
|
||||
"""
|
||||
self.agent = agent
|
||||
self.name = name or getattr(agent, "name", "agent")
|
||||
self.description = description or getattr(agent, "description", "")
|
||||
|
||||
self.config = AgentConfig(
|
||||
state_schema=state_schema,
|
||||
predict_state_config=predict_state_config,
|
||||
use_service_session=use_service_session,
|
||||
require_confirmation=require_confirmation,
|
||||
snapshot_store=snapshot_store,
|
||||
)
|
||||
|
||||
# Server-side Approval State. Populated when approval requests are emitted
|
||||
# and consumed when resume decisions arrive.
|
||||
self._approval_state_store = InMemoryAGUIApprovalStateStore()
|
||||
self._pending_approvals = cast(
|
||||
dict[PendingApprovalKey, PendingApprovalEntry],
|
||||
self._approval_state_store.pending_approvals,
|
||||
)
|
||||
|
||||
@property
|
||||
def snapshot_store(self) -> AGUIThreadSnapshotStore | None:
|
||||
"""Configured AG-UI Thread Snapshot store, if any."""
|
||||
return self.config.snapshot_store
|
||||
|
||||
async def run(
|
||||
self,
|
||||
input_data: dict[str, Any],
|
||||
) -> AsyncGenerator[BaseEvent, None]:
|
||||
"""Run the wrapped agent and yield AG-UI events.
|
||||
|
||||
Args:
|
||||
input_data: The AG-UI run input containing messages, state, etc.
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
async for event in run_agent_stream(
|
||||
input_data,
|
||||
self.agent,
|
||||
self.config,
|
||||
pending_approvals=self._pending_approvals,
|
||||
approval_state_store=self._approval_state_store,
|
||||
):
|
||||
yield event
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Server-side AG-UI approval state storage."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import OrderedDict
|
||||
from typing import Any
|
||||
|
||||
ApprovalScope = str
|
||||
"""Application-defined scope for server-side AG-UI Approval State."""
|
||||
|
||||
DEFAULT_MAX_APPROVAL_STATES = 10_000
|
||||
_APPROVAL_SCOPE_INPUT_KEY = "__ag_ui_approval_scope"
|
||||
_APPROVAL_THREAD_SEPARATOR = "\x1f"
|
||||
|
||||
|
||||
def approval_state_thread_id(*, scope: object | None, thread_id: str) -> str:
|
||||
"""Return the storage thread key for Approval State.
|
||||
|
||||
``None`` is the only unscoped value. A provided scope must be a non-empty
|
||||
string so accidental empty or malformed scopes cannot collapse into the
|
||||
unscoped namespace.
|
||||
"""
|
||||
if scope is None:
|
||||
return thread_id
|
||||
if not isinstance(scope, str) or not scope:
|
||||
raise ValueError("scope must be a non-empty string when provided.")
|
||||
return f"{scope}{_APPROVAL_THREAD_SEPARATOR}{thread_id}"
|
||||
|
||||
|
||||
class InMemoryAGUIApprovalStateStore:
|
||||
"""Bounded process-local server-side store for AG-UI Approval State.
|
||||
|
||||
The default store keeps only pending approval entries. It does not store
|
||||
general ``AgentSession.state`` or AG-UI Thread Snapshots.
|
||||
"""
|
||||
|
||||
def __init__(self, *, max_entries: int = DEFAULT_MAX_APPROVAL_STATES) -> None:
|
||||
"""Initialize the process-local Approval State store.
|
||||
|
||||
Keyword Args:
|
||||
max_entries: Maximum pending approval entries to retain.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``max_entries`` is less than 1.
|
||||
"""
|
||||
if max_entries < 1:
|
||||
raise ValueError("max_entries must be greater than 0.")
|
||||
self.max_entries = max_entries
|
||||
self.pending_approvals: OrderedDict[tuple[str, str], Any] = OrderedDict()
|
||||
self.tool_approval_states: OrderedDict[str, dict[str, Any]] = OrderedDict()
|
||||
|
||||
def evict_oldest(self) -> None:
|
||||
"""Evict oldest pending approval entries until the store is within bounds."""
|
||||
while len(self.pending_approvals) > self.max_entries:
|
||||
self.pending_approvals.popitem(last=False)
|
||||
while len(self.tool_approval_states) > self.max_entries:
|
||||
self.tool_approval_states.popitem(last=False)
|
||||
@@ -0,0 +1,468 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AG-UI Chat Client implementation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import uuid
|
||||
from collections.abc import AsyncIterable, Awaitable, Mapping, MutableSequence, Sequence
|
||||
from functools import wraps
|
||||
from typing import TYPE_CHECKING, Any, Generic, TypedDict, cast
|
||||
|
||||
import httpx
|
||||
from agent_framework import (
|
||||
BaseChatClient,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Content,
|
||||
FunctionTool,
|
||||
Message,
|
||||
ResponseStream,
|
||||
)
|
||||
from agent_framework._middleware import ChatMiddlewareLayer
|
||||
from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
|
||||
from agent_framework.observability import ChatTelemetryLayer
|
||||
|
||||
from ._event_converters import AGUIEventConverter
|
||||
from ._http_service import AGUIHttpService, _serialize_available_interrupts, _serialize_resume
|
||||
from ._message_adapters import agent_framework_messages_to_agui
|
||||
from ._utils import convert_tools_to_agui_format
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # pragma: no cover
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # pragma: no cover
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self, TypedDict # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import Self, TypedDict # pragma: no cover
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agent_framework._middleware import ChatAndFunctionMiddlewareTypes
|
||||
|
||||
from ._types import AGUIChatOptions
|
||||
|
||||
logger: logging.Logger = logging.getLogger("agent_framework.ag_ui")
|
||||
|
||||
|
||||
def _unwrap_server_function_call_contents(contents: MutableSequence[Content | dict[str, Any]]) -> None:
|
||||
"""Replace server_function_call instances with their underlying call content."""
|
||||
for idx, content in enumerate(contents):
|
||||
if content.type == "server_function_call": # type: ignore[union-attr]
|
||||
contents[idx] = content.function_call # type: ignore[assignment, union-attr]
|
||||
|
||||
|
||||
BaseChatClientT = TypeVar("BaseChatClientT", bound=type[BaseChatClient[Any]])
|
||||
|
||||
AGUIChatOptionsT = TypeVar(
|
||||
"AGUIChatOptionsT",
|
||||
bound=TypedDict, # type: ignore[valid-type]
|
||||
default="AGUIChatOptions",
|
||||
covariant=True,
|
||||
)
|
||||
|
||||
|
||||
def _apply_server_function_call_unwrap(client: BaseChatClientT) -> BaseChatClientT:
|
||||
"""Class decorator that unwraps server-side function calls after tool handling."""
|
||||
|
||||
original_get_response = client.get_response
|
||||
|
||||
@wraps(original_get_response)
|
||||
def response_wrapper(
|
||||
self, *args: Any, stream: bool = False, **kwargs: Any
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
if stream:
|
||||
stream_response = original_get_response(self, *args, stream=True, **kwargs)
|
||||
if isinstance(stream_response, ResponseStream):
|
||||
return stream_response.with_transform_hook(_map_update)
|
||||
return ResponseStream(_stream_wrapper_impl(stream_response))
|
||||
return _response_wrapper_impl(self, original_get_response, *args, **kwargs)
|
||||
|
||||
async def _response_wrapper_impl(self, original_func: Any, *args: Any, **kwargs: Any) -> ChatResponse:
|
||||
"""Non-streaming wrapper implementation."""
|
||||
response = await original_func(self, *args, stream=False, **kwargs)
|
||||
if response.messages:
|
||||
for message in response.messages:
|
||||
_unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], message.contents))
|
||||
return response
|
||||
|
||||
async def _stream_wrapper_impl(stream: Any) -> AsyncIterable[ChatResponseUpdate]:
|
||||
"""Streaming wrapper implementation."""
|
||||
if isinstance(stream, Awaitable):
|
||||
stream = await stream
|
||||
async for update in stream:
|
||||
_unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], update.contents))
|
||||
yield update
|
||||
|
||||
def _map_update(update: ChatResponseUpdate) -> ChatResponseUpdate:
|
||||
_unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], update.contents))
|
||||
return update
|
||||
|
||||
client.get_response = response_wrapper # type: ignore[assignment]
|
||||
return client
|
||||
|
||||
|
||||
@_apply_server_function_call_unwrap
|
||||
class AGUIChatClient(
|
||||
FunctionInvocationLayer[AGUIChatOptionsT],
|
||||
ChatMiddlewareLayer[AGUIChatOptionsT],
|
||||
ChatTelemetryLayer[AGUIChatOptionsT],
|
||||
BaseChatClient[AGUIChatOptionsT],
|
||||
Generic[AGUIChatOptionsT],
|
||||
):
|
||||
"""Chat client for communicating with AG-UI compliant servers.
|
||||
|
||||
This client implements the BaseChatClient interface and automatically handles:
|
||||
- Thread ID management for conversation continuity
|
||||
- State synchronization between client and server
|
||||
- Server-Sent Events (SSE) streaming
|
||||
- Event conversion to Agent Framework types
|
||||
- MiddlewareTypes, telemetry, and function invocation support
|
||||
|
||||
Important: Message History Management
|
||||
This client sends exactly the messages it receives to the server. It does NOT
|
||||
automatically maintain conversation history. The server must handle history via thread_id.
|
||||
|
||||
For stateless servers: Use Agent wrapper which will send full message history on each
|
||||
request. However, even with Agent, the server must echo back all context for the
|
||||
agent to maintain history across turns.
|
||||
|
||||
Important: Tool Handling (Hybrid Execution - matches .NET)
|
||||
1. Client tool metadata sent to server - LLM knows about both client and server tools
|
||||
2. Server has its own tools that execute server-side
|
||||
3. When LLM calls a client tool, function invocation executes it locally
|
||||
4. Both client and server tools work together (hybrid pattern)
|
||||
|
||||
The wrapping Agent's function invocation handles client tool execution
|
||||
automatically when the server's LLM decides to call them.
|
||||
|
||||
Examples:
|
||||
Direct usage (server manages thread history):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.ag_ui import AGUIChatClient
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
|
||||
# First message - thread ID auto-generated
|
||||
response = await client.get_response("Hello!")
|
||||
thread_id = response.additional_properties.get("thread_id")
|
||||
|
||||
# Second message - server retrieves history using thread_id
|
||||
response2 = await client.get_response(
|
||||
"How are you?",
|
||||
metadata={"thread_id": thread_id}
|
||||
)
|
||||
|
||||
Recommended usage with Agent (client manages history):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Agent
|
||||
from agent_framework.ag_ui import AGUIChatClient
|
||||
|
||||
client = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
agent = Agent(name="assistant", client=client)
|
||||
session = agent.create_session()
|
||||
|
||||
# Agent automatically maintains history and sends full context
|
||||
response = await agent.run("Hello!", session=session)
|
||||
response2 = await agent.run("How are you?", session=session)
|
||||
|
||||
Streaming usage:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
async for update in client.get_response("Tell me a story", stream=True):
|
||||
if update.contents:
|
||||
for content in update.contents:
|
||||
if hasattr(content, "text"):
|
||||
print(content.text, end="", flush=True)
|
||||
|
||||
Context manager:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
async with AGUIChatClient(endpoint="http://localhost:8888/") as client:
|
||||
response = await client.get_response("Hello!")
|
||||
print(response.messages[0].text)
|
||||
|
||||
Using custom ChatOptions with type safety:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from typing import TypedDict
|
||||
from agent_framework_ag_ui import AGUIChatClient, AGUIChatOptions
|
||||
|
||||
class MyOptions(AGUIChatOptions, total=False):
|
||||
my_custom_option: str
|
||||
|
||||
client: AGUIChatClient[MyOptions] = AGUIChatClient(endpoint="http://localhost:8888/")
|
||||
response = await client.get_response("Hello", options={"my_custom_option": "value"})
|
||||
"""
|
||||
|
||||
OTEL_PROVIDER_NAME = "agui"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
endpoint: str,
|
||||
http_client: httpx.AsyncClient | None = None,
|
||||
timeout: float = 60.0,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
|
||||
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
|
||||
) -> None:
|
||||
"""Initialize the AG-UI chat client.
|
||||
|
||||
Args:
|
||||
endpoint: The AG-UI server endpoint URL (e.g., "http://localhost:8888/")
|
||||
http_client: Optional httpx.AsyncClient instance. If None, one will be created.
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
additional_properties: Additional properties to store
|
||||
middleware: Optional middleware to apply to the client.
|
||||
function_invocation_configuration: Optional function invocation configuration override.
|
||||
"""
|
||||
super().__init__(
|
||||
additional_properties=additional_properties,
|
||||
middleware=middleware,
|
||||
function_invocation_configuration=function_invocation_configuration,
|
||||
)
|
||||
self._http_service = AGUIHttpService(
|
||||
endpoint=endpoint,
|
||||
http_client=http_client,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the HTTP client."""
|
||||
await self._http_service.close()
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
"""Enter async context manager."""
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
"""Exit async context manager."""
|
||||
await self.close()
|
||||
|
||||
def _register_server_tool_placeholder(self, tool_name: str) -> None:
|
||||
"""Register a declaration-only placeholder so function invocation skips execution."""
|
||||
|
||||
config = getattr(self, "function_invocation_configuration", None)
|
||||
if not isinstance(config, dict):
|
||||
return
|
||||
additional_tools = list(config.get("additional_tools", []))
|
||||
if any(getattr(tool, "name", None) == tool_name for tool in additional_tools):
|
||||
return
|
||||
|
||||
placeholder: FunctionTool = FunctionTool(
|
||||
name=tool_name,
|
||||
description="Server-managed tool placeholder (AG-UI)",
|
||||
func=None,
|
||||
)
|
||||
additional_tools.append(placeholder)
|
||||
config["additional_tools"] = additional_tools
|
||||
registered: set[str] = getattr(self, "_registered_server_tools", set())
|
||||
registered.add(tool_name)
|
||||
self._registered_server_tools = registered
|
||||
logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}")
|
||||
|
||||
def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[list[Message], dict[str, Any] | None]:
|
||||
"""Extract state from last message if present.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages
|
||||
|
||||
Returns:
|
||||
Tuple of (messages_without_state, state_dict)
|
||||
"""
|
||||
if not messages:
|
||||
return list(messages), None
|
||||
|
||||
last_message = messages[-1]
|
||||
|
||||
for content in last_message.contents:
|
||||
if isinstance(content, Content) and content.type == "data" and content.media_type == "application/json":
|
||||
try:
|
||||
uri = content.uri
|
||||
if uri.startswith("data:application/json;base64,"): # type: ignore[union-attr]
|
||||
import base64
|
||||
|
||||
encoded_data = uri.split(",", 1)[1] # type: ignore[union-attr]
|
||||
decoded_bytes = base64.b64decode(encoded_data)
|
||||
state = json.loads(decoded_bytes.decode("utf-8"))
|
||||
|
||||
messages_without_state = list(messages[:-1]) if len(messages) > 1 else []
|
||||
return messages_without_state, state
|
||||
except (json.JSONDecodeError, ValueError, KeyError) as e:
|
||||
logger.warning(f"Failed to extract state from message: {e}")
|
||||
|
||||
return list(messages), None
|
||||
|
||||
def _convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
|
||||
"""Convert Agent Framework messages to AG-UI format.
|
||||
|
||||
Args:
|
||||
messages: List of Message objects
|
||||
|
||||
Returns:
|
||||
List of AG-UI formatted message dictionaries
|
||||
"""
|
||||
return agent_framework_messages_to_agui(messages)
|
||||
|
||||
def _get_thread_id(self, options: Mapping[str, Any]) -> str:
|
||||
"""Get or generate thread ID from chat options.
|
||||
|
||||
Args:
|
||||
options: Chat options containing metadata
|
||||
|
||||
Returns:
|
||||
Thread ID string
|
||||
"""
|
||||
thread_id = None
|
||||
metadata = options.get("metadata")
|
||||
if metadata:
|
||||
thread_id = metadata.get("thread_id")
|
||||
|
||||
if not thread_id:
|
||||
thread_id = f"thread_{uuid.uuid4().hex}"
|
||||
|
||||
return thread_id
|
||||
|
||||
@override
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
stream: bool = False,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
"""Internal method to get non-streaming response.
|
||||
|
||||
Keyword Args:
|
||||
messages: List of chat messages
|
||||
stream: Whether to stream the response.
|
||||
options: Chat options for the request
|
||||
**kwargs: Additional keyword arguments
|
||||
|
||||
Returns:
|
||||
ChatResponse object
|
||||
"""
|
||||
if stream:
|
||||
return ResponseStream(
|
||||
self._streaming_impl(
|
||||
messages=messages,
|
||||
options=options,
|
||||
**kwargs,
|
||||
),
|
||||
finalizer=ChatResponse.from_updates,
|
||||
)
|
||||
|
||||
async def _get_response() -> ChatResponse:
|
||||
return await ChatResponse.from_update_generator(
|
||||
self._streaming_impl(
|
||||
messages=messages,
|
||||
options=options,
|
||||
**kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
return _get_response()
|
||||
|
||||
async def _streaming_impl(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
"""Internal method to get streaming response.
|
||||
|
||||
Keyword Args:
|
||||
messages: Sequence of chat messages
|
||||
options: Chat options for the request
|
||||
**kwargs: Additional keyword arguments
|
||||
|
||||
Yields:
|
||||
ChatResponseUpdate objects
|
||||
"""
|
||||
messages_to_send, state = self._extract_state_from_messages(messages)
|
||||
|
||||
thread_id = self._get_thread_id(options)
|
||||
run_id = f"run_{uuid.uuid4().hex}"
|
||||
|
||||
agui_messages = self._convert_messages_to_agui_format(messages_to_send)
|
||||
|
||||
# Send client tools to server so LLM knows about them
|
||||
# Client tools execute via Agent's function invocation wrapper
|
||||
agui_tools = convert_tools_to_agui_format(options.get("tools"))
|
||||
|
||||
# Build set of client tool names (matches .NET clientToolSet)
|
||||
# Used to distinguish client vs server tools in response stream
|
||||
client_tool_set: set[str] = set()
|
||||
tools = options.get("tools")
|
||||
if tools:
|
||||
for tool in tools:
|
||||
if hasattr(tool, "name"):
|
||||
client_tool_set.add(tool.name)
|
||||
self._last_client_tool_set = client_tool_set
|
||||
|
||||
logger.debug(
|
||||
"[AGUIChatClient] Preparing request",
|
||||
extra={
|
||||
"thread_id": thread_id,
|
||||
"run_id": run_id,
|
||||
"client_tools": list(client_tool_set),
|
||||
"messages": [msg.text for msg in messages_to_send if msg.text],
|
||||
},
|
||||
)
|
||||
logger.debug(f"[AGUIChatClient] Client tool set: {client_tool_set}")
|
||||
|
||||
converter = AGUIEventConverter()
|
||||
|
||||
available_interrupts = options.get("available_interrupts", options.get("availableInterrupts"))
|
||||
|
||||
async for event in self._http_service.post_run(
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
messages=agui_messages,
|
||||
state=state,
|
||||
tools=agui_tools,
|
||||
available_interrupts=_serialize_available_interrupts(cast(Sequence[Any] | None, available_interrupts)),
|
||||
resume=_serialize_resume(options.get("resume")),
|
||||
):
|
||||
logger.debug(f"[AGUIChatClient] Raw AG-UI event: {event}")
|
||||
update = converter.convert_event(event)
|
||||
if update is not None:
|
||||
logger.debug(
|
||||
"[AGUIChatClient] Converted update",
|
||||
extra={"role": update.role, "contents": [type(c).__name__ for c in update.contents]},
|
||||
)
|
||||
# Distinguish client vs server tools
|
||||
for i, content in enumerate(update.contents):
|
||||
if content.type == "function_call":
|
||||
logger.debug(
|
||||
f"[AGUIChatClient] Function call: {content.name}, in client_tool_set: {content.name in client_tool_set}"
|
||||
)
|
||||
if content.name in client_tool_set:
|
||||
# Client tool - let function invocation execute it
|
||||
if not content.additional_properties:
|
||||
content.additional_properties = {}
|
||||
content.additional_properties["agui_thread_id"] = thread_id
|
||||
else:
|
||||
# Server tool - wrap so function invocation ignores it
|
||||
logger.debug(f"[AGUIChatClient] Wrapping server tool: {content.name}")
|
||||
self._register_server_tool_placeholder(content.name) # type: ignore[arg-type]
|
||||
update.contents[i] = Content(type="server_function_call", function_call=content) # type: ignore
|
||||
|
||||
yield update
|
||||
@@ -0,0 +1,257 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""FastAPI endpoint creation for AG-UI agents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from collections.abc import AsyncGenerator, Sequence
|
||||
from inspect import isawaitable
|
||||
from typing import Any, cast
|
||||
|
||||
from ag_ui.core import RunErrorEvent
|
||||
from ag_ui.encoder import EventEncoder
|
||||
from agent_framework import SupportsAgentRun, Workflow
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.params import Depends
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from ._agent import AgentFrameworkAgent
|
||||
from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY
|
||||
from ._snapshots import (
|
||||
_DEFAULT_STATE_INPUT_KEY,
|
||||
_SNAPSHOT_SCOPE_INPUT_KEY,
|
||||
AGUIThreadSnapshotStore,
|
||||
SnapshotScopeResolver,
|
||||
)
|
||||
from ._types import AGUIRequest
|
||||
from ._workflow import AgentFrameworkWorkflow
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_KEEPALIVE_COMMENT = "keepalive"
|
||||
|
||||
|
||||
def _get_snapshot_store(
|
||||
protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow,
|
||||
) -> AGUIThreadSnapshotStore | None:
|
||||
if isinstance(protocol_runner, AgentFrameworkAgent):
|
||||
return protocol_runner.config.snapshot_store
|
||||
return protocol_runner.snapshot_store
|
||||
|
||||
|
||||
def _set_snapshot_store(
|
||||
protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow,
|
||||
snapshot_store: AGUIThreadSnapshotStore,
|
||||
) -> None:
|
||||
if isinstance(protocol_runner, AgentFrameworkAgent):
|
||||
protocol_runner.config.snapshot_store = snapshot_store
|
||||
return
|
||||
protocol_runner.snapshot_store = snapshot_store
|
||||
|
||||
|
||||
def _configure_snapshot_persistence(
|
||||
protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow,
|
||||
*,
|
||||
snapshot_store: AGUIThreadSnapshotStore | None,
|
||||
snapshot_scope_resolver: SnapshotScopeResolver | None,
|
||||
) -> None:
|
||||
existing_snapshot_store = _get_snapshot_store(protocol_runner)
|
||||
if snapshot_store is not None:
|
||||
if existing_snapshot_store is not None and existing_snapshot_store is not snapshot_store:
|
||||
raise ValueError("snapshot_store is already configured on the AG-UI runner.")
|
||||
if existing_snapshot_store is None:
|
||||
_set_snapshot_store(protocol_runner, snapshot_store)
|
||||
existing_snapshot_store = snapshot_store
|
||||
|
||||
if existing_snapshot_store is not None and snapshot_scope_resolver is None:
|
||||
raise ValueError(
|
||||
"snapshot_scope_resolver is required when snapshot_store is configured. "
|
||||
"AG-UI Thread ids identify threads but do not authorize snapshot access; "
|
||||
"provide a resolver that returns an explicit Snapshot Scope."
|
||||
)
|
||||
|
||||
|
||||
def _validate_keepalive_seconds(keepalive_seconds: float | None) -> None:
|
||||
if keepalive_seconds is not None and not keepalive_seconds > 0:
|
||||
raise ValueError("keepalive_seconds must be positive or None.")
|
||||
|
||||
|
||||
def add_agent_framework_fastapi_endpoint(
|
||||
app: FastAPI,
|
||||
agent: SupportsAgentRun | AgentFrameworkAgent | Workflow | AgentFrameworkWorkflow,
|
||||
path: str = "/",
|
||||
state_schema: Any | None = None,
|
||||
predict_state_config: dict[str, dict[str, str]] | None = None,
|
||||
allow_origins: list[str] | None = None,
|
||||
default_state: dict[str, Any] | None = None,
|
||||
tags: list[str] | None = None,
|
||||
dependencies: Sequence[Depends] | None = None,
|
||||
snapshot_store: AGUIThreadSnapshotStore | None = None,
|
||||
snapshot_scope_resolver: SnapshotScopeResolver | None = None,
|
||||
keepalive_seconds: float | None = 15,
|
||||
) -> None:
|
||||
"""Add an AG-UI endpoint to a FastAPI app.
|
||||
|
||||
Args:
|
||||
app: The FastAPI application
|
||||
agent: The agent to expose (can be raw SupportsAgentRun or wrapped)
|
||||
path: The endpoint path
|
||||
state_schema: Optional state schema for shared state management; accepts dict or Pydantic model/class
|
||||
predict_state_config: Optional predictive state update configuration.
|
||||
Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
|
||||
allow_origins: CORS origins (not yet implemented)
|
||||
default_state: Optional initial state to seed when the client does not provide state keys
|
||||
tags: OpenAPI tags for endpoint categorization (defaults to ["AG-UI"])
|
||||
dependencies: Optional FastAPI dependencies for authentication/authorization.
|
||||
These dependencies run before the endpoint handler. Use this to add
|
||||
authentication checks, rate limiting, or other middleware-like behavior.
|
||||
Example: `dependencies=[Depends(verify_api_key)]`
|
||||
snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence is opt-in and requires an
|
||||
explicit Snapshot Scope resolver.
|
||||
snapshot_scope_resolver: Optional resolver for the application-defined Snapshot Scope. Required whenever
|
||||
a snapshot store is configured because an AG-UI Thread id is not an authorization boundary.
|
||||
keepalive_seconds: Endpoint SSE keepalive interval in seconds. Defaults to 15. Positive values emit fixed
|
||||
SSE comments while the stream is open. None disables keepalive and preserves the non-keepalive response
|
||||
path. Keepalive comments are transport traffic and do not change AG-UI events.
|
||||
"""
|
||||
_validate_keepalive_seconds(keepalive_seconds)
|
||||
|
||||
protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow
|
||||
if isinstance(agent, AgentFrameworkWorkflow):
|
||||
protocol_runner = agent
|
||||
elif isinstance(agent, AgentFrameworkAgent):
|
||||
protocol_runner = agent
|
||||
elif isinstance(agent, Workflow):
|
||||
protocol_runner = AgentFrameworkWorkflow(workflow=agent)
|
||||
elif isinstance(agent, SupportsAgentRun):
|
||||
protocol_runner = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema=state_schema,
|
||||
predict_state_config=predict_state_config,
|
||||
snapshot_store=snapshot_store,
|
||||
)
|
||||
else:
|
||||
raise TypeError("agent must be SupportsAgentRun, Workflow, AgentFrameworkAgent, or AgentFrameworkWorkflow.")
|
||||
|
||||
_configure_snapshot_persistence(
|
||||
protocol_runner,
|
||||
snapshot_store=snapshot_store,
|
||||
snapshot_scope_resolver=snapshot_scope_resolver,
|
||||
)
|
||||
|
||||
@app.post(path, tags=tags or ["AG-UI"], dependencies=dependencies, response_model=None) # type: ignore[arg-type]
|
||||
async def agent_endpoint(request_body: AGUIRequest) -> Response:
|
||||
"""Handle AG-UI agent requests.
|
||||
|
||||
Note: Function is accessed via FastAPI's decorator registration,
|
||||
despite appearing unused to static analysis.
|
||||
"""
|
||||
try:
|
||||
input_data = request_body.model_dump(exclude_none=True)
|
||||
snapshot_persistence_active = False
|
||||
if snapshot_scope_resolver is not None:
|
||||
snapshot_scope = snapshot_scope_resolver(request_body)
|
||||
if isawaitable(snapshot_scope):
|
||||
snapshot_scope = await snapshot_scope
|
||||
input_data[_APPROVAL_SCOPE_INPUT_KEY] = snapshot_scope
|
||||
if _get_snapshot_store(protocol_runner) is not None:
|
||||
input_data[_SNAPSHOT_SCOPE_INPUT_KEY] = snapshot_scope
|
||||
snapshot_persistence_active = True
|
||||
if default_state:
|
||||
if snapshot_persistence_active:
|
||||
# Defer default application to the runner so defaults only fill keys
|
||||
# missing from both the stored snapshot state and the request state.
|
||||
input_data[_DEFAULT_STATE_INPUT_KEY] = copy.deepcopy(default_state)
|
||||
else:
|
||||
state = input_data.setdefault("state", {})
|
||||
for key, value in default_state.items():
|
||||
if key not in state:
|
||||
state[key] = copy.deepcopy(value)
|
||||
logger.debug(
|
||||
f"[{path}] Received request - Run ID: {input_data.get('run_id', 'no-run-id')}, "
|
||||
f"Thread ID: {input_data.get('thread_id', 'no-thread-id')}, "
|
||||
f"Messages: {len(input_data.get('messages', []))}"
|
||||
)
|
||||
logger.info(f"Received request at {path}: {input_data.get('run_id', 'no-run-id')}")
|
||||
|
||||
keepalive_enabled = keepalive_seconds is not None
|
||||
|
||||
def prepare_frame(encoded: str) -> str | bytes:
|
||||
if keepalive_enabled:
|
||||
return encoded.encode("utf-8")
|
||||
return encoded
|
||||
|
||||
async def event_generator() -> AsyncGenerator[str | bytes]:
|
||||
encoder = EventEncoder()
|
||||
event_count = 0
|
||||
try:
|
||||
async for event in protocol_runner.run(input_data):
|
||||
event_count += 1
|
||||
event_type_name = getattr(event, "type", type(event).__name__)
|
||||
# Log important events at INFO level
|
||||
if "TOOL_CALL" in str(event_type_name) or "RUN" in str(event_type_name):
|
||||
if hasattr(event, "model_dump"):
|
||||
event_data = event.model_dump(exclude_none=True)
|
||||
logger.info(f"[{path}] Event {event_count}: {event_type_name} - {event_data}")
|
||||
else:
|
||||
logger.info(f"[{path}] Event {event_count}: {event_type_name}")
|
||||
|
||||
try:
|
||||
encoded = encoder.encode(event)
|
||||
except Exception as encode_error:
|
||||
logger.exception("[%s] Failed to encode event %s", path, event_type_name)
|
||||
run_error = RunErrorEvent(
|
||||
message="An internal error has occurred while streaming events.",
|
||||
code=type(encode_error).__name__,
|
||||
)
|
||||
try:
|
||||
yield prepare_frame(encoder.encode(run_error))
|
||||
except Exception:
|
||||
logger.exception("[%s] Failed to encode RUN_ERROR event", path)
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
f"[{path}] Encoded as: {encoded[:200]}..."
|
||||
if len(encoded) > 200
|
||||
else f"[{path}] Encoded as: {encoded}"
|
||||
)
|
||||
yield prepare_frame(encoded)
|
||||
|
||||
logger.info(f"[{path}] Completed streaming {event_count} events")
|
||||
except Exception as stream_error:
|
||||
logger.exception("[%s] Streaming failed", path)
|
||||
run_error = RunErrorEvent(
|
||||
message="An internal error has occurred while streaming events.",
|
||||
code=type(stream_error).__name__,
|
||||
)
|
||||
try:
|
||||
yield prepare_frame(encoder.encode(run_error))
|
||||
except Exception:
|
||||
logger.exception("[%s] Failed to encode RUN_ERROR event", path)
|
||||
|
||||
headers = {
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
}
|
||||
if keepalive_seconds is not None:
|
||||
from sse_starlette.event import ServerSentEvent
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
|
||||
return EventSourceResponse(
|
||||
event_generator(),
|
||||
ping=cast(int, keepalive_seconds),
|
||||
ping_message_factory=lambda: ServerSentEvent(comment=_KEEPALIVE_COMMENT),
|
||||
headers=headers,
|
||||
media_type="text/event-stream",
|
||||
)
|
||||
return StreamingResponse(
|
||||
event_generator(),
|
||||
media_type="text/event-stream",
|
||||
headers=headers,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in agent endpoint: {e}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail="An internal error has occurred.") from e
|
||||
@@ -0,0 +1,279 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Event converter for AG-UI protocol events to Agent Framework types."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import (
|
||||
ChatResponseUpdate,
|
||||
Content,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AGUIEventConverter:
|
||||
"""Converter for AG-UI events to Agent Framework types.
|
||||
|
||||
Handles conversion of AG-UI protocol events to ChatResponseUpdate objects
|
||||
while maintaining state, aggregating content, and tracking metadata.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize the converter with fresh state."""
|
||||
self.current_message_id: str | None = None
|
||||
self.current_tool_call_id: str | None = None
|
||||
self.current_tool_name: str | None = None
|
||||
self.accumulated_tool_args: str = ""
|
||||
self.thread_id: str | None = None
|
||||
self.run_id: str | None = None
|
||||
|
||||
@staticmethod
|
||||
def _get_tool_call_id(event: dict[str, Any]) -> str | None:
|
||||
"""Return the tool call ID from either AG-UI field spelling."""
|
||||
tool_call_id = event.get("toolCallId")
|
||||
if tool_call_id is None:
|
||||
tool_call_id = event.get("tool_call_id")
|
||||
if tool_call_id is None:
|
||||
return None
|
||||
return str(tool_call_id)
|
||||
|
||||
def convert_event(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
|
||||
"""Convert a single AG-UI event to ChatResponseUpdate.
|
||||
|
||||
Args:
|
||||
event: AG-UI event dictionary
|
||||
|
||||
Returns:
|
||||
ChatResponseUpdate if event produces content, None otherwise
|
||||
|
||||
Examples:
|
||||
RUN_STARTED event:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
converter = AGUIEventConverter()
|
||||
event = {"type": "RUN_STARTED", "threadId": "t1", "runId": "r1"}
|
||||
update = converter.convert_event(event)
|
||||
assert update.additional_properties["thread_id"] == "t1"
|
||||
|
||||
TEXT_MESSAGE_CONTENT event:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
event = {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "Hello"}
|
||||
update = converter.convert_event(event)
|
||||
assert update.contents[0].text == "Hello"
|
||||
"""
|
||||
raw_event_type = str(event.get("type", ""))
|
||||
event_type = raw_event_type.upper()
|
||||
|
||||
if event_type == "RUN_STARTED":
|
||||
return self._handle_run_started(event)
|
||||
elif event_type == "TEXT_MESSAGE_START":
|
||||
return self._handle_text_message_start(event)
|
||||
elif event_type == "TEXT_MESSAGE_CONTENT":
|
||||
return self._handle_text_message_content(event)
|
||||
elif event_type == "TEXT_MESSAGE_END":
|
||||
return self._handle_text_message_end(event)
|
||||
elif event_type == "TOOL_CALL_START":
|
||||
return self._handle_tool_call_start(event)
|
||||
elif event_type == "TOOL_CALL_ARGS":
|
||||
return self._handle_tool_call_args(event)
|
||||
elif event_type == "TOOL_CALL_END":
|
||||
return self._handle_tool_call_end(event)
|
||||
elif event_type == "TOOL_CALL_RESULT":
|
||||
return self._handle_tool_call_result(event)
|
||||
elif event_type == "RUN_FINISHED":
|
||||
return self._handle_run_finished(event)
|
||||
elif event_type == "RUN_ERROR":
|
||||
return self._handle_run_error(event)
|
||||
elif event_type in {"CUSTOM", "CUSTOM_EVENT"}:
|
||||
return self._handle_custom_event(event, raw_event_type)
|
||||
|
||||
return None
|
||||
|
||||
def _handle_run_started(self, event: dict[str, Any]) -> ChatResponseUpdate:
|
||||
"""Handle RUN_STARTED event."""
|
||||
self.thread_id = event.get("threadId")
|
||||
self.run_id = event.get("runId")
|
||||
|
||||
return ChatResponseUpdate(
|
||||
role="assistant",
|
||||
contents=[],
|
||||
additional_properties={
|
||||
"thread_id": self.thread_id,
|
||||
"run_id": self.run_id,
|
||||
},
|
||||
)
|
||||
|
||||
def _handle_text_message_start(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
|
||||
"""Handle TEXT_MESSAGE_START event."""
|
||||
self.current_message_id = event.get("messageId")
|
||||
return ChatResponseUpdate(
|
||||
role="assistant",
|
||||
message_id=self.current_message_id,
|
||||
contents=[],
|
||||
)
|
||||
|
||||
def _handle_text_message_content(self, event: dict[str, Any]) -> ChatResponseUpdate:
|
||||
"""Handle TEXT_MESSAGE_CONTENT event."""
|
||||
message_id = event.get("messageId")
|
||||
delta = event.get("delta", "")
|
||||
|
||||
if message_id != self.current_message_id:
|
||||
self.current_message_id = message_id
|
||||
|
||||
return ChatResponseUpdate(
|
||||
role="assistant",
|
||||
message_id=self.current_message_id,
|
||||
contents=[Content.from_text(text=delta)],
|
||||
)
|
||||
|
||||
def _handle_text_message_end(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
|
||||
"""Handle TEXT_MESSAGE_END event."""
|
||||
return None
|
||||
|
||||
def _handle_tool_call_start(self, event: dict[str, Any]) -> ChatResponseUpdate:
|
||||
"""Handle TOOL_CALL_START event."""
|
||||
self.current_tool_call_id = self._get_tool_call_id(event)
|
||||
self.current_tool_name = event.get("toolName") or event.get("toolCallName") or event.get("tool_call_name")
|
||||
self.accumulated_tool_args = ""
|
||||
|
||||
return ChatResponseUpdate(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id=self.current_tool_call_id or "",
|
||||
name=self.current_tool_name or "",
|
||||
arguments="",
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
def _handle_tool_call_args(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
|
||||
"""Handle TOOL_CALL_ARGS event."""
|
||||
event_tool_call_id = self._get_tool_call_id(event)
|
||||
if event_tool_call_id is not None:
|
||||
if self.current_tool_call_id and event_tool_call_id != self.current_tool_call_id:
|
||||
logger.warning(
|
||||
"Ignoring TOOL_CALL_ARGS for toolCallId=%s while current toolCallId=%s",
|
||||
event_tool_call_id,
|
||||
self.current_tool_call_id,
|
||||
)
|
||||
return None
|
||||
if not self.current_tool_call_id:
|
||||
self.current_tool_call_id = event_tool_call_id
|
||||
|
||||
delta = event.get("delta", "")
|
||||
self.accumulated_tool_args += delta
|
||||
|
||||
return ChatResponseUpdate(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id=self.current_tool_call_id or "",
|
||||
name=self.current_tool_name or "",
|
||||
arguments=delta,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
def _handle_tool_call_end(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
|
||||
"""Handle TOOL_CALL_END event."""
|
||||
event_tool_call_id = self._get_tool_call_id(event)
|
||||
if (
|
||||
self.current_tool_call_id is None
|
||||
or event_tool_call_id is None
|
||||
or event_tool_call_id == self.current_tool_call_id
|
||||
):
|
||||
self.current_tool_call_id = None
|
||||
self.current_tool_name = None
|
||||
self.accumulated_tool_args = ""
|
||||
return None
|
||||
|
||||
def _handle_tool_call_result(self, event: dict[str, Any]) -> ChatResponseUpdate:
|
||||
"""Handle TOOL_CALL_RESULT event."""
|
||||
tool_call_id = event.get("toolCallId", "")
|
||||
result = event.get("result") if event.get("result") is not None else event.get("content")
|
||||
|
||||
return ChatResponseUpdate(
|
||||
role="tool",
|
||||
contents=[
|
||||
Content.from_function_result(
|
||||
call_id=tool_call_id,
|
||||
result=result,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
def _handle_run_finished(self, event: dict[str, Any]) -> ChatResponseUpdate:
|
||||
"""Handle RUN_FINISHED event."""
|
||||
additional_properties: dict[str, Any] = {
|
||||
"thread_id": self.thread_id,
|
||||
"run_id": self.run_id,
|
||||
}
|
||||
if "interrupt" in event:
|
||||
additional_properties["interrupt"] = event.get("interrupt")
|
||||
if "outcome" in event:
|
||||
outcome = event.get("outcome")
|
||||
additional_properties["outcome"] = outcome
|
||||
if not isinstance(outcome, dict):
|
||||
logger.warning(
|
||||
"RUN_FINISHED outcome should be an object; got %s. Preserving raw outcome.",
|
||||
type(outcome).__name__,
|
||||
)
|
||||
elif outcome.get("type") == "interrupt":
|
||||
interrupts = outcome.get("interrupts")
|
||||
if isinstance(interrupts, list):
|
||||
additional_properties["interrupts"] = interrupts
|
||||
if "result" in event:
|
||||
additional_properties["result"] = event.get("result")
|
||||
|
||||
return ChatResponseUpdate(
|
||||
role="assistant",
|
||||
finish_reason="stop",
|
||||
contents=[],
|
||||
additional_properties=additional_properties,
|
||||
)
|
||||
|
||||
def _handle_run_error(self, event: dict[str, Any]) -> ChatResponseUpdate:
|
||||
"""Handle RUN_ERROR event."""
|
||||
error_message = event.get("message", "Unknown error")
|
||||
|
||||
return ChatResponseUpdate(
|
||||
role="assistant",
|
||||
finish_reason="content_filter",
|
||||
contents=[
|
||||
Content.from_error(
|
||||
message=error_message,
|
||||
error_code="RUN_ERROR",
|
||||
)
|
||||
],
|
||||
additional_properties={
|
||||
"thread_id": self.thread_id,
|
||||
"run_id": self.run_id,
|
||||
},
|
||||
)
|
||||
|
||||
def _handle_custom_event(self, event: dict[str, Any], raw_event_type: str) -> ChatResponseUpdate:
|
||||
"""Handle CUSTOM/CUSTOM_EVENT events.
|
||||
|
||||
Custom events are surfaced as metadata so callers can inspect protocol-specific payloads.
|
||||
"""
|
||||
return ChatResponseUpdate(
|
||||
role="assistant",
|
||||
contents=[],
|
||||
additional_properties={
|
||||
"thread_id": self.thread_id,
|
||||
"run_id": self.run_id,
|
||||
"ag_ui_custom_event": {
|
||||
"name": event.get("name"),
|
||||
"value": event.get("value"),
|
||||
"raw_type": raw_event_type,
|
||||
},
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,265 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""HTTP service for AG-UI protocol communication."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterable, Mapping, Sequence
|
||||
from typing import Any, cast
|
||||
|
||||
import httpx
|
||||
from ag_ui.core import Interrupt, ResumeEntry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _json_safe_protocol_value(value: Any) -> Any:
|
||||
"""Convert protocol values to JSON-compatible data using AG-UI aliases."""
|
||||
model_dump = getattr(value, "model_dump", None)
|
||||
if callable(model_dump):
|
||||
return _json_safe_protocol_value(model_dump(by_alias=True, exclude_none=True))
|
||||
if isinstance(value, Mapping):
|
||||
return {key: _json_safe_protocol_value(item) for key, item in value.items()}
|
||||
if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
|
||||
return [_json_safe_protocol_value(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
def _serialize_available_interrupts(available_interrupts: Sequence[Any] | None) -> list[dict[str, Any]] | None:
|
||||
"""Serialize typed or compatible interrupt inputs to canonical AG-UI JSON."""
|
||||
if available_interrupts is None:
|
||||
return None
|
||||
serialized: list[dict[str, Any]] = []
|
||||
for interrupt in available_interrupts:
|
||||
if isinstance(interrupt, Mapping) and "reason" not in interrupt:
|
||||
interrupt = dict(interrupt)
|
||||
interrupt_type = interrupt.pop("type", None)
|
||||
if interrupt_type == "request_info" or interrupt_type is None:
|
||||
interrupt["reason"] = "input_required"
|
||||
elif isinstance(interrupt_type, str):
|
||||
interrupt["reason"] = interrupt_type
|
||||
serialized.append(
|
||||
cast(dict[str, Any], Interrupt.model_validate(interrupt).model_dump(by_alias=True, exclude_none=True))
|
||||
)
|
||||
return serialized
|
||||
|
||||
|
||||
def _serialize_resume_entry(entry: Any) -> dict[str, Any]:
|
||||
"""Serialize one typed or legacy resume entry to canonical AG-UI JSON."""
|
||||
model_dump = getattr(entry, "model_dump", None)
|
||||
if callable(model_dump):
|
||||
entry = model_dump(by_alias=True, exclude_none=True)
|
||||
|
||||
if not isinstance(entry, Mapping):
|
||||
raise ValueError("Each resume entry must be an object.")
|
||||
|
||||
entry_dict = cast(Mapping[str, Any], entry)
|
||||
interrupt_id = (
|
||||
entry_dict.get("interruptId")
|
||||
or entry_dict.get("interrupt_id")
|
||||
or entry_dict.get("id")
|
||||
or entry_dict.get("toolCallId")
|
||||
)
|
||||
if not interrupt_id:
|
||||
raise ValueError("Each resume entry must include interruptId.")
|
||||
|
||||
status = entry_dict.get("status") or "resolved"
|
||||
payload = (
|
||||
entry_dict.get("payload")
|
||||
if "payload" in entry_dict
|
||||
else entry_dict.get("value")
|
||||
if "value" in entry_dict
|
||||
else entry_dict.get("response")
|
||||
if "response" in entry_dict
|
||||
else {
|
||||
key: value
|
||||
for key, value in entry_dict.items()
|
||||
if key not in {"id", "interruptId", "interrupt_id", "toolCallId", "type", "status"}
|
||||
}
|
||||
)
|
||||
|
||||
serialized: dict[str, Any] = {"interruptId": str(interrupt_id), "status": str(status)}
|
||||
if status != "cancelled" or payload:
|
||||
serialized["payload"] = _json_safe_protocol_value(payload)
|
||||
return cast(dict[str, Any], ResumeEntry.model_validate(serialized).model_dump(by_alias=True, exclude_none=True))
|
||||
|
||||
|
||||
def _serialize_resume(resume: Any) -> Any: # noqa: ANN401
|
||||
"""Serialize typed or compatible resume inputs to canonical AG-UI JSON."""
|
||||
if resume is None:
|
||||
return None
|
||||
if isinstance(resume, Sequence) and not isinstance(resume, (str, bytes, bytearray)):
|
||||
return [_serialize_resume_entry(entry) for entry in resume]
|
||||
if isinstance(resume, Mapping):
|
||||
resume_dict = cast(Mapping[str, Any], resume)
|
||||
if isinstance(resume_dict.get("interrupts"), Sequence) and not isinstance(
|
||||
resume_dict.get("interrupts"), (str, bytes, bytearray)
|
||||
):
|
||||
return [_serialize_resume_entry(entry) for entry in cast(Sequence[Any], resume_dict["interrupts"])]
|
||||
if isinstance(resume_dict.get("interrupt"), Sequence) and not isinstance(
|
||||
resume_dict.get("interrupt"), (str, bytes, bytearray)
|
||||
):
|
||||
return [_serialize_resume_entry(entry) for entry in cast(Sequence[Any], resume_dict["interrupt"])]
|
||||
if any(key in resume_dict for key in ("interruptId", "interrupt_id", "id", "toolCallId")):
|
||||
return [_serialize_resume_entry(resume_dict)]
|
||||
return _json_safe_protocol_value(resume)
|
||||
|
||||
|
||||
class AGUIHttpService:
|
||||
"""HTTP service for AG-UI protocol communication.
|
||||
|
||||
Handles HTTP POST requests and Server-Sent Events (SSE) stream parsing
|
||||
for the AG-UI protocol.
|
||||
|
||||
Examples:
|
||||
Basic usage:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
service = AGUIHttpService("http://localhost:8888/")
|
||||
async for event in service.post_run(
|
||||
thread_id="thread_123",
|
||||
run_id="run_456",
|
||||
messages=[{"role": "user", "content": "Hello"}]
|
||||
):
|
||||
print(event["type"])
|
||||
|
||||
With context manager:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
async with AGUIHttpService("http://localhost:8888/") as service:
|
||||
async for event in service.post_run(...):
|
||||
print(event)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
endpoint: str,
|
||||
http_client: httpx.AsyncClient | None = None,
|
||||
timeout: float = 60.0,
|
||||
) -> None:
|
||||
"""Initialize the HTTP service.
|
||||
|
||||
Args:
|
||||
endpoint: AG-UI server endpoint URL (e.g., "http://localhost:8888/")
|
||||
http_client: Optional httpx AsyncClient. If None, creates a new one.
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
"""
|
||||
self.endpoint = endpoint.rstrip("/")
|
||||
self._owns_client = http_client is None
|
||||
self.http_client = http_client or httpx.AsyncClient(timeout=timeout)
|
||||
|
||||
async def post_run(
|
||||
self,
|
||||
thread_id: str,
|
||||
run_id: str,
|
||||
messages: list[dict[str, Any]],
|
||||
state: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
available_interrupts: Sequence[Any] | None = None,
|
||||
resume: Any = None,
|
||||
) -> AsyncIterable[dict[str, Any]]:
|
||||
"""Post a run request and stream AG-UI events.
|
||||
|
||||
Args:
|
||||
thread_id: Thread identifier for conversation continuity
|
||||
run_id: Unique run identifier
|
||||
messages: List of messages in AG-UI format
|
||||
state: Optional state object to send to server
|
||||
tools: Optional list of tools available to the agent
|
||||
available_interrupts: Optional list of interrupt descriptors available for resumption
|
||||
resume: Optional resume payload to continue a paused run
|
||||
|
||||
Yields:
|
||||
AG-UI event dictionaries parsed from SSE stream
|
||||
|
||||
Raises:
|
||||
httpx.HTTPStatusError: If the HTTP request fails
|
||||
ValueError: If SSE parsing encounters invalid data
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
service = AGUIHttpService("http://localhost:8888/")
|
||||
async for event in service.post_run(
|
||||
thread_id="thread_abc",
|
||||
run_id="run_123",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
state={"user_context": {"name": "Alice"}}
|
||||
):
|
||||
if event["type"] == "TEXT_MESSAGE_CONTENT":
|
||||
print(event["delta"])
|
||||
"""
|
||||
# Build request payload
|
||||
request_data: dict[str, Any] = {
|
||||
"thread_id": thread_id,
|
||||
"run_id": run_id,
|
||||
"messages": messages,
|
||||
}
|
||||
|
||||
if state is not None:
|
||||
request_data["state"] = state
|
||||
|
||||
if tools is not None:
|
||||
request_data["tools"] = tools
|
||||
|
||||
serialized_available_interrupts = _serialize_available_interrupts(available_interrupts)
|
||||
if serialized_available_interrupts is not None:
|
||||
request_data["availableInterrupts"] = serialized_available_interrupts
|
||||
|
||||
serialized_resume = _serialize_resume(resume)
|
||||
if serialized_resume is not None:
|
||||
request_data["resume"] = serialized_resume
|
||||
|
||||
logger.debug(
|
||||
f"Posting run to {self.endpoint}: thread_id={thread_id}, run_id={run_id}, "
|
||||
f"messages={len(messages)}, has_state={state is not None}, has_tools={tools is not None}, "
|
||||
f"has_available_interrupts={available_interrupts is not None}, has_resume={resume is not None}"
|
||||
)
|
||||
|
||||
# Stream the response using SSE
|
||||
async with self.http_client.stream(
|
||||
"POST",
|
||||
self.endpoint,
|
||||
json=request_data,
|
||||
headers={"Accept": "text/event-stream"},
|
||||
) as response:
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.error(f"HTTP request failed: {e.response.status_code} - {e.response.text}")
|
||||
raise
|
||||
|
||||
async for line in response.aiter_lines():
|
||||
# Parse Server-Sent Events format
|
||||
if line.startswith("data: "):
|
||||
data = line[6:] # Remove "data: " prefix
|
||||
try:
|
||||
event = json.loads(data)
|
||||
logger.debug(f"Received event: {event.get('type', 'UNKNOWN')}")
|
||||
yield event
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"Failed to parse SSE data: {data}. Error: {e}")
|
||||
# Continue processing other events instead of failing
|
||||
continue
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the HTTP client if owned by this service.
|
||||
|
||||
Only closes the client if it was created by this service instance.
|
||||
If an external client was provided, it remains the caller's
|
||||
responsibility to close it.
|
||||
"""
|
||||
if self._owns_client and self.http_client:
|
||||
await self.http_client.aclose()
|
||||
|
||||
async def __aenter__(self) -> AGUIHttpService:
|
||||
"""Enter async context manager."""
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
"""Exit async context manager and clean up resources."""
|
||||
await self.close()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
@@ -0,0 +1,247 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Helper functions for orchestration logic.
|
||||
|
||||
This module retains utilities that may be useful for testing or extensions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import (
|
||||
Content,
|
||||
Message,
|
||||
)
|
||||
|
||||
from .._utils import get_role_value
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def pending_tool_call_ids(messages: list[Message]) -> set[str]:
|
||||
"""Get IDs of tool calls without corresponding results.
|
||||
|
||||
Args:
|
||||
messages: List of messages to scan
|
||||
|
||||
Returns:
|
||||
Set of pending tool call IDs
|
||||
"""
|
||||
pending_ids: set[str] = set()
|
||||
resolved_ids: set[str] = set()
|
||||
for msg in messages:
|
||||
for content in msg.contents:
|
||||
if content.type == "function_call" and content.call_id:
|
||||
pending_ids.add(str(content.call_id))
|
||||
elif content.type == "function_result" and content.call_id:
|
||||
resolved_ids.add(str(content.call_id))
|
||||
return pending_ids - resolved_ids
|
||||
|
||||
|
||||
def is_state_context_message(message: Message) -> bool:
|
||||
"""Check if a message is a state context system message.
|
||||
|
||||
Args:
|
||||
message: Message to check
|
||||
|
||||
Returns:
|
||||
True if this is a state context message
|
||||
"""
|
||||
if get_role_value(message) != "system":
|
||||
return False
|
||||
for content in message.contents:
|
||||
if content.type == "text" and content.text.startswith("Current state of the application:"): # type: ignore[union-attr]
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def ensure_tool_call_entry(
|
||||
tool_call_id: str,
|
||||
tool_calls_by_id: dict[str, dict[str, Any]],
|
||||
pending_tool_calls: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""Get or create a tool call entry in the tracking dicts.
|
||||
|
||||
Args:
|
||||
tool_call_id: The tool call ID
|
||||
tool_calls_by_id: Dict mapping IDs to tool call entries
|
||||
pending_tool_calls: List of pending tool calls
|
||||
|
||||
Returns:
|
||||
The tool call entry dict
|
||||
"""
|
||||
entry = tool_calls_by_id.get(tool_call_id)
|
||||
if entry is None:
|
||||
entry = {
|
||||
"id": tool_call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "",
|
||||
"arguments": "",
|
||||
},
|
||||
}
|
||||
tool_calls_by_id[tool_call_id] = entry
|
||||
pending_tool_calls.append(entry)
|
||||
return entry
|
||||
|
||||
|
||||
def tool_name_for_call_id(
|
||||
tool_calls_by_id: dict[str, dict[str, Any]],
|
||||
tool_call_id: str,
|
||||
) -> str | None:
|
||||
"""Get the tool name for a given call ID.
|
||||
|
||||
Args:
|
||||
tool_calls_by_id: Dict mapping IDs to tool call entries
|
||||
tool_call_id: The tool call ID to look up
|
||||
|
||||
Returns:
|
||||
Tool name or None if not found
|
||||
"""
|
||||
entry = tool_calls_by_id.get(tool_call_id)
|
||||
if not entry:
|
||||
return None
|
||||
function = entry.get("function")
|
||||
if not isinstance(function, dict):
|
||||
return None
|
||||
name = function.get("name")
|
||||
return str(name) if name else None
|
||||
|
||||
|
||||
def schema_has_steps(schema: Any) -> bool:
|
||||
"""Check if a schema has a steps array property.
|
||||
|
||||
Args:
|
||||
schema: JSON schema to check
|
||||
|
||||
Returns:
|
||||
True if schema has steps array
|
||||
"""
|
||||
if not isinstance(schema, dict):
|
||||
return False
|
||||
properties = schema.get("properties")
|
||||
if not isinstance(properties, dict):
|
||||
return False
|
||||
steps_schema = properties.get("steps")
|
||||
if not isinstance(steps_schema, dict):
|
||||
return False
|
||||
return steps_schema.get("type") == "array"
|
||||
|
||||
|
||||
def select_approval_tool_name(client_tools: list[Any] | None) -> str | None:
|
||||
"""Select appropriate approval tool from client tools.
|
||||
|
||||
Args:
|
||||
client_tools: List of client tool definitions
|
||||
|
||||
Returns:
|
||||
Name of approval tool, or None if not found
|
||||
"""
|
||||
if not client_tools:
|
||||
return None
|
||||
for tool in client_tools:
|
||||
tool_name = getattr(tool, "name", None)
|
||||
if not tool_name:
|
||||
continue
|
||||
params_fn = getattr(tool, "parameters", None)
|
||||
if not callable(params_fn):
|
||||
continue
|
||||
schema = params_fn()
|
||||
if schema_has_steps(schema):
|
||||
return str(tool_name)
|
||||
return None
|
||||
|
||||
|
||||
def build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any]:
|
||||
"""Build metadata dict with truncated string values for Azure compatibility.
|
||||
|
||||
Azure has a 512 character limit per metadata value.
|
||||
|
||||
Args:
|
||||
thread_metadata: Raw metadata dict
|
||||
|
||||
Returns:
|
||||
Metadata with string values truncated to 512 chars
|
||||
"""
|
||||
if not thread_metadata:
|
||||
return {}
|
||||
safe_metadata: dict[str, Any] = {}
|
||||
for key, value in thread_metadata.items():
|
||||
value_str = value if isinstance(value, str) else json.dumps(value)
|
||||
if len(value_str) > 512:
|
||||
value_str = value_str[:512]
|
||||
safe_metadata[key] = value_str
|
||||
return safe_metadata
|
||||
|
||||
|
||||
def latest_approval_response(messages: list[Message]) -> Content | None:
|
||||
"""Get the latest approval response from messages.
|
||||
|
||||
Args:
|
||||
messages: Messages to search
|
||||
|
||||
Returns:
|
||||
Latest approval response or None
|
||||
"""
|
||||
if not messages:
|
||||
return None
|
||||
last_message = messages[-1]
|
||||
for content in last_message.contents:
|
||||
if content.type == "function_approval_response":
|
||||
return content
|
||||
return None
|
||||
|
||||
|
||||
def approval_steps(approval: Content) -> list[Any]:
|
||||
"""Extract steps from an approval response.
|
||||
|
||||
Args:
|
||||
approval: Approval response content
|
||||
|
||||
Returns:
|
||||
List of steps, or empty list if none
|
||||
"""
|
||||
state_args = approval.additional_properties.get("ag_ui_state_args", None)
|
||||
if isinstance(state_args, dict):
|
||||
steps = state_args.get("steps")
|
||||
if isinstance(steps, list):
|
||||
return steps
|
||||
|
||||
if approval.function_call:
|
||||
parsed_args = approval.function_call.parse_arguments()
|
||||
if isinstance(parsed_args, dict):
|
||||
steps = parsed_args.get("steps")
|
||||
if isinstance(steps, list):
|
||||
return steps
|
||||
|
||||
return []
|
||||
|
||||
|
||||
def is_step_based_approval(
|
||||
approval: Content,
|
||||
predict_state_config: dict[str, dict[str, str]] | None,
|
||||
) -> bool:
|
||||
"""Check if an approval is step-based.
|
||||
|
||||
Args:
|
||||
approval: Approval response to check
|
||||
predict_state_config: Predictive state configuration
|
||||
|
||||
Returns:
|
||||
True if this is a step-based approval
|
||||
"""
|
||||
steps = approval_steps(approval)
|
||||
if steps:
|
||||
return True
|
||||
if not approval.function_call:
|
||||
return False
|
||||
if not predict_state_config:
|
||||
return False
|
||||
tool_name = approval.function_call.name
|
||||
for config in predict_state_config.values():
|
||||
if config.get("tool") == tool_name and config.get("tool_argument") == "steps":
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,232 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Predictive state handling utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from ag_ui.core import StateDeltaEvent
|
||||
|
||||
from .._utils import safe_json_parse
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PredictiveStateHandler:
|
||||
"""Handles predictive state updates from streaming tool calls."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
predict_state_config: dict[str, dict[str, str]] | None = None,
|
||||
current_state: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Initialize the handler.
|
||||
|
||||
Args:
|
||||
predict_state_config: Configuration mapping state keys to tool/argument pairs
|
||||
current_state: Reference to current state dict
|
||||
"""
|
||||
self.predict_state_config = predict_state_config or {}
|
||||
self.current_state = current_state or {}
|
||||
self.streaming_tool_args: str = ""
|
||||
self.last_emitted_state: dict[str, Any] = {}
|
||||
self.state_delta_count: int = 0
|
||||
self.pending_state_updates: dict[str, Any] = {}
|
||||
|
||||
def reset_streaming(self) -> None:
|
||||
"""Reset streaming state for a new tool call."""
|
||||
self.streaming_tool_args = ""
|
||||
self.state_delta_count = 0
|
||||
|
||||
def extract_state_value(
|
||||
self,
|
||||
tool_name: str,
|
||||
args: dict[str, Any] | str | None,
|
||||
) -> tuple[str, Any] | None:
|
||||
"""Extract state value from tool arguments based on config.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool being called
|
||||
args: Tool arguments (dict or JSON string)
|
||||
|
||||
Returns:
|
||||
Tuple of (state_key, state_value) or None if no match
|
||||
"""
|
||||
if not self.predict_state_config:
|
||||
return None
|
||||
|
||||
parsed_args = safe_json_parse(args) if isinstance(args, str) else args
|
||||
if not parsed_args:
|
||||
return None
|
||||
|
||||
for state_key, config in self.predict_state_config.items():
|
||||
if config["tool"] != tool_name:
|
||||
continue
|
||||
tool_arg_name = config["tool_argument"]
|
||||
if tool_arg_name == "*":
|
||||
return (state_key, parsed_args)
|
||||
if tool_arg_name in parsed_args:
|
||||
return (state_key, parsed_args[tool_arg_name])
|
||||
|
||||
return None
|
||||
|
||||
def is_predictive_tool(self, tool_name: str | None) -> bool:
|
||||
"""Check if a tool is configured for predictive state.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool to check
|
||||
|
||||
Returns:
|
||||
True if tool is in predictive state config
|
||||
"""
|
||||
if not tool_name or not self.predict_state_config:
|
||||
return False
|
||||
for config in self.predict_state_config.values():
|
||||
if config["tool"] == tool_name:
|
||||
return True
|
||||
return False
|
||||
|
||||
def emit_streaming_deltas(
|
||||
self,
|
||||
tool_name: str | None,
|
||||
argument_chunk: str,
|
||||
) -> list[StateDeltaEvent]:
|
||||
"""Process streaming argument chunk and emit state deltas.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the current tool
|
||||
argument_chunk: New chunk of JSON arguments
|
||||
|
||||
Returns:
|
||||
List of state delta events to emit
|
||||
"""
|
||||
events: list[StateDeltaEvent] = []
|
||||
if not tool_name or not self.predict_state_config:
|
||||
return events
|
||||
|
||||
self.streaming_tool_args += argument_chunk
|
||||
logger.debug(
|
||||
"Predictive state: accumulated %s chars for tool '%s'",
|
||||
len(self.streaming_tool_args),
|
||||
tool_name,
|
||||
)
|
||||
|
||||
# Try to parse complete JSON first
|
||||
parsed_args = None
|
||||
try:
|
||||
parsed_args = json.loads(self.streaming_tool_args)
|
||||
except json.JSONDecodeError:
|
||||
# Fall back to regex matching for partial JSON
|
||||
events.extend(self._emit_partial_deltas(tool_name))
|
||||
|
||||
if parsed_args:
|
||||
events.extend(self._emit_complete_deltas(tool_name, parsed_args))
|
||||
|
||||
return events
|
||||
|
||||
def _emit_partial_deltas(self, tool_name: str) -> list[StateDeltaEvent]:
|
||||
"""Emit deltas from partial JSON using regex matching.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the current tool
|
||||
|
||||
Returns:
|
||||
List of state delta events
|
||||
"""
|
||||
events: list[StateDeltaEvent] = []
|
||||
|
||||
for state_key, config in self.predict_state_config.items():
|
||||
if config["tool"] != tool_name:
|
||||
continue
|
||||
tool_arg_name = config["tool_argument"]
|
||||
pattern = rf'"{re.escape(tool_arg_name)}":\s*"([^"]*)'
|
||||
match = re.search(pattern, self.streaming_tool_args)
|
||||
|
||||
if match:
|
||||
partial_value = match.group(1).replace("\\n", "\n").replace('\\"', '"').replace("\\\\", "\\")
|
||||
|
||||
if state_key not in self.last_emitted_state or self.last_emitted_state[state_key] != partial_value:
|
||||
event = self._create_delta_event(state_key, partial_value)
|
||||
events.append(event)
|
||||
self.last_emitted_state[state_key] = partial_value
|
||||
self.pending_state_updates[state_key] = partial_value
|
||||
|
||||
return events
|
||||
|
||||
def _emit_complete_deltas(
|
||||
self,
|
||||
tool_name: str,
|
||||
parsed_args: dict[str, Any],
|
||||
) -> list[StateDeltaEvent]:
|
||||
"""Emit deltas from complete parsed JSON.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the current tool
|
||||
parsed_args: Fully parsed arguments dict
|
||||
|
||||
Returns:
|
||||
List of state delta events
|
||||
"""
|
||||
events: list[StateDeltaEvent] = []
|
||||
|
||||
for state_key, config in self.predict_state_config.items():
|
||||
if config["tool"] != tool_name:
|
||||
continue
|
||||
tool_arg_name = config["tool_argument"]
|
||||
|
||||
if tool_arg_name == "*":
|
||||
state_value = parsed_args
|
||||
elif tool_arg_name in parsed_args:
|
||||
state_value = parsed_args[tool_arg_name]
|
||||
else:
|
||||
continue
|
||||
|
||||
if state_key not in self.last_emitted_state or self.last_emitted_state[state_key] != state_value:
|
||||
event = self._create_delta_event(state_key, state_value)
|
||||
events.append(event)
|
||||
self.last_emitted_state[state_key] = state_value
|
||||
self.pending_state_updates[state_key] = state_value
|
||||
|
||||
return events
|
||||
|
||||
def _create_delta_event(self, state_key: str, value: Any) -> StateDeltaEvent:
|
||||
"""Create a state delta event with logging.
|
||||
|
||||
Args:
|
||||
state_key: The state key being updated
|
||||
value: The new value
|
||||
|
||||
Returns:
|
||||
StateDeltaEvent instance
|
||||
"""
|
||||
self.state_delta_count += 1
|
||||
if self.state_delta_count % 10 == 1:
|
||||
logger.info(
|
||||
"StateDeltaEvent #%s for '%s': op=replace, path=/%s, value_length=%s",
|
||||
self.state_delta_count,
|
||||
state_key,
|
||||
state_key,
|
||||
len(str(value)),
|
||||
)
|
||||
elif self.state_delta_count % 100 == 0:
|
||||
logger.info(f"StateDeltaEvent #{self.state_delta_count} emitted")
|
||||
|
||||
return StateDeltaEvent(
|
||||
delta=[
|
||||
{
|
||||
"op": "replace",
|
||||
"path": f"/{state_key}",
|
||||
"value": value,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
def apply_pending_updates(self) -> None:
|
||||
"""Apply pending updates to current state and clear them."""
|
||||
for key, value in self.pending_state_updates.items():
|
||||
self.current_state[key] = value
|
||||
self.pending_state_updates.clear()
|
||||
@@ -0,0 +1,126 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tool handling helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from agent_framework import BaseChatClient
|
||||
from agent_framework._tools import _append_unique_tools # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agent_framework import SupportsAgentRun
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _collect_mcp_tool_functions(mcp_tools: list[Any]) -> list[Any]:
|
||||
"""Extract functions from connected MCP tools.
|
||||
|
||||
Args:
|
||||
mcp_tools: List of MCP tool instances.
|
||||
|
||||
Returns:
|
||||
Functions from connected MCP tools.
|
||||
"""
|
||||
functions: list[Any] = []
|
||||
for mcp_tool in mcp_tools:
|
||||
if getattr(mcp_tool, "is_connected", False) and hasattr(mcp_tool, "functions"):
|
||||
functions.extend(mcp_tool.functions)
|
||||
return functions
|
||||
|
||||
|
||||
def collect_server_tools(agent: SupportsAgentRun) -> list[Any]:
|
||||
"""Collect server tools from an agent.
|
||||
|
||||
This includes both regular tools from default_options and MCP tools.
|
||||
MCP tools are stored separately for lifecycle management but their
|
||||
functions need to be included for tool execution during approval flows.
|
||||
|
||||
Args:
|
||||
agent: Agent instance to collect tools from. Works with Agent
|
||||
or any agent with default_options and optional mcp_tools attributes.
|
||||
|
||||
Returns:
|
||||
List of tools including both regular tools and connected MCP tool functions.
|
||||
"""
|
||||
# Get tools from default_options
|
||||
default_options = getattr(agent, "default_options", None)
|
||||
if default_options is None:
|
||||
return []
|
||||
|
||||
tools_from_agent = default_options.get("tools") if isinstance(default_options, dict) else None
|
||||
server_tools = list(tools_from_agent) if tools_from_agent else []
|
||||
|
||||
# Include functions from connected MCP tools (only available on Agent)
|
||||
mcp_tools = getattr(agent, "mcp_tools", None)
|
||||
if mcp_tools:
|
||||
_append_unique_tools(
|
||||
server_tools,
|
||||
_collect_mcp_tool_functions(mcp_tools),
|
||||
duplicate_error_message="Tool names must be unique. Consider setting `tool_name_prefix` on the MCPTool.",
|
||||
)
|
||||
|
||||
logger.info(f"[TOOLS] Agent has {len(server_tools)} configured tools")
|
||||
for tool in server_tools:
|
||||
tool_name = getattr(tool, "name", "unknown")
|
||||
approval_mode = getattr(tool, "approval_mode", None)
|
||||
logger.info(f"[TOOLS] - {tool_name}: approval_mode={approval_mode}")
|
||||
return server_tools
|
||||
|
||||
|
||||
def register_additional_client_tools(agent: SupportsAgentRun, client_tools: list[Any] | None) -> None:
|
||||
"""Register client tools as additional declaration-only tools to avoid server execution.
|
||||
|
||||
Args:
|
||||
agent: Agent instance to register tools on. Works with Agent
|
||||
or any agent with a client attribute.
|
||||
client_tools: List of client tools to register.
|
||||
"""
|
||||
if not client_tools:
|
||||
return
|
||||
|
||||
client = getattr(agent, "client", None)
|
||||
if client is None:
|
||||
return
|
||||
|
||||
if isinstance(client, BaseChatClient) and client.function_invocation_configuration is not None: # type: ignore[attr-defined]
|
||||
client.function_invocation_configuration["additional_tools"] = client_tools # type: ignore[attr-defined]
|
||||
logger.debug(f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)")
|
||||
|
||||
|
||||
def _has_approval_tools(tools: list[Any]) -> bool:
|
||||
"""Check if any tools require approval."""
|
||||
return any(getattr(tool, "approval_mode", None) == "always_require" for tool in tools)
|
||||
|
||||
|
||||
def merge_tools(server_tools: list[Any], client_tools: list[Any] | None) -> list[Any] | None:
|
||||
"""Combine server and client tools without overriding server metadata.
|
||||
|
||||
IMPORTANT: When server tools have approval_mode="always_require", we MUST return
|
||||
them so they get passed to the streaming response handler. Otherwise, the approval
|
||||
check in _try_execute_function_calls won't find the tool and won't trigger approval.
|
||||
"""
|
||||
if not client_tools:
|
||||
# Even without client tools, we must pass server tools if any require approval
|
||||
if server_tools and _has_approval_tools(server_tools):
|
||||
logger.info(
|
||||
f"[TOOLS] No client tools but server has approval tools - "
|
||||
f"passing {len(server_tools)} server tools for approval mode"
|
||||
)
|
||||
return server_tools
|
||||
logger.info("[TOOLS] No client tools - not passing tools= parameter (using agent's configured tools)")
|
||||
return None
|
||||
|
||||
combined_tools = _append_unique_tools(
|
||||
list(server_tools),
|
||||
client_tools,
|
||||
duplicate_error_message="Tool names must be unique.",
|
||||
)
|
||||
logger.info(
|
||||
f"[TOOLS] Passing tools= parameter with {len(combined_tools)} tools "
|
||||
f"({len(server_tools)} server + {len(client_tools)} client)"
|
||||
)
|
||||
return combined_tools
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,234 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AG-UI Thread Snapshot storage primitives."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Protocol, TypeAlias, runtime_checkable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._types import AGUIRequest
|
||||
|
||||
SnapshotScope: TypeAlias = str
|
||||
"""Application-defined scope for authorizing access to AG-UI Thread Snapshots."""
|
||||
|
||||
AGUIThreadID: TypeAlias = str
|
||||
"""AG-UI Thread identifier within a Snapshot Scope."""
|
||||
|
||||
SnapshotScopeResolver: TypeAlias = Callable[["AGUIRequest"], str | Awaitable[str]]
|
||||
"""Callable that resolves the Snapshot Scope for an AG-UI endpoint request."""
|
||||
|
||||
_SnapshotKey: TypeAlias = tuple[SnapshotScope, AGUIThreadID]
|
||||
|
||||
DEFAULT_MAX_THREAD_SNAPSHOTS = 1_000
|
||||
_SNAPSHOT_SCOPE_INPUT_KEY = "__ag_ui_snapshot_scope"
|
||||
_DEFAULT_STATE_INPUT_KEY = "__ag_ui_default_state"
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AGUIThreadSnapshot:
|
||||
"""Replayable AG-UI Thread state.
|
||||
|
||||
AG-UI Thread Snapshots intentionally contain only data that can be replayed
|
||||
to a UI: message snapshots, optional Shared State, and optional interruption
|
||||
state. They do not include raw events, request metadata, auth claims,
|
||||
diagnostics, traces, or provider responses.
|
||||
|
||||
Attributes:
|
||||
messages: Replayable AG-UI message snapshots.
|
||||
state: Optional AG-UI Shared State snapshot.
|
||||
interrupt: Optional interruption state from ``RUN_FINISHED.outcome.interrupts``.
|
||||
"""
|
||||
|
||||
messages: list[dict[str, Any]] = field(default_factory=list)
|
||||
state: dict[str, Any] | None = None
|
||||
interrupt: list[dict[str, Any]] | None = None
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class AGUIThreadSnapshotStore(Protocol):
|
||||
"""Async store for latest AG-UI Thread Snapshots keyed by scope and thread id."""
|
||||
|
||||
async def save(
|
||||
self,
|
||||
*,
|
||||
scope: SnapshotScope,
|
||||
thread_id: AGUIThreadID,
|
||||
snapshot: AGUIThreadSnapshot,
|
||||
) -> None:
|
||||
"""Save the latest snapshot for an AG-UI Thread within a Snapshot Scope.
|
||||
|
||||
Args:
|
||||
scope: Application-defined Snapshot Scope. This is part of the
|
||||
storage key and must represent the app's authorization boundary.
|
||||
thread_id: AG-UI Thread id within the scope.
|
||||
snapshot: Snapshot to save.
|
||||
"""
|
||||
...
|
||||
|
||||
async def get(
|
||||
self,
|
||||
*,
|
||||
scope: SnapshotScope,
|
||||
thread_id: AGUIThreadID,
|
||||
) -> AGUIThreadSnapshot | None:
|
||||
"""Get the latest snapshot for an AG-UI Thread within a Snapshot Scope.
|
||||
|
||||
Args:
|
||||
scope: Application-defined Snapshot Scope.
|
||||
thread_id: AG-UI Thread id within the scope.
|
||||
|
||||
Returns:
|
||||
The latest snapshot, or ``None`` when no snapshot exists for the key.
|
||||
"""
|
||||
...
|
||||
|
||||
async def delete(
|
||||
self,
|
||||
*,
|
||||
scope: SnapshotScope,
|
||||
thread_id: AGUIThreadID,
|
||||
) -> bool:
|
||||
"""Delete the latest snapshot for an AG-UI Thread within a Snapshot Scope.
|
||||
|
||||
Args:
|
||||
scope: Application-defined Snapshot Scope.
|
||||
thread_id: AG-UI Thread id within the scope.
|
||||
|
||||
Returns:
|
||||
``True`` when a snapshot was deleted, otherwise ``False``.
|
||||
"""
|
||||
...
|
||||
|
||||
async def clear(self, *, scope: SnapshotScope | None = None) -> None:
|
||||
"""Clear saved snapshots.
|
||||
|
||||
Args:
|
||||
scope: Optional Snapshot Scope to clear. When omitted, all in-memory
|
||||
snapshots are cleared.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class InMemoryAGUIThreadSnapshotStore:
|
||||
"""Bounded memory-only latest snapshot store for local development, demos, and tests.
|
||||
|
||||
This store keeps at most one snapshot per ``(scope, thread_id)`` key. It is
|
||||
process-local and not durable production storage.
|
||||
"""
|
||||
|
||||
def __init__(self, *, max_snapshots: int = DEFAULT_MAX_THREAD_SNAPSHOTS) -> None:
|
||||
"""Initialize the in-memory snapshot store.
|
||||
|
||||
Keyword Args:
|
||||
max_snapshots: Maximum number of scoped thread snapshots to retain.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``max_snapshots`` is less than 1.
|
||||
"""
|
||||
if max_snapshots < 1:
|
||||
raise ValueError("max_snapshots must be greater than 0.")
|
||||
self._max_snapshots = max_snapshots
|
||||
self._snapshots: dict[_SnapshotKey, AGUIThreadSnapshot] = {}
|
||||
|
||||
async def save(
|
||||
self,
|
||||
*,
|
||||
scope: SnapshotScope,
|
||||
thread_id: AGUIThreadID,
|
||||
snapshot: AGUIThreadSnapshot,
|
||||
) -> None:
|
||||
"""Save the latest snapshot for an AG-UI Thread within a Snapshot Scope."""
|
||||
key = self._key(scope=scope, thread_id=thread_id)
|
||||
if key in self._snapshots:
|
||||
del self._snapshots[key]
|
||||
self._snapshots[key] = copy.deepcopy(snapshot)
|
||||
self._evict_oldest()
|
||||
|
||||
async def get(
|
||||
self,
|
||||
*,
|
||||
scope: SnapshotScope,
|
||||
thread_id: AGUIThreadID,
|
||||
) -> AGUIThreadSnapshot | None:
|
||||
"""Get the latest snapshot for an AG-UI Thread within a Snapshot Scope."""
|
||||
snapshot = self._snapshots.get(self._key(scope=scope, thread_id=thread_id))
|
||||
return copy.deepcopy(snapshot) if snapshot is not None else None
|
||||
|
||||
async def delete(
|
||||
self,
|
||||
*,
|
||||
scope: SnapshotScope,
|
||||
thread_id: AGUIThreadID,
|
||||
) -> bool:
|
||||
"""Delete the latest snapshot for an AG-UI Thread within a Snapshot Scope."""
|
||||
key = self._key(scope=scope, thread_id=thread_id)
|
||||
if key not in self._snapshots:
|
||||
return False
|
||||
del self._snapshots[key]
|
||||
return True
|
||||
|
||||
async def clear(self, *, scope: SnapshotScope | None = None) -> None:
|
||||
"""Clear saved snapshots, optionally limited to one Snapshot Scope."""
|
||||
if scope is None:
|
||||
self._snapshots.clear()
|
||||
return
|
||||
|
||||
normalized_scope = self._normalize_key_part(scope, "scope")
|
||||
for key in list(self._snapshots):
|
||||
if key[0] == normalized_scope:
|
||||
del self._snapshots[key]
|
||||
|
||||
@classmethod
|
||||
def _key(cls, *, scope: SnapshotScope, thread_id: AGUIThreadID) -> _SnapshotKey:
|
||||
return (
|
||||
cls._normalize_key_part(scope, "scope"),
|
||||
cls._normalize_key_part(thread_id, "thread_id"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_key_part(value: str, name: str) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise TypeError(f"{name} must be a string.")
|
||||
if not value:
|
||||
raise ValueError(f"{name} must be a non-empty string.")
|
||||
return value
|
||||
|
||||
def _evict_oldest(self) -> None:
|
||||
while len(self._snapshots) > self._max_snapshots:
|
||||
del self._snapshots[next(iter(self._snapshots))]
|
||||
|
||||
|
||||
async def _clear_thread_snapshot_interrupt(
|
||||
*,
|
||||
snapshot_store: AGUIThreadSnapshotStore,
|
||||
scope: SnapshotScope,
|
||||
thread_id: AGUIThreadID,
|
||||
interrupt_ids: set[str] | None = None,
|
||||
) -> None:
|
||||
"""Clear completed interruption state from the latest replayable thread snapshot."""
|
||||
try:
|
||||
snapshot = await snapshot_store.get(scope=scope, thread_id=thread_id)
|
||||
if snapshot is None or snapshot.interrupt is None:
|
||||
return
|
||||
if interrupt_ids is None:
|
||||
snapshot.interrupt = None
|
||||
else:
|
||||
remaining_interrupts = [
|
||||
interrupt
|
||||
for interrupt in snapshot.interrupt
|
||||
if str(interrupt.get("id") or interrupt.get("interruptId")) not in interrupt_ids
|
||||
]
|
||||
snapshot.interrupt = remaining_interrupts or None
|
||||
await snapshot_store.save(scope=scope, thread_id=thread_id, snapshot=snapshot)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to clear AG-UI Thread Snapshot interrupt for scope=%s thread_id=%s; keeping previous snapshot.",
|
||||
scope,
|
||||
thread_id,
|
||||
)
|
||||
@@ -0,0 +1,137 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Deterministic tool-driven AG-UI state updates and display payloads.
|
||||
|
||||
Tools wired into the :mod:`agent_framework_ag_ui` endpoint can push a
|
||||
deterministic state update or a per-call tool result display payload by
|
||||
returning :func:`state_update`. Unlike ``predict_state_config`` — which emits
|
||||
``StateDeltaEvent``s optimistically from LLM-predicted tool call arguments —
|
||||
``state_update`` runs *after* the tool executes, so AG-UI state and display
|
||||
content always reflect the tool's actual return value.
|
||||
|
||||
See issue https://github.com/microsoft/agent-framework/issues/3167 for the
|
||||
motivating discussion.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import Content
|
||||
|
||||
from ._utils import make_json_safe
|
||||
|
||||
__all__ = ["TOOL_RESULT_DISPLAY_KEY", "TOOL_RESULT_STATE_KEY", "state_update"]
|
||||
|
||||
|
||||
TOOL_RESULT_STATE_KEY = "__ag_ui_tool_result_state__"
|
||||
"""Reserved ``Content.additional_properties`` key used to carry a tool-driven
|
||||
state snapshot from a tool return value through to the AG-UI emitter."""
|
||||
|
||||
TOOL_RESULT_DISPLAY_KEY = "__ag_ui_tool_result_display__"
|
||||
"""Reserved ``Content.additional_properties`` key used to carry UI-only tool result display content from a tool return value through to the AG-UI emitter."""
|
||||
|
||||
_UNSET = object()
|
||||
|
||||
|
||||
def _serialize_tool_result(value: Any) -> str: # noqa: ANN401
|
||||
return value if isinstance(value, str) else json.dumps(make_json_safe(value))
|
||||
|
||||
|
||||
def state_update(
|
||||
text: str = "",
|
||||
*,
|
||||
state: Mapping[str, Any] | None = None,
|
||||
tool_result: Any = _UNSET, # noqa: ANN401
|
||||
) -> Content:
|
||||
"""Build a tool return value that updates AG-UI shared state or display content.
|
||||
|
||||
Return the result of this helper from an agent tool to push a state update
|
||||
or UI-only display payload to AG-UI clients using the actual tool output,
|
||||
rather than LLM-predicted tool arguments.
|
||||
|
||||
When the AG-UI endpoint emits the tool result, it will:
|
||||
|
||||
* Forward ``text`` to the LLM as the normal ``function_result`` content.
|
||||
* Use ``tool_result`` as the ``ToolCallResultEvent.content`` payload shown
|
||||
to AG-UI clients, falling back to ``text`` when no display payload is set.
|
||||
* Merge ``state`` into ``FlowState.current_state``.
|
||||
* Emit a deterministic ``StateSnapshotEvent`` after the ``ToolCallResult``
|
||||
event so frontends observe the updated state deterministically. If
|
||||
predictive state is enabled, a predictive snapshot may be emitted first.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Content, tool
|
||||
from agent_framework_ag_ui import state_update
|
||||
|
||||
|
||||
@tool
|
||||
async def get_weather(city: str) -> Content:
|
||||
data = await _fetch_weather(city)
|
||||
return state_update(
|
||||
text=f"Weather in {city}: {data['temp']}°C {data['conditions']}",
|
||||
state={"weather": {"city": city, **data}},
|
||||
)
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Content, tool
|
||||
from agent_framework_ag_ui import state_update
|
||||
|
||||
|
||||
@tool
|
||||
async def get_weather(city: str) -> Content:
|
||||
data = await _fetch_weather(city)
|
||||
return state_update(
|
||||
text=f"{city}: {data['temp']}°C and {data['conditions']}",
|
||||
tool_result={
|
||||
"component": "weather-card",
|
||||
"city": city,
|
||||
"temperature": data["temp"],
|
||||
"conditions": data["conditions"],
|
||||
"humidity": data["humidity"],
|
||||
},
|
||||
state={"weather": {"city": city, **data}},
|
||||
)
|
||||
|
||||
Args:
|
||||
text: Text passed back to the LLM as the ``function_result`` content.
|
||||
Defaults to an empty string for tools whose only output is a state
|
||||
update.
|
||||
state: A mapping merged into the AG-UI shared state via JSON-compatible
|
||||
``dict.update`` semantics. Nested dicts are replaced, not deep-merged.
|
||||
tool_result: JSON-safe payload emitted to AG-UI clients as
|
||||
``ToolCallResultEvent.content`` for frontend rendering. The LLM
|
||||
still receives ``text``. If ``text`` is empty, the serialized
|
||||
display payload is also used as the LLM-bound text fallback.
|
||||
|
||||
Returns:
|
||||
A ``Content`` object with ``type="text"``. The state payload rides in
|
||||
``additional_properties`` under :data:`TOOL_RESULT_STATE_KEY`
|
||||
(``"__ag_ui_tool_result_state__"``), and the display payload rides
|
||||
under :data:`TOOL_RESULT_DISPLAY_KEY`
|
||||
(``"__ag_ui_tool_result_display__"``). Both reserved keys are extracted
|
||||
by the AG-UI emitter.
|
||||
|
||||
Raises:
|
||||
TypeError: If ``state`` is not a ``Mapping``.
|
||||
"""
|
||||
if state is not None and not isinstance(state, Mapping):
|
||||
raise TypeError(f"state_update() 'state' must be a Mapping, got {type(state).__name__}")
|
||||
additional_properties: dict[str, Any] = {}
|
||||
if state is not None:
|
||||
additional_properties[TOOL_RESULT_STATE_KEY] = dict(state)
|
||||
if tool_result is not _UNSET:
|
||||
display_content = _serialize_tool_result(tool_result)
|
||||
additional_properties[TOOL_RESULT_DISPLAY_KEY] = display_content
|
||||
if not text:
|
||||
text = display_content
|
||||
return Content.from_text(
|
||||
text,
|
||||
additional_properties=additional_properties,
|
||||
)
|
||||
@@ -0,0 +1,218 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Type definitions for AG-UI integration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Annotated, Any, Generic
|
||||
|
||||
from ag_ui.core import Interrupt, ResumeEntry
|
||||
from agent_framework import ChatOptions
|
||||
from pydantic import AliasChoices, BaseModel, BeforeValidator, Field
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # pragma: no cover
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypedDict # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypedDict # pragma: no cover
|
||||
|
||||
|
||||
AGUIChatOptionsT = TypeVar("AGUIChatOptionsT", bound=TypedDict, default="AGUIChatOptions", covariant=True) # type: ignore[valid-type]
|
||||
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)
|
||||
|
||||
|
||||
def _coerce_legacy_resume_entry(value: Any) -> Any: # noqa: ANN401
|
||||
if not isinstance(value, dict):
|
||||
return value
|
||||
|
||||
interrupt_id = value.get("interruptId") or value.get("interrupt_id") or value.get("id") or value.get("toolCallId")
|
||||
if not interrupt_id:
|
||||
return value
|
||||
|
||||
if "payload" in value:
|
||||
payload = value.get("payload")
|
||||
elif "value" in value:
|
||||
payload = value.get("value")
|
||||
elif "response" in value:
|
||||
payload = value.get("response")
|
||||
else:
|
||||
payload = {
|
||||
key: item
|
||||
for key, item in value.items()
|
||||
if key not in {"id", "interruptId", "interrupt_id", "toolCallId", "type", "status"}
|
||||
}
|
||||
|
||||
entry: dict[str, Any] = {"interruptId": str(interrupt_id), "status": value.get("status", "resolved")}
|
||||
if payload is not None:
|
||||
entry["payload"] = payload
|
||||
return entry
|
||||
|
||||
|
||||
def _coerce_legacy_resume(value: Any) -> Any: # noqa: ANN401
|
||||
if value is None:
|
||||
return value
|
||||
if isinstance(value, dict):
|
||||
if "interrupts" in value:
|
||||
value = value["interrupts"]
|
||||
elif "interrupt" in value:
|
||||
value = value["interrupt"]
|
||||
elif any(key in value for key in ("interruptId", "interrupt_id", "id", "toolCallId")):
|
||||
value = [value]
|
||||
else:
|
||||
return value
|
||||
if not isinstance(value, list):
|
||||
return value
|
||||
return [_coerce_legacy_resume_entry(entry) for entry in value]
|
||||
|
||||
|
||||
class PredictStateConfig(TypedDict):
|
||||
"""Configuration for predictive state updates."""
|
||||
|
||||
state_key: str
|
||||
tool: str
|
||||
tool_argument: str | None
|
||||
|
||||
|
||||
class RunMetadata(TypedDict):
|
||||
"""Metadata for agent run."""
|
||||
|
||||
run_id: str
|
||||
thread_id: str
|
||||
predict_state: list[PredictStateConfig] | None
|
||||
|
||||
|
||||
class AgentState(TypedDict):
|
||||
"""Base state for AG-UI agents."""
|
||||
|
||||
messages: list[Any] | None
|
||||
|
||||
|
||||
class AGUIRequest(BaseModel):
|
||||
"""Request model for AG-UI endpoints."""
|
||||
|
||||
messages: list[dict[str, Any]] = Field(
|
||||
...,
|
||||
description="AG-UI format messages array",
|
||||
)
|
||||
run_id: str | None = Field(
|
||||
None,
|
||||
validation_alias=AliasChoices("run_id", "runId"),
|
||||
description="Optional run identifier for tracking",
|
||||
)
|
||||
thread_id: str | None = Field(
|
||||
None,
|
||||
validation_alias=AliasChoices("thread_id", "threadId"),
|
||||
description="Optional thread identifier for conversation context",
|
||||
)
|
||||
state: dict[str, Any] | None = Field(
|
||||
None,
|
||||
description="Optional shared state for agentic generative UI",
|
||||
)
|
||||
tools: list[dict[str, Any]] | None = Field(
|
||||
None,
|
||||
description="Client-side tools to advertise to the LLM",
|
||||
)
|
||||
context: list[dict[str, Any]] | None = Field(
|
||||
None,
|
||||
description="List of context objects provided to the agent",
|
||||
)
|
||||
forwarded_props: dict[str, Any] | None = Field(
|
||||
None,
|
||||
validation_alias=AliasChoices("forwarded_props", "forwardedProps"),
|
||||
description="Additional properties forwarded to the agent",
|
||||
)
|
||||
parent_run_id: str | None = Field(
|
||||
None,
|
||||
validation_alias=AliasChoices("parent_run_id", "parentRunId"),
|
||||
description="ID of the run that spawned this run",
|
||||
)
|
||||
available_interrupts: list[Interrupt] | None = Field(
|
||||
None,
|
||||
validation_alias=AliasChoices("availableInterrupts", "available_interrupts"),
|
||||
description="Canonical AG-UI interrupts that can be resumed by the server",
|
||||
)
|
||||
resume: Annotated[list[ResumeEntry], BeforeValidator(_coerce_legacy_resume)] | None = Field(
|
||||
None,
|
||||
description="Resume payload for continuing interrupted runs",
|
||||
)
|
||||
|
||||
|
||||
# region AG-UI Chat Options TypedDict
|
||||
|
||||
|
||||
class AGUIChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], total=False):
|
||||
"""AG-UI protocol-specific chat options dict.
|
||||
|
||||
Extends base ChatOptions for the AG-UI (Agent-UI) protocol.
|
||||
AG-UI is a streaming protocol for connecting AI agents to user interfaces.
|
||||
Options are forwarded to the remote AG-UI server.
|
||||
|
||||
See: https://github.com/ag-ui/ag-ui-protocol
|
||||
|
||||
Keys:
|
||||
# Inherited from ChatOptions (forwarded to remote server):
|
||||
model: The model identifier (forwarded as-is to server).
|
||||
temperature: Sampling temperature.
|
||||
top_p: Nucleus sampling parameter.
|
||||
max_tokens: Maximum tokens to generate.
|
||||
stop: Stop sequences.
|
||||
tools: List of tools - sent to server so LLM knows about client tools.
|
||||
Server executes its own tools; client tools execute locally via
|
||||
function invocation middleware.
|
||||
tool_choice: How the model should use tools.
|
||||
metadata: Metadata dict containing thread_id for conversation continuity.
|
||||
|
||||
# Options with limited support (depends on remote server):
|
||||
frequency_penalty: Forwarded if remote server supports it.
|
||||
presence_penalty: Forwarded if remote server supports it.
|
||||
seed: Forwarded if remote server supports it.
|
||||
response_format: Forwarded if remote server supports it.
|
||||
logit_bias: Forwarded if remote server supports it.
|
||||
user: Forwarded if remote server supports it.
|
||||
|
||||
# Options not typically used in AG-UI:
|
||||
store: Not applicable for AG-UI protocol.
|
||||
allow_multiple_tool_calls: Handled by underlying server.
|
||||
|
||||
# AG-UI-specific options:
|
||||
forward_props: Additional properties to forward to the AG-UI server.
|
||||
Useful for passing custom parameters to specific server implementations.
|
||||
context: Shared context/state to send to the server.
|
||||
|
||||
Note:
|
||||
AG-UI is a protocol bridge - actual option support depends on the
|
||||
remote server implementation. The client sends all options to the
|
||||
server, which decides how to handle them.
|
||||
|
||||
Thread ID management:
|
||||
- Pass ``thread_id`` in ``metadata`` to maintain conversation continuity
|
||||
- If not provided, a new thread ID is auto-generated
|
||||
"""
|
||||
|
||||
# AG-UI-specific options
|
||||
forward_props: dict[str, Any]
|
||||
"""Additional properties to forward to the AG-UI server."""
|
||||
|
||||
context: dict[str, Any]
|
||||
"""Shared context/state to send to the server."""
|
||||
|
||||
available_interrupts: list[Interrupt]
|
||||
"""Canonical AG-UI interrupt descriptors available for resumption."""
|
||||
|
||||
resume: list[ResumeEntry]
|
||||
"""Canonical AG-UI resume entries to continue a paused run."""
|
||||
|
||||
# ChatOptions fields not applicable for AG-UI
|
||||
store: None # type: ignore[misc]
|
||||
"""Not applicable for AG-UI protocol."""
|
||||
|
||||
|
||||
AGUI_OPTION_TRANSLATIONS: dict[str, str] = {}
|
||||
"""Maps ChatOptions keys to AG-UI parameter names (protocol uses standard names)."""
|
||||
|
||||
|
||||
# endregion
|
||||
@@ -0,0 +1,292 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Utility functions for AG-UI integration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable, MutableMapping, Sequence
|
||||
from dataclasses import asdict, is_dataclass
|
||||
from datetime import date, datetime
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentResponseUpdate, ChatResponseUpdate, FunctionTool
|
||||
|
||||
# Role mapping constants
|
||||
AGUI_TO_FRAMEWORK_ROLE: dict[str, str] = {
|
||||
"user": "user",
|
||||
"assistant": "assistant",
|
||||
"system": "system",
|
||||
}
|
||||
|
||||
FRAMEWORK_TO_AGUI_ROLE: dict[str, str] = {
|
||||
"user": "user",
|
||||
"assistant": "assistant",
|
||||
"system": "system",
|
||||
}
|
||||
|
||||
ALLOWED_AGUI_ROLES: set[str] = {"user", "assistant", "system", "tool", "reasoning"}
|
||||
|
||||
|
||||
def generate_event_id() -> str:
|
||||
"""Generate a unique event ID."""
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
def safe_json_parse(value: Any) -> dict[str, Any] | None:
|
||||
"""Safely parse a value as JSON dict.
|
||||
|
||||
Args:
|
||||
value: String or dict to parse
|
||||
|
||||
Returns:
|
||||
Parsed dict or None if parsing fails
|
||||
"""
|
||||
if isinstance(value, dict):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def canonical_function_arguments(function_call: Any) -> str | None:
|
||||
"""Return a stable representation of function-call arguments."""
|
||||
if function_call is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
parsed_arguments = function_call.parse_arguments()
|
||||
except Exception:
|
||||
parsed_arguments = getattr(function_call, "arguments", None)
|
||||
|
||||
if parsed_arguments is None:
|
||||
parsed_arguments = {}
|
||||
|
||||
return json.dumps(make_json_safe(parsed_arguments), sort_keys=True, separators=(",", ":"))
|
||||
|
||||
|
||||
def get_role_value(message: Any) -> str:
|
||||
"""Extract role string from a message object.
|
||||
|
||||
Handles both enum roles (with .value) and string roles.
|
||||
|
||||
Args:
|
||||
message: Message object with role attribute
|
||||
|
||||
Returns:
|
||||
Role as lowercase string, or empty string if not found
|
||||
"""
|
||||
role = getattr(message, "role", None)
|
||||
if role is None:
|
||||
return ""
|
||||
if hasattr(role, "value"):
|
||||
return str(role.value)
|
||||
return str(role)
|
||||
|
||||
|
||||
def normalize_agui_role(raw_role: Any) -> str:
|
||||
"""Normalize an AG-UI role to a standard role string.
|
||||
|
||||
Args:
|
||||
raw_role: Raw role value from AG-UI message
|
||||
|
||||
Returns:
|
||||
Normalized role string (user, assistant, system, tool, or reasoning)
|
||||
"""
|
||||
if not isinstance(raw_role, str):
|
||||
return "user"
|
||||
role = raw_role.lower()
|
||||
if role == "developer":
|
||||
return "system"
|
||||
if role in ALLOWED_AGUI_ROLES:
|
||||
return role
|
||||
return "user"
|
||||
|
||||
|
||||
def extract_state_from_tool_args(
|
||||
args: dict[str, Any] | None,
|
||||
tool_arg_name: str,
|
||||
) -> Any:
|
||||
"""Extract state value from tool arguments based on config.
|
||||
|
||||
Args:
|
||||
args: Parsed tool arguments dict
|
||||
tool_arg_name: Name of the argument to extract, or "*" for entire args
|
||||
|
||||
Returns:
|
||||
Extracted state value, or None if not found
|
||||
"""
|
||||
if not args:
|
||||
return None
|
||||
if tool_arg_name == "*":
|
||||
return args
|
||||
return args.get(tool_arg_name)
|
||||
|
||||
|
||||
def merge_state(current: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Merge state updates.
|
||||
|
||||
Args:
|
||||
current: Current state dictionary
|
||||
update: Update to apply
|
||||
|
||||
Returns:
|
||||
Merged state
|
||||
"""
|
||||
result = copy.deepcopy(current)
|
||||
result.update(update)
|
||||
return result
|
||||
|
||||
|
||||
def make_json_safe(obj: Any) -> Any: # noqa: ANN401
|
||||
"""Make an object JSON serializable.
|
||||
|
||||
Args:
|
||||
obj: Object to make JSON safe
|
||||
|
||||
Returns:
|
||||
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):
|
||||
# asdict may return nested non-dataclass objects, so recursively make them safe
|
||||
return make_json_safe(asdict(obj)) # type: ignore[arg-type]
|
||||
if hasattr(obj, "model_dump"):
|
||||
return make_json_safe(obj.model_dump())
|
||||
if hasattr(obj, "to_dict"):
|
||||
return make_json_safe(obj.to_dict())
|
||||
if hasattr(obj, "dict"):
|
||||
return make_json_safe(obj.dict())
|
||||
if hasattr(obj, "__dict__"):
|
||||
return {key: make_json_safe(value) for key, value in vars(obj).items()} # type: ignore[misc]
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [make_json_safe(item) for item in obj] # type: ignore[misc]
|
||||
if isinstance(obj, dict):
|
||||
return {key: make_json_safe(value) for key, value in obj.items()} # type: ignore[misc]
|
||||
return str(obj)
|
||||
|
||||
|
||||
def convert_agui_tools_to_agent_framework(
|
||||
agui_tools: list[dict[str, Any]] | None,
|
||||
) -> list[FunctionTool] | None:
|
||||
"""Convert AG-UI tool definitions to Agent Framework FunctionTool declarations.
|
||||
|
||||
Creates declaration-only FunctionTool instances (no executable implementation).
|
||||
These are used to tell the LLM about available tools. The actual execution
|
||||
happens on the client side via function invocation mixin.
|
||||
|
||||
CRITICAL: These tools MUST have func=None so that declaration_only returns True.
|
||||
This prevents the server from trying to execute client-side tools.
|
||||
|
||||
Args:
|
||||
agui_tools: List of AG-UI tool definitions with name, description, parameters
|
||||
|
||||
Returns:
|
||||
List of FunctionTool declarations, or None if no tools provided
|
||||
"""
|
||||
if not agui_tools:
|
||||
return None
|
||||
|
||||
result: list[FunctionTool] = []
|
||||
for tool_def in agui_tools:
|
||||
# Create declaration-only FunctionTool (func=None means no implementation)
|
||||
# When func=None, the declaration_only property returns True,
|
||||
# which tells the function invocation mixin to return the function call
|
||||
# without executing it (so it can be sent back to the client)
|
||||
func: FunctionTool = FunctionTool(
|
||||
name=tool_def.get("name", ""),
|
||||
description=tool_def.get("description", ""),
|
||||
func=None, # CRITICAL: Makes declaration_only=True
|
||||
input_model=tool_def.get("parameters", {}),
|
||||
)
|
||||
result.append(func)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def convert_tools_to_agui_format(
|
||||
tools: (
|
||||
FunctionTool
|
||||
| Callable[..., Any]
|
||||
| MutableMapping[str, Any]
|
||||
| Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]]
|
||||
| None
|
||||
),
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""Convert tools to AG-UI format.
|
||||
|
||||
This sends only the metadata (name, description, JSON schema) to the server.
|
||||
The actual executable implementation stays on the client side.
|
||||
The function invocation mixin handles client-side execution when
|
||||
the server requests a function.
|
||||
|
||||
Args:
|
||||
tools: Tools to convert (single tool or sequence of tools)
|
||||
|
||||
Returns:
|
||||
List of tool specifications in AG-UI format, or None if no tools provided
|
||||
"""
|
||||
if not tools:
|
||||
return None
|
||||
|
||||
# Normalize to list
|
||||
if not isinstance(tools, list):
|
||||
tool_list: list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] = [tools] # type: ignore[list-item]
|
||||
else:
|
||||
tool_list = tools # type: ignore[assignment]
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
|
||||
for tool_item in tool_list:
|
||||
if isinstance(tool_item, dict):
|
||||
# Already in dict format, pass through
|
||||
results.append(tool_item) # type: ignore[arg-type]
|
||||
elif isinstance(tool_item, FunctionTool):
|
||||
# Convert FunctionTool to AG-UI tool format
|
||||
results.append(
|
||||
{
|
||||
"name": tool_item.name,
|
||||
"description": tool_item.description,
|
||||
"parameters": tool_item.parameters(),
|
||||
}
|
||||
)
|
||||
elif callable(tool_item):
|
||||
# Convert callable to FunctionTool first, then to AG-UI format
|
||||
from agent_framework import tool
|
||||
|
||||
ai_func = tool(tool_item)
|
||||
results.append(
|
||||
{
|
||||
"name": ai_func.name,
|
||||
"description": ai_func.description,
|
||||
"parameters": ai_func.parameters(),
|
||||
}
|
||||
)
|
||||
# Note: dict-based hosted tools (CodeInterpreter, WebSearch, etc.) are passed through
|
||||
# as-is in the first branch. Non-FunctionTool, non-dict items are skipped.
|
||||
|
||||
return results if results else None
|
||||
|
||||
|
||||
def get_conversation_id_from_update(update: AgentResponseUpdate) -> str | None:
|
||||
"""Extract conversation ID from AgentResponseUpdate metadata.
|
||||
|
||||
Args:
|
||||
update: AgentRunResponseUpdate instance
|
||||
Returns:
|
||||
Conversation ID if present, else None
|
||||
|
||||
"""
|
||||
if isinstance(update.raw_representation, ChatResponseUpdate):
|
||||
return update.raw_representation.conversation_id
|
||||
return None
|
||||
@@ -0,0 +1,404 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Workflow wrapper for AG-UI protocol compatibility."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import AsyncGenerator, Callable
|
||||
from typing import Any, cast
|
||||
|
||||
from ag_ui.core import (
|
||||
BaseEvent,
|
||||
MessagesSnapshotEvent,
|
||||
RunErrorEvent,
|
||||
RunFinishedEvent,
|
||||
RunStartedEvent,
|
||||
StateSnapshotEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
TextMessageStartEvent,
|
||||
ToolCallArgsEvent,
|
||||
ToolCallResultEvent,
|
||||
ToolCallStartEvent,
|
||||
)
|
||||
from agent_framework import Workflow
|
||||
|
||||
from ._message_adapters import agui_messages_to_snapshot_format
|
||||
from ._run_common import (
|
||||
_build_run_finished_event,
|
||||
_extract_resume_payload,
|
||||
_normalize_resume_interrupts,
|
||||
_reconstruct_messages_from_thread_snapshot,
|
||||
)
|
||||
from ._snapshots import (
|
||||
_DEFAULT_STATE_INPUT_KEY,
|
||||
_SNAPSHOT_SCOPE_INPUT_KEY,
|
||||
AGUIThreadSnapshot,
|
||||
AGUIThreadSnapshotStore,
|
||||
_clear_thread_snapshot_interrupt,
|
||||
)
|
||||
from ._utils import generate_event_id, make_json_safe
|
||||
from ._workflow_run import run_workflow_stream
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
WorkflowFactory = Callable[[str], Workflow]
|
||||
|
||||
|
||||
def _cancelled_resume_interrupt_ids(resume_payload: Any) -> set[str]:
|
||||
"""Return cancelled interrupt ids from a resume payload."""
|
||||
return {
|
||||
str(interrupt["id"])
|
||||
for interrupt in _normalize_resume_interrupts(resume_payload)
|
||||
if interrupt.get("status") == "cancelled"
|
||||
}
|
||||
|
||||
|
||||
def _event_messages_to_snapshot_dicts(messages: list[Any]) -> list[dict[str, Any]]:
|
||||
"""Convert AG-UI message event models to plain snapshot dictionaries."""
|
||||
safe_messages = make_json_safe(messages)
|
||||
if not isinstance(safe_messages, list):
|
||||
return []
|
||||
return [cast(dict[str, Any], message) for message in safe_messages if isinstance(message, dict)]
|
||||
|
||||
|
||||
class _WorkflowSnapshotBuilder:
|
||||
"""Capture replayable workflow protocol output without retaining raw events."""
|
||||
|
||||
def __init__(self, raw_messages: list[dict[str, Any]]) -> None:
|
||||
self._synthesized_messages = agui_messages_to_snapshot_format(raw_messages)
|
||||
self._emitted_messages: list[dict[str, Any]] | None = None
|
||||
self._open_text_message: dict[str, Any] | None = None
|
||||
self._tool_call_message: dict[str, Any] | None = None
|
||||
self._tool_calls_by_id: dict[str, dict[str, Any]] = {}
|
||||
self.state: dict[str, Any] | None = None
|
||||
self.interrupt: list[dict[str, Any]] | None = None
|
||||
|
||||
def observe(self, event: BaseEvent) -> None:
|
||||
"""Fold one replayable AG-UI event into the latest snapshot state."""
|
||||
if isinstance(event, StateSnapshotEvent):
|
||||
state = make_json_safe(event.snapshot)
|
||||
if isinstance(state, dict):
|
||||
self.state = cast(dict[str, Any], state)
|
||||
return
|
||||
|
||||
if isinstance(event, MessagesSnapshotEvent):
|
||||
self._emitted_messages = _event_messages_to_snapshot_dicts(list(event.messages))
|
||||
return
|
||||
|
||||
if isinstance(event, RunFinishedEvent):
|
||||
outcome = getattr(event, "outcome", None)
|
||||
interrupt = (
|
||||
make_json_safe(getattr(outcome, "interrupts", None))
|
||||
if getattr(outcome, "type", None) == "interrupt"
|
||||
else None
|
||||
)
|
||||
if isinstance(interrupt, list):
|
||||
self.interrupt = [cast(dict[str, Any], item) for item in interrupt if isinstance(item, dict)]
|
||||
return
|
||||
|
||||
if self._emitted_messages is not None:
|
||||
return
|
||||
|
||||
if isinstance(event, TextMessageStartEvent):
|
||||
self._observe_text_start(event)
|
||||
elif isinstance(event, TextMessageContentEvent):
|
||||
self._observe_text_content(event)
|
||||
elif isinstance(event, TextMessageEndEvent):
|
||||
self._observe_text_end(event)
|
||||
elif isinstance(event, ToolCallStartEvent):
|
||||
self._observe_tool_call_start(event)
|
||||
elif isinstance(event, ToolCallArgsEvent):
|
||||
self._observe_tool_call_args(event)
|
||||
elif isinstance(event, ToolCallResultEvent):
|
||||
self._observe_tool_call_result(event)
|
||||
|
||||
def build(self) -> AGUIThreadSnapshot:
|
||||
"""Return the replayable thread snapshot."""
|
||||
self._flush_open_text_message()
|
||||
messages = self._emitted_messages if self._emitted_messages is not None else self._synthesized_messages
|
||||
return AGUIThreadSnapshot(messages=messages, state=self.state, interrupt=self.interrupt)
|
||||
|
||||
def _observe_text_start(self, event: TextMessageStartEvent) -> None:
|
||||
if self._open_text_message is not None and self._open_text_message.get("id") != event.message_id:
|
||||
self._flush_open_text_message()
|
||||
self._open_text_message = {"id": event.message_id, "role": event.role, "content": ""}
|
||||
|
||||
def _observe_text_content(self, event: TextMessageContentEvent) -> None:
|
||||
if self._open_text_message is None or self._open_text_message.get("id") != event.message_id:
|
||||
self._open_text_message = {"id": event.message_id, "role": "assistant", "content": ""}
|
||||
self._open_text_message["content"] = f"{self._open_text_message.get('content', '')}{event.delta}"
|
||||
|
||||
def _observe_text_end(self, event: TextMessageEndEvent) -> None:
|
||||
if self._open_text_message is None or self._open_text_message.get("id") != event.message_id:
|
||||
return
|
||||
self._flush_open_text_message()
|
||||
|
||||
def _observe_tool_call_start(self, event: ToolCallStartEvent) -> None:
|
||||
parent_message_id = event.parent_message_id
|
||||
if (
|
||||
self._open_text_message is not None
|
||||
and parent_message_id is not None
|
||||
and self._open_text_message.get("id") == parent_message_id
|
||||
and self._open_text_message.get("content")
|
||||
):
|
||||
self._open_text_message["id"] = generate_event_id()
|
||||
self._flush_open_text_message()
|
||||
if self._tool_call_message is None or (
|
||||
parent_message_id is not None and self._tool_call_message.get("id") != parent_message_id
|
||||
):
|
||||
self._tool_call_message = {
|
||||
"id": parent_message_id or generate_event_id(),
|
||||
"role": "assistant",
|
||||
"tool_calls": [],
|
||||
}
|
||||
self._synthesized_messages.append(self._tool_call_message)
|
||||
|
||||
tool_call = {
|
||||
"id": event.tool_call_id,
|
||||
"type": "function",
|
||||
"function": {"name": event.tool_call_name, "arguments": ""},
|
||||
}
|
||||
cast(list[dict[str, Any]], self._tool_call_message["tool_calls"]).append(tool_call)
|
||||
self._tool_calls_by_id[event.tool_call_id] = tool_call
|
||||
|
||||
def _observe_tool_call_args(self, event: ToolCallArgsEvent) -> None:
|
||||
tool_call = self._tool_calls_by_id.get(event.tool_call_id)
|
||||
if tool_call is None:
|
||||
return
|
||||
function_payload = cast(dict[str, Any], tool_call["function"])
|
||||
function_payload["arguments"] = f"{function_payload.get('arguments', '')}{event.delta}"
|
||||
|
||||
def _observe_tool_call_result(self, event: ToolCallResultEvent) -> None:
|
||||
self._synthesized_messages.append(
|
||||
{
|
||||
"id": event.message_id,
|
||||
"role": "tool",
|
||||
"toolCallId": event.tool_call_id,
|
||||
"content": event.content,
|
||||
}
|
||||
)
|
||||
# A result closes the current tool-call group; later tool calls start a new
|
||||
# assistant message so replayed transcripts keep results adjacent to their
|
||||
# tool_calls message, which provider APIs require.
|
||||
self._tool_call_message = None
|
||||
|
||||
def _flush_open_text_message(self) -> None:
|
||||
if self._open_text_message is None:
|
||||
return
|
||||
if self._open_text_message.get("content"):
|
||||
self._synthesized_messages.append(self._open_text_message)
|
||||
# Text between tool calls closes the current tool-call group as well.
|
||||
self._tool_call_message = None
|
||||
self._open_text_message = None
|
||||
|
||||
|
||||
async def _hydrate_workflow_thread_snapshot(
|
||||
*,
|
||||
snapshot_store: AGUIThreadSnapshotStore,
|
||||
scope: str,
|
||||
thread_id: str,
|
||||
run_id: str,
|
||||
) -> AsyncGenerator[BaseEvent]:
|
||||
"""Replay the latest stored workflow AG-UI Thread Snapshot without invoking the workflow."""
|
||||
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
|
||||
snapshot = await snapshot_store.get(scope=scope, thread_id=thread_id)
|
||||
if snapshot is None:
|
||||
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id)
|
||||
return
|
||||
|
||||
if snapshot.state is not None:
|
||||
yield StateSnapshotEvent(snapshot=snapshot.state)
|
||||
if snapshot.messages:
|
||||
yield MessagesSnapshotEvent(messages=snapshot.messages) # type: ignore[arg-type]
|
||||
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=snapshot.interrupt)
|
||||
|
||||
|
||||
class AgentFrameworkWorkflow:
|
||||
"""Base AG-UI workflow wrapper.
|
||||
|
||||
Can wrap a native ``Workflow`` or be subclassed for custom ``run`` behavior.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workflow: Workflow | None = None,
|
||||
*,
|
||||
workflow_factory: WorkflowFactory | None = None,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
snapshot_store: AGUIThreadSnapshotStore | None = None,
|
||||
) -> None:
|
||||
"""Initialize the AG-UI workflow wrapper.
|
||||
|
||||
Args:
|
||||
workflow: Optional workflow instance to expose.
|
||||
workflow_factory: Optional factory for thread-scoped workflow instances.
|
||||
name: Optional workflow name.
|
||||
description: Optional workflow description.
|
||||
snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless
|
||||
endpoint setup also provides an explicit Snapshot Scope resolver.
|
||||
"""
|
||||
if workflow is not None and workflow_factory is not None:
|
||||
raise ValueError("Pass either workflow= or workflow_factory=, not both.")
|
||||
|
||||
self.workflow = workflow
|
||||
self._workflow_factory = workflow_factory
|
||||
# Cache keyed by (snapshot_scope, thread_id): the Snapshot Scope is the
|
||||
# authorization boundary, so the same thread id under different scopes
|
||||
# must never share an in-memory workflow instance.
|
||||
self._workflow_by_thread: dict[tuple[str | None, str], Workflow] = {}
|
||||
self.name = name if name is not None else getattr(workflow, "name", "workflow")
|
||||
self.description = description if description is not None else getattr(workflow, "description", "")
|
||||
self.snapshot_store = snapshot_store
|
||||
|
||||
@staticmethod
|
||||
def _thread_id_from_input(input_data: dict[str, Any]) -> str:
|
||||
"""Resolve a stable thread id from AG-UI input payload."""
|
||||
thread_id = input_data.get("thread_id") or input_data.get("threadId")
|
||||
if thread_id is not None:
|
||||
return str(thread_id)
|
||||
return str(uuid.uuid4())
|
||||
|
||||
def _resolve_workflow(self, thread_id: str, snapshot_scope: str | None = None) -> Workflow:
|
||||
"""Get the workflow instance for the current run."""
|
||||
if self.workflow is not None:
|
||||
return self.workflow
|
||||
|
||||
if self._workflow_factory is None:
|
||||
raise NotImplementedError("No workflow is attached. Override run or pass workflow=/workflow_factory=.")
|
||||
|
||||
cache_key = (snapshot_scope, thread_id)
|
||||
workflow = self._workflow_by_thread.get(cache_key)
|
||||
if workflow is None:
|
||||
workflow = self._workflow_factory(thread_id)
|
||||
if not isinstance(workflow, Workflow):
|
||||
raise TypeError("workflow_factory must return a Workflow instance.")
|
||||
self._workflow_by_thread[cache_key] = workflow
|
||||
return workflow
|
||||
|
||||
def clear_thread_workflow(self, thread_id: str, snapshot_scope: str | None = None) -> None:
|
||||
"""Drop cached workflow instances for a thread, optionally limited to one Snapshot Scope."""
|
||||
if snapshot_scope is not None:
|
||||
self._workflow_by_thread.pop((snapshot_scope, thread_id), None)
|
||||
return
|
||||
for key in [key for key in self._workflow_by_thread if key[1] == thread_id]:
|
||||
del self._workflow_by_thread[key]
|
||||
|
||||
def clear_workflow_cache(self) -> None:
|
||||
"""Drop all cached thread workflow instances."""
|
||||
self._workflow_by_thread.clear()
|
||||
|
||||
async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]:
|
||||
"""Run the wrapped workflow and yield AG-UI events.
|
||||
|
||||
Subclasses may override this to provide custom AG-UI streams.
|
||||
"""
|
||||
thread_id = self._thread_id_from_input(input_data)
|
||||
run_id = str(input_data.get("run_id") or input_data.get("runId") or uuid.uuid4())
|
||||
snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY))
|
||||
raw_messages = list(cast(list[dict[str, Any]], input_data.get("messages", []) or []))
|
||||
resume_payload = _extract_resume_payload(input_data)
|
||||
snapshot_store = self.snapshot_store
|
||||
|
||||
if snapshot_store is not None and snapshot_scope is not None and not raw_messages and resume_payload is None:
|
||||
async for event in _hydrate_workflow_thread_snapshot(
|
||||
snapshot_store=snapshot_store,
|
||||
scope=snapshot_scope,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
):
|
||||
yield event
|
||||
return
|
||||
|
||||
# Load the stored snapshot for follow-up turns so the workflow runs with the
|
||||
# full persisted thread history instead of just the latest request messages.
|
||||
stored_snapshot: AGUIThreadSnapshot | None = None
|
||||
if snapshot_store is not None and snapshot_scope is not None:
|
||||
stored_snapshot = await snapshot_store.get(scope=snapshot_scope, thread_id=thread_id)
|
||||
if stored_snapshot is not None and resume_payload is None:
|
||||
raw_messages = _reconstruct_messages_from_thread_snapshot(
|
||||
stored_messages=stored_snapshot.messages,
|
||||
incoming_messages=raw_messages,
|
||||
stored_interrupt=stored_snapshot.interrupt,
|
||||
)
|
||||
input_data["messages"] = raw_messages
|
||||
|
||||
# Merge stored state with request overrides, then fill endpoint-deferred
|
||||
# defaults only for keys missing from both.
|
||||
request_state = input_data.get("state")
|
||||
deferred_default_state = cast(dict[str, Any] | None, input_data.get(_DEFAULT_STATE_INPUT_KEY))
|
||||
effective_state: dict[str, Any] = {}
|
||||
if stored_snapshot is not None and stored_snapshot.state is not None:
|
||||
effective_state.update(stored_snapshot.state)
|
||||
if isinstance(request_state, dict):
|
||||
effective_state.update(cast(dict[str, Any], request_state))
|
||||
if deferred_default_state:
|
||||
for key, value in deferred_default_state.items():
|
||||
if key not in effective_state:
|
||||
effective_state[key] = copy.deepcopy(value)
|
||||
if effective_state:
|
||||
input_data["state"] = effective_state
|
||||
|
||||
workflow = self._resolve_workflow(thread_id, snapshot_scope)
|
||||
builder_seed_messages = raw_messages
|
||||
if resume_payload is not None and stored_snapshot is not None:
|
||||
# Resume requests carry only the synthesized interrupt response, so seed
|
||||
# the builder with stored history to avoid persisting a truncated thread.
|
||||
builder_seed_messages = [
|
||||
copy.deepcopy(message) for message in stored_snapshot.messages
|
||||
] + builder_seed_messages
|
||||
snapshot_builder = (
|
||||
_WorkflowSnapshotBuilder(builder_seed_messages)
|
||||
if snapshot_store is not None and snapshot_scope is not None
|
||||
else None
|
||||
)
|
||||
if snapshot_builder is not None and effective_state:
|
||||
# Seed builder state so a run that emits no StateSnapshotEvent still
|
||||
# persists the latest known Shared State instead of dropping it.
|
||||
state_snapshot = make_json_safe(effective_state)
|
||||
if isinstance(state_snapshot, dict):
|
||||
snapshot_builder.state = cast(dict[str, Any], state_snapshot)
|
||||
run_error_emitted = False
|
||||
async for event in run_workflow_stream(input_data, workflow):
|
||||
if snapshot_builder is not None:
|
||||
snapshot_builder.observe(event)
|
||||
if isinstance(event, RunErrorEvent):
|
||||
run_error_emitted = True
|
||||
if (
|
||||
getattr(event, "code", None) == "WORKFLOW_RESUME_CANCELLED"
|
||||
and snapshot_store is not None
|
||||
and snapshot_scope is not None
|
||||
):
|
||||
await _clear_thread_snapshot_interrupt(
|
||||
snapshot_store=snapshot_store,
|
||||
scope=snapshot_scope,
|
||||
thread_id=thread_id,
|
||||
interrupt_ids=_cancelled_resume_interrupt_ids(resume_payload),
|
||||
)
|
||||
yield event
|
||||
|
||||
if (
|
||||
snapshot_builder is not None
|
||||
and not run_error_emitted
|
||||
and snapshot_store is not None
|
||||
and snapshot_scope is not None
|
||||
):
|
||||
try:
|
||||
await snapshot_store.save(
|
||||
scope=snapshot_scope,
|
||||
thread_id=thread_id,
|
||||
snapshot=snapshot_builder.build(),
|
||||
)
|
||||
except Exception:
|
||||
# RUN_FINISHED has already been yielded; a store failure must not
|
||||
# surface as a second terminal RUN_ERROR event. The previous
|
||||
# snapshot stays available for hydration.
|
||||
logger.exception(
|
||||
"Failed to save AG-UI Thread Snapshot for scope=%s thread_id=%s; keeping previous snapshot.",
|
||||
snapshot_scope,
|
||||
thread_id,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
# Marker file for PEP 561
|
||||
@@ -0,0 +1,3 @@
|
||||
AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/
|
||||
AZURE_OPENAI_API_KEY=your-api-key-here
|
||||
PORT=8000
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"python.analysis.extraPaths": [
|
||||
"${workspaceFolder}/packages/ag-ui/examples"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,360 @@
|
||||
# Agent Framework AG-UI Integration
|
||||
|
||||
AG-UI protocol integration for Agent Framework, enabling seamless integration with AG-UI's web interface and streaming protocol.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install agent-framework-ag-ui
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Using Example Agents with Any Chat Client
|
||||
|
||||
All example agents are factory functions that accept any `SupportsChatGetResponse`-compatible chat client:
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
from agent_framework_ag_ui_examples.agents import simple_agent, weather_agent
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
# Option 1: Use Azure OpenAI
|
||||
azure_client = OpenAIChatCompletionClient(model="gpt-4")
|
||||
add_agent_framework_fastapi_endpoint(app, simple_agent(azure_client), "/chat")
|
||||
|
||||
# Option 2: Use OpenAI
|
||||
openai_client = OpenAIChatClient(model="gpt-4o")
|
||||
add_agent_framework_fastapi_endpoint(app, weather_agent(openai_client), "/weather")
|
||||
|
||||
# Run with: uvicorn main:app --reload
|
||||
```
|
||||
|
||||
### Creating Your Own Agent
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
|
||||
# Create your agent
|
||||
agent = Agent(
|
||||
name="my_agent",
|
||||
instructions="You are a helpful assistant.",
|
||||
client=OpenAIChatCompletionClient(model="gpt-4o"),
|
||||
)
|
||||
|
||||
# Create FastAPI app and add AG-UI endpoint
|
||||
app = FastAPI()
|
||||
add_agent_framework_fastapi_endpoint(app, agent, "/agent")
|
||||
|
||||
# Run with: uvicorn main:app --reload
|
||||
```
|
||||
|
||||
## Features
|
||||
|
||||
This integration supports all 7 AG-UI features:
|
||||
|
||||
1. **Agentic Chat**: Basic streaming chat with tool calling support
|
||||
2. **Backend Tool Rendering**: Tools executed on backend with results streamed via ToolCallResultEvent
|
||||
3. **Human in the Loop**: Function approval requests for user confirmation before tool execution
|
||||
4. **Agentic Generative UI**: Async tools for long-running operations with progress updates
|
||||
5. **Tool-based Generative UI**: Custom UI components rendered on frontend based on tool calls
|
||||
6. **Shared State**: Bidirectional state sync using StateSnapshotEvent and StateDeltaEvent
|
||||
7. **Predictive State Updates**: Stream tool arguments as optimistic state updates during execution
|
||||
|
||||
## Examples
|
||||
|
||||
All example agents are implemented as **factory functions** that accept any chat client implementing `SupportsChatGetResponse`. This provides maximum flexibility to use Azure OpenAI, OpenAI, Anthropic, or any custom chat client implementation.
|
||||
|
||||
### Available Example Agents
|
||||
|
||||
Complete examples for all AG-UI features are available:
|
||||
|
||||
- `simple_agent(client)` - Basic agentic chat (Feature 1)
|
||||
- `weather_agent(client)` - Backend tool rendering (Feature 2)
|
||||
- `human_in_the_loop_agent(client)` - Human-in-the-loop with step customization (Feature 3)
|
||||
- `task_steps_agent_wrapped(client)` - Agentic generative UI with step execution (Feature 4)
|
||||
- `ui_generator_agent(client)` - Tool-based generative UI (Feature 5)
|
||||
- `recipe_agent(client)` - Shared state management (Feature 6)
|
||||
- `document_writer_agent(client)` - Predictive state updates (Feature 7)
|
||||
- `research_assistant_agent(client)` - Research with progress events
|
||||
- `task_planner_agent(client)` - Task planning with approvals
|
||||
- `subgraphs_agent()` - Deterministic travel-planning subgraphs flow (Dojo `subgraphs` feature)
|
||||
|
||||
### Using Example Agents
|
||||
|
||||
```python
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
from agent_framework_ag_ui_examples.agents import (
|
||||
simple_agent,
|
||||
weather_agent,
|
||||
recipe_agent,
|
||||
)
|
||||
|
||||
# Create a chat client (use any SupportsChatGetResponse implementation)
|
||||
azure_client = OpenAIChatCompletionClient(model="gpt-4")
|
||||
openai_client = OpenAIChatClient(model="gpt-4o")
|
||||
|
||||
# Create agent instances by calling the factory functions
|
||||
agent1 = simple_agent(azure_client)
|
||||
agent2 = weather_agent(openai_client)
|
||||
agent3 = recipe_agent(azure_client)
|
||||
```
|
||||
|
||||
### Running the Example Server
|
||||
|
||||
The example server demonstrates all 7 AG-UI features:
|
||||
|
||||
```bash
|
||||
# Install the package
|
||||
pip install agent-framework-ag-ui
|
||||
|
||||
# Run the example server
|
||||
python -m agent_framework_ag_ui_examples
|
||||
|
||||
# Or with debug logging
|
||||
ENABLE_DEBUG_LOGGING=1 python -m agent_framework_ag_ui_examples
|
||||
```
|
||||
|
||||
The server exposes endpoints at:
|
||||
- `/agentic_chat` - Simple chat with `simple_agent`
|
||||
- `/backend_tool_rendering` - Weather tools with `weather_agent`
|
||||
- `/human_in_the_loop` - Step approval with `human_in_the_loop_agent`
|
||||
- `/agentic_generative_ui` - Task steps with `task_steps_agent_wrapped`
|
||||
- `/tool_based_generative_ui` - Custom UI components with `ui_generator_agent`
|
||||
- `/shared_state` - Recipe management with `recipe_agent`
|
||||
- `/predictive_state_updates` - Document writing with `document_writer_agent`
|
||||
- `/subgraphs` - Travel planner with interrupt-driven flight/hotel choices via `subgraphs_agent`
|
||||
|
||||
### Interrupt and Resume Shape
|
||||
|
||||
Human-in-the-loop and workflow examples use the canonical AG-UI protocol shape. A paused run finishes with
|
||||
`RUN_FINISHED.outcome.type == "interrupt"` and renders prompts from `RUN_FINISHED.outcome.interrupts`; it does not
|
||||
depend on a stable top-level `RUN_FINISHED.interrupt` field.
|
||||
|
||||
Resume interrupted example threads with a canonical `resume` array:
|
||||
|
||||
```json
|
||||
{
|
||||
"threadId": "thread-1",
|
||||
"messages": [],
|
||||
"resume": [
|
||||
{
|
||||
"interruptId": "interrupt_1",
|
||||
"status": "resolved",
|
||||
"payload": {
|
||||
"approved": true
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### Complete FastAPI Example
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
from agent_framework_ag_ui_examples.agents import (
|
||||
simple_agent,
|
||||
weather_agent,
|
||||
human_in_the_loop_agent,
|
||||
task_steps_agent_wrapped,
|
||||
ui_generator_agent,
|
||||
recipe_agent,
|
||||
document_writer_agent,
|
||||
subgraphs_agent,
|
||||
)
|
||||
|
||||
app = FastAPI(title="AG-UI Examples")
|
||||
|
||||
# Create a chat client (shared across all agents, or create individual ones)
|
||||
client = OpenAIChatCompletionClient(model="gpt-4")
|
||||
|
||||
# Add all example endpoints
|
||||
add_agent_framework_fastapi_endpoint(app, simple_agent(client), "/agentic_chat")
|
||||
add_agent_framework_fastapi_endpoint(app, weather_agent(client), "/backend_tool_rendering")
|
||||
add_agent_framework_fastapi_endpoint(app, human_in_the_loop_agent(client), "/human_in_the_loop")
|
||||
add_agent_framework_fastapi_endpoint(app, task_steps_agent_wrapped(client), "/agentic_generative_ui") # type: ignore[arg-type]
|
||||
add_agent_framework_fastapi_endpoint(app, ui_generator_agent(client), "/tool_based_generative_ui")
|
||||
add_agent_framework_fastapi_endpoint(app, recipe_agent(client), "/shared_state")
|
||||
add_agent_framework_fastapi_endpoint(app, document_writer_agent(client), "/predictive_state_updates")
|
||||
add_agent_framework_fastapi_endpoint(app, subgraphs_agent(), "/subgraphs")
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
The package uses a clean, orchestrator-based architecture:
|
||||
|
||||
- **AgentFrameworkAgent**: Lightweight wrapper that delegates to orchestrators
|
||||
- **Orchestrators**: Handle different execution flows (default, human-in-the-loop, etc.)
|
||||
- **Confirmation Strategies**: Domain-specific confirmation messages (extensible)
|
||||
- **AgentFrameworkEventBridge**: Converts AgentResponseUpdate to AG-UI events
|
||||
- **Message Adapters**: Bidirectional conversion between AG-UI and Agent Framework message formats
|
||||
- **FastAPI Endpoint**: Streaming HTTP endpoint with Server-Sent Events (SSE)
|
||||
|
||||
### Key Design Patterns
|
||||
|
||||
- **Orchestrator Pattern**: Separates flow control from protocol translation
|
||||
- **Strategy Pattern**: Pluggable confirmation message strategies
|
||||
- **Context Object**: Lazy-loaded execution context passed to orchestrators
|
||||
- **Event Bridge**: Stateless translation of Agent Framework events to AG-UI events
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
### Creating Custom Agent Factories
|
||||
|
||||
You can create your own agent factories following the same pattern as the examples:
|
||||
|
||||
```python
|
||||
from agent_framework import Agent, tool
|
||||
from agent_framework import SupportsChatGetResponse
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
@tool
|
||||
def my_tool(param: str) -> str:
|
||||
"""My custom tool."""
|
||||
return f"Result: {param}"
|
||||
|
||||
def my_custom_agent(client: SupportsChatGetResponse) -> AgentFrameworkAgent:
|
||||
"""Create a custom agent with the specified chat client.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured AgentFrameworkAgent instance
|
||||
"""
|
||||
agent = Agent(
|
||||
name="my_custom_agent",
|
||||
instructions="Custom instructions here",
|
||||
client=client,
|
||||
tools=[my_tool],
|
||||
)
|
||||
|
||||
return AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
name="MyCustomAgent",
|
||||
description="My custom agent description",
|
||||
)
|
||||
|
||||
# Use it
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
client = OpenAIChatCompletionClient()
|
||||
agent = my_custom_agent(client)
|
||||
```
|
||||
|
||||
### Shared State
|
||||
|
||||
State is injected as system messages and updated via predictive state updates:
|
||||
|
||||
```python
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
# Create your agent
|
||||
agent = Agent(
|
||||
name="recipe_agent",
|
||||
client=OpenAIChatCompletionClient(model="gpt-4o"),
|
||||
)
|
||||
|
||||
state_schema = {
|
||||
"recipe": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"ingredients": {"type": "array"}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Configure which tool updates which state fields
|
||||
predict_state_config = {
|
||||
"recipe": {"tool": "update_recipe", "tool_argument": "recipe_data"}
|
||||
}
|
||||
|
||||
wrapped_agent = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema=state_schema,
|
||||
predict_state_config=predict_state_config,
|
||||
)
|
||||
```
|
||||
|
||||
### Predictive State Updates
|
||||
|
||||
Predictive state updates automatically stream tool arguments as optimistic state updates:
|
||||
|
||||
```python
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
# Create your agent
|
||||
agent = Agent(
|
||||
name="document_writer",
|
||||
client=OpenAIChatCompletionClient(model="gpt-4o"),
|
||||
)
|
||||
|
||||
predict_state_config = {
|
||||
"current_title": {"tool": "write_document", "tool_argument": "title"},
|
||||
"current_content": {"tool": "write_document", "tool_argument": "content"},
|
||||
}
|
||||
|
||||
wrapped_agent = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema={"current_title": {"type": "string"}, "current_content": {"type": "string"}},
|
||||
predict_state_config=predict_state_config,
|
||||
require_confirmation=True, # User can approve/reject changes
|
||||
)
|
||||
```
|
||||
|
||||
### Human in the Loop
|
||||
|
||||
Human-in-the-loop is automatically handled when tools are marked for approval:
|
||||
|
||||
```python
|
||||
from agent_framework import tool
|
||||
|
||||
@tool(approval_mode="always_require")
|
||||
def sensitive_action(param: str) -> str:
|
||||
"""This action requires user approval."""
|
||||
return f"Executed with {param}"
|
||||
|
||||
# The orchestrator automatically detects approval responses and handles them
|
||||
```
|
||||
|
||||
### Custom Orchestrators
|
||||
|
||||
Add custom execution flows by implementing the Orchestrator pattern:
|
||||
|
||||
```python
|
||||
from agent_framework.ag_ui._orchestrators import Orchestrator, ExecutionContext
|
||||
|
||||
class MyCustomOrchestrator(Orchestrator):
|
||||
def can_handle(self, context: ExecutionContext) -> bool:
|
||||
# Return True if this orchestrator should handle the request
|
||||
return context.input_data.get("custom_mode") == True
|
||||
|
||||
async def run(self, context: ExecutionContext):
|
||||
# Custom execution logic
|
||||
yield RunStartedEvent(...)
|
||||
# ... your custom flow
|
||||
yield RunFinishedEvent(...)
|
||||
|
||||
wrapped_agent = AgentFrameworkAgent(
|
||||
agent=your_agent,
|
||||
orchestrators=[MyCustomOrchestrator(), DefaultOrchestrator()],
|
||||
)
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Example agents for AG-UI demonstration."""
|
||||
|
||||
from . import agents
|
||||
|
||||
__all__ = ["agents"]
|
||||
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Entry point for running the AG-UI examples server as a module."""
|
||||
|
||||
from .server.main import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,27 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Example agents for AG-UI demonstration."""
|
||||
|
||||
from .document_writer_agent import document_writer_agent
|
||||
from .human_in_the_loop_agent import human_in_the_loop_agent
|
||||
from .recipe_agent import recipe_agent
|
||||
from .research_assistant_agent import research_assistant_agent
|
||||
from .simple_agent import simple_agent
|
||||
from .subgraphs_agent import subgraphs_agent
|
||||
from .task_planner_agent import task_planner_agent
|
||||
from .task_steps_agent import task_steps_agent_wrapped
|
||||
from .ui_generator_agent import ui_generator_agent
|
||||
from .weather_agent import weather_agent
|
||||
|
||||
__all__ = [
|
||||
"document_writer_agent",
|
||||
"human_in_the_loop_agent",
|
||||
"recipe_agent",
|
||||
"research_assistant_agent",
|
||||
"simple_agent",
|
||||
"subgraphs_agent",
|
||||
"task_planner_agent",
|
||||
"task_steps_agent_wrapped",
|
||||
"ui_generator_agent",
|
||||
"weather_agent",
|
||||
]
|
||||
@@ -0,0 +1,69 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Example agent demonstrating predictive state updates with document writing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from agent_framework import Agent, SupportsChatGetResponse, tool
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
|
||||
@tool(approval_mode="always_require")
|
||||
def write_document(document: str) -> str:
|
||||
"""Write a document. Use markdown formatting to format the document.
|
||||
|
||||
It's good to format the document extensively so it's easy to read.
|
||||
You can use all kinds of markdown.
|
||||
However, do not use italic or strike-through formatting, it's reserved for another purpose.
|
||||
You MUST write the full document, even when changing only a few words.
|
||||
When making edits to the document, try to make them minimal - do not change every word.
|
||||
Keep stories SHORT!
|
||||
|
||||
Args:
|
||||
document: The complete document content in markdown format
|
||||
|
||||
Returns:
|
||||
Confirmation that the document was written
|
||||
"""
|
||||
return "Document written."
|
||||
|
||||
|
||||
_DOCUMENT_WRITER_INSTRUCTIONS = (
|
||||
"You are a helpful assistant for writing documents. "
|
||||
"To write the document, you MUST use the write_document tool. "
|
||||
"You MUST write the full document, even when changing only a few words. "
|
||||
"When you wrote the document, DO NOT repeat it as a message. "
|
||||
"Just briefly summarize the changes you made. 2 sentences max. "
|
||||
"\n\n"
|
||||
"The current state of the document will be provided to you. "
|
||||
"When editing, make minimal changes - do not change every word unless requested."
|
||||
)
|
||||
|
||||
|
||||
def document_writer_agent(client: SupportsChatGetResponse) -> AgentFrameworkAgent:
|
||||
"""Create a document writer agent with predictive state updates.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured AgentFrameworkAgent instance with document writing capabilities
|
||||
"""
|
||||
agent = Agent(
|
||||
name="document_writer",
|
||||
instructions=_DOCUMENT_WRITER_INSTRUCTIONS,
|
||||
client=client,
|
||||
tools=[write_document],
|
||||
)
|
||||
|
||||
return AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
name="DocumentWriter",
|
||||
description="Writes and edits documents with predictive state updates",
|
||||
state_schema={
|
||||
"document": {"type": "string", "description": "The current document content"},
|
||||
},
|
||||
predict_state_config={
|
||||
"document": {"tool": "write_document", "tool_argument": "document"},
|
||||
},
|
||||
)
|
||||
+86
@@ -0,0 +1,86 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Human-in-the-loop agent demonstrating step customization (Feature 5)."""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import Agent, SupportsChatGetResponse, tool
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class StepStatus(str, Enum):
|
||||
"""Status of a task step."""
|
||||
|
||||
ENABLED = "enabled"
|
||||
DISABLED = "disabled"
|
||||
|
||||
|
||||
class TaskStep(BaseModel):
|
||||
"""A single step in a task execution plan."""
|
||||
|
||||
description: str = Field(..., description="The text of the step in imperative form (e.g., 'Dig hole', 'Open door')")
|
||||
status: StepStatus = Field(default=StepStatus.ENABLED, description="Whether the step is enabled or disabled")
|
||||
|
||||
|
||||
@tool(
|
||||
name="generate_task_steps",
|
||||
description="Generate execution steps for a task",
|
||||
approval_mode="always_require",
|
||||
)
|
||||
def generate_task_steps(steps: list[TaskStep]) -> str:
|
||||
"""Make up 10 steps (only a couple of words per step) that are required for a task.
|
||||
|
||||
The step should be in imperative form (i.e. Dig hole, Open door, ...).
|
||||
Each step will have status='enabled' by default.
|
||||
|
||||
Args:
|
||||
steps: An array of 10 step objects, each containing description and status
|
||||
|
||||
Returns:
|
||||
Confirmation message
|
||||
"""
|
||||
return f"Generated {len(steps)} execution steps for the task."
|
||||
|
||||
|
||||
def human_in_the_loop_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
|
||||
"""Create a human-in-the-loop agent using tool-based approach for predictive state.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured Agent instance with human-in-the-loop capabilities
|
||||
"""
|
||||
return Agent(
|
||||
name="human_in_the_loop_agent",
|
||||
instructions="""You are a helpful assistant that can perform any task by breaking it down into steps.
|
||||
|
||||
When asked to perform a task, you MUST call the `generate_task_steps` function with the proper
|
||||
number of steps per the request.
|
||||
|
||||
Rules for steps:
|
||||
- Each step description should be in imperative form (e.g., "Dig hole", "Open door", "Prepare ingredients")
|
||||
- Each step should be brief (only a couple of words)
|
||||
- All steps must have status='enabled' initially
|
||||
|
||||
Example steps for "Build a robot":
|
||||
1. "Design blueprint"
|
||||
2. "Gather components"
|
||||
3. "Assemble frame"
|
||||
4. "Install motors"
|
||||
5. "Wire electronics"
|
||||
6. "Program controller"
|
||||
7. "Test movements"
|
||||
8. "Add sensors"
|
||||
9. "Calibrate systems"
|
||||
10. "Final testing"
|
||||
|
||||
IMPORTANT: When you call generate_task_steps, the user will be shown the steps and asked to approve.
|
||||
Do NOT output any text along with the function call - just call the function.
|
||||
After the user approves and the function executes, THEN provide a brief acknowledgment like:
|
||||
"The plan has been created with X steps selected."
|
||||
""",
|
||||
client=client,
|
||||
tools=[generate_task_steps],
|
||||
)
|
||||
@@ -0,0 +1,134 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Recipe agent example demonstrating shared state management (Feature 3)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import Agent, SupportsChatGetResponse, tool
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SkillLevel(str, Enum):
|
||||
"""The skill level required for the recipe."""
|
||||
|
||||
BEGINNER = "Beginner"
|
||||
INTERMEDIATE = "Intermediate"
|
||||
ADVANCED = "Advanced"
|
||||
|
||||
|
||||
class CookingTime(str, Enum):
|
||||
"""The cooking time of the recipe."""
|
||||
|
||||
FIVE_MIN = "5 min"
|
||||
FIFTEEN_MIN = "15 min"
|
||||
THIRTY_MIN = "30 min"
|
||||
FORTY_FIVE_MIN = "45 min"
|
||||
SIXTY_PLUS_MIN = "60+ min"
|
||||
|
||||
|
||||
class Ingredient(BaseModel):
|
||||
"""An ingredient with its details."""
|
||||
|
||||
icon: str = Field(..., description="Emoji icon representing the ingredient (e.g., 🥕)")
|
||||
name: str = Field(..., description="Name of the ingredient")
|
||||
amount: str = Field(..., description="Amount or quantity of the ingredient")
|
||||
|
||||
|
||||
class Recipe(BaseModel):
|
||||
"""A complete recipe."""
|
||||
|
||||
title: str = Field(..., description="The title of the recipe")
|
||||
skill_level: SkillLevel = Field(..., description="The skill level required")
|
||||
special_preferences: list[str] = Field(
|
||||
default_factory=list, description="Dietary preferences (e.g., Vegetarian, Gluten-free)"
|
||||
)
|
||||
cooking_time: CookingTime = Field(..., description="The estimated cooking time")
|
||||
ingredients: list[Ingredient] = Field(..., description="Complete list of ingredients")
|
||||
instructions: list[str] = Field(..., description="Step-by-step cooking instructions")
|
||||
|
||||
|
||||
@tool
|
||||
def update_recipe(recipe: Recipe) -> str:
|
||||
"""Update the recipe with new or modified content.
|
||||
|
||||
You MUST write the complete recipe with ALL fields, even when changing only a few items.
|
||||
When modifying an existing recipe, include ALL existing ingredients and instructions plus your changes.
|
||||
NEVER delete existing data - only add or modify.
|
||||
|
||||
Args:
|
||||
recipe: The complete recipe object with all details
|
||||
|
||||
Returns:
|
||||
Confirmation that the recipe was updated
|
||||
"""
|
||||
return "Recipe updated."
|
||||
|
||||
|
||||
_RECIPE_INSTRUCTIONS = """You are a helpful recipe assistant that creates and modifies recipes.
|
||||
|
||||
CRITICAL RULES:
|
||||
1. You will receive the current recipe state in the system context
|
||||
2. To update the recipe, you MUST use the update_recipe tool
|
||||
3. When modifying a recipe, ALWAYS include ALL existing data plus your changes in the tool call
|
||||
4. NEVER delete existing ingredients or instructions - only add or modify
|
||||
5. After calling the tool, provide a brief conversational message (1-2 sentences)
|
||||
|
||||
When creating a NEW recipe:
|
||||
- Provide all required fields: title, skill_level, cooking_time, ingredients, instructions
|
||||
- Use actual emojis for ingredient icons (🥕 🧄 🧅 🍅 🌿 🍗 🥩 🧀)
|
||||
- Leave special_preferences empty unless specified
|
||||
- Message: "Here's your recipe!" or similar
|
||||
|
||||
When MODIFYING or IMPROVING an existing recipe:
|
||||
- Include ALL existing ingredients + any new ones
|
||||
- Include ALL existing instructions + any new/modified ones
|
||||
- Update other fields as needed
|
||||
- Message: Explain what you improved (e.g., "I upgraded the ingredients to premium quality")
|
||||
- When asked to "improve", enhance with:
|
||||
* Better ingredients (upgrade quality, add complementary flavors)
|
||||
* More detailed instructions
|
||||
* Professional techniques
|
||||
* Adjust skill_level if complexity changes
|
||||
* Add relevant special_preferences
|
||||
|
||||
Example improvements:
|
||||
- Upgrade "chicken" → "organic free-range chicken breast"
|
||||
- Add herbs: basil, oregano, thyme
|
||||
- Add aromatics: garlic, shallots
|
||||
- Add finishing touches: lemon zest, fresh parsley
|
||||
- Make instructions more detailed and professional
|
||||
"""
|
||||
|
||||
|
||||
def recipe_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
|
||||
"""Create a recipe agent with streaming state updates.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured AgentFrameworkAgent instance with recipe management
|
||||
"""
|
||||
agent = Agent(
|
||||
name="recipe_agent",
|
||||
instructions=_RECIPE_INSTRUCTIONS,
|
||||
client=client,
|
||||
tools=[update_recipe],
|
||||
)
|
||||
|
||||
return AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
name="RecipeAgent",
|
||||
description="Creates and modifies recipes with streaming state updates",
|
||||
state_schema={
|
||||
"recipe": {"type": "object", "description": "The current recipe"},
|
||||
},
|
||||
predict_state_config={
|
||||
"recipe": {"tool": "update_recipe", "tool_argument": "recipe"},
|
||||
},
|
||||
require_confirmation=False,
|
||||
)
|
||||
+111
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Example agent demonstrating agentic generative UI with custom events during execution."""
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import Agent, SupportsChatGetResponse, tool
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
|
||||
@tool
|
||||
async def research_topic(topic: str) -> str:
|
||||
"""Research a topic and generate a comprehensive report.
|
||||
|
||||
Args:
|
||||
topic: The topic to research
|
||||
|
||||
Returns:
|
||||
Research report
|
||||
"""
|
||||
# Simulate multi-step research process
|
||||
steps = [
|
||||
("Searching databases", 1.0),
|
||||
("Analyzing sources", 1.5),
|
||||
("Synthesizing information", 1.0),
|
||||
("Generating report", 0.5),
|
||||
]
|
||||
|
||||
results: list[str] = []
|
||||
for step_name, duration in steps:
|
||||
await asyncio.sleep(duration)
|
||||
results.append(f"- {step_name}: completed")
|
||||
|
||||
return f"Research report on '{topic}':\n" + "\n".join(results)
|
||||
|
||||
|
||||
@tool
|
||||
async def create_presentation(title: str, num_slides: int) -> str:
|
||||
"""Create a presentation with multiple slides.
|
||||
|
||||
Args:
|
||||
title: Presentation title
|
||||
num_slides: Number of slides to create
|
||||
|
||||
Returns:
|
||||
Presentation summary
|
||||
"""
|
||||
# Simulate slide generation
|
||||
slides: list[str] = []
|
||||
for i in range(num_slides):
|
||||
await asyncio.sleep(0.5)
|
||||
slides.append(f"Slide {i + 1}: Content for {title}")
|
||||
|
||||
return f"Created presentation '{title}' with {num_slides} slides:\n" + "\n".join(slides)
|
||||
|
||||
|
||||
@tool
|
||||
async def analyze_data(dataset: str) -> str:
|
||||
"""Analyze a dataset and produce insights.
|
||||
|
||||
Args:
|
||||
dataset: The dataset name to analyze
|
||||
|
||||
Returns:
|
||||
Analysis results
|
||||
"""
|
||||
# Simulate data analysis phases
|
||||
phases = [
|
||||
("Loading data", 0.8),
|
||||
("Cleaning data", 1.0),
|
||||
("Running statistical analysis", 1.2),
|
||||
("Generating visualizations", 0.7),
|
||||
]
|
||||
|
||||
insights: list[str] = []
|
||||
for phase_name, duration in phases:
|
||||
await asyncio.sleep(duration)
|
||||
insights.append(f"- {phase_name}: done")
|
||||
|
||||
return f"Analysis of '{dataset}':\n" + "\n".join(insights)
|
||||
|
||||
|
||||
_RESEARCH_ASSISTANT_INSTRUCTIONS = (
|
||||
"You are a research and analysis assistant. "
|
||||
"You can research topics, create presentations, and analyze data. "
|
||||
"Use the available tools to help users with their research needs."
|
||||
)
|
||||
|
||||
|
||||
def research_assistant_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
|
||||
"""Create a research assistant agent.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured AgentFrameworkAgent instance with research capabilities
|
||||
"""
|
||||
agent = Agent(
|
||||
name="research_assistant",
|
||||
instructions=_RESEARCH_ASSISTANT_INSTRUCTIONS,
|
||||
client=client,
|
||||
tools=[research_topic, create_presentation, analyze_data],
|
||||
)
|
||||
|
||||
return AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
name="ResearchAssistant",
|
||||
description="Research assistant that emits progress events during task execution",
|
||||
)
|
||||
@@ -0,0 +1,23 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Simple agentic chat example (Feature 1: Agentic Chat)."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import Agent, SupportsChatGetResponse
|
||||
|
||||
|
||||
def simple_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
|
||||
"""Create a simple chat agent.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured Agent instance
|
||||
"""
|
||||
return Agent[Any](
|
||||
name="simple_chat_agent",
|
||||
instructions="You are a helpful assistant. Be concise and friendly.",
|
||||
client=client,
|
||||
)
|
||||
@@ -0,0 +1,405 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Subgraphs travel planner built with MAF workflow primitives."""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from ag_ui.core import (
|
||||
BaseEvent,
|
||||
StateSnapshotEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
TextMessageStartEvent,
|
||||
)
|
||||
from agent_framework import (
|
||||
Executor,
|
||||
Message,
|
||||
Workflow,
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
handler,
|
||||
response_handler,
|
||||
)
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkWorkflow
|
||||
|
||||
STATIC_FLIGHTS: list[dict[str, str]] = [
|
||||
{
|
||||
"airline": "KLM",
|
||||
"departure": "Amsterdam (AMS)",
|
||||
"arrival": "San Francisco (SFO)",
|
||||
"price": "$650",
|
||||
"duration": "11h 30m",
|
||||
},
|
||||
{
|
||||
"airline": "United",
|
||||
"departure": "Amsterdam (AMS)",
|
||||
"arrival": "San Francisco (SFO)",
|
||||
"price": "$720",
|
||||
"duration": "12h 15m",
|
||||
},
|
||||
]
|
||||
|
||||
STATIC_HOTELS: list[dict[str, str]] = [
|
||||
{
|
||||
"name": "Hotel Zephyr",
|
||||
"location": "Fisherman's Wharf",
|
||||
"price_per_night": "$280/night",
|
||||
"rating": "4.2 stars",
|
||||
},
|
||||
{
|
||||
"name": "The Ritz-Carlton",
|
||||
"location": "Nob Hill",
|
||||
"price_per_night": "$550/night",
|
||||
"rating": "4.8 stars",
|
||||
},
|
||||
{
|
||||
"name": "Hotel Zoe",
|
||||
"location": "Union Square",
|
||||
"price_per_night": "$320/night",
|
||||
"rating": "4.4 stars",
|
||||
},
|
||||
]
|
||||
|
||||
STATIC_EXPERIENCES: list[dict[str, str]] = [
|
||||
{
|
||||
"name": "Pier 39",
|
||||
"type": "activity",
|
||||
"description": "Iconic waterfront destination with shops and sea lions",
|
||||
"location": "Fisherman's Wharf",
|
||||
},
|
||||
{
|
||||
"name": "Golden Gate Bridge",
|
||||
"type": "activity",
|
||||
"description": "World-famous suspension bridge with stunning views",
|
||||
"location": "Golden Gate",
|
||||
},
|
||||
{
|
||||
"name": "Swan Oyster Depot",
|
||||
"type": "restaurant",
|
||||
"description": "Historic seafood counter serving fresh oysters",
|
||||
"location": "Polk Street",
|
||||
},
|
||||
{
|
||||
"name": "Tartine Bakery",
|
||||
"type": "restaurant",
|
||||
"description": "Artisanal bakery famous for bread and pastries",
|
||||
"location": "Mission District",
|
||||
},
|
||||
]
|
||||
|
||||
_STATE_KEY = "subgraphs_state"
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PresentFlights:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PresentHotels:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PlanExperiences:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FinalizeTrip:
|
||||
pass
|
||||
|
||||
|
||||
def _initial_state() -> dict[str, Any]:
|
||||
return {
|
||||
"itinerary": {},
|
||||
"experiences": [],
|
||||
"flights": [],
|
||||
"hotels": [],
|
||||
"planning_step": "start",
|
||||
"active_agent": "supervisor",
|
||||
}
|
||||
|
||||
|
||||
def _emit_text_events(text: str) -> list[BaseEvent]:
|
||||
message_id = str(uuid.uuid4())
|
||||
return [
|
||||
TextMessageStartEvent(message_id=message_id, role="assistant"),
|
||||
TextMessageContentEvent(message_id=message_id, delta=text),
|
||||
TextMessageEndEvent(message_id=message_id),
|
||||
]
|
||||
|
||||
|
||||
async def _emit_text(ctx: WorkflowContext[Any, BaseEvent], text: str) -> None:
|
||||
for event in _emit_text_events(text):
|
||||
await ctx.yield_output(event)
|
||||
|
||||
|
||||
async def _emit_state_snapshot(ctx: WorkflowContext[Any, BaseEvent], state: dict[str, Any]) -> None:
|
||||
await ctx.yield_output(StateSnapshotEvent(snapshot=deepcopy(state)))
|
||||
|
||||
|
||||
def _flight_interrupt_value() -> dict[str, Any]:
|
||||
return {
|
||||
"message": "Choose the flight you want. I recommend KLM because it is cheaper and usually on time.",
|
||||
"options": deepcopy(STATIC_FLIGHTS),
|
||||
"recommendation": deepcopy(STATIC_FLIGHTS[0]),
|
||||
"agent": "flights",
|
||||
}
|
||||
|
||||
|
||||
def _hotel_interrupt_value() -> dict[str, Any]:
|
||||
return {
|
||||
"message": "Choose your hotel. I recommend Hotel Zoe for the best value in a central location.",
|
||||
"options": deepcopy(STATIC_HOTELS),
|
||||
"recommendation": deepcopy(STATIC_HOTELS[2]),
|
||||
"agent": "hotels",
|
||||
}
|
||||
|
||||
|
||||
def _normalize_flight(value: Any) -> dict[str, str] | None:
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
if isinstance(value, dict) and value.get("airline"):
|
||||
return {
|
||||
"airline": str(value.get("airline", "")),
|
||||
"departure": str(value.get("departure", "")),
|
||||
"arrival": str(value.get("arrival", "")),
|
||||
"price": str(value.get("price", "")),
|
||||
"duration": str(value.get("duration", "")),
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_hotel(value: Any) -> dict[str, str] | None:
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
if isinstance(value, dict) and value.get("name"):
|
||||
return {
|
||||
"name": str(value.get("name", "")),
|
||||
"location": str(value.get("location", "")),
|
||||
"price_per_night": str(value.get("price_per_night", "")),
|
||||
"rating": str(value.get("rating", "")),
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
def _load_state(ctx: WorkflowContext[Any, BaseEvent]) -> dict[str, Any]:
|
||||
state = ctx.get_state(_STATE_KEY)
|
||||
if isinstance(state, dict):
|
||||
return state
|
||||
new_state = _initial_state()
|
||||
ctx.set_state(_STATE_KEY, new_state)
|
||||
return new_state
|
||||
|
||||
|
||||
class _SupervisorExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="supervisor_agent")
|
||||
|
||||
@handler
|
||||
async def start(self, message: list[Message], ctx: WorkflowContext[_PresentFlights, BaseEvent]) -> None:
|
||||
del message
|
||||
state = _initial_state()
|
||||
ctx.set_state(_STATE_KEY, state)
|
||||
await _emit_state_snapshot(ctx, state)
|
||||
|
||||
await _emit_text(
|
||||
ctx,
|
||||
"Supervisor: I will coordinate our specialist agents to plan your San Francisco trip end to end.",
|
||||
)
|
||||
|
||||
state["active_agent"] = "flights"
|
||||
state["planning_step"] = "collecting_flights"
|
||||
state["flights"] = deepcopy(STATIC_FLIGHTS)
|
||||
ctx.set_state(_STATE_KEY, state)
|
||||
await _emit_state_snapshot(ctx, state)
|
||||
|
||||
await ctx.send_message(_PresentFlights(), target_id="flights_agent")
|
||||
|
||||
@handler
|
||||
async def finalize(self, message: _FinalizeTrip, ctx: WorkflowContext[Any, BaseEvent]) -> None:
|
||||
del message
|
||||
state = _load_state(ctx)
|
||||
state["active_agent"] = "supervisor"
|
||||
state["planning_step"] = "complete"
|
||||
ctx.set_state(_STATE_KEY, state)
|
||||
await _emit_state_snapshot(ctx, state)
|
||||
await _emit_text(ctx, "Supervisor: Your travel planning is complete and your itinerary is ready.")
|
||||
|
||||
|
||||
class _FlightsExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="flights_agent")
|
||||
|
||||
@handler
|
||||
async def present_options(self, message: _PresentFlights, ctx: WorkflowContext[_PresentHotels, BaseEvent]) -> None:
|
||||
del message
|
||||
await _emit_text(
|
||||
ctx,
|
||||
"Flights Agent: I found two flight options from Amsterdam to San Francisco. "
|
||||
"KLM is recommended for the best value and schedule.",
|
||||
)
|
||||
await ctx.request_info(_flight_interrupt_value(), dict, request_id="flights-choice")
|
||||
|
||||
@response_handler
|
||||
async def handle_selection(
|
||||
self,
|
||||
original_request: dict,
|
||||
response: dict,
|
||||
ctx: WorkflowContext[_PresentHotels, BaseEvent],
|
||||
) -> None:
|
||||
del original_request
|
||||
state = _load_state(ctx)
|
||||
selected_flight = _normalize_flight(response)
|
||||
|
||||
if selected_flight is None:
|
||||
state["active_agent"] = "flights"
|
||||
state["planning_step"] = "collecting_flights"
|
||||
state["flights"] = deepcopy(STATIC_FLIGHTS)
|
||||
ctx.set_state(_STATE_KEY, state)
|
||||
await _emit_state_snapshot(ctx, state)
|
||||
await _emit_text(ctx, "Flights Agent: Please choose a flight option from the selection card to continue.")
|
||||
await ctx.request_info(_flight_interrupt_value(), dict, request_id="flights-choice")
|
||||
return
|
||||
|
||||
itinerary = state.setdefault("itinerary", {})
|
||||
itinerary["flight"] = selected_flight
|
||||
|
||||
state["active_agent"] = "flights"
|
||||
state["planning_step"] = "booking_flight"
|
||||
ctx.set_state(_STATE_KEY, state)
|
||||
await _emit_state_snapshot(ctx, state)
|
||||
|
||||
await _emit_text(
|
||||
ctx,
|
||||
f"Flights Agent: Great choice. I will book the {selected_flight['airline']} flight. "
|
||||
"Now I am routing you to Hotels Agent for accommodation.",
|
||||
)
|
||||
|
||||
state["active_agent"] = "hotels"
|
||||
state["planning_step"] = "collecting_hotels"
|
||||
state["hotels"] = deepcopy(STATIC_HOTELS)
|
||||
ctx.set_state(_STATE_KEY, state)
|
||||
await _emit_state_snapshot(ctx, state)
|
||||
|
||||
await ctx.send_message(_PresentHotels(), target_id="hotels_agent")
|
||||
|
||||
|
||||
class _HotelsExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="hotels_agent")
|
||||
|
||||
@handler
|
||||
async def present_options(self, message: _PresentHotels, ctx: WorkflowContext[_PlanExperiences, BaseEvent]) -> None:
|
||||
del message
|
||||
await _emit_text(
|
||||
ctx,
|
||||
"Hotels Agent: I found three accommodation options in San Francisco. "
|
||||
"Hotel Zoe is recommended for the best balance of location, quality, and price.",
|
||||
)
|
||||
await ctx.request_info(_hotel_interrupt_value(), dict, request_id="hotels-choice")
|
||||
|
||||
@response_handler
|
||||
async def handle_selection(
|
||||
self,
|
||||
original_request: dict,
|
||||
response: dict,
|
||||
ctx: WorkflowContext[_PlanExperiences, BaseEvent],
|
||||
) -> None:
|
||||
del original_request
|
||||
state = _load_state(ctx)
|
||||
selected_hotel = _normalize_hotel(response)
|
||||
|
||||
if selected_hotel is None:
|
||||
state["active_agent"] = "hotels"
|
||||
state["planning_step"] = "collecting_hotels"
|
||||
state["hotels"] = deepcopy(STATIC_HOTELS)
|
||||
ctx.set_state(_STATE_KEY, state)
|
||||
await _emit_state_snapshot(ctx, state)
|
||||
await _emit_text(ctx, "Hotels Agent: Please choose a hotel option from the selection card to continue.")
|
||||
await ctx.request_info(_hotel_interrupt_value(), dict, request_id="hotels-choice")
|
||||
return
|
||||
|
||||
itinerary = state.setdefault("itinerary", {})
|
||||
itinerary["hotel"] = selected_hotel
|
||||
|
||||
state["active_agent"] = "hotels"
|
||||
state["planning_step"] = "booking_hotel"
|
||||
ctx.set_state(_STATE_KEY, state)
|
||||
await _emit_state_snapshot(ctx, state)
|
||||
|
||||
await _emit_text(
|
||||
ctx,
|
||||
f"Hotels Agent: Excellent, {selected_hotel['name']} is booked. "
|
||||
"I am routing you to Experiences Agent for activities and restaurants.",
|
||||
)
|
||||
|
||||
state["active_agent"] = "experiences"
|
||||
state["planning_step"] = "curating_experiences"
|
||||
state["experiences"] = deepcopy(STATIC_EXPERIENCES)
|
||||
ctx.set_state(_STATE_KEY, state)
|
||||
await _emit_state_snapshot(ctx, state)
|
||||
|
||||
await ctx.send_message(_PlanExperiences(), target_id="experiences_agent")
|
||||
|
||||
|
||||
class _ExperiencesExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="experiences_agent")
|
||||
|
||||
@handler
|
||||
async def plan(self, message: _PlanExperiences, ctx: WorkflowContext[_FinalizeTrip, BaseEvent]) -> None:
|
||||
del message
|
||||
await _emit_text(
|
||||
ctx,
|
||||
"Experiences Agent: I planned activities and restaurants including "
|
||||
"Pier 39, Golden Gate Bridge, Swan Oyster Depot, and Tartine Bakery.",
|
||||
)
|
||||
await ctx.send_message(_FinalizeTrip(), target_id="supervisor_agent")
|
||||
|
||||
|
||||
def _build_subgraphs_workflow() -> Workflow:
|
||||
supervisor = _SupervisorExecutor()
|
||||
flights = _FlightsExecutor()
|
||||
hotels = _HotelsExecutor()
|
||||
experiences = _ExperiencesExecutor()
|
||||
|
||||
return (
|
||||
WorkflowBuilder(
|
||||
name="subgraphs",
|
||||
description="Travel planning supervisor with flights/hotels/experiences subgraphs.",
|
||||
start_executor=supervisor,
|
||||
)
|
||||
.add_edge(supervisor, flights)
|
||||
.add_edge(flights, hotels)
|
||||
.add_edge(hotels, experiences)
|
||||
.add_edge(experiences, supervisor)
|
||||
.build()
|
||||
)
|
||||
|
||||
|
||||
def _build_subgraphs_workflow_for_thread(thread_id: str) -> Workflow:
|
||||
"""Create a workflow instance scoped to a single AG-UI thread."""
|
||||
del thread_id
|
||||
return _build_subgraphs_workflow()
|
||||
|
||||
|
||||
def subgraphs_agent() -> AgentFrameworkWorkflow:
|
||||
"""Create the subgraphs travel planner agent."""
|
||||
return AgentFrameworkWorkflow(
|
||||
workflow_factory=_build_subgraphs_workflow_for_thread,
|
||||
name="subgraphs",
|
||||
description="Travel planning workflow with interrupt-driven selections.",
|
||||
)
|
||||
@@ -0,0 +1,84 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Example agent demonstrating human-in-the-loop with function approvals."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import Agent, SupportsChatGetResponse, tool
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
|
||||
@tool(approval_mode="always_require")
|
||||
def create_calendar_event(title: str, date: str, time: str) -> str:
|
||||
"""Create a calendar event.
|
||||
|
||||
Args:
|
||||
title: The event title
|
||||
date: The event date (YYYY-MM-DD)
|
||||
time: The event time (HH:MM)
|
||||
|
||||
Returns:
|
||||
Confirmation message
|
||||
"""
|
||||
return f"Calendar event '{title}' created for {date} at {time}"
|
||||
|
||||
|
||||
@tool(approval_mode="always_require")
|
||||
def send_email(to: str, subject: str, body: str) -> str:
|
||||
"""Send an email.
|
||||
|
||||
Args:
|
||||
to: Recipient email address
|
||||
subject: Email subject
|
||||
body: Email body text
|
||||
|
||||
Returns:
|
||||
Confirmation message
|
||||
"""
|
||||
return f"Email sent to {to} with subject '{subject}'"
|
||||
|
||||
|
||||
@tool(approval_mode="always_require")
|
||||
def book_meeting_room(room_name: str, date: str, start_time: str, end_time: str) -> str:
|
||||
"""Book a meeting room.
|
||||
|
||||
Args:
|
||||
room_name: The meeting room name
|
||||
date: The booking date (YYYY-MM-DD)
|
||||
start_time: Start time (HH:MM)
|
||||
end_time: End time (HH:MM)
|
||||
|
||||
Returns:
|
||||
Confirmation message
|
||||
"""
|
||||
return f"Meeting room '{room_name}' booked for {date} from {start_time} to {end_time}"
|
||||
|
||||
|
||||
_TASK_PLANNER_INSTRUCTIONS = (
|
||||
"You are a helpful assistant that plans and executes tasks. "
|
||||
"You have access to calendar, email, and meeting room booking functions. "
|
||||
"All of these actions require user approval before execution."
|
||||
)
|
||||
|
||||
|
||||
def task_planner_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
|
||||
"""Create a task planner agent with user approval for actions.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured AgentFrameworkAgent instance with task planning capabilities
|
||||
"""
|
||||
agent = Agent(
|
||||
name="task_planner",
|
||||
instructions=_TASK_PLANNER_INSTRUCTIONS,
|
||||
client=client,
|
||||
tools=[create_calendar_event, send_email, book_meeting_room],
|
||||
)
|
||||
|
||||
return AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
name="TaskPlanner",
|
||||
description="Plans and executes tasks with user approval",
|
||||
)
|
||||
@@ -0,0 +1,338 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Task steps agent demonstrating agentic generative UI (Feature 6)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from ag_ui.core import (
|
||||
EventType,
|
||||
MessagesSnapshotEvent,
|
||||
RunFinishedEvent,
|
||||
StateDeltaEvent,
|
||||
StateSnapshotEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
TextMessageStartEvent,
|
||||
ToolCallStartEvent,
|
||||
)
|
||||
from agent_framework import Agent, Content, Message, SupportsChatGetResponse, tool
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkWorkflow
|
||||
|
||||
|
||||
class StepStatus(str, Enum):
|
||||
"""Status of a task step."""
|
||||
|
||||
PENDING = "pending"
|
||||
COMPLETED = "completed"
|
||||
|
||||
|
||||
class TaskStep(BaseModel):
|
||||
"""A single step in a task."""
|
||||
|
||||
description: str = Field(
|
||||
..., description="The text of the step in gerund form (e.g., 'Digging hole', 'Opening door')"
|
||||
)
|
||||
status: StepStatus = Field(default=StepStatus.PENDING, description="The status of the step")
|
||||
|
||||
|
||||
@tool
|
||||
def generate_task_steps(steps: list[TaskStep]) -> str:
|
||||
"""Generate a list of task steps for completing a task.
|
||||
|
||||
Args:
|
||||
steps: Complete list of task steps with descriptions and status
|
||||
|
||||
Returns:
|
||||
Confirmation that steps were generated
|
||||
"""
|
||||
return "Steps generated."
|
||||
|
||||
|
||||
def _create_task_steps_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
|
||||
"""Create the task steps agent using tool-based approach for streaming.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured AgentFrameworkAgent instance
|
||||
"""
|
||||
agent = Agent[Any](
|
||||
name="task_steps_agent",
|
||||
instructions="""You are a helpful assistant that breaks down tasks into actionable steps.
|
||||
|
||||
When asked to perform a task, you MUST:
|
||||
1. Use the generate_task_steps tool to create the steps
|
||||
2. Pay attention to how many steps the user requests (if specified)
|
||||
3. If no specific number is mentioned, use a reasonable number of steps (typically 5-10)
|
||||
4. Each step description should be in gerund form (e.g., "Designing spacecraft", "Training astronauts")
|
||||
5. Each step should be brief (only 2-4 words)
|
||||
6. All steps must have status='pending'
|
||||
7. After calling the tool, provide a brief conversational message (one sentence) saying you created the plan
|
||||
|
||||
Example steps for "Build a treehouse in 5 steps":
|
||||
- "Selecting location"
|
||||
- "Gathering materials"
|
||||
- "Assembling frame"
|
||||
- "Installing platform"
|
||||
- "Adding finishing touches"
|
||||
""",
|
||||
client=client,
|
||||
tools=[generate_task_steps],
|
||||
)
|
||||
|
||||
return AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
name="TaskStepsAgent",
|
||||
description="Generates task steps with streaming state updates",
|
||||
state_schema={
|
||||
"steps": {"type": "array", "description": "The list of task steps"},
|
||||
},
|
||||
predict_state_config={
|
||||
"steps": {
|
||||
"tool": "generate_task_steps",
|
||||
"tool_argument": "steps",
|
||||
}
|
||||
},
|
||||
require_confirmation=False, # Agentic generative UI updates automatically without confirmation
|
||||
)
|
||||
|
||||
|
||||
# Wrap the agent's run method to add step execution simulation
|
||||
class TaskStepsAgentWithExecution(AgentFrameworkWorkflow):
|
||||
"""Wrapper that adds step execution simulation after plan generation.
|
||||
|
||||
This wrapper delegates to AgentFrameworkAgent but is recognized as compatible
|
||||
by add_agent_framework_fastapi_endpoint since it implements run().
|
||||
"""
|
||||
|
||||
def __init__(self, base_agent: AgentFrameworkAgent):
|
||||
"""Initialize wrapper with base agent."""
|
||||
super().__init__(name=base_agent.name, description=base_agent.description)
|
||||
self._base_agent = base_agent
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
"""Delegate all other attribute access to base agent."""
|
||||
return getattr(self._base_agent, name)
|
||||
|
||||
async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[Any]:
|
||||
"""Run the agent and then simulate step execution."""
|
||||
import logging
|
||||
import uuid
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info("TaskStepsAgentWithExecution.run() called - wrapper is active")
|
||||
|
||||
# First, run the base agent to generate the plan - buffer text messages
|
||||
final_state: dict[str, Any] = {}
|
||||
run_finished_event: Any = None
|
||||
tool_call_id: str | None = None
|
||||
buffered_text_events: list[Any] = [] # Buffer text from first LLM call
|
||||
|
||||
async for event in self._base_agent.run(input_data):
|
||||
event_type_str = str(event.type) if hasattr(event, "type") else type(event).__name__
|
||||
logger.info(f"Processing event: {event_type_str}")
|
||||
|
||||
match event:
|
||||
case StateSnapshotEvent(snapshot=snapshot):
|
||||
final_state = snapshot.copy() if snapshot else {}
|
||||
logger.info(f"Captured STATE_SNAPSHOT event with state: {final_state}")
|
||||
yield event
|
||||
case StateDeltaEvent(delta=delta):
|
||||
# Apply state delta to final_state
|
||||
if delta:
|
||||
for patch in delta:
|
||||
if patch.get("op") == "replace" and patch.get("path") == "/steps":
|
||||
final_state["steps"] = patch.get("value", [])
|
||||
logger.info(
|
||||
f"Applied STATE_DELTA: updated steps to {len(final_state.get('steps', []))} items"
|
||||
)
|
||||
logger.info(f"Yielding event immediately: {event_type_str}")
|
||||
yield event
|
||||
case RunFinishedEvent():
|
||||
run_finished_event = event
|
||||
logger.info("Captured RUN_FINISHED event - will send after step execution and summary")
|
||||
case ToolCallStartEvent(tool_call_id=call_id):
|
||||
tool_call_id = call_id
|
||||
logger.info(f"Captured tool_call_id: {tool_call_id}")
|
||||
yield event
|
||||
case TextMessageStartEvent() | TextMessageContentEvent() | TextMessageEndEvent():
|
||||
buffered_text_events.append(event)
|
||||
logger.info(f"Buffered {event_type_str} from first LLM call")
|
||||
case _:
|
||||
logger.info(f"Yielding event immediately: {event_type_str}")
|
||||
yield event
|
||||
|
||||
logger.info(f"Base agent completed. Final state: {final_state}")
|
||||
|
||||
# Now simulate executing the steps
|
||||
if final_state and "steps" in final_state:
|
||||
steps = final_state["steps"]
|
||||
logger.info(f"Starting step execution simulation for {len(steps)} steps")
|
||||
|
||||
for i in range(len(steps)):
|
||||
logger.info(f"Simulating execution of step {i + 1}/{len(steps)}: {steps[i].get('description')}")
|
||||
await asyncio.sleep(1.0) # Simulate work
|
||||
|
||||
# Update step to completed
|
||||
steps[i]["status"] = "completed"
|
||||
logger.info(f"Step {i + 1} marked as completed")
|
||||
|
||||
# Send delta event with manual JSON patch format
|
||||
delta_event = StateDeltaEvent(
|
||||
type=EventType.STATE_DELTA,
|
||||
delta=[
|
||||
{
|
||||
"op": "replace",
|
||||
"path": f"/steps/{i}/status",
|
||||
"value": "completed",
|
||||
}
|
||||
],
|
||||
)
|
||||
logger.info(f"Yielding StateDeltaEvent for step {i + 1}")
|
||||
yield delta_event
|
||||
|
||||
# Send final snapshot
|
||||
final_snapshot = StateSnapshotEvent(
|
||||
type=EventType.STATE_SNAPSHOT,
|
||||
snapshot={"steps": steps},
|
||||
)
|
||||
logger.info("Yielding final StateSnapshotEvent with all steps completed")
|
||||
yield final_snapshot
|
||||
|
||||
# SECOND LLM call: Stream summary from chat client directly
|
||||
logger.info("Making SECOND LLM call to generate summary after step execution")
|
||||
|
||||
# Get the underlying chat agent and client
|
||||
chat_agent = self._base_agent.agent
|
||||
client = chat_agent.client # type: ignore
|
||||
|
||||
# Build messages for summary call
|
||||
|
||||
original_messages = input_data.get("messages", [])
|
||||
|
||||
# Convert to Message objects if needed
|
||||
messages: list[Message] = []
|
||||
for msg in original_messages:
|
||||
if isinstance(msg, dict):
|
||||
content_str = msg.get("content", "")
|
||||
if isinstance(content_str, str):
|
||||
messages.append(
|
||||
Message(
|
||||
role=msg.get("role", "user"),
|
||||
contents=[Content.from_text(text=content_str)],
|
||||
)
|
||||
)
|
||||
elif isinstance(msg, Message):
|
||||
messages.append(msg)
|
||||
|
||||
# Add completion message
|
||||
messages.append(
|
||||
Message(
|
||||
role="user",
|
||||
contents=[
|
||||
Content.from_text(
|
||||
text="The steps have been successfully executed. Provide a brief one-sentence summary."
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
# Stream the LLM response and manually emit text events
|
||||
logger.info("Calling chat client for summary")
|
||||
|
||||
message_id = str(uuid.uuid4())
|
||||
|
||||
try:
|
||||
# Emit TEXT_MESSAGE_START
|
||||
yield TextMessageStartEvent(
|
||||
type=EventType.TEXT_MESSAGE_START,
|
||||
message_id=message_id,
|
||||
role="assistant",
|
||||
)
|
||||
# Small delay to ensure START event is processed before CONTENT events
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
# Stream completion
|
||||
accumulated_text = ""
|
||||
async for chunk in client.get_response(messages=messages, stream=True):
|
||||
# chunk is ChatResponseUpdate
|
||||
if hasattr(chunk, "text") and chunk.text:
|
||||
accumulated_text += chunk.text
|
||||
# Emit TEXT_MESSAGE_CONTENT
|
||||
yield TextMessageContentEvent(
|
||||
type=EventType.TEXT_MESSAGE_CONTENT,
|
||||
message_id=message_id,
|
||||
delta=chunk.text,
|
||||
)
|
||||
|
||||
# Emit TEXT_MESSAGE_END
|
||||
yield TextMessageEndEvent(
|
||||
type=EventType.TEXT_MESSAGE_END,
|
||||
message_id=message_id,
|
||||
)
|
||||
logger.info(f"Summary complete: {accumulated_text}")
|
||||
|
||||
# Build complete message for persistence
|
||||
summary_message = {
|
||||
"role": "assistant",
|
||||
"content": accumulated_text,
|
||||
"id": message_id,
|
||||
}
|
||||
final_messages = list(original_messages)
|
||||
final_messages.append(summary_message)
|
||||
|
||||
# Emit MessagesSnapshotEvent to persist in history
|
||||
yield MessagesSnapshotEvent(
|
||||
type=EventType.MESSAGES_SNAPSHOT,
|
||||
messages=final_messages,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating summary: {e}")
|
||||
# Generate a new message ID for the error
|
||||
error_message_id = str(uuid.uuid4())
|
||||
# Yield TEXT_MESSAGE_START for error
|
||||
yield TextMessageStartEvent(
|
||||
type=EventType.TEXT_MESSAGE_START,
|
||||
message_id=error_message_id,
|
||||
role="assistant",
|
||||
)
|
||||
# Yield error message content
|
||||
yield TextMessageContentEvent(
|
||||
type=EventType.TEXT_MESSAGE_CONTENT,
|
||||
message_id=error_message_id,
|
||||
delta=f"[Summary generation error: {e!s}]",
|
||||
)
|
||||
# Yield TEXT_MESSAGE_END for error
|
||||
yield TextMessageEndEvent(
|
||||
type=EventType.TEXT_MESSAGE_END,
|
||||
message_id=error_message_id,
|
||||
)
|
||||
else:
|
||||
logger.warning(f"No steps found in final_state to execute. final_state={final_state}")
|
||||
|
||||
# Finally send the original RUN_FINISHED event
|
||||
if run_finished_event:
|
||||
logger.info("Yielding original RUN_FINISHED event")
|
||||
yield run_finished_event
|
||||
|
||||
|
||||
def task_steps_agent_wrapped(client: SupportsChatGetResponse[Any]) -> TaskStepsAgentWithExecution:
|
||||
"""Create a task steps agent with execution simulation.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A wrapped agent instance with step execution simulation
|
||||
"""
|
||||
base_agent = _create_task_steps_agent(client)
|
||||
return TaskStepsAgentWithExecution(base_agent)
|
||||
@@ -0,0 +1,193 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Example agent demonstrating Tool-based Generative UI (Feature 5)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, TypedDict
|
||||
|
||||
from agent_framework import Agent, FunctionTool, SupportsChatGetResponse
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # pragma: no cover
|
||||
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 agent_framework import ChatOptions
|
||||
|
||||
# Declaration-only tools (func=None) - actual rendering happens on the client side
|
||||
generate_haiku = FunctionTool(
|
||||
name="generate_haiku",
|
||||
description="""Generate a haiku with image and gradient background (FRONTEND_RENDER).
|
||||
|
||||
This tool generates UI for displaying a haiku with an image and gradient background.
|
||||
The frontend should render this as a custom haiku component.""",
|
||||
func=None, # Makes declaration_only=True so client renders the UI
|
||||
input_model={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"english": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "English haiku lines (exactly 3 lines)",
|
||||
"minItems": 3,
|
||||
"maxItems": 3,
|
||||
},
|
||||
"japanese": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Japanese haiku lines (exactly 3 lines)",
|
||||
"minItems": 3,
|
||||
"maxItems": 3,
|
||||
},
|
||||
"image_name": {
|
||||
"type": "string",
|
||||
"description": """Image filename for visual accompaniment. Must be one of:
|
||||
- "Osaka_Castle_Turret_Stone_Wall_Pine_Trees_Daytime.jpg"
|
||||
- "Tokyo_Skyline_Night_Tokyo_Tower_Mount_Fuji_View.jpg"
|
||||
- "Itsukushima_Shrine_Miyajima_Floating_Torii_Gate_Sunset_Long_Exposure.jpg"
|
||||
- "Takachiho_Gorge_Waterfall_River_Lush_Greenery_Japan.jpg"
|
||||
- "Bonsai_Tree_Potted_Japanese_Art_Green_Foliage.jpeg"
|
||||
- "Shirakawa-go_Gassho-zukuri_Thatched_Roof_Village_Aerial_View.jpg"
|
||||
- "Ginkaku-ji_Silver_Pavilion_Kyoto_Japanese_Garden_Pond_Reflection.jpg"
|
||||
- "Senso-ji_Temple_Asakusa_Cherry_Blossoms_Kimono_Umbrella.jpg"
|
||||
- "Cherry_Blossoms_Sakura_Night_View_City_Lights_Japan.jpg"
|
||||
- "Mount_Fuji_Lake_Reflection_Cherry_Blossoms_Sakura_Spring.jpg"
|
||||
""",
|
||||
},
|
||||
"gradient": {
|
||||
"type": "string",
|
||||
"description": 'CSS gradient string for background (e.g., "linear-gradient(135deg, #667eea 0%, #764ba2 100%)")',
|
||||
},
|
||||
},
|
||||
"required": ["english", "japanese", "image_name", "gradient"],
|
||||
},
|
||||
)
|
||||
|
||||
create_chart = FunctionTool(
|
||||
name="create_chart",
|
||||
description="""Create an interactive chart (FRONTEND_RENDER).
|
||||
|
||||
This tool creates chart specifications for frontend rendering.
|
||||
The frontend should render this as an interactive chart component.""",
|
||||
func=None, # Makes declaration_only=True so client renders the UI
|
||||
input_model={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"chart_type": {
|
||||
"type": "string",
|
||||
"description": "Type of chart (bar, line, pie, scatter)",
|
||||
},
|
||||
"data_points": {
|
||||
"type": "array",
|
||||
"items": {"type": "object"},
|
||||
"description": "Data points for the chart",
|
||||
},
|
||||
"title": {
|
||||
"type": "string",
|
||||
"description": "Chart title",
|
||||
},
|
||||
},
|
||||
"required": ["chart_type", "data_points", "title"],
|
||||
},
|
||||
)
|
||||
|
||||
display_timeline = FunctionTool(
|
||||
name="display_timeline",
|
||||
description="""Display an interactive timeline (FRONTEND_RENDER).
|
||||
|
||||
This tool creates timeline specifications for frontend rendering.
|
||||
The frontend should render this as an interactive timeline component.""",
|
||||
func=None, # Makes declaration_only=True so client renders the UI
|
||||
input_model={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"events": {
|
||||
"type": "array",
|
||||
"items": {"type": "object"},
|
||||
"description": "Events to display on the timeline",
|
||||
},
|
||||
"start_date": {
|
||||
"type": "string",
|
||||
"description": "Timeline start date",
|
||||
},
|
||||
"end_date": {
|
||||
"type": "string",
|
||||
"description": "Timeline end date",
|
||||
},
|
||||
},
|
||||
"required": ["events", "start_date", "end_date"],
|
||||
},
|
||||
)
|
||||
|
||||
show_comparison_table = FunctionTool(
|
||||
name="show_comparison_table",
|
||||
description="""Show a comparison table (FRONTEND_RENDER).
|
||||
|
||||
This tool creates table specifications for frontend rendering.
|
||||
The frontend should render this as an interactive comparison table.""",
|
||||
func=None, # Makes declaration_only=True so client renders the UI
|
||||
input_model={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": {"type": "object"},
|
||||
"description": "Items to compare",
|
||||
},
|
||||
"columns": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Column names",
|
||||
},
|
||||
},
|
||||
"required": ["items", "columns"],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
_UI_GENERATOR_INSTRUCTIONS = """You MUST use the provided tools to generate content. Never respond with plain text descriptions.
|
||||
|
||||
For haiku requests:
|
||||
- Call generate_haiku tool with all 4 required parameters
|
||||
- English: 3 lines
|
||||
- Japanese: 3 lines
|
||||
- image_name: Choose from available images
|
||||
- gradient: CSS gradient string
|
||||
|
||||
For other requests, use the appropriate tool (create_chart, display_timeline, show_comparison_table).
|
||||
"""
|
||||
|
||||
OptionsT = TypeVar("OptionsT", bound=TypedDict, default="ChatOptions") # type: ignore[valid-type]
|
||||
|
||||
|
||||
def ui_generator_agent(client: SupportsChatGetResponse[OptionsT]) -> AgentFrameworkAgent:
|
||||
"""Create a UI generator agent with custom React component rendering.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured AgentFrameworkAgent instance with UI generation capabilities
|
||||
"""
|
||||
agent = Agent(
|
||||
name="ui_generator",
|
||||
instructions=_UI_GENERATOR_INSTRUCTIONS,
|
||||
client=client,
|
||||
tools=[generate_haiku, create_chart, display_timeline, show_comparison_table],
|
||||
# Force tool usage - the LLM MUST call a tool, cannot respond with plain text
|
||||
default_options={"tool_choice": "required"}, # type: ignore
|
||||
)
|
||||
|
||||
return AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
name="UIGenerator",
|
||||
description="Generates custom UI components through tool calls",
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Weather agent example demonstrating backend tool rendering."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import Agent, SupportsChatGetResponse, tool
|
||||
|
||||
|
||||
@tool
|
||||
def get_weather(location: str) -> dict[str, Any]:
|
||||
"""Get the current weather for a location.
|
||||
|
||||
Args:
|
||||
location: The city or location to get weather for.
|
||||
|
||||
Returns:
|
||||
Weather information as a dictionary with temperatures in Celsius.
|
||||
"""
|
||||
# Simulated weather data with structured format (temperatures in Celsius for dojo UI)
|
||||
weather_data = {
|
||||
"seattle": {"temperature": 11, "conditions": "rainy", "humidity": 75, "wind_speed": 12, "feels_like": 10},
|
||||
"san francisco": {"temperature": 14, "conditions": "foggy", "humidity": 85, "wind_speed": 8, "feels_like": 13},
|
||||
"new york city": {"temperature": 18, "conditions": "sunny", "humidity": 60, "wind_speed": 10, "feels_like": 17},
|
||||
"miami": {"temperature": 29, "conditions": "hot and humid", "humidity": 90, "wind_speed": 5, "feels_like": 32},
|
||||
"chicago": {"temperature": 9, "conditions": "windy", "humidity": 65, "wind_speed": 20, "feels_like": 6},
|
||||
}
|
||||
|
||||
location_lower = location.lower()
|
||||
if location_lower in weather_data:
|
||||
return weather_data[location_lower]
|
||||
|
||||
return {
|
||||
"temperature": 21,
|
||||
"conditions": "partly cloudy",
|
||||
"humidity": 50,
|
||||
"wind_speed": 10,
|
||||
"feels_like": 20,
|
||||
}
|
||||
|
||||
|
||||
@tool
|
||||
def get_forecast(location: str, days: int = 3) -> str:
|
||||
"""Get the weather forecast for a location.
|
||||
|
||||
Args:
|
||||
location: The city or location to get forecast for.
|
||||
days: Number of days to forecast (default: 3).
|
||||
|
||||
Returns:
|
||||
Forecast information string.
|
||||
"""
|
||||
forecast: list[str] = []
|
||||
for day in range(1, min(days, 7) + 1):
|
||||
forecast.append(f"Day {day}: Partly cloudy, {60 + day * 2}°F")
|
||||
|
||||
return f"{days}-day forecast for {location}:\n" + "\n".join(forecast)
|
||||
|
||||
|
||||
def weather_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
|
||||
"""Create a weather agent with get_weather and get_forecast tools.
|
||||
|
||||
Args:
|
||||
client: The chat client to use for the agent
|
||||
|
||||
Returns:
|
||||
A configured Agent instance with weather tools
|
||||
"""
|
||||
return Agent[Any](
|
||||
name="weather_agent",
|
||||
instructions=(
|
||||
"You are a helpful weather assistant. "
|
||||
"Use the get_weather and get_forecast functions to help users with weather information. "
|
||||
"Always provide friendly and informative responses. "
|
||||
"First return the weather result, and then return details about the forecast."
|
||||
),
|
||||
client=client,
|
||||
tools=[get_weather, get_forecast],
|
||||
)
|
||||
@@ -0,0 +1,92 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Deterministic tool-driven AG-UI state example.
|
||||
|
||||
This sample demonstrates how a tool can push a *deterministic* state update
|
||||
to the AG-UI frontend based on its actual return value — in contrast to
|
||||
``predict_state_config`` which fires optimistically from LLM-predicted tool
|
||||
call arguments. See issue https://github.com/microsoft/agent-framework/issues/3167.
|
||||
|
||||
The :func:`agent_framework_ag_ui.state_update` helper wraps a text result
|
||||
together with a state snapshot. When a tool returns one of these, the AG-UI
|
||||
endpoint merges the snapshot into the shared state and emits a
|
||||
``StateSnapshotEvent`` after the tool result.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import Agent, Content, SupportsChatGetResponse, tool
|
||||
from agent_framework.ag_ui import AgentFrameworkAgent
|
||||
|
||||
from agent_framework_ag_ui import state_update
|
||||
|
||||
# Simulated weather database — in the issue's motivating example the tool
|
||||
# would instead call a real weather API.
|
||||
_WEATHER_DB: dict[str, dict[str, Any]] = {
|
||||
"seattle": {"temperature": 11, "conditions": "rainy", "humidity": 75},
|
||||
"san francisco": {"temperature": 14, "conditions": "foggy", "humidity": 85},
|
||||
"new york city": {"temperature": 18, "conditions": "sunny", "humidity": 60},
|
||||
"miami": {"temperature": 29, "conditions": "hot and humid", "humidity": 90},
|
||||
"chicago": {"temperature": 9, "conditions": "windy", "humidity": 65},
|
||||
}
|
||||
|
||||
|
||||
@tool
|
||||
async def get_weather(location: str) -> Content:
|
||||
"""Fetch current weather for a location and push it into AG-UI shared state.
|
||||
|
||||
Unlike ``predict_state_config`` — which derives state optimistically from
|
||||
LLM-predicted tool call arguments — this tool uses ``state_update`` to
|
||||
forward the *actual* fetched weather to the frontend. The ``text`` goes
|
||||
back to the LLM as the normal tool result, and the ``state`` dict is merged
|
||||
into the AG-UI shared state.
|
||||
|
||||
Args:
|
||||
location: City name to look up.
|
||||
|
||||
Returns:
|
||||
A :class:`Content` carrying both the LLM-visible text result and a
|
||||
deterministic state snapshot.
|
||||
"""
|
||||
key = location.lower()
|
||||
data = _WEATHER_DB.get(
|
||||
key,
|
||||
{"temperature": 21, "conditions": "partly cloudy", "humidity": 50},
|
||||
)
|
||||
weather_record = {"location": location, **data}
|
||||
return state_update(
|
||||
text=(
|
||||
f"The weather in {location} is {data['conditions']} at "
|
||||
f"{data['temperature']}°C with {data['humidity']}% humidity."
|
||||
),
|
||||
state={"weather": weather_record},
|
||||
)
|
||||
|
||||
|
||||
def weather_state_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
|
||||
"""Create an AG-UI agent with a deterministic tool-driven state tool."""
|
||||
agent = Agent[Any](
|
||||
name="weather_state_agent",
|
||||
instructions=(
|
||||
"You are a weather assistant. When a user asks about the weather "
|
||||
"in a city, call the get_weather tool and use its output to give a "
|
||||
"friendly, concise reply. The tool also updates the shared UI state "
|
||||
"so the frontend can render a weather card from the `weather` key."
|
||||
),
|
||||
client=client,
|
||||
tools=[get_weather],
|
||||
)
|
||||
|
||||
return AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
name="WeatherStateAgent",
|
||||
description="Weather agent that deterministically updates shared state from tool results.",
|
||||
state_schema={
|
||||
"weather": {
|
||||
"type": "object",
|
||||
"description": "Last fetched weather record",
|
||||
},
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Backend tool rendering endpoint."""
|
||||
|
||||
from typing import Any, cast
|
||||
|
||||
from agent_framework._clients import SupportsChatGetResponse
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from fastapi import FastAPI
|
||||
|
||||
from ...agents.weather_agent import weather_agent
|
||||
|
||||
|
||||
def register_backend_tool_rendering(app: FastAPI) -> None:
|
||||
"""Register the backend tool rendering endpoint.
|
||||
|
||||
Args:
|
||||
app: The FastAPI application.
|
||||
"""
|
||||
# Create a chat client and call the factory function
|
||||
client = cast(SupportsChatGetResponse[Any], OpenAIChatCompletionClient())
|
||||
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app,
|
||||
weather_agent(client),
|
||||
"/backend_tool_rendering",
|
||||
)
|
||||
@@ -0,0 +1,173 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Example FastAPI server with AG-UI endpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, cast
|
||||
|
||||
import uvicorn
|
||||
from agent_framework import ChatOptions
|
||||
from agent_framework._clients import SupportsChatGetResponse
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from ..agents.document_writer_agent import document_writer_agent
|
||||
from ..agents.human_in_the_loop_agent import human_in_the_loop_agent
|
||||
from ..agents.recipe_agent import recipe_agent
|
||||
from ..agents.simple_agent import simple_agent
|
||||
from ..agents.subgraphs_agent import subgraphs_agent
|
||||
from ..agents.task_steps_agent import task_steps_agent_wrapped
|
||||
from ..agents.ui_generator_agent import ui_generator_agent
|
||||
from ..agents.weather_agent import weather_agent
|
||||
from ..agents.weather_state_agent import weather_state_agent
|
||||
|
||||
AnthropicClient: type[Any] | None
|
||||
try:
|
||||
import agent_framework.anthropic as _anthropic_namespace
|
||||
except ImportError:
|
||||
# If the Anthropic client isn't installed, we can still run the server with Azure OpenAI as the default chat client
|
||||
AnthropicClient = None
|
||||
else:
|
||||
AnthropicClient = cast(type[Any] | None, getattr(_anthropic_namespace, "AnthropicClient", None))
|
||||
|
||||
# Configure logging to file and console (disabled by default - set ENABLE_DEBUG_LOGGING=1 to enable)
|
||||
if os.getenv("ENABLE_DEBUG_LOGGING"):
|
||||
log_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "ag_ui_server.log")
|
||||
|
||||
# Remove any existing handlers
|
||||
root_logger = logging.getLogger()
|
||||
for handler in root_logger.handlers[:]:
|
||||
root_logger.removeHandler(handler)
|
||||
|
||||
# Configure new handlers
|
||||
file_handler = logging.FileHandler(log_file, mode="w")
|
||||
file_handler.setLevel(logging.INFO)
|
||||
file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
|
||||
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setLevel(logging.INFO)
|
||||
console_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
|
||||
|
||||
root_logger.addHandler(file_handler)
|
||||
root_logger.addHandler(console_handler)
|
||||
root_logger.setLevel(logging.INFO)
|
||||
|
||||
# Explicitly set log levels for our modules
|
||||
logging.getLogger("agent_framework_ag_ui").setLevel(logging.INFO)
|
||||
logging.getLogger("agent_framework").setLevel(logging.INFO)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"AG-UI Examples Server starting... Logs writing to: {log_file}")
|
||||
|
||||
app = FastAPI(title="Agent Framework AG-UI Example Server")
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Create a shared chat client for all agents
|
||||
# You can use different chat clients for different agents if needed
|
||||
# Set CHAT_CLIENT=anthropic to use Anthropic, defaults to Azure OpenAI
|
||||
client: SupportsChatGetResponse[ChatOptions] = cast(
|
||||
SupportsChatGetResponse[ChatOptions],
|
||||
AnthropicClient()
|
||||
if AnthropicClient is not None and os.getenv("CHAT_CLIENT", "").lower() == "anthropic"
|
||||
else OpenAIChatCompletionClient(),
|
||||
)
|
||||
|
||||
# Agentic Chat - basic chat agent
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app=app,
|
||||
agent=simple_agent(client),
|
||||
path="/agentic_chat",
|
||||
)
|
||||
|
||||
# Backend Tool Rendering - agent with tools
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app=app,
|
||||
agent=weather_agent(client),
|
||||
path="/backend_tool_rendering",
|
||||
)
|
||||
|
||||
# Shared State - recipe agent with structured output
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app=app,
|
||||
agent=recipe_agent(client),
|
||||
path="/shared_state",
|
||||
)
|
||||
|
||||
# Predictive State Updates - document writer with predictive state
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app=app,
|
||||
agent=document_writer_agent(client),
|
||||
path="/predictive_state_updates",
|
||||
)
|
||||
|
||||
# Human in the Loop - human-in-the-loop agent with step customization
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app=app,
|
||||
agent=human_in_the_loop_agent(client),
|
||||
path="/human_in_the_loop",
|
||||
state_schema={"steps": {"type": "array"}},
|
||||
predict_state_config={"steps": {"tool": "generate_task_steps", "tool_argument": "steps"}},
|
||||
)
|
||||
|
||||
# Agentic Generative UI - task steps agent with streaming state updates
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app=app,
|
||||
agent=task_steps_agent_wrapped(client),
|
||||
path="/agentic_generative_ui",
|
||||
)
|
||||
|
||||
# Tool-based Generative UI - UI generator with frontend-rendered tools
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app=app,
|
||||
agent=ui_generator_agent(client),
|
||||
path="/tool_based_generative_ui",
|
||||
)
|
||||
|
||||
# Subgraphs - deterministic travel planner with interrupt-driven selections
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app=app,
|
||||
agent=subgraphs_agent(),
|
||||
path="/subgraphs",
|
||||
)
|
||||
|
||||
# Deterministic Tool-Driven State - tool returns state_update() to push snapshot
|
||||
# from actual tool output (see issue #3167).
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app=app,
|
||||
agent=weather_state_agent(client),
|
||||
path="/deterministic_state",
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""Run the server."""
|
||||
port = int(os.getenv("PORT", "8887"))
|
||||
host = os.getenv("HOST", "127.0.0.1")
|
||||
|
||||
print(f"\nAG-UI Examples Server starting on http://{host}:{port}")
|
||||
print("Set ENABLE_DEBUG_LOGGING=1 for detailed request logging\n")
|
||||
|
||||
# Use log_config=None to prevent uvicorn from reconfiguring logging
|
||||
# This preserves our file + console logging setup
|
||||
uvicorn.run(
|
||||
app,
|
||||
host=host,
|
||||
port=port,
|
||||
log_config=None,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,486 @@
|
||||
# Getting Started with AG-UI (Python)
|
||||
|
||||
The AG-UI (Agent UI) protocol provides a standardized way for client applications to interact with AI agents over HTTP. This tutorial demonstrates how to build both server and client applications using the AG-UI protocol with Python.
|
||||
|
||||
## Quick Start - Client Examples
|
||||
|
||||
If you want to quickly try out the AG-UI client, we provide three ready-to-use examples:
|
||||
|
||||
### Basic Interactive Client (`client.py`)
|
||||
|
||||
A simple command-line chat client that demonstrates:
|
||||
- Streaming responses in real-time
|
||||
- Automatic thread management for conversation continuity
|
||||
- Direct `AGUIChatClient` usage (caller manages message history)
|
||||
|
||||
**Run:**
|
||||
```bash
|
||||
python client.py
|
||||
```
|
||||
|
||||
**Note:** This example sends only the current message to the server. The server is responsible for maintaining conversation history using the thread_id.
|
||||
|
||||
### Advanced Features Client (`client_advanced.py`)
|
||||
|
||||
Demonstrates advanced capabilities:
|
||||
- Tool/function calling
|
||||
- Both streaming and non-streaming responses
|
||||
- Multi-turn conversations
|
||||
- Error handling patterns
|
||||
|
||||
**Run:**
|
||||
```bash
|
||||
python client_advanced.py
|
||||
```
|
||||
|
||||
**Note:** This example shows direct `AGUIChatClient` usage. Tool execution and conversation continuity depend on server-side configuration and capabilities.
|
||||
|
||||
### Agent Integration (`client_with_agent.py`)
|
||||
|
||||
Best practice example using `Agent` wrapper with **AgentThread**
|
||||
- **AgentThread** maintains conversation state
|
||||
- Client-side conversation history management via `thread.message_store`
|
||||
- **Hybrid tool execution**: client-side + server-side tools simultaneously
|
||||
- Full conversation history sent on each request
|
||||
- Tool calling with conversation context
|
||||
|
||||
**To demonstrate hybrid tools:**
|
||||
|
||||
1. **Start server with server-side tool** (Terminal 1):
|
||||
```bash
|
||||
# Server has get_time_zone tool
|
||||
python server.py
|
||||
```
|
||||
|
||||
2. **Run client with client-side tool** (Terminal 2):
|
||||
```bash
|
||||
# Client has get_weather tool
|
||||
python client_with_agent.py
|
||||
```
|
||||
|
||||
All examples require a running AG-UI server (see Step 1 below for setup).
|
||||
|
||||
## Understanding AG-UI Architecture
|
||||
|
||||
### Thread Management
|
||||
|
||||
The AG-UI protocol supports two approaches to conversation history:
|
||||
|
||||
1. **Server-Managed Threads** (client.py, client_advanced.py)
|
||||
- Client sends only the current message + thread_id
|
||||
- Server maintains full conversation history
|
||||
- Requires server to support stateful thread storage
|
||||
- Lighter network payload
|
||||
|
||||
2. **Client-Managed History** (client_with_agent.py)
|
||||
- Client maintains full conversation history locally
|
||||
- Full message history sent with each request
|
||||
- Works with any AG-UI server (stateful or stateless)
|
||||
|
||||
The `Agent` wrapper (used in client_with_agent.py) collects messages from local storage and sends the full history to `AGUIChatClient`, which then forwards everything to the server.
|
||||
|
||||
### Tool/Function Calling
|
||||
|
||||
The AG-UI protocol supports **hybrid tool execution** - both client-side AND server-side tools can coexist in the same conversation.
|
||||
|
||||
**The Hybrid Pattern** (client_with_agent.py):
|
||||
```
|
||||
Client defines: Server defines:
|
||||
- get_weather() - get_current_time()
|
||||
- read_sensors() - get_server_forecast()
|
||||
|
||||
User: "What's the weather in SF and what time is it?"
|
||||
↓
|
||||
Agent sends: full history + tool definitions for get_weather, read_sensors
|
||||
↓
|
||||
Server LLM decides: "I need get_weather('SF') and get_current_time()"
|
||||
↓
|
||||
Server executes get_current_time() → "2025-11-11 14:30:00 UTC"
|
||||
Server sends function call request → get_weather('SF')
|
||||
↓
|
||||
Agent intercepts get_weather call → executes locally
|
||||
↓
|
||||
Client sends result → "Sunny, 72°F"
|
||||
↓
|
||||
Server combines both results → "It's sunny and 72°F in SF, and the current time is 2:30 PM UTC"
|
||||
↓
|
||||
Client receives final response
|
||||
```
|
||||
|
||||
**How it works:**
|
||||
|
||||
1. **Client-Side Tools** (`client_with_agent.py`):
|
||||
- Tools defined in Agent's `tools` parameter execute locally
|
||||
- Tool metadata (name, description, schema) sent to server for planning
|
||||
- When server requests client tool → client intercepts → executes locally → sends result
|
||||
|
||||
2. **Server-Side Tools**:
|
||||
- Defined in server agent's configuration
|
||||
- Server executes directly without client involvement
|
||||
- Results included in server's response
|
||||
|
||||
3. **Hybrid Pattern (Both Together)**:
|
||||
- Server LLM sees ALL tool definitions (client + server)
|
||||
- Decides which to use based on task
|
||||
- Server tools execute server-side
|
||||
- Client tools execute client-side
|
||||
|
||||
**Direct AGUIChatClient Usage** (client_advanced.py):
|
||||
Even without Agent wrapper, client-side tools work:
|
||||
- Tools passed in ChatOptions execute locally
|
||||
- Server can also have its own tools
|
||||
- Hybrid execution works automatically
|
||||
|
||||
### Interrupts and Resume Entries
|
||||
|
||||
Human-in-the-loop approvals and workflow input requests pause by emitting a terminal `RUN_FINISHED` event whose
|
||||
`outcome.type` is `"interrupt"`. Generic AG-UI clients should read prompts from `RUN_FINISHED.outcome.interrupts`
|
||||
and resume the same `threadId` with a canonical `resume` array of `ResumeEntry` values.
|
||||
|
||||
```json
|
||||
{
|
||||
"threadId": "thread-1",
|
||||
"messages": [],
|
||||
"resume": [
|
||||
{
|
||||
"interruptId": "approval_1",
|
||||
"status": "resolved",
|
||||
"payload": {
|
||||
"approved": true
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
`Interrupt` and `ResumeEntry` are AG-UI protocol models from `ag_ui.core`; Agent Framework does not define a
|
||||
separate interrupt model. New interrupted runs use `RUN_FINISHED.outcome.interrupts`, not a stable top-level
|
||||
`RUN_FINISHED.interrupt` field.
|
||||
|
||||
## What is AG-UI?
|
||||
|
||||
AG-UI is a protocol that enables:
|
||||
- **Remote agent hosting**: Host AI agents as web services that can be accessed by multiple clients
|
||||
- **Streaming responses**: Real-time streaming of agent responses using Server-Sent Events (SSE)
|
||||
- **Standardized communication**: Consistent message format for agent interactions
|
||||
- **Thread management**: Maintain conversation context across multiple requests
|
||||
- **Advanced features**: Human-in-the-loop, state management, tool rendering
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Before you begin, ensure you have the following:
|
||||
|
||||
- Python 3.10 or later
|
||||
- Azure OpenAI service endpoint and deployment configured
|
||||
- Azure CLI installed and authenticated (for DefaultAzureCredential)
|
||||
- User has the `Cognitive Services OpenAI Contributor` role for the Azure OpenAI resource
|
||||
|
||||
**Note**: These samples use Azure OpenAI models. For more information, see [how to deploy Azure OpenAI models with Azure AI Foundry](https://learn.microsoft.com/en-us/azure/ai-foundry/how-to/deploy-models-openai).
|
||||
|
||||
**Note**: These samples use `DefaultAzureCredential` for authentication. Make sure you're authenticated with Azure (e.g., via `az login`, or environment variables). For more information, see the [Azure Identity documentation](https://learn.microsoft.com/python/api/azure-identity/azure.identity.defaultazurecredential).
|
||||
|
||||
> **Warning**
|
||||
> The AG-UI protocol is still under development and subject to change.
|
||||
> We will keep these samples updated as the protocol evolves.
|
||||
|
||||
## Step 1: Creating an AG-UI Server
|
||||
|
||||
The AG-UI server hosts your AI agent and exposes it via HTTP endpoints using FastAPI.
|
||||
|
||||
### Install Required Packages
|
||||
|
||||
```bash
|
||||
pip install agent-framework-ag-ui
|
||||
```
|
||||
|
||||
Or using uv:
|
||||
|
||||
```bash
|
||||
uv pip install agent-framework-ag-ui
|
||||
```
|
||||
|
||||
### Server Code
|
||||
|
||||
Create a file named `server.py`:
|
||||
|
||||
```python
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AG-UI server example."""
|
||||
|
||||
import os
|
||||
|
||||
from agent_framework import Agent
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Read required configuration
|
||||
endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
|
||||
model = os.environ.get("AZURE_OPENAI_MODEL")
|
||||
api_key = os.environ.get("AZURE_OPENAI_API_KEY")
|
||||
|
||||
if not endpoint:
|
||||
raise ValueError("AZURE_OPENAI_ENDPOINT environment variable is required")
|
||||
if not model:
|
||||
raise ValueError("AZURE_OPENAI_MODEL environment variable is required")
|
||||
if not api_key:
|
||||
raise ValueError("AZURE_OPENAI_API_KEY environment variable is required")
|
||||
|
||||
# Create the AI agent
|
||||
agent = Agent(
|
||||
name="AGUIAssistant",
|
||||
instructions="You are a helpful assistant.",
|
||||
client=OpenAIChatCompletionClient(
|
||||
azure_endpoint=endpoint,
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
),
|
||||
)
|
||||
|
||||
# Create FastAPI app
|
||||
app = FastAPI(title="AG-UI Server")
|
||||
|
||||
# Register the AG-UI endpoint
|
||||
add_agent_framework_fastapi_endpoint(app, agent, "/")
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host="127.0.0.1", port=5100)
|
||||
```
|
||||
|
||||
### Key Concepts
|
||||
|
||||
- **`add_agent_framework_fastapi_endpoint`**: Registers the AG-UI endpoint with automatic request/response handling and SSE streaming
|
||||
- **`Agent`**: The agent that will handle incoming requests
|
||||
- **FastAPI Integration**: Uses FastAPI's native async support for streaming responses
|
||||
- **Instructions**: The agent is created with default instructions, which can be overridden by client messages
|
||||
- **Configuration**: `OpenAIChatCompletionClient` can read from environment variables (`AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_MODEL`, `AZURE_OPENAI_API_KEY`) or accept parameters directly
|
||||
|
||||
**Alternative (simpler)**: Use environment variables only:
|
||||
|
||||
```python
|
||||
# No need to read environment variables manually
|
||||
agent = Agent(
|
||||
name="AGUIAssistant",
|
||||
instructions="You are a helpful assistant.",
|
||||
client=OpenAIChatCompletionClient(), # Reads from environment automatically
|
||||
)
|
||||
```
|
||||
|
||||
### Configure and Run the Server
|
||||
|
||||
Set the required environment variables:
|
||||
|
||||
```bash
|
||||
export AZURE_OPENAI_ENDPOINT="https://your-resource.openai.azure.com/"
|
||||
export AZURE_OPENAI_MODEL="gpt-4o-mini"
|
||||
# Optional: Set API key if not using DefaultAzureCredential
|
||||
# export AZURE_OPENAI_API_KEY="your-api-key"
|
||||
```
|
||||
|
||||
Run the server:
|
||||
|
||||
```bash
|
||||
python server.py
|
||||
```
|
||||
|
||||
Or using uvicorn directly:
|
||||
|
||||
```bash
|
||||
uvicorn server:app --host 127.0.0.1 --port 5100
|
||||
```
|
||||
|
||||
The server will start listening on `http://127.0.0.1:5100`.
|
||||
|
||||
## Step 2: Creating an AG-UI Client
|
||||
|
||||
The AG-UI client connects to the remote server and displays streaming responses. The `AGUIChatClient` is a built-in implementation that integrates with the Agent Framework's standard chat interface.
|
||||
|
||||
### Install Required Packages
|
||||
|
||||
The `AGUIChatClient` is included in the `agent-framework-ag-ui` package (already installed if you installed the server packages).
|
||||
|
||||
```bash
|
||||
pip install agent-framework-ag-ui
|
||||
```
|
||||
|
||||
### Client Code
|
||||
|
||||
Create a file named `client.py`:
|
||||
|
||||
```python
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AG-UI client example using AGUIChatClient."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agent_framework.ag_ui import AGUIChatClient
|
||||
|
||||
|
||||
async def main():
|
||||
"""Main client loop demonstrating AGUIChatClient usage."""
|
||||
# Get server URL from environment or use default
|
||||
server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
|
||||
print(f"Connecting to AG-UI server at: {server_url}\n")
|
||||
|
||||
# Create client with context manager for automatic cleanup
|
||||
async with AGUIChatClient(endpoint=server_url) as client:
|
||||
thread_id: str | None = None
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Get user input
|
||||
message = input("\nUser (:q or quit to exit): ")
|
||||
if not message.strip():
|
||||
print("Request cannot be empty.")
|
||||
continue
|
||||
|
||||
if message.lower() in (":q", "quit"):
|
||||
break
|
||||
|
||||
# Send message and stream the response
|
||||
print("\nAssistant: ", end="", flush=True)
|
||||
|
||||
# Use metadata to maintain conversation continuity
|
||||
metadata = {"thread_id": thread_id} if thread_id else None
|
||||
|
||||
async for update in client.get_response(message, metadata=metadata, stream=True):
|
||||
# Extract thread ID from first update
|
||||
if not thread_id and update.additional_properties:
|
||||
thread_id = update.additional_properties.get("thread_id")
|
||||
if thread_id:
|
||||
print(f"\n[Thread: {thread_id}]")
|
||||
print("Assistant: ", end="", flush=True)
|
||||
|
||||
# Stream text content as it arrives
|
||||
for content in update.contents:
|
||||
if content.type == "text" and content.text:
|
||||
print(content.text, end="", flush=True)
|
||||
|
||||
print() # New line after response
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\nExiting...")
|
||||
except Exception as e:
|
||||
print(f"\nAn error occurred: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
### Key Concepts
|
||||
|
||||
- **`AGUIChatClient`**: Built-in client that implements the Agent Framework's `BaseChatClient` interface
|
||||
- **Automatic Event Handling**: The client automatically converts AG-UI events to Agent Framework types
|
||||
- **Thread Management**: Pass `thread_id` in metadata to maintain conversation context across requests
|
||||
- **Streaming Responses**: Use `get_response(..., stream=True)` for real-time streaming or `get_response(..., stream=False)` for non-streaming
|
||||
- **Context Manager**: Use `async with` for automatic cleanup of HTTP connections
|
||||
- **Standard Interface**: Works with all Agent Framework patterns (Agent, tools, etc.)
|
||||
- **Hybrid Tool Execution**: Supports both client-side and server-side tools executing together in the same conversation
|
||||
|
||||
### Configure and Run the Client
|
||||
|
||||
Optionally set a custom server URL:
|
||||
|
||||
```bash
|
||||
export AGUI_SERVER_URL="http://127.0.0.1:5100/"
|
||||
```
|
||||
|
||||
Run the client (in a separate terminal):
|
||||
|
||||
```bash
|
||||
python client.py
|
||||
```
|
||||
|
||||
## Step 3: Testing the Complete System
|
||||
|
||||
### Expected Output
|
||||
|
||||
```
|
||||
$ python client.py
|
||||
Connecting to AG-UI server at: http://127.0.0.1:5100/
|
||||
|
||||
User (:q or quit to exit): What is the capital of France?
|
||||
|
||||
[Thread: abc123]
|
||||
Assistant: The capital of France is Paris. It is known for its rich history, culture,
|
||||
and iconic landmarks such as the Eiffel Tower and the Louvre Museum.
|
||||
|
||||
User (:q or quit to exit): Tell me a fun fact about space
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Connection Refused
|
||||
|
||||
Ensure the server is running before starting the client:
|
||||
|
||||
```bash
|
||||
# Terminal 1
|
||||
python server.py
|
||||
|
||||
# Terminal 2 (after server starts)
|
||||
python client.py
|
||||
```
|
||||
|
||||
### Authentication Errors
|
||||
|
||||
Make sure you're authenticated with Azure:
|
||||
|
||||
```bash
|
||||
az login
|
||||
```
|
||||
|
||||
Verify you have the correct role assignment on the Azure OpenAI resource.
|
||||
|
||||
### Streaming Not Working
|
||||
|
||||
Check that your client timeout is sufficient:
|
||||
|
||||
```python
|
||||
httpx.AsyncClient(timeout=60.0) # 60 seconds should be enough
|
||||
```
|
||||
|
||||
For long-running agents, increase the timeout accordingly.
|
||||
|
||||
### No Events Received
|
||||
|
||||
Ensure you're using the correct `Accept` header:
|
||||
|
||||
```python
|
||||
headers={"Accept": "text/event-stream"}
|
||||
```
|
||||
|
||||
And parsing SSE format correctly (lines starting with `data: `).
|
||||
|
||||
### Thread Context Lost
|
||||
|
||||
The client automatically manages thread continuity. If context is lost:
|
||||
|
||||
1. Check that `threadId` is being captured from `RUN_STARTED` events
|
||||
2. Ensure the same client instance is used across messages
|
||||
3. Verify the server is receiving the `thread_id` in subsequent requests
|
||||
|
||||
### Event Type Mismatches
|
||||
|
||||
Remember that event types are UPPERCASE with underscores (`RUN_STARTED`, not `run_started`) and field names are camelCase (`threadId`, not `thread_id`).
|
||||
|
||||
### Import Errors
|
||||
|
||||
Make sure all packages are installed:
|
||||
|
||||
```bash
|
||||
pip install agent-framework-ag-ui agent-framework-core fastapi uvicorn httpx
|
||||
```
|
||||
|
||||
Or check your virtual environment is activated:
|
||||
|
||||
```bash
|
||||
source venv/bin/activate # Linux/macOS
|
||||
venv\Scripts\activate # Windows
|
||||
```
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AG-UI client example using AGUIChatClient.
|
||||
|
||||
This example demonstrates how to use the AGUIChatClient to connect to
|
||||
a remote AG-UI server and interact with it using the Agent Framework's
|
||||
standard chat interface.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from typing import cast
|
||||
|
||||
from agent_framework import ChatResponse, ChatResponseUpdate, Message, ResponseStream
|
||||
from agent_framework.ag_ui import AGUIChatClient
|
||||
|
||||
|
||||
async def main():
|
||||
"""Main client loop demonstrating AGUIChatClient usage."""
|
||||
# Get server URL from environment or use default
|
||||
server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
|
||||
print(f"Connecting to AG-UI server at: {server_url}\n")
|
||||
print("Using AGUIChatClient with automatic thread management and Agent Framework integration.\n")
|
||||
|
||||
# Create client with context manager for automatic cleanup
|
||||
async with AGUIChatClient(endpoint=server_url) as client:
|
||||
thread_id: str | None = None
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Get user input
|
||||
message = input("\nUser (:q or quit to exit): ")
|
||||
if not message.strip():
|
||||
print("Request cannot be empty.")
|
||||
continue
|
||||
|
||||
if message.lower() in (":q", "quit"):
|
||||
break
|
||||
|
||||
# Send message and stream the response
|
||||
print("\nAssistant: ", end="", flush=True)
|
||||
|
||||
# Use metadata to maintain conversation continuity
|
||||
metadata = {"thread_id": thread_id} if thread_id else None
|
||||
|
||||
stream = client.get_response(
|
||||
[Message(role="user", contents=[message])],
|
||||
stream=True,
|
||||
options={"metadata": metadata} if metadata else None,
|
||||
)
|
||||
stream = cast(ResponseStream[ChatResponseUpdate, ChatResponse], stream)
|
||||
async for update in stream:
|
||||
# Extract and display thread ID from first update
|
||||
if not thread_id and update.additional_properties:
|
||||
thread_id = update.additional_properties.get("thread_id")
|
||||
if thread_id:
|
||||
print(f"\n\033[93m[Thread: {thread_id}]\033[0m", end="", flush=True)
|
||||
print("\nAssistant: ", end="", flush=True)
|
||||
|
||||
# Display text content as it streams
|
||||
for content in update.contents:
|
||||
if content.type == "text" and content.text:
|
||||
print(f"\033[96m{content.text}\033[0m", end="", flush=True)
|
||||
|
||||
# Display finish reason if present
|
||||
if update.finish_reason:
|
||||
print(f"\n\033[92m[Finished: {update.finish_reason}]\033[0m", end="", flush=True)
|
||||
|
||||
print() # New line after response
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\nExiting...")
|
||||
except Exception as e:
|
||||
print(f"\n\033[91mAn error occurred: {e}\033[0m")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,245 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Advanced AG-UI client example with tools and features.
|
||||
|
||||
This example demonstrates advanced AGUIChatClient features including:
|
||||
- Tool/function calling
|
||||
- Non-streaming responses
|
||||
- Multiple conversation turns
|
||||
- Error handling
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from typing import cast
|
||||
|
||||
from agent_framework import ChatResponse, ChatResponseUpdate, Message, ResponseStream, tool
|
||||
from agent_framework.ag_ui import AGUIChatClient
|
||||
|
||||
|
||||
@tool
|
||||
def get_weather(location: str) -> str:
|
||||
"""Get the current weather for a location.
|
||||
|
||||
Args:
|
||||
location: The city or location name
|
||||
"""
|
||||
# Simulate weather lookup
|
||||
weather_data = {
|
||||
"seattle": "Rainy, 55°F",
|
||||
"san francisco": "Foggy, 62°F",
|
||||
"new york": "Sunny, 68°F",
|
||||
"london": "Cloudy, 52°F",
|
||||
}
|
||||
return weather_data.get(location.lower(), f"Weather data not available for {location}")
|
||||
|
||||
|
||||
@tool
|
||||
def calculate(a: float, b: float, operation: str) -> str:
|
||||
"""Perform basic arithmetic operations.
|
||||
|
||||
Args:
|
||||
a: First number
|
||||
b: Second number
|
||||
operation: Operation to perform (add, subtract, multiply, divide)
|
||||
"""
|
||||
try:
|
||||
if operation == "add":
|
||||
result = a + b
|
||||
elif operation == "subtract":
|
||||
result = a - b
|
||||
elif operation == "multiply":
|
||||
result = a * b
|
||||
elif operation == "divide":
|
||||
result = a / b
|
||||
else:
|
||||
return f"Unsupported operation: {operation}"
|
||||
return f"The result is: {result}"
|
||||
except Exception as e:
|
||||
return f"Error calculating: {e}"
|
||||
|
||||
|
||||
async def streaming_example(client: AGUIChatClient, thread_id: str | None = None):
|
||||
"""Demonstrate streaming responses."""
|
||||
print("\n" + "=" * 60)
|
||||
print("STREAMING EXAMPLE")
|
||||
print("=" * 60)
|
||||
|
||||
metadata = {"thread_id": thread_id} if thread_id else None
|
||||
|
||||
print("\nUser: Tell me a short joke\n")
|
||||
print("Assistant: ", end="", flush=True)
|
||||
|
||||
stream = client.get_response(
|
||||
[Message(role="user", contents=["Tell me a short joke"])],
|
||||
stream=True,
|
||||
options={"metadata": metadata} if metadata else None,
|
||||
)
|
||||
stream = cast(ResponseStream[ChatResponseUpdate, ChatResponse], stream)
|
||||
async for update in stream:
|
||||
if not thread_id and update.additional_properties:
|
||||
thread_id = update.additional_properties.get("thread_id")
|
||||
|
||||
for content in update.contents:
|
||||
if content.type == "text" and content.text: # type: ignore[attr-defined]
|
||||
print(content.text, end="", flush=True) # type: ignore[attr-defined]
|
||||
|
||||
print("\n")
|
||||
return thread_id
|
||||
|
||||
|
||||
async def non_streaming_example(client: AGUIChatClient, thread_id: str | None = None):
|
||||
"""Demonstrate non-streaming responses."""
|
||||
print("\n" + "=" * 60)
|
||||
print("NON-STREAMING EXAMPLE")
|
||||
print("=" * 60)
|
||||
|
||||
metadata = {"thread_id": thread_id} if thread_id else None
|
||||
|
||||
print("\nUser: What is 2 + 2?\n")
|
||||
|
||||
response = await client.get_response([Message(role="user", contents=["What is 2 + 2?"])], metadata=metadata)
|
||||
|
||||
print(f"Assistant: {response.text}")
|
||||
|
||||
if response.additional_properties:
|
||||
thread_id = response.additional_properties.get("thread_id")
|
||||
print(f"\n[Thread: {thread_id}]")
|
||||
|
||||
return thread_id
|
||||
|
||||
|
||||
async def tool_example(client: AGUIChatClient, thread_id: str | None = None):
|
||||
"""Demonstrate sending tool definitions to the server.
|
||||
|
||||
IMPORTANT: When using AGUIChatClient directly (without Agent wrapper):
|
||||
- Tools are sent as DEFINITIONS only
|
||||
- No automatic client-side execution (no function invocation middleware)
|
||||
- Server must have matching tool implementations to execute them
|
||||
|
||||
For CLIENT-SIDE tool execution (like .NET AGUIClient sample):
|
||||
- Use Agent wrapper with tools
|
||||
- See client_with_agent.py for the hybrid pattern
|
||||
- Agent middleware intercepts and executes client tools locally
|
||||
- Server can have its own tools that execute server-side
|
||||
- Both client and server tools work together in same conversation
|
||||
|
||||
This example sends tool definitions and assumes server-side execution.
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TOOL DEFINITION EXAMPLE")
|
||||
print("=" * 60)
|
||||
|
||||
metadata = {"thread_id": thread_id} if thread_id else None
|
||||
|
||||
print("\nUser: What's the weather in Seattle?\n")
|
||||
print("Sending tool definitions to server...")
|
||||
print("(Server must be configured with matching tools to execute them)\n")
|
||||
|
||||
response = await client.get_response(
|
||||
[Message(role="user", contents=["What's the weather in Seattle?"])],
|
||||
tools=[get_weather, calculate],
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
print(f"Assistant: {response.text}")
|
||||
|
||||
# Show tool calls if any
|
||||
tool_called = False
|
||||
for message in response.messages:
|
||||
for content in message.contents:
|
||||
if content.type == "function_call": # type: ignore[attr-defined]
|
||||
print(f"\n[Tool Called: {content.name}]") # type: ignore[attr-defined]
|
||||
tool_called = True
|
||||
|
||||
if not tool_called:
|
||||
print("\n[Note: No tools were called - server may not be configured for tool execution]")
|
||||
|
||||
if response.additional_properties:
|
||||
thread_id = response.additional_properties.get("thread_id")
|
||||
|
||||
return thread_id
|
||||
|
||||
|
||||
async def conversation_example(client: AGUIChatClient):
|
||||
"""Demonstrate multi-turn conversation.
|
||||
|
||||
Note: Conversation continuity depends on the server maintaining thread state.
|
||||
Some servers may require explicit message history to be sent with each request.
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("MULTI-TURN CONVERSATION EXAMPLE")
|
||||
print("=" * 60)
|
||||
print("\nNote: This example uses thread_id for context. Server must support thread-based state.\n")
|
||||
|
||||
# First turn
|
||||
print("User: My name is Alice\n")
|
||||
response1 = await client.get_response([Message(role="user", contents=["My name is Alice"])])
|
||||
print(f"Assistant: {response1.text}")
|
||||
thread_id = response1.additional_properties.get("thread_id")
|
||||
print(f"\n[Thread: {thread_id}]")
|
||||
|
||||
# Second turn - using same thread
|
||||
print("\nUser: What's my name?\n")
|
||||
response2 = await client.get_response(
|
||||
[Message(role="user", contents=["What's my name?"])], options={"metadata": {"thread_id": thread_id}}
|
||||
)
|
||||
print(f"Assistant: {response2.text}")
|
||||
|
||||
# Check if context was maintained
|
||||
if "alice" not in response2.text.lower():
|
||||
print("\n[Note: Server may not maintain thread context - consider using Agent for history management]")
|
||||
|
||||
# Third turn
|
||||
print("\nUser: Can you also tell me what 10 * 5 is?\n")
|
||||
response3 = await client.get_response(
|
||||
[Message(role="user", contents=["Can you also tell me what 10 * 5 is?"])],
|
||||
options={"metadata": {"thread_id": thread_id}},
|
||||
tools=[calculate],
|
||||
)
|
||||
print(f"Assistant: {response3.text}")
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run all examples."""
|
||||
# Get server URL from environment or use default
|
||||
server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
|
||||
|
||||
print("=" * 60)
|
||||
print("AG-UI Chat Client Advanced Examples")
|
||||
print("=" * 60)
|
||||
print(f"\nServer: {server_url}")
|
||||
print("\nThese examples demonstrate various AGUIChatClient features:")
|
||||
print(" 1. Streaming responses")
|
||||
print(" 2. Non-streaming responses")
|
||||
print(" 3. Tool/function calling")
|
||||
print(" 4. Multi-turn conversations")
|
||||
|
||||
try:
|
||||
async with AGUIChatClient(endpoint=server_url) as client:
|
||||
# Run examples in sequence
|
||||
thread_id = await streaming_example(client)
|
||||
thread_id = await non_streaming_example(client, thread_id)
|
||||
await tool_example(client, thread_id)
|
||||
|
||||
# Separate conversation with new thread
|
||||
await conversation_example(client)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("All examples completed successfully!")
|
||||
print("=" * 60)
|
||||
|
||||
except ConnectionError as e:
|
||||
print(f"\n\033[91mConnection Error: {e}\033[0m")
|
||||
print("\nMake sure an AG-UI server is running at the specified endpoint.")
|
||||
except Exception as e:
|
||||
print(f"\n\033[91mError: {e}\033[0m")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,150 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Example showing Agent with AGUIChatClient for hybrid tool execution.
|
||||
|
||||
This demonstrates the HYBRID pattern matching .NET AGUIClient implementation:
|
||||
|
||||
1. AgentSession Pattern (like .NET):
|
||||
- Create session with agent.create_session()
|
||||
- Pass session to agent.run(stream=True) on each turn
|
||||
- Session maintains conversation context via context providers
|
||||
|
||||
2. Hybrid Tool Execution:
|
||||
- AGUIChatClient uses function invocation mixin
|
||||
- Client-side tools (get_weather) can execute locally when server requests them
|
||||
- Server may also have its own tools that execute server-side
|
||||
- Both work together: server LLM decides which tool to call, decorator handles client execution
|
||||
|
||||
This matches .NET pattern: session maintains state, tools execute on appropriate side.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
|
||||
from agent_framework import Agent, tool
|
||||
from agent_framework.ag_ui import AGUIChatClient
|
||||
|
||||
# Enable debug logging
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@tool(description="Get the current weather for a location.")
|
||||
def get_weather(location: str) -> str:
|
||||
"""Get the current weather for a location.
|
||||
|
||||
Args:
|
||||
location: The city or location name
|
||||
"""
|
||||
print(f"[CLIENT] get_weather tool called with location: {location}")
|
||||
weather_data = {
|
||||
"seattle": "Rainy, 55°F",
|
||||
"san francisco": "Foggy, 62°F",
|
||||
"new york": "Sunny, 68°F",
|
||||
"london": "Cloudy, 52°F",
|
||||
}
|
||||
result = weather_data.get(location.lower(), f"Weather data not available for {location}")
|
||||
print(f"[CLIENT] get_weather returning: {result}")
|
||||
return result
|
||||
|
||||
|
||||
async def main():
|
||||
"""Demonstrate Agent + AGUIChatClient hybrid tool execution.
|
||||
|
||||
This matches the .NET pattern from Program.cs where:
|
||||
- AIAgent agent = chatClient.CreateAIAgent(tools: [...])
|
||||
- AgentSession session = agent.CreateSession()
|
||||
- RunStreamingAsync(messages, session)
|
||||
|
||||
Python equivalent:
|
||||
- agent = Agent(client=AGUIChatClient(...), tools=[...])
|
||||
- session = agent.create_session() # Creates session
|
||||
- agent.run(message, stream=True, session=session) # Session tracks context
|
||||
"""
|
||||
server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
|
||||
|
||||
print("=" * 70)
|
||||
print("Agent + AGUIChatClient: Hybrid Tool Execution")
|
||||
print("=" * 70)
|
||||
print(f"\nServer: {server_url}")
|
||||
print("\nThis example demonstrates:")
|
||||
print(" 1. AgentSession maintains conversation state (like .NET)")
|
||||
print(" 2. Client-side tools execute locally via function invocation mixin")
|
||||
print(" 3. Server may have additional tools that execute server-side")
|
||||
print(" 4. HYBRID: Client and server tools work together simultaneously\n")
|
||||
|
||||
try:
|
||||
# Create remote client in async context manager
|
||||
async with AGUIChatClient(endpoint=server_url) as remote_client:
|
||||
# Wrap in Agent for conversation history management
|
||||
agent = Agent(
|
||||
name="remote_assistant",
|
||||
instructions="You are a helpful assistant. Remember user information across the conversation.",
|
||||
client=remote_client,
|
||||
tools=[get_weather],
|
||||
)
|
||||
|
||||
# Create a session to maintain conversation state (like .NET AgentSession)
|
||||
session = agent.create_session()
|
||||
|
||||
print("=" * 70)
|
||||
print("CONVERSATION WITH HISTORY")
|
||||
print("=" * 70)
|
||||
|
||||
# Turn 1: Introduce
|
||||
print("\nUser: My name is Alice and I live in Seattle\n")
|
||||
async for chunk in agent.run("My name is Alice and I live in Seattle", stream=True, session=session):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="", flush=True)
|
||||
print("\n")
|
||||
|
||||
# Turn 2: Ask about name (tests history)
|
||||
print("User: What's my name?\n")
|
||||
async for chunk in agent.run("What's my name?", stream=True, session=session):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="", flush=True)
|
||||
print("\n")
|
||||
|
||||
# Turn 3: Ask about location (tests history)
|
||||
print("User: Where do I live?\n")
|
||||
async for chunk in agent.run("Where do I live?", stream=True, session=session):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="", flush=True)
|
||||
print("\n")
|
||||
|
||||
# Turn 4: Test client-side tool (get_weather is client-side)
|
||||
print("User: What's the weather forecast for today in Seattle?\n")
|
||||
async for chunk in agent.run(
|
||||
"What's the weather forecast for today in Seattle?",
|
||||
stream=True,
|
||||
session=session,
|
||||
):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="", flush=True)
|
||||
print("\n")
|
||||
|
||||
# Turn 5: Test server-side tool (get_time_zone is server-side only)
|
||||
print("User: What time zone is Seattle in?\n")
|
||||
async for chunk in agent.run("What time zone is Seattle in?", stream=True, session=session):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="", flush=True)
|
||||
print("\n")
|
||||
|
||||
except ConnectionError as e:
|
||||
print(f"\n\033[91mConnection Error: {e}\033[0m")
|
||||
print("\nMake sure an AG-UI server is running at the specified endpoint.")
|
||||
except Exception as e:
|
||||
print(f"\n\033[91mError: {e}\033[0m")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,144 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AG-UI server example with server-side tools."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from agent_framework import Agent, tool
|
||||
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
|
||||
from agent_framework.openai import OpenAIChatCompletionClient
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import Depends, FastAPI, HTTPException, Security
|
||||
from fastapi.security import APIKeyHeader
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Enable debug logging
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Read required configuration
|
||||
endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
|
||||
model = os.environ.get("AZURE_OPENAI_MODEL")
|
||||
|
||||
if not endpoint:
|
||||
raise ValueError("AZURE_OPENAI_ENDPOINT environment variable is required")
|
||||
if not model:
|
||||
raise ValueError("AZURE_OPENAI_MODEL environment variable is required")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# AUTHENTICATION EXAMPLE
|
||||
# ============================================================================
|
||||
# This demonstrates how to secure the AG-UI endpoint with API key authentication.
|
||||
# In production, you should use a more robust authentication mechanism such as:
|
||||
# - OAuth 2.0 / OpenID Connect
|
||||
# - JWT tokens with proper validation
|
||||
# - Azure AD / Entra ID integration
|
||||
# - Your organization's identity provider
|
||||
#
|
||||
# The API key should be stored securely (e.g., Azure Key Vault, environment variables)
|
||||
# and rotated regularly.
|
||||
# ============================================================================
|
||||
|
||||
# API key header configuration
|
||||
API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)
|
||||
|
||||
# Get the expected API key from environment variable
|
||||
# In production, use a secrets manager like Azure Key Vault
|
||||
EXPECTED_API_KEY = os.environ.get("AG_UI_API_KEY")
|
||||
|
||||
|
||||
async def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> None:
|
||||
"""Verify the API key provided in the request header.
|
||||
|
||||
Args:
|
||||
api_key: The API key from the X-API-Key header
|
||||
|
||||
Raises:
|
||||
HTTPException: If the API key is missing or invalid
|
||||
"""
|
||||
if not EXPECTED_API_KEY:
|
||||
# If no API key is configured, log a warning but allow the request
|
||||
# This maintains backward compatibility but warns about the security risk
|
||||
logger.warning(
|
||||
"AG_UI_API_KEY environment variable not set. "
|
||||
"The endpoint is accessible without authentication. "
|
||||
"Set AG_UI_API_KEY to enable API key authentication."
|
||||
)
|
||||
return
|
||||
|
||||
if not api_key:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Missing API key. Provide X-API-Key header.",
|
||||
)
|
||||
|
||||
if api_key != EXPECTED_API_KEY:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Invalid API key.",
|
||||
)
|
||||
|
||||
|
||||
# Server-side tool (executes on server)
|
||||
@tool(description="Get the time zone for a location.")
|
||||
def get_time_zone(location: str) -> str:
|
||||
"""Get the time zone for a location.
|
||||
|
||||
Args:
|
||||
location: The city or location name
|
||||
"""
|
||||
print(f"[SERVER] get_time_zone tool called with location: {location}")
|
||||
timezone_data = {
|
||||
"seattle": "Pacific Time (UTC-8)",
|
||||
"san francisco": "Pacific Time (UTC-8)",
|
||||
"new york": "Eastern Time (UTC-5)",
|
||||
"london": "Greenwich Mean Time (UTC+0)",
|
||||
}
|
||||
result = timezone_data.get(location.lower(), f"Time zone data not available for {location}")
|
||||
print(f"[SERVER] get_time_zone returning: {result}")
|
||||
return result
|
||||
|
||||
|
||||
# Create the AI agent with ONLY server-side tools
|
||||
# IMPORTANT: Do NOT include tools that the client provides!
|
||||
# In this example:
|
||||
# - get_time_zone: SERVER-ONLY tool (only server has this)
|
||||
# - get_weather: CLIENT-ONLY tool (client provides this, server should NOT include it)
|
||||
# The client will send get_weather tool metadata so the LLM knows about it,
|
||||
# and the function invocation mixin on AGUIChatClient will execute it client-side.
|
||||
# This matches the .NET AG-UI hybrid execution pattern.
|
||||
agent = Agent(
|
||||
name="AGUIAssistant",
|
||||
instructions="You are a helpful assistant. Use get_weather for weather and get_time_zone for time zones.",
|
||||
client=OpenAIChatCompletionClient(
|
||||
azure_endpoint=endpoint,
|
||||
model=model,
|
||||
),
|
||||
tools=[get_time_zone], # ONLY server-side tools
|
||||
)
|
||||
|
||||
# Create FastAPI app
|
||||
app = FastAPI(title="AG-UI Server")
|
||||
|
||||
# Register the AG-UI endpoint with authentication
|
||||
# The dependencies parameter accepts FastAPI Depends() objects that run before the handler
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app,
|
||||
agent,
|
||||
"/",
|
||||
dependencies=[Depends(verify_api_key)],
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host="127.0.0.1", port=5100, log_level="debug", access_log=True)
|
||||
@@ -0,0 +1,82 @@
|
||||
[project]
|
||||
name = "agent-framework-ag-ui"
|
||||
version = "1.0.0rc8"
|
||||
description = "AG-UI protocol integration for Agent Framework"
|
||||
readme = "README.md"
|
||||
license-files = ["LICENSE"]
|
||||
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
|
||||
requires-python = ">=3.10"
|
||||
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 :: 4 - Beta",
|
||||
"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",
|
||||
"Typing :: Typed",
|
||||
]
|
||||
dependencies = [
|
||||
"agent-framework-core>=1.11.0,<2",
|
||||
"ag-ui-protocol>=0.1.19,<0.2",
|
||||
"fastapi>=0.121.0,<0.138.1",
|
||||
"sse-starlette>=3.4.5,<4",
|
||||
"uvicorn[standard]>=0.30.0,<1"
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest==9.1.1",
|
||||
"httpx==0.28.1",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["agent_framework_ag_ui", "agent_framework_ag_ui_examples"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests/ag_ui"]
|
||||
pythonpath = [".", "tests/ag_ui"]
|
||||
markers = [
|
||||
"integration: marks tests as integration tests that require external services",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
target-version = "py311"
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "N", "W"]
|
||||
ignore = ["E501"]
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.11"
|
||||
warn_return_any = true
|
||||
warn_unused_configs = true
|
||||
disallow_untyped_defs = false
|
||||
|
||||
[tool.pyright]
|
||||
include = ["agent_framework_ag_ui"]
|
||||
exclude = ["tests", "tests/ag_ui", "examples"]
|
||||
typeCheckingMode = "basic"
|
||||
|
||||
[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_ag_ui"
|
||||
|
||||
[tool.poe.tasks.test]
|
||||
help = "Run the default unit test suite for this package."
|
||||
cmd = 'pytest -m "not integration" --cov=agent_framework_ag_ui --cov-report=term-missing:skip-covered -n auto --dist worksteal tests/ag_ui'
|
||||
@@ -0,0 +1,350 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Shared test fixtures and stubs for AG-UI tests."""
|
||||
|
||||
import sys
|
||||
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, MutableSequence, Sequence
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Generic, Literal, TypedDict, cast, overload # noqa: F401
|
||||
|
||||
import pytest
|
||||
from agent_framework import (
|
||||
AgentResponse,
|
||||
AgentResponseUpdate,
|
||||
AgentSession,
|
||||
BaseChatClient,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Content,
|
||||
Message,
|
||||
ServiceSessionId,
|
||||
SupportsAgentRun,
|
||||
SupportsChatGetResponse,
|
||||
)
|
||||
from agent_framework._clients import OptionsCoT
|
||||
from agent_framework._middleware import ChatMiddlewareLayer
|
||||
from agent_framework._tools import FunctionInvocationLayer
|
||||
from agent_framework._types import ResponseStream
|
||||
from agent_framework.observability import ChatTelemetryLayer
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # type: ignore[import] # pragma: no cover
|
||||
|
||||
StreamFn = Callable[..., AsyncIterable[ChatResponseUpdate]]
|
||||
ResponseFn = Callable[..., Awaitable[ChatResponse]]
|
||||
|
||||
|
||||
def pytest_configure() -> None:
|
||||
"""Ensure this test directory is on sys.path so helper modules can be imported by name."""
|
||||
test_dir = str(Path(__file__).resolve().parent)
|
||||
if test_dir not in sys.path:
|
||||
sys.path.insert(0, test_dir)
|
||||
|
||||
|
||||
class StreamingChatClientStub(
|
||||
FunctionInvocationLayer[OptionsCoT],
|
||||
ChatMiddlewareLayer[OptionsCoT],
|
||||
ChatTelemetryLayer[OptionsCoT],
|
||||
BaseChatClient[OptionsCoT],
|
||||
Generic[OptionsCoT],
|
||||
):
|
||||
"""Typed streaming stub that satisfies SupportsChatGetResponse."""
|
||||
|
||||
def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
|
||||
super().__init__(middleware=[])
|
||||
self._stream_fn = stream_fn
|
||||
self._response_fn = response_fn
|
||||
self.last_session: AgentSession | None = None
|
||||
self.last_service_session_id: str | ServiceSessionId | None = None
|
||||
|
||||
@overload
|
||||
def get_response(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
*,
|
||||
stream: Literal[False] = ...,
|
||||
options: ChatOptions[Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse[Any]]: ...
|
||||
|
||||
@overload
|
||||
def get_response(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
*,
|
||||
stream: Literal[False] = ...,
|
||||
options: OptionsCoT | ChatOptions[None] | None = ...,
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse[Any]]: ...
|
||||
|
||||
@overload
|
||||
def get_response(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
*,
|
||||
stream: Literal[True],
|
||||
options: OptionsCoT | ChatOptions[Any] | None = ...,
|
||||
**kwargs: Any,
|
||||
) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...
|
||||
|
||||
def get_response(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
*,
|
||||
stream: bool = False,
|
||||
options: OptionsCoT | ChatOptions[Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
|
||||
client_kwargs = kwargs.get("client_kwargs")
|
||||
if isinstance(client_kwargs, Mapping):
|
||||
self.last_session = cast(AgentSession | None, client_kwargs.get("session"))
|
||||
else:
|
||||
self.last_session = None
|
||||
self.last_service_session_id = self.last_session.service_session_id if self.last_session else None
|
||||
if stream:
|
||||
return super().get_response(
|
||||
messages=messages,
|
||||
stream=True,
|
||||
options=options,
|
||||
**kwargs,
|
||||
)
|
||||
return super().get_response(
|
||||
messages=messages,
|
||||
stream=False,
|
||||
options=options,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@override
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
stream: bool = False,
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
if stream:
|
||||
|
||||
def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
|
||||
return ChatResponse.from_updates(updates)
|
||||
|
||||
return ResponseStream(self._stream_fn(messages, options, **kwargs), finalizer=_finalize)
|
||||
|
||||
return self._get_response_impl(messages, options, **kwargs)
|
||||
|
||||
async def _get_response_impl(
|
||||
self, messages: Sequence[Message], options: Mapping[str, Any], **kwargs: Any
|
||||
) -> ChatResponse:
|
||||
"""Non-streaming implementation."""
|
||||
if self._response_fn is not None:
|
||||
return await self._response_fn(messages, options, **kwargs)
|
||||
|
||||
contents: list[Any] = []
|
||||
async for update in self._stream_fn(list(messages), dict(options), **kwargs):
|
||||
contents.extend(update.contents)
|
||||
|
||||
return ChatResponse(
|
||||
messages=[Message(role="assistant", contents=contents)],
|
||||
response_id="stub-response",
|
||||
)
|
||||
|
||||
|
||||
def stream_from_updates(updates: list[ChatResponseUpdate]) -> StreamFn:
|
||||
"""Create a stream function that yields from a static list of updates."""
|
||||
|
||||
async def _stream(
|
||||
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
|
||||
) -> AsyncIterator[ChatResponseUpdate]:
|
||||
for update in updates:
|
||||
yield update
|
||||
|
||||
return _stream
|
||||
|
||||
|
||||
class StubAgent(SupportsAgentRun):
|
||||
"""Minimal SupportsAgentRun stub for orchestrator tests."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
updates: list[AgentResponseUpdate] | None = None,
|
||||
*,
|
||||
agent_id: str = "stub-agent",
|
||||
agent_name: str | None = "stub-agent",
|
||||
default_options: Any | None = None,
|
||||
client: Any | None = None,
|
||||
) -> None:
|
||||
self.id = agent_id
|
||||
self.name = agent_name
|
||||
self.description = "stub agent"
|
||||
self.updates = updates or [AgentResponseUpdate(contents=[Content.from_text(text="response")], role="assistant")]
|
||||
self.default_options: dict[str, Any] = (
|
||||
default_options if isinstance(default_options, dict) else {"tools": None, "response_format": None}
|
||||
)
|
||||
self.client = client or SimpleNamespace(function_invocation_configuration=None)
|
||||
self.messages_received: list[Any] = []
|
||||
self.tools_received: list[Any] | None = None
|
||||
self.last_session: AgentSession | None = None
|
||||
|
||||
@overload
|
||||
def run(
|
||||
self,
|
||||
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
|
||||
*,
|
||||
stream: Literal[False] = ...,
|
||||
session: AgentSession | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[AgentResponse[Any]]: ...
|
||||
|
||||
@overload
|
||||
def run(
|
||||
self,
|
||||
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
|
||||
*,
|
||||
stream: Literal[True],
|
||||
session: AgentSession | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
|
||||
|
||||
def run(
|
||||
self,
|
||||
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
|
||||
*,
|
||||
stream: bool = False,
|
||||
session: AgentSession | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
|
||||
if stream:
|
||||
|
||||
async def _stream() -> AsyncIterator[AgentResponseUpdate]:
|
||||
if messages is None:
|
||||
self.messages_received = []
|
||||
elif isinstance(messages, (str, Content, Message)):
|
||||
self.messages_received = [messages]
|
||||
else:
|
||||
self.messages_received = list(messages)
|
||||
self.last_session = session
|
||||
self.tools_received = kwargs.get("tools")
|
||||
for update in self.updates:
|
||||
yield update
|
||||
|
||||
def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse:
|
||||
return AgentResponse.from_updates(updates)
|
||||
|
||||
return ResponseStream(_stream(), finalizer=_finalize)
|
||||
|
||||
async def _get_response() -> AgentResponse[Any]:
|
||||
return AgentResponse(messages=[], response_id="stub-response")
|
||||
|
||||
return _get_response()
|
||||
|
||||
def create_session(self, **kwargs: Any) -> AgentSession:
|
||||
return AgentSession(session_id=kwargs.get("session_id"))
|
||||
|
||||
def get_session(self, service_session_id: str | ServiceSessionId, *, session_id: str | None = None) -> AgentSession:
|
||||
return AgentSession(session_id=session_id, service_session_id=service_session_id)
|
||||
|
||||
|
||||
# Fixtures
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def streaming_chat_client_stub() -> type[SupportsChatGetResponse]:
|
||||
"""Return the StreamingChatClientStub class for creating test instances."""
|
||||
return StreamingChatClientStub # type: ignore[return-value]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stream_from_updates_fixture() -> Callable[[list[ChatResponseUpdate]], StreamFn]:
|
||||
"""Return the stream_from_updates helper function."""
|
||||
return stream_from_updates
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stub_agent() -> type[SupportsAgentRun]:
|
||||
"""Return the StubAgent class for creating test instances."""
|
||||
return StubAgent # type: ignore[return-value]
|
||||
|
||||
|
||||
# ── Fixtures for golden / integration tests ──
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def collect_events() -> Callable[..., Any]:
|
||||
"""Return an async helper that collects all events from an async generator."""
|
||||
|
||||
async def _collect(async_gen: AsyncIterable[Any]) -> list[Any]:
|
||||
return [event async for event in async_gen]
|
||||
|
||||
return _collect
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def make_agent_wrapper() -> Callable[..., Any]:
|
||||
"""Factory that builds an AgentFrameworkAgent from a stream function.
|
||||
|
||||
Usage::
|
||||
|
||||
agent = make_agent_wrapper(
|
||||
stream_fn=stream_from_updates(updates),
|
||||
state_schema=...,
|
||||
)
|
||||
events = [e async for e in agent.run(payload)]
|
||||
"""
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
def _factory(
|
||||
stream_fn: StreamFn,
|
||||
*,
|
||||
state_schema: Any | None = None,
|
||||
predict_state_config: dict[str, dict[str, str]] | None = None,
|
||||
require_confirmation: bool = True,
|
||||
) -> Any:
|
||||
client = StreamingChatClientStub(stream_fn)
|
||||
stub = StubAgent(client=client)
|
||||
return AgentFrameworkAgent(
|
||||
agent=stub,
|
||||
state_schema=state_schema,
|
||||
predict_state_config=predict_state_config,
|
||||
require_confirmation=require_confirmation,
|
||||
)
|
||||
|
||||
return _factory
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def make_app() -> Callable[..., Any]:
|
||||
"""Factory that builds a FastAPI app with an AG-UI endpoint.
|
||||
|
||||
Usage::
|
||||
|
||||
app = make_app(agent_or_wrapper, path="/test")
|
||||
"""
|
||||
from fastapi import FastAPI
|
||||
|
||||
from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
|
||||
|
||||
def _factory(
|
||||
agent: Any,
|
||||
*,
|
||||
path: str = "/",
|
||||
state_schema: Any | None = None,
|
||||
predict_state_config: dict[str, dict[str, str]] | None = None,
|
||||
default_state: dict[str, Any] | None = None,
|
||||
) -> FastAPI:
|
||||
app = FastAPI()
|
||||
add_agent_framework_fastapi_endpoint(
|
||||
app,
|
||||
agent,
|
||||
path=path,
|
||||
state_schema=state_schema,
|
||||
predict_state_config=predict_state_config,
|
||||
default_state=default_state,
|
||||
)
|
||||
return app
|
||||
|
||||
return _factory
|
||||
@@ -0,0 +1,210 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""EventStream assertion helper for AG-UI regression tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
class EventStream:
|
||||
"""Wraps a list of AG-UI events with structured assertion methods.
|
||||
|
||||
Usage:
|
||||
events = [event async for event in agent.run(payload)]
|
||||
stream = EventStream(events)
|
||||
stream.assert_bookends()
|
||||
stream.assert_text_messages_balanced()
|
||||
"""
|
||||
|
||||
def __init__(self, events: list[Any]) -> None:
|
||||
self.events = events
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.events)
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.events)
|
||||
|
||||
def types(self) -> list[str]:
|
||||
"""Return ordered list of event type strings."""
|
||||
return [self._type_str(e) for e in self.events]
|
||||
|
||||
def get(self, event_type: str) -> list[Any]:
|
||||
"""Filter events matching the given type string."""
|
||||
return [e for e in self.events if self._type_str(e) == event_type]
|
||||
|
||||
def first(self, event_type: str) -> Any:
|
||||
"""Return the first event matching the given type, or raise."""
|
||||
matches = self.get(event_type)
|
||||
if not matches:
|
||||
raise ValueError(f"No event of type {event_type!r} found. Available: {self.types()}")
|
||||
return matches[0]
|
||||
|
||||
def last(self, event_type: str) -> Any:
|
||||
"""Return the last event matching the given type, or raise."""
|
||||
matches = self.get(event_type)
|
||||
if not matches:
|
||||
raise ValueError(f"No event of type {event_type!r} found. Available: {self.types()}")
|
||||
return matches[-1]
|
||||
|
||||
def snapshot(self) -> dict[str, Any]:
|
||||
"""Return the latest StateSnapshotEvent snapshot dict."""
|
||||
return self.last("STATE_SNAPSHOT").snapshot
|
||||
|
||||
def messages_snapshot(self) -> list[Any]:
|
||||
"""Return the latest MessagesSnapshotEvent messages list."""
|
||||
return self.last("MESSAGES_SNAPSHOT").messages
|
||||
|
||||
def run_finished_interrupts(self, event: Any | None = None) -> list[dict[str, Any]]:
|
||||
"""Return canonical interrupts from a RUN_FINISHED event."""
|
||||
target = event or self.last("RUN_FINISHED")
|
||||
dumped = self._event_dump(target)
|
||||
assert "interrupt" not in dumped
|
||||
outcome = dumped.get("outcome")
|
||||
assert isinstance(outcome, dict), f"Expected RUN_FINISHED.outcome, got {dumped}"
|
||||
assert outcome.get("type") == "interrupt"
|
||||
interrupts = outcome.get("interrupts")
|
||||
assert isinstance(interrupts, list), f"Expected outcome.interrupts, got {outcome}"
|
||||
return interrupts
|
||||
|
||||
@staticmethod
|
||||
def interrupt_metadata_value(interrupt: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return Agent Framework interruption details from canonical interrupt metadata."""
|
||||
metadata = interrupt.get("metadata")
|
||||
assert isinstance(metadata, dict)
|
||||
agent_framework_metadata = metadata.get("agent_framework")
|
||||
assert isinstance(agent_framework_metadata, dict)
|
||||
value = agent_framework_metadata.get("value")
|
||||
assert isinstance(value, dict)
|
||||
return value
|
||||
|
||||
# ── Structural assertions ──
|
||||
|
||||
def assert_bookends(self) -> None:
|
||||
"""Assert first event is RUN_STARTED and last is RUN_FINISHED."""
|
||||
types = self.types()
|
||||
assert types, "Event stream is empty"
|
||||
assert types[0] == "RUN_STARTED", f"Expected RUN_STARTED first, got {types[0]}"
|
||||
assert types[-1] == "RUN_FINISHED", f"Expected RUN_FINISHED last, got {types[-1]}"
|
||||
|
||||
def assert_has_run_lifecycle(self) -> None:
|
||||
"""Assert RUN_STARTED is first and RUN_FINISHED exists (may not be last).
|
||||
|
||||
Use this instead of assert_bookends() for workflow resume streams where
|
||||
_drain_open_message() can emit TEXT_MESSAGE_END after RUN_FINISHED.
|
||||
"""
|
||||
types = self.types()
|
||||
assert types, "Event stream is empty"
|
||||
assert types[0] == "RUN_STARTED", f"Expected RUN_STARTED first, got {types[0]}"
|
||||
assert "RUN_FINISHED" in types, f"Expected RUN_FINISHED in stream. Types: {types}"
|
||||
|
||||
def assert_strict_types(self, expected: list[str]) -> None:
|
||||
"""Assert exact type sequence match."""
|
||||
actual = self.types()
|
||||
assert actual == expected, f"Event type mismatch.\nExpected: {expected}\nActual: {actual}"
|
||||
|
||||
def assert_ordered_types(self, expected: list[str]) -> None:
|
||||
"""Assert expected types appear as a subsequence (in order, not necessarily contiguous)."""
|
||||
actual = self.types()
|
||||
actual_idx = 0
|
||||
for expected_type in expected:
|
||||
found = False
|
||||
while actual_idx < len(actual):
|
||||
if actual[actual_idx] == expected_type:
|
||||
actual_idx += 1
|
||||
found = True
|
||||
break
|
||||
actual_idx += 1
|
||||
if not found:
|
||||
raise AssertionError(
|
||||
f"Expected subsequence type {expected_type!r} not found after index {actual_idx}.\n"
|
||||
f"Expected subsequence: {expected}\n"
|
||||
f"Actual types: {actual}"
|
||||
)
|
||||
|
||||
def assert_text_messages_balanced(self) -> None:
|
||||
"""Assert every TEXT_MESSAGE_START has a matching TEXT_MESSAGE_END with the same message_id."""
|
||||
starts: dict[str, int] = {}
|
||||
ends: set[str] = set()
|
||||
for i, event in enumerate(self.events):
|
||||
t = self._type_str(event)
|
||||
if t == "TEXT_MESSAGE_START":
|
||||
mid = event.message_id
|
||||
assert mid not in starts, f"Duplicate TEXT_MESSAGE_START for message_id={mid}"
|
||||
starts[mid] = i
|
||||
elif t == "TEXT_MESSAGE_END":
|
||||
mid = event.message_id
|
||||
assert mid in starts, f"TEXT_MESSAGE_END for unknown message_id={mid}"
|
||||
assert mid not in ends, f"Duplicate TEXT_MESSAGE_END for message_id={mid}"
|
||||
ends.add(mid)
|
||||
|
||||
unclosed = set(starts.keys()) - ends
|
||||
assert not unclosed, f"Unclosed text messages: {unclosed}"
|
||||
|
||||
def assert_tool_calls_balanced(self) -> None:
|
||||
"""Assert every TOOL_CALL_START has a matching TOOL_CALL_END with the same tool_call_id."""
|
||||
starts: dict[str, int] = {}
|
||||
ends: set[str] = set()
|
||||
for i, event in enumerate(self.events):
|
||||
t = self._type_str(event)
|
||||
if t == "TOOL_CALL_START":
|
||||
tid = event.tool_call_id
|
||||
assert tid not in starts, f"Duplicate TOOL_CALL_START for tool_call_id={tid}"
|
||||
starts[tid] = i
|
||||
elif t == "TOOL_CALL_END":
|
||||
tid = event.tool_call_id
|
||||
assert tid in starts, f"TOOL_CALL_END for unknown tool_call_id={tid}"
|
||||
assert tid not in ends, f"Duplicate TOOL_CALL_END for tool_call_id={tid}"
|
||||
ends.add(tid)
|
||||
|
||||
unclosed = set(starts.keys()) - ends
|
||||
assert not unclosed, f"Unclosed tool calls: {unclosed}"
|
||||
|
||||
def assert_no_run_error(self) -> None:
|
||||
"""Assert no RUN_ERROR events exist."""
|
||||
errors = self.get("RUN_ERROR")
|
||||
if errors:
|
||||
messages = [getattr(e, "message", str(e)) for e in errors]
|
||||
raise AssertionError(f"Found {len(errors)} RUN_ERROR event(s): {messages}")
|
||||
|
||||
def assert_has_type(self, event_type: str) -> None:
|
||||
"""Assert at least one event of the given type exists."""
|
||||
assert event_type in self.types(), f"Expected {event_type!r} in stream. Available: {self.types()}"
|
||||
|
||||
def assert_message_ids_consistent(self) -> None:
|
||||
"""Assert TEXT_MESSAGE_CONTENT events reference valid, open message_ids."""
|
||||
open_messages: set[str] = set()
|
||||
for event in self.events:
|
||||
t = self._type_str(event)
|
||||
if t == "TEXT_MESSAGE_START":
|
||||
open_messages.add(event.message_id)
|
||||
elif t == "TEXT_MESSAGE_END":
|
||||
open_messages.discard(event.message_id)
|
||||
elif t == "TEXT_MESSAGE_CONTENT":
|
||||
mid = event.message_id
|
||||
assert mid in open_messages, f"TEXT_MESSAGE_CONTENT references message_id={mid} which is not open"
|
||||
|
||||
# ── Internal ──
|
||||
|
||||
@staticmethod
|
||||
def _type_str(event: Any) -> str:
|
||||
"""Extract event type as a plain string."""
|
||||
t = getattr(event, "type", None)
|
||||
if t is None:
|
||||
return type(event).__name__
|
||||
if isinstance(t, str):
|
||||
return t
|
||||
return getattr(t, "value", str(t))
|
||||
|
||||
@staticmethod
|
||||
def _event_dump(event: Any) -> dict[str, Any]:
|
||||
"""Serialize an event object or return a raw event dict."""
|
||||
if isinstance(event, dict):
|
||||
return event
|
||||
if hasattr(event, "model_dump"):
|
||||
return event.model_dump(by_alias=True, exclude_none=True)
|
||||
raw = getattr(event, "raw", None)
|
||||
if isinstance(raw, dict):
|
||||
return raw
|
||||
raise TypeError(f"Unsupported event type: {type(event).__name__}")
|
||||
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Conftest for golden tests — ensures parent test dir is importable."""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def pytest_configure() -> None:
|
||||
"""Ensure parent test directory is on sys.path for helper module imports."""
|
||||
parent_test_dir = str(Path(__file__).resolve().parent.parent)
|
||||
if parent_test_dir not in sys.path:
|
||||
sys.path.insert(0, parent_test_dir)
|
||||
@@ -0,0 +1,140 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Golden event-stream tests for the basic agentic chat scenario."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentResponseUpdate, Content
|
||||
from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
|
||||
def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
|
||||
stub = StubAgent(updates=updates)
|
||||
return AgentFrameworkAgent(agent=stub, **kwargs)
|
||||
|
||||
|
||||
async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in agent.run(payload)])
|
||||
|
||||
|
||||
BASIC_PAYLOAD: dict[str, Any] = {
|
||||
"thread_id": "thread-chat",
|
||||
"run_id": "run-chat",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
}
|
||||
|
||||
|
||||
def _text_update(text: str) -> AgentResponseUpdate:
|
||||
return AgentResponseUpdate(contents=[Content.from_text(text=text)], role="assistant")
|
||||
|
||||
|
||||
def _snapshot_role(msg: Any) -> str:
|
||||
"""Extract role string from a snapshot message (Pydantic model or dict)."""
|
||||
role = getattr(msg, "role", None) or (msg.get("role") if isinstance(msg, dict) else None)
|
||||
if role is None:
|
||||
return ""
|
||||
return str(getattr(role, "value", role))
|
||||
|
||||
|
||||
def _snapshot_content(msg: Any) -> str:
|
||||
"""Extract content string from a snapshot message."""
|
||||
content = getattr(msg, "content", None) or (msg.get("content") if isinstance(msg, dict) else "")
|
||||
return str(content) if content else ""
|
||||
|
||||
|
||||
# ── Golden stream tests ──
|
||||
|
||||
|
||||
async def test_basic_chat_golden_event_sequence() -> None:
|
||||
"""Assert the exact event type sequence for a single text response."""
|
||||
agent = _build_agent([_text_update("Hi there!")])
|
||||
stream = await _run(agent, BASIC_PAYLOAD)
|
||||
|
||||
stream.assert_strict_types(
|
||||
[
|
||||
"RUN_STARTED",
|
||||
"TEXT_MESSAGE_START",
|
||||
"TEXT_MESSAGE_CONTENT",
|
||||
"TEXT_MESSAGE_END",
|
||||
"MESSAGES_SNAPSHOT",
|
||||
"RUN_FINISHED",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def test_basic_chat_bookends() -> None:
|
||||
"""RUN_STARTED is first, RUN_FINISHED is last."""
|
||||
agent = _build_agent([_text_update("reply")])
|
||||
stream = await _run(agent, BASIC_PAYLOAD)
|
||||
stream.assert_bookends()
|
||||
|
||||
|
||||
async def test_basic_chat_text_messages_balanced() -> None:
|
||||
"""Every TEXT_MESSAGE_START has a matching TEXT_MESSAGE_END."""
|
||||
agent = _build_agent([_text_update("reply")])
|
||||
stream = await _run(agent, BASIC_PAYLOAD)
|
||||
stream.assert_text_messages_balanced()
|
||||
|
||||
|
||||
async def test_basic_chat_no_errors() -> None:
|
||||
"""No RUN_ERROR events in a normal flow."""
|
||||
agent = _build_agent([_text_update("reply")])
|
||||
stream = await _run(agent, BASIC_PAYLOAD)
|
||||
stream.assert_no_run_error()
|
||||
|
||||
|
||||
async def test_basic_chat_message_id_consistency() -> None:
|
||||
"""All text events reference the same message_id."""
|
||||
agent = _build_agent([_text_update("reply")])
|
||||
stream = await _run(agent, BASIC_PAYLOAD)
|
||||
|
||||
start = stream.first("TEXT_MESSAGE_START")
|
||||
content = stream.first("TEXT_MESSAGE_CONTENT")
|
||||
end = stream.first("TEXT_MESSAGE_END")
|
||||
assert start.message_id == content.message_id == end.message_id
|
||||
|
||||
|
||||
async def test_multi_chunk_text_golden_sequence() -> None:
|
||||
"""Streaming multiple chunks produces START + multiple CONTENT + END."""
|
||||
agent = _build_agent([_text_update("Hello "), _text_update("world!")])
|
||||
stream = await _run(agent, BASIC_PAYLOAD)
|
||||
|
||||
stream.assert_strict_types(
|
||||
[
|
||||
"RUN_STARTED",
|
||||
"TEXT_MESSAGE_START",
|
||||
"TEXT_MESSAGE_CONTENT",
|
||||
"TEXT_MESSAGE_CONTENT",
|
||||
"TEXT_MESSAGE_END",
|
||||
"MESSAGES_SNAPSHOT",
|
||||
"RUN_FINISHED",
|
||||
]
|
||||
)
|
||||
stream.assert_text_messages_balanced()
|
||||
stream.assert_message_ids_consistent()
|
||||
|
||||
|
||||
async def test_messages_snapshot_contains_assistant_reply() -> None:
|
||||
"""MessagesSnapshotEvent includes the assistant's accumulated text."""
|
||||
agent = _build_agent([_text_update("Hello there")])
|
||||
stream = await _run(agent, BASIC_PAYLOAD)
|
||||
|
||||
snapshot = stream.messages_snapshot()
|
||||
assistant_msgs = [m for m in snapshot if _snapshot_role(m) == "assistant"]
|
||||
assert assistant_msgs, "No assistant message in snapshot"
|
||||
assert any("Hello there" in _snapshot_content(m) for m in assistant_msgs)
|
||||
|
||||
|
||||
async def test_empty_messages_produces_start_and_finish() -> None:
|
||||
"""Empty message list still produces RUN_STARTED and RUN_FINISHED."""
|
||||
agent = _build_agent([_text_update("reply")])
|
||||
payload = {"thread_id": "t1", "run_id": "r1", "messages": []}
|
||||
stream = await _run(agent, payload)
|
||||
|
||||
stream.assert_bookends()
|
||||
assert "TEXT_MESSAGE_START" not in stream.types()
|
||||
@@ -0,0 +1,236 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Golden event-stream tests for the backend (server-side) tools scenario."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentResponseUpdate, Content
|
||||
from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
|
||||
def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
|
||||
stub = StubAgent(updates=updates)
|
||||
return AgentFrameworkAgent(agent=stub, **kwargs)
|
||||
|
||||
|
||||
async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in agent.run(payload)])
|
||||
|
||||
|
||||
PAYLOAD: dict[str, Any] = {
|
||||
"thread_id": "thread-tools",
|
||||
"run_id": "run-tools",
|
||||
"messages": [{"role": "user", "content": "What's the weather?"}],
|
||||
}
|
||||
|
||||
|
||||
# ── Golden stream tests ──
|
||||
|
||||
|
||||
async def test_tool_call_lifecycle_golden_sequence() -> None:
|
||||
"""Assert the full event sequence for a tool call → result → text response."""
|
||||
updates = [
|
||||
# LLM calls the tool
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city": "SF"}')],
|
||||
role="assistant",
|
||||
),
|
||||
# Tool result comes back
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-1", result="72°F and sunny")],
|
||||
role="assistant",
|
||||
),
|
||||
# LLM responds with text
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="It's 72°F and sunny in SF!")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_ordered_types(
|
||||
[
|
||||
"RUN_STARTED",
|
||||
"TEXT_MESSAGE_START", # Synthetic start for tool-only message
|
||||
"TOOL_CALL_START",
|
||||
"TOOL_CALL_ARGS",
|
||||
"TOOL_CALL_END",
|
||||
"TOOL_CALL_RESULT",
|
||||
"TEXT_MESSAGE_END", # End of synthetic message
|
||||
"TEXT_MESSAGE_START", # New message for text response
|
||||
"TEXT_MESSAGE_CONTENT",
|
||||
"TEXT_MESSAGE_END",
|
||||
"MESSAGES_SNAPSHOT",
|
||||
"RUN_FINISHED",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def test_tool_calls_balanced() -> None:
|
||||
"""Every TOOL_CALL_START has a matching TOOL_CALL_END."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city": "SF"}')],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-1", result="72°F")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="It's 72°F!")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_tool_calls_balanced()
|
||||
|
||||
|
||||
async def test_text_messages_balanced_with_tools() -> None:
|
||||
"""Text messages are properly balanced even around tool calls."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city": "SF"}')],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-1", result="72°F")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="It's 72°F!")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_text_messages_balanced()
|
||||
|
||||
|
||||
async def test_tool_call_id_matches_result() -> None:
|
||||
"""TOOL_CALL_START and TOOL_CALL_RESULT reference the same tool_call_id."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments="{}")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-1", result="72°F")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
start = stream.first("TOOL_CALL_START")
|
||||
result = stream.first("TOOL_CALL_RESULT")
|
||||
assert start.tool_call_id == result.tool_call_id == "call-1"
|
||||
|
||||
|
||||
async def test_tool_result_content_preserved() -> None:
|
||||
"""TOOL_CALL_RESULT event carries the tool's result content."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments="{}")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-1", result="72°F and sunny")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
result = stream.first("TOOL_CALL_RESULT")
|
||||
assert result.content == "72°F and sunny"
|
||||
|
||||
|
||||
async def test_no_run_error_on_tool_flow() -> None:
|
||||
"""Tool call flow doesn't produce RUN_ERROR."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments="{}")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-1", result="72°F")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_no_run_error()
|
||||
stream.assert_bookends()
|
||||
|
||||
|
||||
async def test_multiple_sequential_tool_calls() -> None:
|
||||
"""Multiple sequential tool calls each produce balanced START/END pairs."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="tool_a", call_id="call-a", arguments="{}")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-a", result="result-a")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="tool_b", call_id="call-b", arguments="{}")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-b", result="result-b")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Done!")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_tool_calls_balanced()
|
||||
stream.assert_text_messages_balanced()
|
||||
stream.assert_bookends()
|
||||
|
||||
# Both tool calls should appear
|
||||
starts = stream.get("TOOL_CALL_START")
|
||||
assert len(starts) == 2
|
||||
assert {s.tool_call_name for s in starts} == {"tool_a", "tool_b"}
|
||||
|
||||
|
||||
async def test_messages_snapshot_includes_tool_calls() -> None:
|
||||
"""MessagesSnapshotEvent includes tool call and result messages."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city":"SF"}')],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-1", result="72°F")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="It's warm!")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_has_type("MESSAGES_SNAPSHOT")
|
||||
snapshot = stream.messages_snapshot()
|
||||
# Should have: user message, assistant with tool_calls, tool result, assistant text
|
||||
assert len(snapshot) >= 3
|
||||
@@ -0,0 +1,364 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Golden event-stream tests for the deterministic tool-driven state scenario.
|
||||
|
||||
Covers issue https://github.com/microsoft/agent-framework/issues/3167 — a tool
|
||||
returning :func:`agent_framework_ag_ui.state_update` must push a deterministic
|
||||
``StateSnapshotEvent`` derived from its actual return value, orthogonal to the
|
||||
optimistic ``predict_state_config`` path. These golden tests pin the user-visible
|
||||
event stream so additive changes cannot silently regress it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentResponseUpdate, Content
|
||||
from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent, state_update
|
||||
|
||||
STATE_SCHEMA = {
|
||||
"weather": {"type": "object", "description": "Last fetched weather"},
|
||||
}
|
||||
|
||||
|
||||
def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
|
||||
stub = StubAgent(updates=updates)
|
||||
kwargs.setdefault("state_schema", STATE_SCHEMA)
|
||||
return AgentFrameworkAgent(agent=stub, **kwargs)
|
||||
|
||||
|
||||
async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in agent.run(payload)])
|
||||
|
||||
|
||||
PAYLOAD: dict[str, Any] = {
|
||||
"thread_id": "thread-det-state",
|
||||
"run_id": "run-det-state",
|
||||
"messages": [{"role": "user", "content": "What's the weather in SF?"}],
|
||||
"state": {"weather": {}},
|
||||
}
|
||||
|
||||
|
||||
def _tool_call(call_id: str, name: str, arguments: str) -> AgentResponseUpdate:
|
||||
return AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name=name, call_id=call_id, arguments=arguments)],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
|
||||
def _tool_result_with_state(call_id: str, text: str, state: dict[str, Any]) -> AgentResponseUpdate:
|
||||
"""Build a function_result update whose inner item carries a state marker.
|
||||
|
||||
This mirrors what the core framework produces when a real ``@tool`` returns
|
||||
:func:`state_update`: ``parse_result`` keeps the ``Content`` as-is, and
|
||||
``Content.from_function_result`` preserves its ``additional_properties``
|
||||
inside ``items``.
|
||||
"""
|
||||
return AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_result(
|
||||
call_id=call_id,
|
||||
result=[state_update(text=text, state=state)],
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
|
||||
def _tool_result_with_display(call_id: str, text: str, tool_result: Any, **kwargs: Any) -> AgentResponseUpdate:
|
||||
"""Build a function_result update carrying an optional UI display marker."""
|
||||
return AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_result(
|
||||
call_id=call_id,
|
||||
result=[state_update(text=text, tool_result=tool_result, **kwargs)],
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
|
||||
# ── Golden stream tests ──
|
||||
|
||||
|
||||
async def test_deterministic_state_emits_snapshot_after_tool_result() -> None:
|
||||
"""The happy path: STATE_SNAPSHOT follows TOOL_CALL_RESULT in order."""
|
||||
updates = [
|
||||
_tool_call("call-1", "get_weather", '{"city": "SF"}'),
|
||||
_tool_result_with_state(
|
||||
"call-1",
|
||||
text="Weather in SF: 14°C foggy",
|
||||
state={"weather": {"city": "SF", "temp": 14, "conditions": "foggy"}},
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="It's 14°C and foggy in SF.")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
stream.assert_tool_calls_balanced()
|
||||
stream.assert_text_messages_balanced()
|
||||
|
||||
# Ordered subsequence: the deterministic STATE_SNAPSHOT must follow the
|
||||
# TOOL_CALL_RESULT. This is the central contract for #3167.
|
||||
stream.assert_ordered_types(
|
||||
[
|
||||
"RUN_STARTED",
|
||||
"TOOL_CALL_START",
|
||||
"TOOL_CALL_ARGS",
|
||||
"TOOL_CALL_END",
|
||||
"TOOL_CALL_RESULT",
|
||||
"STATE_SNAPSHOT",
|
||||
"RUN_FINISHED",
|
||||
]
|
||||
)
|
||||
|
||||
# The final STATE_SNAPSHOT must carry the tool-driven state.
|
||||
snapshot = stream.snapshot()
|
||||
assert snapshot["weather"] == {"city": "SF", "temp": 14, "conditions": "foggy"}
|
||||
|
||||
|
||||
async def test_deterministic_state_does_not_fire_for_plain_tool_result() -> None:
|
||||
"""Regression guard: tools returning plain strings must NOT emit a new STATE_SNAPSHOT.
|
||||
|
||||
The initial STATE_SNAPSHOT fires once from the schema + initial payload
|
||||
state. A plain (non-state_update) tool result must not add another one.
|
||||
"""
|
||||
updates = [
|
||||
_tool_call("call-1", "get_weather", '{"city": "SF"}'),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_result(call_id="call-1", result="14°C foggy")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="It's 14°C and foggy.")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
snapshots = stream.get("STATE_SNAPSHOT")
|
||||
# Only the initial snapshot (from state_schema + payload state) should exist.
|
||||
# No deterministic snapshot should have been added by the plain tool result.
|
||||
assert len(snapshots) == 1, (
|
||||
f"Expected exactly 1 STATE_SNAPSHOT (initial only) for plain tool result; "
|
||||
f"got {len(snapshots)}. Snapshots: {[s.snapshot for s in snapshots]}"
|
||||
)
|
||||
|
||||
|
||||
async def test_deterministic_state_merges_into_initial_state() -> None:
|
||||
"""The tool-driven snapshot must merge into, not replace, pre-existing state keys."""
|
||||
payload = dict(PAYLOAD)
|
||||
payload["state"] = {"weather": {}, "user_preferences": {"unit": "C"}}
|
||||
|
||||
updates = [
|
||||
_tool_call("call-1", "get_weather", '{"city": "SF"}'),
|
||||
_tool_result_with_state(
|
||||
"call-1",
|
||||
text="Weather: 14°C",
|
||||
state={"weather": {"city": "SF", "temp": 14}},
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates, state_schema={**STATE_SCHEMA, "user_preferences": {"type": "object"}})
|
||||
stream = await _run(agent, payload)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
final_snapshot = stream.snapshot()
|
||||
assert final_snapshot["weather"] == {"city": "SF", "temp": 14}
|
||||
assert final_snapshot["user_preferences"] == {"unit": "C"}, (
|
||||
"Pre-existing state keys must survive the deterministic merge"
|
||||
)
|
||||
|
||||
|
||||
async def test_deterministic_state_llm_visible_text_is_clean() -> None:
|
||||
"""The LLM-visible TOOL_CALL_RESULT content must not leak the state marker key."""
|
||||
updates = [
|
||||
_tool_call("call-1", "get_weather", '{"city": "SF"}'),
|
||||
_tool_result_with_state(
|
||||
"call-1",
|
||||
text="Weather in SF: 14°C foggy",
|
||||
state={"weather": {"city": "SF", "temp": 14}},
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
result = stream.first("TOOL_CALL_RESULT")
|
||||
assert result.content == "Weather in SF: 14°C foggy"
|
||||
# The marker key must never appear in the content sent back to the LLM.
|
||||
assert "__ag_ui_tool_result_state__" not in result.content
|
||||
assert "weather" not in result.content # not as a raw state dump
|
||||
|
||||
|
||||
async def test_deterministic_state_multiple_tools_merge_in_order() -> None:
|
||||
"""Two state-updating tools in one run merge in order; later wins on key collisions."""
|
||||
updates = [
|
||||
_tool_call("call-a", "get_weather", '{"city": "SF"}'),
|
||||
_tool_result_with_state(
|
||||
"call-a",
|
||||
text="First result",
|
||||
state={"weather": {"city": "SF", "temp": 14}, "source": "primary"},
|
||||
),
|
||||
_tool_call("call-b", "get_weather_refined", '{"city": "SF"}'),
|
||||
_tool_result_with_state(
|
||||
"call-b",
|
||||
text="Refined result",
|
||||
state={"source": "refined"},
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Here you go.")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(
|
||||
updates,
|
||||
state_schema={**STATE_SCHEMA, "source": {"type": "string"}},
|
||||
)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_tool_calls_balanced()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
# Two tool-driven snapshots emitted (one per tool) plus the initial snapshot.
|
||||
snapshots = stream.get("STATE_SNAPSHOT")
|
||||
assert len(snapshots) >= 2, f"Expected at least 2 STATE_SNAPSHOTs; got {len(snapshots)}"
|
||||
|
||||
final = stream.snapshot()
|
||||
assert final["weather"] == {"city": "SF", "temp": 14}
|
||||
# Later tool must override earlier tool on the shared key.
|
||||
assert final["source"] == "refined"
|
||||
|
||||
|
||||
async def test_deterministic_state_coexists_with_predict_state_config() -> None:
|
||||
"""Predictive state and deterministic state must coexist without clobbering each other."""
|
||||
predict_config = {
|
||||
"draft": {
|
||||
"tool": "write_draft",
|
||||
"tool_argument": "body",
|
||||
}
|
||||
}
|
||||
updates = [
|
||||
# Predictive tool: its argument "body" populates state.draft optimistically.
|
||||
_tool_call("call-1", "write_draft", '{"body": "Hello world"}'),
|
||||
# Then a deterministic tool result landing a different key.
|
||||
_tool_result_with_state(
|
||||
"call-1",
|
||||
text="Draft saved",
|
||||
state={"weather": {"city": "SF", "temp": 14}},
|
||||
),
|
||||
]
|
||||
agent = _build_agent(
|
||||
updates,
|
||||
state_schema={**STATE_SCHEMA, "draft": {"type": "string"}},
|
||||
predict_state_config=predict_config,
|
||||
require_confirmation=False,
|
||||
)
|
||||
payload = dict(PAYLOAD)
|
||||
payload["state"] = {"weather": {}, "draft": ""}
|
||||
stream = await _run(agent, payload)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
stream.assert_tool_calls_balanced()
|
||||
|
||||
# The final observed state must contain both the deterministic and predictive contributions.
|
||||
final = stream.snapshot()
|
||||
assert final["weather"] == {"city": "SF", "temp": 14}, f"Deterministic state missing from final snapshot: {final}"
|
||||
|
||||
|
||||
async def test_tool_result_display_payload_reaches_ui_event_only() -> None:
|
||||
"""Rich display payload overrides TOOL_CALL_RESULT without leaking marker keys."""
|
||||
updates = [
|
||||
_tool_call("call-1", "get_weather", '{"city": "SF"}'),
|
||||
_tool_result_with_display(
|
||||
"call-1",
|
||||
text="Weather in SF: 14°C foggy",
|
||||
tool_result={"city": "SF", "temp": 14, "conditions": "foggy"},
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
stream.assert_tool_calls_balanced()
|
||||
|
||||
result = stream.first("TOOL_CALL_RESULT")
|
||||
assert result.content == '{"city": "SF", "temp": 14, "conditions": "foggy"}'
|
||||
assert "__ag_ui_tool_result_display__" not in result.content
|
||||
assert "__ag_ui_tool_result_state__" not in result.content
|
||||
|
||||
|
||||
async def test_tool_result_display_falls_back_to_text_when_unset() -> None:
|
||||
"""Without a display marker, the UI event keeps the existing text content."""
|
||||
updates = [
|
||||
_tool_call("call-1", "get_weather", '{"city": "SF"}'),
|
||||
_tool_result_with_state(
|
||||
"call-1",
|
||||
text="Weather in SF: 14°C foggy",
|
||||
state={"weather": {"city": "SF", "temp": 14}},
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
stream.assert_tool_calls_balanced()
|
||||
|
||||
result = stream.first("TOOL_CALL_RESULT")
|
||||
assert result.content == "Weather in SF: 14°C foggy"
|
||||
assert "__ag_ui_tool_result_display__" not in result.content
|
||||
assert "__ag_ui_tool_result_state__" not in result.content
|
||||
|
||||
|
||||
async def test_tool_result_display_coexists_with_state_snapshot() -> None:
|
||||
"""Display and durable state markers produce one deterministic state snapshot."""
|
||||
updates = [
|
||||
_tool_call("call-1", "get_weather", '{"city": "SF"}'),
|
||||
_tool_result_with_display(
|
||||
"call-1",
|
||||
text="Weather in SF: 14°C foggy",
|
||||
tool_result={"city": "SF", "temp": 14, "conditions": "foggy"},
|
||||
state={"weather": {"city": "SF", "temp": 14, "conditions": "foggy"}},
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
stream.assert_tool_calls_balanced()
|
||||
stream.assert_ordered_types(["TOOL_CALL_RESULT", "STATE_SNAPSHOT", "RUN_FINISHED"])
|
||||
|
||||
result = stream.first("TOOL_CALL_RESULT")
|
||||
assert result.content == '{"city": "SF", "temp": 14, "conditions": "foggy"}'
|
||||
|
||||
result_idx = stream.events.index(result)
|
||||
deterministic_snapshots = [
|
||||
event
|
||||
for event in stream.events[result_idx + 1 :]
|
||||
if getattr(getattr(event, "type", None), "value", getattr(event, "type", None)) == "STATE_SNAPSHOT"
|
||||
]
|
||||
assert len(deterministic_snapshots) == 1
|
||||
assert deterministic_snapshots[0].snapshot["weather"] == {
|
||||
"city": "SF",
|
||||
"temp": 14,
|
||||
"conditions": "foggy",
|
||||
}
|
||||
assert "__ag_ui_tool_result_display__" not in str(deterministic_snapshots[0].snapshot)
|
||||
assert "__ag_ui_tool_result_state__" not in str(deterministic_snapshots[0].snapshot)
|
||||
@@ -0,0 +1,90 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Golden event-stream tests for the generative UI (workflow-as-agent) scenario."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import WorkflowBuilder, WorkflowContext, executor
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkWorkflow
|
||||
|
||||
|
||||
async def _run(wrapper: AgentFrameworkWorkflow, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in wrapper.run(payload)])
|
||||
|
||||
|
||||
PAYLOAD: dict[str, Any] = {
|
||||
"thread_id": "thread-gen-ui-agent",
|
||||
"run_id": "run-gen-ui-agent",
|
||||
"messages": [{"role": "user", "content": "Generate a UI"}],
|
||||
}
|
||||
|
||||
|
||||
# ── Golden stream tests ──
|
||||
|
||||
|
||||
async def test_workflow_agent_golden_sequence() -> None:
|
||||
"""Workflow-as-agent: emits step events and text content."""
|
||||
|
||||
@executor(id="generator")
|
||||
async def generator(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output("Here is your generated UI content!")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=generator).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
stream.assert_text_messages_balanced()
|
||||
|
||||
# Should have step events for the executor
|
||||
stream.assert_has_type("STEP_STARTED")
|
||||
stream.assert_has_type("STEP_FINISHED")
|
||||
|
||||
# Should have text message content
|
||||
stream.assert_has_type("TEXT_MESSAGE_CONTENT")
|
||||
|
||||
|
||||
async def test_workflow_agent_step_names_match() -> None:
|
||||
"""Step started/finished events reference the executor name."""
|
||||
|
||||
@executor(id="my_executor")
|
||||
async def my_executor(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output("Done!")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=my_executor).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, PAYLOAD)
|
||||
|
||||
started = [e for e in stream.get("STEP_STARTED") if getattr(e, "step_name", "") == "my_executor"]
|
||||
finished = [e for e in stream.get("STEP_FINISHED") if getattr(e, "step_name", "") == "my_executor"]
|
||||
assert started, "Expected STEP_STARTED for 'my_executor'"
|
||||
assert finished, "Expected STEP_FINISHED for 'my_executor'"
|
||||
|
||||
|
||||
async def test_workflow_agent_ordered_events() -> None:
|
||||
"""Workflow events follow expected ordering: RUN_STARTED → STEP_STARTED → content → STEP_FINISHED → RUN_FINISHED."""
|
||||
|
||||
@executor(id="my_step")
|
||||
async def my_step(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output("Generated content")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=my_step).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, PAYLOAD)
|
||||
|
||||
stream.assert_ordered_types(
|
||||
[
|
||||
"RUN_STARTED",
|
||||
"STEP_STARTED",
|
||||
"TEXT_MESSAGE_START",
|
||||
"TEXT_MESSAGE_CONTENT",
|
||||
"STEP_FINISHED",
|
||||
"TEXT_MESSAGE_END",
|
||||
"RUN_FINISHED",
|
||||
]
|
||||
)
|
||||
@@ -0,0 +1,135 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Golden event-stream tests for the client-side (declaration-only) tools scenario."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentResponseUpdate, Content
|
||||
from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
|
||||
def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
|
||||
stub = StubAgent(updates=updates)
|
||||
return AgentFrameworkAgent(agent=stub, **kwargs)
|
||||
|
||||
|
||||
async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in agent.run(payload)])
|
||||
|
||||
|
||||
PAYLOAD: dict[str, Any] = {
|
||||
"thread_id": "thread-gen-ui-tool",
|
||||
"run_id": "run-gen-ui-tool",
|
||||
"messages": [{"role": "user", "content": "Show me a chart"}],
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "render_chart",
|
||||
"description": "Render a chart in the UI",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"data": {"type": "array"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
# ── Golden stream tests ──
|
||||
|
||||
|
||||
async def test_declaration_only_tool_golden_sequence() -> None:
|
||||
"""Declaration-only tool: TOOL_CALL_START/ARGS emitted, TOOL_CALL_END at stream end."""
|
||||
# The LLM calls a client-side tool (no server-side execution)
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
name="render_chart",
|
||||
call_id="call-chart",
|
||||
arguments='{"data": [1, 2, 3]}',
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
# Tool call start and args should be present
|
||||
stream.assert_has_type("TOOL_CALL_START")
|
||||
stream.assert_has_type("TOOL_CALL_ARGS")
|
||||
|
||||
# TOOL_CALL_END should be emitted (via get_pending_without_end)
|
||||
stream.assert_has_type("TOOL_CALL_END")
|
||||
stream.assert_tool_calls_balanced()
|
||||
|
||||
|
||||
async def test_declaration_only_tool_no_tool_call_result() -> None:
|
||||
"""Declaration-only tools should NOT produce TOOL_CALL_RESULT events."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
name="render_chart",
|
||||
call_id="call-chart",
|
||||
arguments='{"data": [1, 2, 3]}',
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
assert "TOOL_CALL_RESULT" not in stream.types(), "Declaration-only tools should not have TOOL_CALL_RESULT"
|
||||
|
||||
|
||||
async def test_declaration_only_tool_text_messages_balanced() -> None:
|
||||
"""Text messages remain balanced even with declaration-only tools."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
name="render_chart",
|
||||
call_id="call-chart",
|
||||
arguments='{"data": [1, 2, 3]}',
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_text_messages_balanced()
|
||||
|
||||
|
||||
async def test_declaration_only_tool_messages_snapshot() -> None:
|
||||
"""MessagesSnapshotEvent includes the tool call for declaration-only tools."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
name="render_chart",
|
||||
call_id="call-chart",
|
||||
arguments='{"data": [1, 2, 3]}',
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_has_type("MESSAGES_SNAPSHOT")
|
||||
@@ -0,0 +1,194 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Golden event-stream tests for the HITL (human-in-the-loop) approval scenario."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentResponseUpdate, Content
|
||||
from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
PREDICT_CONFIG = {
|
||||
"tasks": {
|
||||
"tool": "generate_task_steps",
|
||||
"tool_argument": "steps",
|
||||
}
|
||||
}
|
||||
|
||||
STATE_SCHEMA = {
|
||||
"tasks": {"type": "array", "items": {"type": "object"}},
|
||||
}
|
||||
|
||||
|
||||
def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
|
||||
stub = StubAgent(updates=updates)
|
||||
return AgentFrameworkAgent(
|
||||
agent=stub,
|
||||
state_schema=STATE_SCHEMA,
|
||||
predict_state_config=PREDICT_CONFIG,
|
||||
require_confirmation=True,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in agent.run(payload)])
|
||||
|
||||
|
||||
STEPS = [
|
||||
{"description": "Step 1: Plan", "status": "enabled"},
|
||||
{"description": "Step 2: Execute", "status": "enabled"},
|
||||
]
|
||||
|
||||
|
||||
PAYLOAD: dict[str, Any] = {
|
||||
"thread_id": "thread-hitl",
|
||||
"run_id": "run-hitl",
|
||||
"messages": [{"role": "user", "content": "Plan my tasks"}],
|
||||
"state": {"tasks": []},
|
||||
}
|
||||
|
||||
|
||||
# ── Turn 1: Tool call → confirm_changes → interrupt ──
|
||||
|
||||
|
||||
async def test_hitl_turn1_golden_sequence() -> None:
|
||||
"""Turn 1 emits tool call, confirm_changes, and finishes with interrupt."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
name="generate_task_steps",
|
||||
call_id="call-steps",
|
||||
arguments=json.dumps({"steps": STEPS}),
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
# Should have: tool call start/args/end for the primary tool,
|
||||
# then TOOL_CALL_END, STATE_SNAPSHOT, confirm_changes cycle
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
# confirm_changes tool call should be present
|
||||
tool_starts = stream.get("TOOL_CALL_START")
|
||||
tool_names = [getattr(s, "tool_call_name", None) for s in tool_starts]
|
||||
assert "generate_task_steps" in tool_names
|
||||
assert "confirm_changes" in tool_names
|
||||
|
||||
# RUN_FINISHED should have interrupt metadata
|
||||
interrupt = stream.run_finished_interrupts()
|
||||
assert len(interrupt) > 0
|
||||
|
||||
|
||||
async def test_hitl_turn1_tool_calls_balanced() -> None:
|
||||
"""All tool calls in turn 1 (primary + confirm_changes) are balanced."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
name="generate_task_steps",
|
||||
call_id="call-steps",
|
||||
arguments=json.dumps({"steps": STEPS}),
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_tool_calls_balanced()
|
||||
|
||||
|
||||
async def test_hitl_turn1_text_messages_balanced() -> None:
|
||||
"""Text messages are balanced even in the approval flow."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
name="generate_task_steps",
|
||||
call_id="call-steps",
|
||||
arguments=json.dumps({"steps": STEPS}),
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_text_messages_balanced()
|
||||
|
||||
|
||||
# ── Turn 2: Resume with approval → confirmation message → no interrupt ──
|
||||
|
||||
|
||||
async def test_hitl_turn2_resume_with_approval() -> None:
|
||||
"""Resuming with confirm_changes result emits confirmation text and finishes cleanly."""
|
||||
# Turn 2: user sends confirm_changes result as resume
|
||||
# The agent wrapper sees a confirm_changes response and emits a confirmation message
|
||||
confirm_result = json.dumps(
|
||||
{
|
||||
"accepted": True,
|
||||
"steps": STEPS,
|
||||
}
|
||||
)
|
||||
|
||||
# Build payload with resume containing the approval
|
||||
# For confirm_changes, the messages should include the tool result
|
||||
payload: dict[str, Any] = {
|
||||
"thread_id": "thread-hitl",
|
||||
"run_id": "run-hitl-2",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Plan my tasks"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "confirm-id-1",
|
||||
"type": "function",
|
||||
"function": {"name": "confirm_changes", "arguments": json.dumps({"steps": STEPS})},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"toolCallId": "confirm-id-1",
|
||||
"content": confirm_result,
|
||||
},
|
||||
],
|
||||
"state": {"tasks": []},
|
||||
}
|
||||
|
||||
# In turn 2, the agent sees the confirm_changes result and emits a confirmation text
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="Tasks confirmed!")],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, payload)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_text_messages_balanced()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
# Should have text message content (the confirmation message)
|
||||
text_events = stream.get("TEXT_MESSAGE_CONTENT")
|
||||
assert text_events, "Expected confirmation text message"
|
||||
|
||||
# RUN_FINISHED should NOT have interrupt (approval completed)
|
||||
finished = stream.last("RUN_FINISHED")
|
||||
dumped = finished.model_dump(by_alias=True, exclude_none=True)
|
||||
assert "outcome" not in dumped, f"Expected no interrupt after approval, got {dumped.get('outcome')}"
|
||||
@@ -0,0 +1,130 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Golden event-stream tests for the predictive state scenario."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentResponseUpdate, Content
|
||||
from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
PREDICT_CONFIG = {
|
||||
"document": {
|
||||
"tool": "update_document",
|
||||
"tool_argument": "content",
|
||||
}
|
||||
}
|
||||
|
||||
STATE_SCHEMA = {
|
||||
"document": {"type": "string"},
|
||||
}
|
||||
|
||||
|
||||
def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
|
||||
stub = StubAgent(updates=updates)
|
||||
return AgentFrameworkAgent(
|
||||
agent=stub,
|
||||
state_schema=STATE_SCHEMA,
|
||||
predict_state_config=PREDICT_CONFIG,
|
||||
require_confirmation=False,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in agent.run(payload)])
|
||||
|
||||
|
||||
PAYLOAD: dict[str, Any] = {
|
||||
"thread_id": "thread-predict",
|
||||
"run_id": "run-predict",
|
||||
"messages": [{"role": "user", "content": "Write a document"}],
|
||||
"state": {"document": ""},
|
||||
}
|
||||
|
||||
|
||||
# ── Golden stream tests ──
|
||||
|
||||
|
||||
async def test_predictive_state_emits_deltas_during_tool_args() -> None:
|
||||
"""STATE_DELTA events are emitted as tool arguments stream in."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="update_document", call_id="call-1", arguments="")],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(name="update_document", call_id="call-1", arguments='{"content": "Hello')
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_function_call(name="update_document", call_id="call-1", arguments=' world"}')],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
# PredictState custom event should be present
|
||||
custom_events = stream.get("CUSTOM")
|
||||
predict_events = [e for e in custom_events if getattr(e, "name", None) == "PredictState"]
|
||||
assert predict_events, "Expected PredictState custom event"
|
||||
|
||||
# STATE_DELTA events should be emitted during tool arg streaming
|
||||
assert "STATE_DELTA" in stream.types(), "Expected STATE_DELTA events during predictive streaming"
|
||||
|
||||
|
||||
async def test_predictive_state_snapshot_after_tool_end() -> None:
|
||||
"""STATE_SNAPSHOT is emitted when a predictive tool completes (no confirmation)."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
name="update_document", call_id="call-1", arguments='{"content": "Final text"}'
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
|
||||
# Should have initial state snapshot + updated snapshot after tool completion
|
||||
snapshots = stream.get("STATE_SNAPSHOT")
|
||||
assert len(snapshots) >= 1, "Expected at least one STATE_SNAPSHOT"
|
||||
|
||||
|
||||
async def test_predictive_state_ordered_events() -> None:
|
||||
"""Event ordering: RUN_STARTED → PredictState → STATE_SNAPSHOT → TOOL_CALL_* → STATE_SNAPSHOT → RUN_FINISHED."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(name="update_document", call_id="call-1", arguments='{"content": "doc"}')
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_ordered_types(
|
||||
[
|
||||
"RUN_STARTED",
|
||||
"CUSTOM", # PredictState
|
||||
"STATE_SNAPSHOT", # Initial state
|
||||
"TOOL_CALL_START",
|
||||
"TOOL_CALL_ARGS",
|
||||
"RUN_FINISHED",
|
||||
]
|
||||
)
|
||||
@@ -0,0 +1,110 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Golden event-stream tests for the shared state (structured output) scenario."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentResponseUpdate, Content
|
||||
from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkAgent
|
||||
|
||||
|
||||
class RecipeState(BaseModel):
|
||||
recipe_title: str = ""
|
||||
ingredients: list[str] = []
|
||||
message: str = ""
|
||||
|
||||
|
||||
def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
|
||||
stub = StubAgent(
|
||||
updates=updates,
|
||||
default_options={"tools": None, "response_format": RecipeState},
|
||||
)
|
||||
return AgentFrameworkAgent(
|
||||
agent=stub,
|
||||
state_schema={
|
||||
"recipe_title": {"type": "string"},
|
||||
"ingredients": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in agent.run(payload)])
|
||||
|
||||
|
||||
PAYLOAD: dict[str, Any] = {
|
||||
"thread_id": "thread-state",
|
||||
"run_id": "run-state",
|
||||
"messages": [{"role": "user", "content": "Give me a pasta recipe"}],
|
||||
"state": {"recipe_title": "", "ingredients": []},
|
||||
}
|
||||
|
||||
|
||||
# ── Golden stream tests ──
|
||||
|
||||
|
||||
async def test_shared_state_emits_state_snapshot() -> None:
|
||||
"""Structured output agent emits STATE_SNAPSHOT with parsed model fields."""
|
||||
# The structured output agent gets a response that the framework parses as RecipeState
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_text(
|
||||
text='{"recipe_title": "Pasta Carbonara", "ingredients": ["pasta", "eggs", "cheese"], "message": "Here is your recipe!"}'
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
# Should have STATE_SNAPSHOT with the initial state at minimum
|
||||
stream.assert_has_type("STATE_SNAPSHOT")
|
||||
|
||||
|
||||
async def test_shared_state_initial_snapshot_on_first_update() -> None:
|
||||
"""When state_schema and state are provided, initial STATE_SNAPSHOT is emitted after RUN_STARTED."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[Content.from_text(text='{"recipe_title": "Test", "ingredients": [], "message": "hi"}')],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
# RUN_STARTED should be followed by STATE_SNAPSHOT (initial state)
|
||||
stream.assert_ordered_types(["RUN_STARTED", "STATE_SNAPSHOT"])
|
||||
|
||||
|
||||
async def test_shared_state_text_emitted_from_message_field() -> None:
|
||||
"""Structured output's 'message' field is emitted as text message events."""
|
||||
updates = [
|
||||
AgentResponseUpdate(
|
||||
contents=[
|
||||
Content.from_text(
|
||||
text='{"recipe_title": "Pasta", "ingredients": ["pasta"], "message": "Enjoy your pasta!"}'
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
),
|
||||
]
|
||||
agent = _build_agent(updates)
|
||||
stream = await _run(agent, PAYLOAD)
|
||||
|
||||
# Text should be emitted from the message field
|
||||
text_contents = stream.get("TEXT_MESSAGE_CONTENT")
|
||||
if text_contents:
|
||||
combined = "".join(getattr(e, "delta", "") for e in text_contents)
|
||||
assert "Enjoy your pasta!" in combined
|
||||
@@ -0,0 +1,207 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Golden event-stream tests for the workflow HITL (subgraphs) scenario.
|
||||
|
||||
Extends the existing test_subgraphs_example_agent.py with EventStream assertions
|
||||
on full event ordering, balancing, and interrupt structure.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
from agent_framework_ag_ui_examples.agents.subgraphs_agent import subgraphs_agent
|
||||
|
||||
|
||||
async def _run(agent: Any, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in agent.run(payload)])
|
||||
|
||||
|
||||
# ── Turn 1: Initial request → flight interrupt ──
|
||||
|
||||
|
||||
async def test_subgraphs_turn1_golden_bookends() -> None:
|
||||
"""Turn 1 starts with RUN_STARTED and ends with RUN_FINISHED."""
|
||||
agent = subgraphs_agent()
|
||||
stream = await _run(
|
||||
agent,
|
||||
{
|
||||
"thread_id": "thread-sub-golden-1",
|
||||
"run_id": "run-1",
|
||||
"messages": [{"role": "user", "content": "Plan a trip to San Francisco"}],
|
||||
},
|
||||
)
|
||||
stream.assert_bookends()
|
||||
|
||||
|
||||
async def test_subgraphs_turn1_no_errors() -> None:
|
||||
"""Turn 1 completes without errors."""
|
||||
agent = subgraphs_agent()
|
||||
stream = await _run(
|
||||
agent,
|
||||
{
|
||||
"thread_id": "thread-sub-golden-2",
|
||||
"run_id": "run-1",
|
||||
"messages": [{"role": "user", "content": "Plan a trip"}],
|
||||
},
|
||||
)
|
||||
stream.assert_no_run_error()
|
||||
|
||||
|
||||
async def test_subgraphs_turn1_has_step_events() -> None:
|
||||
"""Turn 1 emits STEP_STARTED and STEP_FINISHED for workflow executors."""
|
||||
agent = subgraphs_agent()
|
||||
stream = await _run(
|
||||
agent,
|
||||
{
|
||||
"thread_id": "thread-sub-golden-3",
|
||||
"run_id": "run-1",
|
||||
"messages": [{"role": "user", "content": "Plan a trip"}],
|
||||
},
|
||||
)
|
||||
stream.assert_has_type("STEP_STARTED")
|
||||
stream.assert_has_type("STEP_FINISHED")
|
||||
|
||||
|
||||
async def test_subgraphs_turn1_interrupt_structure() -> None:
|
||||
"""Turn 1 RUN_FINISHED carries flight interrupt with correct structure."""
|
||||
agent = subgraphs_agent()
|
||||
stream = await _run(
|
||||
agent,
|
||||
{
|
||||
"thread_id": "thread-sub-golden-4",
|
||||
"run_id": "run-1",
|
||||
"messages": [{"role": "user", "content": "Plan a trip to SF"}],
|
||||
},
|
||||
)
|
||||
|
||||
interrupt = stream.run_finished_interrupts()
|
||||
assert len(interrupt) > 0
|
||||
interrupt_value = stream.interrupt_metadata_value(interrupt[0])
|
||||
assert interrupt_value["agent"] == "flights"
|
||||
assert len(interrupt_value["options"]) == 2
|
||||
|
||||
|
||||
async def test_subgraphs_turn1_text_messages_balanced() -> None:
|
||||
"""All text messages in turn 1 are properly balanced."""
|
||||
agent = subgraphs_agent()
|
||||
stream = await _run(
|
||||
agent,
|
||||
{
|
||||
"thread_id": "thread-sub-golden-5",
|
||||
"run_id": "run-1",
|
||||
"messages": [{"role": "user", "content": "Plan a trip"}],
|
||||
},
|
||||
)
|
||||
stream.assert_text_messages_balanced()
|
||||
|
||||
|
||||
async def test_subgraphs_turn1_ordered_flow() -> None:
|
||||
"""Turn 1 event ordering: RUN_STARTED → STATE_SNAPSHOT → STEP_* → TOOL_CALL_* → RUN_FINISHED."""
|
||||
agent = subgraphs_agent()
|
||||
stream = await _run(
|
||||
agent,
|
||||
{
|
||||
"thread_id": "thread-sub-golden-6",
|
||||
"run_id": "run-1",
|
||||
"messages": [{"role": "user", "content": "Plan a trip"}],
|
||||
},
|
||||
)
|
||||
stream.assert_ordered_types(
|
||||
[
|
||||
"RUN_STARTED",
|
||||
"STATE_SNAPSHOT",
|
||||
"STEP_STARTED",
|
||||
"RUN_FINISHED",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
# ── Multi-turn: Flight selection → hotel interrupt → completion ──
|
||||
|
||||
|
||||
async def test_subgraphs_full_flow_event_ordering() -> None:
|
||||
"""Complete 3-turn flow maintains proper event ordering throughout."""
|
||||
agent = subgraphs_agent()
|
||||
thread_id = "thread-sub-golden-full"
|
||||
|
||||
# Turn 1
|
||||
stream1 = await _run(
|
||||
agent,
|
||||
{
|
||||
"thread_id": thread_id,
|
||||
"run_id": "run-1",
|
||||
"messages": [{"role": "user", "content": "Plan a trip to SF from Amsterdam"}],
|
||||
},
|
||||
)
|
||||
stream1.assert_bookends()
|
||||
stream1.assert_no_run_error()
|
||||
|
||||
# Extract flight interrupt
|
||||
interrupt1 = stream1.run_finished_interrupts()[0]
|
||||
|
||||
# Turn 2: Select flight
|
||||
stream2 = await _run(
|
||||
agent,
|
||||
{
|
||||
"thread_id": thread_id,
|
||||
"run_id": "run-2",
|
||||
"resume": {
|
||||
"interrupts": [
|
||||
{
|
||||
"id": interrupt1["id"],
|
||||
"value": json.dumps(
|
||||
{
|
||||
"airline": "United",
|
||||
"departure": "Amsterdam (AMS)",
|
||||
"arrival": "San Francisco (SFO)",
|
||||
"price": "$720",
|
||||
"duration": "12h 15m",
|
||||
}
|
||||
),
|
||||
}
|
||||
]
|
||||
},
|
||||
},
|
||||
)
|
||||
stream2.assert_bookends()
|
||||
stream2.assert_no_run_error()
|
||||
|
||||
# Should now have hotel interrupt
|
||||
interrupt2 = stream2.run_finished_interrupts()
|
||||
assert stream2.interrupt_metadata_value(interrupt2[0])["agent"] == "hotels"
|
||||
|
||||
# Turn 3: Select hotel
|
||||
stream3 = await _run(
|
||||
agent,
|
||||
{
|
||||
"thread_id": thread_id,
|
||||
"run_id": "run-3",
|
||||
"resume": {
|
||||
"interrupts": [
|
||||
{
|
||||
"id": interrupt2[0]["id"],
|
||||
"value": json.dumps(
|
||||
{
|
||||
"name": "The Ritz-Carlton",
|
||||
"location": "Nob Hill",
|
||||
"price_per_night": "$550/night",
|
||||
"rating": "4.8 stars",
|
||||
}
|
||||
),
|
||||
}
|
||||
]
|
||||
},
|
||||
},
|
||||
)
|
||||
stream3.assert_bookends()
|
||||
stream3.assert_no_run_error()
|
||||
stream3.assert_text_messages_balanced()
|
||||
|
||||
# Final turn should not have interrupt
|
||||
finished3 = stream3.last("RUN_FINISHED")
|
||||
final_dump = finished3.model_dump(by_alias=True, exclude_none=True)
|
||||
assert "outcome" not in final_dump, f"Expected no interrupt after completion, got {final_dump.get('outcome')}"
|
||||
@@ -0,0 +1,945 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Comprehensive golden event-stream tests for AgentFrameworkWorkflow.
|
||||
|
||||
Covers the full matrix of workflow-specific AG-UI patterns:
|
||||
- request_info → TOOL_CALL lifecycle and balancing
|
||||
- Executor step events and activity snapshots
|
||||
- Text output, dict output, BaseEvent passthrough, AgentResponse output
|
||||
- Text deduplication across workflow outputs
|
||||
- Workflow error handling → RUN_ERROR
|
||||
- Multi-turn interrupt/resume round-trips
|
||||
- Empty turns with pending requests
|
||||
- Custom workflow events
|
||||
- Text message draining on request_info and executor boundaries
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, cast
|
||||
|
||||
from ag_ui.core import EventType, StateSnapshotEvent
|
||||
from agent_framework import (
|
||||
AgentResponse,
|
||||
Content,
|
||||
Executor,
|
||||
Message,
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
WorkflowEvent,
|
||||
executor,
|
||||
handler,
|
||||
response_handler,
|
||||
)
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
from agent_framework_ag_ui import AgentFrameworkWorkflow
|
||||
|
||||
|
||||
async def _run(wrapper: AgentFrameworkWorkflow, payload: dict[str, Any]) -> EventStream:
|
||||
return EventStream([event async for event in wrapper.run(payload)])
|
||||
|
||||
|
||||
def _payload(
|
||||
msg: str = "go",
|
||||
*,
|
||||
thread_id: str = "thread-wf",
|
||||
run_id: str = "run-wf",
|
||||
**extra: Any,
|
||||
) -> dict[str, Any]:
|
||||
return {"thread_id": thread_id, "run_id": run_id, "messages": [{"role": "user", "content": msg}], **extra}
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 1. Basic workflow text output
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_text_output_golden_sequence() -> None:
|
||||
"""Simple text output: RUN_STARTED → STEP_STARTED → TEXT_* → STEP_FINISHED → TEXT_MESSAGE_END → RUN_FINISHED."""
|
||||
|
||||
@executor(id="greeter")
|
||||
async def greeter(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output("Hello from workflow!")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=greeter).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
stream.assert_text_messages_balanced()
|
||||
stream.assert_has_type("TEXT_MESSAGE_START")
|
||||
stream.assert_has_type("TEXT_MESSAGE_CONTENT")
|
||||
stream.assert_has_type("TEXT_MESSAGE_END")
|
||||
|
||||
# Verify actual content
|
||||
deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
|
||||
assert "Hello from workflow!" in deltas
|
||||
|
||||
|
||||
async def test_workflow_text_output_message_id_consistency() -> None:
|
||||
"""All text events for a single output share the same message_id."""
|
||||
|
||||
@executor(id="echo")
|
||||
async def echo(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output("echo reply")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=echo).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_message_ids_consistent()
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 2. Executor step events and activity snapshots
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_executor_lifecycle_events() -> None:
|
||||
"""Executor invocation produces STEP_STARTED, ACTIVITY_SNAPSHOT, STEP_FINISHED."""
|
||||
|
||||
@executor(id="worker")
|
||||
async def worker(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output("done")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=worker).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
# Step events with executor ID
|
||||
started = [e for e in stream.get("STEP_STARTED") if getattr(e, "step_name", "") == "worker"]
|
||||
finished = [e for e in stream.get("STEP_FINISHED") if getattr(e, "step_name", "") == "worker"]
|
||||
assert started, "Expected STEP_STARTED for 'worker'"
|
||||
assert finished, "Expected STEP_FINISHED for 'worker'"
|
||||
|
||||
# Activity snapshots
|
||||
activities = stream.get("ACTIVITY_SNAPSHOT")
|
||||
assert activities, "Expected ACTIVITY_SNAPSHOT events"
|
||||
# Check one of them has executor payload
|
||||
executor_activities = [a for a in activities if getattr(a, "activity_type", None) == "executor"]
|
||||
assert executor_activities, "Expected executor-type activity snapshots"
|
||||
|
||||
|
||||
async def test_workflow_executor_step_ordering() -> None:
|
||||
"""STEP_STARTED comes before content, STEP_FINISHED comes after."""
|
||||
|
||||
@executor(id="orderer")
|
||||
async def orderer(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output("ordered output")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=orderer).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_ordered_types(
|
||||
[
|
||||
"RUN_STARTED",
|
||||
"STEP_STARTED",
|
||||
"TEXT_MESSAGE_START",
|
||||
"TEXT_MESSAGE_CONTENT",
|
||||
"STEP_FINISHED",
|
||||
"RUN_FINISHED",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 3. Dict output → CUSTOM workflow_output
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_dict_output_maps_to_custom_event() -> None:
|
||||
"""Non-chat dict output is emitted as CUSTOM workflow_output event."""
|
||||
|
||||
@executor(id="structured")
|
||||
async def structured(message: Any, ctx: WorkflowContext[Any, dict[str, int]]) -> None:
|
||||
await ctx.yield_output({"count": 42, "status": 1})
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=structured).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
customs = [e for e in stream.get("CUSTOM") if getattr(e, "name", None) == "workflow_output"]
|
||||
assert len(customs) == 1
|
||||
assert customs[0].value == {"count": 42, "status": 1}
|
||||
|
||||
# Should NOT have TEXT_MESSAGE events for dict output
|
||||
assert "TEXT_MESSAGE_CONTENT" not in stream.types()
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 4. BaseEvent passthrough
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_base_event_passthrough() -> None:
|
||||
"""AG-UI BaseEvent outputs are yielded directly, not wrapped."""
|
||||
|
||||
@executor(id="stateful")
|
||||
async def stateful(message: Any, ctx: WorkflowContext[Any, StateSnapshotEvent]) -> None:
|
||||
await ctx.yield_output(StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot={"active_agent": "flights"}))
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=stateful).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_bookends()
|
||||
snapshots = stream.get("STATE_SNAPSHOT")
|
||||
assert len(snapshots) == 1
|
||||
assert snapshots[0].snapshot["active_agent"] == "flights"
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 5. AgentResponse output (conversation payload)
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_agent_response_output_extracts_latest_assistant() -> None:
|
||||
"""AgentResponse output uses only the latest assistant message, not full history."""
|
||||
|
||||
@executor(id="responder")
|
||||
async def responder(message: Any, ctx: WorkflowContext[Any, AgentResponse]) -> None:
|
||||
response = AgentResponse(
|
||||
messages=[
|
||||
Message(role="user", contents=[Content.from_text("My order is damaged")]),
|
||||
Message(role="assistant", contents=[Content.from_text("I'll process your replacement.")]),
|
||||
]
|
||||
)
|
||||
await ctx.yield_output(response)
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=responder).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_text_messages_balanced()
|
||||
|
||||
deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
|
||||
assert deltas == ["I'll process your replacement."]
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 6. Custom workflow events
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class ProgressEvent(WorkflowEvent):
|
||||
"""Custom workflow event for testing CUSTOM event mapping."""
|
||||
|
||||
def __init__(self, progress: int) -> None:
|
||||
super().__init__(cast(Any, "custom_progress"), data={"progress": progress})
|
||||
|
||||
|
||||
async def test_workflow_custom_events() -> None:
|
||||
"""Custom workflow events are mapped to CUSTOM AG-UI events."""
|
||||
|
||||
@executor(id="progress_tracker")
|
||||
async def progress_tracker(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.add_event(ProgressEvent(25))
|
||||
await ctx.yield_output("In progress...")
|
||||
await ctx.add_event(ProgressEvent(100))
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=progress_tracker).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
progress_events = [e for e in stream.get("CUSTOM") if getattr(e, "name", None) == "custom_progress"]
|
||||
assert len(progress_events) == 2
|
||||
assert progress_events[0].value == {"progress": 25}
|
||||
assert progress_events[1].value == {"progress": 100}
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 7. request_info → TOOL_CALL lifecycle
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_request_info_tool_call_lifecycle() -> None:
|
||||
"""request_info emits TOOL_CALL_START/ARGS/END cycle plus CUSTOM request_info."""
|
||||
|
||||
@executor(id="requester")
|
||||
async def requester(message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.request_info("Need approval", str, request_id="req-1")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=requester).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
# Tool call lifecycle
|
||||
stream.assert_ordered_types(
|
||||
[
|
||||
"RUN_STARTED",
|
||||
"TOOL_CALL_START",
|
||||
"TOOL_CALL_ARGS",
|
||||
"TOOL_CALL_END",
|
||||
"CUSTOM", # request_info
|
||||
"RUN_FINISHED",
|
||||
]
|
||||
)
|
||||
|
||||
# Verify tool call details
|
||||
start = stream.first("TOOL_CALL_START")
|
||||
assert start.tool_call_id == "req-1"
|
||||
assert start.tool_call_name == "request_info"
|
||||
|
||||
# TOOL_CALL_ARGS should contain the request payload
|
||||
args = stream.first("TOOL_CALL_ARGS")
|
||||
assert args.tool_call_id == "req-1"
|
||||
parsed_args = json.loads(args.delta)
|
||||
assert parsed_args["request_id"] == "req-1"
|
||||
|
||||
# Tool calls should be balanced
|
||||
stream.assert_tool_calls_balanced()
|
||||
|
||||
|
||||
async def test_workflow_request_info_interrupt_in_run_finished() -> None:
|
||||
"""request_info populates RUN_FINISHED.outcome.interrupts with the request metadata."""
|
||||
|
||||
@executor(id="requester")
|
||||
async def requester(message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.request_info(
|
||||
{"message": "Choose a flight", "options": [{"airline": "KLM"}], "agent": "flights"},
|
||||
dict,
|
||||
request_id="flights-choice",
|
||||
)
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=requester).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
interrupt = stream.run_finished_interrupts()
|
||||
assert len(interrupt) == 1
|
||||
assert interrupt[0]["id"] == "flights-choice"
|
||||
assert stream.interrupt_metadata_value(interrupt[0])["agent"] == "flights"
|
||||
|
||||
|
||||
async def test_workflow_request_info_emits_interrupt_card_event() -> None:
|
||||
"""request_info with dict data emits a WorkflowInterruptEvent custom event."""
|
||||
|
||||
@executor(id="requester")
|
||||
async def requester(message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.request_info(
|
||||
{"message": "Pick one", "options": ["A", "B"]},
|
||||
dict,
|
||||
request_id="pick-1",
|
||||
)
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=requester).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
interrupt_cards = [e for e in stream.get("CUSTOM") if getattr(e, "name", None) == "WorkflowInterruptEvent"]
|
||||
assert interrupt_cards, "Expected WorkflowInterruptEvent custom event"
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 8. Text message draining on request_info boundary
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_text_drained_before_request_info() -> None:
|
||||
"""Open text message is closed (TEXT_MESSAGE_END) before request_info tool calls begin."""
|
||||
|
||||
@executor(id="text_then_request")
|
||||
async def text_then_request(message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.yield_output("Please confirm this action.") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
|
||||
await ctx.request_info("Need approval", str, request_id="approval-1")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=text_then_request).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_text_messages_balanced()
|
||||
stream.assert_tool_calls_balanced()
|
||||
|
||||
# TEXT_MESSAGE_END must appear before TOOL_CALL_START
|
||||
types = stream.types()
|
||||
text_end_idx = types.index("TEXT_MESSAGE_END")
|
||||
tool_start_idx = types.index("TOOL_CALL_START")
|
||||
assert text_end_idx < tool_start_idx, (
|
||||
f"TEXT_MESSAGE_END (idx={text_end_idx}) must come before TOOL_CALL_START (idx={tool_start_idx})"
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 9. Text deduplication
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_skips_duplicate_text_from_snapshot() -> None:
|
||||
"""Duplicate text from AgentResponse snapshot is not re-emitted."""
|
||||
|
||||
@executor(id="deduper")
|
||||
async def deduper(message: Any, ctx: WorkflowContext[Any, Any]) -> None:
|
||||
text = "Order processed successfully."
|
||||
await ctx.yield_output(text)
|
||||
# Snapshot repeats the same text
|
||||
await ctx.yield_output(
|
||||
AgentResponse(
|
||||
messages=[
|
||||
Message(role="user", contents=[Content.from_text("process order")]),
|
||||
Message(role="assistant", contents=[Content.from_text(text)]),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=deduper).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_text_messages_balanced()
|
||||
deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
|
||||
# Text should appear only once
|
||||
assert deltas == ["Order processed successfully."]
|
||||
|
||||
|
||||
async def test_workflow_skips_consecutive_duplicate_outputs() -> None:
|
||||
"""Consecutive identical text outputs are deduplicated."""
|
||||
|
||||
@executor(id="repeater")
|
||||
async def repeater(message: Any, ctx: WorkflowContext[Any, Any]) -> None:
|
||||
text = "Done!"
|
||||
await ctx.yield_output(text)
|
||||
await ctx.yield_output(text)
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=repeater).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_text_messages_balanced()
|
||||
deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
|
||||
assert deltas == ["Done!"]
|
||||
|
||||
|
||||
async def test_workflow_emits_distinct_consecutive_outputs() -> None:
|
||||
"""Distinct text outputs are all emitted, not incorrectly deduplicated."""
|
||||
|
||||
@executor(id="multisayer")
|
||||
async def multisayer(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output("First part. ")
|
||||
await ctx.yield_output("Second part.")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=multisayer).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_text_messages_balanced()
|
||||
deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
|
||||
assert deltas == ["First part. ", "Second part."]
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 10. Workflow error handling → RUN_ERROR
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_error_emits_run_error_event() -> None:
|
||||
"""Exceptions during workflow streaming produce RUN_ERROR events."""
|
||||
|
||||
class FailingWorkflow:
|
||||
def run(self, **kwargs: Any):
|
||||
async def _stream():
|
||||
raise RuntimeError("workflow exploded")
|
||||
yield # pragma: no cover
|
||||
|
||||
return _stream()
|
||||
|
||||
wrapper = AgentFrameworkWorkflow(workflow=cast(Any, FailingWorkflow()))
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
# Should still have RUN_STARTED
|
||||
stream.assert_has_type("RUN_STARTED")
|
||||
# Should have RUN_ERROR
|
||||
stream.assert_has_type("RUN_ERROR")
|
||||
error = stream.first("RUN_ERROR")
|
||||
assert "workflow exploded" in error.message
|
||||
|
||||
|
||||
async def test_workflow_error_preserves_bookend_structure() -> None:
|
||||
"""Even on error, RUN_STARTED is the first event."""
|
||||
|
||||
class FailingWorkflow:
|
||||
def run(self, **kwargs: Any):
|
||||
async def _stream():
|
||||
raise ValueError("bad input")
|
||||
yield # pragma: no cover
|
||||
|
||||
return _stream()
|
||||
|
||||
wrapper = AgentFrameworkWorkflow(workflow=cast(Any, FailingWorkflow()))
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
types = stream.types()
|
||||
assert types[0] == "RUN_STARTED"
|
||||
assert "RUN_ERROR" in types
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 11. Multi-turn request_info interrupt/resume
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_interrupt_resume_round_trip() -> None:
|
||||
"""Turn 1: request_info → interrupt. Turn 2: resume → completion."""
|
||||
|
||||
class RequesterExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="requester")
|
||||
|
||||
@handler
|
||||
async def start(self, message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.request_info("Choose an option", str, request_id="choice-1")
|
||||
|
||||
@response_handler
|
||||
async def handle_choice(self, original: str, response: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.yield_output(f"You chose: {response}") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=RequesterExecutor()).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
|
||||
# Turn 1
|
||||
stream1 = await _run(wrapper, _payload(thread_id="thread-resume", run_id="run-1"))
|
||||
stream1.assert_bookends()
|
||||
stream1.assert_no_run_error()
|
||||
stream1.assert_tool_calls_balanced()
|
||||
|
||||
interrupt1 = stream1.run_finished_interrupts()
|
||||
assert interrupt1, "Expected interrupt"
|
||||
assert interrupt1[0]["id"] == "choice-1"
|
||||
|
||||
# Turn 2: resume
|
||||
stream2 = await _run(
|
||||
wrapper,
|
||||
{
|
||||
"thread_id": "thread-resume",
|
||||
"run_id": "run-2",
|
||||
"messages": [],
|
||||
"resume": {"interrupts": [{"id": "choice-1", "value": "Option A"}]},
|
||||
},
|
||||
)
|
||||
stream2.assert_has_run_lifecycle()
|
||||
stream2.assert_no_run_error()
|
||||
stream2.assert_text_messages_balanced()
|
||||
|
||||
# Should have the response text
|
||||
deltas = [e.delta for e in stream2.get("TEXT_MESSAGE_CONTENT")]
|
||||
assert any("Option A" in d for d in deltas), f"Expected 'Option A' in deltas: {deltas}"
|
||||
|
||||
# No interrupt after resume
|
||||
finished2 = stream2.last("RUN_FINISHED")
|
||||
dumped2 = finished2.model_dump(by_alias=True, exclude_none=True)
|
||||
assert "outcome" not in dumped2
|
||||
|
||||
|
||||
async def test_workflow_forwarded_props_resume() -> None:
|
||||
"""forwarded_props.command.resume should resume with canonical resume entries."""
|
||||
|
||||
@executor(id="requester")
|
||||
async def requester(message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.request_info({"options": [{"name": "A"}]}, dict, request_id="pick")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=requester).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
|
||||
# Turn 1
|
||||
await _run(wrapper, _payload(thread_id="thread-fwd", run_id="run-1"))
|
||||
|
||||
# Turn 2 via forwarded_props
|
||||
stream2 = await _run(
|
||||
wrapper,
|
||||
{
|
||||
"thread_id": "thread-fwd",
|
||||
"run_id": "run-2",
|
||||
"messages": [],
|
||||
"forwarded_props": {
|
||||
"command": {"resume": [{"interruptId": "pick", "status": "resolved", "payload": {"name": "A"}}]}
|
||||
},
|
||||
},
|
||||
)
|
||||
stream2.assert_bookends()
|
||||
stream2.assert_no_run_error()
|
||||
|
||||
finished = stream2.last("RUN_FINISHED")
|
||||
assert "outcome" not in finished.model_dump(by_alias=True, exclude_none=True)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 12. Empty turns with pending requests
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_empty_turn_with_pending_request_emits_run_error() -> None:
|
||||
"""An empty turn with a pending request must provide canonical resume entries."""
|
||||
|
||||
@executor(id="requester")
|
||||
async def requester(message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.request_info({"prompt": "choose"}, dict, request_id="pick-one")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=requester).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
|
||||
# Turn 1: trigger the request
|
||||
await _run(wrapper, _payload(thread_id="thread-empty", run_id="run-1"))
|
||||
|
||||
# Turn 2: empty messages, no resume
|
||||
stream2 = await _run(
|
||||
wrapper,
|
||||
{
|
||||
"thread_id": "thread-empty",
|
||||
"run_id": "run-2",
|
||||
"messages": [],
|
||||
},
|
||||
)
|
||||
stream2.assert_has_type("RUN_STARTED")
|
||||
run_error = stream2.last("RUN_ERROR")
|
||||
assert run_error.code == "WORKFLOW_RESUME_REQUIRED"
|
||||
|
||||
|
||||
async def test_workflow_empty_turn_no_pending_requests() -> None:
|
||||
"""Empty turn with no pending requests produces clean bookends."""
|
||||
|
||||
@executor(id="noop")
|
||||
async def noop(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output("done")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=noop).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
|
||||
# Run once to completion
|
||||
await _run(wrapper, _payload(thread_id="thread-empty-clean", run_id="run-1"))
|
||||
|
||||
# Empty turn
|
||||
stream2 = await _run(
|
||||
wrapper,
|
||||
{
|
||||
"thread_id": "thread-empty-clean",
|
||||
"run_id": "run-2",
|
||||
"messages": [],
|
||||
},
|
||||
)
|
||||
stream2.assert_bookends()
|
||||
stream2.assert_no_run_error()
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 13. Usage content as CUSTOM event
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_usage_output_maps_to_custom_event() -> None:
|
||||
"""Usage Content outputs are surfaced as custom usage events."""
|
||||
|
||||
@executor(id="usage_reporter")
|
||||
async def usage_reporter(message: Any, ctx: WorkflowContext[Any, Content]) -> None:
|
||||
await ctx.yield_output(
|
||||
Content.from_usage({"input_token_count": 100, "output_token_count": 50, "total_token_count": 150})
|
||||
)
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=usage_reporter).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
stream = await _run(wrapper, _payload())
|
||||
|
||||
stream.assert_bookends()
|
||||
stream.assert_no_run_error()
|
||||
|
||||
usage_events = [e for e in stream.get("CUSTOM") if getattr(e, "name", None) == "usage"]
|
||||
assert len(usage_events) == 1
|
||||
assert usage_events[0].value["input_token_count"] == 100
|
||||
assert usage_events[0].value["total_token_count"] == 150
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 14. Approval flow (Content-based request_info)
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_approval_flow_round_trip() -> None:
|
||||
"""function_approval_request via request_info, then resume with approval response."""
|
||||
|
||||
class ApprovalExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="approval_exec")
|
||||
|
||||
@handler
|
||||
async def start(self, message: Any, ctx: WorkflowContext) -> None:
|
||||
function_call = Content.from_function_call(
|
||||
call_id="refund-call",
|
||||
name="submit_refund",
|
||||
arguments={"order_id": "12345", "amount": "$89.99"},
|
||||
)
|
||||
approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call)
|
||||
await ctx.request_info(approval_request, Content, request_id="approval-1")
|
||||
|
||||
@response_handler
|
||||
async def handle_approval(self, original_request: Content, response: Content, ctx: WorkflowContext) -> None:
|
||||
status = "approved" if bool(response.approved) else "rejected"
|
||||
await ctx.yield_output(f"Refund {status}.") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
|
||||
# Turn 1: request approval
|
||||
stream1 = await _run(wrapper, _payload(thread_id="thread-approval", run_id="run-1"))
|
||||
stream1.assert_bookends()
|
||||
stream1.assert_no_run_error()
|
||||
|
||||
interrupt1 = stream1.run_finished_interrupts()
|
||||
assert interrupt1, "Expected approval interrupt"
|
||||
interrupt_value = stream1.interrupt_metadata_value(interrupt1[0])
|
||||
|
||||
# Turn 2: approve
|
||||
stream2 = await _run(
|
||||
wrapper,
|
||||
{
|
||||
"thread_id": "thread-approval",
|
||||
"run_id": "run-2",
|
||||
"messages": [],
|
||||
"resume": {
|
||||
"interrupts": [
|
||||
{
|
||||
"id": "approval-1",
|
||||
"value": {
|
||||
"type": "function_approval_response",
|
||||
"approved": True,
|
||||
"id": interrupt_value.get("id", "approval-1"),
|
||||
"function_call": interrupt_value.get("function_call"),
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
},
|
||||
)
|
||||
stream2.assert_has_run_lifecycle()
|
||||
stream2.assert_no_run_error()
|
||||
stream2.assert_text_messages_balanced()
|
||||
|
||||
deltas = [e.delta for e in stream2.get("TEXT_MESSAGE_CONTENT")]
|
||||
assert any("approved" in d for d in deltas)
|
||||
|
||||
# No more interrupt
|
||||
finished2 = stream2.last("RUN_FINISHED")
|
||||
assert "outcome" not in finished2.model_dump(by_alias=True, exclude_none=True)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 15. Message list request/response coercion
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_message_list_resume() -> None:
|
||||
"""Resume with list[Message] payload coerces correctly into workflow response."""
|
||||
|
||||
class MessageRequestExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="msg_request")
|
||||
|
||||
@handler
|
||||
async def start(self, message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.request_info({"prompt": "Need follow-up"}, list[Message], request_id="handoff")
|
||||
|
||||
@response_handler
|
||||
async def handle_input(self, original: dict, response: list[Message], ctx: WorkflowContext) -> None:
|
||||
user_text = response[0].text if response else ""
|
||||
await ctx.yield_output(f"Got: {user_text}") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=MessageRequestExecutor()).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
|
||||
# Turn 1
|
||||
await _run(wrapper, _payload(thread_id="thread-msg", run_id="run-1"))
|
||||
|
||||
# Turn 2: resume with message list
|
||||
stream2 = await _run(
|
||||
wrapper,
|
||||
{
|
||||
"thread_id": "thread-msg",
|
||||
"run_id": "run-2",
|
||||
"messages": [],
|
||||
"resume": {
|
||||
"interrupts": [
|
||||
{
|
||||
"id": "handoff",
|
||||
"value": [
|
||||
{"role": "user", "contents": [{"type": "text", "text": "Ship a replacement"}]},
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
},
|
||||
)
|
||||
stream2.assert_has_run_lifecycle()
|
||||
stream2.assert_no_run_error()
|
||||
stream2.assert_text_messages_balanced()
|
||||
|
||||
deltas = [e.delta for e in stream2.get("TEXT_MESSAGE_CONTENT")]
|
||||
assert any("replacement" in d for d in deltas)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 16. Plain text follow-up does NOT infer interrupt response
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_plain_text_does_not_resume_pending_dict_request() -> None:
|
||||
"""Plain text user follow-up should fail instead of being coerced into a dict response."""
|
||||
|
||||
@executor(id="requester")
|
||||
async def requester(message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.request_info(
|
||||
{"message": "Choose a flight", "options": [{"airline": "KLM"}], "agent": "flights"},
|
||||
dict,
|
||||
request_id="flights-choice",
|
||||
)
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=requester).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
|
||||
# Turn 1
|
||||
await _run(wrapper, _payload(thread_id="thread-nocoerce", run_id="run-1"))
|
||||
|
||||
# Turn 2: plain text follow-up with request_info tool call in history
|
||||
stream2 = await _run(
|
||||
wrapper,
|
||||
{
|
||||
"thread_id": "thread-nocoerce",
|
||||
"run_id": "run-2",
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "flights-choice",
|
||||
"type": "function",
|
||||
"function": {"name": "request_info", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "I prefer KLM please"},
|
||||
],
|
||||
},
|
||||
)
|
||||
stream2.assert_has_type("RUN_STARTED")
|
||||
run_error = stream2.last("RUN_ERROR")
|
||||
assert run_error.code == "WORKFLOW_RESUME_REQUIRED"
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 17. Workflow factory (thread-scoped workflows)
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_factory_thread_scoping() -> None:
|
||||
"""workflow_factory creates separate workflow instances per thread_id."""
|
||||
|
||||
def make_workflow(thread_id: str):
|
||||
@executor(id="echo")
|
||||
async def echo(message: Any, ctx: WorkflowContext[Any, str]) -> None:
|
||||
await ctx.yield_output(f"Thread: {thread_id}")
|
||||
|
||||
return WorkflowBuilder(start_executor=echo).build()
|
||||
|
||||
wrapper = AgentFrameworkWorkflow(workflow_factory=make_workflow)
|
||||
|
||||
stream_a = await _run(wrapper, _payload(thread_id="thread-a", run_id="run-a"))
|
||||
stream_b = await _run(wrapper, _payload(thread_id="thread-b", run_id="run-b"))
|
||||
|
||||
stream_a.assert_bookends()
|
||||
stream_b.assert_bookends()
|
||||
|
||||
deltas_a = [e.delta for e in stream_a.get("TEXT_MESSAGE_CONTENT")]
|
||||
deltas_b = [e.delta for e in stream_b.get("TEXT_MESSAGE_CONTENT")]
|
||||
assert any("thread-a" in d for d in deltas_a)
|
||||
assert any("thread-b" in d for d in deltas_b)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# 18. Multiple request_info calls in sequence
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def test_workflow_sequential_request_info_interrupts() -> None:
|
||||
"""Two chained executors each requesting info: first triggers interrupt, resume, then second triggers interrupt.
|
||||
|
||||
This mirrors the subgraphs_agent pattern where separate executors handle sequential interactions.
|
||||
"""
|
||||
|
||||
class NameRequester(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="name_requester")
|
||||
|
||||
@handler
|
||||
async def start(self, message: Any, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.request_info("What's your name?", str, request_id="name-req")
|
||||
|
||||
@response_handler
|
||||
async def handle_name(self, original: str, response: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(response)
|
||||
|
||||
class DestRequester(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="dest_requester")
|
||||
|
||||
@handler
|
||||
async def start(self, message: str, ctx: WorkflowContext[str]) -> None:
|
||||
self._name = message
|
||||
await ctx.request_info("Where to?", str, request_id="dest-req")
|
||||
|
||||
@response_handler
|
||||
async def handle_dest(self, original: str, response: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.yield_output(f"Booking for {self._name} to {response}") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
|
||||
|
||||
name_requester = NameRequester()
|
||||
dest_requester = DestRequester()
|
||||
workflow = WorkflowBuilder(start_executor=name_requester).add_chain([name_requester, dest_requester]).build()
|
||||
wrapper = AgentFrameworkWorkflow(workflow=workflow)
|
||||
|
||||
# Turn 1
|
||||
stream1 = await _run(wrapper, _payload(thread_id="thread-seq", run_id="run-1"))
|
||||
stream1.assert_bookends()
|
||||
stream1.assert_tool_calls_balanced()
|
||||
interrupt1 = stream1.run_finished_interrupts()
|
||||
assert interrupt1[0]["id"] == "name-req"
|
||||
|
||||
# Turn 2: answer name → triggers second executor's request_info
|
||||
stream2 = await _run(
|
||||
wrapper,
|
||||
{
|
||||
"thread_id": "thread-seq",
|
||||
"run_id": "run-2",
|
||||
"messages": [],
|
||||
"resume": {"interrupts": [{"id": "name-req", "value": "Alice"}]},
|
||||
},
|
||||
)
|
||||
stream2.assert_has_run_lifecycle()
|
||||
stream2.assert_tool_calls_balanced()
|
||||
interrupt2 = stream2.run_finished_interrupts()
|
||||
assert interrupt2[0]["id"] == "dest-req"
|
||||
|
||||
# Turn 3: answer destination → completion
|
||||
stream3 = await _run(
|
||||
wrapper,
|
||||
{
|
||||
"thread_id": "thread-seq",
|
||||
"run_id": "run-3",
|
||||
"messages": [],
|
||||
"resume": {"interrupts": [{"id": "dest-req", "value": "Paris"}]},
|
||||
},
|
||||
)
|
||||
stream3.assert_has_run_lifecycle()
|
||||
stream3.assert_no_run_error()
|
||||
stream3.assert_text_messages_balanced()
|
||||
|
||||
deltas = [e.delta for e in stream3.get("TEXT_MESSAGE_CONTENT")]
|
||||
assert any("Alice" in d and "Paris" in d for d in deltas)
|
||||
assert "outcome" not in stream3.last("RUN_FINISHED").model_dump(by_alias=True, exclude_none=True)
|
||||
@@ -0,0 +1,72 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""SSE parsing helpers for AG-UI HTTP round-trip tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from event_stream import EventStream # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
|
||||
|
||||
|
||||
def parse_sse_response(response_content: bytes) -> list[dict[str, Any]]:
|
||||
"""Parse raw SSE bytes from TestClient into a list of event dicts.
|
||||
|
||||
Each SSE event is a ``data: {...}`` line followed by a blank line.
|
||||
"""
|
||||
text = response_content.decode("utf-8")
|
||||
events: list[dict[str, Any]] = []
|
||||
decode_errors: list[str] = []
|
||||
for line in text.splitlines():
|
||||
if line.startswith("data: "):
|
||||
payload = line[6:]
|
||||
try:
|
||||
events.append(json.loads(payload))
|
||||
except json.JSONDecodeError as exc:
|
||||
decode_errors.append(f"payload={payload!r}, error={exc}")
|
||||
continue
|
||||
if decode_errors:
|
||||
joined = "; ".join(decode_errors)
|
||||
raise AssertionError(f"Failed to decode one or more SSE data lines: {joined}")
|
||||
return events
|
||||
|
||||
|
||||
def parse_sse_to_event_stream(response_content: bytes) -> EventStream:
|
||||
"""Parse SSE bytes and wrap in EventStream for structured assertions.
|
||||
|
||||
Returns an EventStream over lightweight SimpleNamespace objects that
|
||||
mirror AG-UI event attributes (type, message_id, tool_call_id, etc.)
|
||||
so that EventStream assertion methods work.
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
|
||||
raw_events = parse_sse_response(response_content)
|
||||
events: list[Any] = []
|
||||
for raw in raw_events:
|
||||
# Normalize camelCase keys to snake_case attributes that EventStream expects
|
||||
ns = SimpleNamespace()
|
||||
ns.type = raw.get("type", "")
|
||||
ns.raw = raw
|
||||
# Map common camelCase fields
|
||||
for camel, snake in _FIELD_MAP.items():
|
||||
if camel in raw:
|
||||
setattr(ns, snake, raw[camel])
|
||||
# Also keep camelCase as attributes for direct access
|
||||
for key, value in raw.items():
|
||||
if not hasattr(ns, key):
|
||||
setattr(ns, key, value)
|
||||
events.append(ns)
|
||||
return EventStream(events)
|
||||
|
||||
|
||||
_FIELD_MAP: dict[str, str] = {
|
||||
"messageId": "message_id",
|
||||
"runId": "run_id",
|
||||
"threadId": "thread_id",
|
||||
"toolCallId": "tool_call_id",
|
||||
"toolCallName": "tool_call_name",
|
||||
"toolName": "tool_call_name",
|
||||
"parentMessageId": "parent_message_id",
|
||||
"stepName": "step_name",
|
||||
}
|
||||
@@ -0,0 +1,564 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for AGUIChatClient."""
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, Awaitable, MutableSequence
|
||||
from typing import Any
|
||||
|
||||
from ag_ui.core import Interrupt, ResumeEntry
|
||||
from agent_framework import (
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Content,
|
||||
Message,
|
||||
ResponseStream,
|
||||
tool,
|
||||
)
|
||||
from pytest import MonkeyPatch
|
||||
|
||||
from agent_framework_ag_ui._client import AGUIChatClient
|
||||
from agent_framework_ag_ui._http_service import AGUIHttpService
|
||||
|
||||
|
||||
class StubAGUIChatClient(AGUIChatClient):
|
||||
"""Testable wrapper exposing protected helpers."""
|
||||
|
||||
@property
|
||||
def http_service(self) -> AGUIHttpService:
|
||||
"""Expose http service for monkeypatching."""
|
||||
return self._http_service
|
||||
|
||||
def extract_state_from_messages(self, messages: list[Message]) -> tuple[list[Message], dict[str, Any] | None]:
|
||||
"""Expose state extraction helper."""
|
||||
return self._extract_state_from_messages(messages)
|
||||
|
||||
def convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
|
||||
"""Expose message conversion helper."""
|
||||
return self._convert_messages_to_agui_format(messages)
|
||||
|
||||
def get_thread_id(self, options: ChatOptions[Any] | dict[str, Any] | None) -> str:
|
||||
"""Expose thread id helper."""
|
||||
return self._get_thread_id(options) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
|
||||
|
||||
def inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[Message],
|
||||
options: ChatOptions[Any] | dict[str, Any] | None,
|
||||
stream: bool = False,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
"""Proxy to protected response call."""
|
||||
return self._inner_get_response(messages=messages, options=options, stream=stream) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
|
||||
|
||||
|
||||
class TestAGUIChatClient:
|
||||
"""Test suite for AGUIChatClient."""
|
||||
|
||||
async def test_client_initialization(self) -> None:
|
||||
"""Test client initialization."""
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
|
||||
assert client.http_service is not None
|
||||
assert client.http_service.endpoint.startswith("http://localhost:8888")
|
||||
|
||||
async def test_client_context_manager(self) -> None:
|
||||
"""Test client as async context manager."""
|
||||
async with StubAGUIChatClient(endpoint="http://localhost:8888/") as client:
|
||||
assert client is not None
|
||||
|
||||
async def test_extract_state_from_messages_no_state(self) -> None:
|
||||
"""Test state extraction when no state is present."""
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
messages = [
|
||||
Message(role="user", contents=["Hello"]),
|
||||
Message(role="assistant", contents=["Hi there"]),
|
||||
]
|
||||
|
||||
result_messages, state = client.extract_state_from_messages(messages)
|
||||
|
||||
assert result_messages == messages
|
||||
assert state is None
|
||||
|
||||
async def test_extract_state_from_messages_with_state(self) -> None:
|
||||
"""Test state extraction from last message."""
|
||||
import base64
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
|
||||
state_data = {"key": "value", "count": 42}
|
||||
state_json = json.dumps(state_data)
|
||||
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
|
||||
|
||||
messages = [
|
||||
Message(role="user", contents=["Hello"]),
|
||||
Message(
|
||||
role="user",
|
||||
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
|
||||
),
|
||||
]
|
||||
|
||||
result_messages, state = client.extract_state_from_messages(messages)
|
||||
|
||||
assert len(result_messages) == 1
|
||||
assert result_messages[0].text == "Hello"
|
||||
assert state == state_data
|
||||
|
||||
async def test_extract_state_invalid_json(self) -> None:
|
||||
"""Test state extraction with invalid JSON."""
|
||||
import base64
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
|
||||
invalid_json = "not valid json"
|
||||
state_b64 = base64.b64encode(invalid_json.encode("utf-8")).decode("utf-8")
|
||||
|
||||
messages = [
|
||||
Message(
|
||||
role="user",
|
||||
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
|
||||
),
|
||||
]
|
||||
|
||||
result_messages, state = client.extract_state_from_messages(messages)
|
||||
|
||||
assert result_messages == messages
|
||||
assert state is None
|
||||
|
||||
async def test_convert_messages_to_agui_format(self) -> None:
|
||||
"""Test message conversion to AG-UI format."""
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
messages = [
|
||||
Message(role="user", contents=["What is the weather?"]),
|
||||
Message(role="assistant", contents=["Let me check."], message_id="msg_123"),
|
||||
]
|
||||
|
||||
agui_messages = client.convert_messages_to_agui_format(messages)
|
||||
|
||||
assert len(agui_messages) == 2
|
||||
assert agui_messages[0]["role"] == "user"
|
||||
assert agui_messages[0]["content"] == "What is the weather?"
|
||||
assert agui_messages[1]["role"] == "assistant"
|
||||
assert agui_messages[1]["content"] == "Let me check."
|
||||
assert agui_messages[1]["id"] == "msg_123"
|
||||
|
||||
async def test_get_thread_id_from_metadata(self) -> None:
|
||||
"""Test thread ID extraction from metadata."""
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
chat_options = ChatOptions(metadata={"thread_id": "existing_thread_123"})
|
||||
|
||||
thread_id = client.get_thread_id(chat_options)
|
||||
|
||||
assert thread_id == "existing_thread_123"
|
||||
|
||||
async def test_get_thread_id_generation(self) -> None:
|
||||
"""Test automatic thread ID generation."""
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
chat_options = ChatOptions()
|
||||
|
||||
thread_id = client.get_thread_id(chat_options)
|
||||
|
||||
assert thread_id.startswith("thread_")
|
||||
assert len(thread_id) > 7
|
||||
|
||||
async def test_get_response_streaming(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Test streaming response method."""
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " world"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [Message(role="user", contents=["Test message"])]
|
||||
chat_options = ChatOptions()
|
||||
|
||||
updates: list[ChatResponseUpdate] = []
|
||||
stream = client.inner_get_response(messages=messages, stream=True, options=chat_options)
|
||||
assert isinstance(stream, ResponseStream)
|
||||
async for update in stream:
|
||||
updates.append(update)
|
||||
|
||||
assert len(updates) == 4
|
||||
assert updates[0].additional_properties is not None
|
||||
assert updates[0].additional_properties["thread_id"] == "thread_1"
|
||||
|
||||
first_content = updates[1].contents[0]
|
||||
second_content = updates[2].contents[0]
|
||||
assert first_content.type == "text"
|
||||
assert second_content.type == "text"
|
||||
assert first_content.text == "Hello"
|
||||
assert second_content.text == " world"
|
||||
|
||||
async def test_get_response_non_streaming(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Test non-streaming response method."""
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Complete response"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [Message(role="user", contents=["Test message"])]
|
||||
chat_options: dict[str, Any] = {}
|
||||
|
||||
response = await client.inner_get_response(messages=messages, options=chat_options)
|
||||
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert "Complete response" in response.text
|
||||
|
||||
async def test_tool_handling(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Test that client tool metadata is sent to server.
|
||||
|
||||
Client tool metadata (name, description, schema) is sent to server for planning.
|
||||
When server requests a client function, function invocation mixin
|
||||
intercepts and executes it locally. This matches .NET AG-UI implementation.
|
||||
"""
|
||||
from agent_framework import tool
|
||||
|
||||
@tool
|
||||
def test_tool(param: str) -> str:
|
||||
"""Test tool."""
|
||||
return "result"
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
# Client tool metadata should be sent to server
|
||||
tools: list[dict[str, Any]] | None = kwargs.get("tools")
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
tool_entry = tools[0]
|
||||
assert tool_entry["name"] == "test_tool"
|
||||
assert tool_entry["description"] == "Test tool."
|
||||
assert "parameters" in tool_entry
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [Message(role="user", contents=["Test with tools"])]
|
||||
chat_options = ChatOptions(tools=[test_tool])
|
||||
|
||||
response = await client.inner_get_response(messages=messages, options=chat_options)
|
||||
|
||||
assert response is not None
|
||||
|
||||
async def test_server_tool_calls_unwrapped_after_invocation(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Ensure server-side tool calls are exposed as FunctionCallContent after processing."""
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_time_zone"},
|
||||
{"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"location": "Seattle"}'},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [Message(role="user", contents=["Test server tool execution"])]
|
||||
|
||||
updates: list[ChatResponseUpdate] = []
|
||||
async for update in client.get_response(messages, stream=True):
|
||||
updates.append(update)
|
||||
|
||||
function_calls = [
|
||||
content for update in updates for content in update.contents if content.type == "function_call"
|
||||
]
|
||||
assert function_calls
|
||||
assert function_calls[0].name == "get_time_zone"
|
||||
|
||||
assert not any(content.type == "server_function_call" for update in updates for content in update.contents)
|
||||
|
||||
async def test_server_tool_calls_not_executed_locally(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Server tools should not trigger local function invocation even when client tools exist."""
|
||||
|
||||
@tool
|
||||
def client_tool() -> str:
|
||||
"""Client tool stub."""
|
||||
return "client"
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_time_zone"},
|
||||
{"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"location": "Seattle"}'},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
async def fake_auto_invoke(*args: object, **kwargs: Any) -> None:
|
||||
function_call = kwargs.get("function_call_content") or args[0]
|
||||
raise AssertionError(f"Unexpected local execution of server tool: {getattr(function_call, 'name', '?')}")
|
||||
|
||||
monkeypatch.setattr("agent_framework._tools._auto_invoke_function", fake_auto_invoke)
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [Message(role="user", contents=["Test server tool execution"])]
|
||||
|
||||
async for _ in client.get_response(
|
||||
messages, stream=True, options={"tool_choice": "auto", "tools": [client_tool]}
|
||||
):
|
||||
pass
|
||||
|
||||
async def test_state_transmission(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Test state is properly transmitted to server."""
|
||||
import base64
|
||||
|
||||
state_data = {"user_id": "123", "session": "abc"}
|
||||
state_json = json.dumps(state_data)
|
||||
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
|
||||
|
||||
messages = [
|
||||
Message(role="user", contents=["Hello"]),
|
||||
Message(
|
||||
role="user",
|
||||
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
|
||||
),
|
||||
]
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
assert kwargs.get("state") == state_data
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
chat_options = ChatOptions()
|
||||
|
||||
response = await client.inner_get_response(messages=messages, options=chat_options)
|
||||
|
||||
assert response is not None
|
||||
|
||||
async def test_extract_state_from_empty_messages(self) -> None:
|
||||
"""Empty messages list returns empty list and None state."""
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
result_messages, state = client.extract_state_from_messages([])
|
||||
assert result_messages == []
|
||||
assert state is None
|
||||
|
||||
async def test_register_server_tool_non_dict_config(self) -> None:
|
||||
"""Non-dict function_invocation_configuration is a no-op."""
|
||||
client = StubAGUIChatClient(
|
||||
endpoint="http://localhost:8888/",
|
||||
function_invocation_configuration=None, # type: ignore[arg-type]
|
||||
)
|
||||
# Should not raise
|
||||
client._register_server_tool_placeholder("some_tool")
|
||||
|
||||
async def test_non_streaming_response(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Non-streaming path collects updates into ChatResponse."""
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [Message(role="user", contents=["Test"])]
|
||||
response = await client.inner_get_response(messages=messages, options={}, stream=False)
|
||||
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
|
||||
async def test_client_tool_sets_additional_properties(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Client tool content gets agui_thread_id additional property."""
|
||||
|
||||
@tool
|
||||
def my_tool(param: str) -> str:
|
||||
"""My tool."""
|
||||
return "result"
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "my_tool"},
|
||||
{"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"param": "test"}'},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [Message(role="user", contents=["Test"])]
|
||||
updates: list[ChatResponseUpdate] = []
|
||||
stream = client.inner_get_response(messages=messages, stream=True, options={"tools": [my_tool]})
|
||||
assert isinstance(stream, ResponseStream)
|
||||
async for update in stream:
|
||||
updates.append(update)
|
||||
|
||||
# Find the function_call content - it should have agui_thread_id
|
||||
found = False
|
||||
for update in updates:
|
||||
for content in update.contents:
|
||||
if content.type == "function_call" and content.name == "my_tool":
|
||||
assert content.additional_properties is not None
|
||||
assert "agui_thread_id" in content.additional_properties
|
||||
found = True
|
||||
break
|
||||
assert found, "Expected to find function_call content for my_tool"
|
||||
|
||||
async def test_tool_call_args_id_mismatch_does_not_execute_current_client_tool(
|
||||
self, monkeypatch: MonkeyPatch
|
||||
) -> None:
|
||||
"""Mismatched TOOL_CALL_ARGS must not be rebound to the latest client tool."""
|
||||
executed: list[int] = []
|
||||
|
||||
@tool
|
||||
def danger_tool(amount: int) -> str:
|
||||
"""Record an invocation for the regression assertion."""
|
||||
executed.append(amount)
|
||||
return f"danger={amount}"
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "safe", "toolName": "safe_tool"},
|
||||
{"type": "TOOL_CALL_START", "toolCallId": "danger", "toolName": "danger_tool"},
|
||||
{"type": "TOOL_CALL_ARGS", "toolCallId": "safe", "delta": '{"amount": 100}'},
|
||||
{"type": "TOOL_CALL_END", "toolCallId": "safe"},
|
||||
{"type": "TOOL_CALL_END", "toolCallId": "danger"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
else:
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_2"},
|
||||
{"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "done"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_2"},
|
||||
]
|
||||
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
response = await client.get_response(
|
||||
[Message(role="user", contents=["Test"])],
|
||||
options={"tools": [danger_tool]},
|
||||
)
|
||||
|
||||
assert response.text == "done"
|
||||
assert executed == []
|
||||
|
||||
async def test_interrupt_options_transmission(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Interrupt option fields are forwarded to the HTTP service."""
|
||||
available_interrupts = [{"id": "req_1", "type": "request_info"}]
|
||||
expected_available_interrupts = [{"id": "req_1", "reason": "input_required"}]
|
||||
resume_payload = {"interrupts": [{"id": "req_1", "value": "approved"}]}
|
||||
expected_resume_payload = [{"interruptId": "req_1", "status": "resolved", "payload": "approved"}]
|
||||
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
assert kwargs.get("available_interrupts") == expected_available_interrupts
|
||||
assert kwargs.get("resume") == expected_resume_payload
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
messages = [Message(role="user", contents=["continue"])]
|
||||
options = {
|
||||
"available_interrupts": available_interrupts,
|
||||
"resume": resume_payload,
|
||||
}
|
||||
|
||||
response = await client.inner_get_response(messages=messages, options=options)
|
||||
assert response is not None
|
||||
|
||||
async def test_typed_interrupt_options_forward_canonical_protocol_shape(self, monkeypatch: MonkeyPatch) -> None:
|
||||
"""Typed interrupt options are forwarded as canonical protocol JSON."""
|
||||
mock_events = [
|
||||
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
|
||||
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
|
||||
]
|
||||
|
||||
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
|
||||
assert kwargs.get("available_interrupts") == [
|
||||
{
|
||||
"id": "approval_1",
|
||||
"reason": "tool_call",
|
||||
"toolCallId": "call_1",
|
||||
"responseSchema": {"type": "object"},
|
||||
}
|
||||
]
|
||||
assert kwargs.get("resume") == [
|
||||
{"interruptId": "approval_1", "status": "resolved", "payload": {"approved": True}}
|
||||
]
|
||||
for event in mock_events:
|
||||
yield event
|
||||
|
||||
client = StubAGUIChatClient(endpoint="http://localhost:8888/")
|
||||
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
|
||||
|
||||
options: dict[str, Any] = {
|
||||
"available_interrupts": [
|
||||
Interrupt(
|
||||
id="approval_1",
|
||||
reason="tool_call",
|
||||
tool_call_id="call_1",
|
||||
response_schema={"type": "object"},
|
||||
)
|
||||
],
|
||||
"resume": [ResumeEntry(interrupt_id="approval_1", status="resolved", payload={"approved": True})],
|
||||
}
|
||||
|
||||
response = await client.inner_get_response(
|
||||
messages=[Message(role="user", contents=["continue"])],
|
||||
options=options,
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user