Skip to content

Commit e603f68

Browse files
alexmojakialexmojaki
authored
Replace all_messages_events with pydantic_ai.all_messages for InstrumentationSettings(version=2) (#2658)
1 parent 057c0bc commit e603f68

File tree

2 files changed

+228
-117
lines changed

2 files changed

+228
-117
lines changed

‎pydantic_ai_slim/pydantic_ai/agent/__init__.py‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -677,16 +677,23 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
677677
def _run_span_end_attributes(
678678
self, state: _agent_graph.GraphAgentState, usage: _usage.RunUsage, settings: InstrumentationSettings
679679
):
680+
if settings.version == 1:
681+
attr_name = 'all_messages_events'
682+
value = [
683+
InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(state.message_history)
684+
]
685+
else:
686+
attr_name = 'pydantic_ai.all_messages'
687+
value = settings.messages_to_otel_messages(state.message_history)
688+
680689
return {
681690
**usage.opentelemetry_attributes(),
682-
'all_messages_events': json.dumps(
683-
[InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(state.message_history)]
684-
),
691+
attr_name: json.dumps(value),
685692
'logfire.json_schema': json.dumps(
686693
{
687694
'type': 'object',
688695
'properties': {
689-
'all_messages_events': {'type': 'array'},
696+
attr_name: {'type': 'array'},
690697
'final_result': {'type': 'object'},
691698
},
692699
}

‎tests/test_logfire.py‎

Lines changed: 217 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -116,66 +116,114 @@ async def my_ret(x: int) -> str:
116116
}
117117
]
118118
)
119-
assert summary.attributes[0] == snapshot(
120-
{
121-
'model_name': 'test',
122-
'agent_name': 'my_agent',
123-
'logfire.msg': 'my_agent run',
124-
'logfire.span_type': 'span',
125-
'gen_ai.usage.input_tokens': 103,
126-
'gen_ai.usage.output_tokens': 12,
127-
'all_messages_events': IsJson(
128-
snapshot(
129-
[
130-
{
131-
'content': 'Hello',
132-
'role': 'user',
133-
'gen_ai.message.index': 0,
134-
'event.name': 'gen_ai.user.message',
135-
},
136-
{
137-
'role': 'assistant',
138-
'tool_calls': [
139-
{
140-
'id': IsStr(),
141-
'type': 'function',
142-
'function': {
119+
120+
if isinstance(instrument, InstrumentationSettings) and instrument.version == 2:
121+
assert summary.attributes[0] == snapshot(
122+
{
123+
'model_name': 'test',
124+
'agent_name': 'my_agent',
125+
'logfire.msg': 'my_agent run',
126+
'logfire.span_type': 'span',
127+
'final_result': '{"my_ret":"1"}',
128+
'gen_ai.usage.input_tokens': 103,
129+
'gen_ai.usage.output_tokens': 12,
130+
'pydantic_ai.all_messages': IsJson(
131+
snapshot(
132+
[
133+
{'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
134+
{
135+
'role': 'assistant',
136+
'parts': [
137+
{
138+
'type': 'tool_call',
139+
'id': IsStr(),
143140
'name': 'my_ret',
144141
'arguments': {'x': 0},
145-
},
146-
}
147-
],
148-
'gen_ai.message.index': 1,
149-
'event.name': 'gen_ai.assistant.message',
150-
},
151-
{
152-
'content': '1',
153-
'role': 'tool',
154-
'id': IsStr(),
155-
'name': 'my_ret',
156-
'gen_ai.message.index': 2,
157-
'event.name': 'gen_ai.tool.message',
158-
},
142+
}
143+
],
144+
},
145+
{
146+
'role': 'user',
147+
'parts': [
148+
{
149+
'type': 'tool_call_response',
150+
'id': IsStr(),
151+
'name': 'my_ret',
152+
'result': '1',
153+
}
154+
],
155+
},
156+
{'role': 'assistant', 'parts': [{'type': 'text', 'content': '{"my_ret":"1"}'}]},
157+
]
158+
)
159+
),
160+
'logfire.json_schema': '{"type": "object", "properties": {"pydantic_ai.all_messages": {"type": "array"}, "final_result": {"type": "object"}}}',
161+
}
162+
)
163+
else:
164+
assert summary.attributes[0] == snapshot(
165+
{
166+
'model_name': 'test',
167+
'agent_name': 'my_agent',
168+
'logfire.msg': 'my_agent run',
169+
'logfire.span_type': 'span',
170+
'gen_ai.usage.input_tokens': 103,
171+
'gen_ai.usage.output_tokens': 12,
172+
'all_messages_events': IsJson(
173+
snapshot(
174+
[
175+
{
176+
'content': 'Hello',
177+
'role': 'user',
178+
'gen_ai.message.index': 0,
179+
'event.name': 'gen_ai.user.message',
180+
},
181+
{
182+
'role': 'assistant',
183+
'tool_calls': [
184+
{
185+
'id': IsStr(),
186+
'type': 'function',
187+
'function': {
188+
'name': 'my_ret',
189+
'arguments': {'x': 0},
190+
},
191+
}
192+
],
193+
'gen_ai.message.index': 1,
194+
'event.name': 'gen_ai.assistant.message',
195+
},
196+
{
197+
'content': '1',
198+
'role': 'tool',
199+
'id': IsStr(),
200+
'name': 'my_ret',
201+
'gen_ai.message.index': 2,
202+
'event.name': 'gen_ai.tool.message',
203+
},
204+
{
205+
'role': 'assistant',
206+
'content': '{"my_ret":"1"}',
207+
'gen_ai.message.index': 3,
208+
'event.name': 'gen_ai.assistant.message',
209+
},
210+
]
211+
)
212+
),
213+
'final_result': '{"my_ret":"1"}',
214+
'logfire.json_schema': IsJson(
215+
snapshot(
159216
{
160-
'role': 'assistant',
161-
'content': '{"my_ret":"1"}',
162-
'gen_ai.message.index': 3,
163-
'event.name': 'gen_ai.assistant.message',
164-
},
165-
]
166-
)
167-
),
168-
'final_result': '{"my_ret":"1"}',
169-
'logfire.json_schema': IsJson(
170-
snapshot(
171-
{
172-
'type': 'object',
173-
'properties': {'all_messages_events': {'type': 'array'}, 'final_result': {'type': 'object'}},
174-
}
175-
)
176-
),
177-
}
178-
)
217+
'type': 'object',
218+
'properties': {
219+
'all_messages_events': {'type': 'array'},
220+
'final_result': {'type': 'object'},
221+
},
222+
}
223+
)
224+
),
225+
}
226+
)
179227
chat_span_attributes = summary.attributes[1]
180228
if instrument is True or instrument.event_mode == 'attributes':
181229
if hasattr(capfire, 'get_collected_metrics'): # pragma: no branch
@@ -368,64 +416,68 @@ class MyOutput:
368416
assert result.output == MyOutput(content='a')
369417

370418
summary = get_logfire_summary()
371-
assert summary.attributes[0] == snapshot(
372-
{
373-
'model_name': 'test',
374-
'agent_name': 'my_agent',
375-
'logfire.msg': 'my_agent run',
376-
'logfire.span_type': 'span',
377-
'gen_ai.usage.input_tokens': 51,
378-
'gen_ai.usage.output_tokens': 5,
379-
'all_messages_events': IsJson(
380-
snapshot(
381-
[
382-
{
383-
'content': 'Here are some instructions',
384-
'role': 'system',
385-
'event.name': 'gen_ai.system.message',
386-
},
387-
{
388-
'content': 'Hello',
389-
'role': 'user',
390-
'gen_ai.message.index': 0,
391-
'event.name': 'gen_ai.user.message',
392-
},
393-
{
394-
'role': 'assistant',
395-
'tool_calls': [
396-
{
397-
'id': IsStr(),
398-
'type': 'function',
399-
'function': {'name': 'final_result', 'arguments': {'content': 'a'}},
400-
}
401-
],
402-
'gen_ai.message.index': 1,
403-
'event.name': 'gen_ai.assistant.message',
404-
},
405-
{
406-
'content': 'Final result processed.',
407-
'role': 'tool',
408-
'id': IsStr(),
409-
'name': 'final_result',
410-
'gen_ai.message.index': 2,
411-
'event.name': 'gen_ai.tool.message',
412-
},
413-
]
414-
)
415-
),
416-
'final_result': '{"content": "a"}',
417-
'logfire.json_schema': IsJson(
418-
snapshot(
419-
{
420-
'type': 'object',
421-
'properties': {'all_messages_events': {'type': 'array'}, 'final_result': {'type': 'object'}},
422-
}
423-
)
424-
),
425-
}
426-
)
427419
chat_span_attributes = summary.attributes[1]
428420
if instrument.version == 1:
421+
assert summary.attributes[0] == snapshot(
422+
{
423+
'model_name': 'test',
424+
'agent_name': 'my_agent',
425+
'logfire.msg': 'my_agent run',
426+
'logfire.span_type': 'span',
427+
'gen_ai.usage.input_tokens': 51,
428+
'gen_ai.usage.output_tokens': 5,
429+
'all_messages_events': IsJson(
430+
snapshot(
431+
[
432+
{
433+
'content': 'Here are some instructions',
434+
'role': 'system',
435+
'event.name': 'gen_ai.system.message',
436+
},
437+
{
438+
'content': 'Hello',
439+
'role': 'user',
440+
'gen_ai.message.index': 0,
441+
'event.name': 'gen_ai.user.message',
442+
},
443+
{
444+
'role': 'assistant',
445+
'tool_calls': [
446+
{
447+
'id': IsStr(),
448+
'type': 'function',
449+
'function': {'name': 'final_result', 'arguments': {'content': 'a'}},
450+
}
451+
],
452+
'gen_ai.message.index': 1,
453+
'event.name': 'gen_ai.assistant.message',
454+
},
455+
{
456+
'content': 'Final result processed.',
457+
'role': 'tool',
458+
'id': IsStr(),
459+
'name': 'final_result',
460+
'gen_ai.message.index': 2,
461+
'event.name': 'gen_ai.tool.message',
462+
},
463+
]
464+
)
465+
),
466+
'final_result': '{"content": "a"}',
467+
'logfire.json_schema': IsJson(
468+
snapshot(
469+
{
470+
'type': 'object',
471+
'properties': {
472+
'all_messages_events': {'type': 'array'},
473+
'final_result': {'type': 'object'},
474+
},
475+
}
476+
)
477+
),
478+
}
479+
)
480+
429481
assert chat_span_attributes['events'] == IsJson(
430482
snapshot(
431483
[
@@ -461,6 +513,58 @@ class MyOutput:
461513
)
462514
)
463515
else:
516+
assert summary.attributes[0] == snapshot(
517+
{
518+
'model_name': 'test',
519+
'agent_name': 'my_agent',
520+
'logfire.msg': 'my_agent run',
521+
'logfire.span_type': 'span',
522+
'gen_ai.usage.input_tokens': 51,
523+
'gen_ai.usage.output_tokens': 5,
524+
'pydantic_ai.all_messages': IsJson(
525+
snapshot(
526+
[
527+
{'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
528+
{
529+
'role': 'assistant',
530+
'parts': [
531+
{
532+
'type': 'tool_call',
533+
'id': IsStr(),
534+
'name': 'final_result',
535+
'arguments': {'content': 'a'},
536+
}
537+
],
538+
},
539+
{
540+
'role': 'user',
541+
'parts': [
542+
{
543+
'type': 'tool_call_response',
544+
'id': IsStr(),
545+
'name': 'final_result',
546+
'result': 'Final result processed.',
547+
}
548+
],
549+
},
550+
]
551+
)
552+
),
553+
'final_result': '{"content": "a"}',
554+
'logfire.json_schema': IsJson(
555+
snapshot(
556+
{
557+
'type': 'object',
558+
'properties': {
559+
'pydantic_ai.all_messages': {'type': 'array'},
560+
'final_result': {'type': 'object'},
561+
},
562+
}
563+
)
564+
),
565+
}
566+
)
567+
464568
assert chat_span_attributes['gen_ai.input.messages'] == IsJson(
465569
snapshot([{'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}])
466570
)

0 commit comments

Comments
 (0)