q-runtime-core 0.9.0__py3-none-any.whl

qruntime/budget.py

@@ -0,0 +1,295 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from qruntime.spend import (
+    SpendSummary,
+    _normalize_agent_id,
+    _normalize_session_id,
+)
+from qruntime.policy import _normalize_origin
+
+
+@dataclass(frozen=True, slots=True)
+class BudgetLimit:
+    network: str
+    asset: str
+    max_amount_atomic: int
+    origin: str | None = None
+    agent_id: str | None = None
+    session_id: str | None = None
+
+    def __post_init__(self) -> None:
+        if (
+            not isinstance(self.network, str)
+            or not self.network.strip()
+        ):
+            raise ValueError(
+                "budget network must be "
+                "a non-empty string"
+            )
+
+        if (
+            not isinstance(self.asset, str)
+            or not self.asset.strip()
+        ):
+            raise ValueError(
+                "budget asset must be "
+                "a non-empty string"
+            )
+
+        if (
+            isinstance(
+                self.max_amount_atomic,
+                bool,
+            )
+            or not isinstance(
+                self.max_amount_atomic,
+                int,
+            )
+            or self.max_amount_atomic <= 0
+        ):
+            raise ValueError(
+                "budget maximum must be "
+                "a positive integer"
+            )
+
+        if self.origin is not None:
+            object.__setattr__(
+                self,
+                "origin",
+                _normalize_origin(
+                    self.origin
+                ),
+            )
+
+        if self.agent_id is not None:
+            object.__setattr__(
+                self,
+                "agent_id",
+                _normalize_agent_id(
+                    self.agent_id
+                ),
+            )
+
+        if self.session_id is not None:
+            object.__setattr__(
+                self,
+                "session_id",
+                _normalize_session_id(
+                    self.session_id
+                ),
+            )
+
+
+@dataclass(frozen=True, slots=True)
+class BudgetDecision:
+    allowed: bool
+
+    network: str
+    asset: str
+
+    limit_atomic: int
+
+    confirmed_atomic: int
+    pending_atomic: int
+    requested_atomic: int
+
+    used_atomic: int
+    projected_atomic: int
+    remaining_before_atomic: int
+    remaining_after_atomic: int
+
+    reason: str | None = None
+    reserved_atomic: int = 0
+    origin: str | None = None
+    agent_id: str | None = None
+    session_id: str | None = None
+
+
+def evaluate_budget(
+    limit: BudgetLimit,
+    *,
+    summary: SpendSummary,
+    requested_amount_atomic: int,
+) -> BudgetDecision:
+    if summary.origin != limit.origin:
+        raise ValueError(
+            "spend summary origin scope "
+            "does not match budget origin"
+        )
+
+    if summary.agent_id != limit.agent_id:
+        raise ValueError(
+            "spend summary agent scope "
+            "does not match budget agent"
+        )
+
+    if summary.session_id != limit.session_id:
+        raise ValueError(
+            "spend summary session scope "
+            "does not match budget session"
+        )
+
+    if (
+        isinstance(
+            requested_amount_atomic,
+            bool,
+        )
+        or not isinstance(
+            requested_amount_atomic,
+            int,
+        )
+        or requested_amount_atomic <= 0
+    ):
+        raise ValueError(
+            "requested amount must be "
+            "a positive integer"
+        )
+
+    confirmed = summary.confirmed_amount(
+        network=limit.network,
+        asset=limit.asset,
+    )
+
+    pending = (
+        summary.pending_exposure_amount(
+            network=limit.network,
+            asset=limit.asset,
+        )
+    )
+
+    reserved = (
+        summary.reserved_exposure_amount(
+            network=limit.network,
+            asset=limit.asset,
+        )
+    )
+
+    used = (
+        confirmed
+        + pending
+        + reserved
+    )
+
+    projected = (
+        used
+        + requested_amount_atomic
+    )
+
+    remaining_before = max(
+        limit.max_amount_atomic - used,
+        0,
+    )
+
+    remaining_after = max(
+        limit.max_amount_atomic
+        - projected,
+        0,
+    )
+
+    allowed = (
+        projected
+        <= limit.max_amount_atomic
+    )
+
+    reason = None
+
+    if not allowed:
+        reason = "aggregate_budget_exceeded"
+
+    return BudgetDecision(
+        allowed=allowed,
+        network=limit.network,
+        asset=limit.asset,
+        limit_atomic=(
+            limit.max_amount_atomic
+        ),
+        confirmed_atomic=confirmed,
+        pending_atomic=pending,
+        reserved_atomic=reserved,
+        origin=limit.origin,
+        agent_id=limit.agent_id,
+        session_id=limit.session_id,
+        requested_atomic=(
+            requested_amount_atomic
+        ),
+        used_atomic=used,
+        projected_atomic=projected,
+        remaining_before_atomic=(
+            remaining_before
+        ),
+        remaining_after_atomic=(
+            remaining_after
+        ),
+        reason=reason,
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class BudgetSet:
+    """Multiple budget constraints that must all allow."""
+
+    limits: tuple[BudgetLimit, ...]
+
+    def __post_init__(self) -> None:
+        normalized = tuple(self.limits)
+
+        if not normalized:
+            raise ValueError(
+                "budget set must contain "
+                "at least one limit"
+            )
+
+        for limit in normalized:
+            if not isinstance(
+                limit,
+                BudgetLimit,
+            ):
+                raise TypeError(
+                    "budget set entries must "
+                    "be BudgetLimit instances"
+                )
+
+        seen = set()
+
+        for limit in normalized:
+            key = (
+                limit.network,
+                limit.asset,
+                limit.origin,
+                limit.agent_id,
+                limit.session_id,
+            )
+
+            if key in seen:
+                raise ValueError(
+                    "budget set contains duplicate "
+                    "budget scope"
+                )
+
+            seen.add(key)
+
+        object.__setattr__(
+            self,
+            "limits",
+            normalized,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class BudgetSetDecision:
+    """Result of evaluating all limits in one budget set."""
+
+    allowed: bool
+    decisions: tuple[BudgetDecision, ...]
+
+    @property
+    def first_denied(
+        self,
+    ) -> BudgetDecision | None:
+        for decision in self.decisions:
+            if not decision.allowed:
+                return decision
+
+        return None