arcus-cli 0.1.0__py3-none-any.whl

arcus/eval/regret.py

@@ -0,0 +1,96 @@
+import random
+from dataclasses import dataclass
+from typing import Callable
+
+import numpy as np
+
+from arcus.adapters.arc_adapter import ArcModel
+from arcus.routing.bandit import Bandit, EpsilonGreedyBandit, RandomBandit, ThompsonSamplingBandit, UCB1Bandit
+
+ARMS = [m.value for m in ArcModel]
+
+
+@dataclass(frozen=True)
+class ArmDistribution:
+    mean: float
+    std: float
+
+
+# illustrative synthetic reward distributions, not measured. regret
+# benchmarking needs a known ground-truth mean per arm to compute regret
+# against, and that's exactly the thing live traffic can never give us:
+# a real request only ever explores one arm per round, so we'd have no
+# way to know what the *other* three would have scored that round. this
+# is the standard way to study a bandit algorithm's exploration behavior
+# in isolation from actual model quality, not a stand-in for missing
+# data.
+DEFAULT_ARM_DISTRIBUTIONS: dict[str, ArmDistribution] = {
+    "gpt-oss-120b": ArmDistribution(mean=0.72, std=0.10),
+    "GLM-5.3": ArmDistribution(mean=0.68, std=0.12),
+    "Kimi-K3": ArmDistribution(mean=0.75, std=0.08),
+    "DeepSeek-V4-Flash": ArmDistribution(mean=0.70, std=0.15),
+}
+
+
+def _sample_reward(dist: ArmDistribution, rng: random.Random) -> float:
+    # clipped to [0, 1], real compute_reward() output is bounded there too
+    return min(1.0, max(0.0, rng.gauss(dist.mean, dist.std)))
+
+
+def simulate_regret(
+    bandit_factory: Callable[[], Bandit],
+    arm_distributions: dict[str, ArmDistribution],
+    n_rounds: int,
+    seed: int,
+) -> list[float]:
+    """Runs one bandit for n_rounds against a known synthetic reward
+    environment and returns its cumulative regret curve. Regret per
+    round is the gap between the best arm's true mean and the pulled
+    arm's true mean, not the noisy realized reward, so the curve tracks
+    the algorithm's actual exploration cost instead of being dominated
+    by sampling noise.
+    """
+    # Thompson sampling draws from numpy's global RNG (see bandit.py),
+    # the other algorithms only touch the stdlib random module, seeding
+    # both keeps every algorithm's run reproducible regardless of which
+    # one is being simulated.
+    rng = random.Random(seed)
+    np.random.seed(seed)
+
+    bandit = bandit_factory()
+    best_mean = max(dist.mean for dist in arm_distributions.values())
+
+    cumulative_regret = 0.0
+    curve = []
+    for _ in range(n_rounds):
+        arm = bandit.select_arm()
+        reward = _sample_reward(arm_distributions[arm], rng)
+        bandit.update(arm, reward)
+
+        cumulative_regret += best_mean - arm_distributions[arm].mean
+        curve.append(cumulative_regret)
+
+    return curve
+
+
+def run_regret_benchmark(
+    arm_distributions: dict[str, ArmDistribution] = DEFAULT_ARM_DISTRIBUTIONS,
+    n_rounds: int = 2000,
+    seed: int = 42,
+) -> dict[str, list[float]]:
+    """Cumulative regret curves for all three real algorithms plus the
+    random baseline, demonstrating this is a real algorithm family and
+    not just "a bandit works."
+    """
+    arms = list(arm_distributions.keys())
+    factories: dict[str, Callable[[], Bandit]] = {
+        "epsilon_greedy": lambda: EpsilonGreedyBandit(arms),
+        "ucb1": lambda: UCB1Bandit(arms),
+        "thompson": lambda: ThompsonSamplingBandit(arms),
+        "random": lambda: RandomBandit(arms),
+    }
+
+    return {
+        name: simulate_regret(factory, arm_distributions, n_rounds, seed)
+        for name, factory in factories.items()
+    }

arcus/quality/gate.py

@@ -0,0 +1,225 @@
+import re
+import time
+from dataclasses import dataclass
+
+from openai import APIError
+from openai.types.chat import ChatCompletion
+from pydantic import BaseModel, ValidationError
+
+from arcus.adapters.arc_adapter import ArcAdapter
+from arcus.routing.bandit import ContextualBandit
+from arcus.routing.reward import compute_reward
+
+_REFUSAL_MARKERS = re.compile(
+    r"i cannot assist|i can't assist"
+    r"|i cannot help|i can't help"
+    r"|as an ai language model"
+    r"|i'm sorry,? but i can'?t"
+    r"|i am not able to|i'm not able to"
+    r"|i won'?t be able to"
+    r"|i must decline"
+    r"|against my guidelines",
+    re.IGNORECASE,
+)
+
+_REPETITION_THRESHOLD = 0.5
+# below this many words there aren't enough trigrams for the duplication
+# ratio to mean anything, a two-sentence answer isn't "looping" just
+# because it reuses a word.
+_MIN_WORDS_FOR_REPETITION_CHECK = 9
+
+
+@dataclass(frozen=True)
+class QualityIssue:
+    kind: str
+    detail: str
+
+
+@dataclass(frozen=True)
+class QualityCheckResult:
+    passed: bool
+    issues: list[QualityIssue]
+
+
+def check_truncation(finish_reason: str | None) -> QualityIssue | None:
+    if finish_reason == "length":
+        return QualityIssue("truncated", "response was cut off before finishing (finish_reason=length)")
+    return None
+
+
+def check_empty(content: str | None) -> QualityIssue | None:
+    # deliberately not a fuzzy "too short" length check, a correct answer
+    # can legitimately be one word. this only catches actually-empty output.
+    if not content or not content.strip():
+        return QualityIssue("empty", "response was empty or whitespace only")
+    return None
+
+
+def check_repetition(content: str | None, threshold: float = _REPETITION_THRESHOLD) -> QualityIssue | None:
+    if not content:
+        return None
+
+    words = content.split()
+    if len(words) < _MIN_WORDS_FOR_REPETITION_CHECK:
+        return None
+
+    trigrams = [tuple(words[i : i + 3]) for i in range(len(words) - 2)]
+    duplication_ratio = 1 - len(set(trigrams)) / len(trigrams)
+
+    if duplication_ratio > threshold:
+        return QualityIssue("repetitive", f"trigram duplication ratio {duplication_ratio:.2f} exceeds {threshold}")
+    return None
+
+
+def check_refusal(content: str | None) -> QualityIssue | None:
+    if not content:
+        return None
+    if _REFUSAL_MARKERS.search(content):
+        return QualityIssue("refusal", "response matched a known refusal phrase")
+    return None
+
+
+def check_schema(content: str | None, schema: type[BaseModel] | None) -> QualityIssue | None:
+    if schema is None:
+        return None
+    if not content:
+        return QualityIssue("schema_invalid", "no content to validate against schema")
+    try:
+        schema.model_validate_json(content)
+    except ValidationError as e:
+        return QualityIssue("schema_invalid", str(e))
+    except ValueError as e:
+        # model_validate_json raises a plain ValueError (from the underlying
+        # json parse) when content isn't valid JSON at all, not a
+        # ValidationError, both count as a schema failure here
+        return QualityIssue("schema_invalid", str(e))
+    return None
+
+
+def check_response(
+    content: str | None,
+    finish_reason: str | None,
+    schema: type[BaseModel] | None = None,
+) -> QualityCheckResult:
+    checks = [
+        check_truncation(finish_reason),
+        check_empty(content),
+        check_repetition(content),
+        check_refusal(content),
+        check_schema(content, schema),
+    ]
+    issues = [issue for issue in checks if issue is not None]
+    return QualityCheckResult(passed=len(issues) == 0, issues=issues)
+
+
+@dataclass(frozen=True)
+class AttemptDetail:
+    model: str
+    passed: bool
+    reward: float
+    latency_ms: float
+    propensity: float
+    issues: list[QualityIssue]
+
+
+@dataclass(frozen=True)
+class QualityGateOutcome:
+    # None only when every arm errored out at the API level rather than
+    # returning a bad answer, there was nothing to hand back at all
+    response: ChatCompletion | None
+    model_used: str
+    passed: bool
+    attempts: list[AttemptDetail]
+    issues: list[QualityIssue]
+
+
+def call_with_quality_gate(
+    adapter: ArcAdapter,
+    bandit: ContextualBandit,
+    context_key: str,
+    messages: list[dict],
+    schema: type[BaseModel] | None = None,
+    max_attempts: int | None = None,
+) -> QualityGateOutcome:
+    """Calls ARC, runs the response through the quality gate, and retries
+    with a different arm on failure, feeding a real (not hardcoded)
+    negative reward back to the bandit each time so it actually learns
+    from the failure instead of just being told "not this one." Every
+    attempt gets logged, not just the winning one, so a caller can later
+    tell how often each model got caught by the gate, not only which
+    model ended up serving the response.
+    """
+    max_attempts = min(max_attempts or len(bandit.arms), len(bandit.arms))
+
+    already_tried: set[str] = set()
+    attempts: list[AttemptDetail] = []
+    completion = None
+    arm = None
+    last_issues: list[QualityIssue] = []
+
+    for _ in range(max_attempts):
+        arm = bandit.select_arm(context_key, exclude=already_tried)
+        # propensity has to be read off the bandit's state right here,
+        # before update() below changes the counts it's computed from
+        propensity = bandit.propensity(context_key, arm, exclude=already_tried)
+
+        start = time.monotonic()
+        try:
+            completion = adapter.chat(arm, messages)
+        except APIError as e:
+            # a network blip, rate limit, or ARC-side outage on this one
+            # model isn't a reason to give up on the whole request, treat
+            # it exactly like a failed quality check: log it, penalize
+            # this arm for this context, and let the loop try the next
+            # one instead of crashing the CLI with a raw traceback.
+            latency_ms = (time.monotonic() - start) * 1000
+            issue = QualityIssue("api_error", str(e))
+            reward = compute_reward(latency_ms=latency_ms, model=arm, quality_score=0.0)
+            bandit.update(context_key, arm, reward)
+
+            attempts.append(
+                AttemptDetail(
+                    model=arm, passed=False, reward=reward, latency_ms=latency_ms,
+                    propensity=propensity, issues=[issue],
+                )
+            )
+            last_issues = [issue]
+            completion = None
+            already_tried.add(arm)
+            continue
+
+        latency_ms = (time.monotonic() - start) * 1000
+
+        choice = completion.choices[0]
+        result = check_response(choice.message.content, choice.finish_reason, schema)
+        last_issues = result.issues
+
+        quality_score = 1.0 if result.passed else 0.0
+        reward = compute_reward(latency_ms=latency_ms, model=arm, quality_score=quality_score)
+        bandit.update(context_key, arm, reward)
+
+        attempts.append(
+            AttemptDetail(
+                model=arm,
+                passed=result.passed,
+                reward=reward,
+                latency_ms=latency_ms,
+                propensity=propensity,
+                issues=result.issues,
+            )
+        )
+
+        if result.passed:
+            return QualityGateOutcome(
+                response=completion, model_used=arm, passed=True, attempts=attempts, issues=[]
+            )
+
+        already_tried.add(arm)
+
+    # every arm got tried and none passed (or errored out), hand back the
+    # last attempt and let the caller decide what to do. response may be
+    # None here if the very last arm failed with an API error rather than
+    # a bad answer, callers need to handle that.
+    return QualityGateOutcome(
+        response=completion, model_used=arm, passed=False, attempts=attempts, issues=last_issues
+    )

arcus/routing/bandit.py

@@ -0,0 +1,226 @@
+import math
+import random
+from typing import Callable, Protocol
+
+import numpy as np
+
+
+class Bandit(Protocol):
+    def select_arm(self, exclude: set[str] | None = None) -> str: ...
+    def update(self, arm: str, reward: float) -> None: ...
+    def propensity(self, arm: str, exclude: set[str] | None = None) -> float: ...
+
+
+class _CountBasedBandit:
+    """Shared bookkeeping for epsilon-greedy and UCB1. Both track a running
+    per-arm pull count and reward sum, and both need every arm tried once
+    before their selection formula is even defined (UCB1's confidence bound
+    divides by pull count, epsilon-greedy has nothing to compare averages
+    against on arm zero).
+    """
+
+    def __init__(self, arms: list[str]) -> None:
+        self.arms = list(arms)
+        self._pulls = {arm: 0 for arm in self.arms}
+        self._reward_sums = {arm: 0.0 for arm in self.arms}
+
+    def _candidates(self, exclude: set[str] | None) -> list[str]:
+        if not exclude:
+            return self.arms
+        return [arm for arm in self.arms if arm not in exclude]
+
+    def _untried_arm(self, exclude: set[str] | None) -> str | None:
+        for arm in self._candidates(exclude):
+            if self._pulls[arm] == 0:
+                return arm
+        return None
+
+    def _mean_reward(self, arm: str) -> float:
+        return self._reward_sums[arm] / self._pulls[arm]
+
+    def update(self, arm: str, reward: float) -> None:
+        self._pulls[arm] += 1
+        self._reward_sums[arm] += reward
+
+    def _cold_start_propensity(self, arm: str, exclude: set[str] | None) -> float | None:
+        # returns None once every candidate arm has at least one pull, so
+        # the caller knows to fall through to its own formula. while an
+        # arm is still untried, select_arm() always returns the first
+        # untried one it finds, so the propensity for that specific arm
+        # is 1.0 and everything else is 0, not a probability distribution
+        # spread across the untried arms.
+        untried = self._untried_arm(exclude)
+        if untried is None:
+            return None
+        return 1.0 if arm == untried else 0.0
+
+
+class EpsilonGreedyBandit(_CountBasedBandit):
+    def __init__(self, arms: list[str], epsilon: float = 0.1) -> None:
+        super().__init__(arms)
+        self.epsilon = epsilon
+
+    def select_arm(self, exclude: set[str] | None = None) -> str:
+        untried = self._untried_arm(exclude)
+        if untried is not None:
+            return untried
+
+        candidates = self._candidates(exclude)
+
+        if random.random() < self.epsilon:
+            return random.choice(candidates)
+
+        best_mean = max(self._mean_reward(arm) for arm in candidates)
+        best_arms = [arm for arm in candidates if self._mean_reward(arm) == best_mean]
+        return random.choice(best_arms)
+
+    def propensity(self, arm: str, exclude: set[str] | None = None) -> float:
+        cold_start = self._cold_start_propensity(arm, exclude)
+        if cold_start is not None:
+            return cold_start
+
+        candidates = self._candidates(exclude)
+        if arm not in candidates:
+            return 0.0
+
+        best_mean = max(self._mean_reward(a) for a in candidates)
+        best_arms = [a for a in candidates if self._mean_reward(a) == best_mean]
+
+        # epsilon/n chance of landing here through the random branch, plus
+        # a share of the (1 - epsilon) greedy branch if this arm is one of
+        # the (possibly tied) best ones.
+        base = self.epsilon / len(candidates)
+        if arm in best_arms:
+            return base + (1 - self.epsilon) / len(best_arms)
+        return base
+
+
+class UCB1Bandit(_CountBasedBandit):
+    def select_arm(self, exclude: set[str] | None = None) -> str:
+        untried = self._untried_arm(exclude)
+        if untried is not None:
+            return untried
+
+        candidates = self._candidates(exclude)
+
+        # textbook UCB1, no tunable exploration constant on purpose, that's
+        # the whole point of this algorithm over epsilon-greedy: provable
+        # logarithmic regret without anything to hand-tune.
+        total_pulls = sum(self._pulls.values())
+
+        def ucb_score(arm: str) -> float:
+            return self._mean_reward(arm) + math.sqrt(2 * math.log(total_pulls) / self._pulls[arm])
+
+        best_score = max(ucb_score(arm) for arm in candidates)
+        best_arms = [arm for arm in candidates if ucb_score(arm) == best_score]
+        return random.choice(best_arms)
+
+    def propensity(self, arm: str, exclude: set[str] | None = None) -> float:
+        cold_start = self._cold_start_propensity(arm, exclude)
+        if cold_start is not None:
+            return cold_start
+
+        candidates = self._candidates(exclude)
+        if arm not in candidates:
+            return 0.0
+
+        total_pulls = sum(self._pulls.values())
+
+        def ucb_score(a: str) -> float:
+            return self._mean_reward(a) + math.sqrt(2 * math.log(total_pulls) / self._pulls[a])
+
+        best_score = max(ucb_score(a) for a in candidates)
+        best_arms = [a for a in candidates if ucb_score(a) == best_score]
+        # UCB1 is deterministic except for ties, so the propensity is just
+        # 1 over however many arms are tied for the top score.
+        return 1 / len(best_arms) if arm in best_arms else 0.0
+
+
+_THOMPSON_PROPENSITY_SAMPLES = 2000
+
+
+class ThompsonSamplingBandit:
+    def __init__(self, arms: list[str]) -> None:
+        self.arms = list(arms)
+        # Beta(1, 1) is a uniform prior. no cold-start branch needed here
+        # like the other two, a fresh arm just has wide uncertainty and
+        # naturally gets sampled a lot until its posterior narrows, that's
+        # Thompson sampling's actual advantage over the count-based ones.
+        self._alpha = {arm: 1.0 for arm in self.arms}
+        self._beta = {arm: 1.0 for arm in self.arms}
+
+    def select_arm(self, exclude: set[str] | None = None) -> str:
+        candidates = [arm for arm in self.arms if not exclude or arm not in exclude]
+        samples = {arm: np.random.beta(self._alpha[arm], self._beta[arm]) for arm in candidates}
+        return max(samples, key=samples.get)
+
+    def update(self, arm: str, reward: float) -> None:
+        # standard Beta-Bernoulli Thompson sampling assumes 0/1 rewards.
+        # ours are continuous in [0, 1], so the reward itself gets treated
+        # as a fractional pseudo-observation rather than a hard success or
+        # failure. this is the usual way to stretch Beta-Bernoulli to
+        # continuous rewards in that range.
+        self._alpha[arm] += reward
+        self._beta[arm] += 1 - reward
+
+    def propensity(self, arm: str, exclude: set[str] | None = None) -> float:
+        candidates = [a for a in self.arms if not exclude or a not in exclude]
+        if arm not in candidates:
+            return 0.0
+
+        # there's no closed form for the probability that one independent
+        # Beta draw beats a handful of others, so this estimates it the
+        # same way select_arm() actually decides: draw a batch of samples
+        # per arm and see how often this arm comes out on top. this is a
+        # Monte Carlo estimate, not an exact value, that's expected and
+        # documented, not a bug.
+        draws = np.stack(
+            [np.random.beta(self._alpha[a], self._beta[a], size=_THOMPSON_PROPENSITY_SAMPLES) for a in candidates]
+        )
+        winners = draws.argmax(axis=0)
+        return float((winners == candidates.index(arm)).mean())
+
+
+class RandomBandit(_CountBasedBandit):
+    """Uniform random arm selection, ignores everything it's learned. This
+    is the 'B' side of the A/B mode flag: hand this to a
+    ContextualBandit instead of one of the real algorithms and you get a
+    random-routing baseline to compare the real bandits against. update()
+    still records pulls and reward sums like the other count-based
+    bandits do, not because selection uses them, but so stats aggregation
+    can report what random selection would have earned per arm.
+    """
+
+    def select_arm(self, exclude: set[str] | None = None) -> str:
+        return random.choice(self._candidates(exclude))
+
+    def propensity(self, arm: str, exclude: set[str] | None = None) -> float:
+        candidates = self._candidates(exclude)
+        return 1 / len(candidates) if arm in candidates else 0.0
+
+
+class ContextualBandit:
+    """One independent bandit instance per context bucket, the 'disjoint'
+    contextual bandit approach, rather than one bandit shared across
+    every kind of request. Instances are created lazily the first time a
+    given context key shows up.
+    """
+
+    def __init__(self, algorithm_factory: Callable[[], Bandit], arms: list[str]) -> None:
+        self._algorithm_factory = algorithm_factory
+        self.arms = list(arms)
+        self._bandits: dict[str, Bandit] = {}
+
+    def _get_bandit(self, context_key: str) -> Bandit:
+        if context_key not in self._bandits:
+            self._bandits[context_key] = self._algorithm_factory()
+        return self._bandits[context_key]
+
+    def select_arm(self, context_key: str, exclude: set[str] | None = None) -> str:
+        return self._get_bandit(context_key).select_arm(exclude=exclude)
+
+    def update(self, context_key: str, arm: str, reward: float) -> None:
+        self._get_bandit(context_key).update(arm, reward)
+
+    def propensity(self, context_key: str, arm: str, exclude: set[str] | None = None) -> float:
+        return self._get_bandit(context_key).propensity(arm, exclude=exclude)