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

solace_agent_mesh/orchestrator/action_manager.py

@@ -0,0 +1,225 @@
+"""The Action Manager tracks all pending actions and manages their execution. It
+does the following things:
+
+1. For each LLM response that has any actions, it create a unique ActionRequestList object.
+2. The ActionRequestList object holds the list of actions to be executed.
+3. As each action response is received, the Action Manager updates the ActionRequestList object.
+4. Once all actions are received, the Action Manager sends the full response to the Orchestrator.
+5. A periodic timer checks to see if any actions should be timed out. The timer is externally
+   managed - this class just has a method to call to check for timeouts.
+
+"""
+
+from uuid import uuid4
+from datetime import datetime
+
+from solace_ai_connector.common.log import log
+from ..common.utils import format_agent_response
+
+ACTION_REQUEST_TIMEOUT = 180
+
+
+class ActionManager:
+    """This class manages all the ActionRequests that are pending"""
+
+    def __new__(cls, kv_store, lock_manager):
+        lock = lock_manager.get_lock("action_manager")
+        with lock:
+            instance = kv_store.get("action_manager_instance")
+            if instance is None:
+                instance = super(ActionManager, cls).__new__(cls)
+                kv_store.set("action_manager_instance", instance)
+        return instance
+
+    def __init__(self, kv_store, lock_manager):
+        self.action_requests = {}
+        self.lock = lock_manager.get_lock("action_manager")
+
+        with self.lock:
+            action_requests = kv_store.get("action_requests")
+            if not action_requests:
+                action_requests = {}
+                kv_store.set("action_requests", action_requests)
+        self.action_requests = action_requests
+
+    def add_action_request(self, action_requestlist, user_properties):
+        """Add an action request to the list"""
+        uuid = str(uuid4())
+
+        # Add the uuid to each action
+        for action in action_requestlist:
+            action["action_list_id"] = uuid
+        with self.lock:
+            if uuid in self.action_requests:
+                log.error("Action request with UUID %s already exists", uuid)
+
+            arl = ActionRequestList(uuid, action_requestlist, user_properties)
+            self.action_requests[uuid] = arl
+
+    def delete_action_request(self, action_list_id):
+        """Delete an action request from the list"""
+        with self.lock:
+            if action_list_id in self.action_requests:
+                del self.action_requests[action_list_id]
+
+    def get_action_info(self, action_list_id, action_name, action_idx):
+        """Get an action request"""
+        with self.lock:
+            action_list = self.action_requests.get(action_list_id)
+            if action_list is None:
+                return None
+
+            action = action_list.get_action(action_name, action_idx)
+            if action is None:
+                return None
+
+        return action
+
+    def add_action_response(self, action_response_obj, response_text_and_files):
+        """Add an action response to the list"""
+        action_list_id = action_response_obj.get("action_list_id")
+        with self.lock:
+            action_list = self.action_requests.get(action_list_id)
+            if action_list is None:
+                log.error(
+                    "Action request %s not found. Maybe it had already timed out",
+                    action_list_id,
+                )
+                return None
+
+            action_list.add_response(action_response_obj, response_text_and_files)
+
+        return action_list
+
+    def do_timeout_check(self):
+        """Check for any actions that have timed out"""
+        events = []
+        with self.lock:
+            action_list_ids = list(self.action_requests.keys())
+            for action_list_id in action_list_ids:
+                action_requestlist = self.action_requests.get(action_list_id)
+                if action_requestlist.has_timed_out():
+                    log.info("Action request %s has timed out", action_list_id)
+                    events.extend(self.do_timeout(action_requestlist, action_list_id))
+        return events
+
+    def do_timeout(self, action_requestlist, action_list_id):
+        """Handle a timeout for a specific action request"""
+        events = action_requestlist.get_action_response_timeout_events()
+        if len(events) == 0:
+            timeout_count = action_requestlist.get_timeout_count()
+            if timeout_count > 10:
+                log.error(
+                    "Action request %s has timed out too many times", action_list_id
+                )
+                del self.action_requests[action_list_id]
+            else:
+                action_requestlist.increment_timeout_count()
+        return events
+
+
+class ActionRequestList:
+    """This class holds the list of actions to be executed for a single LLM response"""
+
+    def __init__(self, action_list_id, actions, user_properties):
+        self.action_list_id = action_list_id
+        self.actions = actions
+        self.user_properties = user_properties
+        self.num_pending_actions = len(actions)
+        self.create_time = datetime.now()
+        self.timeout_count = 0
+        self.responses = {}
+
+    def has_timed_out(self):
+        """Check if the action request has timed out"""
+        return (
+            datetime.now() - self.create_time
+        ).total_seconds() > ACTION_REQUEST_TIMEOUT
+
+    def get_timeout_count(self):
+        """Get the number of times this action request has timed out"""
+        return self.timeout_count
+
+    def increment_timeout_count(self):
+        """Increment the number of times this action request has timed out"""
+        self.timeout_count += 1
+
+    def get_user_properties(self):
+        """Get the user properties"""
+        return self.user_properties
+
+    def get_action_response_timeout_events(self):
+        """Go through all the action requests and return a timeout event
+        for each one that does not yet have a response"""
+        timeout_events = []
+        for action in self.actions:
+            if "response" not in action:
+                timeout_events.append(
+                    {
+                        "message": "Action response timed out",
+                        "action_list_id": self.action_list_id,
+                        "action_idx": action.get("action_idx"),
+                        "action_name": action.get("action_name"),
+                        "user_properties": self.user_properties,
+                    }
+                )
+            action["response"] = {"text": "Action response timed out"}
+            action["timed_out"] = True
+
+        return timeout_events
+
+    def get_action(self, action_name, action_idx):
+        """Get the action with the specified name and index"""
+        if action_idx is None:
+            log.error("Action Response does not have an action_idx")
+            return None
+
+        if action_idx >= len(self.actions):
+            log.error(
+                "Action Response has an invalid action_idx. Max is %s, index is %s",
+                len(self.actions),
+                action_idx,
+            )
+            return None
+
+        action = self.actions[action_idx]
+        if action.get("action_name") != action_name:
+            log.error(
+                "Action Response has an invalid action name. Expected %s but got %s",
+                action.get("action_name"),
+                action_name,
+            )
+            return None
+
+        return action
+
+    def add_response(self, action_response_obj, response_text_and_files):
+        """Add a response to the list - returns True if all actions are complete"""
+        action_idx = action_response_obj.get("action_idx")
+        action_name = action_response_obj.get("action_name")
+
+        action = self.get_action(action_name, action_idx)
+
+        if "response" in action and not action.get("timed_out"):
+            log.error("Action Response already has a response")
+        else:
+            action["response"] = response_text_and_files
+            self.num_pending_actions -= 1
+
+        if self.is_complete():
+            log.info("Action request %s is complete", self.action_list_id)
+            return True
+
+        return False
+
+    def is_complete(self):
+        """Check if all actions have been completed"""
+        return self.num_pending_actions == 0
+
+    def get_responses(self):
+        """Get all the responses"""
+        return self.actions
+
+    def format_ai_response(self):
+        """Format the action response for the AI"""
+        return format_agent_response(self.actions)

solace_agent_mesh/orchestrator/components/orchestrator_action_manager_timeout_component.py

@@ -0,0 +1,54 @@
+"""This is a custom component that handles the action_manager timer going off"""
+
+import os
+
+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 ..action_manager import ActionManager
+
+info = {
+    "class_name": "OrchestratorActionManagerTimeoutComponent",
+    "description": ("This component handles the action_manager timer going off"),
+    "config_parameters": [],
+    "input_schema": {
+        "type": "none",
+    },
+}
+
+
+class OrchestratorActionManagerTimeoutComponent(ComponentBase):
+
+    def __init__(self, **kwargs):
+        super().__init__(info, **kwargs)
+        self.action_manager = ActionManager(self.flow_kv_store, self.flow_lock_manager)
+
+    def invoke(self, message: Message, data):
+        """Called when the timer goes off"""
+
+        # Need to go through all the active action_requests and check if any of them have timed out
+        timeout_events = self.action_manager.do_timeout_check()
+
+        # Now turn these into messages
+        messages = []
+        for event in timeout_events:
+            action_name = event.get("action_name")
+            if action_name is None:
+                log.error("Action name not found in event")
+                continue
+            action_parts = action_name.split(".")
+            if len(action_parts) == 1:
+                action_parts.append(action_parts[0])
+                action_parts[0] = "global"
+            payload = event
+            user_properties = event.get("user_properties")
+            del payload["user_properties"]
+            new_message = {
+                "topic": f"{os.getenv('SOLACE_AGENT_MESH_NAMESPACE')}solace-agent-mesh/v1/actionResponse/agent/{action_parts[0]}/{action_parts[1]}/timeout",
+                "payload": payload,
+                "user_properties": user_properties,
+            }
+            messages.append(new_message)
+
+        return messages

solace_agent_mesh/orchestrator/components/orchestrator_action_response_component.py

@@ -0,0 +1,179 @@
+"""This is the component that handles processing all action responses from agents"""
+
+import os
+from time import time
+
+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.event import Event, EventType
+from ...orchestrator.orchestrator_main import OrchestratorState
+from ...orchestrator.orchestrator_prompt import BasicRagPrompt, ContextQueryPrompt
+from ..action_manager import ActionManager
+
+info = {
+    "class_name": "OrchestratorActionResponseComponent",
+    "description": ("This component handles all action responses from agents"),
+    "config_parameters": [],
+    "input_schema": {
+        # A action response object - it doesn't have a fixed schema
+        "type": "object",
+        "additionalProperties": True,
+    },
+    # We will output a list of topics and messages to send to either the
+    # orchestrator stimulus or to the originator via slack
+    "output_schema": {
+        "type": "array",
+        "items": {
+            "type": "object",
+            "properties": {
+                "topic": {"type": "string"},
+                "payload": {"type": "object"},
+            },
+            "required": ["topic", "payload"],
+        },
+    },
+}
+
+
+class OrchestratorActionResponseComponent(ComponentBase):
+    """This is the component that handles processing all action responses from agents"""
+
+    def __init__(self, **kwargs):
+        super().__init__(info, **kwargs)
+        with self.get_lock("orchestrator_state"):
+            self.orchestrator_state = self.kv_store_get("orchestrator_state")
+            if not self.orchestrator_state:
+                self.orchestrator_state = OrchestratorState()
+                self.kv_store_set("orchestrator_state", self.orchestrator_state)
+
+        self.action_manager = ActionManager(self.flow_kv_store, self.flow_lock_manager)
+
+    def invoke(self, message: Message, data):
+        """Handle action responses from agents"""
+
+        user_properties = message.get_user_properties()
+        user_properties['timestamp_start'] = time()
+        message.set_user_properties(user_properties)
+
+        if not data:
+            log.error("No data received from agent")
+            self.discard_current_message()
+            return None
+
+        if not isinstance(data, dict):
+            log.error("Data received from agent is not a dictionary")
+            self.discard_current_message()
+            return None
+
+        session_id = data.get("session_id")
+        user_response = {}
+
+        if data.get("message"):
+            user_response["text"] = data.get("message")
+
+        user_response["files"] = []
+        if data.get("files"):
+            user_response["files"] = data.get("files")
+
+        if data.get("clear_history"):
+            user_properties = message.get_user_properties()
+            keep_depth = data.get("history_depth_to_keep")
+            if type(keep_depth) != int:
+                if keep_depth.isdigit():
+                    keep_depth = int(keep_depth)
+                else:
+                    keep_depth = 0
+            user_properties["clear_gateway_history"] = [True, keep_depth]
+            message.set_user_properties(user_properties)
+
+        if data.get("error_info"):
+            if self.error_queue is not None:
+                error_info = data.get("error_info", {})
+                action_name = data.get("action_name")
+                action_list_id = data.get("action_list_id")
+                action_idx = data.get("action_idx")
+                action = (
+                    self.action_manager.get_action_info(
+                        action_list_id, action_name, action_idx
+                    )
+                    or {}
+                )
+                agent_name = action.get("agent_name", "Unknown")
+
+                source = ""
+                if agent_name == "global" and action_name == "error_action":
+                    source = "Error raised by the Orchestrator"
+                else:
+                    source = f"Agent: {agent_name}, Action: {action_name}, Params: {action.get('action_params', {})}"
+
+                self.error_queue.put(
+                    Event(
+                        EventType.MESSAGE,
+                        Message(
+                            payload={
+                                "error_message": error_info.get("error_message"),
+                                "source": source,
+                            },
+                            user_properties=message.get_user_properties()
+                            if message
+                            else {},
+                        ),
+                    )
+                )
+
+        if data.get("agent_state_change"):
+            agent_state_change = data.get("agent_state_change")
+            agent_name = agent_state_change.get("agent_name")
+            state = agent_state_change.get("new_state")
+            self.orchestrator_state.update_agent_state(agent_name, state, session_id)
+            if state == "open":
+                user_response["text"] = (
+                    f"Opened {agent_name}. Now you can interact with it."
+                )
+
+        if data.get("context_query"):
+            cq = data.get("context_query")
+            query = cq.get("query")
+            context = cq.get("context")
+            context_type = cq.get("context_type")
+            if context_type == "raw":
+                user_response["text"] = query
+            else:
+                user_response["text"] = ContextQueryPrompt(query=query, context=context)
+
+        events = []
+
+        # Tell the ActionManager about the result - if it is complete, then
+        # it will send the result back to the model
+        action_list = self.action_manager.add_action_response(data, user_response)
+        if action_list and action_list.is_complete():
+            response_text, files = action_list.format_ai_response()
+
+            if response_text:
+                # Send the result back to the model
+                user_properties = message.get_user_properties()
+                events.append(
+                    {
+                        "topic": f"{os.getenv('SOLACE_AGENT_MESH_NAMESPACE')}solace-agent-mesh/v1/stimulus/orchestrator/reinvokeModel",
+                        "payload": {
+                            "text": response_text,
+                            "files": files,
+                            "identity": user_properties.get("identity"),
+                            "channel": user_properties.get("channel"),
+                            "thread_ts": user_properties.get("thread_ts"),
+                            "action_response_reinvoke": True,
+                        },
+                    }
+                )
+            self.action_manager.delete_action_request(action_list.action_list_id)
+
+        if len(events) == 0:
+            self.discard_current_message()
+            return None
+
+        user_properties = message.get_user_properties()
+        user_properties['timestamp_end'] = time()
+        message.set_user_properties(user_properties)
+
+        return events

solace_agent_mesh/orchestrator/components/orchestrator_register_component.py

@@ -0,0 +1,107 @@
+"""This is a custom component that handles registrations from the distributed agents"""
+
+# from solace_ai_connector.common.log import log
+from solace_ai_connector.components.component_base import ComponentBase
+from solace_ai_connector.common.message import Message
+from ...orchestrator.orchestrator_main import OrchestratorState
+
+info = {
+    "class_name": "OrchestratorRegister",
+    "description": ("This component handles registrations from the distributed agents"),
+    "config_parameters": [
+        {
+            "name": "agent_ttl_ms",
+            "required": False,
+            "description": "The time-to-live for agent registrations in milliseconds. There must be a registration from an agent within this time period, otherwise the agent will be considered offline.",
+            "default": 60000,
+        },
+    ],
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            "agent_name": {
+                "type": "string",
+                "description": "The name of the agent.",
+            },
+            "description": {
+                "type": "string",
+                "description": "A description of the application.",
+            },
+            "actions": {
+                "type": "array",
+                "description": "A list of actions the application can perform.",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "name": {
+                            "type": "string",
+                            "description": "The name of the action.",
+                        },
+                        "description": {
+                            "type": "string",
+                            "description": "A description of what the action does.",
+                        },
+                        "params": {
+                            "type": "array",
+                            "description": "A list of parameters for the action.",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "name": {
+                                        "type": "string",
+                                        "description": "The name of the parameter.",
+                                    },
+                                    "description": {
+                                        "type": "string",
+                                        "description": "A description of the parameter.",
+                                    },
+                                    "type": {
+                                        "type": "string",
+                                        "description": "The type of the parameter.",
+                                    },
+                                },
+                                "required": ["name", "description", "type"],
+                            },
+                        },
+                        "examples": {
+                            "type": "array",
+                            "description": "A list of examples for the action.",
+                            "items": {
+                                "type": "string",
+                                "description": "An example for the action.",
+                            },
+                        },
+                        "required_scopes": {
+                            "type": "array",
+                            "description": "A list of required scopes for the action.",
+                            "items": {
+                                "type": "string",
+                                "description": "A required scope for the action.",
+                            },
+                        },
+                    },
+                    "required": ["name", "description", "params"],
+                },
+            },
+        },
+        "required": ["agent_name", "description", "actions"],
+    },
+}
+
+
+class OrchestratorRegister(ComponentBase):
+    def __init__(self, **kwargs):
+        super().__init__(info, **kwargs)
+        self.agent_ttl_ms = self.get_config("agent_ttl_ms")
+        OrchestratorState.set_config({"agent_ttl_ms": self.agent_ttl_ms})
+        with self.get_lock("orchestrator_state"):
+            self.orchestrator_state = self.kv_store_get("orchestrator_state")
+            if not self.orchestrator_state:
+                self.orchestrator_state = OrchestratorState()
+                self.kv_store_set("orchestrator_state", self.orchestrator_state)
+
+    def invoke(self, message: Message, data):
+        """Receive a registration from an agent and store it in the component's state."""
+        # log.info("Received registration from agent")
+        self.orchestrator_state.register_agent(data)
+        return data