ChatCohere with Agentic RAG and Langgraph

[jimmyshah83]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

def agent(state):
    """
    Invokes the agent model to generate a response based on the current state. Given
    the question, it will decide to retrieve using the retriever tool, or simply end.

    Args:
        state (messages): The current state

    Returns:
        dict: The updated state with the agent response appended to messages
    """
    print("---CALL AGENT---")
    messages = state["messages"]
    model = AzureChatOpenAI(model="gpt-4o", temperature=0, streaming=True)
    # model = ChatCohere(model="command-r-plus-08-2024", temperature=0, streaming=True)
    model = model.bind_tools(tools)
    response = model.invoke(messages)
    # We return a list, because this will get added to the existing list
    return {"messages": [response]}

Description

I am building an Agentic RAG system with ChatCohere but when the agent tries to Transform a query, it keeps failing with the below error:
The last message is not an ToolMessage or HumanMessage

Looking at the class ChatCohere, it seems like it might be a design choice:

# check if the last message is a tool message or human message
    if not (
        isinstance(messages[-1], ToolMessage) or isinstance(messages[-1], HumanMessage)
    ):
        raise ValueError("The last message is not an ToolMessage or HumanMessage")

Does anyone know why is this the case? In terms of Agentic AI, the agent should have the ability to rewrite the question to generate a more optimal output but this code base does not allow rewrites.

Note that it works perfectly fine with AzureChatOpenAI

System Info

python-dotenv
pandas
langchain
langchainhub
langchain-cohere
azure-search-documents==11.6.0b7
azure-identity
langgraph
langchain_openai
cohere