feat: Support Qwen3 models on ChatBedrock (+ gpt-oss streaming) (#679)

michaelnchin · web-flow · commit aa296a820839 · 2025-10-21T21:08:24.000-07:00
Adding support for using the new [Qwen3](https://www.aboutamazon.com/news/aws/alibaba-qwen3-deepseek-v3-amazon-bedrock) serverless models with the InvokeModel APIs via ChatBedrock. This PR also implements ChatBedrock streaming support for both Qwen3 and OpenAI GPT-OSS models.
diff --git a/libs/aws/langchain_aws/chat_models/bedrock.py b/libs/aws/langchain_aws/chat_models/bedrock.py
@@ -279,7 +279,7 @@ def convert_messages_to_prompt_writer(messages: List[BaseMessage]) -> str:
     """Convert a list of messages to a prompt for Writer."""
 
     return "\n".join(
-        [_convert_one_message_to_text_llama(message) for message in messages]
+        [_convert_one_message_to_text_writer(message) for message in messages]
     )
 
 
@@ -741,7 +741,7 @@ def format_messages(
     ]:
         if provider == "anthropic":
             return _format_anthropic_messages(messages)
-        elif provider == "openai":
+        elif provider in ("openai", "qwen"):
             return cast(List[Dict[str, Any]], convert_to_openai_messages(messages))
         raise NotImplementedError(
             f"Provider {provider} not supported for format_messages"
@@ -914,7 +914,7 @@ def _stream(
                     system = self.system_prompt_with_tools
             else:
                 system = system_str
-        elif provider == "openai":
+        elif provider in ("openai", "qwen"):
             formatted_messages = cast(
                 List[Dict[str, Any]],
                 ChatPromptAdapter.format_messages(provider, messages),
@@ -1060,7 +1060,7 @@ def _generate(
                 else:
                     system = system_str
                 citations_enabled = _citations_enabled(formatted_messages)
-            elif provider == "openai":
+            elif provider in ("openai", "qwen"):
                 formatted_messages = cast(
                     List[Dict[str, Any]],
                     ChatPromptAdapter.format_messages(provider, messages),
diff --git a/libs/aws/langchain_aws/llms/bedrock.py b/libs/aws/langchain_aws/llms/bedrock.py
@@ -185,12 +185,16 @@ def _stream_response_to_generation_chunk(
         if k
         not in [output_key, "prompt_token_count", "generation_token_count", "created"]
     }
+
+    if provider in ["mistral", "deepseek", "writer"]:
+        text = stream_response[output_key][0]["text"]
+    elif provider in ["openai", "qwen"]:
+        text = stream_response[output_key][0]["delta"].get("content", "")
+    else:
+        text = stream_response[output_key]
+
     return GenerationChunk(
-        text=(
-            stream_response[output_key]
-            if provider not in ["mistral", "deepseek", "writer"]
-            else stream_response[output_key][0]["text"]
-        ),
+        text=text,
         generation_info=generation_info,
     )
 
@@ -297,6 +301,8 @@ class LLMInputOutputAdapter:
         "deepseek": "choices",
         "meta": "generation",
         "mistral": "outputs",
+        "openai": "choices",
+        "qwen": "choices",
         "writer": "choices",
     }
 
@@ -402,14 +408,19 @@ def prepare_input(
                     input_body["max_tokens"] = max_tokens
                 elif provider == "writer":
                     input_body["max_tokens"] = max_tokens
-                elif provider == "openai":
-                    input_body["max_output_tokens"] = max_tokens
                 else:
                     # TODO: Add AI21 support, param depends on specific model.
                     pass
             if temperature is not None:
                 input_body["temperature"] = temperature
 
+        elif provider in ("openai", "qwen"):
+            input_body["messages"] = messages
+            if max_tokens:
+                input_body["max_tokens"] = max_tokens
+            if temperature is not None:
+                input_body["temperature"] = temperature
+
         elif provider == "amazon":
             input_body = dict()
             input_body["inputText"] = prompt
@@ -419,12 +430,6 @@ def prepare_input(
             if temperature is not None:
                 input_body["textGenerationConfig"]["temperature"] = temperature
 
-        elif provider == "openai":
-            input_body["messages"] = messages
-            if max_tokens:
-                input_body["max_tokens"] = max_tokens
-            if temperature is not None:
-                input_body["temperature"] = temperature
         else:
             input_body["inputText"] = prompt
 
@@ -478,6 +483,8 @@ def prepare_output(cls, provider: str, response: Any) -> dict:
                 text = response_body.get("outputs")[0].get("text")
             elif provider == "openai":
                 text = response_body.get("choices")[0].get("message").get("content")
+            elif provider == "qwen":
+                text = response_body.get("choices")[0].get("message").get("content")
             else:
                 text = response_body.get("results")[0].get("outputText")
 
@@ -576,6 +583,14 @@ def prepare_output_stream(
                 yield _get_invocation_metrics_chunk(chunk_obj)
                 return
 
+            elif (
+                provider in ("qwen", "openai")
+                and chunk_obj.get(output_key, [{}])[0].get("finish_reason", "")
+                == "stop"
+            ):
+                yield _get_invocation_metrics_chunk(chunk_obj)
+                return
+
             elif messages_api and (chunk_obj.get("type") == "message_stop"):
                 yield _get_invocation_metrics_chunk(chunk_obj)
                 return
@@ -619,6 +634,14 @@ async def aprepare_output_stream(
             ):
                 return
 
+            elif (
+                provider in ("qwen", "openai")
+                and chunk_obj.get(output_key, [{}])[0].get("finish_reason", "")
+                == "stop"
+            ):
+                yield _get_invocation_metrics_chunk(chunk_obj)
+                return
+
             generation_chunk = _stream_response_to_generation_chunk(
                 chunk_obj,
                 provider=provider,
@@ -1219,7 +1242,7 @@ def _prepare_input_and_invoke_stream(
             provider,
             response,
             stop,
-            True if messages else False,
+            True if (messages and provider == "anthropic") else False,
             coerce_content_to_string=coerce_content_to_string,
         ):
             yield chunk
@@ -1288,7 +1311,7 @@ async def _aprepare_input_and_invoke_stream(
             provider,
             response,
             stop,
-            True if messages else False,
+            True if (messages and provider == "anthropic") else False,
         ):
             yield chunk
 
diff --git a/libs/aws/tests/unit_tests/llms/test_bedrock.py b/libs/aws/tests/unit_tests/llms/test_bedrock.py
@@ -304,6 +304,63 @@ def test__human_assistant_format() -> None:
     {"chunk": {"bytes": b'"[DONE]"'}},
 ]
 
+MOCK_STREAMING_RESPONSE_QWEN = [
+    {
+        "chunk": {
+            "bytes": b'{"choices": [{"delta": {"content": "", "role": "assistant"}, '
+            b'"finish_reason": null, "index": 0}], '
+            b'"created": 1759875373, '
+            b'"id": "chatcmpl-a069cbda08ce4599afae798c4d2de095", '
+            b'"model": "qwen.qwen3-32b-v1:0", '
+            b'"object": "chat.completion.chunk", '
+            b'"service_tier": "auto"}'
+        }
+    },
+    {
+        "chunk": {
+            "bytes": b'{"choices": [{"delta": {"content": "Hello.  \\nGoodbye."}, '
+            b'"finish_reason": "stop", "index": 0}], '
+            b'"created": 1759875373, '
+            b'"id": "chatcmpl-a069cbda08ce4599afae798c4d2de095", '
+            b'"model": "qwen.qwen3-32b-v1:0", '
+            b'"object": "chat.completion.chunk", '
+            b'"service_tier": "auto", '
+            b'"amazon-bedrock-invocationMetrics": {'
+            b'"inputTokenCount": 35, "outputTokenCount": 7, '
+            b'"invocationLatency": 225, "firstByteLatency": 191}}'
+        }
+    },
+]
+
+MOCK_STREAMING_RESPONSE_OPENAI = [
+    {
+        "chunk": {
+            "bytes": b'{"choices": [{"delta": {"content": "Hello."}, '
+            b'"finish_reason": null, "index": 0}], '
+            b'"created": 1759813667, '
+            b'"id": "chatcmpl-fa6fb768b71046eeb3880cbb4a1b07c1", '
+            b'"model": "openai.gpt-oss-20b-1:0", '
+            b'"object": "chat.completion.chunk", "service_tier": "auto"}'
+        }
+    },
+    {
+        "chunk": {
+            "bytes": b'{"choices": [{"delta": {}, '
+            b'"finish_reason": "stop", "index": 0}],'
+            b' "created": 1759813667, '
+            b'"id": "chatcmpl-fa6fb768b71046eeb3880cbb4a1b07c1", '
+            b'"model": "openai.gpt-oss-20b-1:0", '
+            b'"object": "chat.completion.chunk", '
+            b'"service_tier": "auto", '
+            b'"amazon-bedrock-invocationMetrics": {'
+            b'"inputTokenCount": 84, '
+            b'"outputTokenCount": 87, '
+            b'"invocationLatency": 3981, '
+            b'"firstByteLatency": 3615}}'
+        }
+    },
+]
+
 
 async def async_gen_mock_streaming_response() -> AsyncGenerator[Dict, None]:
     for item in MOCK_STREAMING_RESPONSE:
@@ -421,6 +478,56 @@ def writer_streaming_response():
     return response
 
 
+@pytest.fixture
+def qwen_response():
+    body = MagicMock()
+    body.read.return_value = json.dumps(
+        {"choices": [{"message": {"content": "This is the Qwen output text."}}]}
+    ).encode()
+    response = dict(
+        body=body,
+        ResponseMetadata={
+            "HTTPHeaders": {
+                "x-amzn-bedrock-input-token-count": "35",
+                "x-amzn-bedrock-output-token-count": "42",
+            }
+        },
+    )
+
+    return response
+
+
+@pytest.fixture
+def qwen_streaming_response():
+    response = dict(body=MOCK_STREAMING_RESPONSE_QWEN)
+    return response
+
+
+@pytest.fixture
+def openai_response():
+    body = MagicMock()
+    body.read.return_value = json.dumps(
+        {"choices": [{"message": {"content": "This is the OpenAI output text."}}]}
+    ).encode()
+    response = dict(
+        body=body,
+        ResponseMetadata={
+            "HTTPHeaders": {
+                "x-amzn-bedrock-input-token-count": "85",
+                "x-amzn-bedrock-output-token-count": "80",
+            }
+        },
+    )
+
+    return response
+
+
+@pytest.fixture
+def openai_streaming_response():
+    response = dict(body=MOCK_STREAMING_RESPONSE_OPENAI)
+    return response
+
+
 @pytest.fixture
 def cohere_response():
     body = MagicMock()
@@ -556,6 +663,48 @@ def test_prepare_output_stream_for_writer(writer_streaming_response) -> None:
     assert results[1] == "lo."
 
 
+def test_prepare_output_for_qwen(qwen_response):
+    result = LLMInputOutputAdapter.prepare_output("qwen", qwen_response)
+    assert result["text"] == "This is the Qwen output text."
+    assert result["usage"]["prompt_tokens"] == 35
+    assert result["usage"]["completion_tokens"] == 42
+    assert result["usage"]["total_tokens"] == 77
+    assert result["stop_reason"] is None
+
+
+def test_prepare_output_stream_for_qwen(qwen_streaming_response) -> None:
+    results = [
+        chunk.text
+        for chunk in LLMInputOutputAdapter.prepare_output_stream(
+            "qwen", qwen_streaming_response
+        )
+    ]
+
+    assert results[0] == ""
+    assert results[1] == "Hello.  \nGoodbye."
+
+
+def test_prepare_output_for_openai(openai_response):
+    result = LLMInputOutputAdapter.prepare_output("openai", openai_response)
+    assert result["text"] == "This is the OpenAI output text."
+    assert result["usage"]["prompt_tokens"] == 85
+    assert result["usage"]["completion_tokens"] == 80
+    assert result["usage"]["total_tokens"] == 165
+    assert result["stop_reason"] is None
+
+
+def test_prepare_output_stream_for_openai(openai_streaming_response) -> None:
+    results = [
+        chunk.text
+        for chunk in LLMInputOutputAdapter.prepare_output_stream(
+            "openai", openai_streaming_response
+        )
+    ]
+
+    assert results[0] == "Hello."
+    assert results[1] == ""
+
+
 def test_prepare_output_for_cohere(cohere_response):
     result = LLMInputOutputAdapter.prepare_output("cohere", cohere_response)
     assert result["text"] == "This is the Cohere output text."
