botocore: add basic tracing for bedrock InvokeAgent #3404
Conversation
stanek-michal
force-pushed
the
invoke-agent
branch
from
3ca64a2 to
a193a11
Compare
instrumentation/opentelemetry-instrumentation-botocore/test-requirements-0.txt
Outdated
Show resolved
Hide resolved
instrumentation/opentelemetry-instrumentation-botocore/test-requirements-1.txt
Outdated
Show resolved
Hide resolved
instrumentation/opentelemetry-instrumentation-botocore/tests/bedrock_utils.py
Outdated
Show resolved
Hide resolved
| def test_invoke_agent_with_completion( | ||
| span_exporter, | ||
| metric_reader, | ||
| bedrock_runtime_client, |
There was a problem hiding this comment.
Do you still need the bedrock runtime client?
There was a problem hiding this comment.
no, though the create_agent scaffolding will need it still (for s3 and such) but it will be done elsewhere
instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_bedrock.py
Outdated
Show resolved
Hide resolved
...y-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/__init__.py
Outdated
Show resolved
Hide resolved
| self.client = boto3.client('bedrock-agent-runtime') | ||
|
|
||
| # Replace these with your actual agent IDs | ||
| self.agent_id = "TFMZVIWXR7" |
There was a problem hiding this comment.
changed it on both standalone demo and tests
(so tests now need export TOX_OVERRIDE=testenv.pass_env=AWS_ACCESS_KEY_ID,AWS_SECRET_ACCESS_KEY,AWS_SESSION_TOKEN,AWS_DEFAULT_REGION,BEDROCK_AGENT_ID,BEDROCK_AGENT_ALIAS_ID
|
@stanek-michal please sign the CLA! |
| event_stream = result["completion"] | ||
|
|
||
| # Drain the stream so we can instrument AND keep events | ||
| all_events = list(event_stream) |
There was a problem hiding this comment.
Is there a reason we don't want to follow the same approach as we do with InvokeModelWithResponseStream?
It looks like we use a response stream wrapper that is an iterator.
I wonder what memory implications draining the stream all at once would have?
There was a problem hiding this comment.
I thought it was simpler but you're right, apart of memory use I'm worried about introducing additional latency. For enableTrace=False there is just one event in my tests, but for enabled tracing there can be quite a few of them, so it matters. I'll change it to a wrapper but not sure if in this PR or next one with create agent
| all_events = list(event_stream) | ||
|
|
||
| # A replay generator so user code can still iterate | ||
| result["completion"] = _replay_events(all_events) |
There was a problem hiding this comment.
if we use a wrapper that's an iterator, as mentioned above, I'm assuming we wouldn't need this line
...ry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py
Outdated
Show resolved
Hide resolved
stanek-michal
force-pushed
the
invoke-agent
branch
from
a193a11 to
f4e876d
Compare
stanek-michal
force-pushed
the
invoke-agent
branch
from
f4e876d to
4c4c419
Compare
| span.update_name(f"{operation_name} {request_model}") | ||
| if self._call_context.operation == "InvokeAgent": | ||
| if operation_name: | ||
| span.update_name(f"{operation_name}") |
There was a problem hiding this comment.
The **span name** SHOULD be `invoke_agent {gen_ai.agent.name}` if `gen_ai.agent.name` is readily available.
Do we have the agent name at hand?
There was a problem hiding this comment.
It looks like InvokeAgent doesn't specify the agent name in request or response, just agentAlias. Seems like a mandatory step to set up an alias in order to call InvokeAgent API (according to my understanding of 1 and 2):
To deploy your agent, you must create an alias. During alias creation, Amazon Bedrock creates a version of your agent automatically. The alias points to this newly created version. Alternatively, you can point the alias to a previously created version of your agent. Then, you configure your application to make API calls to that alias.
Since it's a required request param, I think we can use it here.
| api_input = input_item["apiInvocationInput"] | ||
| action_group = api_input.get("actionGroup") | ||
| span.set_attribute(GEN_AI_TOOL_NAME, action_group) | ||
| span.set_attribute(GEN_AI_TOOL_TYPE, "extension") |
There was a problem hiding this comment.
Is this called directly by the agent though? From the examples here it looks like these are still called from client side https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_InvokeAgent.html#API_agent-runtime_InvokeAgent_ResponseSyntax
Semconv:
Extension: A tool executed on the agent-side to directly call external APIs, bridging the gap between the agent and real-world systems. Agent-side operations involve actions that are performed by the agent on the server or within the agent's controlled environment.
There was a problem hiding this comment.
apiInvocation is untested for this PR but that's correct - if there is returnControl that means client executes the tool, not the Agent, so in this case it should be function not extension (but overall Agent has the capability to call lambdas on its own if it's configured for it)
| span.update_name(f"{operation_name} {request_model}") | ||
| if self._call_context.operation == "InvokeAgent": | ||
| if operation_name: | ||
| span.update_name(f"{operation_name}") |
There was a problem hiding this comment.
It looks like InvokeAgent doesn't specify the agent name in request or response, just agentAlias. Seems like a mandatory step to set up an alias in order to call InvokeAgent API (according to my understanding of 1 and 2):
To deploy your agent, you must create an alias. During alias creation, Amazon Bedrock creates a version of your agent automatically. The alias points to this newly created version. Alternatively, you can point the alias to a previously created version of your agent. Then, you configure your application to make API calls to that alias.
Since it's a required request param, I think we can use it here.
| @@ -32,6 +32,9 @@ def loader(): | |||
|
|
|||
|
|
|||
| _KNOWN_EXTENSIONS = { | |||
| "bedrock-agent-runtime": _lazy_load( | |||
There was a problem hiding this comment.
Proposal that I don't think is needed for this PR, but more of a future suggestion (and up for discussion):
It might be worth building a separate extension for bedrock-agent-runtime. So far, there seems to be a separate extension for every distinct botocore client (dynamodb, lambda, etc.) and bedrock-runtime is separate from bedrock-agent-runtime. Perhaps it's worth following that pattern and avoid some of the branching logic that's likely to grow as more Bedrock APIs get added.
Furthermore, to introduce agent creation/deletion means adding bedrock-agent client as well, which may again result in another extension or combining it with bedrock-agent-runtime and having one dedicated just to agents.
Description
Add basic support for InvokeAgent API. For now only returnControl function invocation, non-streaming mode. Assumes the agent was created beforehand in Bedrock platform - proper Agent creation/deletion for testcases will be added in another PR.
Type of change
Please delete options that are not relevant.
How Has This Been Tested?
Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration
Does This PR Require a Core Repo Change?
Checklist:
See contributing.md for styleguide, changelog guidelines, and more.