fix(gemini.py): support streaming

2025-04-25 10:44:24 +00:00 · 2024-01-19 19:26:23 -08:00 · 2024-01-19 19:26:23 -08:00 · b07677c6be
commit b07677c6be
parent b59e67f099
4 changed files with 67 additions and 13 deletions
--- a/litellm/main.py
+++ b/litellm/main.py
@ -1382,6 +1382,18 @@ def completion(
                acompletion=acompletion,
                custom_prompt_dict=custom_prompt_dict,
            )
+            if (
+                "stream" in optional_params
+                and optional_params["stream"] == True
+                and acompletion == False
+            ):
+                response = CustomStreamWrapper(
+                    iter(model_response),
+                    model,
+                    custom_llm_provider="gemini",
+                    logging_obj=logging,
+                )
+                return response
            response = model_response
        elif custom_llm_provider == "vertex_ai":
            vertex_ai_project = litellm.vertex_project or get_secret("VERTEXAI_PROJECT")