chore(llmobs): refactor google utility functions (#14025)

[MLOB-2935]

This PR extracts google utils from burgeoning utils.py file and combines
relevant methods from gemini, vertex, and genai integrations into a
single google_utils file. In addition, it aims to clean or dedupe
similar methods.

Note: get_generation_config_google was deleted as it is not used
anywhere in the codebase.

testing:
exisiting tests properly capture this


## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)


[MLOB-2935]:
https://datadoghq.atlassian.net/browse/MLOB-2935?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

---------

Co-authored-by: Yun Kim <[email protected]>

Author: maxzhangdd

ddtrace/contrib/internal/google_genai/_utils.py

@@ -6,7 +6,7 @@
 
 import wrapt
 
-from ddtrace.llmobs._integrations.google_genai_utils import DEFAULT_MODEL_ROLE
+from ddtrace.llmobs._integrations.google_utils import GOOGLE_GENAI_DEFAULT_MODEL_ROLE
 from ddtrace.llmobs._utils import _get_attr
 
 
@@ -32,7 +32,7 @@ def _join_chunks(chunks: List[Any]) -> Optional[Dict[str, Any]]:
                     continue
 
                 if role is None:
-                    role = _get_attr(content, "role", DEFAULT_MODEL_ROLE)
+                    role = _get_attr(content, "role", GOOGLE_GENAI_DEFAULT_MODEL_ROLE)
 
                 parts = _get_attr(content, "parts", [])
                 for part in parts:

ddtrace/contrib/internal/google_genai/patch.py

@@ -9,7 +9,7 @@
 from ddtrace.contrib.internal.trace_utils import with_traced_module
 from ddtrace.contrib.internal.trace_utils import wrap
 from ddtrace.llmobs._integrations import GoogleGenAIIntegration
-from ddtrace.llmobs._integrations.google_genai_utils import extract_provider_and_model_name
+from ddtrace.llmobs._integrations.google_utils import extract_provider_and_model_name
 from ddtrace.trace import Pin
 
 
@@ -27,7 +27,7 @@ def get_version() -> str:
 @with_traced_module
 def traced_generate(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     with integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -46,7 +46,7 @@ def traced_generate(genai, pin, func, instance, args, kwargs):
 @with_traced_module
 async def traced_async_generate(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     with integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -65,7 +65,7 @@ async def traced_async_generate(genai, pin, func, instance, args, kwargs):
 @with_traced_module
 def traced_generate_stream(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -86,7 +86,7 @@ def traced_generate_stream(genai, pin, func, instance, args, kwargs):
 @with_traced_module
 async def traced_async_generate_stream(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -107,7 +107,7 @@ async def traced_async_generate_stream(genai, pin, func, instance, args, kwargs)
 @with_traced_module
 def traced_embed_content(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     with integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -126,7 +126,7 @@ def traced_embed_content(genai, pin, func, instance, args, kwargs):
 @with_traced_module
 async def traced_async_embed_content(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     with integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),

ddtrace/contrib/internal/google_generativeai/patch.py

@@ -11,7 +11,7 @@
 from ddtrace.contrib.internal.trace_utils import with_traced_module
 from ddtrace.contrib.internal.trace_utils import wrap
 from ddtrace.llmobs._integrations import GeminiIntegration
-from ddtrace.llmobs._integrations.utils import extract_model_name_google
+from ddtrace.llmobs._integrations.google_utils import extract_provider_and_model_name
 from ddtrace.trace import Pin
 
 
@@ -40,11 +40,12 @@ def traced_generate(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
     stream = kwargs.get("stream", False)
     generations = None
+    provider_name, model_name = extract_provider_and_model_name(instance=instance, model_name_attr="model_name")
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
-        provider="google",
-        model=extract_model_name_google(instance, "model_name"),
+        provider=provider_name,
+        model=model_name,
         submit_to_llmobs=True,
     )
     try:
@@ -68,11 +69,12 @@ async def traced_agenerate(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
     stream = kwargs.get("stream", False)
     generations = None
+    provider_name, model_name = extract_provider_and_model_name(instance=instance, model_name_attr="model_name")
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
-        provider="google",
-        model=extract_model_name_google(instance, "model_name"),
+        provider=provider_name,
+        model=model_name,
         submit_to_llmobs=True,
     )
     try:

ddtrace/contrib/internal/vertexai/patch.py

@@ -14,7 +14,7 @@
 from ddtrace.contrib.internal.vertexai._utils import TracedAsyncVertexAIStreamResponse
 from ddtrace.contrib.internal.vertexai._utils import TracedVertexAIStreamResponse
 from ddtrace.llmobs._integrations import VertexAIIntegration
-from ddtrace.llmobs._integrations.utils import extract_model_name_google
+from ddtrace.llmobs._integrations.google_utils import extract_provider_and_model_name
 from ddtrace.trace import Pin
 
 
@@ -60,11 +60,12 @@ def _traced_generate(vertexai, pin, func, instance, args, kwargs, model_instance
     integration = vertexai._datadog_integration
     stream = kwargs.get("stream", False)
     generations = None
+    provider_name, model_name = extract_provider_and_model_name(instance=model_instance, model_name_attr="_model_name")
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
-        provider="google",
-        model=extract_model_name_google(model_instance, "_model_name"),
+        provider=provider_name,
+        model=model_name,
         submit_to_llmobs=True,
     )
     # history must be copied since it is modified during the LLM interaction
@@ -92,11 +93,12 @@ async def _traced_agenerate(vertexai, pin, func, instance, args, kwargs, model_i
     integration = vertexai._datadog_integration
     stream = kwargs.get("stream", False)
     generations = None
+    provider_name, model_name = extract_provider_and_model_name(instance=model_instance, model_name_attr="_model_name")
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
-        provider="google",
-        model=extract_model_name_google(model_instance, "_model_name"),
+        provider=provider_name,
+        model=model_name,
         submit_to_llmobs=True,
     )
     # history must be copied since it is modified during the LLM interaction

ddtrace/llmobs/_integrations/gemini.py

@@ -16,9 +16,9 @@
 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 extract_message_from_part_google
-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._integrations.google_utils import extract_message_from_part_gemini_vertexai
+from ddtrace.llmobs._integrations.google_utils import get_system_instructions_gemini_vertexai
+from ddtrace.llmobs._integrations.google_utils import llmobs_get_metadata_gemini_vertexai
 from ddtrace.llmobs._utils import _get_attr
 from ddtrace.trace import Span
 
@@ -43,9 +43,9 @@ def _llmobs_set_tags(
         operation: str = "",
     ) -> None:
         instance = kwargs.get("instance", None)
-        metadata = llmobs_get_metadata_google(kwargs, instance)
+        metadata = llmobs_get_metadata_gemini_vertexai(kwargs, instance)
 
-        system_instruction = get_system_instructions_from_google_model(instance)
+        system_instruction = get_system_instructions_gemini_vertexai(instance)
         input_contents = get_argument_value(args, kwargs, 0, "contents")
         input_messages = self._extract_input_message(input_contents, system_instruction)
 
@@ -95,7 +95,7 @@ def _extract_input_message(self, contents, system_instruction=None):
                 messages.append(message)
                 continue
             for part in parts:
-                message = extract_message_from_part_google(part, role)
+                message = extract_message_from_part_gemini_vertexai(part, role)
                 messages.append(message)
         return messages
 
@@ -107,7 +107,7 @@ def _extract_output_message(self, generations):
             role = content.get("role", "model")
             parts = content.get("parts", [])
             for part in parts:
-                message = extract_message_from_part_google(part, role)
+                message = extract_message_from_part_gemini_vertexai(part, role)
                 output_messages.append(message)
         return output_messages
 

ddtrace/llmobs/_integrations/google_genai.py

@@ -14,12 +14,12 @@
 from ddtrace.llmobs._constants import OUTPUT_VALUE
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._integrations.base import BaseLLMIntegration
-from ddtrace.llmobs._integrations.google_genai_utils import DEFAULT_MODEL_ROLE
-from ddtrace.llmobs._integrations.google_genai_utils import extract_embedding_metrics_google_genai
-from ddtrace.llmobs._integrations.google_genai_utils import extract_generation_metrics_google_genai
-from ddtrace.llmobs._integrations.google_genai_utils import extract_message_from_part_google_genai
-from ddtrace.llmobs._integrations.google_genai_utils import extract_provider_and_model_name
-from ddtrace.llmobs._integrations.google_genai_utils import normalize_contents
+from ddtrace.llmobs._integrations.google_utils import GOOGLE_GENAI_DEFAULT_MODEL_ROLE
+from ddtrace.llmobs._integrations.google_utils import extract_embedding_metrics_google_genai
+from ddtrace.llmobs._integrations.google_utils import extract_generation_metrics_google_genai
+from ddtrace.llmobs._integrations.google_utils import extract_message_from_part_google_genai
+from ddtrace.llmobs._integrations.google_utils import extract_provider_and_model_name
+from ddtrace.llmobs._integrations.google_utils import normalize_contents_google_genai
 from ddtrace.llmobs._utils import _get_attr
 from ddtrace.llmobs.utils import Document
 
@@ -71,7 +71,7 @@ def _llmobs_set_tags(
         response: Optional[Any] = None,
         operation: str = "",
     ) -> None:
-        provider_name, model_name = extract_provider_and_model_name(kwargs)
+        provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
         span._set_ctx_items(
             {
                 SPAN_KIND: operation,
@@ -120,23 +120,23 @@ def _extract_input_messages(self, args: List[Any], kwargs: Dict[str, Any], confi
 
     def _extract_messages_from_contents(self, contents, default_role: str) -> List[Dict[str, Any]]:
         messages = []
-        for content in normalize_contents(contents):
+        for content in normalize_contents_google_genai(contents):
             role = content.get("role") or default_role
             for part in content.get("parts", []):
                 messages.append(extract_message_from_part_google_genai(part, role))
         return messages
 
     def _extract_output_messages(self, response) -> List[Dict[str, Any]]:
         if not response:
-            return [{"content": "", "role": DEFAULT_MODEL_ROLE}]
+            return [{"content": "", "role": GOOGLE_GENAI_DEFAULT_MODEL_ROLE}]
         messages = []
         candidates = _get_attr(response, "candidates", [])
         for candidate in candidates:
             content = _get_attr(candidate, "content", None)
             if not content:
                 continue
             parts = _get_attr(content, "parts", [])
-            role = _get_attr(content, "role", DEFAULT_MODEL_ROLE)
+            role = _get_attr(content, "role", GOOGLE_GENAI_DEFAULT_MODEL_ROLE)
             for part in parts:
                 message = extract_message_from_part_google_genai(part, role)
                 messages.append(message)