mito-ai 0.1.50__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/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+

mito_ai/tests/utils/test_anthropic_utils.py

@@ -0,0 +1,162 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+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 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
+
+
+# Mock the get_user_field and set_user_field functions
+@pytest.fixture(autouse=True)
+def mock_user_functions(monkeypatch):
+    def mock_get_field(field):
+        if field == UJ_USER_EMAIL:
+            return "test@example.com"
+        elif field == UJ_STATIC_USER_ID:
+            return "test_user_id"
+        return None
+
+    def mock_set_field(field, value):
+        # Do nothing in tests
+        pass
+
+    monkeypatch.setattr("mito_ai.utils.anthropic_utils.get_user_field", mock_get_field)
+    monkeypatch.setattr("mito_ai.utils.server_limits.set_user_field", mock_set_field)
+
+
+def test_basic_request_preparation():
+    """Test basic request preparation with minimal parameters"""
+    model = "claude-3-sonnet"
+    max_tokens = 100
+    temperature = 0.7
+    # Use NotGiven to ensure system is not included in inner_data
+    system = anthropic.Omit()
+    messages: List[MessageParam] = [{"role": "user", "content": "Hello"}]
+    message_type = MessageType.CHAT
+
+    data, headers = _prepare_anthropic_request_data_and_headers(
+        model=model,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        system=system,
+        messages=messages,
+        message_type=message_type,
+        tools=None,
+        tool_choice=None,
+        stream=None
+    )
+
+    assert headers == {"Content-Type": "application/json"}
+    assert data["timeout"] == ANTHROPIC_TIMEOUT
+    assert data["max_retries"] == 1
+    assert data["email"] == "test@example.com"
+    assert data["user_id"] == "test_user_id"
+
+    inner_data = data["data"]
+    assert inner_data["model"] == model
+    assert inner_data["max_tokens"] == max_tokens
+    assert inner_data["temperature"] == temperature
+    assert inner_data["messages"] == messages
+    # When system is NotGiven, it should not be included in inner_data
+    assert "system" not in inner_data
+
+
+def test_system_message_handling():
+    """Test handling of system message when provided"""
+    system = "You are a helpful assistant"
+    messages: List[MessageParam] = [{"role": "user", "content": "Hello"}]
+
+    data, _ = _prepare_anthropic_request_data_and_headers(
+        model="claude-3-sonnet",
+        max_tokens=100,
+        temperature=0.7,
+        system=system,
+        messages=messages,
+        message_type=MessageType.CHAT,
+        tools=None,
+        tool_choice=None,
+        stream=None
+    )
+
+    assert data["data"]["system"] == system
+
+
+def test_tools_and_tool_choice():
+    """Test handling of tools and tool_choice parameters"""
+    tools = cast(List[ToolUnionParam], [{
+        "type": "function",
+        "function": {
+            "name": "test_function",
+            "description": "A test function",
+            "parameters": {
+                "type": "object",
+                "properties": {},
+                "required": []
+            }
+        }
+    }])
+    tool_choice: Dict[str, Any] = {"type": "function", "function": {"name": "test_function"}}
+
+    data, _ = _prepare_anthropic_request_data_and_headers(
+        model="claude-3-sonnet",
+        max_tokens=100,
+        temperature=0.7,
+        system=anthropic.Omit(),
+        messages=[{"role": "user", "content": "Hello"}],
+        message_type=MessageType.CHAT,
+        tools=tools,
+        tool_choice=tool_choice,
+        stream=None
+    )
+
+    assert data["data"]["tools"] == tools
+    assert data["data"]["tool_choice"] == tool_choice
+
+
+def test_stream_parameter():
+    """Test handling of stream parameter"""
+    data, _ = _prepare_anthropic_request_data_and_headers(
+        model="claude-3-sonnet",
+        max_tokens=100,
+        temperature=0.7,
+        system=anthropic.Omit(),
+        messages=[{"role": "user", "content": "Hello"}],
+        message_type=MessageType.CHAT,
+        tools=None,
+        tool_choice=None,
+        stream=True
+    )
+
+    assert data["data"]["stream"] is True
+
+
+def test_missing_user_info(monkeypatch):
+    """Test behavior when user email and ID are not available"""
+
+    def mock_get_field(field):
+        return None
+
+    # Override the autouse fixture for this specific test
+    monkeypatch.setattr("mito_ai.utils.anthropic_utils.get_user_field", mock_get_field)
+    monkeypatch.setattr("mito_ai.utils.anthropic_utils.__user_email", None)
+    monkeypatch.setattr("mito_ai.utils.anthropic_utils.__user_id", None)
+
+    data, _ = _prepare_anthropic_request_data_and_headers(
+        model="claude-3-sonnet",
+        max_tokens=100,
+        temperature=0.7,
+        system=anthropic.Omit(),
+        messages=[{"role": "user", "content": "Hello"}],
+        message_type=MessageType.CHAT,
+        tools=None,
+        tool_choice=None,
+        stream=None
+    )
+
+    assert data["email"] is None
+    assert data["user_id"] is None 

mito_ai/tests/utils/test_gemini_utils.py

@@ -0,0 +1,98 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+from mito_ai.utils.gemini_utils import _prepare_gemini_request_data_and_headers
+from mito_ai.completions.models import MessageType
+
+TEST_CONTENTS = [
+    {'role': 'system', 'content': 'You are Mito Data Copilot, an AI assistant for Jupyter.'},
+    {'role': 'user', 'content': 'Help me complete the following task. print 10'},
+    {'role': 'assistant', 'content': 'python\nprint(10)\n\nPrinted the number 10 '},
+    {'role': 'user', 'content': 'Update to print 11'}
+]
+
+
+def test_basic_request_preparation():
+    """Test basic request preparation with minimal parameters."""
+    model = "gemini-pro"
+    message_type = MessageType.CHAT
+
+    data, headers = _prepare_gemini_request_data_and_headers(
+        model=model,
+        contents=TEST_CONTENTS,
+        message_type=message_type
+    )
+
+    assert headers == {"Content-Type": "application/json"}
+    assert data["timeout"] == 30
+    assert data["max_retries"] == 1
+    assert data["data"]["model"] == model
+    assert data["data"]["contents"] == TEST_CONTENTS
+    assert data["data"]["message_type"] == message_type.value
+
+
+def test_request_with_config():
+    """Test request preparation with additional config parameters."""
+    model = "gemini-pro"
+    message_type = MessageType.CHAT
+    config = {
+        "temperature": 0.7,
+        "max_tokens": 100
+    }
+
+    data, headers = _prepare_gemini_request_data_and_headers(
+        model=model,
+        contents=TEST_CONTENTS,
+        message_type=message_type,
+        config=config
+    )
+
+    assert data["data"]["config"] == config
+
+
+def test_request_with_response_format():
+    """Test request preparation with response format info."""
+
+    class TestFormat:
+        name = "test_format"
+        format = "json"
+
+    model = "gemini-pro"
+    message_type = MessageType.CHAT
+    response_format_info = TestFormat()
+
+    data, headers = _prepare_gemini_request_data_and_headers(
+        model=model,
+        contents=TEST_CONTENTS,
+        message_type=message_type,
+        response_format_info=response_format_info
+    )
+
+    expected_format = {
+        "name": "test_format",
+        "format": "json"
+    }
+    assert data["data"]["response_format_info"] == '{"name": "test_format", "format": "json"}'
+
+
+def test_request_with_complex_config():
+    """Test request preparation with complex nested config."""
+    model = "gemini-pro"
+    message_type = MessageType.CHAT
+    config = {
+        "temperature": 0.7,
+        "nested": {
+            "key": "value",
+            "array": [1, 2, 3]
+        }
+    }
+
+    data, headers = _prepare_gemini_request_data_and_headers(
+        model=model,
+        contents=TEST_CONTENTS,
+        message_type=message_type,
+        config=config
+    )
+
+    assert data["data"]["config"] == config

mito_ai/tests/version_check_test.py

@@ -0,0 +1,169 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import time
+import unittest
+from unittest.mock import patch, MagicMock, PropertyMock
+from typing import Any, Optional, Tuple, List, Dict, Callable, Union, cast
+
+import sys
+import os
+
+# Add the project root to path
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+from mito_ai.version_check import VersionCheckHandler
+
+class TestVersionCheckHandler(unittest.TestCase):
+    
+    def setUp(self) -> None:
+        # Clear cache before each test
+        if hasattr(VersionCheckHandler, '_get_latest_version'):
+            VersionCheckHandler._get_latest_version.cache_clear()
+    
+    @patch("mito_ai.version_check.requests.get")
+    def test_get_latest_version_method(self, mock_requests_get: MagicMock) -> None:
+        # Mock successful response
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "info": {"version": "1.2.3"}
+        }
+        mock_requests_get.return_value = mock_response
+        
+        # Call the method
+        version, cache_time = VersionCheckHandler._get_latest_version()
+        
+        # Verify results
+        self.assertEqual(version, "1.2.3")
+        self.assertIsInstance(cache_time, float)
+        
+        # Call it again to use cache
+        version2, cache_time2 = VersionCheckHandler._get_latest_version()
+        
+        # Should return same results without calling requests.get again
+        self.assertEqual(version2, "1.2.3")
+        self.assertEqual(cache_time, cache_time2)
+        mock_requests_get.assert_called_once()
+    
+    @patch("mito_ai.version_check.requests.get")
+    def test_get_latest_version_error(self, mock_requests_get: MagicMock) -> None:
+        # Mock error response
+        mock_requests_get.side_effect = Exception("Connection error")
+        
+        try:
+            # The actual implementation lets the exception propagate when called directly
+            # So we need to catch it in our test
+            version, cache_time = VersionCheckHandler._get_latest_version()
+            self.fail("Expected an exception but none was raised")
+        except Exception as e:
+            # Just verify the exception was raised
+            self.assertEqual(str(e), "Connection error")
+    
+    @patch("mito_ai.version_check.__version__", "1.0.0")
+    @patch("mito_ai.version_check.requests.get")
+    def test_successful_version_fetch(self, mock_requests_get: MagicMock) -> None:
+        # Mock the responses
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "info": {"version": "1.1.0"}
+        }
+        mock_requests_get.return_value = mock_response
+        
+        # Create handler instance
+        handler = self._create_mocked_handler()
+        
+        # Call the get method directly
+        handler.get()
+        
+        # Get the response body and verify
+        response_body = json.loads(handler._write_buffer[0])
+        self.assertEqual(response_body["current_version"], "1.0.0")
+        self.assertEqual(response_body["latest_version"], "1.1.0")
+        self.assertIn("cache_age_seconds", response_body)
+        
+        # Verify the requests were made correctly
+        mock_requests_get.assert_called_once_with(
+            "https://pypi.org/pypi/mito-ai/json", timeout=3
+        )
+    
+    @patch("mito_ai.version_check.__version__", "1.0.0")
+    @patch("mito_ai.version_check.requests.get")
+    def test_cache_behavior(self, mock_requests_get: MagicMock) -> None:
+        # Mock the responses
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "info": {"version": "1.1.0"}
+        }
+        mock_requests_get.return_value = mock_response
+        
+        # First request should make an external API call
+        handler1 = self._create_mocked_handler()
+        handler1.get()
+        
+        # Second request should use the cache
+        handler2 = self._create_mocked_handler()
+        handler2.get()
+        
+        # Verify requests.get was called only once
+        mock_requests_get.assert_called_once()
+    
+    @patch("mito_ai.version_check.__version__", "1.0.0")
+    @patch("mito_ai.version_check.requests.get")
+    def test_pypi_request_failure(self, mock_requests_get: MagicMock) -> None:
+        # Mock the response
+        mock_requests_get.side_effect = Exception("Connection error")
+        
+        # Create handler instance
+        handler = self._create_mocked_handler()
+        
+        # Call the get method directly
+        handler.get()
+        
+        # Get the response body and verify
+        response_body = json.loads(handler._write_buffer[0])
+        # When the request fails, the handler returns an error response
+        self.assertIn("error", response_body)
+        self.assertEqual(handler._status_code, 500)
+    
+    @patch("mito_ai.version_check.__version__", "1.0.0")
+    def test_general_exception_handling(self) -> None:
+        # Create handler instance with a mock that will cause an exception
+        handler = self._create_mocked_handler()
+        
+        with patch.object(VersionCheckHandler, '_get_latest_version', side_effect=Exception("Test error")):
+            # Call the get method directly
+            handler.get()
+            
+            # Get the response body and verify
+            response_body = json.loads(handler._write_buffer[0])
+            self.assertIn("error", response_body)
+            self.assertEqual(handler._status_code, 500)
+    
+    def _create_mocked_handler(self) -> VersionCheckHandler:
+        """Create a mocked RequestHandler instance for testing."""
+        handler = VersionCheckHandler(MagicMock(), MagicMock())
+        handler._status_code = 200
+        handler._write_buffer = []
+        
+        # We need to assign the mock methods to new variables first to avoid
+        # the mypy "Cannot assign to a method" error
+        def mock_write(chunk: str) -> None:
+            # Convert string to bytes to match expected type
+            encoded_chunk = chunk.encode('utf-8') if isinstance(chunk, str) else chunk
+            handler._write_buffer.append(encoded_chunk)
+        
+        def mock_set_status(status_code: int) -> None:
+            handler._status_code = status_code
+        
+        def mock_set_header(name: str, value: str) -> None:
+            pass
+        
+        # Use setattr instead of direct assignment to avoid mypy errors
+        setattr(handler, 'write', mock_write)
+        setattr(handler, 'set_status', mock_set_status)
+        setattr(handler, 'set_header', mock_set_header)
+        
+        return handler 

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/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+