Fix propagation parent ID for integration spans

ddtrace/contrib/openai/_endpoint_hooks.py

@@ -8,7 +8,6 @@
 from ddtrace.contrib.openai.utils import _loop_handler
 from ddtrace.contrib.openai.utils import _process_finished_stream
 from ddtrace.contrib.openai.utils import _tag_tool_calls
-from ddtrace.ext import SpanTypes
 from ddtrace.internal.utils.version import parse_version
 from ddtrace.llmobs._constants import SPAN_KIND
 
@@ -190,7 +189,6 @@ class _CompletionHook(_BaseCompletionHook):
 
     def _record_request(self, pin, integration, span, args, kwargs):
         super()._record_request(pin, integration, span, args, kwargs)
-        span.span_type = SpanTypes.LLM
         if integration.is_pc_sampled_llmobs(span):
             span.set_tag_str(SPAN_KIND, "llm")
         if integration.is_pc_sampled_span(span):
@@ -249,7 +247,6 @@ class _ChatCompletionHook(_BaseCompletionHook):
 
     def _record_request(self, pin, integration, span, args, kwargs):
         super()._record_request(pin, integration, span, args, kwargs)
-        span.span_type = SpanTypes.LLM
         if integration.is_pc_sampled_llmobs(span):
             span.set_tag_str(SPAN_KIND, "llm")
         for idx, m in enumerate(kwargs.get("messages", [])):

ddtrace/llmobs/_integrations/base.py

@@ -16,8 +16,11 @@
 from ddtrace.internal.dogstatsd import get_dogstatsd_client
 from ddtrace.internal.hostname import get_hostname
 from ddtrace.internal.utils.formats import asbool
+from ddtrace.llmobs._constants import PARENT_ID_KEY
+from ddtrace.llmobs._constants import PROPAGATED_PARENT_ID_KEY
 from ddtrace.llmobs._llmobs import LLMObs
 from ddtrace.llmobs._log_writer import V2LogWriter
+from ddtrace.llmobs._utils import _get_llmobs_parent_id
 from ddtrace.sampler import RateSampler
 from ddtrace.settings import IntegrationConfig
 
@@ -123,6 +126,12 @@ def trace(self, pin: Pin, operation_id: str, submit_to_llmobs: bool = False, **k
         self._set_base_span_tags(span, **kwargs)
         if submit_to_llmobs:
             span.span_type = SpanTypes.LLM
+            if span.get_tag(PROPAGATED_PARENT_ID_KEY) is None:
+                # For non-distributed traces or spans in the first service of a distributed trace,
+                # The LLMObs parent ID tag is not set at span start time. We need to manually set the parent ID tag now
+                # in these cases to avoid conflicting with the later propagated tags.
+                parent_id = _get_llmobs_parent_id(span) or "undefined"
+                span.set_tag_str(PARENT_ID_KEY, str(parent_id))
         return span
 
     @classmethod

ddtrace/llmobs/_integrations/openai.py

@@ -17,6 +17,7 @@
 from ddtrace.llmobs._constants import OUTPUT_MESSAGES
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._integrations.base import BaseLLMIntegration
+from ddtrace.pin import Pin
 
 
 class OpenAIIntegration(BaseLLMIntegration):
@@ -44,6 +45,11 @@ def user_api_key(self, value: str) -> None:
         # Match the API key representation that OpenAI uses in their UI.
         self._user_api_key = "sk-...%s" % value[-4:]
 
+    def trace(self, pin: Pin, operation_id: str, submit_to_llmobs: bool = False, **kwargs: Dict[str, Any]) -> Span:
+        if operation_id.endswith("Completion"):
+            submit_to_llmobs = True
+        return super().trace(pin, operation_id, submit_to_llmobs, **kwargs)
+
     def _set_base_span_tags(self, span: Span, **kwargs) -> None:
         span.set_tag_str(COMPONENT, self.integration_config.integration_name)
         if self._user_api_key is not None:

ddtrace/llmobs/_trace_processor.py

@@ -27,6 +27,7 @@
 from ddtrace.llmobs._constants import OUTPUT_DOCUMENTS
 from ddtrace.llmobs._constants import OUTPUT_MESSAGES
 from ddtrace.llmobs._constants import OUTPUT_VALUE
+from ddtrace.llmobs._constants import PARENT_ID_KEY
 from ddtrace.llmobs._constants import SESSION_ID
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._constants import TAGS
@@ -100,11 +101,12 @@ def _llmobs_span_event(self, span: Span) -> Dict[str, Any]:
         span.set_tag_str(ML_APP, ml_app)
         session_id = _get_session_id(span)
         span.set_tag_str(SESSION_ID, session_id)
-
+        parent_id = str(_get_llmobs_parent_id(span) or "undefined")
+        span._meta.pop(PARENT_ID_KEY, None)
         return {
             "trace_id": "{:x}".format(span.trace_id),
             "span_id": str(span.span_id),
-            "parent_id": str(_get_llmobs_parent_id(span) or "undefined"),
+            "parent_id": parent_id,
             "session_id": session_id,
             "name": _get_span_name(span),
             "tags": self._llmobs_tags(span, ml_app=ml_app, session_id=session_id),