Skip to content

Commit 87a5c75

Browse files
authored
fix(llmobs): fix content arg extraction for vertex ai integration [backport #12034 to 2.19] (#12061)
Backports #12034 to 2.19. In [MLOS-42](https://datadoghq.atlassian.net/browse/MLOS-42) a customer was experiencing the following error: ``` ddtrace.internal.utils.ArgumentError: contents (at position 0) ``` where the `content` argument was not being extracted properly from the list of keyword arguments inputted into the `chat.send_message` method. This is because in the Vertex AI integration, we look for the `contents` keyword argument. However, the content field is titled `content` in the [send_message](https://github.com/google-gemini/generative-ai-python/blob/main/google/generativeai/generative_models.py#L514) method and `contents` in the [generate_content](https://github.com/google-gemini/generative-ai-python/blob/main/google/generativeai/generative_models.py#L239) method, so it is necessary to differentiate between these two cases. This PR is a small fix that corrects this error by differentiating between chat and completion requests in order to extract either `content` or `contents` respectively. ## Testing ### Automatic Testing I edited some of the currently existing tests to use the keyword argument extraction rather than the positional argument extraction to get the content in order to confirm that this fix resolves the error. ### Manual Testing Running the following code reproduced the error; furthermore, I confirmed that with this fix, the error is no longer present and the request completes successfully. ``` llm = GenerativeModel("gemini-1.5-flash") chat = llm.start_chat() resp = chat.send_message(content="hello") ``` I also verified that the following code which uses the generate_content method is not impacted (continues to work as before) as a result of this fix. ``` llm = GenerativeModel("gemini-1.5-flash") resp = llm.generate_content(contents="hello") ``` ## 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) [MLOS-42]: https://datadoghq.atlassian.net/browse/MLOS-42?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ --------- Co-authored-by: Yun Kim <[email protected]> (cherry picked from commit 67a5a9c) ## 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) [MLOS-42]: https://datadoghq.atlassian.net/browse/MLOS-42?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ [MLOS-42]: https://datadoghq.atlassian.net/browse/MLOS-42?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ
1 parent 88bc85b commit 87a5c75

File tree

6 files changed

+22
-13
lines changed

6 files changed

+22
-13
lines changed

ddtrace/contrib/internal/vertexai/_utils.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -177,13 +177,13 @@ def _tag_request_content(span, integration, content, content_idx):
177177
tag_request_content_part_google("vertexai", span, integration, part, part_idx, content_idx)
178178

179179

180-
def tag_request(span, integration, instance, args, kwargs):
180+
def tag_request(span, integration, instance, args, kwargs, is_chat):
181181
"""Tag the generation span with request details.
182182
Includes capturing generation configuration, system prompt, and messages.
183183
"""
184184
# instance is either a chat session or a model itself
185185
model_instance = instance if isinstance(instance, GenerativeModel) else instance._model
186-
contents = get_argument_value(args, kwargs, 0, "contents")
186+
contents = get_argument_value(args, kwargs, 0, "content" if is_chat else "contents")
187187
history = _get_attr(instance, "_history", [])
188188
if history:
189189
if isinstance(contents, list):

ddtrace/contrib/internal/vertexai/patch.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ def _traced_generate(vertexai, pin, func, instance, args, kwargs, model_instance
6464
# history must be copied since it is modified during the LLM interaction
6565
history = getattr(instance, "history", [])[:]
6666
try:
67-
tag_request(span, integration, instance, args, kwargs)
67+
tag_request(span, integration, instance, args, kwargs, is_chat)
6868
generations = func(*args, **kwargs)
6969
if stream:
7070
return TracedVertexAIStreamResponse(
@@ -99,7 +99,7 @@ async def _traced_agenerate(vertexai, pin, func, instance, args, kwargs, model_i
9999
# history must be copied since it is modified during the LLM interaction
100100
history = getattr(instance, "history", [])[:]
101101
try:
102-
tag_request(span, integration, instance, args, kwargs)
102+
tag_request(span, integration, instance, args, kwargs, is_chat)
103103
generations = await func(*args, **kwargs)
104104
if stream:
105105
return TracedAsyncVertexAIStreamResponse(

ddtrace/llmobs/_integrations/vertexai.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from typing import Optional
66

77
from ddtrace import Span
8+
from ddtrace.internal.utils import ArgumentError
89
from ddtrace.internal.utils import get_argument_value
910
from ddtrace.llmobs._constants import INPUT_MESSAGES
1011
from ddtrace.llmobs._constants import METADATA
@@ -45,7 +46,11 @@ def _llmobs_set_tags(
4546
metadata = llmobs_get_metadata_google(kwargs, instance)
4647

4748
system_instruction = get_system_instructions_from_google_model(instance)
48-
input_contents = get_argument_value(args, kwargs, 0, "contents")
49+
input_contents = None
50+
try:
51+
input_contents = get_argument_value(args, kwargs, 0, "content")
52+
except ArgumentError:
53+
input_contents = get_argument_value(args, kwargs, 0, "contents")
4954
input_messages = self._extract_input_message(input_contents, history, system_instruction)
5055

5156
output_messages = [{"content": ""}]
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
---
2+
fixes:
3+
- |
4+
vertexai: Resolves an issue with ``chat.send_message()`` where the content keyword argument was not parsed correctly.

tests/contrib/vertexai/test_vertexai.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ def test_vertexai_completion(vertexai):
4242
llm = vertexai.generative_models.GenerativeModel("gemini-1.5-flash")
4343
llm._prediction_client.responses["generate_content"].append(_mock_completion_response(MOCK_COMPLETION_SIMPLE_1))
4444
llm.generate_content(
45-
"Why do bears hibernate?",
45+
contents="Why do bears hibernate?",
4646
generation_config=vertexai.generative_models.GenerationConfig(
4747
stop_sequences=["x"], max_output_tokens=30, temperature=1.0
4848
),
@@ -118,7 +118,7 @@ def test_vertexai_completion_stream(vertexai):
118118
(_mock_completion_stream_chunk(chunk) for chunk in MOCK_COMPLETION_STREAM_CHUNKS)
119119
]
120120
response = llm.generate_content(
121-
"How big is the solar system?",
121+
contents="How big is the solar system?",
122122
generation_config=vertexai.generative_models.GenerationConfig(
123123
stop_sequences=["x"], max_output_tokens=30, temperature=1.0
124124
),
@@ -278,7 +278,7 @@ def test_vertexai_chat(vertexai):
278278
llm._prediction_client.responses["generate_content"].append(_mock_completion_response(MOCK_COMPLETION_SIMPLE_1))
279279
chat = llm.start_chat()
280280
chat.send_message(
281-
"Why do bears hibernate?",
281+
content="Why do bears hibernate?",
282282
generation_config=vertexai.generative_models.GenerationConfig(
283283
stop_sequences=["x"], max_output_tokens=30, temperature=1.0
284284
),
@@ -371,7 +371,7 @@ def test_vertexai_chat_stream(vertexai):
371371
]
372372
chat = llm.start_chat()
373373
response = chat.send_message(
374-
"How big is the solar system?",
374+
content="How big is the solar system?",
375375
generation_config=vertexai.generative_models.GenerationConfig(
376376
stop_sequences=["x"], max_output_tokens=30, temperature=1.0
377377
),

tests/contrib/vertexai/test_vertexai_llmobs.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ def test_completion(self, vertexai, mock_llmobs_writer, mock_tracer):
2121
llm = vertexai.generative_models.GenerativeModel("gemini-1.5-flash")
2222
llm._prediction_client.responses["generate_content"].append(_mock_completion_response(MOCK_COMPLETION_SIMPLE_1))
2323
llm.generate_content(
24-
"Why do bears hibernate?",
24+
contents="Why do bears hibernate?",
2525
generation_config=vertexai.generative_models.GenerationConfig(
2626
stop_sequences=["x"], max_output_tokens=30, temperature=1.0
2727
),
@@ -126,7 +126,7 @@ def test_completion_stream(self, vertexai, mock_llmobs_writer, mock_tracer):
126126
(_mock_completion_stream_chunk(chunk) for chunk in MOCK_COMPLETION_STREAM_CHUNKS)
127127
]
128128
response = llm.generate_content(
129-
"How big is the solar system?",
129+
contents="How big is the solar system?",
130130
generation_config=vertexai.generative_models.GenerationConfig(
131131
stop_sequences=["x"], max_output_tokens=30, temperature=1.0
132132
),
@@ -293,7 +293,7 @@ def test_chat(self, vertexai, mock_llmobs_writer, mock_tracer):
293293
llm._prediction_client.responses["generate_content"].append(_mock_completion_response(MOCK_COMPLETION_SIMPLE_1))
294294
chat = llm.start_chat()
295295
chat.send_message(
296-
"Why do bears hibernate?",
296+
content="Why do bears hibernate?",
297297
generation_config=vertexai.generative_models.GenerationConfig(
298298
stop_sequences=["x"], max_output_tokens=30, temperature=1.0
299299
),
@@ -389,7 +389,7 @@ def test_chat_stream(self, vertexai, mock_llmobs_writer, mock_tracer):
389389
]
390390
chat = llm.start_chat()
391391
response = chat.send_message(
392-
"How big is the solar system?",
392+
content="How big is the solar system?",
393393
generation_config=vertexai.generative_models.GenerationConfig(
394394
stop_sequences=["x"], max_output_tokens=30, temperature=1.0
395395
),

0 commit comments

Comments
 (0)