diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index 3a977260edb..5204fa28bb1 100644 --- a/ddtrace/llmobs/_llmobs.py +++ b/ddtrace/llmobs/_llmobs.py @@ -123,6 +123,7 @@ from ddtrace.llmobs.types import _Meta from ddtrace.llmobs.types import _MetaIO from ddtrace.llmobs.types import _SpanField +from ddtrace.llmobs.types import _SpanLink from ddtrace.llmobs.utils import Documents from ddtrace.llmobs.utils import Messages from ddtrace.llmobs.utils import extract_tool_definitions @@ -2072,6 +2073,30 @@ def activate_distributed_headers(cls, request_headers: Dict[str, str]) -> None: context = HTTPPropagator.extract(request_headers) cls._instance.tracer.context_provider.activate(context) cls._instance._activate_llmobs_distributed_context(request_headers, context, _soft_fail=False) + + @classmethod + def link_span(cls, span: Span, exported_span: ExportedLLMObsSpan, attributes: Dict[str, str] = None) -> None: + """ + Links the exported_span to the provided active span, including the provided attributes + on the span link. Defaults to an output to input span link unless otherwise specified. + """ + if cls.enabled is False: + return + if not attributes: + attributes = {} + if not attributes.get("from"): + attributes["from"] = "output" + if not attributes.get("to"): + attributes["to"] = "input" + span_link = _SpanLink( + span_id=exported_span["span_id"], + trace_id=exported_span["trace_id"], + attributes=attributes, + ) + current_span_links = span._get_ctx_item(SPAN_LINKS) or [] + current_span_links.append(span_link) + span._set_ctx_item(SPAN_LINKS, current_span_links) + # initialize the default llmobs instance