From 78bc819742c04369879d7bef50f3d33dd2da530a Mon Sep 17 00:00:00 2001
From: Hk-Gosuto <yominogosuto@gmail.com>
Date: Fri, 29 Dec 2023 14:10:51 +0800
Subject: [PATCH] fix: fix gemini stream output

---
 app/api/google/[...path]/route.ts |  4 ++--
 app/client/api.ts                 | 11 ++++++---
 app/client/platforms/google.ts    | 39 +++++++++++++++++++++++--------
 app/constant.ts                   |  3 +--
 4 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/app/api/google/[...path]/route.ts b/app/api/google/[...path]/route.ts
index ef148d6d7..37cb08800 100644
--- a/app/api/google/[...path]/route.ts
+++ b/app/api/google/[...path]/route.ts
@@ -46,7 +46,7 @@ async function handle(
     });
   }
 
-  const bearToken = req.headers.get("Authorization") ?? "";
+  const bearToken = req.headers.get("x-goog-api-key") ?? "";
   const token = bearToken.trim().replaceAll("Bearer ", "").trim();
 
   const key = token ? token : serverConfig.googleApiKey;
@@ -63,7 +63,7 @@ async function handle(
     );
   }
 
-  const fetchUrl = `${baseUrl}/${path}?key=${key}&alt=sse`;
+  const fetchUrl = `${baseUrl}/${path}?key=${key}`;
   const fetchOptions: RequestInit = {
     headers: {
       "Content-Type": "application/json",
diff --git a/app/client/api.ts b/app/client/api.ts
index 6d287b52f..d12e013d8 100644
--- a/app/client/api.ts
+++ b/app/client/api.ts
@@ -174,20 +174,25 @@ export function getHeaders(ignoreHeaders?: boolean) {
     headers = {
       "Content-Type": "application/json",
       "x-requested-with": "XMLHttpRequest",
-      "Accept": "application/json",
+      Accept: "application/json",
     };
   }
   const modelConfig = useChatStore.getState().currentSession().mask.modelConfig;
   const isGoogle = modelConfig.model === "gemini-pro";
   const isAzure = accessStore.provider === ServiceProvider.Azure;
-  const authHeader = isAzure ? "api-key" : "Authorization";
+  const authHeader = isGoogle
+    ? "x-goog-api-key"
+    : isAzure
+    ? "api-key"
+    : "Authorization";
   const apiKey = isGoogle
     ? accessStore.googleApiKey
     : isAzure
     ? accessStore.azureApiKey
     : accessStore.openaiApiKey;
 
-  const makeBearer = (s: string) => `${isAzure ? "" : "Bearer "}${s.trim()}`;
+  const makeBearer = (s: string) =>
+    `${isGoogle || isAzure ? "" : "Bearer "}${s.trim()}`;
   const validString = (x: string) => x && x.length > 0;
 
   // use user's api key first
diff --git a/app/client/platforms/google.ts b/app/client/platforms/google.ts
index 60d93543d..3d28c89be 100644
--- a/app/client/platforms/google.ts
+++ b/app/client/platforms/google.ts
@@ -1,4 +1,9 @@
-import { Google, REQUEST_TIMEOUT_MS } from "@/app/constant";
+import {
+  ApiPath,
+  Google,
+  REQUEST_TIMEOUT_MS,
+  ServiceProvider,
+} from "@/app/constant";
 import {
   AgentChatOptions,
   ChatOptions,
@@ -8,14 +13,7 @@ import {
   LLMUsage,
 } from "../api";
 import { useAccessStore, useAppConfig, useChatStore } from "@/app/store";
-import {
-  EventStreamContentType,
-  fetchEventSource,
-} from "@fortaine/fetch-event-source";
-import { prettyObject } from "@/app/utils/format";
-import { getClientConfig } from "@/app/config/client";
-import Locale from "../../locales";
-import { getServerSideConfig } from "@/app/config/server";
+
 export class GeminiProApi implements LLMApi {
   toolAgentChat(options: AgentChatOptions): Promise<void> {
     throw new Error("Method not implemented.");
@@ -188,7 +186,28 @@ export class GeminiProApi implements LLMApi {
     return [];
   }
   path(path: string): string {
-    return "/api/google/" + path;
+    const accessStore = useAccessStore.getState();
+    const isGoogle = accessStore.provider === ServiceProvider.Google;
+
+    if (isGoogle && !accessStore.isValidGoogle()) {
+      throw Error(
+        "incomplete google config, please check it in your settings page",
+      );
+    }
+
+    let baseUrl = isGoogle ? accessStore.googleBaseUrl : ApiPath.GoogleAI;
+
+    if (baseUrl.length === 0) {
+      baseUrl = ApiPath.GoogleAI;
+    }
+    if (baseUrl.endsWith("/")) {
+      baseUrl = baseUrl.slice(0, baseUrl.length - 1);
+    }
+    if (!baseUrl.startsWith("http") && !baseUrl.startsWith(ApiPath.GoogleAI)) {
+      baseUrl = "https://" + baseUrl;
+    }
+
+    return [baseUrl, path].join("/");
   }
 }
 
diff --git a/app/constant.ts b/app/constant.ts
index ba21e4ef9..43ad8f5cc 100644
--- a/app/constant.ts
+++ b/app/constant.ts
@@ -93,8 +93,7 @@ export const Azure = {
 };
 
 export const Google = {
-  ExampleEndpoint:
-    "https://generativelanguage.googleapis.com/v1/models/gemini-pro:generateContent",
+  ExampleEndpoint: "https://generativelanguage.googleapis.com",
   ChatPath: "v1/models/gemini-pro:generateContent",
 
   // /api/openai/v1/chat/completions