mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
test(test_anthropic_experimental_pass_through_messages_handler.py): add unit test confirming custom llm respected
This commit is contained in:
parent
1ecf0315e5
commit
ded1734d34
1 changed files with 30 additions and 0 deletions
|
|
@ -34,3 +34,33 @@ def test_anthropic_experimental_pass_through_messages_handler():
|
|||
print(f"Error: {e}")
|
||||
mock_completion.assert_called_once()
|
||||
mock_completion.call_args.kwargs["api_key"] == "test-api-key"
|
||||
|
||||
|
||||
def test_anthropic_experimental_pass_through_messages_handler_custom_llm_provider():
|
||||
"""
|
||||
Test that litellm.completion is called when a custom LLM provider is given
|
||||
"""
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
|
||||
anthropic_messages_handler,
|
||||
)
|
||||
|
||||
with patch("litellm.completion", return_value="test-response") as mock_completion:
|
||||
try:
|
||||
anthropic_messages_handler(
|
||||
max_tokens=100,
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
model="my-custom-model",
|
||||
custom_llm_provider="my-custom-llm",
|
||||
api_key="test-api-key",
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
|
||||
# Assert that litellm.completion was called when using a custom LLM provider
|
||||
mock_completion.assert_called_once()
|
||||
|
||||
# Verify that the custom provider was passed through
|
||||
call_kwargs = mock_completion.call_args.kwargs
|
||||
assert call_kwargs["custom_llm_provider"] == "my-custom-llm"
|
||||
assert call_kwargs["model"] == "my-custom-llm/my-custom-model"
|
||||
assert call_kwargs["api_key"] == "test-api-key"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue