diff options
Diffstat (limited to 'g4f/Provider/base_provider.py')
-rw-r--r-- | g4f/Provider/base_provider.py | 107 |
1 files changed, 98 insertions, 9 deletions
diff --git a/g4f/Provider/base_provider.py b/g4f/Provider/base_provider.py index 98ad3514..e667819a 100644 --- a/g4f/Provider/base_provider.py +++ b/g4f/Provider/base_provider.py @@ -1,24 +1,28 @@ +from __future__ import annotations + +import asyncio from abc import ABC, abstractmethod -from ..typing import Any, CreateResult +import browser_cookie3 + +from ..typing import Any, AsyncGenerator, CreateResult, Union class BaseProvider(ABC): url: str - working = False - needs_auth = False - supports_stream = False + working = False + needs_auth = False + supports_stream = False supports_gpt_35_turbo = False - supports_gpt_4 = False + supports_gpt_4 = False @staticmethod @abstractmethod def create_completion( model: str, messages: list[dict[str, str]], - stream: bool, - **kwargs: Any, - ) -> CreateResult: + stream: bool, **kwargs: Any) -> CreateResult: + raise NotImplementedError() @classmethod @@ -30,4 +34,89 @@ class BaseProvider(ABC): ("stream", "bool"), ] param = ", ".join([": ".join(p) for p in params]) - return f"g4f.provider.{cls.__name__} supports: ({param})"
\ No newline at end of file + 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] = {} + + for cookie in browser_cookie3.load(cookie_domain): + _cookies[cookie_domain][cookie.name] = cookie.value + + 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): + @classmethod + def create_completion( + cls, + model: str, + messages: list[dict[str, str]], + stream: bool = False, **kwargs: Any) -> CreateResult: + + yield asyncio.run(cls.create_async(model, messages, **kwargs)) + + @staticmethod + @abstractmethod + async def create_async( + model: str, + messages: list[dict[str, str]], **kwargs: Any) -> str: + raise NotImplementedError() + + +class AsyncGeneratorProvider(AsyncProvider): + @classmethod + def create_completion( + cls, + model: str, + messages: list[dict[str, str]], + stream: bool = True, + **kwargs + ) -> CreateResult: + yield from run_generator(cls.create_async_generator(model, messages, stream=stream, **kwargs)) + + @classmethod + async def create_async( + cls, + model: str, + 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) + + @staticmethod + @abstractmethod + def create_async_generator( + 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__() + + while True: + try: + yield loop.run_until_complete(gen.__anext__()) + + except StopAsyncIteration: + break |