unable to execute tool in langgraph

[cris-m]

I was following langgraph example. I use Cohere model and bind TavilySearchResults but I can not execute it.

1. Create Agent Executor

from langchain import hub
from langchain.agents import create_openai_functions_agent
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.prompts import MessagesPlaceholder

tools = [TavilySearchResults(max_results=1)]
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a helpful assistant."),
        ("human", "{input}"),
         MessagesPlaceholder("chat_history"),
    ]
)
llm = ChatCohere(model="command-r-plus", temperature=0).bind_tools(tools)
agent_runnable = prompt | llm

2. Define the state

import operator
from typing import Annotated, TypedDict, Union
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.messages import BaseMessage

class AgentState(TypedDict):
    input: str
    chat_history: list[BaseMessage]
    agent_outcome: Union[AgentAction, AgentFinish, None]
    intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]

3. Nodes

from langchain_core.agents import AgentFinish
from langgraph.prebuilt import ToolNode

tool_executor = ToolNode(tools)

def run_agent(data):
    agent_outcome = agent_runnable.invoke(data)
    return {"agent_outcome": agent_outcome}

def execute_tools(data):
    agent_action = data["agent_outcome"]
    output = tool_executor.invoke(agent_action)
    return {"intermediate_steps": [(agent_action, str(output))]}

# Define logic that will be used to determine which conditional edge to go down
def should_continue(data):
    if isinstance(data["agent_outcome"], AgentFinish):
        return "end"
    else:
        return "continue"

4. Workflow

from langgraph.graph import END, StateGraph, START

workflow = StateGraph(AgentState)
workflow.add_node("agent", run_agent)
workflow.add_node("action", execute_tools)
workflow.add_edge(START, "agent")

workflow.add_conditional_edges(
    "agent",
    should_continue,
    {
        "continue": "action",
        "end": END,
    },
)

workflow.add_edge("action", "agent")
app = workflow.compile()

5. Execution

inputs = {
    "input": "what is the weather in sf", 
    "chat_history": []
}

for s in app.stream(inputs):
    print(list(s.values())[0])
    print("----")

The following error is thrown:

[/usr/local/lib/python3.10/dist-packages/langgraph/prebuilt/tool_node.py](https://58fpipkfh0o-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240830-060131_RC00_669227538#) in _parse_input(self, input)
    141             output_type = "list"
    142             message: AnyMessage = input[-1]
--> 143         elif messages := input.get("messages", []):
    144             output_type = "dict"
    145             message = messages[-1]

AttributeError: 'AIMessage' object has no attribute 'get'

The first stream is printed and it come AIMessage instead of AgentActionMessageLog

{'agent_outcome': AIMessage(content='I will search for the weather in SF and write an answer based on what I find.', additional_kwargs={'documents': None, 'citations': None, 'search_results': None, 'search_queries': None, 'is_search_required': None, 'generation_id': 'fa8183a2-423b-4542-848b-f7554971210e', 'tool_calls': [{'id': '526b85fb4f394960a3757997495a22e8', 'function': {'name': 'tavily_search_results_json', 'arguments': '{"query": "weather in sf"}'}, 'type': 'function'}], 'token_count': {'input_tokens': 935.0, 'output_tokens': 68.0}}, response_metadata={'documents': None, 'citations': None, 'search_results': None, 'search_queries': None, 'is_search_required': None, 'generation_id': 'fa8183a2-423b-4542-848b-f7554971210e', 'tool_calls': [{'id': '526b85fb4f394960a3757997495a22e8', 'function': {'name': 'tavily_search_results_json', 'arguments': '{"query": "weather in sf"}'}, 'type': 'function'}], 'token_count': {'input_tokens': 935.0, 'output_tokens': 68.0}}, id='run-fcdc461f-6180-4ed2-b7b9-7efda6f916ca-0', tool_calls=[{'name': 'tavily_search_results_json', 'args': {'query': 'weather in sf'}, 'id': 'a4cb20bfdd5544baa372207af59453c8', 'type': 'tool_call'}], usage_metadata={'input_tokens': 935, 'output_tokens': 68, 'total_tokens': 1003})}