mito-ai 0.1.57__py3-none-any.whl → 0.1.59__py3-none-any.whl

mito_ai/utils/model_utils.py

@@ -0,0 +1,257 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple, Union, Optional, cast
+from mito_ai import constants
+from mito_ai.utils.version_utils import is_enterprise
+from mito_ai.enterprise.utils import is_abacus_configured
+
+# Model ordering: [fastest, ..., slowest] for each provider
+ANTHROPIC_MODEL_ORDER = [
+    "claude-haiku-4-5-20251001",  # Fastest
+    "claude-sonnet-4-5-20250929",  # Slower
+]
+
+OPENAI_MODEL_ORDER = [
+    "gpt-4.1",      # Fastest
+    "gpt-5",
+    "gpt-5.2",      # Slower
+]
+
+GEMINI_MODEL_ORDER = [
+    "gemini-3-flash-preview",  # Fastest
+    "gemini-3-pro-preview",     # Slower
+]
+
+# Standard model names (used when not in enterprise mode or when LiteLLM is not configured)
+STANDARD_MODELS = [
+    "gpt-4.1",
+    "gpt-5.2",
+    "claude-sonnet-4-5-20250929",
+    "claude-haiku-4-5-20251001",
+    "gemini-3-flash-preview",
+    "gemini-3-pro-preview",
+]
+
+
+def get_available_models() -> List[str]:
+    """
+    Determine which models are available based on enterprise mode and router configuration.
+    
+    Priority order:
+    1. Abacus (if configured)
+    2. LiteLLM (if configured)
+    3. Standard models
+    
+    Returns:
+        List of available model names with appropriate prefixes.
+    """
+    # Check if enterprise mode is enabled AND Abacus is configured (highest priority)
+    if is_abacus_configured():
+        # Return Abacus models (with Abacus/ prefix)
+        return constants.ABACUS_MODELS
+    # Check if enterprise mode is enabled AND LiteLLM is configured
+    elif is_enterprise() and constants.LITELLM_BASE_URL and constants.LITELLM_MODELS:
+        # Return LiteLLM models (with LiteLLM/provider/ prefix or legacy provider/ prefix)
+        return constants.LITELLM_MODELS
+    else:
+        # Return standard models
+        return STANDARD_MODELS
+
+
+def get_fast_model_for_selected_model(selected_model: str) -> str:
+    """
+    Get the fastest model for the client of the selected model.
+    
+    - For standard providers, returns the first (fastest) model from that provider's order.
+    - For enterprise router models (Abacus/LiteLLM), finds the fastest available model by comparing indices.
+    """
+    # Check if this is an enterprise router model (has "/" or router prefix)
+    if "/" in selected_model or selected_model.lower().startswith(('abacus/', 'litellm/')):
+        # Find the fastest model from available models
+        available_models = get_available_models()
+        if not available_models:
+            return selected_model
+        
+        # Filter to only router models (those with "/")
+        router_models = [model for model in available_models if "/" in model]
+        if not router_models:
+            return selected_model
+        
+        # Extract provider/model pairs for ordering
+        pairs_with_indices = []
+        for model in router_models:
+            # Strip router prefix to get underlying model info
+            model_without_router = strip_router_prefix(model)
+            
+            # For Abacus: model_without_router is just the model name (e.g., "gpt-4.1")
+            # For LiteLLM: model_without_router is "provider/model" (e.g., "openai/gpt-4.1")
+            if "/" in model_without_router:
+                # LiteLLM format: provider/model
+                pair = model_without_router.split("/", 1)
+            else:
+                # Abacus format: just model name, need to determine provider
+                provider = get_underlying_model_provider(model)
+                if provider:
+                    pair = [provider, model_without_router]
+                else:
+                    continue
+            
+            index = get_model_order_index(pair)
+            if index is not None:
+                pairs_with_indices.append((model, index))
+        
+        if not pairs_with_indices:
+            return selected_model
+
+        # Find the model with the minimum index (fastest model)
+        fastest_model, _ = min(pairs_with_indices, key=lambda x: x[1])
+        
+        return fastest_model
+    
+    # Standard provider logic - ensure we return a model from the same provider
+    model_lower = selected_model.lower()
+    
+    # Determine provider and get fastest model
+    if model_lower.startswith('claude'):
+        return ANTHROPIC_MODEL_ORDER[0]
+    elif model_lower.startswith('gpt'):
+        return OPENAI_MODEL_ORDER[0]
+    elif model_lower.startswith('gemini'):
+        return GEMINI_MODEL_ORDER[0]
+
+    return selected_model
+
+def get_smartest_model_for_selected_model(selected_model: str) -> str:
+    """
+    Get the smartest model for the client of the selected model.
+    
+    - For standard providers, returns the last (smartest) model from that provider's order.
+    - For enterprise router models (Abacus/LiteLLM), finds the smartest available model by comparing indices.
+    """
+    # Check if this is an enterprise router model (has "/" or router prefix)
+    if "/" in selected_model or selected_model.lower().startswith(('abacus/', 'litellm/')):
+        # Extract underlying provider from selected model
+        selected_provider = get_underlying_model_provider(selected_model)
+        if not selected_provider:
+            return selected_model
+        
+        # Find the smartest model from available models
+        available_models = get_available_models()
+        if not available_models:
+            return selected_model
+        
+        # Filter to only router models with the same underlying provider
+        router_models = []
+        for model in available_models:
+            if "/" in model:
+                model_provider = get_underlying_model_provider(model)
+                if model_provider == selected_provider:
+                    router_models.append(model)
+        
+        if not router_models:
+            return selected_model
+        
+        # Extract provider/model pairs for ordering
+        pairs_with_indices = []
+        for model in router_models:
+            # Strip router prefix to get underlying model info
+            model_without_router = strip_router_prefix(model)
+            
+            # For Abacus: model_without_router is just the model name (e.g., "gpt-4.1")
+            # For LiteLLM: model_without_router is "provider/model" (e.g., "openai/gpt-4.1")
+            if "/" in model_without_router:
+                # LiteLLM format: provider/model
+                pair = model_without_router.split("/", 1)
+            else:
+                # Abacus format: just model name, provider already determined
+                pair = [selected_provider, model_without_router]
+            
+            index = get_model_order_index(pair)
+            if index is not None:
+                pairs_with_indices.append((model, index))
+        
+        if not pairs_with_indices:
+            return selected_model
+
+        # Find the model with the maximum index (smartest model)
+        smartest_model, _ = max(pairs_with_indices, key=lambda x: x[1])
+        
+        return smartest_model
+    
+    # Standard provider logic
+    model_lower = selected_model.lower()
+    
+    # Determine provider and get smartest model
+    if model_lower.startswith('claude'):
+        return ANTHROPIC_MODEL_ORDER[-1]
+    elif model_lower.startswith('gpt'):
+        return OPENAI_MODEL_ORDER[-1]
+    elif model_lower.startswith('gemini'):
+        return GEMINI_MODEL_ORDER[-1]
+
+    return selected_model
+
+def strip_router_prefix(model: str) -> str:
+    """
+    Strip router prefix from model name.
+    
+    Examples:
+    - "Abacus/gpt-4.1" -> "gpt-4.1"
+    - "LiteLLM/openai/gpt-4.1" -> "openai/gpt-4.1"
+    - "gpt-4.1" -> "gpt-4.1" (no prefix, return as-is)
+    """
+    if model.lower().startswith('abacus/'):
+        return model[7:]  # Strip "Abacus/"
+    elif model.lower().startswith('litellm/'):
+        return model[8:]  # Strip "LiteLLM/"
+    return model
+
+def get_underlying_model_provider(full_model_provider_id: str) -> Optional[str]:
+    """
+    Determine the underlying AI provider from a model identifier.
+    
+    For Abacus models (Abacus/model), determine the provider from model name pattern.
+    For LiteLLM models (LiteLLM/provider/model), extract the provider from the prefix.
+    
+    Returns:
+        Provider name ("openai", "anthropic", "google") or None if cannot determine.
+    """
+    # Strip router prefix first
+    model_without_router = strip_router_prefix(full_model_provider_id)
+    
+    # Check if it's a LiteLLM format (provider/model)
+    if "/" in model_without_router:
+        provider, _ = model_without_router.split("/", 1)
+        return provider.lower()
+    
+    # For Abacus models without provider prefix, determine from model name
+    model_lower = model_without_router.lower()
+    if model_lower.startswith('gpt'):
+        return 'openai'
+    elif model_lower.startswith('claude'):
+        return 'anthropic'
+    elif model_lower.startswith('gemini'):
+        return 'google'
+    
+    return None
+
+def get_model_order_index(pair: List[str]) -> Optional[int]:
+    provider, model_name = pair
+    if provider == "openai":
+        try:
+            return OPENAI_MODEL_ORDER.index(model_name)
+        except ValueError:
+            return None
+    elif provider == "anthropic":
+        try:
+            return ANTHROPIC_MODEL_ORDER.index(model_name)
+        except ValueError:
+            return None
+    elif provider == "google":
+        try:
+            return GEMINI_MODEL_ORDER.index(model_name)
+        except ValueError:
+            return None
+    else:
+        return None

mito_ai/utils/open_ai_utils.py

@@ -6,31 +6,18 @@
 
 # Copyright (c) Saga Inc.
 
-import asyncio
-import json
-import time
-from typing import Any, Dict, List, Optional, Final, Union, AsyncGenerator, Tuple, Callable
+from typing import Any, Dict, List, Optional, 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.completions.models import 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.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.enterprise.utils import is_abacus_configured
 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],
@@ -153,19 +140,12 @@ async def stream_ai_completion_from_mito_server(
 
 
 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,
@@ -179,25 +159,33 @@ def get_open_ai_completion_function_params(
     # 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. 
+    # Note: Abacus only supports {"type": "json"} format, not the full JSON schema format.
     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
+        # Check if we're using Abacus - it only supports simple {"type": "json"} format
+        if is_abacus_configured() or model.lower().startswith('abacus/'):
+            completion_function_params["response_format"] = {
+                "type": "json"
+            }
+        else:
+            # For OpenAI and other providers, use the full JSON schema format
+            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

@@ -8,13 +8,25 @@ 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
+    Determine the model type based on the model name prefix.
+    
+    Priority order:
+    1. Check for router prefixes (Abacus/, LiteLLM/)
+    2. Check for legacy LiteLLM format (provider/model)
+    3. Check for standard model name patterns
     """
     if not model:
         return None
 
     model_lower = model.lower()
 
+    # Check for router prefixes first (highest priority)
+    if model_lower.startswith('abacus/'):
+        return 'abacus'
+    elif model_lower.startswith('litellm/'):
+        return 'litellm'    
+
+    # Check for standard model name patterns
     if model_lower.startswith('claude'):
         return 'claude'
     elif model_lower.startswith('gemini'):
@@ -25,32 +37,4 @@ def get_model_provider(model: str) -> Union[str, None]:
         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_SCRATCHPAD_RESULT,
-        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/telemetry_utils.py

@@ -4,7 +4,7 @@
 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.version_utils import MITOSHEET_HELPER_PRIVATE, is_pro, is_enterprise
 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__
@@ -83,9 +83,13 @@ 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
     """
+    # Enterprise mode disables all telemetry
+    if is_enterprise():
+        return False
+    
     # If the user is on the Mito server, then they are sending
     # us their information already
-    if key_type == 'mito_server_key':
+    if key_type == MITO_SERVER_KEY:
         return True
     
     # If private helper is installed, then we don't log anything
@@ -109,6 +113,10 @@ def identify(key_type: Optional[str] = None, is_electron: Optional[bool] = None)
     Helper function for identifying a user. We just take
     their python version, mito version, and email.
     """
+    # Skip entirely if enterprise mode is enabled
+    if is_enterprise():
+        return
+    
     if not telemetry_turned_on(key_type):
         return
 
@@ -208,6 +216,10 @@ def log(
         del final_params[param_name]
     final_params.update(params_to_add)
             
+    # Skip entirely if enterprise mode is enabled
+    if is_enterprise():
+        return
+    
     # 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):