bynn-decisionnode 0.3.1__py3-none-any.whl

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.
@@ -0,0 +1,2269 @@
1
+ """The clients: :class:`DecisionNode` (blocking) and :class:`AsyncDecisionNode` (asyncio)."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import asyncio
6
+ import base64
7
+ import ssl
8
+ import threading
9
+ import time
10
+ import uuid
11
+ from concurrent.futures import ThreadPoolExecutor
12
+ from types import TracebackType
13
+ from typing import (
14
+ Any,
15
+ AsyncIterator,
16
+ Awaitable,
17
+ Callable,
18
+ Dict,
19
+ Iterable,
20
+ List,
21
+ Mapping,
22
+ Optional,
23
+ Sequence,
24
+ Tuple,
25
+ Type,
26
+ TypeVar,
27
+ Union,
28
+ cast,
29
+ )
30
+ from urllib.parse import quote, urlencode, urlsplit
31
+
32
+ import httpx
33
+
34
+ from . import _base
35
+ from ._base import DEFAULT, IDEMPOTENT, MUTATION, OCR, SESSION_OPEN, UNSET, ImageInput, Settings
36
+ from ._batches import (
37
+ NDJSON,
38
+ AsyncBatchResults,
39
+ Batch,
40
+ BatchPage,
41
+ BatchResults,
42
+ BatchRun,
43
+ check_poll,
44
+ line_bytes,
45
+ parse_batch,
46
+ parse_page,
47
+ poll_wait,
48
+ split_lines,
49
+ )
50
+ from ._boosters import BoosterInput, booster_wire, names_ocr
51
+ from ._errors import BadRequestError, ConnectionError, DecisionNodeError, ServerError, TimeoutError
52
+ from ._images import Image
53
+ from ._questions import Question
54
+ from ._questions import choice as _choice
55
+ from ._questions import number as _number
56
+ from ._questions import score as _score
57
+ from ._questions import truth as _truth
58
+ from ._sessions import (
59
+ AsyncSession,
60
+ FrameError,
61
+ Reply,
62
+ Session,
63
+ SessionClosed,
64
+ asupervise,
65
+ supervise,
66
+ )
67
+ from ._types import (
68
+ AnyAnswer,
69
+ ChoiceAnswer,
70
+ Decision,
71
+ ModelInfo,
72
+ NumberAnswer,
73
+ ScoreAnswer,
74
+ TruthAnswer,
75
+ parse_decision,
76
+ parse_model,
77
+ )
78
+ from ._uploads import (
79
+ LOCK_STRIPES,
80
+ CacheKey,
81
+ ImagePlan,
82
+ UploadCache,
83
+ UrlExpired,
84
+ cache_key,
85
+ check_storage_url,
86
+ plan_images,
87
+ presign_body,
88
+ presigned_entries,
89
+ storage_error,
90
+ stripe,
91
+ )
92
+ from ._validate import (
93
+ MAX_UPLOAD_FILES,
94
+ check_body_size,
95
+ check_collection,
96
+ check_enroll_image,
97
+ check_frame,
98
+ check_images,
99
+ check_item,
100
+ check_lines,
101
+ check_list_query,
102
+ check_request,
103
+ check_session_open,
104
+ count_entries,
105
+ usage_error,
106
+ )
107
+ from ._videos import VideoInput, check_videos, video_list
108
+
109
+ __all__ = ["DecisionNode", "AsyncDecisionNode", "DecideRequest", "DecisionResult", "BatchRequest"]
110
+
111
+ DecideRequest = Mapping[str, Any]
112
+ """The arguments of one ``decide`` call for ``decide_many``: ``state``, ``questions``,
113
+ and optionally ``images``, ``videos``, ``model``, ``boosters`` and ``fallback``."""
114
+
115
+ BatchRequest = Mapping[str, Any]
116
+ """One request of a batch: ``custom_id``, ``state``, ``questions``, and optionally
117
+ ``images``, ``videos``, ``model``, ``boosters`` and ``fallback``."""
118
+
119
+ DecisionResult = Union[Decision, DecisionNodeError]
120
+ """One result of ``decide_many``: the :class:`Decision`, or the error that request ended with."""
121
+
122
+ _A = TypeVar("_A", TruthAnswer, ChoiceAnswer, ScoreAnswer, NumberAnswer)
123
+ _Response = Tuple[int, Dict[str, str], str]
124
+ SESSION_KEYS = frozenset({"state", "questions", "instructions", "model", "window", "ttl_seconds"})
125
+ DEFAULT_LIMITS = httpx.Limits(max_connections=100, max_keepalive_connections=100, keepalive_expiry=30.0)
126
+
127
+
128
+ def _single(decision: Decision, question_id: str, kind: Type[_A], status: int, headers: Mapping[str, str]) -> _A:
129
+ """The helper's own answer, or :class:`ServerError` when it is missing or malformed.
130
+
131
+ ``check``, ``classify`` and ``score`` promise a typed answer, so an answer that does
132
+ not parse as one (``decide`` would return it as :class:`UnknownAnswer`) is an error.
133
+ """
134
+ answer: Optional[AnyAnswer] = decision.answers.get(question_id)
135
+ if not isinstance(answer, kind):
136
+ what = "no" if answer is None else "a malformed"
137
+ raise ServerError(
138
+ "the response has %s %s answer for %r" % (what, kind.__name__, question_id),
139
+ status=status,
140
+ request_id=decision.request_id,
141
+ reason=headers.get("x-decisionnode-error"),
142
+ body=decision.raw,
143
+ headers=headers,
144
+ )
145
+ return answer
146
+
147
+
148
+ # Transports known not to follow redirects themselves: in httpx redirects are the
149
+ # client's business, and the client is told never to follow one.
150
+ _SYNC_TRANSPORTS: Tuple[type, ...] = (httpx.HTTPTransport, httpx.MockTransport)
151
+ _ASYNC_TRANSPORTS: Tuple[type, ...] = (httpx.AsyncHTTPTransport, httpx.MockTransport)
152
+
153
+
154
+ def _check_http_client(client: Any, client_type: type, transports: Tuple[type, ...], base_url: str) -> None:
155
+ """Refuse a caller's HTTP client that cannot be shown to be safe, with ``ValueError``.
156
+
157
+ The client must be an ``httpx`` client of exactly the expected type (a subclass
158
+ could follow redirects whatever it is told), every transport it may send through
159
+ must be one of httpx's own, and for an ``https`` base URL none of them may have
160
+ certificate or hostname checks turned off.
161
+ """
162
+ if type(client) is not client_type:
163
+ raise ValueError(
164
+ "http_client must be an %s.%s, got %s: a client of another type may follow redirects"
165
+ % (client_type.__module__.split(".")[0], client_type.__name__, type(client).__name__)
166
+ )
167
+ candidates = [getattr(client, "_transport", None)]
168
+ candidates.extend(getattr(client, "_mounts", {}).values())
169
+ for transport in candidates:
170
+ if transport is None:
171
+ continue
172
+ if type(transport) not in transports:
173
+ raise ValueError(
174
+ "http_client uses a transport of type %s, which may follow redirects; use httpx's own transport"
175
+ % type(transport).__name__
176
+ )
177
+ if base_url.lower().startswith("https:") and not isinstance(transport, httpx.MockTransport):
178
+ ctx = getattr(getattr(transport, "_pool", None), "_ssl_context", None)
179
+ if not isinstance(ctx, ssl.SSLContext) or ctx.verify_mode != ssl.CERT_REQUIRED or not ctx.check_hostname:
180
+ raise ValueError(
181
+ "http_client turns off certificate or hostname checks, which an https base_url requires"
182
+ )
183
+
184
+
185
+ def _transport_error(exc: httpx.TransportError, timeout: float) -> DecisionNodeError:
186
+ if isinstance(exc, httpx.TimeoutException):
187
+ return _timeout_error(timeout)
188
+ return ConnectionError("connection error: %s" % (exc or exc.__class__.__name__))
189
+
190
+
191
+ def _timeout_error(timeout: float) -> DecisionNodeError:
192
+ return TimeoutError("no response within %ss" % ("%g" % timeout))
193
+
194
+
195
+ def _retry_without_answer(policy: str, error: DecisionNodeError) -> bool:
196
+ """Whether an attempt that got no answer (a timeout, a dropped connection) is retried.
197
+ A batch mutation never is; a decision naming ocr is not after a timeout, since the
198
+ server finishes that attempt and may bill it."""
199
+ if policy == MUTATION:
200
+ return False
201
+ return not (policy == OCR and isinstance(error, TimeoutError))
202
+
203
+
204
+ def _idempotency_key() -> Dict[str, str]:
205
+ """A new Idempotency-Key for one call: a random UUID (version 4), sent unchanged on
206
+ every attempt of that call."""
207
+ return {"Idempotency-Key": str(uuid.uuid4())}
208
+
209
+
210
+ def _ocr_send(settings: Settings, boosters: Any) -> Tuple[str, Optional[float]]:
211
+ """The retry policy and the attempt timeout of a decision: a decision whose boosters
212
+ name ocr waits OCR_TIMEOUT unless the caller set a timeout, and is never retried after
213
+ a timeout."""
214
+ if not names_ocr(boosters):
215
+ return DEFAULT, None
216
+ return OCR, None if settings.timeout_given else _base.OCR_TIMEOUT
217
+
218
+
219
+ class _Attempt:
220
+ """One blocking attempt, run on its own thread so the caller can stop waiting at the
221
+ deadline, whatever network step the attempt is in.
222
+
223
+ The caller gets a :class:`TimeoutError` as soon as the whole attempt (sending, the
224
+ response headers and the last byte of the body) has taken longer than the timeout.
225
+ An attempt given up on stops reading at its next chunk and closes its response;
226
+ every network step on its thread is bounded by the same timeout, so it ends soon after.
227
+ """
228
+
229
+ def __init__(
230
+ self,
231
+ http: httpx.Client,
232
+ method: str,
233
+ url: str,
234
+ content: Optional[bytes],
235
+ headers: Dict[str, str],
236
+ timeout: float,
237
+ request: Optional[httpx.Request] = None,
238
+ ) -> None:
239
+ self._args = (http, method, url, content, headers, timeout)
240
+ # A request built whole by the caller (the storage PUT, which carries only its
241
+ # presigned headers) is sent as it is.
242
+ self._request = request
243
+ self._done = threading.Event()
244
+ self._abandoned = threading.Event()
245
+ self._result: Optional[Tuple[httpx.Response, bytes]] = None
246
+ self._error: Optional[BaseException] = None
247
+
248
+ def run(self) -> Tuple[httpx.Response, bytes]:
249
+ timeout = self._args[5]
250
+ thread = threading.Thread(target=self._work, name="decisionnode-attempt", daemon=True)
251
+ thread.start()
252
+ if not self._done.wait(timeout):
253
+ self._abandoned.set()
254
+ raise _timeout_error(timeout)
255
+ if self._error is not None:
256
+ raise self._error
257
+ assert self._result is not None
258
+ return self._result
259
+
260
+ def _work(self) -> None:
261
+ http, method, url, content, headers, timeout = self._args
262
+ try:
263
+ if self._request is not None:
264
+ response = http.send(self._request, stream=True, follow_redirects=False)
265
+ else:
266
+ response = http.send(
267
+ http.build_request(method, url, content=content, headers=headers, timeout=timeout),
268
+ stream=True,
269
+ # Never followed, also on a client passed in: a redirect would send the
270
+ # key to another URL, possibly over plain HTTP.
271
+ follow_redirects=False,
272
+ )
273
+ try:
274
+ chunks: List[bytes] = []
275
+ for chunk in response.iter_bytes():
276
+ if self._abandoned.is_set():
277
+ return
278
+ chunks.append(chunk)
279
+ self._result = (response, b"".join(chunks))
280
+ finally:
281
+ response.close()
282
+ except BaseException as exc: # handed to the caller's thread
283
+ self._error = exc
284
+ finally:
285
+ self._done.set()
286
+
287
+
288
+ def _batch_path(batch_id: Any, suffix: str = "") -> str:
289
+ if not isinstance(batch_id, str) or not batch_id:
290
+ raise ValueError("batch_id must be a batch id such as batch_01J9ZB8Q3T, got %r" % (batch_id,))
291
+ return "%s/%s%s" % (_base.BATCHES_PATH, quote(batch_id, safe=""), suffix)
292
+
293
+
294
+ def _collection_path(collection_id: Any, suffix: str = "") -> str:
295
+ if not isinstance(collection_id, str) or not collection_id:
296
+ raise ValueError(
297
+ "collection_id must be a collection id such as col_... or a public collection's name, got %r"
298
+ % (collection_id,)
299
+ )
300
+ return "%s/%s%s" % (_base.COLLECTIONS_PATH, quote(collection_id, safe=""), suffix)
301
+
302
+
303
+ def _item_path(collection_id: Any, item_id: Any) -> str:
304
+ if not isinstance(item_id, str) or not item_id:
305
+ raise ValueError("item_id must be an item id such as itm_..., got %r" % (item_id,))
306
+ return _collection_path(collection_id, "/items/%s" % quote(item_id, safe=""))
307
+
308
+
309
+ def _collection_list(body: Any) -> List[Dict[str, Any]]:
310
+ found = body.get("collections") if isinstance(body, Mapping) else None
311
+ return list(found) if isinstance(found, list) else []
312
+
313
+
314
+ def _list_query(offset: Any, limit: Any) -> Dict[str, int]:
315
+ check_list_query(offset, limit)
316
+ query: Dict[str, int] = {}
317
+ for name, value in (("offset", offset), ("limit", limit)):
318
+ if value is None:
319
+ continue
320
+ if isinstance(value, bool) or not isinstance(value, int):
321
+ raise ValueError("%s must be a whole number, got %r" % (name, value))
322
+ query[name] = value
323
+ return query
324
+
325
+
326
+ def _session_body(
327
+ client: _Common, state: Any, questions: Any, instructions: Any, model: Optional[str], window: Any, ttl_seconds: Any
328
+ ) -> bytes:
329
+ qs = _base.check_question_map(questions)
330
+ body: Dict[str, Any] = {"model": model or client._settings.model}
331
+ if instructions is not None:
332
+ body["instructions"] = instructions
333
+ if state is not UNSET and state is not None:
334
+ body["state"] = state
335
+ body["questions"] = qs
336
+ if window is not None:
337
+ body["window"] = window
338
+ if ttl_seconds is not None:
339
+ body["ttl_seconds"] = ttl_seconds
340
+ check_session_open(body, client._limits(body["model"]))
341
+ return _base.encode(body)
342
+
343
+
344
+ def _check_supervise(body: Any) -> Mapping[str, Any]:
345
+ if not isinstance(body, Mapping):
346
+ raise ValueError("body must be a mapping of open_session() arguments")
347
+ unknown = set(body) - SESSION_KEYS
348
+ if unknown:
349
+ raise ValueError("body has unknown keys %s" % ", ".join(sorted(map(str, unknown))))
350
+ _base.check_question_map(body.get("questions"))
351
+ return body
352
+
353
+
354
+ class _Prepared:
355
+ """One decide-like request, checked and planned before anything is sent.
356
+
357
+ ``content`` is the final body when no image has to be uploaded; otherwise the body
358
+ is written once the uploads are done.
359
+ """
360
+
361
+ __slots__ = ("path", "state", "questions", "images", "videos", "model", "boosters", "fallback", "plan", "content")
362
+
363
+ def __init__(
364
+ self,
365
+ path: str,
366
+ state: Any,
367
+ questions: Dict[str, Any],
368
+ images: List[Any],
369
+ model: Optional[str],
370
+ videos: List[Dict[str, Any]],
371
+ boosters: Any = None,
372
+ fallback: Optional[bool] = None,
373
+ ) -> None:
374
+ self.path = path
375
+ self.state = state
376
+ self.questions = questions
377
+ self.images = images
378
+ self.videos = videos
379
+ self.model = model
380
+ # The boosters field as sent, or None for none; the fallback option as given.
381
+ self.boosters = boosters
382
+ self.fallback = fallback
383
+ self.plan: Optional[ImagePlan] = None
384
+ self.content: Optional[bytes] = None
385
+
386
+
387
+ class _Line:
388
+ """One batch request, checked and planned. ``content`` is the line as sent when none of
389
+ its images has to be uploaded; otherwise it is written once the uploads are done."""
390
+
391
+ __slots__ = ("req", "questions", "videos", "boosters", "plan", "content")
392
+
393
+ def __init__(
394
+ self, req: Any, questions: Dict[str, Any], videos: List[Dict[str, Any]], boosters: Any, plan: ImagePlan
395
+ ) -> None:
396
+ self.req = req
397
+ self.questions = questions
398
+ self.videos = videos
399
+ self.boosters = boosters
400
+ self.plan = plan
401
+ self.content: Optional[bytes] = None
402
+
403
+ def encode(self, wires: Sequence[Mapping[str, Any]], model: str) -> bytes:
404
+ return line_bytes(self.req, self.questions, wires, model, self.videos, self.boosters)
405
+
406
+ def measure(self, model: str) -> None:
407
+ """The API's body limit on the line as sent (its custom_id included, its line
408
+ break not), its uploads named by a placeholder."""
409
+ if self.plan.uploads:
410
+ check_body_size(len(self.encode(self.plan.project(), model)) - 1)
411
+ else:
412
+ self.content = self.encode(self.plan.finish({}), model)
413
+ check_body_size(len(self.content) - 1)
414
+
415
+
416
+ class _Common:
417
+ _settings: Settings
418
+ _cache: UploadCache
419
+ _http: Any
420
+ # The time uploads, their URLs and cache entries are judged by. Tests replace it.
421
+ _clock: Callable[[], float]
422
+ # The models() entries by name, once models() has been called: the limits only
423
+ # /v1/models states are checked locally from then on.
424
+ _held: Dict[str, Dict[str, Any]]
425
+
426
+ def _limits(self, model: Optional[str] = None) -> Dict[str, Any]:
427
+ """The held ``models()`` entry of the model a request names: by the entry's name,
428
+ else by the entry that states the name as its alias or pinned version. A name the
429
+ held list does not state holds nothing, and its checks are the server's."""
430
+ name = model or self._settings.model
431
+ entry = self._held.get(name)
432
+ if entry is None:
433
+ entry = next(
434
+ (e for e in self._held.values() if name in (e.get("alias"), e.get("version"))),
435
+ None,
436
+ )
437
+ return dict(entry, model=name) if entry else {}
438
+
439
+ def _hold(self, entries: Any) -> None:
440
+ self._held = {
441
+ m["name"]: dict(m) for m in entries or [] if isinstance(m, Mapping) and isinstance(m.get("name"), str)
442
+ }
443
+
444
+ def _plan(self, images: Sequence[Any], strip_id: bool = False) -> ImagePlan:
445
+ return plan_images(images, self._cache, self._settings.inline_threshold, strip_id=strip_id, now=self._clock())
446
+
447
+ def _check_frame(self, session: Any, frame: Any, image: Any, questions: Any, reset: Any, seq: Any) -> Any:
448
+ """A frame's fixed rules and its image's, before anything is sent. Returns the
449
+ questions as a list."""
450
+ if questions is not None:
451
+ if isinstance(questions, (str, bytes)) or not isinstance(questions, Iterable):
452
+ raise ValueError("questions must be a list of the session's question ids")
453
+ questions = list(questions)
454
+ known = session.raw.get("questions")
455
+ check_frame(frame, image, questions, reset, seq, known if isinstance(known, list) else None)
456
+ if image is not None:
457
+ probe = dict(image, id=image.get("id") or "frame") if isinstance(image, Mapping) else image
458
+ check_images([probe], self._limits(session.model), id_length=False)
459
+ return questions
460
+
461
+ @property
462
+ def api_key(self) -> Optional[str]:
463
+ return self._settings.api_key
464
+
465
+ @property
466
+ def base_url(self) -> str:
467
+ return self._settings.base_url
468
+
469
+ @property
470
+ def model(self) -> str:
471
+ """The model sent when a call names none."""
472
+ return self._settings.model
473
+
474
+ @property
475
+ def timeout(self) -> float:
476
+ """Seconds one whole attempt may take, from sending to the last byte of the body."""
477
+ return self._settings.timeout
478
+
479
+ @property
480
+ def max_retries(self) -> int:
481
+ return self._settings.max_retries
482
+
483
+ @property
484
+ def upload_timeout(self) -> float:
485
+ """Seconds one storage PUT attempt of an image may take."""
486
+ return self._settings.upload_timeout
487
+
488
+ @property
489
+ def inline_threshold(self) -> int:
490
+ """Local images of one request totalling fewer bytes than this go inline."""
491
+ return self._settings.inline_threshold
492
+
493
+ def _prepare(
494
+ self,
495
+ state: Any,
496
+ questions: Any,
497
+ images: Any,
498
+ model: Optional[str],
499
+ path: str,
500
+ videos: Any = None,
501
+ boosters: Any = None,
502
+ fallback: Optional[bool] = None,
503
+ ) -> _Prepared:
504
+ """Check a request against the API's fixed rules and plan its images, all before
505
+ any request. Argument problems raise ``ValueError``; content the API would refuse
506
+ raises its error class."""
507
+ qs, imgs, vids, bsts = self._check_decision(
508
+ state, questions, images, model, path, videos=videos, boosters=boosters
509
+ )
510
+ prep = _Prepared(path, state, qs, imgs, model, vids, bsts, _base.check_fallback(fallback))
511
+ plan = self._plan(imgs)
512
+ prep.plan = plan
513
+ if plan.uploads:
514
+ # Written once the uploads are done; measured now with the uploads named by a
515
+ # placeholder, so a body over the limit, or a value that cannot be sent as
516
+ # JSON, fails before anything is uploaded.
517
+ check_body_size(len(self._body(prep, plan.project())))
518
+ else:
519
+ prep.content = self._body(prep, plan.finish({}))
520
+ check_body_size(len(prep.content))
521
+ return prep
522
+
523
+ def _check_decision(
524
+ self,
525
+ state: Any,
526
+ questions: Any,
527
+ images: Any,
528
+ model: Optional[str],
529
+ path: str,
530
+ videos: Any = None,
531
+ boosters: Any = None,
532
+ ) -> Tuple[Dict[str, Any], List[Any], List[Dict[str, Any]], Any]:
533
+ """Every local check of a decide-like request; the questions as sent, the images,
534
+ the videos and the boosters in wire form. Boosters are the server's to judge: they
535
+ count only toward the nesting."""
536
+ qs = _base.check_question_map(questions)
537
+ imgs = _base.image_list(images)
538
+ vids = video_list(videos)
539
+ bsts = booster_wire(boosters)
540
+ # The body as it will be sent, its images and videos as placeholders so the
541
+ # nesting counts them.
542
+ probe = _base.decide_body(
543
+ state, qs, [{"id": ""} for _ in imgs], model, self._settings.model, [{"id": ""} for _ in vids], bsts
544
+ )
545
+ limits = self._limits(probe["model"])
546
+ check_request(probe, limits, image_ids=_image_ids(imgs), tail=count_entries(imgs, vids, limits))
547
+ check_images(imgs, limits)
548
+ check_videos(vids, imgs)
549
+ return qs, imgs, vids, bsts
550
+
551
+ def _body(self, prep: _Prepared, wires: Sequence[Mapping[str, Any]]) -> bytes:
552
+ return _base.encode(
553
+ _base.decide_body(
554
+ prep.state,
555
+ prep.questions,
556
+ wires,
557
+ prep.model,
558
+ self._settings.model,
559
+ prep.videos,
560
+ prep.boosters,
561
+ prep.fallback,
562
+ )
563
+ )
564
+
565
+ def _prepare_many(self, requests: Iterable[DecideRequest], concurrency: int) -> List[_Prepared]:
566
+ # Every request is checked before the first is sent, so a bad one fails the call
567
+ # and nothing is sent.
568
+ out = []
569
+ for i, req in enumerate(_base.check_many(requests, concurrency)):
570
+ try:
571
+ out.append(
572
+ self._prepare(
573
+ req.get("state", UNSET),
574
+ req.get("questions"),
575
+ req.get("images"),
576
+ req.get("model"),
577
+ _base.DECIDE_PATH,
578
+ req.get("videos"),
579
+ req.get("boosters"),
580
+ req.get("fallback"),
581
+ )
582
+ )
583
+ except ValueError as exc:
584
+ raise ValueError("decide_many: request %d: %s" % (i, exc)) from exc
585
+ return out
586
+
587
+ def _stream_options(self, url: str) -> Dict[str, Any]:
588
+ """The proxy and TLS settings the HTTP client would use for the stream's host,
589
+ for the session's WebSocket: the same proxy (or none), the same certificates."""
590
+ parts = urlsplit(url)
591
+ secure = parts.scheme.lower() == "wss"
592
+ http_url = ("https" if secure else "http") + url[len(parts.scheme) :]
593
+ try:
594
+ transport = self._http._transport_for_url(httpx.URL(http_url))
595
+ except Exception: # noqa: BLE001 - an HTTP client without per-URL transports
596
+ transport = getattr(self._http, "_transport", None)
597
+ pool = getattr(transport, "_pool", None)
598
+ options: Dict[str, Any] = {"proxy": _proxy_url(pool)}
599
+ ctx = getattr(pool, "_ssl_context", None)
600
+ if secure and isinstance(ctx, ssl.SSLContext):
601
+ options["ssl"] = ctx
602
+ proxy_ctx = getattr(pool, "_proxy_ssl_context", None)
603
+ if options["proxy"] is not None and isinstance(proxy_ctx, ssl.SSLContext):
604
+ options["proxy_ssl"] = proxy_ctx
605
+ return options
606
+
607
+ def check_request(self, method: str, **arguments: Any) -> None:
608
+ """Run every local check the call ``method`` would run before sending, with the
609
+ same arguments, and raise the same error; send and upload nothing.
610
+
611
+ ``method`` is ``"decide"``, ``"open_session"``, ``"batches.create"`` or
612
+ ``"batches.add"`` (its ``batch_id`` plays no part).
613
+ Limits held from an earlier ``models()`` call are checked too. Returns ``None``
614
+ when the request would be sent. Never makes a request, also on the async client.
615
+ """
616
+ if method == "decide":
617
+ unknown = set(arguments) - _base.DECIDE_KEYS
618
+ if unknown:
619
+ raise TypeError("%s() got unexpected arguments %s" % (method, ", ".join(sorted(unknown))))
620
+ # Prepared as decide prepares it (planned, encoded and measured), never sent.
621
+ self._prepare(
622
+ arguments.get("state", UNSET),
623
+ arguments.get("questions"),
624
+ arguments.get("images"),
625
+ arguments.get("model"),
626
+ _base.DECIDE_PATH,
627
+ arguments.get("videos"),
628
+ arguments.get("boosters"),
629
+ arguments.get("fallback"),
630
+ )
631
+ elif method == "open_session":
632
+ unknown = set(arguments) - SESSION_KEYS
633
+ if unknown:
634
+ raise TypeError("open_session() got unexpected arguments %s" % ", ".join(sorted(unknown)))
635
+ _session_body(
636
+ self,
637
+ arguments.get("state", UNSET),
638
+ arguments.get("questions"),
639
+ arguments.get("instructions"),
640
+ arguments.get("model"),
641
+ arguments.get("window"),
642
+ arguments.get("ttl_seconds"),
643
+ )
644
+ elif method in ("batches.create", "batches.add"):
645
+ arguments.pop("batch_id", None)
646
+ unknown = set(arguments) - {"requests"}
647
+ if unknown:
648
+ raise TypeError("%s() got unexpected arguments %s" % (method, ", ".join(sorted(unknown))))
649
+ self._check_lines(list(arguments.get("requests") or ()), method == "batches.create")
650
+ else:
651
+ raise ValueError(
652
+ "check_request() takes decide, open_session, batches.create or batches.add, got %r" % (method,)
653
+ )
654
+
655
+ def _check_lines(self, items: List[Any], allow_empty: bool) -> List[_Line]:
656
+ """The local checks of a batch call's requests, before anything is sent or uploaded:
657
+ each line's fixed rules, then its size as it will be sent."""
658
+ check_lines(items, allow_empty, self._limits())
659
+ checked = []
660
+ for i, req in enumerate(items, start=1):
661
+ try:
662
+ qs = _base.check_question_map(req.get("questions"))
663
+ imgs = _base.image_list(req.get("images"))
664
+ vids = video_list(req.get("videos"))
665
+ _base.check_fallback(req.get("fallback"))
666
+ except ValueError as exc:
667
+ raise ValueError("batch request %d: %s" % (i, exc)) from exc
668
+ bsts = booster_wire(req.get("boosters"))
669
+ model = req.get("model") or self._settings.model
670
+ probe: Dict[str, Any] = {"model": model, "state": req.get("state"), "questions": qs}
671
+ if imgs:
672
+ probe["images"] = [{"id": ""} for _ in imgs]
673
+ if vids:
674
+ probe["videos"] = [{"id": ""} for _ in vids]
675
+ if bsts is not None:
676
+ probe["boosters"] = bsts
677
+ limits = self._limits(model)
678
+ check_request(probe, limits, image_ids=_image_ids(imgs), tail=count_entries(imgs, vids, limits))
679
+ check_images(imgs, limits)
680
+ check_videos(vids, imgs)
681
+ line = _Line(req, qs, vids, bsts, self._plan(imgs))
682
+ line.measure(self._settings.model)
683
+ checked.append(line)
684
+ return checked
685
+
686
+ def __repr__(self) -> str:
687
+ return "%s(%r)" % (self.__class__.__name__, self._settings)
688
+
689
+
690
+ def _upload_list(images: Any) -> List[Any]:
691
+ """The images of ``upload``: one presign takes 1 to 16 files, so more is that call's
692
+ own 400, raised before anything is sent."""
693
+ imgs = _base.image_list(images)
694
+ if len(imgs) > MAX_UPLOAD_FILES:
695
+ raise usage_error("files: 1 to %d files, got %d." % (MAX_UPLOAD_FILES, len(imgs)))
696
+ return imgs
697
+
698
+
699
+ def _image_ids(images: Sequence[Any]) -> List[Any]:
700
+ """The ids of a request's images as given, for the point and box rules."""
701
+ return [img.id if isinstance(img, Image) else img.get("id") for img in images]
702
+
703
+
704
+ def _proxy_url(pool: Any) -> Optional[str]:
705
+ """The proxy an httpcore pool sends through, as a URL with its credentials, or ``None``."""
706
+ url = getattr(pool, "_proxy_url", None)
707
+ if url is None:
708
+ return None
709
+ try:
710
+ scheme = url.scheme.decode("ascii")
711
+ host = url.host.decode("ascii")
712
+ except (AttributeError, UnicodeDecodeError):
713
+ return None
714
+ if ":" in host:
715
+ host = "[%s]" % host
716
+ userinfo = ""
717
+ for name, value in getattr(pool, "_proxy_headers", None) or []:
718
+ if name.lower() == b"proxy-authorization" and value[:6].lower() == b"basic ":
719
+ try:
720
+ user, _, password = base64.b64decode(value[6:]).decode("utf-8").partition(":")
721
+ except (ValueError, UnicodeDecodeError):
722
+ continue
723
+ userinfo = "%s:%s@" % (quote(user, safe=""), quote(password, safe=""))
724
+ port = ":%d" % url.port if url.port else ""
725
+ return "%s://%s%s%s" % (scheme, userinfo, host, port)
726
+
727
+
728
+ def _http2_available() -> None:
729
+ try:
730
+ import h2 # noqa: F401
731
+ except ImportError:
732
+ raise ValueError('http2=True needs the h2 package: pip install "bynn-decisionnode[http2]"') from None
733
+
734
+
735
+ def _url(settings: Settings, path: str, query: Optional[Mapping[str, Any]]) -> str:
736
+ url = settings.url(path)
737
+ if query:
738
+ url += "?" + urlencode(query)
739
+ return url
740
+
741
+
742
+ def _storage_path(url: str) -> str:
743
+ # Logged without the query, which holds the signature.
744
+ try:
745
+ return urlsplit(url).path
746
+ except ValueError:
747
+ return "storage"
748
+
749
+
750
+ class DecisionNode(_Common):
751
+ """The DecisionNode client.
752
+
753
+ Every argument is optional, and an empty string counts as not given::
754
+
755
+ dn = DecisionNode() # reads DECISIONNODE_API_KEY, DECISIONNODE_BASE_URL, DECISIONNODE_MODEL
756
+
757
+ Args:
758
+ api_key: The API key. Default: ``DECISIONNODE_API_KEY``.
759
+ base_url: The API's root URL. Default: ``DECISIONNODE_BASE_URL``, else
760
+ ``https://api.decisionnode.com``.
761
+ model: The model sent when a call names none. Default: ``DECISIONNODE_MODEL``,
762
+ else ``decisionnode-moe``.
763
+ timeout: Seconds one whole attempt may take, from sending the request to the
764
+ last byte of the response (default 10; 30 for a decision whose boosters name
765
+ ``ocr``). Must be above 0. A timeout you set is used as given on every call.
766
+ max_retries: Retries after the first attempt (default 2). Must be 0 or more.
767
+ headers: Extra headers sent on every API request and session handshake (never
768
+ to image storage).
769
+ http_client: An ``httpx.Client`` to send with, for proxies or custom TLS. The
770
+ client does not close a client you pass in. Redirects are never followed,
771
+ whatever the ``httpx.Client`` is set to; a subclass, a transport of another
772
+ type, or turned-off certificate checks for an ``https`` base URL are refused.
773
+ inline_threshold: Local images of one request that total fewer bytes than this
774
+ go inline as base64; at or above it they are uploaded first (default
775
+ ``DEFAULT_INLINE_THRESHOLD``, 524,288).
776
+ upload_timeout: Seconds one storage PUT attempt of an image may take, apart from
777
+ ``timeout`` (default 60). Must be above 0.
778
+ http2: Use HTTP/2 for the API connection (needs the ``http2`` extra). Ignored
779
+ with an ``http_client``, which keeps its own setting.
780
+
781
+ Raises:
782
+ ValueError: ``base_url`` is not ``https`` (plain ``http`` is accepted only for
783
+ ``localhost``, ``127.0.0.1`` and ``::1``), ``timeout`` is not above 0,
784
+ ``max_retries`` or ``inline_threshold`` is below 0, or ``http_client`` is not
785
+ an ``httpx.Client`` with httpx's own transports, or turns off certificate
786
+ checks for an ``https`` base URL.
787
+
788
+ A client is safe to share between threads, keeps its connections open between
789
+ calls, and remembers its image uploads for reuse. Close it with :meth:`close` or use
790
+ it as a context manager.
791
+ """
792
+
793
+ def __init__(
794
+ self,
795
+ api_key: Optional[str] = None,
796
+ base_url: Optional[str] = None,
797
+ model: Optional[str] = None,
798
+ timeout: Optional[float] = None,
799
+ max_retries: Optional[int] = _base.DEFAULT_MAX_RETRIES,
800
+ headers: Optional[Mapping[str, str]] = None,
801
+ http_client: Optional[httpx.Client] = None,
802
+ inline_threshold: Optional[int] = None,
803
+ http2: bool = False,
804
+ upload_timeout: Optional[float] = _base.DEFAULT_UPLOAD_TIMEOUT,
805
+ ) -> None:
806
+ self._settings = Settings(
807
+ api_key, base_url, model, timeout, max_retries, headers, inline_threshold, upload_timeout
808
+ )
809
+ if http_client is not None:
810
+ _check_http_client(http_client, httpx.Client, _SYNC_TRANSPORTS, self._settings.base_url)
811
+ elif http2:
812
+ _http2_available()
813
+ self._owns_http = http_client is None
814
+ self._http = http_client or httpx.Client(
815
+ timeout=httpx.Timeout(self._settings.timeout),
816
+ cookies=_base.no_cookies(),
817
+ follow_redirects=False,
818
+ limits=DEFAULT_LIMITS,
819
+ http2=bool(http2),
820
+ )
821
+ self._cache = UploadCache()
822
+ self._upload_locks = [threading.Lock() for _ in range(LOCK_STRIPES)]
823
+ # The wait between attempts, PUT retries, batch polls and session reopens. Tests
824
+ # replace it to record the waits without sleeping.
825
+ self._sleep: Callable[[float], None] = time.sleep
826
+ self._clock = time.time
827
+ self._held = {}
828
+ self.batches = Batches(self)
829
+ self.collections = Collections(self)
830
+
831
+ # ---- lifecycle
832
+
833
+ def close(self) -> None:
834
+ """Close the connections, unless the ``httpx.Client`` was passed in."""
835
+ if self._owns_http:
836
+ self._http.close()
837
+
838
+ def __enter__(self) -> DecisionNode:
839
+ return self
840
+
841
+ def __exit__(
842
+ self,
843
+ exc_type: Optional[Type[BaseException]],
844
+ exc: Optional[BaseException],
845
+ tb: Optional[TracebackType],
846
+ ) -> None:
847
+ self.close()
848
+
849
+ # ---- transport
850
+
851
+ def _send(
852
+ self,
853
+ method: str,
854
+ path: str,
855
+ content: Optional[bytes],
856
+ policy: str = DEFAULT,
857
+ content_type: str = "application/json",
858
+ query: Optional[Mapping[str, Any]] = None,
859
+ timeout: Optional[float] = None,
860
+ extra_headers: Optional[Mapping[str, str]] = None,
861
+ ) -> _Response:
862
+ settings = self._settings
863
+ headers = settings.request_headers(content is not None, content_type)
864
+ if extra_headers:
865
+ # The same on every attempt: an Idempotency-Key is the call's, not the attempt's.
866
+ headers.update(extra_headers)
867
+ url = _url(settings, path, query)
868
+ limit = timeout or settings.timeout
869
+ attempt = 0
870
+ while True:
871
+ started = time.monotonic()
872
+ resp_headers: Optional[Dict[str, str]] = None
873
+ try:
874
+ response, raw = _Attempt(self._http, method, url, content, headers, limit).run()
875
+ except (httpx.TransportError, TimeoutError) as exc:
876
+ error = _transport_error(exc, limit) if isinstance(exc, httpx.TransportError) else exc
877
+ _base.log_response(method, path, type(error).__name__, None, started, attempt)
878
+ retry = _retry_without_answer(policy, error)
879
+ else:
880
+ resp_headers = _base.lower_headers(response.headers)
881
+ status = response.status_code
882
+ text = raw.decode("utf-8", errors="replace")
883
+ _base.log_response(method, path, status, resp_headers.get("x-request-id"), started, attempt)
884
+ if 200 <= status < 300:
885
+ return status, resp_headers, text
886
+ error = _base.make_error(status, resp_headers, text)
887
+ retry = _base.should_retry(policy, status, error)
888
+ if not retry or attempt >= settings.max_retries:
889
+ raise error
890
+ self._sleep(_base.retry_wait(attempt, resp_headers))
891
+ attempt += 1
892
+
893
+ def _json(
894
+ self,
895
+ method: str,
896
+ path: str,
897
+ content: Optional[bytes],
898
+ policy: str = DEFAULT,
899
+ content_type: str = "application/json",
900
+ query: Optional[Mapping[str, Any]] = None,
901
+ extra_headers: Optional[Mapping[str, str]] = None,
902
+ ) -> Tuple[int, Dict[str, str], Any]:
903
+ status, headers, text = self._send(
904
+ method, path, content, policy, content_type, query, extra_headers=extra_headers
905
+ )
906
+ return status, headers, _base.parse_json(status, headers, text)
907
+
908
+ def _open_stream(self, path: str) -> httpx.Response:
909
+ """A ``GET`` whose body is read by the caller as it arrives. Retried as any read;
910
+ each network step is bounded by the timeout, not the whole body."""
911
+ settings = self._settings
912
+ headers = settings.request_headers(False)
913
+ url = settings.url(path)
914
+ attempt = 0
915
+ while True:
916
+ started = time.monotonic()
917
+ resp_headers: Optional[Dict[str, str]] = None
918
+ try:
919
+ request = self._http.build_request("GET", url, headers=headers, timeout=settings.timeout)
920
+ response = self._http.send(request, stream=True, follow_redirects=False)
921
+ except httpx.TransportError as exc:
922
+ error = _transport_error(exc, settings.timeout)
923
+ _base.log_response("GET", path, type(error).__name__, None, started, attempt)
924
+ retry = True
925
+ else:
926
+ resp_headers = _base.lower_headers(response.headers)
927
+ status = response.status_code
928
+ _base.log_response("GET", path, status, resp_headers.get("x-request-id"), started, attempt)
929
+ if 200 <= status < 300:
930
+ return response # type: ignore[no-any-return]
931
+ try:
932
+ text = response.read().decode("utf-8", errors="replace")
933
+ except httpx.TransportError:
934
+ text = ""
935
+ finally:
936
+ response.close()
937
+ error = _base.make_error(status, resp_headers, text)
938
+ retry = status in _base.RETRY_STATUSES
939
+ if not retry or attempt >= settings.max_retries:
940
+ raise error
941
+ self._sleep(_base.retry_wait(attempt, resp_headers))
942
+ attempt += 1
943
+
944
+ # ---- uploads
945
+
946
+ def _put(
947
+ self, img: Image, entry: Mapping[str, Any], presigned_at: float, url_ttl: float, may_presign: bool
948
+ ) -> None:
949
+ """One storage PUT with exactly the presigned headers and the bytes: no key, none
950
+ of the client's headers, no redirect followed."""
951
+ settings = self._settings
952
+ path = _storage_path(entry["url"])
953
+ attempt = 0
954
+ while True:
955
+ started = time.monotonic()
956
+ resp_headers: Optional[Dict[str, str]] = None
957
+ request = httpx.Request(
958
+ "PUT",
959
+ entry["url"],
960
+ headers=entry["headers"],
961
+ content=img.data,
962
+ extensions={"timeout": httpx.Timeout(settings.upload_timeout).as_dict()},
963
+ )
964
+ try:
965
+ # Its own deadline, for the whole attempt: the upload_timeout.
966
+ response, raw = _Attempt(
967
+ self._http, "PUT", entry["url"], None, {}, settings.upload_timeout, request=request
968
+ ).run()
969
+ except (httpx.TransportError, TimeoutError) as exc:
970
+ error = _transport_error(exc, settings.upload_timeout) if isinstance(exc, httpx.TransportError) else exc
971
+ _base.log_response("PUT", path, type(error).__name__, None, started, attempt)
972
+ retry = True
973
+ else:
974
+ status = response.status_code
975
+ resp_headers = _base.lower_headers(response.headers)
976
+ _base.log_response("PUT", path, status, None, started, attempt)
977
+ # A 412 on a retry: an earlier attempt stored the object, which is write-once.
978
+ if 200 <= status < 300 or (status == 412 and attempt > 0):
979
+ return
980
+ if status == 403 and may_presign and self._clock() - presigned_at >= url_ttl:
981
+ raise UrlExpired()
982
+ error = storage_error(img.id, status, resp_headers, raw.decode("utf-8", errors="replace"))
983
+ retry = 500 <= status <= 599
984
+ if not retry or attempt >= settings.max_retries:
985
+ raise error
986
+ self._sleep(_base.retry_wait(attempt, resp_headers))
987
+ attempt += 1
988
+
989
+ def _presign_and_put(self, todo: List[Image], may_presign: bool = True) -> Dict[CacheKey, str]:
990
+ _, _, body = self._json("POST", _base.UPLOADS_PATH, _base.encode(presign_body(todo)))
991
+ entries = presigned_entries(body, len(todo))
992
+ for entry in entries:
993
+ check_storage_url(entry["url"])
994
+ url_ttl = float(body.get("expires_in") or 300)
995
+ presigned_at = self._clock()
996
+ out: Dict[CacheKey, str] = {}
997
+ expired: List[Image] = []
998
+ errors: List[BaseException] = []
999
+
1000
+ def one(img: Image, entry: Mapping[str, Any]) -> None:
1001
+ try:
1002
+ self._put(img, entry, presigned_at, url_ttl, may_presign)
1003
+ except UrlExpired:
1004
+ expired.append(img)
1005
+ return
1006
+ except BaseException as exc: # handed to the caller below
1007
+ errors.append(exc)
1008
+ return
1009
+ key = cache_key(img)
1010
+ self._cache.put(key, entry["upload_id"], _base.parse_time(entry.get("expires_at")))
1011
+ out[key] = entry["upload_id"]
1012
+
1013
+ pairs = list(zip(todo, entries))
1014
+ if len(pairs) == 1:
1015
+ one(*pairs[0])
1016
+ else:
1017
+ with ThreadPoolExecutor(max_workers=len(pairs), thread_name_prefix="decisionnode-put") as pool:
1018
+ list(pool.map(lambda pair: one(*pair), pairs))
1019
+ if errors:
1020
+ raise errors[0]
1021
+ if expired:
1022
+ # The URLs' minutes passed before the PUT: presign those once more.
1023
+ out.update(self._presign_and_put(expired, may_presign=False))
1024
+ return out
1025
+
1026
+ def _upload(self, images: Sequence[Image]) -> Dict[CacheKey, str]:
1027
+ """The upload id of each image's bytes: from the cache, else one presign for all
1028
+ the rest and a PUT each, in parallel. Concurrent calls with the same bytes upload
1029
+ them once."""
1030
+ keys = [cache_key(img) for img in images]
1031
+ locks = [self._upload_locks[i] for i in sorted({stripe(k) for k in keys})]
1032
+ for lock in locks:
1033
+ lock.acquire()
1034
+ try:
1035
+ out: Dict[CacheKey, str] = {}
1036
+ todo: List[Image] = []
1037
+ for img, key in zip(images, keys):
1038
+ hit = self._cache.get(key, self._clock())
1039
+ if hit is not None:
1040
+ out[key] = hit
1041
+ elif key not in out and all(cache_key(t) != key for t in todo):
1042
+ todo.append(img)
1043
+ if todo:
1044
+ out.update(self._presign_and_put(todo))
1045
+ return out
1046
+ finally:
1047
+ for lock in reversed(locks):
1048
+ lock.release()
1049
+
1050
+ def _finish(self, plan: ImagePlan) -> List[Dict[str, Any]]:
1051
+ return plan.finish(self._upload(plan.uploads) if plan.uploads else {})
1052
+
1053
+ def _frame_image(self, image: Any) -> Dict[str, Any]:
1054
+ plan = self._plan([image], strip_id=True)
1055
+ return self._finish(plan)[0]
1056
+
1057
+ # ---- decisions
1058
+
1059
+ def _execute(self, prep: _Prepared) -> _Response:
1060
+ """Send a prepared request. A stale reference to an upload this client placed is
1061
+ replaced once: drop it, upload again, send again. That is not a retry."""
1062
+ plan = prep.plan
1063
+ assert plan is not None
1064
+ content = prep.content
1065
+ for attempt in (0, 1):
1066
+ if content is None:
1067
+ content = self._body(prep, self._finish(plan))
1068
+ try:
1069
+ policy, timeout = _ocr_send(self._settings, prep.boosters)
1070
+ return self._send("POST", prep.path, content, policy, timeout=timeout)
1071
+ except BadRequestError as exc:
1072
+ if attempt == 0 and plan.placed and exc.error_type in _base.STALE_UPLOAD:
1073
+ self._cache.drop(plan.placed)
1074
+ plan = self._plan(prep.images)
1075
+ content = None
1076
+ continue
1077
+ raise
1078
+ raise AssertionError("unreachable")
1079
+
1080
+ def _decision(self, prep: _Prepared) -> Decision:
1081
+ status, headers, text = self._execute(prep)
1082
+ return parse_decision(_base.parse_json(status, headers, text), headers)
1083
+
1084
+ def _ask(
1085
+ self,
1086
+ question_id: str,
1087
+ question: Question,
1088
+ kind: Type[_A],
1089
+ state: Any,
1090
+ images: Any,
1091
+ model: Any,
1092
+ videos: Any,
1093
+ boosters: Any,
1094
+ fallback: Optional[bool] = None,
1095
+ ) -> _A:
1096
+ prep = self._prepare(
1097
+ state, {question_id: question}, images, model, _base.DECIDE_PATH, videos, boosters, fallback
1098
+ )
1099
+ status, headers, text = self._execute(prep)
1100
+ decision = parse_decision(_base.parse_json(status, headers, text), headers)
1101
+ return _single(decision, question_id, kind, status, headers)
1102
+
1103
+ # ---- the surface
1104
+
1105
+ def decide(
1106
+ self,
1107
+ state: Any = UNSET,
1108
+ questions: Optional[Mapping[str, Question]] = None,
1109
+ images: Optional[Sequence[ImageInput]] = None,
1110
+ model: Optional[str] = None,
1111
+ *,
1112
+ videos: Optional[Sequence[VideoInput]] = None,
1113
+ boosters: Optional[Sequence[BoosterInput]] = None,
1114
+ fallback: Optional[bool] = None,
1115
+ ) -> Decision:
1116
+ """Ask one or more questions about a state in one request.
1117
+
1118
+ Args:
1119
+ state: What is judged: a string, or any JSON-serialisable value. Left out,
1120
+ the questions are answered from their own text.
1121
+ questions: Question id to question, built with :func:`truth`,
1122
+ :func:`choice`, :func:`score` and :func:`number` or written in wire form.
1123
+ images: Up to 16 :class:`Image` objects every question can see.
1124
+ videos: Up to 4 :class:`Video` objects (frames with their times) every question
1125
+ can see.
1126
+ model: The model for this call. Default (also for ``""``): the client's model.
1127
+ boosters: Booster names, or :class:`Booster` objects for a booster that needs
1128
+ a setting, read before the decision; their readings come back in
1129
+ ``evidence``.
1130
+ fallback: ``False``: a ``decisionnode-moe`` request that no worker can take at
1131
+ once gets the 529 :class:`OverloadedError` instead of an answer from
1132
+ Flash. Left out, the API's default (``True``) applies.
1133
+
1134
+ Returns:
1135
+ A :class:`Decision`. Read each answer by its id: ``decision["route"]``.
1136
+
1137
+ Raises:
1138
+ BadRequestError, ValidationError: before anything is sent, for a request the
1139
+ API would refuse by its fixed rules (an image, a question's limits).
1140
+ """
1141
+ return self._decision(
1142
+ self._prepare(state, questions, images, model, _base.DECIDE_PATH, videos, boosters, fallback)
1143
+ )
1144
+
1145
+ def decide_many(
1146
+ self,
1147
+ requests: Iterable[DecideRequest],
1148
+ concurrency: int = _base.DEFAULT_CONCURRENCY,
1149
+ ) -> List[DecisionResult]:
1150
+ """Run many ``decide`` calls, at most ``concurrency`` at a time (default 8).
1151
+
1152
+ Each request is a mapping of :meth:`decide` arguments. Every request is checked
1153
+ before the first is sent: a bad argument in any of them raises ``ValueError``, a
1154
+ request the API would refuse raises its error class, and nothing is sent. Returns
1155
+ one result per request, in input order: the :class:`Decision`, or the
1156
+ :class:`DecisionNodeError` that request ended with. One failure does not stop the
1157
+ others. Images with the same bytes are uploaded once.
1158
+ """
1159
+ preps = self._prepare_many(requests, concurrency)
1160
+ if not preps:
1161
+ return []
1162
+
1163
+ def run(prep: _Prepared) -> DecisionResult:
1164
+ try:
1165
+ return self._decision(prep)
1166
+ except DecisionNodeError as exc:
1167
+ return exc
1168
+
1169
+ if len(preps) == 1:
1170
+ return [run(preps[0])]
1171
+ with ThreadPoolExecutor(max_workers=min(concurrency, len(preps)), thread_name_prefix="decisionnode") as pool:
1172
+ return list(pool.map(run, preps))
1173
+
1174
+ def check(
1175
+ self,
1176
+ state: Any = UNSET,
1177
+ instructions: Any = None,
1178
+ if_true: Any = None,
1179
+ if_false: Any = None,
1180
+ *,
1181
+ images: Optional[Sequence[ImageInput]] = None,
1182
+ videos: Optional[Sequence[VideoInput]] = None,
1183
+ model: Optional[str] = None,
1184
+ boosters: Optional[Sequence[BoosterInput]] = None,
1185
+ fallback: Optional[bool] = None,
1186
+ ) -> TruthAnswer:
1187
+ """Ask one ``truth`` question. Returns its :class:`TruthAnswer`.
1188
+
1189
+ Raises :class:`ServerError` when that answer is missing or malformed.
1190
+ """
1191
+ return self._ask(
1192
+ "check",
1193
+ _truth(instructions, if_true, if_false),
1194
+ TruthAnswer,
1195
+ state,
1196
+ images,
1197
+ model,
1198
+ videos,
1199
+ boosters,
1200
+ fallback,
1201
+ )
1202
+
1203
+ def classify(
1204
+ self,
1205
+ state: Any = UNSET,
1206
+ instructions: Any = None,
1207
+ options: Union[Mapping[str, Any], Iterable[str], None] = None,
1208
+ *,
1209
+ images: Optional[Sequence[ImageInput]] = None,
1210
+ videos: Optional[Sequence[VideoInput]] = None,
1211
+ model: Optional[str] = None,
1212
+ boosters: Optional[Sequence[BoosterInput]] = None,
1213
+ fallback: Optional[bool] = None,
1214
+ ) -> ChoiceAnswer:
1215
+ """Ask one ``choice`` question. Returns its :class:`ChoiceAnswer`.
1216
+
1217
+ Raises :class:`ServerError` when that answer is missing or malformed.
1218
+ """
1219
+ return self._ask(
1220
+ "classify", _choice(instructions, options), ChoiceAnswer, state, images, model, videos, boosters, fallback
1221
+ )
1222
+
1223
+ def score(
1224
+ self,
1225
+ state: Any = UNSET,
1226
+ instructions: Any = None,
1227
+ levels: Optional[Sequence[Any]] = None,
1228
+ *,
1229
+ images: Optional[Sequence[ImageInput]] = None,
1230
+ videos: Optional[Sequence[VideoInput]] = None,
1231
+ model: Optional[str] = None,
1232
+ boosters: Optional[Sequence[BoosterInput]] = None,
1233
+ fallback: Optional[bool] = None,
1234
+ ) -> ScoreAnswer:
1235
+ """Ask one ``score`` question. Returns its :class:`ScoreAnswer`.
1236
+
1237
+ Raises :class:`ServerError` when that answer is missing or malformed.
1238
+ """
1239
+ return self._ask(
1240
+ "score", _score(instructions, levels), ScoreAnswer, state, images, model, videos, boosters, fallback
1241
+ )
1242
+
1243
+ def number(
1244
+ self,
1245
+ state: Any = UNSET,
1246
+ instructions: Any = None,
1247
+ min: Union[int, float, None] = None, # noqa: A002 - the wire name
1248
+ max: Union[int, float, None] = None, # noqa: A002 - the wire name
1249
+ step: Union[int, float, None] = None,
1250
+ *,
1251
+ images: Optional[Sequence[ImageInput]] = None,
1252
+ videos: Optional[Sequence[VideoInput]] = None,
1253
+ model: Optional[str] = None,
1254
+ boosters: Optional[Sequence[BoosterInput]] = None,
1255
+ fallback: Optional[bool] = None,
1256
+ ) -> NumberAnswer:
1257
+ """Ask one ``number`` question. Returns its :class:`NumberAnswer`.
1258
+
1259
+ Raises :class:`ServerError` when that answer is missing or malformed.
1260
+ """
1261
+ return self._ask(
1262
+ "number",
1263
+ _number(instructions, min, max, step),
1264
+ NumberAnswer,
1265
+ state,
1266
+ images,
1267
+ model,
1268
+ videos,
1269
+ boosters,
1270
+ fallback,
1271
+ )
1272
+
1273
+ def gate(
1274
+ self,
1275
+ action: Any,
1276
+ context: Any = None,
1277
+ *,
1278
+ images: Optional[Sequence[ImageInput]] = None,
1279
+ videos: Optional[Sequence[VideoInput]] = None,
1280
+ model: Optional[str] = None,
1281
+ boosters: Optional[Sequence[BoosterInput]] = None,
1282
+ fallback: Optional[bool] = None,
1283
+ ) -> Decision:
1284
+ """Judge an action before it runs.
1285
+
1286
+ Sends three questions about ``action`` (and ``context``, when given): ``risk``,
1287
+ a score from ``none`` to ``severe``; ``touches_production``, a truth; and
1288
+ ``verdict``, a choice of ``allow``, ``confirm`` or ``block``. What to do with
1289
+ the answers is up to your code.
1290
+ """
1291
+ return self.decide(
1292
+ _base.gate_state(action, context),
1293
+ _base.gate_questions(),
1294
+ images,
1295
+ model,
1296
+ videos=videos,
1297
+ boosters=boosters,
1298
+ fallback=fallback,
1299
+ )
1300
+
1301
+ def models(self) -> List[ModelInfo]:
1302
+ """The models this key may use, with the limits each one publishes."""
1303
+ _, _, body = self._json("GET", _base.MODELS_PATH, None)
1304
+ entries = body.get("models") if isinstance(body, Mapping) else None
1305
+ self._hold(entries)
1306
+ return [parse_model(m) for m in entries or []]
1307
+
1308
+ def upload(self, images: Sequence[ImageInput]) -> List[str]:
1309
+ """Upload local images whatever their size and return the ``upload_id`` of each,
1310
+ in order: one presign for the images not uploaded before, then a PUT each. An
1311
+ upload id can be sent as ``Image(id, upload_id=...)`` for the next hour."""
1312
+ imgs = _upload_list(images)
1313
+ check_images(imgs, self._limits())
1314
+ local = []
1315
+ for img in imgs:
1316
+ if not isinstance(img, Image) or (img.data is None and img.upload_id is None):
1317
+ raise ValueError("upload() takes Image objects made from a path, bytes or an upload_id")
1318
+ if img.data is not None:
1319
+ local.append(img)
1320
+ ids = self._upload(local) if local else {}
1321
+ return [img.upload_id if img.data is None else ids[cache_key(img)] for img in imgs]
1322
+
1323
+ def decide_raw(self, body: Any) -> Any:
1324
+ """Send ``body`` to ``POST /v1/decide`` exactly as given and return the parsed response body.
1325
+
1326
+ The way to use an API field this version of the SDK does not know yet. Nothing in
1327
+ it is checked or uploaded. Retries and errors work as for :meth:`decide`.
1328
+ """
1329
+ _, _, parsed = self._json("POST", _base.DECIDE_PATH, _base.encode(body))
1330
+ return parsed
1331
+
1332
+ # ---- sessions
1333
+
1334
+ def open_session(
1335
+ self,
1336
+ state: Any = UNSET,
1337
+ questions: Optional[Mapping[str, Question]] = None,
1338
+ instructions: Any = None,
1339
+ model: Optional[str] = None,
1340
+ window: Optional[int] = None,
1341
+ ttl_seconds: Optional[int] = None,
1342
+ ) -> Session:
1343
+ """Open a session: the questions are fixed, then every frame sent on its stream is
1344
+ answered with them. Returns a :class:`Session`; call :meth:`Session.connect`.
1345
+
1346
+ Raises:
1347
+ SessionLimitError: the key holds as many open sessions as it may (never retried).
1348
+ """
1349
+ content = _session_body(self, state, questions, instructions, model, window, ttl_seconds)
1350
+ _, headers, body = self._json("POST", _base.SESSIONS_PATH, content, policy=SESSION_OPEN)
1351
+ return Session(self, body if isinstance(body, dict) else {}, headers)
1352
+
1353
+ def supervise(
1354
+ self,
1355
+ body: Mapping[str, Any],
1356
+ feed: Iterable[Mapping[str, Any]],
1357
+ on_reply: Callable[[Union[Reply, FrameError]], Any],
1358
+ ) -> Optional[SessionClosed]:
1359
+ """Keep one session answering a feed of readings, reopening it when it closes.
1360
+
1361
+ ``body`` holds the :meth:`open_session` arguments; each reading of ``feed`` holds
1362
+ the :meth:`Session.send` arguments; ``on_reply`` gets every :class:`Reply` and
1363
+ :class:`FrameError`. On a close code that calls for it the session is reopened
1364
+ (at once on ``4408`` and ``4410``, after a backoff of 1 to 30 seconds with jitter
1365
+ otherwise; connected to again on ``1013``); on any other code it stops. When the
1366
+ feed ends it waits for the replies still owed, ends the session, and returns the
1367
+ last :class:`SessionClosed` seen, or ``None``.
1368
+ """
1369
+ return supervise(self, _check_supervise(body), feed, on_reply, self._settings.timeout)
1370
+
1371
+ # ---- batches
1372
+
1373
+ def run_batch(
1374
+ self,
1375
+ requests: Iterable[BatchRequest],
1376
+ poll_interval: float = 60.0,
1377
+ timeout: float = 26 * 3600.0,
1378
+ ) -> BatchRun:
1379
+ """Run a batch from start to end: create a draft, add the requests, finalize, wait
1380
+ for it to end, and return the ended :class:`Batch` with its results.
1381
+
1382
+ The batch is read every ``poll_interval`` seconds (at least 10), or after the
1383
+ server's ``Retry-After`` when that is longer. Past ``timeout`` seconds (default 26
1384
+ hours) it raises :class:`TimeoutError`; the batch keeps running. If adding the
1385
+ requests fails, the draft is cancelled and the error raised.
1386
+ """
1387
+ check_poll(poll_interval, timeout)
1388
+ lines = self.batches._lines(requests)
1389
+ draft = self.batches.create(())
1390
+ if lines:
1391
+ try:
1392
+ self.batches._add_lines(draft.id, lines)
1393
+ except BaseException:
1394
+ try:
1395
+ self.batches.cancel(draft.id)
1396
+ except DecisionNodeError:
1397
+ pass
1398
+ raise
1399
+ batch, headers = self.batches._finalize(draft.id)
1400
+ started = time.monotonic()
1401
+ waited = 0.0
1402
+ while not batch.ended:
1403
+ remaining = timeout - max(time.monotonic() - started, waited)
1404
+ if remaining <= 0:
1405
+ raise TimeoutError("batch %s did not end within %gs; it keeps running" % (batch.id, timeout))
1406
+ wait = min(poll_wait(poll_interval, headers), remaining)
1407
+ self._sleep(wait)
1408
+ waited += wait
1409
+ batch, headers = self.batches._get(batch.id)
1410
+ return BatchRun(batch, self.batches.results(batch.id))
1411
+
1412
+
1413
+ class Batches:
1414
+ """The Batch API of the blocking client, as ``dn.batches``."""
1415
+
1416
+ def __init__(self, client: DecisionNode) -> None:
1417
+ self._client = client
1418
+
1419
+ def _lines(self, requests: Iterable[BatchRequest], allow_empty: bool = True) -> List[bytes]:
1420
+ client = self._client
1421
+ model = client._settings.model
1422
+ out = []
1423
+ for line in client._check_lines(list(requests), allow_empty):
1424
+ if line.content is None:
1425
+ line.content = line.encode(client._finish(line.plan), model)
1426
+ out.append(line.content)
1427
+ return out
1428
+
1429
+ def create(self, requests: Iterable[BatchRequest] = ()) -> Batch:
1430
+ """Create a draft batch holding these requests (none makes an empty draft)."""
1431
+ content = b"".join(self._lines(requests))
1432
+ _, _, body = self._client._json("POST", _base.BATCHES_PATH, content, MUTATION, NDJSON)
1433
+ return parse_batch(body)
1434
+
1435
+ def _add_lines(self, batch_id: str, lines: List[bytes]) -> Batch:
1436
+ batch: Optional[Batch] = None
1437
+ for content in split_lines(lines):
1438
+ _, _, body = self._client._json("POST", _batch_path(batch_id, "/requests"), content, MUTATION, NDJSON)
1439
+ batch = parse_batch(body)
1440
+ assert batch is not None
1441
+ return batch
1442
+
1443
+ def add(self, batch_id: str, requests: Iterable[BatchRequest]) -> Batch:
1444
+ """Add requests to a draft, in calls of at most 10 MiB, in order. Returns the draft."""
1445
+ _batch_path(batch_id)
1446
+ lines = self._lines(requests, allow_empty=False)
1447
+ return self._add_lines(batch_id, lines)
1448
+
1449
+ def _finalize(self, batch_id: str) -> Tuple[Batch, Optional[Dict[str, str]]]:
1450
+ try:
1451
+ _, headers, body = self._client._json("POST", _batch_path(batch_id, "/finalize"), None, MUTATION)
1452
+ return parse_batch(body), headers
1453
+ except DecisionNodeError as exc:
1454
+ # An earlier finalize may have gone through without its answer arriving.
1455
+ if exc.status == 409 and exc.error_type == "invalid_batch_state":
1456
+ batch, headers = self._get(batch_id)
1457
+ if batch.state != "draft":
1458
+ return batch, headers
1459
+ raise
1460
+
1461
+ def finalize(self, batch_id: str) -> Batch:
1462
+ """Queue a draft to run. Nothing runs and nothing is billed before this call."""
1463
+ return self._finalize(batch_id)[0]
1464
+
1465
+ def _get(self, batch_id: str) -> Tuple[Batch, Dict[str, str]]:
1466
+ _, headers, body = self._client._json("GET", _batch_path(batch_id), None)
1467
+ return parse_batch(body), headers
1468
+
1469
+ def get(self, batch_id: str) -> Batch:
1470
+ """The batch as it is now."""
1471
+ return self._get(batch_id)[0]
1472
+
1473
+ def list(self, offset: Optional[int] = None, limit: Optional[int] = None) -> BatchPage:
1474
+ """One page of the key's batches, newest first."""
1475
+ _, _, body = self._client._json("GET", _base.BATCHES_PATH, None, query=_list_query(offset, limit))
1476
+ return parse_page(body)
1477
+
1478
+ def results(self, batch_id: str) -> BatchResults:
1479
+ """The result lines of an ended batch, read as they arrive, in upload order."""
1480
+ response = self._client._open_stream(_batch_path(batch_id, "/results"))
1481
+ return BatchResults(response, self._client._settings.timeout)
1482
+
1483
+ def cancel(self, batch_id: str) -> Batch:
1484
+ """Cancel a batch: lines not started are cancelled, lines running finish."""
1485
+ _, _, body = self._client._json("POST", _batch_path(batch_id, "/cancel"), None)
1486
+ return parse_batch(body)
1487
+
1488
+
1489
+ class Collections:
1490
+ """Collections of faces, as ``dn.collections``: a workspace's own, and the public
1491
+ ones it may search. Every call returns the API's object as sent.
1492
+
1493
+ ``create`` and ``enroll`` send an ``Idempotency-Key``, a new UUID per call and the
1494
+ same on every attempt of it, so the API creates, enrolls and bills once however often
1495
+ the call is tried; they then retry like any call, and also while the API reports the
1496
+ first attempt still running. Reads and deletes carry no key.
1497
+ """
1498
+
1499
+ def __init__(self, client: DecisionNode) -> None:
1500
+ self._client = client
1501
+
1502
+ def create(self, name: str, type: str = "face") -> Dict[str, Any]: # noqa: A002 - the API's field
1503
+ """Create a collection of the workspace."""
1504
+ check_collection(type, name)
1505
+ content = _base.encode({"type": type, "name": name})
1506
+ _, _, body = self._client._json(
1507
+ "POST", _base.COLLECTIONS_PATH, content, IDEMPOTENT, extra_headers=_idempotency_key()
1508
+ )
1509
+ return cast(Dict[str, Any], body)
1510
+
1511
+ def list(self) -> List[Dict[str, Any]]:
1512
+ """The workspace's own collections, newest first, then the public ones it may use."""
1513
+ _, _, body = self._client._json("GET", _base.COLLECTIONS_PATH, None)
1514
+ return _collection_list(body)
1515
+
1516
+ def get(self, collection_id: str) -> Dict[str, Any]:
1517
+ """One collection, by its id or a public collection's name."""
1518
+ _, _, body = self._client._json("GET", _collection_path(collection_id), None)
1519
+ return cast(Dict[str, Any], body)
1520
+
1521
+ def delete(self, collection_id: str) -> Dict[str, Any]:
1522
+ """Delete a collection and every item in it."""
1523
+ _, _, body = self._client._json("DELETE", _collection_path(collection_id), None)
1524
+ return cast(Dict[str, Any], body)
1525
+
1526
+ def enroll(
1527
+ self, collection_id: str, image: ImageInput, reference: str, label: Optional[str] = None
1528
+ ) -> Dict[str, Any]:
1529
+ """Enroll the face in ``image`` (the largest, when it shows several), sent as a
1530
+ decision's image is: a URL as given, local bytes inline or uploaded."""
1531
+ path = _collection_path(collection_id, "/items")
1532
+ check_item(image, reference, label)
1533
+ client = self._client
1534
+ imgs = _base.image_list(image)
1535
+ check_enroll_image(imgs[0])
1536
+ wires = client._finish(client._plan(imgs))
1537
+ item: Dict[str, Any] = {"image": wires[0], "reference": reference}
1538
+ if label is not None:
1539
+ item["label"] = label
1540
+ _, _, body = client._json("POST", path, _base.encode(item), IDEMPOTENT, extra_headers=_idempotency_key())
1541
+ return cast(Dict[str, Any], body)
1542
+
1543
+ def list_items(
1544
+ self, collection_id: str, offset: Optional[int] = None, limit: Optional[int] = None
1545
+ ) -> Dict[str, Any]:
1546
+ """One page of a collection's items, newest first."""
1547
+ path = _collection_path(collection_id, "/items")
1548
+ _, _, body = self._client._json("GET", path, None, query=_list_query(offset, limit))
1549
+ return cast(Dict[str, Any], body)
1550
+
1551
+ def delete_item(self, collection_id: str, item_id: str) -> Dict[str, Any]:
1552
+ """Delete one item of a collection."""
1553
+ _, _, body = self._client._json("DELETE", _item_path(collection_id, item_id), None)
1554
+ return cast(Dict[str, Any], body)
1555
+
1556
+
1557
+ class AsyncDecisionNode(_Common):
1558
+ """The DecisionNode client for asyncio. Same arguments and methods as :class:`DecisionNode`;
1559
+ every method is a coroutine.
1560
+
1561
+ Args:
1562
+ http_client: An ``httpx.AsyncClient`` to send with. The client does not close
1563
+ a client you pass in. Redirects are never followed, whatever the
1564
+ ``httpx.AsyncClient`` is set to.
1565
+ """
1566
+
1567
+ def __init__(
1568
+ self,
1569
+ api_key: Optional[str] = None,
1570
+ base_url: Optional[str] = None,
1571
+ model: Optional[str] = None,
1572
+ timeout: Optional[float] = None,
1573
+ max_retries: Optional[int] = _base.DEFAULT_MAX_RETRIES,
1574
+ headers: Optional[Mapping[str, str]] = None,
1575
+ http_client: Optional[httpx.AsyncClient] = None,
1576
+ inline_threshold: Optional[int] = None,
1577
+ http2: bool = False,
1578
+ upload_timeout: Optional[float] = _base.DEFAULT_UPLOAD_TIMEOUT,
1579
+ ) -> None:
1580
+ self._settings = Settings(
1581
+ api_key, base_url, model, timeout, max_retries, headers, inline_threshold, upload_timeout
1582
+ )
1583
+ if http_client is not None:
1584
+ _check_http_client(http_client, httpx.AsyncClient, _ASYNC_TRANSPORTS, self._settings.base_url)
1585
+ elif http2:
1586
+ _http2_available()
1587
+ self._owns_http = http_client is None
1588
+ self._http = http_client or httpx.AsyncClient(
1589
+ timeout=httpx.Timeout(self._settings.timeout),
1590
+ cookies=_base.no_cookies(),
1591
+ follow_redirects=False,
1592
+ limits=DEFAULT_LIMITS,
1593
+ http2=bool(http2),
1594
+ )
1595
+ self._cache = UploadCache()
1596
+ self._upload_locks: Optional[List[asyncio.Lock]] = None
1597
+ # The wait between attempts, PUT retries, batch polls and session reopens. Tests
1598
+ # replace it to record the waits without sleeping.
1599
+ self._sleep: Callable[[float], Awaitable[None]] = asyncio.sleep
1600
+ self._clock = time.time
1601
+ self._held = {}
1602
+ self.batches = AsyncBatches(self)
1603
+ self.collections = AsyncCollections(self)
1604
+
1605
+ # ---- lifecycle
1606
+
1607
+ async def close(self) -> None:
1608
+ """Close the connections, unless the ``httpx.AsyncClient`` was passed in."""
1609
+ if self._owns_http:
1610
+ await self._http.aclose()
1611
+
1612
+ async def __aenter__(self) -> AsyncDecisionNode:
1613
+ return self
1614
+
1615
+ async def __aexit__(
1616
+ self,
1617
+ exc_type: Optional[Type[BaseException]],
1618
+ exc: Optional[BaseException],
1619
+ tb: Optional[TracebackType],
1620
+ ) -> None:
1621
+ await self.close()
1622
+
1623
+ # ---- transport
1624
+
1625
+ async def _send(
1626
+ self,
1627
+ method: str,
1628
+ path: str,
1629
+ content: Optional[bytes],
1630
+ policy: str = DEFAULT,
1631
+ content_type: str = "application/json",
1632
+ query: Optional[Mapping[str, Any]] = None,
1633
+ timeout: Optional[float] = None,
1634
+ extra_headers: Optional[Mapping[str, str]] = None,
1635
+ ) -> _Response:
1636
+ settings = self._settings
1637
+ headers = settings.request_headers(content is not None, content_type)
1638
+ if extra_headers:
1639
+ # The same on every attempt: an Idempotency-Key is the call's, not the attempt's.
1640
+ headers.update(extra_headers)
1641
+ url = _url(settings, path, query)
1642
+ limit = timeout or settings.timeout
1643
+ attempt = 0
1644
+ while True:
1645
+ started = time.monotonic()
1646
+ resp_headers: Optional[Dict[str, str]] = None
1647
+ try:
1648
+ # The timeout covers the whole attempt, up to the last byte of the body.
1649
+ response = await asyncio.wait_for(
1650
+ self._http.request(
1651
+ method,
1652
+ url,
1653
+ content=content,
1654
+ headers=headers,
1655
+ timeout=limit,
1656
+ # Never followed, also on a client passed in: a redirect would send
1657
+ # the key to another URL, possibly over plain HTTP.
1658
+ follow_redirects=False,
1659
+ ),
1660
+ limit,
1661
+ )
1662
+ except httpx.TransportError as exc:
1663
+ error = _transport_error(exc, limit)
1664
+ _base.log_response(method, path, type(error).__name__, None, started, attempt)
1665
+ retry = _retry_without_answer(policy, error)
1666
+ except asyncio.TimeoutError:
1667
+ error = _timeout_error(limit)
1668
+ _base.log_response(method, path, type(error).__name__, None, started, attempt)
1669
+ retry = _retry_without_answer(policy, error)
1670
+ else:
1671
+ resp_headers = _base.lower_headers(response.headers)
1672
+ status = response.status_code
1673
+ text = response.content.decode("utf-8", errors="replace")
1674
+ _base.log_response(method, path, status, resp_headers.get("x-request-id"), started, attempt)
1675
+ if 200 <= status < 300:
1676
+ return status, resp_headers, text
1677
+ error = _base.make_error(status, resp_headers, text)
1678
+ retry = _base.should_retry(policy, status, error)
1679
+ if not retry or attempt >= settings.max_retries:
1680
+ raise error
1681
+ await self._sleep(_base.retry_wait(attempt, resp_headers))
1682
+ attempt += 1
1683
+
1684
+ async def _json(
1685
+ self,
1686
+ method: str,
1687
+ path: str,
1688
+ content: Optional[bytes],
1689
+ policy: str = DEFAULT,
1690
+ content_type: str = "application/json",
1691
+ query: Optional[Mapping[str, Any]] = None,
1692
+ extra_headers: Optional[Mapping[str, str]] = None,
1693
+ ) -> Tuple[int, Dict[str, str], Any]:
1694
+ status, headers, text = await self._send(
1695
+ method, path, content, policy, content_type, query, extra_headers=extra_headers
1696
+ )
1697
+ return status, headers, _base.parse_json(status, headers, text)
1698
+
1699
+ async def _open_stream(self, path: str) -> httpx.Response:
1700
+ settings = self._settings
1701
+ headers = settings.request_headers(False)
1702
+ url = settings.url(path)
1703
+ attempt = 0
1704
+ while True:
1705
+ started = time.monotonic()
1706
+ resp_headers: Optional[Dict[str, str]] = None
1707
+ try:
1708
+ request = self._http.build_request("GET", url, headers=headers, timeout=settings.timeout)
1709
+ response = await self._http.send(request, stream=True, follow_redirects=False)
1710
+ except httpx.TransportError as exc:
1711
+ error = _transport_error(exc, settings.timeout)
1712
+ _base.log_response("GET", path, type(error).__name__, None, started, attempt)
1713
+ retry = True
1714
+ else:
1715
+ resp_headers = _base.lower_headers(response.headers)
1716
+ status = response.status_code
1717
+ _base.log_response("GET", path, status, resp_headers.get("x-request-id"), started, attempt)
1718
+ if 200 <= status < 300:
1719
+ return response # type: ignore[no-any-return]
1720
+ try:
1721
+ text = (await response.aread()).decode("utf-8", errors="replace")
1722
+ except httpx.TransportError:
1723
+ text = ""
1724
+ finally:
1725
+ await response.aclose()
1726
+ error = _base.make_error(status, resp_headers, text)
1727
+ retry = status in _base.RETRY_STATUSES
1728
+ if not retry or attempt >= settings.max_retries:
1729
+ raise error
1730
+ await self._sleep(_base.retry_wait(attempt, resp_headers))
1731
+ attempt += 1
1732
+
1733
+ # ---- uploads
1734
+
1735
+ async def _put(
1736
+ self, img: Image, entry: Mapping[str, Any], presigned_at: float, url_ttl: float, may_presign: bool
1737
+ ) -> None:
1738
+ settings = self._settings
1739
+ path = _storage_path(entry["url"])
1740
+ attempt = 0
1741
+ while True:
1742
+ started = time.monotonic()
1743
+ resp_headers: Optional[Dict[str, str]] = None
1744
+ request = httpx.Request(
1745
+ "PUT",
1746
+ entry["url"],
1747
+ headers=entry["headers"],
1748
+ content=img.data,
1749
+ extensions={"timeout": httpx.Timeout(settings.upload_timeout).as_dict()},
1750
+ )
1751
+ try:
1752
+ # Its own deadline, for the whole attempt: the upload_timeout.
1753
+ response = await asyncio.wait_for(
1754
+ self._http.send(request, follow_redirects=False), settings.upload_timeout
1755
+ )
1756
+ except httpx.TransportError as exc:
1757
+ error = _transport_error(exc, settings.upload_timeout)
1758
+ _base.log_response("PUT", path, type(error).__name__, None, started, attempt)
1759
+ retry = True
1760
+ except asyncio.TimeoutError:
1761
+ error = _timeout_error(settings.upload_timeout)
1762
+ _base.log_response("PUT", path, type(error).__name__, None, started, attempt)
1763
+ retry = True
1764
+ else:
1765
+ status = response.status_code
1766
+ resp_headers = _base.lower_headers(response.headers)
1767
+ _base.log_response("PUT", path, status, None, started, attempt)
1768
+ if 200 <= status < 300 or (status == 412 and attempt > 0):
1769
+ return
1770
+ if status == 403 and may_presign and self._clock() - presigned_at >= url_ttl:
1771
+ raise UrlExpired()
1772
+ error = storage_error(img.id, status, resp_headers, response.text)
1773
+ retry = 500 <= status <= 599
1774
+ if not retry or attempt >= settings.max_retries:
1775
+ raise error
1776
+ await self._sleep(_base.retry_wait(attempt, resp_headers))
1777
+ attempt += 1
1778
+
1779
+ async def _presign_and_put(self, todo: List[Image], may_presign: bool = True) -> Dict[CacheKey, str]:
1780
+ _, _, body = await self._json("POST", _base.UPLOADS_PATH, _base.encode(presign_body(todo)))
1781
+ entries = presigned_entries(body, len(todo))
1782
+ for entry in entries:
1783
+ check_storage_url(entry["url"])
1784
+ url_ttl = float(body.get("expires_in") or 300)
1785
+ presigned_at = self._clock()
1786
+ out: Dict[CacheKey, str] = {}
1787
+ expired: List[Image] = []
1788
+ results = await asyncio.gather(
1789
+ *(self._put(img, entry, presigned_at, url_ttl, may_presign) for img, entry in zip(todo, entries)),
1790
+ return_exceptions=True,
1791
+ )
1792
+ errors: List[BaseException] = []
1793
+ for img, entry, result in zip(todo, entries, results):
1794
+ if isinstance(result, UrlExpired):
1795
+ expired.append(img)
1796
+ elif isinstance(result, BaseException):
1797
+ errors.append(result)
1798
+ else:
1799
+ key = cache_key(img)
1800
+ self._cache.put(key, entry["upload_id"], _base.parse_time(entry.get("expires_at")))
1801
+ out[key] = entry["upload_id"]
1802
+ if errors:
1803
+ raise errors[0]
1804
+ if expired:
1805
+ out.update(await self._presign_and_put(expired, may_presign=False))
1806
+ return out
1807
+
1808
+ async def _upload(self, images: Sequence[Image]) -> Dict[CacheKey, str]:
1809
+ if self._upload_locks is None:
1810
+ self._upload_locks = [asyncio.Lock() for _ in range(LOCK_STRIPES)]
1811
+ keys = [cache_key(img) for img in images]
1812
+ locks = [self._upload_locks[i] for i in sorted({stripe(k) for k in keys})]
1813
+ held: List[asyncio.Lock] = []
1814
+ try:
1815
+ for lock in locks:
1816
+ await lock.acquire()
1817
+ held.append(lock)
1818
+ out: Dict[CacheKey, str] = {}
1819
+ todo: List[Image] = []
1820
+ for img, key in zip(images, keys):
1821
+ hit = self._cache.get(key, self._clock())
1822
+ if hit is not None:
1823
+ out[key] = hit
1824
+ elif key not in out and all(cache_key(t) != key for t in todo):
1825
+ todo.append(img)
1826
+ if todo:
1827
+ out.update(await self._presign_and_put(todo))
1828
+ return out
1829
+ finally:
1830
+ for lock in reversed(held):
1831
+ lock.release()
1832
+
1833
+ async def _finish(self, plan: ImagePlan) -> List[Dict[str, Any]]:
1834
+ return plan.finish(await self._upload(plan.uploads) if plan.uploads else {})
1835
+
1836
+ async def _frame_image(self, image: Any) -> Dict[str, Any]:
1837
+ plan = self._plan([image], strip_id=True)
1838
+ return (await self._finish(plan))[0]
1839
+
1840
+ # ---- decisions
1841
+
1842
+ async def _execute(self, prep: _Prepared) -> _Response:
1843
+ plan = prep.plan
1844
+ assert plan is not None
1845
+ content = prep.content
1846
+ for attempt in (0, 1):
1847
+ if content is None:
1848
+ content = self._body(prep, await self._finish(plan))
1849
+ try:
1850
+ policy, timeout = _ocr_send(self._settings, prep.boosters)
1851
+ return await self._send("POST", prep.path, content, policy, timeout=timeout)
1852
+ except BadRequestError as exc:
1853
+ if attempt == 0 and plan.placed and exc.error_type in _base.STALE_UPLOAD:
1854
+ self._cache.drop(plan.placed)
1855
+ plan = self._plan(prep.images)
1856
+ content = None
1857
+ continue
1858
+ raise
1859
+ raise AssertionError("unreachable")
1860
+
1861
+ async def _decision(self, prep: _Prepared) -> Decision:
1862
+ status, headers, text = await self._execute(prep)
1863
+ return parse_decision(_base.parse_json(status, headers, text), headers)
1864
+
1865
+ async def _ask(
1866
+ self,
1867
+ question_id: str,
1868
+ question: Question,
1869
+ kind: Type[_A],
1870
+ state: Any,
1871
+ images: Any,
1872
+ model: Any,
1873
+ videos: Any,
1874
+ boosters: Any,
1875
+ fallback: Optional[bool] = None,
1876
+ ) -> _A:
1877
+ prep = self._prepare(
1878
+ state, {question_id: question}, images, model, _base.DECIDE_PATH, videos, boosters, fallback
1879
+ )
1880
+ status, headers, text = await self._execute(prep)
1881
+ decision = parse_decision(_base.parse_json(status, headers, text), headers)
1882
+ return _single(decision, question_id, kind, status, headers)
1883
+
1884
+ # ---- the surface
1885
+
1886
+ async def decide(
1887
+ self,
1888
+ state: Any = UNSET,
1889
+ questions: Optional[Mapping[str, Question]] = None,
1890
+ images: Optional[Sequence[ImageInput]] = None,
1891
+ model: Optional[str] = None,
1892
+ *,
1893
+ videos: Optional[Sequence[VideoInput]] = None,
1894
+ boosters: Optional[Sequence[BoosterInput]] = None,
1895
+ fallback: Optional[bool] = None,
1896
+ ) -> Decision:
1897
+ """Ask one or more questions about a state in one request. See :meth:`DecisionNode.decide`."""
1898
+ return await self._decision(
1899
+ self._prepare(state, questions, images, model, _base.DECIDE_PATH, videos, boosters, fallback)
1900
+ )
1901
+
1902
+ async def decide_many(
1903
+ self,
1904
+ requests: Iterable[DecideRequest],
1905
+ concurrency: int = _base.DEFAULT_CONCURRENCY,
1906
+ ) -> List[DecisionResult]:
1907
+ """Run many ``decide`` calls, at most ``concurrency`` at a time (default 8).
1908
+
1909
+ Returns one result per request, in input order: the :class:`Decision`, or the
1910
+ :class:`DecisionNodeError` that request ended with.
1911
+ """
1912
+ preps = self._prepare_many(requests, concurrency)
1913
+ limit = asyncio.Semaphore(concurrency)
1914
+
1915
+ async def run(prep: _Prepared) -> DecisionResult:
1916
+ async with limit:
1917
+ try:
1918
+ return await self._decision(prep)
1919
+ except DecisionNodeError as exc:
1920
+ return exc
1921
+
1922
+ return list(await asyncio.gather(*(run(prep) for prep in preps)))
1923
+
1924
+ async def check(
1925
+ self,
1926
+ state: Any = UNSET,
1927
+ instructions: Any = None,
1928
+ if_true: Any = None,
1929
+ if_false: Any = None,
1930
+ *,
1931
+ images: Optional[Sequence[ImageInput]] = None,
1932
+ videos: Optional[Sequence[VideoInput]] = None,
1933
+ model: Optional[str] = None,
1934
+ boosters: Optional[Sequence[BoosterInput]] = None,
1935
+ fallback: Optional[bool] = None,
1936
+ ) -> TruthAnswer:
1937
+ """Ask one ``truth`` question. Returns its :class:`TruthAnswer`.
1938
+
1939
+ Raises :class:`ServerError` when that answer is missing or malformed.
1940
+ """
1941
+ return await self._ask(
1942
+ "check",
1943
+ _truth(instructions, if_true, if_false),
1944
+ TruthAnswer,
1945
+ state,
1946
+ images,
1947
+ model,
1948
+ videos,
1949
+ boosters,
1950
+ fallback,
1951
+ )
1952
+
1953
+ async def classify(
1954
+ self,
1955
+ state: Any = UNSET,
1956
+ instructions: Any = None,
1957
+ options: Union[Mapping[str, Any], Iterable[str], None] = None,
1958
+ *,
1959
+ images: Optional[Sequence[ImageInput]] = None,
1960
+ videos: Optional[Sequence[VideoInput]] = None,
1961
+ model: Optional[str] = None,
1962
+ boosters: Optional[Sequence[BoosterInput]] = None,
1963
+ fallback: Optional[bool] = None,
1964
+ ) -> ChoiceAnswer:
1965
+ """Ask one ``choice`` question. Returns its :class:`ChoiceAnswer`.
1966
+
1967
+ Raises :class:`ServerError` when that answer is missing or malformed.
1968
+ """
1969
+ return await self._ask(
1970
+ "classify", _choice(instructions, options), ChoiceAnswer, state, images, model, videos, boosters, fallback
1971
+ )
1972
+
1973
+ async def score(
1974
+ self,
1975
+ state: Any = UNSET,
1976
+ instructions: Any = None,
1977
+ levels: Optional[Sequence[Any]] = None,
1978
+ *,
1979
+ images: Optional[Sequence[ImageInput]] = None,
1980
+ videos: Optional[Sequence[VideoInput]] = None,
1981
+ model: Optional[str] = None,
1982
+ boosters: Optional[Sequence[BoosterInput]] = None,
1983
+ fallback: Optional[bool] = None,
1984
+ ) -> ScoreAnswer:
1985
+ """Ask one ``score`` question. Returns its :class:`ScoreAnswer`.
1986
+
1987
+ Raises :class:`ServerError` when that answer is missing or malformed.
1988
+ """
1989
+ return await self._ask(
1990
+ "score", _score(instructions, levels), ScoreAnswer, state, images, model, videos, boosters, fallback
1991
+ )
1992
+
1993
+ async def number(
1994
+ self,
1995
+ state: Any = UNSET,
1996
+ instructions: Any = None,
1997
+ min: Union[int, float, None] = None, # noqa: A002 - the wire name
1998
+ max: Union[int, float, None] = None, # noqa: A002 - the wire name
1999
+ step: Union[int, float, None] = None,
2000
+ *,
2001
+ images: Optional[Sequence[ImageInput]] = None,
2002
+ videos: Optional[Sequence[VideoInput]] = None,
2003
+ model: Optional[str] = None,
2004
+ boosters: Optional[Sequence[BoosterInput]] = None,
2005
+ fallback: Optional[bool] = None,
2006
+ ) -> NumberAnswer:
2007
+ """Ask one ``number`` question. Returns its :class:`NumberAnswer`.
2008
+
2009
+ Raises :class:`ServerError` when that answer is missing or malformed.
2010
+ """
2011
+ return await self._ask(
2012
+ "number",
2013
+ _number(instructions, min, max, step),
2014
+ NumberAnswer,
2015
+ state,
2016
+ images,
2017
+ model,
2018
+ videos,
2019
+ boosters,
2020
+ fallback,
2021
+ )
2022
+
2023
+ async def gate(
2024
+ self,
2025
+ action: Any,
2026
+ context: Any = None,
2027
+ *,
2028
+ images: Optional[Sequence[ImageInput]] = None,
2029
+ videos: Optional[Sequence[VideoInput]] = None,
2030
+ model: Optional[str] = None,
2031
+ boosters: Optional[Sequence[BoosterInput]] = None,
2032
+ fallback: Optional[bool] = None,
2033
+ ) -> Decision:
2034
+ """Judge an action before it runs. See :meth:`DecisionNode.gate`."""
2035
+ return await self.decide(
2036
+ _base.gate_state(action, context),
2037
+ _base.gate_questions(),
2038
+ images,
2039
+ model,
2040
+ videos=videos,
2041
+ boosters=boosters,
2042
+ fallback=fallback,
2043
+ )
2044
+
2045
+ async def models(self) -> List[ModelInfo]:
2046
+ """The models this key may use, with the limits each one publishes."""
2047
+ _, _, body = await self._json("GET", _base.MODELS_PATH, None)
2048
+ entries = body.get("models") if isinstance(body, Mapping) else None
2049
+ self._hold(entries)
2050
+ return [parse_model(m) for m in entries or []]
2051
+
2052
+ async def upload(self, images: Sequence[ImageInput]) -> List[str]:
2053
+ """Upload local images whatever their size and return the ``upload_id`` of each.
2054
+ See :meth:`DecisionNode.upload`."""
2055
+ imgs = _upload_list(images)
2056
+ check_images(imgs, self._limits())
2057
+ local = []
2058
+ for img in imgs:
2059
+ if not isinstance(img, Image) or (img.data is None and img.upload_id is None):
2060
+ raise ValueError("upload() takes Image objects made from a path, bytes or an upload_id")
2061
+ if img.data is not None:
2062
+ local.append(img)
2063
+ ids = await self._upload(local) if local else {}
2064
+ return [img.upload_id if img.data is None else ids[cache_key(img)] for img in imgs]
2065
+
2066
+ async def decide_raw(self, body: Any) -> Any:
2067
+ """Send ``body`` to ``POST /v1/decide`` exactly as given and return the parsed response body."""
2068
+ _, _, parsed = await self._json("POST", _base.DECIDE_PATH, _base.encode(body))
2069
+ return parsed
2070
+
2071
+ # ---- sessions
2072
+
2073
+ async def open_session(
2074
+ self,
2075
+ state: Any = UNSET,
2076
+ questions: Optional[Mapping[str, Question]] = None,
2077
+ instructions: Any = None,
2078
+ model: Optional[str] = None,
2079
+ window: Optional[int] = None,
2080
+ ttl_seconds: Optional[int] = None,
2081
+ ) -> AsyncSession:
2082
+ """Open a session. See :meth:`DecisionNode.open_session`."""
2083
+ content = _session_body(self, state, questions, instructions, model, window, ttl_seconds)
2084
+ _, headers, body = await self._json("POST", _base.SESSIONS_PATH, content, policy=SESSION_OPEN)
2085
+ return AsyncSession(self, body if isinstance(body, dict) else {}, headers)
2086
+
2087
+ async def supervise(
2088
+ self,
2089
+ body: Mapping[str, Any],
2090
+ feed: Union[Iterable[Mapping[str, Any]], AsyncIterator[Mapping[str, Any]]],
2091
+ on_reply: Callable[[Union[Reply, FrameError]], Any],
2092
+ ) -> Optional[SessionClosed]:
2093
+ """Keep one session answering a feed of readings. ``feed`` may be an iterable or an
2094
+ async iterable, ``on_reply`` a function or a coroutine function. See
2095
+ :meth:`DecisionNode.supervise`."""
2096
+ return await asupervise(self, _check_supervise(body), feed, on_reply, self._settings.timeout)
2097
+
2098
+ # ---- batches
2099
+
2100
+ async def run_batch(
2101
+ self,
2102
+ requests: Iterable[BatchRequest],
2103
+ poll_interval: float = 60.0,
2104
+ timeout: float = 26 * 3600.0,
2105
+ ) -> BatchRun:
2106
+ """Run a batch from start to end. See :meth:`DecisionNode.run_batch`; the results
2107
+ are an :class:`AsyncBatchResults`."""
2108
+ check_poll(poll_interval, timeout)
2109
+ lines = await self.batches._lines(requests)
2110
+ draft = await self.batches.create(())
2111
+ if lines:
2112
+ try:
2113
+ await self.batches._add_lines(draft.id, lines)
2114
+ except BaseException:
2115
+ try:
2116
+ await self.batches.cancel(draft.id)
2117
+ except DecisionNodeError:
2118
+ pass
2119
+ raise
2120
+ batch, headers = await self.batches._finalize(draft.id)
2121
+ started = time.monotonic()
2122
+ waited = 0.0
2123
+ while not batch.ended:
2124
+ remaining = timeout - max(time.monotonic() - started, waited)
2125
+ if remaining <= 0:
2126
+ raise TimeoutError("batch %s did not end within %gs; it keeps running" % (batch.id, timeout))
2127
+ wait = min(poll_wait(poll_interval, headers), remaining)
2128
+ await self._sleep(wait)
2129
+ waited += wait
2130
+ batch, headers = await self.batches._get(batch.id)
2131
+ return BatchRun(batch, await self.batches.results(batch.id))
2132
+
2133
+
2134
+ class AsyncBatches:
2135
+ """The Batch API of the asyncio client, as ``dn.batches``. Same methods as
2136
+ :class:`Batches`; every method is a coroutine."""
2137
+
2138
+ def __init__(self, client: AsyncDecisionNode) -> None:
2139
+ self._client = client
2140
+
2141
+ async def _lines(self, requests: Iterable[BatchRequest], allow_empty: bool = True) -> List[bytes]:
2142
+ client = self._client
2143
+ model = client._settings.model
2144
+ out = []
2145
+ for line in client._check_lines(list(requests), allow_empty):
2146
+ if line.content is None:
2147
+ line.content = line.encode(await client._finish(line.plan), model)
2148
+ out.append(line.content)
2149
+ return out
2150
+
2151
+ async def create(self, requests: Iterable[BatchRequest] = ()) -> Batch:
2152
+ """Create a draft batch holding these requests (none makes an empty draft)."""
2153
+ content = b"".join(await self._lines(requests))
2154
+ _, _, body = await self._client._json("POST", _base.BATCHES_PATH, content, MUTATION, NDJSON)
2155
+ return parse_batch(body)
2156
+
2157
+ async def _add_lines(self, batch_id: str, lines: List[bytes]) -> Batch:
2158
+ batch: Optional[Batch] = None
2159
+ for content in split_lines(lines):
2160
+ _, _, body = await self._client._json("POST", _batch_path(batch_id, "/requests"), content, MUTATION, NDJSON)
2161
+ batch = parse_batch(body)
2162
+ assert batch is not None
2163
+ return batch
2164
+
2165
+ async def add(self, batch_id: str, requests: Iterable[BatchRequest]) -> Batch:
2166
+ """Add requests to a draft, in calls of at most 10 MiB, in order. Returns the draft."""
2167
+ _batch_path(batch_id)
2168
+ lines = await self._lines(requests, allow_empty=False)
2169
+ return await self._add_lines(batch_id, lines)
2170
+
2171
+ async def _finalize(self, batch_id: str) -> Tuple[Batch, Optional[Dict[str, str]]]:
2172
+ try:
2173
+ _, headers, body = await self._client._json("POST", _batch_path(batch_id, "/finalize"), None, MUTATION)
2174
+ return parse_batch(body), headers
2175
+ except DecisionNodeError as exc:
2176
+ if exc.status == 409 and exc.error_type == "invalid_batch_state":
2177
+ batch, headers = await self._get(batch_id)
2178
+ if batch.state != "draft":
2179
+ return batch, headers
2180
+ raise
2181
+
2182
+ async def finalize(self, batch_id: str) -> Batch:
2183
+ """Queue a draft to run."""
2184
+ return (await self._finalize(batch_id))[0]
2185
+
2186
+ async def _get(self, batch_id: str) -> Tuple[Batch, Dict[str, str]]:
2187
+ _, headers, body = await self._client._json("GET", _batch_path(batch_id), None)
2188
+ return parse_batch(body), headers
2189
+
2190
+ async def get(self, batch_id: str) -> Batch:
2191
+ """The batch as it is now."""
2192
+ return (await self._get(batch_id))[0]
2193
+
2194
+ async def list(self, offset: Optional[int] = None, limit: Optional[int] = None) -> BatchPage:
2195
+ """One page of the key's batches, newest first."""
2196
+ _, _, body = await self._client._json("GET", _base.BATCHES_PATH, None, query=_list_query(offset, limit))
2197
+ return parse_page(body)
2198
+
2199
+ async def results(self, batch_id: str) -> AsyncBatchResults:
2200
+ """The result lines of an ended batch: ``async for result in await dn.batches.results(id)``."""
2201
+ response = await self._client._open_stream(_batch_path(batch_id, "/results"))
2202
+ return AsyncBatchResults(response, self._client._settings.timeout)
2203
+
2204
+ async def cancel(self, batch_id: str) -> Batch:
2205
+ """Cancel a batch."""
2206
+ _, _, body = await self._client._json("POST", _batch_path(batch_id, "/cancel"), None)
2207
+ return parse_batch(body)
2208
+
2209
+
2210
+ class AsyncCollections:
2211
+ """Collections of faces, as ``dn.collections`` on the async client. Every call
2212
+ returns the API's object as sent; ``create`` and ``enroll`` are retried only on 429
2213
+ and 529."""
2214
+
2215
+ def __init__(self, client: AsyncDecisionNode) -> None:
2216
+ self._client = client
2217
+
2218
+ async def create(self, name: str, type: str = "face") -> Dict[str, Any]: # noqa: A002 - the API's field
2219
+ """Create a collection of the workspace."""
2220
+ check_collection(type, name)
2221
+ content = _base.encode({"type": type, "name": name})
2222
+ _, _, body = await self._client._json(
2223
+ "POST", _base.COLLECTIONS_PATH, content, IDEMPOTENT, extra_headers=_idempotency_key()
2224
+ )
2225
+ return cast(Dict[str, Any], body)
2226
+
2227
+ async def list(self) -> List[Dict[str, Any]]:
2228
+ """The workspace's own collections, newest first, then the public ones it may use."""
2229
+ _, _, body = await self._client._json("GET", _base.COLLECTIONS_PATH, None)
2230
+ return _collection_list(body)
2231
+
2232
+ async def get(self, collection_id: str) -> Dict[str, Any]:
2233
+ """One collection, by its id or a public collection's name."""
2234
+ _, _, body = await self._client._json("GET", _collection_path(collection_id), None)
2235
+ return cast(Dict[str, Any], body)
2236
+
2237
+ async def delete(self, collection_id: str) -> Dict[str, Any]:
2238
+ """Delete a collection and every item in it."""
2239
+ _, _, body = await self._client._json("DELETE", _collection_path(collection_id), None)
2240
+ return cast(Dict[str, Any], body)
2241
+
2242
+ async def enroll(
2243
+ self, collection_id: str, image: ImageInput, reference: str, label: Optional[str] = None
2244
+ ) -> Dict[str, Any]:
2245
+ """Enroll the face in ``image``, sent as a decision's image is."""
2246
+ path = _collection_path(collection_id, "/items")
2247
+ check_item(image, reference, label)
2248
+ client = self._client
2249
+ imgs = _base.image_list(image)
2250
+ check_enroll_image(imgs[0])
2251
+ wires = await client._finish(client._plan(imgs))
2252
+ item: Dict[str, Any] = {"image": wires[0], "reference": reference}
2253
+ if label is not None:
2254
+ item["label"] = label
2255
+ _, _, body = await client._json("POST", path, _base.encode(item), IDEMPOTENT, extra_headers=_idempotency_key())
2256
+ return cast(Dict[str, Any], body)
2257
+
2258
+ async def list_items(
2259
+ self, collection_id: str, offset: Optional[int] = None, limit: Optional[int] = None
2260
+ ) -> Dict[str, Any]:
2261
+ """One page of a collection's items, newest first."""
2262
+ path = _collection_path(collection_id, "/items")
2263
+ _, _, body = await self._client._json("GET", path, None, query=_list_query(offset, limit))
2264
+ return cast(Dict[str, Any], body)
2265
+
2266
+ async def delete_item(self, collection_id: str, item_id: str) -> Dict[str, Any]:
2267
+ """Delete one item of a collection."""
2268
+ _, _, body = await self._client._json("DELETE", _item_path(collection_id, item_id), None)
2269
+ return cast(Dict[str, Any], body)