From 496ec099059fef4654ab6913a124989e850ab921 Mon Sep 17 00:00:00 2001
From: user <user@example.com>
Date: Sat, 4 Jan 2025 19:11:02 +0000
Subject: [PATCH] Add v1alpha support (needed for 2.0 flash thinking with the
 new 'thought' parameter), already used by ST

---
 src/proxy/add-v1.ts                                        | 2 +-
 src/proxy/google-ai.ts                                     | 2 +-
 src/proxy/middleware/common.ts                             | 6 ++++--
 src/proxy/middleware/request/mutators/add-google-ai-key.ts | 2 +-
 src/server.ts                                              | 2 +-
 src/shared/api-schemas/google-ai.ts                        | 6 +++++-
 6 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/proxy/add-v1.ts b/src/proxy/add-v1.ts
index d0d8620..6ddda96 100644
--- a/src/proxy/add-v1.ts
+++ b/src/proxy/add-v1.ts
@@ -2,7 +2,7 @@ import { NextFunction, Request, Response } from "express";
 
 export function addV1(req: Request, res: Response, next: NextFunction) {
   // Clients don't consistently use the /v1 prefix so we'll add it for them.
-  if (!req.path.startsWith("/v1/") && !req.path.startsWith("/v1beta/")) {
+  if (!req.path.startsWith("/v1/") && !req.path.match(/^\/(v1alpha|v1beta)\//)) {
     req.url = `/v1${req.url}`;
   }
   next();
diff --git a/src/proxy/google-ai.ts b/src/proxy/google-ai.ts
index 9ccc1d2..d4fb3bf 100644
--- a/src/proxy/google-ai.ts
+++ b/src/proxy/google-ai.ts
@@ -117,7 +117,7 @@ googleAIRouter.get("/v1/models", handleModelRequest);
 
 // Native Google AI chat completion endpoint
 googleAIRouter.post(
-  "/v1beta/models/:modelId:(generateContent|streamGenerateContent)",
+  "/:apiVersion(v1alpha|v1beta)/models/:modelId:(generateContent|streamGenerateContent)",
   ipLimiter,
   createPreprocessorMiddleware(
     { inApi: "google-ai", outApi: "google-ai", service: "google-ai" },
diff --git a/src/proxy/middleware/common.ts b/src/proxy/middleware/common.ts
index d8d0e2f..cb83021 100644
--- a/src/proxy/middleware/common.ts
+++ b/src/proxy/middleware/common.ts
@@ -16,7 +16,8 @@ const ANTHROPIC_COMPLETION_ENDPOINT = "/v1/complete";
 const ANTHROPIC_MESSAGES_ENDPOINT = "/v1/messages";
 const ANTHROPIC_SONNET_COMPAT_ENDPOINT = "/v1/sonnet";
 const ANTHROPIC_OPUS_COMPAT_ENDPOINT = "/v1/opus";
-const GOOGLE_AI_COMPLETION_ENDPOINT = "/v1beta/models";
+const GOOGLE_AI_ALPHA_COMPLETION_ENDPOINT = "/v1alpha/models";
+const GOOGLE_AI_BETA_COMPLETION_ENDPOINT = "/v1beta/models";
 
 export function isTextGenerationRequest(req: Request) {
   return (
@@ -28,7 +29,8 @@ export function isTextGenerationRequest(req: Request) {
       ANTHROPIC_MESSAGES_ENDPOINT,
       ANTHROPIC_SONNET_COMPAT_ENDPOINT,
       ANTHROPIC_OPUS_COMPAT_ENDPOINT,
-      GOOGLE_AI_COMPLETION_ENDPOINT,
+      GOOGLE_AI_ALPHA_COMPLETION_ENDPOINT,
+      GOOGLE_AI_BETA_COMPLETION_ENDPOINT,
     ].some((endpoint) => req.path.startsWith(endpoint))
   );
 }
diff --git a/src/proxy/middleware/request/mutators/add-google-ai-key.ts b/src/proxy/middleware/request/mutators/add-google-ai-key.ts
index 15a25a5..8b160d3 100644
--- a/src/proxy/middleware/request/mutators/add-google-ai-key.ts
+++ b/src/proxy/middleware/request/mutators/add-google-ai-key.ts
@@ -32,7 +32,7 @@ export const addGoogleAIKey: ProxyReqMutator = (manager) => {
     method: "POST",
     protocol: "https:",
     hostname: "generativelanguage.googleapis.com",
-    path: `/v1beta/models/${model}:${
+    path: `/${req.params.apiVersion}/models/${model}:${
       req.isStreaming ? "streamGenerateContent?alt=sse&" : "generateContent?"
     }key=${key.key}`,
     headers: {
diff --git a/src/server.ts b/src/server.ts
index 5b180e8..50760b5 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -92,7 +92,7 @@ app.use("/admin", adminRouter);
 app.use((req, _, next) => {
   // For whatever reason SillyTavern just ignores the path a user provides
   // when using Google AI with reverse proxy.  We'll fix it here.
-  if (req.path.startsWith("/v1beta/models/")) {
+  if (req.path.match(/^\/v1(alpha|beta)\/models\//)) {
     req.url = `${config.proxyEndpointRoute}/google-ai${req.url}`;
     return next();
   }
diff --git a/src/shared/api-schemas/google-ai.ts b/src/shared/api-schemas/google-ai.ts
index 7561b2b..c59e818 100644
--- a/src/shared/api-schemas/google-ai.ts
+++ b/src/shared/api-schemas/google-ai.ts
@@ -5,7 +5,11 @@ import {
 } from "./openai";
 import { APIFormatTransformer } from "./index";
 
-const TextPartSchema = z.object({ text: z.string() });
+const TextPartSchema = z.object({ 
+  text: z.string(),
+  thought: z.boolean().optional()
+});
+
 const InlineDataPartSchema = z.object({
   inlineData: z.object({
     mimeType: z.string(),