agent-wait 0.5.0__tar.gz → 0.6.0__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.
Files changed (27) hide show
  1. {agent_wait-0.5.0 → agent_wait-0.6.0}/PKG-INFO +64 -14
  2. {agent_wait-0.5.0 → agent_wait-0.6.0}/README.md +58 -11
  3. {agent_wait-0.5.0 → agent_wait-0.6.0}/pyproject.toml +5 -3
  4. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/__init__.py +5 -4
  5. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/framework.py +1 -1
  6. agent_wait-0.6.0/src/agent_wait/pydantic_ai/__init__.py +122 -0
  7. {agent_wait-0.5.0 → agent_wait-0.6.0}/.gitignore +0 -0
  8. {agent_wait-0.5.0 → agent_wait-0.6.0}/LICENSE +0 -0
  9. {agent_wait-0.5.0 → agent_wait-0.6.0}/reference/wait-sim/README.md +0 -0
  10. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/announce/__init__.py +0 -0
  11. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/announce/base.py +0 -0
  12. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/announce/composite.py +0 -0
  13. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/announce/log.py +0 -0
  14. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/announce/memory.py +0 -0
  15. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/announce/webhook.py +0 -0
  16. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/aws/__init__.py +0 -0
  17. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/aws/dynamodb.py +0 -0
  18. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/aws/eventbridge.py +0 -0
  19. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/aws/sns.py +0 -0
  20. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/aws/sqs.py +0 -0
  21. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/errors.py +0 -0
  22. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/langgraph/__init__.py +0 -0
  23. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/model.py +0 -0
  24. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/policy.py +0 -0
  25. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/publish.py +0 -0
  26. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/py.typed +0 -0
  27. {agent_wait-0.5.0 → agent_wait-0.6.0}/src/agent_wait/wait.py +0 -0
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: agent-wait
3
- Version: 0.5.0
4
- Summary: Human-in-the-loop for LangGraph agents -- and anything else a run has to wait for. Get the interrupt out of the process (SNS, SQS, EventBridge, DynamoDB, webhook) with everything needed to answer it; get the answer back in.
3
+ Version: 0.6.0
4
+ Summary: Human-in-the-loop for LangGraph and Pydantic AI agents -- and anything else a run has to wait for. Get the interrupt out of the process (SNS, SQS, EventBridge, DynamoDB, webhook) with everything needed to answer it; get the answer back in.
5
5
  Project-URL: Homepage, https://skamalj.github.io/agent-wait/
6
6
  Project-URL: Documentation, https://skamalj.github.io/agent-wait/
7
7
  Project-URL: Source, https://github.com/skamalj/agent-wait
@@ -10,7 +10,7 @@ Project-URL: Changelog, https://github.com/skamalj/agent-wait/blob/main/CHANGELO
10
10
  Author-email: Kamaljeet Singh <skamalj@gmail.com>
11
11
  License-Expression: MIT
12
12
  License-File: LICENSE
13
- Keywords: agents,approval,callback,durable,dynamodb,eventbridge,hitl,human-in-the-loop,interrupt,lambda,langchain,langgraph,serverless,sns,sqs,webhook
13
+ Keywords: agents,approval,callback,durable,dynamodb,eventbridge,hitl,human-in-the-loop,interrupt,lambda,langchain,langgraph,pydantic-ai,serverless,sns,sqs,webhook
14
14
  Classifier: Development Status :: 4 - Beta
15
15
  Classifier: Intended Audience :: Developers
16
16
  Classifier: License :: OSI Approved :: MIT License
@@ -23,10 +23,13 @@ Requires-Python: >=3.12
23
23
  Provides-Extra: all
24
24
  Requires-Dist: boto3>=1.35; extra == 'all'
25
25
  Requires-Dist: langgraph<2,>=1.2; extra == 'all'
26
+ Requires-Dist: pydantic-ai-slim<3,>=2.43; extra == 'all'
26
27
  Provides-Extra: aws
27
28
  Requires-Dist: boto3>=1.35; extra == 'aws'
28
29
  Provides-Extra: langgraph
29
30
  Requires-Dist: langgraph<2,>=1.2; extra == 'langgraph'
31
+ Provides-Extra: pydantic-ai
32
+ Requires-Dist: pydantic-ai-slim<3,>=2.43; extra == 'pydantic-ai'
30
33
  Description-Content-Type: text/markdown
31
34
 
32
35
  # agent-wait
@@ -43,12 +46,13 @@ deploy. But the same shape covers a vendor callback, a payment processor, a KYC
43
46
  another agent — anything the run cannot continue without, and nothing inside the process
44
47
  can supply.
45
48
 
46
- LangGraph has `interrupt()` for this. When a node calls it the graph pauses and the
47
- interrupt is handed to whatever called `invoke()` — and that is where LangGraph stops.
48
- There is no built-in way to tell anyone *else* that a question was asked, and no built-in
49
- way for anyone else to answer it. The moment the question has to reach Slack, an approvals
50
- dashboard, a ticket queue or another service, you are writing that code yourself — whether
51
- your agent is a server that runs for a year or a Lambda that is gone in seconds.
49
+ LangGraph has `interrupt()` for this; Pydantic AI has `ApprovalRequired`. In both, the
50
+ run pauses and the question is handed to whatever called it — and that is where the
51
+ framework stops. There is no built-in way to tell anyone *else* that a question was asked,
52
+ and no built-in way for anyone else to answer it. The moment the question has to reach
53
+ Slack, an approvals dashboard, a ticket queue or another service, you are writing that
54
+ code yourself — whether your agent is a server that runs for a year or a Lambda that is
55
+ gone in seconds.
52
56
 
53
57
  agent-wait is that code. It takes the interrupt and puts it somewhere people (or systems)
54
58
  can see it — a topic, a queue, a webhook, a database row — with everything needed to answer
@@ -56,6 +60,7 @@ it in one envelope, and documents the shape of the answer so the return leg is o
56
60
 
57
61
  ```bash
58
62
  pip install "agent-wait[langgraph]" # @wait + publish_interrupts for LangGraph
63
+ pip install "agent-wait[pydantic-ai]" # the same two names for Pydantic AI
59
64
  pip install "agent-wait[langgraph,aws]" # + SNS / SQS / EventBridge / DynamoDB announcers
60
65
  ```
61
66
 
@@ -145,6 +150,50 @@ publish_interrupts(result, message["thread_id"], announce)
145
150
  That is the complete integration. A duplicate answer runs nothing — LangGraph ignores a
146
151
  resume for a question the thread has moved past — so there is nothing to check.
147
152
 
153
+ ## The same two names on Pydantic AI
154
+
155
+ ```python
156
+ from agent_wait.pydantic_ai import publish_interrupts, wait
157
+ from pydantic_ai import Agent, DeferredToolRequests, DeferredToolResults, RunContext, ToolApproved, ToolDenied
158
+
159
+ agent = Agent("openai:gpt-5.4", output_type=[str, DeferredToolRequests])
160
+
161
+
162
+ @agent.tool
163
+ @wait(FINANCE)
164
+ def issue_refund(ctx: RunContext[None], order_id: str, amount: int) -> str:
165
+ payments.refund(order_id, amount)
166
+ return "refunded"
167
+
168
+
169
+ result = agent.run_sync("refund order 4471", message_history=history)
170
+ publish_interrupts(result, thread_id, announce=[SnsAnnounce(topic_arn)])
171
+ store(thread_id, result.all_messages()) # Pydantic AI has no checkpointer; the history is yours
172
+ ```
173
+
174
+ The run ends with `DeferredToolRequests` in `result.output`; `question_id` is the
175
+ `tool_call_id`. The answer goes back through the framework's own resume, with the answer
176
+ itself riding in `metadata` so the `@wait` rules apply unchanged:
177
+
178
+ ```python
179
+ approved = answer.get("action") == "approve"
180
+ result = agent.run_sync(
181
+ None,
182
+ message_history=load(thread_id),
183
+ deferred_tool_results=DeferredToolResults(
184
+ approvals={question_id: ToolApproved() if approved else ToolDenied(str(answer))},
185
+ metadata={question_id: answer},
186
+ ),
187
+ )
188
+ store(thread_id, result.all_messages()) # before acknowledging the answer message
189
+ ```
190
+
191
+ Two rules the framework leaves to you: name the `RunContext` parameter `ctx`, and store
192
+ the post-run history *before* you acknowledge an answer — a duplicate answer is refused
193
+ only if the history you pass already holds the tool's return. `requires_approval=True`
194
+ tools and `CallDeferred` calls (a job handed to an external system) are published too,
195
+ with the default policy.
196
+
148
197
  ## Async mode — the thread does not park
149
198
 
150
199
  ```python
@@ -240,16 +289,17 @@ One distribution, one import root. Frameworks and providers are subpackages behi
240
289
  ```
241
290
  agent_wait core: WaitPolicy, Question, the envelope, publish(), BaseAnnounce,
242
291
  Webhook/Log/InMemory announcers, and the Framework interface. No deps.
243
- agent_wait.langgraph [langgraph] @wait and publish_interrupts() bound to LangGraph.
244
- agent_wait.aws [aws] SnsAnnounce, SqsAnnounce, EventBridgeAnnounce, DynamoDbAnnounce.
292
+ agent_wait.langgraph [langgraph] @wait and publish_interrupts() bound to LangGraph.
293
+ agent_wait.pydantic_ai [pydantic-ai] the same two names bound to Pydantic AI's deferred tools.
294
+ agent_wait.aws [aws] SnsAnnounce, SqsAnnounce, EventBridgeAnnounce, DynamoDbAnnounce.
245
295
  examples/refund_agent a graph, a host written for Lambda behind SQS, and its CDK stack.
246
296
  docs/ message contract, architecture, announcer guide, consumer guide.
247
297
  ```
248
298
 
249
299
  `@wait` and `publish_interrupts` are written once against a three-method `Framework`
250
- interface; `agent_wait.langgraph` is one implementor of it. Announcers are the same
251
- split: `BaseAnnounce` is the interface, each provider subpackage the implementors. A
252
- `strands` extra follows the same pattern — see
300
+ interface; `agent_wait.langgraph` and `agent_wait.pydantic_ai` are implementors of it,
301
+ about fifty lines each. Announcers are the same split: `BaseAnnounce` is the interface,
302
+ each provider subpackage the implementors. See
253
303
  [Architecture](https://skamalj.github.io/agent-wait/architecture/).
254
304
 
255
305
  MIT. Issues and PRs at [github.com/skamalj/agent-wait](https://github.com/skamalj/agent-wait).
@@ -12,12 +12,13 @@ deploy. But the same shape covers a vendor callback, a payment processor, a KYC
12
12
  another agent — anything the run cannot continue without, and nothing inside the process
13
13
  can supply.
14
14
 
15
- LangGraph has `interrupt()` for this. When a node calls it the graph pauses and the
16
- interrupt is handed to whatever called `invoke()` — and that is where LangGraph stops.
17
- There is no built-in way to tell anyone *else* that a question was asked, and no built-in
18
- way for anyone else to answer it. The moment the question has to reach Slack, an approvals
19
- dashboard, a ticket queue or another service, you are writing that code yourself — whether
20
- your agent is a server that runs for a year or a Lambda that is gone in seconds.
15
+ LangGraph has `interrupt()` for this; Pydantic AI has `ApprovalRequired`. In both, the
16
+ run pauses and the question is handed to whatever called it — and that is where the
17
+ framework stops. There is no built-in way to tell anyone *else* that a question was asked,
18
+ and no built-in way for anyone else to answer it. The moment the question has to reach
19
+ Slack, an approvals dashboard, a ticket queue or another service, you are writing that
20
+ code yourself — whether your agent is a server that runs for a year or a Lambda that is
21
+ gone in seconds.
21
22
 
22
23
  agent-wait is that code. It takes the interrupt and puts it somewhere people (or systems)
23
24
  can see it — a topic, a queue, a webhook, a database row — with everything needed to answer
@@ -25,6 +26,7 @@ it in one envelope, and documents the shape of the answer so the return leg is o
25
26
 
26
27
  ```bash
27
28
  pip install "agent-wait[langgraph]" # @wait + publish_interrupts for LangGraph
29
+ pip install "agent-wait[pydantic-ai]" # the same two names for Pydantic AI
28
30
  pip install "agent-wait[langgraph,aws]" # + SNS / SQS / EventBridge / DynamoDB announcers
29
31
  ```
30
32
 
@@ -114,6 +116,50 @@ publish_interrupts(result, message["thread_id"], announce)
114
116
  That is the complete integration. A duplicate answer runs nothing — LangGraph ignores a
115
117
  resume for a question the thread has moved past — so there is nothing to check.
116
118
 
119
+ ## The same two names on Pydantic AI
120
+
121
+ ```python
122
+ from agent_wait.pydantic_ai import publish_interrupts, wait
123
+ from pydantic_ai import Agent, DeferredToolRequests, DeferredToolResults, RunContext, ToolApproved, ToolDenied
124
+
125
+ agent = Agent("openai:gpt-5.4", output_type=[str, DeferredToolRequests])
126
+
127
+
128
+ @agent.tool
129
+ @wait(FINANCE)
130
+ def issue_refund(ctx: RunContext[None], order_id: str, amount: int) -> str:
131
+ payments.refund(order_id, amount)
132
+ return "refunded"
133
+
134
+
135
+ result = agent.run_sync("refund order 4471", message_history=history)
136
+ publish_interrupts(result, thread_id, announce=[SnsAnnounce(topic_arn)])
137
+ store(thread_id, result.all_messages()) # Pydantic AI has no checkpointer; the history is yours
138
+ ```
139
+
140
+ The run ends with `DeferredToolRequests` in `result.output`; `question_id` is the
141
+ `tool_call_id`. The answer goes back through the framework's own resume, with the answer
142
+ itself riding in `metadata` so the `@wait` rules apply unchanged:
143
+
144
+ ```python
145
+ approved = answer.get("action") == "approve"
146
+ result = agent.run_sync(
147
+ None,
148
+ message_history=load(thread_id),
149
+ deferred_tool_results=DeferredToolResults(
150
+ approvals={question_id: ToolApproved() if approved else ToolDenied(str(answer))},
151
+ metadata={question_id: answer},
152
+ ),
153
+ )
154
+ store(thread_id, result.all_messages()) # before acknowledging the answer message
155
+ ```
156
+
157
+ Two rules the framework leaves to you: name the `RunContext` parameter `ctx`, and store
158
+ the post-run history *before* you acknowledge an answer — a duplicate answer is refused
159
+ only if the history you pass already holds the tool's return. `requires_approval=True`
160
+ tools and `CallDeferred` calls (a job handed to an external system) are published too,
161
+ with the default policy.
162
+
117
163
  ## Async mode — the thread does not park
118
164
 
119
165
  ```python
@@ -209,16 +255,17 @@ One distribution, one import root. Frameworks and providers are subpackages behi
209
255
  ```
210
256
  agent_wait core: WaitPolicy, Question, the envelope, publish(), BaseAnnounce,
211
257
  Webhook/Log/InMemory announcers, and the Framework interface. No deps.
212
- agent_wait.langgraph [langgraph] @wait and publish_interrupts() bound to LangGraph.
213
- agent_wait.aws [aws] SnsAnnounce, SqsAnnounce, EventBridgeAnnounce, DynamoDbAnnounce.
258
+ agent_wait.langgraph [langgraph] @wait and publish_interrupts() bound to LangGraph.
259
+ agent_wait.pydantic_ai [pydantic-ai] the same two names bound to Pydantic AI's deferred tools.
260
+ agent_wait.aws [aws] SnsAnnounce, SqsAnnounce, EventBridgeAnnounce, DynamoDbAnnounce.
214
261
  examples/refund_agent a graph, a host written for Lambda behind SQS, and its CDK stack.
215
262
  docs/ message contract, architecture, announcer guide, consumer guide.
216
263
  ```
217
264
 
218
265
  `@wait` and `publish_interrupts` are written once against a three-method `Framework`
219
- interface; `agent_wait.langgraph` is one implementor of it. Announcers are the same
220
- split: `BaseAnnounce` is the interface, each provider subpackage the implementors. A
221
- `strands` extra follows the same pattern — see
266
+ interface; `agent_wait.langgraph` and `agent_wait.pydantic_ai` are implementors of it,
267
+ about fifty lines each. Announcers are the same split: `BaseAnnounce` is the interface,
268
+ each provider subpackage the implementors. See
222
269
  [Architecture](https://skamalj.github.io/agent-wait/architecture/).
223
270
 
224
271
  MIT. Issues and PRs at [github.com/skamalj/agent-wait](https://github.com/skamalj/agent-wait).
@@ -1,7 +1,7 @@
1
1
  [project]
2
2
  name = "agent-wait"
3
- version = "0.5.0"
4
- description = "Human-in-the-loop for LangGraph agents -- and anything else a run has to wait for. Get the interrupt out of the process (SNS, SQS, EventBridge, DynamoDB, webhook) with everything needed to answer it; get the answer back in."
3
+ version = "0.6.0"
4
+ description = "Human-in-the-loop for LangGraph and Pydantic AI agents -- and anything else a run has to wait for. Get the interrupt out of the process (SNS, SQS, EventBridge, DynamoDB, webhook) with everything needed to answer it; get the answer back in."
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.12"
7
7
  license = "MIT"
@@ -13,6 +13,7 @@ keywords = [
13
13
  "approval",
14
14
  "langgraph",
15
15
  "langchain",
16
+ "pydantic-ai",
16
17
  "agents",
17
18
  "interrupt",
18
19
  "durable",
@@ -39,8 +40,9 @@ dependencies = []
39
40
 
40
41
  [project.optional-dependencies]
41
42
  langgraph = ["langgraph>=1.2,<2"]
43
+ pydantic-ai = ["pydantic-ai-slim>=2.43,<3"]
42
44
  aws = ["boto3>=1.35"]
43
- all = ["agent-wait[langgraph,aws]"]
45
+ all = ["agent-wait[langgraph,pydantic-ai,aws]"]
44
46
 
45
47
  [project.urls]
46
48
  Homepage = "https://skamalj.github.io/agent-wait/"
@@ -12,9 +12,10 @@
12
12
  publish_interrupts(result, thread_id, announce=[SnsAnnounce(topic_arn)])
13
13
 
14
14
  The core is framework- and provider-neutral: the policy, the question, the envelope,
15
- `publish()`, and the announcer interface. `agent_wait.langgraph` implements the
16
- `Framework` interface for LangGraph; `agent_wait.aws` implements announcers for AWS.
17
- Install the extras you need: `pip install agent-wait[langgraph,aws]`.
15
+ `publish()`, and the announcer interface. `agent_wait.langgraph` and
16
+ `agent_wait.pydantic_ai` implement the `Framework` interface for LangGraph and Pydantic
17
+ AI; `agent_wait.aws` implements announcers for AWS. Install the extras you need:
18
+ `pip install agent-wait[langgraph,aws]` or `agent-wait[pydantic-ai,aws]`.
18
19
  """
19
20
 
20
21
  from .announce import (
@@ -47,7 +48,7 @@ from .policy import WaitPolicy, parse_duration
47
48
  from .publish import build_envelope, publish
48
49
  from .wait import make_publish_interrupts, make_wait, question_id_for
49
50
 
50
- __version__ = "0.5.0"
51
+ __version__ = "0.6.0"
51
52
 
52
53
  __all__ = [
53
54
  "MAX_QUESTION_BYTES",
@@ -8,7 +8,7 @@ agent-wait needs exactly three things from an agent framework, and nothing else:
8
8
 
9
9
  Everything users touch -- `@wait`, `publish_interrupts`, the envelope, the announcers --
10
10
  is built once in `agent_wait.wait` on top of this interface. A framework subpackage
11
- (`agent_wait.langgraph`, `agent_wait.strands`, ...) is one implementor of it and two
11
+ (`agent_wait.langgraph`, `agent_wait.pydantic_ai`, ...) is one implementor of it and two
12
12
  lines of binding. Users never see this class; they import the bound names.
13
13
 
14
14
  The same split holds for announcers: `announce.base.BaseAnnounce` is the interface, and
@@ -0,0 +1,122 @@
1
+ """Pydantic AI implementor of `agent_wait.framework.Framework`, and the bound names.
2
+
3
+ from agent_wait.pydantic_ai import wait, publish_interrupts
4
+
5
+ Install with `pip install agent-wait[pydantic-ai]`.
6
+
7
+ ## What Pydantic AI provides, and what we use
8
+
9
+ - **Park:** a tool raises `ApprovalRequired(metadata=...)`. The run ends and
10
+ `result.output` is a `DeferredToolRequests` whose `approvals` list the parked calls
11
+ (`ToolCallPart.tool_call_id`, `.tool_name`, `.args`) and whose `metadata[id]` is
12
+ whatever the tool attached. The agent must list `DeferredToolRequests` in its
13
+ `output_type`, or the raise is an error.
14
+ - **Read back:** the same `result.output`. `stream()` ends with the same object.
15
+ - **Thread:** none. Pydantic AI has no checkpointer; `message_history` is a list the
16
+ host owns and stores under a name of its choosing. For `mode="async"` the host puts
17
+ that name in `deps` as `thread_id` (an attribute or a mapping key).
18
+ - **Resume** (the host's call, not ours):
19
+
20
+ agent.run_sync(None, message_history=history, deferred_tool_results=DeferredToolResults(
21
+ approvals={question_id: ToolApproved()}, # or ToolDenied(reason)
22
+ metadata={question_id: answer}, # the answer, verbatim, if it is a dict
23
+ ))
24
+
25
+ On approval the framework calls the tool again with `ctx.tool_call_approved = True`
26
+ and `ctx.tool_call_metadata = answer`; the `@wait` wrapper then applies the usual
27
+ rules (a `decision` parameter receives the answer; `answer["args"]` edits the call).
28
+ On denial the framework does not call the tool; the model sees the denial message.
29
+ Pass the answer through `metadata`, not `override_args` -- `override_args` replaces
30
+ the whole argument set and skips the wrapper's merge.
31
+
32
+ ## Three things to know
33
+
34
+ - **The host owns idempotency.** A duplicate answer is rejected only if the
35
+ `message_history` passed on resume already contains the tool's return -- so store
36
+ the post-run history *before* acknowledging the answer message. Resuming twice from
37
+ the pre-resume history runs the tool twice.
38
+ - **Name the `RunContext` parameter `ctx`.** The wrapper hides it from the published
39
+ `args` by name, and Pydantic AI finds it by annotation -- both hold under
40
+ `functools.wraps`.
41
+ - **`CallDeferred` calls are published too**, with the default policy: they are the
42
+ non-human wait (a job handed to an external system). The answer for those goes in
43
+ `DeferredToolResults.calls`, not `approvals`.
44
+ """
45
+
46
+ from __future__ import annotations
47
+
48
+ from collections.abc import Mapping
49
+ from typing import Any, cast
50
+
51
+ try:
52
+ from pydantic_ai import ApprovalRequired, DeferredToolRequests, RunContext
53
+ except ImportError as _err: # pragma: no cover - exercised only without the extra
54
+ raise ImportError(
55
+ "agent_wait.pydantic_ai needs Pydantic AI: pip install 'agent-wait[pydantic-ai]'"
56
+ ) from _err
57
+
58
+ from ..framework import Framework
59
+ from ..wait import WAIT_KEY, make_publish_interrupts, make_wait
60
+
61
+ __all__ = ["PydanticAIFramework", "publish_interrupts", "questions_in", "wait"]
62
+
63
+
64
+ def _run_context(call_args: Mapping[str, Any]) -> RunContext[Any] | None:
65
+ for value in call_args.values():
66
+ if isinstance(value, RunContext):
67
+ return value # pyright: ignore[reportUnknownVariableType]
68
+ return None
69
+
70
+
71
+ class PydanticAIFramework(Framework):
72
+ name = "pydantic_ai"
73
+ hidden_params = ("ctx",)
74
+
75
+ def interrupt(self, value: Mapping[str, Any], call_args: Mapping[str, Any]) -> Any:
76
+ ctx = _run_context(call_args)
77
+ if ctx is None:
78
+ # Without the context the re-run after approval is indistinguishable from
79
+ # the first call, and the tool would park forever.
80
+ raise TypeError("a @wait tool on Pydantic AI needs a `ctx: RunContext` parameter")
81
+ if ctx.tool_call_approved:
82
+ # The re-run after approval. The answer rides in the metadata the host
83
+ # attached; without one, approval is the whole answer.
84
+ answer: Any = ctx.tool_call_metadata
85
+ return answer if answer is not None else {"action": "approve"}
86
+ raise ApprovalRequired(metadata=dict(value))
87
+
88
+ def interrupts_in(self, result: Any) -> list[tuple[str, Any]]:
89
+ output: Any = getattr(result, "output", result)
90
+ if not isinstance(output, DeferredToolRequests):
91
+ return []
92
+ out: list[tuple[str, Any]] = []
93
+ for part in [*output.approvals, *output.calls]:
94
+ attached = output.metadata.get(part.tool_call_id)
95
+ if isinstance(attached, Mapping) and WAIT_KEY in attached:
96
+ out.append((part.tool_call_id, dict(attached))) # pyright: ignore[reportUnknownArgumentType]
97
+ else:
98
+ # A tool that never heard of this library: `requires_approval=True`, a bare
99
+ # `ApprovalRequired`, or a `CallDeferred`. The call is the question.
100
+ bare: dict[str, Any] = {"function": part.tool_name, "args": part.args_as_dict()}
101
+ if attached:
102
+ bare["metadata"] = dict(attached) # pyright: ignore[reportUnknownArgumentType]
103
+ out.append((part.tool_call_id, bare))
104
+ return out
105
+
106
+ def current_thread_id(self, call_args: Mapping[str, Any]) -> str:
107
+ ctx = _run_context(call_args)
108
+ deps: Any = ctx.deps if ctx is not None else None
109
+ if isinstance(deps, Mapping):
110
+ thread_id: Any = cast(Mapping[str, Any], deps).get("thread_id")
111
+ else:
112
+ thread_id = getattr(deps, "thread_id", None)
113
+ if not thread_id:
114
+ raise RuntimeError("@wait(mode='async') on Pydantic AI needs `thread_id` in the run's deps")
115
+ return str(thread_id)
116
+
117
+
118
+ _framework = PydanticAIFramework()
119
+
120
+ wait = make_wait(_framework)
121
+ publish_interrupts = make_publish_interrupts(_framework)
122
+ questions_in = publish_interrupts.questions_in # type: ignore[attr-defined]
File without changes
File without changes