wehub-resource-sync
989 lines
33 KiB
Python
989 lines
33 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.gateway.config import EndpointConfig
|
|
from mlflow.gateway.providers.base import PassthroughAction
|
|
from mlflow.gateway.providers.litellm import LiteLLMAdapter, LiteLLMProvider
|
|
from mlflow.gateway.schemas import chat, embeddings
|
|
|
|
TEST_MESSAGE = "This is a test"
|
|
|
|
|
|
def chat_config():
|
|
return {
|
|
"name": "chat",
|
|
"endpoint_type": "llm/v1/chat",
|
|
"model": {
|
|
"provider": "litellm",
|
|
"name": "claude-3-5-sonnet-20241022",
|
|
"config": {
|
|
"litellm_auth_config": {"api_key": "test-key"},
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def chat_config_with_api_base():
|
|
return {
|
|
"name": "chat",
|
|
"endpoint_type": "llm/v1/chat",
|
|
"model": {
|
|
"provider": "litellm",
|
|
"name": "custom-model",
|
|
"config": {
|
|
"litellm_auth_config": {
|
|
"api_key": "test-key",
|
|
"api_base": "https://custom-api.example.com",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def chat_config_with_provider():
|
|
return {
|
|
"name": "chat",
|
|
"endpoint_type": "llm/v1/chat",
|
|
"model": {
|
|
"provider": "litellm",
|
|
"name": "claude-3-5-sonnet-20241022",
|
|
"config": {
|
|
"litellm_provider": "anthropic",
|
|
"litellm_auth_config": {"api_key": "test-key"},
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def embeddings_config():
|
|
return {
|
|
"name": "embeddings",
|
|
"endpoint_type": "llm/v1/embeddings",
|
|
"model": {
|
|
"provider": "litellm",
|
|
"name": "text-embedding-3-small",
|
|
"config": {
|
|
"litellm_auth_config": {"api_key": "test-key"},
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def mock_litellm_chat_response():
|
|
"""Create a mock LiteLLM chat response object."""
|
|
response = mock.MagicMock()
|
|
response.id = "litellm-chat-id"
|
|
response.object = "chat.completion"
|
|
response.created = 1234567890
|
|
response.model = "claude-3-5-sonnet-20241022"
|
|
|
|
choice = mock.MagicMock()
|
|
choice.index = 0
|
|
choice.message = mock.MagicMock()
|
|
choice.message.role = "assistant"
|
|
choice.message.content = TEST_MESSAGE
|
|
choice.message.tool_calls = None
|
|
choice.finish_reason = "stop"
|
|
|
|
response.choices = [choice]
|
|
response.usage = mock.MagicMock()
|
|
response.usage.prompt_tokens = 10
|
|
response.usage.completion_tokens = 20
|
|
response.usage.total_tokens = 30
|
|
response.usage.prompt_tokens_details = None
|
|
|
|
return response
|
|
|
|
|
|
def mock_litellm_embeddings_response():
|
|
"""Create a mock LiteLLM embeddings response object."""
|
|
response = mock.MagicMock()
|
|
response.model = "text-embedding-3-small"
|
|
|
|
data = mock.MagicMock()
|
|
data.__getitem__ = lambda self, key: [0.1, 0.2, 0.3] if key == "embedding" else None
|
|
response.data = [data]
|
|
|
|
response.usage = mock.MagicMock()
|
|
response.usage.prompt_tokens = 5
|
|
response.usage.total_tokens = 5
|
|
|
|
return response
|
|
|
|
|
|
def test_litellm_not_installed():
|
|
config = chat_config()
|
|
with mock.patch("importlib.util.find_spec", return_value=None):
|
|
with pytest.raises(MlflowException, match="pip install litellm"):
|
|
LiteLLMProvider(EndpointConfig(**config))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat():
|
|
config = chat_config()
|
|
mock_response = mock_litellm_chat_response()
|
|
|
|
with mock.patch("litellm.acompletion", return_value=mock_response) as mock_completion:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {
|
|
"messages": [{"role": "user", "content": TEST_MESSAGE}],
|
|
"temperature": 0.7,
|
|
"max_tokens": 100,
|
|
}
|
|
response = await provider.chat(chat.RequestPayload(**payload))
|
|
|
|
assert response.id == "litellm-chat-id"
|
|
assert response.object == "chat.completion"
|
|
assert response.model == "claude-3-5-sonnet-20241022"
|
|
assert len(response.choices) == 1
|
|
assert response.choices[0].message.content == TEST_MESSAGE
|
|
assert response.usage.prompt_tokens == 10
|
|
assert response.usage.completion_tokens == 20
|
|
assert response.usage.total_tokens == 30
|
|
|
|
# Verify litellm was called with correct parameters
|
|
mock_completion.assert_called_once()
|
|
call_kwargs = mock_completion.call_args[1]
|
|
assert call_kwargs["model"] == "claude-3-5-sonnet-20241022"
|
|
assert call_kwargs["messages"] == [{"role": "user", "content": TEST_MESSAGE}]
|
|
assert call_kwargs["temperature"] == 0.7
|
|
assert call_kwargs["max_tokens"] == 100
|
|
assert call_kwargs["api_key"] == "test-key"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_with_api_base():
|
|
config = chat_config_with_api_base()
|
|
mock_response = mock_litellm_chat_response()
|
|
|
|
with mock.patch("litellm.acompletion", return_value=mock_response) as mock_completion:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {"messages": [{"role": "user", "content": TEST_MESSAGE}]}
|
|
await provider.chat(chat.RequestPayload(**payload))
|
|
|
|
# Verify API base is passed
|
|
call_kwargs = mock_completion.call_args[1]
|
|
assert call_kwargs["api_base"] == "https://custom-api.example.com"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_with_provider_prefix():
|
|
config = chat_config_with_provider()
|
|
mock_response = mock_litellm_chat_response()
|
|
|
|
with mock.patch("litellm.acompletion", return_value=mock_response) as mock_completion:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {"messages": [{"role": "user", "content": TEST_MESSAGE}]}
|
|
await provider.chat(chat.RequestPayload(**payload))
|
|
|
|
# Verify model name includes provider prefix
|
|
call_kwargs = mock_completion.call_args[1]
|
|
assert call_kwargs["model"] == "anthropic/claude-3-5-sonnet-20241022"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_stream():
|
|
config = chat_config()
|
|
|
|
# Create mock streaming chunks
|
|
async def mock_stream():
|
|
chunk1 = mock.MagicMock()
|
|
chunk1.id = "chunk-1"
|
|
chunk1.object = "chat.completion.chunk"
|
|
chunk1.created = 1234567890
|
|
chunk1.model = "claude-3-5-sonnet-20241022"
|
|
chunk1.usage = None
|
|
choice1 = mock.MagicMock()
|
|
choice1.index = 0
|
|
choice1.delta = mock.MagicMock(spec=["role", "content"])
|
|
choice1.delta.role = "assistant"
|
|
choice1.delta.content = "Hello"
|
|
choice1.finish_reason = None
|
|
chunk1.choices = [choice1]
|
|
yield chunk1
|
|
|
|
chunk2 = mock.MagicMock()
|
|
chunk2.id = "chunk-2"
|
|
chunk2.object = "chat.completion.chunk"
|
|
chunk2.created = 1234567890
|
|
chunk2.model = "claude-3-5-sonnet-20241022"
|
|
chunk2.usage = None
|
|
choice2 = mock.MagicMock()
|
|
choice2.index = 0
|
|
choice2.delta = mock.MagicMock(spec=["content"])
|
|
choice2.delta.content = " world"
|
|
choice2.finish_reason = "stop"
|
|
chunk2.choices = [choice2]
|
|
yield chunk2
|
|
|
|
# Final usage-only chunk (litellm emits this when include_usage is set).
|
|
chunk3 = mock.MagicMock()
|
|
chunk3.id = "chunk-3"
|
|
chunk3.object = "chat.completion.chunk"
|
|
chunk3.created = 1234567890
|
|
chunk3.model = "claude-3-5-sonnet-20241022"
|
|
chunk3.choices = []
|
|
chunk3.usage = mock.MagicMock()
|
|
chunk3.usage.prompt_tokens = 10
|
|
chunk3.usage.completion_tokens = 20
|
|
chunk3.usage.total_tokens = 30
|
|
chunk3.usage.prompt_tokens_details = None
|
|
yield chunk3
|
|
|
|
with mock.patch("litellm.acompletion", return_value=mock_stream()) as mock_completion:
|
|
provider = LiteLLMProvider(EndpointConfig(**config), enable_tracing=True)
|
|
payload = {
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"stream": True,
|
|
}
|
|
|
|
chunks = [chunk async for chunk in provider.chat_stream(chat.RequestPayload(**payload))]
|
|
|
|
assert len(chunks) == 3
|
|
assert chunks[0].choices[0].delta.content == "Hello"
|
|
assert chunks[1].choices[0].delta.content == " world"
|
|
assert chunks[1].choices[0].finish_reason == "stop"
|
|
# Final chunk forwards token usage instead of content.
|
|
assert chunks[2].usage.prompt_tokens == 10
|
|
assert chunks[2].usage.completion_tokens == 20
|
|
assert chunks[2].usage.total_tokens == 30
|
|
# No cache details reported by the provider => field is omitted.
|
|
assert chunks[2].usage.prompt_tokens_details is None
|
|
|
|
# Verify stream parameter was set
|
|
call_kwargs = mock_completion.call_args[1]
|
|
assert call_kwargs["stream"] is True
|
|
assert call_kwargs["stream_options"]["include_usage"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_stream_forwards_zero_token_usage():
|
|
config = chat_config()
|
|
|
|
async def mock_stream():
|
|
chunk = mock.MagicMock()
|
|
chunk.id = "chunk-1"
|
|
chunk.object = "chat.completion.chunk"
|
|
chunk.created = 1234567890
|
|
chunk.model = "claude-3-5-sonnet-20241022"
|
|
chunk.choices = []
|
|
chunk.usage = mock.MagicMock()
|
|
chunk.usage.prompt_tokens = 0
|
|
chunk.usage.completion_tokens = 0
|
|
chunk.usage.total_tokens = 0
|
|
chunk.usage.prompt_tokens_details = None
|
|
yield chunk
|
|
|
|
with mock.patch("litellm.acompletion", return_value=mock_stream()):
|
|
provider = LiteLLMProvider(EndpointConfig(**config), enable_tracing=True)
|
|
payload = {"messages": [{"role": "user", "content": "Hi"}], "stream": True}
|
|
chunks = [chunk async for chunk in provider.chat_stream(chat.RequestPayload(**payload))]
|
|
|
|
# A valid all-zero usage payload must still be forwarded, not dropped.
|
|
assert chunks[-1].usage.total_tokens == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_stream_forwards_cached_tokens():
|
|
config = chat_config()
|
|
|
|
async def mock_stream():
|
|
chunk = mock.MagicMock()
|
|
chunk.id = "chunk-1"
|
|
chunk.object = "chat.completion.chunk"
|
|
chunk.created = 1234567890
|
|
chunk.model = "claude-3-5-sonnet-20241022"
|
|
chunk.choices = []
|
|
chunk.usage = mock.MagicMock()
|
|
chunk.usage.prompt_tokens = 100
|
|
chunk.usage.completion_tokens = 10
|
|
chunk.usage.total_tokens = 110
|
|
chunk.usage.prompt_tokens_details = mock.MagicMock()
|
|
chunk.usage.prompt_tokens_details.cached_tokens = 80
|
|
yield chunk
|
|
|
|
with mock.patch("litellm.acompletion", return_value=mock_stream()):
|
|
provider = LiteLLMProvider(EndpointConfig(**config), enable_tracing=True)
|
|
payload = {"messages": [{"role": "user", "content": "Hi"}], "stream": True}
|
|
chunks = [chunk async for chunk in provider.chat_stream(chat.RequestPayload(**payload))]
|
|
|
|
# Cache-read tokens (a subset of prompt_tokens) are forwarded so clients can price cache hits.
|
|
assert chunks[-1].usage.prompt_tokens == 100
|
|
assert chunks[-1].usage.prompt_tokens_details.cached_tokens == 80
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embeddings():
|
|
config = embeddings_config()
|
|
mock_response = mock_litellm_embeddings_response()
|
|
|
|
with mock.patch("litellm.aembedding", return_value=mock_response) as mock_embedding:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {"input": "Hello world"}
|
|
response = await provider.embeddings(embeddings.RequestPayload(**payload))
|
|
|
|
assert response.model == "text-embedding-3-small"
|
|
assert len(response.data) == 1
|
|
assert response.data[0].embedding == [0.1, 0.2, 0.3]
|
|
assert response.usage.prompt_tokens == 5
|
|
assert response.usage.total_tokens == 5
|
|
|
|
# Verify litellm was called with correct parameters
|
|
mock_embedding.assert_called_once()
|
|
call_kwargs = mock_embedding.call_args[1]
|
|
assert call_kwargs["model"] == "text-embedding-3-small"
|
|
assert call_kwargs["input"] == "Hello world"
|
|
assert call_kwargs["api_key"] == "test-key"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embeddings_batch():
|
|
config = embeddings_config()
|
|
|
|
# Create mock response for batch
|
|
response = mock.MagicMock()
|
|
response.model = "text-embedding-3-small"
|
|
|
|
data1 = mock.MagicMock()
|
|
data1.__getitem__ = lambda self, key: [0.1, 0.2, 0.3] if key == "embedding" else None
|
|
data2 = mock.MagicMock()
|
|
data2.__getitem__ = lambda self, key: [0.4, 0.5, 0.6] if key == "embedding" else None
|
|
response.data = [data1, data2]
|
|
|
|
response.usage = mock.MagicMock()
|
|
response.usage.prompt_tokens = 10
|
|
response.usage.total_tokens = 10
|
|
|
|
with mock.patch("litellm.aembedding", return_value=response):
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {"input": ["Hello", "World"]}
|
|
response_payload = await provider.embeddings(embeddings.RequestPayload(**payload))
|
|
|
|
assert len(response_payload.data) == 2
|
|
assert response_payload.data[0].embedding == [0.1, 0.2, 0.3]
|
|
assert response_payload.data[1].embedding == [0.4, 0.5, 0.6]
|
|
|
|
|
|
def test_adapter_chat_to_model():
|
|
config = EndpointConfig(**chat_config())
|
|
payload = {
|
|
"messages": [{"role": "user", "content": TEST_MESSAGE}],
|
|
"temperature": 0.7,
|
|
}
|
|
|
|
result = LiteLLMAdapter.chat_to_model(payload, config)
|
|
|
|
assert result["model"] == "claude-3-5-sonnet-20241022"
|
|
assert result["messages"] == [{"role": "user", "content": TEST_MESSAGE}]
|
|
assert result["temperature"] == 0.7
|
|
|
|
|
|
def test_adapter_embeddings_to_model():
|
|
config = EndpointConfig(**embeddings_config())
|
|
payload = {"input": TEST_MESSAGE}
|
|
|
|
result = LiteLLMAdapter.embeddings_to_model(payload, config)
|
|
|
|
assert result["model"] == "text-embedding-3-small"
|
|
assert result["input"] == TEST_MESSAGE
|
|
|
|
|
|
def test_adapter_chat_to_model_with_provider():
|
|
config = EndpointConfig(**chat_config_with_provider())
|
|
payload = {
|
|
"messages": [{"role": "user", "content": TEST_MESSAGE}],
|
|
"temperature": 0.7,
|
|
}
|
|
|
|
result = LiteLLMAdapter.chat_to_model(payload, config)
|
|
|
|
assert result["model"] == "anthropic/claude-3-5-sonnet-20241022"
|
|
assert result["messages"] == [{"role": "user", "content": TEST_MESSAGE}]
|
|
assert result["temperature"] == 0.7
|
|
|
|
|
|
def test_adapter_model_to_chat():
|
|
config = EndpointConfig(**chat_config())
|
|
resp = {
|
|
"id": "test-id",
|
|
"object": "chat.completion",
|
|
"created": 1234567890,
|
|
"model": "test-model",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": TEST_MESSAGE, "tool_calls": None},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
|
}
|
|
|
|
result = LiteLLMAdapter.model_to_chat(resp, config)
|
|
|
|
assert result.id == "test-id"
|
|
assert result.model == "test-model"
|
|
assert len(result.choices) == 1
|
|
assert result.choices[0].message.content == TEST_MESSAGE
|
|
assert result.usage.prompt_tokens == 10
|
|
|
|
|
|
def test_adapter_model_to_embeddings():
|
|
config = EndpointConfig(**embeddings_config())
|
|
resp = {
|
|
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
|
|
"model": "test-model",
|
|
"usage": {"prompt_tokens": 5, "total_tokens": 5},
|
|
}
|
|
|
|
result = LiteLLMAdapter.model_to_embeddings(resp, config)
|
|
|
|
assert result.model == "test-model"
|
|
assert len(result.data) == 1
|
|
assert result.data[0].embedding == [0.1, 0.2, 0.3]
|
|
assert result.data[0].index == 0
|
|
assert result.usage.prompt_tokens == 5
|
|
|
|
|
|
# Passthrough tests
|
|
|
|
|
|
def mock_response_with_model_dump():
|
|
"""Create a mock response object that supports model_dump()."""
|
|
response = mock.MagicMock()
|
|
response.model_dump.return_value = {
|
|
"id": "test-response-id",
|
|
"output": "Test response output",
|
|
"model": "test-model",
|
|
}
|
|
return response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passthrough_openai_responses():
|
|
config = chat_config()
|
|
mock_response = mock_response_with_model_dump()
|
|
|
|
with mock.patch("litellm.aresponses", return_value=mock_response) as mock_aresponses:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {"input": "Hello, world!"}
|
|
|
|
result = await provider.passthrough(
|
|
PassthroughAction.OPENAI_RESPONSES,
|
|
payload,
|
|
headers=None,
|
|
)
|
|
|
|
assert result["id"] == "test-response-id"
|
|
assert result["output"] == "Test response output"
|
|
mock_aresponses.assert_called_once()
|
|
call_kwargs = mock_aresponses.call_args[1]
|
|
assert call_kwargs["model"] == "claude-3-5-sonnet-20241022"
|
|
assert call_kwargs["input"] == "Hello, world!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passthrough_openai_responses_streaming():
|
|
config = chat_config()
|
|
|
|
async def mock_stream():
|
|
for i in range(2):
|
|
chunk = mock.MagicMock()
|
|
chunk.model_dump.return_value = {"chunk": i, "content": f"part{i}"}
|
|
yield chunk
|
|
|
|
with mock.patch("litellm.aresponses", return_value=mock_stream()):
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {"input": "Hello, world!", "stream": True}
|
|
|
|
result = await provider.passthrough(
|
|
PassthroughAction.OPENAI_RESPONSES,
|
|
payload,
|
|
headers=None,
|
|
)
|
|
|
|
chunks = [chunk async for chunk in result]
|
|
assert len(chunks) == 2
|
|
assert b"data:" in chunks[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passthrough_anthropic_messages():
|
|
config = chat_config_with_provider()
|
|
mock_response = mock.MagicMock()
|
|
mock_response.model_dump.return_value = {
|
|
"id": "msg-test-id",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "Hello!"}],
|
|
}
|
|
|
|
with mock.patch(
|
|
"litellm.anthropic.messages.acreate", return_value=mock_response
|
|
) as mock_acreate:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"max_tokens": 100,
|
|
}
|
|
|
|
result = await provider.passthrough(
|
|
PassthroughAction.ANTHROPIC_MESSAGES,
|
|
payload,
|
|
headers=None,
|
|
)
|
|
|
|
assert result["id"] == "msg-test-id"
|
|
assert result["type"] == "message"
|
|
mock_acreate.assert_called_once()
|
|
call_kwargs = mock_acreate.call_args[1]
|
|
assert call_kwargs["model"] == "anthropic/claude-3-5-sonnet-20241022"
|
|
assert call_kwargs["max_tokens"] == 100
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passthrough_anthropic_messages_streaming():
|
|
config = chat_config_with_provider()
|
|
|
|
async def mock_stream():
|
|
# LiteLLM returns raw SSE bytes for Anthropic streaming
|
|
yield b'event: message_start\ndata: {"type":"message_start"}\n\n'
|
|
yield b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n' # noqa: E501
|
|
yield b'event: message_stop\ndata: {"type":"message_stop"}\n\n'
|
|
|
|
with mock.patch(
|
|
"litellm.anthropic.messages.acreate", return_value=mock_stream()
|
|
) as mock_acreate:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"max_tokens": 100,
|
|
"stream": True,
|
|
}
|
|
|
|
result = await provider.passthrough(
|
|
PassthroughAction.ANTHROPIC_MESSAGES,
|
|
payload,
|
|
headers=None,
|
|
)
|
|
|
|
chunks = [chunk async for chunk in result]
|
|
assert len(chunks) == 3
|
|
assert b"message_start" in chunks[0]
|
|
assert b"content_block_delta" in chunks[1]
|
|
assert b"message_stop" in chunks[2]
|
|
|
|
mock_acreate.assert_called_once()
|
|
call_kwargs = mock_acreate.call_args[1]
|
|
assert call_kwargs["stream"] is True
|
|
assert call_kwargs["model"] == "anthropic/claude-3-5-sonnet-20241022"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passthrough_gemini_generate_content():
|
|
config = chat_config()
|
|
mock_response = mock.MagicMock()
|
|
mock_response.model_dump.return_value = {
|
|
"candidates": [{"content": {"parts": [{"text": "Generated content"}]}}],
|
|
}
|
|
|
|
with mock.patch(
|
|
"litellm.google_genai.agenerate_content", return_value=mock_response
|
|
) as mock_agenerate:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {
|
|
"contents": [{"parts": [{"text": "Generate something"}]}],
|
|
}
|
|
|
|
result = await provider.passthrough(
|
|
PassthroughAction.GEMINI_GENERATE_CONTENT,
|
|
payload,
|
|
headers=None,
|
|
)
|
|
|
|
assert "candidates" in result
|
|
mock_agenerate.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passthrough_gemini_stream_generate_content():
|
|
config = chat_config()
|
|
|
|
async def mock_stream():
|
|
for i in range(2):
|
|
chunk = mock.MagicMock()
|
|
chunk.model_dump.return_value = {
|
|
"candidates": [{"content": {"parts": [{"text": f"chunk{i}"}]}}]
|
|
}
|
|
yield chunk
|
|
|
|
# agenerate_content is called with stream=True for streaming
|
|
with mock.patch("litellm.google_genai.agenerate_content", return_value=mock_stream()):
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {
|
|
"contents": [{"parts": [{"text": "Generate something"}]}],
|
|
}
|
|
|
|
result = provider._passthrough_gemini_stream_generate_content({
|
|
"model": "claude-3-5-sonnet-20241022",
|
|
**payload,
|
|
})
|
|
|
|
chunks = [chunk async for chunk in result]
|
|
assert len(chunks) == 2
|
|
assert chunks[0].model_dump() == {
|
|
"candidates": [{"content": {"parts": [{"text": "chunk0"}]}}]
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passthrough_openai_chat():
|
|
config = chat_config()
|
|
mock_response = mock_response_with_model_dump()
|
|
mock_response.model_dump.return_value = {
|
|
"id": "chatcmpl-test",
|
|
"object": "chat.completion",
|
|
"choices": [{"message": {"role": "assistant", "content": "Hello!"}}],
|
|
}
|
|
|
|
with mock.patch("litellm.acompletion", return_value=mock_response) as mock_completion:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
}
|
|
|
|
result = await provider.passthrough(
|
|
PassthroughAction.OPENAI_CHAT,
|
|
payload,
|
|
headers=None,
|
|
)
|
|
|
|
assert result["id"] == "chatcmpl-test"
|
|
assert result["object"] == "chat.completion"
|
|
mock_completion.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passthrough_openai_embeddings():
|
|
config = embeddings_config()
|
|
mock_response = mock.MagicMock()
|
|
mock_response.model_dump.return_value = {
|
|
"object": "list",
|
|
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
|
|
"model": "text-embedding-3-small",
|
|
}
|
|
|
|
with mock.patch("litellm.aembedding", return_value=mock_response) as mock_embedding:
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
payload = {"input": "Hello, world!"}
|
|
|
|
result = await provider.passthrough(
|
|
PassthroughAction.OPENAI_EMBEDDINGS,
|
|
payload,
|
|
headers=None,
|
|
)
|
|
|
|
assert result["object"] == "list"
|
|
assert len(result["data"]) == 1
|
|
mock_embedding.assert_called_once()
|
|
|
|
|
|
def test_response_to_dict_with_model_dump():
|
|
config = chat_config()
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
|
|
response = mock.MagicMock()
|
|
response.model_dump.return_value = {"key": "value"}
|
|
|
|
result = provider._response_to_dict(response)
|
|
assert result == {"key": "value"}
|
|
|
|
|
|
def test_response_to_dict_with_dict_input():
|
|
config = chat_config()
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
|
|
result = provider._response_to_dict({"key": "value"})
|
|
assert result == {"key": "value"}
|
|
|
|
|
|
def test_response_to_dict_with_unknown_type_raises():
|
|
config = chat_config()
|
|
provider = LiteLLMProvider(EndpointConfig(**config))
|
|
|
|
with pytest.raises(TypeError, match="Unexpected response type"):
|
|
provider._response_to_dict("string value")
|
|
|
|
|
|
# Token extraction tests
|
|
|
|
|
|
def test_litellm_extract_passthrough_token_usage_openai_format():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
result = {
|
|
"id": "chatcmpl-123",
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 30,
|
|
},
|
|
}
|
|
token_usage = provider._extract_passthrough_token_usage(PassthroughAction.OPENAI_CHAT, result)
|
|
assert token_usage == {
|
|
"input_tokens": 10,
|
|
"output_tokens": 20,
|
|
"total_tokens": 30,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_passthrough_token_usage_anthropic_format():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
result = {
|
|
"id": "msg_123",
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
},
|
|
}
|
|
token_usage = provider._extract_passthrough_token_usage(
|
|
PassthroughAction.ANTHROPIC_MESSAGES, result
|
|
)
|
|
assert token_usage == {
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"total_tokens": 150,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_passthrough_token_usage_gemini_format():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
result = {
|
|
"candidates": [{"content": {"parts": [{"text": "Hello"}]}}],
|
|
"usageMetadata": {
|
|
"promptTokenCount": 10,
|
|
"candidatesTokenCount": 20,
|
|
"totalTokenCount": 30,
|
|
},
|
|
}
|
|
token_usage = provider._extract_passthrough_token_usage(
|
|
PassthroughAction.GEMINI_GENERATE_CONTENT, result
|
|
)
|
|
assert token_usage == {
|
|
"input_tokens": 10,
|
|
"output_tokens": 20,
|
|
"total_tokens": 30,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_passthrough_token_usage_gemini_with_cached_tokens():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
result = {
|
|
"candidates": [{"content": {"parts": [{"text": "Hello"}]}}],
|
|
"usageMetadata": {
|
|
"promptTokenCount": 10,
|
|
"candidatesTokenCount": 20,
|
|
"totalTokenCount": 30,
|
|
"cachedContentTokenCount": 5,
|
|
},
|
|
}
|
|
token_usage = provider._extract_passthrough_token_usage(
|
|
PassthroughAction.GEMINI_GENERATE_CONTENT, result
|
|
)
|
|
assert token_usage == {
|
|
"input_tokens": 10,
|
|
"output_tokens": 20,
|
|
"total_tokens": 30,
|
|
"cache_read_input_tokens": 5,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_passthrough_token_usage_no_usage():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
result = {"id": "chatcmpl-123", "choices": []}
|
|
token_usage = provider._extract_passthrough_token_usage(PassthroughAction.OPENAI_CHAT, result)
|
|
assert token_usage is None
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_openai_format():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = (
|
|
b'data: {"id":"chatcmpl-123","usage":'
|
|
b'{"prompt_tokens":10,"completion_tokens":20,"total_tokens":30}}\n'
|
|
)
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {
|
|
"input_tokens": 10,
|
|
"output_tokens": 20,
|
|
"total_tokens": 30,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_anthropic_message_start():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = b'data: {"type":"message_start","message":{"usage":{"input_tokens":100}}}\n'
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {"input_tokens": 100}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_anthropic_message_delta():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = b'data: {"type":"message_delta","usage":{"output_tokens":50}}\n'
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
# Method only returns chunk's usage; total is calculated by _stream_passthrough_with_usage
|
|
assert result == {"output_tokens": 50}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_gemini_format():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = (
|
|
b'data: {"candidates":[{"content":{}}],"usageMetadata":'
|
|
b'{"promptTokenCount":10,"candidatesTokenCount":20,"totalTokenCount":30}}\n'
|
|
)
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {
|
|
"input_tokens": 10,
|
|
"output_tokens": 20,
|
|
"total_tokens": 30,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_gemini_with_cached_tokens():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = (
|
|
b'data: {"candidates":[{"content":{}}],"usageMetadata":'
|
|
b'{"promptTokenCount":10,"candidatesTokenCount":20,"totalTokenCount":30,"cachedContentTokenCount":5}}\n'
|
|
)
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {
|
|
"input_tokens": 10,
|
|
"output_tokens": 20,
|
|
"total_tokens": 30,
|
|
"cache_read_input_tokens": 5,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_empty_chunk():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = b""
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_done_chunk():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = b"data: [DONE]\n"
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_invalid_json():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = b"data: {invalid json}\n"
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_responses_api():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
# Responses API returns usage in data.response.usage with input_tokens/output_tokens
|
|
chunk = (
|
|
b'data: {"type":"response.completed","response":{"id":"resp_123",'
|
|
b'"usage":{"input_tokens":9,"output_tokens":65,"total_tokens":74}}}\n'
|
|
)
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {
|
|
"input_tokens": 9,
|
|
"output_tokens": 65,
|
|
"total_tokens": 74,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_passthrough_token_usage_openai_with_cached_tokens():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
result = {
|
|
"id": "chatcmpl-123",
|
|
"usage": {
|
|
"prompt_tokens": 50,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 70,
|
|
"prompt_tokens_details": {"cached_tokens": 30},
|
|
},
|
|
}
|
|
token_usage = provider._extract_passthrough_token_usage(PassthroughAction.OPENAI_CHAT, result)
|
|
assert token_usage == {
|
|
"input_tokens": 50,
|
|
"output_tokens": 20,
|
|
"total_tokens": 70,
|
|
"cache_read_input_tokens": 30,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_passthrough_token_usage_anthropic_with_cached_tokens():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
result = {
|
|
"id": "msg_123",
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"cache_read_input_tokens": 25,
|
|
"cache_creation_input_tokens": 15,
|
|
},
|
|
}
|
|
token_usage = provider._extract_passthrough_token_usage(
|
|
PassthroughAction.ANTHROPIC_MESSAGES, result
|
|
)
|
|
assert token_usage == {
|
|
"input_tokens": 140,
|
|
"output_tokens": 50,
|
|
"total_tokens": 190,
|
|
"cache_read_input_tokens": 25,
|
|
"cache_creation_input_tokens": 15,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_openai_with_cached_tokens():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = (
|
|
b'data: {"id":"chatcmpl-123","usage":'
|
|
b'{"prompt_tokens":50,"completion_tokens":20,"total_tokens":70,'
|
|
b'"prompt_tokens_details":{"cached_tokens":30}}}\n'
|
|
)
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {
|
|
"input_tokens": 50,
|
|
"output_tokens": 20,
|
|
"total_tokens": 70,
|
|
"cache_read_input_tokens": 30,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_anthropic_message_start_with_cached_tokens():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = (
|
|
b'data: {"type":"message_start","message":{"usage":'
|
|
b'{"input_tokens":100,"cache_read_input_tokens":25}}}\n'
|
|
)
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {
|
|
"input_tokens": 100,
|
|
"cache_read_input_tokens": 25,
|
|
}
|
|
|
|
|
|
def test_litellm_extract_streaming_token_usage_responses_api_with_cached_tokens():
|
|
provider = LiteLLMProvider(EndpointConfig(**chat_config()))
|
|
chunk = (
|
|
b'data: {"type":"response.completed","response":{"id":"resp_123",'
|
|
b'"usage":{"input_tokens":100,"output_tokens":50,"total_tokens":150,'
|
|
b'"input_tokens_details":{"cached_tokens":40}}}}\n'
|
|
)
|
|
result = provider._extract_streaming_token_usage(chunk)
|
|
assert result == {
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"total_tokens": 150,
|
|
"cache_read_input_tokens": 40,
|
|
}
|