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/agent_framework/_agents.py
+32 -32
View File
@@ -9,11 +9,11 @@ from uuid import uuid4
from pydantic import BaseModel, Field, PrivateAttr
from ._clients import ChatClient
from ._mcp import McpTool
from ._clients import ChatClientProtocol
from ._mcp import MCPTool
from ._pydantic import AFBaseModel
from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages
from ._tools import AITool
from ._tools import ToolProtocol
from ._types import (
AgentRunResponse,
AgentRunResponseUpdate,
@@ -21,8 +21,8 @@ from ._types import (
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
ChatToolMode,
Role,
)
from .exceptions import AgentExecutionException
from .telemetry import use_agent_telemetry
@@ -34,14 +34,14 @@ else:
TThreadType = TypeVar("TThreadType", bound="AgentThread")
__all__ = ["AIAgent", "AgentBase", "ChatClientAgent"]
__all__ = ["AgentProtocol", "BaseAgent", "ChatAgent"]
# region Agent Protocol
@runtime_checkable
class AIAgent(Protocol):
class AgentProtocol(Protocol):
"""A protocol for an agent that can be invoked."""
@property
@@ -93,7 +93,7 @@ class AIAgent(Protocol):
"""
...
def run_streaming(
def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
@@ -122,10 +122,10 @@ class AIAgent(Protocol):
...
# region AgentBase
# region BaseAgent
class AgentBase(AFBaseModel):
class BaseAgent(AFBaseModel):
"""Base class for all Agent Framework agents.
Attributes:
@@ -167,24 +167,24 @@ class AgentBase(AFBaseModel):
return thread
# region ChatClientAgent
# region ChatAgent
@use_agent_telemetry
class ChatClientAgent(AgentBase):
class ChatAgent(BaseAgent):
"""A Chat Client Agent."""
AGENT_SYSTEM_NAME: ClassVar[str] = "microsoft.agent_framework"
chat_client: ChatClient
chat_client: ChatClientProtocol
instructions: str | None = None
chat_options: ChatOptions
chat_message_store_factory: Callable[[], ChatMessageStore] | None = None
_local_mcp_tools: list[McpTool] = PrivateAttr(default_factory=list) # type: ignore[reportUnknownVariableType]
_local_mcp_tools: list[MCPTool] = PrivateAttr(default_factory=list) # type: ignore[reportUnknownVariableType]
_async_exit_stack: AsyncExitStack = PrivateAttr(default_factory=AsyncExitStack)
def __init__(
self,
chat_client: ChatClient,
chat_client: ChatClientProtocol,
instructions: str | None = None,
*,
id: str | None = None,
@@ -202,10 +202,10 @@ class ChatClientAgent(AgentBase):
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tools: AITool
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
@@ -213,7 +213,7 @@ class ChatClientAgent(AgentBase):
chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
**kwargs: Any,
) -> None:
"""Create a ChatClientAgent.
"""Create a ChatAgent.
Remarks:
The set of attributes from frequency_penalty to additional_properties are used to
@@ -253,8 +253,8 @@ class ChatClientAgent(AgentBase):
# We ignore the MCP Servers here and store them separately,
# we add their functions to the tools list at runtime
normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
local_mcp_tools = [tool for tool in normalized_tools if isinstance(tool, McpTool)]
final_tools = [tool for tool in normalized_tools if not isinstance(tool, McpTool)]
local_mcp_tools = [tool for tool in normalized_tools if isinstance(tool, MCPTool)]
final_tools = [tool for tool in normalized_tools if not isinstance(tool, MCPTool)]
args: dict[str, Any] = {
"chat_client": chat_client,
"chat_message_store_factory": chat_message_store_factory,
@@ -337,8 +337,8 @@ class ChatClientAgent(AgentBase):
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
tools: AITool
| list[AITool]
tools: ToolProtocol
| list[ToolProtocol]
| Callable[..., Any]
| list[Callable[..., Any]]
| MutableMapping[str, Any]
@@ -384,11 +384,11 @@ class ChatClientAgent(AgentBase):
agent_name = self._get_agent_name()
# Resolve final tool list (runtime provided tools + local MCP server tools)
final_tools: list[AITool | Callable[..., Any] | dict[str, Any]] = []
final_tools: list[ToolProtocol | Callable[..., Any] | dict[str, Any]] = []
# Normalize tools argument to a list without mutating the original parameter
normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
for tool in normalized_tools:
if isinstance(tool, McpTool):
if isinstance(tool, MCPTool):
final_tools.extend(tool.functions) # type: ignore
else:
final_tools.append(tool) # type: ignore
@@ -442,7 +442,7 @@ class ChatClientAgent(AgentBase):
additional_properties=response.additional_properties,
)
async def run_streaming(
async def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
@@ -459,10 +459,10 @@ class ChatClientAgent(AgentBase):
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
tools: AITool
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
@@ -472,7 +472,7 @@ class ChatClientAgent(AgentBase):
"""Stream the agent with the given messages and options.
Remarks:
Since you won't always call the agent.run_streaming directly, but it get's called
Since you won't always call the agent.run_stream directly, but it get's called
through orchestration, it is advised to set your default values for
all the chat client parameters in the agent constructor.
If both parameters are used, the ones passed to the run methods take precedence.
@@ -506,11 +506,11 @@ class ChatClientAgent(AgentBase):
response_updates: list[ChatResponseUpdate] = []
# Resolve final tool list (runtime provided tools + local MCP server tools)
final_tools: list[AITool | MutableMapping[str, Any] | Callable[..., Any]] = []
final_tools: list[ToolProtocol | MutableMapping[str, Any] | Callable[..., Any]] = []
# Normalize tools argument to a list without mutating the original parameter
normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
for tool in normalized_tools:
if isinstance(tool, McpTool):
if isinstance(tool, MCPTool):
final_tools.extend(tool.functions) # type: ignore
else:
final_tools.append(tool)
@@ -627,7 +627,7 @@ class ChatClientAgent(AgentBase):
messages: list[ChatMessage] = []
if self.instructions:
messages.append(ChatMessage(role=ChatRole.SYSTEM, text=self.instructions))
messages.append(ChatMessage(role=Role.SYSTEM, text=self.instructions))
if thread.message_store:
messages.extend(await thread.message_store.list_messages() or [])
messages.extend(input_messages or [])
@@ -641,12 +641,12 @@ class ChatClientAgent(AgentBase):
return []
if isinstance(messages, str):
return [ChatMessage(role=ChatRole.USER, text=messages)]
return [ChatMessage(role=Role.USER, text=messages)]
if isinstance(messages, ChatMessage):
return [messages]
return [ChatMessage(role=ChatRole.USER, text=msg) if isinstance(msg, str) else msg for msg in messages]
return [ChatMessage(role=Role.USER, text=msg) if isinstance(msg, str) else msg for msg in messages]
def _get_agent_name(self) -> str:
return self.name or "UnnamedAgent"
python/packages/main/agent_framework/_clients.py
+29 -29
View File
@@ -11,31 +11,31 @@ from pydantic import BaseModel
from ._logging import get_logger
from ._pydantic import AFBaseModel
from ._threads import ChatMessageStore
from ._tools import AIFunction, AITool
from ._tools import AIFunction, ToolProtocol
from ._types import (
AIContents,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatToolMode,
Contents,
FunctionCallContent,
FunctionResultContent,
GeneratedEmbeddings,
)
if TYPE_CHECKING:
from ._agents import ChatClientAgent
from ._agents import ChatAgent
TInput = TypeVar("TInput", contravariant=True)
TEmbedding = TypeVar("TEmbedding")
TChatClientBase = TypeVar("TChatClientBase", bound="ChatClientBase")
TBaseChatClient = TypeVar("TBaseChatClient", bound="BaseChatClient")
logger = get_logger()
__all__ = [
"ChatClient",
"ChatClientBase",
"BaseChatClient",
"ChatClientProtocol",
"EmbeddingGenerator",
"use_tool_calling",
]
@@ -50,7 +50,7 @@ async def _auto_invoke_function(
tool_map: dict[str, AIFunction[BaseModel, Any]],
sequence_index: int | None = None,
request_index: int | None = None,
) -> AIContents:
) -> Contents:
"""Invoke a function call requested by the agent, applying filters that are defined in the agent."""
tool: AIFunction[BaseModel, Any] | None = tool_map.get(function_call_content.name)
if tool is None:
@@ -81,7 +81,7 @@ def _tool_call_non_streaming(
@wraps(func)
async def wrapper(
self: "ChatClientBase",
self: "BaseChatClient",
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
@@ -157,7 +157,7 @@ def _tool_call_streaming(
@wraps(func)
async def wrapper(
self: "ChatClientBase",
self: "BaseChatClient",
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
@@ -217,11 +217,11 @@ def _tool_call_streaming(
return wrapper
def use_tool_calling(cls: type[TChatClientBase]) -> type[TChatClientBase]:
def use_tool_calling(cls: type[TBaseChatClient]) -> type[TBaseChatClient]:
"""Class decorator that enables tool calling for a chat client.
Remarks:
This only works on classes that derive from ChatClientBase
This only works on classes that derive from BaseChatClient
and the `_inner_get_response`
and `_inner_get_streaming_response` methods.
It also sets a `__maximum_iterations_per_request` attribute on the class.
@@ -247,11 +247,11 @@ def use_tool_calling(cls: type[TChatClientBase]) -> type[TChatClientBase]:
return cls
# region ChatClient Protocol
# region ChatClientProtocol Protocol
@runtime_checkable
class ChatClient(Protocol):
class ChatClientProtocol(Protocol):
"""A protocol for a chat client that can generate responses."""
async def get_response(
@@ -270,10 +270,10 @@ class ChatClient(Protocol):
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tools: AITool
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
@@ -327,10 +327,10 @@ class ChatClient(Protocol):
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tools: AITool
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
@@ -370,7 +370,7 @@ class ChatClient(Protocol):
...
class ChatClientBase(AFBaseModel, ABC):
class BaseChatClient(AFBaseModel, ABC):
"""Base class for chat clients."""
MODEL_PROVIDER_NAME: str = "unknown"
@@ -457,10 +457,10 @@ class ChatClientBase(AFBaseModel, ABC):
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tools: AITool
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
@@ -537,10 +537,10 @@ class ChatClientBase(AFBaseModel, ABC):
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tools: AITool
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
@@ -633,14 +633,14 @@ class ChatClientBase(AFBaseModel, ABC):
*,
name: str | None = None,
instructions: str | None = None,
tools: AITool
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
**kwargs: Any,
) -> "ChatClientAgent":
) -> "ChatAgent":
"""Create an agent with the given name and instructions.
Args:
@@ -650,14 +650,14 @@ class ChatClientBase(AFBaseModel, ABC):
chat_message_store_factory: Factory function to create an instance of ChatMessageStore. If not provided,
the default in-memory store will be used.
**kwargs: Additional keyword arguments to pass to the agent.
See ChatClientAgent for all the available options.
See ChatAgent for all the available options.
Returns:
An instance of ChatClientAgent.
An instance of ChatAgent.
"""
from ._agents import ChatClientAgent
from ._agents import ChatAgent
return ChatClientAgent(
return ChatAgent(
chat_client=self,
name=name,
instructions=instructions,
python/packages/main/agent_framework/_mcp.py
+24 -24
View File
@@ -22,7 +22,7 @@ from mcp.shared.session import RequestResponder
from pydantic import BaseModel, create_model
from ._tools import AIFunction
from ._types import AIContents, ChatMessage, ChatRole, DataContent, TextContent, UriContent
from ._types import ChatMessage, Contents, DataContent, Role, TextContent, UriContent
from .exceptions import ToolException, ToolExecutionException
if sys.version_info >= (3, 11):
@@ -31,7 +31,7 @@ else:
from typing_extensions import Self # pragma: no cover
if TYPE_CHECKING:
from ._clients import ChatClient
from ._clients import ChatClientProtocol
logger = logging.getLogger(__name__)
@@ -49,10 +49,10 @@ LOG_LEVEL_MAPPING: dict[types.LoggingLevel, int] = {
}
__all__ = [
"McpSseTools",
"McpStdioTool",
"McpStreamableHttpTool",
"McpWebsocketTool",
"MCPSseTools",
"MCPStdioTool",
"MCPStreamableHTTPTool",
"MCPWebsocketTool",
]
@@ -61,7 +61,7 @@ def _mcp_prompt_message_to_chat_message(
) -> ChatMessage:
"""Convert a MCP container type to a Agent Framework type."""
return ChatMessage(
role=ChatRole(value=mcp_type.role),
role=Role(value=mcp_type.role),
contents=[_mcp_type_to_ai_content(mcp_type.content)], # type: ignore[call-arg]
raw_representation=mcp_type,
)
@@ -69,14 +69,14 @@ def _mcp_prompt_message_to_chat_message(
def _mcp_call_tool_result_to_ai_contents(
mcp_type: types.CallToolResult,
) -> list[AIContents]:
) -> list[Contents]:
"""Convert a MCP container type to a Agent Framework type."""
return [_mcp_type_to_ai_content(item) for item in mcp_type.content]
def _mcp_type_to_ai_content(
mcp_type: types.ImageContent | types.TextContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink,
) -> AIContents:
) -> Contents:
"""Convert a MCP type to a Agent Framework type."""
match mcp_type:
case types.TextContent():
@@ -105,9 +105,9 @@ def _mcp_type_to_ai_content(
def _ai_content_to_mcp_types(
content: AIContents,
content: Contents,
) -> types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink | None:
"""Convert a AIContent type to a MCP type."""
"""Convert a BaseContent type to a MCP type."""
match content:
case TextContent():
return types.TextContent(type="text", text=content.text)
@@ -223,7 +223,7 @@ def _normalize_mcp_name(name: str) -> str:
# region: MCP Plugin
class McpTool:
class MCPTool:
"""Base class with the MCP logic."""
def __init__(
@@ -235,7 +235,7 @@ class McpTool:
load_prompts: bool = True,
session: ClientSession | None = None,
request_timeout: int | None = None,
chat_client: "ChatClient | None" = None,
chat_client: "ChatClientProtocol | None" = None,
) -> None:
"""Initialize the MCP Plugin Base."""
self.name = name
@@ -250,7 +250,7 @@ class McpTool:
self.functions: list[AIFunction[Any, Any]] = []
def __str__(self) -> str:
return f"McpTool(name={self.name}, description={self.description})"
return f"MCPTool(name={self.name}, description={self.description})"
async def connect(self) -> None:
"""Connect to the MCP server."""
@@ -424,7 +424,7 @@ class McpTool:
local_name = _normalize_mcp_name(tool.name)
input_model = _get_input_model_from_mcp_tool(tool)
# Create AIFunctions out of each tool
func: AIFunction[BaseModel, list[AIContents]] = AIFunction(
func: AIFunction[BaseModel, list[Contents]] = AIFunction(
func=partial(self.call_tool, tool.name),
name=local_name,
description=tool.description or "",
@@ -442,7 +442,7 @@ class McpTool:
"""Get an MCP client."""
pass
async def call_tool(self, tool_name: str, **kwargs: Any) -> list[AIContents]:
async def call_tool(self, tool_name: str, **kwargs: Any) -> list[Contents]:
"""Call a tool with the given arguments."""
if not self.session:
raise ToolExecutionException("MCP server not connected, please call connect() before using this method.")
@@ -494,7 +494,7 @@ class McpTool:
# region: MCP Plugin Implementations
class McpStdioTool(McpTool):
class MCPStdioTool(MCPTool):
"""MCP stdio server configuration."""
def __init__(
@@ -511,7 +511,7 @@ class McpStdioTool(McpTool):
args: list[str] | None = None,
env: dict[str, str] | None = None,
encoding: str | None = None,
chat_client: "ChatClient | None" = None,
chat_client: "ChatClientProtocol | None" = None,
**kwargs: Any,
) -> None:
"""Initialize the MCP stdio plugin.
@@ -567,7 +567,7 @@ class McpStdioTool(McpTool):
return stdio_client(server=StdioServerParameters(**args))
class McpSseTools(McpTool):
class MCPSseTools(MCPTool):
"""MCP sse server configuration."""
def __init__(
@@ -584,7 +584,7 @@ class McpSseTools(McpTool):
headers: dict[str, Any] | None = None,
timeout: float | None = None,
sse_read_timeout: float | None = None,
chat_client: "ChatClient | None" = None,
chat_client: "ChatClientProtocol | None" = None,
**kwargs: Any,
) -> None:
"""Initialize the MCP sse plugin.
@@ -643,7 +643,7 @@ class McpSseTools(McpTool):
return sse_client(**args)
class McpStreamableHttpTool(McpTool):
class MCPStreamableHTTPTool(MCPTool):
"""MCP streamable http server configuration."""
def __init__(
@@ -661,7 +661,7 @@ class McpStreamableHttpTool(McpTool):
timeout: float | None = None,
sse_read_timeout: float | None = None,
terminate_on_close: bool | None = None,
chat_client: "ChatClient | None" = None,
chat_client: "ChatClientProtocol | None" = None,
**kwargs: Any,
) -> None:
"""Initialize the MCP streamable http plugin.
@@ -723,7 +723,7 @@ class McpStreamableHttpTool(McpTool):
return streamablehttp_client(**args)
class McpWebsocketTool(McpTool):
class MCPWebsocketTool(MCPTool):
"""MCP websocket server configuration."""
def __init__(
@@ -737,7 +737,7 @@ class McpWebsocketTool(McpTool):
session: ClientSession | None = None,
description: str | None = None,
additional_properties: dict[str, Any] | None = None,
chat_client: "ChatClient | None" = None,
chat_client: "ChatClientProtocol | None" = None,
**kwargs: Any,
) -> None:
"""Initialize the MCP websocket plugin.
python/packages/main/agent_framework/_pydantic.py
+2 -2
View File
@@ -7,9 +7,9 @@ from pydantic import BaseModel, ConfigDict, Field, UrlConstraints
from pydantic.networks import AnyUrl
from pydantic_settings import BaseSettings, SettingsConfigDict
HttpsUrl = Annotated[AnyUrl, UrlConstraints(max_length=2083, allowed_schemes=["https"])]
HTTPsUrl = Annotated[AnyUrl, UrlConstraints(max_length=2083, allowed_schemes=["https"])]
__all__ = ["AFBaseModel", "AFBaseSettings", "HttpsUrl"]
__all__ = ["AFBaseModel", "AFBaseSettings", "HTTPsUrl"]
class AFBaseModel(BaseModel):
python/packages/main/agent_framework/_tools.py
+22 -22
View File
@@ -24,7 +24,7 @@ from ._pydantic import AFBaseModel
from .telemetry import GenAIAttributes, start_as_current_span
if TYPE_CHECKING:
from ._types import AIContents
from ._types import Contents
tracer: trace.Tracer = trace.get_tracer("agent_framework")
meter: metrics.Meter = metrics.get_meter_provider().get_meter("agent_framework")
@@ -32,24 +32,24 @@ logger = get_logger()
__all__ = [
"AIFunction",
"AITool",
"HostedCodeInterpreterTool",
"HostedFileSearchTool",
"HostedWebSearchTool",
"ToolProtocol",
"ai_function",
]
def _parse_inputs(
inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None",
) -> list["AIContents"]:
"""Parse the inputs for a tool, ensuring they are of type AIContents."""
inputs: "Contents | dict[str, Any] | str | list[Contents | dict[str, Any] | str] | None",
) -> list["Contents"]:
"""Parse the inputs for a tool, ensuring they are of type Contents."""
if inputs is None:
return []
from ._types import AIContent, DataContent, HostedFileContent, HostedVectorStoreContent, UriContent
from ._types import BaseContent, DataContent, HostedFileContent, HostedVectorStoreContent, UriContent
parsed_inputs: list["AIContents"] = []
parsed_inputs: list["Contents"] = []
if not isinstance(inputs, list):
inputs = [inputs]
for input_item in inputs:
@@ -75,15 +75,15 @@ def _parse_inputs(
parsed_inputs.append(DataContent(**input_item))
else:
raise ValueError(f"Unsupported input type: {input_item}")
elif isinstance(input_item, AIContent):
elif isinstance(input_item, BaseContent):
parsed_inputs.append(input_item)
else:
raise TypeError(f"Unsupported input type: {type(input_item).__name__}. Expected AIContents or dict.")
raise TypeError(f"Unsupported input type: {type(input_item).__name__}. Expected Contents or dict.")
return parsed_inputs
@runtime_checkable
class AITool(Protocol):
class ToolProtocol(Protocol):
"""Represents a generic tool that can be specified to an AI service.
Attributes:
@@ -111,7 +111,7 @@ ArgsT = TypeVar("ArgsT", bound=BaseModel)
ReturnT = TypeVar("ReturnT")
class AIToolBase(AFBaseModel):
class BaseTool(AFBaseModel):
"""Base class for AI tools, providing common attributes and methods.
Args:
@@ -131,7 +131,7 @@ class AIToolBase(AFBaseModel):
return f"{self.__class__.__name__}(name={self.name})"
class HostedCodeInterpreterTool(AIToolBase):
class HostedCodeInterpreterTool(BaseTool):
"""Represents a hosted tool that can be specified to an AI service to enable it to execute generated code.
This tool does not implement code interpretation itself. It serves as a marker to inform a service
@@ -143,7 +143,7 @@ class HostedCodeInterpreterTool(AIToolBase):
def __init__(
self,
*,
inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None" = None,
inputs: "Contents | dict[str, Any] | str | list[Contents | dict[str, Any] | str] | None" = None,
description: str | None = None,
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
@@ -155,8 +155,8 @@ class HostedCodeInterpreterTool(AIToolBase):
This should mostly be HostedFileContent or HostedVectorStoreContent.
Can also be DataContent, depending on the service used.
When supplying a list, it can contain:
- AIContents instances
- dicts with properties for AIContents (e.g., {"uri": "http://example.com", "media_type": "text/html"})
- Contents instances
- dicts with properties for Contents (e.g., {"uri": "http://example.com", "media_type": "text/html"})
- strings (which will be converted to UriContent with media_type "text/plain").
If None, defaults to an empty list.
description: A description of the tool.
@@ -177,7 +177,7 @@ class HostedCodeInterpreterTool(AIToolBase):
super().__init__(**args, **kwargs)
class HostedWebSearchTool(AIToolBase):
class HostedWebSearchTool(BaseTool):
"""Represents a web search tool that can be specified to an AI service to enable it to perform web searches."""
def __init__(
@@ -206,7 +206,7 @@ class HostedWebSearchTool(AIToolBase):
super().__init__(**args, **kwargs)
class HostedFileSearchTool(AIToolBase):
class HostedFileSearchTool(BaseTool):
"""Represents a file search tool that can be specified to an AI service to enable it to perform file searches."""
inputs: list[Any] | None = None
@@ -214,7 +214,7 @@ class HostedFileSearchTool(AIToolBase):
def __init__(
self,
inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None" = None,
inputs: "Contents | dict[str, Any] | str | list[Contents | dict[str, Any] | str] | None" = None,
max_results: int | None = None,
description: str | None = None,
additional_properties: dict[str, Any] | None = None,
@@ -226,8 +226,8 @@ class HostedFileSearchTool(AIToolBase):
inputs: A list of contents that the tool can accept as input. Defaults to None.
This should be one or more HostedVectorStoreContents.
When supplying a list, it can contain:
- AIContents instances
- dicts with properties for AIContents (e.g., {"uri": "http://example.com", "media_type": "text/html"})
- Contents instances
- dicts with properties for Contents (e.g., {"uri": "http://example.com", "media_type": "text/html"})
- strings (which will be converted to UriContent with media_type "text/plain").
If None, defaults to an empty list.
max_results: The maximum number of results to return from the file search.
@@ -252,8 +252,8 @@ class HostedFileSearchTool(AIToolBase):
super().__init__(**args, **kwargs)
class AIFunction(AIToolBase, Generic[ArgsT, ReturnT]):
"""A AITool that is callable as code.
class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
"""A ToolProtocol that is callable as code.
Args:
name: The name of the function.
python/packages/main/agent_framework/_types.py
+89 -106
View File
@@ -28,7 +28,7 @@ from pydantic import (
from ._logging import get_logger
from ._pydantic import AFBaseModel
from ._tools import AITool, ai_function
from ._tools import ToolProtocol, ai_function
from .exceptions import AgentFrameworkException
if sys.version_info >= (3, 11):
@@ -77,29 +77,28 @@ KNOWN_MEDIA_TYPES = [
__all__ = [
"AIAnnotation",
"AIAnnotations",
"AIContent",
"AIContents",
"AgentRunResponse",
"AgentRunResponseUpdate",
"AnnotatedRegion",
"AnnotatedRegions",
"ChatFinishReason",
"Annotations",
"BaseAnnotation",
"BaseContent",
"ChatMessage",
"ChatOptions",
"ChatResponse",
"ChatResponseUpdate",
"ChatRole",
"ChatToolMode",
"CitationAnnotation",
"Contents",
"DataContent",
"ErrorContent",
"FinishReason",
"FunctionCallContent",
"FunctionResultContent",
"GeneratedEmbeddings",
"HostedFileContent",
"HostedVectorStoreContent",
"Role",
"SpeechToTextOptions",
"TextContent",
"TextReasoningContent",
@@ -231,7 +230,7 @@ def _process_update(
is_new_message = True
if is_new_message:
message = ChatMessage(role=ChatRole.ASSISTANT, contents=[])
message = ChatMessage(role=Role.ASSISTANT, contents=[])
response.messages.append(message)
else:
message = response.messages[-1]
@@ -278,12 +277,12 @@ def _process_update(
def _coalesce_text_content(
contents: list["AIContents"], type_: type["TextContent"] | type["TextReasoningContent"]
contents: list["Contents"], type_: type["TextContent"] | type["TextReasoningContent"]
) -> None:
"""Take any subsequence Text or TextReasoningContent items and coalesce them into a single item."""
if not contents:
return
coalesced_contents: list["AIContents"] = []
coalesced_contents: list["Contents"] = []
first_new_content: Any | None = None
for content in contents:
if isinstance(content, type_):
@@ -313,22 +312,10 @@ def _finalize_response(response: "ChatResponse | AgentRunResponse") -> None:
_coalesce_text_content(msg.contents, TextReasoningContent)
# region AIAnnotation
# region BaseAnnotation
class AnnotatedRegion(AFBaseModel):
"""Represents a collection of annotated regions.
Attributes:
regions: A list of regions that have been annotated.
additional_properties: Optional additional properties associated with the content.
raw_representation: Optional raw representation of the content from an underlying implementation.
"""
type: Literal["annotated_regions"] = "annotated_regions" # type: ignore[assignment]
class TextSpanRegion(AnnotatedRegion):
class TextSpanRegion(AFBaseModel):
"""Represents a region of text that has been annotated."""
type: Literal["text_span"] = "text_span" # type: ignore[assignment]
@@ -337,28 +324,26 @@ class TextSpanRegion(AnnotatedRegion):
AnnotatedRegions = Annotated[
TextSpanRegion | AnnotatedRegion,
TextSpanRegion,
Field(discriminator="type"),
]
class AIAnnotation(AFBaseModel):
class BaseAnnotation(AFBaseModel):
"""Base class for all AI Annotation types.
Args:
type: The type of content, which is always "ai_annotation" for this class.
additional_properties: Optional additional properties associated with the content.
raw_representation: Optional raw representation of the content from an underlying implementation.
"""
type: Literal["ai_annotation"] = "ai_annotation"
annotated_regions: list[AnnotatedRegions] | None = None
additional_properties: dict[str, Any] | None = None
raw_representation: Any | None = Field(default=None, repr=False)
class CitationAnnotation(AIAnnotation):
class CitationAnnotation(BaseAnnotation):
"""Represents a citation annotation.
Attributes:
@@ -381,33 +366,31 @@ class CitationAnnotation(AIAnnotation):
snippet: str | None = None
AIAnnotations = Annotated[
CitationAnnotation | AIAnnotation,
Annotations = Annotated[
CitationAnnotation,
Field(discriminator="type"),
]
# region AIContent
# region BaseContent
class AIContent(AFBaseModel):
class BaseContent(AFBaseModel):
"""Represents content used by AI services.
Attributes:
type: The type of content, which is always "ai" for this class.
annotations: Optional annotations associated with the content.
additional_properties: Optional additional properties associated with the content.
raw_representation: Optional raw representation of the content from an underlying implementation.
"""
type: Literal["ai"] = "ai"
annotations: list[AIAnnotations] | None = None
annotations: list[Annotations] | None = None
additional_properties: dict[str, Any] | None = None
raw_representation: Any | None = Field(default=None, repr=False, exclude=True)
class TextContent(AIContent):
class TextContent(BaseContent):
"""Represents text content in a chat.
Attributes:
@@ -508,7 +491,7 @@ class TextContent(AIContent):
return self
class TextReasoningContent(AIContent):
class TextReasoningContent(BaseContent):
"""Represents text reasoning content in a chat.
Remarks:
@@ -609,7 +592,7 @@ class TextReasoningContent(AIContent):
return self
class DataContent(AIContent):
class DataContent(BaseContent):
"""Represents binary data content with an associated media type (also known as a MIME type).
Attributes:
@@ -632,7 +615,7 @@ class DataContent(AIContent):
self,
*,
uri: str,
annotations: list[AIAnnotations] | None = None,
annotations: list[Annotations] | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
**kwargs: Any,
@@ -658,7 +641,7 @@ class DataContent(AIContent):
*,
data: bytes,
media_type: str,
annotations: list[AIAnnotations] | None = None,
annotations: list[Annotations] | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
**kwargs: Any,
@@ -685,7 +668,7 @@ class DataContent(AIContent):
uri: str | None = None,
data: bytes | None = None,
media_type: str | None = None,
annotations: list[AIAnnotations] | None = None,
annotations: list[Annotations] | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
**kwargs: Any,
@@ -739,7 +722,7 @@ class DataContent(AIContent):
return _has_top_level_media_type(self.media_type, top_level_media_type)
class UriContent(AIContent):
class UriContent(BaseContent):
"""Represents a URI content.
Remarks:
@@ -765,7 +748,7 @@ class UriContent(AIContent):
uri: str,
media_type: str,
*,
annotations: list[AIAnnotations] | None = None,
annotations: list[Annotations] | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
**kwargs: Any,
@@ -807,7 +790,7 @@ def _has_top_level_media_type(media_type: str | None, top_level_media_type: str)
return span.lower() == top_level_media_type.lower()
class ErrorContent(AIContent):
class ErrorContent(BaseContent):
"""Represents an error.
Remarks:
@@ -837,7 +820,7 @@ class ErrorContent(AIContent):
message: str | None = None,
error_code: str | None = None,
details: str | None = None,
annotations: list[AIAnnotations] | None = None,
annotations: list[Annotations] | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
**kwargs: Any,
@@ -868,7 +851,7 @@ class ErrorContent(AIContent):
return f"Error {self.error_code}: {self.message}" if self.error_code else self.message or "Unknown error"
class FunctionCallContent(AIContent):
class FunctionCallContent(BaseContent):
"""Represents a function call request.
Attributes:
@@ -896,7 +879,7 @@ class FunctionCallContent(AIContent):
name: str,
arguments: str | dict[str, Any | None] | None = None,
exception: Exception | None = None,
annotations: list[AIAnnotations] | None = None,
annotations: list[Annotations] | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
**kwargs: Any,
@@ -962,7 +945,7 @@ class FunctionCallContent(AIContent):
)
class FunctionResultContent(AIContent):
class FunctionResultContent(BaseContent):
"""Represents the result of a function call.
Attributes:
@@ -987,7 +970,7 @@ class FunctionResultContent(AIContent):
call_id: str,
result: Any | None = None,
exception: Exception | None = None,
annotations: list[AIAnnotations] | None = None,
annotations: list[Annotations] | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
**kwargs: Any,
@@ -1014,7 +997,7 @@ class FunctionResultContent(AIContent):
)
class UsageContent(AIContent):
class UsageContent(BaseContent):
"""Represents usage information associated with a chat request and response.
Attributes:
@@ -1033,7 +1016,7 @@ class UsageContent(AIContent):
self,
details: UsageDetails,
*,
annotations: list[AIAnnotations] | None = None,
annotations: list[Annotations] | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
**kwargs: Any,
@@ -1048,7 +1031,7 @@ class UsageContent(AIContent):
)
class HostedFileContent(AIContent):
class HostedFileContent(BaseContent):
"""Represents a hosted file content.
Attributes:
@@ -1079,7 +1062,7 @@ class HostedFileContent(AIContent):
)
class HostedVectorStoreContent(AIContent):
class HostedVectorStoreContent(BaseContent):
"""Represents a hosted vector store content.
Attributes:
@@ -1110,7 +1093,7 @@ class HostedVectorStoreContent(AIContent):
)
AIContents = Annotated[
Contents = Annotated[
TextContent
| DataContent
| TextReasoningContent
@@ -1127,7 +1110,7 @@ AIContents = Annotated[
# region Chat Response constants
class ChatRole(AFBaseModel):
class Role(AFBaseModel):
"""Describes the intended purpose of a message within a chat interaction.
Attributes:
@@ -1157,19 +1140,19 @@ class ChatRole(AFBaseModel):
def __repr__(self) -> str:
"""Returns the string representation of the role."""
return f"ChatRole(value={self.value!r})"
return f"Role(value={self.value!r})"
# Note: ClassVar is used to indicate that these are class-level constants, not instance attributes.
# The type: ignore[assignment] is used to suppress the type checker warning about assigning to a ClassVar,
# it gets assigned immediately after the class definition.
ChatRole.SYSTEM = ChatRole(value="system") # type: ignore[assignment]
ChatRole.USER = ChatRole(value="user") # type: ignore[assignment]
ChatRole.ASSISTANT = ChatRole(value="assistant") # type: ignore[assignment]
ChatRole.TOOL = ChatRole(value="tool") # type: ignore[assignment]
Role.SYSTEM = Role(value="system") # type: ignore[assignment]
Role.USER = Role(value="user") # type: ignore[assignment]
Role.ASSISTANT = Role(value="assistant") # type: ignore[assignment]
Role.TOOL = Role(value="tool") # type: ignore[assignment]
class ChatFinishReason(AFBaseModel):
class FinishReason(AFBaseModel):
"""Represents the reason a chat response completed.
Attributes:
@@ -1179,21 +1162,21 @@ class ChatFinishReason(AFBaseModel):
value: str
CONTENT_FILTER: ClassVar[Self] # type: ignore[assignment]
"""A ChatFinishReason representing the model filtering content, whether for safety, prohibited content,
"""A FinishReason representing the model filtering content, whether for safety, prohibited content,
sensitive content, or other such issues."""
LENGTH: ClassVar[Self] # type: ignore[assignment]
"""A ChatFinishReason representing the model reaching the maximum length allowed for the request and/or
"""A FinishReason representing the model reaching the maximum length allowed for the request and/or
response (typically in terms of tokens)."""
STOP: ClassVar[Self] # type: ignore[assignment]
"""A ChatFinishReason representing the model encountering a natural stop point or provided stop sequence."""
"""A FinishReason representing the model encountering a natural stop point or provided stop sequence."""
TOOL_CALLS: ClassVar[Self] # type: ignore[assignment]
"""A ChatFinishReason representing the model requesting the use of a tool that was defined in the request."""
"""A FinishReason representing the model requesting the use of a tool that was defined in the request."""
ChatFinishReason.CONTENT_FILTER = ChatFinishReason(value="content_filter") # type: ignore[assignment]
ChatFinishReason.LENGTH = ChatFinishReason(value="length") # type: ignore[assignment]
ChatFinishReason.STOP = ChatFinishReason(value="stop") # type: ignore[assignment]
ChatFinishReason.TOOL_CALLS = ChatFinishReason(value="tool_calls") # type: ignore[assignment]
FinishReason.CONTENT_FILTER = FinishReason(value="content_filter") # type: ignore[assignment]
FinishReason.LENGTH = FinishReason(value="length") # type: ignore[assignment]
FinishReason.STOP = FinishReason(value="stop") # type: ignore[assignment]
FinishReason.TOOL_CALLS = FinishReason(value="tool_calls") # type: ignore[assignment]
# region ChatMessage
@@ -1211,9 +1194,9 @@ class ChatMessage(AFBaseModel):
"""
role: ChatRole
role: Role
"""The role of the author of the message."""
contents: list[AIContents]
contents: list[Contents]
"""The chat message content items."""
author_name: str | None
"""The name of the author of the message."""
@@ -1227,7 +1210,7 @@ class ChatMessage(AFBaseModel):
@overload
def __init__(
self,
role: ChatRole | Literal["system", "user", "assistant", "tool"],
role: Role | Literal["system", "user", "assistant", "tool"],
*,
text: str,
author_name: str | None = None,
@@ -1249,9 +1232,9 @@ class ChatMessage(AFBaseModel):
@overload
def __init__(
self,
role: ChatRole | Literal["system", "user", "assistant", "tool"],
role: Role | Literal["system", "user", "assistant", "tool"],
*,
contents: MutableSequence[AIContents],
contents: MutableSequence[Contents],
author_name: str | None = None,
message_id: str | None = None,
additional_properties: dict[str, Any] | None = None,
@@ -1261,7 +1244,7 @@ class ChatMessage(AFBaseModel):
Args:
role: The role of the author of the message.
contents: Optional list of AIContent items to include in the message.
contents: Optional list of BaseContent items to include in the message.
author_name: Optional name of the author of the message.
message_id: Optional ID of the chat message.
additional_properties: Optional additional properties associated with the chat message.
@@ -1270,10 +1253,10 @@ class ChatMessage(AFBaseModel):
def __init__(
self,
role: ChatRole | Literal["system", "user", "assistant", "tool"],
role: Role | Literal["system", "user", "assistant", "tool"],
*,
text: str | None = None,
contents: MutableSequence[AIContents] | None = None,
contents: MutableSequence[Contents] | None = None,
author_name: str | None = None,
message_id: str | None = None,
additional_properties: dict[str, Any] | None = None,
@@ -1284,7 +1267,7 @@ class ChatMessage(AFBaseModel):
if text is not None:
contents.append(TextContent(text=text))
if isinstance(role, str):
role = ChatRole(value=role)
role = Role(value=role)
super().__init__(
role=role, # type: ignore[reportCallIssue]
contents=contents, # type: ignore[reportCallIssue]
@@ -1334,7 +1317,7 @@ class ChatResponse(AFBaseModel):
"""The model ID used in the creation of the chat response."""
created_at: CreatedAtT | None = None # use a datetimeoffset type?
"""A timestamp for the chat response."""
finish_reason: ChatFinishReason | None = None
finish_reason: FinishReason | None = None
"""The reason for the chat response."""
usage_details: UsageDetails | None = None
"""The usage details for the chat response."""
@@ -1354,7 +1337,7 @@ class ChatResponse(AFBaseModel):
conversation_id: str | None = None,
model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: ChatFinishReason | None = None,
finish_reason: FinishReason | None = None,
usage_details: UsageDetails | None = None,
value: Any | None = None,
response_format: type[BaseModel] | None = None,
@@ -1389,7 +1372,7 @@ class ChatResponse(AFBaseModel):
conversation_id: str | None = None,
model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: ChatFinishReason | None = None,
finish_reason: FinishReason | None = None,
usage_details: UsageDetails | None = None,
value: Any | None = None,
response_format: type[BaseModel] | None = None,
@@ -1424,7 +1407,7 @@ class ChatResponse(AFBaseModel):
conversation_id: str | None = None,
model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: ChatFinishReason | None = None,
finish_reason: FinishReason | None = None,
usage_details: UsageDetails | None = None,
value: Any | None = None,
response_format: type[BaseModel] | None = None,
@@ -1440,7 +1423,7 @@ class ChatResponse(AFBaseModel):
if text is not None:
if isinstance(text, str):
text = TextContent(text=text)
messages.append(ChatMessage(role=ChatRole.ASSISTANT, contents=[text]))
messages.append(ChatMessage(role=Role.ASSISTANT, contents=[text]))
super().__init__(
messages=messages, # type: ignore[reportCallIssue]
@@ -1528,10 +1511,10 @@ class ChatResponseUpdate(AFBaseModel):
"""
contents: list[AIContents]
contents: list[Contents]
"""The chat response update content items."""
role: ChatRole | None = None
role: Role | None = None
"""The role of the author of the response update."""
author_name: str | None = None
"""The name of the author of the response update."""
@@ -1546,7 +1529,7 @@ class ChatResponseUpdate(AFBaseModel):
"""The model ID associated with this response update."""
created_at: CreatedAtT | None = None # use a datetimeoffset type?
"""A timestamp for the chat response update."""
finish_reason: ChatFinishReason | None = None
finish_reason: FinishReason | None = None
"""The finish reason for the operation."""
additional_properties: dict[str, Any] | None = None
@@ -1558,15 +1541,15 @@ class ChatResponseUpdate(AFBaseModel):
def __init__(
self,
*,
contents: list[AIContents],
role: ChatRole | Literal["system", "user", "assistant", "tool"] | None = None,
contents: list[Contents],
role: Role | Literal["system", "user", "assistant", "tool"] | None = None,
author_name: str | None = None,
response_id: str | None = None,
message_id: str | None = None,
conversation_id: str | None = None,
ai_model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: ChatFinishReason | None = None,
finish_reason: FinishReason | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
) -> None:
@@ -1577,14 +1560,14 @@ class ChatResponseUpdate(AFBaseModel):
self,
*,
text: TextContent | str,
role: ChatRole | Literal["system", "user", "assistant", "tool"] | None = None,
role: Role | Literal["system", "user", "assistant", "tool"] | None = None,
author_name: str | None = None,
response_id: str | None = None,
message_id: str | None = None,
conversation_id: str | None = None,
ai_model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: ChatFinishReason | None = None,
finish_reason: FinishReason | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
) -> None:
@@ -1593,16 +1576,16 @@ class ChatResponseUpdate(AFBaseModel):
def __init__(
self,
*,
contents: list[AIContents] | None = None,
contents: list[Contents] | None = None,
text: TextContent | str | None = None,
role: ChatRole | Literal["system", "user", "assistant", "tool"] | None = None,
role: Role | Literal["system", "user", "assistant", "tool"] | None = None,
author_name: str | None = None,
response_id: str | None = None,
message_id: str | None = None,
conversation_id: str | None = None,
ai_model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: ChatFinishReason | None = None,
finish_reason: FinishReason | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
) -> None:
@@ -1614,7 +1597,7 @@ class ChatResponseUpdate(AFBaseModel):
text = TextContent(text=text)
contents.append(text)
if role and isinstance(role, str):
role = ChatRole(value=role)
role = Role(value=role)
super().__init__(
contents=contents, # type: ignore[reportCallIssue]
additional_properties=additional_properties, # type: ignore[reportCallIssue]
@@ -1637,7 +1620,7 @@ class ChatResponseUpdate(AFBaseModel):
def __str__(self) -> str:
return self.text
def with_(self, contents: list[AIContent] | None = None, message_id: str | None = None) -> Self:
def with_(self, contents: list[BaseContent] | None = None, message_id: str | None = None) -> Self:
"""Returns a new instance with the specified contents and message_id."""
if contents is None:
contents = []
@@ -1709,7 +1692,7 @@ class ChatOptions(AFBaseModel):
store: bool | None = None
temperature: Annotated[float | None, Field(ge=0.0, le=2.0)] = None
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | Mapping[str, Any] | None = None
tools: list[AITool | MutableMapping[str, Any]] | None = None
tools: list[ToolProtocol | MutableMapping[str, Any]] | None = None
top_p: Annotated[float | None, Field(ge=0.0, le=1.0)] = None
user: str | None = None
@@ -1718,21 +1701,21 @@ class ChatOptions(AFBaseModel):
def _validate_tools(
cls,
tools: (
AITool
ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None
),
) -> list[AITool | MutableMapping[str, Any]] | None:
) -> list[ToolProtocol | MutableMapping[str, Any]] | None:
"""Parse the tools field."""
if not tools:
return None
if not isinstance(tools, list):
tools = [tools] # type: ignore[reportAssignmentType, assignment]
for idx, tool in enumerate(tools): # type: ignore[reportArgumentType, arg-type]
if not isinstance(tool, (AITool, MutableMapping)):
# Convert to AITool if it's a function or callable
if not isinstance(tool, (ToolProtocol, MutableMapping)):
# Convert to ToolProtocol if it's a function or callable
tools[idx] = ai_function(tool) # type: ignore[reportIndexIssues, reportCallIssue, reportArgumentType, index, call-overload, arg-type]
return tools # type: ignore[reportReturnType, return-value]
@@ -2006,8 +1989,8 @@ class AgentRunResponse(AFBaseModel):
class AgentRunResponseUpdate(AFBaseModel):
"""Represents a single streaming response chunk from an Agent."""
contents: list[AIContents] = Field(default_factory=list[AIContents])
role: ChatRole | None = None
contents: list[Contents] = Field(default_factory=list[Contents])
role: Role | None = None
author_name: str | None = None
response_id: str | None = None
message_id: str | None = None
python/packages/main/agent_framework/openai/_assistants_client.py
+13 -13
View File
@@ -20,18 +20,18 @@ from openai.types.beta.threads.run_submit_tool_outputs_params import ToolOutput
from openai.types.beta.threads.runs import RunStep
from pydantic import Field, PrivateAttr, SecretStr, ValidationError
from .._clients import ChatClientBase, use_tool_calling
from .._clients import BaseChatClient, use_tool_calling
from .._tools import AIFunction, HostedCodeInterpreterTool, HostedFileSearchTool
from .._types import (
AIContents,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
ChatToolMode,
Contents,
FunctionCallContent,
FunctionResultContent,
Role,
TextContent,
UriContent,
UsageContent,
@@ -39,7 +39,7 @@ from .._types import (
)
from ..exceptions import ServiceInitializationError
from ..telemetry import use_telemetry
from ._shared import OpenAIConfigBase, OpenAISettings
from ._shared import OpenAIConfigMixin, OpenAISettings
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
@@ -52,7 +52,7 @@ __all__ = ["OpenAIAssistantsClient"]
@use_telemetry
@use_tool_calling
class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient):
"""OpenAI Assistants client."""
assistant_id: str | None = Field(default=None)
@@ -274,13 +274,13 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
message_id=response_id,
raw_representation=response.data,
response_id=response_id,
role=ChatRole.ASSISTANT,
role=Role.ASSISTANT,
)
elif response.event == "thread.run.step.created" and isinstance(response.data, RunStep):
response_id = response.data.run_id
elif response.event == "thread.message.delta" and isinstance(response.data, MessageDeltaEvent):
delta = response.data.delta
role = ChatRole.USER if delta.role == "user" else ChatRole.ASSISTANT
role = Role.USER if delta.role == "user" else Role.ASSISTANT
for delta_block in delta.content or []:
if isinstance(delta_block, TextDeltaBlock) and delta_block.text and delta_block.text.value:
@@ -296,7 +296,7 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
contents = self._create_function_call_contents(response.data, response_id)
if contents:
yield ChatResponseUpdate(
role=ChatRole.ASSISTANT,
role=Role.ASSISTANT,
contents=contents,
conversation_id=thread_id,
message_id=response_id,
@@ -317,7 +317,7 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
)
)
yield ChatResponseUpdate(
role=ChatRole.ASSISTANT,
role=Role.ASSISTANT,
contents=[usage_content],
conversation_id=thread_id,
message_id=response_id,
@@ -331,12 +331,12 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
message_id=response_id,
raw_representation=response.data,
response_id=response_id,
role=ChatRole.ASSISTANT,
role=Role.ASSISTANT,
)
def _create_function_call_contents(self, event_data: Run, response_id: str | None) -> list[AIContents]:
def _create_function_call_contents(self, event_data: Run, response_id: str | None) -> list[Contents]:
"""Create function call contents from a tool action event."""
contents: list[AIContents] = []
contents: list[Contents] = []
if event_data.required_action is not None:
for tool_call in event_data.required_action.submit_tool_outputs.tool_calls:
@@ -437,7 +437,7 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
additional_messages = []
additional_messages.append(
AdditionalMessage(
role="assistant" if chat_message.role == ChatRole.ASSISTANT else "user",
role="assistant" if chat_message.role == Role.ASSISTANT else "user",
content=message_contents,
)
)
python/packages/main/agent_framework/openai/_chat_client.py
+24 -24
View File
@@ -15,19 +15,19 @@ from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from openai.types.chat.chat_completion_message_custom_tool_call import ChatCompletionMessageCustomToolCall
from pydantic import BaseModel, SecretStr, ValidationError
from .._clients import ChatClientBase, use_tool_calling
from .._clients import BaseChatClient, use_tool_calling
from .._logging import get_logger
from .._tools import AIFunction, AITool, HostedWebSearchTool
from .._tools import AIFunction, HostedWebSearchTool, ToolProtocol
from .._types import (
AIContents,
ChatFinishReason,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
Contents,
FinishReason,
FunctionCallContent,
FunctionResultContent,
Role,
TextContent,
UsageContent,
UsageDetails,
@@ -39,7 +39,7 @@ from ..exceptions import (
)
from ..telemetry import use_telemetry
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIConfigBase, OpenAIHandler, OpenAISettings, prepare_function_call_results
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings, prepare_function_call_results
__all__ = ["OpenAIChatClient"]
@@ -49,7 +49,7 @@ logger = get_logger("agent_framework.openai")
# region Base Client
@use_telemetry
@use_tool_calling
class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
"""OpenAI Chat completion class."""
async def _inner_get_response(
@@ -112,10 +112,10 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
# region content creation
def _chat_to_tool_spec(self, tools: list[AITool | MutableMapping[str, Any]]) -> list[dict[str, Any]]:
def _chat_to_tool_spec(self, tools: list[ToolProtocol | MutableMapping[str, Any]]) -> list[dict[str, Any]]:
chat_tools: list[dict[str, Any]] = []
for tool in tools:
if isinstance(tool, AITool):
if isinstance(tool, ToolProtocol):
match tool:
case AIFunction():
chat_tools.append(tool.to_json_schema_spec())
@@ -125,7 +125,7 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
chat_tools.append(tool if isinstance(tool, dict) else dict(tool))
return chat_tools
def _process_web_search_tool(self, tools: list[AITool | MutableMapping[str, Any]]) -> dict[str, Any] | None:
def _process_web_search_tool(self, tools: list[ToolProtocol | MutableMapping[str, Any]]) -> dict[str, Any] | None:
for tool in tools:
if isinstance(tool, HostedWebSearchTool):
# Web search tool requires special handling
@@ -173,12 +173,12 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
"""Create a chat message content object from a choice."""
response_metadata = self._get_metadata_from_chat_response(response)
messages: list[ChatMessage] = []
finish_reason: ChatFinishReason | None = None
finish_reason: FinishReason | None = None
for choice in response.choices:
response_metadata.update(self._get_metadata_from_chat_choice(choice))
if choice.finish_reason:
finish_reason = ChatFinishReason(value=choice.finish_reason)
contents: list[AIContents] = []
finish_reason = FinishReason(value=choice.finish_reason)
contents: list[Contents] = []
if parsed_tool_calls := [tool for tool in self._get_tool_calls_from_chat_choice(choice)]:
contents.extend(parsed_tool_calls)
if text_content := self._parse_text_from_choice(choice):
@@ -203,27 +203,27 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
chunk_metadata = self._get_metadata_from_streaming_chat_response(chunk)
if chunk.usage:
return ChatResponseUpdate(
role=ChatRole.ASSISTANT,
role=Role.ASSISTANT,
contents=[UsageContent(details=self._usage_details_from_openai(chunk.usage), raw_representation=chunk)],
ai_model_id=chunk.model,
additional_properties=chunk_metadata,
response_id=chunk.id,
message_id=chunk.id,
)
contents: list[AIContents] = []
finish_reason: ChatFinishReason | None = None
contents: list[Contents] = []
finish_reason: FinishReason | None = None
for choice in chunk.choices:
chunk_metadata.update(self._get_metadata_from_chat_choice(choice))
contents.extend(self._get_tool_calls_from_chat_choice(choice))
if choice.finish_reason:
finish_reason = ChatFinishReason(value=choice.finish_reason)
finish_reason = FinishReason(value=choice.finish_reason)
if text_content := self._parse_text_from_choice(choice):
contents.append(text_content)
return ChatResponseUpdate(
created_at=datetime.fromtimestamp(chunk.created).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
contents=contents,
role=ChatRole.ASSISTANT,
role=Role.ASSISTANT,
ai_model_id=chunk.model,
additional_properties=chunk_metadata,
finish_reason=finish_reason,
@@ -266,9 +266,9 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
"logprobs": getattr(choice, "logprobs", None),
}
def _get_tool_calls_from_chat_choice(self, choice: Choice | ChunkChoice) -> list[AIContents]:
def _get_tool_calls_from_chat_choice(self, choice: Choice | ChunkChoice) -> list[Contents]:
"""Get tool calls from a chat choice."""
resp: list[AIContents] = []
resp: list[Contents] = []
content = choice.message if isinstance(choice, Choice) else choice.delta
if content and content.tool_calls:
for tool in content.tool_calls:
@@ -295,7 +295,7 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
Allowing customization of the key names for role/author, and optionally overriding the role.
ChatRole.TOOL messages need to be formatted different than system/user/assistant messages:
Role.TOOL messages need to be formatted different than system/user/assistant messages:
They require a "tool_call_id" and (function) "name" key, and the "metadata" key should
be removed. The "encoding" key should also be removed.
@@ -320,7 +320,7 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
all_messages: list[dict[str, Any]] = []
for content in message.contents:
args: dict[str, Any] = {
"role": message.role.value if isinstance(message.role, ChatRole) else message.role,
"role": message.role.value if isinstance(message.role, Role) else message.role,
}
if message.additional_properties:
args["metadata"] = message.additional_properties
@@ -344,7 +344,7 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
all_messages.append(args)
return all_messages
def _openai_content_parser(self, content: AIContents) -> dict[str, Any]:
def _openai_content_parser(self, content: Contents) -> dict[str, Any]:
"""Parse contents into the openai format."""
match content:
case FunctionCallContent():
@@ -376,7 +376,7 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
TOpenAIChatClient = TypeVar("TOpenAIChatClient", bound="OpenAIChatClient")
class OpenAIChatClient(OpenAIConfigBase, OpenAIChatClientBase):
class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient):
"""OpenAI Chat completion class."""
def __init__(
python/packages/main/agent_framework/openai/_responses_client.py
+24 -24
View File
@@ -31,22 +31,22 @@ from openai.types.responses.web_search_tool_param import UserLocation as WebSear
from openai.types.responses.web_search_tool_param import WebSearchToolParam
from pydantic import BaseModel, SecretStr, ValidationError
from .._clients import ChatClientBase, use_tool_calling
from .._clients import BaseChatClient, use_tool_calling
from .._logging import get_logger
from .._tools import AIFunction, AITool, HostedCodeInterpreterTool, HostedFileSearchTool, HostedWebSearchTool
from .._tools import AIFunction, HostedCodeInterpreterTool, HostedFileSearchTool, HostedWebSearchTool, ToolProtocol
from .._types import (
AIContents,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
CitationAnnotation,
Contents,
DataContent,
FunctionCallContent,
FunctionResultContent,
HostedFileContent,
HostedVectorStoreContent,
Role,
TextContent,
TextReasoningContent,
TextSpanRegion,
@@ -61,7 +61,7 @@ from ..exceptions import (
)
from ..telemetry import use_telemetry
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIConfigBase, OpenAIHandler, OpenAISettings, prepare_function_call_results
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings, prepare_function_call_results
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
@@ -81,7 +81,7 @@ __all__ = ["OpenAIResponsesClient"]
# region ResponsesClient
class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
"""Base class for all OpenAI Responses based API's."""
FILE_SEARCH_MAX_RESULTS: int = 50
@@ -110,10 +110,10 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
store: bool | None = None,
temperature: float | None = None,
tool_choice: "ChatToolMode" | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tools: AITool
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
@@ -200,10 +200,10 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
store: bool | None = None,
temperature: float | None = None,
tool_choice: "ChatToolMode" | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tools: AITool
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[AITool | Callable[..., Any] | MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
@@ -365,11 +365,11 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
# region Prep methods
def _chat_to_response_tool_spec(
self, tools: list[AITool | MutableMapping[str, Any]]
self, tools: list[ToolProtocol | MutableMapping[str, Any]]
) -> list[ToolParam | dict[str, Any]]:
response_tools: list[ToolParam | dict[str, Any]] = []
for tool in tools:
if isinstance(tool, AITool):
if isinstance(tool, ToolProtocol):
match tool:
case HostedCodeInterpreterTool():
tool_args: dict[str, Any] = {"type": "auto"}
@@ -471,7 +471,7 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
Allowing customization of the key names for role/author, and optionally overriding the role.
ChatRole.TOOL messages need to be formatted different than system/user/assistant messages:
Role.TOOL messages need to be formatted different than system/user/assistant messages:
They require a "tool_call_id" and (function) "name" key, and the "metadata" key should
be removed. The "encoding" key should also be removed.
@@ -507,7 +507,7 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
structured_response: BaseModel | None = response.output_parsed if isinstance(response, ParsedResponse) else None # type: ignore[reportUnknownMemberType]
metadata: dict[str, Any] = response.metadata or {}
contents: list[AIContents] = []
contents: list[Contents] = []
for item in response.output: # type: ignore[reportUnknownMemberType]
match item.type:
# types:
@@ -517,12 +517,12 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
# ResponseFunctionWebSearch |
# ResponseComputerToolCall |
# ResponseReasoningItem |
# McpCall |
# McpApprovalRequest |
# MCPCall |
# MCPApprovalRequest |
# ImageGenerationCall |
# LocalShellCall |
# LocalShellCallAction |
# McpListTools |
# MCPListTools |
# ResponseCodeInterpreterToolCall |
# ResponseCustomToolCall |
# ParsedResponseOutputMessage[BaseModel] |
@@ -677,7 +677,7 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
) -> ChatResponseUpdate:
"""Create a streaming chat message content object from a choice."""
metadata: dict[str, Any] = {}
items: list[AIContents] = []
items: list[Contents] = []
conversation_id: str | None = None
model = self.ai_model_id
# TODO(peterychang): Add support for other content types
@@ -720,7 +720,7 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
return ChatResponseUpdate(
contents=items,
conversation_id=conversation_id,
role=ChatRole.ASSISTANT,
role=Role.ASSISTANT,
ai_model_id=model,
additional_properties=metadata,
raw_representation=event,
@@ -746,7 +746,7 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
"""Parse a chat message into the openai format."""
all_messages: list[dict[str, Any]] = []
args: dict[str, Any] = {
"role": message.role.value if isinstance(message.role, ChatRole) else message.role,
"role": message.role.value if isinstance(message.role, Role) else message.role,
}
if message.additional_properties:
args["metadata"] = message.additional_properties
@@ -769,8 +769,8 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
def _openai_content_parser(
self,
role: ChatRole,
content: AIContents,
role: Role,
content: Contents,
call_id_to_id: dict[str, str],
) -> dict[str, Any]:
"""Parse contents into the openai format."""
@@ -794,7 +794,7 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
return args
case TextContent():
return {
"type": "output_text" if role == ChatRole.ASSISTANT else "input_text",
"type": "output_text" if role == Role.ASSISTANT else "input_text",
"text": content.text,
}
# TODO(peterychang): We'll probably need to specialize the other content types as well
@@ -815,7 +815,7 @@ TOpenAIResponsesClient = TypeVar("TOpenAIResponsesClient", bound="OpenAIResponse
@use_telemetry
@use_tool_calling
class OpenAIResponsesClient(OpenAIConfigBase, OpenAIResponsesClientBase):
class OpenAIResponsesClient(OpenAIConfigMixin, OpenAIBaseResponsesClient):
"""OpenAI Responses client class."""
def __init__(
python/packages/main/agent_framework/openai/_shared.py
+4 -4
View File
@@ -22,7 +22,7 @@ from pydantic.types import StringConstraints
from .._logging import get_logger
from .._pydantic import AFBaseModel, AFBaseSettings
from .._types import AIContents, ChatOptions, SpeechToTextOptions, TextToSpeechOptions
from .._types import ChatOptions, Contents, SpeechToTextOptions, TextToSpeechOptions
from ..exceptions import ServiceInitializationError
from ..telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
@@ -50,7 +50,7 @@ __all__ = [
]
def prepare_function_call_results(content: AIContents | Any | list[AIContents | Any]) -> str | list[str]:
def prepare_function_call_results(content: Contents | Any | list[Contents | Any]) -> str | list[str]:
"""Prepare the values of the function call results."""
if isinstance(content, list):
results: list[str] = []
@@ -117,14 +117,14 @@ class OpenAISettings(AFBaseSettings):
realtime_model_id: str | None = None
class OpenAIHandler(AFBaseModel):
class OpenAIBase(AFBaseModel):
"""Base class for OpenAI Clients."""
client: AsyncOpenAI
ai_model_id: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
class OpenAIConfigBase(OpenAIHandler):
class OpenAIConfigMixin(OpenAIBase):
"""Internal class for configuring a connection to an OpenAI service."""
MODEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
python/packages/main/agent_framework/telemetry.py
+18 -18
View File
@@ -18,8 +18,8 @@ from ._pydantic import AFBaseSettings
if TYPE_CHECKING: # pragma: no cover
from opentelemetry.util._decorator import _AgnosticContextManager # type: ignore[reportPrivateUsage]
from ._agents import AIAgent, ChatClientAgent
from ._clients import ChatClientBase
from ._agents import AgentProtocol, ChatAgent
from ._clients import BaseChatClient
from ._threads import AgentThread
from ._tools import AIFunction
from ._types import (
@@ -31,8 +31,8 @@ if TYPE_CHECKING: # pragma: no cover
ChatResponseUpdate,
)
TChatClientBase = TypeVar("TChatClientBase", bound="ChatClientBase")
TChatClientAgent = TypeVar("TChatClientAgent", bound="ChatClientAgent")
TBaseChatClient = TypeVar("TBaseChatClient", bound="BaseChatClient")
TChatClientAgent = TypeVar("TChatClientAgent", bound="ChatAgent")
tracer = get_tracer("agent_framework")
logger = get_logger()
@@ -269,7 +269,7 @@ def _set_error(span: Span, error: Exception) -> None:
span.set_status(StatusCode.ERROR, repr(error))
# region ChatClient
# region ChatClientProtocol
def _trace_chat_get_response(
@@ -283,7 +283,7 @@ def _trace_chat_get_response(
@functools.wraps(completion_func)
async def wrap_inner_get_response(
self: "ChatClientBase",
self: "BaseChatClient",
*,
messages: MutableSequence["ChatMessage"],
chat_options: "ChatOptions",
@@ -334,7 +334,7 @@ def _trace_chat_get_streaming_response(
@functools.wraps(completion_func)
async def wrap_inner_get_streaming_response(
self: "ChatClientBase", *, messages: MutableSequence["ChatMessage"], chat_options: "ChatOptions", **kwargs: Any
self: "BaseChatClient", *, messages: MutableSequence["ChatMessage"], chat_options: "ChatOptions", **kwargs: Any
) -> AsyncIterable["ChatResponseUpdate"]:
if not MODEL_DIAGNOSTICS_SETTINGS.ENABLED:
# If model diagnostics are not enabled, just return the completion
@@ -375,11 +375,11 @@ def _trace_chat_get_streaming_response(
return wrap_inner_get_streaming_response
def use_telemetry(cls: type[TChatClientBase]) -> type[TChatClientBase]:
def use_telemetry(cls: type[TBaseChatClient]) -> type[TBaseChatClient]:
"""Class decorator that enables telemetry for a chat client.
Remarks:
This only works on classes that derive from ChatClientBase
This only works on classes that derive from BaseChatClient
and the _inner_get_response
and _inner_get_streaming_response methods.
It also relies on the presence of the MODEL_PROVIDER_NAME class variable.
@@ -520,7 +520,7 @@ def _trace_agent_run(
@functools.wraps(run_func)
async def wrap_run(
self: "ChatClientAgent",
self: "ChatAgent",
messages: "str | ChatMessage | list[str] | list[ChatMessage] | None" = None,
*,
thread: "AgentThread | None" = None,
@@ -560,7 +560,7 @@ def _trace_agent_run(
return wrap_run
def _trace_agent_run_streaming(
def _trace_agent_run_stream(
run_func: Callable[..., AsyncIterable["AgentRunResponseUpdate"]],
) -> Callable[..., AsyncIterable["AgentRunResponseUpdate"]]:
"""Decorator to trace streaming agent run activities.
@@ -570,8 +570,8 @@ def _trace_agent_run_streaming(
"""
@functools.wraps(run_func)
async def wrap_run_streaming(
self: "ChatClientAgent",
async def wrap_run_stream(
self: "ChatAgent",
messages: "str | ChatMessage | list[str] | list[ChatMessage] | None" = None,
*,
thread: "AgentThread | None" = None,
@@ -610,23 +610,23 @@ def _trace_agent_run_streaming(
raise
# Mark the wrapper decorator as a streaming agent run decorator
wrap_run_streaming.__model_diagnostics_streaming_agent_run__ = True # type: ignore
return wrap_run_streaming
wrap_run_stream.__model_diagnostics_streaming_agent_run__ = True # type: ignore
return wrap_run_stream
def use_agent_telemetry(cls: type[TChatClientAgent]) -> type[TChatClientAgent]:
"""Class decorator that enables telemetry for an agent."""
if run := getattr(cls, "run", None):
cls.run = _trace_agent_run(run) # type: ignore
if run_streaming := getattr(cls, "run_streaming", None):
cls.run_streaming = _trace_agent_run_streaming(run_streaming) # type: ignore
if run_stream := getattr(cls, "run_stream", None):
cls.run_stream = _trace_agent_run_stream(run_stream) # type: ignore
return cls
def _get_agent_run_span(
*,
operation_name: str,
agent: "AIAgent",
agent: "AgentProtocol",
system: str,
thread: "AgentThread | None",
**kwargs: Any,