chore: 添加虚拟环境到仓库
- 添加 backend_service/venv 虚拟环境 - 包含所有Python依赖包 - 注意：虚拟环境约393MB，包含12655个文件
This commit is contained in:
@@ -0,0 +1,24 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The agent base class."""
|
||||
from ._agent_base import AgentBase
|
||||
from ._react_agent_base import ReActAgentBase
|
||||
from ._react_agent import ReActAgent
|
||||
from ._user_input import (
|
||||
UserInputBase,
|
||||
UserInputData,
|
||||
TerminalUserInput,
|
||||
StudioUserInput,
|
||||
)
|
||||
from ._user_agent import UserAgent
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AgentBase",
|
||||
"ReActAgentBase",
|
||||
"ReActAgent",
|
||||
"UserInputData",
|
||||
"UserInputBase",
|
||||
"TerminalUserInput",
|
||||
"StudioUserInput",
|
||||
"UserAgent",
|
||||
]
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,703 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The agent base class in agentscope."""
|
||||
import asyncio
|
||||
import io
|
||||
import json
|
||||
from asyncio import Task, Queue
|
||||
from collections import OrderedDict
|
||||
from copy import deepcopy
|
||||
from typing import Callable, Any
|
||||
import base64
|
||||
import shortuuid
|
||||
import numpy as np
|
||||
from typing_extensions import deprecated
|
||||
|
||||
from ._agent_meta import _AgentMeta
|
||||
from .._logging import logger
|
||||
from ..module import StateModule
|
||||
from ..message import (
|
||||
Msg,
|
||||
AudioBlock,
|
||||
ToolUseBlock,
|
||||
ToolResultBlock,
|
||||
ImageBlock,
|
||||
VideoBlock,
|
||||
)
|
||||
from ..types import AgentHookTypes
|
||||
|
||||
|
||||
class AgentBase(StateModule, metaclass=_AgentMeta):
|
||||
"""Base class for asynchronous agents."""
|
||||
|
||||
id: str
|
||||
"""The agent's unique identifier, generated using shortuuid."""
|
||||
|
||||
supported_hook_types: list[str] = [
|
||||
"pre_reply",
|
||||
"post_reply",
|
||||
"pre_print",
|
||||
"post_print",
|
||||
"pre_observe",
|
||||
"post_observe",
|
||||
]
|
||||
"""Supported hook types for the agent base class."""
|
||||
|
||||
_class_pre_reply_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"AgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
],
|
||||
dict[str, Any] | None, # The modified kwargs or None
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level hook functions that will be called before the reply
|
||||
function, taking `self` object, the input arguments as input, and
|
||||
generating the modified arguments (if needed). Then input arguments of the
|
||||
reply function will be re-organized into a keyword arguments dictionary.
|
||||
If the one hook returns a new dictionary, the modified arguments will be
|
||||
passed to the next hook or the original reply function."""
|
||||
|
||||
_class_post_reply_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"AgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
Msg, # output, the output message
|
||||
],
|
||||
Msg | None,
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level hook functions that will be called after the reply
|
||||
function, which takes the `self` object and deep copied
|
||||
positional and keyword arguments (args and kwargs), and the output message
|
||||
as input. If the hook returns a message, the new message will be passed
|
||||
to the next hook or the original reply function. Otherwise, the original
|
||||
output will be passed instead."""
|
||||
|
||||
_class_pre_print_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"AgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
],
|
||||
dict[str, Any] | None, # The modified kwargs or None
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level hook functions that will be called before printing,
|
||||
which takes the `self` object, a deep copied arguments dictionary as input,
|
||||
and output the modified arguments (if needed). """
|
||||
|
||||
_class_post_print_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"AgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
Any, # output, `None` if no output
|
||||
],
|
||||
Any,
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level hook functions that will be called after the speak
|
||||
function, which takes the `self` object as input."""
|
||||
|
||||
_class_pre_observe_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"AgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
],
|
||||
dict[str, Any] | None, # The modified kwargs or None
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level hook functions that will be called before the observe
|
||||
function, which takes the `self` object and a deep copied input
|
||||
arguments dictionary as input. To change the input arguments, the hook
|
||||
function needs to output the modified arguments dictionary, which will be
|
||||
used as the input of the next hook function or the original observe
|
||||
function."""
|
||||
|
||||
_class_post_observe_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"AgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
None, # The output, `None` if no output
|
||||
],
|
||||
None,
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level hook functions that will be called after the observe
|
||||
function, which takes the `self` object as input."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize the agent."""
|
||||
super().__init__()
|
||||
|
||||
self.id = shortuuid.uuid()
|
||||
|
||||
# The replying task and identify of the current replying
|
||||
self._reply_task: Task | None = None
|
||||
self._reply_id: str | None = None
|
||||
|
||||
# Initialize the instance-level hooks
|
||||
self._instance_pre_print_hooks = OrderedDict()
|
||||
self._instance_post_print_hooks = OrderedDict()
|
||||
|
||||
self._instance_pre_reply_hooks = OrderedDict()
|
||||
self._instance_post_reply_hooks = OrderedDict()
|
||||
|
||||
self._instance_pre_observe_hooks = OrderedDict()
|
||||
self._instance_post_observe_hooks = OrderedDict()
|
||||
|
||||
# The prefix used in streaming printing, which will save the
|
||||
# accumulated text and audio streaming data for each message id.
|
||||
# e.g. {"text": "xxx", "audio": (stream_obj, "{base64_data}")}
|
||||
self._stream_prefix = {}
|
||||
|
||||
# The subscribers that will receive the reply message by their
|
||||
# `observe` method. The key is the MsgHub id, and the value is the
|
||||
# list of agents.
|
||||
self._subscribers: dict[str, list[AgentBase]] = {}
|
||||
|
||||
# We add this variable in case developers want to disable the console
|
||||
# output of the agent, e.g., in a production environment.
|
||||
self._disable_console_output: bool = False
|
||||
|
||||
# The streaming message queue used to export the messages as a
|
||||
# generator
|
||||
self._disable_msg_queue: bool = True
|
||||
self.msg_queue = None
|
||||
|
||||
async def observe(self, msg: Msg | list[Msg] | None) -> None:
|
||||
"""Receive the given message(s) without generating a reply.
|
||||
|
||||
Args:
|
||||
msg (`Msg | list[Msg] | None`):
|
||||
The message(s) to be observed.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"The observe function is not implemented in"
|
||||
f" {self.__class__.__name__} class.",
|
||||
)
|
||||
|
||||
async def reply(self, *args: Any, **kwargs: Any) -> Msg:
|
||||
"""The main logic of the agent, which generates a reply based on the
|
||||
current state and input arguments."""
|
||||
raise NotImplementedError(
|
||||
"The reply function is not implemented in "
|
||||
f"{self.__class__.__name__} class.",
|
||||
)
|
||||
|
||||
async def print(self, msg: Msg, last: bool = True) -> None:
|
||||
"""The function to display the message.
|
||||
|
||||
Args:
|
||||
msg (`Msg`):
|
||||
The message object to be printed.
|
||||
last (`bool`, defaults to `True`):
|
||||
Whether this is the last one in streaming messages. For
|
||||
non-streaming message, this should always be `True`.
|
||||
"""
|
||||
if not self._disable_msg_queue:
|
||||
await self.msg_queue.put((deepcopy(msg), last))
|
||||
|
||||
if self._disable_console_output:
|
||||
return
|
||||
|
||||
# The accumulated textual content to print, including the text blocks
|
||||
# and the thinking blocks
|
||||
thinking_and_text_to_print = []
|
||||
|
||||
for block in msg.get_content_blocks():
|
||||
if block["type"] == "audio":
|
||||
self._process_audio_block(msg.id, block)
|
||||
|
||||
elif block["type"] == "text":
|
||||
self._print_text_block(
|
||||
msg.id,
|
||||
name_prefix=msg.name,
|
||||
text_content=block["text"],
|
||||
thinking_and_text_to_print=thinking_and_text_to_print,
|
||||
)
|
||||
|
||||
elif block["type"] == "thinking":
|
||||
self._print_text_block(
|
||||
msg.id,
|
||||
name_prefix=f"{msg.name}(thinking)",
|
||||
text_content=block["thinking"],
|
||||
thinking_and_text_to_print=thinking_and_text_to_print,
|
||||
)
|
||||
|
||||
elif last:
|
||||
self._print_last_block(block, msg)
|
||||
|
||||
# Clean up resources if this is the last message in streaming
|
||||
if last and msg.id in self._stream_prefix:
|
||||
if "audio" in self._stream_prefix[msg.id]:
|
||||
player, _ = self._stream_prefix[msg.id]["audio"]
|
||||
# Close the miniaudio player
|
||||
player.close()
|
||||
stream_prefix = self._stream_prefix.pop(msg.id)
|
||||
if "text" in stream_prefix and not stream_prefix["text"].endswith(
|
||||
"\n",
|
||||
):
|
||||
print()
|
||||
|
||||
def _process_audio_block(
|
||||
self,
|
||||
msg_id: str,
|
||||
audio_block: AudioBlock,
|
||||
) -> None:
|
||||
"""Process audio block content.
|
||||
|
||||
Args:
|
||||
msg_id (`str`):
|
||||
The unique identifier of the message
|
||||
audio_block (`AudioBlock`):
|
||||
The audio content block
|
||||
"""
|
||||
if "source" not in audio_block:
|
||||
raise ValueError(
|
||||
"The audio block must contain the 'source' field.",
|
||||
)
|
||||
|
||||
if audio_block["source"]["type"] == "url":
|
||||
import urllib.request
|
||||
import wave
|
||||
import sounddevice as sd
|
||||
|
||||
url = audio_block["source"]["url"]
|
||||
try:
|
||||
with urllib.request.urlopen(url) as response:
|
||||
audio_data = response.read()
|
||||
|
||||
with wave.open(io.BytesIO(audio_data), "rb") as wf:
|
||||
samplerate = wf.getframerate()
|
||||
n_frames = wf.getnframes()
|
||||
audio_frames = wf.readframes(n_frames)
|
||||
|
||||
# Convert byte data to numpy array
|
||||
audio_np = np.frombuffer(audio_frames, dtype=np.int16)
|
||||
|
||||
# Play audio
|
||||
sd.play(audio_np, samplerate)
|
||||
sd.wait()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to play audio from url %s: %s",
|
||||
url,
|
||||
str(e),
|
||||
)
|
||||
|
||||
elif audio_block["source"]["type"] == "base64":
|
||||
data = audio_block["source"]["data"]
|
||||
|
||||
if msg_id not in self._stream_prefix:
|
||||
self._stream_prefix[msg_id] = {}
|
||||
|
||||
audio_prefix = self._stream_prefix[msg_id].get("audio", None)
|
||||
|
||||
import sounddevice as sd
|
||||
|
||||
# The player and the prefix data is cached for streaming audio
|
||||
if audio_prefix:
|
||||
player, audio_prefix_data = audio_prefix
|
||||
else:
|
||||
player = sd.OutputStream(
|
||||
samplerate=24000,
|
||||
channels=1,
|
||||
dtype=np.float32,
|
||||
blocksize=1024,
|
||||
latency="low",
|
||||
)
|
||||
player.start()
|
||||
audio_prefix_data = ""
|
||||
|
||||
# play the audio data
|
||||
new_audio_data = data[len(audio_prefix_data) :]
|
||||
if new_audio_data:
|
||||
audio_bytes = base64.b64decode(new_audio_data)
|
||||
audio_np = np.frombuffer(audio_bytes, dtype=np.int16)
|
||||
audio_float = audio_np.astype(np.float32) / 32768.0
|
||||
|
||||
# Write to the audio output stream
|
||||
player.write(audio_float)
|
||||
|
||||
# save the player and the prefix data
|
||||
self._stream_prefix[msg_id]["audio"] = (
|
||||
player,
|
||||
data,
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
"Unsupported audio source type: "
|
||||
f"{audio_block['source']['type']}",
|
||||
)
|
||||
|
||||
def _print_text_block(
|
||||
self,
|
||||
msg_id: str,
|
||||
name_prefix: str,
|
||||
text_content: str,
|
||||
thinking_and_text_to_print: list[str],
|
||||
) -> None:
|
||||
"""Print the text block and thinking block content.
|
||||
|
||||
Args:
|
||||
msg_id (`str`):
|
||||
The unique identifier of the message
|
||||
name_prefix (`str`):
|
||||
The prefix for the message, e.g. "{name}: " for text block and
|
||||
"{name}(thinking): " for thinking block.
|
||||
text_content (`str`):
|
||||
The textual content to be printed.
|
||||
thinking_and_text_to_print (`list[str]`):
|
||||
A list of textual content to be printed together. Here we
|
||||
gather the text and thinking blocks to print them together.
|
||||
"""
|
||||
thinking_and_text_to_print.append(
|
||||
f"{name_prefix}: {text_content}",
|
||||
)
|
||||
# The accumulated text and thinking blocks to print
|
||||
to_print = "\n".join(thinking_and_text_to_print)
|
||||
|
||||
# The text prefix that has been printed
|
||||
if msg_id not in self._stream_prefix:
|
||||
self._stream_prefix[msg_id] = {}
|
||||
|
||||
text_prefix = self._stream_prefix[msg_id].get("text", "")
|
||||
|
||||
# Only print when there is new text content
|
||||
if len(to_print) > len(text_prefix):
|
||||
print(to_print[len(text_prefix) :], end="")
|
||||
|
||||
# Save the printed text prefix
|
||||
self._stream_prefix[msg_id]["text"] = to_print
|
||||
|
||||
def _print_last_block(
|
||||
self,
|
||||
block: ToolUseBlock | ToolResultBlock | ImageBlock | VideoBlock,
|
||||
msg: Msg,
|
||||
) -> None:
|
||||
"""Process and print the last content block, and the block type
|
||||
is not audio, text, or thinking.
|
||||
|
||||
Args:
|
||||
block (`ToolUseBlock | ToolResultBlock | ImageBlock | VideoBlock`):
|
||||
The content block to be printed
|
||||
msg (`Msg`):
|
||||
The message object
|
||||
"""
|
||||
text_prefix = self._stream_prefix.get(msg.id, {}).get("text", "")
|
||||
|
||||
if text_prefix:
|
||||
# Add a newline to separate from previous text content
|
||||
print_newline = "" if text_prefix.endswith("\n") else "\n"
|
||||
print(
|
||||
f"{print_newline}"
|
||||
f"{json.dumps(block, indent=4, ensure_ascii=False)}",
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"{msg.name}:"
|
||||
f" {json.dumps(block, indent=4, ensure_ascii=False)}",
|
||||
)
|
||||
|
||||
async def __call__(self, *args: Any, **kwargs: Any) -> Msg:
|
||||
"""Call the reply function with the given arguments."""
|
||||
self._reply_id = shortuuid.uuid()
|
||||
|
||||
reply_msg: Msg | None = None
|
||||
try:
|
||||
self._reply_task = asyncio.current_task()
|
||||
reply_msg = await self.reply(*args, **kwargs)
|
||||
|
||||
# The interruption is triggered by calling the interrupt method
|
||||
except asyncio.CancelledError:
|
||||
reply_msg = await self.handle_interrupt(*args, **kwargs)
|
||||
|
||||
finally:
|
||||
# Broadcast the reply message to all subscribers
|
||||
if reply_msg:
|
||||
await self._broadcast_to_subscribers(reply_msg)
|
||||
self._reply_task = None
|
||||
|
||||
return reply_msg
|
||||
|
||||
async def _broadcast_to_subscribers(
|
||||
self,
|
||||
msg: Msg | list[Msg] | None,
|
||||
) -> None:
|
||||
"""Broadcast the message to all subscribers."""
|
||||
for subscribers in self._subscribers.values():
|
||||
for subscriber in subscribers:
|
||||
await subscriber.observe(msg)
|
||||
|
||||
async def handle_interrupt(
|
||||
self,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Msg:
|
||||
"""The post-processing logic when the reply is interrupted by the
|
||||
user or something else."""
|
||||
raise NotImplementedError(
|
||||
f"The handle_interrupt function is not implemented in "
|
||||
f"{self.__class__.__name__}",
|
||||
)
|
||||
|
||||
async def interrupt(self, msg: Msg | list[Msg] | None = None) -> None:
|
||||
"""Interrupt the current reply process."""
|
||||
if self._reply_task and not self._reply_task.done():
|
||||
self._reply_task.cancel(msg)
|
||||
|
||||
def register_instance_hook(
|
||||
self,
|
||||
hook_type: AgentHookTypes,
|
||||
hook_name: str,
|
||||
hook: Callable,
|
||||
) -> None:
|
||||
"""Register a hook to the agent instance, which only takes effect
|
||||
for the current instance.
|
||||
|
||||
Args:
|
||||
hook_type (`str`):
|
||||
The type of the hook, indicating where the hook is to be
|
||||
triggered.
|
||||
hook_name (`str`):
|
||||
The name of the hook. If the name is already registered, the
|
||||
hook will be overwritten.
|
||||
hook (`Callable`):
|
||||
The hook function.
|
||||
"""
|
||||
if not isinstance(self, AgentBase):
|
||||
raise TypeError(
|
||||
"The register_instance_hook method should be called on an "
|
||||
f"instance of AsyncAgentBase, but got {self} of "
|
||||
f"type {type(self)}.",
|
||||
)
|
||||
hooks = getattr(self, f"_instance_{hook_type}_hooks")
|
||||
hooks[hook_name] = hook
|
||||
|
||||
def remove_instance_hook(
|
||||
self,
|
||||
hook_type: AgentHookTypes,
|
||||
hook_name: str,
|
||||
) -> None:
|
||||
"""Remove an instance-level hook from the agent instance.
|
||||
|
||||
Args:
|
||||
hook_type (`AgentHookTypes`):
|
||||
The type of the hook, indicating where the hook is to be
|
||||
triggered.
|
||||
hook_name (`str`):
|
||||
The name of the hook to remove.
|
||||
"""
|
||||
if not isinstance(self, AgentBase):
|
||||
raise TypeError(
|
||||
"The remove_instance_hook method should be called on an "
|
||||
f"instance of AsyncAgentBase, but got {self} of "
|
||||
f"type {type(self)}.",
|
||||
)
|
||||
hooks = getattr(self, f"_instance_{hook_type}_hooks")
|
||||
if hook_name in hooks:
|
||||
del hooks[hook_name]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Hook '{hook_name}' not found in '{hook_type}' hooks of "
|
||||
f"{self.__class__.__name__} instance.",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def register_class_hook(
|
||||
cls,
|
||||
hook_type: AgentHookTypes,
|
||||
hook_name: str,
|
||||
hook: Callable,
|
||||
) -> None:
|
||||
"""The universal function to register a hook to the agent class, which
|
||||
will take effect for all instances of the class.
|
||||
|
||||
Args:
|
||||
hook_type (`AgentHookTypes`):
|
||||
The type of the hook, indicating where the hook is to be
|
||||
triggered.
|
||||
hook_name (`str`):
|
||||
The name of the hook. If the name is already registered, the
|
||||
hook will be overwritten.
|
||||
hook (`Callable`):
|
||||
The hook function.
|
||||
"""
|
||||
|
||||
assert (
|
||||
hook_type in cls.supported_hook_types
|
||||
), f"Invalid hook type: {hook_type}"
|
||||
|
||||
hooks = getattr(cls, f"_class_{hook_type}_hooks")
|
||||
hooks[hook_name] = hook
|
||||
|
||||
@classmethod
|
||||
def remove_class_hook(
|
||||
cls,
|
||||
hook_type: AgentHookTypes,
|
||||
hook_name: str,
|
||||
) -> None:
|
||||
"""Remove a class-level hook from the agent class.
|
||||
|
||||
Args:
|
||||
hook_type (`AgentHookTypes`):
|
||||
The type of the hook, indicating where the hook is to be
|
||||
triggered.
|
||||
hook_name (`str`):
|
||||
The name of the hook to remove.
|
||||
"""
|
||||
|
||||
assert (
|
||||
hook_type in cls.supported_hook_types
|
||||
), f"Invalid hook type: {hook_type}"
|
||||
hooks = getattr(cls, f"_class_{hook_type}_hooks")
|
||||
if hook_name in hooks:
|
||||
del hooks[hook_name]
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Hook '{hook_name}' not found in '{hook_type}' hooks of "
|
||||
f"{cls.__name__} class.",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def clear_class_hooks(
|
||||
cls,
|
||||
hook_type: AgentHookTypes | None = None,
|
||||
) -> None:
|
||||
"""Clear all class-level hooks.
|
||||
|
||||
Args:
|
||||
hook_type (`AgentHookTypes`, optional):
|
||||
The type of the hook to clear. If not specified, all
|
||||
class-level hooks will be cleared.
|
||||
"""
|
||||
|
||||
if hook_type is None:
|
||||
for typ in cls.supported_hook_types:
|
||||
hooks = getattr(cls, f"_class_{typ}_hooks")
|
||||
hooks.clear()
|
||||
else:
|
||||
assert (
|
||||
hook_type in cls.supported_hook_types
|
||||
), f"Invalid hook type: {hook_type}"
|
||||
hooks = getattr(cls, f"_class_{hook_type}_hooks")
|
||||
hooks.clear()
|
||||
|
||||
def clear_instance_hooks(
|
||||
self,
|
||||
hook_type: AgentHookTypes | None = None,
|
||||
) -> None:
|
||||
"""If `hook_type` is not specified, clear all instance-level hooks.
|
||||
Otherwise, clear the specified type of instance-level hooks."""
|
||||
if hook_type is None:
|
||||
for typ in self.supported_hook_types:
|
||||
if not hasattr(self, f"_instance_{typ}_hooks"):
|
||||
raise ValueError(
|
||||
f"Call super().__init__() in the constructor "
|
||||
f"to initialize the instance-level hooks for "
|
||||
f"{self.__class__.__name__}.",
|
||||
)
|
||||
hooks = getattr(self, f"_instance_{typ}_hooks")
|
||||
hooks.clear()
|
||||
|
||||
else:
|
||||
assert (
|
||||
hook_type in self.supported_hook_types
|
||||
), f"Invalid hook type: {hook_type}"
|
||||
if not hasattr(self, f"_instance_{hook_type}_hooks"):
|
||||
raise ValueError(
|
||||
f"Call super().__init__() in the constructor "
|
||||
f"to initialize the instance-level hooks for "
|
||||
f"{self.__class__.__name__}.",
|
||||
)
|
||||
hooks = getattr(self, f"_instance_{hook_type}_hooks")
|
||||
hooks.clear()
|
||||
|
||||
def reset_subscribers(
|
||||
self,
|
||||
msghub_name: str,
|
||||
subscribers: list["AgentBase"],
|
||||
) -> None:
|
||||
"""Reset the subscribers of the agent.
|
||||
|
||||
Args:
|
||||
msghub_name (`str`):
|
||||
The name of the MsgHub that manages the subscribers.
|
||||
subscribers (`list[AgentBase]`):
|
||||
A list of agents that will receive the reply message from
|
||||
this agent via their `observe` method.
|
||||
"""
|
||||
self._subscribers[msghub_name] = [_ for _ in subscribers if _ != self]
|
||||
|
||||
def remove_subscribers(self, msghub_name: str) -> None:
|
||||
"""Remove the msghub subscribers by the given msg hub name.
|
||||
|
||||
Args:
|
||||
msghub_name (`str`):
|
||||
The name of the MsgHub that manages the subscribers.
|
||||
"""
|
||||
if msghub_name not in self._subscribers:
|
||||
logger.warning(
|
||||
"MsgHub named '%s' not found",
|
||||
msghub_name,
|
||||
)
|
||||
else:
|
||||
self._subscribers.pop(msghub_name)
|
||||
|
||||
@deprecated("Please use set_console_output_enabled() instead.")
|
||||
def disable_console_output(self) -> None:
|
||||
"""This function will disable the console output of the agent, e.g.
|
||||
in a production environment to avoid messy logs."""
|
||||
self._disable_console_output = True
|
||||
|
||||
def set_console_output_enabled(self, enabled: bool) -> None:
|
||||
"""Enable or disable the console output of the agent. E.g. in a
|
||||
production environment, you may want to disable the console output to
|
||||
avoid messy logs.
|
||||
|
||||
Args:
|
||||
enabled (`bool`):
|
||||
If `True`, enable the console output. If `False`, disable
|
||||
the console output.
|
||||
"""
|
||||
self._disable_console_output = not enabled
|
||||
|
||||
def set_msg_queue_enabled(
|
||||
self,
|
||||
enabled: bool,
|
||||
queue: Queue | None = None,
|
||||
) -> None:
|
||||
"""Enable or disable the message queue for streaming outputs.
|
||||
|
||||
Args:
|
||||
enabled (`bool`):
|
||||
If `True`, enable the message queue to allow streaming
|
||||
outputs. If `False`, disable the message queue.
|
||||
queue (`Queue | None`, optional):
|
||||
The queue instance that will be used to initialize the
|
||||
message queue when `enable` is `True`.
|
||||
"""
|
||||
if enabled:
|
||||
if queue is None:
|
||||
if self.msg_queue is None:
|
||||
self.msg_queue = asyncio.Queue(maxsize=100)
|
||||
else:
|
||||
self.msg_queue = queue
|
||||
else:
|
||||
self.msg_queue = None
|
||||
|
||||
self._disable_msg_queue = not enabled
|
||||
@@ -0,0 +1,180 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The metaclass for agents in agentscope."""
|
||||
import inspect
|
||||
from copy import deepcopy
|
||||
from functools import wraps
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
TYPE_CHECKING,
|
||||
Callable,
|
||||
)
|
||||
|
||||
from .._utils._common import _execute_async_or_sync_func
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._agent_base import AgentBase
|
||||
else:
|
||||
AgentBase = "AgentBase"
|
||||
|
||||
|
||||
def _normalize_to_kwargs(
|
||||
func: Callable,
|
||||
self: Any,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> dict:
|
||||
"""Normalize the provided positional and keyword arguments into a
|
||||
keyword arguments dictionary that matches the function signature."""
|
||||
sig = inspect.signature(func)
|
||||
try:
|
||||
# Bind the provided arguments to the function signature
|
||||
bound = sig.bind(self, *args, **kwargs)
|
||||
# Apply the default values for parameters
|
||||
bound.apply_defaults()
|
||||
|
||||
# Return the arguments in a dictionary format
|
||||
res = dict(bound.arguments)
|
||||
res.pop("self")
|
||||
return res
|
||||
|
||||
except TypeError as e:
|
||||
# If failed to bind, we raise a TypeError with more context
|
||||
param_names = list(sig.parameters.keys())
|
||||
provided_args = len(args)
|
||||
provided_kwargs = list(kwargs.keys())
|
||||
|
||||
raise TypeError(
|
||||
f"Failed to bind parameters for function '{func.__name__}': {e}\n"
|
||||
f"Expected parameters: {param_names}\n"
|
||||
f"Provided {provided_args} positional args and kwargs: "
|
||||
f"{provided_kwargs}",
|
||||
) from e
|
||||
|
||||
|
||||
def _wrap_with_hooks(
|
||||
original_func: Callable,
|
||||
) -> Callable:
|
||||
"""A decorator to wrap the original async function with pre- and post-hooks
|
||||
|
||||
Args:
|
||||
original_func (`Callable`):
|
||||
The original async function to be wrapped with hooks.
|
||||
"""
|
||||
func_name = original_func.__name__.replace("_", "")
|
||||
|
||||
@wraps(original_func)
|
||||
async def async_wrapper(
|
||||
self: AgentBase,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""The wrapped function, which call the pre- and post-hooks before and
|
||||
after the original function."""
|
||||
|
||||
# Unify all positional and keyword arguments into a keyword arguments
|
||||
normalized_kwargs = _normalize_to_kwargs(
|
||||
original_func,
|
||||
self,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
current_normalized_kwargs = normalized_kwargs
|
||||
assert (
|
||||
hasattr(self, f"_instance_pre_{func_name}_hooks")
|
||||
and hasattr(self, f"_instance_post_{func_name}_hooks")
|
||||
and hasattr(self.__class__, f"_class_pre_{func_name}_hooks")
|
||||
and hasattr(self.__class__, f"_class_post_{func_name}_hooks")
|
||||
), f"Hooks for {func_name} not found in {self.__class__.__name__}"
|
||||
|
||||
# pre-hooks
|
||||
pre_hooks = list(
|
||||
getattr(self, f"_instance_pre_{func_name}_hooks").values(),
|
||||
) + list(
|
||||
getattr(self, f"_class_pre_{func_name}_hooks").values(),
|
||||
)
|
||||
for pre_hook in pre_hooks:
|
||||
modified_keywords = await _execute_async_or_sync_func(
|
||||
pre_hook,
|
||||
self,
|
||||
deepcopy(current_normalized_kwargs),
|
||||
)
|
||||
if modified_keywords is not None:
|
||||
assert isinstance(modified_keywords, dict), (
|
||||
f"Pre-hook must return a dict of keyword arguments, rather"
|
||||
f" than {type(modified_keywords)} from hook "
|
||||
f"{pre_hook.__name__}"
|
||||
)
|
||||
current_normalized_kwargs = modified_keywords
|
||||
|
||||
# original function
|
||||
# handle positional and keyword arguments specifically
|
||||
args = current_normalized_kwargs.get("args", [])
|
||||
kwargs = current_normalized_kwargs.get("kwargs", {})
|
||||
others = {
|
||||
k: v
|
||||
for k, v in current_normalized_kwargs.items()
|
||||
if k not in ["args", "kwargs"]
|
||||
}
|
||||
current_output = await original_func(
|
||||
self,
|
||||
*args,
|
||||
**others,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# post_hooks
|
||||
post_hooks = list(
|
||||
getattr(self, f"_instance_post_{func_name}_hooks").values(),
|
||||
) + list(
|
||||
getattr(self, f"_class_post_{func_name}_hooks").values(),
|
||||
)
|
||||
for post_hook in post_hooks:
|
||||
modified_output = await _execute_async_or_sync_func(
|
||||
post_hook,
|
||||
self,
|
||||
deepcopy(current_normalized_kwargs),
|
||||
deepcopy(current_output),
|
||||
)
|
||||
if modified_output is not None:
|
||||
current_output = modified_output
|
||||
return current_output
|
||||
|
||||
return async_wrapper
|
||||
|
||||
|
||||
class _AgentMeta(type):
|
||||
"""The agent metaclass that wraps the agent's reply, observe and print
|
||||
functions with pre- and post-hooks."""
|
||||
|
||||
def __new__(mcs, name: Any, bases: Any, attrs: Dict) -> Any:
|
||||
"""Wrap the agent's functions with hooks."""
|
||||
|
||||
for func_name in [
|
||||
"reply",
|
||||
"print",
|
||||
"observe",
|
||||
]:
|
||||
if func_name in attrs:
|
||||
attrs[func_name] = _wrap_with_hooks(attrs[func_name])
|
||||
|
||||
return super().__new__(mcs, name, bases, attrs)
|
||||
|
||||
|
||||
class _ReActAgentMeta(_AgentMeta):
|
||||
"""The ReAct metaclass that adds pre- and post-hooks for the _reasoning
|
||||
and _acting functions."""
|
||||
|
||||
def __new__(mcs, name: Any, bases: Any, attrs: Dict) -> Any:
|
||||
"""Wrap the ReAct agent's _reasoning and _acting functions with
|
||||
hooks."""
|
||||
|
||||
for func_name in [
|
||||
"_reasoning",
|
||||
"_acting",
|
||||
]:
|
||||
if func_name in attrs:
|
||||
attrs[func_name] = _wrap_with_hooks(attrs[func_name])
|
||||
|
||||
return super().__new__(mcs, name, bases, attrs)
|
||||
@@ -0,0 +1,767 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# pylint: disable=not-an-iterable
|
||||
# mypy: disable-error-code="list-item"
|
||||
"""ReAct agent class in agentscope."""
|
||||
import asyncio
|
||||
from typing import Type, Any, AsyncGenerator, Literal
|
||||
|
||||
import shortuuid
|
||||
from pydantic import BaseModel, ValidationError, Field
|
||||
|
||||
from ._react_agent_base import ReActAgentBase
|
||||
from .._logging import logger
|
||||
from ..formatter import FormatterBase
|
||||
from ..memory import MemoryBase, LongTermMemoryBase, InMemoryMemory
|
||||
from ..message import Msg, ToolUseBlock, ToolResultBlock, TextBlock
|
||||
from ..model import ChatModelBase
|
||||
from ..rag import KnowledgeBase, Document
|
||||
from ..plan import PlanNotebook
|
||||
from ..tool import Toolkit, ToolResponse
|
||||
from ..tracing import trace_reply
|
||||
|
||||
|
||||
class _QueryRewriteModel(BaseModel):
|
||||
"""The structured model used for query rewriting."""
|
||||
|
||||
rewritten_query: str = Field(
|
||||
description=(
|
||||
"The rewritten query, which should be specific and concise. "
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def finish_function_pre_print_hook(
|
||||
self: "ReActAgent",
|
||||
kwargs: dict[str, Any],
|
||||
) -> dict[str, Any] | None:
|
||||
"""A pre-speak hook function that check if finish_function is called. If
|
||||
so, it will wrap the response argument into a message and return it to
|
||||
replace the original message. By this way, the calling of the finish
|
||||
function will be displayed as a text reply instead of a tool call."""
|
||||
|
||||
msg = kwargs["msg"]
|
||||
|
||||
if isinstance(msg.content, str):
|
||||
return None
|
||||
|
||||
if isinstance(msg.content, list):
|
||||
for i, block in enumerate(msg.content):
|
||||
if (
|
||||
block["type"] == "tool_use"
|
||||
and block["name"] == self.finish_function_name
|
||||
):
|
||||
# Convert the response argument into a text block for
|
||||
# displaying
|
||||
try:
|
||||
msg.content[i] = TextBlock(
|
||||
type="text",
|
||||
text=block["input"].get("response", ""),
|
||||
)
|
||||
return kwargs
|
||||
except Exception:
|
||||
print("Error in block input", block["input"])
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class ReActAgent(ReActAgentBase):
|
||||
"""A ReAct agent implementation in AgentScope, which supports
|
||||
|
||||
- Realtime steering
|
||||
- API-based (parallel) tool calling
|
||||
- Hooks around reasoning, acting, reply, observe and print functions
|
||||
- Structured output generation
|
||||
"""
|
||||
|
||||
finish_function_name: str = "generate_response"
|
||||
"""The function name used to finish replying and return a response to
|
||||
the user."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
sys_prompt: str,
|
||||
model: ChatModelBase,
|
||||
formatter: FormatterBase,
|
||||
toolkit: Toolkit | None = None,
|
||||
memory: MemoryBase | None = None,
|
||||
long_term_memory: LongTermMemoryBase | None = None,
|
||||
long_term_memory_mode: Literal[
|
||||
"agent_control",
|
||||
"static_control",
|
||||
"both",
|
||||
] = "both",
|
||||
enable_meta_tool: bool = False,
|
||||
parallel_tool_calls: bool = False,
|
||||
knowledge: KnowledgeBase | list[KnowledgeBase] | None = None,
|
||||
enable_rewrite_query: bool = True,
|
||||
plan_notebook: PlanNotebook | None = None,
|
||||
print_hint_msg: bool = False,
|
||||
max_iters: int = 10,
|
||||
) -> None:
|
||||
"""Initialize the ReAct agent
|
||||
|
||||
Args:
|
||||
name (`str`):
|
||||
The name of the agent.
|
||||
sys_prompt (`str`):
|
||||
The system prompt of the agent.
|
||||
model (`ChatModelBase`):
|
||||
The chat model used by the agent.
|
||||
formatter (`FormatterBase`):
|
||||
The formatter used to format the messages into the required
|
||||
format of the model API provider.
|
||||
toolkit (`Toolkit | None`, optional):
|
||||
A `Toolkit` object that contains the tool functions. If not
|
||||
provided, a default empty `Toolkit` will be created.
|
||||
memory (`MemoryBase | None`, optional):
|
||||
The memory used to store the dialogue history. If not provided,
|
||||
a default `InMemoryMemory` will be created, which stores
|
||||
messages in a list in memory.
|
||||
long_term_memory (`LongTermMemoryBase | None`, optional):
|
||||
The optional long-term memory, which will provide two tool
|
||||
functions: `retrieve_from_memory` and `record_to_memory`, and
|
||||
will attach the retrieved information to the system prompt
|
||||
before each reply.
|
||||
enable_meta_tool (`bool`, defaults to `False`):
|
||||
If `True`, a meta tool function `reset_equipped_tools` will be
|
||||
added to the toolkit, which allows the agent to manage its
|
||||
equipped tools dynamically.
|
||||
long_term_memory_mode (`Literal['agent_control', 'static_control',\
|
||||
'both']`, defaults to `both`):
|
||||
The mode of the long-term memory. If `agent_control`, two
|
||||
tool functions `retrieve_from_memory` and `record_to_memory`
|
||||
will be registered in the toolkit to allow the agent to
|
||||
manage the long-term memory. If `static_control`, retrieving
|
||||
and recording will happen in the beginning and end of
|
||||
each reply respectively.
|
||||
parallel_tool_calls (`bool`, defaults to `False`):
|
||||
When LLM generates multiple tool calls, whether to execute
|
||||
them in parallel.
|
||||
knowledge (`KnowledgeBase | list[KnowledgeBase] | None`, optional):
|
||||
The knowledge object(s) used by the agent to retrieve
|
||||
relevant documents at the beginning of each reply.
|
||||
enable_rewrite_query (`bool`, defaults to `True`):
|
||||
Whether ask the agent to rewrite the user input query before
|
||||
retrieving from the knowledge base(s), e.g. rewrite "Who am I"
|
||||
to "{user's name}" to get more relevant documents. Only works
|
||||
when the knowledge base(s) is provided.
|
||||
plan_notebook (`PlanNotebook | None`, optional):
|
||||
The plan notebook instance, allow the agent to finish the
|
||||
complex task by decomposing it into a sequence of subtasks.
|
||||
print_hint_msg (`bool`, defaults to `False`):
|
||||
Whether to print the hint messages, including the reasoning
|
||||
hint from the plan notebook, the retrieved information from
|
||||
the long-term memory and knowledge base(s).
|
||||
max_iters (`int`, defaults to `10`):
|
||||
The maximum number of iterations of the reasoning-acting loops.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
assert long_term_memory_mode in [
|
||||
"agent_control",
|
||||
"static_control",
|
||||
"both",
|
||||
]
|
||||
|
||||
# Static variables in the agent
|
||||
self.name = name
|
||||
self._sys_prompt = sys_prompt
|
||||
self.max_iters = max_iters
|
||||
self.model = model
|
||||
self.formatter = formatter
|
||||
|
||||
# -------------- Memory management --------------
|
||||
# Record the dialogue history in the memory
|
||||
self.memory = memory or InMemoryMemory()
|
||||
# If provide the long-term memory, it will be used to retrieve info
|
||||
# in the beginning of each reply, and the result will be added to the
|
||||
# system prompt
|
||||
self.long_term_memory = long_term_memory
|
||||
|
||||
# The long-term memory mode
|
||||
self._static_control = long_term_memory and long_term_memory_mode in [
|
||||
"static_control",
|
||||
"both",
|
||||
]
|
||||
self._agent_control = long_term_memory and long_term_memory_mode in [
|
||||
"agent_control",
|
||||
"both",
|
||||
]
|
||||
|
||||
# -------------- Tool management --------------
|
||||
# If None, a default Toolkit will be created
|
||||
self.toolkit = toolkit or Toolkit()
|
||||
self.toolkit.register_tool_function(
|
||||
getattr(self, self.finish_function_name),
|
||||
)
|
||||
if self._agent_control:
|
||||
# Adding two tool functions into the toolkit to allow self-control
|
||||
self.toolkit.register_tool_function(
|
||||
long_term_memory.record_to_memory,
|
||||
)
|
||||
self.toolkit.register_tool_function(
|
||||
long_term_memory.retrieve_from_memory,
|
||||
)
|
||||
# Add a meta tool function to allow agent-controlled tool management
|
||||
if enable_meta_tool or plan_notebook:
|
||||
self.toolkit.register_tool_function(
|
||||
self.toolkit.reset_equipped_tools,
|
||||
)
|
||||
|
||||
self.parallel_tool_calls = parallel_tool_calls
|
||||
|
||||
# -------------- RAG management --------------
|
||||
# The knowledge base(s) used by the agent
|
||||
if isinstance(knowledge, KnowledgeBase):
|
||||
knowledge = [knowledge]
|
||||
self.knowledge: list[KnowledgeBase] = knowledge or []
|
||||
self.enable_rewrite_query = enable_rewrite_query
|
||||
|
||||
# -------------- Plan management --------------
|
||||
# Equipped the plan-related tools provided by the plan notebook as
|
||||
# a tool group named "plan_related". So that the agent can activate
|
||||
# the plan tools by the meta tool function
|
||||
self.plan_notebook = None
|
||||
if plan_notebook:
|
||||
self.plan_notebook = plan_notebook
|
||||
# When enable_meta_tool is True, plan tools are in plan_related
|
||||
# group and active by agent.
|
||||
# Otherwise, plan tools in bassic group and always active.
|
||||
if enable_meta_tool:
|
||||
self.toolkit.create_tool_group(
|
||||
"plan_related",
|
||||
description=self.plan_notebook.description,
|
||||
)
|
||||
for tool in plan_notebook.list_tools():
|
||||
self.toolkit.register_tool_function(
|
||||
tool,
|
||||
group_name="plan_related",
|
||||
)
|
||||
else:
|
||||
for tool in plan_notebook.list_tools():
|
||||
self.toolkit.register_tool_function(
|
||||
tool,
|
||||
)
|
||||
|
||||
# If print the reasoning hint messages
|
||||
self.print_hint_msg = print_hint_msg
|
||||
|
||||
# The maximum number of iterations of the reasoning-acting loops
|
||||
self.max_iters = max_iters
|
||||
|
||||
# The hint messages that will be attached to the prompt to guide the
|
||||
# agent's behavior before each reasoning step, and cleared after
|
||||
# each reasoning step, meaning the hint messages is one-time use only.
|
||||
# We use an InMemoryMemory instance to store the hint messages
|
||||
self._reasoning_hint_msgs = InMemoryMemory()
|
||||
|
||||
# Variables to record the intermediate state
|
||||
|
||||
# If required structured output model is provided
|
||||
self._required_structured_model: Type[BaseModel] | None = None
|
||||
|
||||
# -------------- State registration and hooks --------------
|
||||
# Register the status variables
|
||||
self.register_state("name")
|
||||
self.register_state("_sys_prompt")
|
||||
|
||||
self.register_instance_hook(
|
||||
"pre_print",
|
||||
"finish_function_pre_print_hook",
|
||||
finish_function_pre_print_hook,
|
||||
)
|
||||
|
||||
@property
|
||||
def sys_prompt(self) -> str:
|
||||
"""The dynamic system prompt of the agent."""
|
||||
return self._sys_prompt
|
||||
|
||||
@trace_reply
|
||||
async def reply(
|
||||
self,
|
||||
msg: Msg | list[Msg] | None = None,
|
||||
structured_model: Type[BaseModel] | None = None,
|
||||
) -> Msg:
|
||||
"""Generate a reply based on the current state and input arguments.
|
||||
|
||||
Args:
|
||||
msg (`Msg | list[Msg] | None`, optional):
|
||||
The input message(s) to the agent.
|
||||
structured_model (`Type[BaseModel] | None`, optional):
|
||||
The required structured output model. If provided, the agent
|
||||
is expected to generate structured output in the `metadata`
|
||||
field of the output message.
|
||||
|
||||
Returns:
|
||||
`Msg`:
|
||||
The output message generated by the agent.
|
||||
"""
|
||||
# Record the input message(s) in the memory
|
||||
await self.memory.add(msg)
|
||||
|
||||
# Retrieve relevant records from the long-term memory if activated
|
||||
await self._retrieve_from_long_term_memory(msg)
|
||||
# Retrieve relevant documents from the knowledge base(s) if any
|
||||
await self._retrieve_from_knowledge(msg)
|
||||
|
||||
self._required_structured_model = structured_model
|
||||
# Record structured output model if provided
|
||||
if structured_model:
|
||||
self.toolkit.set_extended_model(
|
||||
self.finish_function_name,
|
||||
structured_model,
|
||||
)
|
||||
|
||||
# The reasoning-acting loop
|
||||
reply_msg = None
|
||||
for _ in range(self.max_iters):
|
||||
msg_reasoning = await self._reasoning()
|
||||
|
||||
futures = [
|
||||
self._acting(tool_call)
|
||||
for tool_call in msg_reasoning.get_content_blocks(
|
||||
"tool_use",
|
||||
)
|
||||
]
|
||||
|
||||
# Parallel tool calls or not
|
||||
if self.parallel_tool_calls:
|
||||
acting_responses = await asyncio.gather(*futures)
|
||||
|
||||
else:
|
||||
# Sequential tool calls
|
||||
acting_responses = [await _ for _ in futures]
|
||||
|
||||
# Find the first non-None replying message from the acting
|
||||
for acting_msg in acting_responses:
|
||||
reply_msg = reply_msg or acting_msg
|
||||
|
||||
if reply_msg:
|
||||
break
|
||||
|
||||
# When the maximum iterations are reached
|
||||
if reply_msg is None:
|
||||
reply_msg = await self._summarizing()
|
||||
|
||||
# Post-process the memory, long-term memory
|
||||
if self._static_control:
|
||||
await self.long_term_memory.record(
|
||||
[
|
||||
*([*msg] if isinstance(msg, list) else [msg]),
|
||||
*await self.memory.get_memory(),
|
||||
reply_msg,
|
||||
],
|
||||
)
|
||||
|
||||
await self.memory.add(reply_msg)
|
||||
return reply_msg
|
||||
|
||||
async def _reasoning(
|
||||
self,
|
||||
) -> Msg:
|
||||
"""Perform the reasoning process."""
|
||||
if self.plan_notebook:
|
||||
# Insert the reasoning hint from the plan notebook
|
||||
hint_msg = await self.plan_notebook.get_current_hint()
|
||||
if self.print_hint_msg and hint_msg:
|
||||
await self.print(hint_msg)
|
||||
await self._reasoning_hint_msgs.add(hint_msg)
|
||||
|
||||
# Convert Msg objects into the required format of the model API
|
||||
prompt = await self.formatter.format(
|
||||
msgs=[
|
||||
Msg("system", self.sys_prompt, "system"),
|
||||
*await self.memory.get_memory(),
|
||||
# The hint messages to guide the agent's behavior, maybe empty
|
||||
*await self._reasoning_hint_msgs.get_memory(),
|
||||
],
|
||||
)
|
||||
# Clear the hint messages after use
|
||||
await self._reasoning_hint_msgs.clear()
|
||||
|
||||
res = await self.model(
|
||||
prompt,
|
||||
tools=self.toolkit.get_json_schemas(),
|
||||
)
|
||||
|
||||
# handle output from the model
|
||||
interrupted_by_user = False
|
||||
msg = None
|
||||
try:
|
||||
if self.model.stream:
|
||||
msg = Msg(self.name, [], "assistant")
|
||||
async for content_chunk in res:
|
||||
msg.content = content_chunk.content
|
||||
await self.print(msg, False)
|
||||
await self.print(msg, True)
|
||||
|
||||
# Add a tiny sleep to yield the last message object in the
|
||||
# message queue
|
||||
await asyncio.sleep(0.001)
|
||||
|
||||
else:
|
||||
msg = Msg(self.name, list(res.content), "assistant")
|
||||
await self.print(msg, True)
|
||||
|
||||
except asyncio.CancelledError as e:
|
||||
interrupted_by_user = True
|
||||
raise e from None
|
||||
|
||||
finally:
|
||||
if msg and not msg.has_content_blocks("tool_use"):
|
||||
# Turn plain text response into a tool call of the finish
|
||||
# function
|
||||
msg = Msg.from_dict(msg.to_dict())
|
||||
msg.content = [
|
||||
ToolUseBlock(
|
||||
id=shortuuid.uuid(),
|
||||
type="tool_use",
|
||||
name=self.finish_function_name,
|
||||
input={"response": msg.get_text_content()},
|
||||
),
|
||||
]
|
||||
|
||||
# None will be ignored by the memory
|
||||
await self.memory.add(msg)
|
||||
|
||||
# Post-process for user interruption
|
||||
if interrupted_by_user and msg:
|
||||
# Fake tool results
|
||||
tool_use_blocks: list = msg.get_content_blocks(
|
||||
"tool_use",
|
||||
)
|
||||
for tool_call in tool_use_blocks:
|
||||
msg_res = Msg(
|
||||
"system",
|
||||
[
|
||||
ToolResultBlock(
|
||||
type="tool_result",
|
||||
id=tool_call["id"],
|
||||
name=tool_call["name"],
|
||||
output="The tool call has been interrupted "
|
||||
"by the user.",
|
||||
),
|
||||
],
|
||||
"system",
|
||||
)
|
||||
await self.memory.add(msg_res)
|
||||
await self.print(msg_res, True)
|
||||
return msg
|
||||
|
||||
async def _acting(self, tool_call: ToolUseBlock) -> Msg | None:
|
||||
"""Perform the acting process.
|
||||
|
||||
Args:
|
||||
tool_call (`ToolUseBlock`):
|
||||
The tool use block to be executed.
|
||||
|
||||
Returns:
|
||||
`Union[Msg, None]`:
|
||||
Return a message to the user if the `finish_function` is
|
||||
called, otherwise return `None`.
|
||||
"""
|
||||
|
||||
tool_res_msg = Msg(
|
||||
"system",
|
||||
[
|
||||
ToolResultBlock(
|
||||
type="tool_result",
|
||||
id=tool_call["id"],
|
||||
name=tool_call["name"],
|
||||
output=[],
|
||||
),
|
||||
],
|
||||
"system",
|
||||
)
|
||||
try:
|
||||
# Execute the tool call
|
||||
tool_res = await self.toolkit.call_tool_function(tool_call)
|
||||
|
||||
response_msg = None
|
||||
# Async generator handling
|
||||
async for chunk in tool_res:
|
||||
# Turn into a tool result block
|
||||
tool_res_msg.content[0][ # type: ignore[index]
|
||||
"output"
|
||||
] = chunk.content
|
||||
|
||||
# Skip the printing of the finish function call
|
||||
if (
|
||||
tool_call["name"] != self.finish_function_name
|
||||
or tool_call["name"] == self.finish_function_name
|
||||
and (
|
||||
chunk.metadata is None
|
||||
or not chunk.metadata.get("success")
|
||||
)
|
||||
):
|
||||
await self.print(tool_res_msg, chunk.is_last)
|
||||
|
||||
# Raise the CancelledError to handle the interruption in the
|
||||
# handle_interrupt function
|
||||
if chunk.is_interrupted:
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
# Return message if generate_response is called successfully
|
||||
if (
|
||||
tool_call["name"] == self.finish_function_name
|
||||
and chunk.metadata
|
||||
and chunk.metadata.get(
|
||||
"success",
|
||||
True,
|
||||
)
|
||||
):
|
||||
response_msg = chunk.metadata.get("response_msg")
|
||||
|
||||
return response_msg
|
||||
|
||||
finally:
|
||||
# Record the tool result message in the memory
|
||||
await self.memory.add(tool_res_msg)
|
||||
|
||||
async def observe(self, msg: Msg | list[Msg] | None) -> None:
|
||||
"""Receive observing message(s) without generating a reply.
|
||||
|
||||
Args:
|
||||
msg (`Msg | list[Msg] | None`):
|
||||
The message or messages to be observed.
|
||||
"""
|
||||
await self.memory.add(msg)
|
||||
|
||||
async def _summarizing(self) -> Msg:
|
||||
"""Generate a response when the agent fails to solve the problem in
|
||||
the maximum iterations."""
|
||||
hint_msg = Msg(
|
||||
"user",
|
||||
"You have failed to generate response within the maximum "
|
||||
"iterations. Now respond directly by summarizing the current "
|
||||
"situation.",
|
||||
role="user",
|
||||
)
|
||||
|
||||
# Generate a reply by summarizing the current situation
|
||||
prompt = await self.formatter.format(
|
||||
[
|
||||
Msg("system", self.sys_prompt, "system"),
|
||||
*await self.memory.get_memory(),
|
||||
hint_msg,
|
||||
],
|
||||
)
|
||||
# TODO: handle the structured output here, maybe force calling the
|
||||
# finish_function here
|
||||
res = await self.model(prompt)
|
||||
|
||||
res_msg = Msg(self.name, [], "assistant")
|
||||
if isinstance(res, AsyncGenerator):
|
||||
async for chunk in res:
|
||||
res_msg.content = chunk.content
|
||||
await self.print(res_msg, False)
|
||||
await self.print(res_msg, True)
|
||||
|
||||
else:
|
||||
res_msg.content = res.content
|
||||
await self.print(res_msg, True)
|
||||
|
||||
return res_msg
|
||||
|
||||
async def handle_interrupt(
|
||||
self,
|
||||
_msg: Msg | list[Msg] | None = None,
|
||||
) -> Msg:
|
||||
"""The post-processing logic when the reply is interrupted by the
|
||||
user or something else."""
|
||||
|
||||
response_msg = Msg(
|
||||
self.name,
|
||||
"I noticed that you have interrupted me. What can I "
|
||||
"do for you?",
|
||||
"assistant",
|
||||
metadata={
|
||||
# Expose this field to indicate the interruption
|
||||
"is_interrupted": True,
|
||||
},
|
||||
)
|
||||
|
||||
await self.print(response_msg, True)
|
||||
await self.memory.add(response_msg)
|
||||
return response_msg
|
||||
|
||||
def generate_response(
|
||||
self,
|
||||
response: str,
|
||||
**kwargs: Any,
|
||||
) -> ToolResponse:
|
||||
"""Generate a response. Note only the input argument `response` is
|
||||
visible to the others, you should include all the necessary
|
||||
information in the `response` argument.
|
||||
|
||||
Args:
|
||||
response (`str`):
|
||||
Your response to the user.
|
||||
"""
|
||||
response_msg = Msg(
|
||||
self.name,
|
||||
response,
|
||||
"assistant",
|
||||
)
|
||||
|
||||
# Prepare structured output
|
||||
if self._required_structured_model:
|
||||
try:
|
||||
# Use the metadata field of the message to store the
|
||||
# structured output
|
||||
response_msg.metadata = (
|
||||
self._required_structured_model.model_validate(
|
||||
kwargs,
|
||||
).model_dump()
|
||||
)
|
||||
|
||||
except ValidationError as e:
|
||||
return ToolResponse(
|
||||
content=[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=f"Arguments Validation Error: {e}",
|
||||
),
|
||||
],
|
||||
metadata={
|
||||
"success": False,
|
||||
"response_msg": None,
|
||||
},
|
||||
)
|
||||
|
||||
return ToolResponse(
|
||||
content=[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text="Successfully generated response.",
|
||||
),
|
||||
],
|
||||
metadata={
|
||||
"success": True,
|
||||
"response_msg": response_msg,
|
||||
},
|
||||
is_last=True,
|
||||
)
|
||||
|
||||
async def _retrieve_from_long_term_memory(
|
||||
self,
|
||||
msg: Msg | list[Msg] | None,
|
||||
) -> None:
|
||||
"""Insert the retrieved information from the long-term memory into
|
||||
the short-term memory as a Msg object.
|
||||
|
||||
Args:
|
||||
msg (`Msg | list[Msg] | None`):
|
||||
The input message to the agent.
|
||||
"""
|
||||
if self._static_control and msg:
|
||||
# Retrieve information from the long-term memory if available
|
||||
retrieved_info = await self.long_term_memory.retrieve(msg)
|
||||
if retrieved_info:
|
||||
retrieved_msg = Msg(
|
||||
name="long_term_memory",
|
||||
content="<long_term_memory>The content below are "
|
||||
"retrieved from long-term memory, which maybe "
|
||||
f"useful:\n{retrieved_info}</long_term_memory>",
|
||||
role="user",
|
||||
)
|
||||
if self.print_hint_msg:
|
||||
await self.print(retrieved_msg, True)
|
||||
await self.memory.add(retrieved_msg)
|
||||
|
||||
async def _retrieve_from_knowledge(
|
||||
self,
|
||||
msg: Msg | list[Msg] | None,
|
||||
) -> None:
|
||||
"""Insert the retrieved documents from the RAG knowledge base(s) if
|
||||
available.
|
||||
|
||||
Args:
|
||||
msg (`Msg | list[Msg] | None`):
|
||||
The input message to the agent.
|
||||
"""
|
||||
if self.knowledge and msg:
|
||||
# Prepare the user input query
|
||||
query = None
|
||||
if isinstance(msg, Msg):
|
||||
query = msg.get_text_content()
|
||||
elif isinstance(msg, list):
|
||||
query = "\n".join(_.get_text_content() for _ in msg)
|
||||
|
||||
# Skip if the query is empty
|
||||
if not query:
|
||||
return
|
||||
|
||||
# Rewrite the query by the LLM if enabled
|
||||
if self.enable_rewrite_query:
|
||||
try:
|
||||
rewrite_prompt = await self.formatter.format(
|
||||
msgs=[
|
||||
Msg("system", self.sys_prompt, "system"),
|
||||
*await self.memory.get_memory(),
|
||||
Msg(
|
||||
"user",
|
||||
"<system-hint>Now you need to rewrite "
|
||||
"the above user query to be more specific and "
|
||||
"concise for knowledge retrieval. For "
|
||||
"example, rewrite the query 'what happened "
|
||||
"last day' to 'what happened on 2023-10-01' "
|
||||
"(assuming today is 2023-10-02)."
|
||||
"</system-hint>",
|
||||
"user",
|
||||
),
|
||||
],
|
||||
)
|
||||
stream_tmp = self.model.stream
|
||||
self.model.stream = False
|
||||
res = await self.model(
|
||||
rewrite_prompt,
|
||||
structured_model=_QueryRewriteModel,
|
||||
)
|
||||
self.model.stream = stream_tmp
|
||||
if res.metadata and res.metadata.get("rewritten_query"):
|
||||
query = res.metadata["rewritten_query"]
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Skipping the query rewriting due to error: %s",
|
||||
str(e),
|
||||
)
|
||||
|
||||
docs: list[Document] = []
|
||||
for kb in self.knowledge:
|
||||
# retrieve the user input query
|
||||
docs.extend(
|
||||
await kb.retrieve(query=query),
|
||||
)
|
||||
if docs:
|
||||
# Rerank by the relevance score
|
||||
docs = sorted(
|
||||
docs,
|
||||
key=lambda doc: doc.score or 0.0,
|
||||
reverse=True,
|
||||
)
|
||||
# Prepare the retrieved knowledge string
|
||||
retrieved_msg = Msg(
|
||||
name="user",
|
||||
content=[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=(
|
||||
"<retrieved_knowledge>Use the following "
|
||||
"content from the knowledge base(s) if it's "
|
||||
"helpful:\n"
|
||||
),
|
||||
),
|
||||
*[_.metadata.content for _ in docs],
|
||||
TextBlock(
|
||||
type="text",
|
||||
text="</retrieved_knowledge>",
|
||||
),
|
||||
],
|
||||
role="user",
|
||||
)
|
||||
if self.print_hint_msg:
|
||||
await self.print(retrieved_msg, True)
|
||||
await self.memory.add(retrieved_msg)
|
||||
@@ -0,0 +1,116 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The base class for ReAct agent in agentscope."""
|
||||
from abc import abstractmethod
|
||||
from collections import OrderedDict
|
||||
from typing import Callable, Any
|
||||
|
||||
from ._agent_base import AgentBase
|
||||
from ._agent_meta import _ReActAgentMeta
|
||||
from ..message import Msg
|
||||
|
||||
|
||||
class ReActAgentBase(AgentBase, metaclass=_ReActAgentMeta):
|
||||
"""The ReAct agent base class.
|
||||
|
||||
To support ReAct algorithm, this class extends the AgentBase class by
|
||||
adding two abstract interfaces: reasoning and acting, while supporting
|
||||
hook functions at four positions: pre-reasoning, post-reasoning,
|
||||
pre-acting, and post-acting by the `_ReActAgentMeta` metaclass.
|
||||
"""
|
||||
|
||||
supported_hook_types: list[str] = [
|
||||
"pre_reply",
|
||||
"post_reply",
|
||||
"pre_print",
|
||||
"post_print",
|
||||
"pre_observe",
|
||||
"post_observe",
|
||||
"pre_reasoning",
|
||||
"post_reasoning",
|
||||
"pre_acting",
|
||||
"post_acting",
|
||||
]
|
||||
"""Supported hook types for the agent base class."""
|
||||
|
||||
_class_pre_reasoning_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"ReActAgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
],
|
||||
dict[str, Any] | None, # The modified kwargs or None
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level pre-reasoning hooks, taking `self` object, the input
|
||||
arguments as input"""
|
||||
|
||||
_class_post_reasoning_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"ReActAgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
Any, # output
|
||||
],
|
||||
Msg | None, # the modified output message or None
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level post-reasoning hooks, taking `self` object, the input
|
||||
arguments and the output message as input, and return the modified output
|
||||
message or None if no modification is needed."""
|
||||
|
||||
_class_pre_acting_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"ReActAgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
],
|
||||
dict[str, Any] | None, # The modified kwargs or None
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level pre-acting hooks, taking `self` object, the input
|
||||
arguments as input, and return the modified input arguments or None if no
|
||||
modification is needed."""
|
||||
|
||||
_class_post_acting_hooks: dict[
|
||||
str,
|
||||
Callable[
|
||||
[
|
||||
"ReActAgentBase", # self
|
||||
dict[str, Any], # kwargs
|
||||
Any, # output
|
||||
],
|
||||
Msg | None, # the modified output message or None
|
||||
],
|
||||
] = OrderedDict()
|
||||
"""The class-level post-acting hooks, taking `self` object, the input
|
||||
arguments and the output message as input, and return the modified output
|
||||
message or None if no modification is needed."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
) -> None:
|
||||
"""Initialize the ReAct agent base class."""
|
||||
super().__init__()
|
||||
|
||||
# Init reasoning and acting hooks
|
||||
self._instance_pre_reasoning_hooks = OrderedDict()
|
||||
self._instance_post_reasoning_hooks = OrderedDict()
|
||||
self._instance_pre_acting_hooks = OrderedDict()
|
||||
self._instance_post_acting_hooks = OrderedDict()
|
||||
|
||||
@abstractmethod
|
||||
async def _reasoning(
|
||||
self,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""The reasoning process of the ReAct agent, which will be wrapped
|
||||
with pre- and post-hooks."""
|
||||
|
||||
@abstractmethod
|
||||
async def _acting(self, *args: Any, **kwargs: Any) -> Any:
|
||||
"""The acting process of the ReAct agent, which will be wrapped with
|
||||
pre- and post-hooks."""
|
||||
@@ -0,0 +1,128 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The user agent class."""
|
||||
from typing import Type, Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ._agent_base import AgentBase
|
||||
from ._user_input import UserInputBase, TerminalUserInput
|
||||
from ..message import Msg
|
||||
|
||||
|
||||
class UserAgent(AgentBase):
|
||||
"""The class for user interaction, allowing developers to handle the user
|
||||
input from different sources, such as web UI, cli, and other interfaces.
|
||||
"""
|
||||
|
||||
_input_method: UserInputBase = TerminalUserInput()
|
||||
"""The user input method, can be overridden by calling the
|
||||
`register_instance/class_input_method` function."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
) -> None:
|
||||
"""Initialize the user agent with a name."""
|
||||
super().__init__()
|
||||
|
||||
self.name = name
|
||||
|
||||
async def reply(
|
||||
self,
|
||||
msg: Msg | list[Msg] | None = None,
|
||||
structured_model: Type[BaseModel] | None = None,
|
||||
) -> Msg:
|
||||
"""Receive input message(s) and generate a reply message from the user.
|
||||
|
||||
Args:
|
||||
msg (`Msg | list[Msg] | None`, defaults to `None`):
|
||||
The message(s) to be replied. If `None`, the agent will wait
|
||||
for user input.
|
||||
structured_model (`Type[BaseModel] | None`, defaults to `None`):
|
||||
A child class of `pydantic.BaseModel` that defines the
|
||||
structured output format. If provided, the user will be
|
||||
prompted to fill in the required fields.
|
||||
|
||||
Returns:
|
||||
`Msg`:
|
||||
The reply message generated by the user.
|
||||
"""
|
||||
|
||||
# Get the input from the specified input method.
|
||||
input_data = await self._input_method(
|
||||
agent_id=self.id,
|
||||
agent_name=self.name,
|
||||
structured_model=structured_model,
|
||||
)
|
||||
|
||||
blocks_input = input_data.blocks_input
|
||||
if (
|
||||
blocks_input
|
||||
and len(blocks_input) == 1
|
||||
and blocks_input[0].get("type") == "text"
|
||||
):
|
||||
# Turn blocks_input into a string if only one text block exists
|
||||
blocks_input = blocks_input[0].get("text")
|
||||
|
||||
msg = Msg(
|
||||
self.name,
|
||||
content=blocks_input,
|
||||
role="user",
|
||||
metadata=input_data.structured_input,
|
||||
)
|
||||
|
||||
await self.print(msg)
|
||||
|
||||
return msg
|
||||
|
||||
def override_instance_input_method(
|
||||
self,
|
||||
input_method: UserInputBase,
|
||||
) -> None:
|
||||
"""Override the input method of the current UserAgent instance.
|
||||
|
||||
Args:
|
||||
input_method (`UserInputBase`):
|
||||
The callable input method, which should be an object of a
|
||||
class that inherits from `UserInputBase`.
|
||||
"""
|
||||
if not isinstance(input_method, UserInputBase):
|
||||
raise ValueError(
|
||||
f"The input method should be an instance of the child class "
|
||||
f"of `UserInputBase`, but got {type(input_method)} instead.",
|
||||
)
|
||||
self._input_method = input_method
|
||||
|
||||
@classmethod
|
||||
def override_class_input_method(
|
||||
cls,
|
||||
input_method: UserInputBase,
|
||||
) -> None:
|
||||
"""Override the input method of the current UserAgent class.
|
||||
|
||||
Args:
|
||||
input_method (`UserInputBase`):
|
||||
The callable input method, which should be an object of a
|
||||
class that inherits from `UserInputBase`.
|
||||
"""
|
||||
if not isinstance(input_method, UserInputBase):
|
||||
raise ValueError(
|
||||
f"The input method should be an instance of the child class "
|
||||
f"of `UserInputBase`, but got {type(input_method)} instead.",
|
||||
)
|
||||
cls._input_method = input_method
|
||||
|
||||
async def handle_interrupt(
|
||||
self,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Msg:
|
||||
"""The post-processing logic when the reply is interrupted by the
|
||||
user or something else."""
|
||||
raise NotImplementedError(
|
||||
f"The handle_interrupt function is not implemented in "
|
||||
f"{self.__class__.__name__}",
|
||||
)
|
||||
|
||||
async def observe(self, msg: Msg | list[Msg] | None) -> None:
|
||||
"""Observe the message(s) from the other agents or the environment."""
|
||||
@@ -0,0 +1,411 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The user input related classes."""
|
||||
import json.decoder
|
||||
import time
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from queue import Queue
|
||||
from threading import Event
|
||||
from typing import Any, Type, List
|
||||
|
||||
import jsonschema
|
||||
import requests
|
||||
import shortuuid
|
||||
import socketio
|
||||
from pydantic import BaseModel
|
||||
import json5
|
||||
|
||||
from .. import _config
|
||||
from .._logging import logger
|
||||
from ..message import (
|
||||
TextBlock,
|
||||
VideoBlock,
|
||||
AudioBlock,
|
||||
ImageBlock,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class UserInputData:
|
||||
"""The user input data."""
|
||||
|
||||
blocks_input: List[TextBlock | ImageBlock | AudioBlock | VideoBlock] = None
|
||||
"""The text input from the user"""
|
||||
|
||||
structured_input: dict[str, Any] | None = None
|
||||
"""The structured input from the user"""
|
||||
|
||||
|
||||
class UserInputBase:
|
||||
"""The base class used to handle the user input from different sources."""
|
||||
|
||||
@abstractmethod
|
||||
async def __call__(
|
||||
self,
|
||||
agent_id: str,
|
||||
agent_name: str,
|
||||
*args: Any,
|
||||
structured_model: Type[BaseModel] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> UserInputData:
|
||||
"""The user input method, which returns the user input and the
|
||||
required structured data.
|
||||
|
||||
Args:
|
||||
agent_id (`str`):
|
||||
The agent identifier.
|
||||
agent_name (`str`):
|
||||
The agent name.
|
||||
structured_model (`Type[BaseModel] | None`, optional):
|
||||
A base model class that defines the structured input format.
|
||||
|
||||
Returns:
|
||||
`UserInputData`:
|
||||
The user input data.
|
||||
"""
|
||||
|
||||
|
||||
class TerminalUserInput(UserInputBase):
|
||||
"""The terminal user input."""
|
||||
|
||||
def __init__(self, input_hint: str = "User Input: ") -> None:
|
||||
"""Initialize the terminal user input with a hint."""
|
||||
self.input_hint = input_hint
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
agent_id: str,
|
||||
agent_name: str,
|
||||
*args: Any,
|
||||
structured_model: Type[BaseModel] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> UserInputData:
|
||||
"""Handle the user input from the terminal.
|
||||
|
||||
Args:
|
||||
agent_id (`str`):
|
||||
The agent identifier.
|
||||
agent_name (`str`):
|
||||
The agent name.
|
||||
structured_model (`Type[BaseModel] | None`, optional):
|
||||
A base model class that defines the structured input format.
|
||||
|
||||
Returns:
|
||||
`UserInputData`:
|
||||
The user input data.
|
||||
"""
|
||||
|
||||
text_input = input(self.input_hint)
|
||||
|
||||
structured_input = None
|
||||
if structured_model is not None:
|
||||
structured_input = {}
|
||||
|
||||
json_schema = structured_model.model_json_schema()
|
||||
required = json_schema.get("required", [])
|
||||
print("Structured input (press Enter to skip for optional):)")
|
||||
|
||||
for key, item in json_schema.get("properties").items():
|
||||
requirements = {**item}
|
||||
requirements.pop("title")
|
||||
|
||||
while True:
|
||||
res = input(f"\t{key} ({requirements}): ")
|
||||
|
||||
if res == "":
|
||||
if key in required:
|
||||
print(f"Key {key} is required.")
|
||||
continue
|
||||
|
||||
res = item.get("default", None)
|
||||
|
||||
if item.get("type").lower() == "integer":
|
||||
try:
|
||||
res = json5.loads(res)
|
||||
except json.decoder.JSONDecodeError as e:
|
||||
print(
|
||||
"\033[31mInvalid input with error:\n"
|
||||
"```\n"
|
||||
f"{e}\n"
|
||||
"```\033[0m",
|
||||
)
|
||||
continue
|
||||
|
||||
try:
|
||||
jsonschema.validate(res, item)
|
||||
structured_input[key] = res
|
||||
break
|
||||
except jsonschema.ValidationError as e:
|
||||
print(
|
||||
f"\033[31mValidation error:\n```\n{e}\n```\033[0m",
|
||||
)
|
||||
time.sleep(0.5)
|
||||
|
||||
return UserInputData(
|
||||
blocks_input=[TextBlock(type="text", text=text_input)],
|
||||
structured_input=structured_input,
|
||||
)
|
||||
|
||||
|
||||
class StudioUserInput(UserInputBase):
|
||||
"""The class that host the user input on the AgentScope Studio."""
|
||||
|
||||
_websocket_namespace: str = "/python"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
studio_url: str,
|
||||
run_id: str,
|
||||
max_retries: int = 3,
|
||||
reconnect_attempts: int = 3,
|
||||
reconnection_delay: int = 1,
|
||||
reconnection_delay_max: int = 5,
|
||||
) -> None:
|
||||
"""Initialize the StudioUserInput object.
|
||||
|
||||
Args:
|
||||
studio_url (`str`):
|
||||
The URL of the AgentScope Studio.
|
||||
run_id (`str`):
|
||||
The current run identity.
|
||||
max_retries (`int`, defaults to `3`):
|
||||
The maximum number of retries to get user input.
|
||||
"""
|
||||
self._is_connected = False
|
||||
self._is_reconnecting = False
|
||||
|
||||
self.studio_url = studio_url
|
||||
self.run_id = run_id
|
||||
self.max_retries = max_retries
|
||||
|
||||
# Init Websocket
|
||||
self.sio = socketio.Client(
|
||||
reconnection=True,
|
||||
reconnection_attempts=reconnect_attempts,
|
||||
reconnection_delay=reconnection_delay,
|
||||
reconnection_delay_max=reconnection_delay_max,
|
||||
)
|
||||
self.input_queues = {}
|
||||
self.input_events = {}
|
||||
|
||||
@self.sio.on("connect", namespace=self._websocket_namespace)
|
||||
def on_connect() -> None:
|
||||
self._is_connected = True
|
||||
logger.info(
|
||||
'Connected to AgentScope Studio at "%s" with '
|
||||
'run name "%s".',
|
||||
self.studio_url,
|
||||
run_id,
|
||||
)
|
||||
logger.info(
|
||||
"View the run at: %s/dashboard/projects/%s",
|
||||
self.studio_url,
|
||||
_config.project,
|
||||
)
|
||||
|
||||
@self.sio.on("disconnect", namespace=self._websocket_namespace)
|
||||
def on_disconnect() -> None:
|
||||
self._is_connected = False
|
||||
logger.info(
|
||||
"Disconnected from AgentScope Studio at %s",
|
||||
self.studio_url,
|
||||
)
|
||||
|
||||
@self.sio.on("reconnect", namespace=self._websocket_namespace)
|
||||
def on_reconnect(attempt_number: int) -> None:
|
||||
self._is_connected = True
|
||||
self._is_reconnecting = False
|
||||
logger.info(
|
||||
"Reconnected to AgentScope Studio at %s with run_id %s after "
|
||||
"%d attempts",
|
||||
self.studio_url,
|
||||
self.run_id,
|
||||
attempt_number,
|
||||
)
|
||||
|
||||
@self.sio.on("reconnect_attempt", namespace=self._websocket_namespace)
|
||||
def on_reconnect_attempt(attempt_number: int) -> None:
|
||||
self._is_reconnecting = True
|
||||
logger.info(
|
||||
"Attempting to reconnect to AgentScope Studio at %s "
|
||||
"(attempt %d)",
|
||||
self.studio_url,
|
||||
attempt_number,
|
||||
)
|
||||
|
||||
@self.sio.on("reconnect_failed", namespace=self._websocket_namespace)
|
||||
def on_reconnect_failed() -> None:
|
||||
self._is_reconnecting = False
|
||||
logger.error(
|
||||
"Failed to reconnect to AgentScope Studio at %s",
|
||||
self.studio_url,
|
||||
)
|
||||
|
||||
@self.sio.on("reconnect_error", namespace=self._websocket_namespace)
|
||||
def on_reconnect_error(error: Any) -> None:
|
||||
logger.error(
|
||||
"Error while reconnecting to AgentScope Studio at %s: %s",
|
||||
self.studio_url,
|
||||
str(error),
|
||||
)
|
||||
|
||||
# The AgentScope Studio backend send the "sendUserInput" event to
|
||||
# the current python run
|
||||
@self.sio.on("forwardUserInput", namespace=self._websocket_namespace)
|
||||
def receive_user_input(
|
||||
request_id: str,
|
||||
blocks_input: List[
|
||||
TextBlock | ImageBlock | AudioBlock | VideoBlock
|
||||
],
|
||||
structured_input: dict[str, Any],
|
||||
) -> None:
|
||||
if request_id in self.input_queues:
|
||||
self.input_queues[request_id].put(
|
||||
UserInputData(
|
||||
blocks_input=blocks_input,
|
||||
structured_input=structured_input,
|
||||
),
|
||||
)
|
||||
self.input_events[request_id].set()
|
||||
|
||||
try:
|
||||
self.sio.connect(
|
||||
f"{self.studio_url}",
|
||||
namespaces=["/python"],
|
||||
auth={"run_id": self.run_id},
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to connect to AgentScope Studio at {self.studio_url}",
|
||||
) from e
|
||||
|
||||
def _ensure_connected(
|
||||
self,
|
||||
timeout: float = 30.0,
|
||||
check_interval: float = 5.0,
|
||||
) -> None:
|
||||
"""Ensure the connection is established or wait for reconnection.
|
||||
|
||||
Args:
|
||||
timeout (`float`):
|
||||
Maximum time to wait for reconnection in seconds. Defaults
|
||||
to 30.0.
|
||||
check_interval (`float`):
|
||||
Interval between connection checks in seconds. Defaults to 1.0.
|
||||
|
||||
Raises:
|
||||
`RuntimeError`:
|
||||
If connection cannot be established within timeout.
|
||||
"""
|
||||
if self._is_connected:
|
||||
return
|
||||
|
||||
if self._is_reconnecting:
|
||||
start_time = time.time()
|
||||
while self._is_reconnecting:
|
||||
# Check timeout
|
||||
elapsed_time = time.time() - start_time
|
||||
if elapsed_time > timeout:
|
||||
raise RuntimeError(
|
||||
f"Reconnection timeout after {elapsed_time} seconds",
|
||||
)
|
||||
|
||||
# Log status
|
||||
logger.info(
|
||||
"Waiting for reconnection... (%.1fs / %.1fs)",
|
||||
elapsed_time,
|
||||
timeout,
|
||||
)
|
||||
|
||||
# Wait for next check
|
||||
time.sleep(check_interval)
|
||||
|
||||
# After reconnection attempt completed, check final status
|
||||
if self._is_connected:
|
||||
return
|
||||
|
||||
# Not connected and not reconnecting
|
||||
raise RuntimeError(
|
||||
f"Not connected to AgentScope Studio at {self.studio_url}.",
|
||||
)
|
||||
|
||||
async def __call__( # type: ignore[override]
|
||||
self,
|
||||
agent_id: str,
|
||||
agent_name: str,
|
||||
*args: Any,
|
||||
structured_model: Type[BaseModel] | None = None,
|
||||
) -> UserInputData:
|
||||
"""Get the user input from AgentScope Studio.
|
||||
|
||||
Args:
|
||||
agent_id (`str`):
|
||||
The identity of the agent.
|
||||
agent_name (`str`):
|
||||
The name of the agent.
|
||||
structured_model (`Type[BaseModel] | None`, optional):
|
||||
The base model class of the structured input.
|
||||
|
||||
Raises:
|
||||
`RuntimeError`:
|
||||
Failed to get user input from AgentScope Studio.
|
||||
|
||||
Returns:
|
||||
`UserInputData`:
|
||||
The user input.
|
||||
"""
|
||||
self._ensure_connected()
|
||||
|
||||
request_id = shortuuid.uuid()
|
||||
|
||||
self.input_queues[request_id] = Queue()
|
||||
self.input_events[request_id] = Event()
|
||||
|
||||
if structured_model is None:
|
||||
structured_input = None
|
||||
else:
|
||||
structured_input = structured_model.model_json_schema()
|
||||
|
||||
n_retry = 0
|
||||
while True:
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{self.studio_url}/trpc/requestUserInput",
|
||||
json={
|
||||
"requestId": request_id,
|
||||
"runId": self.run_id,
|
||||
"agentId": agent_id,
|
||||
"agentName": agent_name,
|
||||
"structuredInput": structured_input,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
break
|
||||
except Exception as e:
|
||||
if n_retry < self.max_retries:
|
||||
n_retry += 1
|
||||
continue
|
||||
|
||||
raise RuntimeError(
|
||||
"Failed to get user input from AgentScope Studio",
|
||||
) from e
|
||||
|
||||
try:
|
||||
self.input_events[request_id].wait()
|
||||
response_data = self.input_queues[request_id].get()
|
||||
return response_data
|
||||
|
||||
finally:
|
||||
self.input_queues.pop(request_id, None)
|
||||
self.input_events.pop(request_id, None)
|
||||
|
||||
def __del__(self) -> None:
|
||||
"""Cleanup socket connection when object it destroyed"""
|
||||
try:
|
||||
self.sio.disconnect()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to disconnect from AgentScope Studio at %s: %s",
|
||||
self.studio_url,
|
||||
str(e),
|
||||
)
|
||||
Reference in New Issue
Block a user