Skip to content

Commit

Permalink
[bug fix][AIC-py] sort of do exception handling during streaming
Browse files Browse the repository at this point in the history
This piggybacks on the cancellation during stream flow, which works as far as I've tested.
The frontend gets the message.

<img width="396" alt="image" src="https://github.com/lastmile-ai/aiconfig/assets/148090348/5e5f603f-4ac3-405e-aead-a3dbfb78d163">
  • Loading branch information
jonathanlastmileai committed Jan 10, 2024
1 parent bcd2921 commit 44fbded
Showing 1 changed file with 14 additions and 4 deletions.
18 changes: 14 additions & 4 deletions python/src/aiconfig/editor/server/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 "["
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 44fbded

Please sign in to comment.