swarms 7.7.8__py3-none-any.whl → 7.8.0__py3-none-any.whl

swarms/communication/sqlite_wrap.py

@@ -1,15 +1,19 @@
 import sqlite3
 import json
 import datetime
-from typing import List, Optional, Union, Dict
+from typing import List, Optional, Union, Dict, Any
 from pathlib import Path
 import threading
 from contextlib import contextmanager
 import logging
-from dataclasses import dataclass
-from enum import Enum
 import uuid
 import yaml
+from swarms.communication.base_communication import (
+    BaseCommunication,
+    Message,
+    MessageType,
+)
+from typing import Callable
 
 try:
     from loguru import logger
@@ -19,32 +23,7 @@ except ImportError:
     LOGURU_AVAILABLE = False
 
 
-class MessageType(Enum):
-    """Enum for different types of messages in the conversation."""
-
-    SYSTEM = "system"
-    USER = "user"
-    ASSISTANT = "assistant"
-    FUNCTION = "function"
-    TOOL = "tool"
-
-
-@dataclass
-class Message:
-    """Data class representing a message in the conversation."""
-
-    role: str
-    content: Union[str, dict, list]
-    timestamp: Optional[str] = None
-    message_type: Optional[MessageType] = None
-    metadata: Optional[Dict] = None
-    token_count: Optional[int] = None
-
-    class Config:
-        arbitrary_types_allowed = True
-
-
-class SQLiteConversation:
+class SQLiteConversation(BaseCommunication):
     """
     A production-grade SQLite wrapper class for managing conversation history.
     This class provides persistent storage for conversations with various features
@@ -63,7 +42,21 @@ class SQLiteConversation:
 
     def __init__(
         self,
-        db_path: str = "conversations.db",
+        system_prompt: Optional[str] = None,
+        time_enabled: bool = False,
+        autosave: bool = False,
+        save_filepath: str = None,
+        tokenizer: Any = None,
+        context_length: int = 8192,
+        rules: str = None,
+        custom_rules_prompt: str = None,
+        user: str = "User:",
+        auto_save: bool = True,
+        save_as_yaml: bool = True,
+        save_as_json_bool: bool = False,
+        token_count: bool = True,
+        cache_enabled: bool = True,
+        db_path: Union[str, Path] = None,
         table_name: str = "conversations",
         enable_timestamps: bool = True,
         enable_logging: bool = True,
@@ -72,19 +65,31 @@ class SQLiteConversation:
         connection_timeout: float = 5.0,
         **kwargs,
     ):
-        """
-        Initialize the SQLite conversation manager.
+        super().__init__(
+            system_prompt=system_prompt,
+            time_enabled=time_enabled,
+            autosave=autosave,
+            save_filepath=save_filepath,
+            tokenizer=tokenizer,
+            context_length=context_length,
+            rules=rules,
+            custom_rules_prompt=custom_rules_prompt,
+            user=user,
+            auto_save=auto_save,
+            save_as_yaml=save_as_yaml,
+            save_as_json_bool=save_as_json_bool,
+            token_count=token_count,
+            cache_enabled=cache_enabled,
+        )
 
-        Args:
-            db_path (str): Path to the SQLite database file
-            table_name (str): Name of the table to store conversations
-            enable_timestamps (bool): Whether to track message timestamps
-            enable_logging (bool): Whether to enable logging
-            use_loguru (bool): Whether to use loguru for logging
-            max_retries (int): Maximum number of retries for database operations
-            connection_timeout (float): Timeout for database connections
-        """
+        # Calculate default db_path if not provided
+        if db_path is None:
+            db_path = self.get_default_db_path("conversations.sqlite")
         self.db_path = Path(db_path)
+
+        # Ensure parent directory exists
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+
         self.table_name = table_name
         self.enable_timestamps = enable_timestamps
         self.enable_logging = enable_logging
@@ -92,9 +97,7 @@ class SQLiteConversation:
         self.max_retries = max_retries
         self.connection_timeout = connection_timeout
         self._lock = threading.Lock()
-        self.current_conversation_id = (
-            self._generate_conversation_id()
-        )
+        self.tokenizer = tokenizer
 
         # Setup logging
         if self.enable_logging:
@@ -112,6 +115,7 @@ class SQLiteConversation:
 
         # Initialize database
         self._init_db()
+        self.start_new_conversation()
 
     def _generate_conversation_id(self) -> str:
         """Generate a unique conversation ID using UUID and timestamp."""
@@ -811,3 +815,502 @@ class SQLiteConversation:
                 "total_tokens": row["total_tokens"],
                 "roles": self.count_messages_by_role(),
             }
+
+    def delete(self, index: str):
+        """Delete a message from the conversation history."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"DELETE FROM {self.table_name} WHERE id = ? AND conversation_id = ?",
+                (index, self.current_conversation_id),
+            )
+            conn.commit()
+
+    def update(
+        self, index: str, role: str, content: Union[str, dict]
+    ):
+        """Update a message in the conversation history."""
+        if isinstance(content, (dict, list)):
+            content = json.dumps(content)
+
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"""
+                UPDATE {self.table_name}
+                SET role = ?, content = ?
+                WHERE id = ? AND conversation_id = ?
+                """,
+                (role, content, index, self.current_conversation_id),
+            )
+            conn.commit()
+
+    def query(self, index: str) -> Dict:
+        """Query a message in the conversation history."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"""
+                SELECT * FROM {self.table_name}
+                WHERE id = ? AND conversation_id = ?
+                """,
+                (index, self.current_conversation_id),
+            )
+            row = cursor.fetchone()
+
+            if not row:
+                return {}
+
+            content = row["content"]
+            try:
+                content = json.loads(content)
+            except json.JSONDecodeError:
+                pass
+
+            return {
+                "role": row["role"],
+                "content": content,
+                "timestamp": row["timestamp"],
+                "message_type": row["message_type"],
+                "metadata": (
+                    json.loads(row["metadata"])
+                    if row["metadata"]
+                    else None
+                ),
+                "token_count": row["token_count"],
+            }
+
+    def search(self, keyword: str) -> List[Dict]:
+        """Search for messages containing a keyword."""
+        return self.search_messages(keyword)
+
+    def display_conversation(self, detailed: bool = False):
+        """Display the conversation history."""
+        print(self.get_str())
+
+    def export_conversation(self, filename: str):
+        """Export the conversation history to a file."""
+        self.save_as_json(filename)
+
+    def import_conversation(self, filename: str):
+        """Import a conversation history from a file."""
+        self.load_from_json(filename)
+
+    def return_history_as_string(self) -> str:
+        """Return the conversation history as a string."""
+        return self.get_str()
+
+    def clear(self):
+        """Clear the conversation history."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"DELETE FROM {self.table_name} WHERE conversation_id = ?",
+                (self.current_conversation_id,),
+            )
+            conn.commit()
+
+    def get_conversation_timeline_dict(self) -> Dict[str, List[Dict]]:
+        """Get the conversation organized by timestamps."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"""
+                SELECT 
+                    DATE(timestamp) as date,
+                    role,
+                    content,
+                    timestamp,
+                    message_type,
+                    metadata,
+                    token_count
+                FROM {self.table_name}
+                WHERE conversation_id = ?
+                ORDER BY timestamp ASC
+                """,
+                (self.current_conversation_id,),
+            )
+
+            timeline_dict = {}
+            for row in cursor.fetchall():
+                date = row["date"]
+                content = row["content"]
+                try:
+                    content = json.loads(content)
+                except json.JSONDecodeError:
+                    pass
+
+                message = {
+                    "role": row["role"],
+                    "content": content,
+                    "timestamp": row["timestamp"],
+                    "message_type": row["message_type"],
+                    "metadata": (
+                        json.loads(row["metadata"])
+                        if row["metadata"]
+                        else None
+                    ),
+                    "token_count": row["token_count"],
+                }
+
+                if date not in timeline_dict:
+                    timeline_dict[date] = []
+                timeline_dict[date].append(message)
+
+            return timeline_dict
+
+    def truncate_memory_with_tokenizer(self):
+        """Truncate the conversation history based on token count."""
+        if not self.tokenizer:
+            return
+
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"""
+                SELECT id, content, token_count
+                FROM {self.table_name}
+                WHERE conversation_id = ?
+                ORDER BY id ASC
+                """,
+                (self.current_conversation_id,),
+            )
+
+            total_tokens = 0
+            ids_to_keep = []
+
+            for row in cursor.fetchall():
+                token_count = row[
+                    "token_count"
+                ] or self.tokenizer.count_tokens(row["content"])
+                if total_tokens + token_count <= self.context_length:
+                    total_tokens += token_count
+                    ids_to_keep.append(row["id"])
+                else:
+                    break
+
+            if ids_to_keep:
+                ids_str = ",".join(map(str, ids_to_keep))
+                cursor.execute(
+                    f"""
+                    DELETE FROM {self.table_name}
+                    WHERE conversation_id = ?
+                    AND id NOT IN ({ids_str})
+                    """,
+                    (self.current_conversation_id,),
+                )
+                conn.commit()
+
+    def get_conversation_metadata_dict(self) -> Dict:
+        """Get detailed metadata about the conversation."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            # Get basic statistics
+            stats = self.get_statistics()
+
+            # Get message type distribution
+            cursor.execute(
+                f"""
+                SELECT message_type, COUNT(*) as count
+                FROM {self.table_name}
+                WHERE conversation_id = ?
+                GROUP BY message_type
+            """,
+                (self.current_conversation_id,),
+            )
+            type_dist = cursor.fetchall()
+
+            # Get average tokens per message
+            cursor.execute(
+                f"""
+                SELECT AVG(token_count) as avg_tokens
+                FROM {self.table_name}
+                WHERE conversation_id = ? AND token_count IS NOT NULL
+            """,
+                (self.current_conversation_id,),
+            )
+            avg_tokens = cursor.fetchone()
+
+            # Get message frequency by hour
+            cursor.execute(
+                f"""
+                SELECT 
+                    strftime('%H', timestamp) as hour,
+                    COUNT(*) as count
+                FROM {self.table_name}
+                WHERE conversation_id = ?
+                GROUP BY hour
+                ORDER BY hour
+            """,
+                (self.current_conversation_id,),
+            )
+            hourly_freq = cursor.fetchall()
+
+            return {
+                "conversation_id": self.current_conversation_id,
+                "basic_stats": stats,
+                "message_type_distribution": {
+                    row["message_type"]: row["count"]
+                    for row in type_dist
+                    if row["message_type"]
+                },
+                "average_tokens_per_message": (
+                    avg_tokens["avg_tokens"]
+                    if avg_tokens["avg_tokens"] is not None
+                    else 0
+                ),
+                "hourly_message_frequency": {
+                    row["hour"]: row["count"] for row in hourly_freq
+                },
+                "role_distribution": self.count_messages_by_role(),
+            }
+
+    def get_conversation_by_role_dict(self) -> Dict[str, List[Dict]]:
+        """Get the conversation organized by roles."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"""
+                SELECT role, content, timestamp, message_type, metadata, token_count
+                FROM {self.table_name}
+                WHERE conversation_id = ?
+                ORDER BY id ASC
+            """,
+                (self.current_conversation_id,),
+            )
+
+            role_dict = {}
+            for row in cursor.fetchall():
+                role = row["role"]
+                content = row["content"]
+                try:
+                    content = json.loads(content)
+                except json.JSONDecodeError:
+                    pass
+
+                message = {
+                    "content": content,
+                    "timestamp": row["timestamp"],
+                    "message_type": row["message_type"],
+                    "metadata": (
+                        json.loads(row["metadata"])
+                        if row["metadata"]
+                        else None
+                    ),
+                    "token_count": row["token_count"],
+                }
+
+                if role not in role_dict:
+                    role_dict[role] = []
+                role_dict[role].append(message)
+
+            return role_dict
+
+    def get_conversation_as_dict(self) -> Dict:
+        """Get the entire conversation as a dictionary with messages and metadata."""
+        messages = self.get_messages()
+        stats = self.get_statistics()
+
+        return {
+            "conversation_id": self.current_conversation_id,
+            "messages": messages,
+            "metadata": {
+                "total_messages": stats["total_messages"],
+                "unique_roles": stats["unique_roles"],
+                "total_tokens": stats["total_tokens"],
+                "first_message": stats["first_message"],
+                "last_message": stats["last_message"],
+                "roles": self.count_messages_by_role(),
+            },
+        }
+
+    def get_visible_messages(
+        self, agent: Callable, turn: int
+    ) -> List[Dict]:
+        """
+        Get the visible messages for a given agent and turn.
+
+        Args:
+            agent (Agent): The agent.
+            turn (int): The turn number.
+
+        Returns:
+            List[Dict]: The list of visible messages.
+        """
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"""
+                SELECT * FROM {self.table_name}
+                WHERE conversation_id = ?
+                AND json_extract(metadata, '$.turn') < ?
+                ORDER BY id ASC
+                """,
+                (self.current_conversation_id, turn),
+            )
+
+            visible_messages = []
+            for row in cursor.fetchall():
+                metadata = (
+                    json.loads(row["metadata"])
+                    if row["metadata"]
+                    else {}
+                )
+                visible_to = metadata.get("visible_to", "all")
+
+                if visible_to == "all" or (
+                    agent and agent.agent_name in visible_to
+                ):
+                    content = row["content"]
+                    try:
+                        content = json.loads(content)
+                    except json.JSONDecodeError:
+                        pass
+
+                    message = {
+                        "role": row["role"],
+                        "content": content,
+                        "visible_to": visible_to,
+                        "turn": metadata.get("turn"),
+                    }
+                    visible_messages.append(message)
+
+            return visible_messages
+
+    def return_messages_as_list(self) -> List[str]:
+        """Return the conversation messages as a list of formatted strings.
+
+        Returns:
+            list: List of messages formatted as 'role: content'.
+        """
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"""
+                SELECT role, content FROM {self.table_name}
+                WHERE conversation_id = ?
+                ORDER BY id ASC
+                """,
+                (self.current_conversation_id,),
+            )
+
+            return [
+                f"{row['role']}: {json.loads(row['content']) if isinstance(row['content'], str) and row['content'].startswith('{') else row['content']}"
+                for row in cursor.fetchall()
+            ]
+
+    def return_messages_as_dictionary(self) -> List[Dict]:
+        """Return the conversation messages as a list of dictionaries.
+
+        Returns:
+            list: List of dictionaries containing role and content of each message.
+        """
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"""
+                SELECT role, content FROM {self.table_name}
+                WHERE conversation_id = ?
+                ORDER BY id ASC
+                """,
+                (self.current_conversation_id,),
+            )
+
+            messages = []
+            for row in cursor.fetchall():
+                content = row["content"]
+                try:
+                    content = json.loads(content)
+                except json.JSONDecodeError:
+                    pass
+
+                messages.append(
+                    {
+                        "role": row["role"],
+                        "content": content,
+                    }
+                )
+            return messages
+
+    def add_tool_output_to_agent(self, role: str, tool_output: dict):
+        """Add a tool output to the conversation history.
+
+        Args:
+            role (str): The role of the tool.
+            tool_output (dict): The output from the tool to be added.
+        """
+        self.add(role, tool_output, message_type=MessageType.TOOL)
+
+    def get_final_message(self) -> str:
+        """Return the final message from the conversation history.
+
+        Returns:
+            str: The final message formatted as 'role: content'.
+        """
+        last_message = self.get_last_message()
+        if not last_message:
+            return ""
+        return f"{last_message['role']}: {last_message['content']}"
+
+    def get_final_message_content(self) -> Union[str, dict]:
+        """Return the content of the final message from the conversation history.
+
+        Returns:
+            Union[str, dict]: The content of the final message.
+        """
+        last_message = self.get_last_message()
+        if not last_message:
+            return ""
+        return last_message["content"]
+
+    def return_all_except_first(self) -> List[Dict]:
+        """Return all messages except the first one.
+
+        Returns:
+            list: List of messages except the first one.
+        """
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                f"""
+                SELECT role, content, timestamp, message_type, metadata, token_count
+                FROM {self.table_name}
+                WHERE conversation_id = ?
+                ORDER BY id ASC
+                LIMIT -1 OFFSET 2
+                """,
+                (self.current_conversation_id,),
+            )
+
+            messages = []
+            for row in cursor.fetchall():
+                content = row["content"]
+                try:
+                    content = json.loads(content)
+                except json.JSONDecodeError:
+                    pass
+
+                message = {
+                    "role": row["role"],
+                    "content": content,
+                }
+                if row["timestamp"]:
+                    message["timestamp"] = row["timestamp"]
+                if row["message_type"]:
+                    message["message_type"] = row["message_type"]
+                if row["metadata"]:
+                    message["metadata"] = json.loads(row["metadata"])
+                if row["token_count"]:
+                    message["token_count"] = row["token_count"]
+
+                messages.append(message)
+            return messages
+
+    def return_all_except_first_string(self) -> str:
+        """Return all messages except the first one as a string.
+
+        Returns:
+            str: All messages except the first one as a string.
+        """
+        messages = self.return_all_except_first()
+        return "\n".join(f"{msg['content']}" for msg in messages)