arcus-cli 0.1.0__py3-none-any.whl

arcus/cli.py

@@ -0,0 +1,335 @@
+import sys
+import threading
+from uuid import uuid4
+
+import typer
+from rich.console import Console
+from rich.table import Table
+from sqlmodel import Session, select
+
+from arcus.adapters.arc_adapter import ArcAdapter, ArcModel
+from arcus.cache.semantic_cache import lookup as cache_lookup
+from arcus.cache.semantic_cache import store as cache_store
+from arcus.config import ArcusConfig, BanditAlgorithm, config_path, load_config, save_config
+from arcus.embeddings import get_embedding_model
+from arcus.quality.gate import call_with_quality_gate
+from arcus.routing.bandit import (
+    Bandit,
+    ContextualBandit,
+    EpsilonGreedyBandit,
+    RandomBandit,
+    ThompsonSamplingBandit,
+    UCB1Bandit,
+)
+from arcus.routing.context import Context, classify
+from arcus.routing.warm_start import replay_history
+from arcus.storage.db import RequestLog, get_engine, log_request
+from arcus.storage.stats import aggregate_by_arm_and_mode
+
+ARMS = [m.value for m in ArcModel]
+
+_ALGORITHM_FACTORIES: dict[BanditAlgorithm, "type[Bandit]"] = {
+    "epsilon_greedy": EpsilonGreedyBandit,
+    "ucb1": UCB1Bandit,
+    "thompson": ThompsonSamplingBandit,
+}
+
+
+def main(argv: list[str] | None = None) -> None:
+    """Entry point for the `arcus` command. Dispatches by hand on argv[0]
+    rather than using typer's subcommand machinery. the whole point is for
+    `arcus "some question"` to just work with no subcommand at all, and a
+    real subcommand parser (click underneath typer) fights that: it wants
+    to treat the prompt text itself as an unrecognized command. `stats` is
+    the one reserved word, everything else is prompt text.
+    """
+    argv = list(sys.argv[1:] if argv is None else argv)
+
+    if argv and argv[0] == "stats":
+        run_stats()
+        return
+
+    if argv and argv[0] == "chat":
+        random_mode = "--random" in argv
+        run_chat(random_mode=random_mode)
+        return
+
+    random_mode = "--random" in argv
+    argv = [arg for arg in argv if arg != "--random"]
+
+    prompt = _build_prompt(argv)
+    if not prompt:
+        _print_usage()
+        return
+
+    run_ask(prompt, random_mode=random_mode)
+
+
+def _build_prompt(args: list[str]) -> str:
+    # if stdin isn't a tty, something got piped in (a traceback, a test
+    # failure, whatever), treat that as the main
+    # content and let any positional args add an explicit instruction
+    # on top of it, e.g. `python broken.py 2>&1 | arcus "why is this
+    # failing"`.
+    piped = sys.stdin.read().strip() if not sys.stdin.isatty() else ""
+    positional = " ".join(args).strip()
+
+    if piped and positional:
+        return f"{piped}\n\n{positional}"
+    return piped or positional
+
+
+def _print_usage() -> None:
+    Console().print('usage: arcus "<question>"   or   arcus chat   or   arcus stats')
+
+
+def _ensure_config() -> ArcusConfig:
+    try:
+        return load_config()
+    except FileNotFoundError:
+        return _run_setup_wizard()
+
+
+def _run_setup_wizard() -> ArcusConfig:
+    console = Console()
+    console.print("[bold]no ARC key found, let's get you set up.[/bold]")
+    console.print(
+        "grab one from llm.arc.vt.edu under "
+        "User profile > Settings > Account > API keys.\n"
+    )
+
+    api_key = typer.prompt("ARC API key", hide_input=True)
+
+    console.print("checking that key works...")
+    try:
+        ArcAdapter(api_key=api_key).chat(
+            ArcModel.GPT_OSS_120B,
+            [{"role": "user", "content": "reply with just the word ok"}],
+            max_tokens=5,
+        )
+    except Exception as e:
+        console.print(f"[red]that key didn't work:[/red] {e}")
+        raise SystemExit(1) from e
+
+    config = ArcusConfig(arc_api_key=api_key)
+    save_config(config)
+    console.print(f"[green]saved to {config_path()}[/green]\n")
+    return config
+
+
+def run_ask(prompt: str, random_mode: bool = False) -> None:
+    console = Console()
+
+    # the embedding model backs both the context classifier's fallback
+    # path and the semantic cache lookup below, and loading it costs a
+    # few seconds the first time any process touches it. kick that load
+    # off now, in the background, so it's warm (or at least warming) by
+    # the time either one actually calls embed(), instead of eating that
+    # cost inline and in serial with everything else.
+    #
+    # this has to start AFTER building the ArcAdapter, not before.
+    # constructing the openai client is the first time this process
+    # touches httpx internals, and openai does that lazily, on client
+    # construction, not on import. starting the embedding thread first
+    # let it race that first-time httpx touch against sentence-transformers'
+    # own (torch/huggingface_hub) import chain, which also reaches into
+    # httpx, and that produced a real, reliably reproducible crash:
+    # "partially initialized module 'httpx' ... circular import". building
+    # the adapter first means the main thread finishes touching httpx
+    # before any second thread gets a chance to.
+    config = _ensure_config()
+    engine = get_engine()
+    adapter = ArcAdapter(api_key=config.arc_api_key)
+
+    threading.Thread(target=get_embedding_model, daemon=True).start()
+
+    context = classify(prompt)
+    mode = "random" if random_mode else "bandit"
+
+    cached = cache_lookup(prompt, engine=engine)
+    if cached.hit:
+        console.print(cached.response)
+        log_request(
+            prompt=prompt,
+            task_type=context.task_type.value,
+            length_bucket=context.length_bucket.value,
+            model=cached.model,
+            cache_hit=True,
+            quality_passed=True,
+            engine=engine,
+        )
+        return
+
+    algorithm_factory = RandomBandit if random_mode else _ALGORITHM_FACTORIES[config.bandit_algorithm]
+    bandit = ContextualBandit(lambda: algorithm_factory(ARMS), arms=ARMS)
+    # rebuild what this bandit already learned from past requests in this
+    # same mode, otherwise every invocation starts back at zero since
+    # there's no daemon holding it in memory between runs.
+    replay_history(bandit, engine, mode=mode)
+
+    content, model_used, passed = _route_and_answer(
+        adapter, bandit, context, prompt, [{"role": "user", "content": prompt}], mode, engine
+    )
+
+    console.print(content if content else "[red]no usable response from any model.[/red]")
+
+    if passed and content:
+        cache_store(prompt, content, model=model_used, engine=engine)
+
+
+def _route_and_answer(
+    adapter: ArcAdapter,
+    bandit: ContextualBandit,
+    context: Context,
+    prompt: str,
+    messages: list[dict],
+    mode: str,
+    engine,
+    conversation_id: str | None = None,
+    turn_index: int | None = None,
+) -> tuple[str | None, str, bool]:
+    """Runs one turn through the quality gate, logs every attempt, and
+    returns (content, model_used, passed). content can be non-None even
+    when passed is False (the last attempt still produced text, it just
+    didn't clear the gate); content is None only when every arm errored
+    out with nothing to show at all.
+    """
+    outcome = call_with_quality_gate(adapter, bandit, context.key, messages)
+
+    for attempt in outcome.attempts:
+        log_request(
+            prompt=prompt,
+            task_type=context.task_type.value,
+            length_bucket=context.length_bucket.value,
+            model=attempt.model,
+            propensity=attempt.propensity,
+            latency_ms=attempt.latency_ms,
+            mode=mode,
+            reward=attempt.reward,
+            quality_passed=attempt.passed,
+            conversation_id=conversation_id,
+            turn_index=turn_index,
+            engine=engine,
+        )
+
+    # response is None when every arm errored out at the API level (see
+    # call_with_quality_gate), not just returned a bad answer
+    content = outcome.response.choices[0].message.content if outcome.response else None
+    return content, outcome.model_used, outcome.passed
+
+
+_MAX_CHAT_MESSAGES = 20  # ~10 exchanges, a first-guess cap like the
+# reward weights, not tuned against anything yet
+
+
+def _trim_history(messages: list[dict], max_messages: int = _MAX_CHAT_MESSAGES) -> list[dict]:
+    if len(messages) <= max_messages:
+        return messages
+    # drop whole oldest turns (user+assistant pairs), not a single
+    # message, so the transcript never gets left with an orphaned
+    # question and no answer in front of it
+    excess = len(messages) - max_messages
+    excess += excess % 2
+    return messages[excess:]
+
+
+def run_chat(random_mode: bool = False) -> None:
+    console = Console()
+
+    config = _ensure_config()
+    engine = get_engine()
+    adapter = ArcAdapter(api_key=config.arc_api_key)  # before the thread, same reason as run_ask
+
+    threading.Thread(target=get_embedding_model, daemon=True).start()
+
+    mode = "random" if random_mode else "bandit"
+    algorithm_factory = RandomBandit if random_mode else _ALGORITHM_FACTORIES[config.bandit_algorithm]
+    bandit = ContextualBandit(lambda: algorithm_factory(ARMS), arms=ARMS)
+    replay_history(bandit, engine, mode=mode)
+
+    conversation_id = str(uuid4())
+    messages: list[dict] = []
+    turn_index = 0
+
+    console.print("[bold]chatting with arcus, type 'exit' or ctrl-d to leave.[/bold]\n")
+
+    while True:
+        try:
+            user_input = input("you: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            console.print()
+            break
+
+        if not user_input:
+            continue
+        if user_input.lower() in ("exit", "quit"):
+            break
+
+        messages.append({"role": "user", "content": user_input})
+        context = classify(user_input)
+
+        content, model_used, passed = _route_and_answer(
+            adapter,
+            bandit,
+            context,
+            user_input,
+            messages,
+            mode,
+            engine,
+            conversation_id=conversation_id,
+            turn_index=turn_index,
+        )
+
+        if content:
+            messages.append({"role": "assistant", "content": content})
+            console.print(f"[green]arcus ({model_used}):[/green] {content}\n")
+        else:
+            # nothing usable came back, don't leave an unanswered
+            # question sitting in history for the next turn to trip over
+            messages.pop()
+            console.print("[red]no usable response from any model, try rephrasing.[/red]\n")
+
+        messages = _trim_history(messages)
+        turn_index += 1
+
+
+def run_stats() -> None:
+    console = Console()
+    engine = get_engine()
+
+    with Session(engine) as session:
+        rows = session.exec(select(RequestLog)).all()
+
+    if not rows:
+        console.print("no requests logged yet, go ask arcus something.")
+        return
+
+    table = Table(title="arcus stats")
+    table.add_column("model")
+    table.add_column("mode")
+    table.add_column("requests", justify="right")
+    table.add_column("avg reward", justify="right")
+    table.add_column("avg latency (ms)", justify="right")
+    table.add_column("cost score", justify="right")
+
+    for summary in aggregate_by_arm_and_mode(engine):
+        table.add_row(
+            summary.model,
+            summary.mode,
+            str(summary.request_count),
+            f"{summary.avg_reward:.3f}" if summary.avg_reward is not None else "-",
+            f"{summary.avg_latency_ms:.0f}" if summary.avg_latency_ms is not None else "-",
+            f"{summary.cost_score:.2f}" if summary.cost_score is not None else "-",
+        )
+
+    console.print(table)
+
+    cache_hits = sum(1 for row in rows if row.cache_hit)
+    gate_catches = sum(1 for row in rows if not row.quality_passed)
+
+    console.print(f"\ncache hit rate: {cache_hits / len(rows):.1%} ({cache_hits}/{len(rows)})")
+    console.print(f"quality gate catches: {gate_catches} attempt(s) failed and got retried")
+
+
+if __name__ == "__main__":
+    main()

arcus/config.py

@@ -0,0 +1,46 @@
+import tomllib
+from pathlib import Path
+from typing import Literal
+
+from platformdirs import user_config_dir
+from pydantic_settings import BaseSettings
+
+BanditAlgorithm = Literal["epsilon_greedy", "ucb1", "thompson"]
+
+
+class ArcusConfig(BaseSettings):
+    arc_api_key: str
+    # thompson sampling is the default, it's the one that needs the least
+    # hand-tuning (no epsilon to pick) and adapts the fastest early on.
+    # the algorithm is swappable here rather than hardcoded, so a user
+    # can pick a different one without touching code.
+    bandit_algorithm: BanditAlgorithm = "thompson"
+
+
+def config_path() -> Path:
+    # computed inside a function rather than a module constant so tests
+    # can point it somewhere disposable the same way storage/db.py's data
+    # dir already gets monkeypatched.
+    return Path(user_config_dir("arcus")) / "config.toml"
+
+
+def load_config(path: Path | None = None) -> ArcusConfig:
+    path = path or config_path()
+    with open(path, "rb") as f:
+        data = tomllib.load(f)
+    return ArcusConfig(**data)
+
+
+def save_config(config: ArcusConfig, path: Path | None = None) -> None:
+    path = path or config_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    # only two fields right now, not worth pulling in a TOML-writing
+    # dependency just to serialize a couple of key/value lines.
+    path.write_text(
+        f'arc_api_key = "{config.arc_api_key}"\n'
+        f'bandit_algorithm = "{config.bandit_algorithm}"\n'
+    )
+    # the API key lives on disk in plain text, chmod 600 so it's at least
+    # not readable by other users on the same machine.
+    path.chmod(0o600)

arcus/embeddings.py

@@ -0,0 +1,34 @@
+import threading
+
+import numpy as np
+
+_model_cache = None
+_model_lock = threading.Lock()
+
+
+def get_embedding_model():
+    # imported lazily so importing this module doesn't drag torch in for
+    # callers that never actually need embeddings. cached as a singleton
+    # once loaded since the load itself (not the import) is the slow
+    # part, worth paying once per process, not once per call site.
+    #
+    # the lock matters once the CLI started warming this up on a
+    # background thread (see cli.py): without it, a foreground call
+    # landing while the warm-up thread is mid-load would see an empty
+    # cache and kick off a second, redundant load instead of just
+    # waiting for the first one to finish.
+    global _model_cache
+    if _model_cache is not None:
+        return _model_cache
+
+    with _model_lock:
+        if _model_cache is None:
+            from sentence_transformers import SentenceTransformer
+
+            _model_cache = SentenceTransformer("all-MiniLM-L6-v2")
+    return _model_cache
+
+
+def embed(texts: list[str]) -> np.ndarray:
+    model = get_embedding_model()
+    return model.encode(texts, normalize_embeddings=True)

arcus/eval/offline.py

@@ -0,0 +1,219 @@
+import random
+from dataclasses import dataclass
+from statistics import mean
+from typing import Callable
+
+from sqlmodel import Session, select
+
+from arcus.storage.db import RequestLog
+
+Policy = Callable[[str], str]
+
+
+@dataclass(frozen=True)
+class LoggedExample:
+    context_key: str
+    arm: str
+    propensity: float
+    reward: float
+
+
+def load_logged_examples(engine, mode: str = "bandit") -> list[LoggedExample]:
+    """Pulls rows out of the request log that are actually usable for
+    offline evaluation: propensity has to be set (the whole point of
+    logging it from day one) and reward has to be set (a row that errored
+    out before a reward was computed can't tell an estimator anything).
+    """
+    with Session(engine) as session:
+        rows = session.exec(
+            select(RequestLog)
+            .where(RequestLog.mode == mode)
+            .where(RequestLog.propensity.is_not(None))
+            .where(RequestLog.reward.is_not(None))
+        ).all()
+
+    return [
+        LoggedExample(
+            context_key=f"{row.task_type}:{row.length_bucket}",
+            arm=row.model,
+            propensity=row.propensity,
+            reward=row.reward,
+        )
+        for row in rows
+    ]
+
+
+def policy_always(arm: str) -> Policy:
+    return lambda context_key: arm
+
+
+def greedy_policy_from_log(examples: list[LoggedExample]) -> Policy:
+    """Builds a deterministic policy that, for each context bucket, picks
+    whichever arm had the best average logged reward. Fit and evaluated
+    on the same data (no held-out split), so its reported value will run
+    a little optimistic, this is a known limitation of the simple
+    version of this technique, not a bug. A proper cross-fitted version
+    would split the log before fitting.
+    """
+    sums: dict[tuple[str, str], float] = {}
+    counts: dict[tuple[str, str], int] = {}
+    for ex in examples:
+        key = (ex.context_key, ex.arm)
+        sums[key] = sums.get(key, 0.0) + ex.reward
+        counts[key] = counts.get(key, 0) + 1
+
+    best_arm_per_context: dict[str, str] = {}
+    for (context_key, arm), total in sums.items():
+        avg = total / counts[(context_key, arm)]
+        current_best = best_arm_per_context.get(context_key)
+        if current_best is None or avg > sums[(context_key, current_best)] / counts[(context_key, current_best)]:
+            best_arm_per_context[context_key] = arm
+
+    fallback_arm = max(counts, key=counts.get)[1] if counts else None
+
+    def policy(context_key: str) -> str:
+        return best_arm_per_context.get(context_key, fallback_arm)
+
+    return policy
+
+
+def ips_estimate(examples: list[LoggedExample], policy: Policy) -> float:
+    """Inverse propensity scoring. For each logged row, if the target
+    policy would have picked the same arm the logging policy actually
+    picked, that row's reward counts, reweighted by how unlikely the
+    logging policy was to pick it (rarer picks get boosted more, that's
+    what corrects for the logging policy's own selection bias). Rows
+    where the policies disagree contribute nothing, we simply never
+    observed what would have happened.
+    """
+    if not examples:
+        return 0.0
+
+    weighted = [
+        ex.reward / ex.propensity if policy(ex.context_key) == ex.arm else 0.0
+        for ex in examples
+    ]
+    return mean(weighted)
+
+
+def _fit_reward_model(examples: list[LoggedExample]) -> dict[tuple[str, str], float]:
+    groups: dict[tuple[str, str], list[float]] = {}
+    for ex in examples:
+        groups.setdefault((ex.context_key, ex.arm), []).append(ex.reward)
+    return {key: mean(rewards) for key, rewards in groups.items()}
+
+
+def dr_estimate(
+    examples: list[LoggedExample],
+    policy: Policy,
+    reward_model: dict[tuple[str, str], float] | None = None,
+) -> float:
+    """Doubly robust estimator: a direct-method estimate (the reward
+    model's guess at what the target policy would earn) plus an IPS
+    correction term that only kicks in on rows where the logged arm
+    happens to match the target policy. Stays unbiased if either the
+    reward model or the propensities are right, not both, that's the
+    "doubly" part. Same same-data-fit caveat as greedy_policy_from_log
+    applies to the default reward model here.
+    """
+    if not examples:
+        return 0.0
+
+    reward_model = reward_model if reward_model is not None else _fit_reward_model(examples)
+    overall_mean = mean(ex.reward for ex in examples)
+
+    total = 0.0
+    for ex in examples:
+        target_arm = policy(ex.context_key)
+        q_target = reward_model.get((ex.context_key, target_arm), overall_mean)
+
+        if target_arm == ex.arm:
+            q_logged = reward_model.get((ex.context_key, ex.arm), overall_mean)
+            correction = (ex.reward - q_logged) / ex.propensity
+        else:
+            correction = 0.0
+
+        total += q_target + correction
+
+    return total / len(examples)
+
+
+def bootstrap_ci(
+    examples: list[LoggedExample],
+    estimator: Callable[[list[LoggedExample]], float],
+    n_resamples: int = 1000,
+    confidence: float = 0.95,
+    seed: int | None = None,
+) -> tuple[float, float, float]:
+    """Percentile bootstrap: resample the log with replacement a bunch of
+    times, run the estimator on each resample, and read the confidence
+    interval off the sorted results. Plain stdlib random, no scipy, the
+    log sizes here don't need anything fancier.
+    """
+    point = estimator(examples)
+    if not examples:
+        return point, point, point
+
+    rng = random.Random(seed)
+    n = len(examples)
+    resampled = [estimator([examples[rng.randrange(n)] for _ in range(n)]) for _ in range(n_resamples)]
+    resampled.sort()
+
+    lower_idx = int((1 - confidence) / 2 * n_resamples)
+    upper_idx = min(int((1 + confidence) / 2 * n_resamples), n_resamples - 1)
+
+    return point, resampled[lower_idx], resampled[upper_idx]
+
+
+@dataclass(frozen=True)
+class PolicyEvaluation:
+    name: str
+    ips_estimate: float
+    ips_ci: tuple[float, float]
+    dr_estimate: float
+    dr_ci: tuple[float, float]
+
+
+def evaluate_policies(
+    engine,
+    policies: dict[str, Policy],
+    mode: str = "bandit",
+    n_resamples: int = 1000,
+    seed: int | None = None,
+) -> list[PolicyEvaluation]:
+    """Compares the logged policy against offline-estimated alternatives,
+    each with a confidence interval.
+    """
+    examples = load_logged_examples(engine, mode=mode)
+
+    logged_point, logged_low, logged_high = bootstrap_ci(
+        examples, lambda ex: mean(e.reward for e in ex) if ex else 0.0, n_resamples, seed=seed
+    )
+    rows = [
+        PolicyEvaluation(
+            name="logged (as-run)",
+            ips_estimate=logged_point,
+            ips_ci=(logged_low, logged_high),
+            dr_estimate=logged_point,
+            dr_ci=(logged_low, logged_high),
+        )
+    ]
+
+    for name, policy in policies.items():
+        ips_point, ips_low, ips_high = bootstrap_ci(
+            examples, lambda ex, p=policy: ips_estimate(ex, p), n_resamples, seed=seed
+        )
+        dr_point, dr_low, dr_high = bootstrap_ci(
+            examples, lambda ex, p=policy: dr_estimate(ex, p), n_resamples, seed=seed
+        )
+        rows.append(
+            PolicyEvaluation(
+                name=name,
+                ips_estimate=ips_point,
+                ips_ci=(ips_low, ips_high),
+                dr_estimate=dr_point,
+                dr_ci=(dr_low, dr_high),
+            )
+        )
+
+    return rows