Agent Executor astream_events with asyncGenerator tools.

[Lin-jun-xiang]

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

async def chat_completion(
    client,
    message: str | List,
    history: List[List[str]],
    client_name: str = 'OpenAI',
) -> AsyncGenerator:
    """Use OpenAI/Anthropic API to get `streaming` response."""
    try:
        formatted_history = [{"role": "system", "content": SYS_PROMPT}]
        for user, assistant in history:
            # history: [[user_query, assistant_response], [user_query, assistant_response], ...]
            formatted_history.append({"role": "user", "content": user })
            formatted_history.append({"role": "assistant", "content": assistant})
        if isinstance(message, str):
            formatted_history.append({"role": "user", "content": message})
        elif isinstance(message, list):
            formatted_history.extend(message)

        if client_name == 'OpenAI':
            stream = await client.chat.completions.create(
                model="gpt-4o-2024-08-06",
                messages=formatted_history,
                stream=True,
                temperature=0.,
            )
            async for chunk in stream:
                yield chunk.choices[0].delta.content or ""

@tool
async def get_image_tool() -> str:
    """Get image"""
    global snapshot
    with Image.open('myimage.png') as image:
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        snapshot = base64.b64encode(buffered.getvalue()).decode('utf-8')
    return "Get Image Successfully"


@tool
async def describe_image() -> AsyncGenerator:
    """Describe the image"""
    global snapshot_session
    message=[
        {
            "role": "user",
            "content": [
                {"type": "text", "text": 'describe the image'},
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{snapshot}"}}
            ]
        }
    ]
    response = chat_completion(openai_client, message, [])
    async for chunk in response:
        yield chunk


tools = [
    get_image_tool,
    describe_image,
]
openai_client = AsyncOpenAI(api_key=OPENAI_API_KEY)
langchain_llm_client = ChatOpenAI(
    model='gpt-4o',
    temperature=0.,
    api_key=OPENAI_API_KEY,
    streaming=True,
    max_tokens=None,
    verbose=VERBOSE,
    max_retries=5
)
agent = create_tool_calling_agent(langchain_llm_client, tools, AGENT_PROMPT)
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=VERBOSE,
    return_intermediate_steps=True
)

async main():
    tool_names = [tool.name for tool in tools]

    async for event in agent_executor.astream_events(
        {
            "input": formatted_history_query,
            "tool_names": tool_names,
            "agent_scratchpad": lambda x: format_to_openai_tool_messages(x["intermediate_steps"]),
        },
        version='v2'
    ):
        kind = event['event']
        if kind == "on_chain_start":
            pass
        elif kind == "on_chat_model_stream":
            content = event["data"]["chunk"].content
            if content:
                yield content
        elif kind == "on_tool_start":
            pass
        elif kind == "on_tool_end":
            if isinstance(event['data'].get('output'), AsyncGenerator):
                async for event_chunk in event['data'].get('output'):
                    yield event_chunk
            else:
                yield(
                    f"{event['data'].get('output')}\n"
                )
        elif kind == "on_chain_end":
            pass

Description

I attempted to use the langchain agent executor with astream_events. When I gave the following instruction:

"Please take a photo, describe it, and tell me if there are any white clouds in the photo?"

I usually receive a response like the following (output in streaming mode):

"There is a big sun in the photo, ... Get Image Successfully..., I cannot determine the contents of the photo..., there are white clouds and an airplane..."

I have the following questions:

1. From the response, it seems that astream_events tends to execute each tool (e.g., get_image_tool, describe_image) concurrently due to its async nature.
2. Since describe_image is an asyncGenerator, the agent begins summarizing before the photo description is complete, leading to an incorrect understanding, such as "I cannot determine the contents of the photo."

I would like to know if there is a way to restrict astream_event to wait for specific tools (asyncGenerators) to finish their tasks before proceeding to the next step or summarizing.

System Info

System Information

OS: Windows
OS Version: 10.0.22631
Python Version: 3.11.8 (tags/v3.11.8:db85d51, Feb 6 2024, 22:03:32) [MSC v.1937 64 bit (AMD64)]

Package Information

langchain_core: 0.2.30
langchain: 0.2.11
langchain_community: 0.2.10
langsmith: 0.1.86
langchain_anthropic: 0.1.20
langchain_openai: 0.1.17
langchain_text_splitters: 0.2.2
langgraph: 0.2.3

Optional packages not installed

langserve

Other Dependencies

aiohttp: 3.10.5
anthropic: 0.31.2
async-timeout: 4.0.3
dataclasses-json: 0.6.7
defusedxml: 0.7.1
jsonpatch: 1.33
langgraph-checkpoint: 1.0.2
numpy: 1.26.4
openai: 1.40.3
orjson: 3.10.6
packaging: 24.1
pydantic: 2.8.2
PyYAML: 6.0.1
requests: 2.32.3
SQLAlchemy: 2.0.31
tenacity: 8.5.0
tiktoken: 0.7.0
typing-extensions: 4.12.2

[dosubot[bot]]

I found some similar discussions and issues that might be relevant to your question:

1. how could agent step by step with streaming? ^[1]
2. Some tool function should be wait when using astream_event ^[2]
3. Agent Executor astream events sometimes gets blocked. ^[3]

To address your specific question about restricting astream_events to wait for specific tools (asyncGenerators) to finish their tasks before proceeding, you can use asyncio.gather to ensure that all asynchronous tasks are awaited and completed before moving to the next step. Here is a relevant code snippet:

# If the tool chosen is the finishing tool, then we end and return.
if isinstance(output, AgentFinish):
    yield output
    return

actions: List[AgentAction]
if isinstance(output, AgentAction):
    actions = [output]
else:
    actions = output
for agent_action in actions:
    yield agent_action

# Use asyncio.gather to run multiple tool.arun() calls concurrently
result = await asyncio.gather(
    *[
        self._aperform_agent_action(
            name_to_tool_map, color_mapping, agent_action, run_manager
        )
        for agent_action in actions
    ],
)

# TODO This could yield each result as it becomes available
for chunk in result:
    yield chunk

This code ensures that all the asynchronous tasks (tool.arun()) are awaited and completed before proceeding to the next step ^[4][5].

Additionally, ensure that the RunnableConfig object is properly propagated to child runnables, especially if you are running Python versions <= 3.10. This is crucial for emitting events from custom runnables or tools. For example:

from langchain_core.runnables import RunnableConfig

@tool
async def special_summarization_tool_with_config(
    long_text: str, config: RunnableConfig
) -> str:
    """A tool that summarizes input text using advanced techniques."""
    prompt = ChatPromptTemplate.from_template(
        "You are an expert writer. Summarize the following text in 10 words or less:\n\n{long_text}"
    )

    def reverse(x: str):
        return x[::-1]

    chain = prompt | model | StrOutputParser() | reverse
    # Pass the "config" object as an argument to any executed runnables
    summary = await chain.ainvoke({"long_text": long_text}, config=config)
    return summary

By ensuring the proper propagation of the RunnableConfig, you can effectively manage the execution and streaming of events from your tools ^[6].

To continue talking to Dosu, mention @dosu.