feat: capture agent-turn input and spoken output

livekit-agents/livekit/agents/voice/agent_activity.py

@@ -286,6 +286,9 @@ def _agent_turn(
         speech_handle._agent_turn_context = trace.set_span_in_context(span)
         speech_handle._agent_turn_started_at = time.perf_counter()
         speech_handle._agent_turn_agent_name = agent_label
+        speech_handle._agent_turn_record_content = (
+            span.is_recording() and gen_ai_telemetry.capture_content_enabled()
+        )
 
     generation_attrs: dict[str, Any] = {
         trace_types.ATTR_AGENT_TURN_ID: speech_handle._generation_id
@@ -3321,6 +3324,7 @@ def _on_first_frame(fut: asyncio.Future[float] | asyncio.Future[None]) -> None:
             else:
                 forwarded_text = ""
         current_span.set_attribute(trace_types.ATTR_RESPONSE_TEXT, forwarded_text)
+        speech_handle._record_agent_turn_output(forwarded_text)
 
         assistant_metrics: llm.MetricsReport = {}
 
@@ -3446,6 +3450,13 @@ async def _pipeline_reply_task_impl(
             current_span.set_attribute(
                 trace_types.ATTR_USER_INPUT, new_message.raw_text_content or ""
             )
+            if speech_handle._agent_turn_record_content:
+                gen_ai_telemetry.set_content_attributes(
+                    current_span,
+                    input_messages=gen_ai_telemetry.to_input_messages(
+                        llm.ChatContext([new_message])
+                    ),
+                )
 
         if (room_io := self._session._room_io) and room_io.room.isconnected():
             _set_participant_attributes(current_span, room_io.room.local_participant)
@@ -3840,6 +3851,7 @@ async def _next_segment() -> _SpeechSegment | None:
             )
 
         forwarded_text = "".join(out.forwarded_text for out in segment_outputs)
+        speech_handle._record_agent_turn_output(forwarded_text)
         if speech_handle.interrupted:
             # forward_generation already cleared the buffer and waited for playout
             await utils.aio.cancel_and_wait(*tasks)
@@ -4157,6 +4169,7 @@ async def _realtime_reply_task(
                 generation_ev=generation_ev,
                 model_settings=model_settings,
                 instructions=instructions,
+                user_input=user_input,
             )
         finally:
             # reset tool_choice and tools
@@ -4180,12 +4193,20 @@ async def _realtime_generation_task(
         generation_ev: llm.GenerationCreatedEvent,
         model_settings: ModelSettings,
         instructions: str | None = None,
+        user_input: str | None = None,
     ) -> None:
         with _agent_turn(
             speech_handle,
             root_context=self._session._root_span_context,
             agent_label=self._agent.label,
-        ):
+        ) as turn_span:
+            if user_input is not None and speech_handle._agent_turn_record_content:
+                gen_ai_telemetry.set_content_attributes(
+                    turn_span,
+                    input_messages=gen_ai_telemetry.to_input_messages(
+                        llm.ChatContext([llm.ChatMessage(role="user", content=[user_input])])
+                    ),
+                )
             inference_span = tracer.start_span("realtime_inference")
             try:
                 await self._realtime_generation_task_impl(
@@ -4555,6 +4576,7 @@ def _create_assistant_message(
                 continue
 
             trace_text_parts.append(forwarded_text)
+            speech_handle._record_agent_turn_output(forwarded_text)
             chat_msg = _create_assistant_message(
                 message_id=entry.msg.message_id,
                 forwarded_text=forwarded_text,

livekit-agents/livekit/agents/voice/speech_handle.py

@@ -11,7 +11,7 @@
 
 from .. import llm, utils
 from ..log import logger
-from ..telemetry import trace_types
+from ..telemetry import gen_ai, trace_types
 
 INTERRUPTION_TIMEOUT = 5.0  # seconds
 
@@ -64,6 +64,8 @@ def __init__(
         self._agent_turn_context: otel_context.Context | None = None
         self._agent_turn_started_at: float | None = None
         self._agent_turn_agent_name: str | None = None
+        self._agent_turn_record_content = False
+        self._agent_turn_output: list[str] = []
         self._scheduled_at: float | None = None
         self._authorized_at: float | None = None
         self._interrupt_source: InterruptionSource | None = None  # first interrupt's cause
@@ -395,6 +397,15 @@ def _continue_agent_turn(
         self._agent_turn_context = trace.set_span_in_context(span)
         self._agent_turn_started_at = started_at
         self._agent_turn_agent_name = agent_name
+        self._agent_turn_record_content = discarded._agent_turn_record_content
+        self._agent_turn_output = discarded._agent_turn_output
+        discarded._agent_turn_output = []
+
+    def _record_agent_turn_output(self, text: str) -> None:
+        if not gen_ai.capture_content_enabled():
+            self._agent_turn_record_content = False
+        if self._agent_turn_record_content and text:
+            self._agent_turn_output.append(text)
 
     def _end_agent_turn(self, error: BaseException | None) -> None:
         """Close the speech's ``agent_turn`` span: the speech is done, whatever step it was on."""
@@ -413,6 +424,16 @@ def _end_agent_turn(self, error: BaseException | None) -> None:
             return
         if isinstance(error, Exception):
             trace_utils.record_exception(span, error)
+        if self._agent_turn_record_content and gen_ai.capture_content_enabled():
+            gen_ai.set_content_attributes(
+                span,
+                output_messages=[
+                    message
+                    for text in self._agent_turn_output
+                    for message in gen_ai.to_output_messages(text=text)
+                ],
+            )
+        self._agent_turn_output.clear()
         span.end()
 
     def _mark_scheduled(self) -> None:

tests/test_agent_turn_span.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import asyncio
+import json
 from collections.abc import Iterator
 
 import pytest
@@ -16,10 +17,11 @@
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
-from livekit.agents import Agent, RunContext, function_tool
+from livekit.agents import Agent, AgentSession, RunContext, function_tool, llm, utils
 from livekit.agents.llm import FunctionToolCall
-from livekit.agents.telemetry import set_tracer_provider, trace_types, tracer
+from livekit.agents.telemetry import gen_ai, set_tracer_provider, trace_types, tracer
 
+from .fake_realtime import FakeRealtimeModel, fake_capabilities
 from .fake_session import FakeActions, create_session, run_session
 from .trace_schema import assert_trace_well_formed
 
@@ -63,15 +65,20 @@ async def get_weather(self, context: RunContext, location: str) -> str:
         return f"sunny in {location}"
 
 
-async def test_tool_call_is_one_agent_turn(span_exporter: InMemorySpanExporter) -> None:
+@pytest.mark.parametrize("preface", ["", "Let me check."])
+async def test_tool_call_is_one_agent_turn(
+    span_exporter: InMemorySpanExporter, preface: str
+) -> None:
     actions = FakeActions()
     actions.add_user_speech(0.5, 2.0, "What's the weather in Tokyo?")
     actions.add_llm(
-        content="",
+        content=preface,
         tool_calls=[
             FunctionToolCall(name="get_weather", arguments='{"location": "Tokyo"}', call_id="1")
         ],
     )
+    if preface:
+        actions.add_tts(0.5)
     actions.add_llm(content="It is sunny in Tokyo.", input="sunny in Tokyo")
     actions.add_tts(1.0)
 
@@ -85,6 +92,13 @@ async def test_tool_call_is_one_agent_turn(span_exporter: InMemorySpanExporter)
     assert turn.parent is not None and turn.parent.span_id == root.context.span_id
 
     attrs = turn.attributes or {}
+    assert json.loads(attrs["gen_ai.input.messages"]) == [
+        {"role": "user", "parts": [{"type": "text", "content": "What's the weather in Tokyo?"}]}
+    ]
+    assert json.loads(attrs["gen_ai.output.messages"]) == [
+        {"role": "assistant", "parts": [{"type": "text", "content": text}]}
+        for text in ([preface] if preface else []) + ["It is sunny in Tokyo."]
+    ]
     speech_id = attrs[trace_types.ATTR_SPEECH_ID]
     assert attrs[trace_types.ATTR_GENERATION_COUNT] == 2
     assert attrs[trace_types.ATTR_AGENT_TURN_ID] == f"{speech_id}_2"
@@ -101,8 +115,8 @@ async def test_tool_call_is_one_agent_turn(span_exporter: InMemorySpanExporter)
     # both generations' inference, the tool between them, and the speech all nest under it
     assert len(_children(span_exporter, turn, "llm_node")) == 2
     [tool] = _children(span_exporter, turn, "function_tool")
-    [tts] = _children(span_exporter, turn, "tts_node")
-    [speaking] = _children(span_exporter, turn, "agent_speaking")
+    tts = _children(span_exporter, turn, "tts_node")[-1]
+    speaking = _children(span_exporter, turn, "agent_speaking")[-1]
     assert tool.start_time < tts.start_time
     # and the turn covers everything, ending with the speech rather than with the first step
     for child in (tool, tts, speaking):
@@ -228,3 +242,76 @@ def test_sampled_out_turn_is_still_handed_to_the_successor() -> None:
     assert reply._agent_turn_span is not None
     assert reply._agent_turn_started_at == 1.0
     assert reply._agent_turn_agent_name == "a"
+
+
+@pytest.mark.virtual_time
+@pytest.mark.parametrize("capture", [True, False])
+@pytest.mark.parametrize("add_to_chat_ctx", [True, False])
+async def test_say_greeting_records_output_without_llm(
+    span_exporter: InMemorySpanExporter,
+    monkeypatch: pytest.MonkeyPatch,
+    capture: bool,
+    add_to_chat_ctx: bool,
+) -> None:
+    monkeypatch.setattr(gen_ai, "_capture_content", capture)
+    actions = FakeActions()
+    actions.add_tts(0.5, input="Hello there.")
+
+    class GreetingAgent(Agent):
+        async def on_enter(self) -> None:
+            await self.session.say("Hello there.", add_to_chat_ctx=add_to_chat_ctx)
+
+    session = create_session(actions, with_stt=False)
+    await asyncio.wait_for(run_session(session, GreetingAgent(instructions="test")), 60)
+    [turn] = _spans(span_exporter, "agent_turn")
+    assert not _spans(span_exporter, "llm_request")
+    assert "gen_ai.input.messages" not in turn.attributes
+    if capture:
+        assert json.loads(turn.attributes["gen_ai.output.messages"]) == [
+            {"role": "assistant", "parts": [{"type": "text", "content": "Hello there."}]}
+        ]
+    else:
+        assert "gen_ai.output.messages" not in turn.attributes
+
+
+async def test_realtime_turn_records_explicit_input_and_forwarded_output(
+    span_exporter: InMemorySpanExporter,
+) -> None:
+    model = FakeRealtimeModel(capabilities=fake_capabilities(audio_output=False))
+    async with AgentSession(llm=model) as session:
+        await session.start(Agent(instructions="test"))
+        handle = session.generate_reply(user_input="Hi")
+        while not model.active_session._reply_futs:
+            await asyncio.sleep(0)
+        message_ch = utils.aio.Chan[llm.MessageGeneration]()
+        function_ch = utils.aio.Chan[llm.FunctionCall]()
+        text_ch = utils.aio.Chan[str]()
+        audio_ch = utils.aio.Chan()
+        modalities = asyncio.Future()
+        modalities.set_result(["text"])
+        message_ch.send_nowait(
+            llm.MessageGeneration(
+                message_id="reply",
+                text_stream=text_ch,
+                audio_stream=audio_ch,
+                modalities=modalities,
+            )
+        )
+        text_ch.send_nowait("Hello")
+        text_ch.close()
+        audio_ch.close()
+        message_ch.close()
+        function_ch.close()
+        model.active_session._reply_futs[0].set_result(
+            llm.GenerationCreatedEvent(
+                message_stream=message_ch, function_stream=function_ch, user_initiated=True
+            )
+        )
+        await asyncio.wait_for(handle, 5)
+    [turn] = _spans(span_exporter, "agent_turn")
+    assert json.loads(turn.attributes["gen_ai.input.messages"]) == [
+        {"role": "user", "parts": [{"type": "text", "content": "Hi"}]}
+    ]
+    assert json.loads(turn.attributes["gen_ai.output.messages"]) == [
+        {"role": "assistant", "parts": [{"type": "text", "content": "Hello"}]}
+    ]

tests/test_coverage_spans.py

@@ -5,6 +5,7 @@
 
 import asyncio
 import contextlib
+import json
 import time
 from collections.abc import Iterator
 from typing import Any
@@ -87,6 +88,9 @@ async def test_barge_in_records_source_and_playout_position(
     assert len(interrupted) == 1
     turn = interrupted[0]
     assert (turn.attributes or {})[trace_types.ATTR_INTERRUPTION_SOURCE] == "audio_activity"
+    output = json.loads(turn.attributes["gen_ai.output.messages"])
+    assert output[0]["parts"][0]["content"] == turn.attributes[trace_types.ATTR_RESPONSE_TEXT]
+    assert output[0]["parts"][0]["content"] != "Here is a long story for you ... the end."
     position = (turn.attributes or {})[trace_types.ATTR_PLAYOUT_POSITION]
     assert isinstance(position, float)
     # ~2 s of the 10 s story had played (5.5 - 3.5), scaled by the speed factor