fix: db session not commit before long llm call running (#1251)

This commit is contained in:
takatost 2023-09-27 21:40:26 +08:00 committed by GitHub
parent d38eac959b
commit a31466d34e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,6 +1,7 @@
import json import json
import os import os
import re import re
import time
from abc import abstractmethod from abc import abstractmethod
from typing import List, Optional, Any, Union, Tuple from typing import List, Optional, Any, Union, Tuple
import decimal import decimal
@ -20,6 +21,8 @@ from core.prompt.prompt_template import JinjaPromptTemplate
from core.third_party.langchain.llms.fake import FakeLLM from core.third_party.langchain.llms.fake import FakeLLM
import logging import logging
from extensions.ext_database import db
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -129,6 +132,8 @@ class BaseLLM(BaseProviderModel):
if self.deduct_quota: if self.deduct_quota:
self.model_provider.check_quota_over_limit() self.model_provider.check_quota_over_limit()
db.session.commit()
if not callbacks: if not callbacks:
callbacks = self.callbacks callbacks = self.callbacks
else: else: