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 <github@vanvalkenburg.eu>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Eric Zhu
2025-10-13 13:34:27 -07:00
committed by GitHub
Unverified
parent fc12ab9fed
commit 9148392d00
5 changed files with 618 additions and 128 deletions
@@ -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"(?<!^)(?=[A-Z])")
class SerializationProtocol(Protocol):
"""Protocol for objects that support serialization and deserialization.
This protocol defines the interface for objects that can be converted to and from dictionaries.
This protocol defines the interface that classes must implement to be compatible
with the agent framework's serialization system. Any class implementing both
``to_dict()`` and ``from_dict()`` methods will automatically satisfy this protocol
and can be used seamlessly with other serializable components.
The protocol enables type safety and duck typing for serializable objects,
ensuring consistent behavior across the framework.
Examples:
The framework's ``ChatMessage`` class demonstrates the protocol in action:
.. code-block:: python
from agent_framework import SerializationProtocol
from agent_framework import ChatMessage
from agent_framework._serialization import SerializationProtocol
class MySerializable:
def __init__(self, value: str):
self.value = value
# ChatMessage implements SerializationProtocol via SerializationMixin
user_msg = ChatMessage(role="user", text="What's the weather like today?")
def to_dict(self, **kwargs):
return {"value": self.value}
# Serialize to dictionary - automatic type identification and nested serialization
msg_dict = user_msg.to_dict()
# Result: {
# "type": "chat_message",
# "role": {"type": "role", "value": "user"},
# "contents": [{"type": "text_content", "text": "What's the weather like today?"}],
# "message_id": "...",
# "additional_properties": {}
# }
@classmethod
def from_dict(cls, value, **kwargs):
return cls(value["value"])
# Deserialize back to ChatMessage instance - automatic type reconstruction
restored_msg = ChatMessage.from_dict(msg_dict)
print(restored_msg.text) # "What's the weather like today?"
print(restored_msg.role.value) # "user"
# Verify protocol compliance (useful for type checking and validation)
assert isinstance(user_msg, SerializationProtocol)
assert isinstance(restored_msg, SerializationProtocol)
# Verify it implements the protocol
assert isinstance(MySerializable("test"), SerializationProtocol)
The protocol is also implemented by simpler classes like ``UsageDetails``:
.. code-block:: python
from agent_framework import UsageDetails
# Create usage tracking instance
usage = UsageDetails(input_token_count=150, output_token_count=75, total_token_count=225)
# Seamless serialization with type preservation
usage_dict = usage.to_dict()
restored_usage = UsageDetails.from_dict(usage_dict)
# Both satisfy the SerializationProtocol
assert isinstance(usage, SerializationProtocol)
assert restored_usage.total_token_count == 225
The protocol ensures consistent serialization behavior across all framework components,
enabling reliable data persistence, API communication, and object reconstruction
throughout the agent framework ecosystem.
"""
def to_dict(self, **kwargs: Any) -> 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:
- "<type>.<parameter>"
- "<type>.<dict-parameter>.<key>"
where <type> 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),
<parameter> is the name of the parameter in the __init__ method,
<dict-parameter> is the name of a parameter that is a dict,
and <key> 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 ``"<type>.<parameter>"`` or ``"<type>.<dict-parameter>.<key>"``.
dependencies: A nested dictionary mapping type identifiers to their injectable dependencies.
The structure varies based on injection pattern:
- **Simple injection**: ``{"<type>": {"<parameter>": value}}``
- **Dict parameter injection**: ``{"<type>": {"<dict-parameter>": {"<key>": value}}}``
- **Instance-specific injection**: ``{"<type>": {"<field>:<value>": {"<parameter>": 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: <type>.<parameter>
kwargs[param_name] = dep_value
elif len(parts) == 3:
# Dict parameter: <type>.<dict-parameter>.<key>
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 ``"<type>.<parameter>"`` or ``"<type>.<dict-parameter>.<key>"``.
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()
+103 -19
View File
@@ -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)
@@ -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
@@ -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
@@ -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())