From b6f98e408f7a70290ae8a5e6a9bc633fabfe12fa Mon Sep 17 00:00:00 2001
From: jdhuang <jdhuang@amazon.com>
Date: Tue, 2 Apr 2024 20:14:37 +0800
Subject: [PATCH] Add sync iterator

---
 litellm/llms/bedrock.py | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py
index ae78d3ec48..8d5669b388 100644
--- a/litellm/llms/bedrock.py
+++ b/litellm/llms/bedrock.py
@@ -1008,7 +1008,7 @@ def completion(
                         )
                         streaming_choice.delta = delta_obj
                         streaming_model_response.choices = [streaming_choice]
-                        completion_stream = model_response_iterator(
+                        completion_stream = ModelResponseIterator(
                             model_response=streaming_model_response
                         )
                         print_verbose(
@@ -1108,10 +1108,30 @@ def completion(
 
             raise BedrockError(status_code=500, message=traceback.format_exc())
 
+class ModelResponseIterator:
+    def __init__(self, model_response):
+        self.model_response = model_response
+        self.is_done = False
 
-async def model_response_iterator(model_response):
-    yield model_response
+    # Sync iterator
+    def __iter__(self):
+        return self
 
+    def __next__(self):
+        if self.is_done:
+            raise StopIteration
+        self.is_done = True
+        return self.model_response
+
+    # Async iterator
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self.is_done:
+            raise StopAsyncIteration
+        self.is_done = True
+        return self.model_response
 
 def _embedding_func_single(
     model: str,