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

2026-06-16 21:04:09 +08:00 · 2025-12-11 16:30:19 +09:00
parent 4c6a5d4aa1
commit 3481914981
10 changed files with 785 additions and 13 deletions
@@ -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,
@@ -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:
@@ -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)

@@ -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")
@@ -25,6 +25,7 @@ This directory contains samples demonstrating the capabilities of Microsoft Agen
 | [`getting_started/agents/azure_ai_agent/azure_ai_with_azure_ai_search.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_azure_ai_search.py) | Azure AI Agent with Azure AI Search Example |
 | [`getting_started/agents/azure_ai_agent/azure_ai_with_bing_grounding.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_bing_grounding.py) | Azure AI agent with Bing Grounding search for real-time web information |
 | [`getting_started/agents/azure_ai_agent/azure_ai_with_code_interpreter.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_code_interpreter.py) | Azure AI Agent with Code Interpreter Example |
+| [`getting_started/agents/azure_ai_agent/azure_ai_with_code_interpreter_file_generation.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_code_interpreter_file_generation.py) | Azure AI Agent with Code Interpreter File Generation Example |
 | [`getting_started/agents/azure_ai_agent/azure_ai_with_existing_agent.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_existing_agent.py) | Azure AI Agent with Existing Agent Example |
 | [`getting_started/agents/azure_ai_agent/azure_ai_with_existing_thread.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_existing_thread.py) | Azure AI Agent with Existing Thread Example |
 | [`getting_started/agents/azure_ai_agent/azure_ai_with_explicit_settings.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_explicit_settings.py) | Azure AI Agent with Explicit Settings Example |
@@ -47,6 +48,7 @@ This directory contains samples demonstrating the capabilities of Microsoft Agen
 | [`getting_started/agents/azure_ai/azure_ai_with_bing_custom_search.py`](./getting_started/agents/azure_ai/azure_ai_with_bing_custom_search.py) | Azure AI Agent with Bing Custom Search Example |
 | [`getting_started/agents/azure_ai/azure_ai_with_browser_automation.py`](./getting_started/agents/azure_ai/azure_ai_with_browser_automation.py) | Azure AI Agent with Browser Automation Example |
 | [`getting_started/agents/azure_ai/azure_ai_with_code_interpreter.py`](./getting_started/agents/azure_ai/azure_ai_with_code_interpreter.py) | Azure AI Agent with Code Interpreter Example |
+| [`getting_started/agents/azure_ai/azure_ai_with_code_interpreter_file_generation.py`](./getting_started/agents/azure_ai/azure_ai_with_code_interpreter_file_generation.py) | Azure AI Agent with Code Interpreter File Generation Example |
 | [`getting_started/agents/azure_ai/azure_ai_with_existing_agent.py`](./getting_started/agents/azure_ai/azure_ai_with_existing_agent.py) | Azure AI Agent with Existing Agent Example |
 | [`getting_started/agents/azure_ai/azure_ai_with_existing_conversation.py`](./getting_started/agents/azure_ai/azure_ai_with_existing_conversation.py) | Azure AI Agent with Existing Conversation Example |
 | [`getting_started/agents/azure_ai/azure_ai_with_explicit_settings.py`](./getting_started/agents/azure_ai/azure_ai_with_explicit_settings.py) | Azure AI Agent with Explicit Settings Example |
@@ -14,6 +14,7 @@ This folder contains examples demonstrating different ways to create and use age
 | [`azure_ai_with_bing_custom_search.py`](azure_ai_with_bing_custom_search.py) | Shows how to use Bing Custom Search with Azure AI agents to search custom search instances and provide responses with relevant results. Requires a Bing Custom Search connection and instance configured in your Azure AI project. |
 | [`azure_ai_with_browser_automation.py`](azure_ai_with_browser_automation.py) | Shows how to use Browser Automation with Azure AI agents to perform automated web browsing tasks and provide responses based on web interactions. Requires a Browser Automation connection configured in your Azure AI project. |
 | [`azure_ai_with_code_interpreter.py`](azure_ai_with_code_interpreter.py) | Shows how to use the `HostedCodeInterpreterTool` with Azure AI agents to write and execute Python code for mathematical problem solving and data analysis. |
+| [`azure_ai_with_code_interpreter_file_generation.py`](azure_ai_with_code_interpreter_file_generation.py) | Shows how to retrieve file IDs from code interpreter generated files using both streaming and non-streaming approaches. |
 | [`azure_ai_with_existing_agent.py`](azure_ai_with_existing_agent.py) | Shows how to work with a pre-existing agent by providing the agent name and version to the Azure AI client. Demonstrates agent reuse patterns for production scenarios. |
 | [`azure_ai_with_existing_conversation.py`](azure_ai_with_existing_conversation.py) | Demonstrates how to use an existing conversation created on the service side with Azure AI agents. Shows two approaches: specifying conversation ID at the client level and using AgentThread with an existing conversation ID. |
 | [`azure_ai_with_application_endpoint.py`](azure_ai_with_application_endpoint.py) | Demonstrates calling the Azure AI application-scoped endpoint. |
@@ -0,0 +1,111 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+
+from agent_framework import (
+    CitationAnnotation,
+    HostedCodeInterpreterTool,
+    HostedFileContent,
+    TextContent,
+)
+from agent_framework._agents import AgentRunResponseUpdate
+from agent_framework.azure import AzureAIClient
+from azure.identity.aio import AzureCliCredential
+
+"""
+Azure AI V2 Code Interpreter File Generation Sample
+
+This sample demonstrates how the V2 AzureAIClient handles file annotations
+when code interpreter generates text files. It shows both non-streaming
+and streaming approaches to verify file ID extraction.
+"""
+
+QUERY = (
+    "Write a simple Python script that creates a text file called 'sample.txt' containing "
+    "'Hello from the code interpreter!' and save it to disk."
+)
+
+
+async def test_non_streaming() -> None:
+    """Test non-streaming response - should have annotations on TextContent."""
+    print("=== Testing Non-Streaming Response ===")
+
+    async with (
+        AzureCliCredential() as credential,
+        AzureAIClient(credential=credential).create_agent(
+            name="V2CodeInterpreterFileAgent",
+            instructions="You are a helpful assistant that can write and execute Python code to create files.",
+            tools=HostedCodeInterpreterTool(),
+        ) as agent,
+    ):
+        print(f"User: {QUERY}\n")
+
+        result = await agent.run(QUERY)
+        print(f"Agent: {result.text}\n")
+
+        # Check for annotations in the response
+        annotations_found: list[str] = []
+        # AgentRunResponse has messages property, which contains ChatMessage objects
+        for message in result.messages:
+            for content in message.contents:
+                if isinstance(content, TextContent) and content.annotations:
+                    for annotation in content.annotations:
+                        if isinstance(annotation, CitationAnnotation) and annotation.file_id:
+                            annotations_found.append(annotation.file_id)
+                            print(f"Found file annotation: file_id={annotation.file_id}")
+
+        if annotations_found:
+            print(f"SUCCESS: Found {len(annotations_found)} file annotation(s)")
+        else:
+            print("WARNING: No file annotations found in non-streaming response")
+
+
+async def test_streaming() -> None:
+    """Test streaming response - check if file content is captured via HostedFileContent."""
+    print("\n=== Testing Streaming Response ===")
+
+    async with (
+        AzureCliCredential() as credential,
+        AzureAIClient(credential=credential).create_agent(
+            name="V2CodeInterpreterFileAgentStreaming",
+            instructions="You are a helpful assistant that can write and execute Python code to create files.",
+            tools=HostedCodeInterpreterTool(),
+        ) as agent,
+    ):
+        print(f"User: {QUERY}\n")
+        annotations_found: list[str] = []
+        text_chunks: list[str] = []
+        file_ids_found: list[str] = []
+
+        async for update in agent.run_stream(QUERY):
+            if isinstance(update, AgentRunResponseUpdate):
+                for content in update.contents:
+                    if isinstance(content, TextContent):
+                        if content.text:
+                            text_chunks.append(content.text)
+                        if content.annotations:
+                            for annotation in content.annotations:
+                                if isinstance(annotation, CitationAnnotation) and annotation.file_id:
+                                    annotations_found.append(annotation.file_id)
+                                    print(f"Found streaming annotation: file_id={annotation.file_id}")
+                    elif isinstance(content, HostedFileContent):
+                        file_ids_found.append(content.file_id)
+                        print(f"Found streaming HostedFileContent: file_id={content.file_id}")
+
+        print(f"\nAgent response: {''.join(text_chunks)[:200]}...")
+
+        if annotations_found or file_ids_found:
+            total = len(annotations_found) + len(file_ids_found)
+            print(f"SUCCESS: Found {total} file reference(s) in streaming")
+        else:
+            print("WARNING: No file annotations found in streaming response")
+
+
+async def main() -> None:
+    print("AzureAIClient Code Interpreter File Generation Test\n")
+    await test_non_streaming()
+    await test_streaming()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
@@ -9,6 +9,7 @@ This folder contains examples demonstrating different ways to create and use age
 | [`azure_ai_basic.py`](azure_ai_basic.py) | The simplest way to create an agent using `ChatAgent` with `AzureAIAgentClient`. It automatically handles all configuration using environment variables. |
 | [`azure_ai_with_bing_custom_search.py`](azure_ai_with_bing_custom_search.py) | Shows how to use Bing Custom Search with Azure AI agents to find real-time information from the web using custom search configurations. Demonstrates how to set up and use HostedWebSearchTool with custom search instances. |
 | [`azure_ai_with_bing_grounding.py`](azure_ai_with_bing_grounding.py) | Shows how to use Bing Grounding search with Azure AI agents to find real-time information from the web. Demonstrates web search capabilities with proper source citations and comprehensive error handling. |
+| [`azure_ai_with_code_interpreter_file_generation.py`](azure_ai_with_code_interpreter_file_generation.py) | Shows how to retrieve file IDs from code interpreter generated files using both streaming and non-streaming approaches. |
 | [`azure_ai_with_code_interpreter.py`](azure_ai_with_code_interpreter.py) | Shows how to use the HostedCodeInterpreterTool with Azure AI agents to write and execute Python code. Includes helper methods for accessing code interpreter data from response chunks. |
 | [`azure_ai_with_existing_agent.py`](azure_ai_with_existing_agent.py) | Shows how to work with a pre-existing agent by providing the agent ID to the Azure AI chat client. This example also demonstrates proper cleanup of manually created agents. |
 | [`azure_ai_with_existing_thread.py`](azure_ai_with_existing_thread.py) | Shows how to work with a pre-existing thread by providing the thread ID to the Azure AI chat client. This example also demonstrates proper cleanup of manually created threads. |
@@ -0,0 +1,102 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+
+from agent_framework import AgentRunResponseUpdate, ChatAgent, HostedCodeInterpreterTool, HostedFileContent
+from agent_framework.azure import AzureAIAgentClient
+from azure.identity.aio import AzureCliCredential
+
+"""
+Azure AI Agent Code Interpreter File Generation Example
+
+This sample demonstrates using HostedCodeInterpreterTool with AzureAIAgentClient
+to generate a text file and then retrieve it.
+
+The test flow:
+1. Create an agent with code interpreter tool
+2. Ask the agent to generate a txt file using Python code
+3. Capture the file_id from HostedFileContent in the response
+4. Retrieve the file using the agents_client.files API
+"""
+
+
+async def main() -> None:
+    """Test file generation and retrieval with code interpreter."""
+
+    async with AzureCliCredential() as credential:
+        client = AzureAIAgentClient(credential=credential)
+
+        try:
+            async with ChatAgent(
+                chat_client=client,
+                instructions=(
+                    "You are a Python code execution assistant. "
+                    "ALWAYS use the code interpreter tool to execute Python code when asked to create files. "
+                    "Write actual Python code to create files, do not just describe what you would do."
+                ),
+                tools=[HostedCodeInterpreterTool()],
+            ) as agent:
+                # Be very explicit about wanting code execution and a download link
+                query = (
+                    "Use the code interpreter to execute this Python code and then provide me "
+                    "with a download link for the generated file:\n"
+                    "```python\n"
+                    "with open('/mnt/data/sample.txt', 'w') as f:\n"
+                    "    f.write('Hello, World! This is a test file.')\n"
+                    "'/mnt/data/sample.txt'\n"  # Return the path so it becomes downloadable
+                    "```"
+                )
+                print(f"User: {query}\n")
+                print("=" * 60)
+
+                # Collect file_ids from the response
+                file_ids: list[str] = []
+
+                async for chunk in agent.run_stream(query):
+                    if not isinstance(chunk, AgentRunResponseUpdate):
+                        continue
+
+                    for content in chunk.contents:
+                        if content.type == "text":
+                            print(content.text, end="", flush=True)
+                        elif content.type == "hosted_file":
+                            if isinstance(content, HostedFileContent):
+                                file_ids.append(content.file_id)
+                                print(f"\n[File generated: {content.file_id}]")
+
+                print("\n" + "=" * 60)
+
+                # Attempt to retrieve discovered files
+                if file_ids:
+                    print(f"\nAttempting to retrieve {len(file_ids)} file(s):")
+                    for file_id in file_ids:
+                        try:
+                            file_info = await client.agents_client.files.get(file_id)
+                            print(f"  File {file_id}: Retrieved successfully")
+                            print(f"    Filename: {file_info.filename}")
+                            print(f"    Purpose: {file_info.purpose}")
+                            print(f"    Bytes: {file_info.bytes}")
+                        except Exception as e:
+                            print(f"  File {file_id}: FAILED to retrieve - {e}")
+                else:
+                    print("No file IDs were captured from the response.")
+
+                # List all files to see if any exist
+                print("\nListing all files in the agent service:")
+                try:
+                    files_list = await client.agents_client.files.list()
+                    count = 0
+                    for file_info in files_list.data:
+                        count += 1
+                        print(f"  - {file_info.id}: {file_info.filename} ({file_info.purpose})")
+                    if count == 0:
+                        print("  No files found.")
+                except Exception as e:
+                    print(f"  Failed to list files: {e}")
+
+        finally:
+            await client.close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
@@ -0,0 +1,241 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Handoff Workflow with Code Interpreter File Generation Sample
+
+This sample demonstrates retrieving file IDs from code interpreter output
+in a handoff workflow context. A triage agent routes to a code specialist
+that generates a text file, and we verify the file_id is captured correctly
+from the streaming AgentRunUpdateEvent events.
+
+Verifies GitHub issue #2718: files generated by code interpreter in
+HandoffBuilder workflows can be properly retrieved.
+
+Toggle USE_V2_CLIENT to switch between:
+    - V1: AzureAIAgentClient (azure-ai-agents SDK)
+    - V2: AzureAIClient (azure-ai-projects 2.x with Responses API)
+
+IMPORTANT: When using V2 AzureAIClient with HandoffBuilder, each agent must
+have its own client instance. The V2 client binds to a single server-side
+agent name, so sharing a client between agents causes routing issues.
+
+Prerequisites:
+    - `az login` (Azure CLI authentication)
+    - V1: AZURE_AI_AGENT_PROJECT_CONNECTION_STRING
+    - V2: AZURE_AI_PROJECT_ENDPOINT, AZURE_AI_MODEL_DEPLOYMENT_NAME
+"""
+
+import asyncio
+from collections.abc import AsyncIterable
+from contextlib import asynccontextmanager
+from collections.abc import AsyncIterator
+
+from agent_framework import (
+    AgentRunUpdateEvent,
+    ChatAgent,
+    HandoffBuilder,
+    HandoffUserInputRequest,
+    HostedCodeInterpreterTool,
+    HostedFileContent,
+    RequestInfoEvent,
+    TextContent,
+    WorkflowEvent,
+    WorkflowRunState,
+    WorkflowStatusEvent,
+)
+from azure.identity.aio import AzureCliCredential
+
+# Toggle between V1 (AzureAIAgentClient) and V2 (AzureAIClient)
+USE_V2_CLIENT = False
+
+
+async def _drain(stream: AsyncIterable[WorkflowEvent]) -> list[WorkflowEvent]:
+    """Collect all events from an async stream."""
+    return [event async for event in stream]
+
+
+def _handle_events(events: list[WorkflowEvent]) -> tuple[list[RequestInfoEvent], list[str]]:
+    """Process workflow events and extract file IDs and pending requests.
+
+    Returns:
+        Tuple of (pending_requests, file_ids_found)
+    """
+    requests: list[RequestInfoEvent] = []
+    file_ids: list[str] = []
+
+    for event in events:
+        if isinstance(event, WorkflowStatusEvent):
+            if event.state in {WorkflowRunState.IDLE, WorkflowRunState.IDLE_WITH_PENDING_REQUESTS}:
+                print(f"[status] {event.state.name}")
+
+        elif isinstance(event, RequestInfoEvent):
+            if isinstance(event.data, HandoffUserInputRequest):
+                print("\n=== Conversation So Far ===")
+                for msg in event.data.conversation:
+                    speaker = msg.author_name or msg.role.value
+                    text = msg.text or ""
+                    txt = text[:200] + "..." if len(text) > 200 else text
+                    print(f"- {speaker}: {txt}")
+                print("===========================\n")
+            requests.append(event)
+
+        elif isinstance(event, AgentRunUpdateEvent):
+            update = event.data
+            if update is None:
+                continue
+            for content in update.contents:
+                if isinstance(content, HostedFileContent):
+                    file_ids.append(content.file_id)
+                    print(f"[Found HostedFileContent: file_id={content.file_id}]")
+                elif isinstance(content, TextContent) and content.annotations:
+                    for annotation in content.annotations:
+                        if hasattr(annotation, "file_id") and annotation.file_id:
+                            file_ids.append(annotation.file_id)
+                            print(f"[Found file annotation: file_id={annotation.file_id}]")
+
+    return requests, file_ids
+
+
+@asynccontextmanager
+async def create_agents_v1(credential: AzureCliCredential) -> AsyncIterator[tuple[ChatAgent, ChatAgent]]:
+    """Create agents using V1 AzureAIAgentClient."""
+    from agent_framework.azure import AzureAIAgentClient
+
+    async with AzureAIAgentClient(credential=credential) as client:
+        triage = client.create_agent(
+            name="triage_agent",
+            instructions=(
+                "You are a triage agent. Route code-related requests to the code_specialist. "
+                "When the user asks to create or generate files, hand off to code_specialist "
+                "by calling handoff_to_code_specialist."
+            ),
+        )
+
+        code_specialist = client.create_agent(
+            name="code_specialist",
+            instructions=(
+                "You are a Python code specialist. Use the code interpreter to execute Python code "
+                "and create files when requested. Always save files to /mnt/data/ directory."
+            ),
+            tools=[HostedCodeInterpreterTool()],
+        )
+
+        yield triage, code_specialist
+
+
+@asynccontextmanager
+async def create_agents_v2(credential: AzureCliCredential) -> AsyncIterator[tuple[ChatAgent, ChatAgent]]:
+    """Create agents using V2 AzureAIClient.
+
+    Each agent needs its own client instance because the V2 client binds
+    to a single server-side agent name.
+    """
+    from agent_framework.azure import AzureAIClient
+
+    async with (
+        AzureAIClient(credential=credential) as triage_client,
+        AzureAIClient(credential=credential) as code_client,
+    ):
+        triage = triage_client.create_agent(
+            name="TriageAgent",
+            instructions=(
+                "You are a triage agent. Your ONLY job is to route requests to the appropriate specialist. "
+                "For code or file creation requests, call handoff_to_CodeSpecialist immediately. "
+                "Do NOT try to complete tasks yourself. Just hand off."
+            ),
+        )
+
+        code_specialist = code_client.create_agent(
+            name="CodeSpecialist",
+            instructions=(
+                "You are a Python code specialist. You have access to a code interpreter tool. "
+                "Use the code interpreter to execute Python code and create files. "
+                "Always save files to /mnt/data/ directory. "
+                "Do NOT discuss handoffs or routing - just complete the coding task directly."
+            ),
+            tools=[HostedCodeInterpreterTool()],
+        )
+
+        yield triage, code_specialist
+
+
+async def main() -> None:
+    """Run a simple handoff workflow with code interpreter file generation."""
+    client_version = "V2 (AzureAIClient)" if USE_V2_CLIENT else "V1 (AzureAIAgentClient)"
+    print(f"=== Handoff Workflow with Code Interpreter File Generation [{client_version}] ===\n")
+
+    async with AzureCliCredential() as credential:
+        create_agents = create_agents_v2 if USE_V2_CLIENT else create_agents_v1
+
+        async with create_agents(credential) as (triage, code_specialist):
+            workflow = (
+                HandoffBuilder()
+                .participants([triage, code_specialist])
+                .set_coordinator(triage)
+                .with_termination_condition(lambda conv: sum(1 for msg in conv if msg.role.value == "user") >= 2)
+                .build()
+            )
+
+            user_inputs = [
+                "Please create a text file called hello.txt with 'Hello from handoff workflow!' inside it.",
+                "exit",
+            ]
+            input_index = 0
+            all_file_ids: list[str] = []
+
+            print(f"User: {user_inputs[0]}")
+            events = await _drain(workflow.run_stream(user_inputs[0]))
+            requests, file_ids = _handle_events(events)
+            all_file_ids.extend(file_ids)
+            input_index += 1
+
+            while requests:
+                request = requests[0]
+                if input_index >= len(user_inputs):
+                    break
+                user_input = user_inputs[input_index]
+                print(f"\nUser: {user_input}")
+
+                responses = {request.request_id: user_input}
+                events = await _drain(workflow.send_responses_streaming(responses))
+                requests, file_ids = _handle_events(events)
+                all_file_ids.extend(file_ids)
+                input_index += 1
+
+            print("\n" + "=" * 50)
+            if all_file_ids:
+                print(f"SUCCESS: Found {len(all_file_ids)} file ID(s) in handoff workflow:")
+                for fid in all_file_ids:
+                    print(f"  - {fid}")
+            else:
+                print("WARNING: No file IDs captured from the handoff workflow.")
+            print("=" * 50)
+
+            """
+            Sample Output:
+
+            User: Please create a text file called hello.txt with 'Hello from handoff workflow!' inside it.
+            [Found HostedFileContent: file_id=assistant-JT1sA...]
+
+            === Conversation So Far ===
+            - user: Please create a text file called hello.txt with 'Hello from handoff workflow!' inside it.
+            - triage_agent: I am handing off your request to create the text file "hello.txt" with the specified content to the code specialist. They will assist you shortly.
+            - code_specialist: The file "hello.txt" has been created with the content "Hello from handoff workflow!". You can download it using the link below:
+
+            [hello.txt](sandbox:/mnt/data/hello.txt)
+            ===========================
+
+            [status] IDLE_WITH_PENDING_REQUESTS
+
+            User: exit
+            [status] IDLE
+
+            ==================================================
+            SUCCESS: Found 1 file ID(s) in handoff workflow:
+            - assistant-JT1sA...
+            ==================================================
+            """  # noqa: E501
+
+
+if __name__ == "__main__":
+    asyncio.run(main())