mito-ai 0.1.45__py3-none-any.whl → 0.1.47__py3-none-any.whl

mito_ai/__init__.py

@@ -5,6 +5,7 @@ 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.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
@@ -16,6 +17,8 @@ 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.
@@ -62,13 +65,17 @@ 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"),
@@ -100,6 +107,8 @@ def _load_jupyter_server_extension(server_app) -> None: # 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.45'
+__version__ = VERSION = '0.1.47'

mito_ai/anthropic_client.py

@@ -5,9 +5,9 @@ import json
 import anthropic
 from typing import Dict, Any, Optional, Tuple, Union, Callable, List, cast
 
-from anthropic.types import Message, MessageParam
-from mito_ai.completions.models import CompletionError, ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
-from mito_ai.utils.mito_server_utils import ProviderCompletionException
+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
 
@@ -125,6 +125,90 @@ 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.
@@ -149,7 +233,7 @@ class AnthropicClient:
         """
         Get a response from Claude or the Mito server that adheres to the AgentResponse format.
         """
-        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,
@@ -166,6 +250,7 @@ class AnthropicClient:
             # 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
@@ -192,7 +277,7 @@ class AnthropicClient:
     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:

mito_ai/app_deploy/handlers.py

@@ -4,10 +4,11 @@
 import os
 import time
 import logging
-from typing import Any, Union, List
+from typing import Any, Union, List, Optional
 import tempfile
-from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
+from mito_ai.path_utils import AbsoluteAppPath, does_app_path_exists, get_absolute_app_path, get_absolute_notebook_dir_path, get_absolute_notebook_path
 from mito_ai.utils.create import initialize_user
+from mito_ai.utils.error_classes import StreamlitDeploymentError
 from mito_ai.utils.version_utils import is_pro
 from mito_ai.utils.websocket_base import BaseWebSocketHandler
 from mito_ai.app_deploy.app_deploy_utils import  add_files_to_zip
@@ -16,11 +17,13 @@ from mito_ai.app_deploy.models import (
     AppDeployError,
     DeployAppRequest,
     ErrorMessage,
-    MessageType
 )
+from mito_ai.completions.models import MessageType
 from mito_ai.logger import get_logger
 from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
+from mito_ai.utils.telemetry_utils import log_streamlit_app_deployment_failure
 import requests
+import traceback
 
 
 class AppDeployHandler(BaseWebSocketHandler):
@@ -67,41 +70,46 @@ class AppDeployHandler(BaseWebSocketHandler):
         self.log.debug("App builder message received: %s", message)
         
         try:
-            # Ensure message is a string before parsing
-            if not isinstance(message, str):
-                raise ValueError("Message must be a string")
-
             parsed_message = self.parse_message(message)
             message_type = parsed_message.get('type')
             
             if message_type == MessageType.DEPLOY_APP.value:
                 # Handle build app request
                 deploy_app_request = DeployAppRequest(**parsed_message)
-                await self._handle_deploy_app(deploy_app_request)
+                response = await self._handle_deploy_app(deploy_app_request)
+                self.reply(response)
             else:
                 self.log.error(f"Unknown message type: {message_type}")
                 error = AppDeployError(
                     error_type="InvalidRequest",
-                    title=f"Unknown message type: {message_type}"
+                    message=f"Unknown message type: {message_type}",
+                    error_code=400
                 )
-                self.reply(ErrorMessage(**error.__dict__))
-                
-        except ValueError as e:
+                raise StreamlitDeploymentError(error)
+
+        except StreamlitDeploymentError as e:
             self.log.error("Invalid app builder request", exc_info=e)
-            error = AppDeployError.from_exception(e)
-            self.reply(ErrorMessage(**error.__dict__))
+            log_streamlit_app_deployment_failure('mito_server_key', MessageType.DEPLOY_APP, e.error.__dict__)
+            self.reply(
+                DeployAppReply(
+                        parent_id=e.message_id,
+                        url="",
+                        error=ErrorMessage(**e.error.__dict__)
+                    )
+            )
         except Exception as e:
             self.log.error("Error handling app builder message", exc_info=e)
             error = AppDeployError.from_exception(
                 e, 
                 hint="An error occurred while building the app. Please check the logs for details."
             )
+            log_streamlit_app_deployment_failure('mito_server_key', MessageType.DEPLOY_APP, error.__dict__)
             self.reply(ErrorMessage(**error.__dict__))
             
         latency_ms = round((time.time() - start) * 1000)
         self.log.info(f"App builder handler processed in {latency_ms} ms.")
     
-    async def _handle_deploy_app(self, message: DeployAppRequest) -> None:
+    async def _handle_deploy_app(self, message: DeployAppRequest) -> DeployAppReply:
         """Handle a build app request.
         
         Args:
@@ -114,19 +122,22 @@ class AppDeployHandler(BaseWebSocketHandler):
         
         if not message_id:
             self.log.error("Missing message_id in request")
-            return
-        
+            error = AppDeployError(
+                error_type="BadRequest",
+                message="Missing message_id in request",
+                error_code=400,
+                message_id=message_id
+            )
+            raise StreamlitDeploymentError(error)
+
         if not notebook_path:
             error = AppDeployError(
                 error_type="InvalidRequest",
-                title="Missing 'notebook_path' parameter"
+                message="Missing 'notebook_path' parameter",
+                error_code=400,
+                message_id=message_id
             )
-            self.reply(DeployAppReply(
-                parent_id=message_id,
-                url="",
-                error=error
-            ))
-            return
+            raise StreamlitDeploymentError(error)
 
         # Validate JWT token if provided
         token_preview = jwt_token[:20] if jwt_token else "No token provided"
@@ -136,55 +147,42 @@ class AppDeployHandler(BaseWebSocketHandler):
             self.log.error("JWT token validation failed")
             error = AppDeployError(
                 error_type="Unauthorized",
-                title="Invalid authentication token",
-                hint="Please sign in again to deploy your app."
+                message="Invalid authentication token",
+                hint="Please sign in again to deploy your app.",
+                error_code=401,
+                message_id=message_id
             )
-            self.reply(DeployAppReply(
-                parent_id=message_id,
-                url="",
-                error=error
-            ))
-            return
+            raise StreamlitDeploymentError(error)
         else:
             self.log.info("JWT token validation successful")
-            
-        try:
-            notebook_path = str(notebook_path) if notebook_path else ""
 
-            app_directory = os.path.dirname(notebook_path)
-            
-            # Check if the app.py file exists
-            app_path = get_app_path(app_directory)
-            if app_path is None:
-                error = AppDeployError(
-                    error_type="AppNotFound",
-                    title="App not found",
-                    hint="Please make sure the app.py file exists in the same directory as the notebook."
-                )
-                self.reply(DeployAppReply(
-                    parent_id=message_id,
-                    url="",
-                    error=error
-                ))
-            
-            # Finally, deploy the app
-            deploy_url = await self._deploy_app(app_directory, files_to_upload, jwt_token)
+        notebook_path = str(notebook_path) if notebook_path else ""
+        absolute_notebook_path = get_absolute_notebook_path(notebook_path)
+        absolute_app_directory = get_absolute_notebook_dir_path(absolute_notebook_path)
+        app_path = get_absolute_app_path(absolute_app_directory)
+
+        # Check if the app.py file exists
+        app_path_exists = does_app_path_exists(app_path)
+        if not app_path_exists:
+            error = AppDeployError(
+                error_type="AppNotFound",
+                message="App not found",
+                hint="Please make sure the app.py file exists in the same directory as the notebook.",
+                error_code=400,
+                message_id=message_id
+            )
+            raise StreamlitDeploymentError(error)
+
+        # Finally, deploy the app
+        deploy_url = await self._deploy_app(app_path, files_to_upload, message_id, jwt_token)
+
+        # Send the response
+        return DeployAppReply(
+            parent_id=message_id,
+            url=deploy_url if deploy_url else ""
+        )
+        
 
-            # Send the response
-            self.reply(DeployAppReply(
-                parent_id=message_id,
-                url=deploy_url
-            ))
-            
-        except Exception as e:
-            self.log.error(f"Error building app: {e}", exc_info=e)
-            error = AppDeployError.from_exception(e)
-            self.reply(DeployAppReply(
-                parent_id=message_id,
-                url="",
-                error=error
-            ))
-    
     def _validate_jwt_token(self, token: str) -> bool:
         """Basic JWT token validation logic.
 
@@ -219,7 +217,7 @@ class AppDeployHandler(BaseWebSocketHandler):
             return False
 
 
-    async def _deploy_app(self, app_path: str, files_to_upload:List[str], jwt_token: str = '') -> str:
+    async def _deploy_app(self, app_path: AbsoluteAppPath, files_to_upload:List[str], message_id: str, jwt_token: str = '') -> Optional[str]:
         """Deploy the app using pre-signed URLs.
         
         Args:
@@ -255,7 +253,7 @@ class AppDeployHandler(BaseWebSocketHandler):
             expected_app_url = url_data['expected_app_url']
             
             self.log.info(f"Received pre-signed URL. App will be available at: {expected_app_url}")
-            
+
             # Step 2: Create a zip file of the app.
             temp_zip_path = None
             try:
@@ -269,7 +267,14 @@ class AppDeployHandler(BaseWebSocketHandler):
                 upload_response = await self._upload_app_to_s3(temp_zip_path, presigned_url)
             except Exception as e:
                 self.log.error(f"Error zipping app: {e}")
-                raise
+                error = AppDeployError(
+                    error_type="ZippingError",
+                    message=f"Error zipping app: {e}",
+                    traceback=traceback.format_exc(),
+                    error_code=500,
+                    message_id=message_id
+                )
+                raise StreamlitDeploymentError(error)
             finally:
                 # Clean up
                 if temp_zip_path is not None:
@@ -278,20 +283,35 @@ class AppDeployHandler(BaseWebSocketHandler):
             self.log.info(f"Upload successful! Status code: {upload_response.status_code}")
             
             self.log.info(f"Deployment initiated. App will be available at: {expected_app_url}")
-            return expected_app_url # type: ignore
+            return str(expected_app_url)
             
         except requests.exceptions.RequestException as e:
             self.log.error(f"Error during API request: {e}")
             if hasattr(e, 'response') and e.response is not None:
                 error_detail = e.response.json()
+                error_message = error_detail.get('error', "")
                 self.log.error(f"Server error details: {error_detail}")
-                if 'error' in error_detail:
-                    raise Exception(error_detail['error'])
-                raise
+            else:
+                error_message = str(e)
+            
+            error = AppDeployError(
+                error_type="APIException",
+                message=str(error_message),
+                traceback=traceback.format_exc(),
+                error_code=500,
+                message_id=message_id
+            )
+            raise StreamlitDeploymentError(error)
         except Exception as e:
             self.log.error(f"Error during deployment: {str(e)}")
-            raise
-        raise RuntimeError("Unexpected error in _deploy_app")
+            error = AppDeployError(
+                error_type="DeploymentException",
+                message=str(e),
+                traceback=traceback.format_exc(),
+                error_code=500,
+                message_id=message_id
+            )
+            raise StreamlitDeploymentError(error)
 
     async def _upload_app_to_s3(self, app_path: str, presigned_url: str) -> requests.Response:
         """Upload the app to S3 using the presigned URL."""

mito_ai/app_deploy/models.py

@@ -2,15 +2,10 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 from dataclasses import dataclass
-from enum import Enum
+import traceback
 from typing import Literal, Optional, List
 
 
-class MessageType(str, Enum):
-    """Types of app deploy messages."""
-    DEPLOY_APP = "deploy-app"
-
-
 @dataclass(frozen=True)
 class AppDeployError:
     """Error information for app deploy operations."""
@@ -19,7 +14,13 @@ class AppDeployError:
     error_type: str
     
     # Error title.
-    title: str
+    message: str
+
+    #ID of parent to resolve response
+    message_id: Optional[str] = "InvalidMessageID"
+
+    # Error code
+    error_code: Optional[int] = 500
     
     # Error traceback.
     traceback: Optional[str] = None
@@ -28,21 +29,24 @@ class AppDeployError:
     hint: Optional[str] = None
     
     @classmethod
-    def from_exception(cls, e: Exception, hint: Optional[str] = None) -> "AppDeployError":
+    def from_exception(cls, e: Exception, hint: Optional[str] = None, error_code: Optional[int] = 500) -> "AppDeployError":
         """Create an error from an exception.
 
         Args:
             e: The exception.
             hint: Optional hint to fix the error.
+            error_code: Optional error code which defaults to 500
 
         Returns:
             The app builder error.
         """
+        tb_str = "".join(traceback.format_exception(type(e), e, e.__traceback__))
         return cls(
             error_type=type(e).__name__,
-            title=str(e),
-            traceback=getattr(e, "__traceback__", None) and str(e.__traceback__),
+            message=str(e),
+            traceback=tb_str,
             hint=hint,
+            error_code=error_code
         )
 
 
@@ -59,7 +63,7 @@ class DeployAppRequest:
     """Request to deploy an app."""
     
     # Request type.
-    type: Literal["deploy-app"]
+    type: Literal["deploy_app"]
     
     # Message ID.
     message_id: str
@@ -88,4 +92,4 @@ class DeployAppReply:
     error: Optional[AppDeployError] = None
     
     # Type of reply.
-    type: Literal["deploy-app"] = "deploy-app" 
+    type: Literal["deploy_app"] = "deploy_app"

mito_ai/chat_history/handlers.py

@@ -0,0 +1,63 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import tornado
+from typing import List, Any
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.completions.message_history import GlobalMessageHistory
+from mito_ai.completions.models import ChatThreadMetadata
+
+
+class ChatHistoryHandler(APIHandler):
+    """
+    Endpoints for working with chat history threads.
+    """
+
+    def initialize(self, message_history: GlobalMessageHistory) -> None:
+        """Initialize the handler with the global message history instance."""
+        super().initialize()
+        self._message_history = message_history
+
+    @tornado.web.authenticated
+    def get(self, *args: Any, **kwargs: Any) -> None:
+        """Get all chat history threads or a specific thread by ID."""
+        try:
+            # Check if a specific thread ID is provided in the URL
+            thread_id = kwargs.get("thread_id")
+
+            if thread_id:
+                # Get specific thread
+                if thread_id in self._message_history._chat_threads:
+                    thread = self._message_history._chat_threads[thread_id]
+                    thread_data = {
+                        "thread_id": thread.thread_id,
+                        "name": thread.name,
+                        "creation_ts": thread.creation_ts,
+                        "last_interaction_ts": thread.last_interaction_ts,
+                        "display_history": thread.display_history,
+                        "ai_optimized_history": thread.ai_optimized_history,
+                    }
+                    self.finish(thread_data)
+                else:
+                    self.set_status(404)
+                    self.finish({"error": f"Thread with ID {thread_id} not found"})
+            else:
+                # Get all threads
+                threads: List[ChatThreadMetadata] = self._message_history.get_threads()
+
+                # Convert to dict format for JSON serialization
+                threads_data = [
+                    {
+                        "thread_id": thread.thread_id,
+                        "name": thread.name,
+                        "creation_ts": thread.creation_ts,
+                        "last_interaction_ts": thread.last_interaction_ts,
+                    }
+                    for thread in threads
+                ]
+
+                self.finish({"threads": threads_data})
+
+        except Exception as e:
+            self.set_status(500)
+            self.finish({"error": str(e)})

mito_ai/chat_history/urls.py

@@ -0,0 +1,32 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple, Any
+from jupyter_server.utils import url_path_join
+from mito_ai.chat_history.handlers import ChatHistoryHandler
+from mito_ai.completions.message_history import GlobalMessageHistory
+
+
+def get_chat_history_urls(base_url: str, message_history: GlobalMessageHistory) -> List[Tuple[str, Any, dict]]:
+    """Get all chat history related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+        message_history: The global message history instance
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai/chat-history"
+    return [
+        (
+            url_path_join(BASE_URL, "threads"),
+            ChatHistoryHandler,
+            {"message_history": message_history},
+        ),
+        (
+            url_path_join(BASE_URL, "threads", "(?P<thread_id>[^/]+)"),
+            ChatHistoryHandler,
+            {"message_history": message_history},
+        ),
+    ]