fix(llmobs): fix content arg extraction for vertex ai integration [backport 2.18]

ddtrace/contrib/internal/vertexai/_utils.py

@@ -177,13 +177,13 @@ def _tag_request_content(span, integration, content, content_idx):
         tag_request_content_part_google("vertexai", span, integration, part, part_idx, content_idx)
 
 
-def tag_request(span, integration, instance, args, kwargs):
+def tag_request(span, integration, instance, args, kwargs, is_chat):
     """Tag the generation span with request details.
     Includes capturing generation configuration, system prompt, and messages.
     """
     # instance is either a chat session or a model itself
     model_instance = instance if isinstance(instance, GenerativeModel) else instance._model
-    contents = get_argument_value(args, kwargs, 0, "contents")
+    contents = get_argument_value(args, kwargs, 0, "content" if is_chat else "contents")
     history = _get_attr(instance, "_history", [])
     if history:
         if isinstance(contents, list):

ddtrace/contrib/internal/vertexai/patch.py

@@ -64,7 +64,7 @@ def _traced_generate(vertexai, pin, func, instance, args, kwargs, model_instance
     # history must be copied since it is modified during the LLM interaction
     history = getattr(instance, "history", [])[:]
     try:
-        tag_request(span, integration, instance, args, kwargs)
+        tag_request(span, integration, instance, args, kwargs, is_chat)
         generations = func(*args, **kwargs)
         if stream:
             return TracedVertexAIStreamResponse(
@@ -99,7 +99,7 @@ async def _traced_agenerate(vertexai, pin, func, instance, args, kwargs, model_i
     # history must be copied since it is modified during the LLM interaction
     history = getattr(instance, "history", [])[:]
     try:
-        tag_request(span, integration, instance, args, kwargs)
+        tag_request(span, integration, instance, args, kwargs, is_chat)
         generations = await func(*args, **kwargs)
         if stream:
             return TracedAsyncVertexAIStreamResponse(

ddtrace/llmobs/_integrations/vertexai.py

@@ -5,6 +5,7 @@
 from typing import Optional
 
 from ddtrace import Span
+from ddtrace.internal.utils import ArgumentError
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.llmobs._constants import INPUT_MESSAGES
 from ddtrace.llmobs._constants import METADATA
@@ -51,7 +52,11 @@ def _llmobs_set_tags(
         span.set_tag_str(METADATA, safe_json(metadata))
 
         system_instruction = get_system_instructions_from_google_model(instance)
-        input_contents = get_argument_value(args, kwargs, 0, "contents")
+        input_contents = None
+        try:
+            input_contents = get_argument_value(args, kwargs, 0, "content")
+        except ArgumentError:
+            input_contents = get_argument_value(args, kwargs, 0, "contents")
         input_messages = self._extract_input_message(input_contents, history, system_instruction)
         span.set_tag_str(INPUT_MESSAGES, safe_json(input_messages))
 

releasenotes/notes/fix-vertexai-content-extraction-b216207bd8192e5f.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    vertexai: Resolves an issue with ``chat.send_message()`` where the content keyword argument was not parsed correctly.

tests/contrib/vertexai/test_vertexai.py

@@ -42,7 +42,7 @@ def test_vertexai_completion(vertexai):
     llm = vertexai.generative_models.GenerativeModel("gemini-1.5-flash")
     llm._prediction_client.responses["generate_content"].append(_mock_completion_response(MOCK_COMPLETION_SIMPLE_1))
     llm.generate_content(
-        "Why do bears hibernate?",
+        contents="Why do bears hibernate?",
         generation_config=vertexai.generative_models.GenerationConfig(
             stop_sequences=["x"], max_output_tokens=30, temperature=1.0
         ),
@@ -118,7 +118,7 @@ def test_vertexai_completion_stream(vertexai):
         (_mock_completion_stream_chunk(chunk) for chunk in MOCK_COMPLETION_STREAM_CHUNKS)
     ]
     response = llm.generate_content(
-        "How big is the solar system?",
+        contents="How big is the solar system?",
         generation_config=vertexai.generative_models.GenerationConfig(
             stop_sequences=["x"], max_output_tokens=30, temperature=1.0
         ),
@@ -278,7 +278,7 @@ def test_vertexai_chat(vertexai):
     llm._prediction_client.responses["generate_content"].append(_mock_completion_response(MOCK_COMPLETION_SIMPLE_1))
     chat = llm.start_chat()
     chat.send_message(
-        "Why do bears hibernate?",
+        content="Why do bears hibernate?",
         generation_config=vertexai.generative_models.GenerationConfig(
             stop_sequences=["x"], max_output_tokens=30, temperature=1.0
         ),
@@ -371,7 +371,7 @@ def test_vertexai_chat_stream(vertexai):
     ]
     chat = llm.start_chat()
     response = chat.send_message(
-        "How big is the solar system?",
+        content="How big is the solar system?",
         generation_config=vertexai.generative_models.GenerationConfig(
             stop_sequences=["x"], max_output_tokens=30, temperature=1.0
         ),

tests/contrib/vertexai/test_vertexai_llmobs.py

@@ -21,7 +21,7 @@ def test_completion(self, vertexai, mock_llmobs_writer, mock_tracer):
         llm = vertexai.generative_models.GenerativeModel("gemini-1.5-flash")
         llm._prediction_client.responses["generate_content"].append(_mock_completion_response(MOCK_COMPLETION_SIMPLE_1))
         llm.generate_content(
-            "Why do bears hibernate?",
+            contents="Why do bears hibernate?",
             generation_config=vertexai.generative_models.GenerationConfig(
                 stop_sequences=["x"], max_output_tokens=30, temperature=1.0
             ),
@@ -126,7 +126,7 @@ def test_completion_stream(self, vertexai, mock_llmobs_writer, mock_tracer):
             (_mock_completion_stream_chunk(chunk) for chunk in MOCK_COMPLETION_STREAM_CHUNKS)
         ]
         response = llm.generate_content(
-            "How big is the solar system?",
+            contents="How big is the solar system?",
             generation_config=vertexai.generative_models.GenerationConfig(
                 stop_sequences=["x"], max_output_tokens=30, temperature=1.0
             ),
@@ -293,7 +293,7 @@ def test_chat(self, vertexai, mock_llmobs_writer, mock_tracer):
         llm._prediction_client.responses["generate_content"].append(_mock_completion_response(MOCK_COMPLETION_SIMPLE_1))
         chat = llm.start_chat()
         chat.send_message(
-            "Why do bears hibernate?",
+            content="Why do bears hibernate?",
             generation_config=vertexai.generative_models.GenerationConfig(
                 stop_sequences=["x"], max_output_tokens=30, temperature=1.0
             ),
@@ -389,7 +389,7 @@ def test_chat_stream(self, vertexai, mock_llmobs_writer, mock_tracer):
         ]
         chat = llm.start_chat()
         response = chat.send_message(
-            "How big is the solar system?",
+            content="How big is the solar system?",
             generation_config=vertexai.generative_models.GenerationConfig(
                 stop_sequences=["x"], max_output_tokens=30, temperature=1.0
             ),