mito-ai 0.1.50__py3-none-any.whl

mito_ai/utils/mito_server_utils.py

@@ -0,0 +1,242 @@
+# 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, Optional, Callable, Union, AsyncGenerator
+from mito_ai.completions.models import MessageType, CompletionReply, CompletionStreamChunk, CompletionItem
+from mito_ai.utils.server_limits import check_mito_server_quota, update_mito_server_quota
+from tornado.httpclient import HTTPResponse
+from mito_ai.constants import MITO_GEMINI_URL
+from mito_ai.utils.utils import _create_http_client
+
+MITO_ERROR_MARKER = "MITO_ERROR_MARKER:"
+
+class ProviderCompletionException(Exception):
+    """Custom exception for Mito server errors that converts well to CompletionError."""
+    
+    def __init__(self, error_message: str, provider_name: str = "LLM Provider", error_type: str = "LLMProviderError"):
+        self.error_message = error_message
+        self.provider_name = provider_name
+        self.error_type = error_type
+        
+        # Create user-friendly title and hint
+        self.user_friendly_title = f"{provider_name} Error: {error_message}"
+        self.user_friendly_hint = f"There was a problem with {provider_name}. Try switching to a different model and trying again."
+
+        # Set args[0] for fallback compatibility
+        super().__init__(self.user_friendly_title)
+
+    def __str__(self) -> str:
+        return f"{self.provider_name} Error: {self.error_message}"
+
+
+async def get_response_from_mito_server(
+    url: str, 
+    headers: dict, 
+    data: Dict[str, Any],
+    timeout: int, 
+    max_retries: int,
+    message_type: MessageType,
+    provider_name: str = "Mito Server"
+) -> str:
+    """
+    Get a response from the Mito server.
+    
+    Raises:
+        ProviderCompletionException: When the server returns an error or invalid response
+        Exception: For network/HTTP errors (let these bubble up to be handled by retry logic)
+    """
+    # First check the mito server quota. If the user has reached the limit, we raise an exception.
+    check_mito_server_quota(message_type)
+
+    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
+    start_time = time.time()
+    
+    try:
+        res = await http_client.fetch(
+            url,
+            method="POST",
+            headers=headers,
+            body=json.dumps(data),
+            request_timeout=http_client_timeout
+        )
+        print(f"Mito server request completed in {time.time() - start_time:.2f} seconds")
+        
+        # Parse and validate response
+        try:
+            content = json.loads(res.body.decode("utf-8"))
+                        
+            if "completion" in content:
+                return content["completion"] # type: ignore
+            elif "error" in content:
+                # Server returned an error
+                raise ProviderCompletionException(content['error'], provider_name=provider_name)
+            else:
+                # Invalid response format
+                raise ProviderCompletionException(f"No completion found in response: {content}", provider_name=provider_name)
+        except ProviderCompletionException:
+            # Re-raise ProviderCompletionException as-is
+            raise
+        except Exception as e:
+            raise ProviderCompletionException(f"Error parsing response: {str(e)}", provider_name=provider_name)
+            
+    finally:
+        try:
+            # We always update the quota, even if there is an error
+            update_mito_server_quota(message_type)
+        except Exception as e:
+            pass
+        
+        http_client.close()
+
+
+async def stream_response_from_mito_server(
+    url: str,
+    headers: Dict[str, str],
+    data: Dict[str, Any],
+    timeout: int,
+    max_retries: int,
+    message_type: MessageType,
+    reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+    message_id: str,
+    chunk_processor: Optional[Callable[[str], str]] = None,
+    provider_name: str = "Mito Server",
+) -> AsyncGenerator[str, None]:
+    """
+    Stream responses from the Mito server.
+    
+    This is a unified streaming function that can be used by all providers (OpenAI, Anthropic, Gemini).
+    
+    Args:
+        url: The Mito server URL to stream from
+        headers: Request headers
+        data: Request data
+        timeout: Request timeout in seconds
+        max_retries: Maximum number of retries
+        message_type: The message type for quota tracking
+        provider_name: Name of the provider for error messages
+        reply_fn: Optional function to call with each chunk for streaming replies
+        message_id: The message ID to track the request
+        chunk_processor: Optional function to process chunks before yielding (e.g., for Gemini's special processing)
+        
+    Yields:
+        Chunks of text from the streaming response
+    """
+    # Check the mito server quota
+    check_mito_server_quota(message_type)
+    
+    # Create HTTP client with appropriate timeout settings
+    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
+    
+    # Set up streaming infrastructure
+    start_time = time.time()
+    chunk_queue: asyncio.Queue[str] = asyncio.Queue()
+    fetch_complete = False
+    
+    # Define a callback to process chunks and add them to the queue
+    def chunk_callback(chunk: bytes) -> None:
+        try:
+            chunk_str = chunk.decode('utf-8')
+            asyncio.create_task(chunk_queue.put(chunk_str))
+        except Exception as e:
+            print(f"Error processing {provider_name} streaming chunk: {str(e)}")
+    
+    # Execute the streaming request
+    fetch_future = None
+    try:
+        fetch_future = http_client.fetch(
+            url, 
+            method="POST", 
+            headers=headers, 
+            body=json.dumps(data), 
+            request_timeout=http_client_timeout,
+            streaming_callback=chunk_callback
+        )
+        
+        # Create a task to wait for the fetch to complete
+        async def wait_for_fetch() -> None:
+            try:
+                await fetch_future
+                nonlocal fetch_complete
+                fetch_complete = True
+                print(f"{provider_name} fetch completed")
+            except Exception as e:
+                print(f"Error in {provider_name} fetch: {str(e)}")
+                raise
+        
+        # Start the task to wait for fetch completion
+        fetch_task = asyncio.create_task(wait_for_fetch())
+        
+        # Yield chunks as they arrive
+        while not (fetch_complete and chunk_queue.empty()):
+            try:
+                # Wait for a chunk with a timeout to prevent deadlocks
+                chunk = await asyncio.wait_for(chunk_queue.get(), timeout=0.1)
+                
+                # Process chunk if processor is provided
+                processed_chunk = chunk
+                if chunk_processor:
+                    processed_chunk = chunk_processor(chunk)
+
+                # Check if this chunk contains an error marker
+                if processed_chunk.startswith(MITO_ERROR_MARKER):
+                    error_message = processed_chunk[len(MITO_ERROR_MARKER):]
+                    print(f"Detected error in {provider_name} stream: {error_message}")
+                    raise ProviderCompletionException(error_message, provider_name=provider_name)
+
+                if reply_fn is not None and message_id is not None:
+                    # Send the chunk directly to the frontend
+                    reply_fn(CompletionStreamChunk(
+                        parent_id=message_id,
+                        chunk=CompletionItem(
+                            content=processed_chunk,
+                            isIncomplete=True,
+                            token=message_id,
+                        ),
+                        done=False,
+                    ))
+                
+                yield chunk
+            except asyncio.TimeoutError:
+                # No chunk available within timeout, check if fetch is complete
+                if fetch_complete and chunk_queue.empty():
+                    break
+                
+                # Otherwise continue waiting for chunks
+                continue
+                
+        print(f"\n{provider_name} stream completed in {time.time() - start_time:.2f} seconds")
+        
+        if reply_fn is not None and message_id is not None:
+            # Send a final chunk to indicate completion
+            reply_fn(CompletionStreamChunk(
+                parent_id=message_id,
+                chunk=CompletionItem(
+                    content="",
+                    isIncomplete=False,
+                    token=message_id,
+                ),
+                done=True,
+            ))
+    except Exception as e:
+        print(f"\n{provider_name} stream failed after {time.time() - start_time:.2f} seconds with error: {str(e)}")
+        # If an exception occurred, ensure the fetch future is awaited to properly clean up
+        if fetch_future:
+            try:
+                await fetch_future
+            except Exception:
+                pass
+        raise
+    finally:
+        # Clean up resources
+        try:
+            # We always update the quota, even if there is an error
+            update_mito_server_quota(message_type)
+        except Exception as e:
+            pass
+        
+        http_client.close()
+        
+    

mito_ai/utils/open_ai_utils.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python
+# coding: utf-8
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+
+# Copyright (c) Saga Inc.
+
+import asyncio
+import json
+import time
+from typing import Any, Dict, List, Optional, Final, Union, AsyncGenerator, Tuple, Callable
+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 tornado.httpclient import AsyncHTTPClient
+from openai.types.chat import ChatCompletionMessageParam
+
+from mito_ai.utils.utils import is_running_test
+from mito_ai.completions.models import MessageType, ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem
+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.utils.version_utils import is_pro
+from mito_ai.utils.server_limits import check_mito_server_quota
+from mito_ai.utils.telemetry_utils import log_ai_completion_success
+from .utils import _create_http_client
+from mito_ai.constants import MITO_OPENAI_URL
+
+
+__user_email: Optional[str] = None
+__user_id: Optional[str] = None
+
+FAST_OPENAI_MODEL = "gpt-4.1-nano"
+
+def _prepare_request_data_and_headers(
+    last_message_content: Union[str, None],
+    ai_completion_data: Dict[str, Any],
+    timeout: int,
+    max_retries: int,
+    message_type: MessageType,
+) -> Tuple[Dict[str, Any], Dict[str, str]]:
+    """
+    Prepare request data and headers for Mito server API calls.
+    
+    Args:
+        last_message_content: The last message content
+        ai_completion_data: The AI completion data
+        timeout: The timeout in seconds
+        max_retries: The maximum number of retries
+        message_type: The message type
+        
+    Returns:
+        A tuple containing the request data and headers
+    """    
+    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)
+
+    data = {
+        "timeout": timeout,
+        "max_retries": max_retries,
+        "email": __user_email,
+        "user_id": __user_id,
+        "data": ai_completion_data,
+        "user_input": last_message_content or "",  # We add this just for logging purposes
+    }
+
+    headers = {
+        "Content-Type": "application/json",
+    }
+    
+    return data, headers
+
+async def get_ai_completion_from_mito_server(
+    last_message_content: Union[str, None],
+    ai_completion_data: Dict[str, Any],
+    timeout: int,
+    max_retries: int,
+    message_type: MessageType,
+) -> str:
+    
+    # Prepare request data and headers
+    data, headers = _prepare_request_data_and_headers(
+        last_message_content, 
+        ai_completion_data, 
+        timeout, 
+        max_retries, 
+        message_type
+    )
+    
+    return await get_response_from_mito_server(
+        MITO_OPENAI_URL, 
+        headers, 
+        data, 
+        timeout, 
+        max_retries, 
+        message_type,
+        provider_name="OpenAI"
+    )
+
+async def stream_ai_completion_from_mito_server(
+    last_message_content: Union[str, None],
+    ai_completion_data: Dict[str, Any],
+    timeout: int,
+    max_retries: int,
+    message_type: MessageType,
+    reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+    message_id: str,
+) -> AsyncGenerator[str, None]:
+    """
+    Stream AI completions from the Mito server.
+    
+    This function is similar to get_ai_completion_from_mito_server but handles streaming responses.
+    It yields the streamed content as it arrives.
+    
+    Args:
+        last_message_content: The last message content
+        ai_completion_data: The AI completion data
+        timeout: The timeout in seconds
+        max_retries: The maximum number of retries
+        message_type: The message type
+        reply_fn: Optional function to call with each chunk for streaming replies
+        message_id: The message ID to track the request
+        
+    Yields:
+        Chunks of text from the streaming response
+    """
+    # Prepare request data and headers
+    data, headers = _prepare_request_data_and_headers(
+        last_message_content, 
+        ai_completion_data, 
+        timeout, 
+        max_retries, 
+        message_type
+    )
+    
+    # Use the unified streaming function
+    async for chunk in stream_response_from_mito_server(
+        url=MITO_OPENAI_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=None,
+        provider_name="OpenAI",
+    ):
+        yield chunk
+
+
+def get_open_ai_completion_function_params(
+    message_type: MessageType,
+    model: str, 
+    messages: List[ChatCompletionMessageParam], 
+    stream: bool,
+    response_format_info: Optional[ResponseFormatInfo] = None,
+) -> Dict[str, Any]:
+    
+    print("MESSAGE TYPE: ", message_type)
+    message_requires_fast_model = does_message_require_fast_model(message_type)
+    model = FAST_OPENAI_MODEL if message_requires_fast_model else model
+    
+    print(f"model: {model}")
+    
+    completion_function_params = {
+        "model": model,
+        "stream": stream,
+        "messages": messages,
+    }
+    
+    # If a response format is provided, we need to convert it to a json schema.
+    # Pydantic models are supported by the OpenAI API, however, we need to be able to 
+    # serialize it for requests that are going to be sent to the mito server. 
+    # OpenAI expects a very specific schema as seen below. 
+    if response_format_info:
+        json_schema = response_format_info.format.schema()
+        
+        # Add additionalProperties: False to the top-level schema
+        json_schema["additionalProperties"] = False
+        
+        # Nested object definitions in $defs need to have additionalProperties set to False also
+        if "$defs" in json_schema:
+            for def_name, def_schema in json_schema["$defs"].items():
+                if def_schema.get("type") == "object":
+                    def_schema["additionalProperties"] = False
+        
+        completion_function_params["response_format"] = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": f"{response_format_info.name}",
+                "schema": json_schema,
+                "strict": True
+            }
+        }
+
+    return completion_function_params

mito_ai/utils/provider_utils.py

@@ -0,0 +1,49 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Union
+
+from mito_ai.completions.models import MessageType
+
+
+def get_model_provider(model: str) -> Union[str, None]:
+    """
+    Determine the model type based on the model name prefix
+    """
+    if not model:
+        return None
+
+    model_lower = model.lower()
+
+    if model_lower.startswith('claude'):
+        return 'claude'
+    elif model_lower.startswith('gemini'):
+        return 'gemini'
+    elif model_lower.startswith('ollama'):
+        return 'ollama'
+    elif model_lower.startswith('gpt'):
+        return 'openai'
+
+    return None
+
+
+def does_message_require_fast_model(message_type: MessageType) -> bool:
+    """
+    Determines if a message requires the fast model.
+    
+    The fast model is used for messages that are not chat messages.
+    For example, inline completions and chat name generation need to be fast
+    so they don't slow down the user's experience.
+    """
+    
+    if message_type in (MessageType.CHAT, MessageType.SMART_DEBUG, MessageType.CODE_EXPLAIN, MessageType.AGENT_EXECUTION, MessageType.AGENT_AUTO_ERROR_FIXUP):
+        return False
+    elif message_type in (MessageType.INLINE_COMPLETION, MessageType.CHAT_NAME_GENERATION):
+        return True
+    elif message_type in (MessageType.START_NEW_CHAT, MessageType.FETCH_HISTORY, MessageType.GET_THREADS, MessageType.DELETE_THREAD, MessageType.UPDATE_MODEL_CONFIG):
+        # These messages don't use any model, but we add them here for type safety
+        return True
+    else:
+        raise ValueError(f"Invalid message type: {message_type}")
+    
+    

mito_ai/utils/schema.py

@@ -0,0 +1,86 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+from datetime import datetime
+from .utils import get_random_id
+
+# Current field names
+# Some helpful constants
+GITHUB_ACTION_ID = 'github_action'
+GITHUB_ACTION_EMAIL = 'github@action.com'
+
+# Old field names
+UJ_INTENDED_BEHAVIOR = 'intended_behavior'
+UJ_CLOSED_FEEDBACK = 'closed_feedback'
+UJ_MITOSHEET_LAST_FIVE_USAGES = 'mitosheet_last_five_usages'
+
+# Current field names
+UJ_USER_JSON_VERSION = 'user_json_version'
+UJ_STATIC_USER_ID = 'static_user_id'
+UJ_USER_SALT = 'user_salt'
+UJ_USER_EMAIL = 'user_email'
+UJ_RECEIVED_TOURS = 'received_tours'
+UJ_FEEDBACKS = 'feedbacks'
+UJ_FEEDBACKS_V2 = 'feedbacks_v2'
+UJ_MITOSHEET_CURRENT_VERSION = 'mitosheet_current_version'
+UJ_MITOSHEET_LAST_UPGRADED_DATE = 'mitosheet_last_upgraded_date'
+UJ_MITOSHEET_LAST_FIFTY_USAGES = 'mitosheet_last_fifty_usages'
+UJ_MITOSHEET_TELEMETRY = 'mitosheet_telemetry'
+UJ_MITOSHEET_PRO = 'mitosheet_pro'
+UJ_MITOSHEET_ENTERPRISE = 'mitosheet_enterprise'
+UJ_EXPERIMENT = 'experiment'
+UJ_RECEIVED_CHECKLISTS = 'received_checklists'
+UJ_AI_PRIVACY_POLICY = 'ai_privacy_policy'
+UJ_AI_MITO_API_NUM_USAGES = 'ai_mito_api_num_usages'
+UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES = 'ai_mito_autocomplete_num_usages'
+UJ_MITO_AI_FIRST_USAGE_DATE = 'mito_ai_first_usage_date'
+UJ_MITO_AI_LAST_RESET_DATE = 'mito_ai_last_reset_date'
+
+MITO_CONFIG_KEY_HOME_FOLDER = 'MITO_CONFIG_HOME_FOLDER'
+if MITO_CONFIG_KEY_HOME_FOLDER in os.environ:
+    HOME_FOLDER = os.path.expanduser(os.environ[MITO_CONFIG_KEY_HOME_FOLDER])
+else:
+    HOME_FOLDER = os.path.expanduser('~')
+
+# Where all global .mito files are stored
+MITO_FOLDER = os.path.join(HOME_FOLDER, ".mito")
+
+"""
+The most up to date version of the user.json object
+"""
+USER_JSON_VERSION_10 = {
+    # The new version of the user json object
+    UJ_USER_JSON_VERSION: 10,
+    # The static id of the user
+    UJ_STATIC_USER_ID: get_random_id(),
+    # A random secret that the user can use as salt when hashing things
+    UJ_USER_SALT: get_random_id(),
+    # Email of the user
+    UJ_USER_EMAIL: '',
+    # Tours that the user has received
+    UJ_RECEIVED_TOURS: [],
+    # A list of all the feedback the user has given
+    UJ_FEEDBACKS: [],
+    UJ_FEEDBACKS_V2: {},
+    UJ_MITOSHEET_CURRENT_VERSION: 0,
+    UJ_MITOSHEET_LAST_UPGRADED_DATE: datetime.today().strftime('%Y-%m-%d'),
+    UJ_MITOSHEET_LAST_FIFTY_USAGES: [datetime.today().strftime('%Y-%m-%d')],
+    UJ_MITOSHEET_TELEMETRY: True,
+    UJ_MITOSHEET_PRO: False,
+    UJ_MITOSHEET_ENTERPRISE: False,
+    UJ_EXPERIMENT: {
+        'experiment_id': 'installer_communication_and_time_to_value',
+        'variant': 'B',
+    },
+    UJ_RECEIVED_CHECKLISTS: {},
+    UJ_AI_PRIVACY_POLICY: False,
+    UJ_AI_MITO_API_NUM_USAGES: 0,
+    UJ_AI_MITO_AUTOCOMPLETE_NUM_USAGES: 0,
+    UJ_MITO_AI_FIRST_USAGE_DATE: datetime.today().strftime('%Y-%m-%d'),
+    UJ_MITO_AI_LAST_RESET_DATE: datetime.today().strftime('%Y-%m-%d'),
+}
+
+# This is the most up to date user json, and you must update it when
+# you add a new schema
+USER_JSON_DEFAULT = USER_JSON_VERSION_10