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

mito_ai/tests/user/test_user.py

@@ -0,0 +1,120 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import os
+import requests
+import tempfile
+from unittest.mock import patch
+import pytest
+from mito_ai.tests.conftest import TOKEN
+
+
+@pytest.fixture
+def mock_user_json():
+    """Fixture that creates a temporary user.json file with test data"""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Create the .mito directory
+        mito_dir = os.path.join(temp_dir, ".mito")
+        os.makedirs(mito_dir, exist_ok=True)
+
+        # Create a user.json file with test data
+        user_json_path = os.path.join(mito_dir, "user.json")
+        user_data = {
+            "user_email": "test@mail.com",
+            "static_user_id": "test_user_123",
+        }
+
+        with open(user_json_path, "w") as f:
+            json.dump(user_data, f)
+
+        yield user_json_path
+
+
+# --- GET USER KEY ---
+
+
+def test_get_user_with_mocked_data_success(
+    jp_base_url: str, mock_user_json: str
+) -> None:
+    """Test successful GET user endpoint with mocked data"""
+    with patch("mito_ai.utils.db.USER_JSON_PATH", mock_user_json):
+        response = requests.get(
+            jp_base_url + f"/mito-ai/user/user_email",
+            headers={"Authorization": f"token {TOKEN}"},
+        )
+        assert response.status_code == 200
+
+        response_json = response.json()
+        assert response_json["key"] == "user_email"
+        assert response_json["value"] == "test@mail.com"
+
+
+def test_get_user_with_mocked_data_not_found(
+    jp_base_url: str, mock_user_json: str
+) -> None:
+    """Test GET user endpoint with mocked data for non-existent key"""
+    with patch("mito_ai.utils.db.USER_JSON_PATH", mock_user_json):
+        response = requests.get(
+            jp_base_url + "/mito-ai/user/non_existent_key",
+            headers={"Authorization": f"token {TOKEN}"},
+        )
+        assert response.status_code == 404
+
+        response_json = response.json()
+        assert (
+            response_json["error"] == "User field with key 'non_existent_key' not found"
+        )
+
+
+def test_get_user_with_no_auth(jp_base_url: str) -> None:
+    response = requests.get(
+        jp_base_url + f"/mito-ai/user/user_email",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_user_with_incorrect_auth(jp_base_url: str) -> None:
+    response = requests.get(
+        jp_base_url + f"/mito-ai/user/user_email",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+# --- PUT USER KEY ---
+
+
+def test_put_user_with_mocked_data_success(
+    jp_base_url: str, mock_user_json: str
+) -> None:
+    """Test successful PUT user endpoint with mocked data"""
+    with patch("mito_ai.utils.db.USER_JSON_PATH", mock_user_json):
+        response = requests.put(
+            jp_base_url + f"/mito-ai/user/user_email",
+            headers={"Authorization": f"token {TOKEN}"},
+            json={"value": "jdoe@mail.com"},
+        )
+        assert response.status_code == 200
+
+        response_json = response.json()
+        assert response_json["status"] == "success"
+        assert response_json["key"] == "user_email"
+        assert response_json["value"] == "jdoe@mail.com"
+
+
+def test_put_user_with_no_auth(jp_base_url: str) -> None:
+    response = requests.put(
+        jp_base_url + f"/mito-ai/user/user_email",
+        json={"value": "jdoe@mail.com"},
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_put_user_with_incorrect_auth(jp_base_url: str) -> None:
+    response = requests.put(
+        jp_base_url + f"/mito-ai/user/user_email",
+        headers={"Authorization": f"token incorrect-token"},
+        json={"value": "jdoe@mail.com"},
+    )
+    assert response.status_code == 403  # Forbidden

mito_ai/tests/utils/test_anthropic_utils.py

@@ -5,7 +5,7 @@ import pytest
 import anthropic
 from typing import List, Dict, Any, Tuple, Union, cast
 from anthropic.types import MessageParam, ToolUnionParam, ToolParam
-from mito_ai.utils.anthropic_utils import _prepare_anthropic_request_data_and_headers
+from mito_ai.utils.anthropic_utils import ANTHROPIC_TIMEOUT, _prepare_anthropic_request_data_and_headers
 from mito_ai.completions.models import MessageType
 from mito_ai.utils.schema import UJ_STATIC_USER_ID, UJ_USER_EMAIL
 from mito_ai.utils.db import get_user_field
@@ -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
 
@@ -52,7 +52,7 @@ def test_basic_request_preparation():
     )
 
     assert headers == {"Content-Type": "application/json"}
-    assert data["timeout"] == 30
+    assert data["timeout"] == ANTHROPIC_TIMEOUT
     assert data["max_retries"] == 1
     assert data["email"] == "test@example.com"
     assert data["user_id"] == "test_user_id"
@@ -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/user/handlers.py

@@ -0,0 +1,45 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import tornado
+from typing import Any, Optional
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.utils.db import get_user_field, set_user_field
+from mito_ai.utils.telemetry_utils import identify
+from mito_ai.utils.version_utils import is_pro
+
+
+class UserHandler(APIHandler):
+    """Handler for operations on a specific user"""
+
+    @tornado.web.authenticated
+    def get(self, key: str) -> None:
+        value: Optional[Any] = None
+
+        if key == "is_pro":
+            # Special case, since we don't store this key 
+            # in the user.json file.
+            value = str(is_pro())
+        else: 
+            value = get_user_field(key)
+
+        if value is None:
+            self.set_status(404)
+            self.finish(json.dumps({"error": f"User field with key '{key}' not found"}))
+        else:
+            self.finish(json.dumps({"key": key, "value": value}))
+
+    @tornado.web.authenticated
+    def put(self, key: str) -> None:
+        data = json.loads(self.request.body)
+        if "value" not in data:
+            self.set_status(400)
+            self.finish(json.dumps({"error": "Value is required"}))
+            return
+
+        set_user_field(key, data["value"])
+        identify()  # Log the new user
+        self.finish(
+            json.dumps({"status": "success", "key": key, "value": data["value"]})
+        )

mito_ai/user/urls.py

@@ -0,0 +1,21 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Any, List, Tuple
+from jupyter_server.utils import url_path_join
+from mito_ai.user.handlers import UserHandler
+
+
+def get_user_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all user related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    return [
+        (url_path_join(BASE_URL, "user/(.*)"), UserHandler, {}),
+    ]

mito_ai/utils/anthropic_utils.py

@@ -1,42 +1,36 @@
 # 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 openai.types.chat import ChatCompletionMessageParam
-from mito_ai.completions.models import AgentResponse, MessageType, ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem
+from typing import Any, Dict, List, Optional, Union, AsyncGenerator, Tuple, Callable
+from anthropic.types import MessageParam, TextBlockParam, 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 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
 __user_id: Optional[str] = None
 
-timeout = 30
+ANTHROPIC_TIMEOUT = 60
 max_retries = 1
-INLINE_COMPLETION_MODEL = "claude-3-5-haiku-latest"
+
+FAST_ANTHROPIC_MODEL = "claude-3-5-haiku-latest"
 
 def _prepare_anthropic_request_data_and_headers(
     model: Union[str, None],
     max_tokens: int,
     temperature: float,
-    system: Union[str, anthropic.NotGiven],
+    system: Union[str, List[TextBlockParam], anthropic.Omit],
     messages: List[MessageParam],
     message_type: MessageType,
     tools: Optional[List[ToolUnionParam]],
     tool_choice: Optional[dict],
     stream: Optional[bool]
 ) -> Tuple[Dict[str, Any], Dict[str, str]]:
-    check_mito_server_quota(message_type)
+    
     global __user_email, __user_id
     if __user_email is None:
         __user_email = get_user_field(UJ_USER_EMAIL)
@@ -47,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
@@ -60,7 +55,7 @@ def _prepare_anthropic_request_data_and_headers(
         inner_data["stream"] = stream
     # Compose the outer data dict
     data = {
-        "timeout": timeout,
+        "timeout": ANTHROPIC_TIMEOUT,
         "max_retries": max_retries,
         "email": __user_email,
         "user_id": __user_id,
@@ -73,42 +68,31 @@ 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],
     message_type: MessageType
-) -> Message:
+) -> str:
     data, headers = _prepare_anthropic_request_data_and_headers(
         model, max_tokens, temperature, system, messages, message_type, tools, tool_choice, None
     )
-    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
-    start_time = time.time()
-    try:
-        res = await http_client.fetch(
-            MITO_ANTHROPIC_URL,
-            method="POST",
-            headers=headers,
-            body=json.dumps(data),
-            request_timeout=http_client_timeout
-        )
-        print(f"Anthropic request completed in {time.time() - start_time:.2f} seconds")
-    except Exception as e:
-        print(f"Anthropic request failed after {time.time() - start_time:.2f} seconds with error: {str(e)}")
-        raise
-    finally:
-        http_client.close()
-    content = json.loads(res.body)
-    # If the response is wrapped in a 'data' field, extract it
-    if isinstance(content, dict) and "data" in content:
-        return cast(Message, content["data"])
-    return cast(Message, content)
+    
+    return await get_response_from_mito_server(
+        MITO_ANTHROPIC_URL, 
+        headers, 
+        data, 
+        ANTHROPIC_TIMEOUT, 
+        max_retries, 
+        message_type, 
+        provider_name="Claude"
+    )
 
 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, List[TextBlockParam], anthropic.Omit],
     messages: List[MessageParam],
     stream: bool,
     message_type: MessageType,
@@ -118,91 +102,46 @@ async def stream_anthropic_completion_from_mito_server(
     data, headers = _prepare_anthropic_request_data_and_headers(
         model, max_tokens, temperature, system, messages, message_type, None, None, stream
     )
-    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
-    start_time = time.time()
-    chunk_queue: asyncio.Queue[str] = asyncio.Queue()
-    fetch_complete = False
-    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 Anthropic streaming chunk: {str(e)}")
-    fetch_future = None
-    try:
-        fetch_future = http_client.fetch(
-            MITO_ANTHROPIC_URL,
-            method="POST",
-            headers=headers,
-            body=json.dumps(data),
-            request_timeout=http_client_timeout,
-            streaming_callback=chunk_callback
-        )
-        async def wait_for_fetch() -> None:
-            try:
-                await fetch_future
-                nonlocal fetch_complete
-                fetch_complete = True
-                print("Anthropic fetch completed")
-            except Exception as e:
-                print(f"Error in Anthropic fetch: {str(e)}")
-                raise
-        fetch_task = asyncio.create_task(wait_for_fetch())
-        while not (fetch_complete and chunk_queue.empty()):
-            try:
-                chunk = await asyncio.wait_for(chunk_queue.get(), timeout=0.1)
-                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:
-                if fetch_complete and chunk_queue.empty():
-                    break
-                continue
-        print(f"\nAnthropic stream completed in {time.time() - start_time:.2f} seconds")
-        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"\nAnthropic stream failed after {time.time() - start_time:.2f} seconds with error: {str(e)}")
-        if fetch_future:
-            try:
-                await fetch_future
-            except Exception:
-                pass
-        raise
-    finally:
-        http_client.close()
-
+    
+    # Use the unified streaming function
+    # If the reply_fn and message_id are empty, this function still handles those requests. This is particularly needed for the streamlit dashboard functionality
+    actual_reply_fn = reply_fn if reply_fn is not None else (lambda x: None)
+    actual_message_id = message_id if message_id is not None else ""
+    async for chunk in stream_response_from_mito_server(
+        url=MITO_ANTHROPIC_URL,
+        headers=headers,
+        data=data,
+        timeout=ANTHROPIC_TIMEOUT,
+        max_retries=max_retries,
+        message_type=message_type,
+        reply_fn=actual_reply_fn,
+        message_id=actual_message_id,
+        chunk_processor=None,
+        provider_name="Claude",
+    ):
+        yield chunk
+    
+    
 def get_anthropic_completion_function_params(
+    message_type: MessageType,
     model: str,
     messages: List[MessageParam],
     max_tokens: int,
-    system: Union[str, anthropic.NotGiven],
+    system: Union[str, List[TextBlockParam], anthropic.Omit],
     temperature: float = 0.0,
     tools: Optional[List[ToolUnionParam]] = None,
     tool_choice: Optional[dict] = None,
     stream: Optional[bool] = None,
     response_format_info: Optional[ResponseFormatInfo] = None,
-) -> Dict[str, Any]:
+) -> Dict[str, Any]:    
     """
     Build the provider_data dict for Anthropic completions, mirroring the OpenAI approach.
     Only includes fields needed for the Anthropic API.
     """
+    
+    message_requires_fast_model = does_message_require_fast_model(message_type)
+    model = FAST_ANTHROPIC_MODEL if message_requires_fast_model else model
+    
     provider_data = {
         "model": model,
         "max_tokens": max_tokens,
@@ -210,11 +149,6 @@ def get_anthropic_completion_function_params(
         "messages": messages,
         "system": system,
     }
-    if response_format_info is not None:
-        # TODO: This should not be here.. the model is resolved in the anthropic client. 
-        # This also means that chat is using the fast model... 
-        # I bet the same bug exists in gemini...
-        provider_data["model"] = INLINE_COMPLETION_MODEL
     if tools:
         provider_data["tools"] = tools
     if response_format_info and response_format_info.name == "agent_response":

mito_ai/utils/create.py

@@ -36,6 +36,17 @@ def is_user_json_exists_and_valid_json() -> bool:
     except:
         return False
 
+def get_temp_user_id() -> Optional[str]:
+    """
+    Looks for a temporary user ID, generated by the desktop app.
+    """
+    temp_user_id_path = os.path.join(MITO_FOLDER, 'temp_user_id.txt')
+
+    if os.path.exists(temp_user_id_path):
+        with open(temp_user_id_path, 'r') as f:
+            return f.read()
+
+    return None
 
 def try_create_user_json_file() -> None:
     
@@ -50,7 +61,12 @@ def try_create_user_json_file() -> None:
         with open(USER_JSON_PATH, 'w+') as f:
             f.write(json.dumps(USER_JSON_DEFAULT))
 
-        # Then, we take special care to put all the testing/CI environments 
+        # Next, look for a temp user id
+        temp_user_id = get_temp_user_id()
+        if temp_user_id:
+            set_user_field(UJ_STATIC_USER_ID, temp_user_id)
+
+        # Finally, we take special care to put all the testing/CI environments 
         # (e.g. Github actions) under one ID and email
         if is_running_test():
             set_user_field(UJ_STATIC_USER_ID, GITHUB_ACTION_ID)

mito_ai/utils/error_classes.py

@@ -0,0 +1,42 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.app_deploy.models import AppDeployError
+
+class MitoAppError(Exception):
+    """Exception raised for custom error in the application."""
+
+    def __init__(self, message: str, error_code: int) -> None:
+        super().__init__(message)
+        self.message = message
+        self.error_code = error_code
+
+class StreamlitPreviewError(MitoAppError):
+    def __str__(self) -> str:
+        return f"[PreviewError]: {self.message} (Error Code: {self.error_code})"
+
+class StreamlitConversionError(MitoAppError):
+    def __str__(self) -> str:
+        return f"[ConversionError]: {self.message} (Error Code: {self.error_code})"
+
+
+class StreamlitDeploymentError(MitoAppError):
+    """Raised when a deployment operation fails."""
+
+    def __init__(self, error: AppDeployError):
+        self.error = error
+        self.error_type = error.error_type
+        self.message_id = getattr(error, "message_id", "ErrorMessageID")
+        self.error_code = getattr(error, "error_code", 500)
+        self.hint = getattr(error, "hint", "")
+        self.traceback = getattr(error, "traceback", "")
+        self.error_type = getattr(error, "error_type", "Error")
+        self.message = error.message
+        print(f"self_message: {self.message}")
+        super().__init__(self.message, self.error_code)
+
+    def __str__(self) -> str:
+        base = f"[DeploymentError]: {self.message} (Error Code: {self.error_code})"
+        if self.hint:
+            base += f"\nHint: {self.hint}"
+        return base