如何用 TokenCountingHandler 度量 LlamaIndex RAG 管线的 LLM 与 embedding token 消耗【免费下载链接】llama_indexLlamaIndex is the leading document agent and OCR platform项目地址: https://gitcode.com/GitHub_Trending/ll/llama_index在 LlamaIndex 里搭建 RAG 管线时索引构建阶段会调用 embedding 模型查询阶段会调用 LLM两者都会产生 token 消耗但没有统计手段时你无法评估管线的成本。LlamaIndex 将 token 计数从旧的llm_predictor/embed_model对象属性迁移到了回调机制TokenCountingHandler回调可以同时跟踪 embedding、LLM prompt、LLM completion 三类 token并按你的节奏重置。本文给出一个完整可执行的路径把TokenCountingHandler挂到全局Settings先建VectorStoreIndex读取 embedding 计数再执行查询读取 LLM 计数并介绍token_budget上限和单个事件的查看方式。准备条件以官方示例 TokenCountingHandler.ipynb 为准需要准备llama-index与 OpenAI 集成包pip install llama-index pip install llama-index-llms-openaitiktoken用于构造 tokenizer示例中直接导入。一个可用的OPENAI_API_KEY环境变量。一份本地文本数据。示例使用 Paul Graham 的随笔存放于data/paul_graham/paul_graham_essay.txt通过SimpleDirectoryReader读取。如果你只是想离线估算消耗、不调用真实模型可改用MockLLM/MockEmbedding见后文的可选分支。把 TokenCountingHandler 挂到全局 Settingstoken 计数现在是一个回调通过CallbackManager注册并赋给全局Settings.callback_manager这样索引和查询都不需要再显式传 handler。tokenizer 传一个输入文本、返回 token 列表的函数示例用 tiktoken 的gpt-3.5-turbo编码import tiktoken from llama_index.core.callbacks import CallbackManager, TokenCountingHandler from llama_index.core import Settings from llama_index.llms.openai import OpenAI token_counter TokenCountingHandler( tokenizertiktoken.encoding_for_model(gpt-3.5-turbo).encode, verboseFalse, # 设为 True 时会把每次 usage 打印到控制台 ) Settings.llm OpenAI(modelgpt-3.5-turbo, temperature0.2) Settings.callback_manager CallbackManager([token_counter])要点tokenizer可以显式指定也可以省略让其落到全局默认 tokenizer见 llama_index.core.utils 中get_tokenizer()的默认逻辑。verboseTrue时每次 embedding 和 LLM 事件结束都会打印Embedding Token Usage/LLM Prompt Token Usage方便肉眼核对不打印时直接读计数器属性。计数器上的所有计数都是累计值只会在你调用token_counter.reset_counts()时才清零这一点直接决定了下面建索引 → 读数 → 重置 → 查询 → 再读数的操作顺序。测量索引构建阶段的 embedding token加载数据并构建向量索引from llama_index.core import SimpleDirectoryReader, VectorStoreIndex documents SimpleDirectoryReader(./data/paul_graham).load_data() index VectorStoreIndex.from_documents(documents)from_documents会对每个节点调用 embedding 模型TokenCountingHandler会监听EMBEDDING事件并按 chunk 累加计数。构建完成后直接读取print(token_counter.total_embedding_token_count)官方 notebook 中这一步的示例输出为20723文档示例实际数值取决于你的数据量与切分参数不代表固定预期。测量查询阶段的 LLM token查询前建议先重置计数把 embedding 与 LLM 的消耗分开度量token_counter.reset_counts() query_engine index.as_query_engine() response query_engine.query(What did the author do growing up?)然后读取四项指标print( Embedding Tokens: , token_counter.total_embedding_token_count, \n, LLM Prompt Tokens: , token_counter.prompt_llm_token_count, \n, LLM Completion Tokens: , token_counter.completion_llm_token_count, \n, Total LLM Token Count: , token_counter.total_llm_token_count, \n, )四个属性对应属性含义total_embedding_token_count所有 embedding chunk 的 token 累计prompt_llm_token_count所有 LLM 调用的 prompt token 累计completion_llm_token_count所有 LLM 调用的 completion token 累计total_llm_token_countprompt completion 的总和官方 notebook 的查询阶段示例输出为文档示例Embedding Tokens: 6 LLM Prompt Tokens: 4563 LLM Completion Tokens: 123 Total LLM Token Count: 4686这里的Embedding Tokens: 6只包含查询文本这一次 embedding 调用建索引时的大量 embedding 已被reset_counts()清零4563的 prompt 之所以大是因为示例查询的similarity_top_k取回 4 个节点、按默认 chunk size 需要拆成 2 次 LLM 调用来读全部检索文本所以llm_token_counts里有 2 个事件。用 token_budget 设置 LLM 消耗上限TokenCountingHandler支持可选的token_budget参数它是这个 handler 生命周期内 LLM token 总用量的上限当前只约束 LLM token超过时抛出ValueError。实现见 token_counting.pydef _check_budget(self) - None: if ( self.token_budget is not None and self.total_llm_token_count self.token_budget ): raise ValueError( fToken budget exceeded! Limit: {self.token_budget}, fCurrent: {self.total_llm_token_count} )检查在每次事件开始on_event_start和每次 LLM 事件结束on_event_end后触发。用法token_counter TokenCountingHandler( tokenizertiktoken.encoding_for_model(gpt-3.5-turbo).encode, token_budget5000, # LLM token 累计超过该值即抛 ValueError )注意超限时抛异常会中断当前任务适合在脚本里做硬上限保护如果你只想测量而不是拦截就不要设置token_budget。查看单次调用明细累计计数之外handler 把每次调用记录成TokenCountingEvent分别存在llm_token_counts和embedding_token_counts两个列表里。事件字段为prompt发给 LLM 或 embedding 的文本、prompt_token_count、completionLLM 回复embedding 不用、completion_token_count、total_token_countprompt completion、event_id与其他回调 handler 对齐的事件 ID。print(Num LLM token count events: , len(token_counter.llm_token_counts)) print( Num Embedding token count events: , len(token_counter.embedding_token_counts), )notebook 示例输出文档示例Num LLM token count events: 2、Num Embedding token count events: 1。再看单个事件print(prompt: , token_counter.llm_token_counts[0].prompt[:100], ...) print( prompt token count: , token_counter.llm_token_counts[0].prompt_token_count, ) print( completion token count: , token_counter.llm_token_counts[0].completion_token_count, ) print(total token count, token_counter.llm_token_counts[0].total_token_count)按事件明细排查哪一次调用最贵比如是检索到的节点太多导致 prompt 过大还是 completion 被写得过长。可选分支用 MockLLM / MockEmbedding 离线估算如果不想真实调用模型Usage Pattern 文档 给出了用 mock 对象做预测的用法from llama_index.core.llms import MockLLM from llama_index.core import MockEmbedding, Settings llm MockLLM(max_tokens256) embed_model MockEmbedding(embed_dim1536) import tiktoken from llama_index.core.callbacks import CallbackManager, TokenCountingHandler token_counter TokenCountingHandler( tokenizertiktoken.encoding_for_model(gpt-3.5-turbo).encode ) Settings.llm llm Settings.embed_model embed_model Settings.callback_manager CallbackManager([token_counter])按同样顺序建索引、读计数、reset_counts()、查询、再读计数即可。MockLLM的max_tokens参数用作最坏情况预测每次 LLM 回复都按该 token 数计不指定max_tokens时则直接预测回 prompt 本身。这条路径适合在真正接入模型之前预估成本数字是预测值而非真实账单。限制与注意事项计数是累计的不重置的话embedding 与 LLM 计数会持续累加。想分阶段度量时在各阶段之间显式调用reset_counts()。token_budget 只管 LLMtoken_budget超出时抛ValueError且目前只适用于 LLM token 计数不覆盖 embedding。LLM 计数的取值来源handler 会优先从 LLM 原始响应的 usage 字段读取 provider 报告的 token 数支持prompt_tokens/input_tokens/prompt_token_count等输入键与completion_tokens/output_tokens/candidates_token_count等输出键取不到时才回退到本地 tokenizer 估算见 get_tokens_from_response。旧实现已废弃之前在llm_predictor/embed_model对象上直接读last_token_usage/total_token_usage、并用静态 gpt-2 tokenizer 打印到控制台的实现已被废弃last_token_usage和total_token_usage属性并非总能被正确维护。如果你仍在使用旧写法请迁移到本文的TokenCountingHandler方式细节见 Token Counting - Migration Guide。事件过滤构造参数event_starts_to_ignore/event_ends_to_ignore可排除不想计入的CBEventType需要精确控制统计口径时使用。完整可运行的 notebook 见 TokenCountingHandler.ipynb成本结构背景不同 index 在构建期/查询期的 LLM 调用次数差异见 Cost Analysis。【免费下载链接】llama_indexLlamaIndex is the leading document agent and OCR platform项目地址: https://gitcode.com/GitHub_Trending/ll/llama_index创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考