From c7fa2f016161d3839bf399847823d385bf5e2d3b Mon Sep 17 00:00:00 2001 From: Nicole Cybul Date: Fri, 3 Jan 2025 16:46:15 -0500 Subject: [PATCH 1/3] create generic llmobs decorator --- ddtrace/llmobs/_llmobs.py | 18 ++++++++++++++++++ ddtrace/llmobs/decorators.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index 49815151118..83de8a4a34d 100644 --- a/ddtrace/llmobs/_llmobs.py +++ b/ddtrace/llmobs/_llmobs.py @@ -572,6 +572,24 @@ 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 undefined_kind( + cls, 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, + + :param str name: The name of the traced operation. If not provided, a default value of "workflow" will be set. + :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("undefined", name=name, session_id=session_id, ml_app=ml_app) + @classmethod def annotate( cls, diff --git a/ddtrace/llmobs/decorators.py b/ddtrace/llmobs/decorators.py index 7e61f9b4e18..1ad0bbbcf6e 100644 --- a/ddtrace/llmobs/decorators.py +++ b/ddtrace/llmobs/decorators.py @@ -253,6 +253,41 @@ 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, + ) + else: + 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 +296,4 @@ def wrapper(*args, **kwargs): tool = _llmobs_decorator("tool") retrieval = _llmobs_decorator("retrieval") agent = _llmobs_decorator("agent") +observe = _generic_decorator("undefined_kind") From c2ea4bb018d7d16b6eeb6b7890e77535d54c39bd Mon Sep 17 00:00:00 2001 From: Nicole Cybul Date: Tue, 7 Jan 2025 11:49:02 -0500 Subject: [PATCH 2/3] small nits --- ddtrace/llmobs/_llmobs.py | 2 +- ddtrace/llmobs/decorators.py | 17 ++++++++--------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index 83de8a4a34d..d41f2bd5b6a 100644 --- a/ddtrace/llmobs/_llmobs.py +++ b/ddtrace/llmobs/_llmobs.py @@ -579,7 +579,7 @@ def undefined_kind( """ Trace a generic operation. Used when the span kind is not specified, - :param str name: The name of the traced operation. If not provided, a default value of "workflow" will be set. + :param str name: The name of the traced operation. If not provided, a default value of "undefined" will be set. :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`. diff --git a/ddtrace/llmobs/decorators.py b/ddtrace/llmobs/decorators.py index 1ad0bbbcf6e..c91e4635b30 100644 --- a/ddtrace/llmobs/decorators.py +++ b/ddtrace/llmobs/decorators.py @@ -274,15 +274,14 @@ def decorator( session_id, ml_app, ) - else: - decorator = _llmobs_decorator(kind or operation_kind) - return decorator( - original_func, - name, - session_id, - ml_app, - _automatic_io_annotation, - ) + decorator = _llmobs_decorator(kind or operation_kind) + return decorator( + original_func, + name, + session_id, + ml_app, + _automatic_io_annotation, + ) return decorator From 354b3897ceacaaa70486fce20bbdf34ab2a7a332 Mon Sep 17 00:00:00 2001 From: Nicole Cybul Date: Wed, 8 Jan 2025 14:01:10 -0500 Subject: [PATCH 3/3] change name ofdecorator and inline methods for undefined span kinds --- ddtrace/llmobs/_llmobs.py | 14 ++++++-------- ddtrace/llmobs/_trace_processor.py | 4 +--- ddtrace/llmobs/decorators.py | 2 +- 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index d41f2bd5b6a..617b12297fe 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 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 c91e4635b30..8133a496c91 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")