scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Merge pull request #25920 from BerriAI/litellm_/amazing-almeida

Browse source 
[Test] Mock Bedrock Moonshot tests + [Fix] TogetherAIConfig recursion
This commit is contained in:
yuneng-jiang 2026-04-16 17:55:58 -07:00 • committed by GitHub
commit 724926fba9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 329 additions and 75 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/llms/together_ai/chat.py
View file

@ -8,7 +8,7 @@ Docs: https://docs.together.ai/reference/completions-1
from typing import Optional
from litellm.utils import get_model_info
from litellm.utils import supports_function_calling
from litellm._logging import verbose_logger
from ..openai.chat.gpt_transformation import OpenAIGPTConfig
@ -21,18 +21,23 @@ class TogetherAIConfig(OpenAIGPTConfig):
Docs: https://docs.together.ai/docs/json-mode
"""
supports_function_calling: Optional[bool] = None
# Use supports_function_calling() — which reads _get_model_info_helper
# directly — instead of get_model_info(). get_model_info() calls
# get_supported_openai_params() as its first step, which routes back
# into this method for together_ai models, creating a recursion that
# only terminates when Python's recursion limit or the "not mapped"
# exception in _get_model_info_helper is hit (~332 deep calls).
supports_fc: Optional[bool] = None
try:
model_info = get_model_info(model, custom_llm_provider="together_ai")
supports_function_calling = model_info.get(
"supports_function_calling", False
supports_fc = supports_function_calling(
model, custom_llm_provider="together_ai"
)
except Exception as e:
verbose_logger.debug(f"Error getting supported openai params: {e}")
pass
optional_params = super().get_supported_openai_params(model)
if supports_function_calling is not True:
if supports_fc is not True:
verbose_logger.debug(
"Only some together models support function calling/response_format. Docs - https://docs.together.ai/docs/function-calling"
)

374
tests/llm_translation/test_bedrock_moonshot.py
View file

@ -16,10 +16,13 @@ import pytest
import sys
import os
import json
from typing import Optional
from unittest.mock import AsyncMock, Mock, patch
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm.llms.bedrock.common_utils import get_bedrock_chat_config
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
class TestBedrockMoonshotInvoke(BaseLLMChatTest):
@ -27,17 +30,255 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest):
Test suite for Bedrock Moonshot via invoke route.
Inherits all standard LLM tests from BaseLLMChatTest.
"""
def get_base_completion_call_args(self) -> dict:
litellm._turn_on_debug()
return {
"model": "bedrock/invoke/moonshot.kimi-k2-thinking",
}
def test_tool_call_no_arguments(self, tool_call_no_arguments):
"""Test that tool calls with no arguments is translated correctly."""
pass
# ---------------------------------------------------------------------
# The overrides below replace inherited BaseLLMChatTest tests that would
# otherwise make live AWS Bedrock calls. The live versions were
# consistently crashing llm_translation xdist workers. Each override
# patches the HTTP client's post() so no network request is sent, and
# asserts on the outgoing request body (and, where needed, parses a
# canned response) — which is what the translation lane is actually
# supposed to cover.
# ---------------------------------------------------------------------
@staticmethod
def _make_moonshot_response(content: str = "Hi!") -> Mock:
"""Build a Mock httpx.Response that AmazonMoonshotConfig.transform_response
(which delegates to MoonshotChatConfig → OpenAI) can parse."""
body = {
"id": "chatcmpl-test",
"object": "chat.completion",
"created": 1234567890,
"model": "moonshot.kimi-k2-thinking",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": content},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15,
},
}
mock_resp = Mock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.text = json.dumps(body)
mock_resp.json = lambda: body
return mock_resp
def _invoke_with_mocked_post(
self,
*,
messages: list,
extra_kwargs: Optional[dict] = None,
response_content: str = "Hi!",
) -> "tuple[Mock, object]":
"""Run a sync litellm.completion() with HTTPHandler.post patched to
return a canned moonshot response. Returns (mock_post, response)."""
client = HTTPHandler()
mock_resp = self._make_moonshot_response(content=response_content)
with patch.object(
client, "post", new=Mock(return_value=mock_resp)
) as mock_post:
response = litellm.completion(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=messages,
aws_access_key_id="fake",
aws_secret_access_key="fake",
aws_region_name="us-west-2",
client=client,
**(extra_kwargs or {}),
)
return mock_post, response
def test_developer_role_translation(self):
"""Verify LiteLLM maps the ``developer`` role to ``system`` on the
outgoing Bedrock invoke request, without hitting the network."""
mock_post, response = self._invoke_with_mocked_post(
messages=[
{"role": "developer", "content": "Be a good bot!"},
{"role": "user", "content": "Hello, how are you?"},
],
)
mock_post.assert_called_once()
body = json.loads(mock_post.call_args.kwargs["data"])
assert body["messages"][0]["role"] == "system"
assert body["messages"][0]["content"] == "Be a good bot!"
assert body["messages"][1]["role"] == "user"
assert response.choices[0].message.content is not None
def test_message_with_name(self):
"""Verify a user message carrying a ``name`` field is serialized into
the outgoing Bedrock invoke request without breaking the call."""
mock_post, response = self._invoke_with_mocked_post(
messages=[{"role": "user", "content": "Hello", "name": "test_name"}],
)
mock_post.assert_called_once()
body = json.loads(mock_post.call_args.kwargs["data"])
assert body["messages"][0]["role"] == "user"
assert body["messages"][0]["content"] == "Hello"
assert response is not None
def test_content_list_handling(self):
"""Verify the inherited content-list-handling test passes against a
mocked moonshot response (no network)."""
mock_post, response = self._invoke_with_mocked_post(
messages=[
{
"role": "user",
"content": [{"type": "text", "text": "Hello, how are you?"}],
}
],
)
mock_post.assert_called_once()
assert response.choices[0].message.content is not None
def test_pydantic_model_input(self):
"""Verify a completion call with a pydantic ``Message`` as input does
not raise and produces a parseable response."""
from litellm import Message
mock_post, response = self._invoke_with_mocked_post(
messages=[Message(content="Hello, how are you?", role="user")],
)
mock_post.assert_called_once()
assert response is not None
@pytest.mark.parametrize("response_format", [{"type": "text"}])
def test_response_format_type_text_with_tool_calls_no_tool_choice(
self, response_format
):
"""Verify response_format + tools + drop_params sends a valid request
and produces a response object."""
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
},
},
"required": ["location"],
},
},
}
]
mock_post, response = self._invoke_with_mocked_post(
messages=[
{"role": "user", "content": "What's the weather like in Boston today?"}
],
extra_kwargs={
"response_format": response_format,
"tools": tools,
"drop_params": True,
},
)
mock_post.assert_called_once()
body = json.loads(mock_post.call_args.kwargs["data"])
assert "tools" in body
assert body["tools"][0]["function"]["name"] == "get_current_weather"
assert response is not None
def test_streaming(self):
"""Verify stream=True routes to the invoke-with-response-stream
endpoint with the messages body. Iteration of the stream itself is
not exercised here — moonshot streaming delegates to the OpenAI
parser and is covered by the OpenAI test suite.
Note: bedrock invoke streaming cannot be intercepted by patching
the caller-supplied client, because ``CustomStreamWrapper.fetch_sync_stream``
at streaming_handler.py invokes the stored ``make_call`` partial with
``client=litellm.module_level_client``, which overrides any client the
caller passed. Patch ``make_sync_call`` at its import site in
``base_invoke_transformation`` so we observe the exact kwargs the
partial was built with at stream-wrapper construction time.
"""
from litellm.utils import CustomStreamWrapper
captured: dict = {}
def fake_make_sync_call(**kwargs):
captured.update(kwargs)
# Return an empty iterator so the stream wrapper's iteration
# doesn't try to parse real bytes.
return iter([])
with patch(
"litellm.llms.bedrock.chat.invoke_transformations."
"base_invoke_transformation.make_sync_call",
new=fake_make_sync_call,
):
response = litellm.completion(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=[
{
"role": "user",
"content": [{"type": "text", "text": "Hello, how are you?"}],
}
],
stream=True,
aws_access_key_id="fake",
aws_secret_access_key="fake",
aws_region_name="us-west-2",
)
assert isinstance(response, CustomStreamWrapper)
# Trigger fetch_sync_stream → make_call(...) → fake_make_sync_call.
try:
next(iter(response))
except StopIteration:
pass
assert captured, "make_sync_call was never invoked"
assert captured["api_base"].endswith("/invoke-with-response-stream")
body = json.loads(captured["data"])
# Bedrock invoke does not put stream=true in the body (the URL
# carries the streaming flag); verify the user message is present.
assert body["messages"][0]["role"] == "user"
async def test_completion_cost(self):
"""Verify LiteLLM computes a positive cost from a mocked Bedrock
Moonshot response, using the local model cost map."""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
mock_response = self._make_moonshot_response()
client = AsyncHTTPHandler()
with patch.object(client, "post", new=AsyncMock(return_value=mock_response)):
response = await litellm.acompletion(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=[{"role": "user", "content": "Hello, how are you?"}],
aws_access_key_id="fake",
aws_secret_access_key="fake",
aws_region_name="us-west-2",
client=client,
)
assert response._hidden_params["response_cost"] > 0
class TestBedrockMoonshotBasic:
"""Unit tests for Bedrock Moonshot configuration and transformations."""
@ -47,7 +288,7 @@ class TestBedrockMoonshotBasic:
config = get_bedrock_chat_config("bedrock/invoke/moonshot.kimi-k2-thinking")
assert config is not None
assert config.__class__.__name__ == "AmazonMoonshotConfig"
def test_provider_detection_converse(self):
"""Test that Bedrock Moonshot converse models are correctly detected."""
config = get_bedrock_chat_config("bedrock/moonshot.kimi-k2-thinking")
@ -62,8 +303,10 @@ class TestBedrockMoonshotBasic:
def test_supported_params(self):
"""Test that supported OpenAI params are correctly defined."""
config = get_bedrock_chat_config("invoke/moonshot.kimi-k2-thinking")
supported_params = config.get_supported_openai_params("moonshot.kimi-k2-thinking")
supported_params = config.get_supported_openai_params(
"moonshot.kimi-k2-thinking"
)
# Should support these params
assert "temperature" in supported_params
assert "max_tokens" in supported_params
@ -71,10 +314,10 @@ class TestBedrockMoonshotBasic:
assert "stream" in supported_params
assert "tools" in supported_params
assert "tool_choice" in supported_params
# Should NOT support stop sequences on Bedrock
assert "stop" not in supported_params
# Should NOT support functions (use tools instead)
assert "functions" not in supported_params
@ -83,20 +326,20 @@ class TestBedrockMoonshotBasic:
from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
AmazonMoonshotConfig,
)
config = AmazonMoonshotConfig()
messages = [{"role": "user", "content": "Hello"}]
# Test that bedrock/invoke/ prefix is stripped
transformed = config.transform_request(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=messages,
optional_params={},
litellm_params={},
headers={}
headers={},
)
# The model ID in the request body should be stripped
assert transformed["model"] == "moonshot.kimi-k2-thinking"
@ -109,21 +352,27 @@ class TestBedrockMoonshotReasoningContent:
from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
AmazonMoonshotConfig,
)
config = AmazonMoonshotConfig()
# Test with reasoning tags
content_with_reasoning = "<reasoning>This is my thought process</reasoning>This is the answer"
reasoning, content = config._extract_reasoning_from_content(content_with_reasoning)
content_with_reasoning = (
"<reasoning>This is my thought process</reasoning>This is the answer"
)
reasoning, content = config._extract_reasoning_from_content(
content_with_reasoning
)
assert reasoning == "This is my thought process"
assert content == "This is the answer"
assert "<reasoning>" not in content
# Test without reasoning tags
content_without_reasoning = "This is just a regular answer"
reasoning, content = config._extract_reasoning_from_content(content_without_reasoning)
reasoning, content = config._extract_reasoning_from_content(
content_without_reasoning
)
assert reasoning is None
assert content == "This is just a regular answer"
@ -134,8 +383,10 @@ class TestBedrockMoonshotToolCalling:
def test_tool_calling_supported(self):
"""Test that tool calling is supported for Kimi K2 Thinking model."""
config = get_bedrock_chat_config("invoke/moonshot.kimi-k2-thinking")
supported_params = config.get_supported_openai_params("moonshot.kimi-k2-thinking")
supported_params = config.get_supported_openai_params(
"moonshot.kimi-k2-thinking"
)
# Kimi K2 Thinking DOES support tool calls (unlike kimi-thinking-preview)
assert "tools" in supported_params
assert "tool_choice" in supported_params
@ -145,13 +396,11 @@ class TestBedrockMoonshotToolCalling:
from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
AmazonMoonshotConfig,
)
config = AmazonMoonshotConfig()
messages = [
{"role": "user", "content": "What's the weather in San Francisco?"}
]
messages = [{"role": "user", "content": "What's the weather in San Francisco?"}]
optional_params = {
"tools": [
{
@ -161,27 +410,25 @@ class TestBedrockMoonshotToolCalling:
"description": "Get the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
},
"required": ["location"]
}
}
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
}
]
}
transformed = config.transform_request(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
headers={},
)
# Verify model ID is stripped
assert transformed["model"] == "moonshot.kimi-k2-thinking"
# Verify tools are included
assert "tools" in transformed
assert len(transformed["tools"]) == 1
@ -193,9 +440,9 @@ class TestBedrockMoonshotToolCalling:
tool_response_message = {
"role": "tool",
"tool_call_id": "call_123",
"content": json.dumps({"temperature": 72, "condition": "sunny"})
"content": json.dumps({"temperature": 72, "condition": "sunny"}),
}
# Verify the message structure
assert tool_response_message["role"] == "tool"
assert "tool_call_id" in tool_response_message
@ -208,8 +455,10 @@ class TestBedrockMoonshotParameterValidation:
def test_stop_sequences_not_supported(self):
"""Test that stop sequences are correctly excluded from supported params."""
config = get_bedrock_chat_config("invoke/moonshot.kimi-k2-thinking")
supported_params = config.get_supported_openai_params("moonshot.kimi-k2-thinking")
supported_params = config.get_supported_openai_params(
"moonshot.kimi-k2-thinking"
)
# Bedrock Moonshot doesn't support stopSequences field
assert "stop" not in supported_params
@ -218,10 +467,12 @@ class TestBedrockMoonshotParameterValidation:
# Moonshot models support temperature 0-1
# This is handled by the parent MoonshotChatConfig class
config = get_bedrock_chat_config("invoke/moonshot.kimi-k2-thinking")
# Verify config exists and can handle temperature
assert config is not None
supported_params = config.get_supported_openai_params("moonshot.kimi-k2-thinking")
supported_params = config.get_supported_openai_params(
"moonshot.kimi-k2-thinking"
)
assert "temperature" in supported_params
@ -233,34 +484,31 @@ class TestBedrockMoonshotTransformations:
from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
AmazonMoonshotConfig,
)
config = AmazonMoonshotConfig()
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello!"}
{"role": "user", "content": "Hello!"},
]
optional_params = {
"temperature": 0.7,
"max_tokens": 100
}
optional_params = {"temperature": 0.7, "max_tokens": 100}
transformed = config.transform_request(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
headers={},
)
# Verify model ID is stripped
assert transformed["model"] == "moonshot.kimi-k2-thinking"
# Verify messages are included
assert "messages" in transformed
assert len(transformed["messages"]) >= 1
# Verify optional params are included
assert transformed["temperature"] == 0.7
assert transformed["max_tokens"] == 100
@ -270,21 +518,21 @@ class TestBedrockMoonshotTransformations:
from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
AmazonMoonshotConfig,
)
config = AmazonMoonshotConfig()
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello!"}
{"role": "user", "content": "Hello!"},
]
transformed = config.transform_request(
model="moonshot.kimi-k2-thinking",
messages=messages,
optional_params={},
litellm_params={},
headers={}
headers={},
)
# System messages should be supported
assert "messages" in transformed

2
tests/llm_translation/test_together_ai.py
View file

@ -20,7 +20,7 @@ import pytest
class TestTogetherAI(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
litellm.set_verbose = True
return {"model": "together_ai/Qwen/Qwen3.5-9B"}
return {"model": "together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo"}
def test_tool_call_no_arguments(self, tool_call_no_arguments):
"""Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833"""

6
tests/local_testing/test_completion.py
View file

@ -65,7 +65,7 @@ def test_completion_custom_provider_model_name():
try:
litellm.cache = None
response = completion(
model="together_ai/Qwen/Qwen3.5-9B",
model="together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo",
messages=messages,
logger_fn=logger_fn,
)
@ -2815,7 +2815,7 @@ def test_customprompt_together_ai():
print(litellm.success_callback)
print(litellm._async_success_callback)
response = completion(
model="together_ai/Qwen/Qwen3.5-9B",
model="together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo",
messages=messages,
roles={
"system": {
@ -3682,7 +3682,7 @@ def test_completion_together_ai_stream():
messages = [{"content": user_message, "role": "user"}]
try:
response = completion(
model="together_ai/Qwen/Qwen3.5-9B",
model="together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo",
messages=messages,
stream=True,
max_tokens=5,

2
tests/local_testing/test_multiple_deployments.py
View file

@ -25,7 +25,7 @@ model_list = [
{
"model_name": "mistral-7b-instruct",
"litellm_params": { # params for litellm completion/embedding call
"model": "together_ai/Qwen/Qwen3.5-9B",
"model": "together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo",
"api_key": os.getenv("TOGETHERAI_API_KEY"),
},
},

2
tests/local_testing/test_text_completion.py
View file

@ -4034,7 +4034,7 @@ def test_async_text_completion_together_ai():
async def test_get_response():
try:
response = await litellm.atext_completion(
model="together_ai/Qwen/Qwen3.5-9B",
model="together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo",
prompt="good morning",
max_tokens=10,
)

1
tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
View file

@ -18,6 +18,7 @@ from litellm.llms.bedrock.common_utils import (
normalize_tool_input_schema_types_for_bedrock_invoke,
remove_custom_field_from_tools,
)
from litellm.constants import BEDROCK_MIN_THINKING_BUDGET_TOKENS
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeMessagesConfig,
AmazonAnthropicClaudeMessagesStreamDecoder,