From 76a9540a2d16282caec9e8b21fd9d8d4e29a9b45 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 17:00:57 +0000 Subject: [PATCH 1/5] fix(agentscope): align integration with sdk-integrations SKILL - move tools/tool_choice/structured_model from LLM span input into metadata.tools (per SKILL span design) - log the Exception instance instead of str(exc) so Braintrust receives the raw error, not a pre-formatted string - record time_to_first_token for streaming model calls - replace class-name-stripping with an explicit provider map (openai / anthropic / google / dashscope / ollama / trinity) - materialize inline multimodal message parts via _normalize_chat_messages - restrict metadata kwarg pass-through to an explicit config allowlist to avoid leaking credentials passed as kwargs - drop fake-based streaming tests, tighten cassette assertions (provider, metrics, metadata.tools), and add new VCR tests for streaming and error propagation (requires cassette re-record) Co-Authored-By: Claude Opus 4.7 --- .../agentscope/test_agentscope.py | 222 +++++++----------- .../integrations/agentscope/tracing.py | 97 ++++++-- .../auto_test_scripts/test_auto_agentscope.py | 1 + 3 files changed, 165 insertions(+), 155 deletions(-) diff --git a/py/src/braintrust/integrations/agentscope/test_agentscope.py b/py/src/braintrust/integrations/agentscope/test_agentscope.py index 372d675d..395b13be 100644 --- a/py/src/braintrust/integrations/agentscope/test_agentscope.py +++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py @@ -36,14 +36,14 @@ 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 = "test-api-key"): from agentscope.model import OpenAIChatModel if hasattr(OpenAIChatModel, "Parameters"): from agentscope.credential import OpenAICredential return OpenAIChatModel( - credential=OpenAICredential(api_key="test-api-key"), + credential=OpenAICredential(api_key=api_key), model="gpt-4o-mini", parameters=OpenAIChatModel.Parameters(temperature=0), stream=stream, @@ -57,7 +57,7 @@ def _make_model(*, stream: bool = False): ) -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 +66,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 +77,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 +89,19 @@ 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") + + @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 +109,11 @@ 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 ( + agent.reply(_make_user_msg("Say hello in exactly two words.")) + if HAS_AGENT_REPLY_API + else agent(_make_user_msg("Say hello in exactly two words.")) + ) assert response is not None @@ -122,15 +124,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 +212,68 @@ 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}"}]} - - return _stream() + model = _make_model(stream=True) + agent = _make_agent("Streamer", "You are concise. Answer in one sentence.", model=model) - stream = await _model_call_wrapper( - wrapped, - FakeModel(), - args=([{"role": "user", "content": "hi"}],), - kwargs={}, + response = await ( + agent.reply(_make_user_msg("Say hi in five words.")) + if HAS_AGENT_REPLY_API + else agent(_make_user_msg("Say hi in five words.")) ) - async for _ in stream: - pass + 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"]["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}" - - return _stream() - - stream = await _toolkit_call_tool_function_wrapper( - wrapped, - FakeToolkit(), - args=(FakeToolCall(),), - kwargs={}, - ) - async for _ in stream: - pass + with pytest.raises(Exception): + await model([{"role": "user", "content": "hello"}]) 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 +305,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.""" + setup_agentscope(project_name=PROJECT_NAME) + setup_agentscope(project_name=PROJECT_NAME) + + from agentscope.model import OpenAIChatModel + + wrapped = OpenAIChatModel.__call__ + assert hasattr(wrapped, "__wrapped__") + # A second layer of wrapping would expose a nested __wrapped__.__wrapped__. + assert not hasattr(wrapped.__wrapped__, "__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..24c8d27c 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -2,10 +2,14 @@ import contextlib import inspect +import time from contextlib import aclosing from typing import Any +from braintrust.integrations.utils import _normalize_chat_messages 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" @@ -18,8 +22,36 @@ 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", + } +) def _args_kwargs_input(args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: @@ -73,9 +105,7 @@ def _extract_metrics(*candidates: Any) -> dict[str, float] | 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]: @@ -93,6 +123,10 @@ def _model_call_input(args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: if messages is None and args: messages = args[0] + return clean_nones({"messages": _normalize_chat_messages(messages)}) + + +def _model_call_metadata(instance: Any, args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: tools = kwargs.get("tools") if tools is None and len(args) > 1: tools = args[1] @@ -105,25 +139,19 @@ def _model_call_input(args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: if structured_model is None and len(args) > 3: structured_model = args[3] + extra = {key: kwargs[key] for key in _METADATA_CONFIG_KEYS if key in kwargs and kwargs[key] is not None} + return clean_nones( { - "messages": messages, + **_model_metadata(instance), "tools": tools, "tool_choice": tool_choice, "structured_model": structured_model, + **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 @@ -182,7 +210,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,19 +239,28 @@ def _is_async_iterator(value: Any) -> bool: return False -def _deferred_stream_trace(result: Any, span: Any, stack: contextlib.ExitStack, log_fn: Any) -> Any: +def _deferred_stream_trace( + result: Any, + span: Any, + stack: contextlib.ExitStack, + log_fn: Any, + request_start_time: float | None = None, +) -> Any: """Wrap an async iterator so the span stays open until the stream is consumed.""" deferred = stack.pop_all() async def _trace(): with deferred: last_chunk = None + time_to_first_token: float | None = None async with aclosing(result) as agen: async for chunk in agen: + if time_to_first_token is None and request_start_time is not None: + time_to_first_token = time.time() - request_start_time last_chunk = chunk yield chunk if last_chunk is not None: - log_fn(span, last_chunk) + log_fn(span, last_chunk, time_to_first_token) return _trace() @@ -250,15 +287,27 @@ 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, _ttft: 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 _log_model_stream_chunk(span: Any, chunk: Any, time_to_first_token: float | None) -> None: + metrics = _extract_metrics(chunk) or {} + if time_to_first_token is not None: + metrics["time_to_first_token"] = time_to_first_token + span.log(output=_model_call_output(chunk), metrics=metrics or None) + + 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 +315,23 @@ 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: + request_start_time = time.time() result = await wrapped(*args, **kwargs) if _is_async_iterator(result): return _deferred_stream_trace( result, span, stack, - lambda s, chunk: s.log(output=_model_call_output(chunk), metrics=_extract_metrics(chunk)), + _log_model_stream_chunk, + request_start_time=request_start_time, ) 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") From 46e04139d8a64c8282528f5c71506f6b4c6ca8a6 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 16 Jul 2026 13:54:07 -0400 Subject: [PATCH 2/5] record cassettes --- ...gentscope_model_call_error_propagates.yaml | 86 ++++++++++ .../test_agentscope_streaming_model_call.yaml | 152 ++++++++++++++++++ .../test_agentscope_streaming_model_call.yaml | 137 ++++++++++++++++ .../integrations/agentscope/patchers.py | 12 ++ .../agentscope/test_agentscope.py | 22 ++- .../integrations/agentscope/tracing.py | 61 ++++++- 6 files changed, 457 insertions(+), 13 deletions(-) create mode 100644 py/src/braintrust/integrations/agentscope/cassettes/1.0.0/test_agentscope_model_call_error_propagates.yaml create mode 100644 py/src/braintrust/integrations/agentscope/cassettes/1.0.0/test_agentscope_streaming_model_call.yaml create mode 100644 py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_streaming_model_call.yaml 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_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 395b13be..18213798 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, api_key: str = "test-api-key"): +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=api_key), + credential=OpenAICredential(api_key=resolved_api_key), model="gpt-4o-mini", parameters=OpenAIChatModel.Parameters(temperature=0), stream=stream, @@ -52,6 +55,7 @@ def _make_model(*, stream: bool = False, api_key: str = "test-api-key"): return OpenAIChatModel( model_name="gpt-4o-mini", + api_key=resolved_api_key, stream=stream, generate_kwargs={"temperature": 0}, ) @@ -251,6 +255,8 @@ async def test_agentscope_streaming_model_call(memory_logger): 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 @@ -307,15 +313,15 @@ def answer(): def test_setup_agentscope_is_idempotent(): """Repeat setup calls must not double-wrap patched targets.""" - setup_agentscope(project_name=PROJECT_NAME) - setup_agentscope(project_name=PROJECT_NAME) - from agentscope.model import OpenAIChatModel - wrapped = OpenAIChatModel.__call__ + wrapped = inspect.getattr_static(OpenAIChatModel, "__call__") assert hasattr(wrapped, "__wrapped__") - # A second layer of wrapping would expose a nested __wrapped__.__wrapped__. - assert not hasattr(wrapped.__wrapped__, "__wrapped__") + + setup_agentscope(project_name=PROJECT_NAME) + setup_agentscope(project_name=PROJECT_NAME) + + assert inspect.getattr_static(OpenAIChatModel, "__call__") is wrapped class TestAutoInstrumentAgentScope: diff --git a/py/src/braintrust/integrations/agentscope/tracing.py b/py/src/braintrust/integrations/agentscope/tracing.py index 24c8d27c..6f3e028b 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -4,6 +4,7 @@ import inspect import time from contextlib import aclosing +from contextvars import ContextVar from typing import Any from braintrust.integrations.utils import _normalize_chat_messages @@ -13,6 +14,7 @@ _INSTRUMENTATION = "agentscope-auto" +_STREAM_METRICS: ContextVar[dict[str, float] | None] = ContextVar("agentscope_stream_metrics", default=None) def start_span(*args, **kwargs): @@ -95,8 +97,12 @@ def _extract_metrics(*candidates: Any) -> dict[str, float] | None: metrics = {} for source_key, target_key in key_map.items(): value = _field_value(data, source_key) - if isinstance(value, (int, float)): + if isinstance(value, (int, float)) and not isinstance(value, bool) and value >= 0: metrics[target_key] = float(value) + + 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 @@ -301,8 +307,43 @@ async def _toolkit_call_tool_function_wrapper(wrapped: Any, instance: Any, args: raise -def _log_model_stream_chunk(span: Any, chunk: Any, time_to_first_token: float | None) -> None: - metrics = _extract_metrics(chunk) or {} +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: + 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) + + +def _log_model_stream_chunk( + span: Any, + chunk: Any, + time_to_first_token: float | None, + captured_metrics: dict[str, float], +) -> None: + metrics = {**captured_metrics, **(_extract_metrics(chunk) or {})} if time_to_first_token is not None: metrics["time_to_first_token"] = time_to_first_token span.log(output=_model_call_output(chunk), metrics=metrics or None) @@ -320,13 +361,23 @@ async def _model_call_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: di ) try: request_start_time = time.time() - result = await wrapped(*args, **kwargs) + captured_stream_metrics: dict[str, float] = {} + 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): return _deferred_stream_trace( result, span, stack, - _log_model_stream_chunk, + lambda s, chunk, ttft: _log_model_stream_chunk( + s, + chunk, + ttft, + captured_stream_metrics, + ), request_start_time=request_start_time, ) From a6ca63decb7f88bb2980edb1d819298fbae7e4ac Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 16 Jul 2026 13:59:48 -0400 Subject: [PATCH 3/5] theres more cassettes --- ...gentscope_model_call_error_propagates.yaml | 86 +++++++++++++++++++ .../agentscope/test_agentscope.py | 7 +- 2 files changed, 91 insertions(+), 2 deletions(-) create mode 100644 py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_model_call_error_propagates.yaml 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/test_agentscope.py b/py/src/braintrust/integrations/agentscope/test_agentscope.py index 18213798..5266fce5 100644 --- a/py/src/braintrust/integrations/agentscope/test_agentscope.py +++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py @@ -268,8 +268,11 @@ async def test_agentscope_model_call_error_propagates(memory_logger): model = _make_model(api_key="sk-invalid-braintrust-test-key") - with pytest.raises(Exception): - await model([{"role": "user", "content": "hello"}]) + messages = [_make_user_msg("hello")] if IS_AGENTSCOPE_V2 else [{"role": "user", "content": "hello"}] + with pytest.raises(Exception) as exc_info: + await model(messages) + + assert type(exc_info.value).__module__.startswith("openai") spans = memory_logger.pop() llm_spans = [span for span in spans if _span_type(span) == SpanTypeAttribute.LLM] From 53433860c07a8d7f3d881def0947b257442968f7 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Thu, 16 Jul 2026 18:04:50 +0000 Subject: [PATCH 4/5] refactor(agentscope): reuse shared usage/metric helpers, thin per-call plumbing - _extract_metrics now delegates to _parse_openai_usage_metrics + _try_to_dict (via the shared helper), so pydantic Usage objects normalize correctly and we pick up cached/reasoning token subkeys for free - _is_supported_metric_value replaces the inline bool/numeric guard - introduce _kw_or_pos to collapse four "kwargs.get else args[idx]" stanzas - iterate kwargs (not the allowlist) for _METADATA_CONFIG_KEYS extraction - drop unnecessary try/except around getattr in _field_value - _deferred_stream_trace now takes an on_first_chunk callback so tool streams no longer accept an unused time_to_first_token param; inline the final-chunk logger closure in _model_call_wrapper - short-circuit _capture_openai_stream_usage_wrapper: only run _extract_metrics when the chunk actually carries a usage field (OpenAI only sets it on the terminal chunk) - add _run_agent test helper to remove duplicated "agent.reply if HAS_AGENT_REPLY_API else agent(...)" conditionals Co-Authored-By: Claude Opus 4.7 --- .../agentscope/test_agentscope.py | 17 +- .../integrations/agentscope/tracing.py | 146 +++++++++--------- 2 files changed, 78 insertions(+), 85 deletions(-) diff --git a/py/src/braintrust/integrations/agentscope/test_agentscope.py b/py/src/braintrust/integrations/agentscope/test_agentscope.py index 5266fce5..ed58819b 100644 --- a/py/src/braintrust/integrations/agentscope/test_agentscope.py +++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py @@ -103,6 +103,11 @@ def _make_user_msg(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): @@ -113,11 +118,7 @@ async def test_agentscope_simple_agent_run(memory_logger): "You are a concise assistant. Answer in one sentence.", ) - response = await ( - agent.reply(_make_user_msg("Say hello in exactly two words.")) - if HAS_AGENT_REPLY_API - else agent(_make_user_msg("Say hello in exactly two words.")) - ) + response = await _run_agent(agent, "Say hello in exactly two words.") assert response is not None @@ -236,11 +237,7 @@ async def test_agentscope_streaming_model_call(memory_logger): model = _make_model(stream=True) agent = _make_agent("Streamer", "You are concise. Answer in one sentence.", model=model) - response = await ( - agent.reply(_make_user_msg("Say hi in five words.")) - if HAS_AGENT_REPLY_API - else agent(_make_user_msg("Say hi in five words.")) - ) + response = await _run_agent(agent, "Say hi in five words.") assert response is not None spans = memory_logger.pop() diff --git a/py/src/braintrust/integrations/agentscope/tracing.py b/py/src/braintrust/integrations/agentscope/tracing.py index 6f3e028b..73efb78f 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -7,7 +7,11 @@ from contextvars import ContextVar from typing import Any -from braintrust.integrations.utils import _normalize_chat_messages +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 @@ -55,6 +59,25 @@ def start_span(*args, **kwargs): } ) +# 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]: return clean_nones( @@ -82,30 +105,20 @@ 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)) and not isinstance(value, bool) and value >= 0: - 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 @@ -125,34 +138,17 @@ 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(messages)}) + return clean_nones({"messages": _normalize_chat_messages(_kw_or_pos(kwargs, "messages", args, 0))}) def _model_call_metadata(instance: Any, args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: - 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] - - extra = {key: kwargs[key] for key in _METADATA_CONFIG_KEYS if key in kwargs and kwargs[key] is not None} - + extra = {k: v for k, v in kwargs.items() if k in _METADATA_CONFIG_KEYS and v is not None} return clean_nones( { **_model_metadata(instance), - "tools": tools, - "tool_choice": tool_choice, - "structured_model": structured_model, + "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, } ) @@ -182,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: @@ -250,23 +243,30 @@ def _deferred_stream_trace( span: Any, stack: contextlib.ExitStack, log_fn: Any, - request_start_time: float | None = None, + on_first_chunk: Any = None, ) -> Any: - """Wrap an async iterator so the span stays open until the stream is consumed.""" + """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 - time_to_first_token: float | None = None + first_seen = False async with aclosing(result) as agen: async for chunk in agen: - if time_to_first_token is None and request_start_time is not None: - time_to_first_token = time.time() - request_start_time + 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: - log_fn(span, last_chunk, time_to_first_token) + log_fn(span, last_chunk) return _trace() @@ -297,7 +297,7 @@ async def _toolkit_call_tool_function_wrapper(wrapped: Any, instance: Any, args: result, span, stack, - lambda s, chunk, _ttft: s.log(output=chunk), + lambda s, chunk: s.log(output=chunk), ) span.log(output=result) @@ -325,9 +325,12 @@ def _capture_openai_stream_usage_wrapper( async def _capture_usage(): async for chunk in response: - chunk_metrics = _extract_metrics(chunk) - if chunk_metrics: - metrics.update(chunk_metrics) + # 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: @@ -337,18 +340,6 @@ async def _capture_usage(): return wrapped(*positional, **kwargs) -def _log_model_stream_chunk( - span: Any, - chunk: Any, - time_to_first_token: float | None, - captured_metrics: dict[str, float], -) -> None: - metrics = {**captured_metrics, **(_extract_metrics(chunk) or {})} - if time_to_first_token is not None: - metrics["time_to_first_token"] = time_to_first_token - span.log(output=_model_call_output(chunk), metrics=metrics or None) - - 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( @@ -362,23 +353,28 @@ async def _model_call_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: di try: 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, ttft: _log_model_stream_chunk( - s, - chunk, - ttft, - captured_stream_metrics, - ), - request_start_time=request_start_time, + _log_final, + on_first_chunk=_stamp_ttft, ) span.log(output=_model_call_output(result), metrics=_extract_metrics(result)) From 02c755bb0ce88431e2b58178af4ee85945d6d338 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 16 Jul 2026 14:06:27 -0400 Subject: [PATCH 5/5] format --- py/src/braintrust/integrations/agentscope/tracing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/py/src/braintrust/integrations/agentscope/tracing.py b/py/src/braintrust/integrations/agentscope/tracing.py index 73efb78f..871b571f 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -360,6 +360,7 @@ async def _model_call_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: di 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