diff --git a/swarms/agents/base.py b/swarms/agents/base.py index b8a692ee..f99bed05 100644 --- a/swarms/agents/base.py +++ b/swarms/agents/base.py @@ -10,13 +10,14 @@ from swarms.agents.prompts.prompt_generator import FINISH_NAME from swarms.agents.models.base import AbstractModel from swarms.agents.prompts.agent_output_parser import AgentOutputParser from swarms.agents.prompts.agent_prompt_auto import PromptConstructor, MessageFormatter +from swarms.agents.prompts.agent_prompt import AIMessage, HumanMessage, SystemMessage from langchain.chains.llm import LLMChain from langchain.memory import ChatMessageHistory -from langchain.schema import (BaseChatMessageHistory, Document,) -from langchain.schema.messages import AIMessage, HumanMessage, SystemMessage +from langchain.schema import (BaseChatMessageHistory, Document) + from langchain.tools.base import BaseTool from langchain.vectorstores.base import VectorStoreRetriever diff --git a/swarms/agents/prompts/chat_prompt.py b/swarms/agents/prompts/chat_prompt.py new file mode 100644 index 00000000..2334ce61 --- /dev/null +++ b/swarms/agents/prompts/chat_prompt.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from abc import abstractmethod +from typing import Any, Dict, List, Sequence + +from pydantic import Field + + +class Message: + """ + The base abstract Message class. + Messages are the inputs and outputs of ChatModels. + """ + def __init__(self, content: str, role: str, additional_kwargs: Dict = None): + self.content = content + self.role = role + self.additional_kwargs = additional_kwargs if additional_kwargs else {} + + @abstractmethod + def get_type(self) -> str: + pass + + +class HumanMessage(Message): + """ + A Message from a human. + """ + def __init__(self, content: str, role: str = "Human", additional_kwargs: Dict = None, example: bool = False): + super().__init__(content, role, additional_kwargs) + self.example = example + + def get_type(self) -> str: + return "human" + + +class AIMessage(Message): + """ + A Message from an AI. + """ + def __init__(self, content: str, role: str = "AI", additional_kwargs: Dict = None, example: bool = False): + super().__init__(content, role, additional_kwargs) + self.example = example + + def get_type(self) -> str: + return "ai" + + +class SystemMessage(Message): + """ + A Message for priming AI behavior, usually passed in as the first of a sequence + of input messages. + """ + def __init__(self, content: str, role: str = "System", additional_kwargs: Dict = None): + super().__init__(content, role, additional_kwargs) + + def get_type(self) -> str: + return "system" + + +class FunctionMessage(Message): + """ + A Message for passing the result of executing a function back to a model. + """ + def __init__(self, content: str, role: str = "Function", name: str, additional_kwargs: Dict = None): + super().__init__(content, role, additional_kwargs) + self.name = name + + def get_type(self) -> str: + return "function" + + +class ChatMessage(Message): + """ + A Message that can be assigned an arbitrary speaker (i.e. role). + """ + def __init__(self, content: str, role: str, additional_kwargs: Dict = None): + super().__init__(content, role, additional_kwargs) + + def get_type(self) -> str: + return "chat" + + +def get_buffer_string( + messages: Sequence[Message], human_prefix: str = "Human", ai_prefix: str = "AI" +) -> str: + string_messages = [] + for m in messages: + message = f"{m.role}: {m.content}" + if isinstance(m, AIMessage) and "function_call" in m.additional_kwargs: + message += f"{m.additional_kwargs['function_call']}" + string_messages.append(message) + + return "\n".join(string_messages) + + +def message_to_dict(message: Message) -> dict: + return {"type": message.get_type(), "data": message.__dict__} + + +def messages_to_dict(messages: Sequence[Message]) -> List[dict]: + return [message_to_dict(m) for m in messages] + + +def message_from_dict(message: dict) -> Message: + _type = message["type"] + if _type == "human": + return HumanMessage(**message["data"]) + elif _type == "ai": + return AIMessage(**message["data"]) + elif _type == "system": + return SystemMessage(**message["data"]) + elif _type == "chat": + return ChatMessage(**message["data"]) + elif _type == "function": + return FunctionMessage(**message["data"]) + else: + raise ValueError(f"Got unexpected message type: {_type}") + + +def messages_from_dict(messages: List[dict]) -> List[Message]: + return [message_from_dict(m) for m in messages]