Skip to content

Commit 6963b97

Browse files
AbhiPrasadAbhiPrasad
authored
feat: trace litellm async moderation (#448)
resolves #402
1 parent de28291 commit 6963b97

5 files changed

Lines changed: 342 additions & 10 deletions

File tree

‎py/src/braintrust/integrations/litellm/cassettes/1.74.0/test_litellm_amoderation.yaml‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
interactions:
2+
- request:
3+
body: '{"input": "This is a test message", "model": "text-moderation-latest"}'
4+
headers:
5+
accept:
6+
- application/json
7+
accept-encoding:
8+
- gzip, deflate, zstd
9+
connection:
10+
- keep-alive
11+
content-length:
12+
- "70"
13+
content-type:
14+
- application/json
15+
host:
16+
- api.openai.com
17+
user-agent:
18+
- OpenAI/Python 1.92.3
19+
x-stainless-arch:
20+
- arm64
21+
x-stainless-async:
22+
- "false"
23+
x-stainless-lang:
24+
- python
25+
x-stainless-os:
26+
- MacOS
27+
x-stainless-package-version:
28+
- 1.92.3
29+
x-stainless-read-timeout:
30+
- "600"
31+
x-stainless-retry-count:
32+
- "0"
33+
x-stainless-runtime:
34+
- CPython
35+
x-stainless-runtime-version:
36+
- 3.11.10
37+
method: POST
38+
uri: https://api.openai.com/v1/moderations
39+
response:
40+
body:
41+
string: !!binary |
42+
H4sIAAAAAAAAA4ySTW/bMAyG7/kVgs+1Q0qiKPW2YrtuCDagh2EYjISJvfpjkJShQ9H/PrhrsDWz
43+
s14IiHwpfjx8WClVtLviWhX9uIvlTbxNm7vdtySy2bwP+zcfP4R3b1sfb5tPXFxN6n7cSTclZLnP
44+
5fSKdW7HoQR4VkRJxy6n4lp9Ximl1MOTVarYd/XhIFO1fd0luTr5t3WWwxhbmXJOaqWKJPfHujuX
45+
K1U0dZZZd6xT6mXIM8Ek3b5s6tjPxqZC674dxpgW6q1zE6XOMrTDYUbyox07GbayPsT6e9NuL3Ww
46+
bof8nybX7ZByPG6nzaaLo76yr1PoOfJ4vvyfX9N2jIsEoAIABO8sezaaKaAx/xJ5kgEaqxmCxuCD
47+
h8B6iZGtHLKFgGACsWUgKd0CM1+x1V4bZ7zxxNZKyRcQ2soDcCDS2gfnwb38eYYoV2A9kweNjGCJ
48+
pAyXAWOlPQRnDFvrjQuAC+3/AW4qNNpaQHJIQITnUyzwx4otaQ1oHbIjYCn9a86BKtQ40aIACJrM
49+
y3p/XcdvdAE0UtAhkCcwfLqW1cl+WT3+AgAA//8DAOLPZJ40BAAA
50+
headers:
51+
CF-RAY:
52+
- 95cb009578ca964b-LAX
53+
Connection:
54+
- keep-alive
55+
Content-Encoding:
56+
- gzip
57+
Content-Type:
58+
- application/json
59+
Date:
60+
- Wed, 09 Jul 2025 21:44:22 GMT
61+
Server:
62+
- cloudflare
63+
Set-Cookie:
64+
- __cf_bm=ldaJXZHcRfZezUdPq5ALaHKdUQy72qnha5eMS7WlkM0-1752097462-1.0.1.1-sr6WeGVyNhiz_KP2Lk_.fA3kLI0cJ.SpHxCFISvi22KKahdPLnnGUU6wFf9er3ccyNjwbi.vPFOgxqQLtOqmXgLWAa3B8SMaNVuJIS0wU78;
65+
path=/; expires=Wed, 09-Jul-25 22:14:22 GMT; domain=.api.openai.com; HttpOnly;
66+
Secure; SameSite=None
67+
- _cfuvid=eCkXR5wyooQ2qFREVv0FqyORwTycuziWmJAxvkwXBCM-1752097462868-0.0.1.1-604800000;
68+
path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
69+
Transfer-Encoding:
70+
- chunked
71+
X-Content-Type-Options:
72+
- nosniff
73+
alt-svc:
74+
- h3=":443"; ma=86400
75+
cf-cache-status:
76+
- DYNAMIC
77+
openai-organization:
78+
- braintrust-data
79+
openai-processing-ms:
80+
- "168"
81+
openai-version:
82+
- "2020-10-01"
83+
strict-transport-security:
84+
- max-age=31536000; includeSubDomains; preload
85+
x-request-id:
86+
- req_906e1e562ce7296e8e4ea6118065535d
87+
status:
88+
code: 200
89+
message: OK
90+
- request:
91+
body: '{"input": "This is a test message", "model": "text-moderation-latest"}'
92+
headers:
93+
accept:
94+
- application/json
95+
accept-encoding:
96+
- gzip, deflate, zstd
97+
connection:
98+
- keep-alive
99+
content-length:
100+
- "70"
101+
content-type:
102+
- application/json
103+
host:
104+
- api.openai.com
105+
user-agent:
106+
- OpenAI/Python 1.92.3
107+
x-stainless-arch:
108+
- arm64
109+
x-stainless-async:
110+
- "false"
111+
x-stainless-lang:
112+
- python
113+
x-stainless-os:
114+
- MacOS
115+
x-stainless-package-version:
116+
- 1.92.3
117+
x-stainless-read-timeout:
118+
- "600"
119+
x-stainless-retry-count:
120+
- "0"
121+
x-stainless-runtime:
122+
- CPython
123+
x-stainless-runtime-version:
124+
- 3.11.10
125+
method: POST
126+
uri: https://api.openai.com/v1/moderations
127+
response:
128+
body:
129+
string: !!binary |
130+
H4sIAAAAAAAAAwAAAP//jJLNbtswDMfveQrB59ohJVGUelsfoDt0wA7DMBiO4mjwRyspQ4qi7z64
131+
nbE1s9NeCIj8U/z48WkjRBF2xbUo+nEXy5v4Nd19aQC7B3fTxz3fhk9J/3z0p8+3d8XVpO7Hne+m
132+
hOxPuZxesc5hHEoAflVEn45dTsW1+LYRQoinFytEse/qtvVTtX3dJX81+5s6+3aMwU85s1qIIvnT
133+
se7O5UIUhzr7RXesU+r9kBeCyXf78lDHfjE2Fdr2YRhjWqm3zYfo6+yHMLQLkl9h7PzQ+G0b6/tD
134+
aC51sA1DfqfJbRhSjsdm2my6OOoH+5pDfyLP58t//JGaMa4SgAoAEKzRbFlJJodK/U/kRQaotGRw
135+
Ep11FhzLNUa6MsgaHIJyxJqBfGlWmNmKtbRSGWWVJdbal3wBoa4sADsiKa0zFszbnxeIcgXaMlmQ
136+
yAiayJfuMmCspAVnlGKtrTIOcKX9v8BVhUpqDUgGCYjwfIoV/lixJikBtUE2BOxL+5FzoAolTrTI
137+
AYIk9bbeP9fxis6BRHLSObIEiudr2cz2++b5NwAAAP//AwC9WqDuNAQAAA==
138+
headers:
139+
CF-RAY:
140+
- 95cb00980e464896-LAX
141+
Connection:
142+
- keep-alive
143+
Content-Encoding:
144+
- gzip
145+
Content-Type:
146+
- application/json
147+
Date:
148+
- Wed, 09 Jul 2025 21:44:24 GMT
149+
Server:
150+
- cloudflare
151+
Set-Cookie:
152+
- __cf_bm=VQ8b7SImQKVNCIRiMnfxH9.VG3iHeyazbvGOjIXA.qM-1752097464-1.0.1.1-XxY1zHj4dDcIvzE.saBV8uG7R62ARV7U24xTVGKz2Avhl0vz3bmuvZajl9t3blNdf9XEN69FSWuNYfMeTGNjIgkwiKRGg3uDzZpq1PobzkU;
153+
path=/; expires=Wed, 09-Jul-25 22:14:24 GMT; domain=.api.openai.com; HttpOnly;
154+
Secure; SameSite=None
155+
- _cfuvid=0pvwsu4lhmDmwRvRF5PRcQD3zZ06mdYREXQ8lSIbpBA-1752097464624-0.0.1.1-604800000;
156+
path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
157+
Transfer-Encoding:
158+
- chunked
159+
X-Content-Type-Options:
160+
- nosniff
161+
alt-svc:
162+
- h3=":443"; ma=86400
163+
cf-cache-status:
164+
- DYNAMIC
165+
openai-organization:
166+
- braintrust-data
167+
openai-processing-ms:
168+
- "1501"
169+
openai-version:
170+
- "2020-10-01"
171+
strict-transport-security:
172+
- max-age=31536000; includeSubDomains; preload
173+
x-request-id:
174+
- req_3dbafe89ead9b5efcfac20878a15ec19
175+
status:
176+
code: 200
177+
message: OK
178+
version: 1

‎py/src/braintrust/integrations/litellm/cassettes/latest/test_litellm_amoderation.yaml‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
interactions:
2+
- request:
3+
body: '{"input":"This is a test message","model":"omni-moderation-latest"}'
4+
headers:
5+
Accept:
6+
- application/json
7+
Accept-Encoding:
8+
- gzip, deflate
9+
Connection:
10+
- keep-alive
11+
Content-Length:
12+
- '67'
13+
Content-Type:
14+
- application/json
15+
Host:
16+
- api.openai.com
17+
User-Agent:
18+
- OpenAI/Python 2.37.0
19+
X-Stainless-Arch:
20+
- arm64
21+
X-Stainless-Async:
22+
- 'false'
23+
X-Stainless-Lang:
24+
- python
25+
X-Stainless-OS:
26+
- MacOS
27+
X-Stainless-Package-Version:
28+
- 2.37.0
29+
X-Stainless-Runtime:
30+
- CPython
31+
X-Stainless-Runtime-Version:
32+
- 3.12.12
33+
x-stainless-read-timeout:
34+
- '600'
35+
x-stainless-retry-count:
36+
- '0'
37+
method: POST
38+
uri: https://api.openai.com/v1/moderations
39+
response:
40+
body:
41+
string: "{\n \"id\": \"modr-5354\",\n \"model\": \"omni-moderation-latest\",\n
42+
\ \"results\": [\n {\n \"flagged\": false,\n \"categories\":
43+
{\n \"harassment\": false,\n \"harassment/threatening\": false,\n
44+
\ \"sexual\": false,\n \"hate\": false,\n \"hate/threatening\":
45+
false,\n \"illicit\": false,\n \"illicit/violent\": false,\n
46+
\ \"self-harm/intent\": false,\n \"self-harm/instructions\":
47+
false,\n \"self-harm\": false,\n \"sexual/minors\": false,\n
48+
\ \"violence\": false,\n \"violence/graphic\": false\n },\n
49+
\ \"category_scores\": {\n \"harassment\": 0.000047285443452031076,\n
50+
\ \"harassment/threatening\": 4.264746818557914e-6,\n \"sexual\":
51+
0.000027803096387751555,\n \"hate\": 0.000010554685795431098,\n \"hate/threatening\":
52+
2.561282210758673e-7,\n \"illicit\": 0.00004108485376346404,\n \"illicit/violent\":
53+
8.750299760661308e-6,\n \"self-harm/intent\": 0.00021166499492301485,\n
54+
\ \"self-harm/instructions\": 1.3846004563753396e-6,\n \"self-harm\":
55+
8.61465062380632e-6,\n \"sexual/minors\": 2.5466403947055455e-6,\n
56+
\ \"violence\": 0.00048297182378774457,\n \"violence/graphic\":
57+
4.832563818725537e-6\n },\n \"category_applied_input_types\": {\n
58+
\ \"harassment\": [\n \"text\"\n ],\n \"harassment/threatening\":
59+
[\n \"text\"\n ],\n \"sexual\": [\n \"text\"\n
60+
\ ],\n \"hate\": [\n \"text\"\n ],\n \"hate/threatening\":
61+
[\n \"text\"\n ],\n \"illicit\": [\n \"text\"\n
62+
\ ],\n \"illicit/violent\": [\n \"text\"\n ],\n
63+
\ \"self-harm/intent\": [\n \"text\"\n ],\n \"self-harm/instructions\":
64+
[\n \"text\"\n ],\n \"self-harm\": [\n \"text\"\n
65+
\ ],\n \"sexual/minors\": [\n \"text\"\n ],\n
66+
\ \"violence\": [\n \"text\"\n ],\n \"violence/graphic\":
67+
[\n \"text\"\n ]\n }\n }\n ]\n}"
68+
headers:
69+
Access-Control-Expose-Headers:
70+
- CF-Ray
71+
CF-RAY:
72+
- 9fe42f3509c3ab8d-YYZ
73+
Connection:
74+
- keep-alive
75+
Content-Type:
76+
- application/json
77+
Date:
78+
- Tue, 19 May 2026 15:37:42 GMT
79+
Server:
80+
- cloudflare
81+
Strict-Transport-Security:
82+
- max-age=31536000; includeSubDomains; preload
83+
Transfer-Encoding:
84+
- chunked
85+
X-Content-Type-Options:
86+
- nosniff
87+
alt-svc:
88+
- h3=":443"; ma=86400
89+
cf-cache-status:
90+
- DYNAMIC
91+
content-length:
92+
- '1971'
93+
openai-organization:
94+
- braintrust-data
95+
openai-processing-ms:
96+
- '194'
97+
openai-project:
98+
- proj_vsCSXafhhByzWOThMrJcZiw9
99+
openai-version:
100+
- '2020-10-01'
101+
set-cookie:
102+
- __cf_bm=pqYA2HZT90k8h4TqB6S59Yjr4mUngnMhMol6ACAtMsY-1779205061.923457-1.0.1.1-hNTI8nbdvEUMYr26oAUuRpiQPLiy5uiYvjnY2lSqDPhHoez7.JjXKi0AJpg0I1nsdRH1OdEy.yrOkjK2KwEWd8LhIJYB9UH9IBvyfyqNH3SnEXIFYrpAsbRYllQegXLX;
103+
HttpOnly; SameSite=None; Secure; Path=/; Domain=api.openai.com; Expires=Tue,
104+
19 May 2026 16:07:42 GMT
105+
x-openai-proxy-wasm:
106+
- v0.1
107+
x-request-id:
108+
- req_051677dbbf2043f0918d7fbdbe2174e7
109+
status:
110+
code: 200
111+
message: OK
112+
version: 1

‎py/src/braintrust/integrations/litellm/patchers.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
_acompletion_wrapper_async,
99
_aembedding_wrapper_async,
1010
_aimage_generation_wrapper_async,
11+
_amoderation_wrapper_async,
1112
_arerank_wrapper_async,
1213
_aresponses_wrapper_async,
1314
_aspeech_wrapper_async,
@@ -96,6 +97,12 @@ class LiteLLMModerationPatcher(FunctionWrapperPatcher):
9697
wrapper = _moderation_wrapper
9798

9899

100+
class LiteLLMAModerationPatcher(FunctionWrapperPatcher):
101+
name = "litellm.amoderation"
102+
target_path = "amoderation"
103+
wrapper = _amoderation_wrapper_async
104+
105+
99106
class LiteLLMSpeechPatcher(FunctionWrapperPatcher):
100107
name = "litellm.speech"
101108
target_path = "speech"
@@ -148,6 +155,7 @@ class LiteLLMArerankPatcher(FunctionWrapperPatcher):
148155
LiteLLMEmbeddingPatcher,
149156
LiteLLMAembeddingPatcher,
150157
LiteLLMModerationPatcher,
158+
LiteLLMAModerationPatcher,
151159
LiteLLMSpeechPatcher,
152160
LiteLLMAspeechPatcher,
153161
LiteLLMTranscriptionPatcher,
@@ -170,9 +178,9 @@ def wrap_litellm(litellm: Any) -> Any:
170178
that exposes the same top-level callables such as ``completion``,
171179
``acompletion``, ``responses``, ``aresponses``, ``image_generation``,
172180
``text_completion``, ``atext_completion``, ``aimage_generation``,
173-
``embedding``, ``aembedding``, ``moderation``, ``speech``, ``aspeech``,
174-
``transcription``, ``atranscription``, ``rerank``, and ``arerank``). Each
175-
patcher is applied idempotently — calling
181+
``embedding``, ``aembedding``, ``moderation``, ``amoderation``, ``speech``,
182+
``aspeech``, ``transcription``, ``atranscription``, ``rerank``, and
183+
``arerank``). Each patcher is applied idempotently — calling
176184
``wrap_litellm`` twice on the same object is safe.
177185
178186
Args:

‎py/src/braintrust/integrations/litellm/test_litellm.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,24 @@ def test_litellm_moderation(memory_logger):
347347
assert "This is a test message" in str(span["input"])
348348

349349

350+
@pytest.mark.vcr
351+
@pytest.mark.asyncio
352+
async def test_litellm_amoderation(memory_logger):
353+
assert not memory_logger.pop()
354+
355+
response = await litellm.amoderation(model="omni-moderation-latest", input="This is a test message")
356+
357+
assert response
358+
assert response.results
359+
360+
spans = memory_logger.pop()
361+
assert len(spans) == 1
362+
span = spans[0]
363+
assert span["metadata"]["model"] == "omni-moderation-latest"
364+
assert span["metadata"]["provider"] == "litellm"
365+
assert "This is a test message" in str(span["input"])
366+
367+
350368
@pytest.mark.vcr
351369
def test_litellm_image_generation(memory_logger):
352370
assert not memory_logger.pop()

0 commit comments

Comments
 (0)