feat: 支持添加图片生成模型(WIP)

Author: liuruibin

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,37 @@
+# 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('是否返回内容'))
+
+
+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,
+                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,100 @@
+# coding=utf-8
+from functools import reduce
+from typing import List
+
+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 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,
+                chat_record_id,
+                **kwargs) -> NodeResult:
+
+        tti_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
+        message_list = self.generate_message_list(question, history_message)
+        self.context['message_list'] = message_list
+        self.context['dialogue_type'] = dialogue_type
+        print(message_list)
+        print(negative_prompt)
+        image_urls = tti_model.generate_image(question, negative_prompt)
+        self.context['image_list'] = image_urls
+        answer = '\n'.join([f"![Image]({path})" for path in image_urls])
+        return NodeResult({'answer': answer, 'chat_model': tti_model, 'message_list': message_list,
+                           '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/workflow_manage.py

@@ -53,7 +53,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/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/tti.py

@@ -0,0 +1,47 @@
+# coding=utf-8
+import base64
+import os
+from typing import Dict
+
+from langchain_core.messages import HumanMessage
+
+from common import forms
+from common.exception.app_exception import AppApiException
+from common.forms import BaseForm
+from setting.models_provider.base_model_provider import BaseModelCredential, ValidCode
+
+
+class OpenAITextToImageModelCredential(BaseForm, BaseModelCredential):
+    api_base = forms.TextInputField('API 域名', required=True)
+    api_key = forms.PasswordInputField('API Key', required=True)
+
+    def is_valid(self, model_type: str, model_name, model_credential: Dict[str, object], provider,
+                 raise_exception=False):
+        model_type_list = provider.get_model_type_list()
+        if not any(list(filter(lambda mt: mt.get('value') == model_type, model_type_list))):
+            raise AppApiException(ValidCode.valid_error.value, f'{model_type} 模型类型不支持')
+
+        for key in ['api_base', 'api_key']:
+            if key not in model_credential:
+                if raise_exception:
+                    raise AppApiException(ValidCode.valid_error.value, f'{key} 字段为必填字段')
+                else:
+                    return False
+        try:
+            model = provider.get_model(model_type, model_name, model_credential)
+            res = model.check_auth()
+            print(res)
+        except Exception as e:
+            if isinstance(e, AppApiException):
+                raise e
+            if raise_exception:
+                raise AppApiException(ValidCode.valid_error.value, f'校验失败,请检查参数是否正确: {str(e)}')
+            else:
+                return False
+        return True
+
+    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

apps/setting/models_provider/impl/openai_model_provider/model/tti.py

@@ -0,0 +1,67 @@
+from typing import Dict
+
+import requests
+from langchain_core.messages import HumanMessage
+from langchain_openai import ChatOpenAI
+from openai import OpenAI
+
+from common.config.tokenizer_manage_config import TokenizerManage
+from common.util.common import bytes_to_uploaded_file
+from dataset.serializers.file_serializers import FileSerializer
+from setting.models_provider.base_model_provider import MaxKBBaseModel
+from setting.models_provider.impl.base_tti import BaseTextToImage
+
+
+def custom_get_token_ids(text: str):
+    tokenizer = TokenizerManage.get_tokenizer()
+    return tokenizer.encode(text)
+
+
+class OpenAITextToImage(MaxKBBaseModel, BaseTextToImage):
+    api_base: str
+    api_key: str
+    model: str
+    params: dict
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.api_key = kwargs.get('api_key')
+        self.api_base = kwargs.get('api_base')
+        self.model = kwargs.get('model')
+        self.params = kwargs.get('params')
+
+    @staticmethod
+    def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
+        optional_params = {'params': {}}
+        for key, value in model_kwargs.items():
+            if key not in ['model_id', 'use_local', 'streaming']:
+                optional_params['params'][key] = value
+        return OpenAITextToImage(
+            model=model_name,
+            api_base=model_credential.get('api_base'),
+            api_key=model_credential.get('api_key'),
+            **optional_params,
+        )
+
+    def check_auth(self):
+        chat = OpenAI(api_key=self.api_key, base_url=self.api_base)
+        response_list = chat.models.with_raw_response.list()
+
+        # self.generate_image('生成一个小猫图片')
+
+    def generate_image(self, prompt: str, negative_prompt: str = None):
+
+        chat = OpenAI(api_key=self.api_key, base_url=self.api_base)
+        res = chat.images.generate(model=self.model, prompt=prompt, **self.params)
+
+        file_urls = []
+        for content in res.data:
+            url = content.url
+            print(url)
+            file_name = 'generated_image.png'
+            file = bytes_to_uploaded_file(requests.get(url).content, file_name)
+            meta = {'debug': True}
+            file_url = FileSerializer(data={'file': file, 'meta': meta}).upload()
+            file_urls.append(file_url)
+
+        return file_urls