Fix completion stream token snapshots

Author: Yun-Kim

tests/contrib/openai/test_openai_v1.py

@@ -921,10 +921,19 @@ def test_span_finish_on_stream_error(openai, openai_vcr, snapshot_tracer):
                 )
 
 
-@pytest.mark.snapshot(
-    token="tests.contrib.openai.test_openai.test_completion_stream",
-    ignores=["metrics.openai.response.usage.completion_tokens", "metrics.openai.response.usage.total_tokens"],
-)
+@pytest.mark.snapshot
+@pytest.mark.skipif(TIKTOKEN_AVAILABLE, reason="This test estimates token counts")
+def test_completion_stream_est_tokens(openai, openai_vcr, mock_metrics, snapshot_tracer):
+    with openai_vcr.use_cassette("completion_streamed.yaml"):
+        with mock.patch("ddtrace.contrib.internal.openai.utils.encoding_for_model", create=True) as mock_encoding:
+            mock_encoding.return_value.encode.side_effect = lambda x: [1, 2]
+            client = openai.OpenAI()
+            resp = client.completions.create(model="ada", prompt="Hello world", stream=True, n=None)
+            _ = [c for c in resp]
+
+
+@pytest.mark.skipif(not TIKTOKEN_AVAILABLE, reason="This test computes token counts using tiktoken")
+@pytest.mark.snapshot(token="tests.contrib.openai.test_openai.test_completion_stream")
 def test_completion_stream(openai, openai_vcr, mock_metrics, snapshot_tracer):
     with openai_vcr.use_cassette("completion_streamed.yaml"):
         with mock.patch("ddtrace.contrib.internal.openai.utils.encoding_for_model", create=True) as mock_encoding:
@@ -934,10 +943,8 @@ def test_completion_stream(openai, openai_vcr, mock_metrics, snapshot_tracer):
             _ = [c for c in resp]
 
 
-@pytest.mark.snapshot(
-    token="tests.contrib.openai.test_openai.test_completion_stream",
-    ignores=["metrics.openai.response.usage.completion_tokens", "metrics.openai.response.usage.total_tokens"],
-)
+@pytest.mark.skipif(not TIKTOKEN_AVAILABLE, reason="This test computes token counts using tiktoken")
+@pytest.mark.snapshot(token="tests.contrib.openai.test_openai.test_completion_stream")
 async def test_completion_async_stream(openai, openai_vcr, mock_metrics, snapshot_tracer):
     with openai_vcr.use_cassette("completion_streamed.yaml"):
         with mock.patch("ddtrace.contrib.internal.openai.utils.encoding_for_model", create=True) as mock_encoding:
@@ -948,13 +955,10 @@ async def test_completion_async_stream(openai, openai_vcr, mock_metrics, snapsho
 
 
 @pytest.mark.skipif(
-    parse_version(openai_module.version.VERSION) < (1, 6, 0),
+    parse_version(openai_module.version.VERSION) < (1, 6, 0) and not TIKTOKEN_AVAILABLE,
     reason="Streamed response context managers are only available v1.6.0+",
 )
-@pytest.mark.snapshot(
-    token="tests.contrib.openai.test_openai.test_completion_stream",
-    ignores=["metrics.openai.response.usage.completion_tokens", "metrics.openai.response.usage.total_tokens"],
-)
+@pytest.mark.snapshot(token="tests.contrib.openai.test_openai.test_completion_stream")
 def test_completion_stream_context_manager(openai, openai_vcr, mock_metrics, snapshot_tracer):
     with openai_vcr.use_cassette("completion_streamed.yaml"):
         with mock.patch("ddtrace.contrib.internal.openai.utils.encoding_for_model", create=True) as mock_encoding:

tests/snapshots/tests.contrib.openai.test_openai_v1.test_completion_stream_est_tokens.json

@@ -0,0 +1,49 @@
+[[
+  {
+    "name": "openai.request",
+    "service": "tests.contrib.openai",
+    "resource": "createCompletion",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "",
+    "error": 0,
+    "meta": {
+      "_dd.p.dm": "-0",
+      "_dd.p.tid": "677c221c00000000",
+      "component": "openai",
+      "language": "python",
+      "openai.base_url": "https://api.openai.com/v1/",
+      "openai.organization.name": "datadog-4",
+      "openai.request.client": "OpenAI",
+      "openai.request.endpoint": "/v1/completions",
+      "openai.request.method": "POST",
+      "openai.request.model": "ada",
+      "openai.request.n": "None",
+      "openai.request.prompt.0": "Hello world",
+      "openai.request.stream": "True",
+      "openai.response.choices.0.finish_reason": "length",
+      "openai.response.choices.0.text": "! ... A page layouts page drawer? ... Interesting. The \"Tools\" is",
+      "openai.response.model": "ada",
+      "openai.user.api_key": "sk-...key>",
+      "runtime-id": "24f8e851c87e4f758c73d6acd0aaf82b"
+    },
+    "metrics": {
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "openai.organization.ratelimit.requests.limit": 3000,
+      "openai.organization.ratelimit.requests.remaining": 2999,
+      "openai.organization.ratelimit.tokens.limit": 250000,
+      "openai.organization.ratelimit.tokens.remaining": 249979,
+      "openai.request.prompt_tokens_estimated": 1,
+      "openai.response.completion_tokens_estimated": 1,
+      "openai.response.usage.completion_tokens": 16,
+      "openai.response.usage.prompt_tokens": 2,
+      "openai.response.usage.total_tokens": 18,
+      "process_id": 47101
+    },
+    "duration": 37957000,
+    "start": 1736188444222291000
+  }]]