Guardrails AI Streaming PII Validation

guardrails_streaming_pii.py

import time

from rich.live import Live
from rich.table import Table

import openai
from guardrails.hub import DetectPII
from guardrails import Guard


guard = Guard().use(DetectPII, pii_entities=['PERSON', 'PHONE_NUMBER'], on_fail="fix")

prompt = """
Write a 5 sentence paragraph where you mention your name and phone number.
"""

msg_history = [
    {"role": "user", "content": prompt},
]


def generate_table(op=None) -> Table:
    """Make a new table."""
    table = Table()

    table.add_column("Raw LLM Output", width=80)
    table.add_column("Validated LLM Output", width=80)

    if op is None:
        return table

    raw, validated = op.raw_llm_output, op.validated_output
    table.add_row(f"[red]{raw}", f"[green]{validated}")
    return table


# Wrap the OpenAI API call with the `guard` object
fragment_generator = guard(
    openai.chat.completions.create,
    msg_history=msg_history,
    max_tokens=1000,
    temperature=0.1,
    stream=True,
)

with Live(generate_table(), refresh_per_second=60) as live:
    for op in fragment_generator:
        time.sleep(0.01)
        live.update(generate_table(op))