How to distinguish AIMessageChunk used for AgentAction from AIMessageChunk used for AgentFinish when using .astream_events?

[mckinseyalexanderfalaleev]

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

query = "What was the average value for `column_name_1` in January 2024? Start with step by step explanation of your actions."


async for event in agent._get_agent_executor().astream_events({"input": query, "chat_history": []}, version="v1"):
    kind = event["event"]
    if kind == "on_chat_model_stream":
        print(event["data"]["chunk"].content, end="")

Description

Context:
I have the AgentExecutor with a PythonREPL-like tool connected to pandas.DataFrame named df and OpenAIToolsAgentOutputParser.

Question:
When using AgentExecutor in astream_events mode, I capture all tokens produced by LLMs, including those that will be used for AgentAction's log through filtering using on_chat_model_stream:

Input:

query = "What was the average value for `column_name_1` in January 2024? Start with step by step explanation of your actions."


async for event in agent_executor.astream_events({"input": query, "chat_history": []}, version="v1"):
    kind = event["event"]
    if kind == "on_chat_model_stream":
        print(event["data"]["chunk"].content, end="")

Output:

To calculate the average value for the feature `column_name_1` in January 2024, I will perform the following steps:

1. Filter the DataFrame `df` to include only the rows where the timestamp falls within January 2024.
2. Calculate the average of the `column_name_1` column for the filtered data.

Let me proceed with these steps and provide you with the result. The average value for `column_name_1` in January 2024 was approximately 123.0.

I only need the capture tokens that will get into AgentFinish later. For example, the "output" field of the agent's response is precisely what I need when invoking agents in non-streaming mode.

Input:

query = "What was the average value for `column_name_1` in January 2024? Start with step by step explanation of your actions."


response = agent_executor.invoke({"input": query, "chat_history": []})

response["output"]

Output:

'The average value for `column_name_1` in January 2024 was approximately 123.0.'

Is it possible to distinguish message chunks corresponding to AgentFinish from message chunks that will later become AgentAction's log while using astream_events method? These events look the same, and filtering rules can not be implemented.

System Info

System Information
------------------
> OS:  Darwin
> OS Version:  Darwin Kernel Version 23.5.0: Wed May  1 20:19:05 PDT 2024; root:xnu-10063.121.3~5/RELEASE_ARM64_T8112
> Python Version:  3.10.14 (main, Mar 21 2024, 11:21:31) [Clang 14.0.6 ]

Package Information
-------------------
> langchain_core: 0.1.37
> langchain: 0.1.14
> langchain_community: 0.0.30
> langsmith: 0.1.38
> langchain_cli: 0.0.21
> langchain_experimental: 0.0.56
> langchain_openai: 0.1.1
> langchain_text_splitters: 0.0.1
> langgraph: 0.0.30
> langserve: 0.0.51

[dosubot[bot]]

Hey @mckinseyalexanderfalaleev! 👋 I'm here to help you with bug fixes, questions, and more. Let's tackle this together!

Yes, it is possible to distinguish message chunks corresponding to AgentFinish from message chunks that will later become AgentAction's log while using the astream_events method. The AgentFinish chunks will have a structure indicating the final output, while AgentAction chunks will contain details about the actions being taken.

For example, in the astream method, the AgentFinish chunk is identified by the presence of the final output:

{
    "output": "curses foiled again",
    "messages": [
        AIMessage(content="Oh well\nFinal Answer: curses foiled again")
    ],
}

In contrast, AgentAction chunks contain details about the tool being invoked and the input provided to the tool:

{
    "actions": [
        AgentAction(
            tool="Search",
            tool_input="misalignment",
            log="FooBarBaz\nAction: Search\nAction Input: misalignment",
        )
    ],
    "messages": [
        AIMessage(
            content="FooBarBaz\nAction: Search\nAction Input: misalignment"
        )
    ],
}

This distinction allows you to differentiate between the final output of the agent and the intermediate actions it takes.

Additionally, the AgentFinish class has a type attribute set to "AgentFinish" and includes return_values and a log for the final output, while the AgentAction class has a type attribute set to "AgentAction" and includes tool, tool_input, and log for the action details. This structural difference in the schema definitions further aids in distinguishing between the two types of message chunks.

In the context of using the astream_events method, you can identify the type of event by checking the event key in the event dictionary. For example:

async for event in agent_executor.astream_events(
    {"messages": [HumanMessage(content="whats the weather in sf?")]}, version="v1"
):
    kind = event["event"]
    if kind == "on_chain_start":
        if event["name"] == "Agent":
            print(f"Starting agent: {event['name']} with input: {event['data'].get('input')}")
    elif kind == "on_chain_end":
        if event["name"] == "Agent":
            print()
            print("--")
            print(f"Done agent: {event['name']} with output: {event['data'].get('output')['output']}")
    if kind == "on_chat_model_stream":
        content = event["data"]["chunk"].content
        if content:
            print(content, end="|")
    elif kind == "on_tool_start":
        print("--")
        print(f"Starting tool: {event['name']} with inputs: {event['data'].get('input')}")
    elif kind == "on_tool_end":
        print(f"Done tool: {event['name']}")
        print(f"Tool output was: {event['data'].get('output')}")
        print("--")

This code snippet demonstrates how to handle different types of events, including those corresponding to AgentFinish and AgentAction ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[mckinseyalexanderfalaleev]

@dosu

This does not solve the issue. I need to use the astream_events method to access intermediate steps, including tool invocation events.

The proposed option for astream_events would not work as I need to distinguish events, where kind equals to on_chat_model_stream in both cases.

[dosubot[bot]]

To distinguish between AIMessageChunk used for AgentAction and AIMessageChunk used for AgentFinish when using the astream_events method, you can leverage the structure and content of the events. Specifically, you can look for markers in the on_chat_model_stream events that indicate whether the chunk is part of an AgentAction or AgentFinish.

Here is an approach to differentiate between the two:

1. Check for Tool Invocation: AgentAction chunks often contain details about the tool being invoked. You can look for specific keywords or patterns in the content that indicate a tool invocation.

2. Final Output Indicators: AgentFinish chunks will typically contain the final output of the agent. You can look for phrases or patterns that indicate a conclusion or final answer.

Here is an example of how you might implement this:

async for event in agent_executor.astream_events(
    {"messages": [HumanMessage(content="What was the average value for `column_name_1` in January 2024? Start with step by step explanation of your actions.")]}, version="v1"
):
    kind = event["event"]
    if kind == "on_chat_model_stream":
        content = event["data"]["chunk"].content
        if "Action:" in content or "Action Input:" in content:
            # This is likely part of an AgentAction
            print(f"AgentAction chunk: {content}")
        elif "Final Answer:" in content or "The average value" in content:
            # This is likely part of an AgentFinish
            print(f"AgentFinish chunk: {content}")
        else:
            # General content, could be part of either
            print(f"General chunk: {content}")

In this example, the presence of keywords like "Action:" or "Action Input:" helps identify AgentAction chunks, while phrases like "Final Answer:" or specific output patterns help identify AgentFinish chunks.

This approach allows you to filter and distinguish between the different types of message chunks while using the astream_events method ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[rodrigop23]

any updates @mckinseyalexanderfalaleev ?

[mckinseyalexanderfalaleev]

any updates @mckinseyalexanderfalaleev ?

I wasn't able to solve this :(