From 6b39c88d543a1f212b050ca5fb9db6bde1697af5 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Fri, 17 Apr 2026 16:26:51 -0600 Subject: [PATCH 1/5] BTX_SPEC_ROOT to override btx test files --- btx/build.gradle | 12 ++++++++++-- .../java/dev/braintrust/sdkspecimpl/SpecLoader.java | 13 ++++++++++--- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/btx/build.gradle b/btx/build.gradle index 27de2956..72f55b56 100644 --- a/btx/build.gradle +++ b/btx/build.gradle @@ -95,6 +95,14 @@ tasks.register('fetchSpec', Exec) { } test { - dependsOn fetchSpec - systemProperty 'btx.spec.root', specOutputDir.get().asFile.absolutePath + '/test/llm_span' + // BTX_SPEC_ROOT env var lets you point at a local checkout of braintrust-spec instead of the + // fetched copy. When set, fetchSpec is skipped and the system property is set to that path. + // Example: BTX_SPEC_ROOT=/Users/you/braintrust/spec/test/llm_span ./gradlew btx:test + def specRootEnv = System.getenv('BTX_SPEC_ROOT') + if (specRootEnv) { + systemProperty 'btx.spec.root', specRootEnv + } else { + dependsOn fetchSpec + systemProperty 'btx.spec.root', specOutputDir.get().asFile.absolutePath + '/test/llm_span' + } } diff --git a/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecLoader.java b/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecLoader.java index 711cdee3..29157745 100644 --- a/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecLoader.java +++ b/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecLoader.java @@ -33,9 +33,16 @@ public class SpecLoader { /** - * The spec root directory. When the {@code btx.spec.root} system property is set (by the {@code - * fetchSpec} Gradle task), that path is used. Otherwise falls back to the legacy in-tree - * location for local development. + * The spec root directory. Resolved in priority order: + * + *
    + *
  1. The {@code btx.spec.root} system property — set by the Gradle {@code test} task to the + * output of {@code fetchSpec}, or to {@code $BTX_SPEC_ROOT} when that env var is set. + *
  2. The fallback {@code btx/spec/llm_span} in-tree path for ad-hoc local runs. + *
+ * + *

To use a local checkout of braintrust-spec: {@code + * BTX_SPEC_ROOT=/path/to/spec/test/llm_span ./gradlew btx:test} */ private static final String SPEC_ROOT = System.getProperty("btx.spec.root", "btx/spec/llm_span"); From fbd1fcf3032976e031c4c2e36309fac7790c52b0 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Sun, 19 Apr 2026 22:31:56 -0600 Subject: [PATCH 2/5] support pinned version directive for muzzle --- .../gradle/muzzle/MuzzleDirective.groovy | 13 ++++++++++- .../gradle/muzzle/MuzzleExtension.groovy | 4 ++-- .../gradle/muzzle/MuzzleTask.groovy | 23 +++++++++++-------- 3 files changed, 28 insertions(+), 12 deletions(-) diff --git a/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleDirective.groovy b/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleDirective.groovy index a282f9da..b4f78eb2 100644 --- a/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleDirective.groovy +++ b/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleDirective.groovy @@ -40,6 +40,16 @@ class MuzzleDirective { */ List ignoredInstrumentation = [] + /** + * Explicit list of versions to check instead of resolving a range from Maven Central. + * When set, {@code versions} is ignored and no network fetch is performed. + */ + List pinnedVersions = [] + + void pinVersions(String... versions) { + pinnedVersions.addAll(versions) + } + void skipVersions(String... versions) { skipVersions.addAll(versions) } @@ -54,6 +64,7 @@ class MuzzleDirective { @Override String toString() { - "${assertPass ? 'pass' : 'fail'} { ${group}:${module}:${versions} }" + def versionStr = pinnedVersions ? "pinned[${pinnedVersions.join(', ')}]" : versions + "${assertPass ? 'pass' : 'fail'} { ${group}:${module}:${versionStr} }" } } diff --git a/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleExtension.groovy b/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleExtension.groovy index f561f541..5fcdd2c4 100644 --- a/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleExtension.groovy +++ b/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleExtension.groovy @@ -38,8 +38,8 @@ class MuzzleExtension { if (!directive.module) { throw new IllegalArgumentException("muzzle directive requires 'module'") } - if (!directive.versions) { - throw new IllegalArgumentException("muzzle directive requires 'versions'") + if (!directive.pinnedVersions && !directive.versions) { + throw new IllegalArgumentException("muzzle directive requires either 'versions' or 'pinVersions'") } } } diff --git a/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleTask.groovy b/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleTask.groovy index 2d667c30..6dc68a0e 100644 --- a/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleTask.groovy +++ b/buildSrc/src/main/groovy/dev/braintrust/gradle/muzzle/MuzzleTask.groovy @@ -47,16 +47,21 @@ class MuzzleTask extends DefaultTask { logger.lifecycle("[muzzle] Checking: ${directive}") List versions - try { - versions = MavenVersions.resolve( - directive.group, directive.module, directive.versions, directive.skipVersions) - } catch (Exception e) { - throw new GradleException("[muzzle] Failed to resolve versions for ${directive}: ${e.message}", e) - } + if (directive.pinnedVersions) { + versions = directive.pinnedVersions + logger.lifecycle("[muzzle] Using ${versions.size()} pinned version(s)") + } else { + try { + versions = MavenVersions.resolve( + directive.group, directive.module, directive.versions, directive.skipVersions) + } catch (Exception e) { + throw new GradleException("[muzzle] Failed to resolve versions for ${directive}: ${e.message}", e) + } - if (versions.isEmpty()) { - logger.warn("[muzzle] No versions found for ${directive.group}:${directive.module} in range ${directive.versions}") - continue + if (versions.isEmpty()) { + logger.warn("[muzzle] No versions found for ${directive.group}:${directive.module} in range ${directive.versions}") + continue + } } logger.lifecycle("[muzzle] Found ${versions.size()} version(s) to check") From d21687a1314bd1482966ec2810e817194cb49aab Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Sun, 19 Apr 2026 21:08:56 -0600 Subject: [PATCH 3/5] bedrock cassettes --- ...-400d471f-3f4d-4d3e-9571-cb3f015011c5.json | 1 + ...-57a96b00-9f78-4336-89e1-75f8e511afd2.json | 1 + ...-d2a2ec99-c9bf-48e6-b0c6-eb53a928e698.json | 1 + ...m-df6a23a1-b3d8-403e-a499-bc5d257af868.txt | Bin 0 -> 4932 bytes ...-bfd89f01-c330-4c6d-b9c0-05ca56ed005f.json | 1 + ...m-c43d383c-79df-4107-996c-7c63b26994fd.txt | Bin 0 -> 1093 bytes ...-400d471f-3f4d-4d3e-9571-cb3f015011c5.json | 30 ++++++++++++ ...-57a96b00-9f78-4336-89e1-75f8e511afd2.json | 30 ++++++++++++ ...-d2a2ec99-c9bf-48e6-b0c6-eb53a928e698.json | 30 ++++++++++++ ...-df6a23a1-b3d8-403e-a499-bc5d257af868.json | 30 ++++++++++++ ...-bfd89f01-c330-4c6d-b9c0-05ca56ed005f.json | 30 ++++++++++++ ...-c43d383c-79df-4107-996c-7c63b26994fd.json | 30 ++++++++++++ ...-1a27eb77-e3d5-4987-ae9f-4c31512eb77e.json | 1 + ...-6cc00d1f-84f6-4674-81b5-47f65f1c93e9.json | 1 + ...-d0e9a3e5-918d-4a9b-a923-a3482ab7fb08.json | 1 + ...-fce651bd-9cb7-4b5d-a488-5b3da8a28a16.json | 1 + ...-1a27eb77-e3d5-4987-ae9f-4c31512eb77e.json | 44 +++++++++++++++++ ...-6cc00d1f-84f6-4674-81b5-47f65f1c93e9.json | 42 ++++++++++++++++ ...-d0e9a3e5-918d-4a9b-a923-a3482ab7fb08.json | 45 ++++++++++++++++++ ...-fce651bd-9cb7-4b5d-a488-5b3da8a28a16.json | 42 ++++++++++++++++ ...-5279636b-7eb9-4959-a3be-de675f1f14e1.json | 39 +++++++++++++++ ...-57070a21-af08-425a-886a-32727f7339dd.json | 39 +++++++++++++++ ...-75dd8273-1f17-47fd-a992-66877387d37e.json | 39 +++++++++++++++ ...-eb149792-2d60-4963-9808-9a8bf6326d31.json | 39 +++++++++++++++ 24 files changed, 517 insertions(+) create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-400d471f-3f4d-4d3e-9571-cb3f015011c5.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-57a96b00-9f78-4336-89e1-75f8e511afd2.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-d2a2ec99-c9bf-48e6-b0c6-eb53a928e698.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-stream-df6a23a1-b3d8-403e-a499-bc5d257af868.txt create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.anthropic.claude-3-haiku-20240307-v10_converse-bfd89f01-c330-4c6d-b9c0-05ca56ed005f.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.anthropic.claude-3-haiku-20240307-v10_converse-stream-c43d383c-79df-4107-996c-7c63b26994fd.txt create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/mappings/model_us.amazon.nova-lite-v10_converse-400d471f-3f4d-4d3e-9571-cb3f015011c5.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/mappings/model_us.amazon.nova-lite-v10_converse-57a96b00-9f78-4336-89e1-75f8e511afd2.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/mappings/model_us.amazon.nova-lite-v10_converse-d2a2ec99-c9bf-48e6-b0c6-eb53a928e698.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/mappings/model_us.amazon.nova-lite-v10_converse-stream-df6a23a1-b3d8-403e-a499-bc5d257af868.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/mappings/model_us.anthropic.claude-3-haiku-20240307-v10_converse-bfd89f01-c330-4c6d-b9c0-05ca56ed005f.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/bedrock/mappings/model_us.anthropic.claude-3-haiku-20240307-v10_converse-stream-c43d383c-79df-4107-996c-7c63b26994fd.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/__files/btql-1a27eb77-e3d5-4987-ae9f-4c31512eb77e.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/__files/btql-6cc00d1f-84f6-4674-81b5-47f65f1c93e9.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/__files/btql-d0e9a3e5-918d-4a9b-a923-a3482ab7fb08.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/__files/btql-fce651bd-9cb7-4b5d-a488-5b3da8a28a16.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/mappings/btql-1a27eb77-e3d5-4987-ae9f-4c31512eb77e.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/mappings/btql-6cc00d1f-84f6-4674-81b5-47f65f1c93e9.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/mappings/btql-d0e9a3e5-918d-4a9b-a923-a3482ab7fb08.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/mappings/btql-fce651bd-9cb7-4b5d-a488-5b3da8a28a16.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/mappings/otel_v1_traces-5279636b-7eb9-4959-a3be-de675f1f14e1.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/mappings/otel_v1_traces-57070a21-af08-425a-886a-32727f7339dd.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/mappings/otel_v1_traces-75dd8273-1f17-47fd-a992-66877387d37e.json create mode 100644 test-harness/src/testFixtures/resources/cassettes/braintrust/mappings/otel_v1_traces-eb149792-2d60-4963-9808-9a8bf6326d31.json diff --git a/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-400d471f-3f4d-4d3e-9571-cb3f015011c5.json b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-400d471f-3f4d-4d3e-9571-cb3f015011c5.json new file mode 100644 index 00000000..c4425296 --- /dev/null +++ b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-400d471f-3f4d-4d3e-9571-cb3f015011c5.json @@ -0,0 +1 @@ +{"metrics":{"latencyMs":641},"output":{"message":{"content":[{"text":"Sorry, I can't describe an image. I can only provide information about the color. However, if you want to know about the color of the image, I can provide information about it. The image is in a red color."}],"role":"assistant"}},"stopReason":"end_turn","usage":{"inputTokens":536,"outputTokens":49,"serverToolUsage":{},"totalTokens":585}} \ No newline at end of file diff --git a/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-57a96b00-9f78-4336-89e1-75f8e511afd2.json b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-57a96b00-9f78-4336-89e1-75f8e511afd2.json new file mode 100644 index 00000000..14bc9a83 --- /dev/null +++ b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-57a96b00-9f78-4336-89e1-75f8e511afd2.json @@ -0,0 +1 @@ +{"metrics":{"latencyMs":888},"output":{"message":{"content":[{"text":"The capital of France is Paris. Paris is not only the capital but also the largest city in France. It is situated in the northern central part of the country, along the Seine River. Paris is renowned for its rich history, culture, and landmarks. Some of its most famous attractions include the Eiffel Tower, the Louvre Museum, Notre-Dame Cathedral, and the Champs-Élysées. It is a major global city and a hub for art, fashion, gastronomy, and diplomacy. Paris is divided into 20 arrondissements (municipalities), each with its own unique character and attractions. The city is also known for its significant contributions to various fields such as philosophy, science, and literature."}],"role":"assistant"}},"stopReason":"end_turn","usage":{"inputTokens":7,"outputTokens":144,"serverToolUsage":{},"totalTokens":151}} \ No newline at end of file diff --git a/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-d2a2ec99-c9bf-48e6-b0c6-eb53a928e698.json b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-d2a2ec99-c9bf-48e6-b0c6-eb53a928e698.json new file mode 100644 index 00000000..56d5b57c --- /dev/null +++ b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-d2a2ec99-c9bf-48e6-b0c6-eb53a928e698.json @@ -0,0 +1 @@ +{"metrics":{"latencyMs":289},"output":{"message":{"content":[{"text":"Paris"}],"role":"assistant"}},"stopReason":"end_turn","usage":{"inputTokens":12,"outputTokens":2,"serverToolUsage":{},"totalTokens":14}} \ No newline at end of file diff --git a/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-stream-df6a23a1-b3d8-403e-a499-bc5d257af868.txt b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.amazon.nova-lite-v10_converse-stream-df6a23a1-b3d8-403e-a499-bc5d257af868.txt new file mode 100644 index 0000000000000000000000000000000000000000..f0c60b8492035e07b412597e6c5e2da778967649 GIT binary patch literal 4932 zcmd6rU2GIp6o5A-*dmbS$3RUCj)$ZMvM$}0c4eNl0`wJ(N50u4V83epN8##9sFg=ZytE>^mKZPA+CK5Vjo z_s;j7b8_d*@pwEZJswY2&zpaCFVIPjQ2s10WJ#T8VIOfEBT2&Cu=yfA#wefsaj}uj zT4v1PCZlW94x@|oiN8;*o_Bw%eUJqmjMrjul1Qe^bjIqVEc?3ca5vYVFYJAx{lyI% zH@&nuv}NnIj!y7_%`75UIF9LXgK`kV?PCN1Kf+gcudbgjhg>p&)o!s^W&^Rf@t;Xc zcM8!s$%EeHgSgxi^nJido+DYMo4N zk#0~DHXMEa3B`oZSe6!JuFP~*P(d8~ef}y0{XxRs={a*+F|5X!R)cfeqTc@Uj;?TI z=dSLaS6*!jG=r#bNPE%eJL&Yvwc-)QPJt%Pu~@%V_!nXz8V~=n1qBAA#K1Sl<|~%a z63o^L%+;ebN=rpRta;>pNqJh5cCG38N-=4`uWhGQCkTjKd*f9^9F>@lpL=wVHGeZta&6P@t(iRj;*`zWDAn4Wuw~%y5rl5J`vaUP@>-^d# z=GxQd0K~mvZ#^Q9N$hZ{cvd-fD`K0u%J|NB_pdIUL#tM1&}y!n{!!ZI*TQCg+I-Y# zYf8jG6!+KtgifE!Y`N(#ZYj@J5S?VNHl-Dai+3*fBJW&j!J_mzzqZj#rfT;L?}GPJ zr2Sq-cu#)ikn#wx_iLSI^_f3YH~gS%K*Y!X7)G5-Qt51c$6Lyk0!?^`PE#{eH26{p zo4RkkOVH%Fw7HQUY*1`dS@ba{p znz9SilI&?Ml(_$v5)!YB@IzPr-z{^?x>gW|@8&gBydgDz-u%H;)tUw^@RUiDvv+F) z(cCENP)rJTul{iDDb;esa)?h$5Ms?n55)L8*T>Q6OX=ku`1PP_uLQPz@)^!>GXAN| zKz#PpJ&ME!@wMy8;(_Pfs-E`iVzy_<*toA#J1uUd2Gt#vvdE_k!9t?^?igJabp%k z8r;A(V~!liSO(6cVucR;&=zd1o*q;`mE(NLl@l$w+N4>RM_7hXx!woN c9I+#eSvw~d;WaqphV{>~U{Hw2g}1-{3v|bofB*mh literal 0 HcmV?d00001 diff --git a/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.anthropic.claude-3-haiku-20240307-v10_converse-bfd89f01-c330-4c6d-b9c0-05ca56ed005f.json b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.anthropic.claude-3-haiku-20240307-v10_converse-bfd89f01-c330-4c6d-b9c0-05ca56ed005f.json new file mode 100644 index 00000000..ad7ae1fe --- /dev/null +++ b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.anthropic.claude-3-haiku-20240307-v10_converse-bfd89f01-c330-4c6d-b9c0-05ca56ed005f.json @@ -0,0 +1 @@ +{"metrics":{"latencyMs":254},"output":{"message":{"content":[{"text":"Paris."}],"role":"assistant"}},"stopReason":"end_turn","usage":{"inputTokens":19,"outputTokens":5,"serverToolUsage":{},"totalTokens":24}} \ No newline at end of file diff --git a/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.anthropic.claude-3-haiku-20240307-v10_converse-stream-c43d383c-79df-4107-996c-7c63b26994fd.txt b/test-harness/src/testFixtures/resources/cassettes/bedrock/__files/model_us.anthropic.claude-3-haiku-20240307-v10_converse-stream-c43d383c-79df-4107-996c-7c63b26994fd.txt new file mode 100644 index 0000000000000000000000000000000000000000..121ff65992d05a4db1325173e9e02deda31b4c6c GIT binary patch literal 1093 zcmZQzU}$4tU#2`j|1QJupLPoNaI1kFx(ECZ8jG0=xWsfobE z4NMlPc`5NFrA2v4we6fWUl Date: Wed, 15 Apr 2026 13:37:07 -0600 Subject: [PATCH 4/5] bedrock example --- examples/build.gradle | 3 + .../braintrust/examples/SpringAIExample.java | 184 +++++++++--------- 2 files changed, 98 insertions(+), 89 deletions(-) diff --git a/examples/build.gradle b/examples/build.gradle index 8bbe1a4d..7c3860c4 100644 --- a/examples/build.gradle +++ b/examples/build.gradle @@ -36,6 +36,7 @@ dependencies { implementation 'com.google.genai:google-genai:1.20.0' // spring ai examples implementation 'org.springframework.ai:spring-ai-anthropic:1.1.0' + implementation 'org.springframework.ai:spring-ai-bedrock-converse:1.1.0' implementation 'org.springframework.ai:spring-ai-google-genai:1.1.0' implementation 'org.springframework.ai:spring-ai-openai:1.1.0' // spring-ai-openai requires spring-webflux (WebClient) at runtime @@ -47,6 +48,8 @@ dependencies { // to run langchain4j examples implementation 'dev.langchain4j:langchain4j:1.9.1' implementation 'dev.langchain4j:langchain4j-open-ai:1.9.1' + // WireMock for stubbing AWS Bedrock Converse API locally + implementation 'org.wiremock:wiremock:3.12.1' } application { diff --git a/examples/src/main/java/dev/braintrust/examples/SpringAIExample.java b/examples/src/main/java/dev/braintrust/examples/SpringAIExample.java index cbe6d7b8..1004dd2f 100644 --- a/examples/src/main/java/dev/braintrust/examples/SpringAIExample.java +++ b/examples/src/main/java/dev/braintrust/examples/SpringAIExample.java @@ -9,9 +9,13 @@ import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Scope; +import java.util.ArrayList; +import java.util.List; import org.springframework.ai.anthropic.AnthropicChatModel; import org.springframework.ai.anthropic.AnthropicChatOptions; import org.springframework.ai.anthropic.api.AnthropicApi; +import org.springframework.ai.bedrock.converse.BedrockChatOptions; +import org.springframework.ai.bedrock.converse.BedrockProxyChatModel; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.google.genai.GoogleGenAiChatModel; @@ -19,13 +23,13 @@ import org.springframework.ai.openai.OpenAiChatModel; import org.springframework.ai.openai.OpenAiChatOptions; import org.springframework.ai.openai.api.OpenAiApi; -import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.boot.CommandLineRunner; import org.springframework.boot.SpringApplication; import org.springframework.boot.autoconfigure.SpringBootApplication; import org.springframework.boot.autoconfigure.http.client.HttpClientAutoConfiguration; import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration; import org.springframework.context.annotation.Bean; +import software.amazon.awssdk.regions.Region; /** Spring Boot application demonstrating Braintrust + Spring AI integration */ @SpringBootApplication( @@ -41,29 +45,23 @@ public static void main(String[] args) { } @Bean - public CommandLineRunner run( - @Qualifier("openAIChatModel") ChatModel openAIChatModel, - @Qualifier("anthropicChatModel") ChatModel anthropicChatModel, - @Qualifier("googleChatModel") ChatModel googleChatModel, - Tracer tracer, - Braintrust braintrust) { + public CommandLineRunner run(List chatModels, Tracer tracer, Braintrust braintrust) { return args -> { Span rootSpan = tracer.spanBuilder("spring-ai-example").startSpan(); try (Scope scope = rootSpan.makeCurrent()) { System.out.println("\n=== Running Spring Boot Example ===\n"); - // Make a simple chat call var prompt = new Prompt("what's the name of the most popular java DI framework?"); - var oaiResponse = openAIChatModel.call(prompt); - var anthropicResponse = anthropicChatModel.call(prompt); - var googleResponse = googleChatModel.call(prompt); - - System.out.println( - "~~~ SPRING AI CHAT RESPONSES: \noat: %s\nanthropic: %s\ngoogle: %s\n" - .formatted( - oaiResponse.getResult().getOutput().getText(), - anthropicResponse.getResult().getOutput().getText(), - googleResponse.getResult().getOutput().getText())); + + System.out.println("~~~ SPRING AI CHAT RESPONSES:"); + for (var model : chatModels) { + var response = model.call(prompt); + System.out.println( + model.getClass().getSimpleName() + + ": " + + response.getResult().getOutput().getText()); + } + System.out.println(); } finally { rootSpan.end(); } @@ -80,6 +78,85 @@ public CommandLineRunner run( }; } + @Bean + public List chatModels(OpenTelemetry openTelemetry) { + var models = new ArrayList(); + + if (System.getenv("OPENAI_API_KEY") != null) { + models.add( + BraintrustSpringAI.wrap( + openTelemetry, + OpenAiChatModel.builder() + .openAiApi( + OpenAiApi.builder() + .apiKey(System.getenv("OPENAI_API_KEY")) + .build()) + .defaultOptions( + OpenAiChatOptions.builder() + .model("gpt-4o-mini") + .temperature(0.0) + .maxTokens(50) + .build())) + .build()); + } + + if (System.getenv("ANTHROPIC_API_KEY") != null) { + models.add( + BraintrustSpringAI.wrap( + openTelemetry, + AnthropicChatModel.builder() + .anthropicApi( + AnthropicApi.builder() + .apiKey( + System.getenv( + "ANTHROPIC_API_KEY")) + .build()) + .defaultOptions( + AnthropicChatOptions.builder() + .model("claude-3-haiku-20240307") + .temperature(0.0) + .maxTokens(50) + .build())) + .build()); + } + + if (System.getenv("GOOGLE_API_KEY") != null || System.getenv("GEMINI_API_KEY") != null) { + models.add( + GoogleGenAiChatModel.builder() + .genAiClient(BraintrustGenAI.wrap(openTelemetry, new Client.Builder())) + .defaultOptions( + GoogleGenAiChatOptions.builder() + .model("gemini-2.0-flash-lite") + .temperature(0.0) + .maxOutputTokens(50) + .build()) + .build()); + } + + if (System.getenv("AWS_ACCESS_KEY_ID") != null + && System.getenv("AWS_SECRET_ACCESS_KEY") != null) { + models.add( + BedrockProxyChatModel.builder() + .region(Region.US_EAST_1) + .defaultOptions( + BedrockChatOptions.builder() + .model("anthropic.claude-3-haiku-20240307-v1:0") + .temperature(0.0) + .maxTokens(50) + .build()) + .build()); + } + + if (models.isEmpty()) { + System.err.println( + "\nWARNING: No API keys found. Set at least one of: OPENAI_API_KEY," + + " ANTHROPIC_API_KEY, GOOGLE_API_KEY/GEMINI_API_KEY, or" + + " AWS_ACCESS_KEY_ID+AWS_SECRET_ACCESS_KEY\n"); + } + + return models; + } + @Bean public Braintrust braintrust() { return Braintrust.get(BraintrustConfig.fromEnvironment()); @@ -94,75 +171,4 @@ public OpenTelemetry openTelemetry(Braintrust braintrust) { public Tracer tracer(OpenTelemetry openTelemetry) { return openTelemetry.getTracer("spring-ai-instrumentation"); } - - @Bean - public ChatModel openAIChatModel(OpenTelemetry openTelemetry) { - if (null == System.getenv("OPENAI_API_KEY")) { - System.err.println( - "\n" - + "WARNING: OPENAI_API_KEY not found. This example will likely" - + " fail.\n" - + "Set it with: export OPENAI_API_KEY='your-key'\n"); - } - return BraintrustSpringAI.wrap( - openTelemetry, - OpenAiChatModel.builder() - .openAiApi( - OpenAiApi.builder() - .apiKey(System.getenv("OPENAI_API_KEY")) - .build()) - .defaultOptions( - OpenAiChatOptions.builder() - .model("gpt-4o-mini") - .temperature(0.0) - .maxTokens(50) - .build())) - .build(); - } - - @Bean - public ChatModel anthropicChatModel(OpenTelemetry openTelemetry) { - if (null == System.getenv("ANTHROPIC_API_KEY")) { - System.err.println( - "\n" - + "WARNING: ANTHROPIC_API_KEY not found. This example will" - + " likely fail.\n" - + "Set it with: export ANTHROPIC_API_KEY='your-key'\n"); - } - return BraintrustSpringAI.wrap( - openTelemetry, - AnthropicChatModel.builder() - .anthropicApi( - AnthropicApi.builder() - .apiKey(System.getenv("ANTHROPIC_API_KEY")) - .build()) - .defaultOptions( - AnthropicChatOptions.builder() - .model("claude-3-haiku-20240307") - .temperature(0.0) - .maxTokens(50) - .build())) - .build(); - } - - @Bean - public ChatModel googleChatModel(OpenTelemetry openTelemetry) { - if (null == System.getenv("GOOGLE_API_KEY") && null == System.getenv("GEMINI_API_KEY")) { - System.err.println( - "\n" - + "WARNING: Neither GOOGLE_API_KEY nor GEMINI_API_KEY found. This" - + " example will likely fail.\n" - + "Set either: export GOOGLE_API_KEY='your-key' (recommended) or" - + " export GEMINI_API_KEY='your-key'\n"); - } - return GoogleGenAiChatModel.builder() - .genAiClient(BraintrustGenAI.wrap(openTelemetry, new Client.Builder())) - .defaultOptions( - GoogleGenAiChatOptions.builder() - .model("gemini-2.0-flash-lite") - .temperature(0.0) - .maxOutputTokens(50) - .build()) - .build(); - } } From c3bfb01a1b2adcff598f2fd90a1fb3afbed277b0 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Thu, 16 Apr 2026 14:28:23 -0600 Subject: [PATCH 5/5] bedrock instrumentation and spec test executor --- .../aws_bedrock_2_30_0/build.gradle | 62 +++ .../v2_30_0/BraintrustAWSBedrock.java | 71 +++ .../v2_30_0/BraintrustBedrockInterceptor.java | 412 ++++++++++++++++++ .../auto/AWSBedrockInstrumentationModule.java | 83 ++++ .../v2_30_0/BraintrustAWSBedrockTest.java | 152 +++++++ .../InstrumentationSemConv.java | 178 ++++++++ btx/build.gradle | 5 + .../braintrust/sdkspecimpl/SpanConverter.java | 25 ++ .../braintrust/sdkspecimpl/SpecExecutor.java | 99 +++++ .../braintrust/sdkspecimpl/SpecLoader.java | 3 +- examples/build.gradle | 4 +- .../braintrust/examples/SpringAIExample.java | 9 +- gradle.properties | 3 +- settings.gradle | 1 + test-harness/build.gradle | 5 + .../dev/braintrust/Bedrock30TestUtils.java | 104 +++++ .../java/dev/braintrust/TestHarness.java | 19 +- .../testFixtures/java/dev/braintrust/VCR.java | 47 +- 18 files changed, 1258 insertions(+), 24 deletions(-) create mode 100644 braintrust-sdk/instrumentation/aws_bedrock_2_30_0/build.gradle create mode 100644 braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustAWSBedrock.java create mode 100644 braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustBedrockInterceptor.java create mode 100644 braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/auto/AWSBedrockInstrumentationModule.java create mode 100644 braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/test/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustAWSBedrockTest.java create mode 100644 test-harness/src/testFixtures/java/dev/braintrust/Bedrock30TestUtils.java diff --git a/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/build.gradle b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/build.gradle new file mode 100644 index 00000000..a58752b1 --- /dev/null +++ b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/build.gradle @@ -0,0 +1,62 @@ +def awsBedrockVersion = '2.30.0' + +muzzle { + pass { + group = 'software.amazon.awssdk' + module = 'bedrockruntime' + // num of aws patch versions is huge so list out minor versions instead + pinVersions '2.30.0', '2.31.0', '2.32.0', '2.33.0', '2.34.0', + '2.35.0', '2.36.0', '2.37.0', '2.38.0', '2.39.0', + '2.40.0', '2.41.0' + // sdk-core is listed explicitly because SdkDefaultClientBuilder (the ByteBuddy + // interception target) lives there, not in bedrockruntime itself + extraDependency 'software.amazon.awssdk:sdk-core' + extraDependency "com.fasterxml.jackson.core:jackson-databind:${jacksonVersion}" + extraDependency "com.fasterxml.jackson.datatype:jackson-datatype-jsr310:${jacksonVersion}" + extraDependency "com.fasterxml.jackson.datatype:jackson-datatype-jdk8:${jacksonVersion}" + } + pass { + group = 'software.amazon.awssdk' + module = 'bedrockruntime' + versions = '[2.42.36,)' // latest version and up + extraDependency 'software.amazon.awssdk:sdk-core' + extraDependency "com.fasterxml.jackson.core:jackson-databind:${jacksonVersion}" + extraDependency "com.fasterxml.jackson.datatype:jackson-datatype-jsr310:${jacksonVersion}" + extraDependency "com.fasterxml.jackson.datatype:jackson-datatype-jdk8:${jacksonVersion}" + } +} + +dependencies { + compileOnly project(':braintrust-java-agent:instrumenter') + implementation "io.opentelemetry:opentelemetry-api:${otelVersion}" + implementation 'com.google.code.findbugs:jsr305:3.0.2' + implementation "org.slf4j:slf4j-api:${slf4jVersion}" + implementation project(':braintrust-sdk') + + // ByteBuddy for ElementMatcher types used in instrumentation definitions + compileOnly 'net.bytebuddy:byte-buddy:1.17.5' + + // Target library — compileOnly because it will be on the app classpath at runtime + compileOnly "software.amazon.awssdk:bedrockruntime:${awsBedrockVersion}" + + // Test dependencies + testImplementation(testFixtures(project(":test-harness"))) + testImplementation project(':braintrust-java-agent:instrumenter') + testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}" + testRuntimeOnly 'org.junit.platform:junit-platform-launcher' + testImplementation 'net.bytebuddy:byte-buddy-agent:1.17.5' + testRuntimeOnly "org.slf4j:slf4j-simple:${slf4jVersion}" + testImplementation "software.amazon.awssdk:bedrockruntime:${awsBedrockVersion}" + testImplementation "software.amazon.awssdk:netty-nio-client:${awsBedrockVersion}" + testImplementation "com.fasterxml.jackson.core:jackson-databind:${jacksonVersion}" +} + +test { + useJUnitPlatform() + workingDir = rootProject.projectDir + testLogging { + events "passed", "skipped", "failed" + showStandardStreams = true + exceptionFormat "full" + } +} diff --git a/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustAWSBedrock.java b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustAWSBedrock.java new file mode 100644 index 00000000..7d5716fa --- /dev/null +++ b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustAWSBedrock.java @@ -0,0 +1,71 @@ +package dev.braintrust.instrumentation.awsbedrock.v2_30_0; + +import io.opentelemetry.api.OpenTelemetry; +import lombok.extern.slf4j.Slf4j; +import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClientBuilder; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClientBuilder; + +/** Braintrust instrumentation for the AWS Bedrock Runtime client. */ +@Slf4j +public class BraintrustAWSBedrock { + + /** + * Wraps a {@link BedrockRuntimeClientBuilder} so that every {@code converse} call made through + * the resulting client is traced via OpenTelemetry. + * + *

Call this method after applying all custom builder settings + * + * @param openTelemetry the OpenTelemetry instance to use for tracing + * @param builder the client builder to instrument + * @return the same builder (for fluent chaining) + */ + public static BedrockRuntimeClientBuilder wrap( + OpenTelemetry openTelemetry, BedrockRuntimeClientBuilder builder) { + try { + // Read existing config so we don't clobber user-registered interceptors + ClientOverrideConfiguration existing = builder.overrideConfiguration(); + ClientOverrideConfiguration.Builder configBuilder = existing.toBuilder(); + for (var interceptor : configBuilder.executionInterceptors()) { + if (interceptor instanceof BraintrustBedrockInterceptor) { + log.info("builder already wrapped. Skipping"); + return builder; + } + } + configBuilder.addExecutionInterceptor(new BraintrustBedrockInterceptor(openTelemetry)); + builder.overrideConfiguration(configBuilder.build()); + } catch (Exception e) { + log.warn("Failed to apply Bedrock instrumentation", e); + } + return builder; + } + + /** + * Wraps a {@link BedrockRuntimeAsyncClientBuilder} so that every {@code converseStream} call + * made through the resulting client is traced via OpenTelemetry. + * + *

Call this method after applying all custom builder settings. + * + * @param openTelemetry the OpenTelemetry instance to use for tracing + * @param builder the async client builder to instrument + * @return the same builder (for fluent chaining) + */ + public static BedrockRuntimeAsyncClientBuilder wrap( + OpenTelemetry openTelemetry, BedrockRuntimeAsyncClientBuilder builder) { + try { + ClientOverrideConfiguration existing = builder.overrideConfiguration(); + ClientOverrideConfiguration.Builder configBuilder = existing.toBuilder(); + for (var interceptor : configBuilder.executionInterceptors()) { + if (interceptor instanceof BraintrustBedrockInterceptor) { + log.info("async builder already wrapped. Skipping"); + return builder; + } + } + configBuilder.addExecutionInterceptor(new BraintrustBedrockInterceptor(openTelemetry)); + builder.overrideConfiguration(configBuilder.build()); + } catch (Exception e) { + log.warn("Failed to apply async Bedrock instrumentation", e); + } + return builder; + } +} diff --git a/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustBedrockInterceptor.java b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustBedrockInterceptor.java new file mode 100644 index 00000000..6f3e0d58 --- /dev/null +++ b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustBedrockInterceptor.java @@ -0,0 +1,412 @@ +package dev.braintrust.instrumentation.awsbedrock.v2_30_0; + +import dev.braintrust.instrumentation.InstrumentationSemConv; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.StringWriter; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import lombok.extern.slf4j.Slf4j; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; +import software.amazon.awssdk.core.SdkRequest; +import software.amazon.awssdk.core.interceptor.Context.BeforeExecution; +import software.amazon.awssdk.core.interceptor.Context.ModifyHttpRequest; +import software.amazon.awssdk.core.interceptor.Context.ModifyHttpResponse; +import software.amazon.awssdk.core.interceptor.ExecutionAttribute; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute; +import software.amazon.awssdk.http.SdkHttpRequest; +import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest; +import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest; +import software.amazon.awssdk.thirdparty.jackson.core.JsonFactory; +import software.amazon.awssdk.thirdparty.jackson.core.JsonParser; +import software.amazon.awssdk.thirdparty.jackson.core.JsonToken; +import software.amazon.eventstream.Message; +import software.amazon.eventstream.MessageDecoder; + +/** + * AWS SDK ExecutionInterceptor that creates OpenTelemetry spans for Bedrock Converse calls, + * capturing the raw request and response bodies via {@link InstrumentationSemConv}. + */ +@Slf4j +class BraintrustBedrockInterceptor implements ExecutionInterceptor { + private static final String INSTRUMENTATION_NAME = "braintrust-aws-bedrock"; + + private static final ExecutionAttribute SPAN_ATTRIBUTE = + new ExecutionAttribute<>("braintrust.span"); + private static final ExecutionAttribute MODEL_ID_ATTRIBUTE = + new ExecutionAttribute<>("braintrust.modelId"); + + private static final JsonFactory JSON_FACTORY = new JsonFactory(); + + private final Tracer tracer; + + BraintrustBedrockInterceptor(OpenTelemetry openTelemetry) { + this.tracer = openTelemetry.getTracer(INSTRUMENTATION_NAME); + } + + private static final Set INSTRUMENTED_OPERATIONS = Set.of("Converse", "ConverseStream"); + + @Override + public void beforeExecution(BeforeExecution context, ExecutionAttributes executionAttributes) { + String operationName = + executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME); + + // Only instrument Converse and ConverseStream — other Bedrock operations + // (InvokeModel, ApplyGuardrail, etc.) are not LLM calls we know how to tag. + if (!INSTRUMENTED_OPERATIONS.contains(operationName)) { + return; + } + + SdkRequest sdkRequest = context.request(); + String modelId = extractModelId(sdkRequest); + + Span span = tracer.spanBuilder(operationName).setParent(Context.current()).startSpan(); + executionAttributes.putAttribute(SPAN_ATTRIBUTE, span); + if (modelId != null) { + executionAttributes.putAttribute(MODEL_ID_ATTRIBUTE, modelId); + } + } + + @Override + public SdkHttpRequest modifyHttpRequest( + ModifyHttpRequest context, ExecutionAttributes executionAttributes) { + Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + if (span == null) { + return context.httpRequest(); + } + + SdkHttpRequest httpRequest = context.httpRequest(); + String modelId = executionAttributes.getAttribute(MODEL_ID_ATTRIBUTE); + + if (modelId == null) { + modelId = extractModelIdFromPath(httpRequest.encodedPath()); + if (modelId != null) { + executionAttributes.putAttribute(MODEL_ID_ATTRIBUTE, modelId); + } + } + + String requestBody = null; + if (context.requestBody().isPresent()) { + try (InputStream is = context.requestBody().get().contentStreamProvider().newStream()) { + requestBody = new String(is.readAllBytes(), StandardCharsets.UTF_8); + } catch (Exception e) { + log.debug("Failed to capture request body", e); + } + } + + String baseUrl = httpRequest.protocol() + "://" + httpRequest.host(); + List pathSegments = + Arrays.stream(httpRequest.encodedPath().split("/")) + .filter(s -> !s.isEmpty()) + .toList(); + + InstrumentationSemConv.tagLLMSpanRequest( + span, + InstrumentationSemConv.PROVIDER_NAME_BEDROCK, + baseUrl, + pathSegments, + "POST", + requestBody, + modelId); + + return httpRequest; + } + + @Override + public Optional modifyHttpResponseContent( + ModifyHttpResponse context, ExecutionAttributes executionAttributes) { + Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + if (span == null) { + return context.responseBody(); + } + + // Only intercept successful responses. On 4xx/5xx the SDK needs to read the body + // itself to parse the AWS error message — consuming it here would swallow that. + if (context.httpResponse().statusCode() >= 300) { + return context.responseBody(); + } + + Optional body = context.responseBody(); + if (body.isPresent()) { + try { + final byte[] bytes = body.get().readAllBytes(); + try { + String responseBodyStr = new String(bytes, StandardCharsets.UTF_8); + InstrumentationSemConv.tagLLMSpanResponse( + span, InstrumentationSemConv.PROVIDER_NAME_BEDROCK, responseBodyStr); + } catch (Exception e) { + log.debug("Failed to capture response body", e); + } + return Optional.of(new ByteArrayInputStream(bytes)); + } catch (IOException e) { + // unlikely this will happen, but if we get here there's no sensible recovery + throw new RuntimeException("failed to ready response body bytes", e); + } + } + return body; + } + + /** + * Intercepts the async response stream for {@code converseStream} calls. Tees the reactive + * {@link Publisher} so that bytes are fed to a {@link MessageDecoder} as they arrive, and on + * completion the decoded event-stream frames are used to tag the span. + */ + @Override + public Optional> modifyAsyncHttpResponseContent( + software.amazon.awssdk.core.interceptor.Context.ModifyHttpResponse context, + ExecutionAttributes executionAttributes) { + Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + Optional> publisherOpt = context.responsePublisher(); + if (span == null || publisherOpt.isEmpty()) { + return publisherOpt; + } + + // Only intercept successful responses — error responses must flow through untouched + // so the SDK can parse the AWS error body. + if (context.httpResponse().statusCode() >= 300) { + return publisherOpt; + } + + Publisher original = publisherOpt.get(); + Publisher teed = + subscriber -> original.subscribe(new TeeingSubscriber(subscriber, span)); + return Optional.of(teed); + } + + @Override + public void afterExecution( + software.amazon.awssdk.core.interceptor.Context.AfterExecution context, + ExecutionAttributes executionAttributes) { + endSpan(executionAttributes, null); + } + + @Override + public void onExecutionFailure( + software.amazon.awssdk.core.interceptor.Context.FailedExecution context, + ExecutionAttributes executionAttributes) { + endSpan(executionAttributes, context.exception()); + } + + private static void endSpan( + ExecutionAttributes executionAttributes, @javax.annotation.Nullable Throwable error) { + Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + if (span == null) { + return; + } + if (error != null) { + InstrumentationSemConv.tagLLMSpanResponse(span, error); + } + span.end(); + } + + private static String extractModelId(SdkRequest request) { + if (request instanceof ConverseRequest r) return r.modelId(); + if (request instanceof ConverseStreamRequest r) return r.modelId(); + return null; + } + + private static String extractModelIdFromPath(String path) { + if (path != null && path.startsWith("/model/")) { + int start = "/model/".length(); + int end = path.indexOf("/", start); + if (end > start) { + return java.net.URLDecoder.decode( + path.substring(start, end), StandardCharsets.UTF_8); + } + } + return null; + } + + /** + * Tees the reactive byte stream into a {@link MessageDecoder}. On completion, decodes each + * event-stream frame using the AWS SDK's shaded Jackson streaming parser, accumulates the + * response content, and hands a synthetic Converse-shaped JSON string to semconv. + */ + private static class TeeingSubscriber implements Subscriber { + private final Subscriber downstream; + private final Span span; + private final MessageDecoder decoder = new MessageDecoder(); + + // Accumulated incrementally in onNext — no message list retained. + private final StringBuilder text = new StringBuilder(); + private String stopReason = null; + private int inputTokens = 0; + private int outputTokens = 0; + private long startNanos; + private Long timeToFirstTokenNanos = null; + + TeeingSubscriber(Subscriber downstream, Span span) { + this.downstream = downstream; + this.span = span; + } + + @Override + public void onSubscribe(Subscription s) { + startNanos = System.nanoTime(); + downstream.onSubscribe(s); + } + + @Override + public void onNext(ByteBuffer buf) { + byte[] copy = new byte[buf.remaining()]; + buf.duplicate().get(copy); + try { + decoder.feed(copy); + for (Message msg : decoder.getDecodedMessages()) { + var h = msg.getHeaders().get(":event-type"); + if (h == null) continue; + String eventType = h.getString(); + byte[] payload = msg.getPayload(); + switch (eventType) { + case "contentBlockDelta" -> { + String t = parseDeltaText(payload); + if (t != null) { + text.append(t); + if (timeToFirstTokenNanos == null) { + timeToFirstTokenNanos = System.nanoTime() - startNanos; + } + } + } + case "messageStop" -> stopReason = parseStopReason(payload); + case "metadata" -> { + int[] tokens = parseTokenUsage(payload); + inputTokens = tokens[0]; + outputTokens = tokens[1]; + } + default -> {} + } + } + } catch (Exception e) { + log.debug("Failed to feed event-stream decoder", e); + } + downstream.onNext(buf); + } + + @Override + public void onError(Throwable t) { + downstream.onError(t); + } + + @Override + public void onComplete() { + try { + InstrumentationSemConv.tagLLMSpanResponse( + span, + InstrumentationSemConv.PROVIDER_NAME_BEDROCK, + buildConverseJson(text.toString(), stopReason, inputTokens, outputTokens), + timeToFirstTokenNanos); + } catch (Exception e) { + log.debug("Failed to tag span from streaming response", e); + } finally { + downstream.onComplete(); + } + } + + /** + * Parses {@code delta.text} from a {@code contentBlockDelta} payload: {@code + * {"contentBlockIndex":0,"delta":{"text":"...","type":"text_delta"}}} + */ + private static String parseDeltaText(byte[] payload) throws Exception { + try (JsonParser p = JSON_FACTORY.createParser(payload)) { + boolean inDelta = false; + while (p.nextToken() != null) { + if (p.currentToken() == JsonToken.FIELD_NAME) { + if ("delta".equals(p.currentName())) { + inDelta = true; + } else if (inDelta && "text".equals(p.currentName())) { + p.nextToken(); + return p.getText(); + } + } else if (p.currentToken() == JsonToken.END_OBJECT) { + inDelta = false; + } + } + } + return null; + } + + /** + * Parses {@code stopReason} from a {@code messageStop} payload: {@code + * {"stopReason":"end_turn"}} + */ + private static String parseStopReason(byte[] payload) throws Exception { + try (JsonParser p = JSON_FACTORY.createParser(payload)) { + while (p.nextToken() != null) { + if (p.currentToken() == JsonToken.FIELD_NAME + && "stopReason".equals(p.currentName())) { + p.nextToken(); + return p.getText(); + } + } + } + return null; + } + + /** + * Parses {@code [inputTokens, outputTokens]} from a {@code metadata} payload: {@code + * {"usage":{"inputTokens":N,"outputTokens":M},"metrics":{...}}} + */ + private static int[] parseTokenUsage(byte[] payload) throws Exception { + int inputTokens = 0; + int outputTokens = 0; + try (JsonParser p = JSON_FACTORY.createParser(payload)) { + while (p.nextToken() != null) { + if (p.currentToken() == JsonToken.FIELD_NAME) { + if ("inputTokens".equals(p.currentName())) { + p.nextToken(); + inputTokens = p.getIntValue(); + } else if ("outputTokens".equals(p.currentName())) { + p.nextToken(); + outputTokens = p.getIntValue(); + } + } + } + } + return new int[] {inputTokens, outputTokens}; + } + + /** + * Builds a synthetic Converse-shaped JSON string matching what {@code tagBedrockResponse} + * expects, using the shaded Jackson generator for correct escaping. + */ + private static String buildConverseJson( + String text, String stopReason, int inputTokens, int outputTokens) + throws Exception { + StringWriter sw = new StringWriter(); + try (var gen = JSON_FACTORY.createGenerator(sw)) { + gen.writeStartObject(); + gen.writeObjectFieldStart("output"); + gen.writeObjectFieldStart("message"); + gen.writeStringField("role", "assistant"); + gen.writeArrayFieldStart("content"); + gen.writeStartObject(); + gen.writeStringField("text", text); + gen.writeEndObject(); + gen.writeEndArray(); + gen.writeEndObject(); // message + gen.writeEndObject(); // output + gen.writeStringField("stopReason", stopReason != null ? stopReason : "end_turn"); + gen.writeObjectFieldStart("usage"); + gen.writeNumberField("inputTokens", inputTokens); + gen.writeNumberField("outputTokens", outputTokens); + gen.writeNumberField("totalTokens", inputTokens + outputTokens); + gen.writeEndObject(); // usage + gen.writeEndObject(); + } + return sw.toString(); + } + } +} diff --git a/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/auto/AWSBedrockInstrumentationModule.java b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/auto/AWSBedrockInstrumentationModule.java new file mode 100644 index 00000000..3b32edc8 --- /dev/null +++ b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/auto/AWSBedrockInstrumentationModule.java @@ -0,0 +1,83 @@ +package dev.braintrust.instrumentation.awsbedrock.v2_30_0.auto; + +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import com.google.auto.service.AutoService; +import dev.braintrust.instrumentation.InstrumentationModule; +import dev.braintrust.instrumentation.TypeInstrumentation; +import dev.braintrust.instrumentation.TypeTransformer; +import dev.braintrust.instrumentation.awsbedrock.v2_30_0.BraintrustAWSBedrock; +import io.opentelemetry.api.GlobalOpenTelemetry; +import java.util.List; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClientBuilder; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClientBuilder; + +/** + * Auto-instruments the AWS Bedrock Runtime sync and async client builders by hooking into {@code + * SdkDefaultClientBuilder.build()} — the single {@code final} method in the AWS SDK builder + * hierarchy that both {@link BedrockRuntimeClientBuilder} and {@link + * BedrockRuntimeAsyncClientBuilder} ultimately call. The advice checks the runtime type of {@code + * this} to limit instrumentation to Bedrock builders only. + */ +@AutoService(InstrumentationModule.class) +public class AWSBedrockInstrumentationModule extends InstrumentationModule { + private static final String MANUAL_INSTRUMENTATION_PACKAGE = + "dev.braintrust.instrumentation.awsbedrock.v2_30_0."; + + public AWSBedrockInstrumentationModule() { + super("aws_bedrock_2_30_0"); + } + + @Override + public List getHelperClassNames() { + return List.of( + MANUAL_INSTRUMENTATION_PACKAGE + "BraintrustAWSBedrock", + MANUAL_INSTRUMENTATION_PACKAGE + "BraintrustBedrockInterceptor", + MANUAL_INSTRUMENTATION_PACKAGE + "BraintrustBedrockInterceptor$TeeingSubscriber", + "dev.braintrust.json.BraintrustJsonMapper", + "dev.braintrust.instrumentation.InstrumentationSemConv"); + } + + @Override + public List typeInstrumentations() { + return List.of(new SdkDefaultClientBuilderInstrumentation()); + } + + /** + * Targets {@code SdkDefaultClientBuilder} — the abstract base that defines the {@code final + * build()} method inherited by all AWS SDK client builders, including both Bedrock variants. + */ + public static class SdkDefaultClientBuilderInstrumentation implements TypeInstrumentation { + @Override + public ElementMatcher typeMatcher() { + return named("software.amazon.awssdk.core.client.builder.SdkDefaultClientBuilder"); + } + + @Override + public void transform(TypeTransformer transformer) { + transformer.applyAdviceToMethod( + named("build").and(takesArguments(0)), + AWSBedrockInstrumentationModule.class.getName() + "$BedrockBuilderAdvice"); + } + } + + /** + * Fires on entry to {@code build()} for any AWS SDK client builder. Uses {@code instanceof} + * checks to limit actual work to Bedrock builders, then calls the idempotent {@code wrap()} + * method to register the Braintrust {@code ExecutionInterceptor} before the client is built. + */ + public static class BedrockBuilderAdvice { + @Advice.OnMethodEnter + public static void build(@Advice.This Object builder) { + if (builder instanceof BedrockRuntimeClientBuilder bedrockBuilder) { + BraintrustAWSBedrock.wrap(GlobalOpenTelemetry.get(), bedrockBuilder); + } else if (builder instanceof BedrockRuntimeAsyncClientBuilder bedrockBuilder) { + BraintrustAWSBedrock.wrap(GlobalOpenTelemetry.get(), bedrockBuilder); + } + } + } +} diff --git a/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/test/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustAWSBedrockTest.java b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/test/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustAWSBedrockTest.java new file mode 100644 index 00000000..4876436a --- /dev/null +++ b/braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/test/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustAWSBedrockTest.java @@ -0,0 +1,152 @@ +package dev.braintrust.instrumentation.awsbedrock.v2_30_0; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import dev.braintrust.Bedrock30TestUtils; +import dev.braintrust.TestHarness; +import dev.braintrust.instrumentation.Instrumenter; +import io.opentelemetry.api.common.AttributeKey; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Stream; +import lombok.SneakyThrows; +import net.bytebuddy.agent.ByteBuddyAgent; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock; +import software.amazon.awssdk.services.bedrockruntime.model.ConversationRole; +import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest; +import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest; +import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler; +import software.amazon.awssdk.services.bedrockruntime.model.Message; + +public class BraintrustAWSBedrockTest { + + private static final ObjectMapper JSON_MAPPER = new ObjectMapper(); + + @BeforeAll + static void beforeAll() { + var instrumentation = ByteBuddyAgent.install(); + Instrumenter.install(instrumentation, BraintrustAWSBedrockTest.class.getClassLoader()); + } + + private TestHarness testHarness; + private Bedrock30TestUtils bedrockUtils; + + @BeforeEach + void beforeEach() { + testHarness = TestHarness.setup(); + bedrockUtils = new Bedrock30TestUtils(testHarness); + } + + static Stream modelProvider() { + return Stream.of( + Arguments.of("us.anthropic.claude-3-haiku-20240307-v1:0"), + Arguments.of("us.amazon.nova-lite-v1:0")); + } + + @ParameterizedTest(name = "converse with {0}") + @MethodSource("modelProvider") + @SneakyThrows + void converseProducesLlmSpan(String modelId) { + try (var client = bedrockUtils.syncClientBuilder().build()) { + var response = + client.converse( + ConverseRequest.builder() + .modelId(modelId) + .messages( + Message.builder() + .role(ConversationRole.USER) + .content( + ContentBlock.fromText( + "What is the capital of France?" + + " Reply in one word.")) + .build()) + .build()); + assertNotNull(response); + assertFalse( + response.output().message().content().isEmpty(), + "response should have content"); + + var spans = testHarness.awaitExportedSpans(1); + assertEquals(1, spans.size(), "expected exactly one span"); + var span = spans.get(0); + + String spanAttributesJson = + span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); + assertNotNull(spanAttributesJson, "braintrust.span_attributes should be set"); + JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); + assertEquals("llm", spanAttributes.get("type").asText()); + assertNotNull( + span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")), + "braintrust.input_json should be set"); + assertNotNull( + span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")), + "braintrust.output_json should be set"); + } + } + + @Test + @SneakyThrows + void converseStreamProducesLlmSpan() { + String modelId = "us.anthropic.claude-3-haiku-20240307-v1:0"; + + try (var client = bedrockUtils.asyncClientBuilder().build()) { + var accumulatedText = new AtomicReference<>(new StringBuilder()); + + var responseHandler = + ConverseStreamResponseHandler.builder() + .subscriber( + ConverseStreamResponseHandler.Visitor.builder() + .onContentBlockDelta( + evt -> { + if (evt.delta().text() != null) { + accumulatedText + .get() + .append(evt.delta().text()); + } + }) + .build()) + .build(); + + client.converseStream( + ConverseStreamRequest.builder() + .modelId(modelId) + .messages( + Message.builder() + .role(ConversationRole.USER) + .content( + ContentBlock.fromText( + "What is the capital of France?" + + " Reply in one word.")) + .build()) + .build(), + responseHandler) + .get(); + + assertFalse( + accumulatedText.get().isEmpty(), "should have received streamed text content"); + + var spans = testHarness.awaitExportedSpans(1); + assertEquals(1, spans.size(), "expected exactly one span"); + var span = spans.get(0); + + String spanAttributesJson = + span.getAttributes().get(AttributeKey.stringKey("braintrust.span_attributes")); + assertNotNull(spanAttributesJson, "braintrust.span_attributes should be set"); + JsonNode spanAttributes = JSON_MAPPER.readTree(spanAttributesJson); + assertEquals("llm", spanAttributes.get("type").asText()); + assertNotNull( + span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")), + "braintrust.input_json should be set"); + assertNotNull( + span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")), + "braintrust.output_json should be set"); + } + } +} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/InstrumentationSemConv.java b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/InstrumentationSemConv.java index 6f47587c..05c04afd 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/InstrumentationSemConv.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/instrumentation/InstrumentationSemConv.java @@ -17,6 +17,7 @@ public class InstrumentationSemConv { public static final String PROVIDER_NAME_OPENAI = "openai"; public static final String PROVIDER_NAME_ANTHROPIC = "anthropic"; + public static final String PROVIDER_NAME_BEDROCK = "bedrock"; public static final String PROVIDER_NAME_OTHER = "generic-ai-provider"; public static final String UNSET_LLM_SPAN_NAME = "llm"; @@ -32,6 +33,25 @@ public static void tagLLMSpanRequest( @Nonnull List pathSegments, @Nonnull String method, @Nullable String requestBody) { + tagLLMSpanRequest(span, providerName, baseUrl, pathSegments, method, requestBody, null); + } + + /** + * Tag a span with LLM request metadata. + * + * @param modelId explicit model identifier — used by providers (e.g. Bedrock) where the model + * is not present in the request body. When {@code null} the model is extracted from the + * request body if possible. + */ + @SneakyThrows + public static void tagLLMSpanRequest( + Span span, + @Nonnull String providerName, + @Nonnull String baseUrl, + @Nonnull List pathSegments, + @Nonnull String method, + @Nullable String requestBody, + @Nullable String modelId) { switch (providerName) { case PROVIDER_NAME_OPENAI -> tagOpenAIRequest( @@ -39,6 +59,15 @@ public static void tagLLMSpanRequest( case PROVIDER_NAME_ANTHROPIC -> tagAnthropicRequest( span, providerName, baseUrl, pathSegments, method, requestBody); + case PROVIDER_NAME_BEDROCK -> + tagBedrockRequest( + span, + providerName, + baseUrl, + pathSegments, + method, + requestBody, + modelId); default -> tagOpenAIRequest( span, providerName, baseUrl, pathSegments, method, requestBody); @@ -61,6 +90,8 @@ public static void tagLLMSpanResponse( tagOpenAIResponse(span, responseBody, timeToFirstTokenNanoseconds); case PROVIDER_NAME_ANTHROPIC -> tagAnthropicResponse(span, responseBody, timeToFirstTokenNanoseconds); + case PROVIDER_NAME_BEDROCK -> + tagBedrockResponse(span, responseBody, timeToFirstTokenNanoseconds); default -> tagOpenAIResponse(span, responseBody, timeToFirstTokenNanoseconds); } } @@ -234,6 +265,102 @@ private static void tagAnthropicResponse( } } + // ------------------------------------------------------------------------- + // AWS Bedrock provider implementation + // ------------------------------------------------------------------------- + + @SneakyThrows + private static void tagBedrockRequest( + Span span, + String providerName, + String baseUrl, + List pathSegments, + String method, + @Nullable String requestBody, + @Nullable String modelId) { + String endpoint = bedrockEndpoint(pathSegments); + span.updateName("bedrock." + endpoint); + span.setAttribute("braintrust.span_attributes", toJson(Map.of("type", "llm"))); + + Map metadata = new HashMap<>(); + metadata.put("provider", "bedrock"); + metadata.put("endpoint", endpoint); + metadata.put("request_path", String.join("/", pathSegments)); + metadata.put("request_base_uri", baseUrl); + metadata.put("request_method", method); + + if (modelId != null) { + metadata.put("model", modelId); + } + + if (requestBody != null) { + JsonNode requestJson = BraintrustJsonMapper.get().readTree(requestBody); + // Extract inference parameters from inferenceConfig + if (requestJson.has("inferenceConfig")) { + JsonNode cfg = requestJson.get("inferenceConfig"); + if (cfg.has("maxTokens")) metadata.put("max_tokens", cfg.get("maxTokens")); + if (cfg.has("temperature")) metadata.put("temperature", cfg.get("temperature")); + if (cfg.has("topP")) metadata.put("top_p", cfg.get("topP")); + if (cfg.has("stopSequences")) + metadata.put("stop_sequences", cfg.get("stopSequences")); + } + // Bedrock Converse uses "messages" with typed content block arrays like + // [{"text":"..."}] + if (requestJson.has("messages")) { + ArrayNode inputArray = BraintrustJsonMapper.get().createArrayNode(); + // Bedrock puts system prompts in a separate top-level "system" array: + // [{"text": "..."}]. Prepend as a synthetic {role:"system", content:[...]} entry. + if (requestJson.has("system") + && requestJson.get("system").isArray() + && !requestJson.get("system").isEmpty()) { + var systemNode = BraintrustJsonMapper.get().createObjectNode(); + systemNode.put("role", "system"); + systemNode.set("content", requestJson.get("system")); + inputArray.add(systemNode); + } + for (JsonNode msg : requestJson.get("messages")) { + inputArray.add(normalizeBedrockMessage(msg)); + } + span.setAttribute("braintrust.input_json", toJson(inputArray)); + } + } + + span.setAttribute("braintrust.metadata", toJson(metadata)); + } + + @SneakyThrows + private static void tagBedrockResponse( + Span span, String responseBody, @Nullable Long timeToFirstTokenNanoseconds) { + JsonNode responseJson = BraintrustJsonMapper.get().readTree(responseBody); + + // Bedrock output lives at output.message. Normalize to a single-element array matching the + // same [{role, content: [...]}] shape as input so the UI can render the LLM thread view. + if (responseJson.has("output") && responseJson.get("output").has("message")) { + JsonNode message = responseJson.get("output").get("message"); + ArrayNode outputArray = BraintrustJsonMapper.get().createArrayNode(); + outputArray.add(normalizeBedrockMessage(message)); + span.setAttribute("braintrust.output_json", toJson(outputArray)); + } + + Map metrics = new HashMap<>(); + if (timeToFirstTokenNanoseconds != null) { + metrics.put("time_to_first_token", timeToFirstTokenNanoseconds / 1_000_000_000.0); + } + + // Bedrock usage uses camelCase: inputTokens, outputTokens, totalTokens + if (responseJson.has("usage")) { + JsonNode usage = responseJson.get("usage"); + if (usage.has("inputTokens")) metrics.put("prompt_tokens", usage.get("inputTokens")); + if (usage.has("outputTokens")) + metrics.put("completion_tokens", usage.get("outputTokens")); + if (usage.has("totalTokens")) metrics.put("tokens", usage.get("totalTokens")); + } + + if (!metrics.isEmpty()) { + span.setAttribute("braintrust.metrics", toJson(metrics)); + } + } + // ------------------------------------------------------------------------- // Shared helpers // ------------------------------------------------------------------------- @@ -263,6 +390,57 @@ private static JsonNode simplifyAnthropicMessage(JsonNode msg) { return msg; } + /** + * Normalizes a Bedrock Converse message so its content blocks are compatible with the UI's + * schema checks. The Converse wire format uses {@code {"text":"..."}} for text blocks, but both + * the OpenAI and Anthropic schemas the UI validates against require an explicit {@code + * "type":"text"} field. This method adds {@code "type"} to any content block that has a + * recognized Bedrock key but is missing it. + */ + private static JsonNode normalizeBedrockMessage(JsonNode msg) { + if (!msg.has("content") || !msg.get("content").isArray()) { + return msg; + } + var mapper = BraintrustJsonMapper.get(); + ArrayNode normalizedContent = mapper.createArrayNode(); + boolean changed = false; + for (JsonNode block : msg.get("content")) { + if (block.isObject() && !block.has("type")) { + var normalized = (com.fasterxml.jackson.databind.node.ObjectNode) block.deepCopy(); + if (block.has("text")) { + normalized.put("type", "text"); + changed = true; + } else if (block.has("toolUse")) { + normalized.put("type", "tool_use"); + changed = true; + } else if (block.has("toolResult")) { + normalized.put("type", "tool_result"); + changed = true; + } else if (block.has("image")) { + normalized.put("type", "image"); + changed = true; + } + normalizedContent.add(normalized); + } else { + normalizedContent.add(block); + } + } + if (!changed) { + return msg; + } + var result = (com.fasterxml.jackson.databind.node.ObjectNode) msg.deepCopy(); + result.set("content", normalizedContent); + return result; + } + + /** Returns the Bedrock endpoint name from the last URL path segment (e.g. "converse"). */ + private static String bedrockEndpoint(List pathSegments) { + if (pathSegments.isEmpty()) { + return "unknown"; + } + return pathSegments.get(pathSegments.size() - 1); + } + private static String getSpanName(String providerName, List pathSegments) { if (pathSegments.isEmpty()) { return UNSET_LLM_SPAN_NAME; diff --git a/btx/build.gradle b/btx/build.gradle index 72f55b56..150600cf 100644 --- a/btx/build.gradle +++ b/btx/build.gradle @@ -21,6 +21,7 @@ dependencies { testImplementation project(':braintrust-sdk:instrumentation:genai_1_18_0') testImplementation project(':braintrust-sdk:instrumentation:langchain_1_8_0') testImplementation project(':braintrust-sdk:instrumentation:springai_1_0_0') + testImplementation project(':braintrust-sdk:instrumentation:aws_bedrock_2_30_0') // Jackson for JSON processing testImplementation 'com.fasterxml.jackson.core:jackson-databind:2.16.1' @@ -31,6 +32,10 @@ dependencies { // Anthropic SDK testImplementation 'com.anthropic:anthropic-java:2.10.0' + // AWS Bedrock SDK + testImplementation 'software.amazon.awssdk:bedrockruntime:2.30.0' + testImplementation 'software.amazon.awssdk:netty-nio-client:2.30.0' + // Gemini SDK testImplementation 'org.springframework.ai:spring-ai-google-genai:1.1.0' diff --git a/btx/src/test/java/dev/braintrust/sdkspecimpl/SpanConverter.java b/btx/src/test/java/dev/braintrust/sdkspecimpl/SpanConverter.java index fa24035f..f4f18ca1 100644 --- a/btx/src/test/java/dev/braintrust/sdkspecimpl/SpanConverter.java +++ b/btx/src/test/java/dev/braintrust/sdkspecimpl/SpanConverter.java @@ -153,6 +153,31 @@ private static Object transformContentPart(Object part) { return part; } + // Bedrock: {type: image, image: {format: "png", source: {bytes: ""}}} + if ("image".equals(p.get("type")) && p.get("image") instanceof Map) { + // TODO: + // The Braintrust backend does not transform Bedrock image bytes into attachments, + // so we pass this shape through unchanged to match what the backend stores. + // + // Map imageMap = (Map) p.get("image"); + // if (imageMap.get("source") instanceof Map) { + // Map source = (Map) imageMap.get("source"); + // String bytes = (String) source.get("bytes"); + // if (bytes != null) { + // String format = (String) imageMap.getOrDefault("format", "png"); + // String mimeType = "image/" + format; + // Map attachment = + // toAttachment("data:" + mimeType + ";base64," + bytes); + // Map newImage = new LinkedHashMap<>(imageMap); + // newImage.put("source", attachment); + // Map newPart = new LinkedHashMap<>(p); + // newPart.put("image", newImage); + // return newPart; + // } + // } + return part; + } + if (!"image_url".equals(p.get("type"))) return part; Object imageUrlObj = p.get("image_url"); diff --git a/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecExecutor.java b/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecExecutor.java index 238ea897..40dda005 100644 --- a/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecExecutor.java +++ b/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecExecutor.java @@ -17,8 +17,10 @@ import com.openai.models.responses.ResponseCreateParams; import com.openai.models.responses.ResponseInputItem; import com.openai.models.responses.ResponseOutputItem; +import dev.braintrust.Bedrock30TestUtils; import dev.braintrust.TestHarness; import dev.braintrust.instrumentation.anthropic.BraintrustAnthropic; +import dev.braintrust.instrumentation.awsbedrock.v2_30_0.BraintrustAWSBedrock; import dev.braintrust.instrumentation.genai.BraintrustGenAI; import dev.braintrust.instrumentation.langchain.BraintrustLangchain; import dev.braintrust.instrumentation.openai.BraintrustOpenAI; @@ -53,6 +55,16 @@ import org.springframework.ai.openai.api.OpenAiApi; import org.springframework.ai.tool.ToolCallback; import org.springframework.ai.tool.function.FunctionToolCallback; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock; +import software.amazon.awssdk.services.bedrockruntime.model.ConversationRole; +import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest; +import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest; +import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler; +import software.amazon.awssdk.services.bedrockruntime.model.ImageBlock; +import software.amazon.awssdk.services.bedrockruntime.model.ImageFormat; +import software.amazon.awssdk.services.bedrockruntime.model.ImageSource; +import software.amazon.awssdk.services.bedrockruntime.model.Message; /** * Executes LLM spec tests in-process using the Braintrust Java SDK instrumentation. @@ -73,6 +85,7 @@ public class SpecExecutor { private final String openAiApiKey; private final String anthropicBaseUrl; private final String anthropicApiKey; + private final Bedrock30TestUtils bedrockUtils; private final io.opentelemetry.api.OpenTelemetry otel; public SpecExecutor(TestHarness harness) { @@ -83,6 +96,7 @@ public SpecExecutor(TestHarness harness) { this.openAiApiKey = harness.openAiApiKey(); this.anthropicBaseUrl = harness.anthropicBaseUrl(); this.anthropicApiKey = harness.anthropicApiKey(); + this.bedrockUtils = new Bedrock30TestUtils(harness); this.openAIClient = BraintrustOpenAI.wrapOpenAI( @@ -155,6 +169,10 @@ private void dispatchRequest( } else { executeAnthropicMessages(request); } + } else if ("bedrock".equals(provider) && endpoint.contains("/converse-stream")) { + executeBedrockConverseStream(request); + } else if ("bedrock".equals(provider) && endpoint.contains("/converse")) { + executeBedrockConverse(request); } else if ("google".equals(provider) && endpoint.contains(":generateContent")) { executeGeminiGenerateContent(request, endpoint); } else { @@ -615,6 +633,87 @@ private void executeAnthropicMessages(Map request) throws Except } } + // ---- AWS Bedrock ------------------------------------------------------------ + + @SuppressWarnings("unchecked") + private void executeBedrockConverse(Map request) { + String modelId = (String) request.get("modelId"); + + // Build messages from the spec YAML format: [{role, content: [{text: ...} | {image: ...}]}] + List messages = new ArrayList<>(); + for (Map msg : (List>) request.get("messages")) { + String role = (String) msg.get("role"); + List contentBlocks = new ArrayList<>(); + for (Map part : (List>) msg.get("content")) { + if (part.containsKey("text")) { + contentBlocks.add(ContentBlock.fromText((String) part.get("text"))); + } else if (part.containsKey("image")) { + contentBlocks.add( + buildBedrockImageBlock((Map) part.get("image"))); + } + } + messages.add( + Message.builder() + .role(ConversationRole.fromValue(role)) + .content(contentBlocks) + .build()); + } + + var builder = BraintrustAWSBedrock.wrap(otel, bedrockUtils.syncClientBuilder()); + try (var client = builder.build()) { + client.converse(ConverseRequest.builder().modelId(modelId).messages(messages).build()); + } + } + + @SuppressWarnings("unchecked") + private void executeBedrockConverseStream(Map request) throws Exception { + String modelId = (String) request.get("modelId"); + + List messages = new ArrayList<>(); + for (Map msg : (List>) request.get("messages")) { + String role = (String) msg.get("role"); + List contentBlocks = new ArrayList<>(); + for (Map part : (List>) msg.get("content")) { + if (part.containsKey("text")) { + contentBlocks.add(ContentBlock.fromText((String) part.get("text"))); + } + } + messages.add( + Message.builder() + .role(ConversationRole.fromValue(role)) + .content(contentBlocks) + .build()); + } + + var asyncBuilder = BraintrustAWSBedrock.wrap(otel, bedrockUtils.asyncClientBuilder()); + try (var client = asyncBuilder.build()) { + client.converseStream( + ConverseStreamRequest.builder() + .modelId(modelId) + .messages(messages) + .build(), + ConverseStreamResponseHandler.builder() + .subscriber( + ConverseStreamResponseHandler.Visitor.builder().build()) + .build()) + .get(); + } + } + + /** Builds a Bedrock {@link ContentBlock} image from the YAML {@code image:} map. */ + @SuppressWarnings("unchecked") + private static ContentBlock buildBedrockImageBlock(Map imageMap) { + String format = (String) imageMap.getOrDefault("format", "png"); + Map sourceMap = (Map) imageMap.get("source"); + String base64 = (String) sourceMap.get("bytes"); + byte[] imageBytes = java.util.Base64.getDecoder().decode(base64); + return ContentBlock.fromImage( + ImageBlock.builder() + .format(ImageFormat.fromValue(format)) + .source(ImageSource.fromBytes(SdkBytes.fromByteArray(imageBytes))) + .build()); + } + // ---- Google Gemini ---------------------------------------------------------- @SuppressWarnings("unchecked") diff --git a/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecLoader.java b/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecLoader.java index 29157745..1eef02ba 100644 --- a/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecLoader.java +++ b/btx/src/test/java/dev/braintrust/sdkspecimpl/SpecLoader.java @@ -55,7 +55,8 @@ public class SpecLoader { static final Map> CLIENTS_BY_PROVIDER = Map.of( "openai", List.of("openai", "langchain-openai", "springai-openai"), - "anthropic", List.of("anthropic", "springai-anthropic")); + "anthropic", List.of("anthropic", "springai-anthropic"), + "bedrock", List.of("bedrock")); /** * Returns the clients to test for the given provider. Defaults to {@code [providerName]} if the diff --git a/examples/build.gradle b/examples/build.gradle index 7c3860c4..6a231b18 100644 --- a/examples/build.gradle +++ b/examples/build.gradle @@ -25,6 +25,7 @@ dependencies { implementation project(':braintrust-sdk:instrumentation:genai_1_18_0') implementation project(':braintrust-sdk:instrumentation:langchain_1_8_0') implementation project(':braintrust-sdk:instrumentation:springai_1_0_0') + implementation project(':braintrust-sdk:instrumentation:aws_bedrock_2_30_0') runtimeOnly "org.slf4j:slf4j-simple:${slf4jVersion}" // To run otel examples implementation "io.opentelemetry:opentelemetry-exporter-otlp:${otelVersion}" @@ -48,8 +49,7 @@ dependencies { // to run langchain4j examples implementation 'dev.langchain4j:langchain4j:1.9.1' implementation 'dev.langchain4j:langchain4j-open-ai:1.9.1' - // WireMock for stubbing AWS Bedrock Converse API locally - implementation 'org.wiremock:wiremock:3.12.1' + } application { diff --git a/examples/src/main/java/dev/braintrust/examples/SpringAIExample.java b/examples/src/main/java/dev/braintrust/examples/SpringAIExample.java index 1004dd2f..e8a659e0 100644 --- a/examples/src/main/java/dev/braintrust/examples/SpringAIExample.java +++ b/examples/src/main/java/dev/braintrust/examples/SpringAIExample.java @@ -3,6 +3,7 @@ import com.google.genai.Client; import dev.braintrust.Braintrust; import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.instrumentation.awsbedrock.v2_30_0.BraintrustAWSBedrock; import dev.braintrust.instrumentation.genai.BraintrustGenAI; import dev.braintrust.instrumentation.springai.v1_0_0.BraintrustSpringAI; import io.opentelemetry.api.OpenTelemetry; @@ -30,6 +31,7 @@ import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration; import org.springframework.context.annotation.Bean; import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient; /** Spring Boot application demonstrating Braintrust + Spring AI integration */ @SpringBootApplication( @@ -135,12 +137,17 @@ public List chatModels(OpenTelemetry openTelemetry) { if (System.getenv("AWS_ACCESS_KEY_ID") != null && System.getenv("AWS_SECRET_ACCESS_KEY") != null) { + var bedrockClient = + BraintrustAWSBedrock.wrap(openTelemetry, BedrockRuntimeClient.builder()) + .build(); models.add( BedrockProxyChatModel.builder() + .bedrockRuntimeClient(bedrockClient) .region(Region.US_EAST_1) .defaultOptions( BedrockChatOptions.builder() - .model("anthropic.claude-3-haiku-20240307-v1:0") + // .model("us.anthropic.claude-haiku-4-5-20251001-v1:0") + .model("us.amazon.nova-lite-v1:0") .temperature(0.0) .maxTokens(50) .build()) diff --git a/gradle.properties b/gradle.properties index 9ce137d5..21de6bcd 100644 --- a/gradle.properties +++ b/gradle.properties @@ -8,8 +8,7 @@ org.gradle.daemon=true org.gradle.warning.mode=summary # braintrust-spec git ref (SHA or tag) used by btx tests -# TODO: update this to point at a tag once specs publish their first release -braintrustSpecRef=v0.0.1 +braintrustSpecRef=v0.0.2 # Let Gradle locate local JDKs and download one if needed org.gradle.java.installations.auto-detect=true diff --git a/settings.gradle b/settings.gradle index 7aeeb942..67e00999 100644 --- a/settings.gradle +++ b/settings.gradle @@ -16,6 +16,7 @@ include 'braintrust-sdk:instrumentation:anthropic_2_2_0' include 'braintrust-sdk:instrumentation:genai_1_18_0' include 'braintrust-sdk:instrumentation:langchain_1_8_0' include 'braintrust-sdk:instrumentation:springai_1_0_0' +include 'braintrust-sdk:instrumentation:aws_bedrock_2_30_0' include 'braintrust-java-agent:smoke-test:test-instrumentation' include 'braintrust-java-agent:smoke-test:dd-agent' include 'braintrust-java-agent:smoke-test:otel-agent' diff --git a/test-harness/build.gradle b/test-harness/build.gradle index 67068d96..9cbf0a74 100644 --- a/test-harness/build.gradle +++ b/test-harness/build.gradle @@ -41,4 +41,9 @@ dependencies { testFixturesImplementation 'com.google.code.findbugs:jsr305:3.0.2' testFixturesImplementation "org.slf4j:slf4j-api:${slf4jVersion}" testFixturesImplementation "org.slf4j:slf4j-simple:${slf4jVersion}" + + // Bedrock client builders used by Bedrock30TestUtils — compileOnly so consumers + // that don't use Bedrock aren't forced to pull in the AWS SDK. + testFixturesCompileOnly 'software.amazon.awssdk:bedrockruntime:2.30.0' + testFixturesCompileOnly 'software.amazon.awssdk:netty-nio-client:2.30.0' } diff --git a/test-harness/src/testFixtures/java/dev/braintrust/Bedrock30TestUtils.java b/test-harness/src/testFixtures/java/dev/braintrust/Bedrock30TestUtils.java new file mode 100644 index 00000000..077a3617 --- /dev/null +++ b/test-harness/src/testFixtures/java/dev/braintrust/Bedrock30TestUtils.java @@ -0,0 +1,104 @@ +package dev.braintrust; + +import java.net.URI; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.http.Protocol; +import software.amazon.awssdk.http.SdkHttpRequest; +import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClient; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClientBuilder; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClientBuilder; + +/** + * Shared test utilities for constructing AWS Bedrock Runtime client builders wired to the test + * harness (WireMock endpoint override, SigV4 host rewrite, replay credentials). + */ +public class Bedrock30TestUtils { + + public static final String BEDROCK_REGION = "us-east-1"; + public static final String BEDROCK_REAL_HOST = + "bedrock-runtime." + BEDROCK_REGION + ".amazonaws.com"; + + private static final StaticCredentialsProvider REPLAY_CREDS = + StaticCredentialsProvider.create(AwsBasicCredentials.create("fake-key", "fake-secret")); + + private final TestHarness testHarness; + + public Bedrock30TestUtils(TestHarness testHarness) { + this.testHarness = testHarness; + } + + /** + * Returns a {@link BedrockRuntimeClientBuilder} pointed at the test harness WireMock endpoint, + * with the SigV4 host-rewrite interceptor applied and fake credentials in replay mode. + */ + public BedrockRuntimeClientBuilder syncClientBuilder() { + var builder = + BedrockRuntimeClient.builder() + .overrideConfiguration( + ClientOverrideConfiguration.builder() + .addExecutionInterceptor( + new HostRewriteInterceptor(BEDROCK_REAL_HOST)) + .build()) + .region(Region.of(BEDROCK_REGION)) + .endpointOverride(URI.create(testHarness.bedrockBaseUrl(BEDROCK_REGION))); + + if (TestHarness.getVcrMode() == VCR.VcrMode.REPLAY) { + builder.credentialsProvider(REPLAY_CREDS); + } + + return builder; + } + + /** + * Returns a {@link BedrockRuntimeAsyncClientBuilder} pointed at the test harness WireMock + * endpoint, with the SigV4 host-rewrite interceptor, HTTP/1.1 forced (so WireMock can + * proxy/replay the event-stream), and fake credentials in replay mode. + */ + public BedrockRuntimeAsyncClientBuilder asyncClientBuilder() { + var builder = + BedrockRuntimeAsyncClient.builder() + .overrideConfiguration( + ClientOverrideConfiguration.builder() + .addExecutionInterceptor( + new HostRewriteInterceptor(BEDROCK_REAL_HOST)) + .build()) + .region(Region.of(BEDROCK_REGION)) + .endpointOverride(URI.create(testHarness.bedrockBaseUrl(BEDROCK_REGION))) + // Force HTTP/1.1 so WireMock can proxy/replay the event-stream + .httpClientBuilder( + NettyNioAsyncHttpClient.builder().protocol(Protocol.HTTP1_1)); + + if (TestHarness.getVcrMode() == VCR.VcrMode.REPLAY) { + builder.credentialsProvider(REPLAY_CREDS); + } + + return builder; + } + + /** + * Rewrites the {@code Host} header to the real AWS hostname before SigV4 signing, so that + * signatures are valid even when the request is sent to the local WireMock proxy via {@code + * endpointOverride}. + */ + public static class HostRewriteInterceptor implements ExecutionInterceptor { + private final String realHost; + + public HostRewriteInterceptor(String realHost) { + this.realHost = realHost; + } + + @Override + public SdkHttpRequest modifyHttpRequest( + software.amazon.awssdk.core.interceptor.Context.ModifyHttpRequest context, + ExecutionAttributes executionAttributes) { + return context.httpRequest().toBuilder().putHeader("Host", realHost).build(); + } + } +} diff --git a/test-harness/src/testFixtures/java/dev/braintrust/TestHarness.java b/test-harness/src/testFixtures/java/dev/braintrust/TestHarness.java index f45bd8ba..5540c06f 100644 --- a/test-harness/src/testFixtures/java/dev/braintrust/TestHarness.java +++ b/test-harness/src/testFixtures/java/dev/braintrust/TestHarness.java @@ -35,7 +35,10 @@ public class TestHarness { getEnv("OPENAI_API_KEY", ""), getEnv("ANTHROPIC_API_KEY", ""), getEnv("GOOGLE_API_KEY", getEnv("GEMINI_API_KEY", "")), - getEnv("BRAINTRUST_API_KEY", "")); + getEnv("BRAINTRUST_API_KEY", ""), + getEnv("AWS_ACCESS_KEY_ID", ""), + getEnv("AWS_SECRET_ACCESS_KEY", ""), + getEnv("AWS_SESSION_TOKEN", "")); vcr = new VCR( @@ -43,7 +46,8 @@ public class TestHarness { "https://api.openai.com/v1", "openai", "https://api.anthropic.com", "anthropic", "https://generativelanguage.googleapis.com", "google", - "https://api.braintrust.dev", "braintrust"), + "https://api.braintrust.dev", "braintrust", + "https://bedrock-runtime.us-east-1.amazonaws.com", "bedrock"), apiKeysToNeverRecord); vcr.start(); UnitTestShutdownHook.addShutdownHook(1, vcr::stop); @@ -157,6 +161,17 @@ public String googleApiKey() { return getEnv("GOOGLE_API_KEY", getEnv("GEMINI_API_KEY", "test-key")); } + /** + * Returns the VCR proxy URL for the Bedrock Runtime endpoint in the given region. The region + * must have been registered in the VCR target map at init time. + */ + public String bedrockBaseUrl(String region) { + if (!"us-east-1".equals(region)) { + throw new RuntimeException("unsupported region: " + region); + } + return vcr.getUrlForTargetBase("https://bedrock-runtime." + region + ".amazonaws.com"); + } + public String braintrustApiBaseUrl() { return braintrust.config().apiUrl(); } diff --git a/test-harness/src/testFixtures/java/dev/braintrust/VCR.java b/test-harness/src/testFixtures/java/dev/braintrust/VCR.java index c5219ffe..26d95d4c 100644 --- a/test-harness/src/testFixtures/java/dev/braintrust/VCR.java +++ b/test-harness/src/testFixtures/java/dev/braintrust/VCR.java @@ -258,10 +258,13 @@ private void createProgrammaticStubFromMapping( // Extract request body pattern for matching JsonNode bodyPatterns = mapping.at("/request/bodyPatterns"); - // Check if this is an SSE response + // Check if this is an SSE response or binary AWS event-stream response JsonNode contentType = mapping.at("/response/headers/Content-Type"); boolean isSse = contentType.isTextual() && contentType.asText().contains("text/event-stream"); + boolean isEventStream = + contentType.isTextual() + && contentType.asText().contains("application/vnd.amazon.eventstream"); // Check if this is a function invoke request with dynamic OTEL parent info // Only handle invoke requests that have parent.row_ids (which contains dynamic trace IDs) @@ -277,12 +280,16 @@ private void createProgrammaticStubFromMapping( } } - if (!isSse && !isFunctionInvokeWithParent) { + if (!isSse && !isEventStream && !isFunctionInvokeWithParent) { return; // Let WireMock handle other responses normally } if (isSse) { log.info("Creating programmatic stub for SSE response: " + mappingPath.getFileName()); + } else if (isEventStream) { + log.info( + "Creating programmatic stub for binary event-stream response: " + + mappingPath.getFileName()); } else { log.info( "Creating programmatic stub for function invoke with parent: " @@ -294,17 +301,6 @@ private void createProgrammaticStubFromMapping( String responseContentType = contentType.isTextual() ? contentType.asText() : "application/json"; - String body; - if (mapping.at("/response/body").isTextual()) { - body = mapping.at("/response/body").asText(); - } else if (mapping.at("/response/bodyFileName").isTextual()) { - String bodyFileName = mapping.at("/response/bodyFileName").asText(); - Path bodyPath = Paths.get(cassettesRoot, mappingsDir, "__files", bodyFileName); - body = Files.readString(bodyPath); - } else { - return; - } - // Create programmatic stub com.github.tomakehurst.wiremock.client.MappingBuilder stub = com.github.tomakehurst.wiremock.client.WireMock.request( @@ -330,10 +326,29 @@ private void createProgrammaticStubFromMapping( com.github.tomakehurst.wiremock.client.ResponseDefinitionBuilder response = com.github.tomakehurst.wiremock.client.WireMock.aResponse() .withStatus(status) - .withHeader("Content-Type", responseContentType) - .withBody(body); + .withHeader("Content-Type", responseContentType); + + // Binary event-stream bodies must be served as raw bytes to avoid UTF-8 corruption + if (isEventStream) { + if (mapping.at("/response/bodyFileName").isTextual()) { + String bodyFileName = mapping.at("/response/bodyFileName").asText(); + Path bodyPath = Paths.get(cassettesRoot, mappingsDir, "__files", bodyFileName); + response.withBody(Files.readAllBytes(bodyPath)); + } else { + return; + } + } else { + if (mapping.at("/response/body").isTextual()) { + response.withBody(mapping.at("/response/body").asText()); + } else if (mapping.at("/response/bodyFileName").isTextual()) { + String bodyFileName = mapping.at("/response/bodyFileName").asText(); + Path bodyPath = Paths.get(cassettesRoot, mappingsDir, "__files", bodyFileName); + response.withBody(Files.readString(bodyPath)); + } else { + return; + } + } - // Use instance method wireMock.stubFor(stub.willReturn(response)); }