feat: application upload file

Author: shaohuzhang1

apps/application/chat_pipeline/step/search_dataset_step/impl/base_search_dataset_step.py

@@ -24,17 +24,7 @@
 from knowledge.models import SearchMode
 from maxkb.conf import PROJECT_DIR
 from models_provider.models import Model
-from models_provider.tools import get_model
-
-
-def get_model_by_id(_id, workspace_id):
-    model = QuerySet(Model).filter(id=_id, model_type="EMBEDDING")
-    get_authorized_model = DatabaseModelManage.get_model("get_authorized_model")
-    if get_authorized_model is not None:
-        model = get_authorized_model(model, workspace_id)
-    if model is None:
-        raise Exception(_("Model does not exist"))
-    return model
+from models_provider.tools import get_model, get_model_by_id
 
 
 def get_embedding_id(knowledge_id_list):
@@ -65,6 +55,8 @@ def execute(self, problem_text: str, knowledge_id_list: list[str], exclude_docum
         exec_problem_text = padding_problem_text if padding_problem_text is not None else problem_text
         model_id = get_embedding_id(knowledge_id_list)
         model = get_model_by_id(model_id, workspace_id)
+        if model.model_type != "EMBEDDING":
+            raise Exception(_("Model does not exist"))
         self.context['model_name'] = model.name
         embedding_model = ModelManage.get_model(model_id, lambda _id: get_model(model))
         embedding_value = embedding_model.embed_query(exec_problem_text)

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

@@ -71,7 +71,7 @@ def save_image(image_list):
 
         for doc in document:
             file = QuerySet(File).filter(id=doc['file_id']).first()
-            buffer = io.BytesIO(file.get_bytes().tobytes())
+            buffer = io.BytesIO(file.get_bytes())
             buffer.name = doc['name']  # this is the important line
 
             for split_handle in (parse_table_handle_list + split_handles):

apps/chat/views/chat.py

@@ -9,6 +9,7 @@
 from django.http import HttpResponse
 from django.utils.translation import gettext_lazy as _
 from drf_spectacular.utils import extend_schema
+from rest_framework.parsers import MultiPartParser
 from rest_framework.request import Request
 from rest_framework.views import APIView
 
@@ -24,6 +25,8 @@
 from common.constants.permission_constants import ChatAuth
 from common.exception.app_exception import AppAuthenticationFailed
 from common.result import result
+from knowledge.models import FileSourceType
+from oss.serializers.file import FileSerializer
 from users.api import CaptchaAPI
 from users.serializers.login import CaptchaSerializer
 
@@ -134,7 +137,7 @@ class CaptchaView(APIView):
                    summary=_("Get Chat captcha"),
                    description=_("Get Chat captcha"),
                    operation_id=_("Get Chat captcha"),  # type: ignore
-                   tags=[_("User Management")],  # type: ignore
+                   tags=[_("Chat")],  # type: ignore
                    responses=CaptchaAPI.get_response())
     def get(self, request: Request):
         return result.success(CaptchaSerializer().generate())
@@ -150,7 +153,7 @@ class SpeechToText(APIView):
         operation_id=_("speech to text"),  # type: ignore
         request=SpeechToTextAPI.get_request(),
         responses=SpeechToTextAPI.get_response(),
-        tags=[_('Application')]  # type: ignore
+        tags=[_('Chat')]  # type: ignore
     )
     def post(self, request: Request):
         return result.success(
@@ -169,10 +172,34 @@ class TextToSpeech(APIView):
         operation_id=_("text to speech"),  # type: ignore
         request=TextToSpeechAPI.get_request(),
         responses=TextToSpeechAPI.get_response(),
-        tags=[_('Application')]  # type: ignore
+        tags=[_('Chat')]  # type: ignore
     )
     def post(self, request: Request):
         byte_data = TextToSpeechSerializers(
             data={'application_id': request.auth.application_id}).text_to_speech(request.data)
         return HttpResponse(byte_data, status=200, headers={'Content-Type': 'audio/mp3',
                                                             'Content-Disposition': 'attachment; filename="abc.mp3"'})
+
+
+class UploadFile(APIView):
+    authentication_classes = [TokenAuth]
+    parser_classes = [MultiPartParser]
+
+    @extend_schema(
+        methods=['POST'],
+        description=_("Upload files"),
+        summary=_("Upload files"),
+        operation_id=_("Upload files"),  # type: ignore
+        request=TextToSpeechAPI.get_request(),
+        responses=TextToSpeechAPI.get_response(),
+        tags=[_('Application')]  # type: ignore
+    )
+    def post(self, request: Request, chat_id: str):
+        files = request.FILES.getlist('file')
+        file_ids = []
+        meta = {}
+        for file in files:
+            file_url = FileSerializer(
+                data={'file': file, 'meta': meta, 'source_id': chat_id, 'source_type': FileSourceType.CHAT, }).upload()
+            file_ids.append({'name': file.name, 'url': file_url, 'file_id': file_url.split('/')[-1]})
+        return result.success(file_ids)

apps/knowledge/models/knowledge.py

@@ -95,8 +95,6 @@ def default_status_meta():
     return {"state_time": {}}
 
 
-
-
 class KnowledgeFolder(MPTTModel, AppModelMixin):
     id = models.CharField(primary_key=True, max_length=64, editable=False, verbose_name="主键id")
     name = models.CharField(max_length=64, verbose_name="文件夹名称")
@@ -133,9 +131,11 @@ class Knowledge(AppModelMixin):
     class Meta:
         db_table = "knowledge"
 
+
 def get_default_status():
     return Status('').__str__()
 
+
 class Document(AppModelMixin):
     """
     文档表
@@ -224,10 +224,12 @@ class FileSourceType(models.TextChoices):
     TOOL = "TOOL"
     # 文档
     DOCUMENT = "DOCUMENT"
+    # 对话
+    CHAT = "CHAT"
     # 临时30分钟 数据30分钟后被清理 source_id 为TEMPORARY_30_MINUTE
     TEMPORARY_30_MINUTE = "TEMPORARY_30_MINUTE"
     # 临时120分钟 数据120分钟后被清理 source_id为TEMPORARY_100_MINUTE
-    TEMPORARY_120_MINUTE = "TEMPORARY_100_MINUTE"
+    TEMPORARY_120_MINUTE = "TEMPORARY_120_MINUTE"
     # 临时1天 数据1天后被清理 source_id为TEMPORARY_1_DAY
     TEMPORARY_1_DAY = "TEMPORARY_1_DAY"
 

apps/maxkb/urls.py

@@ -42,6 +42,7 @@
     path(admin_api_prefix, include("system_manage.urls")),
     path(admin_api_prefix, include("application.urls")),
     path(admin_api_prefix, include("oss.urls")),
+    path(chat_api_prefix, include("oss.urls")),
     path(chat_api_prefix, include("chat.urls")),
     path(f'{admin_ui_prefix[1:]}/', include('oss.retrieval_urls')),
     path(f'{chat_ui_prefix[1:]}/', include('oss.retrieval_urls')),

ui/src/api/application/application.ts

@@ -296,6 +296,30 @@ const getMcpTools: (application_id: String, loading?: Ref<boolean>) => Promise<R
   return get(`${prefix.value}/${application_id}/mcp_tools`, undefined, loading)
 }
 
+/**
+ * 上传文件
+ * @param 参数  file:file
+ */
+const uploadFile: (
+  file: any,
+  sourceId: string,
+  resourceType:
+    | 'KNOWLEDGE'
+    | 'APPLICATION'
+    | 'TOOL'
+    | 'DOCUMENT'
+    | 'CHAT'
+    | 'TEMPORARY_30_MINUTE'
+    | 'TEMPORARY_120_MINUTE'
+    | 'TEMPORARY_1_DAY',
+) => Promise<Result<any>> = (file, sourceId, resourceType) => {
+  const fd = new FormData()
+  fd.append('file', file)
+  fd.append('source_id', sourceId)
+  fd.append('source_type', resourceType)
+  return post(`/oss/file`, fd)
+}
+
 export default {
   getAllApplication,
   getApplication,
@@ -321,4 +345,5 @@ export default {
   postTextToSpeech,
   speechToText,
   getMcpTools,
+  uploadFile,
 }

ui/src/api/chat/chat.ts

@@ -265,30 +265,57 @@ const speechToText: (data: any, loading?: Ref<boolean>) => Promise<Result<any>>
   return post(`speech_to_text`, data, undefined, loading)
 }
 /**
- * 
+ *
  * @param chat_id  对话ID
- * @param loading 
- * @returns 
+ * @param loading
+ * @returns
  */
-const deleteChat: (chat_id:string, loading?: Ref<boolean>) => Promise<Result<any>> = (
+const deleteChat: (chat_id: string, loading?: Ref<boolean>) => Promise<Result<any>> = (
   chat_id,
   loading,
 ) => {
   return del(`historical_conversation/${chat_id}`, loading)
 }
 /**
- * 
+ *
  * @param chat_id 对话id
  * @param data    对话简介
- * @param loading 
- * @returns 
+ * @param loading
+ * @returns
  */
-const modifyChat: (chat_id:string, data:any, loading?: Ref<boolean> ) => Promise<Result<any>> = (
-  chat_id,data,loading
+const modifyChat: (chat_id: string, data: any, loading?: Ref<boolean>) => Promise<Result<any>> = (
+  chat_id,
+  data,
+  loading,
 ) => {
   return put(`historical_conversation/${chat_id}`, data, undefined, loading)
 }
-
+/**
+ * 上传文件
+ * @param file      文件
+ * @param sourceId  资源id
+ * @param resourceType  资源类型
+ * @returns
+ */
+const uploadFile: (
+  file: any,
+  sourceId: string,
+  resourceType:
+    | 'KNOWLEDGE'
+    | 'APPLICATION'
+    | 'TOOL'
+    | 'DOCUMENT'
+    | 'CHAT'
+    | 'TEMPORARY_30_MINUTE'
+    | 'TEMPORARY_120_MINUTE'
+    | 'TEMPORARY_1_DAY',
+) => Promise<Result<any>> = (file, sourceId, sourceType) => {
+  const fd = new FormData()
+  fd.append('file', file)
+  fd.append('source_id', sourceId)
+  fd.append('source_type', sourceType)
+  return post(`/oss/file`, fd)
+}
 export default {
   open,
   chat,
@@ -316,5 +343,6 @@ export default {
   textToSpeech,
   speechToText,
   deleteChat,
-  modifyChat
+  modifyChat,
+  uploadFile,
 }