From f6ef3cb2237d8c336e915ef77ddbe6f37934c4fd Mon Sep 17 00:00:00 2001
From: MIDORIBIN <aquarion123@gmail.com>
Date: Fri, 28 Jul 2023 19:07:17 +0900
Subject: refactor: refactor provider

---
 g4f/__init__.py | 77 ++++++++++++++++++++++++++++++---------------------------
 1 file changed, 40 insertions(+), 37 deletions(-)

(limited to 'g4f/__init__.py')

diff --git a/g4f/__init__.py b/g4f/__init__.py
index 11920356..d90faf06 100644
--- a/g4f/__init__.py
+++ b/g4f/__init__.py
@@ -1,42 +1,45 @@
-import sys
-from . import Provider
-from g4f import models
+from . import models
+from .provider import BaseProvider
+from .typing import Any, CreateResult
 
 logging = False
 
+
 class ChatCompletion:
     @staticmethod
-    def create(model: models.Model or str, messages: list, provider: Provider.Provider = None, stream: bool = False, auth: str = False, **kwargs):
-        kwargs['auth'] = auth
-        if provider and provider.working == False:
-            return f'{provider.__name__} is not working'
-
-        if provider and provider.needs_auth and not auth:
-            print(
-                f'ValueError: {provider.__name__} requires authentication (use auth="cookie or token or jwt ..." param)', file=sys.stderr)
-            sys.exit(1)
-
-        try:
-            if isinstance(model, str):
-                try:
-                    model = models.ModelUtils.convert[model]
-                except KeyError:
-                    raise Exception(f'The model: {model} does not exist')
-
-            engine = model.best_provider if not provider else provider
-
-            if not engine.supports_stream and stream == True:
-                print(
-                    f"ValueError: {engine.__name__} does not support 'stream' argument", file=sys.stderr)
-                sys.exit(1)
-
-            if logging: print(f'Using {engine.__name__} provider')
-
-            return (engine._create_completion(model.name, messages, stream, **kwargs)
-                    if stream else ''.join(engine._create_completion(model.name, messages, stream, **kwargs)))
-        except TypeError as e:
-            print(e)
-            arg: str = str(e).split("'")[1]
-            print(
-                f"ValueError: {engine.__name__} does not support '{arg}' argument", file=sys.stderr)
-            sys.exit(1)
\ No newline at end of file
+    def create(
+        model: models.Model | str,
+        messages: list[dict[str, str]],
+        provider: type[BaseProvider] | None = None,
+        stream: bool = False,
+        auth: str | None = None,
+        **kwargs: Any,
+    ) -> CreateResult | str:
+        if isinstance(model, str):
+            try:
+                model = models.ModelUtils.convert[model]
+            except KeyError:
+                raise Exception(f"The model: {model} does not exist")
+
+        provider = model.best_provider if provider == None else provider
+
+        if not provider.working:
+            raise Exception(f"{provider.__name__} is not working")
+
+        if provider.needs_auth and not auth:
+            raise Exception(
+                f'ValueError: {provider.__name__} requires authentication (use auth="cookie or token or jwt ..." param)'
+            )
+        if provider.needs_auth:
+            kwargs["auth"] = auth
+
+        if not provider.supports_stream and stream:
+            raise Exception(
+                f"ValueError: {provider.__name__} does not support 'stream' argument"
+            )
+
+        if logging:
+            print(f"Using {provider.__name__} provider")
+
+        result = provider.create_completion(model.name, messages, stream, **kwargs)
+        return result if stream else "".join(result)
-- 
cgit v1.2.3