diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index 49815151118..617b12297fe 100644 --- a/ddtrace/llmobs/_llmobs.py +++ b/ddtrace/llmobs/_llmobs.py @@ -572,6 +572,25 @@ def retrieval( log.warning(SPAN_START_WHILE_DISABLED_WARNING) return cls._instance._start_span("retrieval", name=name, session_id=session_id, ml_app=ml_app) + @classmethod + def trace( + cls, kind: Optional[str] = None, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None + ) -> Span: + """ + Trace any operation. Allows for span kind to be specified and otherwise defaults to empty span kind string. + + :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`. + + :returns: The Span object representing the traced operation. + """ + if cls.enabled is False: + log.warning(SPAN_START_WHILE_DISABLED_WARNING) + return cls._instance._start_span(kind or "", name=name, session_id=session_id, ml_app=ml_app) + @classmethod def annotate( cls, @@ -646,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 231d53d7626..1ecb4126c4b 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 7e61f9b4e18..8133a496c91 100644 --- a/ddtrace/llmobs/decorators.py +++ b/ddtrace/llmobs/decorators.py @@ -253,6 +253,40 @@ def wrapper(*args, **kwargs): return decorator +def _generic_decorator(operation_kind): + def decorator( + original_func: Optional[Callable] = None, + kind: Optional[str] = None, + model_name: Optional[str] = None, + model_provider: Optional[str] = None, + name: Optional[str] = None, + session_id: Optional[str] = None, + ml_app: Optional[str] = None, + _automatic_io_annotation: bool = True, + ): + if kind in MODEL_DECORATORS: + decorator = _model_decorator(kind) + return decorator( + original_func, + model_name, + model_provider, + name, + session_id, + ml_app, + ) + decorator = _llmobs_decorator(kind or operation_kind) + return decorator( + original_func, + name, + session_id, + ml_app, + _automatic_io_annotation, + ) + + return decorator + + +MODEL_DECORATORS = ["llm", "embedding"] llm = _model_decorator("llm") embedding = _model_decorator("embedding") @@ -261,3 +295,4 @@ def wrapper(*args, **kwargs): tool = _llmobs_decorator("tool") retrieval = _llmobs_decorator("retrieval") agent = _llmobs_decorator("agent") +wrap = _generic_decorator("trace")