diff --git a/composer/diagnostics/budget.py b/composer/diagnostics/budget.py new file mode 100644 index 00000000..b6df628b --- /dev/null +++ b/composer/diagnostics/budget.py @@ -0,0 +1,214 @@ +from typing import Iterator, Callable, Any, Mapping +from typing_extensions import TypeVar +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass + +from langgraph.graph import MessagesState +from graphcore.graph import StateMonitor, MonitorReturn +from langchain_core.messages import HumanMessage + +StateVar = TypeVar("StateVar", default=MessagesState, bound=MessagesState) + +# The fraction of a budget at which "budget pressure" begins: budget_monitor's +# warning fires and budget_pressure() flips true. One global so every accessor +# agrees on where the wrap-up window starts. +BUDGET_PRESSURE_THRESHOLD = 0.8 + + +class BudgetExceeded(Exception): + """Hard budget stop, raised cooperatively (from a monitor, between agent + turns) once the active budget is blown. Only monitors given an + ``on_overbudget`` callback raise; the workflow that launched the agent + catches this and converts it into its give-up result.""" + + +class BudgetPressureAbort(Exception): + """Raised by ``pressure_abort_monitor`` to terminate an auxiliary agent + (e.g. a feedback judge) whose output is worthless once the main agent is + in its wrap-up window. Caught by the tool that launched the agent.""" + + +@dataclass +class BudgetCounter: + """One node of the caps-over-pool scheme. Leaf counters are per-phase + *caps* whose ``parent`` is the run's shared *pool* (the real budget); + ``token_cost_budget`` creates parentless one-off counters. Cost accrues + up the chain, and both the hard stop and the pressure window trip on + whichever level is tighter — a phase can only starve later phases up to + its cap, while unspent phase money never leaves the pool (rollover is + automatic, not an explicit transfer).""" + total_budget: float + curr_cost: float + parent: "BudgetCounter | None" = None + + def overbudget(self) -> bool: + if self.curr_cost > self.total_budget: + return True + return self.parent.overbudget() if self.parent is not None else False + + def pressured(self, threshold: float = BUDGET_PRESSURE_THRESHOLD) -> bool: + if self.curr_cost >= self.total_budget * threshold: + return True + return self.parent.pressured(threshold) if self.parent is not None else False + +_budget_accumulator = ContextVar[None | BudgetCounter]("_budget_accumulator", default=None) + +_cost_centers = ContextVar[None | dict[str, BudgetCounter]]("_cost_centers", default=None) + + +DEFAULT_BUDGET_PRESSURE_MESSAGE = """ + +You have almost exceeded the token cost budget allotted for this task. + +Finish your task in as orderly a fashion as possible; partial/incomplete results are better +than going over budget. + +""" + +@contextmanager +def total_budget( + total: float, + caps: Mapping[str, float] +) -> Iterator[None]: + """Install the run's budget: ``total`` is the pool (the real bound on + spend) and ``caps`` are per-phase ceilings. Caps need not sum to the + pool — they only bound how much a single phase may hog, so each can be + generous; whatever a phase doesn't spend simply remains in the pool for + later phases.""" + curr = _cost_centers.get() + if curr is not None: + raise RuntimeError("Not good") + pool = BudgetCounter(total_budget=total, curr_cost=0.0) + prev = _cost_centers.set({ + k: BudgetCounter(total_budget=v, curr_cost=0.0, parent=pool) for (k, v) in caps.items() + }) + # Work running outside any named center (e.g. the report phase) accrues + # to — and feels pressure from — the pool directly. + prev_accum = _budget_accumulator.set(pool) + try: + yield None + finally: + _budget_accumulator.reset(prev_accum) + _cost_centers.reset(prev) + +@contextmanager +def named_budget( + nm: str +) -> Iterator[None]: + if (res := _cost_centers.get()) is None: + raise RuntimeError("No costs installed") + if nm not in res: + raise RuntimeError(f"Named budget item not known: {nm}") + prev = _budget_accumulator.set(res[nm]) + try: + yield + finally: + _budget_accumulator.reset(prev) + +@contextmanager +def named_budget_or_nop( + nm: str +) -> Iterator[None]: + if (_cost_centers.get()) is None: + # A @contextmanager generator must yield exactly once even on the + # nop path — a bare return raises "generator didn't yield". + yield + return + with named_budget(nm): + yield + +@contextmanager +def token_cost_budget( + total_cost: float, +) -> Iterator[None]: + if _budget_accumulator.get() is not None: + raise RuntimeError("Nested budgets not supported") + accum = BudgetCounter(total_budget=total_cost, curr_cost=0.0) + prev = _budget_accumulator.set(accum) + try: + yield None + finally: + _budget_accumulator.reset(prev) + +def accumulate_cost( + cost: float +): + # Accrue up the chain: the active center and (through parent) the pool. + accum = _budget_accumulator.get() + while accum is not None: + accum.curr_cost += cost + accum = accum.parent + +def budget_monitor( + *, + warn_threshold: float = BUDGET_PRESSURE_THRESHOLD, + warning_message: str | Callable[[StateVar], str] | None = None, + state_transformer: Callable[[StateVar], dict[str, Any]] | None = None, + on_overbudget: Callable[[], None] | None = None +) -> StateMonitor[StateVar]: + accum = _budget_accumulator.get() + if accum is None: + return lambda _ign: (None, None) + warned = False + def monitor( + curr_state: StateVar + ) -> MonitorReturn: + nonlocal warned + if accum.overbudget() and on_overbudget is not None: + on_overbudget() + if warned or not accum.pressured(warn_threshold): + return (None, None) + warned = True + msg : str + if warning_message is None: + msg = DEFAULT_BUDGET_PRESSURE_MESSAGE + elif isinstance(warning_message, str): + msg = warning_message + else: + msg = warning_message(curr_state) + + state_upd = None + if state_transformer is not None: + state_upd = state_transformer(curr_state) + return ([HumanMessage(msg)], state_upd) + return monitor + +def overbudget() -> bool: + res = _budget_accumulator.get() + if res is None: + return False + return res.overbudget() + + +def raise_budget_exceeded() -> None: + """``on_overbudget`` callback for agents that opt into the hard stop.""" + raise BudgetExceeded( + "Token cost budget exhausted; the agent was cooperatively terminated." + ) + + +def budget_pressure() -> bool: + """Whether the active budget is inside its wrap-up window: accrued cost at + or past ``BUDGET_PRESSURE_THRESHOLD`` of the phase cap *or* of the run + pool, whichever trips first. False when no budget is installed. Use this + to skip launching work that would only be told to immediately pack it in + (e.g. further property-extraction rounds).""" + res = _budget_accumulator.get() + if res is None: + return False + return res.pressured() + + +def pressure_abort_monitor() -> StateMonitor[MessagesState]: + """Monitor for auxiliary agents (feedback judges) that should not outlive + the main agent's wrap-up window: raises ``BudgetPressureAbort`` between + turns once budget pressure sets in. The tool that launched the agent + catches the exception and returns a canned "terminated for budget" + result. Reads the budget at call time, so it can be attached to a graph + compiled outside any budget scope.""" + def monitor(_curr_state: StateVar) -> MonitorReturn: + if budget_pressure(): + raise BudgetPressureAbort() + return (None, None) + return monitor diff --git a/composer/diagnostics/cost_callback.py b/composer/diagnostics/cost_callback.py new file mode 100644 index 00000000..bb8217eb --- /dev/null +++ b/composer/diagnostics/cost_callback.py @@ -0,0 +1,84 @@ +"""LangChain callback that accumulates the running USD cost of an LLM's calls. + +Sibling to :class:`composer.diagnostics.usage_callback.UsageCallback`: attached at +model construction so it fires for *every* ``invoke`` / ``ainvoke`` through the +model. Where ``UsageCallback`` records raw token counts, this one prices each +response through a :data:`~composer.llm.pricing.PriceProvider` (the model's pricing +curve, curried on model name) and adds the result to a running total. + +Implemented as an :class:`~langchain_core.callbacks.AsyncCallbackHandler` with +``run_inline = True``. The async ``on_llm_end`` is awaited on the event-loop thread +either way, but ``run_inline`` also decides *which* context it runs in: without it, +``ahandle_event`` dispatches the handler through ``asyncio.gather`` — each coroutine +wrapped in a ``Task`` against a ``copy_context()`` snapshot, so any ``ContextVar`` +the handler *sets* lands in that throwaway copy and never reaches the caller. With +``run_inline`` the handler is instead awaited directly in the caller's task and +context (``manager.ahandle_event`` line ~437), so its contextvar reads and writes +are visible to the surrounding LLM call. This matters because the accumulator +participates in contextvar state, not just its own counter. On the sync ``invoke`` +path LangChain still drives the coroutine to completion.""" + +from typing import Any + +from langchain_core.callbacks import AsyncCallbackHandler +from langchain_core.messages import AIMessage +from langchain_core.outputs import ChatGeneration, LLMResult + +from graphcore.utils import get_normalized_token_usage +from composer.llm.pricing import PriceProvider +from .budget import accumulate_cost + + +class CostAccumulator(AsyncCallbackHandler): + """Prices each LLM response and accumulates the total into :attr:`total_cost`. + + ``price_provider`` maps a call's input-token count to its per-MTok + :class:`~composer.llm.pricing.PriceTier`. ``long_cache`` selects the 1-hour + cache-write rate over the 5-minute one; the whole conversation is assumed to + share a single cache TTL (see ``builder_for``).""" + + # Route through the direct-await dispatch branch so the handler runs in the + # caller's task/context (not a gather-spawned Task with a copied context): + # required for the handler's contextvar reads/writes to reach the caller. + run_inline = True + + def __init__(self, price_provider: PriceProvider, *, long_cache: bool = False) -> None: + self._price_provider = price_provider + self._long_cache = long_cache + self.total_cost: float = 0.0 + + async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: + try: + generation = response.generations[0][0] + except IndexError: + return + if not isinstance(generation, ChatGeneration): + return + msg = generation.message + if isinstance(msg, AIMessage): + accumulate_cost(self._cost_of(msg)) + + def _cost_of(self, msg: AIMessage) -> float: + """USD cost of a single response, in dollars. Zero for models with no + pricing-table entry.""" + usage = get_normalized_token_usage(msg) + tier = self._price_provider(usage["total_input_tokens"]) + if tier is None: + return 0.0 + + cache_read = usage["cache_read_tokens"] + cache_write = usage["cache_write_tokens"] + # Fresh input is the total minus the two cache buckets, which the tier + # prices separately. Clamp against provider rounding wobble. + fresh_input = max(0, usage["total_input_tokens"] - cache_read - cache_write) + cache_write_rate = tier.cache_write_1h if self._long_cache else tier.cache_write + + # thinking_tokens are a subset of total_output_tokens and bill at the + # output rate, so they need no separate term here. + per_mtok = ( + fresh_input * tier.input + + cache_read * tier.cache_read + + cache_write * cache_write_rate + + usage["total_output_tokens"] * tier.output + ) + return per_mtok / 1_000_000 diff --git a/composer/foundry/author.py b/composer/foundry/author.py index f7596455..4c03440d 100644 --- a/composer/foundry/author.py +++ b/composer/foundry/author.py @@ -43,7 +43,11 @@ from composer.spec.cvl_generation import ( PropertyFeedbackProtocol, RebuttalBase, SkippedProperty, ) -from composer.spec.feedback import PropertyFeedback +from composer.spec.feedback import BUDGET_ABORT_FEEDBACK, PropertyFeedback +from composer.diagnostics.budget import ( + BudgetExceeded, BudgetPressureAbort, budget_monitor, budget_pressure, + pressure_abort_monitor, raise_budget_exceeded, +) from composer.spec.gen_types import TypedTemplate from composer.spec.graph_builder import bind_standard, run_to_completion from composer.spec.types import PropertyFormulation @@ -372,6 +376,15 @@ class FeedbackTool( async def run(self) -> Command | str: if self.state["curr_test"] is None: return "No test written" + if budget_pressure(): + # Don't launch a judge that would be terminated on its first + # monitor tick; the author's budget warning already tells it + # feedback approval is no longer required. + return ( + "Good? False\nFeedback:\nThe feedback judge was not run due to " + "budget constraints. See the system alert: feedback approval is " + "no longer required for this task." + ) with self.tool_deps() as deps: res = await deps.thunk( self.state["curr_test"], @@ -427,6 +440,8 @@ def did_rough_draft_read(s: ST, _) -> str | None: "foundry_property_judge_system_prompt.j2" ).with_tools( [*get_rough_draft_tools(ST), judge_ctx.get_memory_tool(), GetTestTool.as_tool("get_test")] + ).with_monitor( + pressure_abort_monitor() ).compile_async() async def thunk( @@ -458,17 +473,20 @@ async def thunk( f" Addressing: {r.prior_feedback_reference}\n" f" Evidence: {r.evidence}" ) - res = await run_to_completion( - workflow, - TestJudgeInput( - input=input_parts, curr_test=test_source, - memory=None, did_read=False, - ), - thread_id=uniq_thread_id("foundry-feedback"), - recursion_limit=judge_ctx.recursion_limit, - description="Foundry test feedback judge", - within_tool=within_tool, - ) + try: + res = await run_to_completion( + workflow, + TestJudgeInput( + input=input_parts, curr_test=test_source, + memory=None, did_read=False, + ), + thread_id=uniq_thread_id("foundry-feedback"), + recursion_limit=judge_ctx.recursion_limit, + description="Foundry test feedback judge", + within_tool=within_tool, + ) + except BudgetPressureAbort: + return BUDGET_ABORT_FEEDBACK assert "result" in res return res["result"] @@ -514,17 +532,22 @@ class PublishResultTool( async def run(self) -> Command | str: if (err := check_foundry_completion(self.state)) is not None: return err - ran = self.state["last_test_names"] - if ran is None: - # Unreachable in practice — the forge_test stamp required above - # implies a run recorded its test names — but defend anyway. - return "Completion REJECTED: no forge_test run has been recorded." - with self.tool_deps() as titles: - err = validate_property_tests( - self.property_tests, self.state["skipped"], titles, ran, - ) - if err is not None: - return err + if not budget_pressure(): + # The forge-ground-truth cross-check requires a recorded run; in + # the budget wrap-up window (where the agent is told to delete + # failing tests and publish without re-running forge) the declared + # mapping is accepted as-is. + ran = self.state["last_test_names"] + if ran is None: + # Unreachable in practice — the forge_test stamp required above + # implies a run recorded its test names — but defend anyway. + return "Completion REJECTED: no forge_test run has been recorded." + with self.tool_deps() as titles: + err = validate_property_tests( + self.property_tests, self.state["skipped"], titles, ran, + ) + if err is not None: + return err return tool_state_update( self.tool_call_id, "Accepted", @@ -625,6 +648,21 @@ class FoundryPropertyGenParams(TypedDict): "foundry_property_generation_prompt.j2" ) +_BUDGET_WRAPUP_MESSAGE = """ + +You have almost exceeded the token cost budget for this task. Wrap up IMMEDIATELY; +a partial test file is better than going over budget. Concretely: + +- The forge-test and feedback validation requirements on publishing have been lifted. You + no longer need approval from the feedback judge — ignore any pending or future feedback, + including a judge response saying it was terminated. +- Do NOT start new forge runs or research. +- Delete any tests that do not currently compile or pass. +- Skip (`record_skip`) every property you have not gotten to work, citing budget exhaustion. +- Then publish what remains via the `result` tool. + +""" + async def batch_foundry_test_generation( ctx: WorkflowContext[FoundryGeneration], @@ -700,6 +738,11 @@ async def batch_foundry_test_generation( .with_sys_prompt_template("foundry_property_generation_system_prompt.j2") .inject(lambda b: bound_template.render_to(b.with_initial_prompt_template)) .with_summary_config(FoundryGenerationSummaryConfig()) + .with_monitor(budget_monitor( + warning_message=_BUDGET_WRAPUP_MESSAGE, + state_transformer=lambda _s: {"required_validations": []}, + on_overbudget=raise_budget_exceeded, + )) ) graph = builder.compile_async() @@ -716,13 +759,16 @@ async def batch_foundry_test_generation( ) tid, mnem = await ctx.thread_and_mnemonic() - res_state = await run_to_completion( - graph, - init_state, - thread_id=tid, - description=f"{description} ({mnem})", - recursion_limit=ctx.recursion_limit, - ) + try: + res_state = await run_to_completion( + graph, + init_state, + thread_id=tid, + description=f"{description} ({mnem})", + recursion_limit=ctx.recursion_limit, + ) + except BudgetExceeded as e: + return GaveUp(reason=f"Foundry test generation terminated: {e}") assert "result" in res_state assert res_state["failed"] is not None diff --git a/composer/llm/anthropic.py b/composer/llm/anthropic.py index 61b86bc6..8b82dc8d 100644 --- a/composer/llm/anthropic.py +++ b/composer/llm/anthropic.py @@ -13,6 +13,7 @@ from composer.llm.provider import ( ProviderKind, CacheLevel, _ListIter, NoSuchElementError, ) +from composer.llm.pricing import PriceProvider, price_provider_for if TYPE_CHECKING: from langchain_core.language_models.chat_models import BaseChatModel @@ -126,17 +127,21 @@ class AnthropicModelProvider: model_name: str options: ModelConfiguration features: ModelFeatures - provider: ProviderKind = "anthropic" + price_provider: PriceProvider + provider: Literal["anthropic"] = "anthropic" @staticmethod def create(model_name: str, options: ModelConfiguration) -> "AnthropicModelProvider": - return AnthropicModelProvider(model_name, options, _model_parser(model_name)) + return AnthropicModelProvider( + model_name, options, _model_parser(model_name), price_provider_for(model_name) + ) def builder_for( self, *, cache_level: CacheLevel | None = None, disable_thinking: bool = False ) -> "BaseChatModel": from langchain_anthropic import ChatAnthropic from composer.diagnostics.usage_callback import UsageCallback + from composer.diagnostics.cost_callback import CostAccumulator opts = self.options thinking: dict[str, Any] | None @@ -174,5 +179,10 @@ def builder_for( betas=betas, thinking=thinking, model_kwargs=model_kwargs, - callbacks=[UsageCallback()], + callbacks=[ + UsageCallback(), + CostAccumulator( + self.price_provider, long_cache=cache_level == CacheLevel.LONG + ), + ], ) diff --git a/composer/llm/openai.py b/composer/llm/openai.py index ccca867c..0bbdc192 100644 --- a/composer/llm/openai.py +++ b/composer/llm/openai.py @@ -23,6 +23,7 @@ from composer.llm.provider import ( ProviderKind, CacheLevel, _ListIter, NoSuchElementError, ) +from composer.llm.pricing import PriceProvider, price_provider_for if TYPE_CHECKING: from langchain_core.language_models.chat_models import BaseChatModel @@ -175,16 +176,21 @@ class OpenAIModelProvider: model_name: str options: ModelConfiguration features: OpenAIModelFeatures + price_provider: PriceProvider provider: ProviderKind = "openai" @staticmethod def create(model_name: str, options: ModelConfiguration) -> "OpenAIModelProvider": - return OpenAIModelProvider(model_name, options, _model_parser(model_name)) + return OpenAIModelProvider( + model_name, options, _model_parser(model_name), price_provider_for(model_name) + ) def builder_for( self, *, cache_level: CacheLevel | None = None, disable_thinking: bool = False ) -> "BaseChatModel": from langchain_openai import ChatOpenAI + from composer.diagnostics.usage_callback import UsageCallback + from composer.diagnostics.cost_callback import CostAccumulator opts = self.options kwargs: dict[str, Any] = { @@ -205,5 +211,8 @@ def builder_for( temperature=1, timeout=None, max_retries=2, + # OpenAI has no cache-TTL knob, so long_cache stays False; cache_write_1h + # mirrors cache_write in the table anyway. + callbacks=[UsageCallback(), CostAccumulator(self.price_provider)], **kwargs, ) diff --git a/composer/llm/pricing.py b/composer/llm/pricing.py new file mode 100644 index 00000000..dc0c5ace --- /dev/null +++ b/composer/llm/pricing.py @@ -0,0 +1,127 @@ +from dataclasses import dataclass +from typing import Callable + +@dataclass(frozen=True) +class PriceTier: + """Per-million-token prices in USD for one model at one + context tier. + + ``input`` is the price for *fresh* input tokens (the bucket left + after subtracting ``cache_read`` and ``cache_write`` from the + total). ``output`` is the price for output tokens, which on the + OpenAI side already includes reasoning tokens (billed at the + output rate by both providers). ``cache_read`` is the + cache-hit rate. + + Cache writes come in two rates: ``cache_write`` is the 5-minute + ephemeral rate and ``cache_write_1h`` is the 1-hour ephemeral rate + (~2× base input on Anthropic). Which one applies depends on the + conversation's cache TTL. OpenAI has no separate cache-write rate, + so both fields carry the same value there.""" + input: float + output: float + cache_read: float + cache_write: float + cache_write_1h: float + + +@dataclass(frozen=True) +class ModelPricing: + """Pricing entry for one model family. ``long`` is the + long-context tier (used when input token count exceeds the + threshold) and applies only to OpenAI models that publish a + separate >272K-input rate; Anthropic models keep ``long = None`` + and bill everything at ``short`` rates.""" + short: PriceTier + long: PriceTier | None = None + + +# OpenAI's published >272K input-token threshold for long-context +# pricing. Once an individual call's input crosses this, the long +# tier applies *for the full session* per OpenAI's terms; we +# approximate that by switching on a per-message basis (a session +# that drifts above 272K will mostly stay there). +_OPENAI_LONG_CONTEXT_THRESHOLD = 272_000 + + +# Pricing tables transcribed from Anthropic + OpenAI rate cards. +# Sources should be re-checked when new model families ship. +_PRICING: list[tuple[str, ModelPricing]] = [ + # ---- Anthropic ---- + # claude-opus-4.5 / 4.6 / 4.7 share a rate card; older 4 / 4.1 + # are pricier. Matching by prefix-of-prefix so "claude-opus-4-7" + # and "claude-opus-4-7-20260301" both hit the right entry. + ("claude-opus-4-7", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25, 10.00))), + ("claude-opus-4-6", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25, 10.00))), + ("claude-opus-4-5", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25, 10.00))), + ("claude-opus-4-1", ModelPricing(short=PriceTier(15.00, 75.00, 1.50, 18.75, 30.00))), + ("claude-opus-4", ModelPricing(short=PriceTier(15.00, 75.00, 1.50, 18.75, 30.00))), + + ("claude-sonnet-4-6", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75, 6.00))), + ("claude-sonnet-4-5", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75, 6.00))), + ("claude-sonnet-4", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75, 6.00))), + + ("claude-haiku-4-5", ModelPricing(short=PriceTier(1.00, 5.00, 0.10, 1.25, 2.00))), + + # ---- OpenAI ---- + # gpt-5.5 / 5.4 publish short (≤272K input) and long (>272K) tiers. + # Pro variants don't publish a cached-in discount (cache_read = + # base input). Mini/nano don't publish a long tier; we use short + # for everything on those. OpenAI has no separate cache-write rate, + # so cache_write_1h mirrors cache_write on every OpenAI entry. + ("gpt-5.5-pro", ModelPricing( + short=PriceTier(30.00, 180.00, 30.00, 30.00, 30.00), + long=PriceTier(60.00, 270.00, 60.00, 60.00, 60.00), + )), + ("gpt-5.5", ModelPricing( + short=PriceTier(5.00, 30.00, 0.50, 5.00, 5.00), + long=PriceTier(10.00, 45.00, 1.00, 10.00, 10.00), + )), + ("gpt-5.4-pro", ModelPricing( + short=PriceTier(30.00, 180.00, 30.00, 30.00, 30.00), + long=PriceTier(60.00, 270.00, 60.00, 60.00, 60.00), + )), + ("gpt-5.4-mini", ModelPricing(short=PriceTier(0.75, 4.50, 0.075, 0.75, 0.75))), + ("gpt-5.4-nano", ModelPricing(short=PriceTier(0.20, 1.25, 0.02, 0.20, 0.20))), + ("gpt-5.4", ModelPricing( + short=PriceTier(2.50, 15.00, 0.25, 2.50, 2.50), + long=PriceTier(5.00, 22.50, 0.50, 5.00, 5.00), + )), +] + + +def price_per_mtok(model: str | None, input_tokens: int) -> PriceTier | None: + """Look up per-MTok pricing by model name and call size. Returns + ``None`` for models with no table entry (cost contribution becomes + zero — better than guessing). + + Matched by prefix on the lowercased model name so dated revisions + (``claude-opus-4-7-20260301``, ``gpt-5.5-2026-...``) collapse into + the same family entry. Table is searched in order, so list more + specific prefixes (``gpt-5.5-pro``) before less specific + (``gpt-5.5``). For OpenAI models with a long tier, ``input_tokens`` + chooses short vs. long; Anthropic always uses the short tier.""" + if model is None: + return None + m = model.lower() + for prefix, pricing in _PRICING: + if m.startswith(prefix): + if pricing.long is not None and input_tokens > _OPENAI_LONG_CONTEXT_THRESHOLD: + return pricing.long + return pricing.short + return None + + +# A model's pricing curve as a function of a single call's input-token count: +# ``price_per_mtok`` with the model name curried away. The remaining argument +# picks the short/long context tier (OpenAI); Anthropic ignores it. +type PriceProvider = Callable[[int], PriceTier | None] + + +def price_provider_for(model: str | None) -> PriceProvider: + """Curry :func:`price_per_mtok` on ``model``: returns a callable mapping a + call's input-token count to its :class:`PriceTier` (or ``None`` when the + model has no table entry, so its cost contribution is zero).""" + def provider(input_tokens: int) -> PriceTier | None: + return price_per_mtok(model, input_tokens) + return provider diff --git a/composer/pipeline/core.py b/composer/pipeline/core.py index 3c92a585..87142f66 100644 --- a/composer/pipeline/core.py +++ b/composer/pipeline/core.py @@ -17,7 +17,7 @@ import enum import logging from dataclasses import dataclass -from typing import Protocol +from typing import Protocol, cast from abc import ABC, abstractmethod from pydantic import BaseModel @@ -41,8 +41,10 @@ from composer.spec.source.report import build as report_build from composer.spec.source.task_ids import SYSTEM_ANALYSIS_TASK_ID, REPORT_TASK_ID from .ptypes import ( - BackendJob, BackendResult, ComponentOutcome, CorePhases, CorePipelineResult, Delivered, GaveUp, PipelineRun, SystemAnalysisSpec + BackendJob, BackendResult, ComponentOutcome, CorePhases, CorePipelineResult, + Delivered, GaveUp, PipelineRun, SystemAnalysisSpec, RunBudget ) +from composer.diagnostics.budget import total_budget, named_budget_or_nop COMMON_SYSTEM_CACHE_KEY = "system-analysis" @@ -154,32 +156,60 @@ async def run_pipeline[P: enum.Enum, FormT: BackendResult, H, A: ArtifactIdentif interactive: bool = False, threat_model: Document | None = None, max_bug_rounds: int = 3, + budget: RunBudget | None = None +) -> CorePipelineResult[FormT]: + if budget is not None: + with total_budget(budget.total, cast(dict[str, float], budget.caps)): + return await _run_pipeline_inner( + backend, run, interactive=interactive, + max_bug_rounds=max_bug_rounds, threat_model=threat_model + ) + else: + return await _run_pipeline_inner( + backend, run, interactive=interactive, + max_bug_rounds=max_bug_rounds, threat_model=threat_model + ) + +async def _run_pipeline_inner[P: enum.Enum, FormT: BackendResult, H, A: ArtifactIdentifier]( + backend: PipelineBackend[P, FormT, H, A], + run: PipelineRun[P, H], + *, + interactive: bool, + threat_model: Document | None, + max_bug_rounds: int ) -> CorePipelineResult[FormT]: spec, phases = backend.analysis_spec, backend.core_phases source = run.source # 1. System analysis (shared primitive, backend-parameterized; always yields SourceApplication). - analyzed = await run.runner( - TaskInfo(SYSTEM_ANALYSIS_TASK_ID, "System Analysis", phases["analysis"]), - lambda: run_component_analysis( - ty=SourceApplication, child_ctxt=run.ctx.child(CacheKey(spec.analysis_key)), - input=source, env=run.env, extra_input=[ - f"The main entry point of this application has been explicitly identified as {source.contract_name} at relative path {source.relative_path}. " - "Your output MUST contain an explicit contract instance with this solidity identifier.", - *spec.extra_input - ], - expected_main_id=source.contract_name, - ), - ) + with named_budget_or_nop("system_analysis"): + analyzed = await run.runner( + TaskInfo(SYSTEM_ANALYSIS_TASK_ID, "System Analysis", phases["analysis"]), + lambda: run_component_analysis( + ty=SourceApplication, child_ctxt=run.ctx.child(CacheKey(spec.analysis_key)), + input=source, env=run.env, extra_input=[ + f"The main entry point of this application has been explicitly identified as {source.contract_name} at relative path {source.relative_path}. " + "Your output MUST contain an explicit contract instance with this solidity identifier.", + *spec.extra_input + ], + expected_main_id=source.contract_name, + ), + ) if analyzed is None: raise ValueError("System analysis produced no result.") # 2. Backend transform + main-contract location (prover: harness lift; foundry: identity). - prepared = await backend.prepare_system(analyzed, run) + with named_budget_or_nop("system_preparation"): + prepared = await backend.prepare_system(analyzed, run) # 3. Pre-formalization setup runs CONCURRENTLY with extraction (neither needs the other) — # this preserves the prover's autosetup ∥ bug-analysis overlap, generically. - formalizer_task = asyncio.create_task(prepared.prepare_formalization(run)) + # The budget scope is entered inside the coroutine (not around create_task) so the + # cost-center binding lives in the spawned task's own context. + async def _prepare_formalization() -> Formalizer[FormT]: + with named_budget_or_nop("formalization_preparation"): + return await prepared.prepare_formalization(run) + formalizer_task = asyncio.create_task(_prepare_formalization()) batches = await _extract_all( backend.analysis_spec.properties_key, @@ -201,14 +231,15 @@ async def _run(batch: _Batch) -> ComponentOutcome[FormT]: result : FormT | GaveUp if cached_result is None: label = f"{batch.feat.component.name} ({len(batch.props)} properties)" - result : FormT | GaveUp = await run.runner( - TaskInfo( - formalize_task_id(batch.feat.ind), - f"{batch.feat.component.name} ({len(batch.props)} properties)", - phases["formalization"] - ), - lambda: formalizer.formalize(label, batch.feat, batch.props, child, run), - ) + with named_budget_or_nop("formalization"): + result : FormT | GaveUp = await run.runner( + TaskInfo( + formalize_task_id(batch.feat.ind), + f"{batch.feat.component.name} ({len(batch.props)} properties)", + phases["formalization"] + ), + lambda: formalizer.formalize(label, batch.feat, batch.props, child, run), + ) if not isinstance(result, GaveUp): await child.cache_put(result) else: @@ -265,12 +296,13 @@ async def _one(idx: int) -> _Batch | None: feat = ContractComponentInstance(_contract=main, ind=idx) feat_ctx = await prop_ctx.child(_component_cache_key(feat), {"component": feat.component.model_dump()}) - props = await run.runner( - TaskInfo(extract_task_id(idx), feat.component.name, phase), - lambda conv: run_property_inference( - feat_ctx, run.env, feat, refinement=conv if interactive else None, - threat_model=threat_model, max_rounds=max_rounds, backend_guidance=backend_guidance), - ) + with named_budget_or_nop("property_extraction"): + props = await run.runner( + TaskInfo(extract_task_id(idx), feat.component.name, phase), + lambda conv: run_property_inference( + feat_ctx, run.env, feat, refinement=conv if interactive else None, + threat_model=threat_model, max_rounds=max_rounds, backend_guidance=backend_guidance), + ) return _Batch(feat, props, feat_ctx) if props else None got = await asyncio.gather(*[_one(i) for i in range(len(main.contract.components))]) diff --git a/composer/pipeline/ptypes.py b/composer/pipeline/ptypes.py index e665c790..959e8253 100644 --- a/composer/pipeline/ptypes.py +++ b/composer/pipeline/ptypes.py @@ -117,6 +117,24 @@ def all_failed(self) -> bool: driver raises before returning in the no-outcomes case anyway).""" return bool(self.outcomes) and self.n_delivered == 0 +class PhaseBudget(TypedDict): + """Per-phase spending *caps* (USD). Ceilings, not allotments: they bound + how much a single phase may hog, and need not sum to the run total — + unspent phase money stays in the run pool for later phases.""" + formalization_preparation: float + system_analysis: float + system_preparation: float + property_extraction: float + formalization: float + + +@dataclass(frozen=True) +class RunBudget: + """The run's token-cost budget: ``total`` is the pool (the real bound on + overall spend); ``caps`` are the per-phase ceilings drawn against it.""" + total: float + caps: PhaseBudget + __all__ = [ "CorePipelineResult", "ComponentOutcome", diff --git a/composer/prover/core.py b/composer/prover/core.py index cf1a492b..980e3c66 100644 --- a/composer/prover/core.py +++ b/composer/prover/core.py @@ -45,7 +45,7 @@ from composer.prover.analysis import analyze_cex_raw from composer.prover.cloud import CloudJobError, cloud_results -from composer.prover.ptypes import RuleResult +from composer.prover.ptypes import RuleResult, RulePath, StatusCodes from composer.prover.results import read_and_format_run_result from composer.templates.loader import load_jinja_template from composer.prover.prover_protocol import ProverResult @@ -118,10 +118,20 @@ class ProverReport: ``link`` is the prover run's URL (cloud) or local results directory. """ - rule_status: dict[str, bool] + raw_rule_status: dict[RulePath, StatusCodes] + result_str: str link: str + @property + def rule_status(self) -> dict[str, bool]: + to_ret = {} + for (k, v) in self.raw_rule_status.items(): + if k.rule in to_ret and not to_ret[k.rule]: + continue + to_ret[k.rule] = v == "VERIFIED" + return to_ret + @property def all_verified(self) -> bool: return all(self.rule_status.values()) @@ -528,8 +538,12 @@ async def run_prover( continue prover_report[rule_name] = i.status == "VERIFIED" + raw_rule_results : dict[RulePath, StatusCodes] = { + k.path: k.status for k in parsed.values() + } + return ProverReport( - rule_status=prover_report, + raw_rule_status=raw_rule_results, result_str=result_str, link=run_result["link"], ) diff --git a/composer/spec/context.py b/composer/spec/context.py index 623c3f9c..73439014 100644 --- a/composer/spec/context.py +++ b/composer/spec/context.py @@ -21,7 +21,7 @@ from composer.io.mnemonic_store import assign_mnemonic from composer.core.user import user_data_ns from composer.spec.system_model import SolidityIdentifier - +from composer.diagnostics.budget import budget_pressure # --------------------------------------------------------------------------- # Workflow input types @@ -240,6 +240,11 @@ async def cache_get(self, ty: type[K]) -> K | None: async def cache_put(self, value: K) -> None: """Put a typed value in the cache. No-op if caching disabled.""" + if budget_pressure(): + # refuse to cache anything produced under budget pressure + # It's likely incomplete or "rushed", and a future run with more money + # should try again + return None if not isinstance(value, BaseModel): raise ValueError("Caching not allowed for non-basemodel keys") if self.cache_namespace is None: @@ -257,8 +262,3 @@ async def thread_and_mnemonic(self) -> tuple[str, str]: tid = self.thread_id mnem = await assign_mnemonic(tid, self._store, user_data_ns() + MNEMONIC_KEYS) return (tid, mnem) - -# --------------------------------------------------------------------------- -# Utility -# --------------------------------------------------------------------------- - diff --git a/composer/spec/cvl_generation.py b/composer/spec/cvl_generation.py index 3dd1a3d8..e2b6cd9e 100644 --- a/composer/spec/cvl_generation.py +++ b/composer/spec/cvl_generation.py @@ -31,6 +31,7 @@ from composer.spec.graph_builder import run_to_completion from composer.cvl.tools import put_cvl_raw, put_cvl, get_cvl, edit_cvl from composer.ui.tool_display import tool_display, suppress_ack +from composer.diagnostics.budget import budget_pressure class PropertyFeedbackProtocol(Protocol): @property @@ -310,6 +311,16 @@ async def run(self) -> Command: spec = st["curr_spec"] if spec is None: return tool_return(self.tool_call_id, "No spec put yet") + if budget_pressure(): + # Don't launch a judge that would be terminated on its first + # monitor tick; the author's budget warning already tells it + # feedback approval is no longer required. + return tool_return( + self.tool_call_id, + "Good? False\nFeedback The feedback judge was not run due to budget " + "constraints. See the system alert: feedback approval is no longer " + "required for this task.", + ) skipped = st["skipped"] t = await feedback(spec, skipped, self.rebuttals, self.tool_call_id) msg = f"Good? {t.good}\nFeedback {t.feedback}" diff --git a/composer/spec/feedback.py b/composer/spec/feedback.py index 1aceda7a..e6bada56 100644 --- a/composer/spec/feedback.py +++ b/composer/spec/feedback.py @@ -21,6 +21,7 @@ from composer.spec.cvl_generation import FeedbackToolContext, Rebuttal, SkippedProperty from composer.spec.system_model import ContractComponentInstance from composer.spec.util import uniq_thread_id +from composer.diagnostics.budget import BudgetPressureAbort, pressure_abort_monitor class PropertyFeedback(BaseModel): """ @@ -29,6 +30,19 @@ class PropertyFeedback(BaseModel): good: bool = Field(description="Whether the properties are good as is, or if there is room for improvement") feedback: str = Field(description="The feedback on the rule if work is needed. Can be empty if there is no feedback") + +# Canned judge verdict when the judge was not run (or was killed mid-run) +# because the author is in its budget wrap-up window. The author's own budget +# warning tells it feedback approval is no longer required, so a dead judge +# is expected there, not an error. +BUDGET_ABORT_FEEDBACK = PropertyFeedback( + good=False, + feedback=( + "The feedback judge was terminated due to budget constraints. " + "See the system alert: feedback approval is no longer required for this task." + ), +) + class Properties(TypedDict): properties: list[PropertyFormulation] @@ -98,7 +112,9 @@ def did_rough_draft_read(s: ST, _) -> str | None: lambda b: final_prompt.render_to(b.with_initial_prompt_template) ).inject( lambda g: system_prompt.render_to(g.with_sys_prompt_template) - ).with_tools([*rough_draft_tools, mem, get_cvl(ST), ]).compile_async() + ).with_tools([*rough_draft_tools, mem, get_cvl(ST), ]).with_monitor( + pressure_abort_monitor() + ).compile_async() async def the_tool( cvl: str, @@ -134,14 +150,17 @@ async def the_tool( f" Addressing: {r.prior_feedback_reference}\n" f" Evidence: {r.evidence}" ) - res = await run_to_completion( - workflow, - SpecJudgeInput(input=input_parts, curr_spec=cvl, memory=None, did_read=False), - thread_id=uniq_thread_id("feedback"), - recursion_limit=ctx.recursion_limit, - description="Property feedback judge", - within_tool=within_tool, - ) + try: + res = await run_to_completion( + workflow, + SpecJudgeInput(input=input_parts, curr_spec=cvl, memory=None, did_read=False), + thread_id=uniq_thread_id("feedback"), + recursion_limit=ctx.recursion_limit, + description="Property feedback judge", + within_tool=within_tool, + ) + except BudgetPressureAbort: + return BUDGET_ABORT_FEEDBACK assert "result" in res return res["result"] diff --git a/composer/spec/prop_inference.py b/composer/spec/prop_inference.py index 3edc92c0..b13fc6a8 100644 --- a/composer/spec/prop_inference.py +++ b/composer/spec/prop_inference.py @@ -25,6 +25,7 @@ from composer.tools.thinking import RoughDraftState, get_rough_draft_tools from composer.spec.service_host import Sort, ServiceHost from composer.io.conversation import ConversationContextProvider +from composer.diagnostics.budget import budget_monitor, budget_pressure from composer.spec.refinement import refinement_loop, EndConversation, SyncStateUpdateTool from composer.templates.loader import load_jinja_template from composer.ui.tool_display import tool_display @@ -312,6 +313,8 @@ class ST(MessagesState, RoughDraftState): env.analysis_tools ).with_sys_prompt_template( "property_analysis_system_prompt.j2", sort=env.sort + ).with_monitor( + budget_monitor() ).compile_async() flow_input: BugAnalysisInput = BugAnalysisInput( @@ -365,6 +368,11 @@ async def _run_bug_analysis_inner( last_round_convo : list[AnyMessage] | None = None for i in range(0, max_rounds): + # Under budget pressure a fresh round would be told to pack it in on + # its first monitor tick — don't bother launching it. Round 0 always + # runs (the loop's invariants require at least one round's history). + if i > 0 and budget_pressure(): + break next_result = await _run_bug_round( env, component, front_matter_items, agent_component_analysis, i, prev_rounds, backend_guidance, diff --git a/composer/spec/source/author.py b/composer/spec/source/author.py index d0bf3cd3..f092db49 100644 --- a/composer/spec/source/author.py +++ b/composer/spec/source/author.py @@ -1,4 +1,4 @@ -from typing import NotRequired, override, Literal, Annotated +from typing import Callable, NotRequired, override, Literal, Annotated from typing_extensions import TypedDict import json @@ -32,8 +32,12 @@ from langgraph.types import Command from composer.spec.feedback import property_feedback_judge, FeedbackTemplate from composer.ui.tool_display import tool_display +from composer.diagnostics.budget import ( + BudgetExceeded, budget_monitor, raise_budget_exceeded, +) +from .monitor import monitor -from graphcore.graph import FlowInput +from graphcore.graph import FlowInput, MonitorReturn class SourceAuthorExtra(TypedDict): failed: bool | None @@ -331,6 +335,41 @@ async def run(self) -> Command | str: _PropertyGenTemplate = TypedTemplate[PropertyGenParams]("property_generation_prompt.j2") +_BUDGET_WRAPUP_MESSAGE = """ + +You have almost exceeded the token cost budget for this task. Wrap up IMMEDIATELY; +a partial spec is better than going over budget. Concretely: + +- The feedback and prover validation requirements on publishing have been lifted. You no + longer need approval from the feedback judge — ignore any pending or future feedback, + including a judge response saying it was terminated. +- Do NOT start new prover runs or research. +- Delete any rules/invariants that do not currently work. +- Skip (`record_skip`) every property you have not gotten to work, citing budget exhaustion. +- Then publish what remains via the `result` tool. + +""" + + +def _author_monitor() -> "Callable[[SourceCVLGenerationState], MonitorReturn]": + """The author's monitor: budget wrap-up takes precedence; otherwise the + usual reminders-channel drain. On the (single) turn the budget warning + fires any pending reminders are dropped — moot, since the warning tells + the agent to ignore prover/feedback outcomes anyway.""" + b_monitor = budget_monitor( + warning_message=_BUDGET_WRAPUP_MESSAGE, + state_transformer=lambda _s: {"required_validations": []}, + on_overbudget=raise_budget_exceeded, + ) + + def combined(curr_state: SourceCVLGenerationState) -> MonitorReturn: + msgs, upd = b_monitor(curr_state) + if msgs is None and upd is None: + return monitor(curr_state) + return msgs, upd + + return combined + async def batch_cvl_generation( ctx: WorkflowContext[CVLGeneration], init_config: dict, @@ -376,6 +415,8 @@ async def batch_cvl_generation( [prover_tool, ExpectRulePassage.as_tool("expect_rule_passage"), ExpectRuleFailure.as_tool("expect_rule_failure"), GiveUpTool.as_tool("give_up"), PublishResultTool.as_tool("result"), ctx.get_memory_tool()] ).with_state( SourceCVLGenerationState + ).with_monitor( + _author_monitor() ).with_output_key( "result" ).with_input( @@ -395,24 +436,29 @@ async def batch_cvl_generation( }), props ) - res_state = await run_cvl_generator( - ctx = ctx, - d = task_graph, - description=description, - ctxt=feedback_env, - in_state=SourceCVLGenerationInput( - curr_spec=None, - config=init_config, - spec_stem=spec_stem, - input=[], - required_validations=[FEEDBACK_VALIDATION_KEY, PROVER_VALIDATION_KEY], - rule_skips={}, - skipped=[], - property_rules=[], - validations={}, - failed=None, + try: + res_state = await run_cvl_generator( + ctx = ctx, + d = task_graph, + description=description, + ctxt=feedback_env, + in_state=SourceCVLGenerationInput( + curr_spec=None, + config=init_config, + spec_stem=spec_stem, + input=[], + required_validations=[FEEDBACK_VALIDATION_KEY, PROVER_VALIDATION_KEY], + rule_skips={}, + skipped=[], + property_rules=[], + validations={}, + failed=None, + prover_history=[], + reminders_channel=[] + ) ) - ) + except BudgetExceeded as e: + return GaveUp(reason=f"CVL generation terminated: {e}") assert "result" in res_state assert res_state["failed"] is not None diff --git a/composer/spec/source/monitor.py b/composer/spec/source/monitor.py new file mode 100644 index 00000000..2fcded3b --- /dev/null +++ b/composer/spec/source/monitor.py @@ -0,0 +1,20 @@ +from typing import TYPE_CHECKING + +from langgraph.config import get_config +from langchain_core.messages import HumanMessage +from graphcore.graph import MonitorReturn +from composer.prover.ptypes import StatusCodes + +if TYPE_CHECKING: + # Runtime import would be circular: author.py imports this module. + from .author import SourceCVLGenerationState + +def monitor( + curr_state: "SourceCVLGenerationState" +) -> MonitorReturn: + if not curr_state["reminders_channel"]: + return None, None + + return [HumanMessage(f"{'\n'.join(curr_state["reminders_channel"])}")], { + "reminders_channel": None + } diff --git a/composer/spec/source/prover.py b/composer/spec/source/prover.py index 0b170fed..01c07a15 100644 --- a/composer/spec/source/prover.py +++ b/composer/spec/source/prover.py @@ -14,22 +14,24 @@ import time from contextlib import contextmanager, nullcontext from pathlib import Path -from typing import Annotated, Callable, Iterator, override, AsyncContextManager +from typing import Annotated, Callable, Iterator, override, AsyncContextManager, Sequence, Literal from typing_extensions import TypedDict, NotRequired from langchain_core.tools import InjectedToolCallId, tool, BaseTool +from langchain_core.messages import AIMessage from langgraph.prebuilt import InjectedState -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, Discriminator from langgraph.config import get_stream_writer from langgraph.types import Command -from composer.prover.ptypes import RuleResult +from composer.prover.ptypes import RuleResult, RulePath from graphcore.graph import LLM from composer.prover.core import ( - ProverOptions, ProverCallbacks, run_prover, DefaultCexHandler + ProverOptions, run_prover, DefaultCexHandler ) from composer.prover.callbacks import ProverEventCallbacks +from composer.prover.ptypes import StatusCodes from composer.ui.tool_display import tool_display from composer.diagnostics.stream import ( ProverOutputEvent, CloudPollingEvent, RuleAnalysisResult, @@ -40,6 +42,7 @@ from graphcore.graph import tool_state_update from composer.spec.util import temp_certora_file from composer.spec.gen_types import CERTORA_DIR, SPECS_DIR +from composer.spec.util import string_hash _logger = logging.getLogger("composer.prover") @@ -76,6 +79,23 @@ def _merge_rule_skips(left: dict[str, str], right: dict[str, str]) -> dict[str, to_ret[k] = v return to_ret +class ProverRunLog(TypedDict): + tool_call_id: str + prover_results: list[tuple[RulePath, StatusCodes]] + spec_digest: str + rules: list[str] | None + sort: Literal["run"] + +class NagMarker(TypedDict): + nagged_rules: list[RulePath] + sort: Literal["nag"] + +type ProverHistoryItem = Annotated[ProverRunLog | NagMarker, Discriminator("sort")] + +def _merge_prover_history(left: list[ProverHistoryItem], right: list[ProverHistoryItem]) -> list[ProverHistoryItem]: + to_ret = left.copy() + to_ret.extend(right) + return to_ret class ProverStateExtra(TypedDict): rule_skips: Annotated[dict[str, str], _merge_rule_skips] @@ -86,6 +106,9 @@ class ProverStateExtra(TypedDict): # Basename the spec is materialized/persisted under (e.g. "autospec_"). # NotRequired so other ProverStateExtra injectors (e.g. config_edit) needn't set it. spec_stem: NotRequired[str] + prover_history: Annotated[list[ProverHistoryItem], _merge_prover_history] + reminders_channel: list[str] + type ProverEvents = CEXAnalysisStart | CloudPollingEvent | ProverOutputEvent | RuleAnalysisResult | ProverRun | ProverLink | ProverResult @@ -201,15 +224,8 @@ def tmp_spec( def _prover_sem(cloud: bool) -> AsyncContextManager[None]: if not cloud: return asyncio.Semaphore(1) - - class ToRet(): - async def __aenter__(self): - return - - async def __aexit__(self, exc_type, exc, tb): - return - - return ToRet() + else: + return nullcontext() def get_prover_tool( llm: LLM, @@ -235,8 +251,25 @@ async def verify_spec( state: Annotated[StateWithSkips, InjectedState], rules: list[str] | None = None ) -> str | Command: + last_msg = state["messages"][-1] + if isinstance(last_msg, AIMessage) and any( + i["id"] != tool_call_id for i in last_msg.tool_calls + ): + return "Cannot call the verify_spec tool in parallel with other tool calls. verify_spec must be the only tool you call in a turn" + if state["curr_spec"] is None: return "Specification not yet put on VFS" + + spec_hash = string_hash( + state["curr_spec"] + ) + + if len(state['prover_history']) > 0: + last_item = state["prover_history"][-1] + if last_item["sort"] == "run" and any(i == "TIMEOUT" for (_,i) in last_item) and last_item["spec_digest"] == spec_hash: + return "Refusing to re-run prover on identical spec with a known TIMEOUT result; timeouts are not transient " \ + "errors and will not go away by re-running the tool." + conf = state["config"] # With a seeded stem, name the spec/conf after it (so on-disk names match the # dump) under a lock; else fall back to unique uid names (no lock needed). @@ -255,7 +288,7 @@ async def verify_spec( summary = get_run_summary() component = (spec_stem or main_contract).removeprefix("autospec_") - iteration = summary.prover_total_calls + 1 + iteration = len(state["prover_history"]) + 1 config["msg"] = f"{component} iteration number {iteration}" with temp_certora_file( @@ -276,22 +309,101 @@ async def verify_spec( DefaultCexHandler(llm, state, summarization_threshold=10) ) - if isinstance(result, str): - return result - all_verified = True - for (r, stat) in result.rule_status.items(): - if r in state["rule_skips"]: - continue - if not stat: - all_verified = False - break - if rules is None and all_verified: - return tool_state_update( - tool_call_id=tool_call_id, content=result.result_str, - prover_link=result.link, validations=stamper(state), - ) - return tool_state_update( - tool_call_id=tool_call_id, content=result.result_str, prover_link=result.link + if isinstance(result, str): + return result + + all_verified = True + for (r, stat) in result.rule_status.items(): + if r in state["rule_skips"]: + continue + if not stat: + all_verified = False + break + + spec_hash = string_hash( + state["curr_spec"] + ) + + stuck_rules = { + k: v for (k,v) in result.raw_rule_status.items() if v in ("TIMEOUT", "ERROR", "SANITY_FAILED") and k.rule not in state["rule_skips"] + } + + stuck_count = { + k: 1 for k in stuck_rules.keys() + } + + to_warn : set[RulePath] = set() + + known_tc_ids = { + l["id"] + for msg in state["messages"] if isinstance(msg, AIMessage) + for l in msg.tool_calls if l["name"] == "verify_spec" + } + + seen_post_compaction_history = False + + history_ind = len(state["prover_history"]) - 1 + while history_ind >= 0 and len(stuck_count) > 0: + it = state["prover_history"][history_ind] + if it["sort"] == "nag": + for r in it["nagged_rules"]: + del stuck_count[r] + history_ind -= 1 + continue + assert it["sort"] == "run" + if it["tool_call_id"] not in known_tc_ids: + seen_post_compaction_history = True + target_rules = set(it["rules"]) if it["rules"] else None + for k in stuck_count.keys(): + if target_rules is not None and k.rule not in target_rules: + continue + if not any( + rp == k and stuck_rules[k] == stat for (rp, stat) in it["prover_results"] + ): + del stuck_count[k] + else: + stuck_count[k] += 1 + if stuck_count[k] == 3: + to_warn.add(k) + del stuck_count[k] + + prover_update : list[ProverHistoryItem] = [ + ProverRunLog( + tool_call_id=tool_call_id, + prover_results=[(k, v) for (k,v) in result.raw_rule_status.items()], + rules=rules, + spec_digest=spec_hash, + sort="run" + ) + ] + nag_channel = { + + } + if len(to_warn) > 0: + prover_update.append(NagMarker( + sort="nag", + nagged_rules=list(to_warn) + )) + nag_channel["reminders_channel"] = [ + "The following rule(s) have had identical failures on the last 3 runs of the prover:", + *(f"- {it.pprint()}" for it in to_warn), + "You may need to significantly change your approach, or skip the property if this is a persistent issue (you may need to use rebuttals to communicate" + " these failures to the feedback judge)." + ] + if seen_post_compaction_history: + nag_channel["reminders_channel"].append( + "(NB: Some of these prover calls happened before your most recent task history summarization)" ) + + if rules is None and all_verified: + return tool_state_update( + tool_call_id=tool_call_id, content=result.result_str, + prover_link=result.link, validations=stamper(state), + prover_history=prover_update + ) + return tool_state_update( + tool_call_id=tool_call_id, content=result.result_str, prover_link=result.link, + prover_history=prover_update, **nag_channel + ) return verify_spec diff --git a/composer/spec/system_analysis.py b/composer/spec/system_analysis.py index 57953825..3d713a1e 100644 --- a/composer/spec/system_analysis.py +++ b/composer/spec/system_analysis.py @@ -11,6 +11,7 @@ from composer.spec.service_host import ServiceHost from composer.spec.util import slugify_filename from composer.tools.thinking import RoughDraftState, get_rough_draft_tools +from composer.diagnostics.budget import budget_monitor DESCRIPTION = "Component analysis" @@ -132,7 +133,7 @@ def _validation_wrapper( "application_analysis_prompt.j2", sort=env.sort, has_doc=input is not None - ) + ).with_monitor(budget_monitor()) graph = b.compile_async() inputs : list[str | dict] = [] diff --git a/composer/ui/message_renderer.py b/composer/ui/message_renderer.py index f4d4d0f0..f1e0c0ea 100644 --- a/composer/ui/message_renderer.py +++ b/composer/ui/message_renderer.py @@ -24,6 +24,7 @@ from graphcore.graph import INITIAL_NODE, TOOL_RESULT_NODE, TOOLS_NODE from graphcore.utils import NormalizedTokenUsage, get_normalized_token_usage +from composer.llm.pricing import price_per_mtok KNOWN_NODES: set[str] = {INITIAL_NODE, TOOL_RESULT_NODE, TOOLS_NODE} @@ -42,115 +43,6 @@ def dot(style: str, text: Text | str) -> Text: return result -@dataclass(frozen=True) -class _PriceTier: - """Per-million-token prices in USD for one model at one - context tier. - - ``input`` is the price for *fresh* input tokens (the bucket left - after subtracting ``cache_read`` and ``cache_write`` from the - total). ``output`` is the price for output tokens, which on the - OpenAI side already includes reasoning tokens (billed at the - output rate by both providers). ``cache_read`` / ``cache_write`` - are the cache-bucket rates. - - Anthropic-specific note: the ``cache_write`` rate here is the - 5-minute ephemeral rate. Anthropic also has a 1-hour ephemeral - rate that's higher (~2× base input); our normalized usage rolls - both into a single ``cache_write_tokens`` bucket, so workloads - that heavily use 1h caching will be slightly under-billed by - this banner. Approximation is fine for a banner.""" - input: float - output: float - cache_read: float - cache_write: float - - -@dataclass(frozen=True) -class _ModelPricing: - """Pricing entry for one model family. ``long`` is the - long-context tier (used when input token count exceeds the - threshold) and applies only to OpenAI models that publish a - separate >272K-input rate; Anthropic models keep ``long = None`` - and bill everything at ``short`` rates.""" - short: _PriceTier - long: _PriceTier | None = None - - -# OpenAI's published >272K input-token threshold for long-context -# pricing. Once an individual call's input crosses this, the long -# tier applies *for the full session* per OpenAI's terms; we -# approximate that by switching on a per-message basis (a session -# that drifts above 272K will mostly stay there). -_OPENAI_LONG_CONTEXT_THRESHOLD = 272_000 - - -# Pricing tables transcribed from Anthropic + OpenAI rate cards. -# Sources should be re-checked when new model families ship. -_PRICING: list[tuple[str, _ModelPricing]] = [ - # ---- Anthropic ---- - # claude-opus-4.5 / 4.6 / 4.7 share a rate card; older 4 / 4.1 - # are pricier. Matching by prefix-of-prefix so "claude-opus-4-7" - # and "claude-opus-4-7-20260301" both hit the right entry. - ("claude-opus-4-7", _ModelPricing(short=_PriceTier(5.00, 25.00, 0.50, 6.25))), - ("claude-opus-4-6", _ModelPricing(short=_PriceTier(5.00, 25.00, 0.50, 6.25))), - ("claude-opus-4-5", _ModelPricing(short=_PriceTier(5.00, 25.00, 0.50, 6.25))), - ("claude-opus-4-1", _ModelPricing(short=_PriceTier(15.00, 75.00, 1.50, 18.75))), - ("claude-opus-4", _ModelPricing(short=_PriceTier(15.00, 75.00, 1.50, 18.75))), - - ("claude-sonnet-4-6", _ModelPricing(short=_PriceTier(3.00, 15.00, 0.30, 3.75))), - ("claude-sonnet-4-5", _ModelPricing(short=_PriceTier(3.00, 15.00, 0.30, 3.75))), - ("claude-sonnet-4", _ModelPricing(short=_PriceTier(3.00, 15.00, 0.30, 3.75))), - - ("claude-haiku-4-5", _ModelPricing(short=_PriceTier(1.00, 5.00, 0.10, 1.25))), - - # ---- OpenAI ---- - # gpt-5.5 / 5.4 publish short (≤272K input) and long (>272K) tiers. - # Pro variants don't publish a cached-in discount (cache_read = - # base input). Mini/nano don't publish a long tier; we use short - # for everything on those. - ("gpt-5.5-pro", _ModelPricing( - short=_PriceTier(30.00, 180.00, 30.00, 30.00), - long=_PriceTier(60.00, 270.00, 60.00, 60.00), - )), - ("gpt-5.5", _ModelPricing( - short=_PriceTier(5.00, 30.00, 0.50, 5.00), - long=_PriceTier(10.00, 45.00, 1.00, 10.00), - )), - ("gpt-5.4-pro", _ModelPricing( - short=_PriceTier(30.00, 180.00, 30.00, 30.00), - long=_PriceTier(60.00, 270.00, 60.00, 60.00), - )), - ("gpt-5.4-mini", _ModelPricing(short=_PriceTier(0.75, 4.50, 0.075, 0.75))), - ("gpt-5.4-nano", _ModelPricing(short=_PriceTier(0.20, 1.25, 0.02, 0.20))), - ("gpt-5.4", _ModelPricing( - short=_PriceTier(2.50, 15.00, 0.25, 2.50), - long=_PriceTier(5.00, 22.50, 0.50, 5.00), - )), -] - - -def _price_per_mtok(model: str | None, input_tokens: int) -> _PriceTier | None: - """Look up per-MTok pricing by model name and call size. Returns - ``None`` for models with no table entry (cost contribution becomes - zero — better than guessing). - - Matched by prefix on the lowercased model name so dated revisions - (``claude-opus-4-7-20260301``, ``gpt-5.5-2026-...``) collapse into - the same family entry. Table is searched in order, so list more - specific prefixes (``gpt-5.5-pro``) before less specific - (``gpt-5.5``). For OpenAI models with a long tier, ``input_tokens`` - chooses short vs. long; Anthropic always uses the short tier.""" - if model is None: - return None - m = model.lower() - for prefix, pricing in _PRICING: - if m.startswith(prefix): - if pricing.long is not None and input_tokens > _OPENAI_LONG_CONTEXT_THRESHOLD: - return pricing.long - return pricing.short - return None - class TokenStats: """Accumulates token usage across AI messages and updates a display widget. @@ -177,7 +69,7 @@ def __init__(self, display: Static): @staticmethod def _cost_of(usage: NormalizedTokenUsage) -> float: - price = _price_per_mtok(usage["model_name"], usage["total_input_tokens"]) + price = price_per_mtok(usage["model_name"], usage["total_input_tokens"]) if price is None: return 0.0 # Fresh (non-cached) input is the total minus the cached diff --git a/graphcore b/graphcore index 8b9d3203..f4584f66 160000 --- a/graphcore +++ b/graphcore @@ -1 +1 @@ -Subproject commit 8b9d320316a2cfbfa775da0a961501539a4124f0 +Subproject commit f4584f667644f856191bc04dd46e0cd184036f55 diff --git a/pyproject.toml b/pyproject.toml index f999a88f..89d96261 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ authors = [ ] dependencies = [ - "graphcore @ git+ssh://git@github.com/Certora/graphcore.git@8b9d320316a2cfbfa775da0a961501539a4124f0", + "graphcore @ git+ssh://git@github.com/Certora/graphcore.git@f4584f667644f856191bc04dd46e0cd184036f55", "aiohttp>=3.13", "attrs>=26.1", "Jinja2>=3.1",