Refactor TranscriptProcessor into user and assistant processors

This commit is contained in:
Mark Backman
2024-12-16 15:24:58 -05:00
parent 4211664a77
commit 1117c21483
8 changed files with 185 additions and 158 deletions

View File

@@ -207,13 +207,6 @@ class InterimTranscriptionFrame(TextFrame):
return f"{self.name}(user: {self.user_id}, text: [{self.text}], language: {self.language}, timestamp: {self.timestamp})"
@dataclass
class OpenAILLMContextUserTimestampFrame(DataFrame):
"""Timestamp information for user message in LLM context."""
timestamp: str
@dataclass
class OpenAILLMContextAssistantTimestampFrame(DataFrame):
"""Timestamp information for assistant message in LLM context."""

View File

@@ -15,7 +15,6 @@ from pipecat.frames.frames import (
LLMMessagesFrame,
LLMMessagesUpdateFrame,
LLMSetToolsFrame,
OpenAILLMContextUserTimestampFrame,
StartInterruptionFrame,
TextFrame,
TranscriptionFrame,
@@ -27,7 +26,6 @@ from pipecat.processors.aggregators.openai_llm_context import (
OpenAILLMContextFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.utils.time import time_now_iso8601
class LLMResponseAggregator(FrameProcessor):
@@ -291,10 +289,6 @@ class LLMContextAggregator(LLMResponseAggregator):
frame = OpenAILLMContextFrame(self._context)
await self.push_frame(frame)
# Push timestamp frame with current time
timestamp_frame = OpenAILLMContextUserTimestampFrame(timestamp=time_now_iso8601())
await self.push_frame(timestamp_frame)
# Reset our accumulator state.
self._reset()

View File

@@ -4,7 +4,8 @@
# SPDX-License-Identifier: BSD 2-Clause License
#
from typing import List, Optional
from abc import ABC, abstractmethod
from typing import List
from loguru import logger
@@ -12,7 +13,7 @@ from pipecat.frames.frames import (
ErrorFrame,
Frame,
OpenAILLMContextAssistantTimestampFrame,
OpenAILLMContextUserTimestampFrame,
TranscriptionFrame,
TranscriptionMessage,
TranscriptionUpdateFrame,
)
@@ -20,55 +21,72 @@ from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContextFr
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
class TranscriptProcessor(FrameProcessor):
"""Processes LLM context frames to generate timestamped conversation transcripts.
class BaseTranscriptProcessor(FrameProcessor, ABC):
"""Base class for processing conversation transcripts.
This processor monitors OpenAILLMContextFrame frames and their corresponding
timestamp frames to build a chronological conversation transcript. Messages are
stored by role until their matching timestamp frame arrives, then emitted via
TranscriptionUpdateFrame.
Each LLM context (OpenAI, Anthropic, Google) provides conversion to the standard format:
[
{
"role": "user",
"content": [{"type": "text", "text": "Hi, how are you?"}]
},
{
"role": "assistant",
"content": [{"type": "text", "text": "Great! And you?"}]
}
]
Events:
on_transcript_update: Emitted when timestamped messages are available.
Args: TranscriptionUpdateFrame containing timestamped messages.
Example:
```python
transcript_processor = TranscriptProcessor()
@transcript_processor.event_handler("on_transcript_update")
async def on_transcript_update(processor, frame):
for msg in frame.messages:
print(f"[{msg.timestamp}] {msg.role}: {msg.content}")
```
Provides common functionality for handling transcript messages and updates.
"""
def __init__(self, **kwargs):
"""Initialize the transcript processor.
Args:
**kwargs: Additional arguments passed to FrameProcessor
"""
"""Initialize processor with empty message store."""
super().__init__(**kwargs)
self._processed_messages: List[TranscriptionMessage] = []
self._register_event_handler("on_transcript_update")
self._pending_user_messages: List[TranscriptionMessage] = []
async def _emit_update(self, messages: List[TranscriptionMessage]):
"""Emit transcript updates for new messages.
Args:
messages: New messages to emit in update
"""
if messages:
self._processed_messages.extend(messages)
update_frame = TranscriptionUpdateFrame(messages=messages)
await self._call_event_handler("on_transcript_update", update_frame)
await self.push_frame(update_frame)
@abstractmethod
async def process_frame(self, frame: Frame, direction: FrameDirection):
"""Process incoming frames to build conversation transcript.
Args:
frame: Input frame to process
direction: Frame processing direction
"""
await super().process_frame(frame, direction)
class UserTranscriptProcessor(BaseTranscriptProcessor):
"""Processes user transcription frames into timestamped conversation messages."""
async def process_frame(self, frame: Frame, direction: FrameDirection):
"""Process TranscriptionFrames into user conversation messages.
Args:
frame: Input frame to process
direction: Frame processing direction
"""
await super().process_frame(frame, direction)
if isinstance(frame, TranscriptionFrame):
message = TranscriptionMessage(
role="user", content=frame.text, timestamp=frame.timestamp
)
await self._emit_update([message])
await self.push_frame(frame, direction)
class AssistantTranscriptProcessor(BaseTranscriptProcessor):
"""Processes assistant LLM context frames into timestamped conversation messages."""
def __init__(self, **kwargs):
"""Initialize processor with empty message stores."""
super().__init__(**kwargs)
self._pending_assistant_messages: List[TranscriptionMessage] = []
def _extract_messages(self, messages: List[dict]) -> List[TranscriptionMessage]:
"""Extract conversation messages from standard format.
"""Extract assistant messages from the OpenAI standard message format.
Args:
messages: List of messages in OpenAI format, which can be either:
@@ -80,21 +98,14 @@ class TranscriptProcessor(FrameProcessor):
"""
result = []
for msg in messages:
# Only process user and assistant messages
if msg["role"] not in ("user", "assistant"):
continue
if "content" not in msg:
logger.warning(f"Message missing content field: {msg}")
if msg["role"] != "assistant":
continue
content = msg.get("content")
if isinstance(content, str):
# Handle simple string content
if content:
result.append(TranscriptionMessage(role=msg["role"], content=content))
result.append(TranscriptionMessage(role="assistant", content=content))
elif isinstance(content, list):
# Handle structured content
text_parts = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
@@ -102,13 +113,13 @@ class TranscriptProcessor(FrameProcessor):
if text_parts:
result.append(
TranscriptionMessage(role=msg["role"], content=" ".join(text_parts))
TranscriptionMessage(role="assistant", content=" ".join(text_parts))
)
return result
def _find_new_messages(self, current: List[TranscriptionMessage]) -> List[TranscriptionMessage]:
"""Find messages in current that aren't in self._processed_messages.
"""Find unprocessed messages from current list.
Args:
current: List of current messages
@@ -126,28 +137,15 @@ class TranscriptProcessor(FrameProcessor):
return current[processed_len:]
async def process_frame(self, frame: Frame, direction: FrameDirection):
"""Process frames to build a timestamped conversation transcript.
Handles three frame types in sequence:
1. OpenAILLMContextFrame: Contains new messages to be timestamped
2. OpenAILLMContextUserTimestampFrame: Timestamp for user messages
3. OpenAILLMContextAssistantTimestampFrame: Timestamp for assistant messages
Messages are stored by role until their corresponding timestamp frame arrives.
When a timestamp frame is received, the matching messages are timestamped and
emitted in chronological order via TranscriptionUpdateFrame.
"""Process frames into assistant conversation messages.
Args:
frame: The frame to process
frame: Input frame to process
direction: Frame processing direction
Raises:
ErrorFrame: If message processing fails
"""
await super().process_frame(frame, direction)
if isinstance(frame, OpenAILLMContextFrame):
# Extract and store messages by role
standard_messages = []
for msg in frame.context.messages:
converted = frame.context.to_standard_messages(msg)
@@ -155,34 +153,83 @@ class TranscriptProcessor(FrameProcessor):
current_messages = self._extract_messages(standard_messages)
new_messages = self._find_new_messages(current_messages)
# Store new messages by role
for msg in new_messages:
if msg.role == "user":
self._pending_user_messages.append(msg)
elif msg.role == "assistant":
self._pending_assistant_messages.append(msg)
elif isinstance(frame, OpenAILLMContextUserTimestampFrame):
# Process pending user messages with timestamp
if self._pending_user_messages:
for msg in self._pending_user_messages:
msg.timestamp = frame.timestamp
self._processed_messages.extend(self._pending_user_messages)
update_frame = TranscriptionUpdateFrame(messages=self._pending_user_messages)
await self._call_event_handler("on_transcript_update", update_frame)
await self.push_frame(update_frame)
self._pending_user_messages = []
self._pending_assistant_messages.extend(new_messages)
elif isinstance(frame, OpenAILLMContextAssistantTimestampFrame):
# Process pending assistant messages with timestamp
if self._pending_assistant_messages:
for msg in self._pending_assistant_messages:
msg.timestamp = frame.timestamp
self._processed_messages.extend(self._pending_assistant_messages)
update_frame = TranscriptionUpdateFrame(messages=self._pending_assistant_messages)
await self._call_event_handler("on_transcript_update", update_frame)
await self.push_frame(update_frame)
await self._emit_update(self._pending_assistant_messages)
self._pending_assistant_messages = []
await self.push_frame(frame, direction)
class TranscriptProcessor:
"""Factory for creating and managing transcript processors.
Provides unified access to user and assistant transcript processors
with shared event handling.
Example:
```python
transcript = TranscriptProcessor()
pipeline = Pipeline(
[
transport.input(),
stt,
transcript.user(), # User transcripts
context_aggregator.user(),
llm,
tts,
transport.output(),
context_aggregator.assistant(),
transcript.assistant(), # Assistant transcripts
]
)
@transcript.event_handler("on_transcript_update")
async def handle_update(processor, frame):
print(f"New messages: {frame.messages}")
```
"""
def __init__(self, **kwargs):
"""Initialize factory with user and assistant processors."""
self._user_processor = UserTranscriptProcessor(**kwargs)
self._assistant_processor = AssistantTranscriptProcessor(**kwargs)
self._event_handlers = {}
def user(self) -> UserTranscriptProcessor:
"""Get the user transcript processor."""
return self._user_processor
def assistant(self) -> AssistantTranscriptProcessor:
"""Get the assistant transcript processor."""
return self._assistant_processor
def event_handler(self, event_name: str):
"""Register event handler for both processors.
Args:
event_name: Name of event to handle
Returns:
Decorator function that registers handler with both processors
"""
def decorator(handler):
self._event_handlers[event_name] = handler
@self._user_processor.event_handler(event_name)
async def user_handler(processor, frame):
return await handler(processor, frame)
@self._assistant_processor.event_handler(event_name)
async def assistant_handler(processor, frame):
return await handler(processor, frame)
return handler
return decorator

View File

@@ -24,7 +24,6 @@ from pipecat.frames.frames import (
LLMMessagesFrame,
LLMUpdateSettingsFrame,
OpenAILLMContextAssistantTimestampFrame,
OpenAILLMContextUserTimestampFrame,
TextFrame,
TTSAudioRawFrame,
TTSStartedFrame,
@@ -234,10 +233,6 @@ class GoogleUserContextAggregator(OpenAIUserContextAggregator):
frame = OpenAILLMContextFrame(self._context)
await self.push_frame(frame)
# Push timestamp frame with current time
timestamp_frame = OpenAILLMContextUserTimestampFrame(timestamp=time_now_iso8601())
await self.push_frame(timestamp_frame)
# Reset our accumulator state.
self._reset()