diff --git a/api/responder.py b/api/responder.py index d235cc4..79431d5 100644 --- a/api/responder.py +++ b/api/responder.py @@ -7,6 +7,7 @@ import aiohttp import asyncio import starlette +from typing import Any, Coroutine, Set from rich import print from dotenv import load_dotenv @@ -23,6 +24,15 @@ CRITICAL_API_ERRORS = ['invalid_api_key', 'account_deactivated'] keymanager = providerkeys.manager +background_tasks: Set[asyncio.Task[Any]] = set() + + +def create_background_task(coro: Coroutine[Any, Any, Any]) -> None: + task = asyncio.create_task(coro) + background_tasks.add(task) + task.add_done_callback(background_tasks.discard) + + async def respond( path: str='/v1/chat/completions', user: dict=None, @@ -166,7 +176,7 @@ async def respond( if (not is_stream) and server_json_response: yield json.dumps(server_json_response) - asyncio.create_task( + create_background_task( after_request.after_request( incoming_request=incoming_request, target_request=target_request,