# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD 2-Clause License """NVIDIA Context Aggregator. This module provides specialized frame processors and context aggregators for handling NVIDIA's Speculative Speech Processing feature in conversational AI systems. It manages the processing and aggregation of interim and final transcripts, enabling real-time response generation while maintaining conversation coherence. The processors handle: - Interim transcript processing for early response generation - Context management for speculative responses - TTS response caching and timing control for natural turn-taking - Bidirectional conversation state management Also see: pipecat.processors.aggregators.llm_response pipecat.processors.aggregators.openai_llm_context Classes: NvidiaAssistantContextAggregator: Handles assistant-specific context aggregation. NvidiaUserContextAggregator: Manages user context with interim/final transcripts. NvidiaTTSResponseCacher: Controls TTS response timing. NvidiaContextAggregatorPair: Coordinates paired aggregators. Functions: create_nvidia_context_aggregator: Factory for creating aggregator pairs. """ from dataclasses import dataclass from loguru import logger from pipecat.frames.frames import ( Frame, LLMFullResponseEndFrame, LLMFullResponseStartFrame, StartInterruptionFrame, TranscriptionFrame, TTSAudioRawFrame, TTSStartedFrame, TTSStoppedFrame, TTSTextFrame, UserStartedSpeakingFrame, UserStoppedSpeakingFrame, ) from pipecat.processors.aggregators.llm_response import ( LLMAssistantAggregatorParams, LLMAssistantContextAggregator, LLMUserAggregatorParams, LLMUserContextAggregator, ) from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext, OpenAILLMContextFrame from pipecat.processors.frame_processor import FrameDirection, FrameProcessor from nvidia_pipecat.frames.action import StartedPresenceUserActionFrame from nvidia_pipecat.frames.riva import RivaInterimTranscriptionFrame class NvidiaAssistantContextAggregator(LLMAssistantContextAggregator): """Extends LLMAssistantContextAggregator for NVIDIA-specific requirements. Specializes the base aggregator for handling speculative speech processing, managing assistant responses and context updates based on interim/final transcripts. Args: context (OpenAILLMContext): The context object to use. expect_stripped_words (bool): Whether to expect preprocessed words. Defaults to True. **kwargs: Additional arguments passed to parent class. Input Frames: LLMFullResponseStartFrame: Marks response start LLMFullResponseEndFrame: Marks response end TextFrame: Contains response text StartInterruptionFrame: Signals interruption Output Frames: OpenAILLMContextFrame: Updated context with responses """ async def push_aggregation(self): """Updates the context with current aggregation. For speculative processing, may update existing messages rather than append to maintain context coherence with interim transcripts. - If the last message in context has the same role, it updates that message - Otherwise, appends a new message with the current aggregation - After pushing, resets the aggregation state Returns: None Typical usage example: >>> context = OpenAILLMContext() >>> aggregator = NvidiaAssistantContextAggregator(context) >>> # Update existing response >>> context.add_message({"role": "assistant", "content": "initial response"}) >>> aggregator._aggregation = "updated response" >>> await aggregator.push_aggregation() """ if len(self._aggregation) > 0: context_messages = self.context.get_messages() # Update existing message if same role, otherwise append new one if len(context_messages) > 0 and context_messages[-1]["role"] == self._role: context_messages[-1]["content"] = self._aggregation self.context.set_messages(context_messages) else: self.context.add_message({"role": self._role, "content": self._aggregation}) self._aggregation = "" frame = OpenAILLMContextFrame(self.context) await self.push_frame(frame) # Reset our accumulator state. self.reset() class NvidiaUserContextAggregator(LLMUserContextAggregator): """Extends LLMUserContextAggregator for user-specific context handling. Handles speculative speech processing with interim and final transcriptions. Key features for speculative processing: - Processes stable interim transcripts for early response generation - Manages transition from interim to final transcripts - Deduplicates repeated transcripts to prevent context pollution - Maintains conversation history with configurable turn limits - Tracks user speaking state to coordinate with assistant responses Input Frames: TranscriptionFrame: Final transcription RivaInterimTranscriptionFrame: Interim transcription UserStartedSpeakingFrame: User began speaking UserStoppedSpeakingFrame: User stopped speaking StartInterruptionFrame: Conversation interruption Output Frames: OpenAILLMContextFrame: Updated context with transcripts """ def __init__( self, send_interims: bool = True, chat_history_limit: int = 20, **kwargs, ): """Initialize the NvidiaUserContextAggregator. Args: send_interims (bool, optional): Whether to send interim transcription frames. Defaults to True. chat_history_limit (int): Limits the number of turns in chat history, **kwargs: Additional keyword arguments passed to parent LLMUserContextAggregator. """ super().__init__(**kwargs) self.send_interims = send_interims self.chat_history_limit = chat_history_limit self.last_transcript = None self._user_speaking = False self.seen_final = True async def process_frame(self, frame: Frame, direction: FrameDirection): """Process a frame for speculative speech handling. - Processes stable interim transcripts when user is speaking - Manages transition between interim and final transcripts - Handles user speaking state changes - Deduplicates repeated transcripts Args: frame: Frame to process (TranscriptionFrame, RivaInterimTranscriptionFrame, or user state frames) direction: Direction of frame flow in the pipeline Typical usage example: >>> aggregator = NvidiaUserContextAggregator( ... send_interims=True, # Enable interim transcript processing ... chat_history_limit=20 # Keep last 20 conversation turns ... ) >>> # Process final transcript >>> frame = TranscriptionFrame(text="Hello") >>> await aggregator.process_frame(frame, FrameDirection.DOWNSTREAM) >>> >>> # Process interim transcript >>> frame = RivaInterimTranscriptionFrame(text="Hello", stability=1.0) >>> await aggregator.process_frame(frame, FrameDirection.DOWNSTREAM) """ if isinstance(frame, TranscriptionFrame): logger.debug(f"Recieved final transcript at NvidiaUserContextAggregator {frame.text}") # Only process if this is a new transcript if self.last_transcript is None or (self.last_transcript.rstrip() != frame.text.rstrip()): logger.debug(f"Sent final transcript downstream to LLM from NvidiaUserContextAggregator {frame.text}") self._aggregation = frame.text await self.push_aggregation() self.last_transcript = None self.seen_final = True elif isinstance(frame, RivaInterimTranscriptionFrame): # Process stable interim transcriptions during active speech and before first final result if self.send_interims and (self._user_speaking or not self.seen_final) and frame.stability == 1.0: logger.debug(f"Sent interim transcript downstream to LLM from NvidiaUserContextAggregator {frame.text}") self._aggregation = frame.text await self.push_aggregation() self.last_transcript = frame.text elif isinstance(frame, StartInterruptionFrame): self._user_speaking = False await self._start_interruption() await self.stop_all_metrics() await self.push_frame(frame, direction) else: if isinstance(frame, UserStartedSpeakingFrame): self._user_speaking = True self.seen_final = False elif isinstance(frame, UserStoppedSpeakingFrame): self._user_speaking = False await super().process_frame(frame, direction) async def get_truncated_context(self) -> OpenAILLMContext: """Returns a truncated context limited to specified chat history size. - Counts conversation turns based on user-assistant exchanges - Preserves system and function messages regardless of limit - Processes messages in reverse order to maintain recent history Returns: OpenAILLMContext: New context object containing truncated conversation history, preserving system/function messages and most recent turns. Typical usage example: >>> aggregator = NvidiaUserContextAggregator(chat_history_limit=2) >>> # Context with 3 turns >>> context = OpenAILLMContext() >>> # Turn 1 >>> context.add_message({"role": "user", "content": "Turn 1 user"}) >>> context.add_message({"role": "assistant", "content": "Turn 1 assistant"}) >>> # Turn 2 >>> context.add_message({"role": "user", "content": "Turn 2 user"}) >>> context.add_message({"role": "assistant", "content": "Turn 2 assistant"}) >>> # Turn 3 >>> context.add_message({"role": "user", "content": "Turn 3 user"}) >>> context.add_message({"role": "assistant", "content": "Turn 3 assistant"}) >>> # Get truncated context - will only contain most recent 2 turns >>> truncated = await aggregator.get_truncated_context() >>> print(truncated.get_messages()) # Shows turns 2 and 3 only """ truncated_context = self.context if len(self.context.get_messages()) > 0: truncated_context = OpenAILLMContext() truncated_context_messages = [] current_size = 0 for context_message in reversed(self.context.get_messages()): if ( context_message["role"] == "user" or context_message["role"] == "assistant" or context_message["role"] == "developer" or context_message["role"] == "function" or context_message["role"] == "tool" ): if current_size == self.chat_history_limit: continue if context_message["role"] == "user": current_size = current_size + 1 truncated_context_messages.append(context_message) truncated_context.set_messages(reversed(truncated_context_messages)) return truncated_context async def push_aggregation(self): """Pushes aggregation to context and manages conversation flow. Updates or appends current aggregation to conversation context while maintaining turn-taking structure. For speculative responses, may update existing message rather than append new one. - If the last message in context has the same role, it updates that message - Otherwise, appends a new message with the current aggregation - After pushing, resets the aggregation state Output Frames: OpenAILLMContextFrame: downstream after processing. Typical usage example: >>> context = OpenAILLMContext() >>> aggregator = NvidiaUserContextAggregator(context) >>> # Update existing response >>> context.add_message({"role": "user", "content": "initial query"}) >>> aggregator._aggregation = "updated query" >>> await aggregator.push_aggregation() """ if len(self._aggregation) > 0: context_messages = self.context.get_messages() # Update existing message if same role, otherwise append new one if len(context_messages) > 0 and context_messages[-1]["role"] == self._role: context_messages[-1]["content"] = self._aggregation self.context.set_messages(context_messages) else: self.context.add_message({"role": self._role, "content": self._aggregation}) self._aggregation = "" # Get truncated context and send downstream truncated_context = await self.get_truncated_context() frame = OpenAILLMContextFrame(truncated_context) # Send the interruption before the context frame await self.push_frame(StartInterruptionFrame()) logger.debug( f"Sending context downstream to LLM from NvidiaUserContextAggregator {frame.context.get_messages()}" ) await self.push_frame(frame) # Reset our accumulator state self.reset() class NvidiaTTSResponseCacher(FrameProcessor): """Caches TTS responses and controls release timing for speculative speech. Manages text-to-speech response timing by caching responses and controlling their release based on user speaking state. Maintains natural turn-taking and prevents response overlap during speculative processing. Input frames handled: - LLMFullResponseStartFrame: Marks response start - LLMFullResponseEndFrame: Marks response end - TTSAudioRawFrame: TTS audio data - TTSStartedFrame: TTS start marker - TTSStoppedFrame: TTS stop marker - TTSTextFrame: TTS text data - UserStartedSpeakingFrame: Triggers caching - UserStoppedSpeakingFrame: Triggers release - StartInterruptionFrame: Clears cache """ def __init__(self) -> None: """Initialize the NvidiaTTSResponseCacher.""" super().__init__() self._cache: list[Frame] = [] self.user_stopped_speaking: bool = True async def process_frame(self, frame: Frame, direction: FrameDirection) -> None: """Processes frame for TTS response caching and timing control. - Caches TTS responses while user is speaking - Releases cached responses when user stops speaking - Clears cache on interruptions - Maintains conversation flow by coordinating response timing Also see: - NvidiaUserContextAggregator : Handles user context and speech state - NvidiaAssistantContextAggregator : Manages assistant responses Args: frame: Frame to process direction: Direction of frame flow in pipeline Typical usage example: >>> cacher = NvidiaTTSResponseCacher() >>> # User starts speaking >>> await cacher.process_frame(UserStartedSpeakingFrame(), FrameDirection.DOWNSTREAM) >>> # TTS response arrives - will be cached >>> await cacher.process_frame(TTSAudioRawFrame(audio_data), FrameDirection.DOWNSTREAM) >>> # User stops speaking - cached responses will be released >>> await cacher.process_frame(UserStoppedSpeakingFrame(), FrameDirection.DOWNSTREAM) """ await super().process_frame(frame, direction) # Handle response start - cache if user is speaking if isinstance(frame, LLMFullResponseStartFrame): if self.user_stopped_speaking: await self.push_frame(frame, direction) else: self._cache.clear() # Clear existing cache before new response self._cache.append(frame) # Handle TTS frames - cache or forward based on user speaking state elif isinstance(frame, (TTSAudioRawFrame | TTSStartedFrame | TTSStoppedFrame | TTSTextFrame)): if self.user_stopped_speaking: await self.push_frame(frame, direction) else: self._cache.append(frame) # Handle response end - mark user can speak after forwarding elif isinstance(frame, LLMFullResponseEndFrame): if self.user_stopped_speaking: await self.push_frame(frame, direction) self.user_stopped_speaking = False # Allow user to speak after response ends self._cache.append(frame) # Handle interruptions - clear cache and reset state elif isinstance(frame, StartInterruptionFrame | StartedPresenceUserActionFrame): # TODO: This only works if we have a single user in the system. # it also does not work if other "events" should trigger the cache release # (e.g. new frames by new processors). self._cache.clear() # self.user_stopped_speaking = True await self.push_frame(frame, direction) # Handle user stop speaking - release cached responses elif isinstance(frame, UserStoppedSpeakingFrame): self.user_stopped_speaking = True if self._cache: for cached_frame in self._cache: await self.push_frame(cached_frame) self._cache.clear() await self.push_frame(frame, direction) # Handle user start speaking - update state elif isinstance(frame, UserStartedSpeakingFrame): self.user_stopped_speaking = False await self.push_frame(frame, direction) # Forward all other frames unchanged else: await self.push_frame(frame, direction) @dataclass class NvidiaContextAggregatorPair: """A pair of context aggregators for managing bidirectional conversation. Attributes: _user: NvidiaUserContextAggregator for user-side context _assistant: NvidiaAssistantContextAggregator for assistant-side context """ _user: "NvidiaUserContextAggregator" _assistant: "NvidiaAssistantContextAggregator" def user(self) -> "NvidiaUserContextAggregator": """Get the user context aggregator.""" return self._user def assistant(self) -> "NvidiaAssistantContextAggregator": """Get the assistant context aggregator.""" return self._assistant def create_nvidia_context_aggregator( context: OpenAILLMContext, assistant_expect_stripped_words: bool = True, send_interims: bool = False, chat_history_limit: int = 20, ) -> NvidiaContextAggregatorPair: """Creates a pair of context aggregators for speculative speech processing. - Creates synchronized user and assistant aggregators sharing context - User aggregator handles interim/final transcripts - Assistant aggregator manages response generation - Both work together to maintain conversation coherence Also see: - NvidiaUserContextAggregator : Handles user context - NvidiaAssistantContextAggregator : Handles assistant context Args: context: Base context object to initialize aggregators assistant_expect_stripped_words: Whether assistant expects preprocessed words send_interims: Whether to process interim transcriptions chat_history_limit: Maximum number of conversation turns to maintain Returns: NvidiaContextAggregatorPair: A paired set of user and assistant context aggregators configured for speculative speech processing. Typical usage example: >>> context = OpenAILLMContext() >>> # Create aggregators with default settings >>> aggregators = create_nvidia_context_aggregator(context) >>> >>> # Create aggregators with custom settings >>> aggregators = create_nvidia_context_aggregator( ... context, ... send_interims=True, # Enable interim transcript processing ... chat_history_limit=10, # Keep shorter history ... assistant_expect_stripped_words=False # Raw word processing ... ) >>> >>> # Access individual aggregators >>> user_aggregator = aggregators.user() >>> assistant_aggregator = aggregators.assistant() """ # Create user aggregator with specified settings user_params = LLMUserAggregatorParams(aggregation_timeout=0.01) user = NvidiaUserContextAggregator( send_interims=send_interims, context=context, params=user_params, chat_history_limit=chat_history_limit ) # Create assistant aggregator sharing context with user assistant_params = LLMAssistantAggregatorParams(expect_stripped_words=assistant_expect_stripped_words) assistant = NvidiaAssistantContextAggregator(context=user.context, params=assistant_params) return NvidiaContextAggregatorPair(_user=user, _assistant=assistant)