[MLOB-2035] create generic llmobs decorator

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)

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)

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")