feat: forwarder layer exception handling

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

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)
 

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)