From 79ea466cf5bf5b1c0c97cd0e5d33d5143a54198e Mon Sep 17 00:00:00 2001
From: zeeland <zeeland4work@gmail.com>
Date: Wed, 13 Dec 2023 22:57:19 +0800
Subject: [PATCH] refactor: add CustomStreamWrapper return type for completion

---
 litellm/main.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/litellm/main.py b/litellm/main.py
index bf35d2a27d..8848f40f92 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -284,7 +284,7 @@ def completion(
 
     # Optional liteLLM function params
     **kwargs,
-) -> ModelResponse:
+) -> Union[ModelResponse, CustomStreamWrapper]:
     """
     Perform a completion() using any of litellm supported llms (example gpt-4, gpt-3.5-turbo, claude-2, command-nightly)
     Parameters: