Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 14 additions & 3 deletions src/agents/models/openai_responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from typing import TYPE_CHECKING, Any, Literal, cast, overload

from openai import NOT_GIVEN, APIStatusError, AsyncOpenAI, AsyncStream, NotGiven
from openai._models import add_request_id
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Accessing _models module may be broken when the underlying "openai" package makes breaking changes in the future. We've been dealing with incompatible changes on the "openai" package side several time. Accessing underscore prefixed (kind of private) module could be even more risky in the long run.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yea makes sense - i'll look into doing this differently later today

from openai.types import ChatModel
from openai.types.responses import (
Response,
Expand Down Expand Up @@ -180,16 +181,26 @@ async def stream_response(
prompt=prompt,
)

request_id = None
stream_response = getattr(stream, "response", None)
if stream_response is not None:
headers = getattr(stream_response, "headers", None)
if headers is not None:
request_id = headers.get("x-request-id")

final_response: Response | None = None

async for chunk in stream:
if isinstance(chunk, ResponseCompletedEvent):
final_response = chunk.response
yield chunk

if final_response and tracing.include_data():
span_response.span_data.response = final_response
span_response.span_data.input = input
if final_response:
if request_id:
add_request_id(final_response, request_id)
if tracing.include_data():
span_response.span_data.response = final_response
span_response.span_data.input = input

except Exception as e:
span_response.set_error(
Expand Down
65 changes: 62 additions & 3 deletions tests/test_agent_tracing.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,22 @@
from __future__ import annotations

import asyncio
from types import SimpleNamespace

import pytest
from inline_snapshot import snapshot
from openai.types.responses import ResponseCompletedEvent

from agents import Agent, RunConfig, Runner, trace
from agents import Agent, OpenAIResponsesModel, RunConfig, Runner, trace
from agents.tracing import ResponseSpanData

from .fake_model import FakeModel
from .fake_model import FakeModel, get_response_obj
from .test_responses import get_text_message
from .testing_processor import assert_no_traces, fetch_normalized_spans
from .testing_processor import (
assert_no_traces,
fetch_normalized_spans,
fetch_ordered_spans,
)


@pytest.mark.asyncio
Expand Down Expand Up @@ -292,6 +299,58 @@ async def test_streaming_single_run_is_single_trace():
)


@pytest.mark.asyncio
@pytest.mark.allow_call_model_methods
async def test_streamed_response_request_id_recorded():
request_id = "req_test_123"

class DummyStream:
def __init__(self) -> None:
self.response = SimpleNamespace(headers={"x-request-id": request_id})

def __aiter__(self):
async def gen():
yield ResponseCompletedEvent(
type="response.completed",
response=get_response_obj([get_text_message("first_test")]),
sequence_number=0,
)

return gen()

class DummyResponses:
async def create(self, **kwargs):
assert kwargs.get("stream") is True
return DummyStream()

class DummyResponsesClient:
def __init__(self) -> None:
self.responses = DummyResponses()

model = OpenAIResponsesModel(model="gpt-4", openai_client=DummyResponsesClient()) # type: ignore[arg-type]

agent = Agent(
name="test_agent",
model=model,
)

result = Runner.run_streamed(agent, input="first_test")
async for _ in result.stream_events():
pass

response_spans = [
span
for span in fetch_ordered_spans()
if isinstance(span.span_data, ResponseSpanData) and span.span_data.response is not None
]

assert response_spans
assert any(
getattr(span.span_data.response, "_request_id", None) == request_id
for span in response_spans
)


@pytest.mark.asyncio
async def test_multiple_streamed_runs_are_multiple_traces():
model = FakeModel()
Expand Down