Tracking token usage when using RunnableWithMessageHistory and Tool

[lorerave85]

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

import os
from dotenv import load_dotenv
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.agents import AgentExecutor
from langchain_core.runnables.history import RunnableWithMessageHistory
from toolBuilder import ToolBuilder
from langchain_community.chat_message_histories.redis import RedisChatMessageHistory

load_dotenv()

tool_builder = ToolBuilder()
tools = tool_builder.initialize_tools()

redis_host = os.getenv('REDIS_HOST', 'localhost')
redis_port = os.getenv('REDIS_PORT', '6379')

system_prompt="""
You are an assistant, reply in the same language of user.
Respond only using available tools.
"""
human_prompt = "{input}"
prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    MessagesPlaceholder('chat_history'),
    ("human", human_prompt),
    MessagesPlaceholder('agent_scratchpad'),
])

llm = ChatOpenAI(
    model_name=os.getenv('OPENAI_CHAT_MODEL'),
    temperature=os.getenv('OPENAI_TEMPERATURE'),
    max_tokens=os.getenv('OPENAI_MAX_TOKENS')
)  

agent = create_openai_tools_agent(
    llm=llm,
    tools=tools,
    prompt=prompt
)

agent_executor = AgentExecutor(
    agent=agent, 
    tools=tools,
    max_iterations=os.getenv('LANGCHAIN_MAX_ITERATIONS'),
    verbose=True,
    handle_parsing_errors=True)

REDIS_URL = f"redis://{redis_host}:{redis_port}/0"

config={"configurable": {"session_id": "<SESSION_ID>"}}

agent_with_chat_history = RunnableWithMessageHistory(
    agent_executor,
    lambda session_id: RedisChatMessageHistory(session_id, url=REDIS_URL, max_messages=10),
    input_messages_key="input",
    history_messages_key="chat_history",
)

agent_with_chat_history.invoke(input={"input":[{"content":"cambio password","type":"human"}]},config={"configurable": {"session_id": "foobar"}},)

Description

I'm implementing an agent that interacts with the user through tools and history saved on redis.
this chain will then be exposed through langserve, but it is important to know the tokens consumed by each call (I use invoke).

I have already tried to use the various callback methods but none of them return the tokens.
how can I do it? on langsmith I saw that the total token is shown, so this information is present somewhere.

thanks

System Info

langchain==0.3.4
langchain-chroma==0.1.4
langchain-community==0.3.2
langchain-core==0.3.13
langchain-openai==0.2.0
langchain-text-splitters==0.3.0

[TheoTickel]

did you ever get an answer to this?

[lorerave85]

no, but I solved it.

I modified a class of the library to be able to create a dedicated data; in particular langserve/schema.py
Unfortunately, Monkey Patching did not work on that class, so I am going to replace the file in the dockerbuild.

then I used an http middleware to parse and add the tokens of each step:
@app.middleware("http")
async def parse_response_middleware(request: Request, call_next)

it is not a clean solution, but I have not found any alternatives

Lorenzo