Suport for Cohere Command-R and Command-R+

[ksaegusa]

I was wondering if it would be possible to add support for the Cohere Command-R model in ChatBedrock. I think this feature would really enhance the platform and be helpful for users.

We would appreciate your consideration.

[3coins]

@ksaegusa
I would suggest to start implementing this inside ChatBedrockConverse class, as converse API seems to support most capabilities for Cohere R and R+ as noted in this documentation. https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-supported-models-features.html

Some features might already work with the current implementation, I would encourage you to test and report back what is working and not working here, that can help us discover the gaps on fixing the code to support Cohere models fully. For example, these currently work with ChatBedrockConverse.

Comand R

from langchain_aws import ChatBedrockConverse
llm = ChatBedrockConverse(model="cohere.command-r-v1:0")
llm.invoke("What is the capital of France?")

# Output
AIMessage(content="The current capital of France is Paris. It has been the capital since 1871. Paris is also the most populous city in France, with around 2.1 million residents. The city is a global centre for art, fashion, gastronomy, and culture, and it's a major economic hub.\n\nBefore Paris, the capital of France changed frequently, often based on the ruling monarch. Cities such as Versailles, Turin, Bordeaux, and Rheims served as capitals at various points in history.", additional_kwargs={}, response_metadata={'ResponseMetadata': {'RequestId': 'd3fc3a11-4577-4483-8534-000a57888ec7', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Fri, 11 Oct 2024 19:25:17 GMT', 'content-type': 'application/json', 'content-length': '641', 'connection': 'keep-alive', 'x-amzn-requestid': 'd3fc3a11-4577-4483-8534-000a57888ec7'}, 'RetryAttempts': 0}, 'stopReason': 'end_turn', 'metrics': {'latencyMs': 1222}}, id='run-b4778814-435d-45a3-9a70-4cee7a44ad9b-0', usage_metadata={'input_tokens': 7, 'output_tokens': 104, 'total_tokens': 111})

Command R+

llm = ChatBedrockConverse(model="cohere.command-r-plus-v1:0")
llm.invoke("What is the capital of France?")

# output
AIMessage(content='The capital of France is Paris.', additional_kwargs={}, response_metadata={'ResponseMetadata': {'RequestId': '34dda7c3-2294-4980-a309-711d6d223d76', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Fri, 11 Oct 2024 19:26:14 GMT', 'content-type': 'application/json', 'content-length': '210', 'connection': 'keep-alive', 'x-amzn-requestid': '34dda7c3-2294-4980-a309-711d6d223d76'}, 'RetryAttempts': 0}, 'stopReason': 'end_turn', 'metrics': {'latencyMs': 310}}, id='run-612b9cf6-98bd-455d-b0b8-a65d83610a46-0', usage_metadata={'input_tokens': 7, 'output_tokens': 7, 'total_tokens': 14})

Streaming with chain

from langchain_core.output_parsers import StrOutputParser
chain = llm | StrOutputParser()
for chunk in chain.stream("What is the capital of France?"):
     print(chunk)

# output
The
 capital
 of
 France
 is
 Paris
.

[ksaegusa]

@3coins Thank you for your comment! I can confirm that the Command-R+ model works for invoke and stream in　ChatBedrockConverse.

However, when trying to use ChatHistory with Command-R+, it seems necessary to specify a HumanMessage at the beginning, otherwise an error occurs. (This issue did not occur with the Anthropic model, so I believe it is a problem specific to the Command-R model.)

botocore.errorfactory.ValidationException: An error occurred (ValidationException) when calling the Converse operation: A conversation must start with a user message. Try again with a conversation that starts with a user message.

NG:

chat_template = ChatPromptTemplate.from_messages(
    [
        SystemMessage(
            content=(
                "あなたは日本語が得意なボブです"
                "質問に答えてください"
            )
        ),
        AIMessage(
            content=(
                "ボブです。どうぞよろしくお願いします。何かお手伝いできることはありま"
            )
        ),
        HumanMessagePromptTemplate.from_template("{text}"),
    ]
)

SUCCESS:

chat_template = ChatPromptTemplate.from_messages(
    [
        HumanMessage(
            content=(
                "あなたは日本語が得意なボブです"
                "質問に答えてください"
            )
        ),
        AIMessage(
            content=(
                "ボブです。どうぞよろしくお願いします。何かお手伝いできることはありま"
            )
        ),
        HumanMessagePromptTemplate.from_template("{text}"),
    ]
)

[3coins]

@ksaegusa It seems like the Human message is expected order before an AI message in requests to Bedrock. I tested your example with Anthropic haiku model, and received the same error. It also seems like you are trying to manage the chat history in the prompt template rather than using the RunnableWithHistoryMessage object. I would recommend looking at this guide to setup the chat history, you can adapt the code by adding the system and human message only in the prompt template. Can you try that and let me know if that helps. https://python.langchain.com/v0.2/docs/how_to/message_history/

Note: If you are planning to use LangGraph in the near future, see the v0.3 guide which recommends using the LangGraph memory for the history. https://python.langchain.com/docs/how_to/message_history/

[ksaegusa]

@3coins Thank you! It turns out the issue was with the AI message. After trying RunnableWithHistoryMessage, the problem was resolved.

Also, thank you for the information about LangGraph. I am considering using it in the future and will check it out!

[3coins]

Closing, as Cohere models are supported in v0.2.3.