Skip to content

Commit bf98dc9

Browse files
authored
feat(python-sdk): record_usage() accepts **kwargs and handles 204 (#13)
record_usage() / async_record_usage() now match the 1.4c API contract: - accept arbitrary **kwargs; recognised keys (output_tokens, cost, provider, tool_id, correlation_id, metadata) map to camelCase payload fields, unknown keys pass through unchanged - unify sync/async payload construction via _build_usage_payload - error message now points the caller at actionable next steps Adds unit coverage: - 204 No Content success path (sync + async) — SDK must not raise or parse body - kwargs-forwarding test verifies recognised-key mapping and pass-through semantics Refs: GOV-576
1 parent 78f52d0 commit bf98dc9

2 files changed

Lines changed: 98 additions & 23 deletions

File tree

src/governs_ai/client.py

Lines changed: 57 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -192,32 +192,69 @@ async def async_precheck(
192192
# 1.4c — record_usage()
193193
# ------------------------------------------------------------------
194194

195+
_USAGE_KWARG_MAP = {
196+
"output_tokens": "outputTokens",
197+
"tokens_out": "outputTokens",
198+
"provider": "provider",
199+
"cost": "cost",
200+
"tool": "toolId",
201+
"tool_id": "toolId",
202+
"correlation_id": "correlationId",
203+
"metadata": "metadata",
204+
}
205+
206+
def _build_usage_payload(
207+
self,
208+
org_id: str,
209+
user_id: str,
210+
tokens: int,
211+
model: str,
212+
extras: Dict[str, Any],
213+
) -> Dict[str, Any]:
214+
payload: Dict[str, Any] = {
215+
"orgId": org_id or self.org_id,
216+
"userId": user_id,
217+
"inputTokens": tokens,
218+
"outputTokens": 0,
219+
"model": model,
220+
"provider": "openai",
221+
}
222+
for key, value in extras.items():
223+
if value is None:
224+
continue
225+
payload_key = self._USAGE_KWARG_MAP.get(key, key)
226+
payload[payload_key] = value
227+
return payload
228+
195229
def record_usage(
196230
self,
197231
org_id: str,
198232
user_id: str,
199233
tokens: int,
200234
model: str,
201-
*,
202-
provider: str = "openai",
235+
**kwargs: Any,
203236
) -> None:
204237
"""Record token usage for a model request.
205238
239+
Args:
240+
org_id: Organization ID (falls back to ``client.org_id``).
241+
user_id: End-user identifier.
242+
tokens: Input token count for the request.
243+
model: Model identifier (e.g., ``"gpt-4o-mini"``).
244+
**kwargs: Optional extras forwarded to the platform API. Recognised
245+
keys: ``output_tokens``, ``provider``, ``cost``, ``tool_id``,
246+
``correlation_id``, ``metadata``. Unknown keys are passed
247+
through unchanged.
248+
206249
Example::
207250
208251
client.record_usage(
209252
org_id="org-1", user_id="user-123",
210253
tokens=180, model="gpt-4o-mini",
254+
output_tokens=42, cost=0.0012,
211255
)
212256
"""
213-
payload: Dict[str, Any] = {
214-
"orgId": org_id or self.org_id,
215-
"userId": user_id,
216-
"inputTokens": tokens,
217-
"outputTokens": 0,
218-
"model": model,
219-
"provider": provider,
220-
}
257+
payload = self._build_usage_payload(org_id, user_id, tokens, model, kwargs)
221258
with httpx.Client(timeout=self.timeout) as http:
222259
resp = http.post(
223260
f"{self.base_url}/api/v1/usage",
@@ -226,7 +263,8 @@ def record_usage(
226263
)
227264
if resp.status_code >= 400:
228265
raise GovernsAIError(
229-
f"record_usage failed with HTTP {resp.status_code}: {resp.text}",
266+
f"record_usage failed with HTTP {resp.status_code}: {resp.text} — "
267+
f"verify org_id/user_id and that the API key has usage write scope",
230268
status_code=resp.status_code,
231269
)
232270

@@ -236,18 +274,13 @@ async def async_record_usage(
236274
user_id: str,
237275
tokens: int,
238276
model: str,
239-
*,
240-
provider: str = "openai",
277+
**kwargs: Any,
241278
) -> None:
242-
"""Async variant of :meth:`record_usage`."""
243-
payload: Dict[str, Any] = {
244-
"orgId": org_id or self.org_id,
245-
"userId": user_id,
246-
"inputTokens": tokens,
247-
"outputTokens": 0,
248-
"model": model,
249-
"provider": provider,
250-
}
279+
"""Async variant of :meth:`record_usage`.
280+
281+
Accepts the same arguments and kwargs as the sync form.
282+
"""
283+
payload = self._build_usage_payload(org_id, user_id, tokens, model, kwargs)
251284
async with httpx.AsyncClient(timeout=self.timeout) as http:
252285
resp = await http.post(
253286
f"{self.base_url}/api/v1/usage",
@@ -256,7 +289,8 @@ async def async_record_usage(
256289
)
257290
if resp.status_code >= 400:
258291
raise GovernsAIError(
259-
f"record_usage failed with HTTP {resp.status_code}: {resp.text}",
292+
f"record_usage failed with HTTP {resp.status_code}: {resp.text} — "
293+
f"verify org_id/user_id and that the API key has usage write scope",
260294
status_code=resp.status_code,
261295
)
262296

tests/test_record_usage_budget.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,47 @@ async def test_async_record_usage_sends_correct_payload(client):
4141
assert body["inputTokens"] == 50
4242

4343

44+
@respx.mock
45+
def test_record_usage_succeeds_on_204_no_content(client):
46+
"""Platform API returns 204 No Content — SDK must not raise or try to parse body."""
47+
respx.post(f"{BASE}/api/v1/usage").mock(return_value=httpx.Response(204))
48+
client.record_usage(org_id="org-1", user_id="user-123", tokens=10, model="gpt-4o")
49+
50+
51+
@respx.mock
52+
async def test_async_record_usage_succeeds_on_204_no_content(client):
53+
respx.post(f"{BASE}/api/v1/usage").mock(return_value=httpx.Response(204))
54+
await client.async_record_usage(
55+
org_id="org-1", user_id="user-123", tokens=10, model="gpt-4o"
56+
)
57+
58+
59+
@respx.mock
60+
def test_record_usage_forwards_kwargs_to_payload(client):
61+
"""Recognised kwargs are mapped to camelCase platform fields; unknown kwargs pass through."""
62+
route = respx.post(f"{BASE}/api/v1/usage").mock(return_value=httpx.Response(204))
63+
client.record_usage(
64+
org_id="org-1",
65+
user_id="user-123",
66+
tokens=100,
67+
model="gpt-4o",
68+
output_tokens=50,
69+
cost=0.0012,
70+
provider="anthropic",
71+
tool_id="web_search",
72+
correlation_id="req-abc",
73+
metadata={"session": "s1"},
74+
)
75+
body = json.loads(route.calls[0].request.content)
76+
assert body["inputTokens"] == 100
77+
assert body["outputTokens"] == 50
78+
assert body["cost"] == 0.0012
79+
assert body["provider"] == "anthropic"
80+
assert body["toolId"] == "web_search"
81+
assert body["correlationId"] == "req-abc"
82+
assert body["metadata"] == {"session": "s1"}
83+
84+
4485
@respx.mock
4586
def test_budget_check_allowed(client):
4687
respx.get(f"{BASE}/api/v1/budget/context").mock(

0 commit comments

Comments
 (0)