From fea2657b8485552e3d40ba8fb7521da9316040a3 Mon Sep 17 00:00:00 2001 From: Jack Gerrits Date: Mon, 16 Dec 2024 17:23:14 -0500 Subject: [PATCH 1/5] Add wip component impl --- .../autogen-core/docs/src/reference/index.md | 2 +- .../src/reference/python/autogen_core.rst | 5 +- .../framework/component-config.ipynb | 62 +++++ .../core-user-guide/framework/index.md | 1 + .../autogen-core/src/autogen_core/__init__.py | 14 ++ .../src/autogen_core/_component_config.py | 226 ++++++++++++++++++ .../src/autogen_core/models/_model_client.py | 15 +- python/packages/autogen-core/test.py | 19 ++ .../tests/test_component_config.py | 102 ++++++++ .../models/openai/_azure_token_provider.py | 50 ++++ .../models/openai/_openai_client.py | 30 ++- .../models/openai/config/__init__.py | 4 +- .../src/autogen_test_utils/__init__.py | 43 +++- 13 files changed, 562 insertions(+), 11 deletions(-) create mode 100644 python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb create mode 100644 python/packages/autogen-core/src/autogen_core/_component_config.py create mode 100644 python/packages/autogen-core/test.py create mode 100644 python/packages/autogen-core/tests/test_component_config.py create mode 100644 python/packages/autogen-ext/src/autogen_ext/models/openai/_azure_token_provider.py diff --git a/python/packages/autogen-core/docs/src/reference/index.md b/python/packages/autogen-core/docs/src/reference/index.md index 905521debf3a..f6fadd12960f 100644 --- a/python/packages/autogen-core/docs/src/reference/index.md +++ b/python/packages/autogen-core/docs/src/reference/index.md @@ -25,7 +25,7 @@ python/autogen_agentchat.state :hidden: :caption: AutoGen Core -python/autogen_core +autogen_core python/autogen_core.code_executor python/autogen_core.models python/autogen_core.model_context diff --git a/python/packages/autogen-core/docs/src/reference/python/autogen_core.rst b/python/packages/autogen-core/docs/src/reference/python/autogen_core.rst index ee23badec6ee..c3f7bded900d 100644 --- a/python/packages/autogen-core/docs/src/reference/python/autogen_core.rst +++ b/python/packages/autogen-core/docs/src/reference/python/autogen_core.rst @@ -1,8 +1,11 @@ autogen\_core ============= - .. automodule:: autogen_core :members: :undoc-members: :show-inheritance: + :member-order: bysource + +.. For some reason the following isn't picked up by automodule above +.. autoclass:: autogen_core.ComponentType \ No newline at end of file diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb b/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb new file mode 100644 index 000000000000..0e7dfed63263 --- /dev/null +++ b/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb @@ -0,0 +1,62 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Component config\n", + "\n", + "AutoGen components are able to be declaratively configured in a generic fashion. This is to support configuration based experiences, such as AutoGen studio, but it is also useful for many other scenarios.\n", + "\n", + "The system that provides this is called \"component configuration\". In AutoGen, a component is simply something that can be created from a config object and itself can be dumped to a config object. In this way, you can define a component in code and then get the config object from it.\n", + "\n", + "This system is generic and allows for components defined outside of AutoGen itself (such as extensions) to be configured in the same way.\n", + "\n", + "## How does this differ from state?\n", + "\n", + "This is a very important point to clarify. When we talk about serializing an object, we must include *all* data that makes that object itself. Including things like message history etc. When deserializing from serialized state, you must get back the *exact* same object. This is not the case with component configuration.\n", + "\n", + "Component configuration should be thought of as the blueprint for an object, and can be stamped out many times to create many instances of the same configured object.\n", + "\n", + "## Usage\n", + "\n", + "If you have a component in Python ad want to get the config for it, simply call {py:meth}`~autogen_core.ComponentConfig.dump_component` on it. The resulting object can be passed back into {py:meth}`~autogen_core.ComponentLoader.load_component` to get the component back.\n", + "\n", + "## Creating a component\n", + "\n", + "To add component functionality to a given class:\n", + "\n", + "1. Add a call to {py:meth}`~autogen_core.Component` in the class inheritance list.\n", + "2. Implment the {py:meth}`~autogen_core.ComponentConfigImpl._to_config` and {py:meth}`~autogen_core.ComponentConfigImpl._from_config` methods\n", + "\n", + "For example:\n", + "\n", + "```python\n", + "from pydantic import BaseModel\n", + "from autogen_core import Component\n", + "\n", + "class Config(BaseModel):\n", + " value: str\n", + "\n", + "class MyComponent(Component(\"custom\", Config)):\n", + " def __init__(self, value: str):\n", + " self.value = value\n", + "\n", + " def _to_config(self) -> Config:\n", + " return Config(value=self.value)\n", + "\n", + " @classmethod\n", + " def _from_config(cls, config: Config) -> MyComponent:\n", + " return cls(value=config.value)\n", + "```\n" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/index.md b/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/index.md index 21f9cd00d8f2..be8e9058fed3 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/index.md +++ b/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/index.md @@ -27,4 +27,5 @@ logging telemetry command-line-code-executors distributed-agent-runtime +component-config ``` diff --git a/python/packages/autogen-core/src/autogen_core/__init__.py b/python/packages/autogen-core/src/autogen_core/__init__.py index 0f085d29bdfe..1254f294600d 100644 --- a/python/packages/autogen-core/src/autogen_core/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/__init__.py @@ -12,6 +12,14 @@ from ._base_agent import BaseAgent from ._cancellation_token import CancellationToken from ._closure_agent import ClosureAgent, ClosureContext +from ._component_config import ( + Component, + ComponentConfig, + ComponentConfigImpl, + ComponentLoader, + ComponentModel, + ComponentType, +) from ._constants import ( EVENT_LOGGER_NAME as EVENT_LOGGER_NAME_ALIAS, ) @@ -99,4 +107,10 @@ "ROOT_LOGGER_NAME", "EVENT_LOGGER_NAME", "TRACE_LOGGER_NAME", + "Component", + "ComponentConfig", + "ComponentLoader", + "ComponentConfigImpl", + "ComponentModel", + "ComponentType", ] diff --git a/python/packages/autogen-core/src/autogen_core/_component_config.py b/python/packages/autogen-core/src/autogen_core/_component_config.py new file mode 100644 index 000000000000..e3ad9c92bd25 --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/_component_config.py @@ -0,0 +1,226 @@ +from __future__ import annotations +import importlib +import types +from typing import Generic, Literal, Protocol, Type, cast, overload, runtime_checkable + +from pydantic import BaseModel +from typing_extensions import Self, TypeVar + + +ComponentType = Literal["model", "agent", "tool", "termination", "token_provider"] | str +ConfigT = TypeVar("ConfigT", bound=BaseModel) + +T = TypeVar("T", bound=BaseModel, covariant=True) + +class ComponentModel(BaseModel): + """Model class for a component. Contains all information required to instantiate a component.""" + + provider: str + """Describes how the component can be instantiated.""" + component_type: ComponentType + """Logical type of the component.""" + version: int + """Version of the component specification.""" + description: str | None + """Description of the component.""" + + config: BaseModel + """The config field is passed to a given class's implmentation of :py:meth:`autogen_core.ComponentConfigImpl._from_config` to create a new instance of the component class.""" + + +def _type_to_provider_str(t: type) -> str: + return f"{t.__module__}.{t.__qualname__}" + +@runtime_checkable +class ComponentConfigImpl(Protocol[ConfigT]): + """The two methods a class must implement to be a component. + + Args: + Protocol (ConfigT): Type which derives from :py:class:`pydantic.BaseModel`. + """ + + def _to_config(self) -> ConfigT: + """Dump the configuration that would be requite to create a new instance of a component matching the configuration of this instance. + + Returns: + T: The configuration of the component. + + :meta public: + """ + ... + + @classmethod + def _from_config(cls, config: ConfigT) -> Self: + """Create a new instance of the component from a configuration object. + + Args: + config (T): The configuration object. + + Returns: + Self: The new instance of the component. + + :meta public: + """ + ... + + +ExpectedType = TypeVar("ExpectedType") +class ComponentLoader(): + + @overload + @classmethod + def load_component(cls, model: ComponentModel, expected: None = None) -> Self: + ... + + @overload + @classmethod + def load_component(cls, model: ComponentModel, expected: Type[ExpectedType]) -> ExpectedType: + ... + + @classmethod + def load_component(cls, model: ComponentModel, expected: Type[ExpectedType] | None = None) -> Self | ExpectedType: + """Load a component from a model. Intended to be used with the return type of :py:meth:`autogen_core.ComponentConfig.dump_component`. + + Example: + + .. code-block:: python + + from autogen_core import ComponentModel + from autogen_core.models import ChatCompletionClient + + component: ComponentModel = ... # type: ignore + + model_client = ChatCompletionClient.load_component(component) + + Args: + model (ComponentModel): The model to load the component from. + + Returns: + Self: The loaded component. + + Args: + model (ComponentModel): _description_ + expected (Type[ExpectedType] | None, optional): Explicit type only if used directly on ComponentLoader. Defaults to None. + + Raises: + ValueError: If the provider string is invalid. + TypeError: Provider is not a subclass of ComponentConfigImpl, or the expected type does not match. + + Returns: + Self | ExpectedType: The loaded component. + """ + + # Use global and add further type checks + + output = model.provider.rsplit(".", maxsplit=1) + if len(output) != 2: + raise ValueError("Invalid") + + module_path, class_name = output + module = importlib.import_module(module_path) + component_class = cast(ComponentConfigImpl[BaseModel], module.__getattribute__(class_name)) + + if not isinstance(component_class, ComponentConfigImpl): + raise TypeError("Invalid component class") + + # We're allowed to use the private method here + instance = component_class._from_config(model.config) # type: ignore + + if expected is None and not isinstance(instance, cls): + raise TypeError("Expected type does not match") + elif expected is None: + return cast(Self, instance) + elif not isinstance(instance, expected): + raise TypeError("Expected type does not match") + else: + return cast(ExpectedType, instance) + + +class ComponentConfig(ComponentConfigImpl[ConfigT], ComponentLoader, Generic[ConfigT]): + _config_schema: Type[ConfigT] + _component_type: ComponentType + _description: str | None + + def dump_component(self) -> ComponentModel: + """Dump the component to a model that can be loaded back in. + + Raises: + TypeError: If the component is a local class. + + Returns: + ComponentModel: The model representing the component. + """ + provider = _type_to_provider_str(self.__class__) + + if "" in provider: + raise TypeError("Cannot dump component with local class") + + return ComponentModel( + provider=provider, + component_type=self._component_type, + version=1, + description=self._description, + config=self._to_config(), + ) + + @classmethod + def component_config_schema(cls) -> Type[ConfigT]: + """Get the configuration schema for the component. + + Returns: + Type[ConfigT]: The configuration schema for the component. + """ + return cls._config_schema + + @classmethod + def component_type(cls) -> ComponentType: + """Logical type of the component. + + Returns: + ComponentType: The logical type of the component. + """ + return cls._component_type + + +def Component( + component_type: ComponentType, config_schema: type[T], description: str | None = None +) -> type[ComponentConfig[T]]: + """This enables easy creation of Component classes. It provides a type to inherit from to provide the necessary methods to be a component. + + Example: + + .. code-block:: python + + from pydantic import BaseModel + from autogen_core import Component + + class Config(BaseModel): + value: str + + class MyComponent(Component("custom", Config)): + def __init__(self, value: str): + self.value = value + + def _to_config(self) -> Config: + return Config(value=self.value) + + @classmethod + def _from_config(cls, config: Config) -> MyComponent: + return cls(value=config.value) + + + Args: + component_type (ComponentType): What is the logical type of the component. + config_schema (type[T]): Pydantic model class which represents the configuration of the component. + description (str | None, optional): Helpful description of the component. Defaults to None. + + Returns: + type[ComponentConfig[T]]: A class to be directly inherited from. + """ + return types.new_class( + "Component", + (ComponentConfig[T],), + exec_body=lambda ns: ns.update( + {"_config_schema": config_schema, "_component_type": component_type, "_description": description} + ), + ) diff --git a/python/packages/autogen-core/src/autogen_core/models/_model_client.py b/python/packages/autogen-core/src/autogen_core/models/_model_client.py index c141c480dd78..0cba5b07904e 100644 --- a/python/packages/autogen-core/src/autogen_core/models/_model_client.py +++ b/python/packages/autogen-core/src/autogen_core/models/_model_client.py @@ -1,17 +1,18 @@ from __future__ import annotations -from typing import Mapping, Optional, Sequence, runtime_checkable +from abc import ABC, abstractmethod +from typing import Mapping, Optional, Sequence from typing_extensions import ( Any, AsyncGenerator, - Protocol, Required, TypedDict, Union, ) from .. import CancellationToken +from .._component_config import ComponentLoader from ..tools import Tool, ToolSchema from ._types import CreateResult, LLMMessage, RequestUsage @@ -22,9 +23,9 @@ class ModelCapabilities(TypedDict, total=False): json_output: Required[bool] -@runtime_checkable -class ChatCompletionClient(Protocol): +class ChatCompletionClient(ABC, ComponentLoader): # Caching has to be handled internally as they can depend on the create args that were stored in the constructor + @abstractmethod async def create( self, messages: Sequence[LLMMessage], @@ -36,6 +37,7 @@ async def create( cancellation_token: Optional[CancellationToken] = None, ) -> CreateResult: ... + @abstractmethod def create_stream( self, messages: Sequence[LLMMessage], @@ -47,13 +49,18 @@ def create_stream( cancellation_token: Optional[CancellationToken] = None, ) -> AsyncGenerator[Union[str, CreateResult], None]: ... + @abstractmethod def actual_usage(self) -> RequestUsage: ... + @abstractmethod def total_usage(self) -> RequestUsage: ... + @abstractmethod def count_tokens(self, messages: Sequence[LLMMessage], tools: Sequence[Tool | ToolSchema] = []) -> int: ... + @abstractmethod def remaining_tokens(self, messages: Sequence[LLMMessage], tools: Sequence[Tool | ToolSchema] = []) -> int: ... @property + @abstractmethod def capabilities(self) -> ModelCapabilities: ... diff --git a/python/packages/autogen-core/test.py b/python/packages/autogen-core/test.py new file mode 100644 index 000000000000..b9c7af8cd899 --- /dev/null +++ b/python/packages/autogen-core/test.py @@ -0,0 +1,19 @@ +from autogen_ext.models.openai import AzureOpenAIChatCompletionClient +from autogen_ext.models.openai._azure_token_provider import AzureTokenProvider +from azure.identity import DefaultAzureCredential + +from autogen_core.models import ChatCompletionClient + + +az_model_client = AzureOpenAIChatCompletionClient( + azure_deployment="{your-azure-deployment}", + model="gpt-4o", + api_version="2024-06-01", + azure_endpoint="https://{your-custom-endpoint}.openai.azure.com/", + azure_ad_token_provider=AzureTokenProvider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"), +) + +comp = az_model_client.dump_component() +loaded = ChatCompletionClient.load_component(comp) + +print(loaded.__class__.__name__) \ No newline at end of file diff --git a/python/packages/autogen-core/tests/test_component_config.py b/python/packages/autogen-core/tests/test_component_config.py new file mode 100644 index 000000000000..0bcd2f71b2cf --- /dev/null +++ b/python/packages/autogen-core/tests/test_component_config.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from typing_extensions import Self + +import pytest +from autogen_core import Component +from autogen_core._component_config import ComponentLoader +from autogen_core.models._model_client import ChatCompletionClient +from autogen_test_utils import MyInnerComponent, MyOuterComponent +from pydantic import BaseModel + + +class MyConfig(BaseModel): + info: str + + +class MyComponent(Component("custom", MyConfig)): # type: ignore + def __init__(self, info: str): + self.info = info + + def _to_config(self) -> MyConfig: + return MyConfig(info=self.info) + + @classmethod + def _from_config(cls, config: MyConfig) -> MyComponent: + return cls(info=config.info) + + +def test_custom_component(): + comp = MyComponent("test") + comp2 = MyComponent.load_component(comp.dump_component()) + assert comp.info == comp2.info + assert comp.__class__ == comp2.__class__ + + +def test_custom_component_generic_loader(): + comp = MyComponent("test") + comp2 = ComponentLoader.load_component(comp.dump_component(), MyComponent) + assert comp.info == comp2.info + assert comp.__class__ == comp2.__class__ + + +def test_custom_component_incorrect_class(): + comp = MyComponent("test") + + with pytest.raises(TypeError): + _ = ComponentLoader.load_component(comp.dump_component(), str) + + +def test_nested_component_diff_module(): + inner_class = MyInnerComponent("inner") + comp = MyOuterComponent("test", inner_class) + dumped = comp.dump_component() + comp2 = MyOuterComponent.load_component(dumped) + assert comp.__class__ == comp2.__class__ + assert comp.outer_message == comp2.outer_message + assert comp.inner_class.inner_message == comp2.inner_class.inner_message + assert comp.inner_class.__class__ == comp2.inner_class.__class__ + + +def test_cannot_import_locals(): + class InvalidModelClientConfig(BaseModel): + info: str + + class MyInvalidModelClient(Component("model", InvalidModelClientConfig)): + def __init__(self, info: str): + self.info = info + + def _to_config(self) -> InvalidModelClientConfig: + return InvalidModelClientConfig(info=self.info) + + @classmethod + def _from_config(cls, config: InvalidModelClientConfig) -> Self: + return cls(info=config.info) + + comp = MyInvalidModelClient("test") + with pytest.raises(TypeError): + # Fails due to the class not being importable + ChatCompletionClient.load_component(comp.dump_component()) + + +class InvalidModelClientConfig(BaseModel): + info: str + + +class MyInvalidModelClient(Component("model", InvalidModelClientConfig)): + def __init__(self, info: str): + self.info = info + + def _to_config(self) -> InvalidModelClientConfig: + return InvalidModelClientConfig(info=self.info) + + @classmethod + def _from_config(cls, config: InvalidModelClientConfig) -> Self: + return cls(info=config.info) + + +def test_type_error_on_creation(): + comp = MyInvalidModelClient("test") + # Fails due to MyInvalidModelClient not being a + with pytest.raises(TypeError): + ChatCompletionClient.load_component(comp.dump_component()) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_azure_token_provider.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_azure_token_provider.py new file mode 100644 index 000000000000..d6ea30e08786 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_azure_token_provider.py @@ -0,0 +1,50 @@ +from typing import List, Self +from azure.core.credentials import TokenProvider +from azure.identity import DefaultAzureCredential, get_bearer_token_provider +from autogen_core import Component +from pydantic import BaseModel + +class TokenProviderConfig(BaseModel): + provider_kind: str + scopes: List[str] + +class AzureTokenProvider(Component("token", TokenProviderConfig)): + def __init__(self, credential: TokenProvider, *scopes: str): + self.credential = credential + self.scopes = list(scopes) + self.provider = get_bearer_token_provider(self.credential, *self.scopes) + + def __call__(self) -> str: + return self.provider() + + def _to_config(self) -> TokenProviderConfig: + """Dump the configuration that would be requite to create a new instance of a component matching the configuration of this instance. + + Returns: + T: The configuration of the component. + """ + + if isinstance(self.credential, DefaultAzureCredential): + # we are not currently inspecting the chained credentials + return TokenProviderConfig(provider_kind="DefaultAzureCredential", scopes=self.scopes) + else: + raise ValueError("Only DefaultAzureCredential is supported") + + + @classmethod + def _from_config(cls, config: TokenProviderConfig) -> Self: + """Create a new instance of the component from a configuration object. + + Args: + config (T): The configuration object. + + Returns: + Self: The new instance of the component. + """ + + if config.provider_kind == "DefaultAzureCredential": + return cls(DefaultAzureCredential(), *config.scopes) + else: + raise ValueError("Only DefaultAzureCredential is supported") + + diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index b7276888411f..f9c39cfcc994 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py @@ -13,6 +13,7 @@ List, Mapping, Optional, + Self, Sequence, Set, Type, @@ -25,6 +26,7 @@ EVENT_LOGGER_NAME, TRACE_LOGGER_NAME, CancellationToken, + Component, FunctionCall, Image, ) @@ -63,9 +65,11 @@ from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.shared_params import FunctionDefinition, FunctionParameters -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, TypeAdapter, create_model from typing_extensions import Unpack +from autogen_ext.models.openai._azure_token_provider import AzureTokenProvider + from . import _model_info from .config import AzureOpenAIClientConfiguration, OpenAIClientConfiguration @@ -970,8 +974,10 @@ def __setstate__(self, state: Dict[str, Any]) -> None: self.__dict__.update(state) self._client = _openai_client_from_config(state["_raw_config"]) +class ConfigHolder(BaseModel): + values: Dict[str, Any] -class AzureOpenAIChatCompletionClient(BaseOpenAIChatCompletionClient): +class AzureOpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component("model", ConfigHolder)): """Chat completion client for Azure OpenAI hosted models. Args: @@ -1039,3 +1045,23 @@ def __getstate__(self) -> Dict[str, Any]: def __setstate__(self, state: Dict[str, Any]) -> None: self.__dict__.update(state) self._client = _azure_openai_client_from_config(state["_raw_config"]) + + + def _to_config(self) -> ConfigHolder: + copied_config = self._raw_config.copy() + + if "azure_ad_token_provider" in copied_config: + if not isinstance(copied_config["azure_ad_token_provider"], AzureTokenProvider): + raise ValueError("azure_ad_token_provider must be a AzureTokenProvider to be component serialized") + + copied_config["azure_ad_token_provider"] = copied_config["azure_ad_token_provider"].dump_component() + + return ConfigHolder(values=copied_config) + + @classmethod + def _from_config(cls, config: ConfigHolder) -> Self: + copied_config = config.values.copy() + if "azure_ad_token_provider" in copied_config: + copied_config["azure_ad_token_provider"] = AzureTokenProvider.load_component(copied_config["azure_ad_token_provider"]) + + return cls(**copied_config) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py index 70671347a1a2..04a366cfe86b 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/config/__init__.py @@ -3,6 +3,8 @@ from autogen_core.models import ModelCapabilities from typing_extensions import Required, TypedDict +from autogen_ext.models.openai._azure_token_provider import AzureTokenProvider + class ResponseFormat(TypedDict): type: Literal["text", "json_object"] @@ -46,7 +48,7 @@ class AzureOpenAIClientConfiguration(BaseOpenAIClientConfiguration, total=False) azure_deployment: str api_version: Required[str] azure_ad_token: str - azure_ad_token_provider: AsyncAzureADTokenProvider + azure_ad_token_provider: AsyncAzureADTokenProvider | AzureTokenProvider __all__ = ["AzureOpenAIClientConfiguration", "OpenAIClientConfiguration"] diff --git a/python/packages/autogen-test-utils/src/autogen_test_utils/__init__.py b/python/packages/autogen-test-utils/src/autogen_test_utils/__init__.py index f1aeed2292ce..0902a9841662 100644 --- a/python/packages/autogen-test-utils/src/autogen_test_utils/__init__.py +++ b/python/packages/autogen-test-utils/src/autogen_test_utils/__init__.py @@ -1,8 +1,11 @@ +from __future__ import annotations from dataclasses import dataclass -from typing import Any +from typing import Any, Literal from autogen_core import BaseAgent, DefaultTopicId, MessageContext, RoutedAgent, default_subscription, message_handler - +from autogen_core._component_config import ComponentConfig, ComponentLoader, ComponentModel +from pydantic import BaseModel +from autogen_core import Component @dataclass class MessageType: ... @@ -58,3 +61,39 @@ def __init__(self) -> None: async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any: raise NotImplementedError + + + +class MyInnerConfig(BaseModel): + inner_message: str + + +class MyInnerComponent(Component("custom", MyInnerConfig)): + def __init__(self, inner_message: str): + self.inner_message = inner_message + + def _to_config(self) -> MyInnerConfig: + return MyInnerConfig(inner_message=self.inner_message) + + @classmethod + def _from_config(cls, config: MyInnerConfig) -> MyInnerComponent: + return cls(inner_message=config.inner_message) + +class MyOuterConfig(BaseModel): + outer_message: str + inner_class: ComponentModel + + +class MyOuterComponent(Component("custom", MyOuterConfig)): + def __init__(self, outer_message: str, inner_class: MyInnerComponent): + self.outer_message = outer_message + self.inner_class = inner_class + + def _to_config(self) -> MyOuterConfig: + inner_component_config = self.inner_class.dump_component() + return MyOuterConfig(outer_message=self.outer_message, inner_class=inner_component_config) + + @classmethod + def _from_config(cls, config: MyOuterConfig) -> MyOuterComponent: + inner = MyInnerComponent.load_component(config.inner_class) + return cls(outer_message=config.outer_message, inner_class=inner) From 6e47f83b708a3a3de7a27eb2a56efcfb8b150aed Mon Sep 17 00:00:00 2001 From: Jack Gerrits Date: Thu, 19 Dec 2024 09:17:37 -0500 Subject: [PATCH 2/5] finishing touches --- .../autogen-core/docs/src/reference/index.md | 2 +- .../framework/component-config.ipynb | 66 +++++++- .../autogen-core/src/autogen_core/__init__.py | 2 - .../src/autogen_core/_component_config.py | 143 +++++++++--------- .../tests/test_component_config.py | 80 +++++++--- .../models/openai/_azure_token_provider.py | 18 ++- .../models/openai/_openai_client.py | 17 ++- .../replay/_replay_chat_completion_client.py | 2 +- .../src/autogen_test_utils/__init__.py | 31 +++- 9 files changed, 237 insertions(+), 124 deletions(-) diff --git a/python/packages/autogen-core/docs/src/reference/index.md b/python/packages/autogen-core/docs/src/reference/index.md index f6fadd12960f..905521debf3a 100644 --- a/python/packages/autogen-core/docs/src/reference/index.md +++ b/python/packages/autogen-core/docs/src/reference/index.md @@ -25,7 +25,7 @@ python/autogen_agentchat.state :hidden: :caption: AutoGen Core -autogen_core +python/autogen_core python/autogen_core.code_executor python/autogen_core.models python/autogen_core.model_context diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb b/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb index 0e7dfed63263..a3f7efa644fe 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb +++ b/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/component-config.ipynb @@ -29,16 +29,29 @@ "1. Add a call to {py:meth}`~autogen_core.Component` in the class inheritance list.\n", "2. Implment the {py:meth}`~autogen_core.ComponentConfigImpl._to_config` and {py:meth}`~autogen_core.ComponentConfigImpl._from_config` methods\n", "\n", - "For example:\n", + "For example:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations\n", "\n", - "```python\n", - "from pydantic import BaseModel\n", "from autogen_core import Component\n", + "from pydantic import BaseModel\n", + "\n", "\n", "class Config(BaseModel):\n", " value: str\n", "\n", - "class MyComponent(Component(\"custom\", Config)):\n", + "\n", + "class MyComponent(Component[Config]):\n", + " component_type = \"custom\"\n", + " config_schema = Config\n", + "\n", " def __init__(self, value: str):\n", " self.value = value\n", "\n", @@ -47,14 +60,53 @@ "\n", " @classmethod\n", " def _from_config(cls, config: Config) -> MyComponent:\n", - " return cls(value=config.value)\n", - "```\n" + " return cls(value=config.value)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "## Secrets\n", + "\n", + "If a field of a config object is a secret value, it should be marked using [`SecretStr`](https://docs.pydantic.dev/latest/api/types/#pydantic.types.SecretStr), this will ensure that the value will not be dumped to the config object.\n", + "\n", + "For example:\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from pydantic import BaseModel, SecretStr\n", + "\n", + "\n", + "class ClientConfig(BaseModel):\n", + " endpoint: str\n", + " api_key: SecretStr" ] } ], "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, "language_info": { - "name": "python" + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.7" } }, "nbformat": 4, diff --git a/python/packages/autogen-core/src/autogen_core/__init__.py b/python/packages/autogen-core/src/autogen_core/__init__.py index 1254f294600d..9523476e47e6 100644 --- a/python/packages/autogen-core/src/autogen_core/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/__init__.py @@ -14,7 +14,6 @@ from ._closure_agent import ClosureAgent, ClosureContext from ._component_config import ( Component, - ComponentConfig, ComponentConfigImpl, ComponentLoader, ComponentModel, @@ -108,7 +107,6 @@ "EVENT_LOGGER_NAME", "TRACE_LOGGER_NAME", "Component", - "ComponentConfig", "ComponentLoader", "ComponentConfigImpl", "ComponentModel", diff --git a/python/packages/autogen-core/src/autogen_core/_component_config.py b/python/packages/autogen-core/src/autogen_core/_component_config.py index e3ad9c92bd25..34ea34672ad7 100644 --- a/python/packages/autogen-core/src/autogen_core/_component_config.py +++ b/python/packages/autogen-core/src/autogen_core/_component_config.py @@ -1,17 +1,18 @@ from __future__ import annotations + import importlib -import types -from typing import Generic, Literal, Protocol, Type, cast, overload, runtime_checkable +import warnings +from typing import Any, ClassVar, Generic, Literal, Protocol, Type, cast, overload, runtime_checkable from pydantic import BaseModel from typing_extensions import Self, TypeVar - ComponentType = Literal["model", "agent", "tool", "termination", "token_provider"] | str ConfigT = TypeVar("ConfigT", bound=BaseModel) T = TypeVar("T", bound=BaseModel, covariant=True) + class ComponentModel(BaseModel): """Model class for a component. Contains all information required to instantiate a component.""" @@ -31,8 +32,13 @@ class ComponentModel(BaseModel): def _type_to_provider_str(t: type) -> str: return f"{t.__module__}.{t.__qualname__}" + @runtime_checkable class ComponentConfigImpl(Protocol[ConfigT]): + # Ideally would be ClassVar[Type[ConfigT]], but this is disallowed https://github.com/python/typing/discussions/1424 (despite being valid in this context) + config_schema: Type[ConfigT] + component_type: ClassVar[ComponentType] + """The two methods a class must implement to be a component. Args: @@ -65,17 +71,16 @@ def _from_config(cls, config: ConfigT) -> Self: ExpectedType = TypeVar("ExpectedType") -class ComponentLoader(): + +class ComponentLoader: @overload @classmethod - def load_component(cls, model: ComponentModel, expected: None = None) -> Self: - ... + def load_component(cls, model: ComponentModel, expected: None = None) -> Self: ... @overload @classmethod - def load_component(cls, model: ComponentModel, expected: Type[ExpectedType]) -> ExpectedType: - ... + def load_component(cls, model: ComponentModel, expected: Type[ExpectedType]) -> ExpectedType: ... @classmethod def load_component(cls, model: ComponentModel, expected: Type[ExpectedType] | None = None) -> Self | ExpectedType: @@ -88,7 +93,7 @@ def load_component(cls, model: ComponentModel, expected: Type[ExpectedType] | No from autogen_core import ComponentModel from autogen_core.models import ChatCompletionClient - component: ComponentModel = ... # type: ignore + component: ComponentModel = ... # type: ignore model_client = ChatCompletionClient.load_component(component) @@ -123,8 +128,18 @@ def load_component(cls, model: ComponentModel, expected: Type[ExpectedType] | No if not isinstance(component_class, ComponentConfigImpl): raise TypeError("Invalid component class") + # We need to check the schema is valid + if not hasattr(component_class, "config_schema"): + raise AttributeError("config_schema not defined") + + if not hasattr(component_class, "component_type"): + raise AttributeError("component_type not defined") + + schema = component_class.config_schema + validated_config = schema.model_validate(model.config) + # We're allowed to use the private method here - instance = component_class._from_config(model.config) # type: ignore + instance = component_class._from_config(validated_config) # type: ignore if expected is None and not isinstance(instance, cls): raise TypeError("Expected type does not match") @@ -136,56 +151,11 @@ def load_component(cls, model: ComponentModel, expected: Type[ExpectedType] | No return cast(ExpectedType, instance) -class ComponentConfig(ComponentConfigImpl[ConfigT], ComponentLoader, Generic[ConfigT]): - _config_schema: Type[ConfigT] - _component_type: ComponentType - _description: str | None - - def dump_component(self) -> ComponentModel: - """Dump the component to a model that can be loaded back in. +class Component(ComponentConfigImpl[ConfigT], ComponentLoader, Generic[ConfigT]): + """To create a component class, inherit from this class. Then implement two class variables: - Raises: - TypeError: If the component is a local class. - - Returns: - ComponentModel: The model representing the component. - """ - provider = _type_to_provider_str(self.__class__) - - if "" in provider: - raise TypeError("Cannot dump component with local class") - - return ComponentModel( - provider=provider, - component_type=self._component_type, - version=1, - description=self._description, - config=self._to_config(), - ) - - @classmethod - def component_config_schema(cls) -> Type[ConfigT]: - """Get the configuration schema for the component. - - Returns: - Type[ConfigT]: The configuration schema for the component. - """ - return cls._config_schema - - @classmethod - def component_type(cls) -> ComponentType: - """Logical type of the component. - - Returns: - ComponentType: The logical type of the component. - """ - return cls._component_type - - -def Component( - component_type: ComponentType, config_schema: type[T], description: str | None = None -) -> type[ComponentConfig[T]]: - """This enables easy creation of Component classes. It provides a type to inherit from to provide the necessary methods to be a component. + - :py:attr:`config_schema` - A Pydantic model class which represents the configuration of the component. This is also the type parameter of Component. + - :py:attr:`component_type` - What is the logical type of the component. Example: @@ -194,10 +164,15 @@ def Component( from pydantic import BaseModel from autogen_core import Component + class Config(BaseModel): value: str - class MyComponent(Component("custom", Config)): + + class MyComponent(Component[Config]): + component_type = "custom" + config_schema = Config + def __init__(self, value: str): self.value = value @@ -207,20 +182,40 @@ def _to_config(self) -> Config: @classmethod def _from_config(cls, config: Config) -> MyComponent: return cls(value=config.value) + """ + required_class_vars = ["config_schema", "component_type"] - Args: - component_type (ComponentType): What is the logical type of the component. - config_schema (type[T]): Pydantic model class which represents the configuration of the component. - description (str | None, optional): Helpful description of the component. Defaults to None. + def __init_subclass__(cls, **kwargs: Any): + super().__init_subclass__(**kwargs) - Returns: - type[ComponentConfig[T]]: A class to be directly inherited from. - """ - return types.new_class( - "Component", - (ComponentConfig[T],), - exec_body=lambda ns: ns.update( - {"_config_schema": config_schema, "_component_type": component_type, "_description": description} - ), - ) + for var in cls.required_class_vars: + if not hasattr(cls, var): + warnings.warn( + f"Class variable '{var}' must be defined in {cls.__name__} to be a valid component", stacklevel=2 + ) + + def dump_component(self) -> ComponentModel: + """Dump the component to a model that can be loaded back in. + + Raises: + TypeError: If the component is a local class. + + Returns: + ComponentModel: The model representing the component. + """ + provider = _type_to_provider_str(self.__class__) + + if "" in provider: + raise TypeError("Cannot dump component with local class") + + if not hasattr(self, "component_type"): + raise AttributeError("component_type not defined") + + return ComponentModel( + provider=provider, + component_type=self.component_type, + version=1, + description=None, + config=self._to_config(), + ) diff --git a/python/packages/autogen-core/tests/test_component_config.py b/python/packages/autogen-core/tests/test_component_config.py index 0bcd2f71b2cf..081dd5b8124b 100644 --- a/python/packages/autogen-core/tests/test_component_config.py +++ b/python/packages/autogen-core/tests/test_component_config.py @@ -1,21 +1,23 @@ from __future__ import annotations -from typing_extensions import Self - import pytest -from autogen_core import Component -from autogen_core._component_config import ComponentLoader -from autogen_core.models._model_client import ChatCompletionClient +from autogen_core import Component, ComponentLoader, ComponentModel +from autogen_core._component_config import _type_to_provider_str # type: ignore +from autogen_core.models import ChatCompletionClient from autogen_test_utils import MyInnerComponent, MyOuterComponent -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError +from typing_extensions import Self class MyConfig(BaseModel): info: str -class MyComponent(Component("custom", MyConfig)): # type: ignore - def __init__(self, info: str): +class MyComponent(Component[MyConfig]): + config_schema = MyConfig + component_type = "custom" + + def __init__(self, info: str) -> None: self.info = info def _to_config(self) -> MyConfig: @@ -26,28 +28,28 @@ def _from_config(cls, config: MyConfig) -> MyComponent: return cls(info=config.info) -def test_custom_component(): +def test_custom_component() -> None: comp = MyComponent("test") comp2 = MyComponent.load_component(comp.dump_component()) assert comp.info == comp2.info assert comp.__class__ == comp2.__class__ -def test_custom_component_generic_loader(): +def test_custom_component_generic_loader() -> None: comp = MyComponent("test") comp2 = ComponentLoader.load_component(comp.dump_component(), MyComponent) assert comp.info == comp2.info assert comp.__class__ == comp2.__class__ -def test_custom_component_incorrect_class(): +def test_custom_component_incorrect_class() -> None: comp = MyComponent("test") with pytest.raises(TypeError): _ = ComponentLoader.load_component(comp.dump_component(), str) -def test_nested_component_diff_module(): +def test_nested_component_diff_module() -> None: inner_class = MyInnerComponent("inner") comp = MyOuterComponent("test", inner_class) dumped = comp.dump_component() @@ -58,11 +60,14 @@ def test_nested_component_diff_module(): assert comp.inner_class.__class__ == comp2.inner_class.__class__ -def test_cannot_import_locals(): +def test_cannot_import_locals() -> None: class InvalidModelClientConfig(BaseModel): info: str - class MyInvalidModelClient(Component("model", InvalidModelClientConfig)): + class MyInvalidModelClient(Component[InvalidModelClientConfig]): + config_schema = InvalidModelClientConfig + component_type = "model" + def __init__(self, info: str): self.info = info @@ -83,8 +88,11 @@ class InvalidModelClientConfig(BaseModel): info: str -class MyInvalidModelClient(Component("model", InvalidModelClientConfig)): - def __init__(self, info: str): +class MyInvalidModelClient(Component[InvalidModelClientConfig]): + config_schema = InvalidModelClientConfig + component_type = "model" + + def __init__(self, info: str) -> None: self.info = info def _to_config(self) -> InvalidModelClientConfig: @@ -95,8 +103,44 @@ def _from_config(cls, config: InvalidModelClientConfig) -> Self: return cls(info=config.info) -def test_type_error_on_creation(): +def test_type_error_on_creation() -> None: comp = MyInvalidModelClient("test") - # Fails due to MyInvalidModelClient not being a + # Fails due to MyInvalidModelClient not being a model client with pytest.raises(TypeError): ChatCompletionClient.load_component(comp.dump_component()) + + +with pytest.warns(UserWarning): + + class MyInvalidMissingAttrs(Component[InvalidModelClientConfig]): + def __init__(self, info: str): + self.info = info + + def _to_config(self) -> InvalidModelClientConfig: + return InvalidModelClientConfig(info=self.info) + + @classmethod + def _from_config(cls, config: InvalidModelClientConfig) -> Self: + return cls(info=config.info) + + +def test_fails_to_save_on_missing_attributes() -> None: + comp = MyInvalidMissingAttrs("test") # type: ignore + with pytest.raises(AttributeError): + comp.dump_component() + + +def test_schema_validation_fails_on_bad_config() -> None: + class OtherConfig(BaseModel): + other: str + + config = OtherConfig(other="test") + model = ComponentModel( + provider=_type_to_provider_str(MyComponent), + component_type=MyComponent.component_type, + version=1, + description=None, + config=config, + ) + with pytest.raises(ValidationError): + _ = MyComponent.load_component(model) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_azure_token_provider.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_azure_token_provider.py index d6ea30e08786..a3960fbefa53 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_azure_token_provider.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_azure_token_provider.py @@ -1,14 +1,21 @@ -from typing import List, Self +from typing import List + +from autogen_core import Component from azure.core.credentials import TokenProvider from azure.identity import DefaultAzureCredential, get_bearer_token_provider -from autogen_core import Component from pydantic import BaseModel +from typing_extensions import Self + class TokenProviderConfig(BaseModel): provider_kind: str scopes: List[str] -class AzureTokenProvider(Component("token", TokenProviderConfig)): + +class AzureTokenProvider(Component[TokenProviderConfig]): + component_type = "token_provider" + config_schema = TokenProviderConfig + def __init__(self, credential: TokenProvider, *scopes: str): self.credential = credential self.scopes = list(scopes) @@ -25,12 +32,11 @@ def _to_config(self) -> TokenProviderConfig: """ if isinstance(self.credential, DefaultAzureCredential): - # we are not currently inspecting the chained credentials + # NOTE: we are not currently inspecting the chained credentials, so this could result in a loss of information return TokenProviderConfig(provider_kind="DefaultAzureCredential", scopes=self.scopes) else: raise ValueError("Only DefaultAzureCredential is supported") - @classmethod def _from_config(cls, config: TokenProviderConfig) -> Self: """Create a new instance of the component from a configuration object. @@ -46,5 +52,3 @@ def _from_config(cls, config: TokenProviderConfig) -> Self: return cls(DefaultAzureCredential(), *config.scopes) else: raise ValueError("Only DefaultAzureCredential is supported") - - diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index f9c39cfcc994..44aa1a268f86 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py @@ -13,7 +13,6 @@ List, Mapping, Optional, - Self, Sequence, Set, Type, @@ -65,8 +64,8 @@ from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.shared_params import FunctionDefinition, FunctionParameters -from pydantic import BaseModel, ConfigDict, TypeAdapter, create_model -from typing_extensions import Unpack +from pydantic import BaseModel +from typing_extensions import Self, Unpack from autogen_ext.models.openai._azure_token_provider import AzureTokenProvider @@ -974,10 +973,12 @@ def __setstate__(self, state: Dict[str, Any]) -> None: self.__dict__.update(state) self._client = _openai_client_from_config(state["_raw_config"]) + class ConfigHolder(BaseModel): values: Dict[str, Any] -class AzureOpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component("model", ConfigHolder)): + +class AzureOpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component[ConfigHolder]): """Chat completion client for Azure OpenAI hosted models. Args: @@ -1025,6 +1026,9 @@ class AzureOpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component( """ + component_type = "model" + config_schema = ConfigHolder + def __init__(self, **kwargs: Unpack[AzureOpenAIClientConfiguration]): model_capabilities: Optional[ModelCapabilities] = None copied_args = dict(kwargs).copy() @@ -1046,7 +1050,6 @@ def __setstate__(self, state: Dict[str, Any]) -> None: self.__dict__.update(state) self._client = _azure_openai_client_from_config(state["_raw_config"]) - def _to_config(self) -> ConfigHolder: copied_config = self._raw_config.copy() @@ -1062,6 +1065,8 @@ def _to_config(self) -> ConfigHolder: def _from_config(cls, config: ConfigHolder) -> Self: copied_config = config.values.copy() if "azure_ad_token_provider" in copied_config: - copied_config["azure_ad_token_provider"] = AzureTokenProvider.load_component(copied_config["azure_ad_token_provider"]) + copied_config["azure_ad_token_provider"] = AzureTokenProvider.load_component( + copied_config["azure_ad_token_provider"] + ) return cls(**copied_config) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py b/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py index d0cd159a3aee..5b7ac7095c2b 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py @@ -16,7 +16,7 @@ logger = logging.getLogger(EVENT_LOGGER_NAME) -class ReplayChatCompletionClient: +class ReplayChatCompletionClient(ChatCompletionClient): """ A mock chat completion client that replays predefined responses using an index-based approach. diff --git a/python/packages/autogen-test-utils/src/autogen_test_utils/__init__.py b/python/packages/autogen-test-utils/src/autogen_test_utils/__init__.py index 0902a9841662..b4f9332b6cc6 100644 --- a/python/packages/autogen-test-utils/src/autogen_test_utils/__init__.py +++ b/python/packages/autogen-test-utils/src/autogen_test_utils/__init__.py @@ -1,11 +1,20 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import Any, Literal -from autogen_core import BaseAgent, DefaultTopicId, MessageContext, RoutedAgent, default_subscription, message_handler -from autogen_core._component_config import ComponentConfig, ComponentLoader, ComponentModel +from dataclasses import dataclass +from typing import Any + +from autogen_core import ( + BaseAgent, + Component, + DefaultTopicId, + MessageContext, + RoutedAgent, + default_subscription, + message_handler, +) +from autogen_core._component_config import ComponentModel from pydantic import BaseModel -from autogen_core import Component + @dataclass class MessageType: ... @@ -63,12 +72,14 @@ async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any: raise NotImplementedError - class MyInnerConfig(BaseModel): inner_message: str -class MyInnerComponent(Component("custom", MyInnerConfig)): +class MyInnerComponent(Component[MyInnerConfig]): + config_schema = MyInnerConfig + component_type = "custom" + def __init__(self, inner_message: str): self.inner_message = inner_message @@ -79,12 +90,16 @@ def _to_config(self) -> MyInnerConfig: def _from_config(cls, config: MyInnerConfig) -> MyInnerComponent: return cls(inner_message=config.inner_message) + class MyOuterConfig(BaseModel): outer_message: str inner_class: ComponentModel -class MyOuterComponent(Component("custom", MyOuterConfig)): +class MyOuterComponent(Component[MyOuterConfig]): + config_schema = MyOuterConfig + component_type = "custom" + def __init__(self, outer_message: str, inner_class: MyInnerComponent): self.outer_message = outer_message self.inner_class = inner_class From ce0935b394c10e265b7e656225ecbc0f5958532e Mon Sep 17 00:00:00 2001 From: Jack Gerrits Date: Thu, 19 Dec 2024 09:22:41 -0500 Subject: [PATCH 3/5] remove test file --- python/packages/autogen-core/test.py | 19 ------------------- 1 file changed, 19 deletions(-) delete mode 100644 python/packages/autogen-core/test.py diff --git a/python/packages/autogen-core/test.py b/python/packages/autogen-core/test.py deleted file mode 100644 index b9c7af8cd899..000000000000 --- a/python/packages/autogen-core/test.py +++ /dev/null @@ -1,19 +0,0 @@ -from autogen_ext.models.openai import AzureOpenAIChatCompletionClient -from autogen_ext.models.openai._azure_token_provider import AzureTokenProvider -from azure.identity import DefaultAzureCredential - -from autogen_core.models import ChatCompletionClient - - -az_model_client = AzureOpenAIChatCompletionClient( - azure_deployment="{your-azure-deployment}", - model="gpt-4o", - api_version="2024-06-01", - azure_endpoint="https://{your-custom-endpoint}.openai.azure.com/", - azure_ad_token_provider=AzureTokenProvider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"), -) - -comp = az_model_client.dump_component() -loaded = ChatCompletionClient.load_component(comp) - -print(loaded.__class__.__name__) \ No newline at end of file From 7ed39a55fc8fa6949f54f48077e56f3c8ddc1d49 Mon Sep 17 00:00:00 2001 From: Jack Gerrits Date: Thu, 19 Dec 2024 09:49:05 -0500 Subject: [PATCH 4/5] fix json usage --- python/mypy_plugin.py | 23 +++++++++++++ .../src/autogen_core/_component_config.py | 28 ++++++++++------ .../tests/test_component_config.py | 32 ++++++++++++++++++- .../models/openai/_openai_client.py | 2 +- 4 files changed, 74 insertions(+), 11 deletions(-) create mode 100644 python/mypy_plugin.py diff --git a/python/mypy_plugin.py b/python/mypy_plugin.py new file mode 100644 index 000000000000..4e230e732627 --- /dev/null +++ b/python/mypy_plugin.py @@ -0,0 +1,23 @@ +from typing import Callable +from mypy.plugin import Plugin, DynamicClassDefContext, SymbolTableNode +from mypy.nodes import SymbolTableNode + +class CustomPlugin(Plugin): + def get_dynamic_class_hook( + self, fullname: str + ) -> Callable[[DynamicClassDefContext], None] | None: + + def hook(ctx: DynamicClassDefContext) -> None: + if "Component" in fullname: + # We need to generate mypy.nodes.TypeInfo + # to make mypy understand the type of the class + ctx.api.add_symbol_table_node( + fullname, SymbolTableNode( + + ) + ) + return + +def plugin(version: str): + # ignore version argument if the plugin works with all mypy versions. + return CustomPlugin \ No newline at end of file diff --git a/python/packages/autogen-core/src/autogen_core/_component_config.py b/python/packages/autogen-core/src/autogen_core/_component_config.py index 34ea34672ad7..6652ea6cd2aa 100644 --- a/python/packages/autogen-core/src/autogen_core/_component_config.py +++ b/python/packages/autogen-core/src/autogen_core/_component_config.py @@ -2,7 +2,7 @@ import importlib import warnings -from typing import Any, ClassVar, Generic, Literal, Protocol, Type, cast, overload, runtime_checkable +from typing import Any, ClassVar, Dict, Generic, Literal, Protocol, Type, cast, overload, runtime_checkable from pydantic import BaseModel from typing_extensions import Self, TypeVar @@ -25,8 +25,8 @@ class ComponentModel(BaseModel): description: str | None """Description of the component.""" - config: BaseModel - """The config field is passed to a given class's implmentation of :py:meth:`autogen_core.ComponentConfigImpl._from_config` to create a new instance of the component class.""" + config: dict[str, Any] + """The schema validated config field is passed to a given class's implmentation of :py:meth:`autogen_core.ComponentConfigImpl._from_config` to create a new instance of the component class.""" def _type_to_provider_str(t: type) -> str: @@ -76,14 +76,16 @@ def _from_config(cls, config: ConfigT) -> Self: class ComponentLoader: @overload @classmethod - def load_component(cls, model: ComponentModel, expected: None = None) -> Self: ... + def load_component(cls, model: ComponentModel | Dict[str, Any], expected: None = None) -> Self: ... @overload @classmethod - def load_component(cls, model: ComponentModel, expected: Type[ExpectedType]) -> ExpectedType: ... + def load_component(cls, model: ComponentModel | Dict[str, Any], expected: Type[ExpectedType]) -> ExpectedType: ... @classmethod - def load_component(cls, model: ComponentModel, expected: Type[ExpectedType] | None = None) -> Self | ExpectedType: + def load_component( + cls, model: ComponentModel | Dict[str, Any], expected: Type[ExpectedType] | None = None + ) -> Self | ExpectedType: """Load a component from a model. Intended to be used with the return type of :py:meth:`autogen_core.ComponentConfig.dump_component`. Example: @@ -117,7 +119,12 @@ def load_component(cls, model: ComponentModel, expected: Type[ExpectedType] | No # Use global and add further type checks - output = model.provider.rsplit(".", maxsplit=1) + if isinstance(model, dict): + loaded_model = ComponentModel(**model) + else: + loaded_model = model + + output = loaded_model.provider.rsplit(".", maxsplit=1) if len(output) != 2: raise ValueError("Invalid") @@ -136,7 +143,7 @@ def load_component(cls, model: ComponentModel, expected: Type[ExpectedType] | No raise AttributeError("component_type not defined") schema = component_class.config_schema - validated_config = schema.model_validate(model.config) + validated_config = schema.model_validate(loaded_model.config) # We're allowed to use the private method here instance = component_class._from_config(validated_config) # type: ignore @@ -161,6 +168,8 @@ class Component(ComponentConfigImpl[ConfigT], ComponentLoader, Generic[ConfigT]) .. code-block:: python + from __future__ import annotations + from pydantic import BaseModel from autogen_core import Component @@ -212,10 +221,11 @@ def dump_component(self) -> ComponentModel: if not hasattr(self, "component_type"): raise AttributeError("component_type not defined") + obj_config = self._to_config().model_dump() return ComponentModel( provider=provider, component_type=self.component_type, version=1, description=None, - config=self._to_config(), + config=obj_config, ) diff --git a/python/packages/autogen-core/tests/test_component_config.py b/python/packages/autogen-core/tests/test_component_config.py index 081dd5b8124b..6c0df593fd98 100644 --- a/python/packages/autogen-core/tests/test_component_config.py +++ b/python/packages/autogen-core/tests/test_component_config.py @@ -1,5 +1,7 @@ from __future__ import annotations +import json + import pytest from autogen_core import Component, ComponentLoader, ComponentModel from autogen_core._component_config import _type_to_provider_str # type: ignore @@ -42,6 +44,22 @@ def test_custom_component_generic_loader() -> None: assert comp.__class__ == comp2.__class__ +def test_custom_component_json() -> None: + comp = MyComponent("test") + json_str = comp.dump_component().model_dump_json() + comp2 = MyComponent.load_component(json.loads(json_str)) + assert comp.info == comp2.info + assert comp.__class__ == comp2.__class__ + + +def test_custom_component_generic_loader_json() -> None: + comp = MyComponent("test") + json_str = comp.dump_component().model_dump_json() + comp2 = ComponentLoader.load_component(json.loads(json_str), MyComponent) + assert comp.info == comp2.info + assert comp.__class__ == comp2.__class__ + + def test_custom_component_incorrect_class() -> None: comp = MyComponent("test") @@ -60,6 +78,18 @@ def test_nested_component_diff_module() -> None: assert comp.inner_class.__class__ == comp2.inner_class.__class__ +def test_nested_component_diff_module_json() -> None: + inner_class = MyInnerComponent("inner") + comp = MyOuterComponent("test", inner_class) + dumped = comp.dump_component() + json_str = dumped.model_dump_json() + comp2 = MyOuterComponent.load_component(json.loads(json_str)) + assert comp.__class__ == comp2.__class__ + assert comp.outer_message == comp2.outer_message + assert comp.inner_class.inner_message == comp2.inner_class.inner_message + assert comp.inner_class.__class__ == comp2.inner_class.__class__ + + def test_cannot_import_locals() -> None: class InvalidModelClientConfig(BaseModel): info: str @@ -134,7 +164,7 @@ def test_schema_validation_fails_on_bad_config() -> None: class OtherConfig(BaseModel): other: str - config = OtherConfig(other="test") + config = OtherConfig(other="test").model_dump() model = ComponentModel( provider=_type_to_provider_str(MyComponent), component_type=MyComponent.component_type, diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index 44aa1a268f86..5e826e30c9fd 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py @@ -1057,7 +1057,7 @@ def _to_config(self) -> ConfigHolder: if not isinstance(copied_config["azure_ad_token_provider"], AzureTokenProvider): raise ValueError("azure_ad_token_provider must be a AzureTokenProvider to be component serialized") - copied_config["azure_ad_token_provider"] = copied_config["azure_ad_token_provider"].dump_component() + copied_config["azure_ad_token_provider"] = copied_config["azure_ad_token_provider"].dump_component().model_dump() return ConfigHolder(values=copied_config) From d8241a6ee5bb28f1ce136dc224af144622c248f5 Mon Sep 17 00:00:00 2001 From: Jack Gerrits Date: Thu, 19 Dec 2024 09:55:13 -0500 Subject: [PATCH 5/5] Format --- .../src/autogen_ext/models/openai/_openai_client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index 5e826e30c9fd..6706535faddc 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py @@ -1057,7 +1057,9 @@ def _to_config(self) -> ConfigHolder: if not isinstance(copied_config["azure_ad_token_provider"], AzureTokenProvider): raise ValueError("azure_ad_token_provider must be a AzureTokenProvider to be component serialized") - copied_config["azure_ad_token_provider"] = copied_config["azure_ad_token_provider"].dump_component().model_dump() + copied_config["azure_ad_token_provider"] = ( + copied_config["azure_ad_token_provider"].dump_component().model_dump() + ) return ConfigHolder(values=copied_config)