Skip to content

Commit 22865ae

Browse files
seratchseratch
authored
feat: #678 Add a list of per-request usage data to Usage (#686)
1 parent 46df17d commit 22865ae

File tree

9 files changed

+345
-29
lines changed

9 files changed

+345
-29
lines changed

‎.changeset/three-islands-pump.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
---
2+
'@openai/agents-extensions': patch
3+
'@openai/agents-openai': patch
4+
'@openai/agents-core': patch
5+
---
6+
7+
feat: #678 Add a list of per-request usage data to Usage

‎packages/agents-core/src/index.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,7 @@ export type {
176176
StreamEventResponseStarted,
177177
StreamEventGenericItem,
178178
} from './types';
179-
export { Usage } from './usage';
179+
export { RequestUsage, Usage } from './usage';
180180
export type { Session, SessionInputCallback } from './memory/session';
181181
export { MemorySession } from './memory/memorySession';
182182

‎packages/agents-core/src/runState.ts‎

Lines changed: 35 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -72,11 +72,22 @@ const SerializedSpan: z.ZodType<SerializedSpanType> = serializedSpanBase.extend(
7272
},
7373
);
7474

75+
const requestUsageSchema = z.object({
76+
inputTokens: z.number(),
77+
outputTokens: z.number(),
78+
totalTokens: z.number(),
79+
inputTokensDetails: z.record(z.string(), z.number()).optional(),
80+
outputTokensDetails: z.record(z.string(), z.number()).optional(),
81+
});
82+
7583
const usageSchema = z.object({
7684
requests: z.number(),
7785
inputTokens: z.number(),
7886
outputTokens: z.number(),
7987
totalTokens: z.number(),
88+
inputTokensDetails: z.array(z.record(z.string(), z.number())).optional(),
89+
outputTokensDetails: z.array(z.record(z.string(), z.number())).optional(),
90+
requestUsageEntries: z.array(requestUsageSchema).optional(),
8091
});
8192

8293
const modelResponseSchema = z.object({
@@ -287,6 +298,13 @@ export class RunState<TContext, TAgent extends Agent<any, any>> {
287298
* Run context tracking approvals, usage, and other metadata.
288299
*/
289300
public _context: RunContext<TContext>;
301+
302+
/**
303+
* The usage aggregated for this run. This includes per-request breakdowns when available.
304+
*/
305+
get usage(): Usage {
306+
return this._context.usage;
307+
}
290308
/**
291309
* Tracks what tools each agent has used.
292310
*/
@@ -440,6 +458,22 @@ export class RunState<TContext, TAgent extends Agent<any, any>> {
440458
inputTokens: response.usage.inputTokens,
441459
outputTokens: response.usage.outputTokens,
442460
totalTokens: response.usage.totalTokens,
461+
inputTokensDetails: response.usage.inputTokensDetails,
462+
outputTokensDetails: response.usage.outputTokensDetails,
463+
...(response.usage.requestUsageEntries &&
464+
response.usage.requestUsageEntries.length > 0
465+
? {
466+
requestUsageEntries: response.usage.requestUsageEntries.map(
467+
(entry) => ({
468+
inputTokens: entry.inputTokens,
469+
outputTokens: entry.outputTokens,
470+
totalTokens: entry.totalTokens,
471+
inputTokensDetails: entry.inputTokensDetails,
472+
outputTokensDetails: entry.outputTokensDetails,
473+
}),
474+
),
475+
}
476+
: {}),
443477
},
444478
output: response.output as any,
445479
responseId: response.responseId,
@@ -683,11 +717,7 @@ export function deserializeSpan(
683717
export function deserializeModelResponse(
684718
serializedModelResponse: z.infer<typeof modelResponseSchema>,
685719
): ModelResponse {
686-
const usage = new Usage();
687-
usage.requests = serializedModelResponse.usage.requests;
688-
usage.inputTokens = serializedModelResponse.usage.inputTokens;
689-
usage.outputTokens = serializedModelResponse.usage.outputTokens;
690-
usage.totalTokens = serializedModelResponse.usage.totalTokens;
720+
const usage = new Usage(serializedModelResponse.usage);
691721

692722
return {
693723
usage,

‎packages/agents-core/src/types/protocol.ts‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -743,15 +743,36 @@ export type ModelItem = z.infer<typeof ModelItem>;
743743
// Meta data types
744744
// ----------------------------
745745

746-
export const UsageData = z.object({
747-
requests: z.number().optional(),
746+
export const RequestUsageData = z.object({
748747
inputTokens: z.number(),
749748
outputTokens: z.number(),
750749
totalTokens: z.number(),
751750
inputTokensDetails: z.record(z.string(), z.number()).optional(),
752751
outputTokensDetails: z.record(z.string(), z.number()).optional(),
753752
});
754753

754+
export type RequestUsageData = z.infer<typeof RequestUsageData>;
755+
756+
export const UsageData = z.object({
757+
requests: z.number().optional(),
758+
inputTokens: z.number(),
759+
outputTokens: z.number(),
760+
totalTokens: z.number(),
761+
inputTokensDetails: z
762+
.union([
763+
z.record(z.string(), z.number()),
764+
z.array(z.record(z.string(), z.number())),
765+
])
766+
.optional(),
767+
outputTokensDetails: z
768+
.union([
769+
z.record(z.string(), z.number()),
770+
z.array(z.record(z.string(), z.number())),
771+
])
772+
.optional(),
773+
requestUsageEntries: z.array(RequestUsageData).optional(),
774+
});
775+
755776
export type UsageData = z.infer<typeof UsageData>;
756777

757778
// ----------------------------

‎packages/agents-core/src/usage.ts‎

Lines changed: 139 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,80 @@
1-
import { UsageData } from './types/protocol';
1+
import { RequestUsageData, UsageData } from './types/protocol';
22

3-
type UsageInput = Partial<
4-
UsageData & {
3+
type RequestUsageInput = Partial<
4+
RequestUsageData & {
55
input_tokens: number;
66
output_tokens: number;
77
total_tokens: number;
88
input_tokens_details: object;
99
output_tokens_details: object;
1010
}
11-
> & { requests?: number };
11+
>;
12+
13+
type UsageInput = Partial<
14+
UsageData & {
15+
input_tokens: number;
16+
output_tokens: number;
17+
total_tokens: number;
18+
input_tokens_details:
19+
| Record<string, number>
20+
| Array<Record<string, number>>
21+
| object;
22+
output_tokens_details:
23+
| Record<string, number>
24+
| Array<Record<string, number>>
25+
| object;
26+
request_usage_entries: RequestUsageInput[];
27+
}
28+
> & { requests?: number; requestUsageEntries?: RequestUsageInput[] };
29+
30+
/**
31+
* Usage details for a single API request.
32+
*/
33+
export class RequestUsage {
34+
/**
35+
* The number of input tokens used for this request.
36+
*/
37+
public inputTokens: number;
38+
39+
/**
40+
* The number of output tokens used for this request.
41+
*/
42+
public outputTokens: number;
43+
44+
/**
45+
* The total number of tokens sent and received for this request.
46+
*/
47+
public totalTokens: number;
48+
49+
/**
50+
* Details about the input tokens used for this request.
51+
*/
52+
public inputTokensDetails: Record<string, number>;
53+
54+
/**
55+
* Details about the output tokens used for this request.
56+
*/
57+
public outputTokensDetails: Record<string, number>;
58+
59+
constructor(input?: RequestUsageInput) {
60+
this.inputTokens = input?.inputTokens ?? input?.input_tokens ?? 0;
61+
this.outputTokens = input?.outputTokens ?? input?.output_tokens ?? 0;
62+
this.totalTokens =
63+
input?.totalTokens ??
64+
input?.total_tokens ??
65+
this.inputTokens + this.outputTokens;
66+
const inputTokensDetails =
67+
input?.inputTokensDetails ?? input?.input_tokens_details;
68+
this.inputTokensDetails = inputTokensDetails
69+
? (inputTokensDetails as Record<string, number>)
70+
: {};
71+
const outputTokensDetails =
72+
input?.outputTokensDetails ?? input?.output_tokens_details;
73+
this.outputTokensDetails = outputTokensDetails
74+
? (outputTokensDetails as Record<string, number>)
75+
: {};
76+
}
77+
}
1278

1379
/**
1480
* Tracks token usage and request counts for an agent run.
@@ -44,6 +110,11 @@ export class Usage {
44110
*/
45111
public outputTokensDetails: Array<Record<string, number>> = [];
46112

113+
/**
114+
* List of per-request usage entries for detailed cost calculations.
115+
*/
116+
public requestUsageEntries: RequestUsage[] | undefined;
117+
47118
constructor(input?: UsageInput) {
48119
if (typeof input === 'undefined') {
49120
this.requests = 0;
@@ -52,29 +123,58 @@ export class Usage {
52123
this.totalTokens = 0;
53124
this.inputTokensDetails = [];
54125
this.outputTokensDetails = [];
126+
this.requestUsageEntries = undefined;
55127
} else {
56128
this.requests = input?.requests ?? 1;
57129
this.inputTokens = input?.inputTokens ?? input?.input_tokens ?? 0;
58130
this.outputTokens = input?.outputTokens ?? input?.output_tokens ?? 0;
59-
this.totalTokens = input?.totalTokens ?? input?.total_tokens ?? 0;
131+
this.totalTokens =
132+
input?.totalTokens ??
133+
input?.total_tokens ??
134+
this.inputTokens + this.outputTokens;
60135
const inputTokensDetails =
61136
input?.inputTokensDetails ?? input?.input_tokens_details;
62-
this.inputTokensDetails = inputTokensDetails
63-
? [inputTokensDetails as Record<string, number>]
64-
: [];
137+
if (Array.isArray(inputTokensDetails)) {
138+
this.inputTokensDetails = inputTokensDetails as Array<
139+
Record<string, number>
140+
>;
141+
} else {
142+
this.inputTokensDetails = inputTokensDetails
143+
? [inputTokensDetails as Record<string, number>]
144+
: [];
145+
}
65146
const outputTokensDetails =
66147
input?.outputTokensDetails ?? input?.output_tokens_details;
67-
this.outputTokensDetails = outputTokensDetails
68-
? [outputTokensDetails as Record<string, number>]
69-
: [];
148+
if (Array.isArray(outputTokensDetails)) {
149+
this.outputTokensDetails = outputTokensDetails as Array<
150+
Record<string, number>
151+
>;
152+
} else {
153+
this.outputTokensDetails = outputTokensDetails
154+
? [outputTokensDetails as Record<string, number>]
155+
: [];
156+
}
157+
158+
const requestUsageEntries =
159+
input?.requestUsageEntries ?? input?.request_usage_entries;
160+
const normalizedRequestUsageEntries = Array.isArray(requestUsageEntries)
161+
? requestUsageEntries.map((entry) =>
162+
entry instanceof RequestUsage ? entry : new RequestUsage(entry),
163+
)
164+
: undefined;
165+
this.requestUsageEntries =
166+
normalizedRequestUsageEntries &&
167+
normalizedRequestUsageEntries.length > 0
168+
? normalizedRequestUsageEntries
169+
: undefined;
70170
}
71171
}
72172

73173
add(newUsage: Usage) {
74-
this.requests += newUsage.requests;
75-
this.inputTokens += newUsage.inputTokens;
76-
this.outputTokens += newUsage.outputTokens;
77-
this.totalTokens += newUsage.totalTokens;
174+
this.requests += newUsage.requests ?? 0;
175+
this.inputTokens += newUsage.inputTokens ?? 0;
176+
this.outputTokens += newUsage.outputTokens ?? 0;
177+
this.totalTokens += newUsage.totalTokens ?? 0;
78178
if (newUsage.inputTokensDetails) {
79179
// The type does not allow undefined, but it could happen runtime
80180
this.inputTokensDetails.push(...newUsage.inputTokensDetails);
@@ -83,7 +183,30 @@ export class Usage {
83183
// The type does not allow undefined, but it could happen runtime
84184
this.outputTokensDetails.push(...newUsage.outputTokensDetails);
85185
}
186+
187+
if (
188+
Array.isArray(newUsage.requestUsageEntries) &&
189+
newUsage.requestUsageEntries.length > 0
190+
) {
191+
this.requestUsageEntries ??= [];
192+
this.requestUsageEntries.push(
193+
...newUsage.requestUsageEntries.map((entry) =>
194+
entry instanceof RequestUsage ? entry : new RequestUsage(entry),
195+
),
196+
);
197+
} else if (newUsage.requests === 1 && newUsage.totalTokens > 0) {
198+
this.requestUsageEntries ??= [];
199+
this.requestUsageEntries.push(
200+
new RequestUsage({
201+
inputTokens: newUsage.inputTokens,
202+
outputTokens: newUsage.outputTokens,
203+
totalTokens: newUsage.totalTokens,
204+
inputTokensDetails: newUsage.inputTokensDetails?.[0],
205+
outputTokensDetails: newUsage.outputTokensDetails?.[0],
206+
}),
207+
);
208+
}
86209
}
87210
}
88211

89-
export { UsageData };
212+
export { RequestUsageData, UsageData };

‎packages/agents-core/test/run.test.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,40 @@ describe('Runner.run', () => {
8888
expectTypeOf(result.finalOutput).toEqualTypeOf<string | undefined>();
8989
});
9090

91+
it('exposes aggregated usage on run results', async () => {
92+
const model = new FakeModel([
93+
{
94+
output: [fakeModelMessage('hi there')],
95+
usage: new Usage({
96+
requests: 1,
97+
inputTokens: 2,
98+
outputTokens: 3,
99+
totalTokens: 5,
100+
}),
101+
responseId: 'usage-res',
102+
},
103+
]);
104+
const agent = new Agent({
105+
name: 'UsageAgent',
106+
model,
107+
});
108+
109+
const result = await run(agent, 'ping');
110+
111+
expect(result.state.usage.inputTokens).toBe(2);
112+
expect(result.state.usage.outputTokens).toBe(3);
113+
expect(result.state.usage.totalTokens).toBe(5);
114+
expect(result.state.usage.requestUsageEntries).toEqual([
115+
{
116+
inputTokens: 2,
117+
outputTokens: 3,
118+
totalTokens: 5,
119+
inputTokensDetails: {},
120+
outputTokensDetails: {},
121+
},
122+
]);
123+
});
124+
91125
it('sholuld handle structured output', async () => {
92126
const fakeModel = new FakeModel([
93127
{

0 commit comments

Comments
 (0)