diff --git a/litellm/llms/replicate/chat/handler.py b/litellm/llms/replicate/chat/handler.py index e4bb64fed7..4c75db5abc 100644 --- a/litellm/llms/replicate/chat/handler.py +++ b/litellm/llms/replicate/chat/handler.py @@ -213,7 +213,7 @@ def completion( response = httpx_client.get(url=prediction_url, headers=headers) if ( response.status_code == 200 - and response.json().get("status") == "processing" + and response.json().get("status") in ["processing", "starting"] ): continue return litellm.ReplicateConfig().transform_response( @@ -284,7 +284,7 @@ async def async_completion( response = await async_handler.get(url=prediction_url, headers=headers) if ( response.status_code == 200 - and response.json().get("status") == "processing" + and response.json().get("status") in ["processing", "starting"] ): continue return litellm.ReplicateConfig().transform_response( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 43dc57c34c..fa49305218 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24431,6 +24431,300 @@ "output_cost_per_token": 1e-06, "supports_tool_choice": true }, + "replicate/openai/gpt-5": { + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 1.0e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicateopenai/gpt-oss-20b": { + "input_cost_per_token": 9.0e-08, + "output_cost_per_token": 3.6e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/anthropic/claude-4.5-haiku": { + "input_cost_per_token": 1.0e-06, + "output_cost_per_token": 5.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/ibm-granite/granite-3.3-8b-instruct": { + "input_cost_per_token": 3.0e-08, + "output_cost_per_token": 2.5e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/openai/gpt-4o": { + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 1.0e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_audio_input": true, + "supports_audio_output": true + }, + "replicate/openai/o4-mini": { + "input_cost_per_token": 1.0e-06, + "output_cost_per_token": 4.0e-06, + "output_cost_per_reasoning_token": 4.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_reasoning": true, + "supports_system_messages": true + }, + "replicate/openai/o1-mini": { + "input_cost_per_token": 1.1e-06, + "output_cost_per_token": 4.4e-06, + "output_cost_per_reasoning_token": 4.4e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_reasoning": true, + "supports_system_messages": true + }, + "replicate/openai/o1": { + "input_cost_per_token": 1.5e-05, + "output_cost_per_token": 6.0e-05, + "output_cost_per_reasoning_token": 6.0e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_reasoning": true, + "supports_system_messages": true + }, + "replicate/openai/gpt-4o-mini": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6.0e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/qwen/qwen3-235b-a22b-instruct-2507": { + "input_cost_per_token": 2.64e-07, + "output_cost_per_token": 1.06e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/anthropic/claude-4-sonnet": { + "input_cost_per_token": 3.0e-06, + "output_cost_per_token": 1.5e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/deepseek-ai/deepseek-v3": { + "input_cost_per_token": 1.45e-06, + "output_cost_per_token": 1.45e-06, + "litellm_provider": "replicate", + "mode": "chat", + "max_input_tokens": 65536, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/anthropic/claude-3.7-sonnet": { + "input_cost_per_token": 3.0e-06, + "output_cost_per_token": 1.5e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/anthropic/claude-3.5-haiku": { + "input_cost_per_token": 1.0e-06, + "output_cost_per_token": 5.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/anthropic/claude-3.5-sonnet": { + "input_cost_per_token": 3.75e-06, + "output_cost_per_token": 1.875e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/google/gemini-3-pro": { + "input_cost_per_token": 2.0e-06, + "output_cost_per_token": 1.2e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/anthropic/claude-4.5-sonnet": { + "input_cost_per_token": 3.0e-06, + "output_cost_per_token": 1.5e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/openai/gpt-4.1": { + "input_cost_per_token": 2.0e-06, + "output_cost_per_token": 8.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/openai/gpt-4.1-nano": { + "input_cost_per_token": 1.0e-07, + "output_cost_per_token": 4.0e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/openai/gpt-4.1-mini": { + "input_cost_per_token": 4.0e-07, + "output_cost_per_token": 1.6e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/openai/gpt-5-nano": { + "input_cost_per_token": 5.0e-08, + "output_cost_per_token": 4.0e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/openai/gpt-5-mini": { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/google/gemini-2.5-flash": { + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/openai/gpt-oss-120b": { + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 7.2e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/deepseek-ai/deepseek-v3.1": { + "input_cost_per_token": 6.72e-07, + "output_cost_per_token": 2.016e-06, + "litellm_provider": "replicate", + "mode": "chat", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true + }, + "replicate/xai/grok-4": { + "input_cost_per_token": 7.2e-06, + "output_cost_per_token": 3.6e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/deepseek-ai/deepseek-r1": { + "input_cost_per_token": 3.75e-06, + "output_cost_per_token": 1.0e-05, + "output_cost_per_reasoning_token": 1.0e-05, + "litellm_provider": "replicate", + "mode": "chat", + "max_input_tokens": 65536, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_reasoning": true, + "supports_system_messages": true + }, "rerank-english-v2.0": { "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 43dc57c34c..fa49305218 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24431,6 +24431,300 @@ "output_cost_per_token": 1e-06, "supports_tool_choice": true }, + "replicate/openai/gpt-5": { + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 1.0e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicateopenai/gpt-oss-20b": { + "input_cost_per_token": 9.0e-08, + "output_cost_per_token": 3.6e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/anthropic/claude-4.5-haiku": { + "input_cost_per_token": 1.0e-06, + "output_cost_per_token": 5.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/ibm-granite/granite-3.3-8b-instruct": { + "input_cost_per_token": 3.0e-08, + "output_cost_per_token": 2.5e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/openai/gpt-4o": { + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 1.0e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_audio_input": true, + "supports_audio_output": true + }, + "replicate/openai/o4-mini": { + "input_cost_per_token": 1.0e-06, + "output_cost_per_token": 4.0e-06, + "output_cost_per_reasoning_token": 4.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_reasoning": true, + "supports_system_messages": true + }, + "replicate/openai/o1-mini": { + "input_cost_per_token": 1.1e-06, + "output_cost_per_token": 4.4e-06, + "output_cost_per_reasoning_token": 4.4e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_reasoning": true, + "supports_system_messages": true + }, + "replicate/openai/o1": { + "input_cost_per_token": 1.5e-05, + "output_cost_per_token": 6.0e-05, + "output_cost_per_reasoning_token": 6.0e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_reasoning": true, + "supports_system_messages": true + }, + "replicate/openai/gpt-4o-mini": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6.0e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/qwen/qwen3-235b-a22b-instruct-2507": { + "input_cost_per_token": 2.64e-07, + "output_cost_per_token": 1.06e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/anthropic/claude-4-sonnet": { + "input_cost_per_token": 3.0e-06, + "output_cost_per_token": 1.5e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/deepseek-ai/deepseek-v3": { + "input_cost_per_token": 1.45e-06, + "output_cost_per_token": 1.45e-06, + "litellm_provider": "replicate", + "mode": "chat", + "max_input_tokens": 65536, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/anthropic/claude-3.7-sonnet": { + "input_cost_per_token": 3.0e-06, + "output_cost_per_token": 1.5e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/anthropic/claude-3.5-haiku": { + "input_cost_per_token": 1.0e-06, + "output_cost_per_token": 5.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/anthropic/claude-3.5-sonnet": { + "input_cost_per_token": 3.75e-06, + "output_cost_per_token": 1.875e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/google/gemini-3-pro": { + "input_cost_per_token": 2.0e-06, + "output_cost_per_token": 1.2e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/anthropic/claude-4.5-sonnet": { + "input_cost_per_token": 3.0e-06, + "output_cost_per_token": 1.5e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "replicate/openai/gpt-4.1": { + "input_cost_per_token": 2.0e-06, + "output_cost_per_token": 8.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/openai/gpt-4.1-nano": { + "input_cost_per_token": 1.0e-07, + "output_cost_per_token": 4.0e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/openai/gpt-4.1-mini": { + "input_cost_per_token": 4.0e-07, + "output_cost_per_token": 1.6e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/openai/gpt-5-nano": { + "input_cost_per_token": 5.0e-08, + "output_cost_per_token": 4.0e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/openai/gpt-5-mini": { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2.0e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/google/gemini-2.5-flash": { + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_response_schema": true + }, + "replicate/openai/gpt-oss-120b": { + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 7.2e-07, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/deepseek-ai/deepseek-v3.1": { + "input_cost_per_token": 6.72e-07, + "output_cost_per_token": 2.016e-06, + "litellm_provider": "replicate", + "mode": "chat", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true + }, + "replicate/xai/grok-4": { + "input_cost_per_token": 7.2e-06, + "output_cost_per_token": 3.6e-05, + "litellm_provider": "replicate", + "mode": "chat", + "supports_function_calling": true, + "supports_system_messages": true + }, + "replicate/deepseek-ai/deepseek-r1": { + "input_cost_per_token": 3.75e-06, + "output_cost_per_token": 1.0e-05, + "output_cost_per_reasoning_token": 1.0e-05, + "litellm_provider": "replicate", + "mode": "chat", + "max_input_tokens": 65536, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_reasoning": true, + "supports_system_messages": true + }, "rerank-english-v2.0": { "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, diff --git a/tests/llm_translation/test_replicate.py b/tests/llm_translation/test_replicate.py new file mode 100644 index 0000000000..dee4f5969e --- /dev/null +++ b/tests/llm_translation/test_replicate.py @@ -0,0 +1,289 @@ +""" +Unit tests for Replicate provider, particularly testing DeepSeek models +""" + +import asyncio +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm import completion +from litellm.llms.replicate.chat.handler import ( + async_completion, + 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) +@pytest.mark.skip(reason="Requires REPLICATE_API_KEY environment variable") +def test_replicate_deepseek_integration(): + """Integration test with actual DeepSeek model on Replicate""" + try: + response = completion( + model="replicate/deepseek-ai/deepseek-v3", + messages=[{"role": "user", "content": "Say 'Hello World' and nothing else"}], + max_tokens=20, + ) + + assert response is not None + assert response.choices[0].message.content is not None + assert len(response.choices[0].message.content) > 0 + print(f"Response: {response.choices[0].message.content}") + + except Exception as e: + pytest.fail(f"Integration test failed: {e}")