feat(llmobs): support in-code config for llmobs

ddtrace/llmobs/_llmobs.py

@@ -2,16 +2,19 @@
 import os
 from typing import Any
 from typing import Dict
+from typing import List
 from typing import Optional
 from typing import Union
 
 import ddtrace
 from ddtrace import Span
 from ddtrace import config
+from ddtrace import patch
 from ddtrace.ext import SpanTypes
 from ddtrace.internal import atexit
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.service import Service
+from ddtrace.internal.utils.formats import asbool
 from ddtrace.llmobs._constants import INPUT_DOCUMENTS
 from ddtrace.llmobs._constants import INPUT_MESSAGES
 from ddtrace.llmobs._constants import INPUT_PARAMETERS
@@ -44,6 +47,25 @@ class LLMObs(Service):
     _instance = None
     enabled = False
 
+    # Support llmobs python integrations
+    langchain = "langchain"
+    openai = "openai"
+    botocore = "botocore"
+
+    # APM enabled env var config:
+    _apm_env_config = {
+        "DD_LLMOBS_NO_APM": "0",
+    }
+
+    # APM disabled env var config:
+    _no_apm_env_config = {
+        "DD_INSTRUMENTATION_TELEMETRY_ENABLED": "0",
+        "DD_REMOTE_CONFIGURATION_ENABLED": "0",
+        "DD_OPENAI_METRICS_ENABLED": "0",
+        "DD_LANGCHAIN_METRICS_ENABLED": "0",
+        "DD_LLMOBS_NO_APM": "1",
+    }
+
     def __init__(self, tracer=None):
         super(LLMObs, self).__init__()
         self.tracer = tracer or ddtrace.tracer
@@ -75,30 +97,127 @@ def _stop_service(self) -> None:
             log.warning("Failed to shutdown tracer", exc_info=True)
 
     @classmethod
-    def enable(cls, tracer=None):
+    def _toggle(cls, switch: bool) -> None:
+        # _toggle is a helper function to ensure the state of config._llmobs_enabled is in sync with the class attribute
+        # since our integrations depend on config._llmobs_enabled to set llmobs data
+        config._llmobs_enabled = switch
+        cls.enabled = switch
+
+    @classmethod
+    def enable(
+        cls,
+        ml_app: Optional[str] = None,
+        integrations: Optional[List[str]] = None,
+        dd_apm_enabled=True,
+        dd_site: Optional[str] = None,
+        dd_api_key: Optional[str] = None,
+        dd_env: Optional[str] = None,
+        dd_service: Optional[str] = None,
+        tracer=None,
+    ):
+        """
+        Enable LLM Observability tracing.
+
+        :param str ml_app: The name of your ml application.
+        :param List[str] integrations: A list of integrations to enable auto-tracing for.
+        :param str dd_apm_enabled: Whether Datadog Application Performance Monitoring is enabled
+        :param str dd_site: Your datadog site, override by DD_SITE.
+        :param str dd_api_key: Your datadog api key.
+        :param str dd_env: Your environment name.
+        :param str dd_service: Your service name.
+        """
         if cls._instance is not None:
             log.debug("%s already enabled", cls.__name__)
             return
 
+        if os.getenv("DD_LLMOBS_ENABLED") and not asbool(os.getenv("DD_LLMOBS_ENABLED")):
+            LLMObs._toggle(False)
+            log.debug("LLMObs.enable() called when DD_LLMOBS_ENABLED is set to false, not starting LLMObs service")
+            return
+
+        # grab required values for LLMObs
+        # reading from environment variables is needed in case preload script wasn't called.
+        config._dd_site = config._dd_site or os.getenv("DD_SITE") or dd_site
+        config._dd_api_key = config._dd_api_key or os.getenv("DD_API_KEY") or dd_api_key
+        config._llmobs_ml_app = config._llmobs_ml_app or os.getenv("DD_LLMOBS_APP_NAME") or ml_app
+
+        # validate required values for LLMObs
         if not config._dd_api_key:
-            cls.enabled = False
+            LLMObs._toggle(False)
             raise ValueError(
                 "DD_API_KEY is required for sending LLMObs data. "
                 "Ensure this configuration is set before running your application."
             )
+        if not config._dd_site:
+            LLMObs._toggle(False)
+            raise ValueError(
+                "DD_SITE is required for sending LLMObs data. "
+                "Ensure this configuration is set before running your application."
+            )
         if not config._llmobs_ml_app:
-            cls.enabled = False
+            LLMObs._toggle(False)
             raise ValueError(
                 "DD_LLMOBS_APP_NAME is required for sending LLMObs data. "
                 "Ensure this configuration is set before running your application."
             )
 
+        # grab optional values
+        config.env = dd_env or config.env
+        config.service = dd_service or config.service
+
+        env_update = cls._apm_env_config if dd_apm_enabled else cls._no_apm_env_config
+        for k, v in env_update.items():
+            # don't override environment variables are are explicitly set.
+            if not os.getenv(k):
+                os.environ[k] = v
+                # update config object to ensure that the values are set correctly
+                if k == "DD_INSTRUMENTATION_TELEMETRY_ENABLED":
+                    config._telemetry_enabled = False
+                if k == "DD_REMOTE_CONFIGURATION_ENABLED":
+                    config._remote_config_enabled = False
+
+        # enable LLMObs integations
+        llmobs_integrations = {
+            LLMObs.langchain: lambda: LLMObs._patch_langchain(),
+            LLMObs.openai: lambda: LLMObs._patch_openai(),
+            LLMObs.botocore: lambda: LLMObs._patch_bedrock(),
+        }
+        if integrations:
+            for integration in integrations:
+                if integration in llmobs_integrations:
+                    llmobs_integrations[integration]()
+                else:
+                    log.warning(
+                        "%s is unsupported - LLMObs currently supports %s", integration, str(llmobs_integrations)
+                    )
+
         cls._instance = cls(tracer=tracer)
-        cls.enabled = True
+        LLMObs._toggle(True)
         cls._instance.start()
         atexit.register(cls.disable)
         log.debug("%s enabled", cls.__name__)
 
+    @classmethod
+    def _patch_langchain(cls) -> None:
+        if cls._instance is None:
+            log.warning("%s not enabled, cannot patch langchain", cls.__name__)
+            return
+        patch(langchain=True)
+
+    @classmethod
+    def _patch_openai(cls) -> None:
+        if cls._instance is None:
+            log.warning("%s not enabled, cannot patch openai", cls.__name__)
+            return
+        patch(openai=True)
+
+    @classmethod
+    def _patch_bedrock(cls) -> None:
+        if cls._instance is None:
+            log.warning("%s not enabled, cannot patch bedrock", cls.__name__)
+            return
+        patch(bedrock=True)
+
     @classmethod
     def disable(cls) -> None:
         if cls._instance is None:

tests/llmobs/test_llmobs_service.py

@@ -3,6 +3,8 @@
 import mock
 import pytest
 
+from ddtrace._trace.span import Span
+from ddtrace.ext import SpanTypes
 from ddtrace.llmobs import LLMObs as llmobs_service
 from ddtrace.llmobs._constants import INPUT_DOCUMENTS
 from ddtrace.llmobs._constants import INPUT_MESSAGES
@@ -36,6 +38,16 @@ def mock_logs():
         yield mock_logs
 
 
+def run_llmobs_trace_filter(dummy_tracer):
+    for trace_filter in dummy_tracer._filters:
+        if isinstance(trace_filter, LLMObsTraceProcessor):
+            root_llm_span = Span(name="span1", span_type=SpanTypes.LLM)
+            root_llm_span.set_tag_str(SPAN_KIND, "llm")
+            trace1 = [root_llm_span]
+            return trace_filter.process_trace(trace1)
+    raise ValueError("LLMObsTraceProcessor not found in tracer filters.")
+
+
 def test_service_enable():
     with override_global_config(dict(_dd_api_key="<not-a-real-api-key>", _llmobs_ml_app="<ml-app-name>")):
         dummy_tracer = DummyTracer()
@@ -45,6 +57,22 @@ def test_service_enable():
         assert llmobs_service.enabled
         assert llmobs_instance.tracer == dummy_tracer
         assert any(isinstance(tracer_filter, LLMObsTraceProcessor) for tracer_filter in dummy_tracer._filters)
+        assert run_llmobs_trace_filter(dummy_tracer) is not None
+
+        llmobs_service.disable()
+
+
+def test_service_enable_with_apm_disabled():
+    with override_global_config(dict(_dd_api_key="<not-a-real-api-key>", _llmobs_ml_app="<ml-app-name>")):
+        dummy_tracer = DummyTracer()
+        llmobs_service.enable(tracer=dummy_tracer, dd_apm_enabled=False)
+        llmobs_instance = llmobs_service._instance
+        assert llmobs_instance is not None
+        assert llmobs_service.enabled
+        assert llmobs_instance.tracer == dummy_tracer
+        assert any(isinstance(tracer_filter, LLMObsTraceProcessor) for tracer_filter in dummy_tracer._filters)
+        assert run_llmobs_trace_filter(dummy_tracer) is None
+
         llmobs_service.disable()