Spaces:
Running
Running
| """One place where a model call happens — traced, budgeted and telemetered. | |
| Every node goes through :func:`call_model`, so there is exactly one code path | |
| that charges the budget guard, records the Langfuse generation, and publishes the | |
| `model.call` telemetry row. Nothing can call Claude and skip the guard. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from app.core.budget import BudgetExceededError | |
| from app.core.claude import LLMRequest, LLMResult | |
| from app.core.events import EventKind | |
| from app.graph.context import RunContext | |
| from app.models.run import TokenSpend | |
| logger = logging.getLogger(__name__) | |
| class ModelCallOutcome: | |
| """Result of a guarded model call, plus the spend delta for state.""" | |
| __slots__ = ("result", "spend") | |
| def __init__(self, result: LLMResult, spend: TokenSpend) -> None: | |
| self.result = result | |
| self.spend = spend | |
| async def call_model(ctx: RunContext, request: LLMRequest) -> ModelCallOutcome: | |
| """Execute one model call under the budget guard, tracing it end to end. | |
| Raises: | |
| BudgetExceededError: when the run has exhausted its token or USD budget. | |
| Nodes convert this into a ``BUDGET_ABORT`` status rather than crashing. | |
| """ | |
| model = ctx.engine.model_for(request.role) | |
| max_tokens = request.max_tokens or ctx.settings.llm_max_tokens | |
| # Pre-flight: refuse a call whose worst case cannot fit in the remaining budget. | |
| ctx.budget.ensure_headroom(model, max_tokens) | |
| await ctx.emit( | |
| EventKind.MODEL_CALL, | |
| f"{request.role} → {model}", | |
| {"role": str(request.role), "model": model, "hint": str(request.hint)}, | |
| ) | |
| with ctx.tracer.step( | |
| f"model:{request.role}", | |
| run_id=ctx.run_id, | |
| as_type="generation", | |
| input_data={"hint": str(request.hint), "system": request.system[:400]}, | |
| metadata={"model": model, "engine": ctx.engine.name}, | |
| ) as span: | |
| result = await ctx.engine.complete(request) | |
| span.update(output={"text": result.text[:800], "tool_calls": len(result.tool_calls)}) | |
| cost = result.usage.cost_usd(result.model) | |
| try: | |
| snapshot = ctx.budget.charge(result.model, result.usage) | |
| except BudgetExceededError: | |
| await ctx.emit( | |
| EventKind.WARNING, | |
| "Per-run token budget exhausted — aborting", | |
| {"model": result.model}, | |
| ) | |
| raise | |
| ctx.tracer.record_generation( | |
| f"model:{request.role}", | |
| run_id=ctx.run_id, | |
| model=result.model, | |
| input_data={"hint": str(request.hint)}, | |
| output_data=result.text[:2000], | |
| usage=result.usage, | |
| cost_usd=cost, | |
| metadata={"engine": result.engine, "stop_reason": result.stop_reason}, | |
| ) | |
| spend = TokenSpend( | |
| input_tokens=result.usage.input_tokens, | |
| output_tokens=result.usage.output_tokens, | |
| cache_read_tokens=result.usage.cache_read_tokens, | |
| cache_write_tokens=result.usage.cache_write_tokens, | |
| cost_usd=cost, | |
| calls=1, | |
| ) | |
| logger.debug( | |
| "model call %s/%s tokens=%s cost=%.6f budget_remaining=%.4f", | |
| request.role, | |
| request.hint, | |
| result.usage.total_tokens, | |
| cost, | |
| snapshot.remaining_usd, | |
| ) | |
| return ModelCallOutcome(result=result, spend=spend) | |