mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
dfe4df8c8a
* test: delete unconditionally skipped legacy tests * test: move whole-unit legacy test files into tests/unit * test: keep moved legacy tests free of import-time global state * test: keep the module-level invocation scan pointed at tests/local_testing * ci: drop the agent_testing CircleCI job emptied by the move * test: fix moved-test isolation and router coverage * test: add Tinyfish search package marker * test: isolate moved tests from logger state leaks * test: isolate Helicone logging fixture state * test: isolate Vertex pass-through credentials between moved tests * test: cancel S3 periodic flush tasks started by moved tests * ci: restore CircleCI assistant test selection after move * test: prevent Datadog datetime import shadowing * ci: drop the litellm_assistants_api_testing CircleCI job emptied by the move * test: deduplicate imports in rebased unit tests * test: remove duplicate passthrough router patch import * test: remove moved legacy source files after rebase * test: align moved tests with rebased main * test: carry main's legacy-file edits into moved destinations * test: make the moved cost map fallback tests assert the fetch and the backup The four fallback cases only checked the result was non-empty, so they still passed with integrity validation disabled. They now inject a mock client, assert one fetch happened, and assert the result is exactly the local backup with the fallback reason recorded. --------- Co-authored-by: yuneng <yuneng@berri.ai>
266 lines
9.3 KiB
Python
266 lines
9.3 KiB
Python
"""
|
|
Unit tests for Replicate provider, particularly testing DeepSeek models
|
|
"""
|
|
|
|
import json
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.llms.replicate.chat.handler import (
|
|
async_completion,
|
|
)
|
|
from litellm.llms.replicate.chat.handler import (
|
|
completion as replicate_completion,
|
|
)
|
|
|
|
|
|
class TestReplicateStartingStatus:
|
|
"""Test that Replicate handler correctly handles 'starting' status for DeepSeek models"""
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.llms.replicate.chat.handler.get_async_httpx_client")
|
|
async def test_async_completion_handles_starting_status(self, mock_get_client):
|
|
"""Test that async completion polls correctly when status is 'starting'"""
|
|
# Mock the async HTTP client
|
|
mock_client = AsyncMock()
|
|
mock_get_client.return_value = mock_client
|
|
|
|
# Mock the initial POST response (creates prediction)
|
|
post_response = Mock()
|
|
post_response.json.return_value = {
|
|
"id": "test-prediction-id",
|
|
"urls": {
|
|
"get": "https://api.replicate.com/v1/predictions/test-id",
|
|
"cancel": "https://api.replicate.com/v1/predictions/test-id/cancel",
|
|
},
|
|
}
|
|
mock_client.post = AsyncMock(return_value=post_response)
|
|
|
|
# Mock GET responses - first 'starting', then 'processing', then 'succeeded'
|
|
get_response_starting = Mock()
|
|
get_response_starting.status_code = 200
|
|
get_response_starting.json.return_value = {
|
|
"id": "test-prediction-id",
|
|
"status": "starting",
|
|
"output": None,
|
|
}
|
|
|
|
get_response_processing = Mock()
|
|
get_response_processing.status_code = 200
|
|
get_response_processing.json.return_value = {
|
|
"id": "test-prediction-id",
|
|
"status": "processing",
|
|
"output": None,
|
|
}
|
|
|
|
get_response_succeeded = Mock()
|
|
get_response_succeeded.status_code = 200
|
|
get_response_succeeded.json.return_value = {
|
|
"id": "test-prediction-id",
|
|
"status": "succeeded",
|
|
"output": ["Hello", " from", " DeepSeek!"],
|
|
}
|
|
get_response_succeeded.text = json.dumps(
|
|
get_response_succeeded.json.return_value
|
|
)
|
|
get_response_succeeded.headers = {}
|
|
|
|
# Configure mock to return different responses on successive calls
|
|
mock_client.get = AsyncMock(
|
|
side_effect=[
|
|
get_response_starting,
|
|
get_response_processing,
|
|
get_response_succeeded,
|
|
]
|
|
)
|
|
|
|
# Create mock model response
|
|
model_response = litellm.ModelResponse()
|
|
model_response.choices = [litellm.Choices()]
|
|
model_response.choices[0].message = litellm.Message(content="")
|
|
|
|
# Create mock logging object
|
|
mock_logging = Mock()
|
|
mock_logging.post_call = Mock()
|
|
|
|
# Call async_completion
|
|
result = await async_completion(
|
|
model_response=model_response,
|
|
model="deepseek-ai/deepseek-v3",
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
encoding=None,
|
|
optional_params={},
|
|
litellm_params={},
|
|
version_id="deepseek-ai/deepseek-v3",
|
|
input_data={"input": {"prompt": "test"}},
|
|
api_key="test-key",
|
|
api_base="https://api.replicate.com",
|
|
logging_obj=mock_logging,
|
|
print_verbose=print,
|
|
headers={"Authorization": "Token test-key"},
|
|
)
|
|
|
|
# Assert that we got responses
|
|
assert result is not None
|
|
assert result.choices[0].message.content == "Hello from DeepSeek!"
|
|
|
|
# Verify that GET was called 3 times (starting, processing, succeeded)
|
|
assert mock_client.get.call_count == 3
|
|
|
|
@patch("litellm.llms.replicate.chat.handler._get_httpx_client")
|
|
def test_sync_completion_handles_starting_status(self, mock_get_client):
|
|
"""Test that sync completion polls correctly when status is 'starting'"""
|
|
# Mock the sync HTTP client
|
|
mock_client = Mock()
|
|
mock_get_client.return_value = mock_client
|
|
|
|
# Mock the initial POST response
|
|
post_response = Mock()
|
|
post_response.json.return_value = {
|
|
"id": "test-prediction-id",
|
|
"urls": {
|
|
"get": "https://api.replicate.com/v1/predictions/test-id",
|
|
"cancel": "https://api.replicate.com/v1/predictions/test-id/cancel",
|
|
},
|
|
}
|
|
mock_client.post.return_value = post_response
|
|
|
|
# Mock GET responses
|
|
get_response_starting = Mock()
|
|
get_response_starting.status_code = 200
|
|
get_response_starting.json.return_value = {
|
|
"id": "test-prediction-id",
|
|
"status": "starting",
|
|
"output": None,
|
|
}
|
|
|
|
get_response_succeeded = Mock()
|
|
get_response_succeeded.status_code = 200
|
|
get_response_succeeded.json.return_value = {
|
|
"id": "test-prediction-id",
|
|
"status": "succeeded",
|
|
"output": ["Hello", " DeepSeek!"],
|
|
}
|
|
get_response_succeeded.text = json.dumps(
|
|
get_response_succeeded.json.return_value
|
|
)
|
|
get_response_succeeded.headers = {}
|
|
|
|
# Configure mock to return different responses
|
|
mock_client.get.side_effect = [get_response_starting, get_response_succeeded]
|
|
|
|
# Create mock objects
|
|
model_response = litellm.ModelResponse()
|
|
model_response.choices = [litellm.Choices()]
|
|
model_response.choices[0].message = litellm.Message(content="")
|
|
|
|
mock_logging = Mock()
|
|
mock_logging.post_call = Mock()
|
|
|
|
# Call completion with mock_response to avoid actual API call
|
|
with patch("time.sleep"): # Skip sleep delays in test
|
|
result = replicate_completion(
|
|
model="deepseek-ai/deepseek-v3",
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
api_base="https://api.replicate.com",
|
|
model_response=model_response,
|
|
print_verbose=print,
|
|
optional_params={},
|
|
litellm_params={},
|
|
logging_obj=mock_logging,
|
|
api_key="test-key",
|
|
encoding=None,
|
|
headers={},
|
|
)
|
|
|
|
# Assert results
|
|
assert result is not None
|
|
assert result.choices[0].message.content == "Hello DeepSeek!"
|
|
|
|
# Verify GET was called multiple times
|
|
assert mock_client.get.call_count >= 1
|
|
|
|
|
|
class TestReplicateOutputFormats:
|
|
"""Test that Replicate handler handles different output formats from models"""
|
|
|
|
def test_transform_response_list_output(self):
|
|
"""Test standard list output format"""
|
|
from litellm.llms.replicate.chat.transformation import ReplicateConfig
|
|
|
|
config = ReplicateConfig()
|
|
|
|
# Mock response with list output
|
|
mock_response = Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"status": "succeeded",
|
|
"output": ["Hello", " ", "world"],
|
|
}
|
|
mock_response.text = json.dumps(mock_response.json.return_value)
|
|
mock_response.headers = {}
|
|
|
|
model_response = litellm.ModelResponse()
|
|
model_response.choices = [litellm.Choices()]
|
|
model_response.choices[0].message = litellm.Message(content="")
|
|
|
|
mock_logging = Mock()
|
|
mock_logging.post_call = Mock()
|
|
|
|
result = config.transform_response(
|
|
model="meta/llama-2-70b-chat",
|
|
raw_response=mock_response,
|
|
model_response=model_response,
|
|
logging_obj=mock_logging,
|
|
request_data={"input": {"prompt": "test"}},
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
optional_params={},
|
|
litellm_params={},
|
|
encoding=None,
|
|
api_key="test-key",
|
|
)
|
|
|
|
assert result.choices[0].message.content == "Hello world"
|
|
|
|
def test_transform_response_string_output(self):
|
|
"""Test string output format (as used by some DeepSeek models)"""
|
|
from litellm.llms.replicate.chat.transformation import ReplicateConfig
|
|
|
|
config = ReplicateConfig()
|
|
|
|
# Mock response with string output
|
|
mock_response = Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"status": "succeeded",
|
|
"output": "Hello from DeepSeek",
|
|
}
|
|
mock_response.text = json.dumps(mock_response.json.return_value)
|
|
mock_response.headers = {}
|
|
|
|
model_response = litellm.ModelResponse()
|
|
model_response.choices = [litellm.Choices()]
|
|
model_response.choices[0].message = litellm.Message(content="")
|
|
|
|
mock_logging = Mock()
|
|
mock_logging.post_call = Mock()
|
|
|
|
result = config.transform_response(
|
|
model="deepseek-ai/deepseek-v3",
|
|
raw_response=mock_response,
|
|
model_response=model_response,
|
|
logging_obj=mock_logging,
|
|
request_data={"input": {"prompt": "test"}},
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
optional_params={},
|
|
litellm_params={},
|
|
encoding=None,
|
|
api_key="test-key",
|
|
)
|
|
|
|
assert result.choices[0].message.content == "Hello from DeepSeek"
|
|
|
|
|
|
# Integration test (requires actual API key - skip in CI)
|