@telemetry-dev/bedrock 0.1.3 → 0.1.5

package/src/internal.ts

@@ -20,6 +20,7 @@ export type JsonValue =
   | JsonValue[]
   | JsonRecord
   | ((...args: JsonValue[]) => JsonValue);
+
 export interface JsonRecord {
   [key: string]: JsonValue;
 }
@@ -53,6 +54,7 @@ export function createInstrumentedSend(
     const rawRest = rest.map(asJsonValue);
     const name = isRecord(rawCommand) ? rawCommand.constructor?.name : undefined;
     const handler = String(name) === name ? handlers[name] : undefined;
+
     if (!handler) {
       return originalSend.apply(this, [rawCommand, ...rawRest]);
     }
@@ -64,25 +66,32 @@ export function createInstrumentedSend(
     const fields = safe(() => handler.requestFields(input, options)) ?? {};
     const span = startSpan(spanName, { type: spanType, ...fields });
     const callbackIndex = rawRest.findIndex((arg) => arg instanceof Function);
+
     if (callbackIndex !== -1) {
       const callback = rawRest[callbackIndex] as (cause: JsonValue, data?: JsonValue) => void;
       const wrappedRest = [...rawRest];
       wrappedRest[callbackIndex] = <TError, TData>(error: TError, data?: TData) => {
         const rawError: unknown = error;
         const rawData: unknown = data;
+
         if (rawError != null) {
           endSpan(span, { ...awsMetadataFields(metadataOf(rawError)), error: asError(rawError) });
           callback(asJsonValue(rawError), asJsonValue(rawData));
+
           return;
         }
+
         let nextData: JsonValue | object = asJsonValue(rawData);
+
         try {
           nextData = handler.onResult(asJsonValue(rawData), span, t0, input, options);
         } catch {
           endSpan(span, awsMetadataFields(metadataOf(rawData)));
         }
+
         callback(asJsonValue(rawError), asJsonValue(nextData));
       };
+
       try {
         return originalSend.apply(this, [rawCommand, ...wrappedRest]);
       } catch (error) {
@@ -92,6 +101,7 @@ export function createInstrumentedSend(
     }
 
     let result: JsonValue | object | Promise<JsonValue | object>;
+
     try {
       result = originalSend.apply(this, [rawCommand, ...rawRest]);
     } catch (error) {
@@ -105,6 +115,7 @@ export function createInstrumentedSend(
           return handler.onResult(asJsonValue(resolved), span, t0, input, options);
         } catch {
           endSpan(span, awsMetadataFields(metadataOf(resolved)));
+
           return resolved;
         }
       },
@@ -118,6 +129,7 @@ export function createInstrumentedSend(
 
 export function commandInput<T>(command: T): JsonRecord {
   if (!isRecord(command)) return {};
+
   return isRecord(command.input) ? command.input : {};
 }
 
@@ -151,6 +163,7 @@ export function updateSpan(span: SpanHandle, fields: SpanFields): void {
 
 export function metadataOf<T>(value: T): JsonValue {
   if (!isRecord(value)) return undefined;
+
   return value.$metadata;
 }
 
@@ -172,31 +185,42 @@ export function modeledStreamError(
     throttlingException: "ThrottlingException",
     validationException: "ValidationException",
   } as const;
+
   for (const [key, name] of Object.entries(keys)) {
     const exception = isRecord(event[key]) ? event[key] : undefined;
+
     if (!exception) continue;
+
     const error = new Error(
       stringValue(exception.message) ?? stringValue(exception.originalMessage) ?? name,
     );
+
     error.name = stringValue(exception.name) ?? name;
+
     return { error, fields: awsMetadataFields(exception.$metadata) };
   }
+
   return undefined;
 }
 
 export function awsMetadataFields<T>(metadata: T): SpanFields {
   const raw: unknown = metadata;
+
   if (!isRecord(raw)) return {};
   const requestId = stringValue(raw.requestId);
   const attempts = numberValue(raw.attempts);
   const httpStatusCode = numberValue(raw.httpStatusCode);
   const totalRetryDelay = numberValue(raw.totalRetryDelay);
   const attributes: SpanFields["attributes"] = {};
+
   if (httpStatusCode !== undefined) attributes["aws.http.status_code"] = httpStatusCode;
+
   if (attempts !== undefined && attempts > 1) attributes["aws.request.attempts"] = attempts;
+
   if (totalRetryDelay !== undefined) {
     attributes["aws.request.total_retry_delay_ms"] = totalRetryDelay;
   }
+
   return omitUndefined({
     responseId: requestId,
     attributes: Object.keys(attributes).length > 0 ? attributes : undefined,
@@ -205,14 +229,19 @@ export function awsMetadataFields<T>(metadata: T): SpanFields {
 
 export function mergeFields(...fields: Array<SpanFields | undefined>): SpanFields {
   const merged: SpanFields = {};
+
   for (const field of fields) {
     if (!field) continue;
     const { metadata, attributes, usage, ...rest } = field;
     Object.assign(merged, omitUndefined(rest));
+
     if (usage) merged.usage = { ...merged.usage, ...usage };
+
     if (metadata) merged.metadata = { ...merged.metadata, ...metadata };
+
     if (attributes) merged.attributes = { ...merged.attributes, ...attributes };
   }
+
   return merged;
 }
 
@@ -222,11 +251,13 @@ export function omitUndefined<T extends object>(value: T): T {
 
 export function stringValue<T>(value: T): string | undefined {
   const raw: unknown = value;
+
   return String(raw) === raw && raw.length > 0 ? raw : undefined;
 }
 
 export function numberValue<T>(value: T): number | undefined {
   const raw: unknown = value;
+
   return Number(raw) === raw && Number.isFinite(raw) ? raw : undefined;
 }
 
@@ -236,8 +267,11 @@ export function arrayValue<T, TValue = unknown>(value: TValue): T[] | undefined
 
 export function bytesToString<T>(value: T): string | undefined {
   const raw: unknown = value;
+
   if (String(raw) === raw) return raw;
+
   if (raw instanceof Uint8Array) return new TextDecoder().decode(raw);
+
   if (ArrayBuffer.isView(raw)) {
     return new TextDecoder().decode(new Uint8Array(raw.buffer, raw.byteOffset, raw.byteLength));
   }
@@ -245,7 +279,9 @@ export function bytesToString<T>(value: T): string | undefined {
 
 export function parseJson<T>(value: T): JsonValue {
   const text = bytesToString(value);
+
   if (text === undefined) return undefined;
+
   try {
     return asJsonValue(JSON.parse(text));
   } catch {
@@ -254,7 +290,7 @@ export function parseJson<T>(value: T): JsonValue {
 }
 
 export interface StreamState {
-  feed<T>(event: T): void;
+  feed<T>(event: T): boolean | void;
   finish(partial: boolean): SpanFields;
 }
 
@@ -264,15 +300,18 @@ export function wrapAsyncIterable<T>(
   span: SpanHandle,
   t0: number,
   baseFields: SpanFields,
+  trackOutputChunks = true,
 ): T {
   if (!iterable || !isAsyncIterable(iterable)) {
     endSpan(span, baseFields);
+
     return iterable;
   }
 
   let ended = false;
   let first = false;
   const unregisterToken = {};
+
   const endOnce = (fields?: SpanFields) => {
     if (ended) return;
     ended = true;
@@ -282,15 +321,25 @@ export function wrapAsyncIterable<T>(
 
   const generator = async function* () {
     let completed = false;
+
     try {
       for await (const event of iterable) {
+        const receivedAt = performance.now();
+
         if (!first) {
           first = true;
-          updateSpan(span, { timeToFirstChunkMs: performance.now() - t0 });
+          updateSpan(span, { timeToFirstChunkMs: receivedAt - t0 });
+        }
+
+        const stateHasOutput = safe(() => state.feed(event));
+
+        if (trackOutputChunks && (stateHasOutput ?? bedrockEventHasOutput(event))) {
+          safe(() => span.recordOutputChunk?.(receivedAt));
         }
-        safe(() => state.feed(event));
+
         yield event;
       }
+
       completed = true;
     } catch (error) {
       endOnce(
@@ -305,21 +354,90 @@ export function wrapAsyncIterable<T>(
   };
 
   const wrapped: { [key: PropertyKey]: JsonValue | (() => AsyncGenerator<unknown>) } = {};
+
   if (isRecord(iterable)) {
     for (const key of Reflect.ownKeys(iterable)) {
       wrapped[key] = (iterable as { [key: PropertyKey]: JsonValue })[key];
     }
   }
+
   const finishAbandoned = () => endOnce(safe(() => state.finish(true)));
   wrapped[Symbol.asyncIterator] = () => {
     const iterator = generator();
     STREAM_FINALIZER?.register(iterator, finishAbandoned, unregisterToken);
+
     return iterator;
   };
+
   STREAM_FINALIZER?.register(wrapped, finishAbandoned, unregisterToken);
+
   return wrapped as T;
 }
 
+function bedrockEventHasOutput(event: unknown): boolean {
+  const record = isRecord(event) ? event : undefined;
+
+  if (!record) return false;
+  const output = isRecord(record.output) ? record.output : undefined;
+
+  if (typeof output?.text === "string" && output.text.length > 0) return true;
+  const block = isRecord(record.contentBlockDelta) ? record.contentBlockDelta : undefined;
+  const delta = isRecord(block?.delta) ? block.delta : undefined;
+  const reasoning = isRecord(delta?.reasoningContent) ? delta.reasoningContent : undefined;
+  const toolUse = isRecord(delta?.toolUse) ? delta.toolUse : undefined;
+
+  if (
+    (typeof delta?.text === "string" && delta.text.length > 0) ||
+    (typeof reasoning?.text === "string" && reasoning.text.length > 0) ||
+    (typeof toolUse?.input === "string" && toolUse.input.length > 0)
+  )
+    return true;
+  const chunk = isRecord(record.chunk) ? record.chunk : undefined;
+  const raw = chunk?.bytes;
+
+  if (raw === undefined) return false;
+
+  const text =
+    typeof raw === "string"
+      ? raw
+      : raw instanceof Uint8Array
+        ? new TextDecoder().decode(raw)
+        : undefined;
+
+  if (!text) return false;
+  const parsed = parseJson(text);
+
+  if (parsed === undefined) return text.length > 0;
+
+  if (!isRecord(parsed)) return false;
+  const parsedDelta = isRecord(parsed.delta) ? parsed.delta : undefined;
+
+  const contentDelta = isRecord(parsed.contentBlockDelta ?? parsed.content_block_delta)
+    ? (parsed.contentBlockDelta ?? parsed.content_block_delta)
+    : undefined;
+
+  const nested =
+    isRecord(contentDelta) && isRecord(contentDelta.delta) ? contentDelta.delta : undefined;
+
+  const generations = Array.isArray(parsed.generations) ? parsed.generations : [];
+  const generation = isRecord(generations[0]) ? generations[0] : undefined;
+  const outputs = Array.isArray(parsed.outputs) ? parsed.outputs : [];
+  const firstOutput = isRecord(outputs[0]) ? outputs[0] : undefined;
+
+  return [
+    parsed.outputText,
+    parsed.generation,
+    parsed.completion,
+    parsed.text,
+    parsedDelta?.text,
+    parsedDelta?.thinking,
+    parsedDelta?.partial_json,
+    nested?.text,
+    firstOutput?.text,
+    generation?.text,
+  ].some((value) => typeof value === "string" && value.length > 0);
+}
+
 function isAsyncIterable<T>(value: T): value is T & AsyncIterable<unknown> {
   return Object(value) === value && Symbol.asyncIterator in Object(value);
 }
@@ -334,6 +452,8 @@ function asError<T>(value: T): Error {
 
 export function basenameArn<T>(value: T): string | undefined {
   const text = stringValue(value);
+
   if (!text) return undefined;
+
   return text.split(/[/:]/).filter(Boolean).at(-1) ?? text;
 }

package/src/invoke-model.ts

@@ -18,6 +18,7 @@ import {
 
 export function isEmbeddingModel<T>(modelId: T): boolean {
   const model = stringValue(modelId)?.toLowerCase() ?? "";
+
   return (
     model.includes("titan-embed") || model.includes("cohere.embed") || model.includes("-embedding")
   );
@@ -25,12 +26,14 @@ export function isEmbeddingModel<T>(modelId: T): boolean {
 
 export function invokeModelSpanName(input: JsonRecord): string {
   const model = stringValue(input.modelId) ?? "unknown";
+
   return `${isEmbeddingModel(model) ? "embeddings" : "chat"} ${model}`;
 }
 
 export function invokeModelRequestFields(input: JsonRecord): SpanFields {
   const model = stringValue(input.modelId);
   const body = parseBody(input.body, input.contentType);
+
   return omitUndefined({
     provider: PROVIDER,
     model,
@@ -44,20 +47,25 @@ export function invokeModelResponseFields(input: JsonRecord, output: JsonValue):
   const model = stringValue(input.modelId);
   const body = isRecord(output) ? parseBody(output.body, output.contentType) : undefined;
   const metadata = isRecord(output) ? output.$metadata : undefined;
+
   return mergeFields(awsMetadataFields(metadata), nativeResponseFields(model, body, metadata));
 }
 
 export function parseBody<TBody, TContentType>(body: TBody, contentType: TContentType): JsonValue {
   const type = stringValue(contentType) ?? "application/json";
+
   if (!type.includes("json")) return undefined;
+
   return parseJson(body);
 }
 
 function nativeSamplingFields<T>(model: string | undefined, body: T): SpanFields {
   if (!isRecord(body)) return {};
   const lower = model?.toLowerCase() ?? "";
+
   if (lower.includes("amazon.titan")) {
     const cfg = isRecord(body.textGenerationConfig) ? body.textGenerationConfig : {};
+
     return omitUndefined({
       temperature: numberValue(cfg.temperature),
       topP: numberValue(cfg.topP),
@@ -65,8 +73,10 @@ function nativeSamplingFields<T>(model: string | undefined, body: T): SpanFields
       stopSequences: arrayValue<string>(cfg.stopSequences),
     });
   }
+
   if (lower.includes("amazon.nova")) {
     const cfg = isRecord(body.inferenceConfig) ? body.inferenceConfig : {};
+
     return omitUndefined({
       temperature: numberValue(cfg.temperature),
       topP: numberValue(cfg.topP) ?? numberValue(cfg.top_p),
@@ -75,6 +85,7 @@ function nativeSamplingFields<T>(model: string | undefined, body: T): SpanFields
       stopSequences: arrayValue<string>(cfg.stopSequences),
     });
   }
+
   if (lower.includes("anthropic.claude")) {
     return omitUndefined({
       temperature: numberValue(body.temperature),
@@ -84,6 +95,7 @@ function nativeSamplingFields<T>(model: string | undefined, body: T): SpanFields
       stopSequences: arrayValue<string>(body.stop_sequences),
     });
   }
+
   if (lower.includes("meta.llama")) {
     return omitUndefined({
       temperature: numberValue(body.temperature),
@@ -91,6 +103,7 @@ function nativeSamplingFields<T>(model: string | undefined, body: T): SpanFields
       maxTokens: numberValue(body.max_gen_len),
     });
   }
+
   return omitUndefined({
     temperature: numberValue(body.temperature),
     topP: numberValue(body.top_p) ?? numberValue(body.topP),
@@ -105,6 +118,7 @@ function nativeResponseFields(
 ): SpanFields {
   if (!isRecord(body)) return headerUsageFields(metadata);
   const lower = model?.toLowerCase() ?? "";
+
   if (isEmbeddingModel(model)) {
     return withHeaderUsage(
       {
@@ -115,8 +129,10 @@ function nativeResponseFields(
       metadata,
     );
   }
+
   if (lower.includes("amazon.titan")) {
     const first = Array.isArray(body.results) && isRecord(body.results[0]) ? body.results[0] : {};
+
     return omitUndefined({
       output: body,
       finishReason: stringValue(first.completionReason),
@@ -126,8 +142,10 @@ function nativeResponseFields(
       }),
     });
   }
+
   if (lower.includes("amazon.nova")) {
     const usage = isRecord(body.usage) ? body.usage : {};
+
     return omitUndefined({
       output: body,
       finishReason: stringValue(body.stopReason),
@@ -138,8 +156,10 @@ function nativeResponseFields(
       }),
     });
   }
+
   if (lower.includes("anthropic.claude")) {
     const usage = isRecord(body.usage) ? body.usage : {};
+
     return withHeaderUsage(
       {
         output: body,
@@ -154,6 +174,7 @@ function nativeResponseFields(
       metadata,
     );
   }
+
   if (lower.includes("meta.llama")) {
     return omitUndefined({
       output: body,
@@ -164,32 +185,39 @@ function nativeResponseFields(
       }),
     });
   }
+
   if (lower.includes("mistral")) {
     const first = Array.isArray(body.outputs) && isRecord(body.outputs[0]) ? body.outputs[0] : {};
+
     return withHeaderUsage(
       { output: body, finishReason: stringValue(first.stop_reason) },
       metadata,
     );
   }
+
   if (lower.includes("cohere.command-r")) {
     return withHeaderUsage(
       { output: body, finishReason: stringValue(body.finish_reason) },
       metadata,
     );
   }
+
   if (lower.includes("cohere.command")) {
     const first =
       Array.isArray(body.generations) && isRecord(body.generations[0]) ? body.generations[0] : {};
+
     return withHeaderUsage(
       { output: body, finishReason: stringValue(first.finish_reason) },
       metadata,
     );
   }
+
   return withHeaderUsage({ output: body }, metadata);
 }
 
 function withHeaderUsage<T>(fields: SpanFields, metadata: T): SpanFields {
   if (fields.usage && Object.keys(fields.usage).length > 0) return fields;
+
   return mergeFields(fields, headerUsageFields(metadata));
 }
 
@@ -201,19 +229,24 @@ function headerUsageFields<T>(metadata: T): SpanFields {
         ? metadata.HTTPHeaders
         : undefined
     : undefined;
+
   const usage = omitUndefined({
     inputTokens: headerNumber(headers?.["x-amzn-bedrock-input-token-count"]),
     outputTokens: headerNumber(headers?.["x-amzn-bedrock-output-token-count"]),
   });
+
   return { usage: Object.keys(usage).length > 0 ? usage : undefined };
 }
 
 function headerNumber<T>(value: T): number | undefined {
   const text = stringValue(value);
+
   if (text !== undefined) {
     const parsed = Number.parseInt(text, 10);
+
     return Number.isFinite(parsed) ? parsed : undefined;
   }
+
   return numberValue(value);
 }
 
@@ -224,35 +257,49 @@ export class InvokeModelStreamState implements StreamState {
   private finishReason: string | undefined;
   private error: Error | undefined;
   private errorFields: SpanFields | undefined;
-  feed<T>(event: T): void {
-    if (!isRecord(event)) return;
+  feed<T>(event: T): boolean {
+    if (!isRecord(event)) return false;
     const streamError = modeledStreamError(event);
+
     if (streamError) {
       this.error = streamError.error;
       this.errorFields = streamError.fields;
-      return;
+
+      return false;
     }
+
     const chunk = isRecord(event.chunk) ? event.chunk : undefined;
     const parsed = parseJson(chunk?.bytes);
-    if (parsed === undefined) return;
+
+    if (parsed === undefined) {
+      const raw = chunk?.bytes;
+
+      return (typeof raw === "string" || raw instanceof Uint8Array) && raw.length > 0;
+    }
+
     this.chunks.push(parsed);
     this.text += textFromProviderChunk(parsed);
+
     if (isRecord(parsed)) {
       const delta = isRecord(parsed.delta) ? parsed.delta : undefined;
       const message = isRecord(parsed.message) ? parsed.message : undefined;
       const messageUsage = isRecord(message?.usage) ? message?.usage : undefined;
       const usage = isRecord(parsed.usage) ? parsed.usage : undefined;
+
       const metrics = isRecord(parsed["amazon-bedrock-invocationMetrics"])
         ? parsed["amazon-bedrock-invocationMetrics"]
         : undefined;
+
       const outputs =
         Array.isArray(parsed.outputs) && isRecord(parsed.outputs[0])
           ? parsed.outputs[0]
           : undefined;
+
       const generation =
         Array.isArray(parsed.generations) && isRecord(parsed.generations[0])
           ? parsed.generations[0]
           : undefined;
+
       this.finishReason =
         stringValue(delta?.stop_reason) ??
         stringValue(parsed.stop_reason) ??
@@ -279,6 +326,8 @@ export class InvokeModelStreamState implements StreamState {
           this.usage?.cacheCreationInputTokens,
       });
     }
+
+    return providerChunkHasOutput(parsed);
   }
 
   finish(): SpanFields {
@@ -305,16 +354,20 @@ export class InvokeModelStreamState implements StreamState {
 function textFromProviderChunk<T>(value: T): string {
   if (!isRecord(value)) return "";
   const delta = isRecord(value.delta) ? value.delta : undefined;
+
   const contentBlockDelta = isRecord(value.contentBlockDelta)
     ? value.contentBlockDelta
     : isRecord(value.content_block_delta)
       ? value.content_block_delta
       : undefined;
+
   const nestedDelta = isRecord(contentBlockDelta?.delta) ? contentBlockDelta?.delta : undefined;
+
   const generation =
     Array.isArray(value.generations) && isRecord(value.generations[0])
       ? value.generations[0]
       : undefined;
+
   return (
     stringValue(value.outputText) ??
     stringValue(value.generation) ??
@@ -329,3 +382,14 @@ function textFromProviderChunk<T>(value: T): string {
     ""
   );
 }
+
+function providerChunkHasOutput<T>(value: T): boolean {
+  if (!isRecord(value)) return false;
+  const delta = isRecord(value.delta) ? value.delta : undefined;
+
+  return (
+    textFromProviderChunk(value).length > 0 ||
+    (typeof delta?.thinking === "string" && delta.thinking.length > 0) ||
+    (typeof delta?.partial_json === "string" && delta.partial_json.length > 0)
+  );
+}