From f054219e50c65ee1ed7222ace5e82f0229888165 Mon Sep 17 00:00:00 2001 From: Juan Manuel Ciro Torres Date: Mon, 7 Aug 2023 12:15:19 -0500 Subject: [PATCH] new sse implementation --- backend/app/api/endpoints/base/context.py | 21 ++++++--------------- backend/app/domain/services/base/context.py | 2 +- 2 files changed, 7 insertions(+), 16 deletions(-) diff --git a/backend/app/api/endpoints/base/context.py b/backend/app/api/endpoints/base/context.py index 7bd03f4f9..93fd8a76b 100644 --- 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 index cb1f8520c..e7b3a51dd 100644 --- 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"],