Skip to content
42 changes: 38 additions & 4 deletions python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@
FunctionToolDefinition,
ListSortOrder,
McpTool,
MessageAttachment,
MessageDeltaChunk,
MessageDeltaTextContent,
MessageDeltaTextFileCitationAnnotation,
Expand Down Expand Up @@ -87,10 +88,11 @@
ToolApproval,
ToolDefinition,
ToolOutput,
VectorStoreDataSource,
)
from pydantic import BaseModel

from ._shared import AzureAISettings, to_azure_ai_agent_tools
from ._shared import AzureAISettings, resolve_file_ids, to_azure_ai_agent_tools

if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
Expand Down Expand Up @@ -219,9 +221,21 @@ class AzureAIAgentClient(
# region Hosted Tool Factory Methods

@staticmethod
def get_code_interpreter_tool() -> CodeInterpreterTool:
def get_code_interpreter_tool(
*,
file_ids: list[str | Content] | 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 or Content objects to make available to
the code interpreter. Accepts plain strings or Content.from_hosted_file()
instances. The underlying SDK raises ValueError if both file_ids and
data_sources are provided.
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.

Expand All @@ -230,10 +244,21 @@ def get_code_interpreter_tool() -> CodeInterpreterTool:

from agent_framework.azure import AzureAIAgentClient

# Basic code interpreter
tool = AzureAIAgentClient.get_code_interpreter_tool()

# With uploaded file IDs
tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc123"])

# With Content objects
from agent_framework import Content

tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=[Content.from_hosted_file("file-abc123")])

agent = ChatAgent(client, tools=[tool])
"""
return CodeInterpreterTool()
resolved = resolve_file_ids(file_ids)
return CodeInterpreterTool(file_ids=resolved, data_sources=data_sources)

@staticmethod
def get_file_search_tool(
Expand Down Expand Up @@ -1291,11 +1316,19 @@ def _prepare_messages(
continue

message_contents: list[MessageInputContentBlock] = []
attachments: list[MessageAttachment] = []

for content in chat_message.contents:
match content.type:
case "text":
message_contents.append(MessageInputTextBlock(text=content.text)) # type: ignore[arg-type]
case "hosted_file":
attachments.append(
MessageAttachment(
file_id=content.file_id,
tools=CodeInterpreterTool().definitions,
)
)
case "data" | "uri":
if content.has_top_level_media_type("image"):
message_contents.append(
Expand All @@ -1310,13 +1343,14 @@ def _prepare_messages(
if isinstance(content.raw_representation, MessageInputContentBlock):
message_contents.append(content.raw_representation)

if message_contents:
if message_contents or attachments:
if additional_messages is None:
additional_messages = []
additional_messages.append(
ThreadMessageOptions(
role=MessageRole.AGENT if chat_message.role == "assistant" else MessageRole.USER,
content=message_contents,
attachments=attachments if attachments else None,
)
)

Expand Down
11 changes: 7 additions & 4 deletions python/packages/azure-ai/agent_framework_azure_ai/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@
from azure.ai.projects.models import FileSearchTool as ProjectsFileSearchTool
from azure.core.exceptions import ResourceNotFoundError

from ._shared import AzureAISettings, create_text_format_config
from ._shared import AzureAISettings, create_text_format_config, resolve_file_ids

if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
Expand Down Expand Up @@ -830,14 +830,16 @@ def _enrich_update(update: ChatResponseUpdate) -> ChatResponseUpdate:
@staticmethod
def get_code_interpreter_tool( # type: ignore[override]
*,
file_ids: list[str] | None = None,
file_ids: list[str | Content] | None = None,
container: Literal["auto"] | dict[str, Any] = "auto",
**kwargs: Any,
) -> CodeInterpreterTool:
"""Create a code interpreter tool configuration for Azure AI Projects.

Keyword Args:
file_ids: Optional list of file IDs to make available to the code interpreter.
file_ids: Optional list of file IDs or Content objects to make available to
the code interpreter. Accepts plain strings or Content.from_hosted_file()
instances.
container: Container configuration. Use "auto" for automatic container management.
Note: Custom container settings from this parameter are not used by Azure AI Projects;
use file_ids instead.
Expand All @@ -857,7 +859,8 @@ def get_code_interpreter_tool( # type: ignore[override]
# Extract file_ids from container if provided as dict and file_ids not explicitly set
if file_ids is None and isinstance(container, dict):
file_ids = container.get("file_ids")
tool_container = CodeInterpreterToolAuto(file_ids=file_ids if file_ids else None)
resolved = resolve_file_ids(file_ids)
tool_container = CodeInterpreterToolAuto(file_ids=resolved)
return CodeInterpreterTool(container=tool_container, **kwargs)

@staticmethod
Expand Down
40 changes: 40 additions & 0 deletions python/packages/azure-ai/agent_framework_azure_ai/_shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from typing import Any, cast

from agent_framework import (
Content,
FunctionTool,
)
from agent_framework.exceptions import IntegrationInvalidRequestException
Expand Down Expand Up @@ -109,6 +110,45 @@ def _extract_project_connection_id(additional_properties: dict[str, Any] | None)
return None


def resolve_file_ids(file_ids: Sequence[str | Content] | None) -> list[str] | None:
"""Resolve a list of file ID values that may include Content objects.

Accepts plain strings and Content objects with type "hosted_file", extracting
the file_id from each. This enables users to pass Content.from_hosted_file()
alongside plain file ID strings.

Args:
file_ids: Sequence of file ID strings or Content objects, or None.

Returns:
A list of resolved file ID strings, or None if input is None or empty.

Raises:
ValueError: If a Content object has an unsupported type (not "hosted_file").
"""
if not file_ids:
return None

resolved: list[str] = []
for item in file_ids:
if isinstance(item, str):
resolved.append(item)
elif isinstance(item, Content):
if item.type != "hosted_file":
raise ValueError(
f"Unsupported Content type '{item.type}' for code interpreter file_ids. "
"Only Content.from_hosted_file() is supported."
)
if item.file_id is None:
raise ValueError(
"Content.from_hosted_file() item is missing a file_id. "
"Ensure the Content object has a valid file_id before using it in file_ids."
)
resolved.append(item.file_id)

return resolved if resolved else None


def to_azure_ai_agent_tools(
tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None,
run_options: dict[str, Any] | None = None,
Expand Down
153 changes: 153 additions & 0 deletions python/packages/azure-ai/tests/test_azure_ai_agent_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -855,6 +855,95 @@ 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_get_code_interpreter_tool_with_content() -> None:
"""Test get_code_interpreter_tool accepts Content.from_hosted_file in file_ids."""
from agent_framework import Content
from azure.ai.agents.models import CodeInterpreterTool

content = Content.from_hosted_file("file-content-123")
tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=[content])
assert isinstance(tool, CodeInterpreterTool)
assert "file-content-123" in tool.file_ids


async def test_azure_ai_chat_client_get_code_interpreter_tool_with_mixed_file_ids() -> None:
"""Test get_code_interpreter_tool accepts a mix of strings and Content objects."""
from agent_framework import Content
from azure.ai.agents.models import CodeInterpreterTool

content = Content.from_hosted_file("file-from-content")
tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-plain", content])
assert isinstance(tool, CodeInterpreterTool)
assert "file-plain" in tool.file_ids
assert "file-from-content" in tool.file_ids


async def test_azure_ai_chat_client_get_code_interpreter_tool_content_unsupported_type() -> None:
"""Test get_code_interpreter_tool raises ValueError for unsupported Content types."""
from agent_framework import Content

content = Content.from_hosted_vector_store("vs-123")
with pytest.raises(ValueError, match="Unsupported Content type"):
AzureAIAgentClient.get_code_interpreter_tool(file_ids=[content])


async def test_azure_ai_chat_client_create_agent_stream_submit_tool_approvals(
mock_agents_client: MagicMock,
) -> None:
Expand Down Expand Up @@ -2121,6 +2210,70 @@ def test_azure_ai_chat_client_prepare_messages_with_raw_content_block(
assert additional_messages[0].content[0] == raw_block


def test_azure_ai_chat_client_prepare_messages_with_hosted_file_attachment(
mock_agents_client: MagicMock,
) -> None:
"""Test _prepare_messages converts hosted_file content to MessageAttachment."""
client = create_test_azure_ai_chat_client(mock_agents_client)

file_content = Content.from_hosted_file(file_id="file-abc123")
messages = [Message(role="user", contents=["Analyze this CSV.", file_content])]

additional_messages, instructions, required_action_results = client._prepare_messages(messages) # type: ignore

assert additional_messages is not None
assert len(additional_messages) == 1
msg = additional_messages[0]
# Text content should be present
assert len(msg.content) == 1
assert msg.content[0].text == "Analyze this CSV." # type: ignore[union-attr]
# Attachment should be created from hosted_file
assert msg.attachments is not None
assert len(msg.attachments) == 1
assert msg.attachments[0]["file_id"] == "file-abc123"
assert msg.attachments[0]["tools"] == [{"type": "code_interpreter"}]


def test_azure_ai_chat_client_prepare_messages_with_multiple_hosted_files(
mock_agents_client: MagicMock,
) -> None:
"""Test _prepare_messages handles multiple hosted_file contents as separate attachments."""
client = create_test_azure_ai_chat_client(mock_agents_client)

file1 = Content.from_hosted_file(file_id="file-001")
file2 = Content.from_hosted_file(file_id="file-002")
messages = [Message(role="user", contents=["Analyze both files.", file1, file2])]

additional_messages, _, _ = client._prepare_messages(messages) # type: ignore

assert additional_messages is not None
msg = additional_messages[0]
assert msg.attachments is not None
assert len(msg.attachments) == 2
assert msg.attachments[0]["file_id"] == "file-001"
assert msg.attachments[1]["file_id"] == "file-002"


def test_azure_ai_chat_client_prepare_messages_hosted_file_only(
mock_agents_client: MagicMock,
) -> None:
"""Test _prepare_messages creates a message when only hosted_file content is present (no text)."""
client = create_test_azure_ai_chat_client(mock_agents_client)

file_content = Content.from_hosted_file(file_id="file-only")
messages = [Message(role="user", contents=[file_content])]

additional_messages, _, _ = client._prepare_messages(messages) # type: ignore

assert additional_messages is not None
assert len(additional_messages) == 1
msg = additional_messages[0]
assert msg.content == []
assert msg.attachments is not None
assert len(msg.attachments) == 1
assert msg.attachments[0]["file_id"] == "file-only"


async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_mcp_tool(
mock_agents_client: MagicMock,
) -> None:
Expand Down
29 changes: 29 additions & 0 deletions python/packages/azure-ai/tests/test_azure_ai_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1685,6 +1685,35 @@ def test_get_code_interpreter_tool_with_file_ids() -> None:
assert tool["container"]["file_ids"] == ["file-123", "file-456"]


def test_get_code_interpreter_tool_with_content() -> None:
"""Test get_code_interpreter_tool accepts Content.from_hosted_file in file_ids."""
from agent_framework import Content

content = Content.from_hosted_file("file-content-123")
tool = AzureAIClient.get_code_interpreter_tool(file_ids=[content])
assert isinstance(tool, CodeInterpreterTool)
assert tool["container"]["file_ids"] == ["file-content-123"]


def test_get_code_interpreter_tool_with_mixed_file_ids() -> None:
"""Test get_code_interpreter_tool accepts a mix of strings and Content objects."""
from agent_framework import Content

content = Content.from_hosted_file("file-from-content")
tool = AzureAIClient.get_code_interpreter_tool(file_ids=["file-plain", content])
assert isinstance(tool, CodeInterpreterTool)
assert sorted(tool["container"]["file_ids"]) == ["file-from-content", "file-plain"]


def test_get_code_interpreter_tool_content_unsupported_type() -> None:
"""Test get_code_interpreter_tool raises ValueError for unsupported Content types."""
from agent_framework import Content

content = Content.from_hosted_vector_store("vs-123")
with pytest.raises(ValueError, match="Unsupported Content type"):
AzureAIClient.get_code_interpreter_tool(file_ids=[content])


def test_get_file_search_tool_basic() -> None:
"""Test get_file_search_tool returns FileSearchTool."""
tool = AzureAIClient.get_file_search_tool(vector_store_ids=["vs-123"])
Expand Down
Loading