diff --git a/ddtrace/contrib/internal/vertexai/_utils.py b/ddtrace/contrib/internal/vertexai/_utils.py index 129b97fd920..81b0c13df2d 100644 --- a/ddtrace/contrib/internal/vertexai/_utils.py +++ b/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): diff --git a/ddtrace/contrib/internal/vertexai/patch.py b/ddtrace/contrib/internal/vertexai/patch.py index 2dbce060234..c441888573e 100644 --- a/ddtrace/contrib/internal/vertexai/patch.py +++ b/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( diff --git a/ddtrace/llmobs/_integrations/vertexai.py b/ddtrace/llmobs/_integrations/vertexai.py index 69fdc7eb665..88d653d6ff6 100644 --- a/ddtrace/llmobs/_integrations/vertexai.py +++ b/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)) diff --git a/releasenotes/notes/fix-vertexai-content-extraction-b216207bd8192e5f.yaml b/releasenotes/notes/fix-vertexai-content-extraction-b216207bd8192e5f.yaml new file mode 100644 index 00000000000..9bd0bba789b --- /dev/null +++ b/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. diff --git a/tests/contrib/vertexai/test_vertexai.py b/tests/contrib/vertexai/test_vertexai.py index 5f07c6e177f..afcbfea39ba 100644 --- a/tests/contrib/vertexai/test_vertexai.py +++ b/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 ), diff --git a/tests/contrib/vertexai/test_vertexai_llmobs.py b/tests/contrib/vertexai/test_vertexai_llmobs.py index 78a03bc664c..701f709e213 100644 --- a/tests/contrib/vertexai/test_vertexai_llmobs.py +++ b/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 ),