Python: added AIAnnotation types and extra tests (#374)

* added AIAnnotation types and extra tests * fixed typing and such * use copy * fix raw representation for add * handle annotations in add * clarified concat * self to first
2026-06-16 21:04:09 +08:00 · 2025-08-08 22:27:31 +02:00
parent 8f27e63df6
commit 82d39bc1f7
8 changed files with 934 additions and 346 deletions
@@ -39,7 +39,7 @@ __all__ = [
    "use_tool_calling",
 ]

-# region: Tool Calling Functions and Decorators
+# region Tool Calling Functions and Decorators


 async def _auto_invoke_function(
@@ -253,7 +253,7 @@ def use_tool_calling(cls: type[TChatClientBase]) -> type[TChatClientBase]:
    return cls


-# region: ChatClient Protocol
+# region ChatClient Protocol


@runtime_checkable
@@ -675,7 +675,7 @@ class ChatClientBase(AFBaseModel, ABC):
        return ChatClientAgent(chat_client=self, name=name, instructions=instructions, tools=tools, **kwargs)


-# region: Embedding Client
+# region Embedding Client


@runtime_checkable
@@ -14,6 +14,7 @@ from collections.abc import (
    MutableSequence,
    Sequence,
 )
+from copy import deepcopy
 from typing import Annotated, Any, ClassVar, Generic, Literal, TypeVar, overload

 from pydantic import (
@@ -36,7 +37,7 @@ if sys.version_info >= (3, 11):
 else:
    from typing_extensions import Self  # pragma: no cover

-# region: Constants and types
+# region Constants and types
 _T = TypeVar("_T")
 TValue = TypeVar("TValue")
 TEmbedding = TypeVar("TEmbedding")
@@ -76,11 +77,15 @@ KNOWN_MEDIA_TYPES = [


 __all__ = [
+    "AIAnnotation",
+    "AIAnnotations",
    "AIContent",
    "AIContents",
    "AITool",
    "AgentRunResponse",
    "AgentRunResponseUpdate",
+    "AnnotatedRegion",
+    "AnnotatedRegions",
    "ChatFinishReason",
    "ChatMessage",
    "ChatOptions",
@@ -88,6 +93,7 @@ __all__ = [
    "ChatResponseUpdate",
    "ChatRole",
    "ChatToolMode",
+    "CitationAnnotation",
    "DataContent",
    "ErrorContent",
    "FunctionCallContent",
@@ -97,6 +103,7 @@ __all__ = [
    "StructuredResponse",
    "TextContent",
    "TextReasoningContent",
+    "TextSpanRegion",
    "TextToSpeechOptions",
    "UriContent",
    "UsageContent",
@@ -265,27 +272,24 @@ def _coalesce_text_content(
    if not contents:
        return
    coalesced_contents: list["AIContents"] = []
-    current_texts: list[str] = []
-    first_new_content = None
-    for i, content in enumerate(contents):
+    first_new_content: Any | None = None
+    for content in contents:
        if isinstance(content, type_):
-            current_texts.append(content.text)  # type: ignore[union-attr]
            if first_new_content is None:
-                first_new_content = i
+                first_new_content = deepcopy(content)
+            else:
+                first_new_content += content
        else:
-            if first_new_content is not None:
-                new_content = type_(text="".join(current_texts))
-                new_content.raw_representation = contents[first_new_content].raw_representation
-                new_content.additional_properties = contents[first_new_content].additional_properties
-                # Store the replacement node. We inherit the properties of the first text node. We don't
-                # currently propagate additional properties from the subsequent nodes. If we ever need to,
-                # we can add that here.
-                coalesced_contents.append(new_content)
-                current_texts = []
-                first_new_content = None
+            # skip this content, it is not of the right type
+            # so write the existing one to the list and start a new one,
+            # once the right type is found again
+            if first_new_content:
+                coalesced_contents.append(first_new_content)
+            first_new_content = None
+            # but keep the other content in the new list
            coalesced_contents.append(content)
-    if current_texts:
-        coalesced_contents.append(type_(text="".join(current_texts)))
+    if first_new_content:
+        coalesced_contents.append(first_new_content)
    contents.clear()
    contents.extend(coalesced_contents)

@@ -297,7 +301,81 @@ def _finalize_response(response: "ChatResponse | AgentRunResponse") -> None:
        _coalesce_text_content(msg.contents, TextReasoningContent)


-# region: AIContent
+# region AIAnnotation
+
+
+class AnnotatedRegion(AFBaseModel):
+    """Represents a collection of annotated regions.
+
+    Attributes:
+        regions: A list of regions that have been annotated.
+        additional_properties: Optional additional properties associated with the content.
+        raw_representation: Optional raw representation of the content from an underlying implementation.
+    """
+
+    type: Literal["annotated_regions"] = "annotated_regions"  # type: ignore[assignment]
+
+
+class TextSpanRegion(AnnotatedRegion):
+    """Represents a region of text that has been annotated."""
+
+    type: Literal["text_span"] = "text_span"  # type: ignore[assignment]
+    start_index: int | None = None
+    end_index: int | None = None
+
+
+AnnotatedRegions = Annotated[
+    TextSpanRegion | AnnotatedRegion,
+    Field(discriminator="type"),
+]
+
+
+class AIAnnotation(AFBaseModel):
+    """Base class for all AI Annotation types.
+
+    Args:
+        type: The type of content, which is always "ai_annotation" for this class.
+        additional_properties: Optional additional properties associated with the content.
+        raw_representation: Optional raw representation of the content from an underlying implementation.
+
+    """
+
+    type: Literal["ai_annotation"] = "ai_annotation"
+    annotated_regions: list[AnnotatedRegions] | None = None
+    additional_properties: dict[str, Any] | None = None
+    raw_representation: Any | None = Field(default=None, repr=False)
+
+
+class CitationAnnotation(AIAnnotation):
+    """Represents a citation annotation.
+
+    Attributes:
+        type: The type of content, which is always "citation" for this class.
+        title: The title of the cited content.
+        url: The URL of the cited content.
+        file_id: The file identifier of the cited content, if applicable.
+        tool_name: The name of the tool that generated the citation, if applicable.
+        snippet: A snippet of the cited content, if applicable.
+        annotated_regions: A list of regions that have been annotated with this citation.
+        additional_properties: Optional additional properties associated with the content.
+        raw_representation: Optional raw representation of the content from an underlying implementation.
+    """
+
+    type: Literal["citation"] = "citation"  # type: ignore[assignment]
+    title: str | None = None
+    url: str | None = None
+    file_id: str | None = None
+    tool_name: str | None = None
+    snippet: str | None = None
+
+
+AIAnnotations = Annotated[
+    CitationAnnotation | AIAnnotation,
+    Field(discriminator="type"),
+]
+
+
+# region AIContent


 class AIContent(AFBaseModel):
@@ -305,12 +383,14 @@ class AIContent(AFBaseModel):

    Attributes:
        type: The type of content, which is always "ai" for this class.
+        annotations: Optional annotations associated with the content.
        additional_properties: Optional additional properties associated with the content.
        raw_representation: Optional raw representation of the content from an underlying implementation.

    """

    type: Literal["ai"] = "ai"
+    annotations: list[AIAnnotations] | None = None
    additional_properties: dict[str, Any] | None = None
    raw_representation: Any | None = Field(default=None, repr=False)

@@ -321,6 +401,7 @@ class TextContent(AIContent):
    Attributes:
        text: The text content represented by this instance.
        type: The type of content, which is always "text" for this class.
+        annotations: Optional annotations associated with the content.
        additional_properties: Optional additional properties associated with the content.
        raw_representation: Optional raw representation of the content.
    """
@@ -351,6 +432,69 @@ class TextContent(AIContent):
            **kwargs,
        )

+    def __add__(self, other: "TextContent") -> "TextContent":
+        """Concatenate two TextContent instances.
+
+        The following things happen:
+        The text is concatenated.
+        The annotations are combined.
+        The additional properties are merged, with the values of shared keys of the first instance taking precedence.
+        The raw_representations are combined into a list of them, if they both have one.
+        """
+        if not isinstance(other, TextContent):
+            raise TypeError("Incompatible type")
+        if self.raw_representation is None:
+            raw_representation = other.raw_representation
+        elif other.raw_representation is None:
+            raw_representation = self.raw_representation
+        else:
+            raw_representation = (
+                self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation]
+            ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation])
+        if self.annotations is None:
+            annotations = other.annotations
+        elif other.annotations is None:
+            annotations = self.annotations
+        else:
+            annotations = self.annotations + other.annotations
+        return TextContent(
+            text=self.text + other.text,
+            annotations=annotations,
+            additional_properties={
+                **(other.additional_properties or {}),
+                **(self.additional_properties or {}),
+            },
+            raw_representation=raw_representation,
+        )
+
+    def __iadd__(self, other: "TextContent") -> Self:
+        """In-place concatenation of two TextContent instances.
+
+        The following things happen:
+        The text is concatenated.
+        The annotations are combined.
+        The additional properties are merged, with the values of shared keys of the first instance taking precedence.
+        The raw_representations are combined into a list of them, if they both have one.
+        """
+        if not isinstance(other, TextContent):
+            raise TypeError("Incompatible type")
+        self.text += other.text
+        if self.additional_properties is None:
+            self.additional_properties = {}
+        if other.additional_properties:
+            self.additional_properties = {**other.additional_properties, **self.additional_properties}
+        if self.raw_representation is None:
+            self.raw_representation = other.raw_representation
+        elif other.raw_representation is not None:
+            self.raw_representation = (
+                self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation]
+            ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation])
+        if other.annotations:
+            if self.annotations is None:
+                self.annotations = []
+            self.annotations.extend(other.annotations)
+        return self
+

 class TextReasoningContent(AIContent):
    """Represents text reasoning content in a chat.
@@ -361,14 +505,11 @@ class TextReasoningContent(AIContent):
    Attributes:
        text: The text content represented by this instance.
        type: The type of content, which is always "text_reasoning" for this class.
+        annotations: Optional annotations associated with the content.
        additional_properties: Optional additional properties associated with the content.
        raw_representation: Optional raw representation of the content.
-
-
    """

-    # TODO(eavanvalkenburg): Should we merge these two classes, and use a property to distinguish them?
-
    text: str
    type: Literal["text_reasoning"] = "text_reasoning"  # type: ignore[assignment]

@@ -395,6 +536,66 @@ class TextReasoningContent(AIContent):
            **kwargs,
        )

+    def __add__(self, other: "TextReasoningContent") -> "TextReasoningContent":
+        """Concatenate two TextReasoningContent instances.
+
+        The following things happen:
+        The text is concatenated.
+        The annotations are combined.
+        The additional properties are merged, with the values of shared keys of the first instance taking precedence.
+        The raw_representations are combined into a list of them, if they both have one.
+        """
+        if not isinstance(other, TextReasoningContent):
+            raise TypeError("Incompatible type")
+        if self.raw_representation is None:
+            raw_representation = other.raw_representation
+        elif other.raw_representation is None:
+            raw_representation = self.raw_representation
+        else:
+            raw_representation = (
+                self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation]
+            ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation])
+        if self.annotations is None:
+            annotations = other.annotations
+        elif other.annotations is None:
+            annotations = self.annotations
+        else:
+            annotations = self.annotations + other.annotations
+        return TextReasoningContent(
+            text=self.text + other.text,
+            annotations=annotations,
+            additional_properties={**(self.additional_properties or {}), **(other.additional_properties or {})},
+            raw_representation=raw_representation,
+        )
+
+    def __iadd__(self, other: "TextReasoningContent") -> Self:
+        """In-place concatenation of two TextReasoningContent instances.
+
+        The following things happen:
+        The text is concatenated.
+        The annotations are combined.
+        The additional properties are merged, with the values of shared keys of the first instance taking precedence.
+        The raw_representations are combined into a list of them, if they both have one.
+        """
+        if not isinstance(other, TextReasoningContent):
+            raise TypeError("Incompatible type")
+        self.text += other.text
+        if self.additional_properties is None:
+            self.additional_properties = {}
+        if other.additional_properties:
+            self.additional_properties.update(other.additional_properties)
+        if self.raw_representation is None:
+            self.raw_representation = other.raw_representation
+        elif other.raw_representation is not None:
+            self.raw_representation = (
+                self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation]
+            ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation])
+        if other.annotations:
+            if self.annotations is None:
+                self.annotations = []
+            self.annotations.extend(other.annotations)
+        return self
+

 class DataContent(AIContent):
    """Represents binary data content with an associated media type (also known as a MIME type).
@@ -402,8 +603,9 @@ class DataContent(AIContent):
    Attributes:
        uri: The URI of the data represented by this instance, typically in the form of a data URI.
            Should be in the form: "data:{media_type};base64,{base64_data}".
-        type: The type of content, which is always "data" for this class.
        media_type: The media type of the data.
+        type: The type of content, which is always "data" for this class.
+        annotations: Optional annotations associated with the content.
        additional_properties: Optional additional properties associated with the content.
        raw_representation: Optional raw representation of the content.

@@ -418,6 +620,7 @@ class DataContent(AIContent):
        self,
        *,
        uri: str,
+        annotations: list[AIAnnotations] | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
        **kwargs: Any,
@@ -431,6 +634,7 @@ class DataContent(AIContent):
        Args:
            uri: The URI of the data represented by this instance.
                Should be in the form: "data:{media_type};base64,{base64_data}".
+            annotations: Optional annotations associated with the content.
            additional_properties: Optional additional properties associated with the content.
            raw_representation: Optional raw representation of the content.
            **kwargs: Any additional keyword arguments.
@@ -442,6 +646,7 @@ class DataContent(AIContent):
        *,
        data: bytes,
        media_type: str,
+        annotations: list[AIAnnotations] | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
        **kwargs: Any,
@@ -456,6 +661,7 @@ class DataContent(AIContent):
            data: The binary data represented by this instance.
                The data is transformed into a base64-encoded data URI.
            media_type: The media type of the data.
+            annotations: Optional annotations associated with the content.
            additional_properties: Optional additional properties associated with the content.
            raw_representation: Optional raw representation of the content.
            **kwargs: Any additional keyword arguments.
@@ -467,6 +673,7 @@ class DataContent(AIContent):
        uri: str | None = None,
        data: bytes | None = None,
        media_type: str | None = None,
+        annotations: list[AIAnnotations] | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
        **kwargs: Any,
@@ -483,6 +690,7 @@ class DataContent(AIContent):
            data: The binary data represented by this instance.
                The data is transformed into a base64-encoded data URI.
            media_type: The media type of the data.
+            annotations: Optional annotations associated with the content.
            additional_properties: Optional additional properties associated with the content.
            raw_representation: Optional raw representation of the content.
            **kwargs: Any additional keyword arguments.
@@ -494,6 +702,7 @@ class DataContent(AIContent):
        super().__init__(
            uri=uri,  # type: ignore[reportCallIssue]
            media_type=media_type,  # type: ignore[reportCallIssue]
+            annotations=annotations,
            raw_representation=raw_representation,
            additional_properties=additional_properties,
            **kwargs,
@@ -529,6 +738,7 @@ class UriContent(AIContent):
        uri: The URI of the content, e.g., 'https://example.com/image.png'.
        media_type: The media type of the content, e.g., 'image/png', 'application/json', etc.
        type: The type of content, which is always "uri" for this class.
+        annotations: Optional annotations associated with the content.
        additional_properties: Optional additional properties associated with the content.
        raw_representation: Optional raw representation of the content.

@@ -543,6 +753,7 @@ class UriContent(AIContent):
        uri: str,
        media_type: str,
        *,
+        annotations: list[AIAnnotations] | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
        **kwargs: Any,
@@ -556,6 +767,7 @@ class UriContent(AIContent):
        Args:
            uri: The URI of the content.
            media_type: The media type of the content.
+            annotations: Optional annotations associated with the content.
            additional_properties: Optional additional properties associated with the content.
            raw_representation: Optional raw representation of the content.
            **kwargs: Any additional keyword arguments.
@@ -563,6 +775,7 @@ class UriContent(AIContent):
        super().__init__(
            uri=uri,  # type: ignore[reportCallIssue]
            media_type=media_type,  # type: ignore[reportCallIssue]
+            annotations=annotations,
            additional_properties=additional_properties,
            raw_representation=raw_representation,
            **kwargs,
@@ -590,10 +803,11 @@ class ErrorContent(AIContent):
        but the operation was still able to continue.

    Attributes:
-        type: The type of content, which is always "error" for this class.
        error_code: The error code associated with the error.
        details: Additional details about the error.
        message: The error message.
+        type: The type of content, which is always "error" for this class.
+        annotations: Optional annotations associated with the content.
        additional_properties: Optional additional properties associated with the content.
        raw_representation: Optional raw representation of the content.

@@ -611,6 +825,7 @@ class ErrorContent(AIContent):
        message: str | None = None,
        error_code: str | None = None,
        details: str | None = None,
+        annotations: list[AIAnnotations] | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
        **kwargs: Any,
@@ -621,6 +836,7 @@ class ErrorContent(AIContent):
            message: The error message.
            error_code: The error code associated with the error.
            details: Additional details about the error.
+            annotations: Optional annotations associated with the content.
            additional_properties: Optional additional properties associated with the content.
            raw_representation: Optional raw representation of the content.
            **kwargs: Any additional keyword arguments.
@@ -629,6 +845,7 @@ class ErrorContent(AIContent):
            message=message,  # type: ignore[reportCallIssue]
            error_code=error_code,  # type: ignore[reportCallIssue]
            details=details,  # type: ignore[reportCallIssue]
+            annotations=annotations,
            additional_properties=additional_properties,
            raw_representation=raw_representation,
            **kwargs,
@@ -643,11 +860,12 @@ class FunctionCallContent(AIContent):
    """Represents a function call request.

    Attributes:
-        type: The type of content, which is always "function_call" for this class.
        call_id: The function call identifier.
        name: The name of the function requested.
        arguments: The arguments requested to be provided to the function.
        exception: Any exception that occurred while mapping the original function call data to this representation.
+        type: The type of content, which is always "function_call" for this class.
+        annotations: Optional annotations associated with the content.
        additional_properties: Optional additional properties associated with the content.
        raw_representation: Optional raw representation of the content.

@@ -666,6 +884,7 @@ class FunctionCallContent(AIContent):
        name: str,
        arguments: str | dict[str, Any | None] | None = None,
        exception: Exception | None = None,
+        annotations: list[AIAnnotations] | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
        **kwargs: Any,
@@ -678,6 +897,7 @@ class FunctionCallContent(AIContent):
            arguments: The arguments requested to be provided to the function,
                can be a string to allow gradual completion of the args.
            exception: Any exception that occurred while mapping the original function call data to this representation.
+            annotations: Optional annotations associated with the content.
            additional_properties: Optional additional properties associated with the content.
            raw_representation: Optional raw representation of the content.
            **kwargs: Any additional keyword arguments.
@@ -687,6 +907,7 @@ class FunctionCallContent(AIContent):
            name=name,  # type: ignore[reportCallIssue]
            arguments=arguments,  # type: ignore[reportCallIssue]
            exception=exception,  # type: ignore[reportCallIssue]
+            annotations=annotations,
            raw_representation=raw_representation,
            additional_properties=additional_properties,
            **kwargs,
@@ -733,10 +954,11 @@ class FunctionResultContent(AIContent):
    """Represents the result of a function call.

    Attributes:
-        type: The type of content, which is always "function_result" for this class.
        call_id: The identifier of the function call for which this is the result.
        result: The result of the function call, or a generic error message if the function call failed.
        exception: An exception that occurred if the function call failed.
+        type: The type of content, which is always "function_result" for this class.
+        annotations: Optional annotations associated with the content.
        additional_properties: Optional additional properties associated with the content.
        raw_representation: Optional raw representation of the content.

@@ -753,6 +975,7 @@ class FunctionResultContent(AIContent):
        call_id: str,
        result: Any | None = None,
        exception: Exception | None = None,
+        annotations: list[AIAnnotations] | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
        **kwargs: Any,
@@ -763,6 +986,7 @@ class FunctionResultContent(AIContent):
            call_id: The identifier of the function call for which this is the result.
            result: The result of the function call, or a generic error message if the function call failed.
            exception: An exception that occurred if the function call failed.
+            annotations: Optional annotations associated with the content.
            additional_properties: Optional additional properties associated with the content.
            raw_representation: Optional raw representation of the content.
            **kwargs: Any additional keyword arguments.
@@ -771,8 +995,9 @@ class FunctionResultContent(AIContent):
            call_id=call_id,  # type: ignore[reportCallIssue]
            result=result,  # type: ignore[reportCallIssue]
            exception=exception,  # type: ignore[reportCallIssue]
-            raw_representation=raw_representation,
+            annotations=annotations,
            additional_properties=additional_properties,
+            raw_representation=raw_representation,
            **kwargs,
        )

@@ -781,8 +1006,9 @@ class UsageContent(AIContent):
    """Represents usage information associated with a chat request and response.

    Attributes:
-        type: The type of content, which is always "usage" for this class.
        details: The usage information, including input and output token counts, and any additional counts.
+        type: The type of content, which is always "usage" for this class.
+        annotations: Optional annotations associated with the content.
        additional_properties: Optional additional properties associated with the content.
        raw_representation: Optional raw representation of the content.

@@ -795,6 +1021,7 @@ class UsageContent(AIContent):
        self,
        details: UsageDetails,
        *,
+        annotations: list[AIAnnotations] | None = None,
        additional_properties: dict[str, Any] | None = None,
        raw_representation: Any | None = None,
        **kwargs: Any,
@@ -802,8 +1029,9 @@ class UsageContent(AIContent):
        """Initializes a UsageContent instance."""
        super().__init__(
            details=details,  # type: ignore[reportCallIssue]
-            raw_representation=raw_representation,
+            annotations=annotations,
            additional_properties=additional_properties,
+            raw_representation=raw_representation,
            **kwargs,
        )

@@ -820,7 +1048,7 @@ AIContents = Annotated[
    Field(discriminator="type"),
 ]

-# region: Chat Response constants
+# region Chat Response constants


 class ChatRole(AFBaseModel):
@@ -891,7 +1119,7 @@ ChatFinishReason.LENGTH = ChatFinishReason(value="length")  # type: ignore[assig
 ChatFinishReason.STOP = ChatFinishReason(value="stop")  # type: ignore[assignment]
 ChatFinishReason.TOOL_CALLS = ChatFinishReason(value="tool_calls")  # type: ignore[assignment]

-# region: ChatMessage
+# region ChatMessage


 class ChatMessage(AFBaseModel):
@@ -1000,7 +1228,7 @@ class ChatMessage(AFBaseModel):
        return " ".join(content.text for content in self.contents if isinstance(content, TextContent))


-# region: ChatResponse
+# region ChatResponse


 class ChatResponse(AFBaseModel):
@@ -1261,7 +1489,7 @@ class StructuredResponse(ChatResponse, Generic[TValue]):
        )


-# region: ChatResponseUpdate
+# region ChatResponseUpdate


 class ChatResponseUpdate(AFBaseModel):
@@ -1404,7 +1632,7 @@ class ChatResponseUpdate(AFBaseModel):
        )


-# region: ChatOptions
+# region ChatOptions


 class ChatToolMode(AFBaseModel):
@@ -1574,7 +1802,7 @@ class ChatOptions(AFBaseModel):
        return combined


-# region: GeneratedEmbeddings
+# region GeneratedEmbeddings


 class GeneratedEmbeddings(AFBaseModel, MutableSequence[TEmbedding], Generic[TEmbedding]):
@@ -1791,7 +2019,7 @@ class AgentRunResponseUpdate(AFBaseModel):
        return self.text


-# region: SpeechToTextOptions
+# region SpeechToTextOptions


 class SpeechToTextOptions(AFBaseModel):
@@ -1826,7 +2054,7 @@ class SpeechToTextOptions(AFBaseModel):
        return settings


-# region: TextToSpeechOptions
+# region TextToSpeechOptions


 class TextToSpeechOptions(AFBaseModel):
@@ -19,7 +19,7 @@ class AgentExecutionException(AgentException):
    pass


-# region: Service Exceptions
+# region Service Exceptions


 class ServiceException(AgentFrameworkException):
@@ -177,11 +177,11 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
        content = choice.message if isinstance(choice, Choice) else choice.delta
        if content and (tool_calls := getattr(content, "tool_calls", None)) is not None:
            for tool in cast(list[ChatCompletionMessageToolCall] | list[ChoiceDeltaToolCall], tool_calls):
-                if tool.function:
+                if tool.function:  # type: ignore[reportAttributeAccessIssue, union-attr]
                    fcc = FunctionCallContent(
                        call_id=tool.id if tool.id else "",
-                        name=tool.function.name if tool.function.name else "",
-                        arguments=tool.function.arguments if tool.function.arguments else "",
+                        name=tool.function.name if tool.function.name else "",  # type: ignore
+                        arguments=tool.function.arguments if tool.function.arguments else "",  # type: ignore
                    )
                    resp.append(fcc)

@@ -1,32 +1,40 @@
 # Copyright (c) Microsoft. All rights reserved.

-from collections.abc import MutableSequence
+from collections.abc import AsyncIterable, MutableSequence
 from typing import Any

 from pydantic import BaseModel, ValidationError
 from pytest import fixture, mark, raises

-# region: TextContent
 from agent_framework import (
    AgentRunResponse,
    AgentRunResponseUpdate,
+    AIAnnotation,
    AIContent,
    AIContents,
    AITool,
+    AnnotatedRegion,
+    ChatFinishReason,
    ChatMessage,
    ChatOptions,
    ChatResponse,
    ChatResponseUpdate,
    ChatRole,
    ChatToolMode,
+    CitationAnnotation,
    DataContent,
+    ErrorContent,
    FunctionCallContent,
    FunctionResultContent,
    GeneratedEmbeddings,
+    SpeechToTextOptions,
    StructuredResponse,
    TextContent,
    TextReasoningContent,
+    TextSpanRegion,
+    TextToSpeechOptions,
    UriContent,
+    UsageContent,
    UsageDetails,
    ai_function,
 )
@@ -62,6 +70,9 @@ def ai_function_tool() -> AITool:
    return simple_function


+# region TextContent
+
+
 def test_text_content_positional():
    """Test the TextContent class to ensure it initializes correctly and inherits from AIContent."""
    # Create an instance of TextContent
@@ -96,7 +107,7 @@ def test_text_content_keyword():
        content.type = "ai"


-# region: DataContent
+# region DataContent


 def test_data_content_bytes():
@@ -107,6 +118,8 @@ def test_data_content_bytes():
    # Check the type and content
    assert content.type == "data"
    assert content.uri == "data:application/octet-stream;base64,dGVzdA=="
+    assert content.has_top_level_media_type("application") is True
+    assert content.has_top_level_media_type("image") is False
    assert content.additional_properties["version"] == 1

    # Ensure the instance is of type AIContent
@@ -121,6 +134,8 @@ def test_data_content_uri():
    # Check the type and content
    assert content.type == "data"
    assert content.uri == "data:application/octet-stream;base64,dGVzdA=="
+    # media_type attribute is None when created from uri-only
+    assert content.has_top_level_media_type("application") is False
    assert content.additional_properties["version"] == 1

    # Ensure the instance is of type AIContent
@@ -153,7 +168,7 @@ def test_data_content_empty():
        DataContent(uri="")


-# region: UriContent
+# region UriContent


 def test_uri_content():
@@ -164,13 +179,15 @@ def test_uri_content():
    assert content.type == "uri"
    assert content.uri == "http://example.com"
    assert content.media_type == "image/jpg"
+    assert content.has_top_level_media_type("image") is True
+    assert content.has_top_level_media_type("application") is False
    assert content.additional_properties["version"] == 1

    # Ensure the instance is of type AIContent
    assert isinstance(content, AIContent)


-# region: FunctionCallContent
+# region FunctionCallContent


 def test_function_call_content():
@@ -186,7 +203,44 @@ def test_function_call_content():
    assert isinstance(content, AIContent)


-# region: FunctionResultContent
+def test_function_call_content_parse_arguments():
+    c1 = FunctionCallContent(call_id="1", name="f", arguments='{"a": 1, "b": 2}')
+    assert c1.parse_arguments() == {"a": 1, "b": 2}
+    c2 = FunctionCallContent(call_id="1", name="f", arguments="not json")
+    assert c2.parse_arguments() == {"raw": "not json"}
+    c3 = FunctionCallContent(call_id="1", name="f", arguments={"x": None})
+    assert c3.parse_arguments() == {"x": None}
+
+
+def test_function_call_content_add_merging_and_errors():
+    # str + str concatenation
+    a = FunctionCallContent(call_id="1", name="f", arguments="abc")
+    b = FunctionCallContent(call_id="1", name="f", arguments="def")
+    c = a + b
+    assert isinstance(c.arguments, str) and c.arguments == "abcdef"
+
+    # dict + dict merge
+    a = FunctionCallContent(call_id="1", name="f", arguments={"x": 1})
+    b = FunctionCallContent(call_id="1", name="f", arguments={"y": 2})
+    c = a + b
+    assert c.arguments == {"x": 1, "y": 2}
+
+    # incompatible argument types
+    a = FunctionCallContent(call_id="1", name="f", arguments="abc")
+    b = FunctionCallContent(call_id="1", name="f", arguments={"y": 2})
+    with raises(TypeError):
+        _ = a + b
+
+    # incompatible call ids
+    a = FunctionCallContent(call_id="1", name="f", arguments="abc")
+    b = FunctionCallContent(call_id="2", name="f", arguments="def")
+    from agent_framework.exceptions import AgentFrameworkException
+
+    with raises(AgentFrameworkException):
+        _ = a + b
+
+
+# region FunctionResultContent


 def test_function_result_content():
@@ -201,7 +255,7 @@ def test_function_result_content():
    assert isinstance(content, AIContent)


-# region: UsageDetails
+# region UsageDetails


 def test_usage_details():
@@ -247,7 +301,23 @@ def test_usage_details_additional_counts():
    assert usage.additional_counts["test"] == 1


-# region: AIContent Serialization
+def test_usage_details_add_with_none_and_type_errors():
+    u = UsageDetails(input_token_count=1)
+    # __add__ with None returns self (no change)
+    v = u + None
+    assert v is u
+    # __iadd__ with None leaves unchanged
+    u2 = UsageDetails(input_token_count=2)
+    u2 += None
+    assert u2.input_token_count == 2
+    # wrong type raises
+    with raises(ValueError):
+        _ = u + 42  # type: ignore[arg-type]
+    with raises(ValueError):
+        u += 42  # type: ignore[arg-type]
+
+
+# region AIContent Serialization


@mark.parametrize(
@@ -274,7 +344,7 @@ def test_ai_content_serialization(content_type: type[AIContent], args: dict):
    assert isinstance(test_item.content, content_type)


-# region: ChatMessage
+# region ChatMessage


 def test_chat_message_text():
@@ -310,7 +380,13 @@ def test_chat_message_contents():
    assert message.text == "Hello, how are you? I'm fine, thank you!"


-# region: ChatResponse
+def test_chat_message_with_chatrole_instance():
+    m = ChatMessage(role=ChatRole.USER, text="hi")
+    assert m.role == ChatRole.USER
+    assert m.text == "hi"
+
+
+# region ChatResponse


 def test_chat_response():
@@ -325,9 +401,11 @@ def test_chat_response():
    assert response.messages[0].role == ChatRole.ASSISTANT
    assert response.messages[0].text == "I'm doing well, thank you!"
    assert isinstance(response.messages[0], ChatMessage)
+    # __str__ returns text
+    assert str(response) == response.text


-# region: StructuredResponse
+# region StructuredResponse


 def test_structured_response():
@@ -346,9 +424,11 @@ def test_structured_response():
    # Check the type and content
    assert response.value == ResponseModel(content="Hello, world!", action="test")
    assert isinstance(response, StructuredResponse)
+    # text property returns joined messages text (single message present)
+    assert isinstance(response.text, str)


-# region: ChatResponseUpdate
+# region ChatResponseUpdate


 def test_chat_response_update():
@@ -362,6 +442,15 @@ def test_chat_response_update():
    # Check the type and content
    assert response_update.contents[0].text == "I'm doing well, thank you!"
    assert isinstance(response_update.contents[0], TextContent)
+    assert response_update.text == "I'm doing well, thank you!"
+
+
+def test_chat_response_update_with_method():
+    u = ChatResponseUpdate(text="Hello", message_id="1")
+    v = u.with_(contents=[TextContent(" world")])
+    assert v is not u
+    assert v.text == "Hello world"
+    assert v.message_id == "1"


 def test_chat_response_updates_to_chat_response_one():
@@ -438,7 +527,7 @@ def test_chat_response_updates_to_chat_response_multiple():
 def test_chat_response_updates_to_chat_response_multiple_multiple():
    """Test converting ChatResponseUpdate to ChatResponse."""
    # Create a ChatMessage
-    message1 = TextContent("I'm doing well, ")
+    message1 = TextContent("I'm doing well, ", raw_representation="I'm doing well, ")
    message2 = TextContent("thank you!")

    # Create a ChatResponseUpdate with the message
@@ -457,6 +546,7 @@ def test_chat_response_updates_to_chat_response_multiple_multiple():
    assert len(chat_response.messages) == 1
    assert isinstance(chat_response.messages[0], ChatMessage)
    assert chat_response.messages[0].message_id == "1"
+    assert chat_response.messages[0].contents[0].raw_representation is not None

    assert len(chat_response.messages[0].contents) == 3
    assert isinstance(chat_response.messages[0].contents[0], TextContent)
@@ -469,7 +559,17 @@ def test_chat_response_updates_to_chat_response_multiple_multiple():
    assert chat_response.text == "I'm doing well, thank you! More contextFinal part"


-# region: ChatToolMode
+@mark.asyncio
+async def test_chat_response_from_async_generator():
+    async def gen() -> AsyncIterable[ChatResponseUpdate]:
+        yield ChatResponseUpdate(text="Hello", message_id="1")
+        yield ChatResponseUpdate(text=" world", message_id="1")
+
+    resp = await ChatResponse.from_chat_response_generator(gen())
+    assert resp.text == "Hello world"
+
+
+# region ChatToolMode


 def test_chat_tool_mode():
@@ -497,6 +597,8 @@ def test_chat_tool_mode():
    assert isinstance(none_mode, ChatToolMode)

    assert ChatToolMode.REQUIRED("example_function") == ChatToolMode.REQUIRED("example_function")
+    # serializer returns just the mode
+    assert ChatToolMode.REQUIRED_ANY.model_dump() == "required"


 def test_chat_tool_mode_from_dict():
@@ -525,7 +627,7 @@ def test_generated_embeddings():
    assert issubclass(GeneratedEmbeddings, MutableSequence)


-# region: ChatOptions
+# region ChatOptions


 def test_chat_options_init() -> None:
@@ -543,6 +645,10 @@ def test_chat_options_init_with_args(ai_function_tool, ai_tool) -> None:
        frequency_penalty=0.0,
        user="user-123",
        tools=[ai_function_tool, ai_tool],
+        tool_choice="required",
+        additional_properties={"custom": True},
+        logit_bias={"a": 1},
+        metadata={"m": "v"},
    )
    assert options.ai_model_id == "gpt-4"
    assert options.max_tokens == 1024
@@ -557,10 +663,27 @@ def test_chat_options_init_with_args(ai_function_tool, ai_tool) -> None:
        assert tool.description is not None
        assert tool.parameters() is not None

+    settings = options.to_provider_settings()
+    assert settings["model"] == "gpt-4"  # uses alias
+    assert settings["tool_choice"] == "required"  # serialized via model_serializer
+    assert settings["custom"] is True  # from additional_properties
+    assert "additional_properties" not in settings
+
+
+def test_chat_options_tool_choice_validation_errors():
+    with raises((ValidationError, TypeError)):
+        ChatOptions(tool_choice="invalid-choice")
+
+
+def test_chat_options_tool_choice_excluded_when_no_tools():
+    options = ChatOptions(tool_choice="auto")
+    settings = options.to_provider_settings()
+    assert "tool_choice" not in settings
+

 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])
+    options1 = ChatOptions(ai_model_id="gpt-4o", tools=[ai_function_tool], logit_bias={"x": 1}, metadata={"a": "b"})
+    options2 = ChatOptions(ai_model_id="gpt-4.1", tools=[ai_tool], additional_properties={"p": 1})
    assert options1 != options2
    options3 = options1 & options2

@@ -568,6 +691,9 @@ def test_chat_options_and(ai_function_tool, ai_tool) -> None:
    assert len(options3._ai_tools) == 2
    assert options3._ai_tools == [ai_function_tool, ai_tool]
    assert options3.tools == [ai_function_tool, ai_tool]
+    assert options3.logit_bias == {"x": 1}
+    assert options3.metadata == {"a": "b"}
+    assert options3.additional_properties.get("p") == 1


 # region Agent Response Fixtures
@@ -661,3 +787,228 @@ def test_agent_run_response_update_text_property_empty() -> None:
 def test_agent_run_response_update_str_method(text_content: TextContent) -> None:
    update = AgentRunResponseUpdate(contents=[text_content])
    assert str(update) == "Test content"
+
+
+# region ErrorContent
+
+
+def test_error_content_str():
+    e1 = ErrorContent(message="Oops", error_code="E1")
+    assert str(e1) == "Error E1: Oops"
+    e2 = ErrorContent(message="Oops")
+    assert str(e2) == "Oops"
+    e3 = ErrorContent()
+    assert str(e3) == "Unknown error"
+
+
+# region Annotations
+
+
+def test_annotations_models_and_roundtrip():
+    span = TextSpanRegion(start_index=0, end_index=5)
+    base_region = AnnotatedRegion()
+    ann: AIAnnotation = AIAnnotation(annotated_regions=[span, base_region])
+    cit = CitationAnnotation(title="Doc", url="http://example.com", snippet="Snippet", annotated_regions=[span])
+
+    # Attach to content
+    content = TextContent(text="hello", additional_properties={"v": 1})
+    content.annotations = [ann, cit]
+
+    dumped = content.model_dump()
+    loaded = TextContent.model_validate(dumped)
+    assert isinstance(loaded.annotations, list)
+    assert len(loaded.annotations) == 2
+    assert isinstance(loaded.annotations[0], dict) is False  # pydantic parsed into models
+    # discriminators preserved
+    assert any(getattr(a, "type", None) == "citation" for a in loaded.annotations)
+
+
+def test_function_call_merge_in_process_update_and_usage_aggregation():
+    # Two function call chunks with same call_id should merge
+    u1 = ChatResponseUpdate(contents=[FunctionCallContent(call_id="c1", name="f", arguments="{")], message_id="m")
+    u2 = ChatResponseUpdate(contents=[FunctionCallContent(call_id="c1", name="f", arguments="}")], message_id="m")
+    # plus usage
+    u3 = ChatResponseUpdate(contents=[UsageContent(UsageDetails(input_token_count=1, output_token_count=2))])
+
+    resp = ChatResponse.from_chat_response_updates([u1, u2, u3])
+    assert len(resp.messages) == 1
+    last_contents = resp.messages[0].contents
+    assert any(isinstance(c, FunctionCallContent) for c in last_contents)
+    fcs = [c for c in last_contents if isinstance(c, FunctionCallContent)]
+    assert len(fcs) == 1
+    assert fcs[0].arguments == "{}"
+    assert resp.usage_details is not None
+    assert resp.usage_details.input_token_count == 1
+    assert resp.usage_details.output_token_count == 2
+
+
+def test_function_call_incompatible_ids_are_not_merged():
+    u1 = ChatResponseUpdate(contents=[FunctionCallContent(call_id="a", name="f", arguments="x")], message_id="m")
+    u2 = ChatResponseUpdate(contents=[FunctionCallContent(call_id="b", name="f", arguments="y")], message_id="m")
+
+    resp = ChatResponse.from_chat_response_updates([u1, u2])
+    fcs = [c for c in resp.messages[0].contents if isinstance(c, FunctionCallContent)]
+    assert len(fcs) == 2
+
+
+# region Speech/Text To Speech options
+
+
+def test_speech_to_text_options_provider_settings():
+    o = SpeechToTextOptions(ai_model_id="stt", additional_properties={"x": 1})
+    settings = o.to_provider_settings()
+    assert settings["model"] == "stt"
+    assert settings["x"] == 1
+    assert "additional_properties" not in settings
+
+
+def test_text_to_speech_options_provider_settings():
+    o = TextToSpeechOptions(ai_model_id="tts", response_format="wav", speed=1.2, additional_properties={"x": 2})
+    settings = o.to_provider_settings()
+    assert settings["model"] == "tts"
+    assert settings["response_format"] == "wav"
+    assert settings["x"] == 2
+
+
+# region GeneratedEmbeddings operations
+
+
+def test_generated_embeddings_operations():
+    g = GeneratedEmbeddings[int](embeddings=[1, 2, 3])
+    assert 2 in g
+    assert list(iter(g)) == [1, 2, 3]
+    assert len(g) == 3
+    assert list(reversed(g)) == [3, 2, 1]
+    assert g.index(2) == 1
+    assert g.count(2) == 1
+    assert g[0] == 1
+    assert g[0:2] == [1, 2]
+
+    g[1] = 5
+    assert g[1] == 5
+    g[1:3] = [7, 8]
+    assert g[1:] == [7, 8]
+
+    with raises(TypeError):
+        g[0] = [9]  # int index cannot be set with iterable
+    with raises(TypeError):
+        g[0:1] = 9  # slice requires iterable
+
+    del g[0]
+    assert g.embeddings == [7, 8]
+    del g[0:1]
+    assert g.embeddings == [8]
+
+    g.insert(0, 1)
+    g.append(2)
+    g.extend([3, 4])
+    assert g.embeddings == [1, 8, 2, 3, 4]
+    g.reverse()
+    assert g.embeddings == [4, 3, 2, 8, 1]
+    assert g.pop() == 1
+    g.remove(8)
+    assert g.embeddings == [4, 3, 2]
+
+    # iadd with another GeneratedEmbeddings, including usage merge
+    g2 = GeneratedEmbeddings[int](embeddings=[5], usage=UsageDetails(input_token_count=1))
+    g.usage = UsageDetails(input_token_count=2)
+    g += g2
+    assert g.embeddings[-1] == 5
+    assert g.usage.input_token_count == 3
+
+    # clear
+    g.additional_properties = {"a": 1}
+    g.clear()
+    assert g.embeddings == []
+    assert g.usage is None
+    assert g.additional_properties == {}
+
+
+# region ChatRole & ChatFinishReason basics
+
+
+def test_chat_role_str_and_repr():
+    assert str(ChatRole.USER) == "user"
+    assert "ChatRole(value=" in repr(ChatRole.USER)
+
+
+def test_chat_finish_reason_constants():
+    assert ChatFinishReason.STOP.value == "stop"
+
+
+def test_response_update_propagates_fields_and_metadata():
+    upd = ChatResponseUpdate(
+        text="hello",
+        role="assistant",
+        author_name="bot",
+        response_id="rid",
+        message_id="mid",
+        conversation_id="cid",
+        ai_model_id="model-x",
+        created_at="t0",
+        finish_reason=ChatFinishReason.STOP,
+        additional_properties={"k": "v"},
+    )
+    resp = ChatResponse.from_chat_response_updates([upd])
+    assert resp.response_id == "rid"
+    assert resp.created_at == "t0"
+    assert resp.conversation_id == "cid"
+    assert resp.ai_model_id == "model-x"
+    assert resp.finish_reason == ChatFinishReason.STOP
+    assert resp.additional_properties and resp.additional_properties["k"] == "v"
+    assert resp.messages[0].role == ChatRole.ASSISTANT
+    assert resp.messages[0].author_name == "bot"
+    assert resp.messages[0].message_id == "mid"
+
+
+def test_text_coalescing_preserves_first_properties():
+    t1 = TextContent("A", raw_representation={"r": 1}, additional_properties={"p": 1})
+    t2 = TextContent("B")
+    upd1 = ChatResponseUpdate(text=t1, message_id="x")
+    upd2 = ChatResponseUpdate(text=t2, message_id="x")
+    resp = ChatResponse.from_chat_response_updates([upd1, upd2])
+    # After coalescing there should be a single TextContent with merged text and preserved props from first
+    items = [c for c in resp.messages[0].contents if isinstance(c, TextContent)]
+    assert len(items) >= 1
+    assert items[0].text == "AB"
+    assert items[0].raw_representation == {"r": 1}
+    assert items[0].additional_properties == {"p": 1}
+
+
+def test_function_call_content_parse_numeric_or_list():
+    c_num = FunctionCallContent(call_id="1", name="f", arguments="123")
+    assert c_num.parse_arguments() == {"raw": 123}
+    c_list = FunctionCallContent(call_id="1", name="f", arguments="[1,2]")
+    assert c_list.parse_arguments() == {"raw": [1, 2]}
+
+
+def test_chat_tool_mode_eq_with_string():
+    assert ChatToolMode.AUTO == "auto"
+
+
+def test_chat_options_tool_choice_dict_mapping(ai_tool):
+    opts = ChatOptions(tool_choice={"mode": "required", "required_function_name": "fn"}, tools=[ai_tool])
+    assert isinstance(opts.tool_choice, ChatToolMode)
+    assert opts.tool_choice.mode == "required"
+    assert opts.tool_choice.required_function_name == "fn"
+    # provider settings serialize to just the mode
+    settings = opts.to_provider_settings()
+    assert settings["tool_choice"] == "required"
+
+
+# region AgentRunResponse
+
+
+@fixture
+def agent_run_response_async() -> AgentRunResponse:
+    return AgentRunResponse(messages=[ChatMessage(role="user", text="Hello")])
+
+
+@mark.asyncio
+async def test_agent_run_response_from_async_generator():
+    async def gen():
+        yield AgentRunResponseUpdate(contents=[TextContent("A")])
+        yield AgentRunResponseUpdate(contents=[TextContent("B")])
+
+    r = await AgentRunResponse.from_agent_response_generator(gen())
+    assert r.text == "AB"
@@ -4,7 +4,7 @@ from typing import Any
 from pytest import fixture


-# region: Connector Settings fixtures
+# region Connector Settings fixtures
@fixture
 def exclude_list(request: Any) -> list[str]:
    """Fixture that returns a list of environment variables to exclude."""
@@ -19,7 +19,7 @@ from ._events import (
 from ._typing_utils import is_instance_of
 from ._workflow_context import WorkflowContext

-# region: Executor
+# region Executor


 class Executor:
@@ -96,7 +96,7 @@ class Executor:

 # endregion: Executor

-# region: Handler Decorator
+# region Handler Decorator


 ExecutorT = TypeVar("ExecutorT", bound="Executor")
@@ -183,7 +183,7 @@ def handler(

 # endregion: Handler Decorator

-# region: Agent Executor
+# region Agent Executor


@dataclass
@@ -268,7 +268,7 @@ class AgentExecutor(Executor):
 # endregion: Agent Executor


-# region: Request Info Executor
+# region Request Info Executor


@dataclass