@livekit/agents 1.1.0 → 1.2.0

package/src/voice/generation.ts

@@ -3,8 +3,9 @@
 // SPDX-License-Identifier: Apache-2.0
 import type { AudioFrame } from '@livekit/rtc-node';
 import { AudioResampler } from '@livekit/rtc-node';
+import type { Span } from '@opentelemetry/api';
+import { context as otelContext } from '@opentelemetry/api';
 import type { ReadableStream, ReadableStreamDefaultReader } from 'stream/web';
-import { ZodObject } from 'zod';
 import {
   type ChatContext,
   ChatMessage,
@@ -19,20 +20,38 @@ import {
   isFunctionTool,
   isToolError,
 } from '../llm/tool_context.js';
+import { isZodSchema, parseZodSchema } from '../llm/zod-utils.js';
 import { log } from '../log.js';
 import { IdentityTransform } from '../stream/identity_transform.js';
-import { Future, Task, shortuuid, toError } from '../utils.js';
-import { type Agent, type ModelSettings, asyncLocalStorage, isStopResponse } from './agent.js';
+import { traceTypes, tracer } from '../telemetry/index.js';
+import { USERDATA_TIMED_TRANSCRIPT } from '../types.js';
+import { Future, Task, shortuuid, toError, waitForAbort } from '../utils.js';
+import {
+  type Agent,
+  type ModelSettings,
+  _setActivityTaskInfo,
+  functionCallStorage,
+  isStopResponse,
+} from './agent.js';
 import type { AgentSession } from './agent_session.js';
-import type { AudioOutput, LLMNode, TTSNode, TextOutput } from './io.js';
+import {
+  AudioOutput,
+  type LLMNode,
+  type TTSNode,
+  type TextOutput,
+  type TimedString,
+  createTimedString,
+  isTimedString,
+} from './io.js';
 import { RunContext } from './run_context.js';
 import type { SpeechHandle } from './speech_handle.js';
 
-/* @internal */
+/** @internal */
 export class _LLMGenerationData {
   generatedText: string = '';
   generatedToolCalls: FunctionCall[];
   id: string;
+  ttft?: number;
 
   constructor(
     public readonly textStream: ReadableStream<string>,
@@ -43,6 +62,21 @@ export class _LLMGenerationData {
   }
 }
 
+/**
+ * TTS generation data containing audio stream and optional timed transcripts.
+ * @internal
+ */
+export interface _TTSGenerationData {
+  /** Audio frame stream from TTS */
+  audioStream: ReadableStream<AudioFrame>;
+  /**
+   * Future that resolves to a stream of timed transcripts, or null if TTS doesn't support it.
+   */
+  timedTextsFut: Future<ReadableStream<TimedString> | null>;
+  /** Time to first byte (set when first audio frame is received) */
+  ttfb?: number;
+}
+
 // TODO(brian): remove this class in favor of ToolOutput
 export class _ToolOutput {
   output: _JsOutput[];
@@ -383,6 +417,8 @@ export function performLLMInference(
   toolCtx: ToolContext,
   modelSettings: ModelSettings,
   controller: AbortController,
+  model?: string,
+  provider?: string,
 ): [Task<void>, _LLMGenerationData] {
   const textStream = new IdentityTransform<string>();
   const toolCallStream = new IdentityTransform<FunctionCall>();
@@ -391,9 +427,24 @@ export function performLLMInference(
   const toolCallWriter = toolCallStream.writable.getWriter();
   const data = new _LLMGenerationData(textStream.readable, toolCallStream.readable);
 
-  const inferenceTask = async (signal: AbortSignal) => {
+  const _performLLMInferenceImpl = async (signal: AbortSignal, span: Span) => {
+    span.setAttribute(
+      traceTypes.ATTR_CHAT_CTX,
+      JSON.stringify(chatCtx.toJSON({ excludeTimestamp: false })),
+    );
+    span.setAttribute(traceTypes.ATTR_FUNCTION_TOOLS, JSON.stringify(Object.keys(toolCtx)));
+
+    if (model) {
+      span.setAttribute(traceTypes.ATTR_GEN_AI_REQUEST_MODEL, model);
+    }
+    if (provider) {
+      span.setAttribute(traceTypes.ATTR_GEN_AI_PROVIDER_NAME, provider);
+    }
+
     let llmStreamReader: ReadableStreamDefaultReader<string | ChatChunk> | null = null;
     let llmStream: ReadableStream<string | ChatChunk> | null = null;
+    const startTime = performance.now() / 1000; // Convert to seconds
+    let firstTokenReceived = false;
 
     try {
       llmStream = await node(chatCtx, toolCtx, modelSettings);
@@ -402,16 +453,23 @@ export function performLLMInference(
         return;
       }
 
+      const abortPromise = waitForAbort(signal);
+
       // TODO(brian): add support for dynamic tools
 
       llmStreamReader = llmStream.getReader();
       while (true) {
-        if (signal.aborted) {
-          break;
-        }
-        const { done, value: chunk } = await llmStreamReader.read();
-        if (done) {
-          break;
+        if (signal.aborted) break;
+
+        const result = await Promise.race([llmStreamReader.read(), abortPromise]);
+        if (result === undefined) break;
+
+        const { done, value: chunk } = result;
+        if (done) break;
+
+        if (!firstTokenReceived) {
+          firstTokenReceived = true;
+          data.ttft = performance.now() / 1000 - startTime;
         }
 
         if (typeof chunk === 'string') {
@@ -431,6 +489,9 @@ export function performLLMInference(
                 callId: `${data.id}/fnc_${data.generatedToolCalls.length}`,
                 name: tool.name,
                 args: tool.args,
+                // Preserve thought signature for Gemini 3+ thinking mode
+                thoughtSignature: tool.thoughtSignature,
+                extra: tool.extra || {},
               });
 
               data.generatedToolCalls.push(toolCall);
@@ -447,6 +508,11 @@ export function performLLMInference(
         // No need to check if chunk is of type other than ChatChunk or string like in
         // Python since chunk is defined in the type ChatChunk | string in TypeScript
       }
+
+      span.setAttribute(traceTypes.ATTR_RESPONSE_TEXT, data.generatedText);
+      if (data.ttft !== undefined) {
+        span.setAttribute(traceTypes.ATTR_RESPONSE_TTFT, data.ttft);
+      }
     } catch (error) {
       if (error instanceof DOMException && error.name === 'AbortError') {
         // Abort signal was triggered, handle gracefully
@@ -461,6 +527,15 @@ export function performLLMInference(
     }
   };
 
+  // Capture the current context (agent_turn) to ensure llm_node is properly parented
+  const currentContext = otelContext.active();
+
+  const inferenceTask = async (signal: AbortSignal) =>
+    tracer.startActiveSpan(async (span) => _performLLMInferenceImpl(signal, span), {
+      name: 'llm_node',
+      context: currentContext,
+    });
+
   return [
     Task.from((controller) => inferenceTask(controller.signal), controller, 'performLLMInference'),
     data,
@@ -469,35 +544,131 @@ export function performLLMInference(
 
 export function performTTSInference(
   node: TTSNode,
-  text: ReadableStream<string>,
+  text: ReadableStream<string | TimedString>,
   modelSettings: ModelSettings,
   controller: AbortController,
-): [Task<void>, ReadableStream<AudioFrame>] {
+  model?: string,
+  provider?: string,
+): [Task<void>, _TTSGenerationData] {
   const audioStream = new IdentityTransform<AudioFrame>();
   const outputWriter = audioStream.writable.getWriter();
   const audioOutputStream = audioStream.readable;
 
-  const inferenceTask = async (signal: AbortSignal) => {
+  const timedTextsFut = new Future<ReadableStream<TimedString> | null>();
+  const timedTextsStream = new IdentityTransform<TimedString>();
+  const timedTextsWriter = timedTextsStream.writable.getWriter();
+
+  // Transform stream to extract text from TimedString objects
+  const textOnlyStream = new IdentityTransform<string>();
+  const textOnlyWriter = textOnlyStream.writable.getWriter();
+  (async () => {
+    const reader = text.getReader();
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) {
+          break;
+        }
+        const textValue = typeof value === 'string' ? value : value.text;
+        await textOnlyWriter.write(textValue);
+      }
+      await textOnlyWriter.close();
+    } catch (e) {
+      await textOnlyWriter.abort(e as Error);
+    } finally {
+      reader.releaseLock();
+    }
+  })();
+
+  let ttfb: number | undefined;
+
+  const genData: _TTSGenerationData = {
+    audioStream: audioOutputStream,
+    timedTextsFut,
+    ttfb: undefined,
+  };
+
+  const _performTTSInferenceImpl = async (signal: AbortSignal, span: Span) => {
+    if (model) {
+      span.setAttribute(traceTypes.ATTR_GEN_AI_REQUEST_MODEL, model);
+    }
+    if (provider) {
+      span.setAttribute(traceTypes.ATTR_GEN_AI_PROVIDER_NAME, provider);
+    }
+
     let ttsStreamReader: ReadableStreamDefaultReader<AudioFrame> | null = null;
     let ttsStream: ReadableStream<AudioFrame> | null = null;
+    let pushedDuration = 0;
+    const startTime = performance.now() / 1000; // Convert to seconds
+    let firstByteReceived = false;
 
     try {
-      ttsStream = await node(text, modelSettings);
+      ttsStream = await node(textOnlyStream.readable, modelSettings);
       if (ttsStream === null) {
+        timedTextsFut.resolve(null);
         await outputWriter.close();
+        await timedTextsWriter.close();
         return;
       }
 
+      // This is critical: the future must be resolved with the channel/stream before the loop
+      // so that agent_activity can start reading while we write
+      if (!timedTextsFut.done) {
+        timedTextsFut.resolve(timedTextsStream.readable);
+      }
+
       ttsStreamReader = ttsStream.getReader();
+
+      // In Python, perform_tts_inference has a while loop processing multiple input segments
+      // (separated by FlushSentinel), with pushed_duration accumulating across segments.
+      // JS currently only does single inference, so initialPushedDuration is always 0.
+      // TODO: Add FlushSentinel + multi-segment loop
+      const initialPushedDuration = pushedDuration;
+
       while (true) {
         if (signal.aborted) {
           break;
         }
-        const { done, value: chunk } = await ttsStreamReader.read();
+        const { done, value: frame } = await ttsStreamReader.read();
         if (done) {
           break;
         }
-        await outputWriter.write(chunk);
+
+        if (!firstByteReceived) {
+          firstByteReceived = true;
+          ttfb = performance.now() / 1000 - startTime;
+          genData.ttfb = ttfb;
+          span.setAttribute(traceTypes.ATTR_RESPONSE_TTFB, ttfb);
+        }
+
+        // Write the audio frame to the output stream
+        await outputWriter.write(frame);
+
+        const timedTranscripts = frame.userdata[USERDATA_TIMED_TRANSCRIPT] as
+          | TimedString[]
+          | undefined;
+        if (timedTranscripts && timedTranscripts.length > 0) {
+          for (const timedText of timedTranscripts) {
+            // Uses the INITIAL value (from previous inferences), not the accumulated value
+            const adjustedTimedText = createTimedString({
+              text: timedText.text,
+              startTime:
+                timedText.startTime !== undefined
+                  ? timedText.startTime + initialPushedDuration
+                  : undefined,
+              endTime:
+                timedText.endTime !== undefined
+                  ? timedText.endTime + initialPushedDuration
+                  : undefined,
+              confidence: timedText.confidence,
+              startTimeOffset: timedText.startTimeOffset,
+            });
+            await timedTextsWriter.write(adjustedTimedText);
+          }
+        }
+
+        const frameDuration = frame.samplesPerChannel / frame.sampleRate;
+        pushedDuration += frameDuration;
       }
     } catch (error) {
       if (error instanceof DOMException && error.name === 'AbortError') {
@@ -506,15 +677,29 @@ export function performTTSInference(
       }
       throw error;
     } finally {
+      if (!timedTextsFut.done) {
+        // Ensure downstream consumers don't hang on errors.
+        timedTextsFut.resolve(null);
+      }
       ttsStreamReader?.releaseLock();
       await ttsStream?.cancel();
       await outputWriter.close();
+      await timedTextsWriter.close();
     }
   };
 
+  // Capture the current context (agent_turn) to ensure tts_node is properly parented
+  const currentContext = otelContext.active();
+
+  const inferenceTask = async (signal: AbortSignal) =>
+    tracer.startActiveSpan(async (span) => _performTTSInferenceImpl(signal, span), {
+      name: 'tts_node',
+      context: currentContext,
+    });
+
   return [
     Task.from((controller) => inferenceTask(controller.signal), controller, 'performTTSInference'),
-    audioOutputStream,
+    genData,
   ];
 }
 
@@ -524,7 +709,7 @@ export interface _TextOut {
 }
 
 async function forwardText(
-  source: ReadableStream<string>,
+  source: ReadableStream<string | TimedString>,
   out: _TextOut,
   signal: AbortSignal,
   textOutput: TextOutput | null,
@@ -537,8 +722,13 @@ async function forwardText(
       }
       const { done, value: delta } = await reader.read();
       if (done) break;
-      out.text += delta;
+
+      const deltaIsTimedString = isTimedString(delta);
+      const textDelta = deltaIsTimedString ? delta.text : delta;
+
+      out.text += textDelta;
       if (textOutput !== null) {
+        // Pass TimedString to textOutput for synchronized transcription
         await textOutput.captureText(delta);
       }
       if (!out.firstTextFut.done) {
@@ -554,7 +744,7 @@ async function forwardText(
 }
 
 export function performTextForwarding(
-  source: ReadableStream<string>,
+  source: ReadableStream<string | TimedString>,
   controller: AbortController,
   textOutput: TextOutput | null,
 ): [Task<void>, _TextOut] {
@@ -574,19 +764,28 @@ export function performTextForwarding(
 
 export interface _AudioOut {
   audio: Array<AudioFrame>;
-  firstFrameFut: Future;
+  firstFrameFut: Future<number>;
 }
 
 async function forwardAudio(
   ttsStream: ReadableStream<AudioFrame>,
-  audioOuput: AudioOutput,
+  audioOutput: AudioOutput,
   out: _AudioOut,
   signal?: AbortSignal,
 ): Promise<void> {
   const reader = ttsStream.getReader();
   let resampler: AudioResampler | null = null;
 
+  const onPlaybackStarted = (ev: { createdAt: number }) => {
+    if (!out.firstFrameFut.done) {
+      out.firstFrameFut.resolve(ev.createdAt);
+    }
+  };
+
   try {
+    audioOutput.on(AudioOutput.EVENT_PLAYBACK_STARTED, onPlaybackStarted);
+    audioOutput.resume();
+
     while (true) {
       if (signal?.aborted) {
         break;
@@ -599,35 +798,36 @@ async function forwardAudio(
 
       if (
         !out.firstFrameFut.done &&
-        audioOuput.sampleRate &&
-        audioOuput.sampleRate !== frame.sampleRate &&
+        audioOutput.sampleRate &&
+        audioOutput.sampleRate !== frame.sampleRate &&
         !resampler
       ) {
-        resampler = new AudioResampler(frame.sampleRate, audioOuput.sampleRate, 1);
+        resampler = new AudioResampler(frame.sampleRate, audioOutput.sampleRate, 1);
       }
 
       if (resampler) {
         for (const f of resampler.push(frame)) {
-          await audioOuput.captureFrame(f);
+          await audioOutput.captureFrame(f);
         }
       } else {
-        await audioOuput.captureFrame(frame);
-      }
-
-      // set the first frame future if not already set
-      // (after completing the first frame)
-      if (!out.firstFrameFut.done) {
-        out.firstFrameFut.resolve();
+        await audioOutput.captureFrame(frame);
       }
     }
-  } finally {
-    reader?.releaseLock();
+
     if (resampler) {
       for (const f of resampler.flush()) {
-        await audioOuput.captureFrame(f);
+        await audioOutput.captureFrame(f);
       }
     }
-    audioOuput.flush();
+  } finally {
+    audioOutput.off(AudioOutput.EVENT_PLAYBACK_STARTED, onPlaybackStarted);
+
+    if (!out.firstFrameFut.done) {
+      out.firstFrameFut.reject(new Error('audio forwarding cancelled before playback started'));
+    }
+
+    reader?.releaseLock();
+    audioOutput.flush();
   }
 }
 
@@ -636,10 +836,11 @@ export function performAudioForwarding(
   audioOutput: AudioOutput,
   controller: AbortController,
 ): [Task<void>, _AudioOut] {
-  const out = {
+  const out: _AudioOut = {
     audio: [],
-    firstFrameFut: new Future(),
+    firstFrameFut: new Future<number>(),
   };
+
   return [
     Task.from(
       (controller) => forwardAudio(ttsStream, audioOutput, out, controller.signal),
@@ -684,7 +885,7 @@ export function performToolExecutions({
     const signal = controller.signal;
     const reader = toolCallStream.getReader();
 
-    const tasks: Promise<any>[] = [];
+    const tasks: Task<void>[] = [];
     while (!signal.aborted) {
       const { done, value: toolCall } = await reader.read();
       if (signal.aborted) break;
@@ -732,8 +933,13 @@ export function performToolExecutions({
       try {
         const jsonArgs = JSON.parse(toolCall.args);
 
-        if (tool.parameters instanceof ZodObject) {
-          parsedArgs = tool.parameters.parse(jsonArgs);
+        if (isZodSchema(tool.parameters)) {
+          const result = await parseZodSchema<object>(tool.parameters, jsonArgs);
+          if (result.success) {
+            parsedArgs = result.data;
+          } else {
+            throw result.error;
+          }
         } else {
           parsedArgs = jsonArgs;
         }
@@ -772,16 +978,9 @@ export function performToolExecutions({
         'Executing LLM tool call',
       );
 
-      const toolExecution = asyncLocalStorage.run({ functionCall: toolCall }, async () => {
-        return await tool.execute(parsedArgs, {
-          ctx: new RunContext(session, speechHandle, toolCall),
-          toolCallId: toolCall.callId,
-          abortSignal: signal,
-        });
-      });
-
-      const tracableToolExecution = async (toolExecTask: Promise<unknown>) => {
-        // TODO(brian): add tracing
+      const _tracableToolExecutionImpl = async (toolExecTask: Promise<unknown>, span: Span) => {
+        span.setAttribute(traceTypes.ATTR_FUNCTION_TOOL_NAME, toolCall.name);
+        span.setAttribute(traceTypes.ATTR_FUNCTION_TOOL_ARGS, toolCall.args);
 
         // await for task to complete, if task is aborted, set exception
         let toolOutput: ToolExecutionOutput | undefined;
@@ -792,6 +991,17 @@ export function performToolExecutions({
             exception: isAborted ? new Error('tool call was aborted') : undefined,
             output: isAborted ? undefined : result,
           });
+
+          if (toolOutput.toolCallOutput) {
+            span.setAttribute(
+              traceTypes.ATTR_FUNCTION_TOOL_OUTPUT,
+              toolOutput.toolCallOutput.output,
+            );
+            span.setAttribute(
+              traceTypes.ATTR_FUNCTION_TOOL_IS_ERROR,
+              toolOutput.toolCallOutput.isError,
+            );
+          }
         } catch (rawError) {
           logger.error(
             {
@@ -805,17 +1015,61 @@ export function performToolExecutions({
             toolCall,
             exception: toError(rawError),
           });
+
+          if (toolOutput.toolCallOutput) {
+            span.setAttribute(
+              traceTypes.ATTR_FUNCTION_TOOL_OUTPUT,
+              toolOutput.toolCallOutput.output,
+            );
+            span.setAttribute(traceTypes.ATTR_FUNCTION_TOOL_IS_ERROR, true);
+          }
         } finally {
           if (!toolOutput) throw new Error('toolOutput is undefined');
           toolCompleted(toolOutput);
         }
       };
 
+      const tracableToolExecution = (toolExecTask: Promise<unknown>) =>
+        tracer.startActiveSpan(async (span) => _tracableToolExecutionImpl(toolExecTask, span), {
+          name: 'function_tool',
+        });
+
+      const toolTask = Task.from(
+        async () => {
+          // Ensure this task is marked inline before user tool code executes.
+          const currentTask = Task.current();
+          if (currentTask) {
+            _setActivityTaskInfo(currentTask, {
+              speechHandle,
+              functionCall: toolCall,
+              inlineTask: true,
+            });
+          }
+
+          const toolExecution = functionCallStorage.run({ functionCall: toolCall }, async () => {
+            return await tool.execute(parsedArgs, {
+              ctx: new RunContext(session, speechHandle, toolCall),
+              toolCallId: toolCall.callId,
+              abortSignal: signal,
+            });
+          });
+
+          await tracableToolExecution(toolExecution);
+        },
+        controller,
+        `performToolExecution:${toolCall.name}`,
+      );
+
+      _setActivityTaskInfo(toolTask, {
+        speechHandle,
+        functionCall: toolCall,
+        inlineTask: true,
+      });
       // wait, not cancelling all tool calling tasks
-      tasks.push(tracableToolExecution(toolExecution));
+      tasks.push(toolTask);
     }
 
-    await Promise.allSettled(tasks);
+    await Promise.allSettled(tasks.map((task) => task.result));
     if (toolOutput.output.length > 0) {
       logger.debug(
         {