flyteplugins-typesafe-ai 0.0.0a0__py3-none-any.whl

flyteplugins/typesafe_ai/__init__.py

@@ -0,0 +1,105 @@
+"""Run TypeSafe's System One model (Jev) inside durable Flyte tasks.
+
+Jev answers typed questions in parallel instead of generating text, and returns
+calibrated confidence with every answer. This plugin gives those answers a shape
+that crosses a Flyte task boundary — `Choice`, `Score` and `Noul` — and a way to
+ask a whole battery of them in a single request.
+
+```python
+import enum
+from dataclasses import dataclass, field
+
+import flyte
+from flyteplugins.typesafe_ai import Choice, Noul, Score, ask
+
+env = flyte.TaskEnvironment(
+    "triage",
+    secrets=[flyte.Secret(key="TYPESAFE_API_KEY", as_env_var="TYPESAFE_API_KEY")],
+)
+
+
+class Intent(enum.Enum):
+    '''Which intent best fits the ticket?'''
+
+    REFUND = "refund"
+    '''they want money back'''
+    DELIVERY = "delivery"
+    '''they are asking where their order is'''
+
+
+class Severity(enum.IntEnum):
+    '''How badly is this customer affected?'''
+
+    NONE = 0
+    '''no impact; a question or a comment'''
+    MINOR = 1
+    '''inconvenient, but they can carry on'''
+    SERIOUS = 2
+    '''they are blocked'''
+
+
+@dataclass
+class Triage:
+    # The enums above document themselves, so these fields need no metadata at all.
+    intent: Choice[Intent]
+    severity: Score[Severity]
+    # A Noul has no vocabulary to document itself with, so it needs both.
+    hostile: Noul = field(
+        metadata={"question": "Is the customer hostile?", "criteria": {"true": "insults or threats", "false": "civil"}}
+    )
+
+
+@env.task
+async def triage(ticket: str) -> Triage:
+    return await ask(Triage, {"ticket": ticket})
+```
+"""
+
+import functools
+
+from flyte.types import TypeEngine
+
+from ._ask import BatteryError, ask, ask_with_info, compile_questions
+from ._client import API_KEY_ENV, MissingAPIKey, client
+from ._types import CRITERIA_KEY, QUESTION_KEY, CallInfo, Choice, Noul, Score
+
+
+@functools.lru_cache(maxsize=None)
+def register_typesafe_ai_types():
+    """Register Choice, Score and Noul with the Flyte type engine.
+
+    Called automatically via the `flyte.plugins.types` entry point when
+    `flyte.init()` runs with `load_plugin_type_transformers=True` (the default).
+
+    The three answer types are plain dataclasses, so they reuse the existing
+    `DataclassTransformer` rather than introducing one of their own. Registering
+    them anyway makes the resolution explicit: a parameterized `Choice[Intent]`
+    is matched through its origin instead of falling through to the type engine's
+    last-resort dataclass branch.
+    """
+    from flyte.types._type_engine import DataclassTransformer
+
+    transformer = DataclassTransformer()
+    for answer_type in (Choice, Score, Noul):
+        TypeEngine.register_additional_type(transformer, answer_type)
+
+
+# Also register at import time, so the types work without flyte.init() -- a unit
+# test that only round-trips a battery never calls it.
+register_typesafe_ai_types()
+
+__all__ = [
+    "API_KEY_ENV",
+    "CRITERIA_KEY",
+    "QUESTION_KEY",
+    "BatteryError",
+    "CallInfo",
+    "Choice",
+    "MissingAPIKey",
+    "Noul",
+    "Score",
+    "ask",
+    "ask_with_info",
+    "client",
+    "compile_questions",
+]

flyteplugins/typesafe_ai/_ask.py

@@ -0,0 +1,341 @@
+"""Turn typed questions into one System One call, and back again.
+
+The whole point of System One is that questions are answered in parallel and in
+isolation, so asking forty of them costs about what asking three costs. That only
+pays off if asking forty is as easy to write as asking three.
+
+There are three ways to say what you want, and they compile to the same request:
+
+```python
+await ask(Triage, state)                                       # a battery dataclass
+await ask(Choice[Intent], state)                               # a single question
+await ask({"intent": Choice[Intent], "hostile": Noul}, state)  # an ad-hoc battery
+```
+
+In every form the vocabulary documents itself: an enum's class docstring is the
+question and its member docstrings are the criteria. Override either with metadata
+-- `Annotated[Choice[Intent], {"question": ...}]` anywhere, or
+`field(metadata={...})` on a dataclass field, which wins over the annotation
+because it is the more specific place to say it.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import enum
+import time
+import typing
+from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
+
+from ._client import client as make_client
+from ._docs import class_doc, member_docs
+from ._types import CRITERIA_KEY, QUESTION_KEY, CallInfo, Choice, Noul, Score
+
+B = TypeVar("B")
+
+#: What you can hand to ask(): a battery dataclass, one question type, or a mapping.
+Askable = Union[Type[Any], Mapping[str, Any]]
+
+#: The name a single, unnamed question is asked under.
+SINGLE = "answer"
+
+
+class BatteryError(TypeError):
+    """The questions cannot be compiled. Always a problem in your code."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _Spec:
+    """One question, however it was declared."""
+
+    name: str
+    kind: str  # choice | score | noul
+    type: Any  # the bare Choice[...] / Score[...] / Noul
+    enum_cls: Optional[Type[enum.Enum]]
+    meta: Mapping[str, Any]
+
+
+def _unwrap(tp: Any) -> Tuple[Any, Mapping[str, Any]]:
+    """Split Annotated[X, {...}] into X and the metadata mapping it carries."""
+    if typing.get_origin(tp) is not None and hasattr(tp, "__metadata__"):
+        bare = typing.get_args(tp)[0]
+        meta: Dict[str, Any] = {}
+        for extra in tp.__metadata__:
+            if isinstance(extra, Mapping):
+                meta.update(extra)
+            elif isinstance(extra, str):
+                meta.setdefault(QUESTION_KEY, extra)
+        return bare, meta
+    return tp, {}
+
+
+def _kind(tp: Any) -> Optional[str]:
+    origin = typing.get_origin(tp) or tp
+    if origin is Noul:
+        return "noul"
+    if origin is Choice:
+        return "choice"
+    if origin is Score:
+        return "score"
+    return None
+
+
+def _enum_arg(tp: Any) -> Optional[Type[enum.Enum]]:
+    args = typing.get_args(tp)
+    if args and isinstance(args[0], type) and issubclass(args[0], enum.Enum):
+        return args[0]
+    return None
+
+
+def _spec(name: str, tp: Any, extra_meta: Optional[Mapping[str, Any]] = None) -> _Spec:
+    bare, annotated_meta = _unwrap(tp)
+    kind = _kind(bare)
+    if kind is None:
+        raise BatteryError(
+            f"'{name}' is typed {tp!r}, which is not a question. It must be "
+            "Choice[SomeEnum], Score[SomeIntEnum] or Noul."
+        )
+    # field metadata beats the annotation: it is the more specific place to say it.
+    meta = {**annotated_meta, **(extra_meta or {})}
+    enum_cls = _enum_arg(bare)
+    if kind == "choice" and enum_cls is None:
+        raise BatteryError(f"Choice '{name}' needs an Enum type argument, e.g. Choice[Intent].")
+    if kind == "score" and (enum_cls is None or not issubclass(enum_cls, enum.IntEnum)):
+        raise BatteryError(f"Score '{name}' needs an IntEnum type argument, e.g. Score[Severity].")
+    return _Spec(name=name, kind=kind, type=bare, enum_cls=enum_cls, meta=meta)
+
+
+def _specs(askable: Askable) -> list[_Spec]:
+    """Normalise every accepted form into one list of questions."""
+    if isinstance(askable, Mapping):
+        return [_spec(name, tp) for name, tp in askable.items()]
+
+    # Check for a question type before the dataclass branch: the answer types are
+    # themselves dataclasses, so a bare `Noul` would otherwise look like an empty battery.
+    if _kind(_unwrap(askable)[0]) is not None:
+        return [_spec(SINGLE, askable)]
+
+    if dataclasses.is_dataclass(askable) and isinstance(askable, type):
+        hints = typing.get_type_hints(askable, include_extras=True)
+        specs = []
+        for f in dataclasses.fields(askable):
+            tp = hints[f.name]
+            bare, _ = _unwrap(tp)
+            if _kind(bare) is None:
+                if QUESTION_KEY in f.metadata:  # a question was intended, the type is wrong
+                    raise BatteryError(
+                        f"Field '{f.name}' is typed {tp!r}, which is not a question. It must be "
+                        "Choice[SomeEnum], Score[SomeIntEnum] or Noul."
+                    )
+                continue  # an ordinary field, carried along but never asked
+            specs.append(_spec(f.name, tp, f.metadata))
+        if not specs:
+            raise BatteryError(
+                f"{askable.__name__} has no question fields. A question field is typed "
+                "Choice[SomeEnum], Score[SomeIntEnum] or Noul."
+            )
+        # ask() builds the battery out of answers alone, so any other field has to be
+        # able to default itself. Without this the failure is a bare TypeError from
+        # __init__ that never mentions the battery.
+        unfillable = [
+            f.name
+            for f in dataclasses.fields(askable)
+            if f.name not in {sp.name for sp in specs}
+            and f.init
+            and f.default is dataclasses.MISSING
+            and f.default_factory is dataclasses.MISSING
+        ]
+        if unfillable:
+            raise BatteryError(
+                f"{askable.__name__} has non-question field(s) {unfillable} with no default. "
+                "ask() constructs the battery from the answers, so every field it does not ask "
+                "for needs a default."
+            )
+        return specs
+
+    return [_spec(SINGLE, askable)]  # not a battery and not a question: _spec explains why
+
+
+def _instructions(spec: _Spec) -> str:
+    """The question text: what was written here, or what the vocabulary already says."""
+    asked = spec.meta.get(QUESTION_KEY)
+    if isinstance(asked, str) and asked.strip():
+        return asked
+    inherited = class_doc(spec.enum_cls) if spec.enum_cls is not None else None
+    if inherited:
+        return inherited
+    hint = (
+        "A Noul has no vocabulary to document itself with, so it always needs one."
+        if spec.kind == "noul"
+        else f"Give {spec.enum_cls.__name__} a class docstring, or say it here."  # type: ignore[union-attr]
+    )
+    raise BatteryError(f"'{spec.name}' has no question. {hint}")
+
+
+def _described(enum_cls: Type[enum.Enum]) -> Dict[str, Optional[str]]:
+    docs = member_docs(enum_cls)
+    return {m.name: docs.get(m.name) for m in enum_cls}
+
+
+def _choice_criteria(spec: _Spec) -> Dict[str, Optional[str]]:
+    """Keyed by member name -- the same string Flyte puts on the wire for an enum."""
+    enum_cls = typing.cast(Type[enum.Enum], spec.enum_cls)
+    criteria = _described(enum_cls)
+    given = spec.meta.get(CRITERIA_KEY)
+    if given is None:
+        return criteria
+    if not isinstance(given, Mapping):
+        raise BatteryError(f"Choice '{spec.name}' needs its criteria as a mapping of option -> description.")
+    for key, text in given.items():
+        member = key.name if isinstance(key, enum.Enum) else str(key)
+        if member not in criteria:
+            raise BatteryError(f"Choice '{spec.name}': '{member}' is not a member of {enum_cls.__name__}.")
+        criteria[member] = text
+    return criteria
+
+
+def _score_criteria(spec: _Spec) -> list:
+    """Positional, one rung per step -- which is what the SDK's Score takes."""
+    enum_cls = typing.cast(Type[enum.IntEnum], spec.enum_cls)
+    members = sorted(enum_cls, key=lambda m: m.value)
+    expected = list(range(len(members)))
+    if [m.value for m in members] != expected:
+        raise BatteryError(
+            f"Score '{spec.name}' uses {enum_cls.__name__}, whose values are {[m.value for m in members]}. "
+            f"A rubric is positional, so the members must be {expected} -- one rung per step, starting at zero."
+        )
+    given = spec.meta.get(CRITERIA_KEY)
+    if isinstance(given, Mapping):
+        given = {(k.name if isinstance(k, enum.Enum) else str(k)): v for k, v in given.items()}
+        unknown = set(given) - {m.name for m in members}
+        if unknown:
+            raise BatteryError(f"Score '{spec.name}': {sorted(unknown)} are not members of {enum_cls.__name__}.")
+    elif isinstance(given, Sequence) and not isinstance(given, str):
+        if len(given) != len(members):
+            raise BatteryError(
+                f"Score '{spec.name}' has {len(given)} criteria for {len(members)} rungs. "
+                "A positional rubric needs one entry per member."
+            )
+        return list(given)
+    elif given is not None:
+        raise BatteryError(f"Score '{spec.name}' needs its criteria as a sequence of rungs, or a mapping.")
+    else:
+        given = {}
+
+    docs = _described(enum_cls)
+    return [given.get(m.name) or docs.get(m.name) or m.name.replace("_", " ").lower() for m in members]
+
+
+def compile_questions(askable: Askable) -> Dict[str, Any]:
+    """Build the SDK's question objects from any accepted form."""
+    import typesafe_sdk as ts
+
+    questions: Dict[str, Any] = {}
+    for spec in _specs(askable):
+        instructions = _instructions(spec)
+        if spec.kind == "noul":
+            given = spec.meta.get(CRITERIA_KEY)
+            criteria = None
+            if given is not None:
+                if not isinstance(given, Mapping):
+                    raise BatteryError(f"Noul '{spec.name}' needs criteria like {{'true': ..., 'false': ...}}.")
+                criteria = ts.NoulCriteria(true=given.get("true"), false=given.get("false"))
+            questions[spec.name] = ts.Noul(instructions=instructions, criteria=criteria)
+        elif spec.kind == "choice":
+            questions[spec.name] = ts.Choice(instructions=instructions, criteria=_choice_criteria(spec))
+        else:
+            questions[spec.name] = ts.Score(instructions=instructions, criteria=_score_criteria(spec))
+    return questions
+
+
+def _by_member_name(enum_cls: Type[enum.Enum], probabilities: Dict[str, float]) -> Dict[str, float]:
+    """Re-key a rung-indexed distribution by member name, leaving unknown keys alone."""
+    by_value = {str(m.value): m.name for m in enum_cls}
+    return {by_value.get(k, k): v for k, v in probabilities.items()}
+
+
+def _probabilities(raw: Any) -> Dict[str, float]:
+    if not raw:
+        return {}
+    return {str(k): float(v) for k, v in dict(raw).items()}
+
+
+def _answer(spec: _Spec, answer: Any) -> Any:
+    if answer is None:
+        raise BatteryError(f"System 1 returned no answer for '{spec.name}'.")
+    if spec.kind == "noul":
+        return Noul(value=float(answer.noul))
+    enum_cls = typing.cast(Type[enum.Enum], spec.enum_cls)
+    confidence = float(getattr(answer, "confidence", 0.0) or 0.0)
+    probabilities = _probabilities(getattr(answer, "probabilities", None))
+    if spec.kind == "choice":
+        return Choice(value=enum_cls[answer.choice], confidence=confidence, probabilities=probabilities)
+    position = float(answer.score)
+    rung = min(max(round(position), 0), len(list(enum_cls)) - 1)
+    return Score(
+        value=enum_cls(rung),
+        position=position,
+        confidence=confidence,
+        # The SDK keys a Score's distribution by rung index; name them, so that both
+        # Choice.probabilities and Score.probabilities read as {member name: p}.
+        probabilities=_by_member_name(enum_cls, probabilities),
+    )
+
+
+def _assemble(askable: Askable, specs: list[_Spec], answers: Mapping[str, Any]) -> Any:
+    """Shape the answers like the thing that was asked."""
+    built = {spec.name: _answer(spec, answers.get(spec.name)) for spec in specs}
+    if isinstance(askable, Mapping):
+        return built
+    if dataclasses.is_dataclass(askable) and isinstance(askable, type):
+        return askable(**built)
+    return built[SINGLE]
+
+
+def _answers_of(resp: Any) -> Dict[str, Any]:
+    answers = getattr(resp, "answers", None)
+    if answers is None:  # older/newer shapes expose the three groups instead
+        answers = {**(resp.nouls or {}), **(resp.choices or {}), **(resp.scores or {})}
+    return dict(answers)
+
+
+async def ask_with_info(
+    askable: Askable,
+    state: Any,
+    *,
+    model: Optional[str] = None,
+    client: Any = None,
+) -> Tuple[Any, CallInfo]:
+    """Answer everything in one call, and report what the call cost."""
+    specs = _specs(askable)
+    questions = compile_questions(askable)
+    own = client is None
+    client = client or make_client(model=model)
+    started = time.perf_counter()
+    try:
+        resp = await client.system_one(state=state, questions=questions)
+    finally:
+        if own:
+            await client.aclose()
+    usage = getattr(resp, "usage", None)
+    info = CallInfo(
+        model=str(getattr(resp, "model", "") or ""),
+        questions=len(questions),
+        input_tokens=int(getattr(usage, "input_tokens", 0) or 0),
+        output_tokens=int(getattr(usage, "output_tokens", 0) or 0),
+        latency_s=round(time.perf_counter() - started, 3),
+    )
+    return _assemble(askable, specs, _answers_of(resp)), info
+
+
+async def ask(askable: Askable, state: Any, *, model: Optional[str] = None, client: Any = None) -> Any:
+    """Answer a battery, a single question, or a mapping of them -- in one call.
+
+    ```python
+    triage = await ask(Triage, {"ticket": text})                  # -> Triage
+    intent = await ask(Choice[Intent], {"ticket": text})          # -> Choice[Intent]
+    both = await ask({"intent": Choice[Intent], "hot": Noul}, s)  # -> dict
+    ```
+    """
+    answered, _ = await ask_with_info(askable, state, model=model, client=client)
+    return answered

flyteplugins/typesafe_ai/_client.py

@@ -0,0 +1,50 @@
+"""Reaching TypeSafe from a Flyte task, and saying so clearly when you can't."""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Optional
+
+#: The environment variable the TypeSafe SDK reads. Mount your secret as this:
+#:
+#:     flyte.Secret(key="TYPESAFE_API_KEY", as_env_var="TYPESAFE_API_KEY")
+#:
+#: `flyte.Secret` derives `as_env_var` from the key by upper-casing it and
+#: swapping `-` for `_`, so a key named TYPESAFE_API_KEY mounts correctly on its
+#: own -- spelling it out just makes the env var greppable.
+API_KEY_ENV = "TYPESAFE_API_KEY"
+
+
+class MissingAPIKey(RuntimeError):
+    """Raised at the point of use, in the task that actually needs the key."""
+
+
+def _explain() -> str:
+    return (
+        f"{API_KEY_ENV} is not set, so System 1 cannot be reached.\n\n"
+        "Declare the secret on the TaskEnvironment that runs this task:\n\n"
+        "    import flyte\n\n"
+        "    env = flyte.TaskEnvironment(\n"
+        '        "triage",\n'
+        f'        secrets=[flyte.Secret(key="{API_KEY_ENV}", as_env_var="{API_KEY_ENV}")],\n'
+        "    )\n\n"
+        "and create it once, if it does not exist yet:\n\n"
+        f"    flyte create secret {API_KEY_ENV} --value <your key>\n\n"
+        "The check happens here, at the point of use, rather than at import: a task that merely "
+        "passes answers along does not need the key, and an import-time failure would take down "
+        "every task in the module."
+    )
+
+
+def client(*, api_key: Optional[str] = None, model: Optional[str] = None, **kwargs: Any):
+    """An `AsyncTypeSafeClient`, or a message that says exactly what to do.
+
+    Without this, a missing key surfaces as a 401 from inside the vendor SDK, which
+    tells you nothing about Flyte secrets.
+    """
+    from typesafe_sdk import AsyncTypeSafeClient
+
+    key = api_key or os.environ.get(API_KEY_ENV, "").strip()
+    if not key:
+        raise MissingAPIKey(_explain())
+    return AsyncTypeSafeClient(api_key=key, model=model, **kwargs)

flyteplugins/typesafe_ai/_docs.py

@@ -0,0 +1,78 @@
+"""Recover an enum's own documentation, so the criteria do not have to be repeated.
+
+An enum already says what its members mean -- in a class docstring, and in the
+string literals underneath each member. The first is a real attribute; the second
+is not. A member assignment followed by a bare string literal leaves nothing on the
+member at runtime (`Severity.NONE.__doc__` returns the *class* docstring it
+inherits), so the only way to read it is to parse the source, which is what pydantic
+does for `use_attribute_docstrings` too.
+
+That makes member docs a best-effort input: they work from a normal module on
+disk and quietly return nothing when the source is not available (a REPL, `exec`,
+some frozen or zipped deployments). Everything here degrades to the member name
+rather than failing, because a missing description is a worse prompt, not a broken
+program.
+"""
+
+from __future__ import annotations
+
+import ast
+import enum
+import inspect
+import textwrap
+from functools import lru_cache
+from typing import Dict, Optional, Type
+
+#: Python <= 3.11 writes this docstring onto an undocumented Enum class itself.
+#: It is not something the author wrote, so it must not be mistaken for a question.
+#: (3.12+ leaves __doc__ as None on the class and only inherits from Enum.)
+_DEFAULT_ENUM_DOC = "An enumeration."
+
+
+@lru_cache(maxsize=None)
+def class_doc(cls: Type) -> Optional[str]:
+    """The enum's *own* class docstring, or None if it does not really have one.
+
+    Deliberately not `inspect.getdoc`: that inherits, and an undocumented enum
+    would then report CPython's own text -- "Create a collection of name/value
+    pairs." on 3.12+, or "Enum where members are also (and must be) ints" for an
+    IntEnum -- which would be sent to the model as the question. Reading the class
+    dict directly means only a docstring written on this enum counts.
+    """
+    doc = cls.__dict__.get("__doc__")
+    if not isinstance(doc, str) or not doc.strip() or doc.strip() == _DEFAULT_ENUM_DOC:
+        return None
+    return " ".join(doc.split())
+
+
+@lru_cache(maxsize=None)
+def member_docs(cls: Type[enum.Enum]) -> Dict[str, str]:
+    """{member name: docstring} for members documented by a literal underneath them."""
+    try:
+        source = inspect.getsource(cls)
+    except (OSError, TypeError):  # no source on disk: interactive, exec'd, frozen
+        return {}
+    try:
+        tree = ast.parse(textwrap.dedent(source)).body[0]
+    except (SyntaxError, IndexError):
+        return {}
+    if not isinstance(tree, ast.ClassDef):
+        return {}
+
+    docs: Dict[str, str] = {}
+    pending: Optional[str] = None
+    for node in tree.body:
+        if isinstance(node, (ast.Assign, ast.AnnAssign)):
+            target = node.targets[0] if isinstance(node, ast.Assign) else node.target
+            pending = target.id if isinstance(target, ast.Name) else None
+        elif (
+            pending
+            and isinstance(node, ast.Expr)
+            and isinstance(node.value, ast.Constant)
+            and isinstance(node.value.value, str)
+        ):
+            docs[pending] = " ".join(node.value.value.split())
+            pending = None
+        else:
+            pending = None
+    return docs

flyteplugins/typesafe_ai/_types.py

@@ -0,0 +1,91 @@
+"""The three answer types, as plain dataclasses.
+
+TypeSafe's System One model ("Jev") answers typed questions instead of writing
+text, and every answer arrives with the calibration that produced it. These
+dataclasses are how that answer crosses a Flyte task boundary: Flyte's
+`DataclassTransformer` carries them with no registration and no pydantic, and
+nesting an `Enum` or `IntEnum` inside one keeps it a real member on the way
+back out rather than degrading it to a string or an int.
+
+Note the deliberate name collision with `typesafe_sdk.Choice` / `Score` /
+`Noul`: those describe the *question* you ask, these hold the *answer* you get.
+You write the ones in this module; the plugin builds the SDK's from your battery.
+"""
+
+from __future__ import annotations
+
+import enum
+from dataclasses import dataclass, field
+from typing import Generic, Optional, TypeVar
+
+C = TypeVar("C", bound=enum.Enum)
+S = TypeVar("S", bound=enum.IntEnum)
+
+#: Metadata keys on a battery field. They are the SDK's own words: a question
+#: has `instructions` and `criteria`, and the shape of the criteria depends
+#: on the question type, exactly as in `typesafe_sdk`.
+QUESTION_KEY = "question"
+CRITERIA_KEY = "criteria"
+
+
+@dataclass
+class Choice(Generic[C]):
+    """One pick from a fixed vocabulary, carrying the calibration it came with."""
+
+    value: C
+    confidence: float = 0.0
+    probabilities: dict[str, float] = field(default_factory=dict)
+
+    def certain(self, threshold: float) -> bool:
+        """Is this pick confident enough to act on without a human?"""
+        return self.confidence >= threshold
+
+    def runner_up(self) -> Optional[tuple[str, float]]:
+        """The second-most-likely option, which is what you show a reviewer."""
+        ranked = sorted(self.probabilities.items(), key=lambda kv: kv[1], reverse=True)
+        return ranked[1] if len(ranked) > 1 else None
+
+
+@dataclass
+class Score(Generic[S]):
+    """A position on a rubric.
+
+    Two representations, because both are useful: `value` is the rung that was
+    picked (an `IntEnum` member, so it compares and orders), and `position`
+    is the unrounded place on the scale that Jev actually returned. Branch on the
+    first, sort and threshold on the second.
+    """
+
+    value: S
+    position: float = 0.0
+    confidence: float = 0.0
+    probabilities: dict[str, float] = field(default_factory=dict)
+
+    def at_least(self, rung: S) -> bool:
+        return self.value >= rung
+
+
+@dataclass
+class Noul:
+    """Truthfulness in 0..1.
+
+    Deliberately no `__bool__`: `if noul:` would make 0.02 and 0.98 alike,
+    and choosing the threshold is the part you want in code where it can be read,
+    reviewed and changed.
+    """
+
+    value: float = 0.0
+
+    def at(self, threshold: float) -> bool:
+        return self.value >= threshold
+
+
+@dataclass
+class CallInfo:
+    """What one `system_one` call cost, for the report and for cost accounting."""
+
+    model: str = ""
+    questions: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    latency_s: float = 0.0

flyteplugins_typesafe_ai-0.0.0a0.dist-info/METADATA

@@ -0,0 +1,275 @@
+Metadata-Version: 2.4
+Name: flyteplugins-typesafe-ai
+Version: 0.0.0a0
+Summary: Run TypeSafe's System One model (Jev) inside durable Flyte tasks
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: typesafe-sdk>=0.5.7
+Requires-Dist: flyte
+
+# flyteplugins-typesafe-ai
+
+Run [TypeSafe](https://docs.typesafe.ai/introduction)'s System One model ("Jev")
+inside durable 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
+that makes it 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.
+
+This plugin supplies the two things that takes on Flyte: a shape for the answers
+that survives a task boundary, and a way to ask a whole battery at once.
+
+```bash
+pip install flyteplugins-typesafe-ai
+```
+
+## Quickstart
+
+```python
+import enum
+from dataclasses import dataclass, field
+
+import flyte
+from flyteplugins.typesafe_ai import Choice, Noul, Score, ask
+
+env = flyte.TaskEnvironment(
+    "triage",
+    secrets=[flyte.Secret(key="TYPESAFE_API_KEY", as_env_var="TYPESAFE_API_KEY")],
+)
+
+
+class Intent(enum.Enum):
+    """What is this customer asking for?"""
+
+    REFUND = "refund"
+    """they want money back"""
+    DELIVERY = "delivery status"
+    """they are asking where their order is"""
+    OTHER = "something else"
+
+
+class Severity(enum.IntEnum):
+    """How badly are they blocked?"""
+
+    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 is involved"""
+    BLOCKING = 3
+    """they cannot use the product at all"""
+
+
+@dataclass
+class Triage:
+    # The enums document themselves: class docstring -> question, member docstrings -> criteria.
+    intent: Choice[Intent]
+    severity: Score[Severity]
+    # A Noul has no vocabulary to document itself with, so it carries both.
+    hostile: Noul = field(
+        metadata={
+            "question": "Is the customer hostile?",
+            "criteria": {"true": "insults or threats", "false": "civil"},
+        }
+    )
+    money_at_stake: Noul = field(metadata={"question": "Is a payment or refund involved?"})
+
+
+@env.task
+async def handle(ticket: str) -> str:
+    t = await ask(Triage, {"ticket": ticket})        # one request, four answers
+
+    if t.hostile.at(0.8):                            # thresholds live in your code
+        return "escalate"
+    if not t.intent.certain(0.85):
+        return "review"
+    return f"auto: {t.intent.value.value}, severity {t.severity.value.name}"
+```
+
+## The three answer types
+
+| Type | Holds | Useful members |
+| --- | --- | --- |
+| `Choice[SomeEnum]` | the picked member, `confidence`, `probabilities` | `.certain(threshold)`, `.runner_up()` |
+| `Score[SomeIntEnum]` | the picked rung, 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 no bespoke type
+transformer exists — they reuse Flyte's built-in `DataclassTransformer`. 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.
+
+## 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:
+
+```python
+@dataclass
+class Triage:
+    intent: Choice[Intent]        # question and criteria both come from Intent
+    severity: Score[Severity]
+```
+
+Override either 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 |
+
+```python
+intent: Choice[Intent] = field(metadata={"question": "asked a different way"})
+```
+
+`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`.
+
+`Noul` is the one that always needs you: it has no vocabulary to document itself
+with, so a `Noul` without a question is an error naming the field.
+
+Member docstrings are not stored on the object at runtime — `Severity.NONE.__doc__`
+returns the *class* docstring it inherits — so they are read by parsing the source,
+the same way pydantic implements `use_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:
+
+```python
+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:
+
+```python
+await ask(Annotated[Score[Harm], {"question": "How much harm would this do?"}], state)
+await ask(Annotated[Noul, "Is this aimed at a specific person?"], state)
+```
+
+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. Use `ask_with_info()` when you want the model name, question
+count, token usage and latency back alongside the answers.
+
+The name collision with `typesafe_sdk.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.
+
+## The API key
+
+The SDK reads the key from `TYPESAFE_API_KEY`, so mount your secret as that env
+var. There is no helper for this — it is a plain `flyte.Secret`:
+
+```python
+env = flyte.TaskEnvironment(
+    "triage",
+    secrets=[flyte.Secret(key="TYPESAFE_API_KEY", as_env_var="TYPESAFE_API_KEY")],
+)
+```
+
+`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.
+
+If your secret is stored under a different name, point the key at it and keep the
+mount:
+
+```python
+flyte.Secret(key="my-org-typesafe-key", as_env_var="TYPESAFE_API_KEY")
+```
+
+Create the secret once:
+
+```bash
+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. It is deliberately not an import-time check: 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.
+
+## Examples
+
+- [`examples/triage.py`](examples/triage.py) — fourteen questions in one call, then
+  confidence-gated routing in ordinary Python
+- [`examples/fanout.py`](examples/fanout.py) — a backlog of tickets, one durable
+  Flyte task each, one System One call inside each
+- [`examples/single.py`](examples/single.py) — `Choice`, `Score` and `Noul` used on
+  their own, without a battery dataclass
+
+Both run against a cluster:
+
+```bash
+flyte run --root-dir plugins/typesafe-ai/examples plugins/typesafe-ai/examples/triage.py handle
+```
+
+While this package is unpublished, set `TYPESAFE_LOCAL_WHEELS=1` and build the
+wheels first with `make dist && make dist-plugins`.
+
+## 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:
+
+```python
+@env.task
+async def classify(message: str) -> Choice[Action]:
+    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 1 produces this", make it a task: you get durability,
+caching and retries, and the call stays visible in the run graph.
+
+## A note on IntEnum
+
+`Score` takes an `IntEnum` because a rubric is ordered. Flyte's enum transformer
+used to accept string-valued enums only; it now supports `IntEnum` as well,
+serialized by member name like every other enum, so `severity: Severity` also works
+as a bare task input or output. `Flag`/`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 second core change makes `Choice[Intent]` work as a task type at all: a
+parameterized dataclass is a generic *alias*, which `dataclasses.is_dataclass()`
+rejects, so it used to fall through the type engine to pickle. The alias now
+resolves to its origin for structural checks while the alias itself is kept for
+decoding, which is what binds the type variable.

flyteplugins_typesafe_ai-0.0.0a0.dist-info/RECORD

@@ -0,0 +1,10 @@
+flyteplugins/typesafe_ai/__init__.py,sha256=Y_bjsLJCsYV5tTa3FG1EZCb02_3CbJVTQbrHXqeVeRQ,3121
+flyteplugins/typesafe_ai/_ask.py,sha256=IY_h9JsRuk54AqrEv3TJ0qS8Pkwwx4hybph3cCJM09o,14171
+flyteplugins/typesafe_ai/_client.py,sha256=gG3ZD31KFLc5vn8jzuGRBGv0ejR2j2Do72uGGsqgv7U,2002
+flyteplugins/typesafe_ai/_docs.py,sha256=upSH87D3TYSGQkpT7EqvbfliJeZ_ejZNY8ed_M1sYUk,3197
+flyteplugins/typesafe_ai/_types.py,sha256=_gQ9HX5JK2EI5rwEhdYCa8j7vbJyjL5TEjlseEnKvJc,3032
+flyteplugins_typesafe_ai-0.0.0a0.dist-info/METADATA,sha256=DXDgECn0DTfrEmszYJwymVBs4LRKt5BFCFF2A7OvhsM,10877
+flyteplugins_typesafe_ai-0.0.0a0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+flyteplugins_typesafe_ai-0.0.0a0.dist-info/entry_points.txt,sha256=vJ0BUqt9Y4HiN6Nbwr7CQtFp0gSJPGoKsS1SCJSqBWY,88
+flyteplugins_typesafe_ai-0.0.0a0.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_typesafe_ai-0.0.0a0.dist-info/RECORD,,

flyteplugins_typesafe_ai-0.0.0a0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_typesafe_ai-0.0.0a0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[flyte.plugins.types]
+typesafe_ai = flyteplugins.typesafe_ai:register_typesafe_ai_types

flyteplugins_typesafe_ai-0.0.0a0.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins