diff --git a/.changeset/bedrock-agent-runtime-instrumentation.md b/.changeset/bedrock-agent-runtime-instrumentation.md new file mode 100644 index 000000000..2164afa59 --- /dev/null +++ b/.changeset/bedrock-agent-runtime-instrumentation.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +feat: Add `@aws-sdk/client-bedrock-agent-runtime` instrumentation diff --git a/e2e/config/pr-comment-scenarios.json b/e2e/config/pr-comment-scenarios.json index 65a85ab6e..8a3f8e99e 100644 --- a/e2e/config/pr-comment-scenarios.json +++ b/e2e/config/pr-comment-scenarios.json @@ -273,6 +273,29 @@ } ] }, + { + "scenarioDirName": "bedrock-agent-runtime-instrumentation", + "label": "Bedrock Agent Runtime Instrumentation", + "metadataScenario": "bedrock-agent-runtime-instrumentation", + "variants": [ + { + "variantKey": "bedrock-agent-runtime-v3-wrapped", + "label": "v3 pinned wrapped" + }, + { + "variantKey": "bedrock-agent-runtime-v3-auto", + "label": "v3 pinned auto-hook" + }, + { + "variantKey": "bedrock-agent-runtime-v3-latest-wrapped", + "label": "v3 latest wrapped" + }, + { + "variantKey": "bedrock-agent-runtime-v3-latest-auto", + "label": "v3 latest auto-hook" + } + ] + }, { "scenarioDirName": "bedrock-runtime-instrumentation", "label": "Bedrock Runtime Instrumentation", diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-auto.span-tree.json b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-auto.span-tree.json new file mode 100644 index 000000000..150b6f3a1 --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-auto.span-tree.json @@ -0,0 +1,255 @@ +{ + "span_tree": [ + { + "name": "bedrock-agent-runtime-instrumentation-root", + "type": "task", + "children": [ + { + "name": "bedrock-agent-operation", + "children": [ + { + "name": "bedrockAgent.invokeAgent", + "type": "task", + "children": [ + { + "name": "bedrockAgent.modelInvocation", + "type": "llm", + "children": [], + "input": "", + "output": "", + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "maximumLength": 128, + "model": "us.amazon.nova-lite-v1:0", + "prompt_type": "ORCHESTRATION", + "provider": "aws-bedrock", + "sessionId": "session-1", + "temperature": 0, + "trace_type": "orchestrationTrace" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "tokens": 16 + } + }, + { + "name": "get_weather", + "type": "tool", + "children": [], + "input": { + "parameters": [ + { + "name": "city", + "value": "Paris" + } + ] + }, + "output": { + "text": "{\"temperature\":18}" + }, + "metadata": { + "action_group": "weather", + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "provider": "aws-bedrock", + "sessionId": "session-1", + "tool_type": "action_group", + "trace_type": "orchestrationTrace" + } + } + ], + "input": "Use the weather tool and reply with AGENT_OK.", + "output": { + "text": "AGENT_OK" + }, + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "command": "InvokeAgentCommand", + "contentType": "application/json", + "enableTrace": true, + "operation": "invokeAgent", + "provider": "aws-bedrock", + "sessionId": "session-1" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "time_to_first_token": 0, + "tokens": 16 + } + } + ], + "metadata": { + "operation": "invoke-agent", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-inline-operation", + "children": [ + { + "name": "bedrockAgent.invokeInlineAgent", + "type": "task", + "children": [ + { + "name": "bedrockAgent.modelInvocation", + "type": "llm", + "children": [], + "input": "", + "output": "", + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "maximumLength": 128, + "model": "us.amazon.nova-lite-v1:0", + "prompt_type": "ORCHESTRATION", + "provider": "aws-bedrock", + "sessionId": "session-1", + "temperature": 0, + "trace_type": "orchestrationTrace" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "tokens": 16 + } + } + ], + "input": "Reply with exactly INLINE_OK.", + "output": { + "text": "INLINE_OK" + }, + "metadata": { + "agentName": "braintrust-e2e-inline-agent", + "command": "InvokeInlineAgentCommand", + "contentType": "application/json", + "enableTrace": true, + "endSession": true, + "model": "us.amazon.nova-lite-v1:0", + "operation": "invokeInlineAgent", + "provider": "aws-bedrock", + "sessionId": "session-inline" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "time_to_first_token": 0, + "tokens": 16 + } + } + ], + "metadata": { + "operation": "invoke-inline-agent", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-rag-operation", + "children": [ + { + "name": "bedrockAgent.retrieveAndGenerate", + "type": "task", + "children": [], + "input": { + "text": "Return the deterministic RAG answer." + }, + "output": { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "text": "RAG_OK" + } + } + } + ], + "text": "RAG_OK" + }, + "metadata": { + "command": "RetrieveAndGenerateCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerate", + "provider": "aws-bedrock", + "sessionId": "session-rag" + } + } + ], + "metadata": { + "operation": "retrieve-and-generate", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-rag-stream-operation", + "children": [ + { + "name": "bedrockAgent.retrieveAndGenerateStream", + "type": "task", + "children": [], + "input": { + "text": "Stream the deterministic RAG answer." + }, + "output": { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "span": { + "end": 0, + "start": 0 + }, + "text": "STREAM_RAG_OK" + } + }, + "retrievedReferences": [ + { + "content": { + "text": "Deterministic knowledge-base reference." + }, + "location": { + "type": "WEB", + "webLocation": { + "url": "https://example.com/reference" + } + } + } + ] + } + ], + "text": "STREAM_RAG_OK" + }, + "metadata": { + "command": "RetrieveAndGenerateStreamCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerateStream", + "provider": "aws-bedrock", + "sessionId": "session-rag-stream" + }, + "metrics": { + "time_to_first_token": 0 + } + } + ], + "metadata": { + "operation": "retrieve-and-generate-stream", + "testRunId": "" + } + } + ], + "metadata": { + "scenario": "bedrock-agent-runtime-instrumentation", + "testRunId": "" + } + } + ] +} diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-auto.span-tree.txt b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-auto.span-tree.txt new file mode 100644 index 000000000..b6eca595c --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-auto.span-tree.txt @@ -0,0 +1,201 @@ +span_tree: +└── bedrock-agent-runtime-instrumentation-root [task] + metadata: { + "scenario": "bedrock-agent-runtime-instrumentation", + "testRunId": "" + } + ├── bedrock-agent-operation + │ metadata: { + │ "operation": "invoke-agent", + │ "testRunId": "" + │ } + │ └── bedrockAgent.invokeAgent [task] + │ input: "Use the weather tool and reply with AGENT_OK." + │ output: { + │ "text": "AGENT_OK" + │ } + │ metadata: { + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "command": "InvokeAgentCommand", + │ "contentType": "application/json", + │ "enableTrace": true, + │ "operation": "invokeAgent", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "time_to_first_token": 0, + │ "tokens": 16 + │ } + │ ├── bedrockAgent.modelInvocation [llm] + │ │ input: "" + │ │ output: "" + │ │ metadata: { + │ │ "agentAliasId": "alias-1", + │ │ "agentId": "", + │ │ "agentVersion": "1", + │ │ "maximumLength": 128, + │ │ "model": "us.amazon.nova-lite-v1:0", + │ │ "prompt_type": "ORCHESTRATION", + │ │ "provider": "aws-bedrock", + │ │ "sessionId": "session-1", + │ │ "temperature": 0, + │ │ "trace_type": "orchestrationTrace" + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 4, + │ │ "prompt_tokens": 12, + │ │ "tokens": 16 + │ │ } + │ └── get_weather [tool] + │ input: { + │ "parameters": [ + │ { + │ "name": "city", + │ "value": "Paris" + │ } + │ ] + │ } + │ output: { + │ "text": "{\"temperature\":18}" + │ } + │ metadata: { + │ "action_group": "weather", + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "agentVersion": "1", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1", + │ "tool_type": "action_group", + │ "trace_type": "orchestrationTrace" + │ } + ├── bedrock-agent-inline-operation + │ metadata: { + │ "operation": "invoke-inline-agent", + │ "testRunId": "" + │ } + │ └── bedrockAgent.invokeInlineAgent [task] + │ input: "Reply with exactly INLINE_OK." + │ output: { + │ "text": "INLINE_OK" + │ } + │ metadata: { + │ "agentName": "braintrust-e2e-inline-agent", + │ "command": "InvokeInlineAgentCommand", + │ "contentType": "application/json", + │ "enableTrace": true, + │ "endSession": true, + │ "model": "us.amazon.nova-lite-v1:0", + │ "operation": "invokeInlineAgent", + │ "provider": "aws-bedrock", + │ "sessionId": "session-inline" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "time_to_first_token": 0, + │ "tokens": 16 + │ } + │ └── bedrockAgent.modelInvocation [llm] + │ input: "" + │ output: "" + │ metadata: { + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "agentVersion": "1", + │ "maximumLength": 128, + │ "model": "us.amazon.nova-lite-v1:0", + │ "prompt_type": "ORCHESTRATION", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1", + │ "temperature": 0, + │ "trace_type": "orchestrationTrace" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "tokens": 16 + │ } + ├── bedrock-agent-rag-operation + │ metadata: { + │ "operation": "retrieve-and-generate", + │ "testRunId": "" + │ } + │ └── bedrockAgent.retrieveAndGenerate [task] + │ input: { + │ "text": "Return the deterministic RAG answer." + │ } + │ output: { + │ "citations": [ + │ { + │ "generatedResponsePart": { + │ "textResponsePart": { + │ "text": "RAG_OK" + │ } + │ } + │ } + │ ], + │ "text": "RAG_OK" + │ } + │ metadata: { + │ "command": "RetrieveAndGenerateCommand", + │ "guardrailAction": "NONE", + │ "knowledgeBaseId": "kb-1", + │ "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + │ "operation": "retrieveAndGenerate", + │ "provider": "aws-bedrock", + │ "sessionId": "session-rag" + │ } + └── bedrock-agent-rag-stream-operation + metadata: { + "operation": "retrieve-and-generate-stream", + "testRunId": "" + } + └── bedrockAgent.retrieveAndGenerateStream [task] + input: { + "text": "Stream the deterministic RAG answer." + } + output: { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "span": { + "end": 0, + "start": 0 + }, + "text": "STREAM_RAG_OK" + } + }, + "retrievedReferences": [ + { + "content": { + "text": "Deterministic knowledge-base reference." + }, + "location": { + "type": "WEB", + "webLocation": { + "url": "https://example.com/reference" + } + } + } + ] + } + ], + "text": "STREAM_RAG_OK" + } + metadata: { + "command": "RetrieveAndGenerateStreamCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerateStream", + "provider": "aws-bedrock", + "sessionId": "session-rag-stream" + } + metrics: { + "time_to_first_token": 0 + } diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-auto.span-tree.json b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-auto.span-tree.json new file mode 100644 index 000000000..150b6f3a1 --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-auto.span-tree.json @@ -0,0 +1,255 @@ +{ + "span_tree": [ + { + "name": "bedrock-agent-runtime-instrumentation-root", + "type": "task", + "children": [ + { + "name": "bedrock-agent-operation", + "children": [ + { + "name": "bedrockAgent.invokeAgent", + "type": "task", + "children": [ + { + "name": "bedrockAgent.modelInvocation", + "type": "llm", + "children": [], + "input": "", + "output": "", + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "maximumLength": 128, + "model": "us.amazon.nova-lite-v1:0", + "prompt_type": "ORCHESTRATION", + "provider": "aws-bedrock", + "sessionId": "session-1", + "temperature": 0, + "trace_type": "orchestrationTrace" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "tokens": 16 + } + }, + { + "name": "get_weather", + "type": "tool", + "children": [], + "input": { + "parameters": [ + { + "name": "city", + "value": "Paris" + } + ] + }, + "output": { + "text": "{\"temperature\":18}" + }, + "metadata": { + "action_group": "weather", + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "provider": "aws-bedrock", + "sessionId": "session-1", + "tool_type": "action_group", + "trace_type": "orchestrationTrace" + } + } + ], + "input": "Use the weather tool and reply with AGENT_OK.", + "output": { + "text": "AGENT_OK" + }, + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "command": "InvokeAgentCommand", + "contentType": "application/json", + "enableTrace": true, + "operation": "invokeAgent", + "provider": "aws-bedrock", + "sessionId": "session-1" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "time_to_first_token": 0, + "tokens": 16 + } + } + ], + "metadata": { + "operation": "invoke-agent", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-inline-operation", + "children": [ + { + "name": "bedrockAgent.invokeInlineAgent", + "type": "task", + "children": [ + { + "name": "bedrockAgent.modelInvocation", + "type": "llm", + "children": [], + "input": "", + "output": "", + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "maximumLength": 128, + "model": "us.amazon.nova-lite-v1:0", + "prompt_type": "ORCHESTRATION", + "provider": "aws-bedrock", + "sessionId": "session-1", + "temperature": 0, + "trace_type": "orchestrationTrace" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "tokens": 16 + } + } + ], + "input": "Reply with exactly INLINE_OK.", + "output": { + "text": "INLINE_OK" + }, + "metadata": { + "agentName": "braintrust-e2e-inline-agent", + "command": "InvokeInlineAgentCommand", + "contentType": "application/json", + "enableTrace": true, + "endSession": true, + "model": "us.amazon.nova-lite-v1:0", + "operation": "invokeInlineAgent", + "provider": "aws-bedrock", + "sessionId": "session-inline" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "time_to_first_token": 0, + "tokens": 16 + } + } + ], + "metadata": { + "operation": "invoke-inline-agent", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-rag-operation", + "children": [ + { + "name": "bedrockAgent.retrieveAndGenerate", + "type": "task", + "children": [], + "input": { + "text": "Return the deterministic RAG answer." + }, + "output": { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "text": "RAG_OK" + } + } + } + ], + "text": "RAG_OK" + }, + "metadata": { + "command": "RetrieveAndGenerateCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerate", + "provider": "aws-bedrock", + "sessionId": "session-rag" + } + } + ], + "metadata": { + "operation": "retrieve-and-generate", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-rag-stream-operation", + "children": [ + { + "name": "bedrockAgent.retrieveAndGenerateStream", + "type": "task", + "children": [], + "input": { + "text": "Stream the deterministic RAG answer." + }, + "output": { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "span": { + "end": 0, + "start": 0 + }, + "text": "STREAM_RAG_OK" + } + }, + "retrievedReferences": [ + { + "content": { + "text": "Deterministic knowledge-base reference." + }, + "location": { + "type": "WEB", + "webLocation": { + "url": "https://example.com/reference" + } + } + } + ] + } + ], + "text": "STREAM_RAG_OK" + }, + "metadata": { + "command": "RetrieveAndGenerateStreamCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerateStream", + "provider": "aws-bedrock", + "sessionId": "session-rag-stream" + }, + "metrics": { + "time_to_first_token": 0 + } + } + ], + "metadata": { + "operation": "retrieve-and-generate-stream", + "testRunId": "" + } + } + ], + "metadata": { + "scenario": "bedrock-agent-runtime-instrumentation", + "testRunId": "" + } + } + ] +} diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-auto.span-tree.txt b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-auto.span-tree.txt new file mode 100644 index 000000000..b6eca595c --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-auto.span-tree.txt @@ -0,0 +1,201 @@ +span_tree: +└── bedrock-agent-runtime-instrumentation-root [task] + metadata: { + "scenario": "bedrock-agent-runtime-instrumentation", + "testRunId": "" + } + ├── bedrock-agent-operation + │ metadata: { + │ "operation": "invoke-agent", + │ "testRunId": "" + │ } + │ └── bedrockAgent.invokeAgent [task] + │ input: "Use the weather tool and reply with AGENT_OK." + │ output: { + │ "text": "AGENT_OK" + │ } + │ metadata: { + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "command": "InvokeAgentCommand", + │ "contentType": "application/json", + │ "enableTrace": true, + │ "operation": "invokeAgent", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "time_to_first_token": 0, + │ "tokens": 16 + │ } + │ ├── bedrockAgent.modelInvocation [llm] + │ │ input: "" + │ │ output: "" + │ │ metadata: { + │ │ "agentAliasId": "alias-1", + │ │ "agentId": "", + │ │ "agentVersion": "1", + │ │ "maximumLength": 128, + │ │ "model": "us.amazon.nova-lite-v1:0", + │ │ "prompt_type": "ORCHESTRATION", + │ │ "provider": "aws-bedrock", + │ │ "sessionId": "session-1", + │ │ "temperature": 0, + │ │ "trace_type": "orchestrationTrace" + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 4, + │ │ "prompt_tokens": 12, + │ │ "tokens": 16 + │ │ } + │ └── get_weather [tool] + │ input: { + │ "parameters": [ + │ { + │ "name": "city", + │ "value": "Paris" + │ } + │ ] + │ } + │ output: { + │ "text": "{\"temperature\":18}" + │ } + │ metadata: { + │ "action_group": "weather", + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "agentVersion": "1", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1", + │ "tool_type": "action_group", + │ "trace_type": "orchestrationTrace" + │ } + ├── bedrock-agent-inline-operation + │ metadata: { + │ "operation": "invoke-inline-agent", + │ "testRunId": "" + │ } + │ └── bedrockAgent.invokeInlineAgent [task] + │ input: "Reply with exactly INLINE_OK." + │ output: { + │ "text": "INLINE_OK" + │ } + │ metadata: { + │ "agentName": "braintrust-e2e-inline-agent", + │ "command": "InvokeInlineAgentCommand", + │ "contentType": "application/json", + │ "enableTrace": true, + │ "endSession": true, + │ "model": "us.amazon.nova-lite-v1:0", + │ "operation": "invokeInlineAgent", + │ "provider": "aws-bedrock", + │ "sessionId": "session-inline" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "time_to_first_token": 0, + │ "tokens": 16 + │ } + │ └── bedrockAgent.modelInvocation [llm] + │ input: "" + │ output: "" + │ metadata: { + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "agentVersion": "1", + │ "maximumLength": 128, + │ "model": "us.amazon.nova-lite-v1:0", + │ "prompt_type": "ORCHESTRATION", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1", + │ "temperature": 0, + │ "trace_type": "orchestrationTrace" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "tokens": 16 + │ } + ├── bedrock-agent-rag-operation + │ metadata: { + │ "operation": "retrieve-and-generate", + │ "testRunId": "" + │ } + │ └── bedrockAgent.retrieveAndGenerate [task] + │ input: { + │ "text": "Return the deterministic RAG answer." + │ } + │ output: { + │ "citations": [ + │ { + │ "generatedResponsePart": { + │ "textResponsePart": { + │ "text": "RAG_OK" + │ } + │ } + │ } + │ ], + │ "text": "RAG_OK" + │ } + │ metadata: { + │ "command": "RetrieveAndGenerateCommand", + │ "guardrailAction": "NONE", + │ "knowledgeBaseId": "kb-1", + │ "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + │ "operation": "retrieveAndGenerate", + │ "provider": "aws-bedrock", + │ "sessionId": "session-rag" + │ } + └── bedrock-agent-rag-stream-operation + metadata: { + "operation": "retrieve-and-generate-stream", + "testRunId": "" + } + └── bedrockAgent.retrieveAndGenerateStream [task] + input: { + "text": "Stream the deterministic RAG answer." + } + output: { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "span": { + "end": 0, + "start": 0 + }, + "text": "STREAM_RAG_OK" + } + }, + "retrievedReferences": [ + { + "content": { + "text": "Deterministic knowledge-base reference." + }, + "location": { + "type": "WEB", + "webLocation": { + "url": "https://example.com/reference" + } + } + } + ] + } + ], + "text": "STREAM_RAG_OK" + } + metadata: { + "command": "RetrieveAndGenerateStreamCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerateStream", + "provider": "aws-bedrock", + "sessionId": "session-rag-stream" + } + metrics: { + "time_to_first_token": 0 + } diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-wrapped.span-tree.json b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-wrapped.span-tree.json new file mode 100644 index 000000000..150b6f3a1 --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-wrapped.span-tree.json @@ -0,0 +1,255 @@ +{ + "span_tree": [ + { + "name": "bedrock-agent-runtime-instrumentation-root", + "type": "task", + "children": [ + { + "name": "bedrock-agent-operation", + "children": [ + { + "name": "bedrockAgent.invokeAgent", + "type": "task", + "children": [ + { + "name": "bedrockAgent.modelInvocation", + "type": "llm", + "children": [], + "input": "", + "output": "", + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "maximumLength": 128, + "model": "us.amazon.nova-lite-v1:0", + "prompt_type": "ORCHESTRATION", + "provider": "aws-bedrock", + "sessionId": "session-1", + "temperature": 0, + "trace_type": "orchestrationTrace" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "tokens": 16 + } + }, + { + "name": "get_weather", + "type": "tool", + "children": [], + "input": { + "parameters": [ + { + "name": "city", + "value": "Paris" + } + ] + }, + "output": { + "text": "{\"temperature\":18}" + }, + "metadata": { + "action_group": "weather", + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "provider": "aws-bedrock", + "sessionId": "session-1", + "tool_type": "action_group", + "trace_type": "orchestrationTrace" + } + } + ], + "input": "Use the weather tool and reply with AGENT_OK.", + "output": { + "text": "AGENT_OK" + }, + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "command": "InvokeAgentCommand", + "contentType": "application/json", + "enableTrace": true, + "operation": "invokeAgent", + "provider": "aws-bedrock", + "sessionId": "session-1" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "time_to_first_token": 0, + "tokens": 16 + } + } + ], + "metadata": { + "operation": "invoke-agent", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-inline-operation", + "children": [ + { + "name": "bedrockAgent.invokeInlineAgent", + "type": "task", + "children": [ + { + "name": "bedrockAgent.modelInvocation", + "type": "llm", + "children": [], + "input": "", + "output": "", + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "maximumLength": 128, + "model": "us.amazon.nova-lite-v1:0", + "prompt_type": "ORCHESTRATION", + "provider": "aws-bedrock", + "sessionId": "session-1", + "temperature": 0, + "trace_type": "orchestrationTrace" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "tokens": 16 + } + } + ], + "input": "Reply with exactly INLINE_OK.", + "output": { + "text": "INLINE_OK" + }, + "metadata": { + "agentName": "braintrust-e2e-inline-agent", + "command": "InvokeInlineAgentCommand", + "contentType": "application/json", + "enableTrace": true, + "endSession": true, + "model": "us.amazon.nova-lite-v1:0", + "operation": "invokeInlineAgent", + "provider": "aws-bedrock", + "sessionId": "session-inline" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "time_to_first_token": 0, + "tokens": 16 + } + } + ], + "metadata": { + "operation": "invoke-inline-agent", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-rag-operation", + "children": [ + { + "name": "bedrockAgent.retrieveAndGenerate", + "type": "task", + "children": [], + "input": { + "text": "Return the deterministic RAG answer." + }, + "output": { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "text": "RAG_OK" + } + } + } + ], + "text": "RAG_OK" + }, + "metadata": { + "command": "RetrieveAndGenerateCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerate", + "provider": "aws-bedrock", + "sessionId": "session-rag" + } + } + ], + "metadata": { + "operation": "retrieve-and-generate", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-rag-stream-operation", + "children": [ + { + "name": "bedrockAgent.retrieveAndGenerateStream", + "type": "task", + "children": [], + "input": { + "text": "Stream the deterministic RAG answer." + }, + "output": { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "span": { + "end": 0, + "start": 0 + }, + "text": "STREAM_RAG_OK" + } + }, + "retrievedReferences": [ + { + "content": { + "text": "Deterministic knowledge-base reference." + }, + "location": { + "type": "WEB", + "webLocation": { + "url": "https://example.com/reference" + } + } + } + ] + } + ], + "text": "STREAM_RAG_OK" + }, + "metadata": { + "command": "RetrieveAndGenerateStreamCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerateStream", + "provider": "aws-bedrock", + "sessionId": "session-rag-stream" + }, + "metrics": { + "time_to_first_token": 0 + } + } + ], + "metadata": { + "operation": "retrieve-and-generate-stream", + "testRunId": "" + } + } + ], + "metadata": { + "scenario": "bedrock-agent-runtime-instrumentation", + "testRunId": "" + } + } + ] +} diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-wrapped.span-tree.txt b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-wrapped.span-tree.txt new file mode 100644 index 000000000..b6eca595c --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-latest-wrapped.span-tree.txt @@ -0,0 +1,201 @@ +span_tree: +└── bedrock-agent-runtime-instrumentation-root [task] + metadata: { + "scenario": "bedrock-agent-runtime-instrumentation", + "testRunId": "" + } + ├── bedrock-agent-operation + │ metadata: { + │ "operation": "invoke-agent", + │ "testRunId": "" + │ } + │ └── bedrockAgent.invokeAgent [task] + │ input: "Use the weather tool and reply with AGENT_OK." + │ output: { + │ "text": "AGENT_OK" + │ } + │ metadata: { + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "command": "InvokeAgentCommand", + │ "contentType": "application/json", + │ "enableTrace": true, + │ "operation": "invokeAgent", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "time_to_first_token": 0, + │ "tokens": 16 + │ } + │ ├── bedrockAgent.modelInvocation [llm] + │ │ input: "" + │ │ output: "" + │ │ metadata: { + │ │ "agentAliasId": "alias-1", + │ │ "agentId": "", + │ │ "agentVersion": "1", + │ │ "maximumLength": 128, + │ │ "model": "us.amazon.nova-lite-v1:0", + │ │ "prompt_type": "ORCHESTRATION", + │ │ "provider": "aws-bedrock", + │ │ "sessionId": "session-1", + │ │ "temperature": 0, + │ │ "trace_type": "orchestrationTrace" + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 4, + │ │ "prompt_tokens": 12, + │ │ "tokens": 16 + │ │ } + │ └── get_weather [tool] + │ input: { + │ "parameters": [ + │ { + │ "name": "city", + │ "value": "Paris" + │ } + │ ] + │ } + │ output: { + │ "text": "{\"temperature\":18}" + │ } + │ metadata: { + │ "action_group": "weather", + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "agentVersion": "1", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1", + │ "tool_type": "action_group", + │ "trace_type": "orchestrationTrace" + │ } + ├── bedrock-agent-inline-operation + │ metadata: { + │ "operation": "invoke-inline-agent", + │ "testRunId": "" + │ } + │ └── bedrockAgent.invokeInlineAgent [task] + │ input: "Reply with exactly INLINE_OK." + │ output: { + │ "text": "INLINE_OK" + │ } + │ metadata: { + │ "agentName": "braintrust-e2e-inline-agent", + │ "command": "InvokeInlineAgentCommand", + │ "contentType": "application/json", + │ "enableTrace": true, + │ "endSession": true, + │ "model": "us.amazon.nova-lite-v1:0", + │ "operation": "invokeInlineAgent", + │ "provider": "aws-bedrock", + │ "sessionId": "session-inline" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "time_to_first_token": 0, + │ "tokens": 16 + │ } + │ └── bedrockAgent.modelInvocation [llm] + │ input: "" + │ output: "" + │ metadata: { + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "agentVersion": "1", + │ "maximumLength": 128, + │ "model": "us.amazon.nova-lite-v1:0", + │ "prompt_type": "ORCHESTRATION", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1", + │ "temperature": 0, + │ "trace_type": "orchestrationTrace" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "tokens": 16 + │ } + ├── bedrock-agent-rag-operation + │ metadata: { + │ "operation": "retrieve-and-generate", + │ "testRunId": "" + │ } + │ └── bedrockAgent.retrieveAndGenerate [task] + │ input: { + │ "text": "Return the deterministic RAG answer." + │ } + │ output: { + │ "citations": [ + │ { + │ "generatedResponsePart": { + │ "textResponsePart": { + │ "text": "RAG_OK" + │ } + │ } + │ } + │ ], + │ "text": "RAG_OK" + │ } + │ metadata: { + │ "command": "RetrieveAndGenerateCommand", + │ "guardrailAction": "NONE", + │ "knowledgeBaseId": "kb-1", + │ "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + │ "operation": "retrieveAndGenerate", + │ "provider": "aws-bedrock", + │ "sessionId": "session-rag" + │ } + └── bedrock-agent-rag-stream-operation + metadata: { + "operation": "retrieve-and-generate-stream", + "testRunId": "" + } + └── bedrockAgent.retrieveAndGenerateStream [task] + input: { + "text": "Stream the deterministic RAG answer." + } + output: { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "span": { + "end": 0, + "start": 0 + }, + "text": "STREAM_RAG_OK" + } + }, + "retrievedReferences": [ + { + "content": { + "text": "Deterministic knowledge-base reference." + }, + "location": { + "type": "WEB", + "webLocation": { + "url": "https://example.com/reference" + } + } + } + ] + } + ], + "text": "STREAM_RAG_OK" + } + metadata: { + "command": "RetrieveAndGenerateStreamCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerateStream", + "provider": "aws-bedrock", + "sessionId": "session-rag-stream" + } + metrics: { + "time_to_first_token": 0 + } diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-wrapped.span-tree.json b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-wrapped.span-tree.json new file mode 100644 index 000000000..150b6f3a1 --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-wrapped.span-tree.json @@ -0,0 +1,255 @@ +{ + "span_tree": [ + { + "name": "bedrock-agent-runtime-instrumentation-root", + "type": "task", + "children": [ + { + "name": "bedrock-agent-operation", + "children": [ + { + "name": "bedrockAgent.invokeAgent", + "type": "task", + "children": [ + { + "name": "bedrockAgent.modelInvocation", + "type": "llm", + "children": [], + "input": "", + "output": "", + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "maximumLength": 128, + "model": "us.amazon.nova-lite-v1:0", + "prompt_type": "ORCHESTRATION", + "provider": "aws-bedrock", + "sessionId": "session-1", + "temperature": 0, + "trace_type": "orchestrationTrace" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "tokens": 16 + } + }, + { + "name": "get_weather", + "type": "tool", + "children": [], + "input": { + "parameters": [ + { + "name": "city", + "value": "Paris" + } + ] + }, + "output": { + "text": "{\"temperature\":18}" + }, + "metadata": { + "action_group": "weather", + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "provider": "aws-bedrock", + "sessionId": "session-1", + "tool_type": "action_group", + "trace_type": "orchestrationTrace" + } + } + ], + "input": "Use the weather tool and reply with AGENT_OK.", + "output": { + "text": "AGENT_OK" + }, + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "command": "InvokeAgentCommand", + "contentType": "application/json", + "enableTrace": true, + "operation": "invokeAgent", + "provider": "aws-bedrock", + "sessionId": "session-1" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "time_to_first_token": 0, + "tokens": 16 + } + } + ], + "metadata": { + "operation": "invoke-agent", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-inline-operation", + "children": [ + { + "name": "bedrockAgent.invokeInlineAgent", + "type": "task", + "children": [ + { + "name": "bedrockAgent.modelInvocation", + "type": "llm", + "children": [], + "input": "", + "output": "", + "metadata": { + "agentAliasId": "alias-1", + "agentId": "", + "agentVersion": "1", + "maximumLength": 128, + "model": "us.amazon.nova-lite-v1:0", + "prompt_type": "ORCHESTRATION", + "provider": "aws-bedrock", + "sessionId": "session-1", + "temperature": 0, + "trace_type": "orchestrationTrace" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "tokens": 16 + } + } + ], + "input": "Reply with exactly INLINE_OK.", + "output": { + "text": "INLINE_OK" + }, + "metadata": { + "agentName": "braintrust-e2e-inline-agent", + "command": "InvokeInlineAgentCommand", + "contentType": "application/json", + "enableTrace": true, + "endSession": true, + "model": "us.amazon.nova-lite-v1:0", + "operation": "invokeInlineAgent", + "provider": "aws-bedrock", + "sessionId": "session-inline" + }, + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 12, + "time_to_first_token": 0, + "tokens": 16 + } + } + ], + "metadata": { + "operation": "invoke-inline-agent", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-rag-operation", + "children": [ + { + "name": "bedrockAgent.retrieveAndGenerate", + "type": "task", + "children": [], + "input": { + "text": "Return the deterministic RAG answer." + }, + "output": { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "text": "RAG_OK" + } + } + } + ], + "text": "RAG_OK" + }, + "metadata": { + "command": "RetrieveAndGenerateCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerate", + "provider": "aws-bedrock", + "sessionId": "session-rag" + } + } + ], + "metadata": { + "operation": "retrieve-and-generate", + "testRunId": "" + } + }, + { + "name": "bedrock-agent-rag-stream-operation", + "children": [ + { + "name": "bedrockAgent.retrieveAndGenerateStream", + "type": "task", + "children": [], + "input": { + "text": "Stream the deterministic RAG answer." + }, + "output": { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "span": { + "end": 0, + "start": 0 + }, + "text": "STREAM_RAG_OK" + } + }, + "retrievedReferences": [ + { + "content": { + "text": "Deterministic knowledge-base reference." + }, + "location": { + "type": "WEB", + "webLocation": { + "url": "https://example.com/reference" + } + } + } + ] + } + ], + "text": "STREAM_RAG_OK" + }, + "metadata": { + "command": "RetrieveAndGenerateStreamCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerateStream", + "provider": "aws-bedrock", + "sessionId": "session-rag-stream" + }, + "metrics": { + "time_to_first_token": 0 + } + } + ], + "metadata": { + "operation": "retrieve-and-generate-stream", + "testRunId": "" + } + } + ], + "metadata": { + "scenario": "bedrock-agent-runtime-instrumentation", + "testRunId": "" + } + } + ] +} diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-wrapped.span-tree.txt b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-wrapped.span-tree.txt new file mode 100644 index 000000000..b6eca595c --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/__snapshots__/bedrock-agent-runtime-v3-wrapped.span-tree.txt @@ -0,0 +1,201 @@ +span_tree: +└── bedrock-agent-runtime-instrumentation-root [task] + metadata: { + "scenario": "bedrock-agent-runtime-instrumentation", + "testRunId": "" + } + ├── bedrock-agent-operation + │ metadata: { + │ "operation": "invoke-agent", + │ "testRunId": "" + │ } + │ └── bedrockAgent.invokeAgent [task] + │ input: "Use the weather tool and reply with AGENT_OK." + │ output: { + │ "text": "AGENT_OK" + │ } + │ metadata: { + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "command": "InvokeAgentCommand", + │ "contentType": "application/json", + │ "enableTrace": true, + │ "operation": "invokeAgent", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "time_to_first_token": 0, + │ "tokens": 16 + │ } + │ ├── bedrockAgent.modelInvocation [llm] + │ │ input: "" + │ │ output: "" + │ │ metadata: { + │ │ "agentAliasId": "alias-1", + │ │ "agentId": "", + │ │ "agentVersion": "1", + │ │ "maximumLength": 128, + │ │ "model": "us.amazon.nova-lite-v1:0", + │ │ "prompt_type": "ORCHESTRATION", + │ │ "provider": "aws-bedrock", + │ │ "sessionId": "session-1", + │ │ "temperature": 0, + │ │ "trace_type": "orchestrationTrace" + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 4, + │ │ "prompt_tokens": 12, + │ │ "tokens": 16 + │ │ } + │ └── get_weather [tool] + │ input: { + │ "parameters": [ + │ { + │ "name": "city", + │ "value": "Paris" + │ } + │ ] + │ } + │ output: { + │ "text": "{\"temperature\":18}" + │ } + │ metadata: { + │ "action_group": "weather", + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "agentVersion": "1", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1", + │ "tool_type": "action_group", + │ "trace_type": "orchestrationTrace" + │ } + ├── bedrock-agent-inline-operation + │ metadata: { + │ "operation": "invoke-inline-agent", + │ "testRunId": "" + │ } + │ └── bedrockAgent.invokeInlineAgent [task] + │ input: "Reply with exactly INLINE_OK." + │ output: { + │ "text": "INLINE_OK" + │ } + │ metadata: { + │ "agentName": "braintrust-e2e-inline-agent", + │ "command": "InvokeInlineAgentCommand", + │ "contentType": "application/json", + │ "enableTrace": true, + │ "endSession": true, + │ "model": "us.amazon.nova-lite-v1:0", + │ "operation": "invokeInlineAgent", + │ "provider": "aws-bedrock", + │ "sessionId": "session-inline" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "time_to_first_token": 0, + │ "tokens": 16 + │ } + │ └── bedrockAgent.modelInvocation [llm] + │ input: "" + │ output: "" + │ metadata: { + │ "agentAliasId": "alias-1", + │ "agentId": "", + │ "agentVersion": "1", + │ "maximumLength": 128, + │ "model": "us.amazon.nova-lite-v1:0", + │ "prompt_type": "ORCHESTRATION", + │ "provider": "aws-bedrock", + │ "sessionId": "session-1", + │ "temperature": 0, + │ "trace_type": "orchestrationTrace" + │ } + │ metrics: { + │ "completion_tokens": 4, + │ "prompt_tokens": 12, + │ "tokens": 16 + │ } + ├── bedrock-agent-rag-operation + │ metadata: { + │ "operation": "retrieve-and-generate", + │ "testRunId": "" + │ } + │ └── bedrockAgent.retrieveAndGenerate [task] + │ input: { + │ "text": "Return the deterministic RAG answer." + │ } + │ output: { + │ "citations": [ + │ { + │ "generatedResponsePart": { + │ "textResponsePart": { + │ "text": "RAG_OK" + │ } + │ } + │ } + │ ], + │ "text": "RAG_OK" + │ } + │ metadata: { + │ "command": "RetrieveAndGenerateCommand", + │ "guardrailAction": "NONE", + │ "knowledgeBaseId": "kb-1", + │ "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + │ "operation": "retrieveAndGenerate", + │ "provider": "aws-bedrock", + │ "sessionId": "session-rag" + │ } + └── bedrock-agent-rag-stream-operation + metadata: { + "operation": "retrieve-and-generate-stream", + "testRunId": "" + } + └── bedrockAgent.retrieveAndGenerateStream [task] + input: { + "text": "Stream the deterministic RAG answer." + } + output: { + "citations": [ + { + "generatedResponsePart": { + "textResponsePart": { + "span": { + "end": 0, + "start": 0 + }, + "text": "STREAM_RAG_OK" + } + }, + "retrievedReferences": [ + { + "content": { + "text": "Deterministic knowledge-base reference." + }, + "location": { + "type": "WEB", + "webLocation": { + "url": "https://example.com/reference" + } + } + } + ] + } + ], + "text": "STREAM_RAG_OK" + } + metadata: { + "command": "RetrieveAndGenerateStreamCommand", + "guardrailAction": "NONE", + "knowledgeBaseId": "kb-1", + "model": "arn:aws:bedrock:us-east-1:000000000000:foundation-model/us.amazon.nova-lite-v1:0", + "operation": "retrieveAndGenerateStream", + "provider": "aws-bedrock", + "sessionId": "session-rag-stream" + } + metrics: { + "time_to_first_token": 0 + } diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/assertions.ts b/e2e/scenarios/bedrock-agent-runtime-instrumentation/assertions.ts new file mode 100644 index 000000000..9eb35997a --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/assertions.ts @@ -0,0 +1,298 @@ +import { beforeAll, describe, expect, test } from "vitest"; +import { resolveFileSnapshotPath } from "../../helpers/file-snapshot"; +import type { CapturedLogEvent } from "../../helpers/mock-braintrust-server"; +import { + withScenarioHarness, + type ScenarioRunContext, +} from "../../helpers/scenario-harness"; +import { + matchSpanTreeSnapshot, + spanTreeFields, + type SpanTreeEntry, +} from "../../helpers/span-tree"; +import { + findChildSpans, + findLatestSpan, + spanInstrumentationName, +} from "../../helpers/trace-selectors"; +import { + AGENT_OPERATION_NAME, + INLINE_OPERATION_NAME, + MODEL, + RAG_OPERATION_NAME, + RAG_STREAM_OPERATION_NAME, + ROOT_NAME, + SCENARIO_NAME, +} from "./constants.mjs"; + +type RunScenario = (harness: { + runNodeScenarioDir: (options: { + entry: string; + env: Record; + nodeArgs: string[]; + runContext?: ScenarioRunContext; + scenarioDir: string; + timeoutMs: number; + }) => Promise; + runScenarioDir: (options: { + entry: string; + env: Record; + runContext?: ScenarioRunContext; + scenarioDir: string; + timeoutMs: number; + }) => Promise; +}) => Promise; + +const OPERATIONS = [ + { + command: "InvokeAgentCommand", + operationName: AGENT_OPERATION_NAME, + spanName: "bedrockAgent.invokeAgent", + }, + { + command: "InvokeInlineAgentCommand", + operationName: INLINE_OPERATION_NAME, + spanName: "bedrockAgent.invokeInlineAgent", + }, + { + command: "RetrieveAndGenerateCommand", + operationName: RAG_OPERATION_NAME, + spanName: "bedrockAgent.retrieveAndGenerate", + }, + { + command: "RetrieveAndGenerateStreamCommand", + operationName: RAG_STREAM_OPERATION_NAME, + spanName: "bedrockAgent.retrieveAndGenerateStream", + }, +] as const; + +function findInstrumentationSpan( + events: CapturedLogEvent[], + operationName: string, + spanName: string, +) { + const operation = findLatestSpan(events, operationName); + return findChildSpans(events, spanName, operation?.span.id).find( + (span) => span.output !== undefined, + ); +} + +function findAgentSpan(events: CapturedLogEvent[]) { + return findInstrumentationSpan( + events, + AGENT_OPERATION_NAME, + "bedrockAgent.invokeAgent", + ); +} + +function spanTreeEvents( + events: CapturedLogEvent[], +): Array { + const root = findLatestSpan(events, ROOT_NAME); + const items: Array = root ? [root] : []; + + for (const definition of OPERATIONS) { + const operation = findLatestSpan(events, definition.operationName); + const instrumentationSpan = findInstrumentationSpan( + events, + definition.operationName, + definition.spanName, + ); + if (operation) { + items.push(operation); + } + if (instrumentationSpan) { + items.push(instrumentationSpan); + for (const childName of ["bedrockAgent.modelInvocation", "get_weather"]) { + for (const child of findChildSpans( + events, + childName, + instrumentationSpan.span.id, + )) { + items.push({ + event: child, + fields: { + ...spanTreeFields(child), + input: + child.span.type === "llm" + ? "" + : child.input, + output: + child.span.type === "llm" + ? "" + : child.output, + }, + }); + } + } + } + } + return items; +} + +export function defineBedrockAgentRuntimeInstrumentationAssertions(options: { + name: string; + runScenario: RunScenario; + snapshotName: string; + testFileUrl: string; + timeoutMs: number; +}): void { + const spanSnapshotPath = resolveFileSnapshotPath( + options.testFileUrl, + `${options.snapshotName}.span-tree.json`, + ); + const testConfig = { + timeout: options.timeoutMs, + }; + + describe(options.name, () => { + let events: CapturedLogEvent[] = []; + + beforeAll(async () => { + await withScenarioHarness(async (harness) => { + await options.runScenario(harness); + events = harness.events(); + }); + }, options.timeoutMs); + + test("captures the scenario root span", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + expect(root).toBeDefined(); + expect(root?.row.metadata).toMatchObject({ + scenario: SCENARIO_NAME, + }); + }); + + test("captures all supported Agent Runtime commands", testConfig, () => { + for (const definition of OPERATIONS) { + const operation = findLatestSpan(events, definition.operationName); + const spans = findChildSpans( + events, + definition.spanName, + operation?.span.id, + ); + expect(spans, definition.command).toHaveLength(1); + const span = spans[0]; + expect(span, definition.command).toBeDefined(); + expect(span?.span.type).toBe("task"); + expect(spanInstrumentationName(span)).toBe("bedrock-agent-runtime"); + expect(span?.row.metadata).toMatchObject({ + command: definition.command, + provider: "aws-bedrock", + }); + } + }); + + test( + "captures managed and inline agent streaming output", + testConfig, + () => { + const agentSpan = findAgentSpan(events); + const inlineSpan = findInstrumentationSpan( + events, + INLINE_OPERATION_NAME, + "bedrockAgent.invokeInlineAgent", + ); + expect(agentSpan?.input).toBe( + "Use the weather tool and reply with AGENT_OK.", + ); + expect(agentSpan?.output).toEqual({ text: "AGENT_OK" }); + expect(inlineSpan?.input).toBe("Reply with exactly INLINE_OK."); + expect(inlineSpan?.row.metadata).toMatchObject({ + model: MODEL, + }); + expect(inlineSpan?.output).toEqual({ text: "INLINE_OK" }); + for (const span of [agentSpan, inlineSpan]) { + expect(span?.metrics).toMatchObject({ + completion_tokens: 4, + prompt_tokens: 12, + time_to_first_token: expect.any(Number), + tokens: 16, + }); + } + }, + ); + + test("captures model and action-group child spans", testConfig, () => { + const agentSpan = findAgentSpan(events); + const modelSpan = findChildSpans( + events, + "bedrockAgent.modelInvocation", + agentSpan?.span.id, + ).find((span) => span.output !== undefined); + const toolSpan = findChildSpans( + events, + "get_weather", + agentSpan?.span.id, + ).find((span) => span.output !== undefined); + + expect(modelSpan).toMatchObject({ + metrics: { + completion_tokens: 4, + prompt_tokens: 12, + tokens: 16, + }, + span: { type: "llm" }, + }); + expect(modelSpan?.row.metadata).toMatchObject({ + model: MODEL, + provider: "aws-bedrock", + }); + expect(spanInstrumentationName(modelSpan)).toBe("bedrock-agent-runtime"); + expect(toolSpan).toMatchObject({ + input: { + parameters: [{ name: "city", value: "Paris" }], + }, + output: { + text: JSON.stringify({ temperature: 18 }), + }, + span: { type: "tool" }, + }); + expect(toolSpan?.row.metadata).toMatchObject({ + action_group: "weather", + provider: "aws-bedrock", + tool_type: "action_group", + }); + expect(spanInstrumentationName(toolSpan)).toBe("bedrock-agent-runtime"); + }); + + test("captures non-streaming and streaming RAG results", testConfig, () => { + const ragSpan = findInstrumentationSpan( + events, + RAG_OPERATION_NAME, + "bedrockAgent.retrieveAndGenerate", + ); + const ragStreamSpan = findInstrumentationSpan( + events, + RAG_STREAM_OPERATION_NAME, + "bedrockAgent.retrieveAndGenerateStream", + ); + + expect(ragSpan?.input).toEqual({ + text: "Return the deterministic RAG answer.", + }); + expect(ragSpan?.output).toMatchObject({ + citations: expect.any(Array), + text: "RAG_OK", + }); + expect(ragSpan?.row.metadata).toMatchObject({ + knowledgeBaseId: "kb-1", + model: expect.stringContaining(MODEL), + }); + expect(ragStreamSpan?.output).toMatchObject({ + citations: expect.any(Array), + text: "STREAM_RAG_OK", + }); + expect(ragStreamSpan?.row.metadata).toMatchObject({ + guardrailAction: "NONE", + }); + expect(ragStreamSpan?.metrics).toMatchObject({ + time_to_first_token: expect.any(Number), + }); + }); + + test("matches span tree snapshot", testConfig, async () => { + await matchSpanTreeSnapshot(spanTreeEvents(events), spanSnapshotPath); + }); + }); +} diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/constants.mjs b/e2e/scenarios/bedrock-agent-runtime-instrumentation/constants.mjs new file mode 100644 index 000000000..8c346e2eb --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/constants.mjs @@ -0,0 +1,8 @@ +export const MODEL = "us.amazon.nova-lite-v1:0"; +export const REGION = "us-east-1"; +export const ROOT_NAME = "bedrock-agent-runtime-instrumentation-root"; +export const SCENARIO_NAME = "bedrock-agent-runtime-instrumentation"; +export const AGENT_OPERATION_NAME = "bedrock-agent-operation"; +export const INLINE_OPERATION_NAME = "bedrock-agent-inline-operation"; +export const RAG_OPERATION_NAME = "bedrock-agent-rag-operation"; +export const RAG_STREAM_OPERATION_NAME = "bedrock-agent-rag-stream-operation"; diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/package.json b/e2e/scenarios/bedrock-agent-runtime-instrumentation/package.json new file mode 100644 index 000000000..607742de0 --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/package.json @@ -0,0 +1,18 @@ +{ + "name": "@braintrust/e2e-bedrock-agent-runtime-instrumentation", + "private": true, + "braintrustScenario": { + "bump": { + "dependencies": { + "bedrock-agent-runtime-sdk-v3-latest": { + "package": "@aws-sdk/client-bedrock-agent-runtime", + "range": "3" + } + } + } + }, + "dependencies": { + "bedrock-agent-runtime-sdk-v3": "npm:@aws-sdk/client-bedrock-agent-runtime@3.1048.0", + "bedrock-agent-runtime-sdk-v3-latest": "npm:@aws-sdk/client-bedrock-agent-runtime@3.1098.0" + } +} diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/pnpm-lock.yaml b/e2e/scenarios/bedrock-agent-runtime-instrumentation/pnpm-lock.yaml new file mode 100644 index 000000000..0d049a330 --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/pnpm-lock.yaml @@ -0,0 +1,396 @@ +lockfileVersion: '9.0' + +settings: + autoInstallPeers: true + excludeLinksFromLockfile: false + +importers: + + .: + dependencies: + bedrock-agent-runtime-sdk-v3: + specifier: npm:@aws-sdk/client-bedrock-agent-runtime@3.1048.0 + version: '@aws-sdk/client-bedrock-agent-runtime@3.1048.0' + bedrock-agent-runtime-sdk-v3-latest: + specifier: npm:@aws-sdk/client-bedrock-agent-runtime@3.1098.0 + version: '@aws-sdk/client-bedrock-agent-runtime@3.1098.0' + +packages: + + '@aws-crypto/sha256-browser@5.2.0': + resolution: {integrity: sha512-AXfN/lGotSQwu6HNcEsIASo7kWXZ5HYWvfOmSNKDsEqC4OashTp8alTmaz+F7TC2L083SFv5RdB+qU3Vs1kZqw==} + + '@aws-crypto/sha256-js@5.2.0': + resolution: {integrity: sha512-FFQQyu7edu4ufvIZ+OadFpHHOt+eSTBaYaki44c+akjg7qZg9oOQeLlk77F6tSYqjDAFClrHJk9tMf0HdVyOvA==} + engines: {node: '>=16.0.0'} + + '@aws-crypto/supports-web-crypto@5.2.0': + resolution: {integrity: sha512-iAvUotm021kM33eCdNfwIN//F77/IADDSs58i+MDaOqFrVjZo9bAal0NK7HurRuWLLpF1iLX7gbWrjHjeo+YFg==} + + '@aws-crypto/util@5.2.0': + resolution: {integrity: sha512-4RkU9EsI6ZpBve5fseQlGNUWKMa1RLPQ1dnjnQoe07ldfIzcsGb5hC5W0Dm7u423KWzawlrpbjXBrXCEv9zazQ==} + + '@aws-sdk/client-bedrock-agent-runtime@3.1048.0': + resolution: {integrity: sha512-X/lrKX9ubWaYal53ikoEynMQ/0fMl9D2ZEPKMcTsF+zPTqedYlmNTP4Y1BV/K7vIbvu4DPuFpQ1T8boaUI/NQQ==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/client-bedrock-agent-runtime@3.1098.0': + resolution: {integrity: sha512-VF8dR/C3Sr01VSKQnEHVwbvUPfHW3hTeKVvF5B6bbFhUyR2HpjqQrolYc3wGumomX6U3hyyo1v5iplH5dopPvQ==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/core@3.977.3': + resolution: {integrity: sha512-6YZTpF5Zzl0KdSrztM0l6ErKdnXrnnodIDYfayHE9SFfZU8Ubxls7H2XnfWT121jgwiG+WP1z5kyfVDsH5V4qA==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/credential-provider-env@3.972.64': + resolution: {integrity: sha512-14kkR5aj1c7D+cYCrEcG2W3xw87wMYAEUeB/tMiQ2j0RVKzzdewd4mJw1+Q0JVLr4IFU1JHGDCyj0WqLrtIixg==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/credential-provider-http@3.972.66': + resolution: {integrity: sha512-BrZxoXScBZ+nTOYy388zHkz1n5cS8C+uGFMozD4w9OKEWMq39Cu/9l/k385Hb54dtv0VIeh12f8h67o3xDNH9g==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/credential-provider-ini@3.973.9': + resolution: {integrity: sha512-d/mMvS+sQjSWNA6+JCldK1YB/jQmtoWM8Izj10mrOZdRCeWNVb8IYQuuyHHKKKXcpGo5Pd0LIK6onHbLvlhuVQ==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/credential-provider-login@3.972.71': + resolution: {integrity: sha512-l93922lXnTs5RjM3QrMCXSmXiEgLFLaQ0Lco9F3tJ08o0mzGdAj6m2nAXDuTMHoUuL0u7RKmrhm+Z8/7GCilyg==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/credential-provider-node@3.972.75': + resolution: {integrity: sha512-sdFhR4E3HWZefGL6OsBhxqXdqtFvxGOh8VLfiu9yttPgVNcpEozEDPmrTYmXA3CIf6npUETVLveQPGQ7GBzVhg==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/credential-provider-process@3.972.64': + resolution: {integrity: sha512-Q5noG5vfFo++bUaLGcjj8OHX5tmwa1O/5nDVMeaSSZgjzgvWbW6tgu/gSa+aENwcJXjoXDQh+8TRZR2zHds/aw==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/credential-provider-sso@3.973.8': + resolution: {integrity: sha512-37okeZyRdVkGrU6nyKBPGqQvyl/7CEiMpUVOb0+BCNCIR0VAuBxNgjMb0mnsR1EXm8dqQFfQZOJ5YW/0sBfa1w==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/credential-provider-web-identity@3.972.70': + resolution: {integrity: sha512-Ps/f9USafScb6CSQTRPNMzdKL5x5XRwDIRgFvnuFWfClGZe7/fI7iCFqXdxKjc9LBxP28E7iUgNloGgfgw406w==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/nested-clients@3.997.38': + resolution: {integrity: sha512-jQDs+nxnmgo4+WPOmY373HIC8jx4KyYYMYDI74FWrzX3Ba9fYsV0dSA+7PNYNkhRbbRCRLre83tQ31F1ACCHkA==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/signature-v4-multi-region@3.996.43': + resolution: {integrity: sha512-lKekx8bLBXSv4O+cslk9Zfnw2XKSkWBs3uWL5QGhH2ZAQfNS7FE0vcSSN2vD/AhxX54ZTywWxR4STThoeOXlBA==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/token-providers@3.1098.0': + resolution: {integrity: sha512-6cFAviffeqYGjrXn4FqGbWcTxB3nbpBpQXUm6MvnrWZThJaVBTHHnespHgSUh4yVLfzhipSh9H/XjzkHQd9P4g==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/types@3.974.2': + resolution: {integrity: sha512-3W6IUtSxFbH6X7Wb7DzGCV5QiFQsd0g8bOfntpmDxQlzBoKWUMBu/JPQR0DwkE+Hpnxd6db1tXbOwdeHddG6cA==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/util-locate-window@3.965.8': + resolution: {integrity: sha512-uUbMs1cBZPafD0ohUj6EwNf0fPZ534NvBxHox4hjX+0Rxq5paSYUem7+hi833pYrzrcnBATKIYpR02MDXT5M9g==} + engines: {node: '>=20.0.0'} + + '@aws-sdk/xml-builder@3.972.37': + resolution: {integrity: sha512-zKq4HQum8JwDyEuyfuI4bbiAcU0KxP6qy+9PR/IsR92IyE/DaBAikzAS50tjxip4bqIIANpCcG+Yyj6CVhXupg==} + engines: {node: '>=20.0.0'} + + '@aws/lambda-invoke-store@0.3.0': + resolution: {integrity: sha512-sl4Bm6yiMNYrZKkqqDFWN0UfnWhlS8ivKxrYl+6t0gCLrqr8y3B2IqZZbFRkfaVVp7C/baApyh71P+LeE1A2sQ==} + engines: {node: '>=18.0.0'} + + '@smithy/core@3.31.1': + resolution: {integrity: sha512-CyogUINxvi7C7LDsh8Syo6hVJOT9ckz4rG8dRZfTJ8r91HkMY59PnNooaj7WcHyxEkxPfBAmbgztZU+xTo76lg==} + engines: {node: '>=18.0.0'} + + '@smithy/credential-provider-imds@4.4.16': + resolution: {integrity: sha512-QfuLWAkLzptffFW980AFeHZFdqds2B64rpEd3uJ6lgs3xVn9QegGMUgUcj+4d7dRrAsya3r58ZKpku97WcFb4w==} + engines: {node: '>=18.0.0'} + + '@smithy/fetch-http-handler@5.6.13': + resolution: {integrity: sha512-4fW86pEUOMbrD5nkbyl/tTvPHHWJFbuB2odl6ps9lWfHoXf9HWh3Q/Smh59qH1g7+c/BSZghX6bbUk4gsiMs8A==} + engines: {node: '>=18.0.0'} + + '@smithy/is-array-buffer@2.2.0': + resolution: {integrity: sha512-GGP3O9QFD24uGeAXYUjwSTXARoqpZykHadOmA8G5vfJPK0/DC67qa//0qvqrJzL1xc8WQWX7/yc7fwudjPHPhA==} + engines: {node: '>=14.0.0'} + + '@smithy/node-http-handler@4.8.1': + resolution: {integrity: sha512-emtXvoky671puri18ETf64AFIQUGIEA093F2drXpBgB0OGnBLjcwNR3CA2mYu62IAqNsS56xa5lnTxAgPq7cjw==} + engines: {node: '>=18.0.0'} + + '@smithy/node-http-handler@4.9.13': + resolution: {integrity: sha512-Nmd/Nl35zfYrd+a6OO2cDJb3GPh9bgTjIUhcM+JFfjpp8/osCgboDV5nCT1I01Pv6R13eSKDKLSoVa5ZB6Zsfw==} + engines: {node: '>=18.0.0'} + + '@smithy/signature-v4@5.6.12': + resolution: {integrity: sha512-I6KLtq3H0qqSuV9vLglfi8puHqzygzWHOnI4z/Rdoo+q50vvo18vBRdPAvvEtcaKROz7Zn6qnPa14kRfPH6PcQ==} + engines: {node: '>=18.0.0'} + + '@smithy/types@4.16.1': + resolution: {integrity: sha512-0JFs3V2y2M9tKW5na/qxe69Zv+uxLMO7QBbhxF/FHu/Gp2NFZAAL9tWl9PU02xxo07pb3G9FTyjNc6D5uZrJIg==} + engines: {node: '>=18.0.0'} + + '@smithy/util-buffer-from@2.2.0': + resolution: {integrity: sha512-IJdWBbTcMQ6DA0gdNhh/BwrLkDR+ADW5Kr1aZmd4k3DIF6ezMV4R2NIAmT08wQJ3yUK82thHWmC/TnK/wpMMIA==} + engines: {node: '>=14.0.0'} + + '@smithy/util-utf8@2.3.0': + resolution: {integrity: sha512-R8Rdn8Hy72KKcebgLiv8jQcQkXoLMOGGv5uI1/k0l+snqkOzQ1R0ChUBCxWMlBsFMekWjq0wRudIweFs7sKT5A==} + engines: {node: '>=14.0.0'} + + bowser@2.14.1: + resolution: {integrity: sha512-tzPjzCxygAKWFOJP011oxFHs57HzIhOEracIgAePE4pqB3LikALKnSzUyU4MGs9/iCEUuHlAJTjTc5M+u7YEGg==} + + tslib@2.8.1: + resolution: {integrity: sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==} + +snapshots: + + '@aws-crypto/sha256-browser@5.2.0': + dependencies: + '@aws-crypto/sha256-js': 5.2.0 + '@aws-crypto/supports-web-crypto': 5.2.0 + '@aws-crypto/util': 5.2.0 + '@aws-sdk/types': 3.974.2 + '@aws-sdk/util-locate-window': 3.965.8 + '@smithy/util-utf8': 2.3.0 + tslib: 2.8.1 + + '@aws-crypto/sha256-js@5.2.0': + dependencies: + '@aws-crypto/util': 5.2.0 + '@aws-sdk/types': 3.974.2 + tslib: 2.8.1 + + '@aws-crypto/supports-web-crypto@5.2.0': + dependencies: + tslib: 2.8.1 + + '@aws-crypto/util@5.2.0': + dependencies: + '@aws-sdk/types': 3.974.2 + '@smithy/util-utf8': 2.3.0 + tslib: 2.8.1 + + '@aws-sdk/client-bedrock-agent-runtime@3.1048.0': + dependencies: + '@aws-crypto/sha256-browser': 5.2.0 + '@aws-crypto/sha256-js': 5.2.0 + '@aws-sdk/core': 3.977.3 + '@aws-sdk/credential-provider-node': 3.972.75 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/fetch-http-handler': 5.6.13 + '@smithy/node-http-handler': 4.8.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/client-bedrock-agent-runtime@3.1098.0': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/credential-provider-node': 3.972.75 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/fetch-http-handler': 5.6.13 + '@smithy/node-http-handler': 4.9.13 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/core@3.977.3': + dependencies: + '@aws-sdk/types': 3.974.2 + '@aws-sdk/xml-builder': 3.972.37 + '@aws/lambda-invoke-store': 0.3.0 + '@smithy/core': 3.31.1 + '@smithy/signature-v4': 5.6.12 + '@smithy/types': 4.16.1 + bowser: 2.14.1 + tslib: 2.8.1 + + '@aws-sdk/credential-provider-env@3.972.64': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/credential-provider-http@3.972.66': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/fetch-http-handler': 5.6.13 + '@smithy/node-http-handler': 4.9.13 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/credential-provider-ini@3.973.9': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/credential-provider-env': 3.972.64 + '@aws-sdk/credential-provider-http': 3.972.66 + '@aws-sdk/credential-provider-login': 3.972.71 + '@aws-sdk/credential-provider-process': 3.972.64 + '@aws-sdk/credential-provider-sso': 3.973.8 + '@aws-sdk/credential-provider-web-identity': 3.972.70 + '@aws-sdk/nested-clients': 3.997.38 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/credential-provider-imds': 4.4.16 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/credential-provider-login@3.972.71': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/nested-clients': 3.997.38 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/credential-provider-node@3.972.75': + dependencies: + '@aws-sdk/credential-provider-env': 3.972.64 + '@aws-sdk/credential-provider-http': 3.972.66 + '@aws-sdk/credential-provider-ini': 3.973.9 + '@aws-sdk/credential-provider-process': 3.972.64 + '@aws-sdk/credential-provider-sso': 3.973.8 + '@aws-sdk/credential-provider-web-identity': 3.972.70 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/credential-provider-imds': 4.4.16 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/credential-provider-process@3.972.64': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/credential-provider-sso@3.973.8': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/nested-clients': 3.997.38 + '@aws-sdk/token-providers': 3.1098.0 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/credential-provider-web-identity@3.972.70': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/nested-clients': 3.997.38 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/nested-clients@3.997.38': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/signature-v4-multi-region': 3.996.43 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/fetch-http-handler': 5.6.13 + '@smithy/node-http-handler': 4.9.13 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/signature-v4-multi-region@3.996.43': + dependencies: + '@aws-sdk/types': 3.974.2 + '@smithy/signature-v4': 5.6.12 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/token-providers@3.1098.0': + dependencies: + '@aws-sdk/core': 3.977.3 + '@aws-sdk/nested-clients': 3.997.38 + '@aws-sdk/types': 3.974.2 + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/types@3.974.2': + dependencies: + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws-sdk/util-locate-window@3.965.8': + dependencies: + tslib: 2.8.1 + + '@aws-sdk/xml-builder@3.972.37': + dependencies: + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@aws/lambda-invoke-store@0.3.0': {} + + '@smithy/core@3.31.1': + dependencies: + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@smithy/credential-provider-imds@4.4.16': + dependencies: + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@smithy/fetch-http-handler@5.6.13': + dependencies: + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@smithy/is-array-buffer@2.2.0': + dependencies: + tslib: 2.8.1 + + '@smithy/node-http-handler@4.8.1': + dependencies: + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@smithy/node-http-handler@4.9.13': + dependencies: + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@smithy/signature-v4@5.6.12': + dependencies: + '@smithy/core': 3.31.1 + '@smithy/types': 4.16.1 + tslib: 2.8.1 + + '@smithy/types@4.16.1': + dependencies: + tslib: 2.8.1 + + '@smithy/util-buffer-from@2.2.0': + dependencies: + '@smithy/is-array-buffer': 2.2.0 + tslib: 2.8.1 + + '@smithy/util-utf8@2.3.0': + dependencies: + '@smithy/util-buffer-from': 2.2.0 + tslib: 2.8.1 + + bowser@2.14.1: {} + + tslib@2.8.1: {} diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.impl.mjs b/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.impl.mjs new file mode 100644 index 000000000..050c72736 --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.impl.mjs @@ -0,0 +1,304 @@ +import { wrapBedrockAgentRuntime } from "braintrust"; +import { + collectAsync, + runOperation, + runTracedScenario, +} from "../../helpers/provider-runtime.mjs"; +import { + AGENT_OPERATION_NAME, + INLINE_OPERATION_NAME, + MODEL, + RAG_OPERATION_NAME, + RAG_STREAM_OPERATION_NAME, + REGION, + ROOT_NAME, + SCENARIO_NAME, +} from "./constants.mjs"; + +export const BEDROCK_AGENT_RUNTIME_SCENARIO_TIMEOUT_MS = 180_000; + +function traceEvent(trace) { + return { + trace: { + agentAliasId: "alias-1", + agentId: "agent-1", + agentVersion: "1", + sessionId: "session-1", + trace, + }, + }; +} + +function modelTraceEvents(traceId, output) { + return [ + traceEvent({ + orchestrationTrace: { + modelInvocationInput: { + foundationModel: MODEL, + inferenceConfiguration: { + maximumLength: 128, + temperature: 0, + }, + text: "deterministic Bedrock agent model prompt", + traceId, + type: "ORCHESTRATION", + }, + }, + }), + traceEvent({ + orchestrationTrace: { + modelInvocationOutput: { + metadata: { + usage: { + inputTokens: 12, + outputTokens: 4, + }, + }, + rawResponse: { + content: JSON.stringify(output), + }, + traceId, + }, + }, + }), + ]; +} + +async function* invokeAgentCompletion() { + for (const event of modelTraceEvents("agent-model", { + action: "get_weather", + })) { + yield event; + } + yield traceEvent({ + orchestrationTrace: { + invocationInput: { + actionGroupInvocationInput: { + actionGroupName: "weather", + function: "get_weather", + parameters: [{ name: "city", value: "Paris" }], + }, + traceId: "agent-tool", + }, + }, + }); + yield traceEvent({ + orchestrationTrace: { + observation: { + actionGroupInvocationOutput: { + text: JSON.stringify({ temperature: 18 }), + }, + traceId: "agent-tool", + type: "ACTION_GROUP", + }, + }, + }); + yield { + chunk: { + bytes: new TextEncoder().encode("AGENT_OK"), + }, + }; +} + +async function* invokeInlineAgentCompletion() { + for (const event of modelTraceEvents("inline-model", { + answer: "INLINE_OK", + })) { + yield event; + } + yield { + chunk: { + bytes: new TextEncoder().encode("INLINE_OK"), + }, + }; +} + +async function* retrieveAndGenerateStream() { + yield { output: { text: "STREAM_" } }; + yield { output: { text: "RAG_OK" } }; + yield { + citation: { + generatedResponsePart: { + textResponsePart: { + span: { start: 0, end: 12 }, + text: "STREAM_RAG_OK", + }, + }, + retrievedReferences: [ + { + content: { text: "Deterministic knowledge-base reference." }, + location: { + type: "WEB", + webLocation: { url: "https://example.com/reference" }, + }, + }, + ], + }, + }; + yield { guardrail: { action: "NONE" } }; +} + +function responseFor(commandName) { + switch (commandName) { + case "InvokeAgentCommand": + return { + $metadata: { httpStatusCode: 200, requestId: "request-agent" }, + completion: invokeAgentCompletion(), + contentType: "application/json", + sessionId: "session-1", + }; + case "InvokeInlineAgentCommand": + return { + $metadata: { httpStatusCode: 200, requestId: "request-inline" }, + completion: invokeInlineAgentCompletion(), + contentType: "application/json", + sessionId: "session-inline", + }; + case "RetrieveAndGenerateCommand": + return { + $metadata: { httpStatusCode: 200, requestId: "request-rag" }, + citations: [ + { + generatedResponsePart: { + textResponsePart: { text: "RAG_OK" }, + }, + }, + ], + guardrailAction: "NONE", + output: { text: "RAG_OK" }, + sessionId: "session-rag", + }; + case "RetrieveAndGenerateStreamCommand": + return { + $metadata: { httpStatusCode: 200, requestId: "request-rag-stream" }, + sessionId: "session-rag-stream", + stream: retrieveAndGenerateStream(), + }; + default: + throw new Error( + `Unexpected Bedrock Agent Runtime command: ${commandName}`, + ); + } +} + +function installDeterministicBackend(client) { + client.middlewareStack.add( + (_next, context) => async () => ({ + output: responseFor(context.commandName), + response: {}, + }), + { + name: "braintrustBedrockAgentRuntimeE2EBackend", + priority: "high", + step: "initialize", + }, + ); +} + +function knowledgeBaseConfiguration() { + return { + knowledgeBaseConfiguration: { + knowledgeBaseId: "kb-1", + modelArn: `arn:aws:bedrock:${REGION}:000000000000:foundation-model/${MODEL}`, + }, + type: "KNOWLEDGE_BASE", + }; +} + +export async function runBedrockAgentRuntimeInstrumentationScenario(options) { + const baseClient = new options.BedrockAgentRuntimeClient({ + credentials: { + accessKeyId: "e2e-access-key", + secretAccessKey: "e2e-secret-key", + }, + region: REGION, + }); + installDeterministicBackend(baseClient); + const client = options.decorateClient + ? options.decorateClient(baseClient) + : baseClient; + + try { + await runTracedScenario({ + callback: async () => { + await runOperation(AGENT_OPERATION_NAME, "invoke-agent", async () => { + const response = await client.send( + new options.InvokeAgentCommand({ + agentAliasId: "alias-1", + agentId: "agent-1", + enableTrace: true, + inputText: "Use the weather tool and reply with AGENT_OK.", + sessionId: "session-1", + }), + ); + await collectAsync(response.completion ?? []); + }); + + await runOperation( + INLINE_OPERATION_NAME, + "invoke-inline-agent", + async () => { + const response = await client.send( + new options.InvokeInlineAgentCommand({ + agentName: "braintrust-e2e-inline-agent", + enableTrace: true, + endSession: true, + foundationModel: MODEL, + inputText: "Reply with exactly INLINE_OK.", + instruction: + "Follow the user's requested response format exactly.", + sessionId: "session-inline", + }), + ); + await collectAsync(response.completion ?? []); + }, + ); + + await runOperation( + RAG_OPERATION_NAME, + "retrieve-and-generate", + async () => { + await client.send( + new options.RetrieveAndGenerateCommand({ + input: { text: "Return the deterministic RAG answer." }, + retrieveAndGenerateConfiguration: knowledgeBaseConfiguration(), + }), + ); + }, + ); + + await runOperation( + RAG_STREAM_OPERATION_NAME, + "retrieve-and-generate-stream", + async () => { + const response = await client.send( + new options.RetrieveAndGenerateStreamCommand({ + input: { text: "Stream the deterministic RAG answer." }, + retrieveAndGenerateConfiguration: knowledgeBaseConfiguration(), + }), + ); + await collectAsync(response.stream ?? []); + }, + ); + }, + metadata: { + scenario: SCENARIO_NAME, + }, + projectNameBase: "e2e-bedrock-agent-runtime-instrumentation", + rootName: ROOT_NAME, + }); + } finally { + baseClient.destroy(); + } +} + +export async function runWrappedBedrockAgentRuntimeInstrumentation(options) { + await runBedrockAgentRuntimeInstrumentationScenario({ + decorateClient: wrapBedrockAgentRuntime, + ...options, + }); +} + +export async function runAutoBedrockAgentRuntimeInstrumentation(options) { + await runBedrockAgentRuntimeInstrumentationScenario(options); +} diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.mjs b/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.mjs new file mode 100644 index 000000000..e789c8718 --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.mjs @@ -0,0 +1,22 @@ +const packageName = + process.env.BEDROCK_AGENT_RUNTIME_PACKAGE_NAME ?? + "bedrock-agent-runtime-sdk-v3-latest"; +const { + BedrockAgentRuntimeClient, + InvokeAgentCommand, + InvokeInlineAgentCommand, + RetrieveAndGenerateCommand, + RetrieveAndGenerateStreamCommand, +} = await import(packageName); +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAutoBedrockAgentRuntimeInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => { + await runAutoBedrockAgentRuntimeInstrumentation({ + BedrockAgentRuntimeClient, + InvokeAgentCommand, + InvokeInlineAgentCommand, + RetrieveAndGenerateCommand, + RetrieveAndGenerateStreamCommand, + }); +}); diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.test.ts b/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.test.ts new file mode 100644 index 000000000..2a63ebeab --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.test.ts @@ -0,0 +1,80 @@ +import { describe } from "vitest"; +import { + prepareScenarioDir, + readInstalledPackageVersion, + resolveScenarioDir, +} from "../../helpers/scenario-harness"; +import { defineBedrockAgentRuntimeInstrumentationAssertions } from "./assertions"; +import { BEDROCK_AGENT_RUNTIME_SCENARIO_TIMEOUT_MS } from "./scenario.impl.mjs"; + +const originalScenarioDir = resolveScenarioDir(import.meta.url); +const scenarioDir = await prepareScenarioDir({ + scenarioDir: originalScenarioDir, +}); +const scenarios = await Promise.all( + [ + { + dependencyName: "bedrock-agent-runtime-sdk-v3", + snapshotName: "bedrock-agent-runtime-v3", + }, + { + dependencyName: "bedrock-agent-runtime-sdk-v3-latest", + snapshotName: "bedrock-agent-runtime-v3-latest", + }, + ].map(async (scenario) => ({ + ...scenario, + version: await readInstalledPackageVersion( + scenarioDir, + scenario.dependencyName, + ), + })), +); + +describe.concurrent("variants", () => { + for (const scenario of scenarios) { + describe.sequential(`bedrock agent runtime sdk ${scenario.version}`, () => { + defineBedrockAgentRuntimeInstrumentationAssertions({ + name: "wrapped instrumentation", + runScenario: async ({ runScenarioDir }) => { + await runScenarioDir({ + entry: "scenario.ts", + env: { + BEDROCK_AGENT_RUNTIME_PACKAGE_NAME: scenario.dependencyName, + }, + runContext: { + cassette: false, + variantKey: scenario.snapshotName, + }, + scenarioDir, + timeoutMs: BEDROCK_AGENT_RUNTIME_SCENARIO_TIMEOUT_MS, + }); + }, + snapshotName: `${scenario.snapshotName}-wrapped`, + testFileUrl: import.meta.url, + timeoutMs: BEDROCK_AGENT_RUNTIME_SCENARIO_TIMEOUT_MS, + }); + + defineBedrockAgentRuntimeInstrumentationAssertions({ + name: "auto-hook instrumentation", + runScenario: async ({ runNodeScenarioDir }) => { + await runNodeScenarioDir({ + entry: "scenario.mjs", + env: { + BEDROCK_AGENT_RUNTIME_PACKAGE_NAME: scenario.dependencyName, + }, + nodeArgs: ["--import", "braintrust/hook.mjs"], + runContext: { + cassette: false, + variantKey: scenario.snapshotName, + }, + scenarioDir, + timeoutMs: BEDROCK_AGENT_RUNTIME_SCENARIO_TIMEOUT_MS, + }); + }, + snapshotName: `${scenario.snapshotName}-auto`, + testFileUrl: import.meta.url, + timeoutMs: BEDROCK_AGENT_RUNTIME_SCENARIO_TIMEOUT_MS, + }); + }); + } +}); diff --git a/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.ts b/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.ts new file mode 100644 index 000000000..97471d72e --- /dev/null +++ b/e2e/scenarios/bedrock-agent-runtime-instrumentation/scenario.ts @@ -0,0 +1,22 @@ +const packageName = + process.env.BEDROCK_AGENT_RUNTIME_PACKAGE_NAME ?? + "bedrock-agent-runtime-sdk-v3-latest"; +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runWrappedBedrockAgentRuntimeInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => { + const { + BedrockAgentRuntimeClient, + InvokeAgentCommand, + InvokeInlineAgentCommand, + RetrieveAndGenerateCommand, + RetrieveAndGenerateStreamCommand, + } = await import(packageName); + await runWrappedBedrockAgentRuntimeInstrumentation({ + BedrockAgentRuntimeClient, + InvokeAgentCommand, + InvokeInlineAgentCommand, + RetrieveAndGenerateCommand, + RetrieveAndGenerateStreamCommand, + }); +}); diff --git a/js/src/auto-instrumentations/configs/all.ts b/js/src/auto-instrumentations/configs/all.ts index 905715e43..4a4bb601d 100644 --- a/js/src/auto-instrumentations/configs/all.ts +++ b/js/src/auto-instrumentations/configs/all.ts @@ -33,6 +33,7 @@ import { strandsAgentSDKConfigs } from "./strands-agent-sdk"; interface InstrumentationConfigGroup { integrations: readonly (keyof InstrumentationIntegrationsConfig)[]; + anyIntegrations?: readonly (keyof InstrumentationIntegrationsConfig)[]; configs: readonly InstrumentationConfig[]; } @@ -45,7 +46,8 @@ const defaultInstrumentationConfigGroups: readonly InstrumentationConfigGroup[] }, { integrations: ["anthropic"], configs: anthropicConfigs }, { - integrations: ["bedrock", "awsBedrock", "awsBedrockRuntime"], + integrations: ["bedrock", "awsBedrock"], + anyIntegrations: ["awsBedrockRuntime", "awsBedrockAgentRuntime"], configs: bedrockRuntimeConfigs, }, { @@ -139,10 +141,21 @@ export function getDefaultInstrumentationConfigs({ return [ ...defaultInstrumentationConfigGroups.flatMap( - ({ configs, integrations }) => - isInstrumentationIntegrationDisabled(disabledConfig, ...integrations) - ? [] - : configs, + ({ anyIntegrations, configs, integrations }) => { + if ( + isInstrumentationIntegrationDisabled(disabledConfig, ...integrations) + ) { + return []; + } + if ( + anyIntegrations?.every( + (integration) => disabledConfig?.[integration] === false, + ) + ) { + return []; + } + return configs; + }, ), ...(additionalInstrumentations ?? []), ]; diff --git a/js/src/auto-instrumentations/configs/bedrock-agent-runtime.test.ts b/js/src/auto-instrumentations/configs/bedrock-agent-runtime.test.ts new file mode 100644 index 000000000..a64aec402 --- /dev/null +++ b/js/src/auto-instrumentations/configs/bedrock-agent-runtime.test.ts @@ -0,0 +1,68 @@ +import { describe, expect, it } from "vitest"; +import { + readDisabledInstrumentationEnvConfig, + type InstrumentationIntegrationsConfig, +} from "../../instrumentation/config"; +import { getDefaultInstrumentationConfigs } from "./all"; +import { bedrockAgentRuntimeConfigs } from "./bedrock-agent-runtime"; +import { bedrockRuntimeConfigs } from "./bedrock-runtime"; + +const smithyClientSendChannel = "client.send"; + +function includesSmithyClientSend( + disabledIntegrationConfig?: InstrumentationIntegrationsConfig, +): boolean { + return getDefaultInstrumentationConfigs({ + disabledIntegrationConfig, + }).some( + (config) => + config.channelName === smithyClientSendChannel && + config.module.name === "@smithy/core", + ); +} + +describe("bedrockAgentRuntimeConfigs", () => { + it("reuses the Smithy Client.send transformations", () => { + expect(bedrockAgentRuntimeConfigs).toBe(bedrockRuntimeConfigs); + expect( + bedrockAgentRuntimeConfigs.some( + (config) => + config.channelName === smithyClientSendChannel && + config.module.name === "@smithy/core", + ), + ).toBe(true); + }); + + it("keeps shared transformations while either Bedrock integration is enabled", () => { + expect( + includesSmithyClientSend({ + awsBedrockAgentRuntime: false, + awsBedrockRuntime: true, + }), + ).toBe(true); + expect( + includesSmithyClientSend({ + awsBedrockAgentRuntime: true, + awsBedrockRuntime: false, + }), + ).toBe(true); + expect( + includesSmithyClientSend({ + awsBedrockAgentRuntime: false, + awsBedrockRuntime: false, + }), + ).toBe(false); + }); + + it("supports Agent Runtime disable aliases", () => { + for (const alias of [ + "aws-bedrock-agent-runtime", + "awsbedrockagentruntime", + "@aws-sdk/client-bedrock-agent-runtime", + ]) { + expect(readDisabledInstrumentationEnvConfig(alias).integrations).toEqual({ + awsBedrockAgentRuntime: false, + }); + } + }); +}); diff --git a/js/src/auto-instrumentations/configs/bedrock-agent-runtime.ts b/js/src/auto-instrumentations/configs/bedrock-agent-runtime.ts new file mode 100644 index 000000000..41ca8b57b --- /dev/null +++ b/js/src/auto-instrumentations/configs/bedrock-agent-runtime.ts @@ -0,0 +1,4 @@ +// Bedrock Runtime and Bedrock Agent Runtime both inherit Client.send from +// Smithy. Reuse the same transformer configs and dispatch by command name in +// their respective plugins. +export { bedrockRuntimeConfigs as bedrockAgentRuntimeConfigs } from "./bedrock-runtime"; diff --git a/js/src/auto-instrumentations/index.ts b/js/src/auto-instrumentations/index.ts index 49f728f82..bb511239c 100644 --- a/js/src/auto-instrumentations/index.ts +++ b/js/src/auto-instrumentations/index.ts @@ -31,6 +31,7 @@ export { openaiConfigs } from "./configs/openai"; export { openAICodexConfigs } from "./configs/openai-codex"; export { anthropicConfigs } from "./configs/anthropic"; +export { bedrockAgentRuntimeConfigs } from "./configs/bedrock-agent-runtime"; export { bedrockRuntimeConfigs } from "./configs/bedrock-runtime"; export { aiSDKConfigs } from "./configs/ai-sdk"; export { claudeAgentSDKConfigs } from "./configs/claude-agent-sdk"; diff --git a/js/src/exports.ts b/js/src/exports.ts index 02104f8c2..4683f0feb 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -223,6 +223,7 @@ export { wrapOpenRouter } from "./wrappers/openrouter"; export { wrapMistral } from "./wrappers/mistral"; export { wrapCohere } from "./wrappers/cohere"; export { wrapGroq } from "./wrappers/groq"; +export { wrapBedrockAgentRuntime } from "./wrappers/bedrock-agent-runtime"; export { wrapBedrockRuntime } from "./wrappers/bedrock-runtime"; export { wrapCopilotClient } from "./wrappers/github-copilot"; export { diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index ec88ddec2..61f456866 100644 --- a/js/src/instrumentation/braintrust-plugin.ts +++ b/js/src/instrumentation/braintrust-plugin.ts @@ -15,6 +15,7 @@ import { MistralPlugin } from "./plugins/mistral-plugin"; import { GoogleADKPlugin } from "./plugins/google-adk-plugin"; import { CoherePlugin } from "./plugins/cohere-plugin"; import { GroqPlugin } from "./plugins/groq-plugin"; +import { BedrockAgentRuntimePlugin } from "./plugins/bedrock-agent-runtime-plugin"; import { BedrockRuntimePlugin } from "./plugins/bedrock-runtime-plugin"; import { GenkitPlugin } from "./plugins/genkit-plugin"; import { GitHubCopilotPlugin } from "./plugins/github-copilot-plugin"; @@ -67,6 +68,7 @@ export class BraintrustPlugin extends BasePlugin { private googleADKPlugin: GoogleADKPlugin | null = null; private coherePlugin: CoherePlugin | null = null; private groqPlugin: GroqPlugin | null = null; + private bedrockAgentRuntimePlugin: BedrockAgentRuntimePlugin | null = null; private bedrockRuntimePlugin: BedrockRuntimePlugin | null = null; private genkitPlugin: GenkitPlugin | null = null; private gitHubCopilotPlugin: GitHubCopilotPlugin | null = null; @@ -181,6 +183,15 @@ export class BraintrustPlugin extends BasePlugin { this.bedrockRuntimePlugin.enable(); } + if ( + integrations.bedrock !== false && + integrations.awsBedrock !== false && + integrations.awsBedrockAgentRuntime !== false + ) { + this.bedrockAgentRuntimePlugin = new BedrockAgentRuntimePlugin(); + this.bedrockAgentRuntimePlugin.enable(); + } + if (integrations.genkit !== false) { this.genkitPlugin = new GenkitPlugin(); this.genkitPlugin.enable(); @@ -321,6 +332,11 @@ export class BraintrustPlugin extends BasePlugin { this.bedrockRuntimePlugin = null; } + if (this.bedrockAgentRuntimePlugin) { + this.bedrockAgentRuntimePlugin.disable(); + this.bedrockAgentRuntimePlugin = null; + } + if (this.genkitPlugin) { this.genkitPlugin.disable(); this.genkitPlugin = null; diff --git a/js/src/instrumentation/config.ts b/js/src/instrumentation/config.ts index c51a5b48d..fdd318425 100644 --- a/js/src/instrumentation/config.ts +++ b/js/src/instrumentation/config.ts @@ -21,6 +21,7 @@ export interface InstrumentationIntegrationsConfig { groq?: boolean; bedrock?: boolean; awsBedrock?: boolean; + awsBedrockAgentRuntime?: boolean; awsBedrockRuntime?: boolean; genkit?: boolean; gitHubCopilot?: boolean; @@ -104,6 +105,9 @@ const envIntegrationAliases: Record< "aws-bedrock-runtime": "awsBedrockRuntime", awsbedrockruntime: "awsBedrockRuntime", "@aws-sdk/client-bedrock-runtime": "awsBedrockRuntime", + "aws-bedrock-agent-runtime": "awsBedrockAgentRuntime", + awsbedrockagentruntime: "awsBedrockAgentRuntime", + "@aws-sdk/client-bedrock-agent-runtime": "awsBedrockAgentRuntime", genkit: "genkit", "firebase-genkit": "genkit", githubcopilot: "gitHubCopilot", @@ -144,6 +148,7 @@ export function getDefaultInstrumentationIntegrations(): Record< groq: true, bedrock: true, awsBedrock: true, + awsBedrockAgentRuntime: true, awsBedrockRuntime: true, genkit: true, gitHubCopilot: true, diff --git a/js/src/instrumentation/plugins/bedrock-agent-runtime-channels.ts b/js/src/instrumentation/plugins/bedrock-agent-runtime-channels.ts new file mode 100644 index 000000000..52630a73e --- /dev/null +++ b/js/src/instrumentation/plugins/bedrock-agent-runtime-channels.ts @@ -0,0 +1,44 @@ +import { INSTRUMENTATION_NAMES } from "../../span-origin"; +import type { + BedrockAgentRuntimeChannelContext, + BedrockAgentRuntimeCommandLike, + BedrockAgentRuntimeSendResult, + BedrockAgentRuntimeStreamEvent, +} from "../../vendor-sdk-types/bedrock-agent-runtime"; +import { channel, defineChannels } from "../core/channel-definitions"; + +const clientSendChannel = channel< + [BedrockAgentRuntimeCommandLike, unknown?], + BedrockAgentRuntimeSendResult, + BedrockAgentRuntimeChannelContext, + BedrockAgentRuntimeStreamEvent +>({ + channelName: "client.send", + kind: "async", +}); + +export const bedrockAgentRuntimeChannels = defineChannels( + "aws-bedrock-agent-runtime", + { + clientSend: clientSendChannel, + }, + { instrumentationName: INSTRUMENTATION_NAMES.BEDROCK_AGENT_RUNTIME }, +); + +// These intentionally resolve to the same diagnostics channels as the Bedrock +// Runtime instrumentation. Smithy owns Client.send for both AWS clients. +export const bedrockAgentSmithyCoreChannels = defineChannels( + "@smithy/core", + { + clientSend: clientSendChannel, + }, + { instrumentationName: INSTRUMENTATION_NAMES.BEDROCK_AGENT_RUNTIME }, +); + +export const bedrockAgentSmithyClientChannels = defineChannels( + "@smithy/smithy-client", + { + clientSend: clientSendChannel, + }, + { instrumentationName: INSTRUMENTATION_NAMES.BEDROCK_AGENT_RUNTIME }, +); diff --git a/js/src/instrumentation/plugins/bedrock-agent-runtime-common.ts b/js/src/instrumentation/plugins/bedrock-agent-runtime-common.ts new file mode 100644 index 000000000..39627e7a7 --- /dev/null +++ b/js/src/instrumentation/plugins/bedrock-agent-runtime-common.ts @@ -0,0 +1,105 @@ +import { isObject } from "../../../util/index"; +import type { + BedrockAgentRuntimeCommandLike, + BedrockAgentRuntimeCommandName, +} from "../../vendor-sdk-types/bedrock-agent-runtime"; + +export type BedrockAgentRuntimeOperation = + | "invokeAgent" + | "invokeInlineAgent" + | "retrieveAndGenerate" + | "retrieveAndGenerateStream"; + +const COMMAND_OPERATIONS: Record< + BedrockAgentRuntimeCommandName, + BedrockAgentRuntimeOperation +> = { + InvokeAgentCommand: "invokeAgent", + InvokeInlineAgentCommand: "invokeInlineAgent", + RetrieveAndGenerateCommand: "retrieveAndGenerate", + RetrieveAndGenerateStreamCommand: "retrieveAndGenerateStream", +}; + +export function getBedrockAgentRuntimeCommandName( + command: unknown, +): BedrockAgentRuntimeCommandName | undefined { + if (!isObject(command) || !isObject(command.constructor)) { + return undefined; + } + + const commandName = command.constructor.name; + const input = (command as BedrockAgentRuntimeCommandLike).input; + if (!isObject(input) || !isBedrockAgentRuntimeCommandName(commandName)) { + return undefined; + } + + if ( + commandName === "InvokeAgentCommand" && + (typeof input.agentId !== "string" || + typeof input.agentAliasId !== "string" || + typeof input.sessionId !== "string") + ) { + return undefined; + } + + if ( + commandName === "InvokeInlineAgentCommand" && + (typeof input.foundationModel !== "string" || + typeof input.instruction !== "string" || + typeof input.sessionId !== "string") + ) { + return undefined; + } + + if ( + (commandName === "RetrieveAndGenerateCommand" || + commandName === "RetrieveAndGenerateStreamCommand") && + (!isObject(input.input) || typeof input.input.text !== "string") + ) { + return undefined; + } + + return commandName; +} + +export function getBedrockAgentRuntimeOperation( + command: unknown, +): BedrockAgentRuntimeOperation | undefined { + const commandName = getBedrockAgentRuntimeCommandName(command); + return commandName ? COMMAND_OPERATIONS[commandName] : undefined; +} + +export function getBedrockAgentRuntimeCommandInput( + command: unknown, +): unknown | undefined { + return isObject(command) + ? (command as BedrockAgentRuntimeCommandLike).input + : undefined; +} + +export function buildBedrockAgentRuntimeSpanInfo(command: unknown): { + name: string; + metadata: Record; +} { + const commandName = getBedrockAgentRuntimeCommandName(command); + const operation = getBedrockAgentRuntimeOperation(command); + + return { + name: operation ? `bedrockAgent.${operation}` : "bedrockAgent.client.send", + metadata: { + ...(commandName ? { command: commandName } : {}), + ...(operation ? { operation } : {}), + }, + }; +} + +function isBedrockAgentRuntimeCommandName( + value: unknown, +): value is BedrockAgentRuntimeCommandName { + return ( + value === "InvokeAgentCommand" || + value === "InvokeInlineAgentCommand" || + value === "RetrieveAndGenerateCommand" || + value === "RetrieveAndGenerateStreamCommand" + ); +} diff --git a/js/src/instrumentation/plugins/bedrock-agent-runtime-plugin.test.ts b/js/src/instrumentation/plugins/bedrock-agent-runtime-plugin.test.ts new file mode 100644 index 000000000..28b9d7af3 --- /dev/null +++ b/js/src/instrumentation/plugins/bedrock-agent-runtime-plugin.test.ts @@ -0,0 +1,431 @@ +import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; +import { _exportsForTestingOnly, initLogger } from "../../logger"; +import { configureNode } from "../../node/config"; +import { bedrockAgentSmithyCoreChannels } from "./bedrock-agent-runtime-channels"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +class InvokeAgentCommand { + constructor(public input: Record) {} +} + +class RetrieveAndGenerateCommand { + constructor(public input: Record) {} +} + +class RetrieveAndGenerateStreamCommand { + constructor(public input: Record) {} +} + +class CreateSessionCommand { + constructor(public input: Record) {} +} + +describe("BedrockAgentRuntimePlugin", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectName: "bedrock-agent-runtime-plugin.test.ts", + projectId: "test-project-id", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + it("captures an agent task with model and tool child spans", async () => { + async function* completion() { + yield { + trace: { + sessionId: "session-1", + trace: { + orchestrationTrace: { + modelInvocationInput: { + foundationModel: "amazon.nova-lite-v1:0", + inferenceConfiguration: { + maximumLength: 32, + temperature: 0, + }, + text: "model prompt", + traceId: "model-1", + type: "ORCHESTRATION", + }, + }, + }, + }, + }; + yield { + trace: { + sessionId: "session-1", + trace: { + orchestrationTrace: { + modelInvocationOutput: { + metadata: { + usage: { + inputTokens: 10, + outputTokens: 4, + }, + }, + rawResponse: { + content: JSON.stringify({ action: "get_weather" }), + }, + traceId: "model-1", + }, + }, + }, + }, + }; + yield { + trace: { + sessionId: "session-1", + trace: { + orchestrationTrace: { + invocationInput: { + actionGroupInvocationInput: { + actionGroupName: "weather", + function: "get_weather", + parameters: [{ name: "city", value: "Paris" }], + }, + traceId: "tool-1", + }, + }, + }, + }, + }; + yield { + trace: { + sessionId: "session-1", + trace: { + orchestrationTrace: { + observation: { + actionGroupInvocationOutput: { + text: JSON.stringify({ temperature: 18 }), + }, + traceId: "tool-1", + type: "ACTION_GROUP", + }, + }, + }, + }, + }; + yield { + chunk: { + bytes: new TextEncoder().encode("It is 18°C."), + }, + }; + } + + const response = + await bedrockAgentSmithyCoreChannels.clientSend.tracePromise( + async () => ({ + completion: completion(), + contentType: "application/json", + sessionId: "session-1", + }), + { + arguments: [ + new InvokeAgentCommand({ + agentAliasId: "alias-1", + agentId: "agent-1", + enableTrace: true, + inputText: "What is the weather in Paris?", + sessionId: "session-1", + }) as any, + ], + }, + ); + + for await (const _event of response.completion) { + // Consume the response so the instrumented span can finish. + } + + const spans = (await backgroundLogger.drain()) as Array< + Record & { + span_attributes?: { name?: string; type?: string }; + span_id?: string; + span_parents?: string[]; + } + >; + const parent = spans.find( + (span) => span.span_attributes?.name === "bedrockAgent.invokeAgent", + ); + const model = spans.find( + (span) => span.span_attributes?.name === "bedrockAgent.modelInvocation", + ); + const tool = spans.find( + (span) => span.span_attributes?.name === "get_weather", + ); + + expect(parent).toMatchObject({ + input: "What is the weather in Paris?", + metadata: { + agentAliasId: "alias-1", + agentId: "agent-1", + command: "InvokeAgentCommand", + operation: "invokeAgent", + provider: "aws-bedrock", + sessionId: "session-1", + }, + metrics: { + completion_tokens: 4, + prompt_tokens: 10, + time_to_first_token: expect.any(Number), + tokens: 14, + }, + output: { + text: "It is 18°C.", + }, + span_attributes: { + name: "bedrockAgent.invokeAgent", + type: "task", + }, + }); + expect(model).toMatchObject({ + input: "model prompt", + metadata: { + maximumLength: 32, + model: "amazon.nova-lite-v1:0", + prompt_type: "ORCHESTRATION", + provider: "aws-bedrock", + temperature: 0, + }, + metrics: { + completion_tokens: 4, + prompt_tokens: 10, + tokens: 14, + }, + output: { + action: "get_weather", + }, + span_attributes: { + name: "bedrockAgent.modelInvocation", + type: "llm", + }, + }); + expect(tool).toMatchObject({ + input: { + parameters: [{ name: "city", value: "Paris" }], + }, + metadata: { + action_group: "weather", + provider: "aws-bedrock", + tool_type: "action_group", + }, + output: { + text: JSON.stringify({ temperature: 18 }), + }, + span_attributes: { + name: "get_weather", + type: "tool", + }, + }); + expect(model?.span_parents).toEqual([parent?.span_id]); + expect(tool?.span_parents).toEqual([parent?.span_id]); + }); + + it("captures non-streaming retrieve-and-generate responses", async () => { + await bedrockAgentSmithyCoreChannels.clientSend.tracePromise( + async () => ({ + citations: [{ generatedResponsePart: { text: "Paris" } }], + guardrailAction: "NONE", + output: { text: "Paris is in France." }, + sessionId: "rag-session", + }), + { + arguments: [ + new RetrieveAndGenerateCommand({ + input: { text: "Where is Paris?" }, + retrieveAndGenerateConfiguration: { + knowledgeBaseConfiguration: { + knowledgeBaseId: "kb-1", + modelArn: "arn:aws:bedrock:model/test", + }, + type: "KNOWLEDGE_BASE", + }, + }) as any, + ], + }, + ); + + const spans = await backgroundLogger.drain(); + expect(spans).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + input: { text: "Where is Paris?" }, + metadata: expect.objectContaining({ + knowledgeBaseId: "kb-1", + model: "arn:aws:bedrock:model/test", + operation: "retrieveAndGenerate", + }), + output: { + citations: [{ generatedResponsePart: { text: "Paris" } }], + text: "Paris is in France.", + }, + span_attributes: expect.objectContaining({ + name: "bedrockAgent.retrieveAndGenerate", + type: "task", + }), + }), + ]), + ); + }); + + it("accumulates retrieve-and-generate stream output and citations", async () => { + async function* stream() { + yield { output: { text: "Paris " } }; + yield { output: { text: "is in France." } }; + yield { + citation: { + generatedResponsePart: { text: "Paris" }, + }, + }; + } + + const response = + await bedrockAgentSmithyCoreChannels.clientSend.tracePromise( + async () => ({ + sessionId: "rag-session", + stream: stream(), + }), + { + arguments: [ + new RetrieveAndGenerateStreamCommand({ + input: { text: "Where is Paris?" }, + }) as any, + ], + }, + ); + + for await (const _event of response.stream) { + // Consume the response so the instrumented span can finish. + } + + const spans = await backgroundLogger.drain(); + expect(spans).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + output: { + citations: [ + { + generatedResponsePart: { text: "Paris" }, + }, + ], + text: "Paris is in France.", + }, + span_attributes: expect.objectContaining({ + name: "bedrockAgent.retrieveAndGenerateStream", + }), + }), + ]), + ); + }); + + it("reports stream exceptions and closes open child spans", async () => { + async function* completion() { + yield { + trace: { + sessionId: "session-error", + trace: { + orchestrationTrace: { + modelInvocationInput: { + foundationModel: "amazon.nova-lite-v1:0", + text: "model prompt", + traceId: "model-error", + type: "ORCHESTRATION", + }, + }, + }, + }, + }; + yield { + throttlingException: { + message: "rate limited", + }, + }; + } + + const response = + await bedrockAgentSmithyCoreChannels.clientSend.tracePromise( + async () => ({ + completion: completion(), + sessionId: "session-error", + }), + { + arguments: [ + new InvokeAgentCommand({ + agentAliasId: "alias-1", + agentId: "agent-1", + inputText: "Trigger an error.", + sessionId: "session-error", + }) as any, + ], + }, + ); + + for await (const _event of response.completion) { + // Consume the response so the instrumented spans can finish. + } + + const spans = (await backgroundLogger.drain()) as Array< + Record & { + span_attributes?: { name?: string }; + } + >; + const parent = spans.find( + (span) => span.span_attributes?.name === "bedrockAgent.invokeAgent", + ); + const model = spans.find( + (span) => span.span_attributes?.name === "bedrockAgent.modelInvocation", + ); + expect(parent).toMatchObject({ + error: "throttlingException: rate limited", + }); + expect(model).toMatchObject({ + error: "throttlingException: rate limited", + }); + }); + + it("ignores unrelated Agent Runtime commands and callback sends", async () => { + const channel = bedrockAgentSmithyCoreChannels.clientSend.tracingChannel(); + await bedrockAgentSmithyCoreChannels.clientSend.tracePromise( + async () => ({ sessionId: "not-traced" }), + { + arguments: [ + new CreateSessionCommand({ + sessionMetadata: { purpose: "not an AI execution" }, + }) as any, + ], + }, + ); + + const callbackEvent: any = { + arguments: [ + new InvokeAgentCommand({ + agentAliasId: "alias-1", + agentId: "agent-1", + inputText: "ignored", + sessionId: "session-1", + }), + () => {}, + ], + }; + channel.start?.publish(callbackEvent); + callbackEvent.result = { completion: [] }; + channel.asyncEnd?.publish(callbackEvent); + + await expect(backgroundLogger.drain()).resolves.toEqual([]); + }); +}); diff --git a/js/src/instrumentation/plugins/bedrock-agent-runtime-plugin.ts b/js/src/instrumentation/plugins/bedrock-agent-runtime-plugin.ts new file mode 100644 index 000000000..6beafb851 --- /dev/null +++ b/js/src/instrumentation/plugins/bedrock-agent-runtime-plugin.ts @@ -0,0 +1,822 @@ +import { SpanTypeAttribute, isObject } from "../../../util/index"; +import { + startSpan as startBaseSpan, + withCurrent, + type Span, +} from "../../logger"; +import { + INSTRUMENTATION_NAMES, + withSpanInstrumentationName, +} from "../../span-origin"; +import { getCurrentUnixTimestamp } from "../../util"; +import type { + BedrockAgentRuntimeInvokeAgentRequest, + BedrockAgentRuntimeInvokeInlineAgentRequest, + BedrockAgentRuntimeResponse, + BedrockAgentRuntimeRetrieveAndGenerateRequest, + BedrockAgentRuntimeStreamEvent, + BedrockAgentRuntimeTracePart, +} from "../../vendor-sdk-types/bedrock-agent-runtime"; +import { BasePlugin } from "../core"; +import type { AnyAsyncChannel } from "../core/channel-definitions"; +import { traceStreamingChannel, unsubscribeAll } from "../core/channel-tracing"; +import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; +import { + bedrockAgentRuntimeChannels, + bedrockAgentSmithyClientChannels, + bedrockAgentSmithyCoreChannels, +} from "./bedrock-agent-runtime-channels"; +import { + getBedrockAgentRuntimeCommandInput, + getBedrockAgentRuntimeCommandName, + getBedrockAgentRuntimeOperation, + type BedrockAgentRuntimeOperation, +} from "./bedrock-agent-runtime-common"; + +type ChildSpanState = { + span: Span; +}; + +type AgentStreamState = { + citations: unknown[]; + error?: Error; + files: unknown[]; + firstChunkTime?: number; + guardrailAction?: string; + llmSpans: Map; + metrics: Record; + parentSpan: Span; + returnControl?: unknown; + text: string; + toolSpans: Map; +}; + +const STREAM_EXCEPTION_KEYS = [ + "internalServerException", + "validationException", + "resourceNotFoundException", + "serviceQuotaExceededException", + "throttlingException", + "accessDeniedException", + "conflictException", + "dependencyFailedException", + "badGatewayException", + "modelNotReadyException", +] as const; + +export class BedrockAgentRuntimePlugin extends BasePlugin { + protected onEnable(): void { + this.unsubscribers.push( + ...[ + bedrockAgentRuntimeChannels.clientSend, + bedrockAgentSmithyCoreChannels.clientSend, + bedrockAgentSmithyClientChannels.clientSend, + ].map((channel) => traceBedrockAgentRuntimeClientSendChannel(channel)), + ); + } + + protected onDisable(): void { + this.unsubscribers = unsubscribeAll(this.unsubscribers); + } +} + +function traceBedrockAgentRuntimeClientSendChannel( + channel: AnyAsyncChannel, +): () => void { + return traceStreamingChannel(channel, { + name: ([command]) => { + const operation = getBedrockAgentRuntimeOperation(command); + return operation + ? `bedrockAgent.${operation}` + : "bedrockAgent.client.send"; + }, + shouldTrace: ([command, optionsOrCb, cb]) => + getBedrockAgentRuntimeOperation(command) !== undefined && + typeof optionsOrCb !== "function" && + typeof cb !== "function", + type: SpanTypeAttribute.TASK, + extractInput: ([command]) => extractBedrockAgentRuntimeInput(command), + extractOutput: (result, endEvent) => + extractBedrockAgentRuntimeOutput(endEvent?.arguments?.[0], result), + extractMetadata: (result) => + extractBedrockAgentRuntimeResponseMetadata(result), + extractMetrics: () => ({}), + patchResult: ({ endEvent, result, span, startTime }) => + patchBedrockAgentRuntimeStreamingResult({ + command: endEvent.arguments?.[0], + result, + span, + startTime, + }), + }); +} + +function extractBedrockAgentRuntimeInput(command: unknown): { + input: unknown; + metadata: Record; +} { + const operation = getBedrockAgentRuntimeOperation(command); + const commandName = getBedrockAgentRuntimeCommandName(command); + const request = getBedrockAgentRuntimeCommandInput(command); + const metadata = { + provider: "aws-bedrock", + ...(commandName ? { command: commandName } : {}), + ...(operation ? { operation } : {}), + ...extractRequestMetadata(request, operation), + }; + + if (operation === "invokeAgent") { + const invokeRequest = request as + | BedrockAgentRuntimeInvokeAgentRequest + | undefined; + return { + input: + invokeRequest?.inputText ?? + sanitizeBedrockAgentValue(invokeRequest?.sessionState), + metadata, + }; + } + + if (operation === "invokeInlineAgent") { + const inlineRequest = request as + | BedrockAgentRuntimeInvokeInlineAgentRequest + | undefined; + return { + input: + inlineRequest?.inputText ?? + sanitizeBedrockAgentValue(inlineRequest?.inlineSessionState), + metadata, + }; + } + + if ( + operation === "retrieveAndGenerate" || + operation === "retrieveAndGenerateStream" + ) { + return { + input: sanitizeBedrockAgentValue( + (request as BedrockAgentRuntimeRetrieveAndGenerateRequest | undefined) + ?.input, + ), + metadata, + }; + } + + return { + input: sanitizeBedrockAgentValue(request), + metadata, + }; +} + +function extractRequestMetadata( + request: unknown, + operation: BedrockAgentRuntimeOperation | undefined, +): Record { + if (!isObject(request)) { + return {}; + } + + const metadata: Record = {}; + for (const key of [ + "agentId", + "agentAliasId", + "agentName", + "sessionId", + "memoryId", + "enableTrace", + "endSession", + ]) { + if (request[key] !== undefined) { + metadata[key] = sanitizeBedrockAgentValue(request[key]); + } + } + + if (typeof request.foundationModel === "string") { + metadata.model = request.foundationModel; + } + if ( + operation === "invokeInlineAgent" && + Array.isArray(request.actionGroups) + ) { + metadata.tools = sanitizeBedrockAgentValue(request.actionGroups); + } + + const retrieveConfiguration = request.retrieveAndGenerateConfiguration; + if (isObject(retrieveConfiguration)) { + const knowledgeBaseConfiguration = isObject( + retrieveConfiguration.knowledgeBaseConfiguration, + ) + ? retrieveConfiguration.knowledgeBaseConfiguration + : undefined; + const externalSourcesConfiguration = isObject( + retrieveConfiguration.externalSourcesConfiguration, + ) + ? retrieveConfiguration.externalSourcesConfiguration + : undefined; + if (typeof knowledgeBaseConfiguration?.knowledgeBaseId === "string") { + metadata.knowledgeBaseId = knowledgeBaseConfiguration.knowledgeBaseId; + } + const modelArn = + knowledgeBaseConfiguration?.modelArn ?? + externalSourcesConfiguration?.modelArn; + if (typeof modelArn === "string") { + metadata.model = modelArn; + } + } + + return metadata; +} + +function extractBedrockAgentRuntimeOutput( + command: unknown, + result: unknown, +): unknown { + const operation = getBedrockAgentRuntimeOperation(command); + if (!isObject(result)) { + return sanitizeBedrockAgentValue(result); + } + + if (operation === "retrieveAndGenerate") { + return { + text: isObject(result.output) ? result.output.text : undefined, + ...(result.citations !== undefined + ? { citations: sanitizeBedrockAgentValue(result.citations) } + : {}), + }; + } + + return sanitizeBedrockAgentValue(result); +} + +function extractBedrockAgentRuntimeResponseMetadata( + result: unknown, +): Record | undefined { + if (!isObject(result)) { + return undefined; + } + + const metadata: Record = {}; + for (const key of [ + "contentType", + "sessionId", + "memoryId", + "guardrailAction", + ]) { + if (result[key] !== undefined) { + metadata[key] = sanitizeBedrockAgentValue(result[key]); + } + } + return Object.keys(metadata).length > 0 ? metadata : undefined; +} + +function patchBedrockAgentRuntimeStreamingResult(args: { + command: unknown; + result: unknown; + span: Span; + startTime: number; +}): boolean { + const operation = getBedrockAgentRuntimeOperation(args.command); + if (!isObject(args.result) || !operation) { + return false; + } + + const response = args.result as BedrockAgentRuntimeResponse; + const stream = + operation === "retrieveAndGenerateStream" + ? response.stream + : operation === "invokeAgent" || operation === "invokeInlineAgent" + ? response.completion + : undefined; + if (!isAsyncIterable(stream)) { + return false; + } + + const state: AgentStreamState = { + citations: [], + files: [], + llmSpans: new Map(), + metrics: {}, + parentSpan: args.span, + text: "", + toolSpans: new Map(), + }; + + patchStreamIfNeeded(stream, { + onChunk: (chunk) => { + const hasGeneratedText = handleStreamEvent(state, chunk); + if (hasGeneratedText && state.firstChunkTime === undefined) { + state.firstChunkTime = getCurrentUnixTimestamp(); + } + }, + onComplete: () => { + finalizeOpenChildSpans(state, state.error); + const metrics = { ...state.metrics }; + if (state.firstChunkTime !== undefined) { + metrics.time_to_first_token = state.firstChunkTime - args.startTime; + } + + args.span.log({ + ...(state.error + ? { error: state.error } + : { output: buildStreamOutput(state) }), + metrics, + metadata: { + ...extractBedrockAgentRuntimeResponseMetadata(response), + ...(state.guardrailAction !== undefined + ? { guardrailAction: state.guardrailAction } + : {}), + }, + }); + args.span.end(); + }, + onError: (error) => { + finalizeOpenChildSpans(state, error); + args.span.log({ error }); + args.span.end(); + }, + }); + return true; +} + +function handleStreamEvent( + state: AgentStreamState, + event: BedrockAgentRuntimeStreamEvent, +): boolean { + const exception = extractStreamException(event); + if (exception && !state.error) { + state.error = exception; + } + + const chunkText = decodeBytes(event.chunk?.bytes); + if (chunkText !== undefined) { + state.text += chunkText; + } + if (typeof event.output?.text === "string") { + state.text += event.output.text; + } + if (event.chunk?.attribution !== undefined) { + state.citations.push(sanitizeBedrockAgentValue(event.chunk.attribution)); + } + if (event.citation !== undefined) { + state.citations.push(sanitizeBedrockAgentValue(event.citation)); + } + if (typeof event.guardrail?.action === "string") { + state.guardrailAction = event.guardrail.action; + } + if (event.returnControl !== undefined) { + state.returnControl = sanitizeBedrockAgentValue(event.returnControl); + } + if (event.files !== undefined) { + state.files.push(sanitizeBedrockAgentValue(event.files)); + } + if (event.trace) { + handleTracePart(state, event.trace); + } + return chunkText !== undefined || typeof event.output?.text === "string"; +} + +function handleTracePart( + state: AgentStreamState, + tracePart: BedrockAgentRuntimeTracePart, +): void { + const trace = tracePart.trace; + if (!isObject(trace)) { + return; + } + + for (const [traceType, traceValue] of Object.entries(trace)) { + if (!isObject(traceValue)) { + continue; + } + + if (isObject(traceValue.modelInvocationInput)) { + startModelSpan( + state, + traceType, + traceValue.modelInvocationInput, + tracePart, + ); + } + if (isObject(traceValue.modelInvocationOutput)) { + finishModelSpan( + state, + traceType, + traceValue.modelInvocationOutput, + tracePart, + ); + } + if (isObject(traceValue.invocationInput)) { + startToolSpan(state, traceType, traceValue.invocationInput, tracePart); + } + if (isObject(traceValue.observation)) { + finishToolSpan(state, traceType, traceValue.observation, tracePart); + } + } + + if ( + isObject(trace.failureTrace) && + typeof trace.failureTrace.failureReason === "string" && + !state.error + ) { + state.error = new Error(trace.failureTrace.failureReason); + } +} + +function startModelSpan( + state: AgentStreamState, + traceType: string, + input: Record, + tracePart: BedrockAgentRuntimeTracePart, +): void { + const traceId = getTraceId(input); + if (!traceId || state.llmSpans.has(traceId)) { + return; + } + + const metadata: Record = { + provider: "aws-bedrock", + trace_type: traceType, + ...extractTracePartMetadata(tracePart), + }; + if (typeof input.foundationModel === "string") { + metadata.model = input.foundationModel; + } + if (typeof input.type === "string") { + metadata.prompt_type = input.type; + } + if (isObject(input.inferenceConfiguration)) { + Object.assign( + metadata, + sanitizeBedrockAgentValue(input.inferenceConfiguration), + ); + } + + const span = withCurrent(state.parentSpan, () => + startBaseSpan( + withSpanInstrumentationName( + { + event: { + input: input.text, + metadata, + }, + name: "bedrockAgent.modelInvocation", + spanAttributes: { type: SpanTypeAttribute.LLM }, + }, + INSTRUMENTATION_NAMES.BEDROCK_AGENT_RUNTIME, + ), + ), + ); + state.llmSpans.set(traceId, { span }); +} + +function finishModelSpan( + state: AgentStreamState, + traceType: string, + output: Record, + tracePart: BedrockAgentRuntimeTracePart, +): void { + const traceId = getTraceId(output); + if (!traceId) { + return; + } + + let child = state.llmSpans.get(traceId); + if (!child) { + startModelSpan(state, traceType, { traceId }, tracePart); + child = state.llmSpans.get(traceId); + } + if (!child) { + return; + } + + const metrics = extractModelMetrics(output); + addMetrics(state.metrics, metrics); + child.span.log({ + output: extractModelOutput(output), + metrics, + }); + child.span.end(); + state.llmSpans.delete(traceId); +} + +function startToolSpan( + state: AgentStreamState, + traceType: string, + input: Record, + tracePart: BedrockAgentRuntimeTracePart, +): void { + const traceId = getTraceId(input); + if (!traceId || state.toolSpans.has(traceId)) { + return; + } + + const tool = extractToolInput(input); + if (!tool) { + return; + } + const span = withCurrent(state.parentSpan, () => + startBaseSpan( + withSpanInstrumentationName( + { + event: { + input: tool.input, + metadata: { + provider: "aws-bedrock", + tool_type: tool.type, + trace_type: traceType, + ...tool.metadata, + ...extractTracePartMetadata(tracePart), + }, + }, + name: tool.name, + spanAttributes: { type: SpanTypeAttribute.TOOL }, + }, + INSTRUMENTATION_NAMES.BEDROCK_AGENT_RUNTIME, + ), + ), + ); + state.toolSpans.set(traceId, { span }); +} + +function finishToolSpan( + state: AgentStreamState, + _traceType: string, + observation: Record, + _tracePart: BedrockAgentRuntimeTracePart, +): void { + const traceId = getTraceId(observation); + if (!traceId) { + return; + } + + const child = state.toolSpans.get(traceId); + if (!child) { + return; + } + child.span.log({ output: extractToolOutput(observation) }); + child.span.end(); + state.toolSpans.delete(traceId); +} + +function extractToolInput(input: Record): + | { + input: unknown; + metadata: Record; + name: string; + type: string; + } + | undefined { + if (isObject(input.actionGroupInvocationInput)) { + const action = input.actionGroupInvocationInput; + const name = + (typeof action.function === "string" && action.function) || + (typeof action.apiPath === "string" && action.apiPath) || + (typeof action.actionGroupName === "string" && action.actionGroupName) || + "bedrockAgent.actionGroup"; + return { + input: sanitizeBedrockAgentValue({ + parameters: action.parameters, + requestBody: action.requestBody, + }), + metadata: { + ...(typeof action.actionGroupName === "string" + ? { action_group: action.actionGroupName } + : {}), + ...(typeof action.verb === "string" ? { verb: action.verb } : {}), + }, + name, + type: "action_group", + }; + } + + if (isObject(input.knowledgeBaseLookupInput)) { + const lookup = input.knowledgeBaseLookupInput; + return { + input: lookup.text, + metadata: { + ...(typeof lookup.knowledgeBaseId === "string" + ? { knowledge_base_id: lookup.knowledgeBaseId } + : {}), + }, + name: "bedrockAgent.knowledgeBaseLookup", + type: "knowledge_base", + }; + } + + if (isObject(input.codeInterpreterInvocationInput)) { + return { + input: sanitizeBedrockAgentValue(input.codeInterpreterInvocationInput), + metadata: {}, + name: "bedrockAgent.codeInterpreter", + type: "code_interpreter", + }; + } + + if (isObject(input.agentCollaboratorInvocationInput)) { + return { + input: sanitizeBedrockAgentValue(input.agentCollaboratorInvocationInput), + metadata: {}, + name: "bedrockAgent.collaborator", + type: "agent_collaborator", + }; + } + + return undefined; +} + +function extractToolOutput(observation: Record): unknown { + for (const key of [ + "actionGroupInvocationOutput", + "knowledgeBaseLookupOutput", + "codeInterpreterInvocationOutput", + "agentCollaboratorInvocationOutput", + ]) { + if (observation[key] !== undefined) { + return sanitizeBedrockAgentValue(observation[key]); + } + } + return sanitizeBedrockAgentValue(observation); +} + +function extractModelOutput(output: Record): unknown { + const rawResponse = isObject(output.rawResponse) + ? output.rawResponse + : undefined; + const parsedResponse = isObject(output.parsedResponse) + ? output.parsedResponse + : undefined; + const content = + rawResponse?.content ?? + parsedResponse?.text ?? + parsedResponse?.isValid ?? + undefined; + const parsedContent = + typeof content === "string" ? (parseJson(content) ?? content) : content; + + if (output.reasoningContent !== undefined) { + return { + content: sanitizeBedrockAgentValue(parsedContent), + reasoning: sanitizeBedrockAgentValue(output.reasoningContent), + }; + } + return sanitizeBedrockAgentValue(parsedContent ?? output); +} + +function extractModelMetrics( + output: Record, +): Record { + const metadata = isObject(output.metadata) ? output.metadata : undefined; + const usage = isObject(metadata?.usage) ? metadata.usage : undefined; + const metrics: Record = {}; + if (typeof usage?.inputTokens === "number") { + metrics.prompt_tokens = usage.inputTokens; + } + if (typeof usage?.outputTokens === "number") { + metrics.completion_tokens = usage.outputTokens; + } + if ( + metrics.prompt_tokens !== undefined || + metrics.completion_tokens !== undefined + ) { + metrics.tokens = + (metrics.prompt_tokens ?? 0) + (metrics.completion_tokens ?? 0); + } + return metrics; +} + +function addMetrics( + target: Record, + metrics: Record, +): void { + for (const key of ["prompt_tokens", "completion_tokens", "tokens"]) { + const value = metrics[key]; + if (typeof value === "number") { + target[key] = (target[key] ?? 0) + value; + } + } +} + +function extractTracePartMetadata( + tracePart: BedrockAgentRuntimeTracePart, +): Record { + const metadata: Record = {}; + for (const key of [ + "agentId", + "agentAliasId", + "agentVersion", + "collaboratorName", + "sessionId", + ]) { + if (tracePart[key] !== undefined) { + metadata[key] = sanitizeBedrockAgentValue(tracePart[key]); + } + } + return metadata; +} + +function finalizeOpenChildSpans(state: AgentStreamState, error?: Error): void { + for (const child of [ + ...state.llmSpans.values(), + ...state.toolSpans.values(), + ]) { + if (error) { + child.span.log({ error }); + } + child.span.end(); + } + state.llmSpans.clear(); + state.toolSpans.clear(); +} + +function buildStreamOutput(state: AgentStreamState): unknown { + return { + text: state.text, + ...(state.citations.length > 0 + ? { citations: state.citations.slice(0, 50) } + : {}), + ...(state.returnControl !== undefined + ? { returnControl: state.returnControl } + : {}), + ...(state.files.length > 0 ? { files: state.files.slice(0, 20) } : {}), + }; +} + +function extractStreamException( + event: BedrockAgentRuntimeStreamEvent, +): Error | undefined { + for (const key of STREAM_EXCEPTION_KEYS) { + const value = event[key]; + if (!isObject(value)) { + continue; + } + const message = + typeof value.message === "string" + ? value.message + : typeof value.name === "string" + ? value.name + : key; + return new Error(`${key}: ${message}`); + } + return undefined; +} + +function getTraceId(value: Record): string | undefined { + return typeof value.traceId === "string" ? value.traceId : undefined; +} + +function decodeBytes(value: unknown): string | undefined { + if (typeof value === "string") { + return value; + } + if (value instanceof Uint8Array) { + return new TextDecoder().decode(value); + } + if (value instanceof ArrayBuffer) { + return new TextDecoder().decode(new Uint8Array(value)); + } + return undefined; +} + +function parseJson(value: string): unknown | undefined { + try { + return sanitizeBedrockAgentValue(JSON.parse(value)); + } catch { + return undefined; + } +} + +function sanitizeBedrockAgentValue(value: unknown, depth = 0): unknown { + if (value === undefined || value === null) { + return value; + } + if (typeof value !== "object") { + return value; + } + if (value instanceof Date) { + return value.toISOString(); + } + if (value instanceof Uint8Array) { + return { byte_length: value.byteLength }; + } + if (value instanceof ArrayBuffer) { + return { byte_length: value.byteLength }; + } + if (depth > 20) { + return "[MaxDepth]"; + } + if (Array.isArray(value)) { + return value.map((entry) => sanitizeBedrockAgentValue(entry, depth + 1)); + } + if (!isObject(value)) { + return String(value); + } + + const output: Record = {}; + for (const [key, nested] of Object.entries(value)) { + if (key === "__proto__" || key === "prototype" || key === "constructor") { + continue; + } + output[key] = sanitizeBedrockAgentValue(nested, depth + 1); + } + return output; +} diff --git a/js/src/instrumentation/plugins/instrumentation-names.test.ts b/js/src/instrumentation/plugins/instrumentation-names.test.ts index 2e342f9b2..3e7e9c6e1 100644 --- a/js/src/instrumentation/plugins/instrumentation-names.test.ts +++ b/js/src/instrumentation/plugins/instrumentation-names.test.ts @@ -3,6 +3,11 @@ import { describe, expect, it } from "vitest"; import { INSTRUMENTATION_NAMES } from "../../span-origin"; import { aiSDKChannels, harnessAgentChannels } from "./ai-sdk-channels"; import { anthropicChannels } from "./anthropic-channels"; +import { + bedrockAgentRuntimeChannels, + bedrockAgentSmithyClientChannels, + bedrockAgentSmithyCoreChannels, +} from "./bedrock-agent-runtime-channels"; import { bedrockRuntimeChannels, smithyClientChannels, @@ -36,6 +41,18 @@ describe("built-in instrumentation provenance names", () => { [aiSDKChannels.generateText, INSTRUMENTATION_NAMES.AI_SDK], [harnessAgentChannels.generate, INSTRUMENTATION_NAMES.AI_SDK], [anthropicChannels.messagesCreate, INSTRUMENTATION_NAMES.ANTHROPIC], + [ + bedrockAgentRuntimeChannels.clientSend, + INSTRUMENTATION_NAMES.BEDROCK_AGENT_RUNTIME, + ], + [ + bedrockAgentSmithyCoreChannels.clientSend, + INSTRUMENTATION_NAMES.BEDROCK_AGENT_RUNTIME, + ], + [ + bedrockAgentSmithyClientChannels.clientSend, + INSTRUMENTATION_NAMES.BEDROCK_AGENT_RUNTIME, + ], [bedrockRuntimeChannels.clientSend, INSTRUMENTATION_NAMES.BEDROCK_RUNTIME], [smithyCoreChannels.clientSend, INSTRUMENTATION_NAMES.BEDROCK_RUNTIME], [smithyClientChannels.clientSend, INSTRUMENTATION_NAMES.BEDROCK_RUNTIME], diff --git a/js/src/span-origin.ts b/js/src/span-origin.ts index 0fd2f5fc8..a4e4b0539 100644 --- a/js/src/span-origin.ts +++ b/js/src/span-origin.ts @@ -10,6 +10,7 @@ export type SpanOriginEnvironment = { export const INSTRUMENTATION_NAMES = { AI_SDK: "ai-sdk", ANTHROPIC: "anthropic", + BEDROCK_AGENT_RUNTIME: "bedrock-agent-runtime", BEDROCK_RUNTIME: "bedrock-runtime", BRAINTRUST_JS_LOGGER: "braintrust-js-logger", CLAUDE_AGENT_SDK: "claude-agent-sdk", diff --git a/js/src/vendor-sdk-types/bedrock-agent-runtime.ts b/js/src/vendor-sdk-types/bedrock-agent-runtime.ts new file mode 100644 index 000000000..d9e089768 --- /dev/null +++ b/js/src/vendor-sdk-types/bedrock-agent-runtime.ts @@ -0,0 +1,124 @@ +import type { ChannelSpanInfo } from "../instrumentation/core/types"; + +export type BedrockAgentRuntimeCommandName = + | "InvokeAgentCommand" + | "InvokeInlineAgentCommand" + | "RetrieveAndGenerateCommand" + | "RetrieveAndGenerateStreamCommand"; + +export interface BedrockAgentRuntimeCommandLike { + input?: unknown; + constructor?: { + name?: string; + }; + [key: string]: unknown; +} + +export interface BedrockAgentRuntimeClient { + send: ( + command: BedrockAgentRuntimeCommandLike, + optionsOrCb?: unknown, + cb?: unknown, + ) => Promise | unknown; + [key: string]: unknown; +} + +export interface BedrockAgentRuntimeChannelContext { + span_info?: ChannelSpanInfo; +} + +export interface BedrockAgentRuntimeInvokeAgentRequest { + agentId?: string; + agentAliasId?: string; + sessionId?: string; + inputText?: string; + enableTrace?: boolean; + endSession?: boolean; + memoryId?: string; + streamingConfigurations?: unknown; + sessionState?: unknown; + [key: string]: unknown; +} + +export interface BedrockAgentRuntimeInvokeInlineAgentRequest { + foundationModel?: string; + instruction?: string; + agentName?: string; + sessionId?: string; + inputText?: string; + enableTrace?: boolean; + endSession?: boolean; + streamingConfigurations?: unknown; + inlineSessionState?: unknown; + [key: string]: unknown; +} + +export interface BedrockAgentRuntimeRetrieveAndGenerateRequest { + input?: { + text?: string; + [key: string]: unknown; + }; + sessionId?: string; + retrieveAndGenerateConfiguration?: unknown; + [key: string]: unknown; +} + +export interface BedrockAgentRuntimeTracePart { + sessionId?: string; + trace?: Record; + agentId?: string; + agentAliasId?: string; + agentVersion?: string; + collaboratorName?: string; + [key: string]: unknown; +} + +export interface BedrockAgentRuntimeStreamEvent { + chunk?: { + bytes?: unknown; + attribution?: unknown; + [key: string]: unknown; + }; + trace?: BedrockAgentRuntimeTracePart; + returnControl?: unknown; + files?: unknown; + output?: { + text?: string; + [key: string]: unknown; + }; + citation?: unknown; + guardrail?: { + action?: string; + [key: string]: unknown; + }; + internalServerException?: unknown; + validationException?: unknown; + resourceNotFoundException?: unknown; + serviceQuotaExceededException?: unknown; + throttlingException?: unknown; + accessDeniedException?: unknown; + conflictException?: unknown; + dependencyFailedException?: unknown; + badGatewayException?: unknown; + modelNotReadyException?: unknown; + [key: string]: unknown; +} + +export interface BedrockAgentRuntimeResponse { + completion?: AsyncIterable; + stream?: AsyncIterable; + contentType?: string; + sessionId?: string; + memoryId?: string; + output?: { + text?: string; + [key: string]: unknown; + }; + citations?: unknown; + guardrailAction?: string; + [key: string]: unknown; +} + +export type BedrockAgentRuntimeSendResult = + | BedrockAgentRuntimeResponse + | unknown; diff --git a/js/src/wrappers/bedrock-agent-runtime.test.ts b/js/src/wrappers/bedrock-agent-runtime.test.ts new file mode 100644 index 000000000..ce0810187 --- /dev/null +++ b/js/src/wrappers/bedrock-agent-runtime.test.ts @@ -0,0 +1,134 @@ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from "vitest"; +import { _exportsForTestingOnly, initLogger } from "../logger"; +import { configureNode } from "../node/config"; +import { wrapBedrockAgentRuntime } from "./bedrock-agent-runtime"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +class InvokeInlineAgentCommand { + constructor(public input: Record) {} +} + +describe("wrapBedrockAgentRuntime", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectName: "bedrock-agent-runtime.test.ts", + projectId: "test-project-id", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + it("returns unsupported clients unchanged", () => { + const value = { request: vi.fn() }; + expect(wrapBedrockAgentRuntime(value)).toBe(value); + }); + + it("is idempotent and preserves client methods", async () => { + async function* completion() { + yield { + chunk: { + bytes: new TextEncoder().encode("OK"), + }, + }; + } + const send = vi.fn( + async ( + _command: InvokeInlineAgentCommand, + _optionsOrCb?: unknown, + _cb?: unknown, + ) => ({ + completion: completion(), + sessionId: "inline-session", + }), + ); + const client = { + send, + invokeInlineAgent( + this: { send: typeof send }, + input: Record, + ) { + return this.send(new InvokeInlineAgentCommand(input)); + }, + }; + + const wrapped = wrapBedrockAgentRuntime(client); + expect(wrapBedrockAgentRuntime(client)).toBe(wrapped); + + const response = await wrapped.invokeInlineAgent({ + foundationModel: "amazon.nova-lite-v1:0", + inputText: "Reply with OK.", + instruction: "Reply to the user concisely and accurately.", + sessionId: "inline-session", + }); + for await (const _event of response.completion) { + // Consume the response so the instrumented span can finish. + } + + expect(send).toHaveBeenCalledTimes(1); + const spans = await backgroundLogger.drain(); + expect(spans).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + input: "Reply with OK.", + metadata: expect.objectContaining({ + command: "InvokeInlineAgentCommand", + model: "amazon.nova-lite-v1:0", + operation: "invokeInlineAgent", + provider: "aws-bedrock", + }), + output: { text: "OK" }, + span_attributes: expect.objectContaining({ + name: "bedrockAgent.invokeInlineAgent", + type: "task", + }), + }), + ]), + ); + }); + + it("passes callback overloads through without tracing", async () => { + const callback = vi.fn(); + const send = vi.fn( + ( + _command: InvokeInlineAgentCommand, + _optionsOrCb?: unknown, + _cb?: unknown, + ) => undefined, + ); + const wrapped = wrapBedrockAgentRuntime({ send }); + const command = new InvokeInlineAgentCommand({ + foundationModel: "amazon.nova-lite-v1:0", + instruction: "Reply to the user concisely and accurately.", + sessionId: "inline-session", + }); + + wrapped.send(command, callback); + + expect(send).toHaveBeenCalledWith(command, callback, undefined); + await expect(backgroundLogger.drain()).resolves.toEqual([]); + }); +}); diff --git a/js/src/wrappers/bedrock-agent-runtime.ts b/js/src/wrappers/bedrock-agent-runtime.ts new file mode 100644 index 000000000..65512509e --- /dev/null +++ b/js/src/wrappers/bedrock-agent-runtime.ts @@ -0,0 +1,142 @@ +import { runWithAutoInstrumentationSuppressed } from "../instrumentation/auto-instrumentation-suppression"; +import { bedrockAgentRuntimeChannels } from "../instrumentation/plugins/bedrock-agent-runtime-channels"; +import { + buildBedrockAgentRuntimeSpanInfo, + getBedrockAgentRuntimeOperation, +} from "../instrumentation/plugins/bedrock-agent-runtime-common"; +import type { + BedrockAgentRuntimeClient, + BedrockAgentRuntimeCommandLike, +} from "../vendor-sdk-types/bedrock-agent-runtime"; + +/** + * Wrap an AWS Bedrock Agent Runtime client with Braintrust tracing. + */ +export function wrapBedrockAgentRuntime(client: T): T { + if (isSupportedBedrockAgentRuntimeClient(client)) { + return bedrockAgentRuntimeProxy(client) as T; + } + + // eslint-disable-next-line no-restricted-properties -- preserving wrapper warning behavior. + console.warn("Unsupported Bedrock Agent Runtime library. Not wrapping."); + return client; +} + +const proxyCache = new WeakMap< + BedrockAgentRuntimeClient, + BedrockAgentRuntimeClient +>(); + +const OPERATION_METHODS = new Set([ + "invokeAgent", + "invokeInlineAgent", + "retrieveAndGenerate", + "retrieveAndGenerateStream", +]); + +function isSupportedBedrockAgentRuntimeClient( + value: unknown, +): value is BedrockAgentRuntimeClient { + return ( + typeof value === "object" && + value !== null && + typeof (value as Record).send === "function" + ); +} + +function bedrockAgentRuntimeProxy( + client: BedrockAgentRuntimeClient, +): BedrockAgentRuntimeClient { + const cached = proxyCache.get(client); + if (cached) { + return cached; + } + + const privateMethodCache = new WeakMap< + (...args: unknown[]) => unknown, + (...args: unknown[]) => unknown + >(); + const operationMethodCache = new WeakMap< + (...args: unknown[]) => unknown, + (...args: unknown[]) => unknown + >(); + + const proxy: BedrockAgentRuntimeClient = new Proxy(client, { + get(target, prop, receiver) { + if (prop === "send") { + return wrapSend(target.send.bind(target)); + } + + const value = Reflect.get(target, prop, receiver); + if (typeof value !== "function") { + return value; + } + + if (typeof prop === "string" && OPERATION_METHODS.has(prop)) { + const cachedValue = operationMethodCache.get(value); + if (cachedValue) { + return cachedValue; + } + const thisBoundValue = function ( + this: unknown, + ...args: unknown[] + ): unknown { + const output = Reflect.apply( + value, + this === proxy ? proxy : this, + args, + ); + return output === target ? proxy : output; + }; + operationMethodCache.set(value, thisBoundValue); + return thisBoundValue; + } + + const cachedValue = privateMethodCache.get(value); + if (cachedValue) { + return cachedValue; + } + const thisBoundValue = function ( + this: unknown, + ...args: unknown[] + ): unknown { + const output = Reflect.apply( + value, + this === proxy ? target : this, + args, + ); + return output === target ? proxy : output; + }; + privateMethodCache.set(value, thisBoundValue); + return thisBoundValue; + }, + }); + + proxyCache.set(client, proxy); + return proxy; +} + +function wrapSend( + send: BedrockAgentRuntimeClient["send"], +): BedrockAgentRuntimeClient["send"] { + return (command, optionsOrCb, cb) => { + if ( + getBedrockAgentRuntimeOperation(command) === undefined || + typeof optionsOrCb === "function" || + typeof cb === "function" + ) { + return send(command, optionsOrCb, cb); + } + + return bedrockAgentRuntimeChannels.clientSend.tracePromise( + () => + runWithAutoInstrumentationSuppressed(() => + send(command, optionsOrCb), + ) as Promise, + { + arguments: [command as BedrockAgentRuntimeCommandLike, optionsOrCb], + span_info: buildBedrockAgentRuntimeSpanInfo(command), + }, + ); + }; +}