Weave custom costs

custom_costs.py

# uv pip install weave langchain-google-genai -qq

import os
import weave
from langchain_google_genai import ChatGoogleGenerativeAI

os.environ["GOOGLE_API_KEY"] = 

GEMINI_MODEL_NAME = 'gemini-2.5-flash'
CUSTOM_MODEL_NAME = 'lc_' + GEMINI_MODEL_NAME 

weave_client = weave.init("morgan/testing")

weave_client.add_cost(
    llm_id=CUSTOM_MODEL_NAME, 
    prompt_token_cost=0.3, 
    completion_token_cost=2.5
)

llm = ChatGoogleGenerativeAI(
    model=GEMINI_MODEL_NAME,
)

@weave.op
def a_great_model(nessages):
  response = llm.invoke(messages)
  input_tokens = response.usage_metadata["input_tokens"]
  output_tokens = response.usage_metadata["output_tokens"]
  
  return {
      "response": response,
      "usage": {
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "total_tokens": input_tokens + output_tokens,
            },
      "model": CUSTOM_MODEL_NAME,
    }


messages = [
    ("system", "You love long poems",),
    ("human", "tell me an incredibly long poem about bees."),
]

a_great_model(messages)