15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206 | class WaitForTool(Tool):
"""Poll until a condition holds or a timeout elapses.
Conditions (``condition``):
- ``url`` — poll ``target`` (a URL) with HTTP GET until it responds;
``status`` controls what counts as success (default ``success`` =
2xx).
- ``redis_key`` — poll ``target`` (a key) in a Redis-compatible
store until it exists.
Args:
condition: ``url`` or ``redis_key``.
target: URL or key to poll.
timeout: Seconds before giving up (default from config, 120.0).
poll_interval: Seconds between checks (default from config, 1.0).
status: For ``url``: ``success`` (2xx), ``any``, or an exact
HTTP status code.
Args (config): ``poll_interval``, ``timeout``, and connection keys
for the ``redis_key`` condition (same as the ``redis`` tool).
"""
name = "wait_for"
description = (
"Poll until a condition holds (URL reachable, Redis key exists) "
"or a timeout elapses"
)
def __init__(self, config: dict | None = None):
cfg = config or {}
self.poll_interval = float(cfg.get("poll_interval", 1.0))
self.timeout = float(cfg.get("timeout", 120.0))
self.url = cfg.get("url", "")
self.host = cfg.get("host", "localhost")
self.port = cfg.get("port", 6379)
self.db = cfg.get("db", 0)
self.password = cfg.get("password", "")
self.username = cfg.get("username", "")
def _redis(self):
try:
import redis
except ImportError as e:
msg = "wait_for redis_key requires the 'redis' package (pip install teff[tools])"
raise ImportError(msg) from e
if self.url:
return redis.Redis.from_url(self.url, decode_responses=True)
kwargs: dict = {}
if self.username:
kwargs["username"] = self.username
if self.password:
kwargs["password"] = self.password
return redis.Redis(
host=self.host,
port=int(self.port),
db=int(self.db),
decode_responses=True,
**kwargs,
)
def run( # type: ignore[override]
self,
condition: str,
target: str = "",
timeout: float | None = None,
poll_interval: float | None = None,
status: str = "success",
) -> str:
if not condition:
raise ValueError("condition is required (url, redis_key)")
if not target:
raise ValueError("target is required")
timeout = float(timeout if timeout is not None else self.timeout)
interval = float(
poll_interval if poll_interval is not None else self.poll_interval
)
start = time.monotonic()
if condition == "url":
self._poll_url(target, timeout, interval, status)
elif condition == "redis_key":
client = self._redis()
try:
self._poll(lambda: bool(client.exists(target)), timeout, interval)
finally:
client.close()
else:
raise ValueError(f"unknown condition: {condition}")
return f"condition met after {time.monotonic() - start:.1f}s"
async def arun( # type: ignore[override]
self,
condition: str,
target: str = "",
timeout: float | None = None,
poll_interval: float | None = None,
status: str = "success",
) -> str:
if not condition:
raise ValueError("condition is required (url, redis_key)")
if not target:
raise ValueError("target is required")
timeout = float(timeout if timeout is not None else self.timeout)
interval = float(
poll_interval if poll_interval is not None else self.poll_interval
)
start = time.monotonic()
if condition == "url":
await self._poll_url_async(target, timeout, interval, status)
elif condition == "redis_key":
client = self._redis()
try:
await self._poll_async(
lambda: bool(client.exists(target)), timeout, interval
)
finally:
client.close()
else:
raise ValueError(f"unknown condition: {condition}")
return f"condition met after {time.monotonic() - start:.1f}s"
def _poll(self, check, timeout: float, interval: float) -> None:
deadline = time.monotonic() + timeout
while True:
try:
if check():
return
except Exception:
pass
if time.monotonic() >= deadline:
raise ValueError(f"timed out after {timeout:.0f}s")
time.sleep(interval)
async def _poll_async(self, check, timeout: float, interval: float) -> None:
import asyncio
deadline = time.monotonic() + timeout
while True:
try:
if check():
return
except Exception:
pass
if time.monotonic() >= deadline:
raise ValueError(f"timed out after {timeout:.0f}s")
await asyncio.sleep(interval)
def _poll_url(self, url: str, timeout: float, interval: float, status: str) -> None:
if status not in ("success", "any") and not str(status).isdigit():
raise ValueError(f"unknown status expectation: {status}")
def check() -> bool:
import httpx
try:
response = httpx.get(url, timeout=interval + 2, follow_redirects=True)
except Exception:
return False
code = response.status_code
if status == "any":
return True
if status == "success":
return 200 <= code < 300
return code == int(status)
self._poll(check, timeout, interval)
async def _poll_url_async(
self, url: str, timeout: float, interval: float, status: str
) -> None:
if status not in ("success", "any") and not str(status).isdigit():
raise ValueError(f"unknown status expectation: {status}")
import httpx
async with httpx.AsyncClient() as client:
async def check() -> bool:
try:
response = await client.get(
url, timeout=interval + 2, follow_redirects=True
)
except Exception:
return False
code = response.status_code
if status == "any":
return True
if status == "success":
return 200 <= code < 300
return code == int(status)
await self._poll_async(check, timeout, interval)
|