Merge remote-tracking branch 'origin/litellm_completions_stream_post_call_rewrite' into litellm_completions_stream_guardrail_scan

litellm/llms/openai/completion/guardrail_translation/handler.py

@@ -5,18 +5,26 @@ This module provides guardrail translation support for OpenAI's text completion
 The handler processes the 'prompt' parameter for guardrails.
 """
 
-from collections.abc import Mapping
+import json
+from collections.abc import Mapping, Sequence
+from itertools import chain
 from typing import TYPE_CHECKING, Final
 
 from litellm._logging import verbose_proxy_logger
-from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
-from litellm.types.utils import GenericGuardrailAPIInputs
+from litellm.llms.base_llm.guardrail_translation.base_translation import (
+    BaseTranslation,
+    StreamingScanKey,
+    StreamTransformSink,
+)
+from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items
+from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse
 
 if TYPE_CHECKING:
+    from fastapi import HTTPException
+
     from litellm.integrations.custom_guardrail import CustomGuardrail
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
     from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
-    from litellm.types.utils import TextCompletionResponse
 
 
 class OpenAITextCompletionHandler(BaseTranslation):
@@ -32,6 +40,9 @@ class OpenAITextCompletionHandler(BaseTranslation):
     - A list of strings (for batch completions)
     """
 
+    delivers_ended_stream_rewrites = True
+    assembles_streamed_response = True
+
     async def process_input_messages(
         self,
         data: dict,
@@ -125,7 +136,7 @@ class OpenAITextCompletionHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
         user_api_key_dict: "UserAPIKeyAuth | None" = None,
-        request_data: dict | None = None,
+        request_data: dict[str, object] | None = None,
     ) -> "TextCompletionResponse":
         """
         Process output response by applying guardrails to completion text.
@@ -162,11 +173,7 @@ class OpenAITextCompletionHandler(BaseTranslation):
                 if "response" not in request_data:
                     request_data["response"] = response
 
-            # Add user API key metadata with prefixed keys
-            if "litellm_metadata" not in request_data:
-                user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
-                if user_metadata:
-                    request_data["litellm_metadata"] = user_metadata
+            self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
 
             inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
             # Include model information from the response if available
@@ -195,3 +202,101 @@ class OpenAITextCompletionHandler(BaseTranslation):
                     )
 
         return response
+
+    async def process_output_streaming_response(
+        self,
+        responses_so_far: list[TextCompletionResponse],  # mutable-ok: rewrites the caller's buffered chunks in place
+        guardrail_to_apply: "CustomGuardrail",
+        litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
+        user_api_key_dict: "UserAPIKeyAuth | None" = None,
+        request_data: dict[str, object] | None = None,  # mutable-ok: the hooks share and mutate the request payload
+        stream_transform_sink: StreamTransformSink | None = None,
+        deliver_ended_stream_rewrites: bool = False,
+    ) -> Sequence[TextCompletionResponse]:
+        if not self._check_streaming_has_ended(responses_so_far):
+            return responses_so_far
+        assembled: Final = self._assemble_ended_stream(responses_so_far)
+        pre_guardrail_texts: Final = _choice_texts(assembled)
+        await self.process_output_response(
+            response=assembled,
+            guardrail_to_apply=guardrail_to_apply,
+            litellm_logging_obj=litellm_logging_obj,
+            user_api_key_dict=user_api_key_dict,
+            request_data=request_data,
+        )
+        if not deliver_ended_stream_rewrites:
+            return responses_so_far
+        rewrites: Final = tuple(
+            after for before, after in zip(pre_guardrail_texts, _choice_texts(assembled)) if after != before
+        )
+        for choice_index, text in rewrites:
+            self._write_choice_text(responses_so_far, choice_index, text)
+        return responses_so_far
+
+    def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
+        chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, TextCompletionResponse))
+        if not self._check_streaming_has_ended(chunks):
+            return None
+        assembled: Final = self._assemble_ended_stream(chunks)
+        return StreamingScanKey(texts=tuple(text for _, text in _choice_texts(assembled)), stream_ended=True)
+
+    def build_stream_error_items(
+        self,
+        exc: "HTTPException",
+        responses_so_far: Sequence[object] | None = None,
+    ) -> Sequence[bytes] | None:
+        from litellm.proxy.common_request_processing import sse_error_payload
+
+        _, error_obj = sse_error_payload(exc)
+        return (f'data: {{"error": {json.dumps(error_obj)}}}\n\n'.encode(),)
+
+    def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
+        streamed_choices: Final = chain.from_iterable(stream_item_items(item, "choices") for item in responses_so_far)
+        return any(stream_item_field(choice, "finish_reason") is not None for choice in streamed_choices)
+
+    @staticmethod
+    def _assemble_ended_stream(responses_so_far: Sequence[TextCompletionResponse]) -> TextCompletionResponse:
+        streamed_choices: Final = tuple(chain.from_iterable(item.choices for item in responses_so_far))
+        choice_indices: Final = tuple(dict.fromkeys(choice.index for choice in streamed_choices))
+        assembled_choices: Final = [
+            TextChoices(
+                index=index,
+                text="".join(_choice_text(choice) for choice in streamed_choices if choice.index == index),
+                finish_reason=next(
+                    (
+                        choice.finish_reason
+                        for choice in streamed_choices
+                        if choice.index == index and choice.finish_reason
+                    ),
+                    None,
+                ),
+            )
+            for index in choice_indices
+        ]
+        first: Final = responses_so_far[0]
+        return TextCompletionResponse(
+            id=first.id,
+            created=first.created,
+            model=first.model,
+            object="text_completion",
+            choices=assembled_choices,
+            usage=next((item.usage for item in reversed(responses_so_far) if item.usage is not None), None),
+        )
+
+    @staticmethod
+    def _write_choice_text(responses_so_far: Sequence[TextCompletionResponse], choice_index: int, text: str) -> None:
+        streamed_choices: Final = chain.from_iterable(item.choices for item in responses_so_far)
+        head, *tail = (choice for choice in streamed_choices if choice.index == choice_index)
+        head.text = text
+        for choice in tail:
+            if isinstance(stream_item_field(choice, "text"), str):
+                choice.text = ""
+
+
+def _choice_text(choice: TextChoices) -> str:
+    text: Final = stream_item_field(choice, "text")
+    return text if isinstance(text, str) else ""
+
+
+def _choice_texts(response: TextCompletionResponse) -> tuple[tuple[int, str], ...]:
+    return tuple((choice.index, _choice_text(choice)) for choice in response.choices)