mito-ai 0.1.50__py3-none-any.whl

mito_ai/utils/server_limits.py

@@ -0,0 +1,152 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Final
+from mito_ai.utils.db import (
+    get_chat_completion_count, 
+    get_autocomplete_count,
+    get_first_completion_date, 
+    get_last_reset_date,
+    get_user_field, 
+    set_user_field
+)
+from mito_ai.utils.schema import (
+    UJ_MITO_AI_FIRST_USAGE_DATE, 
+    UJ_AI_MITO_API_NUM_USAGES, 
+    UJ_MITO_AI_LAST_RESET_DATE,
+    UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES
+)
+from mito_ai.utils.telemetry_utils import (MITO_SERVER_FREE_TIER_LIMIT_REACHED, log)
+from mito_ai.utils.version_utils import is_pro
+from mito_ai.completions.models import MessageType
+from datetime import datetime, timedelta
+
+"""
+IT IS EXPRESSLY AGAINST THE SOFTWARE LICENSE TO MODIFY, BYPASS, OR OTHERWISE
+CIRCUMVENT THE QUOTA LIMITS OF THE FREE TIER. We want to provide users with a 
+free tier, but running AI models is expensive, so we need to limit the usage 
+or we will no longer be able to provide this free tier.
+"""
+# Monthly chat completions limit for free tier users
+OS_MONTHLY_AI_COMPLETIONS_LIMIT: Final[int] = 150
+
+# Monthly autocomplete limit for free tier users
+OS_MONTHLY_AUTOCOMPLETE_LIMIT: Final[int] = 5000
+
+def check_mito_server_quota(message_type: MessageType) -> None:
+    """
+    Checks if the user has exceeded their Mito server quota. Pro users have no limits.
+    Raises PermissionError if the user has exceeded their quota.
+    
+    IT IS EXPRESSLY AGAINST THE SOFTWARE LICENSE TO MODIFY, BYPASS, OR OTHERWISE
+    CIRCUMVENT THE QUOTA LIMITS OF THE FREE TIER. We want to provide users with a 
+    free tier, but running AI models is expensive, so we need to limit the usage 
+    or we will no longer be able to provide this free tier.
+    """
+    if is_pro():
+        return
+
+    # Get current date
+    current_date = datetime.now()
+    
+    # Get the date when the last reset occurred
+    last_reset_date_str = get_last_reset_date()
+        
+    # If no last reset date is found, set it to today and reset both counters
+    if last_reset_date_str is None:
+        set_user_field(UJ_MITO_AI_LAST_RESET_DATE, current_date.strftime("%Y-%m-%d"))
+        set_user_field(UJ_AI_MITO_API_NUM_USAGES, 0)
+        set_user_field(UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES, 0)
+        return
+    
+    # Convert the last reset date string to a datetime object
+    last_reset_date = datetime.strptime(last_reset_date_str, "%Y-%m-%d")
+    
+    # Calculate the date one month from the last reset
+    one_month_later = last_reset_date + timedelta(days=30)  # Approximating a month as 30 days
+    
+    # If it's been more than a month since the last reset, reset both counters
+    if current_date >= one_month_later:
+        set_user_field(UJ_AI_MITO_API_NUM_USAGES, 0)
+        set_user_field(UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES, 0)
+        set_user_field(UJ_MITO_AI_LAST_RESET_DATE, current_date.strftime("%Y-%m-%d"))
+        
+    # Check the appropriate limit based on message type
+    if message_type == MessageType.INLINE_COMPLETION:
+        # Check autocomplete limit
+        autocomplete_count = get_autocomplete_count()
+        
+        if autocomplete_count >= OS_MONTHLY_AUTOCOMPLETE_LIMIT:
+            log(MITO_SERVER_FREE_TIER_LIMIT_REACHED)
+            raise PermissionError(MITO_SERVER_FREE_TIER_LIMIT_REACHED)
+    else:
+        # Check chat completion limit
+        completion_count = get_chat_completion_count()
+        
+        if completion_count >= OS_MONTHLY_AI_COMPLETIONS_LIMIT:
+            log(MITO_SERVER_FREE_TIER_LIMIT_REACHED)
+            raise PermissionError(MITO_SERVER_FREE_TIER_LIMIT_REACHED)
+        
+
+def update_mito_server_quota(message_type: MessageType) -> None:
+    """
+    Update the user's quota for the Mito Server.
+    
+    IT IS EXPRESSLY AGAINST THE SOFTWARE LICENSE TO MODIFY, BYPASS, OR OTHERWISE
+    CIRCUMVENT THE QUOTA LIMITS OF THE FREE TIER. We want to provide users with a 
+    free tier, but running AI models is expensive, so we need to limit the usage 
+    or we will no longer be able to provide this free tier.
+    """
+    
+    if message_type == MessageType.CHAT_NAME_GENERATION:
+        # We do not count the Chat Name Generation message type towards the quota
+        return
+    
+    current_date = datetime.now()
+    first_usage_date = get_first_completion_date()
+    last_reset_date_str = get_last_reset_date()
+    
+    if first_usage_date is None:
+        first_usage_date = current_date.strftime("%Y-%m-%d")
+        set_user_field(UJ_MITO_AI_FIRST_USAGE_DATE, first_usage_date)
+    
+    # Initialize the reset date if it doesn't exist
+    if last_reset_date_str is None:
+        last_reset_date_str = current_date.strftime("%Y-%m-%d")
+        set_user_field(UJ_MITO_AI_LAST_RESET_DATE, last_reset_date_str)
+        
+    # Convert the last reset date string to a datetime object
+    last_reset_date = datetime.strptime(last_reset_date_str, "%Y-%m-%d")
+    
+    # Calculate the date one month from the last reset
+    one_month_later = last_reset_date + timedelta(days=30)
+    
+    # If it's been more than a month since the last reset, reset both counters
+    # and make today's date the new reset date
+    if current_date >= one_month_later:
+        set_user_field(UJ_AI_MITO_API_NUM_USAGES, 0)
+        set_user_field(UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES, 0)
+        last_reset_date_str = current_date.strftime("%Y-%m-%d")
+        set_user_field(UJ_MITO_AI_LAST_RESET_DATE, last_reset_date_str)
+    
+    # Update the appropriate usage counter based on message type
+    if message_type == MessageType.INLINE_COMPLETION:
+        # Increment autocomplete count
+        autocomplete_count = get_autocomplete_count()
+        autocomplete_count = 1 if autocomplete_count is None else autocomplete_count + 1
+        
+        try:
+            set_user_field(UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES, autocomplete_count)
+        except Exception as e:
+            raise e
+    else:
+        # Increment chat completion count
+        completion_count = get_chat_completion_count()
+        completion_count = 1 if completion_count is None else completion_count + 1
+        
+        try:
+            set_user_field(UJ_AI_MITO_API_NUM_USAGES, completion_count)
+            set_user_field(UJ_MITO_AI_FIRST_USAGE_DATE, first_usage_date)
+        except Exception as e:
+            raise e
+        

mito_ai/utils/telemetry_utils.py

@@ -0,0 +1,480 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import os
+from typing import Any, Dict, Literal, Optional, List
+from mito_ai.utils.version_utils import MITOSHEET_HELPER_PRIVATE, is_pro
+from mito_ai.utils.schema import UJ_AI_MITO_API_NUM_USAGES, UJ_MITOSHEET_TELEMETRY, UJ_STATIC_USER_ID, UJ_USER_EMAIL, UJ_FEEDBACKS_V2
+from mito_ai.utils.db import get_user_field
+from mito_ai._version import __version__
+from mito_ai.utils.utils import is_running_test
+from mito_ai.completions.models import MessageType
+import analytics
+
+WRITE_KEY = '6I7ptc5wcIGC4WZ0N1t0NXvvAbjRGUgX' 
+analytics.write_key = WRITE_KEY
+
+# If you want, you can optionally choose to print logs
+# helpful for debugging.
+PRINT_LOGS = False
+
+################################# 
+# Mito AI Completion 
+# Constants for logging the success or error of Mito AI
+MITO_AI_COMPLETION_SUCCESS = 'mito_ai_success'
+MITO_AI_COMPLETION_ERROR = 'mito_ai_error'
+MITO_AI_COMPLETION_RETRY = 'mito_ai_retry'
+
+# Params 
+# - logging the type of key 
+KEY_TYPE_PARAM = 'AI_key_type'
+MITO_SERVER_KEY: Literal['mito_server_key'] = 'mito_server_key'
+USER_KEY: Literal['user_key'] = 'user_key'
+# - logging if the user is in dev mode
+IS_DEV_MODE_PARAM = 'is_dev_mode'
+
+# - logging the number of usages of the Mito server
+MITO_SERVER_NUM_USAGES = 'mito_server_num_usages'
+#################################
+
+################################# 
+# Mito Server Free Tier Reached
+MITO_SERVER_FREE_TIER_LIMIT_REACHED = 'mito_server_free_tier_limit_reached'
+#################################
+
+def is_dev_mode() -> bool:
+    """
+    Check if mito-ai is installed in editable/development mode.
+    
+    This function detects editable installs using the modern PEP 660 standard
+    (pip >= 21.3). Works for most development scenarios where developers use
+    `pip install -e .`
+    
+    Returns:
+        bool: True if running in development mode, False otherwise.
+        
+    Limitations:
+        - Requires pip >= 21.3 for reliable detection
+        - Won't detect manual PYTHONPATH manipulation
+        - Won't detect legacy .egg-link installations (very old pip)
+        
+    Note: This is a best-effort detection. For 100% reliability, consider
+    also setting a MITO_DEVELOPER_MODE environment variable.
+    """
+    try:
+        import importlib.metadata
+        import json
+        
+        dist = importlib.metadata.distribution('mito-ai')
+        direct_url_text = dist.read_text('direct_url.json')
+        if direct_url_text:
+            direct_url = json.loads(direct_url_text)
+            direct_url_dir_info = direct_url.get('dir_info', {})
+            editable = direct_url_dir_info.get('editable', None)
+            return False if editable is None else editable
+    except Exception:
+        pass
+    
+    return False
+
+def telemetry_turned_on(key_type: Optional[str] = None) -> bool:
+    """
+    Helper function that tells you if logging is turned on or
+    turned off on the entire Mito instance
+    """
+    # If the user is on the Mito server, then they are sending
+    # us their information already
+    if key_type == 'mito_server_key':
+        return True
+    
+    # If private helper is installed, then we don't log anything
+    if MITOSHEET_HELPER_PRIVATE:
+        return False
+
+    # TODO: Check if the an enterprise user has turned telemetry to true
+
+    # If Mito Pro is on, then don't log anything
+    if is_pro():
+        return False
+
+    telemetry = get_user_field(UJ_MITOSHEET_TELEMETRY)
+    if telemetry is None:
+        return False
+    
+    return bool(telemetry)
+
+def identify(key_type: Optional[str] = None, is_electron: Optional[bool] = None) -> None:
+    """
+    Helper function for identifying a user. We just take
+    their python version, mito version, and email.
+    """
+    if not telemetry_turned_on(key_type):
+        return
+
+    static_user_id = get_user_field(UJ_STATIC_USER_ID)
+    user_email = get_user_field(UJ_USER_EMAIL)
+    feedbacks_v2 = get_user_field(UJ_FEEDBACKS_V2)
+
+    params = {
+        'version_mitoai': __version__,
+        'email': user_email,
+        'is_pro': is_pro(),
+        'is_jupyterhub': 'True' if 'JUPYTERHUB_API_URL' in os.environ else 'False',
+        'is_mito_jupyterhub': 'True' if os.getenv('MITO_JUPYTERHUB') is not None else 'False',
+        IS_DEV_MODE_PARAM: is_dev_mode(),
+        UJ_FEEDBACKS_V2: feedbacks_v2
+    }
+    
+    if is_electron is not None:
+        # Only update field when we get info from the client. Don't default to False.
+        # becase we are only sending this info to the first completion_handler identify call.
+        params['is_mito_desktop'] = is_electron
+
+    if not is_running_test():
+        # TODO: If the user is in JupyterLite, we need to do some extra work.
+        # You can see this in the mitosheet package. 
+        try:
+            analytics.identify(static_user_id, params)
+        except Exception as e:
+            pass
+
+def chunk_param(param: str, param_name: str, chunk_size: int=250) -> Dict[str, str]:
+    """
+    Split a string into chunks of 250 characters.
+    
+    Args:
+        param: The string to be chunked
+        param_name: The name of the param to be chunked (used as prefix for the chunked keys)
+        chunk_size: The number of characters in each chunk
+
+    Returns:
+        dict: A dictionary with keys 'response_part_1', 'response_part_2', etc.
+    """
+
+    chunks = {}
+
+    if not param:
+        return {}
+    
+    num_chunks = (len(param) + chunk_size - 1) // chunk_size
+
+    for i in range(num_chunks):
+        start = i * chunk_size
+        end = min(start + chunk_size, len(param))
+        chunks[f'{param_name}_part_{i + 1}'] = param[start:end]
+
+    return chunks
+
+def log(
+        log_event: str, 
+        params: Optional[Dict[str, Any]]=None, 
+        error: Optional[BaseException]=None, 
+        key_type: Optional[Literal['mito_server_key', 'user_key']] = None,
+        thread_id: Optional[str] = None
+    ) -> None:
+    """
+    This function is the entry point for all logging. 
+
+    If telemetry is not turned off and we are not running tests,
+    we log the ai event
+    """
+    final_params: Dict[str, Any] = params or {}
+    
+    # Then, make sure to add the user email
+    final_params['email'] = get_user_field(UJ_USER_EMAIL)
+
+    # Add the error if it exists
+    if error is not None:
+        final_params['error'] = str(error)
+
+    if thread_id is not None:
+        final_params['thread_id'] = thread_id
+        
+    # Process parameters that need chunking
+    params_to_remove = []
+    params_to_add = {}
+    
+    for param_name, param_value in final_params.items():
+        if isinstance(param_value, str) and len(param_value) > 250:
+            # Mark for removal
+            params_to_remove.append(param_name)
+            # Get chunked parameters
+            chunked_params = chunk_param(param_value, param_name)
+            params_to_add.update(chunked_params)
+    
+    # Apply the changes
+    for param_name in params_to_remove:
+        del final_params[param_name]
+    final_params.update(params_to_add)
+            
+    # Finally, do the acutal logging. We do not log anything when tests are
+    # running, or if telemetry is turned off
+    if not is_running_test() and telemetry_turned_on(key_type):
+        # TODO: If the user is in JupyterLite, we need to do some extra work.
+        # You can see this in the mitosheet package. 
+        
+        try:
+            analytics.track(
+                get_user_field(UJ_STATIC_USER_ID), 
+                log_event, 
+                final_params
+            )
+        except Exception as e:
+            pass
+        
+
+    # If we want to print the logs for debugging reasons, then we print them as well
+    if PRINT_LOGS:
+        print(
+            log_event, 
+            final_params
+        )
+
+    # TODO: Eventually we want to hook this up to the mito log uploader 
+    # so enterprises can log usage if they want to.
+
+def log_ai_completion_success(
+    key_type: Literal['mito_server_key', 'user_key'],
+    message_type: MessageType,
+    last_message_content: str,
+    user_input: str,
+    response: Dict[str, Any],
+    thread_id: str,
+    model: str
+) -> None:
+    """
+    Logs AI completion success based on the input location.
+
+    Args:
+        key_type: The type of key that was used to get the AI completion
+        message_type: The type of message that was sent to the AI
+        last_message_content: The last message sent to the AI
+        user_input: The user input that was sent to the AI
+        response: The response received from the AI
+        thread_id: The thread ID for the conversation
+        model: The model that was used to get the AI completion
+    """
+
+    # Params that every log has
+    base_params = {
+        KEY_TYPE_PARAM: str(key_type),
+        IS_DEV_MODE_PARAM: is_dev_mode(),
+        'model': model,
+    }
+
+    try:
+        code_cell_input = json.dumps(
+            last_message_content.split("Code in the active code cell:")[-1]
+            .strip()
+            .split("```python")[1]
+            .strip()
+            .split("```")[0]
+        )
+        
+        num_usages = get_user_field(UJ_AI_MITO_API_NUM_USAGES)
+    except:
+        # Most user prompts will have an associated code cell that serves as the input context.
+        # However, types like agent:planning (RIP) do not have a code cell input.
+        code_cell_input = ""
+        num_usages = -1
+
+    # Chunk certain params to work around mixpanel's 255 character limit
+    code_cell_input_chunks = chunk_param(code_cell_input, "code_cell_input")
+    response_chunks = chunk_param(response["completion"], "response")
+
+    for chunk_key, chunk_value in code_cell_input_chunks.items():
+        base_params[chunk_key] = chunk_value
+
+    for chunk_key, chunk_value in response_chunks.items():
+        base_params[chunk_key] = chunk_value
+
+    # Log number of usages (for mito server)
+    if num_usages is not None:
+        base_params[MITO_SERVER_NUM_USAGES] = str(num_usages)
+
+    if message_type == MessageType.SMART_DEBUG:
+        error_message = (
+            last_message_content.split("Error Message:")[-1]
+            .split("ERROR ANALYSIS:")[0]
+            .strip()
+        )
+        error_type = error_message.split(": ")[0]
+
+        final_params = base_params
+        final_params["error_message"] = error_message
+        final_params["error_type"] = error_type
+
+        log("mito_ai_smart_debug_success", params=final_params, key_type=key_type, thread_id=thread_id)
+    elif message_type == MessageType.CODE_EXPLAIN:
+        final_params = base_params
+
+        log("mito_ai_code_explain_success", params=final_params, key_type=key_type, thread_id=thread_id)
+    elif message_type == MessageType.CHAT:
+        final_params = base_params
+
+        # Chunk the user input
+        user_input_chunks = chunk_param(user_input, "user_input")
+        
+        for chunk_key, chunk_value in user_input_chunks.items():
+            final_params[chunk_key] = chunk_value
+
+        log("mito_ai_chat_success", params=final_params, key_type=key_type, thread_id=thread_id)
+    elif message_type == MessageType.AGENT_EXECUTION:
+        final_params = base_params
+
+        # Chunk the user input
+        user_input_chunks = chunk_param(user_input, "user_input")
+        
+        for chunk_key, chunk_value in user_input_chunks.items():
+            final_params[chunk_key] = chunk_value
+            
+        # If the user input is not empty, then this is the user giving the agent a new task
+        new_user_input = 'True' if len(user_input_chunks) > 0 else 'False'
+        final_params["new_user_input"] = new_user_input
+
+        log("mito_ai_agent_execution_success", params=final_params, key_type=key_type, thread_id=thread_id)
+    elif message_type == MessageType.INLINE_COMPLETION:
+        final_params = base_params
+        log("mito_ai_inline_completion_success", params=final_params, key_type=key_type, thread_id=thread_id)
+    elif message_type == MessageType.AGENT_AUTO_ERROR_FIXUP:
+        final_params = base_params
+        log("mito_ai_agent_auto_error_fixup_success", params=final_params, key_type=key_type, thread_id=thread_id)
+    else:
+        final_params = base_params
+        final_params["note"] = (
+            "This input_location has not been accounted for in `telemetry_utils.py`."
+        )
+
+        log(f"mito_ai_{message_type.value}_success", params=final_params, key_type=key_type, thread_id=thread_id)
+        
+def log_db_connection_attempt(connection_type: str) -> None:
+    log("mito_ai_db_connection_attempt", params={"connection_type": connection_type})
+
+def log_db_connection_success(connection_type: str, schema: Dict[str, Any]) -> None:
+    log(
+        "mito_ai_db_connection_success",
+        params={
+            "connection_type": connection_type,
+        },
+    )
+
+def log_db_connection_error(connection_type: str, error_message: str) -> None:
+    log(
+        "mito_ai_db_connection_error", 
+        params={
+            "connection_type": connection_type, 
+            "error_message": error_message,
+        }
+    )
+
+def log_file_upload_attempt(
+    filename: str, file_extension: str, is_chunked: bool, total_chunks: int
+) -> None:
+    log(
+        "mito_ai_file_upload_attempt",
+        params={
+            "filename": filename,
+            "file_extension": file_extension,
+            "is_chunked": is_chunked,
+            "total_chunks": total_chunks,
+        },
+    )
+
+def log_file_upload_failure(error: str) -> None:
+    log("mito_ai_file_upload_failure", params={"error_message": error})
+
+def log_ai_completion_retry(key_type: Literal['mito_server_key', 'user_key'], thread_id: str, message_type: MessageType, error: BaseException) -> None:
+    log(MITO_AI_COMPLETION_RETRY, params={KEY_TYPE_PARAM: key_type, "message_type": message_type.value}, thread_id=thread_id, key_type=key_type, error=error)
+    
+def log_ai_completion_error(
+    key_type: Literal['mito_server_key', 'user_key'],
+    thread_id: str,
+    message_type: MessageType, 
+    error: BaseException
+) -> None:
+    log(MITO_AI_COMPLETION_ERROR, params={KEY_TYPE_PARAM: key_type, "message_type": message_type.value}, thread_id=thread_id, key_type=key_type, error=error)
+    
+def log_mito_server_free_tier_limit_reached(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType) -> None:
+    log(MITO_SERVER_FREE_TIER_LIMIT_REACHED, params={KEY_TYPE_PARAM: key_type, "message_type": message_type.value}, key_type=key_type)
+
+
+#################################
+# Streamlit Logs
+#################################
+
+### 
+# Converting Notebook into Streamlit App code
+###
+
+def log_streamlit_app_conversion_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, edit_prompt: str) -> None:
+    log(
+        "mito_ai_streamlit_app_conversion_success", 
+        key_type=key_type,
+        params={
+            "edit_prompt": edit_prompt,
+            "message_type": message_type.value
+        }
+    )
+    
+def log_streamlit_app_validation_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: List[str]) -> None:
+    log(
+        "mito_ai_streamlit_app_conversion_retry",
+        params={
+            "error_message": error,
+            "message_type": message_type.value
+        },
+        key_type=key_type
+    )
+    
+def log_streamlit_app_conversion_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error_message: str, formatted_traceback: str, edit_prompt: str) -> None:
+    log(
+        "mito_ai_streamlit_app_conversion_error",
+        params={
+            "error_message": error_message,
+            "traceback": formatted_traceback,
+            "edit_prompt": edit_prompt,
+            "message_type": message_type.value
+        },
+        key_type=key_type
+    )
+    
+###
+# Setting up Preview
+###
+
+def log_streamlit_app_preview_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, edit_prompt: str) -> None:
+    log(
+        "mito_ai_streamlit_app_preview_success",
+        key_type=key_type,
+        params={
+            "edit_prompt": edit_prompt,
+            "message_type": message_type.value
+        }
+    )
+    
+def log_streamlit_app_preview_failure(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error_message: str, formatted_traceback: str, edit_prompt: str) -> None:
+    log(
+        "mito_ai_streamlit_app_preview_failure",
+        key_type=key_type,
+        params={
+            "error_message": error_message,
+            "traceback": formatted_traceback,
+            "message_type": message_type.value,
+            "edit_prompt": edit_prompt
+        }
+    )
+    
+###
+# Deploying Streamlit App
+###
+
+def log_streamlit_app_deployment_failure(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: Dict) -> None:
+    log(
+        "mito_ai_streamlit_app_deployment_failure",
+        key_type=key_type,
+        params={
+            "error": error, # Contains all details in app_deploy.models.AppDeployError class
+            "message_type": message_type.value
+        }
+    )