在LangChain的RAG RunnableSequence中获取Gemini大模型调用的Token使用元数据

我太懂你现在的纠结了——用LangChain搭的RAG链（RunnableSequence）调用Gemini时，要么拿不到token使用数据，要么能拿到数据的原生SDK调用方式没法和RAG链兼容。下面给你几个实打实的解决办法，都是亲测能跑通的：

方法一：用LangChain回调函数捕获Token使用数据 #

这个方法不用大改你现有的RAG链结构，通过自定义回调函数，就能在每次LLM调用结束后自动提取token使用信息。

先整个自定义的回调类，继承LangChain的BaseCallbackHandler，重写on_llm_end方法——这个方法会在LLM调用结束后触发，刚好能拿到我们要的token使用元数据：

from langchain_core.callbacks import BaseCallbackHandler
from langchain_google_vertexai import VertexAI
from langchain_core.runnables import RunnableSequence, RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

class TokenUsageCallback(BaseCallbackHandler):
    def on_llm_end(self, response, **kwargs):
        # 从LLM响应的元数据里扒出token使用情况
        if hasattr(response, 'llm_output') and 'usage_metadata' in response.llm_output:
            usage = response.llm_output['usage_metadata']
            print(f"Token使用明细：输入token数 {usage['prompt_token_count']}，输出token数 {usage['candidates_token_count']}，总token数 {usage['total_token_count']}")

# 初始化带回调的LLM实例
llm = VertexAI(
    model_name="gemini-1.5-pro-001",
    callbacks=[TokenUsageCallback()]
)

# 这里用简化版RAG链示例，你直接替换成自己的链结构就行
retriever = ... # 你的检索器（比如VectorStoreRetriever）
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

rag_chain = RunnableSequence(
    RunnablePassthrough.assign(context=retriever | format_docs),
    {"question": lambda x: x["question"], "context": lambda x: x["context"]},
    prompt, # 你的提示词模板（比如ChatPromptTemplate或者PromptTemplate）
    llm,
    StrOutputParser()
)

# 调用链的时候，回调会自动打印token使用数据
result = rag_chain.invoke({"question": "What is the meaning of life"})

只要启动这个回调，每次RAG链调用Gemini时，都会自动输出token的使用情况，完全不影响原来的RAG逻辑。

方法二：把原生Gemini调用封装成Runnable集成到链中 #

如果你想完全掌控Gemini的调用过程，拿到最完整的响应数据（包括safety_rating、finish_reason这些），可以把原生Vertex SDK的调用封装成LangChain的Runnable，直接塞到RunnableSequence里。

代码示例如下：

from vertexai.generative_models import GenerativeModel
from langchain_core.runnables import RunnableLambda, RunnableSequence, RunnablePassthrough

# 把原生Gemini调用封装成一个可被RunnableSequence调用的函数
def gemini_call_with_full_metadata(inputs):
    model = GenerativeModel("gemini-1.5-pro-001")
    # 这里根据你的提示词模板格式化输入，记得把上下文和问题拼对
    prompt_text = f"基于以下上下文回答问题：\n{inputs['context']}\n\n问题：{inputs['question']}"
    response = model.generate_content(prompt_text)
    
    # 提取你需要的所有数据
    token_usage = {
        "prompt_tokens": response.usage_metadata.prompt_token_count,
        "completion_tokens": response.usage_metadata.candidates_token_count,
        "total_tokens": response.usage_metadata.total_token_count
    }
    safety_ratings = response.safety_ratings
    finish_reason = response.finish_reason
    answer = response.text
    
    # 返回包含所有数据的字典
    return {
        "answer": answer,
        "token_usage": token_usage,
        "safety_ratings": safety_ratings,
        "finish_reason": finish_reason
    }

# 构建你的RAG链
retriever = ... # 你的检索器
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

rag_chain = RunnableSequence(
    RunnablePassthrough.assign(context=retriever | format_docs),
    RunnableLambda(gemini_call_with_full_metadata)
)

# 调用链，拿到完整的结果和元数据
result = rag_chain.invoke({"question": "What is the meaning of life"})
print("回答内容：", result["answer"])
print("Token使用情况：", result["token_usage"])
print("安全评级：", result["safety_ratings"])

这种方式的好处是，你能拿到Gemini返回的所有原生数据，而且完全兼容RunnableSequence的调用逻辑，只需要把原来的LLM调用节点替换成这个自定义的RunnableLambda就行。

方法三：开启LangChain VertexAI类的元数据返回参数 #

其实LangChain的VertexAI类本身就支持返回LLM的调用元数据，只是默认没开启而已。你只需要在初始化LLM的时候，把return_llm_output参数设为True，就能拿到包含token使用数据的完整响应了。

代码示例：

from langchain_google_vertexai import VertexAI
from langchain_core.runnables import RunnableSequence, RunnablePassthrough

# 初始化LLM时开启元数据返回
llm = VertexAI(
    model_name="gemini-1.5-pro-001",
    return_llm_output=True
)

# 直接调用LLM的话，会返回包含text和llm_output的字典
test_response = llm.invoke("What is the meaning of life")
print("Token使用数据：", test_response["llm_output"]["usage_metadata"])

# 集成到RAG链里的时候，注意不要用StrOutputParser（它会只提取字符串，丢掉元数据）
retriever = ... # 你的检索器
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

rag_chain = RunnableSequence(
    RunnablePassthrough.assign(context=retriever | format_docs),
    {"question": lambda x: x["question"], "context": lambda x: x["context"]},
    prompt, # 你的提示词模板
    llm
)

# 调用链，拿到包含回答和元数据的结果
result = rag_chain.invoke({"question": "What is the meaning of life"})
print("回答内容：", result["text"])
print("Token使用情况：", result["llm_output"]["usage_metadata"])

这个方法是最简单的，几乎不用改现有代码，只要开启一个参数，就能拿到你要的token使用数据。

总结 #

• 如果你不想大改现有RAG链，优先用方法三，开启return_llm_output参数就行；
• 如果需要在调用过程中自动记录token使用，用方法一的回调函数；
• 如果要拿到Gemini返回的所有原生数据（比如safety_rating这些），用方法二的原生封装方式。

备注：内容来源于stack exchange，提问作者Matheus Torquato