new sse implementation

mlcommons · Aug 7, 2023 · f054219 · f054219
1 parent f8af27b
commit f054219
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 16 deletions.
diff --git a/backend/app/api/endpoints/base/context.py b/backend/app/api/endpoints/base/context.py
@@ -53,21 +53,12 @@ async def get_generative_contexts(model: GetGenerativeContextRequest):
 @router.post("/stream")
 async def stream_images(model_info: GetGenerativeContextRequest):
     async def event_generator():
-        try:
-            for _ in range(3):
-                data = ContextService().get_generative_contexts(
-                    model_info.type, model_info.artifacts
-                )
+        for _ in range(3):
+            data = await ContextService().get_generative_contexts(
+                model_info.type, model_info.artifacts
+            )
+            if data:
                 yield json.dumps(data)
-                await asyncio.sleep(1)
-
-        except asyncio.CancelledError as e:
-            print("CancelledError", e)
-            raise
-        except Exception as e:
-            print("Error:", e)
-            raise
-        finally:
-            print("finally")
+            await asyncio.sleep(1)
 
     return EventSourceResponse(event_generator())
diff --git a/backend/app/domain/services/base/context.py b/backend/app/domain/services/base/context.py
@@ -172,7 +172,7 @@ def get_perdi_contexts(
     async def generate_images_stream(self, model_info):
         yield self.get_generative_contexts(model_info["type"], model_info["artifacts"])
 
-    def get_generative_contexts(self, type: str, artifacts: dict) -> dict:
+    async def get_generative_contexts(self, type: str, artifacts: dict) -> dict:
         if type == "nibbler":
             return self.get_nibbler_contexts(
                 prompt=artifacts["prompt"],