diff --git a/py/src/braintrust/integrations/litellm/cassettes/1.74.0/test_litellm_amoderation.yaml b/py/src/braintrust/integrations/litellm/cassettes/1.74.0/test_litellm_amoderation.yaml new file mode 100644 index 00000000..1e4675c3 --- /dev/null +++ b/py/src/braintrust/integrations/litellm/cassettes/1.74.0/test_litellm_amoderation.yaml @@ -0,0 +1,178 @@ +interactions: + - request: + body: '{"input": "This is a test message", "model": "text-moderation-latest"}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - "70" + content-type: + - application/json + host: + - api.openai.com + user-agent: + - OpenAI/Python 1.92.3 + x-stainless-arch: + - arm64 + x-stainless-async: + - "false" + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 1.92.3 + x-stainless-read-timeout: + - "600" + x-stainless-retry-count: + - "0" + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.11.10 + method: POST + uri: https://api.openai.com/v1/moderations + response: + body: + string: !!binary | + H4sIAAAAAAAAA4ySTW/bMAyG7/kVgs+1Q0qiKPW2YrtuCDagh2EYjISJvfpjkJShQ9H/PrhrsDWz + s14IiHwpfjx8WClVtLviWhX9uIvlTbxNm7vdtySy2bwP+zcfP4R3b1sfb5tPXFxN6n7cSTclZLnP + 5fSKdW7HoQR4VkRJxy6n4lp9Ximl1MOTVarYd/XhIFO1fd0luTr5t3WWwxhbmXJOaqWKJPfHujuX + K1U0dZZZd6xT6mXIM8Ek3b5s6tjPxqZC674dxpgW6q1zE6XOMrTDYUbyox07GbayPsT6e9NuL3Ww + bof8nybX7ZByPG6nzaaLo76yr1PoOfJ4vvyfX9N2jIsEoAIABO8sezaaKaAx/xJ5kgEaqxmCxuCD + h8B6iZGtHLKFgGACsWUgKd0CM1+x1V4bZ7zxxNZKyRcQ2soDcCDS2gfnwb38eYYoV2A9kweNjGCJ + pAyXAWOlPQRnDFvrjQuAC+3/AW4qNNpaQHJIQITnUyzwx4otaQ1oHbIjYCn9a86BKtQ40aIACJrM + y3p/XcdvdAE0UtAhkCcwfLqW1cl+WT3+AgAA//8DAOLPZJ40BAAA + headers: + CF-RAY: + - 95cb009578ca964b-LAX + Connection: + - keep-alive + Content-Encoding: + - gzip + Content-Type: + - application/json + Date: + - Wed, 09 Jul 2025 21:44:22 GMT + Server: + - cloudflare + Set-Cookie: + - __cf_bm=ldaJXZHcRfZezUdPq5ALaHKdUQy72qnha5eMS7WlkM0-1752097462-1.0.1.1-sr6WeGVyNhiz_KP2Lk_.fA3kLI0cJ.SpHxCFISvi22KKahdPLnnGUU6wFf9er3ccyNjwbi.vPFOgxqQLtOqmXgLWAa3B8SMaNVuJIS0wU78; + path=/; expires=Wed, 09-Jul-25 22:14:22 GMT; domain=.api.openai.com; HttpOnly; + Secure; SameSite=None + - _cfuvid=eCkXR5wyooQ2qFREVv0FqyORwTycuziWmJAxvkwXBCM-1752097462868-0.0.1.1-604800000; + path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None + Transfer-Encoding: + - chunked + X-Content-Type-Options: + - nosniff + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + openai-organization: + - braintrust-data + openai-processing-ms: + - "168" + openai-version: + - "2020-10-01" + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + x-request-id: + - req_906e1e562ce7296e8e4ea6118065535d + status: + code: 200 + message: OK + - request: + body: '{"input": "This is a test message", "model": "text-moderation-latest"}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - "70" + content-type: + - application/json + host: + - api.openai.com + user-agent: + - OpenAI/Python 1.92.3 + x-stainless-arch: + - arm64 + x-stainless-async: + - "false" + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 1.92.3 + x-stainless-read-timeout: + - "600" + x-stainless-retry-count: + - "0" + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.11.10 + method: POST + uri: https://api.openai.com/v1/moderations + response: + body: + string: !!binary | + H4sIAAAAAAAAAwAAAP//jJLNbtswDMfveQrB59ohJVGUelsfoDt0wA7DMBiO4mjwRyspQ4qi7z64 + nbE1s9NeCIj8U/z48WkjRBF2xbUo+nEXy5v4Nd19aQC7B3fTxz3fhk9J/3z0p8+3d8XVpO7Hne+m + hOxPuZxesc5hHEoAflVEn45dTsW1+LYRQoinFytEse/qtvVTtX3dJX81+5s6+3aMwU85s1qIIvnT + se7O5UIUhzr7RXesU+r9kBeCyXf78lDHfjE2Fdr2YRhjWqm3zYfo6+yHMLQLkl9h7PzQ+G0b6/tD + aC51sA1DfqfJbRhSjsdm2my6OOoH+5pDfyLP58t//JGaMa4SgAoAEKzRbFlJJodK/U/kRQaotGRw + Ep11FhzLNUa6MsgaHIJyxJqBfGlWmNmKtbRSGWWVJdbal3wBoa4sADsiKa0zFszbnxeIcgXaMlmQ + yAiayJfuMmCspAVnlGKtrTIOcKX9v8BVhUpqDUgGCYjwfIoV/lixJikBtUE2BOxL+5FzoAolTrTI + AYIk9bbeP9fxis6BRHLSObIEiudr2cz2++b5NwAAAP//AwC9WqDuNAQAAA== + headers: + CF-RAY: + - 95cb00980e464896-LAX + Connection: + - keep-alive + Content-Encoding: + - gzip + Content-Type: + - application/json + Date: + - Wed, 09 Jul 2025 21:44:24 GMT + Server: + - cloudflare + Set-Cookie: + - __cf_bm=VQ8b7SImQKVNCIRiMnfxH9.VG3iHeyazbvGOjIXA.qM-1752097464-1.0.1.1-XxY1zHj4dDcIvzE.saBV8uG7R62ARV7U24xTVGKz2Avhl0vz3bmuvZajl9t3blNdf9XEN69FSWuNYfMeTGNjIgkwiKRGg3uDzZpq1PobzkU; + path=/; expires=Wed, 09-Jul-25 22:14:24 GMT; domain=.api.openai.com; HttpOnly; + Secure; SameSite=None + - _cfuvid=0pvwsu4lhmDmwRvRF5PRcQD3zZ06mdYREXQ8lSIbpBA-1752097464624-0.0.1.1-604800000; + path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None + Transfer-Encoding: + - chunked + X-Content-Type-Options: + - nosniff + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + openai-organization: + - braintrust-data + openai-processing-ms: + - "1501" + openai-version: + - "2020-10-01" + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + x-request-id: + - req_3dbafe89ead9b5efcfac20878a15ec19 + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/litellm/cassettes/latest/test_litellm_amoderation.yaml b/py/src/braintrust/integrations/litellm/cassettes/latest/test_litellm_amoderation.yaml new file mode 100644 index 00000000..52f8aabb --- /dev/null +++ b/py/src/braintrust/integrations/litellm/cassettes/latest/test_litellm_amoderation.yaml @@ -0,0 +1,112 @@ +interactions: +- request: + body: '{"input":"This is a test message","model":"omni-moderation-latest"}' + headers: + Accept: + - application/json + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '67' + Content-Type: + - application/json + Host: + - api.openai.com + User-Agent: + - OpenAI/Python 2.37.0 + X-Stainless-Arch: + - arm64 + X-Stainless-Async: + - 'false' + X-Stainless-Lang: + - python + X-Stainless-OS: + - MacOS + X-Stainless-Package-Version: + - 2.37.0 + X-Stainless-Runtime: + - CPython + X-Stainless-Runtime-Version: + - 3.12.12 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + method: POST + uri: https://api.openai.com/v1/moderations + response: + body: + string: "{\n \"id\": \"modr-5354\",\n \"model\": \"omni-moderation-latest\",\n + \ \"results\": [\n {\n \"flagged\": false,\n \"categories\": + {\n \"harassment\": false,\n \"harassment/threatening\": false,\n + \ \"sexual\": false,\n \"hate\": false,\n \"hate/threatening\": + false,\n \"illicit\": false,\n \"illicit/violent\": false,\n + \ \"self-harm/intent\": false,\n \"self-harm/instructions\": + false,\n \"self-harm\": false,\n \"sexual/minors\": false,\n + \ \"violence\": false,\n \"violence/graphic\": false\n },\n + \ \"category_scores\": {\n \"harassment\": 0.000047285443452031076,\n + \ \"harassment/threatening\": 4.264746818557914e-6,\n \"sexual\": + 0.000027803096387751555,\n \"hate\": 0.000010554685795431098,\n \"hate/threatening\": + 2.561282210758673e-7,\n \"illicit\": 0.00004108485376346404,\n \"illicit/violent\": + 8.750299760661308e-6,\n \"self-harm/intent\": 0.00021166499492301485,\n + \ \"self-harm/instructions\": 1.3846004563753396e-6,\n \"self-harm\": + 8.61465062380632e-6,\n \"sexual/minors\": 2.5466403947055455e-6,\n + \ \"violence\": 0.00048297182378774457,\n \"violence/graphic\": + 4.832563818725537e-6\n },\n \"category_applied_input_types\": {\n + \ \"harassment\": [\n \"text\"\n ],\n \"harassment/threatening\": + [\n \"text\"\n ],\n \"sexual\": [\n \"text\"\n + \ ],\n \"hate\": [\n \"text\"\n ],\n \"hate/threatening\": + [\n \"text\"\n ],\n \"illicit\": [\n \"text\"\n + \ ],\n \"illicit/violent\": [\n \"text\"\n ],\n + \ \"self-harm/intent\": [\n \"text\"\n ],\n \"self-harm/instructions\": + [\n \"text\"\n ],\n \"self-harm\": [\n \"text\"\n + \ ],\n \"sexual/minors\": [\n \"text\"\n ],\n + \ \"violence\": [\n \"text\"\n ],\n \"violence/graphic\": + [\n \"text\"\n ]\n }\n }\n ]\n}" + headers: + Access-Control-Expose-Headers: + - CF-Ray + CF-RAY: + - 9fe42f3509c3ab8d-YYZ + Connection: + - keep-alive + Content-Type: + - application/json + Date: + - Tue, 19 May 2026 15:37:42 GMT + Server: + - cloudflare + Strict-Transport-Security: + - max-age=31536000; includeSubDomains; preload + Transfer-Encoding: + - chunked + X-Content-Type-Options: + - nosniff + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + content-length: + - '1971' + openai-organization: + - braintrust-data + openai-processing-ms: + - '194' + openai-project: + - proj_vsCSXafhhByzWOThMrJcZiw9 + openai-version: + - '2020-10-01' + set-cookie: + - __cf_bm=pqYA2HZT90k8h4TqB6S59Yjr4mUngnMhMol6ACAtMsY-1779205061.923457-1.0.1.1-hNTI8nbdvEUMYr26oAUuRpiQPLiy5uiYvjnY2lSqDPhHoez7.JjXKi0AJpg0I1nsdRH1OdEy.yrOkjK2KwEWd8LhIJYB9UH9IBvyfyqNH3SnEXIFYrpAsbRYllQegXLX; + HttpOnly; SameSite=None; Secure; Path=/; Domain=api.openai.com; Expires=Tue, + 19 May 2026 16:07:42 GMT + x-openai-proxy-wasm: + - v0.1 + x-request-id: + - req_051677dbbf2043f0918d7fbdbe2174e7 + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/litellm/patchers.py b/py/src/braintrust/integrations/litellm/patchers.py index 470ff361..75ec6791 100644 --- a/py/src/braintrust/integrations/litellm/patchers.py +++ b/py/src/braintrust/integrations/litellm/patchers.py @@ -8,6 +8,7 @@ _acompletion_wrapper_async, _aembedding_wrapper_async, _aimage_generation_wrapper_async, + _amoderation_wrapper_async, _arerank_wrapper_async, _aresponses_wrapper_async, _aspeech_wrapper_async, @@ -96,6 +97,12 @@ class LiteLLMModerationPatcher(FunctionWrapperPatcher): wrapper = _moderation_wrapper +class LiteLLMAModerationPatcher(FunctionWrapperPatcher): + name = "litellm.amoderation" + target_path = "amoderation" + wrapper = _amoderation_wrapper_async + + class LiteLLMSpeechPatcher(FunctionWrapperPatcher): name = "litellm.speech" target_path = "speech" @@ -148,6 +155,7 @@ class LiteLLMArerankPatcher(FunctionWrapperPatcher): LiteLLMEmbeddingPatcher, LiteLLMAembeddingPatcher, LiteLLMModerationPatcher, + LiteLLMAModerationPatcher, LiteLLMSpeechPatcher, LiteLLMAspeechPatcher, LiteLLMTranscriptionPatcher, @@ -170,9 +178,9 @@ def wrap_litellm(litellm: Any) -> Any: that exposes the same top-level callables such as ``completion``, ``acompletion``, ``responses``, ``aresponses``, ``image_generation``, ``text_completion``, ``atext_completion``, ``aimage_generation``, - ``embedding``, ``aembedding``, ``moderation``, ``speech``, ``aspeech``, - ``transcription``, ``atranscription``, ``rerank``, and ``arerank``). Each - patcher is applied idempotently — calling + ``embedding``, ``aembedding``, ``moderation``, ``amoderation``, ``speech``, + ``aspeech``, ``transcription``, ``atranscription``, ``rerank``, and + ``arerank``). Each patcher is applied idempotently — calling ``wrap_litellm`` twice on the same object is safe. Args: diff --git a/py/src/braintrust/integrations/litellm/test_litellm.py b/py/src/braintrust/integrations/litellm/test_litellm.py index c8ce6820..02e867af 100644 --- a/py/src/braintrust/integrations/litellm/test_litellm.py +++ b/py/src/braintrust/integrations/litellm/test_litellm.py @@ -347,6 +347,24 @@ def test_litellm_moderation(memory_logger): assert "This is a test message" in str(span["input"]) +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_litellm_amoderation(memory_logger): + assert not memory_logger.pop() + + response = await litellm.amoderation(model="omni-moderation-latest", input="This is a test message") + + assert response + assert response.results + + spans = memory_logger.pop() + assert len(spans) == 1 + span = spans[0] + assert span["metadata"]["model"] == "omni-moderation-latest" + assert span["metadata"]["provider"] == "litellm" + assert "This is a test message" in str(span["input"]) + + @pytest.mark.vcr def test_litellm_image_generation(memory_logger): assert not memory_logger.pop() diff --git a/py/src/braintrust/integrations/litellm/tracing.py b/py/src/braintrust/integrations/litellm/tracing.py index 96bde9aa..e10700f3 100644 --- a/py/src/braintrust/integrations/litellm/tracing.py +++ b/py/src/braintrust/integrations/litellm/tracing.py @@ -442,6 +442,16 @@ async def _aembedding_wrapper_async(wrapped, instance, args, kwargs): return embedding_response +def _log_moderation_response(span: Span, moderation_response: Any) -> None: + log_response = _try_to_dict(moderation_response) + usage = log_response.get("usage") + metrics = _parse_metrics_from_usage(usage) + span.log( + metrics=metrics, + output=log_response["results"], + ) + + def _moderation_wrapper(wrapped, instance, args, kwargs): """wrapt wrapper for litellm.moderation.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") @@ -450,13 +460,19 @@ def _moderation_wrapper(wrapped, instance, args, kwargs): **merge_dicts(dict(name="Moderation", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) ) as span: moderation_response = wrapped(*args, **kwargs) - log_response = _try_to_dict(moderation_response) - usage = log_response.get("usage") - metrics = _parse_metrics_from_usage(usage) - span.log( - metrics=metrics, - output=log_response["results"], - ) + _log_moderation_response(span, moderation_response) + return moderation_response + + +async def _amoderation_wrapper_async(wrapped, instance, args, kwargs): + """wrapt wrapper for litellm.amoderation.""" + updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") + + with start_span( + **merge_dicts(dict(name="Moderation", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) + ) as span: + moderation_response = await wrapped(*args, **kwargs) + _log_moderation_response(span, moderation_response) return moderation_response