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

solace_agent_mesh/gateway/components/gateway_base.py

@@ -3,15 +3,17 @@ from solace_ai_connector.common.log import log
 from ...services.history_service import HistoryService
 from ..identity.identity_provider import IdentityProvider
 from ...common.constants import DEFAULT_IDENTITY_KEY_FIELD
+from ...orchestrator.orchestrator_prompt import LONG_TERM_MEMORY_PROMPT
 
 
 class GatewayBase(ComponentBase):
     def __init__(self, info, **kwargs):
         super().__init__(info, **kwargs)
         self.gateway_id = self.get_config("gateway_id", "default-change-me")
+        self.system_purpose_prompt_suffix = ""
         self.history_instance = self._initialize_history()
 
-    def _initialize_history(self):
+    def _initialize_history(self) -> HistoryService:
         self.use_history = self.get_config("retain_history", True)
 
         if not self.use_history:
@@ -19,6 +21,9 @@ class GatewayBase(ComponentBase):
 
         history_config = self.get_config("history_config", {})
 
+        if history_config.get("enable_long_term_memory", False):
+            self.system_purpose_prompt_suffix = LONG_TERM_MEMORY_PROMPT
+
         try:
             return HistoryService(
                 history_config, identifier=self.gateway_id + "_history"
@@ -27,6 +32,7 @@ class GatewayBase(ComponentBase):
             log.error("Failed to load history class: %s", e)
             raise
 
+
     def _initialize_identity_component(self):
         identity_config = self.get_config("identity", {})
         identity_key_field = self.get_config("identity_key_field", DEFAULT_IDENTITY_KEY_FIELD)

solace_agent_mesh/gateway/components/gateway_input.py

@@ -6,7 +6,7 @@ from uuid import uuid4
 from solace_ai_connector.common.message import Message
 from solace_ai_connector.common.log import log
 from ...services.file_service import FileService
-from ...common.constants import DEFAULT_IDENTITY_KEY_FIELD
+from ...common.constants import DEFAULT_IDENTITY_KEY_FIELD, HISTORY_USER_ROLE
 from .gateway_base import GatewayBase
 
 info = {
@@ -130,7 +130,7 @@ class GatewayInput(GatewayBase):
 
     def __init__(self, **kwargs):
         super().__init__(info, **kwargs)
-        self.system_purpose = self.get_config("system_purpose", DEFAULT_SYSTEM_PURPOSE)
+        self.system_purpose = self.get_config("system_purpose", DEFAULT_SYSTEM_PURPOSE) + self.system_purpose_prompt_suffix
         self.interaction_type = self.get_config(
             "interaction_type", DEFAULT_INTERACTION_TYPE
         )
@@ -209,17 +209,20 @@ class GatewayInput(GatewayBase):
 
             copied_data["history"] = []
             if self.use_history:
+                other_history_props = {
+                    "identity": identity_value,
+                }
                 prompt = data.get("text", "")
-                self.history_instance.store_history(session_id, "user", prompt)
+                self.history_instance.store_history(session_id, HISTORY_USER_ROLE, prompt, other_history_props)
 
                 for file in attached_files:
-                    self.history_instance.store_file(session_id, file)
+                    self.history_instance.store_file(session_id, file )
 
                 # retrieve all files for the session
                 available_files = self.history_instance.get_files(session_id)
 
                 # Add history to the data
-                copied_data["history"] = self.history_instance.get_history(session_id)
+                copied_data["history"] = self.history_instance.get_history(session_id, other_history_props)
 
             available_files = json.dumps(available_files)
         except Exception as e:

solace_agent_mesh/gateway/components/gateway_output.py

@@ -5,7 +5,7 @@ from solace_ai_connector.common.log import log
 from .gateway_base import GatewayBase
 from ...services.file_service import FileService
 from ...common.utils import files_to_block_text
-
+from ...common.constants import HISTORY_ASSISTANT_ROLE
 
 info = {
     "class_name": "GatewayOutput",
@@ -170,6 +170,7 @@ class GatewayOutput(GatewayBase):
         file_service = FileService()
         user_properties = message.get_user_properties()
         session_id = user_properties.get("session_id")
+        identity_value = user_properties.get("identity")
         files = data.get("files", [])
 
         # Extract the interface queue ID
@@ -181,6 +182,10 @@ class GatewayOutput(GatewayBase):
         )
 
         if self.use_history and session_id:
+            other_history_props = {
+                "identity": identity_value,
+            }
+
             topic = message.get_topic()
             content = data.get("text") or ""
 
@@ -193,14 +198,18 @@ class GatewayOutput(GatewayBase):
                 and data.get("last_chunk")
                 and "text" in data
             ):
+                actions_called = user_properties.get("actions_called", [])
+                if actions_called:
+                    self.history_instance.store_actions(session_id, actions_called)
+                
                 if content:
                     self.history_instance.store_history(
-                        session_id, "assistant", content
+                        session_id, HISTORY_ASSISTANT_ROLE, content, other_history_props
                     )
 
             for file in files:
                 self.history_instance.store_history(
-                    session_id, "assistant", f'\n[Returned file: {{name: {file.get("name")}, url: {file.get("url")}}}]\n'
+                    session_id, HISTORY_ASSISTANT_ROLE, f'\n[Returned file: {{name: {file.get("name")}, url: {file.get("url")}}}]\n', other_history_props
                 )
                 self.history_instance.store_file(session_id, file)
 

solace_agent_mesh/orchestrator/components/orchestrator_stimulus_processor_component.py

@@ -14,7 +14,7 @@ import yaml
 from solace_ai_connector.common.log import log
 from solace_ai_connector.common.message import Message
 
-from ...common.constants import ORCHESTRATOR_COMPONENT_NAME
+from ...common.constants import ORCHESTRATOR_COMPONENT_NAME, HISTORY_MEMORY_ROLE
 from ...services.llm_service.components.llm_request_component import LLMRequestComponent, info as base_info
 from ...services.middleware_service.middleware_service import MiddlewareService
 from ...services.file_service import FileService
@@ -128,6 +128,18 @@ class OrchestratorStimulusProcessorComponent(LLMRequestComponent):
 
         user_properties = message.get_user_properties()
         user_properties['timestamp_end'] = time()
+
+        actions_called = []
+        if results:
+            for result in results:
+                if result.get("payload", {}).get("action_name"):
+                    actions_called.append({
+                        "agent_name": result.get("payload", {}).get("agent_name"),
+                        "action_name": result.get("payload", {}).get("action_name"),
+                        "action_params": result.get("payload", {}).get("action_params"),
+                    })
+        user_properties['actions_called'] = actions_called
+
         message.set_user_properties(user_properties)
 
         return results
@@ -206,7 +218,7 @@ class OrchestratorStimulusProcessorComponent(LLMRequestComponent):
         }
 
         # Get the prompts
-        gateway_history = self.get_gateway_history(data)
+        gateway_history, memory_history = self.get_gateway_history(data)
         system_prompt = SystemPrompt(full_input, examples)
         if action_response_reinvoke:
             user_prompt = ActionResponsePrompt(
@@ -217,6 +229,9 @@ class OrchestratorStimulusProcessorComponent(LLMRequestComponent):
             user_prompt = UserStimulusPrompt(
                 full_input, gateway_history, errors, has_files
             )
+            if memory_history:
+                self.history.store_history(stimulus_uuid, "system", memory_history)
+
 
         # Store the user prompt in the history
         self.history.store_history(stimulus_uuid, "user", user_prompt)
@@ -375,6 +390,9 @@ class OrchestratorStimulusProcessorComponent(LLMRequestComponent):
 
     def get_gateway_history(self, data):
         gateway_history = data.get("history", [])
+        memory_history = None
+        if gateway_history and gateway_history[0].get("role") == HISTORY_MEMORY_ROLE:
+            memory_history =gateway_history[0].get("content")
         # Returning the history from the last user message
         first_user_idx = None
         last_user_idx = None
@@ -385,13 +403,13 @@ class OrchestratorStimulusProcessorComponent(LLMRequestComponent):
                 last_user_idx = idx
 
         if first_user_idx is None:
-            return []  # No user query found
+            return [], memory_history  # No user query found
 
         if not last_user_idx > first_user_idx:
             # Latest user message is already handled by orchestator history
-            return []
+            return [], memory_history
 
-        return gateway_history[first_user_idx:last_user_idx]
+        return gateway_history[first_user_idx:last_user_idx], memory_history
 
     def get_user_input(self, chat_text):
 

solace_agent_mesh/orchestrator/orchestrator_prompt.py

@@ -2,42 +2,39 @@ from typing import Dict, Any, List
 import yaml
 from langchain_core.messages import HumanMessage
 from ..services.file_service import FS_PROTOCOL, Types, LLM_QUERY_OPTIONS, TRANSFORMERS
+from solace_ai_connector.common.log import log
 
 # Cap the number of examples so we don't overwhelm the LLM
 MAX_SYSTEM_PROMPT_EXAMPLES = 6
 
 # Examples that should always be included in the prompt
 fixed_examples = [
-    """    <example>
-        <example_docstring>
-          This example shows a stimulus from a chatbot gateway in which a user is asking about the top stories on the website hacker news. The web_request is not yet open, so the change_agent_status action is invoked to open the web_request agent.
-        </example_docstring> 
-        <example_stimulus>
-          <{tp}stimulus starting_id="1"/>
-          What is the top story on hacker news?
-          </{tp}stimulus>
-          <{tp}stimulus_metadata>
-          local_time: 2024-11-06 15:58:04 EST-0500 (Wednesday)
-          </{tp}stimulus_metadata>    
-        </example_stimulus>
-        <example_response>
-          <{tp}reasoning>
-          - User is asking for the top story on Hacker News
-          - We need to use the web_request agent to fetch the latest information
-          - The web_request agent is currently closed, so we need to open it first
-          - After opening the agent, we\'ll need to make a web request to Hacker News
-          </{tp}reasoning>
-
-          <{tp}status_update>To get the latest top story from Hacker News, I\'ll need to access the web. I\'m preparing to do that now.</{tp}status_update>
-
-          <{tp}invoke_action agent="global" action="change_agent_status">
-          <{tp}parameter name="agent_name">web_request</{tp}parameter>
-          <{tp}parameter name="new_state">open</{tp}parameter>
-          </{tp}invoke_action>
-        </example_response>
-    </example>
-"""
-]
+                    {
+                        "docstring": "This example shows a stimulus from a chatbot gateway in which a user is asking about the top stories on the website hacker news. The web_request is not yet open, so the change_agent_status action is invoked to open the web_request agent.",
+                        "tag_prefix_placeholder": "{tp}",
+                        "starting_id": "1",
+                        "user_input": "What is the top story on hacker news?",
+                        "metadata": [
+                            "local_time: 2024-11-06 15:58:04 EST-0500 (Wednesday)"
+                        ],
+                        "reasoning": [
+                            "- User is asking for the top story on Hacker News",
+                            "- We need to use the web_request agent to fetch the latest information",
+                            "- The web_request agent is currently closed, so we need to open it first",
+                            "- After opening the agent, we'll need to make a web request to Hacker News"
+                        ],
+                        "response_text": "",
+                        "status_update": "To get the latest top story from Hacker News, I'll need to access the web. I'm preparing to do that now.",
+                        "action": {
+                            "agent": "global",
+                            "name": "change_agent_status",
+                            "parameters": {
+                                "agent_name": "web_request",
+                                "new_state": "open"
+                            }
+                        }
+                    }
+                  ]
 
 
 def get_file_handling_prompt(tp: str) -> str:
@@ -102,6 +99,7 @@ def get_file_handling_prompt(tp: str) -> str:
     For all files, there's `encoding` parameter, this is used to encode the file format. The supported values are `datauri`, `base64`, `zip`, and `gzip`. Use this to convert a file to ZIP or datauri for HTML encoding.
         For example (return a file as zip): `{FS_PROTOCOL}://c27e6908-55d5-4ce0-bc93-a8e28f84be12_annual_report.csv?encoding=zip&resolve=true`
         Example 2 (HTML tag must be datauri encoded): `<img src="{FS_PROTOCOL}://c9183b0f-fd11-48b4-ad8f-df221bff3da9_generated_image.png?encoding=datauri&resolve=true" alt="image">`
+        Example 3 (return a regular image directly): `amfs://a1b2c3d4-5e6f-7g8h-9i0j-1k2l3m4n5o6p_photo.png?resolve=true` - When returning images directly in messaging platforms like Slack, don't use any encoding parameter
 
     For all files, there's `resolve=true` parameter, this is used to resolve the url to the actual data before processing.
         For example: `{FS_PROTOCOL}://577c928c-c126-42d8-8a48-020b93096110_names.csv?resolve=true`
@@ -156,11 +154,15 @@ def create_examples(
         String containing all examples with replaced placeholders
     """
     examples = (fixed_examples + agent_examples)[:MAX_SYSTEM_PROMPT_EXAMPLES]
-    return "\n".join([example.replace("{tp}", tp) for example in examples])
+    formatted_examples = format_examples_by_llm_type(examples)
+    
+    return "\n".join([example.replace("{tp}", tp) for example in formatted_examples])
 
 
 def SystemPrompt(info: Dict[str, Any], action_examples: List[str]) -> str:
-    response_format_prompt = info.get("response_format_prompt", "")
+    tp = info["tag_prefix"]
+    response_format_prompt = info.get("response_format_prompt", "") or ""
+    response_format_prompt = response_format_prompt.replace("{{tag_prefix}}", tp)
     response_guidelines_prompt = (
         f"<response_guidelines>\nConsider the following when generating a response to the originator:\n"
         f"{response_format_prompt}</response_guidelines>"
@@ -183,7 +185,6 @@ def SystemPrompt(info: Dict[str, Any], action_examples: List[str]) -> str:
         )
 
     # Merged
-    tp = info["tag_prefix"]
     examples = create_examples(fixed_examples, action_examples, tp)
 
     handling_files = get_file_handling_prompt(tp)
@@ -216,7 +217,21 @@ The assistant's behavior aligns with the system purpose specified below:
     5. Do not perform any other actions besides opening the required agents in this step.
   - Report generation:
     1. If a report is requested and no format is specified, create the report in an HTML file.
-    2. Generate each section of the report independently and store it in the file service with create_file action. When finishing the report, combine the sections using amfs urls with the resolve=true query parameter to insert the sections into the main document. When generating HTML, create the header first with all the necessary CSS and JS links so that it is clear what css the rest of the document will use.
+    2. Generate each section of the report independently and store it in the file service with create_file action. When finishing the report, combine the sections using amfs urls with the resolve=true query parameter to insert the sections into the main document. When inserting amfs HTML URLs into the HTML document, place them directly in the document without any surrounding tags or brackets. Here is an example of the body section of an HTML report combining multiple sections:
+        <body>
+            <!-- Title -->
+            <h1>Report Title</h1>
+
+            <!-- Section 1 -->
+            amfs://xxxxxx.html?resolve=true
+
+            <!-- Section 2 -->
+            amfs://yyyyyy.html?resolve=true
+
+            <!-- Section 3 -->
+            amfs://zzzzzz.html?resolve=true
+        </body>
+        When generating HTML, create the header first with all the necessary CSS and JS links so that it is clear what css the rest of the document will use.
     3. Images are always very useful in reports, so the assistant will add them when appropriate. If images are embedded in html, they must be resolved and converted to datauri format or they won't render in the final document. This can be done by using the encoding=datauri&resolve=true in the amfs link. For example, <img src="amfs://xxxxxx.png?encoding=datauri&resolve=true". The assistant will take care of the rest. Images can be created in parallel
     4. During report generation in interactive sessions, the assistant will send lots of status messages to indicate what is happening. 
   - Handling stimuli with open agents:
@@ -330,7 +345,7 @@ def UserStimulusPrompt(
         f"\tExample of returning a file as zip: <{info['tag_prefix']}file><url>{FS_PROTOCOL}://519321d8-3506-4f8d-9377-e5d6ce74d917_filename.csv?encoding=zip</url></{info['tag_prefix']}file>.\n"
         f"You can also optionally return a non-persistent temporary file using the format <{info['tag_prefix']}file name=\"filename.csv\" mime_type=\"text/csv\">\n<data> data </data>\n</{info['tag_prefix']}file>.\n"
         f" can't nest `<{info['tag_prefix']}file>` tags inside the `<data>` tag. If you need to address another file within the data, use the URL with the `resolve=true` query parameter.\n"
-        "When using a {FS_PROTOCOL} URL outside of a file block, always include the 'resolve=true' query parameter to ensure the URL is resolved to the actual data.\n"
+        f"When using a {FS_PROTOCOL} URL outside of a file block, always include the 'resolve=true' query parameter to ensure the URL is resolved to the actual data.\n"
         f"</{info['tag_prefix']}file_instructions>"
     )
 
@@ -413,3 +428,108 @@ you should ask the originator for more information. Include links to the source
 {context}
 </context>
 """
+
+
+def format_examples_by_llm_type(examples: list, llm_type: str = "anthropic") -> list:
+    """
+    Render examples based on llm type
+    
+    Args:
+        llm_type (str): The type of LLM to render examples for (default: "anthropic")
+        examples (list): List of examples in model-agnostic format
+        
+    Returns:
+        list: List of examples formatted for the specified LLM
+    """
+    formatted_examples = []
+
+    if llm_type == "anthropic":
+        for example in examples:
+            formatted_example = format_example_for_anthropic(example)
+            formatted_examples.append(formatted_example)
+    else:
+        log.error(f"Unsupported LLM type: {llm_type}")
+
+    return formatted_examples
+
+def format_example_for_anthropic(example: dict) -> str:
+    """
+    Format an example for the Anthropic's LLMs
+    """
+    
+    tag_prefix = example.get("tag_prefix_placeholder", "t123")
+    starting_id = example.get("starting_id", "1")
+    docstring = example.get("docstring", "")
+    user_input = example.get("user_input", "")
+    metadata_lines = example.get("metadata", [])
+    reasoning_lines = example.get("reasoning", [])
+    response_text = example.get("response_text", "")
+    
+    # Start building the XML structure, add the description and user input
+    xml_content = f"""<example>
+        <example_docstring>
+            {docstring}
+        </example_docstring>
+        <example_stimulus>
+            <{tag_prefix}stimulus starting_id="{starting_id}">
+            {user_input}
+            </{tag_prefix}stimulus>
+            <{tag_prefix}stimulus_metadata>
+            """
+    
+    # Add metadata lines
+    for metadata_line in metadata_lines:
+        xml_content += f"{metadata_line}\n"
+    
+    xml_content += f"""</{tag_prefix}stimulus_metadata>
+        </example_stimulus>
+        <example_response>
+            <{tag_prefix}reasoning>
+            """
+    
+    # Add reasoning lines
+    for reasoning_line in reasoning_lines:
+        xml_content += f"{reasoning_line}\n"
+    
+    xml_content += f"""</{tag_prefix}reasoning>
+            {response_text}"""
+    
+    # Add action invocation section
+    if "action" in example:
+        action_data = example.get("action", {})
+        status_update = example.get("status_update", "")
+        agent_name = action_data.get("agent", "")
+        action_name = action_data.get("name", "")
+        
+        xml_content += f"""
+            <{tag_prefix}status_update>{status_update}</{tag_prefix}status_update>
+            <{tag_prefix}invoke_action agent="{agent_name}" action="{action_name}">"""
+        
+        # Handle parameters as dictionary
+        parameter_dict = action_data.get("parameters", {})
+        for param_name, param_value in parameter_dict.items():
+            xml_content += f"""
+            <{tag_prefix}parameter name="{param_name}">"""
+            
+            # Handle parameter names and values (as lists)
+            if isinstance(param_value, list):
+                for line in param_value:
+                    xml_content += f"\n{line}"
+                xml_content += "\n"
+            else:
+                # For simple string values
+                xml_content += f"{param_value}"
+                
+            xml_content += f"</{tag_prefix}parameter>\n"
+    
+        xml_content += f"</{tag_prefix}invoke_action>"
+    
+    # Close the XML structure
+    xml_content += """
+        </example_response>
+    </example>
+    """
+
+    return xml_content
+
+LONG_TERM_MEMORY_PROMPT = " - You are capable of remembering things and have long-term memory, this happens automatically."

solace_agent_mesh/services/file_service/file_service.py

@@ -392,6 +392,11 @@ class FileService(AutoExpiry, metaclass=AutoExpirySingletonMeta):
             url = url[5:-6].strip()
         if url.endswith('"') or url.endswith("'") or url.endswith(","):
             url = url[:-1]
+        if url.endswith(")"):
+            open_parenthesis_count = url.count("(")
+            close_parenthesis_count = url.count(")")
+            if close_parenthesis_count > open_parenthesis_count:
+                url = url[:-1]
         return url
 
     @staticmethod

solace_agent_mesh/services/file_service/file_service_constants.py

@@ -40,7 +40,7 @@ BLOCK_TAG_KEYS = [
 Keys to be treated as tags in the file block, and not file attributes.
 """
 
-FS_URL_REGEX = r"""<url>\s*({protocol}:\/\/.+?)\s*<\/url>|{protocol}:\/\/[^\s\?]+(\s|$)|{protocol}:\/\/[^\n\?]+\?.*?(?=\n|$)|\"({protocol}:\/\/[^\"]+?)(\"|\n)|'({protocol}:\/\/[^']+?)('|\n)""".format(
+FS_URL_REGEX = r"""<url>\s*({protocol}:\/\/.+?)\s*<\/url>|{protocol}:\/\/[^\s\?]+(\s|$)|{protocol}:\/\/[^\n\?]+\?.*?(?=\n|$|>)|\"({protocol}:\/\/[^\"]+?)(\"|\n)|'({protocol}:\/\/[^']+?)('|\n)""".format(
     protocol=FS_PROTOCOL
 )
 """

solace_agent_mesh/services/file_service/file_transformations.py

@@ -82,6 +82,8 @@ def apply_file_transformations(
         return file
     text_mime_type_regex = r"text/.*|.*csv|.*json|.*xml|.*yaml|.*x-yaml|.*txt"
     mime_type = metadata.get("mime_type", "")
+    if mime_type is None:
+        mime_type = ""
     name = metadata.get("name", "unknown")
     other = {
         "mime_type": mime_type,
@@ -126,6 +128,14 @@ def apply_file_transformations(
                 return file
 
         if not isinstance(data, str):
-            data = json.dumps(data)
+            # Convert bytes to string
+            if isinstance(data, bytes):
+                try:
+                    data = data.decode("utf-8")
+                except UnicodeDecodeError:
+                    data = base64.b64encode(data).decode("utf-8")
+                    data = f"data:{mime_type};base64,{data}"
+            else:
+                data = json.dumps(data)
 
         return data

solace_agent_mesh/services/file_service/file_utils.py

@@ -34,6 +34,8 @@ def get_str_type(value: str) -> str:
     Returns:
     - str: The type of the value.
     """
+    if not value:
+        return Types.NULL
     if value.isdigit():
         return Types.INT
     elif value.replace(".", "", 1).isdigit():

solace_agent_mesh/services/history_service/history_providers/base_history_provider.py

@@ -1,79 +1,54 @@
 from abc import ABC, abstractmethod
-from typing import Union
 
 class BaseHistoryProvider(ABC):
 
     def __init__(self, config=None):
-        self.config = config
-        self.max_turns = self.config.get("max_turns")
-        self.max_characters = self.config.get("max_characters")
-        self.enforce_alternate_message_roles = self.config.get("enforce_alternate_message_roles")
-
-    @abstractmethod
-    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.
-        """
-        raise NotImplementedError("Method not implemented")
-
+        self.config = config or {}
+    
     @abstractmethod
-    def get_history(self, session_id: str):
+    def get_all_sessions(self) -> list[str]:
         """
-        Retrieve the entire history.
-
-        :param session_id: The session identifier.
-        :return: The complete history.
+        Retrieve all session identifiers.
         """
         raise NotImplementedError("Method not implemented")
 
     @abstractmethod
-    def store_file(self, session_id: str, file: dict):
+    def get_session(self, session_id: str)->dict:
         """
-        Store a file in the history.
+        Retrieve the session .
 
         :param session_id: The session identifier.
-        :param file: The file metadata to be stored in the history.
+        :return: The session metadata.
         """
         raise NotImplementedError("Method not implemented")
-
+    
     @abstractmethod
-    def get_files(self, session_id: str):
+    def delete_session(self, session_id: str):
         """
-        Retrieve the files for a session.
+        Delete the session.
 
         :param session_id: The session identifier.
-        :return: The files for the session.
         """
         raise NotImplementedError("Method not implemented")
     
     @abstractmethod
-    def get_session_meta(self, session_id: str):
+    def store_session(self, session_id: str, data: dict):
         """
-        Retrieve the session metadata.
+        Store the session metadata.
 
         :param session_id: The session identifier.
-        :return: The session metadata.
+        :param data: The session data to be stored.
         """
         raise NotImplementedError("Method not implemented")
+    
 
-    @abstractmethod
-    def clear_history(self, session_id: str, keep_levels=0, clear_files=True):
-        """
-        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).
-        :param clear_files: Whether to clear associated files. Default is True.
+    def update_session(self, session_id: str, data: dict):
         """
-        raise NotImplementedError("Method not implemented")
+        Update data in the store using the partial data provided.
 
-    @abstractmethod
-    def get_all_sessions(self) -> list[str]:
+        :param key: The key to update the data under.
+        :param data: The data to update.
         """
-        Retrieve all session identifiers.
-        """
-        raise NotImplementedError("Method not implemented")
+        history = self.get_session(session_id).copy()
+        history.update(data)
+        self.store_session(session_id, history)