বাজেট, ক্যাশিং ও ব্যবহারের সীমা

# file: evidencedesk/limits.py
"""Cost controls: budgets, billed retries, a model-swap rule, batching, timeouts, usage limits, hit rate, a summary."""
from __future__ import annotations

import random
from collections.abc import Callable, Sequence
from dataclasses import dataclass, replace
from datetime import datetime, timezone
from decimal import Decimal

from evidencedesk.cache import ResponseCache, cache_key
from evidencedesk.evals import CaseModel, Dataset, run_suite
from evidencedesk.errors import LLMError, RetryPolicy, call_with_retry
from evidencedesk.llm import LLMClient, LLMRequest, LLMResponse, Message
from evidencedesk.rag import Retriever
from evidencedesk.tools import ToolRegistry
from evidencedesk.usage import Ledger, ManualClock, UsageRecord, cost_of, record_usage, request_tokens

# --- 1. Input and output budgets ---
@dataclass(frozen=True)
class Budgets:
    max_input_tokens: int = 600  # what we send (estimated)
    max_output_tokens: int = 150  # what the model may write


BY_CLASS = {"faq": Budgets(500, 120), "summary": Budgets(800, 200)}  # one budget per request class


class OverBudget(ValueError):
    code = "input_too_large"


def apply_budgets(request: LLMRequest, budgets: Budgets) -> LLMRequest:
    """Before the call: refuse an oversized input; cap max_tokens (never raise it)."""
    if request_tokens(request) > budgets.max_input_tokens:
        raise OverBudget(f"input is about {request_tokens(request)} tokens, the limit is {budgets.max_input_tokens}")
    return replace(request, max_tokens=min(request.max_tokens, budgets.max_output_tokens))

from evidencedesk.limits import BY_CLASS, OverBudget, apply_budgets
from evidencedesk.llm import LLMRequest, Message

wants_a_lot = LLMRequest((Message("user", "How long do I have to return a bike?"),), max_tokens=1000)
too_long = LLMRequest((Message("user", "Please read this. " * 200),))
faq = BY_CLASS["faq"]
print("max_tokens:", wants_a_lot.max_tokens, "->", apply_budgets(wants_a_lot, faq).max_tokens)
try:
    apply_budgets(too_long, faq)
except OverBudget as error:
    print(error.code, "-", error)
print("input limits by class:", {name: b.max_input_tokens for name, b in BY_CLASS.items()})

max_tokens: 1000 -> 120
input_too_large - input is about 900 tokens, the limit is 500
input limits by class: {'faq': 500, 'summary': 800}

# file: evidencedesk/limits.py (append)
# --- 2. Billed attempts and bounded retries ---
class Flaky:
    """Fails with LLMError(kind) on the first `fail_times` calls, then answers."""

    def __init__(self, inner: LLMClient, kind: str = "server", fail_times: int = 99) -> None:
        self.inner, self.kind, self.fail_times, self.calls = inner, kind, fail_times, 0

    def complete(self, request: LLMRequest) -> LLMResponse:
        self.calls += 1
        if self.calls <= self.fail_times:
            raise LLMError(self.kind, "simulated failure")
        return self.inner.complete(request)


def call_billed(client: LLMClient, request: LLMRequest, ledger: Ledger, request_id: str, policy: RetryPolicy = RetryPolicy(),
                sleep: Callable[[float], None] = lambda s: None, rng: random.Random | None = None) -> LLMResponse:
    """Call with a bounded retry; every attempt goes in the ledger. Assumption: a failed retryable try cost its input."""
    attempt = 0

    def one() -> LLMResponse:
        nonlocal attempt
        attempt += 1
        try:
            response = client.complete(request)
        except LLMError as error:
            if error.retryable:
                ledger.add(UsageRecord(request_id, request.model, request_tokens(request), 0, attempt=attempt))
            raise
        ledger.add(record_usage(request, response, request_id, attempt=attempt))
        return response

    return call_with_retry(one, policy, sleep, rng or random.Random(0))

from evidencedesk.errors import RetryPolicy
from evidencedesk.limits import Flaky, call_billed
from evidencedesk.llm import FakeClient, LLMRequest, Message
from evidencedesk.usage import Ledger

request = LLMRequest((Message("user", "How long do I have to return a new bike?"),))
scenarios = [  # (label, failure kind, failures before success, attempt limit)
    ("recovers after 2 failures", "server", 2, 3),
    ("outage, bounded at 3", "server", 99, 3),
    ("outage, no real limit (20)", "server", 99, 20),
    ("auth error, bounded at 3", "auth", 99, 3),
]
print(f"{'scenario':<28}{'calls':>6}{'billed':>10}  result")
for label, kind, fails, limit in scenarios:
    ledger, flaky = Ledger(), Flaky(FakeClient(), kind, fails)
    try:
        call_billed(flaky, request, ledger, "r1", RetryPolicy(max_attempts=limit))
        result = "answered"
    except Exception as error:
        result = f"{error.kind} after {error.attempts}"
    print(f"{label:<28}{flaky.calls:>6}{ledger.summary().cost:>10.6f}  {result}")

scenario                     calls    billed  result
recovers after 2 failures        3  0.000411  answered
outage, bounded at 3             3  0.000090  server after 3
outage, no real limit (20)      20  0.000600  server after 20
auth error, bounded at 3         1  0.000000  auth after 1

# file: evidencedesk/limits.py (append)
# --- 3. A cheaper model only if the same eval set agrees ---
class Metered:
    """Relabels a fake's answers with a model name and records their (estimated) usage."""

    def __init__(self, inner: LLMClient, model: str, ledger: Ledger) -> None:
        self.inner, self.model, self.ledger = inner, model, ledger

    def complete(self, request: LLMRequest) -> LLMResponse:
        response = replace(self.inner.complete(request), model=self.model)
        self.ledger.add(record_usage(request, response, f"r{len(self.ledger.records) + 1}"))
        return response


@dataclass(frozen=True)
class Trial:
    model: str
    quality: float  # pass_rate of run_suite on the SAME dataset
    cost: Decimal  # credits for that run
    calls: int
    cases: int


@dataclass(frozen=True)
class Verdict:
    accepted: bool
    reason: str  # accepted | quality_dropped | not_cheaper


def accept_cheaper(baseline: Trial, candidate: Trial, tolerance: float = 0.02) -> Verdict:
    """Accept only if quality >= baseline - tolerance AND cost is lower, on the same dataset."""
    if candidate.cases != baseline.cases:
        raise ValueError("compare runs of the same dataset")
    if candidate.quality < baseline.quality - tolerance:
        return Verdict(False, "quality_dropped")
    return Verdict(True, "accepted") if candidate.cost < baseline.cost else Verdict(False, "not_cheaper")


def run_trial(model: str, replies: dict[str, dict], dataset: Dataset, retriever: Retriever, registry: ToolRegistry) -> Trial:
    """Run the whole eval set with hand-written replies (by case id) and price the model calls."""
    ledger = Ledger()
    client = Metered(CaseModel({c.input: replies[c.id] for c in dataset.cases if c.id in replies}), model, ledger)
    quality = run_suite(dataset, retriever, client, registry)["metrics"]["pass_rate"]
    t = ledger.summary()
    return Trial(model, quality, t.cost, t.calls, len(dataset.cases))

from pathlib import Path

from evidencedesk.evals import build_retriever, load_dataset
from evidencedesk.limits import accept_cheaper, run_trial
from evidencedesk.rag import reply
from evidencedesk.tools import default_registry
from tests.test_citations import REPLIES  # the hand-written replies of the evaluation pages

dataset = load_dataset(Path("data/eval_manifest.json"))
retriever, registry = build_retriever(Path("corpus")), default_registry(Path("corpus"))
weak = {**REPLIES, "rc01": reply("Returns are possible.", "kc-returns-v3:c9"),  # a citation that does not exist
        "rc05": reply("Returns are possible.", "kc-shipping-v1:c0"),  # a source that does not support it
        "rc12": reply("Size M.", "kc-nope:c0")}
runs = {"fake-model": ("fake-model", REPLIES), "fake-small (weak)": ("fake-small", weak), "fake-small (same)": ("fake-small", REPLIES)}
trials = {name: run_trial(model, replies, dataset, retriever, registry) for name, (model, replies) in runs.items()}
base = trials["fake-model"]
print(f"{'configuration':<20}{'quality':>8}{'cost':>10}{'calls':>6}  verdict")
for name, t in trials.items():
    print(f"{name:<20}{t.quality:>8}{t.cost:>10.6f}{t.calls:>6}  {'(baseline)' if t is base else accept_cheaper(base, t).reason}")

configuration        quality      cost calls  verdict
fake-model             0.852  0.013518     9  (baseline)
fake-small (weak)      0.741  0.002636     9  quality_dropped
fake-small (same)      0.852  0.002704     9  accepted

# file: evidencedesk/limits.py (append)
# --- 4. Batching ---
BATCH_FACTOR = Decimal("0.5")  # invented "batch price is half"; not a real discount


class BatchClient:
    """A batch-shaped fake: one submission carries many requests."""

    def __init__(self, inner: LLMClient) -> None:
        self.inner, self.submissions = inner, 0

    def complete_many(self, requests: Sequence[LLMRequest]) -> list[LLMResponse]:
        self.submissions += 1
        return [self.inner.complete(r) for r in requests]


def batch_cost(records: Sequence[UsageRecord], factor: Decimal = BATCH_FACTOR) -> Decimal:
    return sum((cost_of(r) for r in records), Decimal(0)) * factor

from evidencedesk.limits import BatchClient, batch_cost
from evidencedesk.llm import FakeClient, LLMRequest, Message
from evidencedesk.usage import Ledger, record_usage

notes = {f"Summarise note {i}": {"text": f"Note {i} in one line.", "usage": {"input_tokens": 120, "output_tokens": 15}} for i in range(1, 7)}
requests = [LLMRequest((Message("user", q),)) for q in notes]
single, per_item = FakeClient(notes), Ledger()
for n, r in enumerate(requests, 1):  # six separate calls
    per_item.add(record_usage(r, single.complete(r), f"r{n}"))
batch = BatchClient(FakeClient(notes))
answers = batch.complete_many(requests)  # one submission
print("per item:", per_item.summary().calls, "calls, cost", per_item.summary().cost)
print("batched: ", batch.submissions, "submission, cost", batch_cost(per_item.records), "(invented factor 0.5), answers:", len(answers))

per item: 6 calls, cost 0.00324
batched:  1 submission, cost 0.001620 (invented factor 0.5), answers: 6

# file: evidencedesk/limits.py (append)
# --- 5. Per-request timeout (checked on an injected clock) ---
class TimeoutClient:
    """A fake cannot be interrupted: read the clock after the call and discard a late answer."""

    def __init__(self, inner: LLMClient, clock: ManualClock, limit_s: float) -> None:
        self.inner, self.clock, self.limit_s = inner, clock, limit_s

    def complete(self, request: LLMRequest) -> LLMResponse:
        start = self.clock()
        response = self.inner.complete(request)
        if self.clock() - start > self.limit_s:
            raise LLMError("timeout", f"no answer within {self.limit_s} s")
        return response

from evidencedesk.errors import RetryPolicy
from evidencedesk.limits import TimeoutClient, call_billed
from evidencedesk.llm import FakeClient, LLMRequest, Message
from evidencedesk.usage import Ledger, ManualClock, SlowClient

request = LLMRequest((Message("user", "How long do I have to return a new bike?"),))
for limit in (0.5, 1.0):
    clock, ledger = ManualClock(), Ledger()
    client = TimeoutClient(SlowClient(FakeClient(), clock), clock, limit_s=limit)  # each call takes 0.78 s on the fake clock
    try:
        call_billed(client, request, ledger, "r1", RetryPolicy(max_attempts=3))
        outcome = "answered"
    except Exception as error:
        outcome = f"{error.kind} after {error.attempts} tries"
    print(f"limit {limit} s: {outcome}; ledger records: {len(ledger.records)}")

limit 0.5 s: timeout after 3 tries; ledger records: 3
limit 1.0 s: answered; ledger records: 1

# file: evidencedesk/limits.py (append)
# --- 6. A local daily usage limit ---
DAY = 86_400


@dataclass(frozen=True)
class Quota:
    calls: int
    tokens: int


@dataclass(frozen=True)
class Decision:
    allowed: bool
    code: str = "ok"  # "ok" or the stable code "budget_exceeded"
    scope: str = ""  # "user" or "global"
    reset_at: str = ""  # ISO 8601, UTC


class UsageLimiter:
    """Per-user and global counters for one UTC day, on an injected clock (epoch seconds)."""

    def __init__(self, clock: Callable[[], float], per_user: Quota = Quota(5, 2000), overall: Quota = Quota(20, 8000)) -> None:
        self.clock, self.per_user, self.overall = clock, per_user, overall
        self.used: dict[tuple[int, str], list[int]] = {}  # (day, scope) -> [calls, tokens]

    def _day(self) -> int:
        return int(self.clock() // DAY)

    def reset_at(self) -> str:
        return datetime.fromtimestamp((self._day() + 1) * DAY, timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

    def _scopes(self, user_id: str):
        return ((f"user:{user_id}", "user", self.per_user), ("global", "global", self.overall))

    def check(self, user_id: str, tokens: int) -> Decision:
        for key, name, quota in self._scopes(user_id):
            calls, used = self.used.get((self._day(), key), (0, 0))
            if calls + 1 > quota.calls or used + tokens > quota.tokens:
                return Decision(False, "budget_exceeded", name, self.reset_at())
        return Decision(True)

    def record(self, user_id: str, tokens: int) -> None:
        day = self._day()
        self.used = {k: v for k, v in self.used.items() if k[0] == day}  # yesterday's counters are dropped
        for key, _, _ in self._scopes(user_id):
            entry = self.used.setdefault((day, key), [0, 0])
            entry[0] += 1
            entry[1] += tokens

from datetime import datetime, timezone

from evidencedesk.limits import Quota, UsageLimiter
from evidencedesk.usage import ManualClock

clock = ManualClock(datetime(2026, 10, 13, 22, 0, tzinfo=timezone.utc).timestamp())  # 22:00 UTC
limiter = UsageLimiter(clock, per_user=Quota(calls=2, tokens=400), overall=Quota(calls=3, tokens=1000))
for who in ("nadia", "nadia", "nadia", "omar", "aiko"):
    d = limiter.check(who, 100)
    if d.allowed:
        limiter.record(who, 100)
    print(f"{who:<6}{d.code:<16}{d.scope:<7}{d.reset_at}".rstrip())
clock.advance(3 * 3600)  # now 01:00 UTC the next day
print("after midnight:", limiter.check("nadia", 100))

nadia ok
nadia ok
nadia budget_exceeded user   2026-10-14T00:00:00Z
omar  ok
aiko  budget_exceeded global 2026-10-14T00:00:00Z
after midnight: Decision(allowed=True, code='ok', scope='', reset_at='')

# file: evidencedesk/limits.py (append)
# --- 7. Cache hit rate ---
@dataclass(frozen=True)
class HitStats:
    hits: int
    misses: int  # lookups that found nothing (expired entries included)
    bypassed: int  # never looked up: not eligible

    @property
    def eligible(self) -> int:
        return self.hits + self.misses

    @property
    def rate(self) -> float | None:
        return round(self.hits / self.eligible, 3) if self.eligible else None

    @property
    def naive_rate(self) -> float | None:
        """The wrong number: bypasses counted as lookups."""
        total = self.eligible + self.bypassed
        return round(self.hits / total, 3) if total else None


def hit_stats(cache: ResponseCache, bypassed: int = 0) -> HitStats:
    return HitStats(cache.hits, cache.misses, bypassed)

from evidencedesk.cache import FAQ_ONLY, CachedClient, ResponseCache
from evidencedesk.limits import hit_stats
from evidencedesk.llm import FakeClient, LLMRequest, Message
from evidencedesk.usage import ManualClock

cached = CachedClient(FakeClient(), ResponseCache(ttl_s=60, clock=ManualClock()), FAQ_ONLY)
faq = LLMRequest((Message("user", "How long do I have to return a new bike?"),))
account = LLMRequest((Message("user", "How long does standard delivery take?"),))
for _ in range(4):
    cached.ask(faq, request_class="faq")
for _ in range(6):
    cached.ask(account, request_class="account", user_id="u1")  # not on the policy: bypass
stats = hit_stats(cached.cache, cached.bypassed)
print(stats, "eligible:", stats.eligible)
print("hit rate:", stats.rate, "| wrong (bypasses counted):", stats.naive_rate)

HitStats(hits=3, misses=1, bypassed=6) eligible: 4
hit rate: 0.75 | wrong (bypasses counted): 0.3

# file: evidencedesk/limits.py (append)
# --- 8. Summary and pricing uncertainty ---
PRICING_NOTE = "Rates are invented. Any real price comparison: VERIFY - checked 2026-10-12; read the current pricing page."


def cost_latency_summary(ledger: Ledger, latencies: Sequence[float], stats: HitStats | None = None) -> dict:
    """Counts and averages from the ledger and the timer. No percentiles: n is too small."""
    t, n = ledger.summary(), len(latencies)
    calls = max(t.calls, 1)
    out = {
        "calls": t.calls, "retries": t.retries,
        "avg_input_tokens": round(t.input_tokens / calls, 1), "avg_output_tokens": round(t.output_tokens / calls, 1),
        "total_cost": str(t.cost), "avg_cost_per_call": str((t.cost / calls).quantize(Decimal("0.000001"))),
        "timed_requests": n, "avg_latency_s": round(sum(latencies) / n, 2) if n else None, "max_latency_s": round(max(latencies), 2) if n else None,
        "sample_note": f"{t.calls} calls, {n} timings: this describes this run only, not production.",
        "pricing_note": PRICING_NOTE,
    }
    if stats is not None:
        out["cache"] = {"hits": stats.hits, "eligible": stats.eligible, "bypassed": stats.bypassed, "hit_rate": stats.rate}
    return out

# file: evidencedesk/limits.py (append)
# --- 9. Mini-lab: a budget-aware summariser ---
SUMMARY_SYSTEM = "Summarise the text in at most two sentences."


@dataclass(frozen=True)
class Summary:
    ok: bool
    code: str  # ok | input_too_large | budget_exceeded | timeout | server ...
    text: str = ""
    source: str = "none"  # cache | model
    reset_at: str = ""
    attempts: int = 0


def summarise(text: str, *, client: LLMClient, cache: ResponseCache, ledger: Ledger, limiter: UsageLimiter, clock: ManualClock,
              user_id: str, request_id: str, budgets: Budgets = BY_CLASS["summary"], timeout_s: float = 2.0,
              policy: RetryPolicy = RetryPolicy(max_attempts=2)) -> Summary:
    """Budget, cache, usage limit, timeout, bounded retry, ledger: in that order. Expected failures are results."""
    request = LLMRequest((Message("system", SUMMARY_SYSTEM), Message("user", text)), max_tokens=budgets.max_output_tokens)
    try:
        request = apply_budgets(request, budgets)
    except OverBudget as error:
        return Summary(False, error.code)
    key = cache_key(request, prompt_version="summary-v1")  # public text only: one shared scope
    if (hit := cache.get(key)) is not None:  # no call, so no budget used
        return Summary(True, "ok", hit.text, "cache")
    worst_case = request_tokens(request) + request.max_tokens
    decision = limiter.check(user_id, worst_case)
    if not decision.allowed:
        return Summary(False, decision.code, reset_at=decision.reset_at)
    before = len(ledger.records)
    try:
        response = call_billed(TimeoutClient(client, clock, timeout_s), request, ledger, request_id, policy)
    except LLMError as error:
        return Summary(False, error.kind, attempts=error.attempts)
    limiter.record(user_id, response.usage.total_tokens or worst_case)
    if response.stop_reason == "end":
        cache.put(key, response)
    return Summary(True, "ok", response.text, "model", attempts=len(ledger.records) - before)

from datetime import datetime, timezone

from evidencedesk.cache import ResponseCache
from evidencedesk.errors import RetryPolicy
from evidencedesk.limits import Flaky, Quota, UsageLimiter, summarise
from evidencedesk.llm import FakeClient
from evidencedesk.usage import Ledger, ManualClock, SlowClient

NOTES = [f"The bike shop opens at nine and closes at six. Note {i}." for i in range(5)]
fake = FakeClient({n: {"text": "The shop is open nine to six.", "usage": {"input_tokens": 40, "output_tokens": 9}} for n in NOTES})
clock = ManualClock(datetime(2026, 10, 13, 22, 0, tzinfo=timezone.utc).timestamp())
ledger, cache = Ledger(), ResponseCache(ttl_s=300, clock=clock)
limiter = UsageLimiter(clock, per_user=Quota(calls=2, tokens=1000), overall=Quota(calls=10, tokens=5000))
lab = dict(cache=cache, ledger=ledger, limiter=limiter, clock=clock, policy=RetryPolicy(max_attempts=2))
steps = [  # (label, client, text, user)
    ("first call", SlowClient(fake, clock), NOTES[0], "u1"),
    ("same text again", SlowClient(fake, clock), NOTES[0], "u1"),
    ("input too large", fake, "word " * 1000, "u1"),
    ("slow model", SlowClient(fake, clock, base_s=3.0), NOTES[1], "u1"),
    ("second model call", fake, NOTES[2], "u1"),
    ("user budget used up", fake, NOTES[3], "u1"),
    ("another user, outage", Flaky(fake, "server"), NOTES[4], "u2"),
]
latencies = []
for n, (label, client, text, user) in enumerate(steps, 1):
    start = clock()
    r = summarise(text, client=client, user_id=user, request_id=f"r{n}", **lab)
    latencies.append(round(clock() - start, 2))
    print(f"{label:<22}{r.code:<17}{r.source:<7}tries={r.attempts} {r.reset_at}".rstrip())

first call            ok               model  tries=1
same text again       ok               cache  tries=0
input too large       input_too_large  none   tries=0
slow model            timeout          none   tries=2
second model call     ok               model  tries=1
user budget used up   budget_exceeded  none   tries=0 2026-10-14T00:00:00Z
another user, outage  server           none   tries=2

import json

from evidencedesk.limits import cost_latency_summary, hit_stats

print(json.dumps(cost_latency_summary(ledger, latencies, hit_stats(cache)), indent=1))

{
 "calls": 6,
 "retries": 2,
 "avg_input_tokens": 30.0,
 "avg_output_tokens": 3.0,
 "total_cost": "0.000756",
 "avg_cost_per_call": "0.000126",
 "timed_requests": 7,
 "avg_latency_s": 0.99,
 "max_latency_s": 6.36,
 "sample_note": "6 calls, 7 timings: this describes this run only, not production.",
 "pricing_note": "Rates are invented. Any real price comparison: VERIFY - checked 2026-10-12; read the current pricing page.",
 "cache": {
  "hits": 1,
  "eligible": 6,
  "bypassed": 0,
  "hit_rate": 0.167
 }
}

python -m pytest -q tests/test_limits.py | sed 's/ in [0-9.]*s//'

.................                                                        [100%]
17 passed

tiny = UsageLimiter(clock, per_user=Quota(calls=1, tokens=100))
for _ in range(2):
    d = tiny.check("kwame", 40)
    print(f"{d.code} {d.scope} {d.reset_at}".strip())
    if d.allowed:
        tiny.record("kwame", 40)

ok
budget_exceeded user 2026-10-14T00:00:00Z

from evidencedesk.errors import from_status

for status, body in ((503, ""), (429, "credit_balance_exhausted"), (400, "")):
    error = from_status(status, {}, body)
    flaky = Flaky(FakeClient(), error.kind)
    try:
        call_billed(flaky, request, Ledger(), "r1", RetryPolicy(max_attempts=3))
    except Exception:
        pass
    print(status, error.kind, "retryable:", error.retryable, "calls:", flaky.calls)

503 server retryable: True calls: 3
429 quota_exhausted retryable: False calls: 1
400 invalid_request retryable: False calls: 1

slip = {**REPLIES, "rc04": reply("Probably.", "kc-assembly-v1:c9")}  # only rc04 gets worse
candidate = run_trial("fake-small", slip, dataset, retriever, registry)
print(candidate.quality, "vs baseline", base.quality)
for tolerance in (0.02, 0.04):
    print("tolerance", tolerance, "->", accept_cheaper(base, candidate, tolerance).reason)

0.815 vs baseline 0.852
tolerance 0.02 -> quality_dropped
tolerance 0.04 -> accepted