response model from response

Author: shellmayr

sentry_sdk/ai/span_config.py

@@ -62,7 +62,9 @@ def set_request_span_data(span, kwargs, integration, config, span_data=None):
                 set_data_normalized(span, span_attr, value)
 
 
-def set_response_span_data(span, response, include_pii, response_config, collected_text=None):
+def set_response_span_data(
+    span, response, include_pii, response_config, collected_text=None
+):
     # type: (Span, Any, bool, Dict[str, Any], Optional[List[str]]) -> None
     """
     Set response span data from a declarative config.
@@ -81,9 +83,7 @@ def set_response_span_data(span, response, include_pii, response_config, collect
     if include_pii:
         pii_sources = response_config.get("pii_sources")
         if pii_sources:
-            set_span_data_from_sources(
-                span, response, pii_sources, require_truthy=True
-            )
+            set_span_data_from_sources(span, response, pii_sources, require_truthy=True)
         if collected_text:
             set_data_normalized(
                 span, SPANDATA.GEN_AI_RESPONSE_TEXT, ["".join(collected_text)]

sentry_sdk/integrations/cohere/v1.py

@@ -36,6 +36,7 @@
 except ImportError:
     _has_chat_types = False
 
+
 def _extract_response_text(response):
     # type: (Any) -> list[str] | None
     text = getattr(response, "text", None)
@@ -50,6 +51,7 @@ def _extract_response_text(response):
     "extract_messages": lambda kw: _extract_messages(kw),
     "response": {
         "sources": {
+            SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)],
             SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)],
             SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)],
         },

sentry_sdk/integrations/cohere/v2.py

@@ -79,6 +79,7 @@ def _extract_v2_response_text(response):
     "extract_messages": lambda kw: _extract_messages_v2(kw.get("messages", [])),
     "response": {
         "sources": {
+            SPANDATA.GEN_AI_RESPONSE_MODEL: [("model",)],
             SPANDATA.GEN_AI_RESPONSE_ID: [("id",)],
             SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)],
         },
@@ -147,7 +148,6 @@ def new_chat(*args, **kwargs):
                 span_data = {
                     SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming,
                     SPANDATA.GEN_AI_REQUEST_MODEL: model if model else None,
-                    SPANDATA.GEN_AI_RESPONSE_MODEL: model if model else None,
                 }
                 set_request_span_data(
                     span, kwargs, integration, COHERE_V2_CHAT_CONFIG, span_data

tests/integrations/cohere/test_cohere.py

@@ -280,6 +280,7 @@ def test_v2_nonstreaming_chat(
             200,
             json={
                 "id": "resp-123",
+                "model": "some-model",
                 "finish_reason": "COMPLETE",
                 "message": {
                     "role": "assistant",