import chainlit as cl

# Import your agent and streaming function.
from app import agent  # your CodeAgent instance
from Gradio_UI import stream_to_gradio  # re-use the streaming generator


@cl.on_message
async def main(message: cl.Message):
    """
    This function is called every time a user inputs a message in the UI.
    It sends back an intermediate response from the tool, followed by the final answer.

    Args:
        message: The user's message.

    Returns:
        None.
    """

    # Stream responses from the agent.
    # (The generator yields gr.ChatMessage objects; we extract role and content.)
    for msg in stream_to_gradio(agent, message.content, reset_agent_memory=False):
        role = msg.role
        # For non-string content (like images/audio) add additional handling as needed.
        content = msg.content if isinstance(msg.content, str) else str(msg.content)
        await cl.Message(content=content, author=role).send()