@telemetry-dev/bedrock 0.1.0 → 0.1.1

package/src/internal.ts

@@ -11,97 +11,118 @@ export interface BedrockInstrumentationOptions {
   captureAgentTrace?: boolean;
 }
 
-export type SendFn = (...args: unknown[]) => unknown;
+export type JsonValue =
+  | string
+  | number
+  | boolean
+  | null
+  | undefined
+  | JsonValue[]
+  | JsonRecord
+  | ((...args: JsonValue[]) => JsonValue);
+export interface JsonRecord {
+  [key: string]: JsonValue;
+}
+
+export type SendFn = (...args: unknown[]) => JsonValue | object | Promise<JsonValue | object>;
 
 export interface OpHandler {
-  spanName(input: Record<string, unknown>): string;
-  spanType(input: Record<string, unknown>): SpanType;
-  requestFields(input: Record<string, unknown>, options: BedrockInstrumentationOptions): SpanFields;
+  spanName(input: JsonRecord): string;
+  spanType(input: JsonRecord): SpanType;
+  requestFields(input: JsonRecord, options: BedrockInstrumentationOptions): SpanFields;
   onResult(
-    result: unknown,
+    result: JsonValue,
     span: SpanHandle,
     t0: number,
-    input: Record<string, unknown>,
+    input: JsonRecord,
     options: BedrockInstrumentationOptions,
-  ): unknown;
+  ): JsonValue | object;
 }
 
 export function createInstrumentedSend(
   originalSend: SendFn,
   options: BedrockInstrumentationOptions,
-  handlers: Record<string, OpHandler>,
+  handlers: { [key: string]: OpHandler },
 ): SendFn {
-  return function instrumentedSend(this: unknown, command: unknown, ...rest: unknown[]) {
-    const name = command?.constructor?.name;
-    const handler = typeof name === "string" ? handlers[name] : undefined;
+  return function instrumentedSend<TThis, TCommand, TRest extends unknown[]>(
+    this: TThis,
+    command: TCommand,
+    ...rest: TRest
+  ) {
+    const rawCommand = asJsonValue(command);
+    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, [command, ...rest]);
+      return originalSend.apply(this, [rawCommand, ...rawRest]);
     }
 
-    const input = commandInput(command);
+    const input = commandInput(rawCommand);
     const t0 = performance.now();
     const spanType = safe(() => handler.spanType(input)) ?? "span";
     const spanName = safe(() => handler.spanName(input)) ?? name ?? "bedrock";
     const fields = safe(() => handler.requestFields(input, options)) ?? {};
     const span = startSpan(spanName, { type: spanType, ...fields });
-    const callbackIndex = rest.findIndex((arg) => typeof arg === "function");
+    const callbackIndex = rawRest.findIndex((arg) => arg instanceof Function);
     if (callbackIndex !== -1) {
-      const callback = rest[callbackIndex] as (error: unknown, data?: unknown) => void;
-      const wrappedRest = [...rest];
-      wrappedRest[callbackIndex] = (error: unknown, data?: unknown) => {
-        if (error != null) {
-          endSpan(span, { ...awsMetadataFields(metadataOf(error)), error });
-          callback(error, data);
+      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 = data;
+        let nextData: JsonValue | object = asJsonValue(rawData);
         try {
-          nextData = handler.onResult(data, span, t0, input, options);
+          nextData = handler.onResult(asJsonValue(rawData), span, t0, input, options);
         } catch {
-          endSpan(span, awsMetadataFields(metadataOf(data)));
+          endSpan(span, awsMetadataFields(metadataOf(rawData)));
         }
-        callback(error, nextData);
+        callback(asJsonValue(rawError), asJsonValue(nextData));
       };
       try {
-        return originalSend.apply(this, [command, ...wrappedRest]);
+        return originalSend.apply(this, [rawCommand, ...wrappedRest]);
       } catch (error) {
-        endSpan(span, { ...awsMetadataFields(metadataOf(error)), error });
+        endSpan(span, { ...awsMetadataFields(metadataOf(error)), error: asError(error) });
         throw error;
       }
     }
 
-    let result: unknown;
+    let result: JsonValue | object | Promise<JsonValue | object>;
     try {
-      result = originalSend.apply(this, [command, ...rest]);
+      result = originalSend.apply(this, [rawCommand, ...rawRest]);
     } catch (error) {
-      endSpan(span, { ...awsMetadataFields(metadataOf(error)), error });
+      endSpan(span, { ...awsMetadataFields(metadataOf(error)), error: asError(error) });
       throw error;
     }
 
     return Promise.resolve(result).then(
       (resolved) => {
         try {
-          return handler.onResult(resolved, span, t0, input, options);
+          return handler.onResult(asJsonValue(resolved), span, t0, input, options);
         } catch {
           endSpan(span, awsMetadataFields(metadataOf(resolved)));
           return resolved;
         }
       },
       (error) => {
-        endSpan(span, { ...awsMetadataFields(metadataOf(error)), error });
+        endSpan(span, { ...awsMetadataFields(metadataOf(error)), error: asError(error) });
         throw error;
       },
     );
   };
 }
 
-export function commandInput(command: unknown): Record<string, unknown> {
-  const value = (command as { input?: unknown } | null | undefined)?.input;
-  return isRecord(value) ? value : {};
+export function commandInput<T>(command: T): JsonRecord {
+  if (!isRecord(command)) return {};
+  return isRecord(command.input) ? command.input : {};
 }
 
-export function isRecord(value: unknown): value is Record<string, unknown> {
-  return value !== null && typeof value === "object" && !Array.isArray(value);
+export function isRecord<T>(value: T): value is T & JsonRecord {
+  return value !== null && Object(value) === value && !Array.isArray(value);
 }
 
 export function safe<T>(fn: () => T): T | undefined {
@@ -128,15 +149,15 @@ export function updateSpan(span: SpanHandle, fields: SpanFields): void {
   }
 }
 
-export function metadataOf(value: unknown): unknown {
-  if (isRecord(value) && isRecord(value.$metadata)) return value.$metadata;
-  return (value as { $metadata?: unknown } | null | undefined)?.$metadata;
+export function metadataOf<T>(value: T): JsonValue {
+  if (!isRecord(value)) return undefined;
+  return value.$metadata;
 }
 
 export function modeledStreamError(
-  event: Record<string, unknown>,
+  event: JsonRecord,
 ): { error: Error; fields: SpanFields } | undefined {
-  const keys: Record<string, string> = {
+  const keys = {
     accessDeniedException: "AccessDeniedException",
     badGatewayException: "BadGatewayException",
     conflictException: "ConflictException",
@@ -150,7 +171,7 @@ export function modeledStreamError(
     serviceUnavailableException: "ServiceUnavailableException",
     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;
@@ -163,12 +184,13 @@ export function modeledStreamError(
   return undefined;
 }
 
-export function awsMetadataFields(metadata: unknown): SpanFields {
-  if (!isRecord(metadata)) return {};
-  const requestId = stringValue(metadata.requestId);
-  const attempts = numberValue(metadata.attempts);
-  const httpStatusCode = numberValue(metadata.httpStatusCode);
-  const totalRetryDelay = numberValue(metadata.totalRetryDelay);
+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;
@@ -194,41 +216,45 @@ export function mergeFields(...fields: Array<SpanFields | undefined>): SpanField
   return merged;
 }
 
-export function omitUndefined<T extends Record<string, unknown>>(value: T): T {
+export function omitUndefined<T extends object>(value: T): T {
   return Object.fromEntries(Object.entries(value).filter(([, item]) => item !== undefined)) as T;
 }
 
-export function stringValue(value: unknown): string | undefined {
-  return typeof value === "string" && value.length > 0 ? value : undefined;
+export function stringValue<T>(value: T): string | undefined {
+  const raw: unknown = value;
+  return String(raw) === raw && raw.length > 0 ? raw : undefined;
 }
 
-export function numberValue(value: unknown): number | undefined {
-  return typeof value === "number" && Number.isFinite(value) ? value : undefined;
+export function numberValue<T>(value: T): number | undefined {
+  const raw: unknown = value;
+  return Number(raw) === raw && Number.isFinite(raw) ? raw : undefined;
 }
 
-export function arrayValue<T>(value: unknown): T[] | undefined {
+export function arrayValue<T, TValue = unknown>(value: TValue): T[] | undefined {
   return Array.isArray(value) ? (value as T[]) : undefined;
 }
 
-export function bytesToString(value: unknown): string | undefined {
-  if (typeof value === "string") return value;
-  if (value instanceof Uint8Array) return new TextDecoder().decode(value);
-  if (ArrayBuffer.isView(value)) return new TextDecoder().decode(value as Uint8Array);
-  return 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));
+  }
 }
 
-export function parseJson(value: unknown): unknown {
+export function parseJson<T>(value: T): JsonValue {
   const text = bytesToString(value);
   if (text === undefined) return undefined;
   try {
-    return JSON.parse(text) as unknown;
+    return asJsonValue(JSON.parse(text));
   } catch {
     return undefined;
   }
 }
 
 export interface StreamState {
-  feed(event: unknown): void;
+  feed<T>(event: T): void;
   finish(partial: boolean): SpanFields;
 }
 
@@ -239,10 +265,7 @@ export function wrapAsyncIterable<T>(
   t0: number,
   baseFields: SpanFields,
 ): T {
-  if (
-    !iterable ||
-    typeof (iterable as unknown as AsyncIterable<unknown>)[Symbol.asyncIterator] !== "function"
-  ) {
+  if (!iterable || !isAsyncIterable(iterable)) {
     endSpan(span, baseFields);
     return iterable;
   }
@@ -257,11 +280,10 @@ export function wrapAsyncIterable<T>(
     endSpan(span, mergeFields(baseFields, fields));
   };
 
-  const source = iterable as unknown as AsyncIterable<unknown>;
   const generator = async function* () {
     let completed = false;
     try {
-      for await (const event of source) {
+      for await (const event of iterable) {
         if (!first) {
           first = true;
           updateSpan(span, { timeToFirstChunkMs: performance.now() - t0 });
@@ -271,17 +293,21 @@ export function wrapAsyncIterable<T>(
       }
       completed = true;
     } catch (error) {
-      endOnce(mergeFields(state.finish(true), awsMetadataFields(metadataOf(error)), { error }));
+      endOnce(
+        mergeFields(state.finish(true), awsMetadataFields(metadataOf(error)), {
+          error: asError(error),
+        }),
+      );
       throw error;
     } finally {
       if (!ended) endOnce(state.finish(!completed));
     }
   };
 
-  const wrapped: Record<PropertyKey, unknown> = {};
+  const wrapped: { [key: PropertyKey]: JsonValue | (() => AsyncGenerator<unknown>) } = {};
   if (isRecord(iterable)) {
     for (const key of Reflect.ownKeys(iterable)) {
-      wrapped[key] = (iterable as Record<PropertyKey, unknown>)[key];
+      wrapped[key] = (iterable as { [key: PropertyKey]: JsonValue })[key];
     }
   }
   const finishAbandoned = () => endOnce(safe(() => state.finish(true)));
@@ -294,7 +320,19 @@ export function wrapAsyncIterable<T>(
   return wrapped as T;
 }
 
-export function basenameArn(value: unknown): string | undefined {
+function isAsyncIterable<T>(value: T): value is T & AsyncIterable<unknown> {
+  return Object(value) === value && Symbol.asyncIterator in Object(value);
+}
+
+function asJsonValue<T>(value: T): JsonValue {
+  return value as JsonValue;
+}
+
+function asError<T>(value: T): Error {
+  return value instanceof Error ? value : new Error(String(value));
+}
+
+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

@@ -1,6 +1,7 @@
 import type { SpanFields } from "@telemetry-dev/sdk";
 
 import {
+  arrayValue,
   awsMetadataFields,
   isRecord,
   mergeFields,
@@ -10,22 +11,24 @@ import {
   parseJson,
   PROVIDER,
   stringValue,
+  type JsonRecord,
+  type JsonValue,
   type StreamState,
 } from "./internal.ts";
 
-export function isEmbeddingModel(modelId: unknown): boolean {
+export function isEmbeddingModel<T>(modelId: T): boolean {
   const model = stringValue(modelId)?.toLowerCase() ?? "";
   return (
     model.includes("titan-embed") || model.includes("cohere.embed") || model.includes("-embedding")
   );
 }
 
-export function invokeModelSpanName(input: Record<string, unknown>): string {
+export function invokeModelSpanName(input: JsonRecord): string {
   const model = stringValue(input.modelId) ?? "unknown";
   return `${isEmbeddingModel(model) ? "embeddings" : "chat"} ${model}`;
 }
 
-export function invokeModelRequestFields(input: Record<string, unknown>): SpanFields {
+export function invokeModelRequestFields(input: JsonRecord): SpanFields {
   const model = stringValue(input.modelId);
   const body = parseBody(input.body, input.contentType);
   return omitUndefined({
@@ -37,23 +40,20 @@ export function invokeModelRequestFields(input: Record<string, unknown>): SpanFi
   });
 }
 
-export function invokeModelResponseFields(
-  input: Record<string, unknown>,
-  output: unknown,
-): SpanFields {
+export function invokeModelResponseFields(input: JsonRecord, output: JsonValue): SpanFields {
   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(body: unknown, contentType: unknown): unknown {
+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(model: string | undefined, body: unknown): SpanFields {
+function nativeSamplingFields<T>(model: string | undefined, body: T): SpanFields {
   if (!isRecord(body)) return {};
   const lower = model?.toLowerCase() ?? "";
   if (lower.includes("amazon.titan")) {
@@ -62,7 +62,7 @@ function nativeSamplingFields(model: string | undefined, body: unknown): SpanFie
       temperature: numberValue(cfg.temperature),
       topP: numberValue(cfg.topP),
       maxTokens: numberValue(cfg.maxTokenCount),
-      stopSequences: Array.isArray(cfg.stopSequences) ? (cfg.stopSequences as string[]) : undefined,
+      stopSequences: arrayValue<string>(cfg.stopSequences),
     });
   }
   if (lower.includes("amazon.nova")) {
@@ -72,7 +72,7 @@ function nativeSamplingFields(model: string | undefined, body: unknown): SpanFie
       topP: numberValue(cfg.topP) ?? numberValue(cfg.top_p),
       topK: numberValue(cfg.topK),
       maxTokens: numberValue(cfg.maxTokens) ?? numberValue(cfg.max_new_tokens),
-      stopSequences: Array.isArray(cfg.stopSequences) ? (cfg.stopSequences as string[]) : undefined,
+      stopSequences: arrayValue<string>(cfg.stopSequences),
     });
   }
   if (lower.includes("anthropic.claude")) {
@@ -81,9 +81,7 @@ function nativeSamplingFields(model: string | undefined, body: unknown): SpanFie
       topP: numberValue(body.top_p),
       topK: numberValue(body.top_k),
       maxTokens: numberValue(body.max_tokens),
-      stopSequences: Array.isArray(body.stop_sequences)
-        ? (body.stop_sequences as string[])
-        : undefined,
+      stopSequences: arrayValue<string>(body.stop_sequences),
     });
   }
   if (lower.includes("meta.llama")) {
@@ -102,8 +100,8 @@ function nativeSamplingFields(model: string | undefined, body: unknown): SpanFie
 
 function nativeResponseFields(
   model: string | undefined,
-  body: unknown,
-  metadata: unknown,
+  body: JsonValue,
+  metadata: JsonValue,
 ): SpanFields {
   if (!isRecord(body)) return headerUsageFields(metadata);
   const lower = model?.toLowerCase() ?? "";
@@ -190,12 +188,12 @@ function nativeResponseFields(
   return withHeaderUsage({ output: body }, metadata);
 }
 
-function withHeaderUsage(fields: SpanFields, metadata: unknown): SpanFields {
+function withHeaderUsage<T>(fields: SpanFields, metadata: T): SpanFields {
   if (fields.usage && Object.keys(fields.usage).length > 0) return fields;
   return mergeFields(fields, headerUsageFields(metadata));
 }
 
-function headerUsageFields(metadata: unknown): SpanFields {
+function headerUsageFields<T>(metadata: T): SpanFields {
   const headers = isRecord(metadata)
     ? isRecord(metadata.httpHeaders)
       ? metadata.httpHeaders
@@ -210,9 +208,10 @@ function headerUsageFields(metadata: unknown): SpanFields {
   return { usage: Object.keys(usage).length > 0 ? usage : undefined };
 }
 
-function headerNumber(value: unknown): number | undefined {
-  if (typeof value === "string") {
-    const parsed = Number.parseInt(value, 10);
+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);
@@ -225,7 +224,7 @@ export class InvokeModelStreamState implements StreamState {
   private finishReason: string | undefined;
   private error: Error | undefined;
   private errorFields: SpanFields | undefined;
-  feed(event: unknown): void {
+  feed<T>(event: T): void {
     if (!isRecord(event)) return;
     const streamError = modeledStreamError(event);
     if (streamError) {
@@ -303,7 +302,7 @@ export class InvokeModelStreamState implements StreamState {
   }
 }
 
-function textFromProviderChunk(value: unknown): string {
+function textFromProviderChunk<T>(value: T): string {
   if (!isRecord(value)) return "";
   const delta = isRecord(value.delta) ? value.delta : undefined;
   const contentBlockDelta = isRecord(value.contentBlockDelta)