Skip to content

Commit 957705e

Browse files
authored
1 parent 7897d0f commit 957705e

4 files changed

Lines changed: 120 additions & 36 deletions

File tree

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"braintrust": patch
3+
---
4+
5+
fix: Fix Langchain anthropic token metrics

‎e2e/scenarios/cloudflare-agents-instrumentation/scenario.ts‎

Lines changed: 16 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,15 @@
11
import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
22
import { once } from "node:events";
33
import { writeFile } from "node:fs/promises";
4-
import net from "node:net";
54
import path from "node:path";
5+
import { stripVTControlCharacters } from "node:util";
66
import {
77
getTestRunId,
88
runMain,
99
scopedName,
1010
} from "../../helpers/scenario-runtime";
1111

1212
async function main() {
13-
const port = await getFreePort();
1413
const viteBin = path.join(
1514
process.cwd(),
1615
"node_modules",
@@ -21,7 +20,7 @@ async function main() {
2120
await buildWorker(viteBin);
2221
const server = spawn(
2322
viteBin,
24-
["preview", "--host", "127.0.0.1", "--port", String(port), "--strictPort"],
23+
["preview", "--host", "127.0.0.1", "--port", "0", "--strictPort"],
2524
{
2625
cwd: process.cwd(),
2726
env: process.env,
@@ -31,8 +30,7 @@ async function main() {
3130
const output = captureOutput(server);
3231

3332
try {
34-
const baseUrl = `http://127.0.0.1:${port}`;
35-
await waitForServer(baseUrl, server, output);
33+
const baseUrl = await waitForServer(server, output);
3634
const testRunId = getTestRunId();
3735
const projectName = scopedName(
3836
"e2e-cloudflare-agents-instrumentation",
@@ -99,22 +97,6 @@ async function buildWorker(viteBin: string): Promise<void> {
9997
}
10098
}
10199

102-
async function getFreePort(): Promise<number> {
103-
const server = net.createServer();
104-
await new Promise<void>((resolve, reject) => {
105-
server.once("error", reject);
106-
server.listen(0, "127.0.0.1", resolve);
107-
});
108-
const address = server.address();
109-
await new Promise<void>((resolve, reject) => {
110-
server.close((error) => (error ? reject(error) : resolve()));
111-
});
112-
if (!address || typeof address === "string") {
113-
throw new Error("Could not allocate a Vite preview-server port");
114-
}
115-
return address.port;
116-
}
117-
118100
function captureOutput(child: ChildProcessWithoutNullStreams): () => string {
119101
let stdout = "";
120102
let stderr = "";
@@ -128,24 +110,29 @@ function captureOutput(child: ChildProcessWithoutNullStreams): () => string {
128110
}
129111

130112
async function waitForServer(
131-
baseUrl: string,
132113
server: ChildProcessWithoutNullStreams,
133114
output: () => string,
134-
): Promise<void> {
115+
): Promise<string> {
135116
const startedAt = Date.now();
136117
while (Date.now() - startedAt < 60_000) {
137118
if (server.exitCode !== null) {
138119
throw new Error(
139120
`Vite exited early with code ${server.exitCode}\n${output()}`,
140121
);
141122
}
142-
try {
143-
const response = await fetch(`${baseUrl}/health`);
144-
if (response.ok) {
145-
return;
123+
// Vite binds port 0 and reports the assigned URL, avoiding a port reservation race.
124+
const baseUrl = stripVTControlCharacters(output()).match(
125+
/Local:\s+(http:\/\/127\.0\.0\.1:\d+)\//,
126+
)?.[1];
127+
if (baseUrl) {
128+
try {
129+
const response = await fetch(`${baseUrl}/health`);
130+
if (response.ok) {
131+
return baseUrl;
132+
}
133+
} catch {
134+
// Continue until workerd accepts requests.
146135
}
147-
} catch {
148-
// Continue until workerd accepts requests.
149136
}
150137
await new Promise((resolve) => setTimeout(resolve, 250));
151138
}

‎js/src/wrappers/langchain/callback-handler.test.ts‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,91 @@ describe("BraintrustLangChainCallbackHandler metrics", () => {
129129
});
130130
});
131131

132+
it.each([
133+
{
134+
name: "TTL buckets without an aggregate",
135+
details: {
136+
cache_creation: undefined,
137+
ephemeral_5m_input_tokens: 4,
138+
ephemeral_1h_input_tokens: 0,
139+
},
140+
expected: {
141+
prompt_cache_creation_5m_tokens: 4,
142+
prompt_cache_creation_1h_tokens: 0,
143+
},
144+
},
145+
{
146+
name: "both TTL buckets",
147+
details: { ephemeral_5m_input_tokens: 4, ephemeral_1h_input_tokens: 6 },
148+
expected: {
149+
prompt_cache_creation_5m_tokens: 4,
150+
prompt_cache_creation_1h_tokens: 6,
151+
},
152+
},
153+
{
154+
name: "a zero-valued TTL bucket",
155+
details: { ephemeral_5m_input_tokens: 4, ephemeral_1h_input_tokens: 0 },
156+
expected: {
157+
prompt_cache_creation_5m_tokens: 4,
158+
prompt_cache_creation_1h_tokens: 0,
159+
},
160+
},
161+
{
162+
name: "only the 5-minute bucket",
163+
details: { ephemeral_5m_input_tokens: 4 },
164+
expected: { prompt_cache_creation_5m_tokens: 4 },
165+
},
166+
{
167+
name: "only the 1-hour bucket",
168+
details: { ephemeral_1h_input_tokens: 6 },
169+
expected: { prompt_cache_creation_1h_tokens: 6 },
170+
},
171+
{
172+
name: "only a zero-valued bucket",
173+
details: { ephemeral_1h_input_tokens: 0 },
174+
expected: { prompt_cache_creation_1h_tokens: 0 },
175+
},
176+
{
177+
name: "null TTL buckets",
178+
details: {
179+
ephemeral_5m_input_tokens: null,
180+
ephemeral_1h_input_tokens: null,
181+
},
182+
expected: { prompt_cache_creation_tokens: 10 },
183+
},
184+
])(
185+
"preserves cache creation metrics with $name",
186+
async ({ details, expected }) => {
187+
const { endLog } = await finishChatModelRun({
188+
generations: [
189+
[
190+
{
191+
message: {
192+
usage_metadata: {
193+
input_tokens: 20,
194+
output_tokens: 2,
195+
input_token_details: {
196+
cache_creation: 10,
197+
cache_read: 3,
198+
...details,
199+
},
200+
},
201+
},
202+
},
203+
],
204+
],
205+
});
206+
207+
expect(endLog.metrics).toEqual({
208+
prompt_tokens: 20,
209+
completion_tokens: 2,
210+
prompt_cached_tokens: 3,
211+
tokens: 22,
212+
...expected,
213+
});
214+
},
215+
);
216+
132217
it("preserves reasoning metrics from message usage metadata", async () => {
133218
const { endLog } = await finishChatModelRun({
134219
generations: [

‎js/src/wrappers/langchain/callback-handler.ts‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -522,18 +522,25 @@ function getMetricsFromResponse(
522522
continue;
523523
}
524524

525-
const inputTokenDetails = usageMetadata.input_token_details;
525+
const inputTokenDetails = isRecord(usageMetadata.input_token_details)
526+
? usageMetadata.input_token_details
527+
: {};
526528
const outputTokenDetails = usageMetadata.output_token_details;
527529
return normalizeTokenMetrics({
528530
total_tokens: usageMetadata.total_tokens,
529531
prompt_tokens: usageMetadata.input_tokens,
530532
completion_tokens: usageMetadata.output_tokens,
531-
prompt_cache_creation_tokens: isRecord(inputTokenDetails)
532-
? inputTokenDetails.cache_creation
533-
: undefined,
534-
prompt_cached_tokens: isRecord(inputTokenDetails)
535-
? inputTokenDetails.cache_read
536-
: undefined,
533+
// Prefer TTL-specific cache writes over the aggregate when available.
534+
prompt_cache_creation_tokens:
535+
inputTokenDetails.ephemeral_5m_input_tokens == null &&
536+
inputTokenDetails.ephemeral_1h_input_tokens == null
537+
? inputTokenDetails.cache_creation
538+
: undefined,
539+
prompt_cache_creation_5m_tokens:
540+
inputTokenDetails.ephemeral_5m_input_tokens,
541+
prompt_cache_creation_1h_tokens:
542+
inputTokenDetails.ephemeral_1h_input_tokens,
543+
prompt_cached_tokens: inputTokenDetails.cache_read,
537544
completion_reasoning_tokens: isRecord(outputTokenDetails)
538545
? outputTokenDetails.reasoning
539546
: undefined,

0 commit comments

Comments
 (0)