Browse Source

fix: request timeout when openai completion (#1265)

takatost 1 năm trước cách đây
mục cha
commit
8606d80c66

+ 8 - 1
api/core/model_providers/models/llm/openai_model.py

@@ -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)
-        return self._client.generate([prompts], stop, callbacks)
+
+        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:
         """

+ 2 - 0
api/services/completion_service.py

@@ -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: