Add support to iframe chatbot (#3929)

### What problem does this PR solve?

#3909

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

api/db/db_models.py

@@ -934,6 +934,7 @@ class APIToken(DataBaseModel):
     token = CharField(max_length=255, null=False, index=True)
     dialog_id = CharField(max_length=32, null=False, index=True)
     source = CharField(max_length=16, null=True, help_text="none|agent|dialog", index=True)
+    beta = CharField(max_length=255, null=True, index=True)
 
     class Meta:
         db_table = "api_token"
@@ -1083,4 +1084,10 @@ def migrate_db():
             )
         except Exception:
             pass
+        try:
+            migrate(
+                migrator.add_column("api_token", "beta", CharField(max_length=255, null=True, index=True))
+            )
+        except Exception:
+            pass
 

api/db/services/canvas_service.py

@@ -13,8 +13,14 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import json
+from uuid import uuid4
+from agent.canvas import Canvas
 from api.db.db_models import DB, CanvasTemplate, UserCanvas
+from api.db.services.api_service import API4ConversationService
 from api.db.services.common_service import CommonService
+from api.db.services.conversation_service import structure_answer
+from api.utils import get_uuid
 
 
 class CanvasTemplateService(CommonService):
@@ -42,3 +48,115 @@ class UserCanvasService(CommonService):
         agents = agents.paginate(page_number, items_per_page)
 
         return list(agents.dicts())
+
+
+def completion(tenant_id, agent_id, question, session_id=None, stream=True, **kwargs):
+    e, cvs = UserCanvasService.get_by_id(agent_id)
+    assert e, "Agent not found."
+    assert cvs.user_id == tenant_id, "You do not own the agent."
+
+    if not isinstance(cvs.dsl, str):
+        cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
+    canvas = Canvas(cvs.dsl, tenant_id)
+
+    if not session_id:
+        session_id = get_uuid()
+        conv = {
+            "id": session_id,
+            "dialog_id": cvs.id,
+            "user_id": kwargs.get("user_id", ""),
+            "message": [{"role": "assistant", "content": canvas.get_prologue()}],
+            "source": "agent",
+            "dsl": json.loads(cvs.dsl)
+        }
+        API4ConversationService.save(**conv)
+        yield "data:" + json.dumps({"code": 0,
+                                    "message": "",
+                                    "data": {
+                                        "session_id": session_id,
+                                        "answer": canvas.get_prologue(),
+                                        "reference": [],
+                                        "param": canvas.get_preset_param()
+                                      }
+                                    },
+                                   ensure_ascii=False) + "\n\n"
+        yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
+        return
+    else:
+        session_id = session_id
+        e, conv = API4ConversationService.get_by_id(session_id)
+        assert e, "Session not found!"
+        canvas = Canvas(json.dumps(conv.dsl), tenant_id)
+
+    messages = conv.message
+    question = {
+        "role": "user",
+        "content": question,
+        "id": str(uuid4())
+    }
+    messages.append(question)
+    msg = []
+    for m in messages:
+        if m["role"] == "system":
+            continue
+        if m["role"] == "assistant" and not msg:
+            continue
+        msg.append(m)
+    if not msg[-1].get("id"):
+        msg[-1]["id"] = get_uuid()
+    message_id = msg[-1]["id"]
+
+    if not conv.reference:
+        conv.reference = []
+    conv.message.append({"role": "assistant", "content": "", "id": message_id})
+    conv.reference.append({"chunks": [], "doc_aggs": []})
+
+    final_ans = {"reference": [], "content": ""}
+
+    canvas.add_user_input(msg[-1]["content"])
+
+    if stream:
+        try:
+            for ans in canvas.run(stream=stream):
+                if ans.get("running_status"):
+                    yield "data:" + json.dumps({"code": 0, "message": "",
+                                                "data": {"answer": ans["content"],
+                                                         "running_status": True}},
+                                               ensure_ascii=False) + "\n\n"
+                    continue
+                for k in ans.keys():
+                    final_ans[k] = ans[k]
+                ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
+                ans = structure_answer(conv, ans, message_id, session_id)
+                yield "data:" + json.dumps({"code": 0, "message": "", "data": ans},
+                                           ensure_ascii=False) + "\n\n"
+
+            canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
+            canvas.history.append(("assistant", final_ans["content"]))
+            if final_ans.get("reference"):
+                canvas.reference.append(final_ans["reference"])
+            conv.dsl = json.loads(str(canvas))
+            API4ConversationService.append_message(conv.id, conv.to_dict())
+        except Exception as e:
+            conv.dsl = json.loads(str(canvas))
+            API4ConversationService.append_message(conv.id, conv.to_dict())
+            yield "data:" + json.dumps({"code": 500, "message": str(e),
+                                        "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
+                                       ensure_ascii=False) + "\n\n"
+        yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
+
+    else:
+        for answer in canvas.run(stream=False):
+            if answer.get("running_status"):
+                continue
+            final_ans["content"] = "\n".join(answer["content"]) if "content" in answer else ""
+            canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
+            if final_ans.get("reference"):
+                canvas.reference.append(final_ans["reference"])
+            conv.dsl = json.loads(str(canvas))
+
+            result = {"answer": final_ans["content"], "reference": final_ans.get("reference", [])}
+            result = structure_answer(conv, result, message_id, session_id)
+            API4ConversationService.append_message(conv.id, conv.to_dict())
+            yield result
+            break

api/db/services/conversation_service.py

@@ -0,0 +1,221 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+from uuid import uuid4
+from api.db import StatusEnum
+from api.db.db_models import Conversation, DB
+from api.db.services.api_service import API4ConversationService
+from api.db.services.common_service import CommonService
+from api.db.services.dialog_service import DialogService, chat
+from api.utils import get_uuid
+import json
+from copy import deepcopy
+
+
+class ConversationService(CommonService):
+    model = Conversation
+
+    @classmethod
+    @DB.connection_context()
+    def get_list(cls,dialog_id,page_number, items_per_page, orderby, desc, id , name):
+        sessions = cls.model.select().where(cls.model.dialog_id ==dialog_id)
+        if id:
+            sessions = sessions.where(cls.model.id == id)
+        if name:
+            sessions = sessions.where(cls.model.name == name)
+        if desc:
+            sessions = sessions.order_by(cls.model.getter_by(orderby).desc())
+        else:
+            sessions = sessions.order_by(cls.model.getter_by(orderby).asc())
+
+        sessions = sessions.paginate(page_number, items_per_page)
+
+        return list(sessions.dicts())
+
+
+def structure_answer(conv, ans, message_id, session_id):
+    reference = ans["reference"]
+    temp_reference = deepcopy(ans["reference"])
+    if not conv.reference:
+        conv.reference.append(temp_reference)
+    else:
+        conv.reference[-1] = temp_reference
+    conv.message[-1] = {"role": "assistant", "content": ans["answer"], "id": message_id}
+
+    chunk_list = [{
+            "id": chunk["chunk_id"],
+            "content": chunk["content"],
+            "document_id": chunk["doc_id"],
+            "document_name": chunk["docnm_kwd"],
+            "dataset_id": chunk["kb_id"],
+            "image_id": chunk["image_id"],
+            "similarity": chunk["similarity"],
+            "vector_similarity": chunk["vector_similarity"],
+            "term_similarity": chunk["term_similarity"],
+            "positions": chunk["positions"],
+        } for chunk in reference.get("chunks", [])]
+
+    reference["chunks"] = chunk_list
+    ans["id"] = message_id
+    ans["session_id"] = session_id
+
+    return ans
+
+
+def completion(tenant_id, chat_id, question, name="New session", session_id=None, stream=True, **kwargs):
+    assert name, "`name` can not be empty."
+    dia = DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value)
+    assert dia, "You do not own the chat."
+
+    if not session_id:
+        conv = {
+            "id": get_uuid(),
+            "dialog_id": chat_id,
+            "name": name,
+            "message": [{"role": "assistant", "content": dia[0].prompt_config.get("prologue")}]
+        }
+        ConversationService.save(**conv)
+        yield "data:" + json.dumps({"code": 0, "message": "",
+                                    "data": {
+                                        "answer": conv["message"][0]["content"],
+                                        "reference": {},
+                                        "audio_binary": None,
+                                        "id": None,
+                                        "session_id": session_id
+                                    }},
+                                   ensure_ascii=False) + "\n\n"
+        yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
+        return
+
+    conv = ConversationService.query(id=session_id, dialog_id=chat_id)
+    if not conv:
+        raise LookupError("Session does not exist")
+
+    conv = conv[0]
+    msg = []
+    question = {
+        "content": question,
+        "role": "user",
+        "id": str(uuid4())
+    }
+    conv.message.append(question)
+    for m in conv.message:
+        if m["role"] == "system":
+            continue
+        if m["role"] == "assistant" and not msg:
+            continue
+        msg.append(m)
+    message_id = msg[-1].get("id")
+    e, dia = DialogService.get_by_id(conv.dialog_id)
+
+    if not conv.reference:
+        conv.reference = []
+    conv.message.append({"role": "assistant", "content": "", "id": message_id})
+    conv.reference.append({"chunks": [], "doc_aggs": []})
+
+    if stream:
+        try:
+            for ans in chat(dia, msg, True, **kwargs):
+                ans = structure_answer(conv, ans, message_id, session_id)
+                yield "data:" + json.dumps({"code": 0, "data": ans}, ensure_ascii=False) + "\n\n"
+            ConversationService.update_by_id(conv.id, conv.to_dict())
+        except Exception as e:
+            yield "data:" + json.dumps({"code": 500, "message": str(e),
+                                        "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
+                                       ensure_ascii=False) + "\n\n"
+        yield "data:" + json.dumps({"code": 0, "data": True}, ensure_ascii=False) + "\n\n"
+
+    else:
+        answer = None
+        for ans in chat(dia, msg, False, **kwargs):
+            answer = structure_answer(conv, ans, message_id, session_id)
+            ConversationService.update_by_id(conv.id, conv.to_dict())
+            break
+        yield answer
+
+
+def iframe_completion(dialog_id, question, session_id=None, stream=True, **kwargs):
+    e, dia = DialogService.get_by_id(dialog_id)
+    assert e, "Dialog not found"
+    if not session_id:
+        session_id = get_uuid()
+        conv = {
+            "id": session_id,
+            "dialog_id": dialog_id,
+            "user_id": kwargs.get("user_id", ""),
+            "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}]
+        }
+        API4ConversationService.save(**conv)
+        yield "data:" + json.dumps({"code": 0, "message": "",
+                                    "data": {
+                                        "answer": conv["message"][0]["content"],
+                                        "reference": {},
+                                        "audio_binary": None,
+                                        "id": None,
+                                        "session_id": session_id
+                                    }},
+                                   ensure_ascii=False) + "\n\n"
+        yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
+        return
+    else:
+        session_id = session_id
+        e, conv = API4ConversationService.get_by_id(session_id)
+        assert e, "Session not found!"
+
+    messages = conv.message
+    question = {
+        "role": "user",
+        "content": question,
+        "id": str(uuid4())
+    }
+    messages.append(question)
+
+    msg = []
+    for m in messages:
+        if m["role"] == "system":
+            continue
+        if m["role"] == "assistant" and not msg:
+            continue
+        msg.append(m)
+    if not msg[-1].get("id"):
+        msg[-1]["id"] = get_uuid()
+    message_id = msg[-1]["id"]
+
+    if not conv.reference:
+        conv.reference = []
+    conv.message.append({"role": "assistant", "content": "", "id": message_id})
+    conv.reference.append({"chunks": [], "doc_aggs": []})
+
+    if stream:
+        try:
+            for ans in chat(dia, msg, True, **kwargs):
+                ans = structure_answer(conv, ans, message_id, session_id)
+                yield "data:" + json.dumps({"code": 0, "message": "", "data": ans},
+                                           ensure_ascii=False) + "\n\n"
+            API4ConversationService.append_message(conv.id, conv.to_dict())
+        except Exception as e:
+            yield "data:" + json.dumps({"code": 500, "message": str(e),
+                                        "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
+                                       ensure_ascii=False) + "\n\n"
+        yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
+
+    else:
+        answer = None
+        for ans in chat(dia, msg, False, **kwargs):
+            answer = structure_answer(conv, ans, message_id, session_id)
+            API4ConversationService.append_message(conv.id, conv.to_dict())
+            break
+        yield answer
+

api/db/services/dialog_service.py

@@ -23,7 +23,7 @@ from timeit import default_timer as timer
 import datetime
 from datetime import timedelta
 from api.db import LLMType, ParserType,StatusEnum
-from api.db.db_models import Dialog, Conversation,DB
+from api.db.db_models import Dialog, DB
 from api.db.services.common_service import CommonService
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.llm_service import LLMService, TenantLLMService, LLMBundle
@@ -60,27 +60,6 @@ class DialogService(CommonService):
         return list(chats.dicts())
 
 
-class ConversationService(CommonService):
-    model = Conversation
-
-    @classmethod
-    @DB.connection_context()
-    def get_list(cls,dialog_id,page_number, items_per_page, orderby, desc, id , name):
-        sessions = cls.model.select().where(cls.model.dialog_id ==dialog_id)
-        if id:
-            sessions = sessions.where(cls.model.id == id)
-        if name:
-            sessions = sessions.where(cls.model.name == name)
-        if desc:
-            sessions = sessions.order_by(cls.model.getter_by(orderby).desc())
-        else:
-            sessions = sessions.order_by(cls.model.getter_by(orderby).asc())
-
-        sessions = sessions.paginate(page_number, items_per_page)
-
-        return list(sessions.dicts())
-
-
 def message_fit_in(msg, max_length=4000):
     def count():
         nonlocal msg

api/db/services/document_service.py

@@ -425,11 +425,12 @@ def queue_raptor_tasks(doc):
 
 def doc_upload_and_parse(conversation_id, file_objs, user_id):
     from rag.app import presentation, picture, naive, audio, email
-    from api.db.services.dialog_service import ConversationService, DialogService
+    from api.db.services.dialog_service import DialogService
     from api.db.services.file_service import FileService
     from api.db.services.llm_service import LLMBundle
     from api.db.services.user_service import TenantService
     from api.db.services.api_service import API4ConversationService
+    from api.db.services.conversation_service import ConversationService
 
     e, conv = ConversationService.get_by_id(conversation_id)
     if not e: