mito-ai 0.1.43__py3-none-any.whl → 0.1.45__py3-none-any.whl

mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py

@@ -87,21 +87,6 @@ df=pd.read_csv('data.csv')
             validator = StreamlitValidator() 
             errors = validator.get_runtime_errors(app_code, app_path)
             assert errors is None
-            
-            
-    @patch('subprocess.Popen')
-    def test_cleanup_with_process(self, mock_popen):
-        """Test cleanup with running process"""
-        validator = StreamlitValidator()
-        validator.temp_dir = "/tmp/test_dir"
-        
-        # Mock directory exists
-        with patch('os.path.exists', return_value=True):
-            with patch('shutil.rmtree') as mock_rmtree:
-                validator.cleanup()
-                
-                mock_rmtree.assert_called_once()
-
 
     @pytest.mark.parametrize("app_code,expected_has_validation_error,expected_error_message", [
         ("x=5", False, ""),

mito_ai/tests/streamlit_preview/test_streamlit_preview_handler.py

@@ -0,0 +1,88 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import os
+import tempfile
+from unittest.mock import patch, AsyncMock, MagicMock
+from mito_ai.streamlit_preview.utils import ensure_app_exists
+
+
+class TestEnsureAppExists:
+    """Test cases for the ensure_app_exists function."""
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "app_exists,streamlit_handler_success,expected_success,expected_error,streamlit_handler_called,streamlit_handler_return",
+        [
+            # Test case 1: App exists, should use existing file
+            (
+                True,  # app_exists
+                True,  # streamlit_handler_success (not relevant)
+                True,  # expected_success
+                "",    # expected_error
+                False, # streamlit_handler_called
+                None,  # streamlit_handler_return (not used)
+            ),
+            # Test case 2: App doesn't exist, streamlit_handler succeeds
+            (
+                False,  # app_exists
+                True,   # streamlit_handler_success
+                True,   # expected_success
+                "",     # expected_error
+                True,   # streamlit_handler_called
+                (True, "/path/to/app.py", "Success"),  # streamlit_handler_return
+            )
+        ],
+        ids=[
+            "app_exists_uses_existing_file",
+            "app_does_not_exist_generates_new_one_success",
+        ]
+    )
+    async def test_ensure_app_exists(
+        self,
+        app_exists,
+        streamlit_handler_success,
+        expected_success,
+        expected_error,
+        streamlit_handler_called,
+        streamlit_handler_return,
+    ):
+        """Test ensure_app_exists function with various scenarios."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            notebook_path = os.path.join(temp_dir, "test_notebook.ipynb")
+            
+            # Set up app_path based on whether app exists
+            app_path = os.path.join(temp_dir, "app.py") if app_exists else None
+            
+            # Create app.py file if it should exist
+            if app_exists:
+                assert app_path is not None  # Type assertion for mypy
+                with open(app_path, "w") as f:
+                    f.write("import streamlit as st\nst.write('Hello World')")
+            
+            # Mock get_app_path to return the appropriate value
+            with patch('mito_ai.streamlit_preview.utils.get_app_path') as mock_get_app_path:
+                mock_get_app_path.return_value = app_path
+                
+                # Mock streamlit_handler
+                with patch('mito_ai.streamlit_preview.utils.streamlit_handler') as mock_streamlit_handler:
+                    if streamlit_handler_return is not None:
+                        mock_streamlit_handler.return_value = streamlit_handler_return
+                    
+                    success, error_msg = await ensure_app_exists(notebook_path)
+                    
+                    # Assertions
+                    assert success == expected_success
+                    assert error_msg == expected_error
+                    
+                    # Verify get_app_path was called with the correct directory
+                    mock_get_app_path.assert_called_once_with(temp_dir)
+                    
+                    # Verify streamlit_handler was called or not called as expected
+                    if streamlit_handler_called:
+                        mock_streamlit_handler.assert_called_once_with(notebook_path)
+                    else:
+                        mock_streamlit_handler.assert_not_called()
+
+    

mito_ai/tests/streamlit_preview/test_streamlit_preview_manager.py

@@ -9,6 +9,7 @@ import shutil
 import subprocess
 import threading
 import requests
+import json
 from unittest.mock import Mock, patch, MagicMock
 from typing import Any
 
@@ -17,6 +18,8 @@ from mito_ai.streamlit_preview.manager import (
     PreviewProcess, 
     get_preview_manager
 )
+from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
+from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
 
 
 class TestStreamlitPreviewManager:
@@ -299,4 +302,4 @@ st.write("Hello, World!")
             # Verify they're gone
             for preview_id in preview_ids:
                 assert manager.get_preview(preview_id) is None
-    
+

mito_ai/tests/utils/test_anthropic_utils.py

@@ -35,7 +35,7 @@ def test_basic_request_preparation():
     max_tokens = 100
     temperature = 0.7
     # Use NotGiven to ensure system is not included in inner_data
-    system = anthropic.NotGiven()
+    system = anthropic.Omit()
     messages: List[MessageParam] = [{"role": "user", "content": "Hello"}]
     message_type = MessageType.CHAT
 
@@ -106,7 +106,7 @@ def test_tools_and_tool_choice():
         model="claude-3-sonnet",
         max_tokens=100,
         temperature=0.7,
-        system=anthropic.NotGiven(),
+        system=anthropic.Omit(),
         messages=[{"role": "user", "content": "Hello"}],
         message_type=MessageType.CHAT,
         tools=tools,
@@ -124,7 +124,7 @@ def test_stream_parameter():
         model="claude-3-sonnet",
         max_tokens=100,
         temperature=0.7,
-        system=anthropic.NotGiven(),
+        system=anthropic.Omit(),
         messages=[{"role": "user", "content": "Hello"}],
         message_type=MessageType.CHAT,
         tools=None,
@@ -150,7 +150,7 @@ def test_missing_user_info(monkeypatch):
         model="claude-3-sonnet",
         max_tokens=100,
         temperature=0.7,
-        system=anthropic.NotGiven(),
+        system=anthropic.Omit(),
         messages=[{"role": "user", "content": "Hello"}],
         message_type=MessageType.CHAT,
         tools=None,

mito_ai/utils/anthropic_utils.py

@@ -1,23 +1,14 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-import asyncio
-import json
-import time
 import anthropic
-from typing import Any, Dict, List, Optional, Union, AsyncGenerator, Tuple, Callable, cast
-
-from anthropic.types import MessageParam, Message, TextBlock, ToolUnionParam
+from typing import Any, Dict, List, Optional, Union, AsyncGenerator, Tuple, Callable
+from anthropic.types import MessageParam, ToolUnionParam
 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 openai.types.chat import ChatCompletionMessageParam
-from mito_ai.completions.models import AgentResponse, MessageType, ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem
+from mito_ai.completions.models import AgentResponse, 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.utils import is_running_test
-from mito_ai.utils.server_limits import check_mito_server_quota
-from .utils import _create_http_client
-from tornado.httpclient import AsyncHTTPClient
 from mito_ai.constants import MITO_ANTHROPIC_URL
 
 __user_email: Optional[str] = None
@@ -32,7 +23,7 @@ def _prepare_anthropic_request_data_and_headers(
     model: Union[str, None],
     max_tokens: int,
     temperature: float,
-    system: Union[str, anthropic.NotGiven],
+    system: Union[str, anthropic.Omit],
     messages: List[MessageParam],
     message_type: MessageType,
     tools: Optional[List[ToolUnionParam]],
@@ -50,10 +41,11 @@ def _prepare_anthropic_request_data_and_headers(
         "model": model,
         "max_tokens": max_tokens,
         "temperature": temperature,
-        "messages": messages
+        "messages": messages,
+        "betas": ["context-1m-2025-08-07"]
     }
-    # Add system to inner_data only if it is not anthropic.NotGiven
-    if not isinstance(system, anthropic.NotGiven):
+    # Add system to inner_data only if it is not anthropic.Omit
+    if not isinstance(system, anthropic.Omit):
         inner_data["system"] = system
     if tools:
         inner_data["tools"] = tools
@@ -76,7 +68,7 @@ async def get_anthropic_completion_from_mito_server(
     model: Union[str, None],
     max_tokens: int,
     temperature: float,
-    system: Union[str, anthropic.NotGiven],
+    system: Union[str, anthropic.Omit],
     messages: List[MessageParam],
     tools: Optional[List[ToolUnionParam]],
     tool_choice: Optional[dict],
@@ -100,7 +92,7 @@ async def stream_anthropic_completion_from_mito_server(
     model: Union[str, None],
     max_tokens: int,
     temperature: float,
-    system: Union[str, anthropic.NotGiven],
+    system: Union[str, anthropic.Omit],
     messages: List[MessageParam],
     stream: bool,
     message_type: MessageType,
@@ -133,7 +125,7 @@ def get_anthropic_completion_function_params(
     model: str,
     messages: List[MessageParam],
     max_tokens: int,
-    system: Union[str, anthropic.NotGiven],
+    system: Union[str, anthropic.Omit],
     temperature: float = 0.0,
     tools: Optional[List[ToolUnionParam]] = None,
     tool_choice: Optional[dict] = None,

mito_ai/utils/telemetry_utils.py

@@ -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.
@@ -362,11 +367,16 @@ def log_file_upload_attempt(
 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'], 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_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}, thread_id=thread_id, 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_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}, 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}, key_type=key_type)