From f90b8a4955de998eed3af665e5284400569800bf Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 22 Mar 2025 20:20:36 -0700 Subject: [PATCH] docs: document streaming iterator guardrail --- docs/my-website/docs/providers/ollama.md | 2 +- .../docs/proxy/guardrails/custom_guardrail.md | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/docs/my-website/docs/providers/ollama.md b/docs/my-website/docs/providers/ollama.md index 21dfde9982..d59d9dd0ce 100644 --- a/docs/my-website/docs/providers/ollama.md +++ b/docs/my-website/docs/providers/ollama.md @@ -240,7 +240,7 @@ model_list: ```bash litellm --config /path/to/config.yaml --detailed_debug -# RUNNING ON http://0.0.0.0:4000 +# RUNNING ON http://0.0.0.0:4000 ``` 3. Test it! diff --git a/docs/my-website/docs/proxy/guardrails/custom_guardrail.md b/docs/my-website/docs/proxy/guardrails/custom_guardrail.md index 50deac511f..ee1a7a94fe 100644 --- a/docs/my-website/docs/proxy/guardrails/custom_guardrail.md +++ b/docs/my-website/docs/proxy/guardrails/custom_guardrail.md @@ -128,6 +128,21 @@ class myCustomGuardrail(CustomGuardrail): ): raise ValueError("Guardrail failed Coffee Detected") + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + request_data: dict, + ) -> AsyncGenerator[ModelResponseStream, None]: + """ + Passes the entire stream to the guardrail + + This is useful for guardrails that need to see the entire response, such as PII masking. + + See Aim guardrail implementation for an example - https://github.com/BerriAI/litellm/blob/d0e022cfacb8e9ebc5409bb652059b6fd97b45c0/litellm/proxy/guardrails/guardrail_hooks/aim.py#L168 + """ + async for item in response: + yield item ```