praxis/server/assist/guardrail_processor.py

"""LiveAssistGuardrailProcessor — in-loop Pipecat frame processor (D-060 layer 2,
REQ-IDEATE-02, TASK-05-02, REQ-IDEATE-09).

A Pipecat FrameProcessor inserted between `llm` and `tts` in the assist pipeline.
Runs the LiveAssistGuardrail.check() on each LLM response before TTS:
  1. Accumulates TextFrame chunks into the full LLM response.
  2. On LLMFullResponseEndFrame: runs guardrail.check() on the accumulated text.
  3. If allowed → pass the text through to TTS. Log the verdict.
  4. If blocked + retry-eligible → inject RETRY_INSTRUCTION, re-run the LLM.
     If the retry also blocks → CANNED_FALLBACK. Log both verdicts.
  5. If blocked + hard violation → CANNED_FALLBACK immediately (no retry).
  6. Increment session.guardrail_block_count on every block.

REQ-IDEATE-09 (incremental audit-log write): the processor writes the partial
turn (ASR transcript) on TranscriptionFrame, before the LLM response. On
LLMFullResponseEndFrame, it updates the turn with the LLM response + verdict.
This ensures abrupt termination (battery death, power loss mid-turn) still
leaves an audit trail.
"""

from __future__ import annotations

import logging
from typing import Any

from pipecat.frames.frames import (
    Frame,
    LLMFullResponseEndFrame,
    TextFrame,
    TranscriptionFrame,
)
from pipecat.processors.frame_processor import FrameProcessor

from server.guardrails.live_assist import (
    CANNED_FALLBACK,
    RETRY_ELIGIBLE_CATEGORIES,
    RETRY_INSTRUCTION,
    LiveAssistGuardrail,
)
from server.services.base import GuardrailContext

log = logging.getLogger(__name__)


class LiveAssistGuardrailProcessor(FrameProcessor):
    """In-loop guardrail processor (post-LLM, pre-TTS — D-060 layer 2).

    Args:
        guardrail: the LiveAssistGuardrail instance.
        session: the AssistSession (for logging verdicts + block count).
        llm_context: the LLMContext (for injecting retry messages — G-049).
    """

    def __init__(
        self,
        guardrail: LiveAssistGuardrail,
        session: Any | None = None,
        llm_context: Any | None = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.guardrail = guardrail
        self.session = session
        self.llm_context = llm_context
        self._accumulated_text: str = ""
        self._retry_used: bool = False
        self._partial_turn_seq: int | None = None

    async def process_frame(self, frame: Frame, direction) -> None:
        # REQ-IDEATE-09: write the partial turn (ASR) before the LLM response.
        if isinstance(frame, TranscriptionFrame):
            if self.session is not None and frame.text:
                try:
                    self._partial_turn_seq = await self.session.log_assist_turn_partial(
                        frame.text
                    )
                except Exception:
                    log.exception("incremental audit-log: partial turn write failed")
            await self.push_frame(frame, direction)
            return

        # BUFFER LLM text chunks: do NOT push to TTS yet. The guardrail check
        # runs on LLMFullResponseEndFrame (after the full LLM response). Only
        # the allowed text (or CANNED_FALLBACK) is pushed to TTS. This is
        # REQ-ASSIST-03 — the guardrail MUST prevent direct-answer text from
        # reaching the learner's ear before the check completes. Streaming the
        # blocked text through to TTS would defeat the guardrail's purpose
        # (the learner would hear + parrot the direct answer before the canned
        # fallback plays). The latency cost of buffering (~200-500ms for 1-3
        # sentences) is acceptable for safety; the C-8 pilot tolerance (D-072)
        # is flagged for v0.6 hardening if the added latency pushes p95 >650ms.
        if isinstance(frame, TextFrame):
            self._accumulated_text += frame.text
            return

        # On LLM full response end: run the guardrail check on the full text.
        if isinstance(frame, LLMFullResponseEndFrame):
            response_text = self._accumulated_text
            verdict = await self.guardrail.check(
                response_text, GuardrailContext(role="assist")
            )

            if verdict.allowed:
                # Allowed → push the buffered text to TTS + log the verdict.
                # (Buffered, not streamed — REQ-ASSIST-03 requires the guardrail
                # check to complete before any text reaches TTS.)
                if response_text:
                    await self.push_frame(TextFrame(text=response_text), direction)
                await self._log_verdict(verdict, response_text)
                await self.push_frame(frame, direction)
                self._accumulated_text = ""
                self._retry_used = False
                return

            # Blocked.
            # NOTE: guardrail_block_count is incremented by
            # session.log_assist_turn_complete() (which checks the verdict).
            # We do NOT increment it here to avoid double-counting.

            if (
                verdict.category in RETRY_ELIGIBLE_CATEGORIES
                and not self._retry_used
                and self.llm_context is not None
            ):
                # Retry-eligible + retry not yet used → inject RETRY_INSTRUCTION.
                # G-049 validated: LLMContext.add_message supports this.
                self._retry_used = True
                try:
                    self.llm_context.add_message(
                        {"role": "system", "content": RETRY_INSTRUCTION}
                    )
                    log.info(
                        "guardrail blocked (retry-eligible, category=%s) — retrying",
                        verdict.category,
                    )
                except Exception:
                    log.exception("retry injection failed — using canned fallback")
                    await self._emit_canned_fallback(frame, direction, verdict, response_text)
                # The LLM will re-run; we reset the accumulator for the retry response.
                self._accumulated_text = ""
                # We do NOT push the LLMFullResponseEndFrame here — the retry
                # response will produce its own. (In a real pipeline the LLM
                # service re-runs on the updated context.)
                return

            # Hard violation OR retry exhausted → CANNED_FALLBACK.
            await self._emit_canned_fallback(frame, direction, verdict, response_text)
            self._accumulated_text = ""
            self._retry_used = False
            return

        # Non-text frames pass through unchanged.
        await self.push_frame(frame, direction)

    async def _emit_canned_fallback(
        self, frame: Frame, direction, verdict: Any, original_text: str
    ) -> None:
        """Replace the blocked response with CANNED_FALLBACK + log the verdict."""
        # Emit a TextFrame with the canned fallback so TTS speaks it.
        await self.push_frame(TextFrame(text=CANNED_FALLBACK), direction)
        await self._log_verdict(verdict, CANNED_FALLBACK)
        # Pass the LLMFullResponseEndFrame through so TTS knows the response is done.
        await self.push_frame(frame, direction)
        log.info(
            "guardrail blocked (category=%s) — canned fallback emitted",
            verdict.category,
        )

    async def _log_verdict(self, verdict: Any, tts_text: str) -> None:
        """Log the guardrail verdict to the session (REQ-IDEATE-09 incremental audit-log)."""
        if self.session is None:
            return
        try:
            verdict_dict = {
                "allowed": verdict.allowed,
                "reason": verdict.reason,
                "category": verdict.category,
                "filtered_text": verdict.filtered_text,
            }
            if self._partial_turn_seq is not None:
                await self.session.log_assist_turn_complete(
                    self._partial_turn_seq,
                    tts_text=tts_text,
                    guardrail_verdict=verdict_dict,
                )
            else:
                # No partial turn was written (e.g., the turn started before the
                # processor was attached) — log a complete turn.
                await self.session.log_assist_turn(
                    asr_text="",
                    tts_text=tts_text,
                    guardrail_verdict=verdict_dict,
                )
        except Exception:
            log.exception("guardrail verdict log failed")


__all__ = ["LiveAssistGuardrailProcessor"]