mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: name changes executed (#607)
* 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
committed by
GitHub
Unverified
GitHub
parent
6310ca5be0
commit
40ab6e9d67
@@ -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:
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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")]
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user