From 2431caed75b2ab7ddfdfa36ae46f9884bf082bff Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Thu, 6 Feb 2025 17:18:43 -0500 Subject: [PATCH] gate recording instances --- ddtrace/llmobs/_integrations/base.py | 4 ++++ ddtrace/llmobs/_integrations/langchain.py | 8 +++----- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/ddtrace/llmobs/_integrations/base.py b/ddtrace/llmobs/_integrations/base.py index a098c899014..9a60f978f0c 100644 --- a/ddtrace/llmobs/_integrations/base.py +++ b/ddtrace/llmobs/_integrations/base.py @@ -65,6 +65,10 @@ def __init__(self, integration_config: IntegrationConfig) -> None: self.start_log_writer() self._llmobs_pc_sampler = RateSampler(sample_rate=config._llmobs_sample_rate) + @property + def span_linking_enabled(self) -> bool: + return asbool(os.getenv("_DD_LLMOBS_AUTO_SPAN_LINKING_ENABLED", "false")) + @property def metrics_enabled(self) -> bool: """ diff --git a/ddtrace/llmobs/_integrations/langchain.py b/ddtrace/llmobs/_integrations/langchain.py index ed50a658733..a585529283a 100644 --- a/ddtrace/llmobs/_integrations/langchain.py +++ b/ddtrace/llmobs/_integrations/langchain.py @@ -1,6 +1,5 @@ from collections import defaultdict import json -import os from typing import Any from typing import Dict from typing import List @@ -12,7 +11,6 @@ from ddtrace.internal.logger import get_logger from ddtrace.internal.utils import ArgumentError from ddtrace.internal.utils import get_argument_value -from ddtrace.internal.utils.formats import asbool from ddtrace.llmobs import LLMObs from ddtrace.llmobs._constants import INPUT_DOCUMENTS from ddtrace.llmobs._constants import INPUT_MESSAGES @@ -107,7 +105,7 @@ class LangChainIntegration(BaseLLMIntegration): """Maps spans to instances.""" def record_steps(self, instance, span): - if not self.llmobs_enabled: + if not self.llmobs_enabled or not self.span_linking_enabled: return steps = getattr(instance, "steps", []) @@ -117,7 +115,7 @@ def record_steps(self, instance, span): self.record_instance(instance, span) def record_instance(self, instance, span): - if not self.llmobs_enabled: + if not self.llmobs_enabled or not self.span_linking_enabled: return instance = _extract_bound(instance) @@ -140,7 +138,7 @@ def _llmobs_set_tags( log.warning("Unsupported operation : %s", operation) return - if asbool(os.getenv("_DD_LLMOBS_AUTO_SPAN_LINKING_ENABLED")): + if self.span_linking_enabled: self._set_links(span) model_provider = span.get_tag(PROVIDER)