Python telemetry (#223)

* initial work on telemetry

* moved tool operation const

* missing quotes

* working otel with samples

* updated readme and other assets

* added tests

* added tests

* small updates

* updated genaiattributes docs

* updated tests

* additional warning

* cleanup of tests

python/packages/main/tests/main/conftest.py

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft. All rights reserved.
+from typing import Any
+
+from pydantic import BaseModel
+from pytest import fixture
+
+from agent_framework import AITool, ChatMessage, ai_function
+from agent_framework.telemetry import ModelDiagnosticSettings
+
+
+@fixture(scope="function")
+def chat_history() -> list[ChatMessage]:
+    return []
+
+
+@fixture
+def ai_tool() -> AITool:
+    """Returns a generic AITool."""
+
+    class GenericTool(BaseModel):
+        name: str
+        description: str | None = None
+        additional_properties: dict[str, Any] | None = None
+
+        def parameters(self) -> dict[str, Any]:
+            """Return the parameters of the tool as a JSON schema."""
+            return {
+                "name": {"type": "string"},
+            }
+
+    return GenericTool(name="generic_tool", description="A generic tool")
+
+
+@fixture
+def ai_function_tool() -> AITool:
+    """Returns a executable AITool."""
+
+    @ai_function
+    def simple_function(x: int, y: int) -> int:
+        """A simple function that adds two numbers."""
+        return x + y
+
+    return simple_function
+
+
+@fixture
+def model_diagnostic_settings(monkeypatch, request) -> ModelDiagnosticSettings:
+    """Fixture to set environment variables for ModelDiagnosticSettings."""
+    enabled = getattr(request, "param", (None, None))[0]
+    sensitive = getattr(request, "param", (None, None))[1]
+    if enabled is None:
+        monkeypatch.delenv("AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS", raising=False)
+    else:
+        monkeypatch.setenv("AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS", str(enabled).lower())
+    if sensitive is None:
+        monkeypatch.delenv("AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE", raising=False)
+    else:
+        monkeypatch.setenv("AGENT_FRAMEWORK_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE", str(sensitive).lower())
+    return ModelDiagnosticSettings(env_file_path="test.env")

python/packages/main/tests/main/test_agents.py

@@ -0,0 +1,288 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import AsyncIterable, MutableSequence, Sequence
+from typing import Any
+from uuid import uuid4
+
+from pytest import fixture, raises
+
+from agent_framework import (
+    Agent,
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    AgentThread,
+    ChatClient,
+    ChatClientAgent,
+    ChatClientAgentThread,
+    ChatClientAgentThreadType,
+    ChatClientBase,
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    ChatRole,
+    TextContent,
+)
+from agent_framework.exceptions import AgentExecutionException
+
+
+# Mock AgentThread implementation for testing
+class MockAgentThread(AgentThread):
+    async def _on_new_messages(self, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        pass
+
+
+# Mock Agent implementation for testing
+class MockAgent(Agent):
+    @property
+    def id(self) -> str:
+        return str(uuid4())
+
+    @property
+    def name(self) -> str | None:
+        """Returns the name of the agent."""
+        return "Name"
+
+    @property
+    def description(self) -> str | None:
+        return "Description"
+
+    async def run(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        return AgentRunResponse(messages=[ChatMessage(role=ChatRole.ASSISTANT, contents=[TextContent("Response")])])
+
+    async def run_stream(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        yield AgentRunResponseUpdate(contents=[TextContent("Response")])
+
+    def get_new_thread(self) -> AgentThread:
+        return MockAgentThread()
+
+
+# Mock ChatClient implementation for testing
+class MockChatClient(ChatClientBase):
+    _mock_response: ChatResponse | None = None
+
+    def __init__(self, mock_response: ChatResponse | None = None) -> None:
+        self._mock_response = mock_response
+
+    async def _inner_get_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        return (
+            self._mock_response
+            if self._mock_response
+            else ChatResponse(messages=ChatMessage(role=ChatRole.ASSISTANT, text="test response"))
+        )
+
+    async def _inner_get_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        yield ChatResponseUpdate(role=ChatRole.ASSISTANT, text=TextContent(text="test streaming response"))
+
+
+@fixture
+def agent_thread() -> AgentThread:
+    return MockAgentThread()
+
+
+@fixture
+def agent() -> Agent:
+    return MockAgent()
+
+
+@fixture
+def chat_client() -> ChatClientBase:
+    return MockChatClient()
+
+
+def test_agent_thread_type(agent_thread: AgentThread) -> None:
+    assert isinstance(agent_thread, AgentThread)
+
+
+def test_agent_type(agent: Agent) -> None:
+    assert isinstance(agent, Agent)
+
+
+async def test_agent_run(agent: Agent) -> None:
+    response = await agent.run("test")
+    assert response.messages[0].role == ChatRole.ASSISTANT
+    assert response.messages[0].text == "Response"
+
+
+async def test_agent_run_stream(agent: Agent) -> None:
+    async def collect_updates(updates: AsyncIterable[AgentRunResponseUpdate]) -> list[AgentRunResponseUpdate]:
+        return [u async for u in updates]
+
+    updates = await collect_updates(agent.run_stream(messages="test"))
+    assert len(updates) == 1
+    assert updates[0].text == "Response"
+
+
+async def test_chat_client_agent_thread_init_in_memory() -> None:
+    messages = [ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])]
+    thread = ChatClientAgentThread(messages=messages)
+
+    assert thread.storage_location == ChatClientAgentThreadType.IN_MEMORY_MESSAGES
+    assert thread.id is None
+    assert thread.chat_messages == messages
+
+
+async def test_chat_client_agent_thread_empty() -> None:
+    thread = ChatClientAgentThread()
+
+    assert thread.storage_location is None
+    assert thread.id is None
+    assert thread.chat_messages is None
+
+
+async def test_chat_client_agent_thread_init_invalid() -> None:
+    with raises(ValueError, match="Cannot specify both id and messages"):
+        ChatClientAgentThread(id="123", messages=[ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])])
+
+    with raises(ValueError, match="ID cannot be empty or whitespace"):
+        ChatClientAgentThread(id=" ")
+
+
+async def test_chat_client_agent_thread_init_conversation_id() -> None:
+    thread_id = str(uuid4())
+    thread = ChatClientAgentThread(id=thread_id)
+
+    assert thread.storage_location == ChatClientAgentThreadType.CONVERSATION_ID
+    assert thread.id == thread_id
+    assert thread.chat_messages is None
+
+
+async def test_chat_client_agent_thread_get_messages() -> None:
+    messages = [ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])]
+    thread = ChatClientAgentThread(messages=messages)
+
+    result = [msg async for msg in thread.get_messages()]
+    assert result == messages
+
+
+async def test_chat_client_agent_thread_on_new_messages_in_memory() -> None:
+    initial_message = ChatMessage(role=ChatRole.USER, contents=[TextContent("Initial message")])
+    new_message = ChatMessage(role=ChatRole.USER, contents=[TextContent("New message")])
+
+    thread = ChatClientAgentThread(messages=[initial_message])
+
+    await thread._on_new_messages(new_message)  # type: ignore[reportPrivateUsage]
+    assert thread.chat_messages == [initial_message, new_message]
+
+
+def test_chat_client_agent_type(chat_client: ChatClient) -> None:
+    chat_client_agent = ChatClientAgent(chat_client=chat_client)
+    assert isinstance(chat_client_agent, Agent)
+
+
+async def test_chat_client_agent_init(chat_client: ChatClient) -> None:
+    agent_id = str(uuid4())
+    agent = ChatClientAgent(chat_client=chat_client, id=agent_id, description="Test")
+
+    assert agent.id == agent_id
+    assert agent.name == "UnnamedAgent"
+    assert agent.description == "Test"
+
+
+async def test_chat_client_agent_run(chat_client: ChatClient) -> None:
+    agent = ChatClientAgent(chat_client=chat_client)
+
+    result = await agent.run("Hello")
+
+    assert result.text == "test response"
+
+
+async def test_chat_client_agent_run_stream(chat_client: ChatClient) -> None:
+    agent = ChatClientAgent(chat_client=chat_client)
+
+    result = await AgentRunResponse.from_agent_response_generator(agent.run_stream("Hello"))
+
+    assert result.text == "test streaming response"
+
+
+async def test_chat_client_agent_get_new_thread(chat_client: ChatClient) -> None:
+    agent = ChatClientAgent(chat_client=chat_client)
+    thread = agent.get_new_thread()
+
+    assert isinstance(thread, ChatClientAgentThread)
+    assert thread.storage_location is None
+
+
+async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatClient) -> None:
+    agent = ChatClientAgent(chat_client=chat_client)
+    message = ChatMessage(role=ChatRole.USER, text="Hello")
+    thread = ChatClientAgentThread(messages=[message])
+
+    result_thread, result_messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=thread,
+        input_messages=[ChatMessage(role=ChatRole.USER, text="Test")],
+        construct_thread=lambda: ChatClientAgentThread(),
+        expected_type=ChatClientAgentThread,
+    )
+
+    assert result_thread == thread
+    assert len(result_messages) == 2
+    assert result_messages[0] == message
+    assert result_messages[1].text == "Test"
+
+
+async def test_chat_client_agent_update_thread_id() -> None:
+    chat_client = MockChatClient(
+        mock_response=ChatResponse(
+            messages=[ChatMessage(role=ChatRole.ASSISTANT, contents=[TextContent("test response")])],
+            conversation_id="123",
+        )
+    )
+    agent = ChatClientAgent(chat_client=chat_client)
+    thread = agent.get_new_thread()
+
+    result = await agent.run("Hello", thread=thread)
+    assert result.text == "test response"
+
+    assert thread.id == "123"
+    assert isinstance(thread, ChatClientAgentThread)
+    assert thread.storage_location == ChatClientAgentThreadType.CONVERSATION_ID
+
+
+async def test_chat_client_agent_update_thread_messages(chat_client: ChatClient) -> None:
+    agent = ChatClientAgent(chat_client=chat_client)
+    thread = agent.get_new_thread()
+
+    result = await agent.run("Hello", thread=thread)
+    assert result.text == "test response"
+
+    assert thread.id is None
+    assert isinstance(thread, ChatClientAgentThread)
+    assert thread.storage_location == ChatClientAgentThreadType.IN_MEMORY_MESSAGES
+
+    assert thread.chat_messages is not None
+    assert len(thread.chat_messages) == 2
+    assert thread.chat_messages[0].text == "Hello"
+    assert thread.chat_messages[1].text == "test response"
+
+
+async def test_chat_client_agent_update_thread_conversation_id_missing(chat_client: ChatClient) -> None:
+    agent = ChatClientAgent(chat_client=chat_client)
+    thread = ChatClientAgentThread(id="123")
+
+    with raises(AgentExecutionException, match="Service did not return a valid conversation id"):
+        agent._update_thread_with_type_and_conversation_id(thread, None)  # type: ignore[reportPrivateUsage]

python/packages/main/tests/main/test_clients.py

@@ -0,0 +1,352 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import sys
+from collections.abc import AsyncIterable, MutableSequence, Sequence
+from typing import Any
+
+from pydantic import Field
+from pytest import fixture
+
+from agent_framework import (
+    ChatClient,
+    ChatClientBase,
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    ChatRole,
+    EmbeddingGenerator,
+    FunctionCallContent,
+    FunctionResultContent,
+    GeneratedEmbeddings,
+    TextContent,
+    ai_function,
+    use_tool_calling,
+)
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore
+else:
+    from typing_extensions import override  # type: ignore[import]
+
+
+class MockChatClient:
+    """Simple implementation of a chat client."""
+
+    async def get_response(
+        self,
+        messages: ChatMessage | Sequence[ChatMessage],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        # Implement the method
+
+        return ChatResponse(messages=ChatMessage(role="assistant", text="test response"))
+
+    async def get_streaming_response(
+        self,
+        messages: ChatMessage | Sequence[ChatMessage],
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        # Implement the method
+        yield ChatResponseUpdate(text=TextContent(text="test streaming response"), role="assistant")
+        yield ChatResponseUpdate(contents=[TextContent(text="another update")], role="assistant")
+
+
+@use_tool_calling
+class MockChatClientBase(ChatClientBase):
+    """Mock implementation of the ChatClientBase."""
+
+    run_responses: list[ChatResponse] = Field(default_factory=list)
+    streaming_responses: list[list[ChatResponseUpdate]] = Field(default_factory=list)
+
+    @override
+    async def _inner_get_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        """Send a chat request to the AI service.
+
+        Args:
+            messages: The chat messages to send.
+            chat_options: The options for the request.
+            kwargs: Any additional keyword arguments.
+
+        Returns:
+            The chat response contents representing the response(s).
+        """
+        if not self.run_responses or chat_options.tool_choice == "none":
+            return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[0].text}"))
+        return self.run_responses.pop(0)
+
+    @override
+    async def _inner_get_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        if not self.streaming_responses or chat_options.tool_choice == "none":
+            yield ChatResponseUpdate(text=f"update - {messages[0].text}", role="assistant")
+            return
+        response = self.streaming_responses.pop(0)
+        for update in response:
+            yield update
+        await asyncio.sleep(0)
+
+
+class MockEmbeddingGenerator:
+    """Simple implementation of an embedding generator."""
+
+    async def generate(
+        self,
+        input_data: Sequence[str],
+        **kwargs: Any,
+    ) -> GeneratedEmbeddings[list[float]]:
+        # Implement the method
+        embeddings = GeneratedEmbeddings[list[float]]()
+        for i, _ in enumerate(input_data):
+            embeddings.append([0.0 * 1, 0.1 * 1, 0.2 * 1, 0.3 * i, 0.4 * i])
+        return embeddings
+
+
+@fixture
+def chat_client() -> MockChatClient:
+    return MockChatClient()
+
+
+@fixture
+def chat_client_base() -> MockChatClientBase:
+    return MockChatClientBase()
+
+
+@fixture
+def embedding_generator() -> MockEmbeddingGenerator:
+    gen: EmbeddingGenerator[str, list[float]] = MockEmbeddingGenerator()
+    return gen
+
+
+def test_chat_client_type(chat_client: MockChatClient):
+    assert isinstance(chat_client, ChatClient)
+
+
+async def test_chat_client_get_response(chat_client: MockChatClient):
+    response = await chat_client.get_response(ChatMessage(role="user", text="Hello"))
+    assert response.text == "test response"
+    assert response.messages[0].role == ChatRole.ASSISTANT
+
+
+async def test_chat_client_get_streaming_response(chat_client: MockChatClient):
+    async for update in chat_client.get_streaming_response(ChatMessage(role="user", text="Hello")):
+        assert update.text == "test streaming response" or update.text == "another update"
+        assert update.role == ChatRole.ASSISTANT
+
+
+def test_embedding_generator_type(embedding_generator: MockEmbeddingGenerator):
+    assert isinstance(embedding_generator, EmbeddingGenerator)
+
+
+async def test_embedding_generator_generate(embedding_generator: MockEmbeddingGenerator):
+    input_data = ["Hello", "world"]
+    embeddings = await embedding_generator.generate(input_data)
+    assert len(embeddings) == len(input_data)
+    for emb in embeddings:
+        assert len(emb) == 5
+
+
+def test_base_client(chat_client_base: MockChatClientBase):
+    assert isinstance(chat_client_base, ChatClientBase)
+    assert isinstance(chat_client_base, ChatClient)
+
+
+async def test_base_client_get_response(chat_client_base: MockChatClientBase):
+    response = await chat_client_base.get_response(ChatMessage(role="user", text="Hello"))
+    assert response.messages[0].role == ChatRole.ASSISTANT
+    assert response.messages[0].text == "test response - Hello"
+
+
+async def test_base_client_get_streaming_response(chat_client_base: MockChatClientBase):
+    async for update in chat_client_base.get_streaming_response(ChatMessage(role="user", text="Hello")):
+        assert update.text == "update - Hello" or update.text == "another update"
+
+
+async def test_base_client_with_function_calling(chat_client_base: MockChatClientBase):
+    exec_counter = 0
+
+    @ai_function(name="test_function")
+    def ai_func(arg1: str) -> str:
+        nonlocal exec_counter
+        exec_counter += 1
+        return f"Processed {arg1}"
+
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=ChatMessage(
+                role="assistant",
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='{"arg1": "value1"}')],
+            )
+        ),
+        ChatResponse(messages=ChatMessage(role="assistant", text="done")),
+    ]
+    response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
+    assert exec_counter == 1
+    assert len(response.messages) == 3
+    assert response.messages[0].role == ChatRole.ASSISTANT
+    assert isinstance(response.messages[0].contents[0], FunctionCallContent)
+    assert response.messages[0].contents[0].name == "test_function"
+    assert response.messages[0].contents[0].arguments == '{"arg1": "value1"}'
+    assert response.messages[0].contents[0].call_id == "1"
+    assert response.messages[1].role == ChatRole.TOOL
+    assert isinstance(response.messages[1].contents[0], FunctionResultContent)
+    assert response.messages[1].contents[0].call_id == "1"
+    assert response.messages[1].contents[0].result == "Processed value1"
+    assert response.messages[2].role == ChatRole.ASSISTANT
+    assert response.messages[2].text == "done"
+
+
+async def test_base_client_with_function_calling_disabled(chat_client_base: MockChatClientBase):
+    chat_client_base.__maximum_iterations_per_request = 0
+    exec_counter = 0
+
+    @ai_function(name="test_function")
+    def ai_func(arg1: str) -> str:
+        nonlocal exec_counter
+        exec_counter += 1
+        return f"Processed {arg1}"
+
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=ChatMessage(
+                role="assistant",
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='{"arg1": "value1"}')],
+            )
+        ),
+        ChatResponse(messages=ChatMessage(role="assistant", text="done")),
+    ]
+    response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
+    assert exec_counter == 0
+    assert len(response.messages) == 1
+    assert response.messages[0].role == ChatRole.ASSISTANT
+    assert response.messages[0].text == "test response - hello"
+
+
+async def test_base_client_with_streaming_function_calling(chat_client_base: MockChatClientBase):
+    exec_counter = 0
+
+    @ai_function(name="test_function")
+    def ai_func(arg1: str) -> str:
+        nonlocal exec_counter
+        exec_counter += 1
+        return f"Processed {arg1}"
+
+    chat_client_base.streaming_responses = [
+        [
+            ChatResponseUpdate(
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='{"arg1":')],
+                role="assistant",
+            ),
+            ChatResponseUpdate(
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='"value1"}')],
+                role="assistant",
+            ),
+        ],
+        [
+            ChatResponseUpdate(
+                contents=[TextContent(text="Processed value1")],
+                role="assistant",
+            )
+        ],
+    ]
+    updates = []
+    async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[ai_func]):
+        updates.append(update)
+    assert len(updates) == 4  # two updates with the function call, the function result and the final text
+    assert updates[0].contents[0].call_id == "1"
+    assert updates[1].contents[0].call_id == "1"
+    assert updates[2].contents[0].call_id == "1"
+    assert updates[3].text == "Processed value1"
+    assert exec_counter == 1
+
+
+async def test_base_client_with_streaming_function_calling_disabled(chat_client_base: MockChatClientBase):
+    chat_client_base.__maximum_iterations_per_request = 0
+    exec_counter = 0
+
+    @ai_function(name="test_function")
+    def ai_func(arg1: str) -> str:
+        nonlocal exec_counter
+        exec_counter += 1
+        return f"Processed {arg1}"
+
+    chat_client_base.streaming_responses = [
+        [
+            ChatResponseUpdate(
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='{"arg1":')],
+                role="assistant",
+            ),
+            ChatResponseUpdate(
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='"value1"}')],
+                role="assistant",
+            ),
+        ],
+        [
+            ChatResponseUpdate(
+                contents=[TextContent(text="Processed value1")],
+                role="assistant",
+            )
+        ],
+    ]
+    updates = []
+    async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[ai_func]):
+        updates.append(update)
+    assert len(updates) == 1
+    assert exec_counter == 0
+
+
+def test_chat_options_parsing_tools(chat_client_base, ai_function_tool) -> None:
+    """Test that chat options can parse tools correctly."""
+
+    def echo() -> str:
+        """Echo the input."""
+        return "Echo"
+
+    dict_function = {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Retrieves current weather for the given location.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {"type": "string", "description": "City and country e.g. Bogotá, Colombia"},
+                    "units": {
+                        "type": "string",
+                        "enum": ["celsius", "fahrenheit"],
+                        "description": "Units the temperature will be returned in.",
+                    },
+                },
+                "required": ["location", "units"],
+                "additionalProperties": False,
+            },
+            "strict": True,
+        },
+    }
+
+    options = ChatOptions(tools=[ai_function_tool, echo, dict_function], tool_choice="auto")
+    assert len(options.tools) == 3
+    assert options.tools[0] == ai_function_tool
+    assert options.tools[1] != echo
+    assert options.tools[2] == dict_function
+    # after prepare, the tools should be represented as dicts
+    # while ai_tools is still the same.
+    chat_client_base._prepare_tools_and_tool_choice(chat_options=options)
+    assert options._ai_tools[0] == ai_function_tool
+    assert options._ai_tools[2] == dict_function
+    assert len(options.tools) == 3
+    assert options.tools[0]["function"]["name"] == "simple_function"
+    assert options.tools[1]["function"]["name"] == "echo"
+    assert options.tools[2]["function"]["name"] == "get_weather"

python/packages/main/tests/main/test_logging.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+import pytest
+
+from agent_framework import get_logger
+from agent_framework.exceptions import AgentFrameworkException
+
+
+def test_get_logger():
+    """Test that the logger is created with the correct name."""
+    logger = get_logger()
+    assert logger.name == "agent_framework"
+
+
+def test_get_logger_custom_name():
+    """Test that the logger can be created with a custom name."""
+    custom_name = "agent_framework.custom"
+    logger = get_logger(custom_name)
+    assert logger.name == custom_name
+
+
+def test_get_logger_invalid_name():
+    """Test that an exception is raised for an invalid logger name."""
+    with pytest.raises(AgentFrameworkException, match="Logger name must start with 'agent_framework'."):
+        get_logger("invalid_name")
+
+
+def test_log(caplog):
+    """Test that the logger can log messages and adheres to the expected format."""
+    logger = get_logger()
+    with caplog.at_level("DEBUG"):
+        logger.debug("This is a debug message")
+        assert len(caplog.records) == 1
+        record = caplog.records[0]
+        assert record.levelname == "DEBUG"
+        assert record.message == "This is a debug message"
+        assert record.name == "agent_framework"
+        assert record.pathname.endswith("test_logging.py")

python/packages/main/tests/main/test_telemetry.py

@@ -0,0 +1,612 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from collections.abc import AsyncIterable, MutableSequence
+from typing import Any
+from unittest.mock import Mock, patch
+
+import pytest
+
+from agent_framework import (
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    ChatRole,
+    UsageDetails,
+)
+from agent_framework.telemetry import (
+    AGENT_FRAMEWORK_USER_AGENT,
+    ROLE_EVENT_MAP,
+    TELEMETRY_DISABLED_ENV_VAR,
+    USER_AGENT_KEY,
+    ChatMessageListTimestampFilter,
+    GenAIAttributes,
+    prepend_agent_framework_to_user_agent,
+    start_as_current_span,
+    use_telemetry,
+)
+
+# region Test constants
+
+
+def test_telemetry_disabled_env_var():
+    """Test that the telemetry disabled environment variable is correctly defined."""
+    assert TELEMETRY_DISABLED_ENV_VAR == "AZURE_TELEMETRY_DISABLED"
+
+
+def test_user_agent_key():
+    """Test that the user agent key is correctly defined."""
+    assert USER_AGENT_KEY == "User-Agent"
+
+
+def test_agent_framework_user_agent_format():
+    """Test that the agent framework user agent is correctly formatted."""
+    assert AGENT_FRAMEWORK_USER_AGENT.startswith("agent-framework-python/")
+
+
+def test_app_info_when_telemetry_enabled():
+    """Test that APP_INFO is set when telemetry is enabled."""
+    with patch("agent_framework.telemetry.IS_TELEMETRY_ENABLED", True):
+        import importlib
+
+        import agent_framework.telemetry
+
+        importlib.reload(agent_framework.telemetry)
+        from agent_framework.telemetry import APP_INFO
+
+        assert APP_INFO is not None
+        assert "agent-framework-version" in APP_INFO
+        assert APP_INFO["agent-framework-version"].startswith("python/")
+
+
+def test_app_info_when_telemetry_disabled():
+    """Test that APP_INFO is None when telemetry is disabled."""
+    # Test the logic directly since APP_INFO is set at module import time
+    with patch("agent_framework.telemetry.IS_TELEMETRY_ENABLED", False):
+        # Simulate the module's logic for APP_INFO
+        test_app_info = (
+            {
+                "agent-framework-version": "python/test",
+            }
+            if False  # This simulates IS_TELEMETRY_ENABLED being False
+            else None
+        )
+        assert test_app_info is None
+
+
+def test_role_event_map():
+    """Test that ROLE_EVENT_MAP contains expected mappings."""
+    assert ROLE_EVENT_MAP["system"] == GenAIAttributes.SYSTEM_MESSAGE.value
+    assert ROLE_EVENT_MAP["user"] == GenAIAttributes.USER_MESSAGE.value
+    assert ROLE_EVENT_MAP["assistant"] == GenAIAttributes.ASSISTANT_MESSAGE.value
+    assert ROLE_EVENT_MAP["tool"] == GenAIAttributes.TOOL_MESSAGE.value
+
+
+def test_enum_values():
+    """Test that GenAIAttributes enum has expected values."""
+    assert GenAIAttributes.OPERATION.value == "gen_ai.operation.name"
+    assert GenAIAttributes.SYSTEM.value == "gen_ai.system"
+    assert GenAIAttributes.MODEL.value == "gen_ai.request.model"
+    assert GenAIAttributes.CHAT_COMPLETION_OPERATION.value == "chat.completions"
+    assert GenAIAttributes.CHAT_STREAMING_COMPLETION_OPERATION.value == "chat.streaming_completions"
+    assert GenAIAttributes.TOOL_EXECUTION_OPERATION.value == "execute_tool"
+
+
+# region Test prepend_agent_framework_to_user_agent
+
+
+def test_prepend_to_existing_user_agent():
+    """Test prepending to existing User-Agent header."""
+    headers = {"User-Agent": "existing-agent/1.0"}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    assert "User-Agent" in result
+    assert result["User-Agent"].startswith("agent-framework-python/")
+    assert "existing-agent/1.0" in result["User-Agent"]
+
+
+def test_prepend_to_empty_headers():
+    """Test prepending to headers without User-Agent."""
+    headers = {"Content-Type": "application/json"}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    assert "User-Agent" in result
+    assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT
+    assert "Content-Type" in result
+
+
+def test_prepend_to_empty_dict():
+    """Test prepending to empty headers dict."""
+    headers = {}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    assert "User-Agent" in result
+    assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT
+
+
+def test_modifies_original_dict():
+    """Test that the function modifies the original headers dict."""
+    headers = {"Other-Header": "value"}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    assert result is headers  # Same object
+    assert "User-Agent" in headers
+
+
+# region ModelDiagnosticSettings tests
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(None, None)], indirect=True)
+def test_default_values(model_diagnostic_settings):
+    """Test default values for ModelDiagnosticSettings."""
+    assert not model_diagnostic_settings.ENABLED
+    assert not model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(False, False)], indirect=True)
+def test_disabled(model_diagnostic_settings):
+    """Test default values for ModelDiagnosticSettings."""
+    assert not model_diagnostic_settings.ENABLED
+    assert not model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
+def test_non_sensitive_events_enabled(model_diagnostic_settings):
+    """Test loading model_diagnostic_settings from environment variables."""
+    assert model_diagnostic_settings.ENABLED
+    assert not model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(True, True)], indirect=True)
+def test_sensitive_events_enabled(model_diagnostic_settings):
+    """Test loading model_diagnostic_settings from environment variables."""
+    assert model_diagnostic_settings.ENABLED
+    assert model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(False, True)], indirect=True)
+def test_sensitive_events_enabled_only(model_diagnostic_settings):
+    """Test loading sensitive events setting from environment.
+
+    But when sensitive events are enabled, diagnostics are also enabled.
+    """
+    assert model_diagnostic_settings.ENABLED
+    assert model_diagnostic_settings.SENSITIVE_EVENTS_ENABLED
+
+
+# region Test ChatMessageListTimestampFilter
+
+
+def test_filter_without_index_key():
+    """Test filter method when record doesn't have INDEX_KEY."""
+    log_filter = ChatMessageListTimestampFilter()
+    record = logging.LogRecord(
+        name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
+    )
+    original_created = record.created
+
+    result = log_filter.filter(record)
+
+    assert result is True
+    assert record.created == original_created
+
+
+def test_filter_with_index_key():
+    """Test filter method when record has INDEX_KEY."""
+    log_filter = ChatMessageListTimestampFilter()
+    record = logging.LogRecord(
+        name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
+    )
+    original_created = record.created
+
+    # Add the index key
+    setattr(record, ChatMessageListTimestampFilter.INDEX_KEY, 5)
+
+    result = log_filter.filter(record)
+
+    assert result is True
+    # Should increment by 5 microseconds (5 * 1e-6)
+    assert record.created == original_created + 5 * 1e-6
+
+
+def test_index_key_constant():
+    """Test that INDEX_KEY constant is correctly defined."""
+    assert ChatMessageListTimestampFilter.INDEX_KEY == "CHAT_MESSAGE_INDEX"
+
+
+# region Test start_as_current_span
+
+
+def test_start_span_basic():
+    """Test starting a span with basic function info."""
+    mock_tracer = Mock()
+    mock_span = Mock()
+    mock_tracer.start_as_current_span.return_value = mock_span
+
+    # Create a mock function
+    mock_function = Mock()
+    mock_function.name = "test_function"
+    mock_function.description = "Test function description"
+
+    result = start_as_current_span(mock_tracer, mock_function)
+
+    assert result == mock_span
+    mock_tracer.start_as_current_span.assert_called_once()
+
+    call_args = mock_tracer.start_as_current_span.call_args
+    assert call_args[0][0] == "execute_tool test_function"
+
+    attributes = call_args[1]["attributes"]
+    assert attributes[GenAIAttributes.OPERATION.value] == GenAIAttributes.TOOL_EXECUTION_OPERATION.value
+    assert attributes[GenAIAttributes.TOOL_NAME.value] == "test_function"
+    assert attributes[GenAIAttributes.TOOL_DESCRIPTION.value] == "Test function description"
+
+
+def test_start_span_with_metadata():
+    """Test starting a span with metadata containing tool_call_id."""
+    mock_tracer = Mock()
+    mock_span = Mock()
+    mock_tracer.start_as_current_span.return_value = mock_span
+
+    mock_function = Mock()
+    mock_function.name = "test_function"
+    mock_function.description = "Test function"
+
+    metadata = {"tool_call_id": "test_call_123"}
+
+    _ = start_as_current_span(mock_tracer, mock_function, metadata)
+
+    call_args = mock_tracer.start_as_current_span.call_args
+    attributes = call_args[1]["attributes"]
+    assert attributes[GenAIAttributes.TOOL_CALL_ID.value] == "test_call_123"
+
+
+def test_start_span_without_description():
+    """Test starting a span when function has no description."""
+    mock_tracer = Mock()
+    mock_span = Mock()
+    mock_tracer.start_as_current_span.return_value = mock_span
+
+    mock_function = Mock()
+    mock_function.name = "test_function"
+    mock_function.description = None
+
+    start_as_current_span(mock_tracer, mock_function)
+
+    call_args = mock_tracer.start_as_current_span.call_args
+    attributes = call_args[1]["attributes"]
+    assert GenAIAttributes.TOOL_DESCRIPTION.value not in attributes
+
+
+def test_start_span_empty_metadata():
+    """Test starting a span with empty metadata."""
+    mock_tracer = Mock()
+    mock_span = Mock()
+    mock_tracer.start_as_current_span.return_value = mock_span
+
+    mock_function = Mock()
+    mock_function.name = "test_function"
+    mock_function.description = "Test function"
+
+    start_as_current_span(mock_tracer, mock_function, {})
+
+    call_args = mock_tracer.start_as_current_span.call_args
+    attributes = call_args[1]["attributes"]
+    assert GenAIAttributes.TOOL_CALL_ID.value not in attributes
+
+
+# region Test use_telemetry decorator
+
+
+def test_decorator_with_valid_class():
+    """Test that decorator works with a valid ChatClientBase-like class."""
+
+    # Create a mock class with the required methods
+    class MockChatClient:
+        MODEL_PROVIDER_NAME = "test_provider"
+
+        async def _inner_get_response(self, *, messages, chat_options, **kwargs):
+            return Mock()
+
+        async def _inner_get_streaming_response(self, *, messages, chat_options, **kwargs):
+            async def gen():
+                yield Mock()
+
+            return gen()
+
+    # Apply the decorator
+    decorated_class = use_telemetry(MockChatClient)
+
+    # Check that the methods were wrapped
+    assert hasattr(decorated_class._inner_get_response, "__model_diagnostics_chat_client__")
+    assert hasattr(decorated_class._inner_get_streaming_response, "__model_diagnostics_streaming_chat_completion__")
+
+
+def test_decorator_with_missing_methods():
+    """Test that decorator handles classes missing required methods gracefully."""
+
+    class MockChatClient:
+        MODEL_PROVIDER_NAME = "test_provider"
+
+    # Apply the decorator - should not raise an error
+    decorated_class = use_telemetry(MockChatClient)
+
+    # Class should be returned unchanged
+    assert decorated_class is MockChatClient
+
+
+def test_decorator_with_partial_methods():
+    """Test decorator when only one method is present."""
+
+    class MockChatClient:
+        MODEL_PROVIDER_NAME = "test_provider"
+
+        async def _inner_get_response(self, *, messages, chat_options, **kwargs):
+            return Mock()
+
+    decorated_class = use_telemetry(MockChatClient)
+
+    # Only the present method should be wrapped
+    assert hasattr(decorated_class._inner_get_response, "__model_diagnostics_chat_client__")
+    assert not hasattr(decorated_class, "_inner_get_streaming_response")
+
+
+# region Test telemetry decorator with mock client
+
+
+@pytest.fixture
+def mock_chat_client():
+    """Create a mock chat client for testing."""
+
+    class MockChatClient:
+        MODEL_PROVIDER_NAME = "test_provider"
+
+        def __init__(self):
+            self.ai_model_id = "test-model"
+
+        def service_url(self):
+            return "https://test.example.com"
+
+        async def _inner_get_response(
+            self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+        ):
+            return ChatResponse(
+                messages=[ChatMessage(role=ChatRole.ASSISTANT, text="Test response")],
+                usage_details=UsageDetails(input_token_count=10, output_token_count=20),
+                finish_reason=None,
+            )
+
+        async def _inner_get_streaming_response(
+            self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+        ):
+            yield ChatResponseUpdate(text="Hello", role=ChatRole.ASSISTANT)
+            yield ChatResponseUpdate(text=" world", role=ChatRole.ASSISTANT)
+
+    return MockChatClient()
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(False, False)], indirect=True)
+async def test_telemetry_disabled_bypasses_instrumentation(mock_chat_client, model_diagnostic_settings):
+    """Test that when diagnostics are disabled, telemetry is bypassed."""
+    decorated_class = use_telemetry(type(mock_chat_client))
+    client = decorated_class()
+
+    messages = [ChatMessage(role=ChatRole.USER, text="Test message")]
+    chat_options = ChatOptions()
+
+    with (
+        patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
+        patch("agent_framework.telemetry.use_span") as mock_use_span,
+    ):
+        # This should not create any spans
+        response = await client._inner_get_response(messages=messages, chat_options=chat_options)
+        assert response is not None
+        mock_use_span.assert_not_called()
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(True, True)], indirect=True)
+async def test_instrumentation_enabled(mock_chat_client, model_diagnostic_settings):
+    """Test that when diagnostics are enabled, telemetry is applied."""
+    decorated_class = use_telemetry(type(mock_chat_client))
+    client = decorated_class()
+
+    messages = [ChatMessage(role=ChatRole.USER, text="Test message")]
+    chat_options = ChatOptions()
+
+    with (
+        patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
+        patch("agent_framework.telemetry.use_span") as mock_use_span,
+        patch("agent_framework.telemetry.logger") as mock_logger,
+    ):
+        response = await client._inner_get_response(messages=messages, chat_options=chat_options)
+        assert response is not None
+        mock_use_span.assert_called_once()
+        # Check that logger.info was called (telemetry logs input/output)
+        assert mock_logger.info.call_count == 2
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
+async def test_streaming_response_with_diagnostics_enabled_via_decorator(mock_chat_client, model_diagnostic_settings):
+    """Test streaming telemetry through the use_telemetry decorator."""
+    decorated_class = use_telemetry(type(mock_chat_client))
+    client = decorated_class()
+    messages = [ChatMessage(role=ChatRole.USER, text="Test")]
+    chat_options = ChatOptions()
+
+    with (
+        patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
+        patch("agent_framework.telemetry.use_span") as mock_use_span,
+        patch("agent_framework.telemetry._get_chat_response_span") as mock_get_span,
+        patch("agent_framework.telemetry._set_chat_response_input") as mock_set_input,
+        patch("agent_framework.telemetry._set_chat_response_output") as mock_set_output,
+    ):
+        mock_span = Mock()
+        mock_use_span.return_value.__enter__.return_value = mock_span
+        mock_use_span.return_value.__exit__.return_value = None
+
+        # We can't easily mock ChatResponse.from_chat_response_updates since it's imported locally,
+        # but we can verify telemetry calls were made
+
+        # Collect all yielded updates
+        updates = []
+        async for update in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
+            updates.append(update)
+
+        # Verify we got the expected updates
+        assert len(updates) == 2
+
+        # Verify telemetry calls were made
+        mock_get_span.assert_called_once()
+        mock_set_input.assert_called_once_with("test_provider", messages)
+        mock_set_output.assert_called_once()
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
+async def test_streaming_response_with_exception_via_decorator(mock_chat_client, model_diagnostic_settings):
+    """Test streaming telemetry exception handling through decorator."""
+
+    async def _inner_get_streaming_response(
+        self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        yield ChatResponseUpdate(text="Partial", role=ChatRole.ASSISTANT)
+        raise ValueError("Test streaming error")
+
+    type(mock_chat_client)._inner_get_streaming_response = _inner_get_streaming_response
+
+    decorated_class = use_telemetry(type(mock_chat_client))
+    client = decorated_class()
+
+    messages = [ChatMessage(role=ChatRole.USER, text="Test")]
+    chat_options = ChatOptions()
+
+    with (
+        patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
+        patch("agent_framework.telemetry.use_span") as mock_use_span,
+        patch("agent_framework.telemetry._get_chat_response_span"),
+        patch("agent_framework.telemetry._set_chat_response_input"),
+        patch("agent_framework.telemetry._set_chat_response_error") as mock_set_error,
+    ):
+        mock_span = Mock()
+        mock_use_span.return_value.__enter__.return_value = mock_span
+        mock_use_span.return_value.__exit__.return_value = None
+
+        # Should raise the exception and call error handler
+        with pytest.raises(ValueError, match="Test streaming error"):
+            async for _ in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
+                pass
+
+        # Verify error was recorded
+        mock_set_error.assert_called_once()
+        assert isinstance(mock_set_error.call_args[0][1], ValueError)
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(False, False)], indirect=True)
+async def test_streaming_response_diagnostics_disabled_via_decorator(model_diagnostic_settings):
+    """Test streaming response when diagnostics are disabled."""
+    from agent_framework import ChatResponseUpdate
+
+    class MockStreamingClientNoDiagnostics:
+        MODEL_PROVIDER_NAME = "test_provider"
+
+        async def _inner_get_streaming_response(
+            self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+        ) -> AsyncIterable[ChatResponseUpdate]:
+            yield ChatResponseUpdate(text="Test", role=ChatRole.ASSISTANT)
+
+    decorated_class = use_telemetry(MockStreamingClientNoDiagnostics)
+    client = decorated_class()
+
+    messages = [ChatMessage(role=ChatRole.USER, text="Test")]
+    chat_options = ChatOptions()
+
+    with (
+        patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
+        patch("agent_framework.telemetry._get_chat_response_span") as mock_get_span,
+    ):
+        # Should not create spans when diagnostics are disabled
+        updates = []
+        async for update in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
+            updates.append(update)
+
+        assert len(updates) == 1
+        # Should not have called telemetry functions
+        mock_get_span.assert_not_called()
+
+
+# region Test empty streaming response handling
+
+
+@pytest.mark.parametrize("model_diagnostic_settings", [(True, False)], indirect=True)
+async def test_empty_streaming_response_via_decorator(model_diagnostic_settings):
+    """Test streaming wrapper with empty response."""
+
+    class MockEmptyStreamingClient:
+        MODEL_PROVIDER_NAME = "test_provider"
+
+        def __init__(self):
+            self.ai_model_id = "test_model"
+
+        def service_url(self) -> str:
+            return "https://test.com"
+
+        async def _inner_get_streaming_response(
+            self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+        ) -> AsyncIterable[ChatResponseUpdate]:
+            # Return empty stream
+            return
+            yield  # This will never be reached
+
+    decorated_class = use_telemetry(MockEmptyStreamingClient)
+    client = decorated_class()
+
+    messages = [ChatMessage(role=ChatRole.USER, text="Test")]
+    chat_options = ChatOptions()
+
+    with (
+        patch("agent_framework.telemetry.MODEL_DIAGNOSTICS_SETTINGS", model_diagnostic_settings),
+        patch("agent_framework.telemetry.use_span") as mock_use_span,
+        patch("agent_framework.telemetry._get_chat_response_span"),
+        patch("agent_framework.telemetry._set_chat_response_input"),
+        patch("agent_framework.telemetry._set_chat_response_output") as mock_set_output,
+    ):
+        mock_span = Mock()
+        mock_use_span.return_value.__enter__.return_value = mock_span
+        mock_use_span.return_value.__exit__.return_value = None
+
+        # Should handle empty stream gracefully
+        updates = []
+        async for update in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
+            updates.append(update)
+
+        assert len(updates) == 0
+        # Should still call telemetry
+        mock_set_output.assert_called_once()
+
+
+def test_start_as_current_span_with_none_metadata():
+    """Test start_as_current_span with None metadata."""
+    mock_tracer = Mock()
+    mock_span = Mock()
+    mock_tracer.start_as_current_span.return_value = mock_span
+
+    mock_function = Mock()
+    mock_function.name = "test_function"
+    mock_function.description = "Test description"
+
+    result = start_as_current_span(mock_tracer, mock_function, None)
+
+    assert result == mock_span
+    call_args = mock_tracer.start_as_current_span.call_args
+    attributes = call_args[1]["attributes"]
+    assert GenAIAttributes.TOOL_CALL_ID.value not in attributes
+
+
+def test_prepend_user_agent_with_none_value():
+    """Test prepend user agent with None value in headers."""
+    headers = {"User-Agent": None}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    # Should handle None gracefully
+    assert "User-Agent" in result
+    assert AGENT_FRAMEWORK_USER_AGENT in str(result["User-Agent"])

python/packages/main/tests/main/test_tools.py

@@ -0,0 +1,290 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from unittest.mock import Mock, patch
+
+import pytest
+from pydantic import BaseModel
+
+from agent_framework import AIFunction, AITool, ai_function
+from agent_framework.telemetry import GenAIAttributes
+
+
+def test_ai_function_decorator():
+    """Test the ai_function decorator."""
+
+    @ai_function(name="test_tool", description="A test tool")
+    def test_tool(x: int, y: int) -> int:
+        """A simple function that adds two numbers."""
+        return x + y
+
+    assert isinstance(test_tool, AITool)
+    assert isinstance(test_tool, AIFunction)
+    assert test_tool.name == "test_tool"
+    assert test_tool.description == "A test tool"
+    assert test_tool.parameters() == {
+        "properties": {"x": {"title": "X", "type": "integer"}, "y": {"title": "Y", "type": "integer"}},
+        "required": ["x", "y"],
+        "title": "test_tool_input",
+        "type": "object",
+    }
+    assert test_tool(1, 2) == 3
+
+
+def test_ai_function_decorator_without_args():
+    """Test the ai_function decorator."""
+
+    @ai_function
+    def test_tool(x: int, y: int) -> int:
+        """A simple function that adds two numbers."""
+        return x + y
+
+    assert isinstance(test_tool, AITool)
+    assert isinstance(test_tool, AIFunction)
+    assert test_tool.name == "test_tool"
+    assert test_tool.description == "A simple function that adds two numbers."
+    assert test_tool.parameters() == {
+        "properties": {"x": {"title": "X", "type": "integer"}, "y": {"title": "Y", "type": "integer"}},
+        "required": ["x", "y"],
+        "title": "test_tool_input",
+        "type": "object",
+    }
+    assert test_tool(1, 2) == 3
+
+
+async def test_ai_function_decorator_with_async():
+    """Test the ai_function decorator with an async function."""
+
+    @ai_function(name="async_test_tool", description="An async test tool")
+    async def async_test_tool(x: int, y: int) -> int:
+        """An async function that adds two numbers."""
+        return x + y
+
+    assert isinstance(async_test_tool, AITool)
+    assert isinstance(async_test_tool, AIFunction)
+    assert async_test_tool.name == "async_test_tool"
+    assert async_test_tool.description == "An async test tool"
+    assert async_test_tool.parameters() == {
+        "properties": {"x": {"title": "X", "type": "integer"}, "y": {"title": "Y", "type": "integer"}},
+        "required": ["x", "y"],
+        "title": "async_test_tool_input",
+        "type": "object",
+    }
+    assert (await async_test_tool(1, 2)) == 3
+
+
+# Telemetry tests for AIFunction
+async def test_ai_function_invoke_telemetry_enabled():
+    """Test the ai_function invoke method with telemetry enabled."""
+
+    @ai_function(name="telemetry_test_tool", description="A test tool for telemetry")
+    def telemetry_test_tool(x: int, y: int) -> int:
+        """A function that adds two numbers for telemetry testing."""
+        return x + y
+
+    # Mock the tracer and span
+    with (
+        patch("agent_framework._tools.tracer") as mock_tracer,
+        patch("agent_framework._tools.start_as_current_span") as mock_start_span,
+    ):
+        mock_span = Mock()
+        mock_context_manager = Mock()
+        mock_context_manager.__enter__ = Mock(return_value=mock_span)
+        mock_context_manager.__exit__ = Mock(return_value=None)
+        mock_start_span.return_value = mock_context_manager
+
+        # Mock the histogram
+        mock_histogram = Mock()
+        telemetry_test_tool.invocation_duration_histogram = mock_histogram
+
+        # Call invoke
+        result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
+
+        # Verify result
+        assert result == 3
+
+        # Verify telemetry calls
+        mock_start_span.assert_called_once_with(
+            mock_tracer, telemetry_test_tool, metadata={"tool_call_id": "test_call_id", "kwargs": {"x": 1, "y": 2}}
+        )
+
+        # Verify histogram was called with correct attributes
+        mock_histogram.record.assert_called_once()
+        call_args = mock_histogram.record.call_args
+        assert call_args[0][0] > 0  # duration should be positive
+        attributes = call_args[1]["attributes"]
+        assert attributes[GenAIAttributes.MEASUREMENT_FUNCTION_TAG_NAME.value] == "telemetry_test_tool"
+        assert attributes[GenAIAttributes.TOOL_CALL_ID.value] == "test_call_id"
+
+
+async def test_ai_function_invoke_telemetry_with_pydantic_args():
+    """Test the ai_function invoke method with Pydantic model arguments."""
+
+    @ai_function(name="pydantic_test_tool", description="A test tool with Pydantic args")
+    def pydantic_test_tool(x: int, y: int) -> int:
+        """A function that adds two numbers using Pydantic args."""
+        return x + y
+
+    # Create arguments as Pydantic model instance
+    args_model = pydantic_test_tool.input_model(x=5, y=10)
+
+    with (
+        patch("agent_framework._tools.tracer") as mock_tracer,
+        patch("agent_framework._tools.start_as_current_span") as mock_start_span,
+    ):
+        mock_span = Mock()
+        mock_context_manager = Mock()
+        mock_context_manager.__enter__ = Mock(return_value=mock_span)
+        mock_context_manager.__exit__ = Mock(return_value=None)
+        mock_start_span.return_value = mock_context_manager
+
+        mock_histogram = Mock()
+        pydantic_test_tool.invocation_duration_histogram = mock_histogram
+
+        # Call invoke with Pydantic model
+        result = await pydantic_test_tool.invoke(arguments=args_model, tool_call_id="pydantic_call")
+
+        # Verify result
+        assert result == 15
+
+        # Verify telemetry calls
+        mock_start_span.assert_called_once_with(
+            mock_tracer, pydantic_test_tool, metadata={"tool_call_id": "pydantic_call", "kwargs": {"x": 5, "y": 10}}
+        )
+
+
+async def test_ai_function_invoke_telemetry_with_exception():
+    """Test the ai_function invoke method with telemetry when an exception occurs."""
+
+    @ai_function(name="exception_test_tool", description="A test tool that raises an exception")
+    def exception_test_tool(x: int, y: int) -> int:
+        """A function that raises an exception for telemetry testing."""
+        raise ValueError("Test exception for telemetry")
+
+    with (
+        patch("agent_framework._tools.tracer"),
+        patch("agent_framework._tools.start_as_current_span") as mock_start_span,
+    ):
+        mock_span = Mock()
+        mock_context_manager = Mock()
+        mock_context_manager.__enter__ = Mock(return_value=mock_span)
+        mock_context_manager.__exit__ = Mock(return_value=None)
+        mock_start_span.return_value = mock_context_manager
+
+        mock_histogram = Mock()
+        exception_test_tool.invocation_duration_histogram = mock_histogram
+
+        # Call invoke and expect exception
+        with pytest.raises(ValueError, match="Test exception for telemetry"):
+            await exception_test_tool.invoke(x=1, y=2, tool_call_id="exception_call")
+
+        # Verify telemetry calls
+        mock_start_span.assert_called_once()
+
+        # Verify span exception recording
+        mock_span.record_exception.assert_called_once()
+        mock_span.set_attribute.assert_called()
+        mock_span.set_status.assert_called_once()
+
+        # Verify histogram was called with error attributes
+        mock_histogram.record.assert_called_once()
+        call_args = mock_histogram.record.call_args
+        attributes = call_args[1]["attributes"]
+        assert attributes[GenAIAttributes.ERROR_TYPE.value] == "ValueError"
+
+
+async def test_ai_function_invoke_telemetry_async_function():
+    """Test the ai_function invoke method with telemetry on async function."""
+
+    @ai_function(name="async_telemetry_test", description="An async test tool for telemetry")
+    async def async_telemetry_test(x: int, y: int) -> int:
+        """An async function for telemetry testing."""
+        return x * y
+
+    with (
+        patch("agent_framework._tools.tracer") as mock_tracer,
+        patch("agent_framework._tools.start_as_current_span") as mock_start_span,
+    ):
+        mock_span = Mock()
+        mock_context_manager = Mock()
+        mock_context_manager.__enter__ = Mock(return_value=mock_span)
+        mock_context_manager.__exit__ = Mock(return_value=None)
+        mock_start_span.return_value = mock_context_manager
+
+        mock_histogram = Mock()
+        async_telemetry_test.invocation_duration_histogram = mock_histogram
+
+        # Call invoke
+        result = await async_telemetry_test.invoke(x=3, y=4, tool_call_id="async_call")
+
+        # Verify result
+        assert result == 12
+
+        # Verify telemetry calls
+        mock_start_span.assert_called_once_with(
+            mock_tracer, async_telemetry_test, metadata={"tool_call_id": "async_call", "kwargs": {"x": 3, "y": 4}}
+        )
+
+        # Verify histogram recording
+        mock_histogram.record.assert_called_once()
+        call_args = mock_histogram.record.call_args
+        attributes = call_args[1]["attributes"]
+        assert attributes[GenAIAttributes.MEASUREMENT_FUNCTION_TAG_NAME.value] == "async_telemetry_test"
+
+
+async def test_ai_function_invoke_telemetry_no_tool_call_id():
+    """Test the ai_function invoke method with telemetry when no tool_call_id is provided."""
+
+    @ai_function(name="no_id_test_tool", description="A test tool without tool_call_id")
+    def no_id_test_tool(x: int) -> int:
+        """A function for testing without tool_call_id."""
+        return x * 2
+
+    with (
+        patch("agent_framework._tools.tracer") as mock_tracer,
+        patch("agent_framework._tools.start_as_current_span") as mock_start_span,
+    ):
+        mock_span = Mock()
+        mock_context_manager = Mock()
+        mock_context_manager.__enter__ = Mock(return_value=mock_span)
+        mock_context_manager.__exit__ = Mock(return_value=None)
+        mock_start_span.return_value = mock_context_manager
+
+        mock_histogram = Mock()
+        no_id_test_tool.invocation_duration_histogram = mock_histogram
+
+        # Call invoke without tool_call_id
+        result = await no_id_test_tool.invoke(x=5)
+
+        # Verify result
+        assert result == 10
+
+        # Verify telemetry calls
+        mock_start_span.assert_called_once_with(
+            mock_tracer, no_id_test_tool, metadata={"tool_call_id": None, "kwargs": {"x": 5}}
+        )
+
+        # Verify histogram attributes
+        mock_histogram.record.assert_called_once()
+        call_args = mock_histogram.record.call_args
+        attributes = call_args[1]["attributes"]
+        assert attributes[GenAIAttributes.TOOL_CALL_ID.value] is None
+
+
+async def test_ai_function_invoke_invalid_pydantic_args():
+    """Test the ai_function invoke method with invalid Pydantic model arguments."""
+
+    @ai_function(name="invalid_args_test", description="A test tool for invalid args")
+    def invalid_args_test(x: int, y: int) -> int:
+        """A function for testing invalid Pydantic args."""
+        return x + y
+
+    # Create a different Pydantic model
+    class WrongModel(BaseModel):
+        a: str
+        b: str
+
+    wrong_args = WrongModel(a="hello", b="world")
+
+    # Call invoke with wrong model type
+    with pytest.raises(TypeError, match="Expected invalid_args_test_input, got WrongModel"):
+        await invalid_args_test.invoke(arguments=wrong_args)

python/packages/main/tests/main/test_types.py

@@ -0,0 +1,663 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import MutableSequence
+from typing import Any
+
+from pydantic import BaseModel, ValidationError
+from pytest import fixture, mark, raises
+
+# region: TextContent
+from agent_framework import (
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    AIContent,
+    AIContents,
+    AITool,
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    ChatRole,
+    ChatToolMode,
+    DataContent,
+    FunctionCallContent,
+    FunctionResultContent,
+    GeneratedEmbeddings,
+    StructuredResponse,
+    TextContent,
+    TextReasoningContent,
+    UriContent,
+    UsageDetails,
+    ai_function,
+)
+
+
+@fixture
+def ai_tool() -> AITool:
+    """Returns a generic AITool."""
+
+    class GenericTool(BaseModel):
+        name: str
+        description: str | None = None
+        additional_properties: dict[str, Any] | None = None
+
+        def parameters(self) -> dict[str, Any]:
+            """Return the parameters of the tool as a JSON schema."""
+            return {
+                "name": {"type": "string"},
+            }
+
+    return GenericTool(name="generic_tool", description="A generic tool")
+
+
+@fixture
+def ai_function_tool() -> AITool:
+    """Returns a executable AITool."""
+
+    @ai_function
+    def simple_function(x: int, y: int) -> int:
+        """A simple function that adds two numbers."""
+        return x + y
+
+    return simple_function
+
+
+def test_text_content_positional():
+    """Test the TextContent class to ensure it initializes correctly and inherits from AIContent."""
+    # Create an instance of TextContent
+    content = TextContent("Hello, world!", raw_representation="Hello, world!", additional_properties={"version": 1})
+
+    # Check the type and content
+    assert content.type == "text"
+    assert content.text == "Hello, world!"
+    assert content.raw_representation == "Hello, world!"
+    assert content.additional_properties["version"] == 1
+    # Ensure the instance is of type AIContent
+    assert isinstance(content, AIContent)
+    with raises(ValidationError):
+        content.type = "ai"
+
+
+def test_text_content_keyword():
+    """Test the TextContent class to ensure it initializes correctly and inherits from AIContent."""
+    # Create an instance of TextContent
+    content = TextContent(
+        text="Hello, world!", raw_representation="Hello, world!", additional_properties={"version": 1}
+    )
+
+    # Check the type and content
+    assert content.type == "text"
+    assert content.text == "Hello, world!"
+    assert content.raw_representation == "Hello, world!"
+    assert content.additional_properties["version"] == 1
+    # Ensure the instance is of type AIContent
+    assert isinstance(content, AIContent)
+    with raises(ValidationError):
+        content.type = "ai"
+
+
+# region: DataContent
+
+
+def test_data_content_bytes():
+    """Test the DataContent class to ensure it initializes correctly."""
+    # Create an instance of DataContent
+    content = DataContent(data=b"test", media_type="application/octet-stream", additional_properties={"version": 1})
+
+    # Check the type and content
+    assert content.type == "data"
+    assert content.uri == "data:application/octet-stream;base64,dGVzdA=="
+    assert content.additional_properties["version"] == 1
+
+    # Ensure the instance is of type AIContent
+    assert isinstance(content, AIContent)
+
+
+def test_data_content_uri():
+    """Test the DataContent class to ensure it initializes correctly with a URI."""
+    # Create an instance of DataContent with a URI
+    content = DataContent(uri="data:application/octet-stream;base64,dGVzdA==", additional_properties={"version": 1})
+
+    # Check the type and content
+    assert content.type == "data"
+    assert content.uri == "data:application/octet-stream;base64,dGVzdA=="
+    assert content.additional_properties["version"] == 1
+
+    # Ensure the instance is of type AIContent
+    assert isinstance(content, AIContent)
+
+
+def test_data_content_invalid():
+    """Test the DataContent class to ensure it raises an error for invalid initialization."""
+    # Attempt to create an instance of DataContent with invalid data
+    # not a proper uri
+    with raises(ValidationError):
+        DataContent(uri="invalid_uri")
+    # unknown media type
+    with raises(ValidationError):
+        DataContent(uri="data:application/random;base64,dGVzdA==")
+    # not valid base64 data
+
+    with raises(ValidationError):
+        DataContent(uri="data:application/json;base64,dGVzdA&")
+
+
+def test_data_content_empty():
+    """Test the DataContent class to ensure it raises an error for empty data."""
+    # Attempt to create an instance of DataContent with empty data
+    with raises(ValidationError):
+        DataContent(data=b"", media_type="application/octet-stream")
+
+    # Attempt to create an instance of DataContent with empty URI
+    with raises(ValidationError):
+        DataContent(uri="")
+
+
+# region: UriContent
+
+
+def test_uri_content():
+    """Test the UriContent class to ensure it initializes correctly."""
+    content = UriContent(uri="http://example.com", media_type="image/jpg", additional_properties={"version": 1})
+
+    # Check the type and content
+    assert content.type == "uri"
+    assert content.uri == "http://example.com"
+    assert content.media_type == "image/jpg"
+    assert content.additional_properties["version"] == 1
+
+    # Ensure the instance is of type AIContent
+    assert isinstance(content, AIContent)
+
+
+# region: FunctionCallContent
+
+
+def test_function_call_content():
+    """Test the FunctionCallContent class to ensure it initializes correctly."""
+    content = FunctionCallContent(call_id="1", name="example_function", arguments={"param1": "value1"})
+
+    # Check the type and content
+    assert content.type == "function_call"
+    assert content.name == "example_function"
+    assert content.arguments == {"param1": "value1"}
+
+    # Ensure the instance is of type AIContent
+    assert isinstance(content, AIContent)
+
+
+# region: FunctionResultContent
+
+
+def test_function_result_content():
+    """Test the FunctionResultContent class to ensure it initializes correctly."""
+    content = FunctionResultContent(call_id="1", result={"param1": "value1"})
+
+    # Check the type and content
+    assert content.type == "function_result"
+    assert content.result == {"param1": "value1"}
+
+    # Ensure the instance is of type AIContent
+    assert isinstance(content, AIContent)
+
+
+# region: UsageDetails
+
+
+def test_usage_details():
+    usage = UsageDetails(input_token_count=5, output_token_count=10, total_token_count=15)
+    assert usage.input_token_count == 5
+    assert usage.output_token_count == 10
+    assert usage.total_token_count == 15
+    assert usage.additional_counts == {}
+
+
+def test_usage_details_addition():
+    usage1 = UsageDetails(
+        input_token_count=5,
+        output_token_count=10,
+        total_token_count=15,
+        test1=10,
+        test2=20,
+    )
+    usage2 = UsageDetails(
+        input_token_count=3,
+        output_token_count=6,
+        total_token_count=9,
+        test1=10,
+        test3=30,
+    )
+
+    combined_usage = usage1 + usage2
+    assert combined_usage.input_token_count == 8
+    assert combined_usage.output_token_count == 16
+    assert combined_usage.total_token_count == 24
+    assert combined_usage.additional_counts["test1"] == 20
+    assert combined_usage.additional_counts["test2"] == 20
+    assert combined_usage.additional_counts["test3"] == 30
+
+
+def test_usage_details_fail():
+    with raises(ValidationError):
+        UsageDetails(input_token_count=5, output_token_count=10, total_token_count=15, wrong_type="42.923")
+
+
+def test_usage_details_additional_counts():
+    usage = UsageDetails(input_token_count=5, output_token_count=10, total_token_count=15, **{"test": 1})
+    assert usage.additional_counts["test"] == 1
+
+
+# region: AIContent Serialization
+
+
+@mark.parametrize(
+    "content_type, args",
+    [
+        (TextContent, {"text": "Hello, world!"}),
+        (DataContent, {"data": b"Hello, world!", "media_type": "text/plain"}),
+        (UriContent, {"uri": "http://example.com", "media_type": "text/html"}),
+        (FunctionCallContent, {"call_id": "1", "name": "example_function", "arguments": {}}),
+        (FunctionResultContent, {"call_id": "1", "result": {}}),
+    ],
+)
+def test_ai_content_serialization(content_type: type[AIContent], args: dict):
+    content = content_type(**args)
+    serialized = content.model_dump()
+    deserialized = content_type.model_validate(serialized)
+    assert deserialized == content
+
+    class TestModel(BaseModel):
+        content: AIContents
+
+    test_item = TestModel.model_validate({"content": serialized})
+
+    assert isinstance(test_item.content, content_type)
+
+
+# region: ChatMessage
+
+
+def test_chat_message_text():
+    """Test the ChatMessage class to ensure it initializes correctly with text content."""
+    # Create a ChatMessage with a role and text content
+    message = ChatMessage(role="user", text="Hello, how are you?")
+
+    # Check the type and content
+    assert message.role == ChatRole.USER
+    assert len(message.contents) == 1
+    assert isinstance(message.contents[0], TextContent)
+    assert message.contents[0].text == "Hello, how are you?"
+    assert message.text == "Hello, how are you?"
+
+    # Ensure the instance is of type AIContent
+    assert isinstance(message.contents[0], AIContent)
+
+
+def test_chat_message_contents():
+    """Test the ChatMessage class to ensure it initializes correctly with contents."""
+    # Create a ChatMessage with a role and multiple contents
+    content1 = TextContent("Hello, how are you?")
+    content2 = TextContent("I'm fine, thank you!")
+    message = ChatMessage(role="user", contents=[content1, content2])
+
+    # Check the type and content
+    assert message.role == ChatRole.USER
+    assert len(message.contents) == 2
+    assert isinstance(message.contents[0], TextContent)
+    assert isinstance(message.contents[1], TextContent)
+    assert message.contents[0].text == "Hello, how are you?"
+    assert message.contents[1].text == "I'm fine, thank you!"
+    assert message.text == "Hello, how are you? I'm fine, thank you!"
+
+
+# region: ChatResponse
+
+
+def test_chat_response():
+    """Test the ChatResponse class to ensure it initializes correctly with a message."""
+    # Create a ChatMessage
+    message = ChatMessage(role="assistant", text="I'm doing well, thank you!")
+
+    # Create a ChatResponse with the message
+    response = ChatResponse(messages=message)
+
+    # Check the type and content
+    assert response.messages[0].role == ChatRole.ASSISTANT
+    assert response.messages[0].text == "I'm doing well, thank you!"
+    assert isinstance(response.messages[0], ChatMessage)
+
+
+# region: StructuredResponse
+
+
+def test_structured_response():
+    """Test the StructuredResponse class to ensure it initializes correctly with a value."""
+
+    class ResponseModel(BaseModel):
+        content: str
+        action: str
+
+    # Create a StructuredResponse with a value
+    response = StructuredResponse[ResponseModel](
+        value=ResponseModel(content="Hello, world!", action="test"),
+        text="{'content': 'Hello, world!', 'action': 'test'}",
+    )
+
+    # Check the type and content
+    assert response.value == ResponseModel(content="Hello, world!", action="test")
+    assert isinstance(response, StructuredResponse)
+
+
+# region: ChatResponseUpdate
+
+
+def test_chat_response_update():
+    """Test the ChatResponseUpdate class to ensure it initializes correctly with a message."""
+    # Create a ChatMessage
+    message = TextContent(text="I'm doing well, thank you!")
+
+    # Create a ChatResponseUpdate with the message
+    response_update = ChatResponseUpdate(contents=[message])
+
+    # Check the type and content
+    assert response_update.contents[0].text == "I'm doing well, thank you!"
+    assert isinstance(response_update.contents[0], TextContent)
+
+
+def test_chat_response_updates_to_chat_response_one():
+    """Test converting ChatResponseUpdate to ChatResponse."""
+    # Create a ChatMessage
+    message1 = TextContent("I'm doing well, ")
+    message2 = TextContent("thank you!")
+
+    # Create a ChatResponseUpdate with the message
+    response_updates = [
+        ChatResponseUpdate(text=message1, message_id="1"),
+        ChatResponseUpdate(text=message2, message_id="1"),
+    ]
+
+    # Convert to ChatResponse
+    chat_response = ChatResponse.from_chat_response_updates(response_updates)
+
+    # Check the type and content
+    assert len(chat_response.messages) == 1
+    assert chat_response.text == "I'm doing well, thank you!"
+    assert isinstance(chat_response.messages[0], ChatMessage)
+    assert len(chat_response.messages[0].contents) == 1
+    assert chat_response.messages[0].message_id == "1"
+
+
+def test_chat_response_updates_to_chat_response_two():
+    """Test converting ChatResponseUpdate to ChatResponse."""
+    # Create a ChatMessage
+    message1 = TextContent("I'm doing well, ")
+    message2 = TextContent("thank you!")
+
+    # Create a ChatResponseUpdate with the message
+    response_updates = [
+        ChatResponseUpdate(text=message1, message_id="1"),
+        ChatResponseUpdate(text=message2, message_id="2"),
+    ]
+
+    # Convert to ChatResponse
+    chat_response = ChatResponse.from_chat_response_updates(response_updates)
+
+    # Check the type and content
+    assert len(chat_response.messages) == 2
+    assert chat_response.text == "I'm doing well, \nthank you!"
+    assert isinstance(chat_response.messages[0], ChatMessage)
+    assert chat_response.messages[0].message_id == "1"
+    assert isinstance(chat_response.messages[1], ChatMessage)
+    assert chat_response.messages[1].message_id == "2"
+
+
+def test_chat_response_updates_to_chat_response_multiple():
+    """Test converting ChatResponseUpdate to ChatResponse."""
+    # Create a ChatMessage
+    message1 = TextContent("I'm doing well, ")
+    message2 = TextContent("thank you!")
+
+    # Create a ChatResponseUpdate with the message
+    response_updates = [
+        ChatResponseUpdate(text=message1, message_id="1"),
+        ChatResponseUpdate(contents=[TextReasoningContent(text="Additional context")], message_id="1"),
+        ChatResponseUpdate(text=message2, message_id="1"),
+    ]
+
+    # Convert to ChatResponse
+    chat_response = ChatResponse.from_chat_response_updates(response_updates)
+
+    # Check the type and content
+    assert len(chat_response.messages) == 1
+    assert chat_response.text == "I'm doing well,  thank you!"
+    assert isinstance(chat_response.messages[0], ChatMessage)
+    assert len(chat_response.messages[0].contents) == 3
+    assert chat_response.messages[0].message_id == "1"
+
+
+def test_chat_response_updates_to_chat_response_multiple_multiple():
+    """Test converting ChatResponseUpdate to ChatResponse."""
+    # Create a ChatMessage
+    message1 = TextContent("I'm doing well, ")
+    message2 = TextContent("thank you!")
+
+    # Create a ChatResponseUpdate with the message
+    response_updates = [
+        ChatResponseUpdate(text=message1, message_id="1"),
+        ChatResponseUpdate(text=message2, message_id="1"),
+        ChatResponseUpdate(contents=[TextReasoningContent(text="Additional context")], message_id="1"),
+        ChatResponseUpdate(contents=[TextContent(text="More context")], message_id="1"),
+        ChatResponseUpdate(text="Final part", message_id="1"),
+    ]
+
+    # Convert to ChatResponse
+    chat_response = ChatResponse.from_chat_response_updates(response_updates)
+
+    # Check the type and content
+    assert len(chat_response.messages) == 1
+    assert isinstance(chat_response.messages[0], ChatMessage)
+    assert chat_response.messages[0].message_id == "1"
+
+    assert len(chat_response.messages[0].contents) == 3
+    assert isinstance(chat_response.messages[0].contents[0], TextContent)
+    assert chat_response.messages[0].contents[0].text == "I'm doing well, thank you!"
+    assert isinstance(chat_response.messages[0].contents[1], TextReasoningContent)
+    assert chat_response.messages[0].contents[1].text == "Additional context"
+    assert isinstance(chat_response.messages[0].contents[2], TextContent)
+    assert chat_response.messages[0].contents[2].text == "More contextFinal part"
+
+    assert chat_response.text == "I'm doing well, thank you! More contextFinal part"
+
+
+# region: ChatToolMode
+
+
+def test_chat_tool_mode():
+    """Test the ChatToolMode class to ensure it initializes correctly."""
+    # Create instances of ChatToolMode
+    auto_mode = ChatToolMode.AUTO
+    required_any = ChatToolMode.REQUIRED_ANY
+    required_mode = ChatToolMode.REQUIRED("example_function")
+    none_mode = ChatToolMode.NONE
+
+    # Check the type and content
+    assert auto_mode.mode == "auto"
+    assert auto_mode.required_function_name is None
+    assert required_any.mode == "required"
+    assert required_any.required_function_name is None
+    assert required_mode.mode == "required"
+    assert required_mode.required_function_name == "example_function"
+    assert none_mode.mode == "none"
+    assert none_mode.required_function_name is None
+
+    # Ensure the instances are of type ChatToolMode
+    assert isinstance(auto_mode, ChatToolMode)
+    assert isinstance(required_any, ChatToolMode)
+    assert isinstance(required_mode, ChatToolMode)
+    assert isinstance(none_mode, ChatToolMode)
+
+    assert ChatToolMode.REQUIRED("example_function") == ChatToolMode.REQUIRED("example_function")
+
+
+def test_chat_tool_mode_from_dict():
+    """Test creating ChatToolMode from a dictionary."""
+    mode_dict = {"mode": "required", "required_function_name": "example_function"}
+    mode = ChatToolMode(**mode_dict)
+
+    # Check the type and content
+    assert mode.mode == "required"
+    assert mode.required_function_name == "example_function"
+
+    # Ensure the instance is of type ChatToolMode
+    assert isinstance(mode, ChatToolMode)
+
+
+def test_generated_embeddings():
+    """Test the GeneratedEmbeddings class to ensure it initializes correctly."""
+    # Create an instance of GeneratedEmbeddings
+    embeddings = GeneratedEmbeddings(embeddings=[[0.1, 0.2, 0.3]])
+
+    # Check the type and content
+    assert embeddings.embeddings == [[0.1, 0.2, 0.3]]
+
+    # Ensure the instance is of type GeneratedEmbeddings
+    assert isinstance(embeddings, GeneratedEmbeddings)
+    assert issubclass(GeneratedEmbeddings, MutableSequence)
+
+
+# region: ChatOptions
+
+
+def test_chat_options_init() -> None:
+    options = ChatOptions()
+    assert options.ai_model_id is None
+
+
+def test_chat_options_init_with_args(ai_function_tool, ai_tool) -> None:
+    options = ChatOptions(
+        ai_model_id="gpt-4",
+        max_tokens=1024,
+        temperature=0.7,
+        top_p=0.9,
+        presence_penalty=0.0,
+        frequency_penalty=0.0,
+        user="user-123",
+        tools=[ai_function_tool, ai_tool],
+    )
+    assert options.ai_model_id == "gpt-4"
+    assert options.max_tokens == 1024
+    assert options.temperature == 0.7
+    assert options.top_p == 0.9
+    assert options.presence_penalty == 0.0
+    assert options.frequency_penalty == 0.0
+    assert options.user == "user-123"
+    for tool in options._ai_tools:
+        assert isinstance(tool, AITool)
+        assert tool.name is not None
+        assert tool.description is not None
+        assert tool.parameters() is not None
+
+
+def test_chat_options_and(ai_function_tool, ai_tool) -> None:
+    options1 = ChatOptions(ai_model_id="gpt-4o", tools=[ai_function_tool])
+    options2 = ChatOptions(ai_model_id="gpt-4.1", tools=[ai_tool])
+    assert options1 != options2
+    options3 = options1 & options2
+
+    assert options3.ai_model_id == "gpt-4.1"
+    assert len(options3._ai_tools) == 2
+    assert options3._ai_tools == [ai_function_tool, ai_tool]
+    assert options3.tools == [ai_function_tool, ai_tool]
+
+
+# region Agent Response Fixtures
+
+
+@fixture
+def chat_message() -> ChatMessage:
+    return ChatMessage(role=ChatRole.USER, text="Hello")
+
+
+@fixture
+def text_content() -> TextContent:
+    return TextContent(text="Test content")
+
+
+@fixture
+def agent_run_response(chat_message: ChatMessage) -> AgentRunResponse:
+    return AgentRunResponse(messages=chat_message)
+
+
+@fixture
+def agent_run_response_update(text_content: TextContent) -> AgentRunResponseUpdate:
+    return AgentRunResponseUpdate(role=ChatRole.ASSISTANT, contents=[text_content])
+
+
+# region AgentRunResponse
+
+
+def test_agent_run_response_init_single_message(chat_message: ChatMessage) -> None:
+    response = AgentRunResponse(messages=chat_message)
+    assert response.messages == [chat_message]
+
+
+def test_agent_run_response_init_list_messages(chat_message: ChatMessage) -> None:
+    response = AgentRunResponse(messages=[chat_message, chat_message])
+    assert len(response.messages) == 2
+    assert response.messages[0] == chat_message
+
+
+def test_agent_run_response_init_none_messages() -> None:
+    response = AgentRunResponse()
+    assert response.messages == []
+
+
+def test_agent_run_response_text_property(chat_message: ChatMessage) -> None:
+    response = AgentRunResponse(messages=[chat_message, chat_message])
+    assert response.text == "HelloHello"
+
+
+def test_agent_run_response_text_property_empty() -> None:
+    response = AgentRunResponse()
+    assert response.text == ""
+
+
+def test_agent_run_response_from_updates(agent_run_response_update: AgentRunResponseUpdate) -> None:
+    updates = [agent_run_response_update, agent_run_response_update]
+    response = AgentRunResponse.from_agent_run_response_updates(updates)
+    assert len(response.messages) > 0
+    assert response.text == "Test contentTest content"
+
+
+def test_agent_run_response_str_method(chat_message: ChatMessage) -> None:
+    response = AgentRunResponse(messages=chat_message)
+    assert str(response) == "Hello"
+
+
+# region AgentRunResponseUpdate
+
+
+def test_agent_run_response_update_init_content_list(text_content: TextContent) -> None:
+    update = AgentRunResponseUpdate(contents=[text_content, text_content])
+    assert len(update.contents) == 2
+    assert update.contents[0] == text_content
+
+
+def test_agent_run_response_update_init_none_content() -> None:
+    update = AgentRunResponseUpdate()
+    assert update.contents == []
+
+
+def test_agent_run_response_update_text_property(text_content: TextContent) -> None:
+    update = AgentRunResponseUpdate(contents=[text_content, text_content])
+    assert update.text == "Test contentTest content"
+
+
+def test_agent_run_response_update_text_property_empty() -> None:
+    update = AgentRunResponseUpdate()
+    assert update.text == ""
+
+
+def test_agent_run_response_update_str_method(text_content: TextContent) -> None:
+    update = AgentRunResponseUpdate(contents=[text_content])
+    assert str(update) == "Test content"