@telemetry-dev/bedrock 0.1.0

package/src/internal.ts

@@ -0,0 +1,301 @@
+import { startSpan, type SpanFields, type SpanHandle, type SpanType } from "@telemetry-dev/sdk";
+
+export const PROVIDER = "amazon-bedrock";
+
+const STREAM_FINALIZER =
+  typeof FinalizationRegistry === "function"
+    ? new FinalizationRegistry<() => void>((finish) => finish())
+    : undefined;
+
+export interface BedrockInstrumentationOptions {
+  captureAgentTrace?: boolean;
+}
+
+export type SendFn = (...args: unknown[]) => unknown;
+
+export interface OpHandler {
+  spanName(input: Record<string, unknown>): string;
+  spanType(input: Record<string, unknown>): SpanType;
+  requestFields(input: Record<string, unknown>, options: BedrockInstrumentationOptions): SpanFields;
+  onResult(
+    result: unknown,
+    span: SpanHandle,
+    t0: number,
+    input: Record<string, unknown>,
+    options: BedrockInstrumentationOptions,
+  ): unknown;
+}
+
+export function createInstrumentedSend(
+  originalSend: SendFn,
+  options: BedrockInstrumentationOptions,
+  handlers: Record<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;
+    if (!handler) {
+      return originalSend.apply(this, [command, ...rest]);
+    }
+
+    const input = commandInput(command);
+    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");
+    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);
+          return;
+        }
+        let nextData = data;
+        try {
+          nextData = handler.onResult(data, span, t0, input, options);
+        } catch {
+          endSpan(span, awsMetadataFields(metadataOf(data)));
+        }
+        callback(error, nextData);
+      };
+      try {
+        return originalSend.apply(this, [command, ...wrappedRest]);
+      } catch (error) {
+        endSpan(span, { ...awsMetadataFields(metadataOf(error)), error });
+        throw error;
+      }
+    }
+
+    let result: unknown;
+    try {
+      result = originalSend.apply(this, [command, ...rest]);
+    } catch (error) {
+      endSpan(span, { ...awsMetadataFields(metadataOf(error)), error });
+      throw error;
+    }
+
+    return Promise.resolve(result).then(
+      (resolved) => {
+        try {
+          return handler.onResult(resolved, span, t0, input, options);
+        } catch {
+          endSpan(span, awsMetadataFields(metadataOf(resolved)));
+          return resolved;
+        }
+      },
+      (error) => {
+        endSpan(span, { ...awsMetadataFields(metadataOf(error)), 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 isRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}
+
+export function safe<T>(fn: () => T): T | undefined {
+  try {
+    return fn();
+  } catch {
+    return undefined;
+  }
+}
+
+export function endSpan(span: SpanHandle, fields?: SpanFields): void {
+  try {
+    span.end(fields);
+  } catch {
+    // Instrumentation must never throw into caller code.
+  }
+}
+
+export function updateSpan(span: SpanHandle, fields: SpanFields): void {
+  try {
+    span.update(fields);
+  } catch {
+    // Instrumentation must never throw into caller code.
+  }
+}
+
+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 modeledStreamError(
+  event: Record<string, unknown>,
+): { error: Error; fields: SpanFields } | undefined {
+  const keys: Record<string, string> = {
+    accessDeniedException: "AccessDeniedException",
+    badGatewayException: "BadGatewayException",
+    conflictException: "ConflictException",
+    dependencyFailedException: "DependencyFailedException",
+    internalServerException: "InternalServerException",
+    modelNotReadyException: "ModelNotReadyException",
+    modelStreamErrorException: "ModelStreamErrorException",
+    modelTimeoutException: "ModelTimeoutException",
+    resourceNotFoundException: "ResourceNotFoundException",
+    serviceQuotaExceededException: "ServiceQuotaExceededException",
+    serviceUnavailableException: "ServiceUnavailableException",
+    throttlingException: "ThrottlingException",
+    validationException: "ValidationException",
+  };
+  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(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);
+  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,
+  });
+}
+
+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;
+}
+
+export function omitUndefined<T extends Record<string, unknown>>(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 numberValue(value: unknown): number | undefined {
+  return typeof value === "number" && Number.isFinite(value) ? value : undefined;
+}
+
+export function arrayValue<T>(value: unknown): 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 parseJson(value: unknown): unknown {
+  const text = bytesToString(value);
+  if (text === undefined) return undefined;
+  try {
+    return JSON.parse(text) as unknown;
+  } catch {
+    return undefined;
+  }
+}
+
+export interface StreamState {
+  feed(event: unknown): void;
+  finish(partial: boolean): SpanFields;
+}
+
+export function wrapAsyncIterable<T>(
+  iterable: T,
+  state: StreamState,
+  span: SpanHandle,
+  t0: number,
+  baseFields: SpanFields,
+): T {
+  if (
+    !iterable ||
+    typeof (iterable as unknown as AsyncIterable<unknown>)[Symbol.asyncIterator] !== "function"
+  ) {
+    endSpan(span, baseFields);
+    return iterable;
+  }
+
+  let ended = false;
+  let first = false;
+  const unregisterToken = {};
+  const endOnce = (fields?: SpanFields) => {
+    if (ended) return;
+    ended = true;
+    STREAM_FINALIZER?.unregister(unregisterToken);
+    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) {
+        if (!first) {
+          first = true;
+          updateSpan(span, { timeToFirstChunkMs: performance.now() - t0 });
+        }
+        safe(() => state.feed(event));
+        yield event;
+      }
+      completed = true;
+    } catch (error) {
+      endOnce(mergeFields(state.finish(true), awsMetadataFields(metadataOf(error)), { error }));
+      throw error;
+    } finally {
+      if (!ended) endOnce(state.finish(!completed));
+    }
+  };
+
+  const wrapped: Record<PropertyKey, unknown> = {};
+  if (isRecord(iterable)) {
+    for (const key of Reflect.ownKeys(iterable)) {
+      wrapped[key] = (iterable as Record<PropertyKey, unknown>)[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;
+}
+
+export function basenameArn(value: unknown): string | undefined {
+  const text = stringValue(value);
+  if (!text) return undefined;
+  return text.split(/[/:]/).filter(Boolean).at(-1) ?? text;
+}

package/src/invoke-model.ts

@@ -0,0 +1,332 @@
+import type { SpanFields } from "@telemetry-dev/sdk";
+
+import {
+  awsMetadataFields,
+  isRecord,
+  mergeFields,
+  modeledStreamError,
+  numberValue,
+  omitUndefined,
+  parseJson,
+  PROVIDER,
+  stringValue,
+  type StreamState,
+} from "./internal.ts";
+
+export function isEmbeddingModel(modelId: unknown): 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 {
+  const model = stringValue(input.modelId) ?? "unknown";
+  return `${isEmbeddingModel(model) ? "embeddings" : "chat"} ${model}`;
+}
+
+export function invokeModelRequestFields(input: Record<string, unknown>): SpanFields {
+  const model = stringValue(input.modelId);
+  const body = parseBody(input.body, input.contentType);
+  return omitUndefined({
+    provider: PROVIDER,
+    model,
+    input: body,
+    outputType: isEmbeddingModel(model) ? "embedding" : undefined,
+    ...nativeSamplingFields(model, body),
+  });
+}
+
+export function invokeModelResponseFields(
+  input: Record<string, unknown>,
+  output: unknown,
+): 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 {
+  const type = stringValue(contentType) ?? "application/json";
+  if (!type.includes("json")) return undefined;
+  return parseJson(body);
+}
+
+function nativeSamplingFields(model: string | undefined, body: unknown): 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),
+      maxTokens: numberValue(cfg.maxTokenCount),
+      stopSequences: Array.isArray(cfg.stopSequences) ? (cfg.stopSequences as string[]) : undefined,
+    });
+  }
+  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),
+      topK: numberValue(cfg.topK),
+      maxTokens: numberValue(cfg.maxTokens) ?? numberValue(cfg.max_new_tokens),
+      stopSequences: Array.isArray(cfg.stopSequences) ? (cfg.stopSequences as string[]) : undefined,
+    });
+  }
+  if (lower.includes("anthropic.claude")) {
+    return omitUndefined({
+      temperature: numberValue(body.temperature),
+      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,
+    });
+  }
+  if (lower.includes("meta.llama")) {
+    return omitUndefined({
+      temperature: numberValue(body.temperature),
+      topP: numberValue(body.top_p),
+      maxTokens: numberValue(body.max_gen_len),
+    });
+  }
+  return omitUndefined({
+    temperature: numberValue(body.temperature),
+    topP: numberValue(body.top_p) ?? numberValue(body.topP),
+    maxTokens: numberValue(body.max_tokens) ?? numberValue(body.maxTokens),
+  });
+}
+
+function nativeResponseFields(
+  model: string | undefined,
+  body: unknown,
+  metadata: unknown,
+): SpanFields {
+  if (!isRecord(body)) return headerUsageFields(metadata);
+  const lower = model?.toLowerCase() ?? "";
+  if (isEmbeddingModel(model)) {
+    return withHeaderUsage(
+      {
+        output: body,
+        outputType: "embedding",
+        usage: omitUndefined({ inputTokens: numberValue(body.inputTextTokenCount) }),
+      },
+      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),
+      usage: omitUndefined({
+        inputTokens: numberValue(body.inputTextTokenCount),
+        outputTokens: numberValue(first.tokenCount),
+      }),
+    });
+  }
+  if (lower.includes("amazon.nova")) {
+    const usage = isRecord(body.usage) ? body.usage : {};
+    return omitUndefined({
+      output: body,
+      finishReason: stringValue(body.stopReason),
+      usage: omitUndefined({
+        inputTokens: numberValue(usage.inputTokens),
+        outputTokens: numberValue(usage.outputTokens),
+        totalTokens: numberValue(usage.totalTokens),
+      }),
+    });
+  }
+  if (lower.includes("anthropic.claude")) {
+    const usage = isRecord(body.usage) ? body.usage : {};
+    return withHeaderUsage(
+      {
+        output: body,
+        finishReason: stringValue(body.stop_reason),
+        usage: omitUndefined({
+          inputTokens: numberValue(usage.input_tokens),
+          outputTokens: numberValue(usage.output_tokens),
+          cacheReadInputTokens: numberValue(usage.cache_read_input_tokens),
+          cacheCreationInputTokens: numberValue(usage.cache_creation_input_tokens),
+        }),
+      },
+      metadata,
+    );
+  }
+  if (lower.includes("meta.llama")) {
+    return omitUndefined({
+      output: body,
+      finishReason: stringValue(body.stop_reason),
+      usage: omitUndefined({
+        inputTokens: numberValue(body.prompt_token_count),
+        outputTokens: numberValue(body.generation_token_count),
+      }),
+    });
+  }
+  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(fields: SpanFields, metadata: unknown): SpanFields {
+  if (fields.usage && Object.keys(fields.usage).length > 0) return fields;
+  return mergeFields(fields, headerUsageFields(metadata));
+}
+
+function headerUsageFields(metadata: unknown): SpanFields {
+  const headers = isRecord(metadata)
+    ? isRecord(metadata.httpHeaders)
+      ? metadata.httpHeaders
+      : isRecord(metadata.HTTPHeaders)
+        ? 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(value: unknown): number | undefined {
+  if (typeof value === "string") {
+    const parsed = Number.parseInt(value, 10);
+    return Number.isFinite(parsed) ? parsed : undefined;
+  }
+  return numberValue(value);
+}
+
+export class InvokeModelStreamState implements StreamState {
+  private readonly chunks: unknown[] = [];
+  private text = "";
+  private usage: SpanFields["usage"];
+  private finishReason: string | undefined;
+  private error: Error | undefined;
+  private errorFields: SpanFields | undefined;
+  feed(event: unknown): void {
+    if (!isRecord(event)) return;
+    const streamError = modeledStreamError(event);
+    if (streamError) {
+      this.error = streamError.error;
+      this.errorFields = streamError.fields;
+      return;
+    }
+    const chunk = isRecord(event.chunk) ? event.chunk : undefined;
+    const parsed = parseJson(chunk?.bytes);
+    if (parsed === undefined) return;
+    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) ??
+        stringValue(outputs?.stop_reason) ??
+        stringValue(generation?.finish_reason) ??
+        stringValue(parsed.completionReason) ??
+        this.finishReason;
+      this.usage = omitUndefined({
+        ...this.usage,
+        inputTokens:
+          numberValue(messageUsage?.input_tokens) ??
+          numberValue(metrics?.inputTokenCount) ??
+          numberValue(parsed.inputTextTokenCount) ??
+          this.usage?.inputTokens,
+        outputTokens:
+          numberValue(usage?.output_tokens) ??
+          numberValue(metrics?.outputTokenCount) ??
+          numberValue(parsed.totalOutputTextTokenCount) ??
+          this.usage?.outputTokens,
+        cacheReadInputTokens:
+          numberValue(messageUsage?.cache_read_input_tokens) ?? this.usage?.cacheReadInputTokens,
+        cacheCreationInputTokens:
+          numberValue(messageUsage?.cache_creation_input_tokens) ??
+          this.usage?.cacheCreationInputTokens,
+      });
+    }
+  }
+
+  finish(): SpanFields {
+    return mergeFields(
+      omitUndefined({
+        output: this.text
+          ? [
+              {
+                role: "assistant",
+                parts: [{ type: "text", content: this.text }],
+                finish_reason: this.finishReason,
+              },
+            ]
+          : this.chunks,
+        finishReason: this.finishReason,
+        usage: this.usage,
+        error: this.error,
+      }),
+      this.errorFields,
+    );
+  }
+}
+
+function textFromProviderChunk(value: unknown): 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) ??
+    stringValue(value.completion) ??
+    stringValue(value.text) ??
+    stringValue(delta?.text) ??
+    stringValue(nestedDelta?.text) ??
+    (Array.isArray(value.outputs) && isRecord(value.outputs[0])
+      ? stringValue(value.outputs[0].text)
+      : undefined) ??
+    stringValue(generation?.text) ??
+    ""
+  );
+}