Skip to content

Commit

Permalink
Merge branch 'main' into zachg/fix_kinesis_none_return_
Browse files Browse the repository at this point in the history
  • Loading branch information
ZStriker19 authored Sep 19, 2024
2 parents ab100eb + 58049cb commit 915b0f1
Show file tree
Hide file tree
Showing 3 changed files with 202 additions and 4 deletions.
9 changes: 5 additions & 4 deletions ddtrace/llmobs/_integrations/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,17 +59,18 @@ def __init__(self, integration_config: IntegrationConfig) -> None:

@property
def metrics_enabled(self) -> bool:
"""Return whether submitting metrics is enabled for this integration, or global config if not set."""
env_metrics_enabled = asbool(os.getenv("DD_{}_METRICS_ENABLED".format(self._integration_name.upper())))
if not env_metrics_enabled and config._llmobs_agentless_enabled:
"""
Return whether submitting metrics is enabled for this integration. Agentless mode disables submitting metrics.
"""
if config._llmobs_agentless_enabled:
return False
if hasattr(self.integration_config, "metrics_enabled"):
return asbool(self.integration_config.metrics_enabled)
return False

@property
def logs_enabled(self) -> bool:
"""Return whether submitting logs is enabled for this integration, or global config if not set."""
"""Return whether submitting logs is enabled for this integration."""
if hasattr(self.integration_config, "logs_enabled"):
return asbool(self.integration_config.logs_enabled)
return False
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
fixes:
- |
LLM Observability: Fixes an issue where the OpenAI and LangChain integrations would still submit integration metrics even in agentless mode.
Integration metrics are now disabled if using agentless mode via ``LLMObs.enable(agentless_enabled=True)`` or setting ``DD_LLMOBS_AGENTLESS_ENABLED=1``.
192 changes: 192 additions & 0 deletions tests/llmobs/test_llmobs_integrations.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,192 @@
import mock
import pytest

from ddtrace.llmobs._integrations import BaseLLMIntegration
from tests.utils import DummyTracer


@pytest.fixture(scope="function")
def mock_integration_config(ddtrace_global_config):
mock_config = mock.Mock()
mock_config.metrics_enabled = True
mock_config.span_char_limit = 10
mock_config.span_prompt_completion_sample_rate = 1.0
mock_config.log_prompt_completion_sample_rate = 1.0
yield mock_config


@pytest.fixture(scope="function")
def ddtrace_global_config():
with mock.patch("ddtrace.llmobs._integrations.base.config") as mock_global_config:
mock_global_config._llmobs_agentless_enabled = False
mock_global_config._llmobs_sample_rate = 1.0
mock_global_config._dd_api_key = "<not-a-real-key>"
yield mock_global_config


@pytest.fixture(scope="function")
def mock_pin():
mock_pin = mock.Mock()
mock_pin.tracer = DummyTracer()
mock_pin.service = "dummy_service"
yield mock_pin


def test_integration_truncate(mock_integration_config):
integration = BaseLLMIntegration(mock_integration_config)
assert integration.trunc("1" * 128) == "1111111111..."
assert integration.trunc("123") == "123"


def test_integration_metrics_enabled(mock_integration_config, ddtrace_global_config, monkeypatch):
"""Metrics should only ever be submitted if not in agentless mode, and if metrics env var is set to true."""
monkeypatch.setenv("DD_BASELLM_METRICS_ENABLED", "1")
integration = BaseLLMIntegration(mock_integration_config)
assert integration.metrics_enabled is True

monkeypatch.setenv("DD_BASELLM_METRICS_ENABLED", "1")
ddtrace_global_config._llmobs_agentless_enabled = True
integration = BaseLLMIntegration(mock_integration_config)
assert integration.metrics_enabled is False

monkeypatch.setenv("DD_BASELLM_METRICS_ENABLED", "0")
mock_integration_config.metrics_enabled = False
ddtrace_global_config._llmobs_agentless_enabled = True
integration = BaseLLMIntegration(mock_integration_config)
assert integration.metrics_enabled is False

monkeypatch.setenv("DD_BASELLM_METRICS_ENABLED", "0")
mock_integration_config.metrics_enabled = False
ddtrace_global_config._llmobs_agentless_enabled = False
integration = BaseLLMIntegration(mock_integration_config)
assert integration.metrics_enabled is False


def test_integration_logs_enabled(mock_integration_config):
mock_integration_config.logs_enabled = False
integration = BaseLLMIntegration(mock_integration_config)
assert integration.logs_enabled is False

mock_integration_config.logs_enabled = True
integration = BaseLLMIntegration(mock_integration_config)
assert integration.logs_enabled is True


@mock.patch("ddtrace.llmobs._integrations.base.LLMObs")
def test_integration_llmobs_enabled(mock_llmobs, mock_integration_config):
mock_llmobs.enabled = True
integration = BaseLLMIntegration(mock_integration_config)
assert integration.llmobs_enabled is True
mock_llmobs.enabled = False
integration = BaseLLMIntegration(mock_integration_config)
assert integration.llmobs_enabled is False


def test_pc_span_sampling(mock_integration_config, mock_pin):
integration = BaseLLMIntegration(mock_integration_config)
integration.pin = mock_pin
with mock_pin.tracer.trace("Dummy span") as mock_span:
assert integration.is_pc_sampled_span(mock_span) is True

mock_integration_config.span_prompt_completion_sample_rate = 0.0
integration = BaseLLMIntegration(mock_integration_config)
integration.pin = mock_pin
with mock_pin.tracer.trace("Dummy span") as mock_span:
assert integration.is_pc_sampled_span(mock_span) is False


def test_pc_span_sampling_log(mock_integration_config, mock_pin):
mock_integration_config.logs_enabled = True
integration = BaseLLMIntegration(mock_integration_config)
integration.pin = mock_pin
with mock_pin.tracer.trace("Dummy span") as mock_span:
assert integration.is_pc_sampled_log(mock_span) is True

mock_integration_config.log_prompt_completion_sample_rate = 0.0
mock_integration_config.logs_enabled = False
integration = BaseLLMIntegration(mock_integration_config)
integration.pin = mock_pin
with mock_pin.tracer.trace("Dummy span") as mock_span:
assert integration.is_pc_sampled_log(mock_span) is False


@mock.patch("ddtrace.llmobs._integrations.base.LLMObs")
def test_pc_span_sampling_llmobs(mock_llmobs, mock_integration_config, mock_pin):
mock_llmobs.enabled = True
integration = BaseLLMIntegration(mock_integration_config)
integration.pin = mock_pin
with mock_pin.tracer.trace("Dummy span") as mock_span:
assert integration.is_pc_sampled_llmobs(mock_span) is True
mock_llmobs.enabled = False
integration = BaseLLMIntegration(mock_integration_config)
integration.pin = mock_pin
with mock_pin.tracer.trace("Dummy span") as mock_span:
assert integration.is_pc_sampled_llmobs(mock_span) is False


@mock.patch("ddtrace.llmobs._integrations.base.V2LogWriter")
def test_log_writer_started(mock_log_writer, mock_integration_config):
mock_integration_config.logs_enabled = True
_ = BaseLLMIntegration(mock_integration_config)
mock_log_writer().start.assert_called_once()

mock_log_writer.reset_mock()
mock_integration_config.logs_enabled = False
_ = BaseLLMIntegration(mock_integration_config)
mock_log_writer().start.assert_not_called()


@mock.patch("ddtrace.llmobs._integrations.base.V2LogWriter")
def test_log_enqueues_log_writer(mock_log_writer, mock_integration_config):
span = DummyTracer().trace("Dummy span", service="dummy_service")
mock_integration_config.logs_enabled = True
integration = BaseLLMIntegration(mock_integration_config)
integration.log(span, "level", "message", {"DummyKey": "DummyValue"})
mock_log_writer().enqueue.assert_called_once_with(
{
"timestamp": mock.ANY,
"message": "message",
"hostname": mock.ANY,
"ddsource": "baseLLM",
"service": "dummy_service",
"status": "level",
"ddtags": mock.ANY,
"dd.span_id": str(span.span_id),
"dd.trace_id": "{:x}".format(span.trace_id),
"DummyKey": "DummyValue",
}
)


def test_metric_calls_statsd_client(mock_integration_config, monkeypatch):
monkeypatch.setenv("DD_BASELLM_METRICS_ENABLED", "1")
span = DummyTracer().trace("Dummy span", service="dummy_service")
integration = BaseLLMIntegration(mock_integration_config)
mock_statsd = mock.Mock()
integration._statsd = mock_statsd

integration.metric(span, "dist", "name", 1, tags=["tag:value"])
mock_statsd.distribution.assert_called_once_with("name", 1, tags=["tag:value"])

integration.metric(span, "incr", "name", 2, tags=["tag:value"])
mock_statsd.increment.assert_called_once_with("name", 2, tags=["tag:value"])

integration.metric(span, "gauge", "name", 3, tags=["tag:value"])
mock_statsd.gauge.assert_called_once_with("name", 3, tags=["tag:value"])

with pytest.raises(ValueError):
integration.metric(span, "metric_type_does_not_exist", "name", 3, tags=["tag:value"])


def test_integration_trace(mock_integration_config, mock_pin):
integration = BaseLLMIntegration(mock_integration_config)
mock_set_base_span_tags = mock.Mock()
integration._set_base_span_tags = mock_set_base_span_tags
integration.pin = mock_pin
with integration.trace(mock_pin, "dummy_operation_id"):
pass
span = mock_pin.tracer.pop()
assert span is not None
assert span[0].resource == "dummy_operation_id"
assert span[0].service == "dummy_service"
mock_set_base_span_tags.assert_called_once()

0 comments on commit 915b0f1

Please sign in to comment.