import datetime import json from termcolor import colored from swarms.memory.base_db import AbstractDatabase from swarms.structs.base import BaseStructure class Conversation(BaseStructure): """ Conversation class Attributes: time_enabled (bool): whether to enable time conversation_history (list): list of messages in the conversation Examples: >>> conv = Conversation() >>> conv.add("user", "Hello, world!") >>> conv.add("assistant", "Hello, user!") >>> conv.display_conversation() user: Hello, world! """ def __init__( self, time_enabled: bool = False, database: AbstractDatabase = None, autosave: bool = True, save_filepath: str = "/runs/conversation.json", *args, **kwargs, ): super().__init__() self.time_enabled = time_enabled self.database = database self.autosave = autosave self.save_filepath = save_filepath self.conversation_history = [] def add(self, role: str, content: str, *args, **kwargs): """Add a message to the conversation history Args: role (str): The role of the speaker content (str): The content of the message """ if self.time_enabled: now = datetime.datetime.now() timestamp = now.strftime("%Y-%m-%d %H:%M:%S") message = { "role": role, "content": content, "timestamp": timestamp, } else: message = { "role": role, "content": content, } self.conversation_history.append(message) if self.autosave: self.save_as_json(self.save_filepath) def delete(self, index: str): """Delete a message from the conversation history Args: index (str): index of the message to delete """ self.conversation_history.pop(index) def update(self, index: str, role, content): """Update a message in the conversation history Args: index (str): index of the message to update role (_type_): role of the speaker content (_type_): content of the message """ self.conversation_history[index] = { "role": role, "content": content, } def query(self, index: str): """Query a message in the conversation history Args: index (str): index of the message to query Returns: str: the message """ return self.conversation_history[index] def search(self, keyword: str): """Search for a message in the conversation history Args: keyword (str): Keyword to search for Returns: str: description """ return [ msg for msg in self.conversation_history if keyword in msg["content"] ] def display_conversation(self, detailed: bool = False): """Display the conversation history Args: detailed (bool, optional): detailed. Defaults to False. """ role_to_color = { "system": "red", "user": "green", "assistant": "blue", "function": "magenta", } for message in self.conversation_history: print( colored( f"{message['role']}: {message['content']}\n\n", role_to_color[message["role"]], ) ) def export_conversation(self, filename: str, *args, **kwargs): """Export the conversation history to a file Args: filename (str): filename to export to """ with open(filename, "w") as f: for message in self.conversation_history: f.write(f"{message['role']}: {message['content']}\n") def import_conversation(self, filename: str): """Import a conversation history from a file Args: filename (str): filename to import from """ with open(filename, "r") as f: for line in f: role, content = line.split(": ", 1) self.add(role, content.strip()) def count_messages_by_role(self): """Count the number of messages by role""" counts = { "system": 0, "user": 0, "assistant": 0, "function": 0, } for message in self.conversation_history: counts[message["role"]] += 1 return counts def return_history_as_string(self): """Return the conversation history as a string Returns: str: the conversation history """ return "\n".join( [ f"{message['role']}: {message['content']}\n\n" for message in self.conversation_history ] ) def save_as_json(self, filename: str): """Save the conversation history as a JSON file Args: filename (str): Save the conversation history as a JSON file """ # Save the conversation history as a JSON file with open(filename, "w") as f: json.dump(self.conversation_history, f) def load_from_json(self, filename: str): """Load the conversation history from a JSON file Args: filename (str): filename to load from """ # Load the conversation history from a JSON file with open(filename, "r") as f: self.conversation_history = json.load(f) def search_keyword_in_conversation(self, keyword: str): """Search for a keyword in the conversation history Args: keyword (str): keyword to search for Returns: str: description """ return [ msg for msg in self.conversation_history if keyword in msg["content"] ] def pretty_print_conversation(self, messages): """Pretty print the conversation history Args: messages (str): messages to print """ role_to_color = { "system": "red", "user": "green", "assistant": "blue", "tool": "magenta", } for message in messages: if message["role"] == "system": print( colored( f"system: {message['content']}\n", role_to_color[message["role"]], ) ) elif message["role"] == "user": print( colored( f"user: {message['content']}\n", role_to_color[message["role"]], ) ) elif message["role"] == "assistant" and message.get( "function_call" ): print( colored( f"assistant: {message['function_call']}\n", role_to_color[message["role"]], ) ) elif message["role"] == "assistant" and not message.get( "function_call" ): print( colored( f"assistant: {message['content']}\n", role_to_color[message["role"]], ) ) elif message["role"] == "tool": print( colored( ( f"function ({message['name']}):" f" {message['content']}\n" ), role_to_color[message["role"]], ) ) def add_to_database(self, *args, **kwargs): """Add the conversation history to the database""" self.database.add("conversation", self.conversation_history) def query_from_database(self, query, *args, **kwargs): """Query the conversation history from the database""" return self.database.query("conversation", query) def delete_from_database(self, *args, **kwargs): """Delete the conversation history from the database""" self.database.delete("conversation") def update_from_database(self, *args, **kwargs): """Update the conversation history from the database""" self.database.update( "conversation", self.conversation_history ) def get_from_database(self, *args, **kwargs): """Get the conversation history from the database""" return self.database.get("conversation") def execute_query_from_database(self, query, *args, **kwargs): """Execute a query on the database""" return self.database.execute_query(query) def fetch_all_from_database(self, *args, **kwargs): """Fetch all from the database""" return self.database.fetch_all() def fetch_one_from_database(self, *args, **kwargs): """Fetch one from the database""" return self.database.fetch_one()