mito-ai 0.1.33__py3-none-any.whl → 0.1.49__py3-none-any.whl

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, 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
 
@@ -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],
@@ -137,8 +106,8 @@ async def stream_ai_completion_from_mito_server(
     timeout: int,
     max_retries: int,
     message_type: MessageType,
-    reply_fn: Optional[Callable[[Union[CompletionReply, CompletionStreamChunk]], None]] = None,
-    message_id: Optional[str] = None,
+    reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+    message_id: str,
 ) -> AsyncGenerator[str, None]:
     """
     Stream AI completions from the Mito server.
@@ -158,7 +127,7 @@ async def stream_ai_completion_from_mito_server(
     Yields:
         Chunks of text from the streaming response
     """
-    # ===== STEP 1: Prepare request data and headers =====
+    # Prepare request data and headers
     data, headers = _prepare_request_data_and_headers(
         last_message_content, 
         ai_completion_data, 
@@ -167,118 +136,36 @@ async def stream_ai_completion_from_mito_server(
         message_type
     )
     
-    # ===== STEP 2: Create HTTP client with appropriate timeout settings =====
-    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
-    
-    # ===== STEP 3: 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 streaming chunk: {str(e)}")
-    
-    # ===== STEP 4: Execute the streaming request =====
-    fetch_future = None
-    try:
-        # Use fetch with streaming_callback to handle streaming responses.
-        # The streaming_callback is not sent as part of the POST request.
-        # It's a parameter for the Tornado AsyncHTTPClient.fetch() method that specifies
-        # how to handle incoming data chunks as they arrive from the server.
-        # When the server sends data in chunks, this callback function is called each time
-        # a new chunk arrives, allowing for immediate processing without waiting for the
-        # entire response to complete.
-        fetch_future = http_client.fetch(
-            MITO_OPENAI_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("Fetch completed")
-            except Exception as e:
-                print(f"Error in fetch: {str(e)}")
-                raise
-        
-        # Start the task to wait for fetch completion
-        fetch_task = asyncio.create_task(wait_for_fetch())
-        
-        # ===== STEP 5: Yield chunks as they arrive =====
-        while not (fetch_complete and chunk_queue.empty()):
-            try:
-                # Wait for a chunk with a timeout. By setting the timeout, we 1. prevent deadlocks
-                # which could happen if fetch_complete has not been set to true yet, and 2. it enables
-                # periodic checking if the queue has a new chunk.
-                chunk = await asyncio.wait_for(chunk_queue.get(), timeout=0.1)
-                
-                # If reply_fn is provided, send the chunk directly to the frontend
-                if reply_fn and message_id:
-                    reply_fn(CompletionStreamChunk(
-                        parent_id=message_id,
-                        chunk=CompletionItem(
-                            content=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"\nStream completed in {time.time() - start_time:.2f} seconds")
-        
-        # Send a final chunk to indicate completion if reply_fn is provided
-        if reply_fn and message_id:
-            reply_fn(CompletionStreamChunk(
-                parent_id=message_id,
-                chunk=CompletionItem(
-                    content="",
-                    isIncomplete=False,
-                    token=message_id,
-                ),
-                done=True,
-            ))
-    except Exception as e:
-        print(f"\nStream 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:
-        # ===== STEP 6: Clean up resources =====
-        http_client.close()
+    # 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,
@@ -309,9 +196,5 @@ def get_open_ai_completion_function_params(
                 "strict": True
             }
         }
-    
-    # o3-mini will error if we try setting the temperature
-    if model == "gpt-4o-mini":
-        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/server_limits.py

@@ -28,7 +28,7 @@ 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] = 50
+OS_MONTHLY_AI_COMPLETIONS_LIMIT: Final[int] = 150
 
 # Monthly autocomplete limit for free tier users
 OS_MONTHLY_AUTOCOMPLETE_LIMIT: Final[int] = 5000

mito_ai/utils/telemetry_utils.py

@@ -3,7 +3,7 @@
 
 import json
 import os
-from typing import Any, Dict, Literal, Optional
+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
@@ -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':
@@ -104,7 +104,7 @@ def telemetry_turned_on(key_type: Optional[str] = None) -> bool:
     
     return bool(telemetry)
 
-def identify(key_type: Optional[str] = None) -> None:
+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.
@@ -125,6 +125,11 @@ def identify(key_type: Optional[str] = None) -> None:
         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.
@@ -174,7 +179,6 @@ def log(
     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
@@ -186,7 +190,24 @@ def log(
 
     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):
@@ -346,3 +367,114 @@ def log_db_connection_error(connection_type: str, error_message: str) -> None:
             "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
+        }
+    )