mito-ai 0.1.32__py3-none-any.whl → 0.1.34__py3-none-any.whl

mito_ai/tests/providers_test.py

@@ -1,438 +0,0 @@
-# Copyright (c) Saga Inc.
-# Distributed under the terms of the GNU Affero General Public License v3.0 License.
-
-from __future__ import annotations
-from datetime import datetime
-from typing import Any, List
-from unittest.mock import patch, MagicMock, AsyncMock
-
-import pytest
-from traitlets.config import Config
-from mito_ai.completions.providers import OpenAIProvider
-from mito_ai.completions.models import (
-    MessageType,
-    AICapabilities,
-    CompletionReply
-)
-from mito_ai.completions.providers import OpenAIProvider
-from mito_ai.completions.models import MessageType, AICapabilities
-from mito_ai.utils.server_limits import OS_MONTHLY_AI_COMPLETIONS_LIMIT
-from openai.types.chat import ChatCompletionMessageParam
-
-REALLY_OLD_DATE = "2020-01-01"
-TODAY = datetime.now().strftime("%Y-%m-%d")
-FAKE_API_KEY = "sk-1234567890"
-
-@pytest.fixture
-def provider_config() -> Config:
-    """Create a proper Config object for the OpenAIProvider."""
-    config = Config()
-    config.OpenAIProvider = Config()
-    config.OpenAIClient = Config()
-    return config
-
-@pytest.fixture(autouse=True)
-def reset_env_vars(monkeypatch: pytest.MonkeyPatch) -> None:
-    for var in [
-        "OPENAI_API_KEY", "CLAUDE_API_KEY",
-        "GEMINI_API_KEY", "OLLAMA_MODEL",
-        "AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_MODEL"
-    ]:
-        monkeypatch.delenv(var, raising=False)
-
-
-def patch_server_limits(is_pro: bool = False, completion_count: int = 1, first_date: str = TODAY) -> Any:
-    return patch.multiple(
-        "mito_ai.utils.server_limits",
-        get_chat_completion_count=MagicMock(return_value=completion_count),
-        get_first_completion_date=MagicMock(return_value=first_date),
-        is_pro=MagicMock(return_value=is_pro),
-        check_mito_server_quota=MagicMock(return_value=None)
-    )
-
-
-def patch_openai_model_list() -> Any:
-    mock_openai_instance = MagicMock()
-    mock_openai_instance.models.list.return_value = [MagicMock(id="gpt-4o-mini")]
-
-    # Patch the constructor call to return your mock instance
-    return patch("openai.OpenAI", return_value=mock_openai_instance)
-
-
-def mock_openai_client() -> Any:
-    """Mock the OpenAI client with user key capabilities."""
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "gpt-4o-mini"},
-        provider="OpenAI with user key",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "user"
-    mock_client.request_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_completions = AsyncMock(return_value="Test completion")
-    return patch("mito_ai.completions.providers.OpenAIClient", return_value=mock_client)
-
-
-def mock_gemini_client() -> Any:
-    """Mock the Gemini client capabilities."""
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "gemini-2-pro"},
-        provider="Gemini",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "gemini"
-    mock_client.request_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_completions = AsyncMock(return_value="Test completion")
-    return patch("mito_ai.completions.providers.GeminiClient", return_value=mock_client)
-
-
-def mock_azure_openai_client() -> Any:
-    """Mock the Azure OpenAI client capabilities."""
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "gpt-4o"},
-        provider="Azure OpenAI",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "azure"
-    mock_client.request_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_completions = AsyncMock(return_value="Test completion")
-    return patch("mito_ai.completions.providers.OpenAIClient", return_value=mock_client)
-
-
-
-def mock_claude_client() -> Any:
-    """Mock the Claude client capabilities."""
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "claude-3-opus-20240229"},
-        provider="Claude",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "claude"
-    mock_client.request_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_response = AsyncMock(return_value="Test completion")
-    return patch("mito_ai.completions.providers.AnthropicClient", return_value=mock_client)
-
-
-def test_os_user_openai_key_set_below_limit(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", FAKE_API_KEY)
-
-    with (
-        patch_server_limits(is_pro=False, completion_count=1),
-        mock_openai_client()
-    ):
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert "user key" in capabilities.provider
-        assert llm.last_error is None
-
-
-def test_os_user_openai_key_set_above_limit(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", FAKE_API_KEY)
-
-    with (
-        patch_server_limits(is_pro=False, completion_count=OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1),
-        mock_openai_client()
-    ):
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert "user key" in capabilities.provider
-        assert llm.last_error is None
-
-
-def test_pro_user_openai_key_set_below_limit(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", FAKE_API_KEY)
-
-    with (
-        patch_server_limits(is_pro=True, completion_count=1),
-        mock_openai_client()
-    ):
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert "user key" in capabilities.provider
-        assert llm.last_error is None
-
-
-def test_pro_user_openai_key_set_above_limit(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", FAKE_API_KEY)
-
-    with (
-        patch_server_limits(is_pro=True, completion_count=OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1),
-        mock_openai_client()
-    ):
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert "user key" in capabilities.provider
-        assert llm.last_error is None
-
-
-def test_gemini_provider(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("GEMINI_API_KEY", "gemini-key")
-    monkeypatch.setattr("mito_ai.constants.GEMINI_API_KEY", "gemini-key")
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
-
-    with mock_gemini_client():
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert capabilities.provider == "Gemini"
-
-
-def test_azure_openai_provider(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setattr("mito_ai.enterprise.utils.is_enterprise", lambda: True)
-    monkeypatch.setattr("mito_ai.enterprise.utils.is_azure_openai_configured", lambda: True)
-    monkeypatch.setattr("mito_ai.enterprise.utils.AZURE_OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setattr("mito_ai.enterprise.utils.AZURE_OPENAI_ENDPOINT", "https://example.com")
-    monkeypatch.setattr("mito_ai.enterprise.utils.AZURE_OPENAI_MODEL", "gpt-4o")
-    monkeypatch.setattr("mito_ai.enterprise.utils.AZURE_OPENAI_API_VERSION", "2024-12-01-preview")
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
-
-    with mock_azure_openai_client():
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert capabilities.provider == "Azure OpenAI"
-
-
-def test_claude_provider(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("CLAUDE_API_KEY", "claude-key")
-    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", "claude-key")
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
-
-    with mock_claude_client():
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert capabilities.provider == "Claude"
-
-
-def test_provider_priority_order(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    # Set up all possible providers
-    monkeypatch.setattr("mito_ai.enterprise.utils.is_enterprise", lambda: True)
-    monkeypatch.setattr("mito_ai.enterprise.utils.is_azure_openai_configured", lambda: True)
-    monkeypatch.setattr("mito_ai.enterprise.utils.AZURE_OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setattr("mito_ai.enterprise.utils.AZURE_OPENAI_ENDPOINT", "https://example.com")
-    monkeypatch.setattr("mito_ai.enterprise.utils.AZURE_OPENAI_MODEL", "gpt-4o")
-    monkeypatch.setattr("mito_ai.enterprise.utils.AZURE_OPENAI_API_VERSION", "2024-12-01-preview")
-    monkeypatch.setenv("OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setenv("CLAUDE_API_KEY", "claude-key")
-    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", "claude-key")
-
-    # Azure OpenAI should have highest priority when enterprise is enabled
-    # Clear other provider settings to ensure Azure OpenAI is selected
-    monkeypatch.delenv("GEMINI_API_KEY", raising=False)
-    monkeypatch.setattr("mito_ai.constants.GEMINI_API_KEY", None)
-    # Clear Claude settings to ensure Azure OpenAI is selected
-    monkeypatch.delenv("CLAUDE_API_KEY", raising=False)
-    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", None)
-    with mock_azure_openai_client():
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert capabilities.provider == "Azure OpenAI"
-
-    # Without enterprise, OpenAI should have highest priority
-    monkeypatch.setattr("mito_ai.enterprise.utils.is_enterprise", lambda: False)
-    monkeypatch.setattr("mito_ai.enterprise.utils.is_azure_openai_configured", lambda: False)
-    with mock_openai_client():
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert capabilities.provider == "OpenAI with user key"
-
-    # Without OpenAI key, Claude should be used (higher priority than Gemini)
-    monkeypatch.delenv("OPENAI_API_KEY")
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
-    # Ensure provider_config doesn't have an api_key set
-    provider_config.OpenAIProvider.api_key = None
-    # Re-enable Claude settings
-    monkeypatch.setenv("CLAUDE_API_KEY", "claude-key")
-    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", "claude-key")
-    with mock_claude_client():
-        llm = OpenAIProvider(config=provider_config)
-        capabilities = llm.capabilities
-        assert capabilities.provider == "Claude"
-
-
-@pytest.mark.asyncio
-async def test_completion_request(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", FAKE_API_KEY)
-
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "gpt-4o-mini"},
-        provider="OpenAI with user key",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "user"
-    mock_client.request_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_completions = AsyncMock(return_value="Test completion")
-
-    with patch("mito_ai.completions.providers.OpenAIClient", return_value=mock_client):
-        llm = OpenAIProvider(config=provider_config)
-        messages: List[ChatCompletionMessageParam] = [
-            {"role": "user", "content": "Test message"}
-        ]
-
-        completion = await llm.request_completions(
-            message_type=MessageType.CHAT,
-            messages=messages,
-            model="gpt-4o-mini"
-        )
-
-        assert completion == "Test completion"
-        mock_client.request_completions.assert_called_once()
-
-
-@pytest.mark.asyncio
-async def test_stream_completion(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("OPENAI_API_KEY", FAKE_API_KEY)
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", FAKE_API_KEY)
-
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "gpt-4o-mini"},
-        provider="OpenAI with user key",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "user"
-    mock_client.request_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_completions = AsyncMock(return_value="Test completion")
-
-    with patch("mito_ai.completions.providers.OpenAIClient", return_value=mock_client):
-        llm = OpenAIProvider(config=provider_config)
-        messages: List[ChatCompletionMessageParam] = [
-            {"role": "user", "content": "Test message"}
-        ]
-
-        reply_chunks = []
-        def mock_reply(chunk):
-            reply_chunks.append(chunk)
-
-        completion = await llm.stream_completions(
-            message_type=MessageType.CHAT,
-            messages=messages,
-            model="gpt-4o-mini",
-            message_id="test-id",
-            thread_id="test-thread",
-            reply_fn=mock_reply
-        )
-
-        assert completion == "Test completion"
-        mock_client.stream_completions.assert_called_once()
-        assert len(reply_chunks) > 0
-        assert isinstance(reply_chunks[0], CompletionReply)
-
-
-def test_error_handling(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("OPENAI_API_KEY", "invalid-key")
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", "invalid-key")
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "gpt-4o-mini"},
-        provider="OpenAI with user key",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "user"
-    mock_client.request_completions.side_effect = Exception("API error")
-
-    with patch("mito_ai.completions.providers.OpenAIClient", return_value=mock_client):
-        llm = OpenAIProvider(config=provider_config)
-        assert llm.last_error is None  # Error should be None until a request is made
-
-
-@pytest.mark.asyncio
-async def test_claude_completion_request(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("CLAUDE_API_KEY", "claude-key")
-    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", "claude-key")
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
-
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "claude-3-opus-20240229"},
-        provider="Claude",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "claude"
-    mock_client.request_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_completions = AsyncMock(return_value="Test completion")
-
-    with patch("mito_ai.completions.providers.AnthropicClient", return_value=mock_client):
-        llm = OpenAIProvider(config=provider_config)
-        messages: List[ChatCompletionMessageParam] = [
-            {"role": "user", "content": "Test message"}
-        ]
-
-        completion = await llm.request_completions(
-            message_type=MessageType.CHAT,
-            messages=messages,
-            model="claude-3-opus-20240229"
-        )
-
-        assert completion == "Test completion"
-        mock_client.request_completions.assert_called_once()
-
-
-@pytest.mark.asyncio
-async def test_claude_stream_completion(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("CLAUDE_API_KEY", "claude-key")
-    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", "claude-key")
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
-
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "claude-3-opus-20240229"},
-        provider="Claude",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "claude"
-    mock_client.request_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_completions = AsyncMock(return_value="Test completion")
-    mock_client.stream_response = AsyncMock(return_value="Test completion")
-
-    with patch("mito_ai.completions.providers.AnthropicClient", return_value=mock_client):
-        llm = OpenAIProvider(config=provider_config)
-        messages: List[ChatCompletionMessageParam] = [
-            {"role": "user", "content": "Test message"}
-        ]
-
-        reply_chunks = []
-        def mock_reply(chunk):
-            reply_chunks.append(chunk)
-
-        completion = await llm.stream_completions(
-            message_type=MessageType.CHAT,
-            messages=messages,
-            model="claude-3-opus-20240229",
-            message_id="test-id",
-            thread_id="test-thread",
-            reply_fn=mock_reply
-        )
-
-        assert completion == "Test completion"
-        mock_client.stream_response.assert_called_once()
-        assert len(reply_chunks) > 0
-        assert isinstance(reply_chunks[0], CompletionReply)
-
-
-def test_claude_error_handling(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
-    monkeypatch.setenv("CLAUDE_API_KEY", "invalid-key")
-    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", "invalid-key")
-    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
-
-    mock_client = MagicMock()
-    mock_client.capabilities = AICapabilities(
-        configuration={"model": "claude-3-opus-20240229"},
-        provider="Claude",
-        type="ai_capabilities"
-    )
-    mock_client.key_type = "claude"
-    mock_client.request_completions.side_effect = Exception("API error")
-
-    with patch("mito_ai.completions.providers.AnthropicClient", return_value=mock_client):
-        llm = OpenAIProvider(config=provider_config)
-        assert llm.last_error is None  # Error should be None until a request is made