From 2a86472b88a7cc7e25fc832bde724ae9fe06acf6 Mon Sep 17 00:00:00 2001
From: chongchuanbing <chongchuanbing@126.com>
Date: Fri, 11 Oct 2024 16:10:27 +0800
Subject: [PATCH] fix chat and thumbnail bug (#2803)

### What problem does this PR solve?

1. fix white screen issue when chat response
2. thumbnail bug when document not support

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [ ] New Feature (non-breaking change which adds functionality)
- [ ] Documentation Update
- [ ] Refactoring
- [ ] Performance Improvement
- [ ] Other (please describe):

---------

Co-authored-by: chongchuanbing <chongchuanbing@gmail.com>
---
 api/apps/document_app.py                      | 5 +++++
 api/db/services/document_service.py           | 2 +-
 api/db/services/file_service.py               | 6 ++++--
 web/src/pages/chat/markdown-content/index.tsx | 2 +-
 4 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/api/apps/document_app.py b/api/apps/document_app.py
index 450b9ae93..68e05dcf9 100644
--- a/api/apps/document_app.py
+++ b/api/apps/document_app.py
@@ -238,6 +238,11 @@ def thumbnails():
 
     try:
         docs = DocumentService.get_thumbnails(doc_ids)
+
+        for doc_item in docs:
+            if doc_item['thumbnail'] and not doc_item['thumbnail'].startswith(IMG_BASE64_PREFIX):
+                doc_item['thumbnail'] = f"/v1/document/image/{doc_item['kb_id']}-{doc_item['thumbnail']}"
+
         return get_json_result(data={d["id"]: d["thumbnail"] for d in docs})
     except Exception as e:
         return server_error_response(e)
diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py
index 5d28ecd64..f60e760f8 100644
--- a/api/db/services/document_service.py
+++ b/api/db/services/document_service.py
@@ -268,7 +268,7 @@ class DocumentService(CommonService):
     @classmethod
     @DB.connection_context()
     def get_thumbnails(cls, docids):
-        fields = [cls.model.id, cls.model.thumbnail]
+        fields = [cls.model.id, cls.model.kb_id, cls.model.thumbnail]
         return list(cls.model.select(
             *fields).where(cls.model.id.in_(docids)).dicts())
 
diff --git a/api/db/services/file_service.py b/api/db/services/file_service.py
index a63643e13..587e517bd 100644
--- a/api/db/services/file_service.py
+++ b/api/db/services/file_service.py
@@ -358,8 +358,10 @@ class FileService(CommonService):
                 doc_id = get_uuid()
 
                 img = thumbnail_img(filename, blob)
-                thumbnail_location = f'thumbnail_{doc_id}.png'
-                STORAGE_IMPL.put(kb.id, thumbnail_location, img)
+                thumbnail_location = ''
+                if img is not None:
+                    thumbnail_location = f'thumbnail_{doc_id}.png'
+                    STORAGE_IMPL.put(kb.id, thumbnail_location, img)
 
                 doc = {
                     "id": doc_id,
diff --git a/web/src/pages/chat/markdown-content/index.tsx b/web/src/pages/chat/markdown-content/index.tsx
index 12cc3d005..db3d1504e 100644
--- a/web/src/pages/chat/markdown-content/index.tsx
+++ b/web/src/pages/chat/markdown-content/index.tsx
@@ -45,7 +45,7 @@ const MarkdownContent = ({
   }, [content, loading, t]);
 
   useEffect(() => {
-    setDocumentIds(reference?.doc_aggs.map((x) => x.doc_id) ?? []);
+    setDocumentIds(reference?.doc_aggs?.map((x) => x.doc_id) ?? []);
   }, [reference, setDocumentIds]);
 
   const handleDocumentButtonClick = useCallback(