pydantic-ai-slim 1.0.0b1__tar.gz → 1.0.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of pydantic-ai-slim might be problematic. Click here for more details.

Files changed (126) hide show
  1. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/PKG-INFO +9 -8
  2. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_a2a.py +1 -1
  3. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_agent_graph.py +65 -49
  4. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_parts_manager.py +3 -1
  5. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_tool_manager.py +33 -6
  6. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/ag_ui.py +75 -43
  7. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/agent/__init__.py +10 -7
  8. pydantic_ai_slim-1.0.2/pydantic_ai/durable_exec/dbos/__init__.py +6 -0
  9. pydantic_ai_slim-1.0.2/pydantic_ai/durable_exec/dbos/_agent.py +718 -0
  10. pydantic_ai_slim-1.0.2/pydantic_ai/durable_exec/dbos/_mcp_server.py +89 -0
  11. pydantic_ai_slim-1.0.2/pydantic_ai/durable_exec/dbos/_model.py +137 -0
  12. pydantic_ai_slim-1.0.2/pydantic_ai/durable_exec/dbos/_utils.py +10 -0
  13. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/durable_exec/temporal/_agent.py +71 -10
  14. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/exceptions.py +2 -2
  15. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/mcp.py +14 -26
  16. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/messages.py +90 -19
  17. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/__init__.py +9 -0
  18. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/anthropic.py +28 -11
  19. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/bedrock.py +6 -14
  20. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/gemini.py +3 -1
  21. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/google.py +58 -5
  22. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/groq.py +122 -34
  23. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/instrumented.py +29 -11
  24. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/openai.py +84 -29
  25. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/__init__.py +4 -0
  26. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/bedrock.py +11 -3
  27. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/google_vertex.py +2 -1
  28. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/groq.py +21 -2
  29. pydantic_ai_slim-1.0.2/pydantic_ai/providers/litellm.py +134 -0
  30. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/retries.py +42 -2
  31. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/tools.py +18 -7
  32. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/combined.py +2 -2
  33. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/function.py +54 -19
  34. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/usage.py +37 -3
  35. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pyproject.toml +12 -6
  36. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/.gitignore +0 -0
  37. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/LICENSE +0 -0
  38. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/README.md +0 -0
  39. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/__init__.py +0 -0
  40. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/__main__.py +0 -0
  41. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_cli.py +0 -0
  42. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_function_schema.py +0 -0
  43. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_griffe.py +0 -0
  44. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_mcp.py +0 -0
  45. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_otel_messages.py +0 -0
  46. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_output.py +0 -0
  47. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_run_context.py +0 -0
  48. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_system_prompt.py +0 -0
  49. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_thinking_part.py +0 -0
  50. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/_utils.py +0 -0
  51. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/agent/abstract.py +0 -0
  52. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/agent/wrapper.py +0 -0
  53. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/builtin_tools.py +0 -0
  54. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/common_tools/__init__.py +0 -0
  55. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/common_tools/duckduckgo.py +0 -0
  56. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/common_tools/tavily.py +0 -0
  57. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/direct.py +0 -0
  58. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/durable_exec/__init__.py +0 -0
  59. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/durable_exec/temporal/__init__.py +0 -0
  60. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/durable_exec/temporal/_function_toolset.py +0 -0
  61. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/durable_exec/temporal/_logfire.py +0 -0
  62. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/durable_exec/temporal/_mcp_server.py +0 -0
  63. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/durable_exec/temporal/_model.py +0 -0
  64. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/durable_exec/temporal/_run_context.py +0 -0
  65. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/durable_exec/temporal/_toolset.py +0 -0
  66. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/ext/__init__.py +0 -0
  67. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/ext/aci.py +0 -0
  68. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/ext/langchain.py +0 -0
  69. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/format_prompt.py +0 -0
  70. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/cohere.py +0 -0
  71. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/fallback.py +0 -0
  72. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/function.py +0 -0
  73. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/huggingface.py +0 -0
  74. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/mcp_sampling.py +0 -0
  75. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/mistral.py +0 -0
  76. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/test.py +0 -0
  77. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/models/wrapper.py +0 -0
  78. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/output.py +0 -0
  79. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/__init__.py +0 -0
  80. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/_json_schema.py +0 -0
  81. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/amazon.py +0 -0
  82. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/anthropic.py +0 -0
  83. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/cohere.py +0 -0
  84. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/deepseek.py +0 -0
  85. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/google.py +0 -0
  86. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/grok.py +0 -0
  87. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/groq.py +0 -0
  88. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/harmony.py +0 -0
  89. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/meta.py +0 -0
  90. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/mistral.py +0 -0
  91. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/moonshotai.py +0 -0
  92. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/openai.py +0 -0
  93. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/profiles/qwen.py +0 -0
  94. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/anthropic.py +0 -0
  95. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/azure.py +0 -0
  96. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/cerebras.py +0 -0
  97. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/cohere.py +0 -0
  98. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/deepseek.py +0 -0
  99. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/fireworks.py +0 -0
  100. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/github.py +0 -0
  101. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/google.py +0 -0
  102. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/google_gla.py +0 -0
  103. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/grok.py +0 -0
  104. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/heroku.py +0 -0
  105. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/huggingface.py +0 -0
  106. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/mistral.py +0 -0
  107. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/moonshotai.py +0 -0
  108. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/ollama.py +0 -0
  109. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/openai.py +0 -0
  110. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/openrouter.py +0 -0
  111. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/together.py +0 -0
  112. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/providers/vercel.py +0 -0
  113. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/py.typed +0 -0
  114. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/result.py +0 -0
  115. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/run.py +0 -0
  116. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/settings.py +0 -0
  117. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/__init__.py +0 -0
  118. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/_dynamic.py +0 -0
  119. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/abstract.py +0 -0
  120. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/approval_required.py +0 -0
  121. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/external.py +0 -0
  122. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/filtered.py +0 -0
  123. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/prefixed.py +0 -0
  124. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/prepared.py +0 -0
  125. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/renamed.py +0 -0
  126. {pydantic_ai_slim-1.0.0b1 → pydantic_ai_slim-1.0.2}/pydantic_ai/toolsets/wrapper.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pydantic-ai-slim
3
- Version: 1.0.0b1
3
+ Version: 1.0.2
4
4
  Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
5
5
  Project-URL: Homepage, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
6
6
  Project-URL: Source, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
@@ -9,7 +9,7 @@ Project-URL: Changelog, https://github.com/pydantic/pydantic-ai/releases
9
9
  Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>, Douwe Maan <douwe@pydantic.dev>
10
10
  License-Expression: MIT
11
11
  License-File: LICENSE
12
- Classifier: Development Status :: 4 - Beta
12
+ Classifier: Development Status :: 5 - Production/Stable
13
13
  Classifier: Environment :: Console
14
14
  Classifier: Environment :: MacOS X
15
15
  Classifier: Intended Audience :: Developers
@@ -28,13 +28,12 @@ Classifier: Programming Language :: Python :: 3.13
28
28
  Classifier: Topic :: Internet
29
29
  Classifier: Topic :: Software Development :: Libraries :: Python Modules
30
30
  Requires-Python: >=3.10
31
- Requires-Dist: eval-type-backport>=0.2.0
32
31
  Requires-Dist: exceptiongroup; python_version < '3.11'
33
- Requires-Dist: genai-prices>=0.0.22
32
+ Requires-Dist: genai-prices>=0.0.23
34
33
  Requires-Dist: griffe>=1.3.2
35
34
  Requires-Dist: httpx>=0.27
36
35
  Requires-Dist: opentelemetry-api>=1.28.0
37
- Requires-Dist: pydantic-graph==1.0.0b1
36
+ Requires-Dist: pydantic-graph==1.0.2
38
37
  Requires-Dist: pydantic>=2.10
39
38
  Requires-Dist: typing-inspection>=0.4.0
40
39
  Provides-Extra: a2a
@@ -53,10 +52,12 @@ Requires-Dist: pyperclip>=1.9.0; extra == 'cli'
53
52
  Requires-Dist: rich>=13; extra == 'cli'
54
53
  Provides-Extra: cohere
55
54
  Requires-Dist: cohere>=5.16.0; (platform_system != 'Emscripten') and extra == 'cohere'
55
+ Provides-Extra: dbos
56
+ Requires-Dist: dbos>=1.13.0; extra == 'dbos'
56
57
  Provides-Extra: duckduckgo
57
58
  Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
58
59
  Provides-Extra: evals
59
- Requires-Dist: pydantic-evals==1.0.0b1; extra == 'evals'
60
+ Requires-Dist: pydantic-evals==1.0.2; extra == 'evals'
60
61
  Provides-Extra: google
61
62
  Requires-Dist: google-genai>=1.31.0; extra == 'google'
62
63
  Provides-Extra: groq
@@ -66,7 +67,7 @@ Requires-Dist: huggingface-hub[inference]>=0.33.5; extra == 'huggingface'
66
67
  Provides-Extra: logfire
67
68
  Requires-Dist: logfire[httpx]>=3.14.1; extra == 'logfire'
68
69
  Provides-Extra: mcp
69
- Requires-Dist: mcp>=1.12.3; (python_version >= '3.10') and extra == 'mcp'
70
+ Requires-Dist: mcp>=1.12.3; extra == 'mcp'
70
71
  Provides-Extra: mistral
71
72
  Requires-Dist: mistralai>=1.9.2; extra == 'mistral'
72
73
  Provides-Extra: openai
@@ -76,7 +77,7 @@ Requires-Dist: tenacity>=8.2.3; extra == 'retries'
76
77
  Provides-Extra: tavily
77
78
  Requires-Dist: tavily-python>=0.5.0; extra == 'tavily'
78
79
  Provides-Extra: temporal
79
- Requires-Dist: temporalio==1.16.0; extra == 'temporal'
80
+ Requires-Dist: temporalio==1.17.0; extra == 'temporal'
80
81
  Provides-Extra: vertexai
81
82
  Requires-Dist: google-auth>=2.36.0; extra == 'vertexai'
82
83
  Requires-Dist: requests>=2.32.2; extra == 'vertexai'
@@ -272,7 +272,7 @@ class AgentWorker(Worker[list[ModelMessage]], Generic[WorkerOutputT, AgentDepsT]
272
272
  assert_never(part)
273
273
  return model_parts
274
274
 
275
- def _response_parts_to_a2a(self, parts: list[ModelResponsePart]) -> list[Part]:
275
+ def _response_parts_to_a2a(self, parts: Sequence[ModelResponsePart]) -> list[Part]:
276
276
  """Convert pydantic-ai ModelResponsePart objects to A2A Part objects.
277
277
 
278
278
  This handles the conversion from pydantic-ai's internal response parts to
@@ -2,7 +2,8 @@ from __future__ import annotations as _annotations
2
2
 
3
3
  import asyncio
4
4
  import dataclasses
5
- import hashlib
5
+ import inspect
6
+ from asyncio import Task
6
7
  from collections import defaultdict, deque
7
8
  from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
8
9
  from contextlib import asynccontextmanager, contextmanager
@@ -302,16 +303,21 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
302
303
  if self.system_prompt_dynamic_functions:
303
304
  for msg in messages:
304
305
  if isinstance(msg, _messages.ModelRequest):
305
- for i, part in enumerate(msg.parts):
306
+ reevaluated_message_parts: list[_messages.ModelRequestPart] = []
307
+ for part in msg.parts:
306
308
  if isinstance(part, _messages.SystemPromptPart) and part.dynamic_ref:
307
309
  # Look up the runner by its ref
308
310
  if runner := self.system_prompt_dynamic_functions.get( # pragma: lax no cover
309
311
  part.dynamic_ref
310
312
  ):
311
313
  updated_part_content = await runner.run(run_context)
312
- msg.parts[i] = _messages.SystemPromptPart(
313
- updated_part_content, dynamic_ref=part.dynamic_ref
314
- )
314
+ part = _messages.SystemPromptPart(updated_part_content, dynamic_ref=part.dynamic_ref)
315
+
316
+ reevaluated_message_parts.append(part)
317
+
318
+ # Replace message parts with reevaluated ones to prevent mutating parts list
319
+ if reevaluated_message_parts != msg.parts:
320
+ msg.parts = reevaluated_message_parts
315
321
 
316
322
  async def _sys_parts(self, run_context: RunContext[DepsT]) -> list[_messages.ModelRequestPart]:
317
323
  """Build the initial messages for the conversation."""
@@ -650,13 +656,6 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
650
656
  )
651
657
 
652
658
 
653
- def multi_modal_content_identifier(identifier: str | bytes) -> str:
654
- """Generate stable identifier for multi-modal content to help LLM in finding a specific file in tool call responses."""
655
- if isinstance(identifier, str):
656
- identifier = identifier.encode('utf-8')
657
- return hashlib.sha1(identifier).hexdigest()[:6]
658
-
659
-
660
659
  async def process_function_tools( # noqa: C901
661
660
  tool_manager: ToolManager[DepsT],
662
661
  tool_calls: list[_messages.ToolCallPart],
@@ -743,7 +742,6 @@ async def process_function_tools( # noqa: C901
743
742
  deferred_tool_results: dict[str, DeferredToolResult] = {}
744
743
  if build_run_context(ctx).tool_call_approved and ctx.deps.tool_call_results is not None:
745
744
  deferred_tool_results = ctx.deps.tool_call_results
746
-
747
745
  # Deferred tool calls are "run" as well, by reading their value from the tool call results
748
746
  calls_to_run.extend(tool_calls_by_kind['external'])
749
747
  calls_to_run.extend(tool_calls_by_kind['unapproved'])
@@ -764,6 +762,7 @@ async def process_function_tools( # noqa: C901
764
762
  calls_to_run,
765
763
  deferred_tool_results,
766
764
  ctx.deps.tracer,
765
+ ctx.deps.usage_limits,
767
766
  output_parts,
768
767
  deferred_calls,
769
768
  ):
@@ -810,6 +809,7 @@ async def _call_tools(
810
809
  tool_calls: list[_messages.ToolCallPart],
811
810
  deferred_tool_results: dict[str, DeferredToolResult],
812
811
  tracer: Tracer,
812
+ usage_limits: _usage.UsageLimits | None,
813
813
  output_parts: list[_messages.ModelRequestPart],
814
814
  output_deferred_calls: dict[Literal['external', 'unapproved'], list[_messages.ToolCallPart]],
815
815
  ) -> AsyncIterator[_messages.HandleResponseEvent]:
@@ -820,7 +820,6 @@ async def _call_tools(
820
820
  for call in tool_calls:
821
821
  yield _messages.FunctionToolCallEvent(call)
822
822
 
823
- # Run all tool tasks in parallel
824
823
  with tracer.start_as_current_span(
825
824
  'running tools',
826
825
  attributes={
@@ -828,39 +827,58 @@ async def _call_tools(
828
827
  'logfire.msg': f'running {len(tool_calls)} tool{"" if len(tool_calls) == 1 else "s"}',
829
828
  },
830
829
  ):
831
- tasks = [
832
- asyncio.create_task(
833
- _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id)),
834
- name=call.tool_name,
835
- )
836
- for call in tool_calls
837
- ]
838
-
839
- pending = tasks
840
- while pending:
841
- done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
842
- for task in done:
843
- index = tasks.index(task)
844
- try:
845
- tool_part, tool_user_part = task.result()
846
- except exceptions.CallDeferred:
847
- deferred_calls_by_index[index] = 'external'
848
- except exceptions.ApprovalRequired:
849
- deferred_calls_by_index[index] = 'unapproved'
850
- else:
851
- yield _messages.FunctionToolResultEvent(tool_part)
852
830
 
853
- tool_parts_by_index[index] = tool_part
854
- if tool_user_part:
855
- user_parts_by_index[index] = tool_user_part
831
+ async def handle_call_or_result(
832
+ coro_or_task: Awaitable[
833
+ tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]
834
+ ]
835
+ | Task[tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]],
836
+ index: int,
837
+ ) -> _messages.HandleResponseEvent | None:
838
+ try:
839
+ tool_part, tool_user_part = (
840
+ (await coro_or_task) if inspect.isawaitable(coro_or_task) else coro_or_task.result()
841
+ )
842
+ except exceptions.CallDeferred:
843
+ deferred_calls_by_index[index] = 'external'
844
+ except exceptions.ApprovalRequired:
845
+ deferred_calls_by_index[index] = 'unapproved'
846
+ else:
847
+ tool_parts_by_index[index] = tool_part
848
+ if tool_user_part:
849
+ user_parts_by_index[index] = tool_user_part
850
+
851
+ return _messages.FunctionToolResultEvent(tool_part)
852
+
853
+ if tool_manager.should_call_sequentially(tool_calls):
854
+ for index, call in enumerate(tool_calls):
855
+ if event := await handle_call_or_result(
856
+ _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
857
+ index,
858
+ ):
859
+ yield event
860
+
861
+ else:
862
+ tasks = [
863
+ asyncio.create_task(
864
+ _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
865
+ name=call.tool_name,
866
+ )
867
+ for call in tool_calls
868
+ ]
869
+
870
+ pending = tasks
871
+ while pending:
872
+ done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
873
+ for task in done:
874
+ index = tasks.index(task)
875
+ if event := await handle_call_or_result(coro_or_task=task, index=index):
876
+ yield event
856
877
 
857
878
  # We append the results at the end, rather than as they are received, to retain a consistent ordering
858
879
  # This is mostly just to simplify testing
859
- for k in sorted(tool_parts_by_index):
860
- output_parts.append(tool_parts_by_index[k])
861
-
862
- for k in sorted(user_parts_by_index):
863
- output_parts.append(user_parts_by_index[k])
880
+ output_parts.extend([tool_parts_by_index[k] for k in sorted(tool_parts_by_index)])
881
+ output_parts.extend([user_parts_by_index[k] for k in sorted(user_parts_by_index)])
864
882
 
865
883
  for k in sorted(deferred_calls_by_index):
866
884
  output_deferred_calls[deferred_calls_by_index[k]].append(tool_calls[k])
@@ -870,14 +888,15 @@ async def _call_tool(
870
888
  tool_manager: ToolManager[DepsT],
871
889
  tool_call: _messages.ToolCallPart,
872
890
  tool_call_result: DeferredToolResult | None,
891
+ usage_limits: _usage.UsageLimits | None,
873
892
  ) -> tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]:
874
893
  try:
875
894
  if tool_call_result is None:
876
- tool_result = await tool_manager.handle_call(tool_call)
895
+ tool_result = await tool_manager.handle_call(tool_call, usage_limits=usage_limits)
877
896
  elif isinstance(tool_call_result, ToolApproved):
878
897
  if tool_call_result.override_args is not None:
879
898
  tool_call = dataclasses.replace(tool_call, args=tool_call_result.override_args)
880
- tool_result = await tool_manager.handle_call(tool_call)
899
+ tool_result = await tool_manager.handle_call(tool_call, usage_limits=usage_limits)
881
900
  elif isinstance(tool_call_result, ToolDenied):
882
901
  return _messages.ToolReturnPart(
883
902
  tool_name=tool_call.tool_name,
@@ -915,10 +934,7 @@ async def _call_tool(
915
934
  f'`ToolReturn` should be used directly.'
916
935
  )
917
936
  elif isinstance(content, _messages.MultiModalContent):
918
- if isinstance(content, _messages.BinaryContent):
919
- identifier = content.identifier or multi_modal_content_identifier(content.data)
920
- else:
921
- identifier = multi_modal_content_identifier(content.url)
937
+ identifier = content.identifier
922
938
 
923
939
  return_values.append(f'See file {identifier}')
924
940
  user_contents.extend([f'This is file {identifier}:', content])
@@ -154,6 +154,7 @@ class ModelResponsePartsManager:
154
154
  *,
155
155
  vendor_part_id: Hashable | None,
156
156
  content: str | None = None,
157
+ id: str | None = None,
157
158
  signature: str | None = None,
158
159
  ) -> ModelResponseStreamEvent:
159
160
  """Handle incoming thinking content, creating or updating a ThinkingPart in the manager as appropriate.
@@ -167,6 +168,7 @@ class ModelResponsePartsManager:
167
168
  of thinking. If None, a new part will be created unless the latest part is already
168
169
  a ThinkingPart.
169
170
  content: The thinking content to append to the appropriate ThinkingPart.
171
+ id: An optional id for the thinking part.
170
172
  signature: An optional signature for the thinking content.
171
173
 
172
174
  Returns:
@@ -197,7 +199,7 @@ class ModelResponsePartsManager:
197
199
  if content is not None:
198
200
  # There is no existing thinking part that should be updated, so create a new one
199
201
  new_part_index = len(self._parts)
200
- part = ThinkingPart(content=content, signature=signature)
202
+ part = ThinkingPart(content=content, id=id, signature=signature)
201
203
  if vendor_part_id is not None: # pragma: no branch
202
204
  self._vendor_id_to_part_index[vendor_part_id] = new_part_index
203
205
  self._parts.append(part)
@@ -14,6 +14,7 @@ from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior
14
14
  from .messages import ToolCallPart
15
15
  from .tools import ToolDefinition
16
16
  from .toolsets.abstract import AbstractToolset, ToolsetTool
17
+ from .usage import UsageLimits
17
18
 
18
19
 
19
20
  @dataclass
@@ -55,6 +56,10 @@ class ToolManager(Generic[AgentDepsT]):
55
56
 
56
57
  return [tool.tool_def for tool in self.tools.values()]
57
58
 
59
+ def should_call_sequentially(self, calls: list[ToolCallPart]) -> bool:
60
+ """Whether to require sequential tool calls for a list of tool calls."""
61
+ return any(tool_def.sequential for call in calls if (tool_def := self.get_tool_def(call.tool_name)))
62
+
58
63
  def get_tool_def(self, name: str) -> ToolDefinition | None:
59
64
  """Get the tool definition for a given tool name, or `None` if the tool is unknown."""
60
65
  if self.tools is None:
@@ -66,7 +71,11 @@ class ToolManager(Generic[AgentDepsT]):
66
71
  return None
67
72
 
68
73
  async def handle_call(
69
- self, call: ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
74
+ self,
75
+ call: ToolCallPart,
76
+ allow_partial: bool = False,
77
+ wrap_validation_errors: bool = True,
78
+ usage_limits: UsageLimits | None = None,
70
79
  ) -> Any:
71
80
  """Handle a tool call by validating the arguments, calling the tool, and handling retries.
72
81
 
@@ -74,13 +83,14 @@ class ToolManager(Generic[AgentDepsT]):
74
83
  call: The tool call part to handle.
75
84
  allow_partial: Whether to allow partial validation of the tool arguments.
76
85
  wrap_validation_errors: Whether to wrap validation errors in a retry prompt part.
86
+ usage_limits: Optional usage limits to check before executing tools.
77
87
  """
78
88
  if self.tools is None or self.ctx is None:
79
89
  raise ValueError('ToolManager has not been prepared for a run step yet') # pragma: no cover
80
90
 
81
91
  if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
82
- # Output tool calls are not traced
83
- return await self._call_tool(call, allow_partial, wrap_validation_errors)
92
+ # Output tool calls are not traced and not counted
93
+ return await self._call_tool(call, allow_partial, wrap_validation_errors, count_tool_usage=False)
84
94
  else:
85
95
  return await self._call_tool_traced(
86
96
  call,
@@ -88,9 +98,17 @@ class ToolManager(Generic[AgentDepsT]):
88
98
  wrap_validation_errors,
89
99
  self.ctx.tracer,
90
100
  self.ctx.trace_include_content,
101
+ usage_limits,
91
102
  )
92
103
 
93
- async def _call_tool(self, call: ToolCallPart, allow_partial: bool, wrap_validation_errors: bool) -> Any:
104
+ async def _call_tool(
105
+ self,
106
+ call: ToolCallPart,
107
+ allow_partial: bool,
108
+ wrap_validation_errors: bool,
109
+ usage_limits: UsageLimits | None = None,
110
+ count_tool_usage: bool = True,
111
+ ) -> Any:
94
112
  if self.tools is None or self.ctx is None:
95
113
  raise ValueError('ToolManager has not been prepared for a run step yet') # pragma: no cover
96
114
 
@@ -121,7 +139,15 @@ class ToolManager(Generic[AgentDepsT]):
121
139
  else:
122
140
  args_dict = validator.validate_python(call.args or {}, allow_partial=pyd_allow_partial)
123
141
 
124
- return await self.toolset.call_tool(name, args_dict, ctx, tool)
142
+ if usage_limits is not None and count_tool_usage:
143
+ usage_limits.check_before_tool_call(self.ctx.usage)
144
+
145
+ result = await self.toolset.call_tool(name, args_dict, ctx, tool)
146
+
147
+ if count_tool_usage:
148
+ self.ctx.usage.tool_calls += 1
149
+
150
+ return result
125
151
  except (ValidationError, ModelRetry) as e:
126
152
  max_retries = tool.max_retries if tool is not None else 1
127
153
  current_retry = self.ctx.retries.get(name, 0)
@@ -160,6 +186,7 @@ class ToolManager(Generic[AgentDepsT]):
160
186
  wrap_validation_errors: bool,
161
187
  tracer: Tracer,
162
188
  include_content: bool = False,
189
+ usage_limits: UsageLimits | None = None,
163
190
  ) -> Any:
164
191
  """See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>."""
165
192
  span_attributes = {
@@ -189,7 +216,7 @@ class ToolManager(Generic[AgentDepsT]):
189
216
  }
190
217
  with tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
191
218
  try:
192
- tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
219
+ tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors, usage_limits)
193
220
  except ToolRetryError as e:
194
221
  part = e.tool_retry
195
222
  if include_content and span.is_recording():
@@ -68,6 +68,9 @@ try:
68
68
  TextMessageContentEvent,
69
69
  TextMessageEndEvent,
70
70
  TextMessageStartEvent,
71
+ # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
72
+ # ThinkingEndEvent,
73
+ # ThinkingStartEvent,
71
74
  ThinkingTextMessageContentEvent,
72
75
  ThinkingTextMessageEndEvent,
73
76
  ThinkingTextMessageStartEvent,
@@ -392,6 +395,12 @@ async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEve
392
395
  if stream_ctx.part_end: # pragma: no branch
393
396
  yield stream_ctx.part_end
394
397
  stream_ctx.part_end = None
398
+ if stream_ctx.thinking:
399
+ # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
400
+ # yield ThinkingEndEvent(
401
+ # type=EventType.THINKING_END,
402
+ # )
403
+ stream_ctx.thinking = False
395
404
  elif isinstance(node, CallToolsNode):
396
405
  async with node.stream(run.ctx) as handle_stream:
397
406
  async for event in handle_stream:
@@ -400,7 +409,7 @@ async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEve
400
409
  yield msg
401
410
 
402
411
 
403
- async def _handle_model_request_event(
412
+ async def _handle_model_request_event( # noqa: C901
404
413
  stream_ctx: _RequestStreamContext,
405
414
  agent_event: ModelResponseStreamEvent,
406
415
  ) -> AsyncIterator[BaseEvent]:
@@ -420,56 +429,70 @@ async def _handle_model_request_event(
420
429
  stream_ctx.part_end = None
421
430
 
422
431
  part = agent_event.part
423
- if isinstance(part, TextPart):
424
- message_id = stream_ctx.new_message_id()
425
- yield TextMessageStartEvent(
426
- message_id=message_id,
427
- )
428
- if part.content: # pragma: no branch
429
- yield TextMessageContentEvent(
430
- message_id=message_id,
432
+ if isinstance(part, ThinkingPart): # pragma: no branch
433
+ if not stream_ctx.thinking:
434
+ # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
435
+ # yield ThinkingStartEvent(
436
+ # type=EventType.THINKING_START,
437
+ # )
438
+ stream_ctx.thinking = True
439
+
440
+ if part.content:
441
+ yield ThinkingTextMessageStartEvent(
442
+ type=EventType.THINKING_TEXT_MESSAGE_START,
443
+ )
444
+ yield ThinkingTextMessageContentEvent(
445
+ type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
431
446
  delta=part.content,
432
447
  )
433
- stream_ctx.part_end = TextMessageEndEvent(
434
- message_id=message_id,
435
- )
436
- elif isinstance(part, ToolCallPart): # pragma: no branch
437
- message_id = stream_ctx.message_id or stream_ctx.new_message_id()
438
- yield ToolCallStartEvent(
439
- tool_call_id=part.tool_call_id,
440
- tool_call_name=part.tool_name,
441
- parent_message_id=message_id,
442
- )
443
- if part.args:
444
- yield ToolCallArgsEvent(
448
+ stream_ctx.part_end = ThinkingTextMessageEndEvent(
449
+ type=EventType.THINKING_TEXT_MESSAGE_END,
450
+ )
451
+ else:
452
+ if stream_ctx.thinking:
453
+ # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
454
+ # yield ThinkingEndEvent(
455
+ # type=EventType.THINKING_END,
456
+ # )
457
+ stream_ctx.thinking = False
458
+
459
+ if isinstance(part, TextPart):
460
+ message_id = stream_ctx.new_message_id()
461
+ yield TextMessageStartEvent(
462
+ message_id=message_id,
463
+ )
464
+ if part.content: # pragma: no branch
465
+ yield TextMessageContentEvent(
466
+ message_id=message_id,
467
+ delta=part.content,
468
+ )
469
+ stream_ctx.part_end = TextMessageEndEvent(
470
+ message_id=message_id,
471
+ )
472
+ elif isinstance(part, ToolCallPart): # pragma: no branch
473
+ message_id = stream_ctx.message_id or stream_ctx.new_message_id()
474
+ yield ToolCallStartEvent(
475
+ tool_call_id=part.tool_call_id,
476
+ tool_call_name=part.tool_name,
477
+ parent_message_id=message_id,
478
+ )
479
+ if part.args:
480
+ yield ToolCallArgsEvent(
481
+ tool_call_id=part.tool_call_id,
482
+ delta=part.args if isinstance(part.args, str) else json.dumps(part.args),
483
+ )
484
+ stream_ctx.part_end = ToolCallEndEvent(
445
485
  tool_call_id=part.tool_call_id,
446
- delta=part.args if isinstance(part.args, str) else json.dumps(part.args),
447
486
  )
448
- stream_ctx.part_end = ToolCallEndEvent(
449
- tool_call_id=part.tool_call_id,
450
- )
451
-
452
- elif isinstance(part, ThinkingPart): # pragma: no branch
453
- yield ThinkingTextMessageStartEvent(
454
- type=EventType.THINKING_TEXT_MESSAGE_START,
455
- )
456
- # Always send the content even if it's empty, as it may be
457
- # used to indicate the start of thinking.
458
- yield ThinkingTextMessageContentEvent(
459
- type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
460
- delta=part.content,
461
- )
462
- stream_ctx.part_end = ThinkingTextMessageEndEvent(
463
- type=EventType.THINKING_TEXT_MESSAGE_END,
464
- )
465
487
 
466
488
  elif isinstance(agent_event, PartDeltaEvent):
467
489
  delta = agent_event.delta
468
490
  if isinstance(delta, TextPartDelta):
469
- yield TextMessageContentEvent(
470
- message_id=stream_ctx.message_id,
471
- delta=delta.content_delta,
472
- )
491
+ if delta.content_delta: # pragma: no branch
492
+ yield TextMessageContentEvent(
493
+ message_id=stream_ctx.message_id,
494
+ delta=delta.content_delta,
495
+ )
473
496
  elif isinstance(delta, ToolCallPartDelta): # pragma: no branch
474
497
  assert delta.tool_call_id, '`ToolCallPartDelta.tool_call_id` must be set'
475
498
  yield ToolCallArgsEvent(
@@ -478,6 +501,14 @@ async def _handle_model_request_event(
478
501
  )
479
502
  elif isinstance(delta, ThinkingPartDelta): # pragma: no branch
480
503
  if delta.content_delta: # pragma: no branch
504
+ if not isinstance(stream_ctx.part_end, ThinkingTextMessageEndEvent):
505
+ yield ThinkingTextMessageStartEvent(
506
+ type=EventType.THINKING_TEXT_MESSAGE_START,
507
+ )
508
+ stream_ctx.part_end = ThinkingTextMessageEndEvent(
509
+ type=EventType.THINKING_TEXT_MESSAGE_END,
510
+ )
511
+
481
512
  yield ThinkingTextMessageContentEvent(
482
513
  type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
483
514
  delta=delta.content_delta,
@@ -629,6 +660,7 @@ class _RequestStreamContext:
629
660
 
630
661
  message_id: str = ''
631
662
  part_end: BaseEvent | None = None
663
+ thinking: bool = False
632
664
 
633
665
  def new_message_id(self) -> str:
634
666
  """Generate a new message ID for the request stream.
@@ -4,15 +4,15 @@ import dataclasses
4
4
  import inspect
5
5
  import json
6
6
  import warnings
7
+ from asyncio import Lock
7
8
  from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
8
9
  from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager
9
10
  from contextvars import ContextVar
10
11
  from typing import TYPE_CHECKING, Any, ClassVar, cast, overload
11
12
 
12
- import anyio
13
13
  from opentelemetry.trace import NoOpTracer, use_span
14
14
  from pydantic.json_schema import GenerateJsonSchema
15
- from typing_extensions import TypeVar, deprecated
15
+ from typing_extensions import Self, TypeVar, deprecated
16
16
 
17
17
  from pydantic_graph import Graph
18
18
 
@@ -157,7 +157,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
157
157
 
158
158
  _event_stream_handler: EventStreamHandler[AgentDepsT] | None = dataclasses.field(repr=False)
159
159
 
160
- _enter_lock: anyio.Lock = dataclasses.field(repr=False)
160
+ _enter_lock: Lock = dataclasses.field(repr=False)
161
161
  _entered_count: int = dataclasses.field(repr=False)
162
162
  _exit_stack: AsyncExitStack | None = dataclasses.field(repr=False)
163
163
 
@@ -374,7 +374,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
374
374
  _utils.Option[Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]]]
375
375
  ] = ContextVar('_override_tools', default=None)
376
376
 
377
- self._enter_lock = anyio.Lock()
377
+ self._enter_lock = Lock()
378
378
  self._entered_count = 0
379
379
  self._exit_stack = None
380
380
 
@@ -1066,7 +1066,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
1066
1066
  strict: Whether to enforce JSON schema compliance (only affects OpenAI).
1067
1067
  See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
1068
1068
  requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
1069
- See the [tools documentation](../tools.md#human-in-the-loop-tool-approval) for more info.
1069
+ See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
1070
1070
  """
1071
1071
 
1072
1072
  def tool_decorator(
@@ -1119,6 +1119,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
1119
1119
  require_parameter_descriptions: bool = False,
1120
1120
  schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
1121
1121
  strict: bool | None = None,
1122
+ sequential: bool = False,
1122
1123
  requires_approval: bool = False,
1123
1124
  ) -> Any:
1124
1125
  """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
@@ -1164,8 +1165,9 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
1164
1165
  schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
1165
1166
  strict: Whether to enforce JSON schema compliance (only affects OpenAI).
1166
1167
  See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
1168
+ sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
1167
1169
  requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
1168
- See the [tools documentation](../tools.md#human-in-the-loop-tool-approval) for more info.
1170
+ See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
1169
1171
  """
1170
1172
 
1171
1173
  def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
@@ -1180,6 +1182,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
1180
1182
  require_parameter_descriptions,
1181
1183
  schema_generator,
1182
1184
  strict,
1185
+ sequential,
1183
1186
  requires_approval,
1184
1187
  )
1185
1188
  return func_
@@ -1355,7 +1358,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
1355
1358
 
1356
1359
  return schema # pyright: ignore[reportReturnType]
1357
1360
 
1358
- async def __aenter__(self) -> AbstractAgent[AgentDepsT, OutputDataT]:
1361
+ async def __aenter__(self) -> Self:
1359
1362
  """Enter the agent context.
1360
1363
 
1361
1364
  This will start all [`MCPServerStdio`s][pydantic_ai.mcp.MCPServerStdio] registered as `toolsets` so they are ready to be used.
@@ -0,0 +1,6 @@
1
+ from ._agent import DBOSAgent
2
+ from ._mcp_server import DBOSMCPServer
3
+ from ._model import DBOSModel
4
+ from ._utils import StepConfig
5
+
6
+ __all__ = ['DBOSAgent', 'DBOSModel', 'DBOSMCPServer', 'StepConfig']