diff options
Diffstat (limited to '')
| -rw-r--r-- | g4f/Provider/base_provider.py | 81 |
1 files changed, 40 insertions, 41 deletions
diff --git a/g4f/Provider/base_provider.py b/g4f/Provider/base_provider.py index 0f499c8c..79f8f617 100644 --- a/g4f/Provider/base_provider.py +++ b/g4f/Provider/base_provider.py @@ -35,30 +35,6 @@ class BaseProvider(ABC): ] param = ", ".join([": ".join(p) for p in params]) return f"g4f.provider.{cls.__name__} supports: ({param})" - - -_cookies = {} - -def get_cookies(cookie_domain: str) -> dict: - if cookie_domain not in _cookies: - _cookies[cookie_domain] = {} - try: - for cookie in browser_cookie3.load(cookie_domain): - _cookies[cookie_domain][cookie.name] = cookie.value - except: - pass - return _cookies[cookie_domain] - - -def format_prompt(messages: list[dict[str, str]], add_special_tokens=False): - if add_special_tokens or len(messages) > 1: - formatted = "\n".join( - ["%s: %s" % ((message["role"]).capitalize(), message["content"]) for message in messages] - ) - return f"{formatted}\nAssistant:" - else: - return messages.pop()["content"] - class AsyncProvider(BaseProvider): @@ -67,8 +43,9 @@ class AsyncProvider(BaseProvider): cls, model: str, messages: list[dict[str, str]], - stream: bool = False, **kwargs: Any) -> CreateResult: - + stream: bool = False, + **kwargs + ) -> CreateResult: yield asyncio.run(cls.create_async(model, messages, **kwargs)) @staticmethod @@ -90,7 +67,20 @@ class AsyncGeneratorProvider(AsyncProvider): stream: bool = True, **kwargs ) -> CreateResult: - yield from run_generator(cls.create_async_generator(model, messages, stream=stream, **kwargs)) + loop = asyncio.new_event_loop() + try: + asyncio.set_event_loop(loop) + generator = cls.create_async_generator(model, messages, stream=stream, **kwargs) + gen = generator.__aiter__() + while True: + try: + yield loop.run_until_complete(gen.__anext__()) + except StopAsyncIteration: + break + finally: + asyncio.set_event_loop(None) + loop.close() + @classmethod async def create_async( @@ -99,27 +89,36 @@ class AsyncGeneratorProvider(AsyncProvider): messages: list[dict[str, str]], **kwargs ) -> str: - chunks = [chunk async for chunk in cls.create_async_generator(model, messages, stream=False, **kwargs)] - if chunks: - return "".join(chunks) + return "".join([chunk async for chunk in cls.create_async_generator(model, messages, stream=False, **kwargs)]) @staticmethod @abstractmethod def create_async_generator( - model: str, - messages: list[dict[str, str]], - **kwargs - ) -> AsyncGenerator: + model: str, + messages: list[dict[str, str]], + **kwargs + ) -> AsyncGenerator: raise NotImplementedError() -def run_generator(generator: AsyncGenerator[Union[Any, str], Any]): - loop = asyncio.new_event_loop() - gen = generator.__aiter__() +_cookies = {} - while True: +def get_cookies(cookie_domain: str) -> dict: + if cookie_domain not in _cookies: + _cookies[cookie_domain] = {} try: - yield loop.run_until_complete(gen.__anext__()) + for cookie in browser_cookie3.load(cookie_domain): + _cookies[cookie_domain][cookie.name] = cookie.value + except: + pass + return _cookies[cookie_domain] - except StopAsyncIteration: - break + +def format_prompt(messages: list[dict[str, str]], add_special_tokens=False): + if add_special_tokens or len(messages) > 1: + formatted = "\n".join( + ["%s: %s" % ((message["role"]).capitalize(), message["content"]) for message in messages] + ) + return f"{formatted}\nAssistant:" + else: + return messages[0]["content"]
\ No newline at end of file |