64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import queue
|
|
from asyncio import Queue as AsyncQueue
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from autogen.io.base import AsyncIOStreamProtocol, IOStreamProtocol
|
|
|
|
from ..events.agent_events import InputRequestEvent
|
|
from ..events.print_event import PrintEvent
|
|
|
|
|
|
class ThreadIOStream:
|
|
def __init__(self) -> None:
|
|
self._input_stream: queue.Queue = queue.Queue() # type: ignore[type-arg]
|
|
self._output_stream: queue.Queue = queue.Queue() # type: ignore[type-arg]
|
|
|
|
def input(self, prompt: str = "", *, password: bool = False) -> str:
|
|
self.send(InputRequestEvent(prompt=prompt, password=password)) # type: ignore[call-arg]
|
|
return self._output_stream.get() # type: ignore[no-any-return]
|
|
|
|
def print(self, *objects: Any, sep: str = " ", end: str = "\n", flush: bool = False) -> None:
|
|
print_message = PrintEvent(*objects, sep=sep, end=end)
|
|
self.send(print_message)
|
|
|
|
def send(self, message: Any) -> None:
|
|
self._input_stream.put(message)
|
|
|
|
@property
|
|
def input_stream(self) -> queue.Queue: # type: ignore[type-arg]
|
|
return self._input_stream
|
|
|
|
|
|
class AsyncThreadIOStream:
|
|
def __init__(self) -> None:
|
|
self._input_stream: AsyncQueue = AsyncQueue() # type: ignore[type-arg]
|
|
self._output_stream: AsyncQueue = AsyncQueue() # type: ignore[type-arg]
|
|
|
|
async def input(self, prompt: str = "", *, password: bool = False) -> str:
|
|
self.send(InputRequestEvent(prompt=prompt, password=password)) # type: ignore[call-arg]
|
|
return await self._output_stream.get() # type: ignore[no-any-return]
|
|
|
|
def print(self, *objects: Any, sep: str = " ", end: str = "\n", flush: bool = False) -> None:
|
|
print_message = PrintEvent(*objects, sep=sep, end=end)
|
|
self.send(print_message)
|
|
|
|
def send(self, message: Any) -> None:
|
|
self._input_stream.put_nowait(message)
|
|
|
|
@property
|
|
def input_stream(self) -> AsyncQueue[Any]:
|
|
return self._input_stream
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
def check_type_1(x: ThreadIOStream) -> IOStreamProtocol:
|
|
return x
|
|
|
|
def check_type_2(x: AsyncThreadIOStream) -> AsyncIOStreamProtocol:
|
|
return x
|