From 9148392d00dafa47a95178622f3800f89bbf0425 Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Mon, 13 Oct 2025 13:34:27 -0700 Subject: [PATCH] Python: Enhanced documentation for dependency injection and serialization features (#1324) * improvements in dep injection and sample * fix for falsy default * fix mypy * update to use a nested dict instead of a string. * clarify docs * Update python/packages/core/agent_framework/_tools.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/core/agent_framework/_tools.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/core/agent_framework/_tools.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/core/agent_framework/_serialization.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * format --------- Co-authored-by: eavanvalkenburg Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../core/agent_framework/_serialization.py | 498 ++++++++++++++---- .../packages/core/agent_framework/_tools.py | 122 ++++- .../tests/core/test_serializable_mixin.py | 16 +- python/packages/core/tests/core/test_tools.py | 42 ++ .../tools/tool_with_injected_func.py | 68 +++ 5 files changed, 618 insertions(+), 128 deletions(-) create mode 100644 python/samples/getting_started/tools/tool_with_injected_func.py diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index f3cb42185a..1a38d9030a 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -2,7 +2,7 @@ import json import re -from collections.abc import MutableMapping +from collections.abc import Mapping, MutableMapping from typing import Any, ClassVar, Protocol, TypeVar, runtime_checkable from ._logging import get_logger @@ -20,28 +20,65 @@ _CAMEL_TO_SNAKE_PATTERN = re.compile(r"(? dict[str, Any]: @@ -74,73 +111,176 @@ class SerializationProtocol(Protocol): def is_serializable(value: Any) -> bool: """Check if a value is JSON serializable. + This function tests whether a value can be directly serialized to JSON + without custom encoding. It checks for basic Python types that have + direct JSON equivalents. + Args: - value: The value to check. + value: The value to check for JSON serializability. Returns: - True if the value is JSON serializable, False otherwise. + True if the value is one of the basic JSON-serializable types + (str, int, float, bool, None, list, dict), False otherwise. + + Note: + This function only checks for direct JSON compatibility. Complex objects + that implement ``SerializationProtocol`` require conversion via ``to_dict()`` + before JSON serialization. """ return isinstance(value, (str, int, float, bool, type(None), list, dict)) class SerializationMixin: - """Mixin class providing serialization and deserialization capabilities. + """Mixin class providing comprehensive serialization and deserialization capabilities. - Classes using this mixin should handle MutableMapping inputs in their __init__ method - for any parameters that expect SerializationMixin/SerializationProtocol instances. - The __init__ should check if the value is a MutableMapping and call from_dict() to convert it. + .. note:: + SerializationMixin is in active development. The API may change in future versions + as we continue to improve and extend its functionality. - So take the two classes below as an example. The first purely uses base types, strings in this case. - The second has a param that is of the type of the first class. - Because we setup the __init__ method to handle MutableMapping, - we can pass in a dict to the second class and it will convert it to an instance of the first class. + This mixin enables classes to automatically handle complex serialization scenarios + including nested objects, dependency injection, and type conversion. It provides + robust support for converting objects to/from dictionaries and JSON strings while + maintaining object relationships and handling external dependencies. + + **Key Features:** + + - Automatic serialization of nested SerializationProtocol objects + - Support for lists and dictionaries containing serializable objects + - Dependency injection system for non-serializable external dependencies + - Flexible exclusion of fields from serialization + - Type-safe deserialization with automatic type conversion + + **Constructor Pattern for Nested Objects:** + + Classes using this mixin should handle ``MutableMapping`` inputs in their ``__init__`` method + for any parameters that expect ``SerializationMixin`` or ``SerializationProtocol`` instances. + This enables automatic conversion of dictionaries to proper object instances during deserialization. + + **Dependency Injection System:** + + The mixin supports injecting external dependencies (like database connections, API clients, + or configuration objects) that shouldn't be serialized but are needed at runtime. + Fields marked in ``INJECTABLE`` are excluded during serialization and can be provided + during deserialization via the ``dependencies`` parameter. Examples: + **Nested object serialization with agent thread management:** + .. code-block:: python - class SerializableMixinType(SerializationMixin): - def __init__(self, param1: str, param2: int) -> None: - self.param1 = param1 - self.param2 = param2 + from agent_framework import ChatMessage + from agent_framework._threads import AgentThreadState, ChatMessageStoreState - class MyClass(SerializationMixin): - def __init__( - self, - regular_param: str, - param: SerializableMixinType | MutableMapping[str, Any] | None = None, - ) -> None: - if isinstance(param, MutableMapping): - self.param = self.from_dict(param) - else: - self.param = param - self.regular_param = regular_param + # ChatMessageStoreState handles nested ChatMessage serialization + store_state = ChatMessageStoreState( + messages=[ + ChatMessage(role="user", text="Hello agent"), + ChatMessage(role="assistant", text="Hi! How can I help?"), + ] + ) + # Nested serialization: messages are automatically converted to dicts + store_dict = store_state.to_dict() + # Result: { + # "type": "chat_message_store_state", + # "messages": [ + # {"type": "chat_message", "role": {...}, "contents": [...]}, + # {"type": "chat_message", "role": {...}, "contents": [...]} + # ] + # } - instance = MyClass.from_dict({"regular_param": "value", "param": {"param1": "value1", "param2": 42}}) + # AgentThreadState contains nested ChatMessageStoreState + thread_state = AgentThreadState(chat_message_store_state=store_state) - A more complex use case involves an injectable dependency that is not serialized. - In this case, the dependency is passed in via the dependencies parameter to from_dict/from_json. + # Deep serialization: nested SerializationMixin objects are handled automatically + thread_dict = thread_state.to_dict() + # The chat_message_store_state and its nested messages are all serialized + + # Reconstruction from nested dictionaries with automatic type conversion + # The __init__ method handles MutableMapping -> object conversion: + reconstructed = AgentThreadState.from_dict({ + "chat_message_store_state": {"messages": [{"role": "user", "text": "Hello again"}]} + }) + # chat_message_store_state becomes ChatMessageStoreState instance automatically + + **Framework tools with exclusion patterns:** - Examples: .. code-block:: python - from library import Client + from agent_framework._tools import BaseTool - class MyClass(SerializationMixin): - INJECTABLE = {"client"} + class WeatherTool(BaseTool): + \"\"\"Example tool that extends BaseTool with additional properties exclusion.\"\"\" - During serialization, the field listed as INJECTABLE (and also DEFAULT_EXCLUDE) will be excluded from the output. - Then in deserialization, - the dependencies dict is checked for any keys matching the formats: - - "." - - ".." - where is the type identifier for the class (either the value of the 'type' class variable or - the snake_cased class name if 'type' is not present), - is the name of the parameter in the __init__ method, - is the name of a parameter that is a dict, - and is a key in that dict parameter. + # Inherits DEFAULT_EXCLUDE = {"additional_properties"} from BaseTool + + def __init__(self, name: str, api_key: str, **kwargs): + super().__init__(name=name, description="Get weather information", **kwargs) + self.api_key = api_key # Will be serialized + + # Additional properties are excluded from serialization + self.additional_properties = {"version": "1.0", "internal_config": {...}} + + + weather_tool = WeatherTool(name="get_weather", api_key="secret-key") + + # Serialization excludes additional_properties but includes other fields + tool_dict = weather_tool.to_dict() + # Result: { + # "type": "weather_tool", + # "name": "get_weather", + # "description": "Get weather information", + # "api_key": "secret-key" + # # additional_properties excluded due to DEFAULT_EXCLUDE + # } + + **Agent framework with injectable dependencies:** + + .. code-block:: python + + from agent_framework import BaseAgent + + + class CustomAgent(BaseAgent): + \"\"\"Custom agent extending BaseAgent with additional functionality.\"\"\" + + # Inherits DEFAULT_EXCLUDE = {"additional_properties"} from BaseAgent + + def __init__(self, **kwargs): + super().__init__(name="custom-agent", description="A custom agent", **kwargs) + + # additional_properties stores runtime configuration but isn't serialized + self.additional_properties.update({ + "runtime_context": {...}, + "session_data": {...} + }) + + + agent = CustomAgent( + context_providers=[...], + middleware=[...] + ) + + # Serialization captures agent configuration but excludes runtime data + agent_dict = agent.to_dict() + # Result: { + # "type": "custom_agent", + # "id": "...", + # "name": "custom-agent", + # "description": "A custom agent", + # "context_provider": [...], + # "middleware": [...] + # # additional_properties excluded + # } + + # Agent can be reconstructed with the same configuration + restored_agent = CustomAgent.from_dict(agent_dict) + + This approach enables the agent framework to maintain clean separation between + persistent configuration and transient runtime state, allowing agents and tools + to be serialized for storage or transmission while preserving their functionality. """ DEFAULT_EXCLUDE: ClassVar[set[str]] = set() @@ -149,12 +289,22 @@ class SerializationMixin: def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Convert the instance and any nested objects to a dictionary. + This method performs deep serialization, automatically converting nested + ``SerializationProtocol`` objects, lists, and dictionaries containing + serializable objects. Non-serializable objects are skipped with debug logging. + + Fields marked in ``DEFAULT_EXCLUDE`` and ``INJECTABLE`` are automatically + excluded from the output, as are any private attributes (starting with '_'). + Keyword Args: - exclude: The set of field names to exclude from serialization. - exclude_none: Whether to exclude None values from the output. Defaults to True. + exclude: Additional field names to exclude from serialization beyond + the default exclusions (``DEFAULT_EXCLUDE`` and ``INJECTABLE``). + exclude_none: Whether to exclude None values from the output. When True, + None values are omitted from the dictionary. Defaults to True. Returns: - Dictionary representation of the instance. + Dictionary representation of the instance including a 'type' field + for type identification during deserialization (unless 'type' is excluded). """ # Combine exclude sets combined_exclude = set(self.DEFAULT_EXCLUDE) @@ -215,10 +365,17 @@ class SerializationMixin: def to_json(self, *, exclude: set[str] | None = None, exclude_none: bool = True, **kwargs: Any) -> str: """Convert the instance to a JSON string. + This is a convenience method that calls ``to_dict()`` and then serializes + the result using ``json.dumps()``. All the same serialization rules apply + as in ``to_dict()``, including automatic exclusion of injectable dependencies + and deep serialization of nested objects. + Keyword Args: - exclude: The set of field names to exclude from serialization. + exclude: Additional field names to exclude from serialization. exclude_none: Whether to exclude None values from the output. Defaults to True. - **kwargs: passed through to the json.dumps method. + **kwargs: Additional keyword arguments passed through to ``json.dumps()``. + Common options include ``indent`` for pretty-printing and + ``ensure_ascii`` for Unicode handling. Returns: JSON string representation of the instance. @@ -229,56 +386,168 @@ class SerializationMixin: def from_dict( cls: type[TClass], value: MutableMapping[str, Any], /, *, dependencies: MutableMapping[str, Any] | None = None ) -> TClass: - """Create an instance from a dictionary. + """Create an instance from a dictionary with optional dependency injection. + + This method reconstructs an object from its dictionary representation, automatically + handling type conversion and dependency injection. It supports three patterns of + dependency injection to handle different scenarios where external dependencies + need to be provided at deserialization time. Args: value: The dictionary containing the instance data (positional-only). + Must include a 'type' field matching the class type identifier. Keyword Args: - dependencies: The dictionary mapping dependency keys to values. - Keys should be in format ``"."`` or ``".."``. + dependencies: A nested dictionary mapping type identifiers to their injectable dependencies. + The structure varies based on injection pattern: + + - **Simple injection**: ``{"": {"": value}}`` + - **Dict parameter injection**: ``{"": {"": {"": value}}}`` + - **Instance-specific injection**: ``{"": {":": {"": value}}}`` Returns: - New instance of the class. + New instance of the class with injected dependencies. + + Raises: + ValueError: If the 'type' field in the data doesn't match the class type identifier. + + Examples: + **Simple Client Injection** - OpenAI client dependency injection: + + .. code-block:: python + + from agent_framework.openai import OpenAIChatClient + from openai import AsyncOpenAI + + + # OpenAI chat client requires an AsyncOpenAI client instance + # The client is marked as INJECTABLE = {"client"} in OpenAIBase + + # Serialized data contains only the model configuration + client_data = { + "type": "open_ai_chat_client", + "model_id": "gpt-4o-mini", + # client is excluded from serialization + } + + # Provide the OpenAI client during deserialization + openai_client = AsyncOpenAI(api_key="your-api-key") + dependencies = {"open_ai_chat_client": {"client": openai_client}} + + # The chat client is reconstructed with the OpenAI client injected + chat_client = OpenAIChatClient.from_dict(client_data, dependencies=dependencies) + # Now ready to make API calls with the injected client + + **Function Injection for Tools** - AIFunction runtime dependency: + + .. code-block:: python + + from agent_framework import AIFunction + from typing import Annotated + + + # Define a function to be wrapped + async def get_current_weather(location: Annotated[str, "The city name"]) -> str: + # In real implementation, this would call a weather API + return f"Current weather in {location}: 72°F and sunny" + + + # AIFunction has INJECTABLE = {"func"} + function_data = { + "type": "ai_function", + "name": "get_weather", + "description": "Get current weather for a location", + # func is excluded from serialization + } + + # Inject the actual function implementation during deserialization + dependencies = {"ai_function": {"func": get_current_weather}} + + # Reconstruct the AIFunction with the callable injected + weather_func = AIFunction.from_dict(function_data, dependencies=dependencies) + # The function is now callable and ready for agent use + + **Middleware Context Injection** - Agent execution context: + + .. code-block:: python + + from agent_framework._middleware import AgentRunContext + from agent_framework import BaseAgent + + # AgentRunContext has INJECTABLE = {"agent", "result"} + context_data = { + "type": "agent_run_context", + "messages": [{"role": "user", "text": "Hello"}], + "is_streaming": False, + "metadata": {"session_id": "abc123"}, + # agent and result are excluded from serialization + } + + # Inject agent and result during middleware processing + my_agent = BaseAgent(name="test-agent") + dependencies = { + "agent_run_context": { + "agent": my_agent, + "result": None, # Will be populated during execution + } + } + + # Reconstruct context with agent dependency for middleware chain + context = AgentRunContext.from_dict(context_data, dependencies=dependencies) + # Middleware can now access context.agent and process the execution + + This injection system allows the agent framework to maintain clean separation + between serializable configuration and runtime dependencies like API clients, + functions, and execution contexts that cannot or should not be persisted. """ if dependencies is None: dependencies = {} # Get the type identifier - type_id = cls._get_type_identifier() + type_id = cls._get_type_identifier(value) + + if (supplied_type := value.get("type")) and supplied_type != type_id: + raise ValueError(f"Type mismatch: expected '{type_id}', got '{supplied_type}'") # Create a copy of the value dict to work with, filtering out the 'type' key kwargs = {k: v for k, v in value.items() if k != "type"} - # Process dependencies - for dep_key, dep_value in dependencies.items(): - parts = dep_key.split(".") - if len(parts) < 2: - continue - - dep_type = parts[0] - if dep_type != type_id: - continue - - param_name = parts[1] - - # Log debug message if dependency is not in INJECTABLE - if param_name not in cls.INJECTABLE: - logger.debug( - f"Dependency '{param_name}' for type '{type_id}' is not in INJECTABLE set. " - f"Available injectable parameters: {cls.INJECTABLE}" - ) - - if len(parts) == 2: - # Simple parameter: . - kwargs[param_name] = dep_value - elif len(parts) == 3: - # Dict parameter: .. - dict_param_name = parts[1] - key = parts[2] - if dict_param_name not in kwargs: - kwargs[dict_param_name] = {} - kwargs[dict_param_name][key] = dep_value + # Process dependencies using dict-based structure + type_deps = dependencies.get(type_id, {}) + for dep_key, dep_value in type_deps.items(): + # Check if this is an instance-specific dependency (field:name format) + if ":" in dep_key: + field, name = dep_key.split(":", 1) + # Only apply if the instance matches + if kwargs.get(field) == name and isinstance(dep_value, dict): + # Apply instance-specific dependencies + for param_name, param_value in dep_value.items(): + if param_name not in cls.INJECTABLE: + logger.debug( + f"Dependency '{param_name}' for type '{type_id}' is not in INJECTABLE set. " + f"Available injectable parameters: {cls.INJECTABLE}" + ) + # Handle nested dict parameters + if ( + isinstance(param_value, dict) + and param_name in kwargs + and isinstance(kwargs[param_name], dict) + ): + kwargs[param_name].update(param_value) + else: + kwargs[param_name] = param_value + else: + # Regular parameter dependency + if dep_key not in cls.INJECTABLE: + logger.debug( + f"Dependency '{dep_key}' for type '{type_id}' is not in INJECTABLE set. " + f"Available injectable parameters: {cls.INJECTABLE}" + ) + # Handle dict parameters - merge if both are dicts + if isinstance(dep_value, dict) and dep_key in kwargs and isinstance(kwargs[dep_key], dict): + kwargs[dep_key].update(dep_value) + else: + kwargs[dep_key] = dep_value return cls(**kwargs) @@ -286,31 +555,56 @@ class SerializationMixin: def from_json(cls: type[TClass], value: str, /, *, dependencies: MutableMapping[str, Any] | None = None) -> TClass: """Create an instance from a JSON string. + This is a convenience method that parses the JSON string using ``json.loads()`` + and then calls ``from_dict()`` to reconstruct the object. All dependency injection + capabilities are available through the ``dependencies`` parameter. + Args: value: The JSON string containing the instance data (positional-only). + Must be valid JSON that deserializes to a dictionary with a 'type' field. Keyword Args: - dependencies: The dictionary mapping dependency keys to values. - Keys should be in format ``"."`` or ``".."``. + dependencies: A nested dictionary mapping type identifiers to their injectable dependencies. + See :meth:`from_dict` for detailed structure and examples of the three + injection patterns (simple, dict parameter, and instance-specific). Returns: - New instance of the class. + New instance of the class with any specified dependencies injected. + + Raises: + json.JSONDecodeError: If the JSON string is malformed. + ValueError: If the parsed data doesn't contain a valid 'type' field. """ data = json.loads(value) return cls.from_dict(data, dependencies=dependencies) @classmethod - def _get_type_identifier(cls) -> str: + def _get_type_identifier(cls, value: Mapping[str, Any] | None = None) -> str: """Get the type identifier for this class. - Returns the value of the ``type`` class variable if present, - otherwise returns a snake_cased version of the class name. + The type identifier is used in serialized data to enable proper deserialization. + It follows a priority order to determine the identifier: + + 1. If ``value`` contains a 'type' field, return that value (for ``from_dict``) + 2. If the class has a ``type`` attribute, use that value (instance-level) + 3. If the class has a ``TYPE`` attribute, use that value (class-level constant) + 4. Otherwise, convert the class name to snake_case as fallback + + Args: + value: Optional mapping containing serialized data that may have a 'type' field. Returns: - Type identifier string. + Type identifier string used for serialization and dependency injection mapping. """ + # for from_dict + if value and (type_ := value.get("type")) and isinstance(type_, str): + return type_ # type:ignore[no-any-return] + # for todict when defined per instance if (type_ := getattr(cls, "type", None)) and isinstance(type_, str): return type_ # type:ignore[no-any-return] - + # for both when defined on class. + if (type_ := getattr(cls, "TYPE", None)) and isinstance(type_, str): + return type_ # type:ignore[no-any-return] + # Fallback and default # Convert class name to snake_case return _CAMEL_TO_SNAKE_PATTERN.sub("_", cls.__name__).lower() diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index eba8307fb2..22f3dee18a 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -4,7 +4,7 @@ import asyncio import inspect import json import sys -from collections.abc import AsyncIterable, Awaitable, Callable, Collection, MutableMapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Collection, Mapping, MutableMapping, Sequence from functools import wraps from time import perf_counter, time_ns from typing import ( @@ -17,6 +17,7 @@ from typing import ( Literal, Protocol, TypeVar, + cast, get_args, get_origin, runtime_checkable, @@ -24,6 +25,7 @@ from typing import ( from opentelemetry.metrics import Histogram from pydantic import AnyUrl, BaseModel, Field, ValidationError, create_model +from pydantic.fields import FieldInfo from ._logging import get_logger from ._serialization import SerializationMixin @@ -49,9 +51,15 @@ if TYPE_CHECKING: ) if sys.version_info >= (3, 12): - from typing import TypedDict # pragma: no cover + from typing import ( + TypedDict, # pragma: no cover + override, # type: ignore # pragma: no cover + ) else: - from typing_extensions import TypedDict # pragma: no cover + from typing_extensions import ( + TypedDict, # pragma: no cover + override, # type: ignore[import] # pragma: no cover + ) if sys.version_info >= (3, 11): from typing import overload # pragma: no cover @@ -540,6 +548,9 @@ def _default_histogram() -> Histogram: ) +TClass = TypeVar("TClass", bound="SerializationMixin") + + class AIFunction(BaseTool, Generic[ArgsT, ReturnT]): """A tool that wraps a Python function to make it callable by AI models. @@ -593,7 +604,7 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]): approval_mode: Literal["always_require", "never_require"] | None = None, additional_properties: dict[str, Any] | None = None, func: Callable[..., Awaitable[ReturnT] | ReturnT], - input_model: type[ArgsT], + input_model: type[ArgsT] | Mapping[str, Any] | None = None, **kwargs: Any, ) -> None: """Initialize the AIFunction. @@ -606,6 +617,8 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]): additional_properties: Additional properties to set on the function. func: The function to wrap. input_model: The Pydantic model that defines the input parameters for the function. + This can also be a JSON schema dictionary. + If not provided, it will be inferred from the function signature. **kwargs: Additional keyword arguments. """ super().__init__( @@ -615,9 +628,19 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]): **kwargs, ) self.func = func - self.input_model = input_model + self.input_model = self._resolve_input_model(input_model) self.approval_mode = approval_mode or "never_require" self._invocation_duration_histogram = _default_histogram() + self.type: Literal["ai_function"] = "ai_function" + + def _resolve_input_model(self, input_model: type[ArgsT] | Mapping[str, Any] | None) -> type[ArgsT]: + if input_model: + if inspect.isclass(input_model) and issubclass(input_model, BaseModel): + return input_model + if isinstance(input_model, Mapping): + return cast(type[ArgsT], _create_model_from_json_schema(self.name, input_model)) + raise TypeError("input_model must be a Pydantic BaseModel subclass or a JSON schema dict.") + return cast(type[ArgsT], _create_input_model_from_func(self.func, self.name)) def __call__(self, *args: Any, **kwargs: Any) -> ReturnT | Awaitable[ReturnT]: """Call the wrapped function with the provided arguments.""" @@ -725,6 +748,14 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]): }, } + @override + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + as_dict = super().to_dict(exclude=exclude, exclude_none=exclude_none) + if (exclude and "input_model" in exclude) or not self.input_model: + return as_dict + as_dict["input_model"] = self.input_model.model_json_schema() + return as_dict + def _tools_to_dict( tools: ( @@ -802,6 +833,73 @@ def _parse_annotation(annotation: Any) -> Any: return annotation +def _create_input_model_from_func(func: Callable[..., Any], tool_name: str) -> type[BaseModel]: + """Create a Pydantic model from a function's signature.""" + sig = inspect.signature(func) + fields = { + pname: ( + _parse_annotation(param.annotation) if param.annotation is not inspect.Parameter.empty else str, + param.default if param.default is not inspect.Parameter.empty else ..., + ) + for pname, param in sig.parameters.items() + if pname not in {"self", "cls"} + } + return create_model(f"{tool_name}_input", **fields) # type: ignore[call-overload, no-any-return] + + +# Map JSON Schema types to Pydantic types +TYPE_MAPPING = { + "string": str, + "integer": int, + "number": float, + "boolean": bool, + "array": list, + "object": dict, + "null": type(None), +} + + +def _create_model_from_json_schema(tool_name: str, schema_json: Mapping[str, Any]) -> type[BaseModel]: + """Creates a Pydantic model from a given JSON Schema. + + Args: + tool_name: The name of the model to be created. + schema_json: The JSON Schema definition. + + Returns: + The dynamically created Pydantic model class. + """ + # Validate that 'properties' exists and is a dict + if "properties" not in schema_json or not isinstance(schema_json["properties"], dict): + raise ValueError( + f"JSON schema for tool '{tool_name}' must contain a 'properties' key of type dict. " + f"Got: {schema_json.get('properties', None)}" + ) + # Extract field definitions with type annotations + field_definitions: dict[str, tuple[type, FieldInfo]] = {} + for field_name, field_schema in schema_json["properties"].items(): + field_args: dict[str, Any] = {} + if (field_description := field_schema.get("description", None)) is not None: + field_args["description"] = field_description + if (field_default := field_schema.get("default", None)) is not None: + field_args["default"] = field_default + field_type = field_schema.get("type", None) + if field_type is None: + raise ValueError( + f"Missing 'type' for field '{field_name}' in JSON schema. " + f"Got: {field_schema}, Supported types: {list(TYPE_MAPPING.keys())}" + ) + python_type = TYPE_MAPPING.get(field_type) + if python_type is None: + raise ValueError( + f"Unsupported type '{field_type}' for field '{field_name}' in JSON schema. " + f"Got: {field_schema}, Supported types: {list(TYPE_MAPPING.keys())}" + ) + field_definitions[field_name] = (python_type, Field(**field_args)) + + return create_model(f"{tool_name}_input", **field_definitions) # type: ignore[call-overload, no-any-return] + + @overload def ai_function( func: Callable[..., ReturnT | Awaitable[ReturnT]], @@ -895,26 +993,12 @@ def ai_function( def wrapper(f: Callable[..., ReturnT | Awaitable[ReturnT]]) -> AIFunction[Any, ReturnT]: tool_name: str = name or getattr(f, "__name__", "unknown_function") # type: ignore[assignment] tool_desc: str = description or (f.__doc__ or "") - sig = inspect.signature(f) - fields = { - pname: ( - _parse_annotation(param.annotation) if param.annotation is not inspect.Parameter.empty else str, - param.default if param.default is not inspect.Parameter.empty else ..., - ) - for pname, param in sig.parameters.items() - if pname not in {"self", "cls"} - } - input_model: Any = create_model(f"{tool_name}_input", **fields) # type: ignore[call-overload] - if not issubclass(input_model, BaseModel): - raise TypeError(f"Input model for {tool_name} must be a subclass of BaseModel, got {input_model}") - return AIFunction[Any, ReturnT]( name=tool_name, description=tool_desc, approval_mode=approval_mode, additional_properties=additional_properties or {}, func=f, - input_model=input_model, ) return wrapper(func) diff --git a/python/packages/core/tests/core/test_serializable_mixin.py b/python/packages/core/tests/core/test_serializable_mixin.py index 1594cb7c5e..0472f881cf 100644 --- a/python/packages/core/tests/core/test_serializable_mixin.py +++ b/python/packages/core/tests/core/test_serializable_mixin.py @@ -45,7 +45,7 @@ class TestSerializationMixin: with caplog.at_level(logging.DEBUG): obj = TestClass.from_dict( {"type": "test_class", "value": "test"}, - dependencies={"test_class.client": mock_client}, + dependencies={"test_class": {"client": mock_client}}, ) assert obj.value == "test" @@ -68,7 +68,7 @@ class TestSerializationMixin: with caplog.at_level(logging.DEBUG): obj = TestClass.from_dict( {"type": "test_class", "value": "test"}, - dependencies={"test_class.other": mock_other}, + dependencies={"test_class": {"other": mock_other}}, ) assert obj.value == "test" @@ -105,9 +105,11 @@ class TestSerializationMixin: obj = TestClass.from_dict( {"type": "test_class", "value": "test"}, dependencies={ - "test_class.client": mock_client, - "test_class.logger": mock_logger, - "test_class.other": mock_other, + "test_class": { + "client": mock_client, + "logger": mock_logger, + "other": mock_other, + } }, ) @@ -136,7 +138,7 @@ class TestSerializationMixin: with caplog.at_level(logging.DEBUG): obj = TestClass.from_dict( {"type": "test_class", "value": "test"}, - dependencies={"test_class.client": mock_client}, + dependencies={"test_class": {"client": mock_client}}, ) assert obj.value == "test" @@ -184,7 +186,7 @@ class TestSerializationMixin: assert "client" not in data # Excluded from serialization # Deserialize with dependency injection - restored = TestClass.from_dict(data, dependencies={"test_class.client": mock_client}) + restored = TestClass.from_dict(data, dependencies={"test_class": {"client": mock_client}}) assert restored.value == "test" assert restored.number == 42 assert restored.client == mock_client diff --git a/python/packages/core/tests/core/test_tools.py b/python/packages/core/tests/core/test_tools.py index ba70477e8c..e2cf6b8d3d 100644 --- a/python/packages/core/tests/core/test_tools.py +++ b/python/packages/core/tests/core/test_tools.py @@ -299,6 +299,48 @@ async def test_ai_function_invoke_invalid_pydantic_args(): await invalid_args_test.invoke(arguments=wrong_args) +def test_ai_function_serialization(): + """Test AIFunction serialization and deserialization.""" + + def serialize_test(x: int, y: int) -> int: + """A function for testing serialization.""" + return x - y + + serialize_test_ai_function = ai_function(name="serialize_test", description="A test tool for serialization")( + serialize_test + ) + + # Serialize to dict + tool_dict = serialize_test_ai_function.to_dict() + assert tool_dict["type"] == "ai_function" + assert tool_dict["name"] == "serialize_test" + assert tool_dict["description"] == "A test tool for serialization" + assert tool_dict["input_model"] == { + "properties": {"x": {"title": "X", "type": "integer"}, "y": {"title": "Y", "type": "integer"}}, + "required": ["x", "y"], + "title": "serialize_test_input", + "type": "object", + } + + # Deserialize from dict + restored_tool = AIFunction.from_dict(tool_dict, dependencies={"ai_function": {"func": serialize_test}}) + assert isinstance(restored_tool, AIFunction) + assert restored_tool.name == "serialize_test" + assert restored_tool.description == "A test tool for serialization" + assert restored_tool.parameters() == serialize_test_ai_function.parameters() + assert restored_tool(10, 4) == 6 + + # Deserialize from dict with instance name + restored_tool_2 = AIFunction.from_dict( + tool_dict, dependencies={"ai_function": {"name:serialize_test": {"func": serialize_test}}} + ) + assert isinstance(restored_tool_2, AIFunction) + assert restored_tool_2.name == "serialize_test" + assert restored_tool_2.description == "A test tool for serialization" + assert restored_tool_2.parameters() == serialize_test_ai_function.parameters() + assert restored_tool_2(10, 4) == 6 + + # region HostedCodeInterpreterTool and _parse_inputs diff --git a/python/samples/getting_started/tools/tool_with_injected_func.py b/python/samples/getting_started/tools/tool_with_injected_func.py new file mode 100644 index 0000000000..bff59f31d2 --- /dev/null +++ b/python/samples/getting_started/tools/tool_with_injected_func.py @@ -0,0 +1,68 @@ +# Copyright (c) Microsoft. All rights reserved. +# type: ignore +""" +AIFunction Tool with Dependency Injection Example + +This example demonstrates how to create an AIFunction tool using the agent framework's +dependency injection system. Instead of providing the function at initialization time, +the actual callable function is injected during deserialization from a dictionary definition. + +Note: + The serialization and deserialization feature used in this example is currently + in active development. The API may change in future versions as we continue + to improve and extend its functionality. Please refer to the latest documentation + for any updates to the dependency injection patterns. + +Usage: + Run this script to see how an AIFunction tool can be created from a dictionary + definition with the function injected at runtime. The agent will use this tool + to perform arithmetic operations. +""" + +import asyncio + +from agent_framework import AIFunction +from agent_framework.openai import OpenAIResponsesClient + +definition = { + "type": "ai_function", + "name": "add_numbers", + "description": "Add two numbers together.", + "input_model": { + "properties": { + "a": {"description": "The first number", "type": "integer"}, + "b": {"description": "The second number", "type": "integer"}, + }, + "required": ["a", "b"], + "title": "func_input", + "type": "object", + }, +} + + +async def main() -> None: + """Main function demonstrating creating a tool with an injected function.""" + + def func(a, b) -> int: + """Add two numbers together.""" + return a + b + + # Create the AIFunction tool using dependency injection + # The 'definition' dictionary contains the serialized tool configuration, + # while the actual function implementation is provided via dependencies. + # + # Dependency structure: {"ai_function": {"name:add_numbers": {"func": func}}} + # - "ai_function": matches the tool type identifier + # - "name:add_numbers": instance-specific injection targeting tools with name="add_numbers" + # - "func": the parameter name that will receive the injected function + tool = AIFunction.from_dict(definition, dependencies={"ai_function": {"name:add_numbers": {"func": func}}}) + + agent = OpenAIResponsesClient().create_agent( + name="FunctionToolAgent", instructions="You are a helpful assistant.", tools=tool + ) + response = await agent.run("What is 5 + 3?") + print(f"Response: {response.text}") + + +if __name__ == "__main__": + asyncio.run(main())