适配外部机器人[新增聊天接口]

api/controllers/service_api/app/completion.py

@@ -1,7 +1,11 @@
+import json
 import logging
+from typing import Generator
 
+import requests
+from flask import request
 from flask_restful import Resource, reqparse  # type: ignore
-from werkzeug.exceptions import InternalServerError, NotFound
+from werkzeug.exceptions import InternalServerError, NotFound, BadRequest
 
 import services
 from controllers.service_api import api
@@ -154,8 +158,111 @@ class ChatStopApi(Resource):
 
         return {"result": "success"}, 200
 
+class ChatApiForRobot(Resource):
+
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("id", type=str, required=True, location="json")
+        parser.add_argument("enterprise_id", type=str, required=False, location="json")
+        parser.add_argument("device_id", type=str, required=False, location="json")
+        parser.add_argument("messages", type=list, required=True, location="json")
+        parser.add_argument("max_tokens", type=int, required=True, location="json")
+        parser.add_argument("stream", type=bool, required=True, location="json")
+        args = parser.parse_args()
+
+        messages = args["messages"]
+        if messages is None or len(messages) == 0:
+            raise BadRequest("messages is empty.")
+
+        id = args["id"]
+        query = messages[len(messages) - 1]["content"]
+        response_mode = "streaming" if args["stream"] else "blocking"
+        device_id = args["device_id"]
+        data = {
+            "inputs": {},
+            "query": query,
+            "response_mode": response_mode,
+            "conversation_id": "",
+            "user": device_id if device_id else "abc-123",
+            "files": []
+        }
+
+        chat_message_url = request.host_url + "v1/chat-messages"
+        logging.info("Sending request to %s", chat_message_url)
+        response = requests.post(chat_message_url, data = json.dumps(data), headers = request.headers)
+
+        if args["stream"]:
+            def after_response_generator():
+                i = 0
+                for line in response.iter_lines():
+                    line_str = line.decode("utf-8")
+                    if not line_str.startswith('data:'):
+                        continue
+
+                    content = json.loads(line_str[6:])
+                    event = content["event"]
+                    if event not in ["message", "message_end"]:
+                        continue
+
+                    new_content = {
+                        "id": id,
+                        "model": "advanced-chat",
+                        "created": content["created_at"],
+                        "choices": [],
+                    }
+
+                    if i == 0:
+                        choice = {
+                            "index": 0,
+                            "delta": {
+                                "role": "assistant",
+                                "content": ""
+                            },
+                            "finish_reason": None
+                        }
+                        new_content["choices"].append(choice)
+                        yield f"id: {i}\ndata: {json.dumps(new_content)}\n\n"
+                        new_content["choices"].pop()
+
+                    i = i + 1
+                    if content["event"] == "message":
+                        choice = {
+                            "index": 0,
+                            "delta": {
+                                "content": content["answer"]
+                            },
+                            "finish_reason": None
+                        }
+                        new_content["choices"].append(choice)
+                    else:
+                        choice = {
+                            "index": 0,
+                            "delta": {},
+                            "finish_reason": "stop"
+                        }
+                        new_content["choices"].append(choice)
+
+                    yield f"id: {i}\ndata: {json.dumps(new_content)}\n\n"
+
+            new_response = after_response_generator()
+
+            def generate() -> Generator:
+                yield from new_response
+
+            return helper.compact_generate_response(generate())
+        else:
+            content = json.loads(response.text)
+            new_response = {
+                "id": id,
+                "model": "advanced-chat",
+                "created": content["created_at"],
+                "answer": content ["answer"],
+            }
+            return new_response
 
 api.add_resource(CompletionApi, "/completion-messages")
 api.add_resource(CompletionStopApi, "/completion-messages/<string:task_id>/stop")
 api.add_resource(ChatApi, "/chat-messages")
 api.add_resource(ChatStopApi, "/chat-messages/<string:task_id>/stop")
+
+api.add_resource(ChatApiForRobot, "/chat-messages-for-robot")