@telemetry-dev/bedrock 0.1.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.
@@ -0,0 +1,390 @@
1
+ import type { SpanFields } from "@telemetry-dev/sdk";
2
+
3
+ import {
4
+ arrayValue,
5
+ awsMetadataFields,
6
+ isRecord,
7
+ mergeFields,
8
+ modeledStreamError,
9
+ numberValue,
10
+ omitUndefined,
11
+ parseJson,
12
+ PROVIDER,
13
+ stringValue,
14
+ type StreamState,
15
+ } from "./internal.ts";
16
+
17
+ export type MessagePart = Record<string, unknown>;
18
+ export interface NormalizedMessage {
19
+ role: string | undefined;
20
+ parts: MessagePart[];
21
+ finish_reason?: string;
22
+ }
23
+
24
+ export function converseRequestFields(input: Record<string, unknown>): SpanFields {
25
+ const inference = isRecord(input.inferenceConfig) ? input.inferenceConfig : {};
26
+ const guardrail = isRecord(input.guardrailConfig) ? input.guardrailConfig : undefined;
27
+ return omitUndefined({
28
+ provider: PROVIDER,
29
+ model: stringValue(input.modelId),
30
+ input: normalizeMessages(input.messages),
31
+ systemInstructions: normalizeSystem(input.system),
32
+ temperature: numberValue(inference.temperature),
33
+ topP: numberValue(inference.topP),
34
+ maxTokens: numberValue(inference.maxTokens),
35
+ stopSequences: arrayValue<string>(inference.stopSequences),
36
+ metadata: guardrail
37
+ ? omitUndefined({
38
+ guardrail_id: stringValue(guardrail.guardrailIdentifier),
39
+ guardrail_version: stringValue(guardrail.guardrailVersion),
40
+ })
41
+ : undefined,
42
+ });
43
+ }
44
+
45
+ export function converseResponseFields(output: unknown): SpanFields {
46
+ const result = isRecord(output) ? output : {};
47
+ const usage = usageFromConverse(result.usage);
48
+ const metrics = isRecord(result.metrics) ? result.metrics : undefined;
49
+ const trace = isRecord(result.trace) ? result.trace : undefined;
50
+ const promptRouter = isRecord(trace?.promptRouter) ? trace?.promptRouter : undefined;
51
+ const guardrail = isRecord(trace?.guardrail) ? trace?.guardrail : undefined;
52
+ const stopReason = stringValue(result.stopReason);
53
+ const message = isRecord(result.output) ? result.output.message : undefined;
54
+ return mergeFields(awsMetadataFields(result.$metadata), {
55
+ output: normalizeOutputMessage(message, stopReason),
56
+ usage,
57
+ finishReason: stopReason,
58
+ responseModel: stringValue(promptRouter?.invokedModelId),
59
+ metadata: omitUndefined({
60
+ server_latency_ms: numberValue(metrics?.latencyMs),
61
+ guardrail_action: stringValue(result.guardrailAction) ?? stringValue(guardrail?.action),
62
+ guardrail_action_reason: stringValue(guardrail?.actionReason),
63
+ }),
64
+ });
65
+ }
66
+
67
+ export function usageFromConverse(value: unknown): SpanFields["usage"] | undefined {
68
+ if (!isRecord(value)) return undefined;
69
+ return omitUndefined({
70
+ inputTokens: numberValue(value.inputTokens),
71
+ outputTokens: numberValue(value.outputTokens),
72
+ totalTokens: numberValue(value.totalTokens),
73
+ cacheReadInputTokens: numberValue(value.cacheReadInputTokens),
74
+ cacheCreationInputTokens: numberValue(value.cacheWriteInputTokens),
75
+ });
76
+ }
77
+
78
+ export function normalizeMessages(value: unknown): NormalizedMessage[] | undefined {
79
+ const messages = arrayValue<Record<string, unknown>>(value);
80
+ if (!messages) return undefined;
81
+ return messages.map((message) => ({
82
+ role: stringValue(message.role),
83
+ parts: normalizeContentList(message.content),
84
+ }));
85
+ }
86
+
87
+ export function normalizeOutputMessage(
88
+ value: unknown,
89
+ finishReason?: string,
90
+ ): NormalizedMessage[] | undefined {
91
+ if (!isRecord(value)) return undefined;
92
+ return [
93
+ omitUndefined({
94
+ role: stringValue(value.role),
95
+ parts: normalizeContentList(value.content),
96
+ finish_reason: finishReason,
97
+ }) as NormalizedMessage,
98
+ ];
99
+ }
100
+
101
+ export function normalizeSystem(value: unknown): MessagePart[] | undefined {
102
+ const blocks = arrayValue<Record<string, unknown>>(value);
103
+ if (!blocks) return undefined;
104
+ const parts = blocks.flatMap((block) => normalizeContentBlock(block));
105
+ return parts.length > 0 ? parts : undefined;
106
+ }
107
+
108
+ function normalizeContentList(value: unknown): MessagePart[] {
109
+ const blocks = arrayValue<Record<string, unknown>>(value);
110
+ if (!blocks) return [];
111
+ return blocks.flatMap((block) => normalizeContentBlock(block));
112
+ }
113
+
114
+ function normalizeContentBlock(block: Record<string, unknown>): MessagePart[] {
115
+ const text = stringValue(block.text);
116
+ if (text !== undefined)
117
+ return [
118
+ omitUndefined({
119
+ type: "text",
120
+ content: text,
121
+ citations: arrayValue<unknown>(block.citations),
122
+ }),
123
+ ];
124
+
125
+ const toolUse = isRecord(block.toolUse) ? block.toolUse : undefined;
126
+ if (toolUse) {
127
+ return [
128
+ omitUndefined({
129
+ type: "tool_call",
130
+ id: stringValue(toolUse.toolUseId),
131
+ name: stringValue(toolUse.name),
132
+ arguments: toolUse.input,
133
+ }),
134
+ ];
135
+ }
136
+
137
+ const toolResult = isRecord(block.toolResult) ? block.toolResult : undefined;
138
+ if (toolResult) {
139
+ return [
140
+ omitUndefined({
141
+ type: "tool_call_response",
142
+ id: stringValue(toolResult.toolUseId),
143
+ response: simplifyToolResult(toolResult.content),
144
+ }),
145
+ ];
146
+ }
147
+
148
+ const reasoning = isRecord(block.reasoningContent) ? block.reasoningContent : undefined;
149
+ if (reasoning) {
150
+ const reasoningText = isRecord(reasoning.reasoningText) ? reasoning.reasoningText : undefined;
151
+ return [
152
+ {
153
+ type: "reasoning",
154
+ content:
155
+ stringValue(reasoningText?.text) ?? (reasoning.redactedContent ? "[redacted]" : ""),
156
+ },
157
+ ];
158
+ }
159
+
160
+ for (const modality of ["image", "document", "video", "audio"] as const) {
161
+ const media = isRecord(block[modality]) ? block[modality] : undefined;
162
+ if (!media) continue;
163
+ const format = stringValue(media.format);
164
+ const source = isRecord(media.source) ? media.source : undefined;
165
+ if (source?.bytes !== undefined) {
166
+ return [
167
+ omitUndefined({
168
+ type: "blob",
169
+ modality,
170
+ mime_type: format ? `${modality}/${format}` : undefined,
171
+ }),
172
+ ];
173
+ }
174
+ const s3 = isRecord(source?.s3Location) ? source?.s3Location : undefined;
175
+ const uri = stringValue(source?.uri) ?? stringValue(s3?.uri) ?? s3Uri(s3);
176
+ if (uri) return [omitUndefined({ type: "uri", uri, modality, mime_type: format })];
177
+ }
178
+
179
+ const citationsContent = isRecord(block.citationsContent) ? block.citationsContent : undefined;
180
+ if (citationsContent) {
181
+ const contentBlocks = arrayValue<Record<string, unknown>>(citationsContent.content) ?? [];
182
+ const citations = arrayValue<unknown>(citationsContent.citations);
183
+ return contentBlocks.flatMap((contentBlock) =>
184
+ normalizeContentBlock(citations ? { ...contentBlock, citations } : contentBlock),
185
+ );
186
+ }
187
+
188
+ const guardContent = isRecord(block.guardContent) ? block.guardContent : undefined;
189
+ if (guardContent) {
190
+ const guardText = isRecord(guardContent.text) ? guardContent.text : undefined;
191
+ const content = stringValue(guardText?.text) ?? stringValue(guardContent.text);
192
+ if (content) return [{ type: "text", content }];
193
+ }
194
+
195
+ try {
196
+ return [{ type: "text", content: JSON.stringify(block) }];
197
+ } catch {
198
+ return [{ type: "text", content: "[unsupported content block]" }];
199
+ }
200
+ }
201
+
202
+ function simplifyToolResult(content: unknown): unknown {
203
+ const blocks = arrayValue<Record<string, unknown>>(content);
204
+ if (!blocks) return content;
205
+ if (blocks.length === 1 && typeof blocks[0]?.text === "string") return blocks[0].text;
206
+ return blocks.flatMap((block) => normalizeContentBlock(block));
207
+ }
208
+
209
+ function s3Uri(value: Record<string, unknown> | undefined): string | undefined {
210
+ const bucket = stringValue(value?.bucket);
211
+ const key = stringValue(value?.key);
212
+ return bucket && key ? `s3://${bucket}/${key}` : undefined;
213
+ }
214
+
215
+ interface StreamBlock {
216
+ kind: "text" | "tool" | "reasoning" | "content";
217
+ text?: string;
218
+ toolUseId?: string;
219
+ name?: string;
220
+ input?: string;
221
+ content?: Record<string, unknown>;
222
+ citations?: unknown[];
223
+ }
224
+
225
+ export class ConverseStreamState implements StreamState {
226
+ private role = "assistant";
227
+ private readonly blocks = new Map<number, StreamBlock>();
228
+ private stopReason: string | undefined;
229
+ private usage: SpanFields["usage"];
230
+ private metadata: Record<string, unknown> = {};
231
+ private responseModel: string | undefined;
232
+ private error: Error | undefined;
233
+ private errorFields: SpanFields | undefined;
234
+
235
+ feed(event: unknown): void {
236
+ if (!isRecord(event)) return;
237
+ const streamError = modeledStreamError(event);
238
+ if (streamError) {
239
+ this.error = streamError.error;
240
+ this.errorFields = streamError.fields;
241
+ return;
242
+ }
243
+ const messageStart = isRecord(event.messageStart) ? event.messageStart : undefined;
244
+ if (messageStart) this.role = stringValue(messageStart.role) ?? this.role;
245
+
246
+ const blockStart = isRecord(event.contentBlockStart) ? event.contentBlockStart : undefined;
247
+ if (blockStart) {
248
+ const index = numberValue(blockStart.contentBlockIndex) ?? this.blocks.size;
249
+ const start = isRecord(blockStart.start) ? blockStart.start : undefined;
250
+ const toolUse = isRecord(start?.toolUse) ? start?.toolUse : undefined;
251
+ const toolResult = isRecord(start?.toolResult) ? start?.toolResult : undefined;
252
+ const image = isRecord(start?.image) ? start?.image : undefined;
253
+ this.blocks.set(
254
+ index,
255
+ toolUse
256
+ ? {
257
+ kind: "tool",
258
+ toolUseId: stringValue(toolUse.toolUseId),
259
+ name: stringValue(toolUse.name),
260
+ input: "",
261
+ }
262
+ : toolResult
263
+ ? {
264
+ kind: "content",
265
+ content: {
266
+ toolResult: omitUndefined({
267
+ toolUseId: stringValue(toolResult.toolUseId),
268
+ status: stringValue(toolResult.status),
269
+ content: [],
270
+ }),
271
+ },
272
+ }
273
+ : image
274
+ ? { kind: "content", content: { image } }
275
+ : { kind: "text", text: "" },
276
+ );
277
+ }
278
+
279
+ const deltaEvent = isRecord(event.contentBlockDelta) ? event.contentBlockDelta : undefined;
280
+ if (deltaEvent) {
281
+ const index = numberValue(deltaEvent.contentBlockIndex) ?? 0;
282
+ const delta = isRecord(deltaEvent.delta) ? deltaEvent.delta : {};
283
+ const block = this.blocks.get(index) ?? { kind: "text", text: "" };
284
+ if (typeof delta.text === "string") {
285
+ block.kind = "text";
286
+ block.text = `${block.text ?? ""}${delta.text}`;
287
+ }
288
+ const toolUse = isRecord(delta.toolUse) ? delta.toolUse : undefined;
289
+ if (toolUse && typeof toolUse.input === "string") {
290
+ block.kind = "tool";
291
+ block.input = `${block.input ?? ""}${toolUse.input}`;
292
+ }
293
+ const reasoning = isRecord(delta.reasoningContent) ? delta.reasoningContent : undefined;
294
+ if (reasoning) {
295
+ block.kind = "reasoning";
296
+ block.text = `${block.text ?? ""}${stringValue(reasoning.text) ?? (reasoning.redactedContent ? "[redacted]" : "")}`;
297
+ }
298
+ const image = isRecord(delta.image) ? delta.image : undefined;
299
+ if (image) {
300
+ const content = isRecord(block.content) ? block.content : {};
301
+ const current = isRecord(content.image) ? content.image : {};
302
+ block.kind = "content";
303
+ block.content = { image: { ...current, ...image } };
304
+ }
305
+ const toolResult = arrayValue<Record<string, unknown>>(delta.toolResult);
306
+ if (toolResult) {
307
+ const content = isRecord(block.content) ? block.content : {};
308
+ const current = isRecord(content.toolResult) ? content.toolResult : {};
309
+ const currentContent = arrayValue<unknown>(current.content) ?? [];
310
+ block.kind = "content";
311
+ block.content = {
312
+ toolResult: {
313
+ ...current,
314
+ content: [...currentContent, ...toolResult],
315
+ },
316
+ };
317
+ }
318
+ if (isRecord(delta.citation)) {
319
+ block.citations = [...(block.citations ?? []), delta.citation];
320
+ }
321
+ this.blocks.set(index, block);
322
+ }
323
+
324
+ const stop = isRecord(event.messageStop) ? event.messageStop : undefined;
325
+ if (stop) this.stopReason = stringValue(stop.stopReason) ?? this.stopReason;
326
+
327
+ const metadata = isRecord(event.metadata) ? event.metadata : undefined;
328
+ if (metadata) {
329
+ this.usage = usageFromConverse(metadata.usage) ?? this.usage;
330
+ const metrics = isRecord(metadata.metrics) ? metadata.metrics : undefined;
331
+ const trace = isRecord(metadata.trace) ? metadata.trace : undefined;
332
+ const promptRouter = isRecord(trace?.promptRouter) ? trace?.promptRouter : undefined;
333
+ const guardrail = isRecord(trace?.guardrail) ? trace?.guardrail : undefined;
334
+ this.responseModel = stringValue(promptRouter?.invokedModelId) ?? this.responseModel;
335
+ this.metadata = {
336
+ ...this.metadata,
337
+ ...omitUndefined({
338
+ server_latency_ms: numberValue(metrics?.latencyMs),
339
+ guardrail_action: stringValue(guardrail?.action),
340
+ guardrail_action_reason: stringValue(guardrail?.actionReason),
341
+ }),
342
+ };
343
+ }
344
+ }
345
+
346
+ finish(partial: boolean): SpanFields {
347
+ const parts = [...this.blocks.entries()]
348
+ .sort(([left], [right]) => left - right)
349
+ .map(([, block]) => blockToPart(block));
350
+ return mergeFields(
351
+ omitUndefined({
352
+ output:
353
+ parts.length > 0
354
+ ? [
355
+ omitUndefined({
356
+ role: this.role,
357
+ parts,
358
+ finish_reason: partial ? undefined : this.stopReason,
359
+ }),
360
+ ]
361
+ : undefined,
362
+ finishReason: partial ? undefined : this.stopReason,
363
+ usage: this.usage,
364
+ responseModel: this.responseModel,
365
+ metadata: Object.keys(this.metadata).length > 0 ? this.metadata : undefined,
366
+ error: this.error,
367
+ }),
368
+ this.errorFields,
369
+ );
370
+ }
371
+ }
372
+
373
+ function blockToPart(block: StreamBlock): MessagePart {
374
+ if (block.kind === "tool") {
375
+ return omitUndefined({
376
+ type: "tool_call",
377
+ id: block.toolUseId,
378
+ name: block.name,
379
+ arguments: parseJson(block.input) ?? block.input ?? "",
380
+ });
381
+ }
382
+ if (block.kind === "reasoning") return { type: "reasoning", content: block.text ?? "" };
383
+ if (block.kind === "content" && block.content)
384
+ return normalizeContentBlock(block.content)[0] ?? {};
385
+ return omitUndefined({
386
+ type: "text",
387
+ content: block.text ?? "",
388
+ citations: block.citations && block.citations.length > 0 ? block.citations : undefined,
389
+ });
390
+ }
package/src/index.ts ADDED
@@ -0,0 +1,69 @@
1
+ import { BedrockRuntimeClient } from "@aws-sdk/client-bedrock-runtime";
2
+
3
+ import { RUNTIME_HANDLERS } from "./internal-handlers.ts";
4
+ import {
5
+ createInstrumentedSend,
6
+ type BedrockInstrumentationOptions,
7
+ type SendFn,
8
+ } from "./internal.ts";
9
+
10
+ export type { BedrockInstrumentationOptions } from "./internal.ts";
11
+
12
+ export interface WrappableBedrockClient {
13
+ send: (...args: never[]) => unknown;
14
+ }
15
+
16
+ const WRAPPED = new WeakSet<object>();
17
+ let ORIGINAL: { proto: object; hadOwn: boolean; value: unknown; wrapper: SendFn } | undefined;
18
+
19
+ export function wrapBedrock<T extends WrappableBedrockClient>(
20
+ client: T,
21
+ options: BedrockInstrumentationOptions = {},
22
+ ): T {
23
+ if (WRAPPED.has(client as object)) return client;
24
+ const original = client.send as SendFn;
25
+ Object.defineProperty(client, "send", {
26
+ value: createInstrumentedSend(original, options, RUNTIME_HANDLERS),
27
+ configurable: true,
28
+ writable: true,
29
+ });
30
+ WRAPPED.add(client as object);
31
+ return client;
32
+ }
33
+
34
+ export function instrumentBedrock(options: BedrockInstrumentationOptions = {}): void {
35
+ if (ORIGINAL) return;
36
+ const proto = BedrockRuntimeClient.prototype as unknown as { send: SendFn };
37
+ const hadOwn = Object.prototype.hasOwnProperty.call(proto, "send");
38
+ const original = proto.send;
39
+ const instrumentedSend = createInstrumentedSend(original, options, RUNTIME_HANDLERS);
40
+ const wrapper = function telemetryDevBedrockSend(
41
+ this: WrappableBedrockClient,
42
+ ...args: unknown[]
43
+ ) {
44
+ if (WRAPPED.has(this as object)) return original.apply(this, args);
45
+ return instrumentedSend.apply(this, args);
46
+ };
47
+ ORIGINAL = { proto, hadOwn, value: original, wrapper };
48
+ Object.defineProperty(proto, "send", {
49
+ configurable: true,
50
+ writable: true,
51
+ value: wrapper,
52
+ });
53
+ }
54
+
55
+ export function uninstrumentBedrock(): void {
56
+ const original = ORIGINAL;
57
+ if (!original) return;
58
+ if (Object.getOwnPropertyDescriptor(original.proto, "send")?.value !== original.wrapper) return;
59
+ ORIGINAL = undefined;
60
+ if (original.hadOwn) {
61
+ Object.defineProperty(original.proto, "send", {
62
+ configurable: true,
63
+ writable: true,
64
+ value: original.value,
65
+ });
66
+ } else {
67
+ Reflect.deleteProperty(original.proto, "send");
68
+ }
69
+ }
@@ -0,0 +1,230 @@
1
+ import type { OpHandler } from "./internal.ts";
2
+ import {
3
+ AgentStreamState,
4
+ FlowStreamState,
5
+ RagStreamState,
6
+ agentRequestFields,
7
+ flowName,
8
+ flowRequestFields,
9
+ inlineAgentRequestFields,
10
+ ragName,
11
+ ragRequestFields,
12
+ ragResponseFields,
13
+ retrieveRequestFields,
14
+ retrieveResponseFields,
15
+ } from "./agent-handlers.ts";
16
+ import { ConverseStreamState, converseRequestFields, converseResponseFields } from "./converse.ts";
17
+ import {
18
+ invokeModelRequestFields,
19
+ invokeModelResponseFields,
20
+ invokeModelSpanName,
21
+ InvokeModelStreamState,
22
+ isEmbeddingModel,
23
+ } from "./invoke-model.ts";
24
+ import {
25
+ awsMetadataFields,
26
+ mergeFields,
27
+ omitUndefined,
28
+ PROVIDER,
29
+ stringValue,
30
+ wrapAsyncIterable,
31
+ endSpan,
32
+ type BedrockInstrumentationOptions,
33
+ } from "./internal.ts";
34
+
35
+ export const RUNTIME_HANDLERS: Record<string, OpHandler> = {
36
+ ConverseCommand: {
37
+ spanName: (input) => `chat ${stringValue(input.modelId) ?? "unknown"}`,
38
+ spanType: () => "generation",
39
+ requestFields: converseRequestFields,
40
+ onResult(result, span) {
41
+ endSpan(span, converseResponseFields(result));
42
+ return result;
43
+ },
44
+ },
45
+ ConverseStreamCommand: {
46
+ spanName: (input) => `chat ${stringValue(input.modelId) ?? "unknown"}`,
47
+ spanType: () => "generation",
48
+ requestFields: converseRequestFields,
49
+ onResult(result, span, t0) {
50
+ const response = result as { stream?: unknown; $metadata?: unknown };
51
+ if (!response.stream) {
52
+ endSpan(span, awsMetadataFields(response.$metadata));
53
+ return result;
54
+ }
55
+ return {
56
+ ...response,
57
+ stream: wrapAsyncIterable(
58
+ response.stream,
59
+ new ConverseStreamState(),
60
+ span,
61
+ t0,
62
+ awsMetadataFields(response.$metadata),
63
+ ),
64
+ };
65
+ },
66
+ },
67
+ InvokeModelCommand: {
68
+ spanName: invokeModelSpanName,
69
+ spanType: (input) => (isEmbeddingModel(input.modelId) ? "embedding" : "generation"),
70
+ requestFields: invokeModelRequestFields,
71
+ onResult(result, span, _t0, input) {
72
+ endSpan(span, invokeModelResponseFields(input, result));
73
+ return result;
74
+ },
75
+ },
76
+ InvokeModelWithResponseStreamCommand: {
77
+ spanName: invokeModelSpanName,
78
+ spanType: (input) => (isEmbeddingModel(input.modelId) ? "embedding" : "generation"),
79
+ requestFields: invokeModelRequestFields,
80
+ onResult(result, span, t0) {
81
+ const response = result as { body?: unknown; $metadata?: unknown };
82
+ if (!response.body) {
83
+ endSpan(span, awsMetadataFields(response.$metadata));
84
+ return result;
85
+ }
86
+ return {
87
+ ...response,
88
+ body: wrapAsyncIterable(
89
+ response.body,
90
+ new InvokeModelStreamState(),
91
+ span,
92
+ t0,
93
+ awsMetadataFields(response.$metadata),
94
+ ),
95
+ };
96
+ },
97
+ },
98
+ ApplyGuardrailCommand: {
99
+ spanName: (input) => `apply_guardrail ${stringValue(input.guardrailIdentifier) ?? "unknown"}`,
100
+ spanType: () => "span",
101
+ requestFields: (input) => ({
102
+ provider: PROVIDER,
103
+ input: omitUndefined({
104
+ source: input.source,
105
+ content: input.content,
106
+ guardrailIdentifier: input.guardrailIdentifier,
107
+ guardrailVersion: input.guardrailVersion,
108
+ }),
109
+ }),
110
+ onResult(result, span) {
111
+ const response = result as Record<string, unknown>;
112
+ endSpan(
113
+ span,
114
+ mergeFields(awsMetadataFields(response.$metadata), {
115
+ output: omitUndefined({ outputs: response.outputs, assessments: response.assessments }),
116
+ metadata: omitUndefined({
117
+ guardrail_action: stringValue(response.action),
118
+ guardrail_action_reason: stringValue(response.actionReason),
119
+ }),
120
+ }),
121
+ );
122
+ return result;
123
+ },
124
+ },
125
+ };
126
+
127
+ export const AGENT_HANDLERS: Record<string, OpHandler> = {
128
+ InvokeAgentCommand: {
129
+ spanName: (input) => `invoke_agent ${stringValue(input.agentId) ?? "unknown"}`,
130
+ spanType: () => "agent",
131
+ requestFields: agentRequestFields,
132
+ onResult: streamAgentResult("completion"),
133
+ },
134
+ InvokeInlineAgentCommand: {
135
+ spanName: (input) => `invoke_agent ${stringValue(input.agentName) ?? "inline-agent"}`,
136
+ spanType: () => "agent",
137
+ requestFields: inlineAgentRequestFields,
138
+ onResult: streamAgentResult("completion"),
139
+ },
140
+ RetrieveCommand: {
141
+ spanName: (input) => `retrieve ${stringValue(input.knowledgeBaseId) ?? "unknown"}`,
142
+ spanType: () => "span",
143
+ requestFields: retrieveRequestFields,
144
+ onResult(result, span) {
145
+ endSpan(span, retrieveResponseFields(result));
146
+ return result;
147
+ },
148
+ },
149
+ RetrieveAndGenerateCommand: {
150
+ spanName: ragName,
151
+ spanType: () => "generation",
152
+ requestFields: ragRequestFields,
153
+ onResult(result, span) {
154
+ endSpan(span, ragResponseFields(result));
155
+ return result;
156
+ },
157
+ },
158
+ RetrieveAndGenerateStreamCommand: {
159
+ spanName: ragName,
160
+ spanType: () => "generation",
161
+ requestFields: ragRequestFields,
162
+ onResult(result, span, t0) {
163
+ const response = result as { stream?: unknown; $metadata?: unknown; sessionId?: string };
164
+ const base = mergeFields(awsMetadataFields(response.$metadata), {
165
+ metadata: omitUndefined({ bedrock_session_id: stringValue(response.sessionId) }),
166
+ });
167
+ if (!response.stream) {
168
+ endSpan(span, base);
169
+ return result;
170
+ }
171
+ return {
172
+ ...response,
173
+ stream: wrapAsyncIterable(response.stream, new RagStreamState(), span, t0, base),
174
+ };
175
+ },
176
+ },
177
+ InvokeFlowCommand: {
178
+ spanName: flowName,
179
+ spanType: () => "agent",
180
+ requestFields: flowRequestFields,
181
+ onResult(result, span, t0) {
182
+ const response = result as { responseStream?: unknown; $metadata?: unknown };
183
+ if (!response.responseStream) {
184
+ endSpan(span, awsMetadataFields(response.$metadata));
185
+ return result;
186
+ }
187
+ return {
188
+ ...response,
189
+ responseStream: wrapAsyncIterable(
190
+ response.responseStream,
191
+ new FlowStreamState(),
192
+ span,
193
+ t0,
194
+ awsMetadataFields(response.$metadata),
195
+ ),
196
+ };
197
+ },
198
+ },
199
+ };
200
+
201
+ function streamAgentResult(key: "completion") {
202
+ return (
203
+ result: unknown,
204
+ span: Parameters<OpHandler["onResult"]>[1],
205
+ t0: number,
206
+ _input: Record<string, unknown>,
207
+ options: BedrockInstrumentationOptions,
208
+ ): unknown => {
209
+ const response = result as {
210
+ completion?: unknown;
211
+ $metadata?: unknown;
212
+ sessionId?: string;
213
+ memoryId?: string;
214
+ };
215
+ const base = mergeFields(awsMetadataFields(response.$metadata), {
216
+ metadata: omitUndefined({
217
+ bedrock_session_id: stringValue(response.sessionId),
218
+ bedrock_memory_id: stringValue(response.memoryId),
219
+ }),
220
+ });
221
+ if (!response[key]) {
222
+ endSpan(span, base);
223
+ return result;
224
+ }
225
+ return {
226
+ ...response,
227
+ [key]: wrapAsyncIterable(response[key], new AgentStreamState(options), span, t0, base),
228
+ };
229
+ };
230
+ }