diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py index 7590111bac..81c1e74610 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py +++ b/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( diff --git a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py index 1daf611bef..771dceae31 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py +++ b/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: