from pydantic import BaseModel
from typing import List, Union, Dict
from loguru import logger
from .message import Message
from coagent.utils.common_utils import (
save_to_jsonl_file, save_to_json_file, read_json_file, read_jsonl_file
)
class Memory(BaseModel):
messages: List[Message] = []
# def __init__(self, messages: List[Message] = []):
# self.messages = messages
def append(self, message: Message):
self.messages.append(message)
def extend(self, memory: 'Memory'):
self.messages.extend(memory.messages)
def update(self, role_name: str, role_type: str, role_content: str):
self.messages.append(Message(role_name, role_type, role_content, role_content))
def clear(self, ):
self.messages = []
def delete(self, ):
pass
def get_messages(self, k=0) -> List[Message]:
"""Return the most recent k memories, return all when k=0"""
return self.messages[-k:]
def split_by_role_type(self) -> List[Dict[str, 'Memory']]:
"""
Split messages into rounds of conversation based on role_type.
Each round consists of consecutive messages of the same role_type.
User messages form a single round, while assistant and function messages are combined into a single round.
Each round is represented by a dict with 'role' and 'memory' keys, with assistant and function messages
labeled as 'assistant'.
"""
rounds = []
current_memory = Memory()
current_role = None
logger.debug(len(self.messages))
for msg in self.messages:
# Determine the message's role, considering 'function' as 'assistant'
message_role = 'assistant' if msg.role_type in ['assistant', 'function'] else 'user'
# If the current memory is empty or the current message is of the same role_type as current_role, add to current memory
if not current_memory.messages or current_role == message_role:
current_memory.append(msg)
else:
# Finish the current memory and start a new one
rounds.append({'role': current_role, 'memory': current_memory})
current_memory = Memory()
current_memory.append(msg)
# Update the current_role, considering 'function' as 'assistant'
current_role = message_role
# Don't forget to add the last memory if it exists
if current_memory.messages:
rounds.append({'role': current_role, 'memory': current_memory})
logger.debug(rounds)
return rounds
def format_rounds_to_html(self) -> str:
formatted_html_str = ""
rounds = self.split_by_role_type()
for round in rounds:
role = round['role']
memory = round['memory']
# 转换当前round的Memory为字符串
messages_str = memory.to_str_messages()
# 根据角色类型添加相应的HTML标签
if role == 'user':
formatted_html_str += f"\n{messages_str}\n\n"
else: # 对于'assistant'和'function'角色，我们将其视为'assistant'
formatted_html_str += f"\n{messages_str}\n\n"
return formatted_html_str
def filter_by_role_type(self, role_types: List[str]) -> List[Message]:
# Filter messages based on role types
return [message for message in self.messages if message.role_type not in role_types]
def select_by_role_type(self, role_types: List[str]) -> List[Message]:
# Select messages based on role types
return [message for message in self.messages if message.role_type in role_types]
def to_tuple_messages(self, return_all: bool = True, content_key="role_content", filter_roles=[]):
# Convert messages to tuples based on parameters
# logger.debug(f"{[message.to_tuple_message(return_all, content_key) for message in self.messages ]}")
return [
message.to_tuple_message(return_all, content_key) for message in self.messages
if message.role_name not in filter_roles
]
def to_dict_messages(self, filter_roles=[]):
# Convert messages to dictionaries based on filter roles
return [
message.to_dict_message() for message in self.messages
if message.role_name not in filter_roles
]
def to_str_messages(self, return_all: bool = True, content_key="role_content", filter_roles=[], with_tag=False):
# Convert messages to strings based on parameters
# for message in self.messages:
# logger.debug(f"{message.role_name}: {message.to_str_content(return_all, content_key, with_tag=with_tag)}")
# logger.debug(f"{[message.to_tuple_message(return_all, content_key) for message in self.messages ]}")
return "\n\n".join([message.to_str_content(return_all, content_key, with_tag=with_tag) for message in self.messages
if message.role_name not in filter_roles
])
def get_parserd_output(self, ):
return [message.parsed_output for message in self.messages]
def get_parserd_output_list(self, ):
# for message in self.messages:
# logger.debug(f"{message.role_name}: {message.parsed_output_list}")
# return [parsed_output for message in self.messages for parsed_output in message.parsed_output_list[1:]]
return [parsed_output for message in self.messages for parsed_output in message.parsed_output_list]
def get_spec_parserd_output(self, ):
return [message.spec_parsed_output for message in self.messages]
def get_rolenames(self, ):
''''''
return [message.role_name for message in self.messages]
@classmethod
def from_memory_list(cls, memorys: List['Memory']) -> 'Memory':
return cls(messages=[message for memory in memorys for message in memory.get_messages()])
def __len__(self, ):
return len(self.messages)
def __str__(self) -> str:
return "\n".join([":".join(i) for i in self.to_tuple_messages()])
def __add__(self, other: Union[Message, 'Memory']) -> 'Memory':
if isinstance(other, Message):
return Memory(messages=self.messages + [other])
elif isinstance(other, Memory):
return Memory(messages=self.messages + other.messages)
else:
raise ValueError(f"cant add unspecified type like as {type(other)}")