diff --git a/.github/workflows/parallel-python-tests.yml b/.github/workflows/parallel-python-tests.yml index 553945dd8..97445cf90 100644 --- a/.github/workflows/parallel-python-tests.yml +++ b/.github/workflows/parallel-python-tests.yml @@ -37,7 +37,7 @@ jobs: distribution: wheel - python: '3.13' parallel: '1.3.7' - openai: '3.42.0' + openai: '3.42.1' distribution: sdist env: FI_BASE_URL: http://127.0.0.1:9 @@ -65,7 +65,8 @@ jobs: packages=(python/frameworks/{parallel,openai}/dist/*.tar.gz) fi .parallel-test-venv/bin/pip install "${packages[@]}" \ - "parallel-web==$PARALLEL_VERSION" "openai==$OPENAI_VERSION" pytest pytest-asyncio + "parallel-web==$PARALLEL_VERSION" "openai==$OPENAI_VERSION" \ + 'fi-instrumentation-otel>=1.2.2' pytest pytest-asyncio - name: Test with only loopback networking shell: bash run: | diff --git a/python/examples/perplexity/tests/test_perplexity_recipe.py b/python/examples/perplexity/tests/test_perplexity_recipe.py index 6ab2c93a2..b93530b1a 100644 --- a/python/examples/perplexity/tests/test_perplexity_recipe.py +++ b/python/examples/perplexity/tests/test_perplexity_recipe.py @@ -37,6 +37,19 @@ == Path(CORE_DISTRIBUTION.locate_file("fi_instrumentation/__init__.py")).resolve() ) +# The recipe CI deliberately retains the old published SDK as a compatibility +# control. Only that exact installed distribution exports echoed instructions; +# current source must not reintroduce them to satisfy legacy expectations. +try: + OPENAI_DISTRIBUTION = distribution("traceAI-openai") +except PackageNotFoundError: + LEGACY_OUTPUT = False +else: + LEGACY_OUTPUT = OPENAI_DISTRIBUTION.version == "0.1.10" and ( + Path(traceai_openai.__file__).resolve() + == Path(OPENAI_DISTRIBUTION.locate_file("traceai_openai/__init__.py")).resolve() + ) + TEST_DIR = Path(__file__).resolve().parent RECIPE_DIR = TEST_DIR.parent VENDOR_KEY = "placeholder-perplexity-key" @@ -254,7 +267,10 @@ def test_responses_at_documented_host(tracing): assert raw["usage"]["input_tokens_details"][key] == value assert raw["usage"]["output_tokens_details"] == fake.USAGE["output_tokens_details"] assert raw["usage"]["tool_calls_details"] == fake.USAGE["tool_calls_details"] - assert fake.INSTRUCTIONS_MARKER in raw["instructions"] + if LEGACY_OUTPUT: + assert fake.INSTRUCTIONS_MARKER in raw["instructions"] + else: + assert "instructions" not in raw assert fake.SEARCH_URL in attrs["output.value"] assert fake.SEARCH_SNIPPET in attrs["output.value"] @@ -349,9 +365,11 @@ def handler(request): assert _call(handler, prompt=marker) == fake.ANSWER span, _ = _assert_export(receiver, provider, project, hide_inputs=hidden) assert (marker in json.dumps(span)) is (not hidden) - # Output echoes remain visible even when the input itself is hidden. - for output_marker in (fake.QUERY_MARKER, fake.INSTRUCTIONS_MARKER): - assert output_marker in _attributes(span)["output.value"] + # Search queries are response-owned output; instructions are input. + output = _attributes(span)["output.value"] + assert fake.QUERY_MARKER in output + assert (fake.INSTRUCTIONS_MARKER in output) is LEGACY_OUTPUT + assert fake.ANSWER in output @pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) @@ -363,8 +381,11 @@ def test_hide_outputs_with_control(tracing, monkeypatch, stream): span, attrs = _assert_export(receiver, provider, project) _assert_usage(attrs) assert (fake.ANSWER in json.dumps(span)) is (not hidden) - for marker in (fake.SEARCH_URL, fake.SEARCH_SNIPPET, fake.QUERY_MARKER, fake.INSTRUCTIONS_MARKER): + for marker in (fake.SEARCH_URL, fake.SEARCH_SNIPPET, fake.QUERY_MARKER): assert (marker in json.dumps(span)) is (not hidden and not stream) + assert (fake.INSTRUCTIONS_MARKER in json.dumps(span)) is ( + LEGACY_OUTPUT and not hidden and not stream + ) @pytest.mark.parametrize("hidden", [False, True], ids=["visible", "hidden"]) @@ -381,8 +402,9 @@ def test_empty_text_stream_raw_fallback(tracing, monkeypatch, hidden): list(client.responses.create(model=fake.REQUEST_MODEL, input="empty stream", stream=True)) span, attrs = _assert_export(receiver, provider, project) _assert_usage(attrs) - for marker in (fake.SEARCH_URL, fake.QUERY_MARKER, fake.INSTRUCTIONS_MARKER): + for marker in (fake.SEARCH_URL, fake.QUERY_MARKER): assert (marker in json.dumps(span)) is (not hidden) + assert (fake.INSTRUCTIONS_MARKER in json.dumps(span)) is (LEGACY_OUTPUT and not hidden) def test_cache_details_require_output_details(tracing): @@ -457,7 +479,7 @@ def test_subprocess_app(tmp_path, stream, cli_model): assert (fake.SEARCH_URL in json.dumps(attrs)) is (not stream) assert (fake.SEARCH_SNIPPET in json.dumps(attrs)) is (not stream) assert (fake.QUERY_MARKER in json.dumps(attrs)) is (not stream) - assert (fake.INSTRUCTIONS_MARKER in json.dumps(attrs)) is (not stream) + assert (fake.INSTRUCTIONS_MARKER in json.dumps(attrs)) is (LEGACY_OUTPUT and not stream) [export] = receiver.requests() assert export["path"] == "/tracer/v1/traces" assert export["headers"]["x-api-key"] == FI_KEY diff --git a/python/fi_instrumentation/instrumentation/_tracers.py b/python/fi_instrumentation/instrumentation/_tracers.py index 779854676..24f2db811 100644 --- a/python/fi_instrumentation/instrumentation/_tracers.py +++ b/python/fi_instrumentation/instrumentation/_tracers.py @@ -203,7 +203,21 @@ def start_as_current_span( record_exception=record_exception, set_status_on_exception=set_status_on_exception, ) as current_span: - yield cast(FiSpan, current_span) + try: + yield cast(FiSpan, current_span) + except BaseException as error: + # OTel use_span handles Exception, but deliberately excludes + # cancellation, KeyboardInterrupt and other local aborts. Record + # those before use_span ends the span; leave ordinary errors to + # OTel so they are not recorded twice. FiSpan applies privacy. + if not isinstance(error, Exception) and span.is_recording(): + if record_exception: + span.record_exception(error, escaped=True) + if set_status_on_exception: + span.set_status(Status( + StatusCode.ERROR, f"{type(error).__name__}: {error}" + )) + raise def start_span( self, diff --git a/python/frameworks/openai/pyproject.toml b/python/frameworks/openai/pyproject.toml index 3a189f628..690651844 100644 --- a/python/frameworks/openai/pyproject.toml +++ b/python/frameworks/openai/pyproject.toml @@ -28,7 +28,7 @@ classifiers = [ [tool.poetry.dependencies] python = ">3.9,<3.14" -fi-instrumentation-otel = ">=0.1.11" +fi-instrumentation-otel = ">=1.2.2" openai = ">=1.69.0" [tool.poetry.urls] diff --git a/python/frameworks/openai/tests/test_chat_stream_completion.py b/python/frameworks/openai/tests/test_chat_stream_completion.py new file mode 100644 index 000000000..695bfe8f8 --- /dev/null +++ b/python/frameworks/openai/tests/test_chat_stream_completion.py @@ -0,0 +1,181 @@ +"""HTTP EOF is not evidence that every requested chat choice completed.""" +import asyncio +import json + +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig +from opentelemetry.trace import StatusCode + +from .test_openai_compatible_spans import ( + REQUEST_MODEL, async_client, client, make_tracer, only_span, +) + + +def chunk(index=0, finish_reason=None, content=None): + return { + "id": "synthetic", "object": "chat.completion.chunk", "created": 1, + "model": REQUEST_MODEL, + "choices": [{"index": index, "delta": {"content": content} if content else {}, + "finish_reason": finish_reason}], + } + + +def stream_response(chunks, done): + body = "".join(f"data: {json.dumps(c)}\n\n" for c in chunks) + if done: + body += "data: [DONE]\n\n" + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body.encode()) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("done", [False, True]) +@pytest.mark.parametrize("case", [ + "empty", "partial", "usage_only", "stop", "length", "tool_calls", "content_filter", + "function_call", "two_complete", "second_partial", "second_missing", "usage_tail", +]) +def test_chat_completion_requires_terminal_choices(make_tracer, asynchronous, done, case): + exporter = make_tracer(TraceConfig(hide_inputs=True, hide_outputs=True, hide_error_details=True)) + chunks = [chunk(content="synthetic-private-output")] + n = 1 + complete = case in {"stop", "length", "tool_calls", "content_filter", "function_call", "two_complete", "usage_tail"} + usage = {**chunk(), "choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}} + if case == "empty": + chunks = [] + elif case == "usage_only": + chunks = [usage] + elif case in {"two_complete", "second_partial", "second_missing"}: + n = 2 + chunks.append(chunk(finish_reason="stop")) + if case != "second_missing": + chunks.append(chunk(index=1, content="second")) + if case == "two_complete": + chunks.append(chunk(index=1, finish_reason="stop")) + elif case == "usage_tail": + chunks.extend([chunk(finish_reason="stop"), usage]) + elif complete: + chunks.append(chunk(finish_reason=case)) + params = dict(model=REQUEST_MODEL, messages=[{"role": "user", "content": "synthetic-input"}], stream=True, n=n) + + async def run_async(): + async with async_client(lambda _: stream_response(chunks, done)) as sdk: + stream = await sdk.chat.completions.create(**params) + actual = [c async for c in stream] + assert len(actual) == len(chunks) + await stream.close() + + if asynchronous: + asyncio.run(run_async()) + else: + with client(lambda _: stream_response(chunks, done)) as sdk: + stream = sdk.chat.completions.create(**params) + actual = list(stream) + assert len(actual) == len(chunks) + stream.close() + span, attrs = only_span(exporter) + assert span.status.status_code == (StatusCode.OK if complete else StatusCode.ERROR) + assert "synthetic-private-output" not in str(attrs) + if case == "usage_tail": + assert attrs["gen_ai.usage.total_tokens"] == 5 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("declared_n,override_n,completed_choices,expected_status", [ + pytest.param(2, 1, 1, StatusCode.OK, id="override-down"), + pytest.param(1, 2, 1, StatusCode.ERROR, id="override-up-incomplete"), + pytest.param(1, 2, 2, StatusCode.OK, id="override-up-complete"), + pytest.param(2, openai.Omit(), 1, StatusCode.OK, id="remove-n"), + pytest.param(2, None, 1, StatusCode.OK, id="null-n"), + pytest.param(openai.NOT_GIVEN, 2, 1, StatusCode.ERROR, id="extra-body-only-n"), +]) +def test_completion_count_uses_effective_body( + make_tracer, asynchronous, declared_n, override_n, completed_choices, expected_status, +): + exporter = make_tracer() + chunks = [chunk(index=i, finish_reason="stop") for i in range(completed_choices)] + wire_bodies = [] + + def handler(request): + body = json.loads(request.content) + wire_bodies.append(body) + if isinstance(override_n, openai.Omit): + assert "n" not in body + else: + assert body["n"] == override_n + return stream_response(chunks, done=True) + + extra_body = {"n": override_n, "synthetic_extra": "not-exported"} + params = dict( + model=REQUEST_MODEL, messages=[{"role": "user", "content": "synthetic-input"}], + stream=True, n=declared_n, extra_body=extra_body, + ) + + async def run_async(): + async with async_client(handler) as sdk: + stream = await sdk.chat.completions.create(**params) + assert len([c async for c in stream]) == completed_choices + await stream.close() + + if asynchronous: + asyncio.run(run_async()) + else: + with client(handler) as sdk: + stream = sdk.chat.completions.create(**params) + assert len(list(stream)) == completed_choices + stream.close() + + assert len(wire_bodies) == 1 + assert extra_body == {"n": override_n, "synthetic_extra": "not-exported"} + span, attrs = only_span(exporter) + assert span.status.status_code == expected_status + # This fix must not expand exported metadata to arbitrary extra_body fields. + parameters = json.loads(attrs["gen_ai.request.parameters"]) + if declared_n is openai.NOT_GIVEN: + assert "n" not in parameters + else: + assert parameters["n"] == declared_n + assert "synthetic_extra" not in str(attrs) + assert "not-exported" not in str(attrs) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("terminal", [False, True]) +@pytest.mark.parametrize("consumer_error", [False, True]) +def test_context_exit_uses_observed_terminal_state(make_tracer, asynchronous, terminal, consumer_error): + exporter = make_tracer() + chunks = [chunk(content="partial", finish_reason="stop" if terminal else None)] + params = dict(model=REQUEST_MODEL, messages=[{"role": "user", "content": "hi"}], stream=True) + error = ValueError("synthetic consumer error") + + async def run_async(): + async with async_client(lambda _: stream_response(chunks, True)) as sdk: + stream = await sdk.chat.completions.create(**params) + async with stream: + await stream.__anext__() + if consumer_error: + raise error + + def run_sync(): + with client(lambda _: stream_response(chunks, True)) as sdk: + stream = sdk.chat.completions.create(**params) + with stream: + next(stream) + if consumer_error: + raise error + + def run(): + if asynchronous: + asyncio.run(run_async()) + else: + run_sync() + + if consumer_error: + with pytest.raises(ValueError) as caught: + run() + assert caught.value is error + else: + run() + span, _ = only_span(exporter) + assert span.status.status_code == (StatusCode.OK if terminal and not consumer_error else StatusCode.ERROR) + assert len([e for e in span.events if e.name == "exception"]) == int(consumer_error) diff --git a/python/frameworks/openai/tests/test_core_dependency_privacy.py b/python/frameworks/openai/tests/test_core_dependency_privacy.py new file mode 100644 index 000000000..160b0a30b --- /dev/null +++ b/python/frameworks/openai/tests/test_core_dependency_privacy.py @@ -0,0 +1,42 @@ +"""Synthetic fixture to verify declared core compatibility and hidden error details.""" +import asyncio +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from traceai_openai import OpenAIInstrumentor + +MESSAGE = 'synthetic-private-upstream-error-483129' + +@pytest.mark.parametrize('async_', [False, True]) +def test_hidden_failed_response(async_, monkeypatch): + monkeypatch.setenv('FI_HIDE_ERROR_DETAILS', 'true') + exporter = InMemorySpanExporter() + provider = register(project_name='core-floor-probe', project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False) + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True)) + body = {'id': 'resp_probe', 'object': 'response', 'created_at': 1, 'model': 'synthetic', + 'status': 'failed', 'output': [], 'error': {'code': 'server_error', 'message': MESSAGE}, + 'usage': {'input_tokens': 1, 'output_tokens': 0, 'total_tokens': 1}} + def handler(req): return httpx.Response(200, json=body) + def make_client(): + http = (httpx.AsyncClient if async_ else httpx.Client)(transport=httpx.MockTransport(handler), trust_env=False) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)(api_key='synthetic-key', base_url='https://api.parallel.ai/v1', max_retries=0, http_client=http) + try: + if async_: + async def run(): + async with make_client() as client: + assert (await client.responses.create(model='synthetic', input='private input')).status == 'failed' + asyncio.run(run()) + else: + with make_client() as client: + assert client.responses.create(model='synthetic', input='private input').status == 'failed' + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert MESSAGE not in spans[0].to_json() + finally: + instrumentor.uninstrument() + provider.shutdown() diff --git a/python/frameworks/openai/tests/test_effective_body_and_raw_consumption.py b/python/frameworks/openai/tests/test_effective_body_and_raw_consumption.py new file mode 100644 index 000000000..74a5c053d --- /dev/null +++ b/python/frameworks/openai/tests/test_effective_body_and_raw_consumption.py @@ -0,0 +1,260 @@ +"""Effective credentials and lazy body lifecycles through real SDK transports.""" + +import asyncio +from copy import deepcopy +import json +from types import SimpleNamespace + +try: + import httpx2 as httpx +except ImportError: + import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor +from traceai_openai._privacy import request_credentials + + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register( + project_name="body-consumption", + project_type=ProjectType.OBSERVE, + span_exporter=exporter, + batch=False, + verbose=False, + ) + OpenAIInstrumentor().instrument(tracer_provider=provider, config=TraceConfig()) + try: + yield provider, exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.force_flush() + provider.shutdown() + + +def sdk(handler, async_=False): + http = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), + trust_env=False, + ) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key="synthetic-static-key", + max_retries=0, + http_client=http, + ) + + +def response_body(status="completed"): + return { + "id": "resp_review", + "object": "response", + "created_at": 1, + "model": "synthetic", + "status": status, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "output": [], + "error": None, + "usage": { + "input_tokens": 5, + "output_tokens": 2, + "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("extra", [False, True]) +@pytest.mark.parametrize("credential", ["authorization", "header"]) +def test_mcp_credentials_in_effective_body(tracing, async_, extra, credential): + _, exporter = tracing + secret = 'synthetic-extra-mcp-credential-17498-"quoted"' + tool = { + "type": "mcp", + "server_label": "synthetic", + "server_url": "https://synthetic.invalid/mcp", + **( + {"authorization": secret} + if credential == "authorization" + else {"headers": {"X-Private-Token": secret}} + ), + } + original_tool = deepcopy(tool) + captured = [] + + def handler(request): + captured.append(json.loads(request.content)) + assert captured[-1]["tools"][0] == tool + return httpx.Response( + 400, + json={ + "error": { + "code": "invalid_request", + "message": "MCP authentication failed for " + secret, + } + }, + ) + + params = {"model": "synthetic", "input": "question"} + params.update( + {"tools": [], "extra_body": {"tools": [tool]}} if extra else {"tools": [tool]} + ) + + async def run(): + async with sdk(handler, True) as client: + await client.responses.create(**params) + + with pytest.raises(openai.BadRequestError) as caught: + if async_: + asyncio.run(run()) + else: + with sdk(handler) as client: + client.responses.create(**params) + assert captured and tool == original_tool + assert secret in caught.value.body["message"] + (span,) = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + assert "synthetic-extra-mcp-credential" not in span.to_json() + if extra: + assert "synthetic.invalid/mcp" not in span.to_json() + + +@pytest.mark.parametrize("override", [None, [], openai.Omit()]) +def test_effective_tools_override_original_without_mutation(override): + original = {"tools": [{"type": "mcp", "authorization": "synthetic-original"}]} + extra = {"tools": override} + secrets = request_credentials(None, SimpleNamespace(extra_json=extra), original) + assert "synthetic-original" not in secrets + assert original["tools"][0]["authorization"] == "synthetic-original" + assert extra["tools"] is override + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("mode", ["parse", "json", "text", "read"]) +@pytest.mark.parametrize("status", ["completed", "failed", "incomplete", "cancelled"]) +def test_raw_body_consumption(tracing, async_, mode, status): + _, exporter = tracing + body = response_body(status) + + def handler(request): + return httpx.Response(200, json=body) + + def check(result): + if mode == "parse": + assert result.status == status + elif mode == "json": + assert result == body + else: + assert json.loads(result) == body + + async def run(): + async with sdk(handler, True) as client: + async with client.responses.with_streaming_response.create( + model="synthetic", + input="question", + ) as raw: + check(await getattr(raw, mode)()) + check(await getattr(raw, mode)()) + await raw.close() + + if async_: + asyncio.run(run()) + else: + with sdk(handler) as client: + with client.responses.with_streaming_response.create( + model="synthetic", + input="question", + ) as raw: + check(getattr(raw, mode)()) + check(getattr(raw, mode)()) + raw.close() + (span,) = exporter.get_finished_spans() + assert span.status.status_code is ( + StatusCode.OK if status == "completed" else StatusCode.ERROR + ) + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("mode", ["read", "text", "json"]) +@pytest.mark.parametrize("outcome", ["completed", "unread", "error", "cancel", "sse"]) +def test_raw_consumption_is_lazy_and_errors_finish_once(tracing, async_, mode, outcome): + _, exporter = tracing + reads = [] + closes = [] + error = ( + KeyboardInterrupt("synthetic abort") + if outcome == "cancel" + else RuntimeError("synthetic read failure") + ) + body = json.dumps(response_body()).encode() + + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + reads.append(1) + if outcome in {"error", "cancel"}: + raise error + yield body + + def close(self): + closes.append(1) + + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + reads.append(1) + if outcome in {"error", "cancel"}: + raise error + yield body + + async def aclose(self): + closes.append(1) + + def handler(request): + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + stream=AsyncBody() if async_ else SyncBody(), + ) + + params = {"model": "synthetic", "input": "question", "stream": outcome == "sse"} + + async def run(): + async with sdk(handler, True) as client: + async with client.responses.with_streaming_response.create(**params) as raw: + assert not reads + assert not exporter.get_finished_spans() + if outcome != "unread": + await getattr(raw, mode)() + await raw.close() + + def run_sync(): + with sdk(handler) as client: + with client.responses.with_streaming_response.create(**params) as raw: + assert not reads + assert not exporter.get_finished_spans() + if outcome != "unread": + getattr(raw, mode)() + raw.close() + + if outcome in {"error", "cancel"}: + with pytest.raises(type(error)) as caught: + asyncio.run(run()) if async_ else run_sync() + assert caught.value is error + else: + asyncio.run(run()) if async_ else run_sync() + assert len(reads) == (0 if outcome == "unread" else 1) + assert closes + (span,) = exporter.get_finished_spans() + assert span.status.status_code is ( + StatusCode.OK if outcome == "completed" else StatusCode.ERROR + ) + if outcome in {"error", "cancel"}: + assert len([event for event in span.events if event.name == "exception"]) == 1 diff --git a/python/frameworks/openai/tests/test_fireworks_route.py b/python/frameworks/openai/tests/test_fireworks_route.py index d5698b981..a09194047 100644 --- a/python/frameworks/openai/tests/test_fireworks_route.py +++ b/python/frameworks/openai/tests/test_fireworks_route.py @@ -90,6 +90,21 @@ def response(outcome): } for delta in deltas ] + # A successful generation finishes its choice before the usage-only + # tail. [DONE] alone is transport termination, not model completion. + chunks.append( + { + "id": "c", + "object": "chat.completion.chunk", + "created": 1, + "model": MODEL, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": "tool_calls" if outcome == "tool_stream" else "stop", + }], + } + ) chunks.append( { "id": "c", diff --git a/python/frameworks/openai/tests/test_raw_envelope_contract.py b/python/frameworks/openai/tests/test_raw_envelope_contract.py new file mode 100644 index 000000000..7f8bba9d0 --- /dev/null +++ b/python/frameworks/openai/tests/test_raw_envelope_contract.py @@ -0,0 +1,189 @@ +"""Raw OpenAI SDK response contracts using synthetic in-memory transport.""" + +import asyncio +import json + +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register( + project_name="raw-review", project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False, + ) + config = TraceConfig( + hide_inputs=True, hide_input_messages=True, + hide_llm_invocation_parameters=True, hide_error_details=True, + ) + OpenAIInstrumentor().instrument(tracer_provider=provider, config=config) + try: + yield exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.shutdown() + + +def body(status="completed"): + return { + "id": "resp_review", "object": "response", "created_at": 1, + "model": "synthetic", "status": status, "parallel_tool_calls": True, + "tool_choice": "auto", "tools": [], "output": [], "error": None, + "usage": { + "input_tokens": 5, "output_tokens": 2, "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +def client(payload, async_): + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + yield payload + + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + yield payload + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "application/json"}, + stream=AsyncBody() if async_ else SyncBody(), + ) + + http_client = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False, + ) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key="synthetic-credential", max_retries=0, http_client=http_client, + ) + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("parse_to", ["default", "dict", "str"]) +def test_raw_custom_parse_respects_input_and_error_privacy(tracing, async_, parse_to): + hidden_input = "synthetic-private-input-24193" + hidden_error = "synthetic-private-provider-error-88123" + response = body("failed") + response.update( + instructions=hidden_input, metadata={"private": hidden_input}, + input=hidden_input, error={"code": "server_error", "message": hidden_error}, + ) + payload = json.dumps(response).encode() + kwargs = {} if parse_to == "default" else { + "to": dict[str, object] if parse_to == "dict" else str, + } + + async def run_async(): + async with client(payload, True) as c: + async with c.responses.with_streaming_response.create( + model="synthetic", input=hidden_input, instructions=hidden_input, + ) as raw: + return await raw.parse(**kwargs) + + if async_: + result = asyncio.run(run_async()) + else: + with client(payload, False) as c: + with c.responses.with_streaming_response.create( + model="synthetic", input=hidden_input, instructions=hidden_input, + ) as raw: + result = raw.parse(**kwargs) + + # The application still receives the original unmodified response. + application_data = ( + result if isinstance(result, str) else json.dumps(result) if isinstance(result, dict) + else result.model_dump_json() + ) + assert hidden_error in application_data + span, = tracing.get_finished_spans() + assert hidden_input not in span.to_json() + assert hidden_error not in span.to_json() + assert span.status.status_code is StatusCode.ERROR + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("status", ["completed", "failed"]) +@pytest.mark.parametrize("mode", ["read", "iter_bytes", "iter_text", "iter_lines"]) +def test_non_sse_raw_iteration_uses_body_status(tracing, async_, status, mode): + payload = json.dumps(body(status)).encode() + + async def run_async(): + async with client(payload, True) as c: + async with c.responses.with_streaming_response.create( + model="synthetic", input="question", + ) as raw: + if mode == "read": + assert await raw.read() == payload + else: + chunks = [part async for part in getattr(raw, mode)()] + actual = (b"" if mode == "iter_bytes" else "").join(chunks) + assert actual == (payload if mode == "iter_bytes" else payload.decode()) + + if async_: + asyncio.run(run_async()) + else: + with client(payload, False) as c: + with c.responses.with_streaming_response.create( + model="synthetic", input="question", + ) as raw: + if mode == "read": + assert raw.read() == payload + else: + chunks = list(getattr(raw, mode)()) + actual = (b"" if mode == "iter_bytes" else "").join(chunks) + assert actual == (payload if mode == "iter_bytes" else payload.decode()) + + span, = tracing.get_finished_spans() + expected = StatusCode.OK if status == "completed" else StatusCode.ERROR + assert span.status.status_code is expected + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("instrumented", [False, True]) +@pytest.mark.parametrize("endpoint", ["responses", "chat"]) +def test_legacy_raw_text_remains_string(tracing, async_, instrumented, endpoint): + if not instrumented: + OpenAIInstrumentor().uninstrument() + response = body() + if endpoint == "chat": + response = { + "id": "chat_review", "object": "chat.completion", "created": 1, + "model": "synthetic", "choices": [{ + "index": 0, "message": {"role": "assistant", "content": "answer"}, + "finish_reason": "stop", + }], + } + payload = json.dumps(response).encode() + + def params(c): + if endpoint == "responses": + return c.responses, {"model": "synthetic", "input": "question"} + return c.chat.completions, { + "model": "synthetic", "messages": [{"role": "user", "content": "question"}], + } + + async def run_async(): + async with client(payload, True) as c: + resource, kwargs = params(c) + raw = await resource.with_raw_response.create(**kwargs) + assert raw.content == payload + assert json.loads(raw.text) == response + + if async_: + asyncio.run(run_async()) + else: + with client(payload, False) as c: + resource, kwargs = params(c) + raw = resource.with_raw_response.create(**kwargs) + assert raw.content == payload + assert json.loads(raw.text) == response diff --git a/python/frameworks/openai/tests/test_request_body.py b/python/frameworks/openai/tests/test_request_body.py new file mode 100644 index 000000000..1351fa537 --- /dev/null +++ b/python/frameworks/openai/tests/test_request_body.py @@ -0,0 +1,43 @@ +"""Request parsing follows the SDK's shallow extra_body merge without mutation.""" +from types import SimpleNamespace + +import openai +import pytest + +from traceai_openai._request import _parse_request_args + + +@pytest.mark.parametrize("json_data,extra_json,expected", [ + pytest.param({"n": 2}, {"n": 1}, {"n": 1}, id="override"), + pytest.param({"n": 2}, {"n": openai.Omit()}, {}, id="remove"), + pytest.param({"n": 2}, {"n": None}, {"n": None}, id="preserve-null"), + pytest.param(None, {"n": 2}, {"n": 2}, id="extra-only"), + pytest.param({"n": 2}, None, {"n": 2}, id="no-extra"), + pytest.param(None, None, {}, id="no-body"), + pytest.param( + {"metadata": {"old": "value", "shared": "old"}}, + {"metadata": {"shared": "new"}}, + {"metadata": {"shared": "new"}}, + id="shallow-replacement", + ), +]) +def test_parse_effective_body(json_data, extra_json, expected): + original_json = dict(json_data) if json_data is not None else None + original_extra = dict(extra_json) if extra_json is not None else None + options = SimpleNamespace(json_data=json_data, extra_json=extra_json) + cast_to = openai.types.chat.ChatCompletion + + _, metadata_parameters = _parse_request_args((cast_to, options)) + assert metadata_parameters == (json_data or {}) + actual_type, actual = _parse_request_args((cast_to, options), include_extra_body=True) + + assert actual_type is cast_to + assert actual == expected + assert options.json_data == original_json + assert options.extra_json == original_extra + + +def test_options_without_extra_json_remain_supported(): + options = SimpleNamespace(json_data={"n": 2}) + _, actual = _parse_request_args((openai.types.chat.ChatCompletion, options)) + assert actual == {"n": 2} diff --git a/python/frameworks/openai/tests/test_response_edge_cases.py b/python/frameworks/openai/tests/test_response_edge_cases.py new file mode 100644 index 000000000..de31ae8ed --- /dev/null +++ b/python/frameworks/openai/tests/test_response_edge_cases.py @@ -0,0 +1,185 @@ +"""Round 3 independent real SDK probes; synthetic credentials and MockTransport only.""" +import asyncio +import json +import httpx +import openai +from typing import Any, Dict +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + +SECRET_INPUT = 'synthetic-private-system-instruction-827391' + +def body(status='completed'): + return {'id': 'resp_independent', 'object': 'response', 'created_at': 1, + 'model': 'parallel', 'status': status, 'parallel_tool_calls': True, + 'tool_choice': 'auto', 'tools': [], 'instructions': SECRET_INPUT, + 'output': [{'id': 'msg_1', 'type': 'message', 'role': 'assistant', + 'status': 'completed', 'content': [{'type': 'output_text', + 'text': 'safe public answer', 'annotations': []}]}], + 'error': None, 'usage': {'input_tokens': 5, 'output_tokens': 2, + 'total_tokens': 7, 'input_tokens_details': {'cached_tokens': 0}, + 'output_tokens_details': {'reasoning_tokens': 0}}} + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register(project_name='review3', project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False) + OpenAIInstrumentor().instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True)) + try: + yield provider, exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.force_flush() + provider.shutdown() + +def client(handler, async_=False, base_url='https://api.parallel.ai/v1'): + http = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key='synthetic-api-key-938472', base_url=base_url, + max_retries=0, http_client=http) + +@pytest.mark.parametrize('async_', [False, True]) +@pytest.mark.parametrize('to_dict', [False, True]) +def test_explicit_dict_parse_keeps_inputs_hidden(tracing, async_, to_dict): + provider, exporter = tracing + def handler(request): + return httpx.Response(200, json=body()) + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create( + model='parallel', instructions=SECRET_INPUT, input='question') as raw: + result = await raw.parse(**({'to': Dict[str, Any] if openai.__version__.startswith('1.') else dict} if to_dict else {})) + assert (result['instructions'] if to_dict else result.instructions) == SECRET_INPUT + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create( + model='parallel', instructions=SECRET_INPUT, input='question') as raw: + result = raw.parse(**({'to': Dict[str, Any] if openai.__version__.startswith('1.') else dict} if to_dict else {})) + assert (result['instructions'] if to_dict else result.instructions) == SECRET_INPUT + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert 'safe public answer' in span.attributes['output.value'] + leaked = {k: v for k, v in span.attributes.items() if SECRET_INPUT in str(v)} + assert not leaked, leaked + +@pytest.mark.parametrize('async_', [False, True]) +@pytest.mark.parametrize('method', ['parse', 'json', 'read', 'text', 'iter_lines']) +def test_successful_raw_body_consumption_is_not_error(tracing, async_, method): + provider, exporter = tracing + payload = json.dumps(body()).encode() + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + yield payload + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + yield payload + def handler(request): + return httpx.Response(200, headers={'content-type':'application/json'}, + stream=AsyncBody() if async_ else SyncBody()) + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create(model='parallel', input='question') as raw: + if method == 'iter_lines': + result = [line async for line in raw.iter_lines()] + else: + result = await getattr(raw, method)() + assert result + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create(model='parallel', input='question') as raw: + result = list(raw.iter_lines()) if method == 'iter_lines' else getattr(raw, method)() + assert result + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == StatusCode.OK, span.to_json() + +@pytest.mark.parametrize('async_', [False, True]) +@pytest.mark.parametrize('instrumented', [False, True]) +def test_raw_malformed_body_remains_inspectable(tracing, async_, instrumented): + provider, exporter = tracing + if not instrumented: + OpenAIInstrumentor().uninstrument() + def handler(request): + return httpx.Response(200, headers={'content-type': 'application/json'}, content=b'{"truncated":') + async def run(): + async with client(handler, True) as c: + raw = await c.responses.with_raw_response.create(model='parallel', input='question') + assert raw.status_code == 200 + assert raw.content == b'{"truncated":' + with pytest.raises(json.JSONDecodeError): + raw.parse() + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + raw = c.responses.with_raw_response.create(model='parallel', input='question') + assert raw.status_code == 200 + assert raw.content == b'{"truncated":' + with pytest.raises(json.JSONDecodeError): + raw.parse() + +@pytest.mark.parametrize('async_', [False, True]) +def test_speech_stream_to_file_not_marked_failed(tracing, async_, tmp_path): + from pathlib import Path + provider, exporter = tracing + payload = b'synthetic-audio-data' + target = tmp_path / ('audio-async.bin' if async_ else 'audio-sync.bin') + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + yield payload + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + yield payload + def handler(request): + return httpx.Response(200, headers={'content-type':'audio/mpeg'}, + stream=AsyncBody() if async_ else SyncBody()) + async def run(): + async with client(handler, True, base_url='https://api.openai.com/v1') as c: + async with c.audio.speech.with_streaming_response.create( + model='tts-1', voice='alloy', input='Hello') as raw: + await raw.stream_to_file(target) + if async_: + asyncio.run(run()) + else: + with client(handler, base_url='https://api.openai.com/v1') as c: + with c.audio.speech.with_streaming_response.create( + model='tts-1', voice='alloy', input='Hello') as raw: + raw.stream_to_file(target) + assert target.read_bytes() == payload + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == StatusCode.OK, span.to_json() + + +@pytest.mark.parametrize('async_', [False, True]) +def test_early_aborted_raw_iterator_is_error(tracing, async_): + provider, exporter = tracing + def handler(request): + return httpx.Response(200, headers={'content-type': 'text/plain'}, content=b'line one\nline two\n') + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create(model='parallel', input='question') as raw: + iterator = raw.iter_lines() + assert await iterator.__anext__() == 'line one' + await iterator.aclose() + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create(model='parallel', input='question') as raw: + iterator = raw.iter_lines() + assert next(iterator) == 'line one' + iterator.close() + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == StatusCode.ERROR diff --git a/python/frameworks/openai/tests/test_response_envelopes.py b/python/frameworks/openai/tests/test_response_envelopes.py new file mode 100644 index 000000000..593bd62c6 --- /dev/null +++ b/python/frameworks/openai/tests/test_response_envelopes.py @@ -0,0 +1,455 @@ +"""Independent review probes; all credentials and transports are synthetic.""" + +import asyncio +import json + +import httpx +import openai +import pytest +from fi_instrumentation import register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register( + project_name="review-edges", + project_type=ProjectType.OBSERVE, + span_exporter=exporter, + batch=False, + verbose=False, + ) + OpenAIInstrumentor().instrument(tracer_provider=provider) + try: + yield provider, exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.force_flush() + provider.shutdown() + + +def client(handler, async_=False, key="synthetic-static-key"): + http = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False + ) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key=key, + base_url="https://api.parallel.ai/v1", + max_retries=0, + http_client=http, + ) + + +def body(status="completed"): + return { + "id": "resp_edge", + "object": "response", + "created_at": 1, + "model": "parallel", + "status": status, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "output": [], + "error": None, + "usage": { + "input_tokens": 5, + "output_tokens": 2, + "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +def sse(status): + response = body(status) + if status == "failed": + response["error"] = {"code": "server_error", "message": "synthetic failure"} + kind = "response.created" if status == "in_progress" else "response." + status + return ( + "event: " + + kind + + "\ndata: " + + json.dumps({"type": kind, "sequence_number": 1, "response": response}) + + "\n\n" + ).encode() + + +def test_concurrent_callable_nonstream_credentials(tracing): + if int(openai.__version__.split(".")[0]) < 3: + pytest.skip("callable auth not available at floor") + _provider, exporter = tracing + keys = ["synthetic-alpha-credential-198273", "synthetic-bravo-credential-543215"] + issued = [] + + async def run(): + received = [] + both_sent = asyncio.Event() + + async def token(): + key = keys[len(issued)] + issued.append(key) + return key + + async def handler(request): + key = request.headers["authorization"][7:] + received.append(key) + if len(received) == 2: + both_sent.set() + await asyncio.wait_for(both_sent.wait(), 3) + response = body("failed") + response["error"] = { + "code": "server_error", + "message": "Rejected token " + key, + } + return httpx.Response(200, json=response) + + async with client(handler, True, token) as c: + responses = await asyncio.gather( + *[c.responses.create(model="parallel", input="question") for _ in keys] + ) + assert set(received) == set(keys) + assert {r.error.message for r in responses} == { + "Rejected token " + k for k in keys + } + + asyncio.run(run()) + assert issued == keys + spans = exporter.get_finished_spans() + assert len(spans) == 2 + assert all(s.status.status_code == StatusCode.ERROR for s in spans) + leaked = { + key: [ + name + for span in spans + for name, value in span.attributes.items() + if key in str(value) + ] + for key in keys + } + assert not any(leaked.values()), leaked + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("status", ["failed", "completed", "in_progress"]) +@pytest.mark.parametrize("exhaust", [False, True]) +def test_streaming_response_envelope_lifecycle(tracing, async_, status, exhaust): + provider, exporter = tracing + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=sse(status) + ) + + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=True + ) as raw: + stream = await raw.parse() + if exhaust: + events = [chunk async for chunk in stream] + assert events[-1].response.status == status + else: + assert (await stream.__anext__()).response.status == status + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=True + ) as raw: + stream = raw.parse() + if exhaust: + events = list(stream) + assert events[-1].response.status == status + else: + assert next(stream).response.status == status + provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, f"Expected one finalized span, got {len(spans)}" + (span,) = spans + assert span.status.status_code == ( + StatusCode.OK if status == "completed" else StatusCode.ERROR + ), span.to_json() + if status != "in_progress": + assert span.attributes["gen_ai.usage.total_tokens"] == 7, span.to_json() + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("instrumented", [False, True]) +def test_raw_explicit_stream_parse(tracing, async_, instrumented): + provider, exporter = tracing + if not instrumented: + OpenAIInstrumentor().uninstrument() + from openai.types.responses import ResponseStreamEvent + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=sse("failed") + ) + + async def run(): + async with client(handler, True) as c: + raw = await c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + stream = raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) + assert raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) is stream + events = [event async for event in stream] + assert events[-1].response.status == "failed" + await stream.close() + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + raw = c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + stream = raw.parse(to=openai.Stream[ResponseStreamEvent]) + assert raw.parse(to=openai.Stream[ResponseStreamEvent]) is stream + events = list(stream) + assert events[-1].response.status == "failed" + stream.close() + provider.force_flush() + spans = exporter.get_finished_spans() + if instrumented: + assert len(spans) == 1, f"Expected one finalized span, got {len(spans)}" + assert spans[0].status.status_code is StatusCode.ERROR + assert spans[0].attributes["gen_ai.usage.total_tokens"] == 7 + else: + assert not spans + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("parse", [False, True]) +def test_lazy_envelope_does_not_read_until_requested(tracing, async_, streaming, parse): + provider, exporter = tracing + reads = [] + closes = [] + payload = sse("completed") if streaming else json.dumps(body()).encode() + + class SyncBody(httpx.SyncByteStream): + def __iter__(self): + reads.append("read") + yield payload + + def close(self): + closes.append("close") + + class AsyncBody(httpx.AsyncByteStream): + async def __aiter__(self): + reads.append("read") + yield payload + + async def aclose(self): + closes.append("close") + + def handler(request): + return httpx.Response( + 200, + headers={ + "content-type": "text/event-stream" if streaming else "application/json" + }, + stream=AsyncBody() if async_ else SyncBody(), + ) + + async def run(): + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=streaming + ) as raw: + assert isinstance(raw, openai.AsyncAPIResponse) + assert not reads + assert not exporter.get_finished_spans() + if parse: + result = await raw.parse() + assert await raw.parse() is result + if streaming: + assert not reads + assert (await result.__anext__()).response.status == "completed" + else: + assert result.status == "completed" + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=streaming + ) as raw: + assert isinstance(raw, openai.APIResponse) + assert not reads + assert not exporter.get_finished_spans() + if parse: + result = raw.parse() + assert raw.parse() is result + if streaming: + assert not reads + assert next(result).response.status == "completed" + else: + assert result.status == "completed" + assert reads == (["read"] if parse else []) + assert closes == ["close"] + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == (StatusCode.OK if parse else StatusCode.ERROR) + if parse: + assert span.attributes["gen_ai.usage.total_tokens"] == 7 + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("modern", [False, True]) +@pytest.mark.parametrize("close_http", [False, True]) +def test_explicit_stream_cache_and_envelope_close(tracing, async_, modern, close_http): + provider, exporter = tracing + from openai.types.responses import ResponseStreamEvent + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=sse("completed") + ) + + async def consume(raw): + if modern: + default = await raw.parse() + typed = await raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) + assert await raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) is typed + assert await raw.parse() is default + else: + default = raw.parse() + typed = raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) + assert raw.parse(to=openai.AsyncStream[ResponseStreamEvent]) is typed + assert raw.parse() is default + assert typed is not default + assert (await typed.__anext__()).response.status == "completed" + assert not exporter.get_finished_spans() + if close_http: + await raw.http_response.aclose() + elif modern: + await raw.close() + else: + await typed.close() + # Both cache entries and every close path share one terminal result. + await default.close() + await typed.close() + + async def run(): + async with client(handler, True) as c: + if modern: + async with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=True + ) as raw: + await consume(raw) + else: + raw = await c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + await consume(raw) + + def consume_sync(raw): + default = raw.parse() + typed = raw.parse(to=openai.Stream[ResponseStreamEvent]) + assert raw.parse(to=openai.Stream[ResponseStreamEvent]) is typed + assert raw.parse() is default + assert typed is not default + assert next(typed).response.status == "completed" + assert not exporter.get_finished_spans() + if close_http: + raw.http_response.close() + elif modern: + raw.close() + else: + typed.close() + default.close() + typed.close() + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + if modern: + with c.responses.with_streaming_response.create( + model="parallel", input="question", stream=True + ) as raw: + consume_sync(raw) + else: + raw = c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + consume_sync(raw) + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.OK + assert span.attributes["gen_ai.usage.total_tokens"] == 7 + + +@pytest.mark.parametrize("cancel", [False, True]) +def test_async_lazy_parse_failure_finishes_once(tracing, cancel): + provider, exporter = tracing + entered = None + + class BrokenBody(httpx.AsyncByteStream): + async def __aiter__(self): + entered.set() + if cancel: + await asyncio.Event().wait() + raise httpx.ReadError("synthetic unreadable response") + yield b"" # Make the method an async generator. + + def handler(request): + return httpx.Response( + 200, headers={"content-type": "application/json"}, stream=BrokenBody() + ) + + async def run(): + nonlocal entered + entered = asyncio.Event() + async with client(handler, True) as c: + async with c.responses.with_streaming_response.create( + model="parallel", input="question" + ) as raw: + task = asyncio.create_task(raw.parse()) + await asyncio.wait_for(entered.wait(), 3) + if cancel: + task.cancel() + with pytest.raises( + asyncio.CancelledError if cancel else httpx.ReadError + ): + await task + assert len(exporter.get_finished_spans()) == 1 + + asyncio.run(run()) + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + + +def test_credential_context_is_reset_and_hooks_restore(tracing): + from traceai_openai._request_context import _active_request + + original_sync = OpenAIInstrumentor()._original_process_response + original_async = OpenAIInstrumentor()._original_async_process_response + assert _active_request.get() is None + with client(lambda request: httpx.Response(200, json=body())) as c: + c.responses.create(model="parallel", input="question") + assert _active_request.get() is None + with client( + lambda request: httpx.Response( + 401, json={"error": {"message": "synthetic error"}} + ) + ) as c: + with pytest.raises(openai.AuthenticationError): + c.responses.create(model="parallel", input="question") + assert _active_request.get() is None + OpenAIInstrumentor().uninstrument() + assert openai.OpenAI._process_response is original_sync + assert openai.AsyncOpenAI._process_response is original_async diff --git a/python/frameworks/openai/tests/test_responses_privacy_and_raw_streams.py b/python/frameworks/openai/tests/test_responses_privacy_and_raw_streams.py new file mode 100644 index 000000000..f9a83ae2e --- /dev/null +++ b/python/frameworks/openai/tests/test_responses_privacy_and_raw_streams.py @@ -0,0 +1,411 @@ +"""Responses privacy and raw-envelope contracts using synthetic SDK transports.""" + +import asyncio +import inspect +import json + +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + +KEY = "synthetic-review-key-not-for-export" +ERROR = "synthetic-private-tool-error" + + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register( + project_name="independent-review", + project_type=ProjectType.OBSERVE, + span_exporter=exporter, + batch=False, + verbose=False, + ) + try: + yield provider, exporter + finally: + OpenAIInstrumentor().uninstrument() + provider.force_flush() + provider.shutdown() + + +def client(handler, async_=False, key=KEY): + http = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False + ) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key=key, + base_url="https://api.parallel.ai/v1", + max_retries=0, + http_client=http, + ) + + +def body(status="completed"): + return { + "id": "resp_review", + "object": "response", + "created_at": 1, + "model": "parallel", + "status": status, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "output": [], + "error": None, + "usage": { + "input_tokens": 5, + "output_tokens": 2, + "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +def sse(kind, response): + return ( + "event: " + + kind + + "\ndata: " + + json.dumps({"type": kind, "sequence_number": 1, "response": response}) + + "\n\n" + ).encode() + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("traceback_marker", [False, True]) +def test_key_redaction_survives_core_status_cleaning(tracing, async_, traceback_marker): + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + message = "Invalid credential: " + KEY + if traceback_marker: + message += "\nTraceback (most recent call last):\n synthetic remote failure" + + def handler(request): + assert request.headers["authorization"] == "Bearer " + KEY + return httpx.Response( + 401, json={"error": {"code": "invalid_api_key", "message": message}} + ) + + async def run(): + async with client(handler, True) as c: + await c.responses.create(model="parallel", input="synthetic question") + + with pytest.raises(openai.AuthenticationError) as caught: + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + c.responses.create(model="parallel", input="synthetic question") + (span,) = exporter.get_finished_spans() + assert KEY in str(caught.value) + assert KEY not in span.to_json(), span.to_json() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_callable_credentials_are_scrubbed_without_extra_provider_calls( + tracing, async_ +): + if tuple(map(int, openai.__version__.split(".")[:2])) < (3, 0): + pytest.skip("callable OpenAI credentials are a latest-SDK feature") + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + calls = [] + + def token(): + calls.append(1) + return KEY + + async def atoken(): + return token() + + def handler(request): + assert request.headers["authorization"] == "Bearer " + KEY + return httpx.Response( + 401, + json={ + "error": {"code": "invalid_api_key", "message": "Invalid key: " + KEY} + }, + ) + + async def run(): + async with client(handler, True, atoken) as c: + await c.responses.create(model="parallel", input="synthetic question") + + with pytest.raises(openai.AuthenticationError) as caught: + if async_: + asyncio.run(run()) + else: + with client(handler, key=token) as c: + c.responses.create(model="parallel", input="synthetic question") + (span,) = exporter.get_finished_spans() + assert len(calls) == 1 + assert KEY in str(caught.value) + assert caught.value.request.headers["authorization"] == "Bearer " + KEY + assert KEY not in span.to_json(), span.to_json() + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("hidden", [False, True]) +@pytest.mark.parametrize( + "kind", + ["http_error", "mcp_protocol_error", "mcp_tool_execution_error", "mcp_list_tools"], +) +@pytest.mark.skipif( + "hide_error_details" not in inspect.signature(TraceConfig).parameters, + reason="Error-detail masking requires a core with hide_error_details support", +) +def test_hidden_error_details_cover_mcp_failure_output( + tracing, async_, stream, hidden, kind +): + provider, exporter = tracing + OpenAIInstrumentor().instrument( + tracer_provider=provider, config=TraceConfig(hide_error_details=hidden) + ) + response = body() + response["output"] = [ + { + "id": "mcp_review", + "type": "mcp_call", + "name": "lookup", + "arguments": "{}", + "server_label": "test", + "error": {"type": "http_error", "code": 500, "message": ERROR}, + "output": "ordinary tool result", + } + ] + if kind == "mcp_list_tools": + response["output"] = [ + { + "id": "mcp_list", + "type": kind, + "server_label": "test", + "tools": [], + "error": ERROR, + } + ] + elif kind == "mcp_tool_execution_error": + response["output"][0]["error"] = {"type": kind, "content": {"text": ERROR}} + else: + response["output"][0]["error"]["type"] = kind + + def handler(request): + if stream: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=sse("response.completed", response), + ) + return httpx.Response(200, json=response) + + async def run(): + async with client(handler, True) as c: + result = await c.responses.create( + model="parallel", input="question", stream=stream + ) + if stream: + events = [event async for event in result] + assert ERROR in events[-1].response.model_dump_json() + else: + assert ERROR in result.model_dump_json() + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + result = c.responses.create( + model="parallel", input="question", stream=stream + ) + if stream: + assert ERROR in list(result)[-1].response.model_dump_json() + else: + assert ERROR in result.model_dump_json() + (span,) = exporter.get_finished_spans() + assert (ERROR in span.to_json()) is (not hidden), span.to_json() + output = json.loads(span.attributes["output.value"])["output"][0] + if kind != "mcp_list_tools": + assert output["output"] == "ordinary tool result" + assert output["error"]["type"] == kind + if kind != "mcp_tool_execution_error": + assert output["error"]["code"] == 500 + elif hidden: + assert "error" not in output + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("status", ["failed", "incomplete", "completed", "in_progress"]) +@pytest.mark.parametrize("close_http", [False, True]) +def test_raw_stream_lifecycle_retains_sdk_contract(tracing, async_, status, close_http): + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + response = body(status) + if status == "failed": + response["error"] = {"code": "server_error", "message": "synthetic failure"} + kind = "response.created" if status == "in_progress" else "response." + status + + def handler(request): + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=sse(kind, response), + ) + + async def run(): + async with client(handler, True) as c: + raw = await c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + assert raw.status_code == 200 + assert raw.headers["content-type"] == "text/event-stream" + assert isinstance(raw, openai._legacy_response.LegacyAPIResponse) + assert not exporter.get_finished_spans() + stream = raw.parse() + assert raw.parse() is stream + assert raw.parse(to=openai.types.responses.Response) is stream + assert (await stream.__anext__()).response.status == status + if close_http: + await raw.http_response.aclose() + else: + await stream.close() + await stream.close() + assert raw.http_response.is_closed + + if async_: + asyncio.run(run()) + else: + with client(handler) as c: + raw = c.responses.with_raw_response.create( + model="parallel", input="question", stream=True + ) + assert raw.status_code == 200 + assert raw.headers["content-type"] == "text/event-stream" + assert isinstance(raw, openai._legacy_response.LegacyAPIResponse) + assert not exporter.get_finished_spans() + stream = raw.parse() + assert raw.parse() is stream + assert raw.parse(to=openai.types.responses.Response) is stream + assert next(stream).response.status == status + if close_http: + raw.http_response.close() + else: + stream.close() + stream.close() + assert raw.http_response.is_closed + (span,) = exporter.get_finished_spans() + assert span.status.status_code is ( + StatusCode.OK if status == "completed" else StatusCode.ERROR + ) + if status != "in_progress": + assert span.attributes["gen_ai.usage.input_tokens"] == 5 + assert span.attributes["gen_ai.usage.output_tokens"] == 2 + assert span.attributes["gen_ai.usage.total_tokens"] == 7 + + +def test_concurrent_callable_credentials_use_each_failed_request(tracing): + if tuple(map(int, openai.__version__.split(".")[:2])) < (3, 0): + pytest.skip("callable OpenAI credentials are a latest-SDK feature") + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + keys = [KEY + "-first", KEY + "-second"] + issued = [] + + async def run(): + both_sent = asyncio.Event() + received = [] + + async def token(): + value = keys[len(issued)] + issued.append(value) + return value + + async def handler(request): + credential = request.headers["authorization"][7:] + received.append(credential) + if len(received) == 2: + both_sent.set() + await asyncio.wait_for(both_sent.wait(), 3) + return httpx.Response( + 401, + json={ + "error": { + "code": "invalid_api_key", + "message": "Invalid key: " + credential, + } + }, + ) + + async with client(handler, True, token) as c: + results = await asyncio.gather( + *[c.responses.create(model="parallel", input="question") for _ in keys], + return_exceptions=True, + ) + assert set(received) == set(keys) + assert all(isinstance(result, openai.AuthenticationError) for result in results) + for result in results: + assert result.request.headers["authorization"][7:] in str(result) + + asyncio.run(run()) + assert issued == keys + spans = exporter.get_finished_spans() + assert len(spans) == 2 + for span in spans: + assert span.status.status_code is StatusCode.ERROR + assert all(key not in span.to_json() for key in keys) + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("raw", [False, True]) +def test_callable_credentials_survive_stream_creation(tracing, async_, raw): + if tuple(map(int, openai.__version__.split(".")[:2])) < (3, 0): + pytest.skip("callable OpenAI credentials are a latest-SDK feature") + provider, exporter = tracing + OpenAIInstrumentor().instrument(tracer_provider=provider) + calls = [] + + def token(): + calls.append(1) + return KEY + + async def atoken(): + return token() + + response = body("failed") + response["error"] = {"code": "server_error", "message": "Rejected " + KEY} + + def handler(request): + assert request.headers["authorization"] == "Bearer " + KEY + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=sse("response.failed", response), + ) + + async def run(): + async with client(handler, True, atoken) as c: + api = c.responses.with_raw_response if raw else c.responses + result = await api.create(model="parallel", input="question", stream=True) + stream = result.parse() if raw else result + assert KEY in (await stream.__anext__()).response.error.message + await stream.close() + + if async_: + asyncio.run(run()) + else: + with client(handler, key=token) as c: + api = c.responses.with_raw_response if raw else c.responses + result = api.create(model="parallel", input="question", stream=True) + stream = result.parse() if raw else result + assert KEY in next(stream).response.error.message + stream.close() + assert len(calls) == 1 + (span,) = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + assert KEY not in span.to_json() diff --git a/python/frameworks/openai/tests/test_responses_safety.py b/python/frameworks/openai/tests/test_responses_safety.py new file mode 100644 index 000000000..fb2d773ff --- /dev/null +++ b/python/frameworks/openai/tests/test_responses_safety.py @@ -0,0 +1,477 @@ +"""Responses privacy and lifecycle contracts with real SDKs, no vendor traffic.""" +import asyncio +import inspect +import json + +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from traceai_openai import OpenAIInstrumentor + +KEY = "synthetic-parallel-vendor-key-DO-NOT-EXPORT" +QUESTION = "synthetic private question" + +@pytest.fixture +def tracing(): + exporter = InMemorySpanExporter() + provider = register(project_name="recall-offline", project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False) + try: + yield provider, exporter + finally: + provider.force_flush() + provider.shutdown() + + +@pytest.fixture +def openai_tracing(tracing): + provider, exporter = tracing + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider) + try: + yield provider, exporter + finally: + instrumentor.uninstrument() + + +def sdk_client(handler, async_=False): + http_client = (httpx.AsyncClient if async_ else httpx.Client)( + transport=httpx.MockTransport(handler), trust_env=False) + return (openai.AsyncOpenAI if async_ else openai.OpenAI)( + api_key=KEY, base_url="https://api.parallel.ai/v1", max_retries=0, http_client=http_client) + + +def response_body(status="completed"): + return { + "id": "resp_synthetic", "object": "response", "created_at": 1, "model": "parallel", + "status": status, "parallel_tool_calls": True, "tool_choice": "auto", "tools": [], + "output": [], "error": {"code": "server_error", "message": "synthetic upstream failed"} if status == "failed" else None, + "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, "output_tokens_details": {"reasoning_tokens": 0}}, + } + + +def sse_event(kind, response): + return ("event: " + kind + "\ndata: " + json.dumps({"type": kind, "sequence_number": 1, "response": response}) + "\n\n").encode() + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +def test_failed_response_must_not_export_success(openai_tracing, async_, stream): + _, exporter = openai_tracing + def handler(request): + if stream: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.failed", response_body("failed"))) + return httpx.Response(200, json=response_body("failed")) + async def run(): + async with sdk_client(handler, True) as client: + result = await client.responses.create(model="parallel", input=QUESTION, stream=stream) + if stream: + events = [event async for event in result] + assert events[0].type == "response.failed" + else: + assert result.status == "failed" + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + result = client.responses.create(model="parallel", input=QUESTION, stream=stream) + if stream: + events = list(result) # Exhaustion is part of the lifecycle contract. + assert events[0].type == "response.failed" + else: + assert result.status == "failed" + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["gen_ai.provider.name"] == "parallel" + assert spans[0].status.status_code is StatusCode.ERROR, spans[0].to_json() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_early_stream_close_must_finish_span(openai_tracing, async_): + provider, exporter = openai_tracing + def handler(request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.created", response_body("in_progress")) + + sse_event("response.completed", response_body())) + async def run(): + async with sdk_client(handler, True) as client: + stream = await client.responses.create(model="parallel", input=QUESTION, stream=True) + assert (await stream.__anext__()).type == "response.created" + await stream.close() + await stream.close() + assert stream.response.is_closed + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + stream = client.responses.create(model="parallel", input=QUESTION, stream=True) + assert next(stream).type == "response.created" + stream.close() + stream.close() + assert stream.response.is_closed + provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, f"closed SDK stream exported {len(spans)} spans" + assert spans[0].status.status_code is StatusCode.ERROR + + +@pytest.mark.parametrize("stage", ["request", "stream"]) +def test_async_cancellation_must_finish_span(openai_tracing, stage): + provider, exporter = openai_tracing + async def run(): + waiting = asyncio.Event() + class WaitingStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield sse_event("response.created", response_body("in_progress")) + waiting.set() + await asyncio.Event().wait() + async def handler(request): + if stage == "request": + waiting.set() + await asyncio.Event().wait() + return httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=WaitingStream()) + async with sdk_client(handler, True) as client: + stream = None + async def call(): + nonlocal stream + stream = await client.responses.create(model="parallel", input=QUESTION, stream=stage == "stream") + if stage == "stream": + async for _ in stream: + pass + task = asyncio.create_task(call()) + await asyncio.wait_for(waiting.wait(), timeout=3) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + if stream is not None: + await stream.close() + asyncio.run(run()) + provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, f"cancelled SDK {stage} exported {len(spans)} spans" + assert spans[0].status.status_code is StatusCode.ERROR + + +@pytest.mark.parametrize("async_", [False, True]) +def test_api_key_in_error_must_be_redacted(openai_tracing, async_): + _, exporter = openai_tracing + def handler(request): + assert request.headers["authorization"] == "Bearer " + KEY + return httpx.Response(401, json={"error": {"code": "invalid_api_key", "message": "Invalid key: " + KEY}}) + async def run(): + async with sdk_client(handler, True) as client: + await client.responses.create(model="parallel", input=QUESTION) + with pytest.raises(openai.AuthenticationError) as caught: + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + client.responses.create(model="parallel", input=QUESTION) + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code is StatusCode.ERROR + assert KEY in str(caught.value) # The caller still receives the original SDK exception. + assert KEY not in spans[0].to_json(), "vendor credential exported in error status/event" + + +@pytest.mark.parametrize("async_", [False, True]) +def test_hidden_instructions_not_reexported_in_output(tracing, async_): + provider, exporter = tracing + instruction = "synthetic confidential system instruction DO-NOT-EXPORT" + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True)) + def handler(request): + assert json.loads(request.content)["instructions"] == instruction + body = response_body() + body["instructions"] = instruction + body["output"] = [{"type": "message", "id": "msg_1", "role": "assistant", "status": "completed", + "content": [{"type": "output_text", "text": "safe answer", "annotations": []}]}] + return httpx.Response(200, json=body) + async def run(): + async with sdk_client(handler, True) as client: + response = await client.responses.create(model="parallel", input=QUESTION, instructions=instruction) + assert response.instructions == instruction + try: + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + response = client.responses.create(model="parallel", input=QUESTION, instructions=instruction) + assert response.instructions == instruction + span, = exporter.get_finished_spans() + assert span.attributes["input.value"] == "__REDACTED__" + assert "safe answer" in span.attributes["output.value"] + assert instruction not in span.to_json(), "FI_HIDE_INPUTS instruction leaked via output.value" + finally: + instrumentor.uninstrument() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_mcp_credentials_not_captured_as_tool_definitions(openai_tracing, async_): + _, exporter = openai_tracing + oauth = "synthetic-oauth-credential-DO-NOT-EXPORT" + apikey = "synthetic-mcp-api-key-DO-NOT-EXPORT" + tool = {"type": "mcp", "server_label": "synthetic", "server_url": "https://synthetic.invalid/mcp", + "require_approval": "never", "authorization": oauth, "headers": {"X-Api-Key": apikey}} + def handler(request): + parsed = json.loads(request.content) + assert parsed["tools"][0]["authorization"] == oauth + assert parsed["tools"][0]["headers"]["X-Api-Key"] == apikey + # Parallel explicitly scrubs tool credentials in its response. + body = response_body() + body["tools"] = [{**tool, "authorization": None, "headers": None}] + return httpx.Response(200, json=body) + async def run(): + async with sdk_client(handler, True) as client: + await client.responses.create(model="parallel", input=QUESTION, tools=[tool]) + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + client.responses.create(model="parallel", input=QUESTION, tools=[tool]) + span, = exporter.get_finished_spans() + assert oauth not in span.to_json(), "MCP OAuth token exported in gen_ai.tool.definitions" + assert apikey not in span.to_json(), "MCP API key exported in gen_ai.tool.definitions" + + +@pytest.mark.parametrize("async_", [False, True]) +def test_sdk_responses_stream_context_exit_finishes_span(openai_tracing, async_): + provider, exporter = openai_tracing + def handler(request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.created", response_body("in_progress")) + + sse_event("response.completed", response_body())) + async def run(): + async with sdk_client(handler, True) as client, client.responses.stream( + model="parallel", input=QUESTION + ) as stream: + async for event in stream: + assert event.type == "response.created" + break + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client, client.responses.stream( + model="parallel", input=QUESTION + ) as stream: + for event in stream: + assert event.type == "response.created" + break + provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, f"SDK context exit exported {len(spans)} spans" + assert spans[0].status.status_code is StatusCode.ERROR + + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("mode", ["complete", "failed", "incomplete", "error", "eof", "empty"]) +def test_terminal_status_and_close_are_consistent(openai_tracing, async_, mode): + _, exporter = openai_tracing + status = {"complete": "completed", "failed": "failed", "incomplete": "incomplete"}.get(mode) + body = response_body(status or "in_progress") + if mode == "failed": + body["error"]["message"] = "upstream rejected " + KEY + payload = sse_event("response.created", response_body("in_progress")) + if status: + payload += sse_event("response." + status, body) + elif mode == "error": + payload += ('data: ' + json.dumps({"type": "error", "code": "server_error", + "message": "upstream rejected " + KEY, "param": None, + "sequence_number": 2}) + '\n\n').encode() + + def handler(request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=payload) + + async def run(): + async with sdk_client(handler, True) as client: + stream = await client.responses.create(model="parallel", input=QUESTION, stream=True) + if mode != "empty": + async for event in stream: + if status and event.type == "response." + status: + break # Terminal event received, but no EOF read yet. + await stream.close() + await stream.close() + assert stream.response.is_closed + + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + stream = client.responses.create(model="parallel", input=QUESTION, stream=True) + if mode != "empty": + for event in stream: + if status and event.type == "response." + status: + break + stream.close() + stream.close() + assert stream.response.is_closed + span, = exporter.get_finished_spans() + assert span.status.status_code is (StatusCode.OK if mode == "complete" else StatusCode.ERROR) + assert KEY not in span.to_json() + if status: + assert span.attributes["gen_ai.usage.input_tokens"] == 5 + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("flag", ["hide_inputs", "hide_input_text", "hide_input_messages"]) +def test_echoed_request_fields_never_become_output(tracing, async_, stream, flag): + provider, exporter = tracing + instruction = "synthetic-private-system-text" + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(**{flag: True})) + body = response_body() + # Empty output exercises the accumulator's JSON fallback instead of output_text. + body.update(instructions=instruction, metadata={"prompt": instruction}, + input=instruction, text={"format": {"type": "json_schema", "name": instruction}}) + + def handler(request): + if stream: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.completed", body)) + return httpx.Response(200, json=body) + + async def run(): + async with sdk_client(handler, True) as client: + result = await client.responses.create(model="parallel", input=QUESTION, + instructions=instruction, stream=stream) + if stream: + events = [event async for event in result] + assert events[0].response.instructions == instruction + else: + assert result.instructions == instruction + + try: + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + result = client.responses.create(model="parallel", input=QUESTION, + instructions=instruction, stream=stream) + if stream: + events = list(result) + assert events[0].response.instructions == instruction + else: + assert result.instructions == instruction + span, = exporter.get_finished_spans() + assert instruction not in span.to_json() + output = json.loads(span.attributes["output.value"]) + assert output["status"] == "completed" + assert output["usage"]["input_tokens"] == 5 + finally: + instrumentor.uninstrument() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_raw_failed_response_retains_sdk_contract(openai_tracing, async_): + _, exporter = openai_tracing + + def handler(request): + return httpx.Response(200, json=response_body("failed")) + + async def run(): + async with sdk_client(handler, True) as client: + raw = await client.responses.with_raw_response.create(model="parallel", input=QUESTION) + assert raw.status_code == 200 + assert raw.parse().status == "failed" + + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + raw = client.responses.with_raw_response.create(model="parallel", input=QUESTION) + assert raw.status_code == 200 + assert raw.parse().status == "failed" + span, = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + + +@pytest.mark.parametrize("async_", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.skipif( + "hide_error_details" not in inspect.signature(TraceConfig).parameters, + reason="Error-detail masking requires a core with hide_error_details support", +) +def test_hidden_error_details_apply_to_failed_response_json(tracing, async_, stream): + provider, exporter = tracing + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(hide_error_details=True)) + message = "synthetic-private-error-details" + body = response_body("failed") + body["error"]["message"] = message + + def handler(request): + if stream: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, + content=sse_event("response.failed", body)) + return httpx.Response(200, json=body) + + async def run(): + async with sdk_client(handler, True) as client: + result = await client.responses.create(model="parallel", input=QUESTION, stream=stream) + if stream: + events = [event async for event in result] + assert events[0].response.error.message == message + else: + assert result.error.message == message + + try: + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + result = client.responses.create(model="parallel", input=QUESTION, stream=stream) + if stream: + events = list(result) + assert events[0].response.error.message == message + else: + assert result.error.message == message + span, = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + assert message not in span.to_json() + finally: + instrumentor.uninstrument() + + +@pytest.mark.parametrize("async_", [False, True]) +def test_mcp_wire_and_function_schema_are_not_mutated(openai_tracing, async_): + _, exporter = openai_tracing + token = 'synthetic-token-with-"quote"-and-\\slash' + mcp = {"type": "mcp", "server_label": "synthetic", "server_url": "https://synthetic.invalid/mcp", + "authorization": token, "headers": {"X-Custom-Credential": token}} + function = {"type": "function", "name": "lookup", "parameters": { + "type": "object", "properties": {"authorization": {"type": "string"}}}} + tools = [mcp, function] + original = json.loads(json.dumps(tools)) + + def handler(request): + assert json.loads(request.content)["tools"] == original + body = response_body() + body["tools"] = tools # Don't rely on a provider scrubbing its response. + return httpx.Response(200, json=body) + + async def run(): + async with sdk_client(handler, True) as client: + result = await client.responses.create(model="parallel", input=QUESTION, tools=tools) + assert result.tools[0].authorization == token + + if async_: + asyncio.run(run()) + else: + with sdk_client(handler) as client: + result = client.responses.create(model="parallel", input=QUESTION, tools=tools) + assert result.tools[0].authorization == token + assert tools == original + span, = exporter.get_finished_spans() + assert "synthetic-token-with" not in span.to_json() + definitions = json.loads(span.attributes["gen_ai.tool.definitions"]) + assert "authorization" not in definitions[0] + assert "headers" not in definitions[0] + assert definitions[1] == function diff --git a/python/frameworks/openai/tests/test_stream_context_manager.py b/python/frameworks/openai/tests/test_stream_context_manager.py index 684f57dd9..388024ec9 100644 --- a/python/frameworks/openai/tests/test_stream_context_manager.py +++ b/python/frameworks/openai/tests/test_stream_context_manager.py @@ -31,7 +31,7 @@ def test_sync_stream_context_finishes_once(make_tracer, mode): span, attrs = only_span(exporter) assert attrs["gen_ai.request.model"] == REQUEST_MODEL assert attrs["gen_ai.provider.name"] == "baseten" - assert span.status.status_code == (StatusCode.ERROR if mode == "error" else StatusCode.OK) + assert span.status.status_code == (StatusCode.OK if mode == "full" else StatusCode.ERROR) if mode == "full": assert attrs["output.value"] == "Tokyo" if mode == "early": @@ -66,7 +66,7 @@ async def run(): span, attrs = only_span(exporter) assert attrs["gen_ai.request.model"] == REQUEST_MODEL assert attrs["gen_ai.provider.name"] == "baseten" - assert span.status.status_code == (StatusCode.ERROR if mode == "error" else StatusCode.OK) + assert span.status.status_code == (StatusCode.OK if mode == "full" else StatusCode.ERROR) if mode == "full": assert attrs["output.value"] == "Tokyo" if mode == "early": diff --git a/python/frameworks/openai/traceai_openai/__init__.py b/python/frameworks/openai/traceai_openai/__init__.py index 891928712..e7633a2a3 100644 --- a/python/frameworks/openai/traceai_openai/__init__.py +++ b/python/frameworks/openai/traceai_openai/__init__.py @@ -16,6 +16,7 @@ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor # type: ignore from traceai_openai._client_provider import set_client_provider as set_client_provider from traceai_openai._request import _AsyncRequest, _Request +from traceai_openai._request_context import capture_response_credentials from traceai_openai.package import _instruments from traceai_openai.version import __version__ from wrapt import wrap_function_wrapper @@ -34,6 +35,8 @@ class OpenAIInstrumentor(BaseInstrumentor): # type: ignore __slots__ = ( "_original_request", "_original_async_request", + "_original_process_response", + "_original_async_process_response", ) def instrumentation_dependencies(self) -> Collection[str]: @@ -63,6 +66,14 @@ def _instrument(self, **kwargs: Any) -> None: name="AsyncOpenAI.request", wrapper=_AsyncRequest(tracer=tracer, openai=openai, config=config), ) + self._original_process_response = openai.OpenAI._process_response + self._original_async_process_response = openai.AsyncOpenAI._process_response + for client in ("OpenAI", "AsyncOpenAI"): + wrap_function_wrapper( + module=_MODULE, + name=f"{client}._process_response", + wrapper=capture_response_credentials, + ) if Protect is not None: self._original_protect = Protect.protect wrap_function_wrapper( @@ -77,3 +88,5 @@ def _uninstrument(self, **kwargs: Any) -> None: openai = import_module(_MODULE) openai.OpenAI.request = self._original_request openai.AsyncOpenAI.request = self._original_async_request + openai.OpenAI._process_response = self._original_process_response + openai.AsyncOpenAI._process_response = self._original_async_process_response diff --git a/python/frameworks/openai/traceai_openai/_attributes/_responses_api.py b/python/frameworks/openai/traceai_openai/_attributes/_responses_api.py index f92d6ec07..e4ff64ec5 100644 --- a/python/frameworks/openai/traceai_openai/_attributes/_responses_api.py +++ b/python/frameworks/openai/traceai_openai/_attributes/_responses_api.py @@ -6,7 +6,8 @@ from opentelemetry.util.types import AttributeValue from typing_extensions import assert_never -from fi_instrumentation import safe_json_dumps +from fi_instrumentation import safe_json_dumps +from traceai_openai._privacy import safe_tool from fi_instrumentation.fi_types import ( ImageAttributes, MessageAttributes, @@ -237,7 +238,7 @@ def _get_attributes_from_response_tool( obj: responses.tool.Tool, prefix: str = "", ) -> Iterator[Tuple[str, AttributeValue]]: - yield (f"{prefix}{ToolAttributes.TOOL_JSON_SCHEMA}", obj.model_dump_json()) + yield (f"{prefix}{ToolAttributes.TOOL_JSON_SCHEMA}", safe_json_dumps(safe_tool(obj.model_dump(mode="json")))) @classmethod @stop_on_exception diff --git a/python/frameworks/openai/traceai_openai/_privacy.py b/python/frameworks/openai/traceai_openai/_privacy.py new file mode 100644 index 000000000..3ff5d3ade --- /dev/null +++ b/python/frameworks/openai/traceai_openai/_privacy.py @@ -0,0 +1,116 @@ +"""Telemetry-only copies; never mutate SDK requests, responses or exceptions.""" + +from __future__ import annotations + +import json +from copy import deepcopy +from collections.abc import Iterable, Mapping +from typing import Any + +from fi_instrumentation import TraceConfig + +REDACTED = "__REDACTED__" + + +def redact(value: Any, secrets: Iterable[str]) -> Any: + if isinstance(value, str): + for secret in secrets: + value = value.replace(secret, REDACTED) + return value + if isinstance(value, (list, tuple)): + return type(value)(redact(item, secrets) for item in value) + return value + + +def request_credentials(instance: Any, options: Any, parameters: Mapping) -> tuple[str, ...]: + """Read static credentials only; never invoke a caller's token provider.""" + values = [getattr(instance, "api_key", None)] + for headers in (getattr(instance, "_custom_headers", None), getattr(options, "headers", None)): + if isinstance(headers, Mapping): + for name, value in headers.items(): + if str(name).lower() in {"authorization", "api-key", "x-api-key"}: + values.append(value) + if isinstance(value, str) and value.lower().startswith("bearer "): + values.append(value[7:]) + # extra_body is a shallow SDK override. Inspect its effective tools only + # for private redaction; never add extra_body to exported request metadata. + extra_json = getattr(options, "extra_json", None) + tools = ( + extra_json["tools"] + if isinstance(extra_json, Mapping) and "tools" in extra_json + else parameters.get("tools") + ) + if isinstance(tools, (list, tuple)): + for tool in tools: + if isinstance(tool, Mapping) and tool.get("type") == "mcp": + values.append(tool.get("authorization")) + headers = tool.get("headers") + if isinstance(headers, Mapping): + values.extend(headers.values()) + secrets = set() + for value in values: + if isinstance(value, str) and value: + secrets.update((value, json.dumps(value)[1:-1], repr(value)[1:-1])) + return tuple(sorted(secrets, key=len, reverse=True)) + + +def response_credentials(response: Any) -> tuple[str, ...]: + """Use the request actually sent, not a shared client's last refreshed token. + + Raw responses and streams retain their HTTP response. Reading its request + headers never calls the user's token provider or modifies authentication. + """ + http_response = getattr(response, "http_response", None) + if http_response is None: + http_response = getattr(response, "response", None) + request = getattr(http_response, "request", None) + return request_credentials(None, request, {}) + + +def safe_tool(tool: Any) -> Any: + if isinstance(tool, Mapping) and tool.get("type") == "mcp": + return {key: value for key, value in tool.items() if key not in {"authorization", "headers"}} + return tool + + +def safe_request(parameters: Mapping) -> dict: + """MCP headers are arbitrary authentication material, not tool schemas.""" + result = dict(parameters) + tools = result.get("tools") + if isinstance(tools, (list, tuple)): + result["tools"] = [safe_tool(tool) for tool in tools] + return result + + +def response_output(response: Any, config: TraceConfig | None = None) -> dict: + """Only response-owned data belongs in output, not echoed instructions/tools. + + An allowlist also keeps new SDK request-echo fields from bypassing input + privacy flags. The SDK response returned to the application is unchanged. + """ + fields = {"id", "object", "created_at", "model", "status", "output", "usage", + "error", "incomplete_details"} + if config is not None and getattr(config, "hide_error_details", False): + fields.difference_update({"error", "incomplete_details"}) + if isinstance(response, Mapping): + output = {key: deepcopy(value) for key, value in response.items() + if key in fields and value is not None} + else: + output = response.model_dump(include=fields, exclude_none=True, mode="json") + if config is not None and getattr(config, "hide_error_details", False): + # These provider-owned error fields are not ordinary tool results or + # assistant text. Work on the dump only, leaving SDK objects untouched. + for item in output.get("output", []): + if not isinstance(item, dict) or item.get("type") not in { + "mcp_call", "mcp_list_tools", "image_generation_call", + }: + continue + error = item.pop("error", None) + if isinstance(error, dict) and error.get("type") in { + "http_error", "mcp_protocol_error", "mcp_tool_execution_error", + }: + classification = {"type": error["type"]} + if type(error.get("code")) is int: + classification["code"] = error["code"] + item["error"] = classification + return output diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index c6e8cf426..04494382d 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -1,6 +1,8 @@ +import json import logging from abc import ABC from contextlib import contextmanager +from inspect import iscoroutinefunction from itertools import chain from types import ModuleType from typing import Any, Awaitable, Callable, Iterable, Iterator, Mapping, Optional, Tuple @@ -10,6 +12,8 @@ FiSpanKindValues, SpanAttributes, ) +from openai import Omit +from openai._legacy_response import LegacyAPIResponse from openai.types.create_embedding_response import CreateEmbeddingResponse from opentelemetry import context as context_api from opentelemetry import trace as trace_api @@ -18,15 +22,25 @@ from opentelemetry.util.types import AttributeValue from traceai_openai._client_provider import _get_client_provider from traceai_openai._providers import provider_from_host +from traceai_openai._privacy import redact, request_credentials, response_credentials, response_output, safe_request +from traceai_openai._request_context import capture_request_credentials from traceai_openai._request_attributes_extractor import _RequestAttributesExtractor from traceai_openai._response_accumulator import ( _ChatCompletionAccumulator, _CompletionAccumulator, _ResponsesAccumulator, + _response_status, ) from traceai_openai._response_attributes_extractor import _ResponseAttributesExtractor from traceai_openai._span_io_handler import add_io_to_span_attributes -from traceai_openai._stream import _AsyncStream, _ResponseAccumulator, _Stream +from traceai_openai._stream import ( + _AsyncRawResponse, + _AsyncStream, + _LegacyRawResponse, + _RawResponse, + _ResponseAccumulator, + _Stream, +) from traceai_openai._utils import ( _as_input_attributes, _finish_tracing, @@ -57,12 +71,13 @@ def _start_as_current_span( attributes: Iterable[Tuple[str, AttributeValue]], context_attributes: Iterable[Tuple[str, AttributeValue]], extra_attributes: Iterable[Tuple[str, AttributeValue]], + secrets: Tuple[str, ...] = (), ) -> Iterator[_WithSpan]: # Keep primary attributes so finalization can restore them after optional # details: OTel's bounded attribute map evicts the oldest entries. initial_attributes = {} try: - initial_attributes = dict(attributes) + initial_attributes = {key: redact(value, secrets) for key, value in attributes} span = self._tracer.start_span(name=span_name, attributes=initial_attributes) except Exception: logger.exception("Failed to start span") @@ -75,6 +90,7 @@ def _start_as_current_span( ) as span: yield _WithSpan( span=span, + secrets=secrets, initial_attributes=initial_attributes, context_attributes=dict(context_attributes), extra_attributes=dict(extra_attributes), @@ -114,6 +130,7 @@ def responses_accumulator(request_parameters: _RequestParameters) -> Any: return _ResponsesAccumulator( request_parameters=request_parameters, chat_completion_type=openai.types.responses.response.Response, + config=self._config, response_attributes_extractor=self._response_attributes_extractor, ) @@ -201,54 +218,94 @@ def _finalize_response( with_span: _WithSpan, cast_to: type, request_parameters: Mapping[str, Any], + request_options: Any = None, ) -> Any: """ Monkey-patch the response object to trace the stream, or finish tracing if the response is not a stream. """ - if hasattr(response, "parse") and callable(response.parse): - try: - response.parse() - except Exception: - logger.exception(f"Failed to parse response of type {type(response)}") - - if self._is_streaming(response): - # Determine if the response is asynchronous - if self._is_async_stream(response): - stream_wrapper = _AsyncStream - else: - stream_wrapper = _Stream - - # Get the response accumulator if available - response_accumulator_factory = self._response_accumulator_factories.get( - cast_to - ) - response_accumulator = ( - response_accumulator_factory(request_parameters) - if response_accumulator_factory - else None + response_accumulator_factory = self._response_accumulator_factories.get(cast_to) + accumulator_parameters = request_parameters + if cast_to is self._openai.types.chat.ChatCompletion and request_options is not None: + # Completion validation needs the on-wire choice count. Keep + # extra_body out of the existing exported request metadata. + _, accumulator_parameters = _parse_request_args( + (cast_to, request_options), include_extra_body=True, ) + response_accumulator = ( + response_accumulator_factory(accumulator_parameters) + if response_accumulator_factory + else None + ) - # Wrap the response with the appropriate stream wrapper - wrapped_stream = stream_wrapper( - stream=response, - with_span=with_span, - response_accumulator=response_accumulator, - ) - - return wrapped_stream - - else: + def on_parse(parsed: Any) -> Any: + if self._is_streaming(parsed): + stream_wrapper = _AsyncStream if self._is_async_stream(parsed) else _Stream + return stream_wrapper( + stream=parsed, + with_span=with_span, + response_accumulator=response_accumulator, + ) + if isinstance(parsed, CreateEmbeddingResponse): + with_span.set_attribute( + SpanAttributes.EMBEDDING_EMBEDDINGS, self._raw_embeddings_json(parsed) + ) + telemetry_response = parsed + status = _response_status(parsed) + if cast_to is self._openai.types.responses.response.Response and not isinstance( + parsed, cast_to + ): + # parse(to=dict/str) returns the original SDK value to the caller, + # but telemetry must pass through the same Response allowlist as + # the default parse. Never serialize the arbitrary raw value. + try: + body = json.loads(parsed) if isinstance(parsed, str) else parsed + if not isinstance(body, Mapping): + raise ValueError("Responses body is not an object") + status = _response_status(body) + telemetry_response = response_output(body, self._config) + except (TypeError, ValueError): + telemetry_response = None + status = trace_api.Status( + trace_api.StatusCode.ERROR, "Response body could not be inspected" + ) + if telemetry_response is not None: + add_io_to_span_attributes( + with_span, None, telemetry_response, config=self._config + ) _finish_tracing( - status=trace_api.Status(status_code=trace_api.StatusCode.OK), + status=status, with_span=with_span, has_attributes=_ResponseAttributes( request_parameters=request_parameters, - response=response, + response=telemetry_response, response_attributes_extractor=self._response_attributes_extractor, ), ) - return response + return parsed + + if hasattr(response, "parse") and callable(response.parse): + # Modern envelopes are deliberately lazy, even for non-SSE responses. + # Do not call async parse from sync I/O/attribute extraction paths. + raw_wrapper = ( + _LegacyRawResponse if isinstance(response, LegacyAPIResponse) + else _AsyncRawResponse if iscoroutinefunction(response.parse) + else _RawResponse + ) + raw = raw_wrapper(response, with_span, response_accumulator, on_parse) + if isinstance(response, LegacyAPIResponse): + # Preserve eager legacy non-stream telemetry. Stream parse caches + # remain SDK-owned, and every later explicit parse type is traced. + try: + parsed = response.parse() + except Exception: + # Optional eager telemetry must not change SDK behavior: callers + # still receive the envelope and may inspect malformed bodies. + logger.debug("Could not eagerly parse raw response for telemetry", exc_info=True) + else: + raw._trace_parsed(parsed) + return raw + return on_parse(response) def _raw_embeddings_json(self, response: CreateEmbeddingResponse) -> str: """The raw embeddings response; vectors are replaced by their dimension @@ -285,6 +342,8 @@ def __call__( return wrapped(*args, **kwargs) try: cast_to, request_parameters = _parse_request_args(args) + secrets = request_credentials(instance, args[1], request_parameters) + request_parameters = safe_request(request_parameters) # Extract input data more robustly input_data = ( request_parameters.get("messages") @@ -299,6 +358,7 @@ def __call__( with self._start_as_current_span( span_name=span_name, + secrets=secrets, attributes=chain( self._get_attributes_from_instance(instance), self._get_attributes_from_request( @@ -318,29 +378,20 @@ def __call__( ) add_io_to_span_attributes(with_span, input_data, None) try: - response = wrapped(*args, **kwargs) - - if isinstance(response, CreateEmbeddingResponse): - with_span.set_attribute( - SpanAttributes.EMBEDDING_EMBEDDINGS, - self._raw_embeddings_json(response), - ) + with capture_request_credentials(instance, with_span): + response = wrapped(*args, **kwargs) + with_span.add_secrets(request_credentials(instance, None, {})) + with_span.add_secrets(response_credentials(response)) - # Add output data to span attributes after getting response - if not self._is_streaming(response): - add_io_to_span_attributes( - with_span, - None, - response, - is_streaming=self._is_streaming(response), - ) return self._finalize_response( response=response, with_span=with_span, cast_to=cast_to, request_parameters=request_parameters, + request_options=args[1], ) - except Exception as e: + except BaseException as e: + with_span.add_secrets(request_credentials(instance, getattr(e, "request", None), {})) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, description=f"{type(e).__name__}: {e}", @@ -362,6 +413,8 @@ async def __call__( return await wrapped(*args, **kwargs) try: cast_to, request_parameters = _parse_request_args(args) + secrets = request_credentials(instance, args[1], request_parameters) + request_parameters = safe_request(request_parameters) # Extract input data more robustly input_data = ( request_parameters.get("messages") @@ -376,6 +429,7 @@ async def __call__( return await wrapped(*args, **kwargs) with self._start_as_current_span( span_name=span_name, + secrets=secrets, attributes=chain( self._get_attributes_from_instance(instance), self._get_attributes_from_request( @@ -395,18 +449,14 @@ async def __call__( ) add_io_to_span_attributes(with_span, input_data, None) try: - response = await wrapped(*args, **kwargs) - - if isinstance(response, CreateEmbeddingResponse): - with_span.set_attribute( - SpanAttributes.EMBEDDING_EMBEDDINGS, - self._raw_embeddings_json(response), - ) - - add_io_to_span_attributes( - with_span, None, response, is_streaming=self._is_streaming(response) - ) - except Exception as exception: + with capture_request_credentials(instance, with_span): + response = await wrapped(*args, **kwargs) + with_span.add_secrets(request_credentials(instance, None, {})) + with_span.add_secrets(response_credentials(response)) + except BaseException as exception: + with_span.add_secrets(request_credentials( + instance, getattr(exception, "request", None), {} + )) with_span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -422,6 +472,7 @@ async def __call__( with_span=with_span, cast_to=cast_to, request_parameters=request_parameters, + request_options=args[1], ) except Exception: logger.exception( @@ -431,18 +482,29 @@ async def __call__( return response -def _parse_request_args(args: Tuple[type, Any]) -> Tuple[type, Mapping[str, Any]]: +def _parse_request_args( + args: Tuple[type, Any], *, include_extra_body: bool = False, +) -> Tuple[type, Mapping[str, Any]]: # We don't use `signature(request).bind()` because `request` could have been monkey-patched # (incorrectly) by others and the signature at runtime may not match the original. # The targeted signature of `request` is here: # https://github.com/openai/openai-python/blob/f1c7d714914e3321ca2e72839fe2d132a8646e7f/src/openai/_base_client.py#L846-L847 # noqa: E501 cast_to: type = args[0] + json_data = getattr(args[1], "json_data", None) + extra_json = getattr(args[1], "extra_json", None) + # Mirror the SDK's shallow extra_body merge without mutating its options. + # Overrides (including Omit removals) affect the number of terminal choices + # required for a successful stream, not the existing metadata contract. + if include_extra_body and extra_json is not None: + if json_data is None: + json_data = extra_json + elif isinstance(json_data, Mapping) and isinstance(extra_json, Mapping): + json_data = { + key: value for key, value in {**json_data, **extra_json}.items() + if not isinstance(value, Omit) + } request_parameters: Mapping[str, Any] = ( - json_data - # See https://github.com/openai/openai-python/blob/f1c7d714914e3321ca2e72839fe2d132a8646e7f/src/openai/_models.py#L427 # noqa: E501 - if hasattr(args[1], "json_data") - and isinstance(json_data := args[1].json_data, Mapping) - else {} + json_data if isinstance(json_data, Mapping) else {} ) # FIXME: Because request parameters is just a Mapping, it can contain any value as long as it # serializes correctly in an HTTP request body. For example, Enum values may be present if a @@ -467,12 +529,6 @@ def __init__( request_parameters: Mapping[str, Any], response_attributes_extractor: _ResponseAttributesExtractor, ) -> None: - if hasattr(response, "parse") and callable(response.parse): - # E.g. see https://github.com/openai/openai-python/blob/f1c7d714914e3321ca2e72839fe2d132a8646e7f/src/openai/_base_client.py#L518 # noqa: E501 - try: - response = response.parse() - except Exception: - logger.exception(f"Failed to parse response of type {type(response)}") self._request_parameters = request_parameters self._response = response self._response_attributes_extractor = response_attributes_extractor diff --git a/python/frameworks/openai/traceai_openai/_request_context.py b/python/frameworks/openai/traceai_openai/_request_context.py new file mode 100644 index 000000000..c81131cf7 --- /dev/null +++ b/python/frameworks/openai/traceai_openai/_request_context.py @@ -0,0 +1,41 @@ +"""Capture request-local credentials before the SDK discards its HTTP envelope.""" + +import logging +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, Iterator, Optional, Tuple + +from traceai_openai._privacy import request_credentials +from traceai_openai._with_span import _WithSpan + +logger = logging.getLogger(__name__) +_active_request: ContextVar[Optional[Tuple[Any, _WithSpan]]] = ContextVar( + "traceai_openai_request", default=None +) + + +@contextmanager +def capture_request_credentials(client: Any, span: _WithSpan) -> Iterator[None]: + token = _active_request.set((client, span)) + try: + yield + finally: + _active_request.reset(token) + + +def capture_response_credentials( + wrapped: Any, instance: Any, args: Any, kwargs: Any +) -> Any: + # Both SDK bounds pass response by keyword to _process_response, before + # parsing removes the request. ContextVars isolate concurrent calls without + # modifying shared clients/transports or invoking a token provider again. + active = _active_request.get() + if active is not None and active[0] is instance: + try: + response = kwargs.get("response") + active[1].add_secrets( + request_credentials(None, getattr(response, "request", None), {}) + ) + except Exception: + logger.warning("Failed to capture request credentials") + return wrapped(*args, **kwargs) diff --git a/python/frameworks/openai/traceai_openai/_response_accumulator.py b/python/frameworks/openai/traceai_openai/_response_accumulator.py index 17417d213..2dfbb7f8b 100644 --- a/python/frameworks/openai/traceai_openai/_response_accumulator.py +++ b/python/frameworks/openai/traceai_openai/_response_accumulator.py @@ -17,16 +17,19 @@ Type, ) -from fi_instrumentation import safe_json_dumps +from fi_instrumentation import TraceConfig, safe_json_dumps from fi_instrumentation.fi_types import FiMimeTypeValues, SpanAttributes +from opentelemetry import trace as trace_api from opentelemetry.util.types import AttributeValue +from opentelemetry.trace import Status, StatusCode +from traceai_openai._privacy import response_output from traceai_openai._utils import _as_output_attributes, _to_dict, _ValueAndType, _io_value_and_type if TYPE_CHECKING: from openai.types import Completion from openai.types.chat import ChatCompletion from openai.types.responses.response import Response - from openai.types.responses.response_completed_event import ResponseCompletedEvent + from openai.types.responses import ResponseStreamEvent __all__ = ( "_CompletionAccumulator", @@ -43,8 +46,21 @@ def get_attributes_from_response( ) -> Iterator[Tuple[str, AttributeValue]]: ... +def _response_status(response: Any) -> Status: + status = ( + response.get("status") + if isinstance(response, Mapping) + else getattr(response, "status", None) + ) + if isinstance(status, str) and status in {"failed", "cancelled", "incomplete"}: + # The vendor's error message can contain credentials or request content. + return Status(StatusCode.ERROR, "Response did not complete") + return Status(StatusCode.OK) + + class _ResponsesAccumulator: __slots__ = ( + "_config", "_is_null", "_values", "_cached_result", @@ -58,16 +74,20 @@ def __init__( request_parameters: Mapping[str, Any], chat_completion_type: Type["Response"], response_attributes_extractor: Optional[_CanGetAttributesFromResponse] = None, + config: Optional[TraceConfig] = None, ) -> None: + self._config = config self._chat_completion_type = chat_completion_type self._request_parameters = request_parameters self._response_attributes_extractor = response_attributes_extractor self._is_null = True self._cached_result: Optional[Dict[str, Any]] = None - self._values: Optional["ResponseCompletedEvent"] = None + self._values: Optional["ResponseStreamEvent"] = None def process_chunk(self, chunk: Any) -> None: - if type(chunk).__name__ == "ResponseCompletedEvent": + if getattr(chunk, "type", None) in { + "response.completed", "response.failed", "response.incomplete", "error" + }: self._is_null = False self._cached_result = None self._values = chunk @@ -76,7 +96,14 @@ def _result(self) -> Any: if self._is_null: return None if self._values: - return self._values.response + return getattr(self._values, "response", None) + + def get_status(self) -> Status: + if self._values is None: + return Status(StatusCode.ERROR, "Response stream ended before a terminal event") + if self._values.type != "response.completed": + return Status(StatusCode.ERROR, "Response stream did not complete") + return _response_status(self._result()) def get_attributes(self) -> Iterator[Tuple[str, AttributeValue]]: if not (result := self._result()): @@ -97,9 +124,9 @@ def get_attributes(self) -> Iterator[Tuple[str, AttributeValue]]: if output_content: yield SpanAttributes.OUTPUT_VALUE, output_content else: - # Fallback to full response if we can't extract content + # Keep response-owned data, never echoed request fields. yield from _as_output_attributes( - _io_value_and_type(result), + _io_value_and_type(response_output(result, self._config)), ) def get_extra_attributes(self) -> Iterator[Tuple[str, AttributeValue]]: @@ -122,6 +149,7 @@ class _ChatCompletionAccumulator: "_chat_completion_type", "_content", "_raw_data", + "_choice_completion", ) def __init__( @@ -136,6 +164,7 @@ def __init__( self._is_null = True self._content: List[str] = [] self._raw_data: List[Dict[str, Any]] = [] + self._choice_completion: Dict[int, bool] = {} self._cached_result: Optional[Dict[str, Any]] = None self._values = _ValuesAccumulator( choices=_IndexedAccumulator( @@ -162,6 +191,13 @@ def __init__( def process_chunk(self, chunk: Any) -> None: raw_chunk = _to_dict(chunk) self._raw_data.append(raw_chunk) + for choice in raw_chunk.get("choices", []): + index = choice.get("index") + if isinstance(index, int): + self._choice_completion[index] = ( + self._choice_completion.get(index, False) + or bool(choice.get("finish_reason")) + ) self._is_null = False self._cached_result = None @@ -185,6 +221,21 @@ def process_chunk(self, chunk: Any) -> None: if content is not None: self._content.append(content) + def get_status(self) -> trace_api.Status: + # The SDK consumes [DONE] internally and treats transport EOF as normal + # exhaustion. Only per-choice finish reasons prove model completion. + # Usage-only tail chunks must not clear an already completed choice. + expected = self._request_parameters.get("n", 1) + if not isinstance(expected, int) or expected < 1: + expected = 1 + if (len(self._choice_completion) >= expected + and all(self._choice_completion.values())): + return trace_api.Status(trace_api.StatusCode.OK) + return trace_api.Status( + trace_api.StatusCode.ERROR, + "Chat stream ended before all choices completed", + ) + def _result(self) -> Optional[Dict[str, Any]]: if self._is_null: return None diff --git a/python/frameworks/openai/traceai_openai/_span_io_handler.py b/python/frameworks/openai/traceai_openai/_span_io_handler.py index 028d2faa5..772db2da3 100644 --- a/python/frameworks/openai/traceai_openai/_span_io_handler.py +++ b/python/frameworks/openai/traceai_openai/_span_io_handler.py @@ -2,8 +2,10 @@ import logging from typing import Any, Optional +from fi_instrumentation import TraceConfig from fi_instrumentation.fi_types import FiMimeTypeValues, SpanAttributes from traceai_openai._with_span import _WithSpan +from traceai_openai._privacy import response_output logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -88,7 +90,7 @@ def _set_mime_type(span: _WithSpan, key: str, value: Any) -> None: span.set_attribute(key, mime) -def _json_output_for(output_data: Any) -> Optional[str]: +def _json_output_for(output_data: Any, config: Optional[TraceConfig] = None) -> Optional[str]: """JSON for responses whose output is not a chat message: the Responses API object as JSON, and a compact summary for embeddings (never the vectors).""" try: @@ -98,7 +100,7 @@ def _json_output_for(output_data: Any) -> Optional[str]: return None if isinstance(output_data, Response): # Lead with the answer text (the SDK's `output_text`) so previews show it. - body = json.loads(output_data.model_dump_json(exclude_none=True)) + body = response_output(output_data, config) output_text = getattr(output_data, "output_text", None) if output_text: body = {"output_text": output_text, **body} @@ -116,7 +118,8 @@ def _json_output_for(output_data: Any) -> Optional[str]: def add_io_to_span_attributes( - span: _WithSpan, input_data: Any, output_data: Any, is_streaming: bool = False + span: _WithSpan, input_data: Any, output_data: Any, is_streaming: bool = False, + config: Optional[TraceConfig] = None ) -> None: """Add input/output data to span attributes. @@ -135,7 +138,7 @@ def add_io_to_span_attributes( if output_data is not None and not is_streaming: if hasattr(output_data, "parse") and not hasattr(output_data, "model_dump"): output_data = output_data.parse() - if (json_output := _json_output_for(output_data)) is not None: + if (json_output := _json_output_for(output_data, config)) is not None: span.set_attribute(SpanAttributes.OUTPUT_VALUE, json_output) span.set_attribute( SpanAttributes.OUTPUT_MIME_TYPE, FiMimeTypeValues.JSON.value diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index a3a60042e..0f50ea1f6 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -1,9 +1,11 @@ +import json import logging -from typing import Any, AsyncIterator, Generator, Iterator, Optional, Protocol, Tuple +from typing import Any, AsyncIterator, Callable, Generator, Iterator, Optional, Protocol, Tuple from openai.types.chat import ChatCompletionChunk from opentelemetry import trace as trace_api from opentelemetry.util.types import AttributeValue +from traceai_openai._response_accumulator import _response_status from traceai_openai._utils import _finish_tracing from traceai_openai._with_span import _WithSpan from wrapt import ObjectProxy @@ -40,6 +42,274 @@ def _finish_on_exit(stream: Any, error: Optional[BaseException]) -> None: stream._finish_tracing(status=status) +class _ResponseCloseProxy(ObjectProxy): + """Responses.stream() closes the HTTP response, not the raw SDK stream. + + Expose a per-call proxy instead of patching HTTPX classes or the client's + transport. Internal SDK reads still use the original response unchanged. + """ + + def __init__(self, response: Any, stream: Any) -> None: + super().__init__(response) + self._self_stream = stream + + def close(self) -> None: + error = None + try: + self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self._self_stream, error) + + async def aclose(self) -> None: + error = None + try: + await self.__wrapped__.aclose() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self._self_stream, error) + + +class _RawResponse(ObjectProxy): + """Trace every parse-cache entry without eagerly reading a lazy envelope.""" + + def __init__( + self, + response: Any, + with_span: _WithSpan, + response_accumulator: Optional[_ResponseAccumulator], + on_parse: Callable[[Any], Any], + ) -> None: + super().__init__(response) + self._self_with_span = with_span + self._self_response_accumulator = response_accumulator + self._self_on_parse = on_parse + self._self_parsed = {} + self._self_body_status: Optional[trace_api.Status] = None + + def _record_body_consumption(self, body_bytes: Optional[bytes] = None) -> None: + # read/text/json consume a non-SSE envelope without calling parse(). + # Iterators do not cache HTTPX content, so inspect their collected bytes. + # Never invoke the SDK post_parser or export the raw body. + if ( + self._self_is_finished + or self._self_body_status is not None + or getattr(self.__wrapped__, "_is_sse_stream", False) + ): + return + try: + body = ( + json.loads(body_bytes) if body_bytes is not None + else self.__wrapped__.http_response.json() + ) + except Exception: + body = None # Successful raw reads need not contain JSON. + self._self_body_status = _response_status(body) + + def _consume_body(self, method: str, *args: Any, **kwargs: Any) -> Any: + try: + result = getattr(self.__wrapped__, method)(*args, **kwargs) + except BaseException as exc: + _finish_on_exit(self, exc) + raise + self._record_body_consumption() + return result + + def read(self, *args: Any, **kwargs: Any) -> Any: + return self._consume_body("read", *args, **kwargs) + + def text(self, *args: Any, **kwargs: Any) -> Any: + return self._consume_body("text", *args, **kwargs) + + def json(self, *args: Any, **kwargs: Any) -> Any: + return self._consume_body("json", *args, **kwargs) + + def _iterate_body(self, method: str, *args: Any, **kwargs: Any) -> Iterator[Any]: + is_sse = getattr(self.__wrapped__, "_is_sse_stream", False) + body = bytearray() + too_large = False + try: + for chunk in getattr(self.__wrapped__, method)(*args, **kwargs): + if not is_sse and not too_large: + encoded = chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") + if method == "iter_lines": + encoded += b"\n" + if len(body) + len(encoded) <= 8 * 1024 * 1024: + body.extend(encoded) + else: + too_large = True + body.clear() + yield chunk + except BaseException as exc: + _finish_on_exit(self, exc) + raise + else: + # Partial iteration and SSE must not report a completed response. + if too_large: + self._self_body_status = trace_api.Status( + trace_api.StatusCode.ERROR, "Response body too large to inspect" + ) + elif not is_sse: + self._record_body_consumption(bytes(body)) + + def iter_bytes(self, *args: Any, **kwargs: Any) -> Iterator[bytes]: + return self._iterate_body("iter_bytes", *args, **kwargs) + + def iter_text(self, *args: Any, **kwargs: Any) -> Iterator[str]: + return self._iterate_body("iter_text", *args, **kwargs) + + def iter_lines(self, *args: Any, **kwargs: Any) -> Iterator[str]: + return self._iterate_body("iter_lines", *args, **kwargs) + + def stream_to_file(self, *args: Any, **kwargs: Any) -> None: + self._consume_body("stream_to_file", *args, **kwargs) + + @property + def _self_is_finished(self) -> bool: + return self._self_with_span.is_finished + + def _trace_parsed(self, parsed: Any) -> Any: + # The SDK owns cache keys (including generic stream types). Cache by + # identity and retain the original through the proxy, so repeated parse + # calls return the same wrapper without collapsing distinct SDK entries. + key = id(parsed) + if key not in self._self_parsed: + self._self_parsed[key] = self._self_on_parse(parsed) + return self._self_parsed[key] + + def parse(self, *args: Any, **kwargs: Any) -> Any: + try: + parsed = self.__wrapped__.parse(*args, **kwargs) + except BaseException as exc: + _finish_on_exit(self, exc) + raise + return self._trace_parsed(parsed) + + @property + def http_response(self) -> Any: + return _ResponseCloseProxy(self.__wrapped__.http_response, self) + + def close(self) -> None: + error = None + try: + self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self, error) + + def _finish_tracing(self, status: Optional[trace_api.Status] = None) -> None: + if self._self_is_finished: + return + accumulator = self._self_response_accumulator + if status is None or status.status_code is not trace_api.StatusCode.ERROR: + get_status = getattr(accumulator, "get_status", None) + if self._self_body_status is not None: + status = self._self_body_status + elif get_status is not None: + status = get_status() + elif not self._self_parsed: + status = trace_api.Status( + trace_api.StatusCode.ERROR, "Response closed before parsing" + ) + if accumulator is not None: + _finish_tracing(self._self_with_span, accumulator, status) + else: + self._self_with_span.finish_tracing(status=status) + + +class _LegacyRawResponse(_RawResponse): + """LegacyAPIResponse.text is a string property, unlike APIResponse.text().""" + + @property + def text(self) -> str: + return self.__wrapped__.text + + +class _AsyncRawResponse(_RawResponse): + """AsyncAPIResponse has awaitable parse/close; LegacyAPIResponse does not.""" + + async def _consume_body(self, method: str, *args: Any, **kwargs: Any) -> Any: + try: + result = await getattr(self.__wrapped__, method)(*args, **kwargs) + except BaseException as exc: + _finish_on_exit(self, exc) + raise + self._record_body_consumption() + return result + + async def read(self, *args: Any, **kwargs: Any) -> Any: + return await self._consume_body("read", *args, **kwargs) + + async def text(self, *args: Any, **kwargs: Any) -> Any: + return await self._consume_body("text", *args, **kwargs) + + async def json(self, *args: Any, **kwargs: Any) -> Any: + return await self._consume_body("json", *args, **kwargs) + + async def _aiterate_body(self, method: str, *args: Any, **kwargs: Any) -> AsyncIterator[Any]: + is_sse = getattr(self.__wrapped__, "_is_sse_stream", False) + body = bytearray() + too_large = False + try: + async for chunk in getattr(self.__wrapped__, method)(*args, **kwargs): + if not is_sse and not too_large: + encoded = chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") + if method == "iter_lines": + encoded += b"\n" + if len(body) + len(encoded) <= 8 * 1024 * 1024: + body.extend(encoded) + else: + too_large = True + body.clear() + yield chunk + except BaseException as exc: + _finish_on_exit(self, exc) + raise + else: + if too_large: + self._self_body_status = trace_api.Status( + trace_api.StatusCode.ERROR, "Response body too large to inspect" + ) + elif not is_sse: + self._record_body_consumption(bytes(body)) + + def iter_bytes(self, *args: Any, **kwargs: Any) -> AsyncIterator[bytes]: + return self._aiterate_body("iter_bytes", *args, **kwargs) + + def iter_text(self, *args: Any, **kwargs: Any) -> AsyncIterator[str]: + return self._aiterate_body("iter_text", *args, **kwargs) + + def iter_lines(self, *args: Any, **kwargs: Any) -> AsyncIterator[str]: + return self._aiterate_body("iter_lines", *args, **kwargs) + + async def stream_to_file(self, *args: Any, **kwargs: Any) -> None: + await self._consume_body("stream_to_file", *args, **kwargs) + + async def parse(self, *args: Any, **kwargs: Any) -> Any: + try: + parsed = await self.__wrapped__.parse(*args, **kwargs) + except BaseException as exc: + _finish_on_exit(self, exc) + raise + return self._trace_parsed(parsed) + + async def close(self) -> None: + error = None + try: + await self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self, error) + + class _Stream(ObjectProxy): # type: ignore __slots__ = ( "_self_with_span", @@ -60,6 +330,13 @@ def __init__( self._self_is_finished = with_span.is_finished self._self_response_accumulator = response_accumulator + @property + def response(self) -> Any: + response = self.__wrapped__.response + if getattr(self._self_response_accumulator, "get_status", None) is not None: + return _ResponseCloseProxy(response, self) + return response + def __enter__(self) -> "_Stream": # ObjectProxy otherwise returns the SDK stream, bypassing our iterator. self.__wrapped__.__enter__() @@ -74,6 +351,16 @@ def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: finally: _finish_on_exit(self, exc) + def close(self) -> None: + error = None + try: + self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self, error) + def __iter__(self) -> Iterator[Any]: return self @@ -85,7 +372,7 @@ def __next__(self) -> Any: status = trace_api.Status(status_code=trace_api.StatusCode.OK) self._finish_tracing(status=status) raise - except Exception as exception: + except BaseException as exception: if not self._self_is_finished: status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -115,6 +402,12 @@ def _finish_tracing( self, status: Optional[trace_api.Status] = None, ) -> None: + if self._self_is_finished: + return + if status is None or status.status_code is not trace_api.StatusCode.ERROR: + get_status = getattr(self._self_response_accumulator, "get_status", None) + if get_status is not None: + status = get_status() _finish_tracing( status=status, with_span=self._self_with_span, @@ -151,6 +444,13 @@ def __init__( self._self_is_finished = with_span.is_finished self._self_response_accumulator = response_accumulator + @property + def response(self) -> Any: + response = self.__wrapped__.response + if getattr(self._self_response_accumulator, "get_status", None) is not None: + return _ResponseCloseProxy(response, self) + return response + async def __aenter__(self) -> "_AsyncStream": await self.__wrapped__.__aenter__() return self @@ -164,6 +464,16 @@ async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: finally: _finish_on_exit(self, exc) + async def close(self) -> None: + error = None + try: + await self.__wrapped__.close() + except BaseException as exc: + error = exc + raise + finally: + _finish_on_exit(self, error) + def __aiter__(self): return self @@ -175,7 +485,7 @@ async def __anext__(self): status = trace_api.Status(status_code=trace_api.StatusCode.OK) self._finish_tracing(status=status) raise - except Exception as exception: + except BaseException as exception: if not self._self_is_finished: status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -205,6 +515,12 @@ def _finish_tracing( self, status: Optional[trace_api.Status] = None, ) -> None: + if self._self_is_finished: + return + if status is None or status.status_code is not trace_api.StatusCode.ERROR: + get_status = getattr(self._self_response_accumulator, "get_status", None) + if get_status is not None: + status = get_status() _finish_tracing( status=status, with_span=self._self_with_span, diff --git a/python/frameworks/openai/traceai_openai/_with_span.py b/python/frameworks/openai/traceai_openai/_with_span.py index 5f5419ad0..c3b336328 100644 --- a/python/frameworks/openai/traceai_openai/_with_span.py +++ b/python/frameworks/openai/traceai_openai/_with_span.py @@ -1,9 +1,11 @@ import logging -from typing import Any, Optional +import traceback +from typing import Any, Optional, Tuple from fi_instrumentation.fi_types import SpanAttributes from opentelemetry import trace as trace_api from opentelemetry.util.types import Attributes +from traceai_openai._privacy import redact logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -22,6 +24,7 @@ class _WithSpan: __slots__ = ( "_span", + "_secrets", "_primary_attributes", "_context_attributes", "_extra_attributes", @@ -34,8 +37,10 @@ def __init__( context_attributes: Attributes = None, extra_attributes: Attributes = None, initial_attributes: Attributes = None, + secrets: Tuple[str, ...] = (), ) -> None: self._span = span + self._secrets = secrets self._primary_attributes = dict(initial_attributes or {}) self._context_attributes = context_attributes self._extra_attributes = extra_attributes @@ -49,11 +54,25 @@ def __init__( def is_finished(self) -> bool: return self._is_finished - def record_exception(self, exception: Exception) -> None: + def add_secrets(self, secrets: Tuple[str, ...]) -> None: + # Credentials resolved by an SDK token provider are available only after + # sending the request. Retain all known values for this span's lifetime. + self._secrets = tuple(sorted(set((*self._secrets, *secrets)), key=len, reverse=True)) + + def record_exception(self, exception: BaseException) -> None: if self._is_finished: return try: - self._span.record_exception(exception) + # Emit the standard event through FITracer's masking API, without + # handing raw credential-bearing exception text to an exporter. + cls = type(exception) + self._span.add_event("exception", attributes={ + "exception.type": f"{cls.__module__}.{cls.__qualname__}", + "exception.message": redact(str(exception), self._secrets), + "exception.stacktrace": redact("".join(traceback.format_exception( + cls, exception, exception.__traceback__)), self._secrets), + "exception.escaped": "False", + }) except Exception: logger.exception("Failed to record exception on span") @@ -96,7 +115,15 @@ def finish_tracing( self.set_attribute(key, value) if status is not None: try: - self._span.set_status(status=status) + description = redact(status.description, self._secrets) + if description: + # Core removes tracebacks by consulting the in-flight SDK + # exception. Its original message may contain credentials. + # Remove the traceback here so core cannot reconstruct that + # unsanitized message; the redacted event keeps the stack. + description = description.partition("Traceback (most recent call last)")[0] + description = description.rstrip() or None + self._span.set_status(status=trace_api.Status(status.status_code, description)) except Exception: logger.exception("Failed to set status code on span") try: @@ -110,6 +137,7 @@ def set_attribute(self, key: str, value: Any) -> None: if self._is_finished: return try: + value = redact(value, self._secrets) self._span.set_attribute(key, value) if key in self._primary_attributes and value is not None: self._primary_attributes[key] = value diff --git a/python/frameworks/parallel/README.md b/python/frameworks/parallel/README.md index ab85d79d7..f6e9448fa 100644 --- a/python/frameworks/parallel/README.md +++ b/python/frameworks/parallel/README.md @@ -234,7 +234,7 @@ def research(client: Parallel, tracer: FITracer, question: str) -> TaskRunResult span.set_attribute("parallel.run_id", run.run_id) span.set_attribute("parallel.processor", "base") result = client.task_run.result(run.run_id, api_timeout=600, timeout=610) - except Exception: + except BaseException: span.set_status(Status(StatusCode.ERROR, "Task request failed")) raise span.set_attribute("parallel.task.status", result.run.status) diff --git a/python/frameworks/parallel/examples/task_workflow.py b/python/frameworks/parallel/examples/task_workflow.py index a3c32c20f..cbaf2b09f 100644 --- a/python/frameworks/parallel/examples/task_workflow.py +++ b/python/frameworks/parallel/examples/task_workflow.py @@ -25,7 +25,7 @@ def research(client: Parallel, tracer: FITracer, question: str) -> TaskRunResult span.set_attribute("parallel.run_id", run.run_id) span.set_attribute("parallel.processor", "base") result = client.task_run.result(run.run_id, api_timeout=600, timeout=610) - except Exception: + except BaseException: span.set_status(Status(StatusCode.ERROR, "Task request failed")) raise span.set_attribute("parallel.task.status", result.run.status) @@ -48,7 +48,7 @@ async def research_async(client: AsyncParallel, tracer: FITracer, question: str) span.set_attribute("parallel.cancelled", True) span.set_status(Status(StatusCode.ERROR, "Task wait cancelled")) raise - except Exception: + except BaseException: span.set_status(Status(StatusCode.ERROR, "Task request failed")) raise span.set_attribute("parallel.task.status", result.run.status) diff --git a/python/frameworks/parallel/tests/test_task_workflow.py b/python/frameworks/parallel/tests/test_task_workflow.py index 52d40bdb4..1e6e09b93 100644 --- a/python/frameworks/parallel/tests/test_task_workflow.py +++ b/python/frameworks/parallel/tests/test_task_workflow.py @@ -205,3 +205,34 @@ def log_message(self, *_): server.shutdown() server.server_close() thread.join(timeout=3) + + +@pytest.mark.parametrize("stage", ["create", "wait"]) +@pytest.mark.parametrize("error_cls", [KeyboardInterrupt, SystemExit]) +def test_task_interrupt_is_error_with_registered_processor(stage, error_cls): + from fi_instrumentation import register + from fi_instrumentation.fi_types import ProjectType + + exporter = InMemorySpanExporter() + provider = register(project_name="task-interrupt", project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=False, verbose=False) + abort = error_cls("synthetic local abort " + KEY) + + def handler(request): + if stage == "create" or request.method == "GET": + raise abort + return httpx.Response(200, json=run_body("queued")) + + try: + with Parallel(api_key=KEY, max_retries=0, http_client=httpx.Client( + transport=httpx.MockTransport(handler), trust_env=False, + )) as client: + with pytest.raises(error_cls) as caught: + example.research(client, FITracer(provider.get_tracer("task-interrupt")), QUESTION) + assert caught.value is abort + span, = exporter.get_finished_spans() + assert span.status.status_code is StatusCode.ERROR + assert KEY not in span.to_json() + assert span.events == () + finally: + provider.shutdown() diff --git a/python/tests/test_tracer_cancellation.py b/python/tests/test_tracer_cancellation.py new file mode 100644 index 000000000..56be5a513 --- /dev/null +++ b/python/tests/test_tracer_cancellation.py @@ -0,0 +1,79 @@ +"""Local aborts must not become successful spans at the FI processor.""" +import asyncio + +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType +from fi_instrumentation.instrumentation._tracers import FITracer +from opentelemetry import trace +from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode + + +@pytest.mark.parametrize("error_type", [KeyboardInterrupt, asyncio.CancelledError, GeneratorExit, ValueError]) +@pytest.mark.parametrize("record_exception", [False, True]) +@pytest.mark.parametrize("set_status", [False, True]) +@pytest.mark.parametrize("hide_details", [False, True]) +@pytest.mark.parametrize("end_on_exit", [False, True]) +def test_context_abort_preserves_flags_and_exception( + error_type, record_exception, set_status, hide_details, end_on_exit, +): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = FITracer(provider.get_tracer(__name__), TraceConfig(hide_error_details=hide_details)) + error = error_type("synthetic-private-abort") + try: + with tracer.start_as_current_span("parent") as parent: + with pytest.raises(error_type) as caught: + with tracer.start_as_current_span( + "child", record_exception=record_exception, + set_status_on_exception=set_status, end_on_exit=end_on_exit, + ) as child: + assert trace.get_current_span() is child + raise error + assert caught.value is error + assert trace.get_current_span() is parent + if not end_on_exit: + assert child.is_recording() + assert not exporter.get_finished_spans() + child.end() + spans = exporter.get_finished_spans() + assert len(spans) == 2 + span = next(s for s in spans if s.name == "child") + assert span.parent.span_id == parent.get_span_context().span_id + assert span.status.status_code == (StatusCode.ERROR if set_status else StatusCode.UNSET) + exceptions = [e for e in span.events if e.name == "exception"] + assert len(exceptions) == int(record_exception) + wire = encode_spans([span]).SerializeToString() + if hide_details: + assert b"synthetic-private-abort" not in wire + elif record_exception or set_status: + assert b"synthetic-private-abort" in wire + finally: + provider.shutdown() + + +@pytest.mark.parametrize("batch", [False, True]) +@pytest.mark.parametrize("error_type", [KeyboardInterrupt, asyncio.CancelledError]) +def test_registered_processor_exports_interrupted_agent_as_error(batch, error_type): + exporter = InMemorySpanExporter() + provider = register( + project_name="synthetic-cancellation", project_type=ProjectType.OBSERVE, + span_exporter=exporter, batch=batch, verbose=False, + ) + tracer = FITracer(provider.get_tracer(__name__), TraceConfig(hide_error_details=True)) + try: + with pytest.raises(error_type): + with tracer.start_as_current_span("order-assistant", fi_span_kind="agent"): + raise error_type("synthetic-private-abort") + provider.force_flush() + (span,) = exporter.get_finished_spans() + assert span.status.status_code == StatusCode.ERROR + assert len([e for e in span.events if e.name == "exception"]) == 1 + assert b"synthetic-private-abort" not in encode_spans([span]).SerializeToString() + finally: + provider.shutdown()