ChatDev/camel/messages/base.py
2023-09-04 19:24:28 +08:00

303 lines
11 KiB
Python

# =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
# Licensed under the Apache License, Version 2.0 (the “License”);
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an “AS IS” BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union
from camel.messages import (
OpenAIAssistantMessage,
OpenAIChatMessage,
OpenAIMessage,
OpenAISystemMessage,
OpenAIUserMessage,
)
from camel.prompts import CodePrompt, TextPrompt
from camel.typing import ModelType, RoleType
@dataclass
class BaseMessage:
r"""Base class for message objects used in CAMEL chat system.
Args:
role_name (str): The name of the user or assistant role.
role_type (RoleType): The type of role, either
:obj:`RoleType.ASSISTANT` or :obj:`RoleType.USER`.
meta_dict (Optional[Dict[str, str]]): Additional metadata dictionary
for the message.
role (str): The role of the message in OpenAI chat system, either
:obj:`"system"`, :obj:`"user"`, or :obj:`"assistant"`.
content (str): The content of the message.
"""
role_name: str
role_type: RoleType
meta_dict: Optional[Dict[str, str]]
role: str
content: str
def __getattribute__(self, name: str) -> Any:
r"""Get attribute override to delegate string methods to the
:obj:`content`.
Args:
name (str): The name of the attribute.
Returns:
Any: The attribute value.
"""
delegate_methods = [
method for method in dir(str) if not method.startswith('_')
]
if name in delegate_methods:
content = super().__getattribute__('content')
if isinstance(content, str):
content_method = getattr(content, name, None)
if callable(content_method):
def modify_arg(arg: Any) -> Any:
r"""Modify the argument for delegate method.
Args:
arg (Any): The argument value.
Returns:
Any: The modified argument value.
"""
if isinstance(arg, BaseMessage):
return arg.content
elif isinstance(arg, (list, tuple)):
return type(arg)(modify_arg(item) for item in arg)
else:
return arg
def wrapper(*args: Any, **kwargs: Any) -> Any:
r"""Wrapper function for delegate method.
Args:
*args (Any): Variable length argument list.
**kwargs (Any): Arbitrary keyword arguments.
Returns:
Any: The result of the delegate method.
"""
modified_args = [modify_arg(arg) for arg in args]
modified_kwargs = {
k: modify_arg(v)
for k, v in kwargs.items()
}
output = content_method(*modified_args,
**modified_kwargs)
return self._create_new_instance(output) if isinstance(
output, str) else output
return wrapper
return super().__getattribute__(name)
def _create_new_instance(self, content: str) -> "BaseMessage":
r"""Create a new instance of the :obj:`BaseMessage` with updated
content.
Args:
content (str): The new content value.
Returns:
BaseMessage: The new instance of :obj:`BaseMessage`.
"""
return self.__class__(role_name=self.role_name,
role_type=self.role_type,
meta_dict=self.meta_dict, role=self.role,
content=content)
def __add__(self, other: Any) -> Union["BaseMessage", Any]:
r"""Addition operator override for :obj:`BaseMessage`.
Args:
other (Any): The value to be added with.
Returns:
Union[BaseMessage, Any]: The result of the addition.
"""
if isinstance(other, BaseMessage):
combined_content = self.content.__add__(other.content)
elif isinstance(other, str):
combined_content = self.content.__add__(other)
else:
raise TypeError(
f"Unsupported operand type(s) for +: '{type(self)}' and "
f"'{type(other)}'")
return self._create_new_instance(combined_content)
def __mul__(self, other: Any) -> Union["BaseMessage", Any]:
r"""Multiplication operator override for :obj:`BaseMessage`.
Args:
other (Any): The value to be multiplied with.
Returns:
Union[BaseMessage, Any]: The result of the multiplication.
"""
if isinstance(other, int):
multiplied_content = self.content.__mul__(other)
return self._create_new_instance(multiplied_content)
else:
raise TypeError(
f"Unsupported operand type(s) for *: '{type(self)}' and "
f"'{type(other)}'")
def __len__(self) -> int:
r"""Length operator override for :obj:`BaseMessage`.
Returns:
int: The length of the content.
"""
return len(self.content)
def __contains__(self, item: str) -> bool:
r"""Contains operator override for :obj:`BaseMessage`.
Args:
item (str): The item to check for containment.
Returns:
bool: :obj:`True` if the item is contained in the content,
:obj:`False` otherwise.
"""
return item in self.content
def token_len(self, model: ModelType = ModelType.GPT_3_5_TURBO) -> int:
r"""Calculate the token length of the message for the specified model.
Args:
model (ModelType, optional): The model type to calculate the token
length. (default: :obj:`ModelType.GPT_3_5_TURBO`)
Returns:
int: The token length of the message.
"""
from camel.utils import num_tokens_from_messages
return num_tokens_from_messages([self.to_openai_chat_message()], model)
def extract_text_and_code_prompts(
self) -> Tuple[List[TextPrompt], List[CodePrompt]]:
r"""Extract text and code prompts from the message content.
Returns:
Tuple[List[TextPrompt], List[CodePrompt]]: A tuple containing a
list of text prompts and a list of code prompts extracted
from the content.
"""
text_prompts: List[TextPrompt] = []
code_prompts: List[CodePrompt] = []
lines = self.content.split("\n")
idx = 0
start_idx = 0
while idx < len(lines):
while idx < len(lines) and (
not lines[idx].lstrip().startswith("```")):
idx += 1
text = "\n".join(lines[start_idx:idx]).strip()
text_prompts.append(TextPrompt(text))
if idx >= len(lines):
break
code_type = lines[idx].strip()[3:].strip()
idx += 1
start_idx = idx
while not lines[idx].lstrip().startswith("```"):
idx += 1
code = "\n".join(lines[start_idx:idx]).strip()
code_prompts.append(CodePrompt(code, code_type=code_type))
idx += 1
start_idx = idx
return text_prompts, code_prompts
def to_openai_message(self, role: Optional[str] = None) -> OpenAIMessage:
r"""Converts the message to an :obj:`OpenAIMessage` object.
Args:
role (Optional[str]): The role of the message in OpenAI chat
system, either :obj:`"system"`, :obj:`"user"`, or
obj:`"assistant"`. (default: :obj:`None`)
Returns:
OpenAIMessage: The converted :obj:`OpenAIMessage` object.
"""
role = role or self.role
if role not in {"system", "user", "assistant"}:
raise ValueError(f"Unrecognized role: {role}")
return {"role": role, "content": self.content}
def to_openai_chat_message(
self,
role: Optional[str] = None,
) -> OpenAIChatMessage:
r"""Converts the message to an :obj:`OpenAIChatMessage` object.
Args:
role (Optional[str]): The role of the message in OpenAI chat
system, either :obj:`"user"`, or :obj:`"assistant"`.
(default: :obj:`None`)
Returns:
OpenAIChatMessage: The converted :obj:`OpenAIChatMessage` object.
"""
role = role or self.role
if role not in {"user", "assistant"}:
raise ValueError(f"Unrecognized role: {role}")
return {"role": role, "content": self.content}
def to_openai_system_message(self) -> OpenAISystemMessage:
r"""Converts the message to an :obj:`OpenAISystemMessage` object.
Returns:
OpenAISystemMessage: The converted :obj:`OpenAISystemMessage`
object.
"""
return {"role": "system", "content": self.content}
def to_openai_user_message(self) -> OpenAIUserMessage:
r"""Converts the message to an :obj:`OpenAIUserMessage` object.
Returns:
OpenAIUserMessage: The converted :obj:`OpenAIUserMessage` object.
"""
return {"role": "user", "content": self.content}
def to_openai_assistant_message(self) -> OpenAIAssistantMessage:
r"""Converts the message to an :obj:`OpenAIAssistantMessage` object.
Returns:
OpenAIAssistantMessage: The converted :obj:`OpenAIAssistantMessage`
object.
"""
return {"role": "assistant", "content": self.content}
def to_dict(self) -> Dict:
r"""Converts the message to a dictionary.
Returns:
dict: The converted dictionary.
"""
return {
"role_name": self.role_name,
"role_type": self.role_type.name,
**(self.meta_dict or {}),
"role": self.role,
"content": self.content,
}