diff --git a/py/src/braintrust/integrations/agentscope/cassettes/1.0.0/test_agentscope_model_call_error_propagates.yaml b/py/src/braintrust/integrations/agentscope/cassettes/1.0.0/test_agentscope_model_call_error_propagates.yaml new file mode 100644 index 00000000..b28f4d33 --- /dev/null +++ b/py/src/braintrust/integrations/agentscope/cassettes/1.0.0/test_agentscope_model_call_error_propagates.yaml @@ -0,0 +1,86 @@ +interactions: +- request: + body: '{"messages":[{"role":"user","content":"hello"}],"model":"gpt-4o-mini","stream":false,"temperature":0}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '101' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.31.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.31.0 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"error\": {\n \"message\": \"Incorrect API key provided: sk-inval******************-key. + You can find your API key at https://platform.openai.com/account/api-keys.\",\n + \ \"type\": \"invalid_request_error\",\n \"code\": \"invalid_api_key\",\n + \ \"param\": null\n },\n \"status\": 401\n}" + headers: + access-control-expose-headers: + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a1c2d7d2bd7844b0-YYZ + connection: + - keep-alive + content-length: + - '276' + content-type: + - text/plain + date: + - Thu, 16 Jul 2026 17:49:21 GMT + server: + - cloudflare + set-cookie: + - __cf_bm=fw5AawKAtiBfybjDIj_R3LMPsv81WLDB5CQTZ3365nQ-1784224161.7172794-1.0.1.1-BWmnjx9aYPksUgoqoRSsMvFdzl21RDhJltuCefOlXneJbta1bJxVpJAmbEAxxINRKJjM1UDFJDJJTBzpctyjPZMtdaQRMxfOKPX5fh4JSHUB3b0a8IpgO3wzcH8KWztJ; + HttpOnly; SameSite=None; Secure; Path=/; Domain=api.openai.com; Expires=Thu, + 16 Jul 2026 18:19:21 GMT + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + x-content-type-options: + - nosniff + x-error-json: + - ewogICJlcnJvciI6IHsKICAgICJtZXNzYWdlIjogIkluY29ycmVjdCBBUEkga2V5IHByb3ZpZGVkOiBzay1pbnZhbCoqKioqKioqKioqKioqKioqKi1rZXkuIFlvdSBjYW4gZmluZCB5b3VyIEFQSSBrZXkgYXQgaHR0cHM6Ly9wbGF0Zm9ybS5vcGVuYWkuY29tL2FjY291bnQvYXBpLWtleXMuIiwKICAgICJ0eXBlIjogImludmFsaWRfcmVxdWVzdF9lcnJvciIsCiAgICAiY29kZSI6ICJpbnZhbGlkX2FwaV9rZXkiLAogICAgInBhcmFtIjogbnVsbAogIH0sCiAgInN0YXR1cyI6IDQwMQp9 + x-openai-authorization-error: + - '401' + x-openai-ide-error-code: + - invalid_api_key + x-openai-internal-caller: + - unknown_through_ide + x-request-id: + - req_ba6a808283f3468a9f18c1f0ce99cd0f + status: + code: 401 + message: Unauthorized +version: 1 diff --git a/py/src/braintrust/integrations/agentscope/cassettes/1.0.0/test_agentscope_streaming_model_call.yaml b/py/src/braintrust/integrations/agentscope/cassettes/1.0.0/test_agentscope_streaming_model_call.yaml new file mode 100644 index 00000000..02e61382 --- /dev/null +++ b/py/src/braintrust/integrations/agentscope/cassettes/1.0.0/test_agentscope_streaming_model_call.yaml @@ -0,0 +1,152 @@ +interactions: +- request: + body: '{"messages":[{"role":"system","name":"system","content":[{"type":"text","text":"You + are concise. Answer in one sentence."}]},{"role":"user","name":"user","content":[{"type":"text","text":"Say + hi in five words."}]}],"model":"gpt-4o-mini","stream":true,"stream_options":{"include_usage":true},"temperature":0,"tools":[{"type":"function","function":{"name":"generate_response","parameters":{"properties":{"response":{"description":"Your + response to the user.","type":"string"}},"required":["response"],"type":"object"},"description":"Generate + a response. Note only the input argument `response` is\n\nvisible to the others, + you should include all the necessary\ninformation in the `response` argument."}}]}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '703' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.31.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.31.0 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: 'data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,"id":"call_CvesWAmJBeVUZOajeyLoUEMO","type":"function","function":{"name":"generate_response","arguments":""}}],"refusal":null},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"KzG8vZhcDXLv92"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\""}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"qV1O0XhHLwa1RU"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"response"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"AGTyAg9vl"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\":\""}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"HGFsjc3CnygV"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"Hello"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"k8JZxxfjsA93"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":" + there"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"A8JADL0IS4m"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":","}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":""} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":" + how"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"RToqAwlXl5fvL"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":" + are"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"IKuJlDqMbr07x"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":" + you"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"krETyUvH3pIss"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"?\""}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"cPVkI5MgwwI7lP"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"}"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":""} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"tool_calls"}],"usage":null,"obfuscation":"sTb5lcMkAJBlnsw"} + + + data: {"id":"chatcmpl-E2KUzZIergSd35up05B7sIbjPkp6D","object":"chat.completion.chunk","created":1784224161,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_349838c1f9","choices":[],"usage":{"prompt_tokens":98,"completion_tokens":20,"total_tokens":118,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}},"obfuscation":"5WnDCGymh"} + + + data: [DONE] + + + ' + headers: + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a1c2d7cdc8a16d93-YYZ + connection: + - keep-alive + content-type: + - text/event-stream; charset=utf-8 + date: + - Thu, 16 Jul 2026 17:49:21 GMT + openai-organization: + - braintrust-data + openai-processing-ms: + - '370' + openai-project: + - proj_vsCSXafhhByzWOThMrJcZiw9 + openai-version: + - '2020-10-01' + server: + - cloudflare + set-cookie: + - __cf_bm=H7UHWDbN4ftdoy5V6eH8uHfYRf1ufxrsNJJMugkzlQ0-1784224160.9227166-1.0.1.1-B7NE72Bmy_e4yLwLm8Brqe11l5C_sf6drK_MFaPSTtjDFgCJ_J9hS0EgYGuHe3Th2QUaJXhHGld3nJRXe6FFv84OBgVixU9IIpFC_AWHrFtidtzJrIoBkIy6nZUXGACK; + HttpOnly; SameSite=None; Secure; Path=/; Domain=api.openai.com; Expires=Thu, + 16 Jul 2026 18:19:21 GMT + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + transfer-encoding: + - chunked + x-content-type-options: + - nosniff + x-openai-proxy-wasm: + - v0.1 + x-ratelimit-limit-requests: + - '30000' + x-ratelimit-limit-tokens: + - '150000000' + x-ratelimit-remaining-requests: + - '29999' + x-ratelimit-remaining-tokens: + - '149999982' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_536096b82554418bb79b86df1a8ea9f0 + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_model_call_error_propagates.yaml b/py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_model_call_error_propagates.yaml new file mode 100644 index 00000000..c92e84c4 --- /dev/null +++ b/py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_model_call_error_propagates.yaml @@ -0,0 +1,86 @@ +interactions: +- request: + body: '{"messages":[{"role":"user","name":"user","content":[{"type":"text","text":"hello"}]}],"model":"gpt-4o-mini","stream":false,"temperature":0.0}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '142' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.31.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.31.0 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"error\": {\n \"message\": \"Incorrect API key provided: sk-inval******************-key. + You can find your API key at https://platform.openai.com/account/api-keys.\",\n + \ \"type\": \"invalid_request_error\",\n \"code\": \"invalid_api_key\",\n + \ \"param\": null\n },\n \"status\": 401\n}" + headers: + access-control-expose-headers: + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a1c2e5cd0db2813d-YYZ + connection: + - keep-alive + content-length: + - '276' + content-type: + - text/plain + date: + - Thu, 16 Jul 2026 17:58:54 GMT + server: + - cloudflare + set-cookie: + - __cf_bm=GqTI1NVehL5HS77TRbwTnqSHvyM6krPBfYSj44NuYbc-1784224734.2447186-1.0.1.1-e8_fn_zGfEsxFuFwst9Z.uCVbmUcFOcd5Jx2yz2vpBsjRrhPW5p3f_cd37Ghur9nLlEn5mKQMHCHfABfk3I3b44vPxF1wEGIks2AHUHRPXXJVrQ8LQtWIiBasxMHKoCg; + HttpOnly; SameSite=None; Secure; Path=/; Domain=api.openai.com; Expires=Thu, + 16 Jul 2026 18:28:54 GMT + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + x-content-type-options: + - nosniff + x-error-json: + - ewogICJlcnJvciI6IHsKICAgICJtZXNzYWdlIjogIkluY29ycmVjdCBBUEkga2V5IHByb3ZpZGVkOiBzay1pbnZhbCoqKioqKioqKioqKioqKioqKi1rZXkuIFlvdSBjYW4gZmluZCB5b3VyIEFQSSBrZXkgYXQgaHR0cHM6Ly9wbGF0Zm9ybS5vcGVuYWkuY29tL2FjY291bnQvYXBpLWtleXMuIiwKICAgICJ0eXBlIjogImludmFsaWRfcmVxdWVzdF9lcnJvciIsCiAgICAiY29kZSI6ICJpbnZhbGlkX2FwaV9rZXkiLAogICAgInBhcmFtIjogbnVsbAogIH0sCiAgInN0YXR1cyI6IDQwMQp9 + x-openai-authorization-error: + - '401' + x-openai-ide-error-code: + - invalid_api_key + x-openai-internal-caller: + - unknown_through_ide + x-request-id: + - req_8cce4e43f50847778cd59ebeef2689dd + status: + code: 401 + message: Unauthorized +version: 1 diff --git a/py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_streaming_model_call.yaml b/py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_streaming_model_call.yaml new file mode 100644 index 00000000..d8fa6517 --- /dev/null +++ b/py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_streaming_model_call.yaml @@ -0,0 +1,137 @@ +interactions: +- request: + body: '{"messages":[{"role":"system","name":"system","content":[{"type":"text","text":"You + are concise. Answer in one sentence."}]},{"role":"user","name":"user","content":[{"type":"text","text":"Say + hi in five words."}]}],"model":"gpt-4o-mini","stream":true,"stream_options":{"include_usage":true},"temperature":0.0}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '309' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.31.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.31.0 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: 'data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[{"index":0,"delta":{"role":"assistant","content":"","refusal":null},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"Qto4LLohK"} + + + data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[{"index":0,"delta":{"content":"Hello"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"FLYHlD"} + + + data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[{"index":0,"delta":{"content":" + there"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"ouAw7"} + + + data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[{"index":0,"delta":{"content":","},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"fr3caJn68i"} + + + data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[{"index":0,"delta":{"content":" + how"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"GjRGJib"} + + + data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[{"index":0,"delta":{"content":" + are"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"A7dA6eb"} + + + data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[{"index":0,"delta":{"content":" + you"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"rOnTk0t"} + + + data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[{"index":0,"delta":{"content":"?"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"cxXYkvpzSd"} + + + data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}],"usage":null,"obfuscation":"qp52P"} + + + data: {"id":"chatcmpl-E2KUl6PCzZQh7uDOl2J2Kz7qmVgk6","object":"chat.completion.chunk","created":1784224147,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_60952bc00e","choices":[],"usage":{"prompt_tokens":29,"completion_tokens":7,"total_tokens":36,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}},"obfuscation":"YgvQcVgypbx"} + + + data: [DONE] + + + ' + headers: + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a1c2d778dd033a53-YYZ + connection: + - keep-alive + content-type: + - text/event-stream; charset=utf-8 + date: + - Thu, 16 Jul 2026 17:49:08 GMT + openai-organization: + - braintrust-data + openai-processing-ms: + - '532' + openai-project: + - proj_vsCSXafhhByzWOThMrJcZiw9 + openai-version: + - '2020-10-01' + server: + - cloudflare + set-cookie: + - __cf_bm=b3FpCq2SwFdftlbVe7fYeu_YxloleHrcAAMKTFbL6JE-1784224147.3343399-1.0.1.1-d.29BaCjF0or5XIQElyp1e_lBP3U2mc3m1J9myxd0gdJOE8Kh1UQFKmGbIk._1leJGzIh.Ac_cy._8ewWefWMH3ZB4TfTPzZm_eosJaKFCYuTtEKxQZpCiAqtxorpDDH; + HttpOnly; SameSite=None; Secure; Path=/; Domain=api.openai.com; Expires=Thu, + 16 Jul 2026 18:19:08 GMT + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + transfer-encoding: + - chunked + x-content-type-options: + - nosniff + x-openai-proxy-wasm: + - v0.1 + x-ratelimit-limit-requests: + - '30000' + x-ratelimit-limit-tokens: + - '150000000' + x-ratelimit-remaining-requests: + - '29999' + x-ratelimit-remaining-tokens: + - '149999982' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_3ec4fbf9fb2349d8bc0fb9d6957c352e + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/agentscope/patchers.py b/py/src/braintrust/integrations/agentscope/patchers.py index 9a2b8188..a70bd716 100644 --- a/py/src/braintrust/integrations/agentscope/patchers.py +++ b/py/src/braintrust/integrations/agentscope/patchers.py @@ -4,6 +4,7 @@ from .tracing import ( _agent_call_wrapper, + _capture_openai_stream_usage_wrapper, _fanout_pipeline_wrapper, _model_call_wrapper, _sequential_pipeline_wrapper, @@ -72,6 +73,16 @@ class _OpenAIChatModelPatcher(FunctionWrapperPatcher): wrapper = _model_call_wrapper +class _OpenAIStreamUsagePatcher(FunctionWrapperPatcher): + """Capture usage chunks that AgentScope 1.x does not yield to callers.""" + + name = "agentscope.model.openai_stream_usage" + target_module = "agentscope.model" + target_path = "OpenAIChatModel._parse_openai_stream_completion_response" + version_spec = ">=1,<2" + wrapper = _capture_openai_stream_usage_wrapper + + class _DashScopeChatModelPatcher(FunctionWrapperPatcher): name = "agentscope.model.dashscope" target_module = "agentscope.model" @@ -113,6 +124,7 @@ class ChatModelPatcher(CompositeFunctionWrapperPatcher): name = "agentscope.model" sub_patchers = ( _OpenAIChatModelPatcher, + _OpenAIStreamUsagePatcher, _DashScopeChatModelPatcher, _AnthropicChatModelPatcher, _OllamaChatModelPatcher, diff --git a/py/src/braintrust/integrations/agentscope/test_agentscope.py b/py/src/braintrust/integrations/agentscope/test_agentscope.py index 372d675d..ed58819b 100644 --- a/py/src/braintrust/integrations/agentscope/test_agentscope.py +++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py @@ -1,6 +1,8 @@ # pylint: disable=import-error,no-name-in-module,no-value-for-parameter,unexpected-keyword-arg,no-member import importlib import importlib.metadata +import inspect +import os import pytest from braintrust import logger @@ -36,14 +38,15 @@ def _span_type(span): return span_type.value if hasattr(span_type, "value") else span_type -def _make_model(*, stream: bool = False): +def _make_model(*, stream: bool = False, api_key: str | None = None): from agentscope.model import OpenAIChatModel + resolved_api_key = api_key if api_key is not None else os.environ["OPENAI_API_KEY"] if hasattr(OpenAIChatModel, "Parameters"): from agentscope.credential import OpenAICredential return OpenAIChatModel( - credential=OpenAICredential(api_key="test-api-key"), + credential=OpenAICredential(api_key=resolved_api_key), model="gpt-4o-mini", parameters=OpenAIChatModel.Parameters(temperature=0), stream=stream, @@ -52,12 +55,13 @@ def _make_model(*, stream: bool = False): return OpenAIChatModel( model_name="gpt-4o-mini", + api_key=resolved_api_key, stream=stream, generate_kwargs={"temperature": 0}, ) -def _make_agent(name: str, sys_prompt: str, *, toolkit=None, multi_agent: bool = False): +def _make_agent(name: str, sys_prompt: str, *, toolkit=None, multi_agent: bool = False, model=None): from agentscope.tool import Toolkit if HAS_AGENT_REPLY_API: @@ -66,7 +70,7 @@ def _make_agent(name: str, sys_prompt: str, *, toolkit=None, multi_agent: bool = agent = Agent( name=name, system_prompt=sys_prompt, - model=_make_model(), + model=model or _make_model(), toolkit=toolkit or Toolkit(), ) else: @@ -77,7 +81,7 @@ def _make_agent(name: str, sys_prompt: str, *, toolkit=None, multi_agent: bool = agent = ReActAgent( name=name, sys_prompt=sys_prompt, - model=_make_model(), + model=model or _make_model(), formatter=OpenAIMultiAgentFormatter() if multi_agent else OpenAIChatFormatter(), toolkit=toolkit or Toolkit(), memory=InMemoryMemory(), @@ -89,11 +93,24 @@ def _make_agent(name: str, sys_prompt: str, *, toolkit=None, multi_agent: bool = return agent +def _make_user_msg(content): + from agentscope.message import Msg + + if HAS_USER_MSG: + from agentscope.message import UserMsg + + return UserMsg("user", content) + return Msg(name="user", content=content, role="user") + + +async def _run_agent(agent, content): + msg = _make_user_msg(content) + return await (agent.reply(msg) if HAS_AGENT_REPLY_API else agent(msg)) + + @pytest.mark.vcr @pytest.mark.asyncio async def test_agentscope_simple_agent_run(memory_logger): - from agentscope.message import Msg - assert not memory_logger.pop() agent = _make_agent( @@ -101,17 +118,7 @@ async def test_agentscope_simple_agent_run(memory_logger): "You are a concise assistant. Answer in one sentence.", ) - if HAS_USER_MSG: - from agentscope.message import UserMsg - - message = UserMsg("user", "Say hello in exactly two words.") - else: - message = Msg( - name="user", - content="Say hello in exactly two words.", - role="user", - ) - response = await (agent.reply(message) if HAS_AGENT_REPLY_API else agent(message)) + response = await _run_agent(agent, "Say hello in exactly two words.") assert response is not None @@ -122,15 +129,20 @@ async def test_agentscope_simple_agent_run(memory_logger): assert agent_span["context"]["span_origin"]["instrumentation"]["name"] == "agentscope-auto" assert _span_type(agent_span) == "task" assert llm_spans - assert llm_spans[0]["metadata"]["model"] == "gpt-4o-mini" - assert "args" not in llm_spans[0]["input"] - assert llm_spans[0]["input"]["messages"][0]["role"] == "system" - assert llm_spans[0]["input"]["messages"][1]["role"] == "user" - assert llm_spans[0]["input"]["messages"][1]["content"][0]["text"] == "Say hello in exactly two words." - assert llm_spans[0]["output"]["role"] == "assistant" - assert llm_spans[0]["output"]["content"][0]["text"] # non-empty LLM response - assert "usage" not in llm_spans[0]["output"] - assert agent_span["span_id"] in llm_spans[0]["span_parents"] + llm_span = llm_spans[0] + assert llm_span["metadata"]["model"] == "gpt-4o-mini" + assert llm_span["metadata"]["provider"] == "openai" + assert "args" not in llm_span["input"] + assert llm_span["input"]["messages"][0]["role"] == "system" + assert llm_span["input"]["messages"][1]["role"] == "user" + assert llm_span["input"]["messages"][1]["content"][0]["text"] == "Say hello in exactly two words." + assert llm_span["output"]["role"] == "assistant" + assert llm_span["output"]["content"][0]["text"] # non-empty LLM response + assert "usage" not in llm_span["output"] + assert llm_span["metrics"]["prompt_tokens"] > 0 + assert llm_span["metrics"]["completion_tokens"] > 0 + assert llm_span["metrics"]["tokens"] > 0 + assert agent_span["span_id"] in llm_span["span_parents"] @pytest.mark.skipif(IS_AGENTSCOPE_V2, reason="AgentScope 2.x removed the pipeline module") @@ -205,130 +217,69 @@ async def test_agentscope_tool_use_creates_tool_span(memory_logger): llm_spans = [span for span in spans if _span_type(span) == SpanTypeAttribute.LLM] assert llm_spans - assert llm_spans[0]["output"]["role"] == "assistant" - assert llm_spans[0]["output"]["content"][0]["type"] == "tool_use" - assert "usage" not in llm_spans[0]["output"] - - -@pytest.mark.asyncio -async def test_model_call_wrapper_stream_logs_final_output_and_metrics(memory_logger): - from braintrust.integrations.agentscope.tracing import _model_call_wrapper - - assert not memory_logger.pop() - - class FakeOpenAIChatModel: - model_name = "gpt-4o-mini" - - async def wrapped(*_args, **_kwargs): - async def _stream(): - yield {"content": [{"type": "text", "text": "Hello"}]} - yield { - "content": [{"type": "text", "text": "Hello there!"}], - "usage": {"prompt_tokens": 29, "completion_tokens": 3, "total_tokens": 32}, - } - - return _stream() - - stream = await _model_call_wrapper( - wrapped, - FakeOpenAIChatModel(), - args=([{"role": "user", "content": "Say hi in two words."}],), - kwargs={}, - ) - - chunks = [chunk async for chunk in stream] - - assert chunks[-1]["content"][0]["text"] == "Hello there!" - - spans = memory_logger.pop() - assert len(spans) == 1 - llm_span = spans[0] - - assert _span_type(llm_span) == SpanTypeAttribute.LLM + llm_span = llm_spans[0] assert llm_span["output"]["role"] == "assistant" - assert llm_span["output"]["content"][0]["text"] == "Hello there!" - assert llm_span["metrics"]["prompt_tokens"] == 29 - assert llm_span["metrics"]["completion_tokens"] == 3 - assert llm_span["metrics"]["tokens"] == 32 + assert llm_span["output"]["content"][0]["type"] == "tool_use" + assert "usage" not in llm_span["output"] + # Tool definitions belong in metadata.tools, NOT input. + assert "tools" not in llm_span["input"] + assert llm_span["metadata"].get("tools") + tool_names = {tool.get("name") or tool.get("function", {}).get("name") for tool in llm_span["metadata"]["tools"]} + assert "execute_python_code" in tool_names +@pytest.mark.vcr @pytest.mark.asyncio -async def test_model_call_wrapper_stream_span_covers_full_stream_duration(memory_logger): - """Span end timestamp must be recorded after the stream is fully consumed, not before.""" - import asyncio - - from braintrust.integrations.agentscope.tracing import _model_call_wrapper - +async def test_agentscope_streaming_model_call(memory_logger): + """A streaming LLM call must produce one span with accumulated output + TTFT.""" assert not memory_logger.pop() - class FakeModel: - model_name = "gpt-4o-mini" - - async def wrapped(*_args, **_kwargs): - async def _stream(): - for i in range(3): - await asyncio.sleep(0.1) - yield {"content": [{"type": "text", "text": f"chunk-{i}"}]} + model = _make_model(stream=True) + agent = _make_agent("Streamer", "You are concise. Answer in one sentence.", model=model) - return _stream() - - stream = await _model_call_wrapper( - wrapped, - FakeModel(), - args=([{"role": "user", "content": "hi"}],), - kwargs={}, - ) - async for _ in stream: - pass + response = await _run_agent(agent, "Say hi in five words.") + assert response is not None spans = memory_logger.pop() - assert len(spans) == 1 - span = spans[0] - m = span.get("metrics", {}) - duration_ms = (m["end"] - m["start"]) * 1000 - # Stream takes ~300ms (3 chunks × 100ms). The span duration must reflect that. - assert duration_ms >= 200, f"Span duration {duration_ms:.0f}ms is too short; span ended before stream was consumed" + llm_spans = [span for span in spans if _span_type(span) == SpanTypeAttribute.LLM] + # One span per API call, not per chunk. + assert len(llm_spans) >= 1 + llm_span = llm_spans[0] -@pytest.mark.asyncio -async def test_toolkit_call_tool_function_wrapper_stream_span_covers_full_stream_duration(memory_logger): - """Tool span end timestamp must be recorded after the stream is fully consumed, not before.""" - import asyncio + assert llm_span["metadata"]["provider"] == "openai" + assert llm_span["metadata"]["model"] == "gpt-4o-mini" + assert llm_span["output"]["role"] == "assistant" + assert llm_span["output"]["content"] + assert llm_span["metrics"]["time_to_first_token"] > 0 + assert llm_span["metrics"]["prompt_tokens"] > 0 + assert llm_span["metrics"]["completion_tokens"] > 0 + assert llm_span["metrics"]["tokens"] > 0 - from braintrust.integrations.agentscope.tracing import _toolkit_call_tool_function_wrapper +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_agentscope_model_call_error_propagates(memory_logger): + """Provider errors must propagate and the span must log the Exception instance.""" assert not memory_logger.pop() - class FakeToolkit: - pass - - class FakeToolCall: - name = "my_tool" + model = _make_model(api_key="sk-invalid-braintrust-test-key") - async def wrapped(*_args, **_kwargs): - async def _stream(): - for i in range(3): - await asyncio.sleep(0.1) - yield f"chunk-{i}" + messages = [_make_user_msg("hello")] if IS_AGENTSCOPE_V2 else [{"role": "user", "content": "hello"}] + with pytest.raises(Exception) as exc_info: + await model(messages) - return _stream() - - stream = await _toolkit_call_tool_function_wrapper( - wrapped, - FakeToolkit(), - args=(FakeToolCall(),), - kwargs={}, - ) - async for _ in stream: - pass + assert type(exc_info.value).__module__.startswith("openai") spans = memory_logger.pop() - assert len(spans) == 1 - span = spans[0] - m = span.get("metrics", {}) - duration_ms = (m["end"] - m["start"]) * 1000 - # Stream takes ~300ms (3 chunks × 100ms). The span duration must reflect that. - assert duration_ms >= 200, f"Span duration {duration_ms:.0f}ms is too short; span ended before stream was consumed" + llm_spans = [span for span in spans if _span_type(span) == SpanTypeAttribute.LLM] + assert llm_spans + llm_span = llm_spans[0] + assert llm_span["metadata"]["provider"] == "openai" + # error is logged (the exact serialized shape is Braintrust's concern; we + # just verify it was populated as a truthy value, i.e. the wrapper called + # span.log(error=exc) rather than swallowing). + assert llm_span.get("error") @pytest.mark.skipif(not IS_AGENTSCOPE_V2, reason="AgentScope 2.x Toolkit.call_tool API") @@ -360,6 +311,19 @@ def answer(): assert tool_span["input"]["tool_name"] == "answer" +def test_setup_agentscope_is_idempotent(): + """Repeat setup calls must not double-wrap patched targets.""" + from agentscope.model import OpenAIChatModel + + wrapped = inspect.getattr_static(OpenAIChatModel, "__call__") + assert hasattr(wrapped, "__wrapped__") + + setup_agentscope(project_name=PROJECT_NAME) + setup_agentscope(project_name=PROJECT_NAME) + + assert inspect.getattr_static(OpenAIChatModel, "__call__") is wrapped + + class TestAutoInstrumentAgentScope: def test_auto_instrument_agentscope(self): verify_autoinstrument_script("test_auto_agentscope.py") diff --git a/py/src/braintrust/integrations/agentscope/tracing.py b/py/src/braintrust/integrations/agentscope/tracing.py index f96e3cfd..871b571f 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -2,13 +2,23 @@ import contextlib import inspect +import time from contextlib import aclosing +from contextvars import ContextVar from typing import Any +from braintrust.integrations.utils import ( + _is_supported_metric_value, + _normalize_chat_messages, + _parse_openai_usage_metrics, +) from braintrust.logger import start_span as _bt_start_span +from braintrust.span_types import SpanTypeAttribute +from braintrust.util import clean_nones _INSTRUMENTATION = "agentscope-auto" +_STREAM_METRICS: ContextVar[dict[str, float] | None] = ContextVar("agentscope_stream_metrics", default=None) def start_span(*args, **kwargs): @@ -18,8 +28,55 @@ def start_span(*args, **kwargs): return _bt_start_span(*args, **kwargs) -from braintrust.span_types import SpanTypeAttribute -from braintrust.util import clean_nones +# Model class → canonical provider slug (matches JS SDK + pricing lookups). +_PROVIDER_BY_MODEL_CLASS = { + "OpenAIChatModel": "openai", + "AnthropicChatModel": "anthropic", + "GeminiChatModel": "google", + "DashScopeChatModel": "dashscope", + "OllamaChatModel": "ollama", + "TrinityChatModel": "trinity", +} + +# Config kwargs safe to surface in metadata. Keeping this explicit avoids +# leaking credentials or other non-config kwargs some providers accept. +_METADATA_CONFIG_KEYS = frozenset( + { + "temperature", + "top_p", + "top_k", + "max_tokens", + "max_output_tokens", + "stop", + "stop_sequences", + "n", + "seed", + "response_format", + "reasoning_effort", + "frequency_penalty", + "presence_penalty", + "stream", + } +) + +# Provider usage payloads that vary between chat- and responses-style APIs. +_USAGE_NAME_MAP = { + "input_tokens": "prompt_tokens", + "output_tokens": "completion_tokens", + "total_tokens": "tokens", +} +_USAGE_PREFIX_MAP = { + "input": "prompt", + "output": "completion", +} + + +def _kw_or_pos(kwargs: dict[str, Any], key: str, args: Any, index: int) -> Any: + """Return ``kwargs[key]`` if set, otherwise ``args[index]`` if in range.""" + value = kwargs.get(key) + if value is not None: + return value + return args[index] if len(args) > index else None def _args_kwargs_input(args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: @@ -48,34 +105,26 @@ def _pipeline_metadata(args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: def _extract_metrics(*candidates: Any) -> dict[str, float] | None: - key_map = { - "prompt_tokens": "prompt_tokens", - "input_tokens": "prompt_tokens", - "completion_tokens": "completion_tokens", - "output_tokens": "completion_tokens", - "total_tokens": "tokens", - "tokens": "tokens", - } - for candidate in candidates: - data = _field_value(candidate, "usage") or candidate - - metrics = {} - for source_key, target_key in key_map.items(): - value = _field_value(data, source_key) - if isinstance(value, (int, float)): - metrics[target_key] = float(value) + usage = _field_value(candidate, "usage") + if usage is None: + continue + parsed = _parse_openai_usage_metrics( + usage, + token_name_map=_USAGE_NAME_MAP, + token_prefix_map=_USAGE_PREFIX_MAP, + ) + metrics = {k: float(v) for k, v in parsed.items() if _is_supported_metric_value(v) and v >= 0} + if "tokens" not in metrics and "prompt_tokens" in metrics and "completion_tokens" in metrics: + metrics["tokens"] = metrics["prompt_tokens"] + metrics["completion_tokens"] if metrics: return metrics - return None def _model_provider_name(instance: Any) -> str: class_name = instance.__class__.__name__ - if class_name.endswith("Model"): - return class_name[: -len("Model")] - return class_name + return _PROVIDER_BY_MODEL_CLASS.get(class_name, class_name) def _model_metadata(instance: Any) -> dict[str, Any]: @@ -89,41 +138,22 @@ def _model_metadata(instance: Any) -> dict[str, Any]: def _model_call_input(args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: - messages = kwargs.get("messages") - if messages is None and args: - messages = args[0] + return clean_nones({"messages": _normalize_chat_messages(_kw_or_pos(kwargs, "messages", args, 0))}) - tools = kwargs.get("tools") - if tools is None and len(args) > 1: - tools = args[1] - - tool_choice = kwargs.get("tool_choice") - if tool_choice is None and len(args) > 2: - tool_choice = args[2] - - structured_model = kwargs.get("structured_model") - if structured_model is None and len(args) > 3: - structured_model = args[3] +def _model_call_metadata(instance: Any, args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: + extra = {k: v for k, v in kwargs.items() if k in _METADATA_CONFIG_KEYS and v is not None} return clean_nones( { - "messages": messages, - "tools": tools, - "tool_choice": tool_choice, - "structured_model": structured_model, + **_model_metadata(instance), + "tools": _kw_or_pos(kwargs, "tools", args, 1), + "tool_choice": _kw_or_pos(kwargs, "tool_choice", args, 2), + "structured_model": _kw_or_pos(kwargs, "structured_model", args, 3), + **extra, } ) -def _model_call_metadata(instance: Any, kwargs: dict[str, Any]) -> dict[str, Any]: - extra_kwargs = { - key: value - for key, value in kwargs.items() - if key not in {"messages", "tools", "tool_choice", "structured_model"} and value is not None - } - return {**_model_metadata(instance), **extra_kwargs} - - def _model_call_output(result: Any) -> Any: if isinstance(result, dict): data = result @@ -148,10 +178,7 @@ def _model_call_output(result: Any) -> Any: def _field_value(data: Any, key: str) -> Any: if isinstance(data, dict): return data.get(key) - try: - return getattr(data, key, None) - except Exception: - return None + return getattr(data, key, None) def _tool_name(tool_call: Any) -> str: @@ -182,7 +209,7 @@ async def _wrapper(wrapped: Any, instance: Any, args: Any, kwargs: dict[str, Any span.log(output=result) return result except Exception as exc: - span.log(error=str(exc)) + span.log(error=exc) raise return _wrapper @@ -211,15 +238,31 @@ def _is_async_iterator(value: Any) -> bool: return False -def _deferred_stream_trace(result: Any, span: Any, stack: contextlib.ExitStack, log_fn: Any) -> Any: - """Wrap an async iterator so the span stays open until the stream is consumed.""" +def _deferred_stream_trace( + result: Any, + span: Any, + stack: contextlib.ExitStack, + log_fn: Any, + on_first_chunk: Any = None, +) -> Any: + """Wrap an async iterator so the span stays open until the stream is consumed. + + ``log_fn(span, last_chunk)`` is invoked once at stream end. ``on_first_chunk`` + (optional) is invoked with no arguments the first time a chunk is yielded, e.g. + to stamp ``time_to_first_token``. + """ deferred = stack.pop_all() async def _trace(): with deferred: last_chunk = None + first_seen = False async with aclosing(result) as agen: async for chunk in agen: + if not first_seen: + first_seen = True + if on_first_chunk is not None: + on_first_chunk() last_chunk = chunk yield chunk if last_chunk is not None: @@ -250,15 +293,53 @@ async def _toolkit_call_tool_function_wrapper(wrapped: Any, instance: Any, args: if inspect.isawaitable(result): result = await result if _is_async_iterator(result): - return _deferred_stream_trace(result, span, stack, lambda s, chunk: s.log(output=chunk)) + return _deferred_stream_trace( + result, + span, + stack, + lambda s, chunk: s.log(output=chunk), + ) span.log(output=result) return result except Exception as exc: - span.log(error=str(exc)) + span.log(error=exc) raise +def _capture_openai_stream_usage_wrapper( + wrapped: Any, + _instance: Any, + args: Any, + kwargs: dict[str, Any], +) -> Any: + """Capture AgentScope 1.x usage-only OpenAI chunks without changing its stream.""" + metrics = _STREAM_METRICS.get() + if metrics is None: + return wrapped(*args, **kwargs) + + positional = list(args) + response = positional[1] if len(positional) > 1 else kwargs.get("response") + if response is None: + return wrapped(*args, **kwargs) + + async def _capture_usage(): + async for chunk in response: + # OpenAI only sets `usage` on the terminal chunk when + # `stream_options.include_usage` is set; skip parsing every chunk. + if _field_value(chunk, "usage") is not None: + chunk_metrics = _extract_metrics(chunk) + if chunk_metrics: + metrics.update(chunk_metrics) + yield chunk + + if len(positional) > 1: + positional[1] = _capture_usage() + else: + kwargs = {**kwargs, "response": _capture_usage()} + return wrapped(*positional, **kwargs) + + async def _model_call_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: dict[str, Any]) -> Any: with contextlib.ExitStack() as stack: span = stack.enter_context( @@ -266,21 +347,39 @@ async def _model_call_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: di name=f"{_model_provider_name(instance)}.call", type=SpanTypeAttribute.LLM, input=_model_call_input(args, kwargs), - metadata=_model_call_metadata(instance, kwargs), + metadata=_model_call_metadata(instance, args, kwargs), ) ) try: - result = await wrapped(*args, **kwargs) + request_start_time = time.time() + captured_stream_metrics: dict[str, float] = {} + time_to_first_token: list[float | None] = [None] + stream_metrics_token = _STREAM_METRICS.set(captured_stream_metrics) + try: + result = await wrapped(*args, **kwargs) + finally: + _STREAM_METRICS.reset(stream_metrics_token) if _is_async_iterator(result): + + def _stamp_ttft() -> None: + time_to_first_token[0] = time.time() - request_start_time + + def _log_final(s: Any, chunk: Any) -> None: + metrics = {**captured_stream_metrics, **(_extract_metrics(chunk) or {})} + if time_to_first_token[0] is not None: + metrics["time_to_first_token"] = time_to_first_token[0] + s.log(output=_model_call_output(chunk), metrics=metrics or None) + return _deferred_stream_trace( result, span, stack, - lambda s, chunk: s.log(output=_model_call_output(chunk), metrics=_extract_metrics(chunk)), + _log_final, + on_first_chunk=_stamp_ttft, ) span.log(output=_model_call_output(result), metrics=_extract_metrics(result)) return result except Exception as exc: - span.log(error=str(exc)) + span.log(error=exc) raise diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_auto_agentscope.py b/py/src/braintrust/integrations/auto_test_scripts/test_auto_agentscope.py index 482ba9e4..0347aead 100644 --- a/py/src/braintrust/integrations/auto_test_scripts/test_auto_agentscope.py +++ b/py/src/braintrust/integrations/auto_test_scripts/test_auto_agentscope.py @@ -98,6 +98,7 @@ assert agent_span["span_attributes"]["type"].value == "task" assert llm_spans, "Should have at least one LLM span" assert llm_spans[0]["metadata"]["model"] == "gpt-4o-mini" + assert llm_spans[0]["metadata"]["provider"] == "openai" assert agent_span["span_id"] in llm_spans[0]["span_parents"] print("SUCCESS")