merge

Author: alexander-alderman-webb

sentry_sdk/integrations/openai.py

@@ -1,5 +1,6 @@
 import sys
 from functools import wraps
+from collections.abc import Iterable
 
 import sentry_sdk
 from sentry_sdk import consts
@@ -32,7 +33,6 @@
         AsyncIterator,
         Iterator,
         Union,
-        Iterable,
     )
     from sentry_sdk.tracing import Span
     from sentry_sdk._types import TextPart
@@ -206,37 +206,31 @@ def _is_system_instruction_completions(message: "ChatCompletionMessageParam") ->
 
 def _get_system_instructions_completions(
     messages: "Iterable[ChatCompletionMessageParam]",
-) -> "list[ChatCompletionSystemMessageParam]":
-    system_instructions = []
-
-    for message in messages:
-        if _is_system_instruction_completions(message):
-            system_instructions.append(message)
+) -> "list[ChatCompletionMessageParam]":
+    if not isinstance(messages, Iterable):
+        return []
 
-    return system_instructions
+    return [
+        message for message in messages if _is_system_instruction_completions(message)
+    ]
 
 
 def _is_system_instruction_responses(message: "ResponseInputItemParam") -> bool:
-    return (
-        isinstance(message, dict)
-        and message.get("type") == "message"
-        and message.get("role") == "system"
-    )
+    if not isinstance(message, dict) or not message.get("role") == "system":
+        return False
+
+    return "type" not in message or message["type"] == "message"
 
 
 def _get_system_instructions_responses(
     messages: "Union[str, ResponseInputParam]",
 ) -> "list[ResponseInputItemParam]":
-    if isinstance(messages, str):
+    if not isinstance(messages, list):
         return []
 
-    system_instructions = []
-
-    for message in messages:
-        if _is_system_instruction_responses(message):
-            system_instructions.append(message)
-
-    return system_instructions
+    return [
+        message for message in messages if _is_system_instruction_responses(message)
+    ]
 
 
 def _transform_system_instructions(
@@ -313,47 +307,63 @@ def _set_responses_api_input_data(
     kwargs: "dict[str, Any]",
     integration: "OpenAIIntegration",
 ) -> None:
+    explicit_instructions: "Union[Optional[str], Omit]" = kwargs.get("instructions")
     messages: "Optional[Union[str, ResponseInputParam]]" = kwargs.get("input")
 
-    if messages is None:
+    if not should_send_default_pii() or not integration.include_prompts:
         set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
         _commmon_set_input_data(span, kwargs)
         return
 
-    explicit_instructions: "Union[Optional[str], Omit]" = kwargs.get("instructions")
-    system_instructions = _get_system_instructions_responses(messages)
     if (
-        (
-            (explicit_instructions is not None and _is_given(explicit_instructions))
-            or len(system_instructions) > 0
-        )
-        and should_send_default_pii()
-        and integration.include_prompts
+        messages is None
+        and explicit_instructions is not None
+        and _is_given(explicit_instructions)
     ):
-        instructions_text_parts: "list[TextPart]" = []
-        if explicit_instructions is not None and _is_given(explicit_instructions):
-            instructions_text_parts.append(
+        set_data_normalized(
+            span,
+            SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
+            [
                 {
                     "type": "text",
                     "content": explicit_instructions,
                 }
-            )
-        # Deliberate use of function accepting completions API type because
-        # of shared structure FOR THIS PURPOSE ONLY.
-        instructions_text_parts += _transform_system_instructions(system_instructions)
+            ],
+            unpack=False,
+        )
+
+        set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
+        _commmon_set_input_data(span, kwargs)
+        return
 
+    if messages is None:
+        set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
+        _commmon_set_input_data(span, kwargs)
+        return
+
+    instructions_text_parts: "list[TextPart]" = []
+    if explicit_instructions is not None and _is_given(explicit_instructions):
+        instructions_text_parts.append(
+            {
+                "type": "text",
+                "content": explicit_instructions,
+            }
+        )
+
+    system_instructions = _get_system_instructions_responses(messages)
+    # Deliberate use of function accepting completions API type because
+    # of shared structure FOR THIS PURPOSE ONLY.
+    instructions_text_parts += _transform_system_instructions(system_instructions)
+
+    if len(instructions_text_parts) > 0:
         set_data_normalized(
             span,
             SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
             instructions_text_parts,
             unpack=False,
         )
 
-    if (
-        isinstance(messages, str)
-        and should_send_default_pii()
-        and integration.include_prompts
-    ):
+    if isinstance(messages, str):
         normalized_messages = normalize_message_roles([messages])  # type: ignore
         scope = sentry_sdk.get_current_scope()
         messages_data = truncate_and_annotate_messages(normalized_messages, span, scope)
@@ -362,22 +372,21 @@ def _set_responses_api_input_data(
                 span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
             )
 
-    elif should_send_default_pii() and integration.include_prompts:
-        non_system_messages = [
-            message
-            for message in messages
-            if not _is_system_instruction_responses(message)
-        ]
-        if len(non_system_messages) > 0:
-            normalized_messages = normalize_message_roles(non_system_messages)  # type: ignore
-            scope = sentry_sdk.get_current_scope()
-            messages_data = truncate_and_annotate_messages(
-                normalized_messages, span, scope
+        set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
+        _commmon_set_input_data(span, kwargs)
+        return
+
+    non_system_messages = [
+        message for message in messages if not _is_system_instruction_responses(message)
+    ]
+    if len(non_system_messages) > 0:
+        normalized_messages = normalize_message_roles(non_system_messages)
+        scope = sentry_sdk.get_current_scope()
+        messages_data = truncate_and_annotate_messages(normalized_messages, span, scope)
+        if messages_data is not None:
+            set_data_normalized(
+                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
             )
-            if messages_data is not None:
-                set_data_normalized(
-                    span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
-                )
 
     set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
     _commmon_set_input_data(span, kwargs)
@@ -392,52 +401,50 @@ def _set_completions_api_input_data(
         "messages"
     )
 
+    if not should_send_default_pii() or not integration.include_prompts:
+        set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses")
+        _commmon_set_input_data(span, kwargs)
+        return
+
     if messages is None:
         set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
         _commmon_set_input_data(span, kwargs)
         return
 
     system_instructions = _get_system_instructions_completions(messages)
-    if (
-        len(system_instructions) > 0
-        and should_send_default_pii()
-        and integration.include_prompts
-    ):
+    if len(system_instructions) > 0:
         set_data_normalized(
             span,
             SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
             _transform_system_instructions(system_instructions),
             unpack=False,
         )
 
-    if (
-        isinstance(messages, str)
-        and should_send_default_pii()
-        and integration.include_prompts
-    ):
+    if isinstance(messages, str):
         normalized_messages = normalize_message_roles([messages])  # type: ignore
         scope = sentry_sdk.get_current_scope()
         messages_data = truncate_and_annotate_messages(normalized_messages, span, scope)
         if messages_data is not None:
             set_data_normalized(
                 span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
             )
-    elif should_send_default_pii() and integration.include_prompts:
-        non_system_messages = [
-            message
-            for message in messages
-            if not _is_system_instruction_completions(message)
-        ]
-        if len(non_system_messages) > 0:
-            normalized_messages = normalize_message_roles(non_system_messages)  # type: ignore
-            scope = sentry_sdk.get_current_scope()
-            messages_data = truncate_and_annotate_messages(
-                normalized_messages, span, scope
+        set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+        _commmon_set_input_data(span, kwargs)
+        return
+
+    non_system_messages = [
+        message
+        for message in messages
+        if not _is_system_instruction_completions(message)
+    ]
+    if len(non_system_messages) > 0:
+        normalized_messages = normalize_message_roles(non_system_messages)  # type: ignore
+        scope = sentry_sdk.get_current_scope()
+        messages_data = truncate_and_annotate_messages(normalized_messages, span, scope)
+        if messages_data is not None:
+            set_data_normalized(
+                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
             )
-            if messages_data is not None:
-                set_data_normalized(
-                    span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
-                )
 
     set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
     _commmon_set_input_data(span, kwargs)