solace-agent-mesh 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

solace_agent_mesh/services/history_service/history_providers/redis_history_provider.py

@@ -0,0 +1,163 @@
+import json
+import time
+from .base_history_provider import BaseHistoryProvider
+
+class RedisHistoryProvider(BaseHistoryProvider):
+    def __init__(self, config=None):
+        super().__init__(config)
+        try:
+            import redis
+        except ImportError:
+            raise ImportError("Please install the redis package to use the RedisHistoryProvider.\n\t$ pip install redis")
+        
+        self.redis_client = redis.Redis(
+            host=self.config.get("redis_host", "localhost"),
+            port=self.config.get("redis_port", 6379),
+            db=self.config.get("redis_db", 0),
+        )
+
+    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)
+
+            # 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 get_session_meta(self, session_id: str):
+        """
+        Retrieve the session metadata.
+
+        :param session_id: The session identifier.
+        :return: The session metadata.
+        """
+        # 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,
+        }
+
+
+    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]

solace_agent_mesh/services/history_service/history_service.py

@@ -0,0 +1,139 @@
+import time
+import importlib
+from typing import Union
+
+from solace_ai_connector.common.log import log
+
+from ...common.time import ONE_HOUR, FIVE_MINUTES
+from ..common import AutoExpiry, AutoExpirySingletonMeta
+from .history_providers.memory_history_provider import MemoryHistoryProvider
+from .history_providers.redis_history_provider import RedisHistoryProvider
+from .history_providers.base_history_provider import BaseHistoryProvider
+
+HISTORY_PROVIDERS = {
+    "redis": RedisHistoryProvider,
+    "memory": MemoryHistoryProvider,
+}
+
+
+DEFAULT_PROVIDER = "memory"
+
+DEFAULT_MAX_TURNS = 40
+DEFAULT_MAX_CHARACTERS = 50_000
+
+DEFAULT_HISTORY_POLICY = {
+    "max_turns": DEFAULT_MAX_TURNS,
+    "max_characters": DEFAULT_MAX_CHARACTERS,
+    "enforce_alternate_message_roles": False,
+}
+
+
+# HistoryService class - Manages history storage and retrieval
+class HistoryService(AutoExpiry, metaclass=AutoExpirySingletonMeta):
+    history_provider: BaseHistoryProvider
+
+    def __init__(self, config={}, identifier=None):
+        self.identifier = identifier
+        self.config = config
+        self.provider_type = self.config.get("type", DEFAULT_PROVIDER)
+        self.time_to_live = self.config.get("time_to_live", ONE_HOUR)
+        self.expiration_check_interval = self.config.get(
+            "expiration_check_interval", FIVE_MINUTES
+        )
+
+        if self.provider_type not in HISTORY_PROVIDERS and not self.config.get(
+            "module_path"
+        ):
+            raise ValueError(
+                f"Unsupported history provider type: {self.provider_type}. No module_path provided."
+            )
+
+        history_policy = {
+            **DEFAULT_HISTORY_POLICY,
+            **self.config.get("history_policy", {}),
+        }
+        if self.provider_type in HISTORY_PROVIDERS:
+            # Load built-in history provider
+            self.history_provider = HISTORY_PROVIDERS[self.provider_type](
+                history_policy
+            )
+        else:
+            try:
+                # Load the provider from the module path
+                module_name = self.provider_type
+                module_path = self.config.get("module_path")
+                module = importlib.import_module(module_path, package=__package__)
+                history_class = getattr(module, module_name)
+                if not issubclass(history_class, BaseHistoryProvider):
+                    raise ValueError(
+                        f"History provider class {history_class} does not inherit from BaseHistoryProvider"
+                    )
+                self.history_provider = history_class(history_policy)
+            except Exception as e:
+                raise ImportError("Unable to load component: " + str(e)) from e
+
+        # Start the background thread for auto-expiry
+        self._start_auto_expiry_thread(self.expiration_check_interval)
+
+    def _delete_expired_items(self):
+        """Checks all history entries and deletes those that have exceeded max_time_to_live."""
+        current_time = time.time()
+        sessions = self.history_provider.get_all_sessions()
+        for session_id in sessions:
+            session = self.history_provider.get_session_meta(session_id)
+            if not session:
+                continue
+            elapsed_time = current_time - session["last_active_time"]
+            if elapsed_time > self.time_to_live:
+                self.history_provider.clear_history(session_id)
+                log.debug(f"History for session {session_id} has expired")
+
+    def store_history(self, session_id: str, role: str, content: Union[str, dict]):
+        """
+        Store a new entry in the history.
+
+        :param session_id: The session identifier.
+        :param role: The role of the entry to be stored in the history.
+        :param content: The content of the entry to be stored in the history.
+        """
+        if not content:
+            return
+        return self.history_provider.store_history(session_id, role, content)
+
+    def get_history(self, session_id:str) -> list:
+        """
+        Retrieve the entire history.
+
+        :param session_id: The session identifier.
+        :return: The complete history.
+        """
+        return self.history_provider.get_history(session_id)
+
+    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 to be stored in the history.
+        """
+        if not file:
+            return
+        return self.history_provider.store_file(session_id, file)
+
+    def get_files(self, session_id:str) -> list:
+        """
+        Retrieve the files for a session.
+
+        :param session_id: The session identifier.
+        :return: The files for the session.
+        """
+        return self.history_provider.get_files(session_id)
+
+    def clear_history(self, session_id:str, keep_levels=0):
+        """
+        Clear the history and files, 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).
+        """
+        return self.history_provider.clear_history(session_id, keep_levels)

solace_agent_mesh/services/llm_service/components/llm_request_component.py

@@ -0,0 +1,293 @@
+"""LLM Request Component for performing LLM service requests."""
+
+import uuid
+from typing import Dict, Any
+
+from solace_ai_connector.components.component_base import ComponentBase
+from solace_ai_connector.common.log import log
+from solace_ai_connector.common.message import Message
+from solace_ai_connector.common.utils import ensure_slash_on_end
+
+info = {
+    "class_name": "LLMRequestComponent",
+    "description": "Component that performs LLM service requests",
+    "config_parameters": [
+        {
+            "name": "llm_service_topic",
+            "required": True,
+            "description": "The topic for the LLM service",
+        },
+        {
+            "name": "stream_to_flow",
+            "required": False,
+            "description": (
+                "Name the flow to stream the output to - this must be configured for "
+                "llm_mode='stream'. This is mutually exclusive with stream_to_next_component."
+            ),
+            "default": "",
+        },
+        {
+            "name": "stream_to_next_component",
+            "required": False,
+            "description": (
+                "Whether to stream the output to the next component in the flow. "
+                "This is mutually exclusive with stream_to_flow."
+            ),
+            "default": False,
+        },
+        {
+            "name": "llm_mode",
+            "required": False,
+            "description": (
+                "The mode for streaming results: 'sync' or 'stream'. 'stream' "
+                "will just stream the results to the named flow. 'none' will "
+                "wait for the full response."
+            ),
+            "default": "none",
+        },
+        {
+            "name": "stream_batch_size",
+            "required": False,
+            "description": "The minimum number of words in a single streaming result.",
+            "default": 15,
+        },
+    ],
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            "messages": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "role": {
+                            "type": "string",
+                            "enum": ["system", "user", "assistant"],
+                        },
+                        "content": {"type": "string"},
+                    },
+                    "required": ["role", "content"],
+                },
+            },
+            "source_info": {
+                "type": "object",
+                "properties": {
+                    "type": {"type": "string"},
+                },
+                "required": ["type"],
+            },
+        },
+        "required": ["messages"],
+    },
+    "output_schema": {
+        "type": "object",
+        "properties": {
+            "content": {
+                "type": "string",
+                "description": "The generated response from the model",
+            },
+            "chunk": {
+                "type": "string",
+                "description": "The current chunk of the response",
+            },
+            "response_uuid": {
+                "type": "string",
+                "description": "The UUID of the response",
+            },
+            "first_chunk": {
+                "type": "boolean",
+                "description": "Whether this is the first chunk of the response",
+            },
+            "last_chunk": {
+                "type": "boolean",
+                "description": "Whether this is the last chunk of the response",
+            },
+            "streaming": {
+                "type": "boolean",
+                "description": "Whether this is a streaming response",
+            },
+        },
+        "required": ["content"],
+    },
+}
+
+
+class LLMRequestComponent(ComponentBase):
+    """Component that performs LLM service requests."""
+
+    def __init__(self, child_info=None, **kwargs):
+        super().__init__(child_info or info, **kwargs)
+        self.init()
+
+    def init(self):
+        """Initialize the component with configuration parameters."""
+        self.llm_service_topic = ensure_slash_on_end(
+            self.get_config("llm_service_topic")
+        )
+        self.stream_to_flow = self.get_config("stream_to_flow")
+        self.stream_to_next_component = self.get_config("stream_to_next_component")
+        self.llm_mode = self.get_config("llm_mode")
+        self.stream_batch_size = self.get_config("stream_batch_size")
+
+        if self.stream_to_flow and self.stream_to_next_component:
+            raise ValueError(
+                "stream_to_flow and stream_to_next_component are mutually exclusive"
+            )
+
+        if not self.is_broker_request_response_enabled():
+            raise ValueError(
+                "LLM service topic is set, but the component does not "
+                f"have its broker request/response enabled, {self.__class__.__name__}"
+            )
+
+    def invoke(self, message: Message, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Invoke the LLM service request.
+
+        Args:
+            message (Message): The input message.
+            data (Dict[str, Any]): The input data containing the messages.
+
+        Returns:
+            Dict[str, Any]: The response from the LLM service.
+        """
+        messages = data.get("messages", [])
+        source_info = data.get("source_info", {})
+        llm_message = self._create_llm_message(message, messages, source_info)
+        response_uuid = str(uuid.uuid4())
+
+        try:
+            if self.llm_mode == "stream":
+                return self._handle_streaming(message, llm_message, response_uuid)
+            else:
+                return self._handle_sync(llm_message)
+        except Exception as e:
+            log.error("Error invoking LLM service: %s", e, exc_info=True)
+            raise
+
+    def _handle_sync(self, llm_message: Message) -> Dict[str, Any]:
+        """
+        Handle synchronous LLM service request.
+
+        Args:
+            llm_message (Message): The message to send to the LLM service.
+
+        Returns:
+            Dict[str, Any]: The response from the LLM service.
+        """
+        response = self.do_broker_request_response(llm_message)
+        return response.get_payload()
+
+    def _handle_streaming(
+        self, input_message: Message, llm_message: Message, response_uuid: str
+    ) -> Dict[str, Any]:
+        """
+        Handle streaming LLM service request.
+
+        Args:
+            input_message (Message): The original input message.
+            llm_message (Message): The message to send to the LLM service.
+            response_uuid (str): The UUID for the response.
+
+        Returns:
+            Dict[str, Any]: The final response from the LLM service.
+        """
+        aggregate_result = ""
+        current_batch = ""
+        first_chunk = True
+
+        for response_message, last_message in self.do_broker_request_response(
+            llm_message,
+            stream=True,
+            streaming_complete_expression="input.payload:last_chunk",
+        ):
+            payload = response_message.get_payload()
+            content = payload.get("chunk", "")
+            aggregate_result += content
+            current_batch += content
+
+            if len(current_batch.split()) >= self.stream_batch_size or last_message:
+                self._send_streaming_chunk(
+                    input_message,
+                    current_batch,
+                    aggregate_result,
+                    response_uuid,
+                    first_chunk,
+                    last_message,
+                )
+                current_batch = ""
+                first_chunk = False
+
+            if last_message:
+                break
+
+        return {
+            "content": aggregate_result,
+            "response_uuid": response_uuid,
+            "streaming": True,
+            "last_chunk": True,
+        }
+
+    def _create_llm_message(self, message: Message, messages: list, source_info: dict) -> Message:
+        """
+        Create a message for the LLM service request.
+
+        Args:
+            message (Message): The original input message.
+            messages (list): The list of messages to send to the LLM service.
+            source_info (dict): Information about the caller to help montoring LLM requests.
+
+        Returns:
+            Message: The created message for the LLM service.
+        """
+        user_properties = message.get_user_properties().copy()
+        stimulus_uuid = user_properties.get("stimulus_uuid", str(uuid.uuid4()))
+        session_id = user_properties.get("session_id", "x")
+        originator_id = user_properties.get("originator_id", "x")
+        user_properties["llm_request_source_info"] = source_info
+
+        topic = f"{self.llm_service_topic}{stimulus_uuid}/{session_id}/{originator_id}"
+
+        return Message(
+            topic=topic,
+            payload={"messages": messages, "stream": True},
+            user_properties=user_properties,
+        )
+
+    def _send_streaming_chunk(
+        self,
+        input_message: Message,
+        chunk: str,
+        aggregate_result: str,
+        response_uuid: str,
+        first_chunk: bool,
+        last_chunk: bool,
+    ):
+        """
+        Send a streaming chunk to the specified flow or next component.
+
+        Args:
+            input_message (Message): The original input message.
+            chunk (str): The current chunk of the response.
+            aggregate_result (str): The aggregated result so far.
+            response_uuid (str): The UUID of the response.
+            first_chunk (bool): Whether this is the first chunk.
+            last_chunk (bool): Whether this is the last chunk.
+        """
+        payload = {
+            "chunk": chunk,
+            "content": aggregate_result,
+            "response_uuid": response_uuid,
+            "first_chunk": first_chunk,
+            "last_chunk": last_chunk,
+            "streaming": True,
+        }
+        message = Message(
+            payload=payload,
+            user_properties=input_message.get_user_properties(),
+        )
+
+        if self.stream_to_flow:
+            self.send_to_flow(self.stream_to_flow, message)
+        elif self.stream_to_next_component:
+            self.send_message(message)