litmus-python-sdk 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.
- litmus_python_sdk-0.1.0/PKG-INFO +46 -0
- litmus_python_sdk-0.1.0/README.md +34 -0
- litmus_python_sdk-0.1.0/pyproject.toml +33 -0
- litmus_python_sdk-0.1.0/src/litmus/__init__.py +11 -0
- litmus_python_sdk-0.1.0/src/litmus/client.py +489 -0
- litmus_python_sdk-0.1.0/src/litmus/consumer.py +158 -0
- litmus_python_sdk-0.1.0/src/litmus/request.py +126 -0
- litmus_python_sdk-0.1.0/src/litmus/version.py +1 -0
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: litmus-python-sdk
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Litmus Python SDK - implicit evals for AI products
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Requires-Dist: httpx>=0.27.0
|
|
7
|
+
Requires-Dist: pytest>=8.0 ; extra == 'dev'
|
|
8
|
+
Requires-Dist: ruff>=0.4 ; extra == 'dev'
|
|
9
|
+
Requires-Python: >=3.12
|
|
10
|
+
Provides-Extra: dev
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
|
|
13
|
+
# litmus-sdk
|
|
14
|
+
|
|
15
|
+
Python SDK for [Litmus](https://trylitmus.com) - implicit evals for AI products.
|
|
16
|
+
|
|
17
|
+
## Install
|
|
18
|
+
|
|
19
|
+
```bash
|
|
20
|
+
pip install litmus-sdk
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
## Quick start
|
|
24
|
+
|
|
25
|
+
```python
|
|
26
|
+
from litmus import LitmusClient
|
|
27
|
+
|
|
28
|
+
client = LitmusClient(api_key="ltm_pk_live_...")
|
|
29
|
+
|
|
30
|
+
# Track a generation and user signals
|
|
31
|
+
gen = client.generation("session-123", prompt_id="content_gen")
|
|
32
|
+
gen.accept()
|
|
33
|
+
gen.edit(edit_distance=0.3)
|
|
34
|
+
|
|
35
|
+
# Flush before exit (serverless, scripts, etc.)
|
|
36
|
+
client.shutdown()
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
## How it works
|
|
40
|
+
|
|
41
|
+
Events are queued in memory and shipped to the Litmus ingest API on a
|
|
42
|
+
background thread. Batches are sent every 0.5s or when 100 events
|
|
43
|
+
accumulate (both configurable). The consumer retries transient failures
|
|
44
|
+
with exponential backoff.
|
|
45
|
+
|
|
46
|
+
For serverless environments, pass `sync_mode=True` to send inline.
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
# litmus-sdk
|
|
2
|
+
|
|
3
|
+
Python SDK for [Litmus](https://trylitmus.com) - implicit evals for AI products.
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install litmus-sdk
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Quick start
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
from litmus import LitmusClient
|
|
15
|
+
|
|
16
|
+
client = LitmusClient(api_key="ltm_pk_live_...")
|
|
17
|
+
|
|
18
|
+
# Track a generation and user signals
|
|
19
|
+
gen = client.generation("session-123", prompt_id="content_gen")
|
|
20
|
+
gen.accept()
|
|
21
|
+
gen.edit(edit_distance=0.3)
|
|
22
|
+
|
|
23
|
+
# Flush before exit (serverless, scripts, etc.)
|
|
24
|
+
client.shutdown()
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
## How it works
|
|
28
|
+
|
|
29
|
+
Events are queued in memory and shipped to the Litmus ingest API on a
|
|
30
|
+
background thread. Batches are sent every 0.5s or when 100 events
|
|
31
|
+
accumulate (both configurable). The consumer retries transient failures
|
|
32
|
+
with exponential backoff.
|
|
33
|
+
|
|
34
|
+
For serverless environments, pass `sync_mode=True` to send inline.
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "litmus-python-sdk"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Litmus Python SDK - implicit evals for AI products"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.12"
|
|
7
|
+
license = "MIT"
|
|
8
|
+
dependencies = [
|
|
9
|
+
"httpx>=0.27.0",
|
|
10
|
+
]
|
|
11
|
+
|
|
12
|
+
[project.optional-dependencies]
|
|
13
|
+
dev = [
|
|
14
|
+
"pytest>=8.0",
|
|
15
|
+
"ruff>=0.4",
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
[build-system]
|
|
19
|
+
requires = ["uv_build>=0.11.3,<0.12"]
|
|
20
|
+
build-backend = "uv_build"
|
|
21
|
+
|
|
22
|
+
[tool.uv.build-backend]
|
|
23
|
+
module-name = "litmus"
|
|
24
|
+
|
|
25
|
+
[tool.ruff]
|
|
26
|
+
target-version = "py312"
|
|
27
|
+
line-length = 100
|
|
28
|
+
|
|
29
|
+
[tool.ruff.lint]
|
|
30
|
+
select = ["E", "F", "I", "N", "W", "UP"]
|
|
31
|
+
|
|
32
|
+
[tool.pytest.ini_options]
|
|
33
|
+
testpaths = ["tests"]
|
|
@@ -0,0 +1,489 @@
|
|
|
1
|
+
"""Litmus Python SDK client.
|
|
2
|
+
|
|
3
|
+
Drop-in event tracking with background batching, modeled after
|
|
4
|
+
PostHog's producer/consumer architecture. Events go into a
|
|
5
|
+
thread-safe queue, a daemon thread drains them in batches, and
|
|
6
|
+
batch_post() ships them to /v1/events.
|
|
7
|
+
|
|
8
|
+
from litmus import LitmusClient
|
|
9
|
+
|
|
10
|
+
client = LitmusClient(api_key="ltm_pk_live_...")
|
|
11
|
+
gen = client.generation("session-123", prompt_id="content_gen")
|
|
12
|
+
gen.accept()
|
|
13
|
+
gen.edit(edit_distance=0.3)
|
|
14
|
+
client.shutdown()
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
import atexit
|
|
20
|
+
import logging
|
|
21
|
+
import queue
|
|
22
|
+
from collections.abc import Callable
|
|
23
|
+
from datetime import UTC, datetime
|
|
24
|
+
from uuid import uuid4
|
|
25
|
+
|
|
26
|
+
from litmus.consumer import Consumer
|
|
27
|
+
from litmus.request import DEFAULT_HOST, batch_post
|
|
28
|
+
from litmus.version import VERSION
|
|
29
|
+
|
|
30
|
+
log = logging.getLogger("litmus")
|
|
31
|
+
|
|
32
|
+
# Same system events the TS SDK knows about
|
|
33
|
+
SYSTEM_EVENTS = frozenset(
|
|
34
|
+
[
|
|
35
|
+
"$generation",
|
|
36
|
+
"$regenerate",
|
|
37
|
+
"$copy",
|
|
38
|
+
"$edit",
|
|
39
|
+
"$abandon",
|
|
40
|
+
"$accept",
|
|
41
|
+
"$view",
|
|
42
|
+
"$partial_copy",
|
|
43
|
+
"$refine",
|
|
44
|
+
"$followup",
|
|
45
|
+
"$rephrase",
|
|
46
|
+
"$undo",
|
|
47
|
+
"$share",
|
|
48
|
+
"$flag",
|
|
49
|
+
"$rate",
|
|
50
|
+
"$escalate",
|
|
51
|
+
"$switch_model",
|
|
52
|
+
"$retry_context",
|
|
53
|
+
"$post_accept_edit",
|
|
54
|
+
"$blur",
|
|
55
|
+
"$return",
|
|
56
|
+
"$scroll_regression",
|
|
57
|
+
"$navigate",
|
|
58
|
+
"$interrupt",
|
|
59
|
+
]
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
class Generation:
|
|
64
|
+
"""Handle for a single AI generation. Lets you record behavioral
|
|
65
|
+
signals without re-threading IDs on every call.
|
|
66
|
+
|
|
67
|
+
gen = client.generation("session-123")
|
|
68
|
+
gen.accept()
|
|
69
|
+
gen.edit(edit_distance=0.3)
|
|
70
|
+
"""
|
|
71
|
+
|
|
72
|
+
__slots__ = ("id", "_session_id", "_defaults", "_client")
|
|
73
|
+
|
|
74
|
+
def __init__(
|
|
75
|
+
self,
|
|
76
|
+
client: LitmusClient,
|
|
77
|
+
session_id: str,
|
|
78
|
+
generation_id: str,
|
|
79
|
+
defaults: dict,
|
|
80
|
+
):
|
|
81
|
+
self._client = client
|
|
82
|
+
self._session_id = session_id
|
|
83
|
+
self.id = generation_id
|
|
84
|
+
self._defaults = defaults
|
|
85
|
+
|
|
86
|
+
def _emit(self, event_type: str, metadata: dict | None = None) -> None:
|
|
87
|
+
merged = {**self._defaults.get("metadata", {}), **(metadata or {})}
|
|
88
|
+
self._client.track(
|
|
89
|
+
event_type=event_type,
|
|
90
|
+
session_id=self._session_id,
|
|
91
|
+
user_id=self._defaults.get("user_id"),
|
|
92
|
+
prompt_id=self._defaults.get("prompt_id"),
|
|
93
|
+
prompt_version=self._defaults.get("prompt_version"),
|
|
94
|
+
generation_id=self.id,
|
|
95
|
+
metadata=merged if merged else None,
|
|
96
|
+
)
|
|
97
|
+
|
|
98
|
+
def accept(self, metadata: dict | None = None) -> None:
|
|
99
|
+
self._emit("$accept", metadata)
|
|
100
|
+
|
|
101
|
+
def edit(
|
|
102
|
+
self,
|
|
103
|
+
edit_distance: float | None = None,
|
|
104
|
+
metadata: dict | None = None,
|
|
105
|
+
) -> None:
|
|
106
|
+
m = {**(metadata or {})}
|
|
107
|
+
if edit_distance is not None:
|
|
108
|
+
m["edit_distance"] = edit_distance
|
|
109
|
+
self._emit("$edit", m)
|
|
110
|
+
|
|
111
|
+
def regenerate(self, metadata: dict | None = None) -> None:
|
|
112
|
+
self._emit("$regenerate", metadata)
|
|
113
|
+
|
|
114
|
+
def copy(self, metadata: dict | None = None) -> None:
|
|
115
|
+
self._emit("$copy", metadata)
|
|
116
|
+
|
|
117
|
+
def abandon(self, metadata: dict | None = None) -> None:
|
|
118
|
+
self._emit("$abandon", metadata)
|
|
119
|
+
|
|
120
|
+
def view(self, metadata: dict | None = None) -> None:
|
|
121
|
+
self._emit("$view", metadata)
|
|
122
|
+
|
|
123
|
+
def refine(
|
|
124
|
+
self,
|
|
125
|
+
refinement_type: str | None = None,
|
|
126
|
+
metadata: dict | None = None,
|
|
127
|
+
) -> None:
|
|
128
|
+
m = {**(metadata or {})}
|
|
129
|
+
if refinement_type is not None:
|
|
130
|
+
m["refinement_type"] = refinement_type
|
|
131
|
+
self._emit("$refine", m)
|
|
132
|
+
|
|
133
|
+
def followup(self, metadata: dict | None = None) -> None:
|
|
134
|
+
self._emit("$followup", metadata)
|
|
135
|
+
|
|
136
|
+
def rephrase(self, metadata: dict | None = None) -> None:
|
|
137
|
+
self._emit("$rephrase", metadata)
|
|
138
|
+
|
|
139
|
+
def undo(self, metadata: dict | None = None) -> None:
|
|
140
|
+
self._emit("$undo", metadata)
|
|
141
|
+
|
|
142
|
+
def share(
|
|
143
|
+
self,
|
|
144
|
+
channel: str | None = None,
|
|
145
|
+
edited_before_share: bool | None = None,
|
|
146
|
+
metadata: dict | None = None,
|
|
147
|
+
) -> None:
|
|
148
|
+
m = {**(metadata or {})}
|
|
149
|
+
if channel is not None:
|
|
150
|
+
m["channel"] = channel
|
|
151
|
+
if edited_before_share is not None:
|
|
152
|
+
m["edited_before_share"] = edited_before_share
|
|
153
|
+
self._emit("$share", m)
|
|
154
|
+
|
|
155
|
+
def flag(
|
|
156
|
+
self,
|
|
157
|
+
reason: str | None = None,
|
|
158
|
+
metadata: dict | None = None,
|
|
159
|
+
) -> None:
|
|
160
|
+
m = {**(metadata or {})}
|
|
161
|
+
if reason is not None:
|
|
162
|
+
m["reason"] = reason
|
|
163
|
+
self._emit("$flag", m)
|
|
164
|
+
|
|
165
|
+
def rate(
|
|
166
|
+
self,
|
|
167
|
+
value: float,
|
|
168
|
+
scale: str = "binary",
|
|
169
|
+
metadata: dict | None = None,
|
|
170
|
+
) -> None:
|
|
171
|
+
m = {"value": value, "scale": scale, **(metadata or {})}
|
|
172
|
+
self._emit("$rate", m)
|
|
173
|
+
|
|
174
|
+
def escalate(self, metadata: dict | None = None) -> None:
|
|
175
|
+
self._emit("$escalate", metadata)
|
|
176
|
+
|
|
177
|
+
def post_accept_edit(
|
|
178
|
+
self,
|
|
179
|
+
edit_distance: float | None = None,
|
|
180
|
+
time_since_accept_ms: int | None = None,
|
|
181
|
+
metadata: dict | None = None,
|
|
182
|
+
) -> None:
|
|
183
|
+
m = {**(metadata or {})}
|
|
184
|
+
if edit_distance is not None:
|
|
185
|
+
m["edit_distance"] = edit_distance
|
|
186
|
+
if time_since_accept_ms is not None:
|
|
187
|
+
m["time_since_accept_ms"] = time_since_accept_ms
|
|
188
|
+
self._emit("$post_accept_edit", m)
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
class Feature:
|
|
192
|
+
"""Scoped handle for an AI feature. Carries defaults so you don't
|
|
193
|
+
repeat prompt_id/model/user_id on every generation.
|
|
194
|
+
|
|
195
|
+
summarizer = client.feature("summarizer", model="gpt-4o")
|
|
196
|
+
gen = summarizer.generation("session-123")
|
|
197
|
+
gen.accept()
|
|
198
|
+
"""
|
|
199
|
+
|
|
200
|
+
__slots__ = ("_client", "_defaults", "name")
|
|
201
|
+
|
|
202
|
+
def __init__(self, client: LitmusClient, name: str, defaults: dict):
|
|
203
|
+
self._client = client
|
|
204
|
+
self.name = name
|
|
205
|
+
self._defaults = {**defaults, "prompt_id": defaults.get("prompt_id", name)}
|
|
206
|
+
|
|
207
|
+
def generation(
|
|
208
|
+
self,
|
|
209
|
+
session_id: str,
|
|
210
|
+
user_id: str | None = None,
|
|
211
|
+
prompt_version: str | None = None,
|
|
212
|
+
metadata: dict | None = None,
|
|
213
|
+
) -> Generation:
|
|
214
|
+
base_meta: dict = {"feature": self.name}
|
|
215
|
+
model = self._defaults.get("model")
|
|
216
|
+
if model:
|
|
217
|
+
base_meta["model"] = model
|
|
218
|
+
|
|
219
|
+
merged = {
|
|
220
|
+
**self._defaults,
|
|
221
|
+
"user_id": user_id or self._defaults.get("user_id"),
|
|
222
|
+
"prompt_version": prompt_version or self._defaults.get("prompt_version"),
|
|
223
|
+
"metadata": {
|
|
224
|
+
**base_meta,
|
|
225
|
+
**self._defaults.get("metadata", {}),
|
|
226
|
+
**(metadata or {}),
|
|
227
|
+
},
|
|
228
|
+
}
|
|
229
|
+
return self._client.generation(session_id, **merged)
|
|
230
|
+
|
|
231
|
+
def track(
|
|
232
|
+
self,
|
|
233
|
+
event_type: str,
|
|
234
|
+
session_id: str,
|
|
235
|
+
user_id: str | None = None,
|
|
236
|
+
metadata: dict | None = None,
|
|
237
|
+
**kwargs: object,
|
|
238
|
+
) -> str | None:
|
|
239
|
+
return self._client.track(
|
|
240
|
+
event_type=event_type,
|
|
241
|
+
session_id=session_id,
|
|
242
|
+
user_id=user_id or self._defaults.get("user_id"),
|
|
243
|
+
prompt_id=kwargs.get("prompt_id") or self._defaults.get("prompt_id"),
|
|
244
|
+
metadata={
|
|
245
|
+
**self._defaults.get("metadata", {}),
|
|
246
|
+
"feature": self.name,
|
|
247
|
+
**(metadata or {}),
|
|
248
|
+
},
|
|
249
|
+
**{k: v for k, v in kwargs.items() if k != "prompt_id"},
|
|
250
|
+
)
|
|
251
|
+
|
|
252
|
+
|
|
253
|
+
class LitmusClient:
|
|
254
|
+
"""Litmus event client with background batching.
|
|
255
|
+
|
|
256
|
+
Args:
|
|
257
|
+
api_key: Your Litmus API key (ltm_pk_live_... or ltm_pk_test_...).
|
|
258
|
+
host: Ingest endpoint URL. Defaults to https://ingest.trylitmus.com.
|
|
259
|
+
max_queue_size: Max events buffered in memory before dropping. Default: 10000.
|
|
260
|
+
on_error: Callback(exception, batch) invoked on send failure.
|
|
261
|
+
flush_at: Batch size threshold that triggers an upload. Default: 100.
|
|
262
|
+
flush_interval: Seconds to wait before flushing a partial batch. Default: 0.5.
|
|
263
|
+
gzip: Compress payloads with gzip. Default: False.
|
|
264
|
+
max_retries: Max retry attempts per batch. Default: 3.
|
|
265
|
+
sync_mode: Send events inline (no background thread). Useful for
|
|
266
|
+
serverless or testing. Default: False.
|
|
267
|
+
timeout: HTTP timeout in seconds. Default: 15.
|
|
268
|
+
threads: Number of consumer threads. Default: 1.
|
|
269
|
+
send: Actually send events (set False for dry-run). Default: True.
|
|
270
|
+
debug: Enable DEBUG-level logging. Default: False.
|
|
271
|
+
disabled: Silently drop all events. Default: False.
|
|
272
|
+
"""
|
|
273
|
+
|
|
274
|
+
log = logging.getLogger("litmus")
|
|
275
|
+
|
|
276
|
+
def __init__(
|
|
277
|
+
self,
|
|
278
|
+
api_key: str,
|
|
279
|
+
host: str | None = None,
|
|
280
|
+
max_queue_size: int = 10_000,
|
|
281
|
+
on_error: Callable[[Exception, list[dict]], None] | None = None,
|
|
282
|
+
flush_at: int = 100,
|
|
283
|
+
flush_interval: float = 0.5,
|
|
284
|
+
gzip: bool = False,
|
|
285
|
+
max_retries: int = 3,
|
|
286
|
+
sync_mode: bool = False,
|
|
287
|
+
timeout: int = 15,
|
|
288
|
+
threads: int = 1,
|
|
289
|
+
send: bool = True,
|
|
290
|
+
debug: bool = False,
|
|
291
|
+
disabled: bool = False,
|
|
292
|
+
):
|
|
293
|
+
self._queue: queue.Queue[dict] = queue.Queue(max_queue_size)
|
|
294
|
+
self.api_key = api_key
|
|
295
|
+
self.host = (host or DEFAULT_HOST).rstrip("/")
|
|
296
|
+
self.on_error = on_error
|
|
297
|
+
self.send = send
|
|
298
|
+
self.sync_mode = sync_mode
|
|
299
|
+
self.gzip = gzip
|
|
300
|
+
self.timeout = timeout
|
|
301
|
+
self.disabled = disabled
|
|
302
|
+
self.consumers: list[Consumer] | None = None
|
|
303
|
+
|
|
304
|
+
if debug:
|
|
305
|
+
logging.basicConfig()
|
|
306
|
+
self.log.setLevel(logging.DEBUG)
|
|
307
|
+
|
|
308
|
+
if sync_mode:
|
|
309
|
+
self.consumers = None
|
|
310
|
+
else:
|
|
311
|
+
if send:
|
|
312
|
+
atexit.register(self.join)
|
|
313
|
+
|
|
314
|
+
self.consumers = []
|
|
315
|
+
for _ in range(threads):
|
|
316
|
+
consumer = Consumer(
|
|
317
|
+
queue=self._queue,
|
|
318
|
+
api_key=self.api_key,
|
|
319
|
+
host=self.host,
|
|
320
|
+
on_error=on_error,
|
|
321
|
+
flush_at=flush_at,
|
|
322
|
+
flush_interval=flush_interval,
|
|
323
|
+
use_gzip=gzip,
|
|
324
|
+
retries=max_retries,
|
|
325
|
+
timeout=timeout,
|
|
326
|
+
)
|
|
327
|
+
self.consumers.append(consumer)
|
|
328
|
+
if send:
|
|
329
|
+
consumer.start()
|
|
330
|
+
|
|
331
|
+
# -- public API ----------------------------------------------------------
|
|
332
|
+
|
|
333
|
+
def track(
|
|
334
|
+
self,
|
|
335
|
+
event_type: str,
|
|
336
|
+
session_id: str,
|
|
337
|
+
user_id: str | None = None,
|
|
338
|
+
prompt_id: str | None = None,
|
|
339
|
+
prompt_version: str | None = None,
|
|
340
|
+
generation_id: str | None = None,
|
|
341
|
+
metadata: dict | None = None,
|
|
342
|
+
timestamp: datetime | None = None,
|
|
343
|
+
) -> str | None:
|
|
344
|
+
"""Enqueue a single event. Returns the event UUID or None if dropped."""
|
|
345
|
+
if self.disabled:
|
|
346
|
+
return None
|
|
347
|
+
|
|
348
|
+
ts = timestamp or datetime.now(tz=UTC)
|
|
349
|
+
event_id = str(uuid4())
|
|
350
|
+
|
|
351
|
+
msg: dict = {
|
|
352
|
+
"id": event_id,
|
|
353
|
+
"type": event_type,
|
|
354
|
+
"session_id": session_id,
|
|
355
|
+
"timestamp": ts.isoformat(),
|
|
356
|
+
}
|
|
357
|
+
if user_id:
|
|
358
|
+
msg["user_id"] = user_id
|
|
359
|
+
if prompt_id:
|
|
360
|
+
msg["prompt_id"] = prompt_id
|
|
361
|
+
if prompt_version:
|
|
362
|
+
msg["prompt_version"] = prompt_version
|
|
363
|
+
if generation_id:
|
|
364
|
+
msg["generation_id"] = generation_id
|
|
365
|
+
|
|
366
|
+
props = {"$lib": "litmus-python", "$lib_version": VERSION}
|
|
367
|
+
if metadata:
|
|
368
|
+
props.update(metadata)
|
|
369
|
+
msg["metadata"] = props
|
|
370
|
+
|
|
371
|
+
self.log.debug("queueing: %s", msg)
|
|
372
|
+
|
|
373
|
+
if not self.send:
|
|
374
|
+
return event_id
|
|
375
|
+
|
|
376
|
+
if self.sync_mode:
|
|
377
|
+
batch_post(
|
|
378
|
+
self.api_key,
|
|
379
|
+
host=self.host,
|
|
380
|
+
use_gzip=self.gzip,
|
|
381
|
+
timeout=self.timeout,
|
|
382
|
+
batch=[msg],
|
|
383
|
+
)
|
|
384
|
+
return event_id
|
|
385
|
+
|
|
386
|
+
try:
|
|
387
|
+
self._queue.put(msg, block=False)
|
|
388
|
+
return event_id
|
|
389
|
+
except queue.Full:
|
|
390
|
+
self.log.warning("litmus queue is full, event dropped")
|
|
391
|
+
return None
|
|
392
|
+
|
|
393
|
+
def generation(
|
|
394
|
+
self,
|
|
395
|
+
session_id: str,
|
|
396
|
+
user_id: str | None = None,
|
|
397
|
+
prompt_id: str | None = None,
|
|
398
|
+
prompt_version: str | None = None,
|
|
399
|
+
model: str | None = None,
|
|
400
|
+
metadata: dict | None = None,
|
|
401
|
+
) -> Generation:
|
|
402
|
+
"""Create a generation and return a handle for recording signals."""
|
|
403
|
+
generation_id = str(uuid4())
|
|
404
|
+
defaults = {
|
|
405
|
+
"user_id": user_id,
|
|
406
|
+
"prompt_id": prompt_id,
|
|
407
|
+
"prompt_version": prompt_version,
|
|
408
|
+
"model": model,
|
|
409
|
+
"metadata": metadata or {},
|
|
410
|
+
}
|
|
411
|
+
|
|
412
|
+
self.track(
|
|
413
|
+
event_type="$generation",
|
|
414
|
+
session_id=session_id,
|
|
415
|
+
user_id=user_id,
|
|
416
|
+
generation_id=generation_id,
|
|
417
|
+
prompt_id=prompt_id,
|
|
418
|
+
prompt_version=prompt_version,
|
|
419
|
+
metadata=metadata,
|
|
420
|
+
)
|
|
421
|
+
|
|
422
|
+
return Generation(self, session_id, generation_id, defaults)
|
|
423
|
+
|
|
424
|
+
def attach(
|
|
425
|
+
self,
|
|
426
|
+
generation_id: str,
|
|
427
|
+
session_id: str,
|
|
428
|
+
user_id: str | None = None,
|
|
429
|
+
prompt_id: str | None = None,
|
|
430
|
+
prompt_version: str | None = None,
|
|
431
|
+
metadata: dict | None = None,
|
|
432
|
+
) -> Generation:
|
|
433
|
+
"""Attach to an existing generation (e.g. one created by a frontend SDK).
|
|
434
|
+
|
|
435
|
+
Returns a Generation handle for recording signals without
|
|
436
|
+
re-emitting the $generation event. Use this when the generation
|
|
437
|
+
was already created by another SDK and you have the generation_id.
|
|
438
|
+
|
|
439
|
+
# Frontend created the generation, backend received the ID
|
|
440
|
+
gen = client.attach(request.generation_id, session_id)
|
|
441
|
+
gen.accept() # backend-side signal
|
|
442
|
+
"""
|
|
443
|
+
defaults = {
|
|
444
|
+
"user_id": user_id,
|
|
445
|
+
"prompt_id": prompt_id,
|
|
446
|
+
"prompt_version": prompt_version,
|
|
447
|
+
"metadata": metadata or {},
|
|
448
|
+
}
|
|
449
|
+
return Generation(self, session_id, generation_id, defaults)
|
|
450
|
+
|
|
451
|
+
def feature(
|
|
452
|
+
self,
|
|
453
|
+
name: str,
|
|
454
|
+
model: str | None = None,
|
|
455
|
+
user_id: str | None = None,
|
|
456
|
+
prompt_version: str | None = None,
|
|
457
|
+
metadata: dict | None = None,
|
|
458
|
+
) -> Feature:
|
|
459
|
+
"""Create a scoped feature handle that carries defaults."""
|
|
460
|
+
defaults = {
|
|
461
|
+
"prompt_id": name,
|
|
462
|
+
"model": model,
|
|
463
|
+
"user_id": user_id,
|
|
464
|
+
"prompt_version": prompt_version,
|
|
465
|
+
"metadata": metadata or {},
|
|
466
|
+
}
|
|
467
|
+
return Feature(self, name, defaults)
|
|
468
|
+
|
|
469
|
+
def flush(self) -> None:
|
|
470
|
+
"""Block until the queue is fully drained."""
|
|
471
|
+
size = self._queue.qsize()
|
|
472
|
+
self._queue.join()
|
|
473
|
+
self.log.debug("flushed ~%d items", size)
|
|
474
|
+
|
|
475
|
+
def join(self) -> None:
|
|
476
|
+
"""Stop consumer threads (call after flush)."""
|
|
477
|
+
if self.consumers:
|
|
478
|
+
for consumer in self.consumers:
|
|
479
|
+
consumer.pause()
|
|
480
|
+
try:
|
|
481
|
+
consumer.join()
|
|
482
|
+
except RuntimeError:
|
|
483
|
+
pass
|
|
484
|
+
|
|
485
|
+
def shutdown(self) -> None:
|
|
486
|
+
"""Flush all pending events and stop consumer threads.
|
|
487
|
+
Call this before process exit in serverless environments."""
|
|
488
|
+
self.flush()
|
|
489
|
+
self.join()
|
|
@@ -0,0 +1,158 @@
|
|
|
1
|
+
"""Background consumer thread that drains the event queue in batches.
|
|
2
|
+
|
|
3
|
+
A daemon thread that wakes on a flush interval, collects up to
|
|
4
|
+
flush_at items from the queue, and POSTs them as a single batch.
|
|
5
|
+
Retries with exponential backoff on transient failures.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import json
|
|
11
|
+
import logging
|
|
12
|
+
import time
|
|
13
|
+
from collections.abc import Callable
|
|
14
|
+
from queue import Empty
|
|
15
|
+
from threading import Thread
|
|
16
|
+
from typing import TYPE_CHECKING
|
|
17
|
+
|
|
18
|
+
from litmus.request import APIError, batch_post
|
|
19
|
+
|
|
20
|
+
if TYPE_CHECKING:
|
|
21
|
+
from queue import Queue
|
|
22
|
+
|
|
23
|
+
log = logging.getLogger("litmus")
|
|
24
|
+
|
|
25
|
+
# Hard ceiling per message to avoid blowing up the batch
|
|
26
|
+
MAX_MSG_SIZE = 900 * 1024 # 900 KiB
|
|
27
|
+
BATCH_SIZE_LIMIT = 5 * 1024 * 1024 # 5 MiB total batch
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class Consumer(Thread):
|
|
31
|
+
"""Drains the client's queue and ships batches to the ingest API."""
|
|
32
|
+
|
|
33
|
+
def __init__(
|
|
34
|
+
self,
|
|
35
|
+
queue: Queue[dict],
|
|
36
|
+
api_key: str,
|
|
37
|
+
host: str | None = None,
|
|
38
|
+
on_error: Callable[[Exception, list[dict]], None] | None = None,
|
|
39
|
+
flush_at: int = 100,
|
|
40
|
+
flush_interval: float = 0.5,
|
|
41
|
+
use_gzip: bool = False,
|
|
42
|
+
retries: int = 10,
|
|
43
|
+
timeout: int = 15,
|
|
44
|
+
):
|
|
45
|
+
super().__init__()
|
|
46
|
+
self.daemon = True
|
|
47
|
+
self.queue = queue
|
|
48
|
+
self.api_key = api_key
|
|
49
|
+
self.host = host
|
|
50
|
+
self.on_error = on_error
|
|
51
|
+
self.flush_at = flush_at
|
|
52
|
+
self.flush_interval = flush_interval
|
|
53
|
+
self.use_gzip = use_gzip
|
|
54
|
+
self.retries = retries
|
|
55
|
+
self.timeout = timeout
|
|
56
|
+
self.running = True
|
|
57
|
+
|
|
58
|
+
def run(self) -> None:
|
|
59
|
+
log.debug("consumer thread started")
|
|
60
|
+
while self.running:
|
|
61
|
+
self.upload()
|
|
62
|
+
log.debug("consumer thread exited")
|
|
63
|
+
|
|
64
|
+
def pause(self) -> None:
|
|
65
|
+
self.running = False
|
|
66
|
+
|
|
67
|
+
def upload(self) -> bool:
|
|
68
|
+
"""Pull the next batch off the queue and send it. Returns success."""
|
|
69
|
+
batch = self._next_batch()
|
|
70
|
+
if not batch:
|
|
71
|
+
return False
|
|
72
|
+
|
|
73
|
+
try:
|
|
74
|
+
self._send_with_retries(batch)
|
|
75
|
+
return True
|
|
76
|
+
except Exception as exc:
|
|
77
|
+
log.error("error uploading: %s", exc)
|
|
78
|
+
if self.on_error:
|
|
79
|
+
try:
|
|
80
|
+
self.on_error(exc, batch)
|
|
81
|
+
except Exception as handler_err:
|
|
82
|
+
log.error("on_error callback failed: %s", handler_err)
|
|
83
|
+
return False
|
|
84
|
+
finally:
|
|
85
|
+
for _ in batch:
|
|
86
|
+
self.queue.task_done()
|
|
87
|
+
|
|
88
|
+
# -- internal ------------------------------------------------------------
|
|
89
|
+
|
|
90
|
+
def _next_batch(self) -> list[dict]:
|
|
91
|
+
"""Collect items from the queue until we hit flush_at count,
|
|
92
|
+
flush_interval timeout, or the batch size limit."""
|
|
93
|
+
items: list[dict] = []
|
|
94
|
+
total_size = 0
|
|
95
|
+
start = time.monotonic()
|
|
96
|
+
|
|
97
|
+
while len(items) < self.flush_at:
|
|
98
|
+
elapsed = time.monotonic() - start
|
|
99
|
+
if elapsed >= self.flush_interval:
|
|
100
|
+
break
|
|
101
|
+
try:
|
|
102
|
+
item = self.queue.get(
|
|
103
|
+
block=True,
|
|
104
|
+
timeout=self.flush_interval - elapsed,
|
|
105
|
+
)
|
|
106
|
+
item_size = len(json.dumps(item).encode())
|
|
107
|
+
if item_size > MAX_MSG_SIZE:
|
|
108
|
+
log.error("event exceeds 900 KiB limit, dropping")
|
|
109
|
+
self.queue.task_done()
|
|
110
|
+
continue
|
|
111
|
+
items.append(item)
|
|
112
|
+
total_size += item_size
|
|
113
|
+
if total_size >= BATCH_SIZE_LIMIT:
|
|
114
|
+
log.debug("hit batch size limit (%d bytes)", total_size)
|
|
115
|
+
break
|
|
116
|
+
except Empty:
|
|
117
|
+
break
|
|
118
|
+
|
|
119
|
+
return items
|
|
120
|
+
|
|
121
|
+
def _send_with_retries(self, batch: list[dict]) -> None:
|
|
122
|
+
"""Attempt to POST the batch, retrying transient errors."""
|
|
123
|
+
last_exc: Exception | None = None
|
|
124
|
+
|
|
125
|
+
for attempt in range(self.retries + 1):
|
|
126
|
+
try:
|
|
127
|
+
batch_post(
|
|
128
|
+
self.api_key,
|
|
129
|
+
host=self.host,
|
|
130
|
+
use_gzip=self.use_gzip,
|
|
131
|
+
timeout=self.timeout,
|
|
132
|
+
batch=batch,
|
|
133
|
+
)
|
|
134
|
+
return
|
|
135
|
+
except Exception as exc:
|
|
136
|
+
last_exc = exc
|
|
137
|
+
if not self._is_retryable(exc):
|
|
138
|
+
raise
|
|
139
|
+
if attempt < self.retries:
|
|
140
|
+
retry_after = getattr(exc, "retry_after", None)
|
|
141
|
+
if retry_after and retry_after > 0:
|
|
142
|
+
time.sleep(retry_after)
|
|
143
|
+
else:
|
|
144
|
+
time.sleep(min(2**attempt, 30))
|
|
145
|
+
|
|
146
|
+
if last_exc:
|
|
147
|
+
raise last_exc
|
|
148
|
+
|
|
149
|
+
@staticmethod
|
|
150
|
+
def _is_retryable(exc: Exception) -> bool:
|
|
151
|
+
if isinstance(exc, APIError):
|
|
152
|
+
if exc.status == "N/A":
|
|
153
|
+
return False
|
|
154
|
+
status = int(exc.status)
|
|
155
|
+
# 4xx is not retryable except 408 (timeout) and 429 (rate limit)
|
|
156
|
+
if 400 <= status < 500 and status not in (408, 429):
|
|
157
|
+
return False
|
|
158
|
+
return True
|
|
@@ -0,0 +1,126 @@
|
|
|
1
|
+
"""HTTP transport for batch event ingestion.
|
|
2
|
+
|
|
3
|
+
Handles gzip compression, retries with exponential backoff, and
|
|
4
|
+
Retry-After header parsing.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import gzip as gzip_mod
|
|
10
|
+
import json
|
|
11
|
+
import logging
|
|
12
|
+
from datetime import UTC, datetime
|
|
13
|
+
from io import BytesIO
|
|
14
|
+
|
|
15
|
+
import httpx
|
|
16
|
+
|
|
17
|
+
from litmus.version import VERSION
|
|
18
|
+
|
|
19
|
+
log = logging.getLogger("litmus")
|
|
20
|
+
|
|
21
|
+
DEFAULT_HOST = "https://ingest.trylitmus.com"
|
|
22
|
+
USER_AGENT = f"litmus-python/{VERSION}"
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class APIError(Exception):
|
|
26
|
+
def __init__(
|
|
27
|
+
self,
|
|
28
|
+
status: int | str,
|
|
29
|
+
message: str,
|
|
30
|
+
retry_after: float | None = None,
|
|
31
|
+
):
|
|
32
|
+
self.status = status
|
|
33
|
+
self.message = message
|
|
34
|
+
self.retry_after = retry_after
|
|
35
|
+
|
|
36
|
+
def __str__(self) -> str:
|
|
37
|
+
return f"[Litmus] {self.message} ({self.status})"
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
_client: httpx.Client | None = None
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def _get_client() -> httpx.Client:
|
|
44
|
+
global _client
|
|
45
|
+
if _client is None:
|
|
46
|
+
_client = httpx.Client(
|
|
47
|
+
transport=httpx.HTTPTransport(retries=2),
|
|
48
|
+
follow_redirects=True,
|
|
49
|
+
)
|
|
50
|
+
return _client
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def batch_post(
|
|
54
|
+
api_key: str,
|
|
55
|
+
host: str | None = None,
|
|
56
|
+
use_gzip: bool = False,
|
|
57
|
+
timeout: int = 15,
|
|
58
|
+
batch: list[dict] | None = None,
|
|
59
|
+
) -> None:
|
|
60
|
+
"""POST a batch of events to /v1/events.
|
|
61
|
+
|
|
62
|
+
This is the only network call the SDK makes. Everything else
|
|
63
|
+
feeds into a queue that eventually calls this.
|
|
64
|
+
"""
|
|
65
|
+
url = (host or DEFAULT_HOST).rstrip("/") + "/v1/events"
|
|
66
|
+
|
|
67
|
+
body = {
|
|
68
|
+
"events": batch or [],
|
|
69
|
+
"sent_at": datetime.now(tz=UTC).isoformat(),
|
|
70
|
+
}
|
|
71
|
+
data = json.dumps(body, default=_json_serializer)
|
|
72
|
+
|
|
73
|
+
headers = {
|
|
74
|
+
"Content-Type": "application/json",
|
|
75
|
+
"User-Agent": USER_AGENT,
|
|
76
|
+
"Authorization": f"Bearer {api_key}",
|
|
77
|
+
}
|
|
78
|
+
|
|
79
|
+
if use_gzip:
|
|
80
|
+
headers["Content-Encoding"] = "gzip"
|
|
81
|
+
buf = BytesIO()
|
|
82
|
+
with gzip_mod.GzipFile(fileobj=buf, mode="w") as gz:
|
|
83
|
+
gz.write(data.encode("utf-8"))
|
|
84
|
+
content = buf.getvalue()
|
|
85
|
+
else:
|
|
86
|
+
content = data.encode("utf-8")
|
|
87
|
+
|
|
88
|
+
res = _get_client().post(url, content=content, headers=headers, timeout=timeout)
|
|
89
|
+
|
|
90
|
+
if res.status_code in (200, 202):
|
|
91
|
+
log.debug("batch of %d events uploaded", len(batch or []))
|
|
92
|
+
return
|
|
93
|
+
|
|
94
|
+
# Parse Retry-After for the consumer's backoff logic
|
|
95
|
+
retry_after = _parse_retry_after(res)
|
|
96
|
+
|
|
97
|
+
try:
|
|
98
|
+
payload = res.json()
|
|
99
|
+
detail = payload.get("error", res.text)
|
|
100
|
+
except (ValueError, KeyError):
|
|
101
|
+
detail = res.text
|
|
102
|
+
|
|
103
|
+
raise APIError(res.status_code, detail, retry_after=retry_after)
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def _parse_retry_after(res: httpx.Response) -> float | None:
|
|
107
|
+
header = res.headers.get("Retry-After")
|
|
108
|
+
if not header:
|
|
109
|
+
return None
|
|
110
|
+
try:
|
|
111
|
+
return float(header)
|
|
112
|
+
except (ValueError, TypeError):
|
|
113
|
+
pass
|
|
114
|
+
try:
|
|
115
|
+
from email.utils import parsedate_to_datetime
|
|
116
|
+
|
|
117
|
+
target = parsedate_to_datetime(header)
|
|
118
|
+
return max(0.0, (target - datetime.now(UTC)).total_seconds())
|
|
119
|
+
except (ValueError, TypeError):
|
|
120
|
+
return None
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
def _json_serializer(obj: object) -> str:
|
|
124
|
+
if isinstance(obj, datetime):
|
|
125
|
+
return obj.isoformat()
|
|
126
|
+
raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
VERSION = "0.1.0"
|