agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: [BREAKING] Moved to a single get_response and run API (#3379)

Browse Source 
* WIP

* big update to new ResponseStream model

* fixed tests and typing

* fixed tests and typing

* fixed tools typevar import

* fix

* mypy fix

* mypy fixes and some cleanup

* fix missing quoted names

* and client

* fix  imports agui

* fix anthropic override

* fix agui

* fix ag ui

* fix import

* fix anthropic types

* fix mypy

* refactoring

* updated typing

* fix 3.11

* fixes

* redid layering of chat clients and agents

* redid layering of chat clients and agents

* Fix lint, type, and test issues after rebase

- Add @overload decorators to AgentProtocol.run() for type compatibility
- Add missing docstring params (middleware, function_invocation_configuration)
- Fix TODO format (TD002) by adding author tags
- Fix broken observability tests from upstream:
  - Replace non-existent use_instrumentation with direct instantiation
  - Replace non-existent use_agent_instrumentation with AgentTelemetryLayer mixin
  - Fix get_streaming_response to use get_response(stream=True)
  - Add AgentInitializationError import
  - Update streaming exception tests to match actual behavior

* Fix AgentExecutionException import error in test_agents.py

- Replace non-existent AgentExecutionException with AgentRunException

* Fix test import and asyncio deprecation issues

- Add 'tests' to pythonpath in ag-ui pyproject.toml for utils_test_ag_ui import
- Replace deprecated asyncio.get_event_loop().run_until_complete with asyncio.run

* Fix azure-ai test failures

- Update _prepare_options patching to use correct class path
- Fix test_to_azure_ai_agent_tools_web_search_missing_connection to clear env vars

* Convert ag-ui utils_test_ag_ui.py to conftest.py

- Move test utilities to conftest.py for proper pytest discovery
- Update all test imports to use conftest instead of utils_test_ag_ui
- Remove old utils_test_ag_ui.py file
- Revert pythonpath change in pyproject.toml

* fix: use relative imports for ag-ui test utilities

* fix agui

* Rename Bare*Client to Raw*Client and BaseChatClient

- Renamed BareChatClient to BaseChatClient (abstract base class)
- Renamed BareOpenAIChatClient to RawOpenAIChatClient
- Renamed BareOpenAIResponsesClient to RawOpenAIResponsesClient
- Renamed BareAzureAIClient to RawAzureAIClient
- Added warning docstrings to Raw* classes about layer ordering
- Updated README in samples/getting_started/agents/custom with layer docs
- Added test for span ordering with function calling

* Fix layer ordering: FunctionInvocationLayer before ChatTelemetryLayer

This ensures each inner LLM call gets its own telemetry span, resulting in
the correct span sequence: chat -> execute_tool -> chat

Updated all production clients and test mocks to use correct ordering:
- ChatMiddlewareLayer (first)
- FunctionInvocationLayer (second)
- ChatTelemetryLayer (third)
- BaseChatClient/Raw...Client (fourth)

* Remove run_stream usage

* Fix conversation_id propagation

* Python: Add BaseAgent implementation for Claude Agent SDK (#3509)

* Added ClaudeAgent implementation

* Updated streaming logic

* Small updates

* Small update

* Fixes

* Small fix

* Naming improvements

* Updated imports

* Addressed comments

* Updated package versions

* Update Claude agent connector layering

* fix test and plugin

* Store function middleware in invocation layer

* Fix telemetry streaming and ag-ui tests

* Remove legacy ag-ui tests folder

* updates

* Remove terminate flag from FunctionInvocationContext, use MiddlewareTermination instead

- Remove terminate attribute from FunctionInvocationContext
- Add result attribute to MiddlewareTermination to carry function results
- FunctionMiddlewarePipeline.execute() now lets MiddlewareTermination propagate
- _auto_invoke_function captures context.result in exception before re-raising
- _try_execute_function_calls catches MiddlewareTermination and sets should_terminate
- Fix handoff middleware to append to chat_client.function_middleware directly
- Update tests to use raise MiddlewareTermination instead of context.terminate
- Add middleware flow documentation in samples/concepts/tools/README.md
- Fix ag-ui to use FunctionMiddlewarePipeline instead of removed create_function_middleware_pipeline

* fix: remove references to removed terminate flag in purview tests, add type ignore

* fix: move _test_utils.py from package to test folder

* fix: call get_final_response() to trigger context provider notification in streaming test

* fix: correct broken links in tools README

* docs: clarify default middleware behavior in summary table

* fix: ensure inner stream result hooks are called when using map()/from_awaitable()

* Fix mypy type errors

* Address PR review comments on observability.py

- Remove TODO comment about unconsumed streams, add explanatory note instead
- Remove redundant _close_span cleanup hook (already called in _finalize_stream)
- Clarify behavior: cleanup hooks run after stream iteration, if stream is not
  consumed the span remains open until garbage collected

* Remove gen_ai.client.operation.duration from span attributes

Duration is a metrics-only attribute per OpenTelemetry semantic conventions.
It should be recorded to the histogram but not set as a span attribute.

* Remove duration from _get_response_attributes, pass directly to _capture_response

Duration is a metrics-only attribute. It's now passed directly to _capture_response
instead of being included in the attributes dict that gets set on the span.

* Remove redundant _close_span cleanup hook in AgentTelemetryLayer

_finalize_stream already calls _close_span() in its finally block,
so adding it as a separate cleanup hook is redundant.

* Use weakref.finalize to close span when stream is garbage collected

If a user creates a streaming response but never consumes it, the cleanup
hooks won't run. Now we register a weak reference finalizer that will close
the span when the stream object is garbage collected, ensuring spans don't
leak in this scenario.

* Fix _get_finalizers_from_stream to use _result_hooks attribute

Renamed function to _get_result_hooks_from_stream and fixed it to
look for the _result_hooks attribute which is the correct name in
ResponseStream class.

* Add missing asyncio import in test_request_info_mixin.py

* Fix leftover merge conflict marker in image_generation sample

* Update integration tests

* Fix integration tests: increase max_iterations from 1 to 2

Tests with tool_choice options require at least 2 iterations:
1. First iteration to get function call and execute the tool
2. Second iteration to get the final text response

With max_iterations=1, streaming tests would return early with only
the function call/result but no final text content.

* Fix duplicate function call error in conversation-based APIs

When using conversation_id (for Responses/Assistants APIs), the server
already has the function call message from the previous response. We
should only send the new function result message, not all messages
including the function call which would cause a duplicate ID error.

Fix: When conversation_id is set, only send the last message (the tool
result) instead of all response.messages.

* Add regression test for conversation_id propagation between tool iterations

Port test from PR #3664 with updates for new streaming API pattern.
Tests that conversation_id is properly updated in options dict during
function invocation loop iterations.

* Fix tool_choice=required to return after tool execution

When tool_choice is 'required', the user's intent is to force exactly one
tool call. After the tool executes, return immediately with the function
call and result - don't continue to call the model again.

This fixes integration tests that were failing with empty text responses
because with tool_choice=required, the model would keep returning function
calls instead of text.

Also adds regression tests for:
- conversation_id propagation between tool iterations (from PR #3664)
- tool_choice=required returns after tool execution

* Document tool_choice behavior in tools README

- Add table explaining tool_choice values (auto, none, required)
- Explain why tool_choice=required returns immediately after tool execution
- Add code example showing the difference between required and auto
- Update flow diagram to show the early return path for tool_choice=required

* Fix tool_choice=None behavior - don't default to 'auto'

Remove the hardcoded default of 'auto' for tool_choice in ChatAgent init.
When tool_choice is not specified (None), it will now not be sent to the
API, allowing the API's default behavior to be used.

Users who want tool_choice='auto' can still explicitly set it either in
default_options or at runtime.

Fixes #3585

* Fix tool_choice=none should not remove tools

In OpenAI Assistants client, tools were not being sent when
tool_choice='none'. This was incorrect - tool_choice='none' means
the model won't call tools, but tools should still be available
in the request (they may be used later in the conversation).

Fixes #3585

* Add test for tool_choice=none preserving tools

Adds a regression test to ensure that when tool_choice='none' is set but
tools are provided, the tools are still sent to the API. This verifies
the fix for #3585.

* Fix tool_choice=none should not remove tools in all clients

Apply the same fix to OpenAI Responses client and Azure AI client:
- OpenAI Responses: Remove else block that popped tool_choice/parallel_tool_calls
- Azure AI: Remove tool_choice != 'none' check when adding tools

When tool_choice='none', the model won't call tools, but tools should
still be sent to the API so they're available for future turns.

Also update README to clarify tool_choice=required supports multiple tools.

Fixes #3585

* Keep tool_choice even when tools is None

Move tool_choice processing outside of the 'if tools' block in OpenAI
Responses client so tool_choice is sent to the API even when no tools
are provided.

* Update test to match new parallel_tool_calls behavior

Changed test_prepare_options_removes_parallel_tool_calls_when_no_tools to
test_prepare_options_preserves_parallel_tool_calls_when_no_tools to reflect
that parallel_tool_calls is now preserved even when no tools are present,
consistent with the tool_choice behavior.

* Fix ChatMessage API and Role enum usage after rebase

- Update ChatMessage instantiation to use keyword args (role=, text=, contents=)
- Fix Role enum comparisons to use .value for string comparison
- Add created_at to AgentResponse in error handling
- Fix AgentResponse.from_updates -> from_agent_run_response_updates
- Fix DurableAgentStateMessage.from_chat_message to convert Role enum to string
- Add Role import where needed

* Fix additional ChatMessage API and method name changes

- Fix ChatMessage usage in workflow files (use text= instead of contents= for strings)
- Fix AgentResponse.from_updates -> from_agent_run_response_updates in workflow files
- Fix test files for ChatMessage and Role enum usage

* Fix remaining ChatMessage API usage in test files

* Fix more ChatMessage and Role API changes in source and test files

- Fix ChatMessage in _magentic.py replan method
- Fix Role enum comparison in test assertions
- Fix remaining test files with old ChatMessage syntax

* Fix ChatMessage and Role API changes across packages

- Add Role import where missing
- Fix ChatMessage signature: positional args to keyword args (role=, text=, contents=)
- Fix Role enum comparisons: .role.value instead of .role string
- Fix FinishReason enum usage in ag-ui event converters
- Rename AgentResponse.from_updates to from_agent_run_response_updates in ag-ui

Fixes API compatibility after Types API Review improvements merge

* Fix ChatMessage and Role API changes in github_copilot tests

* Fix ChatMessage and Role API changes in redis and github_copilot packages

- Fix redis provider: Role enum comparison using .value
- Fix redis tests: ChatMessage signature and Role comparisons
- Fix github_copilot tests: ChatMessage signature and Role comparisons
- Update docstring examples in redis chat message store

* Fix ChatMessage and Role API changes in devui package

- Fix executor: ChatMessage signature change
- Fix conversations: Role enum to string conversion in two places
- Fix tests: ChatMessage signatures and Role comparisons

* Fix ChatMessage and Role API changes in a2a and lab packages

- Fix a2a tests: Role comparisons and ChatMessage signatures
- Fix lab tau2 source: Role enum comparison in flip_messages, log_messages, sliding_window
- Fix lab tau2 tests: ChatMessage signatures and Role comparisons

* Remove duplicate test files from ag-ui/tests (tests are in ag_ui_tests)

* Fix ChatMessage and Role API changes across packages

After rebasing on upstream/main which merged PR #3647 (Types API Review
improvements), fix all packages to use the new API:

- ChatMessage: Use keyword args (role=, text=, contents=) instead of
  positional args
- Role: Compare using .value attribute since it's now an enum

Packages fixed:
- ag-ui: Fixed Role value extraction bugs in _message_adapters.py
- anthropic: Fixed ChatMessage and Role comparisons in tests
- azure-ai: Fixed Role comparison in _client.py
- azure-ai-search: Fixed ChatMessage and Role in source/tests
- bedrock: Fixed ChatMessage signatures in tests
- chatkit: Fixed ChatMessage and Role in source/tests
- copilotstudio: Fixed ChatMessage and Role in tests
- declarative: Fixed ChatMessage in _executors_agents.py
- mem0: Fixed ChatMessage and Role in source/tests
- purview: Fixed ChatMessage in source/tests

* Fix mypy errors for ChatMessage and Role API changes

- durabletask: Use str() fallback in role value extraction
- core: Fix ChatMessage in _orchestrator_helpers.py to use keyword args
- core: Add type ignore for _conversation_state.py contents deserialization
- ag-ui: Fix type ignore comments (call-overload instead of arg-type)
- azure-ai-search: Fix get_role_value type hint to accept Any
- lab: Move get_role_value to module level with Any type hint

* Improve CI test timeout configuration

- Increase job timeout from 10 to 15 minutes
- Reduce per-test timeout to 60s (was 900s/300s)
- Add --timeout_method thread for better timeout handling
- Add --timeout-verbose to see which tests are slow
- Reduce retries from 3 to 2 and delay from 10s to 5s

This ensures individual test timeouts are shorter than the job
timeout, providing better visibility when tests hang.

With 60s timeout and 2 retries, worst case per test is ~180s.

* Fix ChatMessage API usage in docstrings and source

- Fix ChatMessage positional args in docstrings: _serialization.py, _threads.py, _middleware.py
- Fix ChatMessage in tau2 runner.py
- Fix role comparison in _orchestrator_helpers.py to use .value
- Fix role comparison in _group_chat.py docstring example
- Fix role assertions in test_durable_entities.py to use .value

* Revert tool_choice/parallel_tool_calls changes - must be removed when no tools

OpenAI API requires tool_choice and parallel_tool_calls to only be
present when tools are specified. Restored the logic that removes
these options when there are no tools.

- Restored check in _chat_client.py to remove tool_choice and
  parallel_tool_calls when no tools present
- Restored same logic in _responses_client.py
- Reverted test to expect the correct behavior

* fixed issue in tests

* fix: resolve merge conflict markers in ag-ui tests

* fix: restructure ag-ui tests and fix Role/FinishReason to use string types

* fix: streaming function invocation and middleware termination

- Refactor streaming function invocation to use get_final_response() on inner streams
- Fix MiddlewareTermination to accept result parameter for passing results
- Fix _AutoHandoffMiddleware to use MiddlewareTermination instead of context.terminate
- Fix AgentMiddlewareLayer.run() to properly forward function/chat middleware
- Remove duplicate middleware registration in AgentMiddlewareLayer.__init__
- Fix exception handling in _auto_invoke_function to properly capture termination
- Fix mypy errors in core package
- Update tests to use stream=True parameter for unified run API

* fix all tests command

* Refactor integration tests to use pytest fixtures

- Merge testutils.py into conftest.py for azurefunctions integration tests
- Merge dt_testutils.py into conftest.py for durabletask integration tests
- Convert all integration tests to use fixtures instead of direct imports
  (fixes ModuleNotFoundError with --import-mode=importlib)
- Add sample_helper fixture for azurefunctions tests
- Add agent_client_factory and orchestration_helper fixtures for durabletask
- Integration tests now skip with descriptive messages when services unavailable
- Restructure devui tests into tests/devui/ with proper conftest.py
- Add test organization guidelines to CODING_STANDARD.md
- Remove __init__.py from test directories per pytest best practices

* Fix pytest_collection_modifyitems to only skip integration tests

The hook was skipping all tests in the test session, not just
integration tests. Now it only skips items in the integration_tests
directory.

* Fix mem0 tests failing on Python 3.13

Use patch.object on the imported module instead of @patch with string
path to ensure the mock takes effect regardless of import timing.

* fix mem0

* another attempt for mem0

* fix for mem0

* fix mem0

* Increase worker initialization wait time in durabletask tests

Increase from 2 to 8 seconds to allow time for:
- Python startup and module imports
- Azure OpenAI client creation
- Agent registration with DTS worker
- Worker connection to DTS

This helps prevent test failures in CI where the first tests may run
before the worker is fully ready to process requests.

* Fix streaming test to use ResponseStream with finalizer

The _consume_stream method now expects a ResponseStream that can provide
a final AgentResponse via get_final_response(). Update the test to use
ResponseStream with AgentResponse.from_updates as the finalizer.

* Fix MockToolCallingAgent to use new ResponseStream API and update samples

* small updates to run_stream to run

* fix sub workflow

* temp fix for az func test

---------

Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2026-02-05 20:09:58 +00:00
committed by GitHub
co-authored by Dmytro Struk
parent d1205896a1
commit 3dc59c83b5
372 changed files with 11583 additions and 9465 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/.cspell.json
+2
View File
@@ -38,6 +38,8 @@
"endregion",
"entra",
"faiss",
"finalizer",
"finalizers",
"genai",
"generativeai",
"hnsw",
python/CODING_STANDARD.md
+53
View File
@@ -484,3 +484,56 @@ 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`
python/packages/a2a/agent_framework_a2a/_agent.py
+69 -21
View File
@@ -4,8 +4,8 @@ import base64
import json
import re
import uuid
from collections.abc import AsyncIterable, Sequence
from typing import Any, Final, cast
from collections.abc import AsyncIterable, Awaitable, Sequence
from typing import Any, Final, Literal, cast, overload
import httpx
from a2a.client import Client, ClientConfig, ClientFactory, minimal_agent_card
@@ -32,10 +32,11 @@ from agent_framework import (
BaseAgent,
ChatMessage,
Content,
ResponseStream,
normalize_messages,
prepend_agent_framework_to_user_agent,
)
from agent_framework.observability import use_agent_instrumentation
from agent_framework.observability import AgentTelemetryLayer
__all__ = ["A2AAgent"]
@@ -56,8 +57,7 @@ def _get_uri_data(uri: str) -> str:
return match.group("base64_data")
@use_agent_instrumentation
class A2AAgent(BaseAgent):
class A2AAgent(AgentTelemetryLayer, BaseAgent):
"""Agent2Agent (A2A) protocol implementation.
Wraps an A2A Client to connect the Agent Framework with external A2A-compliant agents
@@ -184,44 +184,92 @@ class A2AAgent(BaseAgent):
if self._http_client is not None and self._close_http_client:
await self._http_client.aclose()
async def run(
@overload
def run(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Get a response from the agent.
This method returns the final result of the agent's execution
as a single AgentResponse object. The caller is blocked until
the final result is available.
as a single AgentResponse object when stream=False. When stream=True,
it returns a ResponseStream that yields AgentResponseUpdate objects.
Args:
messages: The message(s) to send to the agent.
Keyword Args:
stream: Whether to stream the response. Defaults to False.
thread: The conversation thread associated with the message(s).
kwargs: Additional keyword arguments.
Returns:
An agent response item.
When stream=False: An Awaitable[AgentResponse].
When stream=True: A ResponseStream of AgentResponseUpdate items.
"""
if stream:
return self._run_stream_impl(messages=messages, thread=thread, **kwargs)
return self._run_impl(messages=messages, thread=thread, **kwargs)
async def _run_impl(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse[Any]:
"""Non-streaming implementation of run."""
# Collect all updates and use framework to consolidate updates into response
updates = [update async for update in self.run_stream(messages, thread=thread, **kwargs)]
updates: list[AgentResponseUpdate] = []
async for update in self._stream_updates(messages, thread=thread, **kwargs):
updates.append(update)
return AgentResponse.from_updates(updates)
async def run_stream(
def _run_stream_impl(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Streaming implementation of run."""
def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]:
return AgentResponse.from_updates(list(updates))
return ResponseStream(self._stream_updates(messages, thread=thread, **kwargs), finalizer=_finalize)
async def _stream_updates(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
"""Run the agent as a stream.
This method will return the intermediate steps and final results of the
agent's execution as a stream of AgentResponseUpdate objects to the caller.
"""Internal method to stream updates from the A2A agent.
Args:
messages: The message(s) to send to the agent.
@@ -231,10 +279,10 @@ class A2AAgent(BaseAgent):
kwargs: Additional keyword arguments.
Yields:
An agent response item.
AgentResponseUpdate items from the A2A agent.
"""
messages = normalize_messages(messages)
a2a_message = self._prepare_message_for_a2a(messages[-1])
normalized_messages = normalize_messages(messages)
a2a_message = self._prepare_message_for_a2a(normalized_messages[-1])
response_stream = self.client.send_message(a2a_message)
python/packages/a2a/tests/test_a2a_agent.py
+7 -7
View File
@@ -295,7 +295,7 @@ def test_prepare_message_for_a2a_with_error_content(a2a_agent: A2AAgent) -> None
# Create ChatMessage with ErrorContent
error_content = Content.from_error(message="Test error message")
message = ChatMessage("user", [error_content])
message = ChatMessage(role="user", contents=[error_content])
# Convert to A2A message
a2a_message = a2a_agent._prepare_message_for_a2a(message)
@@ -310,7 +310,7 @@ def test_prepare_message_for_a2a_with_uri_content(a2a_agent: A2AAgent) -> None:
# Create ChatMessage with UriContent
uri_content = Content.from_uri(uri="http://example.com/file.pdf", media_type="application/pdf")
message = ChatMessage("user", [uri_content])
message = ChatMessage(role="user", contents=[uri_content])
# Convert to A2A message
a2a_message = a2a_agent._prepare_message_for_a2a(message)
@@ -326,7 +326,7 @@ def test_prepare_message_for_a2a_with_data_content(a2a_agent: A2AAgent) -> None:
# Create ChatMessage with DataContent (base64 data URI)
data_content = Content.from_uri(uri="data:text/plain;base64,SGVsbG8gV29ybGQ=", media_type="text/plain")
message = ChatMessage("user", [data_content])
message = ChatMessage(role="user", contents=[data_content])
# Convert to A2A message
a2a_message = a2a_agent._prepare_message_for_a2a(message)
@@ -340,20 +340,20 @@ def test_prepare_message_for_a2a_with_data_content(a2a_agent: A2AAgent) -> None:
def test_prepare_message_for_a2a_empty_contents_raises_error(a2a_agent: A2AAgent) -> None:
"""Test _prepare_message_for_a2a with empty contents raises ValueError."""
# Create ChatMessage with no contents
message = ChatMessage("user", [])
message = ChatMessage(role="user", contents=[])
# Should raise ValueError for empty contents
with raises(ValueError, match="ChatMessage.contents is empty"):
a2a_agent._prepare_message_for_a2a(message)
async def test_run_stream_with_message_response(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
"""Test run_stream() method with immediate Message response."""
async def test_run_streaming_with_message_response(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
"""Test run(stream=True) method with immediate Message response."""
mock_a2a_client.add_message_response("msg-stream-123", "Streaming response from agent!", "agent")
# Collect streaming updates
updates: list[AgentResponseUpdate] = []
async for update in a2a_agent.run_stream("Hello agent"):
async for update in a2a_agent.run("Hello agent", stream=True):
updates.append(update)
# Verify streaming response
python/packages/ag-ui/README.md
+1 -1
View File
@@ -46,7 +46,7 @@ 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_streaming_response("Hello!"):
async for update in client.get_response("Hello!", stream=True):
for content in update.contents:
if isinstance(content, TextContent):
print(content.text, end="", flush=True)
python/packages/ag-ui/agent_framework_ag_ui/_client.py
+94 -49
View File
@@ -6,9 +6,9 @@ import json
import logging
import sys
import uuid
from collections.abc import AsyncIterable, MutableSequence
from collections.abc import AsyncIterable, Awaitable, Mapping, MutableSequence, Sequence
from functools import wraps
from typing import TYPE_CHECKING, Any, Generic, cast
from typing import TYPE_CHECKING, Any, Generic, TypedDict, cast
import httpx
from agent_framework import (
@@ -18,10 +18,11 @@ from agent_framework import (
ChatResponseUpdate,
Content,
FunctionTool,
use_chat_middleware,
use_function_invocation,
ResponseStream,
)
from agent_framework.observability import use_instrumentation
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
@@ -42,6 +43,8 @@ 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(__name__)
@@ -67,35 +70,51 @@ TAGUIChatOptions = TypeVar(
def _apply_server_function_call_unwrap(chat_client: TBaseChatClient) -> TBaseChatClient:
"""Class decorator that unwraps server-side function calls after tool handling."""
original_get_streaming_response = chat_client.get_streaming_response
@wraps(original_get_streaming_response)
async def streaming_wrapper(self: Any, *args: Any, **kwargs: Any) -> AsyncIterable[ChatResponseUpdate]:
async for update in original_get_streaming_response(self, *args, **kwargs):
_unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], update.contents))
yield update
chat_client.get_streaming_response = streaming_wrapper # type: ignore[assignment]
original_get_response = chat_client.get_response
@wraps(original_get_response)
async def response_wrapper(self: Any, *args: Any, **kwargs: Any) -> ChatResponse:
response: ChatResponse[Any] = await original_get_response(self, *args, **kwargs) # type: ignore[var-annotated]
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
return response # type: ignore[no-any-return]
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
chat_client.get_response = response_wrapper # type: ignore[assignment]
return chat_client
@_apply_server_function_call_unwrap
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]):
class AGUIChatClient(
ChatMiddlewareLayer[TAGUIChatOptions],
FunctionInvocationLayer[TAGUIChatOptions],
ChatTelemetryLayer[TAGUIChatOptions],
BaseChatClient[TAGUIChatOptions],
Generic[TAGUIChatOptions],
):
"""Chat client for communicating with AG-UI compliant servers.
This client implements the BaseChatClient interface and automatically handles:
@@ -103,6 +122,7 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
- 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
@@ -115,10 +135,10 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
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, @use_function_invocation executes it locally
3. When LLM calls a client tool, function invocation executes it locally
4. Both client and server tools work together (hybrid pattern)
The wrapping ChatAgent's @use_function_invocation handles client tool execution
The wrapping ChatAgent's function invocation handles client tool execution
automatically when the server's LLM decides to call them.
Examples:
@@ -159,7 +179,7 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
.. code-block:: python
async for update in client.get_streaming_response("Tell me a story"):
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"):
@@ -196,6 +216,8 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
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,
**kwargs: Any,
) -> None:
"""Initialize the AG-UI chat client.
@@ -205,9 +227,16 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
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.
**kwargs: Additional arguments passed to BaseChatClient
"""
super().__init__(additional_properties=additional_properties, **kwargs)
super().__init__(
additional_properties=additional_properties,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
self._http_service = AGUIHttpService(
endpoint=endpoint,
http_client=http_client,
@@ -230,9 +259,10 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
"""Register a declaration-only placeholder so function invocation skips execution."""
config = getattr(self, "function_invocation_configuration", None)
if not config:
if not isinstance(config, dict):
return
if any(getattr(tool, "name", None) == tool_name for tool in config.additional_tools):
additional_tools = list(config.get("additional_tools", []))
if any(getattr(tool, "name", None) == tool_name for tool in additional_tools):
return
placeholder: FunctionTool[Any, Any] = FunctionTool(
@@ -240,7 +270,8 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
description="Server-managed tool placeholder (AG-UI)",
func=None,
)
config.additional_tools = list(config.additional_tools) + [placeholder]
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 # type: ignore[attr-defined]
@@ -250,7 +281,7 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}")
def _extract_state_from_messages(
self, messages: MutableSequence[ChatMessage]
self, messages: Sequence[ChatMessage]
) -> tuple[list[ChatMessage], dict[str, Any] | None]:
"""Extract state from last message if present.
@@ -297,7 +328,7 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
"""
return agent_framework_messages_to_agui(messages)
def _get_thread_id(self, options: dict[str, Any]) -> str:
def _get_thread_id(self, options: Mapping[str, Any]) -> str:
"""Get or generate thread ID from chat options.
Args:
@@ -317,43 +348,57 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
return thread_id
@override
async def _inner_get_response(
def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
stream: bool = False,
options: Mapping[str, Any],
**kwargs: Any,
) -> ChatResponse:
) -> 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
"""
return await ChatResponse.from_update_generator(
self._inner_get_streaming_response(
messages=messages,
options=options,
**kwargs,
if stream:
return ResponseStream(
self._streaming_impl(
messages=messages,
options=options,
**kwargs,
),
finalizer=ChatResponse.from_updates,
)
)
@override
async def _inner_get_streaming_response(
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: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
"""Internal method to get streaming response.
Keyword Args:
messages: List of chat messages
messages: Sequence of chat messages
options: Chat options for the request
**kwargs: Additional keyword arguments
@@ -368,7 +413,7 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
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 ChatAgent's @use_function_invocation wrapper
# Client tools execute via ChatAgent's function invocation wrapper
agui_tools = convert_tools_to_agui_format(options.get("tools"))
# Build set of client tool names (matches .NET clientToolSet)
@@ -415,12 +460,12 @@ class AGUIChatClient(BaseChatClient[TAGUIChatOptions], Generic[TAGUIChatOptions]
f"[AGUIChatClient] Function call: {content.name}, in client_tool_set: {content.name in client_tool_set}" # type: ignore[attr-defined]
)
if content.name in client_tool_set: # type: ignore[attr-defined]
# Client tool - let @use_function_invocation execute it
# Client tool - let function invocation execute it
if not content.additional_properties: # type: ignore[attr-defined]
content.additional_properties = {} # type: ignore[attr-defined]
content.additional_properties["agui_thread_id"] = thread_id # type: ignore[attr-defined]
else:
# Server tool - wrap so @use_function_invocation ignores it
# Server tool - wrap so function invocation ignores it
logger.debug(f"[AGUIChatClient] Wrapping server tool: {content.name}") # type: ignore[union-attr]
self._register_server_tool_placeholder(content.name) # type: ignore[arg-type]
update.contents[i] = Content(type="server_function_call", function_call=content) # type: ignore
python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
+6 -5
View File
@@ -590,7 +590,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
arguments=arguments,
)
)
chat_msg = ChatMessage("assistant", contents)
chat_msg = ChatMessage(role="assistant", contents=contents)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
@@ -620,14 +620,14 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
)
approval_contents.append(approval_response)
chat_msg = ChatMessage(role, approval_contents) # type: ignore[arg-type]
chat_msg = ChatMessage(role=role, contents=approval_contents) # type: ignore[call-overload]
else:
# Regular text message
content = msg.get("content", "")
if isinstance(content, str):
chat_msg = ChatMessage(role, [Content.from_text(text=content)])
chat_msg = ChatMessage(role=role, contents=[Content.from_text(text=content)]) # type: ignore[call-overload]
else:
chat_msg = ChatMessage(role, [Content.from_text(text=str(content))])
chat_msg = ChatMessage(role=role, contents=[Content.from_text(text=str(content))]) # type: ignore[call-overload]
if "id" in msg:
chat_msg.message_id = msg["id"]
@@ -671,7 +671,8 @@ def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str
continue
# Convert ChatMessage to AG-UI format
role = FRAMEWORK_TO_AGUI_ROLE.get(msg.role, "user")
role_value: str = msg.role if hasattr(msg.role, "value") else msg.role # type: ignore[assignment]
role = FRAMEWORK_TO_AGUI_ROLE.get(role_value, "user")
content_text = ""
tool_calls: list[dict[str, Any]] = []
python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_tooling.py
+2 -2
View File
@@ -79,8 +79,8 @@ def register_additional_client_tools(agent: "AgentProtocol", client_tools: list[
if chat_client is None:
return
if isinstance(chat_client, BaseChatClient) and chat_client.function_invocation_configuration is not None:
chat_client.function_invocation_configuration.additional_tools = client_tools
if isinstance(chat_client, BaseChatClient) and chat_client.function_invocation_configuration is not None: # type: ignore[attr-defined]
chat_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)")
python/packages/ag-ui/agent_framework_ag_ui/_run.py
+21 -6
View File
@@ -5,8 +5,9 @@
import json
import logging
import uuid
from collections.abc import Awaitable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, cast
from ag_ui.core import (
BaseEvent,
@@ -30,13 +31,15 @@ from agent_framework import (
Content,
prepare_function_call_results,
)
from agent_framework._middleware import extract_and_merge_function_middleware
from agent_framework._middleware import FunctionMiddlewarePipeline
from agent_framework._tools import (
FunctionInvocationConfiguration,
_collect_approval_responses, # type: ignore
_replace_approval_contents_with_results, # type: ignore
_try_execute_function_calls, # type: ignore
normalize_function_invocation_configuration,
)
from agent_framework._types import ResponseStream
from agent_framework.exceptions import AgentExecutionException
from ._message_adapters import normalize_agui_input_messages
from ._orchestration._predictive_state import PredictiveStateHandler
@@ -601,8 +604,13 @@ async def _resolve_approval_responses(
# Execute approved tool calls
if approved_responses and tools:
chat_client = getattr(agent, "chat_client", None)
config = getattr(chat_client, "function_invocation_configuration", None) or FunctionInvocationConfiguration()
middleware_pipeline = extract_and_merge_function_middleware(chat_client, run_kwargs)
config = normalize_function_invocation_configuration(
getattr(chat_client, "function_invocation_configuration", None)
)
middleware_pipeline = FunctionMiddlewarePipeline(
*getattr(chat_client, "function_middleware", ()),
*run_kwargs.get("middleware", ()),
)
# Filter out AG-UI-specific kwargs that should not be passed to tool execution
tool_kwargs = {k: v for k, v in run_kwargs.items() if k != "options"}
try:
@@ -862,7 +870,14 @@ async def run_agent_stream(
# Stream from agent - emit RunStarted after first update to get service IDs
run_started_emitted = False
all_updates: list[Any] = [] # Collect for structured output processing
async for update in agent.run_stream(messages, **run_kwargs):
response_stream = agent.run(messages, stream=True, **run_kwargs)
if isinstance(response_stream, ResponseStream):
stream = response_stream
else:
stream = await cast(Awaitable[ResponseStream[Any, Any]], response_stream)
if not isinstance(stream, ResponseStream):
raise AgentExecutionException("Chat client did not return a ResponseStream.")
async for update in stream:
# Collect updates for structured output processing
if response_format is not None:
all_updates.append(update)
python/packages/ag-ui/agent_framework_ag_ui/_types.py
+1 -1
View File
@@ -102,7 +102,7 @@ class AGUIChatOptions(ChatOptions[TResponseModel], Generic[TResponseModel], tota
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
@use_function_invocation middleware.
function invocation middleware.
tool_choice: How the model should use tools.
metadata: Metadata dict containing thread_id for conversation continuity.
python/packages/ag-ui/agent_framework_ag_ui/_utils.py
+3 -3
View File
@@ -165,7 +165,7 @@ def convert_agui_tools_to_agent_framework(
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 @use_function_invocation.
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.
@@ -183,7 +183,7 @@ def convert_agui_tools_to_agent_framework(
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 @use_function_invocation to return the function call
# 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[Any, Any] = FunctionTool(
name=tool_def.get("name", ""),
@@ -209,7 +209,7 @@ def convert_tools_to_agui_format(
This sends only the metadata (name, description, JSON schema) to the server.
The actual executable implementation stays on the client side.
The @use_function_invocation decorator handles client-side execution when
The function invocation mixin handles client-side execution when
the server requests a function.
Args:
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_steps_agent.py
+1 -1
View File
@@ -268,7 +268,7 @@ class TaskStepsAgentWithExecution:
# Stream completion
accumulated_text = ""
async for chunk in chat_client.get_streaming_response(messages=messages):
async for chunk in chat_client.get_response(messages=messages, stream=True):
# chunk is ChatResponseUpdate
if hasattr(chunk, "text") and chunk.text:
accumulated_text += chunk.text
python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/backend_tool_rendering.py
+4 -1
View File
@@ -2,6 +2,9 @@
"""Backend tool rendering endpoint."""
from typing import Any, cast
from agent_framework._clients import ChatClientProtocol
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
from agent_framework.azure import AzureOpenAIChatClient
from fastapi import FastAPI
@@ -16,7 +19,7 @@ def register_backend_tool_rendering(app: FastAPI) -> None:
app: The FastAPI application.
"""
# Create a chat client and call the factory function
chat_client = AzureOpenAIChatClient()
chat_client = cast(ChatClientProtocol[Any], AzureOpenAIChatClient())
add_agent_framework_fastapi_endpoint(
app,
python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py
+5 -3
View File
@@ -4,10 +4,11 @@
import logging
import os
from typing import cast
import uvicorn
from agent_framework import ChatOptions
from agent_framework._clients import BaseChatClient
from agent_framework._clients import ChatClientProtocol
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
from agent_framework.anthropic import AnthropicClient
from agent_framework.azure import AzureOpenAIChatClient
@@ -64,8 +65,9 @@ app.add_middleware(
# 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
chat_client: BaseChatClient[ChatOptions] = (
AnthropicClient() if os.getenv("CHAT_CLIENT", "").lower() == "anthropic" else AzureOpenAIChatClient()
chat_client: ChatClientProtocol[ChatOptions] = cast(
ChatClientProtocol[ChatOptions],
AnthropicClient() if os.getenv("CHAT_CLIENT", "").lower() == "anthropic" else AzureOpenAIChatClient(),
)
# Agentic Chat - basic chat agent
python/packages/ag-ui/getting_started/README.md
+2 -2
View File
@@ -323,7 +323,7 @@ async def main():
# Use metadata to maintain conversation continuity
metadata = {"thread_id": thread_id} if thread_id else None
async for update in client.get_streaming_response(message, metadata=metadata):
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")
@@ -353,7 +353,7 @@ if __name__ == "__main__":
- **`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_streaming_response()` for real-time streaming or `get_response()` for non-streaming
- **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 (ChatAgent, tools, etc.)
- **Hybrid Tool Execution**: Supports both client-side and server-side tools executing together in the same conversation
python/packages/ag-ui/getting_started/client.py
+11 -3
View File
@@ -9,7 +9,9 @@ standard chat interface.
import asyncio
import os
from typing import cast
from agent_framework import ChatResponse, ChatResponseUpdate, ResponseStream
from agent_framework.ag_ui import AGUIChatClient
@@ -41,7 +43,13 @@ async def main():
# Use metadata to maintain conversation continuity
metadata = {"thread_id": thread_id} if thread_id else None
async for update in client.get_streaming_response(message, metadata=metadata):
stream = client.get_response(
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")
@@ -51,8 +59,8 @@ async def main():
# Display text content as it streams
for content in update.contents:
if hasattr(content, "text") and content.text: # type: ignore[attr-defined]
print(f"\033[96m{content.text}\033[0m", end="", flush=True) # type: ignore[attr-defined]
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:
python/packages/ag-ui/getting_started/client_advanced.py
+9 -2
View File
@@ -11,8 +11,9 @@ This example demonstrates advanced AGUIChatClient features including:
import asyncio
import os
from typing import cast
from agent_framework import tool
from agent_framework import ChatResponse, ChatResponseUpdate, ResponseStream, tool
from agent_framework.ag_ui import AGUIChatClient
@@ -69,7 +70,13 @@ async def streaming_example(client: AGUIChatClient, thread_id: str | None = None
print("\nUser: Tell me a short joke\n")
print("Assistant: ", end="", flush=True)
async for update in client.get_streaming_response("Tell me a short joke", metadata=metadata):
stream = client.get_response(
"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")
python/packages/ag-ui/getting_started/client_with_agent.py
+13 -9
View File
@@ -6,11 +6,11 @@ This demonstrates the HYBRID pattern matching .NET AGUIClient implementation:
1. AgentThread Pattern (like .NET):
- Create thread with agent.get_new_thread()
- Pass thread to agent.run_stream() on each turn
- Pass thread to agent.run(stream=True) on each turn
- Thread automatically maintains conversation history via message_store
2. Hybrid Tool Execution:
- AGUIChatClient has @use_function_invocation decorator
- 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
@@ -63,7 +63,7 @@ async def main():
Python equivalent:
- agent = ChatAgent(chat_client=AGUIChatClient(...), tools=[...])
- thread = agent.get_new_thread() # Creates thread with message_store
- agent.run_stream(message, thread=thread) # Thread accumulates history
- agent.run(message, stream=True, thread=thread) # Thread accumulates history
"""
server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
@@ -73,7 +73,7 @@ async def main():
print(f"\nServer: {server_url}")
print("\nThis example demonstrates:")
print(" 1. AgentThread maintains conversation state (like .NET)")
print(" 2. Client-side tools execute locally via @use_function_invocation")
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")
@@ -97,35 +97,39 @@ async def main():
# Turn 1: Introduce
print("\nUser: My name is Alice and I live in Seattle\n")
async for chunk in agent.run_stream("My name is Alice and I live in Seattle", thread=thread):
async for chunk in agent.run("My name is Alice and I live in Seattle", stream=True, thread=thread):
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_stream("What's my name?", thread=thread):
async for chunk in agent.run("What's my name?", stream=True, thread=thread):
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_stream("Where do I live?", thread=thread):
async for chunk in agent.run("Where do I live?", stream=True, thread=thread):
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_stream("What's the weather forecast for today in Seattle?", thread=thread):
async for chunk in agent.run(
"What's the weather forecast for today in Seattle?",
stream=True,
thread=thread,
):
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_stream("What time zone is Seattle in?", thread=thread):
async for chunk in agent.run("What time zone is Seattle in?", stream=True, thread=thread):
if chunk.text:
print(chunk.text, end="", flush=True)
print("\n")
python/packages/ag-ui/getting_started/server.py
+1 -1
View File
@@ -112,7 +112,7 @@ def get_time_zone(location: str) -> str:
# - 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 @use_function_invocation on AGUIChatClient will execute it client-side.
# and the function invocation mixin on AGUIChatClient will execute it client-side.
# This matches the .NET AG-UI hybrid execution pattern.
agent = ChatAgent(
name="AGUIAssistant",
python/packages/ag-ui/pyproject.toml
+3 -4
View File
@@ -31,7 +31,6 @@ dependencies = [
[project.optional-dependencies]
dev = [
"pytest>=8.0.0",
"pytest-asyncio>=0.24.0",
"httpx>=0.27.0",
]
@@ -44,7 +43,7 @@ packages = ["agent_framework_ag_ui", "agent_framework_ag_ui_examples"]
[tool.pytest.ini_options]
asyncio_mode = "auto"
testpaths = ["tests"]
testpaths = ["tests/ag_ui"]
pythonpath = ["."]
[tool.ruff]
@@ -62,7 +61,7 @@ warn_unused_configs = true
disallow_untyped_defs = false
[tool.pyright]
exclude = ["tests", "examples"]
exclude = ["tests", "tests/ag_ui", "examples"]
typeCheckingMode = "basic"
[tool.poe]
@@ -71,4 +70,4 @@ include = "../../shared_tasks.toml"
[tool.poe.tasks]
mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_ag_ui"
test = "pytest --cov=agent_framework_ag_ui --cov-report=term-missing:skip-covered tests"
test = "pytest --cov=agent_framework_ag_ui --cov-report=term-missing:skip-covered tests/ag_ui"
python/packages/ag-ui/tests/ag_ui/conftest.py
+243
View File
@@ -0,0 +1,243 @@
# 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 types import SimpleNamespace
from typing import Any, Generic, Literal, cast, overload
import pytest
from agent_framework import (
AgentProtocol,
AgentResponse,
AgentResponseUpdate,
AgentThread,
BaseChatClient,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
)
from agent_framework._clients import TOptions_co
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]]
class StreamingChatClientStub(
ChatMiddlewareLayer[TOptions_co],
FunctionInvocationLayer[TOptions_co],
ChatTelemetryLayer[TOptions_co],
BaseChatClient[TOptions_co],
Generic[TOptions_co],
):
"""Typed streaming stub that satisfies ChatClientProtocol."""
def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
super().__init__(function_middleware=[])
self._stream_fn = stream_fn
self._response_fn = response_fn
self.last_thread: AgentThread | None = None
self.last_service_thread_id: str | None = None
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: Literal[False] = ...,
options: ChatOptions[Any],
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]]: ...
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: Literal[False] = ...,
options: TOptions_co | ChatOptions[None] | None = ...,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]]: ...
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: Literal[True],
options: TOptions_co | ChatOptions[Any] | None = ...,
**kwargs: Any,
) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: bool = False,
options: TOptions_co | ChatOptions[Any] | None = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
self.last_thread = kwargs.get("thread")
self.last_service_thread_id = self.last_thread.service_thread_id if self.last_thread else None
return cast(
Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
super().get_response(
messages=messages,
stream=cast(Literal[True, False], stream),
options=options,
**kwargs,
),
)
@override
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
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[ChatMessage], 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=[ChatMessage(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[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
for update in updates:
yield update
return _stream
class StubAgent(AgentProtocol):
"""Minimal AgentProtocol 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,
chat_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.chat_client = chat_client or SimpleNamespace(function_invocation_configuration=None)
self.messages_received: list[Any] = []
self.tools_received: list[Any] | None = None
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
if stream:
async def _stream() -> AsyncIterator[AgentResponseUpdate]:
self.messages_received = [] if messages is None else list(messages) # type: ignore[arg-type]
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 get_new_thread(self, **kwargs: Any) -> AgentThread:
return AgentThread()
# Fixtures
@pytest.fixture
def streaming_chat_client_stub() -> type[ChatClientProtocol]:
"""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[AgentProtocol]:
"""Return the StubAgent class for creating test instances."""
return StubAgent # type: ignore[return-value]
python/packages/ag-ui/tests/test_ag_ui_client.py → python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py
+24 -28
View File
@@ -3,7 +3,7 @@
"""Tests for AGUIChatClient."""
import json
from collections.abc import AsyncGenerator, AsyncIterable, MutableSequence
from collections.abc import AsyncGenerator, Awaitable, MutableSequence
from typing import Any
from agent_framework import (
@@ -12,6 +12,7 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
ResponseStream,
tool,
)
from pytest import MonkeyPatch
@@ -42,18 +43,11 @@ class TestableAGUIChatClient(AGUIChatClient):
"""Expose thread id helper."""
return self._get_thread_id(options)
async def inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any]
) -> AsyncIterable[ChatResponseUpdate]:
"""Proxy to protected streaming call."""
async for update in self._inner_get_streaming_response(messages=messages, options=options):
yield update
async def inner_get_response(
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any]
) -> ChatResponse:
def inner_get_response(
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], stream: bool = False
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
"""Proxy to protected response call."""
return await self._inner_get_response(messages=messages, options=options)
return self._inner_get_response(messages=messages, options=options, stream=stream)
class TestAGUIChatClient:
@@ -75,8 +69,8 @@ class TestAGUIChatClient:
"""Test state extraction when no state is present."""
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
messages = [
ChatMessage("user", ["Hello"]),
ChatMessage("assistant", ["Hi there"]),
ChatMessage(role="user", text="Hello"),
ChatMessage(role="assistant", text="Hi there"),
]
result_messages, state = client.extract_state_from_messages(messages)
@@ -95,7 +89,7 @@ class TestAGUIChatClient:
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
messages = [
ChatMessage("user", ["Hello"]),
ChatMessage(role="user", text="Hello"),
ChatMessage(
role="user",
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
@@ -133,8 +127,8 @@ class TestAGUIChatClient:
"""Test message conversion to AG-UI format."""
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
messages = [
ChatMessage("user", ["What is the weather?"]),
ChatMessage("assistant", ["Let me check."], message_id="msg_123"),
ChatMessage(role="user", text="What is the weather?"),
ChatMessage(role="assistant", text="Let me check.", message_id="msg_123"),
]
agui_messages = client.convert_messages_to_agui_format(messages)
@@ -165,7 +159,7 @@ class TestAGUIChatClient:
assert thread_id.startswith("thread_")
assert len(thread_id) > 7
async def test_get_streaming_response(self, monkeypatch: MonkeyPatch) -> None:
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"},
@@ -181,11 +175,11 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage("user", ["Test message"])]
messages = [ChatMessage(role="user", text="Test message")]
chat_options = ChatOptions()
updates: list[ChatResponseUpdate] = []
async for update in client.inner_get_streaming_response(messages=messages, options=chat_options):
async for update in client._inner_get_response(messages=messages, stream=True, options=chat_options):
updates.append(update)
assert len(updates) == 4
@@ -214,7 +208,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage("user", ["Test message"])]
messages = [ChatMessage(role="user", text="Test message")]
chat_options = {}
response = await client.inner_get_response(messages=messages, options=chat_options)
@@ -227,7 +221,7 @@ class TestAGUIChatClient:
"""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, @use_function_invocation decorator
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
@@ -257,7 +251,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage("user", ["Test with tools"])]
messages = [ChatMessage(role="user", text="Test with tools")]
chat_options = ChatOptions(tools=[test_tool])
response = await client.inner_get_response(messages=messages, options=chat_options)
@@ -281,10 +275,10 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage("user", ["Test server tool execution"])]
messages = [ChatMessage(role="user", text="Test server tool execution")]
updates: list[ChatResponseUpdate] = []
async for update in client.get_streaming_response(messages):
async for update in client.get_response(messages, stream=True):
updates.append(update)
function_calls = [
@@ -323,9 +317,11 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage("user", ["Test server tool execution"])]
messages = [ChatMessage(role="user", text="Test server tool execution")]
async for _ in client.get_streaming_response(messages, options={"tool_choice": "auto", "tools": [client_tool]}):
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:
@@ -337,7 +333,7 @@ class TestAGUIChatClient:
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
messages = [
ChatMessage("user", ["Hello"]),
ChatMessage(role="user", text="Hello"),
ChatMessage(
role="user",
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
python/packages/ag-ui/tests/test_agent_wrapper_comprehensive.py → python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py
+53 -68
View File
@@ -3,20 +3,15 @@
"""Comprehensive tests for AgentFrameworkAgent (_agent.py)."""
import json
import sys
from collections.abc import AsyncIterator, MutableSequence
from pathlib import Path
from typing import Any
import pytest
from agent_framework import ChatAgent, ChatMessage, ChatOptions, ChatResponseUpdate, Content
from pydantic import BaseModel
sys.path.insert(0, str(Path(__file__).parent))
from utils_test_ag_ui import StreamingChatClientStub
async def test_agent_initialization_basic():
async def test_agent_initialization_basic(streaming_chat_client_stub):
"""Test basic agent initialization without state schema."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -26,7 +21,7 @@ async def test_agent_initialization_basic():
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent[ChatOptions](
chat_client=StreamingChatClientStub(stream_fn),
chat_client=streaming_chat_client_stub(stream_fn),
name="test_agent",
instructions="Test",
)
@@ -38,7 +33,7 @@ async def test_agent_initialization_basic():
assert wrapper.config.predict_state_config == {}
async def test_agent_initialization_with_state_schema():
async def test_agent_initialization_with_state_schema(streaming_chat_client_stub):
"""Test agent initialization with state_schema."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -47,14 +42,14 @@ async def test_agent_initialization_with_state_schema():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
state_schema: dict[str, dict[str, Any]] = {"document": {"type": "string"}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
assert wrapper.config.state_schema == state_schema
async def test_agent_initialization_with_predict_state_config():
async def test_agent_initialization_with_predict_state_config(streaming_chat_client_stub):
"""Test agent initialization with predict_state_config."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -63,14 +58,14 @@ async def test_agent_initialization_with_predict_state_config():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
wrapper = AgentFrameworkAgent(agent=agent, predict_state_config=predict_config)
assert wrapper.config.predict_state_config == predict_config
async def test_agent_initialization_with_pydantic_state_schema():
async def test_agent_initialization_with_pydantic_state_schema(streaming_chat_client_stub):
"""Test agent initialization when state_schema is provided as Pydantic model/class."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -83,7 +78,7 @@ async def test_agent_initialization_with_pydantic_state_schema():
document: str
tags: list[str] = []
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper_class_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState)
wrapper_instance_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState(document="hi"))
@@ -93,7 +88,7 @@ async def test_agent_initialization_with_pydantic_state_schema():
assert wrapper_instance_schema.config.state_schema == expected_properties
async def test_run_started_event_emission():
async def test_run_started_event_emission(streaming_chat_client_stub):
"""Test RunStartedEvent is emitted at start of run."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -102,7 +97,7 @@ async def test_run_started_event_emission():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -117,7 +112,7 @@ async def test_run_started_event_emission():
assert events[0].thread_id is not None
async def test_predict_state_custom_event_emission():
async def test_predict_state_custom_event_emission(streaming_chat_client_stub):
"""Test PredictState CustomEvent is emitted when predict_state_config is present."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -126,7 +121,7 @@ async def test_predict_state_custom_event_emission():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
predict_config = {
"document": {"tool": "write_doc", "tool_argument": "content"},
"summary": {"tool": "summarize", "tool_argument": "text"},
@@ -149,7 +144,7 @@ async def test_predict_state_custom_event_emission():
assert {"state_key": "summary", "tool": "summarize", "tool_argument": "text"} in predict_value
async def test_initial_state_snapshot_with_schema():
async def test_initial_state_snapshot_with_schema(streaming_chat_client_stub):
"""Test initial StateSnapshotEvent emission when state_schema present."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -158,7 +153,7 @@ async def test_initial_state_snapshot_with_schema():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
state_schema = {"document": {"type": "string"}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
@@ -179,7 +174,7 @@ async def test_initial_state_snapshot_with_schema():
assert snapshot_events[0].snapshot == {"document": "Initial content"}
async def test_state_initialization_object_type():
async def test_state_initialization_object_type(streaming_chat_client_stub):
"""Test state initialization with object type in schema."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -188,7 +183,7 @@ async def test_state_initialization_object_type():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
state_schema: dict[str, dict[str, Any]] = {"recipe": {"type": "object", "properties": {}}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
@@ -206,7 +201,7 @@ async def test_state_initialization_object_type():
assert snapshot_events[0].snapshot == {"recipe": {}}
async def test_state_initialization_array_type():
async def test_state_initialization_array_type(streaming_chat_client_stub):
"""Test state initialization with array type in schema."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -215,7 +210,7 @@ async def test_state_initialization_array_type():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
state_schema: dict[str, dict[str, Any]] = {"steps": {"type": "array", "items": {}}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
@@ -233,7 +228,7 @@ async def test_state_initialization_array_type():
assert snapshot_events[0].snapshot == {"steps": []}
async def test_run_finished_event_emission():
async def test_run_finished_event_emission(streaming_chat_client_stub):
"""Test RunFinishedEvent is emitted at end of run."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -242,7 +237,7 @@ async def test_run_finished_event_emission():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -255,7 +250,7 @@ async def test_run_finished_event_emission():
assert events[-1].type == "RUN_FINISHED"
async def test_tool_result_confirm_changes_accepted():
async def test_tool_result_confirm_changes_accepted(streaming_chat_client_stub):
"""Test confirm_changes tool result handling when accepted."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -264,7 +259,7 @@ async def test_tool_result_confirm_changes_accepted():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Document updated")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(
agent=agent,
state_schema={"document": {"type": "string"}},
@@ -302,7 +297,7 @@ async def test_tool_result_confirm_changes_accepted():
assert confirmation_found, f"No confirmation in deltas: {[e.delta for e in text_content_events]}"
async def test_tool_result_confirm_changes_rejected():
async def test_tool_result_confirm_changes_rejected(streaming_chat_client_stub):
"""Test confirm_changes tool result handling when rejected."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -311,7 +306,7 @@ async def test_tool_result_confirm_changes_rejected():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="OK")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Simulate tool result message with rejection
@@ -336,7 +331,7 @@ async def test_tool_result_confirm_changes_rejected():
assert any("what would you like me to change" in e.delta.lower() for e in text_content_events)
async def test_tool_result_function_approval_accepted():
async def test_tool_result_function_approval_accepted(streaming_chat_client_stub):
"""Test function approval tool result when steps are accepted."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -345,7 +340,7 @@ async def test_tool_result_function_approval_accepted():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="OK")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Simulate tool result with multiple steps
@@ -382,7 +377,7 @@ async def test_tool_result_function_approval_accepted():
assert "create calendar event" in full_text.lower()
async def test_tool_result_function_approval_rejected():
async def test_tool_result_function_approval_rejected(streaming_chat_client_stub):
"""Test function approval tool result when rejected."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -391,7 +386,7 @@ async def test_tool_result_function_approval_rejected():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="OK")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Simulate tool result rejection with steps
@@ -419,7 +414,7 @@ async def test_tool_result_function_approval_rejected():
assert any("what would you like me to change about the plan" in e.delta.lower() for e in text_content_events)
async def test_thread_metadata_tracking():
async def test_thread_metadata_tracking(streaming_chat_client_stub):
"""Test that thread metadata includes ag_ui_thread_id and ag_ui_run_id.
AG-UI internal metadata is stored in thread.metadata for orchestration,
@@ -427,21 +422,16 @@ async def test_thread_metadata_tracking():
"""
from agent_framework.ag_ui import AgentFrameworkAgent
captured_thread: dict[str, Any] = {}
captured_options: dict[str, Any] = {}
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
# Capture the thread object from kwargs
thread = kwargs.get("thread")
if thread and hasattr(thread, "metadata"):
captured_thread["metadata"] = thread.metadata
# Capture options to verify internal keys are NOT passed to chat client
captured_options.update(options)
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {
@@ -455,7 +445,8 @@ async def test_thread_metadata_tracking():
events.append(event)
# AG-UI internal metadata should be stored in thread.metadata
thread_metadata = captured_thread.get("metadata", {})
thread = agent.chat_client.last_thread
thread_metadata = thread.metadata if thread and hasattr(thread, "metadata") else {}
assert thread_metadata.get("ag_ui_thread_id") == "test_thread_123"
assert thread_metadata.get("ag_ui_run_id") == "test_run_456"
@@ -465,7 +456,7 @@ async def test_thread_metadata_tracking():
assert "ag_ui_run_id" not in options_metadata
async def test_state_context_injection():
async def test_state_context_injection(streaming_chat_client_stub):
"""Test that current state is injected into thread metadata.
AG-UI internal metadata (including current_state) is stored in thread.metadata
@@ -473,21 +464,16 @@ async def test_state_context_injection():
"""
from agent_framework_ag_ui import AgentFrameworkAgent
captured_thread: dict[str, Any] = {}
captured_options: dict[str, Any] = {}
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
# Capture the thread object from kwargs
thread = kwargs.get("thread")
if thread and hasattr(thread, "metadata"):
captured_thread["metadata"] = thread.metadata
# Capture options to verify internal keys are NOT passed to chat client
captured_options.update(options)
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(
agent=agent,
state_schema={"document": {"type": "string"}},
@@ -503,7 +489,8 @@ async def test_state_context_injection():
events.append(event)
# Current state should be stored in thread.metadata
thread_metadata = captured_thread.get("metadata", {})
thread = agent.chat_client.last_thread
thread_metadata = thread.metadata if thread and hasattr(thread, "metadata") else {}
current_state = thread_metadata.get("current_state")
if isinstance(current_state, str):
current_state = json.loads(current_state)
@@ -514,7 +501,7 @@ async def test_state_context_injection():
assert "current_state" not in options_metadata
async def test_no_messages_provided():
async def test_no_messages_provided(streaming_chat_client_stub):
"""Test handling when no messages are provided."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -523,7 +510,7 @@ async def test_no_messages_provided():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data: dict[str, Any] = {"messages": []}
@@ -538,7 +525,7 @@ async def test_no_messages_provided():
assert events[-1].type == "RUN_FINISHED"
async def test_message_end_event_emission():
async def test_message_end_event_emission(streaming_chat_client_stub):
"""Test TextMessageEndEvent is emitted for assistant messages."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -547,7 +534,7 @@ async def test_message_end_event_emission():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello world")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -566,7 +553,7 @@ async def test_message_end_event_emission():
assert end_index < finished_index
async def test_error_handling_with_exception():
async def test_error_handling_with_exception(streaming_chat_client_stub):
"""Test that exceptions during agent execution are re-raised."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -577,7 +564,7 @@ async def test_error_handling_with_exception():
yield ChatResponseUpdate(contents=[])
raise RuntimeError("Simulated failure")
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -587,7 +574,7 @@ async def test_error_handling_with_exception():
pass
async def test_json_decode_error_in_tool_result():
async def test_json_decode_error_in_tool_result(streaming_chat_client_stub):
"""Test handling of orphaned tool result - should be sanitized out."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -598,7 +585,7 @@ async def test_json_decode_error_in_tool_result():
yield ChatResponseUpdate(contents=[])
raise AssertionError("ChatClient should not be called with orphaned tool result")
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Send invalid JSON as tool result without preceding tool call
@@ -624,7 +611,7 @@ async def test_json_decode_error_in_tool_result():
assert len(tool_events) == 0
async def test_agent_with_use_service_thread_is_false():
async def test_agent_with_use_service_thread_is_false(streaming_chat_client_stub):
"""Test that when use_service_thread is False, the AgentThread used to run the agent is NOT set to the service thread ID."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -633,14 +620,11 @@ async def test_agent_with_use_service_thread_is_false():
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
nonlocal request_service_thread_id
thread = kwargs.get("thread")
request_service_thread_id = thread.service_thread_id if thread else None
yield ChatResponseUpdate(
contents=[Content.from_text(text="Response")], response_id="resp_67890", conversation_id="conv_12345"
)
agent = ChatAgent(chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=False)
input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
@@ -651,7 +635,7 @@ async def test_agent_with_use_service_thread_is_false():
assert request_service_thread_id is None # type: ignore[attr-defined] (service_thread_id should be set)
async def test_agent_with_use_service_thread_is_true():
async def test_agent_with_use_service_thread_is_true(streaming_chat_client_stub):
"""Test that when use_service_thread is True, the AgentThread used to run the agent is set to the service thread ID."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -667,7 +651,7 @@ async def test_agent_with_use_service_thread_is_true():
contents=[Content.from_text(text="Response")], response_id="resp_67890", conversation_id="conv_12345"
)
agent = ChatAgent(chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(chat_client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=True)
input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
@@ -675,10 +659,11 @@ async def test_agent_with_use_service_thread_is_true():
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
request_service_thread_id = agent.chat_client.last_service_thread_id
assert request_service_thread_id == "conv_123456" # type: ignore[attr-defined] (service_thread_id should be set)
async def test_function_approval_mode_executes_tool():
async def test_function_approval_mode_executes_tool(streaming_chat_client_stub):
"""Test that function approval with approval_mode='always_require' sends the correct messages."""
from agent_framework import tool
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -702,7 +687,7 @@ async def test_function_approval_mode_executes_tool():
yield ChatResponseUpdate(contents=[Content.from_text(text="Processing completed")])
agent = ChatAgent(
chat_client=StreamingChatClientStub(stream_fn),
chat_client=streaming_chat_client_stub(stream_fn),
name="test_agent",
instructions="Test",
tools=[get_datetime],
@@ -769,7 +754,7 @@ async def test_function_approval_mode_executes_tool():
)
async def test_function_approval_mode_rejection():
async def test_function_approval_mode_rejection(streaming_chat_client_stub):
"""Test that function approval rejection creates a rejection response."""
from agent_framework import tool
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -795,7 +780,7 @@ async def test_function_approval_mode_rejection():
agent = ChatAgent(
name="test_agent",
instructions="Test",
chat_client=StreamingChatClientStub(stream_fn),
chat_client=streaming_chat_client_stub(stream_fn),
tools=[delete_all_data],
)
wrapper = AgentFrameworkAgent(agent=agent)
python/packages/ag-ui/tests/test_endpoint.py → python/packages/ag-ui/tests/ag_ui/test_endpoint.py
+31 -30
View File
@@ -3,9 +3,8 @@
"""Tests for FastAPI endpoint creation (_endpoint.py)."""
import json
import sys
from pathlib import Path
import pytest
from agent_framework import ChatAgent, ChatResponseUpdate, Content
from fastapi import FastAPI, Header, HTTPException
from fastapi.params import Depends
@@ -14,17 +13,19 @@ from fastapi.testclient import TestClient
from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
from agent_framework_ag_ui._agent import AgentFrameworkAgent
sys.path.insert(0, str(Path(__file__).parent))
from utils_test_ag_ui import StreamingChatClientStub, stream_from_updates
def build_chat_client(response_text: str = "Test response") -> StreamingChatClientStub:
@pytest.fixture
def build_chat_client(streaming_chat_client_stub, stream_from_updates_fixture):
"""Create a typed chat client stub for endpoint tests."""
updates = [ChatResponseUpdate(contents=[Content.from_text(text=response_text)])]
return StreamingChatClientStub(stream_from_updates(updates))
def _build(response_text: str = "Test response"):
updates = [ChatResponseUpdate(contents=[Content.from_text(text=response_text)])]
return streaming_chat_client_stub(stream_from_updates_fixture(updates))
return _build
async def test_add_endpoint_with_agent_protocol():
async def test_add_endpoint_with_agent_protocol(build_chat_client):
"""Test adding endpoint with raw AgentProtocol."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -38,7 +39,7 @@ async def test_add_endpoint_with_agent_protocol():
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
async def test_add_endpoint_with_wrapped_agent():
async def test_add_endpoint_with_wrapped_agent(build_chat_client):
"""Test adding endpoint with pre-wrapped AgentFrameworkAgent."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -53,7 +54,7 @@ async def test_add_endpoint_with_wrapped_agent():
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
async def test_endpoint_with_state_schema():
async def test_endpoint_with_state_schema(build_chat_client):
"""Test endpoint with state_schema parameter."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -69,7 +70,7 @@ async def test_endpoint_with_state_schema():
assert response.status_code == 200
async def test_endpoint_with_default_state_seed():
async def test_endpoint_with_default_state_seed(build_chat_client):
"""Test endpoint seeds default state when client omits it."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -96,7 +97,7 @@ async def test_endpoint_with_default_state_seed():
assert snapshots[0]["snapshot"]["proverbs"] == default_state["proverbs"]
async def test_endpoint_with_predict_state_config():
async def test_endpoint_with_predict_state_config(build_chat_client):
"""Test endpoint with predict_state_config parameter."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -110,7 +111,7 @@ async def test_endpoint_with_predict_state_config():
assert response.status_code == 200
async def test_endpoint_request_logging():
async def test_endpoint_request_logging(build_chat_client):
"""Test that endpoint logs request details."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -130,7 +131,7 @@ async def test_endpoint_request_logging():
assert response.status_code == 200
async def test_endpoint_event_streaming():
async def test_endpoint_event_streaming(build_chat_client):
"""Test that endpoint streams events correctly."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client("Streamed response"))
@@ -164,7 +165,7 @@ async def test_endpoint_event_streaming():
assert found_run_finished
async def test_endpoint_error_handling():
async def test_endpoint_error_handling(build_chat_client):
"""Test endpoint error handling during request parsing."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -180,7 +181,7 @@ async def test_endpoint_error_handling():
assert response.status_code == 422
async def test_endpoint_multiple_paths():
async def test_endpoint_multiple_paths(build_chat_client):
"""Test adding multiple endpoints with different paths."""
app = FastAPI()
agent1 = ChatAgent(name="agent1", instructions="First agent", chat_client=build_chat_client("Response 1"))
@@ -198,7 +199,7 @@ async def test_endpoint_multiple_paths():
assert response2.status_code == 200
async def test_endpoint_default_path():
async def test_endpoint_default_path(build_chat_client):
"""Test endpoint with default path."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -211,7 +212,7 @@ async def test_endpoint_default_path():
assert response.status_code == 200
async def test_endpoint_response_headers():
async def test_endpoint_response_headers(build_chat_client):
"""Test that endpoint sets correct response headers."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -227,7 +228,7 @@ async def test_endpoint_response_headers():
assert response.headers["cache-control"] == "no-cache"
async def test_endpoint_empty_messages():
async def test_endpoint_empty_messages(build_chat_client):
"""Test endpoint with empty messages list."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -240,7 +241,7 @@ async def test_endpoint_empty_messages():
assert response.status_code == 200
async def test_endpoint_complex_input():
async def test_endpoint_complex_input(build_chat_client):
"""Test endpoint with complex input data."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -265,7 +266,7 @@ async def test_endpoint_complex_input():
assert response.status_code == 200
async def test_endpoint_openapi_schema():
async def test_endpoint_openapi_schema(build_chat_client):
"""Test that endpoint generates proper OpenAPI schema with request model."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -309,7 +310,7 @@ async def test_endpoint_openapi_schema():
assert "messages" in agui_request_schema["required"]
async def test_endpoint_default_tags():
async def test_endpoint_default_tags(build_chat_client):
"""Test that endpoint uses default 'AG-UI' tag."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -327,7 +328,7 @@ async def test_endpoint_default_tags():
assert endpoint_spec["tags"] == ["AG-UI"]
async def test_endpoint_custom_tags():
async def test_endpoint_custom_tags(build_chat_client):
"""Test that endpoint accepts custom tags."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -345,7 +346,7 @@ async def test_endpoint_custom_tags():
assert endpoint_spec["tags"] == ["Custom", "Agent"]
async def test_endpoint_missing_required_field():
async def test_endpoint_missing_required_field(build_chat_client):
"""Test that endpoint validates required fields with Pydantic."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -362,7 +363,7 @@ async def test_endpoint_missing_required_field():
assert "detail" in error_detail
async def test_endpoint_internal_error_handling():
async def test_endpoint_internal_error_handling(build_chat_client):
"""Test endpoint error handling when an exception occurs before streaming starts."""
from unittest.mock import patch
@@ -383,7 +384,7 @@ async def test_endpoint_internal_error_handling():
assert response.json() == {"error": "An internal error has occurred."}
async def test_endpoint_with_dependencies_blocks_unauthorized():
async def test_endpoint_with_dependencies_blocks_unauthorized(build_chat_client):
"""Test that endpoint blocks requests when authentication dependency fails."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -402,7 +403,7 @@ async def test_endpoint_with_dependencies_blocks_unauthorized():
assert response.json()["detail"] == "Unauthorized"
async def test_endpoint_with_dependencies_allows_authorized():
async def test_endpoint_with_dependencies_allows_authorized(build_chat_client):
"""Test that endpoint allows requests when authentication dependency passes."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -425,7 +426,7 @@ async def test_endpoint_with_dependencies_allows_authorized():
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
async def test_endpoint_with_multiple_dependencies():
async def test_endpoint_with_multiple_dependencies(build_chat_client):
"""Test that endpoint supports multiple dependencies."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
@@ -453,7 +454,7 @@ async def test_endpoint_with_multiple_dependencies():
assert "second" in execution_order
async def test_endpoint_without_dependencies_is_accessible():
async def test_endpoint_without_dependencies_is_accessible(build_chat_client):
"""Test that endpoint without dependencies remains accessible (backward compatibility)."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
python/packages/ag-ui/tests/test_event_converters.py → python/packages/ag-ui/tests/ag_ui/test_event_converters.py
View File
python/packages/ag-ui/tests/test_helpers.py → python/packages/ag-ui/tests/ag_ui/test_helpers.py
+5 -5
View File
@@ -29,8 +29,8 @@ class TestPendingToolCallIds:
def test_no_tool_calls(self):
"""Returns empty set when no tool calls in messages."""
messages = [
ChatMessage("user", [Content.from_text("Hello")]),
ChatMessage("assistant", [Content.from_text("Hi there")]),
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text("Hi there")]),
]
result = pending_tool_call_ids(messages)
assert result == set()
@@ -114,7 +114,7 @@ class TestIsStateContextMessage:
def test_empty_contents(self):
"""Returns False for message with empty contents."""
message = ChatMessage("system", [])
message = ChatMessage(role="system", contents=[])
assert is_state_context_message(message) is False
@@ -342,7 +342,7 @@ class TestLatestApprovalResponse:
def test_no_approval_response(self):
"""Returns None when no approval response in last message."""
messages = [
ChatMessage("assistant", [Content.from_text("Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text("Hello")]),
]
result = latest_approval_response(messages)
assert result is None
@@ -357,7 +357,7 @@ class TestLatestApprovalResponse:
function_call=fc,
)
messages = [
ChatMessage("user", [approval_content]),
ChatMessage(role="user", contents=[approval_content]),
]
result = latest_approval_response(messages)
assert result is approval_content
python/packages/ag-ui/tests/test_http_service.py → python/packages/ag-ui/tests/ag_ui/test_http_service.py
View File
python/packages/ag-ui/tests/test_message_adapters.py → python/packages/ag-ui/tests/ag_ui/test_message_adapters.py
+3 -3
View File
@@ -24,7 +24,7 @@ def sample_agui_message():
@pytest.fixture
def sample_agent_framework_message():
"""Create a sample Agent Framework message."""
return ChatMessage("user", [Content.from_text(text="Hello")], message_id="msg-123")
return ChatMessage(role="user", contents=[Content.from_text(text="Hello")], message_id="msg-123")
def test_agui_to_agent_framework_basic(sample_agui_message):
@@ -484,7 +484,7 @@ def test_agent_framework_to_agui_multiple_text_contents():
def test_agent_framework_to_agui_no_message_id():
"""Test message without message_id - should auto-generate ID."""
msg = ChatMessage("user", [Content.from_text(text="Hello")])
msg = ChatMessage(role="user", contents=[Content.from_text(text="Hello")])
messages = agent_framework_messages_to_agui([msg])
@@ -496,7 +496,7 @@ def test_agent_framework_to_agui_no_message_id():
def test_agent_framework_to_agui_system_role():
"""Test system role conversion."""
msg = ChatMessage("system", [Content.from_text(text="System")])
msg = ChatMessage(role="system", contents=[Content.from_text(text="System")])
messages = agent_framework_messages_to_agui([msg])
python/packages/ag-ui/tests/test_message_hygiene.py → python/packages/ag-ui/tests/ag_ui/test_message_hygiene.py
+6 -12
View File
@@ -33,14 +33,12 @@ def test_sanitize_tool_history_filters_out_confirm_changes_only_message() -> Non
# Assistant message with only confirm_changes should be filtered out
assistant_messages = [
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
]
assert len(assistant_messages) == 0
# No synthetic tool result should be injected since confirm_changes was filtered out
tool_messages = [
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
]
tool_messages = [msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "tool"]
assert len(tool_messages) == 0
@@ -182,7 +180,7 @@ def test_sanitize_tool_history_filters_confirm_changes_keeps_other_tools() -> No
# Find the assistant message
assistant_messages = [
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
]
assert len(assistant_messages) == 1
@@ -192,9 +190,7 @@ def test_sanitize_tool_history_filters_confirm_changes_keeps_other_tools() -> No
assert "confirm_changes" not in function_call_names
# Only one tool message (for call_1), no synthetic for confirm_changes
tool_messages = [
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
]
tool_messages = [msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "tool"]
assert len(tool_messages) == 1
assert str(tool_messages[0].contents[0].call_id) == "call_1"
@@ -249,7 +245,7 @@ def test_sanitize_tool_history_filters_confirm_changes_from_assistant_messages()
# Find the assistant message in sanitized output
assistant_messages = [
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
]
assert len(assistant_messages) == 1
@@ -261,9 +257,7 @@ def test_sanitize_tool_history_filters_confirm_changes_from_assistant_messages()
assert "confirm_changes" not in function_call_names
# No synthetic tool result for confirm_changes (it was filtered from the message)
tool_messages = [
msg for msg in sanitized if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
]
tool_messages = [msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "tool"]
# No tool results expected since there are no completed tool calls
# (the approval response is handled separately by the framework)
tool_call_ids = {str(msg.contents[0].call_id) for msg in tool_messages}
python/packages/ag-ui/tests/test_predictive_state.py → python/packages/ag-ui/tests/ag_ui/test_predictive_state.py
View File
python/packages/ag-ui/tests/test_run.py → python/packages/ag-ui/tests/ag_ui/test_run.py
+7 -7
View File
@@ -212,7 +212,7 @@ class TestInjectStateContext:
def test_no_state_message(self):
"""Returns original messages when no state context needed."""
messages = [ChatMessage("user", [Content.from_text("Hello")])]
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
result = _inject_state_context(messages, {}, {})
assert result == messages
@@ -224,8 +224,8 @@ class TestInjectStateContext:
def test_last_message_not_user(self):
"""Returns original messages when last message is not from user."""
messages = [
ChatMessage("user", [Content.from_text("Hello")]),
ChatMessage("assistant", [Content.from_text("Hi")]),
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text("Hi")]),
]
state = {"key": "value"}
schema = {"properties": {"key": {"type": "string"}}}
@@ -237,8 +237,8 @@ class TestInjectStateContext:
"""Injects state context before last user message."""
messages = [
ChatMessage("system", [Content.from_text("You are helpful")]),
ChatMessage("user", [Content.from_text("Hello")]),
ChatMessage(role="system", contents=[Content.from_text("You are helpful")]),
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
]
state = {"document": "content"}
schema = {"properties": {"document": {"type": "string"}}}
@@ -405,7 +405,7 @@ def test_extract_approved_state_updates_no_handler():
"""Test _extract_approved_state_updates returns empty with no handler."""
from agent_framework_ag_ui._run import _extract_approved_state_updates
messages = [ChatMessage("user", [Content.from_text("Hello")])]
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
result = _extract_approved_state_updates(messages, None)
assert result == {}
@@ -416,7 +416,7 @@ def test_extract_approved_state_updates_no_approval():
from agent_framework_ag_ui._run import _extract_approved_state_updates
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "content"}})
messages = [ChatMessage("user", [Content.from_text("Hello")])]
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
result = _extract_approved_state_updates(messages, handler)
assert result == {}
python/packages/ag-ui/tests/test_service_thread_id.py → python/packages/ag-ui/tests/ag_ui/test_service_thread_id.py
+6 -11
View File
@@ -2,19 +2,14 @@
"""Tests for service-managed thread IDs, and service-generated response ids."""
import sys
from pathlib import Path
from typing import Any
from ag_ui.core import RunFinishedEvent, RunStartedEvent
from agent_framework import Content
from agent_framework._types import AgentResponseUpdate, ChatResponseUpdate
sys.path.insert(0, str(Path(__file__).parent))
from utils_test_ag_ui import StubAgent
async def test_service_thread_id_when_there_are_updates():
async def test_service_thread_id_when_there_are_updates(stub_agent):
"""Test that service-managed thread IDs (conversation_id) are correctly set as the thread_id in events."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -29,7 +24,7 @@ async def test_service_thread_id_when_there_are_updates():
),
)
]
agent = StubAgent(updates=updates)
agent = stub_agent(updates=updates)
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {
@@ -46,12 +41,12 @@ async def test_service_thread_id_when_there_are_updates():
assert isinstance(events[-1], RunFinishedEvent)
async def test_service_thread_id_when_no_user_message():
async def test_service_thread_id_when_no_user_message(stub_agent):
"""Test when user submits no messages, emitted events still have with a thread_id"""
from agent_framework.ag_ui import AgentFrameworkAgent
updates: list[AgentResponseUpdate] = []
agent = StubAgent(updates=updates)
agent = stub_agent(updates=updates)
wrapper = AgentFrameworkAgent(agent=agent)
input_data: dict[str, list[dict[str, str]]] = {
@@ -68,12 +63,12 @@ async def test_service_thread_id_when_no_user_message():
assert isinstance(events[-1], RunFinishedEvent)
async def test_service_thread_id_when_user_supplied_thread_id():
async def test_service_thread_id_when_user_supplied_thread_id(stub_agent):
"""Test that user-supplied thread IDs are preserved in emitted events."""
from agent_framework.ag_ui import AgentFrameworkAgent
updates: list[AgentResponseUpdate] = []
agent = StubAgent(updates=updates)
agent = stub_agent(updates=updates)
wrapper = AgentFrameworkAgent(agent=agent)
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}], "threadId": "conv_12345"}
python/packages/ag-ui/tests/test_structured_output.py → python/packages/ag-ui/tests/ag_ui/test_structured_output.py
+14 -19
View File
@@ -3,17 +3,12 @@
"""Tests for structured output handling in _agent.py."""
import json
import sys
from collections.abc import AsyncIterator, MutableSequence
from pathlib import Path
from typing import Any
from agent_framework import ChatAgent, ChatMessage, ChatOptions, ChatResponseUpdate, Content
from pydantic import BaseModel
sys.path.insert(0, str(Path(__file__).parent))
from utils_test_ag_ui import StreamingChatClientStub, stream_from_updates
class RecipeOutput(BaseModel):
"""Test Pydantic model for recipe output."""
@@ -35,7 +30,7 @@ class GenericOutput(BaseModel):
data: dict[str, Any]
async def test_structured_output_with_recipe():
async def test_structured_output_with_recipe(streaming_chat_client_stub, stream_from_updates_fixture):
"""Test structured output processing with recipe state."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -46,7 +41,7 @@ async def test_structured_output_with_recipe():
contents=[Content.from_text(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
)
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=RecipeOutput)
wrapper = AgentFrameworkAgent(
@@ -73,7 +68,7 @@ async def test_structured_output_with_recipe():
assert any("Here is your recipe" in e.delta for e in text_events)
async def test_structured_output_with_steps():
async def test_structured_output_with_steps(streaming_chat_client_stub, stream_from_updates_fixture):
"""Test structured output processing with steps state."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -88,7 +83,7 @@ async def test_structured_output_with_steps():
}
yield ChatResponseUpdate(contents=[Content.from_text(text=json.dumps(steps_data))])
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=StepsOutput)
wrapper = AgentFrameworkAgent(
@@ -113,7 +108,7 @@ async def test_structured_output_with_steps():
assert steps_snapshots[0].snapshot["steps"][0]["id"] == "1"
async def test_structured_output_with_no_schema_match():
async def test_structured_output_with_no_schema_match(streaming_chat_client_stub, stream_from_updates_fixture):
"""Test structured output when response fields don't match state_schema keys."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -122,7 +117,7 @@ async def test_structured_output_with_no_schema_match():
]
agent = ChatAgent(
name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_from_updates(updates))
name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_from_updates_fixture(updates))
)
agent.default_options = ChatOptions(response_format=GenericOutput)
@@ -143,7 +138,7 @@ async def test_structured_output_with_no_schema_match():
assert len(snapshot_events) >= 1
async def test_structured_output_without_schema():
async def test_structured_output_without_schema(streaming_chat_client_stub, stream_from_updates_fixture):
"""Test structured output without state_schema treats all fields as state."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -158,7 +153,7 @@ async def test_structured_output_without_schema():
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text='{"data": {"key": "value"}, "info": "processed"}')])
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=DataOutput)
wrapper = AgentFrameworkAgent(
@@ -181,7 +176,7 @@ async def test_structured_output_without_schema():
assert snapshot_events[0].snapshot["info"] == "processed"
async def test_no_structured_output_when_no_response_format():
async def test_no_structured_output_when_no_response_format(streaming_chat_client_stub, stream_from_updates_fixture):
"""Test that structured output path is skipped when no response_format."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -190,7 +185,7 @@ async def test_no_structured_output_when_no_response_format():
agent = ChatAgent(
name="test",
instructions="Test",
chat_client=StreamingChatClientStub(stream_from_updates(updates)),
chat_client=streaming_chat_client_stub(stream_from_updates_fixture(updates)),
)
# No response_format set
@@ -208,7 +203,7 @@ async def test_no_structured_output_when_no_response_format():
assert text_events[0].delta == "Regular text"
async def test_structured_output_with_message_field():
async def test_structured_output_with_message_field(streaming_chat_client_stub, stream_from_updates_fixture):
"""Test structured output that includes a message field."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -218,7 +213,7 @@ async def test_structured_output_with_message_field():
output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
yield ChatResponseUpdate(contents=[Content.from_text(text=json.dumps(output_data))])
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=RecipeOutput)
wrapper = AgentFrameworkAgent(
@@ -243,7 +238,7 @@ async def test_structured_output_with_message_field():
assert len(end_events) >= 1
async def test_empty_updates_no_structured_processing():
async def test_empty_updates_no_structured_processing(streaming_chat_client_stub, stream_from_updates_fixture):
"""Test that empty updates don't trigger structured output processing."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -253,7 +248,7 @@ async def test_empty_updates_no_structured_processing():
if False:
yield ChatResponseUpdate(contents=[])
agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=RecipeOutput)
wrapper = AgentFrameworkAgent(agent=agent)
python/packages/ag-ui/tests/test_tooling.py → python/packages/ag-ui/tests/ag_ui/test_tooling.py
+3 -3
View File
@@ -54,17 +54,17 @@ def test_merge_tools_filters_duplicates() -> None:
def test_register_additional_client_tools_assigns_when_configured() -> None:
"""register_additional_client_tools should set additional_tools on the chat client."""
from agent_framework import BaseChatClient, FunctionInvocationConfiguration
from agent_framework import BaseChatClient, normalize_function_invocation_configuration
mock_chat_client = MagicMock(spec=BaseChatClient)
mock_chat_client.function_invocation_configuration = FunctionInvocationConfiguration()
mock_chat_client.function_invocation_configuration = normalize_function_invocation_configuration(None)
agent = ChatAgent(chat_client=mock_chat_client)
tools = [DummyTool("x")]
register_additional_client_tools(agent, tools)
assert mock_chat_client.function_invocation_configuration.additional_tools == tools
assert mock_chat_client.function_invocation_configuration["additional_tools"] == tools
def test_collect_server_tools_includes_mcp_tools_when_connected() -> None:
python/packages/ag-ui/tests/test_types.py → python/packages/ag-ui/tests/ag_ui/test_types.py
View File
python/packages/ag-ui/tests/test_utils.py → python/packages/ag-ui/tests/ag_ui/test_utils.py
+1 -1
View File
@@ -408,7 +408,7 @@ def test_get_role_value_with_enum():
from agent_framework_ag_ui._utils import get_role_value
message = ChatMessage("user", [Content.from_text("test")])
message = ChatMessage(role="user", contents=[Content.from_text("test")])
result = get_role_value(message)
assert result == "user"
python/packages/ag-ui/tests/utils_test_ag_ui.py
-124
View File
@@ -1,124 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Shared test stubs for AG-UI tests."""
import sys
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, MutableSequence
from types import SimpleNamespace
from typing import Any, Generic
from agent_framework import (
AgentProtocol,
AgentResponse,
AgentResponseUpdate,
AgentThread,
BaseChatClient,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Content,
)
from agent_framework._clients import TOptions_co
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[..., AsyncIterator[ChatResponseUpdate]]
ResponseFn = Callable[..., Awaitable[ChatResponse]]
class StreamingChatClientStub(BaseChatClient[TOptions_co], Generic[TOptions_co]):
"""Typed streaming stub that satisfies ChatClientProtocol."""
def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
super().__init__()
self._stream_fn = stream_fn
self._response_fn = response_fn
@override
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
async for update in self._stream_fn(messages, options, **kwargs):
yield update
@override
async def _inner_get_response(
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> ChatResponse:
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(messages, options, **kwargs):
contents.extend(update.contents)
return ChatResponse(
messages=[ChatMessage("assistant", 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[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
for update in updates:
yield update
return _stream
class StubAgent(AgentProtocol):
"""Minimal AgentProtocol 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,
chat_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.chat_client = chat_client or SimpleNamespace(function_invocation_configuration=None)
self.messages_received: list[Any] = []
self.tools_received: list[Any] | None = None
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
return AgentResponse(messages=[], response_id="stub-response")
def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
async def _stream() -> AsyncIterator[AgentResponseUpdate]:
self.messages_received = [] if messages is None else list(messages) # type: ignore[arg-type]
self.tools_received = kwargs.get("tools")
for update in self.updates:
yield update
return _stream()
def get_new_thread(self, **kwargs: Any) -> AgentThread:
return AgentThread()
python/packages/anthropic/agent_framework_anthropic/_chat_client.py
+57 -42
View File
@@ -1,32 +1,37 @@
# Copyright (c) Microsoft. All rights reserved.
import sys
from collections.abc import AsyncIterable, MutableMapping, MutableSequence, Sequence
from typing import Any, ClassVar, Final, Generic, Literal
from collections.abc import AsyncIterable, Awaitable, Mapping, MutableMapping, Sequence
from typing import Any, ClassVar, Final, Generic, Literal, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
Annotation,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
ChatMessage,
ChatMiddlewareLayer,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
FinishReasonLiteral,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
HostedCodeInterpreterTool,
HostedMCPTool,
HostedWebSearchTool,
ResponseStream,
TextSpanRegion,
UsageDetails,
get_logger,
prepare_function_call_results,
use_chat_middleware,
use_function_invocation,
)
from agent_framework._pydantic import AFBaseSettings
from agent_framework._types import _get_data_bytes_as_str # type: ignore
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.observability import use_instrumentation
from agent_framework.observability import ChatTelemetryLayer
from anthropic import AsyncAnthropic
from anthropic.types.beta import (
BetaContentBlock,
@@ -58,6 +63,7 @@ if sys.version_info >= (3, 12):
else:
from typing_extensions import override # type: ignore # pragma: no cover
__all__ = [
"AnthropicChatOptions",
"AnthropicClient",
@@ -177,7 +183,7 @@ ROLE_MAP: dict[str, str] = {
"tool": "user",
}
FINISH_REASON_MAP: dict[str, str] = {
FINISH_REASON_MAP: dict[str, FinishReasonLiteral] = {
"stop_sequence": "stop",
"max_tokens": "length",
"tool_use": "tool_calls",
@@ -223,11 +229,14 @@ class AnthropicSettings(AFBaseSettings):
chat_model_id: str | None = None
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptions]):
"""Anthropic Chat client."""
class AnthropicClient(
ChatMiddlewareLayer[TAnthropicOptions],
FunctionInvocationLayer[TAnthropicOptions],
ChatTelemetryLayer[TAnthropicOptions],
BaseChatClient[TAnthropicOptions],
Generic[TAnthropicOptions],
):
"""Anthropic Chat client with middleware, telemetry, and function invocation support."""
OTEL_PROVIDER_NAME: ClassVar[str] = "anthropic" # type: ignore[reportIncompatibleVariableOverride, misc]
@@ -238,6 +247,8 @@ class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptio
model_id: str | None = None,
anthropic_client: AsyncAnthropic | None = None,
additional_beta_flags: list[str] | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
@@ -252,6 +263,8 @@ class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptio
For instance if you need to set a different base_url for testing or private deployments.
additional_beta_flags: Additional beta flags to enable on the client.
Default flags are: "mcp-client-2025-04-04", "code-execution-2025-08-25".
middleware: Optional middleware to apply to the client.
function_invocation_configuration: Optional function invocation configuration override.
env_file_path: Path to environment file for loading settings.
env_file_encoding: Encoding of the environment file.
kwargs: Additional keyword arguments passed to the parent class.
@@ -322,7 +335,11 @@ class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptio
)
# Initialize parent
super().__init__(**kwargs)
super().__init__(
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
# Initialize instance variables
self.anthropic_client = anthropic_client
@@ -334,42 +351,40 @@ class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptio
# region Get response methods
@override
async def _inner_get_response(
def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
) -> ChatResponse:
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
# prepare
run_options = self._prepare_options(messages, options, **kwargs)
# execute
message = await self.anthropic_client.beta.messages.create(**run_options, stream=False)
# process
return self._process_message(message, options)
@override
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
# prepare
run_options = self._prepare_options(messages, options, **kwargs)
# execute and process
async for chunk in await self.anthropic_client.beta.messages.create(**run_options, stream=True):
parsed_chunk = self._process_stream_event(chunk)
if parsed_chunk:
yield parsed_chunk
if stream:
# Streaming mode
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
async for chunk in await self.anthropic_client.beta.messages.create(**run_options, stream=True):
parsed_chunk = self._process_stream_event(chunk)
if parsed_chunk:
yield parsed_chunk
return self._build_response_stream(_stream(), response_format=options.get("response_format"))
# Non-streaming mode
async def _get_response() -> ChatResponse:
message = await self.anthropic_client.beta.messages.create(**run_options, stream=False)
return self._process_message(message, options)
return _get_response()
# region Prep methods
def _prepare_options(
self,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
**kwargs: Any,
) -> dict[str, Any]:
"""Create run options for the Anthropic client based on messages and options.
@@ -443,7 +458,7 @@ class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptio
run_options.update(kwargs)
return run_options
def _prepare_betas(self, options: dict[str, Any]) -> set[str]:
def _prepare_betas(self, options: Mapping[str, Any]) -> set[str]:
"""Prepare the beta flags for the Anthropic API request.
Args:
@@ -493,7 +508,7 @@ class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptio
"schema": schema,
}
def _prepare_messages_for_anthropic(self, messages: MutableSequence[ChatMessage]) -> list[dict[str, Any]]:
def _prepare_messages_for_anthropic(self, messages: Sequence[ChatMessage]) -> list[dict[str, Any]]:
"""Prepare a list of ChatMessages for the Anthropic client.
This skips the first message if it is a system message,
@@ -525,7 +540,7 @@ class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptio
a_content.append({
"type": "image",
"source": {
"data": content.get_data_bytes_as_str(), # type: ignore[attr-defined]
"data": _get_data_bytes_as_str(content), # type: ignore[attr-defined]
"media_type": content.media_type,
"type": "base64",
},
@@ -564,7 +579,7 @@ class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptio
"content": a_content,
}
def _prepare_tools_for_anthropic(self, options: dict[str, Any]) -> dict[str, Any] | None:
def _prepare_tools_for_anthropic(self, options: Mapping[str, Any]) -> dict[str, Any] | None:
"""Prepare tools and tool choice configuration for the Anthropic API request.
Args:
@@ -657,7 +672,7 @@ class AnthropicClient(BaseChatClient[TAnthropicOptions], Generic[TAnthropicOptio
# region Response Processing Methods
def _process_message(self, message: BetaMessage, options: dict[str, Any]) -> ChatResponse:
def _process_message(self, message: BetaMessage, options: Mapping[str, Any]) -> ChatResponse:
"""Process the response from the Anthropic client.
Args:
python/packages/anthropic/tests/test_anthropic_client.py
+33 -33
View File
@@ -148,7 +148,7 @@ def test_anthropic_client_service_url(mock_anthropic_client: MagicMock) -> None:
def test_prepare_message_for_anthropic_text(mock_anthropic_client: MagicMock) -> None:
"""Test converting text message to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
message = ChatMessage("user", ["Hello, world!"])
message = ChatMessage(role="user", text="Hello, world!")
result = chat_client._prepare_message_for_anthropic(message)
@@ -227,8 +227,8 @@ def test_prepare_messages_for_anthropic_with_system(mock_anthropic_client: Magic
"""Test converting messages list with system message."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [
ChatMessage("system", ["You are a helpful assistant."]),
ChatMessage("user", ["Hello!"]),
ChatMessage(role="system", text="You are a helpful assistant."),
ChatMessage(role="user", text="Hello!"),
]
result = chat_client._prepare_messages_for_anthropic(messages)
@@ -243,8 +243,8 @@ def test_prepare_messages_for_anthropic_without_system(mock_anthropic_client: Ma
"""Test converting messages list without system message."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [
ChatMessage("user", ["Hello!"]),
ChatMessage("assistant", ["Hi there!"]),
ChatMessage(role="user", text="Hello!"),
ChatMessage(role="assistant", text="Hi there!"),
]
result = chat_client._prepare_messages_for_anthropic(messages)
@@ -372,7 +372,7 @@ async def test_prepare_options_basic(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with basic ChatOptions."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options = ChatOptions(max_tokens=100, temperature=0.7)
run_options = chat_client._prepare_options(messages, chat_options)
@@ -388,8 +388,8 @@ async def test_prepare_options_with_system_message(mock_anthropic_client: MagicM
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [
ChatMessage("system", ["You are helpful."]),
ChatMessage("user", ["Hello"]),
ChatMessage(role="system", text="You are helpful."),
ChatMessage(role="user", text="Hello"),
]
chat_options = ChatOptions()
@@ -403,7 +403,7 @@ async def test_prepare_options_with_tool_choice_auto(mock_anthropic_client: Magi
"""Test _prepare_options with auto tool choice."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options = ChatOptions(tool_choice="auto")
run_options = chat_client._prepare_options(messages, chat_options)
@@ -415,7 +415,7 @@ async def test_prepare_options_with_tool_choice_required(mock_anthropic_client:
"""Test _prepare_options with required tool choice."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
# For required with specific function, need to pass as dict
chat_options = ChatOptions(tool_choice={"mode": "required", "required_function_name": "get_weather"})
@@ -429,7 +429,7 @@ async def test_prepare_options_with_tool_choice_none(mock_anthropic_client: Magi
"""Test _prepare_options with none tool choice."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options = ChatOptions(tool_choice="none")
run_options = chat_client._prepare_options(messages, chat_options)
@@ -446,7 +446,7 @@ async def test_prepare_options_with_tools(mock_anthropic_client: MagicMock) -> N
"""Get weather for a location."""
return f"Weather for {location}"
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options = ChatOptions(tools=[get_weather])
run_options = chat_client._prepare_options(messages, chat_options)
@@ -459,7 +459,7 @@ async def test_prepare_options_with_stop_sequences(mock_anthropic_client: MagicM
"""Test _prepare_options with stop sequences."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options = ChatOptions(stop=["STOP", "END"])
run_options = chat_client._prepare_options(messages, chat_options)
@@ -471,7 +471,7 @@ async def test_prepare_options_with_top_p(mock_anthropic_client: MagicMock) -> N
"""Test _prepare_options with top_p."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options = ChatOptions(top_p=0.9)
run_options = chat_client._prepare_options(messages, chat_options)
@@ -666,7 +666,7 @@ async def test_inner_get_response(mock_anthropic_client: MagicMock) -> None:
mock_anthropic_client.beta.messages.create.return_value = mock_message
messages = [ChatMessage("user", ["Hi"])]
messages = [ChatMessage(role="user", text="Hi")]
chat_options = ChatOptions(max_tokens=10)
response = await chat_client._inner_get_response( # type: ignore[attr-defined]
@@ -678,8 +678,8 @@ async def test_inner_get_response(mock_anthropic_client: MagicMock) -> None:
assert len(response.messages) == 1
async def test_inner_get_streaming_response(mock_anthropic_client: MagicMock) -> None:
"""Test _inner_get_streaming_response method."""
async def test_inner_get_response_streaming(mock_anthropic_client: MagicMock) -> None:
"""Test _inner_get_response method with streaming."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
# Create mock streaming response
@@ -690,12 +690,12 @@ async def test_inner_get_streaming_response(mock_anthropic_client: MagicMock) ->
mock_anthropic_client.beta.messages.create.return_value = mock_stream()
messages = [ChatMessage("user", ["Hi"])]
messages = [ChatMessage(role="user", text="Hi")]
chat_options = ChatOptions(max_tokens=10)
chunks: list[ChatResponseUpdate] = []
async for chunk in chat_client._inner_get_streaming_response( # type: ignore[attr-defined]
messages=messages, options=chat_options
async for chunk in chat_client._inner_get_response( # type: ignore[attr-defined]
messages=messages, options=chat_options, stream=True
):
if chunk:
chunks.append(chunk)
@@ -721,7 +721,7 @@ async def test_anthropic_client_integration_basic_chat() -> None:
"""Integration test for basic chat completion."""
client = AnthropicClient()
messages = [ChatMessage("user", ["Say 'Hello, World!' and nothing else."])]
messages = [ChatMessage(role="user", text="Say 'Hello, World!' and nothing else.")]
response = await client.get_response(messages=messages, options={"max_tokens": 50})
@@ -738,10 +738,10 @@ async def test_anthropic_client_integration_streaming_chat() -> None:
"""Integration test for streaming chat completion."""
client = AnthropicClient()
messages = [ChatMessage("user", ["Count from 1 to 5."])]
messages = [ChatMessage(role="user", text="Count from 1 to 5.")]
chunks = []
async for chunk in client.get_streaming_response(messages=messages, options={"max_tokens": 50}):
async for chunk in client.get_response(messages=messages, stream=True, options={"max_tokens": 50}):
chunks.append(chunk)
assert len(chunks) > 0
@@ -754,7 +754,7 @@ async def test_anthropic_client_integration_function_calling() -> None:
"""Integration test for function calling."""
client = AnthropicClient()
messages = [ChatMessage("user", ["What's the weather in San Francisco?"])]
messages = [ChatMessage(role="user", text="What's the weather in San Francisco?")]
tools = [get_weather]
response = await client.get_response(
@@ -774,7 +774,7 @@ async def test_anthropic_client_integration_hosted_tools() -> None:
"""Integration test for hosted tools."""
client = AnthropicClient()
messages = [ChatMessage("user", ["What tools do you have available?"])]
messages = [ChatMessage(role="user", text="What tools do you have available?")]
tools = [
HostedWebSearchTool(),
HostedCodeInterpreterTool(),
@@ -801,8 +801,8 @@ async def test_anthropic_client_integration_with_system_message() -> None:
client = AnthropicClient()
messages = [
ChatMessage("system", ["You are a pirate. Always respond like a pirate."]),
ChatMessage("user", ["Hello!"]),
ChatMessage(role="system", text="You are a pirate. Always respond like a pirate."),
ChatMessage(role="user", text="Hello!"),
]
response = await client.get_response(messages=messages, options={"max_tokens": 50})
@@ -817,7 +817,7 @@ async def test_anthropic_client_integration_temperature_control() -> None:
"""Integration test with temperature control."""
client = AnthropicClient()
messages = [ChatMessage("user", ["Say hello."])]
messages = [ChatMessage(role="user", text="Say hello.")]
response = await client.get_response(
messages=messages,
@@ -835,11 +835,11 @@ async def test_anthropic_client_integration_ordering() -> None:
client = AnthropicClient()
messages = [
ChatMessage("user", ["Say hello."]),
ChatMessage("user", ["Then say goodbye."]),
ChatMessage("assistant", ["Thank you for chatting!"]),
ChatMessage("assistant", ["Let me know if I can help."]),
ChatMessage("user", ["Just testing things."]),
ChatMessage(role="user", text="Say hello."),
ChatMessage(role="user", text="Then say goodbye."),
ChatMessage(role="assistant", text="Thank you for chatting!"),
ChatMessage(role="assistant", text="Let me know if I can help."),
ChatMessage(role="user", text="Just testing things."),
]
response = await client.get_response(messages=messages)
python/packages/azure-ai-search/agent_framework_azure_ai_search/_search_provider.py
+8 -3
View File
@@ -524,8 +524,13 @@ class AzureAISearchContextProvider(ContextProvider):
# Convert to list and filter to USER/ASSISTANT messages with text only
messages_list = [messages] if isinstance(messages, ChatMessage) else list(messages)
def get_role_value(role: str | Any) -> str:
return role.value if hasattr(role, "value") else str(role)
filtered_messages = [
msg for msg in messages_list if msg and msg.text and msg.text.strip() and msg.role in ["user", "assistant"]
msg
for msg in messages_list
if msg and msg.text and msg.text.strip() and get_role_value(msg.role) in ["user", "assistant"]
]
if not filtered_messages:
@@ -546,8 +551,8 @@ class AzureAISearchContextProvider(ContextProvider):
return Context()
# Create context messages: first message with prompt, then one message per result part
context_messages = [ChatMessage("user", [self.context_prompt])]
context_messages.extend([ChatMessage("user", [part]) for part in search_result_parts])
context_messages = [ChatMessage(role="user", text=self.context_prompt)]
context_messages.extend([ChatMessage(role="user", text=part) for part in search_result_parts])
return Context(messages=context_messages)
python/packages/azure-ai-search/tests/test_search_provider.py
+10 -10
View File
@@ -39,7 +39,7 @@ def mock_index_client() -> AsyncMock:
def sample_messages() -> list[ChatMessage]:
"""Create sample chat messages for testing."""
return [
ChatMessage("user", ["What is in the documents?"]),
ChatMessage(role="user", text="What is in the documents?"),
]
@@ -318,7 +318,7 @@ class TestSemanticSearch:
)
# Empty message
context = await provider.invoking([ChatMessage("user", [""])])
context = await provider.invoking([ChatMessage(role="user", text="")])
assert isinstance(context, Context)
assert len(context.messages) == 0
@@ -520,10 +520,10 @@ class TestMessageFiltering:
# Mix of message types
messages = [
ChatMessage("system", ["System message"]),
ChatMessage("user", ["User message"]),
ChatMessage("assistant", ["Assistant message"]),
ChatMessage("tool", ["Tool message"]),
ChatMessage(role="system", text="System message"),
ChatMessage(role="user", text="User message"),
ChatMessage(role="assistant", text="Assistant message"),
ChatMessage(role="tool", text="Tool message"),
]
context = await provider.invoking(messages)
@@ -548,9 +548,9 @@ class TestMessageFiltering:
# Messages with empty/whitespace text
messages = [
ChatMessage("user", [""]),
ChatMessage("user", [" "]),
ChatMessage("user", [None]),
ChatMessage(role="user", text=""),
ChatMessage(role="user", text=" "),
ChatMessage(role="user", text=""), # ChatMessage with None text becomes empty string
]
context = await provider.invoking(messages)
@@ -581,7 +581,7 @@ class TestCitations:
mode="semantic",
)
context = await provider.invoking([ChatMessage("user", ["test query"])])
context = await provider.invoking([ChatMessage(role="user", text="test query")])
# Check that citation is included
assert isinstance(context, Context)
python/packages/azure-ai/agent_framework_azure_ai/__init__.py
+2 -1
View File
@@ -4,7 +4,7 @@ import importlib.metadata
from ._agent_provider import AzureAIAgentsProvider
from ._chat_client import AzureAIAgentClient, AzureAIAgentOptions
from ._client import AzureAIClient, AzureAIProjectAgentOptions
from ._client import AzureAIClient, AzureAIProjectAgentOptions, RawAzureAIClient
from ._project_provider import AzureAIProjectAgentProvider
from ._shared import AzureAISettings
@@ -21,5 +21,6 @@ __all__ = [
"AzureAIProjectAgentOptions",
"AzureAIProjectAgentProvider",
"AzureAISettings",
"RawAzureAIClient",
"__version__",
]
python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py
+5 -5
View File
@@ -9,7 +9,7 @@ from agent_framework import (
ChatAgent,
ContextProvider,
FunctionTool,
Middleware,
MiddlewareTypes,
ToolProtocol,
normalize_tools,
)
@@ -175,7 +175,7 @@ class AzureAIAgentsProvider(Generic[TOptions_co]):
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Create a new agent on the Azure AI service and return a ChatAgent.
@@ -272,7 +272,7 @@ class AzureAIAgentsProvider(Generic[TOptions_co]):
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Retrieve an existing agent from the service and return a ChatAgent.
@@ -328,7 +328,7 @@ class AzureAIAgentsProvider(Generic[TOptions_co]):
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Wrap an existing Agent SDK object as a ChatAgent without making HTTP calls.
@@ -381,7 +381,7 @@ class AzureAIAgentsProvider(Generic[TOptions_co]):
agent: Agent,
provided_tools: Sequence[ToolProtocol | MutableMapping[str, Any]] | None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Create a ChatAgent from an Agent SDK object.
python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
+71 -43
View File
@@ -5,37 +5,41 @@ import json
import os
import re
import sys
from collections.abc import AsyncIterable, Callable, Mapping, MutableMapping, MutableSequence, Sequence
from typing import Any, ClassVar, Generic
from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence
from typing import Any, ClassVar, Generic, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
Annotation,
BaseChatClient,
ChatAgent,
ChatAndFunctionMiddlewareTypes,
ChatMessage,
ChatMessageStoreProtocol,
ChatMiddlewareLayer,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
ContextProvider,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
HostedCodeInterpreterTool,
HostedFileSearchTool,
HostedMCPTool,
HostedWebSearchTool,
Middleware,
MiddlewareTypes,
ResponseStream,
Role,
TextSpanRegion,
ToolProtocol,
UsageDetails,
get_logger,
prepare_function_call_results,
use_chat_middleware,
use_function_invocation,
)
from agent_framework.exceptions import ServiceInitializationError, ServiceInvalidRequestError, ServiceResponseException
from agent_framework.observability import use_instrumentation
from agent_framework.observability import ChatTelemetryLayer
from azure.ai.agents.aio import AgentsClient
from azure.ai.agents.models import (
Agent,
@@ -198,11 +202,14 @@ TAzureAIAgentOptions = TypeVar(
# endregion
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIAgentOptions]):
"""Azure AI Agent Chat client."""
class AzureAIAgentClient(
ChatMiddlewareLayer[TAzureAIAgentOptions],
FunctionInvocationLayer[TAzureAIAgentOptions],
ChatTelemetryLayer[TAzureAIAgentOptions],
BaseChatClient[TAzureAIAgentOptions],
Generic[TAzureAIAgentOptions],
):
"""Azure AI Agent Chat client with middleware, telemetry, and function invocation support."""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai" # type: ignore[reportIncompatibleVariableOverride, misc]
@@ -218,6 +225,8 @@ class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIA
model_deployment_name: str | None = None,
credential: AsyncTokenCredential | None = None,
should_cleanup_agent: bool = True,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
@@ -242,6 +251,8 @@ class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIA
should_cleanup_agent: Whether to cleanup (delete) agents created by this client when
the client is closed or context is exited. Defaults to True. Only affects agents
created by this client instance; existing agents passed via agent_id are never deleted.
middleware: Optional sequence of middlewares to include.
function_invocation_configuration: Optional function invocation configuration.
env_file_path: Path to environment file for loading settings.
env_file_encoding: Encoding of the environment file.
kwargs: Additional keyword arguments passed to the parent class.
@@ -316,7 +327,11 @@ class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIA
should_close_client = True
# Initialize parent
super().__init__(**kwargs)
super().__init__(
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
# Initialize instance variables
self.agents_client = agents_client
@@ -345,35 +360,48 @@ class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIA
await self._close_client_if_needed()
@override
async def _inner_get_response(
def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
**kwargs: Any,
) -> ChatResponse:
return await ChatResponse.from_update_generator(
updates=self._inner_get_streaming_response(messages=messages, options=options, **kwargs),
output_format_type=options.get("response_format"),
)
@override
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
# prepare
run_options, required_action_results = await self._prepare_options(messages, options, **kwargs)
agent_id = await self._get_agent_id_or_create(run_options)
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
if stream:
# Streaming mode - return the async generator directly
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
# prepare
run_options, required_action_results = await self._prepare_options(messages, options, **kwargs)
agent_id = await self._get_agent_id_or_create(run_options)
# execute and process
async for update in self._process_stream(
*(await self._create_agent_stream(agent_id, run_options, required_action_results))
):
yield update
# execute and process
async for update in self._process_stream(
*(await self._create_agent_stream(agent_id, run_options, required_action_results))
):
yield update
return self._build_response_stream(_stream(), response_format=options.get("response_format"))
# Non-streaming mode - collect updates and convert to response
async def _get_response() -> ChatResponse:
async def _get_streaming() -> AsyncIterable[ChatResponseUpdate]:
# prepare
run_options, required_action_results = await self._prepare_options(messages, options, **kwargs)
agent_id = await self._get_agent_id_or_create(run_options)
# execute and process
async for update in self._process_stream(
*(await self._create_agent_stream(agent_id, run_options, required_action_results))
):
yield update
return await ChatResponse.from_update_generator(
updates=_get_streaming(),
output_format_type=options.get("response_format"),
)
return _get_response()
async def _get_agent_id_or_create(self, run_options: dict[str, Any] | None = None) -> str:
"""Determine which agent to use and create if needed.
@@ -637,7 +665,7 @@ class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIA
match event_data:
case MessageDeltaChunk():
# only one event_type: AgentStreamEvent.THREAD_MESSAGE_DELTA
role = "user" if event_data.delta.role == "user" else "assistant"
role: Role = "user" if event_data.delta.role == "user" else "assistant" # type: ignore[assignment]
# Extract URL citations from the delta chunk
url_citations = self._extract_url_citations(event_data, azure_search_tool_calls)
@@ -876,7 +904,7 @@ class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIA
async def _prepare_options(
self,
messages: MutableSequence[ChatMessage],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
**kwargs: Any,
) -> tuple[dict[str, Any], list[Content] | None]:
@@ -1004,10 +1032,10 @@ class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIA
if agent_definition.tool_resources:
run_options["tool_resources"] = agent_definition.tool_resources
# Add run tools if tool_choice allows
tool_choice = options.get("tool_choice")
# Add run tools - always include tools if provided, regardless of tool_choice
# tool_choice="none" means the model won't call tools, but tools should still be available
tools = options.get("tools")
if tool_choice is not None and tool_choice != "none" and tools:
if tools:
tool_definitions.extend(to_azure_ai_agent_tools(tools, run_options))
# Handle MCP tool resources
@@ -1056,7 +1084,7 @@ class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIA
return mcp_resources
def _prepare_messages(
self, messages: MutableSequence[ChatMessage]
self, messages: Sequence[ChatMessage]
) -> tuple[
list[ThreadMessageOptions] | None,
list[str],
@@ -1271,7 +1299,7 @@ class AzureAIAgentClient(BaseChatClient[TAzureAIAgentOptions], Generic[TAzureAIA
default_options: TAzureAIAgentOptions | Mapping[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> ChatAgent[TAzureAIAgentOptions]:
"""Convert this chat client to a ChatAgent.
python/packages/azure-ai/agent_framework_azure_ai/_client.py
+143 -20
View File
@@ -1,26 +1,28 @@
# Copyright (c) Microsoft. All rights reserved.
import sys
from collections.abc import Callable, Mapping, MutableMapping, MutableSequence, Sequence
from typing import Any, ClassVar, Generic, TypeVar, cast
from collections.abc import Callable, Mapping, MutableMapping, Sequence
from typing import Any, ClassVar, Generic, TypedDict, TypeVar, cast
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatAgent,
ChatAndFunctionMiddlewareTypes,
ChatMessage,
ChatMessageStoreProtocol,
ChatMiddlewareLayer,
ContextProvider,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
HostedMCPTool,
Middleware,
MiddlewareTypes,
ToolProtocol,
get_logger,
use_chat_middleware,
use_function_invocation,
)
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.observability import use_instrumentation
from agent_framework.observability import ChatTelemetryLayer
from agent_framework.openai import OpenAIResponsesOptions
from agent_framework.openai._responses_client import OpenAIBaseResponsesClient
from agent_framework.openai._responses_client import RawOpenAIResponsesClient
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import MCPTool, PromptAgentDefinition, PromptAgentDefinitionText, RaiConfig, Reasoning
from azure.core.credentials_async import AsyncTokenCredential
@@ -64,11 +66,21 @@ TAzureAIClientOptions = TypeVar(
)
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class AzureAIClient(OpenAIBaseResponsesClient[TAzureAIClientOptions], Generic[TAzureAIClientOptions]):
"""Azure AI Agent client."""
class RawAzureAIClient(RawOpenAIResponsesClient[TAzureAIClientOptions], Generic[TAzureAIClientOptions]):
"""Raw Azure AI client without middleware, telemetry, or function invocation layers.
Warning:
**This class should not normally be used directly.** It does not include middleware,
telemetry, or function invocation support that you most likely need. If you do use it,
you should consider which additional layers to apply. There is a defined ordering that
you should follow:
1. **ChatMiddlewareLayer** - Should be applied first as it also prepares function middleware
2. **FunctionInvocationLayer** - Handles tool/function calling loop
3. **ChatTelemetryLayer** - Must be inside the function calling loop for correct per-call telemetry
Use ``AzureAIClient`` instead for a fully-featured client with all layers applied.
"""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai" # type: ignore[reportIncompatibleVariableOverride, misc]
@@ -88,7 +100,10 @@ class AzureAIClient(OpenAIBaseResponsesClient[TAzureAIClientOptions], Generic[TA
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize an Azure AI Agent client.
"""Initialize a bare Azure AI client.
This is the core implementation without middleware, telemetry, or function invocation layers.
For most use cases, prefer :class:`AzureAIClient` which includes all standard layers.
Keyword Args:
project_client: An existing AIProjectClient to use. If not provided, one will be created.
@@ -379,8 +394,8 @@ class AzureAIClient(OpenAIBaseResponsesClient[TAzureAIClientOptions], Generic[TA
@override
async def _prepare_options(
self,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
**kwargs: Any,
) -> dict[str, Any]:
"""Take ChatOptions and create the specific options for Azure AI."""
@@ -468,13 +483,11 @@ class AzureAIClient(OpenAIBaseResponsesClient[TAzureAIClientOptions], Generic[TA
return transformed
@override
def _get_current_conversation_id(self, options: dict[str, Any], **kwargs: Any) -> str | None:
def _get_current_conversation_id(self, options: Mapping[str, Any], **kwargs: Any) -> str | None:
"""Get the current conversation ID from chat options or kwargs."""
return options.get("conversation_id") or kwargs.get("conversation_id") or self.conversation_id
def _prepare_messages_for_azure_ai(
self, messages: MutableSequence[ChatMessage]
) -> tuple[list[ChatMessage], str | None]:
def _prepare_messages_for_azure_ai(self, messages: Sequence[ChatMessage]) -> tuple[list[ChatMessage], str | None]:
"""Prepare input from messages and convert system/developer messages to instructions."""
result: list[ChatMessage] = []
instructions_list: list[str] = []
@@ -558,7 +571,7 @@ class AzureAIClient(OpenAIBaseResponsesClient[TAzureAIClientOptions], Generic[TA
default_options: TAzureAIClientOptions | Mapping[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> ChatAgent[TAzureAIClientOptions]:
"""Convert this chat client to a ChatAgent.
@@ -597,3 +610,113 @@ class AzureAIClient(OpenAIBaseResponsesClient[TAzureAIClientOptions], Generic[TA
middleware=middleware,
**kwargs,
)
class AzureAIClient(
ChatMiddlewareLayer[TAzureAIClientOptions],
FunctionInvocationLayer[TAzureAIClientOptions],
ChatTelemetryLayer[TAzureAIClientOptions],
RawAzureAIClient[TAzureAIClientOptions],
Generic[TAzureAIClientOptions],
):
"""Azure AI client with middleware, telemetry, and function invocation support.
This is the recommended client for most use cases. It includes:
- Chat middleware support for request/response interception
- OpenTelemetry-based telemetry for observability
- Automatic function/tool invocation handling
For a minimal implementation without these features, use :class:`RawAzureAIClient`.
"""
def __init__(
self,
*,
project_client: AIProjectClient | None = None,
agent_name: str | None = None,
agent_version: str | None = None,
agent_description: str | None = None,
conversation_id: str | None = None,
project_endpoint: str | None = None,
model_deployment_name: str | None = None,
credential: AsyncTokenCredential | None = None,
use_latest_version: bool | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize an Azure AI client with full layer support.
Keyword Args:
project_client: An existing AIProjectClient to use. If not provided, one will be created.
agent_name: The name to use when creating new agents or using existing agents.
agent_version: The version of the agent to use.
agent_description: The description to use when creating new agents.
conversation_id: Default conversation ID to use for conversations. Can be overridden by
conversation_id property when making a request.
project_endpoint: The Azure AI Project endpoint URL.
Can also be set via environment variable AZURE_AI_PROJECT_ENDPOINT.
Ignored when a project_client is passed.
model_deployment_name: The model deployment name to use for agent creation.
Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
credential: Azure async credential to use for authentication.
use_latest_version: Boolean flag that indicates whether to use latest agent version
if it exists in the service.
middleware: Optional sequence of chat middlewares to include.
function_invocation_configuration: Optional function invocation configuration.
env_file_path: Path to environment file for loading settings.
env_file_encoding: Encoding of the environment file.
kwargs: Additional keyword arguments passed to the parent class.
Examples:
.. code-block:: python
from agent_framework_azure_ai import AzureAIClient
from azure.identity.aio import DefaultAzureCredential
# Using environment variables
# Set AZURE_AI_PROJECT_ENDPOINT=https://your-project.cognitiveservices.azure.com
# Set AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4
credential = DefaultAzureCredential()
client = AzureAIClient(credential=credential)
# Or passing parameters directly
client = AzureAIClient(
project_endpoint="https://your-project.cognitiveservices.azure.com",
model_deployment_name="gpt-4",
credential=credential,
)
# Or loading from a .env file
client = AzureAIClient(credential=credential, env_file_path="path/to/.env")
# Using custom ChatOptions with type safety:
from typing import TypedDict
from agent_framework import ChatOptions
class MyOptions(ChatOptions, total=False):
my_custom_option: str
client: AzureAIClient[MyOptions] = AzureAIClient(credential=credential)
response = await client.get_response("Hello", options={"my_custom_option": "value"})
"""
super().__init__(
project_client=project_client,
agent_name=agent_name,
agent_version=agent_version,
agent_description=agent_description,
conversation_id=conversation_id,
project_endpoint=project_endpoint,
model_deployment_name=model_deployment_name,
credential=credential,
use_latest_version=use_latest_version,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
**kwargs,
)
python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
+5 -5
View File
@@ -9,7 +9,7 @@ from agent_framework import (
ChatAgent,
ContextProvider,
FunctionTool,
Middleware,
MiddlewareTypes,
ToolProtocol,
get_logger,
normalize_tools,
@@ -166,7 +166,7 @@ class AzureAIProjectAgentProvider(Generic[TOptions_co]):
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Create a new agent on the Azure AI service and return a local ChatAgent wrapper.
@@ -268,7 +268,7 @@ class AzureAIProjectAgentProvider(Generic[TOptions_co]):
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Retrieve an existing agent from the Azure AI service and return a local ChatAgent wrapper.
@@ -328,7 +328,7 @@ class AzureAIProjectAgentProvider(Generic[TOptions_co]):
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Wrap an SDK agent version object into a ChatAgent without making HTTP calls.
@@ -368,7 +368,7 @@ class AzureAIProjectAgentProvider(Generic[TOptions_co]):
details: AgentVersionDetails,
provided_tools: Sequence[ToolProtocol | MutableMapping[str, Any]] | None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Create a ChatAgent from an AgentVersionDetails.
python/packages/azure-ai/tests/test_azure_ai_agent_client.py
+38 -28
View File
@@ -91,6 +91,17 @@ def create_test_azure_ai_chat_client(
client._azure_search_tool_calls = [] # Add the new instance variable
client.additional_properties = {}
client.middleware = None
client.chat_middleware = []
client.function_middleware = []
client.otel_provider_name = "azure.ai"
client.function_invocation_configuration = {
"enabled": True,
"max_iterations": 5,
"max_consecutive_errors_per_request": 0,
"terminate_on_unknown_calls": False,
"additional_tools": [],
"include_detailed_errors": False,
}
return client
@@ -308,10 +319,10 @@ async def test_azure_ai_chat_client_thread_management_through_public_api(mock_ag
mock_stream.__aenter__ = AsyncMock(return_value=empty_async_iter())
mock_stream.__aexit__ = AsyncMock(return_value=None)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
# Call without existing thread - should create new one
response = chat_client.get_streaming_response(messages)
response = chat_client.get_response(messages, stream=True)
# Consume the generator to trigger the method execution
async for _ in response:
pass
@@ -335,7 +346,7 @@ async def test_azure_ai_chat_client_prepare_options_basic(mock_agents_client: Ma
"""Test _prepare_options with basic ChatOptions."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options: ChatOptions = {"max_tokens": 100, "temperature": 0.7}
run_options, tool_results = await chat_client._prepare_options(messages, chat_options) # type: ignore
@@ -348,7 +359,7 @@ async def test_azure_ai_chat_client_prepare_options_no_chat_options(mock_agents_
"""Test _prepare_options with default ChatOptions."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
run_options, tool_results = await chat_client._prepare_options(messages, {}) # type: ignore
@@ -365,7 +376,7 @@ async def test_azure_ai_chat_client_prepare_options_with_image_content(mock_agen
mock_agents_client.get_agent = AsyncMock(return_value=None)
image_content = Content.from_uri(uri="https://example.com/image.jpg", media_type="image/jpeg")
messages = [ChatMessage("user", [image_content])]
messages = [ChatMessage(role="user", contents=[image_content])]
run_options, _ = await chat_client._prepare_options(messages, {}) # type: ignore
@@ -454,8 +465,8 @@ async def test_azure_ai_chat_client_prepare_options_with_messages(mock_agents_cl
# Test with system message (becomes instruction)
messages = [
ChatMessage("system", ["You are a helpful assistant"]),
ChatMessage("user", ["Hello"]),
ChatMessage(role="system", text="You are a helpful assistant"),
ChatMessage(role="user", text="Hello"),
]
run_options, _ = await chat_client._prepare_options(messages, {}) # type: ignore
@@ -477,7 +488,7 @@ async def test_azure_ai_chat_client_prepare_options_with_instructions_from_optio
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
mock_agents_client.get_agent = AsyncMock(return_value=None)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options: ChatOptions = {
"instructions": "You are a thoughtful reviewer. Give brief feedback.",
}
@@ -500,8 +511,8 @@ async def test_azure_ai_chat_client_prepare_options_merges_instructions_from_mes
mock_agents_client.get_agent = AsyncMock(return_value=None)
messages = [
ChatMessage("system", ["Context: You are reviewing marketing copy."]),
ChatMessage("user", ["Review this tagline"]),
ChatMessage(role="system", text="Context: You are reviewing marketing copy."),
ChatMessage(role="user", text="Review this tagline"),
]
chat_options: ChatOptions = {
"instructions": "Be concise and constructive in your feedback.",
@@ -519,20 +530,18 @@ async def test_azure_ai_chat_client_prepare_options_merges_instructions_from_mes
async def test_azure_ai_chat_client_inner_get_response(mock_agents_client: MagicMock) -> None:
"""Test _inner_get_response method."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
messages = [ChatMessage("user", ["Hello"])]
chat_options: ChatOptions = {}
async def mock_streaming_response():
yield ChatResponseUpdate(role="assistant", text="Hello back")
yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("Hello back")])
with (
patch.object(chat_client, "_inner_get_streaming_response", return_value=mock_streaming_response()),
patch.object(chat_client, "_inner_get_response", return_value=mock_streaming_response()),
patch("agent_framework.ChatResponse.from_update_generator") as mock_from_generator,
):
mock_response = ChatResponse(messages=ChatMessage("assistant", ["Hello back"]))
mock_response = ChatResponse(messages=[ChatMessage(role="assistant", text="Hello back")])
mock_from_generator.return_value = mock_response
result = await chat_client._inner_get_response(messages=messages, options=chat_options) # type: ignore
result = await ChatResponse.from_update_generator(mock_streaming_response())
assert result is mock_response
mock_from_generator.assert_called_once()
@@ -672,7 +681,7 @@ async def test_azure_ai_chat_client_prepare_options_tool_choice_required_specifi
dict_tool = {"type": "function", "function": {"name": "test_function"}}
chat_options = {"tools": [dict_tool], "tool_choice": required_tool_mode}
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
run_options, _ = await chat_client._prepare_options(messages, chat_options) # type: ignore
@@ -717,7 +726,7 @@ async def test_azure_ai_chat_client_prepare_options_mcp_never_require(mock_agent
mcp_tool = HostedMCPTool(name="Test MCP Tool", url="https://example.com/mcp", approval_mode="never_require")
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options: ChatOptions = {"tools": [mcp_tool], "tool_choice": "auto"}
with patch("agent_framework_azure_ai._shared.McpTool") as mock_mcp_tool_class:
@@ -749,7 +758,7 @@ async def test_azure_ai_chat_client_prepare_options_mcp_with_headers(mock_agents
name="Test MCP Tool", url="https://example.com/mcp", headers=headers, approval_mode="never_require"
)
messages = [ChatMessage("user", ["Hello"])]
messages = [ChatMessage(role="user", text="Hello")]
chat_options: ChatOptions = {"tools": [mcp_tool], "tool_choice": "auto"}
with patch("agent_framework_azure_ai._shared.McpTool") as mock_mcp_tool_class:
@@ -1408,7 +1417,7 @@ async def test_azure_ai_chat_client_get_response() -> None:
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage("user", ["What's the weather like today?"]))
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
# Test that the agents_client can be used to get a response
response = await azure_ai_chat_client.get_response(messages=messages)
@@ -1426,7 +1435,7 @@ async def test_azure_ai_chat_client_get_response_tools() -> None:
assert isinstance(azure_ai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage("user", ["What's the weather like in Seattle?"]))
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
# Test that the agents_client can be used to get a response
response = await azure_ai_chat_client.get_response(
@@ -1454,10 +1463,10 @@ async def test_azure_ai_chat_client_streaming() -> None:
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage("user", ["What's the weather like today?"]))
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
# Test that the agents_client can be used to get a response
response = azure_ai_chat_client.get_streaming_response(messages=messages)
response = azure_ai_chat_client.get_response(messages=messages, stream=True)
full_message: str = ""
async for chunk in response:
@@ -1478,11 +1487,12 @@ async def test_azure_ai_chat_client_streaming_tools() -> None:
assert isinstance(azure_ai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage("user", ["What's the weather like in Seattle?"]))
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
# Test that the agents_client can be used to get a response
response = azure_ai_chat_client.get_streaming_response(
response = azure_ai_chat_client.get_response(
messages=messages,
stream=True,
options={"tools": [get_weather], "tool_choice": "auto"},
)
full_message: str = ""
@@ -1522,7 +1532,7 @@ async def test_azure_ai_chat_client_agent_basic_run_streaming() -> None:
) as agent:
# Run streaming query
full_message: str = ""
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
async for chunk in agent.run("Please respond with exactly: 'This is a streaming response test.'", stream=True):
assert chunk is not None
assert isinstance(chunk, AgentResponseUpdate)
if chunk.text:
@@ -2097,7 +2107,7 @@ def test_azure_ai_chat_client_prepare_messages_with_function_result(
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
function_result = Content.from_function_result(call_id='["run_123", "call_456"]', result="test result")
messages = [ChatMessage("user", [function_result])]
messages = [ChatMessage(role="user", contents=[function_result])]
additional_messages, instructions, required_action_results = chat_client._prepare_messages(messages) # type: ignore
@@ -2117,7 +2127,7 @@ def test_azure_ai_chat_client_prepare_messages_with_raw_content_block(
# Create content with raw_representation that is a MessageInputContentBlock
raw_block = MessageInputTextBlock(text="Raw block text")
custom_content = Content(type="custom", raw_representation=raw_block)
messages = [ChatMessage("user", [custom_content])]
messages = [ChatMessage(role="user", contents=[custom_content])]
additional_messages, instructions, required_action_results = chat_client._prepare_messages(messages) # type: ignore
python/packages/azure-ai/tests/test_azure_ai_client.py
+55 -32
View File
@@ -298,9 +298,9 @@ async def test_prepare_messages_for_azure_ai_with_system_messages(
client = create_test_azure_ai_client(mock_project_client)
messages = [
ChatMessage("system", [Content.from_text(text="You are a helpful assistant.")]),
ChatMessage("user", [Content.from_text(text="Hello")]),
ChatMessage("assistant", [Content.from_text(text="System response")]),
ChatMessage(role="system", contents=[Content.from_text(text="You are a helpful assistant.")]),
ChatMessage(role="user", contents=[Content.from_text(text="Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text(text="System response")]),
]
result_messages, instructions = client._prepare_messages_for_azure_ai(messages) # type: ignore
@@ -318,8 +318,8 @@ async def test_prepare_messages_for_azure_ai_no_system_messages(
client = create_test_azure_ai_client(mock_project_client)
messages = [
ChatMessage("user", [Content.from_text(text="Hello")]),
ChatMessage("assistant", [Content.from_text(text="Hi there!")]),
ChatMessage(role="user", contents=[Content.from_text(text="Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text(text="Hi there!")]),
]
result_messages, instructions = client._prepare_messages_for_azure_ai(messages) # type: ignore
@@ -419,10 +419,13 @@ async def test_prepare_options_basic(mock_project_client: MagicMock) -> None:
"""Test prepare_options basic functionality."""
client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0")
messages = [ChatMessage("user", [Content.from_text(text="Hello")])]
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
with (
patch.object(client.__class__.__bases__[0], "_prepare_options", return_value={"model": "test-model"}),
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
return_value={"model": "test-model"},
),
patch.object(
client,
"_get_agent_reference_or_create",
@@ -453,10 +456,13 @@ async def test_prepare_options_with_application_endpoint(
agent_version="1",
)
messages = [ChatMessage("user", [Content.from_text(text="Hello")])]
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
with (
patch.object(client.__class__.__bases__[0], "_prepare_options", return_value={"model": "test-model"}),
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
return_value={"model": "test-model"},
),
patch.object(
client,
"_get_agent_reference_or_create",
@@ -492,10 +498,13 @@ async def test_prepare_options_with_application_project_client(
agent_version="1",
)
messages = [ChatMessage("user", [Content.from_text(text="Hello")])]
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
with (
patch.object(client.__class__.__bases__[0], "_prepare_options", return_value={"model": "test-model"}),
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
return_value={"model": "test-model"},
),
patch.object(
client,
"_get_agent_reference_or_create",
@@ -968,13 +977,12 @@ async def test_prepare_options_excludes_response_format(
"""Test that prepare_options excludes response_format, text, and text_format from final run options."""
client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0")
messages = [ChatMessage("user", [Content.from_text(text="Hello")])]
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
chat_options: ChatOptions = {}
with (
patch.object(
client.__class__.__bases__[0],
"_prepare_options",
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
return_value={
"model": "test-model",
"response_format": ResponseFormatModel,
@@ -1299,7 +1307,8 @@ async def client() -> AsyncGenerator[AzureAIClient, None]:
)
try:
assert client.function_invocation_configuration
client.function_invocation_configuration.max_iterations = 1
# Need at least 2 iterations for tool_choice tests: one to get function call, one to get final response
client.function_invocation_configuration["max_iterations"] = 2
yield client
finally:
await project_client.agents.delete(agent_name=agent_name)
@@ -1354,10 +1363,10 @@ async def test_integration_options(
# Prepare test message
if option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [ChatMessage("user", ["What is the weather in Seattle?"])]
messages = [ChatMessage(role="user", text="What is the weather in Seattle?")]
else:
# Generic prompt for simple options
messages = [ChatMessage("user", ["Say 'Hello World' briefly."])]
messages = [ChatMessage(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options: dict[str, Any] = {option_name: option_value, "tools": [get_weather]}
@@ -1365,13 +1374,13 @@ async def test_integration_options(
for streaming in [False, True]:
if streaming:
# Test streaming mode
response_gen = client.get_streaming_response(
response_stream = client.get_response(
messages=messages,
stream=True,
options=options,
)
output_format = option_value if option_name == "response_format" else None
response = await ChatResponse.from_update_generator(response_gen, output_format_type=output_format)
response = await response_stream.get_final_response()
else:
# Test non-streaming mode
response = await client.get_response(
@@ -1381,12 +1390,26 @@ async def test_integration_options(
assert response is not None
assert isinstance(response, ChatResponse)
assert response.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# For tool_choice="required", we return after tool execution without a model text response
is_required_tool_choice = option_name == "tool_choice" and (
option_value == "required" or (isinstance(option_value, dict) and option_value.get("mode") == "required")
)
if is_required_tool_choice:
# Response should have function call and function result, but no text from model
assert len(response.messages) >= 2, f"Expected function call + result for {option_name}"
has_function_call = any(c.type == "function_call" for msg in response.messages for c in msg.contents)
has_function_result = any(c.type == "function_result" for msg in response.messages for c in msg.contents)
assert has_function_call, f"No function call in response for {option_name}"
assert has_function_result, f"No function result in response for {option_name}"
else:
assert response.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# Validate based on option type
if needs_validation:
if option_name.startswith("tool_choice"):
if option_name.startswith("tool_choice") and not is_required_tool_choice:
# Should have called the weather function
text = response.text.lower()
assert "sunny" in text or "seattle" in text, f"Tool not invoked for {option_name}"
@@ -1457,24 +1480,24 @@ async def test_integration_agent_options(
# Prepare test message
if option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [ChatMessage("user", ["The weather in Seattle is sunny"])]
messages.append(ChatMessage("user", ["What is the weather in Seattle?"]))
messages = [ChatMessage(role="user", text="The weather in Seattle is sunny")]
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
else:
# Generic prompt for simple options
messages = [ChatMessage("user", ["Say 'Hello World' briefly."])]
messages = [ChatMessage(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options = {option_name: option_value}
if streaming:
# Test streaming mode
response_gen = client.get_streaming_response(
response_stream = client.get_response(
messages=messages,
stream=True,
options=options,
)
output_format = option_value if option_name.startswith("response_format") else None
response = await ChatResponse.from_update_generator(response_gen, output_format_type=output_format)
response = await response_stream.get_final_response()
else:
# Test non-streaming mode
response = await client.get_response(
@@ -1516,7 +1539,7 @@ async def test_integration_web_search() -> None:
},
}
if streaming:
response = await ChatResponse.from_update_generator(client.get_streaming_response(**content))
response = await client.get_response(stream=True, **content).get_final_response()
else:
response = await client.get_response(**content)
@@ -1541,7 +1564,7 @@ async def test_integration_web_search() -> None:
},
}
if streaming:
response = await ChatResponse.from_update_generator(client.get_streaming_response(**content))
response = await client.get_response(stream=True, **content).get_final_response()
else:
response = await client.get_response(**content)
assert response.text is not None
python/packages/azure-ai/tests/test_shared.py
+20 -3
View File
@@ -1,6 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
from unittest.mock import MagicMock
import os
from unittest.mock import MagicMock, patch
import pytest
from agent_framework import (
@@ -78,8 +79,24 @@ def test_to_azure_ai_agent_tools_code_interpreter() -> None:
def test_to_azure_ai_agent_tools_web_search_missing_connection() -> None:
"""Test HostedWebSearchTool raises without connection info."""
tool = HostedWebSearchTool()
with pytest.raises(ServiceInitializationError, match="Bing search tool requires"):
to_azure_ai_agent_tools([tool])
# Clear any environment variables that could provide connection info
with patch.dict(
os.environ,
{"BING_CONNECTION_ID": "", "BING_CUSTOM_CONNECTION_ID": "", "BING_CUSTOM_INSTANCE_NAME": ""},
clear=False,
):
# Also need to unset the keys if they exist
env_backup = {}
for key in ["BING_CONNECTION_ID", "BING_CUSTOM_CONNECTION_ID", "BING_CUSTOM_INSTANCE_NAME"]:
env_backup[key] = os.environ.pop(key, None)
try:
with pytest.raises(ServiceInitializationError, match="Bing search tool requires"):
to_azure_ai_agent_tools([tool])
finally:
# Restore environment
for key, value in env_backup.items():
if value is not None:
os.environ[key] = value
def test_to_azure_ai_agent_tools_dict_passthrough() -> None:
python/packages/azurefunctions/pyproject.toml
+1
View File
@@ -43,6 +43,7 @@ environments = [
fallback-version = "0.0.0"
[tool.pytest.ini_options]
testpaths = 'tests'
pythonpath = ["tests/integration_tests"]
addopts = "-ra -q -r fEX"
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "function"
python/packages/azurefunctions/tests/integration_tests/conftest.py
+479 -25
View File
@@ -1,34 +1,468 @@
# Copyright (c) Microsoft. All rights reserved.
"""
Pytest configuration for Durable Agent Framework tests.
Pytest configuration for Azure Functions integration tests.
This module provides fixtures and configuration for pytest.
This module provides fixtures, configuration, and test utilities for pytest.
"""
import os
import shutil
import socket
import subprocess
import sys
import time
import uuid
from collections.abc import Iterator, Mapping
from contextlib import suppress
from pathlib import Path
from typing import Any
import pytest
import requests
# Add the integration_tests directory to the path so testutils can be imported
sys.path.insert(0, str(Path(__file__).parent))
# =============================================================================
# Configuration Constants
# =============================================================================
from testutils import (
FunctionAppStartupError,
build_base_url,
cleanup_function_app,
find_available_port,
get_sample_path_from_marker,
load_and_validate_env,
start_function_app,
wait_for_function_app_ready,
TIMEOUT = 30 # seconds
ORCHESTRATION_TIMEOUT = 180 # seconds for orchestrations
_DEFAULT_HOST = "localhost"
# Emulator ports (match CI workflow configuration)
_AZURITE_BLOB_PORT = 10000
_DTS_EMULATOR_PORT = 8080
# =============================================================================
# Exceptions
# =============================================================================
class FunctionAppStartupError(RuntimeError):
"""Raised when the Azure Functions host fails to start reliably."""
pass
# =============================================================================
# Environment and Service Checks
# =============================================================================
def _load_env_file_if_present() -> None:
"""Load environment variables from the local .env file when available."""
env_file = Path(__file__).parent / ".env"
if not env_file.exists():
return
try:
from dotenv import load_dotenv
load_dotenv(env_file)
except ImportError:
# python-dotenv not available; rely on existing environment
pass
def _check_func_cli_available() -> bool:
"""Check if Azure Functions Core Tools (func) is installed and available."""
return shutil.which("func") is not None
def _check_port_listening(port: int, host: str = _DEFAULT_HOST) -> bool:
"""Check if a service is listening on the given port."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.settimeout(1)
return sock.connect_ex((host, port)) == 0
def _check_azurite_available() -> bool:
"""Check if Azurite (Azure Storage emulator) is available on the expected port."""
return _check_port_listening(_AZURITE_BLOB_PORT)
def _check_dts_emulator_available() -> bool:
"""Check if Durable Task Scheduler emulator is available on the expected port."""
return _check_port_listening(_DTS_EMULATOR_PORT)
def _should_skip_azure_functions_integration_tests() -> tuple[bool, str]:
"""Determine whether Azure Functions integration tests should be skipped."""
_load_env_file_if_present()
run_integration_tests = os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
if not run_integration_tests:
return (
True,
"Integration tests are disabled. Set RUN_INTEGRATION_TESTS=true to enable Azure Functions sample tests.",
)
# Check for Azure Functions Core Tools
if not _check_func_cli_available():
return (
True,
"Azure Functions Core Tools (func) not installed. Install with: npm install -g azure-functions-core-tools@4", # noqa: E501
)
# Check for Azurite (Azure Storage emulator)
if not _check_azurite_available():
return (
True,
f"Azurite not running on port {_AZURITE_BLOB_PORT}. Start with: docker run -d -p 10000:10000 -p 10001:10001 -p 10002:10002 mcr.microsoft.com/azure-storage/azurite", # noqa: E501
)
# Check for Durable Task Scheduler emulator
if not _check_dts_emulator_available():
return (
True,
f"Durable Task Scheduler emulator not running on port {_DTS_EMULATOR_PORT}. Start with: docker run -d -p 8080:8080 -p 8082:8082 mcr.microsoft.com/dts/dts-emulator:latest", # noqa: E501
)
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "").strip()
if not endpoint or endpoint == "https://your-resource.openai.azure.com/":
return True, "No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
deployment_name = os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", "").strip()
if not deployment_name or deployment_name == "your-deployment-name":
return True, "No real AZURE_OPENAI_CHAT_DEPLOYMENT_NAME provided; skipping integration tests."
return False, "Integration tests enabled."
_SKIP_AZURE_FUNCTIONS_INTEGRATION_TESTS, _AZURE_FUNCTIONS_SKIP_REASON = _should_skip_azure_functions_integration_tests()
skip_if_azure_functions_integration_tests_disabled = pytest.mark.skipif(
_SKIP_AZURE_FUNCTIONS_INTEGRATION_TESTS,
reason=_AZURE_FUNCTIONS_SKIP_REASON,
)
# =============================================================================
# Test Helper Class
# =============================================================================
class SampleTestHelper:
"""Helper class for testing samples."""
@staticmethod
def post_json(url: str, data: dict[str, Any], timeout: int = TIMEOUT) -> requests.Response:
"""POST JSON data to a URL."""
return requests.post(url, json=data, headers={"Content-Type": "application/json"}, timeout=timeout)
@staticmethod
def post_text(url: str, text: str, timeout: int = TIMEOUT) -> requests.Response:
"""POST plain text to a URL."""
return requests.post(url, data=text, headers={"Content-Type": "text/plain"}, timeout=timeout)
@staticmethod
def get(url: str, timeout: int = TIMEOUT) -> requests.Response:
"""GET request to a URL."""
return requests.get(url, timeout=timeout)
@staticmethod
def wait_for_orchestration(
status_url: str, max_wait: int = ORCHESTRATION_TIMEOUT, poll_interval: int = 2
) -> dict[str, Any]:
"""Wait for an orchestration to complete.
Args:
status_url: URL to poll for orchestration status
max_wait: Maximum seconds to wait
poll_interval: Seconds between polls
Returns:
Final orchestration status
Raises:
TimeoutError: If orchestration doesn't complete in time
"""
start_time = time.time()
while time.time() - start_time < max_wait:
response = requests.get(status_url, timeout=TIMEOUT)
response.raise_for_status()
status = response.json()
runtime_status = status.get("runtimeStatus", "")
if runtime_status in ["Completed", "Failed", "Terminated"]:
return status
time.sleep(poll_interval)
raise TimeoutError(f"Orchestration did not complete within {max_wait} seconds")
@staticmethod
def wait_for_orchestration_with_output(
status_url: str, max_wait: int = ORCHESTRATION_TIMEOUT, poll_interval: int = 2
) -> dict[str, Any]:
"""Wait for an orchestration to complete and have output available.
This is a specialized version of wait_for_orchestration that also
ensures the output field is present, handling timing race conditions.
Args:
status_url: URL to poll for orchestration status
max_wait: Maximum seconds to wait
poll_interval: Seconds between polls
Returns:
Final orchestration status with output
Raises:
TimeoutError: If orchestration doesn't complete with output in time
"""
start_time = time.time()
while time.time() - start_time < max_wait:
response = requests.get(status_url, timeout=TIMEOUT)
response.raise_for_status()
status = response.json()
runtime_status = status.get("runtimeStatus", "")
if runtime_status in ["Failed", "Terminated"]:
return status
if runtime_status == "Completed" and status.get("output"):
return status
# If completed but no output, continue polling for a bit more to
# handle the race condition where output has not been persisted yet.
time.sleep(poll_interval)
# Provide detailed error message based on final status
final_response = requests.get(status_url, timeout=TIMEOUT)
final_response.raise_for_status()
final_status = final_response.json()
final_runtime_status = final_status.get("runtimeStatus", "Unknown")
if final_runtime_status == "Completed":
if "output" not in final_status:
raise TimeoutError(
"Orchestration completed but 'output' field is missing after "
f"{max_wait} seconds. Final status: {final_status}"
)
if not final_status["output"]:
raise TimeoutError(
"Orchestration completed but output is empty after "
f"{max_wait} seconds. Final status: {final_status}"
)
raise TimeoutError(
"Orchestration completed with output but validation failed after "
f"{max_wait} seconds. Final status: {final_status}"
)
raise TimeoutError(
"Orchestration did not complete within "
f"{max_wait} seconds. Final status: {final_runtime_status}, "
f"Full status: {final_status}"
)
# =============================================================================
# Function App Lifecycle Management
# =============================================================================
def _resolve_repo_root() -> Path:
"""Resolve the repository root, preferring GITHUB_WORKSPACE when available."""
workspace = os.getenv("GITHUB_WORKSPACE")
if workspace:
candidate = Path(workspace).expanduser()
if not (candidate / "samples").exists() and (candidate / "python" / "samples").exists():
return (candidate / "python").resolve()
return candidate.resolve()
# If `GITHUB_WORKSPACE` is not set,
# go up from conftest.py -> integration_tests -> tests -> azurefunctions -> packages -> python
return Path(__file__).resolve().parents[4]
def _get_sample_path_from_marker(request: pytest.FixtureRequest) -> tuple[Path | None, str | None]:
"""Get sample path from @pytest.mark.sample() marker.
Returns a tuple of (sample_path, error_message).
If successful, error_message is None.
If failed, sample_path is None and error_message contains the reason.
"""
marker = request.node.get_closest_marker("sample")
if not marker:
return (
None,
(
"No @pytest.mark.sample() marker found on test. Add pytestmark with "
"@pytest.mark.sample('sample_name') to the test module."
),
)
if not marker.args:
return (
None,
"@pytest.mark.sample() marker found but no sample name provided. Use @pytest.mark.sample('sample_name').",
)
sample_name = marker.args[0]
repo_root = _resolve_repo_root()
sample_path = repo_root / "samples" / "getting_started" / "azure_functions" / sample_name
if not sample_path.exists():
return None, f"Sample directory does not exist: {sample_path}"
return sample_path, None
def _find_available_port(host: str = _DEFAULT_HOST) -> int:
"""Find an available TCP port on the given host."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind((host, 0))
return sock.getsockname()[1]
def _build_base_url(port: int, host: str = _DEFAULT_HOST) -> str:
"""Construct a base URL for the Azure Functions host."""
return f"http://{host}:{port}"
def _is_port_in_use(port: int, host: str = _DEFAULT_HOST) -> bool:
"""Check if a port is already in use.
Returns True if the port is in use, False otherwise.
"""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
return sock.connect_ex((host, port)) == 0
def _load_and_validate_env() -> None:
"""Load .env file from current directory if it exists, then validate required environment variables.
Raises pytest.fail if required environment variables are missing.
"""
_load_env_file_if_present()
# Required environment variables for Azure Functions samples
# These match the variables defined in .env.example
required_env_vars = [
"AZURE_OPENAI_ENDPOINT",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME",
"AzureWebJobsStorage",
"DURABLE_TASK_SCHEDULER_CONNECTION_STRING",
"FUNCTIONS_WORKER_RUNTIME",
]
# Check if required env vars are set
missing_vars = [var for var in required_env_vars if not os.environ.get(var)]
if missing_vars:
pytest.fail(
f"Missing required environment variables: {', '.join(missing_vars)}. "
"Please create a .env file in tests/integration_tests/ based on .env.example or "
"set these variables in your environment."
)
def _start_function_app(sample_path: Path, port: int) -> subprocess.Popen[Any]:
"""Start a function app in the specified sample directory.
Returns the subprocess.Popen object for the running process.
"""
env = os.environ.copy()
# Use a unique TASKHUB_NAME for each test run to ensure test isolation.
# This prevents conflicts between parallel or repeated test runs, as Durable Functions
# use the task hub name to separate orchestration state.
env["TASKHUB_NAME"] = f"test{uuid.uuid4().hex[:8]}"
# On Windows, use CREATE_NEW_PROCESS_GROUP to allow proper termination
# shell=True only on Windows to handle PATH resolution
if sys.platform == "win32":
return subprocess.Popen(
["func", "start", "--port", str(port)],
cwd=str(sample_path),
creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
shell=True,
env=env,
)
# On Unix, don't use shell=True to avoid shell wrapper issues
return subprocess.Popen(["func", "start", "--port", str(port)], cwd=str(sample_path), env=env)
def _wait_for_function_app_ready(func_process: subprocess.Popen[Any], port: int, max_wait: int = 60) -> None:
"""Block until the Azure Functions host responds healthy or fail fast."""
start_time = time.time()
health_url = f"{_build_base_url(port)}/api/health"
last_error: Exception | None = None
while time.time() - start_time < max_wait:
# If the process exited early, capture any previously seen error and fail fast.
if func_process.poll() is not None:
raise FunctionAppStartupError(
f"Function app process exited with code {func_process.returncode} before becoming healthy"
) from last_error
if _is_port_in_use(port):
try:
response = requests.get(health_url, timeout=5)
if response.status_code == 200:
return
last_error = RuntimeError(f"Health check returned {response.status_code}")
except requests.RequestException as exc:
last_error = exc
time.sleep(1)
raise FunctionAppStartupError(
f"Function app did not become healthy on port {port} within {max_wait} seconds"
) from last_error
def _cleanup_function_app(func_process: subprocess.Popen[Any]) -> None:
"""Clean up the function app process and all its children.
Uses psutil if available for more thorough cleanup, falls back to basic termination.
"""
try:
import psutil
if func_process.poll() is None: # Process still running
# Get parent process
parent = psutil.Process(func_process.pid)
# Get all child processes recursively
children = parent.children(recursive=True)
# Kill children first
for child in children:
with suppress(psutil.NoSuchProcess, psutil.AccessDenied):
child.kill()
# Kill parent
with suppress(psutil.NoSuchProcess, psutil.AccessDenied):
parent.kill()
# Wait for all to terminate
_gone, alive = psutil.wait_procs(children + [parent], timeout=3)
# Force kill any remaining
for proc in alive:
with suppress(psutil.NoSuchProcess, psutil.AccessDenied):
proc.kill()
except ImportError:
# Fallback if psutil not available
try:
if func_process.poll() is None:
func_process.kill()
func_process.wait()
except Exception:
# Ignore all exceptions during fallback cleanup; best effort to terminate process.
pass
except Exception:
pass # Best effort cleanup
# Give the port time to be released
time.sleep(2)
# =============================================================================
# Pytest Configuration
# =============================================================================
def pytest_configure(config: pytest.Config) -> None:
"""Register custom markers."""
config.addinivalue_line("markers", "orchestration: marks tests that use orchestrations (require Azurite)")
@@ -38,10 +472,25 @@ def pytest_configure(config: pytest.Config) -> None:
)
def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
"""Skip integration tests in this directory if prerequisites are not met."""
should_skip, reason = _should_skip_azure_functions_integration_tests()
if should_skip:
skip_marker = pytest.mark.skip(reason=reason)
for item in items:
# Only skip items that are in this integration_tests directory
if "integration_tests" in str(item.fspath):
item.add_marker(skip_marker)
# =============================================================================
# Pytest Fixtures
# =============================================================================
@pytest.fixture(scope="session")
def function_app_running() -> bool:
"""
Check if the function app is running on localhost:7071.
"""Check if the function app is running on localhost:7071.
This fixture can be used to skip tests if the function app is not available.
"""
@@ -61,8 +510,7 @@ def skip_if_no_function_app(function_app_running: bool) -> None:
@pytest.fixture(scope="module")
def function_app_for_test(request: pytest.FixtureRequest) -> Iterator[dict[str, int | str]]:
"""
Start the function app for the corresponding sample based on marker.
"""Start the function app for the corresponding sample based on marker.
This fixture:
1. Determines which sample to run from @pytest.mark.sample()
@@ -78,14 +526,14 @@ def function_app_for_test(request: pytest.FixtureRequest) -> Iterator[dict[str,
...
"""
# Get sample path from marker
sample_path, error_message = get_sample_path_from_marker(request)
sample_path, error_message = _get_sample_path_from_marker(request)
if error_message:
pytest.fail(error_message)
assert sample_path is not None, "Sample path must be resolved before starting the function app"
# Load .env file if it exists and validate required env vars
load_and_validate_env()
_load_and_validate_env()
max_attempts = 3
last_error: Exception | None = None
@@ -94,17 +542,17 @@ def function_app_for_test(request: pytest.FixtureRequest) -> Iterator[dict[str,
port = 0
for _ in range(max_attempts):
port = find_available_port()
base_url = build_base_url(port)
func_process = start_function_app(sample_path, port)
port = _find_available_port()
base_url = _build_base_url(port)
func_process = _start_function_app(sample_path, port)
try:
wait_for_function_app_ready(func_process, port)
_wait_for_function_app_ready(func_process, port)
last_error = None
break
except FunctionAppStartupError as exc:
last_error = exc
cleanup_function_app(func_process)
_cleanup_function_app(func_process)
func_process = None
if func_process is None:
@@ -117,10 +565,16 @@ def function_app_for_test(request: pytest.FixtureRequest) -> Iterator[dict[str,
yield {"base_url": base_url, "port": port}
finally:
if func_process is not None:
cleanup_function_app(func_process)
_cleanup_function_app(func_process)
@pytest.fixture(scope="module")
def base_url(function_app_for_test: Mapping[str, int | str]) -> str:
"""Expose the function app's base URL to tests."""
return str(function_app_for_test["base_url"])
@pytest.fixture(scope="session")
def sample_helper() -> type[SampleTestHelper]:
"""Provide the SampleTestHelper class for tests."""
return SampleTestHelper
python/packages/azurefunctions/tests/integration_tests/test_01_single_agent.py
+11 -12
View File
@@ -16,13 +16,11 @@ Usage:
import pytest
from agent_framework_durabletask import THREAD_ID_HEADER
from testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
# Module-level markers - applied to all tests in this file
pytestmark = [
pytest.mark.sample("01_single_agent"),
pytest.mark.usefixtures("function_app_for_test"),
skip_if_azure_functions_integration_tests_disabled,
]
@@ -30,20 +28,21 @@ class TestSampleSingleAgent:
"""Tests for 01_single_agent sample."""
@pytest.fixture(autouse=True)
def _set_base_url(self, base_url: str) -> None:
"""Provide agent-specific base URL for the tests."""
def _setup(self, base_url: str, sample_helper) -> None:
"""Provide agent-specific base URL and helper for the tests."""
self.base_url = f"{base_url}/api/agents/Joker"
self.helper = sample_helper
def test_health_check(self, base_url: str) -> None:
def test_health_check(self, base_url: str, sample_helper) -> None:
"""Test health check endpoint."""
response = SampleTestHelper.get(f"{base_url}/api/health")
response = sample_helper.get(f"{base_url}/api/health")
assert response.status_code == 200
data = response.json()
assert data["status"] == "healthy"
def test_simple_message_json(self) -> None:
"""Test sending a simple message with JSON payload."""
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{self.base_url}/run",
{"message": "Tell me a short joke about cloud computing.", "thread_id": "test-simple-json"},
)
@@ -62,7 +61,7 @@ class TestSampleSingleAgent:
def test_simple_message_plain_text(self) -> None:
"""Test sending a message with plain text payload."""
response = SampleTestHelper.post_text(f"{self.base_url}/run", "Tell me a short joke about networking.")
response = self.helper.post_text(f"{self.base_url}/run", "Tell me a short joke about networking.")
assert response.status_code in [200, 202]
# Agent responded with plain text when the request body was text/plain.
@@ -71,7 +70,7 @@ class TestSampleSingleAgent:
def test_thread_id_in_query(self) -> None:
"""Test using thread_id in query parameter."""
response = SampleTestHelper.post_text(
response = self.helper.post_text(
f"{self.base_url}/run?thread_id=test-query-thread", "Tell me a short joke about weather in Texas."
)
assert response.status_code in [200, 202]
@@ -84,7 +83,7 @@ class TestSampleSingleAgent:
thread_id = "test-continuity"
# First message
response1 = SampleTestHelper.post_json(
response1 = self.helper.post_json(
f"{self.base_url}/run",
{"message": "Tell me a short joke about weather in Seattle.", "thread_id": thread_id},
)
@@ -95,7 +94,7 @@ class TestSampleSingleAgent:
assert data1["message_count"] == 2 # Initial + reply
# Second message in same session
response2 = SampleTestHelper.post_json(
response2 = self.helper.post_json(
f"{self.base_url}/run", {"message": "What about San Francisco?", "thread_id": thread_id}
)
assert response2.status_code == 200
@@ -104,7 +103,7 @@ class TestSampleSingleAgent:
else:
# In async mode, we can't easily test message count
# Just verify we can make multiple calls
response2 = SampleTestHelper.post_json(
response2 = self.helper.post_json(
f"{self.base_url}/run", {"message": "What about Texas?", "thread_id": thread_id}
)
assert response2.status_code == 202
python/packages/azurefunctions/tests/integration_tests/test_02_multi_agent.py
+4 -5
View File
@@ -15,13 +15,11 @@ Usage:
"""
import pytest
from testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
# Module-level markers - applied to all tests in this file
pytestmark = [
pytest.mark.sample("02_multi_agent"),
pytest.mark.usefixtures("function_app_for_test"),
skip_if_azure_functions_integration_tests_disabled,
]
@@ -29,14 +27,15 @@ class TestSampleMultiAgent:
"""Tests for 02_multi_agent sample."""
@pytest.fixture(autouse=True)
def _set_agent_urls(self, base_url: str) -> None:
def _setup(self, base_url: str, sample_helper) -> None:
"""Configure base URLs for Weather and Math agents."""
self.weather_base_url = f"{base_url}/api/agents/WeatherAgent"
self.math_base_url = f"{base_url}/api/agents/MathAgent"
self.helper = sample_helper
def test_weather_agent(self) -> None:
"""Test WeatherAgent endpoint."""
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{self.weather_base_url}/run",
{"message": "What is the weather in Seattle?"},
)
@@ -47,7 +46,7 @@ class TestSampleMultiAgent:
def test_math_agent(self) -> None:
"""Test MathAgent endpoint."""
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{self.math_base_url}/run",
{"message": "Calculate a 20% tip on a $50 bill", "wait_for_response": False},
)
python/packages/azurefunctions/tests/integration_tests/test_03_reliable_streaming.py
+7 -10
View File
@@ -19,16 +19,12 @@ import time
import pytest
import requests
from testutils import (
SampleTestHelper,
skip_if_azure_functions_integration_tests_disabled,
)
# Module-level markers - applied to all tests in this file
pytestmark = [
pytest.mark.sample("03_reliable_streaming"),
pytest.mark.usefixtures("function_app_for_test"),
skip_if_azure_functions_integration_tests_disabled,
pytest.mark.skip(reason="Temp disabled to fix test instability - needs investigation into root cause"),
]
@@ -36,16 +32,17 @@ class TestSampleReliableStreaming:
"""Tests for 03_reliable_streaming sample."""
@pytest.fixture(autouse=True)
def _set_base_url(self, base_url: str) -> None:
"""Provide the base URL for each test."""
def _setup(self, base_url: str, sample_helper) -> None:
"""Provide the base URL and helper for each test."""
self.base_url = base_url
self.agent_url = f"{base_url}/api/agents/TravelPlanner"
self.stream_url = f"{base_url}/api/agent/stream"
self.helper = sample_helper
def test_agent_run_and_stream(self) -> None:
"""Test agent execution with Redis streaming."""
# Start agent run
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{self.agent_url}/run",
{"message": "Plan a 1-day trip to Seattle in 1 sentence", "wait_for_response": False},
)
@@ -69,7 +66,7 @@ class TestSampleReliableStreaming:
def test_stream_with_sse_format(self) -> None:
"""Test streaming with Server-Sent Events format."""
# Start agent run
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{self.agent_url}/run",
{"message": "What's the weather like?", "wait_for_response": False},
)
@@ -113,7 +110,7 @@ class TestSampleReliableStreaming:
def test_health_endpoint(self) -> None:
"""Test health check endpoint."""
response = SampleTestHelper.get(f"{self.base_url}/api/health")
response = self.helper.get(f"{self.base_url}/api/health")
assert response.status_code == 200
data = response.json()
assert data["status"] == "healthy"
python/packages/azurefunctions/tests/integration_tests/test_04_single_agent_orchestration_chaining.py
+7 -4
View File
@@ -19,13 +19,11 @@ Usage:
"""
import pytest
from testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
# Module-level markers - applied to all tests in this file
pytestmark = [
pytest.mark.sample("04_single_agent_orchestration_chaining"),
pytest.mark.usefixtures("function_app_for_test"),
skip_if_azure_functions_integration_tests_disabled,
]
@@ -33,17 +31,22 @@ pytestmark = [
class TestSampleOrchestrationChaining:
"""Tests for 04_single_agent_orchestration_chaining sample."""
@pytest.fixture(autouse=True)
def _setup(self, sample_helper) -> None:
"""Provide the helper for each test."""
self.helper = sample_helper
def test_orchestration_chaining(self, base_url: str) -> None:
"""Test sequential agent calls in orchestration."""
# Start orchestration
response = SampleTestHelper.post_json(f"{base_url}/api/singleagent/run", {})
response = self.helper.post_json(f"{base_url}/api/singleagent/run", {})
assert response.status_code == 202
data = response.json()
assert "instanceId" in data
assert "statusQueryGetUri" in data
# Wait for completion with output available
status = SampleTestHelper.wait_for_orchestration_with_output(data["statusQueryGetUri"])
status = self.helper.wait_for_orchestration_with_output(data["statusQueryGetUri"])
assert status["runtimeStatus"] == "Completed"
assert "output" in status
python/packages/azurefunctions/tests/integration_tests/test_05_multi_agent_orchestration_concurrency.py
+7 -4
View File
@@ -19,31 +19,34 @@ Usage:
"""
import pytest
from testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
# Module-level markers - applied to all tests in this file
pytestmark = [
pytest.mark.orchestration,
pytest.mark.sample("05_multi_agent_orchestration_concurrency"),
pytest.mark.usefixtures("function_app_for_test"),
skip_if_azure_functions_integration_tests_disabled,
]
class TestSampleMultiAgentConcurrency:
"""Tests for 05_multi_agent_orchestration_concurrency sample."""
@pytest.fixture(autouse=True)
def _setup(self, sample_helper) -> None:
"""Provide the helper for each test."""
self.helper = sample_helper
def test_concurrent_agents(self, base_url: str) -> None:
"""Test multiple agents running concurrently."""
# Start orchestration
response = SampleTestHelper.post_text(f"{base_url}/api/multiagent/run", "What is temperature?")
response = self.helper.post_text(f"{base_url}/api/multiagent/run", "What is temperature?")
assert response.status_code == 202
data = response.json()
assert "instanceId" in data
assert "statusQueryGetUri" in data
# Wait for completion
status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
status = self.helper.wait_for_orchestration(data["statusQueryGetUri"])
assert status["runtimeStatus"] == "Completed"
output = status["output"]
assert "physicist" in output
python/packages/azurefunctions/tests/integration_tests/test_06_multi_agent_orchestration_conditionals.py
+9 -6
View File
@@ -19,23 +19,26 @@ Usage:
"""
import pytest
from testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
# Module-level markers - applied to all tests in this file
pytestmark = [
pytest.mark.orchestration,
pytest.mark.sample("06_multi_agent_orchestration_conditionals"),
pytest.mark.usefixtures("function_app_for_test"),
skip_if_azure_functions_integration_tests_disabled,
]
class TestSampleMultiAgentConditionals:
"""Tests for 06_multi_agent_orchestration_conditionals sample."""
@pytest.fixture(autouse=True)
def _setup(self, sample_helper) -> None:
"""Provide the helper for each test."""
self.helper = sample_helper
def test_legitimate_email(self, base_url: str) -> None:
"""Test conditional logic with legitimate email."""
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{base_url}/api/spamdetection/run",
{
"email_id": "email-test-001",
@@ -48,13 +51,13 @@ class TestSampleMultiAgentConditionals:
assert "statusQueryGetUri" in data
# Wait for completion
status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
status = self.helper.wait_for_orchestration(data["statusQueryGetUri"])
assert status["runtimeStatus"] == "Completed"
assert "Email sent:" in status["output"]
def test_spam_email(self, base_url: str) -> None:
"""Test conditional logic with spam email."""
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{base_url}/api/spamdetection/run",
{"email_id": "email-test-002", "email_content": "URGENT! You have won $1,000,000! Click here now!"},
)
@@ -63,7 +66,7 @@ class TestSampleMultiAgentConditionals:
assert "instanceId" in data
# Wait for completion
status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
status = self.helper.wait_for_orchestration(data["statusQueryGetUri"])
assert status["runtimeStatus"] == "Completed"
assert "Email marked as spam:" in status["output"]
python/packages/azurefunctions/tests/integration_tests/test_07_single_agent_orchestration_hitl.py
+20 -21
View File
@@ -21,13 +21,11 @@ Usage:
import time
import pytest
from testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
# Module-level markers - applied to all tests in this file
pytestmark = [
pytest.mark.sample("07_single_agent_orchestration_hitl"),
pytest.mark.usefixtures("function_app_for_test"),
skip_if_azure_functions_integration_tests_disabled,
]
@@ -36,14 +34,15 @@ class TestSampleHITLOrchestration:
"""Tests for 07_single_agent_orchestration_hitl sample."""
@pytest.fixture(autouse=True)
def _set_hitl_base_url(self, base_url: str) -> None:
"""Prepare the HITL API base URL for the module's tests."""
def _setup(self, base_url: str, sample_helper) -> None:
"""Provide the helper and base URL for each test."""
self.hitl_base_url = f"{base_url}/api/hitl"
self.helper = sample_helper
def test_hitl_orchestration_approval(self) -> None:
"""Test HITL orchestration with human approval."""
# Start orchestration
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{self.hitl_base_url}/run",
{"topic": "artificial intelligence", "max_review_attempts": 3, "approval_timeout_hours": 1.0},
)
@@ -58,13 +57,13 @@ class TestSampleHITLOrchestration:
time.sleep(5)
# Check status to ensure it's waiting for approval
status_response = SampleTestHelper.get(data["statusQueryGetUri"])
status_response = self.helper.get(data["statusQueryGetUri"])
assert status_response.status_code == 200
status = status_response.json()
assert status["runtimeStatus"] in ["Running", "Pending"]
# Send approval
approval_response = SampleTestHelper.post_json(
approval_response = self.helper.post_json(
f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""}
)
assert approval_response.status_code == 200
@@ -72,7 +71,7 @@ class TestSampleHITLOrchestration:
assert approval_data["approved"] is True
# Wait for orchestration to complete
status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
status = self.helper.wait_for_orchestration(data["statusQueryGetUri"])
assert status["runtimeStatus"] == "Completed"
assert "output" in status
assert "content" in status["output"]
@@ -80,7 +79,7 @@ class TestSampleHITLOrchestration:
def test_hitl_orchestration_rejection_with_feedback(self) -> None:
"""Test HITL orchestration with rejection and subsequent approval."""
# Start orchestration
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{self.hitl_base_url}/run",
{"topic": "machine learning", "max_review_attempts": 3, "approval_timeout_hours": 1.0},
)
@@ -92,7 +91,7 @@ class TestSampleHITLOrchestration:
time.sleep(5)
# Send rejection with feedback
rejection_response = SampleTestHelper.post_json(
rejection_response = self.helper.post_json(
f"{self.hitl_base_url}/approve/{instance_id}",
{"approved": False, "feedback": "Please make it more concise and focus on practical applications."},
)
@@ -102,25 +101,25 @@ class TestSampleHITLOrchestration:
time.sleep(5)
# Check status - should still be running
status_response = SampleTestHelper.get(data["statusQueryGetUri"])
status_response = self.helper.get(data["statusQueryGetUri"])
assert status_response.status_code == 200
status = status_response.json()
assert status["runtimeStatus"] in ["Running", "Pending"]
# Now approve the revised content
approval_response = SampleTestHelper.post_json(
approval_response = self.helper.post_json(
f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""}
)
assert approval_response.status_code == 200
# Wait for completion
status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
status = self.helper.wait_for_orchestration(data["statusQueryGetUri"])
assert status["runtimeStatus"] == "Completed"
assert "output" in status
def test_hitl_orchestration_missing_topic(self) -> None:
"""Test HITL orchestration with missing topic."""
response = SampleTestHelper.post_json(f"{self.hitl_base_url}/run", {"max_review_attempts": 3})
response = self.helper.post_json(f"{self.hitl_base_url}/run", {"max_review_attempts": 3})
assert response.status_code == 400
data = response.json()
assert "error" in data
@@ -128,7 +127,7 @@ class TestSampleHITLOrchestration:
def test_hitl_get_status(self) -> None:
"""Test getting orchestration status."""
# Start orchestration
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{self.hitl_base_url}/run",
{"topic": "quantum computing", "max_review_attempts": 2, "approval_timeout_hours": 1.0},
)
@@ -137,7 +136,7 @@ class TestSampleHITLOrchestration:
instance_id = data["instanceId"]
# Get status
status_response = SampleTestHelper.get(f"{self.hitl_base_url}/status/{instance_id}")
status_response = self.helper.get(f"{self.hitl_base_url}/status/{instance_id}")
assert status_response.status_code == 200
status = status_response.json()
assert "instanceId" in status
@@ -146,12 +145,12 @@ class TestSampleHITLOrchestration:
# Cleanup: approve to complete orchestration
time.sleep(5)
SampleTestHelper.post_json(f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""})
self.helper.post_json(f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""})
def test_hitl_approval_invalid_payload(self) -> None:
"""Test sending approval with invalid payload."""
# Start orchestration first
response = SampleTestHelper.post_json(
response = self.helper.post_json(
f"{self.hitl_base_url}/run",
{"topic": "test topic", "max_review_attempts": 1, "approval_timeout_hours": 1.0},
)
@@ -162,7 +161,7 @@ class TestSampleHITLOrchestration:
time.sleep(3)
# Send approval without 'approved' field
approval_response = SampleTestHelper.post_json(
approval_response = self.helper.post_json(
f"{self.hitl_base_url}/approve/{instance_id}", {"feedback": "Some feedback"}
)
assert approval_response.status_code == 400
@@ -170,11 +169,11 @@ class TestSampleHITLOrchestration:
assert "error" in error_data
# Cleanup
SampleTestHelper.post_json(f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""})
self.helper.post_json(f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""})
def test_hitl_status_invalid_instance(self) -> None:
"""Test getting status for non-existent instance."""
response = SampleTestHelper.get(f"{self.hitl_base_url}/status/invalid-instance-id")
response = self.helper.get(f"{self.hitl_base_url}/status/invalid-instance-id")
assert response.status_code == 404
data = response.json()
assert "error" in data
python/packages/azurefunctions/tests/integration_tests/testutils.py
-397
View File
@@ -1,397 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""
Shared test helper utilities for sample integration tests.
This module provides common utilities for testing Azure Functions samples.
"""
import os
import socket
import subprocess
import sys
import time
import uuid
from contextlib import suppress
from pathlib import Path
from typing import Any
import pytest
import requests
# Configuration
TIMEOUT = 30 # seconds
ORCHESTRATION_TIMEOUT = 180 # seconds for orchestrations
_DEFAULT_HOST = "localhost"
class FunctionAppStartupError(RuntimeError):
"""Raised when the Azure Functions host fails to start reliably."""
pass
def _load_env_file_if_present() -> None:
"""Load environment variables from the local .env file when available."""
env_file = Path(__file__).parent / ".env"
if not env_file.exists():
return
try:
from dotenv import load_dotenv
load_dotenv(env_file)
except ImportError:
# python-dotenv not available; rely on existing environment
pass
def _should_skip_azure_functions_integration_tests() -> tuple[bool, str]:
"""Determine whether Azure Functions integration tests should be skipped."""
_load_env_file_if_present()
run_integration_tests = os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
if not run_integration_tests:
return (
True,
"Integration tests are disabled. Set RUN_INTEGRATION_TESTS=true to enable Azure Functions sample tests.",
)
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "").strip()
if not endpoint or endpoint == "https://your-resource.openai.azure.com/":
return True, "No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
deployment_name = os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", "").strip()
if not deployment_name or deployment_name == "your-deployment-name":
return True, "No real AZURE_OPENAI_CHAT_DEPLOYMENT_NAME provided; skipping integration tests."
return False, "Integration tests enabled."
_SKIP_AZURE_FUNCTIONS_INTEGRATION_TESTS, _AZURE_FUNCTIONS_SKIP_REASON = _should_skip_azure_functions_integration_tests()
skip_if_azure_functions_integration_tests_disabled = pytest.mark.skipif(
_SKIP_AZURE_FUNCTIONS_INTEGRATION_TESTS,
reason=_AZURE_FUNCTIONS_SKIP_REASON,
)
class SampleTestHelper:
"""Helper class for testing samples."""
@staticmethod
def post_json(url: str, data: dict[str, Any], timeout: int = TIMEOUT) -> requests.Response:
"""POST JSON data to a URL."""
return requests.post(url, json=data, headers={"Content-Type": "application/json"}, timeout=timeout)
@staticmethod
def post_text(url: str, text: str, timeout: int = TIMEOUT) -> requests.Response:
"""POST plain text to a URL."""
return requests.post(url, data=text, headers={"Content-Type": "text/plain"}, timeout=timeout)
@staticmethod
def get(url: str, timeout: int = TIMEOUT) -> requests.Response:
"""GET request to a URL."""
return requests.get(url, timeout=timeout)
@staticmethod
def wait_for_orchestration(
status_url: str, max_wait: int = ORCHESTRATION_TIMEOUT, poll_interval: int = 2
) -> dict[str, Any]:
"""
Wait for an orchestration to complete.
Args:
status_url: URL to poll for orchestration status
max_wait: Maximum seconds to wait
poll_interval: Seconds between polls
Returns:
Final orchestration status
Raises:
TimeoutError: If orchestration doesn't complete in time
"""
start_time = time.time()
while time.time() - start_time < max_wait:
response = requests.get(status_url, timeout=TIMEOUT)
response.raise_for_status()
status = response.json()
runtime_status = status.get("runtimeStatus", "")
if runtime_status in ["Completed", "Failed", "Terminated"]:
return status
time.sleep(poll_interval)
raise TimeoutError(f"Orchestration did not complete within {max_wait} seconds")
@staticmethod
def wait_for_orchestration_with_output(
status_url: str, max_wait: int = ORCHESTRATION_TIMEOUT, poll_interval: int = 2
) -> dict[str, Any]:
"""
Wait for an orchestration to complete and have output available.
This is a specialized version of wait_for_orchestration that also
ensures the output field is present, handling timing race conditions.
Args:
status_url: URL to poll for orchestration status
max_wait: Maximum seconds to wait
poll_interval: Seconds between polls
Returns:
Final orchestration status with output
Raises:
TimeoutError: If orchestration doesn't complete with output in time
"""
start_time = time.time()
while time.time() - start_time < max_wait:
response = requests.get(status_url, timeout=TIMEOUT)
response.raise_for_status()
status = response.json()
runtime_status = status.get("runtimeStatus", "")
if runtime_status in ["Failed", "Terminated"]:
return status
if runtime_status == "Completed" and status.get("output"):
return status
# If completed but no output, continue polling for a bit more to
# handle the race condition where output has not been persisted yet.
time.sleep(poll_interval)
# Provide detailed error message based on final status
final_response = requests.get(status_url, timeout=TIMEOUT)
final_response.raise_for_status()
final_status = final_response.json()
final_runtime_status = final_status.get("runtimeStatus", "Unknown")
if final_runtime_status == "Completed":
if "output" not in final_status:
raise TimeoutError(
"Orchestration completed but 'output' field is missing after "
f"{max_wait} seconds. Final status: {final_status}"
)
if not final_status["output"]:
raise TimeoutError(
"Orchestration completed but output is empty after "
f"{max_wait} seconds. Final status: {final_status}"
)
raise TimeoutError(
"Orchestration completed with output but validation failed after "
f"{max_wait} seconds. Final status: {final_status}"
)
raise TimeoutError(
"Orchestration did not complete within "
f"{max_wait} seconds. Final status: {final_runtime_status}, "
f"Full status: {final_status}"
)
# Function App Lifecycle Management Helpers
def _resolve_repo_root() -> Path:
"""Resolve the repository root, preferring GITHUB_WORKSPACE when available."""
workspace = os.getenv("GITHUB_WORKSPACE")
if workspace:
candidate = Path(workspace).expanduser()
if not (candidate / "samples").exists() and (candidate / "python" / "samples").exists():
return (candidate / "python").resolve()
return candidate.resolve()
# If `GITHUB_WORKSPACE` is not set,
# go up from testutils.py -> integration_tests -> tests -> azurefunctions -> packages -> python
return Path(__file__).resolve().parents[4]
def get_sample_path_from_marker(request) -> tuple[Path | None, str | None]:
"""
Get sample path from @pytest.mark.sample() marker.
Returns a tuple of (sample_path, error_message).
If successful, error_message is None.
If failed, sample_path is None and error_message contains the reason.
"""
marker = request.node.get_closest_marker("sample")
if not marker:
return (
None,
(
"No @pytest.mark.sample() marker found on test. Add pytestmark with "
"@pytest.mark.sample('sample_name') to the test module."
),
)
if not marker.args:
return (
None,
"@pytest.mark.sample() marker found but no sample name provided. Use @pytest.mark.sample('sample_name').",
)
sample_name = marker.args[0]
repo_root = _resolve_repo_root()
sample_path = repo_root / "samples" / "getting_started" / "azure_functions" / sample_name
if not sample_path.exists():
return None, f"Sample directory does not exist: {sample_path}"
return sample_path, None
def find_available_port(host: str = _DEFAULT_HOST) -> int:
"""Find an available TCP port on the given host."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind((host, 0))
return sock.getsockname()[1]
def build_base_url(port: int, host: str = _DEFAULT_HOST) -> str:
"""Construct a base URL for the Azure Functions host."""
return f"http://{host}:{port}"
def is_port_in_use(port: int, host: str = _DEFAULT_HOST) -> bool:
"""
Check if a port is already in use.
Returns True if the port is in use, False otherwise.
"""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
return sock.connect_ex((host, port)) == 0
def load_and_validate_env() -> None:
"""
Load .env file from current directory if it exists,
then validate that required environment variables are present.
Raises pytest.fail if required environment variables are missing.
"""
_load_env_file_if_present()
# Required environment variables for Azure Functions samples
# These match the variables defined in .env.example
required_env_vars = [
"AZURE_OPENAI_ENDPOINT",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME",
"AzureWebJobsStorage",
"DURABLE_TASK_SCHEDULER_CONNECTION_STRING",
"FUNCTIONS_WORKER_RUNTIME",
]
# Check if required env vars are set
missing_vars = [var for var in required_env_vars if not os.environ.get(var)]
if missing_vars:
pytest.fail(
f"Missing required environment variables: {', '.join(missing_vars)}. "
"Please create a .env file in tests/integration_tests/ based on .env.example or "
"set these variables in your environment."
)
def start_function_app(sample_path: Path, port: int) -> subprocess.Popen:
"""
Start a function app in the specified sample directory.
Returns the subprocess.Popen object for the running process.
"""
env = os.environ.copy()
# Use a unique TASKHUB_NAME for each test run to ensure test isolation.
# This prevents conflicts between parallel or repeated test runs, as Durable Functions
# use the task hub name to separate orchestration state.
env["TASKHUB_NAME"] = f"test{uuid.uuid4().hex[:8]}"
# On Windows, use CREATE_NEW_PROCESS_GROUP to allow proper termination
# shell=True only on Windows to handle PATH resolution
if sys.platform == "win32":
return subprocess.Popen(
["func", "start", "--port", str(port)],
cwd=str(sample_path),
creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
shell=True,
env=env,
)
# On Unix, don't use shell=True to avoid shell wrapper issues
return subprocess.Popen(["func", "start", "--port", str(port)], cwd=str(sample_path), env=env)
def wait_for_function_app_ready(func_process: subprocess.Popen, port: int, max_wait: int = 60) -> None:
"""Block until the Azure Functions host responds healthy or fail fast."""
start_time = time.time()
health_url = f"{build_base_url(port)}/api/health"
last_error: Exception | None = None
while time.time() - start_time < max_wait:
# If the process exited early, capture any previously seen error and fail fast.
if func_process.poll() is not None:
raise FunctionAppStartupError(
f"Function app process exited with code {func_process.returncode} before becoming healthy"
) from last_error
if is_port_in_use(port):
try:
response = requests.get(health_url, timeout=5)
if response.status_code == 200:
return
last_error = RuntimeError(f"Health check returned {response.status_code}")
except requests.RequestException as exc:
last_error = exc
time.sleep(1)
raise FunctionAppStartupError(
f"Function app did not become healthy on port {port} within {max_wait} seconds"
) from last_error
def cleanup_function_app(func_process: subprocess.Popen) -> None:
"""
Clean up the function app process and all its children.
Uses psutil if available for more thorough cleanup, falls back to basic termination.
"""
try:
import psutil
if func_process.poll() is None: # Process still running
# Get parent process
parent = psutil.Process(func_process.pid)
# Get all child processes recursively
children = parent.children(recursive=True)
# Kill children first
for child in children:
with suppress(psutil.NoSuchProcess, psutil.AccessDenied):
child.kill()
# Kill parent
with suppress(psutil.NoSuchProcess, psutil.AccessDenied):
parent.kill()
# Wait for all to terminate
_gone, alive = psutil.wait_procs(children + [parent], timeout=3)
# Force kill any remaining
for proc in alive:
with suppress(psutil.NoSuchProcess, psutil.AccessDenied):
proc.kill()
except ImportError:
# Fallback if psutil not available
try:
if func_process.poll() is None:
func_process.kill()
func_process.wait()
except Exception:
# Ignore all exceptions during fallback cleanup; best effort to terminate process.
pass
except Exception:
pass # Best effort cleanup
# Give the port time to be released
time.sleep(2)
python/packages/azurefunctions/tests/test_app.py
+15 -5
View File
@@ -355,7 +355,9 @@ class TestAgentEntityOperations:
async def test_entity_run_agent_operation(self) -> None:
"""Test that entity can run agent operation."""
mock_agent = Mock()
mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[ChatMessage("assistant", ["Test response"])]))
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Test response")])
)
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="test-conv-123"))
@@ -371,7 +373,9 @@ class TestAgentEntityOperations:
async def test_entity_stores_conversation_history(self) -> None:
"""Test that the entity stores conversation history."""
mock_agent = Mock()
mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[ChatMessage("assistant", ["Response 1"])]))
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Response 1")])
)
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
@@ -403,7 +407,9 @@ class TestAgentEntityOperations:
async def test_entity_increments_message_count(self) -> None:
"""Test that the entity increments the message count."""
mock_agent = Mock()
mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[ChatMessage("assistant", ["Response"])]))
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Response")])
)
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
@@ -442,7 +448,9 @@ class TestAgentEntityFactory:
def test_entity_function_handles_run_operation(self) -> None:
"""Test that the entity function handles the run operation."""
mock_agent = Mock()
mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[ChatMessage("assistant", ["Response"])]))
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Response")])
)
entity_function = create_agent_entity(mock_agent)
@@ -467,7 +475,9 @@ class TestAgentEntityFactory:
def test_entity_function_handles_run_agent_operation(self) -> None:
"""Test that the entity function handles the deprecated run_agent operation for backward compatibility."""
mock_agent = Mock()
mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[ChatMessage("assistant", ["Response"])]))
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Response")])
)
entity_function = create_agent_entity(mock_agent)
python/packages/azurefunctions/tests/test_entities.py
+1 -1
View File
@@ -19,7 +19,7 @@ TFunc = TypeVar("TFunc", bound=Callable[..., Any])
def _agent_response(text: str | None) -> AgentResponse:
"""Create an AgentResponse with a single assistant message."""
message = ChatMessage("assistant", [text]) if text is not None else ChatMessage("assistant", [])
message = ChatMessage(role="assistant", text=text) if text is not None else ChatMessage(role="assistant", text="")
return AgentResponse(messages=[message])
python/packages/azurefunctions/tests/test_orchestration.py
+2 -2
View File
@@ -136,7 +136,7 @@ class TestAgentResponseHelpers:
# Simulate successful entity task completion
entity_task.state = TaskState.SUCCEEDED
entity_task.result = AgentResponse(messages=[ChatMessage("assistant", ["Test response"])]).to_dict()
entity_task.result = AgentResponse(messages=[ChatMessage(role="assistant", text="Test response")]).to_dict()
# Clear pending_tasks to simulate that parent has processed the child
task.pending_tasks.clear()
@@ -178,7 +178,7 @@ class TestAgentResponseHelpers:
# Simulate successful entity task with JSON response
entity_task.state = TaskState.SUCCEEDED
entity_task.result = AgentResponse(messages=[ChatMessage("assistant", ['{"answer": "42"}'])]).to_dict()
entity_task.result = AgentResponse(messages=[ChatMessage(role="assistant", text='{"answer": "42"}')]).to_dict()
# Clear pending_tasks to simulate that parent has processed the child
task.pending_tasks.clear()
python/packages/bedrock/agent_framework_bedrock/_chat_client.py
+60 -41
View File
@@ -4,30 +4,34 @@ import asyncio
import json
import sys
from collections import deque
from collections.abc import AsyncIterable, MutableMapping, MutableSequence, Sequence
from typing import Any, ClassVar, Generic, Literal
from collections.abc import AsyncIterable, Awaitable, Mapping, MutableMapping, Sequence
from typing import Any, ClassVar, Generic, Literal, TypedDict
from uuid import uuid4
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
ChatMessage,
ChatMiddlewareLayer,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
FinishReasonLiteral,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
ResponseStream,
ToolProtocol,
UsageDetails,
get_logger,
prepare_function_call_results,
use_chat_middleware,
use_function_invocation,
validate_tool_mode,
)
from agent_framework._pydantic import AFBaseSettings
from agent_framework.exceptions import ServiceInitializationError, ServiceInvalidResponseError
from agent_framework.observability import use_instrumentation
from agent_framework.observability import ChatTelemetryLayer
from boto3.session import Session as Boto3Session
from botocore.client import BaseClient
from botocore.config import Config as BotoConfig
@@ -190,7 +194,7 @@ ROLE_MAP: dict[str, str] = {
"tool": "user",
}
FINISH_REASON_MAP: dict[str, str] = {
FINISH_REASON_MAP: dict[str, FinishReasonLiteral] = {
"end_turn": "stop",
"stop_sequence": "stop",
"max_tokens": "length",
@@ -212,11 +216,14 @@ class BedrockSettings(AFBaseSettings):
session_token: SecretStr | None = None
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class BedrockChatClient(BaseChatClient[TBedrockChatOptions], Generic[TBedrockChatOptions]):
"""Async chat client for Amazon Bedrock's Converse API."""
class BedrockChatClient(
ChatMiddlewareLayer[TBedrockChatOptions],
FunctionInvocationLayer[TBedrockChatOptions],
ChatTelemetryLayer[TBedrockChatOptions],
BaseChatClient[TBedrockChatOptions],
Generic[TBedrockChatOptions],
):
"""Async chat client for Amazon Bedrock's Converse API with middleware, telemetry, and function invocation."""
OTEL_PROVIDER_NAME: ClassVar[str] = "aws.bedrock" # type: ignore[reportIncompatibleVariableOverride, misc]
@@ -230,6 +237,8 @@ class BedrockChatClient(BaseChatClient[TBedrockChatOptions], Generic[TBedrockCha
session_token: str | None = None,
client: BaseClient | None = None,
boto3_session: Boto3Session | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
@@ -244,6 +253,8 @@ class BedrockChatClient(BaseChatClient[TBedrockChatOptions], Generic[TBedrockCha
session_token: Optional AWS session token for temporary credentials.
client: Preconfigured Bedrock runtime client; when omitted a boto3 session is created.
boto3_session: Custom boto3 session used to build the runtime client if provided.
middleware: Optional sequence of middlewares to include.
function_invocation_configuration: Optional function invocation configuration
env_file_path: Optional .env file path used by ``BedrockSettings`` to load defaults.
env_file_encoding: Encoding for the optional .env file.
kwargs: Additional arguments forwarded to ``BaseChatClient``.
@@ -289,7 +300,11 @@ class BedrockChatClient(BaseChatClient[TBedrockChatOptions], Generic[TBedrockCha
config=BotoConfig(user_agent_extra=AGENT_FRAMEWORK_USER_AGENT),
)
super().__init__(**kwargs)
super().__init__(
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
self._bedrock_client = client
self.model_id = settings.chat_model_id
self.region = settings.region
@@ -305,41 +320,45 @@ class BedrockChatClient(BaseChatClient[TBedrockChatOptions], Generic[TBedrockCha
return Boto3Session(**session_kwargs)
@override
async def _inner_get_response(
def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
) -> ChatResponse:
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
request = self._prepare_options(messages, options, **kwargs)
raw_response = await asyncio.to_thread(self._bedrock_client.converse, **request)
return self._process_converse_response(raw_response)
@override
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
response = await self._inner_get_response(messages=messages, options=options, **kwargs)
contents = list(response.messages[0].contents if response.messages else [])
if response.usage_details:
contents.append(Content.from_usage(usage_details=response.usage_details)) # type: ignore[arg-type]
yield ChatResponseUpdate(
response_id=response.response_id,
contents=contents,
model_id=response.model_id,
finish_reason=response.finish_reason,
raw_representation=response.raw_representation,
)
if stream:
# Streaming mode - simulate streaming by yielding a single update
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
response = await asyncio.to_thread(self._bedrock_client.converse, **request)
parsed_response = self._process_converse_response(response)
contents = list(parsed_response.messages[0].contents if parsed_response.messages else [])
if parsed_response.usage_details:
contents.append(Content.from_usage(usage_details=parsed_response.usage_details)) # type: ignore[arg-type]
yield ChatResponseUpdate(
response_id=parsed_response.response_id,
contents=contents,
model_id=parsed_response.model_id,
finish_reason=parsed_response.finish_reason,
raw_representation=parsed_response.raw_representation,
)
return self._build_response_stream(_stream())
# Non-streaming mode
async def _get_response() -> ChatResponse:
raw_response = await asyncio.to_thread(self._bedrock_client.converse, **request)
return self._process_converse_response(raw_response)
return _get_response()
def _prepare_options(
self,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
**kwargs: Any,
) -> dict[str, Any]:
model_id = options.get("model_id") or self.model_id
@@ -572,7 +591,7 @@ class BedrockChatClient(BaseChatClient[TBedrockChatOptions], Generic[TBedrockCha
message = output.get("message", {})
content_blocks = message.get("content", []) or []
contents = self._parse_message_contents(content_blocks)
chat_message = ChatMessage("assistant", contents, raw_representation=message)
chat_message = ChatMessage(role="assistant", contents=contents, raw_representation=message)
usage_details = self._parse_usage(response.get("usage") or output.get("usage"))
finish_reason = self._map_finish_reason(output.get("completionReason") or response.get("stopReason"))
response_id = response.get("responseId") or message.get("id")
@@ -640,7 +659,7 @@ class BedrockChatClient(BaseChatClient[TBedrockChatOptions], Generic[TBedrockCha
logger.debug("Ignoring unsupported Bedrock content block: %s", block)
return contents
def _map_finish_reason(self, reason: str | None) -> str | None:
def _map_finish_reason(self, reason: str | None) -> FinishReasonLiteral | None:
if not reason:
return None
return FINISH_REASON_MAP.get(reason.lower())
python/packages/bedrock/tests/test_bedrock_client.py
+5 -6
View File
@@ -2,7 +2,6 @@
from __future__ import annotations
import asyncio
from typing import Any
import pytest
@@ -33,7 +32,7 @@ class _StubBedrockRuntime:
}
def test_get_response_invokes_bedrock_runtime() -> None:
async def test_get_response_invokes_bedrock_runtime() -> None:
stub = _StubBedrockRuntime()
client = BedrockChatClient(
model_id="amazon.titan-text",
@@ -42,11 +41,11 @@ def test_get_response_invokes_bedrock_runtime() -> None:
)
messages = [
ChatMessage("system", [Content.from_text(text="You are concise.")]),
ChatMessage("user", [Content.from_text(text="hello")]),
ChatMessage(role="system", contents=[Content.from_text(text="You are concise.")]),
ChatMessage(role="user", contents=[Content.from_text(text="hello")]),
]
response = asyncio.run(client.get_response(messages=messages, options={"max_tokens": 32}))
response = await client.get_response(messages=messages, options={"max_tokens": 32})
assert stub.calls, "Expected the runtime client to be called"
payload = stub.calls[0]
@@ -63,7 +62,7 @@ def test_build_request_requires_non_system_messages() -> None:
client=_StubBedrockRuntime(),
)
messages = [ChatMessage("system", [Content.from_text(text="Only system text")])]
messages = [ChatMessage(role="system", contents=[Content.from_text(text="Only system text")])]
with pytest.raises(ServiceInitializationError):
client._prepare_options(messages, {})
python/packages/bedrock/tests/test_bedrock_settings.py
+2 -2
View File
@@ -46,7 +46,7 @@ def test_build_request_includes_tool_config() -> None:
"tools": [tool],
"tool_choice": {"mode": "required", "required_function_name": "get_weather"},
}
messages = [ChatMessage("user", [Content.from_text(text="hi")])]
messages = [ChatMessage(role="user", contents=[Content.from_text(text="hi")])]
request = client._prepare_options(messages, options)
@@ -58,7 +58,7 @@ def test_build_request_serializes_tool_history() -> None:
client = _build_client()
options: ChatOptions = {}
messages = [
ChatMessage("user", [Content.from_text(text="how's weather?")]),
ChatMessage(role="user", contents=[Content.from_text(text="how's weather?")]),
ChatMessage(
role="assistant",
contents=[
python/packages/chatkit/README.md
+1 -1
View File
@@ -104,7 +104,7 @@ class MyChatKitServer(ChatKitServer[dict[str, Any]]):
agent_messages = await simple_to_agent_input(thread_items_page.data)
# Run the agent and stream responses
response_stream = agent.run_stream(agent_messages)
response_stream = agent.run(agent_messages, stream=True)
# Convert agent responses back to ChatKit events
async for event in stream_agent_response(response_stream, thread.id):
python/packages/chatkit/agent_framework_chatkit/_converter.py
+10 -10
View File
@@ -100,21 +100,21 @@ class ThreadItemConverter:
# If only text and no attachments, use text parameter for simplicity
if text_content.strip() and not data_contents:
user_message = ChatMessage("user", [text_content.strip()])
user_message = ChatMessage(role="user", text=text_content.strip())
else:
# Build contents list with both text and attachments
contents: list[Content] = []
if text_content.strip():
contents.append(Content.from_text(text=text_content.strip()))
contents.extend(data_contents)
user_message = ChatMessage("user", contents)
user_message = ChatMessage(role="user", contents=contents)
# Handle quoted text if this is the last message
messages = [user_message]
if item.quoted_text and is_last_message:
quoted_context = ChatMessage(
"user",
[f"The user is referring to this in particular:\n{item.quoted_text}"],
role="user",
text=f"The user is referring to this in particular:\n{item.quoted_text}",
)
# Prepend quoted context before the main message
messages.insert(0, quoted_context)
@@ -213,7 +213,7 @@ class ThreadItemConverter:
message = converter.hidden_context_to_input(hidden_item)
# Returns: ChatMessage(role=SYSTEM, text="<HIDDEN_CONTEXT>User's email: ...</HIDDEN_CONTEXT>")
"""
return ChatMessage("system", [f"<HIDDEN_CONTEXT>{item.content}</HIDDEN_CONTEXT>"])
return ChatMessage(role="system", text=f"<HIDDEN_CONTEXT>{item.content}</HIDDEN_CONTEXT>")
def tag_to_message_content(self, tag: UserMessageTagContent) -> Content:
"""Convert a ChatKit tag (@-mention) to Agent Framework content.
@@ -292,7 +292,7 @@ class ThreadItemConverter:
f"A message was displayed to the user that the following task was performed:\n<Task>\n{task_text}\n</Task>"
)
return ChatMessage("user", [text])
return ChatMessage(role="user", text=text)
def workflow_to_input(self, item: WorkflowItem) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit WorkflowItem to Agent Framework ChatMessage(s).
@@ -347,7 +347,7 @@ class ThreadItemConverter:
f"<Task>\n{task_text}\n</Task>"
)
messages.append(ChatMessage("user", [text]))
messages.append(ChatMessage(role="user", text=text))
return messages if messages else None
@@ -389,7 +389,7 @@ class ThreadItemConverter:
try:
widget_json = item.widget.model_dump_json(exclude_unset=True, exclude_none=True)
text = f"The following graphical UI widget (id: {item.id}) was displayed to the user:{widget_json}"
return ChatMessage("user", [text])
return ChatMessage(role="user", text=text)
except Exception:
# If JSON serialization fails, skip the widget
return None
@@ -415,7 +415,7 @@ class ThreadItemConverter:
if not text_parts:
return None
return ChatMessage("assistant", ["".join(text_parts)])
return ChatMessage(role="assistant", text="".join(text_parts))
async def client_tool_call_to_input(self, item: ClientToolCallItem) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit ClientToolCallItem to Agent Framework ChatMessage(s).
@@ -563,7 +563,7 @@ class ThreadItemConverter:
from agent_framework import ChatAgent
agent = ChatAgent(...)
response = await agent.run_stream(messages)
response = await agent.run(messages)
"""
thread_items = list(thread_items) if isinstance(thread_items, Sequence) else [thread_items]
python/packages/claude/agent_framework_claude/_agent.py
+60 -34
View File
@@ -2,9 +2,9 @@
import contextlib
import sys
from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
from collections.abc import AsyncIterable, Awaitable, Callable, MutableMapping, Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Generic
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, overload
from agent_framework import (
AgentMiddlewareTypes,
@@ -175,7 +175,7 @@ class ClaudeAgent(BaseAgent, Generic[TOptions]):
.. code-block:: python
async with ClaudeAgent() as agent:
async for update in agent.run_stream("Write a poem"):
async for update in agent.run("Write a poem"):
print(update.text, end="", flush=True)
With session management:
@@ -552,7 +552,59 @@ class ClaudeAgent(BaseAgent, Generic[TOptions]):
return ""
return "\n".join([msg.text or "" for msg in messages])
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
options: TOptions | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]: ...
@overload
async def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
options: TOptions | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AgentResponse[Any]: ...
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
options: TOptions | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate] | Awaitable[AgentResponse[Any]]:
"""Run the agent with the given messages.
Args:
messages: The messages to process.
Keyword Args:
stream: If True, returns an async iterable of updates. If False (default),
returns an awaitable AgentResponse.
thread: The conversation thread. If thread has service_thread_id set,
the agent will resume that session.
options: Runtime options (model, permission_mode can be changed per-request).
kwargs: Additional keyword arguments.
Returns:
When stream=True: An AsyncIterable[AgentResponseUpdate] for streaming updates.
When stream=False: An Awaitable[AgentResponse] with the complete response.
"""
if stream:
return self._run_streaming(messages, thread=thread, options=options, **kwargs)
return self._run_non_streaming(messages, thread=thread, options=options, **kwargs)
async def _run_non_streaming(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
@@ -560,26 +612,13 @@ class ClaudeAgent(BaseAgent, Generic[TOptions]):
options: TOptions | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AgentResponse[Any]:
"""Run the agent with the given messages.
Args:
messages: The messages to process.
Keyword Args:
thread: The conversation thread. If thread has service_thread_id set,
the agent will resume that session.
options: Runtime options (model, permission_mode can be changed per-request).
kwargs: Additional keyword arguments.
Returns:
AgentResponse with the agent's response.
"""
"""Internal non-streaming implementation."""
thread = thread or self.get_new_thread()
return await AgentResponse.from_agent_response_generator(
self.run_stream(messages, thread=thread, options=options, **kwargs)
return await AgentResponse.from_update_generator(
self._run_streaming(messages, thread=thread, options=options, **kwargs)
)
async def run_stream(
async def _run_streaming(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
@@ -587,20 +626,7 @@ class ClaudeAgent(BaseAgent, Generic[TOptions]):
options: TOptions | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
"""Stream the agent's response.
Args:
messages: The messages to process.
Keyword Args:
thread: The conversation thread. If thread has service_thread_id set,
the agent will resume that session.
options: Runtime options (model, permission_mode can be changed per-request).
kwargs: Additional keyword arguments.
Yields:
AgentResponseUpdate objects containing chunks of the response.
"""
"""Internal streaming implementation."""
thread = thread or self.get_new_thread()
# Ensure we're connected to the right session
python/packages/claude/tests/__init__.py
-1
View File
@@ -1 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
python/packages/claude/tests/test_claude_agent.py
+11 -11
View File
@@ -312,7 +312,7 @@ class TestClaudeAgentRun:
class TestClaudeAgentRunStream:
"""Tests for ClaudeAgent run_stream method."""
"""Tests for ClaudeAgent streaming run method."""
@staticmethod
async def _create_async_generator(items: list[Any]) -> Any:
@@ -332,7 +332,7 @@ class TestClaudeAgentRunStream:
return mock_client
async def test_run_stream_yields_updates(self) -> None:
"""Test run_stream yields AgentResponseUpdate objects."""
"""Test run(stream=True) yields AgentResponseUpdate objects."""
from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
from claude_agent_sdk.types import StreamEvent
@@ -371,16 +371,16 @@ class TestClaudeAgentRunStream:
with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client):
agent = ClaudeAgent()
updates: list[AgentResponseUpdate] = []
async for update in agent.run_stream("Hello"):
async for update in agent.run("Hello", stream=True):
updates.append(update)
# StreamEvent yields text deltas
# StreamEvent yields text deltas (2 events)
assert len(updates) == 2
assert updates[0].role == "assistant"
assert updates[0].text == "Streaming "
assert updates[1].text == "response"
async def test_run_stream_raises_on_assistant_message_error(self) -> None:
"""Test run_stream raises ServiceException when AssistantMessage has an error."""
"""Test run raises ServiceException when AssistantMessage has an error."""
from agent_framework.exceptions import ServiceException
from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
@@ -404,13 +404,13 @@ class TestClaudeAgentRunStream:
with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client):
agent = ClaudeAgent()
with pytest.raises(ServiceException) as exc_info:
async for _ in agent.run_stream("Hello"):
async for _ in agent.run("Hello", stream=True):
pass
assert "Invalid request to Claude API" in str(exc_info.value)
assert "Error details from API" in str(exc_info.value)
async def test_run_stream_raises_on_result_message_error(self) -> None:
"""Test run_stream raises ServiceException when ResultMessage.is_error is True."""
"""Test run raises ServiceException when ResultMessage.is_error is True."""
from agent_framework.exceptions import ServiceException
from claude_agent_sdk import ResultMessage
@@ -430,7 +430,7 @@ class TestClaudeAgentRunStream:
with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client):
agent = ClaudeAgent()
with pytest.raises(ServiceException) as exc_info:
async for _ in agent.run_stream("Hello"):
async for _ in agent.run("Hello", stream=True):
pass
assert "Model 'claude-sonnet-4.5' not found" in str(exc_info.value)
@@ -697,9 +697,9 @@ class TestFormatPrompt:
"""Test formatting multiple messages."""
agent = ClaudeAgent()
messages = [
ChatMessage("user", [Content.from_text(text="Hi")]),
ChatMessage("assistant", [Content.from_text(text="Hello!")]),
ChatMessage("user", [Content.from_text(text="How are you?")]),
ChatMessage(role="user", contents=[Content.from_text(text="Hi")]),
ChatMessage(role="assistant", contents=[Content.from_text(text="Hello!")]),
ChatMessage(role="user", contents=[Content.from_text(text="How are you?")]),
]
result = agent._format_prompt(messages) # type: ignore[reportPrivateUsage]
assert "Hi" in result
python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py
+78 -56
View File
@@ -1,7 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import AsyncIterable
from typing import Any, ClassVar
from collections.abc import AsyncIterable, Awaitable, Sequence
from typing import Any, ClassVar, Literal, overload
from agent_framework import (
AgentMiddlewareTypes,
@@ -12,6 +12,7 @@ from agent_framework import (
ChatMessage,
Content,
ContextProvider,
ResponseStream,
normalize_messages,
)
from agent_framework._pydantic import AFBaseSettings
@@ -204,35 +205,64 @@ class CopilotStudioAgent(BaseAgent):
self.token_cache = token_cache
self.scopes = scopes
async def run(
@overload
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
stream: Literal[False] = False,
thread: AgentThread | None = None,
**kwargs: Any,
) -> "Awaitable[AgentResponse]": ...
@overload
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ...
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
**kwargs: Any,
) -> "Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]":
"""Get a response from the agent.
This method returns the final result of the agent's execution
as a single AgentResponse object. When stream=True, it returns
a ResponseStream that yields AgentResponseUpdate objects.
Args:
messages: The message(s) to send to the agent.
Keyword Args:
stream: Whether to stream the response. Defaults to False.
thread: The conversation thread associated with the message(s).
kwargs: Additional keyword arguments.
Returns:
When stream=False: An Awaitable[AgentResponse].
When stream=True: A ResponseStream of AgentResponseUpdate items.
"""
if stream:
return self._run_stream_impl(messages=messages, thread=thread, **kwargs)
return self._run_impl(messages=messages, thread=thread, **kwargs)
async def _run_impl(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
"""Get a response from the agent.
This method returns the final result of the agent's execution
as a single AgentResponse object. The caller is blocked until
the final result is available.
Note: For streaming responses, use the run_stream method, which returns
intermediate steps and the final result as a stream of AgentResponseUpdate
objects. Streaming only the final result is not feasible because the timing of
the final result's availability is unknown, and blocking the caller until then
is undesirable in streaming scenarios.
Args:
messages: The message(s) to send to the agent.
Keyword Args:
thread: The conversation thread associated with the message(s).
kwargs: Additional keyword arguments.
Returns:
An agent response item.
"""
"""Non-streaming implementation of run."""
if not thread:
thread = self.get_new_thread()
thread.service_thread_id = await self._start_new_conversation()
@@ -250,49 +280,41 @@ class CopilotStudioAgent(BaseAgent):
return AgentResponse(messages=response_messages, response_id=response_id)
async def run_stream(
def _run_stream_impl(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
"""Run the agent as a stream.
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
"""Streaming implementation of run."""
This method will return the intermediate steps and final results of the
agent's execution as a stream of AgentResponseUpdate objects to the caller.
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
nonlocal thread
if not thread:
thread = self.get_new_thread()
thread.service_thread_id = await self._start_new_conversation()
Note: An AgentResponseUpdate object contains a chunk of a message.
input_messages = normalize_messages(messages)
Args:
messages: The message(s) to send to the agent.
question = "\n".join([message.text for message in input_messages])
Keyword Args:
thread: The conversation thread associated with the message(s).
kwargs: Additional keyword arguments.
activities = self.client.ask_question(question, thread.service_thread_id)
Yields:
An agent response item.
"""
if not thread:
thread = self.get_new_thread()
thread.service_thread_id = await self._start_new_conversation()
async for message in self._process_activities(activities, streaming=True):
yield AgentResponseUpdate(
role=message.role,
contents=message.contents,
author_name=message.author_name,
raw_representation=message.raw_representation,
response_id=message.message_id,
message_id=message.message_id,
)
input_messages = normalize_messages(messages)
def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse[None]:
return AgentResponse.from_updates(updates)
question = "\n".join([message.text for message in input_messages])
activities = self.client.ask_question(question, thread.service_thread_id)
async for message in self._process_activities(activities, streaming=True):
yield AgentResponseUpdate(
role=message.role,
contents=message.contents,
author_name=message.author_name,
raw_representation=message.raw_representation,
response_id=message.message_id,
message_id=message.message_id,
)
return ResponseStream(_stream(), finalizer=_finalize)
async def _start_new_conversation(self) -> str:
"""Start a new conversation with the Copilot Studio agent.
python/packages/copilotstudio/tests/test_copilot_agent.py
+13 -13
View File
@@ -143,7 +143,7 @@ class TestCopilotStudioAgent:
mock_copilot_client.start_conversation.return_value = create_async_generator([conversation_activity])
mock_copilot_client.ask_question.return_value = create_async_generator([mock_activity])
chat_message = ChatMessage("user", [Content.from_text("test message")])
chat_message = ChatMessage(role="user", contents=[Content.from_text("test message")])
response = await agent.run(chat_message)
assert isinstance(response, AgentResponse)
@@ -179,8 +179,8 @@ class TestCopilotStudioAgent:
with pytest.raises(ServiceException, match="Failed to start a new conversation"):
await agent.run("test message")
async def test_run_stream_with_string_message(self, mock_copilot_client: MagicMock) -> None:
"""Test run_stream method with string message."""
async def test_run_streaming_with_string_message(self, mock_copilot_client: MagicMock) -> None:
"""Test run(stream=True) method with string message."""
agent = CopilotStudioAgent(client=mock_copilot_client)
conversation_activity = MagicMock()
@@ -196,7 +196,7 @@ class TestCopilotStudioAgent:
mock_copilot_client.ask_question.return_value = create_async_generator([typing_activity])
response_count = 0
async for response in agent.run_stream("test message"):
async for response in agent.run("test message", stream=True):
assert isinstance(response, AgentResponseUpdate)
content = response.contents[0]
assert content.type == "text"
@@ -205,8 +205,8 @@ class TestCopilotStudioAgent:
assert response_count == 1
async def test_run_stream_with_thread(self, mock_copilot_client: MagicMock) -> None:
"""Test run_stream method with existing thread."""
async def test_run_streaming_with_thread(self, mock_copilot_client: MagicMock) -> None:
"""Test run(stream=True) method with existing thread."""
agent = CopilotStudioAgent(client=mock_copilot_client)
thread = AgentThread()
@@ -223,7 +223,7 @@ class TestCopilotStudioAgent:
mock_copilot_client.ask_question.return_value = create_async_generator([typing_activity])
response_count = 0
async for response in agent.run_stream("test message", thread=thread):
async for response in agent.run("test message", thread=thread, stream=True):
assert isinstance(response, AgentResponseUpdate)
content = response.contents[0]
assert content.type == "text"
@@ -233,8 +233,8 @@ class TestCopilotStudioAgent:
assert response_count == 1
assert thread.service_thread_id == "test-conversation-id"
async def test_run_stream_no_typing_activity(self, mock_copilot_client: MagicMock) -> None:
"""Test run_stream method with non-typing activity."""
async def test_run_streaming_no_typing_activity(self, mock_copilot_client: MagicMock) -> None:
"""Test run(stream=True) method with non-typing activity."""
agent = CopilotStudioAgent(client=mock_copilot_client)
conversation_activity = MagicMock()
@@ -249,7 +249,7 @@ class TestCopilotStudioAgent:
mock_copilot_client.ask_question.return_value = create_async_generator([message_activity])
response_count = 0
async for _response in agent.run_stream("test message"):
async for _response in agent.run("test message", stream=True):
response_count += 1
assert response_count == 0
@@ -297,12 +297,12 @@ class TestCopilotStudioAgent:
assert isinstance(response, AgentResponse)
assert len(response.messages) == 1
async def test_run_stream_start_conversation_failure(self, mock_copilot_client: MagicMock) -> None:
"""Test run_stream method when conversation start fails."""
async def test_run_streaming_start_conversation_failure(self, mock_copilot_client: MagicMock) -> None:
"""Test run(stream=True) method when conversation start fails."""
agent = CopilotStudioAgent(client=mock_copilot_client)
mock_copilot_client.start_conversation.return_value = create_async_generator([])
with pytest.raises(ServiceException, match="Failed to start a new conversation"):
async for _ in agent.run_stream("test message"):
async for _ in agent.run("test message", stream=True):
pass
python/packages/core/agent_framework/_agents.py
+349 -257
View File
@@ -3,15 +3,17 @@
import inspect
import re
import sys
from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
from contextlib import AbstractAsyncContextManager, AsyncExitStack
from copy import deepcopy
from functools import partial
from itertools import chain
from typing import (
TYPE_CHECKING,
Any,
ClassVar,
Generic,
Literal,
Protocol,
cast,
overload,
@@ -28,21 +30,26 @@ from ._clients import BaseChatClient, ChatClientProtocol
from ._logging import get_logger
from ._mcp import LOG_LEVEL_MAPPING, MCPTool
from ._memory import Context, ContextProvider
from ._middleware import Middleware, use_agent_middleware
from ._middleware import AgentMiddlewareLayer, MiddlewareTypes
from ._serialization import SerializationMixin
from ._threads import AgentThread, ChatMessageStoreProtocol
from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, FunctionTool, ToolProtocol
from ._tools import (
FunctionInvocationLayer,
FunctionTool,
ToolProtocol,
)
from ._types import (
AgentResponse,
AgentResponseUpdate,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Content,
ResponseStream,
map_chat_to_agent_update,
normalize_messages,
)
from .exceptions import AgentExecutionException, AgentInitializationError
from .observability import use_agent_instrumentation
from .observability import AgentTelemetryLayer
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
@@ -71,7 +78,7 @@ TThreadType = TypeVar("TThreadType", bound="AgentThread")
TOptions_co = TypeVar(
"TOptions_co",
bound=TypedDict, # type: ignore[valid-type]
default="ChatOptions",
default="ChatOptions[None]",
covariant=True,
)
@@ -146,7 +153,17 @@ def _sanitize_agent_name(agent_name: str | None) -> str | None:
return sanitized
__all__ = ["AgentProtocol", "BaseAgent", "ChatAgent"]
class _RunContext(TypedDict):
thread: AgentThread
input_messages: list[ChatMessage]
thread_messages: list[ChatMessage]
agent_name: str
chat_options: dict[str, Any]
filtered_kwargs: dict[str, Any]
finalize_kwargs: dict[str, Any]
__all__ = ["AgentProtocol", "BareAgent", "BaseAgent", "ChatAgent", "RawChatAgent"]
# region Agent Protocol
@@ -179,20 +196,20 @@ class AgentProtocol(Protocol):
self.name = "Custom Agent"
self.description = "A fully custom agent implementation"
async def run(self, messages=None, *, thread=None, **kwargs):
# Your custom implementation
from agent_framework import AgentResponse
async def run(self, messages=None, *, stream=False, thread=None, **kwargs):
if stream:
# Your custom streaming implementation
async def _stream():
from agent_framework import AgentResponseUpdate
return AgentResponse(messages=[], response_id="custom-response")
yield AgentResponseUpdate()
def run_stream(self, messages=None, *, thread=None, **kwargs):
# Your custom streaming implementation
async def _stream():
from agent_framework import AgentResponseUpdate
return _stream()
else:
# Your custom implementation
from agent_framework import AgentResponse
yield AgentResponseUpdate()
return _stream()
return AgentResponse(messages=[], response_id="custom-response")
def get_new_thread(self, **kwargs):
# Return your own thread implementation
@@ -208,60 +225,56 @@ class AgentProtocol(Protocol):
name: str | None
description: str | None
async def run(
@overload
def run(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
) -> Awaitable[AgentResponse[Any]]:
"""Get a response from the agent (non-streaming)."""
...
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Get a streaming response from the agent."""
...
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Get a response from the agent.
This method returns the final result of the agent's execution
as a single AgentResponse object. The caller is blocked until
the final result is available.
Note: For streaming responses, use the run_stream method, which returns
intermediate steps and the final result as a stream of AgentResponseUpdate
objects. Streaming only the final result is not feasible because the timing of
the final result's availability is unknown, and blocking the caller until then
is undesirable in streaming scenarios.
This method can return either a complete response or stream partial updates
depending on the stream parameter. Streaming returns a ResponseStream that
can be iterated for updates and finalized for the full response.
Args:
messages: The message(s) to send to the agent.
Keyword Args:
stream: Whether to stream the response. Defaults to False.
thread: The conversation thread associated with the message(s).
kwargs: Additional keyword arguments.
Returns:
An agent response item.
"""
...
def run_stream(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
"""Run the agent as a stream.
This method will return the intermediate steps and final results of the
agent's execution as a stream of AgentResponseUpdate objects to the caller.
Note: An AgentResponseUpdate object contains a chunk of a message.
Args:
messages: The message(s) to send to the agent.
Keyword Args:
thread: The conversation thread associated with the message(s).
kwargs: Additional keyword arguments.
Yields:
An agent response item.
When stream=False: An AgentResponse with the final result.
When stream=True: A ResponseStream of AgentResponseUpdate items with
``get_final_response()`` for the final AgentResponse.
"""
...
@@ -276,12 +289,15 @@ class AgentProtocol(Protocol):
class BaseAgent(SerializationMixin):
"""Base class for all Agent Framework agents.
This is the minimal base class without middleware or telemetry layers.
For most use cases, prefer :class:`ChatAgent` which includes all standard layers.
This class provides core functionality for agent implementations, including
context providers, middleware support, and thread management.
Note:
BaseAgent cannot be instantiated directly as it doesn't implement the
``run()``, ``run_stream()``, and other methods required by AgentProtocol.
``run()`` and other methods required by AgentProtocol.
Use a concrete implementation like ChatAgent or create a subclass.
Examples:
@@ -292,16 +308,17 @@ class BaseAgent(SerializationMixin):
# Create a concrete subclass that implements the protocol
class SimpleAgent(BaseAgent):
async def run(self, messages=None, *, thread=None, **kwargs):
# Custom implementation
return AgentResponse(messages=[], response_id="simple-response")
async def run(self, messages=None, *, stream=False, thread=None, **kwargs):
if stream:
def run_stream(self, messages=None, *, thread=None, **kwargs):
async def _stream():
# Custom streaming implementation
yield AgentResponseUpdate()
async def _stream():
# Custom streaming implementation
yield AgentResponseUpdate()
return _stream()
return _stream()
else:
# Custom implementation
return AgentResponse(messages=[], response_id="simple-response")
# Now instantiate the concrete subclass
@@ -328,7 +345,7 @@ class BaseAgent(SerializationMixin):
name: str | None = None,
description: str | None = None,
context_provider: ContextProvider | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
additional_properties: MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> None:
@@ -350,8 +367,8 @@ class BaseAgent(SerializationMixin):
self.name = name
self.description = description
self.context_provider = context_provider
self.middleware: list[Middleware] | None = (
cast(list[Middleware], middleware) if middleware is not None else None
self.middleware: list[MiddlewareTypes] | None = (
cast(list[MiddlewareTypes], middleware) if middleware is not None else None
)
# Merge kwargs into additional_properties
@@ -428,7 +445,7 @@ class BaseAgent(SerializationMixin):
arg_name: The name of the function argument (default: "task").
arg_description: The description for the function argument.
If None, defaults to "Task for {tool_name}".
stream_callback: Optional callback for streaming responses. If provided, uses run_stream.
stream_callback: Optional callback for streaming responses. If provided, uses run(..., stream=True).
Returns:
A FunctionTool that can be used as a tool by other agents.
@@ -475,15 +492,15 @@ class BaseAgent(SerializationMixin):
input_text = kwargs.get(arg_name, "")
# Forward runtime context kwargs, excluding arg_name and conversation_id.
forwarded_kwargs = {k: v for k, v in kwargs.items() if k not in (arg_name, "conversation_id")}
forwarded_kwargs = {k: v for k, v in kwargs.items() if k not in (arg_name, "conversation_id", "options")}
if stream_callback is None:
# Use non-streaming mode
return (await self.run(input_text, **forwarded_kwargs)).text
return (await self.run(input_text, stream=False, **forwarded_kwargs)).text
# Use streaming mode - accumulate updates and create final response
response_updates: list[AgentResponseUpdate] = []
async for update in self.run_stream(input_text, **forwarded_kwargs):
async for update in self.run(input_text, stream=True, **forwarded_kwargs):
response_updates.append(update)
if is_async_callback:
await stream_callback(update) # type: ignore[misc]
@@ -504,13 +521,18 @@ class BaseAgent(SerializationMixin):
return agent_tool
# Backward compatibility alias
BareAgent = BaseAgent
# region ChatAgent
@use_agent_middleware
@use_agent_instrumentation(capture_usage=False) # type: ignore[arg-type,misc]
class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
"""A Chat Client Agent.
class RawChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
"""A Chat Client Agent without middleware or telemetry layers.
This is the core chat agent implementation. For most use cases,
prefer :class:`ChatAgent` which includes all standard layers.
This is the primary agent implementation that uses a chat client to interact
with language models. It supports tools, context providers, middleware, and
@@ -554,8 +576,10 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
)
# Use streaming responses
async for update in agent.run_stream("What's the weather in Paris?"):
stream = agent.run("What's the weather in Paris?", stream=True)
async for update in stream:
print(update.text, end="")
final = await stream.get_final_response()
With typed options for IDE autocomplete:
@@ -601,7 +625,6 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
default_options: TOptions_co | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
middleware: Sequence[Middleware] | None = None,
**kwargs: Any,
) -> None:
"""Initialize a ChatAgent instance.
@@ -625,7 +648,7 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
tool_choice, and provider-specific options like reasoning_effort.
You can also create your own TypedDict for custom chat clients.
Note: response_format typing does not flow into run outputs when set via default_options.
These can be overridden at runtime via the ``options`` parameter of ``run()`` and ``run_stream()``.
These can be overridden at runtime via the ``options`` parameter of ``run()``.
tools: The tools to use for the request.
kwargs: Any additional keyword arguments. Will be stored as ``additional_properties``.
@@ -642,7 +665,7 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
"Use conversation_id for service-managed threads or chat_message_store_factory for local storage."
)
if not hasattr(chat_client, FUNCTION_INVOKING_CHAT_CLIENT_MARKER) and isinstance(chat_client, BaseChatClient):
if not isinstance(chat_client, FunctionInvocationLayer) and isinstance(chat_client, BaseChatClient):
logger.warning(
"The provided chat client does not support function invoking, this might limit agent capabilities."
)
@@ -652,10 +675,9 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
name=name,
description=description,
context_provider=context_provider,
middleware=middleware,
**kwargs,
)
self.chat_client: ChatClientProtocol[TOptions_co] = chat_client
self.chat_client = chat_client
self.chat_message_store_factory = chat_message_store_factory
# Get tools from options or named parameter (named param takes precedence)
@@ -754,10 +776,11 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
self.chat_client._update_agent_name_and_description(self.name, self.description) # type: ignore[reportAttributeAccessIssue, attr-defined]
@overload
async def run(
def run(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
tools: ToolProtocol
| Callable[..., Any]
@@ -766,36 +789,54 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
| None = None,
options: "ChatOptions[TResponseModelT]",
**kwargs: Any,
) -> AgentResponse[TResponseModelT]: ...
) -> Awaitable[AgentResponse[TResponseModelT]]: ...
@overload
async def run(
def run(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
options: TOptions_co | Mapping[str, Any] | "ChatOptions[Any]" | None = None,
options: "TOptions_co | ChatOptions[None] | None" = None,
**kwargs: Any,
) -> AgentResponse[Any]: ...
) -> Awaitable[AgentResponse[Any]]: ...
async def run(
@overload
def run(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
options: TOptions_co | Mapping[str, Any] | "ChatOptions[Any]" | None = None,
options: "TOptions_co | ChatOptions[Any] | None" = None,
**kwargs: Any,
) -> AgentResponse[Any]:
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
options: "TOptions_co | ChatOptions[Any] | None" = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Run the agent with the given messages and options.
Note:
@@ -806,6 +847,7 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
Args:
messages: The messages to process.
stream: Whether to stream the response. Defaults to False.
Keyword Args:
thread: The thread to use for the agent.
@@ -818,34 +860,154 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
Will only be passed to functions that are called.
Returns:
An AgentResponse containing the agent's response.
When stream=False: An Awaitable[AgentResponse] containing the agent's response.
When stream=True: A ResponseStream of AgentResponseUpdate items with
``get_final_response()`` for the final AgentResponse.
"""
# Build options dict from provided options
if not stream:
async def _run_non_streaming() -> AgentResponse[Any]:
ctx = await self._prepare_run_context(
messages=messages,
thread=thread,
tools=tools,
options=options,
kwargs=kwargs,
)
response = await self.chat_client.get_response( # type: ignore[call-overload]
messages=ctx["thread_messages"],
stream=False,
options=ctx["chat_options"],
**ctx["filtered_kwargs"],
)
if not response:
raise AgentExecutionException("Chat client did not return a response.")
await self._finalize_response_and_update_thread(
response=response,
agent_name=ctx["agent_name"],
thread=ctx["thread"],
input_messages=ctx["input_messages"],
kwargs=ctx["finalize_kwargs"],
)
response_format = ctx["chat_options"].get("response_format")
if not (
response_format is not None
and isinstance(response_format, type)
and issubclass(response_format, BaseModel)
):
response_format = None
return AgentResponse(
messages=response.messages,
response_id=response.response_id,
created_at=response.created_at,
usage_details=response.usage_details,
value=response.value,
response_format=response_format,
raw_representation=response,
additional_properties=response.additional_properties,
)
return _run_non_streaming()
# Use a holder to capture the context created during stream initialization
ctx_holder: dict[str, _RunContext | None] = {"ctx": None}
async def _post_hook(response: AgentResponse) -> None:
ctx = ctx_holder["ctx"]
if ctx is None:
return # No context available (shouldn't happen in normal flow)
# Update thread with conversation_id
await self._update_thread_with_type_and_conversation_id(ctx["thread"], response.response_id)
# Ensure author names are set for all messages
for message in response.messages:
if message.author_name is None:
message.author_name = ctx["agent_name"]
# Notify thread of new messages
await self._notify_thread_of_new_messages(
ctx["thread"],
ctx["input_messages"],
response.messages,
**{k: v for k, v in ctx["finalize_kwargs"].items() if k != "thread"},
)
async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]:
ctx_holder["ctx"] = await self._prepare_run_context(
messages=messages,
thread=thread,
tools=tools,
options=options,
kwargs=kwargs,
)
ctx: _RunContext = ctx_holder["ctx"] # type: ignore[assignment] # Safe: we just assigned it
return self.chat_client.get_response( # type: ignore[call-overload, no-any-return]
messages=ctx["thread_messages"],
stream=True,
options=ctx["chat_options"],
**ctx["filtered_kwargs"],
)
return (
ResponseStream
.from_awaitable(_get_stream())
.map(
transform=partial(
map_chat_to_agent_update,
agent_name=self.name,
),
finalizer=partial(
self._finalize_response_updates, response_format=options.get("response_format") if options else None
),
)
.with_result_hook(_post_hook)
)
def _finalize_response_updates(
self,
updates: Sequence[AgentResponseUpdate],
*,
response_format: Any | None = None,
) -> AgentResponse:
"""Finalize response updates into a single AgentResponse."""
output_format_type = response_format if isinstance(response_format, type) else None
return AgentResponse.from_updates(updates, output_format_type=output_format_type)
async def _prepare_run_context(
self,
*,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None,
thread: AgentThread | None,
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None,
options: Mapping[str, Any] | None,
kwargs: dict[str, Any],
) -> _RunContext:
opts = dict(options) if options else {}
# Get tools from options or named parameter (named param takes precedence)
tools_ = tools if tools is not None else opts.pop("tools", None)
tools_ = cast(
ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None,
tools_,
)
input_messages = normalize_messages(messages)
thread, run_chat_options, thread_messages = await self._prepare_thread_and_messages(
thread=thread, input_messages=input_messages, **kwargs
)
normalized_tools: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = ( # type:ignore[reportUnknownVariableType]
# Normalize tools
normalized_tools: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = (
[] if tools_ is None else tools_ if isinstance(tools_, list) else [tools_]
)
agent_name = self._get_agent_name()
# Resolve final tool list (runtime provided tools + local MCP server tools)
final_tools: list[ToolProtocol | Callable[..., Any] | dict[str, Any]] = []
# Normalize tools argument to a list without mutating the original parameter
for tool in normalized_tools:
if isinstance(tool, MCPTool):
if not tool.is_connected:
@@ -864,6 +1026,7 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
"model_id": opts.pop("model_id", None),
"conversation_id": thread.service_thread_id,
"allow_multiple_tool_calls": opts.pop("allow_multiple_tool_calls", None),
"additional_function_arguments": opts.pop("additional_function_arguments", None),
"frequency_penalty": opts.pop("frequency_penalty", None),
"logit_bias": opts.pop("logit_bias", None),
"max_tokens": opts.pop("max_tokens", None),
@@ -885,15 +1048,38 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
co = _merge_options(run_chat_options, run_opts)
# Ensure thread is forwarded in kwargs for tool invocation
kwargs["thread"] = thread
finalize_kwargs = dict(kwargs)
finalize_kwargs["thread"] = thread
# Filter chat_options from kwargs to prevent duplicate keyword argument
filtered_kwargs = {k: v for k, v in kwargs.items() if k != "chat_options"}
response = await self.chat_client.get_response(
messages=thread_messages,
options=co, # type: ignore[arg-type]
**filtered_kwargs,
)
filtered_kwargs = {k: v for k, v in finalize_kwargs.items() if k != "chat_options"}
return {
"thread": thread,
"input_messages": input_messages,
"thread_messages": thread_messages,
"agent_name": agent_name,
"chat_options": co,
"filtered_kwargs": filtered_kwargs,
"finalize_kwargs": finalize_kwargs,
}
async def _finalize_response_and_update_thread(
self,
response: ChatResponse,
agent_name: str,
thread: AgentThread,
input_messages: list[ChatMessage],
kwargs: dict[str, Any],
) -> None:
"""Finalize response by updating thread and setting author names.
Args:
response: The chat response to finalize.
agent_name: The name of the agent to set as author.
thread: The conversation thread.
input_messages: The input messages.
kwargs: Additional keyword arguments.
"""
await self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
# Ensure that the author name is set for each message in the response.
@@ -909,150 +1095,6 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
response.messages,
**{k: v for k, v in kwargs.items() if k != "thread"},
)
response_format = co.get("response_format")
if not (
response_format is not None and isinstance(response_format, type) and issubclass(response_format, BaseModel)
):
response_format = None
return AgentResponse(
messages=response.messages,
response_id=response.response_id,
created_at=response.created_at,
usage_details=response.usage_details,
value=response.value,
response_format=response_format,
raw_representation=response,
additional_properties=response.additional_properties,
)
async def run_stream(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
options: TOptions_co | Mapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
"""Stream the agent with the given messages and options.
Note:
Since you won't always call ``agent.run_stream()`` directly (it gets called
through orchestration), it is advised to set your default values for
all the chat client parameters in the agent constructor.
If both parameters are used, the ones passed to the run methods take precedence.
Args:
messages: The messages to process.
Keyword Args:
thread: The thread to use for the agent.
tools: The tools to use for this specific run (merged with agent-level tools).
options: A TypedDict containing chat options. When using a typed agent like
``ChatAgent[OpenAIChatOptions]``, this enables IDE autocomplete for
provider-specific options including temperature, max_tokens, model_id,
tool_choice, and provider-specific options like reasoning_effort.
kwargs: Additional keyword arguments for the agent.
Will only be passed to functions that are called.
Yields:
AgentResponseUpdate objects containing chunks of the agent's response.
"""
# Build options dict from provided options
opts = dict(options) if options else {}
# Get tools from options or named parameter (named param takes precedence)
tools_ = tools if tools is not None else opts.pop("tools", None)
input_messages = normalize_messages(messages)
thread, run_chat_options, thread_messages = await self._prepare_thread_and_messages(
thread=thread, input_messages=input_messages, **kwargs
)
agent_name = self._get_agent_name()
# Resolve final tool list (runtime provided tools + local MCP server tools)
final_tools: list[ToolProtocol | MutableMapping[str, Any] | Callable[..., Any]] = []
normalized_tools: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = ( # type: ignore[reportUnknownVariableType]
[] if tools_ is None else tools_ if isinstance(tools_, list) else [tools_]
)
# Normalize tools argument to a list without mutating the original parameter
for tool in normalized_tools:
if isinstance(tool, MCPTool):
if not tool.is_connected:
await self._async_exit_stack.enter_async_context(tool)
final_tools.extend(tool.functions) # type: ignore
else:
final_tools.append(tool)
for mcp_server in self.mcp_tools:
if not mcp_server.is_connected:
await self._async_exit_stack.enter_async_context(mcp_server)
final_tools.extend(mcp_server.functions)
# Build options dict from run_stream() options merged with provided options
run_opts: dict[str, Any] = {
"model_id": opts.pop("model_id", None),
"conversation_id": thread.service_thread_id,
"allow_multiple_tool_calls": opts.pop("allow_multiple_tool_calls", None),
"frequency_penalty": opts.pop("frequency_penalty", None),
"logit_bias": opts.pop("logit_bias", None),
"max_tokens": opts.pop("max_tokens", None),
"metadata": opts.pop("metadata", None),
"presence_penalty": opts.pop("presence_penalty", None),
"response_format": opts.pop("response_format", None),
"seed": opts.pop("seed", None),
"stop": opts.pop("stop", None),
"store": opts.pop("store", None),
"temperature": opts.pop("temperature", None),
"tool_choice": opts.pop("tool_choice", None),
"tools": final_tools,
"top_p": opts.pop("top_p", None),
"user": opts.pop("user", None),
**opts, # Remaining options are provider-specific
}
# Remove None values and merge with chat_options
run_opts = {k: v for k, v in run_opts.items() if v is not None}
co = _merge_options(run_chat_options, run_opts)
# Ensure thread is forwarded in kwargs for tool invocation
kwargs["thread"] = thread
# Filter chat_options from kwargs to prevent duplicate keyword argument
filtered_kwargs = {k: v for k, v in kwargs.items() if k != "chat_options"}
response_updates: list[ChatResponseUpdate] = []
async for update in self.chat_client.get_streaming_response(
messages=thread_messages,
options=co, # type: ignore[arg-type]
**filtered_kwargs,
):
response_updates.append(update)
if update.author_name is None:
update.author_name = agent_name
yield AgentResponseUpdate(
contents=update.contents,
role=update.role,
author_name=update.author_name,
response_id=update.response_id,
message_id=update.message_id,
created_at=update.created_at,
additional_properties=update.additional_properties,
raw_representation=update,
)
response = ChatResponse.from_updates(response_updates, output_format_type=co.get("response_format"))
await self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
await self._notify_thread_of_new_messages(
thread,
input_messages,
response.messages,
**{k: v for k, v in kwargs.items() if k != "thread"},
)
@override
def get_new_thread(
@@ -1326,3 +1368,53 @@ class ChatAgent(BaseAgent, Generic[TOptions_co]): # type: ignore[misc]
The agent's name, or 'UnnamedAgent' if no name is set.
"""
return self.name or "UnnamedAgent"
class ChatAgent(
AgentTelemetryLayer,
AgentMiddlewareLayer,
RawChatAgent[TOptions_co],
Generic[TOptions_co],
):
"""A Chat Client Agent with middleware, telemetry, and full layer support.
This is the recommended agent class for most use cases. It includes:
- Agent middleware support for request/response interception
- OpenTelemetry-based telemetry for observability
For a minimal implementation without these features, use :class:`RawChatAgent`.
"""
def __init__(
self,
chat_client: ChatClientProtocol[TOptions_co],
instructions: str | None = None,
*,
id: str | None = None,
name: str | None = None,
description: str | None = None,
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: TOptions_co | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> None:
"""Initialize a ChatAgent instance."""
super().__init__(
chat_client=chat_client,
instructions=instructions,
id=id,
name=name,
description=description,
tools=tools,
default_options=default_options,
chat_message_store_factory=chat_message_store_factory,
context_provider=context_provider,
middleware=middleware,
**kwargs,
)
python/packages/core/agent_framework/_clients.py
+168 -150
View File
@@ -1,14 +1,13 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import sys
from abc import ABC, abstractmethod
from collections.abc import (
AsyncIterable,
Awaitable,
Callable,
Mapping,
MutableMapping,
MutableSequence,
Sequence,
)
from typing import (
@@ -16,6 +15,7 @@ from typing import (
Any,
ClassVar,
Generic,
Literal,
Protocol,
TypedDict,
cast,
@@ -27,17 +27,9 @@ from pydantic import BaseModel
from ._logging import get_logger
from ._memory import ContextProvider
from ._middleware import (
ChatMiddleware,
ChatMiddlewareCallable,
FunctionMiddleware,
FunctionMiddlewareCallable,
Middleware,
)
from ._serialization import SerializationMixin
from ._threads import ChatMessageStoreProtocol
from ._tools import (
FUNCTION_INVOKING_CHAT_CLIENT_MARKER,
FunctionInvocationConfiguration,
ToolProtocol,
)
@@ -45,7 +37,7 @@ from ._types import (
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Content,
ResponseStream,
prepare_messages,
validate_chat_options,
)
@@ -58,10 +50,14 @@ else:
if TYPE_CHECKING:
from ._agents import ChatAgent
from ._middleware import (
MiddlewareTypes,
)
from ._types import ChatOptions
TInput = TypeVar("TInput", contravariant=True)
TEmbedding = TypeVar("TEmbedding")
TBaseChatClient = TypeVar("TBaseChatClient", bound="BaseChatClient")
@@ -79,13 +75,16 @@ __all__ = [
TOptions_contra = TypeVar(
"TOptions_contra",
bound=TypedDict, # type: ignore[valid-type]
default="ChatOptions",
default="ChatOptions[None]",
contravariant=True,
)
# Used for the overloads that capture the response model type from options
TResponseModelT = TypeVar("TResponseModelT", bound=BaseModel)
@runtime_checkable
class ChatClientProtocol(Protocol[TOptions_contra]): #
class ChatClientProtocol(Protocol[TOptions_contra]):
"""A protocol for a chat client that can generate responses.
This protocol defines the interface that all chat clients must implement,
@@ -107,17 +106,22 @@ class ChatClientProtocol(Protocol[TOptions_contra]): #
# Any class implementing the required methods is compatible
class CustomChatClient:
async def get_response(self, messages, **kwargs):
# Your custom implementation
return ChatResponse(messages=[], response_id="custom")
additional_properties: dict = {}
def get_streaming_response(self, messages, **kwargs):
async def _stream():
from agent_framework import ChatResponseUpdate
def get_response(self, messages, *, stream=False, **kwargs):
if stream:
from agent_framework import ChatResponseUpdate, ResponseStream
yield ChatResponseUpdate()
async def _stream():
yield ChatResponseUpdate()
return _stream()
return ResponseStream(_stream())
else:
async def _response():
return ChatResponse(messages=[], response_id="custom")
return _response()
# Verify the instance satisfies the protocol
@@ -128,56 +132,60 @@ class ChatClientProtocol(Protocol[TOptions_contra]): #
additional_properties: dict[str, Any]
@overload
async def get_response(
def get_response(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage],
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: Literal[False] = ...,
options: "ChatOptions[TResponseModelT]",
**kwargs: Any,
) -> "ChatResponse[TResponseModelT]": ...
) -> Awaitable[ChatResponse[TResponseModelT]]: ...
@overload
async def get_response(
def get_response(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage],
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
options: TOptions_contra | None = None,
stream: Literal[False] = ...,
options: "TOptions_contra | ChatOptions[None] | None" = None,
**kwargs: Any,
) -> ChatResponse:
) -> Awaitable[ChatResponse[Any]]: ...
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: Literal[True],
options: "TOptions_contra | ChatOptions[Any] | None" = None,
**kwargs: Any,
) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: bool = False,
options: "TOptions_contra | ChatOptions[Any] | None" = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
"""Send input and return the response.
Args:
messages: The sequence of input messages to send.
stream: Whether to stream the response. Defaults to False.
options: Chat options as a TypedDict.
**kwargs: Additional chat options.
Returns:
The response messages generated by the client.
When stream=False: An awaitable ChatResponse from the client.
When stream=True: A ResponseStream yielding partial updates.
Raises:
ValueError: If the input message sequence is ``None``.
"""
...
def get_streaming_response(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage],
*,
options: TOptions_contra | None = None,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
"""Send input messages and stream the response.
Args:
messages: The sequence of input messages to send.
options: Chat options as a TypedDict.
**kwargs: Additional chat options.
Yields:
ChatResponseUpdate: Partial response updates as they're generated.
"""
...
# endregion
@@ -188,27 +196,30 @@ class ChatClientProtocol(Protocol[TOptions_contra]): #
TOptions_co = TypeVar(
"TOptions_co",
bound=TypedDict, # type: ignore[valid-type]
default="ChatOptions",
default="ChatOptions[None]",
covariant=True,
)
TResponseModel = TypeVar("TResponseModel", bound=BaseModel | None, default=None, covariant=True)
TResponseModelT = TypeVar("TResponseModelT", bound=BaseModel)
class BaseChatClient(SerializationMixin, ABC, Generic[TOptions_co]):
"""Base class for chat clients.
"""Abstract base class for chat clients without middleware wrapping.
This abstract base class provides core functionality for chat client implementations,
including middleware support, message preparation, and tool normalization.
including message preparation and tool normalization, but without middleware,
telemetry, or function invocation support.
The generic type parameter TOptions specifies which options TypedDict this client
accepts. This enables IDE autocomplete and type checking for provider-specific options
when using the typed overloads of get_response and get_streaming_response.
when using the typed overloads of get_response.
Note:
BaseChatClient cannot be instantiated directly as it's an abstract base class.
Subclasses must implement ``_inner_get_response()`` and ``_inner_get_streaming_response()``.
Subclasses must implement ``_inner_get_response()`` with a stream parameter to handle both
streaming and non-streaming responses.
For full-featured clients with middleware, telemetry, and function invocation support,
use the public client classes (e.g., ``OpenAIChatClient``, ``OpenAIResponsesClient``)
which compose these layers correctly.
Examples:
.. code-block:: python
@@ -218,15 +229,20 @@ class BaseChatClient(SerializationMixin, ABC, Generic[TOptions_co]):
class CustomChatClient(BaseChatClient):
async def _inner_get_response(self, *, messages, options, **kwargs):
# Your custom implementation
return ChatResponse(messages=[ChatMessage("assistant", ["Hello!"])], response_id="custom-response")
async def _inner_get_response(self, *, messages, stream, options, **kwargs):
if stream:
# Streaming implementation
from agent_framework import ChatResponseUpdate
async def _inner_get_streaming_response(self, *, messages, options, **kwargs):
# Your custom streaming implementation
from agent_framework import ChatResponseUpdate
async def _stream():
yield ChatResponseUpdate(role="assistant", contents=[{"type": "text", "text": "Hello!"}])
yield ChatResponseUpdate(role="assistant", contents=[{"type": "text", "text": "Hello!"}])
return _stream()
else:
# Non-streaming implementation
return ChatResponse(
messages=[ChatMessage(role="assistant", text="Hello!")], response_id="custom-response"
)
# Create an instance of your custom client
@@ -234,6 +250,9 @@ class BaseChatClient(SerializationMixin, ABC, Generic[TOptions_co]):
# Use the client to get responses
response = await client.get_response("Hello, how are you?")
# Or stream responses
async for update in client.get_response("Hello!", stream=True):
print(update)
"""
OTEL_PROVIDER_NAME: ClassVar[str] = "unknown"
@@ -243,28 +262,17 @@ class BaseChatClient(SerializationMixin, ABC, Generic[TOptions_co]):
def __init__(
self,
*,
middleware: (
Sequence[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable] | None
) = None,
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
"""Initialize a BaseChatClient instance.
Keyword Args:
middleware: Middleware for the client.
additional_properties: Additional properties for the client.
kwargs: Additional keyword arguments (merged into additional_properties).
"""
# Merge kwargs into additional_properties
self.additional_properties = additional_properties or {}
self.additional_properties.update(kwargs)
self.middleware = middleware
self.function_invocation_configuration = (
FunctionInvocationConfiguration() if hasattr(self.__class__, FUNCTION_INVOKING_CHAT_CLIENT_MARKER) else None
)
super().__init__(**kwargs)
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
"""Convert the instance to a dictionary.
@@ -287,121 +295,128 @@ class BaseChatClient(SerializationMixin, ABC, Generic[TOptions_co]):
return result
# region Internal methods to be implemented by the derived classes
async def _validate_options(self, options: Mapping[str, Any]) -> dict[str, Any]:
"""Validate and normalize chat options.
Subclasses should call this at the start of _inner_get_response to validate options.
Args:
options: The raw options dict.
Returns:
The validated and normalized options dict.
"""
return await validate_chat_options(dict(options))
def _finalize_response_updates(
self,
updates: Sequence[ChatResponseUpdate],
*,
response_format: Any | None = None,
) -> ChatResponse:
"""Finalize response updates into a single ChatResponse."""
output_format_type = response_format if isinstance(response_format, type) else None
return ChatResponse.from_updates(updates, output_format_type=output_format_type)
def _build_response_stream(
self,
stream: AsyncIterable[ChatResponseUpdate] | Awaitable[AsyncIterable[ChatResponseUpdate]],
*,
response_format: Any | None = None,
) -> ResponseStream[ChatResponseUpdate, ChatResponse]:
"""Create a ResponseStream with the standard finalizer."""
return ResponseStream(
stream,
finalizer=lambda updates: self._finalize_response_updates(updates, response_format=response_format),
)
# region Internal method to be implemented by derived classes
@abstractmethod
async def _inner_get_response(
def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
) -> ChatResponse:
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
"""Send a chat request to the AI service.
Subclasses must implement this method to handle both streaming and non-streaming
responses based on the stream parameter. Implementations should call
``await self._validate_options(options)`` at the start to validate options.
Keyword Args:
messages: The chat messages to send.
options: The options dict for the request.
messages: The prepared chat messages to send.
stream: Whether to stream the response.
options: The options dict for the request (call _validate_options first).
kwargs: Any additional keyword arguments.
Returns:
The chat response contents representing the response(s).
When stream=False: An Awaitable ChatResponse from the model.
When stream=True: A ResponseStream of ChatResponseUpdate instances.
"""
@abstractmethod
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
"""Send a streaming chat request to the AI service.
Keyword Args:
messages: The chat messages to send.
options: The options dict for the request.
kwargs: Any additional keyword arguments.
Yields:
ChatResponseUpdate: The streaming chat message contents.
"""
# Below is needed for mypy: https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
if False:
yield
await asyncio.sleep(0) # pragma: no cover
# This is a no-op, but it allows the method to be async and return an AsyncIterable.
# The actual implementation should yield ChatResponseUpdate instances as needed.
# endregion
# region Public method
@overload
async def get_response(
def get_response(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage],
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: Literal[False] = ...,
options: "ChatOptions[TResponseModelT]",
**kwargs: Any,
) -> ChatResponse[TResponseModelT]: ...
) -> Awaitable[ChatResponse[TResponseModelT]]: ...
@overload
async def get_response(
def get_response(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage],
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
options: TOptions_co | None = None,
stream: Literal[False] = ...,
options: "TOptions_co | ChatOptions[None] | None" = None,
**kwargs: Any,
) -> ChatResponse: ...
) -> Awaitable[ChatResponse[Any]]: ...
async def get_response(
@overload
def get_response(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage],
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
options: TOptions_co | "ChatOptions[Any]" | None = None,
stream: Literal[True],
options: "TOptions_co | ChatOptions[Any] | None" = None,
**kwargs: Any,
) -> ChatResponse[Any]:
) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
*,
stream: bool = False,
options: "TOptions_co | ChatOptions[Any] | None" = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
"""Get a response from a chat client.
Args:
messages: The message or messages to send to the model.
stream: Whether to stream the response. Defaults to False.
options: Chat options as a TypedDict.
**kwargs: Other keyword arguments, can be used to pass function specific parameters.
Returns:
A chat response from the model.
When streaming a response stream of ChatResponseUpdates, otherwise an Awaitable ChatResponse.
"""
return await self._inner_get_response(
messages=prepare_messages(messages),
options=await validate_chat_options(dict(options) if options else {}),
prepared_messages = prepare_messages(messages)
return self._inner_get_response(
messages=prepared_messages,
stream=stream,
options=options or {}, # type: ignore[arg-type]
**kwargs,
)
async def get_streaming_response(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage],
*,
options: TOptions_co | None = None,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
"""Get a streaming response from a chat client.
Args:
messages: The message or messages to send to the model.
options: Chat options as a TypedDict.
**kwargs: Other keyword arguments, can be used to pass function specific parameters.
Yields:
ChatResponseUpdate: A stream representing the response(s) from the LLM.
"""
async for update in self._inner_get_streaming_response(
messages=prepare_messages(messages),
options=await validate_chat_options(dict(options) if options else {}),
**kwargs,
):
yield update
def service_url(self) -> str:
"""Get the URL of the service.
@@ -428,7 +443,8 @@ class BaseChatClient(SerializationMixin, ABC, Generic[TOptions_co]):
default_options: TOptions_co | Mapping[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence["MiddlewareTypes"] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> "ChatAgent[TOptions_co]":
"""Create a ChatAgent with this client.
@@ -452,6 +468,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[TOptions_co]):
If not provided, the default in-memory store will be used.
context_provider: Context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
function_invocation_configuration: Optional function invocation configuration override.
kwargs: Any additional keyword arguments. Will be stored as ``additional_properties``.
Returns:
@@ -488,5 +505,6 @@ class BaseChatClient(SerializationMixin, ABC, Generic[TOptions_co]):
chat_message_store_factory=chat_message_store_factory,
context_provider=context_provider,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
python/packages/core/agent_framework/_middleware.py
+481 -743
View File
File diff suppressed because it is too large Load Diff
python/packages/core/agent_framework/_serialization.py
+6 -6
View File
@@ -38,7 +38,7 @@ class SerializationProtocol(Protocol):
# ChatMessage implements SerializationProtocol via SerializationMixin
user_msg = ChatMessage("user", ["What's the weather like today?"])
user_msg = ChatMessage(role="user", text="What's the weather like today?")
# Serialize to dictionary - automatic type identification and nested serialization
msg_dict = user_msg.to_dict()
@@ -175,8 +175,8 @@ class SerializationMixin:
# ChatMessageStoreState handles nested ChatMessage serialization
store_state = ChatMessageStoreState(
messages=[
ChatMessage("user", ["Hello agent"]),
ChatMessage("assistant", ["Hi! How can I help?"]),
ChatMessage(role="user", text="Hello agent"),
ChatMessage(role="assistant", text="Hi! How can I help?"),
]
)
@@ -473,7 +473,7 @@ class SerializationMixin:
weather_func = FunctionTool.from_dict(function_data, dependencies=dependencies)
# The function is now callable and ready for agent use
**Middleware Context Injection** - Agent execution context:
**MiddlewareTypes Context Injection** - Agent execution context:
.. code-block:: python
@@ -484,7 +484,7 @@ class SerializationMixin:
context_data = {
"type": "agent_run_context",
"messages": [{"role": "user", "text": "Hello"}],
"is_streaming": False,
"stream": False,
"metadata": {"session_id": "abc123"},
# agent and result are excluded from serialization
}
@@ -500,7 +500,7 @@ class SerializationMixin:
# Reconstruct context with agent dependency for middleware chain
context = AgentRunContext.from_dict(context_data, dependencies=dependencies)
# Middleware can now access context.agent and process the execution
# MiddlewareTypes can now access context.agent and process the execution
This injection system allows the agent framework to maintain clean separation
between serializable configuration and runtime dependencies like API clients,
python/packages/core/agent_framework/_threads.py
+1 -1
View File
@@ -202,7 +202,7 @@ class ChatMessageStore:
store = ChatMessageStore()
# Add messages
message = ChatMessage("user", ["Hello"])
message = ChatMessage(role="user", text="Hello")
await store.add_messages([message])
# Retrieve messages
python/packages/core/agent_framework/_tools.py
+653 -639
View File
File diff suppressed because it is too large Load Diff
python/packages/core/agent_framework/_types.py
+428 -77
View File
@@ -1,15 +1,12 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import base64
import json
import re
import sys
from collections.abc import (
AsyncIterable,
Callable,
Mapping,
MutableMapping,
Sequence,
)
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence
from copy import deepcopy
from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NewType, cast, overload
@@ -40,13 +37,19 @@ __all__ = [
"Content",
"FinishReason",
"FinishReasonLiteral",
"ResponseStream",
"Role",
"RoleLiteral",
"TFinal",
"TOuterFinal",
"TOuterUpdate",
"TUpdate",
"TextSpanRegion",
"ToolMode",
"UsageDetails",
"add_usage_details",
"detect_media_type_from_base64",
"map_chat_to_agent_update",
"merge_chat_options",
"normalize_messages",
"normalize_tools",
@@ -63,7 +66,7 @@ logger = get_logger("agent_framework")
# region Content Parsing Utilities
def _parse_content_list(contents_data: Sequence[Any]) -> list["Content"]:
def _parse_content_list(contents_data: Sequence[Any]) -> list[Content]:
"""Parse a list of content data into appropriate Content objects.
Args:
@@ -72,7 +75,7 @@ def _parse_content_list(contents_data: Sequence[Any]) -> list["Content"]:
Returns:
List of Content objects with unknown types logged and ignored
"""
contents: list["Content"] = []
contents: list[Content] = []
for content_data in contents_data:
if content_data is None:
continue
@@ -184,7 +187,7 @@ def detect_media_type_from_base64(
return None
def _get_data_bytes_as_str(content: "Content") -> str | None:
def _get_data_bytes_as_str(content: Content) -> str | None:
"""Extract base64 data string from data URI.
Args:
@@ -213,7 +216,7 @@ def _get_data_bytes_as_str(content: "Content") -> str | None:
return data # type: ignore[return-value, no-any-return]
def _get_data_bytes(content: "Content") -> bytes | None:
def _get_data_bytes(content: Content) -> bytes | None:
"""Extract and decode binary data from data URI.
Args:
@@ -484,8 +487,8 @@ class Content:
file_id: str | None = None,
vector_store_id: str | None = None,
# Code interpreter tool fields
inputs: list["Content"] | None = None,
outputs: list["Content"] | Any | None = None,
inputs: list[Content] | None = None,
outputs: list[Content] | Any | None = None,
# Image generation tool fields
image_id: str | None = None,
# MCP server tool fields
@@ -494,7 +497,7 @@ class Content:
output: Any = None,
# Function approval fields
id: str | None = None,
function_call: "Content | None" = None,
function_call: Content | None = None,
user_input_request: bool | None = None,
approved: bool | None = None,
# Common fields
@@ -845,7 +848,7 @@ class Content:
cls: type[TContent],
*,
call_id: str | None = None,
inputs: Sequence["Content"] | None = None,
inputs: Sequence[Content] | None = None,
annotations: Sequence[Annotation] | None = None,
additional_properties: MutableMapping[str, Any] | None = None,
raw_representation: Any = None,
@@ -865,7 +868,7 @@ class Content:
cls: type[TContent],
*,
call_id: str | None = None,
outputs: Sequence["Content"] | None = None,
outputs: Sequence[Content] | None = None,
annotations: Sequence[Annotation] | None = None,
additional_properties: MutableMapping[str, Any] | None = None,
raw_representation: Any = None,
@@ -966,7 +969,7 @@ class Content:
def from_function_approval_request(
cls: type[TContent],
id: str,
function_call: "Content",
function_call: Content,
*,
annotations: Sequence[Annotation] | None = None,
additional_properties: MutableMapping[str, Any] | None = None,
@@ -988,7 +991,7 @@ class Content:
cls: type[TContent],
approved: bool,
id: str,
function_call: "Content",
function_call: Content,
*,
annotations: Sequence[Annotation] | None = None,
additional_properties: MutableMapping[str, Any] | None = None,
@@ -1008,7 +1011,7 @@ class Content:
def to_function_approval_response(
self,
approved: bool,
) -> "Content":
) -> Content:
"""Convert a function approval request content to a function approval response content."""
if self.type != "function_approval_request":
raise ContentError(
@@ -1125,7 +1128,7 @@ class Content:
**remaining,
)
def __add__(self, other: "Content") -> "Content":
def __add__(self, other: Content) -> Content:
"""Concatenate or merge two Content instances."""
if not isinstance(other, Content):
raise TypeError(f"Incompatible type: Cannot add Content with {type(other).__name__}")
@@ -1143,7 +1146,7 @@ class Content:
return self._add_usage_content(other)
raise ContentError(f"Addition not supported for content type: {self.type}")
def _add_text_content(self, other: "Content") -> "Content":
def _add_text_content(self, other: Content) -> Content:
"""Add two TextContent instances."""
# Merge raw representations
if self.raw_representation is None:
@@ -1174,7 +1177,7 @@ class Content:
raw_representation=raw_representation,
)
def _add_text_reasoning_content(self, other: "Content") -> "Content":
def _add_text_reasoning_content(self, other: Content) -> Content:
"""Add two TextReasoningContent instances."""
# Merge raw representations
if self.raw_representation is None:
@@ -1214,7 +1217,7 @@ class Content:
raw_representation=raw_representation,
)
def _add_function_call_content(self, other: "Content") -> "Content":
def _add_function_call_content(self, other: Content) -> Content:
"""Add two FunctionCallContent instances."""
other_call_id = getattr(other, "call_id", None)
self_call_id = getattr(self, "call_id", None)
@@ -1258,7 +1261,7 @@ class Content:
raw_representation=raw_representation,
)
def _add_usage_content(self, other: "Content") -> "Content":
def _add_usage_content(self, other: Content) -> Content:
"""Add two UsageContent instances by combining their usage details."""
self_details = getattr(self, "usage_details", {})
other_details = getattr(other, "usage_details", {})
@@ -1372,7 +1375,7 @@ class Content:
# endregion
def _prepare_function_call_results_as_dumpable(content: "Content | Any | list[Content | Any]") -> Any:
def _prepare_function_call_results_as_dumpable(content: Content | Any | list[Content | Any]) -> Any:
if isinstance(content, list):
# Particularly deal with lists of Content
return [_prepare_function_call_results_as_dumpable(item) for item in content]
@@ -1388,7 +1391,7 @@ def _prepare_function_call_results_as_dumpable(content: "Content | Any | list[Co
return content
def prepare_function_call_results(content: "Content | Any | list[Content | Any]") -> str:
def prepare_function_call_results(content: Content | Any | list[Content | Any]) -> str:
"""Prepare the values of the function call results."""
if isinstance(content, Content):
# For BaseContent objects, use to_dict and serialize to JSON
@@ -1510,7 +1513,7 @@ class ChatMessage(SerializationMixin):
def __init__(
self,
role: RoleLiteral | str,
contents: "Sequence[Content | str | Mapping[str, Any]] | None" = None,
contents: Sequence[Content | str | Mapping[str, Any]] | None = None,
*,
text: str | None = None,
author_name: str | None = None,
@@ -1684,9 +1687,7 @@ def prepend_instructions_to_messages(
# region ChatResponse
def _process_update(
response: "ChatResponse | AgentResponse", update: "ChatResponseUpdate | AgentResponseUpdate"
) -> None:
def _process_update(response: ChatResponse | AgentResponse, update: ChatResponseUpdate | AgentResponseUpdate) -> None:
"""Processes a single update and modifies the response in place."""
is_new_message = False
if (
@@ -1760,11 +1761,11 @@ def _process_update(
response.model_id = update.model_id
def _coalesce_text_content(contents: list["Content"], type_str: Literal["text", "text_reasoning"]) -> None:
def _coalesce_text_content(contents: list[Content], type_str: Literal["text", "text_reasoning"]) -> None:
"""Take any subsequence Text or TextReasoningContent items and coalesce them into a single item."""
if not contents:
return
coalesced_contents: list["Content"] = []
coalesced_contents: list[Content] = []
first_new_content: Any | None = None
for content in contents:
if content.type == type_str:
@@ -1787,7 +1788,7 @@ def _coalesce_text_content(contents: list["Content"], type_str: Literal["text",
contents.extend(coalesced_contents)
def _finalize_response(response: "ChatResponse | AgentResponse") -> None:
def _finalize_response(response: ChatResponse | AgentResponse) -> None:
"""Finalizes the response by performing any necessary post-processing."""
for msg in response.messages:
_coalesce_text_content(msg.contents, "text")
@@ -1855,7 +1856,7 @@ class ChatResponse(SerializationMixin, Generic[TResponseModel]):
conversation_id: str | None = None,
model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: FinishReasonLiteral | str | None = None,
finish_reason: FinishReasonLiteral | FinishReason | None = None,
usage_details: UsageDetails | None = None,
value: TResponseModel | None = None,
response_format: type[BaseModel] | None = None,
@@ -1896,7 +1897,10 @@ class ChatResponse(SerializationMixin, Generic[TResponseModel]):
self.conversation_id = conversation_id
self.model_id = model_id
self.created_at = created_at
self.finish_reason: str | None = finish_reason
# Handle legacy dict format for finish_reason
if isinstance(finish_reason, dict) and "value" in finish_reason:
finish_reason = finish_reason["value"]
self.finish_reason = finish_reason
self.usage_details = usage_details
self._value: TResponseModel | None = value
self._response_format: type[BaseModel] | None = response_format
@@ -1907,25 +1911,25 @@ class ChatResponse(SerializationMixin, Generic[TResponseModel]):
@overload
@classmethod
def from_updates(
cls: type["ChatResponse[Any]"],
updates: Sequence["ChatResponseUpdate"],
cls: type[ChatResponse[Any]],
updates: Sequence[ChatResponseUpdate],
*,
output_format_type: type[TResponseModelT],
) -> "ChatResponse[TResponseModelT]": ...
) -> ChatResponse[TResponseModelT]: ...
@overload
@classmethod
def from_updates(
cls: type["ChatResponse[Any]"],
updates: Sequence["ChatResponseUpdate"],
cls: type[ChatResponse[Any]],
updates: Sequence[ChatResponseUpdate],
*,
output_format_type: None = None,
) -> "ChatResponse[Any]": ...
) -> ChatResponse[Any]: ...
@classmethod
def from_updates(
cls: type[TChatResponse],
updates: Sequence["ChatResponseUpdate"],
updates: Sequence[ChatResponseUpdate],
*,
output_format_type: type[BaseModel] | None = None,
) -> TChatResponse:
@@ -1962,25 +1966,25 @@ class ChatResponse(SerializationMixin, Generic[TResponseModel]):
@overload
@classmethod
async def from_update_generator(
cls: type["ChatResponse[Any]"],
updates: AsyncIterable["ChatResponseUpdate"],
cls: type[ChatResponse[Any]],
updates: AsyncIterable[ChatResponseUpdate],
*,
output_format_type: type[TResponseModelT],
) -> "ChatResponse[TResponseModelT]": ...
) -> ChatResponse[TResponseModelT]: ...
@overload
@classmethod
async def from_update_generator(
cls: type["ChatResponse[Any]"],
updates: AsyncIterable["ChatResponseUpdate"],
cls: type[ChatResponse[Any]],
updates: AsyncIterable[ChatResponseUpdate],
*,
output_format_type: None = None,
) -> "ChatResponse[Any]": ...
) -> ChatResponse[Any]: ...
@classmethod
async def from_update_generator(
cls: type[TChatResponse],
updates: AsyncIterable["ChatResponseUpdate"],
updates: AsyncIterable[ChatResponseUpdate],
*,
output_format_type: type[BaseModel] | None = None,
) -> TChatResponse:
@@ -2096,14 +2100,14 @@ class ChatResponseUpdate(SerializationMixin):
self,
*,
contents: Sequence[Content] | None = None,
role: RoleLiteral | str | None = None,
role: RoleLiteral | Role | None = None,
author_name: str | None = None,
response_id: str | None = None,
message_id: str | None = None,
conversation_id: str | None = None,
model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: FinishReasonLiteral | str | None = None,
finish_reason: FinishReasonLiteral | FinishReason | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
) -> None:
@@ -2138,20 +2142,14 @@ class ChatResponseUpdate(SerializationMixin):
processed_contents.append(c)
self.contents = processed_contents
# Handle legacy dict formats for role and finish_reason
if isinstance(role, dict) and "value" in role:
role = role["value"]
if isinstance(finish_reason, dict) and "value" in finish_reason:
finish_reason = finish_reason["value"]
self.role: str | None = role
self.role = role
self.author_name = author_name
self.response_id = response_id
self.message_id = message_id
self.conversation_id = conversation_id
self.model_id = model_id
self.created_at = created_at
self.finish_reason: str | None = finish_reason
self.finish_reason = finish_reason
self.additional_properties = additional_properties
self.raw_representation = raw_representation
@@ -2304,25 +2302,25 @@ class AgentResponse(SerializationMixin, Generic[TResponseModel]):
@overload
@classmethod
def from_updates(
cls: type["AgentResponse[Any]"],
updates: Sequence["AgentResponseUpdate"],
cls: type[AgentResponse[Any]],
updates: Sequence[AgentResponseUpdate],
*,
output_format_type: type[TResponseModelT],
) -> "AgentResponse[TResponseModelT]": ...
) -> AgentResponse[TResponseModelT]: ...
@overload
@classmethod
def from_updates(
cls: type["AgentResponse[Any]"],
updates: Sequence["AgentResponseUpdate"],
cls: type[AgentResponse[Any]],
updates: Sequence[AgentResponseUpdate],
*,
output_format_type: None = None,
) -> "AgentResponse[Any]": ...
) -> AgentResponse[Any]: ...
@classmethod
def from_updates(
cls: type[TAgentRunResponse],
updates: Sequence["AgentResponseUpdate"],
updates: Sequence[AgentResponseUpdate],
*,
output_format_type: type[BaseModel] | None = None,
) -> TAgentRunResponse:
@@ -2342,26 +2340,26 @@ class AgentResponse(SerializationMixin, Generic[TResponseModel]):
@overload
@classmethod
async def from_agent_response_generator(
cls: type["AgentResponse[Any]"],
updates: AsyncIterable["AgentResponseUpdate"],
async def from_update_generator(
cls: type[AgentResponse[Any]],
updates: AsyncIterable[AgentResponseUpdate],
*,
output_format_type: type[TResponseModelT],
) -> "AgentResponse[TResponseModelT]": ...
) -> AgentResponse[TResponseModelT]: ...
@overload
@classmethod
async def from_agent_response_generator(
cls: type["AgentResponse[Any]"],
updates: AsyncIterable["AgentResponseUpdate"],
async def from_update_generator(
cls: type[AgentResponse[Any]],
updates: AsyncIterable[AgentResponseUpdate],
*,
output_format_type: None = None,
) -> "AgentResponse[Any]": ...
) -> AgentResponse[Any]: ...
@classmethod
async def from_agent_response_generator(
async def from_update_generator(
cls: type[TAgentRunResponse],
updates: AsyncIterable["AgentResponseUpdate"],
updates: AsyncIterable[AgentResponseUpdate],
*,
output_format_type: type[BaseModel] | None = None,
) -> TAgentRunResponse:
@@ -2504,6 +2502,353 @@ class AgentResponseUpdate(SerializationMixin):
return self.text
# region ResponseStream
def map_chat_to_agent_update(update: ChatResponseUpdate, agent_name: str | None) -> AgentResponseUpdate:
return AgentResponseUpdate(
contents=update.contents,
role=update.role,
author_name=update.author_name or agent_name,
response_id=update.response_id,
message_id=update.message_id,
created_at=update.created_at,
additional_properties=update.additional_properties,
raw_representation=update,
)
# Type variables for ResponseStream
TUpdate = TypeVar("TUpdate")
TFinal = TypeVar("TFinal")
TOuterUpdate = TypeVar("TOuterUpdate")
TOuterFinal = TypeVar("TOuterFinal")
class ResponseStream(AsyncIterable[TUpdate], Generic[TUpdate, TFinal]):
"""Async stream wrapper that supports iteration and deferred finalization."""
def __init__(
self,
stream: AsyncIterable[TUpdate] | Awaitable[AsyncIterable[TUpdate]],
*,
finalizer: Callable[[Sequence[TUpdate]], TFinal | Awaitable[TFinal]] | None = None,
transform_hooks: list[Callable[[TUpdate], TUpdate | Awaitable[TUpdate] | None]] | None = None,
cleanup_hooks: list[Callable[[], Awaitable[None] | None]] | None = None,
result_hooks: list[Callable[[TFinal], TFinal | Awaitable[TFinal | None] | None]] | None = None,
) -> None:
"""A Async Iterable stream of updates.
Args:
stream: An async iterable or awaitable that resolves to an async iterable of updates.
Keyword Args:
finalizer: An optional callable that takes the list of all updates and produces a final result.
transform_hooks: Optional list of callables that transform each update as it is yielded.
cleanup_hooks: Optional list of callables that run after the stream is fully consumed (before finalizer).
result_hooks: Optional list of callables that transform the final result (after finalizer).
"""
self._stream_source = stream
self._finalizer = finalizer
self._stream: AsyncIterable[TUpdate] | None = None
self._iterator: AsyncIterator[TUpdate] | None = None
self._updates: list[TUpdate] = []
self._consumed: bool = False
self._finalized: bool = False
self._final_result: TFinal | None = None
self._transform_hooks: list[Callable[[TUpdate], TUpdate | Awaitable[TUpdate] | None]] = (
transform_hooks if transform_hooks is not None else []
)
self._result_hooks: list[Callable[[TFinal], TFinal | Awaitable[TFinal | None] | None]] = (
result_hooks if result_hooks is not None else []
)
self._cleanup_hooks: list[Callable[[], Awaitable[None] | None]] = (
cleanup_hooks if cleanup_hooks is not None else []
)
self._cleanup_run: bool = False
self._inner_stream: ResponseStream[Any, Any] | None = None
self._inner_stream_source: ResponseStream[Any, Any] | Awaitable[ResponseStream[Any, Any]] | None = None
self._wrap_inner: bool = False
self._map_update: Callable[[Any], Any | Awaitable[Any]] | None = None
def map(
self,
transform: Callable[[TUpdate], TOuterUpdate | Awaitable[TOuterUpdate]],
finalizer: Callable[[Sequence[TOuterUpdate]], TOuterFinal | Awaitable[TOuterFinal]],
) -> ResponseStream[TOuterUpdate, TOuterFinal]:
"""Create a new stream that transforms each update.
The returned stream delegates iteration to this stream, ensuring single consumption.
Each update is transformed by the provided function before being yielded.
Since the update type changes, a new finalizer MUST be provided that works with
the transformed update type. The inner stream's finalizer cannot be used as it
expects the original update type.
When ``get_final_response()`` is called on the mapped stream:
1. The inner stream's finalizer runs first (on the original updates)
2. The inner stream's result_hooks run (on the inner final result)
3. The outer stream's finalizer runs (on the transformed updates)
4. The outer stream's result_hooks run (on the outer final result)
This ensures that post-processing hooks registered on the inner stream (e.g.,
context provider notifications, telemetry) are still executed.
Args:
transform: Function to transform each update to a new type.
finalizer: Function to convert collected (transformed) updates to the final type.
This is required because the inner stream's finalizer won't work with
the new update type.
Returns:
A new ResponseStream with transformed update and final types.
Example:
>>> chat_stream.map(
... lambda u: AgentResponseUpdate(...),
... AgentResponse.from_updates,
... )
"""
stream: ResponseStream[Any, Any] = ResponseStream(self, finalizer=finalizer)
stream._inner_stream_source = self
stream._wrap_inner = True
stream._map_update = transform
return stream # type: ignore[return-value]
def with_finalizer(
self,
finalizer: Callable[[Sequence[TUpdate]], TOuterFinal | Awaitable[TOuterFinal]],
) -> ResponseStream[TUpdate, TOuterFinal]:
"""Create a new stream with a different finalizer.
The returned stream delegates iteration to this stream, ensuring single consumption.
When `get_final_response()` is called, the new finalizer is used instead of any
existing finalizer.
**IMPORTANT**: The inner stream's finalizer and result_hooks are NOT called when
a new finalizer is provided via this method.
Args:
finalizer: Function to convert collected updates to the final response type.
Returns:
A new ResponseStream with the new final type.
Example:
>>> stream.with_finalizer(AgentResponse.from_updates)
"""
stream: ResponseStream[Any, Any] = ResponseStream(self, finalizer=finalizer)
stream._inner_stream_source = self
stream._wrap_inner = True
return stream # type: ignore[return-value]
@classmethod
def from_awaitable(
cls,
awaitable: Awaitable[ResponseStream[TUpdate, TFinal]],
) -> ResponseStream[TUpdate, TFinal]:
"""Create a ResponseStream from an awaitable that resolves to a ResponseStream.
This is useful when you have an async function that returns a ResponseStream
and you want to wrap it to add hooks or use it in a pipeline.
The returned stream delegates to the inner stream once it resolves, using the
inner stream's finalizer if no new finalizer is provided.
Args:
awaitable: An awaitable that resolves to a ResponseStream.
Returns:
A new ResponseStream that wraps the awaitable.
Example:
>>> async def get_stream() -> ResponseStream[Update, Response]: ...
>>> stream = ResponseStream.from_awaitable(get_stream())
"""
stream: ResponseStream[Any, Any] = cls(awaitable) # type: ignore[arg-type]
stream._inner_stream_source = awaitable # type: ignore[assignment]
stream._wrap_inner = True
return stream # type: ignore[return-value]
async def _get_stream(self) -> AsyncIterable[TUpdate]:
if self._stream is None:
if hasattr(self._stream_source, "__aiter__"):
self._stream = self._stream_source # type: ignore[assignment]
else:
self._stream = await self._stream_source # type: ignore[assignment]
if isinstance(self._stream, ResponseStream) and self._wrap_inner:
self._inner_stream = self._stream
return self._stream
return self._stream # type: ignore[return-value]
def __aiter__(self) -> ResponseStream[TUpdate, TFinal]:
return self
async def __anext__(self) -> TUpdate:
if self._iterator is None:
stream = await self._get_stream()
self._iterator = stream.__aiter__()
try:
update = await self._iterator.__anext__()
except StopAsyncIteration:
self._consumed = True
await self._run_cleanup_hooks()
raise
except Exception:
await self._run_cleanup_hooks()
raise
if self._map_update is not None:
mapped = self._map_update(update)
if isinstance(mapped, Awaitable):
update = await mapped
else:
update = mapped # type: ignore[assignment]
self._updates.append(update)
for hook in self._transform_hooks:
hooked = hook(update)
if isinstance(hooked, Awaitable):
update = await hooked
elif hooked is not None:
update = hooked # type: ignore[assignment]
return update
def __await__(self) -> Any:
async def _wrap() -> ResponseStream[TUpdate, TFinal]:
await self._get_stream()
return self
return _wrap().__await__()
async def get_final_response(self) -> TFinal:
"""Get the final response by applying the finalizer to all collected updates.
If a finalizer is configured, it receives the list of updates and returns the final type.
Result hooks are then applied in order to transform the result.
If no finalizer is configured, returns the collected updates as Sequence[TUpdate].
For wrapped streams (created via .map() or .from_awaitable()):
- The inner stream's finalizer is called first to produce the inner final result.
- The inner stream's result_hooks are then applied to that inner result.
- The outer stream's finalizer is called to convert the outer (mapped) updates to the final type.
- The outer stream's result_hooks are then applied to transform the outer result.
This ensures that post-processing hooks registered on the inner stream (e.g., context
provider notifications) are still executed even when the stream is wrapped/mapped.
"""
if self._wrap_inner:
if self._inner_stream is None:
if self._inner_stream_source is None:
raise ValueError("No inner stream configured for this stream.")
if isinstance(self._inner_stream_source, ResponseStream):
self._inner_stream = self._inner_stream_source
else:
self._inner_stream = await self._inner_stream_source
if not self._finalized:
# Consume outer stream (which delegates to inner) if not already consumed
if not self._consumed:
async for _ in self:
pass
# First, finalize the inner stream and run its result hooks
# This ensures inner post-processing (e.g., context provider notifications) runs
if self._inner_stream._finalizer is not None:
inner_result: Any = self._inner_stream._finalizer(self._inner_stream._updates)
if isinstance(inner_result, Awaitable):
inner_result = await inner_result
else:
inner_result = self._inner_stream._updates
# Run inner stream's result hooks
for hook in self._inner_stream._result_hooks:
hooked = hook(inner_result)
if isinstance(hooked, Awaitable):
hooked = await hooked
if hooked is not None:
inner_result = hooked
self._inner_stream._final_result = inner_result
self._inner_stream._finalized = True
# Now finalize the outer stream with its own finalizer
# If outer has no finalizer, use inner's result (preserves from_awaitable behavior)
if self._finalizer is not None:
result: Any = self._finalizer(self._updates)
if isinstance(result, Awaitable):
result = await result
else:
# No outer finalizer - use inner's finalized result
result = inner_result
# Apply outer's result_hooks
for hook in self._result_hooks:
hooked = hook(result)
if isinstance(hooked, Awaitable):
hooked = await hooked
if hooked is not None:
result = hooked
self._final_result = result
self._finalized = True
return self._final_result # type: ignore[return-value]
if not self._finalized:
if not self._consumed:
async for _ in self:
pass
# Use finalizer if configured, otherwise return collected updates
if self._finalizer is not None:
result = self._finalizer(self._updates)
if isinstance(result, Awaitable):
result = await result
else:
result = self._updates
for hook in self._result_hooks:
hooked = hook(result)
if isinstance(hooked, Awaitable):
hooked = await hooked
if hooked is not None:
result = hooked
self._final_result = result
self._finalized = True
return self._final_result # type: ignore[return-value]
def with_transform_hook(
self,
hook: Callable[[TUpdate], TUpdate | Awaitable[TUpdate] | None],
) -> ResponseStream[TUpdate, TFinal]:
"""Register a transform hook executed for each update during iteration."""
self._transform_hooks.append(hook)
return self
def with_result_hook(
self,
hook: Callable[[TFinal], TFinal | Awaitable[TFinal | None] | None],
) -> ResponseStream[TUpdate, TFinal]:
"""Register a result hook executed after finalization."""
self._result_hooks.append(hook)
self._finalized = False
self._final_result = None
return self
def with_cleanup_hook(
self,
hook: Callable[[], Awaitable[None] | None],
) -> ResponseStream[TUpdate, TFinal]:
"""Register a cleanup hook executed after stream consumption (before finalizer)."""
self._cleanup_hooks.append(hook)
return self
async def _run_cleanup_hooks(self) -> None:
if self._cleanup_run:
return
self._cleanup_run = True
for hook in self._cleanup_hooks:
result = hook()
if isinstance(result, Awaitable):
await result
@property
def updates(self) -> Sequence[TUpdate]:
return self._updates
# region ChatOptions
@@ -2570,7 +2915,13 @@ class _ChatOptionsBase(TypedDict, total=False):
presence_penalty: float
# Tool configuration (forward reference to avoid circular import)
tools: "ToolProtocol | Callable[..., Any] | MutableMapping[str, Any] | Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] | None" # noqa: E501
tools: (
ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None
)
tool_choice: ToolMode | Literal["auto", "required", "none"]
allow_multiple_tool_calls: bool
python/packages/core/agent_framework/_workflows/_agent.py
+65 -37
View File
@@ -4,10 +4,10 @@ import json
import logging
import sys
import uuid
from collections.abc import AsyncIterable
from collections.abc import AsyncIterable, Awaitable
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, ClassVar, cast
from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload
from agent_framework import (
AgentResponse,
@@ -124,24 +124,49 @@ class WorkflowAgent(BaseAgent):
# region Run Methods
async def run(
@overload
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AgentResponse:
"""Get a response from the workflow agent (non-streaming).
) -> AsyncIterable[AgentResponseUpdate]: ...
This method runs the workflow in non-streaming mode.
@overload
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AgentResponse: ...
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate] | Awaitable[AgentResponse]:
"""Get a response from the workflow agent.
Args:
messages: The message(s) to send to the workflow. Required for new runs,
should be None when resuming from checkpoint.
Keyword Args:
stream: If True, returns an async iterable of updates. If False (default),
returns an awaitable AgentResponse.
thread: The conversation thread. If None, a new thread will be created.
checkpoint_id: ID of checkpoint to restore from. If provided, the workflow
resumes from this checkpoint instead of starting fresh.
@@ -152,12 +177,35 @@ class WorkflowAgent(BaseAgent):
and tool functions.
Returns:
An AgentResponse representing the workflow execution results. The response
includes all output events and requests emitted during the workflow run.
WorkflowOutputEvents will be converted to ChatMessages in the response.
RequestInfoEvents will be converted to function call and approval request contents
in the response.
When stream=True: An AsyncIterable[AgentResponseUpdate] for streaming updates.
When stream=False: An Awaitable[AgentResponse] with the complete response.
"""
if stream:
return self._run_streaming(
messages=messages,
thread=thread,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
**kwargs,
)
return self._run_non_streaming(
messages=messages,
thread=thread,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
**kwargs,
)
async def _run_non_streaming(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AgentResponse:
"""Internal non-streaming implementation."""
input_messages = normalize_messages_input(messages)
thread = thread or self.get_new_thread()
response_id = str(uuid.uuid4())
@@ -171,7 +219,7 @@ class WorkflowAgent(BaseAgent):
return response
async def run_stream(
async def _run_streaming(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
@@ -180,29 +228,7 @@ class WorkflowAgent(BaseAgent):
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
"""Stream response updates from the workflow agent.
Args:
messages: The message(s) to send to the workflow. Required for new runs,
should be None when resuming from checkpoint.
Keyword Args:
thread: The conversation thread. If None, a new thread will be created.
checkpoint_id: ID of checkpoint to restore from. If provided, the workflow
resumes from this checkpoint instead of starting fresh.
checkpoint_storage: Runtime checkpoint storage. When provided with checkpoint_id,
used to load and restore the checkpoint. When provided without checkpoint_id,
enables checkpointing for this run.
**kwargs: Additional keyword arguments passed through to underlying workflow
and tool functions.
Yields:
AgentResponseUpdate objects representing the workflow execution progress.
Updates include output events and requests emitted during the workflow run.
WorkflowOutputEvents will be converted to AgentResponseUpdate objects.
RequestInfoEvents will be converted to function call and approval request contents
in the updates.
"""
"""Internal streaming implementation."""
input_messages = normalize_messages_input(messages)
thread = thread or self.get_new_thread()
response_updates: list[AgentResponseUpdate] = []
@@ -322,8 +348,9 @@ class WorkflowAgent(BaseAgent):
# Resume from checkpoint - don't prepend thread history since workflow state
# is being restored from the checkpoint
if streaming:
async for event in self.workflow.run_stream(
async for event in self.workflow.run(
message=None,
stream=True,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
**kwargs,
@@ -344,8 +371,9 @@ class WorkflowAgent(BaseAgent):
conversation_messages = await self._build_conversation_messages(thread, input_messages)
if streaming:
async for event in self.workflow.run_stream(
async for event in self.workflow.run(
message=conversation_messages,
stream=True,
checkpoint_storage=checkpoint_storage,
**kwargs,
):
python/packages/core/agent_framework/_workflows/_agent_executor.py
+5 -3
View File
@@ -65,7 +65,7 @@ class AgentExecutor(Executor):
"""built-in executor that wraps an agent for handling messages.
AgentExecutor adapts its behavior based on the workflow execution mode:
- run_stream(): Emits incremental WorkflowOutputEvents as the agent produces tokens
- run(stream=True): Emits incremental WorkflowOutputEvents as the agent produces tokens
- run(): Emits a single WorkflowOutputEvent containing the complete response
Use `with_output_from` in WorkflowBuilder to control whether the AgentResponse
@@ -195,7 +195,7 @@ class AgentExecutor(Executor):
if not self._pending_agent_requests:
# All pending requests have been resolved; resume agent execution
self._cache = normalize_messages_input(ChatMessage("user", self._pending_responses_to_agent))
self._cache = normalize_messages_input(ChatMessage(role="user", contents=self._pending_responses_to_agent))
self._pending_responses_to_agent.clear()
await self._run_agent_and_emit(ctx)
@@ -334,6 +334,7 @@ class AgentExecutor(Executor):
response = await self._agent.run(
self._cache,
stream=False,
thread=self._agent_thread,
**run_kwargs,
)
@@ -361,8 +362,9 @@ class AgentExecutor(Executor):
updates: list[AgentResponseUpdate] = []
user_input_requests: list[Content] = []
async for update in self._agent.run_stream(
async for update in self._agent.run(
self._cache,
stream=True,
thread=self._agent_thread,
**run_kwargs,
):
python/packages/core/agent_framework/_workflows/_base_group_chat_orchestrator.py
+6 -6
View File
@@ -214,7 +214,7 @@ class BaseGroupChatOrchestrator(Executor, ABC):
Usage:
workflow.run("Write a blog post about AI agents")
"""
await self._handle_messages([ChatMessage("user", [task])], ctx)
await self._handle_messages([ChatMessage(role="user", text=task)], ctx)
@handler
async def handle_message(
@@ -231,7 +231,7 @@ class BaseGroupChatOrchestrator(Executor, ABC):
ctx: Workflow context
Usage:
workflow.run(ChatMessage("user", ["Write a blog post about AI agents"]))
workflow.run(ChatMessage(role="user", text="Write a blog post about AI agents"))
"""
await self._handle_messages([task], ctx)
@@ -250,8 +250,8 @@ class BaseGroupChatOrchestrator(Executor, ABC):
ctx: Workflow context
Usage:
workflow.run([
ChatMessage("user", ["Write a blog post about AI agents"]),
ChatMessage("user", ["Make it engaging and informative."])
ChatMessage(role="user", text="Write a blog post about AI agents"),
ChatMessage(role="user", text="Make it engaging and informative.")
])
"""
if not task:
@@ -401,7 +401,7 @@ class BaseGroupChatOrchestrator(Executor, ABC):
Returns:
ChatMessage with completion content
"""
return ChatMessage("assistant", [message], author_name=self._name)
return ChatMessage(role="assistant", text=message, author_name=self._name)
# Participant routing (shared across all patterns)
@@ -465,7 +465,7 @@ class BaseGroupChatOrchestrator(Executor, ABC):
# AgentExecutors receive simple message list
messages: list[ChatMessage] = []
if additional_instruction:
messages.append(ChatMessage("user", [additional_instruction]))
messages.append(ChatMessage(role="user", text=additional_instruction))
request = AgentExecutorRequest(messages=messages, should_respond=True)
await ctx.send_message(request, target_id=target)
await ctx.add_event(
python/packages/core/agent_framework/_workflows/_const.py
+1 -1
View File
@@ -11,7 +11,7 @@ INTERNAL_SOURCE_PREFIX = "internal"
# State key for storing run kwargs that should be passed to agent invocations.
# Used by all orchestration patterns (Sequential, Concurrent, GroupChat, Handoff, Magentic)
# to pass kwargs from workflow.run_stream() through to agent.run_stream() and @tool functions.
# to pass kwargs from workflow.run() through to agent.run() and @tool functions.
WORKFLOW_RUN_KWARGS_KEY = "_workflow_run_kwargs"
python/packages/core/agent_framework/_workflows/_conversation_state.py
+1 -1
View File
@@ -64,7 +64,7 @@ def decode_chat_messages(payload: Iterable[dict[str, Any]]) -> list[ChatMessage]
additional[key] = decode_checkpoint_value(value)
restored.append(
ChatMessage(
ChatMessage( # type: ignore[call-overload]
role=role,
contents=contents,
author_name=item.get("author_name"),
python/packages/core/agent_framework/_workflows/_message_utils.py
+2 -2
View File
@@ -22,7 +22,7 @@ def normalize_messages_input(
return []
if isinstance(messages, str):
return [ChatMessage("user", [messages])]
return [ChatMessage(role="user", text=messages)]
if isinstance(messages, ChatMessage):
return [messages]
@@ -30,7 +30,7 @@ def normalize_messages_input(
normalized: list[ChatMessage] = []
for item in messages:
if isinstance(item, str):
normalized.append(ChatMessage("user", [item]))
normalized.append(ChatMessage(role="user", text=item))
elif isinstance(item, ChatMessage):
normalized.append(item)
else:
python/packages/core/agent_framework/_workflows/_orchestration_request_info.py
+1 -1
View File
@@ -72,7 +72,7 @@ class AgentRequestInfoResponse:
Returns:
AgentRequestInfoResponse instance.
"""
return AgentRequestInfoResponse(messages=[ChatMessage("user", [text]) for text in texts])
return AgentRequestInfoResponse(messages=[ChatMessage(role="user", text=text) for text in texts])
@staticmethod
def approve() -> "AgentRequestInfoResponse":
python/packages/core/agent_framework/_workflows/_orchestrator_helpers.py
+2 -2
View File
@@ -89,7 +89,7 @@ def create_completion_message(
"""
message_text = text or f"Conversation {reason}."
return ChatMessage(
"assistant",
[message_text],
role="assistant",
text=message_text,
author_name=author_name,
)
python/packages/core/agent_framework/_workflows/_runner_context.py
+3 -3
View File
@@ -203,7 +203,7 @@ class RunnerContext(Protocol):
"""Set whether agents should stream incremental updates.
Args:
streaming: True for streaming mode (run_stream), False for non-streaming (run).
streaming: True for streaming mode (stream=True), False for non-streaming (stream=False).
"""
...
@@ -301,7 +301,7 @@ class InProcRunnerContext:
self._runtime_checkpoint_storage: CheckpointStorage | None = None
self._workflow_id: str | None = None
# Streaming flag - set by workflow's run_stream() vs run()
# Streaming flag - set by workflow's run(..., stream=True) vs run(..., stream=False)
self._streaming: bool = False
# region Messaging and Events
@@ -442,7 +442,7 @@ class InProcRunnerContext:
"""Set whether agents should stream incremental updates.
Args:
streaming: True for streaming mode (run_stream), False for non-streaming (run).
streaming: True for streaming mode (run(stream=True)), False for non-streaming.
"""
self._streaming = streaming
python/packages/core/agent_framework/_workflows/_workflow.py
+78 -141
View File
@@ -8,7 +8,7 @@ import logging
import types
import uuid
from collections.abc import AsyncIterable, Awaitable, Callable
from typing import Any
from typing import Any, Literal, overload
from ..observability import OtelAttr, capture_exception, create_workflow_span
from ._agent import WorkflowAgent
@@ -129,7 +129,7 @@ class Workflow(DictConvertible):
The workflow provides two primary execution APIs, each supporting multiple scenarios:
- **run()**: Execute to completion, returns WorkflowRunResult with all events
- **run_stream()**: Returns async generator yielding events as they occur
- **run(..., stream=True)**: Returns ResponseStream yielding events as they occur
Both methods support:
- Initial workflow runs: Provide `message` parameter
@@ -138,7 +138,7 @@ class Workflow(DictConvertible):
- Runtime checkpointing: Provide `checkpoint_storage` to enable/override checkpointing for this run
## State Management
Workflow instances contain states and states are preserved across calls to `run` and `run_stream`.
Workflow instances contain states and states are preserved across calls to `run`.
To execute multiple independent runs, create separate Workflow instances via WorkflowBuilder.
## External Input Requests
@@ -156,7 +156,7 @@ class Workflow(DictConvertible):
Build-time (via WorkflowBuilder):
workflow = WorkflowBuilder().with_checkpointing(storage).build()
Runtime (via run/run_stream parameters):
Runtime (via run parameters):
result = await workflow.run(message, checkpoint_storage=runtime_storage)
When enabled, checkpoints are created at the end of each superstep, capturing:
@@ -447,7 +447,77 @@ class Workflow(DictConvertible):
source_span_ids=None,
)
async def run_stream(
@overload
def run(
self,
message: Any | None = None,
*,
stream: Literal[True],
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AsyncIterable[WorkflowEvent]: ...
@overload
async def run(
self,
message: Any | None = None,
*,
stream: Literal[False] = ...,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
include_status_events: bool = False,
**kwargs: Any,
) -> WorkflowRunResult: ...
def run(
self,
message: Any | None = None,
*,
stream: bool = False,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
include_status_events: bool = False,
**kwargs: Any,
) -> AsyncIterable[WorkflowEvent] | Awaitable[WorkflowRunResult]:
"""Run the workflow, optionally streaming events.
Unified interface supporting initial runs and checkpoint restoration.
Args:
message: Initial message for the start executor. Required for new workflow runs,
should be None when resuming from checkpoint.
stream: If True, returns an async iterable of events. If False (default),
returns an awaitable WorkflowRunResult.
checkpoint_id: ID of checkpoint to restore from. If provided, the workflow resumes
from this checkpoint instead of starting fresh.
checkpoint_storage: Runtime checkpoint storage.
include_status_events: Whether to include WorkflowStatusEvent instances (non-streaming only).
**kwargs: Additional keyword arguments to pass through to agent invocations.
Returns:
When stream=True: An AsyncIterable[WorkflowEvent] for streaming events.
When stream=False: An Awaitable[WorkflowRunResult] with all events.
Raises:
ValueError: If both message and checkpoint_id are provided, or if neither is provided.
"""
if stream:
return self._run_streaming(
message=message,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
**kwargs,
)
return self._run_non_streaming(
message=message,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
include_status_events=include_status_events,
**kwargs,
)
async def _run_streaming(
self,
message: Any | None = None,
*,
@@ -455,75 +525,7 @@ class Workflow(DictConvertible):
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AsyncIterable[WorkflowEvent]:
"""Run the workflow and stream events.
Unified streaming interface supporting initial runs and checkpoint restoration.
Args:
message: Initial message for the start executor. Required for new workflow runs,
should be None when resuming from checkpoint.
checkpoint_id: ID of checkpoint to restore from. If provided, the workflow resumes
from this checkpoint instead of starting fresh. When resuming, checkpoint_storage
must be provided (either at build time or runtime) to load the checkpoint.
checkpoint_storage: Runtime checkpoint storage with two behaviors:
- With checkpoint_id: Used to load and restore the specified checkpoint
- Without checkpoint_id: Enables checkpointing for this run, overriding
build-time configuration
**kwargs: Additional keyword arguments to pass through to agent invocations.
These are stored in State and accessible in @tool functions
via the **kwargs parameter.
Yields:
WorkflowEvent: Events generated during workflow execution.
Raises:
ValueError: If both message and checkpoint_id are provided, or if neither is provided.
ValueError: If checkpoint_id is provided but no checkpoint storage is available
(neither at build time nor runtime).
RuntimeError: If checkpoint restoration fails.
Examples:
Initial run:
.. code-block:: python
async for event in workflow.run_stream("start message"):
process(event)
With custom context for tools:
.. code-block:: python
async for event in workflow.run_stream(
"analyze data",
custom_data={"endpoint": "https://api.example.com"},
user_token={"user": "alice"},
):
process(event)
Enable checkpointing at runtime:
.. code-block:: python
storage = FileCheckpointStorage("./checkpoints")
async for event in workflow.run_stream("start", checkpoint_storage=storage):
process(event)
Resume from checkpoint (storage provided at build time):
.. code-block:: python
async for event in workflow.run_stream(checkpoint_id="cp_123"):
process(event)
Resume from checkpoint (storage provided at runtime):
.. code-block:: python
storage = FileCheckpointStorage("./checkpoints")
async for event in workflow.run_stream(checkpoint_id="cp_123", checkpoint_storage=storage):
process(event)
"""
"""Internal streaming implementation."""
# Validate mutually exclusive parameters BEFORE setting running flag
if message is not None and checkpoint_id is not None:
raise ValueError("Cannot provide both 'message' and 'checkpoint_id'. Use one or the other.")
@@ -583,7 +585,7 @@ class Workflow(DictConvertible):
finally:
self._reset_running_flag()
async def run(
async def _run_non_streaming(
self,
message: Any | None = None,
*,
@@ -592,72 +594,7 @@ class Workflow(DictConvertible):
include_status_events: bool = False,
**kwargs: Any,
) -> WorkflowRunResult:
"""Run the workflow to completion and return all events.
Unified non-streaming interface supporting initial runs and checkpoint restoration.
Args:
message: Initial message for the start executor. Required for new workflow runs,
should be None when resuming from checkpoint.
checkpoint_id: ID of checkpoint to restore from. If provided, the workflow resumes
from this checkpoint instead of starting fresh. When resuming, checkpoint_storage
must be provided (either at build time or runtime) to load the checkpoint.
checkpoint_storage: Runtime checkpoint storage with two behaviors:
- With checkpoint_id: Used to load and restore the specified checkpoint
- Without checkpoint_id: Enables checkpointing for this run, overriding
build-time configuration
include_status_events: Whether to include WorkflowStatusEvent instances in the result list.
**kwargs: Additional keyword arguments to pass through to agent invocations.
These are stored in State and accessible in @tool functions
via the **kwargs parameter.
Returns:
A WorkflowRunResult instance containing events generated during workflow execution.
Raises:
ValueError: If both message and checkpoint_id are provided, or if neither is provided.
ValueError: If checkpoint_id is provided but no checkpoint storage is available
(neither at build time nor runtime).
RuntimeError: If checkpoint restoration fails.
Examples:
Initial run:
.. code-block:: python
result = await workflow.run("start message")
outputs = result.get_outputs()
With custom context for tools:
.. code-block:: python
result = await workflow.run(
"analyze data",
custom_data={"endpoint": "https://api.example.com"},
user_token={"user": "alice"},
)
Enable checkpointing at runtime:
.. code-block:: python
storage = FileCheckpointStorage("./checkpoints")
result = await workflow.run("start", checkpoint_storage=storage)
Resume from checkpoint (storage provided at build time):
.. code-block:: python
result = await workflow.run(checkpoint_id="cp_123")
Resume from checkpoint (storage provided at runtime):
.. code-block:: python
storage = FileCheckpointStorage("./checkpoints")
result = await workflow.run(checkpoint_id="cp_123", checkpoint_storage=storage)
"""
"""Internal non-streaming implementation."""
# Validate mutually exclusive parameters BEFORE setting running flag
if message is not None and checkpoint_id is not None:
raise ValueError("Cannot provide both 'message' and 'checkpoint_id'. Use one or the other.")
python/packages/core/agent_framework/_workflows/_workflow_context.py
+1 -1
View File
@@ -460,6 +460,6 @@ class WorkflowContext(Generic[OutT, W_OutT]):
"""Check if the workflow is running in streaming mode.
Returns:
True if the workflow was started with run_stream(), False if started with run().
True if the workflow was started with stream=True, False otherwise.
"""
return self._runner_context.is_streaming()
python/packages/core/agent_framework/ag_ui/__init__.py
+1
View File
@@ -8,6 +8,7 @@ PACKAGE_NAME = "agent-framework-ag-ui"
_IMPORTS = [
"__version__",
"AgentFrameworkAgent",
"AGUIThread",
"add_agent_framework_fastapi_endpoint",
"AGUIChatClient",
"AGUIEventConverter",
python/packages/core/agent_framework/azure/_chat_client.py
+26 -13
View File
@@ -3,8 +3,8 @@
import json
import logging
import sys
from collections.abc import Mapping
from typing import Any, Generic
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Generic
from azure.core.credentials import TokenCredential
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
@@ -14,15 +14,17 @@ from pydantic import BaseModel, ValidationError
from agent_framework import (
Annotation,
ChatMiddlewareLayer,
ChatResponse,
ChatResponseUpdate,
Content,
use_chat_middleware,
use_function_invocation,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
)
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.observability import use_instrumentation
from agent_framework.openai._chat_client import OpenAIBaseChatClient, OpenAIChatOptions
from agent_framework.observability import ChatTelemetryLayer
from agent_framework.openai import OpenAIChatOptions
from agent_framework.openai._chat_client import RawOpenAIChatClient
from ._shared import (
AzureOpenAIConfigMixin,
@@ -42,6 +44,9 @@ if sys.version_info >= (3, 11):
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from agent_framework._middleware import MiddlewareTypes
logger: logging.Logger = logging.getLogger(__name__)
__all__ = ["AzureOpenAIChatClient", "AzureOpenAIChatOptions", "AzureUserSecurityContext"]
@@ -143,13 +148,15 @@ TChatResponse = TypeVar("TChatResponse", ChatResponse, ChatResponseUpdate)
TAzureOpenAIChatClient = TypeVar("TAzureOpenAIChatClient", bound="AzureOpenAIChatClient")
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class AzureOpenAIChatClient(
AzureOpenAIConfigMixin, OpenAIBaseChatClient[TAzureOpenAIChatOptions], Generic[TAzureOpenAIChatOptions]
class AzureOpenAIChatClient( # type: ignore[misc]
AzureOpenAIConfigMixin,
ChatMiddlewareLayer[TAzureOpenAIChatOptions],
FunctionInvocationLayer[TAzureOpenAIChatOptions],
ChatTelemetryLayer[TAzureOpenAIChatOptions],
RawOpenAIChatClient[TAzureOpenAIChatOptions],
Generic[TAzureOpenAIChatOptions],
):
"""Azure OpenAI Chat completion class."""
"""Azure OpenAI Chat completion class with middleware, telemetry, and function invocation support."""
def __init__(
self,
@@ -168,6 +175,8 @@ class AzureOpenAIChatClient(
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: str | None = None,
middleware: Sequence["MiddlewareTypes"] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
"""Initialize an Azure OpenAI Chat completion client.
@@ -199,6 +208,8 @@ class AzureOpenAIChatClient(
env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'.
instruction_role: The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`.
middleware: Optional sequence of middleware to apply to requests.
function_invocation_configuration: Optional configuration for function invocation behavior.
kwargs: Other keyword parameters.
Examples:
@@ -269,6 +280,8 @@ class AzureOpenAIChatClient(
default_headers=default_headers,
client=async_client,
instruction_role=instruction_role,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
@@ -276,7 +289,7 @@ class AzureOpenAIChatClient(
def _parse_text_from_openai(self, choice: Choice | ChunkChoice) -> Content | None:
"""Parse the choice into a Content object with type='text'.
Overwritten from OpenAIBaseChatClient to deal with Azure On Your Data function.
Overwritten from RawOpenAIChatClient to deal with Azure On Your Data function.
For docs see:
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/references/on-your-data?tabs=python#context
"""
python/packages/core/agent_framework/azure/_responses_client.py
+18 -11
View File
@@ -1,7 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
import sys
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Generic
from urllib.parse import urljoin
@@ -9,11 +9,11 @@ from azure.core.credentials import TokenCredential
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
from pydantic import ValidationError
from .._middleware import use_chat_middleware
from .._tools import use_function_invocation
from .._middleware import ChatMiddlewareLayer
from .._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
from ..exceptions import ServiceInitializationError
from ..observability import use_instrumentation
from ..openai._responses_client import OpenAIBaseResponsesClient
from ..observability import ChatTelemetryLayer
from ..openai._responses_client import RawOpenAIResponsesClient
from ._shared import (
AzureOpenAIConfigMixin,
AzureOpenAISettings,
@@ -33,6 +33,7 @@ else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from .._middleware import MiddlewareTypes
from ..openai._responses_client import OpenAIResponsesOptions
__all__ = ["AzureOpenAIResponsesClient"]
@@ -46,15 +47,15 @@ TAzureOpenAIResponsesOptions = TypeVar(
)
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class AzureOpenAIResponsesClient(
class AzureOpenAIResponsesClient( # type: ignore[misc]
AzureOpenAIConfigMixin,
OpenAIBaseResponsesClient[TAzureOpenAIResponsesOptions],
ChatMiddlewareLayer[TAzureOpenAIResponsesOptions],
FunctionInvocationLayer[TAzureOpenAIResponsesOptions],
ChatTelemetryLayer[TAzureOpenAIResponsesOptions],
RawOpenAIResponsesClient[TAzureOpenAIResponsesOptions],
Generic[TAzureOpenAIResponsesOptions],
):
"""Azure Responses completion class."""
"""Azure Responses completion class with middleware, telemetry, and function invocation support."""
def __init__(
self,
@@ -73,6 +74,8 @@ class AzureOpenAIResponsesClient(
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: str | None = None,
middleware: Sequence["MiddlewareTypes"] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
"""Initialize an Azure OpenAI Responses client.
@@ -104,6 +107,8 @@ class AzureOpenAIResponsesClient(
env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'.
instruction_role: The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`.
middleware: Optional sequence of middleware to apply to requests.
function_invocation_configuration: Optional configuration for function invocation behavior.
kwargs: Additional keyword arguments.
Examples:
@@ -184,6 +189,8 @@ class AzureOpenAIResponsesClient(
default_headers=default_headers,
client=async_client,
instruction_role=instruction_role,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
)
@override
python/packages/core/agent_framework/observability.py
+391 -518
View File
File diff suppressed because it is too large Load Diff
python/packages/core/agent_framework/openai/__init__.py
-1
View File
@@ -1,6 +1,5 @@
# Copyright (c) Microsoft. All rights reserved.
from ._assistant_provider import * # noqa: F403
from ._assistants_client import * # noqa: F403
from ._chat_client import * # noqa: F403
python/packages/core/agent_framework/openai/_assistant_provider.py
+9 -9
View File
@@ -10,7 +10,7 @@ from pydantic import BaseModel, SecretStr, ValidationError
from .._agents import ChatAgent
from .._memory import ContextProvider
from .._middleware import Middleware
from .._middleware import MiddlewareTypes
from .._tools import FunctionTool, ToolProtocol
from .._types import normalize_tools
from ..exceptions import ServiceInitializationError
@@ -204,7 +204,7 @@ class OpenAIAssistantProvider(Generic[TOptions_co]):
tools: _ToolsType | None = None,
metadata: dict[str, str] | None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Create a new assistant on OpenAI and return a ChatAgent.
@@ -226,7 +226,7 @@ class OpenAIAssistantProvider(Generic[TOptions_co]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
Include ``response_format`` here for structured output responses.
middleware: Middleware for the ChatAgent.
middleware: MiddlewareTypes for the ChatAgent.
context_provider: Context provider for the ChatAgent.
Returns:
@@ -312,7 +312,7 @@ class OpenAIAssistantProvider(Generic[TOptions_co]):
tools: _ToolsType | None = None,
instructions: str | None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Retrieve an existing assistant by ID and return a ChatAgent.
@@ -331,7 +331,7 @@ class OpenAIAssistantProvider(Generic[TOptions_co]):
instructions: Override the assistant's instructions (optional).
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: Middleware for the ChatAgent.
middleware: MiddlewareTypes for the ChatAgent.
context_provider: Context provider for the ChatAgent.
Returns:
@@ -378,7 +378,7 @@ class OpenAIAssistantProvider(Generic[TOptions_co]):
tools: _ToolsType | None = None,
instructions: str | None = None,
default_options: TOptions_co | None = None,
middleware: Sequence[Middleware] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> "ChatAgent[TOptions_co]":
"""Wrap an existing SDK Assistant object as a ChatAgent.
@@ -396,7 +396,7 @@ class OpenAIAssistantProvider(Generic[TOptions_co]):
instructions: Override the assistant's instructions (optional).
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: Middleware for the ChatAgent.
middleware: MiddlewareTypes for the ChatAgent.
context_provider: Context provider for the ChatAgent.
Returns:
@@ -520,7 +520,7 @@ class OpenAIAssistantProvider(Generic[TOptions_co]):
assistant: Assistant,
tools: list[ToolProtocol | MutableMapping[str, Any]] | None,
instructions: str | None,
middleware: Sequence[Middleware] | None,
middleware: Sequence[MiddlewareTypes] | None,
context_provider: ContextProvider | None,
default_options: TOptions_co | None = None,
**kwargs: Any,
@@ -531,7 +531,7 @@ class OpenAIAssistantProvider(Generic[TOptions_co]):
assistant: The OpenAI Assistant object.
tools: Tools for the agent.
instructions: Instructions override.
middleware: Middleware for the agent.
middleware: MiddlewareTypes for the agent.
context_provider: Context provider for the agent.
default_options: Default chat options for the agent (may include response_format).
**kwargs: Additional arguments passed to ChatAgent.
python/packages/core/agent_framework/openai/_assistants_client.py
+63 -44
View File
@@ -8,9 +8,9 @@ from collections.abc import (
Callable,
Mapping,
MutableMapping,
MutableSequence,
Sequence,
)
from typing import Any, Generic, Literal, cast
from typing import TYPE_CHECKING, Any, Generic, Literal, TypedDict, cast
from openai import AsyncOpenAI
from openai.types.beta.threads import (
@@ -28,12 +28,13 @@ from openai.types.beta.threads.runs import RunStep
from pydantic import BaseModel, ValidationError
from .._clients import BaseChatClient
from .._middleware import use_chat_middleware
from .._middleware import ChatMiddlewareLayer
from .._tools import (
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
HostedCodeInterpreterTool,
HostedFileSearchTool,
use_function_invocation,
)
from .._types import (
ChatMessage,
@@ -41,11 +42,12 @@ from .._types import (
ChatResponse,
ChatResponseUpdate,
Content,
ResponseStream,
UsageDetails,
prepare_function_call_results,
)
from ..exceptions import ServiceInitializationError
from ..observability import use_instrumentation
from ..observability import ChatTelemetryLayer
from ._shared import OpenAIConfigMixin, OpenAISettings
if sys.version_info >= (3, 13):
@@ -63,6 +65,8 @@ if sys.version_info >= (3, 11):
else:
from typing_extensions import Self, TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from .._middleware import MiddlewareTypes
__all__ = [
"AssistantToolResources",
@@ -198,15 +202,15 @@ TOpenAIAssistantsOptions = TypeVar(
# endregion
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class OpenAIAssistantsClient(
class OpenAIAssistantsClient( # type: ignore[misc]
OpenAIConfigMixin,
ChatMiddlewareLayer[TOpenAIAssistantsOptions],
FunctionInvocationLayer[TOpenAIAssistantsOptions],
ChatTelemetryLayer[TOpenAIAssistantsOptions],
BaseChatClient[TOpenAIAssistantsOptions],
Generic[TOpenAIAssistantsOptions],
):
"""OpenAI Assistants client."""
"""OpenAI Assistants client with middleware, telemetry, and function invocation support."""
def __init__(
self,
@@ -223,6 +227,8 @@ class OpenAIAssistantsClient(
async_client: AsyncOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
middleware: Sequence["MiddlewareTypes"] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
"""Initialize an OpenAI Assistants client.
@@ -249,6 +255,8 @@ class OpenAIAssistantsClient(
env_file_path: Use the environment settings file as a fallback
to environment variables.
env_file_encoding: The encoding of the environment settings file.
middleware: Optional sequence of middleware to apply to requests.
function_invocation_configuration: Optional configuration for function invocation behavior.
kwargs: Other keyword parameters.
Examples:
@@ -308,6 +316,8 @@ class OpenAIAssistantsClient(
default_headers=default_headers,
client=async_client,
base_url=openai_settings.base_url,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
)
self.assistant_id: str | None = assistant_id
self.assistant_name: str | None = assistant_name
@@ -337,44 +347,51 @@ class OpenAIAssistantsClient(
object.__setattr__(self, "_should_delete_assistant", False)
@override
async def _inner_get_response(
def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
) -> ChatResponse:
return await ChatResponse.from_update_generator(
updates=self._inner_get_streaming_response(messages=messages, options=options, **kwargs),
output_format_type=options.get("response_format"),
)
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
if stream:
# Streaming mode - return the async generator directly
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
# prepare
run_options, tool_results = self._prepare_options(messages, options, **kwargs)
@override
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
# prepare
run_options, tool_results = self._prepare_options(messages, options, **kwargs)
# Get the thread ID
thread_id: str | None = options.get(
"conversation_id", run_options.get("conversation_id", self.thread_id)
)
# Get the thread ID
thread_id: str | None = options.get("conversation_id", run_options.get("conversation_id", self.thread_id))
if thread_id is None and tool_results is not None:
raise ValueError("No thread ID was provided, but chat messages includes tool results.")
if thread_id is None and tool_results is not None:
raise ValueError("No thread ID was provided, but chat messages includes tool results.")
# Determine which assistant to use and create if needed
assistant_id = await self._get_assistant_id_or_create()
# Determine which assistant to use and create if needed
assistant_id = await self._get_assistant_id_or_create()
# execute
stream_obj, thread_id = await self._create_assistant_stream(
thread_id, assistant_id, run_options, tool_results
)
# execute
stream, thread_id = await self._create_assistant_stream(thread_id, assistant_id, run_options, tool_results)
# process
async for update in self._process_stream_events(stream_obj, thread_id):
yield update
# process
async for update in self._process_stream_events(stream, thread_id):
yield update
return self._build_response_stream(_stream(), response_format=options.get("response_format"))
# Non-streaming mode - collect updates and convert to response
async def _get_response() -> ChatResponse:
stream_result = self._inner_get_response(messages=messages, options=options, stream=True, **kwargs)
return await ChatResponse.from_update_generator(
updates=stream_result, # type: ignore[arg-type]
output_format_type=options.get("response_format"), # type: ignore[arg-type]
)
return _get_response()
async def _get_assistant_id_or_create(self) -> str:
"""Determine which assistant to use and create if needed.
@@ -489,8 +506,8 @@ class OpenAIAssistantsClient(
for delta_block in delta.content or []:
if isinstance(delta_block, TextDeltaBlock) and delta_block.text and delta_block.text.value:
yield ChatResponseUpdate(
role=role,
contents=[Content.from_text(text=delta_block.text.value)],
role=role, # type: ignore[arg-type]
contents=[Content.from_text(delta_block.text.value)],
conversation_id=thread_id,
message_id=response_id,
raw_representation=response.data,
@@ -586,8 +603,8 @@ class OpenAIAssistantsClient(
def _prepare_options(
self,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
**kwargs: Any,
) -> tuple[dict[str, Any], list[Content] | None]:
from .._types import validate_tool_mode
@@ -618,7 +635,9 @@ class OpenAIAssistantsClient(
tool_mode = validate_tool_mode(tool_choice)
tool_definitions: list[MutableMapping[str, Any]] = []
if tool_mode["mode"] != "none" and tools is not None:
# Always include tools if provided, regardless of tool_choice
# tool_choice="none" means the model won't call tools, but tools should still be available
if tools is not None:
for tool in tools:
if isinstance(tool, FunctionTool):
tool_definitions.append(tool.to_json_schema_spec()) # type: ignore[reportUnknownArgumentType]
python/packages/core/agent_framework/openai/_chat_client.py
+112 -76
View File
@@ -2,7 +2,7 @@
import json
import sys
from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, MutableSequence, Sequence
from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence
from datetime import datetime, timezone
from itertools import chain
from typing import Any, Generic, Literal
@@ -18,14 +18,22 @@ from pydantic import BaseModel, ValidationError
from .._clients import BaseChatClient
from .._logging import get_logger
from .._middleware import use_chat_middleware
from .._tools import FunctionTool, HostedWebSearchTool, ToolProtocol, use_function_invocation
from .._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer
from .._tools import (
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
HostedWebSearchTool,
ToolProtocol,
)
from .._types import (
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
FinishReason,
ResponseStream,
UsageDetails,
prepare_function_call_results,
)
@@ -34,7 +42,7 @@ from ..exceptions import (
ServiceInvalidRequestError,
ServiceResponseException,
)
from ..observability import use_instrumentation
from ..observability import ChatTelemetryLayer
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings
@@ -124,74 +132,91 @@ OPTION_TRANSLATIONS: dict[str, str] = {
# region Base Client
class OpenAIBaseChatClient(OpenAIBase, BaseChatClient[TOpenAIChatOptions], Generic[TOpenAIChatOptions]):
"""OpenAI Chat completion class."""
class RawOpenAIChatClient( # type: ignore[misc]
OpenAIBase,
BaseChatClient[TOpenAIChatOptions],
Generic[TOpenAIChatOptions],
):
"""Raw OpenAI Chat completion class without middleware, telemetry, or function invocation.
Warning:
**This class should not normally be used directly.** It does not include middleware,
telemetry, or function invocation support that you most likely need. If you do use it,
you should consider which additional layers to apply. There is a defined ordering that
you should follow:
1. **ChatMiddlewareLayer** - Should be applied first as it also prepares function middleware
2. **FunctionInvocationLayer** - Handles tool/function calling loop
3. **ChatTelemetryLayer** - Must be inside the function calling loop for correct per-call telemetry
Use ``OpenAIChatClient`` instead for a fully-featured client with all layers applied.
"""
@override
async def _inner_get_response(
def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
) -> ChatResponse:
client = await self._ensure_client()
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
# prepare
options_dict = self._prepare_options(messages, options)
try:
# execute and process
return self._parse_response_from_openai(
await client.chat.completions.create(stream=False, **options_dict), options
)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error: {ex}",
inner_exception=ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
@override
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
client = await self._ensure_client()
# prepare
options_dict = self._prepare_options(messages, options)
options_dict["stream_options"] = {"include_usage": True}
try:
# execute and process
async for chunk in await client.chat.completions.create(stream=True, **options_dict):
if len(chunk.choices) == 0 and chunk.usage is None:
continue
yield self._parse_response_update_from_openai(chunk)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error: {ex}",
if stream:
# Streaming mode
options_dict["stream_options"] = {"include_usage": True}
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
client = await self._ensure_client()
try:
async for chunk in await client.chat.completions.create(stream=True, **options_dict):
if len(chunk.choices) == 0 and chunk.usage is None:
continue
yield self._parse_response_update_from_openai(chunk)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error: {ex}",
inner_exception=ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
return self._build_response_stream(_stream(), response_format=options.get("response_format"))
# Non-streaming mode
async def _get_response() -> ChatResponse:
client = await self._ensure_client()
try:
return self._parse_response_from_openai(
await client.chat.completions.create(stream=False, **options_dict), options
)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error: {ex}",
inner_exception=ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
return _get_response()
# region content creation
@@ -217,7 +242,7 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient[TOpenAIChatOptions], Gener
case _:
logger.debug("Unsupported tool passed (type: %s), ignoring", type(tool))
else:
chat_tools.append(tool if isinstance(tool, dict) else dict(tool))
chat_tools.append(tool) # type: ignore[arg-type]
ret_dict: dict[str, Any] = {}
if chat_tools:
ret_dict["tools"] = chat_tools
@@ -225,7 +250,7 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient[TOpenAIChatOptions], Gener
ret_dict["web_search_options"] = web_search_options
return ret_dict
def _prepare_options(self, messages: MutableSequence[ChatMessage], options: dict[str, Any]) -> dict[str, Any]:
def _prepare_options(self, messages: Sequence[ChatMessage], options: Mapping[str, Any]) -> dict[str, Any]:
# Prepend instructions from options if they exist
from .._types import prepend_instructions_to_messages, validate_tool_mode
@@ -256,10 +281,11 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient[TOpenAIChatOptions], Gener
tools = options.get("tools")
if tools is not None:
run_options.update(self._prepare_tools_for_openai(tools))
# Only include tool_choice and parallel_tool_calls if tools are present
if not run_options.get("tools"):
run_options.pop("parallel_tool_calls", None)
run_options.pop("tool_choice", None)
if tool_choice := run_options.pop("tool_choice", None):
elif tool_choice := run_options.pop("tool_choice", None):
tool_mode = validate_tool_mode(tool_choice)
if (mode := tool_mode.get("mode")) == "required" and (
func_name := tool_mode.get("required_function_name")
@@ -279,15 +305,15 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient[TOpenAIChatOptions], Gener
run_options["response_format"] = type_to_response_format_param(response_format)
return run_options
def _parse_response_from_openai(self, response: ChatCompletion, options: dict[str, Any]) -> "ChatResponse":
def _parse_response_from_openai(self, response: ChatCompletion, options: Mapping[str, Any]) -> "ChatResponse":
"""Parse a response from OpenAI into a ChatResponse."""
response_metadata = self._get_metadata_from_chat_response(response)
messages: list[ChatMessage] = []
finish_reason: str | None = None
finish_reason: FinishReason | None = None
for choice in response.choices:
response_metadata.update(self._get_metadata_from_chat_choice(choice))
if choice.finish_reason:
finish_reason = choice.finish_reason
finish_reason = choice.finish_reason # type: ignore[assignment]
contents: list[Content] = []
if text_content := self._parse_text_from_openai(choice):
contents.append(text_content)
@@ -295,7 +321,7 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient[TOpenAIChatOptions], Gener
contents.extend(parsed_tool_calls)
if reasoning_details := getattr(choice.message, "reasoning_details", None):
contents.append(Content.from_text_reasoning(protected_data=json.dumps(reasoning_details)))
messages.append(ChatMessage("assistant", contents))
messages.append(ChatMessage(role="assistant", contents=contents))
return ChatResponse(
response_id=response.id,
created_at=datetime.fromtimestamp(response.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
@@ -327,12 +353,12 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient[TOpenAIChatOptions], Gener
message_id=chunk.id,
)
contents: list[Content] = []
finish_reason: str | None = None
finish_reason: FinishReason | None = None
for choice in chunk.choices:
chunk_metadata.update(self._get_metadata_from_chat_choice(choice))
contents.extend(self._parse_tool_calls_from_openai(choice))
if choice.finish_reason:
finish_reason = choice.finish_reason
finish_reason = choice.finish_reason # type: ignore[assignment]
if text_content := self._parse_text_from_openai(choice):
contents.append(text_content)
@@ -563,11 +589,15 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient[TOpenAIChatOptions], Gener
# region Public client
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient[TOpenAIChatOptions], Generic[TOpenAIChatOptions]):
"""OpenAI Chat completion class."""
class OpenAIChatClient( # type: ignore[misc]
OpenAIConfigMixin,
ChatMiddlewareLayer[TOpenAIChatOptions],
FunctionInvocationLayer[TOpenAIChatOptions],
ChatTelemetryLayer[TOpenAIChatOptions],
RawOpenAIChatClient[TOpenAIChatOptions],
Generic[TOpenAIChatOptions],
):
"""OpenAI Chat completion class with middleware, telemetry, and function invocation support."""
def __init__(
self,
@@ -579,6 +609,8 @@ class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient[TOpenAIChatOption
async_client: AsyncOpenAI | None = None,
instruction_role: str | None = None,
base_url: str | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
@@ -599,6 +631,8 @@ class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient[TOpenAIChatOption
base_url: The base URL to use. If provided will override
the standard value for an OpenAI connector, the env vars or .env file value.
Can also be set via environment variable OPENAI_BASE_URL.
middleware: Optional sequence of ChatAndFunctionMiddlewareTypes to apply to requests.
function_invocation_configuration: Optional configuration for function invocation support.
env_file_path: Use the environment settings file as a fallback
to environment variables.
env_file_encoding: The encoding of the environment settings file.
@@ -661,4 +695,6 @@ class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient[TOpenAIChatOption
default_headers=default_headers,
client=async_client,
instruction_role=instruction_role,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
)
python/packages/core/agent_framework/openai/_responses_client.py
+118 -89
View File
@@ -7,12 +7,11 @@ from collections.abc import (
Callable,
Mapping,
MutableMapping,
MutableSequence,
Sequence,
)
from datetime import datetime, timezone
from itertools import chain
from typing import Any, Generic, Literal, cast
from typing import TYPE_CHECKING, Any, Generic, Literal, NoReturn, TypedDict, cast
from openai import AsyncOpenAI, BadRequestError
from openai.types.responses.file_search_tool_param import FileSearchToolParam
@@ -36,8 +35,10 @@ from pydantic import BaseModel, ValidationError
from .._clients import BaseChatClient
from .._logging import get_logger
from .._middleware import use_chat_middleware
from .._middleware import ChatMiddlewareLayer
from .._tools import (
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
HostedCodeInterpreterTool,
HostedFileSearchTool,
@@ -45,7 +46,6 @@ from .._tools import (
HostedMCPTool,
HostedWebSearchTool,
ToolProtocol,
use_function_invocation,
)
from .._types import (
Annotation,
@@ -54,6 +54,8 @@ from .._types import (
ChatResponse,
ChatResponseUpdate,
Content,
ResponseStream,
Role,
TextSpanRegion,
UsageDetails,
detect_media_type_from_base64,
@@ -66,7 +68,7 @@ from ..exceptions import (
ServiceInvalidRequestError,
ServiceResponseException,
)
from ..observability import use_instrumentation
from ..observability import ChatTelemetryLayer
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings
@@ -83,10 +85,18 @@ if sys.version_info >= (3, 11):
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from .._middleware import (
ChatMiddleware,
ChatMiddlewareCallable,
FunctionMiddleware,
FunctionMiddlewareCallable,
)
logger = get_logger("agent_framework.openai")
__all__ = ["OpenAIResponsesClient", "OpenAIResponsesOptions"]
__all__ = ["OpenAIResponsesClient", "OpenAIResponsesOptions", "RawOpenAIResponsesClient"]
# region OpenAI Responses Options TypedDict
@@ -193,95 +203,105 @@ TOpenAIResponsesOptions = TypeVar(
# region ResponsesClient
class OpenAIBaseResponsesClient(
class RawOpenAIResponsesClient( # type: ignore[misc]
OpenAIBase,
BaseChatClient[TOpenAIResponsesOptions],
Generic[TOpenAIResponsesOptions],
):
"""Base class for all OpenAI Responses based API's."""
"""Raw OpenAI Responses client without middleware, telemetry, or function invocation.
Warning:
**This class should not normally be used directly.** It does not include middleware,
telemetry, or function invocation support that you most likely need. If you do use it,
you should consider which additional layers to apply. There is a defined ordering that
you should follow:
1. **ChatMiddlewareLayer** - Should be applied first as it also prepares function middleware
2. **FunctionInvocationLayer** - Handles tool/function calling loop
3. **ChatTelemetryLayer** - Must be inside the function calling loop for correct per-call telemetry
Use ``OpenAIResponsesClient`` instead for a fully-featured client with all layers applied.
"""
FILE_SEARCH_MAX_RESULTS: int = 50
# region Inner Methods
@override
async def _inner_get_response(
async def _prepare_request(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
**kwargs: Any,
) -> ChatResponse:
) -> tuple[AsyncOpenAI, dict[str, Any], dict[str, Any]]:
"""Validate options and prepare the request.
Returns:
Tuple of (client, run_options, validated_options).
"""
client = await self._ensure_client()
# prepare
run_options = await self._prepare_options(messages, options, **kwargs)
try:
# execute and process
if "text_format" in run_options:
response = await client.responses.parse(stream=False, **run_options)
else:
response = await client.responses.create(stream=False, **run_options)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error: {ex}",
inner_exception=ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
validated_options = await self._validate_options(options)
run_options = await self._prepare_options(messages, validated_options, **kwargs)
return client, run_options, validated_options
def _handle_request_error(self, ex: Exception) -> NoReturn:
"""Convert exceptions to appropriate service exceptions. Always raises."""
if isinstance(ex, BadRequestError) and ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error: {ex}",
inner_exception=ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
return self._parse_response_from_openai(response, options=options)
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
@override
async def _inner_get_streaming_response(
def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
client = await self._ensure_client()
# prepare
run_options = await self._prepare_options(messages, options, **kwargs)
function_call_ids: dict[int, tuple[str, str]] = {} # output_index: (call_id, name)
try:
# execute and process
if "text_format" not in run_options:
async for chunk in await client.responses.create(stream=True, **run_options):
yield self._parse_chunk_from_openai(
chunk,
options=options,
function_call_ids=function_call_ids,
)
return
async with client.responses.stream(**run_options) as response:
async for chunk in response:
yield self._parse_chunk_from_openai(
chunk,
options=options,
function_call_ids=function_call_ids,
)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error: {ex}",
inner_exception=ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
if stream:
function_call_ids: dict[int, tuple[str, str]] = {}
validated_options: dict[str, Any] | None = None
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
nonlocal validated_options
client, run_options, validated_options = await self._prepare_request(messages, options, **kwargs)
try:
if "text_format" in run_options:
async with client.responses.stream(**run_options) as response:
async for chunk in response:
yield self._parse_chunk_from_openai(
chunk, options=validated_options, function_call_ids=function_call_ids
)
else:
async for chunk in await client.responses.create(stream=True, **run_options):
yield self._parse_chunk_from_openai(
chunk, options=validated_options, function_call_ids=function_call_ids
)
except Exception as ex:
self._handle_request_error(ex)
response_format = validated_options.get("response_format") if validated_options else None
return self._build_response_stream(_stream(), response_format=response_format)
# Non-streaming
async def _get_response() -> ChatResponse:
client, run_options, validated_options = await self._prepare_request(messages, options, **kwargs)
try:
if "text_format" in run_options:
response = await client.responses.parse(stream=False, **run_options)
else:
response = await client.responses.create(stream=False, **run_options)
except Exception as ex:
self._handle_request_error(ex)
return self._parse_response_from_openai(response, options=validated_options)
return _get_response()
def _prepare_response_and_text_format(
self,
@@ -499,8 +519,8 @@ class OpenAIBaseResponsesClient(
async def _prepare_options(
self,
messages: MutableSequence[ChatMessage],
options: dict[str, Any],
messages: Sequence[ChatMessage],
options: Mapping[str, Any],
**kwargs: Any,
) -> dict[str, Any]:
"""Take options dict and create the specific options for Responses API."""
@@ -596,7 +616,7 @@ class OpenAIBaseResponsesClient(
raise ValueError("model_id must be a non-empty string")
options["model"] = self.model_id
def _get_current_conversation_id(self, options: dict[str, Any], **kwargs: Any) -> str | None:
def _get_current_conversation_id(self, options: Mapping[str, Any], **kwargs: Any) -> str | None:
"""Get the current conversation ID, preferring kwargs over options.
This ensures runtime-updated conversation IDs (for example, from tool execution
@@ -651,10 +671,10 @@ class OpenAIBaseResponsesClient(
continue
case "function_result":
new_args: dict[str, Any] = {}
new_args.update(self._prepare_content_for_openai(message.role, content, call_id_to_id))
new_args.update(self._prepare_content_for_openai(message.role, content, call_id_to_id)) # type: ignore[arg-type]
all_messages.append(new_args)
case "function_call":
function_call = self._prepare_content_for_openai(message.role, content, call_id_to_id)
function_call = self._prepare_content_for_openai(message.role, content, call_id_to_id) # type: ignore[arg-type]
all_messages.append(function_call) # type: ignore
case "function_approval_response" | "function_approval_request":
all_messages.append(self._prepare_content_for_openai(message.role, content, call_id_to_id)) # type: ignore
@@ -668,7 +688,7 @@ class OpenAIBaseResponsesClient(
def _prepare_content_for_openai(
self,
role: str,
role: Role,
content: Content,
call_id_to_id: dict[str, str],
) -> dict[str, Any]:
@@ -1026,7 +1046,7 @@ class OpenAIBaseResponsesClient(
)
case _:
logger.debug("Unparsed output of type: %s: %s", item.type, item)
response_message = ChatMessage("assistant", contents)
response_message = ChatMessage(role="assistant", contents=contents)
args: dict[str, Any] = {
"response_id": response.id,
"created_at": datetime.fromtimestamp(response.created_at, tz=timezone.utc).strftime(
@@ -1413,15 +1433,15 @@ class OpenAIBaseResponsesClient(
return {}
@use_function_invocation
@use_instrumentation
@use_chat_middleware
class OpenAIResponsesClient(
class OpenAIResponsesClient( # type: ignore[misc]
OpenAIConfigMixin,
OpenAIBaseResponsesClient[TOpenAIResponsesOptions],
ChatMiddlewareLayer[TOpenAIResponsesOptions],
FunctionInvocationLayer[TOpenAIResponsesOptions],
ChatTelemetryLayer[TOpenAIResponsesOptions],
RawOpenAIResponsesClient[TOpenAIResponsesOptions],
Generic[TOpenAIResponsesOptions],
):
"""OpenAI Responses client class."""
"""OpenAI Responses client class with middleware, telemetry, and function invocation support."""
def __init__(
self,
@@ -1435,6 +1455,10 @@ class OpenAIResponsesClient(
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
middleware: (
Sequence["ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable"] | None
) = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
"""Initialize an OpenAI Responses client.
@@ -1456,6 +1480,8 @@ class OpenAIResponsesClient(
env_file_path: Use the environment settings file as a fallback
to environment variables.
env_file_encoding: The encoding of the environment settings file.
middleware: Optional middleware to apply to the client.
function_invocation_configuration: Optional function invocation configuration override.
kwargs: Other keyword parameters.
Examples:
@@ -1516,4 +1542,7 @@ class OpenAIResponsesClient(
client=async_client,
instruction_role=instruction_role,
base_url=openai_settings.base_url,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
python/packages/core/agent_framework/openai/_shared.py
+6 -3
View File
@@ -138,11 +138,12 @@ class OpenAIBase(SerializationMixin):
if model_id:
self.model_id = model_id.strip()
# Call super().__init__() to continue MRO chain (e.g., BaseChatClient)
# Call super().__init__() to continue MRO chain (e.g., RawChatClient)
# Extract known kwargs that belong to other base classes
additional_properties = kwargs.pop("additional_properties", None)
middleware = kwargs.pop("middleware", None)
instruction_role = kwargs.pop("instruction_role", None)
function_invocation_configuration = kwargs.pop("function_invocation_configuration", None)
# Build super().__init__() args
super_kwargs = {}
@@ -150,6 +151,8 @@ class OpenAIBase(SerializationMixin):
super_kwargs["additional_properties"] = additional_properties
if middleware is not None:
super_kwargs["middleware"] = middleware
if function_invocation_configuration is not None:
super_kwargs["function_invocation_configuration"] = function_invocation_configuration
# Call super().__init__() with filtered kwargs
super().__init__(**super_kwargs)
@@ -273,8 +276,8 @@ class OpenAIConfigMixin(OpenAIBase):
if instruction_role:
args["instruction_role"] = instruction_role
# Ensure additional_properties and middleware are passed through kwargs to BaseChatClient
# These are consumed by BaseChatClient.__init__ via kwargs
# Ensure additional_properties and middleware are passed through kwargs to RawChatClient
# These are consumed by RawChatClient.__init__ via kwargs
super().__init__(**args, **kwargs)

Some files were not shown because too many files have changed in this diff Show More