Skip to content

Commit f22f8b1

Browse files
author
Jeel Mehta
committed
Gen-AI python implementation
1 parent 95e31e3 commit f22f8b1

File tree

3 files changed

+172
-182
lines changed

3 files changed

+172
-182
lines changed

‎aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_span_processing_util.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
GEN_AI_USAGE_INPUT_TOKENS: str = "gen_ai.usage.input_tokens"
3737
GEN_AI_USAGE_OUTPUT_TOKENS: str = "gen_ai.usage.output_tokens"
3838

39+
3940
# Get dialect keywords retrieved from dialect_keywords.json file.
4041
# Only meant to be invoked by SQL_KEYWORD_PATTERN and unit tests
4142
def _get_dialect_keywords() -> List[str]:

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

Lines changed: 101 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,11 @@
22
# SPDX-License-Identifier: Apache-2.0
33
import abc
44
import inspect
5-
from typing import Any, Dict, Optional
65
import json
7-
from botocore.response import StreamingBody
86
import math
7+
from typing import Any, Dict, Optional
8+
9+
from botocore.response import StreamingBody
910

1011
from amazon.opentelemetry.distro._aws_attribute_keys import (
1112
AWS_BEDROCK_AGENT_ID,
@@ -14,7 +15,16 @@
1415
AWS_BEDROCK_GUARDRAIL_ID,
1516
AWS_BEDROCK_KNOWLEDGE_BASE_ID,
1617
)
17-
from amazon.opentelemetry.distro._aws_span_processing_util import GEN_AI_REQUEST_MODEL, GEN_AI_SYSTEM, GEN_AI_REQUEST_MAX_TOKENS, GEN_AI_REQUEST_TEMPERATURE, GEN_AI_REQUEST_TOP_P, GEN_AI_RESPONSE_FINISH_REASONS, GEN_AI_USAGE_INPUT_TOKENS, GEN_AI_USAGE_OUTPUT_TOKENS
18+
from amazon.opentelemetry.distro._aws_span_processing_util import (
19+
GEN_AI_REQUEST_MAX_TOKENS,
20+
GEN_AI_REQUEST_MODEL,
21+
GEN_AI_REQUEST_TEMPERATURE,
22+
GEN_AI_REQUEST_TOP_P,
23+
GEN_AI_RESPONSE_FINISH_REASONS,
24+
GEN_AI_SYSTEM,
25+
GEN_AI_USAGE_INPUT_TOKENS,
26+
GEN_AI_USAGE_OUTPUT_TOKENS,
27+
)
1828
from opentelemetry.instrumentation.botocore.extensions.types import (
1929
_AttributeMapT,
2030
_AwsSdkCallContext,
@@ -241,100 +251,101 @@ def extract_attributes(self, attributes: _AttributeMapT):
241251
attributes[GEN_AI_SYSTEM] = _AWS_BEDROCK_SYSTEM
242252

243253
model_id = self._call_context.params.get(_MODEL_ID)
244-
#attributes["Testing"]= "Test"
254+
# attributes["Testing"]= "Test"
245255
if model_id:
246-
attributes[GEN_AI_REQUEST_MODEL] = model_id
247-
256+
attributes[GEN_AI_REQUEST_MODEL] = model_id
257+
248258
# Get the request body if it exists
249-
body = self._call_context.params.get('body')
250-
#print("This is the body :",body)
259+
body = self._call_context.params.get("body")
260+
# print("This is the body :",body)
251261
if body:
252262
try:
253263
request_body = json.loads(body)
254-
255-
if 'amazon.titan' in model_id:
264+
265+
if "amazon.titan" in model_id:
256266
self._extract_titan_attributes(attributes, request_body)
257-
elif 'anthropic.claude' in model_id:
267+
elif "anthropic.claude" in model_id:
258268
self._extract_claude_attributes(attributes, request_body)
259-
elif 'meta.llama' in model_id:
269+
elif "meta.llama" in model_id:
260270
self._extract_llama_attributes(attributes, request_body)
261-
elif 'cohere.command' in model_id:
271+
elif "cohere.command" in model_id:
262272
self._extract_cohere_attributes(attributes, request_body)
263-
elif 'ai21.jamba' in model_id:
273+
elif "ai21.jamba" in model_id:
264274
self._extract_ai21_attributes(attributes, request_body)
265-
elif 'mistral' in model_id:
275+
elif "mistral" in model_id:
266276
self._extract_mistral_attributes(attributes, request_body)
267-
277+
268278
except json.JSONDecodeError:
269279
print("Error: Unable to parse the body as JSON")
280+
270281
def _extract_titan_attributes(self, attributes, request_body):
271-
config = request_body.get('textGenerationConfig', {})
272-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, config.get('temperature'))
273-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, config.get('topP'))
274-
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, config.get('maxTokenCount'))
282+
config = request_body.get("textGenerationConfig", {})
283+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, config.get("temperature"))
284+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, config.get("topP"))
285+
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, config.get("maxTokenCount"))
275286

276287
def _extract_claude_attributes(self, attributes, request_body):
277-
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get('max_tokens'))
278-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get('temperature'))
279-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get('top_p'))
288+
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens"))
289+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature"))
290+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get("top_p"))
280291

281292
def _extract_cohere_attributes(self, attributes, request_body):
282-
prompt = request_body.get('message')
293+
prompt = request_body.get("message")
283294
if prompt:
284295
attributes[GEN_AI_USAGE_INPUT_TOKENS] = math.ceil(len(prompt) / 6)
285-
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get('max_tokens'))
286-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get('temperature'))
287-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get('p'))
296+
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens"))
297+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature"))
298+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get("p"))
288299

289300
def _extract_ai21_attributes(self, attributes, request_body):
290-
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get('max_tokens'))
291-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get('temperature'))
292-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get('top_p'))
301+
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens"))
302+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature"))
303+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get("top_p"))
293304

294305
def _extract_llama_attributes(self, attributes, request_body):
295-
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get('max_gen_len'))
296-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get('temperature'))
297-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get('top_p'))
306+
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_gen_len"))
307+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature"))
308+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get("top_p"))
298309

299310
def _extract_mistral_attributes(self, attributes, request_body):
300311
print("This is the request body:", request_body)
301-
prompt = request_body.get('prompt')
312+
prompt = request_body.get("prompt")
302313
if prompt:
303314
attributes[GEN_AI_USAGE_INPUT_TOKENS] = math.ceil(len(prompt) / 6)
304-
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get('max_tokens'))
305-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get('temperature'))
306-
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get('top_p'))
315+
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens"))
316+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature"))
317+
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, request_body.get("top_p"))
307318

308319
@staticmethod
309320
def _set_if_not_none(attributes, key, value):
310321
if value is not None:
311322
attributes[key] = value
312-
323+
313324
def on_success(self, span: Span, result: Dict[str, Any]):
314325
super().on_success(span, result)
315-
326+
316327
model_id = self._call_context.params.get(_MODEL_ID)
317328
if not model_id:
318329
return
319330

320-
if 'body' in result and isinstance(result['body'], StreamingBody):
331+
if "body" in result and isinstance(result["body"], StreamingBody):
321332
try:
322333
# Read the entire content of the StreamingBody
323-
body_content = result['body'].read()
334+
body_content = result["body"].read()
324335
# Decode the bytes to string and parse as JSON
325-
response_body = json.loads(body_content.decode('utf-8'))
326-
327-
if 'amazon.titan' in model_id:
336+
response_body = json.loads(body_content.decode("utf-8"))
337+
338+
if "amazon.titan" in model_id:
328339
self._handle_amazon_titan_response(span, response_body)
329-
elif 'anthropic.claude' in model_id:
340+
elif "anthropic.claude" in model_id:
330341
self._handle_anthropic_claude_response(span, response_body)
331-
elif 'meta.llama' in model_id:
342+
elif "meta.llama" in model_id:
332343
self._handle_meta_llama_response(span, response_body)
333-
elif 'cohere.command' in model_id:
344+
elif "cohere.command" in model_id:
334345
self._handle_cohere_command_response(span, response_body)
335-
elif 'ai21.jamba' in model_id:
346+
elif "ai21.jamba" in model_id:
336347
self._handle_ai21_jamba_response(span, response_body)
337-
elif 'mistral' in model_id:
348+
elif "mistral" in model_id:
338349
self._handle_mistral_mistral_response(span, response_body)
339350

340351
except json.JSONDecodeError:
@@ -343,60 +354,60 @@ def on_success(self, span: Span, result: Dict[str, Any]):
343354
print(f"Error processing response: {str(e)}")
344355
finally:
345356
# Make sure to close the stream
346-
result['body'].close()
357+
result["body"].close()
347358

348359
def _handle_amazon_titan_response(self, span: Span, response_body: Dict[str, Any]):
349-
if 'inputTextTokenCount' in response_body:
350-
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, response_body['inputTextTokenCount'])
351-
352-
result = response_body['results'][0]
353-
if 'tokenCount' in result:
354-
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, result['tokenCount'])
355-
if 'completionReason' in result:
356-
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [result['completionReason']])
357-
360+
if "inputTextTokenCount" in response_body:
361+
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, response_body["inputTextTokenCount"])
362+
363+
result = response_body["results"][0]
364+
if "tokenCount" in result:
365+
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, result["tokenCount"])
366+
if "completionReason" in result:
367+
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [result["completionReason"]])
368+
358369
def _handle_anthropic_claude_response(self, span: Span, response_body: Dict[str, Any]):
359-
if 'usage' in response_body:
360-
usage = response_body['usage']
361-
if 'input_tokens' in usage:
362-
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, usage['input_tokens'])
363-
if 'output_tokens' in usage:
364-
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, usage['output_tokens'])
365-
if 'stop_reason' in response_body:
366-
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body['stop_reason']])
370+
if "usage" in response_body:
371+
usage = response_body["usage"]
372+
if "input_tokens" in usage:
373+
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, usage["input_tokens"])
374+
if "output_tokens" in usage:
375+
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, usage["output_tokens"])
376+
if "stop_reason" in response_body:
377+
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stop_reason"]])
367378

368379
def _handle_cohere_command_response(self, span: Span, response_body: Dict[str, Any]):
369380
# Output tokens: Approximate from the response text
370-
if 'text' in response_body:
371-
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, math.ceil(len(response_body['text']) / 6))
372-
if 'finish_reason' in response_body:
373-
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body['finish_reason']])
381+
if "text" in response_body:
382+
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, math.ceil(len(response_body["text"]) / 6))
383+
if "finish_reason" in response_body:
384+
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body["finish_reason"]])
374385

375386
def _handle_ai21_jamba_response(self, span: Span, response_body: Dict[str, Any]):
376-
if 'usage' in response_body:
377-
usage = response_body['usage']
378-
if 'prompt_tokens' in usage:
379-
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, usage['prompt_tokens'])
380-
if 'completion_tokens' in usage:
381-
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, usage['completion_tokens'])
382-
if 'choices' in response_body:
383-
choices = response_body['choices'][0]
384-
if 'finish_reason' in choices:
385-
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [choices['finish_reason']])
387+
if "usage" in response_body:
388+
usage = response_body["usage"]
389+
if "prompt_tokens" in usage:
390+
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, usage["prompt_tokens"])
391+
if "completion_tokens" in usage:
392+
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, usage["completion_tokens"])
393+
if "choices" in response_body:
394+
choices = response_body["choices"][0]
395+
if "finish_reason" in choices:
396+
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [choices["finish_reason"]])
386397

387398
def _handle_meta_llama_response(self, span: Span, response_body: Dict[str, Any]):
388-
if 'prompt_token_count' in response_body:
389-
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, response_body['prompt_token_count'])
390-
if 'generation_token_count' in response_body:
391-
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, response_body['generation_token_count'])
392-
if 'stop_reason' in response_body:
393-
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body['stop_reason']])
394-
399+
if "prompt_token_count" in response_body:
400+
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, response_body["prompt_token_count"])
401+
if "generation_token_count" in response_body:
402+
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, response_body["generation_token_count"])
403+
if "stop_reason" in response_body:
404+
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stop_reason"]])
405+
395406
def _handle_mistral_mistral_response(self, span: Span, response_body: Dict[str, Any]):
396407
print("This is the response body :", response_body)
397408
if "outputs" in response_body:
398409
outputs = response_body["outputs"][0]
399410
if "text" in outputs:
400411
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, math.ceil(len(outputs["text"]) / 6))
401-
if 'stop_reason' in outputs:
402-
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [outputs['stop_reason']])
412+
if "stop_reason" in outputs:
413+
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [outputs["stop_reason"]])

0 commit comments

Comments
 (0)