fix: db not commit when streaming output (#1266)

This commit is contained in:
takatost 2023-10-01 16:41:52 +08:00 committed by GitHub
parent 8606d80c66
commit 86a9dea428
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 7 additions and 5 deletions

View File

@ -52,3 +52,4 @@ pandas==1.5.3
xinference==0.4.2 xinference==0.4.2
safetensors==0.3.2 safetensors==0.3.2
zhipuai==1.0.7 zhipuai==1.0.7
werkzeug==2.3.7

View File

@ -200,17 +200,14 @@ class CompletionService:
except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError, except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError,
ModelCurrentlyNotSupportError) as e: ModelCurrentlyNotSupportError) as e:
db.session.rollback()
PubHandler.pub_error(user, generate_task_id, e) PubHandler.pub_error(user, generate_task_id, e)
except LLMAuthorizationError: except LLMAuthorizationError:
db.session.rollback()
PubHandler.pub_error(user, generate_task_id, LLMAuthorizationError('Incorrect API key provided')) PubHandler.pub_error(user, generate_task_id, LLMAuthorizationError('Incorrect API key provided'))
except Exception as e: except Exception as e:
db.session.rollback()
logging.exception("Unknown Error in completion") logging.exception("Unknown Error in completion")
PubHandler.pub_error(user, generate_task_id, e) PubHandler.pub_error(user, generate_task_id, e)
finally: finally:
db.session.close() db.session.commit()
@classmethod @classmethod
def countdown_and_close(cls, worker_thread, pubsub, user, generate_task_id) -> threading.Thread: def countdown_and_close(cls, worker_thread, pubsub, user, generate_task_id) -> threading.Thread:
@ -388,6 +385,8 @@ class CompletionService:
logging.exception(e) logging.exception(e)
raise raise
finally: finally:
db.session.commit()
try: try:
pubsub.unsubscribe(generate_channel) pubsub.unsubscribe(generate_channel)
except ConnectionError: except ConnectionError:
@ -425,6 +424,8 @@ class CompletionService:
logging.exception(e) logging.exception(e)
raise raise
finally: finally:
db.session.commit()
try: try:
pubsub.unsubscribe(generate_channel) pubsub.unsubscribe(generate_channel)
except ConnectionError: except ConnectionError: