Refactor Gemini, GeminiCLI, GeminiPro, and QwenCode providers to streamline model handling and track live instances

This commit is contained in:
hlohaus 2025-09-06 12:32:13 +02:00
parent 6607c82d3c
commit 9fbff6347c
4 changed files with 11 additions and 29 deletions

View file

@ -91,10 +91,6 @@ class Gemini(AsyncGeneratorProvider, ProviderModelMixin):
models = [
default_model, "gemini-2.5-flash", "gemini-2.5-pro"
]
model_aliases = {
"gemini-2.0": "",
"gemini-2.0-flash": ["gemini-2.0-flash", "gemini-2.0-flash", "gemini-2.0-flash-exp"],
}
synthesize_content_type = "audio/vnd.wav"
@ -106,29 +102,6 @@ class Gemini(AsyncGeneratorProvider, ProviderModelMixin):
refresh_interval = 540
rotate_tasks = {}
@classmethod
def get_model(cls, model: str) -> str:
"""Get the internal model name from the user-provided model name."""
if not model:
return cls.default_model
# Check if the model exists directly in our models list
if model in cls.models:
return model
# Check if there's an alias for this model
if model in cls.model_aliases:
alias = cls.model_aliases[model]
# If the alias is a list, randomly select one of the options
if isinstance(alias, list):
selected_model = random.choice(alias)
debug.log(f"Gemini: Selected model '{selected_model}' from alias '{model}'")
return selected_model
debug.log(f"Gemini: Using model '{alias}' for alias '{model}'")
return alias
raise ModelNotFoundError(f"Model {model} not found")
@classmethod
async def nodriver_login(cls, proxy: str = None) -> AsyncIterator[str]:
if not has_nodriver:

View file

@ -510,6 +510,15 @@ class GeminiCLI(AsyncGeneratorProvider, ProviderModelMixin):
auth_manager: AuthManager = None
@classmethod
def get_models(cls, **kwargs):
if cls.live == 0:
if cls.auth_manager is None:
cls.auth_manager = AuthManager(env=os.environ)
if cls.auth_manager.get_access_token() is not None:
cls.live += 1
return cls.models
@classmethod
async def create_async_generator(
cls,

View file

@ -59,6 +59,7 @@ class GeminiPro(AsyncGeneratorProvider, ProviderModelMixin):
if "generateContent" in model.get("supportedGenerationMethods")
]
cls.models.sort()
cls.live += 1
except Exception as e:
debug.error(e)
if api_key is not None:

View file

@ -15,12 +15,11 @@ class QwenCode(OpenaiTemplate):
needs_auth = True
active_by_default = True
default_model = "qwen3-coder-plus"
models = [default_model]
client = QwenContentGenerator(QwenOAuth2Client())
@classmethod
def get_models(cls, **kwargs):
if not cls.models:
cls.models = [cls.default_model]
if cls.live == 0 and cls.client.shared_manager.isTokenValid(cls.client.shared_manager.getCurrentCredentials()):
cls.live += 1
return cls.models