Coverage for tests/test_tools.py: 99.42%

1import json (empty)

2from dataclasses import dataclass (empty)

3from typing import Annotated, Any, Callable, Literal, Union (empty)

4 

5import pydantic_core (empty)

6import pytest (empty)

7from _pytest.logging import LogCaptureFixture (empty)

8from inline_snapshot import snapshot (empty)

9from pydantic import BaseModel, Field, WithJsonSchema (empty)

10from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue (empty)

11from pydantic_core import PydanticSerializationError, core_schema (empty)

12 

13from pydantic_ai import Agent, RunContext, Tool, UserError (empty)

14from pydantic_ai.messages import ( (empty)

15 ModelMessage, 

16 ModelRequest, 

17 ModelResponse, 

18 TextPart, 

19 ToolCallPart, 

20 ToolReturnPart, 

21) 

22from pydantic_ai.models.function import AgentInfo, FunctionModel (empty)

23from pydantic_ai.models.test import TestModel (empty)

24from pydantic_ai.tools import ToolDefinition (empty)

25 

26 

27def test_tool_no_ctx(): (empty)

28 agent = Agent(TestModel()) (empty)

29 

30 with pytest.raises(UserError) as exc_info: (empty)

31 

32 @agent.tool # pyright: ignore[reportArgumentType] (empty)

33 def invalid_tool(x: int) -> str: # pragma: no cover (empty)

34 return 'Hello' 

35 

36 assert str(exc_info.value) == snapshot( (empty)

37 'Error generating schema for test_tool_no_ctx.<locals>.invalid_tool:\n' 

38 ' First parameter of tools that take context must be annotated with RunContext[...]' 

39 ) 

40 

41 

42def test_tool_plain_with_ctx(): (empty)

43 agent = Agent(TestModel()) (empty)

44 

45 with pytest.raises(UserError) as exc_info: (empty)

46 

47 @agent.tool_plain (empty)

48 async def invalid_tool(ctx: RunContext[None]) -> str: # pragma: no cover (empty)

49 return 'Hello' 

50 

51 assert str(exc_info.value) == snapshot( (empty)

52 'Error generating schema for test_tool_plain_with_ctx.<locals>.invalid_tool:\n' 

53 ' RunContext annotations can only be used with tools that take context' 

54 ) 

55 

56 

57def test_tool_ctx_second(): (empty)

58 agent = Agent(TestModel()) (empty)

59 

60 with pytest.raises(UserError) as exc_info: (empty)

61 

62 @agent.tool # pyright: ignore[reportArgumentType] (empty)

63 def invalid_tool(x: int, ctx: RunContext[None]) -> str: # pragma: no cover (empty)

64 return 'Hello' 

65 

66 assert str(exc_info.value) == snapshot( (empty)

67 'Error generating schema for test_tool_ctx_second.<locals>.invalid_tool:\n' 

68 ' First parameter of tools that take context must be annotated with RunContext[...]\n' 

69 ' RunContext annotations can only be used as the first argument' 

70 ) 

71 

72 

73async def google_style_docstring(foo: int, bar: str) -> str: # pragma: no cover (empty)

74 """Do foobar stuff, a lot. 

75 

76 Args: 

77 foo: The foo thing. 

78 bar: The bar thing. 

79 """ 

80 return f'{foo} {bar}' 

81 

82 

83async def get_json_schema(_messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: (empty)

84 if len(info.function_tools) == 1: (empty)

85 r = info.function_tools[0] (empty)

86 return ModelResponse(parts=[TextPart(pydantic_core.to_json(r).decode())]) (empty)

87 else: 

88 return ModelResponse(parts=[TextPart(pydantic_core.to_json(info.function_tools).decode())]) (empty)

89 

90 

91@pytest.mark.parametrize('docstring_format', ['google', 'auto']) (empty)

92def test_docstring_google(docstring_format: Literal['google', 'auto']): (empty)

93 agent = Agent(FunctionModel(get_json_schema)) (empty)

94 agent.tool_plain(docstring_format=docstring_format)(google_style_docstring) (empty)

95 

96 result = agent.run_sync('Hello') (empty)

97 json_schema = json.loads(result.data) (empty)

98 assert json_schema == snapshot( (empty)

99 { 

100 'name': 'google_style_docstring', 

101 'description': 'Do foobar stuff, a lot.', 

102 'parameters_json_schema': { 

103 'properties': { 

104 'foo': {'description': 'The foo thing.', 'type': 'integer'}, 

105 'bar': {'description': 'The bar thing.', 'type': 'string'}, 

106 }, 

107 'required': ['foo', 'bar'], 

108 'type': 'object', 

109 'additionalProperties': False, 

110 }, 

111 'outer_typed_dict_key': None, 

112 } 

113 ) 

114 keys = list(json_schema.keys()) (empty)

115 # name should be the first key 

116 assert keys[0] == 'name' (empty)

117 # description should be the second key 

118 assert keys[1] == 'description' (empty)

119 

120 

121def sphinx_style_docstring(foo: int, /) -> str: # pragma: no cover (empty)

122 """Sphinx style docstring. 

123 

124 :param foo: The foo thing. 

125 """ 

126 return str(foo) 

127 

128 

129@pytest.mark.parametrize('docstring_format', ['sphinx', 'auto']) (empty)

130def test_docstring_sphinx(docstring_format: Literal['sphinx', 'auto']): (empty)

131 agent = Agent(FunctionModel(get_json_schema)) (empty)

132 agent.tool_plain(docstring_format=docstring_format)(sphinx_style_docstring) (empty)

133 

134 result = agent.run_sync('Hello') (empty)

135 json_schema = json.loads(result.data) (empty)

136 assert json_schema == snapshot( (empty)

137 { 

138 'name': 'sphinx_style_docstring', 

139 'description': 'Sphinx style docstring.', 

140 'parameters_json_schema': { 

141 'properties': {'foo': {'description': 'The foo thing.', 'type': 'integer'}}, 

142 'required': ['foo'], 

143 'type': 'object', 

144 'additionalProperties': False, 

145 }, 

146 'outer_typed_dict_key': None, 

147 } 

148 ) 

149 

150 

151def numpy_style_docstring(*, foo: int, bar: str) -> str: # pragma: no cover (empty)

152 """Numpy style docstring. 

153 

154 Parameters 

155 ---------- 

156 foo : int 

157 The foo thing. 

158 bar : str 

159 The bar thing. 

160 """ 

161 return f'{foo} {bar}' 

162 

163 

164@pytest.mark.parametrize('docstring_format', ['numpy', 'auto']) (empty)

165def test_docstring_numpy(docstring_format: Literal['numpy', 'auto']): (empty)

166 agent = Agent(FunctionModel(get_json_schema)) (empty)

167 agent.tool_plain(docstring_format=docstring_format)(numpy_style_docstring) (empty)

168 

169 result = agent.run_sync('Hello') (empty)

170 json_schema = json.loads(result.data) (empty)

171 assert json_schema == snapshot( (empty)

172 { 

173 'name': 'numpy_style_docstring', 

174 'description': 'Numpy style docstring.', 

175 'parameters_json_schema': { 

176 'properties': { 

177 'foo': {'description': 'The foo thing.', 'type': 'integer'}, 

178 'bar': {'description': 'The bar thing.', 'type': 'string'}, 

179 }, 

180 'required': ['foo', 'bar'], 

181 'type': 'object', 

182 'additionalProperties': False, 

183 }, 

184 'outer_typed_dict_key': None, 

185 } 

186 ) 

187 

188 

189def test_google_style_with_returns(): (empty)

190 agent = Agent(FunctionModel(get_json_schema)) (empty)

191 

192 def my_tool(x: int) -> str: # pragma: no cover (empty)

193 """A function that does something. 

194 

195 Args: 

196 x: The input value. 

197 

198 Returns: 

199 str: The result as a string. 

200 """ 

201 return str(x) 

202 

203 agent.tool_plain(my_tool) (empty)

204 result = agent.run_sync('Hello') (empty)

205 json_schema = json.loads(result.data) (empty)

206 assert json_schema == snapshot( (empty)

207 { 

208 'name': 'my_tool', 

209 'description': """\ 

210<summary>A function that does something.</summary> 

211<returns> 

212<type>str</type> 

213<description>The result as a string.</description> 

214</returns>\ 

215""", 

216 'parameters_json_schema': { 

217 'additionalProperties': False, 

218 'properties': {'x': {'description': 'The input value.', 'type': 'integer'}}, 

219 'required': ['x'], 

220 'type': 'object', 

221 }, 

222 'outer_typed_dict_key': None, 

223 } 

224 ) 

225 

226 

227def test_sphinx_style_with_returns(): (empty)

228 agent = Agent(FunctionModel(get_json_schema)) (empty)

229 

230 def my_tool(x: int) -> str: # pragma: no cover (empty)

231 """A sphinx function with returns. 

232 

233 :param x: The input value. 

234 :rtype: str 

235 :return: The result as a string with type. 

236 """ 

237 return str(x) 

238 

239 agent.tool_plain(docstring_format='sphinx')(my_tool) (empty)

240 result = agent.run_sync('Hello') (empty)

241 json_schema = json.loads(result.data) (empty)

242 assert json_schema == snapshot( (empty)

243 { 

244 'name': 'my_tool', 

245 'description': """\ 

246<summary>A sphinx function with returns.</summary> 

247<returns> 

248<type>str</type> 

249<description>The result as a string with type.</description> 

250</returns>\ 

251""", 

252 'parameters_json_schema': { 

253 'additionalProperties': False, 

254 'properties': {'x': {'description': 'The input value.', 'type': 'integer'}}, 

255 'required': ['x'], 

256 'type': 'object', 

257 }, 

258 'outer_typed_dict_key': None, 

259 } 

260 ) 

261 

262 

263def test_numpy_style_with_returns(): (empty)

264 agent = Agent(FunctionModel(get_json_schema)) (empty)

265 

266 def my_tool(x: int) -> str: # pragma: no cover (empty)

267 """A numpy function with returns. 

268 

269 Parameters 

270 ---------- 

271 x : int 

272 The input value. 

273 

274 Returns 

275 ------- 

276 str 

277 The result as a string with type. 

278 """ 

279 return str(x) 

280 

281 agent.tool_plain(docstring_format='numpy')(my_tool) (empty)

282 result = agent.run_sync('Hello') (empty)

283 json_schema = json.loads(result.data) (empty)

284 assert json_schema == snapshot( (empty)

285 { 

286 'name': 'my_tool', 

287 'description': """\ 

288<summary>A numpy function with returns.</summary> 

289<returns> 

290<type>str</type> 

291<description>The result as a string with type.</description> 

292</returns>\ 

293""", 

294 'parameters_json_schema': { 

295 'additionalProperties': False, 

296 'properties': {'x': {'description': 'The input value.', 'type': 'integer'}}, 

297 'required': ['x'], 

298 'type': 'object', 

299 }, 

300 'outer_typed_dict_key': None, 

301 } 

302 ) 

303 

304 

305def only_returns_type() -> str: # pragma: no cover (empty)

306 """ 

307 

308 Returns: 

309 str: The result as a string. 

310 """ 

311 return 'foo' 

312 

313 

314def test_only_returns_type(): (empty)

315 agent = Agent(FunctionModel(get_json_schema)) (empty)

316 agent.tool_plain(only_returns_type) (empty)

317 

318 result = agent.run_sync('Hello') (empty)

319 json_schema = json.loads(result.data) (empty)

320 assert json_schema == snapshot( (empty)

321 { 

322 'name': 'only_returns_type', 

323 'description': """\ 

324<returns> 

325<type>str</type> 

326<description>The result as a string.</description> 

327</returns>\ 

328""", 

329 'parameters_json_schema': {'additionalProperties': False, 'properties': {}, 'type': 'object'}, 

330 'outer_typed_dict_key': None, 

331 } 

332 ) 

333 

334 

335def unknown_docstring(**kwargs: int) -> str: # pragma: no cover (empty)

336 """Unknown style docstring.""" 

337 return str(kwargs) 

338 

339 

340def test_docstring_unknown(): (empty)

341 agent = Agent(FunctionModel(get_json_schema)) (empty)

342 agent.tool_plain(unknown_docstring) (empty)

343 

344 result = agent.run_sync('Hello') (empty)

345 json_schema = json.loads(result.data) (empty)

346 assert json_schema == snapshot( (empty)

347 { 

348 'name': 'unknown_docstring', 

349 'description': 'Unknown style docstring.', 

350 'parameters_json_schema': {'properties': {}, 'type': 'object', 'additionalProperties': True}, 

351 'outer_typed_dict_key': None, 

352 } 

353 ) 

354 

355 

356# fmt: off 

357async def google_style_docstring_no_body( (empty)

358 foo: int, bar: Annotated[str, Field(description='from fields')] (empty)

359) -> str: # pragma: no cover (empty)

360 """ 

361 Args: 

362 foo: The foo thing. 

363 bar: The bar thing. 

364 """ 

365 

366 return f'{foo} {bar}' 

367# fmt: on 

368 

369 

370@pytest.mark.parametrize('docstring_format', ['google', 'auto']) (empty)

371def test_docstring_google_no_body(docstring_format: Literal['google', 'auto']): (empty)

372 agent = Agent(FunctionModel(get_json_schema)) (empty)

373 agent.tool_plain(docstring_format=docstring_format)(google_style_docstring_no_body) (empty)

374 

375 result = agent.run_sync('') (empty)

376 json_schema = json.loads(result.data) (empty)

377 assert json_schema == snapshot( (empty)

378 { 

379 'name': 'google_style_docstring_no_body', 

380 'description': '', 

381 'parameters_json_schema': { 

382 'properties': { 

383 'foo': {'description': 'The foo thing.', 'type': 'integer'}, 

384 'bar': {'description': 'from fields', 'type': 'string'}, 

385 }, 

386 'required': ['foo', 'bar'], 

387 'type': 'object', 

388 'additionalProperties': False, 

389 }, 

390 'outer_typed_dict_key': None, 

391 } 

392 ) 

393 

394 

395class Foo(BaseModel): (empty)

396 x: int (empty)

397 y: str (empty)

398 

399 

400def test_takes_just_model(): (empty)

401 agent = Agent() (empty)

402 

403 @agent.tool_plain (empty)

404 def takes_just_model(model: Foo) -> str: (empty)

405 return f'{model.x} {model.y}' (empty)

406 

407 result = agent.run_sync('', model=FunctionModel(get_json_schema)) (empty)

408 json_schema = json.loads(result.data) (empty)

409 assert json_schema == snapshot( (empty)

410 { 

411 'name': 'takes_just_model', 

412 'description': None, 

413 'parameters_json_schema': { 

414 'properties': { 

415 'x': {'type': 'integer'}, 

416 'y': {'type': 'string'}, 

417 }, 

418 'required': ['x', 'y'], 

419 'title': 'Foo', 

420 'type': 'object', 

421 }, 

422 'outer_typed_dict_key': None, 

423 } 

424 ) 

425 

426 result = agent.run_sync('', model=TestModel()) (empty)

427 assert result.data == snapshot('{"takes_just_model":"0 a"}') (empty)

428 

429 

430def test_takes_model_and_int(): (empty)

431 agent = Agent() (empty)

432 

433 @agent.tool_plain (empty)

434 def takes_just_model(model: Foo, z: int) -> str: (empty)

435 return f'{model.x} {model.y} {z}' (empty)

436 

437 result = agent.run_sync('', model=FunctionModel(get_json_schema)) (empty)

438 json_schema = json.loads(result.data) (empty)

439 assert json_schema == snapshot( (empty)

440 { 

441 'name': 'takes_just_model', 

442 'description': '', 

443 'parameters_json_schema': { 

444 '$defs': { 

445 'Foo': { 

446 'properties': { 

447 'x': {'type': 'integer'}, 

448 'y': {'type': 'string'}, 

449 }, 

450 'required': ['x', 'y'], 

451 'title': 'Foo', 

452 'type': 'object', 

453 } 

454 }, 

455 'properties': { 

456 'model': {'$ref': '#/$defs/Foo'}, 

457 'z': {'type': 'integer'}, 

458 }, 

459 'required': ['model', 'z'], 

460 'type': 'object', 

461 'additionalProperties': False, 

462 }, 

463 'outer_typed_dict_key': None, 

464 } 

465 ) 

466 

467 result = agent.run_sync('', model=TestModel()) (empty)

468 assert result.data == snapshot('{"takes_just_model":"0 a 0"}') (empty)

469 

470 

471# pyright: reportPrivateUsage=false 

472def test_init_tool_plain(): (empty)

473 call_args: list[int] = [] (empty)

474 

475 def plain_tool(x: int) -> int: (empty)

476 call_args.append(x) (empty)

477 return x + 1 (empty)

478 

479 agent = Agent('test', tools=[Tool(plain_tool)], retries=7) (empty)

480 result = agent.run_sync('foobar') (empty)

481 assert result.data == snapshot('{"plain_tool":1}') (empty)

482 assert call_args == snapshot([0]) (empty)

483 assert agent._function_tools['plain_tool'].takes_ctx is False (empty)

484 assert agent._function_tools['plain_tool'].max_retries == 7 (empty)

485 

486 agent_infer = Agent('test', tools=[plain_tool], retries=7) (empty)

487 result = agent_infer.run_sync('foobar') (empty)

488 assert result.data == snapshot('{"plain_tool":1}') (empty)

489 assert call_args == snapshot([0, 0]) (empty)

490 assert agent_infer._function_tools['plain_tool'].takes_ctx is False (empty)

491 assert agent_infer._function_tools['plain_tool'].max_retries == 7 (empty)

492 

493 

494def ctx_tool(ctx: RunContext[int], x: int) -> int: (empty)

495 return x + ctx.deps (empty)

496 

497 

498# pyright: reportPrivateUsage=false 

499def test_init_tool_ctx(): (empty)

500 agent = Agent('test', tools=[Tool(ctx_tool, takes_ctx=True, max_retries=3)], deps_type=int, retries=7) (empty)

501 result = agent.run_sync('foobar', deps=5) (empty)

502 assert result.data == snapshot('{"ctx_tool":5}') (empty)

503 assert agent._function_tools['ctx_tool'].takes_ctx is True (empty)

504 assert agent._function_tools['ctx_tool'].max_retries == 3 (empty)

505 

506 agent_infer = Agent('test', tools=[ctx_tool], deps_type=int) (empty)

507 result = agent_infer.run_sync('foobar', deps=6) (empty)

508 assert result.data == snapshot('{"ctx_tool":6}') (empty)

509 assert agent_infer._function_tools['ctx_tool'].takes_ctx is True (empty)

510 

511 

512def test_repeat_tool(): (empty)

513 with pytest.raises(UserError, match="Tool name conflicts with existing tool: 'ctx_tool'"): (empty)

514 Agent('test', tools=[Tool(ctx_tool), ctx_tool], deps_type=int) (empty)

515 

516 

517def test_tool_return_conflict(): (empty)

518 # this is okay 

519 Agent('test', tools=[ctx_tool], deps_type=int) (empty)

520 # this is also okay 

521 Agent('test', tools=[ctx_tool], deps_type=int, result_type=int) (empty)

522 # this raises an error 

523 with pytest.raises(UserError, match="Tool name conflicts with result schema name: 'ctx_tool'"): (empty)

524 Agent('test', tools=[ctx_tool], deps_type=int, result_type=int, result_tool_name='ctx_tool') (empty)

525 

526 

527def test_init_ctx_tool_invalid(): (empty)

528 def plain_tool(x: int) -> int: # pragma: no cover (empty)

529 return x + 1 

530 

531 m = r'First parameter of tools that take context must be annotated with RunContext\[\.\.\.\]' (empty)

532 with pytest.raises(UserError, match=m): (empty)

533 Tool(plain_tool, takes_ctx=True) (empty)

534 

535 

536def test_init_plain_tool_invalid(): (empty)

537 with pytest.raises(UserError, match='RunContext annotations can only be used with tools that take context'): (empty)

538 Tool(ctx_tool, takes_ctx=False) (empty)

539 

540 

541def test_return_pydantic_model(): (empty)

542 agent = Agent('test') (empty)

543 

544 @agent.tool_plain (empty)

545 def return_pydantic_model(x: int) -> Foo: (empty)

546 return Foo(x=x, y='a') (empty)

547 

548 result = agent.run_sync('') (empty)

549 assert result.data == snapshot('{"return_pydantic_model":{"x":0,"y":"a"}}') (empty)

550 

551 

552def test_return_bytes(): (empty)

553 agent = Agent('test') (empty)

554 

555 @agent.tool_plain (empty)

556 def return_pydantic_model() -> bytes: (empty)

557 return '🐈 Hello'.encode() (empty)

558 

559 result = agent.run_sync('') (empty)

560 assert result.data == snapshot('{"return_pydantic_model":"🐈 Hello"}') (empty)

561 

562 

563def test_return_bytes_invalid(): (empty)

564 agent = Agent('test') (empty)

565 

566 @agent.tool_plain (empty)

567 def return_pydantic_model() -> bytes: (empty)

568 return b'\00 \x81' (empty)

569 

570 with pytest.raises(PydanticSerializationError, match='invalid utf-8 sequence of 1 bytes from index 2'): (empty)

571 agent.run_sync('') (empty)

572 

573 

574def test_return_unknown(): (empty)

575 agent = Agent('test') (empty)

576 

577 class Foobar: (empty)

578 pass (empty)

579 

580 @agent.tool_plain (empty)

581 def return_pydantic_model() -> Foobar: (empty)

582 return Foobar() (empty)

583 

584 with pytest.raises(PydanticSerializationError, match='Unable to serialize unknown type:'): (empty)

585 agent.run_sync('') (empty)

586 

587 

588def test_dynamic_cls_tool(): (empty)

589 @dataclass (empty)

590 class MyTool(Tool[int]): (empty)

591 spam: int (empty)

592 

593 def __init__(self, spam: int = 0, **kwargs: Any): (empty)

594 self.spam = spam (empty)

595 kwargs.update(function=self.tool_function, takes_ctx=False) (empty)

596 super().__init__(**kwargs) (empty)

597 

598 def tool_function(self, x: int, y: str) -> str: (empty)

599 return f'{self.spam} {x} {y}' (empty)

600 

601 async def prepare_tool_def(self, ctx: RunContext[int]) -> Union[ToolDefinition, None]: (empty)

602 if ctx.deps != 42: (empty)

603 return await super().prepare_tool_def(ctx) (empty)

604 

605 agent = Agent('test', tools=[MyTool(spam=777)], deps_type=int) (empty)

606 r = agent.run_sync('', deps=1) (empty)

607 assert r.data == snapshot('{"tool_function":"777 0 a"}') (empty)

608 

609 r = agent.run_sync('', deps=42) (empty)

610 assert r.data == snapshot('success (no tool calls)') (empty)

611 

612 

613def test_dynamic_plain_tool_decorator(): (empty)

614 agent = Agent('test', deps_type=int) (empty)

615 

616 async def prepare_tool_def(ctx: RunContext[int], tool_def: ToolDefinition) -> Union[ToolDefinition, None]: (empty)

617 if ctx.deps != 42: (empty)

618 return tool_def (empty)

619 

620 @agent.tool_plain(prepare=prepare_tool_def) (empty)

621 def foobar(x: int, y: str) -> str: (empty)

622 return f'{x} {y}' (empty)

623 

624 r = agent.run_sync('', deps=1) (empty)

625 assert r.data == snapshot('{"foobar":"0 a"}') (empty)

626 

627 r = agent.run_sync('', deps=42) (empty)

628 assert r.data == snapshot('success (no tool calls)') (empty)

629 

630 

631def test_dynamic_tool_decorator(): (empty)

632 agent = Agent('test', deps_type=int) (empty)

633 

634 async def prepare_tool_def(ctx: RunContext[int], tool_def: ToolDefinition) -> Union[ToolDefinition, None]: (empty)

635 if ctx.deps != 42: (empty)

636 return tool_def (empty)

637 

638 @agent.tool(prepare=prepare_tool_def) (empty)

639 def foobar(ctx: RunContext[int], x: int, y: str) -> str: (empty)

640 return f'{ctx.deps} {x} {y}' (empty)

641 

642 r = agent.run_sync('', deps=1) (empty)

643 assert r.data == snapshot('{"foobar":"1 0 a"}') (empty)

644 

645 r = agent.run_sync('', deps=42) (empty)

646 assert r.data == snapshot('success (no tool calls)') (empty)

647 

648 

649def test_plain_tool_name(): (empty)

650 agent = Agent(FunctionModel(get_json_schema)) (empty)

651 

652 def my_tool(arg: str) -> str: ... 652 ↛ exitline 652 didn't return from function 'my_tool' because (empty)

653 

654 agent.tool_plain(name='foo_tool')(my_tool) (empty)

655 result = agent.run_sync('Hello') (empty)

656 json_schema = json.loads(result.data) (empty)

657 assert json_schema['name'] == 'foo_tool' (empty)

658 

659 

660def test_tool_name(): (empty)

661 agent = Agent(FunctionModel(get_json_schema)) (empty)

662 

663 def my_tool(ctx: RunContext, arg: str) -> str: ... 663 ↛ exitline 663 didn't return from function 'my_tool' because (empty)

664 

665 agent.tool(name='foo_tool')(my_tool) (empty)

666 result = agent.run_sync('Hello') (empty)

667 json_schema = json.loads(result.data) (empty)

668 assert json_schema['name'] == 'foo_tool' (empty)

669 

670 

671def test_dynamic_tool_use_messages(): (empty)

672 async def repeat_call_foobar(_messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: (empty)

673 if info.function_tools: (empty)

674 tool = info.function_tools[0] (empty)

675 return ModelResponse(parts=[ToolCallPart(tool.name, {'x': 42, 'y': 'a'})]) (empty)

676 else: 

677 return ModelResponse(parts=[TextPart('done')]) (empty)

678 

679 agent = Agent(FunctionModel(repeat_call_foobar), deps_type=int) (empty)

680 

681 async def prepare_tool_def(ctx: RunContext[int], tool_def: ToolDefinition) -> Union[ToolDefinition, None]: (empty)

682 if len(ctx.messages) < 5: (empty)

683 return tool_def (empty)

684 

685 @agent.tool(prepare=prepare_tool_def) (empty)

686 def foobar(ctx: RunContext[int], x: int, y: str) -> str: (empty)

687 return f'{ctx.deps} {x} {y}' (empty)

688 

689 r = agent.run_sync('', deps=1) (empty)

690 assert r.data == snapshot('done') (empty)

691 message_part_kinds = [(m.kind, [p.part_kind for p in m.parts]) for m in r.all_messages()] (empty)

692 assert message_part_kinds == snapshot( (empty)

693 [ 

694 ('request', ['user-prompt']), 

695 ('response', ['tool-call']), 

696 ('request', ['tool-return']), 

697 ('response', ['tool-call']), 

698 ('request', ['tool-return']), 

699 ('response', ['text']), 

700 ] 

701 ) 

702 

703 

704def test_future_run_context(create_module: Callable[[str], Any]): (empty)

705 mod = create_module(""" (empty)

706from __future__ import annotations 

707 

708from pydantic_ai import Agent, RunContext 

709 

710def ctx_tool(ctx: RunContext[int], x: int) -> int: 

711 return x + ctx.deps 

712 

713agent = Agent('test', tools=[ctx_tool], deps_type=int) 

714 """) 

715 result = mod.agent.run_sync('foobar', deps=5) (empty)

716 assert result.data == snapshot('{"ctx_tool":5}') (empty)

717 

718 

719async def tool_without_return_annotation_in_docstring() -> str: # pragma: no cover (empty)

720 """A tool that documents what it returns but doesn't have a return annotation in the docstring.""" 

721 

722 return '' 

723 

724 

725def test_suppress_griffe_logging(caplog: LogCaptureFixture): (empty)

726 # This would cause griffe to emit a warning log if we didn't suppress the griffe logging. 

727 

728 agent = Agent(FunctionModel(get_json_schema)) (empty)

729 agent.tool_plain(tool_without_return_annotation_in_docstring) (empty)

730 

731 result = agent.run_sync('') (empty)

732 json_schema = json.loads(result.data) (empty)

733 assert json_schema == snapshot( (empty)

734 { 

735 'description': "A tool that documents what it returns but doesn't have a return annotation in the docstring.", 

736 'name': 'tool_without_return_annotation_in_docstring', 

737 'outer_typed_dict_key': None, 

738 'parameters_json_schema': {'additionalProperties': False, 'properties': {}, 'type': 'object'}, 

739 } 

740 ) 

741 

742 # Without suppressing griffe logging, we get: 

743 # assert caplog.messages == snapshot(['<module>:4: No type or annotation for returned value 1']) 

744 assert caplog.messages == snapshot([]) (empty)

745 

746 

747async def missing_parameter_descriptions_docstring(foo: int, bar: str) -> str: # pragma: no cover (empty)

748 """Describes function ops, but missing parameter descriptions.""" 

749 return f'{foo} {bar}' 

750 

751 

752def test_enforce_parameter_descriptions() -> None: (empty)

753 agent = Agent(FunctionModel(get_json_schema)) (empty)

754 

755 with pytest.raises(UserError) as exc_info: (empty)

756 agent.tool_plain(require_parameter_descriptions=True)(missing_parameter_descriptions_docstring) (empty)

757 

758 error_reason = exc_info.value.args[0] (empty)

759 error_parts = [ (empty)

760 'Error generating schema for missing_parameter_descriptions_docstring', 

761 'Missing parameter descriptions for ', 

762 'foo', 

763 'bar', 

764 ] 

765 assert all(err_part in error_reason for err_part in error_parts) (empty)

766 

767 

768def test_json_schema_required_parameters(set_event_loop: None): (empty)

769 agent = Agent(FunctionModel(get_json_schema)) (empty)

770 

771 @agent.tool (empty)

772 def my_tool(ctx: RunContext[None], a: int, b: int = 1) -> int: (empty)

773 raise NotImplementedError 

774 

775 @agent.tool_plain (empty)

776 def my_tool_plain(*, a: int = 1, b: int) -> int: (empty)