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

@ -51,4 +51,5 @@ stripe~=5.5.0
pandas==1.5.3
xinference==0.4.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,
LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError,
ModelCurrentlyNotSupportError) as e:
db.session.rollback()
PubHandler.pub_error(user, generate_task_id, e)
except LLMAuthorizationError:
db.session.rollback()
PubHandler.pub_error(user, generate_task_id, LLMAuthorizationError('Incorrect API key provided'))
except Exception as e:
db.session.rollback()
logging.exception("Unknown Error in completion")
PubHandler.pub_error(user, generate_task_id, e)
finally:
db.session.close()
db.session.commit()
@classmethod
def countdown_and_close(cls, worker_thread, pubsub, user, generate_task_id) -> threading.Thread:
@ -388,6 +385,8 @@ class CompletionService:
logging.exception(e)
raise
finally:
db.session.commit()
try:
pubsub.unsubscribe(generate_channel)
except ConnectionError:
@ -425,6 +424,8 @@ class CompletionService:
logging.exception(e)
raise
finally:
db.session.commit()
try:
pubsub.unsubscribe(generate_channel)
except ConnectionError: