From 44fbded312c090bd7f20c6f8d92bd0d906288dda Mon Sep 17 00:00:00 2001 From: Jonathan Lessinger Date: Wed, 10 Jan 2024 14:01:40 -0500 Subject: [PATCH] [bug fix][AIC-py] sort of do exception handling during streaming This piggybacks on the cancellation during stream flow, which works as far as I've tested. The frontend gets the message. image --- python/src/aiconfig/editor/server/server.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/python/src/aiconfig/editor/server/server.py b/python/src/aiconfig/editor/server/server.py index 9c257804a..fc091cd51 100644 --- a/python/src/aiconfig/editor/server/server.py +++ b/python/src/aiconfig/editor/server/server.py @@ -234,12 +234,19 @@ def generate(cancellation_token_id: str): # type: ignore # Use multi-threading so that we don't block run command from # displaying the streamed output (if streaming is supported) def run_async_config_in_thread(): - asyncio.run(aiconfig.run(prompt_name=prompt_name, params=params, run_with_dependencies=False, options=inference_options)) # type: ignore + try: + asyncio.run(aiconfig.run(prompt_name=prompt_name, params=params, run_with_dependencies=False, options=inference_options)) # type: ignore + except Exception as e: + output_text_queue.put(e) + output_text_queue.put(STOP_STREAMING_SIGNAL) # type: ignore - def create_cancellation_payload(): + def create_error_payload(message: str, code: int): aiconfig_json = aiconfig_deep_copy.model_dump(exclude=EXCLUDE_OPTIONS) if aiconfig_deep_copy is not None else None - return json.dumps({"error": {"message": "The task was cancelled.", "code": 499, "data": aiconfig_json}}) + return json.dumps({"error": {"message": message, "code": code, "data": aiconfig_json}}) + + def create_cancellation_payload(): + return create_error_payload(message="The task was cancelled.", code=499) def handle_cancellation(): yield "[" @@ -307,7 +314,10 @@ def kill_thread(thread_id: int | None): yield from handle_cancellation() return - if isinstance(text, str): + if isinstance(text, Exception): + yield from create_error_payload(message=f"Exception: {text}", code=500) + return + elif isinstance(text, str): accumulated_output_text += text elif isinstance(text, dict) and "content" in text: # TODO: Fix streaming output format so that it returns text