Skip to content

Commit 551fa85

Browse files
author
Jeel Mehta
committed
Gen-AI python implementation
1 parent dcc616a commit 551fa85

File tree

2 files changed

+79
-42
lines changed

2 files changed

+79
-42
lines changed

‎aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py‎

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -279,6 +279,9 @@ def _extract_claude_attributes(self, attributes, request_body):
279279
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get('top_p'))
280280

281281
def _extract_cohere_attributes(self, attributes, request_body):
282+
prompt = request_body.get('message')
283+
if prompt:
284+
attributes[GEN_AI_USAGE_INPUT_TOKENS] = math.ceil(len(prompt) / 6)
282285
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get('max_tokens'))
283286
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get('temperature'))
284287
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get('p'))
@@ -294,6 +297,7 @@ def _extract_llama_attributes(self, attributes, request_body):
294297
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get('top_p'))
295298

296299
def _extract_mistral_attributes(self, attributes, request_body):
300+
print("This is the request body:", request_body)
297301
prompt = request_body.get('prompt')
298302
if prompt:
299303
attributes[GEN_AI_USAGE_INPUT_TOKENS] = math.ceil(len(prompt) / 6)
@@ -342,7 +346,6 @@ def on_success(self, span: Span, result: Dict[str, Any]):
342346
result['body'].close()
343347

344348
def _handle_amazon_titan_response(self, span: Span, response_body: Dict[str, Any]):
345-
#print("This is the response body :", response_body)
346349
if 'inputTextTokenCount' in response_body:
347350
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, response_body['inputTextTokenCount'])
348351

@@ -353,7 +356,6 @@ def _handle_amazon_titan_response(self, span: Span, response_body: Dict[str, Any
353356
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [result['completionReason']])
354357

355358
def _handle_anthropic_claude_response(self, span: Span, response_body: Dict[str, Any]):
356-
#print("This is the response body :", response_body)
357359
if 'usage' in response_body:
358360
usage = response_body['usage']
359361
if 'input_tokens' in usage:
@@ -364,20 +366,13 @@ def _handle_anthropic_claude_response(self, span: Span, response_body: Dict[str,
364366
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body['stop_reason']])
365367

366368
def _handle_cohere_command_response(self, span: Span, response_body: Dict[str, Any]):
367-
print("This is the response body :", response_body)
368-
# Input tokens: Approximate from the user's message in chat history
369-
if 'chat_history' in response_body:
370-
user_messages = [msg['message'] for msg in response_body['chat_history'] if msg['role'] == 'USER']
371-
input_text = ' '.join(user_messages)
372-
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, math.ceil(len(input_text) / 6))
373369
# Output tokens: Approximate from the response text
374370
if 'text' in response_body:
375371
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, math.ceil(len(response_body['text']) / 6))
376372
if 'finish_reason' in response_body:
377373
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body['finish_reason']])
378374

379375
def _handle_ai21_jamba_response(self, span: Span, response_body: Dict[str, Any]):
380-
print("This is the response body :", response_body)
381376
if 'usage' in response_body:
382377
usage = response_body['usage']
383378
if 'prompt_tokens' in usage:
@@ -390,7 +385,6 @@ def _handle_ai21_jamba_response(self, span: Span, response_body: Dict[str, Any])
390385
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [choices['finish_reason']])
391386

392387
def _handle_meta_llama_response(self, span: Span, response_body: Dict[str, Any]):
393-
print("This is the response body :", response_body)
394388
if 'prompt_token_count' in response_body:
395389
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, response_body['prompt_token_count'])
396390
if 'generation_token_count' in response_body:

‎aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py‎

Lines changed: 75 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from io import BytesIO
88
import json
99
from botocore.response import StreamingBody
10+
import math
1011

1112
import gevent.monkey
1213
import pkg_resources
@@ -290,39 +291,42 @@ def _test_patched_botocore_instrumentation(self):
290291
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.output_tokens"], 36)
291292
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.response.finish_reasons"], ["end_turn"])
292293

293-
#BedrockRuntime - Cohere Command Models _testing Pending
294-
# self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
295-
# request_body = {
296-
# "max_tokens": 512,
297-
# "temperature": 0.5,
298-
# "p":0.75,
299-
# }
294+
#BedrockRuntime - Cohere Command Models
295+
self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
296+
request_body = {
297+
'message': "Describe the purpose of a 'hello world' program in one line.",
298+
"max_tokens": 512,
299+
"temperature": 0.5,
300+
"p":0.75,
301+
}
300302

301-
# bedrock_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
302-
# "bedrock-runtime",
303-
# model_id="cohere.command",
304-
# request_body=json.dumps(request_body)
305-
# )
306-
# self.assertEqual(len(bedrock_runtime_attributes), 5)
307-
# self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM)
308-
# self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], "cohere.command")
309-
# self.assertEqual(bedrock_runtime_attributes["gen_ai.request.max_tokens"], 512)
310-
# self.assertEqual(bedrock_runtime_attributes["gen_ai.request.temperature"], 0.5)
311-
# self.assertEqual(bedrock_runtime_attributes["gen_ai.request.top_p"], 0.75)
312-
# response_body = {
313-
# 'finish_reason': 'COMPLETE'
314-
# }
315-
# json_bytes = json.dumps(response_body).encode('utf-8')
316-
# body_bytes = BytesIO(json_bytes)
317-
# streaming_body = StreamingBody(body_bytes, len(json_bytes))
318-
# bedrock_runtime_success_attributes: Dict[str, str] = _do_on_success_bedrock(
319-
# "bedrock-runtime",
320-
# model_id="cohere.command",
321-
# streaming_body=streaming_body
322-
# )
323-
# self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.input_tokens"], 23)
324-
# self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.output_tokens"], 36)
325-
# self.assertEqual(bedrock_runtime_success_attributes["gen_ai.response.finish_reasons"], ["COMPLETE"])
303+
bedrock_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
304+
"bedrock-runtime",
305+
model_id="cohere.command",
306+
request_body=json.dumps(request_body)
307+
)
308+
self.assertEqual(len(bedrock_runtime_attributes), 6)
309+
self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM)
310+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], "cohere.command")
311+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.max_tokens"], 512)
312+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.temperature"], 0.5)
313+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.top_p"], 0.75)
314+
self.assertEqual(bedrock_runtime_attributes["gen_ai.usage.input_tokens"], math.ceil(len(request_body['message'])/6))
315+
response_body = {
316+
'text': 'A "hello world" program serves as a simple introduction to programming, helping developers confirm their setup and test their coding environment.',
317+
'finish_reason': 'COMPLETE'
318+
319+
}
320+
json_bytes = json.dumps(response_body).encode('utf-8')
321+
body_bytes = BytesIO(json_bytes)
322+
streaming_body = StreamingBody(body_bytes, len(json_bytes))
323+
bedrock_runtime_success_attributes: Dict[str, str] = _do_on_success_bedrock(
324+
"bedrock-runtime",
325+
model_id="cohere.command",
326+
streaming_body=streaming_body
327+
)
328+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.output_tokens"], math.ceil(len(response_body['text'])/6))
329+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.response.finish_reasons"], ["COMPLETE"])
326330

327331
#BedrockRuntime - AI21 Jamba Models
328332
self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
@@ -397,6 +401,45 @@ def _test_patched_botocore_instrumentation(self):
397401
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.output_tokens"], 36)
398402
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.response.finish_reasons"], ["stop"])
399403

404+
#BedrockRuntime - Mistral Models
405+
self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
406+
msg="Hello, World"
407+
formatted_prompt = f"<s>[INST] {msg} [/INST]"
408+
request_body = {
409+
'prompt': formatted_prompt,
410+
"max_tokens": 512,
411+
"temperature": 0.5,
412+
"top_p": 0.9,
413+
}
414+
415+
bedrock_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
416+
"bedrock-runtime",
417+
model_id="mistral",
418+
request_body=json.dumps(request_body)
419+
)
420+
self.assertEqual(len(bedrock_runtime_attributes), 6)
421+
self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM)
422+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], "mistral")
423+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.max_tokens"], 512)
424+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.temperature"], 0.5)
425+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.top_p"], 0.9)
426+
self.assertEqual(bedrock_runtime_attributes["gen_ai.usage.input_tokens"], math.ceil(len(request_body['prompt'])/6))
427+
response_body = {
428+
'outputs':[{
429+
'text': 'Goodbye, World',
430+
'stop_reason': 'stop'}]
431+
}
432+
json_bytes = json.dumps(response_body).encode('utf-8')
433+
body_bytes = BytesIO(json_bytes)
434+
streaming_body = StreamingBody(body_bytes, len(json_bytes))
435+
bedrock_runtime_success_attributes: Dict[str, str] = _do_on_success_bedrock(
436+
"bedrock-runtime",
437+
model_id="mistral",
438+
streaming_body=streaming_body
439+
)
440+
#self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.input_tokens"], 31) Srill have concerns regarging these lines
441+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.output_tokens"], math.ceil(len(response_body['outputs'][0]['text'])/6))
442+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.response.finish_reasons"], ["stop"])
400443

401444
# SecretsManager
402445
self.assertTrue("secretsmanager" in _KNOWN_EXTENSIONS)

0 commit comments

Comments
 (0)