如何在Llama-Index查询引擎中用TikToken追踪Anthropic模型Token用量

针对你遇到的聊天补全能获取Token计数、但查询引擎无法获取的问题，可通过以下方案解决：

方案1：利用LlamaIndex回调系统统计Token #

通过自定义回调处理器，在LLM调用的生命周期节点中，用Anthropic的tokenizer手动统计输入输出的Token数量。

from llama_index.core.callbacks.base_handler import BaseCallbackHandler
from llama_index.core.callbacks.schema import CBEventType, EventPayload
from llama_index.llms.anthropic import Anthropic
from llama_index.core import Settings
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.response_synthesizers import CompactAndRefine
from llama_index.core import VectorStoreIndex

class TokenCountHandler(BaseCallbackHandler):
    def __init__(self, tokenizer):
        super().__init__(event_starts_to_ignore=[], event_ends_to_ignore=[])
        self.tokenizer = tokenizer
        self.total_input_tokens = 0
        self.total_output_tokens = 0

    def on_event_start(self, event_type: CBEventType, payload: dict | None = None, **kwargs) -> None:
        if event_type == CBEventType.LLM:
            # 统计输入prompt的Token数
            prompt = payload.get(EventPayload.PROMPT, "")
            self.total_input_tokens += len(self.tokenizer.encode(prompt))

    def on_event_end(self, event_type: CBEventType, payload: dict | None = None, **kwargs) -> None:
        if event_type == CBEventType.LLM:
            # 统计输出响应的Token数
            response = payload.get(EventPayload.RESPONSE, "")
            self.total_output_tokens += len(self.tokenizer.encode(response.text))

def generate_response(question, db_name, collection, usecase_id, llm, master_prompt):
    llm = Anthropic(model=llm, temperature=0.5)
    tokenizer = Anthropic().tokenizer
    Settings.tokenizer = tokenizer
    
    # 初始化Token统计处理器并注册到全局设置
    token_handler = TokenCountHandler(tokenizer)
    Settings.callbacks = [token_handler]
    
    embed_model = OpenAIEmbedding(model="text-embedding-3-small")
    vector_store = get_vectordb(db_name, collection)
    Settings.llm = llm
    Settings.embed_model = embed_model
    
    index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
    vector_retriever = index.as_retriever(
        vector_store_query_mode="default",
        similarity_top_k=5,
    )
    text_retriever = index.as_retriever(
        vector_store_query_mode="sparse",
        similarity_top_k=5,
    )
    retriever = QueryFusionRetriever(
        [vector_retriever, text_retriever],
        similarity_top_k=5,
        num_queries=1,
        mode="relative_score",
        use_async=False,
    )

    response_synthesizer = CompactAndRefine()

    # 移除冗余的index.as_query_engine()赋值，避免覆盖自定义配置
    query_engine = RetrieverQueryEngine(
        retriever=retriever,
        response_synthesizer=response_synthesizer,
    )
    
    response = query_engine.query(question)
    
    # 输出统计结果
    print(f"输入Token数: {token_handler.total_input_tokens}")
    print(f"输出Token数: {token_handler.total_output_tokens}")
    
    return response

方案2：从API原始响应中提取官方Token计数 #

Anthropic API会在响应中返回usage字段，包含官方统计的Token用量，可直接从查询引擎返回的响应对象中提取：

# 在generate_response函数的返回前添加以下代码
response = query_engine.query(question)
# 提取API返回的Token用量
if hasattr(response, 'raw') and 'usage' in response.raw:
    input_tokens = response.raw['usage']['input_tokens']
    output_tokens = response.raw['usage']['output_tokens']
    print(f"API统计输入Token: {input_tokens}, 输出Token: {output_tokens}")

关键注意事项 #

• 你原代码中存在冗余逻辑：先创建了自定义RetrieverQueryEngine，随后又用index.as_query_engine()覆盖配置，导致之前设置的检索器和合成器失效，务必移除这行冗余代码。
• 确保Settings.tokenizer正确设置为Anthropic的tokenizer，保障LlamaIndex内部Token统计逻辑正常运行。

内容的提问来源于stack exchange，提问作者Mohil