ag2 0.9.1__py3-none-any.whl → 0.9.1.post0__py3-none-any.whl

autogen/interop/pydantic_ai/pydantic_ai.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+import sys
+import warnings
+from functools import wraps
+from inspect import signature
+from typing import Any, Callable, Optional
+
+from ...doc_utils import export_module
+from ...import_utils import optional_import_block, require_optional_import
+from ...tools import Tool
+from ..registry import register_interoperable_class
+
+__all__ = ["PydanticAIInteroperability"]
+
+with optional_import_block():
+    from pydantic_ai import RunContext
+    from pydantic_ai.tools import Tool as PydanticAITool
+    from pydantic_ai.usage import Usage
+
+
+@register_interoperable_class("pydanticai")
+@export_module("autogen.interop")
+class PydanticAIInteroperability:
+    """A class implementing the `Interoperable` protocol for converting Pydantic AI tools
+    into a general `Tool` format.
+
+    This class takes a `PydanticAITool` and converts it into a standard `Tool` object,
+    ensuring compatibility between Pydantic AI tools and other systems that expect
+    the `Tool` format. It also provides a mechanism for injecting context parameters
+    into the tool's function.
+    """
+
+    @staticmethod
+    @require_optional_import("pydantic_ai", "interop-pydantic-ai")
+    def inject_params(
+        ctx: Any,
+        tool: Any,
+    ) -> Callable[..., Any]:
+        """Wraps the tool's function to inject context parameters and handle retries.
+
+        This method ensures that context parameters are properly passed to the tool
+        when invoked and that retries are managed according to the tool's settings.
+
+        Args:
+            ctx (Optional[RunContext[Any]]): The run context, which may include dependencies and retry information.
+            tool (PydanticAITool): The Pydantic AI tool whose function is to be wrapped.
+
+        Returns:
+            Callable[..., Any]: A wrapped function that includes context injection and retry handling.
+
+        Raises:
+            ValueError: If the tool fails after the maximum number of retries.
+        """
+        ctx_typed: Optional[RunContext[Any]] = ctx  # type: ignore[no-any-unimported]
+        tool_typed: PydanticAITool[Any] = tool  # type: ignore[no-any-unimported]
+
+        max_retries = tool_typed.max_retries if tool_typed.max_retries is not None else 1
+        f = tool_typed.function
+
+        @wraps(f)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if tool_typed.current_retry >= max_retries:
+                raise ValueError(f"{tool_typed.name} failed after {max_retries} retries")
+
+            try:
+                if ctx_typed is not None:
+                    kwargs.pop("ctx", None)
+                    ctx_typed.retry = tool_typed.current_retry
+                    result = f(**kwargs, ctx=ctx_typed)  # type: ignore[call-arg]
+                else:
+                    result = f(**kwargs)  # type: ignore[call-arg]
+                tool_typed.current_retry = 0
+            except Exception as e:
+                tool_typed.current_retry += 1
+                raise e
+
+            return result
+
+        sig = signature(f)
+        if ctx_typed is not None:
+            new_params = [param for name, param in sig.parameters.items() if name != "ctx"]
+        else:
+            new_params = list(sig.parameters.values())
+
+        wrapper.__signature__ = sig.replace(parameters=new_params)  # type: ignore[attr-defined]
+
+        return wrapper
+
+    @classmethod
+    @require_optional_import("pydantic_ai", "interop-pydantic-ai")
+    def convert_tool(cls, tool: Any, deps: Any = None, **kwargs: Any) -> Tool:
+        """Converts a given Pydantic AI tool into a general `Tool` format.
+
+        This method verifies that the provided tool is a valid `PydanticAITool`,
+        handles context dependencies if necessary, and returns a standardized `Tool` object.
+
+        Args:
+            tool (Any): The tool to convert, expected to be an instance of `PydanticAITool`.
+            deps (Any, optional): The dependencies to inject into the context, required if
+                                   the tool takes a context. Defaults to None.
+            **kwargs (Any): Additional arguments that are not used in this method.
+
+        Returns:
+            Tool: A standardized `Tool` object converted from the Pydantic AI tool.
+
+        Raises:
+            ValueError: If the provided tool is not an instance of `PydanticAITool`, or if
+                        dependencies are missing for tools that require a context.
+            UserWarning: If the `deps` argument is provided for a tool that does not take a context.
+        """
+        if not isinstance(tool, PydanticAITool):
+            raise ValueError(f"Expected an instance of `pydantic_ai.tools.Tool`, got {type(tool)}")
+
+        # needed for type checking
+        pydantic_ai_tool: PydanticAITool[Any] = tool  # type: ignore[no-any-unimported]
+
+        if tool.takes_ctx and deps is None:
+            raise ValueError("If the tool takes a context, the `deps` argument must be provided")
+        if not tool.takes_ctx and deps is not None:
+            warnings.warn(
+                "The `deps` argument is provided but will be ignored because the tool does not take a context.",
+                UserWarning,
+            )
+
+        ctx = (
+            RunContext(
+                model=None,  # type: ignore [arg-type]
+                usage=Usage(),
+                prompt="",
+                deps=deps,
+                retry=0,
+                # All messages send to or returned by a model.
+                # This is mostly used on pydantic_ai Agent level.
+                messages=[],  # TODO: check in the future if this is needed on Tool level
+                tool_name=pydantic_ai_tool.name,
+            )
+            if tool.takes_ctx
+            else None
+        )
+
+        func = PydanticAIInteroperability.inject_params(
+            ctx=ctx,
+            tool=pydantic_ai_tool,
+        )
+
+        return Tool(
+            name=pydantic_ai_tool.name,
+            description=pydantic_ai_tool.description,
+            func_or_tool=func,
+            parameters_json_schema=pydantic_ai_tool._parameters_json_schema,
+        )
+
+    @classmethod
+    def get_unsupported_reason(cls) -> Optional[str]:
+        if sys.version_info < (3, 9):
+            return "This submodule is only supported for Python versions 3.9 and above"
+
+        with optional_import_block() as result:
+            import pydantic_ai.tools  # noqa: F401
+
+        if not result.is_successful:
+            return "Please install `interop-pydantic-ai` extra to use this module:\n\n\tpip install ag2[interop-pydantic-ai]"
+
+        return None

autogen/interop/registry.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Callable, TypeVar
+
+from ..doc_utils import export_module
+from .interoperable import Interoperable
+
+__all__ = ["InteroperableRegistry", "register_interoperable_class"]
+
+InteroperableClass = TypeVar("InteroperableClass", bound=type[Interoperable])
+
+
+class InteroperableRegistry:
+    def __init__(self) -> None:
+        self._registry: dict[str, type[Interoperable]] = {}
+
+    def register(self, short_name: str, cls: InteroperableClass) -> InteroperableClass:
+        if short_name in self._registry:
+            raise ValueError(f"Duplicate registration for {short_name}")
+
+        self._registry[short_name] = cls
+
+        return cls
+
+    def get_short_names(self) -> list[str]:
+        return sorted(self._registry.keys())
+
+    def get_supported_types(self) -> list[str]:
+        short_names = self.get_short_names()
+        supported_types = [name for name in short_names if self._registry[name].get_unsupported_reason() is None]
+        return supported_types
+
+    def get_class(self, short_name: str) -> type[Interoperable]:
+        return self._registry[short_name]
+
+    @classmethod
+    def get_instance(cls) -> "InteroperableRegistry":
+        return _register
+
+
+# global registry
+_register = InteroperableRegistry()
+
+
+# register decorator
+@export_module("autogen.interop")
+def register_interoperable_class(short_name: str) -> Callable[[InteroperableClass], InteroperableClass]:
+    """Register an Interoperable class in the global registry.
+
+    Returns:
+        Callable[[InteroperableClass], InteroperableClass]: Decorator function
+
+    Example:
+        ```python
+        @register_interoperable_class("myinterop")
+        class MyInteroperability(Interoperable):
+            def convert_tool(self, tool: Any) -> Tool:
+                # implementation
+                ...
+        ```
+    """
+
+    def inner(cls: InteroperableClass) -> InteroperableClass:
+        global _register
+        return _register.register(short_name, cls)
+
+    return inner

autogen/io/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+from .base import IOStream, InputStream, OutputStream
+from .console import IOConsole
+from .websockets import IOWebsockets
+
+# Set the default input/output stream to the console
+IOStream.set_global_default(IOConsole())
+IOStream.set_default(IOConsole())
+
+__all__ = ("IOConsole", "IOStream", "IOWebsockets", "InputStream", "OutputStream")

autogen/io/base.py

@@ -0,0 +1,151 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Any, Optional, Protocol, Union, runtime_checkable
+
+from ..doc_utils import export_module
+from ..events.base_event import BaseEvent
+
+__all__ = ("IOStream", "InputStream", "OutputStream")
+
+logger = logging.getLogger(__name__)
+
+
+@runtime_checkable
+@export_module("autogen.io")
+class OutputStream(Protocol):
+    def print(self, *objects: Any, sep: str = " ", end: str = "\n", flush: bool = False) -> None:
+        """Print data to the output stream.
+
+        Args:
+            objects (any): The data to print.
+            sep (str, optional): The separator between objects. Defaults to " ".
+            end (str, optional): The end of the output. Defaults to "\n".
+            flush (bool, optional): Whether to flush the output. Defaults to False.
+        """
+        ...  # pragma: no cover
+
+    def send(self, message: BaseEvent) -> None:
+        """Send data to the output stream.
+
+        Args:
+            message (BaseEvent): BaseEvent from autogen.messages.base_message
+        """
+        ...
+
+
+@runtime_checkable
+@export_module("autogen.io")
+class InputStream(Protocol):
+    def input(self, prompt: str = "", *, password: bool = False) -> str:
+        """Read a line from the input stream.
+
+        Args:
+            prompt (str, optional): The prompt to display. Defaults to "".
+            password (bool, optional): Whether to read a password. Defaults to False.
+
+        Returns:
+            str: The line read from the input stream.
+
+        """
+        ...  # pragma: no cover
+
+
+@runtime_checkable
+@export_module("autogen.io")
+class AsyncInputStream(Protocol):
+    async def input(self, prompt: str = "", *, password: bool = False) -> str:
+        """Read a line from the input stream.
+
+        Args:
+            prompt (str, optional): The prompt to display. Defaults to "".
+            password (bool, optional): Whether to read a password. Defaults to False.
+
+        Returns:
+            str: The line read from the input stream.
+
+        """
+        ...  # pragma: no cover
+
+
+@runtime_checkable
+@export_module("autogen.io")
+class IOStreamProtocol(InputStream, OutputStream, Protocol):
+    """A protocol for input/output streams."""
+
+
+@runtime_checkable
+@export_module("autogen.io")
+class AsyncIOStreamProtocol(AsyncInputStream, OutputStream, Protocol):
+    """A protocol for input/output streams."""
+
+
+iostream_union = Union[IOStreamProtocol, AsyncIOStreamProtocol]
+
+
+@export_module("autogen.io")
+class IOStream:
+    """A protocol for input/output streams."""
+
+    # ContextVar must be used in multithreaded or async environments
+    _default_io_stream: ContextVar[Optional[iostream_union]] = ContextVar("default_iostream", default=None)
+    _default_io_stream.set(None)
+    _global_default: Optional[iostream_union] = None
+
+    @staticmethod
+    def set_global_default(stream: iostream_union) -> None:
+        """Set the default input/output stream.
+
+        Args:
+            stream (IOStream): The input/output stream to set as the default.
+        """
+        IOStream._global_default = stream
+
+    @staticmethod
+    def get_global_default() -> iostream_union:
+        """Get the default input/output stream.
+
+        Returns:
+            IOStream: The default input/output stream.
+        """
+        if IOStream._global_default is None:
+            raise RuntimeError("No global default IOStream has been set")
+        return IOStream._global_default
+
+    @staticmethod
+    def get_default() -> iostream_union:
+        """Get the default input/output stream.
+
+        Returns:
+            IOStream: The default input/output stream.
+        """
+        iostream = IOStream._default_io_stream.get()
+        if iostream is None:
+            iostream = IOStream.get_global_default()
+            # Set the default IOStream of the current context (thread/cooroutine)
+            IOStream.set_default(iostream)
+        return iostream
+
+    @staticmethod
+    @contextmanager
+    def set_default(stream: Optional[iostream_union]) -> Iterator[None]:
+        """Set the default input/output stream.
+
+        Args:
+            stream (IOStream): The input/output stream to set as the default.
+        """
+        global _default_io_stream
+        try:
+            token = IOStream._default_io_stream.set(stream)
+            yield
+        finally:
+            IOStream._default_io_stream.reset(token)
+
+        return

autogen/io/console.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+import getpass
+from typing import Any
+
+from ..doc_utils import export_module
+from ..events.base_event import BaseEvent
+from ..events.print_event import PrintEvent
+from .base import IOStream
+
+__all__ = ("IOConsole",)
+
+
+@export_module("autogen.io")
+class IOConsole(IOStream):
+    """A console input/output stream."""
+
+    def print(self, *objects: Any, sep: str = " ", end: str = "\n", flush: bool = False) -> None:
+        """Print data to the output stream.
+
+        Args:
+            objects (any): The data to print.
+            sep (str, optional): The separator between objects. Defaults to " ".
+            end (str, optional): The end of the output. Defaults to "\n".
+            flush (bool, optional): Whether to flush the output. Defaults to False.
+        """
+        print_message = PrintEvent(*objects, sep=sep, end=end)
+        self.send(print_message)
+        # print(*objects, sep=sep, end=end, flush=flush)
+
+    def send(self, message: BaseEvent) -> None:
+        """Send a message to the output stream.
+
+        Args:
+            message (Any): The message to send.
+        """
+        message.print()
+
+    def input(self, prompt: str = "", *, password: bool = False) -> str:
+        """Read a line from the input stream.
+
+        Args:
+            prompt (str, optional): The prompt to display. Defaults to "".
+            password (bool, optional): Whether to read a password. Defaults to False.
+
+        Returns:
+            str: The line read from the input stream.
+
+        """
+        if password:
+            return getpass.getpass(prompt if prompt != "" else "Password: ")
+        return input(prompt)

autogen/io/processors/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+from .base import AsyncEventProcessorProtocol, EventProcessorProtocol
+from .console_event_processor import AsyncConsoleEventProcessor, ConsoleEventProcessor
+
+__all__ = [
+    "AsyncConsoleEventProcessor",
+    "AsyncEventProcessorProtocol",
+    "ConsoleEventProcessor",
+    "EventProcessorProtocol",
+]

autogen/io/processors/base.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+from typing import TYPE_CHECKING, Protocol
+
+from ...doc_utils import export_module
+
+if TYPE_CHECKING:
+    from ..run_response import AsyncRunResponseProtocol, RunResponseProtocol
+
+__all__ = ["AsyncEventProcessorProtocol", "EventProcessorProtocol"]
+
+
+@export_module("autogen.io")
+class EventProcessorProtocol(Protocol):
+    def process(self, response: "RunResponseProtocol") -> None: ...
+
+
+@export_module("autogen.io")
+class AsyncEventProcessorProtocol(Protocol):
+    async def process(self, response: "AsyncRunResponseProtocol") -> None: ...

autogen/io/processors/console_event_processor.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+import getpass
+from typing import TYPE_CHECKING
+
+from ...doc_utils import export_module
+from ...events.agent_events import InputRequestEvent
+from ...events.base_event import BaseEvent
+
+if TYPE_CHECKING:
+    from ..run_response import AsyncRunResponseProtocol, RunResponseProtocol
+    from .base import AsyncEventProcessorProtocol, EventProcessorProtocol
+
+
+@export_module("autogen.io")
+class ConsoleEventProcessor:
+    def process(self, response: "RunResponseProtocol") -> None:
+        for event in response.events:
+            self.process_event(event)
+
+    def process_event(self, event: BaseEvent) -> None:
+        if isinstance(event, InputRequestEvent):
+            prompt = event.content.prompt  # type: ignore[attr-defined]
+            if event.content.password:  # type: ignore[attr-defined]
+                result = getpass.getpass(prompt if prompt != "" else "Password: ")
+            result = input(prompt)
+            event.content.respond(result)  # type: ignore[attr-defined]
+        else:
+            event.print()
+
+
+@export_module("autogen.io")
+class AsyncConsoleEventProcessor:
+    async def process(self, response: "AsyncRunResponseProtocol") -> None:
+        async for event in response.events:
+            await self.process_event(event)
+
+    async def process_event(self, event: BaseEvent) -> None:
+        if isinstance(event, InputRequestEvent):
+            prompt = event.content.prompt  # type: ignore[attr-defined]
+            if event.content.password:  # type: ignore[attr-defined]
+                result = getpass.getpass(prompt if prompt != "" else "Password: ")
+            result = input(prompt)
+            await event.content.respond(result)  # type: ignore[attr-defined]
+        else:
+            event.print()
+
+
+if TYPE_CHECKING:
+
+    def check_type_1(x: ConsoleEventProcessor) -> EventProcessorProtocol:
+        return x
+
+    def check_type_2(x: AsyncConsoleEventProcessor) -> AsyncEventProcessorProtocol:
+        return x