noburn 0.1.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.
- noburn/__init__.py +4 -0
- noburn/guard.py +569 -0
- noburn/py.typed +0 -0
- noburn-0.1.0.dist-info/METADATA +112 -0
- noburn-0.1.0.dist-info/RECORD +6 -0
- noburn-0.1.0.dist-info/WHEEL +4 -0
noburn/__init__.py
ADDED
noburn/guard.py
ADDED
|
@@ -0,0 +1,569 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import re
|
|
4
|
+
import time
|
|
5
|
+
import json
|
|
6
|
+
import queue
|
|
7
|
+
import socket
|
|
8
|
+
import threading
|
|
9
|
+
import urllib.request
|
|
10
|
+
import urllib.error
|
|
11
|
+
from dataclasses import dataclass, field
|
|
12
|
+
from typing import Any
|
|
13
|
+
|
|
14
|
+
DEFAULT_BASE_URL = "https://noburn.dev"
|
|
15
|
+
RULES_TTL_SECONDS = 60
|
|
16
|
+
_MAX_EVENT_QUEUE = 1000
|
|
17
|
+
_MAX_EVENT_RETRIES = 3
|
|
18
|
+
|
|
19
|
+
_PRICING: list[tuple[str, float, float]] = [
|
|
20
|
+
("gpt-4o-mini", 0.15, 0.60),
|
|
21
|
+
("gpt-4o", 2.50, 10.00),
|
|
22
|
+
("gpt-4-turbo", 10.00, 30.00),
|
|
23
|
+
("gpt-4", 30.00, 60.00),
|
|
24
|
+
("gpt-3.5-turbo", 0.50, 1.50),
|
|
25
|
+
("o3-mini", 1.10, 4.40),
|
|
26
|
+
("o1-mini", 3.00, 12.00),
|
|
27
|
+
("o1", 15.00, 60.00),
|
|
28
|
+
("claude-haiku", 0.80, 4.00),
|
|
29
|
+
("claude-sonnet", 3.00, 15.00),
|
|
30
|
+
("claude-opus", 15.00, 75.00),
|
|
31
|
+
("gemini-2.0-flash", 0.10, 0.40),
|
|
32
|
+
("gemini-1.5-flash", 0.075, 0.30),
|
|
33
|
+
("gemini-1.5-pro", 1.25, 5.00),
|
|
34
|
+
("gemini-1.0-pro", 0.50, 1.50),
|
|
35
|
+
("deepseek", 0.14, 0.28),
|
|
36
|
+
("mistral-large", 2.00, 6.00),
|
|
37
|
+
("mistral-small", 0.20, 0.60),
|
|
38
|
+
("llama-3", 0.20, 0.20),
|
|
39
|
+
]
|
|
40
|
+
|
|
41
|
+
_OP_RE = re.compile(r"^(>=|<=|>|<|==?)\s*([0-9.]+)$")
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _estimate_cost(model: str, tokens_in: int, tokens_out: int) -> float:
|
|
45
|
+
lower = model.lower()
|
|
46
|
+
entry = next((p for p in _PRICING if lower.startswith(p[0])), None)
|
|
47
|
+
if entry is None:
|
|
48
|
+
entry = next((p for p in _PRICING if p[0] in lower), None)
|
|
49
|
+
if entry is None and lower.startswith("claude"):
|
|
50
|
+
for family in ("haiku", "sonnet", "opus"):
|
|
51
|
+
if family in lower:
|
|
52
|
+
entry = next((p for p in _PRICING if p[0] == f"claude-{family}"), None)
|
|
53
|
+
if entry:
|
|
54
|
+
break
|
|
55
|
+
if entry is None:
|
|
56
|
+
return 0.0
|
|
57
|
+
return (tokens_in * entry[1] + tokens_out * entry[2]) / 1_000_000
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def _eval_operator(condition_value: str, runtime_value: float) -> bool:
|
|
61
|
+
m = _OP_RE.match(condition_value.strip())
|
|
62
|
+
if not m:
|
|
63
|
+
return False
|
|
64
|
+
op, num_str = m.group(1), m.group(2)
|
|
65
|
+
threshold = float(num_str)
|
|
66
|
+
return {
|
|
67
|
+
">=": runtime_value >= threshold,
|
|
68
|
+
"<=": runtime_value <= threshold,
|
|
69
|
+
">": runtime_value > threshold,
|
|
70
|
+
"<": runtime_value < threshold,
|
|
71
|
+
"=": runtime_value == threshold,
|
|
72
|
+
"==": runtime_value == threshold,
|
|
73
|
+
}.get(op, False)
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _matches(
|
|
77
|
+
when: dict[str, str],
|
|
78
|
+
context: dict[str, str],
|
|
79
|
+
numerics: dict[str, float],
|
|
80
|
+
) -> bool:
|
|
81
|
+
for key, value in when.items():
|
|
82
|
+
if key in numerics:
|
|
83
|
+
if not _eval_operator(value, numerics[key]):
|
|
84
|
+
return False
|
|
85
|
+
elif context.get(key) != value:
|
|
86
|
+
return False
|
|
87
|
+
return True
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
# ─── Errors ───────────────────────────────────────────────────────────────────
|
|
91
|
+
|
|
92
|
+
class NoburnAuthError(Exception):
|
|
93
|
+
"""Raised when the SDK key is invalid or unauthorised."""
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
class NoburnTimeoutError(Exception):
|
|
97
|
+
"""Raised when the noburn API does not respond within the timeout."""
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
# ─── Result types ─────────────────────────────────────────────────────────────
|
|
101
|
+
|
|
102
|
+
@dataclass
|
|
103
|
+
class CheckResult:
|
|
104
|
+
blocked: bool
|
|
105
|
+
block_reason: str | None
|
|
106
|
+
forced_model_tier: str | None # 'small' | 'big' | None
|
|
107
|
+
spend_usd: float
|
|
108
|
+
budget_cap_usd: float | None
|
|
109
|
+
remaining_usd: float | None
|
|
110
|
+
user_spend_usd: float | None = field(default=None)
|
|
111
|
+
user_cap_usd: float | None = field(default=None)
|
|
112
|
+
user_remaining_usd: float | None = field(default=None)
|
|
113
|
+
run_spend_usd: float | None = field(default=None)
|
|
114
|
+
run_cap_usd: float | None = field(default=None)
|
|
115
|
+
run_remaining_usd: float | None = field(default=None)
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
@dataclass
|
|
119
|
+
class StartRunResult:
|
|
120
|
+
run_id: str
|
|
121
|
+
budget_cap_usd: float | None
|
|
122
|
+
started_at: str
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
@dataclass
|
|
126
|
+
class EndRunResult:
|
|
127
|
+
run_id: str
|
|
128
|
+
status: str # 'completed' | 'budget_exceeded'
|
|
129
|
+
spend_usd: float
|
|
130
|
+
budget_cap_usd: float | None
|
|
131
|
+
started_at: str
|
|
132
|
+
ended_at: str
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
# ─── NoburnGuard ─────────────────────────────────────────────────────────────
|
|
136
|
+
|
|
137
|
+
class NoburnGuard:
|
|
138
|
+
"""
|
|
139
|
+
Drop-in LLM spend enforcement for Python.
|
|
140
|
+
|
|
141
|
+
Usage::
|
|
142
|
+
|
|
143
|
+
guard = NoburnGuard(api_key="sk-nb-...", project_id="...", budget_cap_usd=10.0)
|
|
144
|
+
|
|
145
|
+
# Agent run with per-run budget cap
|
|
146
|
+
run = guard.start_run(budget_cap_usd=0.50)
|
|
147
|
+
|
|
148
|
+
result = guard.check(
|
|
149
|
+
model="gpt-4o",
|
|
150
|
+
estimated_tokens_in=1000,
|
|
151
|
+
estimated_tokens_out=300,
|
|
152
|
+
run_id=run.run_id,
|
|
153
|
+
)
|
|
154
|
+
if result.blocked:
|
|
155
|
+
raise RuntimeError(result.block_reason)
|
|
156
|
+
|
|
157
|
+
# ... LLM call ...
|
|
158
|
+
|
|
159
|
+
guard.record(model="gpt-4o", tokens_in=980, tokens_out=290,
|
|
160
|
+
cost_usd=0.005, was_blocked=False, run_id=run.run_id)
|
|
161
|
+
|
|
162
|
+
guard.end_run(run.run_id)
|
|
163
|
+
"""
|
|
164
|
+
|
|
165
|
+
def __init__(
|
|
166
|
+
self,
|
|
167
|
+
api_key: str,
|
|
168
|
+
project_id: str,
|
|
169
|
+
budget_cap_usd: float | None = None,
|
|
170
|
+
base_url: str = DEFAULT_BASE_URL,
|
|
171
|
+
timeout: float = 2.0,
|
|
172
|
+
on_error: str = "allow",
|
|
173
|
+
) -> None:
|
|
174
|
+
if not api_key:
|
|
175
|
+
raise ValueError("api_key is required")
|
|
176
|
+
if not project_id:
|
|
177
|
+
raise ValueError("project_id is required")
|
|
178
|
+
if on_error not in ("allow", "block"):
|
|
179
|
+
raise ValueError("on_error must be 'allow' or 'block'")
|
|
180
|
+
|
|
181
|
+
self._api_key = api_key
|
|
182
|
+
self._project_id = project_id
|
|
183
|
+
self._budget_cap_usd = budget_cap_usd
|
|
184
|
+
self._base_url = base_url.rstrip("/")
|
|
185
|
+
self._timeout = timeout
|
|
186
|
+
# When the noburn API is unreachable and the local fallback has no
|
|
187
|
+
# opinion: 'allow' fails open (default — an outage never breaks your
|
|
188
|
+
# app); 'block' fails closed (blocks with reason 'noburn_unreachable').
|
|
189
|
+
self._on_error = on_error
|
|
190
|
+
|
|
191
|
+
self._rules: list[dict[str, Any]] = []
|
|
192
|
+
self._rules_loaded_at: float = 0.0
|
|
193
|
+
self._spend_usd: float = 0.0
|
|
194
|
+
# Guards _spend_usd / _reserved_usd against concurrent threads.
|
|
195
|
+
self._lock = threading.Lock()
|
|
196
|
+
# In-flight cost reserved by local-fallback checks, not yet recorded.
|
|
197
|
+
self._reserved_usd: float = 0.0
|
|
198
|
+
# Single background worker drains a bounded retry queue, so record()
|
|
199
|
+
# never blocks, never loses an event to a transient blip, and never
|
|
200
|
+
# spawns an unbounded number of threads under load.
|
|
201
|
+
self._event_queue: "queue.Queue[tuple[bytes, int]]" = queue.Queue(maxsize=_MAX_EVENT_QUEUE)
|
|
202
|
+
self._dropped_events: int = 0
|
|
203
|
+
self._worker_started = False
|
|
204
|
+
self._worker_lock = threading.Lock()
|
|
205
|
+
|
|
206
|
+
# ─── Agent runs ───────────────────────────────────────────────────────────
|
|
207
|
+
|
|
208
|
+
def start_run(
|
|
209
|
+
self,
|
|
210
|
+
budget_cap_usd: float | None = None,
|
|
211
|
+
end_user_id: str | None = None,
|
|
212
|
+
metadata: dict[str, Any] | None = None,
|
|
213
|
+
) -> StartRunResult:
|
|
214
|
+
"""
|
|
215
|
+
Start an agent run. Returns a StartRunResult with run_id to attach to
|
|
216
|
+
every check() and record() call in the run.
|
|
217
|
+
Raises on network or auth error — a run must be registered to proceed.
|
|
218
|
+
"""
|
|
219
|
+
payload = json.dumps({
|
|
220
|
+
"project_id": self._project_id,
|
|
221
|
+
"budget_cap_usd": budget_cap_usd,
|
|
222
|
+
"end_user_id": end_user_id,
|
|
223
|
+
"metadata": metadata,
|
|
224
|
+
}).encode()
|
|
225
|
+
|
|
226
|
+
req = urllib.request.Request(
|
|
227
|
+
f"{self._base_url}/api/v1/runs",
|
|
228
|
+
data=payload,
|
|
229
|
+
headers={
|
|
230
|
+
"Authorization": f"Bearer {self._api_key}",
|
|
231
|
+
"Content-Type": "application/json",
|
|
232
|
+
},
|
|
233
|
+
method="POST",
|
|
234
|
+
)
|
|
235
|
+
try:
|
|
236
|
+
with urllib.request.urlopen(req, timeout=self._timeout) as resp:
|
|
237
|
+
body = json.loads(resp.read())
|
|
238
|
+
return StartRunResult(
|
|
239
|
+
run_id=body["run_id"],
|
|
240
|
+
budget_cap_usd=body.get("budget_cap_usd"),
|
|
241
|
+
started_at=body["started_at"],
|
|
242
|
+
)
|
|
243
|
+
except urllib.error.HTTPError as exc:
|
|
244
|
+
if exc.code in (401, 403):
|
|
245
|
+
raise NoburnAuthError("Invalid SDK key") from exc
|
|
246
|
+
raise
|
|
247
|
+
except (TimeoutError, socket.timeout) as exc:
|
|
248
|
+
raise NoburnTimeoutError("noburn API timed out") from exc
|
|
249
|
+
except urllib.error.URLError as exc:
|
|
250
|
+
if isinstance(exc.reason, (TimeoutError, socket.timeout)):
|
|
251
|
+
raise NoburnTimeoutError("noburn API timed out") from exc
|
|
252
|
+
raise
|
|
253
|
+
|
|
254
|
+
def end_run(self, run_id: str) -> EndRunResult | None:
|
|
255
|
+
"""
|
|
256
|
+
End a run. Fire-and-forget — returns None on any network error.
|
|
257
|
+
"""
|
|
258
|
+
payload = json.dumps({"project_id": self._project_id}).encode()
|
|
259
|
+
req = urllib.request.Request(
|
|
260
|
+
f"{self._base_url}/api/v1/runs/{run_id}/end",
|
|
261
|
+
data=payload,
|
|
262
|
+
headers={
|
|
263
|
+
"Authorization": f"Bearer {self._api_key}",
|
|
264
|
+
"Content-Type": "application/json",
|
|
265
|
+
},
|
|
266
|
+
method="POST",
|
|
267
|
+
)
|
|
268
|
+
try:
|
|
269
|
+
with urllib.request.urlopen(req, timeout=self._timeout) as resp:
|
|
270
|
+
body = json.loads(resp.read())
|
|
271
|
+
return EndRunResult(
|
|
272
|
+
run_id=body["run_id"],
|
|
273
|
+
status=body["status"],
|
|
274
|
+
spend_usd=float(body.get("spend_usd", 0)),
|
|
275
|
+
budget_cap_usd=body.get("budget_cap_usd"),
|
|
276
|
+
started_at=body["started_at"],
|
|
277
|
+
ended_at=body["ended_at"],
|
|
278
|
+
)
|
|
279
|
+
except urllib.error.HTTPError as exc:
|
|
280
|
+
if exc.code in (401, 403):
|
|
281
|
+
raise NoburnAuthError("Invalid SDK key") from exc
|
|
282
|
+
return None
|
|
283
|
+
except Exception:
|
|
284
|
+
return None
|
|
285
|
+
|
|
286
|
+
# ─── Internal ─────────────────────────────────────────────────────────────
|
|
287
|
+
|
|
288
|
+
def _refresh_rules(self) -> None:
|
|
289
|
+
if time.monotonic() - self._rules_loaded_at < RULES_TTL_SECONDS:
|
|
290
|
+
return
|
|
291
|
+
|
|
292
|
+
url = f"{self._base_url}/api/v1/policy?project_id={self._project_id}"
|
|
293
|
+
req = urllib.request.Request(
|
|
294
|
+
url,
|
|
295
|
+
headers={"Authorization": f"Bearer {self._api_key}"},
|
|
296
|
+
)
|
|
297
|
+
try:
|
|
298
|
+
with urllib.request.urlopen(req, timeout=self._timeout) as resp:
|
|
299
|
+
body = json.loads(resp.read())
|
|
300
|
+
self._rules = body.get("rules", [])
|
|
301
|
+
self._rules_loaded_at = time.monotonic()
|
|
302
|
+
except urllib.error.HTTPError as exc:
|
|
303
|
+
if exc.code in (401, 403):
|
|
304
|
+
raise NoburnAuthError("Invalid SDK key") from exc
|
|
305
|
+
except Exception:
|
|
306
|
+
pass
|
|
307
|
+
|
|
308
|
+
def _server_check(
|
|
309
|
+
self,
|
|
310
|
+
model: str,
|
|
311
|
+
estimated_cost: float,
|
|
312
|
+
end_user_id: str | None,
|
|
313
|
+
run_id: str | None,
|
|
314
|
+
context: dict[str, str] | None,
|
|
315
|
+
tokens_in: int = 0,
|
|
316
|
+
tokens_out: int = 0,
|
|
317
|
+
) -> CheckResult | None:
|
|
318
|
+
payload = json.dumps({
|
|
319
|
+
"project_id": self._project_id,
|
|
320
|
+
"model": model,
|
|
321
|
+
"end_user_id": end_user_id,
|
|
322
|
+
"run_id": run_id,
|
|
323
|
+
"estimated_cost_usd": estimated_cost,
|
|
324
|
+
# Send raw token counts so the server can price the call
|
|
325
|
+
# authoritatively rather than trusting estimated_cost_usd.
|
|
326
|
+
"estimated_tokens_in": max(0, tokens_in),
|
|
327
|
+
"estimated_tokens_out": max(0, tokens_out),
|
|
328
|
+
"context": context,
|
|
329
|
+
}).encode()
|
|
330
|
+
|
|
331
|
+
req = urllib.request.Request(
|
|
332
|
+
f"{self._base_url}/api/v1/check",
|
|
333
|
+
data=payload,
|
|
334
|
+
headers={
|
|
335
|
+
"Authorization": f"Bearer {self._api_key}",
|
|
336
|
+
"Content-Type": "application/json",
|
|
337
|
+
},
|
|
338
|
+
method="POST",
|
|
339
|
+
)
|
|
340
|
+
try:
|
|
341
|
+
with urllib.request.urlopen(req, timeout=self._timeout) as resp:
|
|
342
|
+
body = json.loads(resp.read())
|
|
343
|
+
return CheckResult(
|
|
344
|
+
blocked=not body["allowed"],
|
|
345
|
+
block_reason=body.get("reason"),
|
|
346
|
+
forced_model_tier=body.get("forced_model_tier"),
|
|
347
|
+
spend_usd=float(body.get("project_spend_usd", 0)),
|
|
348
|
+
budget_cap_usd=body.get("project_cap_usd"),
|
|
349
|
+
remaining_usd=body.get("project_remaining_usd"),
|
|
350
|
+
user_spend_usd=float(body.get("user_spend_usd", 0)),
|
|
351
|
+
user_cap_usd=body.get("user_cap_usd"),
|
|
352
|
+
user_remaining_usd=body.get("user_remaining_usd"),
|
|
353
|
+
run_spend_usd=float(body.get("run_spend_usd", 0)),
|
|
354
|
+
run_cap_usd=body.get("run_cap_usd"),
|
|
355
|
+
run_remaining_usd=body.get("run_remaining_usd"),
|
|
356
|
+
)
|
|
357
|
+
except urllib.error.HTTPError as exc:
|
|
358
|
+
if exc.code in (401, 403):
|
|
359
|
+
raise NoburnAuthError("Invalid SDK key") from exc
|
|
360
|
+
return None
|
|
361
|
+
except Exception:
|
|
362
|
+
return None
|
|
363
|
+
|
|
364
|
+
# ─── Public API ───────────────────────────────────────────────────────────
|
|
365
|
+
|
|
366
|
+
def check(
|
|
367
|
+
self,
|
|
368
|
+
model: str,
|
|
369
|
+
estimated_tokens_in: int,
|
|
370
|
+
estimated_tokens_out: int,
|
|
371
|
+
end_user_id: str | None = None,
|
|
372
|
+
run_id: str | None = None,
|
|
373
|
+
cost_usd: float | None = None,
|
|
374
|
+
context: dict[str, str] | None = None,
|
|
375
|
+
) -> CheckResult:
|
|
376
|
+
"""
|
|
377
|
+
Check whether a planned LLM call should proceed.
|
|
378
|
+
Call this *before* every LLM API call.
|
|
379
|
+
Never raises except NoburnAuthError (bad key).
|
|
380
|
+
|
|
381
|
+
Every call hits POST /api/v1/check for a server-authoritative decision.
|
|
382
|
+
Falls back to local check if the server is unreachable.
|
|
383
|
+
"""
|
|
384
|
+
# Clamp to non-negative so a bad caller can't drive the cost negative.
|
|
385
|
+
tokens_in = max(0, estimated_tokens_in)
|
|
386
|
+
tokens_out = max(0, estimated_tokens_out)
|
|
387
|
+
estimated_cost = cost_usd if cost_usd is not None else _estimate_cost(
|
|
388
|
+
model, tokens_in, tokens_out
|
|
389
|
+
)
|
|
390
|
+
|
|
391
|
+
server_result = self._server_check(
|
|
392
|
+
model, estimated_cost, end_user_id, run_id, context, tokens_in, tokens_out
|
|
393
|
+
)
|
|
394
|
+
if server_result is not None:
|
|
395
|
+
return server_result
|
|
396
|
+
|
|
397
|
+
# Server unreachable. Enforce against actual + in-flight reserved spend
|
|
398
|
+
# under a lock so concurrent fallback checks (threads) don't all pass the
|
|
399
|
+
# same pre-record budget and don't lose updates to the counter.
|
|
400
|
+
with self._lock:
|
|
401
|
+
effective = self._spend_usd + self._reserved_usd
|
|
402
|
+
spend_snapshot = self._spend_usd
|
|
403
|
+
remaining = (
|
|
404
|
+
self._budget_cap_usd - effective
|
|
405
|
+
if self._budget_cap_usd is not None else None
|
|
406
|
+
)
|
|
407
|
+
if self._on_error == "block":
|
|
408
|
+
return CheckResult(
|
|
409
|
+
blocked=True, block_reason="noburn_unreachable", forced_model_tier=None,
|
|
410
|
+
spend_usd=spend_snapshot, budget_cap_usd=self._budget_cap_usd,
|
|
411
|
+
remaining_usd=remaining,
|
|
412
|
+
)
|
|
413
|
+
utilization = (
|
|
414
|
+
effective / self._budget_cap_usd
|
|
415
|
+
if self._budget_cap_usd and self._budget_cap_usd > 0 else 0.0
|
|
416
|
+
)
|
|
417
|
+
|
|
418
|
+
try:
|
|
419
|
+
self._refresh_rules()
|
|
420
|
+
except NoburnAuthError:
|
|
421
|
+
raise
|
|
422
|
+
|
|
423
|
+
ctx: dict[str, str] = {"model": model, **(context or {})}
|
|
424
|
+
if end_user_id is not None:
|
|
425
|
+
ctx["end_user_id"] = end_user_id
|
|
426
|
+
if run_id is not None:
|
|
427
|
+
ctx["run_id"] = run_id
|
|
428
|
+
numerics = {"utilization": utilization}
|
|
429
|
+
|
|
430
|
+
blocked = False
|
|
431
|
+
block_reason: str | None = None
|
|
432
|
+
forced_tier: str | None = None
|
|
433
|
+
|
|
434
|
+
for rule in self._rules:
|
|
435
|
+
when: dict[str, str] = rule.get("when", {})
|
|
436
|
+
then: str = rule.get("then", "")
|
|
437
|
+
if not _matches(when, ctx, numerics):
|
|
438
|
+
continue
|
|
439
|
+
if then == "allow":
|
|
440
|
+
break
|
|
441
|
+
if then == "block":
|
|
442
|
+
blocked, block_reason = True, "policy_rule"
|
|
443
|
+
break
|
|
444
|
+
if then in ("force_small", "force_big"):
|
|
445
|
+
forced_tier = "small" if then == "force_small" else "big"
|
|
446
|
+
break
|
|
447
|
+
|
|
448
|
+
if not blocked and forced_tier is None and \
|
|
449
|
+
self._budget_cap_usd is not None and effective + estimated_cost > self._budget_cap_usd:
|
|
450
|
+
blocked, block_reason = True, "budget_exceeded"
|
|
451
|
+
|
|
452
|
+
# Reserve the estimated cost for any call we let through; record()
|
|
453
|
+
# releases the reservation when the actual outcome arrives.
|
|
454
|
+
if not blocked:
|
|
455
|
+
with self._lock:
|
|
456
|
+
self._reserved_usd += estimated_cost
|
|
457
|
+
|
|
458
|
+
return CheckResult(
|
|
459
|
+
blocked=blocked, block_reason=block_reason, forced_model_tier=forced_tier,
|
|
460
|
+
spend_usd=spend_snapshot, budget_cap_usd=self._budget_cap_usd,
|
|
461
|
+
remaining_usd=max(0.0, remaining or 0.0) if block_reason == "budget_exceeded" else remaining,
|
|
462
|
+
)
|
|
463
|
+
|
|
464
|
+
def record(
|
|
465
|
+
self,
|
|
466
|
+
model: str,
|
|
467
|
+
tokens_in: int,
|
|
468
|
+
tokens_out: int,
|
|
469
|
+
cost_usd: float,
|
|
470
|
+
was_blocked: bool,
|
|
471
|
+
end_user_id: str | None = None,
|
|
472
|
+
run_id: str | None = None,
|
|
473
|
+
latency_ms: int | None = None,
|
|
474
|
+
block_reason: str | None = None,
|
|
475
|
+
timestamp: str | None = None,
|
|
476
|
+
) -> None:
|
|
477
|
+
"""
|
|
478
|
+
Record the actual outcome of an LLM call. Fire-and-forget: returns
|
|
479
|
+
immediately and never raises — the event is queued and delivered (with
|
|
480
|
+
retry) by a single background worker, so it never adds latency, never
|
|
481
|
+
loses an event to a transient blip, and never spawns a thread per call.
|
|
482
|
+
"""
|
|
483
|
+
with self._lock:
|
|
484
|
+
if not was_blocked:
|
|
485
|
+
self._spend_usd += cost_usd
|
|
486
|
+
# Release the reservation made by the corresponding fallback check().
|
|
487
|
+
self._reserved_usd = max(0.0, self._reserved_usd - cost_usd)
|
|
488
|
+
|
|
489
|
+
payload = json.dumps({
|
|
490
|
+
"project_id": self._project_id,
|
|
491
|
+
"model": model,
|
|
492
|
+
"tokens_in": tokens_in,
|
|
493
|
+
"tokens_out": tokens_out,
|
|
494
|
+
"cost_usd": cost_usd,
|
|
495
|
+
"was_blocked": was_blocked,
|
|
496
|
+
"block_reason": block_reason,
|
|
497
|
+
"end_user_id": end_user_id,
|
|
498
|
+
"run_id": run_id,
|
|
499
|
+
"latency_ms": latency_ms,
|
|
500
|
+
"timestamp": timestamp or time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
|
501
|
+
}).encode()
|
|
502
|
+
|
|
503
|
+
self._enqueue_event(payload)
|
|
504
|
+
|
|
505
|
+
@property
|
|
506
|
+
def dropped_event_count(self) -> int:
|
|
507
|
+
"""Events dropped via queue overflow or retry exhaustion."""
|
|
508
|
+
return self._dropped_events
|
|
509
|
+
|
|
510
|
+
def flush(self, timeout: float | None = None) -> None:
|
|
511
|
+
"""Block until the event queue is drained. Call before process exit."""
|
|
512
|
+
try:
|
|
513
|
+
# queue.join() has no timeout; approximate by polling unfinished tasks.
|
|
514
|
+
deadline = None if timeout is None else time.monotonic() + timeout
|
|
515
|
+
while self._event_queue.unfinished_tasks:
|
|
516
|
+
if deadline is not None and time.monotonic() > deadline:
|
|
517
|
+
return
|
|
518
|
+
time.sleep(0.01)
|
|
519
|
+
except Exception:
|
|
520
|
+
pass
|
|
521
|
+
|
|
522
|
+
def _enqueue_event(self, payload: bytes) -> None:
|
|
523
|
+
self._ensure_worker()
|
|
524
|
+
try:
|
|
525
|
+
self._event_queue.put_nowait((payload, 0))
|
|
526
|
+
except queue.Full:
|
|
527
|
+
self._dropped_events += 1
|
|
528
|
+
|
|
529
|
+
def _ensure_worker(self) -> None:
|
|
530
|
+
with self._worker_lock:
|
|
531
|
+
if self._worker_started:
|
|
532
|
+
return
|
|
533
|
+
self._worker_started = True
|
|
534
|
+
threading.Thread(target=self._worker, daemon=True).start()
|
|
535
|
+
|
|
536
|
+
def _worker(self) -> None:
|
|
537
|
+
while True:
|
|
538
|
+
payload, attempt = self._event_queue.get()
|
|
539
|
+
try:
|
|
540
|
+
if not self._send_event(payload):
|
|
541
|
+
if attempt + 1 < _MAX_EVENT_RETRIES:
|
|
542
|
+
time.sleep(0.2 * (2 ** attempt))
|
|
543
|
+
try:
|
|
544
|
+
self._event_queue.put_nowait((payload, attempt + 1))
|
|
545
|
+
except queue.Full:
|
|
546
|
+
self._dropped_events += 1
|
|
547
|
+
else:
|
|
548
|
+
self._dropped_events += 1
|
|
549
|
+
finally:
|
|
550
|
+
self._event_queue.task_done()
|
|
551
|
+
|
|
552
|
+
def _send_event(self, payload: bytes) -> bool:
|
|
553
|
+
req = urllib.request.Request(
|
|
554
|
+
f"{self._base_url}/api/v1/events",
|
|
555
|
+
data=payload,
|
|
556
|
+
headers={
|
|
557
|
+
"Authorization": f"Bearer {self._api_key}",
|
|
558
|
+
"Content-Type": "application/json",
|
|
559
|
+
},
|
|
560
|
+
method="POST",
|
|
561
|
+
)
|
|
562
|
+
try:
|
|
563
|
+
with urllib.request.urlopen(req, timeout=self._timeout):
|
|
564
|
+
return True
|
|
565
|
+
except urllib.error.HTTPError as exc:
|
|
566
|
+
# 4xx won't succeed on retry (bad key / bad body) — stop retrying.
|
|
567
|
+
return 400 <= exc.code < 500
|
|
568
|
+
except Exception:
|
|
569
|
+
return False
|
noburn/py.typed
ADDED
|
File without changes
|
|
@@ -0,0 +1,112 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: noburn
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Hard LLM spend enforcement for Python — pre-call budget checks and per-user metering.
|
|
5
|
+
Project-URL: Homepage, https://noburn.dev
|
|
6
|
+
Project-URL: Documentation, https://noburn.dev/docs
|
|
7
|
+
Project-URL: Repository, https://github.com/orvi2014/noburn
|
|
8
|
+
License: MIT
|
|
9
|
+
Keywords: ai,anthropic,budget,cost-control,langchain,llm,openai
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Intended Audience :: Developers
|
|
12
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
19
|
+
Requires-Python: >=3.9
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# noburn
|
|
23
|
+
|
|
24
|
+
Hard LLM spend enforcement for Python — pre-call budget checks and per-user metering.
|
|
25
|
+
|
|
26
|
+
Check a planned LLM call against project-, per-user-, and per-run budget caps
|
|
27
|
+
*before* you make it, then record what it actually cost.
|
|
28
|
+
|
|
29
|
+
## Install
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
pip install noburn
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
Requires Python ≥ 3.9. No third-party dependencies. Ships type hints (`py.typed`).
|
|
36
|
+
|
|
37
|
+
## Quick start
|
|
38
|
+
|
|
39
|
+
```python
|
|
40
|
+
import os
|
|
41
|
+
from noburn import NoburnGuard
|
|
42
|
+
|
|
43
|
+
guard = NoburnGuard(
|
|
44
|
+
api_key=os.environ["NOBURN_API_KEY"],
|
|
45
|
+
project_id=os.environ["NOBURN_PROJECT_ID"],
|
|
46
|
+
budget_cap_usd=10.0, # optional local fallback cap
|
|
47
|
+
on_error="allow", # "allow" (default, fail-open) | "block" (fail-closed)
|
|
48
|
+
)
|
|
49
|
+
|
|
50
|
+
decision = guard.check(
|
|
51
|
+
model="gpt-4o",
|
|
52
|
+
estimated_tokens_in=1000,
|
|
53
|
+
estimated_tokens_out=300,
|
|
54
|
+
end_user_id="user_123", # optional — enforces that user's cap
|
|
55
|
+
)
|
|
56
|
+
|
|
57
|
+
if decision.blocked:
|
|
58
|
+
raise RuntimeError(f"Blocked by noburn: {decision.block_reason}")
|
|
59
|
+
|
|
60
|
+
# ... make your LLM call ...
|
|
61
|
+
|
|
62
|
+
guard.record(
|
|
63
|
+
model="gpt-4o",
|
|
64
|
+
tokens_in=980,
|
|
65
|
+
tokens_out=290,
|
|
66
|
+
cost_usd=0.005,
|
|
67
|
+
was_blocked=False,
|
|
68
|
+
end_user_id="user_123",
|
|
69
|
+
)
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
`record()` is fire-and-forget — it returns immediately (delivery happens on a
|
|
73
|
+
background daemon thread) and never raises, so it adds no latency to your call
|
|
74
|
+
path.
|
|
75
|
+
|
|
76
|
+
## Agent runs
|
|
77
|
+
|
|
78
|
+
Bound a single agent invocation with its own cap:
|
|
79
|
+
|
|
80
|
+
```python
|
|
81
|
+
run = guard.start_run(budget_cap_usd=0.5)
|
|
82
|
+
|
|
83
|
+
decision = guard.check(
|
|
84
|
+
model="gpt-4o",
|
|
85
|
+
estimated_tokens_in=1000,
|
|
86
|
+
estimated_tokens_out=300,
|
|
87
|
+
run_id=run.run_id,
|
|
88
|
+
)
|
|
89
|
+
|
|
90
|
+
guard.record(model="gpt-4o", tokens_in=980, tokens_out=290,
|
|
91
|
+
cost_usd=0.005, was_blocked=False, run_id=run.run_id)
|
|
92
|
+
|
|
93
|
+
summary = guard.end_run(run.run_id) # EndRunResult(status=..., spend_usd=..., ...)
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
## Behavior
|
|
97
|
+
|
|
98
|
+
- **`check()` never raises** except `NoburnAuthError` (bad key). Every call asks
|
|
99
|
+
the server for an authoritative decision; if the server is unreachable it
|
|
100
|
+
falls back to a local in-memory check.
|
|
101
|
+
- **`on_error`** controls that fallback when nothing local blocks the call:
|
|
102
|
+
`"allow"` (default) fails open so a noburn outage never breaks your app;
|
|
103
|
+
`"block"` fails closed (`block_reason="noburn_unreachable"`).
|
|
104
|
+
- **`start_run()` raises** on network/auth error (`NoburnTimeoutError` on
|
|
105
|
+
timeout) — a run must be registered before you proceed.
|
|
106
|
+
|
|
107
|
+
> Note: `check()`/`start_run()` use blocking HTTP. In async apps, run them in a
|
|
108
|
+
> thread executor (e.g. `await asyncio.to_thread(guard.check, ...)`).
|
|
109
|
+
|
|
110
|
+
## License
|
|
111
|
+
|
|
112
|
+
MIT
|
|
@@ -0,0 +1,6 @@
|
|
|
1
|
+
noburn/__init__.py,sha256=NztIfn9xj2NFbv6OnOQ2MSMl7LPX8_irfNBNYs8VmcM,186
|
|
2
|
+
noburn/guard.py,sha256=7EwAiIS_zpov41oNvyLx7G2yaW6ljS3kkNcs6UGBgDk,21641
|
|
3
|
+
noburn/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
4
|
+
noburn-0.1.0.dist-info/METADATA,sha256=_U26hrOUrxyJZtDomnr37RPmSaE0NVPHJgmBurW2Nu8,3385
|
|
5
|
+
noburn-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
|
|
6
|
+
noburn-0.1.0.dist-info/RECORD,,
|