How to Include Token Usage and Cost Details in LangServe Response JSON?

[quartermaine]

I'm working on a LangServe implementation and would like to enhance the response JSON to include metadata related to token usage and cost details for a chain. Specifically, I want to track the following metrics:

Total Tokens
Prompt Tokens
Completion Tokens
Total Cost (USD)

For context, here’s an example of my current chain setup:

chain = (  
    sql_chain  
    | RunnableLambda(lambda x: add_attribute(x, 'command', get_query_memory()))  
).with_types(input_type=InputType, output_type=OutputType)

I’d like to integrate these metrics into the response_metadata of the JSON returned by the chain.
In my previous standalone scripts using OpenAI’s API, I’ve used the get_openai_callback() context manager to extract these details, as shown below:

with get_openai_callback() as cb:  
    llm_res = llm_chain.run(res)  
    print(f"Response: {llm_res}")  
    print(f"Total Tokens: {cb.total_tokens}")  
    print(f"Prompt Tokens: {cb.prompt_tokens}")  
    print(f"Completion Tokens: {cb.completion_tokens}")  
    print(f"Total Cost (USD): ${cb.total_cost}")

However, I’m not sure how to incorporate similar logic into my LangServe chain and ensure that these details are added to the response JSON.