feat: Support image generate model

apps/application/flow/step_node/__init__.py

@@ -18,14 +18,15 @@
 
 from .document_extract_node import *
 from .image_understand_step_node import *
+from .image_generate_step_node import *
 
 from .search_dataset_node import *
 from .start_node import *
 
 node_list = [BaseStartStepNode, BaseChatNode, BaseSearchDatasetNode, BaseQuestionNode, BaseConditionNode, BaseReplyNode,
              BaseFunctionNodeNode, BaseFunctionLibNodeNode, BaseRerankerNode, BaseApplicationNode,
              BaseDocumentExtractNode,
-             BaseImageUnderstandNode, BaseFormNode]
+             BaseImageUnderstandNode, BaseImageGenerateNode, BaseFormNode]
 
 
 def get_node(node_type):

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

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

apps/application/flow/step_node/image_generate_step_node/i_image_generate_node.py

@@ -0,0 +1,40 @@
+# 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 ImageGenerateNodeSerializer(serializers.Serializer):
+    model_id = serializers.CharField(required=True, error_messages=ErrMessage.char("模型id"))
+
+    prompt = serializers.CharField(required=True, error_messages=ErrMessage.char("提示词(正向)"))
+
+    negative_prompt = serializers.CharField(required=False, default='', error_messages=ErrMessage.char("提示词(负向)"))
+    # 多轮对话数量
+    dialogue_number = serializers.IntegerField(required=True, error_messages=ErrMessage.integer("多轮对话数量"))
+
+    dialogue_type = serializers.CharField(required=True, error_messages=ErrMessage.char("对话存储类型"))
+
+    is_result = serializers.BooleanField(required=False, error_messages=ErrMessage.boolean('是否返回内容'))
+
+    model_params_setting = serializers.JSONField(required=False, default=dict, error_messages=ErrMessage.json("模型参数设置"))
+
+
+class IImageGenerateNode(INode):
+    type = 'image-generate-node'
+
+    def get_node_params_serializer_class(self) -> Type[serializers.Serializer]:
+        return ImageGenerateNodeSerializer
+
+    def _run(self):
+        return self.execute(**self.node_params_serializer.data, **self.flow_params_serializer.data)
+
+    def execute(self, model_id, prompt, negative_prompt, dialogue_number, dialogue_type, history_chat_record, chat_id,
+                model_params_setting,
+                chat_record_id,
+                **kwargs) -> NodeResult:
+        pass

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

@@ -0,0 +1,3 @@
+# coding=utf-8
+
+from .base_image_generate_node import BaseImageGenerateNode

apps/application/flow/step_node/image_generate_step_node/impl/base_image_generate_node.py

@@ -0,0 +1,117 @@
+# coding=utf-8
+from functools import reduce
+from typing import List
+
+import requests
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
+
+from application.flow.i_step_node import NodeResult
+from application.flow.step_node.image_generate_step_node.i_image_generate_node import IImageGenerateNode
+from common.util.common import bytes_to_uploaded_file
+from dataset.serializers.file_serializers import FileSerializer
+from setting.models_provider.tools import get_model_instance_by_model_user_id
+
+
+class BaseImageGenerateNode(IImageGenerateNode):
+    def save_context(self, details, workflow_manage):
+        self.context['answer'] = details.get('answer')
+        self.context['question'] = details.get('question')
+        self.answer_text = details.get('answer')
+
+    def execute(self, model_id, prompt, negative_prompt, dialogue_number, dialogue_type, history_chat_record, chat_id,
+                model_params_setting,
+                chat_record_id,
+                **kwargs) -> NodeResult:
+        print(model_params_setting)
+        application = self.workflow_manage.work_flow_post_handler.chat_info.application
+        tti_model = get_model_instance_by_model_user_id(model_id, self.flow_params_serializer.data.get('user_id'), **model_params_setting)
+        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
+        message_list = self.generate_message_list(question, history_message)
+        self.context['message_list'] = message_list
+        self.context['dialogue_type'] = dialogue_type
+        print(message_list)
+        image_urls = tti_model.generate_image(question, negative_prompt)
+        # 保存图片
+        file_urls = []
+        for image_url in image_urls:
+            file_name = 'generated_image.png'
+            file = bytes_to_uploaded_file(requests.get(image_url).content, file_name)
+            meta = {
+                'debug': False if application.id else True,
+                'chat_id': chat_id,
+                'application_id': str(application.id) if application.id else None,
+            }
+            file_url = FileSerializer(data={'file': file, 'meta': meta}).upload()
+            file_urls.append(file_url)
+        self.context['image_list'] = file_urls
+        answer = '\n'.join([f"![Image]({path})" for path in file_urls])
+        return NodeResult({'answer': answer, 'chat_model': tti_model, 'message_list': message_list,
+                           'image': [{'file_id': path.split('/')[-1], 'file_url': path} for path in file_urls],
+                           'history_message': history_message, 'question': question}, {})
+
+    def generate_history_ai_message(self, chat_record):
+        for val in chat_record.details.values():
+            if self.node.id == val['node_id'] and 'image_list' in val:
+                if val['dialogue_type'] == 'WORKFLOW':
+                    return chat_record.get_ai_message()
+                return AIMessage(content=val['answer'])
+        return chat_record.get_ai_message()
+
+    def get_history_message(self, history_chat_record, dialogue_number):
+        start_index = len(history_chat_record) - dialogue_number
+        history_message = reduce(lambda x, y: [*x, *y], [
+            [self.generate_history_human_message(history_chat_record[index]),
+             self.generate_history_ai_message(history_chat_record[index])]
+            for index in
+            range(start_index if start_index > 0 else 0, len(history_chat_record))], [])
+        return history_message
+
+    def generate_history_human_message(self, chat_record):
+
+        for data in chat_record.details.values():
+            if self.node.id == data['node_id'] and 'image_list' in data:
+                image_list = data['image_list']
+                if len(image_list) == 0 or data['dialogue_type'] == 'WORKFLOW':
+                    return HumanMessage(content=chat_record.problem_text)
+                return HumanMessage(content=data['question'])
+        return HumanMessage(content=chat_record.problem_text)
+
+    def generate_prompt_question(self, prompt):
+        return self.workflow_manage.generate_prompt(prompt)
+
+    def generate_message_list(self, question: str, history_message):
+        return [
+            *history_message,
+            question
+        ]
+
+    @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'),
+            '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'),
+            'dialogue_type': self.context.get('dialogue_type')
+        }

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

@@ -22,6 +22,9 @@ class ImageUnderstandNodeSerializer(serializers.Serializer):
 
     image_list = serializers.ListField(required=False, error_messages=ErrMessage.list("图片"))
 
+    model_params_setting = serializers.JSONField(required=False, default=dict, error_messages=ErrMessage.json("模型参数设置"))
+
+
 
 class IImageUnderstandNode(INode):
     type = 'image-understand-node'
@@ -35,6 +38,7 @@ def _run(self):
         return self.execute(image=res, **self.node_params_serializer.data, **self.flow_params_serializer.data)
 
     def execute(self, model_id, system, prompt, dialogue_number, dialogue_type, history_chat_record, stream, chat_id,
+                model_params_setting,
                 chat_record_id,
                 image,
                 **kwargs) -> NodeResult:

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

@@ -12,6 +12,7 @@
 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
+from imghdr import what
 
 
 def _write_context(node_variable: Dict, workflow_variable: Dict, node: INode, workflow, answer: str):
@@ -59,8 +60,9 @@ def write_context(node_variable: Dict, workflow_variable: Dict, node: INode, wor
 
 def file_id_to_base64(file_id: str):
     file = QuerySet(File).filter(id=file_id).first()
-    base64_image = base64.b64encode(file.get_byte()).decode("utf-8")
-    return base64_image
+    file_bytes = file.get_byte()
+    base64_image = base64.b64encode(file_bytes).decode("utf-8")
+    return [base64_image, what(None, file_bytes.tobytes())]
 
 
 class BaseImageUnderstandNode(IImageUnderstandNode):
@@ -70,14 +72,15 @@ def save_context(self, details, workflow_manage):
         self.answer_text = details.get('answer')
 
     def execute(self, model_id, system, prompt, dialogue_number, dialogue_type, history_chat_record, stream, chat_id,
+                model_params_setting,
                 chat_record_id,
                 image,
                 **kwargs) -> NodeResult:
         # 处理不正确的参数
         if image is None or not isinstance(image, list):
             image = []
-
-        image_model = get_model_instance_by_model_user_id(model_id, self.flow_params_serializer.data.get('user_id'))
+        print(model_params_setting)
+        image_model = get_model_instance_by_model_user_id(model_id, self.flow_params_serializer.data.get('user_id'), **model_params_setting)
         # 执行详情中的历史消息不需要图片内容
         history_message = self.get_history_message_for_details(history_chat_record, dialogue_number)
         self.context['history_message'] = history_message
@@ -151,7 +154,7 @@ def generate_history_human_message(self, chat_record):
                 return HumanMessage(
                     content=[
                         {'type': 'text', 'text': data['question']},
-                        *[{'type': 'image_url', 'image_url': {'url': f'data:image/jpeg;base64,{base64_image}'}} for
+                        *[{'type': 'image_url', 'image_url': {'url': f'data:image/{base64_image[1]};base64,{base64_image[0]}'}} for
                           base64_image in image_base64_list]
                     ])
         return HumanMessage(content=chat_record.problem_text)
@@ -166,8 +169,10 @@ def generate_message_list(self, image_model, system: str, prompt: str, history_m
             for img in image:
                 file_id = img['file_id']
                 file = QuerySet(File).filter(id=file_id).first()
-                base64_image = base64.b64encode(file.get_byte()).decode("utf-8")
-                images.append({'type': 'image_url', 'image_url': {'url': f'data:image/jpeg;base64,{base64_image}'}})
+                image_bytes = file.get_byte()
+                base64_image = base64.b64encode(image_bytes).decode("utf-8")
+                image_format = what(None, image_bytes.tobytes())
+                images.append({'type': 'image_url', 'image_url': {'url': f'data:image/{image_format};base64,{base64_image}'}})
             messages = [HumanMessage(
                 content=[
                     {'type': 'text', 'text': self.workflow_manage.generate_prompt(prompt)},

apps/application/flow/workflow_manage.py

@@ -54,7 +54,7 @@ def __init__(self, _id: str, _type: str, x: int, y: int, properties: dict, **kwa
 
 
 end_nodes = ['ai-chat-node', 'reply-node', 'function-node', 'function-lib-node', 'application-node',
-             'image-understand-node']
+             'image-understand-node', 'image-generate-node']
 
 
 class Flow:

apps/common/forms/text_input_field.py

@@ -8,6 +8,7 @@
 """
 from typing import Dict
 
+from common.forms import BaseLabel
 from common.forms.base_field import BaseField, TriggerType
 
 
@@ -16,7 +17,7 @@ class TextInputField(BaseField):
     文本输入框
     """
 
-    def __init__(self, label: str,
+    def __init__(self, label: str or BaseLabel,
                  required: bool = False,
                  default_value=None,
                  relation_show_field_dict: Dict = None,

apps/common/util/common.py

@@ -8,9 +8,12 @@
 """
 import hashlib
 import importlib
+import mimetypes
+import io
 from functools import reduce
 from typing import Dict, List
 
+from django.core.files.uploadedfile import InMemoryUploadedFile
 from django.db.models import QuerySet
 
 from ..exception.app_exception import AppApiException
@@ -111,3 +114,25 @@ def bulk_create_in_batches(model, data, batch_size=1000):
         batch = data[i:i + batch_size]
         model.objects.bulk_create(batch)
 
+
+def bytes_to_uploaded_file(file_bytes, file_name="file.txt"):
+    content_type, _ = mimetypes.guess_type(file_name)
+    if content_type is None:
+        # 如果未能识别，设置为默认的二进制文件类型
+        content_type = "application/octet-stream"
+    # 创建一个内存中的字节流对象
+    file_stream = io.BytesIO(file_bytes)
+
+    # 获取文件大小
+    file_size = len(file_bytes)
+
+    # 创建 InMemoryUploadedFile 对象
+    uploaded_file = InMemoryUploadedFile(
+        file=file_stream,
+        field_name=None,
+        name=file_name,
+        content_type=content_type,
+        size=file_size,
+        charset=None,
+    )
+    return uploaded_file

apps/setting/models_provider/base_model_provider.py

@@ -150,6 +150,7 @@ class ModelTypeConst(Enum):
     STT = {'code': 'STT', 'message': '语音识别'}
     TTS = {'code': 'TTS', 'message': '语音合成'}
     IMAGE = {'code': 'IMAGE', 'message': '图片理解'}
+    TTI = {'code': 'TTI', 'message': '图片生成'}
     RERANKER = {'code': 'RERANKER', 'message': '重排模型'}
 
 

apps/setting/models_provider/impl/base_tti.py

@@ -0,0 +1,14 @@
+# coding=utf-8
+from abc import abstractmethod
+
+from pydantic import BaseModel
+
+
+class BaseTextToImage(BaseModel):
+    @abstractmethod
+    def check_auth(self):
+        pass
+
+    @abstractmethod
+    def generate_image(self, prompt: str, negative_prompt: str = None):
+        pass

apps/setting/models_provider/impl/openai_model_provider/credential/image.py

@@ -7,9 +7,26 @@
 
 from common import forms
 from common.exception.app_exception import AppApiException
-from common.forms import BaseForm
+from common.forms import BaseForm, TooltipLabel
 from setting.models_provider.base_model_provider import BaseModelCredential, ValidCode
 
+class OpenAIImageModelParams(BaseForm):
+    temperature = forms.SliderField(TooltipLabel('温度', '较高的数值会使输出更加随机，而较低的数值会使其更加集中和确定'),
+                                    required=True, default_value=0.7,
+                                    _min=0.1,
+                                    _max=1.0,
+                                    _step=0.01,
+                                    precision=2)
+
+    max_tokens = forms.SliderField(
+        TooltipLabel('输出最大Tokens', '指定模型可生成的最大token个数'),
+        required=True, default_value=800,
+        _min=1,
+        _max=100000,
+        _step=1,
+        precision=0)
+
+
 
 class OpenAIImageModelCredential(BaseForm, BaseModelCredential):
     api_base = forms.TextInputField('API 域名', required=True)
@@ -45,4 +62,4 @@ def encryption_dict(self, model: Dict[str, object]):
         return {**model, 'api_key': super().encryption(model.get('api_key', ''))}
 
     def get_model_params_setting_form(self, model_name):
-        pass
+        return OpenAIImageModelParams()