Repository navigation
feat: 支持上报模型返回的额外字段 - #328
Conversation
AI Code Review审查结论通过 审查范围:提交 c98a61d..e5ea854「feat: 支持上报模型返回的额外字段」,共 9 个文件(212 增 / 8 删),核心为 OpenAIModel 新增 response_metadata_extractor 提取器(四条生成路径接入)、_trace.py 上报 provider_response_metadata span 属性、示例与 2 个测试文件。计划符合性:功能主体(非流式/流式/Responses 流四路接入、JSON 可序列化校验、异常吞并、trace 上报、示例与测试)均已落地;_runner_manager.py 的 await 修复与 Runner.close 的 async 签名一致,正确。主要风险:1) 流式路径仅取最后一帧提取元数据,真实后端元数据帧位置不保证为末帧(尤其 Responses 流会被 terminal 事件覆盖),存在静默丢失;2) 非流式提取仅见顶层 dict,嵌套层厂商字段拿不到;3) 提取器返回关键信息缺失时仅 warn,无透出。测试充分性:新增 3 个测试覆盖非流式顶层提取、流式末帧提取、trace 属性上报,但未覆盖异常/非 dict 返回、元数据位于首/中帧、Responses 流路径、嵌套字段位置。门禁结论:无 SEVERE 级阻断性问题(无数据损坏、崩溃、安全漏洞;核心丢失风险与位置契约问题以 MODERATE 呈现),state 为 PASSED。 发现的问题中等
问题: 流式路径(Chat Completions 触发条件: 响应头/首个 chunk 携带元数据、或 SSE 流有多个收尾事件(真实后端常见),末帧不含目标字段。 实际影响: 依赖该功能的可观测性场景(如 Venus 追踪链路)中元数据静默丢失,最终 修正方向: 累积候选事件(如保留首个携带且提取器返回非空的 中等
问题: 触发条件: 后端把目标字段放在 实际影响: 嵌套位置携带的厂商字段无法上报,链路观测数据缺失;且提取器失败仅产生一条 warning,用户难以定位。 修正方向: 向提取器传入完整响应 dict 的同时文档明确约定字段位置,或提取时对 较低
问题: 新测试 触发条件: 回归测试执行时,凡依赖上述未覆盖分支的用户场景(厂商字段早到、提取器异常、Responses API 流)都没有防回归保障。 实际影响: 异常/非 dict 返回被静默吞掉、末帧覆盖中间帧等缺陷无法被 CI 发现,且将来改动易引入回归。 修正方向: 补齐 较低
问题: 触发条件: 后续扩展在 实际影响: 当前无实害, 修正方向: 改为 |
| provider_response_metadata = self._extract_provider_response_metadata(last_event_dict) | ||
| yield self._attach_provider_response_metadata(final_response, provider_response_metadata) |
There was a problem hiding this comment.
问题: 流式路径(Chat Completions _generate_stream 与 Responses API _generate_responses_stream)仅在 async for 结束后用最后一块 chunk 调用一次提取器。OpenAI 兼容后端在 SSE 流终止前往往还会发送额外的元数据事件(如无 choices 的 usage 收尾块、response.completed 等),这些事件在真实流中是最后一条;而测试 test_streaming_extracts_provider_metadata_from_usage_chunk 中最后一块恰好是携带 venusMarker 的 usage-only chunk——该字段在该块中才出现。即实测中只有当后端调度恰好把元数据放在最后一帧时才成功;若元数据早到(如首个非空 chunk 或中间 chunk 携带),或最后还有一个不含该字段的收尾 chunk,提取器见不到它,元数据被丢弃。Responses 流中 last_event_dict 也常被 terminal 的 response.completed 事件(response_data 中通常不含顶层元数据)覆盖。
触发条件: 响应头/首个 chunk 携带元数据、或 SSE 流有多个收尾事件(真实后端常见),末帧不含目标字段。
实际影响: 依赖该功能的可观测性场景(如 Venus 追踪链路)中元数据静默丢失,最终 LlmResponse.custom_metadata 及 trace span 中缺少本应上报的 provider_response_metadata。
修正方向: 累积候选事件(如保留首个携带且提取器返回非空的 last_event_dict,即第一次提取成功后不再覆盖),或在每个 chunk 到达时调用提取器并合并非空结果,而不只在循环结束后用最后一帧。
| provider_response_metadata = self._extract_provider_response_metadata(response_dict) | ||
| return self._attach_provider_response_metadata(llm_response, provider_response_metadata) |
There was a problem hiding this comment.
问题: _generate_single(非流式)中提取器收到的是 response.model_dump() 的顶层 dict。多数 OpenAI 兼容后端把厂商字段放在响应体顶层(如 providerRequestId、some_field),测试也如此构造;但部分后端(如 OpenAI Responses 结构或 gateway 实现)把元数据放在 choices[].message 或 response 嵌套对象里,提取器只拿到顶层 dict,看不到嵌套字段。已有测试 test_non_streaming_extracts_provider_metadata 只覆盖顶层场景。
触发条件: 后端把目标字段放在 choices[0].message、response 或其他嵌套层而非顶层时。
实际影响: 嵌套位置携带的厂商字段无法上报,链路观测数据缺失;且提取器失败仅产生一条 warning,用户难以定位。
修正方向: 向提取器传入完整响应 dict 的同时文档明确约定字段位置,或提取时对 choices[0].message 等常见嵌套位置做兜底查找;并补充非顶层字段的测试。
| async def test_streaming_extracts_provider_metadata_from_usage_chunk(self): | ||
| """Provider metadata survives a final usage-only chunk.""" | ||
|
|
||
| def extract_metadata(response_data): | ||
| marker = response_data.get("venusMarker") | ||
| if not marker: | ||
| return None | ||
| return {"venus_marker": {"span_id": marker["spanId"]}} | ||
|
|
||
| model = _model(response_metadata_extractor=extract_metadata) | ||
| request = _request([Content(parts=[Part.from_text(text="hi")], role="user")]) | ||
|
|
||
| content_chunk = Mock() | ||
| content_chunk.model_dump.return_value = { | ||
| "id": "resp_1", | ||
| "choices": [{ | ||
| "delta": { | ||
| "content": "hello" | ||
| }, | ||
| "finish_reason": "stop", | ||
| }], | ||
| "usage": None, | ||
| } | ||
| usage_chunk = Mock() | ||
| usage_chunk.model_dump.return_value = { | ||
| "id": "resp_1", | ||
| "choices": [], | ||
| "usage": { | ||
| "prompt_tokens": 1, | ||
| "completion_tokens": 1, | ||
| "total_tokens": 2, | ||
| }, | ||
| "venusMarker": { | ||
| "spanId": "9d3e43a402a76a5b" | ||
| }, | ||
| } | ||
|
|
||
| async def mock_stream(): | ||
| yield content_chunk | ||
| yield usage_chunk | ||
|
|
||
| with patch.object(model, "_create_async_client") as mock_factory: | ||
| mock_client = AsyncMock() | ||
| mock_client.chat.completions.create = AsyncMock(return_value=mock_stream()) | ||
| mock_client.close = AsyncMock() | ||
| mock_factory.return_value = mock_client | ||
|
|
||
| responses = [] | ||
| async for response in model.generate_async(request, stream=True): | ||
| responses.append(response) | ||
|
|
||
| final_response = next(response for response in responses if not response.partial) | ||
| assert final_response.custom_metadata == { | ||
| "stream_complete": True, | ||
| "provider_response_metadata": { | ||
| "venus_marker": { | ||
| "span_id": "9d3e43a402a76a5b" | ||
| } | ||
| }, | ||
| } | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_streaming_null_response_raises(self): |
There was a problem hiding this comment.
问题: 新测试 test_streaming_extracts_provider_metadata_from_usage_chunk 覆盖了 final usage-chunk 场景,但未覆盖提取器抛异常、返回非 dict、返回 None 及元数据在首个/中间 chunk(非末帧)的路径,未覆盖 _generate_responses_stream(Responses API)与 _generate_single 异常分支。这些分支的行为各不相同。
触发条件: 回归测试执行时,凡依赖上述未覆盖分支的用户场景(厂商字段早到、提取器异常、Responses API 流)都没有防回归保障。
实际影响: 异常/非 dict 返回被静默吞掉、末帧覆盖中间帧等缺陷无法被 CI 发现,且将来改动易引入回归。
修正方向: 补齐 except Exception 返回空 dict、非 dict 类型、首/中帧携带元数据、Responses 流提取等用例的断言。
| """Attach extracted metadata under a stable, provider-neutral namespace.""" | ||
| if not metadata: | ||
| return response | ||
| custom_metadata = dict(response.custom_metadata or {}) | ||
| custom_metadata[const.PROVIDER_RESPONSE_METADATA] = metadata | ||
| response.custom_metadata = custom_metadata | ||
| return response |
There was a problem hiding this comment.
问题: _attach_provider_response_metadata 直接原地 response.custom_metadata[const.PROVIDER_RESPONSE_METADATA] = metadata,没有走 model_copy。终态 LlmResponse 是非流式 _generate_single 中 yield 前的最后一步,流式路径也只作用于新建的 final response,属安全;但 _responses_error 与非流式 _generate_single 中 custom_metadata 存在其他分支赋值(如 _retry.py 的 error response),若后续在此基类上复用该方法,原地修改可能污染同一实例。
触发条件: 后续扩展在 _attach_provider_response_metadata 被调用时引用同一 LlmResponse 的 custom_metadata 的其他字段。
实际影响: 当前无实害,custom_metadata 的赋值本就要求整体兼容;仅提示未来改动与 partial 响应路径共用一个实例时可能产生隐式共享。
修正方向: 改为 response.model_copy(update={'custom_metadata': {**(response.custom_metadata or {}), const.PROVIDER_RESPONSE_METADATA: metadata}}) 或先 dict() 再赋值,保持不变性。
e5ea854 to
5fc5b15
Compare
AI Code Review审查结论通过 审查范围:c98a61d..5fc5b15(9 个文件,+471/-8),计划为「支持上报模型返回的额外字段」。主要改动:OpenAIModel 新增 response_metadata_extractor 回调,在四条生成路径(chat completions 非流式/流式、responses API 非流式/流式)的终态响应上通过 _attach_provider_response_metadata 把提取结果挂到 LlmResponse.custom_metadata["provider_response_metadata"],并由 trace_call_llm 写入 span 属性 trpc.python.agent.provider_response_metadata;示例与测试同步更新。计划符合性:各路径接线一致、提取器异常被隔离、流式场景按"首个事件提取一次"的文档语义实现,trace 使用 _safe_json_serialize 且仅对 dict 生效;事件经 _llm_processor._create_event_from_response 原样透传 custom_metadata,示例读取路径可达。测试覆盖:非流式/流式/Responses 流式/提取器异常与非法返回值/流式去重均有用例,trace 断言了属性写入。验证后保留 3 个低严重度问题:(1) 客户端取消(GeneratorExit)时差值对还在进行模型调用,提取器仍在事件循环汇总运行,其内部的 json.dumps 验证与 logger 调用为阻塞同步操作,long-poll 流式场景首个 chunk 处理在事件循环热路径上新增同步 JSON 序列化(默认提取器为 None 时开销为一次属性读取,可忽略);(2) 错误路径(openai 抛异常、重试层重建响应、Responses 流式 error 事件)下 provider 元数据不上报,可能导致调用端观察不到失败请求的追踪字段,属设计取舍;(3) trace 测试里 llm_response 为 MagicMock 而非真实 LlmResponse,若 custom_metadata 含目标字段也能通过。未发现高置信的正确性、安全或数据完整性缺陷,门禁结论为通过。 发现的问题较低
问题: 流式路径把提取器调用和内部的 触发条件: 配置了 实际影响: stream TTFB 和事件循环吞吐出现毫秒级抖动;对海量 token 场景,首个 chunk 的序列化开销被放大到每次会话请求上。 修正方向: 把 |
|
|
||
| chunk_dict: dict = chunk.model_dump() | ||
| logger.debug("🔥 RAW LLM CHUNK: %s", json.dumps(chunk_dict, ensure_ascii=False)) | ||
| if not metadata_extraction_attempted: | ||
| provider_response_metadata = self._extract_provider_response_metadata(chunk_dict) | ||
| metadata_extraction_attempted = True |
There was a problem hiding this comment.
问题: 流式路径把提取器调用和内部的 json.dumps(metadata) 校验放在了外部 async for 循环捕获每个 chunk 后、继续处理该 chunk 之前的位置,等价于把阻塞的同步序列化插进了事件循环每次 chunk 处理的热路径(_generate_responses_stream 的 2318-2322 行同理)。
触发条件: 配置了 response_metadata_extractor 且流式响应第一个 chunk 较大(含 usage 完整统计或大段文本)时触发;长连接流式场景下该 chunk 的后续 API 数据、其他并发连接的事件循环任务都因此被同步 JSON 序列化和 logger 调用阻塞。
实际影响: stream TTFB 和事件循环吞吐出现毫秒级抖动;对海量 token 场景,首个 chunk 的序列化开销被放大到每次会话请求上。
修正方向: 把 _extract_provider_response_metadata 的同步 json.dumps 校验和日志与数据处理解耦:仅对提取器返回值做浅度 isinstance 校验,将 JSON 可序列化性校验延迟到 _attach_provider_response_metadata 或留给 trace 层 _safe_json_serialize 兜底;或把提取器调用移出 async for 之前的首次平面检查(先取事件、再后台/低优先级执行)。
| assistant_started = True | ||
|
|
||
| async for event in runner.run_async(user_id=user_id, session_id=session_id, new_message=user_content): | ||
| provider_metadata = (event.custom_metadata or {}).get("provider_response_metadata") |
5fc5b15 to
6844420
Compare
6844420 to
db552c3
Compare
AI Code Review审查结论通过 审查范围:c98a61d..6844420「feat: 支持上报模型返回的额外字段」,18 个文件(+947/-8)。核心实现:OpenAIModel 新增 response_metadata_extractor 回调,在 4 条调用路径(非流式 chat、非流式 Responses、流式 chat、流式 Responses)提取白名单字段,附加到 LlmResponse.custom_metadata 的 provider_response_metadata 命名空间并上报 trace span;另含文档、示例(llmagent_with_model_extra_fields)、测试(test_openai_model_ext.py +342 行、test_trace.py +20 行)及两处示例修复(_runner_manager.py 补 await、model_retry 示例补 import json,均为正确修复)。计划已实现:提取失败兜底(异常/非 dict/不可序列化均忽略)与文档一致,主路径测试覆盖充分(提取成功、嵌套字段、流式首 chunk、trace 属性)。主要风险:(1) MODERATE——流式 Chat Completions 提取器仅对首 chunk 调用一次且不验证是否提取成功,真实服务首 chunk 通常仅为 role 空 chunk,导致后续携带厂商字段的 chunk 被跳过、最终无元数据,测试 mock 均默认首 chunk 带字段未覆盖该场景;(2) LOW——重试层错误响应不携带元数据,失败场景无法获得链路标识;(3) LOW——trace 属性无大小限制,超出 OTel 属性上限静默截断;(4) LOW——非流式 Responses 路径无测试、流式「跳过提取器」测试未真正覆盖时序分支。无 SEVERE 级问题,state=PASSED,建议修复流式首 chunk 提取逻辑后合入。 发现的问题中等
问题: 流式 Chat Completions 路径中,提取器只在首个 chunk 上调用一次,但不验证首次调用是否真正提取到了元数据。OpenAI 兼容服务的流式响应首个 chunk 通常是仅含 触发条件: 使用 实际影响: 流式调用下上报功能静默失效——事件 修正方向: 将 较低
问题: 新增的 provider 元数据 span 属性对提取结果大小无任何限制,直接 触发条件: 用户在提取器中返回较大字段(如长 trace-id 列表、嵌套对象),或误按文档警告返回接近完整响应的大字典。 实际影响: 属性被 span SDK 或 collector 静默截断或丢弃,trace 中元数据不完整而调用方无法感知;同时该属性与 修正方向: 对序列化结果设置明确长度上限(如截断至 4KB 并附截断标记),或仅当元数据大小低于阈值时才设置属性;与 较低
问题: 新增的元数据附加只发生在四条成功路径(非流式/流式 × chat/Responses)的 attach 调用上;SDK 重试层 触发条件: 配置了 实际影响: 失败场景下用户无法通过统一渠道获取最后一次请求的厂商链路标识(如 span-id),与“每次模型调用都上报元数据”的预期不符,影响链路排障。 修正方向: 在重试层构造最终错误响应时,将最后一次实际请求的提取结果合并进 较低
问题: 新增测试对 4 条调用路径的覆盖不均衡: 触发条件: 流式响应首个 chunk 不带厂商字段(真实服务常见)时该分支无测试保护;Responses 非流式路径回归无护栏。 实际影响: 上述两类缺陷(流式元数据丢失、Responses 非流式路径回归)均无法被现有测试发现,测试给出的通过信号具有误导性。 修正方向: 补充 |
| async for chunk in response: | ||
| if chunk is None: | ||
| continue | ||
|
|
||
| chunk_dict: dict = chunk.model_dump() | ||
| logger.debug("🔥 RAW LLM CHUNK: %s", json.dumps(chunk_dict, ensure_ascii=False)) | ||
| if not metadata_extraction_attempted: | ||
| provider_response_metadata = self._extract_provider_response_metadata(chunk_dict) | ||
| metadata_extraction_attempted = True |
There was a problem hiding this comment.
问题: 流式 Chat Completions 路径中,提取器只在首个 chunk 上调用一次,但不验证首次调用是否真正提取到了元数据。OpenAI 兼容服务的流式响应首个 chunk 通常是仅含 role 的占位 chunk(choices[0].delta 无内容),此时 _extract_provider_response_metadata 返回 {},而 metadata_extraction_attempted 已置为 True,后续携带厂商字段(如示例中的 venusMarker)的 chunk 被永久跳过,最终响应的 custom_metadata 中缺失 provider_response_metadata。
触发条件: 使用 response_metadata_extractor + stream=True 调用 generate_async,且厂商字段不在首 chunk 中(真实服务的常见行为;测试 mock 均人为让首 chunk 携带字段,未暴露该场景)。
实际影响: 流式调用下上报功能静默失效——事件 custom_metadata 与 trace span 中无任何厂商元数据,用户从非流式切到流式后发现字段丢失,且日志无提示,难以定位。
修正方向: 将 metadata_extraction_attempted 的置位推迟到首次提取到非空结果之后(循环内仅当 provider_response_metadata 仍为空时才再次尝试),并补充「首 chunk 无字段、后续 chunk 有字段」的流式测试用例。
| custom_metadata = llm_response.custom_metadata | ||
| if isinstance(custom_metadata, dict): | ||
| provider_metadata = custom_metadata.get(PROVIDER_RESPONSE_METADATA) | ||
| if isinstance(provider_metadata, dict) and provider_metadata: | ||
| span.set_attribute( | ||
| f"{_trpc_agent_span_name}.{PROVIDER_RESPONSE_METADATA}", | ||
| _safe_json_serialize(provider_metadata), | ||
| ) |
There was a problem hiding this comment.
问题: 新增的 provider 元数据 span 属性对提取结果大小无任何限制,直接 _safe_json_serialize 后 set_attribute。OTel span 属性通常有默认大小上限(常见 SDK 约 4096 字节),超限属性会被静默丢弃或截断,且调用方(response_metadata_extractor)由用户配置,可返回任意大小数据。
触发条件: 用户在提取器中返回较大字段(如长 trace-id 列表、嵌套对象),或误按文档警告返回接近完整响应的大字典。
实际影响: 属性被 span SDK 或 collector 静默截断或丢弃,trace 中元数据不完整而调用方无法感知;同时该属性与 llm_response 属性内容重复,徒增 span 体积和采集成本(每个模型调用都会增大)。
修正方向: 对序列化结果设置明确长度上限(如截断至 4KB 并附截断标记),或仅当元数据大小低于阈值时才设置属性;与 llm_response 属性的截断策略保持一致。
| return self._attach_provider_response_metadata( | ||
| llm_response, | ||
| self._extract_provider_response_metadata(response_dict), | ||
| ) |
There was a problem hiding this comment.
问题: 新增的元数据附加只发生在四条成功路径(非流式/流式 × chat/Responses)的 attach 调用上;SDK 重试层 retry_model_call 重试耗尽后返回的最终错误 LlmResponse(custom_metadata={"error": str(ex)},见未变更的 _retry.py:_build_error_response)不经过此 attach,且该响应经 _llm_processor.py:271 直接透传 custom_metadata,最终事件中同样不含 provider_response_metadata。
触发条件: 配置了 response_metadata_extractor,发生可重试异常且重试全部失败(retry_model_call 返回错误响应),或模型返回非 2xx 状态码后重试耗尽。
实际影响: 失败场景下用户无法通过统一渠道获取最后一次请求的厂商链路标识(如 span-id),与“每次模型调用都上报元数据”的预期不符,影响链路排障。
修正方向: 在重试层构造最终错误响应时,将最后一次实际请求的提取结果合并进 custom_metadata(如在 _generate_async_impl 层保留上一份提取结果并附加),使错误响应与成功响应一致携带 provider 元数据。
| thought_partials = [r for r in partial_responses if r.content and r.content.parts[0].thought] | ||
| assert len(thought_partials) >= 1 | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_streaming_extracts_provider_metadata_from_first_chunk(self): | ||
| """Provider metadata is extracted once from the first stream chunk.""" | ||
|
|
||
| def extract_metadata(response_data): | ||
| marker = response_data.get("venusMarker") | ||
| if not marker: | ||
| return None | ||
| return {"venus_marker": {"span_id": marker["spanId"]}} | ||
|
|
||
| model = _model(response_metadata_extractor=extract_metadata) | ||
| request = _request([Content(parts=[Part.from_text(text="hi")], role="user")]) | ||
|
|
||
| content_chunk = Mock() | ||
| content_chunk.model_dump.return_value = { | ||
| "id": "resp_1", | ||
| "choices": [{ | ||
| "delta": { | ||
| "content": "hello" | ||
| }, | ||
| "finish_reason": "stop", | ||
| }], | ||
| "usage": None, | ||
| "venusMarker": { | ||
| "spanId": "9d3e43a402a76a5b" | ||
| }, | ||
| } | ||
| usage_chunk = Mock() | ||
| usage_chunk.model_dump.return_value = { | ||
| "id": "resp_1", | ||
| "choices": [], | ||
| "usage": { | ||
| "prompt_tokens": 1, | ||
| "completion_tokens": 1, | ||
| "total_tokens": 2, | ||
| }, | ||
| "venusMarker": { | ||
| "spanId": "9d3e43a402a76a5b" | ||
| }, | ||
| } | ||
|
|
||
| async def mock_stream(): | ||
| yield content_chunk | ||
| yield usage_chunk | ||
|
|
||
| with patch.object(model, "_create_async_client") as mock_factory: | ||
| mock_client = AsyncMock() | ||
| mock_client.chat.completions.create = AsyncMock(return_value=mock_stream()) | ||
| mock_client.close = AsyncMock() | ||
| mock_factory.return_value = mock_client | ||
|
|
||
| responses = [] | ||
| async for response in model.generate_async(request, stream=True): | ||
| responses.append(response) | ||
|
|
||
| final_response = next(response for response in responses if not response.partial) | ||
| assert final_response.custom_metadata == { | ||
| "stream_complete": True, | ||
| "provider_response_metadata": { | ||
| "venus_marker": { | ||
| "span_id": "9d3e43a402a76a5b" | ||
| } | ||
| }, | ||
| } | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_streaming_keeps_metadata_from_earlier_chunk(self): | ||
| """A later usage-only chunk without metadata must not drop earlier fields.""" | ||
|
|
||
| def extract_metadata(response_data): | ||
| marker = response_data.get("venusMarker") | ||
| if not marker: | ||
| return None | ||
| return {"venus_marker": {"span_id": marker["spanId"]}} | ||
|
|
||
| model = _model(response_metadata_extractor=extract_metadata) | ||
| request = _request([Content(parts=[Part.from_text(text="hi")], role="user")]) | ||
|
|
||
| content_chunk = Mock() | ||
| content_chunk.model_dump.return_value = { | ||
| "id": "resp_1", | ||
| "choices": [{ | ||
| "delta": { | ||
| "content": "hello" | ||
| }, | ||
| "finish_reason": "stop", | ||
| }], | ||
| "usage": None, | ||
| "venusMarker": { | ||
| "spanId": "early-span" | ||
| }, | ||
| } | ||
| usage_chunk = Mock() | ||
| usage_chunk.model_dump.return_value = { | ||
| "id": "resp_1", | ||
| "choices": [], | ||
| "usage": { | ||
| "prompt_tokens": 1, | ||
| "completion_tokens": 1, | ||
| "total_tokens": 2, | ||
| }, | ||
| } | ||
|
|
||
| async def mock_stream(): | ||
| yield content_chunk | ||
| yield usage_chunk | ||
|
|
||
| with patch.object(model, "_create_async_client") as mock_factory: | ||
| mock_client = AsyncMock() | ||
| mock_client.chat.completions.create = AsyncMock(return_value=mock_stream()) | ||
| mock_client.close = AsyncMock() | ||
| mock_factory.return_value = mock_client | ||
|
|
||
| responses = [] | ||
| async for response in model.generate_async(request, stream=True): | ||
| responses.append(response) | ||
|
|
||
| final_response = next(response for response in responses if not response.partial) | ||
| assert final_response.custom_metadata == { | ||
| "stream_complete": True, | ||
| "provider_response_metadata": { | ||
| "venus_marker": { | ||
| "span_id": "early-span" | ||
| } | ||
| }, | ||
| } | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_responses_streaming_keeps_metadata_from_earlier_event(self): | ||
| """Responses stream merges metadata even if response.completed lacks it.""" | ||
|
|
||
| def extract_metadata(response_data): | ||
| marker = response_data.get("venusMarker") | ||
| if not marker: | ||
| return None | ||
| return {"venus_marker": {"span_id": marker["spanId"]}} | ||
|
|
||
| model = _model(use_responses_api=True, response_metadata_extractor=extract_metadata) | ||
| request = _request([Content(parts=[Part.from_text(text="hi")], role="user")]) | ||
|
|
||
| async def stream_events(): | ||
| yield { | ||
| "type": "response.created", | ||
| "response": { | ||
| "id": "resp_meta", | ||
| }, | ||
| "venusMarker": { | ||
| "spanId": "responses-span" | ||
| }, | ||
| } | ||
| yield {"type": "response.output_text.delta", "delta": "hello"} | ||
| yield { | ||
| "type": "response.completed", | ||
| "response": { | ||
| "id": "resp_meta", | ||
| "status": "completed", | ||
| "output": [{ | ||
| "type": "message", | ||
| "role": "assistant", | ||
| "content": [{ | ||
| "type": "output_text", | ||
| "text": "hello" | ||
| }], | ||
| }], | ||
| }, | ||
| } | ||
|
|
||
| with patch.object(model, "_create_async_client") as mock_factory: | ||
| mock_client = AsyncMock() | ||
| mock_client.responses.create = AsyncMock(return_value=stream_events()) | ||
| mock_client.close = AsyncMock() | ||
| mock_factory.return_value = mock_client | ||
|
|
||
| responses = [] | ||
| async for response in model.generate_async(request, stream=True): | ||
| responses.append(response) | ||
|
|
||
| final_response = next(response for response in responses if not response.partial) | ||
| assert final_response.custom_metadata["provider_response_metadata"] == { | ||
| "venus_marker": { | ||
| "span_id": "responses-span" | ||
| } | ||
| } | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_streaming_skips_extractor_after_first_nonempty(self): | ||
| """Later chunks must not rerun the extractor once metadata is found.""" | ||
| calls = {"count": 0} | ||
|
|
||
| def extract_metadata(response_data): | ||
| calls["count"] += 1 | ||
| marker = response_data.get("venusMarker") | ||
| if not marker: | ||
| return None | ||
| return {"venus_marker": {"span_id": marker["spanId"]}} | ||
|
|
||
| model = _model(response_metadata_extractor=extract_metadata) | ||
| request = _request([Content(parts=[Part.from_text(text="hi")], role="user")]) | ||
|
|
||
| first_chunk = Mock() | ||
| first_chunk.model_dump.return_value = { | ||
| "id": "resp_1", | ||
| "choices": [{ | ||
| "delta": { | ||
| "content": "hello" | ||
| }, | ||
| "finish_reason": "stop", | ||
| }], | ||
| "usage": None, | ||
| "venusMarker": { | ||
| "spanId": "first-span" | ||
| }, | ||
| } | ||
| later_chunk = Mock() | ||
| later_chunk.model_dump.return_value = { | ||
| "id": "resp_1", | ||
| "choices": [], | ||
| "usage": { | ||
| "prompt_tokens": 1, | ||
| "completion_tokens": 1, | ||
| "total_tokens": 2, | ||
| }, | ||
| "venusMarker": { | ||
| "spanId": "later-span" | ||
| }, | ||
| } | ||
|
|
||
| async def mock_stream(): | ||
| yield first_chunk | ||
| yield later_chunk | ||
|
|
||
| with patch.object(model, "_create_async_client") as mock_factory: | ||
| mock_client = AsyncMock() | ||
| mock_client.chat.completions.create = AsyncMock(return_value=mock_stream()) | ||
| mock_client.close = AsyncMock() | ||
| mock_factory.return_value = mock_client | ||
|
|
||
| responses = [] | ||
| async for response in model.generate_async(request, stream=True): | ||
| responses.append(response) | ||
|
|
||
| final_response = next(response for response in responses if not response.partial) | ||
| assert final_response.custom_metadata["provider_response_metadata"] == { | ||
| "venus_marker": { | ||
| "span_id": "first-span" | ||
| } | ||
| } | ||
| assert calls["count"] == 1 |
There was a problem hiding this comment.
问题: 新增测试对 4 条调用路径的覆盖不均衡:use_responses_api=True 的非流式路径(_generate_responses_single,本文件 1806-1809 行的新 attach 逻辑)完全没有测试;且 test_streaming_skips_extractor_after_first_nonempty 的 mock 中首 chunk 已携带 venusMarker,即“首次尝试即成功”,并未真正覆盖断言名所述的「首个 chunk 无字段、后续 chunk 才调用提取器」的时序分支,而该分支正是流式元数据丢失风险所在。
触发条件: 流式响应首个 chunk 不带厂商字段(真实服务常见)时该分支无测试保护;Responses 非流式路径回归无护栏。
实际影响: 上述两类缺陷(流式元数据丢失、Responses 非流式路径回归)均无法被现有测试发现,测试给出的通过信号具有误导性。
修正方向: 补充 use_responses_api=True 非流式元数据附加测试;扩展流式测试覆盖「首 chunk 提取返回空 → 后续 chunk 再尝试提取」的时序,并断言提取器调用次数符合预期。
AI Code Review审查结论通过 审查范围与结论范围: 计划符合性:功能按计划落地。白名单回调、JSON 可序列化校验、异常吞并、流式仅首次事件提取均有实现与测试; 主要风险:流式路径"仅第一个事件提取一次"的闩锁在首个 chunk/事件不含元数据时永久失效——真实网关常把标记放在后续事件(如 Responses 的 测试充分性:7 个新测试覆盖非流式/两种流式的提取、返回值非法与异常容错、"元数据在首个 chunk"的情形;但未覆盖"首个 chunk 无标记、后续 chunk 有标记"这一流式核心失败模式, 门禁结论:PASSED。无阻断合入的高置信缺陷;1 个 MODERATE、3 个 LOW,均有界且可操作,建议合入前修复流式首次事件提取的闩锁语义并补充对应测试。 发现的问题中等
问题: 两条流式路径(chat 流 触发条件: 配置了 实际影响: 功能在流式场景对这类网关完全静默失效: 修正方向: 将闩锁语义改为"提取到非空结果后才停止",首个事件提取为空时继续尝试后续事件,或在流结束时用最终事件( 较低
问题: 触发条件: 用户自定义工具、after_model_callback 或下游代码在 实际影响: 用户数据被静默覆盖并随事件转发到 AG-UI/A2A 与会话持久化,丢失后无任何告警。 修正方向: 在文档中明确保留该命名空间并做冲突检测:合并时若已存在同名键且值不同,记录 warning(或改用嵌套子命名空间避免与用户顶层键冲突)。 较低
问题: 触发条件: 流式模式下最终文本只出现在终结事件而无前置 partial(如 实际影响: 最终回答文本被静默丢弃不打印,示例演示输出缺失答案。 修正方向: 移除初始化时的 True(改用 None/False 表达"尚未开始"),或在 non-partial 文本分支中先打印再置位,确保最终文本事件必然被输出。 较低
问题: 触发条件: 任何人在修闩锁语义时以该测试为回归依据——它会在"提取空结果后仍在后续 chunk 重试"的实现下同样通过(count 仍为 1),无法防止回归。 实际影响: 核心契约(元数据出现在非首个事件时是否还能上报)无测试保护,后续改动可能保持测试全绿但功能仍静默丢失元数据。 修正方向: 新增用例:首个 chunk 不含标记、第二个 chunk 含标记,断言最终响应成功带上元数据(按预期契约应提取到);并将本测试重命名为反映实际语义"首次事件后不再重复调用"。 |
| async for chunk in response: | ||
| if chunk is None: | ||
| continue | ||
|
|
||
| chunk_dict: dict = chunk.model_dump() | ||
| logger.debug("🔥 RAW LLM CHUNK: %s", json.dumps(chunk_dict, ensure_ascii=False)) | ||
| if not metadata_extraction_attempted: | ||
| provider_response_metadata = self._extract_provider_response_metadata(chunk_dict) | ||
| metadata_extraction_attempted = True |
There was a problem hiding this comment.
问题: 两条流式路径(chat 流 _generate_stream 2496-2498 行、Responses 流 _generate_responses_stream 2320-2322 行)的 metadata_extraction_attempted 闩锁在第一个 chunk/事件后无条件置位,与提取结果无关;且 chat 流的提取先于 _is_thinking_event/keep-alive 事件处理执行。只要第一个事件不含厂商字段(如 stream_server.event thinking 事件、role-only delta、仅 usage 的空 chunk,或 Responses 的 response.created 信封),后续事件即使携带标记也永远不会传给 extractor,最终响应静默缺失 provider_response_metadata。
触发条件: 配置了 response_metadata_extractor 并使用默认流式模式(RunConfig.streaming 默认 True);兼容网关只在后续事件(如 response.completed 或首个文本 chunk)注入 venusMarker 等字段——这是真实网关最常见的注入位置。
实际影响: 功能在流式场景对这类网关完全静默失效:final_response.custom_metadata 不含提取结果,调用方(如示例 run_agent.py)读不到任何厂商元数据,无任何日志或报错提示,且与同一模型非流式调用的行为不一致。
修正方向: 将闩锁语义改为"提取到非空结果后才停止",首个事件提取为空时继续尝试后续事件,或在流结束时用最终事件(response.completed/usage chunk)补一次提取;同时把提取移到 thinking 事件分支之后,并补充"首个 chunk 无标记、后续 chunk 有标记"的测试。
| def _attach_provider_response_metadata( | ||
| response: LlmResponse, | ||
| metadata: dict[str, Any], | ||
| ) -> LlmResponse: | ||
| """Attach extracted metadata under a stable, provider-neutral namespace.""" | ||
| if not metadata: | ||
| return response | ||
| custom_metadata = dict(response.custom_metadata or {}) | ||
| custom_metadata[const.PROVIDER_RESPONSE_METADATA] = metadata | ||
| response.custom_metadata = custom_metadata | ||
| return response |
There was a problem hiding this comment.
问题: _attach_provider_response_metadata 在合并时无条件执行 custom_metadata[const.PROVIDER_RESPONSE_METADATA] = metadata,对用户或 SDK 其他层预先写入的同名键静默覆盖,没有任何保留/冲突检查。custom_metadata 是面向用户的自由字典(abc/_response.py 中注释为 JSON 可序列化即可),SDK 内部多处也在写入(如 streaming_tool_call、tool_progress、advanced_memory_request_context_fingerprint、error)。
触发条件: 用户自定义工具、after_model_callback 或下游代码在 custom_metadata 中使用了 provider_response_metadata 这一键(该键随本次变更被导出为公共常量并成为事实上的保留命名)。
实际影响: 用户数据被静默覆盖并随事件转发到 AG-UI/A2A 与会话持久化,丢失后无任何告警。
修正方向: 在文档中明确保留该命名空间并做冲突检测:合并时若已存在同名键且值不同,记录 warning(或改用嵌套子命名空间避免与用户顶层键冲突)。
| assistant_started = True | ||
|
|
||
| async for event in runner.run_async(user_id=user_id, session_id=session_id, new_message=user_content): | ||
| provider_metadata = (event.custom_metadata or {}).get(PROVIDER_RESPONSE_METADATA) | ||
| if provider_metadata: | ||
| print("\nProvider metadata: ", f"{json.dumps(provider_metadata, ensure_ascii=False)}") | ||
| if event.is_error(): | ||
| if assistant_started: | ||
| print() | ||
| assistant_started = False | ||
| print(f"Error: {event.error_code}: {event.error_message}") | ||
| continue | ||
|
|
||
| if not event.content or not event.content.parts: | ||
| continue | ||
|
|
||
| if event.partial: | ||
| for part in event.content.parts: | ||
| if part.text and not part.thought: | ||
| if not assistant_started: | ||
| print("Assistant: ", end="", flush=True) | ||
| assistant_started = True | ||
| print(part.text, end="", flush=True) | ||
| continue | ||
|
|
||
| for part in event.content.parts: | ||
| if part.thought: | ||
| continue | ||
| if part.function_call: | ||
| print(f"\nInvoke Tool: {part.function_call.name}({part.function_call.args})") | ||
| assistant_started = False | ||
| elif part.function_response: | ||
| print(f"Tool Result: {part.function_response.response}") | ||
| elif part.text and not assistant_started: |
There was a problem hiding this comment.
问题: assistant_started = True 在进入事件循环前置位(第 45 行),而 function_response 分支(第 74-75 行)不会将其复位为 False,只有 function_call 分支复位——因此第 78 行 elif part.text and not assistant_started: 在"模型直接给出最终文本且此前未出现 function_call"的所有场景下为 False。
触发条件: 流式模式下最终文本只出现在终结事件而无前置 partial(如 add_tools_to_prompt 过滤缓冲,或 run_config.streaming=False 的非流式运行),且事件序列中没有 function_call 复位标志。
实际影响: 最终回答文本被静默丢弃不打印,示例演示输出缺失答案。
修正方向: 移除初始化时的 True(改用 None/False 表达"尚未开始"),或在 non-partial 文本分支中先打印再置位,确保最终文本事件必然被输出。
No description provided.