Skip to content

Commit aeab196

Browse files
xrmxxrmx
committed
Add tracing for ConverseStream
1 parent 2756c1e commit aeab196

File tree

8 files changed

+367
-8
lines changed

8 files changed

+367
-8
lines changed

‎instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,11 +188,15 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
188188
}
189189

190190
_safe_invoke(extension.extract_attributes, attributes)
191+
end_span_on_exit = extension.should_end_span_on_exit()
191192

192193
with self._tracer.start_as_current_span(
193194
call_context.span_name,
194195
kind=call_context.span_kind,
195196
attributes=attributes,
197+
# tracing streaming services require to close the span manually
198+
# at a later time after the stream has been consumed
199+
end_on_exit=end_span_on_exit,
196200
) as span:
197201
_safe_invoke(extension.before_service_call, span)
198202
self._call_request_hook(span, call_context)

‎instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py‎

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,12 @@
2323
import logging
2424
from typing import Any
2525

26+
from botocore.eventstream import EventStream
2627
from botocore.response import StreamingBody
2728

29+
from opentelemetry.instrumentation.botocore.extensions.bedrock_utils import (
30+
ConverseStreamWrapper,
31+
)
2832
from opentelemetry.instrumentation.botocore.extensions.types import (
2933
_AttributeMapT,
3034
_AwsSdkExtension,
@@ -62,7 +66,14 @@ class _BedrockRuntimeExtension(_AwsSdkExtension):
6266
Amazon Bedrock Runtime</a>.
6367
"""
6468

65-
_HANDLED_OPERATIONS = {"Converse", "InvokeModel"}
69+
_HANDLED_OPERATIONS = {"Converse", "ConverseStream", "InvokeModel"}
70+
_DONT_CLOSE_SPAN_ON_END_OPERATIONS = {"ConverseStream"}
71+
72+
def should_end_span_on_exit(self):
73+
return (
74+
self._call_context.operation
75+
not in self._DONT_CLOSE_SPAN_ON_END_OPERATIONS
76+
)
6677

6778
def extract_attributes(self, attributes: _AttributeMapT):
6879
if self._call_context.operation not in self._HANDLED_OPERATIONS:
@@ -77,7 +88,7 @@ def extract_attributes(self, attributes: _AttributeMapT):
7788
GenAiOperationNameValues.CHAT.value
7889
)
7990

80-
# Converse
91+
# Converse / ConverseStream
8192
if inference_config := self._call_context.params.get(
8293
"inferenceConfig"
8394
):
@@ -251,10 +262,25 @@ def on_success(self, span: Span, result: dict[str, Any]):
251262
return
252263

253264
if not span.is_recording():
265+
if not self.should_end_span_on_exit():
266+
span.end()
254267
return
255268

256-
# Converse
257-
self._converse_on_success(span, result)
269+
# ConverseStream
270+
if "stream" in result and isinstance(result["stream"], EventStream):
271+
272+
def stream_done_callback(response):
273+
self._converse_on_success(span, response)
274+
print("span.attr", span.attributes)
275+
span.end()
276+
277+
result["stream"] = ConverseStreamWrapper(
278+
result["stream"], stream_done_callback
279+
)
280+
return
281+
else:
282+
# Converse
283+
self._converse_on_success(span, result)
258284

259285
model_id = self._call_context.params.get(_MODEL_ID_KEY)
260286
if not model_id:
@@ -328,3 +354,6 @@ def on_error(self, span: Span, exception: _BotoClientErrorT):
328354
span.set_status(Status(StatusCode.ERROR, str(exception)))
329355
if span.is_recording():
330356
span.set_attribute(ERROR_TYPE, type(exception).__qualname__)
357+
358+
if not self.should_end_span_on_exit():
359+
span.end()

‎instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock_utils.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# Includes work from:
16+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
17+
# SPDX-License-Identifier: Apache-2.0
18+
19+
from __future__ import annotations
20+
21+
from botocore.eventstream import EventStream
22+
from wrapt import ObjectProxy
23+
24+
25+
class ConverseStreamWrapper(ObjectProxy):
26+
"""Wrapper for botocor.event.EventStream"""
27+
28+
def __init__(
29+
self,
30+
stream: EventStream,
31+
stream_done_callback,
32+
):
33+
super().__init__(stream)
34+
35+
self._stream_done_callback = stream_done_callback
36+
# accumulating things in the same shape of non-streaming version
37+
# {"usage": {"inputTokens": 0, "outputTokens": 0}, "stopReason": "finish"}
38+
self._response = {}
39+
40+
def __iter__(self):
41+
for event in self.__wrapped__:
42+
self._process_event(event)
43+
yield event
44+
45+
def _process_event(self, event):
46+
if "messageStart" in event:
47+
# {'messageStart': {'role': 'assistant'}}
48+
pass
49+
50+
if "contentBlockDelta" in event:
51+
# {'contentBlockDelta': {'delta': {'text': "Hello"}, 'contentBlockIndex': 0}}
52+
pass
53+
54+
if "contentBlockStop" in event:
55+
# {'contentBlockStop': {'contentBlockIndex': 0}}
56+
pass
57+
58+
if "messageStop" in event:
59+
# {'messageStop': {'stopReason': 'end_turn'}}
60+
if stop_reason := event["messageStop"].get("stopReason"):
61+
self._response["stopReason"] = stop_reason
62+
63+
if "metadata" in event:
64+
# {'metadata': {'usage': {'inputTokens': 12, 'outputTokens': 15, 'totalTokens': 27}, 'metrics': {'latencyMs': 2980}}}
65+
if usage := event["metadata"].get("usage"):
66+
self._response["usage"] = {}
67+
if input_tokens := usage.get("inputTokens"):
68+
self._response["usage"]["inputTokens"] = input_tokens
69+
70+
if output_tokens := usage.get("outputTokens"):
71+
self._response["usage"]["outputTokens"] = output_tokens
72+
73+
self._stream_done_callback(self._response)

‎instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/types.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,14 @@ def should_trace_service_call(self) -> bool: # pylint:disable=no-self-use
101101
"""
102102
return True
103103

104+
def should_end_span_on_exit(self) -> bool: # pylint:disable=no-self-use
105+
"""Returns if the span should be close automatically on exit
106+
107+
Extensions might override this function to disable automatic closing
108+
of the span if they need to close it at a later time themselves.
109+
"""
110+
return True
111+
104112
def extract_attributes(self, attributes: _AttributeMapT):
105113
"""Callback which gets invoked before the span is created.
106114

‎instrumentation/opentelemetry-instrumentation-botocore/tests/bedrock_utils.py‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ def assert_completion_attributes_from_streaming_body(
9191
)
9292

9393

94-
def assert_completion_attributes(
94+
def assert_converse_completion_attributes(
9595
span: ReadableSpan,
9696
request_model: str,
9797
response: dict[str, Any] | None,
@@ -128,6 +128,34 @@ def assert_completion_attributes(
128128
)
129129

130130

131+
def assert_converse_stream_completion_attributes(
132+
span: ReadableSpan,
133+
request_model: str,
134+
input_tokens: int | None = None,
135+
output_tokens: int | None = None,
136+
finish_reason: tuple[str] | None = None,
137+
operation_name: str = "chat",
138+
request_top_p: int | None = None,
139+
request_temperature: int | None = None,
140+
request_max_tokens: int | None = None,
141+
request_stop_sequences: list[str] | None = None,
142+
):
143+
return assert_all_attributes(
144+
span,
145+
request_model,
146+
input_tokens,
147+
output_tokens,
148+
finish_reason,
149+
operation_name,
150+
request_top_p,
151+
request_temperature,
152+
request_max_tokens,
153+
tuple(request_stop_sequences)
154+
if request_stop_sequences is not None
155+
else request_stop_sequences,
156+
)
157+
158+
131159
def assert_equal_or_not_present(value, attribute_name, span):
132160
if value is not None:
133161
assert value == span.attributes[attribute_name]

‎instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_converse_stream_with_content.yaml‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
interactions:
2+
- request:
3+
body: '{"messages": [{"role": "user", "content": [{"text": "Say this is a test"}]}],
4+
"inferenceConfig": {"maxTokens": 10, "temperature": 0.8, "topP": 1, "stopSequences":
5+
["|"]}}'
6+
headers:
7+
Content-Length:
8+
- '170'
9+
Content-Type:
10+
- !!binary |
11+
YXBwbGljYXRpb24vanNvbg==
12+
User-Agent:
13+
- !!binary |
14+
Qm90bzMvMS4zNS41NiBtZC9Cb3RvY29yZSMxLjM1LjU2IHVhLzIuMCBvcy9saW51eCM2LjEuMC0x
15+
MDM0LW9lbSBtZC9hcmNoI3g4Nl82NCBsYW5nL3B5dGhvbiMzLjEwLjEyIG1kL3B5aW1wbCNDUHl0
16+
aG9uIGNmZy9yZXRyeS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM1LjU2
17+
X-Amz-Date:
18+
- !!binary |
19+
MjAyNTAxMjNUMDk1MTU2Wg==
20+
X-Amz-Security-Token:
21+
- test_aws_security_token
22+
X-Amzn-Trace-Id:
23+
- !!binary |
24+
Um9vdD0xLTA0YmY4MjVjLTAxMTY5NjdhYWM1NmIxM2RlMDI1N2QwMjtQYXJlbnQ9MDdkM2U3N2Rl
25+
OGFjMzJhNDtTYW1wbGVkPTE=
26+
amz-sdk-invocation-id:
27+
- !!binary |
28+
ZGQ1MTZiNTEtOGU1Yi00NGYyLTk5MzMtZjAwYzBiOGFkYWYw
29+
amz-sdk-request:
30+
- !!binary |
31+
YXR0ZW1wdD0x
32+
authorization:
33+
- Bearer test_aws_authorization
34+
method: POST
35+
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-lite-v1/converse-stream
36+
response:
37+
body:
38+
string: !!binary |
39+
AAAAlAAAAFLEwW5hCzpldmVudC10eXBlBwAMbWVzc2FnZVN0YXJ0DTpjb250ZW50LXR5cGUHABBh
40+
cHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVldmVudHsicCI6ImFiY2RlZmdoaWprbG1u
41+
b3BxcnN0dXZ3Iiwicm9sZSI6ImFzc2lzdGFudCJ9P+wfRAAAAMQAAABXjLhVJQs6ZXZlbnQtdHlw
42+
ZQcAEWNvbnRlbnRCbG9ja0RlbHRhDTpjb250ZW50LXR5cGUHABBhcHBsaWNhdGlvbi9qc29uDTpt
43+
ZXNzYWdlLXR5cGUHAAVldmVudHsiY29udGVudEJsb2NrSW5kZXgiOjAsImRlbHRhIjp7InRleHQi
44+
OiJIaSEgSG93IGNhbiBJIGhlbHAgeW91In0sInAiOiJhYmNkZWZnaGlqa2xtbm9wcXJzdHUifeBJ
45+
9mIAAACJAAAAVlvc+UsLOmV2ZW50LXR5cGUHABBjb250ZW50QmxvY2tTdG9wDTpjb250ZW50LXR5
46+
cGUHABBhcHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVldmVudHsiY29udGVudEJsb2Nr
47+
SW5kZXgiOjAsInAiOiJhYmNkZSJ95xzwrwAAAKcAAABRu0n9jQs6ZXZlbnQtdHlwZQcAC21lc3Nh
48+
Z2VTdG9wDTpjb250ZW50LXR5cGUHABBhcHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVl
49+
dmVudHsicCI6ImFiY2RlZmdoaWprbG1ub3BxcnN0dXZ3eHl6QUJDREVGR0hJSiIsInN0b3BSZWFz
50+
b24iOiJtYXhfdG9rZW5zIn1LR3pNAAAAygAAAE5X40OECzpldmVudC10eXBlBwAIbWV0YWRhdGEN
51+
OmNvbnRlbnQtdHlwZQcAEGFwcGxpY2F0aW9uL2pzb24NOm1lc3NhZ2UtdHlwZQcABWV2ZW50eyJt
52+
ZXRyaWNzIjp7ImxhdGVuY3lNcyI6NjA4fSwicCI6ImFiY2RlZmdoaWprIiwidXNhZ2UiOnsiaW5w
53+
dXRUb2tlbnMiOjgsIm91dHB1dFRva2VucyI6MTAsInRvdGFsVG9rZW5zIjoxOH19iiQr+w==
54+
headers:
55+
Connection:
56+
- keep-alive
57+
Content-Type:
58+
- application/vnd.amazon.eventstream
59+
Date:
60+
- Thu, 23 Jan 2025 09:51:56 GMT
61+
Set-Cookie: test_set_cookie
62+
Transfer-Encoding:
63+
- chunked
64+
x-amzn-RequestId:
65+
- 2b74a5d3-615a-4f81-b00f-f0b10a618e23
66+
status:
67+
code: 200
68+
message: OK
69+
version: 1

‎instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_converse_stream_with_invalid_model.yaml‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
interactions:
2+
- request:
3+
body: '{"messages": [{"role": "user", "content": [{"text": "Say this is a test"}]}]}'
4+
headers:
5+
Content-Length:
6+
- '77'
7+
Content-Type:
8+
- !!binary |
9+
YXBwbGljYXRpb24vanNvbg==
10+
User-Agent:
11+
- !!binary |
12+
Qm90bzMvMS4zNS41NiBtZC9Cb3RvY29yZSMxLjM1LjU2IHVhLzIuMCBvcy9saW51eCM2LjEuMC0x
13+
MDM0LW9lbSBtZC9hcmNoI3g4Nl82NCBsYW5nL3B5dGhvbiMzLjEwLjEyIG1kL3B5aW1wbCNDUHl0
14+
aG9uIGNmZy9yZXRyeS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM1LjU2
15+
X-Amz-Date:
16+
- !!binary |
17+
MjAyNTAxMjNUMDk1MTU3Wg==
18+
X-Amz-Security-Token:
19+
- test_aws_security_token
20+
X-Amzn-Trace-Id:
21+
- !!binary |
22+
Um9vdD0xLTI5NzA1OTZhLTEyZWI5NDk2ODA1ZjZhYzE5YmU3ODM2NztQYXJlbnQ9Y2M0OTA0YWE2
23+
ZjQ2NmYxYTtTYW1wbGVkPTE=
24+
amz-sdk-invocation-id:
25+
- !!binary |
26+
MjQzZWY2ZDgtNGJhNy00YTVlLWI0MGEtYThiNDE2ZDIzYjhk
27+
amz-sdk-request:
28+
- !!binary |
29+
YXR0ZW1wdD0x
30+
authorization:
31+
- Bearer test_aws_authorization
32+
method: POST
33+
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/does-not-exist/converse-stream
34+
response:
35+
body:
36+
string: '{"message":"The provided model identifier is invalid."}'
37+
headers:
38+
Connection:
39+
- keep-alive
40+
Content-Length:
41+
- '55'
42+
Content-Type:
43+
- application/json
44+
Date:
45+
- Thu, 23 Jan 2025 09:51:57 GMT
46+
Set-Cookie: test_set_cookie
47+
x-amzn-ErrorType:
48+
- ValidationException:http://internal.amazon.com/coral/com.amazon.bedrock/
49+
x-amzn-RequestId:
50+
- 358b122c-d045-4d8f-a5bb-b0bd8cf6ee59
51+
status:
52+
code: 400
53+
message: Bad Request
54+
version: 1

0 commit comments

Comments
 (0)