Skip to content

Commit ab17f7c

Browse files
ezhang6811ezhang6811
committed
more minor fixes
1 parent 9c88101 commit ab17f7c

File tree

3 files changed

+111
-43
lines changed

3 files changed

+111
-43
lines changed

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

+75-32
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,9 @@
3131
ConverseStreamWrapper,
3232
InvokeModelWithResponseStreamWrapper,
3333
_Choice,
34+
estimate_token_count,
3435
genai_capture_message_content,
3536
message_to_event,
36-
estimate_token_count,
3737
)
3838
from opentelemetry.instrumentation.botocore.extensions.types import (
3939
_AttributeMapT,
@@ -106,6 +106,7 @@
106106

107107
_MODEL_ID_KEY: str = "modelId"
108108

109+
109110
class _BedrockRuntimeExtension(_AwsSdkExtension):
110111
"""
111112
This class is an extension for <a
@@ -255,7 +256,9 @@ def _extract_titan_attributes(self, attributes, request_body):
255256
attributes, GEN_AI_REQUEST_MAX_TOKENS, config.get("maxTokenCount")
256257
)
257258
self._set_if_not_none(
258-
attributes, GEN_AI_REQUEST_STOP_SEQUENCES, config.get("stopSequences")
259+
attributes,
260+
GEN_AI_REQUEST_STOP_SEQUENCES,
261+
config.get("stopSequences"),
259262
)
260263

261264
def _extract_nova_attributes(self, attributes, request_body):
@@ -270,21 +273,29 @@ def _extract_nova_attributes(self, attributes, request_body):
270273
attributes, GEN_AI_REQUEST_MAX_TOKENS, config.get("max_new_tokens")
271274
)
272275
self._set_if_not_none(
273-
attributes, GEN_AI_REQUEST_STOP_SEQUENCES, config.get("stopSequences")
276+
attributes,
277+
GEN_AI_REQUEST_STOP_SEQUENCES,
278+
config.get("stopSequences"),
274279
)
275280

276281
def _extract_claude_attributes(self, attributes, request_body):
277282
self._set_if_not_none(
278-
attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens")
283+
attributes,
284+
GEN_AI_REQUEST_MAX_TOKENS,
285+
request_body.get("max_tokens"),
279286
)
280287
self._set_if_not_none(
281-
attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature")
288+
attributes,
289+
GEN_AI_REQUEST_TEMPERATURE,
290+
request_body.get("temperature"),
282291
)
283292
self._set_if_not_none(
284293
attributes, GEN_AI_REQUEST_TOP_P, request_body.get("top_p")
285294
)
286295
self._set_if_not_none(
287-
attributes, GEN_AI_REQUEST_STOP_SEQUENCES, request_body.get("stop_sequences")
296+
attributes,
297+
GEN_AI_REQUEST_STOP_SEQUENCES,
298+
request_body.get("stop_sequences"),
288299
)
289300

290301
def _extract_command_r_attributes(self, attributes, request_body):
@@ -293,16 +304,22 @@ def _extract_command_r_attributes(self, attributes, request_body):
293304
attributes, GEN_AI_USAGE_INPUT_TOKENS, estimate_token_count(prompt)
294305
)
295306
self._set_if_not_none(
296-
attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens")
307+
attributes,
308+
GEN_AI_REQUEST_MAX_TOKENS,
309+
request_body.get("max_tokens"),
297310
)
298311
self._set_if_not_none(
299-
attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature")
312+
attributes,
313+
GEN_AI_REQUEST_TEMPERATURE,
314+
request_body.get("temperature"),
300315
)
301316
self._set_if_not_none(
302317
attributes, GEN_AI_REQUEST_TOP_P, request_body.get("p")
303318
)
304319
self._set_if_not_none(
305-
attributes, GEN_AI_REQUEST_STOP_SEQUENCES, request_body.get("stop_sequences")
320+
attributes,
321+
GEN_AI_REQUEST_STOP_SEQUENCES,
322+
request_body.get("stop_sequences"),
306323
)
307324

308325
def _extract_command_attributes(self, attributes, request_body):
@@ -311,24 +328,34 @@ def _extract_command_attributes(self, attributes, request_body):
311328
attributes, GEN_AI_USAGE_INPUT_TOKENS, estimate_token_count(prompt)
312329
)
313330
self._set_if_not_none(
314-
attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens")
331+
attributes,
332+
GEN_AI_REQUEST_MAX_TOKENS,
333+
request_body.get("max_tokens"),
315334
)
316335
self._set_if_not_none(
317-
attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature")
336+
attributes,
337+
GEN_AI_REQUEST_TEMPERATURE,
338+
request_body.get("temperature"),
318339
)
319340
self._set_if_not_none(
320341
attributes, GEN_AI_REQUEST_TOP_P, request_body.get("p")
321342
)
322343
self._set_if_not_none(
323-
attributes, GEN_AI_REQUEST_STOP_SEQUENCES, request_body.get("stop_sequences")
344+
attributes,
345+
GEN_AI_REQUEST_STOP_SEQUENCES,
346+
request_body.get("stop_sequences"),
324347
)
325348

326349
def _extract_llama_attributes(self, attributes, request_body):
327350
self._set_if_not_none(
328-
attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_gen_len")
351+
attributes,
352+
GEN_AI_REQUEST_MAX_TOKENS,
353+
request_body.get("max_gen_len"),
329354
)
330355
self._set_if_not_none(
331-
attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature")
356+
attributes,
357+
GEN_AI_REQUEST_TEMPERATURE,
358+
request_body.get("temperature"),
332359
)
333360
self._set_if_not_none(
334361
attributes, GEN_AI_REQUEST_TOP_P, request_body.get("top_p")
@@ -339,13 +366,19 @@ def _extract_mistral_attributes(self, attributes, request_body):
339366
prompt = request_body.get("prompt")
340367
if prompt:
341368
self._set_if_not_none(
342-
attributes, GEN_AI_USAGE_INPUT_TOKENS, estimate_token_count(prompt)
369+
attributes,
370+
GEN_AI_USAGE_INPUT_TOKENS,
371+
estimate_token_count(prompt),
343372
)
344373
self._set_if_not_none(
345-
attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens")
374+
attributes,
375+
GEN_AI_REQUEST_MAX_TOKENS,
376+
request_body.get("max_tokens"),
346377
)
347378
self._set_if_not_none(
348-
attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature")
379+
attributes,
380+
GEN_AI_REQUEST_TEMPERATURE,
381+
request_body.get("temperature"),
349382
)
350383
self._set_if_not_none(
351384
attributes, GEN_AI_REQUEST_TOP_P, request_body.get("top_p")
@@ -361,7 +394,6 @@ def _set_if_not_none(attributes, key, value):
361394

362395
def _get_request_messages(self):
363396
"""Extracts and normalize system and user / assistant messages"""
364-
input_text = None
365397
if system := self._call_context.params.get("system", []):
366398
system_messages = [{"role": "system", "content": system}]
367399
else:
@@ -390,20 +422,23 @@ def _get_request_messages(self):
390422
messages = self._get_messages_from_input_text(
391423
decoded_body, "message"
392424
)
393-
elif "cohere.command" in model_id or "meta.llama" in model_id or "mistral.mistral" in model_id:
425+
elif (
426+
"cohere.command" in model_id
427+
or "meta.llama" in model_id
428+
or "mistral.mistral" in model_id
429+
):
394430
messages = self._get_messages_from_input_text(
395431
decoded_body, "prompt"
396432
)
397433

398434
return system_messages + messages
399435

436+
# pylint: disable=no-self-use
400437
def _get_messages_from_input_text(
401438
self, decoded_body: dict[str, Any], input_name: str
402439
):
403440
if input_text := decoded_body.get(input_name):
404-
return [
405-
{"role": "user", "content": [{"text": input_text}]}
406-
]
441+
return [{"role": "user", "content": [{"text": input_text}]}]
407442
return []
408443

409444
def before_service_call(
@@ -843,11 +878,13 @@ def _handle_cohere_command_r_response(
843878
):
844879
if "text" in response_body:
845880
span.set_attribute(
846-
GEN_AI_USAGE_OUTPUT_TOKENS, estimate_token_count(response_body["text"])
881+
GEN_AI_USAGE_OUTPUT_TOKENS,
882+
estimate_token_count(response_body["text"]),
847883
)
848884
if "finish_reason" in response_body:
849885
span.set_attribute(
850-
GEN_AI_RESPONSE_FINISH_REASONS, [response_body["finish_reason"]]
886+
GEN_AI_RESPONSE_FINISH_REASONS,
887+
[response_body["finish_reason"]],
851888
)
852889

853890
event_logger = instrumentor_context.event_logger
@@ -867,11 +904,13 @@ def _handle_cohere_command_response(
867904
generations = response_body["generations"][0]
868905
if "text" in generations:
869906
span.set_attribute(
870-
GEN_AI_USAGE_OUTPUT_TOKENS, estimate_token_count(generations["text"])
907+
GEN_AI_USAGE_OUTPUT_TOKENS,
908+
estimate_token_count(generations["text"]),
871909
)
872910
if "finish_reason" in generations:
873911
span.set_attribute(
874-
GEN_AI_RESPONSE_FINISH_REASONS, [generations["finish_reason"]]
912+
GEN_AI_RESPONSE_FINISH_REASONS,
913+
[generations["finish_reason"]],
875914
)
876915

877916
event_logger = instrumentor_context.event_logger
@@ -893,17 +932,16 @@ def _handle_meta_llama_response(
893932
)
894933
if "generation_token_count" in response_body:
895934
span.set_attribute(
896-
GEN_AI_USAGE_OUTPUT_TOKENS, response_body["generation_token_count"],
935+
GEN_AI_USAGE_OUTPUT_TOKENS,
936+
response_body["generation_token_count"],
897937
)
898938
if "stop_reason" in response_body:
899939
span.set_attribute(
900940
GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stop_reason"]]
901941
)
902942

903943
event_logger = instrumentor_context.event_logger
904-
choice = _Choice.from_invoke_meta_llama(
905-
response_body, capture_content
906-
)
944+
choice = _Choice.from_invoke_meta_llama(response_body, capture_content)
907945
event_logger.emit(choice.to_choice_event())
908946

909947
def _handle_mistral_ai_response(
@@ -916,9 +954,14 @@ def _handle_mistral_ai_response(
916954
if "outputs" in response_body:
917955
outputs = response_body["outputs"][0]
918956
if "text" in outputs:
919-
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, estimate_token_count(outputs["text"]))
957+
span.set_attribute(
958+
GEN_AI_USAGE_OUTPUT_TOKENS,
959+
estimate_token_count(outputs["text"]),
960+
)
920961
if "stop_reason" in outputs:
921-
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [outputs["stop_reason"]])
962+
span.set_attribute(
963+
GEN_AI_RESPONSE_FINISH_REASONS, [outputs["stop_reason"]]
964+
)
922965

923966
event_logger = instrumentor_context.event_logger
924967
choice = _Choice.from_invoke_mistral_mistral(

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

+2-3
Original file line numberDiff line numberDiff line change
@@ -358,9 +358,8 @@ def _process_anthropic_claude_chunk(self, chunk):
358358
self._stream_done_callback(self._response)
359359
return
360360

361-
def estimate_token_count(
362-
message: str
363-
) -> int:
361+
362+
def estimate_token_count(message: str) -> int:
364363
# https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
365364
# use 6 chars per token to approximate token count when not provided in response body
366365
return math.ceil(len(message) / 6)

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

+34-8
Original file line numberDiff line numberDiff line change
@@ -1178,15 +1178,15 @@ def set_if_not_none(config, key, value):
11781178
"max_tokens": max_tokens,
11791179
"temperature": temperature,
11801180
"p": top_p,
1181-
"stop_sequences": stop_sequences
1181+
"stop_sequences": stop_sequences,
11821182
}
11831183
elif "cohere.command" in llm_model:
11841184
body = {
11851185
"prompt": prompt,
11861186
"max_tokens": max_tokens,
11871187
"temperature": temperature,
11881188
"p": top_p,
1189-
"stop_sequences": stop_sequences
1189+
"stop_sequences": stop_sequences,
11901190
}
11911191
elif "meta.llama" in llm_model:
11921192
body = {
@@ -1201,7 +1201,7 @@ def set_if_not_none(config, key, value):
12011201
"max_tokens": max_tokens,
12021202
"temperature": temperature,
12031203
"top_p": top_p,
1204-
"stop": stop_sequences
1204+
"stop": stop_sequences,
12051205
}
12061206
else:
12071207
raise ValueError(f"No config for {llm_model}")
@@ -1224,7 +1224,15 @@ def get_model_name_from_family(llm_model):
12241224

12251225
@pytest.mark.parametrize(
12261226
"model_family",
1227-
["amazon.nova", "amazon.titan", "anthropic.claude", "cohere.command-r", "cohere.command", "meta.llama", "mistral.mistral"],
1227+
[
1228+
"amazon.nova",
1229+
"amazon.titan",
1230+
"anthropic.claude",
1231+
"cohere.command-r",
1232+
"cohere.command",
1233+
"meta.llama",
1234+
"mistral.mistral",
1235+
],
12281236
)
12291237
@pytest.mark.vcr()
12301238
def test_invoke_model_with_content(
@@ -1236,7 +1244,12 @@ def test_invoke_model_with_content(
12361244
):
12371245
# pylint:disable=too-many-locals
12381246
llm_model_value = get_model_name_from_family(model_family)
1239-
max_tokens, temperature, top_p, stop_sequences = 10, 0.8, 0.99 if model_family == "cohere.command-r" else 1, ["|"]
1247+
max_tokens, temperature, top_p, stop_sequences = (
1248+
10,
1249+
0.8,
1250+
0.99 if model_family == "cohere.command-r" else 1,
1251+
["|"],
1252+
)
12401253
body = get_invoke_model_body(
12411254
llm_model_value, max_tokens, temperature, top_p, stop_sequences
12421255
)
@@ -1284,7 +1297,7 @@ def test_invoke_model_with_content(
12841297
finish_reason = "max_tokens"
12851298
elif model_family == "cohere.command-r":
12861299
message = {
1287-
"content": "This is a test. How are you doing today",
1300+
"content": "This is a test. How are you doing today",
12881301
}
12891302
finish_reason = "MAX_TOKENS"
12901303
elif model_family == "cohere.command":
@@ -1835,7 +1848,15 @@ def test_invoke_model_with_content_tool_call(
18351848

18361849
@pytest.mark.parametrize(
18371850
"model_family",
1838-
["amazon.nova", "amazon.titan", "anthropic.claude", "cohere.command-r", "cohere.command", "meta.llama", "mistral.mistral"],
1851+
[
1852+
"amazon.nova",
1853+
"amazon.titan",
1854+
"anthropic.claude",
1855+
"cohere.command-r",
1856+
"cohere.command",
1857+
"meta.llama",
1858+
"mistral.mistral",
1859+
],
18391860
)
18401861
@pytest.mark.vcr()
18411862
def test_invoke_model_no_content(
@@ -1847,7 +1868,12 @@ def test_invoke_model_no_content(
18471868
):
18481869
# pylint:disable=too-many-locals
18491870
llm_model_value = get_model_name_from_family(model_family)
1850-
max_tokens, temperature, top_p, stop_sequences = 10, 0.8, 0.99 if model_family == "cohere.command-r" else 1, ["|"]
1871+
max_tokens, temperature, top_p, stop_sequences = (
1872+
10,
1873+
0.8,
1874+
0.99 if model_family == "cohere.command-r" else 1,
1875+
["|"],
1876+
)
18511877
body = get_invoke_model_body(
18521878
llm_model_value, max_tokens, temperature, top_p, stop_sequences
18531879
)

0 commit comments

Comments
 (0)