[Bug]: ReActAgent.astream_chat() throws AssertionError after updating to the lastest version

[dzmitrysyrtsou]

Bug Description

After updating to the latest version, ReActAgent.astream_chat() throws AssertionError: assert isinstance(chat_response, StreamingAgentChatResponse).
This seems to be related to a new assertion added to llama-index-core/llama_index/core/agent/runner/base.py, namely:

Isn't the logic now conflicting?

Looks like a bug....

Version

latest

Steps to Reproduce

Create a ReActAgent and call astream_chat() method.

Relevant Logs/Tracbacks

No response

[dosubot]

Hey @dzmitrysyrtsou! 👋 I'm here to help you with any bugs, questions, or contributions you have. Let's squash those bugs together!

The AssertionError in ReActAgent.astream_chat() is due to a new assertion added in the llama-index-core/llama_index/core/agent/runner/base.py file, which ensures that the chat_response is an instance of StreamingAgentChatResponse. This change was introduced in the commit titled "Add mypy to core (#14883)" ^[1][2].

To resolve this issue, ensure that the astream_chat method returns a StreamingAgentChatResponse and that the chat history is correctly maintained. The test case test_astream_chat_basic provides a good reference for the expected behavior:

@pytest.mark.asyncio()
async def test_astream_chat_basic(
    add_tool: FunctionTool,
) -> None:
    mock_llm = MockStreamChatLLM(
        responses=[
            ChatMessage(
                content=MOCK_ACTION_RESPONSE,
                role=MessageRole.ASSISTANT,
            ),
            ChatMessage(
                content=MOCK_STREAM_FINAL_RESPONSE,
                role=MessageRole.ASSISTANT,
            ),
        ]
    )

    agent = ReActAgent.from_tools(
        tools=[add_tool],
        llm=mock_llm,
    )
    response = await agent.astream_chat("What is 1 + 1?")
    assert isinstance(response, StreamingAgentChatResponse)

    # exhaust stream
    async for delta in response.async_response_gen():
        continue
    expected_answer = MOCK_STREAM_FINAL_RESPONSE.split("Answer: ")[-1].strip()
    assert response.response == expected_answer
    assert agent.chat_history == [
        ChatMessage(
            content="What is 1 + 1?",
            role=MessageRole.USER,
        ),
        ChatMessage(
            content="2 is the final answer.",
            role=MessageRole.ASSISTANT,
        ),
    ]

Ensure that your implementation aligns with this test case to avoid the AssertionError.

To continue talking to Dosu, mention @dosu.

[dzmitrysyrtsou]

Hey, I'm pretty sure we're using the Agent correctly. Here is a chunk of our code:

 async def get_assistant_response(self, request_text: str):
        agent = self._create_agent()
        chat_history = self._get_chat_history()
        llm_response = (
            await agent.astream_chat(
                message=request_text, chat_history=chat_history
            )
        )
        response_stream = llm_response.async_response_gen()

We are deeply convinced that the problem is not with the way the agent is used, but rather with the implementation of astream_chat() method.

[dosubot]

The issue you're encountering is likely due to the recent changes in the astream_chat method, specifically the new assertion that ensures chat_response is an instance of StreamingAgentChatResponse. This assertion was added in a recent commit titled "Add mypy to core (#14883)" ^[1][2].

Here is the exact implementation of the astream_chat method for reference:

@trace_method("chat")
async def astream_chat(
    self,
    message: str,
    chat_history: Optional[List[ChatMessage]] = None,
    tool_choice: Union[str, dict] = "auto",
) -> StreamingAgentChatResponse:
    with self.callback_manager.event(
        CBEventType.AGENT_STEP,
        payload={EventPayload.MESSAGES: [message]},
    ) as e:
        chat_response = await self._achat(
            message, chat_history, tool_choice, mode=ChatResponseMode.STREAM
        )
        assert isinstance(chat_response, StreamingAgentChatResponse)
        e.on_end(payload={EventPayload.RESPONSE: chat_response})
    return chat_response

The chat_response must be an instance of StreamingAgentChatResponse. If the _achat method returns a different type, this assertion will fail.

Ensure that the _achat method is correctly returning a StreamingAgentChatResponse. If _achat is returning a different type, you may need to adjust its implementation or the way it is called to ensure it returns the correct type.

If you need further assistance, please provide more details about the _achat method and how it is implemented. This will help in diagnosing the issue more accurately.

To continue talking to Dosu, mention @dosu.

[dzmitrysyrtsou]

We're not using a custom implementation of _achat() method. We're using a built-in implementation from Llama.

[logan-markewich]

Its definitely a bug lol