diff options
Diffstat (limited to '')
-rw-r--r-- | g4f/Provider/base_provider.py | 41 |
1 files changed, 23 insertions, 18 deletions
diff --git a/g4f/Provider/base_provider.py b/g4f/Provider/base_provider.py index 95f1b0b2..bc47a1fa 100644 --- a/g4f/Provider/base_provider.py +++ b/g4f/Provider/base_provider.py @@ -5,8 +5,8 @@ from asyncio import AbstractEventLoop from concurrent.futures import ThreadPoolExecutor from abc import abstractmethod from inspect import signature, Parameter -from .helper import get_event_loop, get_cookies, format_prompt -from ..typing import CreateResult, AsyncResult, Messages +from .helper import get_cookies, format_prompt +from ..typing import CreateResult, AsyncResult, Messages, Union from ..base_provider import BaseProvider from ..errors import NestAsyncioError @@ -20,6 +20,17 @@ if sys.platform == 'win32': if isinstance(asyncio.get_event_loop_policy(), asyncio.WindowsProactorEventLoopPolicy): asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) +def get_running_loop() -> Union[AbstractEventLoop, None]: + try: + loop = asyncio.get_running_loop() + if not hasattr(loop.__class__, "_nest_patched"): + raise NestAsyncioError( + 'Use "create_async" instead of "create" function in a running event loop. Or use "nest_asyncio" package.' + ) + return loop + except RuntimeError: + pass + class AbstractProvider(BaseProvider): """ Abstract class for providing asynchronous functionality to derived classes. @@ -56,7 +67,7 @@ class AbstractProvider(BaseProvider): return await asyncio.wait_for( loop.run_in_executor(executor, create_func), - timeout=kwargs.get("timeout", 0) + timeout=kwargs.get("timeout") ) @classmethod @@ -118,14 +129,7 @@ class AsyncProvider(AbstractProvider): Returns: CreateResult: The result of the completion creation. """ - try: - loop = asyncio.get_running_loop() - if not hasattr(loop.__class__, "_nest_patched"): - raise NestAsyncioError( - 'Use "create_async" instead of "create" function in a running event loop. Or use "nest_asyncio" package.' - ) - except RuntimeError: - pass + get_running_loop() yield asyncio.run(cls.create_async(model, messages, **kwargs)) @staticmethod @@ -180,15 +184,12 @@ class AsyncGeneratorProvider(AsyncProvider): Returns: CreateResult: The result of the streaming completion creation. """ - try: - loop = asyncio.get_running_loop() - if not hasattr(loop.__class__, "_nest_patched"): - raise NestAsyncioError( - 'Use "create_async" instead of "create" function in a running event loop. Or use "nest_asyncio" package.' - ) - except RuntimeError: + loop = get_running_loop() + new_loop = False + if not loop: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) + new_loop = True generator = cls.create_async_generator(model, messages, stream=stream, **kwargs) gen = generator.__aiter__() @@ -199,6 +200,10 @@ class AsyncGeneratorProvider(AsyncProvider): except StopAsyncIteration: break + if new_loop: + loop.close() + asyncio.set_event_loop(None) + @classmethod async def create_async( cls, |