@telemetry-dev/bedrock 0.1.2 → 0.1.5

package/src/converse.ts

@@ -17,6 +17,7 @@ import {
 } from "./internal.ts";
 
 export type MessagePart = object;
+
 export interface NormalizedMessage {
   role: string | undefined;
   parts: MessagePart[];
@@ -26,6 +27,7 @@ export interface NormalizedMessage {
 export function converseRequestFields(input: JsonRecord): SpanFields {
   const inference = isRecord(input.inferenceConfig) ? input.inferenceConfig : {};
   const guardrail = isRecord(input.guardrailConfig) ? input.guardrailConfig : undefined;
+
   return omitUndefined({
     provider: PROVIDER,
     model: stringValue(input.modelId),
@@ -53,6 +55,7 @@ export function converseResponseFields<T>(output: T): SpanFields {
   const guardrail = isRecord(trace?.guardrail) ? trace?.guardrail : undefined;
   const stopReason = stringValue(result.stopReason);
   const message = isRecord(result.output) ? result.output.message : undefined;
+
   return mergeFields(awsMetadataFields(result.$metadata), {
     output: normalizeOutputMessage(message, stopReason),
     usage,
@@ -68,6 +71,7 @@ export function converseResponseFields<T>(output: T): SpanFields {
 
 export function usageFromConverse<T>(value: T): SpanFields["usage"] | undefined {
   if (!isRecord(value)) return undefined;
+
   return omitUndefined({
     inputTokens: numberValue(value.inputTokens),
     outputTokens: numberValue(value.outputTokens),
@@ -79,7 +83,9 @@ export function usageFromConverse<T>(value: T): SpanFields["usage"] | undefined
 
 export function normalizeMessages<T>(value: T): NormalizedMessage[] | undefined {
   const messages = arrayValue<JsonRecord>(value);
+
   if (!messages) return undefined;
+
   return messages.map((message) => ({
     role: stringValue(message.role),
     parts: normalizeContentList(message.content),
@@ -91,29 +97,36 @@ export function normalizeOutputMessage<T>(
   finishReason?: string,
 ): NormalizedMessage[] | undefined {
   if (!isRecord(value)) return undefined;
+
   const message = omitUndefined({
     role: stringValue(value.role),
     parts: normalizeContentList(value.content),
     finish_reason: finishReason,
   });
+
   return [message];
 }
 
 export function normalizeSystem<T>(value: T): MessagePart[] | undefined {
   const blocks = arrayValue<JsonRecord>(value);
+
   if (!blocks) return undefined;
   const parts = blocks.flatMap((block) => normalizeContentBlock(block));
+
   return parts.length > 0 ? parts : undefined;
 }
 
 function normalizeContentList<T>(value: T): MessagePart[] {
   const blocks = arrayValue<JsonRecord>(value);
+
   if (!blocks) return [];
+
   return blocks.flatMap((block) => normalizeContentBlock(block));
 }
 
 function normalizeContentBlock(block: JsonRecord): MessagePart[] {
   const text = stringValue(block.text);
+
   if (text !== undefined)
     return [
       omitUndefined({
@@ -124,6 +137,7 @@ function normalizeContentBlock(block: JsonRecord): MessagePart[] {
     ];
 
   const toolUse = isRecord(block.toolUse) ? block.toolUse : undefined;
+
   if (toolUse) {
     return [
       omitUndefined({
@@ -136,6 +150,7 @@ function normalizeContentBlock(block: JsonRecord): MessagePart[] {
   }
 
   const toolResult = isRecord(block.toolResult) ? block.toolResult : undefined;
+
   if (toolResult) {
     return [
       omitUndefined({
@@ -147,8 +162,10 @@ function normalizeContentBlock(block: JsonRecord): MessagePart[] {
   }
 
   const reasoning = isRecord(block.reasoningContent) ? block.reasoningContent : undefined;
+
   if (reasoning) {
     const reasoningText = isRecord(reasoning.reasoningText) ? reasoning.reasoningText : undefined;
+
     return [
       {
         type: "reasoning",
@@ -160,9 +177,11 @@ function normalizeContentBlock(block: JsonRecord): MessagePart[] {
 
   for (const modality of ["image", "document", "video", "audio"] as const) {
     const media = isRecord(block[modality]) ? block[modality] : undefined;
+
     if (!media) continue;
     const format = stringValue(media.format);
     const source = isRecord(media.source) ? media.source : undefined;
+
     if (source?.bytes !== undefined) {
       return [
         omitUndefined({
@@ -172,25 +191,32 @@ function normalizeContentBlock(block: JsonRecord): MessagePart[] {
         }),
       ];
     }
+
     const s3 = isRecord(source?.s3Location) ? source?.s3Location : undefined;
     const uri = stringValue(source?.uri) ?? stringValue(s3?.uri) ?? s3Uri(s3);
+
     if (uri) return [omitUndefined({ type: "uri", uri, modality, mime_type: format })];
   }
 
   const citationsContent = isRecord(block.citationsContent) ? block.citationsContent : undefined;
+
   if (citationsContent) {
     const contentBlocks = arrayValue<JsonRecord>(citationsContent.content) ?? [];
     const citations = arrayValue<JsonValue>(citationsContent.citations);
+
     return contentBlocks.flatMap((contentBlock) => {
       if (!citations) return normalizeContentBlock(contentBlock);
+
       return normalizeContentBlock({ ...contentBlock, citations });
     });
   }
 
   const guardContent = isRecord(block.guardContent) ? block.guardContent : undefined;
+
   if (guardContent) {
     const guardText = isRecord(guardContent.text) ? guardContent.text : undefined;
     const content = stringValue(guardText?.text) ?? stringValue(guardContent.text);
+
     if (content) return [{ type: "text", content }];
   }
 
@@ -203,15 +229,19 @@ function normalizeContentBlock(block: JsonRecord): MessagePart[] {
 
 function simplifyToolResult<T>(content: T): JsonValue {
   const blocks = arrayValue<JsonRecord>(content);
+
   if (!blocks) return isJsonValue(content) ? content : undefined;
   const text = stringValue(blocks[0]?.text);
+
   if (blocks.length === 1 && text !== undefined) return text;
+
   return blocks.flatMap((block) => normalizeContentBlock(block)) as JsonValue[];
 }
 
 function s3Uri(value: JsonRecord | undefined): string | undefined {
   const bucket = stringValue(value?.bucket);
   const key = stringValue(value?.key);
+
   return bucket && key ? `s3://${bucket}/${key}` : undefined;
 }
 
@@ -238,15 +268,20 @@ export class ConverseStreamState implements StreamState {
   feed<T>(event: T): void {
     if (!isRecord(event)) return;
     const streamError = modeledStreamError(event);
+
     if (streamError) {
       this.error = streamError.error;
       this.errorFields = streamError.fields;
+
       return;
     }
+
     const messageStart = isRecord(event.messageStart) ? event.messageStart : undefined;
+
     if (messageStart) this.role = stringValue(messageStart.role) ?? this.role;
 
     const blockStart = isRecord(event.contentBlockStart) ? event.contentBlockStart : undefined;
+
     if (blockStart) {
       const index = numberValue(blockStart.contentBlockIndex) ?? this.blocks.size;
       const start = isRecord(blockStart.start) ? blockStart.start : undefined;
@@ -280,34 +315,44 @@ export class ConverseStreamState implements StreamState {
     }
 
     const deltaEvent = isRecord(event.contentBlockDelta) ? event.contentBlockDelta : undefined;
+
     if (deltaEvent) {
       const index = numberValue(deltaEvent.contentBlockIndex) ?? 0;
       const delta = isRecord(deltaEvent.delta) ? deltaEvent.delta : {};
       const block = this.blocks.get(index) ?? { kind: "text", text: "" };
       const text = stringValue(delta.text);
+
       if (text !== undefined) {
         block.kind = "text";
         block.text = `${block.text ?? ""}${text}`;
       }
+
       const toolUse = isRecord(delta.toolUse) ? delta.toolUse : undefined;
       const input = stringValue(toolUse?.input);
+
       if (toolUse && input !== undefined) {
         block.kind = "tool";
         block.input = `${block.input ?? ""}${input}`;
       }
+
       const reasoning = isRecord(delta.reasoningContent) ? delta.reasoningContent : undefined;
+
       if (reasoning) {
         block.kind = "reasoning";
         block.text = `${block.text ?? ""}${stringValue(reasoning.text) ?? (reasoning.redactedContent ? "[redacted]" : "")}`;
       }
+
       const image = isRecord(delta.image) ? delta.image : undefined;
+
       if (image) {
         const content = isRecord(block.content) ? block.content : {};
         const current: JsonRecord = isRecord(content.image) ? content.image : {};
         block.kind = "content";
         block.content = { image: { ...current, ...image } };
       }
+
       const toolResult = arrayValue<JsonRecord>(delta.toolResult);
+
       if (toolResult) {
         const content = isRecord(block.content) ? block.content : {};
         const current: JsonRecord = isRecord(content.toolResult) ? content.toolResult : {};
@@ -320,16 +365,20 @@ export class ConverseStreamState implements StreamState {
           },
         };
       }
+
       if (isRecord(delta.citation)) {
         block.citations = [...(block.citations ?? []), delta.citation];
       }
+
       this.blocks.set(index, block);
     }
 
     const stop = isRecord(event.messageStop) ? event.messageStop : undefined;
+
     if (stop) this.stopReason = stringValue(stop.stopReason) ?? this.stopReason;
 
     const metadata = isRecord(event.metadata) ? event.metadata : undefined;
+
     if (metadata) {
       this.usage = usageFromConverse(metadata.usage) ?? this.usage;
       const metrics = isRecord(metadata.metrics) ? metadata.metrics : undefined;
@@ -352,6 +401,7 @@ export class ConverseStreamState implements StreamState {
     const parts = [...this.blocks.entries()]
       .sort(([left], [right]) => left - right)
       .map(([, block]) => blockToPart(block));
+
     return mergeFields(
       omitUndefined({
         output:
@@ -395,9 +445,12 @@ function blockToPart(block: StreamBlock) {
       arguments: parseJson(block.input) ?? block.input ?? "",
     });
   }
+
   if (block.kind === "reasoning") return { type: "reasoning", content: block.text ?? "" };
+
   if (block.kind === "content" && block.content)
     return normalizeContentBlock(block.content)[0] ?? {};
+
   return omitUndefined({
     type: "text",
     content: block.text ?? "",

package/src/index.ts

@@ -28,6 +28,7 @@ export function wrapBedrock<T extends WrappableBedrockClient>(
     writable: true,
   });
   WRAPPED.add(client as object);
+
   return client;
 }
 
@@ -37,13 +38,16 @@ export function instrumentBedrock(options: BedrockInstrumentationOptions = {}):
   const hadOwn = Object.prototype.hasOwnProperty.call(proto, "send");
   const original = proto.send;
   const instrumentedSend = createInstrumentedSend(original, options, RUNTIME_HANDLERS);
+
   const wrapper = function telemetryDevBedrockSend(
     this: WrappableBedrockClient,
     ...args: unknown[]
   ) {
     if (WRAPPED.has(this as object)) return original.apply(this, args);
+
     return instrumentedSend.apply(this, args);
   };
+
   ORIGINAL = { proto, hadOwn, value: original, wrapper };
   Object.defineProperty(proto, "send", {
     configurable: true,
@@ -54,9 +58,12 @@ export function instrumentBedrock(options: BedrockInstrumentationOptions = {}):
 
 export function uninstrumentBedrock(): void {
   const original = ORIGINAL;
+
   if (!original) return;
+
   if (Object.getOwnPropertyDescriptor(original.proto, "send")?.value !== original.wrapper) return;
   ORIGINAL = undefined;
+
   if (original.hadOwn) {
     Object.defineProperty(original.proto, "send", {
       configurable: true,

package/src/internal-handlers.ts

@@ -39,6 +39,7 @@ export const RUNTIME_HANDLERS = {
     requestFields: converseRequestFields,
     onResult(result, span) {
       endSpan(span, converseResponseFields(result));
+
       return result;
     },
   },
@@ -48,10 +49,13 @@ export const RUNTIME_HANDLERS = {
     requestFields: converseRequestFields,
     onResult(result, span, t0) {
       const response = result as { stream?: unknown; $metadata?: unknown };
+
       if (!response.stream) {
         endSpan(span, awsMetadataFields(response.$metadata));
+
         return result;
       }
+
       return {
         ...response,
         stream: wrapAsyncIterable(
@@ -70,6 +74,7 @@ export const RUNTIME_HANDLERS = {
     requestFields: invokeModelRequestFields,
     onResult(result, span, _t0, input) {
       endSpan(span, invokeModelResponseFields(input, result));
+
       return result;
     },
   },
@@ -79,10 +84,13 @@ export const RUNTIME_HANDLERS = {
     requestFields: invokeModelRequestFields,
     onResult(result, span, t0) {
       const response = result as { body?: unknown; $metadata?: unknown };
+
       if (!response.body) {
         endSpan(span, awsMetadataFields(response.$metadata));
+
         return result;
       }
+
       return {
         ...response,
         body: wrapAsyncIterable(
@@ -119,6 +127,7 @@ export const RUNTIME_HANDLERS = {
           }),
         }),
       );
+
       return result;
     },
   },
@@ -143,6 +152,7 @@ export const AGENT_HANDLERS = {
     requestFields: retrieveRequestFields,
     onResult(result, span) {
       endSpan(span, retrieveResponseFields(result));
+
       return result;
     },
   },
@@ -152,6 +162,7 @@ export const AGENT_HANDLERS = {
     requestFields: ragRequestFields,
     onResult(result, span) {
       endSpan(span, ragResponseFields(result));
+
       return result;
     },
   },
@@ -161,16 +172,20 @@ export const AGENT_HANDLERS = {
     requestFields: ragRequestFields,
     onResult(result, span, t0) {
       const response = result as { stream?: unknown; $metadata?: unknown; sessionId?: string };
+
       const base = mergeFields(awsMetadataFields(response.$metadata), {
         metadata: omitUndefined({ bedrock_session_id: stringValue(response.sessionId) }),
       });
+
       if (!response.stream) {
         endSpan(span, base);
+
         return result;
       }
+
       return {
         ...response,
-        stream: wrapAsyncIterable(response.stream, new RagStreamState(), span, t0, base),
+        stream: wrapAsyncIterable(response.stream, new RagStreamState(), span, t0, base, false),
       };
     },
   },
@@ -180,10 +195,13 @@ export const AGENT_HANDLERS = {
     requestFields: flowRequestFields,
     onResult(result, span, t0) {
       const response = result as { responseStream?: unknown; $metadata?: unknown };
+
       if (!response.responseStream) {
         endSpan(span, awsMetadataFields(response.$metadata));
+
         return result;
       }
+
       return {
         ...response,
         responseStream: wrapAsyncIterable(
@@ -192,6 +210,7 @@ export const AGENT_HANDLERS = {
           span,
           t0,
           awsMetadataFields(response.$metadata),
+          false,
         ),
       };
     },
@@ -212,19 +231,23 @@ function streamAgentResult(key: "completion") {
       sessionId?: string;
       memoryId?: string;
     };
+
     const base = mergeFields(awsMetadataFields(response.$metadata), {
       metadata: omitUndefined({
         bedrock_session_id: stringValue(response.sessionId),
         bedrock_memory_id: stringValue(response.memoryId),
       }),
     });
+
     if (!response[key]) {
       endSpan(span, base);
+
       return result;
     }
+
     return {
       ...response,
-      [key]: wrapAsyncIterable(response[key], new AgentStreamState(options), span, t0, base),
+      [key]: wrapAsyncIterable(response[key], new AgentStreamState(options), span, t0, base, false),
     };
   };
 }