fix(openai): ensure embeddings input gets pc sampled correctly [backport to 1.14] (#6074)

(cherry picked from commit 93d700e)

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/contributing.html#Release-Note-Guidelines)
are followed.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.

Author: Yun-Kim

ddtrace/contrib/openai/patch.py

@@ -623,27 +623,26 @@ def handle_request(self, pin, integration, span, args, kwargs):
         return self._handle_response(pin, span, integration, resp)
 
 
-class _EmbeddingHook(_EndpointHook):
+class _EmbeddingHook(_BaseCompletionHook):
+    _request_tag_attrs = ["model", "user"]
+
     def handle_request(self, pin, integration, span, args, kwargs):
-        for kw_attr in ["model", "input", "user"]:
-            if kw_attr in kwargs:
-                if kw_attr == "input" and integration.is_pc_sampled_span(span):
-                    if isinstance(kwargs["input"], list):
-                        for idx, inp in enumerate(kwargs["input"]):
-                            span.set_tag_str("openai.request.input.%d" % idx, integration.trunc(str(inp)))
-                    else:
-                        span.set_tag("openai.request.%s" % kw_attr, kwargs[kw_attr])
-                else:
-                    span.set_tag("openai.request.%s" % kw_attr, kwargs[kw_attr])
+        embedding_input = kwargs.get("input", "")
+        if integration.is_pc_sampled_span(span):
+            if isinstance(embedding_input, list):
+                for idx, inp in enumerate(embedding_input):
+                    span.set_tag_str("openai.request.input.%d" % idx, integration.trunc(str(inp)))
+            else:
+                span.set_tag("openai.request.input", embedding_input)
+
+        self._record_request(span, kwargs)
 
         resp, error = yield
 
         if resp:
             if "data" in resp:
                 span.set_tag("openai.response.data.num-embeddings", len(resp["data"]))
                 span.set_tag("openai.response.data.embedding-length", len(resp["data"][0]["embedding"]))
-            if "object" in kwargs:
-                span.set_tag("openai.response.%s" % kw_attr, kwargs[kw_attr])
             integration.record_usage(span, resp.get("usage"))
         return resp
 

releasenotes/notes/fix-openai-embeddings-pc-sample-rate-52a8433eb22833e1.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    openai: This fix resolves an issue where embeddings inputs were always tagged regardless of the
+    configured prompt-completion sample rate.

tests/contrib/openai/test_openai.py

@@ -951,6 +951,27 @@ def test_completion_truncation(openai, openai_vcr, mock_tracer):
                     assert len(completion.replace("...", "")) == limit
 
 
+@pytest.mark.parametrize(
+    "ddtrace_config_openai",
+    [
+        dict(
+            _api_key="<not-real-but-it's-something>",
+            span_prompt_completion_sample_rate=0,
+        )
+    ],
+)
+def test_embedding_unsampled_prompt_completion(openai, openai_vcr, ddtrace_config_openai, mock_logs, mock_tracer):
+    if not hasattr(openai, "Embedding"):
+        pytest.skip("embedding not supported for this version of openai")
+    with openai_vcr.use_cassette("embedding.yaml"):
+        openai.Embedding.create(input="hello world", model="text-embedding-ada-002")
+    logs = mock_logs.enqueue.call_count
+    traces = mock_tracer.pop_traces()
+    assert len(traces) == 1
+    assert traces[0][0].get_tag("openai.request.input") is None
+    assert logs == 0
+
+
 @pytest.mark.parametrize(
     "ddtrace_config_openai",
     [