sentry-sdk 0.7.5__py2.py3-none-any.whl → 2.46.0__py2.py3-none-any.whl

sentry_sdk/integrations/cloud_resource_context.py

@@ -0,0 +1,280 @@
+import json
+import urllib3
+
+from sentry_sdk.integrations import Integration
+from sentry_sdk.api import set_context
+from sentry_sdk.utils import logger
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Dict
+
+
+CONTEXT_TYPE = "cloud_resource"
+
+HTTP_TIMEOUT = 2.0
+
+AWS_METADATA_HOST = "169.254.169.254"
+AWS_TOKEN_URL = "http://{}/latest/api/token".format(AWS_METADATA_HOST)
+AWS_METADATA_URL = "http://{}/latest/dynamic/instance-identity/document".format(
+    AWS_METADATA_HOST
+)
+
+GCP_METADATA_HOST = "metadata.google.internal"
+GCP_METADATA_URL = "http://{}/computeMetadata/v1/?recursive=true".format(
+    GCP_METADATA_HOST
+)
+
+
+class CLOUD_PROVIDER:  # noqa: N801
+    """
+    Name of the cloud provider.
+    see https://opentelemetry.io/docs/reference/specification/resource/semantic_conventions/cloud/
+    """
+
+    ALIBABA = "alibaba_cloud"
+    AWS = "aws"
+    AZURE = "azure"
+    GCP = "gcp"
+    IBM = "ibm_cloud"
+    TENCENT = "tencent_cloud"
+
+
+class CLOUD_PLATFORM:  # noqa: N801
+    """
+    The cloud platform.
+    see https://opentelemetry.io/docs/reference/specification/resource/semantic_conventions/cloud/
+    """
+
+    AWS_EC2 = "aws_ec2"
+    GCP_COMPUTE_ENGINE = "gcp_compute_engine"
+
+
+class CloudResourceContextIntegration(Integration):
+    """
+    Adds cloud resource context to the Senty scope
+    """
+
+    identifier = "cloudresourcecontext"
+
+    cloud_provider = ""
+
+    aws_token = ""
+    http = urllib3.PoolManager(timeout=HTTP_TIMEOUT)
+
+    gcp_metadata = None
+
+    def __init__(self, cloud_provider=""):
+        # type: (str) -> None
+        CloudResourceContextIntegration.cloud_provider = cloud_provider
+
+    @classmethod
+    def _is_aws(cls):
+        # type: () -> bool
+        try:
+            r = cls.http.request(
+                "PUT",
+                AWS_TOKEN_URL,
+                headers={"X-aws-ec2-metadata-token-ttl-seconds": "60"},
+            )
+
+            if r.status != 200:
+                return False
+
+            cls.aws_token = r.data.decode()
+            return True
+
+        except urllib3.exceptions.TimeoutError:
+            logger.debug(
+                "AWS metadata service timed out after %s seconds", HTTP_TIMEOUT
+            )
+            return False
+        except Exception as e:
+            logger.debug("Error checking AWS metadata service: %s", str(e))
+            return False
+
+    @classmethod
+    def _get_aws_context(cls):
+        # type: () -> Dict[str, str]
+        ctx = {
+            "cloud.provider": CLOUD_PROVIDER.AWS,
+            "cloud.platform": CLOUD_PLATFORM.AWS_EC2,
+        }
+
+        try:
+            r = cls.http.request(
+                "GET",
+                AWS_METADATA_URL,
+                headers={"X-aws-ec2-metadata-token": cls.aws_token},
+            )
+
+            if r.status != 200:
+                return ctx
+
+            data = json.loads(r.data.decode("utf-8"))
+
+            try:
+                ctx["cloud.account.id"] = data["accountId"]
+            except Exception:
+                pass
+
+            try:
+                ctx["cloud.availability_zone"] = data["availabilityZone"]
+            except Exception:
+                pass
+
+            try:
+                ctx["cloud.region"] = data["region"]
+            except Exception:
+                pass
+
+            try:
+                ctx["host.id"] = data["instanceId"]
+            except Exception:
+                pass
+
+            try:
+                ctx["host.type"] = data["instanceType"]
+            except Exception:
+                pass
+
+        except urllib3.exceptions.TimeoutError:
+            logger.debug(
+                "AWS metadata service timed out after %s seconds", HTTP_TIMEOUT
+            )
+        except Exception as e:
+            logger.debug("Error fetching AWS metadata: %s", str(e))
+
+        return ctx
+
+    @classmethod
+    def _is_gcp(cls):
+        # type: () -> bool
+        try:
+            r = cls.http.request(
+                "GET",
+                GCP_METADATA_URL,
+                headers={"Metadata-Flavor": "Google"},
+            )
+
+            if r.status != 200:
+                return False
+
+            cls.gcp_metadata = json.loads(r.data.decode("utf-8"))
+            return True
+
+        except urllib3.exceptions.TimeoutError:
+            logger.debug(
+                "GCP metadata service timed out after %s seconds", HTTP_TIMEOUT
+            )
+            return False
+        except Exception as e:
+            logger.debug("Error checking GCP metadata service: %s", str(e))
+            return False
+
+    @classmethod
+    def _get_gcp_context(cls):
+        # type: () -> Dict[str, str]
+        ctx = {
+            "cloud.provider": CLOUD_PROVIDER.GCP,
+            "cloud.platform": CLOUD_PLATFORM.GCP_COMPUTE_ENGINE,
+        }
+
+        try:
+            if cls.gcp_metadata is None:
+                r = cls.http.request(
+                    "GET",
+                    GCP_METADATA_URL,
+                    headers={"Metadata-Flavor": "Google"},
+                )
+
+                if r.status != 200:
+                    return ctx
+
+                cls.gcp_metadata = json.loads(r.data.decode("utf-8"))
+
+            try:
+                ctx["cloud.account.id"] = cls.gcp_metadata["project"]["projectId"]
+            except Exception:
+                pass
+
+            try:
+                ctx["cloud.availability_zone"] = cls.gcp_metadata["instance"][
+                    "zone"
+                ].split("/")[-1]
+            except Exception:
+                pass
+
+            try:
+                # only populated in google cloud run
+                ctx["cloud.region"] = cls.gcp_metadata["instance"]["region"].split("/")[
+                    -1
+                ]
+            except Exception:
+                pass
+
+            try:
+                ctx["host.id"] = cls.gcp_metadata["instance"]["id"]
+            except Exception:
+                pass
+
+        except urllib3.exceptions.TimeoutError:
+            logger.debug(
+                "GCP metadata service timed out after %s seconds", HTTP_TIMEOUT
+            )
+        except Exception as e:
+            logger.debug("Error fetching GCP metadata: %s", str(e))
+
+        return ctx
+
+    @classmethod
+    def _get_cloud_provider(cls):
+        # type: () -> str
+        if cls._is_aws():
+            return CLOUD_PROVIDER.AWS
+
+        if cls._is_gcp():
+            return CLOUD_PROVIDER.GCP
+
+        return ""
+
+    @classmethod
+    def _get_cloud_resource_context(cls):
+        # type: () -> Dict[str, str]
+        cloud_provider = (
+            cls.cloud_provider
+            if cls.cloud_provider != ""
+            else CloudResourceContextIntegration._get_cloud_provider()
+        )
+        if cloud_provider in context_getters.keys():
+            return context_getters[cloud_provider]()
+
+        return {}
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        cloud_provider = CloudResourceContextIntegration.cloud_provider
+        unsupported_cloud_provider = (
+            cloud_provider != "" and cloud_provider not in context_getters.keys()
+        )
+
+        if unsupported_cloud_provider:
+            logger.warning(
+                "Invalid value for cloud_provider: %s (must be in %s). Falling back to autodetection...",
+                CloudResourceContextIntegration.cloud_provider,
+                list(context_getters.keys()),
+            )
+
+        context = CloudResourceContextIntegration._get_cloud_resource_context()
+        if context != {}:
+            set_context(CONTEXT_TYPE, context)
+
+
+# Map with the currently supported cloud providers
+# mapping to functions extracting the context
+context_getters = {
+    CLOUD_PROVIDER.AWS: CloudResourceContextIntegration._get_aws_context,
+    CLOUD_PROVIDER.GCP: CloudResourceContextIntegration._get_gcp_context,
+}

sentry_sdk/integrations/cohere.py

@@ -0,0 +1,274 @@
+from functools import wraps
+
+from sentry_sdk import consts
+from sentry_sdk.ai.monitoring import record_token_usage
+from sentry_sdk.consts import SPANDATA
+from sentry_sdk.ai.utils import set_data_normalized
+
+from typing import TYPE_CHECKING
+
+from sentry_sdk.tracing_utils import set_span_errored
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Iterator
+    from sentry_sdk.tracing import Span
+
+import sentry_sdk
+from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
+
+try:
+    from cohere.client import Client
+    from cohere.base_client import BaseCohere
+    from cohere import (
+        ChatStreamEndEvent,
+        NonStreamedChatResponse,
+    )
+
+    if TYPE_CHECKING:
+        from cohere import StreamedChatResponse
+except ImportError:
+    raise DidNotEnable("Cohere not installed")
+
+try:
+    # cohere 5.9.3+
+    from cohere import StreamEndStreamedChatResponse
+except ImportError:
+    from cohere import StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse
+
+
+COLLECTED_CHAT_PARAMS = {
+    "model": SPANDATA.AI_MODEL_ID,
+    "k": SPANDATA.AI_TOP_K,
+    "p": SPANDATA.AI_TOP_P,
+    "seed": SPANDATA.AI_SEED,
+    "frequency_penalty": SPANDATA.AI_FREQUENCY_PENALTY,
+    "presence_penalty": SPANDATA.AI_PRESENCE_PENALTY,
+    "raw_prompting": SPANDATA.AI_RAW_PROMPTING,
+}
+
+COLLECTED_PII_CHAT_PARAMS = {
+    "tools": SPANDATA.AI_TOOLS,
+    "preamble": SPANDATA.AI_PREAMBLE,
+}
+
+COLLECTED_CHAT_RESP_ATTRS = {
+    "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": 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,
+}
+
+
+class CohereIntegration(Integration):
+    identifier = "cohere"
+    origin = f"auto.ai.{identifier}"
+
+    def __init__(self, include_prompts=True):
+        # type: (CohereIntegration, bool) -> None
+        self.include_prompts = include_prompts
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        BaseCohere.chat = _wrap_chat(BaseCohere.chat, streaming=False)
+        Client.embed = _wrap_embed(Client.embed)
+        BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True)
+
+
+def _capture_exception(exc):
+    # type: (Any) -> None
+    set_span_errored()
+
+    event, hint = event_from_exception(
+        exc,
+        client_options=sentry_sdk.get_client().options,
+        mechanism={"type": "cohere", "handled": False},
+    )
+    sentry_sdk.capture_event(event, hint=hint)
+
+
+def _wrap_chat(f, streaming):
+    # type: (Callable[..., Any], bool) -> Callable[..., Any]
+
+    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.AI_RESPONSES,
+                    [res.text],
+                )
+            for pii_attr in COLLECTED_PII_CHAT_RESP_ATTRS:
+                if hasattr(res, pii_attr):
+                    set_data_normalized(span, "ai." + pii_attr, getattr(res, pii_attr))
+
+        for attr in COLLECTED_CHAT_RESP_ATTRS:
+            if hasattr(res, attr):
+                set_data_normalized(span, "ai." + attr, 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,
+                )
+
+            if hasattr(res.meta, "warnings"):
+                set_data_normalized(span, SPANDATA.AI_WARNINGS, res.meta.warnings)
+
+    @wraps(f)
+    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
+            or not isinstance(kwargs.get("message"), str)
+        ):
+            return f(*args, **kwargs)
+
+        message = kwargs.get("message")
+
+        span = sentry_sdk.start_span(
+            op=consts.OP.COHERE_CHAT_COMPLETIONS_CREATE,
+            name="cohere.client.Chat",
+            origin=CohereIntegration.origin,
+        )
+        span.__enter__()
+        try:
+            res = f(*args, **kwargs)
+        except Exception as e:
+            _capture_exception(e)
+            span.__exit__(None, None, None)
+            raise e from None
+
+        with capture_internal_exceptions():
+            if should_send_default_pii() and integration.include_prompts:
+                set_data_normalized(
+                    span,
+                    SPANDATA.AI_INPUT_MESSAGES,
+                    list(
+                        map(
+                            lambda x: {
+                                "role": getattr(x, "role", "").lower(),
+                                "content": getattr(x, "message", ""),
+                            },
+                            kwargs.get("chat_history", []),
+                        )
+                    )
+                    + [{"role": "user", "content": message}],
+                )
+                for k, v in COLLECTED_PII_CHAT_PARAMS.items():
+                    if k in kwargs:
+                        set_data_normalized(span, v, kwargs[k])
+
+            for k, v in COLLECTED_CHAT_PARAMS.items():
+                if k in kwargs:
+                    set_data_normalized(span, v, kwargs[k])
+            set_data_normalized(span, SPANDATA.AI_STREAMING, False)
+
+            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
+
+                    span.__exit__(None, None, None)
+
+                return new_iterator()
+            elif isinstance(res, NonStreamedChatResponse):
+                collect_chat_response_fields(
+                    span,
+                    res,
+                    include_pii=should_send_default_pii()
+                    and integration.include_prompts,
+                )
+                span.__exit__(None, None, None)
+            else:
+                set_data_normalized(span, "unknown_response", True)
+                span.__exit__(None, None, None)
+            return res
+
+    return new_chat
+
+
+def _wrap_embed(f):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+
+    @wraps(f)
+    def new_embed(*args, **kwargs):
+        # type: (*Any, **Any) -> Any
+        integration = sentry_sdk.get_client().get_integration(CohereIntegration)
+        if integration is None:
+            return f(*args, **kwargs)
+
+        with sentry_sdk.start_span(
+            op=consts.OP.COHERE_EMBEDDINGS_CREATE,
+            name="Cohere Embedding Creation",
+            origin=CohereIntegration.origin,
+        ) as span:
+            if "texts" in kwargs and (
+                should_send_default_pii() and integration.include_prompts
+            ):
+                if isinstance(kwargs["texts"], str):
+                    set_data_normalized(span, SPANDATA.AI_TEXTS, [kwargs["texts"]])
+                elif (
+                    isinstance(kwargs["texts"], list)
+                    and len(kwargs["texts"]) > 0
+                    and isinstance(kwargs["texts"][0], str)
+                ):
+                    set_data_normalized(
+                        span, SPANDATA.AI_INPUT_MESSAGES, kwargs["texts"]
+                    )
+
+            if "model" in kwargs:
+                set_data_normalized(span, SPANDATA.AI_MODEL_ID, kwargs["model"])
+            try:
+                res = f(*args, **kwargs)
+            except Exception as e:
+                _capture_exception(e)
+                raise e from None
+            if (
+                hasattr(res, "meta")
+                and hasattr(res.meta, "billed_units")
+                and hasattr(res.meta.billed_units, "input_tokens")
+            ):
+                record_token_usage(
+                    span,
+                    input_tokens=res.meta.billed_units.input_tokens,
+                    total_tokens=res.meta.billed_units.input_tokens,
+                )
+            return res
+
+    return new_embed

sentry_sdk/integrations/dedupe.py

@@ -1,13 +1,17 @@
-from sentry_sdk.hub import Hub
-from sentry_sdk.utils import ContextVar
+import weakref
+
+import sentry_sdk
+from sentry_sdk.utils import ContextVar, logger
 from sentry_sdk.integrations import Integration
 from sentry_sdk.scope import add_global_event_processor
 
-if False:
-    from typing import Any
-    from typing import Dict
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
     from typing import Optional
 
+    from sentry_sdk._types import Event, Hint
+
 
 class DedupeIntegration(Integration):
     identifier = "dedupe"
@@ -21,13 +25,43 @@ class DedupeIntegration(Integration):
         # type: () -> None
         @add_global_event_processor
         def processor(event, hint):
-            # type: (Dict[str, Any], Dict[str, Any]) -> Optional[Dict[str, Any]]
-            integration = Hub.current.get_integration(DedupeIntegration)
-            if integration is not None:
-                exc_info = hint.get("exc_info", None)
-                if exc_info is not None:
-                    exc = exc_info[1]
-                    if integration._last_seen.get(None) is exc:
-                        return None
-                    integration._last_seen.set(exc)
+            # type: (Event, Optional[Hint]) -> Optional[Event]
+            if hint is None:
+                return event
+
+            integration = sentry_sdk.get_client().get_integration(DedupeIntegration)
+            if integration is None:
+                return event
+
+            exc_info = hint.get("exc_info", None)
+            if exc_info is None:
+                return event
+
+            last_seen = integration._last_seen.get(None)
+            if last_seen is not None:
+                # last_seen is either a weakref or the original instance
+                last_seen = (
+                    last_seen() if isinstance(last_seen, weakref.ref) else last_seen
+                )
+
+            exc = exc_info[1]
+            if last_seen is exc:
+                logger.info("DedupeIntegration dropped duplicated error event %s", exc)
+                return None
+
+            # we can only weakref non builtin types
+            try:
+                integration._last_seen.set(weakref.ref(exc))
+            except TypeError:
+                integration._last_seen.set(exc)
+
             return event
+
+    @staticmethod
+    def reset_last_seen():
+        # type: () -> None
+        integration = sentry_sdk.get_client().get_integration(DedupeIntegration)
+        if integration is None:
+            return
+
+        integration._last_seen.set(None)