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

api/controller/task/ChatCompleteTask.py

@@ -50,6 +50,7 @@ class ChatCompleteTask:
         self.chat_complete_service = ChatCompleteService(self.dbs, self.page_title)
         self.chat_complete = await self.chat_complete_service.__aenter__()
         
+        try:
             if await self.chat_complete.page_index_exists():
                 question_tokens = await self.tiktoken.get_tokens(question)
 
@@ -68,8 +69,10 @@ class ChatCompleteTask:
                 
                 return res
             else:
-            await self._exit()
                 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,6 +240,7 @@ class ChatCompleteService:
         message_log = []
         if self.conversation_chunk is not None:
             for message in self.conversation_chunk.message_data:
+                if message["role"] in ["user", "assistant"]:
                     message_log.append(
                         {
                             "role": message["role"],

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