121 lines
3.4 KiB
Python
121 lines
3.4 KiB
Python
import dataclasses
|
|
import json
|
|
import logging
|
|
import os
|
|
import traceback
|
|
|
|
from opentelemetry import context as context_api
|
|
from opentelemetry.instrumentation.haystack.config import Config
|
|
from opentelemetry.semconv_ai import SpanAttributes
|
|
|
|
|
|
class EnhancedJSONEncoder(json.JSONEncoder):
|
|
def default(self, o):
|
|
if dataclasses.is_dataclass(o):
|
|
return dataclasses.asdict(o)
|
|
if hasattr(o, "to_json"):
|
|
return o.to_json()
|
|
return super().default(o)
|
|
|
|
|
|
def should_send_prompts():
|
|
return (
|
|
os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
|
|
).lower() == "true" or context_api.get_value("override_enable_content_tracing")
|
|
|
|
|
|
def dont_throw(func):
|
|
"""
|
|
A decorator that wraps the passed in function and logs exceptions instead of throwing them.
|
|
|
|
@param func: The function to wrap
|
|
@return: The wrapper function
|
|
"""
|
|
# Obtain a logger specific to the function's module
|
|
logger = logging.getLogger(func.__module__)
|
|
|
|
def wrapper(*args, **kwargs):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except Exception as e:
|
|
logger.debug(
|
|
"OpenLLMetry failed to trace in %s, error: %s", func.__name__, str(e)
|
|
)
|
|
if Config.exception_logger:
|
|
Config.exception_logger(e)
|
|
|
|
return wrapper
|
|
|
|
|
|
@dont_throw
|
|
def process_request(span, args, kwargs):
|
|
if should_send_prompts():
|
|
kwargs_to_serialize = kwargs.copy()
|
|
for arg in args:
|
|
if arg and isinstance(arg, dict):
|
|
for key, value in arg.items():
|
|
kwargs_to_serialize[key] = value
|
|
args_to_serialize = [arg for arg in args if not isinstance(arg, dict)]
|
|
input_entity = {"args": args_to_serialize, "kwargs": kwargs_to_serialize}
|
|
span.set_attribute(
|
|
SpanAttributes.TRACELOOP_ENTITY_INPUT,
|
|
json.dumps(input_entity, cls=EnhancedJSONEncoder),
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def process_response(span, response):
|
|
if should_send_prompts():
|
|
span.set_attribute(
|
|
SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
|
|
json.dumps(response, cls=EnhancedJSONEncoder),
|
|
)
|
|
|
|
|
|
def set_span_attribute(span, name, value):
|
|
if value is not None:
|
|
if value != "":
|
|
span.set_attribute(name, value)
|
|
return
|
|
|
|
|
|
def with_tracer_wrapper(func):
|
|
"""Helper for providing tracer for wrapper functions."""
|
|
|
|
def _with_tracer(tracer, to_wrap):
|
|
def wrapper(wrapped, instance, args, kwargs):
|
|
# prevent double wrapping
|
|
if hasattr(wrapped, "__wrapped__"):
|
|
return wrapped(*args, **kwargs)
|
|
|
|
return func(tracer, to_wrap, wrapped, instance, args, kwargs)
|
|
|
|
return wrapper
|
|
|
|
return _with_tracer
|
|
|
|
|
|
def dont_throw(func):
|
|
"""
|
|
A decorator that wraps the passed in function and logs exceptions instead of throwing them.
|
|
|
|
@param func: The function to wrap
|
|
@return: The wrapper function
|
|
"""
|
|
# Obtain a logger specific to the function's module
|
|
logger = logging.getLogger(func.__module__)
|
|
|
|
def wrapper(*args, **kwargs):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except Exception as e:
|
|
logger.debug(
|
|
"OpenLLMetry failed to trace in %s, error: %s",
|
|
func.__name__,
|
|
traceback.format_exc(),
|
|
)
|
|
if Config.exception_logger:
|
|
Config.exception_logger(e)
|
|
|
|
return wrapper
|