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

ddtrace/llmobs/_llmobs.py

@@ -8,6 +8,7 @@
 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
@@ -44,6 +45,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 +95,103 @@ def _stop_service(self) -> None:
             log.warning("Failed to shutdown tracer", exc_info=True)
 
     @classmethod
-    def enable(cls, tracer=None):
+    def enable(
+        cls,
+        ml_app: Optional[str] = None,
+        integrations: Optional[list[str]] = None,
+        apm_enabled=False,
+        dd_env: Optional[str] = None,
+        dd_service: Optional[str] = None,
+        dd_site: Optional[str] = None,
+        dd_api_key: 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_site Your datadog site.
+        :param str dd_api_key: Your datadog api key.
+        :param str dd_apm_enabled: Whether Datadog Application Performance Monitoring is enabled.
+        """
         if cls._instance is not None:
             log.debug("%s already enabled", cls.__name__)
             return
 
+        if dd_api_key:
+            config._dd_api_key = dd_api_key
+
         if not config._dd_api_key:
             cls.enabled = False
             raise ValueError(
                 "DD_API_KEY is required for sending LLMObs data. "
                 "Ensure this configuration is set before running your application."
             )
+
+        # update environment config based on APM enabled/disabled
+        os.environ.update(cls._apm_env_config if apm_enabled else cls._no_apm_env_config)
+
+        if ml_app:
+            config._llmobs_ml_app = ml_app
+        if dd_site:
+            config._dd_site = dd_site
+        if dd_api_key:
+            config._dd_api_key = dd_api_key
+        if dd_env:
+            config.env = dd_env
+        if dd_service:
+            config.service = dd_service
+
         if not config._llmobs_ml_app:
             cls.enabled = False
             raise ValueError(
                 "DD_LLMOBS_APP_NAME is required for sending LLMObs data. "
                 "Ensure this configuration is set before running your application."
             )
 
+        # 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
         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/conftest.py

@@ -69,6 +69,6 @@ def LLMObs(mock_llmobs_span_writer, mock_llmobs_eval_metric_writer, ddtrace_glob
     global_config.update(ddtrace_global_config)
     with override_global_config(global_config):
         dummy_tracer = DummyTracer()
-        llmobs_service.enable(tracer=dummy_tracer)
+        llmobs_service.enable(apm_enabled=True, tracer=dummy_tracer)
         yield llmobs_service
         llmobs_service.disable()

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 None
+
+        llmobs_service.disable()
+
+
+def test_service_enable_with_apm():
+    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, apm_enabled=True)
+        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 not None
+
         llmobs_service.disable()
 
 

tests/llmobs/test_llmobs_trace_processor.py

@@ -29,13 +29,14 @@ def mock_logs():
         yield mock_logs
 
 
-def test_processor_returns_all_traces_by_default(monkeypatch):
-    """Test that the LLMObsTraceProcessor returns all traces by default."""
+def test_processor_does_not_submit_apm_traces_if_no_apm_env_var_is_true(monkeypatch):
+    """Test that the LLMObsTraceProcessor does not submit APM traces by default."""
+    monkeypatch.setenv("DD_LLMOBS_NO_APM", "1")
     trace_filter = LLMObsTraceProcessor(llmobs_span_writer=mock.MagicMock())
     root_llm_span = Span(name="span1", span_type=SpanTypes.LLM)
     root_llm_span.set_tag_str(SPAN_KIND, "llm")
     trace1 = [root_llm_span]
-    assert trace_filter.process_trace(trace1) == trace1
+    assert trace_filter.process_trace(trace1) is None
 
 
 def test_processor_returns_all_traces_if_no_apm_env_var_is_false(monkeypatch):