From 24c5502d7601ba3f2c5ab459015c74002b3af5c6 Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <hlohaus@users.noreply.github.com>
Date: Wed, 15 May 2024 21:07:49 +0200
Subject: Add needs auth to provierds, Add PerplexityApi provider Add proxy
 support to nodriver

---
 g4f/Provider/needs_auth/Gemini.py        |  9 ++++++---
 g4f/Provider/needs_auth/OpenaiChat.py    |  9 ++++++---
 g4f/Provider/needs_auth/PerplexityApi.py | 31 +++++++++++++++++++++++++++++++
 g4f/Provider/needs_auth/__init__.py      |  3 ++-
 4 files changed, 45 insertions(+), 7 deletions(-)
 create mode 100644 g4f/Provider/needs_auth/PerplexityApi.py

(limited to 'g4f/Provider/needs_auth')

diff --git a/g4f/Provider/needs_auth/Gemini.py b/g4f/Provider/needs_auth/Gemini.py
index 25ad1c6e..e468f64a 100644
--- a/g4f/Provider/needs_auth/Gemini.py
+++ b/g4f/Provider/needs_auth/Gemini.py
@@ -59,7 +59,7 @@ class Gemini(AsyncGeneratorProvider):
     _cookies: Cookies = None
 
     @classmethod
-    async def nodriver_login(cls) -> AsyncIterator[str]:
+    async def nodriver_login(cls, proxy: str = None) -> AsyncIterator[str]:
         try:
             import nodriver as uc
         except ImportError:
@@ -71,7 +71,10 @@ class Gemini(AsyncGeneratorProvider):
             user_data_dir = None
         if debug.logging:
             print(f"Open nodriver with user_dir: {user_data_dir}")
-        browser = await uc.start(user_data_dir=user_data_dir)
+        browser = await uc.start(
+            user_data_dir=user_data_dir,
+            browser_args=None if proxy is None else [f"--proxy-server={proxy}"],
+        )
         login_url = os.environ.get("G4F_LOGIN_URL")
         if login_url:
             yield f"Please login: [Google Gemini]({login_url})\n\n"
@@ -134,7 +137,7 @@ class Gemini(AsyncGeneratorProvider):
         ) as session:
             snlm0e  = await cls.fetch_snlm0e(session, cls._cookies) if cls._cookies else None
             if not snlm0e:
-                async for chunk in cls.nodriver_login():
+                async for chunk in cls.nodriver_login(proxy):
                     yield chunk
                 if cls._cookies is None:
                     async for chunk in cls.webdriver_login(proxy):
diff --git a/g4f/Provider/needs_auth/OpenaiChat.py b/g4f/Provider/needs_auth/OpenaiChat.py
index 03ea4539..b4b8bb02 100644
--- a/g4f/Provider/needs_auth/OpenaiChat.py
+++ b/g4f/Provider/needs_auth/OpenaiChat.py
@@ -403,7 +403,7 @@ class OpenaiChat(AsyncGeneratorProvider, ProviderModelMixin):
                 except NoValidHarFileError as e:
                     error = e
                 if cls._api_key is None:
-                    await cls.nodriver_access_token()
+                    await cls.nodriver_access_token(proxy)
                 if cls._api_key is None and cls.needs_auth:
                     raise error
                 cls.default_model = cls.get_model(await cls.get_default_model(session, cls._headers))
@@ -625,7 +625,7 @@ this.fetch = async (url, options) => {
         cls._update_cookie_header()
 
     @classmethod
-    async def nodriver_access_token(cls):
+    async def nodriver_access_token(cls, proxy: str = None):
         try:
             import nodriver as uc
         except ImportError:
@@ -637,7 +637,10 @@ this.fetch = async (url, options) => {
             user_data_dir = None
         if debug.logging:
             print(f"Open nodriver with user_dir: {user_data_dir}")
-        browser = await uc.start(user_data_dir=user_data_dir)
+        browser = await uc.start(
+            user_data_dir=user_data_dir,
+            browser_args=None if proxy is None else [f"--proxy-server={proxy}"],
+        )
         page = await browser.get("https://chatgpt.com/")
         await page.select("[id^=headlessui-menu-button-]", 240)
         api_key = await page.evaluate(
diff --git a/g4f/Provider/needs_auth/PerplexityApi.py b/g4f/Provider/needs_auth/PerplexityApi.py
new file mode 100644
index 00000000..35d8d9d6
--- /dev/null
+++ b/g4f/Provider/needs_auth/PerplexityApi.py
@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from .Openai import Openai
+from ...typing import AsyncResult, Messages
+
+class PerplexityApi(Openai):
+    label = "Perplexity API"
+    url = "https://www.perplexity.ai"
+    working = True
+    default_model = "llama-3-sonar-large-32k-online"
+    models = [
+        "llama-3-sonar-small-32k-chat",
+        "llama-3-sonar-small-32k-online",
+        "llama-3-sonar-large-32k-chat",
+        "llama-3-sonar-large-32k-online",
+        "llama-3-8b-instruct",
+        "llama-3-70b-instruct",
+        "mixtral-8x7b-instruct"
+    ]
+
+    @classmethod
+    def create_async_generator(
+        cls,
+        model: str,
+        messages: Messages,
+        api_base: str = "https://api.perplexity.ai",
+        **kwargs
+    ) -> AsyncResult:
+        return super().create_async_generator(
+            model, messages, api_base=api_base, **kwargs
+        )
\ No newline at end of file
diff --git a/g4f/Provider/needs_auth/__init__.py b/g4f/Provider/needs_auth/__init__.py
index 805d9fca..b5463b71 100644
--- a/g4f/Provider/needs_auth/__init__.py
+++ b/g4f/Provider/needs_auth/__init__.py
@@ -7,4 +7,5 @@ from .Poe           import Poe
 from .Openai        import Openai
 from .Groq          import Groq
 from .OpenRouter    import OpenRouter
-from .OpenaiAccount import OpenaiAccount
\ No newline at end of file
+from .OpenaiAccount import OpenaiAccount
+from .PerplexityApi import PerplexityApi
\ No newline at end of file
-- 
cgit v1.2.3