mito-ai 0.1.33__py3-none-any.whl → 0.1.49__py3-none-any.whl

mito_ai/__init__.py

@@ -5,12 +5,33 @@ from typing import List, Dict
 from jupyter_server.utils import url_path_join
 from mito_ai.completions.handlers import CompletionHandler
 from mito_ai.completions.providers import OpenAIProvider
-from mito_ai.app_builder.handlers import AppBuilderHandler
+from mito_ai.completions.message_history import GlobalMessageHistory
+from mito_ai.app_deploy.handlers import AppDeployHandler
+from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
 from mito_ai.log.urls import get_log_urls
 from mito_ai.version_check import VersionCheckHandler
 from mito_ai.db.urls import get_db_urls
 from mito_ai.settings.urls import get_settings_urls
 from mito_ai.rules.urls import get_rules_urls
+from mito_ai.auth.urls import get_auth_urls
+from mito_ai.streamlit_preview.urls import get_streamlit_preview_urls
+from mito_ai.app_manager.handlers import AppManagerHandler
+from mito_ai.file_uploads.urls import get_file_uploads_urls
+from mito_ai.user.urls import get_user_urls
+from mito_ai.chat_history.urls import get_chat_history_urls
+
+# Force Matplotlib to use the Jupyter inline backend.
+# Background: importing Streamlit sets os.environ["MPLBACKEND"] = "Agg" very early.
+# In a Jupyter kernel, that selects a non‑interactive canvas and can trigger:
+#   "UserWarning: FigureCanvasAgg is non-interactive, and thus cannot be shown"
+# which prevents figures from rendering in notebook outputs.
+# We preempt this by selecting the canonical Jupyter inline backend BEFORE any
+# Matplotlib import, so figures render inline reliably. This must run very early.
+# See: https://github.com/streamlit/streamlit/issues/9640
+
+import os
+os.environ["MPLBACKEND"] = "module://matplotlib_inline.backend_inline"
+
 try:
     from _version import __version__
 except ImportError:
@@ -44,31 +65,50 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     base_url = web_app.settings["base_url"]
 
     open_ai_provider = OpenAIProvider(config=server_app.config)
+    
+    # Create a single GlobalMessageHistory instance for the entire server
+    # This ensures thread-safe access to the .mito/ai-chats directory
+    global_message_history = GlobalMessageHistory()
 
     # WebSocket handlers
     handlers = [
         (
             url_path_join(base_url, "mito-ai", "completions"),
             CompletionHandler,
-            {"llm": open_ai_provider},
+            {"llm": open_ai_provider, "message_history": global_message_history},
+        ),
+        (
+            url_path_join(base_url, "mito-ai", "app-deploy"),
+            AppDeployHandler,
+            {}
         ),
         (
-            url_path_join(base_url, "mito-ai", "app-builder"),
-            AppBuilderHandler,
+            url_path_join(base_url, "mito-ai", "streamlit-preview"),
+            StreamlitPreviewHandler,
             {}
         ),
         (
             url_path_join(base_url, "mito-ai", "version-check"),
             VersionCheckHandler,
             {},
+        ),
+        (
+            url_path_join(base_url, "mito-ai", "app-manager"),
+            AppManagerHandler,
+            {}
         )
     ]
     
     # REST API endpoints
-    handlers.extend(get_db_urls(base_url))
-    handlers.extend(get_settings_urls(base_url))
-    handlers.extend(get_rules_urls(base_url))
-    handlers.extend(get_log_urls(base_url))
-    
+    handlers.extend(get_db_urls(base_url))  # type: ignore
+    handlers.extend(get_settings_urls(base_url))  # type: ignore
+    handlers.extend(get_rules_urls(base_url))  # type: ignore
+    handlers.extend(get_log_urls(base_url, open_ai_provider.key_type))  # type: ignore
+    handlers.extend(get_auth_urls(base_url))  # type: ignore
+    handlers.extend(get_streamlit_preview_urls(base_url))  # type: ignore
+    handlers.extend(get_file_uploads_urls(base_url)) # type: ignore
+    handlers.extend(get_user_urls(base_url)) # type: ignore
+    handlers.extend(get_chat_history_urls(base_url, global_message_history)) # type: ignore
+
     web_app.add_handlers(host_pattern, handlers)
     server_app.log.info("Loaded the mito_ai server extension")

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.33'
+__version__ = VERSION = '0.1.49'

mito_ai/anthropic_client.py

@@ -5,8 +5,9 @@ import json
 import anthropic
 from typing import Dict, Any, Optional, Tuple, Union, Callable, List, cast
 
-from anthropic.types import Message, MessageParam
+from anthropic.types import Message, MessageParam, TextBlockParam
 from mito_ai.completions.models import ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
+from mito_ai.constants import MESSAGE_HISTORY_TRIM_THRESHOLD
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.utils.anthropic_utils import get_anthropic_completion_from_mito_server, stream_anthropic_completion_from_mito_server, get_anthropic_completion_function_params
 
@@ -15,8 +16,6 @@ from mito_ai.utils.anthropic_utils import get_anthropic_completion_from_mito_ser
 # 8192 is the maximum allowed number of output tokens for claude-3-5-haiku-20241022
 MAX_TOKENS = 8_000
 
-ANTHROPIC_FAST_MODEL = "claude-3-5-haiku-latest"
-
 def extract_and_parse_anthropic_json_response(response: Message) -> Union[object, Any]:
     """
     Extracts and parses the JSON response from the Claude API.
@@ -53,12 +52,12 @@ def extract_and_parse_anthropic_json_response(response: Message) -> Union[object
 
 
 def get_anthropic_system_prompt_and_messages(messages: List[ChatCompletionMessageParam]) -> Tuple[
-    Union[str, anthropic.NotGiven], List[MessageParam]]:
+    Union[str, anthropic.Omit], List[MessageParam]]:
     """
     Convert a list of OpenAI messages to a list of Anthropic messages.
     """
     
-    system_prompt: Union[str, anthropic.NotGiven] = anthropic.NotGiven()
+    system_prompt: Union[str, anthropic.Omit] = anthropic.Omit()
     anthropic_messages: List[MessageParam] = []
 
     for message in messages:
@@ -126,14 +125,97 @@ def get_anthropic_system_prompt_and_messages(messages: List[ChatCompletionMessag
     return system_prompt, anthropic_messages
 
 
+def add_cache_control_to_message(message: MessageParam) -> MessageParam:
+    """
+    Adds cache_control to a message's content.
+    Handles both string content and list of content blocks.
+    """
+    content = message.get("content")
+    
+    if isinstance(content, str):
+        # Simple string content - convert to list format with cache_control
+        return {
+            "role": message["role"],
+            "content": [
+                {
+                    "type": "text",
+                    "text": content,
+                    "cache_control": {"type": "ephemeral"}
+                }
+            ]
+        }
+    
+    elif isinstance(content, list) and len(content) > 0:
+        # List of content blocks - add cache_control to last block
+        content_blocks = content.copy()
+        last_block = content_blocks[-1].copy()
+        last_block["cache_control"] = {"type": "ephemeral"}
+        content_blocks[-1] = last_block
+        
+        return {
+            "role": message["role"],
+            "content": content_blocks
+        }
+    
+    else:
+        # Edge case: empty or malformed content
+        return message
+
+
+def get_anthropic_system_prompt_and_messages_with_caching(messages: List[ChatCompletionMessageParam]) -> Tuple[
+    Union[str, List[TextBlockParam], anthropic.Omit], List[MessageParam]]:
+    """
+    Convert a list of OpenAI messages to a list of Anthropic messages with caching applied.
+    
+    Caching Strategy:
+    1. System prompt (static) → Always cached
+    2. Stable conversation history → Cache at keep_recent boundary
+    3. Recent messages → Never cached (always fresh)
+    
+    The keep_recent parameter determines which messages are stable and won't be trimmed.
+    We cache at the keep_recent boundary because those messages are guaranteed to be stable.
+    """
+    
+    # Get the base system prompt and messages
+    system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+    
+    # 1. Cache the system prompt always
+    # If the system prompt is something like anthropic.Omit, we don't need to cache it
+    cached_system_prompt: Union[str, List[TextBlockParam], anthropic.Omit] = system_prompt
+    if isinstance(system_prompt, str):
+        cached_system_prompt = [{
+            "type": "text",
+            "text": system_prompt,
+            "cache_control": {"type": "ephemeral"}
+        }]
+    
+    # 2. Cache conversation history at the boundary where the messages are stable.
+    # Messages are stable after they are more than MESSAGE_HISTORY_TRIM_THRESHOLD old.
+    # At this point, the messages are not edited anymore, so they will not invalidate the cache.
+    # If we included the messages before the boundary in the cache, then every time we send a new
+    # message, we would invalidate the cache and we would never get a cache hit except for the system prompt.
+    messages_with_cache = []
+    
+    if len(anthropic_messages) > 0:
+        cache_boundary = len(anthropic_messages) - MESSAGE_HISTORY_TRIM_THRESHOLD - 1
+        
+        # Add all messages, but only add cache_control to the message at the boundary
+        for i, msg in enumerate(anthropic_messages):
+            if i == cache_boundary:
+                messages_with_cache.append(add_cache_control_to_message(msg))
+            else:
+                messages_with_cache.append(msg)
+    
+    return cached_system_prompt, messages_with_cache
+
+
 class AnthropicClient:
     """
     A client for interacting with the Anthropic API or the Mito server fallback.
     """
 
-    def __init__(self, api_key: Optional[str], model: str, timeout: int = 30, max_retries: int = 1):
+    def __init__(self, api_key: Optional[str], timeout: int = 30, max_retries: int = 1):
         self.api_key = api_key
-        self.model = model
         self.timeout = timeout
         self.max_retries = max_retries
         self.client: Optional[anthropic.Anthropic]
@@ -142,66 +224,66 @@ class AnthropicClient:
         else:
             self.client = None
 
-    async def request_completions(self, messages: List[ChatCompletionMessageParam],
-                                  response_format_info: Optional[ResponseFormatInfo] = None,
-                                  message_type: MessageType = MessageType.CHAT) -> Any:
+    async def request_completions(
+        self, messages: List[ChatCompletionMessageParam],
+        model: str,
+        response_format_info: Optional[ResponseFormatInfo] = None,
+        message_type: MessageType = MessageType.CHAT
+    ) -> Any:
         """
         Get a response from Claude or the Mito server that adheres to the AgentResponse format.
         """
-        try:
-            anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+        anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
+        
+        provider_data = get_anthropic_completion_function_params(
+            message_type=message_type,
+            model=model,
+            messages=anthropic_messages,
+            max_tokens=MAX_TOKENS,
+            temperature=0,
+            system=anthropic_system_prompt,
+            stream=None,
+            response_format_info=response_format_info
+        )
+        
+        if self.api_key:
+            # Unpack provider_data for direct API call
+            assert self.client is not None
+            response = self.client.messages.create(**provider_data)
             
-            provider_data = get_anthropic_completion_function_params(
-                model=self.model if response_format_info else ANTHROPIC_FAST_MODEL,
-                messages=anthropic_messages,
-                max_tokens=MAX_TOKENS,
-                temperature=0,
-                system=anthropic_system_prompt,
-                stream=None,
-                response_format_info=response_format_info
-            )
-            
-            if self.api_key:
-                # Unpack provider_data for direct API call
-                assert self.client is not None
-                response = self.client.messages.create(**provider_data)
-                if provider_data.get("tool_choice") is not None:
-                    result = extract_and_parse_anthropic_json_response(response)
-                    return json.dumps(result) if not isinstance(result, str) else result
-                else:
-                    content = response.content
-                    if content[0].type == "text":
-                        return content[0].text
-                    else:
-                        return ""
+            if provider_data.get("tool_choice") is not None:
+                result = extract_and_parse_anthropic_json_response(response)
+                return json.dumps(result) if not isinstance(result, str) else result
             else:
-                # Only pass provider_data to the server
-                response = await get_anthropic_completion_from_mito_server(
-                    model=provider_data["model"],
-                    max_tokens=provider_data["max_tokens"],
-                    temperature=provider_data["temperature"],
-                    system=provider_data["system"],
-                    messages=provider_data["messages"],
-                    tools=provider_data.get("tools"),
-                    tool_choice=provider_data.get("tool_choice"),
-                    message_type=message_type
-                )
-                return response
-        except anthropic.RateLimitError:
-            raise Exception("Rate limit exceeded. Please try again later or reduce your request frequency.")
-        except Exception as e:
-            return f"Error streaming content: {str(e)}"
+                content = response.content
+                if content[0].type == "text":
+                    return content[0].text
+                else:
+                    return ""
+        else:
+            # Only pass provider_data to the server
+            response = await get_anthropic_completion_from_mito_server(
+                model=provider_data["model"],
+                max_tokens=provider_data["max_tokens"],
+                temperature=provider_data["temperature"],
+                system=provider_data["system"],
+                messages=provider_data["messages"],
+                tools=provider_data.get("tools"),
+                tool_choice=provider_data.get("tool_choice"),
+                message_type=message_type
+            )
+            return response
 
-    async def stream_response(self, messages: List[ChatCompletionMessageParam], message_id: str, message_type: MessageType,
+    async def stream_completions(self, messages: List[ChatCompletionMessageParam], model: str, message_id: str, message_type: MessageType,
                               reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]) -> str:
         try:
-            anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+            anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
             accumulated_response = ""
 
             if self.api_key:
                 assert self.client is not None
                 stream = self.client.messages.create(
-                    model=self.model,
+                    model=model,
                     max_tokens=MAX_TOKENS,
                     temperature=0,
                     system=anthropic_system_prompt,
@@ -209,7 +291,6 @@ class AnthropicClient:
                     stream=True
                 )
 
-
                 for chunk in stream:
                     if chunk.type == "content_block_delta" and chunk.delta.type == "text_delta":
                         content = chunk.delta.text
@@ -229,24 +310,17 @@ class AnthropicClient:
 
             else:
                 async for stram_chunk in stream_anthropic_completion_from_mito_server(
-                    model=self.model,
+                    model=model,
                     max_tokens=MAX_TOKENS,
                     temperature=0,
                     system=anthropic_system_prompt,
                     messages=anthropic_messages,
                     stream=True,
-                    message_type=message_type
+                    message_type=message_type,
+                    reply_fn=reply_fn,
+                    message_id=message_id
                 ):
                     accumulated_response += stram_chunk
-                    reply_fn(CompletionStreamChunk(
-                        parent_id=message_id,
-                        chunk=CompletionItem(
-                            content=stram_chunk,
-                            isIncomplete=True,
-                            token=message_id,
-                        ),
-                        done=False,
-                    ))
 
             return accumulated_response
 
@@ -254,6 +328,7 @@ class AnthropicClient:
             raise Exception("Rate limit exceeded. Please try again later or reduce your request frequency.")
 
         except Exception as e:
-            return f"Error streaming content: {str(e)}"
+            print(f"Error streaming content: {str(e)}")
+            raise e
 
 

mito_ai/{app_builder → app_deploy}/__init__.py

@@ -3,4 +3,4 @@
 
 """App builder module for Mito AI."""
 
-from .handlers import AppBuilderHandler 
+from .handlers import AppDeployHandler 

mito_ai/app_deploy/app_deploy_utils.py

@@ -0,0 +1,44 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import zipfile
+import logging
+from typing import List, Optional
+
+from mito_ai.path_utils import AbsoluteNotebookDirPath
+
+def add_files_to_zip(
+    zip_path: str, 
+    notebook_dir_path: AbsoluteNotebookDirPath,
+    files_to_add: List[str], 
+    app_file_name: str,
+    logger: Optional[logging.Logger] = None
+) -> None:
+    """Create a zip file at zip_path and add the selected files/folders."""
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for file_to_add_rel_path in files_to_add:
+            
+            file_to_add_abs_path = os.path.join(notebook_dir_path, file_to_add_rel_path)
+
+            if os.path.isfile(file_to_add_abs_path):
+                basename = os.path.basename(file_to_add_abs_path)
+                
+                if basename == app_file_name:
+                    # For the actual app file, we want to write it just as app.py 
+                    # so our infra can always deploy using `streamlit run app.py`
+                    # without having to account for different app names
+                    zipf.write(file_to_add_abs_path, arcname='app.py')
+                else:
+                    # otherwise we want to keep the name as is so all references
+                    # to it from the app are correct
+                    zipf.write(file_to_add_abs_path, arcname=file_to_add_rel_path)
+            elif os.path.isdir(file_to_add_abs_path):
+                for root, _, files in os.walk(file_to_add_abs_path):
+                    for file in files:
+                        file_abs = os.path.join(root, file)
+                        arcname = os.path.relpath(file_abs, notebook_dir_path)
+                        zipf.write(file_abs, arcname=arcname)
+            else:
+                if logger:
+                    logger.warning(f"Skipping missing file: {file_to_add_abs_path}")