agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
This commit is contained in:
Eduard van Valkenburg
2025-08-08 22:27:31 +02:00
committed by GitHub
Unverified
parent 8f27e63df6
commit 82d39bc1f7
8 changed files with 934 additions and 346 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/agent_framework/_clients.py
+3 -3
View File
@@ -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
python/packages/main/agent_framework/_types.py
+266 -38
View File
@@ -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):
python/packages/main/agent_framework/exceptions.py
+1 -1
View File
@@ -19,7 +19,7 @@ class AgentExecutionException(AgentException):
pass
# region: Service Exceptions
# region Service Exceptions
class ServiceException(AgentFrameworkException):
python/packages/main/agent_framework/openai/_chat_client.py
+3 -3
View File
@@ -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)