|
| 1 | +package dev.braintrust.instrumentation.anthropic.v2_2_0; |
| 2 | + |
| 3 | +import static org.junit.jupiter.api.Assertions.*; |
| 4 | + |
| 5 | +import com.anthropic.client.AnthropicClient; |
| 6 | +import com.anthropic.client.okhttp.AnthropicOkHttpClient; |
| 7 | +import com.anthropic.core.JsonValue; |
| 8 | +import com.anthropic.models.messages.CacheControlEphemeral; |
| 9 | +import com.anthropic.models.messages.MessageCreateParams; |
| 10 | +import com.anthropic.models.messages.TextBlockParam; |
| 11 | +import com.fasterxml.jackson.databind.JsonNode; |
| 12 | +import com.fasterxml.jackson.databind.ObjectMapper; |
| 13 | +import dev.braintrust.TestHarness; |
| 14 | +import dev.braintrust.instrumentation.Instrumenter; |
| 15 | +import io.opentelemetry.api.common.AttributeKey; |
| 16 | +import java.util.List; |
| 17 | +import java.util.UUID; |
| 18 | +import lombok.SneakyThrows; |
| 19 | +import net.bytebuddy.agent.ByteBuddyAgent; |
| 20 | +import org.junit.jupiter.api.BeforeAll; |
| 21 | +import org.junit.jupiter.api.BeforeEach; |
| 22 | +import org.junit.jupiter.api.Test; |
| 23 | + |
| 24 | +/** |
| 25 | + * Tests that prompt caching metrics (per-TTL breakdown) are correctly extracted from Anthropic API |
| 26 | + * responses and attached to spans. |
| 27 | + * |
| 28 | + * <p>Uses a cache-buster nonce in the system prompt to guarantee cache misses, ensuring {@code |
| 29 | + * cache_creation_input_tokens} is always positive. In VCR replay mode the nonce is a fixed string |
| 30 | + * so cassette matching still works. |
| 31 | + */ |
| 32 | +public class BraintrustAnthropicPromptCachingTest { |
| 33 | + private static final String TEST_MODEL = "claude-sonnet-4-5-20250929"; |
| 34 | + private static final ObjectMapper JSON_MAPPER = new ObjectMapper(); |
| 35 | + |
| 36 | + /** |
| 37 | + * Nonce injected into system prompts to bust Anthropic's server-side prompt cache. Random UUID |
| 38 | + * when running live ({@code VCR_MODE=off}), fixed string otherwise so VCR cassettes match. |
| 39 | + */ |
| 40 | + private static final String VCR_NONCE; |
| 41 | + |
| 42 | + static { |
| 43 | + String vcrMode = System.getenv().getOrDefault("VCR_MODE", "replay").toUpperCase(); |
| 44 | + VCR_NONCE = "OFF".equals(vcrMode) ? UUID.randomUUID().toString() : "vcr-mode"; |
| 45 | + } |
| 46 | + |
| 47 | + @BeforeAll |
| 48 | + public static void beforeAll() { |
| 49 | + var instrumentation = ByteBuddyAgent.install(); |
| 50 | + Instrumenter.install( |
| 51 | + instrumentation, BraintrustAnthropicPromptCachingTest.class.getClassLoader()); |
| 52 | + } |
| 53 | + |
| 54 | + private TestHarness testHarness; |
| 55 | + |
| 56 | + @BeforeEach |
| 57 | + void beforeEach() { |
| 58 | + testHarness = TestHarness.setup(); |
| 59 | + } |
| 60 | + |
| 61 | + /** |
| 62 | + * Sends a single request with two system content blocks — one cached at 5m TTL and one at 1h |
| 63 | + * TTL — and verifies that the per-TTL prompt cache creation metrics are present on the span. |
| 64 | + */ |
| 65 | + @Test |
| 66 | + @SneakyThrows |
| 67 | + void testPromptCachingDualTtl() { |
| 68 | + AnthropicClient client = |
| 69 | + AnthropicOkHttpClient.builder() |
| 70 | + .baseUrl(testHarness.anthropicBaseUrl()) |
| 71 | + .apiKey(testHarness.anthropicApiKey()) |
| 72 | + .build(); |
| 73 | + |
| 74 | + // 1h TTL block — must come before 5m (Anthropic requires descending TTL order). |
| 75 | + // Requires the extended-cache-ttl beta header. |
| 76 | + // The text must exceed Claude Sonnet's minimum cacheable size (~1024 tokens / ~4000 chars). |
| 77 | + CacheControlEphemeral cacheControl1h = |
| 78 | + CacheControlEphemeral.builder() |
| 79 | + .putAdditionalProperty("ttl", JsonValue.from("1h")) |
| 80 | + .build(); |
| 81 | + |
| 82 | + TextBlockParam systemBlock1h = |
| 83 | + TextBlockParam.builder() |
| 84 | + .text( |
| 85 | + buildPaddedSystemText( |
| 86 | + "1h-block", |
| 87 | + "Reference: capitals include Paris, Berlin, Rome, Madrid," |
| 88 | + + " Lisbon.")) |
| 89 | + .cacheControl(cacheControl1h) |
| 90 | + .build(); |
| 91 | + |
| 92 | + // 5m TTL block — default ephemeral cache. |
| 93 | + CacheControlEphemeral cacheControl5m = |
| 94 | + CacheControlEphemeral.builder() |
| 95 | + .putAdditionalProperty("ttl", JsonValue.from("5m")) |
| 96 | + .build(); |
| 97 | + |
| 98 | + TextBlockParam systemBlock5m = |
| 99 | + TextBlockParam.builder() |
| 100 | + .text( |
| 101 | + buildPaddedSystemText( |
| 102 | + "5m-block", |
| 103 | + "You are a helpful geography assistant. Answer in one" |
| 104 | + + " sentence.")) |
| 105 | + .cacheControl(cacheControl5m) |
| 106 | + .build(); |
| 107 | + |
| 108 | + var request = |
| 109 | + MessageCreateParams.builder() |
| 110 | + .model(TEST_MODEL) |
| 111 | + .systemOfTextBlockParams(List.of(systemBlock1h, systemBlock5m)) |
| 112 | + .addUserMessage("What is the capital of France?") |
| 113 | + .maxTokens(128) |
| 114 | + .temperature(0.0) |
| 115 | + .putAdditionalHeader("anthropic-beta", "extended-cache-ttl-2025-04-11") |
| 116 | + .build(); |
| 117 | + |
| 118 | + var response = client.messages().create(request); |
| 119 | + |
| 120 | + // Basic response sanity |
| 121 | + assertNotNull(response); |
| 122 | + assertNotNull(response.id()); |
| 123 | + var contentBlock = response.content().get(0); |
| 124 | + assertTrue(contentBlock.isText()); |
| 125 | + assertFalse(contentBlock.asText().text().isEmpty()); |
| 126 | + |
| 127 | + // Verify span metrics |
| 128 | + var spans = testHarness.awaitExportedSpans(); |
| 129 | + assertEquals(1, spans.size()); |
| 130 | + var span = spans.get(0); |
| 131 | + |
| 132 | + String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics")); |
| 133 | + assertNotNull(metricsJson, "metrics should be present"); |
| 134 | + JsonNode metrics = JSON_MAPPER.readTree(metricsJson); |
| 135 | + |
| 136 | + // Standard token metrics |
| 137 | + assertPositive(metrics, "prompt_tokens"); |
| 138 | + assertPositive(metrics, "completion_tokens"); |
| 139 | + assertPositive(metrics, "tokens"); |
| 140 | + |
| 141 | + // Per-TTL breakdown — both should be positive on a cold cache |
| 142 | + assertFalse( |
| 143 | + metrics.has("prompt_cache_creation_tokens"), |
| 144 | + "anthropic cache tokens must be set INSTEAD of the aggregate meteric"); |
| 145 | + assertPositive(metrics, "prompt_cache_creation_5m_tokens"); |
| 146 | + assertPositive(metrics, "prompt_cache_creation_1h_tokens"); |
| 147 | + assertEquals( |
| 148 | + response.usage().cacheCreationInputTokens().get(), |
| 149 | + (metrics.get("prompt_cache_creation_5m_tokens").intValue() |
| 150 | + + metrics.get("prompt_cache_creation_1h_tokens").intValue()), |
| 151 | + "ttl tokens should sum to total token count"); |
| 152 | + |
| 153 | + // Cache read may be 0 on cold cache, but should be present and non-negative |
| 154 | + assertTrue(metrics.has("prompt_cached_tokens"), "prompt_cached_tokens should be present"); |
| 155 | + assertTrue( |
| 156 | + metrics.get("prompt_cached_tokens").asDouble() >= 0, |
| 157 | + "prompt_cached_tokens should be non-negative"); |
| 158 | + } |
| 159 | + |
| 160 | + private static void assertPositive(JsonNode metrics, String field) { |
| 161 | + assertTrue(metrics.has(field), field + " should be present"); |
| 162 | + assertTrue( |
| 163 | + metrics.get(field).isNumber(), |
| 164 | + field + " should be a number, got: " + metrics.get(field)); |
| 165 | + assertTrue( |
| 166 | + metrics.get(field).asDouble() > 0, |
| 167 | + field + " should be positive, got: " + metrics.get(field).asDouble()); |
| 168 | + } |
| 169 | + |
| 170 | + /** |
| 171 | + * Build a system prompt text padded to exceed the minimum cacheable size (~4000 ASCII chars). |
| 172 | + * Includes the VCR nonce and a block identifier for cache isolation. |
| 173 | + */ |
| 174 | + private String buildPaddedSystemText(String blockId, String coreInstruction) { |
| 175 | + StringBuilder sb = new StringBuilder(); |
| 176 | + sb.append("[cache-buster: ").append(blockId).append(" ").append(VCR_NONCE).append("]\n"); |
| 177 | + sb.append(coreInstruction).append("\n\n"); |
| 178 | + |
| 179 | + // Pad with numbered guidelines to exceed the token minimum. |
| 180 | + for (int i = 1; i <= 80; i++) { |
| 181 | + sb.append(i) |
| 182 | + .append(". This is guideline number ") |
| 183 | + .append(i) |
| 184 | + .append(". It exists to pad the system prompt past the minimum cacheable ") |
| 185 | + .append("token threshold for Claude Sonnet models. Each guideline adds ") |
| 186 | + .append("approximately fifty tokens of padding text to the prompt.\n"); |
| 187 | + } |
| 188 | + return sb.toString(); |
| 189 | + } |
| 190 | +} |
0 commit comments