124 lines
4.4 KiB
Python
124 lines
4.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""An agent that replies in a dictionary format."""
|
|
from typing import Optional, Union, Sequence
|
|
|
|
from ..message import Msg
|
|
from .agent import AgentBase
|
|
from ..parsers import ParserBase
|
|
|
|
|
|
class DictDialogAgent(AgentBase):
|
|
"""An agent that generates response in a dict format, where user can
|
|
specify the required fields in the response via specifying the parser
|
|
|
|
About parser, please refer to our
|
|
[tutorial](https://modelscope.github.io/agentscope/en/tutorial/203-parser.html)
|
|
|
|
For usage example, please refer to the example of werewolf in
|
|
`examples/game_werewolf`"""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
sys_prompt: str,
|
|
model_config_name: str,
|
|
use_memory: bool = True,
|
|
memory_config: Optional[dict] = None,
|
|
max_retries: Optional[int] = 3,
|
|
) -> None:
|
|
"""Initialize the dict dialog agent.
|
|
|
|
Arguments:
|
|
name (`str`):
|
|
The name of the agent.
|
|
sys_prompt (`Optional[str]`, defaults to `None`):
|
|
The system prompt of the agent, which can be passed by args
|
|
or hard-coded in the agent.
|
|
model_config_name (`str`, defaults to None):
|
|
The name of the model config, which is used to load model from
|
|
configuration.
|
|
use_memory (`bool`, defaults to `True`):
|
|
Whether the agent has memory.
|
|
memory_config (`Optional[dict]`, defaults to `None`):
|
|
The config of memory.
|
|
max_retries (`Optional[int]`, defaults to `None`):
|
|
The maximum number of retries when failed to parse the model
|
|
output.
|
|
""" # noqa
|
|
super().__init__(
|
|
name=name,
|
|
sys_prompt=sys_prompt,
|
|
model_config_name=model_config_name,
|
|
use_memory=use_memory,
|
|
memory_config=memory_config,
|
|
)
|
|
|
|
self.parser = None
|
|
self.max_retries = max_retries
|
|
|
|
def set_parser(self, parser: ParserBase) -> None:
|
|
"""Set response parser, which will provide 1) format instruction; 2)
|
|
response parsing; 3) filtering fields when returning message, storing
|
|
message in memory. So developers only need to change the
|
|
parser, and the agent will work as expected.
|
|
"""
|
|
self.parser = parser
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
"""Reply function of the agent.
|
|
Processes the input data, generates a prompt using the current
|
|
dialogue memory and system prompt, and invokes the language
|
|
model to produce a response. The response is then formatted
|
|
and added to the dialogue memory.
|
|
|
|
Args:
|
|
x (`Optional[Union[Msg, Sequence[Msg]]]`, defaults to `None`):
|
|
The input message(s) to the agent, which also can be omitted if
|
|
the agent doesn't need any input.
|
|
|
|
|
|
Returns:
|
|
`Msg`: The output message generated by the agent.
|
|
|
|
Raises:
|
|
`json.decoder.JSONDecodeError`:
|
|
If the response from the language model is not valid JSON,
|
|
it defaults to treating the response as plain text.
|
|
"""
|
|
# record the input if needed
|
|
if self.memory:
|
|
self.memory.add(x)
|
|
|
|
# prepare prompt
|
|
prompt = self.model.format(
|
|
Msg("system", self.sys_prompt, role="system"),
|
|
self.memory
|
|
and self.memory.get_memory()
|
|
or x, # type: ignore[arg-type]
|
|
Msg("system", self.parser.format_instruction, "system"),
|
|
)
|
|
|
|
# call llm
|
|
res = self.model(
|
|
prompt,
|
|
parse_func=self.parser.parse,
|
|
max_retries=self.max_retries,
|
|
)
|
|
|
|
# Filter the parsed response by keys for storing in memory, returning
|
|
# in the reply function, and feeding into the metadata field in the
|
|
# returned message object.
|
|
self.memory.add(
|
|
Msg(self.name, self.parser.to_memory(res.parsed), "assistant"),
|
|
)
|
|
|
|
msg = Msg(
|
|
self.name,
|
|
content=self.parser.to_content(res.parsed),
|
|
role="assistant",
|
|
metadata=self.parser.to_metadata(res.parsed),
|
|
)
|
|
self.speak(msg)
|
|
|
|
return msg
|