Model

async def after_invoke(self, response: InvocationResponse) -> None:
    """LLMeter Callback.after_invoke hook

    Calls calculate_request_cost() with `save=True` to save the cost on the InvocationResponse.
    """
    await self.calculate_request_cost(response, save=True)

async def after_run(self, result: Result) -> None:
    """LLMeter Callback.after_run hook

    Calls calculate_run_cost() with `save=True` to save the cost on the Result.
    """
    await self.calculate_run_cost(
        result, recalculate_request_costs=False, save=True
    )

async def before_invoke(self, payload: dict) -> None:
    """This LLMeter Callback hook is a no-op for CostModel"""
    pass

async def before_run(self, run_config: _RunConfig) -> None:
    """Initialize state for all run-level cost dimensions in the model"""
    for dim in self.run_dims.values():
        await dim.before_run_start(run_config)

async def calculate_request_cost(
    self,
    response: InvocationResponse,
    save: bool = False,
) -> CalculatedCostWithDimensions:
    """Calculate the costs of a single FM invocation (excluding any session-level costs)

    Args:
        response: The InvocationResponse to estimate costs for
        save: Set `True` to also store the result in `response.cost`, in addition to returning
            it. Defaults to `False`
    """
    dim_costs = CalculatedCostWithDimensions(
        **{
            name: await dim.calculate(response)
            for name, dim in self.request_dims.items()
        }
    )
    if save:
        dim_costs.save_on_namespace(response, key_prefix="cost_")

    return dim_costs

async def calculate_run_cost(
    self,
    result: Result,
    recalculate_request_costs: bool = True,
    save: bool = False,
) -> CalculatedCostWithDimensions:
    """Calculate the run-level costs of a test (including any request-level costs)

    NOTE: Depending on the types of `run_dims` in your model, this may throw an error if
    `before_run` was not called first to initialise the state.

    Args:
        result: An LLMeter Run result
        recalculate_request_costs: Set `False` if your `result.responses` have already been
            annotated with costs in line with the current cost model. By default (`True`),
            the costs for each response will be re-calculated.
        save: Set `True` to also store the result in `result.cost`, in addition to returning
            it. Defaults to `False`.
    """
    run_cost = CalculatedCostWithDimensions(
        **{name: await dim.calculate(result) for name, dim in self.run_dims.items()}
    )

    if recalculate_request_costs:
        resp_costs = [
            await self.calculate_request_cost(r, save=save)
            for r in result.responses
        ]
    else:
        resp_costs = list(
            filter(
                lambda c: c,  # Skip responses where no cost data was found at all
                (
                    CalculatedCostWithDimensions.load_from_namespace(
                        r, key_prefix="cost_"
                    )
                    for r in result.responses
                ),
            ),
        )
    # Merge the total request-level costs into the run-level costs:
    # (Unless requests is empty, because sum([]) = 0 and not a CalculatedCostWithDimensions)
    if len(resp_costs):
        run_cost.merge(sum(resp_costs))  # type: ignore

    if save:
        # Save the overall run cost and breakdown on the main result object:
        run_cost.save_on_namespace(result, key_prefix="cost_")
        # Contribute both 1/ the summary stats of request-level costs, and 2/ the overall run
        # cost+breakdown, to result.stats:
        stats = CalculatedCostWithDimensions.summary_statistics(
            resp_costs,
            key_prefix="cost_",
            key_dim_name_suffix="_per_request",
            # cost_total_per_request would be confusing, so skip 'total':
            key_total_name_and_suffix="per_request",
        )
        run_cost.save_on_namespace(stats, key_prefix="cost_")
        result._update_contributed_stats(stats)

    return run_cost

def save_to_file(self, path: str) -> None:
    """Save the cost model (including all dimensions) to a JSON file"""
    with open(path, "w") as f:
        f.write(self.to_json())