From 5531b92f0537043c93831ed299574951cb44a669 Mon Sep 17 00:00:00 2001 From: Yun Kim Date: Mon, 6 Jan 2025 13:35:42 -0500 Subject: [PATCH] Fix completion stream token snapshots --- tests/contrib/openai/test_openai_v1.py | 30 +++++++----- ..._v1.test_completion_stream_est_tokens.json | 49 +++++++++++++++++++ 2 files changed, 66 insertions(+), 13 deletions(-) create mode 100644 tests/snapshots/tests.contrib.openai.test_openai_v1.test_completion_stream_est_tokens.json diff --git a/tests/contrib/openai/test_openai_v1.py b/tests/contrib/openai/test_openai_v1.py index e7b188610a6..91737d9e5eb 100644 --- a/tests/contrib/openai/test_openai_v1.py +++ b/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) or 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: diff --git a/tests/snapshots/tests.contrib.openai.test_openai_v1.test_completion_stream_est_tokens.json b/tests/snapshots/tests.contrib.openai.test_openai_v1.test_completion_stream_est_tokens.json new file mode 100644 index 00000000000..445dc39db98 --- /dev/null +++ b/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 + }]]