Jev, or its local stand-in

uv init jev-lab && cd jev-lab
uv add --dev pytest

"""jevlocal: a local stand-in for the TypeSafe Python SDK (typesafe_sdk).
 
It mirrors the subset of the real SDK used in this course, so exercise code
written against it runs unchanged on Jev once you have early access:
 
    TypeSafeClient(...).system_one(state=..., questions={...})
    Choice / Score / Noul / NoulCriteria
    response.answers / .choices / .scores / .nouls / .usage / .model
 
Three backends, picked with the JEV_BACKEND environment variable or the
`backend=` argument:
 
- "fake"      (default) deterministic and offline. Scores options by word
              overlap between the state and each option's text, or uses the
              answers you pass in `fake_answers`. Good for tests and for
              wiring up control flow. It is NOT intelligent.
- "claude"    asks Claude Code headless (`claude -p`) for a probability
              distribution per question. Runs on your Claude subscription.
              These are verbalised probabilities from a generative model,
              not calibrated ones: measuring that gap is part of the course.
- "typesafe"  calls the real Jev API through `typesafe_sdk` (needs
              TYPESAFE_API_KEY and `pip install typesafe-sdk`).
 
Confidence is computed with the formulas TypeSafe documents at
https://docs.typesafe.ai/confidence (Choice, Score; Noul has none).
"""
 
from __future__ import annotations
 
import json
import math
import os
import re
import shutil
import subprocess
from dataclasses import dataclass, field
from typing import Any, Callable, Literal, Mapping, Union
 
__all__ = [
    "Choice", "Score", "Noul", "NoulCriteria",
    "ChoiceAnswer", "ScoreAnswer", "NoulAnswer", "Usage", "SystemOneResponse",
    "TypeSafeClient", "choice_confidence", "score_confidence", "noul_confidence",
]
 
Text = Union[str, Mapping[str, Any], list]
 
 
# ---------------------------------------------------------------- questions
 
@dataclass(frozen=True)
class Choice:
    instructions: Text
    criteria: Mapping[str, Text]
    type: Literal["choice"] = "choice"
 
    def __post_init__(self) -> None:
        if not 2 <= len(self.criteria) <= 255:
            raise ValueError("A Choice needs between 2 and 255 options")
 
 
@dataclass(frozen=True)
class Score:
    instructions: Text
    criteria: list
    type: Literal["score"] = "score"
 
    def __post_init__(self) -> None:
        if not 2 <= len(self.criteria) <= 10:
            raise ValueError("A Score needs between 2 and 10 levels, ordered low to high")
 
 
@dataclass(frozen=True)
class NoulCriteria:
    true: Text
    false: Text
 
 
@dataclass(frozen=True)
class Noul:
    instructions: Text
    criteria: NoulCriteria | None = None
    type: Literal["noul"] = "noul"
 
 
Question = Union[Choice, Score, Noul]
 
 
# ------------------------------------------------------------------ answers
 
@dataclass(frozen=True)
class ChoiceAnswer:
    choice: str
    confidence: float
    probabilities: dict[str, float]
    type: Literal["choice"] = "choice"
 
 
@dataclass(frozen=True)
class ScoreAnswer:
    score: float
    confidence: float
    legend: dict[int, Any]
    probabilities: dict[int, float]
    type: Literal["score"] = "score"
 
 
@dataclass(frozen=True)
class NoulAnswer:
    noul: float
    type: Literal["noul"] = "noul"
 
 
Answer = Union[ChoiceAnswer, ScoreAnswer, NoulAnswer]
 
 
@dataclass(frozen=True)
class Usage:
    input_tokens: int | None = None
    output_tokens: int | None = None
 
 
@dataclass(frozen=True)
class SystemOneResponse:
    model: str
    answers: dict[str, Answer]
    usage: Usage = field(default_factory=Usage)
 
    @property
    def choices(self) -> dict[str, ChoiceAnswer]:
        return {k: v for k, v in self.answers.items() if isinstance(v, ChoiceAnswer)}
 
    @property
    def scores(self) -> dict[str, ScoreAnswer]:
        return {k: v for k, v in self.answers.items() if isinstance(v, ScoreAnswer)}
 
    @property
    def nouls(self) -> dict[str, NoulAnswer]:
        return {k: v for k, v in self.answers.items() if isinstance(v, NoulAnswer)}
 
 
# --------------------------------------------------------------- confidence
 
def choice_confidence(probs: list[float]) -> float:
    """(p_max - 1/n) / (1 - 1/n): 0 for a uniform spread, 1 for a single peak."""
    n = len(probs)
    return max(0.0, (max(probs) - 1 / n) / (1 - 1 / n))
 
 
def score_confidence(probs: list[float]) -> float:
    """1 - (expected distance from the most likely level) / (mean distance of a uniform
    spread from the middle level), floored at 0."""
    n = len(probs)
    mode = max(range(n), key=lambda i: probs[i])
    spread = sum(p * abs(i - mode) for i, p in enumerate(probs))
    uniform = sum(abs(i - (n - 1) / 2) for i in range(n)) / n
    return max(0.0, 1 - spread / uniform)
 
 
def noul_confidence(p: float) -> float:
    """|2p - 1|. Jev returns no confidence for a Noul; this is the documented equivalent."""
    return abs(2 * p - 1)
 
 
def _normalise(values: list[float]) -> list[float]:
    values = [max(0.0, float(v)) for v in values]
    total = sum(values)
    return [1 / len(values)] * len(values) if total == 0 else [v / total for v in values]
 
 
def _build_answer(question: Question, dist: Any) -> Answer:
    """Turn a raw distribution (list, dict or float) into a typed answer."""
    if isinstance(question, Noul):
        p = float(dist if not isinstance(dist, Mapping) else dist.get("true", dist.get("yes", 0.5)))
        return NoulAnswer(noul=min(1.0, max(0.0, p)))
    if isinstance(question, Choice):
        keys = list(question.criteria)
        raw = [dist.get(k, 0.0) for k in keys] if isinstance(dist, Mapping) else list(dist)
        probs = _normalise(raw)
        best = max(range(len(keys)), key=lambda i: probs[i])
        return ChoiceAnswer(choice=keys[best], confidence=choice_confidence(probs),
                            probabilities=dict(zip(keys, probs)))
    levels = list(range(len(question.criteria)))
    raw = [dist.get(i, dist.get(str(i), 0.0)) for i in levels] if isinstance(dist, Mapping) else list(dist)
    probs = _normalise(raw)
    return ScoreAnswer(score=sum(i * p for i, p in zip(levels, probs)),
                       confidence=score_confidence(probs),
                       legend=dict(zip(levels, question.criteria)),
                       probabilities=dict(zip(levels, probs)))
 
 
# ----------------------------------------------------------------- backends
 
def _text(value: Any) -> str:
    return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, sort_keys=True)
 
 
def _words(value: Any) -> set[str]:
    return set(re.findall(r"[a-z0-9]{3,}", _text(value).lower()))
 
 
def _fake_distribution(state: Any, question: Question) -> Any:
    """Word-overlap heuristic: deterministic, offline, deliberately dumb."""
    words = _words(state)
    if isinstance(question, Noul):
        target = _words(question.criteria.true if question.criteria else question.instructions)
        overlap = len(words & target) / max(1, len(target))
        return 1 / (1 + math.exp(-6 * (overlap - 0.3)))
    options = (list(question.criteria.values()) if isinstance(question, Choice) else question.criteria)
    logits = [3.0 * len(words & _words(o)) / max(1, len(_words(o))) for o in options]
    exp = [math.exp(x) for x in logits]
    return [e / sum(exp) for e in exp]
 
 
def _claude_schema(questions: Mapping[str, Question]) -> dict:
    props: dict[str, Any] = {}
    for qid, q in questions.items():
        if isinstance(q, Noul):
            props[qid] = {"type": "number", "minimum": 0, "maximum": 1}
        else:
            keys = list(q.criteria) if isinstance(q, Choice) else [str(i) for i in range(len(q.criteria))]
            props[qid] = {"type": "object", "additionalProperties": False, "required": keys,
                          "properties": {k: {"type": "number", "minimum": 0, "maximum": 1} for k in keys}}
    return {"type": "object", "additionalProperties": False, "required": list(questions), "properties": props}
 
 
def _claude_prompt(state: Any, questions: Mapping[str, Question]) -> str:
    lines = ["Answer each question about the STATE below. Do not write prose.",
             "For a yes/no question give the probability that the answer is yes.",
             "For a choice or score question give a probability for every option; they must sum to 1.",
             "Use probabilities that reflect your real uncertainty.", "", "STATE:", _text(state), "", "QUESTIONS:"]
    for qid, q in questions.items():
        lines.append(f"- {qid} ({q.type}): {_text(q.instructions)}")
        if isinstance(q, Choice):
            lines += [f"    option {k}: {_text(v)}" for k, v in q.criteria.items()]
        elif isinstance(q, Score):
            lines += [f"    level {i}: {_text(v)}" for i, v in enumerate(q.criteria)]
        elif q.criteria:
            lines += [f"    yes means: {_text(q.criteria.true)}", f"    no means: {_text(q.criteria.false)}"]
    return "\n".join(lines)
 
 
def _run_claude(state: Any, questions: Mapping[str, Question], model: str) -> tuple[dict, Usage]:
    if shutil.which("claude") is None:
        raise RuntimeError("JEV_BACKEND=claude needs Claude Code on your PATH (see /setup/claude-code)")
    args = ["claude", "-p", _claude_prompt(state, questions), "--output-format", "json",
            "--json-schema", json.dumps(_claude_schema(questions)), "--model", model]
    out = subprocess.run(args, capture_output=True, text=True, check=True, timeout=300).stdout
    envelope = json.loads(out)
    usage = envelope.get("usage") or {}
    return envelope["structured_output"], Usage(usage.get("input_tokens"), usage.get("output_tokens"))
 
 
def _run_typesafe(state: Any, questions: Mapping[str, Question], model: str) -> SystemOneResponse:
    import typesafe_sdk as ts  # pip install typesafe-sdk
 
    real = {}
    for qid, q in questions.items():
        if isinstance(q, Choice):
            real[qid] = ts.Choice(instructions=q.instructions, criteria=dict(q.criteria))
        elif isinstance(q, Score):
            real[qid] = ts.Score(instructions=q.instructions, criteria=list(q.criteria))
        else:
            real[qid] = ts.Noul(instructions=q.instructions)
    with ts.TypeSafeClient(model=model) as client:
        r = client.system_one(state=state, questions=real)
    answers = {qid: _build_answer(questions[qid], a.noul if a.type == "noul" else a.probabilities)
               for qid, a in r.answers.items()}
    return SystemOneResponse(model=r.model, answers=answers,
                             usage=Usage(r.usage.input_tokens, r.usage.output_tokens))
 
 
# ------------------------------------------------------------------- client
 
FakeAnswers = Union[Mapping[str, Any], Callable[[Any, str, Question], Any]]
 
 
class TypeSafeClient:
    """Same call shape as typesafe_sdk.TypeSafeClient.system_one."""
 
    def __init__(self, model: str | None = None, backend: str | None = None,
                 fake_answers: FakeAnswers | None = None, claude_model: str = "haiku") -> None:
        self.backend = backend or os.environ.get("JEV_BACKEND", "fake")
        if self.backend not in ("fake", "claude", "typesafe"):
            raise ValueError(f"Unknown backend {self.backend!r}")
        self.model = model or os.environ.get("TYPESAFE_DEFAULT_MODEL", "jev-latest")
        self.fake_answers = fake_answers
        self.claude_model = claude_model
 
    def __enter__(self) -> "TypeSafeClient":
        return self
 
    def __exit__(self, *exc: object) -> None:
        return None
 
    def system_one(self, state: Any, questions: Mapping[str, Question], model: str | None = None) -> SystemOneResponse:
        if not questions:
            raise ValueError("Ask at least one question")
        if self.backend == "typesafe":
            return _run_typesafe(state, questions, model or self.model)
        if self.backend == "claude":
            raw, usage = _run_claude(state, questions, self.claude_model)
            answers = {qid: _build_answer(q, raw[qid]) for qid, q in questions.items()}
            return SystemOneResponse(model=f"claude-{self.claude_model} (via claude -p)", answers=answers, usage=usage)
        answers = {}
        for qid, q in questions.items():
            if callable(self.fake_answers):
                dist = self.fake_answers(state, qid, q)
            elif self.fake_answers is not None and qid in self.fake_answers:
                dist = self.fake_answers[qid]
            else:
                dist = _fake_distribution(state, q)
            answers[qid] = _build_answer(q, dist)
        tokens = len(_text(state)) // 4 + sum(len(_text(q.instructions)) // 4 for q in questions.values())
        return SystemOneResponse(model="jevlocal-fake", answers=answers, usage=Usage(tokens, 0))

uv add typesafe-sdk
echo "TYPESAFE_API_KEY=your-key-here" >> .env
echo ".env" >> .gitignore

JEV_BACKEND=typesafe uv run --env-file .env python your_script.py

from jevlocal import Choice, Noul, TypeSafeClient
 
with TypeSafeClient() as client:
    r = client.system_one(
        state="I was charged twice for my invoice this month.",
        questions={
            "team": Choice(instructions="Which team should handle this?",
                           criteria={"billing": "charged invoice payment",
                                     "shipping": "delivery parcel courier"}),
            "refund": Noul(instructions="Is the customer charged twice?"),
        },
    )
print(r.model, r.choices["team"].choice, round(r.choices["team"].confidence, 2))

jevlocal-fake billing 0.76