mito-ai 0.1.44__py3-none-any.whl → 0.1.46__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

@@ -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/user/handlers.py

@@ -0,0 +1,33 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import tornado
+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
+
+
+class UserHandler(APIHandler):
+    """Handler for operations on a specific user"""
+
+    @tornado.web.authenticated
+    def get(self, key: str) -> None:
+        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,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, 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 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, List[TextBlockParam], 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, List[TextBlockParam], anthropic.Omit],
     messages: List[MessageParam],
     stream: bool,
     message_type: MessageType,
@@ -110,6 +102,7 @@ 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
     )
+    
     # 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)
@@ -127,19 +120,20 @@ async def stream_anthropic_completion_from_mito_server(
         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.

mito_ai/utils/message_history_utils.py

@@ -3,6 +3,7 @@
 
 import re
 from typing import List
+from mito_ai.constants import MESSAGE_HISTORY_TRIM_THRESHOLD
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.prompt_builders.prompt_constants import (
     ACTIVE_CELL_ID_SECTION_HEADING,
@@ -43,18 +44,18 @@ def trim_sections_from_message_content(content: str) -> str:
     return content
 
 
-def trim_old_messages(messages: List[ChatCompletionMessageParam], keep_recent: int = 3) -> List[ChatCompletionMessageParam]:
+def trim_old_messages(messages: List[ChatCompletionMessageParam]) -> List[ChatCompletionMessageParam]:
     """
     Trims metadata sections from messages that are older than the specified number of recent messages.
     We do this in order to reduce the token count of the messages, which helps us stay under the token limit for the LLM.
     """
-    if len(messages) <= keep_recent:
+    if len(messages) <= MESSAGE_HISTORY_TRIM_THRESHOLD:
         return messages
         
     # Process all messages except the keep_recent most recent ones. 
     # Only trim user messages, which is where this metadata lives. 
     # We want to not edit the system messages, as they contain important information / examples.
-    for i in range(len(messages) - keep_recent):
+    for i in range(len(messages) - MESSAGE_HISTORY_TRIM_THRESHOLD):
         content = messages[i].get("content")
         
         is_user_message = messages[i].get("role") == "user"

mito_ai/utils/telemetry_utils.py

@@ -386,10 +386,13 @@ def log_mito_server_free_tier_limit_reached(key_type: Literal['mito_server_key',
 # Streamlit Conversion
 #################################
 
-def log_streamlit_app_creation_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType) -> None:
+def log_streamlit_app_creation_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, edit_prompt: str) -> None:
     log(
         "mito_ai_streamlit_app_creation_success", 
-        key_type=key_type
+        key_type=key_type,
+        params={
+            "edit_prompt": edit_prompt,
+        }
     )
     
 def log_streamlit_app_creation_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str) -> None:
@@ -401,12 +404,12 @@ def log_streamlit_app_creation_retry(key_type: Literal['mito_server_key', 'user_
         key_type=key_type
     )
     
-def log_streamlit_app_creation_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str) -> None:
+def log_streamlit_app_creation_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str, edit_prompt: str) -> None:
     log(
         "mito_ai_streamlit_app_creation_error", 
         params={
             "error_message": error,
+            "edit_prompt": edit_prompt,
         },
         key_type=key_type
     )
-