feat: 高级编排支持文件上传(WIP)

Author: liuruibin

apps/application/flow/step_node/__init__.py

@@ -16,9 +16,12 @@
 from .function_lib_node import *
 from .function_node import *
 from .reranker_node import *
+from .document_extract_node import *
+from .image_understand_step_node import *
 
 node_list = [BaseStartStepNode, BaseChatNode, BaseSearchDatasetNode, BaseQuestionNode, BaseConditionNode, BaseReplyNode,
-             BaseFunctionNodeNode, BaseFunctionLibNodeNode, BaseRerankerNode, BaseApplicationNode]
+             BaseFunctionNodeNode, BaseFunctionLibNodeNode, BaseRerankerNode, BaseApplicationNode, BaseDocumentExtractNode,
+             BaseImageUnderstandNode]
 
 
 def get_node(node_type):

apps/application/flow/step_node/document_extract_node/__init__.py

@@ -0,0 +1 @@
+from .impl import *

apps/application/flow/step_node/document_extract_node/i_document_extract_node.py

@@ -0,0 +1,30 @@
+# coding=utf-8
+
+from typing import Type
+
+from rest_framework import serializers
+
+from application.flow.i_step_node import INode, NodeResult
+from common.util.field_message import ErrMessage
+
+
+class DocumentExtractNodeSerializer(serializers.Serializer):
+    # 需要查询的数据集id列表
+    file_list = serializers.ListField(required=True, child=serializers.UUIDField(required=True),
+                                      error_messages=ErrMessage.list("数据集id列表"))
+
+    def is_valid(self, *, raise_exception=False):
+        super().is_valid(raise_exception=True)
+
+
+class IDocumentExtractNode(INode):
+    type = 'document-extract-node'
+
+    def get_node_params_serializer_class(self) -> Type[serializers.Serializer]:
+        return DocumentExtractNodeSerializer
+
+    def _run(self):
+        return self.execute(**self.flow_params_serializer.data)
+
+    def execute(self, file_list, **kwargs) -> NodeResult:
+        pass

apps/application/flow/step_node/document_extract_node/impl/__init__.py

@@ -0,0 +1 @@
+from .base_document_extract_node import BaseDocumentExtractNode

apps/application/flow/step_node/document_extract_node/impl/base_document_extract_node.py

@@ -0,0 +1,11 @@
+# coding=utf-8
+
+from application.flow.step_node.document_extract_node.i_document_extract_node import IDocumentExtractNode
+
+
+class BaseDocumentExtractNode(IDocumentExtractNode):
+    def execute(self, file_list, **kwargs):
+        pass
+
+    def get_details(self, index: int, **kwargs):
+        pass

apps/application/flow/step_node/image_understand_step_node/__init__.py

@@ -0,0 +1,3 @@
+# coding=utf-8
+
+from .impl import *

apps/application/flow/step_node/image_understand_step_node/i_image_understand_node.py

@@ -0,0 +1,39 @@
+# coding=utf-8
+
+from typing import Type
+
+from rest_framework import serializers
+
+from application.flow.i_step_node import INode, NodeResult
+from common.util.field_message import ErrMessage
+
+
+class ImageUnderstandNodeSerializer(serializers.Serializer):
+    model_id = serializers.CharField(required=True, error_messages=ErrMessage.char("模型id"))
+    system = serializers.CharField(required=False, allow_blank=True, allow_null=True,
+                                   error_messages=ErrMessage.char("角色设定"))
+    prompt = serializers.CharField(required=True, error_messages=ErrMessage.char("提示词"))
+    # 多轮对话数量
+    dialogue_number = serializers.IntegerField(required=True, error_messages=ErrMessage.integer("多轮对话数量"))
+
+    is_result = serializers.BooleanField(required=False, error_messages=ErrMessage.boolean('是否返回内容'))
+
+    image_list = serializers.ListField(required=False, error_messages=ErrMessage.list("图片仅1张"))
+
+
+class IImageUnderstandNode(INode):
+    type = 'image-understand-node'
+
+    def get_node_params_serializer_class(self) -> Type[serializers.Serializer]:
+        return ImageUnderstandNodeSerializer
+
+    def _run(self):
+        res = self.workflow_manage.get_reference_field(self.node_params_serializer.data.get('image_list')[0],
+                                                       self.node_params_serializer.data.get('image_list')[1:])
+        return self.execute(image=res, **self.node_params_serializer.data, **self.flow_params_serializer.data)
+
+    def execute(self, model_id, system, prompt, dialogue_number, history_chat_record, stream, chat_id,
+                chat_record_id,
+                image,
+                **kwargs) -> NodeResult:
+        pass

apps/application/flow/step_node/image_understand_step_node/impl/__init__.py

@@ -0,0 +1,3 @@
+# coding=utf-8
+
+from .base_image_understand_node import BaseImageUnderstandNode

apps/application/flow/step_node/image_understand_step_node/impl/base_image_understand_node.py

@@ -0,0 +1,147 @@
+# coding=utf-8
+import base64
+import os
+import time
+from functools import reduce
+from typing import List, Dict
+
+from django.db.models import QuerySet
+from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
+
+from application.flow.i_step_node import NodeResult, INode
+from application.flow.step_node.image_understand_step_node.i_image_understand_node import IImageUnderstandNode
+from dataset.models import File
+from setting.models_provider.tools import get_model_instance_by_model_user_id
+
+
+def _write_context(node_variable: Dict, workflow_variable: Dict, node: INode, workflow, answer: str):
+    chat_model = node_variable.get('chat_model')
+    message_tokens = node_variable['usage_metadata']['output_tokens'] if 'usage_metadata' in node_variable else 0
+    answer_tokens = chat_model.get_num_tokens(answer)
+    node.context['message_tokens'] = message_tokens
+    node.context['answer_tokens'] = answer_tokens
+    node.context['answer'] = answer
+    node.context['history_message'] = node_variable['history_message']
+    node.context['question'] = node_variable['question']
+    node.context['run_time'] = time.time() - node.context['start_time']
+    if workflow.is_result(node, NodeResult(node_variable, workflow_variable)):
+        workflow.answer += answer
+
+
+def write_context_stream(node_variable: Dict, workflow_variable: Dict, node: INode, workflow):
+    """
+    写入上下文数据 (流式)
+    @param node_variable:      节点数据
+    @param workflow_variable:  全局数据
+    @param node:               节点
+    @param workflow:           工作流管理器
+    """
+    response = node_variable.get('result')
+    answer = ''
+    for chunk in response:
+        answer += chunk.content
+        yield chunk.content
+    _write_context(node_variable, workflow_variable, node, workflow, answer)
+
+
+def write_context(node_variable: Dict, workflow_variable: Dict, node: INode, workflow):
+    """
+    写入上下文数据
+    @param node_variable:      节点数据
+    @param workflow_variable:  全局数据
+    @param node:               节点实例对象
+    @param workflow:           工作流管理器
+    """
+    response = node_variable.get('result')
+    answer = response.content
+    _write_context(node_variable, workflow_variable, node, workflow, answer)
+
+
+class BaseImageUnderstandNode(IImageUnderstandNode):
+    def execute(self, model_id, system, prompt, dialogue_number, history_chat_record, stream, chat_id, chat_record_id,
+                image,
+                **kwargs) -> NodeResult:
+        image_model = get_model_instance_by_model_user_id(model_id, self.flow_params_serializer.data.get('user_id'))
+        history_message = self.get_history_message(history_chat_record, dialogue_number)
+        self.context['history_message'] = history_message
+        question = self.generate_prompt_question(prompt)
+        self.context['question'] = question.content
+        # todo 处理上传图片
+        message_list = self.generate_message_list(image_model, system, prompt, history_message, image)
+        self.context['message_list'] = message_list
+        self.context['image_list'] = image
+        if stream:
+            r = image_model.stream(message_list)
+            return NodeResult({'result': r, 'chat_model': image_model, 'message_list': message_list,
+                               'history_message': history_message, 'question': question.content}, {},
+                              _write_context=write_context_stream)
+        else:
+            r = image_model.invoke(message_list)
+            return NodeResult({'result': r, 'chat_model': image_model, 'message_list': message_list,
+                               'history_message': history_message, 'question': question.content}, {},
+                              _write_context=write_context)
+
+    @staticmethod
+    def get_history_message(history_chat_record, dialogue_number):
+        start_index = len(history_chat_record) - dialogue_number
+        history_message = reduce(lambda x, y: [*x, *y], [
+            [history_chat_record[index].get_human_message(), history_chat_record[index].get_ai_message()]
+            for index in
+            range(start_index if start_index > 0 else 0, len(history_chat_record))], [])
+        return history_message
+
+    def generate_prompt_question(self, prompt):
+        return HumanMessage(self.workflow_manage.generate_prompt(prompt))
+
+    def generate_message_list(self, image_model, system: str, prompt: str, history_message, image):
+        if image is not None and len(image) > 0:
+            file_id = image[0]['file_id']
+            file = QuerySet(File).filter(id=file_id).first()
+            base64_image = base64.b64encode(file.get_byte()).decode("utf-8")
+            messages = [HumanMessage(
+                content=[
+                    {'type': 'text', 'text': self.workflow_manage.generate_prompt(prompt)},
+                    {'type': 'image_url', 'image_url': {'url': f'data:image/jpeg;base64,{base64_image}'}},
+                ])]
+        else:
+            messages = [HumanMessage(self.workflow_manage.generate_prompt(prompt))]
+
+        if system is not None and len(system) > 0:
+            return [
+                SystemMessage(self.workflow_manage.generate_prompt(system)),
+                *history_message,
+                *messages
+            ]
+        else:
+            return [
+                *history_message,
+                *messages
+            ]
+
+    @staticmethod
+    def reset_message_list(message_list: List[BaseMessage], answer_text):
+        result = [{'role': 'user' if isinstance(message, HumanMessage) else 'ai', 'content': message.content} for
+                  message
+                  in
+                  message_list]
+        result.append({'role': 'ai', 'content': answer_text})
+        return result
+
+    def get_details(self, index: int, **kwargs):
+        return {
+            'name': self.node.properties.get('stepName'),
+            "index": index,
+            'run_time': self.context.get('run_time'),
+            'system': self.node_params.get('system'),
+            'history_message': [{'content': message.content, 'role': message.type} for message in
+                                (self.context.get('history_message') if self.context.get(
+                                    'history_message') is not None else [])],
+            'question': self.context.get('question'),
+            'answer': self.context.get('answer'),
+            'type': self.node.type,
+            'message_tokens': self.context.get('message_tokens'),
+            'answer_tokens': self.context.get('answer_tokens'),
+            'status': self.status,
+            'err_message': self.err_message,
+            'image_list': self.context.get('image_list')
+        }

apps/application/flow/step_node/start_node/impl/base_start_node.py

@@ -41,7 +41,7 @@ def execute(self, question, **kwargs) -> NodeResult:
         """
         开始节点 初始化全局变量
         """
-        return NodeResult({'question': question},
+        return NodeResult({'question': question, 'image': self.workflow_manage.image_list},
                           workflow_variable)
 
     def get_details(self, index: int, **kwargs):
@@ -61,5 +61,6 @@ def get_details(self, index: int, **kwargs):
             'type': self.node.type,
             'status': self.status,
             'err_message': self.err_message,
+            'image_list': self.context.get('image'),
             'global_fields': global_fields
         }