Skip to content

Commit 4da62ce

Browse files
committed
fix: #885
1 parent e10cb23 commit 4da62ce

File tree

6 files changed

+3154
-2597
lines changed

6 files changed

+3154
-2597
lines changed

pydantic_ai_slim/pydantic_ai/messages.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -858,6 +858,12 @@ class ModelResponse:
858858
provider_request_id: str | None = None
859859
"""request ID as specified by the model provider. This can be used to track the specific request to the model."""
860860

861+
id: str | None = None
862+
"""Response ID as specified by the model provider. Used to populate gen_ai.response.id in OpenTelemetry."""
863+
864+
finish_reason: str | None = None
865+
"""Reason the model finished generating the response. Used to populate gen_ai.response.finish_reasons in OpenTelemetry."""
866+
861867
def otel_events(self, settings: InstrumentationSettings) -> list[Event]:
862868
"""Return OpenTelemetry events for the response."""
863869
result: list[Event] = []

pydantic_ai_slim/pydantic_ai/models/__init__.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -544,6 +544,8 @@ class StreamedResponse(ABC):
544544

545545
model_request_parameters: ModelRequestParameters
546546
final_result_event: FinalResultEvent | None = field(default=None, init=False)
547+
_id: str | None = field(default=None, init=False)
548+
_finish_reason: str | None = field(default=None, init=False)
547549

548550
_parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
549551
_event_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
@@ -598,6 +600,8 @@ def get(self) -> ModelResponse:
598600
model_name=self.model_name,
599601
timestamp=self.timestamp,
600602
usage=self.usage(),
603+
id=self._id,
604+
finish_reason=self._finish_reason,
601605
)
602606

603607
def usage(self) -> RequestUsage:

pydantic_ai_slim/pydantic_ai/models/instrumented.py

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -299,22 +299,27 @@ def _record_metrics():
299299

300300
events = self.instrumentation_settings.messages_to_otel_events(messages)
301301
for event in self.instrumentation_settings.messages_to_otel_events([response]):
302+
choice_body: dict[str, Any] = {
303+
'index': 0,
304+
'message': event.body,
305+
}
306+
if response.finish_reason is not None:
307+
choice_body['finish_reason'] = response.finish_reason
302308
events.append(
303309
Event(
304310
'gen_ai.choice',
305-
body={
306-
# TODO finish_reason
307-
'index': 0,
308-
'message': event.body,
309-
},
311+
body=choice_body,
310312
)
311313
)
312-
span.set_attributes(
313-
{
314-
**response.usage.opentelemetry_attributes(),
315-
'gen_ai.response.model': response_model,
316-
}
317-
)
314+
response_attributes = {
315+
**response.usage.opentelemetry_attributes(),
316+
'gen_ai.response.model': response_model,
317+
}
318+
if response.id is not None:
319+
response_attributes['gen_ai.response.id'] = response.id
320+
if response.finish_reason is not None:
321+
response_attributes['gen_ai.response.finish_reasons'] = [response.finish_reason]
322+
span.set_attributes(response_attributes)
318323
span.update_name(f'{operation} {request_model}')
319324
for event in events:
320325
event.attributes = {

pydantic_ai_slim/pydantic_ai/models/openai.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -499,6 +499,8 @@ def _process_response(self, response: chat.ChatCompletion | str) -> ModelRespons
499499
timestamp=timestamp,
500500
provider_details=vendor_details,
501501
provider_request_id=response.id,
502+
id=response.id,
503+
finish_reason=choice.finish_reason,
502504
)
503505

504506
async def _process_streamed_response(
@@ -801,6 +803,8 @@ def _process_response(self, response: responses.Response) -> ModelResponse:
801803
model_name=response.model,
802804
provider_request_id=response.id,
803805
timestamp=timestamp,
806+
id=response.id,
807+
finish_reason=response.status,
804808
)
805809

806810
async def _process_streamed_response(
@@ -1140,11 +1144,19 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
11401144
async for chunk in self._response:
11411145
self._usage += _map_usage(chunk)
11421146

1147+
# Capture the response ID from the chunk
1148+
if chunk.id and self._id is None:
1149+
self._id = chunk.id
1150+
11431151
try:
11441152
choice = chunk.choices[0]
11451153
except IndexError:
11461154
continue
11471155

1156+
# Capture the finish_reason when it becomes available
1157+
if choice.finish_reason and self._finish_reason is None:
1158+
self._finish_reason = choice.finish_reason
1159+
11481160
# Handle the text part of the response
11491161
content = choice.delta.content
11501162
if content is not None:
@@ -1197,6 +1209,11 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
11971209
# NOTE: You can inspect the builtin tools used checking the `ResponseCompletedEvent`.
11981210
if isinstance(chunk, responses.ResponseCompletedEvent):
11991211
self._usage += _map_usage(chunk.response)
1212+
# Capture id and finish_reason from completed response
1213+
if chunk.response.id and self._id is None:
1214+
self._id = chunk.response.id
1215+
if chunk.response.status and self._finish_reason is None:
1216+
self._finish_reason = chunk.response.status
12001217

12011218
elif isinstance(chunk, responses.ResponseContentPartAddedEvent):
12021219
pass # there's nothing we need to do here
@@ -1205,7 +1222,9 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
12051222
pass # there's nothing we need to do here
12061223

12071224
elif isinstance(chunk, responses.ResponseCreatedEvent):
1208-
pass # there's nothing we need to do here
1225+
# Capture id from created response
1226+
if chunk.response.id and self._id is None:
1227+
self._id = chunk.response.id
12091228

12101229
elif isinstance(chunk, responses.ResponseFailedEvent): # pragma: no cover
12111230
self._usage += _map_usage(chunk.response)

tests/models/test_openai.py

Lines changed: 96 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -151,13 +151,18 @@ def get_mock_chat_completion_kwargs(async_open_ai: AsyncOpenAI) -> list[dict[str
151151

152152

153153
def completion_message(
154-
message: ChatCompletionMessage, *, usage: CompletionUsage | None = None, logprobs: ChoiceLogprobs | None = None
154+
message: ChatCompletionMessage,
155+
*,
156+
usage: CompletionUsage | None = None,
157+
logprobs: ChoiceLogprobs | None = None,
158+
response_id: str = '123',
159+
finish_reason: str = 'stop'
155160
) -> chat.ChatCompletion:
156-
choices = [Choice(finish_reason='stop', index=0, message=message)]
161+
choices = [Choice(finish_reason=finish_reason, index=0, message=message)]
157162
if logprobs:
158-
choices = [Choice(finish_reason='stop', index=0, message=message, logprobs=logprobs)]
163+
choices = [Choice(finish_reason=finish_reason, index=0, message=message, logprobs=logprobs)]
159164
return chat.ChatCompletion(
160-
id='123',
165+
id=response_id,
161166
choices=choices,
162167
created=1704067200, # 2024-01-01
163168
model='gpt-4o-123',
@@ -192,13 +197,17 @@ async def test_request_simple_success(allow_model_requests: None):
192197
model_name='gpt-4o-123',
193198
timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
194199
provider_request_id='123',
200+
id='123',
201+
finish_reason='stop',
195202
),
196203
ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
197204
ModelResponse(
198205
parts=[TextPart(content='world')],
199206
model_name='gpt-4o-123',
200207
timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
201208
provider_request_id='123',
209+
id='123',
210+
finish_reason='stop',
202211
),
203212
]
204213
)
@@ -242,6 +251,36 @@ async def test_request_simple_usage(allow_model_requests: None):
242251
)
243252

244253

254+
async def test_id_and_finish_reason_fields(allow_model_requests: None):
255+
"""Test that id and finish_reason fields are properly populated in ModelResponse."""
256+
# Test with different finish reasons
257+
test_cases = [
258+
('stop', 'response-id-1'),
259+
('length', 'response-id-2'),
260+
('tool_calls', 'response-id-3'),
261+
]
262+
263+
for finish_reason, response_id in test_cases:
264+
c = completion_message(
265+
ChatCompletionMessage(content='test response', role='assistant'),
266+
response_id=response_id,
267+
finish_reason=finish_reason,
268+
)
269+
mock_client = MockOpenAI.create_mock(c)
270+
m = OpenAIModel('gpt-4o', provider=OpenAIProvider(openai_client=mock_client))
271+
agent = Agent(m)
272+
273+
result = await agent.run('test')
274+
assert result.output == 'test response'
275+
276+
# Check that the ModelResponse contains the correct id and finish_reason
277+
messages = result.all_messages()
278+
model_response = messages[1] # Second message should be the model response
279+
assert isinstance(model_response, ModelResponse)
280+
assert model_response.id == response_id
281+
assert model_response.finish_reason == finish_reason
282+
283+
245284
async def test_request_structured_response(allow_model_requests: None):
246285
c = completion_message(
247286
ChatCompletionMessage(
@@ -422,9 +461,9 @@ async def get_location(loc_name: str) -> str:
422461
FinishReason = Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call']
423462

424463

425-
def chunk(delta: list[ChoiceDelta], finish_reason: FinishReason | None = None) -> chat.ChatCompletionChunk:
464+
def chunk(delta: list[ChoiceDelta], finish_reason: FinishReason | None = None, chunk_id: str = 'x') -> chat.ChatCompletionChunk:
426465
return chat.ChatCompletionChunk(
427-
id='x',
466+
id=chunk_id,
428467
choices=[
429468
ChunkChoice(index=index, delta=delta, finish_reason=finish_reason) for index, delta in enumerate(delta)
430469
],
@@ -435,8 +474,8 @@ def chunk(delta: list[ChoiceDelta], finish_reason: FinishReason | None = None) -
435474
)
436475

437476

438-
def text_chunk(text: str, finish_reason: FinishReason | None = None) -> chat.ChatCompletionChunk:
439-
return chunk([ChoiceDelta(content=text, role='assistant')], finish_reason=finish_reason)
477+
def text_chunk(text: str, finish_reason: FinishReason | None = None, chunk_id: str = 'x') -> chat.ChatCompletionChunk:
478+
return chunk([ChoiceDelta(content=text, role='assistant')], finish_reason=finish_reason, chunk_id=chunk_id)
440479

441480

442481
async def test_stream_text(allow_model_requests: None):
@@ -552,6 +591,55 @@ async def test_stream_structured_finish_reason(allow_model_requests: None):
552591
assert result.is_complete
553592

554593

594+
async def test_stream_id_and_finish_reason_fields(allow_model_requests: None):
595+
"""Test that streaming responses properly track id and finish_reason fields."""
596+
# Test streaming text response
597+
stream = [
598+
text_chunk('hello ', chunk_id='stream-response-123'),
599+
text_chunk('world', chunk_id='stream-response-123'),
600+
text_chunk('!', finish_reason='stop', chunk_id='stream-response-123'),
601+
]
602+
mock_client = MockOpenAI.create_mock_stream(stream)
603+
m = OpenAIModel('gpt-4o', provider=OpenAIProvider(openai_client=mock_client))
604+
agent = Agent(m)
605+
606+
async with agent.run_stream('test') as result:
607+
assert not result.is_complete
608+
text_chunks = [c async for c in result.stream_text(debounce_by=None)]
609+
assert text_chunks == ['hello ', 'hello world', 'hello world!']
610+
assert result.is_complete
611+
612+
# Get the final messages and check the ModelResponse
613+
messages = result.all_messages()
614+
model_response = messages[1] # Second message should be the model response
615+
assert isinstance(model_response, ModelResponse)
616+
assert model_response.id == 'stream-response-123'
617+
assert model_response.finish_reason == 'stop'
618+
619+
# Test streaming with structured output and different finish reason
620+
stream = [
621+
struc_chunk('final_result', None),
622+
chunk([ChoiceDelta(tool_calls=[ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(name=None, arguments='{"first": "Test"'))])], chunk_id='struct-response-456'),
623+
chunk([ChoiceDelta(tool_calls=[ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(name=None, arguments='}'))])], finish_reason='length', chunk_id='struct-response-456'),
624+
]
625+
mock_client = MockOpenAI.create_mock_stream(stream)
626+
m = OpenAIModel('gpt-4o', provider=OpenAIProvider(openai_client=mock_client))
627+
agent = Agent(m, output_type=MyTypedDict)
628+
629+
async with agent.run_stream('test') as result:
630+
assert not result.is_complete
631+
structured_chunks = [dict(c) async for c in result.stream(debounce_by=None)]
632+
assert structured_chunks == [{'first': 'Test'}, {'first': 'Test'}]
633+
assert result.is_complete
634+
635+
# Get the final messages and check the ModelResponse
636+
messages = result.all_messages()
637+
model_response = messages[1] # Second message should be the model response
638+
assert isinstance(model_response, ModelResponse)
639+
assert model_response.id == 'struct-response-456'
640+
assert model_response.finish_reason == 'length'
641+
642+
555643
async def test_stream_native_output(allow_model_requests: None):
556644
stream = [
557645
chunk([]),

0 commit comments

Comments
 (0)