mito-ai 0.1.50__py3-none-any.whl

mito_ai/__init__.py

@@ -0,0 +1,114 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+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
+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:
+    # Fallback when using the package in dev mode without installing in editable mode with pip. It is highly recommended to install
+    # the package from a stable release or in editable mode: https://pip.pypa.io/en/stable/topics/local-project-installs/#editable-installs
+    import warnings
+    
+    warnings.warn("Importing 'mito_ai' outside a proper installation.")
+    __version__ = "dev"
+
+def _jupyter_labextension_paths() -> List[Dict[str, str]]:
+    return [{"src": "labextension", "dest": "mito_ai"}]
+
+
+def _jupyter_server_extension_points() -> List[Dict[str, str]]:
+    """
+    Returns a list of dictionaries with metadata describing
+    where to find the `_load_jupyter_server_extension` function.
+    """
+    return [{"module": "mito_ai"}]
+
+
+# Jupyter Server is the backend used by JupyterLab. A sever extension lets
+# us add new API's to the backend, so we can do some processing that we don't
+# want to exist in the users's javascript.
+# For a further explanation of the Jupyter architecture watch the first 35 minutes
+# of this video: https://www.youtube.com/watch?v=9_-siU-_XoI
+def _load_jupyter_server_extension(server_app) -> None: # type: ignore
+    host_pattern = ".*$"
+    web_app = server_app.web_app
+    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, "message_history": global_message_history},
+        ),
+        (
+            url_path_join(base_url, "mito-ai", "app-deploy"),
+            AppDeployHandler,
+            {}
+        ),
+        (
+            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))  # 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

@@ -0,0 +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.50'

mito_ai/anthropic_client.py

@@ -0,0 +1,334 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import anthropic
+from typing import Dict, Any, Optional, Tuple, Union, Callable, List, cast
+
+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
+
+# Max tokens is a required parameter for the Anthropic API. 
+# We set it to a high number so that we can edit large code cells
+# 8192 is the maximum allowed number of output tokens for claude-3-5-haiku-20241022
+MAX_TOKENS = 8_000
+
+def extract_and_parse_anthropic_json_response(response: Message) -> Union[object, Any]:
+    """
+    Extracts and parses the JSON response from the Claude API.
+    """
+    try:
+        # Check for tool use in the response
+        for content_block in response.content:
+            if content_block.type == "tool_use" and content_block.name == "agent_response":
+                result = content_block.input
+                return result
+
+        # If no tool use was found, try to parse the text response
+        text_response = None
+        for content_block in response.content:
+            if content_block.type == "text":
+                text_response = content_block.text
+                break
+
+        if text_response:
+            # Try to extract JSON from the text response
+            import re
+            json_pattern = r'(\{.*\})'
+            match = re.search(json_pattern, text_response, re.DOTALL)
+            if match:
+                try:
+                    json_response = json.loads(match.group(0))
+                    return json_response
+                except json.JSONDecodeError:
+                    pass
+
+        raise Exception("No valid AgentResponse format found in the response")
+    except Exception as e:
+        raise Exception(f"Failed to parse response: {e}")
+
+
+def get_anthropic_system_prompt_and_messages(messages: List[ChatCompletionMessageParam]) -> Tuple[
+    Union[str, anthropic.Omit], List[MessageParam]]:
+    """
+    Convert a list of OpenAI messages to a list of Anthropic messages.
+    """
+    
+    system_prompt: Union[str, anthropic.Omit] = anthropic.Omit()
+    anthropic_messages: List[MessageParam] = []
+
+    for message in messages:
+        if 'content' not in message:
+            continue
+
+        # We assume that the conversation only has one system message.
+        # Or if there are multiple, we take the last one.
+        if message['role'] == 'system':
+            system_prompt = str(message['content'])
+
+        # Construct the messages for the user and assistant in Anthropic format.
+        if message['role'] == 'user':
+            content = message['content']
+
+            # Handle mixed content (text + images)
+            if isinstance(content, list):
+                anthropic_content = []
+
+                for item in content:
+                    if isinstance(item, dict):
+                        item_dict = cast(Dict[str, Any], item)
+                        if item_dict.get('type') == 'text':
+                            # Add text content
+                            text_content = item_dict.get('text', '')
+                            anthropic_content.append({
+                                "type": "text",
+                                "text": text_content
+                            })
+                        elif item_dict.get('type') == 'image_url':
+                            # Convert OpenAI image format to Anthropic format
+                            image_url_obj = item_dict.get('image_url', {})
+                            if isinstance(image_url_obj, dict):
+                                image_url = image_url_obj.get('url', '')
+                            else:
+                                image_url = str(image_url_obj)
+
+                            # Extract media type and base64 data
+                            if image_url.startswith('data:'):
+                                # Format: data:image/png;base64,<base64_data>
+                                header, base64_data = image_url.split(',', 1)
+                                media_type = header.split(';')[0].split(':')[1]  # Extract image/png or image/jpeg
+                            else:
+                                # If it's not a data URL, assume it's direct base64 and default to image/png
+                                media_type = "image/png"
+                                base64_data = image_url
+
+                            anthropic_content.append({
+                                "type": "image",
+                                "source": {
+                                    "type": "base64",
+                                    "media_type": media_type,
+                                    "data": base64_data
+                                }
+                            })
+
+                anthropic_messages.append(MessageParam(role='user', content=cast(Any, anthropic_content)))
+            else:
+                # Handle simple text content
+                anthropic_messages.append(MessageParam(role='user', content=str(content)))
+
+        elif message['role'] == 'assistant':
+            anthropic_messages.append(MessageParam(role='assistant', content=str(message['content'])))
+
+    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], timeout: int = 30, max_retries: int = 1):
+        self.api_key = api_key
+        self.timeout = timeout
+        self.max_retries = max_retries
+        self.client: Optional[anthropic.Anthropic]
+        if api_key:
+            self.client = anthropic.Anthropic(api_key=api_key)
+        else:
+            self.client = None
+
+    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.
+        """
+        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)
+            
+            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 ""
+        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_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_with_caching(messages)
+            accumulated_response = ""
+
+            if self.api_key:
+                assert self.client is not None
+                stream = self.client.messages.create(
+                    model=model,
+                    max_tokens=MAX_TOKENS,
+                    temperature=0,
+                    system=anthropic_system_prompt,
+                    messages=anthropic_messages,
+                    stream=True
+                )
+
+                for chunk in stream:
+                    if chunk.type == "content_block_delta" and chunk.delta.type == "text_delta":
+                        content = chunk.delta.text
+                        accumulated_response += content
+
+                        is_finished = chunk.type == "message_stop"
+
+                        reply_fn(CompletionStreamChunk(
+                            parent_id=message_id,
+                            chunk=CompletionItem(
+                                content=content,
+                                isIncomplete=not is_finished,
+                                token=message_id,
+                            ),
+                            done=is_finished,
+                        ))
+
+            else:
+                async for stram_chunk in stream_anthropic_completion_from_mito_server(
+                    model=model,
+                    max_tokens=MAX_TOKENS,
+                    temperature=0,
+                    system=anthropic_system_prompt,
+                    messages=anthropic_messages,
+                    stream=True,
+                    message_type=message_type,
+                    reply_fn=reply_fn,
+                    message_id=message_id
+                ):
+                    accumulated_response += stram_chunk
+
+            return accumulated_response
+
+        except anthropic.RateLimitError:
+            raise Exception("Rate limit exceeded. Please try again later or reduce your request frequency.")
+
+        except Exception as e:
+            print(f"Error streaming content: {str(e)}")
+            raise e
+
+

mito_ai/app_deploy/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""App builder module for Mito AI."""
+
+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}")