TypeSafe AI

The flyteplugins-typesafe-ai plugin runs TypeSafe’s System One model (“Jev”) inside Flyte tasks.

Jev does not write text. It answers typed questions — in parallel, in isolation, and with calibrated confidence attached to every answer. The documented property worth building around is that adding questions barely changes the response time, so the right move is to ask many small questions in one call and compose the result in code you can read and change.

The plugin supplies the two things that Flyte needs: a shape for the answers that survives a task boundary, and a way to ask a whole battery at once.

pip install flyteplugins-typesafe-ai

When to use this plugin

You want a fast, model-based I/O guard in front of a generative model — filtering hostile input, or parsing one raw request into clean typed fields for several downstream steps.
Your agent’s control flow is a set of small decisions (which intent, which tool, is there enough to answer) that you would rather branch on directly than fish back out of prose.
You want the routing thresholds in reviewable Python, and the answers to carry calibrated confidence so abstention is a real option.

For the pattern in full — composing verdicts from atomic questions, confidence-gated routing, and typed agent loops — see System one types in the user guide.

For a complete pipeline built on this plugin — three patterns and a measured A/B against a one-shot generative baseline — see Typed decisions for agentic pipelines.

The three answer types

Type	Holds	Useful members
`Choice`	the picked member of your enum, `confidence`, `probabilities`	`.certain(threshold)`, `.runner_up()`
`Score`	the picked rung of your `IntEnum`, the unrounded `position`, `confidence`	`.at_least(rung)`
`Noul`	truthfulness in 0..1	`.at(threshold)`

Choice comes back as the enum member, not a string, and Score keeps both representations on purpose: value is the rung you branch on, position is where on the scale the answer actually landed, which is what you sort and threshold by.

Noul deliberately has no __bool__. if noul: would treat 0.02 and 0.98 alike, and picking the threshold is the part that belongs in reviewable code.

These are plain dataclasses, so pydantic is not required and there is no bespoke type transformer — they reuse Flyte’s built-in dataclass handling. The plugin registers them with the type engine through the standard flyte.plugins.types entry point, so flyte.init() picks them up, and importing the package registers them too.

Setup

The TypeSafe SDK reads the key from TYPESAFE_API_KEY, so mount your secret as that environment variable. There is no helper for this — it is a plain flyte.Secret:

_env.py
                
image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "flyteplugins-typesafe-ai",
)

# The TypeSafe SDK reads the key from TYPESAFE_API_KEY, so mount the secret as
# that env var. `as_env_var` is spelled out on purpose: it is the string you
# will grep for when a task cannot find the key.
env = flyte.TaskEnvironment(
    name="typesafe-ai",
    image=image,
    secrets=[flyte.Secret(key="TYPESAFE_API_KEY", as_env_var="TYPESAFE_API_KEY")],
    resources=flyte.Resources(cpu=1, memory="1Gi"),
)

flyte.Secret derives as_env_var from the key by upper-casing it and swapping - for _, so a secret named TYPESAFE_API_KEY mounts correctly from flyte.Secret(key="TYPESAFE_API_KEY") alone. Spelling as_env_var out is worth the extra words: it is the string you will grep for when a task cannot find the key.

Create the secret once:

flyte create secret TYPESAFE_API_KEY --value <your key>

If the key is missing, the failure happens at the point of use — in the task that actually calls System One — with a message naming the declaration and the CLI command. That is deliberate: a module’s tasks are imported together on the dataplane, so an import-time raise would take down tasks that never touch System One, and a task that merely passes answers along needs no key at all.

Declaring questions

The vocabulary documents itself. An enum’s class docstring is the question and its member docstrings are the criteria, so a documented enum needs nothing at the call site:

triage.py
                
                    
                
            
                
            

                
                    
                
            
class Intent(enum.Enum):
    """What is this customer actually asking for?"""

    REFUND = "refund"
    """they want money back for something already paid for"""
    DELIVERY_STATUS = "delivery status"
    """they want to know where an order is, or when it will arrive"""
    TECHNICAL_ISSUE = "technical issue"
    """something in the product is not working"""
    ACCOUNT_ACCESS = "account access"
    """they cannot get into their account"""
    OTHER = "something else"
    """none of the above fits"""

class Severity(enum.IntEnum):
    """How badly is this customer blocked right now?"""

    NONE = 0
    """no impact; a question or a comment"""
    MINOR = 1
    """inconvenient, but they can carry on"""
    SERIOUS = 2
    """they are blocked and a deadline or payment is involved"""
    BLOCKING = 3
    """they cannot use the product at all, or money is already lost"""

Group the questions into a dataclass. A Noul has no vocabulary to document itself with, so it always carries its own question — and, where a bare question would be ambiguous, its own criteria:

triage.py
                
                    
                
            
                
            

                
                    
                
            
@dataclass
class Triage:
    """One ticket, fourteen typed answers, one request."""

    # --- the three that decide what happens -------------------------------
    # Intent and Severity document themselves above -- class docstring for the
    # question, member docstrings for the criteria -- so these need no metadata.
    intent: Choice[Intent]
    severity: Score[Severity]
    hostile: Noul = field(
        metadata={
            "question": "Is the customer hostile or abusive?",
            "criteria": {"true": "insults, threats or slurs", "false": "civil, even if angry"},
        }
    )
    # --- the eleven a human reviewer wants anyway --------------------------
    has_order_reference: Noul = field(
        metadata={"question": "Does the ticket name a specific order or reference number?"}
    )
    money_at_stake: Noul = field(metadata={"question": "Is a payment, refund or charge involved?"})
    deadline_mentioned: Noul = field(
        metadata={"question": "Does the customer mention a deadline or an event they need this for?"}
    )
    already_contacted: Noul = field(metadata={"question": "Have they contacted support about this before?"})
    asks_for_human: Noul = field(metadata={"question": "Are they explicitly asking for a human agent?"})
    threatens_chargeback: Noul = field(
        metadata={"question": "Do they threaten a chargeback, a review or legal action?"}
    )
    needs_account_change: Noul = field(metadata={"question": "Would resolving this require changing their account?"})
    reports_bug: Noul = field(metadata={"question": "Are they reporting something that looks like a product defect?"})
    mentions_competitor: Noul = field(metadata={"question": "Do they mention leaving for a competitor?"})
    non_english: Noul = field(metadata={"question": "Is the ticket written in a language other than English?"})
    resolvable_now: Noul = field(metadata={"question": "Could a well-informed agent resolve this in a single reply?"})

Override either the question or the criteria in ordinary dataclasses.field metadata, under two keys named after the SDK’s own arguments:

key	meaning
`question`	the instructions for this question
`criteria`	the same shape `typesafe_sdk` takes for that question type

criteria follows the SDK exactly: a mapping keyed by enum member name for a Choice, a positional sequence of rungs for a Score (so the IntEnum must number its rungs 0..n-1 — a gap is rejected with an error that says so), and {"true": ..., "false": ...} for a Noul.

Member docstrings are not stored on the object at runtime, so they are read by parsing the source, the same way pydantic implements attribute docstrings. That makes them best-effort: where the source is not available (a REPL, exec, some frozen deployments) the criterion falls back to the member name rather than failing.

Three ways to ask

ask() takes any of these and compiles them into a single system_one call:

        
    
triage = await ask(Triage, state)                      # a battery dataclass -> Triage
intent = await ask(Choice[Intent], state)              # one question        -> Choice[Intent]
answers = await ask({"intent": Choice[Intent],         # an ad-hoc battery   -> dict
                     "hostile": Noul}, state)

Outside a dataclass there is no field to hang metadata on, so Annotated carries it instead — a mapping, or a bare string when all you have is the question:

answer_types.py
                
                    
                
            
                
            

                
                    
                
            
@env.task
async def moderate(message: str = MESSAGES[1]) -> str:
    """Three standalone questions, each asked on its own terms."""

    # 1. A bare question type, answered in its own durable task. The enum's class
    #    docstring is the question and its member docstrings are the criteria.
    action: Choice[Action] = await classify(message)

    # 2. Annotated, when you want to ask this vocabulary a different question than
    #    the one its docstring states.
    harm: Score[Harm] = await ask(
        Annotated[Score[Harm], {"question": "How much harm would this message do if published as-is?"}],
        {"message": message},
    )

    # 3. A Noul has no vocabulary to document itself with, so it always carries its
    #    own question -- and here, its own criteria too.
    directed: Noul = await ask(
        Annotated[
            Noul,
            {
                "question": "Is this aimed at a specific person?",
                "criteria": {"true": "addressed at an individual", "false": "general or about the product"},
            },
        ],
        {"message": message},
    )

    return (
        f"action={action.value.value} (p={action.confidence:.2f}) "
        f"harm={harm.value.name}@{harm.position:.1f} directed={directed.value:.2f}"
    )

Both forms work on a dataclass field too. If a field has metadata and an Annotated annotation, the field metadata wins — it is the more specific place to say it.

Prefer one call to several. Three separate ask() calls are three round trips, while a battery or a mapping asks everything at once, which is the property the whole design rests on:

answer_types.py
                
                    
                
            
                
            

                
                    
                
            
@env.task
async def moderate_in_one_call(message: str = MESSAGES[2]) -> str:
    """The same three questions, assembled as a mapping -- so they cost one request.

    Three separate `ask()` calls are three round trips. When the questions are known
    together, hand them over together: System One answers them in parallel, and the
    whole point is that the second and third are nearly free.
    """
    answers = await ask(
        {
            "action": Choice[Action],
            "harm": Score[Harm],
            "directed": Annotated[Noul, "Is this aimed at a specific person?"],
        },
        {"message": message},
    )
    return (
        f"action={answers['action'].value.value} "
        f"harm={answers['harm'].value.name} "
        f"directed={answers['directed'].value:.2f}"
    )

Use ask_with_info() when you want the model name, question count, token usage and latency back alongside the answers:

triage.py
                
@env.task
async def triage(ticket: str) -> Triage:
    """One request, fourteen answers. `Triage` crosses the task boundary as a struct."""
    answered, info = await ask_with_info(Triage, {"ticket": ticket})
    print(f"{info.questions} questions, one call: {info.latency_s}s, {info.input_tokens}in/{info.output_tokens}out")
    return answered

The name collision with typesafe_sdk’s own Choice / Score / Noul is deliberate and one-directional: those describe the question, these hold the answer. You write the ones in this package; the plugin builds the SDK’s from your battery.

Branching on the answers

Thresholds live in your code, not in a prompt:

triage.py
                
                    
                
            
                
            

                
                    
                
            
@env.task
async def handle(ticket: str = SAMPLE) -> str:
    """Route the ticket in ordinary code, on calibrated numbers rather than vibes."""
    t = await triage(ticket)

    # Guard first: what must never be auto-answered.
    if t.hostile.at(0.8) or t.threatens_chargeback.at(0.7):
        route = "escalate"
    # Then confidence: act only when the pick is clear.
    elif not t.intent.certain(0.85) or t.severity.at_least(Severity.BLOCKING):
        route = "review"
    else:
        route = "auto"

    # Every facet that came back true -- free to read, because they rode along in the same call.
    fired = sorted(name for name, answer in vars(t).items() if isinstance(answer, Noul) and answer.at(0.5))
    return (
        f"route={route} intent={t.intent.value.value} (p={t.intent.confidence:.2f}) "
        f"severity={t.severity.value.name}@{t.severity.position:.1f} fired={fired}"
    )

Answers as task inputs and outputs

Choice, Score and Noul are plain dataclasses, so Flyte carries them with nothing registered — including on their own, not just inside a battery:

answer_types.py
                
@env.task
async def classify(message: str) -> Choice[Action]:
    """An answer type is a perfectly good task output.

    `Choice[Action]` is a parameterized dataclass, so it crosses the boundary as a
    struct -- the picked member, its confidence and the whole distribution -- and
    the caller gets a real `Choice` back, not a dict.
    """
    return await ask(Choice[Action], {"message": message})

They also get the dict coercion every dataclass input gets, so a caller may pass {"value": "refund", "confidence": 0.91} where a Choice[Intent] is expected, and omitted fields fall back to their defaults. Note that an enum nested in a dataclass is spelled by its value ("refund"), not its name — that is mashumaro’s convention for dataclass fields, and it differs from the name-based spelling Flyte uses for a bare enum at the top level.

Inference is never implicit. A dict arriving for a Choice[Intent] is the serialized answer, never a state to go ask about — the two are indistinguishable by shape, and replay depends on the serialized reading winning. If you want the interface to say “System One produces this”, make it a task: you get durability, caching and retries, and the call stays visible in the run graph.

Two kinds of parallelism

Flyte runs one durable task per item across the cluster; inside each task, System One answers the whole battery in a single request. The second one is why per-item latency barely moves when you add questions.

fanout.py
                
                    
                
            
                
            

                
                    
                
            
@env.task
async def triage_one(ticket: str) -> tuple[Triage, CallInfo]:
    return await ask_with_info(Triage, {"ticket": ticket})

@env.task(report=True)
async def backlog() -> str:
    """One durable task per ticket; one System One call inside each."""
    results = await asyncio.gather(*[triage_one(t) for t in BACKLOG])

    rows = []
    for ticket, (t, info) in zip(BACKLOG, results):
        fired = sum(1 for answer in vars(t).values() if isinstance(answer, Noul) and answer.at(0.5))
        urgent = t.severity.at_least(Severity.SERIOUS)
        rows.append(
            f"<tr><td>{ticket[:54]}…</td><td>{t.intent.value.value}</td>"
            f"<td>{t.intent.confidence:.2f}</td><td>{t.severity.value.name}</td>"
            f"<td>{'yes' if urgent else 'no'}</td><td>{fired}</td>"
            f"<td>{info.questions}</td><td>{info.latency_s:.2f}s</td></tr>"
        )

    calls = [info for _, info in results]
    total_q = sum(c.questions for c in calls)
    slowest = max(c.latency_s for c in calls)
    tokens = sum(c.input_tokens + c.output_tokens for c in calls)

    flyte.report.get_tab("Backlog").log(
        f"<p>{len(BACKLOG)} tickets, <b>{total_q}</b> typed answers, "
        f"<b>{len(calls)}</b> System One calls — one per ticket, not one per question. "
        f"Slowest call {slowest:.2f}s; {tokens:,} tokens total.</p>"
        "<table><thead><tr><th>ticket</th><th>intent</th><th>conf</th><th>severity</th>"
        "<th>urgent</th><th>facets fired</th><th>questions</th><th>latency</th></tr></thead>"
        f"<tbody>{''.join(rows)}</tbody></table>"
    )
    await flyte.report.flush.aio()

    urgent = [t for t, _ in results if t.severity.at_least(Severity.SERIOUS)]
    hostile = [t for t, _ in results if t.hostile.at(0.8)]
    return (
        f"{len(BACKLOG)} tickets -> {total_q} typed answers in {len(calls)} calls; "
        f"{len(urgent)} urgent, {len(hostile)} hostile; {tokens:,} tokens; slowest call {slowest:.2f}s"
    )

A note on `IntEnum`

Score takes an IntEnum because a rubric is ordered. Flyte’s enum transformer accepts IntEnum as well as string-valued enums, serialized by member name like every other enum, so a bare severity: Severity also works as a task input or output. Flag and IntFlag remain unsupported, with a message explaining why: a composite member like READ|WRITE has a name but cannot be looked up by it, so it cannot come back.

A parameterized dataclass such as Choice[Intent] is a generic alias, which dataclasses.is_dataclass() rejects. The type engine resolves the alias to its origin for structural checks while keeping the alias itself for decoding, which is what binds the type variable — so Choice[Intent] works as a task type directly.

Next steps

System one types: the guard, composition and confidence-gating patterns, end to end.
Typed decisions for agentic pipelines: three pipelines built on this plugin, plus a measured A/B against a one-shot generative baseline.
Agent frameworks: running a framework’s own agent loop as durable Flyte tasks.