mito-ai 0.1.32__py3-none-any.whl → 0.1.34__py3-none-any.whl

mito_ai/utils/gemini_utils.py

@@ -5,14 +5,17 @@ import asyncio
 import json
 import time
 from typing import Any, Dict, List, Optional, Callable, Union, AsyncGenerator, Tuple
-from tornado.httpclient import AsyncHTTPClient
+from mito_ai.utils.mito_server_utils import get_response_from_mito_server
 from mito_ai.completions.models import AgentResponse, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
-from .utils import _create_http_client
 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]],
@@ -62,25 +65,15 @@ async def get_gemini_completion_from_mito_server(
     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)
-    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
-    start_time = time.time()
-    try:
-        res = await http_client.fetch(
-            MITO_GEMINI_URL,
-            method="POST",
-            headers=headers,
-            body=json.dumps(data),
-            request_timeout=http_client_timeout
-        )
-        print(f"Gemini request completed in {time.time() - start_time:.2f} seconds")
-    except Exception as e:
-        print(f"Gemini request failed after {time.time() - start_time:.2f} seconds with error: {str(e)}")
-        raise
-    finally:
-        http_client.close()
-    
-    # The response is a string
-    return res.body.decode("utf-8")
+    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,
@@ -163,15 +156,18 @@ async def stream_gemini_completion_from_mito_server(
         http_client.close()
 
 def get_gemini_completion_function_params(
+    message_type: MessageType,
     model: str,
     contents: list[dict[str, Any]],
-    message_type: MessageType,
     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,

mito_ai/utils/mito_server_utils.py

@@ -0,0 +1,92 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.completions.models import MessageType
+from mito_ai.utils.server_limits import check_mito_server_quota, update_mito_server_quota
+from tornado.httpclient import HTTPResponse
+import time
+import json
+from typing import Any, Dict, Optional
+from mito_ai.constants import MITO_GEMINI_URL
+from mito_ai.utils.utils import _create_http_client
+
+
+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()
+        
+    

mito_ai/utils/open_ai_utils.py

@@ -10,6 +10,8 @@ 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
+from mito_ai.utils.provider_utils import does_message_require_fast_model
 from tornado.httpclient import AsyncHTTPClient
 from openai.types.chat import ChatCompletionMessageParam
 
@@ -27,6 +29,8 @@ 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],
@@ -46,10 +50,7 @@ def _prepare_request_data_and_headers(
         
     Returns:
         A tuple containing the request data and headers
-    """
-    # Check that the user is allowed to use the Mito Server
-    check_mito_server_quota(message_type)
-    
+    """    
     global __user_email, __user_id
 
     if __user_email is None:
@@ -89,47 +90,15 @@ async def get_ai_completion_from_mito_server(
         message_type
     )
     
-    # Create HTTP client with appropriate timeout settings
-    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
-    
-    # There are several types of timeout errors that can happen here. 
-    # == 504 Timeout (tornado.httpclient.HTTPClientError: 504) ==  
-    # The server (AWS Lambda) took too long to process your request
-    # == 599 Timeout (tornado.httpclient.HTTPClientError: 599) ==  
-    # The client (Tornado) gave up waiting for a response
-    
-    start_time = time.time()
-    try:
-        res = await http_client.fetch(
-            # Important: DO NOT CHANGE MITO_AI_URL. If you want to use the dev endpoint, 
-            # go to the top of this file and change MITO_AI_URL to MITO_AI_DEV_URL. We 
-            # have a pytest that ensures that the MITO_AI_URL is always set to MITO_AI_PROD_URL 
-            # before merging into dev. So if you change which variable we are using here, the 
-            # test will not catch our mistakes.
-            MITO_OPENAI_URL, 
-            method="POST", 
-            headers=headers, 
-            body=json.dumps(data), 
-            # For some reason, we need to add the request_timeout here as well
-            request_timeout=http_client_timeout
-        )
-        print(f"Request completed in {time.time() - start_time:.2f} seconds")
-    except Exception as e:
-        print(f"Request failed after {time.time() - start_time:.2f} seconds with error: {str(e)}")
-        raise
-    finally:
-        http_client.close()
-
-    # The lambda function returns a dictionary with a completion entry in it,
-    # so we just return that.
-    content = json.loads(res.body)
-    
-    if "completion" in content:
-        return content["completion"] # type: ignore
-    elif "error" in content:
-        raise Exception(f"{content['error']}")
-    else:
-        raise Exception(f"No completion found in response: {content}")
+    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],
@@ -273,12 +242,19 @@ 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,
@@ -311,7 +287,7 @@ def get_open_ai_completion_function_params(
         }
     
     # o3-mini will error if we try setting the temperature
-    if model == "gpt-4o-mini":
+    if not model.startswith("o3"):
         completion_function_params["temperature"] = 0.0
 
     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/telemetry_utils.py

@@ -24,6 +24,7 @@ PRINT_LOGS = False
 # 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 
@@ -82,7 +83,6 @@ 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':
@@ -338,6 +338,15 @@ def log_db_connection_success(connection_type: str, schema: Dict[str, Any]) -> N
         },
     )
 
+def log_ai_completion_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: BaseException) -> None:
+    log(MITO_AI_COMPLETION_RETRY, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, key_type=key_type, error=error)
+    
+def log_ai_completion_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: BaseException) -> None:
+    log(MITO_AI_COMPLETION_ERROR, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, 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}, key_type=key_type)
+
 def log_db_connection_error(connection_type: str, error_message: str) -> None:
     log(
         "mito_ai_db_connection_error", 
@@ -346,3 +355,4 @@ def log_db_connection_error(connection_type: str, error_message: str) -> None:
             "error_message": error_message,
         }
     )
+

{mito_ai-0.1.32.data → mito_ai-0.1.34.data}/data/share/jupyter/labextensions/mito_ai/build_log.json

@@ -710,7 +710,7 @@
             "semver": {},
             "vscode-diff": {},
             "mito_ai": {
-              "version": "0.1.32",
+              "version": "0.1.34",
               "singleton": true,
               "import": "/home/runner/work/mito/mito/mito-ai/lib/index.js"
             }

{mito_ai-0.1.32.data → mito_ai-0.1.34.data}/data/share/jupyter/labextensions/mito_ai/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.32",
+  "version": "0.1.34",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -138,7 +138,7 @@
     "outputDir": "mito_ai/labextension",
     "schemaDir": "schema",
     "_build": {
-      "load": "static/remoteEntry.a711c58b58423173bd24.js",
+      "load": "static/remoteEntry.51d07439b02aaa830975.js",
       "extension": "./extension",
       "style": "./style"
     }

{mito_ai-0.1.32.data → mito_ai-0.1.34.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/package.json.orig

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.32",
+  "version": "0.1.34",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",