Is the token usage aggregated across a chain of subchains?

[daniel-deychakiwsky]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate, HumanMessagePromptTemplate
from langchain_core.output_parsers import StrOutputParser

model = ChatOpenAI(
    model="gpt-4o",
    temperature=0.0,
    max_tokens=500,
    top_p=1.0,
    api_key="..."
)

prompt_one = ChatPromptTemplate.from_messages([HumanMessage(content="Hey! How are you?")])
prompt_two = ChatPromptTemplate.from_messages([HumanMessagePromptTemplate(
    prompt=PromptTemplate(
        template="You said: {ai_response}. Can you explain?",
        input_variables=["ai_response"],
        )
    )
])

chain_one = prompt_one | model | StrOutputParser()
chain_two = prompt_two | model

chain_all = chain_one | chain_two

ai_message = chain_all.invoke({})
ai_message.usage_metadata

Description

Is the token usage aggregated across a chain of subchains? In the provided code snippet, I have two sub chains (chain_one and chain_two) which are chained together as sub chains to form a final chain (chain_all). I can intercept the token usage from each subchain by using listener functions on each of them however I'm curious to clarify if langchain already outputs an aggregate token usage from both subchains on the composed chain output (the ai_message which is returned from chain_all.invoke())?

System Info

System Information

OS: Linux
OS Version: #1 SMP Fri Oct 18 12:02:02 UTC 2024
Python Version: 3.10.12 (main, Sep 11 2024, 15:47:36) [GCC 11.4.0]

Package Information

langchain_core: 0.3.15
langchain: 0.3.7
langsmith: 0.1.139
langchain_google_vertexai: 2.0.7
langchain_openai: 0.2.5
langchain_text_splitters: 0.3.2

Optional packages not installed

langgraph
langserve

Other Dependencies

aiohttp: 3.10.10
anthropic[vertexai]: Installed. No version info available.
async-timeout: 4.0.3
google-cloud-aiplatform: 1.70.0
google-cloud-storage: 2.18.2
httpx: 0.27.2
httpx-sse: 0.4.0
jsonpatch: 1.33
langchain-mistralai: Installed. No version info available.
numpy: 1.26.4
openai: 1.54.0
orjson: 3.10.11
packaging: 24.1
pydantic: 2.9.2
PyYAML: 6.0.2
requests: 2.32.3
requests-toolbelt: 1.0.0
SQLAlchemy: 2.0.36
tenacity: 9.0.0
tiktoken: 0.8.0
typing-extensions: 4.12.2