直接调用Gemini API时LangSmith Token追踪异常，求类wrap_openai工具

核心思路：手动封装Gemini客户端并注入追踪逻辑 #

LangSmith的@traceable注解默认不会自动捕获Gemini API的Token消耗，需要手动对Gemini的客户端或请求方法进行封装，模拟wrap_openai的实现逻辑，把Token统计和追踪信息绑定到LangSmith的trace中。

具体实现步骤 #

1. 提取Gemini API的Token消耗数据
   Gemini API的响应里包含usage_metadata字段，其中prompt_token_count和candidates_token_count分别对应输入和输出的Token数量，可直接从返回结果中提取。

2. 封装Gemini调用函数并关联LangSmith Trace
   自行编写包装函数，在调用Gemini API后手动将Token数据添加到当前trace中，示例代码如下：

   from langsmith import traceable
   import google.generativeai as genai
   
   def wrap_gemini_generate(model):
       def wrapped_generate(**kwargs):
           # 执行原始Gemini生成请求
           response = model.generate_content(**kwargs)
           # 提取Token消耗数据
           usage = response.usage_metadata
           prompt_tokens = usage.prompt_token_count
           completion_tokens = usage.candidates_token_count
           total_tokens = prompt_tokens + completion_tokens
           
           # 获取当前活跃的trace，写入Token信息
           from langsmith.run_tracker import get_current_run
           current_run = get_current_run()
           if current_run:
               current_run.update(
                   extra={
                       "token_usage": {
                           "prompt_tokens": prompt_tokens,
                           "completion_tokens": completion_tokens,
                           "total_tokens": total_tokens
                       }
                   }
               )
           return response
       return wrapped_generate
   
   # 初始化Gemini模型
   genai.configure(api_key="YOUR_API_KEY")
   model = genai.GenerativeModel("gemini-pro")
   # 包装生成方法
   model.generate_content = wrap_gemini_generate(model)
   
   # 用@traceable标记业务函数
   @traceable
   def my_gemini_call(prompt):
       return model.generate_content(prompt)
   

3. 验证追踪效果
   调用my_gemini_call后，进入LangSmith控制台查看对应trace，即可看到自定义添加的token_usage字段，其中包含准确的Token统计数据。

关键说明 #

• 若使用Gemini的chat等其他调用方法，可按相同逻辑包装对应方法，提取响应中的usage_metadata即可。
• 需确保LangSmith环境变量（LANGCHAIN_API_KEY、LANGCHAIN_TRACING_V2）配置正确，否则get_current_run会返回None，无法写入Token信息。

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