diff --git a/python/packages/main/agent_framework/_tools.py b/python/packages/main/agent_framework/_tools.py index 5bea36015d..fd4c76a7bc 100644 --- a/python/packages/main/agent_framework/_tools.py +++ b/python/packages/main/agent_framework/_tools.py @@ -4,7 +4,7 @@ import inspect from collections.abc import Awaitable, Callable from functools import wraps from time import perf_counter -from typing import Annotated, Any, Generic, Protocol, TypeVar, get_args, get_origin, runtime_checkable +from typing import TYPE_CHECKING, Annotated, Any, Generic, Protocol, TypeVar, get_args, get_origin, runtime_checkable from opentelemetry import metrics, trace from pydantic import BaseModel, Field, create_model @@ -12,6 +12,9 @@ from pydantic import BaseModel, Field, create_model from ._logging import get_logger from .telemetry import GenAIAttributes, start_as_current_span +if TYPE_CHECKING: + from ._types import AIContents + tracer: trace.Tracer = trace.get_tracer("agent_framework") meter: metrics.Meter = metrics.get_meter_provider().get_meter("agent_framework") logger = get_logger() @@ -216,6 +219,45 @@ def ai_function( return decorator(func) if func else decorator # type: ignore[reportReturnType, return-value] +def _parse_inputs( + inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None", +) -> list["AIContents"]: + """Parse the inputs for a tool, ensuring they are of type AIContents.""" + if inputs is None: + return [] + + from ._types import AIContent, DataContent, HostedFileContent, HostedVectorStoreContent, UriContent + + parsed_inputs: list["AIContents"] = [] + if not isinstance(inputs, list): + inputs = [inputs] + for input_item in inputs: + if isinstance(input_item, str): + # If it's a string, we assume it's a URI or similar identifier. + # Convert it to a UriContent or similar type as needed. + parsed_inputs.append(UriContent(uri=input_item, media_type="text/plain")) + elif isinstance(input_item, dict): + # If it's a dict, we assume it contains properties for a specific content type. + # we check if the required keys are present to determine the type. + if "uri" in input_item: + parsed_inputs.append( + UriContent(**input_item) if "media_type" in input_item else DataContent(**input_item) + ) + elif "file_id" in input_item: + parsed_inputs.append(HostedFileContent(**input_item)) + elif "vector_store_id" in input_item: + parsed_inputs.append(HostedVectorStoreContent(**input_item)) + elif "data" in input_item: + parsed_inputs.append(DataContent(**input_item)) + else: + raise ValueError(f"Unsupported input type: {input_item}") + elif isinstance(input_item, AIContent): + parsed_inputs.append(input_item) + else: + raise TypeError(f"Unsupported input type: {type(input_item).__name__}. Expected AIContents or dict.") + return parsed_inputs + + class HostedCodeInterpreterTool(AITool): """Represents a hosted tool that can be specified to an AI service to enable it to execute generated code. @@ -226,6 +268,7 @@ class HostedCodeInterpreterTool(AITool): def __init__( self, name: str = "code_interpreter", + inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None" = None, description: str | None = None, additional_properties: dict[str, Any] | None = None, ): @@ -233,10 +276,19 @@ class HostedCodeInterpreterTool(AITool): Args: name: The name of the tool. Defaults to "code_interpreter". + inputs: A list of contents that the tool can accept as input. Defaults to None. + This should mostly be HostedFileContent or HostedVectorStoreContent. + Can also be DataContent, depending on the service used. + When supplying a list, it can contain: + - AIContents instances + - dicts with properties for AIContents (e.g., {"uri": "http://example.com", "media_type": "text/html"}) + - strings (which will be converted to UriContent with media_type "text/plain"). + If None, defaults to an empty list. description: A description of the tool. additional_properties: Additional properties associated with the tool, specific to the service used. """ self.name = name + self.inputs = _parse_inputs(inputs) self.description = description self.additional_properties = additional_properties diff --git a/python/packages/main/agent_framework/_types.py b/python/packages/main/agent_framework/_types.py index 1a0cc83d67..f20b17590a 100644 --- a/python/packages/main/agent_framework/_types.py +++ b/python/packages/main/agent_framework/_types.py @@ -99,6 +99,8 @@ __all__ = [ "FunctionCallContent", "FunctionResultContent", "GeneratedEmbeddings", + "HostedFileContent", + "HostedVectorStoreContent", "SpeechToTextOptions", "StructuredResponse", "TextContent", @@ -1036,6 +1038,68 @@ class UsageContent(AIContent): ) +class HostedFileContent(AIContent): + """Represents a hosted file content. + + Attributes: + file_id: The identifier of the hosted file. + type: The type of content, which is always "hosted_file" for this class. + additional_properties: Optional additional properties associated with the content. + raw_representation: Optional raw representation of the content. + + """ + + type: Literal["hosted_file"] = "hosted_file" # type: ignore[assignment] + file_id: str + + def __init__( + self, + file_id: str, + *, + additional_properties: dict[str, Any] | None = None, + raw_representation: Any | None = None, + **kwargs: Any, + ) -> None: + """Initializes a HostedFileContent instance.""" + super().__init__( + file_id=file_id, # type: ignore[reportCallIssue] + additional_properties=additional_properties, + raw_representation=raw_representation, + **kwargs, + ) + + +class HostedVectorStoreContent(AIContent): + """Represents a hosted vector store content. + + Attributes: + vector_store_id: The identifier of the hosted vector store. + type: The type of content, which is always "hosted_vector_store" for this class. + additional_properties: Optional additional properties associated with the content. + raw_representation: Optional raw representation of the content. + + """ + + type: Literal["hosted_vector_store"] = "hosted_vector_store" # type: ignore[assignment] + vector_store_id: str + + def __init__( + self, + vector_store_id: str, + *, + additional_properties: dict[str, Any] | None = None, + raw_representation: Any | None = None, + **kwargs: Any, + ) -> None: + """Initializes a HostedVectorStoreContent instance.""" + super().__init__( + vector_store_id=vector_store_id, # type: ignore[reportCallIssue] + additional_properties=additional_properties, + raw_representation=raw_representation, + **kwargs, + ) + + AIContents = Annotated[ TextContent | DataContent @@ -1044,7 +1108,9 @@ AIContents = Annotated[ | FunctionCallContent | FunctionResultContent | ErrorContent - | UsageContent, + | UsageContent + | HostedFileContent + | HostedVectorStoreContent, Field(discriminator="type"), ] diff --git a/python/packages/main/tests/main/test_tools.py b/python/packages/main/tests/main/test_tools.py index be8cfb5f14..38918cf814 100644 --- a/python/packages/main/tests/main/test_tools.py +++ b/python/packages/main/tests/main/test_tools.py @@ -5,7 +5,8 @@ from unittest.mock import Mock, patch import pytest from pydantic import BaseModel -from agent_framework import AIFunction, AITool, ai_function +from agent_framework import AIFunction, AITool, HostedCodeInterpreterTool, ai_function +from agent_framework._tools import _parse_inputs from agent_framework.telemetry import GenAIAttributes @@ -288,3 +289,231 @@ async def test_ai_function_invoke_invalid_pydantic_args(): # Call invoke with wrong model type with pytest.raises(TypeError, match="Expected invalid_args_test_input, got WrongModel"): await invalid_args_test.invoke(arguments=wrong_args) + + +# Tests for HostedCodeInterpreterTool and _parse_inputs + + +def test_hosted_code_interpreter_tool_default(): + """Test HostedCodeInterpreterTool with default parameters.""" + tool = HostedCodeInterpreterTool() + + assert tool.name == "code_interpreter" + assert tool.inputs == [] + assert tool.description is None + assert tool.additional_properties is None + assert str(tool) == "HostedCodeInterpreterTool(name=code_interpreter)" + + +def test_hosted_code_interpreter_tool_custom_name(): + """Test HostedCodeInterpreterTool with custom name.""" + tool = HostedCodeInterpreterTool(name="custom_interpreter") + + assert tool.name == "custom_interpreter" + assert tool.inputs == [] + assert str(tool) == "HostedCodeInterpreterTool(name=custom_interpreter)" + + +def test_hosted_code_interpreter_tool_with_description(): + """Test HostedCodeInterpreterTool with description and additional properties.""" + tool = HostedCodeInterpreterTool( + name="test_interpreter", + description="A test code interpreter", + additional_properties={"version": "1.0", "language": "python"}, + ) + + assert tool.name == "test_interpreter" + assert tool.description == "A test code interpreter" + assert tool.additional_properties == {"version": "1.0", "language": "python"} + + +def test_parse_inputs_none(): + """Test _parse_inputs with None input.""" + result = _parse_inputs(None) + assert result == [] + + +def test_parse_inputs_string(): + """Test _parse_inputs with string input.""" + from agent_framework import UriContent + + result = _parse_inputs("http://example.com") + assert len(result) == 1 + assert isinstance(result[0], UriContent) + assert result[0].uri == "http://example.com" + assert result[0].media_type == "text/plain" + + +def test_parse_inputs_list_of_strings(): + """Test _parse_inputs with list of strings.""" + from agent_framework import UriContent + + inputs = ["http://example.com", "https://test.org"] + result = _parse_inputs(inputs) + + assert len(result) == 2 + assert all(isinstance(item, UriContent) for item in result) + assert result[0].uri == "http://example.com" + assert result[1].uri == "https://test.org" + assert all(item.media_type == "text/plain" for item in result) + + +def test_parse_inputs_uri_dict(): + """Test _parse_inputs with URI dictionary.""" + from agent_framework import UriContent + + input_dict = {"uri": "http://example.com", "media_type": "application/json"} + result = _parse_inputs(input_dict) + + assert len(result) == 1 + assert isinstance(result[0], UriContent) + assert result[0].uri == "http://example.com" + assert result[0].media_type == "application/json" + + +def test_parse_inputs_hosted_file_dict(): + """Test _parse_inputs with hosted file dictionary.""" + from agent_framework import HostedFileContent + + input_dict = {"file_id": "file-123"} + result = _parse_inputs(input_dict) + + assert len(result) == 1 + assert isinstance(result[0], HostedFileContent) + assert result[0].file_id == "file-123" + + +def test_parse_inputs_hosted_vector_store_dict(): + """Test _parse_inputs with hosted vector store dictionary.""" + from agent_framework import HostedVectorStoreContent + + input_dict = {"vector_store_id": "vs-789"} + result = _parse_inputs(input_dict) + + assert len(result) == 1 + assert isinstance(result[0], HostedVectorStoreContent) + assert result[0].vector_store_id == "vs-789" + + +def test_parse_inputs_data_dict(): + """Test _parse_inputs with data dictionary.""" + from agent_framework import DataContent + + input_dict = {"data": b"test data", "media_type": "application/octet-stream"} + result = _parse_inputs(input_dict) + + assert len(result) == 1 + assert isinstance(result[0], DataContent) + assert result[0].uri == "data:application/octet-stream;base64,dGVzdCBkYXRh" + assert result[0].media_type == "application/octet-stream" + + +def test_parse_inputs_ai_contents_instance(): + """Test _parse_inputs with AIContents instance.""" + from agent_framework import TextContent + + text_content = TextContent(text="Hello, world!") + result = _parse_inputs(text_content) + + assert len(result) == 1 + assert isinstance(result[0], TextContent) + assert result[0].text == "Hello, world!" + + +def test_parse_inputs_mixed_list(): + """Test _parse_inputs with mixed input types.""" + from agent_framework import HostedFileContent, TextContent, UriContent + + inputs = [ + "http://example.com", # string + {"uri": "https://test.org", "media_type": "text/html"}, # URI dict + {"file_id": "file-456"}, # hosted file dict + TextContent(text="Hello"), # AIContents instance + ] + + result = _parse_inputs(inputs) + + assert len(result) == 4 + assert isinstance(result[0], UriContent) + assert result[0].uri == "http://example.com" + assert isinstance(result[1], UriContent) + assert result[1].uri == "https://test.org" + assert result[1].media_type == "text/html" + assert isinstance(result[2], HostedFileContent) + assert result[2].file_id == "file-456" + assert isinstance(result[3], TextContent) + assert result[3].text == "Hello" + + +def test_parse_inputs_unsupported_dict(): + """Test _parse_inputs with unsupported dictionary format.""" + input_dict = {"unsupported_key": "value"} + + with pytest.raises(ValueError, match="Unsupported input type"): + _parse_inputs(input_dict) + + +def test_parse_inputs_unsupported_type(): + """Test _parse_inputs with unsupported input type.""" + with pytest.raises(TypeError, match="Unsupported input type: int"): + _parse_inputs(123) + + +def test_hosted_code_interpreter_tool_with_string_input(): + """Test HostedCodeInterpreterTool with string input.""" + from agent_framework import UriContent + + tool = HostedCodeInterpreterTool(inputs="http://example.com") + + assert len(tool.inputs) == 1 + assert isinstance(tool.inputs[0], UriContent) + assert tool.inputs[0].uri == "http://example.com" + + +def test_hosted_code_interpreter_tool_with_dict_inputs(): + """Test HostedCodeInterpreterTool with dictionary inputs.""" + from agent_framework import HostedFileContent, UriContent + + inputs = [{"uri": "http://example.com", "media_type": "text/html"}, {"file_id": "file-123"}] + + tool = HostedCodeInterpreterTool(inputs=inputs) + + assert len(tool.inputs) == 2 + assert isinstance(tool.inputs[0], UriContent) + assert tool.inputs[0].uri == "http://example.com" + assert tool.inputs[0].media_type == "text/html" + assert isinstance(tool.inputs[1], HostedFileContent) + assert tool.inputs[1].file_id == "file-123" + + +def test_hosted_code_interpreter_tool_with_ai_contents(): + """Test HostedCodeInterpreterTool with AIContents instances.""" + from agent_framework import DataContent, TextContent + + inputs = [TextContent(text="Hello, world!"), DataContent(data=b"test", media_type="text/plain")] + + tool = HostedCodeInterpreterTool(inputs=inputs) + + assert len(tool.inputs) == 2 + assert isinstance(tool.inputs[0], TextContent) + assert tool.inputs[0].text == "Hello, world!" + assert isinstance(tool.inputs[1], DataContent) + assert tool.inputs[1].media_type == "text/plain" + + +def test_hosted_code_interpreter_tool_with_single_input(): + """Test HostedCodeInterpreterTool with single input (not in list).""" + from agent_framework import HostedFileContent + + input_dict = {"file_id": "file-single"} + tool = HostedCodeInterpreterTool(inputs=input_dict) + + assert len(tool.inputs) == 1 + assert isinstance(tool.inputs[0], HostedFileContent) + assert tool.inputs[0].file_id == "file-single" + + +def test_hosted_code_interpreter_tool_with_unknown_input(): + """Test HostedCodeInterpreterTool with single unknown input.""" + with pytest.raises(ValueError, match="Unsupported input type"): + HostedCodeInterpreterTool(inputs={"hosted_file": "file-single"}) diff --git a/python/packages/main/tests/main/test_types.py b/python/packages/main/tests/main/test_types.py index fb51cf6de5..0f4ed6fd0e 100644 --- a/python/packages/main/tests/main/test_types.py +++ b/python/packages/main/tests/main/test_types.py @@ -27,6 +27,8 @@ from agent_framework import ( FunctionCallContent, FunctionResultContent, GeneratedEmbeddings, + HostedFileContent, + HostedVectorStoreContent, SpeechToTextOptions, StructuredResponse, TextContent, @@ -187,6 +189,68 @@ def test_uri_content(): assert isinstance(content, AIContent) +# region: HostedFileContent + + +def test_hosted_file_content(): + """Test the HostedFileContent class to ensure it initializes correctly.""" + content = HostedFileContent(file_id="file-123", additional_properties={"version": 1}) + + # Check the type and content + assert content.type == "hosted_file" + assert content.file_id == "file-123" + assert content.additional_properties["version"] == 1 + + # Ensure the instance is of type AIContent + assert isinstance(content, AIContent) + + +def test_hosted_file_content_minimal(): + """Test the HostedFileContent class with minimal parameters.""" + content = HostedFileContent(file_id="file-456") + + # Check the type and content + assert content.type == "hosted_file" + assert content.file_id == "file-456" + assert content.additional_properties is None + assert content.raw_representation is None + + # Ensure the instance is of type AIContent + assert isinstance(content, AIContent) + + +# region: HostedVectorStoreContent + + +def test_hosted_vector_store_content(): + """Test the HostedVectorStoreContent class to ensure it initializes correctly.""" + content = HostedVectorStoreContent(vector_store_id="vs-789", additional_properties={"version": 1}) + + # Check the type and content + assert content.type == "hosted_vector_store" + assert content.vector_store_id == "vs-789" + assert content.additional_properties["version"] == 1 + + # Ensure the instance is of type AIContent + assert isinstance(content, HostedVectorStoreContent) + assert isinstance(content, AIContent) + + +def test_hosted_vector_store_content_minimal(): + """Test the HostedVectorStoreContent class with minimal parameters.""" + content = HostedVectorStoreContent(vector_store_id="vs-101112") + + # Check the type and content + assert content.type == "hosted_vector_store" + assert content.vector_store_id == "vs-101112" + assert content.additional_properties is None + assert content.raw_representation is None + + # Ensure the instance is of type AIContent + assert isinstance(content, HostedVectorStoreContent) + assert isinstance(content, AIContent) + + # region FunctionCallContent @@ -328,6 +392,8 @@ def test_usage_details_add_with_none_and_type_errors(): (UriContent, {"uri": "http://example.com", "media_type": "text/html"}), (FunctionCallContent, {"call_id": "1", "name": "example_function", "arguments": {}}), (FunctionResultContent, {"call_id": "1", "result": {}}), + (HostedFileContent, {"file_id": "file-123"}), + (HostedVectorStoreContent, {"vector_store_id": "vs-789"}), ], ) def test_ai_content_serialization(content_type: type[AIContent], args: dict):