weave_custom_costs_hack.py

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("cost-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 hack_usage(usage):
  return usage

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



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

a_great_model(messages)