diff --git a/python/docs/generate_docs.py b/python/docs/generate_docs.py index 15325ee7c9..e3c3d9c45d 100644 --- a/python/docs/generate_docs.py +++ b/python/docs/generate_docs.py @@ -54,7 +54,6 @@ async def generate_af_docs(root_path: Path): "undoc-members": True, "show-inheritance": True, "imported-members": True, - "inherited-members": 'AFBaseModel', }, }, } diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 597e7b2da4..066afee08b 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -5,7 +5,7 @@ import json import re import uuid from collections.abc import AsyncIterable, Sequence -from typing import Any +from typing import Any, cast import httpx from a2a.client import Client, ClientConfig, ClientFactory, minimal_agent_card @@ -307,7 +307,7 @@ class A2AAgent(BaseAgent): role=A2ARole("user"), parts=parts, message_id=message.message_id or uuid.uuid4().hex, - metadata=message.additional_properties, + metadata=cast(dict[str, Any], message.additional_properties), ) def _a2a_parts_to_contents(self, parts: Sequence[A2APart]) -> list[Contents]: diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py index b64ea62d3e..28a81e715b 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py @@ -84,7 +84,7 @@ from azure.ai.projects.aio import AIProjectClient from azure.ai.projects.models import ConnectionType from azure.core.credentials_async import AsyncTokenCredential from azure.core.exceptions import HttpResponseError, ResourceNotFoundError -from pydantic import BaseModel, Field, PrivateAttr, ValidationError +from pydantic import BaseModel, ValidationError if sys.version_info >= (3, 11): from typing import Self # pragma: no cover @@ -128,14 +128,6 @@ class AzureAIAgentClient(BaseChatClient): """Azure AI Agent Chat client.""" OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai" # type: ignore[reportIncompatibleVariableOverride, misc] - project_client: AIProjectClient = Field(...) - credential: AsyncTokenCredential | None = Field(...) - agent_id: str | None = Field(default=None) - agent_name: str | None = Field(default=None) - ai_model_id: str | None = Field(default=None) - thread_id: str | None = Field(default=None) - _should_delete_agent: bool = PrivateAttr(default=False) # Track whether we should delete the agent - _should_close_client: bool = PrivateAttr(default=False) # Track whether we should close client connection def __init__( self, @@ -205,16 +197,18 @@ class AzureAIAgentClient(BaseChatClient): ) should_close_client = True - super().__init__( - project_client=project_client, # type: ignore[reportCallIssue] - credential=async_credential, # type: ignore[reportCallIssue] - agent_id=agent_id, # type: ignore[reportCallIssue] - thread_id=thread_id, # type: ignore[reportCallIssue] - agent_name=agent_name, # type: ignore[reportCallIssue] - ai_model_id=azure_ai_settings.model_deployment_name, # type: ignore[reportCallIssue] - **kwargs, - ) - self._should_close_client = should_close_client + # Initialize parent + super().__init__(**kwargs) + + # Initialize instance variables + self.project_client = project_client + self.credential = async_credential + self.agent_id = agent_id + self.agent_name = agent_name + self.model_id = azure_ai_settings.model_deployment_name + self.thread_id = thread_id + self._should_delete_agent = False # Track whether we should delete the agent + self._should_close_client = should_close_client # Track whether we should close client connection async def setup_azure_ai_observability(self, enable_sensitive_data: bool | None = None) -> None: """Use this method to setup tracing in your Azure AI Project. @@ -251,7 +245,7 @@ class AzureAIAgentClient(BaseChatClient): await self._close_client_if_needed() @classmethod - def from_dict(cls: type[TAzureAIAgentClient], settings: dict[str, Any]) -> TAzureAIAgentClient: + def from_settings(cls: type[TAzureAIAgentClient], settings: dict[str, Any]) -> TAzureAIAgentClient: """Initialize a AzureAIAgentClient from a dictionary of settings. Args: @@ -317,11 +311,11 @@ class AzureAIAgentClient(BaseChatClient): """ # If no agent_id is provided, create a temporary agent if self.agent_id is None: - if not self.ai_model_id: + if not self.model_id: raise ServiceInitializationError("Model deployment name is required for agent creation.") agent_name: str = self.agent_name or "UnnamedAgent" - args: dict[str, Any] = {"model": self.ai_model_id, "name": agent_name} + args: dict[str, Any] = {"model": self.model_id, "name": agent_name} if run_options: if "tools" in run_options: args["tools"] = run_options["tools"] @@ -866,8 +860,8 @@ class AzureAIAgentClient(BaseChatClient): if isinstance(content, FunctionResultContent): if tool_outputs is None: tool_outputs = [] - result_contents: list[Any] = ( # type: ignore - content.result if isinstance(content.result, list) else [content.result] # type: ignore + result_contents: list[Any] = ( + content.result if isinstance(content.result, list) else [content.result] ) results: list[Any] = [] for item in result_contents: diff --git a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py index 8c23e209bd..05028b81bb 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py +++ b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py @@ -50,17 +50,26 @@ def create_test_azure_ai_chat_client( azure_ai_settings: AzureAISettings | None = None, should_delete_agent: bool = False, ) -> AzureAIAgentClient: - """Helper function to create AzureAIAgentClient instances for testing, bypassing Pydantic validation.""" + """Helper function to create AzureAIAgentClient instances for testing, bypassing normal validation.""" if azure_ai_settings is None: azure_ai_settings = AzureAISettings(env_file_path="test.env") - return AzureAIAgentClient.model_construct( - project_client=mock_ai_project_client, - agent_id=agent_id, - thread_id=thread_id, - _should_delete_agent=should_delete_agent, - ai_model_id=azure_ai_settings.model_deployment_name, - ) + # Create client instance directly + client = object.__new__(AzureAIAgentClient) + + # Set attributes directly + client.project_client = mock_ai_project_client + client.credential = None + client.agent_id = agent_id + client.agent_name = None + client.model_id = azure_ai_settings.model_deployment_name + client.thread_id = thread_id + client._should_delete_agent = should_delete_agent + client._should_close_client = False + client.additional_properties = {} + client.middleware = None + + return client def test_azure_ai_settings_init(azure_ai_unit_test_env: dict[str, str]) -> None: @@ -101,14 +110,19 @@ def test_azure_ai_chat_client_init_auto_create_client( ) -> None: """Test AzureAIAgentClient initialization with auto-created project_client.""" azure_ai_settings = AzureAISettings(**azure_ai_unit_test_env) # type: ignore - chat_client = AzureAIAgentClient.model_construct( - project_client=mock_ai_project_client, - agent_id=None, - thread_id=None, - _should_delete_agent=False, - _azure_ai_settings=azure_ai_settings, - credential=None, - ) + + # Create client instance directly + chat_client = object.__new__(AzureAIAgentClient) + chat_client.project_client = mock_ai_project_client + chat_client.agent_id = None + chat_client.thread_id = None + chat_client._should_delete_agent = False + chat_client._should_close_client = False + chat_client.credential = None + chat_client.model_id = azure_ai_settings.model_deployment_name + chat_client.agent_name = None + chat_client.additional_properties = {} + chat_client.middleware = None assert chat_client.project_client is mock_ai_project_client assert chat_client.agent_id is None @@ -169,7 +183,6 @@ def test_azure_ai_chat_client_from_dict(mock_ai_project_client: MagicMock) -> No azure_ai_settings = AzureAISettings( project_endpoint=settings["project_endpoint"], model_deployment_name=settings["model_deployment_name"], - agent_name=settings["agent_name"], ) chat_client: AzureAIAgentClient = create_test_azure_ai_chat_client( @@ -229,9 +242,7 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_create_new( azure_ai_unit_test_env: dict[str, str], ) -> None: """Test _get_agent_id_or_create when creating a new agent.""" - azure_ai_settings = AzureAISettings( - model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"], agent_name="TestAgent" - ) + azure_ai_settings = AzureAISettings(model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]) chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, azure_ai_settings=azure_ai_settings) agent_id = await chat_client._get_agent_id_or_create() # type: ignore @@ -519,9 +530,7 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_with_run_options( mock_ai_project_client: MagicMock, azure_ai_unit_test_env: dict[str, str] ) -> None: """Test _get_agent_id_or_create with run_options containing tools and instructions.""" - azure_ai_settings = AzureAISettings( - model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"], agent_name="TestAgent" - ) + azure_ai_settings = AzureAISettings(model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]) chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, azure_ai_settings=azure_ai_settings) run_options = { @@ -618,6 +627,7 @@ def get_weather( return f"The weather in {location} is sunny with a high of 25°C." +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_get_response() -> None: """Test Azure AI Chat Client response.""" @@ -642,6 +652,7 @@ async def test_azure_ai_chat_client_get_response() -> None: assert any(word in response.text.lower() for word in ["sunny", "25"]) +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_get_response_tools() -> None: """Test Azure AI Chat Client response with tools.""" @@ -663,6 +674,7 @@ async def test_azure_ai_chat_client_get_response_tools() -> None: assert any(word in response.text.lower() for word in ["sunny", "25"]) +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_streaming() -> None: """Test Azure AI Chat Client streaming response.""" @@ -693,6 +705,7 @@ async def test_azure_ai_chat_client_streaming() -> None: assert any(word in full_message.lower() for word in ["sunny", "25"]) +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_streaming_tools() -> None: """Test Azure AI Chat Client streaming response with tools.""" @@ -719,6 +732,7 @@ async def test_azure_ai_chat_client_streaming_tools() -> None: assert any(word in full_message.lower() for word in ["sunny", "25"]) +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_agent_basic_run() -> None: """Test ChatAgent basic run functionality with AzureAIAgentClient.""" @@ -735,6 +749,7 @@ async def test_azure_ai_chat_client_agent_basic_run() -> None: assert "Hello World" in response.text +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_agent_basic_run_streaming() -> None: """Test ChatAgent basic streaming functionality with AzureAIAgentClient.""" @@ -754,6 +769,7 @@ async def test_azure_ai_chat_client_agent_basic_run_streaming() -> None: assert "streaming response test" in full_message.lower() +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_agent_thread_persistence() -> None: """Test ChatAgent thread persistence across runs with AzureAIAgentClient.""" @@ -779,6 +795,7 @@ async def test_azure_ai_chat_client_agent_thread_persistence() -> None: assert "42" in second_response.text +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_agent_existing_thread_id() -> None: """Test ChatAgent existing thread ID functionality with AzureAIAgentClient.""" @@ -816,6 +833,7 @@ async def test_azure_ai_chat_client_agent_existing_thread_id() -> None: assert "alice" in response2.text.lower() +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_agent_code_interpreter(): """Test ChatAgent with code interpreter through AzureAIAgentClient.""" @@ -835,6 +853,7 @@ async def test_azure_ai_chat_client_agent_code_interpreter(): assert "120" in response.text or "factorial" in response.text.lower() +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_agent_with_mcp_tools() -> None: """Test MCP tools defined at agent creation with AzureAIAgentClient.""" @@ -856,6 +875,7 @@ async def test_azure_ai_chat_client_agent_with_mcp_tools() -> None: assert any(term in response.text.lower() for term in ["app service", "azure", "web", "application"]) +@pytest.mark.flaky @skip_if_azure_ai_integration_tests_disabled async def test_azure_ai_chat_client_agent_level_tool_persistence(): """Test that agent-level tools persist across multiple runs with AzureAIAgentClient.""" diff --git a/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py b/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py index 164a31e059..4636e72a81 100644 --- a/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py +++ b/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py @@ -268,7 +268,6 @@ class CopilotStudioAgent(BaseAgent): yield AgentRunResponseUpdate( role=message.role, contents=message.contents, - additional_properties=message.additional_properties, author_name=message.author_name, raw_representation=message.raw_representation, response_id=message.message_id, diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 58a71d3552..6142c52ddd 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -16,6 +16,7 @@ from ._logging import get_logger from ._mcp import MCPTool from ._memory import AggregateContextProvider, Context, ContextProvider from ._middleware import Middleware, use_agent_middleware +from ._serialization import SerializationMixin from ._threads import AgentThread, ChatMessageStoreProtocol from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, AIFunction, ToolProtocol from ._types import ( @@ -135,26 +136,20 @@ class AgentProtocol(Protocol): # region BaseAgent -class BaseAgent: - """Base class for all Agent Framework agents. +class BaseAgent(SerializationMixin): + """Base class for all Agent Framework agents.""" - Attributes: - id: The unique identifier of the agent If no id is provided, - a new UUID will be generated. - name: The name of the agent, can be None. - description: The description of the agent. - display_name: The display name of the agent, which is either the name or id. - context_providers: The collection of multiple context providers to include during agent invocation. - middleware: List of middleware to intercept agent and function invocations. - """ + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"additional_properties"} def __init__( self, + *, id: str | None = None, name: str | None = None, description: str | None = None, context_providers: ContextProvider | Sequence[ContextProvider] | None = None, middleware: Middleware | Sequence[Middleware] | None = None, + additional_properties: MutableMapping[str, Any] | None = None, **kwargs: Any, ) -> None: """Base class for all Agent Framework agents. @@ -167,7 +162,8 @@ class BaseAgent: display_name: The display name of the agent, which is either the name or id. context_providers: The collection of multiple context providers to include during agent invocation. middleware: List of middleware to intercept agent and function invocations. - kwargs: will be stored in `additional_properties` + additional_properties: Additional properties set on the agent. + kwargs: Additional keyword arguments (merged into additional_properties). """ if id is None: id = str(uuid4()) @@ -179,7 +175,10 @@ class BaseAgent: self.middleware: list[Middleware] | None = cast(list[Middleware], middleware) if middleware else None else: self.middleware = [middleware] - self.additional_properties = kwargs + + # Merge kwargs into additional_properties + self.additional_properties: dict[str, Any] = cast(dict[str, Any], additional_properties or {}) + self.additional_properties.update(kwargs) async def _notify_thread_of_new_messages( self, @@ -213,7 +212,7 @@ class BaseAgent: async def deserialize_thread(self, serialized_thread: Any, **kwargs: Any) -> AgentThread: """Deserializes the thread.""" thread: AgentThread = self.get_new_thread() - await thread.deserialize(serialized_thread, **kwargs) + await thread.update_from_thread_state(serialized_thread, **kwargs) return thread def as_tool( @@ -700,7 +699,7 @@ class ChatAgent(BaseAgent): store=store, temperature=temperature, tool_choice=tool_choice, - tools=final_tools, # type: ignore[reportArgumentType] + tools=final_tools, top_p=top_p, user=user, additional_properties=additional_properties or {}, diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 4b2e26bb65..693fb3202f 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -3,9 +3,9 @@ import asyncio from abc import ABC, abstractmethod from collections.abc import AsyncIterable, Callable, MutableMapping, MutableSequence, Sequence -from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, runtime_checkable +from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol, TypeVar, runtime_checkable -from pydantic import BaseModel, Field +from pydantic import BaseModel from ._logging import get_logger from ._mcp import MCPTool @@ -17,7 +17,7 @@ from ._middleware import ( FunctionMiddlewareCallable, Middleware, ) -from ._pydantic import AFBaseModel +from ._serialization import SerializationMixin from ._threads import ChatMessageStoreProtocol from ._tools import ToolProtocol from ._types import ( @@ -189,21 +189,61 @@ def prepare_messages(messages: str | ChatMessage | list[str] | list[ChatMessage] return return_messages -class BaseChatClient(AFBaseModel, ABC): +class BaseChatClient(SerializationMixin, ABC): """Base class for chat clients.""" - additional_properties: dict[str, Any] = Field(default_factory=dict) - middleware: ( - ChatMiddleware - | ChatMiddlewareCallable - | FunctionMiddleware - | FunctionMiddlewareCallable - | list[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable] - | None - ) = None - OTEL_PROVIDER_NAME: str = "unknown" + OTEL_PROVIDER_NAME: ClassVar[str] = "unknown" + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"additional_properties"} # This is used for OTel setup, should be overridden in subclasses + def __init__( + self, + *, + middleware: ( + ChatMiddleware + | ChatMiddlewareCallable + | FunctionMiddleware + | FunctionMiddlewareCallable + | list[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable] + | None + ) = None, + additional_properties: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Initialize BaseChatClient. + + Args: + additional_properties: Additional properties for the client. + middleware: Middleware for the client. + **kwargs: Additional keyword arguments (merged into additional_properties). + """ + # Merge kwargs into additional_properties + self.additional_properties = additional_properties or {} + self.additional_properties.update(kwargs) + + self.middleware = middleware + + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + """Convert the instance to a dictionary. + + Extracts additional_properties fields to the root level. + + Args: + exclude: Set of field names to exclude from serialization. + exclude_none: Whether to exclude None values from the output. Defaults to True. + + Returns: + Dictionary representation of the instance. + """ + # Get the base dict from SerializationMixin + result = super().to_dict(exclude=exclude, exclude_none=exclude_none) + + # Extract additional_properties to root level + if self.additional_properties: + result.update(self.additional_properties) + + return result + def prepare_messages( self, messages: str | ChatMessage | list[str] | list[ChatMessage], chat_options: ChatOptions ) -> MutableSequence[ChatMessage]: diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index f487c28099..af48f137d3 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -60,7 +60,7 @@ def _mcp_prompt_message_to_chat_message( """Convert a MCP container type to a Agent Framework type.""" return ChatMessage( role=Role(value=mcp_type.role), - contents=[_mcp_type_to_ai_content(mcp_type.content)], # type: ignore[call-arg] + contents=[_mcp_type_to_ai_content(mcp_type.content)], raw_representation=mcp_type, ) diff --git a/python/packages/core/agent_framework/_pydantic.py b/python/packages/core/agent_framework/_pydantic.py index 077da03650..8aac34e02f 100644 --- a/python/packages/core/agent_framework/_pydantic.py +++ b/python/packages/core/agent_framework/_pydantic.py @@ -3,19 +3,13 @@ from typing import Annotated, Any, ClassVar, TypeVar -from pydantic import BaseModel, ConfigDict, Field, UrlConstraints +from pydantic import Field, UrlConstraints from pydantic.networks import AnyUrl from pydantic_settings import BaseSettings, SettingsConfigDict HTTPsUrl = Annotated[AnyUrl, UrlConstraints(max_length=2083, allowed_schemes=["https"])] -__all__ = ["AFBaseModel", "AFBaseSettings", "HTTPsUrl"] - - -class AFBaseModel(BaseModel): - """Base class for all pydantic models in the Agent Framework.""" - - model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True, validate_assignment=True) +__all__ = ["AFBaseSettings", "HTTPsUrl"] TSettings = TypeVar("TSettings", bound="AFBaseSettings") diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py new file mode 100644 index 0000000000..e2078efd75 --- /dev/null +++ b/python/packages/core/agent_framework/_serialization.py @@ -0,0 +1,275 @@ +# Copyright (c) Microsoft. All rights reserved. + +import json +import re +from collections.abc import MutableMapping +from typing import Any, ClassVar, Protocol, TypeVar, runtime_checkable + +from ._logging import get_logger + +logger = get_logger() + +TClass = TypeVar("TClass", bound="SerializationMixin") +TProtocol = TypeVar("TProtocol", bound="SerializationProtocol") + +# Regex pattern for converting CamelCase to snake_case +_CAMEL_TO_SNAKE_PATTERN = re.compile(r"(? dict[str, Any]: + """Convert the instance to a dictionary.""" + ... + + @classmethod + def from_dict(cls: type[TProtocol], value: MutableMapping[str, Any], /, **kwargs: Any) -> TProtocol: + """Create an instance from a dictionary.""" + ... + + +def is_serializable(value: Any) -> bool: + """Check if a value is JSON serializable.""" + return isinstance(value, (str, int, float, bool, type(None), list, dict)) + + +class SerializationMixin: + """Mixin class providing serialization and deserialization capabilities. + + Classes using this mixin should handle MutableMapping inputs in their __init__ method + for any parameters that expect SerializationMixin/SerializationProtocol instances. + The __init__ should check if the value is a MutableMapping and call from_dict() to convert it. + + So take the two classes below as an example. The first purely uses base types, strings in this case. + The second has a param that is of the type of the first class. + Because we setup the __init__ method to handle MutableMapping, + we can pass in a dict to the second class and it will convert it to an instance of the first class. + + Example: + + .. code-block:: python + + class SerializableMixinType(SerializationMixin): + def __init__(self, param1: str, param2: int) -> None: + self.param1 = param1 + self.param2 = param2 + + + class MyClass(SerializationMixin): + def __init__( + self, + regular_param: str, + param: SerializableMixinType | MutableMapping[str, Any] | None = None, + ) -> None: + if isinstance(param, MutableMapping): + self.param = self.from_dict(param) + else: + self.param = param + self.regular_param = regular_param + + + instance = MyClass.from_dict({"regular_param": "value", "param": {"param1": "value1", "param2": 42}}) + + A more complex use case involves a injectable dependency that is not serialized. + In this case, the dependency is passed in via the dependencies parameter to from_dict/from_json. + + Example: + + .. code-block:: python + + from libary import Client + + + class MyClass(SerializationMixin): + INJECTABLE = {"client"} + + def __init__(self, regular_param: str, client: Client) -> None: + self.client = client + self.regular_param = regular_param + + + json_of_class = MyClass(regular_param="value", client=Client()).to_json() + # this looks like: {"type": "my_class", "regular_param": "value"} + + instance = MyClass.from_dict(json_of_class, dependencies={"my_class.client": Client()}) + + During serialization, the field listed as INJECTABLE (and also DEFAULT_EXCLUDE) will be excluded from the output. + Then in deserialization, + the dependencies dict is checked for any keys matching the formats: + - "." + - ".." + where is the type identifier for the class (either the value of the 'type' class variable or + the snake_cased class name if 'type' is not present), + is the name of the parameter in the __init__ method, + is the name of a parameter that is a dict, + and is a key in that dict parameter. + """ + + DEFAULT_EXCLUDE: ClassVar[set[str]] = set() + INJECTABLE: ClassVar[set[str]] = set() + + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + """Convert the instance and any nested objects to a dictionary. + + Args: + exclude: Set of field names to exclude from serialization. + exclude_none: Whether to exclude None values from the output. Defaults to True. + + Returns: + Dictionary representation of the instance. + """ + # Combine exclude sets + combined_exclude = set(self.DEFAULT_EXCLUDE) + if exclude: + combined_exclude.update(exclude) + combined_exclude.update(self.INJECTABLE) + + # Get all instance attributes + result: dict[str, Any] = {"type": self._get_type_identifier()} + for key, value in self.__dict__.items(): + if key not in combined_exclude and not key.startswith("_"): + if exclude_none and value is None: + continue + # Recursively serialize SerializationProtocol objects + if isinstance(value, SerializationProtocol): + result[key] = value.to_dict(exclude=exclude, exclude_none=exclude_none) + continue + # Handle lists containing SerializationProtocol objects + if isinstance(value, list): + value_as_list: list[Any] = [] + for item in value: + if isinstance(item, SerializationProtocol): + value_as_list.append(item.to_dict(exclude=exclude, exclude_none=exclude_none)) + continue + if is_serializable(item): + value_as_list.append(item) + continue + logger.debug( + f"Skipping non-serializable item in list attribute '{key}' of type {type(item).__name__}" + ) + result[key] = value_as_list + continue + # Handle dicts containing SerializationProtocol values + if isinstance(value, dict): + serialized_dict: dict[str, Any] = {} + for k, v in value.items(): + if isinstance(v, SerializationProtocol): + serialized_dict[k] = v.to_dict(exclude=exclude, exclude_none=exclude_none) + continue + # Check if the value is JSON serializable + if is_serializable(v): + serialized_dict[k] = v + continue + logger.debug( + f"Skipping non-serializable value for key '{k}' in dict attribute '{key}' " + f"of type {type(v).__name__}" + ) + result[key] = serialized_dict + continue + # Directly include JSON serializable values + if is_serializable(value): + result[key] = value + continue + logger.debug(f"Skipping non-serializable attribute '{key}' of type {type(value).__name__}") + + return result + + def to_json(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> str: + """Convert the instance to a JSON string. + + Args: + exclude: Set of field names to exclude from serialization. + exclude_none: Whether to exclude None values from the output. Defaults to True. + + Returns: + JSON string representation of the instance. + """ + return json.dumps(self.to_dict(exclude=exclude, exclude_none=exclude_none)) + + @classmethod + def from_dict( + cls: type[TClass], value: MutableMapping[str, Any], /, dependencies: MutableMapping[str, Any] | None = None + ) -> TClass: + """Create an instance from a dictionary. + + Args: + value: Dictionary containing the instance data (positional-only). + dependencies: Dictionary mapping dependency keys to values. + Keys should be in format "." or "..". + + Returns: + New instance of the class. + """ + if dependencies is None: + dependencies = {} + + # Get the type identifier + type_id = cls._get_type_identifier() + + # Create a copy of the value dict to work with, filtering out the 'type' key + kwargs = {k: v for k, v in value.items() if k != "type"} + + # Process dependencies + for dep_key, dep_value in dependencies.items(): + parts = dep_key.split(".") + if len(parts) < 2: + continue + + dep_type = parts[0] + if dep_type != type_id: + continue + + param_name = parts[1] + + # Log debug message if dependency is not in INJECTABLE + if param_name not in cls.INJECTABLE: + logger.debug( + f"Dependency '{param_name}' for type '{type_id}' is not in INJECTABLE set. " + f"Available injectable parameters: {cls.INJECTABLE}" + ) + + if len(parts) == 2: + # Simple parameter: . + kwargs[param_name] = dep_value + elif len(parts) == 3: + # Dict parameter: .. + dict_param_name = parts[1] + key = parts[2] + if dict_param_name not in kwargs: + kwargs[dict_param_name] = {} + kwargs[dict_param_name][key] = dep_value + + return cls(**kwargs) + + @classmethod + def from_json(cls: type[TClass], value: str, /, dependencies: MutableMapping[str, Any] | None = None) -> TClass: + """Create an instance from a JSON string. + + Args: + value: JSON string containing the instance data (positional-only). + dependencies: Dictionary mapping dependency keys to values. + Keys should be in format "." or "..". + + Returns: + New instance of the class. + """ + data = json.loads(value) + return cls.from_dict(data, dependencies=dependencies) + + @classmethod + def _get_type_identifier(cls) -> str: + """Get the type identifier for this class. + + Returns the value of the 'type' class variable if present, + otherwise returns a snake_cased version of the class name. + + Returns: + Type identifier string. + """ + if (type_ := getattr(cls, "type", None)) and isinstance(type_, str): + return type_ # type:ignore[no-any-return] + + # Convert class name to snake_case + return _CAMEL_TO_SNAKE_PATTERN.sub("_", cls.__name__).lower() diff --git a/python/packages/core/agent_framework/_threads.py b/python/packages/core/agent_framework/_threads.py index c334d76538..378efd5699 100644 --- a/python/packages/core/agent_framework/_threads.py +++ b/python/packages/core/agent_framework/_threads.py @@ -3,10 +3,9 @@ from collections.abc import Sequence from typing import Any, Protocol, TypeVar -from pydantic import model_validator +from pydantic import BaseModel, ConfigDict, model_validator from ._memory import AggregateContextProvider -from ._pydantic import AFBaseModel from ._types import ChatMessage from .exceptions import AgentThreadException @@ -64,7 +63,19 @@ class ChatMessageStoreProtocol(Protocol): ... -class AgentThreadState(AFBaseModel): +class ChatMessageStoreState(BaseModel): + """State model for serializing and deserializing chat message store data. + + Attributes: + messages: List of chat messages stored in the message store. + """ + + messages: list[ChatMessage] + + model_config = ConfigDict(arbitrary_types_allowed=True) + + +class AgentThreadState(BaseModel): """State model for serializing and deserializing thread information. Attributes: @@ -73,7 +84,9 @@ class AgentThreadState(AFBaseModel): """ service_thread_id: str | None = None - chat_message_store_state: Any | None = None + chat_message_store_state: ChatMessageStoreState | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True) @model_validator(mode="before") def validate_only_one(cls, values: dict[str, Any]) -> dict[str, Any]: @@ -86,16 +99,6 @@ class AgentThreadState(AFBaseModel): return values -class ChatMessageStoreState(AFBaseModel): - """State model for serializing and deserializing chat message store data. - - Attributes: - messages: List of chat messages stored in the message store. - """ - - messages: list[ChatMessage] - - TChatMessageStore = TypeVar("TChatMessageStore", bound="ChatMessageStore") diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index e7860470f1..46110215b3 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -11,6 +11,7 @@ from typing import ( TYPE_CHECKING, Annotated, Any, + ClassVar, Final, Generic, Literal, @@ -22,10 +23,10 @@ from typing import ( ) from opentelemetry.metrics import Histogram -from pydantic import AnyUrl, BaseModel, Field, PrivateAttr, ValidationError, create_model, field_validator +from pydantic import AnyUrl, BaseModel, Field, ValidationError, create_model from ._logging import get_logger -from ._pydantic import AFBaseModel +from ._serialization import SerializationMixin from .exceptions import ChatClientInitializationError, ToolException from .observability import ( OPERATION_DURATION_BUCKET_BOUNDARIES, @@ -150,7 +151,7 @@ class ToolProtocol(Protocol): ... -class BaseTool(AFBaseModel): +class BaseTool(SerializationMixin): """Base class for AI tools, providing common attributes and methods. Args: @@ -159,9 +160,29 @@ class BaseTool(AFBaseModel): additional_properties: Additional properties associated with the tool. """ - name: str = Field(..., kw_only=False) - description: str = "" - additional_properties: dict[str, Any] | None = None + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"additional_properties"} + + def __init__( + self, + *, + name: str, + description: str = "", + additional_properties: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Initialize the BaseTool. + + Args: + name: The name of the tool. + description: A description of the tool. + additional_properties: Additional properties associated with the tool. + **kwargs: Additional keyword arguments. + """ + self.name = name + self.description = description + self.additional_properties = additional_properties + for key, value in kwargs.items(): + setattr(self, key, value) def __str__(self) -> str: """Return a string representation of the tool.""" @@ -177,8 +198,6 @@ class HostedCodeInterpreterTool(BaseTool): that it is allowed to execute generated code if the service is capable of doing so. """ - inputs: list[Any] = Field(default_factory=list) - def __init__( self, *, @@ -202,18 +221,17 @@ class HostedCodeInterpreterTool(BaseTool): additional_properties: Additional properties associated with the tool. **kwargs: Additional keyword arguments to pass to the base class. """ - args: dict[str, Any] = { - "name": "code_interpreter", - } - if inputs: - args["inputs"] = _parse_inputs(inputs) - if description is not None: - args["description"] = description - if additional_properties is not None: - args["additional_properties"] = additional_properties if "name" in kwargs: raise ValueError("The 'name' argument is reserved for the HostedCodeInterpreterTool and cannot be set.") - super().__init__(**args, **kwargs) + + self.inputs = _parse_inputs(inputs) if inputs else [] + + super().__init__( + name="code_interpreter", + description=description or "", + additional_properties=additional_properties, + **kwargs, + ) class HostedWebSearchTool(BaseTool): @@ -263,11 +281,6 @@ class HostedMCPSpecificApproval(TypedDict, total=False): class HostedMCPTool(BaseTool): """Represents a MCP tool that is managed and executed by the service.""" - url: AnyUrl - approval_mode: Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None = None - allowed_tools: set[str] | None = None - headers: dict[str, str] | None = None - def __init__( self, *, @@ -296,43 +309,45 @@ class HostedMCPTool(BaseTool): additional_properties: Additional properties to include in the tool definition. **kwargs: Additional keyword arguments to pass to the base class. """ - args: dict[str, Any] = { - "name": name, - "url": url, - } - if allowed_tools is not None: - args["allowed_tools"] = allowed_tools - if approval_mode is not None: - args["approval_mode"] = approval_mode - if headers is not None: - args["headers"] = headers - if description is not None: - args["description"] = description - if additional_properties is not None: - args["additional_properties"] = additional_properties try: - super().__init__(**args, **kwargs) - except ValidationError as err: - raise ToolException(f"Error initializing HostedMCPTool: {err}", inner_exception=err) from err + # Validate approval_mode + if approval_mode is not None: + if isinstance(approval_mode, str): + if approval_mode not in ("always_require", "never_require"): + raise ValueError( + f"Invalid approval_mode: {approval_mode}. " + "Must be 'always_require', 'never_require', or a dict with 'always_require_approval' " + "or 'never_require_approval' keys." + ) + elif isinstance(approval_mode, dict): + # Validate that the dict has sets + for key, value in approval_mode.items(): + if not isinstance(value, set): + approval_mode[key] = set(value) # type: ignore - @field_validator("approval_mode") - def validate_approval_mode(cls, approval_mode: str | dict[str, Any] | None) -> str | dict[str, Any] | None: - """Validate the approval_mode field to ensure it is one of the accepted values.""" - if approval_mode is None or not isinstance(approval_mode, dict): - return approval_mode - # Validate that the dict has sets - for key, value in approval_mode.items(): - if not isinstance(value, set): - approval_mode[key] = set(value) # Convert to set if it's a list or other collection - return approval_mode + # Validate allowed_tools + if allowed_tools is not None and isinstance(allowed_tools, dict): + raise TypeError( + f"allowed_tools must be a sequence of strings, not a dict. Got: {type(allowed_tools).__name__}" + ) + + super().__init__( + name=name, + description=description or "", + additional_properties=additional_properties, + **kwargs, + ) + self.url = url if isinstance(url, AnyUrl) else AnyUrl(url) + self.approval_mode = approval_mode + self.allowed_tools = set(allowed_tools) if allowed_tools else None + self.headers = headers + except (ValidationError, ValueError, TypeError) as err: + raise ToolException(f"Error initializing HostedMCPTool: {err}", inner_exception=err) from err 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 - max_results: int | None = None - def __init__( self, inputs: "Contents | dict[str, Any] | str | list[Contents | dict[str, Any] | str] | None" = None, @@ -357,20 +372,18 @@ class HostedFileSearchTool(BaseTool): additional_properties: Additional properties associated with the tool. **kwargs: Additional keyword arguments to pass to the base class. """ - args: dict[str, Any] = { - "name": "file_search", - } - if inputs: - args["inputs"] = _parse_inputs(inputs) - if max_results: - args["max_results"] = max_results - if description is not None: - args["description"] = description - if additional_properties is not None: - args["additional_properties"] = additional_properties if "name" in kwargs: raise ValueError("The 'name' argument is reserved for the HostedFileSearchTool and cannot be set.") - super().__init__(**args, **kwargs) + + self.inputs = _parse_inputs(inputs) if inputs else None + self.max_results = max_results + + super().__init__( + name="file_search", + description=description or "", + additional_properties=additional_properties, + **kwargs, + ) def _default_histogram() -> Histogram: @@ -406,9 +419,38 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]): input_model: The Pydantic model that defines the input parameters for the function. """ - func: Callable[..., Awaitable[ReturnT] | ReturnT] - input_model: type[ArgsT] - _invocation_duration_histogram: Histogram = PrivateAttr(default_factory=_default_histogram) + INJECTABLE: ClassVar[set[str]] = {"func"} + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"input_model", "_invocation_duration_histogram"} + + def __init__( + self, + *, + name: str, + description: str = "", + additional_properties: dict[str, Any] | None = None, + func: Callable[..., Awaitable[ReturnT] | ReturnT], + input_model: type[ArgsT], + **kwargs: Any, + ) -> None: + """Initialize the AIFunction. + + Args: + name: The name of the function. + description: A description of the function. + additional_properties: Additional properties to set on the function. + func: The function to wrap. + input_model: The Pydantic model that defines the input parameters for the function. + **kwargs: Additional keyword arguments. + """ + super().__init__( + name=name, + description=description, + additional_properties=additional_properties, + **kwargs, + ) + self.func = func + self.input_model = input_model + self._invocation_duration_histogram = _default_histogram() def __call__(self, *args: Any, **kwargs: Any) -> ReturnT | Awaitable[ReturnT]: """Call the wrapped function with the provided arguments.""" @@ -668,7 +710,7 @@ def _get_tool_map( tools: "ToolProtocol \ | Callable[..., Any] \ | MutableMapping[str, Any] \ - | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]", + | Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]", ) -> dict[str, AIFunction[Any, Any]]: ai_function_list: dict[str, AIFunction[Any, Any]] = {} for tool in tools if isinstance(tools, list) else [tools]: @@ -689,7 +731,7 @@ async def execute_function_calls( tools: "ToolProtocol \ | Callable[..., Any] \ | MutableMapping[str, Any] \ - | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]", + | Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]", middleware_pipeline: Any = None, # Optional MiddlewarePipeline to avoid circular imports ) -> list["Contents"]: tool_map = _get_tool_map(tools) @@ -778,7 +820,7 @@ def _handle_function_calls_response( # Use the stored middleware pipeline instead of extracting from kwargs # because kwargs may have been modified by the underlying function middleware_pipeline = stored_middleware_pipeline - function_results = await execute_function_calls( + function_call_results: list[Contents] = await execute_function_calls( custom_args=kwargs, attempt_idx=attempt_idx, function_calls=function_calls, @@ -786,7 +828,7 @@ def _handle_function_calls_response( middleware_pipeline=middleware_pipeline, ) # add a single ChatMessage to the response with the results - result_message = ChatMessage(role="tool", contents=function_results) # type: ignore[call-overload] + result_message = ChatMessage(role="tool", contents=function_call_results) response.messages.append(result_message) # response should contain 2 messages after this, # one with function call contents @@ -891,7 +933,7 @@ def _handle_function_calls_streaming_response( update_conversation_id(kwargs, response.conversation_id) prepped_messages = [] - tools = kwargs.get("tools") + tools: Sequence[ToolProtocol | MutableMapping[str, Any]] | None = kwargs.get("tools") if not tools and (chat_options := kwargs.get("chat_options")) and isinstance(chat_options, ChatOptions): tools = chat_options.tools @@ -903,7 +945,7 @@ def _handle_function_calls_streaming_response( custom_args=kwargs, attempt_idx=attempt_idx, function_calls=function_calls, - tools=tools, # type: ignore[reportArgumentType] + tools=tools, middleware_pipeline=middleware_pipeline, ) function_result_msg = ChatMessage(role="tool", contents=function_results) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 9a6fd7558e..aeb06f1b18 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -15,12 +15,10 @@ from collections.abc import ( from copy import deepcopy from typing import Any, ClassVar, Literal, TypeVar, overload -from pydantic import ( - BaseModel, - ValidationError, -) +from pydantic import BaseModel, ValidationError from ._logging import get_logger +from ._serialization import SerializationMixin from ._tools import ToolProtocol, ai_function from .exceptions import AdditionItemMismatch @@ -29,6 +27,39 @@ if sys.version_info >= (3, 11): else: from typing_extensions import Self # pragma: no cover + +__all__ = [ + "AgentRunResponse", + "AgentRunResponseUpdate", + "AnnotatedRegions", + "Annotations", + "BaseAnnotation", + "BaseContent", + "ChatMessage", + "ChatOptions", + "ChatResponse", + "ChatResponseUpdate", + "CitationAnnotation", + "Contents", + "DataContent", + "ErrorContent", + "FinishReason", + "FunctionApprovalRequestContent", + "FunctionApprovalResponseContent", + "FunctionCallContent", + "FunctionResultContent", + "HostedFileContent", + "HostedVectorStoreContent", + "Role", + "TextContent", + "TextReasoningContent", + "TextSpanRegion", + "ToolMode", + "UriContent", + "UsageContent", + "UsageDetails", +] + logger = get_logger("agent_framework") @@ -57,7 +88,46 @@ class EnumLike(type): return cls -def _parse_content_list(contents_data: list[Any]) -> list["Contents"]: +def _parse_content(content_data: MutableMapping[str, Any]) -> "Contents": + """Parse a single content data dictionary into the appropriate Content object. + + Args: + content_data: Content data (dict) + + Returns: + Content object or raises ValidationError if parsing fails + """ + content_type = str(content_data.get("type")) + match content_type: + case "text": + return TextContent.from_dict(content_data) + case "data": + return DataContent.from_dict(content_data) + case "uri": + return UriContent.from_dict(content_data) + case "error": + return ErrorContent.from_dict(content_data) + case "function_call": + return FunctionCallContent.from_dict(content_data) + case "function_result": + return FunctionResultContent.from_dict(content_data) + case "usage": + return UsageContent.from_dict(content_data) + case "hosted_file": + return HostedFileContent.from_dict(content_data) + case "hosted_vector_store": + return HostedVectorStoreContent.from_dict(content_data) + case "function_approval_request": + return FunctionApprovalRequestContent.from_dict(content_data) + case "function_approval_response": + return FunctionApprovalResponseContent.from_dict(content_data) + case "text_reasoning": + return TextReasoningContent.from_dict(content_data) + case _: + raise ValidationError([f"Unknown content type '{content_type}'"], model=Contents) # type: ignore + + +def _parse_content_list(contents_data: Sequence[Any]) -> list["Contents"]: """Parse a list of content data dictionaries into appropriate Content objects. Args: @@ -67,38 +137,13 @@ def _parse_content_list(contents_data: list[Any]) -> list["Contents"]: List of Content objects with unknown types logged and ignored """ contents: list["Contents"] = [] - for content_data in contents_data: if isinstance(content_data, dict): - # Determine the content type and create the appropriate class - content_type = str(content_data.get("type")) - if content_type == "text": - contents.append(TextContent.from_dict(content_data)) - elif content_type == "data": - contents.append(DataContent.from_dict(content_data)) - elif content_type == "uri": - contents.append(UriContent.from_dict(content_data)) - elif content_type == "error": - contents.append(ErrorContent.from_dict(content_data)) - elif content_type == "function_call": - contents.append(FunctionCallContent.from_dict(content_data)) - elif content_type == "function_result": - contents.append(FunctionResultContent.from_dict(content_data)) - elif content_type == "usage": - contents.append(UsageContent.from_dict(content_data)) - elif content_type == "hosted_file": - contents.append(HostedFileContent.from_dict(content_data)) - elif content_type == "hosted_vector_store": - contents.append(HostedVectorStoreContent.from_dict(content_data)) - elif content_type == "function_approval_request": - contents.append(FunctionApprovalRequestContent.from_dict(content_data)) - elif content_type == "function_approval_response": - contents.append(FunctionApprovalResponseContent.from_dict(content_data)) - elif content_type == "text_reasoning": - contents.append(TextReasoningContent.from_dict(content_data)) - else: - # Log unknown content types and ignore them - logger.warning(f"Unknown content type '{content_type}', ignoring: {content_data}") + try: + content = _parse_content(content_data) + contents.append(content) + except ValidationError as ve: + logger.warning(f"Skipping unknown content type or invalid content: {ve}") else: # If it's already a content object, keep it as is contents.append(content_data) @@ -147,40 +192,7 @@ KNOWN_MEDIA_TYPES = [ ] -__all__ = [ - "AgentRunResponse", - "AgentRunResponseUpdate", - "AnnotatedRegions", - "Annotations", - "BaseAnnotation", - "BaseContent", - "ChatMessage", - "ChatOptions", - "ChatResponse", - "ChatResponseUpdate", - "CitationAnnotation", - "Contents", - "DataContent", - "ErrorContent", - "FinishReason", - "FunctionApprovalRequestContent", - "FunctionApprovalResponseContent", - "FunctionCallContent", - "FunctionResultContent", - "HostedFileContent", - "HostedVectorStoreContent", - "Role", - "TextContent", - "TextReasoningContent", - "TextSpanRegion", - "ToolMode", - "UriContent", - "UsageContent", - "UsageDetails", -] - - -class UsageDetails: +class UsageDetails(SerializationMixin): """Provides usage details about a request/response. Attributes: @@ -190,6 +202,8 @@ class UsageDetails: additional_counts: A dictionary of additional token counts, can be set by passing kwargs. """ + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"_extra_counts"} + def __init__( self, input_token_count: int | None = None, @@ -217,19 +231,7 @@ class UsageDetails: raise ValueError(f"Additional counts must be integers, got {type(value).__name__}") self._extra_counts[key] = value - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "UsageDetails": - """Create a UsageDetails instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - UsageDetails instance created from the dictionary. - """ - return cls(**data) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: + def to_dict(self, *, exclude_none: bool = True, exclude: set[str] | None = None) -> dict[str, Any]: """Convert the UsageDetails instance to a dictionary. Args: @@ -239,24 +241,13 @@ class UsageDetails: Returns: Dictionary representation of the UsageDetails instance. """ + # Get the base dict from parent class + result = super().to_dict(exclude_none=exclude_none, exclude=exclude) + + # Add additional counts (extra fields) if exclude is None: exclude = set() - result: dict[str, Any] = {} - - # Handle main fields - for field_name, value in [ - ("input_token_count", self.input_token_count), - ("output_token_count", self.output_token_count), - ("total_token_count", self.total_token_count), - ]: - if field_name in exclude: - continue - if exclude_none and value is None: - continue - result[field_name] = value - - # Add additional counts (extra fields) for key, value in self._extra_counts.items(): if key in exclude: continue @@ -268,9 +259,7 @@ class UsageDetails: def __str__(self) -> str: """Returns a string representation of the usage details.""" - import json - - return json.dumps(self.to_dict(exclude_none=True), indent=4) + return self.to_json() @property def additional_counts(self) -> dict[str, int]: @@ -339,113 +328,10 @@ class UsageDetails: ) -def _process_update( - response: "ChatResponse | AgentRunResponse", update: "ChatResponseUpdate | AgentRunResponseUpdate" -) -> None: - """Processes a single update and modifies the response in place.""" - is_new_message = False - if ( - not response.messages - or ( - update.message_id - and response.messages[-1].message_id - and response.messages[-1].message_id != update.message_id - ) - or (update.role and response.messages[-1].role != update.role) - ): - is_new_message = True - - if is_new_message: - message = ChatMessage(role=Role.ASSISTANT, contents=[]) - response.messages.append(message) - else: - message = response.messages[-1] - # Incorporate the update's properties into the message. - if update.author_name is not None: - message.author_name = update.author_name - if update.role is not None: - message.role = update.role - if update.message_id: - message.message_id = update.message_id - for content in update.contents: - if ( - isinstance(content, FunctionCallContent) - and len(message.contents) > 0 - and isinstance(message.contents[-1], FunctionCallContent) - ): - try: - message.contents[-1] += content - except AdditionItemMismatch: - message.contents.append(content) - elif isinstance(content, UsageContent): - if response.usage_details is None: - response.usage_details = UsageDetails() - response.usage_details += content.details - else: - message.contents.append(content) - # Incorporate the update's properties into the response. - if update.response_id: - response.response_id = update.response_id - if update.created_at is not None: - response.created_at = update.created_at - if update.additional_properties is not None: - if response.additional_properties is None: - response.additional_properties = {} - response.additional_properties.update(update.additional_properties) - if response.raw_representation is None: - response.raw_representation = [] - if not isinstance(response.raw_representation, list): - response.raw_representation = [response.raw_representation] - response.raw_representation.append(update.raw_representation) - if isinstance(response, ChatResponse) and isinstance(update, ChatResponseUpdate): - if update.conversation_id is not None: - response.conversation_id = update.conversation_id - if update.finish_reason is not None: - response.finish_reason = update.finish_reason - if update.model_id is not None: - response.model_id = update.model_id - - -def _coalesce_text_content( - 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["Contents"] = [] - first_new_content: Any | None = None - for content in contents: - if isinstance(content, type_): - if first_new_content is None: - first_new_content = deepcopy(content) - else: - first_new_content += content - else: - # skip this content, it is not of the right type - # so write the existing one to the list and start a new one, - # once the right type is found again - if first_new_content: - coalesced_contents.append(first_new_content) - first_new_content = None - # but keep the other content in the new list - coalesced_contents.append(content) - if first_new_content: - coalesced_contents.append(first_new_content) - contents.clear() - contents.extend(coalesced_contents) - - -def _finalize_response(response: "ChatResponse | AgentRunResponse") -> None: - """Finalizes the response by performing any necessary post-processing.""" - for msg in response.messages: - _coalesce_text_content(msg.contents, TextContent) - _coalesce_text_content(msg.contents, TextReasoningContent) - - # region BaseAnnotation -class TextSpanRegion: +class TextSpanRegion(SerializationMixin): """Represents a region of text that has been annotated.""" def __init__( @@ -471,46 +357,11 @@ class TextSpanRegion: if not hasattr(self, key): setattr(self, key, value) - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "TextSpanRegion": - """Create a TextSpanRegion instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - TextSpanRegion instance created from the dictionary. - """ - return cls(**data) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the TextSpanRegion instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - - Returns: - Dictionary representation of the TextSpanRegion instance. - """ - if exclude is None: - exclude = set() - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - result[key] = value - - return result - AnnotatedRegions = TextSpanRegion -class BaseAnnotation: +class BaseAnnotation(SerializationMixin): """Base class for all AI Annotation types. Args: @@ -519,10 +370,12 @@ class BaseAnnotation: """ + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"raw_representation", "additional_properties"} + def __init__( self, *, - annotated_regions: list[AnnotatedRegions] | None = None, + annotated_regions: list[AnnotatedRegions] | list[MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -530,85 +383,50 @@ class BaseAnnotation: """Initialize BaseAnnotation. Args: - annotated_regions: A list of regions that have been annotated. + annotated_regions: A list of regions that have been annotated. Can be region objects or dicts. additional_properties: Optional additional properties associated with the content. raw_representation: Optional raw representation of the content from an underlying implementation. - **kwargs: Additional keyword arguments (for compatibility with subclasses). + **kwargs: Additional keyword arguments (merged into additional_properties). """ - self.annotated_regions = annotated_regions - self.additional_properties = additional_properties + # Handle annotated_regions conversion from dict format (for SerializationMixin support) + self.annotated_regions: list[AnnotatedRegions] | None = None + if annotated_regions is not None: + converted_regions: list[AnnotatedRegions] = [] + for region_data in annotated_regions: + if isinstance(region_data, MutableMapping): + if region_data.get("type", "") == "text_span": + converted_regions.append(TextSpanRegion.from_dict(region_data)) + else: + logger.warning(f"Unknown region type: {region_data.get('type', '')} in {region_data}") + else: + # Already a region object, keep as is + converted_regions.append(region_data) + self.annotated_regions = converted_regions + + # Merge kwargs into additional_properties + self.additional_properties = additional_properties or {} + self.additional_properties.update(kwargs) + self.raw_representation = raw_representation - # Handle any additional kwargs that weren't consumed by this class - for key, value in kwargs.items(): - if not hasattr(self, key): - setattr(self, key, value) + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + """Convert the instance to a dictionary. - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "BaseAnnotation": - """Create a BaseAnnotation instance from a dictionary. + Extracts additional_properties fields to the root level. Args: - data: Dictionary containing the data to create the instance from. + exclude: Set of field names to exclude from serialization. + exclude_none: Whether to exclude None values from the output. Defaults to True. Returns: - BaseAnnotation instance created from the dictionary. + Dictionary representation of the instance. """ - data_copy = dict(data).copy() + # Get the base dict from SerializationMixin + result = super().to_dict(exclude=exclude, exclude_none=exclude_none) - # Handle annotated_regions - convert from list of dicts to list of AnnotatedRegions objects if needed - if annotated_regions := data_copy.get("annotated_regions"): - regions = [] - for region_data in annotated_regions: - if isinstance(region_data, dict): - region_type = region_data.get("type") - if region_type == "text_span": - regions.append(TextSpanRegion.from_dict(region_data)) - else: - logger.warning(f"Unknown region type: {region_type} in {region_data}") - else: - # Already an object, keep as is - regions.append(region_data) - data_copy["annotated_regions"] = regions - - return cls(**data_copy) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the BaseAnnotation instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - 'raw_representation' is always excluded as per original Pydantic config. - - Returns: - Dictionary representation of the BaseAnnotation instance. - """ - if exclude is None: - exclude = set() - - # Always exclude raw_representation as it was marked with exclude=True in Pydantic - exclude = exclude | {"raw_representation"} - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - - # Special handling for annotated_regions - if key == "annotated_regions" and value is not None: - regions_list = [] - for region in value: - if hasattr(region, "to_dict"): - regions_list.append(region.to_dict(exclude_none=exclude_none)) - else: - # Fallback for non-object regions - regions_list.append(region) - result[key] = regions_list - else: - result[key] = value + # Extract additional_properties to root level + if self.additional_properties: + result.update(self.additional_properties) return result @@ -676,7 +494,7 @@ Annotations = CitationAnnotation TContents = TypeVar("TContents", bound="BaseContent") -class BaseContent: +class BaseContent(SerializationMixin): """Represents content used by AI services. Attributes: @@ -686,10 +504,12 @@ class BaseContent: """ + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"raw_representation", "additional_properties"} + def __init__( self, *, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -697,86 +517,52 @@ class BaseContent: """Initialize BaseContent. Args: - annotations: Optional annotations associated with the content. + annotations: Optional annotations associated with the content. Can be annotation objects or dicts. additional_properties: Optional additional properties associated with the content. raw_representation: Optional raw representation of the content from an underlying implementation. - **kwargs: Additional keyword arguments (for compatibility with subclasses). + **kwargs: Additional keyword arguments (merged into additional_properties). """ - self.annotations = annotations - self.additional_properties = additional_properties + self.annotations: list[Annotations] | None = None + # Handle annotations conversion from dict format (for SerializationMixin support) + if annotations is not None: + converted_annotations: list[Annotations] = [] + for annotation_data in annotations: + if isinstance(annotation_data, Annotations): + # If it's already an annotation object, keep it as is + converted_annotations.append(annotation_data) + elif isinstance(annotation_data, MutableMapping) and annotation_data.get("type", "") == "citation": + converted_annotations.append(CitationAnnotation.from_dict(annotation_data)) + else: + logger.debug( + f"Unknown annotation found: {annotation_data.get('type', 'no_type')}" + f" with data: {annotation_data}" + ) + self.annotations = converted_annotations + + # Merge kwargs into additional_properties + self.additional_properties = additional_properties or {} + self.additional_properties.update(kwargs) + self.raw_representation = raw_representation - # Handle any additional kwargs that weren't consumed by this class - for key, value in kwargs.items(): - if not hasattr(self, key): - setattr(self, key, value) + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + """Convert the instance to a dictionary. - @classmethod - def from_dict(cls: type[TContents], data: dict[str, Any]) -> TContents: - """Create a BaseContent instance from a dictionary. + Extracts additional_properties fields to the root level. Args: - data: Dictionary containing the data to create the instance from. + exclude: Set of field names to exclude from serialization. + exclude_none: Whether to exclude None values from the output. Defaults to True. Returns: - BaseContent instance created from the dictionary. + Dictionary representation of the instance. """ - # Handle annotations conversion from dict format - if "annotations" in data and data["annotations"] is not None: - annotations = [] - for annotation_data in data["annotations"]: - if isinstance(annotation_data, dict): - # Determine the annotation type and create the appropriate class - annotation_type = annotation_data.get("type") - if annotation_type == "citation": - annotations.append(CitationAnnotation.from_dict(annotation_data)) - else: - # Fallback to BaseAnnotation for unknown types - annotations.append(BaseAnnotation.from_dict(annotation_data)) - else: - # If it's already an annotation object, keep it as is - annotations.append(annotation_data) - data = data.copy() - data["annotations"] = annotations + # Get the base dict from SerializationMixin + result = super().to_dict(exclude=exclude, exclude_none=exclude_none) - return cls(**data) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the BaseContent instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - 'raw_representation' is always excluded as per original Pydantic config. - - Returns: - Dictionary representation of the BaseContent instance. - """ - if exclude is None: - exclude = set() - - # Always exclude raw_representation as it was marked with exclude=True in Pydantic - exclude = exclude | {"raw_representation"} - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - - # Handle annotations conversion to dict format - if key == "annotations" and value is not None: - annotations_list = [] - for annotation in value: - if hasattr(annotation, "to_dict"): - annotations_list.append(annotation.to_dict(exclude_none=exclude_none)) - else: - # If it's already a dict or other serializable type, keep it as is - annotations_list.append(annotation) - result[key] = annotations_list - else: - result[key] = value + # Extract additional_properties to root level + if self.additional_properties: + result.update(self.additional_properties) return result @@ -798,7 +584,7 @@ class TextContent(BaseContent): *, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, **kwargs: Any, ): """Initializes a TextContent instance. @@ -830,6 +616,8 @@ class TextContent(BaseContent): """ if not isinstance(other, TextContent): raise TypeError("Incompatible type") + + # Merge raw representations if self.raw_representation is None: raw_representation = other.raw_representation elif other.raw_representation is None: @@ -838,21 +626,27 @@ class TextContent(BaseContent): raw_representation = ( self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation] ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation]) + + # Merge annotations if self.annotations is None: annotations = other.annotations elif other.annotations is None: annotations = self.annotations else: annotations = self.annotations + other.annotations - return TextContent( - text=self.text + other.text, - annotations=annotations, - additional_properties={ + + # Create new instance using from_dict for proper deserialization + result_dict = { + "text": self.text + other.text, + "type": "text", + "annotations": [ann.to_dict(exclude_none=False) for ann in annotations] if annotations else None, + "additional_properties": { **(other.additional_properties or {}), **(self.additional_properties or {}), }, - raw_representation=raw_representation, - ) + "raw_representation": raw_representation, + } + return TextContent.from_dict(result_dict) def __iadd__(self, other: "TextContent") -> Self: """In-place concatenation of two TextContent instances. @@ -865,21 +659,33 @@ class TextContent(BaseContent): """ if not isinstance(other, TextContent): raise TypeError("Incompatible type") + + # Concatenate text self.text += other.text + + # Merge additional properties (self takes precedence) if self.additional_properties is None: self.additional_properties = {} if other.additional_properties: - self.additional_properties = {**other.additional_properties, **self.additional_properties} + # Update from other first, then restore self's values to maintain precedence + self_props = self.additional_properties.copy() + self.additional_properties.update(other.additional_properties) + self.additional_properties.update(self_props) + + # Merge raw representations if self.raw_representation is None: self.raw_representation = other.raw_representation elif other.raw_representation is not None: self.raw_representation = ( self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation] ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation]) + + # Merge annotations if other.annotations: if self.annotations is None: self.annotations = [] self.annotations.extend(other.annotations) + return self @@ -903,7 +709,7 @@ class TextReasoningContent(BaseContent): *, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, **kwargs: Any, ): """Initializes a TextReasoningContent instance. @@ -935,6 +741,8 @@ class TextReasoningContent(BaseContent): """ if not isinstance(other, TextReasoningContent): raise TypeError("Incompatible type") + + # Merge raw representations if self.raw_representation is None: raw_representation = other.raw_representation elif other.raw_representation is None: @@ -943,18 +751,24 @@ class TextReasoningContent(BaseContent): raw_representation = ( self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation] ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation]) + + # Merge annotations if self.annotations is None: annotations = other.annotations elif other.annotations is None: annotations = self.annotations else: annotations = self.annotations + other.annotations - return TextReasoningContent( - text=self.text + other.text, - annotations=annotations, - additional_properties={**(self.additional_properties or {}), **(other.additional_properties or {})}, - raw_representation=raw_representation, - ) + + # Create new instance using from_dict for proper deserialization + result_dict = { + "text": self.text + other.text, + "type": "text_reasoning", + "annotations": [ann.to_dict(exclude_none=False) for ann in annotations] if annotations else None, + "additional_properties": {**(self.additional_properties or {}), **(other.additional_properties or {})}, + "raw_representation": raw_representation, + } + return TextReasoningContent.from_dict(result_dict) def __iadd__(self, other: "TextReasoningContent") -> Self: """In-place concatenation of two TextReasoningContent instances. @@ -967,21 +781,33 @@ class TextReasoningContent(BaseContent): """ if not isinstance(other, TextReasoningContent): raise TypeError("Incompatible type") + + # Concatenate text self.text += other.text + + # Merge additional properties (self takes precedence) if self.additional_properties is None: self.additional_properties = {} if other.additional_properties: + # Update from other first, then restore self's values to maintain precedence + self_props = self.additional_properties.copy() self.additional_properties.update(other.additional_properties) + self.additional_properties.update(self_props) + + # Merge raw representations if self.raw_representation is None: self.raw_representation = other.raw_representation elif other.raw_representation is not None: self.raw_representation = ( self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation] ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation]) + + # Merge annotations if other.annotations: if self.annotations is None: self.annotations = [] self.annotations.extend(other.annotations) + return self @@ -1004,7 +830,7 @@ class DataContent(BaseContent): self, *, uri: str, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1030,7 +856,7 @@ class DataContent(BaseContent): *, data: bytes, media_type: str, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1057,7 +883,7 @@ class DataContent(BaseContent): uri: str | None = None, data: bytes | None = None, media_type: str | None = None, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1141,7 +967,7 @@ class UriContent(BaseContent): uri: str, media_type: str, *, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1211,7 +1037,7 @@ class ErrorContent(BaseContent): message: str | None = None, error_code: str | None = None, details: str | None = None, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1265,7 +1091,7 @@ class FunctionCallContent(BaseContent): name: str, arguments: str | dict[str, Any | None] | None = None, exception: Exception | None = None, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1352,7 +1178,7 @@ class FunctionResultContent(BaseContent): call_id: str, result: Any | None = None, exception: Exception | None = None, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1394,9 +1220,9 @@ class UsageContent(BaseContent): def __init__( self, - details: UsageDetails, + details: UsageDetails | MutableMapping[str, Any], *, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1408,64 +1234,12 @@ class UsageContent(BaseContent): raw_representation=raw_representation, **kwargs, ) + # Convert dict to UsageDetails if needed + if isinstance(details, MutableMapping): + details = UsageDetails.from_dict(details) self.details = details self.type: Literal["usage"] = "usage" - @classmethod - def from_dict(cls: type[TContents], data: dict[str, Any]) -> TContents: - """Create a UsageContent instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - UsageContent instance created from the dictionary. - """ - data_copy = data.copy() - - # Handle details - convert from dict to UsageDetails if needed - if "details" in data_copy and isinstance(data_copy["details"], dict): - data_copy["details"] = UsageDetails.from_dict(data_copy["details"]) - - # Handle annotations via BaseContent logic - if "annotations" in data_copy and data_copy["annotations"] is not None: - annotations = [] - for annotation_data in data_copy["annotations"]: - if isinstance(annotation_data, dict): - # Determine the annotation type and create the appropriate class - annotation_type = annotation_data.get("type") - if annotation_type == "citation": - annotations.append(CitationAnnotation.from_dict(annotation_data)) - else: - # Fallback to BaseAnnotation for unknown types - annotations.append(BaseAnnotation.from_dict(annotation_data)) - else: - # If it's already an annotation object, keep it as is - annotations.append(annotation_data) - data_copy["annotations"] = annotations - - return cls(**data_copy) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the UsageContent instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - - Returns: - Dictionary representation of the UsageContent instance. - """ - result = super().to_dict(exclude_none=exclude_none, exclude=exclude) - - # Handle details - convert to dict if it's a UsageDetails object - if hasattr(self.details, "to_dict"): - result["details"] = self.details.to_dict(exclude_none=exclude_none) - else: - result["details"] = self.details - - return result - class HostedFileContent(BaseContent): """Represents a hosted file content. @@ -1532,7 +1306,7 @@ class BaseUserInputRequest(BaseContent): self, *, id: str, - annotations: list[Annotations] | None = None, + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1566,8 +1340,8 @@ class FunctionApprovalResponseContent(BaseContent): approved: bool, *, id: str, - function_call: FunctionCallContent, - annotations: list[Annotations] | None = None, + function_call: FunctionCallContent | MutableMapping[str, Any], + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1577,7 +1351,7 @@ class FunctionApprovalResponseContent(BaseContent): Args: approved: Whether the function call was approved. id: The unique identifier for the request. - function_call: The function call content to be approved. + function_call: The function call content to be approved. Can be a FunctionCallContent object or dict. annotations: Optional list of annotations for the request. additional_properties: Optional additional properties for the request. raw_representation: Optional raw representation of the request. @@ -1591,48 +1365,14 @@ class FunctionApprovalResponseContent(BaseContent): ) self.id = id self.approved = approved - self.function_call = function_call + # Convert dict to FunctionCallContent if needed (for SerializationMixin support) + if isinstance(function_call, MutableMapping): + self.function_call = FunctionCallContent.from_dict(function_call) + else: + self.function_call = function_call # Override the type for this specific subclass self.type: Literal["function_approval_response"] = "function_approval_response" - @classmethod - def from_dict(cls: type[TContents], data: dict[str, Any]) -> TContents: - """Create a FunctionApprovalResponseContent instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - FunctionApprovalResponseContent instance created from the dictionary. - """ - data_copy = data.copy() - - # Handle function_call - convert from dict to FunctionCallContent if needed - if (function_call := data_copy.get("function_call")) and isinstance(function_call, dict): - data_copy["function_call"] = FunctionCallContent.from_dict(function_call) - - return cls(**data_copy) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the FunctionApprovalResponseContent instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - - Returns: - Dictionary representation of the FunctionApprovalResponseContent instance. - """ - result = super().to_dict(exclude_none=exclude_none, exclude=exclude) - - # Handle function_call - convert to dict if it's a content object - if hasattr(self.function_call, "to_dict"): - result["function_call"] = self.function_call.to_dict(exclude_none=exclude_none) - else: - result["function_call"] = self.function_call - - return result - class FunctionApprovalRequestContent(BaseContent): """Represents a request for user approval of a function call.""" @@ -1641,8 +1381,8 @@ class FunctionApprovalRequestContent(BaseContent): self, *, id: str, - function_call: FunctionCallContent, - annotations: list[Annotations] | None = None, + function_call: FunctionCallContent | MutableMapping[str, Any], + annotations: list[Annotations | MutableMapping[str, Any]] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, **kwargs: Any, @@ -1651,7 +1391,7 @@ class FunctionApprovalRequestContent(BaseContent): Args: id: The unique identifier for the request. - function_call: The function call content to be approved. + function_call: The function call content to be approved. Can be a FunctionCallContent object or dict. annotations: Optional list of annotations for the request. additional_properties: Optional additional properties for the request. raw_representation: Optional raw representation of the request. @@ -1664,48 +1404,14 @@ class FunctionApprovalRequestContent(BaseContent): **kwargs, ) self.id = id - self.function_call = function_call + # Convert dict to FunctionCallContent if needed (for SerializationMixin support) + if isinstance(function_call, MutableMapping): + self.function_call = FunctionCallContent.from_dict(function_call) + else: + self.function_call = function_call # Override the type for this specific subclass self.type: Literal["function_approval_request"] = "function_approval_request" - @classmethod - def from_dict(cls: type[TContents], data: dict[str, Any]) -> TContents: - """Create a FunctionApprovalRequestContent instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - FunctionApprovalRequestContent instance created from the dictionary. - """ - data_copy = data.copy() - - # Handle function_call - convert from dict to FunctionCallContent if needed - if "function_call" in data_copy and isinstance(data_copy["function_call"], dict): - data_copy["function_call"] = FunctionCallContent.from_dict(data_copy["function_call"]) - - return cls(**data_copy) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the FunctionApprovalRequestContent instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - - Returns: - Dictionary representation of the FunctionApprovalRequestContent instance. - """ - result = super().to_dict(exclude_none=exclude_none, exclude=exclude) - - # Handle function_call - convert to dict if it's a content object - if hasattr(self.function_call, "to_dict"): - result["function_call"] = self.function_call.to_dict(exclude_none=exclude_none) - else: - result["function_call"] = self.function_call - - return result - def create_response(self, approved: bool) -> "FunctionApprovalResponseContent": """Create a response for the function approval request.""" return FunctionApprovalResponseContent( @@ -1736,7 +1442,7 @@ Contents = ( # region Chat Response constants -class Role(metaclass=EnumLike): +class Role(SerializationMixin, metaclass=EnumLike): """Describes the intended purpose of a message within a chat interaction. Attributes: @@ -1771,41 +1477,6 @@ class Role(metaclass=EnumLike): """ self.value = value - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "Role": - """Create a Role instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - Role instance created from the dictionary. - """ - return cls(**data) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the Role instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - - Returns: - Dictionary representation of the Role instance. - """ - if exclude is None: - exclude = set() - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - result[key] = value - - return result - def __str__(self) -> str: """Returns the string representation of the role.""" return self.value @@ -1825,7 +1496,7 @@ class Role(metaclass=EnumLike): return hash(self.value) -class FinishReason(metaclass=EnumLike): +class FinishReason(SerializationMixin, metaclass=EnumLike): """Represents the reason a chat response completed. Attributes: @@ -1854,41 +1525,6 @@ class FinishReason(metaclass=EnumLike): """ self.value = value - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "FinishReason": - """Create a FinishReason instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - FinishReason instance created from the dictionary. - """ - return cls(**data) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the FinishReason instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - - Returns: - Dictionary representation of the FinishReason instance. - """ - if exclude is None: - exclude = set() - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - result[key] = value - - return result - def __eq__(self, other: object) -> bool: """Check if two FinishReason instances are equal.""" if not isinstance(other, FinishReason): @@ -1911,7 +1547,7 @@ class FinishReason(metaclass=EnumLike): # region ChatMessage -class ChatMessage: +class ChatMessage(SerializationMixin): """Represents a chat message. Attributes: @@ -1924,6 +1560,8 @@ class ChatMessage: """ + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"raw_representation"} + @overload def __init__( self, @@ -1932,8 +1570,9 @@ class ChatMessage: text: str, author_name: str | None = None, message_id: str | None = None, - additional_properties: dict[str, Any] | None = None, + additional_properties: MutableMapping[str, Any] | None = None, raw_representation: Any | None = None, + **kwargs: Any, ) -> None: """Initializes a ChatMessage with a role and text content. @@ -1944,6 +1583,7 @@ class ChatMessage: message_id: Optional ID of the chat message. additional_properties: Optional additional properties associated with the chat message. raw_representation: Optional raw representation of the chat message. + **kwargs: Additional keyword arguments. """ @overload @@ -1951,11 +1591,12 @@ class ChatMessage: self, role: Role | Literal["system", "user", "assistant", "tool"], *, - contents: MutableSequence[Contents], + contents: Sequence[Contents | Mapping[str, Any]], author_name: str | None = None, message_id: str | None = None, - additional_properties: dict[str, Any] | None = None, + additional_properties: MutableMapping[str, Any] | None = None, raw_representation: Any | None = None, + **kwargs: Any, ) -> None: """Initializes a ChatMessage with a role and optional contents. @@ -1966,110 +1607,53 @@ class ChatMessage: message_id: Optional ID of the chat message. additional_properties: Optional additional properties associated with the chat message. raw_representation: Optional raw representation of the chat message. + **kwargs: Additional keyword arguments. """ def __init__( self, - role: Role | Literal["system", "user", "assistant", "tool"], + role: Role | Literal["system", "user", "assistant", "tool"] | dict[str, Any], *, text: str | None = None, - contents: MutableSequence[Contents] | None = None, + contents: Sequence[Contents | Mapping[str, Any]] | None = None, author_name: str | None = None, message_id: str | None = None, - additional_properties: dict[str, Any] | None = None, + additional_properties: MutableMapping[str, Any] | None = None, raw_representation: Any | None = None, + **kwargs: Any, ) -> None: """Initialize ChatMessage. Args: - role: The role of the author of the message. + role: The role of the author of the message (Role, string, or dict). text: Optional text content of the message. - contents: Optional list of BaseContent items to include in the message. + contents: Optional list of BaseContent items or dicts 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. raw_representation: Optional raw representation of the chat message. + kwargs: will be combined with additional_properties if provided. """ - if contents is None: - contents = [] - if text is not None: - contents.append(TextContent(text=text)) - if isinstance(role, str): + # Handle role conversion + if isinstance(role, dict): + role = Role.from_dict(role) + elif isinstance(role, str): role = Role(value=role) + # Handle contents conversion + parsed_contents = [] if contents is None else _parse_content_list(contents) + + if text is not None: + parsed_contents.append(TextContent(text=text)) + self.role = role - self.contents = list(contents) # Convert to list to ensure it's mutable + self.contents = parsed_contents self.author_name = author_name self.message_id = message_id - self.additional_properties = additional_properties + self.additional_properties = additional_properties or {} + self.additional_properties.update(kwargs or {}) self.raw_representation = raw_representation - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "ChatMessage": - """Create a ChatMessage instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - ChatMessage instance created from the dictionary. - """ - data_copy = dict(data).copy() - - # Handle role - convert from dict to Role if needed - if role := data.get("role"): - data_copy["role"] = Role.from_dict(role) if isinstance(role, dict) else Role(value=role) - # Handle contents - convert from list of dicts to list of Contents objects if needed - if contents := data.get("contents"): - data_copy["contents"] = _parse_content_list(contents) - - return cls(**data_copy) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the ChatMessage instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - 'raw_representation' is always excluded as per original Pydantic config. - - Returns: - Dictionary representation of the ChatMessage instance. - """ - if exclude is None: - exclude = set() - - # Always exclude raw_representation as it was marked with exclude=True in Pydantic - exclude = exclude | {"raw_representation"} - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - - # Handle role conversion to dict format - if key == "role" and value is not None: - if hasattr(value, "to_dict"): - result[key] = value.to_dict(exclude_none=exclude_none) - else: - result[key] = value - # Handle contents conversion to dict format - elif key == "contents" and value is not None: - contents_list = [] - for content in value: - if hasattr(content, "to_dict"): - contents_list.append(content.to_dict(exclude_none=exclude_none)) - else: - # If it's already a dict or other serializable type, keep it as is - contents_list.append(content) - result[key] = contents_list - else: - result[key] = value - - return result - @property def text(self) -> str: """Returns the text content of the message. @@ -2083,7 +1667,116 @@ class ChatMessage: # region ChatResponse -class ChatResponse: +def _process_update( + response: "ChatResponse | AgentRunResponse", update: "ChatResponseUpdate | AgentRunResponseUpdate" +) -> None: + """Processes a single update and modifies the response in place.""" + is_new_message = False + if ( + not response.messages + or ( + update.message_id + and response.messages[-1].message_id + and response.messages[-1].message_id != update.message_id + ) + or (update.role and response.messages[-1].role != update.role) + ): + is_new_message = True + + if is_new_message: + message = ChatMessage(role=Role.ASSISTANT, contents=[]) + response.messages.append(message) + else: + message = response.messages[-1] + # Incorporate the update's properties into the message. + if update.author_name is not None: + message.author_name = update.author_name + if update.role is not None: + message.role = update.role + if update.message_id: + message.message_id = update.message_id + for content in update.contents: + if ( + isinstance(content, FunctionCallContent) + and len(message.contents) > 0 + and isinstance(message.contents[-1], FunctionCallContent) + ): + try: + message.contents[-1] += content + except AdditionItemMismatch: + message.contents.append(content) + elif isinstance(content, UsageContent): + if response.usage_details is None: + response.usage_details = UsageDetails() + response.usage_details += content.details + elif isinstance(content, (dict, MutableMapping)): + try: + cont = _parse_content(content) + message.contents.append(cont) + except ValidationError as ve: + logger.warning(f"Skipping unknown content type or invalid content: {ve}") + else: + message.contents.append(content) + # Incorporate the update's properties into the response. + if update.response_id: + response.response_id = update.response_id + if update.created_at is not None: + response.created_at = update.created_at + if update.additional_properties is not None: + if response.additional_properties is None: + response.additional_properties = {} + response.additional_properties.update(update.additional_properties) + if response.raw_representation is None: + response.raw_representation = [] + if not isinstance(response.raw_representation, list): + response.raw_representation = [response.raw_representation] + response.raw_representation.append(update.raw_representation) + if isinstance(response, ChatResponse) and isinstance(update, ChatResponseUpdate): + if update.conversation_id is not None: + response.conversation_id = update.conversation_id + if update.finish_reason is not None: + response.finish_reason = update.finish_reason + if update.model_id is not None: + response.model_id = update.model_id + + +def _coalesce_text_content( + 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["Contents"] = [] + first_new_content: Any | None = None + for content in contents: + if isinstance(content, type_): + if first_new_content is None: + first_new_content = deepcopy(content) + else: + first_new_content += content + else: + # skip this content, it is not of the right type + # so write the existing one to the list and start a new one, + # once the right type is found again + if first_new_content: + coalesced_contents.append(first_new_content) + first_new_content = None + # but keep the other content in the new list + coalesced_contents.append(content) + if first_new_content: + coalesced_contents.append(first_new_content) + contents.clear() + contents.extend(coalesced_contents) + + +def _finalize_response(response: "ChatResponse | AgentRunResponse") -> None: + """Finalizes the response by performing any necessary post-processing.""" + for msg in response.messages: + _coalesce_text_content(msg.contents, TextContent) + _coalesce_text_content(msg.contents, TextReasoningContent) + + +class ChatResponse(SerializationMixin): """Represents the response to a chat request. Attributes: @@ -2099,6 +1792,8 @@ class ChatResponse: raw_representation: The raw representation of the chat response from an underlying implementation. """ + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"raw_representation", "additional_properties"} + @overload def __init__( self, @@ -2172,14 +1867,14 @@ class ChatResponse: def __init__( self, *, - messages: ChatMessage | MutableSequence[ChatMessage] | None = None, + messages: ChatMessage | MutableSequence[ChatMessage] | list[dict[str, Any]] | None = None, text: TextContent | str | None = None, response_id: str | None = None, conversation_id: str | None = None, model_id: str | None = None, created_at: CreatedAtT | None = None, - finish_reason: FinishReason | None = None, - usage_details: UsageDetails | None = None, + finish_reason: FinishReason | dict[str, Any] | None = None, + usage_details: UsageDetails | dict[str, Any] | None = None, value: Any | None = None, response_format: type[BaseModel] | None = None, additional_properties: dict[str, Any] | None = None, @@ -2187,15 +1882,34 @@ class ChatResponse: **kwargs: Any, ) -> None: """Initializes a ChatResponse with the provided parameters.""" + # Handle messages conversion if messages is None: messages = [] elif not isinstance(messages, MutableSequence): messages = [messages] + else: + # Convert any dicts in messages list to ChatMessage objects + converted_messages: list[ChatMessage] = [] + for msg in messages: + if isinstance(msg, dict): + converted_messages.append(ChatMessage.from_dict(msg)) + else: + converted_messages.append(msg) + messages = converted_messages + if text is not None: if isinstance(text, str): text = TextContent(text=text) messages.append(ChatMessage(role=Role.ASSISTANT, contents=[text])) + # Handle finish_reason conversion + if isinstance(finish_reason, dict): + finish_reason = FinishReason.from_dict(finish_reason) + + # Handle usage_details conversion + if isinstance(usage_details, dict): + usage_details = UsageDetails.from_dict(usage_details) + self.messages = list(messages) self.response_id = response_id self.conversation_id = conversation_id @@ -2204,93 +1918,13 @@ class ChatResponse: self.finish_reason = finish_reason self.usage_details = usage_details self.value = value - self.additional_properties = additional_properties - self.raw_representation = raw_representation - - # Handle any additional kwargs - for key, value in kwargs.items(): - setattr(self, key, value) + self.additional_properties = additional_properties or {} + self.additional_properties.update(kwargs or {}) + self.raw_representation: Any | list[Any] | None = raw_representation if response_format: self.try_parse_value(output_format_type=response_format) - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "ChatResponse": - """Create a ChatResponse instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - ChatResponse instance created from the dictionary. - """ - data_copy = dict(data).copy() - - # Handle messages - convert from list of dicts to list of ChatMessage objects if needed - if messages := data_copy.get("messages"): - parsed_messages = [] - for message_data in messages: - if isinstance(message_data, ChatMessage): - parsed_messages.append(message_data) - elif isinstance(message_data, dict): - parsed_messages.append(ChatMessage.from_dict(message_data)) - else: - logger.warning(f"Unknown type for message: {message_data}") - data_copy["messages"] = parsed_messages - - # Handle finish_reason - convert from dict to FinishReason if needed - if finish_reason := data_copy.get("finish_reason"): - data_copy["finish_reason"] = FinishReason.from_dict(finish_reason) - - # Handle usage_details - convert from dict to UsageDetails if needed - if (usage_details := data_copy.get("usage_details")) and isinstance(usage_details, dict): - data_copy["usage_details"] = UsageDetails.from_dict(usage_details) - return cls(**data_copy) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the ChatResponse instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - 'raw_representation' is always excluded as per original Pydantic config. - - Returns: - Dictionary representation of the ChatResponse instance. - """ - if exclude is None: - exclude = set() - - # Always exclude raw_representation as it was marked with exclude=True in Pydantic - exclude = exclude | {"raw_representation"} - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - - # Handle messages conversion to dict format - if key == "messages" and value is not None: - messages_list = [] - for message in value: - if hasattr(message, "to_dict"): - messages_list.append(message.to_dict(exclude_none=exclude_none)) - else: - messages_list.append(message) - result[key] = messages_list - # Handle finish_reason conversion to dict format - elif (key == "finish_reason" and value is not None) or (key == "usage_details" and value is not None): - if hasattr(value, "to_dict"): - result[key] = value.to_dict(exclude_none=exclude_none) - else: - result[key] = value - else: - result[key] = value - - return result - @classmethod def from_chat_response_updates( cls: type[TChatResponse], @@ -2343,8 +1977,8 @@ class ChatResponse: # region ChatResponseUpdate -class ChatResponseUpdate: - """Represents a single streaming response chunk from a `ModelClient`. +class ChatResponseUpdate(SerializationMixin): + """Represents a single streaming response chunk from a `ChatClient`. Attributes: contents: The chat response update content items. @@ -2361,68 +1995,62 @@ class ChatResponseUpdate: """ - @overload - def __init__( - self, - *, - 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, - model_id: str | None = None, - created_at: CreatedAtT | None = None, - finish_reason: FinishReason | None = None, - additional_properties: dict[str, Any] | None = None, - raw_representation: Any | None = None, - ) -> None: - """Initializes a ChatResponseUpdate with the provided parameters.""" - - @overload - def __init__( - self, - *, - text: TextContent | str, - 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, - model_id: str | None = None, - created_at: CreatedAtT | None = None, - finish_reason: FinishReason | None = None, - additional_properties: dict[str, Any] | None = None, - raw_representation: Any | None = None, - ) -> None: - """Initializes a ChatResponseUpdate with the provided parameters.""" + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"raw_representation"} def __init__( self, *, - contents: list[Contents] | None = None, + contents: Sequence[Contents | dict[str, Any]] | None = None, text: TextContent | str | None = None, - role: Role | Literal["system", "user", "assistant", "tool"] | None = None, + role: Role | Literal["system", "user", "assistant", "tool"] | dict[str, Any] | None = None, author_name: str | None = None, response_id: str | None = None, message_id: str | None = None, conversation_id: str | None = None, model_id: str | None = None, created_at: CreatedAtT | None = None, - finish_reason: FinishReason | None = None, + finish_reason: FinishReason | dict[str, Any] | None = None, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, + **kwargs: Any, ) -> None: - """Initializes a ChatResponseUpdate with the provided parameters.""" - if contents is None: - contents = [] + """Initializes a ChatResponseUpdate with the provided parameters. + + Args: + contents: Optional list of BaseContent items or dicts to include in the update. + text: Optional text content to include in the update. + role: Optional role of the author of the response update (Role, string, or dict + author_name: Optional name of the author of the response update. + response_id: Optional ID of the response of which this update is a part. + message_id: Optional ID of the message of which this update is a part. + conversation_id: Optional identifier for the state of the conversation of which this update is a part + model_id: Optional model ID associated with this response update. + created_at: Optional timestamp for the chat response update. + finish_reason: Optional finish reason for the operation. + additional_properties: Optional additional properties associated with the chat response update. + raw_representation: Optional raw representation of the chat response update + from an underlying implementation. + **kwargs: Any additional keyword arguments. + + """ + # Handle contents conversion + contents = [] if contents is None else _parse_content_list(contents) + if text is not None: if isinstance(text, str): text = TextContent(text=text) contents.append(text) - if role and isinstance(role, str): + + # Handle role conversion + if isinstance(role, dict): + role = Role.from_dict(role) + elif isinstance(role, str): role = Role(value=role) + # Handle finish_reason conversion + if isinstance(finish_reason, dict): + finish_reason = FinishReason.from_dict(finish_reason) + self.contents = list(contents) self.role = role self.author_name = author_name @@ -2435,81 +2063,6 @@ class ChatResponseUpdate: self.additional_properties = additional_properties self.raw_representation = raw_representation - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "ChatResponseUpdate": - """Create a ChatResponseUpdate instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - ChatResponseUpdate instance created from the dictionary. - """ - data_copy = dict(data).copy() - - # Handle contents - convert from list of dicts to list of Contents objects if needed - if "contents" in data_copy and data_copy["contents"] is not None: - data_copy["contents"] = _parse_content_list(data_copy["contents"]) - # Handle role - convert from dict to Role if needed - if role := data.get("role"): - data_copy["role"] = Role.from_dict(role) if isinstance(role, dict) else Role(value=role) - # Handle contents - convert from list of dicts to list of Contents objects if needed - if contents := data.get("contents"): - data_copy["contents"] = _parse_content_list(contents) - - # Handle finish_reason - convert from dict to FinishReason if needed - if finish_reason := data.get("finish_reason"): - data_copy["finish_reason"] = ( - FinishReason.from_dict(finish_reason) - if isinstance(finish_reason, dict) - else FinishReason(value=finish_reason) - ) - return cls(**data_copy) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the ChatResponseUpdate instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - 'raw_representation' is always excluded as per original Pydantic config. - - Returns: - Dictionary representation of the ChatResponseUpdate instance. - """ - if exclude is None: - exclude = set() - - # Always exclude raw_representation as it was marked with exclude=True in Pydantic - exclude = exclude | {"raw_representation"} - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - - # Handle contents conversion to dict format - if key == "contents" and value is not None: - contents_list = [] - for content in value: - if hasattr(content, "to_dict"): - contents_list.append(content.to_dict(exclude_none=exclude_none)) - else: - contents_list.append(content) - result[key] = contents_list - # Handle role conversion to dict format - elif (key == "role" and value is not None) or (key == "finish_reason" and value is not None): - if hasattr(value, "to_dict"): - result[key] = value.to_dict(exclude_none=exclude_none) - else: - result[key] = value - else: - result[key] = value - - return result - @property def text(self) -> str: """Returns the concatenated text of all contents in the update.""" @@ -2533,10 +2086,214 @@ class ChatResponseUpdate: return ChatResponseUpdate.from_dict(current_data) +# region AgentRunResponse + + +class AgentRunResponse(SerializationMixin): + """Represents the response to an Agent run request. + + Provides one or more response messages and metadata about the response. + A typical response will contain a single message, but may contain multiple + messages in scenarios involving function calls, RAG retrievals, or complex logic. + """ + + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"raw_representation"} + + def __init__( + self, + messages: ChatMessage + | list[ChatMessage] + | MutableMapping[str, Any] + | list[MutableMapping[str, Any]] + | None = None, + response_id: str | None = None, + created_at: CreatedAtT | None = None, + usage_details: UsageDetails | MutableMapping[str, Any] | None = None, + value: Any | None = None, + raw_representation: Any | None = None, + additional_properties: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Initialize an AgentRunResponse. + + Attributes: + messages: The list of chat messages in the response. + response_id: The ID of the chat response. + created_at: A timestamp for the chat response. + usage_details: The usage details for the chat response. + value: The structured output of the agent run response, if applicable. + additional_properties: Any additional properties associated with the chat response. + raw_representation: The raw representation of the chat response from an underlying implementation. + **kwargs: Additional properties to set on the response. + """ + processed_messages: list[ChatMessage] = [] + if messages is not None: + if isinstance(messages, ChatMessage): + processed_messages.append(messages) + elif isinstance(messages, list): + for message_data in messages: + if isinstance(message_data, ChatMessage): + processed_messages.append(message_data) + elif isinstance(message_data, MutableMapping): + processed_messages.append(ChatMessage.from_dict(message_data)) + else: + logger.warning(f"Unknown message content: {message_data}") + elif isinstance(messages, MutableMapping): + processed_messages.append(ChatMessage.from_dict(messages)) + + # Convert usage_details from dict if needed (for SerializationMixin support) + if isinstance(usage_details, MutableMapping): + usage_details = UsageDetails.from_dict(usage_details) + + self.messages = processed_messages + self.response_id = response_id + self.created_at = created_at + self.usage_details = usage_details + self.value = value + self.additional_properties = additional_properties or {} + self.additional_properties.update(kwargs or {}) + self.raw_representation = raw_representation + + @property + def text(self) -> str: + """Get the concatenated text of all messages.""" + return "".join(msg.text for msg in self.messages) if self.messages else "" + + @property + def user_input_requests(self) -> list[UserInputRequestContents]: + """Get all BaseUserInputRequest messages from the response.""" + return [ + content + for msg in self.messages + for content in msg.contents + if isinstance(content, UserInputRequestContents) + ] + + @classmethod + def from_agent_run_response_updates( + cls: type[TAgentRunResponse], + updates: Sequence["AgentRunResponseUpdate"], + *, + output_format_type: type[BaseModel] | None = None, + ) -> TAgentRunResponse: + """Joins multiple updates into a single AgentRunResponse.""" + msg = cls(messages=[]) + for update in updates: + _process_update(msg, update) + _finalize_response(msg) + if output_format_type: + msg.try_parse_value(output_format_type) + return msg + + @classmethod + async def from_agent_response_generator( + cls: type[TAgentRunResponse], + updates: AsyncIterable["AgentRunResponseUpdate"], + *, + output_format_type: type[BaseModel] | None = None, + ) -> TAgentRunResponse: + """Joins multiple updates into a single AgentRunResponse.""" + msg = cls(messages=[]) + async for update in updates: + _process_update(msg, update) + _finalize_response(msg) + if output_format_type: + msg.try_parse_value(output_format_type) + return msg + + def __str__(self) -> str: + return self.text + + def try_parse_value(self, output_format_type: type[BaseModel]) -> None: + """If there is a value, does nothing, otherwise tries to parse the text into the value.""" + if self.value is None: + try: + self.value = output_format_type.model_validate_json(self.text) # type: ignore[reportUnknownMemberType] + except ValidationError as ex: + logger.debug("Failed to parse value from agent run response text: %s", ex) + + +# region AgentRunResponseUpdate + + +class AgentRunResponseUpdate(SerializationMixin): + """Represents a single streaming response chunk from an Agent.""" + + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"raw_representation"} + + def __init__( + self, + *, + contents: Sequence[Contents | MutableMapping[str, Any]] | None = None, + text: TextContent | str | None = None, + role: Role | MutableMapping[str, Any] | str | None = None, + author_name: str | None = None, + response_id: str | None = None, + message_id: str | None = None, + created_at: CreatedAtT | None = None, + additional_properties: MutableMapping[str, Any] | None = None, + raw_representation: Any | None = None, + **kwargs: Any, + ) -> None: + """Initialize an AgentRunResponseUpdate. + + Args: + contents: Optional list of BaseContent items or dicts to include in the update. + text: Optional text content of the update. + role: The role of the author of the response update (Role, string, or dict + author_name: Optional name of the author of the response update. + response_id: Optional ID of the response of which this update is a part. + message_id: Optional ID of the message of which this update is a part. + created_at: Optional timestamp for the chat response update. + additional_properties: Optional additional properties associated with the chat response update. + raw_representation: Optional raw representation of the chat response update. + kwargs: will be combined with additional_properties if provided. + + """ + contents = [] if contents is None else _parse_content_list(contents) + + if text is not None: + if isinstance(text, str): + text = TextContent(text=text) + contents.append(text) + + # Convert role from dict if needed (for SerializationMixin support) + if isinstance(role, MutableMapping): + role = Role.from_dict(role) + elif isinstance(role, str): + role = Role(value=role) + + self.contents = contents + self.role = role + self.author_name = author_name + self.response_id = response_id + self.message_id = message_id + self.created_at = created_at + self.additional_properties = additional_properties + self.raw_representation: Any | list[Any] | None = raw_representation + + @property + def text(self) -> str: + """Get the concatenated text of all TextContent objects in contents.""" + return ( + "".join(content.text for content in self.contents if isinstance(content, TextContent)) + if self.contents + else "" + ) + + @property + def user_input_requests(self) -> list[UserInputRequestContents]: + """Get all BaseUserInputRequest messages from the response.""" + return [content for content in self.contents if isinstance(content, UserInputRequestContents)] + + def __str__(self) -> str: + return self.text + + # region ChatOptions -class ToolMode(metaclass=EnumLike): +class ToolMode(SerializationMixin, metaclass=EnumLike): """Defines if and how tools are used in a chat request.""" # Constants configuration for EnumLike metaclass @@ -2565,41 +2322,6 @@ class ToolMode(metaclass=EnumLike): self.mode = mode self.required_function_name = required_function_name - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "ToolMode": - """Create a ToolMode instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - ToolMode instance created from the dictionary. - """ - return cls(**data) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the ToolMode instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - - Returns: - Dictionary representation of the ToolMode instance. - """ - if exclude is None: - exclude = set() - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - result[key] = value - - return result - @classmethod def REQUIRED(cls, function_name: str | None = None) -> "ToolMode": """Returns a ToolMode that requires the specified function to be called.""" @@ -2632,9 +2354,11 @@ class ToolMode(metaclass=EnumLike): return f"ToolMode(mode={self.mode!r})" -class ChatOptions: +class ChatOptions(SerializationMixin): """Common request settings for AI services.""" + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"_tools"} # Internal field, use .tools property + def __init__( self, *, @@ -2742,46 +2466,6 @@ class ChatOptions: """Set the tools.""" self._tools = self._validate_tools(new_tools) - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "ChatOptions": - """Create a ChatOptions instance from a dictionary.""" - return cls(**data) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the ChatOptions instance to a dictionary.""" - exclude = exclude or set() - result: dict[str, Any] = {} - - for key, value in [ - ("additional_properties", self.additional_properties), - ("model_id", self.model_id), - ("allow_multiple_tool_calls", self.allow_multiple_tool_calls), - ("conversation_id", self.conversation_id), - ("frequency_penalty", self.frequency_penalty), - ("instructions", self.instructions), - ("logit_bias", self.logit_bias), - ("max_tokens", self.max_tokens), - ("metadata", self.metadata), - ("presence_penalty", self.presence_penalty), - ("response_format", self.response_format), - ("seed", self.seed), - ("stop", self.stop), - ("store", self.store), - ("temperature", self.temperature), - ("tool_choice", self.tool_choice), - ("tools", self.tools), - ("top_p", self.top_p), - ("user", self.user), - ]: - if key not in exclude and (not exclude_none or value is not None): - if isinstance(value, ToolMode): - result[key] = value.to_dict() - elif isinstance(value, list) and value and hasattr(value[0], "to_dict"): - result[key] = [item.to_dict() if hasattr(item, "to_dict") else item for item in value] # type: ignore[assignment] - else: - result[key] = value - return result - @classmethod def _validate_tools( cls, @@ -2820,7 +2504,7 @@ class ChatOptions: case _: raise ValidationError(f"Invalid tool choice: {tool_choice}") if isinstance(tool_choice, (dict, Mapping)): - return ToolMode.from_dict(tool_choice) + return ToolMode.from_dict(tool_choice) # type: ignore return tool_choice def to_provider_settings(self, by_alias: bool = True, exclude: set[str] | None = None) -> dict[str, Any]: @@ -2833,7 +2517,7 @@ class ChatOptions: Returns: Dictionary of settings for provider. """ - default_exclude = {"additional_properties"} + default_exclude = {"additional_properties", "type"} # 'type' is for serialization, not API calls # No tool choice if no tools are defined if self.tools is None or len(self.tools) == 0: default_exclude.add("tool_choice") @@ -2901,307 +2585,3 @@ class ChatOptions: if tool not in combined.tools: combined.tools.append(tool) return combined - - -# region AgentRunResponse - - -class AgentRunResponse: - """Represents the response to an Agent run request. - - Provides one or more response messages and metadata about the response. - A typical response will contain a single message, but may contain multiple - messages in scenarios involving function calls, RAG retrievals, or complex logic. - """ - - def __init__( - self, - messages: ChatMessage | list[ChatMessage] | None = None, - response_id: str | None = None, - created_at: CreatedAtT | None = None, - usage_details: UsageDetails | None = None, - value: Any | None = None, - raw_representation: Any | None = None, - additional_properties: dict[str, Any] | None = None, - **kwargs: Any, - ) -> None: - """Initialize an AgentRunResponse. - - Attributes: - messages: The list of chat messages in the response. - response_id: The ID of the chat response. - created_at: A timestamp for the chat response. - usage_details: The usage details for the chat response. - value: The structured output of the agent run response, if applicable. - additional_properties: Any additional properties associated with the chat response. - raw_representation: The raw representation of the chat response from an underlying implementation. - **kwargs: Additional properties to set on the response. - """ - processed_messages: list[ChatMessage] = [] - if messages is not None: - if isinstance(messages, ChatMessage): - processed_messages.append(messages) - elif isinstance(messages, list): - processed_messages.extend(messages) - - self.messages = processed_messages - self.response_id = response_id - self.created_at = created_at - self.usage_details = usage_details - self.value = value - self.additional_properties = additional_properties - self.raw_representation = raw_representation - - # Handle any additional kwargs - for key, value in kwargs.items(): - setattr(self, key, value) - - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "AgentRunResponse": - """Create an AgentRunResponse instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - AgentRunResponse instance created from the dictionary. - """ - data_copy = dict(data).copy() - - # Handle messages - convert from list of dicts to list of ChatMessage objects if needed - if messages := data_copy.get("messages"): - parsed_messages: list[ChatMessage] = [] - for message_data in messages: - if isinstance(message_data, ChatMessage): - parsed_messages.append(message_data) - elif isinstance(message_data, dict): - parsed_messages.append(ChatMessage.from_dict(message_data)) - else: - logger.warning(f"Unknown message content: {message_data}") - data_copy["messages"] = parsed_messages - - # Handle usage_details - convert from dict to UsageDetails if needed - if "usage_details" in data_copy and isinstance(data_copy["usage_details"], dict): - data_copy["usage_details"] = UsageDetails.from_dict(data_copy["usage_details"]) - - return cls(**data_copy) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the AgentRunResponse instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - 'raw_representation' is always excluded as per original Pydantic config. - - Returns: - Dictionary representation of the AgentRunResponse instance. - """ - if exclude is None: - exclude = set() - - # Always exclude raw_representation as it was marked with exclude=True in Pydantic - exclude = exclude | {"raw_representation"} - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - - # Handle messages conversion to dict format - if key == "messages" and value is not None: - messages_list = [] - for message in value: - if hasattr(message, "to_dict"): - messages_list.append(message.to_dict(exclude_none=exclude_none)) - else: - messages_list.append(message) - result[key] = messages_list - # Handle usage_details conversion to dict format - elif key == "usage_details" and value is not None: - if hasattr(value, "to_dict"): - result[key] = value.to_dict(exclude_none=exclude_none) - else: - result[key] = value - else: - result[key] = value - - return result - - @property - def text(self) -> str: - """Get the concatenated text of all messages.""" - return "".join(msg.text for msg in self.messages) if self.messages else "" - - @property - def user_input_requests(self) -> list[UserInputRequestContents]: - """Get all BaseUserInputRequest messages from the response.""" - return [ - content - for msg in self.messages - for content in msg.contents - if isinstance(content, UserInputRequestContents) - ] - - @classmethod - def from_agent_run_response_updates( - cls: type[TAgentRunResponse], - updates: Sequence["AgentRunResponseUpdate"], - *, - output_format_type: type[BaseModel] | None = None, - ) -> TAgentRunResponse: - """Joins multiple updates into a single AgentRunResponse.""" - msg = cls(messages=[]) - for update in updates: - _process_update(msg, update) - _finalize_response(msg) - if output_format_type: - msg.try_parse_value(output_format_type) - return msg - - @classmethod - async def from_agent_response_generator( - cls: type[TAgentRunResponse], - updates: AsyncIterable["AgentRunResponseUpdate"], - *, - output_format_type: type[BaseModel] | None = None, - ) -> TAgentRunResponse: - """Joins multiple updates into a single AgentRunResponse.""" - msg = cls(messages=[]) - async for update in updates: - _process_update(msg, update) - _finalize_response(msg) - if output_format_type: - msg.try_parse_value(output_format_type) - return msg - - def __str__(self) -> str: - return self.text - - def try_parse_value(self, output_format_type: type[BaseModel]) -> None: - """If there is a value, does nothing, otherwise tries to parse the text into the value.""" - if self.value is None: - try: - self.value = output_format_type.model_validate_json(self.text) # type: ignore[reportUnknownMemberType] - except ValidationError as ex: - logger.debug("Failed to parse value from agent run response text: %s", ex) - - -# region AgentRunResponseUpdate - - -class AgentRunResponseUpdate: - """Represents a single streaming response chunk from an Agent.""" - - def __init__( - self, - *, - contents: list[Contents] | None = None, - text: TextContent | str | None = None, - role: Role | None = None, - author_name: str | None = None, - response_id: str | None = None, - message_id: str | None = None, - created_at: CreatedAtT | None = None, - additional_properties: dict[str, Any] | None = None, - raw_representation: Any | None = None, - ) -> None: - """Initialize an AgentRunResponseUpdate.""" - if contents is None: - contents = [] - if text is not None: - if isinstance(text, str): - text = TextContent(text=text) - contents.append(text) - - self.contents = contents - self.role = role - self.author_name = author_name - self.response_id = response_id - self.message_id = message_id - self.created_at = created_at - self.additional_properties = additional_properties - self.raw_representation = raw_representation - - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> "AgentRunResponseUpdate": - """Create an AgentRunResponseUpdate instance from a dictionary. - - Args: - data: Dictionary containing the data to create the instance from. - - Returns: - AgentRunResponseUpdate instance created from the dictionary. - """ - data_copy = dict(data).copy() - - # Handle contents - convert from list of dicts to list of Contents objects if needed - if contents := data_copy.get("contents"): - data_copy["contents"] = _parse_content_list(contents) - # Handle role - convert from dict to Role if needed - if role := data.get("role"): - data_copy["role"] = Role.from_dict(role) if isinstance(role, dict) else Role(value=role) - return cls(**data_copy) - - def to_dict(self, *, exclude_none: bool = False, exclude: set[str] | None = None) -> dict[str, Any]: - """Convert the AgentRunResponseUpdate instance to a dictionary. - - Args: - exclude_none: Whether to exclude None values from the output. - exclude: Set of field names to exclude from the output. - 'raw_representation' is always excluded as per original Pydantic config. - - Returns: - Dictionary representation of the AgentRunResponseUpdate instance. - """ - if exclude is None: - exclude = set() - - # Always exclude raw_representation as it was marked with exclude=True in Pydantic - exclude = exclude | {"raw_representation"} - - result: dict[str, Any] = {} - for key, value in self.__dict__.items(): - if key in exclude: - continue - if exclude_none and value is None: - continue - - # Handle contents conversion to dict format - if key == "contents" and value is not None: - contents_list = [] - for content in value: - if hasattr(content, "to_dict"): - contents_list.append(content.to_dict(exclude_none=exclude_none)) - else: - contents_list.append(content) - result[key] = contents_list - # Handle role conversion to dict format - elif key == "role" and value is not None: - if hasattr(value, "to_dict"): - result[key] = value.to_dict(exclude_none=exclude_none) - else: - result[key] = value - else: - result[key] = value - - return result - - @property - def text(self) -> str: - """Get the concatenated text of all TextContent objects in contents.""" - return ( - "".join(content.text for content in self.contents if isinstance(content, TextContent)) - if self.contents - else "" - ) - - @property - def user_input_requests(self) -> list[UserInputRequestContents]: - """Get all BaseUserInputRequest messages from the response.""" - return [content for content in self.contents if isinstance(content, UserInputRequestContents)] - - def __str__(self) -> str: - return self.text diff --git a/python/packages/core/agent_framework/azure/_assistants_client.py b/python/packages/core/agent_framework/azure/_assistants_client.py index 3668ae58b6..370b3e3c89 100644 --- a/python/packages/core/agent_framework/azure/_assistants_client.py +++ b/python/packages/core/agent_framework/azure/_assistants_client.py @@ -127,9 +127,10 @@ class AzureOpenAIAssistantsClient(OpenAIAssistantsClient): async_client = AsyncAzureOpenAI(**client_params) super().__init__( - ai_model_id=azure_openai_settings.chat_deployment_name, + model_id=azure_openai_settings.chat_deployment_name, assistant_id=assistant_id, assistant_name=assistant_name, thread_id=thread_id, async_client=async_client, # type: ignore[reportArgumentType] + default_headers=default_headers, ) diff --git a/python/packages/core/agent_framework/azure/_chat_client.py b/python/packages/core/agent_framework/azure/_chat_client.py index 5bdfa707ec..6f48082b25 100644 --- a/python/packages/core/agent_framework/azure/_chat_client.py +++ b/python/packages/core/agent_framework/azure/_chat_client.py @@ -48,6 +48,7 @@ class AzureOpenAIChatClient(AzureOpenAIConfigMixin, OpenAIBaseChatClient): def __init__( self, + *, api_key: str | None = None, deployment_name: str | None = None, endpoint: str | None = None, @@ -62,6 +63,7 @@ class AzureOpenAIChatClient(AzureOpenAIConfigMixin, OpenAIBaseChatClient): env_file_path: str | None = None, env_file_encoding: str | None = None, instruction_role: str | None = None, + **kwargs: Any, ) -> None: """Initialize an AzureChatCompletion service. @@ -87,6 +89,7 @@ class AzureOpenAIChatClient(AzureOpenAIConfigMixin, OpenAIBaseChatClient): env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'. instruction_role: The role to use for 'instruction' messages, for example, summarization prompts could use `developer` or `system`. (Optional) + kwargs: Other keyword parameters. """ try: # Filter out any None values from the arguments @@ -123,27 +126,7 @@ class AzureOpenAIChatClient(AzureOpenAIConfigMixin, OpenAIBaseChatClient): default_headers=default_headers, client=async_client, instruction_role=instruction_role, - ) - - @classmethod - def from_dict(cls: type[TAzureOpenAIChatClient], settings: dict[str, Any]) -> TAzureOpenAIChatClient: - """Initialize an Azure OpenAI service from a dictionary of settings. - - Args: - settings: A dictionary of settings for the service. - should contain keys: service_id, and optionally: - ad_auth, ad_token_provider, default_headers - """ - return cls( - api_key=settings.get("api_key"), - deployment_name=settings.get("deployment_name"), - endpoint=settings.get("endpoint"), - base_url=settings.get("base_url"), - api_version=settings.get("api_version"), - ad_token=settings.get("ad_token"), - ad_token_provider=settings.get("ad_token_provider"), - default_headers=settings.get("default_headers"), - env_file_path=settings.get("env_file_path"), + **kwargs, ) @override diff --git a/python/packages/core/agent_framework/azure/_responses_client.py b/python/packages/core/agent_framework/azure/_responses_client.py index d5ef19f927..4c688cd79a 100644 --- a/python/packages/core/agent_framework/azure/_responses_client.py +++ b/python/packages/core/agent_framework/azure/_responses_client.py @@ -43,6 +43,7 @@ class AzureOpenAIResponsesClient(AzureOpenAIConfigMixin, OpenAIBaseResponsesClie env_file_path: str | None = None, env_file_encoding: str | None = None, instruction_role: str | None = None, + **kwargs: Any, ) -> None: """Initialize an AzureResponses service. @@ -68,7 +69,10 @@ class AzureOpenAIResponsesClient(AzureOpenAIConfigMixin, OpenAIBaseResponsesClie env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'. instruction_role: The role to use for 'instruction' messages, for example, summarization prompts could use `developer` or `system`. (Optional) + kwargs: Additional keyword arguments. """ + if model_id := kwargs.pop("model_id", None) and not deployment_name: + deployment_name = str(model_id) try: azure_openai_settings = AzureOpenAISettings( # pydantic settings will see if there is a value, if not, will try the env var or .env file @@ -115,22 +119,3 @@ class AzureOpenAIResponsesClient(AzureOpenAIConfigMixin, OpenAIBaseResponsesClie client=async_client, instruction_role=instruction_role, ) - - @classmethod - def from_dict(cls: type[TAzureOpenAIResponsesClient], settings: dict[str, Any]) -> TAzureOpenAIResponsesClient: - """Initialize an Open AI service from a dictionary of settings. - - Args: - settings: A dictionary of settings for the service. - """ - return cls( - api_key=settings.get("api_key"), - deployment_name=settings.get("deployment_name"), - endpoint=settings.get("endpoint"), - base_url=settings.get("base_url"), - api_version=settings.get("api_version"), - ad_token=settings.get("ad_token"), - ad_token_provider=settings.get("ad_token_provider"), - default_headers=settings.get("default_headers"), - env_file_path=settings.get("env_file_path"), - ) diff --git a/python/packages/core/agent_framework/azure/_shared.py b/python/packages/core/agent_framework/azure/_shared.py index 2b80efd711..bc455600e7 100644 --- a/python/packages/core/agent_framework/azure/_shared.py +++ b/python/packages/core/agent_framework/azure/_shared.py @@ -8,10 +8,10 @@ from typing import Any, ClassVar, Final from azure.core.credentials import TokenCredential from openai.lib.azure import AsyncAzureOpenAI -from pydantic import ConfigDict, SecretStr, model_validator, validate_call +from pydantic import SecretStr, model_validator from .._pydantic import AFBaseSettings, HTTPsUrl -from .._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent +from .._telemetry import APP_INFO, prepend_agent_framework_to_user_agent from ..exceptions import ServiceInitializationError from ..openai._shared import OpenAIBase from ._entra_id_authentication import get_entra_auth_token @@ -124,9 +124,9 @@ class AzureOpenAISettings(AFBaseSettings): class AzureOpenAIConfigMixin(OpenAIBase): """Internal class for configuring a connection to an Azure OpenAI service.""" - OTEL_PROVIDER_NAME: ClassVar[str] = "azure_openai" # type: ignore[reportIncompatibleVariableOverride, misc] + OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.openai" + # Note: INJECTABLE = {"client"} is inherited from OpenAIBase - @validate_call(config=ConfigDict(arbitrary_types_allowed=True)) def __init__( self, deployment_name: str, @@ -207,36 +207,20 @@ class AzureOpenAIConfigMixin(OpenAIBase): args["websocket_base_url"] = kwargs.pop("websocket_base_url") client = AsyncAzureOpenAI(**args) - args = { - "ai_model_id": deployment_name, - "client": client, - } - if instruction_role: - args["instruction_role"] = instruction_role - super().__init__(**args, **kwargs) - def to_dict(self) -> dict[str, Any]: - """Convert the configuration to a dictionary.""" - client_settings = { - "base_url": str(self.client.base_url), - "api_version": self.client._custom_query["api-version"], # type: ignore - "api_key": self.client.api_key, - "ad_token": getattr(self.client, "_azure_ad_token", None), - "ad_token_provider": getattr(self.client, "_azure_ad_token_provider", None), - "default_headers": {k: v for k, v in self.client.default_headers.items() if k != USER_AGENT_KEY}, - } - base = self.model_dump( - exclude={ - "prompt_tokens", - "completion_tokens", - "total_tokens", - "api_type", - "org_id", - "service_id", - "client", - }, - by_alias=True, - exclude_none=True, - ) - base.update(client_settings) - return base + # Store configuration as instance attributes for serialization + self.endpoint = str(endpoint) + self.base_url = str(base_url) + self.api_version = api_version + self.deployment_name = deployment_name + self.instruction_role = instruction_role + # Store default_headers but filter out USER_AGENT_KEY for serialization + if default_headers: + from .._telemetry import USER_AGENT_KEY + + def_headers = {k: v for k, v in default_headers.items() if k != USER_AGENT_KEY} + else: + def_headers = None + self.default_headers = def_headers + + super().__init__(model_id=deployment_name, client=client) diff --git a/python/packages/core/agent_framework/openai/_assistants_client.py b/python/packages/core/agent_framework/openai/_assistants_client.py index d31b09baa3..a7a44279cc 100644 --- a/python/packages/core/agent_framework/openai/_assistants_client.py +++ b/python/packages/core/agent_framework/openai/_assistants_client.py @@ -18,7 +18,7 @@ from openai.types.beta.threads import ( from openai.types.beta.threads.run_create_params import AdditionalMessage 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 pydantic import ValidationError from .._clients import BaseChatClient from .._middleware import use_chat_middleware @@ -57,28 +57,25 @@ __all__ = ["OpenAIAssistantsClient"] class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient): """OpenAI Assistants client.""" - assistant_id: str | None = Field(default=None) - assistant_name: str | None = Field(default=None) - thread_id: str | None = Field(default=None) - _should_delete_assistant: bool = PrivateAttr(default=False) # Track whether we should delete the assistant - def __init__( self, - ai_model_id: str | None = None, + model_id: str | None = None, assistant_id: str | None = None, assistant_name: str | None = None, thread_id: str | None = None, api_key: str | None = None, org_id: str | None = None, + base_url: str | None = None, default_headers: Mapping[str, str] | None = None, async_client: AsyncOpenAI | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, + **kwargs: Any, ) -> None: """Initialize an OpenAI Assistants client. Args: - ai_model_id: OpenAI model name, see + model_id: OpenAI model name, see https://platform.openai.com/docs/models assistant_id: The ID of an OpenAI assistant to use. If not provided, a new assistant will be created (and deleted after the request). @@ -90,18 +87,21 @@ class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient): the env vars or .env file value. org_id: The optional org ID to use. If provided will override, the env vars or .env file value. + base_url: The optional base URL to use. If provided will override, default_headers: The default headers mapping of string keys to string values for HTTP requests. (Optional) async_client: An existing client to use. (Optional) env_file_path: Use the environment settings file as a fallback to environment variables. (Optional) env_file_encoding: The encoding of the environment settings file. (Optional) + kwargs: Other keyword parameters. """ try: openai_settings = OpenAISettings( - api_key=SecretStr(api_key) if api_key else None, + api_key=api_key, # type: ignore[reportArgumentType] + base_url=base_url, org_id=org_id, - chat_model_id=ai_model_id, + chat_model_id=model_id, env_file_path=env_file_path, env_file_encoding=env_file_encoding, ) @@ -119,15 +119,17 @@ class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient): ) super().__init__( - ai_model_id=openai_settings.chat_model_id, - assistant_id=assistant_id, # type: ignore[reportCallIssue] - assistant_name=assistant_name, # type: ignore[reportCallIssue] - thread_id=thread_id, # type: ignore[reportCallIssue] + model_id=openai_settings.chat_model_id, api_key=openai_settings.api_key.get_secret_value() if openai_settings.api_key else None, org_id=openai_settings.org_id, default_headers=default_headers, client=async_client, + base_url=openai_settings.base_url, ) + self.assistant_id: str | None = assistant_id + self.assistant_name: str | None = assistant_name + self.thread_id: str | None = thread_id + self._should_delete_assistant: bool = False async def __aenter__(self) -> "Self": """Async context manager entry.""" @@ -141,8 +143,8 @@ class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient): """Clean up any assistants we created.""" if self._should_delete_assistant and self.assistant_id is not None: await self.client.beta.assistants.delete(self.assistant_id) - self.assistant_id = None - self._should_delete_assistant = False + object.__setattr__(self, "assistant_id", None) + object.__setattr__(self, "_should_delete_assistant", False) async def _inner_get_response( self, @@ -194,10 +196,7 @@ class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient): """ # If no assistant is provided, create a temporary assistant if self.assistant_id is None: - created_assistant = await self.client.beta.assistants.create( - name=self.assistant_name, model=self.ai_model_id - ) - + created_assistant = await self.client.beta.assistants.create(name=self.assistant_name, model=self.model_id) self.assistant_id = created_assistant.id self._should_delete_assistant = True @@ -495,4 +494,4 @@ class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient): # This is a no-op in the base class, but can be overridden by subclasses # to update the agent name in the client. if agent_name and not self.assistant_name: - self.assistant_name = agent_name + object.__setattr__(self, "assistant_name", agent_name) diff --git a/python/packages/core/agent_framework/openai/_chat_client.py b/python/packages/core/agent_framework/openai/_chat_client.py index 0199ec4ecd..8f772fb83c 100644 --- a/python/packages/core/agent_framework/openai/_chat_client.py +++ b/python/packages/core/agent_framework/openai/_chat_client.py @@ -14,7 +14,7 @@ from openai.types.chat.chat_completion import ChatCompletion, Choice from openai.types.chat.chat_completion_chunk import ChatCompletionChunk 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 pydantic import BaseModel, ValidationError from .._clients import BaseChatClient from .._logging import get_logger @@ -172,7 +172,7 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient): options_dict.pop("tool_choice", None) if "model" not in options_dict: - options_dict["model"] = self.ai_model_id + options_dict["model"] = self.model_id if ( chat_options.response_format and isinstance(chat_options.response_format, type) @@ -464,7 +464,7 @@ class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient): def __init__( self, - ai_model_id: str | None = None, + model_id: str | None = None, api_key: str | None = None, org_id: str | None = None, default_headers: Mapping[str, str] | None = None, @@ -477,7 +477,7 @@ class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient): """Initialize an OpenAIChatCompletion service. Args: - ai_model_id: OpenAI model name, see + model_id: OpenAI model name, see https://platform.openai.com/docs/models api_key: The optional API key to use. If provided will override, the env vars or .env file value. @@ -497,10 +497,10 @@ class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient): """ try: openai_settings = OpenAISettings( - api_key=SecretStr(api_key) if api_key else None, + api_key=api_key, # type: ignore[reportArgumentType] base_url=base_url, org_id=org_id, - chat_model_id=ai_model_id, + chat_model_id=model_id, env_file_path=env_file_path, env_file_encoding=env_file_encoding, ) @@ -514,11 +514,11 @@ class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient): if not openai_settings.chat_model_id: raise ServiceInitializationError( "OpenAI model ID is required. " - "Set via 'ai_model_id' parameter or 'OPENAI_CHAT_MODEL_ID' environment variable." + "Set via 'model_id' parameter or 'OPENAI_CHAT_MODEL_ID' environment variable." ) super().__init__( - ai_model_id=openai_settings.chat_model_id, + model_id=openai_settings.chat_model_id, api_key=openai_settings.api_key.get_secret_value() if openai_settings.api_key else None, base_url=openai_settings.base_url if openai_settings.base_url else None, org_id=openai_settings.org_id, @@ -526,15 +526,3 @@ class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient): client=async_client, instruction_role=instruction_role, ) - - @classmethod - def from_dict(cls: type[TOpenAIChatClient], settings: dict[str, Any]) -> TOpenAIChatClient: - """Initialize an Open AI Chat Client from a dictionary of settings. - - Args: - settings: A dictionary of settings for the service. - """ - return cls(**settings) - - -# endregion diff --git a/python/packages/core/agent_framework/openai/_responses_client.py b/python/packages/core/agent_framework/openai/_responses_client.py index 261e35dbeb..f5560ea767 100644 --- a/python/packages/core/agent_framework/openai/_responses_client.py +++ b/python/packages/core/agent_framework/openai/_responses_client.py @@ -22,7 +22,7 @@ from openai.types.responses.tool_param import ( ) from openai.types.responses.web_search_tool_param import UserLocation as WebSearchUserLocation from openai.types.responses.web_search_tool_param import WebSearchToolParam -from pydantic import BaseModel, SecretStr, ValidationError +from pydantic import BaseModel, ValidationError from .._clients import BaseChatClient from .._logging import get_logger @@ -333,7 +333,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): if chat_options.conversation_id: options_dict["previous_response_id"] = chat_options.conversation_id if "model" not in options_dict: - options_dict["model"] = self.ai_model_id + options_dict["model"] = self.model_id return options_dict def _prepare_chat_messages_for_request(self, chat_messages: Sequence[ChatMessage]) -> list[dict[str, Any]]: @@ -731,7 +731,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): metadata: dict[str, Any] = {} contents: list[Contents] = [] conversation_id: str | None = None - model = self.ai_model_id + model = self.model_id # TODO(peterychang): Add support for other content types match event.type: # types: @@ -942,24 +942,27 @@ class OpenAIResponsesClient(OpenAIConfigMixin, OpenAIBaseResponsesClient): def __init__( self, - ai_model_id: str | None = None, + model_id: str | None = None, api_key: str | None = None, org_id: str | None = None, + base_url: str | None = None, default_headers: Mapping[str, str] | None = None, async_client: AsyncOpenAI | None = None, instruction_role: str | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, + **kwargs: Any, ) -> None: """Initialize an OpenAIChatCompletion service. Args: - ai_model_id: OpenAI model name, see + model_id: OpenAI model name, see https://platform.openai.com/docs/models api_key: The optional API key to use. If provided will override, the env vars or .env file value. org_id: The optional org ID to use. If provided will override, the env vars or .env file value. + base_url: The optional base URL to use. If provided will override, default_headers: The default headers mapping of string keys to string values for HTTP requests. (Optional) async_client: An existing client to use. (Optional) @@ -968,12 +971,14 @@ class OpenAIResponsesClient(OpenAIConfigMixin, OpenAIBaseResponsesClient): env_file_path: Use the environment settings file as a fallback to environment variables. (Optional) env_file_encoding: The encoding of the environment settings file. (Optional) + kwargs: Other keyword parameters. """ try: openai_settings = OpenAISettings( - api_key=SecretStr(api_key) if api_key else None, + api_key=api_key, # type: ignore[reportArgumentType] org_id=org_id, - responses_model_id=ai_model_id, + base_url=base_url, + responses_model_id=model_id, env_file_path=env_file_path, env_file_encoding=env_file_encoding, ) @@ -987,26 +992,15 @@ class OpenAIResponsesClient(OpenAIConfigMixin, OpenAIBaseResponsesClient): if not openai_settings.responses_model_id: raise ServiceInitializationError( "OpenAI model ID is required. " - "Set via 'ai_model_id' parameter or 'OPENAI_RESPONSES_MODEL_ID' environment variable." + "Set via 'model_id' parameter or 'OPENAI_RESPONSES_MODEL_ID' environment variable." ) super().__init__( - ai_model_id=openai_settings.responses_model_id, + model_id=openai_settings.responses_model_id, api_key=openai_settings.api_key.get_secret_value() if openai_settings.api_key else None, org_id=openai_settings.org_id, default_headers=default_headers, client=async_client, instruction_role=instruction_role, + base_url=openai_settings.base_url, ) - - @classmethod - def from_dict(cls: type[TOpenAIResponsesClient], settings: dict[str, Any]) -> TOpenAIResponsesClient: - """Initialize an Open AI service from a dictionary of settings. - - Args: - settings: A dictionary of settings for the service. - """ - return cls(**settings) - - -# endregion diff --git a/python/packages/core/agent_framework/openai/_shared.py b/python/packages/core/agent_framework/openai/_shared.py index 9383e68b89..e99de779ee 100644 --- a/python/packages/core/agent_framework/openai/_shared.py +++ b/python/packages/core/agent_framework/openai/_shared.py @@ -4,7 +4,7 @@ import json import logging from collections.abc import Mapping from copy import copy -from typing import Annotated, Any, ClassVar, Union +from typing import Any, ClassVar, Union from openai import ( AsyncOpenAI, @@ -17,11 +17,11 @@ from openai.types.chat import ChatCompletion, ChatCompletionChunk from openai.types.images_response import ImagesResponse from openai.types.responses.response import Response from openai.types.responses.response_stream_event import ResponseStreamEvent -from pydantic import ConfigDict, Field, SecretStr, validate_call -from pydantic.types import StringConstraints +from pydantic import SecretStr from .._logging import get_logger -from .._pydantic import AFBaseModel, AFBaseSettings +from .._pydantic import AFBaseSettings +from .._serialization import SerializationMixin from .._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent from .._types import ChatOptions, Contents from ..exceptions import ServiceInitializationError @@ -106,11 +106,45 @@ class OpenAISettings(AFBaseSettings): responses_model_id: str | None = None -class OpenAIBase(AFBaseModel): +class OpenAIBase(SerializationMixin): """Base class for OpenAI Clients.""" - client: AsyncOpenAI - ai_model_id: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] + INJECTABLE: ClassVar[set[str]] = {"client"} + + def __init__(self, *, client: AsyncOpenAI, model_id: str, **kwargs: Any) -> None: + """Initialize OpenAIBase. + + Args: + client: The AsyncOpenAI client instance. + model_id: The AI model ID to use (non-empty, whitespace stripped). + **kwargs: Additional keyword arguments. + """ + if not model_id or not model_id.strip(): + raise ValueError("model_id must be a non-empty string") + self.client = client + self.model_id = model_id.strip() + + # Call super().__init__() to continue MRO chain (e.g., BaseChatClient) + # Extract known kwargs that belong to other base classes + additional_properties = kwargs.pop("additional_properties", None) + middleware = kwargs.pop("middleware", None) + instruction_role = kwargs.pop("instruction_role", None) + + # Build super().__init__() args + super_kwargs = {} + if additional_properties is not None: + super_kwargs["additional_properties"] = additional_properties + if middleware is not None: + super_kwargs["middleware"] = middleware + + # Call super().__init__() with filtered kwargs + super().__init__(**super_kwargs) + + # Store instruction_role and any remaining kwargs as instance attributes + if instruction_role is not None: + self.instruction_role = instruction_role + for key, value in kwargs.items(): + setattr(self, key, value) class OpenAIConfigMixin(OpenAIBase): @@ -118,11 +152,10 @@ class OpenAIConfigMixin(OpenAIBase): OTEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc] - @validate_call(config=ConfigDict(arbitrary_types_allowed=True)) def __init__( self, - ai_model_id: str = Field(min_length=1), - api_key: str | None = Field(min_length=1), + model_id: str, + api_key: str | None = None, org_id: str | None = None, default_headers: Mapping[str, str] | None = None, client: AsyncOpenAI | None = None, @@ -136,7 +169,7 @@ class OpenAIConfigMixin(OpenAIBase): different types of AI model interactions, like chat or text completion. Args: - ai_model_id: OpenAI model identifier. Must be non-empty. + model_id: OpenAI model identifier. Must be non-empty. Default to a preset value. api_key: OpenAI API key for authentication. Must be non-empty. (Optional) @@ -167,32 +200,25 @@ class OpenAIConfigMixin(OpenAIBase): if base_url: args["base_url"] = base_url client = AsyncOpenAI(**args) + + # Store configuration as instance attributes for serialization + self.org_id = org_id + self.base_url = str(base_url) + # Store default_headers but filter out USER_AGENT_KEY for serialization + if default_headers: + self.default_headers: dict[str, Any] | None = { + k: v for k, v in default_headers.items() if k != USER_AGENT_KEY + } + else: + self.default_headers = None + args = { - "ai_model_id": ai_model_id, + "model_id": model_id, "client": client, } if instruction_role: args["instruction_role"] = instruction_role - super().__init__(**args, **kwargs) - def to_dict(self) -> dict[str, Any]: - """Create a dict of the service settings.""" - client_settings = { - "api_key": self.client.api_key, - "default_headers": {k: v for k, v in self.client.default_headers.items() if k != USER_AGENT_KEY}, - } - if self.client.organization: - client_settings["org_id"] = self.client.organization - base = self.model_dump( - exclude={ - "prompt_tokens", - "completion_tokens", - "total_tokens", - "api_type", - "client", - }, - by_alias=True, - exclude_none=True, - ) - base.update(client_settings) - return base + # Ensure additional_properties and middleware are passed through kwargs to BaseChatClient + # These are consumed by BaseChatClient.__init__ via kwargs + super().__init__(**args, **kwargs) diff --git a/python/packages/core/tests/azure/test_azure_assistants_client.py b/python/packages/core/tests/azure/test_azure_assistants_client.py index bdab12970f..758be68d3b 100644 --- a/python/packages/core/tests/azure/test_azure_assistants_client.py +++ b/python/packages/core/tests/azure/test_azure_assistants_client.py @@ -40,17 +40,20 @@ def create_test_azure_assistants_client( thread_id: str | None = None, should_delete_assistant: bool = False, ) -> AzureOpenAIAssistantsClient: - """Helper function to create AzureOpenAIAssistantsClient instances for testing, bypassing Pydantic validation.""" - return AzureOpenAIAssistantsClient.model_construct( - ai_model_id=deployment_name or "test_chat_deployment", + """Helper function to create AzureOpenAIAssistantsClient instances for testing.""" + client = AzureOpenAIAssistantsClient( + deployment_name=deployment_name or "test_chat_deployment", assistant_id=assistant_id, assistant_name=assistant_name, thread_id=thread_id, api_key="test-api-key", endpoint="https://test-endpoint.com", - client=mock_async_azure_openai, - _should_delete_assistant=should_delete_assistant, + async_client=mock_async_azure_openai, ) + # Set the _should_delete_assistant flag directly if needed + if should_delete_assistant: + object.__setattr__(client, "_should_delete_assistant", True) + return client @pytest.fixture @@ -88,7 +91,7 @@ def test_azure_assistants_client_init_with_client(mock_async_azure_openai: Magic ) assert chat_client.client is mock_async_azure_openai - assert chat_client.ai_model_id == "test_chat_deployment" + assert chat_client.model_id == "test_chat_deployment" assert chat_client.assistant_id == "existing-assistant-id" assert chat_client.thread_id == "test-thread-id" assert not chat_client._should_delete_assistant # type: ignore @@ -100,19 +103,16 @@ def test_azure_assistants_client_init_auto_create_client( mock_async_azure_openai: MagicMock, ) -> None: """Test AzureOpenAIAssistantsClient initialization with auto-created client.""" - chat_client = AzureOpenAIAssistantsClient.model_construct( - ai_model_id=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"], - assistant_id=None, + chat_client = AzureOpenAIAssistantsClient( + deployment_name=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"], assistant_name="TestAssistant", - thread_id=None, api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"], endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"], - client=mock_async_azure_openai, - _should_delete_assistant=False, + async_client=mock_async_azure_openai, ) assert chat_client.client is mock_async_azure_openai - assert chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"] + assert chat_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"] assert chat_client.assistant_id is None assert chat_client.assistant_name == "TestAssistant" assert not chat_client._should_delete_assistant # type: ignore @@ -145,7 +145,7 @@ def test_azure_assistants_client_init_with_default_headers(azure_openai_unit_tes default_headers=default_headers, ) - assert chat_client.ai_model_id == "test_chat_deployment" + assert chat_client.model_id == "test_chat_deployment" assert isinstance(chat_client, ChatClientProtocol) # Assert that the default header we added is present in the client's default headers @@ -241,11 +241,10 @@ def test_azure_assistants_client_serialize(azure_openai_unit_test_env: dict[str, dumped_settings = chat_client.to_dict() - assert dumped_settings["ai_model_id"] == "test_chat_deployment" + assert dumped_settings["model_id"] == "test_chat_deployment" assert dumped_settings["assistant_id"] == "test-assistant-id" assert dumped_settings["assistant_name"] == "TestAssistant" assert dumped_settings["thread_id"] == "test-thread-id" - assert dumped_settings["api_key"] == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"] # Assert that the default header we added is present in the dumped_settings default headers for key, value in default_headers.items(): @@ -262,6 +261,7 @@ def get_weather( return f"The weather in {location} is sunny with a high of 25°C." +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_client_get_response() -> None: """Test Azure Assistants Client response.""" @@ -286,6 +286,7 @@ async def test_azure_assistants_client_get_response() -> None: assert any(word in response.text.lower() for word in ["sunny", "25", "weather", "seattle"]) +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_client_get_response_tools() -> None: """Test Azure Assistants Client response with tools.""" @@ -307,6 +308,7 @@ async def test_azure_assistants_client_get_response_tools() -> None: assert any(word in response.text.lower() for word in ["sunny", "25", "weather"]) +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_client_streaming() -> None: """Test Azure Assistants Client streaming response.""" @@ -337,6 +339,7 @@ async def test_azure_assistants_client_streaming() -> None: assert any(word in full_message.lower() for word in ["sunny", "25", "weather", "seattle"]) +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_client_streaming_tools() -> None: """Test Azure Assistants Client streaming response with tools.""" @@ -363,6 +366,7 @@ async def test_azure_assistants_client_streaming_tools() -> None: assert any(word in full_message.lower() for word in ["sunny", "25", "weather"]) +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_client_with_existing_assistant() -> None: """Test Azure Assistants Client with existing assistant ID.""" @@ -390,6 +394,7 @@ async def test_azure_assistants_client_with_existing_assistant() -> None: assert len(response.text) > 0 +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_agent_basic_run(): """Test ChatAgent basic run functionality with AzureOpenAIAssistantsClient.""" @@ -406,6 +411,7 @@ async def test_azure_assistants_agent_basic_run(): assert "Hello World" in response.text +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_agent_basic_run_streaming(): """Test ChatAgent basic streaming functionality with AzureOpenAIAssistantsClient.""" @@ -425,6 +431,7 @@ async def test_azure_assistants_agent_basic_run_streaming(): assert "streaming response test" in full_message.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_agent_thread_persistence(): """Test ChatAgent thread persistence across runs with AzureOpenAIAssistantsClient.""" @@ -453,6 +460,7 @@ async def test_azure_assistants_agent_thread_persistence(): assert thread.service_thread_id is not None +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_agent_existing_thread_id(): """Test ChatAgent with existing thread ID to continue conversations across agent instances.""" @@ -497,6 +505,7 @@ async def test_azure_assistants_agent_existing_thread_id(): assert "paris" in response2.text.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_agent_code_interpreter(): """Test ChatAgent with code interpreter through AzureOpenAIAssistantsClient.""" @@ -516,6 +525,7 @@ async def test_azure_assistants_agent_code_interpreter(): assert "120" in response.text or "factorial" in response.text.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_assistants_client_agent_level_tool_persistence(): """Test that agent-level tools persist across multiple runs with Azure Assistants Client.""" diff --git a/python/packages/core/tests/azure/test_azure_chat_client.py b/python/packages/core/tests/azure/test_azure_chat_client.py index e7c7cb5046..aad231ac98 100644 --- a/python/packages/core/tests/azure/test_azure_chat_client.py +++ b/python/packages/core/tests/azure/test_azure_chat_client.py @@ -53,7 +53,7 @@ def test_init(azure_openai_unit_test_env: dict[str, str]) -> None: assert azure_chat_client.client is not None assert isinstance(azure_chat_client.client, AsyncAzureOpenAI) - assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"] + assert azure_chat_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"] assert isinstance(azure_chat_client, BaseChatClient) @@ -76,7 +76,7 @@ def test_init_base_url(azure_openai_unit_test_env: dict[str, str]) -> None: assert azure_chat_client.client is not None assert isinstance(azure_chat_client.client, AsyncAzureOpenAI) - assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"] + assert azure_chat_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"] assert isinstance(azure_chat_client, BaseChatClient) for key, value in default_headers.items(): assert key in azure_chat_client.client.default_headers @@ -89,7 +89,7 @@ def test_init_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None: assert azure_chat_client.client is not None assert isinstance(azure_chat_client.client, AsyncAzureOpenAI) - assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"] + assert azure_chat_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"] assert isinstance(azure_chat_client, BaseChatClient) @@ -130,10 +130,11 @@ def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None: azure_chat_client = AzureOpenAIChatClient.from_dict(settings) dumped_settings = azure_chat_client.to_dict() - assert dumped_settings["ai_model_id"] == settings["deployment_name"] - assert str(settings["deployment_name"]) in str(dumped_settings["base_url"]) - assert settings["api_key"] == dumped_settings["api_key"] + assert dumped_settings["model_id"] == settings["deployment_name"] + assert str(settings["endpoint"]) in str(dumped_settings["endpoint"]) + assert str(settings["deployment_name"]) == str(dumped_settings["deployment_name"]) assert settings["api_version"] == dumped_settings["api_version"] + assert "api_key" not in dumped_settings # Assert that the default header we added is present in the dumped_settings default headers for key, value in default_headers.items(): @@ -608,6 +609,7 @@ def get_weather(location: str) -> str: return f"The weather in {location} is sunny and 72°F." +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_openai_chat_client_response() -> None: """Test Azure OpenAI chat completion responses.""" @@ -637,6 +639,7 @@ async def test_azure_openai_chat_client_response() -> None: ) +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_openai_chat_client_response_tools() -> None: """Test AzureOpenAI chat completion responses.""" @@ -658,6 +661,7 @@ async def test_azure_openai_chat_client_response_tools() -> None: assert "scientists" in response.text +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_openai_chat_client_streaming() -> None: """Test Azure OpenAI chat completion responses.""" @@ -692,6 +696,7 @@ async def test_azure_openai_chat_client_streaming() -> None: assert "scientists" in full_message +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_openai_chat_client_streaming_tools() -> None: """Test AzureOpenAI chat completion responses.""" @@ -718,6 +723,7 @@ async def test_azure_openai_chat_client_streaming_tools() -> None: assert "scientists" in full_message +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_openai_chat_client_agent_basic_run(): """Test Azure OpenAI chat client agent basic run functionality with AzureOpenAIChatClient.""" @@ -733,6 +739,7 @@ async def test_azure_openai_chat_client_agent_basic_run(): assert "hello world" in response.text.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_openai_chat_client_agent_basic_run_streaming(): """Test Azure OpenAI chat client agent basic streaming functionality with AzureOpenAIChatClient.""" @@ -750,6 +757,7 @@ async def test_azure_openai_chat_client_agent_basic_run_streaming(): assert "streaming response test" in full_text.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_openai_chat_client_agent_thread_persistence(): """Test Azure OpenAI chat client agent thread persistence across runs with AzureOpenAIChatClient.""" @@ -774,6 +782,7 @@ async def test_azure_openai_chat_client_agent_thread_persistence(): assert "alice" in response2.text.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_openai_chat_client_agent_existing_thread(): """Test Azure OpenAI chat client agent with existing thread to continue conversations across agent instances.""" @@ -808,6 +817,7 @@ async def test_azure_openai_chat_client_agent_existing_thread(): assert "alice" in second_response.text.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_chat_client_agent_level_tool_persistence(): """Test that agent-level tools persist across multiple runs with Azure Chat Client.""" diff --git a/python/packages/core/tests/azure/test_azure_responses_client.py b/python/packages/core/tests/azure/test_azure_responses_client.py index 5e51fc8bcf..a495d05837 100644 --- a/python/packages/core/tests/azure/test_azure_responses_client.py +++ b/python/packages/core/tests/azure/test_azure_responses_client.py @@ -76,7 +76,7 @@ def test_init(azure_openai_unit_test_env: dict[str, str]) -> None: # Test successful initialization azure_responses_client = AzureOpenAIResponsesClient() - assert azure_responses_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"] + assert azure_responses_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"] assert isinstance(azure_responses_client, ChatClientProtocol) @@ -86,12 +86,12 @@ def test_init_validation_fail() -> None: AzureOpenAIResponsesClient(api_key="34523", deployment_name={"test": "dict"}) # type: ignore -def test_init_ai_model_id_constructor(azure_openai_unit_test_env: dict[str, str]) -> None: +def test_init_model_id_constructor(azure_openai_unit_test_env: dict[str, str]) -> None: # Test successful initialization - ai_model_id = "test_model_id" - azure_responses_client = AzureOpenAIResponsesClient(deployment_name=ai_model_id) + model_id = "test_model_id" + azure_responses_client = AzureOpenAIResponsesClient(deployment_name=model_id) - assert azure_responses_client.ai_model_id == ai_model_id + assert azure_responses_client.model_id == model_id assert isinstance(azure_responses_client, ChatClientProtocol) @@ -103,7 +103,7 @@ def test_init_with_default_header(azure_openai_unit_test_env: dict[str, str]) -> default_headers=default_headers, ) - assert azure_responses_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"] + assert azure_responses_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"] assert isinstance(azure_responses_client, ChatClientProtocol) # Assert that the default header we added is present in the client's default headers @@ -124,15 +124,15 @@ def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None: default_headers = {"X-Unit-Test": "test-guid"} settings = { - "ai_model_id": azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"], + "deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"], "api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"], "default_headers": default_headers, } azure_responses_client = AzureOpenAIResponsesClient.from_dict(settings) dumped_settings = azure_responses_client.to_dict() - assert dumped_settings["ai_model_id"] == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"] - assert dumped_settings["api_key"] == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"] + assert dumped_settings["deployment_name"] == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"] + assert "api_key" not in dumped_settings # Assert that the default header we added is present in the dumped_settings default headers for key, value in default_headers.items(): assert key in dumped_settings["default_headers"] @@ -141,6 +141,7 @@ def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None: assert "User-Agent" not in dumped_settings["default_headers"] +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_response() -> None: """Test azure responses client responses.""" @@ -184,6 +185,7 @@ async def test_azure_responses_client_response() -> None: assert "sunny" in structured_response.value.weather.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_response_tools() -> None: """Test azure responses client tools.""" @@ -223,6 +225,7 @@ async def test_azure_responses_client_response_tools() -> None: assert "sunny" in structured_response.value.weather.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_streaming() -> None: """Test Azure azure responses client streaming responses.""" @@ -273,6 +276,7 @@ async def test_azure_responses_client_streaming() -> None: assert "sunny" in structured_response.value.weather.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_streaming_tools() -> None: """Test azure responses client streaming tools.""" @@ -320,6 +324,7 @@ async def test_azure_responses_client_streaming_tools() -> None: assert "sunny" in output.weather.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_basic_run(): """Test Azure Responses Client agent basic run functionality with AzureOpenAIResponsesClient.""" @@ -336,6 +341,7 @@ async def test_azure_responses_client_agent_basic_run(): assert "hello world" in response.text.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_basic_run_streaming(): """Test Azure Responses Client agent basic streaming functionality with AzureOpenAIResponsesClient.""" @@ -353,6 +359,7 @@ async def test_azure_responses_client_agent_basic_run_streaming(): assert "streaming response test" in full_text.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_thread_persistence(): """Test Azure Responses Client agent thread persistence across runs with AzureOpenAIResponsesClient.""" @@ -376,6 +383,7 @@ async def test_azure_responses_client_agent_thread_persistence(): assert second_response.text is not None +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_thread_storage_with_store_true(): """Test Azure Responses Client agent with store=True to verify service_thread_id is returned.""" @@ -407,6 +415,7 @@ async def test_azure_responses_client_agent_thread_storage_with_store_true(): assert len(thread.service_thread_id) > 0 +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_existing_thread(): """Test Azure Responses Client agent with existing thread to continue conversations across agent instances.""" @@ -441,6 +450,7 @@ async def test_azure_responses_client_agent_existing_thread(): assert "photography" in second_response.text.lower() +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_hosted_code_interpreter_tool(): """Test Azure Responses Client agent with HostedCodeInterpreterTool through AzureOpenAIResponsesClient.""" @@ -462,6 +472,7 @@ async def test_azure_responses_client_agent_hosted_code_interpreter_tool(): assert contains_relevant_content or len(response.text.strip()) > 10 +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_level_tool_persistence(): """Test that agent-level tools persist across multiple runs with Azure Responses Client.""" @@ -488,6 +499,7 @@ async def test_azure_responses_client_agent_level_tool_persistence(): assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"]) +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_chat_options_run_level() -> None: """Integration test for comprehensive ChatOptions parameter coverage with Azure Response Agent.""" @@ -511,6 +523,7 @@ async def test_azure_responses_client_agent_chat_options_run_level() -> None: assert len(response.text) > 0 +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_chat_options_agent_level() -> None: """Integration test for comprehensive ChatOptions parameter coverage with Azure Response Agent.""" @@ -534,6 +547,7 @@ async def test_azure_responses_client_agent_chat_options_agent_level() -> None: assert len(response.text) > 0 +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled async def test_azure_responses_client_agent_hosted_mcp_tool() -> None: """Integration test for HostedMCPTool with Azure Response Agent using Microsoft Learn MCP.""" @@ -562,6 +576,7 @@ async def test_azure_responses_client_agent_hosted_mcp_tool() -> None: assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"]) +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled @pytest.mark.skip(reason="File search requires API key auth, subscription only allows token auth") async def test_azure_responses_client_file_search() -> None: @@ -588,6 +603,7 @@ async def test_azure_responses_client_file_search() -> None: assert "75" in response.text +@pytest.mark.flaky @skip_if_azure_integration_tests_disabled @pytest.mark.skip(reason="File search requires API key auth, subscription only allows token auth") async def test_azure_responses_client_file_search_streaming() -> None: diff --git a/python/packages/core/tests/main/__init__.py b/python/packages/core/tests/core/__init__.py similarity index 100% rename from python/packages/core/tests/main/__init__.py rename to python/packages/core/tests/core/__init__.py diff --git a/python/packages/core/tests/main/conftest.py b/python/packages/core/tests/core/conftest.py similarity index 96% rename from python/packages/core/tests/main/conftest.py rename to python/packages/core/tests/core/conftest.py index d9980e02cc..ca524a4144 100644 --- a/python/packages/core/tests/main/conftest.py +++ b/python/packages/core/tests/core/conftest.py @@ -8,7 +8,7 @@ from typing import Any from unittest.mock import patch from uuid import uuid4 -from pydantic import BaseModel, Field +from pydantic import BaseModel from pytest import fixture from agent_framework import ( @@ -116,9 +116,11 @@ class MockChatClient: 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) - call_count: int = Field(default=0) + def __init__(self, **kwargs: Any): + super().__init__(**kwargs) + self.run_responses: list[ChatResponse] = [] + self.streaming_responses: list[list[ChatResponseUpdate]] = [] + self.call_count: int = 0 @override async def _inner_get_response( diff --git a/python/packages/core/tests/main/test_agents.py b/python/packages/core/tests/core/test_agents.py similarity index 100% rename from python/packages/core/tests/main/test_agents.py rename to python/packages/core/tests/core/test_agents.py diff --git a/python/packages/core/tests/main/test_clients.py b/python/packages/core/tests/core/test_clients.py similarity index 100% rename from python/packages/core/tests/main/test_clients.py rename to python/packages/core/tests/core/test_clients.py diff --git a/python/packages/core/tests/main/test_logging.py b/python/packages/core/tests/core/test_logging.py similarity index 100% rename from python/packages/core/tests/main/test_logging.py rename to python/packages/core/tests/core/test_logging.py diff --git a/python/packages/core/tests/main/test_mcp.py b/python/packages/core/tests/core/test_mcp.py similarity index 99% rename from python/packages/core/tests/main/test_mcp.py rename to python/packages/core/tests/core/test_mcp.py index c9a2c3c5fd..df768fee8d 100644 --- a/python/packages/core/tests/main/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -510,6 +510,7 @@ def test_local_mcp_streamable_http_tool_init(): # Integration test +@pytest.mark.flaky @skip_if_mcp_integration_tests_disabled async def test_streamable_http_integration(): """Test MCP StreamableHTTP integration.""" diff --git a/python/packages/core/tests/main/test_memory.py b/python/packages/core/tests/core/test_memory.py similarity index 100% rename from python/packages/core/tests/main/test_memory.py rename to python/packages/core/tests/core/test_memory.py diff --git a/python/packages/core/tests/main/test_middleware.py b/python/packages/core/tests/core/test_middleware.py similarity index 100% rename from python/packages/core/tests/main/test_middleware.py rename to python/packages/core/tests/core/test_middleware.py diff --git a/python/packages/core/tests/main/test_middleware_context_result.py b/python/packages/core/tests/core/test_middleware_context_result.py similarity index 100% rename from python/packages/core/tests/main/test_middleware_context_result.py rename to python/packages/core/tests/core/test_middleware_context_result.py diff --git a/python/packages/core/tests/main/test_middleware_with_agent.py b/python/packages/core/tests/core/test_middleware_with_agent.py similarity index 100% rename from python/packages/core/tests/main/test_middleware_with_agent.py rename to python/packages/core/tests/core/test_middleware_with_agent.py diff --git a/python/packages/core/tests/main/test_middleware_with_chat.py b/python/packages/core/tests/core/test_middleware_with_chat.py similarity index 100% rename from python/packages/core/tests/main/test_middleware_with_chat.py rename to python/packages/core/tests/core/test_middleware_with_chat.py diff --git a/python/packages/core/tests/main/test_observability.py b/python/packages/core/tests/core/test_observability.py similarity index 100% rename from python/packages/core/tests/main/test_observability.py rename to python/packages/core/tests/core/test_observability.py diff --git a/python/packages/core/tests/core/test_serializable_mixin.py b/python/packages/core/tests/core/test_serializable_mixin.py new file mode 100644 index 0000000000..1594cb7c5e --- /dev/null +++ b/python/packages/core/tests/core/test_serializable_mixin.py @@ -0,0 +1,190 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for SerializationMixin functionality.""" + +import logging +from typing import Any + +from agent_framework._serialization import SerializationMixin + + +class TestSerializationMixin: + """Test SerializationMixin serialization, deserialization, and dependency injection.""" + + def test_basic_serialization(self): + """Test basic to_dict and from_dict functionality.""" + + class TestClass(SerializationMixin): + def __init__(self, value: str, number: int): + self.value = value + self.number = number + + obj = TestClass(value="test", number=42) + data = obj.to_dict() + + assert data["type"] == "test_class" + assert data["value"] == "test" + assert data["number"] == 42 + + restored = TestClass.from_dict(data) + assert restored.value == "test" + assert restored.number == 42 + + def test_injectable_dependency_no_warning(self, caplog): + """Test that injectable dependencies don't trigger debug logging.""" + + class TestClass(SerializationMixin): + INJECTABLE = {"client"} + + def __init__(self, value: str, client: Any = None): + self.value = value + self.client = client + + mock_client = "mock_client_instance" + + with caplog.at_level(logging.DEBUG): + obj = TestClass.from_dict( + {"type": "test_class", "value": "test"}, + dependencies={"test_class.client": mock_client}, + ) + + assert obj.value == "test" + assert obj.client == mock_client + # No debug message should be logged for injectable dependency + assert not any("is not in INJECTABLE set" in record.message for record in caplog.records) + + def test_non_injectable_dependency_logs_debug(self, caplog): + """Test that non-injectable dependencies trigger debug logging.""" + + class TestClass(SerializationMixin): + INJECTABLE = {"client"} + + def __init__(self, value: str, other: Any = None): + self.value = value + self.other = other + + mock_other = "mock_other_instance" + + with caplog.at_level(logging.DEBUG): + obj = TestClass.from_dict( + {"type": "test_class", "value": "test"}, + dependencies={"test_class.other": mock_other}, + ) + + assert obj.value == "test" + assert obj.other == mock_other + # Debug message should be logged for non-injectable dependency + debug_messages = [record.message for record in caplog.records if record.levelname == "DEBUG"] + assert any("is not in INJECTABLE set" in msg for msg in debug_messages) + assert any("other" in msg for msg in debug_messages) + assert any("client" in msg for msg in debug_messages) # Should mention available injectable + + def test_multiple_dependencies_mixed_injectable(self, caplog): + """Test with both injectable and non-injectable dependencies.""" + + class TestClass(SerializationMixin): + INJECTABLE = {"client", "logger"} + + def __init__( + self, + value: str, + client: Any = None, + logger: Any = None, + other: Any = None, + ): + self.value = value + self.client = client + self.logger = logger + self.other = other + + mock_client = "mock_client" + mock_logger = "mock_logger" + mock_other = "mock_other" + + with caplog.at_level(logging.DEBUG): + obj = TestClass.from_dict( + {"type": "test_class", "value": "test"}, + dependencies={ + "test_class.client": mock_client, + "test_class.logger": mock_logger, + "test_class.other": mock_other, + }, + ) + + assert obj.value == "test" + assert obj.client == mock_client + assert obj.logger == mock_logger + assert obj.other == mock_other + + # Only 'other' should trigger debug logging + debug_messages = [record.message for record in caplog.records if record.levelname == "DEBUG"] + assert any("other" in msg and "is not in INJECTABLE set" in msg for msg in debug_messages) + # 'client' and 'logger' should not be mentioned as non-injectable dependencies + assert not any("Dependency 'client'" in msg and "is not in INJECTABLE set" in msg for msg in debug_messages) + assert not any("Dependency 'logger'" in msg and "is not in INJECTABLE set" in msg for msg in debug_messages) + + def test_no_injectable_set_defined(self, caplog): + """Test behavior when INJECTABLE is not defined (empty set default).""" + + class TestClass(SerializationMixin): + def __init__(self, value: str, client: Any = None): + self.value = value + self.client = client + + mock_client = "mock_client" + + with caplog.at_level(logging.DEBUG): + obj = TestClass.from_dict( + {"type": "test_class", "value": "test"}, + dependencies={"test_class.client": mock_client}, + ) + + assert obj.value == "test" + assert obj.client == mock_client + # Should log debug message since INJECTABLE is empty by default + debug_messages = [record.message for record in caplog.records if record.levelname == "DEBUG"] + assert any("client" in msg and "is not in INJECTABLE set" in msg for msg in debug_messages) + + def test_default_exclude_serialization(self): + """Test that DEFAULT_EXCLUDE fields are not included in to_dict().""" + + class TestClass(SerializationMixin): + DEFAULT_EXCLUDE = {"secret"} + + def __init__(self, value: str, secret: str): + self.value = value + self.secret = secret + + obj = TestClass(value="test", secret="hidden") + data = obj.to_dict() + + assert "value" in data + assert "secret" not in data + assert data["value"] == "test" + + def test_roundtrip_with_injectable_dependency(self): + """Test full roundtrip serialization/deserialization with injectable dependency.""" + + class TestClass(SerializationMixin): + INJECTABLE = {"client"} + DEFAULT_EXCLUDE = {"client"} + + def __init__(self, value: str, number: int, client: Any = None): + self.value = value + self.number = number + self.client = client + + mock_client = "mock_client" + obj = TestClass(value="test", number=42, client=mock_client) + + # Serialize + data = obj.to_dict() + assert data["value"] == "test" + assert data["number"] == 42 + assert "client" not in data # Excluded from serialization + + # Deserialize with dependency injection + restored = TestClass.from_dict(data, dependencies={"test_class.client": mock_client}) + assert restored.value == "test" + assert restored.number == 42 + assert restored.client == mock_client diff --git a/python/packages/core/tests/main/test_telemetry.py b/python/packages/core/tests/core/test_telemetry.py similarity index 100% rename from python/packages/core/tests/main/test_telemetry.py rename to python/packages/core/tests/core/test_telemetry.py diff --git a/python/packages/core/tests/main/test_threads.py b/python/packages/core/tests/core/test_threads.py similarity index 99% rename from python/packages/core/tests/main/test_threads.py rename to python/packages/core/tests/core/test_threads.py index f9fe7c13ed..c04cab577e 100644 --- a/python/packages/core/tests/main/test_threads.py +++ b/python/packages/core/tests/core/test_threads.py @@ -377,10 +377,10 @@ class TestThreadState: def test_init_with_chat_message_store_state(self) -> None: """Test AgentThreadState initialization with chat_message_store_state.""" store_data: dict[str, Any] = {"messages": []} - state = AgentThreadState(chat_message_store_state=store_data) + state = AgentThreadState.model_validate({"chat_message_store_state": store_data}) assert state.service_thread_id is None - assert state.chat_message_store_state == store_data + assert state.chat_message_store_state.messages == [] def test_init_with_both(self) -> None: """Test AgentThreadState initialization with both parameters.""" diff --git a/python/packages/core/tests/main/test_tools.py b/python/packages/core/tests/core/test_tools.py similarity index 100% rename from python/packages/core/tests/main/test_tools.py rename to python/packages/core/tests/core/test_tools.py diff --git a/python/packages/core/tests/main/test_types.py b/python/packages/core/tests/core/test_types.py similarity index 78% rename from python/packages/core/tests/main/test_types.py rename to python/packages/core/tests/core/test_types.py index f1d3f32df2..1572ed471e 100644 --- a/python/packages/core/tests/main/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterable from typing import Any +import pytest from pydantic import BaseModel, ValidationError from pytest import fixture, mark, raises @@ -10,7 +11,6 @@ from agent_framework import ( AgentRunResponse, AgentRunResponseUpdate, AIFunction, - BaseAnnotation, BaseContent, ChatMessage, ChatOptions, @@ -209,7 +209,7 @@ def test_hosted_file_content_minimal(): # Check the type and content assert content.type == "hosted_file" assert content.file_id == "file-456" - assert content.additional_properties is None + assert content.additional_properties == {} assert content.raw_representation is None # Ensure the instance is of type BaseContent @@ -240,7 +240,7 @@ def test_hosted_vector_store_content_minimal(): # Check the type and content assert content.type == "hosted_vector_store" assert content.vector_store_id == "vs-101112" - assert content.additional_properties is None + assert content.additional_properties == {} assert content.raw_representation is None # Ensure the instance is of type BaseContent @@ -1204,33 +1204,6 @@ def test_text_content_add_comprehensive_coverage(): assert result.raw_representation == ["raw1", "raw2", "raw3"] -def test_text_content_add_annotations_coverage(): - """Test TextContent __add__ method with annotation combinations to improve coverage.""" - - ann1 = BaseAnnotation() - ann2 = BaseAnnotation() - - # Test first has annotations, second has None - t1 = TextContent("Hello", annotations=[ann1]) - t2 = TextContent(" World", annotations=None) - result = t1 + t2 - assert result.annotations == [ann1] - - # Test first has None, second has annotations - t1 = TextContent("Hello", annotations=None) - t2 = TextContent(" World", annotations=[ann2]) - result = t1 + t2 - assert result.annotations == [ann2] - - # Test both have annotations - t1 = TextContent("Hello", annotations=[ann1]) - t2 = TextContent(" World", annotations=[ann2]) - result = t1 + t2 - assert len(result.annotations) == 2 - assert ann1 in result.annotations - assert ann2 in result.annotations - - def test_text_content_iadd_coverage(): """Test TextContent __iadd__ method for better coverage.""" @@ -1277,7 +1250,7 @@ def test_comprehensive_to_dict_exclude_options(): text_content = TextContent("Hello", raw_representation=None, additional_properties={"prop": "val"}) text_dict = text_content.to_dict(exclude_none=True) assert "raw_representation" not in text_dict - assert text_dict["additional_properties"] == {"prop": "val"} + assert text_dict["prop"] == "val" # Test with custom exclude set text_dict_exclude = text_content.to_dict(exclude={"additional_properties"}) @@ -1328,8 +1301,6 @@ def test_chat_message_from_dict_with_mixed_content(): {"type": "text", "text": "Hello"}, {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {"arg": "val"}}, {"type": "function_result", "call_id": "call1", "result": "success"}, - # Test with unknown type that falls back to BaseContent - {"type": "unknown_type", "raw_representation": "something"}, ], } @@ -1383,7 +1354,7 @@ def test_comprehensive_serialization_methods(): text_data = { "text": "Hello world", "raw_representation": {"key": "value"}, - "additional_properties": {"prop": "val"}, + "prop": "val", "annotations": None, } text_content = TextContent.from_dict(text_data) @@ -1394,7 +1365,7 @@ def test_comprehensive_serialization_methods(): # Test round-trip text_dict = text_content.to_dict() assert text_dict["text"] == "Hello world" - assert text_dict["additional_properties"] == {"prop": "val"} + assert text_dict["prop"] == "val" # Note: raw_representation is always excluded from to_dict() output # Test with exclude_none @@ -1464,18 +1435,19 @@ def test_usage_content_serialization_with_details(): # Test from_dict with details as dict usage_data = { - "details": {"input_token_count": 10, "output_token_count": 20, "total_token_count": 30}, - "annotations": [ - {"type": "citation", "start": 0, "end": 5, "citation": "source1"}, - {"type": "unknown", "custom_field": "value"}, # Tests fallback to BaseAnnotation - ], + "type": "usage", + "details": { + "type": "usage_details", + "input_token_count": 10, + "output_token_count": 20, + "total_token_count": 30, + "custom_count": 5, + }, } usage_content = UsageContent.from_dict(usage_data) assert isinstance(usage_content.details, UsageDetails) assert usage_content.details.input_token_count == 10 - assert len(usage_content.annotations) == 2 - assert isinstance(usage_content.annotations[0], CitationAnnotation) - assert isinstance(usage_content.annotations[1], BaseAnnotation) + assert usage_content.details.additional_counts["custom_count"] == 5 # Test to_dict with UsageDetails object usage_dict = usage_content.to_dict() @@ -1488,9 +1460,15 @@ def test_function_approval_response_content_serialization(): # Test from_dict with function_call as dict response_data = { + "type": "function_approval_response", "id": "response123", "approved": True, - "function_call": {"call_id": "call123", "name": "test_func", "arguments": {"param": "value"}}, + "function_call": { + "type": "function_call", + "call_id": "call123", + "name": "test_func", + "arguments": {"param": "value"}, + }, } response_content = FunctionApprovalResponseContent.from_dict(response_data) assert isinstance(response_content.function_call, FunctionCallContent) @@ -1512,7 +1490,12 @@ def test_chat_response_complex_serialization(): {"role": "assistant", "contents": [{"type": "text", "text": "Hi there"}]}, ], "finish_reason": {"value": "stop"}, - "usage_details": {"input_token_count": 5, "output_token_count": 8, "total_token_count": 13}, + "usage_details": { + "type": "usage_details", + "input_token_count": 5, + "output_token_count": 8, + "total_token_count": 13, + }, "model_id": "gpt-4", # Test alias handling } @@ -1543,22 +1526,21 @@ def test_chat_response_update_all_content_types(): {"type": "error", "error": "An error occurred"}, {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}}, {"type": "function_result", "call_id": "call1", "result": "success"}, - {"type": "usage", "details": {"input_token_count": 1}}, + {"type": "usage", "details": {"type": "usage_details", "input_token_count": 1}}, {"type": "hosted_file", "file_id": "file123"}, {"type": "hosted_vector_store", "vector_store_id": "vs123"}, { "type": "function_approval_request", "id": "req1", - "function_call": {"call_id": "call1", "name": "func", "arguments": {}}, + "function_call": {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}}, }, { "type": "function_approval_response", "id": "resp1", "approved": True, - "function_call": {"call_id": "call1", "name": "func", "arguments": {}}, + "function_call": {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}}, }, {"type": "text_reasoning", "text": "reasoning"}, - {"type": "unknown_type", "custom_field": "value"}, # Tests fallback ] } @@ -1586,7 +1568,12 @@ def test_agent_run_response_complex_serialization(): {"role": "user", "contents": [{"type": "text", "text": "Hello"}]}, {"role": "assistant", "contents": [{"type": "text", "text": "Hi"}]}, ], - "usage_details": {"input_token_count": 3, "output_token_count": 2, "total_token_count": 5}, + "usage_details": { + "type": "usage_details", + "input_token_count": 3, + "output_token_count": 2, + "total_token_count": 5, + }, } response = AgentRunResponse.from_dict(response_data) @@ -1612,22 +1599,21 @@ def test_agent_run_response_update_all_content_types(): {"type": "error", "error": "An error occurred"}, {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}}, {"type": "function_result", "call_id": "call1", "result": "success"}, - {"type": "usage", "details": {"input_token_count": 1}}, + {"type": "usage", "details": {"type": "usage_details", "input_token_count": 1}}, {"type": "hosted_file", "file_id": "file123"}, {"type": "hosted_vector_store", "vector_store_id": "vs123"}, { "type": "function_approval_request", "id": "req1", - "function_call": {"call_id": "call1", "name": "func", "arguments": {}}, + "function_call": {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}}, }, { "type": "function_approval_response", "id": "resp1", "approved": True, - "function_call": {"call_id": "call1", "name": "func", "arguments": {}}, + "function_call": {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}}, }, {"type": "text_reasoning", "text": "reasoning"}, - {"type": "unknown_type", "custom_field": "value"}, # Tests fallback ], "role": {"value": "assistant"}, # Test role as dict } @@ -1648,3 +1634,394 @@ def test_agent_run_response_update_all_content_types(): update_str = AgentRunResponseUpdate.from_dict(update_data_str_role) assert isinstance(update_str.role, Role) assert update_str.role.value == "user" + + +# region Serialization + + +@mark.parametrize( + "content_class,init_kwargs", + [ + pytest.param( + TextContent, + { + "type": "text", + "text": "Hello world", + "raw_representation": "raw", + }, + id="text_content", + ), + pytest.param( + TextReasoningContent, + { + "type": "text_reasoning", + "text": "Reasoning text", + "raw_representation": "raw", + }, + id="text_reasoning_content", + ), + pytest.param( + DataContent, + { + "type": "data", + "uri": "data:text/plain;base64,dGVzdCBkYXRh", + }, + id="data_content_with_uri", + ), + pytest.param( + DataContent, + { + "type": "data", + "data": b"test data", + "media_type": "text/plain", + }, + id="data_content_with_bytes", + ), + pytest.param( + UriContent, + { + "type": "uri", + "uri": "http://example.com", + "media_type": "text/html", + }, + id="uri_content", + ), + pytest.param( + HostedFileContent, + {"type": "hosted_file", "file_id": "file-123"}, + id="hosted_file_content", + ), + pytest.param( + HostedVectorStoreContent, + { + "type": "hosted_vector_store", + "vector_store_id": "vs-789", + }, + id="hosted_vector_store_content", + ), + pytest.param( + FunctionCallContent, + { + "type": "function_call", + "call_id": "call-1", + "name": "test_func", + "arguments": {"arg": "val"}, + }, + id="function_call_content", + ), + pytest.param( + FunctionResultContent, + { + "type": "function_result", + "call_id": "call-1", + "result": "success", + }, + id="function_result_content", + ), + pytest.param( + ErrorContent, + { + "type": "error", + "message": "Error occurred", + "error_code": "E001", + }, + id="error_content", + ), + pytest.param( + UsageContent, + { + "type": "usage", + "details": { + "type": "usage_details", + "input_token_count": 10, + "output_token_count": 20, + "reasoning_tokens": 5, + }, + }, + id="usage_content", + ), + pytest.param( + FunctionApprovalRequestContent, + { + "type": "function_approval_request", + "id": "req-1", + "function_call": {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}}, + }, + id="function_approval_request", + ), + pytest.param( + FunctionApprovalResponseContent, + { + "type": "function_approval_response", + "id": "resp-1", + "approved": True, + "function_call": {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}}, + }, + id="function_approval_response", + ), + pytest.param( + ChatMessage, + { + "role": {"type": "role", "value": "user"}, + "contents": [ + {"type": "text", "text": "Hello"}, + {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}}, + ], + "message_id": "msg-123", + "author_name": "User", + }, + id="chat_message", + ), + pytest.param( + ChatResponse, + { + "type": "chat_response", + "messages": [ + { + "type": "chat_message", + "role": {"type": "role", "value": "user"}, + "contents": [{"type": "text", "text": "Hello"}], + }, + { + "type": "chat_message", + "role": {"type": "role", "value": "assistant"}, + "contents": [{"type": "text", "text": "Hi there"}], + }, + ], + "finish_reason": {"type": "finish_reason", "value": "stop"}, + "usage_details": { + "type": "usage_details", + "input_token_count": 10, + "output_token_count": 20, + "total_token_count": 30, + }, + "response_id": "resp-123", + "model_id": "gpt-4", + }, + id="chat_response", + ), + pytest.param( + ChatResponseUpdate, + { + "contents": [ + {"type": "text", "text": "Hello"}, + {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}}, + ], + "role": {"type": "role", "value": "assistant"}, + "finish_reason": {"type": "finish_reason", "value": "stop"}, + "message_id": "msg-123", + "response_id": "resp-123", + }, + id="chat_response_update", + ), + pytest.param( + AgentRunResponse, + { + "messages": [ + { + "role": {"type": "role", "value": "user"}, + "contents": [{"type": "text", "text": "Question"}], + }, + { + "role": {"type": "role", "value": "assistant"}, + "contents": [{"type": "text", "text": "Answer"}], + }, + ], + "response_id": "run-123", + "usage_details": { + "type": "usage_details", + "input_token_count": 5, + "output_token_count": 3, + "total_token_count": 8, + }, + }, + id="agent_run_response", + ), + pytest.param( + AgentRunResponseUpdate, + { + "contents": [ + {"type": "text", "text": "Streaming"}, + {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}}, + ], + "role": {"type": "role", "value": "assistant"}, + "message_id": "msg-123", + "response_id": "run-123", + "author_name": "Agent", + }, + id="agent_run_response_update", + ), + ], +) +def test_content_roundtrip_serialization(content_class: type[BaseContent], init_kwargs: dict[str, Any]): + """Test to_dict/from_dict roundtrip for all content types.""" + # Create instance + content = content_class(**init_kwargs) + + # Serialize to dict + content_dict = content.to_dict() + + # Verify type key is in serialized dict + assert "type" in content_dict + if hasattr(content, "type"): + assert content_dict["type"] == content.type # type: ignore[attr-defined] + + # Deserialize from dict + reconstructed = content_class.from_dict(content_dict) + + # Verify type + assert isinstance(reconstructed, content_class) + # Check type attribute dynamically + if hasattr(content, "type"): + assert reconstructed.type == content.type # type: ignore[attr-defined] + + # Verify key attributes (excluding raw_representation which is not serialized) + for key, value in init_kwargs.items(): + if key == "type": + continue + if key == "raw_representation": + # raw_representation is intentionally excluded from serialization + continue + + # Special handling for DataContent created with 'data' parameter + if content_class == DataContent and key == "data": + # DataContent converts 'data' to 'uri', so we skip checking 'data' attribute + # Instead we verify that uri and media_type are set correctly + assert hasattr(reconstructed, "uri") + assert hasattr(reconstructed, "media_type") + assert reconstructed.media_type == init_kwargs.get("media_type") + # Verify the uri contains the encoded data + assert reconstructed.uri.startswith(f"data:{init_kwargs.get('media_type')};base64,") + continue + + reconstructed_value = getattr(reconstructed, key) + + # Special handling for nested SerializationMixin objects + if hasattr(value, "to_dict"): + # Compare the serialized forms + assert reconstructed_value.to_dict() == value.to_dict() + # Special handling for lists that may contain dicts converted to objects + elif isinstance(value, list) and value and isinstance(reconstructed_value, list): + # Check if this is a list of objects that were created from dicts + if isinstance(value[0], dict) and hasattr(reconstructed_value[0], "to_dict"): + # Compare each item by serializing the reconstructed object + assert len(reconstructed_value) == len(value) + + else: + assert reconstructed_value == value + # Special handling for dicts that get converted to objects (like UsageDetails, FunctionCallContent) + elif isinstance(value, dict) and hasattr(reconstructed_value, "to_dict"): + # Compare the dict with the serialized form of the object, excluding 'type' key + reconstructed_dict = reconstructed_value.to_dict() + assert len(reconstructed_dict) == len(value) + else: + assert reconstructed_value == value + + +def test_text_content_with_annotations_serialization(): + """Test TextContent with CitationAnnotation and TextSpanRegion roundtrip serialization.""" + # Create TextSpanRegion + region = TextSpanRegion(start_index=0, end_index=5) + + # Create CitationAnnotation with region + citation = CitationAnnotation( + title="Test Citation", + url="http://example.com/citation", + file_id="file-123", + tool_name="test_tool", + snippet="This is a test snippet", + annotated_regions=[region], + additional_properties={"custom": "value"}, + ) + + # Create TextContent with annotation + content = TextContent( + text="Hello world", annotations=[citation], additional_properties={"content_key": "content_val"} + ) + + # Serialize to dict + content_dict = content.to_dict() + + # Verify structure + assert content_dict["type"] == "text" + assert content_dict["text"] == "Hello world" + assert content_dict["content_key"] == "content_val" + assert len(content_dict["annotations"]) == 1 + + # Verify annotation structure + annotation_dict = content_dict["annotations"][0] + assert annotation_dict["type"] == "citation" + assert annotation_dict["title"] == "Test Citation" + assert annotation_dict["url"] == "http://example.com/citation" + assert annotation_dict["file_id"] == "file-123" + assert annotation_dict["tool_name"] == "test_tool" + assert annotation_dict["snippet"] == "This is a test snippet" + assert annotation_dict["custom"] == "value" + + # Verify region structure + assert len(annotation_dict["annotated_regions"]) == 1 + region_dict = annotation_dict["annotated_regions"][0] + assert region_dict["type"] == "text_span" + assert region_dict["start_index"] == 0 + assert region_dict["end_index"] == 5 + + # Deserialize from dict + reconstructed = TextContent.from_dict(content_dict) + + # Verify reconstructed content + assert isinstance(reconstructed, TextContent) + assert reconstructed.text == "Hello world" + assert reconstructed.type == "text" + assert reconstructed.additional_properties == {"content_key": "content_val"} + + # Verify reconstructed annotation + assert len(reconstructed.annotations) == 1 # type: ignore[arg-type] + recon_annotation = reconstructed.annotations[0] # type: ignore[index] + assert isinstance(recon_annotation, CitationAnnotation) + assert recon_annotation.title == "Test Citation" + assert recon_annotation.url == "http://example.com/citation" + assert recon_annotation.file_id == "file-123" + assert recon_annotation.tool_name == "test_tool" + assert recon_annotation.snippet == "This is a test snippet" + assert recon_annotation.additional_properties == {"custom": "value"} + + # Verify reconstructed region + assert len(recon_annotation.annotated_regions) == 1 # type: ignore[arg-type] + recon_region = recon_annotation.annotated_regions[0] # type: ignore[index] + assert isinstance(recon_region, TextSpanRegion) + assert recon_region.start_index == 0 + assert recon_region.end_index == 5 + assert recon_region.type == "text_span" + + +def test_text_content_with_multiple_annotations_serialization(): + """Test TextContent with multiple annotations roundtrip serialization.""" + # Create multiple regions + region1 = TextSpanRegion(start_index=0, end_index=5) + region2 = TextSpanRegion(start_index=6, end_index=11) + + # Create multiple citations + citation1 = CitationAnnotation(title="Citation 1", url="http://example.com/1", annotated_regions=[region1]) + + citation2 = CitationAnnotation(title="Citation 2", url="http://example.com/2", annotated_regions=[region2]) + + # Create TextContent with multiple annotations + content = TextContent(text="Hello world", annotations=[citation1, citation2]) + + # Serialize + content_dict = content.to_dict() + + # Verify we have 2 annotations + assert len(content_dict["annotations"]) == 2 + assert content_dict["annotations"][0]["title"] == "Citation 1" + assert content_dict["annotations"][1]["title"] == "Citation 2" + + # Deserialize + reconstructed = TextContent.from_dict(content_dict) + + # Verify reconstruction + assert len(reconstructed.annotations) == 2 + assert all(isinstance(ann, CitationAnnotation) for ann in reconstructed.annotations) + assert reconstructed.annotations[0].title == "Citation 1" + assert reconstructed.annotations[1].title == "Citation 2" + assert all(isinstance(ann.annotated_regions[0], TextSpanRegion) for ann in reconstructed.annotations) diff --git a/python/packages/core/tests/main/utils.py b/python/packages/core/tests/core/utils.py similarity index 100% rename from python/packages/core/tests/main/utils.py rename to python/packages/core/tests/core/utils.py diff --git a/python/packages/core/tests/openai/test_openai_assistants_client.py b/python/packages/core/tests/openai/test_openai_assistants_client.py index 9d2a1bf5ce..fdaeb06a95 100644 --- a/python/packages/core/tests/openai/test_openai_assistants_client.py +++ b/python/packages/core/tests/openai/test_openai_assistants_client.py @@ -46,23 +46,26 @@ skip_if_openai_integration_tests_disabled = pytest.mark.skipif( def create_test_openai_assistants_client( mock_async_openai: MagicMock, - ai_model_id: str | None = None, + model_id: str | None = None, assistant_id: str | None = None, assistant_name: str | None = None, thread_id: str | None = None, should_delete_assistant: bool = False, ) -> OpenAIAssistantsClient: - """Helper function to create OpenAIAssistantsClient instances for testing, bypassing Pydantic validation.""" - return OpenAIAssistantsClient.model_construct( - ai_model_id=ai_model_id or "gpt-4", + """Helper function to create OpenAIAssistantsClient instances for testing.""" + client = OpenAIAssistantsClient( + model_id=model_id or "gpt-4", assistant_id=assistant_id, assistant_name=assistant_name, thread_id=thread_id, api_key="test-api-key", org_id="test-org-id", - client=mock_async_openai, - _should_delete_assistant=should_delete_assistant, + async_client=mock_async_openai, ) + # Set the _should_delete_assistant flag directly if needed + if should_delete_assistant: + object.__setattr__(client, "_should_delete_assistant", True) + return client async def create_vector_store(client: OpenAIAssistantsClient) -> tuple[str, HostedVectorStoreContent]: @@ -117,11 +120,11 @@ def mock_async_openai() -> MagicMock: def test_openai_assistants_client_init_with_client(mock_async_openai: MagicMock) -> None: """Test OpenAIAssistantsClient initialization with existing client.""" chat_client = create_test_openai_assistants_client( - mock_async_openai, ai_model_id="gpt-4", assistant_id="existing-assistant-id", thread_id="test-thread-id" + mock_async_openai, model_id="gpt-4", assistant_id="existing-assistant-id", thread_id="test-thread-id" ) assert chat_client.client is mock_async_openai - assert chat_client.ai_model_id == "gpt-4" + assert chat_client.model_id == "gpt-4" assert chat_client.assistant_id == "existing-assistant-id" assert chat_client.thread_id == "test-thread-id" assert not chat_client._should_delete_assistant # type: ignore @@ -133,19 +136,16 @@ def test_openai_assistants_client_init_auto_create_client( mock_async_openai: MagicMock, ) -> None: """Test OpenAIAssistantsClient initialization with auto-created client.""" - chat_client = OpenAIAssistantsClient.model_construct( - ai_model_id=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"], - assistant_id=None, + chat_client = OpenAIAssistantsClient( + model_id=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"], assistant_name="TestAssistant", - thread_id=None, api_key=openai_unit_test_env["OPENAI_API_KEY"], org_id=openai_unit_test_env["OPENAI_ORG_ID"], - client=mock_async_openai, - _should_delete_assistant=False, + async_client=mock_async_openai, ) assert chat_client.client is mock_async_openai - assert chat_client.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] + assert chat_client.model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] assert chat_client.assistant_id is None assert chat_client.assistant_name == "TestAssistant" assert not chat_client._should_delete_assistant # type: ignore @@ -155,7 +155,7 @@ def test_openai_assistants_client_init_validation_fail() -> None: """Test OpenAIAssistantsClient initialization with validation failure.""" with pytest.raises(ServiceInitializationError): # Force failure by providing invalid model ID type - this should cause validation to fail - OpenAIAssistantsClient(ai_model_id=123, api_key="valid-key") # type: ignore + OpenAIAssistantsClient(model_id=123, api_key="valid-key") # type: ignore @pytest.mark.parametrize("exclude_list", [["OPENAI_CHAT_MODEL_ID"]], indirect=True) @@ -171,7 +171,7 @@ def test_openai_assistants_client_init_missing_model_id(openai_unit_test_env: di def test_openai_assistants_client_init_missing_api_key(openai_unit_test_env: dict[str, str]) -> None: """Test OpenAIAssistantsClient initialization with missing API key.""" with pytest.raises(ServiceInitializationError): - OpenAIAssistantsClient(ai_model_id="gpt-4", env_file_path="nonexistent.env") + OpenAIAssistantsClient(model_id="gpt-4", env_file_path="nonexistent.env") def test_openai_assistants_client_init_with_default_headers(openai_unit_test_env: dict[str, str]) -> None: @@ -179,12 +179,12 @@ def test_openai_assistants_client_init_with_default_headers(openai_unit_test_env default_headers = {"X-Unit-Test": "test-guid"} chat_client = OpenAIAssistantsClient( - ai_model_id="gpt-4", + model_id="gpt-4", api_key=openai_unit_test_env["OPENAI_API_KEY"], default_headers=default_headers, ) - assert chat_client.ai_model_id == "gpt-4" + assert chat_client.model_id == "gpt-4" assert isinstance(chat_client, ChatClientProtocol) # Assert that the default header we added is present in the client's default headers @@ -211,7 +211,7 @@ async def test_openai_assistants_client_get_assistant_id_or_create_create_new( ) -> None: """Test _get_assistant_id_or_create when creating a new assistant.""" chat_client = create_test_openai_assistants_client( - mock_async_openai, ai_model_id="gpt-4", assistant_name="TestAssistant" + mock_async_openai, model_id="gpt-4", assistant_name="TestAssistant" ) assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore @@ -269,7 +269,7 @@ def test_openai_assistants_client_serialize(openai_unit_test_env: dict[str, str] # Test basic initialization and to_dict chat_client = OpenAIAssistantsClient( - ai_model_id="gpt-4", + model_id="gpt-4", assistant_id="test-assistant-id", assistant_name="TestAssistant", thread_id="test-thread-id", @@ -280,11 +280,10 @@ def test_openai_assistants_client_serialize(openai_unit_test_env: dict[str, str] dumped_settings = chat_client.to_dict() - assert dumped_settings["ai_model_id"] == "gpt-4" + assert dumped_settings["model_id"] == "gpt-4" assert dumped_settings["assistant_id"] == "test-assistant-id" assert dumped_settings["assistant_name"] == "TestAssistant" assert dumped_settings["thread_id"] == "test-thread-id" - assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"] assert dumped_settings["org_id"] == openai_unit_test_env["OPENAI_ORG_ID"] # Assert that the default header we added is present in the dumped_settings default headers @@ -915,6 +914,7 @@ def get_weather( return f"The weather in {location} is sunny with a high of 25°C." +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_client_get_response() -> None: """Test OpenAI Assistants Client response.""" @@ -939,6 +939,7 @@ async def test_openai_assistants_client_get_response() -> None: assert any(word in response.text.lower() for word in ["sunny", "25", "weather", "seattle"]) +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_client_get_response_tools() -> None: """Test OpenAI Assistants Client response with tools.""" @@ -960,6 +961,7 @@ async def test_openai_assistants_client_get_response_tools() -> None: assert any(word in response.text.lower() for word in ["sunny", "25", "weather"]) +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_client_streaming() -> None: """Test OpenAI Assistants Client streaming response.""" @@ -990,6 +992,7 @@ async def test_openai_assistants_client_streaming() -> None: assert any(word in full_message.lower() for word in ["sunny", "25", "weather", "seattle"]) +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_client_streaming_tools() -> None: """Test OpenAI Assistants Client streaming response with tools.""" @@ -1016,6 +1019,7 @@ async def test_openai_assistants_client_streaming_tools() -> None: assert any(word in full_message.lower() for word in ["sunny", "25", "weather"]) +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_client_with_existing_assistant() -> None: """Test OpenAI Assistants Client with existing assistant ID.""" @@ -1028,7 +1032,7 @@ async def test_openai_assistants_client_with_existing_assistant() -> None: # Now test using the existing assistant async with OpenAIAssistantsClient( - ai_model_id="gpt-4o-mini", assistant_id=assistant_id + model_id="gpt-4o-mini", assistant_id=assistant_id ) as openai_assistants_client: assert isinstance(openai_assistants_client, ChatClientProtocol) assert openai_assistants_client.assistant_id == assistant_id @@ -1043,6 +1047,7 @@ async def test_openai_assistants_client_with_existing_assistant() -> None: assert len(response.text) > 0 +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled @pytest.mark.skip(reason="OpenAI file search functionality is currently broken - tracked in GitHub issue") async def test_openai_assistants_client_file_search() -> None: @@ -1066,6 +1071,7 @@ async def test_openai_assistants_client_file_search() -> None: assert any(word in response.text.lower() for word in ["sunny", "25", "weather"]) +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled @pytest.mark.skip(reason="OpenAI file search functionality is currently broken - tracked in GitHub issue") async def test_openai_assistants_client_file_search_streaming() -> None: @@ -1096,6 +1102,7 @@ async def test_openai_assistants_client_file_search_streaming() -> None: assert any(word in full_message.lower() for word in ["sunny", "25", "weather"]) +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_agent_basic_run(): """Test ChatAgent basic run functionality with OpenAIAssistantsClient.""" @@ -1112,6 +1119,7 @@ async def test_openai_assistants_agent_basic_run(): assert "Hello World" in response.text +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_agent_basic_run_streaming(): """Test ChatAgent basic streaming functionality with OpenAIAssistantsClient.""" @@ -1131,6 +1139,7 @@ async def test_openai_assistants_agent_basic_run_streaming(): assert "streaming response test" in full_message.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_agent_thread_persistence(): """Test ChatAgent thread persistence across runs with OpenAIAssistantsClient.""" @@ -1159,6 +1168,7 @@ async def test_openai_assistants_agent_thread_persistence(): assert thread.service_thread_id is not None +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_agent_existing_thread_id(): """Test ChatAgent with existing thread ID to continue conversations across agent instances.""" @@ -1203,6 +1213,7 @@ async def test_openai_assistants_agent_existing_thread_id(): assert "paris" in response2.text.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_agent_code_interpreter(): """Test ChatAgent with code interpreter through OpenAIAssistantsClient.""" @@ -1222,6 +1233,7 @@ async def test_openai_assistants_agent_code_interpreter(): assert "120" in response.text or "factorial" in response.text.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_assistants_client_agent_level_tool_persistence(): """Test that agent-level tools persist across multiple runs with OpenAI Assistants Client.""" diff --git a/python/packages/core/tests/openai/test_openai_chat_client.py b/python/packages/core/tests/openai/test_openai_chat_client.py index 5099769fa3..0c03eb775e 100644 --- a/python/packages/core/tests/openai/test_openai_chat_client.py +++ b/python/packages/core/tests/openai/test_openai_chat_client.py @@ -41,22 +41,22 @@ def test_init(openai_unit_test_env: dict[str, str]) -> None: # Test successful initialization open_ai_chat_completion = OpenAIChatClient() - assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] + assert open_ai_chat_completion.model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] assert isinstance(open_ai_chat_completion, ChatClientProtocol) def test_init_validation_fail() -> None: # Test successful initialization with pytest.raises(ServiceInitializationError): - OpenAIChatClient(api_key="34523", ai_model_id={"test": "dict"}) # type: ignore + OpenAIChatClient(api_key="34523", model_id={"test": "dict"}) # type: ignore -def test_init_ai_model_id_constructor(openai_unit_test_env: dict[str, str]) -> None: +def test_init_model_id_constructor(openai_unit_test_env: dict[str, str]) -> None: # Test successful initialization - ai_model_id = "test_model_id" - open_ai_chat_completion = OpenAIChatClient(ai_model_id=ai_model_id) + model_id = "test_model_id" + open_ai_chat_completion = OpenAIChatClient(model_id=model_id) - assert open_ai_chat_completion.ai_model_id == ai_model_id + assert open_ai_chat_completion.model_id == model_id assert isinstance(open_ai_chat_completion, ChatClientProtocol) @@ -68,7 +68,7 @@ def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None: default_headers=default_headers, ) - assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] + assert open_ai_chat_completion.model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] assert isinstance(open_ai_chat_completion, ChatClientProtocol) # Assert that the default header we added is present in the client's default headers @@ -95,7 +95,7 @@ def test_init_base_url_from_settings_env() -> None: }, ): client = OpenAIChatClient() - assert client.ai_model_id == "gpt-5" + assert client.model_id == "gpt-5" assert str(client.client.base_url) == "https://custom-openai-endpoint.com/v1/" @@ -109,11 +109,11 @@ def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None: @pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True) def test_init_with_empty_api_key(openai_unit_test_env: dict[str, str]) -> None: - ai_model_id = "test_model_id" + model_id = "test_model_id" with pytest.raises(ServiceInitializationError): OpenAIChatClient( - ai_model_id=ai_model_id, + model_id=model_id, env_file_path="test.env", ) @@ -122,15 +122,14 @@ def test_serialize(openai_unit_test_env: dict[str, str]) -> None: default_headers = {"X-Unit-Test": "test-guid"} settings = { - "ai_model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"], + "model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"], "api_key": openai_unit_test_env["OPENAI_API_KEY"], "default_headers": default_headers, } open_ai_chat_completion = OpenAIChatClient.from_dict(settings) dumped_settings = open_ai_chat_completion.to_dict() - assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] - assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"] + assert dumped_settings["model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] # Assert that the default header we added is present in the dumped_settings default headers for key, value in default_headers.items(): assert key in dumped_settings["default_headers"] @@ -141,18 +140,17 @@ def test_serialize(openai_unit_test_env: dict[str, str]) -> None: def test_serialize_with_org_id(openai_unit_test_env: dict[str, str]) -> None: settings = { - "ai_model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"], + "model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"], "api_key": openai_unit_test_env["OPENAI_API_KEY"], "org_id": openai_unit_test_env["OPENAI_ORG_ID"], } open_ai_chat_completion = OpenAIChatClient.from_dict(settings) dumped_settings = open_ai_chat_completion.to_dict() - assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] - assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"] + assert dumped_settings["model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"] assert dumped_settings["org_id"] == openai_unit_test_env["OPENAI_ORG_ID"] # Assert that the 'User-Agent' header is not present in the dumped_settings default headers - assert "User-Agent" not in dumped_settings["default_headers"] + assert "User-Agent" not in dumped_settings.get("default_headers", {}) async def test_content_filter_exception_handling(openai_unit_test_env: dict[str, str]) -> None: @@ -210,6 +208,7 @@ def get_weather(location: str) -> str: return f"The weather in {location} is sunny and 72°F." +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_completion_response() -> None: """Test OpenAI chat completion responses.""" @@ -237,6 +236,7 @@ async def test_openai_chat_completion_response() -> None: assert "scientists" in response.text +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_completion_response_tools() -> None: """Test OpenAI chat completion responses.""" @@ -259,6 +259,7 @@ async def test_openai_chat_completion_response_tools() -> None: assert "scientists" in response.text +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_streaming() -> None: """Test Azure OpenAI chat completion responses.""" @@ -294,6 +295,7 @@ async def test_openai_chat_client_streaming() -> None: assert "scientists" in full_message +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_streaming_tools() -> None: """Test AzureOpenAI chat completion responses.""" @@ -321,10 +323,11 @@ async def test_openai_chat_client_streaming_tools() -> None: assert "scientists" in full_message +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_web_search() -> None: # Currently only a select few models support web search tool calls - openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview") + openai_chat_client = OpenAIChatClient(model_id="gpt-4o-search-preview") assert isinstance(openai_chat_client, ChatClientProtocol) @@ -361,9 +364,10 @@ async def test_openai_chat_client_web_search() -> None: assert response.text is not None +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_web_search_streaming() -> None: - openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview") + openai_chat_client = OpenAIChatClient(model_id="gpt-4o-search-preview") assert isinstance(openai_chat_client, ChatClientProtocol) @@ -414,11 +418,12 @@ async def test_openai_chat_client_web_search_streaming() -> None: assert full_message is not None +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_agent_basic_run(): """Test OpenAI chat client agent basic run functionality with OpenAIChatClient.""" async with ChatAgent( - chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"), + chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"), ) as agent: # Test basic run response = await agent.run("Hello! Please respond with 'Hello World' exactly.") @@ -429,11 +434,12 @@ async def test_openai_chat_client_agent_basic_run(): assert "hello world" in response.text.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_agent_basic_run_streaming(): """Test OpenAI chat client agent basic streaming functionality with OpenAIChatClient.""" async with ChatAgent( - chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"), + chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"), ) as agent: # Test streaming run full_text = "" @@ -446,11 +452,12 @@ async def test_openai_chat_client_agent_basic_run_streaming(): assert "streaming response test" in full_text.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_agent_thread_persistence(): """Test OpenAI chat client agent thread persistence across runs with OpenAIChatClient.""" async with ChatAgent( - chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"), + chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"), instructions="You are a helpful assistant with good memory.", ) as agent: # Create a new thread that will be reused @@ -470,6 +477,7 @@ async def test_openai_chat_client_agent_thread_persistence(): assert "alice" in response2.text.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_agent_existing_thread(): """Test OpenAI chat client agent with existing thread to continue conversations across agent instances.""" @@ -477,7 +485,7 @@ async def test_openai_chat_client_agent_existing_thread(): preserved_thread = None async with ChatAgent( - chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"), + chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"), instructions="You are a helpful assistant with good memory.", ) as first_agent: # Start a conversation and capture the thread @@ -493,7 +501,7 @@ async def test_openai_chat_client_agent_existing_thread(): # Second conversation - reuse the thread in a new agent instance if preserved_thread: async with ChatAgent( - chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"), + chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"), instructions="You are a helpful assistant with good memory.", ) as second_agent: # Reuse the preserved thread @@ -504,12 +512,13 @@ async def test_openai_chat_client_agent_existing_thread(): assert "alice" in second_response.text.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_agent_level_tool_persistence(): """Test that agent-level tools persist across multiple runs with OpenAI Chat Client.""" async with ChatAgent( - chat_client=OpenAIChatClient(ai_model_id="gpt-4.1"), + chat_client=OpenAIChatClient(model_id="gpt-4.1"), instructions="You are a helpful assistant that uses available tools.", tools=[get_weather], # Agent-level tool ) as agent: @@ -530,6 +539,7 @@ async def test_openai_chat_client_agent_level_tool_persistence(): assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"]) +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_chat_client_run_level_tool_isolation(): """Test that run-level tools are isolated to specific runs and don't persist with OpenAI Chat Client.""" @@ -544,7 +554,7 @@ async def test_openai_chat_client_run_level_tool_isolation(): return f"The weather in {location} is sunny and 72°F." async with ChatAgent( - chat_client=OpenAIChatClient(ai_model_id="gpt-4.1"), + chat_client=OpenAIChatClient(model_id="gpt-4.1"), instructions="You are a helpful assistant.", ) as agent: # First run - use run-level tool @@ -571,7 +581,7 @@ async def test_openai_chat_client_run_level_tool_isolation(): async def test_exception_message_includes_original_error_details() -> None: """Test that exception messages include original error details in the new format.""" - client = OpenAIChatClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIChatClient(model_id="test-model", api_key="test-key") messages = [ChatMessage(role="user", text="test message")] mock_response = MagicMock() diff --git a/python/packages/core/tests/openai/test_openai_responses_client.py b/python/packages/core/tests/openai/test_openai_responses_client.py index c8392a7831..305c943141 100644 --- a/python/packages/core/tests/openai/test_openai_responses_client.py +++ b/python/packages/core/tests/openai/test_openai_responses_client.py @@ -96,22 +96,22 @@ def test_init(openai_unit_test_env: dict[str, str]) -> None: # Test successful initialization openai_responses_client = OpenAIResponsesClient() - assert openai_responses_client.ai_model_id == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"] + assert openai_responses_client.model_id == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"] assert isinstance(openai_responses_client, ChatClientProtocol) def test_init_validation_fail() -> None: # Test successful initialization with pytest.raises(ServiceInitializationError): - OpenAIResponsesClient(api_key="34523", ai_model_id={"test": "dict"}) # type: ignore + OpenAIResponsesClient(api_key="34523", model_id={"test": "dict"}) # type: ignore -def test_init_ai_model_id_constructor(openai_unit_test_env: dict[str, str]) -> None: +def test_init_model_id_constructor(openai_unit_test_env: dict[str, str]) -> None: # Test successful initialization - ai_model_id = "test_model_id" - openai_responses_client = OpenAIResponsesClient(ai_model_id=ai_model_id) + model_id = "test_model_id" + openai_responses_client = OpenAIResponsesClient(model_id=model_id) - assert openai_responses_client.ai_model_id == ai_model_id + assert openai_responses_client.model_id == model_id assert isinstance(openai_responses_client, ChatClientProtocol) @@ -123,7 +123,7 @@ def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None: default_headers=default_headers, ) - assert openai_responses_client.ai_model_id == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"] + assert openai_responses_client.model_id == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"] assert isinstance(openai_responses_client, ChatClientProtocol) # Assert that the default header we added is present in the client's default headers @@ -142,11 +142,11 @@ def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None: @pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True) def test_init_with_empty_api_key(openai_unit_test_env: dict[str, str]) -> None: - ai_model_id = "test_model_id" + model_id = "test_model_id" with pytest.raises(ServiceInitializationError): OpenAIResponsesClient( - ai_model_id=ai_model_id, + model_id=model_id, env_file_path="test.env", ) @@ -155,15 +155,14 @@ def test_serialize(openai_unit_test_env: dict[str, str]) -> None: default_headers = {"X-Unit-Test": "test-guid"} settings = { - "ai_model_id": openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"], + "model_id": openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"], "api_key": openai_unit_test_env["OPENAI_API_KEY"], "default_headers": default_headers, } openai_responses_client = OpenAIResponsesClient.from_dict(settings) dumped_settings = openai_responses_client.to_dict() - assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"] - assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"] + assert dumped_settings["model_id"] == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"] # Assert that the default header we added is present in the dumped_settings default headers for key, value in default_headers.items(): assert key in dumped_settings["default_headers"] @@ -174,24 +173,23 @@ def test_serialize(openai_unit_test_env: dict[str, str]) -> None: def test_serialize_with_org_id(openai_unit_test_env: dict[str, str]) -> None: settings = { - "ai_model_id": openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"], + "model_id": openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"], "api_key": openai_unit_test_env["OPENAI_API_KEY"], "org_id": openai_unit_test_env["OPENAI_ORG_ID"], } openai_responses_client = OpenAIResponsesClient.from_dict(settings) dumped_settings = openai_responses_client.to_dict() - assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"] - assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"] + assert dumped_settings["model_id"] == openai_unit_test_env["OPENAI_RESPONSES_MODEL_ID"] assert dumped_settings["org_id"] == openai_unit_test_env["OPENAI_ORG_ID"] # Assert that the 'User-Agent' header is not present in the dumped_settings default headers - assert "User-Agent" not in dumped_settings["default_headers"] + assert "User-Agent" not in dumped_settings.get("default_headers", {}) def test_get_response_with_invalid_input() -> None: """Test get_response with invalid inputs to trigger exception handling.""" - client = OpenAIResponsesClient(ai_model_id="invalid-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="invalid-model", api_key="test-key") # Test with empty messages which should trigger ServiceInvalidRequestError with pytest.raises(ServiceInvalidRequestError, match="Messages are required"): @@ -200,7 +198,7 @@ def test_get_response_with_invalid_input() -> None: def test_get_response_with_all_parameters() -> None: """Test get_response with all possible parameters to cover parameter handling logic.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test with comprehensive parameter set - should fail due to invalid API key with pytest.raises(ServiceResponseException): @@ -232,7 +230,7 @@ def test_get_response_with_all_parameters() -> None: def test_web_search_tool_with_location() -> None: """Test HostedWebSearchTool with location parameters.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test web search tool with location web_search_tool = HostedWebSearchTool( @@ -254,7 +252,7 @@ def test_web_search_tool_with_location() -> None: def test_file_search_tool_with_invalid_inputs() -> None: """Test HostedFileSearchTool with invalid vector store inputs.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test with invalid inputs type (should trigger ValueError) file_search_tool = HostedFileSearchTool(inputs=[HostedFileContent(file_id="invalid")]) @@ -268,7 +266,7 @@ def test_file_search_tool_with_invalid_inputs() -> None: def test_code_interpreter_tool_variations() -> None: """Test HostedCodeInterpreterTool with and without file inputs.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test code interpreter without files code_tool_empty = HostedCodeInterpreterTool() @@ -293,7 +291,7 @@ def test_code_interpreter_tool_variations() -> None: def test_content_filter_exception() -> None: """Test that content filter errors in get_response are properly handled.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Mock a BadRequestError with content_filter code mock_error = BadRequestError( @@ -313,7 +311,7 @@ def test_content_filter_exception() -> None: def test_hosted_file_search_tool_validation() -> None: """Test get_response HostedFileSearchTool validation.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test HostedFileSearchTool without inputs (should raise ValueError) empty_file_search_tool = HostedFileSearchTool() @@ -326,7 +324,7 @@ def test_hosted_file_search_tool_validation() -> None: def test_chat_message_parsing_with_function_calls() -> None: """Test get_response message preparation with function call and result content types in conversation flow.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Create messages with function call and result content function_call = FunctionCallContent( @@ -351,7 +349,7 @@ def test_chat_message_parsing_with_function_calls() -> None: async def test_response_format_parse_path() -> None: """Test get_response response_format parsing path.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Mock successful parse response mock_parsed_response = MagicMock() @@ -375,7 +373,7 @@ async def test_response_format_parse_path() -> None: async def test_bad_request_error_non_content_filter() -> None: """Test get_response BadRequestError without content_filter.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Mock a BadRequestError without content_filter code mock_error = BadRequestError( @@ -396,7 +394,7 @@ async def test_bad_request_error_non_content_filter() -> None: async def test_streaming_content_filter_exception_handling() -> None: """Test that content filter errors in get_streaming_response are properly handled.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Mock the OpenAI client to raise a BadRequestError with content_filter code with patch.object(client.client.responses, "create") as mock_create: @@ -413,10 +411,11 @@ async def test_streaming_content_filter_exception_handling() -> None: break +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_get_streaming_response_with_all_parameters() -> None: """Test get_streaming_response with all possible parameters.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Should fail due to invalid API key with pytest.raises(ServiceResponseException): @@ -449,7 +448,7 @@ async def test_get_streaming_response_with_all_parameters() -> None: def test_response_content_creation_with_annotations() -> None: """Test _create_response_content with different annotation types.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Create a mock response with annotated text content mock_response = MagicMock() @@ -489,7 +488,7 @@ def test_response_content_creation_with_annotations() -> None: def test_response_content_creation_with_refusal() -> None: """Test _create_response_content with refusal content.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Create a mock response with refusal content mock_response = MagicMock() @@ -519,7 +518,7 @@ def test_response_content_creation_with_refusal() -> None: def test_response_content_creation_with_reasoning() -> None: """Test _create_response_content with reasoning content.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Create a mock response with reasoning content mock_response = MagicMock() @@ -550,7 +549,7 @@ def test_response_content_creation_with_reasoning() -> None: def test_response_content_creation_with_code_interpreter() -> None: """Test _create_response_content with code interpreter outputs.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Create a mock response with code interpreter outputs mock_response = MagicMock() @@ -588,7 +587,7 @@ def test_response_content_creation_with_code_interpreter() -> None: def test_response_content_creation_with_function_call() -> None: """Test _create_response_content with function call content.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Create a mock response with function call mock_response = MagicMock() @@ -620,7 +619,7 @@ def test_response_content_creation_with_function_call() -> None: def test_tools_to_response_tools_with_hosted_mcp() -> None: """Test that HostedMCPTool is converted to the correct response tool dict.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") tool = HostedMCPTool( name="My MCP", @@ -650,7 +649,7 @@ def test_tools_to_response_tools_with_hosted_mcp() -> None: def test_create_response_content_with_mcp_approval_request() -> None: """Test that a non-streaming mcp_approval_request is parsed into FunctionApprovalRequestContent.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") mock_response = MagicMock() mock_response.output_parsed = None @@ -681,7 +680,7 @@ def test_create_response_content_with_mcp_approval_request() -> None: def test_tools_to_response_tools_with_raw_image_generation() -> None: """Test that raw image_generation tool dict is handled correctly with parameter mapping.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test with raw tool dict using user-friendly parameter names tool = { @@ -710,7 +709,7 @@ def test_tools_to_response_tools_with_raw_image_generation() -> None: def test_tools_to_response_tools_with_raw_image_generation_openai_responses_params() -> None: """Test raw image_generation tool with OpenAI-specific parameters.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test with OpenAI-specific parameters tool = { @@ -742,7 +741,7 @@ def test_tools_to_response_tools_with_raw_image_generation_openai_responses_para def test_tools_to_response_tools_with_raw_image_generation_minimal() -> None: """Test raw image_generation tool with minimal configuration.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test with minimal parameters (just type) tool = {"type": "image_generation"} @@ -760,7 +759,7 @@ def test_tools_to_response_tools_with_raw_image_generation_minimal() -> None: def test_create_streaming_response_content_with_mcp_approval_request() -> None: """Test that a streaming mcp_approval_request event is parsed into FunctionApprovalRequestContent.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") chat_options = ChatOptions() function_call_ids: dict[int, tuple[str, str]] = {} @@ -787,7 +786,7 @@ def test_end_to_end_mcp_approval_flow(span_exporter) -> None: """End-to-end mocked test: model issues an mcp_approval_request, user approves, client sends mcp_approval_response. """ - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # First mocked response: model issues an mcp_approval_request mock_response1 = MagicMock() @@ -851,7 +850,7 @@ def test_end_to_end_mcp_approval_flow(span_exporter) -> None: def test_usage_details_basic() -> None: """Test _usage_details_from_openai without cached or reasoning tokens.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") mock_usage = MagicMock() mock_usage.input_tokens = 100 @@ -869,7 +868,7 @@ def test_usage_details_basic() -> None: def test_usage_details_with_cached_tokens() -> None: """Test _usage_details_from_openai with cached input tokens.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") mock_usage = MagicMock() mock_usage.input_tokens = 200 @@ -887,7 +886,7 @@ def test_usage_details_with_cached_tokens() -> None: def test_usage_details_with_reasoning_tokens() -> None: """Test _usage_details_from_openai with reasoning tokens.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") mock_usage = MagicMock() mock_usage.input_tokens = 150 @@ -905,7 +904,7 @@ def test_usage_details_with_reasoning_tokens() -> None: def test_get_metadata_from_response() -> None: """Test the _get_metadata_from_response method.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test with logprobs mock_output_with_logprobs = MagicMock() @@ -925,7 +924,7 @@ def test_get_metadata_from_response() -> None: def test_streaming_response_basic_structure() -> None: """Test that _create_streaming_response_content returns proper structure.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") chat_options = ChatOptions(store=True) function_call_ids: dict[int, tuple[str, str]] = {} @@ -942,6 +941,7 @@ def test_streaming_response_basic_structure() -> None: assert response.raw_representation is mock_event +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_response() -> None: """Test OpenAI chat completion responses.""" @@ -986,6 +986,7 @@ async def test_openai_responses_client_response() -> None: assert output.weather is not None +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_response_tools() -> None: """Test OpenAI chat completion responses.""" @@ -1025,6 +1026,7 @@ async def test_openai_responses_client_response_tools() -> None: assert "sunny" in output.weather.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_streaming() -> None: """Test OpenAI chat completion responses.""" @@ -1071,6 +1073,7 @@ async def test_openai_responses_client_streaming() -> None: assert output.weather is not None +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_streaming_tools() -> None: """Test OpenAI chat completion responses.""" @@ -1118,6 +1121,7 @@ async def test_openai_responses_client_streaming_tools() -> None: assert "sunny" in output.weather.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_web_search() -> None: openai_responses_client = OpenAIResponsesClient() @@ -1157,6 +1161,7 @@ async def test_openai_responses_client_web_search() -> None: assert response.text is not None +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_web_search_streaming() -> None: openai_responses_client = OpenAIResponsesClient() @@ -1210,6 +1215,7 @@ async def test_openai_responses_client_web_search_streaming() -> None: assert full_message is not None +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_file_search() -> None: openai_responses_client = OpenAIResponsesClient() @@ -1234,6 +1240,7 @@ async def test_openai_responses_client_file_search() -> None: assert "75" in response.text +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_streaming_file_search() -> None: openai_responses_client = OpenAIResponsesClient() @@ -1268,6 +1275,7 @@ async def test_openai_responses_client_streaming_file_search() -> None: assert "75" in full_message +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_basic_run(): """Test OpenAI Responses Client agent basic run functionality with OpenAIResponsesClient.""" @@ -1284,6 +1292,7 @@ async def test_openai_responses_client_agent_basic_run(): assert "hello world" in response.text.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_basic_run_streaming(): """Test OpenAI Responses Client agent basic streaming functionality with OpenAIResponsesClient.""" @@ -1301,6 +1310,7 @@ async def test_openai_responses_client_agent_basic_run_streaming(): assert "streaming response test" in full_text.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_thread_persistence(): """Test OpenAI Responses Client agent thread persistence across runs with OpenAIResponsesClient.""" @@ -1324,6 +1334,7 @@ async def test_openai_responses_client_agent_thread_persistence(): assert second_response.text is not None +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_thread_storage_with_store_true(): """Test OpenAI Responses Client agent with store=True to verify service_thread_id is returned.""" @@ -1355,6 +1366,7 @@ async def test_openai_responses_client_agent_thread_storage_with_store_true(): assert len(thread.service_thread_id) > 0 +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_existing_thread(): """Test OpenAI Responses Client agent with existing thread to continue conversations across agent instances.""" @@ -1389,6 +1401,7 @@ async def test_openai_responses_client_agent_existing_thread(): assert "photography" in second_response.text.lower() +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_hosted_code_interpreter_tool(): """Test OpenAI Responses Client agent with HostedCodeInterpreterTool through OpenAIResponsesClient.""" @@ -1410,6 +1423,7 @@ async def test_openai_responses_client_agent_hosted_code_interpreter_tool(): assert contains_relevant_content or len(response.text.strip()) > 10 +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_raw_image_generation_tool(): """Test OpenAI Responses Client agent with raw image_generation tool through OpenAIResponsesClient.""" @@ -1446,6 +1460,7 @@ async def test_openai_responses_client_agent_raw_image_generation_tool(): assert image_content_found, "Expected to find image content in response" +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_level_tool_persistence(): """Test that agent-level tools persist across multiple runs with OpenAI Responses Client.""" @@ -1472,6 +1487,7 @@ async def test_openai_responses_client_agent_level_tool_persistence(): assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"]) +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_run_level_tool_isolation(): """Test that run-level tools are isolated to specific runs and don't persist with OpenAI Responses Client.""" @@ -1511,6 +1527,7 @@ async def test_openai_responses_client_run_level_tool_isolation(): assert call_count == 1 +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_chat_options_run_level() -> None: """Integration test for comprehensive ChatOptions parameter coverage with OpenAI Response Agent.""" @@ -1534,6 +1551,7 @@ async def test_openai_responses_client_agent_chat_options_run_level() -> None: assert len(response.text) > 0 +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_chat_options_agent_level() -> None: """Integration test for comprehensive ChatOptions parameter coverage with OpenAI Response Agent.""" @@ -1557,6 +1575,7 @@ async def test_openai_responses_client_agent_chat_options_agent_level() -> None: assert len(response.text) > 0 +@pytest.mark.flaky @skip_if_openai_integration_tests_disabled async def test_openai_responses_client_agent_hosted_mcp_tool() -> None: """Integration test for HostedMCPTool with OpenAI Response Agent using Microsoft Learn MCP.""" @@ -1587,7 +1606,7 @@ async def test_openai_responses_client_agent_hosted_mcp_tool() -> None: def test_service_response_exception_includes_original_error_details() -> None: """Test that ServiceResponseException messages include original error details in the new format.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") messages = [ChatMessage(role="user", text="test message")] mock_response = MagicMock() @@ -1612,7 +1631,7 @@ def test_service_response_exception_includes_original_error_details() -> None: def test_get_streaming_response_with_response_format() -> None: """Test get_streaming_response with response_format.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") messages = [ChatMessage(role="user", text="Test streaming with format")] # It will fail due to invalid API key, but exercises the code path @@ -1627,7 +1646,7 @@ def test_get_streaming_response_with_response_format() -> None: def test_openai_content_parser_image_content() -> None: """Test _openai_content_parser with image content variations.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test image content with detail parameter and file_id image_content_with_detail = UriContent( @@ -1651,7 +1670,7 @@ def test_openai_content_parser_image_content() -> None: def test_openai_content_parser_audio_content() -> None: """Test _openai_content_parser with audio content variations.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test WAV audio content wav_content = UriContent(uri="data:audio/wav;base64,abc123", media_type="audio/wav") @@ -1669,7 +1688,7 @@ def test_openai_content_parser_audio_content() -> None: def test_openai_content_parser_unsupported_content() -> None: """Test _openai_content_parser with unsupported content types.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test unsupported audio format unsupported_audio = UriContent(uri="data:audio/ogg;base64,ghi789", media_type="audio/ogg") @@ -1684,7 +1703,7 @@ def test_openai_content_parser_unsupported_content() -> None: def test_create_streaming_response_content_code_interpreter() -> None: """Test _create_streaming_response_content with code_interpreter_call.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") chat_options = ChatOptions() function_call_ids: dict[int, tuple[str, str]] = {} @@ -1708,7 +1727,7 @@ def test_create_streaming_response_content_code_interpreter() -> None: def test_create_streaming_response_content_reasoning() -> None: """Test _create_streaming_response_content with reasoning content.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") chat_options = ChatOptions() function_call_ids: dict[int, tuple[str, str]] = {} @@ -1732,7 +1751,7 @@ def test_create_streaming_response_content_reasoning() -> None: def test_openai_content_parser_text_reasoning_comprehensive() -> None: """Test _openai_content_parser with TextReasoningContent all additional properties.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test TextReasoningContent with all additional properties comprehensive_reasoning = TextReasoningContent( @@ -1754,7 +1773,7 @@ def test_openai_content_parser_text_reasoning_comprehensive() -> None: def test_streaming_reasoning_text_delta_event() -> None: """Test reasoning text delta event creates TextReasoningContent.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") chat_options = ChatOptions() function_call_ids: dict[int, tuple[str, str]] = {} @@ -1779,7 +1798,7 @@ def test_streaming_reasoning_text_delta_event() -> None: def test_streaming_reasoning_text_done_event() -> None: """Test reasoning text done event creates TextReasoningContent with complete text.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") chat_options = ChatOptions() function_call_ids: dict[int, tuple[str, str]] = {} @@ -1805,7 +1824,7 @@ def test_streaming_reasoning_text_done_event() -> None: def test_streaming_reasoning_summary_text_delta_event() -> None: """Test reasoning summary text delta event creates TextReasoningContent.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") chat_options = ChatOptions() function_call_ids: dict[int, tuple[str, str]] = {} @@ -1830,7 +1849,7 @@ def test_streaming_reasoning_summary_text_delta_event() -> None: def test_streaming_reasoning_summary_text_done_event() -> None: """Test reasoning summary text done event creates TextReasoningContent with complete text.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") chat_options = ChatOptions() function_call_ids: dict[int, tuple[str, str]] = {} @@ -1856,7 +1875,7 @@ def test_streaming_reasoning_summary_text_done_event() -> None: def test_streaming_reasoning_events_preserve_metadata() -> None: """Test that reasoning events preserve metadata like regular text events.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") chat_options = ChatOptions() function_call_ids: dict[int, tuple[str, str]] = {} @@ -1894,7 +1913,7 @@ def test_streaming_reasoning_events_preserve_metadata() -> None: def test_create_response_content_image_generation_raw_base64(): """Test image generation response parsing with raw base64 string.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Create a mock response with raw base64 image data (PNG signature) mock_response = MagicMock() @@ -1928,7 +1947,7 @@ def test_create_response_content_image_generation_raw_base64(): def test_create_response_content_image_generation_existing_data_uri(): """Test image generation response parsing with existing data URI.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Create a mock response with existing data URI mock_response = MagicMock() @@ -1961,7 +1980,7 @@ def test_create_response_content_image_generation_existing_data_uri(): def test_create_response_content_image_generation_format_detection(): """Test different image format detection from base64 data.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Test JPEG detection jpeg_signature = b"\xff\xd8\xff" @@ -2014,7 +2033,7 @@ def test_create_response_content_image_generation_format_detection(): def test_create_response_content_image_generation_fallback(): """Test image generation with invalid base64 falls back to PNG.""" - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") # Create a mock response with invalid base64 mock_response = MagicMock() @@ -2046,7 +2065,7 @@ def test_create_response_content_image_generation_fallback(): def test_prepare_options_store_parameter_handling() -> None: - client = OpenAIResponsesClient(ai_model_id="test-model", api_key="test-key") + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") messages = [ChatMessage(role="user", text="Test message")] test_conversation_id = "test-conversation-123" diff --git a/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py b/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py index 5e5409a745..dee1a30c25 100644 --- a/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py +++ b/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py @@ -5,8 +5,6 @@ from dataclasses import dataclass, field from datetime import datetime, timezone from typing import Any -import pytest - from agent_framework._workflows._checkpoint import CheckpointStorage, WorkflowCheckpoint from agent_framework._workflows._events import RequestInfoEvent, WorkflowEvent from agent_framework._workflows._executor import ( @@ -295,7 +293,6 @@ def test_restore_state_falls_back_to_base_request_type() -> None: assert isinstance(restored.data, RequestInfoMessage) -@pytest.mark.asyncio async def test_run_persists_pending_requests_in_runner_state() -> None: shared_state = SharedState() runner_ctx = _StubRunnerContext() diff --git a/python/packages/core/tests/workflow/test_sequential.py b/python/packages/core/tests/workflow/test_sequential.py index d734e5d606..54df5b1638 100644 --- a/python/packages/core/tests/workflow/test_sequential.py +++ b/python/packages/core/tests/workflow/test_sequential.py @@ -117,7 +117,6 @@ async def test_sequential_with_custom_executor_summary() -> None: assert msgs[2].role == Role.ASSISTANT and msgs[2].text.startswith("Summary of users:") -@pytest.mark.asyncio async def test_sequential_checkpoint_resume_round_trip() -> None: storage = InMemoryCheckpointStorage() diff --git a/python/packages/devui/frontend/src/types/agent-framework.ts b/python/packages/devui/frontend/src/types/agent-framework.ts index da36247793..b00a590277 100644 --- a/python/packages/devui/frontend/src/types/agent-framework.ts +++ b/python/packages/devui/frontend/src/types/agent-framework.ts @@ -223,7 +223,7 @@ export interface ChatResponseUpdate { response_id?: string; message_id?: string; conversation_id?: string; - ai_model_id?: string; + model_id?: string; created_at?: CreatedAtT; finish_reason?: FinishReason; additional_properties?: Record; diff --git a/python/packages/devui/samples/in_memory_mode.py b/python/packages/devui/samples/in_memory_mode.py index 5260d20c7a..d89b1821bb 100644 --- a/python/packages/devui/samples/in_memory_mode.py +++ b/python/packages/devui/samples/in_memory_mode.py @@ -46,7 +46,7 @@ def main(): "You are a helpful weather and time assistant. Use the available tools to " "provide accurate weather information and current time for any location." ), - chat_client=OpenAIChatClient(ai_model_id="gpt-4o-mini"), + chat_client=OpenAIChatClient(model_id="gpt-4o-mini"), tools=[get_weather, get_time], ) @@ -54,7 +54,7 @@ def main(): name="general-assistant", description="A simple conversational agent", instructions="You are a helpful assistant.", - chat_client=OpenAIChatClient(ai_model_id="gpt-4o-mini"), + chat_client=OpenAIChatClient(model_id="gpt-4o-mini"), ) # Collect entities for serving diff --git a/python/packages/devui/samples/weather_agent/agent.py b/python/packages/devui/samples/weather_agent/agent.py index 7a75cc395b..57c4113274 100644 --- a/python/packages/devui/samples/weather_agent/agent.py +++ b/python/packages/devui/samples/weather_agent/agent.py @@ -42,7 +42,7 @@ agent = ChatAgent( and forecasts for any location. Always be helpful and provide detailed weather information when asked. """, - chat_client=OpenAIChatClient(ai_model_id=os.environ.get("OPENAI_CHAT_MODEL_ID", "gpt-4o")), + chat_client=OpenAIChatClient(model_id=os.environ.get("OPENAI_CHAT_MODEL_ID", "gpt-4o")), tools=[get_weather, get_forecast], ) diff --git a/python/packages/lab/lightning/README.md b/python/packages/lab/lightning/README.md index f7ed398270..71cb69ecbd 100644 --- a/python/packages/lab/lightning/README.md +++ b/python/packages/lab/lightning/README.md @@ -51,7 +51,7 @@ async def math_agent(task: TaskType, llm: LLM) -> float: MCPStdioTool(name="calculator", command="uvx", args=["mcp-server-calculator"]) as mcp_server, ChatAgent( chat_client=OpenAIChatClient( - ai_model_id=llm.model, + model_id=llm.model, api_key="your-api-key", base_url=llm.endpoint, ), diff --git a/python/packages/lab/lightning/samples/train_math_agent.py b/python/packages/lab/lightning/samples/train_math_agent.py index 4b6a64d57d..951448f2e2 100644 --- a/python/packages/lab/lightning/samples/train_math_agent.py +++ b/python/packages/lab/lightning/samples/train_math_agent.py @@ -168,7 +168,7 @@ async def math_agent(task: MathProblem, llm: LLM) -> float: MCPStdioTool(name="calculator", command="uvx", args=["mcp-server-calculator"]) as mcp_server, ChatAgent( chat_client=OpenAIChatClient( - ai_model_id=llm.model, # This is the model being trained + model_id=llm.model, # This is the model being trained api_key=os.getenv("OPENAI_API_KEY") or "dummy", # Can be dummy when connecting to training LLM base_url=llm.endpoint, # vLLM server endpoint provided by agent-lightning ), diff --git a/python/packages/lab/lightning/samples/train_tau2_agent.py b/python/packages/lab/lightning/samples/train_tau2_agent.py index c6c28f7ed8..9b35f05981 100644 --- a/python/packages/lab/lightning/samples/train_tau2_agent.py +++ b/python/packages/lab/lightning/samples/train_tau2_agent.py @@ -103,14 +103,14 @@ class Tau2Agent(LitAgent): assistant_chat_client = OpenAIChatClient( base_url=llm.endpoint, # vLLM endpoint for the model being trained api_key=openai_api_key, - ai_model_id=llm.model, # Model ID being trained + model_id=llm.model, # Model ID being trained ) # User simulator: uses a fixed, capable model for consistent simulation user_simulator_chat_client = OpenAIChatClient( base_url=openai_base_url, # External API endpoint api_key=openai_api_key, - ai_model_id="gpt-4.1", # Fixed model for user simulator + model_id="gpt-4.1", # Fixed model for user simulator ) try: diff --git a/python/packages/lab/lightning/tests/test_lightning.py b/python/packages/lab/lightning/tests/test_lightning.py index 37c34881c8..570368ee4b 100644 --- a/python/packages/lab/lightning/tests/test_lightning.py +++ b/python/packages/lab/lightning/tests/test_lightning.py @@ -18,10 +18,6 @@ from openai.types.chat import ChatCompletion, ChatCompletionMessage from openai.types.chat.chat_completion import Choice -def test_import(): - """Test that the module can be imported.""" - - @pytest.fixture def workflow_two_agents(): """Test a workflow with two OpenAI chat agents where first agent's result passes to second agent.""" @@ -111,14 +107,12 @@ def workflow_two_agents(): yield workflow -@pytest.mark.asyncio async def test_openai_workflow_two_agents(workflow_two_agents): events = await workflow_two_agents.run("Please analyze the quarterly sales data") assert "Based on the analysis 'Analyzed data shows trend upward', I recommend investing" in events.get_outputs() -@pytest.mark.asyncio async def test_observability(workflow_two_agents): r"""Expected trace tree: diff --git a/python/packages/lab/pyproject.toml b/python/packages/lab/pyproject.toml index b411663880..3b908b7d81 100644 --- a/python/packages/lab/pyproject.toml +++ b/python/packages/lab/pyproject.toml @@ -54,6 +54,27 @@ math = [ "sympy>=1.13.0", ] +[dependency-groups] +dev = [ + "uv>=0.8.2,<0.9.0", + "pre-commit >= 3.7", + "ruff>=0.11.8", + "pytest>=8.4.1", + "pytest-asyncio>=1.0.0", + "pytest-cov>=6.2.1", + "pytest-env>=1.1.5", + "pytest-xdist[psutil]>=3.8.0", + "pytest-timeout>=2.3.1", + "pytest-retry>=1", + "mypy>=1.16.1", + "pyright>=1.1.402", + #tasks + "poethepoet>=0.36.0", + "rich", + "tomli", + "tomli-w", +] + [project.scripts] gaia_viewer = "agent_framework_lab_gaia:viewer_main" lightning = "agent_framework_lab_lightning:main" diff --git a/python/packages/lab/tau2/README.md b/python/packages/lab/tau2/README.md index 22aacc6865..d98e45ade1 100644 --- a/python/packages/lab/tau2/README.md +++ b/python/packages/lab/tau2/README.md @@ -61,12 +61,12 @@ async def run_single_task(): assistant_client = OpenAIChatClient( base_url="https://api.openai.com/v1", api_key="your-api-key", - ai_model_id="gpt-4o" + model_id="gpt-4o" ) user_client = OpenAIChatClient( base_url="https://api.openai.com/v1", api_key="your-api-key", - ai_model_id="gpt-4o-mini" + model_id="gpt-4o-mini" ) # Get a task and run it diff --git a/python/packages/lab/tau2/samples/run_benchmark.py b/python/packages/lab/tau2/samples/run_benchmark.py index 53ced32af4..5c2ed14d5e 100644 --- a/python/packages/lab/tau2/samples/run_benchmark.py +++ b/python/packages/lab/tau2/samples/run_benchmark.py @@ -96,14 +96,14 @@ async def run_benchmark(assistant_model: str, user_model: str, debug_task_id: st assistant_chat_client = OpenAIChatClient( base_url=openai_base_url, api_key=openai_api_key, - ai_model_id=assistant_model, + model_id=assistant_model, ) # User simulator: simulates realistic customer behavior and requests user_simulator_chat_client = OpenAIChatClient( base_url=openai_base_url, api_key=openai_api_key, - ai_model_id=user_model, + model_id=user_model, ) # STEP 4: Filter task set for debug mode @@ -133,8 +133,8 @@ async def run_benchmark(assistant_model: str, user_model: str, debug_task_id: st # Initialize result structure for this task result: dict[str, Any] = { "config": { - "assistant": assistant_chat_client.ai_model_id, - "user": user_simulator_chat_client.ai_model_id, + "assistant": assistant_chat_client.model_id, + "user": user_simulator_chat_client.model_id, }, "task": task, } @@ -183,8 +183,8 @@ async def run_benchmark(assistant_model: str, user_model: str, debug_task_id: st # Initialize result structure for this task result: dict[str, Any] = { "config": { - "assistant": assistant_chat_client.ai_model_id, - "user": user_simulator_chat_client.ai_model_id, + "assistant": assistant_chat_client.model_id, + "user": user_simulator_chat_client.model_id, }, "task": task, } diff --git a/python/packages/redis/agent_framework_redis/_chat_message_store.py b/python/packages/redis/agent_framework_redis/_chat_message_store.py index 5243db5198..539d9e1604 100644 --- a/python/packages/redis/agent_framework_redis/_chat_message_store.py +++ b/python/packages/redis/agent_framework_redis/_chat_message_store.py @@ -9,10 +9,10 @@ from uuid import uuid4 import redis.asyncio as redis from agent_framework import ChatMessage -from agent_framework._pydantic import AFBaseModel +from pydantic import BaseModel -class RedisStoreState(AFBaseModel): +class RedisStoreState(BaseModel): """State model for serializing and deserializing Redis chat message store data.""" thread_id: str diff --git a/python/packages/redis/tests/__init__.py b/python/packages/redis/tests/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/python/pyproject.toml b/python/pyproject.toml index 6dc1abea59..c3a1d91938 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -52,7 +52,6 @@ environments = [ [tool.uv.workspace] members = [ "packages/*" ] -exclude = [ "packages/agent_framework_project.egg-info" ] [tool.uv.sources] agent-framework = { workspace = true } @@ -128,7 +127,8 @@ notice-rgx = "^# Copyright \\(c\\) Microsoft\\. All rights reserved\\." min-file-size = 1 [tool.pytest.ini_options] -testpaths = '**/tests' +testpaths = 'packages/**/tests' +norecursedirs = '**/lab/**' addopts = "-ra -q -r fEX" asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" diff --git a/python/samples/getting_started/agents/anthropic/anthropic_with_openai_chat_client.py b/python/samples/getting_started/agents/anthropic/anthropic_with_openai_chat_client.py index 8c1b55faec..3fbf8ff3fd 100644 --- a/python/samples/getting_started/agents/anthropic/anthropic_with_openai_chat_client.py +++ b/python/samples/getting_started/agents/anthropic/anthropic_with_openai_chat_client.py @@ -38,7 +38,7 @@ async def non_streaming_example() -> None: agent = OpenAIChatClient( api_key=os.getenv("ANTHROPIC_API_KEY"), base_url="https://api.anthropic.com/v1/", - ai_model_id=os.getenv("ANTHROPIC_MODEL"), + model_id=os.getenv("ANTHROPIC_MODEL"), ).create_agent( name="WeatherAgent", instructions="You are a helpful weather agent.", @@ -58,7 +58,7 @@ async def streaming_example() -> None: agent = OpenAIChatClient( api_key=os.getenv("ANTHROPIC_API_KEY"), base_url="https://api.anthropic.com/v1/", - ai_model_id=os.getenv("ANTHROPIC_MODEL"), + model_id=os.getenv("ANTHROPIC_MODEL"), ).create_agent( name="WeatherAgent", instructions="You are a helpful weather agent.", diff --git a/python/samples/getting_started/agents/custom/custom_chat_client.py b/python/samples/getting_started/agents/custom/custom_chat_client.py index 89a0bfcae3..df1c59460a 100644 --- a/python/samples/getting_started/agents/custom/custom_chat_client.py +++ b/python/samples/getting_started/agents/custom/custom_chat_client.py @@ -3,7 +3,7 @@ import asyncio import random from collections.abc import AsyncIterable, MutableSequence -from typing import Any +from typing import Any, ClassVar from agent_framework import ( BaseChatClient, @@ -46,9 +46,7 @@ class EchoingChatClient(BaseChatClient): and implementing the required _inner_get_response() and _inner_get_streaming_response() methods. """ - OTEL_PROVIDER_NAME: str = "EchoingChatClient" - - prefix: str = "Echo:" + OTEL_PROVIDER_NAME: ClassVar[str] = "EchoingChatClient" def __init__(self, *, prefix: str = "Echo:", **kwargs: Any) -> None: """Initialize the EchoingChatClient. @@ -57,10 +55,8 @@ class EchoingChatClient(BaseChatClient): prefix: Prefix to add to echoed messages. **kwargs: Additional keyword arguments passed to BaseChatClient. """ - super().__init__( - prefix=prefix, # type: ignore - **kwargs, - ) + super().__init__(**kwargs) + self.prefix = prefix async def _inner_get_response( self, diff --git a/python/samples/getting_started/agents/openai/openai_assistants_with_explicit_settings.py b/python/samples/getting_started/agents/openai/openai_assistants_with_explicit_settings.py index 203debe1e3..bff52fe8d2 100644 --- a/python/samples/getting_started/agents/openai/openai_assistants_with_explicit_settings.py +++ b/python/samples/getting_started/agents/openai/openai_assistants_with_explicit_settings.py @@ -21,7 +21,7 @@ async def main() -> None: print("=== OpenAI Assistants Client with Explicit Settings ===") async with OpenAIAssistantsClient( - ai_model_id=os.environ["OPENAI_CHAT_MODEL_ID"], + model_id=os.environ["OPENAI_CHAT_MODEL_ID"], api_key=os.environ["OPENAI_API_KEY"], ).create_agent( instructions="You are a helpful weather agent.", diff --git a/python/samples/getting_started/agents/openai/openai_chat_client_with_explicit_settings.py b/python/samples/getting_started/agents/openai/openai_chat_client_with_explicit_settings.py index ba773b7809..a870798425 100644 --- a/python/samples/getting_started/agents/openai/openai_chat_client_with_explicit_settings.py +++ b/python/samples/getting_started/agents/openai/openai_chat_client_with_explicit_settings.py @@ -21,7 +21,7 @@ async def main() -> None: print("=== OpenAI Chat Client with Explicit Settings ===") agent = OpenAIChatClient( - ai_model_id=os.environ["OPENAI_CHAT_MODEL_ID"], + model_id=os.environ["OPENAI_CHAT_MODEL_ID"], api_key=os.environ["OPENAI_API_KEY"], ).create_agent( instructions="You are a helpful weather agent.", diff --git a/python/samples/getting_started/agents/openai/openai_chat_client_with_web_search.py b/python/samples/getting_started/agents/openai/openai_chat_client_with_web_search.py index f83ca14aba..6b8ea24740 100644 --- a/python/samples/getting_started/agents/openai/openai_chat_client_with_web_search.py +++ b/python/samples/getting_started/agents/openai/openai_chat_client_with_web_search.py @@ -7,7 +7,7 @@ from agent_framework.openai import OpenAIChatClient async def main() -> None: - client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview") + client = OpenAIChatClient(model_id="gpt-4o-search-preview") message = "What is the current weather? Do not ask for my current location." # Test that the client will use the web search tool with location diff --git a/python/samples/getting_started/agents/openai/openai_responses_client_reasoning.py b/python/samples/getting_started/agents/openai/openai_responses_client_reasoning.py index f9dfb34f16..56c2d03e4b 100644 --- a/python/samples/getting_started/agents/openai/openai_responses_client_reasoning.py +++ b/python/samples/getting_started/agents/openai/openai_responses_client_reasoning.py @@ -10,7 +10,7 @@ async def reasoning_example() -> None: """Example of reasoning response (get results as they are generated).""" print("=== Reasoning Example ===") - agent = OpenAIResponsesClient(ai_model_id="gpt-5").create_agent( + agent = OpenAIResponsesClient(model_id="gpt-5").create_agent( name="MathHelper", instructions="You are a personal math tutor. When asked a math question, " "write and run code using the python tool to answer the question.", diff --git a/python/samples/getting_started/agents/openai/openai_responses_client_with_explicit_settings.py b/python/samples/getting_started/agents/openai/openai_responses_client_with_explicit_settings.py index 19b785c91c..f17df3a409 100644 --- a/python/samples/getting_started/agents/openai/openai_responses_client_with_explicit_settings.py +++ b/python/samples/getting_started/agents/openai/openai_responses_client_with_explicit_settings.py @@ -21,7 +21,7 @@ async def main() -> None: print("=== OpenAI Responses Client with Explicit Settings ===") agent = OpenAIResponsesClient( - ai_model_id=os.environ["OPENAI_RESPONSES_MODEL_ID"], + model_id=os.environ["OPENAI_RESPONSES_MODEL_ID"], api_key=os.environ["OPENAI_API_KEY"], ).create_agent( instructions="You are a helpful weather agent.", diff --git a/python/samples/getting_started/chat_client/chat_response_cancellation.py b/python/samples/getting_started/chat_client/chat_response_cancellation.py index fcd43d2f0b..afbf246bfe 100644 --- a/python/samples/getting_started/chat_client/chat_response_cancellation.py +++ b/python/samples/getting_started/chat_client/chat_response_cancellation.py @@ -11,7 +11,7 @@ async def main() -> None: Creates a task for the chat request, waits briefly, then cancels it to show proper cleanup. Configuration: - - OpenAI model ID: Use "ai_model_id" parameter or "OPENAI_CHAT_MODEL_ID" environment variable + - OpenAI model ID: Use "model_id" parameter or "OPENAI_CHAT_MODEL_ID" environment variable - OpenAI API key: Use "api_key" parameter or "OPENAI_API_KEY" environment variable """ chat_client = OpenAIChatClient() diff --git a/python/samples/getting_started/context_providers/redis/redis_basics.py b/python/samples/getting_started/context_providers/redis/redis_basics.py index e0a649f333..ffa8c32a60 100644 --- a/python/samples/getting_started/context_providers/redis/redis_basics.py +++ b/python/samples/getting_started/context_providers/redis/redis_basics.py @@ -175,7 +175,7 @@ async def main() -> None: ) # Create chat client for the agent - client = OpenAIChatClient(ai_model_id=os.getenv("OPENAI_CHAT_MODEL_ID"), api_key=os.getenv("OPENAI_API_KEY")) + client = OpenAIChatClient(model_id=os.getenv("OPENAI_CHAT_MODEL_ID"), api_key=os.getenv("OPENAI_API_KEY")) # Create agent wired to the Redis context provider. The provider automatically # persists conversational details and surfaces relevant context on each turn. agent = client.create_agent( @@ -219,7 +219,7 @@ async def main() -> None: # Create agent exposing the flight search tool. Tool outputs are captured by the # provider and become retrievable context for later turns. - client = OpenAIChatClient(ai_model_id=os.getenv("OPENAI_CHAT_MODEL_ID"), api_key=os.getenv("OPENAI_API_KEY")) + client = OpenAIChatClient(model_id=os.getenv("OPENAI_CHAT_MODEL_ID"), api_key=os.getenv("OPENAI_API_KEY")) agent = client.create_agent( name="MemoryEnhancedAssistant", instructions=( diff --git a/python/samples/getting_started/context_providers/redis/redis_conversation.py b/python/samples/getting_started/context_providers/redis/redis_conversation.py index 6972d2ad6d..1ca54a4ae6 100644 --- a/python/samples/getting_started/context_providers/redis/redis_conversation.py +++ b/python/samples/getting_started/context_providers/redis/redis_conversation.py @@ -60,7 +60,7 @@ async def main() -> None: ) # Create chat client for the agent - client = OpenAIChatClient(ai_model_id=os.getenv("OPENAI_CHAT_MODEL_ID"), api_key=os.getenv("OPENAI_API_KEY")) + client = OpenAIChatClient(model_id=os.getenv("OPENAI_CHAT_MODEL_ID"), api_key=os.getenv("OPENAI_API_KEY")) # Create agent wired to the Redis context provider. The provider automatically # persists conversational details and surfaces relevant context on each turn. agent = client.create_agent( diff --git a/python/samples/getting_started/context_providers/redis/redis_threads.py b/python/samples/getting_started/context_providers/redis/redis_threads.py index 057b09b858..6a9022895c 100644 --- a/python/samples/getting_started/context_providers/redis/redis_threads.py +++ b/python/samples/getting_started/context_providers/redis/redis_threads.py @@ -47,7 +47,7 @@ async def example_global_thread_scope() -> None: global_thread_id = str(uuid.uuid4()) client = OpenAIChatClient( - ai_model_id=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"), + model_id=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"), api_key=os.getenv("OPENAI_API_KEY"), ) @@ -100,7 +100,7 @@ async def example_per_operation_thread_scope() -> None: print("-" * 40) client = OpenAIChatClient( - ai_model_id=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"), + model_id=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"), api_key=os.getenv("OPENAI_API_KEY"), ) @@ -168,7 +168,7 @@ async def example_multiple_agents() -> None: print("-" * 40) client = OpenAIChatClient( - ai_model_id=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"), + model_id=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"), api_key=os.getenv("OPENAI_API_KEY"), ) diff --git a/python/samples/getting_started/multimodal_input/openai_chat_multimodal.py b/python/samples/getting_started/multimodal_input/openai_chat_multimodal.py index 39af7a7575..c47cf20bcc 100644 --- a/python/samples/getting_started/multimodal_input/openai_chat_multimodal.py +++ b/python/samples/getting_started/multimodal_input/openai_chat_multimodal.py @@ -11,7 +11,7 @@ from agent_framework.openai import OpenAIChatClient async def test_image() -> None: """Test image analysis with OpenAI.""" - client = OpenAIChatClient(ai_model_id="gpt-4o") + client = OpenAIChatClient(model_id="gpt-4o") # Fetch image from httpbin image_url = "https://httpbin.org/image/jpeg" @@ -30,7 +30,7 @@ async def test_image() -> None: async def test_audio() -> None: """Test audio analysis with OpenAI.""" - client = OpenAIChatClient(ai_model_id="gpt-4o-audio-preview") + client = OpenAIChatClient(model_id="gpt-4o-audio-preview") # Create minimal WAV file (0.1 seconds of silence) wav_header = ( diff --git a/python/samples/getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py b/python/samples/getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py index 8f73c213a6..3c869439c2 100644 --- a/python/samples/getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py +++ b/python/samples/getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py @@ -102,7 +102,7 @@ async def main() -> None: # Create executors for the workflow. print("Creating chat client and executors...") - mini_chat_client = OpenAIChatClient(ai_model_id="gpt-4.1-nano") + mini_chat_client = OpenAIChatClient(model_id="gpt-4.1-nano") worker = Worker(id="sub-worker", chat_client=mini_chat_client) request_info_executor = RequestInfoExecutor(id="request_info") reviewer = ReviewerWithHumanInTheLoop(worker_id=worker.id, request_info_id=request_info_executor.id) diff --git a/python/samples/getting_started/workflows/agents/workflow_as_agent_reflection_pattern.py b/python/samples/getting_started/workflows/agents/workflow_as_agent_reflection_pattern.py index 775c927af9..f8840845ac 100644 --- a/python/samples/getting_started/workflows/agents/workflow_as_agent_reflection_pattern.py +++ b/python/samples/getting_started/workflows/agents/workflow_as_agent_reflection_pattern.py @@ -197,8 +197,8 @@ async def main() -> None: # Initialize chat clients and executors. print("Creating chat client and executors...") - mini_chat_client = OpenAIChatClient(ai_model_id="gpt-4.1-nano") - chat_client = OpenAIChatClient(ai_model_id="gpt-4.1") + mini_chat_client = OpenAIChatClient(model_id="gpt-4.1-nano") + chat_client = OpenAIChatClient(model_id="gpt-4.1") reviewer = Reviewer(id="reviewer", chat_client=chat_client) worker = Worker(id="worker", chat_client=mini_chat_client) diff --git a/python/samples/getting_started/workflows/orchestration/magentic.py b/python/samples/getting_started/workflows/orchestration/magentic.py index 14e00b0418..95038cd0e4 100644 --- a/python/samples/getting_started/workflows/orchestration/magentic.py +++ b/python/samples/getting_started/workflows/orchestration/magentic.py @@ -55,7 +55,7 @@ async def main() -> None: # This agent requires the gpt-4o-search-preview model to perform web searches. # Feel free to explore with other agents that support web search, for example, # the `OpenAIResponseAgent` or `AzureAgentProtocol` with bing grounding. - chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"), + chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"), ) coder_agent = ChatAgent( diff --git a/python/samples/getting_started/workflows/orchestration/magentic_human_plan_update.py b/python/samples/getting_started/workflows/orchestration/magentic_human_plan_update.py index fb45f94855..339554d3ec 100644 --- a/python/samples/getting_started/workflows/orchestration/magentic_human_plan_update.py +++ b/python/samples/getting_started/workflows/orchestration/magentic_human_plan_update.py @@ -58,7 +58,7 @@ async def main() -> None: # This agent requires the gpt-4o-search-preview model to perform web searches. # Feel free to explore with other agents that support web search, for example, # the `OpenAIResponseAgent` or `AzureAgentProtocol` with bing grounding. - chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"), + chat_client=OpenAIChatClient(model_id="gpt-4o-search-preview"), ) coder_agent = ChatAgent( diff --git a/python/tests/__init__.py b/python/tests/__init__.py deleted file mode 100644 index 58b29ef7f3..0000000000 --- a/python/tests/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Test utilities for sample testing.""" diff --git a/python/tests/sample_utils.py b/python/tests/sample_utils.py deleted file mode 100644 index 878566e223..0000000000 --- a/python/tests/sample_utils.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -import logging -from collections.abc import Awaitable, Callable -from typing import Any - -logger = logging.getLogger(__name__) - - -async def retry( - func: Callable[[], Awaitable[Any]], - retries: int = 3, - reset: Callable[[], None] | None = None, - name: str | None = None, -) -> None: - """Retry function with reset capability and proper logging. - - Args: - func: The function to retry. - retries: Number of retries. - reset: Function to reset the state of any variables used in the function. - name: Optional name for logging purposes. - """ - func_name = name or func.__module__ - logger.info(f"Running {retries} retries with func: {func_name}") - - for i in range(retries): - logger.info(f" Try {i + 1} for {func_name}") - try: - if reset: - reset() - await func() - return - except Exception as e: - logger.warning(f" On try {i + 1} got this error: {e}") - if i == retries - 1: # Last retry - raise - - # Binary exponential backoff like Semantic Kernel - backoff = 2**i - logger.info(f" Sleeping for {backoff} seconds before retrying") - await asyncio.sleep(backoff) diff --git a/python/tests/samples/__init__.py b/python/tests/samples/__init__.py deleted file mode 100644 index 1857030e0e..0000000000 --- a/python/tests/samples/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Sample tests package.""" diff --git a/python/tests/samples/getting_started/__init__.py b/python/tests/samples/getting_started/__init__.py deleted file mode 100644 index 9c75ed9a9b..0000000000 --- a/python/tests/samples/getting_started/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Getting started sample tests.""" diff --git a/python/tests/samples/getting_started/test_agents.py b/python/tests/samples/getting_started/test_agent_samples.py similarity index 99% rename from python/tests/samples/getting_started/test_agents.py rename to python/tests/samples/getting_started/test_agent_samples.py index 73e90a2330..710a5604fc 100644 --- a/python/tests/samples/getting_started/test_agents.py +++ b/python/tests/samples/getting_started/test_agent_samples.py @@ -146,7 +146,6 @@ from samples.getting_started.agents.openai.openai_responses_client_with_thread i from samples.getting_started.agents.openai.openai_responses_client_with_web_search import ( main as openai_responses_client_with_web_search, ) -from tests.sample_utils import retry # Environment variable for controlling sample tests RUN_SAMPLES_TESTS = "RUN_SAMPLES_TESTS" @@ -579,6 +578,7 @@ agent_samples = [ ] +@pytest.mark.flaky @mark.parametrize("sample, responses", agent_samples) async def test_agent_samples(sample: Callable[..., Awaitable[Any]], responses: list[str], monkeypatch: MonkeyPatch): """Test agent samples with input mocking and retry logic.""" @@ -592,4 +592,4 @@ async def test_agent_samples(sample: Callable[..., Awaitable[Any]], responses: l return responses.pop(0) if responses else "exit" monkeypatch.setattr("builtins.input", mock_input) - await retry(sample, retries=3, reset=reset) + await sample diff --git a/python/tests/samples/getting_started/test_chat_client.py b/python/tests/samples/getting_started/test_chat_client_samples.py similarity index 98% rename from python/tests/samples/getting_started/test_chat_client.py rename to python/tests/samples/getting_started/test_chat_client_samples.py index 740ea919b4..0a699c5908 100644 --- a/python/tests/samples/getting_started/test_chat_client.py +++ b/python/tests/samples/getting_started/test_chat_client_samples.py @@ -32,7 +32,6 @@ from samples.getting_started.chat_client.openai_chat_client import ( from samples.getting_started.chat_client.openai_responses_client import ( main as openai_responses_client, ) -from tests.sample_utils import retry # Environment variable for controlling sample tests RUN_SAMPLES_TESTS = "RUN_SAMPLES_TESTS" @@ -132,4 +131,4 @@ async def test_chat_client_samples( return responses.pop(0) if responses else "exit" monkeypatch.setattr("builtins.input", mock_input) - await retry(sample, retries=3, reset=reset) + await sample diff --git a/python/tests/samples/getting_started/test_threads.py b/python/tests/samples/getting_started/test_threads_samples.py similarity index 95% rename from python/tests/samples/getting_started/test_threads.py rename to python/tests/samples/getting_started/test_threads_samples.py index ec530efa6e..51c9103c39 100644 --- a/python/tests/samples/getting_started/test_threads.py +++ b/python/tests/samples/getting_started/test_threads_samples.py @@ -10,7 +10,6 @@ from pytest import MonkeyPatch, mark, param from samples.getting_started.threads.custom_chat_message_store_thread import main as threads_custom_store from samples.getting_started.threads.suspend_resume_thread import main as threads_suspend_resume -from tests.sample_utils import retry # Environment variable for controlling sample tests RUN_SAMPLES_TESTS = "RUN_SAMPLES_TESTS" @@ -51,4 +50,4 @@ async def test_thread_samples(sample: Callable[..., Awaitable[Any]], responses: return responses.pop(0) if responses else "exit" monkeypatch.setattr("builtins.input", mock_input) - await retry(sample, retries=3, reset=reset) + await sample diff --git a/python/uv.lock b/python/uv.lock index 863b83b575..7791e89e55 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -294,6 +294,26 @@ tau2 = [ { name = "tiktoken", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] +[package.dev-dependencies] +dev = [ + { name = "mypy", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "poethepoet", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pre-commit", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pyright", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pytest", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pytest-asyncio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pytest-cov", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pytest-env", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pytest-retry", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pytest-timeout", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pytest-xdist", extra = ["psutil"], marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "rich", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "ruff", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "tomli", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "tomli-w", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "uv", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, +] + [package.metadata] requires-dist = [ { name = "agent-framework-core", editable = "packages/core" }, @@ -312,6 +332,26 @@ requires-dist = [ ] provides-extras = ["gaia", "lightning", "tau2", "math"] +[package.metadata.requires-dev] +dev = [ + { name = "mypy", specifier = ">=1.16.1" }, + { name = "poethepoet", specifier = ">=0.36.0" }, + { name = "pre-commit", specifier = ">=3.7" }, + { name = "pyright", specifier = ">=1.1.402" }, + { name = "pytest", specifier = ">=8.4.1" }, + { name = "pytest-asyncio", specifier = ">=1.0.0" }, + { name = "pytest-cov", specifier = ">=6.2.1" }, + { name = "pytest-env", specifier = ">=1.1.5" }, + { name = "pytest-retry", specifier = ">=1" }, + { name = "pytest-timeout", specifier = ">=2.3.1" }, + { name = "pytest-xdist", extras = ["psutil"], specifier = ">=3.8.0" }, + { name = "rich" }, + { name = "ruff", specifier = ">=0.11.8" }, + { name = "tomli" }, + { name = "tomli-w" }, + { name = "uv", specifier = ">=0.8.2,<0.9.0" }, +] + [[package]] name = "agent-framework-mem0" version = "1.0.0b251001"