Преглед изворни кода

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

takatost пре 1 година
родитељ
комит
86a9dea428
2 измењених фајлова са 7 додато и 5 уклоњено
  1. 2 1
      api/requirements.txt
  2. 5 4
      api/services/completion_service.py

+ 2 - 1
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
+zhipuai==1.0.7
+werkzeug==2.3.7

+ 5 - 4
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: