64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
|
from opentelemetry.instrumentation.replicate.utils import (
|
||
|
|
dont_throw,
|
||
|
|
should_send_prompts,
|
||
|
|
)
|
||
|
|
from opentelemetry.semconv._incubating.attributes import (
|
||
|
|
gen_ai_attributes as GenAIAttributes,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _set_span_attribute(span, name, value):
|
||
|
|
if value is not None:
|
||
|
|
if value != "":
|
||
|
|
span.set_attribute(name, value)
|
||
|
|
return
|
||
|
|
|
||
|
|
|
||
|
|
@dont_throw
|
||
|
|
def set_input_attributes(span, args, kwargs):
|
||
|
|
if not span.is_recording():
|
||
|
|
return
|
||
|
|
|
||
|
|
input_attribute = kwargs.get("input")
|
||
|
|
if should_send_prompts():
|
||
|
|
_set_span_attribute(
|
||
|
|
span, f"{GenAIAttributes.GEN_AI_PROMPT}.0.user", input_attribute.get("prompt")
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@dont_throw
|
||
|
|
def set_model_input_attributes(span, args, kwargs):
|
||
|
|
if not span.is_recording():
|
||
|
|
return
|
||
|
|
|
||
|
|
if args is not None and len(args) > 0:
|
||
|
|
_set_span_attribute(span, GenAIAttributes.GEN_AI_REQUEST_MODEL, args[0])
|
||
|
|
elif kwargs.get("version"):
|
||
|
|
_set_span_attribute(
|
||
|
|
span, GenAIAttributes.GEN_AI_REQUEST_MODEL, kwargs.get("version").id
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
_set_span_attribute(span, GenAIAttributes.GEN_AI_REQUEST_MODEL, "unknown")
|
||
|
|
|
||
|
|
input_attribute = kwargs.get("input")
|
||
|
|
|
||
|
|
_set_span_attribute(
|
||
|
|
span, GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE, input_attribute.get("temperature")
|
||
|
|
)
|
||
|
|
_set_span_attribute(
|
||
|
|
span, GenAIAttributes.GEN_AI_REQUEST_TOP_P, input_attribute.get("top_p")
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@dont_throw
|
||
|
|
def set_response_attributes(span, response):
|
||
|
|
if should_send_prompts():
|
||
|
|
if isinstance(response, list):
|
||
|
|
for index, item in enumerate(response):
|
||
|
|
prefix = f"{GenAIAttributes.GEN_AI_COMPLETION}.{index}"
|
||
|
|
_set_span_attribute(span, f"{prefix}.content", item)
|
||
|
|
elif isinstance(response, str):
|
||
|
|
_set_span_attribute(
|
||
|
|
span, f"{GenAIAttributes.GEN_AI_COMPLETION}.0.content", response
|
||
|
|
)
|