feat: [backend] vision support (#1510)

Co-authored-by: Garfield Dai <dai.hai@foxmail.com>

api/core/callback_handler/llm_callback_handler.py

@@ -11,7 +11,8 @@ from pydantic import BaseModel
 from core.callback_handler.entity.llm_message import LLMMessage
 from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException, \
     ConversationTaskInterruptException
-from core.model_providers.models.entity.message import to_prompt_messages, PromptMessage
+from core.model_providers.models.entity.message import to_prompt_messages, PromptMessage, LCHumanMessageWithFiles, \
+    ImagePromptMessageFile
 from core.model_providers.models.llm.base import BaseLLM
 from core.moderation.base import ModerationOutputsResult, ModerationAction
 from core.moderation.factory import ModerationFactory
@@ -72,7 +73,12 @@ class LLMCallbackHandler(BaseCallbackHandler):
 
             real_prompts.append({
                 "role": role,
-                "text": message.content
+                "text": message.content,
+                "files": [{
+                    "type": file.type.value,
+                    "data": file.data[:10] + '...[TRUNCATED]...' + file.data[-10:],
+                    "detail": file.detail.value if isinstance(file, ImagePromptMessageFile) else None,
+                } for file in (message.files if isinstance(message, LCHumanMessageWithFiles) else [])]
             })
 
         self.llm_message.prompt = real_prompts

api/core/completion.py

@@ -13,11 +13,12 @@ from core.callback_handler.llm_callback_handler import LLMCallbackHandler
 from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException, \
     ConversationTaskInterruptException
 from core.external_data_tool.factory import ExternalDataToolFactory
+from core.file.file_obj import FileObj
 from core.model_providers.error import LLMBadRequestError
 from core.memory.read_only_conversation_token_db_buffer_shared_memory import \
     ReadOnlyConversationTokenDBBufferSharedMemory
 from core.model_providers.model_factory import ModelFactory
-from core.model_providers.models.entity.message import PromptMessage
+from core.model_providers.models.entity.message import PromptMessage, PromptMessageFile
 from core.model_providers.models.llm.base import BaseLLM
 from core.orchestrator_rule_parser import OrchestratorRuleParser
 from core.prompt.prompt_template import PromptTemplateParser
@@ -30,8 +31,9 @@ from core.moderation.factory import ModerationFactory
 class Completion:
     @classmethod
     def generate(cls, task_id: str, app: App, app_model_config: AppModelConfig, query: str, inputs: dict,
-                 user: Union[Account, EndUser], conversation: Optional[Conversation], streaming: bool,
-                 is_override: bool = False, retriever_from: str = 'dev'):
+                 files: List[FileObj], user: Union[Account, EndUser], conversation: Optional[Conversation],
+                 streaming: bool, is_override: bool = False, retriever_from: str = 'dev',
+                 auto_generate_name: bool = True):
         """
         errors: ProviderTokenNotInitError
         """
@@ -64,16 +66,21 @@ class Completion:
             is_override=is_override,
             inputs=inputs,
             query=query,
+            files=files,
             streaming=streaming,
-            model_instance=final_model_instance
+            model_instance=final_model_instance,
+            auto_generate_name=auto_generate_name
         )
 
+        prompt_message_files = [file.prompt_message_file for file in files]
+
         rest_tokens_for_context_and_memory = cls.get_validate_rest_tokens(
             mode=app.mode,
             model_instance=final_model_instance,
             app_model_config=app_model_config,
             query=query,
-            inputs=inputs
+            inputs=inputs,
+            files=prompt_message_files
         )
 
         # init orchestrator rule parser
@@ -95,6 +102,7 @@ class Completion:
                     app_model_config=app_model_config,
                     query=query,
                     inputs=inputs,
+                    files=prompt_message_files,
                     agent_execute_result=None,
                     conversation_message_task=conversation_message_task,
                     memory=memory,
@@ -146,6 +154,7 @@ class Completion:
                 app_model_config=app_model_config,
                 query=query,
                 inputs=inputs,
+                files=prompt_message_files,
                 agent_execute_result=agent_execute_result,
                 conversation_message_task=conversation_message_task,
                 memory=memory,
@@ -257,6 +266,7 @@ class Completion:
     @classmethod
     def run_final_llm(cls, model_instance: BaseLLM, mode: str, app_model_config: AppModelConfig, query: str,
                       inputs: dict,
+                      files: List[PromptMessageFile],
                       agent_execute_result: Optional[AgentExecuteResult],
                       conversation_message_task: ConversationMessageTask,
                       memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory],
@@ -266,10 +276,12 @@ class Completion:
         # get llm prompt
         if app_model_config.prompt_type == 'simple':
             prompt_messages, stop_words = prompt_transform.get_prompt(
-                mode=mode,
+                app_mode=mode,
+                app_model_config=app_model_config,
                 pre_prompt=app_model_config.pre_prompt,
                 inputs=inputs,
                 query=query,
+                files=files,
                 context=agent_execute_result.output if agent_execute_result else None,
                 memory=memory,
                 model_instance=model_instance
@@ -280,6 +292,7 @@ class Completion:
                 app_model_config=app_model_config,
                 inputs=inputs,
                 query=query,
+                files=files,
                 context=agent_execute_result.output if agent_execute_result else None,
                 memory=memory,
                 model_instance=model_instance
@@ -337,7 +350,7 @@ class Completion:
 
     @classmethod
     def get_validate_rest_tokens(cls, mode: str, model_instance: BaseLLM, app_model_config: AppModelConfig,
-                                 query: str, inputs: dict) -> int:
+                                 query: str, inputs: dict, files: List[PromptMessageFile]) -> int:
         model_limited_tokens = model_instance.model_rules.max_tokens.max
         max_tokens = model_instance.get_model_kwargs().max_tokens
 
@@ -348,15 +361,16 @@ class Completion:
             max_tokens = 0
 
         prompt_transform = PromptTransform()
-        prompt_messages = []
 
         # get prompt without memory and context
         if app_model_config.prompt_type == 'simple':
             prompt_messages, _ = prompt_transform.get_prompt(
-                mode=mode,
+                app_mode=mode,
+                app_model_config=app_model_config,
                 pre_prompt=app_model_config.pre_prompt,
                 inputs=inputs,
                 query=query,
+                files=files,
                 context=None,
                 memory=None,
                 model_instance=model_instance
@@ -367,6 +381,7 @@ class Completion:
                 app_model_config=app_model_config,
                 inputs=inputs,
                 query=query,
+                files=files,
                 context=None,
                 memory=None,
                 model_instance=model_instance

api/core/conversation_message_task.py

@@ -6,8 +6,9 @@ from core.callback_handler.entity.agent_loop import AgentLoop
 from core.callback_handler.entity.dataset_query import DatasetQueryObj
 from core.callback_handler.entity.llm_message import LLMMessage
 from core.callback_handler.entity.chain_result import ChainResult
+from core.file.file_obj import FileObj
 from core.model_providers.model_factory import ModelFactory
-from core.model_providers.models.entity.message import to_prompt_messages, MessageType
+from core.model_providers.models.entity.message import to_prompt_messages, MessageType, PromptMessageFile
 from core.model_providers.models.llm.base import BaseLLM
 from core.prompt.prompt_builder import PromptBuilder
 from core.prompt.prompt_template import PromptTemplateParser
@@ -16,13 +17,14 @@ from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import DatasetQuery
 from models.model import AppModelConfig, Conversation, Account, Message, EndUser, App, MessageAgentThought, \
-    MessageChain, DatasetRetrieverResource
+    MessageChain, DatasetRetrieverResource, MessageFile
 
 
 class ConversationMessageTask:
     def __init__(self, task_id: str, app: App, app_model_config: AppModelConfig, user: Account,
-                 inputs: dict, query: str, streaming: bool, model_instance: BaseLLM,
-                 conversation: Optional[Conversation] = None, is_override: bool = False):
+                 inputs: dict, query: str, files: List[FileObj], streaming: bool,
+                 model_instance: BaseLLM, conversation: Optional[Conversation] = None, is_override: bool = False,
+                 auto_generate_name: bool = True):
         self.start_at = time.perf_counter()
 
         self.task_id = task_id
@@ -35,6 +37,7 @@ class ConversationMessageTask:
         self.user = user
         self.inputs = inputs
         self.query = query
+        self.files = files
         self.streaming = streaming
 
         self.conversation = conversation
@@ -45,6 +48,7 @@ class ConversationMessageTask:
         self.message = None
 
         self.retriever_resource = None
+        self.auto_generate_name = auto_generate_name
 
         self.model_dict = self.app_model_config.model_dict
         self.provider_name = self.model_dict.get('provider')
@@ -100,7 +104,7 @@ class ConversationMessageTask:
                 model_id=self.model_name,
                 override_model_configs=json.dumps(override_model_configs) if override_model_configs else None,
                 mode=self.mode,
-                name='',
+                name='New conversation',
                 inputs=self.inputs,
                 introduction=introduction,
                 system_instruction=system_instruction,
@@ -142,6 +146,19 @@ class ConversationMessageTask:
         db.session.add(self.message)
         db.session.commit()
 
+        for file in self.files:
+            message_file = MessageFile(
+                message_id=self.message.id,
+                type=file.type.value,
+                transfer_method=file.transfer_method.value,
+                url=file.url,
+                upload_file_id=file.upload_file_id,
+                created_by_role=('account' if isinstance(self.user, Account) else 'end_user'),
+                created_by=self.user.id
+            )
+            db.session.add(message_file)
+            db.session.commit()
+
     def append_message_text(self, text: str):
         if text is not None:
             self._pub_handler.pub_text(text)
@@ -176,7 +193,8 @@ class ConversationMessageTask:
         message_was_created.send(
             self.message,
             conversation=self.conversation,
-            is_first_message=self.is_new_conversation
+            is_first_message=self.is_new_conversation,
+            auto_generate_name=self.auto_generate_name
         )
 
         if not by_stopped:

api/core/file/file_obj.py

@@ -0,0 +1,79 @@
+import enum
+from typing import Optional
+
+from pydantic import BaseModel
+
+from core.file.upload_file_parser import UploadFileParser
+from core.model_providers.models.entity.message import PromptMessageFile, ImagePromptMessageFile
+from extensions.ext_database import db
+from models.model import UploadFile
+
+
+class FileType(enum.Enum):
+    IMAGE = 'image'
+
+    @staticmethod
+    def value_of(value):
+        for member in FileType:
+            if member.value == value:
+                return member
+        raise ValueError(f"No matching enum found for value '{value}'")
+
+
+class FileTransferMethod(enum.Enum):
+    REMOTE_URL = 'remote_url'
+    LOCAL_FILE = 'local_file'
+
+    @staticmethod
+    def value_of(value):
+        for member in FileTransferMethod:
+            if member.value == value:
+                return member
+        raise ValueError(f"No matching enum found for value '{value}'")
+
+
+class FileObj(BaseModel):
+    id: Optional[str]
+    tenant_id: str
+    type: FileType
+    transfer_method: FileTransferMethod
+    url: Optional[str]
+    upload_file_id: Optional[str]
+    file_config: dict
+
+    @property
+    def data(self) -> Optional[str]:
+        return self._get_data()
+
+    @property
+    def preview_url(self) -> Optional[str]:
+        return self._get_data(force_url=True)
+
+    @property
+    def prompt_message_file(self) -> PromptMessageFile:
+        if self.type == FileType.IMAGE:
+            image_config = self.file_config.get('image')
+
+            return ImagePromptMessageFile(
+                data=self.data,
+                detail=ImagePromptMessageFile.DETAIL.HIGH
+                if image_config.get("detail") == "high" else ImagePromptMessageFile.DETAIL.LOW
+            )
+
+    def _get_data(self, force_url: bool = False) -> Optional[str]:
+        if self.type == FileType.IMAGE:
+            if self.transfer_method == FileTransferMethod.REMOTE_URL:
+                return self.url
+            elif self.transfer_method == FileTransferMethod.LOCAL_FILE:
+                upload_file = (db.session.query(UploadFile)
+                               .filter(
+                    UploadFile.id == self.upload_file_id,
+                    UploadFile.tenant_id == self.tenant_id
+                ).first())
+
+                return UploadFileParser.get_image_data(
+                    upload_file=upload_file,
+                    force_url=force_url
+                )
+
+        return None

api/core/file/message_file_parser.py

@@ -0,0 +1,180 @@
+from typing import List, Union, Optional, Dict
+
+import requests
+
+from core.file.file_obj import FileObj, FileType, FileTransferMethod
+from core.file.upload_file_parser import SUPPORT_EXTENSIONS
+from extensions.ext_database import db
+from models.account import Account
+from models.model import MessageFile, EndUser, AppModelConfig, UploadFile
+
+
+class MessageFileParser:
+
+    def __init__(self, tenant_id: str, app_id: str) -> None:
+        self.tenant_id = tenant_id
+        self.app_id = app_id
+
+    def validate_and_transform_files_arg(self, files: List[dict], app_model_config: AppModelConfig,
+                                         user: Union[Account, EndUser]) -> List[FileObj]:
+        """
+        validate and transform files arg
+
+        :param files:
+        :param app_model_config:
+        :param user:
+        :return:
+        """
+        file_upload_config = app_model_config.file_upload_dict
+
+        for file in files:
+            if not isinstance(file, dict):
+                raise ValueError('Invalid file format, must be dict')
+            if not file.get('type'):
+                raise ValueError('Missing file type')
+            FileType.value_of(file.get('type'))
+            if not file.get('transfer_method'):
+                raise ValueError('Missing file transfer method')
+            FileTransferMethod.value_of(file.get('transfer_method'))
+            if file.get('transfer_method') == FileTransferMethod.REMOTE_URL.value:
+                if not file.get('url'):
+                    raise ValueError('Missing file url')
+                if not file.get('url').startswith('http'):
+                    raise ValueError('Invalid file url')
+            if file.get('transfer_method') == FileTransferMethod.LOCAL_FILE.value and not file.get('upload_file_id'):
+                raise ValueError('Missing file upload_file_id')
+
+        # transform files to file objs
+        type_file_objs = self._to_file_objs(files, file_upload_config)
+
+        # validate files
+        new_files = []
+        for file_type, file_objs in type_file_objs.items():
+            if file_type == FileType.IMAGE:
+                # parse and validate files
+                image_config = file_upload_config.get('image')
+
+                # check if image file feature is enabled
+                if not image_config['enabled']:
+                    continue
+
+                # Validate number of files
+                if len(files) > image_config['number_limits']:
+                    raise ValueError(f"Number of image files exceeds the maximum limit {image_config['number_limits']}")
+
+                for file_obj in file_objs:
+                    # Validate transfer method
+                    if file_obj.transfer_method.value not in image_config['transfer_methods']:
+                        raise ValueError(f'Invalid transfer method: {file_obj.transfer_method.value}')
+
+                    # Validate file type
+                    if file_obj.type != FileType.IMAGE:
+                        raise ValueError(f'Invalid file type: {file_obj.type}')
+
+                    if file_obj.transfer_method == FileTransferMethod.REMOTE_URL:
+                        # check remote url valid and is image
+                        result, error = self._check_image_remote_url(file_obj.url)
+                        if result is False:
+                            raise ValueError(error)
+                    elif file_obj.transfer_method == FileTransferMethod.LOCAL_FILE:
+                        # get upload file from upload_file_id
+                        upload_file = (db.session.query(UploadFile)
+                                       .filter(
+                            UploadFile.id == file_obj.upload_file_id,
+                            UploadFile.tenant_id == self.tenant_id,
+                            UploadFile.created_by == user.id,
+                            UploadFile.created_by_role == ('account' if isinstance(user, Account) else 'end_user'),
+                            UploadFile.extension.in_(SUPPORT_EXTENSIONS)
+                        ).first())
+
+                        # check upload file is belong to tenant and user
+                        if not upload_file:
+                            raise ValueError('Invalid upload file')
+
+                    new_files.append(file_obj)
+
+        # return all file objs
+        return new_files
+
+    def transform_message_files(self, files: List[MessageFile], app_model_config: Optional[AppModelConfig]) -> List[FileObj]:
+        """
+        transform message files
+
+        :param files:
+        :param app_model_config:
+        :return:
+        """
+        # transform files to file objs
+        type_file_objs = self._to_file_objs(files, app_model_config.file_upload_dict)
+
+        # return all file objs
+        return [file_obj for file_objs in type_file_objs.values() for file_obj in file_objs]
+
+    def _to_file_objs(self, files: List[Union[Dict, MessageFile]],
+                      file_upload_config: dict) -> Dict[FileType, List[FileObj]]:
+        """
+        transform files to file objs
+
+        :param files:
+        :param file_upload_config:
+        :return:
+        """
+        type_file_objs: Dict[FileType, List[FileObj]] = {
+            # Currently only support image
+            FileType.IMAGE: []
+        }
+
+        if not files:
+            return type_file_objs
+
+        # group by file type and convert file args or message files to FileObj
+        for file in files:
+            file_obj = self._to_file_obj(file, file_upload_config)
+            if file_obj.type not in type_file_objs:
+                continue
+
+            type_file_objs[file_obj.type].append(file_obj)
+
+        return type_file_objs
+
+    def _to_file_obj(self, file: Union[dict, MessageFile], file_upload_config: dict) -> FileObj:
+        """
+        transform file to file obj
+
+        :param file:
+        :return:
+        """
+        if isinstance(file, dict):
+            transfer_method = FileTransferMethod.value_of(file.get('transfer_method'))
+            return FileObj(
+                tenant_id=self.tenant_id,
+                type=FileType.value_of(file.get('type')),
+                transfer_method=transfer_method,
+                url=file.get('url') if transfer_method == FileTransferMethod.REMOTE_URL else None,
+                upload_file_id=file.get('upload_file_id') if transfer_method == FileTransferMethod.LOCAL_FILE else None,
+                file_config=file_upload_config
+            )
+        else:
+            return FileObj(
+                id=file.id,
+                tenant_id=self.tenant_id,
+                type=FileType.value_of(file.type),
+                transfer_method=FileTransferMethod.value_of(file.transfer_method),
+                url=file.url,
+                upload_file_id=file.upload_file_id or None,
+                file_config=file_upload_config
+            )
+
+    def _check_image_remote_url(self, url):
+        try:
+            headers = {
+                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+            }
+
+            response = requests.head(url, headers=headers, allow_redirects=True)
+            if response.status_code == 200:
+                return True, ""
+            else:
+                return False, "URL does not exist."
+        except requests.RequestException as e:
+            return False, f"Error checking URL: {e}"

api/core/file/upload_file_parser.py

@@ -0,0 +1,79 @@
+import base64
+import hashlib
+import hmac
+import logging
+import os
+import time
+from typing import Optional
+
+from flask import current_app
+
+from extensions.ext_storage import storage
+
+SUPPORT_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'gif']
+
+
+class UploadFileParser:
+    @classmethod
+    def get_image_data(cls, upload_file, force_url: bool = False) -> Optional[str]:
+        if not upload_file:
+            return None
+
+        if upload_file.extension not in SUPPORT_EXTENSIONS:
+            return None
+
+        if current_app.config['MULTIMODAL_SEND_IMAGE_FORMAT'] == 'url' or force_url:
+            return cls.get_signed_temp_image_url(upload_file)
+        else:
+            # get image file base64
+            try:
+                data = storage.load(upload_file.key)
+            except FileNotFoundError:
+                logging.error(f'File not found: {upload_file.key}')
+                return None
+
+            encoded_string = base64.b64encode(data).decode('utf-8')
+            return f'data:{upload_file.mime_type};base64,{encoded_string}'
+
+    @classmethod
+    def get_signed_temp_image_url(cls, upload_file) -> str:
+        """
+        get signed url from upload file
+
+        :param upload_file: UploadFile object
+        :return:
+        """
+        base_url = current_app.config.get('FILES_URL')
+        image_preview_url = f'{base_url}/files/{upload_file.id}/image-preview'
+
+        timestamp = str(int(time.time()))
+        nonce = os.urandom(16).hex()
+        data_to_sign = f"image-preview|{upload_file.id}|{timestamp}|{nonce}"
+        secret_key = current_app.config['SECRET_KEY'].encode()
+        sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
+        encoded_sign = base64.urlsafe_b64encode(sign).decode()
+
+        return f"{image_preview_url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
+
+    @classmethod
+    def verify_image_file_signature(cls, upload_file_id: str, timestamp: str, nonce: str, sign: str) -> bool:
+        """
+        verify signature
+
+        :param upload_file_id: file id
+        :param timestamp: timestamp
+        :param nonce: nonce
+        :param sign: signature
+        :return:
+        """
+        data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}"
+        secret_key = current_app.config['SECRET_KEY'].encode()
+        recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
+        recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()
+
+        # verify signature
+        if sign != recalculated_encoded_sign:
+            return False
+
+        current_time = int(time.time())
+        return current_time - int(timestamp) <= 300  # expired after 5 minutes

api/core/generator/llm_generator.py

@@ -16,7 +16,7 @@ from core.prompt.prompts import CONVERSATION_TITLE_PROMPT, GENERATOR_QA_PROMPT
 
 class LLMGenerator:
     @classmethod
-    def generate_conversation_name(cls, tenant_id: str, query, answer):
+    def generate_conversation_name(cls, tenant_id: str, query):
         prompt = CONVERSATION_TITLE_PROMPT
 
         if len(query) > 2000:
@@ -40,8 +40,12 @@ class LLMGenerator:
 
         result_dict = json.loads(answer)
         answer = result_dict['Your Output']
+        name = answer.strip()
 
-        return answer.strip()
+        if len(name) > 75:
+            name = name[:75] + '...'
+
+        return name
 
     @classmethod
     def generate_suggested_questions_after_answer(cls, tenant_id: str, histories: str):

api/core/memory/read_only_conversation_token_db_buffer_shared_memory.py

@@ -3,6 +3,7 @@ from typing import Any, List, Dict
 from langchain.memory.chat_memory import BaseChatMemory
 from langchain.schema import get_buffer_string, BaseMessage
 
+from core.file.message_file_parser import MessageFileParser
 from core.model_providers.models.entity.message import PromptMessage, MessageType, to_lc_messages
 from core.model_providers.models.llm.base import BaseLLM
 from extensions.ext_database import db
@@ -21,6 +22,8 @@ class ReadOnlyConversationTokenDBBufferSharedMemory(BaseChatMemory):
     @property
     def buffer(self) -> List[BaseMessage]:
         """String buffer of memory."""
+        app_model = self.conversation.app
+
         # fetch limited messages desc, and return reversed
         messages = db.session.query(Message).filter(
             Message.conversation_id == self.conversation.id,
@@ -28,10 +31,25 @@ class ReadOnlyConversationTokenDBBufferSharedMemory(BaseChatMemory):
         ).order_by(Message.created_at.desc()).limit(self.message_limit).all()
 
         messages = list(reversed(messages))
+        message_file_parser = MessageFileParser(tenant_id=app_model.tenant_id, app_id=self.conversation.app_id)
 
         chat_messages: List[PromptMessage] = []
         for message in messages:
-            chat_messages.append(PromptMessage(content=message.query, type=MessageType.USER))
+            files = message.message_files
+            if files:
+                file_objs = message_file_parser.transform_message_files(
+                    files, message.app_model_config
+                )
+
+                prompt_message_files = [file_obj.prompt_message_file for file_obj in file_objs]
+                chat_messages.append(PromptMessage(
+                    content=message.query,
+                    type=MessageType.USER,
+                    files=prompt_message_files
+                ))
+            else:
+                chat_messages.append(PromptMessage(content=message.query, type=MessageType.USER))
+
             chat_messages.append(PromptMessage(content=message.answer, type=MessageType.ASSISTANT))
 
         if not chat_messages:

api/core/model_providers/models/entity/message.py

@@ -1,4 +1,5 @@
 import enum
+from typing import Any, cast, Union, List, Dict
 
 from langchain.schema import HumanMessage, AIMessage, SystemMessage, BaseMessage, FunctionMessage
 from pydantic import BaseModel
@@ -18,17 +19,53 @@ class MessageType(enum.Enum):
     SYSTEM = 'system'
 
 
+class PromptMessageFileType(enum.Enum):
+    IMAGE = 'image'
+
+    @staticmethod
+    def value_of(value):
+        for member in PromptMessageFileType:
+            if member.value == value:
+                return member
+        raise ValueError(f"No matching enum found for value '{value}'")
+
+
+
+class PromptMessageFile(BaseModel):
+    type: PromptMessageFileType
+    data: Any
+
+
+class ImagePromptMessageFile(PromptMessageFile):
+    class DETAIL(enum.Enum):
+        LOW = 'low'
+        HIGH = 'high'
+
+    type: PromptMessageFileType = PromptMessageFileType.IMAGE
+    detail: DETAIL = DETAIL.LOW
+
+
 class PromptMessage(BaseModel):
     type: MessageType = MessageType.USER
     content: str = ''
+    files: list[PromptMessageFile] = []
     function_call: dict = None
 
 
+class LCHumanMessageWithFiles(HumanMessage):
+    # content: Union[str, List[Union[str, Dict]]]
+    content: str
+    files: list[PromptMessageFile]
+
+
 def to_lc_messages(messages: list[PromptMessage]):
     lc_messages = []
     for message in messages:
         if message.type == MessageType.USER:
-            lc_messages.append(HumanMessage(content=message.content))
+            if not message.files:
+                lc_messages.append(HumanMessage(content=message.content))
+            else:
+                lc_messages.append(LCHumanMessageWithFiles(content=message.content, files=message.files))
         elif message.type == MessageType.ASSISTANT:
             additional_kwargs = {}
             if message.function_call:
@@ -44,7 +81,14 @@ def to_prompt_messages(messages: list[BaseMessage]):
     prompt_messages = []
     for message in messages:
         if isinstance(message, HumanMessage):
-            prompt_messages.append(PromptMessage(content=message.content, type=MessageType.USER))
+            if isinstance(message, LCHumanMessageWithFiles):
+                prompt_messages.append(PromptMessage(
+                    content=message.content,
+                    type=MessageType.USER,
+                    files=message.files
+                ))
+            else:
+                prompt_messages.append(PromptMessage(content=message.content, type=MessageType.USER))
         elif isinstance(message, AIMessage):
             message_kwargs = {
                 'content': message.content,

api/core/model_providers/models/llm/openai_model.py

@@ -1,11 +1,9 @@
-import decimal
 import logging
 from typing import List, Optional, Any
 
 import openai
 from langchain.callbacks.manager import Callbacks
 from langchain.schema import LLMResult
-from openai import api_requestor
 
 from core.model_providers.providers.base import BaseModelProvider
 from core.third_party.langchain.llms.chat_open_ai import EnhanceChatOpenAI

api/core/prompt/prompt_transform.py

@@ -8,7 +8,7 @@ from langchain.memory.chat_memory import BaseChatMemory
 from langchain.schema import BaseMessage
 
 from core.model_providers.models.entity.model_params import ModelMode
-from core.model_providers.models.entity.message import PromptMessage, MessageType, to_prompt_messages
+from core.model_providers.models.entity.message import PromptMessage, MessageType, to_prompt_messages, PromptMessageFile
 from core.model_providers.models.llm.base import BaseLLM
 from core.model_providers.models.llm.baichuan_model import BaichuanModel
 from core.model_providers.models.llm.huggingface_hub_model import HuggingfaceHubModel
@@ -16,32 +16,59 @@ from core.model_providers.models.llm.openllm_model import OpenLLMModel
 from core.model_providers.models.llm.xinference_model import XinferenceModel
 from core.prompt.prompt_builder import PromptBuilder
 from core.prompt.prompt_template import PromptTemplateParser
+from models.model import AppModelConfig
+
 
 class AppMode(enum.Enum):
     COMPLETION = 'completion'
     CHAT = 'chat'
 
+
 class PromptTransform:
-    def get_prompt(self, mode: str,
-                   pre_prompt: str, inputs: dict,
+    def get_prompt(self,
+                   app_mode: str,
+                   app_model_config: AppModelConfig,
+                   pre_prompt: str,
+                   inputs: dict,
                    query: str,
+                   files: List[PromptMessageFile],
                    context: Optional[str],
                    memory: Optional[BaseChatMemory],
                    model_instance: BaseLLM) -> \
             Tuple[List[PromptMessage], Optional[List[str]]]:
-        prompt_rules = self._read_prompt_rules_from_file(self._prompt_file_name(mode, model_instance))
-        prompt, stops = self._get_prompt_and_stop(prompt_rules, pre_prompt, inputs, query, context, memory, model_instance)
-        return [PromptMessage(content=prompt)], stops
+        model_mode = app_model_config.model_dict['mode']
+
+        app_mode_enum = AppMode(app_mode)
+        model_mode_enum = ModelMode(model_mode)
+
+        prompt_rules = self._read_prompt_rules_from_file(self._prompt_file_name(app_mode, model_instance))
+
+        if app_mode_enum == AppMode.CHAT and model_mode_enum == ModelMode.CHAT:
+            stops = None
+
+            prompt_messages = self._get_simple_chat_app_chat_model_prompt_messages(prompt_rules, pre_prompt, inputs,
+                                                                                   query, context, memory,
+                                                                                   model_instance, files)
+        else:
+            stops = prompt_rules.get('stops')
+            if stops is not None and len(stops) == 0:
+                stops = None
+
+            prompt_messages = self._get_simple_others_prompt_messages(prompt_rules, pre_prompt, inputs, query, context,
+                                                                      memory,
+                                                                      model_instance, files)
+        return prompt_messages, stops
+
+    def get_advanced_prompt(self,
+                            app_mode: str,
+                            app_model_config: AppModelConfig,
+                            inputs: dict,
+                            query: str,
+                            files: List[PromptMessageFile],
+                            context: Optional[str],
+                            memory: Optional[BaseChatMemory],
+                            model_instance: BaseLLM) -> List[PromptMessage]:
 
-    def get_advanced_prompt(self, 
-            app_mode: str,
-            app_model_config: str, 
-            inputs: dict,
-            query: str,
-            context: Optional[str],
-            memory: Optional[BaseChatMemory],
-            model_instance: BaseLLM) -> List[PromptMessage]:
-        
         model_mode = app_model_config.model_dict['mode']
 
         app_mode_enum = AppMode(app_mode)
@@ -51,15 +78,20 @@ class PromptTransform:
 
         if app_mode_enum == AppMode.CHAT:
             if model_mode_enum == ModelMode.COMPLETION:
-                prompt_messages = self._get_chat_app_completion_model_prompt_messages(app_model_config, inputs, query, context, memory, model_instance)
+                prompt_messages = self._get_chat_app_completion_model_prompt_messages(app_model_config, inputs, query,
+                                                                                      files, context, memory,
+                                                                                      model_instance)
             elif model_mode_enum == ModelMode.CHAT:
-                prompt_messages =  self._get_chat_app_chat_model_prompt_messages(app_model_config, inputs, query, context, memory, model_instance)
+                prompt_messages = self._get_chat_app_chat_model_prompt_messages(app_model_config, inputs, query, files,
+                                                                                context, memory, model_instance)
         elif app_mode_enum == AppMode.COMPLETION:
             if model_mode_enum == ModelMode.CHAT:
-                prompt_messages =  self._get_completion_app_chat_model_prompt_messages(app_model_config, inputs, context)
+                prompt_messages = self._get_completion_app_chat_model_prompt_messages(app_model_config, inputs,
+                                                                                      files, context)
             elif model_mode_enum == ModelMode.COMPLETION:
-                prompt_messages =  self._get_completion_app_completion_model_prompt_messages(app_model_config, inputs, context)
-            
+                prompt_messages = self._get_completion_app_completion_model_prompt_messages(app_model_config, inputs,
+                                                                                            files, context)
+
         return prompt_messages
 
     def _get_history_messages_from_memory(self, memory: BaseChatMemory,
@@ -71,7 +103,7 @@ class PromptTransform:
         return external_context[memory_key]
 
     def _get_history_messages_list_from_memory(self, memory: BaseChatMemory,
-                                          max_token_limit: int) -> List[PromptMessage]:
+                                               max_token_limit: int) -> List[PromptMessage]:
         """Get memory messages."""
         memory.max_token_limit = max_token_limit
         memory.return_messages = True
@@ -79,7 +111,7 @@ class PromptTransform:
         external_context = memory.load_memory_variables({})
         memory.return_messages = False
         return to_prompt_messages(external_context[memory_key])
-    
+
     def _prompt_file_name(self, mode: str, model_instance: BaseLLM) -> str:
         # baichuan
         if isinstance(model_instance, BaichuanModel):
@@ -94,13 +126,13 @@ class PromptTransform:
             return 'common_completion'
         else:
             return 'common_chat'
-        
+
     def _prompt_file_name_for_baichuan(self, mode: str) -> str:
         if mode == 'completion':
             return 'baichuan_completion'
         else:
             return 'baichuan_chat'
-    
+
     def _read_prompt_rules_from_file(self, prompt_name: str) -> dict:
         # Get the absolute path of the subdirectory
         prompt_path = os.path.join(
@@ -111,12 +143,53 @@ class PromptTransform:
         # Open the JSON file and read its content
         with open(json_file_path, 'r') as json_file:
             return json.load(json_file)
-        
-    def _get_prompt_and_stop(self, prompt_rules: dict, pre_prompt: str, inputs: dict,
-                             query: str,
-                             context: Optional[str],
-                             memory: Optional[BaseChatMemory],
-                             model_instance: BaseLLM) -> Tuple[str, Optional[list]]:
+
+    def _get_simple_chat_app_chat_model_prompt_messages(self, prompt_rules: dict, pre_prompt: str, inputs: dict,
+                                                        query: str,
+                                                        context: Optional[str],
+                                                        memory: Optional[BaseChatMemory],
+                                                        model_instance: BaseLLM,
+                                                        files: List[PromptMessageFile]) -> List[PromptMessage]:
+        prompt_messages = []
+
+        context_prompt_content = ''
+        if context and 'context_prompt' in prompt_rules:
+            prompt_template = PromptTemplateParser(template=prompt_rules['context_prompt'])
+            context_prompt_content = prompt_template.format(
+                {'context': context}
+            )
+
+        pre_prompt_content = ''
+        if pre_prompt:
+            prompt_template = PromptTemplateParser(template=pre_prompt)
+            prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs}
+            pre_prompt_content = prompt_template.format(
+                prompt_inputs
+            )
+
+        prompt = ''
+        for order in prompt_rules['system_prompt_orders']:
+            if order == 'context_prompt':
+                prompt += context_prompt_content
+            elif order == 'pre_prompt':
+                prompt += pre_prompt_content
+
+        prompt = re.sub(r'<\|.*?\|>', '', prompt)
+
+        prompt_messages.append(PromptMessage(type=MessageType.SYSTEM, content=prompt))
+
+        self._append_chat_histories(memory, prompt_messages, model_instance)
+
+        prompt_messages.append(PromptMessage(type=MessageType.USER, content=query, files=files))
+
+        return prompt_messages
+
+    def _get_simple_others_prompt_messages(self, prompt_rules: dict, pre_prompt: str, inputs: dict,
+                                           query: str,
+                                           context: Optional[str],
+                                           memory: Optional[BaseChatMemory],
+                                           model_instance: BaseLLM,
+                                           files: List[PromptMessageFile]) -> List[PromptMessage]:
         context_prompt_content = ''
         if context and 'context_prompt' in prompt_rules:
             prompt_template = PromptTemplateParser(template=prompt_rules['context_prompt'])
@@ -175,16 +248,12 @@ class PromptTransform:
 
         prompt = re.sub(r'<\|.*?\|>', '', prompt)
 
-        stops = prompt_rules.get('stops')
-        if stops is not None and len(stops) == 0:
-            stops = None
+        return [PromptMessage(content=prompt, files=files)]
 
-        return prompt, stops
-    
     def _set_context_variable(self, context: str, prompt_template: PromptTemplateParser, prompt_inputs: dict) -> None:
         if '#context#' in prompt_template.variable_keys:
             if context:
-                prompt_inputs['#context#'] = context    
+                prompt_inputs['#context#'] = context
             else:
                 prompt_inputs['#context#'] = ''
 
@@ -195,17 +264,18 @@ class PromptTransform:
             else:
                 prompt_inputs['#query#'] = ''
 
-    def _set_histories_variable(self, memory: BaseChatMemory, raw_prompt: str, conversation_histories_role: dict, 
-                                prompt_template: PromptTemplateParser, prompt_inputs: dict, model_instance: BaseLLM) -> None:
+    def _set_histories_variable(self, memory: BaseChatMemory, raw_prompt: str, conversation_histories_role: dict,
+                                prompt_template: PromptTemplateParser, prompt_inputs: dict,
+                                model_instance: BaseLLM) -> None:
         if '#histories#' in prompt_template.variable_keys:
             if memory:
                 tmp_human_message = PromptBuilder.to_human_message(
                     prompt_content=raw_prompt,
-                    inputs={ '#histories#': '', **prompt_inputs }
+                    inputs={'#histories#': '', **prompt_inputs}
                 )
 
                 rest_tokens = self._calculate_rest_token(tmp_human_message, model_instance)
-                
+
                 memory.human_prefix = conversation_histories_role['user_prefix']
                 memory.ai_prefix = conversation_histories_role['assistant_prefix']
                 histories = self._get_history_messages_from_memory(memory, rest_tokens)
@@ -213,7 +283,8 @@ class PromptTransform:
             else:
                 prompt_inputs['#histories#'] = ''
 
-    def _append_chat_histories(self, memory: BaseChatMemory, prompt_messages: list[PromptMessage], model_instance: BaseLLM) -> None:
+    def _append_chat_histories(self, memory: BaseChatMemory, prompt_messages: list[PromptMessage],
+                               model_instance: BaseLLM) -> None:
         if memory:
             rest_tokens = self._calculate_rest_token(prompt_messages, model_instance)
 
@@ -242,19 +313,19 @@ class PromptTransform:
         return prompt
 
     def _get_chat_app_completion_model_prompt_messages(self,
-            app_model_config: str,
-            inputs: dict,
-            query: str,
-            context: Optional[str],
-            memory: Optional[BaseChatMemory],
-            model_instance: BaseLLM) -> List[PromptMessage]:
-        
+                                                       app_model_config: AppModelConfig,
+                                                       inputs: dict,
+                                                       query: str,
+                                                       files: List[PromptMessageFile],
+                                                       context: Optional[str],
+                                                       memory: Optional[BaseChatMemory],
+                                                       model_instance: BaseLLM) -> List[PromptMessage]:
+
         raw_prompt = app_model_config.completion_prompt_config_dict['prompt']['text']
         conversation_histories_role = app_model_config.completion_prompt_config_dict['conversation_histories_role']
 
         prompt_messages = []
-        prompt = ''
-        
+
         prompt_template = PromptTemplateParser(template=raw_prompt)
         prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs}
 
@@ -262,28 +333,29 @@ class PromptTransform:
 
         self._set_query_variable(query, prompt_template, prompt_inputs)
 
-        self._set_histories_variable(memory, raw_prompt, conversation_histories_role, prompt_template, prompt_inputs, model_instance)
+        self._set_histories_variable(memory, raw_prompt, conversation_histories_role, prompt_template, prompt_inputs,
+                                     model_instance)
 
         prompt = self._format_prompt(prompt_template, prompt_inputs)
 
-        prompt_messages.append(PromptMessage(type = MessageType(MessageType.USER) ,content=prompt))
+        prompt_messages.append(PromptMessage(type=MessageType.USER, content=prompt, files=files))
 
         return prompt_messages
 
     def _get_chat_app_chat_model_prompt_messages(self,
-            app_model_config: str,
-            inputs: dict,
-            query: str,
-            context: Optional[str],
-            memory: Optional[BaseChatMemory],
-            model_instance: BaseLLM) -> List[PromptMessage]:
+                                                 app_model_config: AppModelConfig,
+                                                 inputs: dict,
+                                                 query: str,
+                                                 files: List[PromptMessageFile],
+                                                 context: Optional[str],
+                                                 memory: Optional[BaseChatMemory],
+                                                 model_instance: BaseLLM) -> List[PromptMessage]:
         raw_prompt_list = app_model_config.chat_prompt_config_dict['prompt']
 
         prompt_messages = []
 
         for prompt_item in raw_prompt_list:
             raw_prompt = prompt_item['text']
-            prompt = ''
 
             prompt_template = PromptTemplateParser(template=raw_prompt)
             prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs}
@@ -292,23 +364,23 @@ class PromptTransform:
 
             prompt = self._format_prompt(prompt_template, prompt_inputs)
 
-            prompt_messages.append(PromptMessage(type = MessageType(prompt_item['role']) ,content=prompt))
-        
+            prompt_messages.append(PromptMessage(type=MessageType(prompt_item['role']), content=prompt))
+
         self._append_chat_histories(memory, prompt_messages, model_instance)
 
-        prompt_messages.append(PromptMessage(type = MessageType.USER ,content=query))
+        prompt_messages.append(PromptMessage(type=MessageType.USER, content=query, files=files))
 
         return prompt_messages
 
     def _get_completion_app_completion_model_prompt_messages(self,
-                   app_model_config: str,
-                   inputs: dict,
-                   context: Optional[str]) -> List[PromptMessage]:
+                                                             app_model_config: AppModelConfig,
+                                                             inputs: dict,
+                                                             files: List[PromptMessageFile],
+                                                             context: Optional[str]) -> List[PromptMessage]:
         raw_prompt = app_model_config.completion_prompt_config_dict['prompt']['text']
 
         prompt_messages = []
-        prompt = ''
-        
+
         prompt_template = PromptTemplateParser(template=raw_prompt)
         prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs}
 
@@ -316,21 +388,21 @@ class PromptTransform:
 
         prompt = self._format_prompt(prompt_template, prompt_inputs)
 
-        prompt_messages.append(PromptMessage(type = MessageType(MessageType.USER) ,content=prompt))
+        prompt_messages.append(PromptMessage(type=MessageType(MessageType.USER), content=prompt, files=files))
 
         return prompt_messages
 
     def _get_completion_app_chat_model_prompt_messages(self,
-                   app_model_config: str,
-                   inputs: dict,
-                   context: Optional[str]) -> List[PromptMessage]:
+                                                       app_model_config: AppModelConfig,
+                                                       inputs: dict,
+                                                       files: List[PromptMessageFile],
+                                                       context: Optional[str]) -> List[PromptMessage]:
         raw_prompt_list = app_model_config.chat_prompt_config_dict['prompt']
 
         prompt_messages = []
 
         for prompt_item in raw_prompt_list:
             raw_prompt = prompt_item['text']
-            prompt = ''
 
             prompt_template = PromptTemplateParser(template=raw_prompt)
             prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs}
@@ -339,6 +411,11 @@ class PromptTransform:
 
             prompt = self._format_prompt(prompt_template, prompt_inputs)
 
-            prompt_messages.append(PromptMessage(type = MessageType(prompt_item['role']) ,content=prompt))
-        
-        return prompt_messages
+            prompt_messages.append(PromptMessage(type=MessageType(prompt_item['role']), content=prompt))
+
+        for prompt_message in prompt_messages[::-1]:
+            if prompt_message.type == MessageType.USER:
+                prompt_message.files = files
+                break
+
+        return prompt_messages

api/core/third_party/langchain/llms/chat_open_ai.py

@@ -1,10 +1,13 @@
 import os
 
-from typing import Dict, Any, Optional, Union, Tuple
+from typing import Dict, Any, Optional, Union, Tuple, List, cast
 
 from langchain.chat_models import ChatOpenAI
+from langchain.schema import BaseMessage, ChatMessage, HumanMessage, AIMessage, SystemMessage, FunctionMessage
 from pydantic import root_validator
 
+from core.model_providers.models.entity.message import LCHumanMessageWithFiles, PromptMessageFileType, ImagePromptMessageFile
+
 
 class EnhanceChatOpenAI(ChatOpenAI):
     request_timeout: Optional[Union[float, Tuple[float, float]]] = (5.0, 300.0)
@@ -48,3 +51,102 @@ class EnhanceChatOpenAI(ChatOpenAI):
             "api_key": self.openai_api_key,
             "organization": self.openai_organization if self.openai_organization else None,
         }
+
+    def _create_message_dicts(
+        self, messages: List[BaseMessage], stop: Optional[List[str]]
+    ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
+        params = self._client_params
+        if stop is not None:
+            if "stop" in params:
+                raise ValueError("`stop` found in both the input and default params.")
+            params["stop"] = stop
+        message_dicts = [self._convert_message_to_dict(m) for m in messages]
+        return message_dicts, params
+
+    def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
+        """Calculate num tokens for gpt-3.5-turbo and gpt-4 with tiktoken package.
+
+        Official documentation: https://github.com/openai/openai-cookbook/blob/
+        main/examples/How_to_format_inputs_to_ChatGPT_models.ipynb"""
+        model, encoding = self._get_encoding_model()
+        if model.startswith("gpt-3.5-turbo-0301"):
+            # every message follows <im_start>{role/name}\n{content}<im_end>\n
+            tokens_per_message = 4
+            # if there's a name, the role is omitted
+            tokens_per_name = -1
+        elif model.startswith("gpt-3.5-turbo") or model.startswith("gpt-4"):
+            tokens_per_message = 3
+            tokens_per_name = 1
+        else:
+            raise NotImplementedError(
+                f"get_num_tokens_from_messages() is not presently implemented "
+                f"for model {model}."
+                "See https://github.com/openai/openai-python/blob/main/chatml.md for "
+                "information on how messages are converted to tokens."
+            )
+        num_tokens = 0
+        messages_dict = [self._convert_message_to_dict(m) for m in messages]
+        for message in messages_dict:
+            num_tokens += tokens_per_message
+            for key, value in message.items():
+                # Cast str(value) in case the message value is not a string
+                # This occurs with function messages
+                # TODO: The current token calculation method for the image type is not implemented,
+                #  which need to download the image and then get the resolution for calculation,
+                #  and will increase the request delay
+                if isinstance(value, list):
+                    text = ''
+                    for item in value:
+                        if isinstance(item, dict) and item['type'] == 'text':
+                            text += item['text']
+
+                    value = text
+                num_tokens += len(encoding.encode(str(value)))
+                if key == "name":
+                    num_tokens += tokens_per_name
+        # every reply is primed with <im_start>assistant
+        num_tokens += 3
+        return num_tokens
+
+    def _convert_message_to_dict(self, message: BaseMessage) -> dict:
+        if isinstance(message, ChatMessage):
+            message_dict = {"role": message.role, "content": message.content}
+        elif isinstance(message, LCHumanMessageWithFiles):
+            content = [
+                {
+                    "type": "text",
+                    "text": message.content
+                }
+            ]
+
+            for file in message.files:
+                if file.type == PromptMessageFileType.IMAGE:
+                    file = cast(ImagePromptMessageFile, file)
+                    content.append({
+                        "type": "image_url",
+                        "image_url": {
+                            "url": file.data,
+                            "detail": file.detail.value
+                        }
+                    })
+
+            message_dict = {"role": "user", "content": content}
+        elif isinstance(message, HumanMessage):
+            message_dict = {"role": "user", "content": message.content}
+        elif isinstance(message, AIMessage):
+            message_dict = {"role": "assistant", "content": message.content}
+            if "function_call" in message.additional_kwargs:
+                message_dict["function_call"] = message.additional_kwargs["function_call"]
+        elif isinstance(message, SystemMessage):
+            message_dict = {"role": "system", "content": message.content}
+        elif isinstance(message, FunctionMessage):
+            message_dict = {
+                "role": "function",
+                "content": message.content,
+                "name": message.name,
+            }
+        else:
+            raise ValueError(f"Got unknown type {message}")
+        if "name" in message.additional_kwargs:
+            message_dict["name"] = message.additional_kwargs["name"]
+        return message_dict