"""Tests for the survey agent's TaskGroup workflow. Demonstrates the best practices documented in https://docs.livekit.io/agents/logic/tasks/#testing-task-groups: - **Initialize ``userdata``** — the ``session`` fixture passes ``userdata=_userdata()`` to every ``AgentSession``. The survey's tasks read ``session.userdata.candidate_name`` and write into ``task_results``; in Python, accessing an unset ``userdata`` raises ``ValueError``. - **Sleep before the first ``session.run()``** — the fixture sleeps 0.5s after ``sess.start()`` (and the full-flow test sleeps again between TaskGroup sub-tasks). ``AgentTask`` transitions briefly leave ``session.llm`` unset, and ``session.run()`` does not fall back to it during that window. - **Drive multiple turns** — the LLM often replies conversationally before invoking a completion tool. ``_drive_until_called`` sends an initial input, then keeps nudging until every expected tool name appears in ``sess.history.items``. This is the spirit of the docs' guidance to prefer ``contains_function_call()`` over ``next_event()``: don't couple the test to a specific event index in a single ``RunResult``. - **Parse ``item.arguments`` with ``json.loads``** — see ``test_commute_task_records_can_commute`` and ``test_experience_task_records_years_and_description``. The other tests assert on ``userdata`` or membership in the tool-call set instead, where argument parsing isn't needed. - **Don't assert on startup output** — none of the tests inspect the ``IntroTask`` greeting produced in ``on_enter``; ``RunResult`` does not capture output generated before the first ``session.run()``. - **Tasks tested in isolation and as a group** — four isolation tests wrap a single ``AgentTask`` in ``_SingleTaskAgent``; ``test_full_task_group_flow`` exercises the full ``TaskGroup`` ordering, ``on_task_completed`` callback, and ``task_results`` keying. Test-only adjustment (not a documented practice, but worth flagging): ``_SurveyAgentForTesting`` sets ``summarize_chat_ctx=False`` so the final ``TaskGroup`` summarization LLM call doesn't run during tests. Run with:: uv add --dev aiocsv aiofiles uv run pytest examples/survey/test_survey_agent.py """ from __future__ import annotations import asyncio import json from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import AsyncExitStack import pytest from livekit.agents import Agent, AgentSession, AgentTask, inference, llm from livekit.agents.beta.workflows import TaskGroup from livekit.agents.llm import FunctionCall from .agent import ( BehavioralTask, CommuteTask, ExperienceTask, IntroTask, Userdata, ) # AgentTask transitions briefly clear `session.llm`; sleep before the first # `sess.run()` and between TaskGroup sub-tasks so the new sub-task can take # over. See https://docs.livekit.io/agents/logic/tasks/#testing-task-groups. _TASK_TRANSITION_DELAY = 0.5 def _llm_model() -> llm.LLM: return inference.LLM( model="openai/gpt-4.1-mini", extra_kwargs={"parallel_tool_calls": False, "temperature": 0.2}, ) def _userdata() -> Userdata: # Initialize userdata: tasks read session.userdata.candidate_name and # write task_results. In Python, accessing an unset userdata raises # ValueError("AgentSession userdata is not set"). return Userdata(filename="results-test.csv", candidate_name="", task_results={}) # --------------------------------------------------------------------------- # Fixtures and helpers # --------------------------------------------------------------------------- SessionStarter = Callable[[Agent], Awaitable[AgentSession]] @pytest.fixture async def session() -> AsyncIterator[SessionStarter]: """Yield a function that starts an `AgentSession` for a given agent. Handles LLM lifecycle, the `_TASK_TRANSITION_DELAY` after `start()`, and a non-draining shutdown on teardown so any in-flight LLM activity is cancelled instead of allowed to drain (which would let stdout output continue past the assertions). """ async with AsyncExitStack() as stack: sessions: list[AgentSession] = [] async def _start(agent: Agent) -> AgentSession: model = await stack.enter_async_context(_llm_model()) sess = AgentSession(llm=model, userdata=_userdata()) await sess.start(agent) sessions.append(sess) await asyncio.sleep(_TASK_TRANSITION_DELAY) return sess try: yield _start finally: for sess in sessions: sess.shutdown(drain=False) await asyncio.wait_for(sess.aclose(), timeout=10) async def _run(sess: AgentSession, user_input: str, *, timeout: float = 30): return await asyncio.wait_for(sess.run(user_input=user_input), timeout=timeout) def _called_tools(sess: AgentSession) -> set[str]: return {item.name for item in sess.history.items if item.type == "function_call"} def _last_calls(sess: AgentSession, names: set[str]) -> dict[str, FunctionCall]: """Most recent function-call item in `sess.history` for each requested name.""" found: dict[str, FunctionCall] = {} for item in reversed(sess.history.items): if item.type == "function_call" and item.name in names and item.name not in found: found[item.name] = item if found.keys() == names: break return found async def _drive_until_called( sess: AgentSession, *, expected: str | set[str], initial: str, nudge: str = "Yes, that's right. Please go ahead and record it.", max_turns: int = 4, ) -> dict[str, FunctionCall]: """Drive a task to completion across multiple turns and return the matching calls. The LLM may respond conversationally before calling a completion tool. We send `initial`, then keep nudging until every tool name in `expected` appears in the session's function-call history (or we run out of turns). Returns a dict mapping each expected name to the most recent matching call. """ required = {expected} if isinstance(expected, str) else set(expected) await _run(sess, initial) for _ in range(max_turns - 1): if required.issubset(_called_tools(sess)): break await _run(sess, nudge) assert required.issubset(_called_tools(sess)), ( f"expected tools {required} not called; got {_called_tools(sess)}" ) return _last_calls(sess, required) class _SingleTaskAgent(Agent): """Thin wrapper that runs a single AgentTask and exits. Used for isolation tests so each task can be exercised without the full TaskGroup orchestration. """ def __init__(self, task: AgentTask) -> None: super().__init__(instructions="Run a single survey task.") self._task = task async def on_enter(self) -> None: await self._task # --------------------------------------------------------------------------- # Isolated task tests # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_intro_task_in_isolation(session: SessionStarter) -> None: """IntroTask records the candidate's name via record_intro.""" sess = await session(_SingleTaskAgent(IntroTask())) await _drive_until_called( sess, expected="record_intro", initial=( "Hi, my name is Alice. I'm a backend engineer with five years of " "experience building APIs at Acme." ), ) # The IntroTask writes the candidate name into userdata. assert sess.userdata.candidate_name.lower() == "alice" @pytest.mark.asyncio async def test_commute_task_records_can_commute(session: SessionStarter) -> None: """CommuteTask records can_commute=True with the chosen method.""" sess = await session(_SingleTaskAgent(CommuteTask())) calls = await _drive_until_called( sess, expected="record_commute_flexibility", initial="Yes, I can commute three days a week. I usually take the subway.", ) # Function call arguments are stored as raw JSON — parse before asserting. args = json.loads(calls["record_commute_flexibility"].arguments) assert args["can_commute"] is True assert args["commute_method"] == "subway" @pytest.mark.asyncio async def test_experience_task_records_years_and_description( session: SessionStarter, ) -> None: """ExperienceTask captures years_of_experience and a description.""" sess = await session(_SingleTaskAgent(ExperienceTask())) calls = await _drive_until_called( sess, expected="record_experience", initial=( "I have five years of experience total. I started as a junior engineer " "at Acme working on data pipelines for two years, then moved to Globex " "as a senior backend engineer for the past three years." ), ) args = json.loads(calls["record_experience"].arguments) assert args["years_of_experience"] == 5 assert "acme" in args["experience_description"].lower() @pytest.mark.asyncio async def test_behavioral_task_completes_after_three_records( session: SessionStarter, ) -> None: """BehavioralTask only completes once strengths, weaknesses, and work style have all been recorded — typically requires multiple turns. """ sess = await session(_SingleTaskAgent(BehavioralTask())) await _drive_until_called( sess, expected={"record_strengths", "record_weaknesses", "record_work_style"}, initial=( "My biggest strength is debugging hard distributed systems issues. " "My main weakness is that I sometimes over-engineer early prototypes. " "I work best as part of a team — that is my work style." ), max_turns=6, ) # --------------------------------------------------------------------------- # Full TaskGroup flow # --------------------------------------------------------------------------- class _SurveyAgentForTesting(Agent): """Mirrors the production ``SurveyAgent`` but disables chat-context summarization and skips the email step. - ``summarize_chat_ctx=False``: summarization issues an additional LLM call at the end of the group; disabling it keeps the test focused on task orchestration and avoids depending on summary quality. - We omit ``GetEmailTask`` and the CSV write to keep the test offline-safe and free of filesystem side effects. The orchestration semantics being verified — sequential ordering, ``task_results`` keying, and ``on_task_completed`` callbacks — are identical to production. `done` is set after `task_results` is assigned, so the test can wait on a single signal instead of polling. """ def __init__(self, completed_ids: list[str], done: asyncio.Event) -> None: super().__init__(instructions="You are a survey agent screening candidates.") self._completed_ids = completed_ids self._done = done async def on_enter(self) -> None: async def _on_task_completed(event): # type: ignore[no-untyped-def] self._completed_ids.append(event.task_id) group = TaskGroup( summarize_chat_ctx=False, on_task_completed=_on_task_completed, ) group.add(lambda: IntroTask(), id="intro", description="Collect name and intro.") group.add(lambda: CommuteTask(), id="commute", description="Ask about commute.") group.add(lambda: ExperienceTask(), id="experience", description="Collect work history.") result = await group self.session.userdata.task_results = result.task_results self._done.set() @pytest.mark.asyncio async def test_full_task_group_flow(session: SessionStarter) -> None: """The TaskGroup runs IntroTask → CommuteTask → ExperienceTask in order, populates ``task_results`` keyed by task id, and fires ``on_task_completed`` exactly once per task. """ completed_ids: list[str] = [] done = asyncio.Event() sess = await session(_SurveyAgentForTesting(completed_ids=completed_ids, done=done)) # Don't assert on startup output — the IntroTask greeting produced in # on_enter is not captured in the RunResult below. # Drive each task to completion before moving on. Each call loops through # additional turns until the expected tool fires — the LLM often replies # conversationally before invoking a completion tool. Sleep between # sub-tasks for the same reason as the initial transition delay. await _drive_until_called( sess, expected="record_intro", initial=( "My name is Bob, I'm a software engineer with eight years of experience " "focused on APIs." ), ) await asyncio.sleep(_TASK_TRANSITION_DELAY) await _drive_until_called( sess, expected="record_commute_flexibility", initial="Yes, I can commute three days a week. I'd be driving in.", ) await asyncio.sleep(_TASK_TRANSITION_DELAY) await _drive_until_called( sess, expected="record_experience", initial=( "I have eight years total — five at Initech on backend systems and the " "last three at Hooli leading an API team." ), ) # Wait for `_SurveyAgentForTesting.on_enter` to finish: TaskGroup # finalization, callback delivery, and `task_results` assignment. await asyncio.wait_for(done.wait(), timeout=10) results = sess.userdata.task_results assert completed_ids == ["intro", "commute", "experience"], ( f"tasks completed out of order: {completed_ids}" ) assert set(results.keys()) == {"intro", "commute", "experience"} assert results["intro"].name.lower() == "bob" assert results["commute"].can_commute is True assert results["commute"].commute_method == "driving" assert results["experience"].years_of_experience == 8