Function test_retrieval_tool

tests/main/test_doc_chat_agent.py:309–335

(
    test_settings: Settings, retrieval_agent, query: str, expected: str
)

    ],
)
def test_retrieval_tool(
    test_settings: Settings, retrieval_agent, query: str, expected: str
):
    set_global(test_settings)
    retrieval_agent.enable_message(RetrievalTool)
    task = Task(
        retrieval_agent,
        restart=True,
        interactive=False,
        system_message=f"""
        To answer user's query, use the `retrieval_tool` to retrieve relevant passages, 
        and ONLY then answer the query. 
        In case the query is simply a topic or search phrase, 
        guess what the user may want to know, and formulate it as a 
        question to be answered, and use this as the `query` field in the 
        `retrieval_tool`. 
                
        When you are ready to show your answer, say {DONE}, followed by the answer.
        """,
    )
    # 3 turns:
    # 1. LLM gen `retrieval_tool` request
    # 2. Agent gen `retrieval_tool` response (i.e. returns relevant passages)
    # 3. LLM gen answer based on passages
    ans = task.run(query, turns=3).content
    expected = [e.strip().lower() for e in expected.split(",")]
    assert all([e in ans.lower() for e in expected])


@pytest.fixture(scope="function")