tokenrail 0.4.0__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.
- tokenrail/__init__.py +32 -0
- tokenrail/catalog.py +320 -0
- tokenrail/client.py +58 -0
- tokenrail/executor.py +249 -0
- tokenrail/monitor.py +457 -0
- tokenrail/providers/__init__.py +3 -0
- tokenrail/providers/base.py +23 -0
- tokenrail/providers/openai.py +452 -0
- tokenrail/py.typed +0 -0
- tokenrail/sinks.py +106 -0
- tokenrail/types.py +258 -0
- tokenrail-0.4.0.dist-info/METADATA +338 -0
- tokenrail-0.4.0.dist-info/RECORD +15 -0
- tokenrail-0.4.0.dist-info/WHEEL +4 -0
- tokenrail-0.4.0.dist-info/licenses/LICENSE +21 -0
tokenrail/__init__.py
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
"""Thin client and batch execution helpers for OpenAI Responses API workloads.
|
|
2
|
+
|
|
3
|
+
tokenrail wraps the OpenAI Responses API with a ``client.responses``-style surface
|
|
4
|
+
and adds thread-based batch execution, structured output parsing, client-side
|
|
5
|
+
RPM/TPM submit throttling, per-model token/cost monitoring, and resumable JSONL /
|
|
6
|
+
per-request result writing.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from .client import RailClient
|
|
10
|
+
from .executor import BatchExecutor, batch_items_from_queries
|
|
11
|
+
from .monitor import RollingMetricsMonitor
|
|
12
|
+
from .providers import OpenAIProvider
|
|
13
|
+
from .sinks import PerRequestJsonSink, ResultsJsonlSink
|
|
14
|
+
from .types import BatchItem, CostBreakdown, NormalizedResponse, StatsSnapshot, UsageBreakdown
|
|
15
|
+
|
|
16
|
+
__version__ = "0.4.0"
|
|
17
|
+
|
|
18
|
+
__all__ = [
|
|
19
|
+
"__version__",
|
|
20
|
+
"BatchExecutor",
|
|
21
|
+
"BatchItem",
|
|
22
|
+
"CostBreakdown",
|
|
23
|
+
"NormalizedResponse",
|
|
24
|
+
"OpenAIProvider",
|
|
25
|
+
"PerRequestJsonSink",
|
|
26
|
+
"RailClient",
|
|
27
|
+
"ResultsJsonlSink",
|
|
28
|
+
"RollingMetricsMonitor",
|
|
29
|
+
"StatsSnapshot",
|
|
30
|
+
"UsageBreakdown",
|
|
31
|
+
"batch_items_from_queries",
|
|
32
|
+
]
|
tokenrail/catalog.py
ADDED
|
@@ -0,0 +1,320 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import re
|
|
4
|
+
import threading
|
|
5
|
+
import warnings
|
|
6
|
+
from collections.abc import Iterable
|
|
7
|
+
from dataclasses import dataclass
|
|
8
|
+
from decimal import Decimal
|
|
9
|
+
from typing import TypeVar
|
|
10
|
+
|
|
11
|
+
from .types import CostBreakdown, UsageBreakdown
|
|
12
|
+
|
|
13
|
+
_T = TypeVar("_T")
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class ModelCatalogFallbackWarning(UserWarning):
|
|
17
|
+
"""Warning emitted when an unregistered model uses fallback catalog data."""
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
@dataclass(frozen=True, slots=True)
|
|
21
|
+
class ModelCapabilities:
|
|
22
|
+
reasoning_effort: bool
|
|
23
|
+
verbosity: bool
|
|
24
|
+
temperature: bool
|
|
25
|
+
top_p: bool
|
|
26
|
+
max_output_tokens: bool
|
|
27
|
+
response_format: bool
|
|
28
|
+
supported_reasoning_efforts: tuple[str, ...] | None = None
|
|
29
|
+
sampling_requires_no_reasoning: bool = False
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
@dataclass(frozen=True, slots=True)
|
|
33
|
+
class ModelPricing:
|
|
34
|
+
input_per_million: Decimal
|
|
35
|
+
cached_input_per_million: Decimal | None
|
|
36
|
+
output_per_million: Decimal
|
|
37
|
+
service_tier: str = "default"
|
|
38
|
+
cache_write_input_per_million: Decimal | None = None
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
_CAPABILITY_RULES: list[tuple[tuple[str, ...], ModelCapabilities]] = [
|
|
42
|
+
(
|
|
43
|
+
("gpt-6-astra", "gpt-6.1-sol"),
|
|
44
|
+
ModelCapabilities(
|
|
45
|
+
reasoning_effort=True,
|
|
46
|
+
verbosity=True,
|
|
47
|
+
temperature=False,
|
|
48
|
+
top_p=False,
|
|
49
|
+
max_output_tokens=True,
|
|
50
|
+
response_format=True,
|
|
51
|
+
supported_reasoning_efforts=("low", "medium", "high", "xhigh", "max"),
|
|
52
|
+
),
|
|
53
|
+
),
|
|
54
|
+
(
|
|
55
|
+
("gpt-6-sol", "gpt-6-luna"),
|
|
56
|
+
ModelCapabilities(
|
|
57
|
+
reasoning_effort=True,
|
|
58
|
+
verbosity=True,
|
|
59
|
+
temperature=True,
|
|
60
|
+
top_p=True,
|
|
61
|
+
max_output_tokens=True,
|
|
62
|
+
response_format=True,
|
|
63
|
+
supported_reasoning_efforts=("none", "low", "medium", "high", "xhigh", "max"),
|
|
64
|
+
sampling_requires_no_reasoning=True,
|
|
65
|
+
),
|
|
66
|
+
),
|
|
67
|
+
(
|
|
68
|
+
("gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.6"),
|
|
69
|
+
ModelCapabilities(
|
|
70
|
+
reasoning_effort=True,
|
|
71
|
+
verbosity=True,
|
|
72
|
+
temperature=True,
|
|
73
|
+
top_p=True,
|
|
74
|
+
max_output_tokens=True,
|
|
75
|
+
response_format=True,
|
|
76
|
+
),
|
|
77
|
+
),
|
|
78
|
+
(
|
|
79
|
+
("gpt-5", "o1", "o3", "o4"),
|
|
80
|
+
ModelCapabilities(
|
|
81
|
+
reasoning_effort=True,
|
|
82
|
+
verbosity=True,
|
|
83
|
+
temperature=True,
|
|
84
|
+
top_p=True,
|
|
85
|
+
max_output_tokens=True,
|
|
86
|
+
response_format=True,
|
|
87
|
+
),
|
|
88
|
+
),
|
|
89
|
+
(
|
|
90
|
+
("gpt-4.1", "gpt-4o"),
|
|
91
|
+
ModelCapabilities(
|
|
92
|
+
reasoning_effort=False,
|
|
93
|
+
verbosity=True,
|
|
94
|
+
temperature=True,
|
|
95
|
+
top_p=True,
|
|
96
|
+
max_output_tokens=True,
|
|
97
|
+
response_format=True,
|
|
98
|
+
),
|
|
99
|
+
),
|
|
100
|
+
]
|
|
101
|
+
|
|
102
|
+
_DEFAULT_CAPABILITIES = ModelCapabilities(
|
|
103
|
+
reasoning_effort=False,
|
|
104
|
+
verbosity=False,
|
|
105
|
+
temperature=True,
|
|
106
|
+
top_p=True,
|
|
107
|
+
max_output_tokens=True,
|
|
108
|
+
response_format=True,
|
|
109
|
+
)
|
|
110
|
+
|
|
111
|
+
_PRICING_RULES: list[tuple[tuple[str, ...], ModelPricing]] = [
|
|
112
|
+
# Standard-tier short-context rates verified on 2026-10-08:
|
|
113
|
+
# https://developers.openai.com/api/docs/pricing
|
|
114
|
+
(
|
|
115
|
+
("gpt-6-astra",),
|
|
116
|
+
ModelPricing(
|
|
117
|
+
Decimal("10.00"),
|
|
118
|
+
Decimal("1.00"),
|
|
119
|
+
Decimal("50.00"),
|
|
120
|
+
cache_write_input_per_million=Decimal("12.50"),
|
|
121
|
+
),
|
|
122
|
+
),
|
|
123
|
+
(
|
|
124
|
+
("gpt-6.1-sol",),
|
|
125
|
+
ModelPricing(
|
|
126
|
+
Decimal("2.00"),
|
|
127
|
+
Decimal("0.10"),
|
|
128
|
+
Decimal("10.00"),
|
|
129
|
+
cache_write_input_per_million=Decimal("2.50"),
|
|
130
|
+
),
|
|
131
|
+
),
|
|
132
|
+
(
|
|
133
|
+
("gpt-6-sol",),
|
|
134
|
+
ModelPricing(
|
|
135
|
+
Decimal("2.00"),
|
|
136
|
+
Decimal("0.20"),
|
|
137
|
+
Decimal("10.00"),
|
|
138
|
+
cache_write_input_per_million=Decimal("2.50"),
|
|
139
|
+
),
|
|
140
|
+
),
|
|
141
|
+
(
|
|
142
|
+
("gpt-6-luna",),
|
|
143
|
+
ModelPricing(
|
|
144
|
+
Decimal("0.10"),
|
|
145
|
+
Decimal("0.01"),
|
|
146
|
+
Decimal("0.50"),
|
|
147
|
+
cache_write_input_per_million=Decimal("0.125"),
|
|
148
|
+
),
|
|
149
|
+
),
|
|
150
|
+
(
|
|
151
|
+
("gpt-5.6-terra",),
|
|
152
|
+
ModelPricing(
|
|
153
|
+
Decimal("2.00"),
|
|
154
|
+
Decimal("0.20"),
|
|
155
|
+
Decimal("12.00"),
|
|
156
|
+
cache_write_input_per_million=Decimal("2.50"),
|
|
157
|
+
),
|
|
158
|
+
),
|
|
159
|
+
(
|
|
160
|
+
("gpt-5.6-luna",),
|
|
161
|
+
ModelPricing(
|
|
162
|
+
Decimal("0.20"),
|
|
163
|
+
Decimal("0.02"),
|
|
164
|
+
Decimal("1.20"),
|
|
165
|
+
cache_write_input_per_million=Decimal("0.25"),
|
|
166
|
+
),
|
|
167
|
+
),
|
|
168
|
+
(
|
|
169
|
+
("gpt-5.6-sol", "gpt-5.6"),
|
|
170
|
+
ModelPricing(
|
|
171
|
+
Decimal("4.00"),
|
|
172
|
+
Decimal("0.40"),
|
|
173
|
+
Decimal("20.00"),
|
|
174
|
+
cache_write_input_per_million=Decimal("5.00"),
|
|
175
|
+
),
|
|
176
|
+
),
|
|
177
|
+
(("gpt-5.5",), ModelPricing(Decimal("5.00"), Decimal("0.50"), Decimal("30.00"))),
|
|
178
|
+
(("gpt-5.4-mini",), ModelPricing(Decimal("0.750"), Decimal("0.075"), Decimal("4.500"))),
|
|
179
|
+
(("gpt-5.4-nano",), ModelPricing(Decimal("0.20"), Decimal("0.02"), Decimal("1.25"))),
|
|
180
|
+
(("gpt-5.4",), ModelPricing(Decimal("2.50"), Decimal("0.25"), Decimal("15.00"))),
|
|
181
|
+
(("gpt-5.2",), ModelPricing(Decimal("1.75"), Decimal("0.175"), Decimal("14.00"))),
|
|
182
|
+
(("gpt-5-mini",), ModelPricing(Decimal("0.25"), Decimal("0.025"), Decimal("2.00"))),
|
|
183
|
+
(("gpt-5-nano",), ModelPricing(Decimal("0.05"), Decimal("0.005"), Decimal("0.40"))),
|
|
184
|
+
(("gpt-5",), ModelPricing(Decimal("1.25"), Decimal("0.125"), Decimal("10.00"))),
|
|
185
|
+
(("gpt-4.1-mini",), ModelPricing(Decimal("0.40"), Decimal("0.10"), Decimal("1.60"))),
|
|
186
|
+
(("gpt-4.1-nano",), ModelPricing(Decimal("0.10"), Decimal("0.025"), Decimal("0.40"))),
|
|
187
|
+
(("gpt-4.1",), ModelPricing(Decimal("2.00"), Decimal("0.50"), Decimal("8.00"))),
|
|
188
|
+
(("gpt-4o-mini",), ModelPricing(Decimal("0.15"), Decimal("0.075"), Decimal("0.60"))),
|
|
189
|
+
(("gpt-4o",), ModelPricing(Decimal("2.50"), Decimal("1.25"), Decimal("10.00"))),
|
|
190
|
+
(("o4-mini",), ModelPricing(Decimal("1.10"), Decimal("0.275"), Decimal("4.40"))),
|
|
191
|
+
(("o3",), ModelPricing(Decimal("2.00"), Decimal("0.50"), Decimal("8.00"))),
|
|
192
|
+
(("o1",), ModelPricing(Decimal("15.00"), Decimal("7.50"), Decimal("60.00"))),
|
|
193
|
+
]
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
_MODEL_NAME_DELIMITERS = {"-", "_", "/", ":", " ", "."}
|
|
197
|
+
_OFFICIAL_SNAPSHOT_SUFFIX = re.compile(r"-\d{4}-\d{2}-\d{2}")
|
|
198
|
+
_WARNED_FALLBACK_MODELS: set[str] = set()
|
|
199
|
+
_WARNING_LOCK = threading.Lock()
|
|
200
|
+
|
|
201
|
+
|
|
202
|
+
def _is_delimited_match(model: str, candidate: str, start: int, end: int) -> bool:
|
|
203
|
+
before = model[start - 1] if start > 0 else None
|
|
204
|
+
after = model[end] if end < len(model) else None
|
|
205
|
+
before_ok = before is None or before in _MODEL_NAME_DELIMITERS
|
|
206
|
+
after_ok = after is None or after in _MODEL_NAME_DELIMITERS
|
|
207
|
+
return before_ok and after_ok
|
|
208
|
+
|
|
209
|
+
|
|
210
|
+
def _match_rule(model: str, rules: Iterable[tuple[tuple[str, ...], _T]]) -> tuple[str, _T] | None:
|
|
211
|
+
matches: list[tuple[int, int, str, _T]] = []
|
|
212
|
+
for rule_index, (prefixes, payload) in enumerate(rules):
|
|
213
|
+
for prefix in prefixes:
|
|
214
|
+
start = model.find(prefix)
|
|
215
|
+
while start != -1:
|
|
216
|
+
end = start + len(prefix)
|
|
217
|
+
if _is_delimited_match(model, prefix, start, end):
|
|
218
|
+
matches.append((len(prefix), rule_index, prefix, payload))
|
|
219
|
+
break
|
|
220
|
+
start = model.find(prefix, start + 1)
|
|
221
|
+
|
|
222
|
+
if not matches:
|
|
223
|
+
return None
|
|
224
|
+
|
|
225
|
+
matches.sort(key=lambda item: (-item[0], item[1]))
|
|
226
|
+
_, _, prefix, payload = matches[0]
|
|
227
|
+
return prefix, payload
|
|
228
|
+
|
|
229
|
+
|
|
230
|
+
def _is_registered_model_name(model: str, pricing_match: tuple[str, ModelPricing] | None) -> bool:
|
|
231
|
+
if pricing_match is None:
|
|
232
|
+
return False
|
|
233
|
+
prefix, _ = pricing_match
|
|
234
|
+
if model == prefix:
|
|
235
|
+
return True
|
|
236
|
+
return model.startswith(prefix) and _OFFICIAL_SNAPSHOT_SUFFIX.fullmatch(model[len(prefix) :]) is not None
|
|
237
|
+
|
|
238
|
+
|
|
239
|
+
def _warn_for_catalog_fallback(
|
|
240
|
+
model: str,
|
|
241
|
+
capability_match: tuple[str, ModelCapabilities] | None,
|
|
242
|
+
pricing_match: tuple[str, ModelPricing] | None,
|
|
243
|
+
) -> None:
|
|
244
|
+
if _is_registered_model_name(model, pricing_match):
|
|
245
|
+
return
|
|
246
|
+
|
|
247
|
+
with _WARNING_LOCK:
|
|
248
|
+
if model in _WARNED_FALLBACK_MODELS:
|
|
249
|
+
return
|
|
250
|
+
_WARNED_FALLBACK_MODELS.add(model)
|
|
251
|
+
|
|
252
|
+
capability_source = repr(capability_match[0]) if capability_match is not None else "the default capabilities"
|
|
253
|
+
pricing_source = repr(pricing_match[0]) if pricing_match is not None else "no pricing entry (cost=None)"
|
|
254
|
+
warnings.warn(
|
|
255
|
+
f"Model {model!r} is not explicitly registered in tokenrail's model catalog; "
|
|
256
|
+
f"using capabilities from {capability_source} and pricing from {pricing_source}.",
|
|
257
|
+
ModelCatalogFallbackWarning,
|
|
258
|
+
stacklevel=3,
|
|
259
|
+
)
|
|
260
|
+
|
|
261
|
+
|
|
262
|
+
def _resolve_model_catalog(
|
|
263
|
+
model: str,
|
|
264
|
+
) -> tuple[tuple[str, ModelCapabilities] | None, tuple[str, ModelPricing] | None]:
|
|
265
|
+
capability_match = _match_rule(model, _CAPABILITY_RULES)
|
|
266
|
+
pricing_match = _match_rule(model, _PRICING_RULES)
|
|
267
|
+
_warn_for_catalog_fallback(model, capability_match, pricing_match)
|
|
268
|
+
return capability_match, pricing_match
|
|
269
|
+
|
|
270
|
+
|
|
271
|
+
def get_model_capabilities(model: str) -> ModelCapabilities:
|
|
272
|
+
"""Return the request-parameter capabilities for ``model``.
|
|
273
|
+
|
|
274
|
+
Matching is delimiter-aware substring matching against the checked-in
|
|
275
|
+
capability registry; unknown models fall back to a conservative default.
|
|
276
|
+
"""
|
|
277
|
+
capability_match, _ = _resolve_model_catalog(model)
|
|
278
|
+
return capability_match[1] if capability_match is not None else _DEFAULT_CAPABILITIES
|
|
279
|
+
|
|
280
|
+
|
|
281
|
+
def get_model_pricing(model: str, service_tier: str = "default") -> ModelPricing | None:
|
|
282
|
+
"""Return per-million-token pricing for ``model``, or ``None`` if unknown.
|
|
283
|
+
|
|
284
|
+
The checked-in registry only carries default-tier base prices; long-context
|
|
285
|
+
premiums and regional uplifts are not applied. For other service tiers the
|
|
286
|
+
default-tier price is returned as an approximation.
|
|
287
|
+
"""
|
|
288
|
+
_, pricing_match = _resolve_model_catalog(model)
|
|
289
|
+
return pricing_match[1] if pricing_match is not None else None
|
|
290
|
+
|
|
291
|
+
|
|
292
|
+
def calculate_cost(
|
|
293
|
+
model: str,
|
|
294
|
+
usage: UsageBreakdown,
|
|
295
|
+
payer: str | None,
|
|
296
|
+
service_tier: str = "default",
|
|
297
|
+
) -> CostBreakdown | None:
|
|
298
|
+
"""Compute the nominal USD cost of ``usage`` and attribute it to a payer.
|
|
299
|
+
|
|
300
|
+
Returns ``None`` when the model has no pricing entry. When ``payer`` is
|
|
301
|
+
``"openai"`` the cost is attributed to OpenAI instead of the developer.
|
|
302
|
+
"""
|
|
303
|
+
pricing = get_model_pricing(model, service_tier=service_tier)
|
|
304
|
+
if pricing is None:
|
|
305
|
+
return None
|
|
306
|
+
|
|
307
|
+
ordinary_input = max(usage.input_tokens - usage.cached_tokens - usage.cache_write_tokens, 0)
|
|
308
|
+
cached_rate = pricing.cached_input_per_million or Decimal("0")
|
|
309
|
+
cache_write_rate = pricing.cache_write_input_per_million or pricing.input_per_million
|
|
310
|
+
total = (
|
|
311
|
+
(Decimal(ordinary_input) * pricing.input_per_million)
|
|
312
|
+
+ (Decimal(usage.cached_tokens) * cached_rate)
|
|
313
|
+
+ (Decimal(usage.cache_write_tokens) * cache_write_rate)
|
|
314
|
+
+ (Decimal(usage.output_tokens) * pricing.output_per_million)
|
|
315
|
+
) / Decimal("1000000")
|
|
316
|
+
|
|
317
|
+
nominal = float(total)
|
|
318
|
+
if payer == "openai":
|
|
319
|
+
return CostBreakdown(nominal_usd=nominal, developer_usd=0.0, openai_usd=nominal, payer=payer)
|
|
320
|
+
return CostBreakdown(nominal_usd=nominal, developer_usd=nominal, openai_usd=0.0, payer=payer)
|
tokenrail/client.py
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Any
|
|
4
|
+
|
|
5
|
+
from .providers.base import BaseProvider
|
|
6
|
+
from .providers.openai import OpenAIProvider
|
|
7
|
+
from .types import NormalizedResponse
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class _ResponsesNamespace:
|
|
11
|
+
def __init__(self, provider: BaseProvider) -> None:
|
|
12
|
+
self._provider = provider
|
|
13
|
+
|
|
14
|
+
def create(self, **kwargs: Any) -> NormalizedResponse:
|
|
15
|
+
return self._provider.create(**kwargs)
|
|
16
|
+
|
|
17
|
+
def parse(self, **kwargs: Any) -> NormalizedResponse:
|
|
18
|
+
return self._provider.parse(**kwargs)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class RailClient:
|
|
22
|
+
"""Provider-agnostic client with a ``client.responses.create(...)`` surface.
|
|
23
|
+
|
|
24
|
+
Wraps a :class:`~tokenrail.providers.base.BaseProvider` and exposes it through
|
|
25
|
+
a ``responses`` namespace that mirrors the OpenAI SDK call shape while
|
|
26
|
+
returning :class:`~tokenrail.types.NormalizedResponse` objects.
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
def __init__(self, provider: BaseProvider) -> None:
|
|
30
|
+
self.provider = provider
|
|
31
|
+
self.responses = _ResponsesNamespace(provider)
|
|
32
|
+
|
|
33
|
+
@classmethod
|
|
34
|
+
def openai(
|
|
35
|
+
cls,
|
|
36
|
+
*,
|
|
37
|
+
api_key: str | None = None,
|
|
38
|
+
organization: str | None = None,
|
|
39
|
+
timeout: float | None = None,
|
|
40
|
+
base_url: str | None = None,
|
|
41
|
+
max_retries: int = 2,
|
|
42
|
+
client: Any | None = None,
|
|
43
|
+
) -> RailClient:
|
|
44
|
+
"""Build a :class:`RailClient` backed by the OpenAI Python SDK.
|
|
45
|
+
|
|
46
|
+
``max_retries`` configures the SDK's built-in retry behavior; tokenrail
|
|
47
|
+
does not add its own retry loop. Pass ``client`` to inject a pre-built
|
|
48
|
+
(or fake) OpenAI client instead of constructing one.
|
|
49
|
+
"""
|
|
50
|
+
provider = OpenAIProvider(
|
|
51
|
+
client=client,
|
|
52
|
+
api_key=api_key,
|
|
53
|
+
organization=organization,
|
|
54
|
+
timeout=timeout,
|
|
55
|
+
base_url=base_url,
|
|
56
|
+
max_retries=max_retries,
|
|
57
|
+
)
|
|
58
|
+
return cls(provider=provider)
|
tokenrail/executor.py
ADDED
|
@@ -0,0 +1,249 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import time
|
|
4
|
+
from collections import deque
|
|
5
|
+
from collections.abc import Callable, Sequence
|
|
6
|
+
from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
|
|
7
|
+
from typing import Any
|
|
8
|
+
|
|
9
|
+
from .monitor import RollingMetricsMonitor
|
|
10
|
+
from .sinks import ResultSink
|
|
11
|
+
from .types import BatchItem, NormalizedResponse, StatsSnapshot, TimingBreakdown, UsageBreakdown
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def batch_items_from_queries(queries: dict[str, Any], **shared_request_kwargs: Any) -> list[BatchItem]:
|
|
15
|
+
"""Build :class:`BatchItem` objects from an ``{id: input}`` mapping.
|
|
16
|
+
|
|
17
|
+
Each value becomes the request ``input``; ``shared_request_kwargs`` (e.g.
|
|
18
|
+
``model``, ``reasoning_effort``) are applied to every item.
|
|
19
|
+
"""
|
|
20
|
+
return [
|
|
21
|
+
BatchItem(id=str(item_id), request_kwargs={"input": messages, **shared_request_kwargs})
|
|
22
|
+
for item_id, messages in queries.items()
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _error_response(item_id: str, model: str, provider: str, error: Exception) -> NormalizedResponse:
|
|
27
|
+
return NormalizedResponse(
|
|
28
|
+
id=item_id,
|
|
29
|
+
model=model,
|
|
30
|
+
provider=provider,
|
|
31
|
+
output_text=None,
|
|
32
|
+
raw_response=None,
|
|
33
|
+
usage=UsageBreakdown.empty(),
|
|
34
|
+
billing=None,
|
|
35
|
+
cost=None,
|
|
36
|
+
timing=TimingBreakdown(started_at=0.0, completed_at=0.0, latency_seconds=0.0),
|
|
37
|
+
error=f"{type(error).__name__}: {error}",
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
class _SubmitRateLimiter:
|
|
42
|
+
def __init__(
|
|
43
|
+
self,
|
|
44
|
+
*,
|
|
45
|
+
max_rpm: int | None,
|
|
46
|
+
max_tpm: int | None,
|
|
47
|
+
window_seconds: float = 60.0,
|
|
48
|
+
time_fn: Callable[[], float] = time.time,
|
|
49
|
+
sleep_fn: Callable[[float], None] = time.sleep,
|
|
50
|
+
) -> None:
|
|
51
|
+
if max_rpm is not None and max_rpm < 1:
|
|
52
|
+
raise ValueError("max_rpm must be at least 1")
|
|
53
|
+
if max_tpm is not None and max_tpm < 1:
|
|
54
|
+
raise ValueError("max_tpm must be at least 1")
|
|
55
|
+
self.max_rpm = max_rpm
|
|
56
|
+
self.max_tpm = max_tpm
|
|
57
|
+
self.window_seconds = window_seconds
|
|
58
|
+
self.time_fn = time_fn
|
|
59
|
+
self.sleep_fn = sleep_fn
|
|
60
|
+
self._submitted_at: deque[float] = deque()
|
|
61
|
+
self._completed_events: deque[tuple[float, int]] = deque()
|
|
62
|
+
self._inflight_estimates: deque[int] = deque()
|
|
63
|
+
self._inflight_estimated_tokens = 0
|
|
64
|
+
self._completed_requests = 0
|
|
65
|
+
self._completed_tokens = 0
|
|
66
|
+
|
|
67
|
+
def _prune(self, now: float) -> None:
|
|
68
|
+
cutoff = now - self.window_seconds
|
|
69
|
+
while self._submitted_at and self._submitted_at[0] <= cutoff:
|
|
70
|
+
self._submitted_at.popleft()
|
|
71
|
+
while self._completed_events and self._completed_events[0][0] <= cutoff:
|
|
72
|
+
self._completed_events.popleft()
|
|
73
|
+
|
|
74
|
+
def _estimated_next_tokens(self) -> int:
|
|
75
|
+
if self._completed_requests == 0:
|
|
76
|
+
return 0
|
|
77
|
+
return (self._completed_tokens + self._completed_requests - 1) // self._completed_requests
|
|
78
|
+
|
|
79
|
+
def _rolling_completed_tokens(self) -> int:
|
|
80
|
+
return sum(tokens for _, tokens in self._completed_events)
|
|
81
|
+
|
|
82
|
+
def can_submit(self) -> bool:
|
|
83
|
+
now = self.time_fn()
|
|
84
|
+
self._prune(now)
|
|
85
|
+
if self.max_rpm is not None and len(self._submitted_at) >= self.max_rpm:
|
|
86
|
+
return False
|
|
87
|
+
if self.max_tpm is not None:
|
|
88
|
+
if self._completed_requests == 0 and self._submitted_at:
|
|
89
|
+
return False
|
|
90
|
+
estimated_next = self._estimated_next_tokens()
|
|
91
|
+
if not self._completed_events and self._inflight_estimated_tokens == 0:
|
|
92
|
+
return True
|
|
93
|
+
if self._rolling_completed_tokens() + self._inflight_estimated_tokens + estimated_next > self.max_tpm:
|
|
94
|
+
return False
|
|
95
|
+
return True
|
|
96
|
+
|
|
97
|
+
def retry_after(self) -> float | None:
|
|
98
|
+
now = self.time_fn()
|
|
99
|
+
self._prune(now)
|
|
100
|
+
waits: list[float] = []
|
|
101
|
+
if self.max_rpm is not None and len(self._submitted_at) >= self.max_rpm:
|
|
102
|
+
waits.append(self._submitted_at[0] + self.window_seconds - now)
|
|
103
|
+
if self.max_tpm is not None and self._completed_events:
|
|
104
|
+
projected = (
|
|
105
|
+
self._rolling_completed_tokens() + self._inflight_estimated_tokens + self._estimated_next_tokens()
|
|
106
|
+
)
|
|
107
|
+
if projected > self.max_tpm:
|
|
108
|
+
waits.append(self._completed_events[0][0] + self.window_seconds - now)
|
|
109
|
+
if waits:
|
|
110
|
+
return max(min(waits), 0.0)
|
|
111
|
+
if not self.can_submit():
|
|
112
|
+
return None
|
|
113
|
+
return 0.0
|
|
114
|
+
|
|
115
|
+
def wait_until_allowed(self) -> None:
|
|
116
|
+
while not self.can_submit():
|
|
117
|
+
self.sleep_fn(self.retry_after() or 0.01)
|
|
118
|
+
|
|
119
|
+
def record_submit(self) -> None:
|
|
120
|
+
now = self.time_fn()
|
|
121
|
+
self._prune(now)
|
|
122
|
+
self._submitted_at.append(now)
|
|
123
|
+
estimated_tokens = self._estimated_next_tokens() if self.max_tpm is not None else 0
|
|
124
|
+
self._inflight_estimates.append(estimated_tokens)
|
|
125
|
+
self._inflight_estimated_tokens += estimated_tokens
|
|
126
|
+
|
|
127
|
+
def record_completion(self, response: NormalizedResponse) -> None:
|
|
128
|
+
now = self.time_fn()
|
|
129
|
+
self._prune(now)
|
|
130
|
+
if self._inflight_estimates:
|
|
131
|
+
self._inflight_estimated_tokens -= self._inflight_estimates.popleft()
|
|
132
|
+
total_tokens = response.usage.total_tokens or (response.usage.input_tokens + response.usage.output_tokens)
|
|
133
|
+
self._completed_events.append((now, total_tokens))
|
|
134
|
+
self._completed_requests += 1
|
|
135
|
+
self._completed_tokens += total_tokens
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
class BatchExecutor:
|
|
139
|
+
"""Thread-based batch runner for :class:`~tokenrail.client.RailClient` requests.
|
|
140
|
+
|
|
141
|
+
Submits items to a thread pool while honoring optional client-side
|
|
142
|
+
``max_rpm`` / ``max_tpm`` submit limits, writes each result to the
|
|
143
|
+
configured sinks, and records metrics on the monitor. Items whose ids are
|
|
144
|
+
already present in the first sink are skipped, which makes re-runs
|
|
145
|
+
resumable. Request errors are captured as error responses rather than
|
|
146
|
+
raised, so a single failing item does not abort the batch. Prompt-cache
|
|
147
|
+
settings belong in each item's request kwargs and are forwarded unchanged.
|
|
148
|
+
"""
|
|
149
|
+
|
|
150
|
+
def __init__(
|
|
151
|
+
self,
|
|
152
|
+
*,
|
|
153
|
+
client: Any,
|
|
154
|
+
max_workers: int = 20,
|
|
155
|
+
max_rpm: int | None = None,
|
|
156
|
+
max_tpm: int | None = None,
|
|
157
|
+
sinks: Sequence[ResultSink] | None = None,
|
|
158
|
+
monitor: RollingMetricsMonitor | None = None,
|
|
159
|
+
) -> None:
|
|
160
|
+
self.client = client
|
|
161
|
+
self.max_workers = max_workers
|
|
162
|
+
self.max_rpm = max_rpm
|
|
163
|
+
self.max_tpm = max_tpm
|
|
164
|
+
self.sinks = list(sinks or [])
|
|
165
|
+
self.monitor = monitor or RollingMetricsMonitor()
|
|
166
|
+
self._time_fn = time.time
|
|
167
|
+
self._sleep_fn = time.sleep
|
|
168
|
+
|
|
169
|
+
def _save(self, response: NormalizedResponse) -> None:
|
|
170
|
+
for sink in self.sinks:
|
|
171
|
+
sink.save(response)
|
|
172
|
+
|
|
173
|
+
def _load_done_ids(self) -> set[str]:
|
|
174
|
+
if not self.sinks:
|
|
175
|
+
return set()
|
|
176
|
+
return self.sinks[0].load_done_ids()
|
|
177
|
+
|
|
178
|
+
def _prepare_items(self, items: Sequence[BatchItem] | dict[str, Any]) -> list[BatchItem]:
|
|
179
|
+
if isinstance(items, dict):
|
|
180
|
+
return batch_items_from_queries(items)
|
|
181
|
+
return [BatchItem(id=str(item.id), request_kwargs=dict(item.request_kwargs)) for item in items]
|
|
182
|
+
|
|
183
|
+
def _request_kwargs(self, item: BatchItem) -> dict[str, Any]:
|
|
184
|
+
request_kwargs = dict(item.request_kwargs)
|
|
185
|
+
request_kwargs.setdefault("request_id", item.id)
|
|
186
|
+
return request_kwargs
|
|
187
|
+
|
|
188
|
+
def _call_single(self, item: BatchItem) -> NormalizedResponse:
|
|
189
|
+
request_kwargs = self._request_kwargs(item)
|
|
190
|
+
try:
|
|
191
|
+
if request_kwargs.get("text_format") is not None:
|
|
192
|
+
return self.client.responses.parse(**request_kwargs)
|
|
193
|
+
return self.client.responses.create(**request_kwargs)
|
|
194
|
+
except Exception as exc:
|
|
195
|
+
model = str(request_kwargs.get("model") or getattr(self.client.provider, "model_id", "unknown"))
|
|
196
|
+
return _error_response(item.id, model=model, provider=self.client.provider.name, error=exc)
|
|
197
|
+
|
|
198
|
+
def _run_threaded(self, items: list[BatchItem]) -> None:
|
|
199
|
+
limiter = _SubmitRateLimiter(
|
|
200
|
+
max_rpm=self.max_rpm,
|
|
201
|
+
max_tpm=self.max_tpm,
|
|
202
|
+
time_fn=self._time_fn,
|
|
203
|
+
sleep_fn=self._sleep_fn,
|
|
204
|
+
)
|
|
205
|
+
next_item = 0
|
|
206
|
+
pending: set[Future[NormalizedResponse]] = set()
|
|
207
|
+
|
|
208
|
+
with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
|
|
209
|
+
while next_item < len(items) or pending:
|
|
210
|
+
while next_item < len(items) and len(pending) < self.max_workers and limiter.can_submit():
|
|
211
|
+
limiter.record_submit()
|
|
212
|
+
pending.add(executor.submit(self._call_single, items[next_item]))
|
|
213
|
+
next_item += 1
|
|
214
|
+
|
|
215
|
+
if not pending:
|
|
216
|
+
if next_item < len(items):
|
|
217
|
+
limiter.wait_until_allowed()
|
|
218
|
+
continue
|
|
219
|
+
|
|
220
|
+
timeout = None
|
|
221
|
+
if next_item < len(items) and len(pending) < self.max_workers and not limiter.can_submit():
|
|
222
|
+
timeout = limiter.retry_after()
|
|
223
|
+
done, pending = wait(pending, timeout=timeout, return_when=FIRST_COMPLETED)
|
|
224
|
+
if not done:
|
|
225
|
+
continue
|
|
226
|
+
|
|
227
|
+
for future in done:
|
|
228
|
+
response = future.result()
|
|
229
|
+
limiter.record_completion(response)
|
|
230
|
+
self._save(response)
|
|
231
|
+
self.monitor.record(response)
|
|
232
|
+
|
|
233
|
+
def run(self, items: Sequence[BatchItem] | dict[str, Any]) -> StatsSnapshot:
|
|
234
|
+
"""Execute ``items`` (a sequence of :class:`BatchItem` or an ``{id: input}``
|
|
235
|
+
dict) and return the final :class:`~tokenrail.types.StatsSnapshot`."""
|
|
236
|
+
self.monitor.reset()
|
|
237
|
+
normalized_items = self._prepare_items(items)
|
|
238
|
+
done_ids = self._load_done_ids()
|
|
239
|
+
todo = [item for item in normalized_items if item.id not in done_ids]
|
|
240
|
+
skipped = len(normalized_items) - len(todo)
|
|
241
|
+
self.monitor.start(
|
|
242
|
+
total_requests=len(normalized_items),
|
|
243
|
+
todo_requests=len(todo),
|
|
244
|
+
skipped_requests=skipped,
|
|
245
|
+
)
|
|
246
|
+
|
|
247
|
+
self._run_threaded(todo)
|
|
248
|
+
|
|
249
|
+
return self.monitor.finalize(total_requests=len(normalized_items), skipped_requests=skipped)
|