Skip to content

Commit dcc616a

Browse files
author
Jeel Mehta
committed
Gen-AI python implementation
1 parent 4d7a014 commit dcc616a

File tree

2 files changed

+196
-11
lines changed

2 files changed

+196
-11
lines changed

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,7 @@ def on_success(self, span: Span, result: Dict[str, Any]):
342342
result['body'].close()
343343

344344
def _handle_amazon_titan_response(self, span: Span, response_body: Dict[str, Any]):
345+
#print("This is the response body :", response_body)
345346
if 'inputTextTokenCount' in response_body:
346347
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, response_body['inputTextTokenCount'])
347348

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

354355
def _handle_anthropic_claude_response(self, span: Span, response_body: Dict[str, Any]):
356+
#print("This is the response body :", response_body)
355357
if 'usage' in response_body:
356358
usage = response_body['usage']
357359
if 'input_tokens' in usage:
@@ -362,6 +364,7 @@ def _handle_anthropic_claude_response(self, span: Span, response_body: Dict[str,
362364
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body['stop_reason']])
363365

364366
def _handle_cohere_command_response(self, span: Span, response_body: Dict[str, Any]):
367+
print("This is the response body :", response_body)
365368
# Input tokens: Approximate from the user's message in chat history
366369
if 'chat_history' in response_body:
367370
user_messages = [msg['message'] for msg in response_body['chat_history'] if msg['role'] == 'USER']
@@ -387,15 +390,16 @@ def _handle_ai21_jamba_response(self, span: Span, response_body: Dict[str, Any])
387390
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [choices['finish_reason']])
388391

389392
def _handle_meta_llama_response(self, span: Span, response_body: Dict[str, Any]):
390-
#print("This is the response body :", response_body)
393+
print("This is the response body :", response_body)
391394
if 'prompt_token_count' in response_body:
392395
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, response_body['prompt_token_count'])
393396
if 'generation_token_count' in response_body:
394397
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, response_body['generation_token_count'])
395398
if 'stop_reason' in response_body:
396-
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, response_body['stop_reason'])
399+
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body['stop_reason']])
397400

398401
def _handle_mistral_mistral_response(self, span: Span, response_body: Dict[str, Any]):
402+
print("This is the response body :", response_body)
399403
if "outputs" in response_body:
400404
outputs = response_body["outputs"][0]
401405
if "text" in outputs:

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

Lines changed: 190 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@
44
from typing import Any, Dict
55
from unittest import TestCase
66
from unittest.mock import MagicMock, patch
7+
from io import BytesIO
8+
import json
9+
from botocore.response import StreamingBody
710

811
import gevent.monkey
912
import pkg_resources
@@ -211,12 +214,189 @@ def _test_patched_botocore_instrumentation(self):
211214
bedrock_agent_runtime_sucess_attributes: Dict[str, str] = _do_on_success_bedrock("bedrock-agent-runtime")
212215
self.assertEqual(len(bedrock_agent_runtime_sucess_attributes), 0)
213216

214-
# BedrockRuntime
217+
# BedrockRuntime - Amazon Titan Models
218+
self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
219+
request_body = {
220+
"textGenerationConfig": {
221+
"maxTokenCount": 512,
222+
"temperature": 0.9,
223+
"topP": 0.75,
224+
}
225+
}
226+
bedrock_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
227+
"bedrock-runtime",
228+
model_id="amazon.titan",
229+
request_body=json.dumps(request_body)
230+
)
231+
self.assertEqual(len(bedrock_runtime_attributes), 5)
232+
self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM)
233+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], "amazon.titan")
234+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.max_tokens"], 512)
235+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.temperature"], 0.9)
236+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.top_p"], 0.75)
237+
response_body = {
238+
"inputTextTokenCount": 123,
239+
"results": [{
240+
"tokenCount": 456,
241+
"outputText": "testing",
242+
"completionReason": "FINISH",
243+
}]
244+
}
245+
json_bytes = json.dumps(response_body).encode('utf-8')
246+
body_bytes = BytesIO(json_bytes)
247+
streaming_body = StreamingBody(body_bytes, len(json_bytes))
248+
bedrock_runtime_success_attributes: Dict[str, str] = _do_on_success_bedrock(
249+
"bedrock-runtime",
250+
model_id="amazon.titan",
251+
streaming_body=streaming_body
252+
)
253+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.input_tokens"], 123)
254+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.output_tokens"], 456)
255+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.response.finish_reasons"], ["FINISH"])
256+
257+
#BedrockRuntime - Anthropic Claude Models
258+
215259
self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
216-
bedrock_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock("bedrock-runtime")
217-
self.assertEqual(len(bedrock_runtime_attributes), 2)
260+
request_body = {
261+
"anthropic_version": "bedrock-2023-05-31",
262+
"max_tokens": 512,
263+
"temperature": 0.5,
264+
"top_p":0.999,
265+
}
266+
267+
bedrock_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
268+
"bedrock-runtime",
269+
model_id="anthropic.claude",
270+
request_body=json.dumps(request_body)
271+
)
272+
self.assertEqual(len(bedrock_runtime_attributes), 5)
218273
self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM)
219-
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], _GEN_AI_REQUEST_MODEL)
274+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], "anthropic.claude")
275+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.max_tokens"], 512)
276+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.temperature"], 0.5)
277+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.top_p"], 0.999)
278+
response_body = {
279+
'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 23, 'output_tokens': 36}
280+
}
281+
json_bytes = json.dumps(response_body).encode('utf-8')
282+
body_bytes = BytesIO(json_bytes)
283+
streaming_body = StreamingBody(body_bytes, len(json_bytes))
284+
bedrock_runtime_success_attributes: Dict[str, str] = _do_on_success_bedrock(
285+
"bedrock-runtime",
286+
model_id="anthropic.claude",
287+
streaming_body=streaming_body
288+
)
289+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.input_tokens"], 23)
290+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.output_tokens"], 36)
291+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.response.finish_reasons"], ["end_turn"])
292+
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+
# }
300+
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"])
326+
327+
#BedrockRuntime - AI21 Jamba Models
328+
self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
329+
request_body = {
330+
"max_tokens": 512,
331+
"temperature": 0.5,
332+
"top_p":0.9,
333+
}
334+
335+
bedrock_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
336+
"bedrock-runtime",
337+
model_id="ai21.jamba",
338+
request_body=json.dumps(request_body)
339+
)
340+
self.assertEqual(len(bedrock_runtime_attributes), 5)
341+
self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM)
342+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], "ai21.jamba")
343+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.max_tokens"], 512)
344+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.temperature"], 0.5)
345+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.top_p"], 0.9)
346+
response_body = {
347+
'choices':[{'finish_reason': "stop"}],
348+
'usage': {'prompt_tokens': 24,
349+
'completion_tokens': 31,
350+
'total_tokens': 55}
351+
}
352+
json_bytes = json.dumps(response_body).encode('utf-8')
353+
body_bytes = BytesIO(json_bytes)
354+
streaming_body = StreamingBody(body_bytes, len(json_bytes))
355+
bedrock_runtime_success_attributes: Dict[str, str] = _do_on_success_bedrock(
356+
"bedrock-runtime",
357+
model_id="ai21.jamba",
358+
streaming_body=streaming_body
359+
)
360+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.input_tokens"], 24)
361+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.output_tokens"], 31)
362+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.response.finish_reasons"], ["stop"])
363+
364+
#BedrockRuntime - Meta LLama Models
365+
self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
366+
request_body = {
367+
"max_gen_len": 512,
368+
"temperature": 0.5,
369+
"top_p":0.9,
370+
}
371+
372+
bedrock_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
373+
"bedrock-runtime",
374+
model_id="meta.llama",
375+
request_body=json.dumps(request_body)
376+
)
377+
self.assertEqual(len(bedrock_runtime_attributes), 5)
378+
self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM)
379+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], "meta.llama")
380+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.max_tokens"], 512)
381+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.temperature"], 0.5)
382+
self.assertEqual(bedrock_runtime_attributes["gen_ai.request.top_p"], 0.9)
383+
response_body = {
384+
'prompt_token_count': 31,
385+
'generation_token_count': 36,
386+
'stop_reason': 'stop'
387+
}
388+
json_bytes = json.dumps(response_body).encode('utf-8')
389+
body_bytes = BytesIO(json_bytes)
390+
streaming_body = StreamingBody(body_bytes, len(json_bytes))
391+
bedrock_runtime_success_attributes: Dict[str, str] = _do_on_success_bedrock(
392+
"bedrock-runtime",
393+
model_id="meta.llama",
394+
streaming_body=streaming_body
395+
)
396+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.input_tokens"], 31)
397+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.usage.output_tokens"], 36)
398+
self.assertEqual(bedrock_runtime_success_attributes["gen_ai.response.finish_reasons"], ["stop"])
399+
220400

221401
# SecretsManager
222402
self.assertTrue("secretsmanager" in _KNOWN_EXTENSIONS)
@@ -385,26 +565,27 @@ def _do_extract_sqs_attributes() -> Dict[str, str]:
385565
return _do_extract_attributes(service_name, params)
386566

387567

388-
def _do_extract_attributes_bedrock(service, operation=None) -> Dict[str, str]:
568+
def _do_extract_attributes_bedrock(service, operation=None, model_id=None, request_body=None) -> Dict[str, str]:
389569
params: Dict[str, Any] = {
390570
"agentId": _BEDROCK_AGENT_ID,
391571
"dataSourceId": _BEDROCK_DATASOURCE_ID,
392572
"knowledgeBaseId": _BEDROCK_KNOWLEDGEBASE_ID,
393573
"guardrailId": _BEDROCK_GUARDRAIL_ID,
394-
"modelId": _GEN_AI_REQUEST_MODEL,
574+
"modelId": model_id,
575+
"body": request_body,
395576
}
396577
return _do_extract_attributes(service, params, operation)
397578

398579

399-
def _do_on_success_bedrock(service, operation=None) -> Dict[str, str]:
580+
def _do_on_success_bedrock(service, operation=None, model_id=None, streaming_body=None) -> Dict[str, str]:
400581
result: Dict[str, Any] = {
401582
"agentId": _BEDROCK_AGENT_ID,
402583
"dataSourceId": _BEDROCK_DATASOURCE_ID,
403584
"knowledgeBaseId": _BEDROCK_KNOWLEDGEBASE_ID,
404585
"guardrailId": _BEDROCK_GUARDRAIL_ID,
405-
"modelId": _GEN_AI_REQUEST_MODEL,
586+
"body": streaming_body,
406587
}
407-
return _do_on_success(service, result, operation)
588+
return _do_on_success(service, result, operation, params={"modelId": model_id})
408589

409590

410591
def _do_extract_secretsmanager_attributes() -> Dict[str, str]:

0 commit comments

Comments
 (0)