arcus-cli 0.1.0__py3-none-any.whl

arcus/routing/context.py

@@ -0,0 +1,218 @@
+import re
+from dataclasses import dataclass
+from enum import Enum
+
+
+class TaskType(str, Enum):
+    CODE = "code"
+    REASONING_MATH = "reasoning_math"
+    WRITING = "writing"
+    LONG_DOCUMENT = "long_document"
+    GENERAL = "general"
+
+
+class LengthBucket(str, Enum):
+    SHORT = "short"
+    MEDIUM = "medium"
+    LONG = "long"
+    VERY_LONG = "very_long"
+
+
+_CODE_MARKERS = re.compile(
+    r"```"
+    r"|traceback \(most recent call last\)"
+    r"|\b(syntax|type|name|index|key|attribute|value|indentation)error\b"
+    r"|\bexception in thread\b"
+    r"|\bnpm err!"
+    r"|^\s*(def|class|import|from|function|public|private|const|let)\s",
+    re.IGNORECASE | re.MULTILINE,
+)
+
+_MATH_MARKERS = re.compile(
+    r"\b(solve|prove|derivative|integral|equation|theorem|factorize|simplify)\b"
+    r"|\btime complexity\b|\bbig[- ]o\b"
+    r"|solve for [a-z]\b",
+    re.IGNORECASE,
+)
+
+_DOC_MARKERS = re.compile(
+    r"\bsummariz(e|ing)\b"
+    r"|\btl;?dr\b"
+    r"|\bthe following (document|article|text|paper)\b"
+    r"|\bgiven the (text|document) below\b",
+    re.IGNORECASE,
+)
+
+_WRITING_MARKERS = re.compile(
+    r"\b(write|draft|compose)\s+(a|an|me a)\b"
+    r"|\b(essay|blog post|cover letter|short story|poem)\b",
+    re.IGNORECASE,
+)
+
+# anchor prompts for the embedding fallback below. GENERAL gets its own
+# examples too, not just the other four, otherwise the fallback would
+# always pick one of code/math/writing/doc even for plain small talk
+# since there'd be nothing else in the running.
+_SIMILARITY_THRESHOLD = 0.35
+
+_TASK_TYPE_EXAMPLES: dict[TaskType, list[str]] = {
+    TaskType.CODE: [
+        "why is my for loop not terminating",
+        "getting a null pointer exception in this function",
+        "how do I fix this compiler error",
+        "review this pull request for bugs",
+        "my api call keeps returning a 500",
+    ],
+    TaskType.REASONING_MATH: [
+        "walk me through the proof of this theorem",
+        "what's the expected value of this dice roll",
+        "help me figure out the recurrence relation for this algorithm",
+        "is this argument logically valid",
+        "explain why this inequality holds",
+    ],
+    TaskType.WRITING: [
+        "help me phrase this paragraph better",
+        "can you make this sound more persuasive",
+        "give me a catchy title for this post",
+        "rewrite this in a more formal tone",
+        "I need a few opening lines for a speech",
+    ],
+    TaskType.LONG_DOCUMENT: [
+        "here's a long report, pull out the key takeaways",
+        "go through this contract and flag anything unusual",
+        "condense this research paper into a few bullet points",
+        "what are the main arguments made across these pages",
+    ],
+    TaskType.GENERAL: [
+        "what's a good movie to watch tonight",
+        "how's the weather looking this weekend",
+        "tell me something interesting",
+        "what time zone is Tokyo in",
+        "hey, how are you",
+    ],
+}
+
+_embedding_classifier_cache = None
+
+
+def _get_embedding_classifier():
+    # unlike the db engine, this one is worth caching as a real
+    # singleton. loading the model is the expensive part (a few hundred
+    # ms to a couple seconds), so paying it once per process is the
+    # whole point. the model itself comes from arcus.embeddings, shared
+    # with the semantic cache, so a request that needs both this fallback
+    # and a cache lookup only pays that load cost once, not twice.
+    global _embedding_classifier_cache
+    if _embedding_classifier_cache is not None:
+        return _embedding_classifier_cache
+
+    from arcus.embeddings import embed
+
+    labels: list[TaskType] = []
+    examples: list[str] = []
+    for task_type, prompts in _TASK_TYPE_EXAMPLES.items():
+        for prompt in prompts:
+            labels.append(task_type)
+            examples.append(prompt)
+
+    anchor_embeddings = embed(examples)
+
+    _embedding_classifier_cache = (labels, anchor_embeddings)
+    return _embedding_classifier_cache
+
+
+def _classify_by_embedding(text: str) -> TaskType:
+    try:
+        labels, anchor_embeddings = _get_embedding_classifier()
+        from arcus.embeddings import embed
+
+        query_embedding = embed([text])[0]
+    except Exception:
+        # no internet on first run, weights not cached yet, whatever the
+        # reason, the regex rules already did their best. don't let a
+        # missing model take classification down entirely.
+        return TaskType.GENERAL
+
+    similarities = anchor_embeddings @ query_embedding
+    best_idx = int(similarities.argmax())
+
+    if similarities[best_idx] < _SIMILARITY_THRESHOLD:
+        return TaskType.GENERAL
+
+    return labels[best_idx]
+
+
+def _looks_like_code(text: str) -> bool:
+    return bool(_CODE_MARKERS.search(text))
+
+
+def _looks_like_math(text: str) -> bool:
+    return bool(_MATH_MARKERS.search(text))
+
+
+def _looks_like_document_task(text: str) -> bool:
+    return bool(_DOC_MARKERS.search(text))
+
+
+def _looks_like_writing_request(text: str) -> bool:
+    return bool(_WRITING_MARKERS.search(text))
+
+
+def classify_task_type(text: str) -> TaskType:
+    # order matters here. code and error output are the most confident
+    # signal we ever get (and piping errors in is the single most common
+    # thing this tool will see), so that gets checked first no matter
+    # what else is in the prompt.
+    if _looks_like_code(text):
+        return TaskType.CODE
+    if _looks_like_math(text):
+        return TaskType.REASONING_MATH
+    if _looks_like_document_task(text):
+        return TaskType.LONG_DOCUMENT
+    if _looks_like_writing_request(text):
+        return TaskType.WRITING
+
+    # nothing hit a keyword. rather than giving up and calling it
+    # GENERAL, check how close it reads to labeled examples of each
+    # category, catches phrasing the regex rules just don't cover.
+    result = _classify_by_embedding(text)
+
+    # a long wall of text with no other signal is still a document task
+    # even if the embedding read came back wishy-washy about it, raw
+    # length beats a low-confidence embedding guess here.
+    if result == TaskType.GENERAL and bucket_length(text) in (LengthBucket.LONG, LengthBucket.VERY_LONG):
+        return TaskType.LONG_DOCUMENT
+
+    return result
+
+
+def bucket_length(text: str) -> LengthBucket:
+    # rough chars-per-token estimate (~4), not a real tokenizer. none of
+    # ARC's four models use OpenAI's tokenizer anyway so an exact count
+    # would just be fake precision. these thresholds are a first guess,
+    # worth retuning once we have real query lengths to look at.
+    approx_tokens = len(text) // 4
+
+    if approx_tokens < 100:
+        return LengthBucket.SHORT
+    if approx_tokens < 500:
+        return LengthBucket.MEDIUM
+    if approx_tokens < 2000:
+        return LengthBucket.LONG
+    return LengthBucket.VERY_LONG
+
+
+@dataclass(frozen=True)
+class Context:
+    task_type: TaskType
+    length_bucket: LengthBucket
+
+    @property
+    def key(self) -> str:
+        # this is the dict key the per-context bandits key off of, one
+        # bandit instance per bucket.
+        return f"{self.task_type.value}:{self.length_bucket.value}"
+
+
+def classify(text: str) -> Context:
+    return Context(task_type=classify_task_type(text), length_bucket=bucket_length(text))

arcus/routing/reward.py

@@ -0,0 +1,91 @@
+from dataclasses import dataclass
+
+# published commercial hosting rates for these open-weight models, one
+# primary source per model, as of August 2026. this is a SIMULATED cost
+# signal, ARC itself is free to the user, nobody's actually being
+# billed these numbers. rates are dollars per 1M tokens.
+#
+# gpt-oss-120b: Groq
+# GLM-5.3: Zhipu/Z.ai official pricing
+# Kimi-K3: Moonshot AI official pricing (cache-miss rate, ignoring the
+#   cache-hit discount, more precision than a simulated signal needs)
+# DeepSeek-V4-Flash: DeepSeek official pricing (off-peak rate, ignoring
+#   the peak-hours surcharge for the same reason)
+MODEL_HOSTING_RATES: dict[str, tuple[float, float]] = {
+    "gpt-oss-120b": (0.15, 0.60),
+    "GLM-5.3": (1.40, 4.40),
+    "Kimi-K3": (3.00, 15.00),
+    "DeepSeek-V4-Flash": (0.22, 0.66),
+}
+
+
+def _blended_rate(model: str) -> float:
+    input_rate, output_rate = MODEL_HOSTING_RATES[model]
+    return (input_rate + output_rate) / 2
+
+
+# computed from the table above rather than hardcoded, so the raw rates
+# stay the actual source of truth. cheapest model scores 1.0, priciest
+# scores 0.0, everything else falls in between.
+def _compute_cost_scores() -> dict[str, float]:
+    blended = {model: _blended_rate(model) for model in MODEL_HOSTING_RATES}
+    cheapest = min(blended.values())
+    priciest = max(blended.values())
+    spread = priciest - cheapest
+
+    return {model: 1 - (rate - cheapest) / spread for model, rate in blended.items()}
+
+
+COST_SCORES = _compute_cost_scores()
+
+
+def normalize_latency(latency_ms: float, ceiling_ms: float = 20_000) -> float:
+    # past the ceiling, slower is just uniformly bad, no reason to keep
+    # penalizing harder. 20s is a reasonable worst-case wait for an
+    # interactive CLI tool.
+    return max(0.0, min(1.0, 1 - latency_ms / ceiling_ms))
+
+
+def basic_quality_score(finish_reason: str | None, content: str | None) -> float:
+    """A minimal quality signal: finished cleanly (finish_reason == "stop")
+    and non-empty content. The real request pipeline uses the fuller
+    checks in quality/gate.py instead.
+    """
+    if finish_reason != "stop":
+        return 0.0
+    if not content or not content.strip():
+        return 0.0
+    return 1.0
+
+
+@dataclass(frozen=True)
+class RewardWeights:
+    quality: float
+    latency: float
+    cost: float
+
+
+# quality weighted highest: a fast, cheap, wrong answer is worth about
+# nothing. latency next: it's the most immediately felt cost on shared,
+# sometimes-contended infra, where quality can degrade during peak load.
+# cost lowest: ARC is free to the user, this dimension exists to
+# demonstrate routing that accounts for cost as an engineering practice,
+# not because real budget pressure exists here.
+DEFAULT_WEIGHTS = RewardWeights(quality=0.5, latency=0.3, cost=0.2)
+
+
+def compute_reward(
+    latency_ms: float,
+    model: str,
+    quality_score: float,
+    weights: RewardWeights = DEFAULT_WEIGHTS,
+) -> float:
+    weight_sum = weights.quality + weights.latency + weights.cost
+    if abs(weight_sum - 1.0) > 1e-6:
+        raise ValueError(f"reward weights must sum to 1.0, got {weight_sum}")
+
+    return (
+        weights.quality * quality_score
+        + weights.latency * normalize_latency(latency_ms)
+        + weights.cost * COST_SCORES[model]
+    )

arcus/routing/warm_start.py

@@ -0,0 +1,35 @@
+from sqlmodel import Session, select
+
+from arcus.routing.bandit import ContextualBandit
+from arcus.storage.db import RequestLog
+
+
+def replay_history(bandit: ContextualBandit, engine, mode: str) -> None:
+    """Rebuilds a bandit's learned state from past requests in the log.
+
+    Every `arcus` invocation is a fresh process, there's no daemon keeping
+    the bandit alive in memory between runs. Without this, each run would
+    start from a blank slate and the router would never actually learn
+    anything. It works because update() on all four algorithms is just an
+    associative accumulation of pull counts and reward sums, so replaying
+    the log in chronological order and feeding each row back through
+    update() lands on the same state as if the process had been running
+    continuously the whole time.
+    """
+    with Session(engine) as session:
+        rows = session.exec(
+            select(RequestLog)
+            .where(RequestLog.mode == mode)
+            .where(RequestLog.reward.is_not(None))
+            .order_by(RequestLog.created_at)
+        ).all()
+
+    for row in rows:
+        if row.model not in bandit.arms:
+            # a model that was live when this row was logged but has
+            # since been retired or renamed by ARC. no arm to credit the
+            # reward to anymore, skip it rather than crash the whole
+            # replay over one stale row.
+            continue
+        context_key = f"{row.task_type}:{row.length_bucket}"
+        bandit.update(context_key, row.model, row.reward)

arcus/storage/db.py

@@ -0,0 +1,117 @@
+import os
+from datetime import UTC, datetime
+from pathlib import Path
+
+from platformdirs import user_data_dir
+from sqlmodel import Field, Session, SQLModel, create_engine
+
+
+class RequestLog(SQLModel, table=True):
+    id: int | None = Field(default=None, primary_key=True)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+    prompt: str
+    task_type: str
+    length_bucket: str
+    model: str
+
+    # nullable since a cache hit skips the bandit and quality gate
+    # entirely, but propensity in particular has to be a column from day
+    # one: there's no way to go back and add propensity logging to rows
+    # that already happened.
+    propensity: float | None = Field(default=None)
+    latency_ms: float | None = Field(default=None)
+    finish_reason: str | None = Field(default=None)
+    error: str | None = Field(default=None)
+
+    # "bandit" or "random", whichever ContextualBandit actually drove
+    # this request. mode plus model is the whole point of the A/B stats
+    # comparison.
+    mode: str | None = Field(default=None)
+    # the actual compute_reward() output at request time, stored rather
+    # than recomputed later. if reward.py's weights or cost table change
+    # down the line, historical stats should still reflect what the
+    # bandit was actually updated with, not whatever the weights happen
+    # to be when someone runs `arcus stats`.
+    reward: float | None = Field(default=None)
+
+    # was this response served from the semantic cache instead of an
+    # actual ARC call. needed for `arcus stats` to report a cache hit
+    # rate, the current schema has no other way to tell a cache hit apart
+    # from a normal request.
+    cache_hit: bool = Field(default=False)
+    # did this specific attempt pass the quality gate. combined with
+    # `model`, this is what makes a per-model catch rate possible instead
+    # of only ever seeing which model eventually succeeded.
+    quality_passed: bool = Field(default=True)
+
+    # groups every turn of one `arcus chat` session together. null for
+    # one-shot `arcus "..."` calls, there's no session to group those into.
+    conversation_id: str | None = Field(default=None, index=True)
+    # position of this turn within its conversation, 0-indexed. null
+    # alongside conversation_id for one-shot calls.
+    turn_index: int | None = Field(default=None)
+
+
+def _database_url() -> str:
+    override = os.environ.get("ARCUS_DATABASE_URL")
+    if override:
+        return override
+
+    data_dir = Path(user_data_dir("arcus"))
+    data_dir.mkdir(parents=True, exist_ok=True)
+    return f"sqlite:///{data_dir / 'arcus.db'}"
+
+
+def get_engine():
+    # deliberately not cached as a module-level singleton. sqlite engine
+    # creation is cheap, and skipping the singleton means there's no
+    # stale global state to worry about resetting between tests (or
+    # between one CLI run and the next, if ARCUS_DATABASE_URL changes).
+    engine = create_engine(_database_url())
+    SQLModel.metadata.create_all(engine)
+    return engine
+
+
+def log_request(
+    prompt: str,
+    task_type: str,
+    length_bucket: str,
+    model: str,
+    propensity: float | None = None,
+    latency_ms: float | None = None,
+    finish_reason: str | None = None,
+    error: str | None = None,
+    mode: str | None = None,
+    reward: float | None = None,
+    cache_hit: bool = False,
+    quality_passed: bool = True,
+    conversation_id: str | None = None,
+    turn_index: int | None = None,
+    engine=None,
+) -> RequestLog:
+    engine = engine or get_engine()
+
+    entry = RequestLog(
+        prompt=prompt,
+        task_type=task_type,
+        length_bucket=length_bucket,
+        model=model,
+        propensity=propensity,
+        latency_ms=latency_ms,
+        finish_reason=finish_reason,
+        error=error,
+        mode=mode,
+        reward=reward,
+        cache_hit=cache_hit,
+        quality_passed=quality_passed,
+        conversation_id=conversation_id,
+        turn_index=turn_index,
+    )
+
+    with Session(engine) as session:
+        session.add(entry)
+        session.commit()
+        session.refresh(entry)
+
+    return entry

arcus/storage/stats.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+from statistics import mean
+
+from sqlmodel import Session, select
+
+from arcus.routing.reward import COST_SCORES
+from arcus.storage.db import RequestLog
+
+
+@dataclass(frozen=True)
+class ArmModeSummary:
+    model: str
+    mode: str
+    request_count: int
+    avg_reward: float | None
+    avg_latency_ms: float | None
+    cost_score: float | None
+
+
+def aggregate_by_arm_and_mode(engine) -> list[ArmModeSummary]:
+    """Groups logged requests by (model, mode) and averages reward/latency
+    per group. Done in plain Python rather than a SQL GROUP BY/AVG, this
+    is a personal local CLI's log, realistically thousands of rows at
+    most, not a scale where hand-rolled SQL aggregation earns its
+    complexity over just reading the rows.
+    """
+    with Session(engine) as session:
+        rows = session.exec(select(RequestLog)).all()
+
+    groups: dict[tuple[str, str], list[RequestLog]] = {}
+    for row in rows:
+        key = (row.model, row.mode or "unknown")
+        groups.setdefault(key, []).append(row)
+
+    summaries = []
+    for (model, mode), entries in groups.items():
+        rewards = [e.reward for e in entries if e.reward is not None]
+        latencies = [e.latency_ms for e in entries if e.latency_ms is not None]
+
+        summaries.append(
+            ArmModeSummary(
+                model=model,
+                mode=mode,
+                request_count=len(entries),
+                avg_reward=mean(rewards) if rewards else None,
+                avg_latency_ms=mean(latencies) if latencies else None,
+                cost_score=COST_SCORES.get(model),
+            )
+        )
+
+    return sorted(summaries, key=lambda s: (s.model, s.mode))