67 lines
1.8 KiB
Python
67 lines
1.8 KiB
Python
|
import json
|
||
|
|
|
||
|
|
from opentelemetry.instrumentation.sagemaker.utils import should_send_prompts
|
||
|
|
from opentelemetry.semconv._incubating.attributes import (
|
||
|
|
gen_ai_attributes as GenAIAttributes,
|
||
|
|
)
|
||
|
|
from opentelemetry.semconv_ai import (
|
||
|
|
SpanAttributes,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _try_parse_json(value):
|
||
|
|
"""Try to decode JSON if it's a string or bytes; fallback to raw value."""
|
||
|
|
try:
|
||
|
|
if isinstance(value, bytes):
|
||
|
|
value = value.decode("utf-8").strip()
|
||
|
|
if isinstance(value, str):
|
||
|
|
return json.loads(value)
|
||
|
|
return value
|
||
|
|
except (json.JSONDecodeError, UnicodeDecodeError):
|
||
|
|
return str(value)
|
||
|
|
|
||
|
|
|
||
|
|
def _set_span_attribute(span, name, value):
|
||
|
|
if value is not None:
|
||
|
|
if value != "":
|
||
|
|
span.set_attribute(name, value)
|
||
|
|
return
|
||
|
|
|
||
|
|
|
||
|
|
def set_stream_response_attributes(span, response_body):
|
||
|
|
if not span.is_recording() or not should_send_prompts():
|
||
|
|
return
|
||
|
|
|
||
|
|
_set_span_attribute(
|
||
|
|
span, SpanAttributes.TRACELOOP_ENTITY_OUTPUT, json.dumps(response_body)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def set_call_span_attributes(span, kwargs, response):
|
||
|
|
if not span.is_recording():
|
||
|
|
return
|
||
|
|
|
||
|
|
endpoint_name = kwargs.get("EndpointName")
|
||
|
|
_set_span_attribute(span, GenAIAttributes.GEN_AI_REQUEST_MODEL, endpoint_name)
|
||
|
|
|
||
|
|
|
||
|
|
def set_call_request_attributes(span, kwargs):
|
||
|
|
if not span.is_recording() or not should_send_prompts():
|
||
|
|
return
|
||
|
|
|
||
|
|
raw_request = kwargs.get("Body")
|
||
|
|
request_body = _try_parse_json(raw_request)
|
||
|
|
_set_span_attribute(
|
||
|
|
span, SpanAttributes.TRACELOOP_ENTITY_INPUT, json.dumps(request_body)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def set_call_response_attributes(span, raw_response):
|
||
|
|
if not span.is_recording() or not should_send_prompts():
|
||
|
|
return
|
||
|
|
response_body = _try_parse_json(raw_response)
|
||
|
|
|
||
|
|
_set_span_attribute(
|
||
|
|
span, SpanAttributes.TRACELOOP_ENTITY_OUTPUT, json.dumps(response_body)
|
||
|
|
)
|