feat(langchain): update integration to use gen_ai.* instead of ai.*

Author: shellmayr

sentry_sdk/ai/monitoring.py

@@ -12,7 +12,7 @@
 if TYPE_CHECKING:
     from typing import Optional, Callable, Any
 
-_ai_pipeline_name = ContextVar("ai_pipeline_name", default=None)
+_ai_pipeline_name = ContextVar(SPANDATA.GEN_AI_PIPELINE_NAME, default=None)
 
 
 def set_ai_pipeline_name(name):

sentry_sdk/consts.py

@@ -474,6 +474,11 @@ class SPANDATA:
     Example: "COMPLETE"
     """
 
+    GEN_AI_RESPONSE_FORMAT = "gen_ai.response.format"
+    """
+    For an AI model call, the format of the response
+    """
+
     GEN_AI_RESPONSE_ID = "gen_ai.response.id"
     """
     Unique identifier for the completion.
@@ -515,6 +520,11 @@ class SPANDATA:
     The frequency penalty parameter used to reduce repetitiveness of generated tokens.
     Example: 0.1
     """
+    GEN_AI_REQUEST_LOGIT_BIAS = "gen_ai.logit_bias"
+    """
+    The logit bias parameter used to control the model's response.
+    Example: {"12345": -100}
+    """
 
     GEN_AI_REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens"
     """
@@ -546,6 +556,12 @@ class SPANDATA:
     Example: "1234567890"
     """
 
+    GEN_AI_REQUEST_TAGS = "gen_ai.request.tags"
+    """
+    The tags passed to the model.
+    Example: {"tag1": "value1", "tag2": "value2"}
+    """
+
     GEN_AI_REQUEST_TEMPERATURE = "gen_ai.request.temperature"
     """
     The temperature parameter used to control randomness in the output.

sentry_sdk/integrations/langchain.py

@@ -32,17 +32,18 @@
 
 
 DATA_FIELDS = {
-    "temperature": SPANDATA.AI_TEMPERATURE,
-    "top_p": SPANDATA.AI_TOP_P,
-    "top_k": SPANDATA.AI_TOP_K,
-    "function_call": SPANDATA.AI_FUNCTION_CALL,
-    "tool_calls": SPANDATA.AI_TOOL_CALLS,
-    "tools": SPANDATA.AI_TOOLS,
-    "response_format": SPANDATA.AI_RESPONSE_FORMAT,
-    "logit_bias": SPANDATA.AI_LOGIT_BIAS,
-    "tags": SPANDATA.AI_TAGS,
+    "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE,
+    "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P,
+    "top_k": SPANDATA.GEN_AI_REQUEST_TOP_K,
+    "function_call": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
+    "tool_calls": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
+    "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
+    "response_format": SPANDATA.GEN_AI_RESPONSE_FORMAT,
+    "logit_bias": SPANDATA.GEN_AI_REQUEST_LOGIT_BIAS,
+    "tags": SPANDATA.GEN_AI_REQUEST_TAGS,
 }
 
+# TODO(shellmayr): is this still the case?
 # To avoid double collecting tokens, we do *not* measure
 # token counts for models for which we have an explicit integration
 NO_COLLECT_TOKEN_MODELS = [
@@ -191,7 +192,7 @@ def on_llm_start(
             )
             span = watched_span.span
             if should_send_default_pii() and self.include_prompts:
-                set_data_normalized(span, SPANDATA.AI_INPUT_MESSAGES, prompts)
+                set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompts)
             for k, v in DATA_FIELDS.items():
                 if k in all_params:
                     set_data_normalized(span, v, all_params[k])
@@ -222,11 +223,11 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
             if not model and "anthropic" in all_params.get("_type"):
                 model = "claude-2"
             if model:
-                span.set_data(SPANDATA.AI_MODEL_ID, model)
+                span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model)
             if should_send_default_pii() and self.include_prompts:
                 set_data_normalized(
                     span,
-                    SPANDATA.AI_INPUT_MESSAGES,
+                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
                     [
                         [self._normalize_langchain_message(x) for x in list_]
                         for list_ in messages
@@ -271,7 +272,7 @@ def on_llm_end(self, response, *, run_id, **kwargs):
             if should_send_default_pii() and self.include_prompts:
                 set_data_normalized(
                     span_data.span,
-                    SPANDATA.AI_RESPONSES,
+                    SPANDATA.GEN_AI_RESPONSE_TEXT,
                     [[x.text for x in list_] for list_ in response.generations],
                 )
 
@@ -317,7 +318,9 @@ def on_chain_start(self, serialized, inputs, *, run_id, **kwargs):
             )
             metadata = kwargs.get("metadata")
             if metadata:
-                set_data_normalized(watched_span.span, SPANDATA.AI_METADATA, metadata)
+                set_data_normalized(
+                    watched_span.span, SPANDATA.GEN_AI_REQUEST_METADATA, metadata
+                )
 
     def on_chain_end(self, outputs, *, run_id, **kwargs):
         # type: (SentryLangchainCallback, Dict[str, Any], UUID, Any) -> Any
@@ -350,7 +353,9 @@ def on_agent_action(self, action, *, run_id, **kwargs):
             )
             if action.tool_input and should_send_default_pii() and self.include_prompts:
                 set_data_normalized(
-                    watched_span.span, SPANDATA.AI_INPUT_MESSAGES, action.tool_input
+                    watched_span.span,
+                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
+                    action.tool_input,
                 )
 
     def on_agent_finish(self, finish, *, run_id, **kwargs):
@@ -364,7 +369,9 @@ def on_agent_finish(self, finish, *, run_id, **kwargs):
                 return
             if should_send_default_pii() and self.include_prompts:
                 set_data_normalized(
-                    span_data.span, SPANDATA.AI_RESPONSES, finish.return_values.items()
+                    span_data.span,
+                    SPANDATA.GEN_AI_RESPONSE_TEXT,
+                    finish.return_values.items(),
                 )
             self._exit_span(span_data, run_id)
 
@@ -384,12 +391,14 @@ def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
             if should_send_default_pii() and self.include_prompts:
                 set_data_normalized(
                     watched_span.span,
-                    SPANDATA.AI_INPUT_MESSAGES,
+                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
                     kwargs.get("inputs", [input_str]),
                 )
                 if kwargs.get("metadata"):
                     set_data_normalized(
-                        watched_span.span, SPANDATA.AI_METADATA, kwargs.get("metadata")
+                        watched_span.span,
+                        SPANDATA.GEN_AI_REQUEST_METADATA,
+                        kwargs.get("metadata"),
                     )
 
     def on_tool_end(self, output, *, run_id, **kwargs):
@@ -403,7 +412,9 @@ def on_tool_end(self, output, *, run_id, **kwargs):
             if not span_data:
                 return
             if should_send_default_pii() and self.include_prompts:
-                set_data_normalized(span_data.span, SPANDATA.AI_RESPONSES, output)
+                set_data_normalized(
+                    span_data.span, SPANDATA.GEN_AI_RESPONSE_TEXT, output
+                )
             self._exit_span(span_data, run_id)
 
     def on_tool_error(self, error, *args, run_id, **kwargs):

tests/integrations/langchain/test_langchain.py

@@ -196,22 +196,24 @@ def test_langchain_agent(
 
     if send_default_pii and include_prompts:
         assert (
-            "You are very powerful" in chat_spans[0]["data"][SPANDATA.AI_INPUT_MESSAGES]
+            "You are very powerful"
+            in chat_spans[0]["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]
         )
-        assert "5" in chat_spans[0]["data"][SPANDATA.AI_RESPONSES]
-        assert "word" in tool_exec_span["data"][SPANDATA.AI_INPUT_MESSAGES]
-        assert 5 == int(tool_exec_span["data"][SPANDATA.AI_RESPONSES])
+        assert "5" in chat_spans[0]["data"][SPANDATA.GEN_AI_RESPONSE_TEXT]
+        assert "word" in tool_exec_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]
+        assert 5 == int(tool_exec_span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT])
         assert (
-            "You are very powerful" in chat_spans[1]["data"][SPANDATA.AI_INPUT_MESSAGES]
+            "You are very powerful"
+            in chat_spans[1]["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]
         )
-        assert "5" in chat_spans[1]["data"][SPANDATA.AI_RESPONSES]
+        assert "5" in chat_spans[1]["data"][SPANDATA.GEN_AI_RESPONSE_TEXT]
     else:
-        assert SPANDATA.AI_INPUT_MESSAGES not in chat_spans[0].get("data", {})
-        assert SPANDATA.AI_RESPONSES not in chat_spans[0].get("data", {})
-        assert SPANDATA.AI_INPUT_MESSAGES not in chat_spans[1].get("data", {})
-        assert SPANDATA.AI_RESPONSES not in chat_spans[1].get("data", {})
-        assert SPANDATA.AI_INPUT_MESSAGES not in tool_exec_span.get("data", {})
-        assert SPANDATA.AI_RESPONSES not in tool_exec_span.get("data", {})
+        assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in chat_spans[0].get("data", {})
+        assert SPANDATA.GEN_AI_RESPONSE_TEXT not in chat_spans[0].get("data", {})
+        assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in chat_spans[1].get("data", {})
+        assert SPANDATA.GEN_AI_RESPONSE_TEXT not in chat_spans[1].get("data", {})
+        assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in tool_exec_span.get("data", {})
+        assert SPANDATA.GEN_AI_RESPONSE_TEXT not in tool_exec_span.get("data", {})
 
 
 def test_langchain_error(sentry_init, capture_events):