agent/agent-framework
2
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

Capture file IDs from code interpreter in streaming responses (#2741)

Browse Source
This commit is contained in:
Evan Mattson
2025-12-11 16:30:19 +09:00
committed by GitHub
Unverified
parent 4c6a5d4aa1
commit 3481914981
10 changed files with 785 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
+47
View File
@@ -63,6 +63,8 @@ from azure.ai.agents.models import (
McpTool,
MessageDeltaChunk,
MessageDeltaTextContent,
MessageDeltaTextFileCitationAnnotation,
MessageDeltaTextFilePathAnnotation,
MessageDeltaTextUrlCitationAnnotation,
MessageImageUrlParam,
MessageInputContentBlock,
@@ -471,6 +473,45 @@ class AzureAIAgentClient(BaseChatClient):
return url_citations
def _extract_file_path_contents(self, message_delta_chunk: MessageDeltaChunk) -> list[HostedFileContent]:
"""Extract file references from MessageDeltaChunk annotations.
Code interpreter generates files that are referenced via file path or file citation
annotations in the message content. This method extracts those file IDs and returns
them as HostedFileContent objects.
Handles two annotation types:
- MessageDeltaTextFilePathAnnotation: Contains file_path.file_id
- MessageDeltaTextFileCitationAnnotation: Contains file_citation.file_id
Args:
message_delta_chunk: The message delta chunk to process
Returns:
List of HostedFileContent objects for any files referenced in annotations
"""
file_contents: list[HostedFileContent] = []
for content in message_delta_chunk.delta.content:
if isinstance(content, MessageDeltaTextContent) and content.text and content.text.annotations:
for annotation in content.text.annotations:
if isinstance(annotation, MessageDeltaTextFilePathAnnotation):
# Extract file_id from the file_path annotation
file_path = getattr(annotation, "file_path", None)
if file_path is not None:
file_id = getattr(file_path, "file_id", None)
if file_id:
file_contents.append(HostedFileContent(file_id=file_id))
elif isinstance(annotation, MessageDeltaTextFileCitationAnnotation):
# Extract file_id from the file_citation annotation
file_citation = getattr(annotation, "file_citation", None)
if file_citation is not None:
file_id = getattr(file_citation, "file_id", None)
if file_id:
file_contents.append(HostedFileContent(file_id=file_id))
return file_contents
def _get_real_url_from_citation_reference(
self, citation_url: str, azure_search_tool_calls: list[dict[str, Any]]
) -> str:
@@ -530,6 +571,9 @@ class AzureAIAgentClient(BaseChatClient):
# Extract URL citations from the delta chunk
url_citations = self._extract_url_citations(event_data, azure_search_tool_calls)
# Extract file path contents from code interpreter outputs
file_contents = self._extract_file_path_contents(event_data)
# Create contents with citations if any exist
citation_content: list[Contents] = []
if event_data.text or url_citations:
@@ -538,6 +582,9 @@ class AzureAIAgentClient(BaseChatClient):
text_content_obj.annotations = url_citations
citation_content.append(text_content_obj)
# Add file contents from file path annotations
citation_content.extend(file_contents)
yield ChatResponseUpdate(
role=role,
contents=citation_content if citation_content else None,
python/packages/azure-ai/tests/test_azure_ai_agent_client.py
+105
View File
@@ -24,6 +24,7 @@ from agent_framework import (
FunctionCallContent,
FunctionResultContent,
HostedCodeInterpreterTool,
HostedFileContent,
HostedFileSearchTool,
HostedMCPTool,
HostedVectorStoreContent,
@@ -42,6 +43,8 @@ from azure.ai.agents.models import (
FileInfo,
MessageDeltaChunk,
MessageDeltaTextContent,
MessageDeltaTextFileCitationAnnotation,
MessageDeltaTextFilePathAnnotation,
MessageDeltaTextUrlCitationAnnotation,
RequiredFunctionToolCall,
RequiredMcpToolCall,
@@ -1362,6 +1365,108 @@ def test_azure_ai_chat_client_extract_url_citations_with_citations(mock_agents_c
assert citation.annotated_regions[0].end_index == 20
def test_azure_ai_chat_client_extract_file_path_contents_with_file_path_annotation(
mock_agents_client: MagicMock,
) -> None:
"""Test _extract_file_path_contents with MessageDeltaChunk containing file path annotation."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
# Create mock file_path annotation
mock_file_path = MagicMock()
mock_file_path.file_id = "assistant-test-file-123"
mock_annotation = MagicMock(spec=MessageDeltaTextFilePathAnnotation)
mock_annotation.file_path = mock_file_path
# Create mock text content with annotations
mock_text = MagicMock()
mock_text.annotations = [mock_annotation]
mock_text_content = MagicMock(spec=MessageDeltaTextContent)
mock_text_content.text = mock_text
# Create mock delta
mock_delta = MagicMock()
mock_delta.content = [mock_text_content]
# Create mock MessageDeltaChunk
mock_chunk = MagicMock(spec=MessageDeltaChunk)
mock_chunk.delta = mock_delta
# Call the method
file_contents = chat_client._extract_file_path_contents(mock_chunk)
# Verify results
assert len(file_contents) == 1
assert isinstance(file_contents[0], HostedFileContent)
assert file_contents[0].file_id == "assistant-test-file-123"
def test_azure_ai_chat_client_extract_file_path_contents_with_file_citation_annotation(
mock_agents_client: MagicMock,
) -> None:
"""Test _extract_file_path_contents with MessageDeltaChunk containing file citation annotation."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
# Create mock file_citation annotation
mock_file_citation = MagicMock()
mock_file_citation.file_id = "cfile_test-citation-456"
mock_annotation = MagicMock(spec=MessageDeltaTextFileCitationAnnotation)
mock_annotation.file_citation = mock_file_citation
# Create mock text content with annotations
mock_text = MagicMock()
mock_text.annotations = [mock_annotation]
mock_text_content = MagicMock(spec=MessageDeltaTextContent)
mock_text_content.text = mock_text
# Create mock delta
mock_delta = MagicMock()
mock_delta.content = [mock_text_content]
# Create mock MessageDeltaChunk
mock_chunk = MagicMock(spec=MessageDeltaChunk)
mock_chunk.delta = mock_delta
# Call the method
file_contents = chat_client._extract_file_path_contents(mock_chunk)
# Verify results
assert len(file_contents) == 1
assert isinstance(file_contents[0], HostedFileContent)
assert file_contents[0].file_id == "cfile_test-citation-456"
def test_azure_ai_chat_client_extract_file_path_contents_empty_annotations(
mock_agents_client: MagicMock,
) -> None:
"""Test _extract_file_path_contents with no annotations returns empty list."""
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
# Create mock text content with no annotations
mock_text = MagicMock()
mock_text.annotations = []
mock_text_content = MagicMock(spec=MessageDeltaTextContent)
mock_text_content.text = mock_text
# Create mock delta
mock_delta = MagicMock()
mock_delta.content = [mock_text_content]
# Create mock MessageDeltaChunk
mock_chunk = MagicMock(spec=MessageDeltaChunk)
mock_chunk.delta = mock_delta
# Call the method
file_contents = chat_client._extract_file_path_contents(mock_chunk)
# Verify results
assert len(file_contents) == 0
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
python/packages/core/agent_framework/openai/_responses_client.py
+71 -13
View File
@@ -3,7 +3,7 @@
from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, MutableSequence, Sequence
from datetime import datetime, timezone
from itertools import chain
from typing import Any, TypeVar
from typing import Any, TypeVar, cast
from openai import AsyncOpenAI, BadRequestError
from openai.types.responses.file_search_tool_param import FileSearchToolParam
@@ -199,7 +199,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
return response_format, prepared_text
if isinstance(response_format, Mapping):
format_config = self._convert_response_format(response_format)
format_config = self._convert_response_format(cast("Mapping[str, Any]", response_format))
if prepared_text is None:
prepared_text = {}
elif "format" in prepared_text and prepared_text["format"] != format_config:
@@ -212,20 +212,21 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
def _convert_response_format(self, response_format: Mapping[str, Any]) -> dict[str, Any]:
"""Convert Chat style response_format into Responses text format config."""
if "format" in response_format and isinstance(response_format["format"], Mapping):
return dict(response_format["format"])
return dict(cast("Mapping[str, Any]", response_format["format"]))
format_type = response_format.get("type")
if format_type == "json_schema":
schema_section = response_format.get("json_schema", response_format)
if not isinstance(schema_section, Mapping):
raise ServiceInvalidRequestError("json_schema response_format must be a mapping.")
schema = schema_section.get("schema")
schema_section_typed = cast("Mapping[str, Any]", schema_section)
schema: Any = schema_section_typed.get("schema")
if schema is None:
raise ServiceInvalidRequestError("json_schema response_format requires a schema.")
name = (
schema_section.get("name")
or schema_section.get("title")
or (schema.get("title") if isinstance(schema, Mapping) else None)
name: str = str(
schema_section_typed.get("name")
or schema_section_typed.get("title")
or (cast("Mapping[str, Any]", schema).get("title") if isinstance(schema, Mapping) else None)
or "response"
)
format_config: dict[str, Any] = {
@@ -532,12 +533,13 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
"text": content.text,
},
}
if content.additional_properties is not None:
if status := content.additional_properties.get("status"):
props: dict[str, Any] | None = getattr(content, "additional_properties", None)
if props:
if status := props.get("status"):
ret["status"] = status
if reasoning_text := content.additional_properties.get("reasoning_text"):
if reasoning_text := props.get("reasoning_text"):
ret["content"] = {"type": "reasoning_text", "text": reasoning_text}
if encrypted_content := content.additional_properties.get("encrypted_content"):
if encrypted_content := props.get("encrypted_content"):
ret["encrypted_content"] = encrypted_content
return ret
case DataContent() | UriContent():
@@ -824,7 +826,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
"raw_representation": response,
}
conversation_id = self.get_conversation_id(response, chat_options.store)
conversation_id = self.get_conversation_id(response, chat_options.store) # type: ignore[reportArgumentType]
if conversation_id:
args["conversation_id"] = conversation_id
@@ -911,6 +913,8 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
metadata.update(self._get_metadata_from_response(event_part))
case "refusal":
contents.append(TextContent(text=event_part.refusal, raw_representation=event))
case _:
pass
case "response.output_text.delta":
contents.append(TextContent(text=event.delta, raw_representation=event))
metadata.update(self._get_metadata_from_response(event))
@@ -1032,6 +1036,60 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
raw_representation=event,
)
)
case "response.output_text.annotation.added":
# Handle streaming text annotations (file citations, file paths, etc.)
annotation: Any = event.annotation
def _get_ann_value(key: str) -> Any:
"""Extract value from annotation (dict or object)."""
if isinstance(annotation, dict):
return cast("dict[str, Any]", annotation).get(key)
return getattr(annotation, key, None)
ann_type = _get_ann_value("type")
ann_file_id = _get_ann_value("file_id")
if ann_type == "file_path":
if ann_file_id:
contents.append(
HostedFileContent(
file_id=str(ann_file_id),
additional_properties={
"annotation_index": event.annotation_index,
"index": _get_ann_value("index"),
},
raw_representation=event,
)
)
elif ann_type == "file_citation":
if ann_file_id:
contents.append(
HostedFileContent(
file_id=str(ann_file_id),
additional_properties={
"annotation_index": event.annotation_index,
"filename": _get_ann_value("filename"),
"index": _get_ann_value("index"),
},
raw_representation=event,
)
)
elif ann_type == "container_file_citation":
if ann_file_id:
contents.append(
HostedFileContent(
file_id=str(ann_file_id),
additional_properties={
"annotation_index": event.annotation_index,
"container_id": _get_ann_value("container_id"),
"filename": _get_ann_value("filename"),
"start_index": _get_ann_value("start_index"),
"end_index": _get_ann_value("end_index"),
},
raw_representation=event,
)
)
else:
logger.debug("Unparsed annotation type in streaming: %s", ann_type)
case _:
logger.debug("Unparsed event of type: %s: %s", event.type, event)
python/packages/core/tests/openai/test_openai_responses_client.py
+104
View File
@@ -993,6 +993,110 @@ def test_streaming_response_basic_structure() -> None:
assert response.raw_representation is mock_event
def test_streaming_annotation_added_with_file_path() -> None:
"""Test streaming annotation added event with file_path type extracts HostedFileContent."""
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
chat_options = ChatOptions()
function_call_ids: dict[int, tuple[str, str]] = {}
mock_event = MagicMock()
mock_event.type = "response.output_text.annotation.added"
mock_event.annotation_index = 0
mock_event.annotation = {
"type": "file_path",
"file_id": "file-abc123",
"index": 42,
}
response = client._create_streaming_response_content(mock_event, chat_options, function_call_ids)
assert len(response.contents) == 1
content = response.contents[0]
assert isinstance(content, HostedFileContent)
assert content.file_id == "file-abc123"
assert content.additional_properties is not None
assert content.additional_properties.get("annotation_index") == 0
assert content.additional_properties.get("index") == 42
def test_streaming_annotation_added_with_file_citation() -> None:
"""Test streaming annotation added event with file_citation type extracts HostedFileContent."""
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
chat_options = ChatOptions()
function_call_ids: dict[int, tuple[str, str]] = {}
mock_event = MagicMock()
mock_event.type = "response.output_text.annotation.added"
mock_event.annotation_index = 1
mock_event.annotation = {
"type": "file_citation",
"file_id": "file-xyz789",
"filename": "sample.txt",
"index": 15,
}
response = client._create_streaming_response_content(mock_event, chat_options, function_call_ids)
assert len(response.contents) == 1
content = response.contents[0]
assert isinstance(content, HostedFileContent)
assert content.file_id == "file-xyz789"
assert content.additional_properties is not None
assert content.additional_properties.get("filename") == "sample.txt"
assert content.additional_properties.get("index") == 15
def test_streaming_annotation_added_with_container_file_citation() -> None:
"""Test streaming annotation added event with container_file_citation type."""
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
chat_options = ChatOptions()
function_call_ids: dict[int, tuple[str, str]] = {}
mock_event = MagicMock()
mock_event.type = "response.output_text.annotation.added"
mock_event.annotation_index = 2
mock_event.annotation = {
"type": "container_file_citation",
"file_id": "file-container123",
"container_id": "container-456",
"filename": "data.csv",
"start_index": 10,
"end_index": 50,
}
response = client._create_streaming_response_content(mock_event, chat_options, function_call_ids)
assert len(response.contents) == 1
content = response.contents[0]
assert isinstance(content, HostedFileContent)
assert content.file_id == "file-container123"
assert content.additional_properties is not None
assert content.additional_properties.get("container_id") == "container-456"
assert content.additional_properties.get("filename") == "data.csv"
assert content.additional_properties.get("start_index") == 10
assert content.additional_properties.get("end_index") == 50
def test_streaming_annotation_added_with_unknown_type() -> None:
"""Test streaming annotation added event with unknown type is ignored."""
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
chat_options = ChatOptions()
function_call_ids: dict[int, tuple[str, str]] = {}
mock_event = MagicMock()
mock_event.type = "response.output_text.annotation.added"
mock_event.annotation_index = 0
mock_event.annotation = {
"type": "url_citation",
"url": "https://example.com",
}
response = client._create_streaming_response_content(mock_event, chat_options, function_call_ids)
# url_citation should not produce HostedFileContent
assert len(response.contents) == 0
def test_service_response_exception_includes_original_error_details() -> None:
"""Test that ServiceResponseException messages include original error details in the new format."""
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")