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 5649184
Showing 1 changed file with 11 additions and 4 deletions.
15 changes: 11 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,16 @@ 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_cancellation_payload(message="The task was cancelled.", code=499):
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": 499, "data": aiconfig_json}})

def handle_cancellation():
yield "["
Expand Down Expand Up @@ -307,7 +311,10 @@ def kill_thread(thread_id: int | None):
yield from handle_cancellation()
return

if isinstance(text, str):
if isinstance(text, Exception):
yield from create_cancellation_payload(message=f"Exception: {text}", code=499)
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 5649184

Please sign in to comment.