From d0aba83db372f655f775c7e56970f04adfff18e4 Mon Sep 17 00:00:00 2001
From: LanQian <5499636+LanQian528@users.noreply.github.com>
Date: Tue, 14 May 2024 02:40:35 +0800
Subject: [PATCH] update gpt-4o

---
 api/models.py           |  4 +++-
 chatgpt/ChatService.py  | 19 ++++++++++---------
 chatgpt/reverseProxy.py |  6 ++++--
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/api/models.py b/api/models.py
index ebdaeff..d759246 100644
--- a/api/models.py
+++ b/api/models.py
@@ -6,6 +6,7 @@ model_proxy = {
     "gpt-4-turbo-preview": "gpt-4-0125-preview",
     "gpt-4-vision-preview": "gpt-4-1106-vision-preview",
     "gpt-4-turbo": "gpt-4-turbo-2024-04-09",
+    "gpt-4o": "gpt-4o-2024-05-13",
     "claude-3-opus": "claude-3-opus-20240229",
     "claude-3-sonnet": "claude-3-sonnet-20240229",
     "claude-3-haiku": "claude-3-haiku-20240307",
@@ -18,5 +19,6 @@ model_system_fingerprint = {
                            "fp_b77cb481ed"],
     "gpt-4-1106-preview": ["fp_e467c31c3d", "fp_d986a8d1ba", "fp_99a5a401bb", "fp_123d5a9f90", "fp_0d1affc7a6",
                            "fp_5c95a4634e"],
-    "gpt-4-turbo-2024-04-09": ["fp_d1bac968b4"]
+    "gpt-4-turbo-2024-04-09": ["fp_d1bac968b4"],
+    "gpt-4o-2024-05-13": ["fp_d1bac968b4"]
 }
diff --git a/chatgpt/ChatService.py b/chatgpt/ChatService.py
index 5a07bb3..f07b2ba 100644
--- a/chatgpt/ChatService.py
+++ b/chatgpt/ChatService.py
@@ -96,12 +96,13 @@ class ChatService:
                 resp = r.json()
                 self.persona = resp.get("persona")
                 if "gpt-4" in self.origin_model and self.persona != "chatgpt-paid":
-                    raise HTTPException(status_code=404, detail={
-                        "message": f"The model `{self.origin_model}` does not exist or you do not have access to it.",
-                        "type": "invalid_request_error",
-                        "param": None,
-                        "code": "model_not_found"
-                    })
+                    if "gpt-4o" not in self.origin_model:
+                        raise HTTPException(status_code=404, detail={
+                            "message": f"The model `{self.origin_model}` does not exist or you do not have access to it.",
+                            "type": "invalid_request_error",
+                            "param": None,
+                            "code": "model_not_found"
+                        })
                 arkose = resp.get('arkose', {})
                 proofofwork = resp.get('proofofwork', {})
                 turnstile = resp.get('turnstile', {})
@@ -166,19 +167,19 @@ class ChatService:
             'Openai-Sentinel-Proof-Token': self.proof_token,
             'Openai-Sentinel-Arkose-Token': self.arkose_token,
         })
+        conversation_mode = {"kind": "primary_assistant"}
         if "gizmo" in self.origin_model:
             model = "gpt-4"
             gizmo_id = self.data.get("model").split("gpt-4-gizmo-")[-1]
             conversation_mode = {"kind": "gizmo_interaction", "gizmo_id": gizmo_id}
         elif "gpt-4-mobile" in self.origin_model:
             model = "gpt-4-mobile"
-            conversation_mode = {"kind": "primary_assistant"}
+        elif "gpt-4o" in self.origin_model:
+            model = "gpt-4o"
         elif "gpt-4" in self.origin_model:
             model = "gpt-4"
-            conversation_mode = {"kind": "primary_assistant"}
         else:
             model = "text-davinci-002-render-sha"
-            conversation_mode = {"kind": "primary_assistant"}
         logger.info(f"Model mapping: {self.origin_model} -> {model}")
         self.chat_request = {
             "action": "next",
diff --git a/chatgpt/reverseProxy.py b/chatgpt/reverseProxy.py
index f985f43..8971bab 100644
--- a/chatgpt/reverseProxy.py
+++ b/chatgpt/reverseProxy.py
@@ -105,7 +105,9 @@ async def chatgpt_reverse_proxy(request: Request, path: str):
         try:
             r = await client.request(request.method, f"{base_url}/{path}", params=params, headers=headers,
                                      cookies=request_cookies, data=data, stream=True, allow_redirects=False)
-            if r.status_code == 307:
+            if r.status_code == 304:
+                return Response(status_code=304)
+            elif r.status_code == 307:
                 if "oai-dm=1" not in r.headers.get("Location"):
                     return Response(status_code=307, headers={
                         "Location": r.headers.get("Location").replace("chat.openai.com", origin_host)
@@ -135,7 +137,7 @@ async def chatgpt_reverse_proxy(request: Request, path: str):
                 for cookie_name in r.cookies:
                     if cookie_name in request_cookies:
                         continue
-                    for cookie_domain in [".chatgpt.com", ".chat.openai.com"]:
+                    for cookie_domain in [".chatgpt.com"]:
                         cookie_value = r.cookies.get(name=cookie_name, domain=cookie_domain)
                         if cookie_name.startswith("__"):
                             response.set_cookie(key=cookie_name, value=cookie_value, secure=True, httponly=True)