Skip to main content
This is an interactive notebook. You can run it locally or use the links below:

​
Set up a custom cost model

This guide shows you how to define a custom model in Weave that reports its own token usage, register a custom per-token cost for that model, and then retrieve calls with cost information attached. Use this when you’re working with a model that isn’t covered by Weave’s built-in cost data and you want to track spending alongside your traces. Weave calculates costs based on the number of tokens used and the model used. Weave grabs this usage and model from the output and associates them with the call. Set up a simple custom model that calculates its own token usage and stores that in Weave.

​
Set up the environment

Install and import all needed packages. Set WANDB_API_KEY in your environment so that you can log in with wandb.login() (this should be given to the colab as a secret). Set the project in W&B you want to log this into in name_of_wandb_project.
name_of_wandb_project can also be in the format of {team_name}/{project_name} to specify a team to log the traces into.
Then fetch a Weave client by calling weave.init(). 
%pip install wandb weave datetime --quiet
python
import os

import wandb
from google.colab import userdata

import weave

os.environ["WANDB_API_KEY"] = userdata.get("WANDB_API_KEY")
name_of_wandb_project = "custom-cost-model"

wandb.login()
python
weave_client = weave.init(name_of_wandb_project)

​
Set up a model with Weave

Next, define a Weave Model subclass that performs its own token accounting. Returning the usage counts and model name in the output dictionary is what lets Weave later associate the call with the custom cost defined in the next section. 
from weave import Model

class YourModel(Model):
    attribute1: str
    attribute2: int

    def simple_token_count(self, text: str) -> int:
        return len(text) // 3

    # This is a custom op that we are defining
    # It takes in a string, and outputs a dict with the usage counts, model name, and the output
    @weave.op()
    def custom_model_generate(self, input_data: str) -> dict:
        # Model logic goes here
        # Here is where you would have a custom generate function
        prediction = self.attribute1 + " " + input_data

        # Usage counts
        prompt_tokens = self.simple_token_count(input_data)
        completion_tokens = self.simple_token_count(prediction)

        # We return a dictionary with the usage counts, model name, and the output
        # Weave will automatically associate this with the trace
        # This object {usage, model, output} matches the output of a OpenAI Call
        return {
            "usage": {
                "input_tokens": prompt_tokens,
                "output_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            },
            "model": "your_model_name",
            "output": prediction,
        }

    # In our predict function we call our custom generate function, and return the output.
    @weave.op()
    def predict(self, input_data: str) -> dict:
        # Here is where you would do any post processing of the data
        outputs = self.custom_model_generate(input_data)
        return outputs["output"]

​
Add a custom cost

Add a custom cost. Once you have added a custom cost and your calls have usage, you can fetch the calls with include_cost and review the costs under summary.weave.costs. 
model = YourModel(attribute1="Hello", attribute2=1)
model.predict("world")

# We then add a custom cost to our project
weave_client.add_cost(
    llm_id="your_model_name", prompt_token_cost=0.1, completion_token_cost=0.2
)

# We can then query for the calls, and with include_costs=True
# we receive the costs back attached to the calls
calls = weave_client.get_calls(filter={"trace_roots_only": True}, include_costs=True)

list(calls)
You now have a custom model that records its own token usage, a custom cost registered against that model’s llm_id, and a way to retrieve calls with per-call cost data attached under summary.weave.costs.