@telemetry-dev/bedrock 0.1.0

package/package.json

@@ -0,0 +1,75 @@
+{
+  "name": "@telemetry-dev/bedrock",
+  "version": "0.1.0",
+  "description": "AWS Bedrock instrumentation for telemetry.dev: Converse, InvokeModel, and Bedrock Agents spans for AWS SDK v3 clients.",
+  "keywords": [
+    "aws",
+    "bedrock",
+    "genai",
+    "llm",
+    "observability",
+    "opentelemetry",
+    "telemetry",
+    "tracing"
+  ],
+  "homepage": "https://telemetry.dev",
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/telemetry-dev/telemetry.dev.git",
+    "directory": "packages/bedrock"
+  },
+  "files": [
+    "dist",
+    "src"
+  ],
+  "type": "module",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.mts",
+      "import": "./dist/index.mjs",
+      "default": "./dist/index.mjs"
+    },
+    "./agents": {
+      "types": "./dist/agents.d.mts",
+      "import": "./dist/agents.mjs",
+      "default": "./dist/agents.mjs"
+    },
+    "./package.json": "./package.json"
+  },
+  "publishConfig": {
+    "access": "public"
+  },
+  "devDependencies": {
+    "@aws-sdk/client-bedrock-agent-runtime": "^3.1000.0",
+    "@aws-sdk/client-bedrock-runtime": "^3.1000.0",
+    "@opentelemetry/api": "^1.9.1",
+    "@opentelemetry/sdk-logs": "^0.218.0",
+    "@opentelemetry/sdk-trace-base": "^2.7.1",
+    "@smithy/eventstream-codec": "^4.2.5",
+    "@smithy/protocol-http": "^5.3.10",
+    "@smithy/types": "^4.12.0",
+    "@types/node": "^25.5.0",
+    "@typescript/native-preview": "7.0.0-dev.20260328.1",
+    "typescript": "^6.0.2",
+    "vite-plus": "0.1.20",
+    "vitest": "npm:@voidzero-dev/vite-plus-test@0.1.20",
+    "@telemetry-dev/sdk": "^0.1.0"
+  },
+  "peerDependencies": {
+    "@aws-sdk/client-bedrock-agent-runtime": ">=3.600.0 <4",
+    "@aws-sdk/client-bedrock-runtime": ">=3.600.0 <4",
+    "@telemetry-dev/sdk": "^0.1.0"
+  },
+  "peerDependenciesMeta": {
+    "@aws-sdk/client-bedrock-agent-runtime": {
+      "optional": true
+    }
+  },
+  "scripts": {
+    "build": "pnpm exec vp pack",
+    "dev": "pnpm exec vp pack --watch",
+    "test": "vp test",
+    "check": "vp check"
+  }
+}

package/src/agent-handlers.ts

@@ -0,0 +1,315 @@
+import type { SpanFields } from "@telemetry-dev/sdk";
+
+import {
+  arrayValue,
+  awsMetadataFields,
+  basenameArn,
+  bytesToString,
+  isRecord,
+  mergeFields,
+  modeledStreamError,
+  numberValue,
+  omitUndefined,
+  PROVIDER,
+  stringValue,
+  type BedrockInstrumentationOptions,
+  type StreamState,
+} from "./internal.ts";
+
+export function agentInputMessages(inputText: unknown): unknown {
+  const text = stringValue(inputText);
+  return text ? [{ role: "user", parts: [{ type: "text", content: text }] }] : undefined;
+}
+
+export function agentRequestFields(input: Record<string, unknown>): SpanFields {
+  return omitUndefined({
+    provider: PROVIDER,
+    input: agentInputMessages(input.inputText),
+    agentId: stringValue(input.agentId),
+    metadata: omitUndefined({
+      bedrock_agent_alias_id: stringValue(input.agentAliasId),
+      bedrock_session_id: stringValue(input.sessionId),
+      bedrock_memory_id: stringValue(input.memoryId),
+    }),
+  });
+}
+
+export function inlineAgentRequestFields(input: Record<string, unknown>): SpanFields {
+  return omitUndefined({
+    input: agentInputMessages(input.inputText),
+    model: stringValue(input.foundationModel),
+    provider: PROVIDER,
+    agentName: stringValue(input.agentName),
+    metadata: omitUndefined({ bedrock_session_id: stringValue(input.sessionId) }),
+  });
+}
+
+export class AgentStreamState implements StreamState {
+  private text = "";
+  private usage: NonNullable<SpanFields["usage"]> = {};
+  private traceEventCount = 0;
+  private failureReason: string | undefined;
+  private guardrailAction: string | undefined;
+  private returnControl: unknown;
+  private readonly traces: unknown[] = [];
+  private error: Error | undefined;
+  private errorFields: SpanFields | undefined;
+  private citationCount = 0;
+  private readonly outputFiles: unknown[] = [];
+
+  constructor(private readonly options: BedrockInstrumentationOptions) {}
+
+  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 bytes = chunk?.bytes;
+    if (bytes !== undefined) this.text += bytesToString(bytes) ?? "";
+    const attribution = isRecord(chunk?.attribution) ? chunk.attribution : undefined;
+    const citations = arrayValue<unknown>(attribution?.citations);
+    if (citations) this.citationCount += citations.length;
+
+    const files = isRecord(event.files) ? arrayValue<unknown>(event.files.files) : undefined;
+    if (files) this.outputFiles.push(...files.map(sanitizeOutputFile));
+
+    const returnControl = isRecord(event.returnControl) ? event.returnControl : undefined;
+    if (returnControl) this.returnControl = returnControl;
+
+    const traceEvent = isRecord(event.trace) ? event.trace : undefined;
+    if (traceEvent) {
+      this.traceEventCount += 1;
+      if (this.options.captureAgentTrace) this.traces.push(traceEvent);
+      collectTrace(traceEvent, (usage) => this.addUsage(usage));
+      const failure = findKey(traceEvent, "failureTrace");
+      if (isRecord(failure))
+        this.failureReason = stringValue(failure.failureReason) ?? this.failureReason;
+      const guardrail = findKey(traceEvent, "guardrailTrace");
+      if (isRecord(guardrail))
+        this.guardrailAction = stringValue(guardrail.action) ?? this.guardrailAction;
+    }
+  }
+
+  finish(): SpanFields {
+    const metadata = omitUndefined({
+      trace_event_count: this.traceEventCount || undefined,
+      failure_reason: this.failureReason,
+      guardrail_action: this.guardrailAction,
+      return_control: this.returnControl ? true : undefined,
+      citation_count: this.citationCount || undefined,
+      output_files: this.outputFiles.length > 0 ? this.outputFiles : undefined,
+      agent_trace:
+        this.options.captureAgentTrace && this.traces.length > 0 ? this.traces : undefined,
+    });
+    return mergeFields(
+      omitUndefined({
+        output: this.returnControl
+          ? { returnControl: this.returnControl }
+          : this.text
+            ? [{ role: "assistant", parts: [{ type: "text", content: this.text }] }]
+            : undefined,
+        usage: Object.keys(this.usage).length > 0 ? this.usage : undefined,
+        metadata,
+        error: this.error,
+      }),
+      this.errorFields,
+    );
+  }
+  private addUsage(value: unknown): void {
+    if (!isRecord(value)) return;
+    const inputTokens = numberValue(value.inputTokens);
+    const outputTokens = numberValue(value.outputTokens);
+    const totalTokens = numberValue(value.totalTokens);
+    if (inputTokens !== undefined)
+      this.usage.inputTokens = (this.usage.inputTokens ?? 0) + inputTokens;
+    if (outputTokens !== undefined)
+      this.usage.outputTokens = (this.usage.outputTokens ?? 0) + outputTokens;
+    if (totalTokens !== undefined)
+      this.usage.totalTokens = (this.usage.totalTokens ?? 0) + totalTokens;
+    else if (inputTokens !== undefined || outputTokens !== undefined) {
+      this.usage.totalTokens = (this.usage.inputTokens ?? 0) + (this.usage.outputTokens ?? 0);
+    }
+  }
+}
+
+function sanitizeOutputFile(value: unknown): unknown {
+  if (!isRecord(value)) return value;
+  return omitUndefined({
+    name: stringValue(value.name),
+    type: stringValue(value.type),
+  });
+}
+
+export function retrieveRequestFields(input: Record<string, unknown>): SpanFields {
+  return omitUndefined({
+    provider: PROVIDER,
+    input: isRecord(input.retrievalQuery) ? input.retrievalQuery : undefined,
+    metadata: { knowledge_base_id: stringValue(input.knowledgeBaseId) },
+  });
+}
+
+export function retrieveResponseFields(output: unknown): SpanFields {
+  const result = isRecord(output) ? output : {};
+  const retrievalResults = arrayValue<unknown>(result.retrievalResults);
+  return mergeFields(awsMetadataFields(result.$metadata), {
+    output: retrievalResults,
+    metadata: omitUndefined({
+      citation_count: retrievalResults?.length,
+      guardrail_action: stringValue(result.guardrailAction),
+    }),
+  });
+}
+
+export function ragModel(input: Record<string, unknown>): string | undefined {
+  const config = isRecord(input.retrieveAndGenerateConfiguration)
+    ? input.retrieveAndGenerateConfiguration
+    : undefined;
+  const kb = isRecord(config?.knowledgeBaseConfiguration)
+    ? config?.knowledgeBaseConfiguration
+    : undefined;
+  const external = isRecord(config?.externalSourcesConfiguration)
+    ? config?.externalSourcesConfiguration
+    : undefined;
+  return stringValue(kb?.modelArn) ?? stringValue(external?.modelArn);
+}
+
+export function ragRequestFields(input: Record<string, unknown>): SpanFields {
+  const text = isRecord(input.input) ? stringValue(input.input.text) : undefined;
+  return omitUndefined({
+    provider: PROVIDER,
+    model: ragModel(input),
+    input: text ? [{ role: "user", parts: [{ type: "text", content: text }] }] : undefined,
+    metadata: omitUndefined({ bedrock_session_id: stringValue(input.sessionId) }),
+  });
+}
+
+export function ragResponseFields(output: unknown): SpanFields {
+  const result = isRecord(output) ? output : {};
+  const text = isRecord(result.output) ? stringValue(result.output.text) : undefined;
+  return mergeFields(awsMetadataFields(result.$metadata), {
+    output: text ? [{ role: "assistant", parts: [{ type: "text", content: text }] }] : undefined,
+    metadata: omitUndefined({
+      citation_count: arrayValue<unknown>(result.citations)?.length,
+      guardrail_action: stringValue(result.guardrailAction),
+      bedrock_session_id: stringValue(result.sessionId),
+    }),
+  });
+}
+
+export class RagStreamState implements StreamState {
+  private text = "";
+  private citationCount = 0;
+  private guardrailAction: 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 output = isRecord(event.output) ? event.output : undefined;
+    this.text += stringValue(output?.text) ?? "";
+    if (event.citation) this.citationCount += 1;
+    const guardrail = isRecord(event.guardrail) ? event.guardrail : undefined;
+    this.guardrailAction = stringValue(guardrail?.action) ?? this.guardrailAction;
+  }
+
+  finish(): SpanFields {
+    return mergeFields(
+      omitUndefined({
+        output: this.text
+          ? [{ role: "assistant", parts: [{ type: "text", content: this.text }] }]
+          : undefined,
+        metadata: omitUndefined({
+          citation_count: this.citationCount || undefined,
+          guardrail_action: this.guardrailAction,
+        }),
+        error: this.error,
+      }),
+      this.errorFields,
+    );
+  }
+}
+
+export function flowRequestFields(input: Record<string, unknown>): SpanFields {
+  return omitUndefined({
+    provider: PROVIDER,
+    input: input.inputs,
+    agentId: stringValue(input.flowIdentifier),
+  });
+}
+
+export class FlowStreamState implements StreamState {
+  private outputs: unknown[] = [];
+  private completionReason: 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 output = isRecord(event.flowOutputEvent) ? event.flowOutputEvent : undefined;
+    if (output) this.outputs.push(output.content);
+    const inputRequest = isRecord(event.flowMultiTurnInputRequestEvent)
+      ? event.flowMultiTurnInputRequestEvent
+      : undefined;
+    if (inputRequest) this.outputs.push(inputRequest.content);
+    const completion = isRecord(event.flowCompletionEvent) ? event.flowCompletionEvent : undefined;
+    this.completionReason = stringValue(completion?.completionReason) ?? this.completionReason;
+  }
+
+  finish(): SpanFields {
+    return mergeFields(
+      omitUndefined({
+        output: this.outputs.length > 0 ? this.outputs : undefined,
+        finishReason: this.completionReason,
+        error: this.error,
+      }),
+      this.errorFields,
+    );
+  }
+}
+
+export function flowName(input: Record<string, unknown>): string {
+  return `invoke_flow ${stringValue(input.flowIdentifier) ?? "unknown"}`;
+}
+
+export function ragName(input: Record<string, unknown>): string {
+  return `retrieve_and_generate ${basenameArn(ragModel(input)) ?? "knowledge-base"}`;
+}
+
+function collectTrace(value: unknown, onUsage: (usage: unknown) => void): void {
+  if (!isRecord(value)) return;
+  const output = isRecord(value.modelInvocationOutput) ? value.modelInvocationOutput : undefined;
+  if (output) {
+    const metadata = isRecord(output.metadata) ? output.metadata : undefined;
+    onUsage(metadata?.usage);
+  }
+  for (const child of Object.values(value)) {
+    if (isRecord(child)) collectTrace(child, onUsage);
+    else if (Array.isArray(child)) child.forEach((item) => collectTrace(item, onUsage));
+  }
+}
+
+function findKey(value: unknown, key: string): unknown {
+  if (!isRecord(value)) return undefined;
+  if (value[key] !== undefined) return value[key];
+  for (const child of Object.values(value)) {
+    const found = findKey(child, key);
+    if (found !== undefined) return found;
+  }
+  return undefined;
+}

package/src/agents.ts

@@ -0,0 +1,69 @@
+import { BedrockAgentRuntimeClient } from "@aws-sdk/client-bedrock-agent-runtime";
+
+import { AGENT_HANDLERS } from "./internal-handlers.ts";
+import {
+  createInstrumentedSend,
+  type BedrockInstrumentationOptions,
+  type SendFn,
+} from "./internal.ts";
+
+export type { BedrockInstrumentationOptions } from "./internal.ts";
+
+export interface WrappableBedrockAgentClient {
+  send: (...args: never[]) => unknown;
+}
+
+const WRAPPED = new WeakSet<object>();
+let ORIGINAL: { proto: object; hadOwn: boolean; value: unknown; wrapper: SendFn } | undefined;
+
+export function wrapBedrockAgents<T extends WrappableBedrockAgentClient>(
+  client: T,
+  options: BedrockInstrumentationOptions = {},
+): T {
+  if (WRAPPED.has(client as object)) return client;
+  const original = client.send as SendFn;
+  Object.defineProperty(client, "send", {
+    value: createInstrumentedSend(original, options, AGENT_HANDLERS),
+    configurable: true,
+    writable: true,
+  });
+  WRAPPED.add(client as object);
+  return client;
+}
+
+export function instrumentBedrockAgents(options: BedrockInstrumentationOptions = {}): void {
+  if (ORIGINAL) return;
+  const proto = BedrockAgentRuntimeClient.prototype as unknown as { send: SendFn };
+  const hadOwn = Object.prototype.hasOwnProperty.call(proto, "send");
+  const original = proto.send;
+  const instrumentedSend = createInstrumentedSend(original, options, AGENT_HANDLERS);
+  const wrapper = function telemetryDevBedrockAgentSend(
+    this: WrappableBedrockAgentClient,
+    ...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,
+    writable: true,
+    value: wrapper,
+  });
+}
+
+export function uninstrumentBedrockAgents(): 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,
+      writable: true,
+      value: original.value,
+    });
+  } else {
+    Reflect.deleteProperty(original.proto, "send");
+  }
+}