fix: request timeout when openai completion (#1265)

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

View File

@ -5,6 +5,7 @@ from typing import List, Optional, Any
import openai
from langchain.callbacks.manager import Callbacks
from langchain.schema import LLMResult
from openai import api_requestor
from core.model_providers.providers.base import BaseModelProvider
from core.third_party.langchain.llms.chat_open_ai import EnhanceChatOpenAI
@ -105,7 +106,13 @@ class OpenAIModel(BaseLLM):
raise ModelCurrentlyNotSupportError("Dify Hosted OpenAI GPT-4 currently not support.")
prompts = self._get_prompt_from_messages(messages)
try:
return self._client.generate([prompts], stop, callbacks)
finally:
thread_context = api_requestor._thread_context
if hasattr(thread_context, "session") and thread_context.session:
thread_context.session.close()
def get_num_tokens(self, messages: List[PromptMessage]) -> int:
"""

View File

@ -209,6 +209,8 @@ class CompletionService:
db.session.rollback()
logging.exception("Unknown Error in completion")
PubHandler.pub_error(user, generate_task_id, e)
finally:
db.session.close()
@classmethod
def countdown_and_close(cls, worker_thread, pubsub, user, generate_task_id) -> threading.Thread: