dedupe _get_input_attributes

Author: antonpirker

sentry_sdk/tracing_utils.py

@@ -995,46 +995,51 @@ def _get_input_attributes(template, send_pii, args, kwargs):
     """
     attributes = {}  # type: dict[str, Any]
 
-    if template in [SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_TOOL, SPANTEMPLATE.AI_CHAT]:
-        for key, value in list(kwargs.items()):
-            if key in ("model", "model_name") and isinstance(value, str):
-                attributes[SPANDATA.GEN_AI_REQUEST_MODEL] = value
+    mapping = {
+        "model": (SPANDATA.GEN_AI_REQUEST_MODEL, str),
+        "model_name": (SPANDATA.GEN_AI_REQUEST_MODEL, str),
+        "agent": (SPANDATA.GEN_AI_AGENT_NAME, str),
+        "agent_name": (SPANDATA.GEN_AI_AGENT_NAME, str),
+        "max_tokens": (SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, int),
+        "frequency_penalty": (SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, float),
+        "presence_penalty": (SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, float),
+        "temperature": (SPANDATA.GEN_AI_REQUEST_TEMPERATURE, float),
+        "top_p": (SPANDATA.GEN_AI_REQUEST_TOP_P, float),
+        "top_k": (SPANDATA.GEN_AI_REQUEST_TOP_K, int),
+    }
 
-            if key in ("agent", "agent_name") and isinstance(value, str):
-                attributes[SPANDATA.GEN_AI_AGENT_NAME] = value
+    def _set_from_key(key, value):
+        # type: (str, Any) -> None
+        if key in mapping:
+            (attribute, data_type) = mapping[key]
+            if value is not None and isinstance(value, data_type):
+                attributes[attribute] = value
 
+    if template in [SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_TOOL, SPANTEMPLATE.AI_CHAT]:
+        for key, value in list(kwargs.items()):
             if key == "prompt" and isinstance(value, str):
                 attributes.setdefault(SPANDATA.GEN_AI_REQUEST_MESSAGES, []).append(
                     {"role": "user", "content": value}
                 )
+                continue
+
             if key == "system_prompt" and isinstance(value, str):
                 attributes.setdefault(SPANDATA.GEN_AI_REQUEST_MESSAGES, []).append(
                     {"role": "system", "content": value}
                 )
+                continue
 
-            if key == "max_tokens" and isinstance(value, int):
-                attributes[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] = value
-            if key == "frequency_penalty" and isinstance(value, float):
-                attributes[SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] = value
-            if key == "presence_penalty" and isinstance(value, float):
-                attributes[SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] = value
-            if key == "temperature" and isinstance(value, float):
-                attributes[SPANDATA.GEN_AI_REQUEST_TEMPERATURE] = value
-            if key == "top_p" and isinstance(value, float):
-                attributes[SPANDATA.GEN_AI_REQUEST_TOP_P] = value
-            if key == "top_k" and isinstance(value, int):
-                attributes[SPANDATA.GEN_AI_REQUEST_TOP_K] = value
+            _set_from_key(key, value)
 
     if SPANDATA.GEN_AI_REQUEST_MESSAGES in attributes:
         attributes[SPANDATA.GEN_AI_REQUEST_MESSAGES] = safe_repr(
             attributes[SPANDATA.GEN_AI_REQUEST_MESSAGES]
         )
 
     if template == SPANTEMPLATE.AI_TOOL and send_pii:
-        if send_pii:
-            attributes[SPANDATA.GEN_AI_TOOL_INPUT] = safe_repr(
-                {"args": args, "kwargs": kwargs}
-            )
+        attributes[SPANDATA.GEN_AI_TOOL_INPUT] = safe_repr(
+            {"args": args, "kwargs": kwargs}
+        )
 
     return attributes
 
@@ -1059,12 +1064,12 @@ def _get_value(source, key):
                     value = None
         return value if isinstance(value, int) else None
 
-    def _set_from_keys(target_key, keys):
+    def _set_from_keys(attribute, keys):
         # type: (str, tuple[str, ...]) -> None
         for key in keys:
             value = _get_value(usage, key)
             if value is not None:
-                attributes[target_key] = value
+                attributes[attribute] = value
 
     _set_from_keys(
         SPANDATA.GEN_AI_USAGE_INPUT_TOKENS,