Add Writer model support to ChatBedrock (#478)

Fixes #472 

Added support for using [Writer
Palmyra](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-writer-palmyra.html)
models with Invoke API via ChatBedrock.

As a note, Writer Palmyra X4/X5 model prompt formatting is not publicly
documented - Llama/Mistral format seems to work without issues for now.

Author: michaelnchin

libs/aws/langchain_aws/chat_models/bedrock.py

@@ -226,6 +226,28 @@ def convert_messages_to_prompt_deepseek(messages: List[BaseMessage]) -> str:
     return prompt
 
 
+def _convert_one_message_to_text_writer(message: BaseMessage) -> str:
+    if isinstance(message, ChatMessage):
+        message_text = f"\n\n{message.role.capitalize()}: {message.content}"
+    elif isinstance(message, HumanMessage):
+        message_text = f"[INST] {message.content} [/INST]"
+    elif isinstance(message, AIMessage):
+        message_text = f"{message.content}"
+    elif isinstance(message, SystemMessage):
+        message_text = f"<<SYS>> {message.content} <</SYS>>"
+    else:
+        raise ValueError(f"Got unknown type {message}")
+    return message_text
+
+
+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]
+    )
+
+
 def _format_image(image_url: str) -> Dict:
     """
     Formats an image of format data:image/jpeg;base64,{b64_string}
@@ -553,6 +575,8 @@ def convert_messages_to_prompt(
                 human_prompt="\n\nUser:",
                 ai_prompt="\n\nBot:",
             )
+        elif provider == "writer":
+            prompt = convert_messages_to_prompt_writer(messages=messages)
         else:
             raise NotImplementedError(
                 f"Provider {provider} model does not support chat."

libs/aws/langchain_aws/llms/bedrock.py

@@ -176,7 +176,7 @@ def _stream_response_to_generation_chunk(
     return GenerationChunk(
         text=(
             stream_response[output_key]
-            if provider not in ["mistral", "deepseek"]
+            if provider not in ["mistral", "deepseek", "writer"]
             else stream_response[output_key][0]["text"]
         ),
         generation_info=generation_info,
@@ -273,6 +273,7 @@ class LLMInputOutputAdapter:
         "deepseek": "choices",
         "meta": "generation",
         "mistral": "outputs",
+        "writer": "choices"
     }
 
     @classmethod
@@ -363,7 +364,7 @@ def prepare_input(
             if temperature is not None:
                 input_body["temperature"] = temperature
 
-        elif provider in ("ai21", "cohere", "meta", "mistral", "deepseek"):
+        elif provider in ("ai21", "cohere", "meta", "mistral", "deepseek", "writer"):
             input_body["prompt"] = prompt
             if max_tokens:
                 if provider == "cohere":
@@ -374,6 +375,8 @@ def prepare_input(
                     input_body["max_tokens"] = max_tokens
                 elif provider == "deepseek":
                     input_body["max_tokens"] = max_tokens
+                elif provider == "writer":
+                    input_body["max_tokens"] = max_tokens
                 else:
                     # TODO: Add AI21 support, param depends on specific model.
                     pass
@@ -429,16 +432,16 @@ def prepare_output(cls, provider: str, response: Any) -> dict:
                     tool_calls = extract_tool_calls(content)
 
         else:
-            if provider == "ai21":
+            if provider in ["deepseek", "writer"]:
+                text = response_body.get("choices")[0].get("text")
+            elif provider == "ai21":
                 text = response_body.get("completions")[0].get("data").get("text")
             elif provider == "cohere":
                 text = response_body.get("generations")[0].get("text")
             elif provider == "meta":
                 text = response_body.get("generation")
             elif provider == "mistral":
                 text = response_body.get("outputs")[0].get("text")
-            elif provider == "deepseek":
-                text = response_body.get("choices")[0].get("text")
             else:
                 text = response_body.get("results")[0].get("outputText")
 
@@ -493,7 +496,10 @@ def prepare_output_stream(
 
             chunk_obj = json.loads(chunk.get("bytes").decode())
 
-            if provider == "cohere" and (
+            if provider == "writer" and chunk_obj == "[DONE]":
+                return
+
+            elif provider == "cohere" and (
                 chunk_obj["is_finished"] or chunk_obj[output_key] == "<EOS_TOKEN>"
             ):
                 return

libs/aws/tests/unit_tests/llms/test_bedrock.py

@@ -280,6 +280,21 @@ def test__human_assistant_format() -> None:
     {"chunk": {"bytes": b'{"choices": [{"text": "you.","stop_reason": "stop"}]}'}},
 ]
 
+MOCK_STREAMING_RESPONSE_WRITER = [
+    {"chunk": {'bytes': b'{"id":"cmpl-ec61121fa19443caa7f614bde08e926c",'
+              b'"object":"text_completion",'
+              b'"created":1747106231,'
+              b'"model":"writer.palmyra-x5-v1:0",'
+              b'"choices":[{"index":0,"text":"Hel","logprobs":null,"finish_reason":null,"stop_reason":null}],'
+              b'"usage":null}'}},
+    {"chunk": {'bytes': b'{"id":"cmpl-ec61121fa19443caa7f614bde08e926c",'
+              b'"object":"text_completion",'
+              b'"created":1747106231,'
+              b'"model":"writer.palmyra-x5-v1:0",'
+              b'"choices":[{"index":0,"text":"lo.","logprobs":null,"finish_reason":"length","stop_reason":null}],'
+              b'"usage":null}'}},
+    {"chunk": {'bytes': b'"[DONE]"'}},
+]
 
 async def async_gen_mock_streaming_response() -> AsyncGenerator[Dict, None]:
     for item in MOCK_STREAMING_RESPONSE:
@@ -372,6 +387,31 @@ def deepseek_streaming_response():
     return response
 
 
+@pytest.fixture
+def writer_response():
+    body = MagicMock()
+    body.read.return_value = json.dumps(
+        {'choices': [{'text': ' This is the Writer output text.'}]}
+    ).encode()
+    response = dict(
+        body=body,
+        ResponseMetadata={
+            "HTTPHeaders": {
+                "x-amzn-bedrock-input-token-count": "17",
+                "x-amzn-bedrock-output-token-count": "8",
+            }
+        },
+    )
+
+    return response
+
+
+@pytest.fixture
+def writer_streaming_response():
+    response = dict(body=MOCK_STREAMING_RESPONSE_WRITER)
+    return response
+
+
 @pytest.fixture
 def cohere_response():
     body = MagicMock()
@@ -486,6 +526,27 @@ def test_prepare_output_stream_for_deepseek(deepseek_streaming_response) -> None
     assert results[1] == "you."
 
 
+def test_prepare_output_for_writer(writer_response):
+    result = LLMInputOutputAdapter.prepare_output("writer", writer_response)
+    assert result["text"] == " This is the Writer output text."
+    assert result["usage"]["prompt_tokens"] == 17
+    assert result["usage"]["completion_tokens"] == 8
+    assert result["usage"]["total_tokens"] == 25
+    assert result["stop_reason"] is None
+
+
+def test_prepare_output_stream_for_writer(writer_streaming_response) -> None:
+    results = [
+        chunk.text
+        for chunk in LLMInputOutputAdapter.prepare_output_stream(
+            "writer", writer_streaming_response
+        )
+    ]
+
+    assert results[0] == "Hel"
+    assert results[1] == "lo."
+
+
 def test_prepare_output_for_cohere(cohere_response):
     result = LLMInputOutputAdapter.prepare_output("cohere", cohere_response)
     assert result["text"] == "This is the Cohere output text."