Skip to content

Commit 4454583

Browse files
g-eojg-eoj
committed
Add tests
1 parent d1b2399 commit 4454583

File tree

1 file changed

+114
-9
lines changed

1 file changed

+114
-9
lines changed

‎tests/test_agent.py‎

Lines changed: 114 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6132,19 +6132,93 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
61326132
return ModelResponse(parts=[TextPart('')])
61336133

61346134
agent = Agent(FunctionModel(llm))
6135-
output_schema = {'type': 'string'}
6136-
assert agent.output_json_schema == output_schema
6135+
assert agent.output_json_schema == snapshot({'type': 'string'})
61376136

61386137

61396138
async def test_tool_output_json_schema():
61406139
def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
61416140
return ModelResponse(parts=[TextPart('')])
61426141

6143-
agent = Agent(FunctionModel(llm), output_type=[ToolOutput(bool)])
6144-
assert agent.output_json_schema
6142+
agent = Agent(
6143+
FunctionModel(llm),
6144+
output_type=[ToolOutput(bool, name='alice', description='Dreaming...')],
6145+
)
6146+
assert agent.output_json_schema == snapshot(
6147+
{
6148+
'type': 'object',
6149+
'properties': {
6150+
'result': {
6151+
'anyOf': [
6152+
{
6153+
'type': 'object',
6154+
'properties': {
6155+
'kind': {'type': 'string', 'const': 'alice'},
6156+
'data': {
6157+
'properties': {'response': {'type': 'boolean'}},
6158+
'required': ['response'],
6159+
'type': 'object',
6160+
},
6161+
},
6162+
'required': ['kind', 'data'],
6163+
'additionalProperties': False,
6164+
'title': 'alice',
6165+
'description': 'Dreaming...',
6166+
}
6167+
]
6168+
}
6169+
},
6170+
'required': ['result'],
6171+
'additionalProperties': False,
6172+
}
6173+
)
61456174

6146-
agent = Agent(FunctionModel(llm), output_type=[ToolOutput(bool), ToolOutput(bool)])
6147-
assert agent.output_json_schema
6175+
agent = Agent(
6176+
FunctionModel(llm),
6177+
output_type=[ToolOutput(bool, name='alice'), ToolOutput(bool, name='bob')],
6178+
)
6179+
assert agent.output_json_schema == snapshot(
6180+
{
6181+
'type': 'object',
6182+
'properties': {
6183+
'result': {
6184+
'anyOf': [
6185+
{
6186+
'type': 'object',
6187+
'properties': {
6188+
'kind': {'type': 'string', 'const': 'alice'},
6189+
'data': {
6190+
'properties': {'response': {'type': 'boolean'}},
6191+
'required': ['response'],
6192+
'type': 'object',
6193+
},
6194+
},
6195+
'required': ['kind', 'data'],
6196+
'additionalProperties': False,
6197+
'title': 'alice',
6198+
'description': 'bool: The final response which ends this conversation',
6199+
},
6200+
{
6201+
'type': 'object',
6202+
'properties': {
6203+
'kind': {'type': 'string', 'const': 'bob'},
6204+
'data': {
6205+
'properties': {'response': {'type': 'boolean'}},
6206+
'required': ['response'],
6207+
'type': 'object',
6208+
},
6209+
},
6210+
'required': ['kind', 'data'],
6211+
'additionalProperties': False,
6212+
'title': 'bob',
6213+
'description': 'bool: The final response which ends this conversation',
6214+
},
6215+
]
6216+
}
6217+
},
6218+
'required': ['result'],
6219+
'additionalProperties': False,
6220+
}
6221+
)
61486222

61496223

61506224
async def test_native_output_json_schema():
@@ -6155,7 +6229,9 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
61556229
FunctionModel(llm),
61566230
output_type=NativeOutput([bool], name='native_output_name', description='native_output_description'),
61576231
)
6158-
assert agent.output_json_schema
6232+
assert agent.output_json_schema == snapshot(
6233+
{'properties': {'response': {'type': 'boolean'}}, 'required': ['response'], 'type': 'object'}
6234+
)
61596235

61606236

61616237
async def test_prompted_output_json_schema():
@@ -6166,7 +6242,9 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
61666242
FunctionModel(llm),
61676243
output_type=PromptedOutput([bool], name='prompted_output_name', description='prompted_output_description'),
61686244
)
6169-
assert agent.output_json_schema
6245+
assert agent.output_json_schema == snapshot(
6246+
{'properties': {'response': {'type': 'boolean'}}, 'required': ['response'], 'type': 'object'}
6247+
)
61706248

61716249

61726250
async def test_custom_output_json_schema():
@@ -6183,4 +6261,31 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
61836261
description='A human with a name and age',
61846262
)
61856263
agent = Agent(FunctionModel(llm), output_type=HumanDict)
6186-
assert agent.output_json_schema
6264+
assert agent.output_json_schema == snapshot(
6265+
{
6266+
'type': 'object',
6267+
'properties': {
6268+
'result': {
6269+
'anyOf': [
6270+
{
6271+
'type': 'object',
6272+
'properties': {
6273+
'kind': {'type': 'string', 'const': 'final_result'},
6274+
'data': {
6275+
'properties': {'name': {'type': 'string'}, 'age': {'type': 'integer'}},
6276+
'required': ['name', 'age'],
6277+
'type': 'object',
6278+
},
6279+
},
6280+
'required': ['kind', 'data'],
6281+
'additionalProperties': False,
6282+
'title': 'final_result',
6283+
'description': 'A human with a name and age',
6284+
}
6285+
]
6286+
}
6287+
},
6288+
'required': ['result'],
6289+
'additionalProperties': False,
6290+
}
6291+
)

0 commit comments

Comments
 (0)