227 lines
8.2 KiB
Python
227 lines
8.2 KiB
Python
import os
|
|
import pytest
|
|
from unittest.mock import patch, Mock
|
|
from langchain_openai.chat_models import ChatOpenAI
|
|
from langchain_anthropic.chat_models import ChatAnthropic
|
|
from dataclasses import dataclass
|
|
|
|
from ra_aid.env import validate_environment
|
|
from ra_aid.llm import initialize_llm, initialize_expert_llm
|
|
|
|
@pytest.fixture
|
|
def clean_env(monkeypatch):
|
|
"""Remove relevant environment variables before each test"""
|
|
env_vars = [
|
|
'ANTHROPIC_API_KEY', 'OPENAI_API_KEY', 'OPENROUTER_API_KEY',
|
|
'OPENAI_API_BASE', 'EXPERT_ANTHROPIC_API_KEY', 'EXPERT_OPENAI_API_KEY',
|
|
'EXPERT_OPENROUTER_API_KEY', 'EXPERT_OPENAI_API_BASE'
|
|
]
|
|
for var in env_vars:
|
|
monkeypatch.delenv(var, raising=False)
|
|
|
|
@pytest.fixture
|
|
def mock_openai():
|
|
"""
|
|
Mock ChatOpenAI class for testing OpenAI provider initialization.
|
|
Prevents actual API calls during testing.
|
|
"""
|
|
with patch('ra_aid.llm.ChatOpenAI') as mock:
|
|
mock.return_value = Mock(spec=ChatOpenAI)
|
|
yield mock
|
|
|
|
def test_initialize_expert_defaults(clean_env, mock_openai, monkeypatch):
|
|
"""Test expert LLM initialization with default parameters."""
|
|
monkeypatch.setenv("EXPERT_OPENAI_API_KEY", "test-key")
|
|
llm = initialize_expert_llm()
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-key",
|
|
model="o1-preview"
|
|
)
|
|
|
|
def test_initialize_expert_openai_custom(clean_env, mock_openai, monkeypatch):
|
|
"""Test expert OpenAI initialization with custom parameters."""
|
|
monkeypatch.setenv("EXPERT_OPENAI_API_KEY", "test-key")
|
|
llm = initialize_expert_llm("openai", "gpt-4-preview")
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-key",
|
|
model="gpt-4-preview"
|
|
)
|
|
|
|
def test_initialize_expert_anthropic(clean_env, mock_anthropic, monkeypatch):
|
|
"""Test expert Anthropic initialization."""
|
|
monkeypatch.setenv("EXPERT_ANTHROPIC_API_KEY", "test-key")
|
|
llm = initialize_expert_llm("anthropic", "claude-3")
|
|
|
|
mock_anthropic.assert_called_once_with(
|
|
api_key="test-key",
|
|
model_name="claude-3"
|
|
)
|
|
|
|
def test_initialize_expert_openrouter(clean_env, mock_openai, monkeypatch):
|
|
"""Test expert OpenRouter initialization."""
|
|
monkeypatch.setenv("EXPERT_OPENROUTER_API_KEY", "test-key")
|
|
llm = initialize_expert_llm("openrouter", "models/mistral-large")
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-key",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
model="models/mistral-large"
|
|
)
|
|
|
|
def test_initialize_expert_openai_compatible(clean_env, mock_openai, monkeypatch):
|
|
"""Test expert OpenAI-compatible initialization."""
|
|
monkeypatch.setenv("EXPERT_OPENAI_API_KEY", "test-key")
|
|
monkeypatch.setenv("EXPERT_OPENAI_API_BASE", "http://test-url")
|
|
llm = initialize_expert_llm("openai-compatible", "local-model")
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-key",
|
|
base_url="http://test-url",
|
|
model="local-model"
|
|
)
|
|
|
|
def test_initialize_expert_unsupported_provider(clean_env):
|
|
"""Test error handling for unsupported provider in expert mode."""
|
|
with pytest.raises(ValueError, match=r"Unsupported provider: unknown"):
|
|
initialize_expert_llm("unknown", "model")
|
|
|
|
def test_initialize_openai(clean_env, mock_openai):
|
|
"""Test OpenAI provider initialization"""
|
|
os.environ["OPENAI_API_KEY"] = "test-key"
|
|
model = initialize_llm("openai", "gpt-4")
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-key",
|
|
model="gpt-4"
|
|
)
|
|
|
|
def test_initialize_anthropic(clean_env, mock_anthropic):
|
|
"""Test Anthropic provider initialization"""
|
|
os.environ["ANTHROPIC_API_KEY"] = "test-key"
|
|
model = initialize_llm("anthropic", "claude-3")
|
|
|
|
mock_anthropic.assert_called_once_with(
|
|
api_key="test-key",
|
|
model_name="claude-3"
|
|
)
|
|
|
|
def test_initialize_openrouter(clean_env, mock_openai):
|
|
"""Test OpenRouter provider initialization"""
|
|
os.environ["OPENROUTER_API_KEY"] = "test-key"
|
|
model = initialize_llm("openrouter", "mistral-large")
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-key",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
model="mistral-large"
|
|
)
|
|
|
|
def test_initialize_openai_compatible(clean_env, mock_openai):
|
|
"""Test OpenAI-compatible provider initialization"""
|
|
os.environ["OPENAI_API_KEY"] = "test-key"
|
|
os.environ["OPENAI_API_BASE"] = "https://custom-endpoint/v1"
|
|
model = initialize_llm("openai-compatible", "local-model")
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-key",
|
|
base_url="https://custom-endpoint/v1",
|
|
model="local-model"
|
|
)
|
|
|
|
def test_initialize_unsupported_provider(clean_env):
|
|
"""Test initialization with unsupported provider raises ValueError"""
|
|
with pytest.raises(ValueError) as exc_info:
|
|
initialize_llm("unsupported", "model")
|
|
assert str(exc_info.value) == "Unsupported provider: unsupported"
|
|
|
|
def test_provider_name_validation():
|
|
"""Test provider name validation and normalization."""
|
|
# Test all supported providers
|
|
providers = ["openai", "anthropic", "openrouter", "openai-compatible"]
|
|
for provider in providers:
|
|
try:
|
|
with patch(f'ra_aid.llm.ChatOpenAI'), patch('ra_aid.llm.ChatAnthropic'):
|
|
initialize_llm(provider, "test-model")
|
|
except ValueError:
|
|
pytest.fail(f"Valid provider {provider} raised ValueError")
|
|
|
|
# Test case sensitivity
|
|
with patch('ra_aid.llm.ChatOpenAI'):
|
|
with pytest.raises(ValueError):
|
|
initialize_llm("OpenAI", "test-model")
|
|
|
|
def test_initialize_llm_cross_provider(clean_env, mock_openai, mock_anthropic, monkeypatch):
|
|
"""Test initializing different providers in sequence."""
|
|
# Initialize OpenAI
|
|
monkeypatch.setenv("OPENAI_API_KEY", "openai-key")
|
|
llm1 = initialize_llm("openai", "gpt-4")
|
|
|
|
# Initialize Anthropic
|
|
monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic-key")
|
|
llm2 = initialize_llm("anthropic", "claude-3")
|
|
|
|
# Verify both were initialized correctly
|
|
mock_openai.assert_called_once_with(
|
|
api_key="openai-key",
|
|
model="gpt-4"
|
|
)
|
|
mock_anthropic.assert_called_once_with(
|
|
api_key="anthropic-key",
|
|
model_name="claude-3"
|
|
)
|
|
|
|
def test_environment_variable_precedence(clean_env, mock_openai, monkeypatch):
|
|
"""Test environment variable precedence and fallback."""
|
|
from ra_aid.env import validate_environment
|
|
from dataclasses import dataclass
|
|
|
|
@dataclass
|
|
class Args:
|
|
provider: str
|
|
expert_provider: str
|
|
model: str = None
|
|
expert_model: str = None
|
|
|
|
# Test expert mode with explicit key
|
|
# Set up base environment first
|
|
monkeypatch.setenv("OPENAI_API_KEY", "base-key")
|
|
monkeypatch.setenv("EXPERT_OPENAI_API_KEY", "expert-key")
|
|
monkeypatch.setenv("TAVILY_API_KEY", "tavily-key")
|
|
args = Args(provider="openai", expert_provider="openai")
|
|
expert_enabled, expert_missing, web_enabled, web_missing = validate_environment(args)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert web_enabled
|
|
assert not web_missing
|
|
|
|
llm = initialize_expert_llm()
|
|
mock_openai.assert_called_with(
|
|
api_key="expert-key",
|
|
model="o1-preview"
|
|
)
|
|
|
|
# Test empty key validation
|
|
monkeypatch.setenv("EXPERT_OPENAI_API_KEY", "")
|
|
monkeypatch.delenv("OPENAI_API_KEY", raising=False) # Remove fallback
|
|
monkeypatch.delenv("TAVILY_API_KEY", raising=False) # Remove web research
|
|
monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic-key") # Add for provider validation
|
|
monkeypatch.setenv("ANTHROPIC_MODEL", "claude-3-haiku-20240307") # Add model for provider validation
|
|
args = Args(provider="anthropic", expert_provider="openai") # Change base provider to avoid validation error
|
|
expert_enabled, expert_missing, web_enabled, web_missing = validate_environment(args)
|
|
assert not expert_enabled
|
|
assert expert_missing
|
|
assert not web_enabled
|
|
assert web_missing
|
|
|
|
@pytest.fixture
|
|
def mock_anthropic():
|
|
"""
|
|
Mock ChatAnthropic class for testing Anthropic provider initialization.
|
|
Prevents actual API calls during testing.
|
|
"""
|
|
with patch('ra_aid.llm.ChatAnthropic') as mock:
|
|
mock.return_value = Mock(spec=ChatAnthropic)
|
|
yield mock
|