mito-ai 0.1.50__py3-none-any.whl

mito_ai/utils/anthropic_utils.py

@@ -0,0 +1,168 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import anthropic
+from typing import Any, Dict, List, Optional, Union, AsyncGenerator, Tuple, Callable
+from anthropic.types import MessageParam, TextBlockParam, ToolUnionParam
+from mito_ai.utils.mito_server_utils import get_response_from_mito_server, stream_response_from_mito_server
+from mito_ai.utils.provider_utils import does_message_require_fast_model
+from mito_ai.completions.models import AgentResponse, MessageType, ResponseFormatInfo, CompletionReply, CompletionStreamChunk
+from mito_ai.utils.schema import UJ_STATIC_USER_ID, UJ_USER_EMAIL
+from mito_ai.utils.db import get_user_field
+from mito_ai.constants import MITO_ANTHROPIC_URL
+
+__user_email: Optional[str] = None
+__user_id: Optional[str] = None
+
+ANTHROPIC_TIMEOUT = 60
+max_retries = 1
+
+FAST_ANTHROPIC_MODEL = "claude-3-5-haiku-latest"
+
+def _prepare_anthropic_request_data_and_headers(
+    model: Union[str, None],
+    max_tokens: int,
+    temperature: float,
+    system: Union[str, List[TextBlockParam], anthropic.Omit],
+    messages: List[MessageParam],
+    message_type: MessageType,
+    tools: Optional[List[ToolUnionParam]],
+    tool_choice: Optional[dict],
+    stream: Optional[bool]
+) -> Tuple[Dict[str, Any], Dict[str, str]]:
+    
+    global __user_email, __user_id
+    if __user_email is None:
+        __user_email = get_user_field(UJ_USER_EMAIL)
+    if __user_id is None:
+        __user_id = get_user_field(UJ_STATIC_USER_ID)
+    # Build the inner data dict (excluding timeout, max_retries, email, user_id)
+    inner_data: Dict[str, Any] = {
+        "model": model,
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "messages": messages,
+        "betas": ["context-1m-2025-08-07"]
+    }
+    # Add system to inner_data only if it is not anthropic.Omit
+    if not isinstance(system, anthropic.Omit):
+        inner_data["system"] = system
+    if tools:
+        inner_data["tools"] = tools
+    if tool_choice:
+        inner_data["tool_choice"] = tool_choice
+    if stream:
+        inner_data["stream"] = stream
+    # Compose the outer data dict
+    data = {
+        "timeout": ANTHROPIC_TIMEOUT,
+        "max_retries": max_retries,
+        "email": __user_email,
+        "user_id": __user_id,
+        "data": inner_data
+    }
+    headers = {"Content-Type": "application/json"}
+    return data, headers
+
+async def get_anthropic_completion_from_mito_server(
+    model: Union[str, None],
+    max_tokens: int,
+    temperature: float,
+    system: Union[str, anthropic.Omit],
+    messages: List[MessageParam],
+    tools: Optional[List[ToolUnionParam]],
+    tool_choice: Optional[dict],
+    message_type: MessageType
+) -> str:
+    data, headers = _prepare_anthropic_request_data_and_headers(
+        model, max_tokens, temperature, system, messages, message_type, tools, tool_choice, None
+    )
+    
+    return await get_response_from_mito_server(
+        MITO_ANTHROPIC_URL, 
+        headers, 
+        data, 
+        ANTHROPIC_TIMEOUT, 
+        max_retries, 
+        message_type, 
+        provider_name="Claude"
+    )
+
+async def stream_anthropic_completion_from_mito_server(
+    model: Union[str, None],
+    max_tokens: int,
+    temperature: float,
+    system: Union[str, List[TextBlockParam], anthropic.Omit],
+    messages: List[MessageParam],
+    stream: bool,
+    message_type: MessageType,
+    reply_fn: Optional[Callable[[Union[CompletionReply, CompletionStreamChunk]], None]] = None,
+    message_id: Optional[str] = None,
+) -> AsyncGenerator[str, None]:
+    data, headers = _prepare_anthropic_request_data_and_headers(
+        model, max_tokens, temperature, system, messages, message_type, None, None, stream
+    )
+    
+    # Use the unified streaming function
+    # If the reply_fn and message_id are empty, this function still handles those requests. This is particularly needed for the streamlit dashboard functionality
+    actual_reply_fn = reply_fn if reply_fn is not None else (lambda x: None)
+    actual_message_id = message_id if message_id is not None else ""
+    async for chunk in stream_response_from_mito_server(
+        url=MITO_ANTHROPIC_URL,
+        headers=headers,
+        data=data,
+        timeout=ANTHROPIC_TIMEOUT,
+        max_retries=max_retries,
+        message_type=message_type,
+        reply_fn=actual_reply_fn,
+        message_id=actual_message_id,
+        chunk_processor=None,
+        provider_name="Claude",
+    ):
+        yield chunk
+    
+    
+def get_anthropic_completion_function_params(
+    message_type: MessageType,
+    model: str,
+    messages: List[MessageParam],
+    max_tokens: int,
+    system: Union[str, List[TextBlockParam], anthropic.Omit],
+    temperature: float = 0.0,
+    tools: Optional[List[ToolUnionParam]] = None,
+    tool_choice: Optional[dict] = None,
+    stream: Optional[bool] = None,
+    response_format_info: Optional[ResponseFormatInfo] = None,
+) -> Dict[str, Any]:    
+    """
+    Build the provider_data dict for Anthropic completions, mirroring the OpenAI approach.
+    Only includes fields needed for the Anthropic API.
+    """
+    
+    message_requires_fast_model = does_message_require_fast_model(message_type)
+    model = FAST_ANTHROPIC_MODEL if message_requires_fast_model else model
+    
+    provider_data = {
+        "model": model,
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "messages": messages,
+        "system": system,
+    }
+    if tools:
+        provider_data["tools"] = tools
+    if response_format_info and response_format_info.name == "agent_response":
+        provider_data["tools"] = [{
+            "name": "agent_response",
+            "description": "Output a structured response following the AgentResponse format",
+            "input_schema": AgentResponse.model_json_schema()
+        }]
+        provider_data["tool_choice"] = {"type": "tool", "name": "agent_response"}
+    
+    
+    if tool_choice:
+        provider_data["tool_choice"] = tool_choice
+    if stream is not None:
+        provider_data["stream"] = stream
+    # Optionally handle response_format_info if Anthropic supports it in the future
+    return provider_data

mito_ai/utils/create.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GPL License.
+
+"""
+Utilities for creating and initializing the user.json
+file with the current schema
+"""
+
+import json
+import os
+from typing import Literal, Optional
+
+from mito_ai.utils.db import (MITO_FOLDER, USER_JSON_PATH, set_user_field)
+from mito_ai.utils.schema import (GITHUB_ACTION_EMAIL, GITHUB_ACTION_ID,
+                                    UJ_STATIC_USER_ID, UJ_USER_EMAIL,
+                                    USER_JSON_DEFAULT)
+from mito_ai.utils.telemetry_utils import identify
+from mito_ai.utils.utils import is_running_test
+
+
+def is_user_json_exists_and_valid_json() -> bool:
+    """
+    Helper function that determines if the current user.json both
+    exists and is valid json
+    """
+    if not os.path.exists(USER_JSON_PATH):
+        return False
+
+    try:
+        with open(USER_JSON_PATH, 'r') as f:
+            json.loads(f.read())
+        return True
+    except:
+        return False
+
+def get_temp_user_id() -> Optional[str]:
+    """
+    Looks for a temporary user ID, generated by the desktop app.
+    """
+    temp_user_id_path = os.path.join(MITO_FOLDER, 'temp_user_id.txt')
+
+    if os.path.exists(temp_user_id_path):
+        with open(temp_user_id_path, 'r') as f:
+            return f.read()
+
+    return None
+
+def try_create_user_json_file() -> None:
+    
+    # Create the mito folder if it does not exist
+    if not os.path.exists(MITO_FOLDER):
+        os.mkdir(MITO_FOLDER)
+
+    # We create a user.json file if it does not exist, or if it
+    # is invalid (e.g. it is not parseable JSON).
+    if not is_user_json_exists_and_valid_json():
+        # First, we write an empty default object
+        with open(USER_JSON_PATH, 'w+') as f:
+            f.write(json.dumps(USER_JSON_DEFAULT))
+
+        # Next, look for a temp user id
+        temp_user_id = get_temp_user_id()
+        if temp_user_id:
+            set_user_field(UJ_STATIC_USER_ID, temp_user_id)
+
+        # Finally, we take special care to put all the testing/CI environments 
+        # (e.g. Github actions) under one ID and email
+        if is_running_test():
+            set_user_field(UJ_STATIC_USER_ID, GITHUB_ACTION_ID)
+            set_user_field(UJ_USER_EMAIL, GITHUB_ACTION_EMAIL)
+
+
+def initialize_user() -> None:
+    """
+    Internal helper function that gets called whenever a ai completion is requested.
+
+    It:
+    1. Creates the user.json if it does not exist (though it usually does, from the installer)
+    2. Identifies the user
+    """
+    # Try to create the user.json file, if it does not already exist
+    try_create_user_json_file()
+
+    # Identify the user 
+    identify()
+
+    # Note, its possible that a user has a previous version of the user.json file if they 
+    # downloaded Mito >1 year ago and have not yet upgraded. In this case, we would like to 
+    # upgrade the user.json to the newest versio. But for simplicity, let's not 
+    # try to upgrade it here. We will either move away from user.json all together
+    # or unify the utilities for mitosheet and mitoai.

mito_ai/utils/db.py

@@ -0,0 +1,74 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""
+Helpers for accessing the user.json file
+"""
+import os
+import json
+from typing import Any, Optional, Final
+from mito_ai.utils.schema import (
+    MITO_FOLDER, 
+    UJ_AI_MITO_API_NUM_USAGES, 
+    UJ_MITO_AI_FIRST_USAGE_DATE, 
+    UJ_MITO_AI_LAST_RESET_DATE,
+    UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES
+)
+
+# The path of the user.json file
+USER_JSON_PATH: Final[str] = os.path.join(MITO_FOLDER, 'user.json')
+
+
+def get_user_field(field: str) -> Optional[Any]:
+    """
+    Returns the value stored at field in the user.json file,
+    but may read a different file if it passed
+    """
+    try:
+        with open(USER_JSON_PATH) as f:
+            return json.load(f)[field]
+    except: 
+        return None
+    
+def set_user_field(field: str, value: Any) -> None:
+    """
+    Updates the value of a specific field in user.json
+    """
+    with open(USER_JSON_PATH, 'r') as user_file_old:
+        old_user_json = json.load(user_file_old)
+        old_user_json[field] = value
+        with open(USER_JSON_PATH, 'w+') as f:
+            f.write(json.dumps(old_user_json))
+            
+def get_chat_completion_count() -> int:
+    """
+    Returns the number of AI completions the user has made in the current 30-day period.
+    
+    Using this helper function lets us mock the number of completions in tests.
+    """
+    return get_user_field(UJ_AI_MITO_API_NUM_USAGES) or 0
+
+def get_autocomplete_count() -> int:
+    """
+    Returns the number of autocomplete completions the user has made in the current 30-day period.
+    
+    Using this helper function lets us mock the number of autocomplete completions in tests.
+    """
+    return get_user_field(UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES) or 0
+
+def get_first_completion_date() -> Optional[str]:
+    """
+    Returns the date of the user's last completion in YYYY-MM-DD format.
+    
+    Using this helper function lets us mock the date of the first completion in tests.
+    """
+    return get_user_field(UJ_MITO_AI_FIRST_USAGE_DATE) or None
+
+def get_last_reset_date() -> Optional[str]:
+    """
+    Returns the date of the user's last reset in YYYY-MM-DD format.
+    Both chat completions and autocomplete share the same reset date.
+    
+    Using this helper function lets us mock the date of the last reset in tests.
+    """
+    return get_user_field(UJ_MITO_AI_LAST_RESET_DATE) or None

mito_ai/utils/error_classes.py

@@ -0,0 +1,42 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.app_deploy.models import AppDeployError
+
+class MitoAppError(Exception):
+    """Exception raised for custom error in the application."""
+
+    def __init__(self, message: str, error_code: int) -> None:
+        super().__init__(message)
+        self.message = message
+        self.error_code = error_code
+
+class StreamlitPreviewError(MitoAppError):
+    def __str__(self) -> str:
+        return f"[PreviewError]: {self.message} (Error Code: {self.error_code})"
+
+class StreamlitConversionError(MitoAppError):
+    def __str__(self) -> str:
+        return f"[ConversionError]: {self.message} (Error Code: {self.error_code})"
+
+
+class StreamlitDeploymentError(MitoAppError):
+    """Raised when a deployment operation fails."""
+
+    def __init__(self, error: AppDeployError):
+        self.error = error
+        self.error_type = error.error_type
+        self.message_id = getattr(error, "message_id", "ErrorMessageID")
+        self.error_code = getattr(error, "error_code", 500)
+        self.hint = getattr(error, "hint", "")
+        self.traceback = getattr(error, "traceback", "")
+        self.error_type = getattr(error, "error_type", "Error")
+        self.message = error.message
+        print(f"self_message: {self.message}")
+        super().__init__(self.message, self.error_code)
+
+    def __str__(self) -> str:
+        base = f"[DeploymentError]: {self.message} (Error Code: {self.error_code})"
+        if self.hint:
+            base += f"\nHint: {self.hint}"
+        return base

mito_ai/utils/gemini_utils.py

@@ -0,0 +1,133 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import asyncio
+import json
+import time
+from typing import Any, Dict, List, Optional, Callable, Union, AsyncGenerator, Tuple
+from mito_ai.utils.mito_server_utils import get_response_from_mito_server, stream_response_from_mito_server
+from mito_ai.completions.models import AgentResponse, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
+from mito_ai.constants import MITO_GEMINI_URL
+from mito_ai.utils.provider_utils import does_message_require_fast_model
+from mito_ai.utils.utils import _create_http_client
+
+timeout = 30
+max_retries = 1
+
+FAST_GEMINI_MODEL = "gemini-2.0-flash-lite"
+
+def _prepare_gemini_request_data_and_headers(
+    model: str,
+    contents: List[Dict[str, Any]],
+    message_type: MessageType,
+    config: Optional[Dict[str, Any]] = None,
+    response_format_info: Optional[Any] = None,
+    stream: bool = False
+) -> Tuple[Dict[str, Any], Dict[str, str]]:
+    
+    inner_data: Dict[str, Any] = {
+        "model": model,
+        "contents": contents,
+        "message_type": message_type.value if hasattr(message_type, 'value') else str(message_type),
+    }
+    
+    if response_format_info:
+        # Ensure the format is a string, not a class
+        format_value = getattr(response_format_info, 'format', None)
+        if isinstance(format_value, type):
+            format_value = format_value.__name__
+        inner_data["response_format_info"] = json.dumps({
+            "name": getattr(response_format_info, 'name', None),
+            "format": format_value
+        })
+        
+    if stream:
+        inner_data["stream"] = True
+        
+    if config:
+        # Ensure config is serializable
+        inner_data["config"] = json.loads(json.dumps(config))
+        
+    data = {
+        "timeout": timeout,
+        "max_retries": max_retries,
+        "data": inner_data
+    }
+    
+    headers = {"Content-Type": "application/json"}
+    return data, headers
+
+async def get_gemini_completion_from_mito_server(
+    model: str,
+    contents: List[Dict[str, Any]],
+    message_type: MessageType,
+    config: Optional[Dict[str, Any]] = None,
+    response_format_info: Optional[Any] = None
+) -> str:
+    data, headers = _prepare_gemini_request_data_and_headers(model, contents, message_type, config, response_format_info, stream=False)
+    return await get_response_from_mito_server(
+        MITO_GEMINI_URL, 
+        headers, 
+        data, 
+        timeout, 
+        max_retries, 
+        message_type, 
+        provider_name="Gemini"
+    )    
+
+async def stream_gemini_completion_from_mito_server(
+    model: str,
+    contents: List[Dict[str, Any]],
+    message_type: MessageType,
+    message_id: str,
+    reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]
+) -> AsyncGenerator[str, None]:
+    data, headers = _prepare_gemini_request_data_and_headers(model, contents, message_type, stream=True)
+    
+    # Define chunk processor for Gemini's special processing
+    def gemini_chunk_processor(chunk: str) -> str:
+        clean_chunk = chunk.strip('"')
+        return clean_chunk.encode().decode('unicode_escape')
+    
+    # Use the unified streaming function with Gemini's chunk processor
+    async for chunk in stream_response_from_mito_server(
+        url=MITO_GEMINI_URL,
+        headers=headers,
+        data=data,
+        timeout=timeout,
+        max_retries=max_retries,
+        message_type=message_type,
+        reply_fn=reply_fn,
+        message_id=message_id,
+        chunk_processor=gemini_chunk_processor,
+        provider_name="Gemini",
+    ):
+        yield chunk
+
+def get_gemini_completion_function_params(
+    message_type: MessageType,
+    model: str,
+    contents: list[dict[str, Any]],
+    response_format_info: Optional[Any] = None,
+) -> Dict[str, Any]:
+    """
+    Build the provider_data dict for Gemini completions, mirroring the OpenAI/Anthropic approach.
+    Only includes fields needed for the Gemini API.
+    """
+    message_requires_fast_model = does_message_require_fast_model(message_type)
+    model = FAST_GEMINI_MODEL if message_requires_fast_model else model
+    
+    provider_data: Dict[str, Any] = {
+        "model": model,
+        "contents": contents,
+        "message_type": message_type.value if hasattr(message_type, 'value') else str(message_type),
+    }
+        
+    # Configure response format if provided
+    if response_format_info:
+        provider_data["config"] = {
+            "response_mime_type": "application/json",
+            "response_schema": AgentResponse.model_json_schema()
+        }
+        
+    return provider_data

mito_ai/utils/message_history_utils.py

@@ -0,0 +1,87 @@
+# Copyright (c) Saga Inc.
+# 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 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
+)
+
+
+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.
+
+    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
+    ]
+    
+    for heading in section_headings:
+        content = re.sub(
+            f"{re.escape(heading)}\n(?:.+\n)+",
+            f"{heading} {CONTENT_REMOVED_PLACEHOLDER}\n",
+            content,
+        )
+
+    return content
+
+
+def trim_old_messages(messages: List[ChatCompletionMessageParam]) -> List[ChatCompletionMessageParam]:
+    """
+    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.
+    """
+    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")
+        
+        is_user_message = messages[i].get("role") == "user"
+        if not is_user_message: 
+            continue
+        
+        content = messages[i].get("content")
+        
+        if content is None:
+            continue
+        
+        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)
+        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
+            text_content = ""
+            for section in content:
+                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)        
+
+
+    return messages