You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
128 lines
3.8 KiB
128 lines
3.8 KiB
import os
|
|
import pytest
|
|
from unittest.mock import Mock, patch
|
|
from swarms.models.anthropic import Anthropic
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_anthropic_env():
|
|
os.environ["ANTHROPIC_API_URL"] = "https://test.anthropic.com"
|
|
os.environ["ANTHROPIC_API_KEY"] = "test_api_key"
|
|
yield
|
|
del os.environ["ANTHROPIC_API_URL"]
|
|
del os.environ["ANTHROPIC_API_KEY"]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_requests_post():
|
|
with patch("requests.post") as mock_post:
|
|
yield mock_post
|
|
|
|
|
|
@pytest.fixture
|
|
def anthropic_instance():
|
|
return Anthropic(model="test-model")
|
|
|
|
|
|
def test_anthropic_init_default_values(anthropic_instance):
|
|
assert anthropic_instance.model == "test-model"
|
|
assert anthropic_instance.max_tokens_to_sample == 256
|
|
assert anthropic_instance.temperature is None
|
|
assert anthropic_instance.top_k is None
|
|
assert anthropic_instance.top_p is None
|
|
assert anthropic_instance.streaming is False
|
|
assert anthropic_instance.default_request_timeout == 600
|
|
assert anthropic_instance.anthropic_api_url == "https://test.anthropic.com"
|
|
assert anthropic_instance.anthropic_api_key == "test_api_key"
|
|
|
|
|
|
def test_anthropic_init_custom_values():
|
|
anthropic_instance = Anthropic(
|
|
model="custom-model",
|
|
max_tokens_to_sample=128,
|
|
temperature=0.8,
|
|
top_k=5,
|
|
top_p=0.9,
|
|
streaming=True,
|
|
default_request_timeout=300,
|
|
)
|
|
assert anthropic_instance.model == "custom-model"
|
|
assert anthropic_instance.max_tokens_to_sample == 128
|
|
assert anthropic_instance.temperature == 0.8
|
|
assert anthropic_instance.top_k == 5
|
|
assert anthropic_instance.top_p == 0.9
|
|
assert anthropic_instance.streaming is True
|
|
assert anthropic_instance.default_request_timeout == 300
|
|
|
|
|
|
def test_anthropic_default_params(anthropic_instance):
|
|
default_params = anthropic_instance._default_params()
|
|
assert default_params == {
|
|
"max_tokens_to_sample": 256,
|
|
"model": "test-model",
|
|
}
|
|
|
|
|
|
def test_anthropic_run(mock_anthropic_env, mock_requests_post, anthropic_instance):
|
|
mock_response = Mock()
|
|
mock_response.json.return_value = {"completion": "Generated text"}
|
|
mock_requests_post.return_value = mock_response
|
|
|
|
task = "Generate text"
|
|
stop = ["stop1", "stop2"]
|
|
|
|
completion = anthropic_instance.run(task, stop)
|
|
|
|
assert completion == "Generated text"
|
|
mock_requests_post.assert_called_once_with(
|
|
"https://test.anthropic.com/completions",
|
|
headers={"Authorization": "Bearer test_api_key"},
|
|
json={
|
|
"prompt": task,
|
|
"stop_sequences": stop,
|
|
"max_tokens_to_sample": 256,
|
|
"model": "test-model",
|
|
},
|
|
timeout=600,
|
|
)
|
|
|
|
|
|
def test_anthropic_call(mock_anthropic_env, mock_requests_post, anthropic_instance):
|
|
mock_response = Mock()
|
|
mock_response.json.return_value = {"completion": "Generated text"}
|
|
mock_requests_post.return_value = mock_response
|
|
|
|
task = "Generate text"
|
|
stop = ["stop1", "stop2"]
|
|
|
|
completion = anthropic_instance(task, stop)
|
|
|
|
assert completion == "Generated text"
|
|
mock_requests_post.assert_called_once_with(
|
|
"https://test.anthropic.com/completions",
|
|
headers={"Authorization": "Bearer test_api_key"},
|
|
json={
|
|
"prompt": task,
|
|
"stop_sequences": stop,
|
|
"max_tokens_to_sample": 256,
|
|
"model": "test-model",
|
|
},
|
|
timeout=600,
|
|
)
|
|
|
|
|
|
def test_anthropic_exception_handling(
|
|
mock_anthropic_env, mock_requests_post, anthropic_instance
|
|
):
|
|
mock_response = Mock()
|
|
mock_response.json.return_value = {"error": "An error occurred"}
|
|
mock_requests_post.return_value = mock_response
|
|
|
|
task = "Generate text"
|
|
stop = ["stop1", "stop2"]
|
|
|
|
with pytest.raises(Exception) as excinfo:
|
|
anthropic_instance(task, stop)
|
|
|
|
assert "An error occurred" in str(excinfo.value)
|