[bot] Bedrock Runtime plugin does not minimize output or capture token metrics for embedding models (Titan Embeddings, Cohere Embed)

[braintrust-bot]

Summary

The @aws-sdk/client-bedrock-runtime InvokeModelCommand surface is used both for text-generation models and for embeddings models (Amazon Titan Text Embeddings V2, Cohere Embed v3/v4 on Bedrock). This repo's Bedrock Runtime plugin treats every invokeModel call identically — it does not detect embeddings responses, so it logs the full raw embedding vector(s) as span output instead of a minimized summary, and it fails to extract token-usage metrics for Titan Embeddings because the plugin's usage-parsing logic only looks for a nested usage/metrics object, which Titan's response does not have.

This is a "materially less detail / less consistent" gap relative to every other embeddings-capable provider already instrumented in this repo: OpenAI, Cohere (direct), Mistral, Groq, and HuggingFace all special-case their embeddings endpoints to log { embedding_length: N } (or equivalent) instead of the raw vector, and to parse token usage from the embeddings response. Bedrock's invokeModel path has no such special-casing at all.

What instrumentation is missing

In js/src/instrumentation/plugins/bedrock-runtime-plugin.ts:

• extractBedrockRuntimeOutput (line 151): for operation === "invokeModel" it does parseJsonBody(response?.body) ?? summarizeBody(response?.body) — fully generic, no branch for embedding responses.
• extractBedrockRuntimeResponseMetrics (line 205) → parseBedrockRuntimeMetrics (line 229): looks for inputTokens / inputTokenCount / input_tokens / prompt_tokens inside a usage (or metadata.usage) object. Titan Embeddings V2's response is { "embedding": [...], "inputTextTokenCount": <int> } — there is no usage object at all, and the field is named inputTextTokenCount, which none of the recognized aliases match. As a result, no token metric is captured for Titan Embeddings calls.
• bedrock-runtime-common.ts's getBedrockRuntimeOperation/buildBedrockRuntimeSpanInfo dispatch purely on AWS SDK command class name (InvokeModelCommand, etc.) and model-agnostic operation strings ("invokeModel"); there is no notion of "this invokeModel call is an embeddings call" anywhere in the plugin.

Concretely, for amazon.titan-embed-text-v2:0 (default 1024-dim, non-normalized-off) or cohere.embed-english-v3 (1024-dim, up to 96 texts per call), the full float array(s) — e.g. {"embedding": [0.0123, -0.0456, ... 1024 floats]} or {"embeddings": [[1024 floats], [1024 floats], ...], "id", "response_type", "texts"} — get logged verbatim as span output, and (for Titan) the inputTextTokenCount field is silently dropped rather than surfaced as a prompt_tokens/tokens metric.

Compare to the OpenAI plugin's embeddings handling (js/src/instrumentation/plugins/openai-plugin.ts, embeddingsCreate channel):

extractOutput: (result) => {
  const embedding = result?.data?.[0]?.embedding;
  return Array.isArray(embedding)
    ? { embedding_length: embedding.length }
    : undefined;
},

No equivalent exists for Bedrock's invokeModel embeddings path.

Braintrust docs status

not_found — https://www.braintrust.dev/docs/integrations/ai-providers/bedrock documents general Bedrock model access (chat models, tool calls, multimodal) but does not mention embeddings models or Titan/Cohere Embed specifically. A broader search of braintrust.dev/docs turned up no page describing Bedrock embeddings tracing/logging behavior. General Bedrock chat/Converse instrumentation is otherwise documented and implemented (see closed issue #1741).

Upstream sources

• Titan Text Embeddings V2 model card (request/response shape, inputTextTokenCount field): https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-amazon-titan-text-embeddings-v2.html
• Titan Embeddings G1 parameters: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-embed-text.html
• Cohere Embed v3 on Bedrock (request/response shape, embeddings/response_type fields): https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-embed-v3.html

Braintrust local repo files inspected

• js/src/instrumentation/plugins/bedrock-runtime-plugin.ts (extractBedrockRuntimeOutput, extractBedrockRuntimeResponseMetrics, parseBedrockRuntimeMetrics)
• js/src/instrumentation/plugins/bedrock-runtime-common.ts (getBedrockRuntimeOperation, buildBedrockRuntimeSpanInfo)
• js/src/instrumentation/plugins/openai-plugin.ts (embeddings comparison baseline, embeddingsCreate channel)
• js/src/instrumentation/plugins/cohere-plugin.ts, mistral-plugin.ts, groq-plugin.ts, huggingface-plugin.ts (confirmed each special-cases embeddings output minimization; grepped for SpanTypeAttribute.EMBEDDING/embedding usage across js/src/instrumentation/plugins/)