agent/agent-framework
1
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: Fix Python pyright package scoping and typing remediation (#4426)

Browse Source 
* Fix Python pyright package scoping and typing remediation

Implements issue #4407 by removing the root pyright include, adding package-level pyright includes, and resolving pyright/mypy typing issues across Python packages. Also cleans unnecessary casts and applies line-level, rule-specific ignores where external libraries are too dynamic.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Reduce pyright cost in handoff cloning

Simplify cloned_options construction in HandoffAgentExecutor to avoid expensive TypedDict narrowing/inference in _handoff.py, which was causing pyright to spend a long time in orchestrations.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix types

* Fix lint and type-check regressions

Resolve current Python package check failures across lint, pyright, and mypy after recent code changes, including purview/declarative pyright issues and multiple ruff simplification findings.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fixed hooks

* Stabilize package tests and test tasks

Resolve cross-package non-integration test failures, simplify streaming type flow, harden locale/culture handling, and standardize package test poe tasks to exclude integration tests where applicable.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* lots of small fixes

* Fix current Python test regressions

Address current failing unit tests in azure-ai, bedrock, and azure-cosmos while keeping Bedrock parsing logic inline (no new static helper methods).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* small fixes

* small fixes

* removed pydantic from json

* final updates

* fix core

* fix tests

* fix obser

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2026-03-05 16:32:24 +01:00
committed by GitHub
Unverified
parent 4a043c6c66
commit 55ddd841b7
122 changed files with 2328 additions and 2407 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/tests/core/test_skills.py
+3 -5
View File
@@ -10,7 +10,7 @@ from unittest.mock import AsyncMock
import pytest
from agent_framework import Skill, SkillResource, SkillsProvider, SessionContext
from agent_framework import SessionContext, Skill, SkillResource, SkillsProvider
from agent_framework._skills import (
DEFAULT_RESOURCE_EXTENSIONS,
_create_instructions,
@@ -1348,9 +1348,7 @@ class TestReadAndParseSkillFile:
def test_valid_file(self, tmp_path: Path) -> None:
skill_dir = tmp_path / "my-skill"
skill_dir.mkdir()
(skill_dir / "SKILL.md").write_text(
"---\nname: my-skill\ndescription: A skill.\n---\nBody.", encoding="utf-8"
)
(skill_dir / "SKILL.md").write_text("---\nname: my-skill\ndescription: A skill.\n---\nBody.", encoding="utf-8")
result = _read_and_parse_skill_file(str(skill_dir))
assert result is not None
name, desc, content = result
@@ -1393,7 +1391,7 @@ class TestCreateResourceElement:
def test_xml_escapes_name(self) -> None:
r = SkillResource(name='ref"special', content="data")
elem = _create_resource_element(r)
assert '&quot;' in elem
assert "&quot;" in elem
def test_xml_escapes_description(self) -> None:
r = SkillResource(name="ref", description='Uses <tags> & "quotes"', content="data")
python/packages/core/tests/core/test_tools.py
+1 -465
View File
@@ -5,7 +5,7 @@ from unittest.mock import Mock
import pytest
from opentelemetry import trace
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from pydantic import BaseModel, ValidationError
from pydantic import BaseModel
from agent_framework import (
Content,
@@ -13,7 +13,6 @@ from agent_framework import (
tool,
)
from agent_framework._tools import (
_build_pydantic_model_from_json_schema,
_parse_annotation,
_parse_inputs,
)
@@ -1001,467 +1000,4 @@ def test_parse_annotation_with_annotated_and_literal():
assert get_args(literal_type) == ("A", "B", "C")
def test_build_pydantic_model_from_json_schema_array_of_objects_issue():
"""Test for Tools with complex input schema (array of objects).
This test verifies that JSON schemas with array properties containing nested objects
are properly parsed, ensuring that the nested object schema is preserved
and not reduced to a bare dict.
Example from issue:
```
const SalesOrderItemSchema = z.object({
customerMaterialNumber: z.string().optional(),
quantity: z.number(),
unitOfMeasure: z.string()
});
const CreateSalesOrderInputSchema = z.object({
contract: z.string(),
items: z.array(SalesOrderItemSchema)
});
```
The issue was that agents only saw:
```
{"contract": "str", "items": "list[dict]"}
```
Instead of the proper nested schema with all fields.
"""
# Schema matching the issue description
schema = {
"type": "object",
"properties": {
"contract": {"type": "string", "description": "Reference contract number"},
"items": {
"type": "array",
"description": "Sales order line items",
"items": {
"type": "object",
"properties": {
"customerMaterialNumber": {
"type": "string",
"description": "Customer's material number",
},
"quantity": {"type": "number", "description": "Order quantity"},
"unitOfMeasure": {
"type": "string",
"description": "Unit of measure (e.g., 'ST', 'KG', 'TO')",
},
},
"required": ["quantity", "unitOfMeasure"],
},
},
},
"required": ["contract", "items"],
}
model = _build_pydantic_model_from_json_schema("create_sales_order", schema)
# Test valid data
valid_data = {
"contract": "CONTRACT-123",
"items": [
{
"customerMaterialNumber": "MAT-001",
"quantity": 10,
"unitOfMeasure": "ST",
},
{"quantity": 5.5, "unitOfMeasure": "KG"},
],
}
instance = model(**valid_data)
# Verify the data was parsed correctly
assert instance.contract == "CONTRACT-123"
assert len(instance.items) == 2
# Verify first item
assert instance.items[0].customerMaterialNumber == "MAT-001"
assert instance.items[0].quantity == 10
assert instance.items[0].unitOfMeasure == "ST"
# Verify second item (optional field not provided)
assert instance.items[1].quantity == 5.5
assert instance.items[1].unitOfMeasure == "KG"
# Verify that items are proper BaseModel instances, not bare dicts
assert isinstance(instance.items[0], BaseModel)
assert isinstance(instance.items[1], BaseModel)
# Verify that the nested object has the expected fields
assert hasattr(instance.items[0], "customerMaterialNumber")
assert hasattr(instance.items[0], "quantity")
assert hasattr(instance.items[0], "unitOfMeasure")
# CRITICAL: Validate using the same methods that actual chat clients use
# This is what would actually be sent to the LLM
# Create a FunctionTool wrapper to access the client-facing APIs
def dummy_func(**kwargs):
return kwargs
test_func = FunctionTool(
func=dummy_func,
name="create_sales_order",
description="Create a sales order",
input_model=model,
)
# Test 1: Anthropic client uses tool.parameters() directly
anthropic_schema = test_func.parameters()
# Verify contract property
assert "contract" in anthropic_schema["properties"]
assert anthropic_schema["properties"]["contract"]["type"] == "string"
# Verify items array property exists
assert "items" in anthropic_schema["properties"]
items_prop = anthropic_schema["properties"]["items"]
assert items_prop["type"] == "array"
# THE KEY TEST for Anthropic: array items must have proper object schema
assert "items" in items_prop, "Array should have 'items' schema definition"
array_items_schema = items_prop["items"]
# Resolve schema if using $ref
if "$ref" in array_items_schema:
ref_path = array_items_schema["$ref"]
assert ref_path.startswith("#/$defs/") or ref_path.startswith("#/definitions/")
ref_name = ref_path.split("/")[-1]
defs = anthropic_schema.get("$defs", anthropic_schema.get("definitions", {}))
assert ref_name in defs, f"Referenced schema '{ref_name}' should exist"
item_schema = defs[ref_name]
else:
item_schema = array_items_schema
# Verify the nested object has all properties defined
assert "properties" in item_schema, "Array items should have properties (not bare dict)"
item_properties = item_schema["properties"]
# All three fields must be present in schema sent to LLM
assert "customerMaterialNumber" in item_properties, "customerMaterialNumber missing from LLM schema"
assert "quantity" in item_properties, "quantity missing from LLM schema"
assert "unitOfMeasure" in item_properties, "unitOfMeasure missing from LLM schema"
# Verify types are correct
assert item_properties["customerMaterialNumber"]["type"] == "string"
assert item_properties["quantity"]["type"] in ["number", "integer"]
assert item_properties["unitOfMeasure"]["type"] == "string"
# Test 2: OpenAI client uses tool.to_json_schema_spec()
openai_spec = test_func.to_json_schema_spec()
assert openai_spec["type"] == "function"
assert "function" in openai_spec
openai_schema = openai_spec["function"]["parameters"]
# Verify the same structure is present in OpenAI format
assert "items" in openai_schema["properties"]
openai_items_prop = openai_schema["properties"]["items"]
assert openai_items_prop["type"] == "array"
assert "items" in openai_items_prop
openai_array_items = openai_items_prop["items"]
if "$ref" in openai_array_items:
ref_path = openai_array_items["$ref"]
ref_name = ref_path.split("/")[-1]
defs = openai_schema.get("$defs", openai_schema.get("definitions", {}))
openai_item_schema = defs[ref_name]
else:
openai_item_schema = openai_array_items
assert "properties" in openai_item_schema
openai_props = openai_item_schema["properties"]
assert "customerMaterialNumber" in openai_props
assert "quantity" in openai_props
assert "unitOfMeasure" in openai_props
# Test validation - missing required quantity
with pytest.raises(ValidationError):
model(
contract="CONTRACT-456",
items=[
{
"customerMaterialNumber": "MAT-002",
"unitOfMeasure": "TO",
# Missing required 'quantity'
}
],
)
# Test validation - missing required unitOfMeasure
with pytest.raises(ValidationError):
model(
contract="CONTRACT-789",
items=[
{
"quantity": 20
# Missing required 'unitOfMeasure'
}
],
)
def test_one_of_discriminator_polymorphism():
"""Test that oneOf with discriminator creates proper polymorphic union types.
Tests that oneOf + discriminator patterns are properly converted to Pydantic discriminated unions.
"""
schema = {
"$defs": {
"CreateProject": {
"description": "Action: Create an Azure DevOps project.",
"properties": {
"name": {
"const": "create_project",
"default": "create_project",
"type": "string",
},
"params": {"$ref": "#/$defs/CreateProjectParams"},
},
"required": ["params"],
"type": "object",
},
"CreateProjectParams": {
"description": "Parameters for the create_project action.",
"properties": {
"orgUrl": {"minLength": 1, "type": "string"},
"projectName": {"minLength": 1, "type": "string"},
"description": {"default": "", "type": "string"},
"template": {"default": "Agile", "type": "string"},
"sourceControl": {
"default": "Git",
"enum": ["Git", "Tfvc"],
"type": "string",
},
"visibility": {"default": "private", "type": "string"},
},
"required": ["orgUrl", "projectName"],
"type": "object",
},
"DeployRequest": {
"description": "Request to deploy Azure DevOps resources.",
"properties": {
"projectName": {"minLength": 1, "type": "string"},
"organization": {"minLength": 1, "type": "string"},
"actions": {
"items": {
"discriminator": {
"mapping": {
"create_project": "#/$defs/CreateProject",
"hello_world": "#/$defs/HelloWorld",
},
"propertyName": "name",
},
"oneOf": [
{"$ref": "#/$defs/HelloWorld"},
{"$ref": "#/$defs/CreateProject"},
],
},
"type": "array",
},
},
"required": ["projectName", "organization"],
"type": "object",
},
"HelloWorld": {
"description": "Action: Prints a greeting message.",
"properties": {
"name": {
"const": "hello_world",
"default": "hello_world",
"type": "string",
},
"params": {"$ref": "#/$defs/HelloWorldParams"},
},
"required": ["params"],
"type": "object",
},
"HelloWorldParams": {
"description": "Parameters for the hello_world action.",
"properties": {
"name": {
"description": "Name to greet",
"minLength": 1,
"type": "string",
}
},
"required": ["name"],
"type": "object",
},
},
"properties": {"params": {"$ref": "#/$defs/DeployRequest"}},
"required": ["params"],
"type": "object",
}
# Build the model
model = _build_pydantic_model_from_json_schema("deploy_tool", schema)
# Verify the model structure
assert model is not None
assert issubclass(model, BaseModel)
# Test with HelloWorld action
hello_world_data = {
"params": {
"projectName": "MyProject",
"organization": "MyOrg",
"actions": [
{
"name": "hello_world",
"params": {"name": "Alice"},
}
],
}
}
instance = model(**hello_world_data)
assert instance.params.projectName == "MyProject"
assert instance.params.organization == "MyOrg"
assert len(instance.params.actions) == 1
assert instance.params.actions[0].name == "hello_world"
assert instance.params.actions[0].params.name == "Alice"
# Test with CreateProject action
create_project_data = {
"params": {
"projectName": "MyProject",
"organization": "MyOrg",
"actions": [
{
"name": "create_project",
"params": {
"orgUrl": "https://dev.azure.com/myorg",
"projectName": "NewProject",
"sourceControl": "Git",
},
}
],
}
}
instance2 = model(**create_project_data)
assert instance2.params.actions[0].name == "create_project"
assert instance2.params.actions[0].params.projectName == "NewProject"
assert instance2.params.actions[0].params.sourceControl == "Git"
# Test with mixed actions
mixed_data = {
"params": {
"projectName": "MyProject",
"organization": "MyOrg",
"actions": [
{"name": "hello_world", "params": {"name": "Bob"}},
{
"name": "create_project",
"params": {
"orgUrl": "https://dev.azure.com/myorg",
"projectName": "AnotherProject",
},
},
],
}
}
instance3 = model(**mixed_data)
assert len(instance3.params.actions) == 2
assert instance3.params.actions[0].name == "hello_world"
assert instance3.params.actions[1].name == "create_project"
def test_const_creates_literal():
"""Test that const in JSON Schema creates Literal type."""
schema = {
"properties": {
"action": {
"const": "create",
"type": "string",
"description": "Action type",
},
"value": {"type": "integer"},
},
"required": ["action", "value"],
}
model = _build_pydantic_model_from_json_schema("test_const", schema)
# Verify valid const value works
instance = model(action="create", value=42)
assert instance.action == "create"
assert instance.value == 42
# Verify incorrect const value fails
with pytest.raises(ValidationError):
model(action="delete", value=42)
def test_enum_creates_literal():
"""Test that enum in JSON Schema creates Literal type."""
schema = {
"properties": {
"status": {
"enum": ["pending", "approved", "rejected"],
"type": "string",
"description": "Status",
},
"priority": {"enum": [1, 2, 3], "type": "integer"},
},
"required": ["status"],
}
model = _build_pydantic_model_from_json_schema("test_enum", schema)
# Verify valid enum values work
instance = model(status="approved", priority=2)
assert instance.status == "approved"
assert instance.priority == 2
# Verify invalid enum value fails
with pytest.raises(ValidationError):
model(status="unknown")
with pytest.raises(ValidationError):
model(status="pending", priority=5)
def test_nested_object_with_const_and_enum():
"""Test that const and enum work in nested objects."""
schema = {
"properties": {
"config": {
"type": "object",
"properties": {
"type": {
"const": "production",
"default": "production",
"type": "string",
},
"level": {"enum": ["low", "medium", "high"], "type": "string"},
},
"required": ["level"],
}
},
"required": ["config"],
}
model = _build_pydantic_model_from_json_schema("test_nested", schema)
# Valid data
instance = model(config={"type": "production", "level": "high"})
assert instance.config.type == "production"
assert instance.config.level == "high"
# Invalid const in nested object
with pytest.raises(ValidationError):
model(config={"type": "development", "level": "low"})
# Invalid enum in nested object
with pytest.raises(ValidationError):
model(config={"type": "production", "level": "critical"})
# endregion
python/packages/core/tests/core/test_types.py
+14 -6
View File
@@ -550,7 +550,6 @@ def test_usage_details():
assert usage["input_token_count"] == 5
assert usage["output_token_count"] == 10
assert usage["total_token_count"] == 15
assert usage.get("additional_counts", {}) == {}
def test_usage_details_addition():
@@ -581,8 +580,8 @@ def test_usage_details_addition():
def test_usage_details_fail():
# TypedDict doesn't validate types at runtime, so this test no longer applies
# Creating UsageDetails with wrong types won't raise ValueError
usage = UsageDetails(input_token_count=5, output_token_count=10, total_token_count=15, wrong_type="42.923") # type: ignore[typeddict-item]
assert usage["wrong_type"] == "42.923" # type: ignore[typeddict-item]
usage = UsageDetails(input_token_count=5, output_token_count=10, total_token_count=15, wrong_type="42.923")
assert usage["wrong_type"] == "42.923"
def test_usage_details_additional_counts():
@@ -601,6 +600,15 @@ def test_usage_details_add_with_none_and_type_errors():
# TypedDict doesn't support + operator, use add_usage_details
def test_usage_details_add_skips_non_int():
u1 = UsageDetails(input_token_count=10, other="test")
u2 = UsageDetails(input_token_count=10, another="test")
u3 = add_usage_details(u1, u2)
assert len(u3.keys()) == 1
assert "input_token_count" in u3
assert u3["input_token_count"] == 20
# region UserInputRequest and Response
@@ -1705,7 +1713,7 @@ def test_chat_response_complex_serialization():
{"role": "user", "contents": [{"type": "text", "text": "Hello"}]},
{"role": "assistant", "contents": [{"type": "text", "text": "Hi there"}]},
],
"finish_reason": {"value": "stop"},
"finish_reason": "stop",
"usage_details": {
"type": "usage_details",
"input_token_count": 5,
@@ -1831,7 +1839,7 @@ def test_agent_run_response_update_all_content_types():
},
{"type": "text_reasoning", "text": "reasoning"},
],
"role": {"value": "assistant"}, # Test role as dict
"role": "assistant", # Test role as dict
}
update = AgentResponseUpdate.from_dict(update_data)
@@ -2394,7 +2402,7 @@ def test_content_add_usage_content_non_integer_values():
result = usage1 + usage2
# Non-integer "model" should take first non-None value
assert result.usage_details["model"] == "gpt-4"
assert "model" not in result.usage_details
# Integer "count" should be summed
assert result.usage_details["count"] == 30
python/packages/core/tests/openai/test_openai_embedding_client.py
+8 -1
View File
@@ -212,7 +212,8 @@ def test_azure_construction_with_existing_client() -> None:
assert client.client is mock_client
def test_azure_construction_missing_deployment_name_raises() -> None:
def test_azure_construction_missing_deployment_name_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", raising=False)
with pytest.raises(ValueError, match="deployment name is required"):
AzureOpenAIEmbeddingClient(
api_key="test-key",
@@ -272,6 +273,7 @@ skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
@skip_if_openai_integration_tests_disabled
@pytest.mark.flaky
@pytest.mark.integration
async def test_integration_openai_get_embeddings() -> None:
"""End-to-end test of OpenAI embedding generation."""
client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
@@ -289,6 +291,7 @@ async def test_integration_openai_get_embeddings() -> None:
@skip_if_openai_integration_tests_disabled
@pytest.mark.flaky
@pytest.mark.integration
async def test_integration_openai_get_embeddings_multiple() -> None:
"""Test embedding generation for multiple inputs."""
client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
@@ -302,6 +305,7 @@ async def test_integration_openai_get_embeddings_multiple() -> None:
@skip_if_openai_integration_tests_disabled
@pytest.mark.flaky
@pytest.mark.integration
async def test_integration_openai_get_embeddings_with_dimensions() -> None:
"""Test embedding generation with custom dimensions."""
client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
@@ -315,6 +319,7 @@ async def test_integration_openai_get_embeddings_with_dimensions() -> None:
@skip_if_azure_openai_integration_tests_disabled
@pytest.mark.flaky
@pytest.mark.integration
async def test_integration_azure_openai_get_embeddings() -> None:
"""End-to-end test of Azure OpenAI embedding generation."""
client = AzureOpenAIEmbeddingClient()
@@ -332,6 +337,7 @@ async def test_integration_azure_openai_get_embeddings() -> None:
@skip_if_azure_openai_integration_tests_disabled
@pytest.mark.flaky
@pytest.mark.integration
async def test_integration_azure_openai_get_embeddings_multiple() -> None:
"""Test Azure OpenAI embedding generation for multiple inputs."""
client = AzureOpenAIEmbeddingClient()
@@ -345,6 +351,7 @@ async def test_integration_azure_openai_get_embeddings_multiple() -> None:
@skip_if_azure_openai_integration_tests_disabled
@pytest.mark.flaky
@pytest.mark.integration
async def test_integration_azure_openai_get_embeddings_with_dimensions() -> None:
"""Test Azure OpenAI embedding generation with custom dimensions."""
client = AzureOpenAIEmbeddingClient()
python/packages/core/tests/workflow/test_agent_executor.py
+13 -38
View File
@@ -5,6 +5,7 @@ from collections.abc import AsyncIterable, Awaitable
from typing import TYPE_CHECKING, Any, Literal, overload
import pytest
from agent_framework import (
AgentExecutor,
AgentResponse,
@@ -59,30 +60,19 @@ class _CountingAgent(BaseAgent):
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> (
Awaitable[AgentResponse[Any]]
| ResponseStream[AgentResponseUpdate, AgentResponse[Any]]
):
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
self.call_count += 1
if stream:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(
contents=[
Content.from_text(
text=f"Response #{self.call_count}: {self.name}"
)
]
contents=[Content.from_text(text=f"Response #{self.call_count}: {self.name}")]
)
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
async def _run() -> AgentResponse:
return AgentResponse(
messages=[
Message("assistant", [f"Response #{self.call_count}: {self.name}"])
]
)
return AgentResponse(messages=[Message("assistant", [f"Response #{self.call_count}: {self.name}"])])
return _run()
@@ -120,10 +110,7 @@ class _StreamingHookAgent(BaseAgent):
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> (
Awaitable[AgentResponse[Any]]
| ResponseStream[AgentResponseUpdate, AgentResponse[Any]]
):
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
if stream:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
@@ -138,9 +125,9 @@ class _StreamingHookAgent(BaseAgent):
self.result_hook_called = True
return response
return ResponseStream(
_stream(), finalizer=AgentResponse.from_updates
).with_result_hook(_mark_result_hook_called)
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates).with_result_hook(
_mark_result_hook_called
)
async def _run() -> AgentResponse:
return AgentResponse(messages=[Message("assistant", ["hook test"])])
@@ -148,9 +135,7 @@ class _StreamingHookAgent(BaseAgent):
return _run()
async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks() -> (
None
):
async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks() -> None:
"""AgentExecutor should call get_final_response() so stream result hooks execute."""
agent = _StreamingHookAgent(id="hook_agent", name="HookAgent")
executor = AgentExecutor(agent, id="hook_exec")
@@ -217,9 +202,7 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
executor_state = executor_states[executor.id] # type: ignore[index]
assert "cache" in executor_state, "Checkpoint should store executor cache state"
assert "agent_session" in executor_state, (
"Checkpoint should store executor session state"
)
assert "agent_session" in executor_state, "Checkpoint should store executor session state"
# Verify session state structure
session_state = executor_state["agent_session"] # type: ignore[index]
@@ -240,15 +223,11 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
assert restored_agent.call_count == 0
# Build new workflow with the restored executor
wf_resume = SequentialBuilder(
participants=[restored_executor], checkpoint_storage=storage
).build()
wf_resume = SequentialBuilder(participants=[restored_executor], checkpoint_storage=storage).build()
# Resume from checkpoint
resumed_output: AgentExecutorResponse | None = None
async for ev in wf_resume.run(
checkpoint_id=restore_checkpoint.checkpoint_id, stream=True
):
async for ev in wf_resume.run(checkpoint_id=restore_checkpoint.checkpoint_id, stream=True):
if ev.type == "output":
resumed_output = ev.data # type: ignore[assignment]
if ev.type == "status" and ev.state in (
@@ -391,11 +370,7 @@ async def test_prepare_agent_run_args_strips_all_reserved_kwargs_at_once(
assert options is not None
assert options["additional_function_arguments"]["custom"] == 1
warned_keys = {
r.message.split("'")[1]
for r in caplog.records
if "reserved" in r.message.lower()
}
warned_keys = {r.message.split("'")[1] for r in caplog.records if "reserved" in r.message.lower()}
assert warned_keys == {"session", "stream", "messages"}
python/packages/core/tests/workflow/test_agent_utils.py
+24 -3
View File
@@ -16,10 +16,31 @@ class MockAgent:
self.description: str | None = None
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(self, messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def create_session(self, **kwargs: Any) -> AgentSession:
"""Creates a new conversation session for the agent."""
python/packages/core/tests/workflow/test_edge.py
+1 -2
View File
@@ -4,9 +4,8 @@ from dataclasses import dataclass
from typing import Any
from unittest.mock import patch
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
import pytest
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from agent_framework import (
Executor,
python/packages/core/tests/workflow/test_executor.py
+30 -97
View File
@@ -3,6 +3,8 @@
from dataclasses import dataclass
import pytest
from typing_extensions import Never
from agent_framework import (
Executor,
Message,
@@ -14,7 +16,6 @@ from agent_framework import (
handler,
response_handler,
)
from typing_extensions import Never
# Module-level types for string forward reference tests
@@ -155,11 +156,7 @@ async def test_executor_invoked_event_contains_input_data():
workflow = WorkflowBuilder(start_executor=upper).add_edge(upper, collector).build()
events = await workflow.run("hello world")
invoked_events = [
e
for e in events
if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
]
invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
assert len(invoked_events) == 2
@@ -193,16 +190,10 @@ async def test_executor_completed_event_contains_sent_messages():
sender = MultiSenderExecutor(id="sender")
collector = CollectorExecutor(id="collector")
workflow = (
WorkflowBuilder(start_executor=sender).add_edge(sender, collector).build()
)
workflow = WorkflowBuilder(start_executor=sender).add_edge(sender, collector).build()
events = await workflow.run("hello")
completed_events = [
e
for e in events
if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
]
completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
# Sender should have completed with the sent messages
sender_completed = next(e for e in completed_events if e.executor_id == "sender")
@@ -210,9 +201,7 @@ async def test_executor_completed_event_contains_sent_messages():
assert sender_completed.data == ["hello-first", "hello-second"]
# Collector should have completed with no sent messages (None)
collector_completed_events = [
e for e in completed_events if e.executor_id == "collector"
]
collector_completed_events = [e for e in completed_events if e.executor_id == "collector"]
# Collector is called twice (once per message from sender)
assert len(collector_completed_events) == 2
for collector_completed in collector_completed_events:
@@ -231,11 +220,7 @@ async def test_executor_completed_event_includes_yielded_outputs():
workflow = WorkflowBuilder(start_executor=executor).build()
events = await workflow.run("test")
completed_events = [
e
for e in events
if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
]
completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
assert len(completed_events) == 1
assert completed_events[0].executor_id == "yielder"
@@ -263,9 +248,7 @@ async def test_executor_events_with_complex_message_types():
class ProcessorExecutor(Executor):
@handler
async def handle(
self, request: Request, ctx: WorkflowContext[Response]
) -> None:
async def handle(self, request: Request, ctx: WorkflowContext[Response]) -> None:
response = Response(results=[request.query.upper()] * request.limit)
await ctx.send_message(response)
@@ -277,23 +260,13 @@ async def test_executor_events_with_complex_message_types():
processor = ProcessorExecutor(id="processor")
collector = CollectorExecutor(id="collector")
workflow = (
WorkflowBuilder(start_executor=processor).add_edge(processor, collector).build()
)
workflow = WorkflowBuilder(start_executor=processor).add_edge(processor, collector).build()
input_request = Request(query="hello", limit=3)
events = await workflow.run(input_request)
invoked_events = [
e
for e in events
if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
]
completed_events = [
e
for e in events
if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
]
invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
# Check processor invoked event has the Request object
processor_invoked = next(e for e in invoked_events if e.executor_id == "processor")
@@ -302,9 +275,7 @@ async def test_executor_events_with_complex_message_types():
assert processor_invoked.data.limit == 3
# Check processor completed event has the Response object
processor_completed = next(
e for e in completed_events if e.executor_id == "processor"
)
processor_completed = next(e for e in completed_events if e.executor_id == "processor")
assert processor_completed.data is not None
assert len(processor_completed.data) == 1
assert isinstance(processor_completed.data[0], Response)
@@ -390,9 +361,7 @@ def test_executor_workflow_output_types_property():
# Test executor with union workflow output types
class UnionWorkflowOutputExecutor(Executor):
@handler
async def handle(
self, text: str, ctx: WorkflowContext[int, str | bool]
) -> None:
async def handle(self, text: str, ctx: WorkflowContext[int, str | bool]) -> None:
pass
executor = UnionWorkflowOutputExecutor(id="union_workflow_output")
@@ -403,15 +372,11 @@ def test_executor_workflow_output_types_property():
# Test executor with multiple handlers having different workflow output types
class MultiHandlerWorkflowExecutor(Executor):
@handler
async def handle_string(
self, text: str, ctx: WorkflowContext[int, str]
) -> None:
async def handle_string(self, text: str, ctx: WorkflowContext[int, str]) -> None:
pass
@handler
async def handle_number(
self, num: int, ctx: WorkflowContext[bool, float]
) -> None:
async def handle_number(self, num: int, ctx: WorkflowContext[bool, float]) -> None:
pass
executor = MultiHandlerWorkflowExecutor(id="multi_workflow")
@@ -465,9 +430,7 @@ def test_executor_output_types_includes_response_handlers():
pass
@response_handler
async def handle_response(
self, original_request: str, response: bool, ctx: WorkflowContext[float]
) -> None:
async def handle_response(self, original_request: str, response: bool, ctx: WorkflowContext[float]) -> None:
pass
executor = RequestResponseExecutor(id="request_response")
@@ -574,9 +537,7 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
"""Test that executor_invoked event (type='executor_invoked').data captures original input, not mutated input."""
@executor(id="Mutator")
async def mutator(
messages: list[Message], ctx: WorkflowContext[list[Message]]
) -> None:
async def mutator(messages: list[Message], ctx: WorkflowContext[list[Message]]) -> None:
# The handler mutates the input list by appending new messages
original_len = len(messages)
messages.append(Message(role="assistant", text="Added by executor"))
@@ -591,11 +552,7 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
events = await workflow.run(input_messages)
# Find the invoked event for the Mutator executor
invoked_events = [
e
for e in events
if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
]
invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
assert len(invoked_events) == 1
mutator_invoked = invoked_events[0]
@@ -672,12 +629,8 @@ class TestHandlerExplicitTypes:
assert handler_func._handler_spec["output_types"] == [list] # pyright: ignore[reportFunctionMemberAccess]
# Verify can_handle
assert exec_instance.can_handle(
WorkflowMessage(data={"key": "value"}, source_id="mock")
)
assert not exec_instance.can_handle(
WorkflowMessage(data="string", source_id="mock")
)
assert exec_instance.can_handle(WorkflowMessage(data={"key": "value"}, source_id="mock"))
assert not exec_instance.can_handle(WorkflowMessage(data="string", source_id="mock"))
def test_handler_with_explicit_union_input_type(self):
"""Test that explicit union input_type is handled correctly."""
@@ -698,9 +651,7 @@ class TestHandlerExplicitTypes:
assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
assert exec_instance.can_handle(WorkflowMessage(data=42, source_id="mock"))
# Cannot handle float
assert not exec_instance.can_handle(
WorkflowMessage(data=3.14, source_id="mock")
)
assert not exec_instance.can_handle(WorkflowMessage(data=3.14, source_id="mock"))
def test_handler_with_explicit_union_output_type(self):
"""Test that explicit union output is normalized to a list."""
@@ -776,9 +727,7 @@ class TestHandlerExplicitTypes:
class OnlyWorkflowOutputExecutor(Executor): # pyright: ignore[reportUnusedClass]
@handler(workflow_output=bool)
async def handle(
self, message: str, ctx: WorkflowContext[int, str]
) -> None:
async def handle(self, message: str, ctx: WorkflowContext[int, str]) -> None:
pass
def test_handler_explicit_input_type_allows_no_message_annotation(self):
@@ -803,9 +752,7 @@ class TestHandlerExplicitTypes:
pass
@handler
async def handle_introspected(
self, message: float, ctx: WorkflowContext[bool]
) -> None:
async def handle_introspected(self, message: float, ctx: WorkflowContext[bool]) -> None:
pass
exec_instance = MixedExecutor(id="mixed")
@@ -831,9 +778,7 @@ class TestHandlerExplicitTypes:
# Should resolve the string to the actual type
assert ForwardRefMessage in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
assert exec_instance.can_handle(
WorkflowMessage(data=ForwardRefMessage("hello"), source_id="mock")
)
assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefMessage("hello"), source_id="mock"))
def test_handler_with_string_forward_reference_union(self):
"""Test that string forward references work with union types."""
@@ -846,12 +791,8 @@ class TestHandlerExplicitTypes:
exec_instance = StringUnionExecutor(id="string_union")
# Should handle both types
assert exec_instance.can_handle(
WorkflowMessage(data=ForwardRefTypeA("hello"), source_id="mock")
)
assert exec_instance.can_handle(
WorkflowMessage(data=ForwardRefTypeB(42), source_id="mock")
)
assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefTypeA("hello"), source_id="mock"))
assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefTypeB(42), source_id="mock"))
def test_handler_with_string_forward_reference_output_type(self):
"""Test that string forward references work for output_type."""
@@ -890,9 +831,7 @@ class TestHandlerExplicitTypes:
class PrecedenceExecutor(Executor):
@handler(input=int, output=float, workflow_output=str)
async def handle(
self, message: int, ctx: WorkflowContext[int, bool]
) -> None:
async def handle(self, message: int, ctx: WorkflowContext[int, bool]) -> None:
pass
exec_instance = PrecedenceExecutor(id="precedence")
@@ -958,9 +897,7 @@ class TestHandlerExplicitTypes:
async def handle(self, message, ctx: WorkflowContext) -> None: # type: ignore[no-untyped-def]
pass
exec_instance = StringUnionWorkflowOutputExecutor(
id="string_union_workflow_output"
)
exec_instance = StringUnionWorkflowOutputExecutor(id="string_union_workflow_output")
# Should resolve both types from string union
assert ForwardRefTypeA in exec_instance.workflow_output_types
@@ -971,14 +908,10 @@ class TestHandlerExplicitTypes:
class IntrospectedWorkflowOutputExecutor(Executor):
@handler
async def handle(
self, message: str, ctx: WorkflowContext[int, bool]
) -> None:
async def handle(self, message: str, ctx: WorkflowContext[int, bool]) -> None:
pass
exec_instance = IntrospectedWorkflowOutputExecutor(
id="introspected_workflow_output"
)
exec_instance = IntrospectedWorkflowOutputExecutor(id="introspected_workflow_output")
# Should use introspected types from WorkflowContext[int, bool]
assert int in exec_instance.output_types
python/packages/core/tests/workflow/test_workflow_agent.py
+32 -4
View File
@@ -717,9 +717,23 @@ class TestWorkflowAgent:
return AgentSession()
@overload
def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
@@ -813,9 +827,23 @@ class TestWorkflowAgent:
return AgentSession()
@overload
def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
python/packages/core/tests/workflow/test_workflow_kwargs.py
+80 -10
View File
@@ -52,9 +52,23 @@ class _KwargsCapturingAgent(BaseAgent):
self.captured_kwargs = []
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
@@ -90,9 +104,23 @@ class _OptionsAwareAgent(BaseAgent):
self.captured_kwargs = []
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
@@ -475,9 +503,23 @@ async def test_kwargs_preserved_on_response_continuation() -> None:
self._asked = False
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
@@ -538,9 +580,23 @@ async def test_kwargs_overridden_on_response_continuation() -> None:
self._asked = False
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
@@ -605,9 +661,23 @@ async def test_kwargs_empty_value_passed_on_continuation() -> None:
self._asked = False
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
python/packages/core/tests/workflow/test_workflow_states.py
+6 -2
View File
@@ -38,7 +38,9 @@ async def test_executor_failed_and_workflow_failed_events_streaming():
events.append(ev)
# executor_failed event (type='executor_failed') should be emitted before workflow failed event
executor_failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
executor_failed_events: list[WorkflowEvent[Any]] = [
e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"
]
assert executor_failed_events, "executor_failed event should be emitted when start executor fails"
assert executor_failed_events[0].executor_id == "f"
assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK
@@ -96,7 +98,9 @@ async def test_executor_failed_event_from_second_executor_in_chain():
events.append(ev)
# executor_failed event should be emitted for the failing executor
executor_failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
executor_failed_events: list[WorkflowEvent[Any]] = [
e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"
]
assert executor_failed_events, "executor_failed event should be emitted when second executor fails"
assert executor_failed_events[0].executor_id == "failing"
assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK