Make use of SPANDATA consistent

sentry_sdk/ai/monitoring.py

@@ -1,6 +1,7 @@
 import inspect
 from functools import wraps
 
+from sentry_sdk.consts import SPANDATA
 import sentry_sdk.utils
 from sentry_sdk import start_span
 from sentry_sdk.tracing import Span
@@ -39,7 +40,7 @@ def sync_wrapped(*args, **kwargs):
                 for k, v in kwargs.pop("sentry_data", {}).items():
                     span.set_data(k, v)
                 if curr_pipeline:
-                    span.set_data("ai.pipeline.name", curr_pipeline)
+                    span.set_data(SPANDATA.AI_PIPELINE_NAME, curr_pipeline)
                     return f(*args, **kwargs)
                 else:
                     _ai_pipeline_name.set(description)
@@ -68,7 +69,7 @@ async def async_wrapped(*args, **kwargs):
                 for k, v in kwargs.pop("sentry_data", {}).items():
                     span.set_data(k, v)
                 if curr_pipeline:
-                    span.set_data("ai.pipeline.name", curr_pipeline)
+                    span.set_data(SPANDATA.AI_PIPELINE_NAME, curr_pipeline)
                     return await f(*args, **kwargs)
                 else:
                     _ai_pipeline_name.set(description)
@@ -100,7 +101,7 @@ def record_token_usage(
     # type: (Span, Optional[int], Optional[int], Optional[int]) -> None
     ai_pipeline_name = get_ai_pipeline_name()
     if ai_pipeline_name:
-        span.set_data("ai.pipeline.name", ai_pipeline_name)
+        span.set_data(SPANDATA.AI_PIPELINE_NAME, ai_pipeline_name)
     if prompt_tokens is not None:
         span.set_measurement("ai_prompt_tokens_used", value=prompt_tokens)
     if completion_tokens is not None:

sentry_sdk/consts.py

@@ -187,7 +187,7 @@ class SPANDATA:
     For an AI model call, the format of the response
     """
 
-    AI_LOGIT_BIAS = "ai.response_format"
+    AI_LOGIT_BIAS = "ai.logit_bias"
     """
     For an AI model call, the logit bias
     """
@@ -204,7 +204,6 @@ class SPANDATA:
     Minimize pre-processing done to the prompt sent to the LLM.
     Example: true
     """
-
     AI_RESPONSES = "ai.responses"
     """
     The responses to an AI model call. Always as a list.
@@ -217,6 +216,66 @@ class SPANDATA:
     Example: 123.45
     """
 
+    AI_CITATIONS = "ai.citations"
+    """
+    References or sources cited by the AI model in its response.
+    Example: ["Smith et al. 2020", "Jones 2019"]
+    """
+
+    AI_DOCUMENTS = "ai.documents"
+    """
+    Documents or content chunks used as context for the AI model.
+    Example: ["doc1.txt", "doc2.pdf"]
+    """
+
+    AI_SEARCH_QUERIES = "ai.search_queries"
+    """
+    Queries used to search for relevant context or documents.
+    Example: ["climate change effects", "renewable energy"]
+    """
+
+    AI_SEARCH_RESULTS = "ai.search_results"
+    """
+    Results returned from search queries for context.
+    Example: ["Result 1", "Result 2"]
+    """
+
+    AI_GENERATION_ID = "ai.generation_id"
+    """
+    Unique identifier for the completion.
+    Example: "gen_123abc"
+    """
+
+    AI_SEARCH_REQUIRED = "ai.is_search_required"
+    """
+    Boolean indicating if the model needs to perform a search.
+    Example: true
+    """
+
+    AI_FINISH_REASON = "ai.finish_reason"
+    """
+    The reason why the model stopped generating.
+    Example: "length"
+    """
+
+    AI_PIPELINE_NAME = "ai.pipeline.name"
+    """
+    Name of the AI pipeline or chain being executed.
+    Example: "qa-pipeline"
+    """
+
+    AI_TEXTS = "ai.texts"
+    """
+    Raw text inputs provided to the model.
+    Example: ["What is machine learning?"]
+    """
+
+    AI_WARNINGS = "ai.warnings"
+    """
+    Warning messages generated during model execution.
+    Example: ["Token limit exceeded"]
+    """
+
     DB_NAME = "db.name"
     """
     The name of the database being accessed. For commands that switch the database, this should be set to the target database (even if the command fails).

sentry_sdk/integrations/cohere.py

@@ -52,17 +52,17 @@
 }
 
 COLLECTED_CHAT_RESP_ATTRS = {
-    "generation_id": "ai.generation_id",
-    "is_search_required": "ai.is_search_required",
-    "finish_reason": "ai.finish_reason",
+    "generation_id": SPANDATA.AI_GENERATION_ID,
+    "is_search_required": SPANDATA.AI_SEARCH_REQUIRED,
+    "finish_reason": SPANDATA.AI_FINISH_REASON,
 }
 
 COLLECTED_PII_CHAT_RESP_ATTRS = {
-    "citations": "ai.citations",
-    "documents": "ai.documents",
-    "search_queries": "ai.search_queries",
-    "search_results": "ai.search_results",
-    "tool_calls": "ai.tool_calls",
+    "citations": SPANDATA.AI_CITATIONS,
+    "documents": SPANDATA.AI_DOCUMENTS,
+    "search_queries": SPANDATA.AI_SEARCH_QUERIES,
+    "search_results": SPANDATA.AI_SEARCH_RESULTS,
+    "tool_calls": SPANDATA.AI_TOOL_CALLS,
 }
 
 
@@ -127,7 +127,7 @@ def collect_chat_response_fields(span, res, include_pii):
                 )
 
             if hasattr(res.meta, "warnings"):
-                set_data_normalized(span, "ai.warnings", res.meta.warnings)
+                set_data_normalized(span, SPANDATA.AI_WARNINGS, res.meta.warnings)
 
     @wraps(f)
     def new_chat(*args, **kwargs):
@@ -238,7 +238,7 @@ def new_embed(*args, **kwargs):
                 should_send_default_pii() and integration.include_prompts
             ):
                 if isinstance(kwargs["texts"], str):
-                    set_data_normalized(span, "ai.texts", [kwargs["texts"]])
+                    set_data_normalized(span, SPANDATA.AI_TEXTS, [kwargs["texts"]])
                 elif (
                     isinstance(kwargs["texts"], list)
                     and len(kwargs["texts"]) > 0

sentry_sdk/integrations/huggingface_hub.py

@@ -97,7 +97,7 @@ def new_text_generation(*args, **kwargs):
                 if should_send_default_pii() and integration.include_prompts:
                     set_data_normalized(
                         span,
-                        "ai.responses",
+                        SPANDATA.AI_RESPONSES,
                         [res],
                     )
                 span.__exit__(None, None, None)
@@ -107,7 +107,7 @@ def new_text_generation(*args, **kwargs):
                 if should_send_default_pii() and integration.include_prompts:
                     set_data_normalized(
                         span,
-                        "ai.responses",
+                        SPANDATA.AI_RESPONSES,
                         [res.generated_text],
                     )
                 if res.details is not None and res.details.generated_tokens > 0:

sentry_sdk/integrations/openai.py

@@ -155,7 +155,7 @@ def _new_chat_completion_common(f, *args, **kwargs):
             if should_send_default_pii() and integration.include_prompts:
                 set_data_normalized(
                     span,
-                    "ai.responses",
+                    SPANDATA.AI_RESPONSES,
                     list(map(lambda x: x.message, res.choices)),
                 )
             _calculate_chat_completion_usage(
@@ -329,15 +329,15 @@ def _new_embeddings_create_common(f, *args, **kwargs):
             should_send_default_pii() and integration.include_prompts
         ):
             if isinstance(kwargs["input"], str):
-                set_data_normalized(span, "ai.input_messages", [kwargs["input"]])
+                set_data_normalized(span, SPANDATA.AI_INPUT_MESSAGES, [kwargs["input"]])
             elif (
                 isinstance(kwargs["input"], list)
                 and len(kwargs["input"]) > 0
                 and isinstance(kwargs["input"][0], str)
             ):
-                set_data_normalized(span, "ai.input_messages", kwargs["input"])
+                set_data_normalized(span, SPANDATA.AI_INPUT_MESSAGES, kwargs["input"])
         if "model" in kwargs:
-            set_data_normalized(span, "ai.model_id", kwargs["model"])
+            set_data_normalized(span, SPANDATA.AI_MODEL_ID, kwargs["model"])
 
         response = yield f, args, kwargs
 

tests/integrations/anthropic/test_anthropic.py

@@ -128,7 +128,7 @@ def test_nonstreaming_create_message(
     assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 10
     assert span["measurements"]["ai_completion_tokens_used"]["value"] == 20
     assert span["measurements"]["ai_total_tokens_used"]["value"] == 30
-    assert span["data"]["ai.streaming"] is False
+    assert span["data"][SPANDATA.AI_STREAMING] is False
 
 
 @pytest.mark.asyncio
@@ -196,7 +196,7 @@ async def test_nonstreaming_create_message_async(
     assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 10
     assert span["measurements"]["ai_completion_tokens_used"]["value"] == 20
     assert span["measurements"]["ai_total_tokens_used"]["value"] == 30
-    assert span["data"]["ai.streaming"] is False
+    assert span["data"][SPANDATA.AI_STREAMING] is False
 
 
 @pytest.mark.parametrize(
@@ -296,7 +296,7 @@ def test_streaming_create_message(
     assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 10
     assert span["measurements"]["ai_completion_tokens_used"]["value"] == 30
     assert span["measurements"]["ai_total_tokens_used"]["value"] == 40
-    assert span["data"]["ai.streaming"] is True
+    assert span["data"][SPANDATA.AI_STREAMING] is True
 
 
 @pytest.mark.asyncio
@@ -399,7 +399,7 @@ async def test_streaming_create_message_async(
     assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 10
     assert span["measurements"]["ai_completion_tokens_used"]["value"] == 30
     assert span["measurements"]["ai_total_tokens_used"]["value"] == 40
-    assert span["data"]["ai.streaming"] is True
+    assert span["data"][SPANDATA.AI_STREAMING] is True
 
 
 @pytest.mark.skipif(
@@ -528,7 +528,7 @@ def test_streaming_create_message_with_input_json_delta(
     assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 366
     assert span["measurements"]["ai_completion_tokens_used"]["value"] == 51
     assert span["measurements"]["ai_total_tokens_used"]["value"] == 417
-    assert span["data"]["ai.streaming"] is True
+    assert span["data"][SPANDATA.AI_STREAMING] is True
 
 
 @pytest.mark.asyncio
@@ -665,7 +665,7 @@ async def test_streaming_create_message_with_input_json_delta_async(
     assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 366
     assert span["measurements"]["ai_completion_tokens_used"]["value"] == 51
     assert span["measurements"]["ai_total_tokens_used"]["value"] == 417
-    assert span["data"]["ai.streaming"] is True
+    assert span["data"][SPANDATA.AI_STREAMING] is True
 
 
 def test_exception_message_create(sentry_init, capture_events):
@@ -810,7 +810,7 @@ def test_add_ai_data_to_span_with_input_json_delta(sentry_init):
         assert span._data.get(SPANDATA.AI_RESPONSES) == [
             {"type": "text", "text": "{'test': 'data','more': 'json'}"}
         ]
-        assert span._data.get("ai.streaming") is True
+        assert span._data.get(SPANDATA.AI_STREAMING) is True
         assert span._measurements.get("ai_prompt_tokens_used")["value"] == 10
         assert span._measurements.get("ai_completion_tokens_used")["value"] == 20
         assert span._measurements.get("ai_total_tokens_used")["value"] == 30

tests/integrations/cohere/test_cohere.py

@@ -5,6 +5,7 @@
 from cohere import Client, ChatMessage
 
 from sentry_sdk import start_transaction
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.cohere import CohereIntegration
 
 from unittest import mock  # python 3.3 and above
@@ -53,15 +54,15 @@ def test_nonstreaming_chat(
     assert tx["type"] == "transaction"
     span = tx["spans"][0]
     assert span["op"] == "ai.chat_completions.create.cohere"
-    assert span["data"]["ai.model_id"] == "some-model"
+    assert span["data"][SPANDATA.AI_MODEL_ID] == "some-model"
 
     if send_default_pii and include_prompts:
-        assert "some context" in span["data"]["ai.input_messages"][0]["content"]
-        assert "hello" in span["data"]["ai.input_messages"][1]["content"]
-        assert "the model response" in span["data"]["ai.responses"]
+        assert "some context" in span["data"][SPANDATA.AI_INPUT_MESSAGES][0]["content"]
+        assert "hello" in span["data"][SPANDATA.AI_INPUT_MESSAGES][1]["content"]
+        assert "the model response" in span["data"][SPANDATA.AI_RESPONSES]
     else:
-        assert "ai.input_messages" not in span["data"]
-        assert "ai.responses" not in span["data"]
+        assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
+        assert SPANDATA.AI_RESPONSES not in span["data"]
 
     assert span["measurements"]["ai_completion_tokens_used"]["value"] == 10
     assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 20
@@ -124,15 +125,15 @@ def test_streaming_chat(sentry_init, capture_events, send_default_pii, include_p
     assert tx["type"] == "transaction"
     span = tx["spans"][0]
     assert span["op"] == "ai.chat_completions.create.cohere"
-    assert span["data"]["ai.model_id"] == "some-model"
+    assert span["data"][SPANDATA.AI_MODEL_ID] == "some-model"
 
     if send_default_pii and include_prompts:
-        assert "some context" in span["data"]["ai.input_messages"][0]["content"]
-        assert "hello" in span["data"]["ai.input_messages"][1]["content"]
-        assert "the model response" in span["data"]["ai.responses"]
+        assert "some context" in span["data"][SPANDATA.AI_INPUT_MESSAGES][0]["content"]
+        assert "hello" in span["data"][SPANDATA.AI_INPUT_MESSAGES][1]["content"]
+        assert "the model response" in span["data"][SPANDATA.AI_RESPONSES]
     else:
-        assert "ai.input_messages" not in span["data"]
-        assert "ai.responses" not in span["data"]
+        assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
+        assert SPANDATA.AI_RESPONSES not in span["data"]
 
     assert span["measurements"]["ai_completion_tokens_used"]["value"] == 10
     assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 20
@@ -194,9 +195,9 @@ def test_embed(sentry_init, capture_events, send_default_pii, include_prompts):
     span = tx["spans"][0]
     assert span["op"] == "ai.embeddings.create.cohere"
     if send_default_pii and include_prompts:
-        assert "hello" in span["data"]["ai.input_messages"]
+        assert "hello" in span["data"][SPANDATA.AI_INPUT_MESSAGES]
     else:
-        assert "ai.input_messages" not in span["data"]
+        assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
 
     assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 10
     assert span["measurements"]["ai_total_tokens_used"]["value"] == 10

tests/integrations/huggingface_hub/test_huggingface_hub.py

@@ -8,6 +8,7 @@
 from huggingface_hub.errors import OverloadedError
 
 from sentry_sdk import start_transaction
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.huggingface_hub import HuggingfaceHubIntegration
 
 
@@ -67,11 +68,11 @@ def test_nonstreaming_chat_completion(
     assert span["op"] == "ai.chat_completions.create.huggingface_hub"
 
     if send_default_pii and include_prompts:
-        assert "hello" in span["data"]["ai.input_messages"]
-        assert "the model response" in span["data"]["ai.responses"]
+        assert "hello" in span["data"][SPANDATA.AI_INPUT_MESSAGES]
+        assert "the model response" in span["data"][SPANDATA.AI_RESPONSES]
     else:
-        assert "ai.input_messages" not in span["data"]
-        assert "ai.responses" not in span["data"]
+        assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
+        assert SPANDATA.AI_RESPONSES not in span["data"]
 
     if details_arg:
         assert span["measurements"]["ai_total_tokens_used"]["value"] == 10
@@ -126,11 +127,11 @@ def test_streaming_chat_completion(
     assert span["op"] == "ai.chat_completions.create.huggingface_hub"
 
     if send_default_pii and include_prompts:
-        assert "hello" in span["data"]["ai.input_messages"]
-        assert "the model response" in span["data"]["ai.responses"]
+        assert "hello" in span["data"][SPANDATA.AI_INPUT_MESSAGES]
+        assert "the model response" in span["data"][SPANDATA.AI_RESPONSES]
     else:
-        assert "ai.input_messages" not in span["data"]
-        assert "ai.responses" not in span["data"]
+        assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
+        assert SPANDATA.AI_RESPONSES not in span["data"]
 
     if details_arg:
         assert span["measurements"]["ai_total_tokens_used"]["value"] == 10