diff options
Diffstat (limited to 'g4f/__init__.py')
| -rw-r--r-- | g4f/__init__.py | 87 |
1 files changed, 52 insertions, 35 deletions
diff --git a/g4f/__init__.py b/g4f/__init__.py index 3b0fcad0..699dc238 100644 --- a/g4f/__init__.py +++ b/g4f/__init__.py @@ -4,49 +4,57 @@ import os from .errors import * from .models import Model, ModelUtils, _all_models -from .Provider import BaseProvider, AsyncGeneratorProvider, RetryProvider, ProviderUtils -from .typing import Messages, CreateResult, AsyncResult, Union, List -from . import debug +from .Provider import AsyncGeneratorProvider, ProviderUtils +from .typing import Messages, CreateResult, AsyncResult, Union +from . import debug, version +from .base_provider import BaseRetryProvider, ProviderType def get_model_and_provider(model : Union[Model, str], - provider : Union[type[BaseProvider], str, None], + provider : Union[ProviderType, str, None], stream : bool, - ignored : List[str] = None, + ignored : list[str] = None, ignore_working: bool = False, - ignore_stream: bool = False) -> tuple[Model, type[BaseProvider]]: + ignore_stream: bool = False) -> tuple[str, ProviderType]: if debug.version_check: debug.version_check = False - debug.check_pypi_version() - + version.utils.check_pypi_version() + if isinstance(provider, str): if provider in ProviderUtils.convert: provider = ProviderUtils.convert[provider] else: raise ProviderNotFoundError(f'Provider not found: {provider}') - if isinstance(model, str): - if model in ModelUtils.convert: - model = ModelUtils.convert[model] - else: - raise ModelNotFoundError(f'The model: {model} does not exist') - if not provider: + if isinstance(model, str): + if model in ModelUtils.convert: + model = ModelUtils.convert[model] + else: + raise ModelNotFoundError(f'Model not found: {model}') provider = model.best_provider - if isinstance(provider, RetryProvider) and ignored: - provider.providers = [p for p in provider.providers if p.__name__ not in ignored] - if not provider: raise ProviderNotFoundError(f'No provider found for model: {model}') - if not provider.working and not ignore_working: + if isinstance(model, Model): + model = model.name + + if ignored and isinstance(provider, BaseRetryProvider): + provider.providers = [p for p in provider.providers if p.__name__ not in ignored] + + if not ignore_working and not provider.working: raise ProviderNotWorkingError(f'{provider.__name__} is not working') if not ignore_stream and not provider.supports_stream and stream: raise StreamNotSupportedError(f'{provider.__name__} does not support "stream" argument') if debug.logging: - print(f'Using {provider.__name__} provider') + if model: + print(f'Using {provider.__name__} provider and {model} model') + else: + print(f'Using {provider.__name__} provider') + + debug.last_provider = provider return model, provider @@ -54,10 +62,10 @@ class ChatCompletion: @staticmethod def create(model : Union[Model, str], messages : Messages, - provider : Union[type[BaseProvider], str, None] = None, + provider : Union[ProviderType, str, None] = None, stream : bool = False, auth : Union[str, None] = None, - ignored : List[str] = None, + ignored : list[str] = None, ignore_working: bool = False, ignore_stream_and_auth: bool = False, **kwargs) -> Union[CreateResult, str]: @@ -75,32 +83,33 @@ class ChatCompletion: if proxy: kwargs['proxy'] = proxy - result = provider.create_completion(model.name, messages, stream, **kwargs) + result = provider.create_completion(model, messages, stream, **kwargs) return result if stream else ''.join(result) @staticmethod - async def create_async(model : Union[Model, str], - messages : Messages, - provider : Union[type[BaseProvider], str, None] = None, - stream : bool = False, - ignored : List[str] = None, - **kwargs) -> Union[AsyncResult, str]: + def create_async(model : Union[Model, str], + messages : Messages, + provider : Union[ProviderType, str, None] = None, + stream : bool = False, + ignored : list[str] = None, + **kwargs) -> Union[AsyncResult, str]: + model, provider = get_model_and_provider(model, provider, False, ignored) if stream: if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider): - return await provider.create_async_generator(model.name, messages, **kwargs) + return provider.create_async_generator(model, messages, **kwargs) raise StreamNotSupportedError(f'{provider.__name__} does not support "stream" argument in "create_async"') - return await provider.create_async(model.name, messages, **kwargs) + return provider.create_async(model, messages, **kwargs) class Completion: @staticmethod def create(model : Union[Model, str], prompt : str, - provider : Union[type[BaseProvider], None] = None, + provider : Union[ProviderType, None] = None, stream : bool = False, - ignored : List[str] = None, **kwargs) -> Union[CreateResult, str]: + ignored : list[str] = None, **kwargs) -> Union[CreateResult, str]: allowed_models = [ 'code-davinci-002', @@ -111,10 +120,18 @@ class Completion: 'text-davinci-003' ] if model not in allowed_models: - raise ModelNotAllowed(f'Can\'t use {model} with Completion.create()') + raise ModelNotAllowedError(f'Can\'t use {model} with Completion.create()') model, provider = get_model_and_provider(model, provider, stream, ignored) - result = provider.create_completion(model.name, [{"role": "user", "content": prompt}], stream, **kwargs) + result = provider.create_completion(model, [{"role": "user", "content": prompt}], stream, **kwargs) - return result if stream else ''.join(result)
\ No newline at end of file + return result if stream else ''.join(result) + +def get_last_provider(as_dict: bool = False) -> Union[ProviderType, dict[str, str]]: + last = debug.last_provider + if isinstance(last, BaseRetryProvider): + last = last.last_provider + if last and as_dict: + return {"name": last.__name__, "url": last.url} + return last
\ No newline at end of file |