update

app/services/file/impl/r2r_file.py

@@ -6,7 +6,7 @@ import aiofiles
 import aiofiles.os
 from fastapi import UploadFile
 from sqlalchemy.ext.asyncio import AsyncSession
-
+import os
 from app.models import File
 from app.providers.r2r import R2R
 from app.providers.storage import storage
@@ -15,9 +15,10 @@ from app.services.file.impl.oss_file import OSSFileService
 # import asyncio
 from pathlib import Path
 
+allowed_formats = [".bmp", ".heic", ".jpeg", ".png", ".tiff"]
 
-class R2RFileService(OSSFileService):
 
+class R2RFileService(OSSFileService):
     @staticmethod
     async def create_file(
         *, session: AsyncSession, purpose: str, file: UploadFile
@@ -131,10 +132,26 @@ class R2RFileService(OSSFileService):
             return files
 
         for doc in search_results:
+            file_extension = os.path.splitext(doc.metadata["title"])[1].lower()
+            if file_extension in allowed_formats:
+                files.append(
+                    {
+                        "id": str(doc.id),
+                        "text": doc.text,
+                        "title": doc.metadata["title"],
+                        "url": "https://r2r.s3.cn-north-1.amazonaws.com.cn/r2r/documents/"
+                        + doc.document_id,
+                    }
+                )
             # print(doc.metadata)
-            files.append(
-                {"id": str(doc.id), "text": doc.text, "title": doc.metadata["title"]}
-            )
+            else:
+                files.append(
+                    {
+                        "id": str(doc.id),
+                        "text": doc.text,
+                        "title": doc.metadata["title"],
+                    }
+                )
             # file_key = doc.metadata.file_key
             # file_key = doc.metadata.title if file_key is None else file_key
             # text = doc.text