diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index d41f2bd5b6..617b12297f 100644 --- a/ddtrace/llmobs/_llmobs.py +++ b/ddtrace/llmobs/_llmobs.py @@ -573,13 +573,14 @@ def retrieval( return cls._instance._start_span("retrieval", name=name, session_id=session_id, ml_app=ml_app) @classmethod - def undefined_kind( - cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None + def trace( + cls, kind: Optional[str] = None, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None ) -> Span: """ - Trace a generic operation. Used when the span kind is not specified, + Trace any operation. Allows for span kind to be specified and otherwise defaults to empty span kind string. - :param str name: The name of the traced operation. If not provided, a default value of "undefined" will be set. + :param str kind: The span kind of the traced operation. If not provided, the empty string will be used. + :param str name: The name of the traced operation. If not provided, the empty string will be used. :param str session_id: The ID of the underlying user session. Required for tracking sessions. :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default value will be set to the value of `DD_LLMOBS_ML_APP`. @@ -588,7 +589,7 @@ def undefined_kind( """ if cls.enabled is False: log.warning(SPAN_START_WHILE_DISABLED_WARNING) - return cls._instance._start_span("undefined", name=name, session_id=session_id, ml_app=ml_app) + return cls._instance._start_span(kind or "", name=name, session_id=session_id, ml_app=ml_app) @classmethod def annotate( @@ -664,9 +665,6 @@ def annotate( span.name = _name if prompt is not None: cls._tag_prompt(span, prompt) - if not span_kind: - log.debug("Span kind not specified, skipping annotation for input/output data") - return if input_data is not None or output_data is not None: if span_kind == "llm": cls._tag_llm_io(span, input_messages=input_data, output_messages=output_data) diff --git a/ddtrace/llmobs/_trace_processor.py b/ddtrace/llmobs/_trace_processor.py index 231d53d762..1ecb4126c4 100644 --- a/ddtrace/llmobs/_trace_processor.py +++ b/ddtrace/llmobs/_trace_processor.py @@ -76,9 +76,7 @@ def submit_llmobs_span(self, span: Span) -> None: def _llmobs_span_event(self, span: Span) -> Tuple[Dict[str, Any], bool]: """Span event object structure.""" - span_kind = span._get_ctx_item(SPAN_KIND) - if not span_kind: - raise KeyError("Span kind not found in span context") + span_kind = span._get_ctx_item(SPAN_KIND) or "" meta: Dict[str, Any] = {"span.kind": span_kind, "input": {}, "output": {}} if span_kind in ("llm", "embedding") and span._get_ctx_item(MODEL_NAME) is not None: meta["model_name"] = span._get_ctx_item(MODEL_NAME) diff --git a/ddtrace/llmobs/decorators.py b/ddtrace/llmobs/decorators.py index c91e4635b3..8133a496c9 100644 --- a/ddtrace/llmobs/decorators.py +++ b/ddtrace/llmobs/decorators.py @@ -295,4 +295,4 @@ def decorator( tool = _llmobs_decorator("tool") retrieval = _llmobs_decorator("retrieval") agent = _llmobs_decorator("agent") -observe = _generic_decorator("undefined_kind") +wrap = _generic_decorator("trace")