« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Decimal price calculation and restricted expression evaluation.
3Computes per-dimension relay charges from normalized ``RelayUsage`` and
4configured price snapshots. Price expressions are ported from the
5relaykit billing conventions and evaluated by a small allow-listed
6parser: numeric literals, named usage dimensions, ``+ - * /``, ``min``,
7``max``, and parentheses. Attribute access, arbitrary calls,
8exponentiation, assignment, imports, and unbounded recursion are
9rejected. All arithmetic uses non-negative ``Decimal`` with explicit
10per-dimension rounding; every safety violation returns a
11``RelayBillingError`` instead of clamping silently.
12"""
14from __future__ import annotations
16from collections.abc import Callable, Mapping
17from dataclasses import dataclass
18from decimal import (
19 ROUND_HALF_UP,
20 Decimal,
21 DecimalException,
22 InvalidOperation,
23)
24import re
26from lexigram.contracts.ai.governance import (
27 RelayBillingError,
28 RelayChargeBreakdown,
29 RelayPriceEstimatorProtocol,
30 charge_overflow,
31 invalid_usage,
32 unknown_price,
33)
34from lexigram.contracts.ai.llm import CostEstimatorProtocol
35from lexigram.contracts.ai.relay import RelayUsage
36from lexigram.contracts.core.result import Err, Ok, Result
38__all__ = [
39 "BREAKDOWN_FIELDS",
40 "DEFAULT_MAX_CHARGE",
41 "DEFAULT_MAX_TOKENS",
42 "DEFAULT_SCALE",
43 "PriceSnapshot",
44 "RelayPricingEngine",
45 "SimpleCostEstimator",
46 "evaluate_expression",
47]
49BREAKDOWN_FIELDS = (
50 "prompt",
51 "cached_prompt",
52 "completion",
53 "reasoning",
54 "audio_input",
55 "audio_output",
56 "image",
57)
59_DIMENSION_TO_USAGE = {
60 "prompt": "prompt_tokens",
61 "cached_prompt": "cache_read_tokens",
62 "completion": "completion_tokens",
63 "reasoning": "reasoning_tokens",
64 "audio_input": "audio_input_tokens",
65 "audio_output": "audio_output_tokens",
66 "image": "image_tokens",
67}
69_ALLOWED_DIMENSION_NAMES = (
70 "prompt_tokens",
71 "completion_tokens",
72 "cache_read_tokens",
73 "cache_creation_tokens",
74 "reasoning_tokens",
75 "audio_input_tokens",
76 "audio_output_tokens",
77 "image_tokens",
78 "input_tokens",
79 "output_tokens",
80 "total_tokens",
81)
83_NUMBER_RE = re.compile(r"\d+(\.\d+)?([eE][+-]?\d+)?")
84_NAME_RE = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*")
86MAX_EXPRESSION_LENGTH = 512
87MAX_PARSE_DEPTH = 64
88DEFAULT_SCALE = 10
89DEFAULT_MAX_CHARGE = Decimal("1_000_000")
90DEFAULT_MAX_TOKENS = 2**31 - 1
91_PER_1M = Decimal(1_000_000)
94class _ExpressionError(Exception):
95 """Internal control flow for expression parsing and evaluation."""
98@dataclass(frozen=True, slots=True)
99class PriceSnapshot:
100 """Configured per-dimension price expressions for one model.
102 Each expression is evaluated against the named usage dimensions of
103 ``RelayUsage``, e.g. ``"prompt_tokens * 0.0000025"``. A dimension
104 without an expression contributes zero while remaining present in
105 the audit breakdown.
107 Attributes:
108 expressions: Mapping of breakdown field name to expression string.
109 """
111 expressions: Mapping[str, str]
113 def __post_init__(self) -> None:
114 """Validate dimension names and expression syntax."""
115 unknown = set(self.expressions) - set(BREAKDOWN_FIELDS)
116 if unknown:
117 names = ", ".join(sorted(unknown))
118 raise ValueError(f"unknown price dimensions: {names}")
119 for field_name, expression in self.expressions.items():
120 if not expression.strip():
121 raise ValueError(f"empty expression for dimension {field_name}")
122 if _syntax_error(expression) is not None:
123 raise ValueError(
124 f"invalid expression for dimension {field_name}: {expression}"
125 )
127 @classmethod
128 def from_per_1m(cls, prices: Mapping[str, Decimal]) -> PriceSnapshot:
129 """Build a snapshot from per-1M-token Decimal prices.
131 Args:
132 prices: Per-1M-token price per breakdown dimension; missing
133 dimensions default to zero.
135 Returns:
136 A snapshot whose expressions multiply each usage dimension
137 by its per-token price.
139 Raises:
140 ValueError: If a price is negative or a dimension is unknown.
141 """
142 unknown = set(prices) - set(BREAKDOWN_FIELDS)
143 if unknown:
144 names = ", ".join(sorted(unknown))
145 raise ValueError(f"unknown price dimensions: {names}")
146 expressions: dict[str, str] = {}
147 for field_name in BREAKDOWN_FIELDS:
148 price = prices.get(field_name, Decimal(0))
149 if price < 0:
150 raise ValueError(f"negative price for dimension {field_name}")
151 if price == 0:
152 expressions[field_name] = "0"
153 continue
154 usage_field = _DIMENSION_TO_USAGE[field_name]
155 per_token = (price / _PER_1M).normalize()
156 expressions[field_name] = f"{usage_field} * {format(per_token, 'f')}"
157 return cls(expressions=expressions)
160@dataclass(frozen=True, slots=True)
161class _Token:
162 """One lexical token of a price expression."""
164 kind: str
165 value: str = ""
168@dataclass(frozen=True, slots=True)
169class _NumberNode:
170 """Literal numeric node."""
172 value: Decimal
175@dataclass(frozen=True, slots=True)
176class _DimensionNode:
177 """Named usage dimension node."""
179 name: str
182@dataclass(frozen=True, slots=True)
183class _BinaryNode:
184 """Binary arithmetic node (``+ - * /``)."""
186 operator: str
187 left: _Node
188 right: _Node
191@dataclass(frozen=True, slots=True)
192class _ExtremeNode:
193 """``min``/``max`` call node."""
195 operator: str
196 left: _Node
197 right: _Node
200_Node = _NumberNode | _DimensionNode | _BinaryNode | _ExtremeNode
203def _syntax_error(expression: str) -> RelayBillingError | None:
204 """Return the parse error for *expression*, or ``None`` if valid."""
205 result = _tokenize(expression)
206 if result.is_err():
207 return result.unwrap_err()
208 parsed = _parse(result.unwrap())
209 if parsed.is_err():
210 return parsed.unwrap_err()
211 return None
214def _tokenize(expression: str) -> Result[list[_Token], RelayBillingError]:
215 """Tokenize a price expression, rejecting unsupported characters."""
216 if len(expression) > MAX_EXPRESSION_LENGTH:
217 return Err(
218 invalid_usage(
219 message=f"price expression exceeds {MAX_EXPRESSION_LENGTH} chars"
220 )
221 )
222 tokens: list[_Token] = []
223 index = 0
224 length = len(expression)
225 while index < length:
226 char = expression[index]
227 if char.isspace():
228 index += 1
229 continue
230 if char.isdigit():
231 match = _NUMBER_RE.match(expression, index)
232 assert match is not None # noqa: S101 # regex pre-guaranteed by isdigit() branch
233 tokens.append(_Token(kind="NUMBER", value=match.group(0)))
234 index = match.end()
235 continue
236 if char.isalpha() or char == "_":
237 match = _NAME_RE.match(expression, index)
238 assert match is not None # noqa: S101 # regex pre-guaranteed by isalpha() branch
239 tokens.append(_Token(kind="NAME", value=match.group(0)))
240 index = match.end()
241 continue
242 if char in "+-*/":
243 tokens.append(_Token(kind="OP", value=char))
244 index += 1
245 continue
246 if char == "(":
247 tokens.append(_Token(kind="LPAREN"))
248 index += 1
249 continue
250 if char == ")":
251 tokens.append(_Token(kind="RPAREN"))
252 index += 1
253 continue
254 if char == ",":
255 tokens.append(_Token(kind="COMMA"))
256 index += 1
257 continue
258 return Err(
259 invalid_usage(message=f"unsupported character {char!r} in price expression")
260 )
261 return Ok(tokens)
264def _parse(tokens: list[_Token]) -> Result[_Node, RelayBillingError]:
265 """Parse tokens into an allow-listed AST."""
267 position = 0
269 def peek() -> _Token | None:
270 """Return the next unconsumed token, if any."""
271 if position < len(tokens):
272 return tokens[position]
273 return None
275 def advance() -> _Token:
276 """Consume and return the next token."""
277 nonlocal position
278 token = tokens[position]
279 position += 1
280 return token
282 def expect(kind: str) -> _Token:
283 """Consume a token of *kind* or raise a parse error."""
284 token = peek()
285 if token is None or token.kind != kind:
286 raise _ExpressionError(f"expected {kind}")
287 return advance()
289 def parse_expression(depth: int) -> _Node:
290 """Parse ``+ -``; depth bounds recursion."""
291 if depth > MAX_PARSE_DEPTH:
292 raise _ExpressionError("expression nested too deeply")
293 node = parse_term(depth + 1)
294 while True:
295 token = peek()
296 if token is not None and token.kind == "OP" and token.value in "+-":
297 advance()
298 node = _BinaryNode(
299 operator=token.value,
300 left=node,
301 right=parse_term(depth + 1),
302 )
303 else:
304 return node
306 def parse_term(depth: int) -> _Node:
307 """Parse ``* /``; depth bounds recursion."""
308 node = parse_factor(depth + 1)
309 while True:
310 token = peek()
311 if token is not None and token.kind == "OP" and token.value in "*/":
312 advance()
313 node = _BinaryNode(
314 operator=token.value,
315 left=node,
316 right=parse_factor(depth + 1),
317 )
318 else:
319 return node
321 def parse_factor(depth: int) -> _Node:
322 """Parse literals, dimensions, parentheses, and min/max calls."""
323 if depth > MAX_PARSE_DEPTH:
324 raise _ExpressionError("expression nested too deeply")
325 token = peek()
326 if token is None:
327 raise _ExpressionError("unexpected end of expression")
328 if token.kind == "NUMBER":
329 advance()
330 try:
331 return _NumberNode(value=Decimal(token.value))
332 except InvalidOperation as exc:
333 raise _ExpressionError("invalid numeric literal") from exc
334 if token.kind == "NAME":
335 advance()
336 if token.value in ("min", "max"):
337 expect("LPAREN")
338 left = parse_expression(depth + 1)
339 expect("COMMA")
340 right = parse_expression(depth + 1)
341 expect("RPAREN")
342 return _ExtremeNode(operator=token.value, left=left, right=right)
343 following = peek()
344 if following is not None and following.kind == "LPAREN":
345 raise _ExpressionError(f"call {token.value!r} not allowed")
346 return _DimensionNode(name=token.value)
347 if token.kind == "LPAREN":
348 advance()
349 node = parse_expression(depth + 1)
350 expect("RPAREN")
351 return node
352 raise _ExpressionError(f"unexpected token {token.value or token.kind!r}")
354 if not tokens:
355 return Err(invalid_usage(message="empty price expression"))
356 try:
357 node = parse_expression(0)
358 if position != len(tokens):
359 token = tokens[position]
360 raise _ExpressionError(
361 f"unexpected trailing token {token.value or token.kind!r}"
362 )
363 return Ok(node)
364 except _ExpressionError as exc:
365 return Err(invalid_usage(message=str(exc)))
368def _evaluate(node: _Node, usage: RelayUsage, max_tokens: int) -> Decimal:
369 """Evaluate an AST against usage, raising on any safety violation."""
370 if isinstance(node, _NumberNode):
371 return node.value
372 if isinstance(node, _DimensionNode):
373 if node.name not in _ALLOWED_DIMENSION_NAMES:
374 raise _ExpressionError(f"unknown usage dimension {node.name!r}")
375 value = getattr(usage, node.name)
376 if value < 0:
377 raise _ExpressionError(f"negative usage dimension {node.name}")
378 if value > max_tokens:
379 raise _ExpressionError(f"usage dimension {node.name} exceeds maximum")
380 return Decimal(value)
381 left = _evaluate(node.left, usage, max_tokens)
382 right = _evaluate(node.right, usage, max_tokens)
383 try:
384 if isinstance(node, _BinaryNode):
385 if node.operator == "+":
386 result = left + right
387 elif node.operator == "-":
388 result = left - right
389 elif node.operator == "*":
390 result = left * right
391 else:
392 if right == 0:
393 raise _ExpressionError("division by zero")
394 result = left / right
395 else:
396 result = min(left, right) if node.operator == "min" else max(left, right)
397 except DecimalException as exc:
398 raise _ExpressionError("decimal arithmetic failure") from exc
399 if not result.is_finite():
400 raise _ExpressionError("non-finite result")
401 if result < 0:
402 raise _ExpressionError("negative result")
403 return result
406def evaluate_expression(
407 expression: str,
408 usage: RelayUsage,
409 *,
410 scale: int = DEFAULT_SCALE,
411 rounding: str = ROUND_HALF_UP,
412 max_tokens: int = DEFAULT_MAX_TOKENS,
413) -> Result[Decimal, RelayBillingError]:
414 """Evaluate a price expression against normalized usage.
416 Args:
417 expression: Price expression over named usage dimensions.
418 usage: Normalized usage the expression is evaluated against.
419 scale: Decimal places retained on the result.
420 rounding: Rounding mode applied to the result.
421 max_tokens: Integer maximum for any usage dimension.
423 Returns:
424 Ok(rounded charge) on success, Err(RelayBillingError) on any
425 syntax or arithmetic safety violation.
426 """
427 if scale < 0:
428 return Err(invalid_usage(message="scale must be non-negative"))
429 tokens = _tokenize(expression)
430 if tokens.is_err():
431 return Err(tokens.unwrap_err())
432 parsed = _parse(tokens.unwrap())
433 if parsed.is_err():
434 return Err(parsed.unwrap_err())
435 try:
436 value = _evaluate(parsed.unwrap(), usage, max_tokens)
437 except _ExpressionError as exc:
438 return Err(invalid_usage(message=str(exc)))
439 try:
440 quantum = Decimal(1).scaleb(-scale)
441 return Ok(value.quantize(quantum, rounding=rounding))
442 except DecimalException:
443 return Err(invalid_usage(message="charge rounding failed"))
446class RelayPricingEngine(RelayPriceEstimatorProtocol):
447 """Detailed per-dimension relay price estimator.
449 Args:
450 price_provider: Callable returning the ``PriceSnapshot`` for a
451 ``(model, provider, channel)`` triple, or ``None`` when the
452 model has no configured price.
453 scale: Default decimal places retained per dimension.
454 rounding: Default rounding mode applied per dimension.
455 max_charge: Maximum total charge; above it fails closed.
456 max_tokens: Integer maximum for any usage dimension.
457 dimension_scales: Per-dimension decimal places overrides.
458 dimension_roundings: Per-dimension rounding mode overrides.
460 Raises:
461 ValueError: If *max_charge* or *scale* is negative, or a
462 dimension override references an unknown breakdown field.
463 """
465 def __init__(
466 self,
467 price_provider: Callable[[str, str, str], PriceSnapshot | None],
468 *,
469 scale: int = DEFAULT_SCALE,
470 rounding: str = ROUND_HALF_UP,
471 max_charge: Decimal = DEFAULT_MAX_CHARGE,
472 max_tokens: int = DEFAULT_MAX_TOKENS,
473 dimension_scales: Mapping[str, int] | None = None,
474 dimension_roundings: Mapping[str, str] | None = None,
475 ) -> None:
476 if scale < 0:
477 raise ValueError("scale must be non-negative")
478 if max_charge < 0:
479 raise ValueError("max_charge must be non-negative")
480 unknown_scales = set(dimension_scales or {}) - set(BREAKDOWN_FIELDS)
481 if unknown_scales:
482 names = ", ".join(sorted(unknown_scales))
483 raise ValueError(f"unknown dimension scales: {names}")
484 unknown_roundings = set(dimension_roundings or {}) - set(BREAKDOWN_FIELDS)
485 if unknown_roundings:
486 names = ", ".join(sorted(unknown_roundings))
487 raise ValueError(f"unknown dimension roundings: {names}")
488 self._price_provider = price_provider
489 self._scale = scale
490 self._rounding = rounding
491 self._max_charge = max_charge
492 self._max_tokens = max_tokens
493 self._dimension_scales = dict(dimension_scales or {})
494 self._dimension_roundings = dict(dimension_roundings or {})
496 def estimate_charge(
497 self,
498 model: str,
499 usage: RelayUsage,
500 *,
501 provider: str = "",
502 channel: str = "",
503 ) -> Result[RelayChargeBreakdown, RelayBillingError]:
504 """Compute the per-dimension charge breakdown for *usage*.
506 Returns:
507 Ok(breakdown) on success, Err(RelayBillingError) on unknown
508 prices, negative usage, overflow, or expression failures.
509 """
510 for name in _ALLOWED_DIMENSION_NAMES:
511 if getattr(usage, name) < 0:
512 return Err(invalid_usage(message=f"negative usage dimension {name}"))
513 snapshot = self._price_provider(model, provider, channel)
514 if snapshot is None:
515 return Err(
516 unknown_price(message=f"no price configured for model {model!r}")
517 )
518 charges: dict[str, Decimal] = {}
519 for field_name in BREAKDOWN_FIELDS:
520 expression = snapshot.expressions.get(field_name, "0")
521 scale = self._dimension_scales.get(field_name, self._scale)
522 rounding = self._dimension_roundings.get(field_name, self._rounding)
523 charge = evaluate_expression(
524 expression,
525 usage,
526 scale=scale,
527 rounding=rounding,
528 max_tokens=self._max_tokens,
529 )
530 if charge.is_err():
531 return Err(charge.unwrap_err())
532 charges[field_name] = charge.unwrap()
533 total = sum(charges.values(), Decimal(0))
534 if total > self._max_charge:
535 return Err(
536 charge_overflow(
537 message=f"charge {total} exceeds configured maximum {self._max_charge}"
538 )
539 )
540 return Ok(RelayChargeBreakdown(**charges, total=total))
543class SimpleCostEstimator(RelayPriceEstimatorProtocol):
544 """Simple prompt/completion price estimator over ``CostEstimatorProtocol``.
546 Priced at the underlying estimator's input and output rates; the
547 detailed relay dimensions are always zero. Unknown models yield a
548 zero charge, matching the estimator's unknown-model-zero policy.
550 Args:
551 estimator: Existing LLM cost estimator reused for the simple path.
552 scale: Decimal places retained on each part of the charge.
553 rounding: Rounding mode applied to each part of the charge.
554 max_charge: Maximum total charge; above it fails closed.
555 """
557 def __init__(
558 self,
559 estimator: CostEstimatorProtocol,
560 *,
561 scale: int = DEFAULT_SCALE,
562 rounding: str = ROUND_HALF_UP,
563 max_charge: Decimal = DEFAULT_MAX_CHARGE,
564 ) -> None:
565 if scale < 0:
566 raise ValueError("scale must be non-negative")
567 if max_charge < 0:
568 raise ValueError("max_charge must be non-negative")
569 self._estimator = estimator
570 self._scale = scale
571 self._rounding = rounding
572 self._max_charge = max_charge
574 def estimate_charge(
575 self,
576 model: str,
577 usage: RelayUsage,
578 *,
579 provider: str = "",
580 channel: str = "",
581 ) -> Result[RelayChargeBreakdown, RelayBillingError]:
582 """Estimate prompt and completion charges from the LLM estimator."""
583 provider_arg = provider or None
584 try:
585 prompt_cost = float(
586 self._estimator.estimate_cost(
587 model,
588 usage.prompt_tokens,
589 provider_arg,
590 prompt_tokens=usage.prompt_tokens,
591 completion_tokens=0,
592 )
593 )
594 completion_cost = float(
595 self._estimator.estimate_cost(
596 model,
597 usage.completion_tokens,
598 provider_arg,
599 prompt_tokens=0,
600 completion_tokens=usage.completion_tokens,
601 )
602 )
603 except (TypeError, ValueError, OverflowError) as exc:
604 return Err(invalid_usage(message=f"estimator failed: {exc}"))
605 prompt = Decimal(str(prompt_cost))
606 completion = Decimal(str(completion_cost))
607 if not prompt.is_finite() or not completion.is_finite():
608 return Err(invalid_usage(message="estimator returned a non-finite cost"))
609 if prompt < 0 or completion < 0:
610 return Err(invalid_usage(message="estimator returned a negative cost"))
611 quantum = Decimal(1).scaleb(-self._scale)
612 try:
613 prompt = prompt.quantize(quantum, rounding=self._rounding)
614 completion = completion.quantize(quantum, rounding=self._rounding)
615 except DecimalException:
616 return Err(invalid_usage(message="charge rounding failed"))
617 zero = Decimal(0)
618 total = prompt + completion
619 if total > self._max_charge:
620 return Err(
621 charge_overflow(
622 message=f"charge {total} exceeds configured maximum {self._max_charge}"
623 )
624 )
625 return Ok(
626 RelayChargeBreakdown(
627 prompt=prompt,
628 cached_prompt=zero,
629 completion=completion,
630 reasoning=zero,
631 audio_input=zero,
632 audio_output=zero,
633 image=zero,
634 total=total,
635 )
636 )