generalize helper for extracting token metrics

DataDog · Jan 22, 2025 · 635935e · 635935e
1 parent 4611816
commit 635935e
Show file tree

Hide file tree

Showing 7 changed files with 41 additions and 41 deletions.
diff --git a/ddtrace/llmobs/_constants.py b/ddtrace/llmobs/_constants.py
@@ -32,6 +32,7 @@
 INPUT_TOKENS_METRIC_KEY = "input_tokens"
 OUTPUT_TOKENS_METRIC_KEY = "output_tokens"
 TOTAL_TOKENS_METRIC_KEY = "total_tokens"
+INTEGRATIONS_USING_INPUT_OUTPUT_TOKENS = {"anthropic"}
 
 EVP_PROXY_AGENT_BASE_PATH = "evp_proxy/v2"
 EVP_PROXY_AGENT_ENDPOINT = "{}/api/v2/llmobs".format(EVP_PROXY_AGENT_BASE_PATH)

diff --git a/ddtrace/llmobs/_integrations/anthropic.py b/ddtrace/llmobs/_integrations/anthropic.py
@@ -18,6 +18,7 @@
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._constants import TOTAL_TOKENS_METRIC_KEY
 from ddtrace.llmobs._integrations.base import BaseLLMIntegration
+from ddtrace.llmobs._integrations.utils import get_llmobs_metrics_tags
 from ddtrace.llmobs._utils import _get_attr
 
 
@@ -77,7 +78,7 @@ def _llmobs_set_tags(
                 INPUT_MESSAGES: input_messages,
                 METADATA: parameters,
                 OUTPUT_MESSAGES: output_messages,
-                METRICS: self._get_llmobs_metrics_tags(span),
+                METRICS: get_llmobs_metrics_tags("anthropic", span),
             }
         )
 
@@ -188,18 +189,3 @@ def record_usage(self, span: Span, usage: Dict[str, Any]) -> None:
             span.set_metric("anthropic.response.usage.output_tokens", output_tokens)
         if input_tokens is not None and output_tokens is not None:
             span.set_metric("anthropic.response.usage.total_tokens", input_tokens + output_tokens)
-
-    @staticmethod
-    def _get_llmobs_metrics_tags(span):
-        usage = {}
-        input_tokens = span.get_metric("anthropic.response.usage.input_tokens")
-        output_tokens = span.get_metric("anthropic.response.usage.output_tokens")
-        total_tokens = span.get_metric("anthropic.response.usage.total_tokens")
-
-        if input_tokens is not None:
-            usage[INPUT_TOKENS_METRIC_KEY] = input_tokens
-        if output_tokens is not None:
-            usage[OUTPUT_TOKENS_METRIC_KEY] = output_tokens
-        if total_tokens is not None:
-            usage[TOTAL_TOKENS_METRIC_KEY] = total_tokens
-        return usage
diff --git a/ddtrace/llmobs/_integrations/bedrock.py b/ddtrace/llmobs/_integrations/bedrock.py
@@ -18,6 +18,7 @@
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._constants import TOTAL_TOKENS_METRIC_KEY
 from ddtrace.llmobs._integrations import BaseLLMIntegration
+from ddtrace.llmobs._integrations.utils import get_llmobs_metrics_tags
 from ddtrace.llmobs._utils import _get_llmobs_parent_id
 
 
@@ -57,22 +58,11 @@ def _llmobs_set_tags(
                 MODEL_PROVIDER: span.get_tag("bedrock.request.model_provider") or "",
                 INPUT_MESSAGES: input_messages,
                 METADATA: parameters,
-                METRICS: self._llmobs_metrics(span, response),
+                METRICS: get_llmobs_metrics_tags("bedrock", span),
                 OUTPUT_MESSAGES: output_messages,
             }
         )
 
-    @staticmethod
-    def _llmobs_metrics(span: Span, response: Optional[Dict[str, Any]]) -> Dict[str, Any]:
-        metrics = {}
-        if response and response.get("text"):
-            prompt_tokens = int(span.get_tag("bedrock.usage.prompt_tokens") or 0)
-            completion_tokens = int(span.get_tag("bedrock.usage.completion_tokens") or 0)
-            metrics[INPUT_TOKENS_METRIC_KEY] = prompt_tokens
-            metrics[OUTPUT_TOKENS_METRIC_KEY] = completion_tokens
-            metrics[TOTAL_TOKENS_METRIC_KEY] = prompt_tokens + completion_tokens
-        return metrics
-
     @staticmethod
     def _extract_input_message(prompt):
         """Extract input messages from the stored prompt.

diff --git a/ddtrace/llmobs/_integrations/gemini.py b/ddtrace/llmobs/_integrations/gemini.py
@@ -15,7 +15,7 @@
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._integrations.base import BaseLLMIntegration
 from ddtrace.llmobs._integrations.utils import extract_message_from_part_google
-from ddtrace.llmobs._integrations.utils import get_llmobs_metrics_tags_google
+from ddtrace.llmobs._integrations.utils import get_llmobs_metrics_tags
 from ddtrace.llmobs._integrations.utils import get_system_instructions_from_google_model
 from ddtrace.llmobs._integrations.utils import llmobs_get_metadata_google
 from ddtrace.llmobs._utils import _get_attr
@@ -59,7 +59,7 @@ def _llmobs_set_tags(
                 METADATA: metadata,
                 INPUT_MESSAGES: input_messages,
                 OUTPUT_MESSAGES: output_messages,
-                METRICS: get_llmobs_metrics_tags_google("google_generativeai", span),
+                METRICS: get_llmobs_metrics_tags("google_generativeai", span),
             }
         )
 

diff --git a/ddtrace/llmobs/_integrations/openai.py b/ddtrace/llmobs/_integrations/openai.py
@@ -22,6 +22,7 @@
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._constants import TOTAL_TOKENS_METRIC_KEY
 from ddtrace.llmobs._integrations.base import BaseLLMIntegration
+from ddtrace.llmobs._integrations.utils import get_llmobs_metrics_tags
 from ddtrace.llmobs._utils import _get_attr
 from ddtrace.llmobs.utils import Document
 from ddtrace.trace import Pin
@@ -275,12 +276,4 @@ def _extract_llmobs_metrics_tags(span: Span, resp: Any) -> Dict[str, Any]:
                 OUTPUT_TOKENS_METRIC_KEY: completion_tokens,
                 TOTAL_TOKENS_METRIC_KEY: prompt_tokens + completion_tokens,
             }
-        prompt_tokens = span.get_metric("openai.response.usage.prompt_tokens")
-        completion_tokens = span.get_metric("openai.response.usage.completion_tokens")
-        if prompt_tokens is None or completion_tokens is None:
-            return {}
-        return {
-            INPUT_TOKENS_METRIC_KEY: prompt_tokens,
-            OUTPUT_TOKENS_METRIC_KEY: completion_tokens,
-            TOTAL_TOKENS_METRIC_KEY: prompt_tokens + completion_tokens,
-        }
+        return get_llmobs_metrics_tags("openai", span)
diff --git a/ddtrace/llmobs/_integrations/utils.py b/ddtrace/llmobs/_integrations/utils.py
@@ -4,6 +4,7 @@
 from ddtrace.llmobs._constants import INPUT_TOKENS_METRIC_KEY
 from ddtrace.llmobs._constants import OUTPUT_TOKENS_METRIC_KEY
 from ddtrace.llmobs._constants import TOTAL_TOKENS_METRIC_KEY
+from ddtrace.llmobs._constants import INTEGRATIONS_USING_INPUT_OUTPUT_TOKENS
 from ddtrace.llmobs._utils import _get_attr
 
 
@@ -132,6 +133,35 @@ def get_llmobs_metrics_tags_google(integration_name, span):
         usage[TOTAL_TOKENS_METRIC_KEY] = total_tokens
     return usage
 
+def get_llmobs_metrics_tags(integration_name, span):
+    usage = {}
+
+    # bedrock integration tags usage under meta instead of metrics
+    if integration_name == "bedrock":
+        input_tokens = int(span.get_tag("bedrock.usage.prompt_tokens") or 0)
+        output_tokens = int(span.get_tag("bedrock.usage.completion_tokens") or 0)
+        usage[INPUT_TOKENS_METRIC_KEY] = input_tokens
+        usage[OUTPUT_TOKENS_METRIC_KEY] = output_tokens
+        usage[TOTAL_TOKENS_METRIC_KEY] = input_tokens + output_tokens
+        return usage
+
+    prompt_tokens_name = "prompt_tokens"
+    completion_tokens_name = "completion_tokens"
+    if integration_name in INTEGRATIONS_USING_INPUT_OUTPUT_TOKENS:
+        prompt_tokens_name = "input_tokens"
+        completion_tokens_name = "output_tokens"
+    input_tokens = span.get_metric("%s.response.usage.%s" % (integration_name, prompt_tokens_name))
+    output_tokens = span.get_metric("%s.response.usage.%s" % (integration_name, completion_tokens_name))
+    total_tokens = span.get_metric("%s.response.usage.total_tokens" % integration_name)
+
+    if input_tokens is not None:
+        usage[INPUT_TOKENS_METRIC_KEY] = input_tokens
+    if output_tokens is not None:
+        usage[OUTPUT_TOKENS_METRIC_KEY] = output_tokens
+    if total_tokens is not None:
+        usage[TOTAL_TOKENS_METRIC_KEY] = total_tokens
+    return usage
+
 
 def get_system_instructions_from_google_model(model_instance):
     """

diff --git a/ddtrace/llmobs/_integrations/vertexai.py b/ddtrace/llmobs/_integrations/vertexai.py
@@ -15,7 +15,7 @@
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._integrations.base import BaseLLMIntegration
 from ddtrace.llmobs._integrations.utils import extract_message_from_part_google
-from ddtrace.llmobs._integrations.utils import get_llmobs_metrics_tags_google
+from ddtrace.llmobs._integrations.utils import get_llmobs_metrics_tags
 from ddtrace.llmobs._integrations.utils import get_system_instructions_from_google_model
 from ddtrace.llmobs._integrations.utils import llmobs_get_metadata_google
 from ddtrace.llmobs._utils import _get_attr
@@ -60,7 +60,7 @@ def _llmobs_set_tags(
                 METADATA: metadata,
                 INPUT_MESSAGES: input_messages,
                 OUTPUT_MESSAGES: output_messages,
-                METRICS: get_llmobs_metrics_tags_google("vertexai", span),
+                METRICS: get_llmobs_metrics_tags("vertexai", span),
             }
         )