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

Author: moonbox3

python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py

@@ -63,6 +63,8 @@
     McpTool,
     MessageDeltaChunk,
     MessageDeltaTextContent,
+    MessageDeltaTextFileCitationAnnotation,
+    MessageDeltaTextFilePathAnnotation,
     MessageDeltaTextUrlCitationAnnotation,
     MessageImageUrlParam,
     MessageInputContentBlock,
@@ -471,6 +473,45 @@ def _extract_url_citations(
 
         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 @@ async def _process_stream(
                         # 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 @@ async def _process_stream(
                                 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

@@ -24,6 +24,7 @@
     FunctionCallContent,
     FunctionResultContent,
     HostedCodeInterpreterTool,
+    HostedFileContent,
     HostedFileSearchTool,
     HostedMCPTool,
     HostedVectorStoreContent,
@@ -42,6 +43,8 @@
     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

@@ -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 @@ def _prepare_text_config(
             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 @@ def _prepare_text_config(
     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 @@ def _openai_content_parser(
                         "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 @@ def _create_response_content(
             "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 @@ def _create_streaming_response_content(
                         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 @@ def _create_streaming_response_content(
                         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)