ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

使用 @mlflow/gemini 为 Google Gemini 接入 MLflow Tracing:自动追踪 generateContent 调用的完整指南

使用 @mlflow/gemini 为 Google Gemini 接入 MLflow Tracing:自动追踪 generateContent 调用的完整指南 使用 mlflow/gemini 为 Google Gemini 接入 MLflow Tracing自动追踪 generateContent 调用的完整指南【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow本文介绍 MLflow TypeScript SDK 家族中的 Gemini 集成包mlflow/gemini。该包基于mlflow/core的 Tracing 能力通过tracedGemini()一行代码即可对 Google Gemini SDKgoogle/genai的generateContent调用进行自动插桩将输入、输出、Token 用量、耗时与异常信息记录为 MLflow Trace并在 MLflow UI 中可视化。读完本文你将掌握从环境准备、包安装、SDK 初始化到运行与结果验证的完整接入方案并理解其 Proxy 插桩机制与 Span 数据结构的底层实现。包概览Gemini 集成包在整个 TypeScript SDK 中的定位mlflow/gemini是 MLflow TypeScript SDK位于 libs/typescript下的集成包之一与anthropic、openai等集成包并列专门为 Google Gemini 提供自动插桩能力。其元信息定义在 package.json项值说明包名mlflow/gemini发布在 NPM 上的 Gemini 集成包版本0.4.0与mlflow/core0.4.0 对齐描述Gemini integration package for MLflow Tracing为 MLflow Tracing 提供 Gemini 集成许可证Apache-2.0与整个 MLflow 项目一致详见 LICENSE.txtNode 版本要求20核心包mlflow/core要求 Node18集成包要求更高该包的运行时依赖关系为mlflow/core^0.4.0peer dependency提供核心 Tracing 功能与手动插桩 API如withSpan、init、flushTraces等google/genai^1.22.0peer dependencyGoogle 官方 Gemini JavaScript SDKtracedGemini包装的就是它的客户端实例。由于两者都是 peer dependencies安装时需要根据所用包管理器决定是否显式安装这两个包详见下一节。安装一条命令安装 Gemini 集成包在项目中安装mlflow/gemininpm install mlflow/gemini安装后请确认mlflow/core与google/genai两个 peer dependencies 已就绪。使用npm7 时peer dependencies 通常会随主包自动安装若使用较旧的npm或其他包管理器可能需要显式安装npm install mlflow/core google/genai从源码构建与质量检查方面package.json 提供了以下脚本npm run buildtsc编译到dist/、npm testJest 测试、npm run lintESLint零警告阈值。发布包仅包含编译产物dist/目录。快速开始5 步完成 Gemini 调用追踪第 1 步启动 MLflow Tracking Server如果你还没有可用的 MLflow 服务先通过 Python 环境启动一个本地 Tracking Serverpip install mlflow mlflow server --backend-store-uri sqlite:///mlruns.db --port 5000自托管 MLflow Server 要求 Python 3.10 或更高版本。如果本地没有 Python 环境也可以使用托管的 MLflow 服务快速体验或者参考仓库中的 Docker 部署方案如 docker/Dockerfile 与 docker-compose/docker-compose.yml。第 2 步初始化 MLflow SDK在你的应用入口处初始化mlflow/core指向刚才启动的 Tracking Server并指定实验Experimentimport * as mlflow from mlflow/core; mlflow.init({ trackingUri: http://localhost:5000, experimentId: experiment-id, });trackingUriMLflow Tracking Server 的地址本地启动即为http://localhost:5000experimentIdTrace 归属的 Experiment ID需要预先在 MLflow UI 中创建或通过MlflowClient.createExperiment获取测试代码 index.test.ts 中即为这种用法。第 3 步用 tracedGemini 包装 Gemini 客户端创建 Google Gemini 客户端后调用tracedGemini(gemini)得到追踪版客户端随后正常调用models.generateContent即可调用方式与原生 SDK 完全一致import { tracedGemini } from mlflow/gemini; import { GoogleGenAI } from google/genai; const gemini new GoogleGenAI({ apiKey: process.env.GEMINI_API_KEY }); const client tracedGemini(gemini); const response await client.models.generateContent({ model: gemini-2.0-flash-001, contents: Hello Gemini, });GEMINI_API_KEY需在 Google AI Studio 申请并写入环境变量。model参数传入 Gemini 模型名示例使用gemini-2.0-flash-001contents传入提示词内容。第 4 步在 MLflow UI 中查看 Trace调用完成后打开 MLflow UI默认地址http://localhost:5000在对应 Experiment 下即可看到本次调用生成的 Trace其中包含Span 名称generateContentSpan 类型LLM输入Inputs传入的model与contents参数输出OutputsGemini 的完整响应对象属性AttributesToken 用量与消息格式标记状态、开始与结束时间。说明原 README 中的追踪 UI 截图托管于项目外部本文以文字描述替代实际界面效果以你本地部署的 MLflow UI 为准。第 5 步确保 Trace 落盘测试与调试场景在测试或需要立即读取 Trace 的场景中可调用mlflow.flushTraces()强制将缓存的 Trace 刷出再通过mlflow.getLastActiveTraceId()拿到最近一次 Trace ID随后用MlflowClient.getTrace(traceId)拉取完整 Trace 数据——这正是 index.test.ts 中验证行为的标准姿势await mlflow.flushTraces(); const traceId mlflow.getLastActiveTraceId(); const trace await client.getTrace(traceId!);工作原理Proxy 自动插桩的实现细节tracedGemini并非对 Gemini SDK 的逐方法重写而是基于 JavaScriptProxy的轻量自动插桩。核心实现位于 src/index.ts可拆解为三层1. 递归代理tracedGeminitracedGemini返回一个Proxy包装的客户端。在get陷阱中属性是函数判断是否需要追踪见下文第 2 层需要则返回wrapWithTracing包装后的函数否则bind(target)后原样返回保证this上下文正确属性是普通对象递归调用tracedGemini继续代理从而支持client.models.generateContent这类多级链式调用其他值原样透传。正是这种递归代理设计使得wrappedGemini.models.generateContent(...)与原生写法逐字相同无需改动业务代码。2. 追踪白名单shouldTraceMethod是否追踪由两个常量决定const SUPPORTED_MODULES [models]; const SUPPORTED_METHODS [generateContent];shouldTraceMethod将模块名客户端构造函数的constructor.name转小写后与SUPPORTED_MODULES比对方法名与SUPPORTED_METHODS比对二者同时命中才执行插桩。这意味着当前版本只追踪models.generateContent文本/多模态内容生成models.embedContent、models.countTokens等其他方法不在白名单内不会被追踪模块名匹配也是大小写不敏感的。测试用例 index.test.ts 中should not trace methods outside SUPPORTED_METHODS专门验证了白名单外的调用不会产生追踪行为。3. 包装与 Span 生成wrapWithTracing命中白名单后方法被withSpan包裹return withSpan( async (span: LiveSpan) { span.setInputs(args[0]); const result await fn.apply(this, args); span.setOutputs(result); // ...token usage 提取与属性设置 return result; }, { name: methodName, spanType }, );withSpan是mlflow/core提供的手动插桩 API见 core/src/core/api.ts其职责包括创建 Span、记录耗时、自动捕获异常并将 Span 状态标记为ERROR。spanType由getSpanType方法映射generateContent对应SpanType.LLM其余方法返回undefined不追踪。Span 数据模型Inputs、Outputs、Token 用量与属性每一次被追踪的generateContent调用都会生成一个LLM类型的 Span其字段可在 UI 或getTrace返回的 Trace 数据中核对以 index.test.ts 中的断言为据字段预期值说明namegenerateContentSpan 名称即被追踪的方法名spanTypeLLM见 core 的 SpanType 枚举LLM 类型用于 UI 特殊渲染logLevelINFO由SpanType.LLM的默认日志级别决定status.statusCodeOK/ERROR成功为 OK抛异常时为 ERRORinputs{ model, contents, ... }即调用generateContent时传入的第一个参数对象outputsGemini 响应对象完整透传不做裁剪startTime/endTimeISO 时间戳由withSpan自动记录attributes[mlflow.chat.tokenUsage]{ input_tokens, output_tokens, total_tokens }Token 用量见下文attributes[mlflow.message.format]gemini标记消息格式供下游 UI 解析使用Token 用量提取Token 用量来自 Gemini 响应的usageMetadata字段提取逻辑见 src/index.ts 中的extractTokenUsageconst usage response?.usageMetadata ?? response?.usage; const input usage.promptTokenCount; const output usage.candidatesTokenCount; const total usage.totalTokenCount;三者齐备时才组装为{ input_tokens, output_tokens, total_tokens }并通过span.setAttribute(SpanAttributeKey.TOKEN_USAGE, usage)写入属性提取失败仅打印 debug 日志不影响调用本身。SpanAttributeKey.TOKEN_USAGE的键名为mlflow.chat.tokenUsage定义于 core/src/core/constants.ts格式约定为{input_tokens: int, output_tokens: int, total_tokens: int}。测试中通过 Mock 服务器返回固定 Token 数输入 10、输出 5、总计 15并逐项断言了三个 Token 键的存在与数值见 index.test.ts。消息格式标记每个 Span 都会设置SpanAttributeKey.MESSAGE_FORMAT gemini键名mlflow.message.format。该属性用于下游如 MLflow UI判断按哪种消息格式解析 Span 内容是实现多厂商 LLM Trace 统一渲染的关键约定。嵌套追踪父 Span 下的子 SpantracedGemini天然支持嵌套场景。当generateContent调用发生在mlflow.withSpan包裹的父 Span 内时Gemini 调用会自动成为父 Span 的子 Span形成一棵 Span 树。测试 index.test.ts 验证了如下结构Trace └── predict (CHAIN) ← 手动创建inputs: Hello from parent span └── generateContent (LLM) ← 自动插桩inputs: { model, contents }父 Span 记录应用的业务语义如predict子 Span 记录 Gemini 调用细节二者构成完整的调用链。这意味着你可以用mlflow.withSpan为整个 Agent/Chain 流程建外层 Span让 Gemini 调用自动挂载其下无需手动管理 Span 的父子关系。错误处理异常调用如何被记录当generateContent抛错如 API 返回 429 限流、鉴权失败等时withSpan会捕获异常并标记 Span 状态。测试用例should handle generateContent errors properlyindex.test.ts模拟了 HTTP 429RESOURCE_EXHAUSTED响应并验证Trace 整体状态trace.info.state为ERRORLLM Span 的status.statusCode为SpanStatusCode.ERRORSpan 的inputs依然保留可回溯是哪个请求失败Span 的outputs为undefined调用未成功返回startTime/endTime仍被记录可计算失败耗时。因此即便 Gemini API 调用失败你依然能在 MLflow UI 中看到完整的失败记录便于排查限流与参数问题。测试验证Mock 服务器驱动的集成测试mlflow/gemini的测试采用MSWMock Service Worker在 Node 侧拦截真实的 Gemini API 请求无需真实 API Key 即可运行tests/mockGeminiServer.ts 注册了https://generativelanguage.googleapis.com/v1beta/models/*:generateContent与v1/models/*:generateContent两个端点的 Mock返回包含candidates、usageMetadata的仿真响应tests/index.test.ts 在每个用例中通过mlflow.init初始化 SDKtrackingUri 指向http://localhost:5000、调用tracedGemini包装客户端、执行generateContent最后用flushTracesgetTrace断言 Span 内容jest.config.js 使用ts-jest预设并通过全局 setup/teardown 脚本管理测试环境超时 30 秒。如果你要在自己的项目里复现这套测试可以参照该结构用 MSW 拦截 Gemini 端点再断言getTrace返回的 Span 字段。已知范围与注意事项追踪范围有限当前仅覆盖models.generateContentmodels模块外的属性和generateContent之外的方法会被透传但不会产生 SpanNode 版本mlflow/gemini要求 Node20核心包为18请确认运行环境Python 版本自托管 MLflow Server 需 Python 3.10peer dependencies安装后请核对mlflow/core与google/genai版本是否满足要求核心包与集成包版本需对齐到 0.4.xGemini SDK 需^1.22.0追踪开销Proxy 插桩发生在客户端对象访问层对业务代码零侵入仅在白名单命中的方法上产生额外 Span 写入开销。至此你已经掌握了mlflow/gemini从安装、初始化、自动追踪到结果核验与故障排查的完整链路可以将任意基于google/genai的 Gemini 应用快速接入 MLflow 的可观测体系。【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表