async def test_conversation_message_format():
    """Test that the conversation includes properly formatted assistant messages."""
    # Create mock config
    config = MagicMock()
    config.stream = True

    # Create mock providers
    llm_provider = MockLLMProvider(
        response_content="This is a test message",
        citations=[]
    )
    db_provider = MockDatabaseProvider()

    # Create mock search results collector
    search_results = {
        "abc1234": {
            "document_id": "doc_abc1234",
            "text": "This is document text for abc1234",
            "metadata": {"source": "source_abc1234"}
        },
        "def5678": {
            "document_id": "doc_def5678",
            "text": "This is document text for def5678",
            "metadata": {"source": "source_def5678"}
        }
    }
    search_results_collector = MockSearchResultsCollector(search_results)

    # Create agent
    agent = MockR2RStreamingAgent(
        database_provider=db_provider,
        llm_provider=llm_provider,
        config=config,
        rag_generation_config=GenerationConfig(model="test/model")
    )

    # Set the search results collector
    agent.search_results_collector = search_results_collector

    # Test a simple query
    messages = [Message(role="user", content="Test query")]

    # Run the agent
    stream = agent.arun(messages=messages)
    await collect_stream_output(stream)

    # Get the last message from the conversation
    last_message = agent.conversation.messages[-1]

    # Verify message format - note that MockR2RStreamingAgent uses a hardcoded response
    assert last_message.role == "assistant", "Last message should be from assistant"
    assert "This is a test response with citations" in last_message.content, "Message content should include response"
    assert "metadata" in last_message.dict(), "Message should include metadata"
    assert "citations" in last_message.metadata, "Message metadata should include citations"
Test that the conversation includes properly formatted assistant messages.
py/tests/unit/agent/test_agent.py