[FEAT][Conversation]

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 <a name="introduction"></a>
+
+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 <a name="installation"></a>
+
+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 <a name="class-conversation"></a>
+
+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 <a name="attributes"></a>
+
+- `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 <a name="methods"></a>
+
+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 <a name="usage-examples"></a>
+
+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 <a name="example-1-creating-a-conversation"></a>
+
+Let's start by creating a Conversation object and enabling timestamps for messages:
+
+```python
+conversation = Conversation(time_enabled=True)
+```
+
+### Example 2: Adding Messages <a name="example-2-adding-messages"></a>
+
+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 <a name="example-3-displaying-and-exporting-conversation"></a>
+
+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 <a name="example-4-counting-messages-by-role"></a>
+
+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 <a name="example-5-loading-and-searching"></a>
+
+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 <a name="additional-information"></a>
+
+- 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 <a name="references"></a>
+
+For more information on the Conversation module and its usage, refer to the official documentation and examples.
+

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"

swarms/models/openai_function_caller.py

@@ -234,10 +234,10 @@ class OpenAIFunctionCaller:
                     )
                 )
 
-    def call(self, prompt: str) -> Dict:
+    def call(self, task: str, *args, **kwargs) -> Dict:
-        response = openai.Completion.create(
+        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:
+    def run(self, task: str, *args, **kwargs) -> str:
-        response = self.call(prompt)
+        response = self.call(task, *args, **kwargs)
         return response["choices"][0]["text"].strip()

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"]

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"]],
+                    )
+                )

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!"