@telemetry-dev/bedrock 0.1.6 → 0.1.8

package/README.md

@@ -13,6 +13,7 @@ This package instruments AWS SDK for JavaScript v3 clients for:
 - `@aws-sdk/client-bedrock-agent-runtime`
   - `InvokeAgent`
   - `InvokeInlineAgent`
+  - `Rerank`
   - `Retrieve`
   - `RetrieveAndGenerate`
   - `RetrieveAndGenerateStream`
@@ -113,6 +114,7 @@ All instrumented AWS responses/errors include `gen_ai.response.id` from the AWS
 | `InvokeModelWithResponseStream`     | `generation` or `embedding` | `chat {modelId}` / `embeddings {modelId}`             | provider-native chunk text where known, optional final `amazon-bedrock-invocationMetrics` usage                                                                                                        |
 | `ApplyGuardrail`                    | `span`                      | `apply_guardrail {guardrailIdentifier}`               | guardrail input/output, action, action reason, request id                                                                                                                                              |
 | `InvokeAgent` / `InvokeInlineAgent` | `agent`                     | `invoke_agent {agentId}` / `invoke_agent {agentName}` | user input, streamed answer, session/memory ids, alias id, trace usage aggregation, trace event count, return-control output                                                                           |
+| `Rerank`                            | `span` (`rerank` operation) | `rerank {modelArn basename}`                          | queries, inline sources, model request fields, ranked results and scores, result counts, pagination metadata, request id                                                                               |
 | `Retrieve`                          | `span`                      | `retrieve {knowledgeBaseId}`                          | query, retrieval results, guardrail action, result count metadata                                                                                                                                      |
 | `RetrieveAndGenerate` / stream      | `generation`                | `retrieve_and_generate {modelArn basename}`           | input text, generated text, citations count, session id, guardrail action                                                                                                                              |
 | `InvokeFlow`                        | `agent`                     | `invoke_flow {flowIdentifier}`                        | inputs, flow output events, completion reason                                                                                                                                                          |
@@ -149,7 +151,6 @@ The following calls pass through unless a supported operation above is used:
 - `InvokeModelWithBidirectionalStream`
 - `StartAsyncInvoke`, `GetAsyncInvoke`, `ListAsyncInvokes`
 - Agent Runtime session CRUD
-- `Rerank`
 - `GenerateQuery`
 - `OptimizePrompt`
 - `AgenticRetrieveStream`

package/dist/agents.mjs

@@ -1,4 +1,4 @@
-import { r as createInstrumentedSend, t as AGENT_HANDLERS } from "./internal-handlers-B47qvHGh.mjs";
+import { r as createInstrumentedSend, t as AGENT_HANDLERS } from "./internal-handlers-Byxp14OZ.mjs";
 import { BedrockAgentRuntimeClient } from "@aws-sdk/client-bedrock-agent-runtime";
 //#region src/agents.ts
 const WRAPPED = /* @__PURE__ */ new WeakSet();

package/dist/index.mjs

@@ -1,4 +1,4 @@
-import { n as RUNTIME_HANDLERS, r as createInstrumentedSend } from "./internal-handlers-B47qvHGh.mjs";
+import { n as RUNTIME_HANDLERS, r as createInstrumentedSend } from "./internal-handlers-Byxp14OZ.mjs";
 import { BedrockRuntimeClient } from "@aws-sdk/client-bedrock-runtime";
 //#region src/index.ts
 const WRAPPED = /* @__PURE__ */ new WeakSet();

package/dist/{internal-handlers-B47qvHGh.mjs → internal-handlers-Byxp14OZ.mjs}

@@ -1,4 +1,4 @@
-import { startSpan } from "@telemetry-dev/sdk";
+import { boundedCapture, captureEnabled, startSpan } from "@telemetry-dev/sdk";
 //#region src/internal.ts
 const PROVIDER = "amazon-bedrock";
 const STREAM_FINALIZER = typeof FinalizationRegistry === "function" ? new FinalizationRegistry((finish) => finish()) : void 0;
@@ -89,9 +89,23 @@ function safe(fn) {
 		return;
 	}
 }
+function withCacheInclusiveInput(fields) {
+	const usage = fields?.usage;
+	if (usage?.inputTokens === void 0) return fields;
+	const inputTokens = usage.inputTokens + (usage.cacheReadInputTokens ?? 0) + (usage.cacheCreationInputTokens ?? 0);
+	const totalTokens = usage.totalTokens !== void 0 && usage.outputTokens !== void 0 ? inputTokens + usage.outputTokens : usage.totalTokens;
+	return {
+		...fields,
+		usage: {
+			...usage,
+			inputTokens,
+			totalTokens
+		}
+	};
+}
 function endSpan(span, fields) {
 	try {
-		span.end(fields);
+		span.end(withCacheInclusiveInput(fields));
 	} catch {}
 }
 function updateSpan(span, fields) {
@@ -431,6 +445,53 @@ function retrieveResponseFields(output) {
 		})
 	});
 }
+function rerankConfiguration(input) {
+	const configuration = isRecord(input.rerankingConfiguration) ? input.rerankingConfiguration : void 0;
+	const bedrock = isRecord(configuration?.bedrockRerankingConfiguration) ? configuration.bedrockRerankingConfiguration : void 0;
+	return {
+		bedrock,
+		model: isRecord(bedrock?.modelConfiguration) ? bedrock.modelConfiguration : void 0
+	};
+}
+function rerankModel(input) {
+	return stringValue(rerankConfiguration(input).model?.modelArn);
+}
+function rerankName(input) {
+	return `rerank ${rerankModel(input)?.split("/").at(-1) ?? "unknown"}`;
+}
+function rerankRequestFields(input) {
+	const { bedrock, model } = rerankConfiguration(input);
+	const capture = captureEnabled("input") ? boundedCapture(omitUndefined({
+		queries: input.queries,
+		sources: input.sources,
+		additionalModelRequestFields: model?.additionalModelRequestFields,
+		nextToken: input.nextToken
+	})) : void 0;
+	return omitUndefined({
+		provider: PROVIDER,
+		model: stringValue(model?.modelArn),
+		input: capture?.value,
+		attributes: {
+			"gen_ai.operation.name": "rerank",
+			...capture?.truncated ? { "telemetry.dev.capture.truncated": true } : void 0
+		},
+		metadata: omitUndefined({ requested_result_count: numberValue(bedrock?.numberOfResults) })
+	});
+}
+function rerankResponseFields(output) {
+	const raw = output;
+	const result = isRecord(raw) ? raw : {};
+	const results = arrayValue(result.results);
+	const capture = captureEnabled("output") ? boundedCapture(results) : void 0;
+	return mergeFields(awsMetadataFields(result.$metadata), {
+		output: capture?.value,
+		attributes: capture?.truncated ? { "telemetry.dev.capture.truncated": true } : void 0,
+		metadata: omitUndefined({
+			result_count: results?.length,
+			has_next_token: stringValue(result.nextToken) ? true : void 0
+		})
+	});
+}
 function ragModel(input) {
 	const config = isRecord(input.retrieveAndGenerateConfiguration) ? input.retrieveAndGenerateConfiguration : void 0;
 	const kb = isRecord(config?.knowledgeBaseConfiguration) ? config?.knowledgeBaseConfiguration : void 0;
@@ -1212,6 +1273,15 @@ const AGENT_HANDLERS = {
 		requestFields: inlineAgentRequestFields,
 		onResult: streamAgentResult("completion")
 	},
+	RerankCommand: {
+		spanName: rerankName,
+		spanType: () => "span",
+		requestFields: rerankRequestFields,
+		onResult(result, span) {
+			endSpan(span, rerankResponseFields(result));
+			return result;
+		}
+	},
 	RetrieveCommand: {
 		spanName: (input) => `retrieve ${stringValue(input.knowledgeBaseId) ?? "unknown"}`,
 		spanType: () => "span",

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@telemetry-dev/bedrock",
-  "version": "0.1.6",
+  "version": "0.1.8",
   "description": "AWS Bedrock instrumentation for telemetry.dev: Converse, InvokeModel, and Bedrock Agents spans for AWS SDK v3 clients.",
   "keywords": [
     "aws",
@@ -49,7 +49,7 @@
     "@smithy/eventstream-codec": "^4.2.5",
     "@smithy/protocol-http": "^5.3.10",
     "@smithy/types": "^4.12.0",
-    "@telemetry-dev/sdk": "^0.1.5",
+    "@telemetry-dev/sdk": "^0.1.6",
     "@types/node": "^25.5.0",
     "@typescript/native-preview": "7.0.0-dev.20260328.1",
     "typescript": "^6.0.2",
@@ -59,7 +59,7 @@
   "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.5"
+    "@telemetry-dev/sdk": "^0.1.6"
   },
   "peerDependenciesMeta": {
     "@aws-sdk/client-bedrock-agent-runtime": {

package/src/agent-handlers.ts

@@ -1,4 +1,4 @@
-import type { SpanFields } from "@telemetry-dev/sdk";
+import { boundedCapture, captureEnabled, type SpanFields } from "@telemetry-dev/sdk";
 
 import {
   arrayValue,
@@ -191,6 +191,75 @@ export function retrieveResponseFields<T>(output: T): SpanFields {
   });
 }
 
+function rerankConfiguration(input: JsonRecord) {
+  const configuration = isRecord(input.rerankingConfiguration)
+    ? input.rerankingConfiguration
+    : undefined;
+
+  const bedrock = isRecord(configuration?.bedrockRerankingConfiguration)
+    ? configuration.bedrockRerankingConfiguration
+    : undefined;
+
+  return {
+    bedrock,
+    model: isRecord(bedrock?.modelConfiguration) ? bedrock.modelConfiguration : undefined,
+  };
+}
+
+export function rerankModel(input: JsonRecord): string | undefined {
+  return stringValue(rerankConfiguration(input).model?.modelArn);
+}
+
+export function rerankName(input: JsonRecord): string {
+  const model = rerankModel(input);
+
+  return `rerank ${model?.split("/").at(-1) ?? "unknown"}`;
+}
+
+export function rerankRequestFields(input: JsonRecord): SpanFields {
+  const { bedrock, model } = rerankConfiguration(input);
+
+  const capture = captureEnabled("input")
+    ? boundedCapture(
+        omitUndefined({
+          queries: input.queries,
+          sources: input.sources,
+          additionalModelRequestFields: model?.additionalModelRequestFields,
+          nextToken: input.nextToken,
+        }),
+      )
+    : undefined;
+
+  return omitUndefined({
+    provider: PROVIDER,
+    model: stringValue(model?.modelArn),
+    input: capture?.value,
+    attributes: {
+      "gen_ai.operation.name": "rerank",
+      ...(capture?.truncated ? { "telemetry.dev.capture.truncated": true } : undefined),
+    },
+    metadata: omitUndefined({
+      requested_result_count: numberValue(bedrock?.numberOfResults),
+    }),
+  });
+}
+
+export function rerankResponseFields<T>(output: T): SpanFields {
+  const raw: unknown = output;
+  const result = isRecord(raw) ? raw : {};
+  const results = arrayValue<unknown>(result.results);
+  const capture = captureEnabled("output") ? boundedCapture(results) : undefined;
+
+  return mergeFields(awsMetadataFields(result.$metadata), {
+    output: capture?.value,
+    attributes: capture?.truncated ? { "telemetry.dev.capture.truncated": true } : undefined,
+    metadata: omitUndefined({
+      result_count: results?.length,
+      has_next_token: stringValue(result.nextToken) ? true : undefined,
+    }),
+  });
+}
+
 export function ragModel(input: JsonRecord): string | undefined {
   const config = isRecord(input.retrieveAndGenerateConfiguration)
     ? input.retrieveAndGenerateConfiguration

package/src/internal-handlers.ts

@@ -10,6 +10,9 @@ import {
   ragName,
   ragRequestFields,
   ragResponseFields,
+  rerankName,
+  rerankRequestFields,
+  rerankResponseFields,
   retrieveRequestFields,
   retrieveResponseFields,
 } from "./agent-handlers.ts";
@@ -146,6 +149,16 @@ export const AGENT_HANDLERS = {
     requestFields: inlineAgentRequestFields,
     onResult: streamAgentResult("completion"),
   },
+  RerankCommand: {
+    spanName: rerankName,
+    spanType: () => "span",
+    requestFields: rerankRequestFields,
+    onResult(result, span) {
+      endSpan(span, rerankResponseFields(result));
+
+      return result;
+    },
+  },
   RetrieveCommand: {
     spanName: (input) => `retrieve ${stringValue(input.knowledgeBaseId) ?? "unknown"}`,
     spanType: () => "span",

package/src/internal.ts

@@ -145,9 +145,24 @@ export function safe<T>(fn: () => T): T | undefined {
   }
 }
 
+// Bedrock reports input tokens without cache reads and writes; span input tokens include them.
+function withCacheInclusiveInput(fields: SpanFields | undefined): SpanFields | undefined {
+  const usage = fields?.usage;
+
+  if (usage?.inputTokens === undefined) return fields;
+  const inputTokens =
+    usage.inputTokens + (usage.cacheReadInputTokens ?? 0) + (usage.cacheCreationInputTokens ?? 0);
+  const totalTokens =
+    usage.totalTokens !== undefined && usage.outputTokens !== undefined
+      ? inputTokens + usage.outputTokens
+      : usage.totalTokens;
+
+  return { ...fields, usage: { ...usage, inputTokens, totalTokens } };
+}
+
 export function endSpan(span: SpanHandle, fields?: SpanFields): void {
   try {
-    span.end(fields);
+    span.end(withCacheInclusiveInput(fields));
   } catch {
     // Instrumentation must never throw into caller code.
   }