refacotr

Author: shellmayr

sentry_sdk/integrations/cohere/__init__.py

@@ -21,17 +21,6 @@
 except ImportError:
     raise DidNotEnable("Cohere not installed")
 
-COLLECTED_CHAT_PARAMS = {
-    "model": SPANDATA.GEN_AI_REQUEST_MODEL,
-    "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE,
-    "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS,
-    "k": SPANDATA.GEN_AI_REQUEST_TOP_K,
-    "p": SPANDATA.GEN_AI_REQUEST_TOP_P,
-    "seed": SPANDATA.GEN_AI_REQUEST_SEED,
-    "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY,
-    "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY,
-}
-
 
 def _normalize_embedding_input(texts):
     # type: (Any) -> Any
@@ -66,7 +55,6 @@ def __init__(self, include_prompts=True):
     def setup_once():
         # type: () -> None
         # Lazy imports to avoid circular dependencies:
-        # v1/v2 import COLLECTED_CHAT_PARAMS and _capture_exception from this module.
         from sentry_sdk.integrations.cohere.v1 import setup_v1
         from sentry_sdk.integrations.cohere.v2 import setup_v2
 

sentry_sdk/integrations/cohere/utils.py

@@ -0,0 +1,35 @@
+from sentry_sdk.ai.utils import set_data_normalized
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+
+_MISSING = object()
+
+
+def transitive_getattr(obj, *attrs):
+    # type: (Any, str) -> Any
+    current = obj
+    for attr in attrs:
+        current = getattr(current, attr, _MISSING)
+        if current is _MISSING:
+            return None
+    return current
+
+
+def get_first_from_sources(obj, source_paths, require_truthy=False):
+    # type: (Any, list[tuple[str, ...]], bool) -> Any
+    for source_path in source_paths:
+        value = transitive_getattr(obj, *source_path)
+        if value if require_truthy else value is not None:
+            return value
+    return None
+
+
+def set_span_data_from_sources(span, obj, target_sources, require_truthy):
+    # type: (Any, Any, dict[str, list[tuple[str, ...]]], bool) -> None
+    for spandata_key, source_paths in target_sources.items():
+        value = get_first_from_sources(obj, source_paths, require_truthy=require_truthy)
+        if value is not None:
+            set_data_normalized(span, spandata_key, value)

sentry_sdk/integrations/cohere/v1.py

@@ -2,32 +2,30 @@
 from functools import wraps
 
 from sentry_sdk.ai.monitoring import record_token_usage
-from sentry_sdk.consts import OP, SPANDATA
-from sentry_sdk.ai.utils import set_data_normalized
 from sentry_sdk.ai.span_config import set_input_span_data
+from sentry_sdk.ai.utils import set_data_normalized, transform_message_content
+from sentry_sdk.consts import OP, SPANDATA
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Any, Callable, Iterator
     from cohere import StreamedChatResponse
-    from sentry_sdk.tracing import Span
 
 import sentry_sdk
-from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.utils import capture_internal_exceptions, reraise
-
 from sentry_sdk.integrations.cohere import (
     CohereIntegration,
-    COLLECTED_CHAT_PARAMS,
     _capture_exception,
 )
+from sentry_sdk.integrations.cohere.utils import (
+    get_first_from_sources,
+    set_span_data_from_sources,
+)
+from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.utils import capture_internal_exceptions, reraise
 
 try:
-    from cohere import (
-        ChatStreamEndEvent,
-        NonStreamedChatResponse,
-    )
+    from cohere import ChatStreamEndEvent, NonStreamedChatResponse
 
     try:
         from cohere import StreamEndStreamedChatResponse
@@ -40,53 +38,32 @@
 except ImportError:
     _has_chat_types = False
 
-COLLECTED_PII_CHAT_PARAMS = {
-    "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
-    "preamble": SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
-}
-
-
-def _extract_messages_v1(kwargs):
-    # type: (dict[str, Any]) -> list[dict[str, str]]
-    """Extract role/content dicts from V1-style chat_history + message."""
-    messages = []
-    for x in kwargs.get("chat_history", []):
-        messages.append(
-            {
-                "role": getattr(x, "role", "").lower(),
-                "content": getattr(x, "message", ""),
-            }
-        )
-    message = kwargs.get("message")
-    if message:
-        messages.append({"role": "user", "content": message})
-    return messages
-
-
-COHERE_V1_CHAT_CONFIG = {
-    "system": "cohere",
-    "operation": "chat",
-    "params": COLLECTED_CHAT_PARAMS,
-    "pii_params": COLLECTED_PII_CHAT_PARAMS,
-    "extract_messages": _extract_messages_v1,
+CHAT_RESPONSE_SOURCES = {
+    SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)],
+    SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)],
 }
-
-COLLECTED_CHAT_RESP_ATTRS = {
-    "generation_id": SPANDATA.GEN_AI_RESPONSE_ID,
-    "finish_reason": SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS,
+PII_CHAT_RESPONSE_SOURCES = {
+    SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("tool_calls",)],
 }
-
-COLLECTED_PII_CHAT_RESP_ATTRS = {
-    "tool_calls": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
+CHAT_RESPONSE_TEXT_SOURCES = [("text",)]
+CHAT_USAGE_TOKEN_SOURCES = {
+    SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: [
+        ("meta", "billed_units", "input_tokens"),
+        ("meta", "tokens", "input_tokens"),
+    ],
+    SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: [
+        ("meta", "billed_units", "output_tokens"),
+        ("meta", "tokens", "output_tokens"),
+    ],
 }
+STREAM_RESPONSE_SOURCES = [("response",)]
 
 
 def setup_v1(wrap_embed_fn):
     # type: (Callable[..., Any]) -> None
-    """Called from CohereIntegration.setup_once() to patch V1 Client methods."""
     try:
-        from cohere.client import Client
         from cohere.base_client import BaseCohere
+        from cohere.client import Client
     except ImportError:
         return
 
@@ -104,7 +81,6 @@ def _wrap_chat(f, streaming):
     def new_chat(*args, **kwargs):
         # type: (*Any, **Any) -> Any
         integration = sentry_sdk.get_client().get_integration(CohereIntegration)
-
         if (
             integration is None
             or "message" not in kwargs
@@ -113,6 +89,7 @@ def new_chat(*args, **kwargs):
             return f(*args, **kwargs)
 
         model = kwargs.get("model", "")
+        include_pii = should_send_default_pii() and integration.include_prompts
 
         with sentry_sdk.start_span(
             op=OP.GEN_AI_CHAT,
@@ -128,75 +105,84 @@ def new_chat(*args, **kwargs):
                 reraise(*exc_info)
 
             with capture_internal_exceptions():
+                if model:
+                    set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MODEL, model)
                 set_input_span_data(
                     span,
                     kwargs,
                     integration,
                     {
-                        **COHERE_V1_CHAT_CONFIG,
+                        "system": "cohere",
+                        "operation": "chat",
+                        "extract_messages": _extract_messages_v1,
                         "extra_static": {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming},
                     },
                 )
 
                 if streaming:
-                    old_iterator = res
-
-                    def new_iterator():
-                        # type: () -> Iterator[StreamedChatResponse]
-                        with capture_internal_exceptions():
-                            for x in old_iterator:
-                                if isinstance(x, ChatStreamEndEvent) or isinstance(
-                                    x, StreamEndStreamedChatResponse
-                                ):
-                                    collect_chat_response_fields(
-                                        span,
-                                        x.response,
-                                        include_pii=should_send_default_pii()
-                                        and integration.include_prompts,
-                                    )
-                                yield x
-
-                    return new_iterator()
-                elif isinstance(res, NonStreamedChatResponse):
-                    collect_chat_response_fields(
-                        span,
-                        res,
-                        include_pii=should_send_default_pii()
-                        and integration.include_prompts,
-                    )
+                    return _iter_v1_stream_events(res, span, include_pii)
+                if isinstance(res, NonStreamedChatResponse):
+                    _collect_v1_response_fields(span, res, include_pii=include_pii)
                 else:
                     set_data_normalized(span, "unknown_response", True)
                 return res
 
-    def collect_chat_response_fields(span, res, include_pii):
-        # type: (Span, NonStreamedChatResponse, bool) -> None
-        if include_pii:
-            if hasattr(res, "text"):
-                set_data_normalized(
-                    span,
-                    SPANDATA.GEN_AI_RESPONSE_TEXT,
-                    [res.text],
-                )
-            for attr, spandata_key in COLLECTED_PII_CHAT_RESP_ATTRS.items():
-                if hasattr(res, attr):
-                    set_data_normalized(span, spandata_key, getattr(res, attr))
+    return new_chat
 
-        for attr, spandata_key in COLLECTED_CHAT_RESP_ATTRS.items():
-            if hasattr(res, attr):
-                set_data_normalized(span, spandata_key, getattr(res, attr))
 
-        if hasattr(res, "meta"):
-            if hasattr(res.meta, "billed_units"):
-                record_token_usage(
-                    span,
-                    input_tokens=res.meta.billed_units.input_tokens,
-                    output_tokens=res.meta.billed_units.output_tokens,
-                )
-            elif hasattr(res.meta, "tokens"):
-                record_token_usage(
-                    span,
-                    input_tokens=res.meta.tokens.input_tokens,
-                    output_tokens=res.meta.tokens.output_tokens,
-                )
+def _extract_messages_v1(kwargs):
+    # type: (dict[str, Any]) -> list[dict[str, str]]
+    messages = []
+    for x in kwargs.get("chat_history", []):
+        messages.append(
+            {
+                "role": getattr(x, "role", "").lower(),
+                "content": transform_message_content(getattr(x, "message", "")),
+            }
+        )
+    message = kwargs.get("message")
+    if message:
+        messages.append({"role": "user", "content": transform_message_content(message)})
+    return messages
 
-    return new_chat
+
+def _iter_v1_stream_events(old_iterator, span, include_pii):
+    # type: (Any, Any, bool) -> Iterator[StreamedChatResponse]
+    with capture_internal_exceptions():
+        for x in old_iterator:
+            if isinstance(x, ChatStreamEndEvent) or isinstance(
+                x, StreamEndStreamedChatResponse
+            ):
+                _collect_v1_stream_end_fields(span, x, include_pii)
+            yield x
+
+
+def _collect_v1_stream_end_fields(span, event, include_pii):
+    # type: (Any, Any, bool) -> None
+    response = get_first_from_sources(event, STREAM_RESPONSE_SOURCES)
+    if response is not None:
+        _collect_v1_response_fields(span, response, include_pii)
+
+
+def _collect_v1_response_fields(span, response, include_pii):
+    # type: (Any, Any, bool) -> None
+    if include_pii:
+        text = get_first_from_sources(response, CHAT_RESPONSE_TEXT_SOURCES)
+        if text is not None:
+            set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, [text])
+        set_span_data_from_sources(
+            span, response, PII_CHAT_RESPONSE_SOURCES, require_truthy=False
+        )
+
+    set_span_data_from_sources(
+        span, response, CHAT_RESPONSE_SOURCES, require_truthy=False
+    )
+    record_token_usage(
+        span,
+        input_tokens=get_first_from_sources(
+            response, CHAT_USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS]
+        ),
+        output_tokens=get_first_from_sources(
+            response, CHAT_USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS]
+        ),
+    )