Skip to content

Commit

Permalink
[bug fix][AIC-py] sort of do exception handling during streaming (#860)
Browse files Browse the repository at this point in the history
[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.

<img width="396" alt="image"
src="https://github.com/lastmile-ai/aiconfig/assets/148090348/5e5f603f-4ac3-405e-aead-a3dbfb78d163">
  • Loading branch information
jonathanlastmileai authored Jan 10, 2024
2 parents 08f7086 + 44fbded commit a036254
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 a036254

Please sign in to comment.