diff --git a/api/requirements.txt b/api/requirements.txt index d616d50972..9901aee85b 100644 --- a/api/requirements.txt +++ b/api/requirements.txt @@ -51,4 +51,5 @@ stripe~=5.5.0 pandas==1.5.3 xinference==0.4.2 safetensors==0.3.2 -zhipuai==1.0.7 \ No newline at end of file +zhipuai==1.0.7 +werkzeug==2.3.7 \ No newline at end of file diff --git a/api/services/completion_service.py b/api/services/completion_service.py index 00034b3f3d..252c85e93f 100644 --- a/api/services/completion_service.py +++ b/api/services/completion_service.py @@ -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: