尝试修复资源未关闭的问题

api/controller/task/ChatCompleteTask.py

@@ -50,26 +50,29 @@ class ChatCompleteTask:
         self.chat_complete_service = ChatCompleteService(self.dbs, self.page_title)
         self.chat_complete = await self.chat_complete_service.__aenter__()
         
-        if await self.chat_complete.page_index_exists():
-            question_tokens = await self.tiktoken.get_tokens(question)
-
-            extract_limit = embedding_search["limit"] or 10
-
-            self.transatcion_id: Optional[str] = None
-            self.point_cost: int = 0
-            if not self.is_system:
-                usage_res = await self.mwapi.ai_toolbox_start_transaction(self.user_id, "chatcomplete",
-                                                                             question_tokens, extract_limit)
-                self.transatcion_id = usage_res["transaction_id"]
-                self.point_cost = usage_res["point_cost"]
-
-            res = await self.chat_complete.prepare_chat_complete(question, conversation_id=conversation_id,
-                            user_id=self.user_id, edit_message_id=edit_message_id, embedding_search=embedding_search)
-            
-            return res
-        else:
-            await self._exit()
-            raise MediaWikiPageNotFoundException("Page %s not found." % self.page_title)
+        try:
+            if await self.chat_complete.page_index_exists():
+                question_tokens = await self.tiktoken.get_tokens(question)
+
+                extract_limit = embedding_search["limit"] or 10
+
+                self.transatcion_id: Optional[str] = None
+                self.point_cost: int = 0
+                if not self.is_system:
+                    usage_res = await self.mwapi.ai_toolbox_start_transaction(self.user_id, "chatcomplete",
+                                                                                question_tokens, extract_limit)
+                    self.transatcion_id = usage_res["transaction_id"]
+                    self.point_cost = usage_res["point_cost"]
+
+                res = await self.chat_complete.prepare_chat_complete(question, conversation_id=conversation_id,
+                                user_id=self.user_id, edit_message_id=edit_message_id, embedding_search=embedding_search)
+                
+                return res
+            else:
+                raise MediaWikiPageNotFoundException("Page %s not found." % self.page_title)
+        except Exception as e:
+            await self.end()
+            raise e
 
     async def _on_message(self, delta_message: str):
         self.chunks.append(delta_message)
@@ -122,9 +125,9 @@ class ChatCompleteTask:
 
             await self._on_error(e)
         finally:
-            await self._exit()
+            await self.end()
     
-    async def _exit(self):
+    async def end(self):
         await self.chat_complete_service.__aexit__(None, None, None)
         del chat_complete_tasks[self.task_id]
         self.is_finished = True

api/model/base.py

@@ -24,7 +24,6 @@ class BaseHelper:
 
     async def __aexit__(self, exc_type, exc, tb):
         await self.session.__aexit__(exc_type, exc, tb)
-        pass
 
 T = TypeVar("T", bound=BaseModel)
 

service/chat_complete.py

@@ -240,12 +240,13 @@ class ChatCompleteService:
         message_log = []
         if self.conversation_chunk is not None:
             for message in self.conversation_chunk.message_data:
-                message_log.append(
-                    {
-                        "role": message["role"],
-                        "content": message["content"],
-                    }
-                )
+                if message["role"] in ["user", "assistant"]:
+                    message_log.append(
+                        {
+                            "role": message["role"],
+                            "content": message["content"],
+                        }
+                    )
 
         if self.extract_doc is not None:
             doc_prompt_content = "\n".join(

service/database.py

@@ -39,8 +39,7 @@ class DatabaseService:
     
     async def init(self):
         loop = local.loop
-        self.pool = asyncpg.create_pool(**config.DATABASE, loop=loop)
-        await self.pool.__aenter__()
+        self.pool = await asyncpg.create_pool(**config.DATABASE, loop=loop)
 
         engine = create_async_engine(get_dsn(), echo=config.DEBUG)
         self.engine = engine