mito-ai 0.1.55__py3-none-any.whl → 0.1.56__py3-none-any.whl

mito_ai/tests/providers/test_anthropic_client.py

@@ -385,12 +385,16 @@ def test_caching_system_prompt_scenarios(messages, expected_system_type, expecte
 @pytest.mark.parametrize("message_count,expected_cache_boundary", [
     (1, None),  # 1 message, No cache boundary
     (3, None),  # 3 messages, No cache boundary
-    (5, 1),  # 5 messages, cache at index 2
-    (10, 6),  # 10 messages, cache at index 6
+    (5, 1),  # 5 messages, cache at index 1 (5 - 3 - 1 = 1)
+    (10, 6),  # 10 messages, cache at index 6 (10 - 3 - 1 = 6)
 ])
-def test_caching_conversation_history(message_count, expected_cache_boundary):
+def test_caching_conversation_history(message_count, expected_cache_boundary, monkeypatch):
     """Test that conversation history is cached at the keep_recent boundary for different message counts."""
     
+    # Mock MAX_TRIM_THRESHOLD to use a fixed value for testing
+    import mito_ai.anthropic_client as anthropic_client_module
+    monkeypatch.setattr(anthropic_client_module, 'MAX_TRIM_THRESHOLD', 3)
+    
     # Create messages based on the parameter
     messages: List[ChatCompletionMessageParam] = [
         ChatCompletionSystemMessageParam(role="system", content="You are helpful.")

mito_ai/utils/message_history_utils.py

@@ -2,47 +2,65 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import re
-from typing import List
-from mito_ai.constants import MESSAGE_HISTORY_TRIM_THRESHOLD
+from typing import List, Dict, Optional
 from openai.types.chat import ChatCompletionMessageParam
-from mito_ai.completions.prompt_builders.prompt_constants import (
-    ACTIVE_CELL_ID_SECTION_HEADING,
-    ACTIVE_CELL_OUTPUT_SECTION_HEADING,
-    GET_CELL_OUTPUT_TOOL_RESPONSE_SECTION_HEADING,
-    FILES_SECTION_HEADING,
-    STREAMLIT_APP_STATUS_SECTION_HEADING,
-    VARIABLES_SECTION_HEADING,
-    JUPYTER_NOTEBOOK_SECTION_HEADING,
-    CONTENT_REMOVED_PLACEHOLDER
-)
+from mito_ai.completions.prompt_builders.prompt_section_registry import get_all_section_classes
 
 
-def trim_sections_from_message_content(content: str) -> str:
-    """
-    Removes specific metadata sections from message content to reduce token count so
-    that users don't exceed the token limit for the LLM.
+def build_section_to_trim_message_index_mapping() -> Dict[str, Optional[int]]:
+    """Build mapping from section names to trim_after_messages thresholds."""
+    mapping = {}
+    for section_class in get_all_section_classes():
+        section_name = section_class.__name__.replace("Section", "")
+        mapping[section_name] = section_class.trim_after_messages
+    return mapping
 
-    These sections are replaced with a placeholder text.
-    """
 
-    # Replace metadata sections with placeholders
-    section_headings = [
-        FILES_SECTION_HEADING,
-        VARIABLES_SECTION_HEADING,
-        JUPYTER_NOTEBOOK_SECTION_HEADING,
-        GET_CELL_OUTPUT_TOOL_RESPONSE_SECTION_HEADING,
-        ACTIVE_CELL_OUTPUT_SECTION_HEADING,
-        ACTIVE_CELL_ID_SECTION_HEADING,
-        STREAMLIT_APP_STATUS_SECTION_HEADING
-    ]
+def trim_message_content(content: str, message_age: int) -> str:
+    """
+    Trims sections from XML string based on age and thresholds.
+    
+    Args:
+        content: The message content as a string (may contain XML tags)
+        message_age: The age of the message (0 = most recent, higher = older)
+    
+    Returns:
+        The content with trimmed sections removed
+    """
+    section_mapping = build_section_to_trim_message_index_mapping()
+    
+    # Special handling for Example sections - remove entirely if they should be trimmed
+    # Match: <Example name="...">...</Example> or <Example>...</Example>
+    example_pattern = r'<Example(?:\s+name="[^"]*")?>.*?</Example>'
+    example_matches = list(re.finditer(example_pattern, content, flags=re.DOTALL))
+    
+    for match in reversed(example_matches):  # Process from end to start to preserve indices
+        example_threshold = section_mapping.get("Example")
+        if example_threshold is not None and message_age >= example_threshold:
+            # Remove the entire Example block
+            content = content[:match.start()] + content[match.end():]
+    
+    # For other sections, parse and trim based on section_mapping
+    # Match XML tags like <SectionName>...</SectionName>
+    # Skip Example sections as they're handled separately above
+    for section_name, threshold in section_mapping.items():
+        if threshold is None:
+            # Never trim sections with None threshold
+            continue
+        
+        if section_name == "Example":
+            # Example sections are handled separately above
+            continue
+        
+        if message_age >= threshold:
+            # Pattern to match the section tag and its content
+            # Handles both self-closing and content tags
+            pattern = rf'<{re.escape(section_name)}>.*?</{re.escape(section_name)}>'
+            matches = list(re.finditer(pattern, content, flags=re.DOTALL))
+            
+            for match in reversed(matches):  # Process from end to start
+                content = content[:match.start()] + content[match.end():]
     
-    for heading in section_headings:
-        content = re.sub(
-            f"{re.escape(heading)}\n(?:.+\n)+",
-            f"{heading} {CONTENT_REMOVED_PLACEHOLDER}\n",
-            content,
-        )
-
     return content
 
 
@@ -50,28 +68,30 @@ def trim_old_messages(messages: List[ChatCompletionMessageParam]) -> List[ChatCo
     """
     Trims metadata sections from messages that are older than the specified number of recent messages.
     We do this in order to reduce the token count of the messages, which helps us stay under the token limit for the LLM.
+    
+    Only trims user messages for now, but the design allows for easy extension to system/assistant messages.
     """
-    if len(messages) <= MESSAGE_HISTORY_TRIM_THRESHOLD:
-        return messages
         
     # Process all messages except the keep_recent most recent ones. 
     # Only trim user messages, which is where this metadata lives. 
     # We want to not edit the system messages, as they contain important information / examples.
-    for i in range(len(messages) - MESSAGE_HISTORY_TRIM_THRESHOLD):
-        content = messages[i].get("content")
+    total_messages = len(messages)
+    for i in range(total_messages):
         
+        # Only trim user messages
         is_user_message = messages[i].get("role") == "user"
         if not is_user_message: 
             continue
         
         content = messages[i].get("content")
-        
         if content is None:
             continue
         
+        message_age = total_messages - i - 1
+        
         if isinstance(content, str):
             # If content is just a string, then we just trim the metadata sections
-            messages[i]["content"] = trim_sections_from_message_content(content)
+            messages[i]["content"] = trim_message_content(content, message_age)
         else: 
             # Otherwise, we get rid of the image_url section and just keep the trimmed text
             # We assume that there is only one text section in the content
@@ -80,8 +100,12 @@ def trim_old_messages(messages: List[ChatCompletionMessageParam]) -> List[ChatCo
                 if section.get("type") == "text" and "text" in section:
                     text_content = section["text"] #type: ignore
                     break
-                
-            messages[i]["content"] = trim_sections_from_message_content(text_content)        
-
+            
+            trimmed_text = trim_message_content(text_content, message_age)
+            # Update the text section with trimmed content
+            for section in content:
+                if section.get("type") == "text" and "text" in section:
+                    section["text"] = trimmed_text  #type: ignore
+                    break
 
     return messages

{mito_ai-0.1.55.data → mito_ai-0.1.56.data}/data/share/jupyter/labextensions/mito_ai/build_log.json

@@ -720,7 +720,7 @@
             "semver": {},
             "vscode-diff": {},
             "mito_ai": {
-              "version": "0.1.55",
+              "version": "0.1.56",
               "singleton": true,
               "import": "/home/runner/work/mito/mito/mito-ai/lib/index.js"
             }

{mito_ai-0.1.55.data → mito_ai-0.1.56.data}/data/share/jupyter/labextensions/mito_ai/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.55",
+  "version": "0.1.56",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -141,7 +141,7 @@
     "schemaDir": "schema",
     "themePath": "style/theme/theme.css",
     "_build": {
-      "load": "static/remoteEntry.9dfbffc3592eb6f0aef9.js",
+      "load": "static/remoteEntry.1e7b5cf362385f109883.js",
       "extension": "./extension",
       "style": "./style"
     }

{mito_ai-0.1.55.data → mito_ai-0.1.56.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/package.json.orig

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.55",
+  "version": "0.1.56",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",