Python: [BREAKING] added SerializationMixin and applied to contents, agents, chat client… (#1012)

* added SerializationMixin and applied to contents, agents, chat clients, removed AFBaseModel

* fix annotations type

* mypy fixes

* fix tests

* fix serializable subvalues and added large docstring

* updated indents in code block

* fixed exported urls

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]:

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:

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."""

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,

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 {},

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]:

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,
     )
 

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")

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"(?<!^)(?=[A-Z])")
+
+
+@runtime_checkable
+class SerializationProtocol(Protocol):
+    """Protocol for objects that support serialization and deserialization."""
+
+    def to_dict(self, **kwargs: Any) -> 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:
+    - "<type>.<parameter>"
+    - "<type>.<dict-parameter>.<key>"
+    where <type> 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),
+    <parameter> is the name of the parameter in the __init__ method,
+    <dict-parameter> is the name of a parameter that is a dict,
+    and <key> 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 "<type>.<parameter>" or "<type>.<dict-parameter>.<key>".
+
+        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: <type>.<parameter>
+                kwargs[param_name] = dep_value
+            elif len(parts) == 3:
+                # Dict parameter: <type>.<dict-parameter>.<key>
+                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 "<type>.<parameter>" or "<type>.<dict-parameter>.<key>".
+
+        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()

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")
 
 

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)

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,
         )

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

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"),
-        )

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)

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)

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

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

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)

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."""

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."""

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:

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(

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."""

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

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."""

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)

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."""

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()

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"

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()

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()
 

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<string, unknown>;

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

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],
 )
 

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,
             ),

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
             ),

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:

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:
 

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"

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

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,
             }

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