From c1b992c3460cb0069127524a9b987a8af475ec14 Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <hlohaus@users.noreply.github.com>
Date: Thu, 8 Feb 2024 22:02:52 +0100
Subject: Add Gemini Provider with image upload and generation

---
 g4f/Provider/needs_auth/ThebApi.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

(limited to 'g4f/Provider/needs_auth/ThebApi.py')

diff --git a/g4f/Provider/needs_auth/ThebApi.py b/g4f/Provider/needs_auth/ThebApi.py
index 8ec7bda8..ea633243 100644
--- a/g4f/Provider/needs_auth/ThebApi.py
+++ b/g4f/Provider/needs_auth/ThebApi.py
@@ -3,7 +3,8 @@ from __future__ import annotations
 import requests
 
 from ...typing import Any, CreateResult, Messages
-from ..base_provider import AbstractProvider
+from ..base_provider import AbstractProvider, ProviderModelMixin
+from ...errors import MissingAuthError
 
 models = {
     "theb-ai": "TheB.AI",
@@ -29,13 +30,16 @@ models = {
     "qwen-7b-chat": "Qwen 7B"
 }
 
-class ThebApi(AbstractProvider):
+class ThebApi(AbstractProvider, ProviderModelMixin):
     url = "https://theb.ai"
     working = True
     needs_auth = True
+    default_model = "gpt-3.5-turbo"
+    models = list(models)
 
-    @staticmethod
+    @classmethod
     def create_completion(
+        cls,
         model: str,
         messages: Messages,
         stream: bool,
@@ -43,8 +47,8 @@ class ThebApi(AbstractProvider):
         proxy: str = None,
         **kwargs
     ) -> CreateResult:
-        if model and model not in models:
-            raise ValueError(f"Model are not supported: {model}")
+        if not auth:
+            raise MissingAuthError("Missing auth")
         headers = {
             'accept': 'application/json',
             'authorization': f'Bearer {auth}',
@@ -54,7 +58,7 @@ class ThebApi(AbstractProvider):
         # models = dict([(m["id"], m["name"]) for m in response])
         # print(json.dumps(models, indent=4))
         data: dict[str, Any] = {
-            "model": model if model else "gpt-3.5-turbo",
+            "model": cls.get_model(model),
             "messages": messages,
             "stream": False,
             "model_params": {
-- 
cgit v1.2.3