Add tests for kimi 2 bedrock model

2026-08-05 20:23:18 +00:00 · 2026-01-08 10:07:33 +05:30
parent ddd52e0e2f
commit af6883712e
3 changed files with 469 additions and 2 deletions
@@ -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
@@ -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) |
@@ -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