feat: Model Client and associated Content Types (#53)

* feat: ModelClient and content types

* refactor: Pythonify ChatResponseFormat and ChatRole

* feat: Add guardrail interfaces

* refactor: Remove CancellationToken

* feat: Solidify the Usage APIs

* Adds well-known keys for additional_counts, and guidance for how to avoid collisions between providers
* Implement sum-aggregation for usage

* refactor: Move AITool out of model_client

* refactor: Copy editing

* fix: CI checks (pyupgrade, ruff, etc.)

* ci: Fix pre-commit to use pyright in  uv venv

The existing pyright precommit hook inside of python-pyright is no longer being maintained by the owner (see  https://github.com/RobertCraigie/pyright-python/issues/265)

The fix is to define the hook ourselves, relying on `uv run` to drive it. In order for that to work right we need to use the "system" language to break out of the sandbox.

* fix: Pyright error fixes

* docs: Update models and types design docs

* Python: Refinement of content types and model client  (#112)

* refinement of structure and buildup
with ports from semantigen

* refined the data and uri contents

* refined chat response and updates

* moved things and added tests

* moved out of src folder

* fixed imports and tests

* small tweaks

* missing build system

* upgrade

* add mypy

* fixed typing for types

* fix tests

* fixed tool

* disable json checks on vscode

* remove print

---------

Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>
Co-authored-by: eavanvalkenburg <github@vanvalkenburg.eu>

python/agent_framework/__init__.py

@@ -10,6 +10,30 @@ except importlib.metadata.PackageNotFoundError:
 
 _IMPORTS = {
     "get_logger": "._logging",
+    "AITool": "._tools",
+    "ai_function": "._tools",
+    "AIContent": "._types",
+    "AIContents": "._types",
+    "TextContent": "._types",
+    "TextReasoningContent": "._types",
+    "DataContent": "._types",
+    "UriContent": "._types",
+    "UsageContent": "._types",
+    "UsageDetails": "._types",
+    "FunctionCallContent": "._types",
+    "FunctionResultContent": "._types",
+    "ChatFinishReason": "._types",
+    "ChatMessage": "._types",
+    "ChatResponse": "._types",
+    "StructuredResponse": "._types",
+    "ChatResponseUpdate": "._types",
+    "ChatRole": "._types",
+    "ErrorContent": "._types",
+    "ModelClient": "._types",
+    "ChatOptions": "._types",
+    "ChatToolMode": "._types",
+    "InputGuardrail": ".guard_rails",
+    "OutputGuardrail": ".guard_rails",
 }
 
 

python/agent_framework/__init__.pyi

@@ -1,11 +1,57 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import importlib.metadata
-
-try:
-    __version__ = importlib.metadata.version(__name__)
-except importlib.metadata.PackageNotFoundError:
-    __version__ = "0.0.0"  # Fallback for development mode
+from . import __version__  # type: ignore[attr-defined]
 from ._logging import get_logger
+from ._tools import AITool, ai_function
+from ._types import (
+    AIContent,
+    AIContents,
+    ChatFinishReason,
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    ChatRole,
+    ChatToolMode,
+    DataContent,
+    ErrorContent,
+    FunctionCallContent,
+    FunctionResultContent,
+    ModelClient,
+    StructuredResponse,
+    TextContent,
+    TextReasoningContent,
+    UriContent,
+    UsageContent,
+    UsageDetails,
+)
+from .guard_rails import InputGuardrail, OutputGuardrail
 
-__all__ = ["__version__", "get_logger"]
+__all__ = [
+    "AIContent",
+    "AIContents",
+    "AITool",
+    "ChatFinishReason",
+    "ChatMessage",
+    "ChatOptions",
+    "ChatResponse",
+    "ChatResponseUpdate",
+    "ChatRole",
+    "ChatToolMode",
+    "DataContent",
+    "ErrorContent",
+    "FunctionCallContent",
+    "FunctionResultContent",
+    "InputGuardrail",
+    "ModelClient",
+    "OutputGuardrail",
+    "StructuredResponse",
+    "TextContent",
+    "TextReasoningContent",
+    "UriContent",
+    "UsageContent",
+    "UsageDetails",
+    "__version__",
+    "ai_function",
+    "get_logger",
+]

python/agent_framework/_cancellation_token.py

@@ -0,0 +1,50 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import threading
+from asyncio import Future
+from collections.abc import Callable
+from typing import Any
+
+
+# from https://github.com/microsoft/autogen/blob/main/python/packages/autogen-core/src/autogen_core/_cancellation_token.py
+class CancellationToken:
+    """A token used to cancel pending async calls."""
+
+    def __init__(self) -> None:
+        self._cancelled: bool = False
+        self._lock: threading.Lock = threading.Lock()
+        self._callbacks: list[Callable[[], None]] = []
+
+    def cancel(self) -> None:
+        """Cancel pending async calls linked to this cancellation token."""
+        with self._lock:
+            if not self._cancelled:
+                self._cancelled = True
+                for callback in self._callbacks:
+                    callback()
+
+    def is_cancelled(self) -> bool:
+        """Check if the CancellationToken has been used."""
+        with self._lock:
+            return self._cancelled
+
+    def add_callback(self, callback: Callable[[], None]) -> None:
+        """Attach a callback that will be called when cancel is invoked."""
+        with self._lock:
+            if self._cancelled:
+                callback()
+            else:
+                self._callbacks.append(callback)
+
+    def link_future(self, future: Future[Any]) -> Future[Any]:
+        """Link a pending async call to a token to allow its cancellation."""
+        with self._lock:
+            if self._cancelled:
+                future.cancel()
+            else:
+
+                def _cancel() -> None:
+                    future.cancel()
+
+                self._callbacks.append(_cancel)
+        return future

python/agent_framework/_pydantic.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from pydantic import BaseModel, ConfigDict
+
+
+class AFBaseModel(BaseModel):
+    """Base class for all pydantic models in the Agent Framework."""
+
+    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True, validate_assignment=True)

python/agent_framework/_tools.py

@@ -0,0 +1,122 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import functools
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import Any, Generic, Protocol, TypeVar, runtime_checkable
+
+from pydantic import BaseModel, create_model
+
+
+@runtime_checkable
+class AITool(Protocol):
+    """Represents a tool that can be specified to an AI service."""
+
+    name: str
+    """The name of the tool."""
+    description: str | None = None
+    """A description of the tool, suitable for use in describing the purpose to a model."""
+    additional_properties: dict[str, Any] | None = None
+    """Additional properties associated with the tool."""
+
+    def __str__(self) -> str:
+        """Return a string representation of the tool."""
+        ...
+
+
+ArgsT = TypeVar("ArgsT", bound=BaseModel)
+ReturnT = TypeVar("ReturnT")
+
+
+class AIFunction(Generic[ArgsT, ReturnT]):
+    """A tool that represents a function that can be called by an AI service."""
+
+    def __init__(
+        self,
+        func: Callable[..., Awaitable[ReturnT] | ReturnT],
+        name: str,
+        description: str,
+        input_model: type[ArgsT],
+        **kwargs: Any,
+    ):
+        """Initialize a FunctionTool.
+
+        Args:
+            func: The function to wrap.
+            name: The name of the tool.
+            description: A description of the tool.
+            input_model: A Pydantic model that defines the input parameters for the function.
+            **kwargs: Additional properties to set on the tool.
+                stored in additional_properties.
+        """
+        self.name = name
+        self.description = description
+        self.input_model = input_model
+        self.additional_properties: dict[str, Any] | None = kwargs
+        self._func = func
+
+    def model_json_schema(self) -> dict[str, Any]:
+        """Return the JSON schema of the input model."""
+        return self.input_model.model_json_schema()
+
+    def __call__(self, *args: Any, **kwargs: Any) -> ReturnT | Awaitable[ReturnT]:
+        """Call the wrapped function with the provided arguments."""
+        return self._func(*args, **kwargs)
+
+    async def invoke(
+        self,
+        *,
+        arguments: ArgsT | None = None,
+        **kwargs: Any,
+    ) -> ReturnT:
+        """Run the AI function with the provided arguments as a Pydantic model.
+
+        Args:
+            arguments: A Pydantic model instance containing the arguments for the function.
+            kwargs: keyword arguments to pass to the function, will not be used if `args` is provided.
+        """
+        if arguments is not None:
+            if not isinstance(arguments, self.input_model):
+                raise TypeError(f"Expected {self.input_model.__name__}, got {type(arguments).__name__}")
+            kwargs = arguments.model_dump(exclude_none=True)
+        res = self.__call__(**kwargs)
+        if inspect.isawaitable(res):
+            return await res
+        return res
+
+
+def ai_function(
+    func: Callable[..., ReturnT | Awaitable[ReturnT]] | None = None,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    additional_properties: dict[str, Any] | None = None,
+) -> AIFunction[Any, ReturnT] | Callable[[Callable[..., ReturnT | Awaitable[ReturnT]]], AIFunction[Any, ReturnT]]:
+    """Create a AIFunction from a function and return the callable tool object."""
+
+    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: (
+                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 = create_model(f"{tool_name}_input", **fields)  # type: ignore[call-overload]
+
+        return functools.update_wrapper(  # type: ignore[return-value]
+            AIFunction[Any, ReturnT](
+                func=f,
+                name=tool_name,
+                description=tool_desc,
+                input_model=input_model,
+                **(additional_properties if additional_properties is not None else {}),
+            ),
+            f,
+        )
+
+    return wrapper(func) if func else wrapper

python/agent_framework/guard_rails.py

@@ -0,0 +1,25 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from typing import Generic, Protocol, TypeVar, runtime_checkable
+
+TInput = TypeVar("TInput")
+TResponse = TypeVar("TResponse")
+
+
+@runtime_checkable
+class InputGuardrail(Protocol, Generic[TInput]):
+    """A protocol for input guardrails that can validate and transform input messages."""
+
+    def __call__(self, message: TInput) -> TInput:
+        """Validate and possibly transform the input message."""
+        ...
+
+
+@runtime_checkable
+class OutputGuardrail(Protocol, Generic[TResponse]):
+    """A protocol for output guardrails that can validate and transform output messages."""
+
+    def __call__(self, message: TResponse) -> TResponse:
+        """Validate and possibly transform the output message."""
+        ...