152 lines
4.6 KiB
Python
152 lines
4.6 KiB
Python
# 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
|