From e39b1d25ee4fb093b96c30558120c79ba535c8a8 Mon Sep 17 00:00:00 2001
From: Junyan Qin <1010553892@qq.com>
Date: Sat, 23 Sep 2023 03:46:35 +0000
Subject: [PATCH] feat: forwarder layer exception handling

---
 free_one_api/entities/exceptions.py      | 24 +++++++++
 free_one_api/impls/adapter/revChatGPT.py |  1 -
 free_one_api/impls/forward/mgr.py        | 69 ++++++++++++++++--------
 3 files changed, 70 insertions(+), 24 deletions(-)
 create mode 100644 free_one_api/entities/exceptions.py

diff --git a/free_one_api/entities/exceptions.py b/free_one_api/entities/exceptions.py
new file mode 100644
index 0000000..9866d0c
--- /dev/null
+++ b/free_one_api/entities/exceptions.py
@@ -0,0 +1,24 @@
+
+
+class QueryHandlingError(Exception):
+    
+    status_code: int
+    code: str
+    message: str
+    type: str
+    param: str
+    
+    def __init__(self, status_code: int, code: str, message: str, type: str=None, param: str=None):
+        """Raise this exception when the query handling failed.
+        
+        Args:
+            code (int): The http status code that should be returned to the client, please refer to the OpenAI API document.
+            message (str): The error message that should be returned to the client.
+            type (str): The error type that should be returned to the client.
+            param (str): The error param that should be returned to the client.
+        """
+        self.status_code = status_code
+        self.code = code
+        self.message = message
+        self.type = type
+        self.param = param
diff --git a/free_one_api/impls/adapter/revChatGPT.py b/free_one_api/impls/adapter/revChatGPT.py
index 30f3766..cfb536e 100644
--- a/free_one_api/impls/adapter/revChatGPT.py
+++ b/free_one_api/impls/adapter/revChatGPT.py
@@ -117,7 +117,6 @@ async def query(self, req: request.Request) -> typing.Generator[response.Respons
                     ]
                 }
             })
-        new_messages
         
         random_int = random.randint(0, 1000000000)
         
diff --git a/free_one_api/impls/forward/mgr.py b/free_one_api/impls/forward/mgr.py
index 924e7cd..e4c419b 100644
--- a/free_one_api/impls/forward/mgr.py
+++ b/free_one_api/impls/forward/mgr.py
@@ -8,7 +8,7 @@
 from ...models.forward import mgr as forwardmgr
 from ...models.channel import mgr as channelmgr
 from ...models.key import mgr as apikeymgr
-from ...entities import channel, apikey, request, response
+from ...entities import channel, apikey, request, response, exceptions
 
 
 class ForwardManager(forwardmgr.AbsForwardManager):
@@ -27,25 +27,35 @@ async def __stream_query(
         
         t = int(time.time())
         async def _gen():
-            async for resp in chan.adapter.query(req):
-                
-                if (resp.normal_message is None or len(resp.normal_message) == 0) and resp.finish_reason == response.FinishReason.NULL:
-                    continue
-                
-                yield "data: {}\n\n".format(json.dumps({
-                    "id": "chatcmpl-"+id_suffix,
-                    "object": "chat.completion.chunk",
-                    "created": t,
-                    "model": req.model,
-                    "choices": [{
-                        "index": 0,
-                        "delta": {
-                            "content": resp.normal_message,
-                        } if resp.normal_message else {},
-                        "finish_reason": resp.finish_reason.value
-                    }]
+            try:
+                async for resp in chan.adapter.query(req):
+                    
+                    if (resp.normal_message is None or len(resp.normal_message) == 0) and resp.finish_reason == response.FinishReason.NULL:
+                        continue
+                    
+                    yield "data: {}\n\n".format(json.dumps({
+                        "id": "chatcmpl-"+id_suffix,
+                        "object": "chat.completion.chunk",
+                        "created": t,
+                        "model": req.model,
+                        "choices": [{
+                            "index": 0,
+                            "delta": {
+                                "content": resp.normal_message,
+                            } if resp.normal_message else {},
+                            "finish_reason": resp.finish_reason.value
+                        }]
+                    }))
+                yield "data: [DONE]\n\n"
+            except exceptions.QueryHandlingError as e:
+                yield "data: {}\n\ndata: [DONE]\n\n".format(json.dumps({
+                    "error": {
+                        "code": e.code,
+                        "message": e.message,
+                        "type": e.type,
+                        "param": e.param,
+                    }
                 }))
-            yield "data: [DONE]\n\n"
         
         spent_ms = int((time.time() - before)*1000)
         
@@ -82,10 +92,23 @@ async def __non_stream_query(
         
         resp_tmp: response.Response = None
         
-        async for resp in chan.adapter.query(req):
-            if resp.normal_message is not None:
-                resp_tmp = resp
-                normal_message += resp.normal_message
+        try:
+        
+            async for resp in chan.adapter.query(req):
+                if resp.normal_message is not None:
+                    resp_tmp = resp
+                    normal_message += resp.normal_message
+                    
+        except exceptions.QueryHandlingError as e:
+            # check for custom error raised by adapter
+            return quart.jsonify({
+                "error": {
+                    "code": e.code,
+                    "message": e.message,
+                    "type": e.type,
+                    "param": e.param,
+                }
+            }), e.status_code
                 
         spent_ms = int((time.time() - before)*1000)