Streaming Transform and Rehydrate (SSE)

customer message ─▶ POST /v1/transform/stream ─▶ safe chunks stream to the agent UI
                                                  terminal event carries session_state
model reply (tokens) ─▶ POST /v1/rehydrate/stream ─▶ restored chunks stream to the customer

data: {"chunk": "Hi {{person:p_001:...}}, your refund for order {{order:o_001:...}} ...", "complete": false}
data: {"complete": true, "session_id": "...", "session_state": "...", "stats": {...}}

pip install -r requirements.txt

# Start the OGuardAI runtime. Person detection needs the NER sidecar (see the note below). The bundled
# stack requires these four secrets, and its client key ships least-privilege (channel user_output
# only), so widen it to the channels this example restores on (customer_email, log_safe).
export GUARDAI_SESSION_SECRET=$(openssl rand -base64 32)
export GUARDAI_API_KEY_1=$(openssl rand -hex 32)
export GUARDAI_DETECTOR_API_KEY=$(openssl rand -hex 32)
export GUARDAI_REDIS_PASSWORD=$(openssl rand -hex 32)
export GUARDAI_API_KEY_1_ALLOWED_OUTPUT_CHANNELS=user_output,customer_email,log_safe
docker compose -f deploy/docker/docker-compose.yml up -d

export GUARDAI_URL=http://localhost:3000
export GUARDAI_API_KEY="$GUARDAI_API_KEY_1"   # the bundled stack runs api_key auth
python main.py

# The streaming client uses requests for line-oriented SSE reading.
requests>=2.31

Subject: Refund status for order ORD-2026-7781

Hi, this is Max Mustermann. I am still waiting on the refund for order
ORD-2026-7781. Please confirm it was sent to my account IBAN
GB29 NWBK 6016 1331 9268 19 and email me an update at max.mustermann@example.com.
You can also reach me on +1 (206) 555-0147.

"""OGuardAI streaming walkthrough, company-grade (SSE over raw HTTP, Python).

Scenario
--------
A support desk streams protected text so its chat UI can render token by token, the way a customer
expects a modern assistant to type. Two streams do the work:

    POST /v1/transform/stream   the inbound customer message is tokenized and streamed out, so the
                                agent UI shows the protected text as it is produced. The terminal
                                event carries the sealed session_state used to restore later.
    POST /v1/rehydrate/stream   the model's token-only reply is restored and streamed back, so the
                                customer sees their real order, name, and contact details typed out
                                live, governed by the output channel.

Both directions are Server-Sent Events: incremental ``{"chunk": ..., "complete": false}`` events,
then a terminal ``{"complete": true, ...}`` event. This example parses that stream correctly (never
emitting a partial ``{{type:id:cap}}`` token to the reader is the runtime's job; the client just has to
handle the event shape), captures the session across the two calls, restores on two channels, and
handles the error and fail-closed cases a real integration must.

No LLM provider is required: the "model reply" is a deterministic stub that echoes the tokens the
model would have seen. The one-line change to a real provider is called out in the README. All PII
in the fixture is fictional.

Prerequisites
-------------
    pip install -r requirements.txt      # requests
    # Start the OGuardAI runtime (person detection needs the NER sidecar):
    #   GUARDAI_SESSION_SECRET=$(openssl rand -base64 32) \
    #     docker compose -f deploy/docker/docker-compose.yml up -d

Usage
-----
    export GUARDAI_URL=http://localhost:3000
    export GUARDAI_API_KEY="<key>"       # only if the server runs with auth
    python main.py

Policy: ``default`` (full restore on user_output, formatted on customer_email, redacted on log_safe,
IBAN masked everywhere).
"""

from __future__ import annotations

import json
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Iterator

import requests

BASE = (os.environ.get("GUARDAI_URL") or os.environ.get("GUARDAI_BASE_URL") or "http://localhost:3000").rstrip("/")
API_KEY = os.environ.get("GUARDAI_API_KEY", "")
POLICY = "default"
LANGUAGE = "en"
HERE = Path(__file__).resolve().parent
MESSAGE_FILE = HERE / "support_message.txt"

_JSON_HEADERS = {"Content-Type": "application/json"}
_SSE_HEADERS = {**_JSON_HEADERS, "Accept": "text/event-stream"}
if API_KEY:
    _JSON_HEADERS["X-API-Key"] = API_KEY
    _SSE_HEADERS["X-API-Key"] = API_KEY

_TOKEN = re.compile(r"\{\{([a-z_]+):[^}]*\}\}")
RULE = "=" * 78

# Raw values from the fixture that must never survive into the text the model sees. The structured
# values are caught by the built-in regex detectors, so a leak is a hard failure on every lane. The
# name needs the NER detector: a leak is a hard failure when NER is active, and a printed NOTE when
# the runtime runs without it.
NER_SENSITIVE = ["Max Mustermann"]
STRUCTURED_SENSITIVE = [
    "max.mustermann@example.com",
    "+1 (206) 555-0147",
    "GB29 NWBK 6016 1331 9268 19",
    "ORD-2026-7781",
]


class HttpError(RuntimeError):
    """A non-2xx HTTP response from the runtime (e.g. a tampered session at stream open)."""

    def __init__(self, status: int, body: dict) -> None:
        self.status = status
        self.code = body.get("code") or f"HTTP_{status}"
        # Surface only OGuardAI's own structured error string; never the raw upstream body, which
        # can echo the request's PII. Fall back to a generic status-only message.
        error = body.get("error")
        message = error if isinstance(error, str) else f"OGuardAI request failed (HTTP {status})"
        super().__init__(f"HTTP {status} {self.code}: {message}")


class StreamError(RuntimeError):
    """An ``error`` event carried inside the SSE stream (e.g. output blocked mid-stream)."""

    def __init__(self, code: str | None, message: str | None) -> None:
        self.code = code
        super().__init__(f"stream error {code}: {message}")


@dataclass
class TransformStream:
    safe_text: str = ""
    session_state: str = ""
    stats: dict = field(default_factory=dict)


@dataclass
class RehydrateStream:
    restored_text: str = ""
    tokens_resolved: int | None = None
    tokens_unresolved: list[str] = field(default_factory=list)
    stats: dict = field(default_factory=dict)


def _raise_for_status(resp: requests.Response) -> None:
    if resp.status_code >= 400:
        # Parse the body only when it is a JSON error envelope; a non-JSON body can be the raw
        # upstream response and echo the request's PII, so it is dropped rather than surfaced.
        body: dict = {}
        if resp.headers.get("content-type", "").startswith("application/json"):
            try:
                parsed = resp.json()
                body = parsed if isinstance(parsed, dict) else {}
            except ValueError:
                body = {}
        raise HttpError(resp.status_code, body)


def _iter_sse(resp: requests.Response) -> Iterator[dict]:
    """Yield each parsed ``data:`` JSON event from an SSE response, stopping at ``[DONE]``."""
    for raw in resp.iter_lines():
        if not raw or not raw.startswith(b"data:"):
            continue
        payload = raw[len(b"data:") :].strip()
        if payload == b"[DONE]":
            return
        yield json.loads(payload)


def _classify(event: dict) -> str:
    """SSE events are untyped JSON; classify by field. Order matters: an error event wins."""
    if event.get("error") or event.get("event") == "error":
        return "error"
    if event.get("complete") or event.get("event") == "complete":
        return "complete"
    return "chunk"


def stream_transform(text: str, *, on_chunk: Callable[[str], None]) -> TransformStream:
    """Stream-tokenize ``text``. ``on_chunk`` receives each safe chunk for live rendering."""
    body = {"input": text, "policy": POLICY, "language": LANGUAGE}
    out = TransformStream()
    with requests.post(f"{BASE}/v1/transform/stream", json=body, headers=_SSE_HEADERS, stream=True, timeout=60) as resp:
        _raise_for_status(resp)
        for event in _iter_sse(resp):
            kind = _classify(event)
            if kind == "error":
                raise StreamError(event.get("code"), event.get("error"))
            if kind == "complete":
                out.session_state = event.get("session_state", "") or ""
                out.stats = event.get("stats", {}) or {}
                continue
            chunk = event.get("chunk") or ""
            out.safe_text += chunk
            on_chunk(chunk)
    return out


def stream_rehydrate(output: str, session_state: str, *, channel: str, on_chunk: Callable[[str], None]) -> RehydrateStream:
    """Stream-restore ``output`` under ``session_state`` on ``channel``, chunk by chunk."""
    body = {"output": output, "session_state": session_state, "output_channel": channel}
    out = RehydrateStream()
    with requests.post(f"{BASE}/v1/rehydrate/stream", json=body, headers=_SSE_HEADERS, stream=True, timeout=60) as resp:
        _raise_for_status(resp)
        for event in _iter_sse(resp):
            kind = _classify(event)
            if kind == "error":
                raise StreamError(event.get("code"), event.get("error"))
            if kind == "complete":
                out.tokens_resolved = event.get("tokens_resolved")
                out.tokens_unresolved = event.get("tokens_unresolved") or []
                out.stats = event.get("stats", {}) or {}
                continue
            chunk = event.get("chunk") or ""
            out.restored_text += chunk
            on_chunk(chunk)
    return out


def _tokens_by_type(safe_text: str) -> dict[str, str]:
    return {m.group(1): m.group(0) for m in _TOKEN.finditer(safe_text)}


def simulated_model_reply(safe_text: str) -> str:
    """Deterministic stand-in for a streaming LLM: a reply that reuses only the tokens it was given.

    It never invents a name, order, or contact detail, so every token resolves on rehydrate. A real
    integration accumulates the provider's streamed completion into the finished reply text, then
    passes that string to ``stream_rehydrate``, which restores it and re-streams the restored text.
    """
    t = _tokens_by_type(safe_text)
    return (
        f"Hi {t.get('person', 'there')}, your refund for order {t.get('order', 'your order')} has been "
        f"sent to the account {t.get('iban', 'on file')}. A written confirmation is on its way to "
        f"{t.get('email', 'your email')}."
    )


def fetch_ner_active() -> bool:
    """Read ner_active from GET /v1/capabilities: the CONFIGURED detector intent, not sidecar
    liveness. It is the one signal that decides whether a person name is expected to tokenize. When
    it cannot be read, or the field is anything but an explicit false, treat NER as available so an
    untokenized name fails rather than passing silently."""
    try:
        resp = requests.get(f"{BASE}/v1/capabilities", headers=_JSON_HEADERS, timeout=10)
        resp.raise_for_status()
        data = resp.json()
    except (requests.RequestException, ValueError):
        return True
    if not isinstance(data, dict):
        return True
    val = data.get("ner_active", True)
    return val is not False


def check_crossing(text: str, where: str, *, ner_available: bool, name_notes: list[str]) -> None:
    """Fail closed on anything sensitive that crossed to the model. Structured PII is caught by the
    built-in regex detectors, so a leak is a hard failure on every lane. A person name is a hard
    failure when NER is active; when the runtime runs without NER it cannot tokenize the name, so the
    leak is reported as a NOTE and collected for the closing summary."""
    structured = [value for value in STRUCTURED_SENSITIVE if value in text]
    if structured:
        raise SystemExit(f"FAIL: structured PII reached {where}: {structured}")
    names = [value for value in NER_SENSITIVE if value in text]
    if not names:
        return
    if ner_available:
        raise SystemExit(f"FAIL: a name reached {where} while NER is active: {names}")
    for name in names:
        print(f"      NOTE: {name} reached {where}; the runtime has no NER to tokenize the name")
        name_notes.append(name)


def live(label: str) -> None:
    print(f"\n{label}")
    print("  ", end="", flush=True)


def main() -> int:
    message = MESSAGE_FILE.read_text(encoding="utf-8").strip()
    print(RULE)
    print(f"OGuardAI streaming walkthrough   runtime: {BASE}   policy: {POLICY}")
    print(RULE)

    ner_available = fetch_ner_active()
    name_notes: list[str] = []

    try:
        # 1. Stream the outbound (masking) direction. The agent UI renders tokenized text live.
        live("[1] transform/stream  (what the model will see, streamed as it is produced):")
        transformed = stream_transform(message, on_chunk=lambda c: print(c, end="", flush=True))
        print(
            f"\n      [complete] entities={transformed.stats.get('entities_transformed')} "
            f"session captured={bool(transformed.session_state)}"
        )
        check_crossing(
            transformed.safe_text,
            "the model input stream",
            ner_available=ner_available,
            name_notes=name_notes,
        )

        if not transformed.session_state:
            # Fail closed: without a session there is nothing to restore against. Never guess.
            print("      no session_state returned; refusing to continue (fail closed)")
            return 1

        # 2. The model's streamed reply, in tokens only. In production this is the provider stream.
        reply = simulated_model_reply(transformed.safe_text)
        print(f"\n[2] the model replies in tokens only:\n      {reply}")
        check_crossing(
            reply, "the model reply stream", ner_available=ner_available, name_notes=name_notes
        )

        # 3. Stream the inbound (restoring) direction for the customer. Same session, full restore.
        live("[3] rehydrate/stream  channel=customer_email  (what the customer sees, typed live):")
        customer = stream_rehydrate(
            reply, transformed.session_state, channel="customer_email", on_chunk=lambda c: print(c, end="", flush=True)
        )
        print(f"\n      [complete] tokens_resolved={customer.tokens_resolved} unresolved={customer.tokens_unresolved}")

        # 4. The same reply, streamed on the audit channel: fully redacted, safe to persist.
        live("[4] rehydrate/stream  channel=log_safe  (the audit copy, no PII):")
        stream_rehydrate(
            reply, transformed.session_state, channel="log_safe", on_chunk=lambda c: print(c, end="", flush=True)
        )
        print()

        # 5. Fail closed: a tampered session must be rejected at stream open, not half-restored.
        print("\n[5] fail-closed: a tampered session is rejected at stream open")
        tampered = transformed.session_state[:-8] + "deadbeef"
        try:
            stream_rehydrate(reply, tampered, channel="customer_email", on_chunk=lambda _c: None)
            print("      FAIL, tampered session was accepted")
            return 1
        except HttpError as exc:
            print(f"      rejected: {exc} (no raw value streamed)")
        except StreamError as exc:
            print(f"      rejected mid-stream: {exc} (no raw value streamed)")

    except HttpError as exc:
        print(f"\nruntime error: {exc}")
        print(f"is the OGuardAI server running at {BASE}? start it with docker compose (see the README).")
        return 1
    except requests.ConnectionError:
        print(f"\ncannot reach the OGuardAI server at {BASE}. start it first (see the README).")
        return 1

    print("\n" + RULE)
    if name_notes:
        reached = ", ".join(dict.fromkeys(name_notes))
        print("Both directions streamed correctly and the session carried across the two calls.")
        print(f"Structured PII stayed inside the runtime; the name {reached} reached the model because")
        print("this runtime runs without NER. Enable the NER detector to tokenize person names too.")
    else:
        print("Both directions streamed correctly; the session carried across the two calls; the customer")
        print("saw full values, the audit log saw none, and a tampered session was refused.")
    print(RULE)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

[1] transform/stream  (what the model will see, streamed as it is produced):
  Subject: Refund status for order {{order:o_001:...}} Hi, this is {{person:p_001:...}}. ...
  ... account IBAN {{iban:ib_001:...}} ... update at {{email:e_001:...}} ... on {{phone:ph_001:...}}.
      [complete] entities=5 session captured=True

[2] the model replies in tokens only:
      Hi {{person:p_001:...}}, your refund for order {{order:o_001:...}} has been sent ...

[3] rehydrate/stream  channel=customer_email  (what the customer sees, typed live):
  Hi Mr. Max Mustermann, your refund for order ORD-2026-7781 has been sent to the account
  IBAN **** (masked). A written confirmation is on its way to max.mustermann@example.com.
      [complete] tokens_resolved=4 unresolved=[]

[4] rehydrate/stream  channel=log_safe  (the audit copy, no PII):
  Hi [REDACTED], your refund for order [REDACTED] has been sent ...

[5] fail-closed: a tampered session is rejected at stream open
      rejected: HTTP 400 GUARDAI_INVALID_INPUT: ... (no raw value streamed)