@outputai/llm 0.1.13-next.f537949.0 → 0.2.1-dev.82acd68.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/src/index.d.ts CHANGED
@@ -1,4 +1,5 @@
1
1
  import type {
2
+ AgentStreamParameters,
2
3
  GenerateTextResult as AIGenerateTextResult,
3
4
  StreamTextResult as AIStreamTextResult,
4
5
  CallSettings,
@@ -122,8 +123,8 @@ export type {
122
123
  TextStreamPart
123
124
  } from 'ai';
124
125
 
125
- // Re-export the tool helper function, Output, smoothStream, and stop condition helpers
126
- export { tool, Output, smoothStream, stepCountIs, hasToolCall } from 'ai';
126
+ // Re-export the tool helper function, Output, smoothStream, stop condition helpers, and jsonSchema
127
+ export { tool, Output, smoothStream, stepCountIs, hasToolCall, jsonSchema } from 'ai';
127
128
 
128
129
  // Web search tool factories
129
130
  export { tavilySearch, tavilyExtract, tavilyCrawl, tavilyMap } from '@tavily/ai-sdk';
@@ -196,6 +197,21 @@ type StreamTextAiSdkOptions<
196
197
  experimental_transform?: StreamTextTransform<Tools> | Array<StreamTextTransform<Tools>>;
197
198
  };
198
199
 
200
+ /**
201
+ * Like {@link StreamTextAiSdkOptions} but `onFinish` receives {@link WrappedStreamTextOnFinishEvent} (adds `cost`).
202
+ */
203
+ type OutputStreamTextAiSdkOptions<
204
+ Tools extends ToolSet = ToolSet,
205
+ Output extends AIOutput<unknown, unknown> = AIOutput<unknown, unknown>
206
+ > = Omit<StreamTextAiSdkOptions<Tools, Output>, 'onFinish'> & {
207
+ onFinish?: WrappedStreamTextOnFinishCallback<Tools>;
208
+ };
209
+
210
+ /** Agent {@link Agent.stream} options: same as AI SDK plus wrapped `onFinish` (adds `cost`). */
211
+ export type OutputAgentStreamParameters = Omit<AgentStreamParameters<never, ToolSet>, 'onFinish'> & {
212
+ onFinish?: WrappedStreamTextOnFinishCallback<ToolSet>;
213
+ };
214
+
199
215
  /** A source extracted from search tool results during multi-step LLM execution. */
200
216
  export type ExtractedSource = {
201
217
  type: 'source';
@@ -205,6 +221,29 @@ export type ExtractedSource = {
205
221
  title: string;
206
222
  };
207
223
 
224
+ /**
225
+ * Cost breakdown from the cost module (`calculateLLMCallCost`). `total` is null when pricing data is missing or calculation fails.
226
+ */
227
+ export type LLMCallCost = {
228
+ total: number | null;
229
+ components?: Array<{
230
+ name: string,
231
+ value: number
232
+ }>;
233
+ message?: string;
234
+ };
235
+
236
+ /**
237
+ * `streamText` and agent `stream` `onFinish` event after the stream response wrapper: same as the AI SDK
238
+ * finish payload plus optional `cost` from pricing.
239
+ */
240
+ export type WrappedStreamTextOnFinishEvent<Tools extends ToolSet = ToolSet> =
241
+ Parameters<StreamTextOnFinishCallback<Tools>>[0] & { cost?: LLMCallCost };
242
+
243
+ export type WrappedStreamTextOnFinishCallback<Tools extends ToolSet = ToolSet> = (
244
+ event: WrappedStreamTextOnFinishEvent<Tools>
245
+ ) => void | PromiseLike<void>;
246
+
208
247
  /**
209
248
  * Result from generateText including full AI SDK response metadata.
210
249
  * Extends AI SDK's GenerateTextResult with a unified `result` field.
@@ -216,6 +255,8 @@ export type GenerateTextResult<
216
255
  > = AIGenerateTextResult<Tools, Output> & {
217
256
  /** Unified field name alias for 'text' - provides consistency across all generate* functions */
218
257
  result: string;
258
+ /** Calculated cost in USD for the LLM call (present after wrapping; `total` may be null if pricing is unavailable) */
259
+ cost?: LLMCallCost;
219
260
  /** Sources extracted from search tool results, merged with any native provider sources */
220
261
  sources: ExtractedSource[];
221
262
  };
@@ -300,7 +341,7 @@ export function generateText<
300
341
  * @param args.prompt - Prompt file name.
301
342
  * @param args.variables - Variables to interpolate.
302
343
  * @param args.onChunk - Callback for each stream chunk (optional).
303
- * @param args.onFinish - Callback when stream finishes (optional).
344
+ * @param args.onFinish - Callback when stream finishes; receives {@link WrappedStreamTextOnFinishEvent} (`cost` optional).
304
345
  * @param args.onError - Callback when stream errors (optional).
305
346
  * @returns AI SDK stream result with textStream, fullStream, and metadata promises.
306
347
  */
@@ -311,7 +352,7 @@ export function streamText<
311
352
  args: {
312
353
  prompt: string,
313
354
  variables?: Record<string, string | number | boolean>
314
- } & StreamTextAiSdkOptions<Tools, Output>
355
+ } & OutputStreamTextAiSdkOptions<Tools, Output>
315
356
  ): AIStreamTextResult<Tools, Output>;
316
357
 
317
358
  export { skill } from './skill.js';
@@ -385,18 +426,21 @@ export declare class Agent extends import( 'ai' ).ToolLoopAgent {
385
426
  maxRetries?: number;
386
427
  } );
387
428
 
388
- /** Run the agent and return when complete. */
429
+ /**
430
+ * Run the agent and return when complete.
431
+ * Same augmented shape as {@link generateText}: `result`, optional `cost`, merged `sources`.
432
+ */
389
433
  generate( options?: {
390
434
  messages?: import( 'ai' ).ModelMessage[];
391
435
  abortSignal?: AbortSignal;
392
436
  onStepFinish?: import( 'ai' ).GenerateTextOnStepFinishCallback<ToolSet>;
393
- } ): Promise<import( 'ai' ).GenerateTextResult<ToolSet, import( 'ai' ).Output<unknown, unknown>>>;
437
+ } ): Promise<GenerateTextResult<ToolSet, import( 'ai' ).Output<unknown, unknown>>>;
394
438
 
395
- /** Stream the agent's response. */
396
- stream( options?: {
397
- messages?: import( 'ai' ).ModelMessage[];
398
- abortSignal?: AbortSignal;
399
- onStepFinish?: import( 'ai' ).StreamTextOnStepFinishCallback<ToolSet>;
400
- experimental_transform?: import( 'ai' ).StreamTextTransform<ToolSet> | import( 'ai' ).StreamTextTransform<ToolSet>[];
401
- } ): Promise<import( 'ai' ).StreamTextResult<ToolSet, import( 'ai' ).Output<unknown, unknown>>>;
439
+ /**
440
+ * Stream the agent's response.
441
+ * `onFinish` receives {@link WrappedStreamTextOnFinishEvent} (`cost` optional), matching {@link streamText}.
442
+ */
443
+ stream( options?: OutputAgentStreamParameters ): Promise<
444
+ AIStreamTextResult<ToolSet, import( 'ai' ).Output<unknown, unknown>>
445
+ >;
402
446
  };
package/src/index.js CHANGED
@@ -5,5 +5,5 @@ export { registerProvider, getRegisteredProviders } from './ai_model.js';
5
5
  export { tavilySearch, tavilyExtract, tavilyCrawl, tavilyMap } from '@tavily/ai-sdk';
6
6
  export { webSearch as exaSearch } from '@exalabs/ai-sdk';
7
7
  export { perplexitySearch } from '@perplexity-ai/ai-sdk';
8
- export { tool, Output, smoothStream, stepCountIs, hasToolCall } from 'ai';
8
+ export { tool, Output, smoothStream, stepCountIs, hasToolCall, jsonSchema } from 'ai';
9
9
  export * as ai from 'ai';
@@ -0,0 +1,29 @@
1
+ import { describe, it, expect } from 'vitest';
2
+ import { ROLE, getContent, isRole } from './message.js';
3
+
4
+ describe( 'message utils', () => {
5
+ describe( 'ROLE', () => {
6
+ it( 'exposes expected role string constants', () => {
7
+ expect( ROLE ).toEqual( {
8
+ SYSTEM: 'system',
9
+ USER: 'user',
10
+ ASSISTANT: 'assistant',
11
+ TOOL: 'tool'
12
+ } );
13
+ } );
14
+ } );
15
+
16
+ describe( 'isRole', () => {
17
+ it( 'returns a predicate that matches messages by role', () => {
18
+ const isUser = isRole( ROLE.USER );
19
+ expect( isUser( { role: 'user', content: 'hi' } ) ).toBe( true );
20
+ expect( isUser( { role: 'assistant', content: 'bye' } ) ).toBe( false );
21
+ } );
22
+ } );
23
+
24
+ describe( 'getContent', () => {
25
+ it( 'returns message content', () => {
26
+ expect( getContent( { role: 'user', content: 'hello' } ) ).toBe( 'hello' );
27
+ } );
28
+ } );
29
+ } );
@@ -0,0 +1,70 @@
1
+ import { combineSources, extractSourcesFromSteps } from './source_extraction.js';
2
+ import { calculateLLMCallCost } from '../cost/index.js';
3
+ import { endTraceWithSuccess } from './trace.js';
4
+
5
+ /**
6
+ * Calculates the cost and wraps an AI SDK text response in a Proxy with shortcut for 'result' and 'cost'
7
+ *
8
+ * Emits the `cost:llm:request` event.
9
+ *
10
+ * Also finishes the trace events.
11
+ *
12
+ * @param {object} args
13
+ * @param {string} args.traceId - id created by the startTrace
14
+ * @param {string} args.modelId - id of the model used
15
+ * @param {object} args.response - AI SDK's text response
16
+ * @returns {object} Proxied response
17
+ */
18
+ export const wrapTextResponse = async ( { traceId, modelId, response } ) => {
19
+ const sourcesFromTools = extractSourcesFromSteps( response.steps );
20
+ const cost = await calculateLLMCallCost( { usage: response.totalUsage, modelId } );
21
+
22
+ endTraceWithSuccess( { traceId, modelId, response, cost, sourcesFromTools } );
23
+
24
+ return new Proxy( response, {
25
+ get( target, prop, receiver ) {
26
+ if ( prop === 'result' ) {
27
+ return target.text;
28
+ }
29
+ if ( prop === 'cost' ) {
30
+ return cost;
31
+ }
32
+ if ( prop === 'sources' && sourcesFromTools.length > 0 ) {
33
+ return combineSources( { sourcesFromTools, sourcesFromResponse: response.sources } );
34
+ }
35
+ return Reflect.get( target, prop, receiver );
36
+ }
37
+ } );
38
+ };
39
+
40
+ /**
41
+ * Wraps the response returned by the onFinish callback from the stream.
42
+ *
43
+ * When the onFinish is triggered, concludes the trace event, calculates cost and emits `cost:llm:request`.
44
+ * Returns a proxy around the response with `cost` property.
45
+ *
46
+ * @param {object} args
47
+ * @param {string} args.traceId - id created by the startTrace
48
+ * @param {string} args.modelId - id of the model used
49
+ * @param {Function} args.onFinish - Original callback to call with the Proxied reponse
50
+ * @returns {object} Proxied response
51
+ */
52
+ export const wrapStreamOnFinishResponse = ( { traceId, modelId, onFinish: _onFinish } ) => ( {
53
+ async onFinish( response ) {
54
+ const cost = await calculateLLMCallCost( { modelId, usage: response.totalUsage } );
55
+
56
+ endTraceWithSuccess( { traceId, modelId, response, cost } );
57
+
58
+ _onFinish?.( new Proxy( response, {
59
+ get( target, prop, receiver ) {
60
+ if ( prop === 'result' ) {
61
+ return target.text;
62
+ }
63
+ if ( prop === 'cost' ) {
64
+ return cost;
65
+ }
66
+ return Reflect.get( target, prop, receiver );
67
+ }
68
+ } ) );
69
+ }
70
+ } );
@@ -0,0 +1,172 @@
1
+ import { describe, it, expect, vi, beforeEach } from 'vitest';
2
+
3
+ const mocks = vi.hoisted( () => ( {
4
+ extractSourcesFromSteps: vi.fn(),
5
+ calculateLLMCallCost: vi.fn(),
6
+ endTraceWithSuccess: vi.fn()
7
+ } ) );
8
+
9
+ vi.mock( './source_extraction.js', async importOriginal => {
10
+ const actual = await importOriginal();
11
+ return {
12
+ ...actual,
13
+ extractSourcesFromSteps: mocks.extractSourcesFromSteps
14
+ };
15
+ } );
16
+
17
+ vi.mock( '../cost/index.js', () => ( {
18
+ calculateLLMCallCost: mocks.calculateLLMCallCost
19
+ } ) );
20
+
21
+ vi.mock( './trace.js', () => ( {
22
+ endTraceWithSuccess: mocks.endTraceWithSuccess
23
+ } ) );
24
+
25
+ import { wrapTextResponse, wrapStreamOnFinishResponse } from './response_wrappers.js';
26
+
27
+ describe( 'wrapTextResponse', () => {
28
+ const traceId = 'trace-1';
29
+ const modelId = 'test-model';
30
+ const mockCost = { total: 0.001, components: [ { name: 'input_tokens', value: 0.001 } ] };
31
+
32
+ beforeEach( () => {
33
+ vi.clearAllMocks();
34
+ mocks.extractSourcesFromSteps.mockReturnValue( [] );
35
+ mocks.calculateLLMCallCost.mockResolvedValue( mockCost );
36
+ } );
37
+
38
+ it( 'proxies result to text and attaches cost', async () => {
39
+ const response = {
40
+ text: 'hello',
41
+ totalUsage: { inputTokens: 10, outputTokens: 5 },
42
+ steps: [],
43
+ sources: []
44
+ };
45
+
46
+ const wrapped = await wrapTextResponse( { traceId, modelId, response } );
47
+
48
+ expect( wrapped.result ).toBe( 'hello' );
49
+ expect( wrapped.text ).toBe( 'hello' );
50
+ expect( wrapped.cost ).toEqual( mockCost );
51
+ expect( mocks.calculateLLMCallCost ).toHaveBeenCalledWith( {
52
+ usage: response.totalUsage,
53
+ modelId
54
+ } );
55
+ expect( mocks.endTraceWithSuccess ).toHaveBeenCalledWith( {
56
+ traceId,
57
+ modelId,
58
+ response,
59
+ cost: mockCost,
60
+ sourcesFromTools: []
61
+ } );
62
+ } );
63
+
64
+ it( 'leaves sources unchanged when no tool-extracted sources (same as raw response)', async () => {
65
+ const nativeSources = [
66
+ { type: 'source', sourceType: 'url', id: 'n1', url: 'https://native.test', title: 'Native' }
67
+ ];
68
+ mocks.extractSourcesFromSteps.mockReturnValue( [] );
69
+
70
+ const response = {
71
+ text: 'x',
72
+ totalUsage: {},
73
+ steps: [],
74
+ sources: nativeSources
75
+ };
76
+
77
+ const wrapped = await wrapTextResponse( { traceId, modelId, response } );
78
+
79
+ expect( wrapped.sources ).toBe( nativeSources );
80
+ } );
81
+
82
+ it( 'merges tool-extracted sources with response sources by url when tool sources exist', async () => {
83
+ const toolSource = {
84
+ type: 'source',
85
+ sourceType: 'url',
86
+ id: 'id-a',
87
+ url: 'https://example.com/a',
88
+ title: 'A'
89
+ };
90
+ const responseSource = {
91
+ type: 'source',
92
+ sourceType: 'url',
93
+ id: 'id-b',
94
+ url: 'https://example.com/b',
95
+ title: 'B'
96
+ };
97
+ mocks.extractSourcesFromSteps.mockReturnValue( [ toolSource ] );
98
+
99
+ const response = {
100
+ text: 'x',
101
+ totalUsage: {},
102
+ steps: [ {} ],
103
+ sources: [ responseSource ]
104
+ };
105
+
106
+ const wrapped = await wrapTextResponse( { traceId, modelId, response } );
107
+
108
+ expect( wrapped.sources ).toEqual( expect.arrayContaining( [ toolSource, responseSource ] ) );
109
+ expect( wrapped.sources ).toHaveLength( 2 );
110
+ } );
111
+
112
+ it( 'when tool sources overlap urls, later entry in merge wins', async () => {
113
+ const url = 'https://example.com/same';
114
+ mocks.extractSourcesFromSteps.mockReturnValue( [
115
+ { type: 'source', sourceType: 'url', id: '1', url, title: 'from-tool' }
116
+ ] );
117
+
118
+ const response = {
119
+ text: 'x',
120
+ totalUsage: {},
121
+ steps: [],
122
+ sources: [
123
+ { type: 'source', sourceType: 'url', id: '2', url, title: 'from-response' }
124
+ ]
125
+ };
126
+
127
+ const wrapped = await wrapTextResponse( { traceId, modelId, response } );
128
+
129
+ expect( wrapped.sources ).toHaveLength( 1 );
130
+ expect( wrapped.sources[0].title ).toBe( 'from-response' );
131
+ } );
132
+ } );
133
+
134
+ describe( 'wrapStreamOnFinishResponse', () => {
135
+ const traceId = 'stream-trace';
136
+ const modelId = 'stream-model';
137
+ const mockCost = { total: 0.002, components: [] };
138
+
139
+ beforeEach( () => {
140
+ vi.clearAllMocks();
141
+ mocks.calculateLLMCallCost.mockResolvedValue( mockCost );
142
+ } );
143
+
144
+ it( 'onFinish finishes trace, proxies result and cost for user callback, and forwards other props', async () => {
145
+ const userOnFinish = vi.fn();
146
+ const response = {
147
+ text: 'done',
148
+ totalUsage: { inputTokens: 1 },
149
+ finishReason: 'stop'
150
+ };
151
+
152
+ const callbacks = wrapStreamOnFinishResponse( {
153
+ traceId,
154
+ modelId,
155
+ onFinish: userOnFinish
156
+ } );
157
+
158
+ await callbacks.onFinish( response );
159
+
160
+ expect( mocks.endTraceWithSuccess ).toHaveBeenCalledWith( {
161
+ traceId,
162
+ modelId,
163
+ response,
164
+ cost: mockCost
165
+ } );
166
+ expect( userOnFinish ).toHaveBeenCalledTimes( 1 );
167
+ const proxied = userOnFinish.mock.calls[0][0];
168
+ expect( proxied.result ).toBe( 'done' );
169
+ expect( proxied.cost ).toEqual( mockCost );
170
+ expect( proxied.finishReason ).toBe( 'stop' );
171
+ } );
172
+ } );
@@ -47,3 +47,17 @@ export function extractSourcesFromSteps( steps ) {
47
47
  return [];
48
48
  }
49
49
  }
50
+
51
+ /**
52
+ * Merge sources used tools and sources from AI SDK response into a single list
53
+ *
54
+ * Deduplicate by url (prefer to keep items from sources from response).
55
+ *
56
+ * @param {object} args
57
+ * @param {object[]} args.sourcesFromTools
58
+ * @param {object[]} args.sourcesFromResponse
59
+ * @returns {object[]} Merged sources
60
+ */
61
+ export const combineSources = ( { sourcesFromTools, sourcesFromResponse } ) =>
62
+ new Map( sourcesFromTools.concat( Array.isArray( sourcesFromResponse ) ? sourcesFromResponse : [] )
63
+ .map( s => [ s.url, s ] ) ).values().toArray();
@@ -1,5 +1,5 @@
1
1
  import { describe, it, expect } from 'vitest';
2
- import { extractSourcesFromSteps } from './source_extraction.js';
2
+ import { combineSources, extractSourcesFromSteps } from './source_extraction.js';
3
3
 
4
4
  describe( 'extractSourcesFromSteps', () => {
5
5
  it( 'returns empty array for undefined/null/empty steps', () => {
@@ -167,3 +167,28 @@ describe( 'extractSourcesFromSteps', () => {
167
167
  expect( s1[0].id ).toHaveLength( 16 );
168
168
  } );
169
169
  } );
170
+
171
+ describe( 'combineSources', () => {
172
+ it( 'treats non-array sourcesFromResponse like empty array', () => {
173
+ const tool = [ { url: 'https://u.test', title: 'T', type: 'source', sourceType: 'url', id: '1' } ];
174
+
175
+ expect(
176
+ combineSources( { sourcesFromTools: tool, sourcesFromResponse: undefined } )
177
+ ).toHaveLength( 1 );
178
+
179
+ expect(
180
+ combineSources( { sourcesFromTools: tool, sourcesFromResponse: {} } )
181
+ ).toHaveLength( 1 );
182
+ } );
183
+
184
+ it( 'dedupes by url with response sources winning after tools', () => {
185
+ const url = 'https://shared.test';
186
+ const tools = [ { url, title: 'from-tool', type: 'source', sourceType: 'url', id: 'a' } ];
187
+ const response = [ { url, title: 'from-response', type: 'source', sourceType: 'url', id: 'b' } ];
188
+
189
+ const merged = combineSources( { sourcesFromTools: tools, sourcesFromResponse: response } );
190
+
191
+ expect( merged ).toHaveLength( 1 );
192
+ expect( merged[0].title ).toBe( 'from-response' );
193
+ } );
194
+ } );
@@ -0,0 +1,18 @@
1
+ import { Tracing, emitEvent } from '@outputai/core/sdk_activity_integration';
2
+
3
+ export const startTrace = ( { name, ...details } ) => {
4
+ const traceId = `${name}-${Date.now()}`;
5
+ Tracing.addEventStart( { kind: 'llm', name, id: traceId, details } );
6
+ return traceId;
7
+ };
8
+
9
+ export const endTraceWithError = ( { traceId, error } ) => {
10
+ Tracing.addEventError( { id: traceId, details: error } );
11
+ };
12
+
13
+ export const endTraceWithSuccess = ( { traceId, modelId, response, cost, ...extra } ) => {
14
+ const { totalUsage: usage, text: result, providerMetadata } = response;
15
+ Tracing.addEventAttribute( { eventId: traceId, name: Tracing.Attribute.COST, value: cost } );
16
+ Tracing.addEventEnd( { id: traceId, details: { result, usage, providerMetadata, ...extra } } );
17
+ emitEvent( 'cost:llm:request', { modelId, cost, usage } );
18
+ };
@@ -0,0 +1,95 @@
1
+ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
2
+
3
+ vi.mock( '@outputai/core/sdk_activity_integration', () => ( {
4
+ Tracing: {
5
+ addEventStart: vi.fn(),
6
+ addEventError: vi.fn(),
7
+ addEventAttribute: vi.fn(),
8
+ addEventEnd: vi.fn(),
9
+ Attribute: {
10
+ COST: 'cost'
11
+ }
12
+ },
13
+ emitEvent: vi.fn()
14
+ } ) );
15
+
16
+ import { Tracing, emitEvent } from '@outputai/core/sdk_activity_integration';
17
+ import { startTrace, endTraceWithError, endTraceWithSuccess } from './trace.js';
18
+
19
+ const tracing = vi.mocked( Tracing, true );
20
+
21
+ describe( 'trace utils', () => {
22
+ beforeEach( () => {
23
+ vi.clearAllMocks();
24
+ } );
25
+
26
+ afterEach( () => {
27
+ vi.restoreAllMocks();
28
+ } );
29
+
30
+ describe( 'startTrace', () => {
31
+ it( 'starts an llm trace with name-based id and passes remaining fields as details', () => {
32
+ vi.spyOn( Date, 'now' ).mockReturnValue( 9_000_000_000 );
33
+
34
+ const traceId = startTrace( { name: 'generateText', prompt: 'p', variables: { k: 1 } } );
35
+
36
+ expect( traceId ).toBe( 'generateText-9000000000' );
37
+ expect( tracing.addEventStart ).toHaveBeenCalledWith( {
38
+ kind: 'llm',
39
+ name: 'generateText',
40
+ id: 'generateText-9000000000',
41
+ details: { prompt: 'p', variables: { k: 1 } }
42
+ } );
43
+ } );
44
+ } );
45
+
46
+ describe( 'endTraceWithError', () => {
47
+ it( 'records an error on the trace event', () => {
48
+ const err = new Error( 'failed' );
49
+
50
+ endTraceWithError( { traceId: 't-1', error: err } );
51
+
52
+ expect( tracing.addEventError ).toHaveBeenCalledWith( { id: 't-1', details: err } );
53
+ } );
54
+ } );
55
+
56
+ describe( 'endTraceWithSuccess', () => {
57
+ it( 'adds cost attribute, ends the trace with response fields and extra details, and emits cost:llm:request', () => {
58
+ const cost = { total: 0.01, components: [] };
59
+ const usage = { inputTokens: 2, outputTokens: 3 };
60
+ const response = {
61
+ text: 'hello',
62
+ totalUsage: usage,
63
+ providerMetadata: { provider: 'x' }
64
+ };
65
+
66
+ endTraceWithSuccess( {
67
+ traceId: 'trace-a',
68
+ modelId: 'my-model',
69
+ response,
70
+ cost,
71
+ sourcesFromTools: [ { url: 'https://u.test', title: '' } ]
72
+ } );
73
+
74
+ expect( tracing.addEventAttribute ).toHaveBeenCalledWith( {
75
+ eventId: 'trace-a',
76
+ name: 'cost',
77
+ value: cost
78
+ } );
79
+ expect( tracing.addEventEnd ).toHaveBeenCalledWith( {
80
+ id: 'trace-a',
81
+ details: {
82
+ result: 'hello',
83
+ usage,
84
+ providerMetadata: { provider: 'x' },
85
+ sourcesFromTools: [ { url: 'https://u.test', title: '' } ]
86
+ }
87
+ } );
88
+ expect( emitEvent ).toHaveBeenCalledWith( 'cost:llm:request', {
89
+ modelId: 'my-model',
90
+ cost,
91
+ usage
92
+ } );
93
+ } );
94
+ } );
95
+ } );
@@ -1,21 +0,0 @@
1
- import { extractSourcesFromSteps } from './source_extraction.js';
2
- import { endTraceWithSuccess } from './trace_utils.js';
3
-
4
- export const wrapInOutputResponse = async ( response, { traceId, modelId } ) => {
5
- const sourcesFromTools = extractSourcesFromSteps( response.steps );
6
- await endTraceWithSuccess( traceId, modelId, response, { sourcesFromTools } );
7
-
8
- return new Proxy( response, {
9
- get( target, prop, receiver ) {
10
- if ( prop === 'result' ) {
11
- return target.text;
12
- }
13
- if ( prop === 'sources' && sourcesFromTools.length > 0 ) {
14
- const responseSources = Array.isArray( target[prop] ) ? target[prop] : [];
15
- const byUrl = new Map( [ ...sourcesFromTools, ...responseSources ].map( s => [ s.url, s ] ) );
16
- return [ ...byUrl.values() ];
17
- }
18
- return Reflect.get( target, prop, receiver );
19
- }
20
- } );
21
- };
@@ -1,30 +0,0 @@
1
- import { Tracing, emitEvent } from '@outputai/core/sdk_activity_integration';
2
- import { calculateLLMCallCost } from './cost/index.js';
3
-
4
- export const startTrace = ( name, details ) => {
5
- const traceId = `${name}-${Date.now()}`;
6
- Tracing.addEventStart( { kind: 'llm', name, id: traceId, details } );
7
- return traceId;
8
- };
9
-
10
- export const endTraceWithError = ( traceId, error ) => {
11
- Tracing.addEventError( { id: traceId, details: error } );
12
- };
13
-
14
- export const endTraceWithSuccess = async ( traceId, modelId, response, extraDetails = {} ) => {
15
- const { text: result, totalUsage: usage, providerMetadata } = response;
16
- const cost = await calculateLLMCallCost( { usage, modelId } );
17
- emitEvent( 'llm:call_cost', { modelId, cost, usage } );
18
- Tracing.addEventEnd( { id: traceId, details: { result, usage, cost, providerMetadata, ...extraDetails } } );
19
- };
20
-
21
- export const traceStreamCallbacks = ( traceId, modelId, { onFinish: userOnFinish, onError: userOnError } = {} ) => ( {
22
- async onFinish( response ) {
23
- await endTraceWithSuccess( traceId, modelId, response );
24
- userOnFinish?.( response );
25
- },
26
- onError( event ) {
27
- Tracing.addEventError( { id: traceId, details: event.error } );
28
- userOnError?.( event );
29
- }
30
- } );
File without changes