diff --git a/.github/workflows/baseten-python-contracts.yml b/.github/workflows/baseten-python-contracts.yml new file mode 100644 index 000000000..2ef37d5de --- /dev/null +++ b/.github/workflows/baseten-python-contracts.yml @@ -0,0 +1,65 @@ +name: Baseten Python contracts + +on: + pull_request: + branches: [dev, main] + paths: + - 'python/frameworks/openai/**' + - 'python/fi_instrumentation/**' + - 'python/examples/baseten/**' + - 'python/tests/harness/**' + - '.github/workflows/baseten-python-contracts.yml' + push: + branches: [dev] + paths: + - 'python/frameworks/openai/**' + - 'python/fi_instrumentation/**' + - 'python/examples/baseten/**' + - 'python/tests/harness/**' + - '.github/workflows/baseten-python-contracts.yml' + +permissions: + contents: read + +jobs: + contracts: + name: Baseten OpenAI ${{ matrix.sdk }} + runs-on: ubuntu-24.04 + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + sdk: ['1.69.0', 'latest'] + env: + PYTHONDONTWRITEBYTECODE: '1' + SDK_VERSION: ${{ matrix.sdk }} + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + ref: ${{ github.event.pull_request.head.sha || github.sha }} + persist-credentials: false + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: '3.12' + - name: Install source packages and the selected SDK + run: | + sdk=openai + if [ "$SDK_VERSION" != latest ]; then sdk="openai==$SDK_VERSION"; fi + python -m pip install -e python -e python/frameworks/openai "$sdk" httpx pytest + python -c 'import openai; print("Resolved OpenAI SDK:", openai.__version__)' + - name: Run network-isolated provider and recipe contracts + run: | + sudo unshare --net -- bash -c 'ip link set lo up; exec "$@"' bash \ + env PATH="$PATH" PYTHONDONTWRITEBYTECODE=1 \ + python -m pytest python/frameworks/openai/tests/test_baseten_provider.py \ + python/frameworks/openai/tests/test_openai_compatible_spans.py \ + python/frameworks/openai/tests/test_stream_lifecycle.py \ + python/frameworks/openai/tests/test_stream_context_manager.py \ + python/frameworks/openai/tests/test_high_level_stream_lifecycle.py \ + python/frameworks/openai/tests/test_high_level_stream_error.py \ + --noconftest -p no:cacheprovider -o addopts='' -q + sudo unshare --net -- bash -c 'ip link set lo up; exec "$@"' bash \ + env PATH="$PATH" PYTHONDONTWRITEBYTECODE=1 \ + PYTHONPATH="python/examples/baseten/src:python/tests" \ + python -m pytest python/examples/baseten/tests \ + --noconftest -p no:cacheprovider -o addopts='' -q diff --git a/python/frameworks/openai/README.md b/python/frameworks/openai/README.md index ea7036cec..319f658d7 100644 --- a/python/frameworks/openai/README.md +++ b/python/frameworks/openai/README.md @@ -100,6 +100,25 @@ with OpenAI() as client: trace_provider.force_flush() # Flush before a short-lived script exits. ``` +## Baseten + +Use the same `OpenAIInstrumentor` with Baseten's OpenAI-compatible clients; no +second instrumentor is needed. After the Quickstart tracer setup, create your +client with `api_key=os.environ["BASETEN_API_KEY"]` and the appropriate base URL: + +- Model APIs: `https://inference.baseten.co/v1`. +- Dedicated deployments: your deployment's OpenAI-compatible base URL, such as + `https://model-abc123.api.baseten.co/environments/production/sync/v1`. + +The hosted endpoint and dedicated `model-*.api.baseten.co` endpoints, including +regional model endpoints, are attributed to `baseten`. Unknown proxies retain +`openai` unless you set an explicit client provider override. Supported API +routes and model names depend on the deployment; use the values in +[Baseten's inference API reference](https://docs.baseten.co/reference/inference-api/overview). +The standalone Model APIs recipe in `python/examples/baseten` accepts only Model +APIs URLs; use the OpenAI client directly for a dedicated deployment. Baseten's +separate Performance Client and non-OpenAI protocols need their own tracing setup. + ## Vercel AI Gateway Use the same `OpenAIInstrumentor` for Chat Completions sent to Vercel AI Gateway; diff --git a/python/frameworks/openai/tests/test_baseten_provider.py b/python/frameworks/openai/tests/test_baseten_provider.py new file mode 100644 index 000000000..0ffe6b817 --- /dev/null +++ b/python/frameworks/openai/tests/test_baseten_provider.py @@ -0,0 +1,224 @@ +"""Baseten attribution through the real OpenAI SDK and a loopback OTLP receiver. + +HTTP model responses and credentials are synthetic. Resource coverage asserts +instrumentor behavior, not that every Baseten deployment serves every API. +""" + +import asyncio +import json +import sys +from pathlib import Path +from urllib.parse import urlsplit + +import httpx +import openai +import pytest +from fi_instrumentation import TraceConfig, register +from fi_instrumentation.fi_types import ProjectType + +from traceai_openai import OpenAIInstrumentor +from traceai_openai._providers import provider_from_host + +from .test_openai_compatible_spans import ( + ANSWER, REQUEST_MODEL, RESPONSE_MODEL, async_client, chat_completion, + chat_stream, client, embeddings_body, responses_body, responses_stream, unauthorized, +) + +# Use the repository's collector harness instead of a second telemetry implementation. +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "tests")) +from harness import Receiver, _flatten_attributes + +BASE_URLS = [ + "https://inference.baseten.co/v1", + "https://model-abc123.api.baseten.co/environments/production/sync/v1", + "https://model-abc123-us-west.api.baseten.co/sync/v1", +] +PROMPT = "synthetic-baseten-private-prompt" + + +@pytest.mark.parametrize("host,expected", [ + ("inference.baseten.co", "baseten"), + ("model-abc123.api.baseten.co", "baseten"), + ("model-abc123-us-west.api.baseten.co", "baseten"), + ("MODEL-Abc123.API.BASETEN.CO.", "baseten"), + ("model-a.api.baseten.co", "baseten"), + ("model-" + "a" * 57 + ".api.baseten.co", "baseten"), + ("model-" + "a" * 58 + ".api.baseten.co", "openai"), + ("model-.api.baseten.co", "openai"), + ("model--abc.api.baseten.co", "openai"), + ("model-abc-.api.baseten.co", "openai"), + ("model-abc_123.api.baseten.co", "openai"), + ("model-abc.123.api.baseten.co", "openai"), + ("model-abc.api.baseten.co.customer.example", "openai"), + ("proxy.model-abc.api.baseten.co", "openai"), + ("notmodel-abc.api.baseten.co", "openai"), + ("chain-abc.api.baseten.co", "openai"), + ("api.baseten.co", "openai"), + ("baseten.co", "openai"), + ("inference.baseten.co.customer.example", "openai"), + ("api.openai.com", "openai"), + (None, "openai"), +]) +def test_baseten_host_boundaries(host, expected): + assert provider_from_host(host) == expected + + +@pytest.fixture(scope="module") +def receiver(): + with Receiver() as sink: + yield sink + + +@pytest.fixture +def wire_tracer(receiver, monkeypatch): + # No inherited telemetry destinations, auth, privacy settings or HTTP proxies. + import os + for name in tuple(os.environ): + if name.startswith(("FI_", "OTEL_", "OPENAI_")) or name.lower() in ( + "http_proxy", "https_proxy", "all_proxy", "no_proxy" + ): + monkeypatch.delenv(name) + monkeypatch.setenv("NO_PROXY", "*") + monkeypatch.setenv("FI_API_KEY", "synthetic-fi-key") + monkeypatch.setenv("FI_SECRET_KEY", "synthetic-fi-secret") + monkeypatch.setenv("FI_BASE_URL", receiver.origin) + receiver.clear() + provider = register( + project_name="baseten-contract", project_type=ProjectType.OBSERVE, + set_global_tracer_provider=False, verbose=False, + ) + instrumentor = OpenAIInstrumentor() + + def configure(hidden=False): + instrumentor.instrument(tracer_provider=provider, config=TraceConfig( + hide_inputs=hidden, hide_outputs=hidden, + hide_llm_invocation_parameters=hidden, + )) + return provider + + try: + yield configure + finally: + instrumentor.uninstrument() + provider.force_flush() + provider.shutdown() + + +def only_span(receiver, provider): + assert provider.force_flush() + spans = receiver.spans() + assert len(spans) == 1 + exports = receiver.requests() + assert len(exports) == 1 + assert exports[0]["path"] == "/tracer/v1/traces" + assert exports[0]["headers"]["x-api-key"] == "synthetic-fi-key" + assert exports[0]["headers"]["x-secret-key"] == "synthetic-fi-secret" + assert exports[0]["resource_attributes"][0]["project_name"] == "baseten-contract" + assert exports[0]["resource_attributes"][0]["project_type"] == "observe" + for secret in ("test-vendor-key", "synthetic-fi-key", "synthetic-fi-secret"): + assert secret not in json.dumps(spans) + return spans[0], _flatten_attributes(spans[0]["attributes"]) + + +@pytest.mark.parametrize("base_url", BASE_URLS, ids=["model-api", "dedicated", "regional"]) +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("resource,outcome", [ + ("chat", "success"), ("chat", "stream"), ("chat", "error"), + ("responses", "success"), ("responses", "stream"), ("responses", "error"), + ("embeddings", "success"), ("embeddings", "error"), +]) +@pytest.mark.parametrize("hidden", [False, True], ids=["content", "masked"]) +def test_baseten_wire_contract(receiver, wire_tracer, base_url, use_async, resource, outcome, hidden): + provider = wire_tracer(hidden) + seen = [] + route = "chat/completions" if resource == "chat" else resource + + def handler(request): + seen.append(request) + assert request.url.path == urlsplit(base_url).path + "/" + route + assert request.headers["authorization"] == "Bearer test-vendor-key" + if outcome == "error": + return unauthorized(request) + if outcome == "stream": + return chat_stream() if resource == "chat" else responses_stream() + body = {"chat": chat_completion, "responses": responses_body, + "embeddings": embeddings_body}[resource]() + return httpx.Response(200, json=body) + + kwargs = {"model": REQUEST_MODEL} + if resource == "chat": + kwargs.update(messages=[{"role": "user", "content": PROMPT}], stream=outcome == "stream") + else: + kwargs["input"] = PROMPT + if resource == "responses": + kwargs["stream"] = outcome == "stream" + else: + kwargs["encoding_format"] = "float" + + def create(sdk): + endpoint = sdk.chat.completions if resource == "chat" else getattr(sdk, resource) + return endpoint.create(**kwargs) + + async def run_async(): + async with async_client(handler, base_url=base_url) as sdk: + response = await create(sdk) + if outcome == "stream": + async for _ in response: + pass + + def run_sync(): + with client(handler, base_url=base_url) as sdk: + response = create(sdk) + if outcome == "stream": + list(response) + + if outcome == "error": + with pytest.raises(openai.AuthenticationError): + asyncio.run(run_async()) if use_async else run_sync() + else: + asyncio.run(run_async()) if use_async else run_sync() + assert len(seen) == 1 + span, attrs = only_span(receiver, provider) + assert attrs["gen_ai.provider.name"] == "baseten" + assert attrs["gen_ai.request.model"] == ( + REQUEST_MODEL if outcome == "error" or (resource == "chat" and outcome == "stream") + else RESPONSE_MODEL + ) + assert span["status"]["code"] == ("STATUS_CODE_ERROR" if outcome == "error" else "STATUS_CODE_OK") + if resource == "chat" and outcome == "success": + assert int(attrs["gen_ai.usage.input_tokens"]) == 5 + assert int(attrs["gen_ai.usage.output_tokens"]) == 2 + if hidden: + assert PROMPT not in json.dumps(span) + assert ANSWER not in json.dumps(span) + + +@pytest.mark.parametrize("base_url,expected", [ + ("https://MODEL-Abc123.API.BASETEN.CO.:443/v1", "baseten"), + ("https://model-abc.api.baseten.co.customer.example/v1", "openai"), + ("https://proxy.model-abc.api.baseten.co/v1", "openai"), + ("https://model-abc.api.baseten.co@customer.example/v1", "openai"), + ("https://customer.example/model-abc.api.baseten.co/v1", "openai"), +]) +def test_baseten_url_controls(receiver, wire_tracer, base_url, expected): + provider = wire_tracer() + with client(lambda _: httpx.Response(200, json=chat_completion()), base_url=base_url) as sdk: + sdk.chat.completions.create(model=REQUEST_MODEL, messages=[{"role": "user", "content": PROMPT}]) + _, attrs = only_span(receiver, provider) + assert attrs["gen_ai.provider.name"] == expected + + +def test_baseten_client_endpoint_changes(receiver, wire_tracer): + provider = wire_tracer() + with client(lambda _: httpx.Response(200, json=chat_completion())) as sdk: + for base_url, expected in [ + *[(url, "baseten") for url in BASE_URLS], + ("https://customer.example/v1", "openai"), + ]: + sdk.base_url = base_url + sdk.chat.completions.create( + model=REQUEST_MODEL, messages=[{"role": "user", "content": PROMPT}], + ) + _, attrs = only_span(receiver, provider) + assert attrs["gen_ai.provider.name"] == expected + receiver.clear() diff --git a/python/frameworks/openai/tests/test_high_level_stream_error.py b/python/frameworks/openai/tests/test_high_level_stream_error.py new file mode 100644 index 000000000..260bbfdfc --- /dev/null +++ b/python/frameworks/openai/tests/test_high_level_stream_error.py @@ -0,0 +1,39 @@ +"""The high-level helper must retain user exceptions on unfinished streams.""" +import asyncio + +import pytest +from opentelemetry.trace import StatusCode + +from .test_high_level_stream_lifecycle import _response_events, _stream_context +from .test_stream_lifecycle import base_url, client, streaming, tracing + + +@pytest.mark.parametrize("resource", ["chat", "responses"]) +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("tracing,hide_errors", [(False, False), (True, True)], indirect=["tracing"]) +def test_high_level_context_exception_finishes_once(tracing, base_url, resource, use_async, hide_errors): + provider, exporter = tracing + handler = streaming if resource == "chat" else _response_events + if use_async: + async def run(): + async with client(handler, base_url, True) as sdk: + with pytest.raises(ValueError, match="synthetic consumer error"): + async with _stream_context(sdk, resource) as stream: + await stream.__anext__() + raise ValueError("synthetic consumer error") + assert stream._response.is_closed + asyncio.run(run()) + else: + with client(handler, base_url) as sdk: + with pytest.raises(ValueError, match="synthetic consumer error"): + with _stream_context(sdk, resource) as stream: + next(stream) + raise ValueError("synthetic consumer error") + assert stream._response.is_closed + assert provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == StatusCode.ERROR + assert spans[0].attributes["gen_ai.provider.name"] == "baseten" + if hide_errors: + assert "synthetic consumer error" not in spans[0].to_json() diff --git a/python/frameworks/openai/tests/test_high_level_stream_lifecycle.py b/python/frameworks/openai/tests/test_high_level_stream_lifecycle.py new file mode 100644 index 000000000..fbc047b1a --- /dev/null +++ b/python/frameworks/openai/tests/test_high_level_stream_lifecycle.py @@ -0,0 +1,67 @@ +"""High-level SDK stream helpers must finish spans when their context closes early.""" +import asyncio +import json + +import pytest +from opentelemetry.trace import StatusCode + +from .test_stream_lifecycle import ARGS, base_url, client, http, streaming, tracing + + +def _response_events(_): + response = { + "id": "resp_synthetic", "object": "response", "created_at": 1, + "model": "served-model", "status": "completed", "parallel_tool_calls": True, + "tool_choice": "auto", "tools": [], "output": [], + "usage": {"input_tokens": 5, "output_tokens": 0, "total_tokens": 5}, + } + events = [ + {"type": "response.created", "sequence_number": 0, + "response": {**response, "status": "in_progress"}}, + {"type": "response.completed", "sequence_number": 1, "response": response}, + ] + content = "".join( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events + ).encode() + return http.Response(200, headers={"content-type": "text/event-stream"}, content=content) + + +def _stream_context(sdk, resource): + if resource == "responses": + return sdk.responses.stream(model="requested-model", input="synthetic prompt") + completions = sdk.chat.completions + if not hasattr(completions, "stream"): + completions = sdk.beta.chat.completions + return completions.stream(**ARGS) + + +@pytest.mark.parametrize("resource", ["chat", "responses"]) +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("consumption", ["empty", "partial", "exhausted"]) +def test_high_level_context_finalizes_on_early_exit(tracing, base_url, resource, use_async, consumption): + provider, exporter = tracing + handler = streaming if resource == "chat" else _response_events + if use_async: + async def run(): + async with client(handler, base_url, True) as sdk: + async with _stream_context(sdk, resource) as stream: + if consumption == "partial": + await stream.__anext__() + elif consumption == "exhausted": + async for _ in stream: + pass + assert stream._response.is_closed + asyncio.run(run()) + else: + with client(handler, base_url) as sdk: + with _stream_context(sdk, resource) as stream: + if consumption == "partial": + next(stream) + elif consumption == "exhausted": + list(stream) + assert stream._response.is_closed + assert provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1, "SDK closed the HTTP response without ending its span" + assert spans[0].status.status_code == StatusCode.OK + assert spans[0].attributes["gen_ai.provider.name"] == "baseten" diff --git a/python/frameworks/openai/tests/test_request_interrupt.py b/python/frameworks/openai/tests/test_request_interrupt.py new file mode 100644 index 000000000..a944ec05c --- /dev/null +++ b/python/frameworks/openai/tests/test_request_interrupt.py @@ -0,0 +1,83 @@ +"""Native HTTP request interruption must finish the span before propagating.""" + +import os + +import openai +import pytest +from fi_instrumentation import TraceConfig +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 +from traceai_openai import OpenAIInstrumentor + +if int(openai.__version__.split(".")[0]) >= 3: + import httpx2 as http +else: + import httpx as http + + +@pytest.mark.parametrize("base_url", [ + "https://model-synthetic.api.baseten.co/sync/v1", + "https://model-synthetic-us-west.api.baseten.co/sync/v1", +]) +@pytest.mark.parametrize("hide_errors", [False, True]) +@pytest.mark.parametrize("interrupt", [False, True], ids=["success-control", "interrupted"]) +def test_sync_request_interrupt_finishes_span(base_url, hide_errors, interrupt, monkeypatch): + for name in tuple(os.environ): + if name.startswith(("FI_", "OTEL_", "OPENAI_")): + monkeypatch.delenv(name) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = OpenAIInstrumentor() + instrumentor.instrument( + tracer_provider=provider, config=TraceConfig(hide_error_details=hide_errors), + ) + error = KeyboardInterrupt("synthetic request interruption") + + def handle(_): + if interrupt: + raise error + return http.Response(200, json={ + "id": "synthetic", "object": "chat.completion", "created": 1, + "model": "served-model", "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + }], + }) + + try: + with openai.OpenAI( + api_key="synthetic", base_url=base_url, max_retries=0, + http_client=http.Client(transport=http.MockTransport(handle), trust_env=False), + ) as client: + if interrupt: + with pytest.raises(KeyboardInterrupt) as caught: + client.chat.completions.create( + model="requested-model", messages=[{"role": "user", "content": "hi"}], + ) + assert caught.value is error + else: + client.chat.completions.create( + model="requested-model", messages=[{"role": "user", "content": "hi"}], + ) + # Must finish before client close; force_flush cannot finish an active span. + assert provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.attributes["gen_ai.provider.name"] == "baseten" + assert span.status.status_code == ( + StatusCode.ERROR if interrupt else StatusCode.OK + ) + if interrupt: + if hide_errors: + assert "synthetic request interruption" not in span.to_json() + else: + assert sum(event.name == "exception" for event in span.events) == 1 + assert exporter.get_finished_spans() == spans + finally: + instrumentor.uninstrument() + provider.shutdown() diff --git a/python/frameworks/openai/tests/test_stream_lifecycle.py b/python/frameworks/openai/tests/test_stream_lifecycle.py new file mode 100644 index 000000000..b1ed19ed5 --- /dev/null +++ b/python/frameworks/openai/tests/test_stream_lifecycle.py @@ -0,0 +1,360 @@ +"""Explicit close and task cancellation through the SDK's native HTTP transport.""" + +import asyncio +from contextlib import aclosing +import inspect +import json +import os +import sys + +import openai +import pytest +from fi_instrumentation import TraceConfig +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 + +from traceai_openai import OpenAIInstrumentor + +# OpenAI 3 uses httpx2; do not exercise it through the legacy httpx adapter. +if int(openai.__version__.split(".")[0]) >= 3: + import httpx2 as http +else: + import httpx as http + +ARGS = {"model": "requested-model", "messages": [{"role": "user", "content": "hello"}]} +CHUNK = { + "id": "test", "object": "chat.completion.chunk", "created": 1, "model": "served-model", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "partial-answer"}, + "finish_reason": None}], +} +CHUNK_BYTES = ("data: " + json.dumps(CHUNK) + "\n\n").encode() + + +@pytest.fixture(params=[ + "https://model-synthetic.api.baseten.co/environments/production/sync/v1", + "https://model-synthetic-us-west.api.baseten.co/sync/v1", +]) +def base_url(request): + return request.param + + +@pytest.fixture +def tracing(monkeypatch, request): + for name in tuple(os.environ): + if name.startswith(("FI_", "OTEL_", "OPENAI_")): + monkeypatch.delenv(name) + config = TraceConfig() + if getattr(request, "param", False): + if "hide_error_details" not in inspect.signature(TraceConfig).parameters: + pytest.skip("The locked legacy core predates error-detail masking") + config = TraceConfig(hide_error_details=True) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=config) + yield provider, exporter + instrumentor.uninstrument() + provider.shutdown() + + +def client(handler, base_url, use_async=False): + sdk_type = openai.AsyncOpenAI if use_async else openai.OpenAI + http_type = http.AsyncClient if use_async else http.Client + return sdk_type( + api_key="synthetic-key", base_url=base_url, max_retries=0, + http_client=http_type(transport=http.MockTransport(handler), trust_env=False), + ) + + +def streaming(_): + return http.Response( + 200, headers={"content-type": "text/event-stream"}, + content=CHUNK_BYTES + b"data: [DONE]\n\n", + ) + + +def cancelled(message): + # Task.cancel(msg) propagates msg to the awaiter only on Python 3.11+. + return pytest.raises( + asyncio.CancelledError, match=message if sys.version_info >= (3, 11) else None, + ) + + +def only_span(tracing, status): + provider, exporter = tracing + assert provider.force_flush() + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.status.status_code == status + assert span.attributes["gen_ai.provider.name"] == "baseten" + return span + + +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("consumption", ["empty", "partial", "exhausted"]) +def test_explicit_close_finishes_once(tracing, base_url, use_async, consumption): + if use_async: + async def run(): + async with client(streaming, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + if consumption == "partial": + await stream.__anext__() + elif consumption == "exhausted": + async for _ in stream: + pass + assert await stream.close() is None + assert stream.response.is_closed + only_span(tracing, StatusCode.OK) + # Repeated close and context-manager cleanup must not end twice. + await stream.close() + await stream.__aexit__(None, None, None) + asyncio.run(run()) + else: + with client(streaming, base_url) as sdk: + stream = sdk.chat.completions.create(**ARGS, stream=True) + if consumption == "partial": + next(stream) + elif consumption == "exhausted": + list(stream) + assert stream.close() is None + assert stream.response.is_closed + only_span(tracing, StatusCode.OK) + stream.close() + stream.__exit__(None, None, None) + span = only_span(tracing, StatusCode.OK) + if consumption != "empty": + assert span.attributes["output.value"] == "partial-answer" + assert span.attributes["gen_ai.request.model"] == ARGS["model"] + + +@pytest.mark.skipif( + not hasattr(openai.AsyncStream, "aclose"), + reason="This SDK version does not expose AsyncStream.aclose", +) +@pytest.mark.parametrize("consumption", ["empty", "partial", "exhausted"]) +@pytest.mark.parametrize("cleanup", ["aclose", "stdlib_aclosing"]) +def test_async_close_alias_finishes_once(tracing, base_url, consumption, cleanup): + async def run(): + async with client(streaming, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + + async def consume(): + if consumption == "partial": + await stream.__anext__() + elif consumption == "exhausted": + async for _ in stream: + pass + + if cleanup == "stdlib_aclosing": + async with aclosing(stream): + await consume() + else: + await consume() + assert await stream.aclose() is None + assert stream.response.is_closed + only_span(tracing, StatusCode.OK) + # Interleaved aliases and context cleanup must not export duplicates. + await stream.aclose() + await stream.close() + await stream.__aexit__(None, None, None) + + asyncio.run(run()) + span = only_span(tracing, StatusCode.OK) + if consumption != "empty": + assert span.attributes["output.value"] == "partial-answer" + assert span.attributes["gen_ai.request.model"] == ARGS["model"] + + +@pytest.mark.parametrize("use_async,close_method", [ + (False, "close"), (True, "close"), + pytest.param(True, "aclose", marks=pytest.mark.skipif( + not hasattr(openai.AsyncStream, "aclose"), + reason="This SDK version does not expose AsyncStream.aclose", + )), +]) +def test_close_failure_preserves_exception_and_finishes( + tracing, base_url, use_async, close_method, monkeypatch, +): + error = RuntimeError("synthetic close failure") + calls = [] + + if use_async: + async def fail_close(): + calls.append(True) + raise error + + async def run(): + async with client(streaming, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + await stream.__anext__() + with monkeypatch.context() as patch: + patch.setattr(stream.response, "aclose", fail_close) + with pytest.raises(RuntimeError) as caught: + await getattr(stream, close_method)() + assert caught.value is error + only_span(tracing, StatusCode.ERROR) + await stream.close() + asyncio.run(run()) + else: + def fail_close(): + calls.append(True) + raise error + + with client(streaming, base_url) as sdk: + stream = sdk.chat.completions.create(**ARGS, stream=True) + next(stream) + with monkeypatch.context() as patch: + patch.setattr(stream.response, "close", fail_close) + with pytest.raises(RuntimeError) as caught: + stream.close() + assert caught.value is error + stream.close() + assert calls == [True] + span = only_span(tracing, StatusCode.ERROR) + assert sum(event.name == "exception" for event in span.events) == 1 + + +@pytest.mark.parametrize("tracing,hide_errors", [(False, False), (True, True)], indirect=["tracing"]) +def test_sync_stream_interrupt_finishes_before_close(tracing, base_url, hide_errors): + class InterruptBody(http.SyncByteStream): + def __iter__(self): + yield CHUNK_BYTES + raise KeyboardInterrupt("synthetic stream interruption") + + def close(self): + pass + + def handler(_): + return http.Response( + 200, headers={"content-type": "text/event-stream"}, stream=InterruptBody(), + ) + + with client(handler, base_url) as sdk: + stream = sdk.chat.completions.create(**ARGS, stream=True) + assert next(stream).choices[0].delta.content == "partial-answer" + with pytest.raises(KeyboardInterrupt, match="synthetic stream interruption"): + next(stream) + span = only_span(tracing, StatusCode.ERROR) + stream.close() + stream.close() + assert stream.response.is_closed + + assert only_span(tracing, StatusCode.ERROR) is span + assert span.attributes["output.value"] == "partial-answer" + if hide_errors: + assert "synthetic stream interruption" not in span.to_json() + else: + assert sum(event.name == "exception" for event in span.events) == 1 + + +@pytest.mark.parametrize("tracing,hide_errors", [(False, False), (True, True)], indirect=["tracing"]) +def test_cancel_pending_request_finishes(tracing, base_url, hide_errors): + async def run(): + reached_transport = asyncio.Event() + + async def handler(_): + reached_transport.set() + await asyncio.Event().wait() + + async with client(handler, base_url, True) as sdk: + task = asyncio.create_task(sdk.chat.completions.create(**ARGS)) + await asyncio.wait_for(reached_transport.wait(), 2) + task.cancel("synthetic request cancellation") + with cancelled("synthetic request cancellation"): + await task + assert task.cancelled() + # Finalization must not depend on subsequent client cleanup. + only_span(tracing, StatusCode.ERROR) + asyncio.run(run()) + span = only_span(tracing, StatusCode.ERROR) + if hide_errors: + assert "synthetic request cancellation" not in span.to_json() + else: + assert any(event.attributes["exception.type"] == "asyncio.exceptions.CancelledError" + for event in span.events if event.name == "exception") + + +@pytest.mark.parametrize("tracing,hide_errors", [(False, False), (True, True)], indirect=["tracing"]) +def test_cancel_stream_iteration_finishes_before_close(tracing, base_url, hide_errors): + async def run(): + waiting_for_chunk = asyncio.Event() + closed = [] + + class SlowBody(http.AsyncByteStream): + async def __aiter__(self): + yield CHUNK_BYTES + waiting_for_chunk.set() + await asyncio.Event().wait() + + async def aclose(self): + closed.append(True) + + def handler(_): + return http.Response( + 200, headers={"content-type": "text/event-stream"}, stream=SlowBody(), + ) + + async with client(handler, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + await stream.__anext__() + task = asyncio.create_task(stream.__anext__()) + await asyncio.wait_for(waiting_for_chunk.wait(), 2) + task.cancel("synthetic stream cancellation") + with cancelled("synthetic stream cancellation"): + await task + assert task.cancelled() + only_span(tracing, StatusCode.ERROR) + await stream.close() + await stream.close() + await stream.__aexit__(None, None, None) + assert stream.response.is_closed + assert closed == [True] + asyncio.run(run()) + span = only_span(tracing, StatusCode.ERROR) + assert span.attributes["output.value"] == "partial-answer" + if hide_errors: + assert "synthetic stream cancellation" not in span.to_json() + else: + assert sum(event.name == "exception" for event in span.events) == 1 + + +@pytest.mark.parametrize("close_method", [ + "close", + pytest.param("aclose", marks=pytest.mark.skipif( + not hasattr(openai.AsyncStream, "aclose"), + reason="This SDK version does not expose AsyncStream.aclose", + )), +]) +def test_cancel_async_close_finishes(tracing, base_url, close_method): + async def run(): + closing = asyncio.Event() + + class SlowCloseBody(http.AsyncByteStream): + async def __aiter__(self): + yield CHUNK_BYTES + + async def aclose(self): + closing.set() + await asyncio.Event().wait() + + def handler(_): + return http.Response( + 200, headers={"content-type": "text/event-stream"}, stream=SlowCloseBody(), + ) + + async with client(handler, base_url, True) as sdk: + stream = await sdk.chat.completions.create(**ARGS, stream=True) + await stream.__anext__() + task = asyncio.create_task(getattr(stream, close_method)()) + await asyncio.wait_for(closing.wait(), 2) + task.cancel("synthetic close cancellation") + with cancelled("synthetic close cancellation"): + await task + assert task.cancelled() + only_span(tracing, StatusCode.ERROR) + asyncio.run(run()) + only_span(tracing, StatusCode.ERROR) diff --git a/python/frameworks/openai/traceai_openai/_providers.py b/python/frameworks/openai/traceai_openai/_providers.py index 2c426e46f..22330501b 100644 --- a/python/frameworks/openai/traceai_openai/_providers.py +++ b/python/frameworks/openai/traceai_openai/_providers.py @@ -6,6 +6,7 @@ instrumentor's historical label, "openai", rather than a guess. """ +import re from typing import Optional from fi_instrumentation.fi_types import FiLLMProviderValues @@ -58,6 +59,10 @@ def provider_from_host(host: Optional[str]) -> str: host = host.lower().rstrip(".") if host in _HOSTS: return _HOSTS[host] + # Dedicated and regional Baseten endpoints use a single model-* DNS label. + # Bound it to 63 characters; do not label arbitrary subdomains or look-alikes. + if re.fullmatch(r"model-[a-z0-9](?:[a-z0-9-]{0,55}[a-z0-9])?\.api\.baseten\.co", host): + return "baseten" for suffix, provider in _SUFFIXES: if host.endswith(suffix): return provider diff --git a/python/frameworks/openai/traceai_openai/_request.py b/python/frameworks/openai/traceai_openai/_request.py index c6e8cf426..e9a60050d 100644 --- a/python/frameworks/openai/traceai_openai/_request.py +++ b/python/frameworks/openai/traceai_openai/_request.py @@ -340,7 +340,8 @@ def __call__( cast_to=cast_to, request_parameters=request_parameters, ) - except Exception as e: + except BaseException as e: + # end_on_exit=False: an interrupted request also needs finalization. status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, description=f"{type(e).__name__}: {e}", @@ -406,7 +407,8 @@ async def __call__( add_io_to_span_attributes( with_span, None, response, is_streaming=self._is_streaming(response) ) - except Exception as exception: + except BaseException as exception: + # end_on_exit=False: cancellation also needs explicit finalization. with_span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, diff --git a/python/frameworks/openai/traceai_openai/_stream.py b/python/frameworks/openai/traceai_openai/_stream.py index a3a60042e..956d4d11f 100644 --- a/python/frameworks/openai/traceai_openai/_stream.py +++ b/python/frameworks/openai/traceai_openai/_stream.py @@ -1,4 +1,7 @@ +import asyncio import logging +import sys +import weakref from typing import Any, AsyncIterator, Generator, Iterator, Optional, Protocol, Tuple from openai.types.chat import ChatCompletionChunk @@ -26,7 +29,7 @@ def get_extra_attributes(self) -> Iterator[Tuple[str, AttributeValue]]: ... def _finish_on_exit(stream: Any, error: Optional[BaseException]) -> None: - """Finish even on an early context exit, without ending an exhausted span twice.""" + """Finish on close, context exit or error without ending a span twice.""" if stream._self_is_finished: return if error is None: @@ -46,6 +49,7 @@ class _Stream(ObjectProxy): # type: ignore "_self_iteration_count", "_self_is_finished", "_self_response_accumulator", + "_self_defer_close", ) def __init__( @@ -59,6 +63,29 @@ def __init__( self._self_iteration_count = 0 self._self_is_finished = with_span.is_finished self._self_response_accumulator = response_accumulator + self._self_defer_close = False + # The SDK's high-level .stream() helpers close the HTTP response + # directly, bypassing raw_stream.close() on early context exit. + response = stream.response + original_close = response.close + stream_ref = weakref.ref(self) + + def close_response() -> Any: + # A high-level context manager closes the response while handling + # the user's exception; preserve that failure on the span. + error: Optional[BaseException] = sys.exc_info()[1] + try: + return original_close() + except BaseException as close_error: + error = close_error + raise + finally: + response.close = original_close + traced = stream_ref() + if traced is not None and not traced._self_defer_close: + _finish_on_exit(traced, error) + + response.close = close_response def __enter__(self) -> "_Stream": # ObjectProxy otherwise returns the SDK stream, bypassing our iterator. @@ -66,33 +93,46 @@ def __enter__(self) -> "_Stream": return self def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: + self._self_defer_close = True try: return self.__wrapped__.__exit__(exc_type, exc, tb) except BaseException as close_error: exc = close_error raise finally: + self._self_defer_close = False _finish_on_exit(self, exc) + def close(self) -> Any: + error: Optional[BaseException] = None + self._self_defer_close = True + try: + return self.__wrapped__.close() + except BaseException as close_error: + error = close_error + raise + finally: + self._self_defer_close = False + _finish_on_exit(self, error) + def __iter__(self) -> Iterator[Any]: return self def __next__(self) -> Any: try: - chunk = self.__wrapped__.__next__() + self._self_defer_close = True + try: + chunk = self.__wrapped__.__next__() + finally: + self._self_defer_close = False except StopIteration: if not self._self_is_finished: status = trace_api.Status(status_code=trace_api.StatusCode.OK) self._finish_tracing(status=status) raise - except Exception as exception: - if not self._self_is_finished: - status = trace_api.Status( - status_code=trace_api.StatusCode.ERROR, - description=f"{type(exception).__name__}: {exception}", - ) - self._self_with_span.record_exception(exception) - self._finish_tracing(status=status) + except BaseException as exception: + # Interruptions also terminate the stream; preserve the error span. + _finish_on_exit(self, exception) raise else: self._process_chunk(chunk) @@ -137,6 +177,7 @@ class _AsyncStream(ObjectProxy): "_self_iteration_count", "_self_is_finished", "_self_response_accumulator", + "_self_defer_close", ) def __init__( @@ -150,39 +191,75 @@ def __init__( self._self_iteration_count = 0 self._self_is_finished = with_span.is_finished self._self_response_accumulator = response_accumulator + self._self_defer_close = None + response = stream.response + original_close = response.aclose + stream_ref = weakref.ref(self) + + async def close_response() -> Any: + error: Optional[BaseException] = sys.exc_info()[1] + try: + return await original_close() + except BaseException as close_error: + error = close_error + raise + finally: + response.aclose = original_close + traced = stream_ref() + if traced is not None and traced._self_defer_close is not asyncio.current_task(): + _finish_on_exit(traced, error) + + response.aclose = close_response async def __aenter__(self) -> "_AsyncStream": await self.__wrapped__.__aenter__() return self async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: + self._self_defer_close = asyncio.current_task() try: return await self.__wrapped__.__aexit__(exc_type, exc, tb) except BaseException as close_error: exc = close_error raise finally: + self._self_defer_close = None _finish_on_exit(self, exc) + async def close(self) -> Any: + error: Optional[BaseException] = None + self._self_defer_close = asyncio.current_task() + try: + return await self.__wrapped__.close() + except BaseException as close_error: + error = close_error + raise + finally: + self._self_defer_close = None + _finish_on_exit(self, error) + + async def aclose(self) -> Any: + # The SDK alias must use our close path, not bypass span finalization. + return await self.close() + def __aiter__(self): return self async def __anext__(self): try: - chunk = await self.__wrapped__.__anext__() + self._self_defer_close = asyncio.current_task() + try: + chunk = await self.__wrapped__.__anext__() + finally: + self._self_defer_close = None except StopAsyncIteration: if not self._self_is_finished: status = trace_api.Status(status_code=trace_api.StatusCode.OK) self._finish_tracing(status=status) raise - except Exception as exception: - if not self._self_is_finished: - status = trace_api.Status( - status_code=trace_api.StatusCode.ERROR, - description=f"{type(exception).__name__}: {exception}", - ) - self._self_with_span.record_exception(exception) - self._finish_tracing(status=status) + except BaseException as exception: + # Cancellation is a BaseException; finalize without swallowing it. + _finish_on_exit(self, exception) raise else: self._process_chunk(chunk) diff --git a/python/frameworks/openai/traceai_openai/_with_span.py b/python/frameworks/openai/traceai_openai/_with_span.py index 5f5419ad0..923a31af8 100644 --- a/python/frameworks/openai/traceai_openai/_with_span.py +++ b/python/frameworks/openai/traceai_openai/_with_span.py @@ -49,7 +49,7 @@ def __init__( def is_finished(self) -> bool: return self._is_finished - def record_exception(self, exception: Exception) -> None: + def record_exception(self, exception: BaseException) -> None: if self._is_finished: return try: