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: name changes executed (#607)

Browse Source 
* name changes executed

* updated adr to accepted

* renamed openai base config

* renamed openai config to mixin

* added renames in user docs

* reverted mcperror

* fix tests

* remove sse from tests
This commit is contained in:
Eduard van Valkenburg
2025-09-04 17:00:38 +02:00
committed by GitHub
Unverified
parent 6310ca5be0
commit 40ab6e9d67
100 changed files with 1223 additions and 1100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/tests/main/conftest.py
+5 -5
View File
@@ -4,7 +4,7 @@ from typing import Any
from pydantic import BaseModel
from pytest import fixture
from agent_framework import AITool, ChatMessage, ai_function
from agent_framework import ChatMessage, ToolProtocol, ai_function
from agent_framework.telemetry import ModelDiagnosticSettings
@@ -14,8 +14,8 @@ def chat_history() -> list[ChatMessage]:
@fixture
def ai_tool() -> AITool:
"""Returns a generic AITool."""
def ai_tool() -> ToolProtocol:
"""Returns a generic ToolProtocol."""
class GenericTool(BaseModel):
name: str
@@ -32,8 +32,8 @@ def ai_tool() -> AITool:
@fixture
def ai_function_tool() -> AITool:
"""Returns a executable AITool."""
def ai_function_tool() -> ToolProtocol:
"""Returns a executable ToolProtocol."""
@ai_function
def simple_function(x: int, y: int) -> int:
python/packages/main/tests/main/test_agents.py
+50 -50
View File
@@ -7,19 +7,19 @@ from uuid import uuid4
from pytest import fixture, raises
from agent_framework import (
AgentProtocol,
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
AIAgent,
ChatClient,
ChatClientAgent,
ChatClientBase,
BaseChatClient,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatMessageList,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
Role,
TextContent,
)
from agent_framework.exceptions import AgentExecutionException
@@ -31,7 +31,7 @@ class MockAgentThread(AgentThread):
# Mock Agent implementation for testing
class MockAgent(AIAgent):
class MockAgent(AgentProtocol):
@property
def id(self) -> str:
return str(uuid4())
@@ -57,9 +57,9 @@ class MockAgent(AIAgent):
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentRunResponse:
return AgentRunResponse(messages=[ChatMessage(role=ChatRole.ASSISTANT, contents=[TextContent("Response")])])
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("Response")])])
async def run_streaming(
async def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
@@ -72,8 +72,8 @@ class MockAgent(AIAgent):
return MockAgentThread()
# Mock ChatClient implementation for testing
class MockChatClient(ChatClientBase):
# Mock ChatClientProtocol implementation for testing
class MockChatClient(BaseChatClient):
_mock_response: ChatResponse | None = None
def __init__(self, mock_response: ChatResponse | None = None) -> None:
@@ -89,7 +89,7 @@ class MockChatClient(ChatClientBase):
return (
self._mock_response
if self._mock_response
else ChatResponse(messages=ChatMessage(role=ChatRole.ASSISTANT, text="test response"))
else ChatResponse(messages=ChatMessage(role=Role.ASSISTANT, text="test response"))
)
async def _inner_get_streaming_response(
@@ -99,7 +99,7 @@ class MockChatClient(ChatClientBase):
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(role=ChatRole.ASSISTANT, text=TextContent(text="test streaming response"))
yield ChatResponseUpdate(role=Role.ASSISTANT, text=TextContent(text="test streaming response"))
@fixture
@@ -108,12 +108,12 @@ def agent_thread() -> AgentThread:
@fixture
def agent() -> AIAgent:
def agent() -> AgentProtocol:
return MockAgent()
@fixture
def chat_client() -> ChatClientBase:
def chat_client() -> BaseChatClient:
return MockChatClient()
@@ -121,33 +121,33 @@ def test_agent_thread_type(agent_thread: AgentThread) -> None:
assert isinstance(agent_thread, AgentThread)
def test_agent_type(agent: AIAgent) -> None:
assert isinstance(agent, AIAgent)
def test_agent_type(agent: AgentProtocol) -> None:
assert isinstance(agent, AgentProtocol)
async def test_agent_run(agent: AIAgent) -> None:
async def test_agent_run(agent: AgentProtocol) -> None:
response = await agent.run("test")
assert response.messages[0].role == ChatRole.ASSISTANT
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == "Response"
async def test_agent_run_streaming(agent: AIAgent) -> None:
async def test_agent_run_streaming(agent: AgentProtocol) -> None:
async def collect_updates(updates: AsyncIterable[AgentRunResponseUpdate]) -> list[AgentRunResponseUpdate]:
return [u async for u in updates]
updates = await collect_updates(agent.run_streaming(messages="test"))
updates = await collect_updates(agent.run_stream(messages="test"))
assert len(updates) == 1
assert updates[0].text == "Response"
def test_chat_client_agent_type(chat_client: ChatClient) -> None:
chat_client_agent = ChatClientAgent(chat_client=chat_client)
assert isinstance(chat_client_agent, AIAgent)
def test_chat_client_agent_type(chat_client: ChatClientProtocol) -> None:
chat_client_agent = ChatAgent(chat_client=chat_client)
assert isinstance(chat_client_agent, AgentProtocol)
async def test_chat_client_agent_init(chat_client: ChatClient) -> None:
async def test_chat_client_agent_init(chat_client: ChatClientProtocol) -> None:
agent_id = str(uuid4())
agent = ChatClientAgent(chat_client=chat_client, id=agent_id, description="Test")
agent = ChatAgent(chat_client=chat_client, id=agent_id, description="Test")
assert agent.id == agent_id
assert agent.name is None
@@ -155,9 +155,9 @@ async def test_chat_client_agent_init(chat_client: ChatClient) -> None:
assert agent.display_name == agent_id # Display name defaults to id if name is None
async def test_chat_client_agent_init_with_name(chat_client: ChatClient) -> None:
async def test_chat_client_agent_init_with_name(chat_client: ChatClientProtocol) -> None:
agent_id = str(uuid4())
agent = ChatClientAgent(chat_client=chat_client, id=agent_id, name="Test Agent", description="Test")
agent = ChatAgent(chat_client=chat_client, id=agent_id, name="Test Agent", description="Test")
assert agent.id == agent_id
assert agent.name == "Test Agent"
@@ -165,37 +165,37 @@ async def test_chat_client_agent_init_with_name(chat_client: ChatClient) -> None
assert agent.display_name == "Test Agent" # Display name is the name if present
async def test_chat_client_agent_run(chat_client: ChatClient) -> None:
agent = ChatClientAgent(chat_client=chat_client)
async def test_chat_client_agent_run(chat_client: ChatClientProtocol) -> None:
agent = ChatAgent(chat_client=chat_client)
result = await agent.run("Hello")
assert result.text == "test response"
async def test_chat_client_agent_run_streaming(chat_client: ChatClient) -> None:
agent = ChatClientAgent(chat_client=chat_client)
async def test_chat_client_agent_run_streaming(chat_client: ChatClientProtocol) -> None:
agent = ChatAgent(chat_client=chat_client)
result = await AgentRunResponse.from_agent_response_generator(agent.run_streaming("Hello"))
result = await AgentRunResponse.from_agent_response_generator(agent.run_stream("Hello"))
assert result.text == "test streaming response"
async def test_chat_client_agent_get_new_thread(chat_client: ChatClient) -> None:
agent = ChatClientAgent(chat_client=chat_client)
async def test_chat_client_agent_get_new_thread(chat_client: ChatClientProtocol) -> None:
agent = ChatAgent(chat_client=chat_client)
thread = agent.get_new_thread()
assert isinstance(thread, AgentThread)
async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatClient) -> None:
agent = ChatClientAgent(chat_client=chat_client)
message = ChatMessage(role=ChatRole.USER, text="Hello")
async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatClientProtocol) -> None:
agent = ChatAgent(chat_client=chat_client)
message = ChatMessage(role=Role.USER, text="Hello")
thread = AgentThread(message_store=ChatMessageList(messages=[message]))
_, result_messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=thread,
input_messages=[ChatMessage(role=ChatRole.USER, text="Test")],
input_messages=[ChatMessage(role=Role.USER, text="Test")],
)
assert len(result_messages) == 2
@@ -206,11 +206,11 @@ async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatCl
async def test_chat_client_agent_update_thread_id() -> None:
chat_client = MockChatClient(
mock_response=ChatResponse(
messages=[ChatMessage(role=ChatRole.ASSISTANT, contents=[TextContent("test response")])],
messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")])],
conversation_id="123",
)
)
agent = ChatClientAgent(chat_client=chat_client)
agent = ChatAgent(chat_client=chat_client)
thread = agent.get_new_thread()
result = await agent.run("Hello", thread=thread)
@@ -219,8 +219,8 @@ async def test_chat_client_agent_update_thread_id() -> None:
assert thread.service_thread_id == "123"
async def test_chat_client_agent_update_thread_messages(chat_client: ChatClient) -> None:
agent = ChatClientAgent(chat_client=chat_client)
async def test_chat_client_agent_update_thread_messages(chat_client: ChatClientProtocol) -> None:
agent = ChatAgent(chat_client=chat_client)
thread = agent.get_new_thread()
result = await agent.run("Hello", thread=thread)
@@ -237,26 +237,26 @@ async def test_chat_client_agent_update_thread_messages(chat_client: ChatClient)
assert chat_messages[1].text == "test response"
async def test_chat_client_agent_update_thread_conversation_id_missing(chat_client: ChatClient) -> None:
agent = ChatClientAgent(chat_client=chat_client)
async def test_chat_client_agent_update_thread_conversation_id_missing(chat_client: ChatClientProtocol) -> None:
agent = ChatAgent(chat_client=chat_client)
thread = AgentThread(service_thread_id="123")
with raises(AgentExecutionException, match="Service did not return a valid conversation id"):
agent._update_thread_with_type_and_conversation_id(thread, None) # type: ignore[reportPrivateUsage]
async def test_chat_client_agent_default_author_name(chat_client: ChatClient) -> None:
async def test_chat_client_agent_default_author_name(chat_client: ChatClientProtocol) -> None:
# Name is not specified here, so default name should be used
agent = ChatClientAgent(chat_client=chat_client)
agent = ChatAgent(chat_client=chat_client)
result = await agent.run("Hello")
assert result.text == "test response"
assert result.messages[0].author_name == "UnnamedAgent"
async def test_chat_client_agent_author_name_as_agent_name(chat_client: ChatClient) -> None:
async def test_chat_client_agent_author_name_as_agent_name(chat_client: ChatClientProtocol) -> None:
# Name is specified here, so it should be used as author name
agent = ChatClientAgent(chat_client=chat_client, name="TestAgent")
agent = ChatAgent(chat_client=chat_client, name="TestAgent")
result = await agent.run("Hello")
assert result.text == "test response"
@@ -267,11 +267,11 @@ async def test_chat_client_agent_author_name_is_used_from_response() -> None:
chat_client = MockChatClient(
mock_response=ChatResponse(
messages=[
ChatMessage(role=ChatRole.ASSISTANT, contents=[TextContent("test response")], author_name="TestAuthor")
ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")], author_name="TestAuthor")
]
)
)
agent = ChatClientAgent(chat_client=chat_client)
agent = ChatAgent(chat_client=chat_client)
result = await agent.run("Hello")
assert result.text == "test response"
python/packages/main/tests/main/test_clients.py
+24 -24
View File
@@ -9,17 +9,17 @@ from pydantic import Field
from pytest import fixture
from agent_framework import (
ChatClient,
ChatClientBase,
BaseChatClient,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
EmbeddingGenerator,
FunctionCallContent,
FunctionResultContent,
GeneratedEmbeddings,
Role,
TextContent,
ai_function,
use_tool_calling,
@@ -54,8 +54,8 @@ class MockChatClient:
@use_tool_calling
class MockChatClientBase(ChatClientBase):
"""Mock implementation of the ChatClientBase."""
class MockBaseChatClient(BaseChatClient):
"""Mock implementation of the BaseChatClient."""
run_responses: list[ChatResponse] = Field(default_factory=list)
streaming_responses: list[list[ChatResponseUpdate]] = Field(default_factory=list)
@@ -120,8 +120,8 @@ def chat_client() -> MockChatClient:
@fixture
def chat_client_base() -> MockChatClientBase:
return MockChatClientBase()
def chat_client_base() -> MockBaseChatClient:
return MockBaseChatClient()
@fixture
@@ -131,19 +131,19 @@ def embedding_generator() -> MockEmbeddingGenerator:
def test_chat_client_type(chat_client: MockChatClient):
assert isinstance(chat_client, ChatClient)
assert isinstance(chat_client, ChatClientProtocol)
async def test_chat_client_get_response(chat_client: MockChatClient):
response = await chat_client.get_response(ChatMessage(role="user", text="Hello"))
assert response.text == "test response"
assert response.messages[0].role == ChatRole.ASSISTANT
assert response.messages[0].role == Role.ASSISTANT
async def test_chat_client_get_streaming_response(chat_client: MockChatClient):
async for update in chat_client.get_streaming_response(ChatMessage(role="user", text="Hello")):
assert update.text == "test streaming response" or update.text == "another update"
assert update.role == ChatRole.ASSISTANT
assert update.role == Role.ASSISTANT
def test_embedding_generator_type(embedding_generator: MockEmbeddingGenerator):
@@ -158,23 +158,23 @@ async def test_embedding_generator_generate(embedding_generator: MockEmbeddingGe
assert len(emb) == 5
def test_base_client(chat_client_base: MockChatClientBase):
assert isinstance(chat_client_base, ChatClientBase)
assert isinstance(chat_client_base, ChatClient)
def test_base_client(chat_client_base: MockBaseChatClient):
assert isinstance(chat_client_base, BaseChatClient)
assert isinstance(chat_client_base, ChatClientProtocol)
async def test_base_client_get_response(chat_client_base: MockChatClientBase):
async def test_base_client_get_response(chat_client_base: MockBaseChatClient):
response = await chat_client_base.get_response(ChatMessage(role="user", text="Hello"))
assert response.messages[0].role == ChatRole.ASSISTANT
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == "test response - Hello"
async def test_base_client_get_streaming_response(chat_client_base: MockChatClientBase):
async def test_base_client_get_streaming_response(chat_client_base: MockBaseChatClient):
async for update in chat_client_base.get_streaming_response(ChatMessage(role="user", text="Hello")):
assert update.text == "update - Hello" or update.text == "another update"
async def test_base_client_with_function_calling(chat_client_base: MockChatClientBase):
async def test_base_client_with_function_calling(chat_client_base: MockBaseChatClient):
exec_counter = 0
@ai_function(name="test_function")
@@ -195,20 +195,20 @@ async def test_base_client_with_function_calling(chat_client_base: MockChatClien
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
assert exec_counter == 1
assert len(response.messages) == 3
assert response.messages[0].role == ChatRole.ASSISTANT
assert response.messages[0].role == Role.ASSISTANT
assert isinstance(response.messages[0].contents[0], FunctionCallContent)
assert response.messages[0].contents[0].name == "test_function"
assert response.messages[0].contents[0].arguments == '{"arg1": "value1"}'
assert response.messages[0].contents[0].call_id == "1"
assert response.messages[1].role == ChatRole.TOOL
assert response.messages[1].role == Role.TOOL
assert isinstance(response.messages[1].contents[0], FunctionResultContent)
assert response.messages[1].contents[0].call_id == "1"
assert response.messages[1].contents[0].result == "Processed value1"
assert response.messages[2].role == ChatRole.ASSISTANT
assert response.messages[2].role == Role.ASSISTANT
assert response.messages[2].text == "done"
async def test_base_client_with_function_calling_disabled(chat_client_base: MockChatClientBase):
async def test_base_client_with_function_calling_disabled(chat_client_base: MockBaseChatClient):
chat_client_base.__maximum_iterations_per_request = 0
exec_counter = 0
@@ -230,11 +230,11 @@ async def test_base_client_with_function_calling_disabled(chat_client_base: Mock
response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
assert exec_counter == 0
assert len(response.messages) == 1
assert response.messages[0].role == ChatRole.ASSISTANT
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == "test response - hello"
async def test_base_client_with_streaming_function_calling(chat_client_base: MockChatClientBase):
async def test_base_client_with_streaming_function_calling(chat_client_base: MockBaseChatClient):
exec_counter = 0
@ai_function(name="test_function")
@@ -272,7 +272,7 @@ async def test_base_client_with_streaming_function_calling(chat_client_base: Moc
assert exec_counter == 1
async def test_base_client_with_streaming_function_calling_disabled(chat_client_base: MockChatClientBase):
async def test_base_client_with_streaming_function_calling_disabled(chat_client_base: MockBaseChatClient):
chat_client_base.__maximum_iterations_per_request = 0
exec_counter = 0
python/packages/main/tests/main/test_mcp.py
+27 -35
View File
@@ -12,19 +12,18 @@ from mcp.shared.exceptions import McpError
from pydantic import AnyUrl, ValidationError
from agent_framework import (
AITool,
ChatMessage,
ChatRole,
DataContent,
McpSseTools,
McpStdioTool,
McpStreamableHttpTool,
McpWebsocketTool,
MCPStdioTool,
MCPStreamableHTTPTool,
MCPWebsocketTool,
Role,
TextContent,
ToolProtocol,
UriContent,
)
from agent_framework._mcp import (
McpTool,
MCPTool,
_ai_content_to_mcp_types,
_chat_message_to_mcp_types,
_get_input_model_from_mcp_prompt,
@@ -275,20 +274,20 @@ def test_get_input_model_from_mcp_prompt():
model(arg2="optional")
# McpTool tests
# MCPTool tests
async def test_local_mcp_server_initialization():
"""Test McpTool initialization."""
server = McpTool(name="test_server")
assert isinstance(server, AITool)
"""Test MCPTool initialization."""
server = MCPTool(name="test_server")
assert isinstance(server, ToolProtocol)
assert server.name == "test_server"
assert server.session is None
assert server.functions == []
async def test_local_mcp_server_context_manager():
"""Test McpTool as context manager."""
"""Test MCPTool as context manager."""
class TestServer(McpTool):
class TestServer(MCPTool):
async def connect(self):
# Mock connection
self.session = Mock(spec=ClientSession)
@@ -306,7 +305,7 @@ async def test_local_mcp_server_context_manager():
async def test_local_mcp_server_load_functions():
"""Test loading functions from MCP server."""
class TestServer(McpTool):
class TestServer(MCPTool):
async def connect(self):
self.session = Mock(spec=ClientSession)
# Mock tools list response
@@ -330,7 +329,7 @@ async def test_local_mcp_server_load_functions():
return None
server = TestServer(name="test_server")
assert isinstance(server, AITool)
assert isinstance(server, ToolProtocol)
async with server:
await server.load_tools()
assert len(server.functions) == 1
@@ -340,7 +339,7 @@ async def test_local_mcp_server_load_functions():
async def test_local_mcp_server_load_prompts():
"""Test loading prompts from MCP server."""
class TestServer(McpTool):
class TestServer(MCPTool):
async def connect(self):
self.session = Mock(spec=ClientSession)
# Mock prompts list response
@@ -369,7 +368,7 @@ async def test_local_mcp_server_load_prompts():
async def test_local_mcp_server_function_execution():
"""Test function execution through MCP server."""
class TestServer(McpTool):
class TestServer(MCPTool):
async def connect(self):
self.session = Mock(spec=ClientSession)
self.session.list_tools = AsyncMock(
@@ -410,7 +409,7 @@ async def test_local_mcp_server_function_execution():
async def test_local_mcp_server_function_execution_error():
"""Test function execution error handling."""
class TestServer(McpTool):
class TestServer(MCPTool):
async def connect(self):
self.session = Mock(spec=ClientSession)
self.session.list_tools = AsyncMock(
@@ -448,7 +447,7 @@ async def test_local_mcp_server_function_execution_error():
async def test_local_mcp_server_prompt_execution():
"""Test prompt execution through MCP server."""
class TestMcpTool(McpTool):
class TestMCPTool(MCPTool):
async def connect(self):
self.session = Mock(spec=ClientSession)
self.session.list_prompts = AsyncMock(
@@ -474,7 +473,7 @@ async def test_local_mcp_server_prompt_execution():
def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
return None
server = TestMcpTool(name="test_server")
server = TestMCPTool(name="test_server")
async with server:
await server.load_prompts()
prompt = server.functions[0]
@@ -482,37 +481,30 @@ async def test_local_mcp_server_prompt_execution():
assert len(result) == 1
assert isinstance(result[0], ChatMessage)
assert result[0].role == ChatRole.USER
assert result[0].role == Role.USER
assert len(result[0].contents) == 1
assert result[0].contents[0].text == "Test message"
# Server implementation tests
def test_local_mcp_stdio_tool_init():
"""Test McpStdioTool initialization."""
tool = McpStdioTool(name="test", command="echo", args=["hello"])
"""Test MCPStdioTool initialization."""
tool = MCPStdioTool(name="test", command="echo", args=["hello"])
assert tool.name == "test"
assert tool.command == "echo"
assert tool.args == ["hello"]
def test_local_mcp_sse_tools_init():
"""Test McpSseTools initialization."""
tool = McpSseTools(name="test", url="http://localhost:8080")
assert tool.name == "test"
assert tool.url == "http://localhost:8080"
def test_local_mcp_websocket_tool_init():
"""Test McpWebsocketTool initialization."""
tool = McpWebsocketTool(name="test", url="ws://localhost:8080")
"""Test MCPWebsocketTool initialization."""
tool = MCPWebsocketTool(name="test", url="ws://localhost:8080")
assert tool.name == "test"
assert tool.url == "ws://localhost:8080"
def test_local_mcp_streamable_http_tool_init():
"""Test McpStreamableHttpTool initialization."""
tool = McpStreamableHttpTool(name="test", url="http://localhost:8080")
"""Test MCPStreamableHTTPTool initialization."""
tool = MCPStreamableHTTPTool(name="test", url="http://localhost:8080")
assert tool.name == "test"
assert tool.url == "http://localhost:8080"
@@ -525,7 +517,7 @@ async def test_streamable_http_integration():
if not url.startswith("http"):
pytest.skip("LOCAL_MCP_URL is not an HTTP URL")
tool = McpStreamableHttpTool(name="integration_test", url=url)
tool = MCPStreamableHTTPTool(name="integration_test", url=url)
async with tool:
# Test that we can connect and load tools
python/packages/main/tests/main/test_telemetry.py
+38 -38
View File
@@ -13,7 +13,7 @@ from agent_framework import (
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
Role,
UsageDetails,
)
from agent_framework.telemetry import (
@@ -301,7 +301,7 @@ def test_start_span_empty_metadata():
def test_decorator_with_valid_class():
"""Test that decorator works with a valid ChatClientBase-like class."""
"""Test that decorator works with a valid BaseChatClient-like class."""
# Create a mock class with the required methods
class MockChatClient:
@@ -373,7 +373,7 @@ def mock_chat_client():
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
):
return ChatResponse(
messages=[ChatMessage(role=ChatRole.ASSISTANT, text="Test response")],
messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")],
usage_details=UsageDetails(input_token_count=10, output_token_count=20),
finish_reason=None,
)
@@ -381,8 +381,8 @@ def mock_chat_client():
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
):
yield ChatResponseUpdate(text="Hello", role=ChatRole.ASSISTANT)
yield ChatResponseUpdate(text=" world", role=ChatRole.ASSISTANT)
yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield ChatResponseUpdate(text=" world", role=Role.ASSISTANT)
return MockChatClient()
@@ -393,7 +393,7 @@ async def test_telemetry_disabled_bypasses_instrumentation(mock_chat_client, mod
decorated_class = use_telemetry(type(mock_chat_client))
client = decorated_class()
messages = [ChatMessage(role=ChatRole.USER, text="Test message")]
messages = [ChatMessage(role=Role.USER, text="Test message")]
chat_options = ChatOptions()
with (
@@ -412,7 +412,7 @@ async def test_instrumentation_enabled(mock_chat_client, model_diagnostic_settin
decorated_class = use_telemetry(type(mock_chat_client))
client = decorated_class()
messages = [ChatMessage(role=ChatRole.USER, text="Test message")]
messages = [ChatMessage(role=Role.USER, text="Test message")]
chat_options = ChatOptions()
with (
@@ -432,7 +432,7 @@ async def test_streaming_response_with_diagnostics_enabled_via_decorator(mock_ch
"""Test streaming telemetry through the use_telemetry decorator."""
decorated_class = use_telemetry(type(mock_chat_client))
client = decorated_class()
messages = [ChatMessage(role=ChatRole.USER, text="Test")]
messages = [ChatMessage(role=Role.USER, text="Test")]
chat_options = ChatOptions()
with (
@@ -470,7 +470,7 @@ async def test_streaming_response_with_exception_via_decorator(mock_chat_client,
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(text="Partial", role=ChatRole.ASSISTANT)
yield ChatResponseUpdate(text="Partial", role=Role.ASSISTANT)
raise ValueError("Test streaming error")
type(mock_chat_client)._inner_get_streaming_response = _inner_get_streaming_response
@@ -478,7 +478,7 @@ async def test_streaming_response_with_exception_via_decorator(mock_chat_client,
decorated_class = use_telemetry(type(mock_chat_client))
client = decorated_class()
messages = [ChatMessage(role=ChatRole.USER, text="Test")]
messages = [ChatMessage(role=Role.USER, text="Test")]
chat_options = ChatOptions()
with (
@@ -513,12 +513,12 @@ async def test_streaming_response_diagnostics_disabled_via_decorator(model_diagn
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(text="Test", role=ChatRole.ASSISTANT)
yield ChatResponseUpdate(text="Test", role=Role.ASSISTANT)
decorated_class = use_telemetry(MockStreamingClientNoDiagnostics)
client = decorated_class()
messages = [ChatMessage(role=ChatRole.USER, text="Test")]
messages = [ChatMessage(role=Role.USER, text="Test")]
chat_options = ChatOptions()
with (
@@ -561,7 +561,7 @@ async def test_empty_streaming_response_via_decorator(model_diagnostic_settings)
decorated_class = use_telemetry(MockEmptyStreamingClient)
client = decorated_class()
messages = [ChatMessage(role=ChatRole.USER, text="Test")]
messages = [ChatMessage(role=Role.USER, text="Test")]
chat_options = ChatOptions()
with (
@@ -617,7 +617,7 @@ def test_prepend_user_agent_with_none_value():
def test_agent_decorator_with_valid_class():
"""Test that agent decorator works with a valid ChatClientAgent-like class."""
"""Test that agent decorator works with a valid ChatAgent-like class."""
from agent_framework.telemetry import use_agent_telemetry
# Create a mock class with the required methods
@@ -633,7 +633,7 @@ def test_agent_decorator_with_valid_class():
async def run(self, messages=None, *, thread=None, **kwargs):
return Mock()
async def run_streaming(self, messages=None, *, thread=None, **kwargs):
async def run_stream(self, messages=None, *, thread=None, **kwargs):
async def gen():
yield Mock()
@@ -644,7 +644,7 @@ def test_agent_decorator_with_valid_class():
# Check that the methods were wrapped
assert hasattr(decorated_class.run, "__model_diagnostics_agent_run__")
assert hasattr(decorated_class.run_streaming, "__model_diagnostics_streaming_agent_run__")
assert hasattr(decorated_class.run_stream, "__model_diagnostics_streaming_agent_run__")
def test_agent_decorator_with_missing_methods():
@@ -680,7 +680,7 @@ def test_agent_decorator_with_partial_methods():
# Only the present method should be wrapped
assert hasattr(decorated_class.run, "__model_diagnostics_agent_run__")
assert not hasattr(decorated_class, "run_streaming")
assert not hasattr(decorated_class, "run_stream")
# region Test agent telemetry decorator with mock agent
@@ -689,7 +689,7 @@ def test_agent_decorator_with_partial_methods():
@pytest.fixture
def mock_chat_client_agent():
"""Create a mock chat client agent for testing."""
from agent_framework import AgentRunResponse, ChatMessage, ChatRole, UsageDetails
from agent_framework import AgentRunResponse, ChatMessage, Role, UsageDetails
class MockChatClientAgent:
AGENT_SYSTEM_NAME = "test_agent_system"
@@ -702,17 +702,17 @@ def mock_chat_client_agent():
async def run(self, messages=None, *, thread=None, **kwargs):
return AgentRunResponse(
messages=[ChatMessage(role=ChatRole.ASSISTANT, text="Agent response")],
messages=[ChatMessage(role=Role.ASSISTANT, text="Agent response")],
usage_details=UsageDetails(input_token_count=15, output_token_count=25),
response_id="test_response_id",
raw_representation=Mock(finish_reason=Mock(value="stop")),
)
async def run_streaming(self, messages=None, *, thread=None, **kwargs):
async def run_stream(self, messages=None, *, thread=None, **kwargs):
from agent_framework import AgentRunResponseUpdate
yield AgentRunResponseUpdate(text="Hello", role=ChatRole.ASSISTANT)
yield AgentRunResponseUpdate(text=" from agent", role=ChatRole.ASSISTANT)
yield AgentRunResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield AgentRunResponseUpdate(text=" from agent", role=Role.ASSISTANT)
return MockChatClientAgent()
@@ -778,7 +778,7 @@ async def test_agent_streaming_response_with_diagnostics_enabled_via_decorator(
# Collect all yielded updates
updates = []
async for update in agent.run_streaming("Test message"):
async for update in agent.run_stream("Test message"):
updates.append(update)
# Verify we got the expected updates
@@ -795,13 +795,13 @@ async def test_agent_streaming_response_with_exception_via_decorator(mock_chat_c
"""Test agent streaming telemetry exception handling through decorator."""
from agent_framework.telemetry import use_agent_telemetry
async def run_streaming(self, messages=None, *, thread=None, **kwargs):
from agent_framework import AgentRunResponseUpdate, ChatRole
async def run_stream(self, messages=None, *, thread=None, **kwargs):
from agent_framework import AgentRunResponseUpdate, Role
yield AgentRunResponseUpdate(text="Partial", role=ChatRole.ASSISTANT)
yield AgentRunResponseUpdate(text="Partial", role=Role.ASSISTANT)
raise ValueError("Test agent streaming error")
type(mock_chat_client_agent).run_streaming = run_streaming
type(mock_chat_client_agent).run_stream = run_stream
decorated_class = use_agent_telemetry(type(mock_chat_client_agent))
agent = decorated_class()
@@ -819,7 +819,7 @@ async def test_agent_streaming_response_with_exception_via_decorator(mock_chat_c
# Should raise the exception and call error handler
with pytest.raises(ValueError, match="Test agent streaming error"):
async for _ in agent.run_streaming("Test message"):
async for _ in agent.run_stream("Test message"):
pass
# Verify error was recorded
@@ -830,7 +830,7 @@ async def test_agent_streaming_response_with_exception_via_decorator(mock_chat_c
@pytest.mark.parametrize("model_diagnostic_settings", [(False, False)], indirect=True)
async def test_agent_streaming_response_diagnostics_disabled_via_decorator(model_diagnostic_settings):
"""Test agent streaming response when diagnostics are disabled."""
from agent_framework import AgentRunResponseUpdate, ChatRole
from agent_framework import AgentRunResponseUpdate, Role
from agent_framework.telemetry import use_agent_telemetry
class MockStreamingAgentNoDiagnostics:
@@ -841,8 +841,8 @@ async def test_agent_streaming_response_diagnostics_disabled_via_decorator(model
self.name = "test_agent"
self.display_name = "Test Agent"
async def run_streaming(self, messages=None, *, thread=None, **kwargs):
yield AgentRunResponseUpdate(text="Test", role=ChatRole.ASSISTANT)
async def run_stream(self, messages=None, *, thread=None, **kwargs):
yield AgentRunResponseUpdate(text="Test", role=Role.ASSISTANT)
decorated_class = use_agent_telemetry(MockStreamingAgentNoDiagnostics)
agent = decorated_class()
@@ -853,7 +853,7 @@ async def test_agent_streaming_response_diagnostics_disabled_via_decorator(model
):
# Should not create spans when diagnostics are disabled
updates = []
async for update in agent.run_streaming("Test message"):
async for update in agent.run_stream("Test message"):
updates.append(update)
assert len(updates) == 1
@@ -874,7 +874,7 @@ async def test_agent_empty_streaming_response_via_decorator(model_diagnostic_set
self.name = "test_agent"
self.display_name = "Test Agent"
async def run_streaming(self, messages=None, *, thread=None, **kwargs):
async def run_stream(self, messages=None, *, thread=None, **kwargs):
# Return empty stream
return
yield # This will never be reached
@@ -895,7 +895,7 @@ async def test_agent_empty_streaming_response_via_decorator(model_diagnostic_set
# Should handle empty stream gracefully
updates = []
async for update in agent.run_streaming("Test message"):
async for update in agent.run_stream("Test message"):
updates.append(update)
assert len(updates) == 0
@@ -943,16 +943,16 @@ async def test_agent_run_with_thread_and_kwargs(mock_chat_client_agent, model_di
@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
async def test_agent_run_with_list_messages(mock_chat_client_agent, model_diagnostic_settings):
"""Test agent run with list of messages."""
from agent_framework import ChatMessage, ChatRole
from agent_framework import ChatMessage, Role
from agent_framework.telemetry import use_agent_telemetry
decorated_class = use_agent_telemetry(type(mock_chat_client_agent))
agent = decorated_class()
messages = [
ChatMessage(role=ChatRole.USER, text="First message"),
ChatMessage(role=ChatRole.ASSISTANT, text="Response"),
ChatMessage(role=ChatRole.USER, text="Second message"),
ChatMessage(role=Role.USER, text="First message"),
ChatMessage(role=Role.ASSISTANT, text="Response"),
ChatMessage(role=Role.USER, text="Second message"),
]
with (
python/packages/main/tests/main/test_threads.py
+6 -6
View File
@@ -5,7 +5,7 @@ from typing import Any
import pytest
from agent_framework import AgentThread, ChatMessage, ChatMessageList, ChatRole
from agent_framework import AgentThread, ChatMessage, ChatMessageList, Role
from agent_framework._threads import StoreState, ThreadState, deserialize_thread_state, thread_on_new_messages
@@ -37,16 +37,16 @@ class MockChatMessageStore:
def sample_messages() -> list[ChatMessage]:
"""Fixture providing sample chat messages for testing."""
return [
ChatMessage(role=ChatRole.USER, text="Hello", message_id="msg1"),
ChatMessage(role=ChatRole.ASSISTANT, text="Hi there!", message_id="msg2"),
ChatMessage(role=ChatRole.USER, text="How are you?", message_id="msg3"),
ChatMessage(role=Role.USER, text="Hello", message_id="msg1"),
ChatMessage(role=Role.ASSISTANT, text="Hi there!", message_id="msg2"),
ChatMessage(role=Role.USER, text="How are you?", message_id="msg3"),
]
@pytest.fixture
def sample_message() -> ChatMessage:
"""Fixture providing a single sample chat message for testing."""
return ChatMessage(role=ChatRole.USER, text="Test message", message_id="test1")
return ChatMessage(role=Role.USER, text="Test message", message_id="test1")
class TestAgentThread:
@@ -171,7 +171,7 @@ class TestAgentThread:
async def test_on_new_messages_with_existing_store(self, sample_message: ChatMessage) -> None:
"""Test _on_new_messages adds to existing message store."""
initial_messages = [ChatMessage(role=ChatRole.USER, text="Initial", message_id="init1")]
initial_messages = [ChatMessage(role=Role.USER, text="Initial", message_id="init1")]
store = ChatMessageList(initial_messages)
thread = AgentThread(message_store=store)
python/packages/main/tests/main/test_tools.py
+7 -7
View File
@@ -5,7 +5,7 @@ from unittest.mock import Mock, patch
import pytest
from pydantic import BaseModel
from agent_framework import AIFunction, AITool, HostedCodeInterpreterTool, ai_function
from agent_framework import AIFunction, HostedCodeInterpreterTool, ToolProtocol, ai_function
from agent_framework._tools import _parse_inputs
from agent_framework.telemetry import GenAIAttributes
@@ -18,7 +18,7 @@ def test_ai_function_decorator():
"""A simple function that adds two numbers."""
return x + y
assert isinstance(test_tool, AITool)
assert isinstance(test_tool, ToolProtocol)
assert isinstance(test_tool, AIFunction)
assert test_tool.name == "test_tool"
assert test_tool.description == "A test tool"
@@ -39,7 +39,7 @@ def test_ai_function_decorator_without_args():
"""A simple function that adds two numbers."""
return x + y
assert isinstance(test_tool, AITool)
assert isinstance(test_tool, ToolProtocol)
assert isinstance(test_tool, AIFunction)
assert test_tool.name == "test_tool"
assert test_tool.description == "A simple function that adds two numbers."
@@ -60,7 +60,7 @@ async def test_ai_function_decorator_with_async():
"""An async function that adds two numbers."""
return x + y
assert isinstance(async_test_tool, AITool)
assert isinstance(async_test_tool, ToolProtocol)
assert isinstance(async_test_tool, AIFunction)
assert async_test_tool.name == "async_test_tool"
assert async_test_tool.description == "An async test tool"
@@ -399,7 +399,7 @@ def test_parse_inputs_data_dict():
def test_parse_inputs_ai_contents_instance():
"""Test _parse_inputs with AIContents instance."""
"""Test _parse_inputs with Contents instance."""
from agent_framework import TextContent
text_content = TextContent(text="Hello, world!")
@@ -418,7 +418,7 @@ def test_parse_inputs_mixed_list():
"http://example.com", # string
{"uri": "https://test.org", "media_type": "text/html"}, # URI dict
{"file_id": "file-456"}, # hosted file dict
TextContent(text="Hello"), # AIContents instance
TextContent(text="Hello"), # Contents instance
]
result = _parse_inputs(inputs)
@@ -477,7 +477,7 @@ def test_hosted_code_interpreter_tool_with_dict_inputs():
def test_hosted_code_interpreter_tool_with_ai_contents():
"""Test HostedCodeInterpreterTool with AIContents instances."""
"""Test HostedCodeInterpreterTool with Contents instances."""
from agent_framework import DataContent, TextContent
inputs = [TextContent(text="Hello, world!"), DataContent(data=b"test", media_type="text/plain")]
python/packages/main/tests/main/test_types.py
+57 -61
View File
@@ -9,32 +9,30 @@ from pytest import fixture, mark, raises
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AIAnnotation,
AIContent,
AIContents,
AIFunction,
AITool,
AnnotatedRegion,
ChatFinishReason,
BaseContent,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
ChatToolMode,
CitationAnnotation,
Contents,
DataContent,
ErrorContent,
FinishReason,
FunctionCallContent,
FunctionResultContent,
GeneratedEmbeddings,
HostedFileContent,
HostedVectorStoreContent,
Role,
SpeechToTextOptions,
TextContent,
TextReasoningContent,
TextSpanRegion,
TextToSpeechOptions,
ToolProtocol,
UriContent,
UsageContent,
UsageDetails,
@@ -43,8 +41,8 @@ from agent_framework import (
@fixture
def ai_tool() -> AITool:
"""Returns a generic AITool."""
def ai_tool() -> ToolProtocol:
"""Returns a generic ToolProtocol."""
class GenericTool(BaseModel):
name: str
@@ -61,8 +59,8 @@ def ai_tool() -> AITool:
@fixture
def ai_function_tool() -> AITool:
"""Returns a executable AITool."""
def ai_function_tool() -> ToolProtocol:
"""Returns a executable ToolProtocol."""
@ai_function
def simple_function(x: int, y: int) -> int:
@@ -76,7 +74,7 @@ def ai_function_tool() -> AITool:
def test_text_content_positional():
"""Test the TextContent class to ensure it initializes correctly and inherits from AIContent."""
"""Test the TextContent class to ensure it initializes correctly and inherits from BaseContent."""
# Create an instance of TextContent
content = TextContent("Hello, world!", raw_representation="Hello, world!", additional_properties={"version": 1})
@@ -85,14 +83,14 @@ def test_text_content_positional():
assert content.text == "Hello, world!"
assert content.raw_representation == "Hello, world!"
assert content.additional_properties["version"] == 1
# Ensure the instance is of type AIContent
assert isinstance(content, AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(content, BaseContent)
with raises(ValidationError):
content.type = "ai"
def test_text_content_keyword():
"""Test the TextContent class to ensure it initializes correctly and inherits from AIContent."""
"""Test the TextContent class to ensure it initializes correctly and inherits from BaseContent."""
# Create an instance of TextContent
content = TextContent(
text="Hello, world!", raw_representation="Hello, world!", additional_properties={"version": 1}
@@ -103,8 +101,8 @@ def test_text_content_keyword():
assert content.text == "Hello, world!"
assert content.raw_representation == "Hello, world!"
assert content.additional_properties["version"] == 1
# Ensure the instance is of type AIContent
assert isinstance(content, AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(content, BaseContent)
with raises(ValidationError):
content.type = "ai"
@@ -124,8 +122,8 @@ def test_data_content_bytes():
assert content.has_top_level_media_type("image") is False
assert content.additional_properties["version"] == 1
# Ensure the instance is of type AIContent
assert isinstance(content, AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(content, BaseContent)
def test_data_content_uri():
@@ -140,8 +138,8 @@ def test_data_content_uri():
assert content.has_top_level_media_type("application") is False
assert content.additional_properties["version"] == 1
# Ensure the instance is of type AIContent
assert isinstance(content, AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(content, BaseContent)
def test_data_content_invalid():
@@ -185,8 +183,8 @@ def test_uri_content():
assert content.has_top_level_media_type("application") is False
assert content.additional_properties["version"] == 1
# Ensure the instance is of type AIContent
assert isinstance(content, AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(content, BaseContent)
# region: HostedFileContent
@@ -201,8 +199,8 @@ def test_hosted_file_content():
assert content.file_id == "file-123"
assert content.additional_properties["version"] == 1
# Ensure the instance is of type AIContent
assert isinstance(content, AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(content, BaseContent)
def test_hosted_file_content_minimal():
@@ -215,8 +213,8 @@ def test_hosted_file_content_minimal():
assert content.additional_properties is None
assert content.raw_representation is None
# Ensure the instance is of type AIContent
assert isinstance(content, AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(content, BaseContent)
# region: HostedVectorStoreContent
@@ -231,9 +229,9 @@ def test_hosted_vector_store_content():
assert content.vector_store_id == "vs-789"
assert content.additional_properties["version"] == 1
# Ensure the instance is of type AIContent
# Ensure the instance is of type BaseContent
assert isinstance(content, HostedVectorStoreContent)
assert isinstance(content, AIContent)
assert isinstance(content, BaseContent)
def test_hosted_vector_store_content_minimal():
@@ -246,9 +244,9 @@ def test_hosted_vector_store_content_minimal():
assert content.additional_properties is None
assert content.raw_representation is None
# Ensure the instance is of type AIContent
# Ensure the instance is of type BaseContent
assert isinstance(content, HostedVectorStoreContent)
assert isinstance(content, AIContent)
assert isinstance(content, BaseContent)
# region FunctionCallContent
@@ -263,8 +261,8 @@ def test_function_call_content():
assert content.name == "example_function"
assert content.arguments == {"param1": "value1"}
# Ensure the instance is of type AIContent
assert isinstance(content, AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(content, BaseContent)
def test_function_call_content_parse_arguments():
@@ -315,8 +313,8 @@ def test_function_result_content():
assert content.type == "function_result"
assert content.result == {"param1": "value1"}
# Ensure the instance is of type AIContent
assert isinstance(content, AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(content, BaseContent)
# region UsageDetails
@@ -381,7 +379,7 @@ def test_usage_details_add_with_none_and_type_errors():
u += 42 # type: ignore[arg-type]
# region AIContent Serialization
# region BaseContent Serialization
@mark.parametrize(
@@ -396,14 +394,14 @@ def test_usage_details_add_with_none_and_type_errors():
(HostedVectorStoreContent, {"vector_store_id": "vs-789"}),
],
)
def test_ai_content_serialization(content_type: type[AIContent], args: dict):
def test_ai_content_serialization(content_type: type[BaseContent], args: dict):
content = content_type(**args)
serialized = content.model_dump()
deserialized = content_type.model_validate(serialized)
assert deserialized == content
class TestModel(BaseModel):
content: AIContents
content: Contents
test_item = TestModel.model_validate({"content": serialized})
@@ -419,14 +417,14 @@ def test_chat_message_text():
message = ChatMessage(role="user", text="Hello, how are you?")
# Check the type and content
assert message.role == ChatRole.USER
assert message.role == Role.USER
assert len(message.contents) == 1
assert isinstance(message.contents[0], TextContent)
assert message.contents[0].text == "Hello, how are you?"
assert message.text == "Hello, how are you?"
# Ensure the instance is of type AIContent
assert isinstance(message.contents[0], AIContent)
# Ensure the instance is of type BaseContent
assert isinstance(message.contents[0], BaseContent)
def test_chat_message_contents():
@@ -437,7 +435,7 @@ def test_chat_message_contents():
message = ChatMessage(role="user", contents=[content1, content2])
# Check the type and content
assert message.role == ChatRole.USER
assert message.role == Role.USER
assert len(message.contents) == 2
assert isinstance(message.contents[0], TextContent)
assert isinstance(message.contents[1], TextContent)
@@ -447,8 +445,8 @@ def test_chat_message_contents():
def test_chat_message_with_chatrole_instance():
m = ChatMessage(role=ChatRole.USER, text="hi")
assert m.role == ChatRole.USER
m = ChatMessage(role=Role.USER, text="hi")
assert m.role == Role.USER
assert m.text == "hi"
@@ -464,7 +462,7 @@ def test_chat_response():
response = ChatResponse(messages=message)
# Check the type and content
assert response.messages[0].role == ChatRole.ASSISTANT
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == "I'm doing well, thank you!"
assert isinstance(response.messages[0], ChatMessage)
# __str__ returns text
@@ -484,7 +482,7 @@ def test_chat_response_with_format():
response = ChatResponse(messages=message)
# Check the type and content
assert response.messages[0].role == ChatRole.ASSISTANT
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == '{"response": "Hello"}'
assert isinstance(response.messages[0], ChatMessage)
assert response.text == '{"response": "Hello"}'
@@ -503,7 +501,7 @@ def test_chat_response_with_format_init():
response = ChatResponse(messages=message, response_format=OutputModel)
# Check the type and content
assert response.messages[0].role == ChatRole.ASSISTANT
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == '{"response": "Hello"}'
assert isinstance(response.messages[0], ChatMessage)
assert response.text == '{"response": "Hello"}'
@@ -767,7 +765,7 @@ def test_chat_options_init_with_args(ai_function_tool, ai_tool) -> None:
assert options.frequency_penalty == 0.0
assert options.user == "user-123"
for tool in options.tools:
assert isinstance(tool, AITool)
assert isinstance(tool, ToolProtocol)
assert tool.name is not None
assert tool.description is not None
if isinstance(tool, AIFunction):
@@ -809,7 +807,7 @@ def test_chat_options_and(ai_function_tool, ai_tool) -> None:
@fixture
def chat_message() -> ChatMessage:
return ChatMessage(role=ChatRole.USER, text="Hello")
return ChatMessage(role=Role.USER, text="Hello")
@fixture
@@ -824,7 +822,7 @@ def agent_run_response(chat_message: ChatMessage) -> AgentRunResponse:
@fixture
def agent_run_response_update(text_content: TextContent) -> AgentRunResponseUpdate:
return AgentRunResponseUpdate(role=ChatRole.ASSISTANT, contents=[text_content])
return AgentRunResponseUpdate(role=Role.ASSISTANT, contents=[text_content])
# region AgentRunResponse
@@ -914,18 +912,16 @@ def test_error_content_str():
def test_annotations_models_and_roundtrip():
span = TextSpanRegion(start_index=0, end_index=5)
base_region = AnnotatedRegion()
ann: AIAnnotation = AIAnnotation(annotated_regions=[span, base_region])
cit = CitationAnnotation(title="Doc", url="http://example.com", snippet="Snippet", annotated_regions=[span])
# Attach to content
content = TextContent(text="hello", additional_properties={"v": 1})
content.annotations = [ann, cit]
content.annotations = [cit]
dumped = content.model_dump()
loaded = TextContent.model_validate(dumped)
assert isinstance(loaded.annotations, list)
assert len(loaded.annotations) == 2
assert len(loaded.annotations) == 1
assert isinstance(loaded.annotations[0], dict) is False # pydantic parsed into models
# discriminators preserved
assert any(getattr(a, "type", None) == "citation" for a in loaded.annotations)
@@ -1032,16 +1028,16 @@ def test_generated_embeddings_operations():
assert g.additional_properties == {}
# region ChatRole & ChatFinishReason basics
# region Role & FinishReason basics
def test_chat_role_str_and_repr():
assert str(ChatRole.USER) == "user"
assert "ChatRole(value=" in repr(ChatRole.USER)
assert str(Role.USER) == "user"
assert "Role(value=" in repr(Role.USER)
def test_chat_finish_reason_constants():
assert ChatFinishReason.STOP.value == "stop"
assert FinishReason.STOP.value == "stop"
def test_response_update_propagates_fields_and_metadata():
@@ -1054,7 +1050,7 @@ def test_response_update_propagates_fields_and_metadata():
conversation_id="cid",
ai_model_id="model-x",
created_at="t0",
finish_reason=ChatFinishReason.STOP,
finish_reason=FinishReason.STOP,
additional_properties={"k": "v"},
)
resp = ChatResponse.from_chat_response_updates([upd])
@@ -1062,9 +1058,9 @@ def test_response_update_propagates_fields_and_metadata():
assert resp.created_at == "t0"
assert resp.conversation_id == "cid"
assert resp.ai_model_id == "model-x"
assert resp.finish_reason == ChatFinishReason.STOP
assert resp.finish_reason == FinishReason.STOP
assert resp.additional_properties and resp.additional_properties["k"] == "v"
assert resp.messages[0].role == ChatRole.ASSISTANT
assert resp.messages[0].role == Role.ASSISTANT
assert resp.messages[0].author_name == "bot"
assert resp.messages[0].message_id == "mid"
python/packages/main/tests/openai/test_openai_assistants_client.py
+40 -40
View File
@@ -14,19 +14,19 @@ from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
ChatClient,
ChatClientAgent,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
ChatToolMode,
FunctionCallContent,
FunctionResultContent,
HostedCodeInterpreterTool,
HostedFileSearchTool,
HostedVectorStoreContent,
Role,
TextContent,
UriContent,
UsageContent,
@@ -125,7 +125,7 @@ def test_openai_assistants_client_init_with_client(mock_async_openai: MagicMock)
assert chat_client.assistant_id == "existing-assistant-id"
assert chat_client.thread_id == "test-thread-id"
assert not chat_client._should_delete_assistant # type: ignore
assert isinstance(chat_client, ChatClient)
assert isinstance(chat_client, ChatClientProtocol)
def test_openai_assistants_client_init_auto_create_client(
@@ -185,7 +185,7 @@ def test_openai_assistants_client_init_with_default_headers(openai_unit_test_env
)
assert chat_client.ai_model_id == "gpt-4"
assert isinstance(chat_client, ChatClient)
assert isinstance(chat_client, ChatClientProtocol)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
@@ -412,7 +412,7 @@ async def test_openai_assistants_client_process_stream_events_thread_run_created
update = updates[0]
assert isinstance(update, ChatResponseUpdate)
assert update.conversation_id == thread_id
assert update.role == ChatRole.ASSISTANT
assert update.role == Role.ASSISTANT
assert update.contents == []
assert update.raw_representation == mock_response.data
@@ -457,7 +457,7 @@ async def test_openai_assistants_client_process_stream_events_message_delta_text
update = updates[0]
assert isinstance(update, ChatResponseUpdate)
assert update.conversation_id == thread_id
assert update.role == ChatRole.ASSISTANT
assert update.role == Role.ASSISTANT
assert update.text == "Hello from assistant"
assert update.raw_representation == mock_message_delta
@@ -497,7 +497,7 @@ async def test_openai_assistants_client_process_stream_events_requires_action(mo
update = updates[0]
assert isinstance(update, ChatResponseUpdate)
assert update.conversation_id == thread_id
assert update.role == ChatRole.ASSISTANT
assert update.role == Role.ASSISTANT
assert len(update.contents) == 1
assert update.contents[0] == test_function_content
assert update.raw_representation == mock_run
@@ -579,7 +579,7 @@ async def test_openai_assistants_client_process_stream_events_run_completed_with
update = updates[0]
assert isinstance(update, ChatResponseUpdate)
assert update.conversation_id == thread_id
assert update.role == ChatRole.ASSISTANT
assert update.role == Role.ASSISTANT
assert len(update.contents) == 1
# Check the usage content
@@ -632,7 +632,7 @@ def test_openai_assistants_client_create_run_options_basic(mock_async_openai: Ma
top_p=0.9,
)
messages = [ChatMessage(role=ChatRole.USER, text="Hello")]
messages = [ChatMessage(role=Role.USER, text="Hello")]
# Call the method
run_options, tool_results = chat_client._create_run_options(messages, chat_options) # type: ignore
@@ -661,7 +661,7 @@ def test_openai_assistants_client_create_run_options_with_ai_function_tool(mock_
tool_choice="auto",
)
messages = [ChatMessage(role=ChatRole.USER, text="Hello")]
messages = [ChatMessage(role=Role.USER, text="Hello")]
# Call the method
run_options, tool_results = chat_client._create_run_options(messages, chat_options) # type: ignore
@@ -686,7 +686,7 @@ def test_openai_assistants_client_create_run_options_with_code_interpreter(mock_
tool_choice="auto",
)
messages = [ChatMessage(role=ChatRole.USER, text="Calculate something")]
messages = [ChatMessage(role=Role.USER, text="Calculate something")]
# Call the method
run_options, tool_results = chat_client._create_run_options(messages, chat_options) # type: ignore
@@ -706,7 +706,7 @@ def test_openai_assistants_client_create_run_options_tool_choice_none(mock_async
tool_choice="none",
)
messages = [ChatMessage(role=ChatRole.USER, text="Hello")]
messages = [ChatMessage(role=Role.USER, text="Hello")]
# Call the method
run_options, tool_results = chat_client._create_run_options(messages, chat_options) # type: ignore
@@ -727,7 +727,7 @@ def test_openai_assistants_client_create_run_options_required_function(mock_asyn
tool_choice=tool_choice,
)
messages = [ChatMessage(role=ChatRole.USER, text="Hello")]
messages = [ChatMessage(role=Role.USER, text="Hello")]
# Call the method
run_options, tool_results = chat_client._create_run_options(messages, chat_options) # type: ignore
@@ -753,7 +753,7 @@ def test_openai_assistants_client_create_run_options_with_file_search_tool(mock_
tool_choice="auto",
)
messages = [ChatMessage(role=ChatRole.USER, text="Search for information")]
messages = [ChatMessage(role=Role.USER, text="Search for information")]
# Call the method
run_options, tool_results = chat_client._create_run_options(messages, chat_options) # type: ignore
@@ -778,7 +778,7 @@ def test_openai_assistants_client_create_run_options_with_mapping_tool(mock_asyn
tool_choice="auto",
)
messages = [ChatMessage(role=ChatRole.USER, text="Use custom tool")]
messages = [ChatMessage(role=Role.USER, text="Use custom tool")]
# Call the method
run_options, tool_results = chat_client._create_run_options(messages, chat_options) # type: ignore
@@ -795,8 +795,8 @@ def test_openai_assistants_client_create_run_options_with_system_message(mock_as
chat_client = create_test_openai_assistants_client(mock_async_openai)
messages = [
ChatMessage(role=ChatRole.SYSTEM, text="You are a helpful assistant."),
ChatMessage(role=ChatRole.USER, text="Hello"),
ChatMessage(role=Role.SYSTEM, text="You are a helpful assistant."),
ChatMessage(role=Role.USER, text="Hello"),
]
# Call the method
@@ -816,7 +816,7 @@ def test_openai_assistants_client_create_run_options_with_image_content(mock_asy
# Create message with image content
image_content = UriContent(uri="https://example.com/image.jpg", media_type="image/jpeg")
messages = [ChatMessage(role=ChatRole.USER, contents=[image_content])]
messages = [ChatMessage(role=Role.USER, contents=[image_content])]
# Call the method
run_options, tool_results = chat_client._create_run_options(messages, None) # type: ignore
@@ -924,7 +924,7 @@ def get_weather(
async def test_openai_assistants_client_get_response() -> None:
"""Test OpenAI Assistants Client response."""
async with OpenAIAssistantsClient() as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClient)
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
@@ -948,7 +948,7 @@ async def test_openai_assistants_client_get_response() -> None:
async def test_openai_assistants_client_get_response_tools() -> None:
"""Test OpenAI Assistants Client response with tools."""
async with OpenAIAssistantsClient() as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClient)
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
@@ -969,7 +969,7 @@ async def test_openai_assistants_client_get_response_tools() -> None:
async def test_openai_assistants_client_streaming() -> None:
"""Test OpenAI Assistants Client streaming response."""
async with OpenAIAssistantsClient() as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClient)
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
@@ -999,7 +999,7 @@ async def test_openai_assistants_client_streaming() -> None:
async def test_openai_assistants_client_streaming_tools() -> None:
"""Test OpenAI Assistants Client streaming response with tools."""
async with OpenAIAssistantsClient() as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClient)
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
@@ -1035,7 +1035,7 @@ async def test_openai_assistants_client_with_existing_assistant() -> None:
async with OpenAIAssistantsClient(
ai_model_id="gpt-4o-mini", assistant_id=assistant_id
) as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClient)
assert isinstance(openai_assistants_client, ChatClientProtocol)
assert openai_assistants_client.assistant_id == assistant_id
messages = [ChatMessage(role="user", text="What can you do?")]
@@ -1052,7 +1052,7 @@ async def test_openai_assistants_client_with_existing_assistant() -> None:
async def test_openai_assistants_client_file_search() -> None:
"""Test OpenAI Assistants Client response."""
async with OpenAIAssistantsClient() as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClient)
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
@@ -1074,7 +1074,7 @@ async def test_openai_assistants_client_file_search() -> None:
async def test_openai_assistants_client_file_search_streaming() -> None:
"""Test OpenAI Assistants Client response."""
async with OpenAIAssistantsClient() as openai_assistants_client:
assert isinstance(openai_assistants_client, ChatClient)
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
@@ -1101,8 +1101,8 @@ async def test_openai_assistants_client_file_search_streaming() -> None:
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_agent_basic_run():
"""Test ChatClientAgent basic run functionality with OpenAIAssistantsClient."""
async with ChatClientAgent(
"""Test ChatAgent basic run functionality with OpenAIAssistantsClient."""
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
) as agent:
# Run a simple query
@@ -1117,13 +1117,13 @@ async def test_openai_assistants_agent_basic_run():
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_agent_basic_run_streaming():
"""Test ChatClientAgent basic streaming functionality with OpenAIAssistantsClient."""
async with ChatClientAgent(
"""Test ChatAgent basic streaming functionality with OpenAIAssistantsClient."""
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
) as agent:
# Run streaming query
full_message: str = ""
async for chunk in agent.run_streaming("Please respond with exactly: 'This is a streaming response test.'"):
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert chunk is not None
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
@@ -1136,8 +1136,8 @@ async def test_openai_assistants_agent_basic_run_streaming():
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_agent_thread_persistence():
"""Test ChatClientAgent thread persistence across runs with OpenAIAssistantsClient."""
async with ChatClientAgent(
"""Test ChatAgent thread persistence across runs with OpenAIAssistantsClient."""
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
instructions="You are a helpful assistant with good memory.",
) as agent:
@@ -1164,11 +1164,11 @@ async def test_openai_assistants_agent_thread_persistence():
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_agent_existing_thread_id():
"""Test ChatClientAgent with existing thread ID to continue conversations across agent instances."""
"""Test ChatAgent with existing thread ID to continue conversations across agent instances."""
# First, create a conversation and capture the thread ID
existing_thread_id = None
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
instructions="You are a helpful weather agent.",
tools=[get_weather],
@@ -1188,7 +1188,7 @@ async def test_openai_assistants_agent_existing_thread_id():
# Now continue with the same thread ID in a new agent instance
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIAssistantsClient(thread_id=existing_thread_id),
instructions="You are a helpful weather agent.",
tools=[get_weather],
@@ -1208,9 +1208,9 @@ async def test_openai_assistants_agent_existing_thread_id():
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_agent_code_interpreter():
"""Test ChatClientAgent with code interpreter through OpenAIAssistantsClient."""
"""Test ChatAgent with code interpreter through OpenAIAssistantsClient."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
instructions="You are a helpful assistant that can write and execute Python code.",
tools=[HostedCodeInterpreterTool()],
@@ -1229,7 +1229,7 @@ async def test_openai_assistants_agent_code_interpreter():
async def test_openai_assistants_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with OpenAI Assistants Client."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
@@ -1264,7 +1264,7 @@ async def test_openai_assistants_client_run_level_tool_isolation():
call_count += 1
return f"The weather in {location} is sunny and 72°F."
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIAssistantsClient(),
instructions="You are a helpful assistant.",
) as agent:
python/packages/main/tests/openai/test_openai_chat_client.py
+24 -24
View File
@@ -10,15 +10,15 @@ from openai import BadRequestError
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AITool,
ChatClient,
ChatClientAgent,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
HostedWebSearchTool,
TextContent,
ToolProtocol,
ai_function,
)
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
@@ -39,7 +39,7 @@ def test_init(openai_unit_test_env: dict[str, str]) -> None:
open_ai_chat_completion = OpenAIChatClient()
assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
assert isinstance(open_ai_chat_completion, ChatClient)
assert isinstance(open_ai_chat_completion, ChatClientProtocol)
def test_init_validation_fail() -> None:
@@ -54,7 +54,7 @@ def test_init_ai_model_id_constructor(openai_unit_test_env: dict[str, str]) -> N
open_ai_chat_completion = OpenAIChatClient(ai_model_id=ai_model_id)
assert open_ai_chat_completion.ai_model_id == ai_model_id
assert isinstance(open_ai_chat_completion, ChatClient)
assert isinstance(open_ai_chat_completion, ChatClientProtocol)
def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None:
@@ -66,7 +66,7 @@ def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None:
)
assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
assert isinstance(open_ai_chat_completion, ChatClient)
assert isinstance(open_ai_chat_completion, ChatClientProtocol)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
@@ -154,15 +154,15 @@ def test_unsupported_tool_handling(openai_unit_test_env: dict[str, str]) -> None
"""Test that unsupported tool types are handled correctly."""
client = OpenAIChatClient()
# Create a mock AITool that's not an AIFunction
unsupported_tool = MagicMock(spec=AITool)
# Create a mock ToolProtocol that's not an AIFunction
unsupported_tool = MagicMock(spec=ToolProtocol)
unsupported_tool.__class__.__name__ = "UnsupportedAITool"
# This should ignore the unsupported AITool and return empty list
# This should ignore the unsupported ToolProtocol and return empty list
result = client._chat_to_tool_spec([unsupported_tool]) # type: ignore
assert result == []
# Also test with a non-AITool that should be converted to dict
# Also test with a non-ToolProtocol that should be converted to dict
dict_tool = {"type": "function", "name": "test"}
result = client._chat_to_tool_spec([dict_tool]) # type: ignore
assert result == [dict_tool]
@@ -190,7 +190,7 @@ async def test_openai_chat_completion_response() -> None:
"""Test OpenAI chat completion responses."""
openai_chat_client = OpenAIChatClient()
assert isinstance(openai_chat_client, ChatClient)
assert isinstance(openai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
@@ -217,7 +217,7 @@ async def test_openai_chat_completion_response_tools() -> None:
"""Test OpenAI chat completion responses."""
openai_chat_client = OpenAIChatClient()
assert isinstance(openai_chat_client, ChatClient)
assert isinstance(openai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
@@ -239,7 +239,7 @@ async def test_openai_chat_client_streaming() -> None:
"""Test Azure OpenAI chat completion responses."""
openai_chat_client = OpenAIChatClient()
assert isinstance(openai_chat_client, ChatClient)
assert isinstance(openai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
@@ -274,7 +274,7 @@ async def test_openai_chat_client_streaming_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
openai_chat_client = OpenAIChatClient()
assert isinstance(openai_chat_client, ChatClient)
assert isinstance(openai_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
@@ -301,7 +301,7 @@ async def test_openai_chat_client_web_search() -> None:
# Currently only a select few models support web search tool calls
openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview")
assert isinstance(openai_chat_client, ChatClient)
assert isinstance(openai_chat_client, ChatClientProtocol)
# Test that the client will use the web search tool
response = await openai_chat_client.get_response(
@@ -340,7 +340,7 @@ async def test_openai_chat_client_web_search() -> None:
async def test_openai_chat_client_web_search_streaming() -> None:
openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview")
assert isinstance(openai_chat_client, ChatClient)
assert isinstance(openai_chat_client, ChatClientProtocol)
# Test that the client will use the web search tool
response = openai_chat_client.get_streaming_response(
@@ -392,7 +392,7 @@ async def test_openai_chat_client_web_search_streaming() -> None:
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_agent_basic_run():
"""Test OpenAI chat client agent basic run functionality with OpenAIChatClient."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
) as agent:
# Test basic run
@@ -407,12 +407,12 @@ async def test_openai_chat_client_agent_basic_run():
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_agent_basic_run_streaming():
"""Test OpenAI chat client agent basic streaming functionality with OpenAIChatClient."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
) as agent:
# Test streaming run
full_text = ""
async for chunk in agent.run_streaming("Please respond with exactly: 'This is a streaming response test.'"):
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
full_text += chunk.text
@@ -424,7 +424,7 @@ async def test_openai_chat_client_agent_basic_run_streaming():
@skip_if_openai_integration_tests_disabled
async def test_openai_chat_client_agent_thread_persistence():
"""Test OpenAI chat client agent thread persistence across runs with OpenAIChatClient."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
instructions="You are a helpful assistant with good memory.",
) as agent:
@@ -451,7 +451,7 @@ async def test_openai_chat_client_agent_existing_thread():
# First conversation - capture the thread
preserved_thread = None
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
@@ -467,7 +467,7 @@ async def test_openai_chat_client_agent_existing_thread():
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
@@ -483,7 +483,7 @@ async def test_openai_chat_client_agent_existing_thread():
async def test_openai_chat_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with OpenAI Chat Client."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIChatClient(ai_model_id="gpt-4.1"),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
@@ -518,7 +518,7 @@ async def test_openai_chat_client_run_level_tool_isolation():
call_count += 1
return f"The weather in {location} is sunny and 72°F."
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIChatClient(ai_model_id="gpt-4.1"),
instructions="You are a helpful assistant.",
) as agent:
python/packages/main/tests/openai/test_openai_responses_client.py
+24 -24
View File
@@ -13,12 +13,11 @@ from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
ChatClient,
ChatClientAgent,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
ChatRole,
FunctionCallContent,
FunctionResultContent,
HostedCodeInterpreterTool,
@@ -26,6 +25,7 @@ from agent_framework import (
HostedFileSearchTool,
HostedVectorStoreContent,
HostedWebSearchTool,
Role,
TextContent,
TextReasoningContent,
UriContent,
@@ -87,7 +87,7 @@ def test_init(openai_unit_test_env: dict[str, str]) -> None:
openai_responses_client = OpenAIResponsesClient()
assert openai_responses_client.ai_model_id == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"]
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
def test_init_validation_fail() -> None:
@@ -102,7 +102,7 @@ def test_init_ai_model_id_constructor(openai_unit_test_env: dict[str, str]) -> N
openai_responses_client = OpenAIResponsesClient(ai_model_id=ai_model_id)
assert openai_responses_client.ai_model_id == ai_model_id
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None:
@@ -114,7 +114,7 @@ def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None:
)
assert openai_responses_client.ai_model_id == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"]
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
@@ -731,7 +731,7 @@ def test_streaming_response_basic_structure() -> None:
# Should get a valid ChatResponseUpdate structure
assert isinstance(response, ChatResponseUpdate)
assert response.role == ChatRole.ASSISTANT
assert response.role == Role.ASSISTANT
assert response.ai_model_id == "test-model"
assert isinstance(response.contents, list)
assert response.raw_representation is mock_event
@@ -742,7 +742,7 @@ async def test_openai_responses_client_response() -> None:
"""Test OpenAI chat completion responses."""
openai_responses_client = OpenAIResponsesClient()
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
@@ -785,7 +785,7 @@ async def test_openai_responses_client_response_tools() -> None:
"""Test OpenAI chat completion responses."""
openai_responses_client = OpenAIResponsesClient()
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What is the weather in New York?"))
@@ -824,7 +824,7 @@ async def test_openai_responses_client_streaming() -> None:
"""Test Azure OpenAI chat completion responses."""
openai_responses_client = OpenAIResponsesClient()
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
@@ -877,7 +877,7 @@ async def test_openai_responses_client_streaming_tools() -> None:
"""Test OpenAI chat completion responses."""
openai_responses_client = OpenAIResponsesClient()
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = [ChatMessage(role="user", text="What is the weather in Seattle?")]
@@ -923,7 +923,7 @@ async def test_openai_responses_client_streaming_tools() -> None:
async def test_openai_responses_client_web_search() -> None:
openai_responses_client = OpenAIResponsesClient()
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
# Test that the client will use the web search tool
response = await openai_responses_client.get_response(
@@ -962,7 +962,7 @@ async def test_openai_responses_client_web_search() -> None:
async def test_openai_responses_client_web_search_streaming() -> None:
openai_responses_client = OpenAIResponsesClient()
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
# Test that the client will use the web search tool
response = openai_responses_client.get_streaming_response(
@@ -1015,7 +1015,7 @@ async def test_openai_responses_client_web_search_streaming() -> None:
async def test_openai_responses_client_file_search() -> None:
openai_responses_client = OpenAIResponsesClient()
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
file_id, vector_store = await create_vector_store(openai_responses_client)
# Test that the client will use the web search tool
@@ -1039,7 +1039,7 @@ async def test_openai_responses_client_file_search() -> None:
async def test_openai_responses_client_streaming_file_search() -> None:
openai_responses_client = OpenAIResponsesClient()
assert isinstance(openai_responses_client, ChatClient)
assert isinstance(openai_responses_client, ChatClientProtocol)
file_id, vector_store = await create_vector_store(openai_responses_client)
# Test that the client will use the web search tool
@@ -1088,12 +1088,12 @@ async def test_openai_responses_client_agent_basic_run():
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_basic_run_streaming():
"""Test OpenAI Responses Client agent basic streaming functionality with OpenAIResponsesClient."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
) as agent:
# Test streaming run
full_text = ""
async for chunk in agent.run_streaming("Please respond with exactly: 'This is a streaming response test.'"):
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
full_text += chunk.text
@@ -1105,7 +1105,7 @@ async def test_openai_responses_client_agent_basic_run_streaming():
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_thread_persistence():
"""Test OpenAI Responses Client agent thread persistence across runs with OpenAIResponsesClient."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant with good memory.",
) as agent:
@@ -1128,7 +1128,7 @@ async def test_openai_responses_client_agent_thread_persistence():
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_thread_storage_with_store_true():
"""Test OpenAI Responses Client agent with store=True to verify service_thread_id is returned."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant.",
) as agent:
@@ -1162,7 +1162,7 @@ async def test_openai_responses_client_agent_existing_thread():
# First conversation - capture the thread
preserved_thread = None
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
@@ -1178,7 +1178,7 @@ async def test_openai_responses_client_agent_existing_thread():
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
@@ -1193,7 +1193,7 @@ async def test_openai_responses_client_agent_existing_thread():
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_hosted_code_interpreter_tool():
"""Test OpenAI Responses Client agent with HostedCodeInterpreterTool through OpenAIResponsesClient."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that can execute Python code.",
tools=[HostedCodeInterpreterTool()],
@@ -1215,7 +1215,7 @@ async def test_openai_responses_client_agent_hosted_code_interpreter_tool():
async def test_openai_responses_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with OpenAI Responses Client."""
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
@@ -1250,7 +1250,7 @@ async def test_openai_responses_client_run_level_tool_isolation():
call_count += 1
return f"The weather in {location} is sunny and 72°F."
async with ChatClientAgent(
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant.",
) as agent: