Skip to content

Commit fb51a5f

Browse files
alexmojakidmontagu
alexmojaki
and
dmontagu
authored
Replace all_messages in agent span with all_messages_events in same format as InstrumentedModel span (#1018)
Co-authored-by: David Montague <[email protected]>
1 parent b63a418 commit fb51a5f

File tree

2 files changed

+59
-125
lines changed

2 files changed

+59
-125
lines changed

‎pydantic_ai_slim/pydantic_ai/_agent_graph.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
result,
2424
usage as _usage,
2525
)
26+
from .models.instrumented import InstrumentedModel
2627
from .result import ResultDataT
2728
from .settings import ModelSettings, merge_model_settings
2829
from .tools import (
@@ -494,7 +495,10 @@ def _handle_final_result(
494495
messages.append(_messages.ModelRequest(parts=tool_responses))
495496

496497
run_span.set_attribute('usage', usage)
497-
run_span.set_attribute('all_messages', messages)
498+
run_span.set_attribute(
499+
'all_messages_events',
500+
[InstrumentedModel.event_to_dict(e) for e in InstrumentedModel.messages_to_otel_events(messages)],
501+
)
498502

499503
# End the run with self.data
500504
return End(final_result)

‎tests/test_logfire.py‎

Lines changed: 54 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from typing import Any, Callable
55

66
import pytest
7-
from dirty_equals import IsInt, IsJson, IsStr
7+
from dirty_equals import IsInt, IsJson
88
from inline_snapshot import snapshot
99
from typing_extensions import NotRequired, TypedDict
1010

@@ -114,136 +114,66 @@ async def my_ret(x: int) -> str:
114114
'logfire.msg_template': '{agent_name} run {prompt=}',
115115
'logfire.msg': 'my_agent run prompt=Hello',
116116
'logfire.span_type': 'span',
117-
'all_messages': IsJson(
118-
[
119-
{
120-
'parts': [
121-
{
122-
'content': 'Hello',
123-
'timestamp': IsStr(regex=r'\d{4}-\d{2}-.+'),
124-
'part_kind': 'user-prompt',
125-
},
126-
],
127-
'kind': 'request',
128-
},
129-
{
130-
'parts': [
131-
{'tool_name': 'my_ret', 'args': {'x': 0}, 'tool_call_id': None, 'part_kind': 'tool-call'}
132-
],
133-
'model_name': 'test',
134-
'timestamp': IsStr(regex=r'\d{4}-\d{2}-.+'),
135-
'kind': 'response',
136-
},
137-
{
138-
'parts': [
139-
{
140-
'tool_name': 'my_ret',
141-
'content': '1',
142-
'tool_call_id': None,
143-
'timestamp': IsStr(regex=r'\d{4}-\d{2}-.+'),
144-
'part_kind': 'tool-return',
145-
},
146-
],
147-
'kind': 'request',
148-
},
149-
{
150-
'parts': [{'content': '{"my_ret":"1"}', 'part_kind': 'text'}],
151-
'model_name': 'test',
152-
'timestamp': IsStr(regex=r'\d{4}-\d{2}-.+'),
153-
'kind': 'response',
154-
},
155-
]
117+
'all_messages_events': IsJson(
118+
snapshot(
119+
[
120+
{
121+
'content': 'Hello',
122+
'role': 'user',
123+
'event.name': 'gen_ai.user.message',
124+
},
125+
{
126+
'role': 'assistant',
127+
'tool_calls': [
128+
{
129+
'id': None,
130+
'type': 'function',
131+
'function': {
132+
'name': 'my_ret',
133+
'arguments': {'x': 0},
134+
},
135+
}
136+
],
137+
'event.name': 'gen_ai.assistant.message',
138+
},
139+
{
140+
'content': '1',
141+
'role': 'tool',
142+
'id': None,
143+
'event.name': 'gen_ai.tool.message',
144+
},
145+
{
146+
'role': 'assistant',
147+
'content': '{"my_ret":"1"}',
148+
'event.name': 'gen_ai.assistant.message',
149+
},
150+
]
151+
)
156152
),
157153
'usage': IsJson(
158154
{'requests': 2, 'request_tokens': 103, 'response_tokens': 12, 'total_tokens': 115, 'details': None}
159155
),
160156
'logfire.json_schema': IsJson(
161-
{
162-
'type': 'object',
163-
'properties': {
164-
'prompt': {},
165-
'agent': {
166-
'type': 'object',
167-
'title': 'Agent',
168-
'x-python-datatype': 'dataclass',
169-
'properties': {
170-
'model': {'type': 'object', 'title': 'TestModel', 'x-python-datatype': 'dataclass'}
171-
},
172-
},
173-
'model_name': {},
174-
'agent_name': {},
175-
'all_messages': {
176-
'type': 'array',
177-
'prefixItems': [
178-
{
179-
'type': 'object',
180-
'title': 'ModelRequest',
181-
'x-python-datatype': 'dataclass',
182-
'properties': {
183-
'parts': {
184-
'type': 'array',
185-
'items': {
186-
'type': 'object',
187-
'title': 'UserPromptPart',
188-
'x-python-datatype': 'dataclass',
189-
'properties': {'timestamp': {'type': 'string', 'format': 'date-time'}},
190-
},
191-
}
192-
},
193-
},
194-
{
195-
'type': 'object',
196-
'title': 'ModelResponse',
197-
'x-python-datatype': 'dataclass',
198-
'properties': {
199-
'parts': {
200-
'type': 'array',
201-
'items': {
202-
'type': 'object',
203-
'title': 'ToolCallPart',
204-
'x-python-datatype': 'dataclass',
205-
},
206-
},
207-
'timestamp': {'type': 'string', 'format': 'date-time'},
208-
},
209-
},
210-
{
211-
'type': 'object',
212-
'title': 'ModelRequest',
213-
'x-python-datatype': 'dataclass',
214-
'properties': {
215-
'parts': {
216-
'type': 'array',
217-
'items': {
218-
'type': 'object',
219-
'title': 'ToolReturnPart',
220-
'x-python-datatype': 'dataclass',
221-
'properties': {'timestamp': {'type': 'string', 'format': 'date-time'}},
222-
},
223-
}
224-
},
225-
},
226-
{
227-
'type': 'object',
228-
'title': 'ModelResponse',
229-
'x-python-datatype': 'dataclass',
230-
'properties': {
231-
'parts': {
232-
'type': 'array',
233-
'items': {
234-
'type': 'object',
235-
'title': 'TextPart',
236-
'x-python-datatype': 'dataclass',
237-
},
238-
},
239-
'timestamp': {'type': 'string', 'format': 'date-time'},
240-
},
157+
snapshot(
158+
{
159+
'type': 'object',
160+
'properties': {
161+
'prompt': {},
162+
'agent': {
163+
'type': 'object',
164+
'title': 'Agent',
165+
'x-python-datatype': 'dataclass',
166+
'properties': {
167+
'model': {'type': 'object', 'title': 'TestModel', 'x-python-datatype': 'dataclass'}
241168
},
242-
],
169+
},
170+
'model_name': {},
171+
'agent_name': {},
172+
'usage': {'type': 'object', 'title': 'Usage', 'x-python-datatype': 'dataclass'},
173+
'all_messages_events': {'type': 'array'},
243174
},
244-
'usage': {'type': 'object', 'title': 'Usage', 'x-python-datatype': 'dataclass'},
245-
},
246-
}
175+
}
176+
)
247177
),
248178
}
249179
)

0 commit comments

Comments
 (0)