LlamaIndex

The OGuardAI Python SDK ships an optional LlamaIndex integration for RAG. It gives you two complementary seams: a node postprocessor that tokenizes retrieved node text so the response-synthesis model only ever sees safe_text, and a query wrapper that transforms the query and restores the answer outside the query engine. Both chain a single sealed session, so a value that appears in both the query and a retrieved chunk gets the same token, and the final rehydrate resolves tokens from either source. Each seam is a thin layer over the SDK transform and rehydrate primitives.

Install

The integration is an optional extra, so LlamaIndex is only pulled in when you use it. The adapter imports llama_index lazily, so importing guardai_sdk never requires LlamaIndex.

pip install "oguardai-sdk[llamaindex]"

The two seams

The exported names are guarded_query, guarded_node_postprocessor, and NodeSessionSink.

• guarded_query(query_engine, client, query_str, ...) is the boundary-safe path. It calls /v1/transform on the query and /v1/rehydrate on the answer outside the query engine. No LlamaIndex callback or tracer observes the raw query or the restored answer. Prefer this whenever a PII-exporting tracer or observability hook is attached to your engine.
• guarded_node_postprocessor(client, ...) returns a BaseNodePostprocessor that replaces each retrieved node's text with tokenized safe_text before the response-synthesis model runs. It guarantees the synthesis model sees only tokens. It runs inside the query engine, after retrieval, so a tracer attached at or above the retrieval stage can still observe the raw retrieved chunk text before tokenization. To keep raw PII out of the vector store and its retrieval callbacks entirely, tokenize documents at ingestion (transform them before you index) so the store holds only safe_text in the first place.

Pair them through a shared NodeSessionSink so the query session and the node sessions are the same sealed session. That lets the final rehydrate resolve tokens produced by both the query and the retrieved nodes.

Python

from guardai_sdk import OGuardAIClient
from guardai_sdk.integrations.llamaindex import (
    NodeSessionSink,
    guarded_node_postprocessor,
    guarded_query,
)

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

# One sealed session shared by the query and every retrieved node.
sink = NodeSessionSink()

# Install the guarded postprocessor on your existing query engine so retrieved node
# text is tokenized before the synthesis model sees it. `index` is any LlamaIndex index.
query_engine = index.as_query_engine(
    node_postprocessors=[
        guarded_node_postprocessor(
            client,
            session_sink=sink,
            policy="german-support",
            language="de",
        ),
    ],
)

# Boundary-safe: the query is tokenized and the answer is restored OUTSIDE the engine.
answer = guarded_query(
    query_engine,
    client,
    "Was haben wir Katharina Brandl zur Kautionsrueckzahlung zugesagt?",
    session_sink=sink,
    policy="german-support",
    language="de",
    output_channel="user_output",
)
print(answer)

The synthesis model only ever sees text like Vorgang {{order:o_001:57e07d900ca4}} fuer {{person:p_001:ad4f97591c16}}, and the returned answer has the real values restored per the output channel.

Session chaining

NodeSessionSink is a small dataclass with a single field, session_state: str | None = None. When you pass the same sink to guarded_query and to the postprocessor installed on its engine:

1. guarded_query transforms the query and seeds the query session into the sink.
2. When the engine runs, the postprocessor reads the sink's session as its starting state, transforms each retrieved node onto it, and writes the final combined session back to the sink.
3. guarded_query reads the final combined session from the sink and rehydrates the answer with it, so tokens from both the query and the retrieved nodes resolve.

Because the session chains, the same raw value in the query and in a chunk yields the same token id. That token stability is what makes cross-document RAG answers restore correctly.

Standalone postprocessor

You can use guarded_node_postprocessor on its own, without guarded_query, when you drive retrieval and rehydration yourself. Pass a NodeSessionSink so you can read the sealed session after the engine runs, then rehydrate the answer with the SDK client:

from guardai_sdk import OGuardAIClient
from guardai_sdk.integrations.llamaindex import NodeSessionSink, guarded_node_postprocessor

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

query_engine = index.as_query_engine(
    node_postprocessors=[guarded_node_postprocessor(client, session_sink=sink)],
)

response = query_engine.query("...tokenized query...")
restored = client.rehydrate(
    str(response),
    session_state=sink.session_state,
    output_channel="user_output",
)
print(restored.restored_text)

If the engine returns a response object whose text is not a plain string or a .response attribute, pass response_extractor to guarded_query to pull the answer text out:

answer = guarded_query(
    query_engine,
    client,
    query_str,
    response_extractor=lambda r: r.response.message.content,
)

How it works

1. guarded_query calls /v1/transform on the query, replacing PII with semantic tokens, and seeds the sealed session into the shared sink.
2. query_engine.query runs on the tokenized query. The guarded postprocessor tokenizes each retrieved node's text, chaining the sealed session across nodes, so the synthesis model only ever sees safe_text.
3. guarded_query calls /v1/rehydrate on the engine's answer with the final combined session, restoring real values per output_channel and restore_mode.

output_channel defaults to user_output and restore_mode defaults to full. Channel rules may only tighten restoration, never widen it, so a value the policy caps (for example an IBAN masked under the german-support policy) stays capped even under restore_mode="full".

What the model and tracer see

• The query engine receives only the tokenized query, never the raw query text.
• Each retrieved node's text is replaced with tokenized safe_text before the response-synthesis model runs, so the model sees tokens and safe metadata, never a raw name, email, or IBAN.
• The engine's answer still holds tokens. Rehydration happens after query_engine.query returns, outside the engine, so no LlamaIndex callback or tracer attached to the guarded_query boundary observes a raw or restored value.

The one place raw retrieved text still exists is between the retriever and the postprocessor, inside the engine. If a retrieval-stage tracer is a concern, tokenize documents at ingestion so the vector store never holds raw PII. Person and company detection needs the NER detector; with the builtin regex detectors alone, names are not tokenized and would remain visible to the synthesis step, so run the full stack when you need name-level protection.

Fail closed

Both seams fail closed before any raw text is exposed. If a transform returns no session_state, guarded_query raises ValidationError and never runs the query, and guarded_node_postprocessor raises ValidationError and never exposes the node's text. A corrupted or expired session makes rehydrate raise a typed error (GUARDAI_SESSION_EXPIRED or GUARDAI_INVALID_INPUT) rather than return unresolved or raw text.