34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
|
import logging
|
||
|
|
from opentelemetry import context as context_api
|
||
|
|
from opentelemetry.context import attach, set_value
|
||
|
|
from opentelemetry.instrumentation.utils import (
|
||
|
|
_SUPPRESS_INSTRUMENTATION_KEY,
|
||
|
|
)
|
||
|
|
from opentelemetry.instrumentation.haystack.utils import (
|
||
|
|
with_tracer_wrapper,
|
||
|
|
process_request,
|
||
|
|
process_response,
|
||
|
|
)
|
||
|
|
from opentelemetry.semconv_ai import SpanAttributes, TraceloopSpanKindValues
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
@with_tracer_wrapper
|
||
|
|
def wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
|
||
|
|
if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
|
||
|
|
return wrapped(*args, **kwargs)
|
||
|
|
name = "haystack_pipeline"
|
||
|
|
attach(set_value("workflow_name", name))
|
||
|
|
with tracer.start_as_current_span(f"{name}.workflow") as span:
|
||
|
|
span.set_attribute(
|
||
|
|
SpanAttributes.TRACELOOP_SPAN_KIND,
|
||
|
|
TraceloopSpanKindValues.WORKFLOW.value,
|
||
|
|
)
|
||
|
|
span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
|
||
|
|
process_request(span, args, kwargs)
|
||
|
|
response = wrapped(*args, **kwargs)
|
||
|
|
process_response(span, response)
|
||
|
|
|
||
|
|
return response
|