apply same fix for gemini

DataDog · Jan 22, 2025 · d124e97 · d124e97
1 parent 7841124
commit d124e97
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/ddtrace/contrib/internal/google_generativeai/_utils.py b/ddtrace/contrib/internal/google_generativeai/_utils.py
@@ -108,7 +108,7 @@ def tag_request(span, integration, instance, args, kwargs):
     """Tag the generation span with request details.
     Includes capturing generation configuration, system prompt, and messages.
     """
-    contents = get_argument_value(args, kwargs, 0, "contents")
+    contents = get_argument_value(args, kwargs, 0, "content")
     generation_config = get_generation_config_google(instance, kwargs)
     system_instruction = get_system_instructions_from_google_model(instance)
     stream = kwargs.get("stream", None)

diff --git a/ddtrace/llmobs/_integrations/gemini.py b/ddtrace/llmobs/_integrations/gemini.py
@@ -44,7 +44,7 @@ def _llmobs_set_tags(
         metadata = llmobs_get_metadata_google(kwargs, instance)
 
         system_instruction = get_system_instructions_from_google_model(instance)
-        input_contents = get_argument_value(args, kwargs, 0, "contents")
+        input_contents = get_argument_value(args, kwargs, 0, "content")
         input_messages = self._extract_input_message(input_contents, system_instruction)
 
         output_messages = [{"content": ""}]