Protect Haystack pipelines with OGuardAI so PII is tokenized before the generator and restored after it

The OGuardAI Python SDK ships an optional Haystack integration: two pipeline components that wrap the framework-agnostic transform / rehydrate client. guarded_transform_component tokenizes the input text into safe_text (plus a sealed session_state) so any downstream component only ever sees tokens. guarded_rehydrate_component restores the model output using that same session_state. You wire the transform component before the generator and the rehydrate component after it.

Unlike a single wrapper function, the Haystack adapter is a pair of ordinary Haystack components. That gives you two ways to restore, with different trust properties. Read the trust boundary below before you pick one.

Install

The integration is an optional extra, so Haystack is only pulled in when you use it. The framework is imported lazily, so importing guardai_sdk never requires Haystack.

pip install "oguardai-sdk[haystack]"

Trust boundary

The core guarantee is the same in both usage patterns: the generator, and every component wired between the transform component and the restore point, only ever sees the tokenized safe_text. The model receives Beantworte diese Anfrage: {{person:p_001:ad4f97591c16}} ..., never the raw name, email, or IBAN. Raw PII exists only inside the OGuardAI runtime, except for values an operator explicitly lists in a policy whitelist (the shipped strict policies have an empty whitelist).

Two honest caveats specific to Haystack:

The transform component receives the raw text at its own input socket, because something has to tokenize it. So a pipeline-wide tracer that instruments every component observes the raw input at the transform component's boundary.
The rehydrate component runs inside the pipeline and produces the restored PII at its output socket. A pipeline-wide tracer therefore observes the restored values at the rehydrate component's boundary.

The two patterns below differ only in where restoration happens relative to the pipeline.

Fail closed: guarded_transform_component refuses to emit text if the server returns no session_state. Its run raises ValidationError (GUARDAI_INVALID_INPUT) before any safe_text leaves the component, so a generator can never run on text that cannot be restored. On the restore side, an invalid or tampered session_state makes the underlying client raise rather than return raw text, so a broken session stops the pipeline instead of leaking.

Recommended: rehydrate outside the pipeline

Put guarded_transform_component and the generator in the pipeline, end the pipeline at the generator, and call client.rehydrate outside it. The restored PII is then never a pipeline value, so no Haystack tracer, logger, or intermediate component ever observes a restored name or IBAN. This is the path to prefer when a PII-exporting tracer is attached to the pipeline.

from haystack import Pipeline
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator

from guardai_sdk import OGuardAIClient
from guardai_sdk.integrations.haystack import guarded_transform_component

client = OGuardAIClient(base_url="http://localhost:3000")

pipe = Pipeline()
pipe.add_component("guard_in", guarded_transform_component(client, policy="german-support", language="de"))
pipe.add_component("prompt", PromptBuilder(template="Beantworte diese Mieteranfrage kurz und formal:\n{{ safe_text }}"))
pipe.add_component("llm", OpenAIGenerator(model="gpt-4o-mini"))

# Token-only text flows into the generator; the sealed session bypasses it.
pipe.connect("guard_in.safe_text", "prompt.safe_text")
pipe.connect("prompt.prompt", "llm.prompt")

result = pipe.run({"guard_in": {"text": "Bitte um Kautionsrueckzahlung, Katharina Brandl, k.brandl@example.de"}})

# The generator's reply still holds tokens. Restore it OUTSIDE the pipeline.
draft = result["llm"]["replies"][0]
session_state = result["guard_in"]["session_state"]
restored = client.rehydrate(draft, session_state=session_state, output_channel="customer_email", restore_mode="full")
print(restored.restored_text)

guarded_transform_component(client, *, policy=None, language=None) returns a component whose run takes text: str and emits safe_text: str and session_state: str. Because session_state is left unconnected here, Haystack surfaces it as a pipeline output alongside the generator's replies.

In-pipeline rehydration

If no PII-exporting tracer is attached, you can keep the whole flow in one pipeline by adding guarded_rehydrate_component after the generator. This is the most compact wiring, at the cost of the restored PII being a pipeline value (see the tracer caveat above).

from haystack import Pipeline
from haystack.components.builders import PromptBuilder
from haystack.components.converters import OutputAdapter
from haystack.components.generators import OpenAIGenerator

from guardai_sdk import OGuardAIClient
from guardai_sdk.integrations.haystack import (
    guarded_rehydrate_component,
    guarded_transform_component,
)

client = OGuardAIClient(base_url="http://localhost:3000")

pipe = Pipeline()
pipe.add_component("guard_in", guarded_transform_component(client, policy="german-support", language="de"))
pipe.add_component("prompt", PromptBuilder(template="Beantworte diese Mieteranfrage kurz und formal:\n{{ safe_text }}"))
pipe.add_component("llm", OpenAIGenerator(model="gpt-4o-mini"))
# The generator emits replies: list[str]; the rehydrate component takes a single text: str.
pipe.add_component("pick", OutputAdapter(template="{{ replies[0] }}", output_type=str))
pipe.add_component("guard_out", guarded_rehydrate_component(client, output_channel="customer_email", restore_mode="full"))

pipe.connect("guard_in.safe_text", "prompt.safe_text")
pipe.connect("guard_in.session_state", "guard_out.session_state")
pipe.connect("prompt.prompt", "llm.prompt")
pipe.connect("llm.replies", "pick.replies")
pipe.connect("pick.output", "guard_out.text")

result = pipe.run({"guard_in": {"text": "Bitte um Kautionsrueckzahlung, Katharina Brandl, k.brandl@example.de"}})
print(result["guard_out"]["restored_text"])

guarded_rehydrate_component(client, *, output_channel="user_output", restore_mode="full") returns a component whose run takes text: str and session_state: str and emits restored_text: str. The output_channel decides what each audience is allowed to see, and restore_mode selects one of full, partial, masked, formatted, abstract, none. Channel rules may only tighten restoration, never widen it.

What the generator and tracer see

The transform component surfaces only safe_text and the sealed session_state. It does not surface entity metadata, so the generator sees the tokens and nothing else:

Beantworte diese Mieteranfrage kurz und formal:
Bitte um Kautionsrueckzahlung, {{person:p_001:ad4f97591c16}}, {{email:e_001:1bcaef1a4aff}}

The real name and email never enter the prompt, the generator, or any component between the two guards. Restoration happens either outside the pipeline (recommended pattern) or in the final component (in-pipeline pattern). The honorific and gender-correct form of a German name are resolved server-side from the sealed session at restore time, not by anything the generator is told.

How it works

guarded_transform_component.run(text) calls /v1/transform, replacing PII with semantic tokens, and emits safe_text plus the sealed session_state. If the server returns no session_state, it raises ValidationError instead of emitting text.
The generator, and any component wired on safe_text, run on tokens only.
Restoration calls /v1/rehydrate with the model output and the sealed session_state, restoring the real values per the output channel and restore mode. In the recommended pattern this is a client.rehydrate call outside the pipeline; in the compact pattern it is guarded_rehydrate_component as the last component.

If the transform fails, no safe_text is emitted and the generator never runs (fail closed).

Haystack