tracing_observer.py

"""
Custom LangSmith Tracing Observer for Pipecat Voice Agent.

This observer captures detailed information from STT, LLM, and TTS services
and adds them as OpenTelemetry span attributes for LangSmith visualization.
"""
import time
from typing import Optional
from loguru import logger

from opentelemetry import trace

from pipecat.observers.base_observer import BaseObserver
from pipecat.frames.frames import (
    Frame,
    TranscriptionFrame,
    InterimTranscriptionFrame,
    TextFrame,
    LLMFullResponseStartFrame,
    LLMFullResponseEndFrame,
    TTSStartedFrame,
    TTSStoppedFrame,
    LLMMessagesFrame,
)
from pipecat.processors.frame_processor import FrameDirection


class DetailedTracingObserver(BaseObserver):
    """
    Observer that captures detailed tracing information for LangSmith.
    
    Captures:
    - STT: transcripts, latency
    - LLM: input messages, output response, token usage, latency
    - TTS: input text, character count, latency
    """
    
    def __init__(self):
        super().__init__()
        self._tracer = trace.get_tracer(__name__)
        
        # State tracking
        self._current_stt_start: float = 0
        self._current_llm_start: float = 0
        self._current_tts_start: float = 0
        self._current_user_message: str = ""
        self._current_llm_response: str = ""
        self._current_tts_text: str = ""
        self._llm_messages: list = []
        
        # Spans
        self._stt_span: Optional[trace.Span] = None
        self._llm_span: Optional[trace.Span] = None
        self._tts_span: Optional[trace.Span] = None
        
    async def on_push_frame(
        self,
        src: "FrameProcessor",
        dst: "FrameProcessor", 
        frame: Frame,
        direction: FrameDirection,
        timestamp: int,
    ):
        """Called when a frame is pushed from one processor to another."""
        
        # STT Transcription received
        if isinstance(frame, TranscriptionFrame):
            self._handle_transcription(frame, src)
            
        # LLM Messages being sent (capture input)
        elif isinstance(frame, LLMMessagesFrame):
            self._handle_llm_messages(frame)
            
        # LLM Response starting
        elif isinstance(frame, LLMFullResponseStartFrame):
            self._handle_llm_start()
            
        # LLM Response text chunks
        elif isinstance(frame, TextFrame):
            if self._llm_span:
                self._current_llm_response += frame.text
                
        # LLM Response ended
        elif isinstance(frame, LLMFullResponseEndFrame):
            self._handle_llm_end()
            
        # TTS Started
        elif isinstance(frame, TTSStartedFrame):
            self._handle_tts_start()
            
        # TTS Stopped
        elif isinstance(frame, TTSStoppedFrame):
            self._handle_tts_end()
    
    def _handle_transcription(self, frame: TranscriptionFrame, src):
        """Handle STT transcription frame."""
        transcript = frame.text
        self._current_user_message = transcript
        
        with self._tracer.start_as_current_span("stt_transcription") as span:
            span.set_attribute("langsmith.span.kind", "LLM")
            span.set_attribute("gen_ai.system", "OpenAI")
            span.set_attribute("gen_ai.request.model", "whisper-1")
            span.set_attribute("gen_ai.operation.name", "stt")
            
            # Input/Output for STT
            span.set_attribute("stt.input.type", "audio")
            span.set_attribute("stt.output.transcript", transcript)
            span.set_attribute("stt.output.word_count", len(transcript.split()))
            span.set_attribute("stt.output.character_count", len(transcript))
            
            # LangSmith specific
            span.set_attribute("gen_ai.completion.0.content", transcript)
            span.set_attribute("gen_ai.completion.0.role", "transcription")
            
        logger.debug(f"Traced STT: '{transcript}'")
    
    def _handle_llm_messages(self, frame: LLMMessagesFrame):
        """Capture LLM input messages."""
        self._llm_messages = frame.messages if hasattr(frame, 'messages') else []
        self._current_llm_start = time.time()
        
    def _handle_llm_start(self):
        """Handle LLM response start."""
        if self._current_llm_start == 0:
            self._current_llm_start = time.time()
        self._current_llm_response = ""
        
        self._llm_span = self._tracer.start_span("llm_generation")
        self._llm_span.set_attribute("langsmith.span.kind", "LLM")
        self._llm_span.set_attribute("gen_ai.system", "AWS Bedrock")
        self._llm_span.set_attribute("gen_ai.request.model", "us.anthropic.claude-3-5-haiku-20241022-v1:0")
        self._llm_span.set_attribute("gen_ai.operation.name", "chat")
        
        # Add input messages
        if self._current_user_message:
            self._llm_span.set_attribute("gen_ai.prompt.0.role", "user")
            self._llm_span.set_attribute("gen_ai.prompt.0.content", self._current_user_message)
    
    def _handle_llm_end(self):
        """Handle LLM response end - record the full response."""
        if self._llm_span:
            latency_ms = (time.time() - self._current_llm_start) * 1000
            
            # Output response
            self._llm_span.set_attribute("gen_ai.completion.0.role", "assistant")
            self._llm_span.set_attribute("gen_ai.completion.0.content", self._current_llm_response)
            
            # Metrics
            self._llm_span.set_attribute("llm.latency_ms", latency_ms)
            self._llm_span.set_attribute("llm.output.character_count", len(self._current_llm_response))
            self._llm_span.set_attribute("llm.output.word_count", len(self._current_llm_response.split()))
            
            # Estimate tokens (rough approximation)
            input_tokens = len(self._current_user_message.split()) * 1.3
            output_tokens = len(self._current_llm_response.split()) * 1.3
            self._llm_span.set_attribute("gen_ai.usage.prompt_tokens", int(input_tokens))
            self._llm_span.set_attribute("gen_ai.usage.completion_tokens", int(output_tokens))
            self._llm_span.set_attribute("gen_ai.usage.total_tokens", int(input_tokens + output_tokens))
            
            self._llm_span.end()
            self._llm_span = None
            
            logger.debug(f"Traced LLM: {latency_ms:.0f}ms, output: {len(self._current_llm_response)} chars")
            
            # Reset for next turn
            self._current_llm_start = 0
    
    def _handle_tts_start(self):
        """Handle TTS start."""
        self._current_tts_start = time.time()
        self._current_tts_text = self._current_llm_response  # TTS will speak the LLM response
        
        self._tts_span = self._tracer.start_span("tts_synthesis")
        self._tts_span.set_attribute("langsmith.span.kind", "LLM")
        self._tts_span.set_attribute("gen_ai.system", "OpenAI")
        self._tts_span.set_attribute("gen_ai.request.model", "tts-1")
        self._tts_span.set_attribute("gen_ai.operation.name", "tts")
        
        # Input text
        self._tts_span.set_attribute("tts.input.text", self._current_tts_text[:500])  # Limit for attribute size
        self._tts_span.set_attribute("tts.input.character_count", len(self._current_tts_text))
        self._tts_span.set_attribute("tts.input.word_count", len(self._current_tts_text.split()))
        self._tts_span.set_attribute("tts.voice", "nova")
    
    def _handle_tts_end(self):
        """Handle TTS end."""
        if self._tts_span:
            latency_ms = (time.time() - self._current_tts_start) * 1000
            
            # Output
            self._tts_span.set_attribute("tts.output.type", "audio")
            self._tts_span.set_attribute("tts.latency_ms", latency_ms)
            
            self._tts_span.end()
            self._tts_span = None
            
            logger.debug(f"Traced TTS: {latency_ms:.0f}ms, {len(self._current_tts_text)} chars")
            
            # Reset
            self._current_tts_start = 0
            self._current_tts_text = ""


def create_detailed_tracing_observer() -> DetailedTracingObserver:
    """Factory function to create the tracing observer."""
    return DetailedTracingObserver()