tripwire-litellm 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- tripwire_litellm-0.1.0/PKG-INFO +9 -0
- tripwire_litellm-0.1.0/README.md +47 -0
- tripwire_litellm-0.1.0/pyproject.toml +19 -0
- tripwire_litellm-0.1.0/setup.cfg +4 -0
- tripwire_litellm-0.1.0/tests/test_handler.py +307 -0
- tripwire_litellm-0.1.0/tripwire_litellm/__init__.py +5 -0
- tripwire_litellm-0.1.0/tripwire_litellm/handler.py +406 -0
- tripwire_litellm-0.1.0/tripwire_litellm.egg-info/PKG-INFO +9 -0
- tripwire_litellm-0.1.0/tripwire_litellm.egg-info/SOURCES.txt +10 -0
- tripwire_litellm-0.1.0/tripwire_litellm.egg-info/dependency_links.txt +1 -0
- tripwire_litellm-0.1.0/tripwire_litellm.egg-info/requires.txt +5 -0
- tripwire_litellm-0.1.0/tripwire_litellm.egg-info/top_level.txt +1 -0
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tripwire-litellm
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Tripwire hard-stop and ingest callback for a customer's LiteLLM proxy
|
|
5
|
+
Requires-Python: >=3.9
|
|
6
|
+
Requires-Dist: httpx
|
|
7
|
+
Provides-Extra: test
|
|
8
|
+
Requires-Dist: pytest; extra == "test"
|
|
9
|
+
Requires-Dist: pytest-asyncio; extra == "test"
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
# tripwire-litellm
|
|
2
|
+
|
|
3
|
+
Pre-call budget check and post-call ingest for a customer's LiteLLM proxy. Product behavior, fail-open, and overshoot are in `docs/hard-stop.md` at the repo root.
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
Python 3.9 or newer. The proxy already has LiteLLM.
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install -e integrations/litellm
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Config
|
|
14
|
+
|
|
15
|
+
`callbacks` must point at the instance:
|
|
16
|
+
|
|
17
|
+
```yaml
|
|
18
|
+
litellm_settings:
|
|
19
|
+
callbacks: ["tripwire_litellm.handler.proxy_handler_instance"]
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
```bash
|
|
23
|
+
TRIPWIRE_INGEST_TOKEN=tw_live_... # required, never log it
|
|
24
|
+
TRIPWIRE_BASE_URL=https://www.tripwire.agency
|
|
25
|
+
TRIPWIRE_FAIL_MODE=open # or closed
|
|
26
|
+
TRIPWIRE_CHECK_TIMEOUT_MS=800
|
|
27
|
+
TRIPWIRE_DENY_STATUS=402 # not 429; SDKs retry 429
|
|
28
|
+
TRIPWIRE_RESERVE_USD=0
|
|
29
|
+
TRIPWIRE_HARD_STOP=on # off = ingest only
|
|
30
|
+
TRIPWIRE_CLIENT_ID=
|
|
31
|
+
TRIPWIRE_PROJECT_ID=
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
## Verified hook
|
|
35
|
+
|
|
36
|
+
Checked against the LiteLLM proxy call-hook docs (fetched 2026-10-05). Minimum proxy version: **1.65.0** (the `async_pre_call_hook` arguments and `UserAPIKeyAuth.metadata` / `team_metadata` are stable there). These unit tests stub `litellm.integrations.custom_logger` when the package is not installed.
|
|
37
|
+
|
|
38
|
+
- `async_pre_call_hook(self, user_api_key_dict, cache, data, call_type)`
|
|
39
|
+
- Request attribution: `data["metadata"]["tripwire_client_id"]` and `tripwire_project_id`
|
|
40
|
+
- Virtual key: `user_api_key_dict.metadata`
|
|
41
|
+
- Team: `user_api_key_dict.team_metadata`
|
|
42
|
+
- Fallback: `TRIPWIRE_CLIENT_ID` and `TRIPWIRE_PROJECT_ID`
|
|
43
|
+
- `async_log_success_event(self, kwargs, response_obj, start_time, end_time)`
|
|
44
|
+
- Call id: `kwargs["litellm_call_id"]`, then `kwargs["litellm_params"]["litellm_call_id"]`, then `response_obj.id`
|
|
45
|
+
- Cost and provider: `kwargs["response_cost"]`, `kwargs["litellm_params"]["custom_llm_provider"]`, `kwargs["model"]`
|
|
46
|
+
|
|
47
|
+
A deny raises `fastapi.HTTPException` 402 with `error.code` `tripwire_budget_exceeded`. Allow returns `data` unchanged.
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "tripwire-litellm"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Tripwire hard-stop and ingest callback for a customer's LiteLLM proxy"
|
|
5
|
+
requires-python = ">=3.9"
|
|
6
|
+
dependencies = ["httpx"]
|
|
7
|
+
|
|
8
|
+
[project.optional-dependencies]
|
|
9
|
+
test = ["pytest", "pytest-asyncio"]
|
|
10
|
+
|
|
11
|
+
[build-system]
|
|
12
|
+
requires = ["setuptools>=68"]
|
|
13
|
+
build-backend = "setuptools.build_meta"
|
|
14
|
+
|
|
15
|
+
[tool.setuptools.packages.find]
|
|
16
|
+
include = ["tripwire_litellm*"]
|
|
17
|
+
|
|
18
|
+
[tool.pytest.ini_options]
|
|
19
|
+
asyncio_mode = "auto"
|
|
@@ -0,0 +1,307 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import json
|
|
3
|
+
import logging
|
|
4
|
+
import os
|
|
5
|
+
|
|
6
|
+
import httpx
|
|
7
|
+
import pytest
|
|
8
|
+
|
|
9
|
+
from tripwire_litellm.handler import TripwireHandler, proxy_handler_instance
|
|
10
|
+
|
|
11
|
+
TOKEN = "tw_live_SUPERSECRET"
|
|
12
|
+
CLIENT = "11111111-1111-4111-8111-111111111111"
|
|
13
|
+
PROJECT = "44444444-4444-4444-8444-444444444444"
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class Key:
|
|
17
|
+
def __init__(self, metadata=None, team_metadata=None):
|
|
18
|
+
self.metadata = metadata or {}
|
|
19
|
+
self.team_metadata = team_metadata or {}
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def allow_body(ttl=15, **extra):
|
|
23
|
+
body = {
|
|
24
|
+
"decision": "allow",
|
|
25
|
+
"enforced": True,
|
|
26
|
+
"reason": "within_budget",
|
|
27
|
+
"scope": "client",
|
|
28
|
+
"ttl_seconds": ttl,
|
|
29
|
+
"spend_usd": 1,
|
|
30
|
+
"budget_usd": 10,
|
|
31
|
+
}
|
|
32
|
+
body.update(extra)
|
|
33
|
+
return body
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def deny_body():
|
|
37
|
+
return {
|
|
38
|
+
"decision": "deny",
|
|
39
|
+
"enforced": True,
|
|
40
|
+
"reason": "budget_exceeded",
|
|
41
|
+
"scope": "project",
|
|
42
|
+
"ttl_seconds": 30,
|
|
43
|
+
"spend_usd": 10,
|
|
44
|
+
"budget_usd": 10,
|
|
45
|
+
}
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
@pytest.fixture
|
|
49
|
+
def env(monkeypatch):
|
|
50
|
+
monkeypatch.setenv("TRIPWIRE_INGEST_TOKEN", TOKEN)
|
|
51
|
+
monkeypatch.setenv("TRIPWIRE_BASE_URL", "https://www.tripwire.agency")
|
|
52
|
+
monkeypatch.setenv("TRIPWIRE_FAIL_MODE", "open")
|
|
53
|
+
monkeypatch.setenv("TRIPWIRE_HARD_STOP", "on")
|
|
54
|
+
monkeypatch.delenv("TRIPWIRE_CLIENT_ID", raising=False)
|
|
55
|
+
monkeypatch.delenv("TRIPWIRE_PROJECT_ID", raising=False)
|
|
56
|
+
monkeypatch.delenv("TRIPWIRE_DENY_STATUS", raising=False)
|
|
57
|
+
monkeypatch.delenv("TRIPWIRE_RESERVE_USD", raising=False)
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def handler(transport, clock=None):
|
|
61
|
+
box = {"t": 0.0}
|
|
62
|
+
|
|
63
|
+
def now():
|
|
64
|
+
return box["t"] if clock is None else clock["t"]
|
|
65
|
+
|
|
66
|
+
return TripwireHandler(transport=transport, now=now), box if clock is None else clock
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def data_for(client=CLIENT, project=""):
|
|
70
|
+
meta = {}
|
|
71
|
+
if client:
|
|
72
|
+
meta["tripwire_client_id"] = client
|
|
73
|
+
if project:
|
|
74
|
+
meta["tripwire_project_id"] = project
|
|
75
|
+
return {"model": "gpt-4o", "metadata": meta}
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
async def test_allow_returns_data_and_deny_raises_402(env):
|
|
79
|
+
calls = {"n": 0}
|
|
80
|
+
|
|
81
|
+
def transport(request: httpx.Request) -> httpx.Response:
|
|
82
|
+
calls["n"] += 1
|
|
83
|
+
body = json.loads(request.content.decode())
|
|
84
|
+
if body.get("estimated_cost_usd") == 0 and "client_id" in body:
|
|
85
|
+
return httpx.Response(200, json=allow_body())
|
|
86
|
+
return httpx.Response(200, json=deny_body())
|
|
87
|
+
|
|
88
|
+
hook, _clock = handler(httpx.MockTransport(transport))
|
|
89
|
+
payload = data_for()
|
|
90
|
+
assert await hook.async_pre_call_hook(Key(), None, payload, "completion") is payload
|
|
91
|
+
|
|
92
|
+
denied = data_for(project=PROJECT)
|
|
93
|
+
# Force the deny response by a separate handler.
|
|
94
|
+
def deny_transport(request: httpx.Request) -> httpx.Response:
|
|
95
|
+
assert request.headers["authorization"] == f"Bearer {TOKEN}"
|
|
96
|
+
assert TOKEN not in request.content.decode()
|
|
97
|
+
return httpx.Response(200, json=deny_body())
|
|
98
|
+
|
|
99
|
+
deny_hook, _clock = handler(httpx.MockTransport(deny_transport))
|
|
100
|
+
with pytest.raises(Exception) as caught:
|
|
101
|
+
await deny_hook.async_pre_call_hook(Key(), None, data_for(project=PROJECT), "completion")
|
|
102
|
+
assert caught.value.status_code == 402
|
|
103
|
+
assert caught.value.detail["error"]["code"] == "tripwire_budget_exceeded"
|
|
104
|
+
assert "monthly budget" in caught.value.detail["error"]["message"]
|
|
105
|
+
assert "$10.00 of $10.00" in caught.value.detail["error"]["message"]
|
|
106
|
+
assert denied is not payload
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
async def test_cache_hit_skips_the_second_request_and_ttl_0_rechecks(env):
|
|
110
|
+
calls = {"n": 0}
|
|
111
|
+
|
|
112
|
+
def transport(request: httpx.Request) -> httpx.Response:
|
|
113
|
+
calls["n"] += 1
|
|
114
|
+
ttl = 15 if calls["n"] == 1 else 0
|
|
115
|
+
return httpx.Response(200, json=allow_body(ttl=ttl if calls["n"] < 3 else 0))
|
|
116
|
+
|
|
117
|
+
hook, clock = handler(httpx.MockTransport(transport))
|
|
118
|
+
await hook.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
119
|
+
await hook.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
120
|
+
assert calls["n"] == 1
|
|
121
|
+
|
|
122
|
+
zero = {"n": 0}
|
|
123
|
+
|
|
124
|
+
def always(request: httpx.Request) -> httpx.Response:
|
|
125
|
+
zero["n"] += 1
|
|
126
|
+
return httpx.Response(200, json=allow_body(ttl=0))
|
|
127
|
+
|
|
128
|
+
fresh, _clock = handler(httpx.MockTransport(always))
|
|
129
|
+
await fresh.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
130
|
+
await fresh.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
131
|
+
assert zero["n"] == 2
|
|
132
|
+
assert clock["t"] == 0
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
async def test_concurrent_misses_share_one_request(env):
|
|
136
|
+
calls = {"n": 0}
|
|
137
|
+
started = asyncio.Event()
|
|
138
|
+
release = asyncio.Event()
|
|
139
|
+
|
|
140
|
+
class Slow(httpx.AsyncBaseTransport):
|
|
141
|
+
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
|
142
|
+
calls["n"] += 1
|
|
143
|
+
started.set()
|
|
144
|
+
await release.wait()
|
|
145
|
+
return httpx.Response(200, json=allow_body(ttl=15))
|
|
146
|
+
|
|
147
|
+
hook, _clock = handler(Slow())
|
|
148
|
+
|
|
149
|
+
async def one():
|
|
150
|
+
return await hook.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
151
|
+
|
|
152
|
+
pending = [asyncio.create_task(one()) for _ in range(50)]
|
|
153
|
+
await asyncio.wait_for(started.wait(), timeout=2)
|
|
154
|
+
await asyncio.sleep(0.05)
|
|
155
|
+
release.set()
|
|
156
|
+
results = await asyncio.gather(*pending)
|
|
157
|
+
assert calls["n"] == 1
|
|
158
|
+
assert all(item["metadata"]["tripwire_client_id"] == CLIENT for item in results)
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
async def test_timeout_allows_when_open_and_503_when_closed(env, monkeypatch):
|
|
162
|
+
def boom(request: httpx.Request) -> httpx.Response:
|
|
163
|
+
raise httpx.ReadTimeout("timed out", request=request)
|
|
164
|
+
|
|
165
|
+
open_hook, _clock = handler(httpx.MockTransport(boom))
|
|
166
|
+
payload = data_for()
|
|
167
|
+
assert await open_hook.async_pre_call_hook(Key(), None, payload, "completion") is payload
|
|
168
|
+
|
|
169
|
+
monkeypatch.setenv("TRIPWIRE_FAIL_MODE", "closed")
|
|
170
|
+
closed, _clock = handler(httpx.MockTransport(boom))
|
|
171
|
+
with pytest.raises(Exception) as caught:
|
|
172
|
+
await closed.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
173
|
+
assert caught.value.status_code == 503
|
|
174
|
+
assert caught.value.detail["error"]["code"] == "tripwire_unavailable"
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
async def test_circuit_breaker_opens_and_recovers(env):
|
|
178
|
+
calls = {"n": 0}
|
|
179
|
+
|
|
180
|
+
def transport(request: httpx.Request) -> httpx.Response:
|
|
181
|
+
calls["n"] += 1
|
|
182
|
+
if calls["n"] <= 5:
|
|
183
|
+
raise httpx.ConnectError("down")
|
|
184
|
+
return httpx.Response(200, json=allow_body(ttl=0))
|
|
185
|
+
|
|
186
|
+
clock = {"t": 100.0}
|
|
187
|
+
hook, _clock = handler(httpx.MockTransport(transport), clock)
|
|
188
|
+
for _ in range(5):
|
|
189
|
+
await hook.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
190
|
+
assert calls["n"] == 5
|
|
191
|
+
await hook.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
192
|
+
assert calls["n"] == 5
|
|
193
|
+
clock["t"] = 130.0
|
|
194
|
+
await hook.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
195
|
+
assert calls["n"] == 6
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
async def test_attribution_precedence_and_missing_warns_once(env, monkeypatch, caplog):
|
|
199
|
+
seen = []
|
|
200
|
+
|
|
201
|
+
def transport(request: httpx.Request) -> httpx.Response:
|
|
202
|
+
seen.append(json.loads(request.content.decode()))
|
|
203
|
+
return httpx.Response(200, json=allow_body(ttl=0))
|
|
204
|
+
|
|
205
|
+
hook, _clock = handler(httpx.MockTransport(transport))
|
|
206
|
+
monkeypatch.setenv("TRIPWIRE_CLIENT_ID", "env-client")
|
|
207
|
+
monkeypatch.setenv("TRIPWIRE_PROJECT_ID", "env-project")
|
|
208
|
+
key = Key(
|
|
209
|
+
metadata={"tripwire_client_id": "key-client", "tripwire_project_id": "key-project"},
|
|
210
|
+
team_metadata={"tripwire_client_id": "team-client", "tripwire_project_id": "team-project"},
|
|
211
|
+
)
|
|
212
|
+
await hook.async_pre_call_hook(
|
|
213
|
+
key,
|
|
214
|
+
None,
|
|
215
|
+
{"metadata": {"tripwire_client_id": "req-client", "tripwire_project_id": "req-project"}},
|
|
216
|
+
"completion",
|
|
217
|
+
)
|
|
218
|
+
await hook.async_pre_call_hook(key, None, {"metadata": {}}, "completion")
|
|
219
|
+
await hook.async_pre_call_hook(
|
|
220
|
+
Key(team_metadata={"tripwire_client_id": "team-client"}),
|
|
221
|
+
None,
|
|
222
|
+
{},
|
|
223
|
+
"completion",
|
|
224
|
+
)
|
|
225
|
+
await hook.async_pre_call_hook(Key(), None, {}, "completion")
|
|
226
|
+
assert seen[0]["client_id"] == "req-client"
|
|
227
|
+
assert seen[0]["project_id"] == "req-project"
|
|
228
|
+
assert seen[1]["client_id"] == "key-client"
|
|
229
|
+
assert seen[2]["client_id"] == "team-client"
|
|
230
|
+
assert seen[2]["project_id"] == "env-project"
|
|
231
|
+
|
|
232
|
+
monkeypatch.delenv("TRIPWIRE_CLIENT_ID", raising=False)
|
|
233
|
+
monkeypatch.delenv("TRIPWIRE_PROJECT_ID", raising=False)
|
|
234
|
+
caplog.set_level(logging.WARNING, logger="tripwire_litellm")
|
|
235
|
+
bare = TripwireHandler(transport=httpx.MockTransport(transport), now=lambda: 0)
|
|
236
|
+
await bare.async_pre_call_hook(Key(), None, {}, "completion")
|
|
237
|
+
await bare.async_pre_call_hook(Key(), None, {}, "completion")
|
|
238
|
+
warnings = [rec.message for rec in caplog.records if "no client or project" in rec.message]
|
|
239
|
+
assert len(warnings) == 1
|
|
240
|
+
|
|
241
|
+
|
|
242
|
+
async def test_success_log_invalidates_the_cache(env):
|
|
243
|
+
calls = {"check": 0, "ingest": 0}
|
|
244
|
+
bodies = []
|
|
245
|
+
|
|
246
|
+
def transport(request: httpx.Request) -> httpx.Response:
|
|
247
|
+
if request.url.path.endswith("/budget/check"):
|
|
248
|
+
calls["check"] += 1
|
|
249
|
+
return httpx.Response(200, json=allow_body(ttl=60))
|
|
250
|
+
calls["ingest"] += 1
|
|
251
|
+
bodies.append(json.loads(request.content.decode()))
|
|
252
|
+
assert TOKEN not in request.content.decode()
|
|
253
|
+
return httpx.Response(201, json={"event": {"id": "evt-1"}})
|
|
254
|
+
|
|
255
|
+
hook, _clock = handler(httpx.MockTransport(transport))
|
|
256
|
+
payload = data_for(project=PROJECT)
|
|
257
|
+
await hook.async_pre_call_hook(Key(), None, payload, "completion")
|
|
258
|
+
await hook.async_pre_call_hook(Key(), None, payload, "completion")
|
|
259
|
+
assert calls["check"] == 1
|
|
260
|
+
|
|
261
|
+
class Usage:
|
|
262
|
+
prompt_tokens = 12
|
|
263
|
+
completion_tokens = 4
|
|
264
|
+
|
|
265
|
+
class Response:
|
|
266
|
+
id = "chatcmpl-1"
|
|
267
|
+
usage = Usage()
|
|
268
|
+
|
|
269
|
+
await hook.async_log_success_event(
|
|
270
|
+
{
|
|
271
|
+
"model": "gpt-4o",
|
|
272
|
+
"response_cost": 0.02,
|
|
273
|
+
"litellm_call_id": "call-99",
|
|
274
|
+
"litellm_params": {
|
|
275
|
+
"custom_llm_provider": "openai",
|
|
276
|
+
"metadata": {"tripwire_client_id": CLIENT, "tripwire_project_id": PROJECT},
|
|
277
|
+
},
|
|
278
|
+
},
|
|
279
|
+
Response(),
|
|
280
|
+
None,
|
|
281
|
+
None,
|
|
282
|
+
)
|
|
283
|
+
assert calls["ingest"] == 1
|
|
284
|
+
assert bodies[0]["external_request_id"] == "call-99"
|
|
285
|
+
assert bodies[0]["estimated_cost_usd"] == 0.02
|
|
286
|
+
assert bodies[0]["client_id"] == CLIENT
|
|
287
|
+
assert bodies[0]["input_tokens"] == 12
|
|
288
|
+
await hook.async_pre_call_hook(Key(), None, payload, "completion")
|
|
289
|
+
assert calls["check"] == 2
|
|
290
|
+
|
|
291
|
+
|
|
292
|
+
async def test_token_never_appears_in_logs(env, caplog):
|
|
293
|
+
def boom(request: httpx.Request) -> httpx.Response:
|
|
294
|
+
raise httpx.ReadTimeout("timed out", request=request)
|
|
295
|
+
|
|
296
|
+
caplog.set_level(logging.DEBUG)
|
|
297
|
+
hook, _clock = handler(httpx.MockTransport(boom))
|
|
298
|
+
await hook.async_pre_call_hook(Key(), None, data_for(), "completion")
|
|
299
|
+
text = "\n".join(record.getMessage() for record in caplog.records)
|
|
300
|
+
assert TOKEN not in text
|
|
301
|
+
assert "SUPERSECRET" not in text
|
|
302
|
+
assert "Bearer" not in text
|
|
303
|
+
|
|
304
|
+
|
|
305
|
+
def test_proxy_handler_instance_is_an_instance():
|
|
306
|
+
assert isinstance(proxy_handler_instance, TripwireHandler)
|
|
307
|
+
assert not isinstance(proxy_handler_instance, type)
|
|
@@ -0,0 +1,406 @@
|
|
|
1
|
+
"""LiteLLM proxy callback: stop a call before it runs, then ingest it.
|
|
2
|
+
|
|
3
|
+
Configure an instance, not the class:
|
|
4
|
+
|
|
5
|
+
litellm_settings:
|
|
6
|
+
callbacks: ["tripwire_litellm.handler.proxy_handler_instance"]
|
|
7
|
+
|
|
8
|
+
Hook signature verified against the LiteLLM proxy call-hook docs
|
|
9
|
+
(async_pre_call_hook(user_api_key_dict, cache, data, call_type)).
|
|
10
|
+
Virtual-key metadata is UserAPIKeyAuth.metadata. Team metadata is
|
|
11
|
+
UserAPIKeyAuth.team_metadata. Request overrides live on
|
|
12
|
+
data["metadata"] as tripwire_client_id and tripwire_project_id.
|
|
13
|
+
Success logging is async_log_success_event(kwargs, response_obj,
|
|
14
|
+
start_time, end_time). The call id is kwargs["litellm_call_id"],
|
|
15
|
+
then litellm_params.litellm_call_id, then response_obj.id.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
import asyncio
|
|
21
|
+
import logging
|
|
22
|
+
import os
|
|
23
|
+
import time
|
|
24
|
+
from typing import Any, Callable, Dict, Optional, Tuple
|
|
25
|
+
|
|
26
|
+
import httpx
|
|
27
|
+
|
|
28
|
+
try:
|
|
29
|
+
from litellm.integrations.custom_logger import CustomLogger
|
|
30
|
+
except ImportError: # LiteLLM is installed on the proxy, not in unit tests.
|
|
31
|
+
|
|
32
|
+
class CustomLogger: # type: ignore[no-redef]
|
|
33
|
+
pass
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
logger = logging.getLogger("tripwire_litellm")
|
|
37
|
+
|
|
38
|
+
Scope = Tuple[str, str]
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
class _Result:
|
|
42
|
+
def __init__(self, kind: str, body: Optional[dict] = None) -> None:
|
|
43
|
+
self.kind = kind
|
|
44
|
+
self.body = body or {}
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
class TripwireHandler(CustomLogger):
|
|
48
|
+
def __init__(
|
|
49
|
+
self,
|
|
50
|
+
transport: Optional[httpx.AsyncBaseTransport] = None,
|
|
51
|
+
now: Optional[Callable[[], float]] = None,
|
|
52
|
+
) -> None:
|
|
53
|
+
self._transport = transport
|
|
54
|
+
self._now = now or time.monotonic
|
|
55
|
+
self._cache: Dict[Scope, Tuple[float, _Result]] = {}
|
|
56
|
+
self._inflight: Dict[Scope, "asyncio.Task[_Result]"] = {}
|
|
57
|
+
self._lock: Optional[asyncio.Lock] = None
|
|
58
|
+
self._failures = 0
|
|
59
|
+
self._open_until = 0.0
|
|
60
|
+
self._warned_missing = False
|
|
61
|
+
self._last_warn = -1e9
|
|
62
|
+
|
|
63
|
+
async def async_pre_call_hook(self, user_api_key_dict: Any, cache: Any, data: dict, call_type: str) -> dict:
|
|
64
|
+
del cache, call_type
|
|
65
|
+
if not isinstance(data, dict):
|
|
66
|
+
data = {}
|
|
67
|
+
if not _hard_stop_enabled():
|
|
68
|
+
return data
|
|
69
|
+
client_id, project_id = resolve_attribution(data, user_api_key_dict)
|
|
70
|
+
if not client_id and not project_id:
|
|
71
|
+
self._warn_missing()
|
|
72
|
+
return data
|
|
73
|
+
result = await self._cached_check(client_id, project_id, _model(data))
|
|
74
|
+
if result.kind == "deny":
|
|
75
|
+
_raise(_deny_status(), _deny_detail(result.body))
|
|
76
|
+
if result.kind == "error":
|
|
77
|
+
if _fail_closed():
|
|
78
|
+
self._warn_limited("Tripwire budget check failed; blocking the call")
|
|
79
|
+
_raise(503, _unavailable_detail())
|
|
80
|
+
self._warn_limited("Tripwire budget check failed; allowing the call")
|
|
81
|
+
return data
|
|
82
|
+
|
|
83
|
+
async def async_log_success_event(self, kwargs: dict, response_obj: Any, start_time: Any, end_time: Any) -> None:
|
|
84
|
+
del start_time, end_time
|
|
85
|
+
if not isinstance(kwargs, dict):
|
|
86
|
+
kwargs = {}
|
|
87
|
+
client_id, project_id = attribution_from_kwargs(kwargs)
|
|
88
|
+
if not client_id and not project_id:
|
|
89
|
+
self._warn_missing()
|
|
90
|
+
return
|
|
91
|
+
self._cache.pop((client_id, project_id), None)
|
|
92
|
+
if not _token():
|
|
93
|
+
self._warn_limited("Tripwire ingest token is not set; usage was not recorded")
|
|
94
|
+
return
|
|
95
|
+
payload = _ingest_payload(kwargs, response_obj, client_id, project_id)
|
|
96
|
+
try:
|
|
97
|
+
async with self._client(10) as client:
|
|
98
|
+
await client.post(_url("/api/v1/ingest/usage"), headers=_headers(), json=payload)
|
|
99
|
+
except Exception:
|
|
100
|
+
self._warn_limited("Tripwire usage ingest failed")
|
|
101
|
+
|
|
102
|
+
async def _cached_check(self, client_id: str, project_id: str, model: Optional[str]) -> _Result:
|
|
103
|
+
key: Scope = (client_id, project_id)
|
|
104
|
+
now = self._now()
|
|
105
|
+
async with self._guard():
|
|
106
|
+
if now < self._open_until:
|
|
107
|
+
return _Result("error")
|
|
108
|
+
cached = self._cache.get(key)
|
|
109
|
+
if cached and cached[0] > now:
|
|
110
|
+
return cached[1]
|
|
111
|
+
task = self._inflight.get(key)
|
|
112
|
+
if task is None:
|
|
113
|
+
task = asyncio.create_task(self._fetch(client_id, project_id, model, key))
|
|
114
|
+
self._inflight[key] = task
|
|
115
|
+
try:
|
|
116
|
+
return await task
|
|
117
|
+
finally:
|
|
118
|
+
async with self._guard():
|
|
119
|
+
if self._inflight.get(key) is task:
|
|
120
|
+
self._inflight.pop(key, None)
|
|
121
|
+
|
|
122
|
+
def _guard(self) -> asyncio.Lock:
|
|
123
|
+
if self._lock is None:
|
|
124
|
+
self._lock = asyncio.Lock()
|
|
125
|
+
return self._lock
|
|
126
|
+
|
|
127
|
+
async def _fetch(self, client_id: str, project_id: str, model: Optional[str], key: Scope) -> _Result:
|
|
128
|
+
result = await self._request(client_id, project_id, model)
|
|
129
|
+
if result.kind == "error":
|
|
130
|
+
self._note_failure()
|
|
131
|
+
return result
|
|
132
|
+
self._failures = 0
|
|
133
|
+
self._open_until = 0.0
|
|
134
|
+
ttl = _ttl(result)
|
|
135
|
+
if ttl > 0:
|
|
136
|
+
self._cache[key] = (self._now() + ttl, result)
|
|
137
|
+
return result
|
|
138
|
+
|
|
139
|
+
async def _request(self, client_id: str, project_id: str, model: Optional[str]) -> _Result:
|
|
140
|
+
if not _token():
|
|
141
|
+
return _Result("error")
|
|
142
|
+
body: Dict[str, Any] = {"estimated_cost_usd": _reserve()}
|
|
143
|
+
if client_id:
|
|
144
|
+
body["client_id"] = client_id
|
|
145
|
+
if project_id:
|
|
146
|
+
body["project_id"] = project_id
|
|
147
|
+
if model:
|
|
148
|
+
body["model"] = model
|
|
149
|
+
timeout = _timeout_seconds()
|
|
150
|
+
try:
|
|
151
|
+
async with self._client(timeout) as client:
|
|
152
|
+
response = await client.post(_url("/api/v1/budget/check"), headers=_headers(), json=body)
|
|
153
|
+
except Exception:
|
|
154
|
+
return _Result("error")
|
|
155
|
+
if response.status_code == 401 or response.status_code >= 500 or response.status_code != 200:
|
|
156
|
+
return _Result("error")
|
|
157
|
+
try:
|
|
158
|
+
payload = response.json()
|
|
159
|
+
except Exception:
|
|
160
|
+
return _Result("error")
|
|
161
|
+
if not isinstance(payload, dict):
|
|
162
|
+
return _Result("error")
|
|
163
|
+
decision = payload.get("decision")
|
|
164
|
+
if decision == "deny":
|
|
165
|
+
return _Result("deny", payload)
|
|
166
|
+
if decision == "allow":
|
|
167
|
+
return _Result("allow", payload)
|
|
168
|
+
return _Result("error")
|
|
169
|
+
|
|
170
|
+
def _note_failure(self) -> None:
|
|
171
|
+
self._failures += 1
|
|
172
|
+
if self._failures >= 5:
|
|
173
|
+
self._open_until = self._now() + 30
|
|
174
|
+
self._failures = 0
|
|
175
|
+
|
|
176
|
+
def _client(self, timeout: float) -> httpx.AsyncClient:
|
|
177
|
+
kwargs: Dict[str, Any] = {"timeout": timeout}
|
|
178
|
+
if self._transport is not None:
|
|
179
|
+
kwargs["transport"] = self._transport
|
|
180
|
+
return httpx.AsyncClient(**kwargs)
|
|
181
|
+
|
|
182
|
+
def _warn_missing(self) -> None:
|
|
183
|
+
if self._warned_missing:
|
|
184
|
+
return
|
|
185
|
+
self._warned_missing = True
|
|
186
|
+
logger.warning("Tripwire has no client or project attribution; allowing the call")
|
|
187
|
+
|
|
188
|
+
def _warn_limited(self, message: str) -> None:
|
|
189
|
+
now = self._now()
|
|
190
|
+
if now - self._last_warn < 30:
|
|
191
|
+
return
|
|
192
|
+
self._last_warn = now
|
|
193
|
+
logger.warning(message)
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
proxy_handler_instance = TripwireHandler()
|
|
197
|
+
|
|
198
|
+
|
|
199
|
+
def resolve_attribution(data: Any, user_api_key_dict: Any) -> Tuple[str, str]:
|
|
200
|
+
"""First source that sets an id wins for that id. Request metadata, then the virtual key, then the team, then env."""
|
|
201
|
+
sources = [
|
|
202
|
+
_meta(data, "metadata") if isinstance(data, dict) else {},
|
|
203
|
+
_obj_dict(user_api_key_dict, "metadata"),
|
|
204
|
+
_obj_dict(user_api_key_dict, "team_metadata"),
|
|
205
|
+
{
|
|
206
|
+
"tripwire_client_id": os.environ.get("TRIPWIRE_CLIENT_ID", ""),
|
|
207
|
+
"tripwire_project_id": os.environ.get("TRIPWIRE_PROJECT_ID", ""),
|
|
208
|
+
},
|
|
209
|
+
]
|
|
210
|
+
client_id = ""
|
|
211
|
+
project_id = ""
|
|
212
|
+
for source in sources:
|
|
213
|
+
if not client_id:
|
|
214
|
+
client_id = _clean(source.get("tripwire_client_id"))
|
|
215
|
+
if not project_id:
|
|
216
|
+
project_id = _clean(source.get("tripwire_project_id"))
|
|
217
|
+
return client_id, project_id
|
|
218
|
+
|
|
219
|
+
|
|
220
|
+
def attribution_from_kwargs(kwargs: dict) -> Tuple[str, str]:
|
|
221
|
+
params = kwargs.get("litellm_params")
|
|
222
|
+
if not isinstance(params, dict):
|
|
223
|
+
params = {}
|
|
224
|
+
metadata = params.get("metadata")
|
|
225
|
+
if not isinstance(metadata, dict):
|
|
226
|
+
raw = kwargs.get("metadata")
|
|
227
|
+
metadata = raw if isinstance(raw, dict) else {}
|
|
228
|
+
key = kwargs.get("user_api_key_dict")
|
|
229
|
+
if key is None:
|
|
230
|
+
key = params.get("user_api_key_dict")
|
|
231
|
+
if key is None:
|
|
232
|
+
key = _DictKey(
|
|
233
|
+
metadata.get("user_api_key_metadata") if isinstance(metadata.get("user_api_key_metadata"), dict) else {},
|
|
234
|
+
metadata.get("user_api_key_team_metadata")
|
|
235
|
+
if isinstance(metadata.get("user_api_key_team_metadata"), dict)
|
|
236
|
+
else {},
|
|
237
|
+
)
|
|
238
|
+
return resolve_attribution({"metadata": metadata}, key)
|
|
239
|
+
|
|
240
|
+
|
|
241
|
+
class _DictKey:
|
|
242
|
+
def __init__(self, metadata: dict, team_metadata: dict) -> None:
|
|
243
|
+
self.metadata = metadata
|
|
244
|
+
self.team_metadata = team_metadata
|
|
245
|
+
|
|
246
|
+
|
|
247
|
+
def _ingest_payload(kwargs: dict, response_obj: Any, client_id: str, project_id: str) -> dict:
|
|
248
|
+
params = kwargs.get("litellm_params") if isinstance(kwargs.get("litellm_params"), dict) else {}
|
|
249
|
+
usage = getattr(response_obj, "usage", None)
|
|
250
|
+
payload: Dict[str, Any] = {
|
|
251
|
+
"source_type": "litellm_proxy",
|
|
252
|
+
"provider_name": params.get("custom_llm_provider") or "unknown",
|
|
253
|
+
"model_name": kwargs.get("model") or "unknown",
|
|
254
|
+
"status": "success",
|
|
255
|
+
"input_tokens": int(getattr(usage, "prompt_tokens", 0) or 0),
|
|
256
|
+
"output_tokens": int(getattr(usage, "completion_tokens", 0) or 0),
|
|
257
|
+
"estimated_cost_usd": float(kwargs.get("response_cost") or 0),
|
|
258
|
+
}
|
|
259
|
+
if client_id:
|
|
260
|
+
payload["client_id"] = client_id
|
|
261
|
+
if project_id:
|
|
262
|
+
payload["project_id"] = project_id
|
|
263
|
+
call_id = _call_id(kwargs, params, response_obj)
|
|
264
|
+
if call_id:
|
|
265
|
+
payload["external_request_id"] = call_id
|
|
266
|
+
return payload
|
|
267
|
+
|
|
268
|
+
|
|
269
|
+
def _call_id(kwargs: dict, params: dict, response_obj: Any) -> str:
|
|
270
|
+
for value in (kwargs.get("litellm_call_id"), params.get("litellm_call_id"), getattr(response_obj, "id", None)):
|
|
271
|
+
if isinstance(value, str) and value.strip():
|
|
272
|
+
return value.strip()
|
|
273
|
+
return ""
|
|
274
|
+
|
|
275
|
+
|
|
276
|
+
def _deny_detail(body: dict) -> dict:
|
|
277
|
+
scope = body.get("scope")
|
|
278
|
+
label = "This project" if scope == "project" else "This client" if scope == "client" else "This budget"
|
|
279
|
+
spend = _money(body.get("spend_usd"))
|
|
280
|
+
budget = _money(body.get("budget_usd"))
|
|
281
|
+
message = f"{label} has used its monthly budget ({spend} of {budget}). Raise the budget in Tripwire to resume."
|
|
282
|
+
return {"error": {"type": "budget_exceeded", "code": "tripwire_budget_exceeded", "message": message}}
|
|
283
|
+
|
|
284
|
+
|
|
285
|
+
def _unavailable_detail() -> dict:
|
|
286
|
+
return {
|
|
287
|
+
"error": {
|
|
288
|
+
"type": "tripwire_unavailable",
|
|
289
|
+
"code": "tripwire_unavailable",
|
|
290
|
+
"message": "Tripwire budget check is unavailable.",
|
|
291
|
+
}
|
|
292
|
+
}
|
|
293
|
+
|
|
294
|
+
|
|
295
|
+
def _raise(status: int, detail: dict) -> None:
|
|
296
|
+
try:
|
|
297
|
+
from fastapi import HTTPException
|
|
298
|
+
except ImportError:
|
|
299
|
+
|
|
300
|
+
class HTTPException(Exception):
|
|
301
|
+
def __init__(self, status_code: int, detail: Any = None) -> None:
|
|
302
|
+
self.status_code = status_code
|
|
303
|
+
self.detail = detail
|
|
304
|
+
super().__init__(str(detail))
|
|
305
|
+
|
|
306
|
+
raise HTTPException(status_code=status, detail=detail)
|
|
307
|
+
|
|
308
|
+
|
|
309
|
+
def _headers() -> Dict[str, str]:
|
|
310
|
+
return {
|
|
311
|
+
"Authorization": f"Bearer {_token()}",
|
|
312
|
+
"Content-Type": "application/json",
|
|
313
|
+
}
|
|
314
|
+
|
|
315
|
+
|
|
316
|
+
def _token() -> str:
|
|
317
|
+
return os.environ.get("TRIPWIRE_INGEST_TOKEN", "").strip()
|
|
318
|
+
|
|
319
|
+
|
|
320
|
+
def _base_url() -> str:
|
|
321
|
+
raw = os.environ.get("TRIPWIRE_BASE_URL", "").strip() or "https://www.tripwire.agency"
|
|
322
|
+
return raw.rstrip("/")
|
|
323
|
+
|
|
324
|
+
|
|
325
|
+
def _url(path: str) -> str:
|
|
326
|
+
return f"{_base_url()}{path}"
|
|
327
|
+
|
|
328
|
+
|
|
329
|
+
def _hard_stop_enabled() -> bool:
|
|
330
|
+
value = os.environ.get("TRIPWIRE_HARD_STOP", "on").strip().lower()
|
|
331
|
+
return value not in {"off", "0", "false", "no"}
|
|
332
|
+
|
|
333
|
+
|
|
334
|
+
def _fail_closed() -> bool:
|
|
335
|
+
return os.environ.get("TRIPWIRE_FAIL_MODE", "open").strip().lower() == "closed"
|
|
336
|
+
|
|
337
|
+
|
|
338
|
+
def _timeout_seconds() -> float:
|
|
339
|
+
raw = os.environ.get("TRIPWIRE_CHECK_TIMEOUT_MS", "").strip()
|
|
340
|
+
try:
|
|
341
|
+
ms = int(raw) if raw else 800
|
|
342
|
+
except ValueError:
|
|
343
|
+
ms = 800
|
|
344
|
+
return max(ms, 1) / 1000
|
|
345
|
+
|
|
346
|
+
|
|
347
|
+
def _deny_status() -> int:
|
|
348
|
+
raw = os.environ.get("TRIPWIRE_DENY_STATUS", "").strip()
|
|
349
|
+
try:
|
|
350
|
+
status = int(raw) if raw else 402
|
|
351
|
+
except ValueError:
|
|
352
|
+
return 402
|
|
353
|
+
return status if 400 <= status <= 599 else 402
|
|
354
|
+
|
|
355
|
+
|
|
356
|
+
def _reserve() -> float:
|
|
357
|
+
raw = os.environ.get("TRIPWIRE_RESERVE_USD", "").strip()
|
|
358
|
+
try:
|
|
359
|
+
value = float(raw) if raw else 0
|
|
360
|
+
except ValueError:
|
|
361
|
+
return 0
|
|
362
|
+
return value if value > 0 else 0
|
|
363
|
+
|
|
364
|
+
|
|
365
|
+
def _ttl(result: _Result) -> float:
|
|
366
|
+
raw = result.body.get("ttl_seconds")
|
|
367
|
+
try:
|
|
368
|
+
ttl = float(raw)
|
|
369
|
+
except (TypeError, ValueError):
|
|
370
|
+
return 60 if result.kind == "allow" else 0
|
|
371
|
+
return ttl if ttl > 0 else 0
|
|
372
|
+
|
|
373
|
+
|
|
374
|
+
def _money(value: Any) -> str:
|
|
375
|
+
try:
|
|
376
|
+
return f"${float(value):.2f}"
|
|
377
|
+
except (TypeError, ValueError):
|
|
378
|
+
return "$0.00"
|
|
379
|
+
|
|
380
|
+
|
|
381
|
+
def _model(data: dict) -> Optional[str]:
|
|
382
|
+
model = data.get("model")
|
|
383
|
+
if isinstance(model, str) and model.strip():
|
|
384
|
+
return model.strip()
|
|
385
|
+
return None
|
|
386
|
+
|
|
387
|
+
|
|
388
|
+
def _meta(data: dict, key: str) -> dict:
|
|
389
|
+
raw = data.get(key)
|
|
390
|
+
return raw if isinstance(raw, dict) else {}
|
|
391
|
+
|
|
392
|
+
|
|
393
|
+
def _obj_dict(obj: Any, name: str) -> dict:
|
|
394
|
+
raw = getattr(obj, name, None)
|
|
395
|
+
if isinstance(raw, dict):
|
|
396
|
+
return raw
|
|
397
|
+
if isinstance(obj, dict):
|
|
398
|
+
value = obj.get(name)
|
|
399
|
+
return value if isinstance(value, dict) else {}
|
|
400
|
+
return {}
|
|
401
|
+
|
|
402
|
+
|
|
403
|
+
def _clean(value: Any) -> str:
|
|
404
|
+
if not isinstance(value, str):
|
|
405
|
+
return ""
|
|
406
|
+
return value.strip()
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tripwire-litellm
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Tripwire hard-stop and ingest callback for a customer's LiteLLM proxy
|
|
5
|
+
Requires-Python: >=3.9
|
|
6
|
+
Requires-Dist: httpx
|
|
7
|
+
Provides-Extra: test
|
|
8
|
+
Requires-Dist: pytest; extra == "test"
|
|
9
|
+
Requires-Dist: pytest-asyncio; extra == "test"
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
tests/test_handler.py
|
|
4
|
+
tripwire_litellm/__init__.py
|
|
5
|
+
tripwire_litellm/handler.py
|
|
6
|
+
tripwire_litellm.egg-info/PKG-INFO
|
|
7
|
+
tripwire_litellm.egg-info/SOURCES.txt
|
|
8
|
+
tripwire_litellm.egg-info/dependency_links.txt
|
|
9
|
+
tripwire_litellm.egg-info/requires.txt
|
|
10
|
+
tripwire_litellm.egg-info/top_level.txt
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
tripwire_litellm
|