From 5249e1c9044db4af50a1012d9fa6fe0873015cf9 Mon Sep 17 00:00:00 2001
From: reanon <>
Date: Wed, 9 Jul 2025 20:14:01 +0200
Subject: [PATCH] gemini -exp whitelist

---
 src/config.ts          | 10 ++++++++++
 src/proxy/google-ai.ts | 17 +++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/config.ts b/src/config.ts
index 4b1c7fd..4d0f17e 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -29,6 +29,15 @@ type Config = {
    * same but the APIs are different. Vertex is the GCP product for enterprise.
    **/
   googleAIKey?: string;
+  /**
+   * Comma-delimited list of Google AI experimental model names that are
+   * allowed to bypass the experimental model block. By default, all models
+   * containing "exp" are blocked, but specific models listed here will be
+   * permitted.
+   * 
+   * @example "gemini-2.0-flash-exp,gemini-exp-1206"
+   */
+  allowedExpModels?: string;
   /**
    * Comma-delimited list of Mistral AI API keys.
    */
@@ -450,6 +459,7 @@ export const config: Config = {
   anthropicKey: getEnvWithDefault("ANTHROPIC_KEY", ""),
   qwenKey: getEnvWithDefault("QWEN_KEY", ""),
   googleAIKey: getEnvWithDefault("GOOGLE_AI_KEY", ""),
+  allowedExpModels: getEnvWithDefault("ALLOWED_EXP_MODELS", ""),
   mistralAIKey: getEnvWithDefault("MISTRAL_AI_KEY", ""),
   deepseekKey: getEnvWithDefault("DEEPSEEK_KEY", ""),
   xaiKey: getEnvWithDefault("XAI_KEY", ""),
diff --git a/src/proxy/google-ai.ts b/src/proxy/google-ai.ts
index d852535..90d29e2 100644
--- a/src/proxy/google-ai.ts
+++ b/src/proxy/google-ai.ts
@@ -235,12 +235,29 @@ function maybeReassignModel(req: Request) {
  * This function is intended to be used as a RequestPreprocessor.
  * It throws an error if an experimental model is detected, which should be
  * caught by the proxy's onError handler.
+ * 
+ * Models can be allowed through the ALLOWED_EXP_MODELS environment variable.
  */
 function checkAndBlockExperimentalModels(req: Request) { // Changed signature
   const modelId = req.body.model as string | undefined;
 
   // Check if the model ID contains "exp" (case-insensitive)
   if (modelId && modelId.toLowerCase().includes("exp")) {
+    // Check if this specific model is in the allowlist
+    const allowedModels = config.allowedExpModels
+      ?.split(",")
+      .map(model => model.trim())
+      .filter(model => model.length > 0) || [];
+    
+    const isAllowed = allowedModels.some(allowedModel => 
+      modelId.toLowerCase() === allowedModel.toLowerCase()
+    );
+    
+    if (isAllowed) {
+      req.log.info({ modelId }, "Allowing experimental Google AI model via allowlist.");
+      return; // Allow the request to proceed
+    }
+    
     req.log.warn({ modelId }, "Blocking request to experimental Google AI model.");
     const err: any = new Error("Experimental models are too unstable to be supported in proxy code. Please use preview models instead.");
     err.statusCode = 400;