Skip to content

Commit 2f022cc

Browse files
SrdjanLLSrdjanLL
committed
Ensure spans end on early stream closure for Bedrock Streaming APIs
1 parent 9bc7764 commit 2f022cc

7 files changed

+528
-14
lines changed

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

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -499,18 +499,20 @@ def _converse_on_success(
499499
[stop_reason],
500500
)
501501

502-
event_logger = instrumentor_context.event_logger
503-
choice = _Choice.from_converse(result, capture_content)
504-
# this path is used by streaming apis, in that case we are already out of the span
505-
# context so need to add the span context manually
506-
span_ctx = span.get_span_context()
507-
event_logger.emit(
508-
choice.to_choice_event(
509-
trace_id=span_ctx.trace_id,
510-
span_id=span_ctx.span_id,
511-
trace_flags=span_ctx.trace_flags,
502+
# In case of an early stream closure, the result may not contain outputs
503+
if "output" in result and "message" in result["output"]:
504+
event_logger = instrumentor_context.event_logger
505+
choice = _Choice.from_converse(result, capture_content)
506+
# this path is used by streaming apis, in that case we are already out of the span
507+
# context so need to add the span context manually
508+
span_ctx = span.get_span_context()
509+
event_logger.emit(
510+
choice.to_choice_event(
511+
trace_id=span_ctx.trace_id,
512+
span_id=span_ctx.span_id,
513+
trace_flags=span_ctx.trace_flags,
514+
)
512515
)
513-
)
514516

515517
metrics = instrumentor_context.metrics
516518
metrics_attributes = self._extract_metrics_attributes()
@@ -781,9 +783,11 @@ def _handle_amazon_nova_response(
781783
GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stopReason"]]
782784
)
783785

784-
event_logger = instrumentor_context.event_logger
785-
choice = _Choice.from_converse(response_body, capture_content)
786-
event_logger.emit(choice.to_choice_event())
786+
# In case of an early stream closure, the result may not contain outputs
787+
if "output" in response_body and "message" in response_body["output"]:
788+
event_logger = instrumentor_context.event_logger
789+
choice = _Choice.from_converse(response_body, capture_content)
790+
event_logger.emit(choice.to_choice_event())
787791

788792
metrics = instrumentor_context.metrics
789793
metrics_attributes = self._extract_metrics_attributes()

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,11 @@ def _process_event(self, event):
138138

139139
return
140140

141+
def close(self):
142+
self.__wrapped__.close()
143+
# Treat the stream as done to ensure the span end.
144+
self._stream_done_callback(self._response)
145+
141146

142147
# pylint: disable=abstract-method
143148
class InvokeModelWithResponseStreamWrapper(ObjectProxy):
@@ -164,6 +169,11 @@ def __init__(
164169
self._tool_json_input_buf = ""
165170
self._record_message = False
166171

172+
def close(self):
173+
self.__wrapped__.close()
174+
# Treat the stream as done to ensure the span end.
175+
self._stream_done_callback(self._response)
176+
167177
def __iter__(self):
168178
try:
169179
for event in self.__wrapped__:

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

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
interactions:
2+
- request:
3+
body: |-
4+
{
5+
"messages": [
6+
{
7+
"role": "user",
8+
"content": [
9+
{
10+
"text": "Say this is a test"
11+
}
12+
]
13+
}
14+
],
15+
"inferenceConfig": {
16+
"maxTokens": 10,
17+
"temperature": 0.8,
18+
"topP": 1,
19+
"stopSequences": [
20+
"|"
21+
]
22+
}
23+
}
24+
headers:
25+
Content-Length:
26+
- '170'
27+
Content-Type:
28+
- application/json
29+
User-Agent:
30+
- Boto3/1.35.56 md/Botocore#1.35.56 ua/2.0 os/macos#24.4.0 md/arch#arm64 lang/python#3.12.0
31+
md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.35.56
32+
X-Amz-Date:
33+
- 20250509T104610Z
34+
X-Amz-Security-Token:
35+
- test_aws_security_token
36+
X-Amzn-Trace-Id:
37+
- Root=1-98fd1391-6f6d1be3cc209d9ba2fb3034;Parent=6dc095915efed44e;Sampled=1
38+
amz-sdk-invocation-id:
39+
- 3a37ec88-d989-4397-bfbe-d463b66a3a08
40+
amz-sdk-request:
41+
- attempt=1
42+
authorization:
43+
- Bearer test_aws_authorization
44+
method: POST
45+
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-lite-v1/converse-stream
46+
response:
47+
body:
48+
string: !!binary |
49+
AAAApAAAAFJl4NbnCzpldmVudC10eXBlBwAMbWVzc2FnZVN0YXJ0DTpjb250ZW50LXR5cGUHABBh
50+
cHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVldmVudHsicCI6ImFiY2RlZmdoaWprbG1u
51+
b3BxcnN0dXZ3eHl6QUJDREVGR0hJSktMTSIsInJvbGUiOiJhc3Npc3RhbnQifSKBqDQAAADZAAAA
52+
VxTIBhYLOmV2ZW50LXR5cGUHABFjb250ZW50QmxvY2tEZWx0YQ06Y29udGVudC10eXBlBwAQYXBw
53+
bGljYXRpb24vanNvbg06bWVzc2FnZS10eXBlBwAFZXZlbnR7ImNvbnRlbnRCbG9ja0luZGV4Ijow
54+
LCJkZWx0YSI6eyJ0ZXh0IjoiSGksIGRpZCB5b3UgaGF2ZSBhIHF1ZXN0aW9uIn0sInAiOiJhYmNk
55+
ZWZnaGlqa2xtbm9wcXJzdHV2d3h5ekFCQ0RFRkdISUpLIn22h6U3AAAAqgAAAFbdvayfCzpldmVu
56+
dC10eXBlBwAQY29udGVudEJsb2NrU3RvcA06Y29udGVudC10eXBlBwAQYXBwbGljYXRpb24vanNv
57+
bg06bWVzc2FnZS10eXBlBwAFZXZlbnR7ImNvbnRlbnRCbG9ja0luZGV4IjowLCJwIjoiYWJjZGVm
58+
Z2hpamtsbW5vcHFyc3R1dnd4eXpBQkNERUZHSElKS0wifQSjvQkAAAChAAAAUTQJCC0LOmV2ZW50
59+
LXR5cGUHAAttZXNzYWdlU3RvcA06Y29udGVudC10eXBlBwAQYXBwbGljYXRpb24vanNvbg06bWVz
60+
c2FnZS10eXBlBwAFZXZlbnR7InAiOiJhYmNkZWZnaGlqa2xtbm9wcXJzdHV2d3h5ekFCQ0QiLCJz
61+
dG9wUmVhc29uIjoibWF4X3Rva2VucyJ9jxPapgAAAMsAAABOaoNqNAs6ZXZlbnQtdHlwZQcACG1l
62+
dGFkYXRhDTpjb250ZW50LXR5cGUHABBhcHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVl
63+
dmVudHsibWV0cmljcyI6eyJsYXRlbmN5TXMiOjYyNX0sInAiOiJhYmNkZWZnaGlqa2wiLCJ1c2Fn
64+
ZSI6eyJpbnB1dFRva2VucyI6OCwib3V0cHV0VG9rZW5zIjoxMCwidG90YWxUb2tlbnMiOjE4fX0B
65+
/Sts
66+
headers:
67+
Connection:
68+
- keep-alive
69+
Content-Type:
70+
- application/vnd.amazon.eventstream
71+
Date:
72+
- Fri, 09 May 2025 10:46:10 GMT
73+
Set-Cookie: test_set_cookie
74+
Transfer-Encoding:
75+
- chunked
76+
x-amzn-RequestId:
77+
- 5aaaa521-d2c6-4980-ba0a-7bba19633a40
78+
status:
79+
code: 200
80+
message: OK
81+
version: 1

‎instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_response_stream_close_before_consumption[amazon.nova].yaml

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
interactions:
2+
- request:
3+
body: |-
4+
{
5+
"messages": [
6+
{
7+
"role": "user",
8+
"content": [
9+
{
10+
"text": "Say this is a test"
11+
}
12+
]
13+
}
14+
],
15+
"inferenceConfig": {
16+
"max_new_tokens": 10,
17+
"temperature": 0.8,
18+
"topP": 1,
19+
"stopSequences": [
20+
"|"
21+
]
22+
},
23+
"schemaVersion": "messages-v1"
24+
}
25+
headers:
26+
Content-Length:
27+
- '207'
28+
User-Agent:
29+
- Boto3/1.35.56 md/Botocore#1.35.56 ua/2.0 os/macos#24.4.0 md/arch#arm64 lang/python#3.12.0
30+
md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.35.56
31+
X-Amz-Date:
32+
- 20250509T104611Z
33+
X-Amz-Security-Token:
34+
- test_aws_security_token
35+
X-Amzn-Trace-Id:
36+
- Root=1-5213a530-965d38a8ed67e35669b3df67;Parent=0a0288bcb5638afa;Sampled=1
37+
amz-sdk-invocation-id:
38+
- 3c22bafe-b2d4-4bb1-95f4-81103f04cfa9
39+
amz-sdk-request:
40+
- attempt=1
41+
authorization:
42+
- Bearer test_aws_authorization
43+
method: POST
44+
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.nova-micro-v1%3A0/invoke-with-response-stream
45+
response:
46+
body:
47+
string: !!binary |
48+
AAAAsgAAAEvuKxwFCzpldmVudC10eXBlBwAFY2h1bmsNOmNvbnRlbnQtdHlwZQcAEGFwcGxpY2F0
49+
aW9uL2pzb24NOm1lc3NhZ2UtdHlwZQcABWV2ZW50eyJieXRlcyI6ImV5SnRaWE56WVdkbFUzUmhj
50+
blFpT25zaWNtOXNaU0k2SW1GemMybHpkR0Z1ZENKOWZRPT0iLCJwIjoiYWJjZGVmZ2hpamtsbW5v
51+
cCJ9Z2/zBQAAAPIAAABLtthETAs6ZXZlbnQtdHlwZQcABWNodW5rDTpjb250ZW50LXR5cGUHABBh
52+
cHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVldmVudHsiYnl0ZXMiOiJleUpqYjI1MFpX
53+
NTBRbXh2WTJ0RVpXeDBZU0k2ZXlKa1pXeDBZU0k2ZXlKMFpYaDBJam9pU1hRaWZTd2lZMjl1ZEdW
54+
dWRFSnNiMk5yU1c1a1pYZ2lPakI5ZlE9PSIsInAiOiJhYmNkZWZnaGlqa2xtbm9wcXJzdHV2d3h5
55+
ekFCQ0RFRkdISUpLTE1OIn3kNVULAAAAzAAAAEuoyUKrCzpldmVudC10eXBlBwAFY2h1bmsNOmNv
56+
bnRlbnQtdHlwZQcAEGFwcGxpY2F0aW9uL2pzb24NOm1lc3NhZ2UtdHlwZQcABWV2ZW50eyJieXRl
57+
cyI6ImV5SmpiMjUwWlc1MFFteHZZMnRUZEc5d0lqcDdJbU52Ym5SbGJuUkNiRzlqYTBsdVpHVjRJ
58+
am93ZlgwPSIsInAiOiJhYmNkZWZnaGlqa2xtbm9wcXJzdHV2d3h5ekFCQ0RFRkdIIn0XD1FEAAAB
59+
IwAAAEvJEcEACzpldmVudC10eXBlBwAFY2h1bmsNOmNvbnRlbnQtdHlwZQcAEGFwcGxpY2F0aW9u
60+
L2pzb24NOm1lc3NhZ2UtdHlwZQcABWV2ZW50eyJieXRlcyI6ImV5SmpiMjUwWlc1MFFteHZZMnRF
61+
Wld4MFlTSTZleUprWld4MFlTSTZleUowWlhoMElqb2lJSE52ZFc1a2N5QnNhV3RsSUhsdmRTZHla
62+
U0JwYm1sMGFXRjBhVzVuSW4wc0ltTnZiblJsYm5SQ2JHOWphMGx1WkdWNElqb3hmWDA9IiwicCI6
63+
ImFiY2RlZmdoaWprbG1ub3BxcnN0dXZ3eHl6QUJDREVGR0hJSktMTU5PUFFSU1RVVldYWVowIn07
64+
YX/vAAAAyQAAAEtgKc3bCzpldmVudC10eXBlBwAFY2h1bmsNOmNvbnRlbnQtdHlwZQcAEGFwcGxp
65+
Y2F0aW9uL2pzb24NOm1lc3NhZ2UtdHlwZQcABWV2ZW50eyJieXRlcyI6ImV5SmpiMjUwWlc1MFFt
66+
eHZZMnRUZEc5d0lqcDdJbU52Ym5SbGJuUkNiRzlqYTBsdVpHVjRJam94ZlgwPSIsInAiOiJhYmNk
67+
ZWZnaGlqa2xtbm9wcXJzdHV2d3h5ekFCQ0RFIn2N06RmAAAA4gAAAEvWONPOCzpldmVudC10eXBl
68+
BwAFY2h1bmsNOmNvbnRlbnQtdHlwZQcAEGFwcGxpY2F0aW9uL2pzb24NOm1lc3NhZ2UtdHlwZQcA
69+
BWV2ZW50eyJieXRlcyI6ImV5SmpiMjUwWlc1MFFteHZZMnRFWld4MFlTSTZleUprWld4MFlTSTZl
70+
eUowWlhoMElqb2lJR0VnZEdWemRDQnZjaUo5TENKamIyNTBaVzUwUW14dlkydEpibVJsZUNJNk1u
71+
MTkiLCJwIjoiYWJjZGVmZ2hpamtsbW5vcCJ9phCwbgAAALgAAABLpJsEpAs6ZXZlbnQtdHlwZQcA
72+
BWNodW5rDTpjb250ZW50LXR5cGUHABBhcHBsaWNhdGlvbi9qc29uDTptZXNzYWdlLXR5cGUHAAVl
73+
dmVudHsiYnl0ZXMiOiJleUpqYjI1MFpXNTBRbXh2WTJ0VGRHOXdJanA3SW1OdmJuUmxiblJDYkc5
74+
amEwbHVaR1Y0SWpveWZYMD0iLCJwIjoiYWJjZGVmZ2hpamtsbW4ifawRkW0AAAC4AAAAS6SbBKQL
75+
OmV2ZW50LXR5cGUHAAVjaHVuaw06Y29udGVudC10eXBlBwAQYXBwbGljYXRpb24vanNvbg06bWVz
76+
c2FnZS10eXBlBwAFZXZlbnR7ImJ5dGVzIjoiZXlKdFpYTnpZV2RsVTNSdmNDSTZleUp6ZEc5d1Vt
77+
VmhjMjl1SWpvaWJXRjRYM1J2YTJWdWN5SjlmUT09IiwicCI6ImFiY2RlZmdoaWprbG1uIn3/LlyA
78+
AAACPAAAAEutNbP9CzpldmVudC10eXBlBwAFY2h1bmsNOmNvbnRlbnQtdHlwZQcAEGFwcGxpY2F0
79+
aW9uL2pzb24NOm1lc3NhZ2UtdHlwZQcABWV2ZW50eyJieXRlcyI6ImV5SnRaWFJoWkdGMFlTSTZl
80+
eUoxYzJGblpTSTZleUpwYm5CMWRGUnZhMlZ1Y3lJNk5Td2liM1YwY0hWMFZHOXJaVzV6SWpveE1D
81+
d2lZMkZqYUdWU1pXRmtTVzV3ZFhSVWIydGxia052ZFc1MElqb3dMQ0pqWVdOb1pWZHlhWFJsU1c1
82+
d2RYUlViMnRsYmtOdmRXNTBJam93ZlN3aWJXVjBjbWxqY3lJNmUzMHNJblJ5WVdObElqcDdmWDBz
83+
SW1GdFlYcHZiaTFpWldSeWIyTnJMV2x1ZG05allYUnBiMjVOWlhSeWFXTnpJanA3SW1sdWNIVjBW
84+
RzlyWlc1RGIzVnVkQ0k2TlN3aWIzVjBjSFYwVkc5clpXNURiM1Z1ZENJNk1UQXNJbWx1ZG05allY
85+
UnBiMjVNWVhSbGJtTjVJam94Tmpjc0ltWnBjbk4wUW5sMFpVeGhkR1Z1WTNraU9qVTRMQ0pqWVdO
86+
b1pWSmxZV1JKYm5CMWRGUnZhMlZ1UTI5MWJuUWlPakFzSW1OaFkyaGxWM0pwZEdWSmJuQjFkRlJ2
87+
YTJWdVEyOTFiblFpT2pCOWZRPT0iLCJwIjoiYWJjZGVmZ2hpamtsbW5vcHFyc3R1dnd4eXoifeNI
88+
JRA=
89+
headers:
90+
Connection:
91+
- keep-alive
92+
Content-Type:
93+
- application/vnd.amazon.eventstream
94+
Date:
95+
- Fri, 09 May 2025 10:46:12 GMT
96+
Set-Cookie: test_set_cookie
97+
Transfer-Encoding:
98+
- chunked
99+
X-Amzn-Bedrock-Content-Type:
100+
- application/json
101+
x-amzn-RequestId:
102+
- 9c19e136-9600-42bd-9ecc-155abfd7b32b
103+
status:
104+
code: 200
105+
message: OK
106+
version: 1

‎instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_response_stream_close_before_consumption[amazon.titan].yaml

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
interactions:
2+
- request:
3+
body: |-
4+
{
5+
"inputText": "Say this is a test",
6+
"textGenerationConfig": {
7+
"maxTokenCount": 10,
8+
"temperature": 0.8,
9+
"topP": 1,
10+
"stopSequences": [
11+
"|"
12+
]
13+
}
14+
}
15+
headers:
16+
Content-Length:
17+
- '137'
18+
User-Agent:
19+
- Boto3/1.35.56 md/Botocore#1.35.56 ua/2.0 os/macos#24.4.0 md/arch#arm64 lang/python#3.12.0
20+
md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.35.56
21+
X-Amz-Date:
22+
- 20250509T104612Z
23+
X-Amz-Security-Token:
24+
- test_aws_security_token
25+
X-Amzn-Trace-Id:
26+
- Root=1-aacd9195-44986c6d8f97a81aa0b3e92f;Parent=c40ea6db4768abf2;Sampled=1
27+
amz-sdk-invocation-id:
28+
- 3f596967-dd06-445e-983e-573250f2a2ba
29+
amz-sdk-request:
30+
- attempt=1
31+
authorization:
32+
- Bearer test_aws_authorization
33+
method: POST
34+
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-lite-v1/invoke-with-response-stream
35+
response:
36+
body:
37+
string: !!binary |
38+
AAACAgAAAEuzJLUaCzpldmVudC10eXBlBwAFY2h1bmsNOmNvbnRlbnQtdHlwZQcAEGFwcGxpY2F0
39+
aW9uL2pzb24NOm1lc3NhZ2UtdHlwZQcABWV2ZW50eyJieXRlcyI6ImV5SnZkWFJ3ZFhSVVpYaDBJ
40+
am9pSUdOdmJXMWxiblJjYmtobGJHeHZJU0JKSUdGdElHRWdZMjl0Y0hWMFpYSWdjSEp2WjNKaGJT
41+
QmtaWE5wWjI1bFpDSXNJbWx1WkdWNElqb3dMQ0owYjNSaGJFOTFkSEIxZEZSbGVIUlViMnRsYmtO
42+
dmRXNTBJam94TUN3aVkyOXRjR3hsZEdsdmJsSmxZWE52YmlJNklreEZUa2RVU0NJc0ltbHVjSFYw
43+
VkdWNGRGUnZhMlZ1UTI5MWJuUWlPalVzSW1GdFlYcHZiaTFpWldSeWIyTnJMV2x1ZG05allYUnBi
44+
MjVOWlhSeWFXTnpJanA3SW1sdWNIVjBWRzlyWlc1RGIzVnVkQ0k2TlN3aWIzVjBjSFYwVkc5clpX
45+
NURiM1Z1ZENJNk1UQXNJbWx1ZG05allYUnBiMjVNWVhSbGJtTjVJam8yTVRZc0ltWnBjbk4wUW5s
46+
MFpVeGhkR1Z1WTNraU9qWXhObjE5IiwicCI6ImFiY2RlZmdoaWprbG1ub3BxcnN0dXZ3eCJ9/QFZ
47+
eg==
48+
headers:
49+
Connection:
50+
- keep-alive
51+
Content-Type:
52+
- application/vnd.amazon.eventstream
53+
Date:
54+
- Fri, 09 May 2025 10:46:12 GMT
55+
Set-Cookie: test_set_cookie
56+
Transfer-Encoding:
57+
- chunked
58+
X-Amzn-Bedrock-Content-Type:
59+
- application/json
60+
x-amzn-RequestId:
61+
- 9fa89295-177e-46a5-a898-1ebacad2b339
62+
status:
63+
code: 200
64+
message: OK
65+
version: 1

0 commit comments

Comments
 (0)