solace-agent-mesh 0.1.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

solace_agent_mesh/services/history_service/history_providers/memory_history_provider.py

@@ -1,167 +1,33 @@
-import time
+"""
+Memory history provider
+"""
+
 from .base_history_provider import BaseHistoryProvider
 
 
 class MemoryHistoryProvider(BaseHistoryProvider):
+    """
+    A history provider that stores history in memory.
+    """
+
     def __init__(self, config=None):
         super().__init__(config)
         self.history = {}
 
-    def store_history(self, session_id: str, role: str, content: str | dict):
-        """
-        Store a new entry in the history.
-
-        :param session_id: The session identifier.
-        :param history_entry: The entry to be stored in the history.
-        """
-
+    def store_session(self, session_id, data):
         if session_id not in self.history:
-            self.history[session_id] = {
-                "history": [],
-                "files": [],
-                "last_active_time": time.time(),
-                "num_characters": 0,
-                "num_turns": 0,
-            }
-
-        # Check if adding another entry would exceed the max_turns
-        if self.history[session_id]["num_turns"] == self.max_turns:
-            # Remove the oldest entry
-            oldest_entry = self.history[session_id]["history"].pop(0)
-            # Subtract the length of the oldest entry from the total length
-            self.history[session_id]["num_characters"] -= len(
-                str(oldest_entry["content"])
-            )
-            self.history[session_id]["num_turns"] -= 1
-
-        if (
-            self.enforce_alternate_message_roles
-            and self.history[session_id]["num_turns"] > 0
-            # Check if the last entry was by the same role
-            and self.history[session_id]["history"]
-            and self.history[session_id]["history"][-1]["role"] == role
-        ):
-            # Append to last entry
-            self.history[session_id]["history"][-1]["content"] += content
-        else:
-            # Add the new entry
-            self.history[session_id]["history"].append(
-                {"role": role, "content": content}
-            )
-            # Update the number of turns
-            self.history[session_id]["num_turns"] += 1
-
-        # Update the length
-        self.history[session_id]["num_characters"] += len(str(content))
-        # Update the last active time
-        self.history[session_id]["last_active_time"] = time.time()
+            self.history[session_id] = {}
+        
+        self.history[session_id].update(data)
 
-        # Check if we have exceeded max_characters
-        if self.max_characters:
-            while (
-                self.history[session_id]["num_characters"] > self.max_characters
-                and self.history[session_id]["num_turns"] > 0
-            ):
-                # Remove the oldest entry
-                oldest_entry = self.history[session_id]["history"].pop(0)
-                # Subtract the length of the oldest entry from the total length
-                self.history[session_id]["num_characters"] -= len(
-                    str(oldest_entry["content"])
-                )
-                self.history[session_id]["num_turns"] -= 1
-
-    def get_history(self, session_id: str):
-        """
-        Retrieve the entire history for a session.
-
-        :param session_id: The session identifier.
-        :return: The complete history for the session.
-        """
-        if session_id not in self.history:
-            return []
-        return self.history.get(session_id)["history"]
-
-    def store_file(self, session_id: str, file: dict):
-        """
-        Store a file in the history.
-
-        :param session_id: The session identifier.
-        :param file: The file metadata to be stored in the history.
-        """
+    def get_session(self, session_id):
         if session_id not in self.history:
-            self.history[session_id] = {
-                "history": [],
-                "files": [],
-                "last_active_time": time.time(),
-                "num_characters": 0,
-                "num_turns": 0,
-            }
-        self.history[session_id]["last_active_time"] = time.time()
-
-        # Check duplicate
-        for f in self.history[session_id]["files"]:
-            if f.get("url") and f.get("url") == file.get("url"):
-                return
-
-        self.history[session_id]["files"].append(file)
-
-    def get_files(self, session_id: str):
-        """
-        Retrieve the files for a session.
-
-        :param session_id: The session identifier.
-        :return: The files for the session.
-        """
-        if session_id not in self.history:
-            return []
-        files = []
-        current_time = time.time()
-        all_files = self.history.get(session_id)["files"].copy()
-        for file in all_files:
-            expiration_timestamp = file.get("expiration_timestamp")
-            if expiration_timestamp and current_time > expiration_timestamp:
-                self.history[session_id]["files"].remove(file)
-                continue
-            files.append(file)
-        return files
-
-    def clear_history(self, session_id: str, keep_levels=0):
-        """
-        Clear the history for a session, optionally keeping a specified number of recent entries.
-
-        :param session_id: The session identifier.
-        :param keep_levels: Number of most recent history entries to keep. Default is 0 (clear all).
-        """
-        if session_id in self.history:
-            if keep_levels <= 0:
-                del self.history[session_id]
-            else:
-                self.history[session_id]["history"] = self.history[session_id][
-                    "history"
-                ][-keep_levels:]
-                # Recalculate the length and num_turns
-                self.history[session_id]["num_characters"] = sum(
-                    len(str(entry)) for entry in self.history[session_id]["history"]
-                )
-                self.history[session_id]["num_turns"] = len(
-                    self.history[session_id]["history"]
-                )
-
-    def get_session_meta(self, session_id: str):
-        """
-        Retrieve the session metadata.
-
-        :param session_id: The session identifier.
-        :return: The session metadata.
-        """
-        if session_id in self.history:
-            session = self.history[session_id]
-            return {
-                "num_characters": session["num_characters"],
-                "num_turns": session["num_turns"],
-                "last_active_time": session["last_active_time"],
-            }
-        return None
+            return {}
+        return self.history[session_id]
 
     def get_all_sessions(self) -> list[str]:
         return list(self.history.keys())
+
+    def delete_session(self, session_id):
+        if session_id in self.history:
+            del self.history[session_id]

solace_agent_mesh/services/history_service/history_providers/mongodb_history_provider.py

@@ -0,0 +1,66 @@
+"""
+MongoDB-based history provider for storing session data.
+"""
+from .base_history_provider import BaseHistoryProvider
+
+class MongoDBHistoryProvider(BaseHistoryProvider):
+    """
+    A MongoDB-based history provider for storing session data.
+    """
+    def __init__(self, config=None):
+        super().__init__(config)
+
+        try:
+            from pymongo import MongoClient
+        except ImportError:
+            raise ImportError("Please install the pymongo package to use the MongoDBHistoryProvider.\n\t$ pip install pymongo")
+        
+        if not self.config.get("mongodb_uri"):
+            raise ValueError("Missing required configuration for MongoDBHistoryProvider, Missing 'mongodb_uri' in 'store_config'.")
+        
+        
+        self.client = MongoClient(self.config.get("mongodb_uri"))
+        self.db = self.client[self.config.get("mongodb_db", "history_db")]
+        self.collection = self.db[self.config.get("mongodb_collection", "sessions")]
+    
+    def _get_key(self, session_id):
+        """
+        Generate a document identifier for a session.
+
+        :param session_id: The session identifier.
+        :return: The session ID as the primary key.
+        """
+        return {"_id": session_id}
+    
+    def store_session(self, session_id: str, data: dict):
+        """
+        Store the session metadata.
+
+        :param session_id: The session identifier.
+        :param data: The session data to be stored.
+        """
+        self.collection.update_one(self._get_key(session_id), {"$set": {"data": data}}, upsert=True)
+    
+    def get_session(self, session_id: str)->dict:
+        """
+        Retrieve the session.
+
+        :param session_id: The session identifier.
+        :return: The session metadata as a dictionary.
+        """
+        document = self.collection.find_one(self._get_key(session_id))
+        return document.get("data") if document else {}
+    
+    def get_all_sessions(self) -> list[str]:
+        """
+        Retrieve all session identifiers.
+        """
+        return [doc["_id"] for doc in self.collection.find({}, {"_id": 1})]
+    
+    def delete_session(self, session_id: str):
+        """
+        Delete the session.
+
+        :param session_id: The session identifier.
+        """
+        self.collection.delete_one(self._get_key(session_id))

solace_agent_mesh/services/history_service/history_providers/redis_history_provider.py

@@ -1,8 +1,13 @@
+"""
+A history provider that stores history in Redis.
+"""
 import json
-import time
 from .base_history_provider import BaseHistoryProvider
 
 class RedisHistoryProvider(BaseHistoryProvider):
+    """
+    A history provider that stores history in Redis.
+    """
     def __init__(self, config=None):
         super().__init__(config)
         try:
@@ -14,150 +19,48 @@ class RedisHistoryProvider(BaseHistoryProvider):
             host=self.config.get("redis_host", "localhost"),
             port=self.config.get("redis_port", 6379),
             db=self.config.get("redis_db", 0),
+            decode_responses=True  # Ensures string output
         )
+    
+    def _get_key(self, session_id):
+        """
+        Generate a Redis key with a specific prefix for a session.
 
-    def _get_history_key(self, session_id: str):
-        return f"session:{session_id}:history"
-
-    def _get_files_key(self, session_id: str):
-        return f"session:{session_id}:files"
-
-    def store_history(self, session_id: str, role: str, content: str | dict):
-        key = self._get_history_key(session_id)
-        entry = {"role": role, "content": content}
-        entry_json = json.dumps(entry)
-
-        # Check if session exists, if not initialize it
-        if not self.redis_client.exists(key):
-            self.redis_client.hset(session_id, mapping={
-                "num_characters": 0,
-                "num_turns": 0,
-                "last_active_time": time.time()
-            })
-
-        # Get current stats
-        session_meta = self.redis_client.hgetall(session_id)
-        num_characters = int(session_meta.get(b"num_characters", 0))
-        num_turns = int(session_meta.get(b"num_turns", 0))
-
-        # Add the new entry
-        if self.enforce_alternate_message_roles and num_turns > 0:
-            last_entry = json.loads(self.redis_client.lindex(key, -1))
-            if last_entry["role"] == role:
-                last_entry["content"] += content
-                self.redis_client.lset(key, -1, json.dumps(last_entry))
-            else:
-                self.redis_client.rpush(key, entry_json)
-                num_turns += 1
-        else:
-            self.redis_client.rpush(key, entry_json)
-            num_turns += 1
-        num_characters += len(str(content))
-
-        # Enforce max_turns by trimming the oldest entry if needed
-        if self.max_turns and num_turns > self.max_turns:
-            oldest_entry = json.loads(self.redis_client.lpop(key))
-            num_characters -= len(str(oldest_entry["content"]))
-            num_turns -= 1
-
-        # Enforce max_characters
-        if self.max_characters:
-            while num_characters > self.max_characters and num_turns > 0:
-                oldest_entry = json.loads(self.redis_client.lpop(key))
-                num_characters -= len(str(oldest_entry["content"]))
-                num_turns -= 1
-
-        # Update metadata and set expiration
-        self.redis_client.hset(session_id, mapping={
-            "num_characters": num_characters,
-            "num_turns": num_turns,
-            "last_active_time": time.time()
-        })
-
-    def get_history(self, session_id: str):
-        key = self._get_history_key(session_id)
-        history = self.redis_client.lrange(key, 0, -1)
-
-        # Decode JSON entries and return a list of dictionaries
-        return [json.loads(entry) for entry in history]
-
-    def store_file(self, session_id: str, file: dict):
-        key = self._get_files_key(session_id)
-        file_entry = json.dumps(file)
-
-        # Avoid duplicate files by checking existing URLs
-        existing_files = self.get_files(session_id)
-        if any(f.get("url") == file.get("url") for f in existing_files):
-            return
-
-        # Add the file and update metadata
-        self.redis_client.rpush(key, file_entry)
-        self.redis_client.hset(session_id, "last_active_time", time.time())
-
-    def get_files(self, session_id: str):
-        key = self._get_files_key(session_id)
-        current_time = time.time()
-        files = self.redis_client.lrange(key, 0, -1)
-
-        valid_files = []
-        for file_json in files:
-            file = json.loads(file_json)
-            expiration_timestamp = file.get("expiration_timestamp")
-
-            # Remove expired files
-            if expiration_timestamp and current_time > expiration_timestamp:
-                self.redis_client.lrem(key, 0, file_json)
-            else:
-                valid_files.append(file)
-
-        return valid_files
-
-    def clear_history(self, session_id: str, keep_levels=0):
-        history_key = self._get_history_key(session_id)
-        files_key = self._get_files_key(session_id)
-
-        if keep_levels > 0:
-            # Keep the latest `keep_levels` entries
-            self.redis_client.ltrim(history_key, -keep_levels, -1)
-
-            # Recalculate session metadata
-            remaining_entries = self.redis_client.lrange(history_key, 0, -1)
-            num_characters = sum(len(str(json.loads(entry)["content"])) for entry in remaining_entries)
-            num_turns = len(remaining_entries)
+        :param session_id: The session identifier.
+        :return: A formatted Redis key string.
+        """
+        return f"sessions:{session_id}:history"
 
-            # Update metadata
-            self.redis_client.hset(session_id, mapping={
-                "num_characters": num_characters,
-                "num_turns": num_turns
-            })
-        else:
-            # Clear all history and files
-            self.redis_client.delete(history_key, files_key, session_id)
+    def store_session(self, session_id: str, data: dict):
+        """
+        Store the session metadata.
 
+        :param session_id: The session identifier.
+        :param data: The session data to be stored.
+        """
+        self.redis_client.set(self._get_key(session_id), json.dumps(data))
 
-    def get_session_meta(self, session_id: str):
+    def get_session(self, session_id: str)->dict:
         """
-        Retrieve the session metadata.
+        Retrieve the session.
 
         :param session_id: The session identifier.
-        :return: The session metadata.
+        :return: The session metadata as a dictionary.
         """
-        # Check if session exists
-        if not self.redis_client.exists(session_id):
-            return None
-        # Get current stats
-        session_meta = self.redis_client.hgetall(session_id)
-        num_characters = int(session_meta.get(b"num_characters", 0))
-        num_turns = int(session_meta.get(b"num_turns", 0))
-        last_active_time = float(session_meta.get(b"last_active_time", 0))
-        return {
-            "num_characters": num_characters,
-            "num_turns": num_turns,
-            "last_active_time": last_active_time,
-        }
+        data = self.redis_client.get(self._get_key(session_id))
+        return json.loads(data) if data else {}
 
+    def get_all_sessions(self) -> list[str]:
+        """
+        Retrieve all session identifiers.
+        """
+        keys = self.redis_client.keys("sessions:*:history")
+        return [key.split(":")[1] for key in keys]
+    
+    def delete_session(self, session_id: str):
+        """
+        Delete the session.
 
-    def get_all_sessions(self)-> list[str]:
-        # List all sessions based on Redis keys
-        session_keys = self.redis_client.scan_iter("session:*:history")
-        return [key.decode().split(":")[1] for key in session_keys]
+        :param session_id: The session identifier.
+        """
+        self.redis_client.delete(self._get_key(session_id))

solace_agent_mesh/services/history_service/history_providers/sql_history_provider.py

@@ -0,0 +1,93 @@
+import json
+
+from .base_history_provider import BaseHistoryProvider
+from ....common.postgres_database import PostgreSQLDatabase
+from ....common.mysql_database import MySQLDatabase
+
+
+class DatabaseFactory:
+    """
+    Factory class to create database instances.
+    """
+    DATABASE_PROVIDERS = {
+        "postgres": PostgreSQLDatabase,
+        "mysql": MySQLDatabase,
+    }
+
+    @staticmethod
+    def get_database(db_type, **kwargs):
+        if db_type not in DatabaseFactory.DATABASE_PROVIDERS:
+            raise ValueError(f"Unsupported database type: {db_type}")
+        return DatabaseFactory.DATABASE_PROVIDERS[db_type](**kwargs)
+
+class SQLHistoryProvider(BaseHistoryProvider):
+    """
+    A history provider that stores session history in a SQL database.
+    """
+    def __init__(self, config=None):
+        super().__init__(config)
+        self.db_type = self.config.get("db_type", "postgres")
+        self.table_name = self.config.get("table_name", "session_history")
+        self.db = DatabaseFactory.get_database(
+            self.db_type,
+            host=self.config.get("sql_host"),
+            user=self.config.get("sql_user"),
+            password=self.config.get("sql_password"),
+            database=self.config.get("sql_database"),
+        )
+        self._ensure_table_exists()
+    
+    def _ensure_table_exists(self):
+        """
+        Ensures the required table exists in the database.
+        """
+        query = f"""
+        CREATE TABLE IF NOT EXISTS {self.table_name} (
+            session_id TEXT PRIMARY KEY,
+            data JSON
+        )
+        """
+        self.db.execute(query)
+    
+    def store_session(self, session_id: str, data: dict):
+        """
+        Store or update session metadata.
+        """
+        query = f"""
+        INSERT INTO {self.table_name} (session_id, data) 
+        VALUES (%s, %s) 
+        ON CONFLICT (session_id) DO UPDATE 
+        SET data = EXCLUDED.data
+        """ if self.db_type == "postgres" else f"""
+        INSERT INTO {self.table_name} (session_id, data) 
+        VALUES (%s, %s) 
+        ON DUPLICATE KEY UPDATE data = VALUES(data)
+        """
+        self.db.execute(query, (session_id, json.dumps(data)))
+    
+    def get_session(self, session_id: str) -> dict:
+        """
+        Retrieve a session by ID.
+        """
+        query = f"SELECT data FROM {self.table_name} WHERE session_id = %s"
+        cursor = self.db.execute(query, (session_id,))
+        row = cursor.fetchone()
+        if not row.get("data"):
+            return {}
+        data = row["data"] if isinstance(row["data"], dict) else json.loads(row["data"])
+        return data
+    
+    def get_all_sessions(self) -> list[str]:
+        """
+        Retrieve all session identifiers.
+        """
+        query = f"SELECT session_id FROM {self.table_name}"
+        cursor = self.db.execute(query)
+        return [row["session_id"] for row in cursor.fetchall()]
+    
+    def delete_session(self, session_id: str):
+        """
+        Delete a session by ID, ensuring only one row is deleted.
+        """
+        query = f"DELETE FROM {self.table_name} WHERE session_id = %s LIMIT 1"
+        self.db.execute(query, (session_id,))