tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-22 18:25:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #18797 from BerriAI/litellm_bedrock_kimi2_model

Browse Source 
[Feat]Add bedrock kimi2 model support
This commit is contained in:
Sameer Kankute
2026-01-08 15:29:12 +05:30
committed by GitHub
parent 6c00f6f342 e5c3fcbb52
commit 4abc94119f
9 changed files with 735 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/providers/bedrock.md
+2 -1
View File
@@ -7,7 +7,7 @@ ALL Bedrock models (Anthropic, Meta, Deepseek, Mistral, Amazon, etc.) are Suppor
| Property | Details |
|-------|-------|
| Description | Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs). |
| Provider Route on LiteLLM | `bedrock/`, [`bedrock/converse/`](#set-converse--invoke-route), [`bedrock/invoke/`](#set-invoke-route), [`bedrock/converse_like/`](#calling-via-internal-proxy), [`bedrock/llama/`](#deepseek-not-r1), [`bedrock/deepseek_r1/`](#deepseek-r1), [`bedrock/qwen3/`](#qwen3-imported-models), [`bedrock/qwen2/`](./bedrock_imported.md#qwen2-imported-models), [`bedrock/openai/`](./bedrock_imported.md#openai-compatible-imported-models-qwen-25-vl-etc) |
| Provider Route on LiteLLM | `bedrock/`, [`bedrock/converse/`](#set-converse--invoke-route), [`bedrock/invoke/`](#set-invoke-route), [`bedrock/converse_like/`](#calling-via-internal-proxy), [`bedrock/llama/`](#deepseek-not-r1), [`bedrock/deepseek_r1/`](#deepseek-r1), [`bedrock/qwen3/`](#qwen3-imported-models), [`bedrock/qwen2/`](./bedrock_imported.md#qwen2-imported-models), [`bedrock/openai/`](./bedrock_imported.md#openai-compatible-imported-models-qwen-25-vl-etc), [`bedrock/moonshot`](./bedrock_imported.md#moonshot-kimi-k2-thinking) |
| Provider Doc | [Amazon Bedrock ↗](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) |
| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/embeddings`, `/images/generations` |
| Rerank Endpoint | `/rerank` |
@@ -1941,6 +1941,7 @@ Here's an example of using a bedrock model with LiteLLM. For a complete list, re
| Mixtral 8x7B Instruct | `completion(model='bedrock/mistral.mixtral-8x7b-instruct-v0:1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| TwelveLabs Pegasus 1.2 (US) | `completion(model='bedrock/us.twelvelabs.pegasus-1-2-v1:0', messages=messages, mediaSource={...})` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| TwelveLabs Pegasus 1.2 (EU) | `completion(model='bedrock/eu.twelvelabs.pegasus-1-2-v1:0', messages=messages, mediaSource={...})` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Moonshot Kimi K2 Thinking | `completion(model='bedrock/moonshot.kimi-k2-thinking', messages=messages)` or `completion(model='bedrock/invoke/moonshot.kimi-k2-thinking', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
## Bedrock Embedding
docs/my-website/docs/providers/bedrock_imported.md
+177 -1
View File
@@ -431,4 +431,180 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
"max_tokens": 300,
"temperature": 0.5
}'
```
```
### Moonshot Kimi K2 Thinking
Moonshot AI's Kimi K2 Thinking model is now available on Amazon Bedrock. This model features advanced reasoning capabilities with automatic reasoning content extraction.
| Property | Details |
|----------|---------|
| Provider Route | `bedrock/moonshot.kimi-k2-thinking`, `bedrock/invoke/moonshot.kimi-k2-thinking` |
| Provider Documentation | [AWS Bedrock Moonshot Announcement ↗](https://aws.amazon.com/about-aws/whats-new/2025/12/amazon-bedrock-fully-managed-open-weight-models/) |
| Supported Parameters | `temperature`, `max_tokens`, `top_p`, `stream`, `tools`, `tool_choice` |
| Special Features | Reasoning content extraction, Tool calling |
#### Supported Features
- **Reasoning Content Extraction**: Automatically extracts `<reasoning>` tags and returns them as `reasoning_content` (similar to OpenAI's o1 models)
- **Tool Calling**: Full support for function/tool calling with tool responses
- **Streaming**: Both streaming and non-streaming responses
- **System Messages**: System message support
#### Basic Usage
<Tabs>
<TabItem value="sdk" label="SDK">
```python title="Moonshot Kimi K2 SDK Usage" showLineNumbers
from litellm import completion
import os
os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-key"
os.environ["AWS_REGION_NAME"] = "us-west-2" # or your preferred region
# Basic completion
response = completion(
model="bedrock/moonshot.kimi-k2-thinking", # or bedrock/invoke/moonshot.kimi-k2-thinking
messages=[
{"role": "user", "content": "What is 2+2? Think step by step."}
],
temperature=0.7,
max_tokens=200
)
print(response.choices[0].message.content)
# Access reasoning content if present
if response.choices[0].message.reasoning_content:
print("Reasoning:", response.choices[0].message.reasoning_content)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml title="config.yaml" showLineNumbers
model_list:
- model_name: kimi-k2
litellm_params:
model: bedrock/moonshot.kimi-k2-thinking
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_region_name: us-west-2
```
**2. Start proxy**
```bash title="Start LiteLLM Proxy" showLineNumbers
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash title="Test Kimi K2 via Proxy" showLineNumbers
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "kimi-k2",
"messages": [
{
"role": "user",
"content": "What is 2+2? Think step by step."
}
],
"temperature": 0.7,
"max_tokens": 200
}'
```
</TabItem>
</Tabs>
#### Tool Calling Example
```python title="Kimi K2 with Tool Calling" showLineNumbers
from litellm import completion
import os
os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-key"
os.environ["AWS_REGION_NAME"] = "us-west-2"
# Tool calling example
response = completion(
model="bedrock/moonshot.kimi-k2-thinking",
messages=[
{"role": "user", "content": "What's the weather in Tokyo?"}
],
tools=[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather in a location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city name"
}
},
"required": ["location"]
}
}
}
]
)
if response.choices[0].message.tool_calls:
tool_call = response.choices[0].message.tool_calls[0]
print(f"Tool called: {tool_call.function.name}")
print(f"Arguments: {tool_call.function.arguments}")
```
#### Streaming Example
```python title="Kimi K2 Streaming" showLineNumbers
from litellm import completion
import os
os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-key"
os.environ["AWS_REGION_NAME"] = "us-west-2"
response = completion(
model="bedrock/moonshot.kimi-k2-thinking",
messages=[
{"role": "user", "content": "Explain quantum computing in simple terms."}
],
stream=True,
temperature=0.7
)
for chunk in response:
if chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="")
# Check for reasoning content in streaming
if hasattr(chunk.choices[0].delta, 'reasoning_content') and chunk.choices[0].delta.reasoning_content:
print(f"\n[Reasoning: {chunk.choices[0].delta.reasoning_content}]")
```
#### Supported Parameters
| Parameter | Type | Description | Supported |
|-----------|------|-------------|-----------|
| `temperature` | float (0-1) | Controls randomness in output | ✅ |
| `max_tokens` | integer | Maximum tokens to generate | ✅ |
| `top_p` | float | Nucleus sampling parameter | ✅ |
| `stream` | boolean | Enable streaming responses | ✅ |
| `tools` | array | Tool/function definitions | ✅ |
| `tool_choice` | string/object | Tool choice specification | ✅ |
| `stop` | array | Stop sequences | ❌ (Not supported on Bedrock) |
litellm/__init__.py
+1
View File
@@ -1343,6 +1343,7 @@ if TYPE_CHECKING:
from .llms.bedrock.chat.invoke_transformations.amazon_llama_transformation import AmazonLlamaConfig as AmazonLlamaConfig
from .llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation import AmazonDeepSeekR1Config as AmazonDeepSeekR1Config
from .llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import AmazonMistralConfig as AmazonMistralConfig
from .llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import AmazonMoonshotConfig as AmazonMoonshotConfig
from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation import AmazonTitanConfig as AmazonTitanConfig
from .llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import AmazonTwelveLabsPegasusConfig as AmazonTwelveLabsPegasusConfig
from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig as AmazonInvokeConfig
litellm/_lazy_imports_registry.py
+2
View File
@@ -165,6 +165,7 @@ LLM_CONFIG_NAMES = (
"AmazonLlamaConfig",
"AmazonDeepSeekR1Config",
"AmazonMistralConfig",
"AmazonMoonshotConfig",
"AmazonTitanConfig",
"AmazonTwelveLabsPegasusConfig",
"AmazonInvokeConfig",
@@ -556,6 +557,7 @@ _LLM_CONFIGS_IMPORT_MAP = {
"AmazonLlamaConfig": (".llms.bedrock.chat.invoke_transformations.amazon_llama_transformation", "AmazonLlamaConfig"),
"AmazonDeepSeekR1Config": (".llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation", "AmazonDeepSeekR1Config"),
"AmazonMistralConfig": (".llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation", "AmazonMistralConfig"),
"AmazonMoonshotConfig": (".llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation", "AmazonMoonshotConfig"),
"AmazonTitanConfig": (".llms.bedrock.chat.invoke_transformations.amazon_titan_transformation", "AmazonTitanConfig"),
"AmazonTwelveLabsPegasusConfig": (".llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation", "AmazonTwelveLabsPegasusConfig"),
"AmazonInvokeConfig": (".llms.bedrock.chat.invoke_transformations.base_invoke_transformation", "AmazonInvokeConfig"),
litellm/constants.py
+1
View File
@@ -909,6 +909,7 @@ BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[
"twelvelabs",
"openai",
"stability",
"moonshot",
]
BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[
litellm/llms/bedrock/base_aws_llm.py
+4
View File
@@ -369,6 +369,10 @@ class BaseAWSLLM:
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
model_id, spec="stability"
)
elif provider == "moonshot" and "moonshot/" in model_id:
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
model_id, spec="moonshot"
)
return model_id
@staticmethod
litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py
+256
View File
@@ -0,0 +1,256 @@
"""
Transformation for Bedrock Moonshot AI (Kimi K2) models.
Supports the Kimi K2 Thinking model available on Amazon Bedrock.
Model format: bedrock/moonshot.kimi-k2-thinking-v1:0
Reference: https://aws.amazon.com/about-aws/whats-new/2025/12/amazon-bedrock-fully-managed-open-weight-models/
"""
from typing import TYPE_CHECKING, Any, List, Optional, Union
import re
import httpx
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
AmazonInvokeConfig,
)
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.moonshot.chat.transformation import MoonshotChatConfig
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.types.utils import ModelResponse
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig):
"""
Configuration for Bedrock Moonshot AI (Kimi K2) models.
Reference:
https://aws.amazon.com/about-aws/whats-new/2025/12/amazon-bedrock-fully-managed-open-weight-models/
https://platform.moonshot.ai/docs/api/chat
Supported Params for the Amazon / Moonshot models:
- `max_tokens` (integer) max tokens
- `temperature` (float) temperature for model (0-1 for Moonshot)
- `top_p` (float) top p for model
- `stream` (bool) whether to stream responses
- `tools` (list) tool definitions (supported on kimi-k2-thinking)
- `tool_choice` (str|dict) tool choice specification (supported on kimi-k2-thinking)
NOT Supported on Bedrock:
- `stop` sequences (Bedrock doesn't support stopSequences field for this model)
Note: The kimi-k2-thinking model DOES support tool calls, unlike kimi-thinking-preview.
"""
def __init__(self, **kwargs):
AmazonInvokeConfig.__init__(self, **kwargs)
MoonshotChatConfig.__init__(self, **kwargs)
@property
def custom_llm_provider(self) -> Optional[str]:
return "bedrock"
def _get_model_id(self, model: str) -> str:
"""
Extract the actual model ID from the LiteLLM model name.
Removes routing prefixes like:
- bedrock/invoke/moonshot.kimi-k2-thinking -> moonshot.kimi-k2-thinking
- invoke/moonshot.kimi-k2-thinking -> moonshot.kimi-k2-thinking
- moonshot.kimi-k2-thinking -> moonshot.kimi-k2-thinking
"""
# Remove bedrock/ prefix if present
if model.startswith("bedrock/"):
model = model[8:]
# Remove invoke/ prefix if present
if model.startswith("invoke/"):
model = model[7:]
# Remove any provider prefix (e.g., moonshot/)
if "/" in model and not model.startswith("arn:"):
parts = model.split("/", 1)
if len(parts) == 2:
model = parts[1]
return model
def get_supported_openai_params(self, model: str) -> List[str]:
"""
Get the supported OpenAI params for Moonshot AI models on Bedrock.
Bedrock-specific limitations:
- stopSequences field is not supported on Bedrock (unlike native Moonshot API)
- functions parameter is not supported (use tools instead)
- tool_choice doesn't support "required" value
Note: kimi-k2-thinking DOES support tool calls (unlike kimi-thinking-preview)
The parent MoonshotChatConfig class handles the kimi-thinking-preview exclusion.
"""
excluded_params: List[str] = ["functions", "stop"] # Bedrock doesn't support stopSequences
base_openai_params = super(MoonshotChatConfig, self).get_supported_openai_params(model=model)
final_params: List[str] = []
for param in base_openai_params:
if param not in excluded_params:
final_params.append(param)
return final_params
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI parameters to Moonshot AI parameters for Bedrock.
Handles Moonshot AI specific limitations:
- tool_choice doesn't support "required" value
- Temperature <0.3 limitation for n>1
- Temperature range is [0, 1] (not [0, 2] like OpenAI)
"""
return MoonshotChatConfig.map_openai_params(
self,
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=drop_params,
)
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform the request for Bedrock Moonshot AI models.
Uses the Moonshot transformation logic which handles:
- Converting content lists to strings (Moonshot doesn't support list format)
- Adding tool_choice="required" message if needed
- Temperature and parameter validation
"""
# Filter out AWS credentials using the existing method from BaseAWSLLM
self._get_boto_credentials_from_optional_params(optional_params, model)
# Strip routing prefixes to get the actual model ID
clean_model_id = self._get_model_id(model)
# Use Moonshot's transform_request which handles message transformation
# and tool_choice="required" workaround
return MoonshotChatConfig.transform_request(
self,
model=clean_model_id,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
def _extract_reasoning_from_content(self, content: str) -> tuple[Optional[str], str]:
"""
Extract reasoning content from <reasoning> tags in the response.
Moonshot AI's Kimi K2 Thinking model returns reasoning in <reasoning> tags.
This method extracts that content and returns it separately.
Args:
content: The full content string from the API response
Returns:
tuple: (reasoning_content, main_content)
"""
if not content:
return None, content
# Match <reasoning>...</reasoning> tags
reasoning_match = re.match(
r"<reasoning>(.*?)</reasoning>\s*(.*)",
content,
re.DOTALL
)
if reasoning_match:
reasoning_content = reasoning_match.group(1).strip()
main_content = reasoning_match.group(2).strip()
return reasoning_content, main_content
return None, content
def transform_response(
self,
model: str,
raw_response: httpx.Response,
model_response: "ModelResponse",
logging_obj: LiteLLMLoggingObj,
request_data: dict,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> "ModelResponse":
"""
Transform the response from Bedrock Moonshot AI models.
Moonshot AI uses OpenAI-compatible response format, but returns reasoning
content in <reasoning> tags. This method:
1. Calls parent class transformation
2. Extracts reasoning content from <reasoning> tags
3. Sets reasoning_content on the message object
"""
# First, get the standard transformation
model_response = MoonshotChatConfig.transform_response(
self,
model=model,
raw_response=raw_response,
model_response=model_response,
logging_obj=logging_obj,
request_data=request_data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
api_key=api_key,
json_mode=json_mode,
)
# Extract reasoning content from <reasoning> tags
if model_response.choices and len(model_response.choices) > 0:
for choice in model_response.choices:
# Only process Choices (not StreamingChoices) which have message attribute
if isinstance(choice, Choices) and choice.message and choice.message.content:
reasoning_content, main_content = self._extract_reasoning_from_content(
choice.message.content
)
if reasoning_content:
# Set the reasoning_content field
choice.message.reasoning_content = reasoning_content
# Update the main content without reasoning tags
choice.message.content = main_content
return model_response
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BedrockError:
"""Return the appropriate error class for Bedrock."""
return BedrockError(status_code=status_code, message=error_message)
litellm/llms/bedrock/common_utils.py
+2
View File
@@ -629,6 +629,8 @@ def get_bedrock_chat_config(model: str):
return litellm.AmazonCohereConfig()
elif bedrock_invoke_provider == "mistral":
return litellm.AmazonMistralConfig()
elif bedrock_invoke_provider == "moonshot":
return litellm.AmazonMoonshotConfig()
elif bedrock_invoke_provider == "deepseek_r1":
return litellm.AmazonDeepSeekR1Config()
elif bedrock_invoke_provider == "nova":
tests/llm_translation/test_bedrock_moonshot.py
+290
View File
@@ -0,0 +1,290 @@
"""
Tests for Bedrock Moonshot (Kimi K2) integration.
This test suite verifies:
1. Basic completion functionality
2. Streaming responses
3. System message support
4. Temperature parameter handling
5. Reasoning content extraction from <reasoning> tags
6. Tool calling support (including tool response handling)
7. Parameter validation (e.g., stop sequences not supported)
"""
from base_llm_unit_tests import BaseLLMChatTest
import pytest
import sys
import os
import json
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm.llms.bedrock.common_utils import get_bedrock_chat_config
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
class TestBedrockMoonshotBasic:
"""Unit tests for Bedrock Moonshot configuration and transformations."""
def test_provider_detection_invoke(self):
"""Test that Bedrock Moonshot invoke models are correctly detected."""
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")
assert config is not None
def test_config_initialization(self):
"""Test that AmazonMoonshotConfig initializes correctly."""
config = get_bedrock_chat_config("invoke/moonshot.kimi-k2-thinking")
assert config is not None
assert config.custom_llm_provider == "bedrock"
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")
# Should support these params
assert "temperature" in supported_params
assert "max_tokens" in supported_params
assert "top_p" in supported_params
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
def test_transform_request_strips_model_prefix(self):
"""Test that model ID prefixes are correctly stripped in transform_request."""
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={}
)
# The model ID in the request body should be stripped
assert transformed["model"] == "moonshot.kimi-k2-thinking"
class TestBedrockMoonshotReasoningContent:
"""Tests for reasoning content extraction."""
def test_reasoning_content_extraction(self):
"""Test that reasoning content is extracted from <reasoning> tags."""
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)
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)
assert reasoning is None
assert content == "This is just a regular answer"
class TestBedrockMoonshotToolCalling:
"""Unit tests for tool calling functionality."""
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")
# Kimi K2 Thinking DOES support tool calls (unlike kimi-thinking-preview)
assert "tools" in supported_params
assert "tool_choice" in supported_params
def test_tool_call_request_format(self):
"""Test that tool call requests are formatted correctly."""
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?"}
]
optional_params = {
"tools": [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather",
"parameters": {
"type": "object",
"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={}
)
# 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
assert transformed["tools"][0]["function"]["name"] == "get_weather"
def test_tool_response_message_format(self):
"""Test that tool response messages are formatted correctly."""
# This tests the proper format for sending tool responses back
tool_response_message = {
"role": "tool",
"tool_call_id": "call_123",
"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
assert "content" in tool_response_message
class TestBedrockMoonshotParameterValidation:
"""Tests for parameter validation and edge cases."""
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")
# Bedrock Moonshot doesn't support stopSequences field
assert "stop" not in supported_params
def test_temperature_range(self):
"""Test that temperature parameter is handled correctly."""
# 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")
assert "temperature" in supported_params
class TestBedrockMoonshotTransformations:
"""Tests for request/response transformations."""
def test_transform_request_basic(self):
"""Test basic request transformation."""
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!"}
]
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={}
)
# 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
def test_transform_request_with_system_message(self):
"""Test request transformation with system message."""
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!"}
]
transformed = config.transform_request(
model="moonshot.kimi-k2-thinking",
messages=messages,
optional_params={},
litellm_params={},
headers={}
)
# System messages should be supported
assert "messages" in transformed