583 lines
22 KiB
Python
583 lines
22 KiB
Python
"""
|
|
Copyright 2024-2026 ReplyPier
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import patch, MagicMock, AsyncMock
|
|
from app.utils import agno_utils
|
|
from fastapi import HTTPException
|
|
|
|
|
|
class TestCreateModel:
|
|
"""Test cases for the create_model function"""
|
|
|
|
def test_create_model_openai_with_response_format(self):
|
|
"""Test creating OpenAI model with response format"""
|
|
with patch("app.utils.agno_utils.OpenAIChat") as mock_openai:
|
|
mock_model = MagicMock()
|
|
mock_openai.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="OPENAI",
|
|
api_key="test-api-key",
|
|
model_name="gpt-4",
|
|
max_tokens=2000,
|
|
response_format={"type": "json_object"}
|
|
)
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="gpt-4",
|
|
max_tokens=2000,
|
|
response_format={"type": "json_object"}
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_openai_without_response_format(self):
|
|
"""Test creating OpenAI model without response format"""
|
|
with patch("app.utils.agno_utils.OpenAIChat") as mock_openai:
|
|
mock_model = MagicMock()
|
|
mock_openai.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="OPENAI",
|
|
api_key="test-api-key",
|
|
model_name="gpt-3.5-turbo"
|
|
)
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="gpt-3.5-turbo",
|
|
max_tokens=1000
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_openai_compatible(self):
|
|
"""Test creating an OPENAI_COMPATIBLE model passes base_url through"""
|
|
with patch("app.utils.agno_utils.OpenAIChat") as mock_openai:
|
|
mock_model = MagicMock()
|
|
mock_openai.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="OPENAI_COMPATIBLE",
|
|
api_key="test-api-key",
|
|
model_name="local-model",
|
|
base_url="https://openrouter.ai/api/v1"
|
|
)
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="local-model",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
max_tokens=1000
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_openai_compatible_strips_chat_completions_suffix(self):
|
|
"""A base_url accidentally including /chat/completions is normalized"""
|
|
with patch("app.utils.agno_utils.OpenAIChat") as mock_openai:
|
|
mock_model = MagicMock()
|
|
mock_openai.return_value = mock_model
|
|
|
|
agno_utils.create_model(
|
|
model_type="OPENAI_COMPATIBLE",
|
|
api_key="test-api-key",
|
|
model_name="local-model",
|
|
base_url="https://openrouter.ai/api/v1/chat/completions/"
|
|
)
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="local-model",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
max_tokens=1000
|
|
)
|
|
|
|
def test_create_model_openai_compatible_requires_base_url(self):
|
|
"""OPENAI_COMPATIBLE without a base_url raises"""
|
|
with pytest.raises(HTTPException):
|
|
agno_utils.create_model(
|
|
model_type="OPENAI_COMPATIBLE",
|
|
api_key="test-api-key",
|
|
model_name="local-model"
|
|
)
|
|
|
|
def test_create_model_replypier(self):
|
|
"""Test creating REPLYPIER model (should use OpenAI)"""
|
|
with patch("app.utils.agno_utils.OpenAIChat") as mock_openai:
|
|
mock_model = MagicMock()
|
|
mock_openai.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="REPLYPIER",
|
|
api_key="test-api-key",
|
|
model_name="replypier-model"
|
|
)
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="replypier-model",
|
|
max_tokens=1000
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_anthropic(self):
|
|
"""Test creating Anthropic model"""
|
|
# Mock the import statement inside the function
|
|
mock_claude_class = MagicMock()
|
|
mock_model = MagicMock()
|
|
mock_claude_class.return_value = mock_model
|
|
|
|
with patch.dict('sys.modules', {'agno.models.anthropic': MagicMock(Claude=mock_claude_class)}):
|
|
result = agno_utils.create_model(
|
|
model_type="ANTHROPIC",
|
|
api_key="test-api-key",
|
|
model_name="claude-3-opus",
|
|
max_tokens=1500
|
|
)
|
|
|
|
mock_claude_class.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="claude-3-opus",
|
|
max_tokens=1500
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_deepseek(self):
|
|
"""Test creating DeepSeek model"""
|
|
with patch("agno.models.deepseek.DeepSeek") as mock_deepseek:
|
|
mock_model = MagicMock()
|
|
mock_deepseek.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="DEEPSEEK",
|
|
api_key="test-api-key",
|
|
model_name="deepseek-chat"
|
|
)
|
|
|
|
mock_deepseek.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="deepseek-chat",
|
|
max_tokens=1000
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_google(self):
|
|
"""Test creating Google Gemini model"""
|
|
import sys
|
|
from types import ModuleType
|
|
import agno.models
|
|
|
|
google_module = ModuleType("agno.models.google")
|
|
google_module.Gemini = MagicMock()
|
|
mock_model = MagicMock()
|
|
google_module.Gemini.return_value = mock_model
|
|
|
|
with patch.dict(sys.modules, {"agno.models.google": google_module}), \
|
|
patch.object(agno.models, "google", google_module, create=True):
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="GOOGLE",
|
|
api_key="test-api-key",
|
|
model_name="gemini-pro"
|
|
)
|
|
|
|
google_module.Gemini.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="gemini-pro",
|
|
max_output_tokens=1000 # Gemini's agno model uses max_output_tokens
|
|
)
|
|
assert result == mock_model
|
|
assert result == mock_model
|
|
|
|
def test_create_model_googlevertex(self):
|
|
"""Test creating Google Vertex AI model (should trigger import error)"""
|
|
with pytest.raises(HTTPException) as excinfo:
|
|
agno_utils.create_model(
|
|
model_type="GOOGLEVERTEX",
|
|
api_key="test-api-key",
|
|
model_name="gemini-pro-vertex"
|
|
)
|
|
|
|
assert excinfo.value.status_code == 400
|
|
assert "is not available in this installation" in excinfo.value.detail
|
|
|
|
def test_create_model_groq_with_response_format(self):
|
|
"""Test creating Groq model with response format (uses PatchedGroq)"""
|
|
result = agno_utils.create_model(
|
|
model_type="GROQ",
|
|
api_key="test-api-key",
|
|
model_name="mixtral-8x7b",
|
|
response_format={"type": "json_object"}
|
|
)
|
|
|
|
# Should return a PatchedGroq instance
|
|
assert result.__class__.__name__ == "PatchedGroq"
|
|
assert result.api_key == "test-api-key"
|
|
assert result.id == "mixtral-8x7b"
|
|
assert result.max_tokens == 1000
|
|
|
|
def test_create_model_groq_without_response_format(self):
|
|
"""Test creating Groq model without response format (uses PatchedGroq)"""
|
|
result = agno_utils.create_model(
|
|
model_type="GROQ",
|
|
api_key="test-api-key",
|
|
model_name="llama2-70b"
|
|
)
|
|
|
|
# Should return a PatchedGroq instance
|
|
assert result.__class__.__name__ == "PatchedGroq"
|
|
assert result.api_key == "test-api-key"
|
|
assert result.id == "llama2-70b"
|
|
assert result.max_tokens == 1000
|
|
|
|
def test_patched_groq_response_format_tools_conflict(self):
|
|
"""Test that PatchedGroq handles response_format + tools conflict correctly"""
|
|
result = agno_utils.create_model(
|
|
model_type="GROQ",
|
|
api_key="test-api-key",
|
|
model_name="llama-3.3-70b-versatile"
|
|
)
|
|
|
|
# Test with tools - should NOT include response_format
|
|
params_with_tools = result.get_request_params(
|
|
response_format={"type": "json_object"},
|
|
tools=[{"type": "function", "function": {"name": "test_tool"}}]
|
|
)
|
|
assert "tools" in params_with_tools
|
|
assert "response_format" not in params_with_tools
|
|
|
|
# Test without tools - should include response_format
|
|
params_without_tools = result.get_request_params(
|
|
response_format={"type": "json_object"}
|
|
)
|
|
assert "response_format" in params_without_tools
|
|
assert "tools" not in params_without_tools
|
|
|
|
def test_create_model_mistral(self):
|
|
"""Test creating Mistral model (mistralai is now a dependency)"""
|
|
with patch("agno.models.mistral.MistralChat") as mock_mistral:
|
|
mock_model = MagicMock()
|
|
mock_mistral.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="MISTRAL",
|
|
api_key="test-api-key",
|
|
model_name="mistral-large-latest"
|
|
)
|
|
|
|
mock_mistral.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="mistral-large-latest",
|
|
max_tokens=1000
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_huggingface(self):
|
|
"""Test creating HuggingFace model"""
|
|
with patch("agno.models.huggingface.HuggingFace") as mock_hf:
|
|
mock_model = MagicMock()
|
|
mock_hf.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="HUGGINGFACE",
|
|
api_key="test-api-key",
|
|
model_name="microsoft/DialoGPT-medium"
|
|
)
|
|
|
|
mock_hf.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="microsoft/DialoGPT-medium",
|
|
max_tokens=1000
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_ollama(self):
|
|
"""Test creating Ollama model (no API key needed)"""
|
|
with patch("agno.models.ollama.Ollama") as mock_ollama:
|
|
mock_model = MagicMock()
|
|
mock_ollama.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="OLLAMA",
|
|
api_key="not-needed",
|
|
model_name="llama2"
|
|
)
|
|
|
|
mock_ollama.assert_called_once_with(id="llama2")
|
|
assert result == mock_model
|
|
|
|
def test_create_model_xai(self):
|
|
"""Test creating xAI model"""
|
|
with patch("agno.models.xai.xAI") as mock_xai:
|
|
mock_model = MagicMock()
|
|
mock_xai.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="XAI",
|
|
api_key="test-api-key",
|
|
model_name="grok-1"
|
|
)
|
|
|
|
mock_xai.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
id="grok-1",
|
|
max_tokens=1000
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_case_insensitive(self):
|
|
"""Test that model type is case insensitive"""
|
|
with patch("app.utils.agno_utils.OpenAIChat") as mock_openai:
|
|
mock_model = MagicMock()
|
|
mock_openai.return_value = mock_model
|
|
|
|
# Test lowercase
|
|
result = agno_utils.create_model(
|
|
model_type="openai",
|
|
api_key="test-api-key",
|
|
model_name="gpt-4"
|
|
)
|
|
assert result == mock_model
|
|
|
|
# Test mixed case
|
|
mock_openai.reset_mock()
|
|
result = agno_utils.create_model(
|
|
model_type="OpenAI",
|
|
api_key="test-api-key",
|
|
model_name="gpt-4"
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_unsupported_type(self):
|
|
"""Test creating an unsupported model type"""
|
|
with pytest.raises(HTTPException) as excinfo:
|
|
agno_utils.create_model(
|
|
model_type="UNSUPPORTED_MODEL",
|
|
api_key="test-api-key",
|
|
model_name="test-model"
|
|
)
|
|
|
|
assert excinfo.value.status_code == 500
|
|
assert "Failed to initialize model" in excinfo.value.detail
|
|
|
|
def test_create_model_import_error(self):
|
|
"""Test handling import errors"""
|
|
# Mock the module to raise ImportError when Claude is accessed
|
|
mock_module = MagicMock()
|
|
mock_module.Claude.side_effect = ImportError("Module not found")
|
|
|
|
with patch.dict('sys.modules', {'agno.models.anthropic': mock_module}):
|
|
with pytest.raises(HTTPException) as excinfo:
|
|
agno_utils.create_model(
|
|
model_type="ANTHROPIC",
|
|
api_key="test-api-key",
|
|
model_name="claude-3"
|
|
)
|
|
|
|
assert excinfo.value.status_code == 400
|
|
assert "is not available in this installation" in excinfo.value.detail
|
|
|
|
def test_create_model_general_exception(self):
|
|
"""Test handling general exceptions during model creation"""
|
|
with patch("app.utils.agno_utils.OpenAIChat", side_effect=ValueError("Invalid configuration")):
|
|
with pytest.raises(HTTPException) as excinfo:
|
|
agno_utils.create_model(
|
|
model_type="OPENAI",
|
|
api_key="test-api-key",
|
|
model_name="gpt-4"
|
|
)
|
|
|
|
assert excinfo.value.status_code == 500
|
|
assert "Failed to initialize model" in excinfo.value.detail
|
|
|
|
|
|
class TestModelApiKey:
|
|
"""Test cases for the test_model_api_key function"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_test_model_api_key_success(self):
|
|
"""Test successful API key validation"""
|
|
with patch("app.utils.agno_utils.create_model") as mock_create_model, \
|
|
patch("app.utils.agno_utils.Agent") as mock_agent_class:
|
|
# Setup mocks
|
|
mock_model = MagicMock()
|
|
mock_create_model.return_value = mock_model
|
|
|
|
# Create an async mock for arun method
|
|
mock_agent = MagicMock()
|
|
mock_agent.arun = AsyncMock(return_value="Test successful.")
|
|
mock_agent_class.return_value = mock_agent
|
|
|
|
# Call the function
|
|
result = await agno_utils.test_model_api_key("test-api-key", "OPENAI", "gpt-4")
|
|
|
|
# Assert the function was called with correct parameters
|
|
mock_create_model.assert_called_once_with("OPENAI", "test-api-key", "gpt-4", base_url=None)
|
|
mock_agent_class.assert_called_once_with(
|
|
name="Test Agent",
|
|
model=mock_model,
|
|
instructions="You are a test agent. Just respond with 'Test successful.'",
|
|
debug_mode=False
|
|
)
|
|
mock_agent.arun.assert_called_once_with(message="This is a test message.")
|
|
assert result is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_test_model_api_key_create_model_failure(self):
|
|
"""Test API key validation when model creation fails"""
|
|
with patch("app.utils.agno_utils.create_model", side_effect=Exception("Invalid API key")):
|
|
result = await agno_utils.test_model_api_key("invalid-api-key", "OPENAI", "gpt-4")
|
|
assert result is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_test_model_api_key_agent_creation_failure(self):
|
|
"""Test API key validation when agent creation fails"""
|
|
with patch("app.utils.agno_utils.create_model") as mock_create_model, \
|
|
patch("app.utils.agno_utils.Agent", side_effect=Exception("Agent creation failed")):
|
|
mock_model = MagicMock()
|
|
mock_create_model.return_value = mock_model
|
|
|
|
result = await agno_utils.test_model_api_key("test-api-key", "OPENAI", "gpt-4")
|
|
assert result is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_test_model_api_key_agent_run_failure(self):
|
|
"""Test API key validation when agent run fails"""
|
|
with patch("app.utils.agno_utils.create_model") as mock_create_model, \
|
|
patch("app.utils.agno_utils.Agent") as mock_agent_class:
|
|
# Setup mocks
|
|
mock_model = MagicMock()
|
|
mock_create_model.return_value = mock_model
|
|
|
|
# Create an async mock that raises an exception
|
|
mock_agent = MagicMock()
|
|
mock_agent.arun = AsyncMock(side_effect=Exception("API request failed"))
|
|
mock_agent_class.return_value = mock_agent
|
|
|
|
# Call the function
|
|
result = await agno_utils.test_model_api_key("invalid-api-key", "OPENAI", "gpt-4")
|
|
|
|
# Assert result
|
|
assert result is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_test_model_api_key_different_model_types(self):
|
|
"""Test API key validation with different model types"""
|
|
test_cases = [
|
|
("ANTHROPIC", "claude-3"),
|
|
("GROQ", "mixtral-8x7b"),
|
|
("GOOGLE", "gemini-pro"),
|
|
("DEEPSEEK", "deepseek-chat")
|
|
]
|
|
|
|
for model_type, model_name in test_cases:
|
|
with patch("app.utils.agno_utils.create_model") as mock_create_model, \
|
|
patch("app.utils.agno_utils.Agent") as mock_agent_class:
|
|
# Setup mocks
|
|
mock_model = MagicMock()
|
|
mock_create_model.return_value = mock_model
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.arun = AsyncMock(return_value="Test successful.")
|
|
mock_agent_class.return_value = mock_agent
|
|
|
|
# Call the function
|
|
result = await agno_utils.test_model_api_key("test-api-key", model_type, model_name)
|
|
|
|
# Assert
|
|
mock_create_model.assert_called_once_with(model_type, "test-api-key", model_name, base_url=None)
|
|
assert result is True
|
|
|
|
|
|
class TestEdgeCases:
|
|
"""Test edge cases and boundary conditions"""
|
|
|
|
def test_create_model_empty_strings(self):
|
|
"""Test create_model with empty strings"""
|
|
with patch("app.utils.agno_utils.OpenAIChat") as mock_openai:
|
|
mock_model = MagicMock()
|
|
mock_openai.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="OPENAI",
|
|
api_key="",
|
|
model_name=""
|
|
)
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="",
|
|
id="",
|
|
max_tokens=1000
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_custom_max_tokens(self):
|
|
"""Test create_model with custom max_tokens"""
|
|
with patch("app.utils.agno_utils.OpenAIChat") as mock_openai:
|
|
mock_model = MagicMock()
|
|
mock_openai.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="OPENAI",
|
|
api_key="test-key",
|
|
model_name="gpt-4",
|
|
max_tokens=4000
|
|
)
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-key",
|
|
id="gpt-4",
|
|
max_tokens=4000
|
|
)
|
|
assert result == mock_model
|
|
|
|
def test_create_model_zero_max_tokens(self):
|
|
"""Test create_model with zero max_tokens"""
|
|
with patch("app.utils.agno_utils.OpenAIChat") as mock_openai:
|
|
mock_model = MagicMock()
|
|
mock_openai.return_value = mock_model
|
|
|
|
result = agno_utils.create_model(
|
|
model_type="OPENAI",
|
|
api_key="test-key",
|
|
model_name="gpt-4",
|
|
max_tokens=0
|
|
)
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-key",
|
|
id="gpt-4",
|
|
max_tokens=0
|
|
)
|
|
assert result == mock_model
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_test_model_api_key_empty_strings(self):
|
|
"""Test test_model_api_key with empty strings"""
|
|
with patch("app.utils.agno_utils.create_model") as mock_create_model, \
|
|
patch("app.utils.agno_utils.Agent") as mock_agent_class:
|
|
mock_model = MagicMock()
|
|
mock_create_model.return_value = mock_model
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.arun = AsyncMock(return_value="Test successful.")
|
|
mock_agent_class.return_value = mock_agent
|
|
|
|
result = await agno_utils.test_model_api_key("", "", "")
|
|
|
|
mock_create_model.assert_called_once_with("", "", "", base_url=None)
|
|
assert result is True
|