Python: Add file_ids support to AzureAIAgentClient.get_code_interpreter_tool()

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

@@ -87,6 +87,7 @@
     ToolApproval,
     ToolDefinition,
     ToolOutput,
+    VectorStoreDataSource,
 )
 from pydantic import BaseModel
 
@@ -219,9 +220,18 @@ class AzureAIAgentClient(
     # region Hosted Tool Factory Methods
 
     @staticmethod
-    def get_code_interpreter_tool() -> CodeInterpreterTool:
+    def get_code_interpreter_tool(
+        *,
+        file_ids: list[str] | None = None,
+        data_sources: list[VectorStoreDataSource] | None = None,
+    ) -> CodeInterpreterTool:
         """Create a code interpreter tool configuration for Azure AI Agents.
 
+        Keyword Args:
+            file_ids: List of uploaded file IDs to make available to the code interpreter.
+            data_sources: List of vector store data sources for enterprise file search.
+                Mutually exclusive with file_ids.
+
         Returns:
             A CodeInterpreterTool instance ready to pass to ChatAgent.
 
@@ -230,10 +240,15 @@ def get_code_interpreter_tool() -> CodeInterpreterTool:
 
                 from agent_framework.azure import AzureAIAgentClient
 
+                # Basic code interpreter
                 tool = AzureAIAgentClient.get_code_interpreter_tool()
+
+                # With uploaded files
+                tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc123"])
+
                 agent = ChatAgent(client, tools=[tool])
         """
-        return CodeInterpreterTool()
+        return CodeInterpreterTool(file_ids=file_ids, data_sources=data_sources)
 
     @staticmethod
     def get_file_search_tool(

python/packages/azure-ai/tests/test_azure_ai_agent_client.py

@@ -858,6 +858,63 @@ async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_file_search_with_
     assert run_options["tool_resources"] == {"file_search": {"vector_store_ids": ["vs-123"]}}
 
 
+async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_code_interpreter_with_file_ids(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_tools_for_azure_ai with CodeInterpreterTool with file_ids from get_code_interpreter_tool()."""
+
+    client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
+
+    code_interpreter_tool = client.get_code_interpreter_tool(file_ids=["file-123", "file-456"])
+
+    run_options: dict[str, Any] = {}
+    result = await client._prepare_tools_for_azure_ai([code_interpreter_tool], run_options)  # type: ignore
+
+    assert len(result) == 1
+    assert result[0] == {"type": "code_interpreter"}
+    assert "tool_resources" in run_options
+    assert "code_interpreter" in run_options["tool_resources"]
+    assert sorted(run_options["tool_resources"]["code_interpreter"]["file_ids"]) == ["file-123", "file-456"]
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_basic() -> None:
+    """Test get_code_interpreter_tool returns CodeInterpreterTool without files."""
+    from azure.ai.agents.models import CodeInterpreterTool
+
+    tool = AzureAIAgentClient.get_code_interpreter_tool()
+    assert isinstance(tool, CodeInterpreterTool)
+    assert len(tool.file_ids) == 0
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_file_ids() -> None:
+    """Test get_code_interpreter_tool forwards file_ids to the SDK."""
+    from azure.ai.agents.models import CodeInterpreterTool
+
+    tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc", "file-def"])
+    assert isinstance(tool, CodeInterpreterTool)
+    assert "file-abc" in tool.file_ids
+    assert "file-def" in tool.file_ids
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_data_sources() -> None:
+    """Test get_code_interpreter_tool forwards data_sources to the SDK."""
+    from azure.ai.agents.models import CodeInterpreterTool, VectorStoreDataSource
+
+    ds = VectorStoreDataSource(asset_identifier="test-asset-id", asset_type="id_asset")
+    tool = AzureAIAgentClient.get_code_interpreter_tool(data_sources=[ds])
+    assert isinstance(tool, CodeInterpreterTool)
+    assert "test-asset-id" in tool.data_sources
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_mutually_exclusive() -> None:
+    """Test get_code_interpreter_tool raises ValueError when both file_ids and data_sources are provided."""
+    from azure.ai.agents.models import VectorStoreDataSource
+
+    ds = VectorStoreDataSource(asset_identifier="test-asset-id", asset_type="id_asset")
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc"], data_sources=[ds])
+
+
 async def test_azure_ai_chat_client_create_agent_stream_submit_tool_approvals(
     mock_agents_client: MagicMock,
 ) -> None: