From 300f26880eaf86ddf4e5cace0e76f2d5562b3ea1 Mon Sep 17 00:00:00 2001 From: Kye Date: Thu, 21 Dec 2023 14:39:18 -0500 Subject: [PATCH] [FEAT][Conversation] --- docs/swarms/structs/conversation.md | 132 +++++++++++++ mkdocs.yml | 1 + swarms/models/openai_function_caller.py | 13 +- swarms/structs/__init__.py | 3 +- swarms/structs/conversation.py | 242 ++++++++++++++++++++++++ tests/structs/test_conversation.py | 241 +++++++++++++++++++++++ 6 files changed, 625 insertions(+), 7 deletions(-) create mode 100644 docs/swarms/structs/conversation.md create mode 100644 swarms/structs/conversation.py create mode 100644 tests/structs/test_conversation.py diff --git a/docs/swarms/structs/conversation.md b/docs/swarms/structs/conversation.md new file mode 100644 index 00000000..723944bf --- /dev/null +++ b/docs/swarms/structs/conversation.md @@ -0,0 +1,132 @@ +# Conversation Module Documentation + +## Table of Contents + +1. [Introduction](#introduction) +2. [Installation](#installation) +3. [Class: Conversation](#class-conversation) + - [Attributes](#attributes) + - [Methods](#methods) +4. [Usage Examples](#usage-examples) + - [Example 1: Creating a Conversation](#example-1-creating-a-conversation) + - [Example 2: Adding Messages](#example-2-adding-messages) + - [Example 3: Displaying and Exporting Conversation](#example-3-displaying-and-exporting-conversation) + - [Example 4: Counting Messages by Role](#example-4-counting-messages-by-role) + - [Example 5: Loading and Searching](#example-5-loading-and-searching) +5. [Additional Information](#additional-information) +6. [References](#references) + +--- + +## 1. Introduction + +The Conversation module provides a versatile and extensible structure for managing and analyzing text-based conversations. Whether you're developing a chatbot, analyzing customer support interactions, or conducting research on dialogues, this module simplifies the process of handling conversation data. + +With the Conversation module, you can add, delete, update, query, and search for messages within a conversation. You can also display, export, and import conversation history, making it an essential tool for various applications. + +## 2. Installation + +To use the Conversation module, you need to have Python installed on your system. Additionally, you can install the required dependencies using pip: + +```bash +pip install termcolor +``` + +Once you have the dependencies installed, you can import the Conversation module into your Python code. + +```python +from swarms.structs.conversation import Conversation +``` + +## 3. Class: Conversation + +The Conversation class is the core of this module. It allows you to create and manipulate conversation histories. Below are the attributes and methods provided by this class. + +### Attributes + +- `time_enabled` (bool): Indicates whether timestamps are enabled for messages in the conversation. +- `conversation_history` (list): A list that stores the conversation history as a collection of messages. + +### Methods + +The Conversation class provides the following methods: + +- `add(role: str, content: str, *args, **kwargs)`: Adds a message to the conversation history. +- `delete(index: str)`: Deletes a message from the conversation history. +- `update(index: str, role, content)`: Updates a message in the conversation history. +- `query(index: str)`: Queries a message in the conversation history. +- `search(keyword: str)`: Searches for messages containing a specific keyword. +- `display_conversation(detailed: bool = False)`: Displays the conversation history. +- `export_conversation(filename: str)`: Exports the conversation history to a file. +- `import_conversation(filename: str)`: Imports a conversation history from a file. +- `count_messages_by_role()`: Counts the number of messages by role. +- `return_history_as_string()`: Returns the conversation history as a string. +- `save_as_json(filename: str)`: Saves the conversation history as a JSON file. +- `load_from_json(filename: str)`: Loads the conversation history from a JSON file. +- `search_keyword_in_conversation(keyword: str)`: Searches for a keyword in the conversation history. +- `pretty_print_conversation(messages)`: Pretty prints the conversation history. + +## 4. Usage Examples + +In this section, we'll provide practical examples of how to use the Conversation module to manage and analyze conversation data. + +### Example 1: Creating a Conversation + +Let's start by creating a Conversation object and enabling timestamps for messages: + +```python +conversation = Conversation(time_enabled=True) +``` + +### Example 2: Adding Messages + +You can add messages to the conversation using the `add` method. Here's how to add a user message and an assistant response: + +```python +conversation.add("user", "Hello, how can I help you?") +conversation.add("assistant", "Hi there! I'm here to assist you.") +``` + +### Example 3: Displaying and Exporting Conversation + +You can display the conversation history and export it to a file. Let's see how to do this: + +```python +# Display the conversation +conversation.display_conversation() + +# Export the conversation to a file +conversation.export_conversation("conversation_history.txt") +``` + +### Example 4: Counting Messages by Role + +You can count the number of messages by role (e.g., user, assistant, system) using the `count_messages_by_role` method: + +```python +message_counts = conversation.count_messages_by_role() +print(message_counts) +``` + +### Example 5: Loading and Searching + +You can load a conversation from a file and search for messages containing a specific keyword: + +```python +# Load conversation from a file +conversation.load_from_json("saved_conversation.json") + +# Search for messages containing the keyword "help" +results = conversation.search("help") +print(results) +``` + +## 5. Additional Information + +- The Conversation module is designed to provide flexibility and ease of use for managing and analyzing text-based conversations. +- You can extend the module by adding custom functionality or integrating it into your chatbot or natural language processing applications. + +## 6. References + +For more information on the Conversation module and its usage, refer to the official documentation and examples. + diff --git a/mkdocs.yml b/mkdocs.yml index eeb64c04..d1620c27 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -103,6 +103,7 @@ nav: - AutoScaler: "swarms/swarms/autoscaler.md" - Agent: "swarms/structs/agent.md" - SequentialWorkflow: 'swarms/structs/sequential_workflow.md' + - Conversation: "swarms/structs/conversation.md" - swarms.memory: - Weaviate: "swarms/memory/weaviate.md" - PineconDB: "swarms/memory/pinecone.md" diff --git a/swarms/models/openai_function_caller.py b/swarms/models/openai_function_caller.py index 6542e457..e6822793 100644 --- a/swarms/models/openai_function_caller.py +++ b/swarms/models/openai_function_caller.py @@ -234,10 +234,10 @@ class OpenAIFunctionCaller: ) ) - def call(self, prompt: str) -> Dict: - response = openai.Completion.create( + def call(self, task: str, *args, **kwargs) -> Dict: + return openai.Completion.create( engine=self.model, - prompt=prompt, + prompt=task, max_tokens=self.max_tokens, temperature=self.temperature, top_p=self.top_p, @@ -253,9 +253,10 @@ class OpenAIFunctionCaller: user=self.user, messages=self.messages, timeout_sec=self.timeout_sec, + *args, + **kwargs, ) - return response - def run(self, prompt: str) -> str: - response = self.call(prompt) + def run(self, task: str, *args, **kwargs) -> str: + response = self.call(task, *args, **kwargs) return response["choices"][0]["text"].strip() diff --git a/swarms/structs/__init__.py b/swarms/structs/__init__.py index e389ed76..b3cb4412 100644 --- a/swarms/structs/__init__.py +++ b/swarms/structs/__init__.py @@ -1,5 +1,6 @@ from swarms.structs.agent import Agent from swarms.structs.sequential_workflow import SequentialWorkflow from swarms.structs.autoscaler import AutoScaler +from swarms.structs.conversation import Conversation -__all__ = ["Agent", "SequentialWorkflow", "AutoScaler"] +__all__ = ["Agent", "SequentialWorkflow", "AutoScaler", "Conversation"] diff --git a/swarms/structs/conversation.py b/swarms/structs/conversation.py new file mode 100644 index 00000000..dcf807ab --- /dev/null +++ b/swarms/structs/conversation.py @@ -0,0 +1,242 @@ +import json +import datetime + +from termcolor import colored + +from swarms.structs.base import BaseStructure + + +class Conversation(BaseStructure): + def __init__(self, time_enabled: bool = False, *args, **kwargs): + super().__init__() + self.time_enabled = time_enabled + 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) + + 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): + """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"]], + ) + ) diff --git a/tests/structs/test_conversation.py b/tests/structs/test_conversation.py new file mode 100644 index 00000000..84673a42 --- /dev/null +++ b/tests/structs/test_conversation.py @@ -0,0 +1,241 @@ +import pytest +from swarms.structs.conversation import Conversation + + +@pytest.fixture +def conversation(): + conv = Conversation() + conv.add("user", "Hello, world!") + conv.add("assistant", "Hello, user!") + return conv + + +def test_add_message(): + conv = Conversation() + conv.add("user", "Hello, world!") + assert len(conv.conversation_history) == 1 + assert conv.conversation_history[0]["role"] == "user" + assert conv.conversation_history[0]["content"] == "Hello, world!" + + +def test_add_message_with_time(): + conv = Conversation(time_enabled=True) + conv.add("user", "Hello, world!") + assert len(conv.conversation_history) == 1 + assert conv.conversation_history[0]["role"] == "user" + assert conv.conversation_history[0]["content"] == "Hello, world!" + assert "timestamp" in conv.conversation_history[0] + + +def test_delete_message(): + conv = Conversation() + conv.add("user", "Hello, world!") + conv.delete(0) + assert len(conv.conversation_history) == 0 + + +def test_delete_message_out_of_bounds(): + conv = Conversation() + conv.add("user", "Hello, world!") + with pytest.raises(IndexError): + conv.delete(1) + + +def test_update_message(): + conv = Conversation() + conv.add("user", "Hello, world!") + conv.update(0, "assistant", "Hello, user!") + assert len(conv.conversation_history) == 1 + assert conv.conversation_history[0]["role"] == "assistant" + assert conv.conversation_history[0]["content"] == "Hello, user!" + + +def test_update_message_out_of_bounds(): + conv = Conversation() + conv.add("user", "Hello, world!") + with pytest.raises(IndexError): + conv.update(1, "assistant", "Hello, user!") + + +def test_return_history_as_string_with_messages(conversation): + result = conversation.return_history_as_string() + assert result is not None + + +def test_return_history_as_string_with_no_messages(): + conv = Conversation() + result = conv.return_history_as_string() + assert result == "" + + +@pytest.mark.parametrize( + "role, content", + [ + ("user", "Hello, world!"), + ("assistant", "Hello, user!"), + ("system", "System message"), + ("function", "Function message"), + ], +) +def test_return_history_as_string_with_different_roles(role, content): + conv = Conversation() + conv.add(role, content) + result = conv.return_history_as_string() + expected = f"{role}: {content}\n\n" + assert result == expected + + +@pytest.mark.parametrize("message_count", range(1, 11)) +def test_return_history_as_string_with_multiple_messages( + message_count, +): + conv = Conversation() + for i in range(message_count): + conv.add("user", f"Message {i + 1}") + result = conv.return_history_as_string() + expected = "".join( + [f"user: Message {i + 1}\n\n" for i in range(message_count)] + ) + assert result == expected + + +@pytest.mark.parametrize( + "content", + [ + "Hello, world!", + "This is a longer message with multiple words.", + "This message\nhas multiple\nlines.", + "This message has special characters: !@#$%^&*()", + "This message has unicode characters: 你好，世界！", + ], +) +def test_return_history_as_string_with_different_contents(content): + conv = Conversation() + conv.add("user", content) + result = conv.return_history_as_string() + expected = f"user: {content}\n\n" + assert result == expected + + +def test_return_history_as_string_with_large_message(conversation): + large_message = "Hello, world! " * 10000 # 10,000 repetitions + conversation.add("user", large_message) + result = conversation.return_history_as_string() + expected = ( + "user: Hello, world!\n\nassistant: Hello, user!\n\nuser:" + f" {large_message}\n\n" + ) + assert result == expected + + +def test_search_keyword_in_conversation(conversation): + result = conversation.search_keyword_in_conversation("Hello") + assert len(result) == 2 + assert result[0]["content"] == "Hello, world!" + assert result[1]["content"] == "Hello, user!" + + +def test_export_import_conversation(conversation, tmp_path): + filename = tmp_path / "conversation.txt" + conversation.export_conversation(filename) + new_conversation = Conversation() + new_conversation.import_conversation(filename) + assert ( + new_conversation.return_history_as_string() + == conversation.return_history_as_string() + ) + + +def test_count_messages_by_role(conversation): + counts = conversation.count_messages_by_role() + assert counts["user"] == 1 + assert counts["assistant"] == 1 + + +def test_display_conversation(capsys, conversation): + conversation.display_conversation() + captured = capsys.readouterr() + assert "user: Hello, world!\n\n" in captured.out + assert "assistant: Hello, user!\n\n" in captured.out + + +def test_display_conversation_detailed(capsys, conversation): + conversation.display_conversation(detailed=True) + captured = capsys.readouterr() + assert "user: Hello, world!\n\n" in captured.out + assert "assistant: Hello, user!\n\n" in captured.out + + +def test_search(): + conv = Conversation() + conv.add("user", "Hello, world!") + conv.add("assistant", "Hello, user!") + results = conv.search("Hello") + assert len(results) == 2 + assert results[0]["content"] == "Hello, world!" + assert results[1]["content"] == "Hello, user!" + + +def test_return_history_as_string(): + conv = Conversation() + conv.add("user", "Hello, world!") + conv.add("assistant", "Hello, user!") + result = conv.return_history_as_string() + expected = "user: Hello, world!\n\nassistant: Hello, user!\n\n" + assert result == expected + + +def test_search_no_results(): + conv = Conversation() + conv.add("user", "Hello, world!") + conv.add("assistant", "Hello, user!") + results = conv.search("Goodbye") + assert len(results) == 0 + + +def test_search_case_insensitive(): + conv = Conversation() + conv.add("user", "Hello, world!") + conv.add("assistant", "Hello, user!") + results = conv.search("hello") + assert len(results) == 2 + assert results[0]["content"] == "Hello, world!" + assert results[1]["content"] == "Hello, user!" + + +def test_search_multiple_occurrences(): + conv = Conversation() + conv.add("user", "Hello, world! Hello, world!") + conv.add("assistant", "Hello, user!") + results = conv.search("Hello") + assert len(results) == 2 + assert results[0]["content"] == "Hello, world! Hello, world!" + assert results[1]["content"] == "Hello, user!" + + +def test_query_no_results(): + conv = Conversation() + conv.add("user", "Hello, world!") + conv.add("assistant", "Hello, user!") + results = conv.query("Goodbye") + assert len(results) == 0 + + +def test_query_case_insensitive(): + conv = Conversation() + conv.add("user", "Hello, world!") + conv.add("assistant", "Hello, user!") + results = conv.query("hello") + assert len(results) == 2 + assert results[0]["content"] == "Hello, world!" + assert results[1]["content"] == "Hello, user!" + + +def test_query_multiple_occurrences(): + conv = Conversation() + conv.add("user", "Hello, world! Hello, world!") + conv.add("assistant", "Hello, user!") + results = conv.query("Hello") + assert len(results) == 2 + assert results[0]["content"] == "Hello, world! Hello, world!" + assert results[1]["content"] == "Hello, user!"