ddtrace 3.11.0rc2__cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl → 3.11.0rc3__cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl

ddtrace/_trace/sampling_rule.py

@@ -60,7 +60,7 @@ class SamplingRule(object):
             number of characters, and "?" meaning any one character. If all tags specified in a SamplingRule are
             matches with a given span, that span is considered to have matching tags with the rule.
         """
-        self.sample_rate = float(min(1, max(0, sample_rate)))
+        self.sample_rate = min(1.0, max(0.0, float(sample_rate)))
         # since span.py converts None to 'None' for tags, and does not accept 'None' for metrics
         # we can just create a GlobMatcher for 'None' and it will match properly
         self._tag_value_matchers = (
@@ -111,46 +111,38 @@ class SamplingRule(object):
         :returns: Whether this span matches or not
         :rtype: :obj:`bool`
         """
-        tags_match = self.tags_match(span)
-        return tags_match and self._matches((span.service, span.name, span.resource))
+        return self.tags_match(span) and self._matches((span.service, span.name, span.resource))
 
     def tags_match(self, span: Span) -> bool:
-        tag_match = True
-        if self._tag_value_matchers:
-            tag_match = self.check_tags(span.get_tags(), span.get_metrics())
-        return tag_match
+        if not self._tag_value_matchers:
+            return True
 
-    def check_tags(self, meta, metrics):
-        if meta is None and metrics is None:
+        meta = span._meta or {}
+        metrics = span._metrics or {}
+        if not meta and not metrics:
             return False
 
-        tag_match = False
-        for tag_key in self._tag_value_matchers.keys():
-            value = meta.get(tag_key)
-            # it's because we're not checking metrics first before continuing
+        for tag_key, pattern in self._tag_value_matchers.items():
+            value = meta.get(tag_key, metrics.get(tag_key))
             if value is None:
-                value = metrics.get(tag_key)
-                if value is None:
-                    continue
-                # Floats: Matching floating point values with a non-zero decimal part is not supported.
-                # For floating point values with a non-zero decimal part, any all * pattern always returns true.
-                # Other patterns always return false.
-                if isinstance(value, float):
-                    if not value.is_integer():
-                        if all(c == "*" for c in self._tag_value_matchers[tag_key].pattern):
-                            tag_match = True
-                            continue
-                        else:
-                            return False
-                    else:
-                        value = int(value)
-
-            tag_match = self._tag_value_matchers[tag_key].match(str(value))
-            # if we don't match with all specified tags for a rule, it's not a match
-            if tag_match is False:
+                # If the tag is not present, we failed the match
+                # (Metrics and meta do not support the value None)
+                return False
+
+            if isinstance(value, float):
+                # Floats: Convert floats that represent integers to int for matching. This is because
+                # SamplingRules only support integers for matfching or glob patterns.
+                if value.is_integer():
+                    value = int(value)
+                elif set(pattern.pattern) - {"?", "*"}:
+                    # Only match floats to patterns that only contain wildcards (ex: * or ?*)
+                    # This is because we do not want to match floats to patterns like `23.*`.
+                    return False
+
+            if not pattern.match(str(value)):
                 return False
 
-        return tag_match
+        return True
 
     def sample(self, span):
         """

ddtrace/_trace/trace_handlers.py

@@ -661,13 +661,9 @@ def _on_botocore_patched_bedrock_api_call_started(ctx, request_params):
 
     span.set_tag_str("bedrock.request.model_provider", ctx["model_provider"])
     span.set_tag_str("bedrock.request.model", ctx["model_name"])
-    for k, v in request_params.items():
-        if k == "prompt":
-            if integration.is_pc_sampled_span(span):
-                v = integration.trunc(str(v))
-        span.set_tag_str("bedrock.request.{}".format(k), str(v))
-        if k == "n":
-            ctx.set_item("num_generations", str(v))
+
+    if "n" in request_params:
+        ctx.set_item("num_generations", str(request_params["n"]))
 
 
 def _on_botocore_patched_bedrock_api_call_exception(ctx, exc_info):
@@ -680,16 +676,6 @@ def _on_botocore_patched_bedrock_api_call_exception(ctx, exc_info):
     span.finish()
 
 
-def _on_botocore_patched_bedrock_api_call_success(ctx, reqid, latency, input_token_count, output_token_count):
-    span = ctx.span
-    span.set_tag_str("bedrock.response.id", reqid)
-    span.set_tag_str("bedrock.response.duration", latency)
-    if input_token_count:
-        span.set_metric("bedrock.response.usage.prompt_tokens", int(input_token_count))
-    if output_token_count:
-        span.set_metric("bedrock.response.usage.completion_tokens", int(output_token_count))
-
-
 def _propagate_context(ctx, headers):
     distributed_tracing_enabled = ctx["integration_config"].distributed_tracing_enabled
     span = ctx.span
@@ -731,38 +717,13 @@ def _on_botocore_bedrock_process_response_converse(
 def _on_botocore_bedrock_process_response(
     ctx: core.ExecutionContext,
     formatted_response: Dict[str, Any],
-    metadata: Dict[str, Any],
-    body: Dict[str, List[Dict]],
-    should_set_choice_ids: bool,
 ) -> None:
-    text = formatted_response["text"]
-    span = ctx.span
-    model_name = ctx["model_name"]
-    if should_set_choice_ids:
-        for i in range(len(text)):
-            span.set_tag_str("bedrock.response.choices.{}.id".format(i), str(body["generations"][i]["id"]))
-    integration = ctx["bedrock_integration"]
-    if metadata is not None:
-        for k, v in metadata.items():
-            if k in ["usage.completion_tokens", "usage.prompt_tokens"] and v:
-                span.set_metric("bedrock.response.{}".format(k), int(v))
-            else:
-                span.set_tag_str("bedrock.{}".format(k), str(v))
-    if "embed" in model_name:
-        span.set_metric("bedrock.response.embedding_length", len(formatted_response["text"][0]))
-        span.finish()
-        return
-    for i in range(len(formatted_response["text"])):
-        if integration.is_pc_sampled_span(span):
-            span.set_tag_str(
-                "bedrock.response.choices.{}.text".format(i),
-                integration.trunc(str(formatted_response["text"][i])),
-            )
-        span.set_tag_str(
-            "bedrock.response.choices.{}.finish_reason".format(i), str(formatted_response["finish_reason"][i])
-        )
-    integration.llmobs_set_tags(span, args=[ctx], kwargs={}, response=formatted_response)
-    span.finish()
+    with ctx.span as span:
+        model_name = ctx["model_name"]
+        integration = ctx["bedrock_integration"]
+        if "embed" in model_name:
+            return
+        integration.llmobs_set_tags(span, args=[ctx], kwargs={}, response=formatted_response)
 
 
 def _on_botocore_sqs_recvmessage_post(
@@ -931,7 +892,6 @@ def listen():
     core.on("botocore.client_context.update_messages", _on_botocore_update_messages)
     core.on("botocore.patched_bedrock_api_call.started", _on_botocore_patched_bedrock_api_call_started)
     core.on("botocore.patched_bedrock_api_call.exception", _on_botocore_patched_bedrock_api_call_exception)
-    core.on("botocore.patched_bedrock_api_call.success", _on_botocore_patched_bedrock_api_call_success)
     core.on("botocore.bedrock.process_response", _on_botocore_bedrock_process_response)
     core.on("botocore.bedrock.process_response_converse", _on_botocore_bedrock_process_response_converse)
     core.on("botocore.sqs.ReceiveMessage.post", _on_botocore_sqs_recvmessage_post)

ddtrace/_trace/utils_botocore/span_tags.py

@@ -12,11 +12,52 @@ from ddtrace.ext import SpanKind
 from ddtrace.ext import aws
 from ddtrace.ext import http
 from ddtrace.internal.constants import COMPONENT
+from ddtrace.internal.serverless import in_aws_lambda
 from ddtrace.internal.utils.formats import deep_getattr
 
 
 _PAYLOAD_TAGGER = AWSPayloadTagging()
 
+SERVICE_MAP = {
+    "eventbridge": "events",
+    "events": "events",
+    "sqs": "sqs",
+    "sns": "sns",
+    "kinesis": "kinesis",
+    "dynamodb": "dynamodb",
+    "dynamodbdocument": "dynamodb",
+}
+
+
+# Helper to build AWS hostname from service, region and parameters
+def _derive_peer_hostname(service: str, region: str, params: Optional[Dict[str, Any]] = None) -> Optional[str]:
+    """Return hostname for given AWS service according to Datadog peer hostname rules.
+
+    Only returns hostnames for specific AWS services:
+        - eventbridge/events -> events.<region>.amazonaws.com
+        - sqs               -> sqs.<region>.amazonaws.com
+        - sns               -> sns.<region>.amazonaws.com
+        - kinesis           -> kinesis.<region>.amazonaws.com
+        - dynamodb          -> dynamodb.<region>.amazonaws.com
+        - s3                -> <bucket>.s3.<region>.amazonaws.com (if Bucket param present)
+                              s3.<region>.amazonaws.com          (otherwise)
+
+    Other services return ``None``.
+    """
+
+    if not region:
+        return None
+
+    aws_service = service.lower()
+
+    if aws_service == "s3":
+        bucket = params.get("Bucket") if params else None
+        return f"{bucket}.s3.{region}.amazonaws.com" if bucket else f"s3.{region}.amazonaws.com"
+
+    mapped = SERVICE_MAP.get(aws_service)
+
+    return f"{mapped}.{region}.amazonaws.com" if mapped else None
+
 
 def set_botocore_patched_api_call_span_tags(span: Span, instance, args, params, endpoint_name, operation):
     span.set_tag_str(COMPONENT, config.botocore.integration_name)
@@ -51,6 +92,13 @@ def set_botocore_patched_api_call_span_tags(span: Span, instance, args, params,
         span.set_tag_str("aws.region", region_name)
         span.set_tag_str("region", region_name)
 
+        # Derive peer hostname only in serverless environments to avoid
+        # unnecessary tag noise in traditional hosts/containers.
+        if in_aws_lambda():
+            hostname = _derive_peer_hostname(endpoint_name, region_name, params)
+            if hostname:
+                span.set_tag_str("peer.service", hostname)
+
 
 def set_botocore_response_metadata_tags(
     span: Span, result: Dict[str, Any], is_error_code_fn: Optional[Callable] = None

ddtrace/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.11.0rc2'
-__version_tuple__ = version_tuple = (3, 11, 0, 'rc2')
+__version__ = version = '3.11.0rc3'
+__version_tuple__ = version_tuple = (3, 11, 0, 'rc3')

ddtrace/appsec/_constants.py

@@ -275,6 +275,12 @@ class API_SECURITY(metaclass=Constant_Class):
     SAMPLE_RATE: Literal["DD_API_SECURITY_REQUEST_SAMPLE_RATE"] = "DD_API_SECURITY_REQUEST_SAMPLE_RATE"
     SAMPLE_DELAY: Literal["DD_API_SECURITY_SAMPLE_DELAY"] = "DD_API_SECURITY_SAMPLE_DELAY"
     MAX_PAYLOAD_SIZE: Literal[0x1000000] = 0x1000000  # 16MB maximum size
+    ENDPOINT_COLLECTION: Literal[
+        "DD_API_SECURITY_ENDPOINT_COLLECTION_ENABLED"
+    ] = "DD_API_SECURITY_ENDPOINT_COLLECTION_ENABLED"
+    ENDPOINT_COLLECTION_LIMIT: Literal[
+        "DD_API_SECURITY_ENDPOINT_COLLECTION_MESSAGE_LIMIT"
+    ] = "DD_API_SECURITY_ENDPOINT_COLLECTION_MESSAGE_LIMIT"
 
 
 class WAF_CONTEXT_NAMES(metaclass=Constant_Class):
@@ -348,6 +354,7 @@ class DEFAULT(metaclass=Constant_Class):
         r"|ey[I-L][\w=-]+\.(ey[I-L][\w=-]+(?:\.[\w.+\/=-]+)?)|[\-]{5}BEGIN[a-z\s]+PRIVATE\sKEY[\-]{5}([^\-]+)[\-]"
         r"{5}END[a-z\s]+PRIVATE\sKEY|ssh-rsa\s*([a-z0-9\/\.+]{100,})"
     )
+    ENDPOINT_COLLECTION_LIMIT = 300
 
 
 class EXPLOIT_PREVENTION(metaclass=Constant_Class):

ddtrace/appsec/_handlers.py

@@ -4,6 +4,7 @@ import json
 from typing import Any
 from typing import Dict
 from typing import Optional
+from typing import Union
 
 import xmltodict
 
@@ -11,6 +12,7 @@ from ddtrace._trace.span import Span
 from ddtrace.appsec._asm_request_context import _call_waf
 from ddtrace.appsec._asm_request_context import _call_waf_first
 from ddtrace.appsec._asm_request_context import get_blocked
+from ddtrace.appsec._asm_request_context import set_body_response
 from ddtrace.appsec._constants import SPAN_DATA_NAMES
 from ddtrace.appsec._http_utils import extract_cookies_from_headers
 from ddtrace.appsec._http_utils import normalize_headers
@@ -157,6 +159,14 @@ def _on_lambda_start_response(
     _call_waf(("aws_lambda",))
 
 
+def _on_lambda_parse_body(
+    response_body: Optional[Union[str, Dict[str, Any]]],
+):
+    if asm_config._api_security_feature_active:
+        if response_body:
+            set_body_response(response_body)
+
+
 # ASGI
 
 
@@ -408,6 +418,7 @@ def listen():
 
     core.on("aws_lambda.start_request", _on_lambda_start_request)
     core.on("aws_lambda.start_response", _on_lambda_start_response)
+    core.on("aws_lambda.parse_body", _on_lambda_parse_body)
 
     core.on("grpc.server.response.message", _on_grpc_server_response)
     core.on("grpc.server.data", _on_grpc_server_data)

ddtrace/appsec/_processor.py

@@ -189,7 +189,7 @@ class AppSecSpanProcessor(SpanProcessor):
             if skip_event:
                 core.discard_item("appsec_skip_next_lambda_event")
                 log.debug(
-                    "appsec: ignoring unsupported lamdba event",
+                    "appsec: ignoring unsupported lambda event",
                 )
                 span.set_metric(APPSEC.UNSUPPORTED_EVENT_TYPE, 1.0)
                 return

ddtrace/contrib/internal/aiobotocore/patch.py

@@ -5,6 +5,7 @@ import aiobotocore.client
 import wrapt
 
 from ddtrace import config
+from ddtrace._trace.utils_botocore.span_tags import _derive_peer_hostname
 from ddtrace.constants import _SPAN_MEASURED_KEY
 from ddtrace.constants import SPAN_KIND
 from ddtrace.contrib.internal.trace_utils import ext_service
@@ -16,6 +17,7 @@ from ddtrace.ext import http
 from ddtrace.internal.constants import COMPONENT
 from ddtrace.internal.schema import schematize_cloud_api_operation
 from ddtrace.internal.schema import schematize_service_name
+from ddtrace.internal.serverless import in_aws_lambda
 from ddtrace.internal.utils import ArgumentError
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.internal.utils.formats import asbool
@@ -145,6 +147,12 @@ async def _wrapped_api_call(original_func, instance, args, kwargs):
 
         region_name = deep_getattr(instance, "meta.region_name")
 
+        if in_aws_lambda():
+            # Derive the peer hostname now that we have both service and region.
+            hostname = _derive_peer_hostname(endpoint_name, region_name, params)
+            if hostname:
+                span.set_tag_str("peer.service", hostname)
+
         meta = {
             "aws.agent": "aiobotocore",
             "aws.operation": operation,

ddtrace/contrib/internal/boto/patch.py

@@ -7,6 +7,7 @@ import boto.connection
 import wrapt
 
 from ddtrace import config
+from ddtrace._trace.utils_botocore.span_tags import _derive_peer_hostname
 from ddtrace.constants import _SPAN_MEASURED_KEY
 from ddtrace.constants import SPAN_KIND
 from ddtrace.ext import SpanKind
@@ -16,6 +17,7 @@ from ddtrace.ext import http
 from ddtrace.internal.constants import COMPONENT
 from ddtrace.internal.schema import schematize_cloud_api_operation
 from ddtrace.internal.schema import schematize_service_name
+from ddtrace.internal.serverless import in_aws_lambda
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.internal.utils.formats import asbool
 from ddtrace.internal.utils.wrappers import unwrap
@@ -122,6 +124,12 @@ def patched_query_request(original_func, instance, args, kwargs):
             meta[aws.REGION] = region_name
             meta[aws.AWSREGION] = region_name
 
+            if in_aws_lambda():
+                # Derive the peer hostname now that we have both service and region.
+                hostname = _derive_peer_hostname(endpoint_name, region_name, params)
+                if hostname:
+                    meta["peer.service"] = hostname
+
         span.set_tags(meta)
 
         # Original func returns a boto.connection.HTTPResponse object
@@ -183,6 +191,12 @@ def patched_auth_request(original_func, instance, args, kwargs):
             meta[aws.REGION] = region_name
             meta[aws.AWSREGION] = region_name
 
+            if in_aws_lambda():
+                # Derive the peer hostname
+                hostname = _derive_peer_hostname(endpoint_name, region_name, None)
+                if hostname:
+                    meta["peer.service"] = hostname
+
         span.set_tags(meta)
 
         # Original func returns a boto.connection.HTTPResponse object

ddtrace/contrib/internal/botocore/services/bedrock.py

@@ -48,12 +48,9 @@ class TracedBotocoreStreamingBody(wrapt.ObjectProxy):
             self._body.append(json.loads(body))
             if self.__wrapped__.tell() == int(self.__wrapped__._content_length):
                 formatted_response = _extract_text_and_response_reason(self._execution_ctx, self._body[0])
-                model_provider = self._execution_ctx["model_provider"]
-                model_name = self._execution_ctx["model_name"]
-                should_set_choice_ids = model_provider == _COHERE and "embed" not in model_name
                 core.dispatch(
                     "botocore.bedrock.process_response",
-                    [self._execution_ctx, formatted_response, None, self._body[0], should_set_choice_ids],
+                    [self._execution_ctx, formatted_response],
                 )
             return body
         except Exception:
@@ -67,12 +64,9 @@ class TracedBotocoreStreamingBody(wrapt.ObjectProxy):
             for line in lines:
                 self._body.append(json.loads(line))
             formatted_response = _extract_text_and_response_reason(self._execution_ctx, self._body[0])
-            model_provider = self._execution_ctx["model_provider"]
-            model_name = self._execution_ctx["model_name"]
-            should_set_choice_ids = model_provider == _COHERE and "embed" not in model_name
             core.dispatch(
                 "botocore.bedrock.process_response",
-                [self._execution_ctx, formatted_response, None, self._body[0], should_set_choice_ids],
+                [self._execution_ctx, formatted_response],
             )
             return lines
         except Exception:
@@ -93,16 +87,10 @@ class TracedBotocoreStreamingBody(wrapt.ObjectProxy):
         finally:
             if exception_raised:
                 return
-            metadata = _extract_streamed_response_metadata(self._execution_ctx, self._body)
             formatted_response = _extract_streamed_response(self._execution_ctx, self._body)
-            model_provider = self._execution_ctx["model_provider"]
-            model_name = self._execution_ctx["model_name"]
-            should_set_choice_ids = (
-                model_provider == _COHERE and "is_finished" not in self._body[0] and "embed" not in model_name
-            )
             core.dispatch(
                 "botocore.bedrock.process_response",
-                [self._execution_ctx, formatted_response, metadata, self._body, should_set_choice_ids],
+                [self._execution_ctx, formatted_response],
             )
 
 
@@ -443,18 +431,6 @@ def handle_bedrock_response(
         safe_token_count(cache_write_tokens),
     )
 
-    # for both converse & invoke, dispatch success event to store basic metrics
-    core.dispatch(
-        "botocore.patched_bedrock_api_call.success",
-        [
-            ctx,
-            str(metadata.get("RequestId", "")),
-            request_latency,
-            str(input_tokens),
-            str(output_tokens),
-        ],
-    )
-
     if ctx["resource"] == "Converse":
         core.dispatch("botocore.bedrock.process_response_converse", [ctx, result])
         return result

ddtrace/contrib/internal/django/patch.py

@@ -49,6 +49,7 @@ from ddtrace.internal.utils.formats import asbool
 from ddtrace.internal.utils.importlib import func_name
 from ddtrace.propagation._database_monitoring import _DBM_Propagator
 from ddtrace.settings.asm import config as asm_config
+from ddtrace.settings.asm import endpoint_collection
 from ddtrace.settings.integration import IntegrationConfig
 from ddtrace.trace import Pin
 from ddtrace.vendor.packaging.version import parse as parse_version
@@ -81,6 +82,7 @@ config._add(
         use_legacy_resource_format=asbool(os.getenv("DD_DJANGO_USE_LEGACY_RESOURCE_FORMAT", default=False)),
         _trace_asgi_websocket=os.getenv("DD_ASGI_TRACE_WEBSOCKET", default=False),
         obfuscate_404_resource=os.getenv("DD_ASGI_OBFUSCATE_404_RESOURCE", default=False),
+        views={},
     ),
 )
 
@@ -598,7 +600,7 @@ def traced_template_render(django, pin, wrapped, instance, args, kwargs):
         return wrapped(*args, **kwargs)
 
 
-def instrument_view(django, view):
+def instrument_view(django, view, path=None):
     """
     Helper to wrap Django views.
 
@@ -608,10 +610,24 @@ def instrument_view(django, view):
         for cls in reversed(getmro(view)):
             _instrument_view(django, cls)
 
-    return _instrument_view(django, view)
+    return _instrument_view(django, view, path=path)
 
 
-def _instrument_view(django, view):
+def extract_request_method_list(view):
+    try:
+        while "view_func" in view.__code__.co_freevars:
+            view = view.__closure__[view.__code__.co_freevars.index("view_func")].cell_contents
+        if "request_method_list" in view.__code__.co_freevars:
+            return view.__closure__[view.__code__.co_freevars.index("request_method_list")].cell_contents
+        return []
+    except Exception:
+        return []
+
+
+_DEFAULT_METHODS = ("get", "delete", "post", "options", "head")
+
+
+def _instrument_view(django, view, path=None):
     """Helper to wrap Django views."""
     from . import utils
 
@@ -620,9 +636,14 @@ def _instrument_view(django, view):
         return view
 
     # Patch view HTTP methods and lifecycle methods
-    http_method_names = getattr(view, "http_method_names", ("get", "delete", "post", "options", "head"))
+
+    http_method_names = getattr(view, "http_method_names", ())
+    request_method_list = extract_request_method_list(view) or http_method_names
+    if path is not None:
+        for method in request_method_list or ["*"]:
+            endpoint_collection.add_endpoint(method, path, operation_name="django.request")
     lifecycle_methods = ("setup", "dispatch", "http_method_not_allowed")
-    for name in list(http_method_names) + list(lifecycle_methods):
+    for name in list(request_method_list or _DEFAULT_METHODS) + list(lifecycle_methods):
         try:
             func = getattr(view, name, None)
             if not func or isinstance(func, wrapt.ObjectProxy):
@@ -665,18 +686,20 @@ def traced_urls_path(django, pin, wrapped, instance, args, kwargs):
     try:
         from_args = False
         view = kwargs.pop("view", None)
+        path = kwargs.pop("path", None)
         if view is None:
             view = args[1]
             from_args = True
+        if path is None:
+            path = args[0]
 
         core.dispatch("service_entrypoint.patch", (unwrap(view),))
-
         if from_args:
             args = list(args)
-            args[1] = instrument_view(django, view)
+            args[1] = instrument_view(django, view, path=path)
             args = tuple(args)
         else:
-            kwargs["view"] = instrument_view(django, view)
+            kwargs["view"] = instrument_view(django, view, path=path)
     except Exception:
         log.debug("Failed to instrument Django url path %r %r", args, kwargs, exc_info=True)
     return wrapped(*args, **kwargs)

ddtrace/contrib/internal/google_genai/_utils.py

@@ -6,7 +6,7 @@ from typing import Optional
 
 import wrapt
 
-from ddtrace.llmobs._integrations.google_genai_utils import DEFAULT_MODEL_ROLE
+from ddtrace.llmobs._integrations.google_utils import GOOGLE_GENAI_DEFAULT_MODEL_ROLE
 from ddtrace.llmobs._utils import _get_attr
 
 
@@ -32,7 +32,7 @@ def _join_chunks(chunks: List[Any]) -> Optional[Dict[str, Any]]:
                     continue
 
                 if role is None:
-                    role = _get_attr(content, "role", DEFAULT_MODEL_ROLE)
+                    role = _get_attr(content, "role", GOOGLE_GENAI_DEFAULT_MODEL_ROLE)
 
                 parts = _get_attr(content, "parts", [])
                 for part in parts:

ddtrace/contrib/internal/google_genai/patch.py

@@ -9,7 +9,7 @@ from ddtrace.contrib.internal.trace_utils import unwrap
 from ddtrace.contrib.internal.trace_utils import with_traced_module
 from ddtrace.contrib.internal.trace_utils import wrap
 from ddtrace.llmobs._integrations import GoogleGenAIIntegration
-from ddtrace.llmobs._integrations.google_genai_utils import extract_provider_and_model_name
+from ddtrace.llmobs._integrations.google_utils import extract_provider_and_model_name
 from ddtrace.trace import Pin
 
 
@@ -27,7 +27,7 @@ def get_version() -> str:
 @with_traced_module
 def traced_generate(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     with integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -46,7 +46,7 @@ def traced_generate(genai, pin, func, instance, args, kwargs):
 @with_traced_module
 async def traced_async_generate(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     with integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -65,7 +65,7 @@ async def traced_async_generate(genai, pin, func, instance, args, kwargs):
 @with_traced_module
 def traced_generate_stream(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -86,7 +86,7 @@ def traced_generate_stream(genai, pin, func, instance, args, kwargs):
 @with_traced_module
 async def traced_async_generate_stream(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -107,7 +107,7 @@ async def traced_async_generate_stream(genai, pin, func, instance, args, kwargs)
 @with_traced_module
 def traced_embed_content(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     with integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
@@ -126,7 +126,7 @@ def traced_embed_content(genai, pin, func, instance, args, kwargs):
 @with_traced_module
 async def traced_async_embed_content(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
-    provider_name, model_name = extract_provider_and_model_name(kwargs)
+    provider_name, model_name = extract_provider_and_model_name(kwargs=kwargs)
     with integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),

ddtrace/contrib/internal/google_generativeai/patch.py

@@ -11,7 +11,7 @@ from ddtrace.contrib.internal.trace_utils import unwrap
 from ddtrace.contrib.internal.trace_utils import with_traced_module
 from ddtrace.contrib.internal.trace_utils import wrap
 from ddtrace.llmobs._integrations import GeminiIntegration
-from ddtrace.llmobs._integrations.utils import extract_model_name_google
+from ddtrace.llmobs._integrations.google_utils import extract_provider_and_model_name
 from ddtrace.trace import Pin
 
 
@@ -40,11 +40,12 @@ def traced_generate(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
     stream = kwargs.get("stream", False)
     generations = None
+    provider_name, model_name = extract_provider_and_model_name(instance=instance, model_name_attr="model_name")
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
-        provider="google",
-        model=extract_model_name_google(instance, "model_name"),
+        provider=provider_name,
+        model=model_name,
         submit_to_llmobs=True,
     )
     try:
@@ -68,11 +69,12 @@ async def traced_agenerate(genai, pin, func, instance, args, kwargs):
     integration = genai._datadog_integration
     stream = kwargs.get("stream", False)
     generations = None
+    provider_name, model_name = extract_provider_and_model_name(instance=instance, model_name_attr="model_name")
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__class__.__name__, func.__name__),
-        provider="google",
-        model=extract_model_name_google(instance, "model_name"),
+        provider=provider_name,
+        model=model_name,
         submit_to_llmobs=True,
     )
     try: