From afba07e42d10457099e3d11e500b47463003a47a Mon Sep 17 00:00:00 2001
From: glide-the <2533736852@qq.com>
Date: Tue, 21 May 2024 14:20:11 +0800
Subject: [PATCH] get_img_base64

---
 chatchat-server/chatchat/configs/__init__.py          |  9 +++++++++
 chatchat-server/chatchat/configs/_basic_config.py     |  5 +++++
 chatchat-server/chatchat/webui.py                     |  4 ++--
 .../chatchat/webui_pages/dialogue/dialogue.py         |  2 +-
 chatchat-server/chatchat/webui_pages/utils.py         | 11 ++++++++---
 5 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/chatchat-server/chatchat/configs/__init__.py b/chatchat-server/chatchat/configs/__init__.py
index 9684db5c..a441d980 100644
--- a/chatchat-server/chatchat/configs/__init__.py
+++ b/chatchat-server/chatchat/configs/__init__.py
@@ -118,6 +118,13 @@ def _import_data_path() -> Any:
 
     return DATA_PATH
 
+def _import_img_dir() -> Any:
+    basic_config_load = CONFIG_IMPORTS.get("_basic_config.py")
+    load_mod = basic_config_load.get("load_mod")
+    IMG_DIR = load_mod(basic_config_load.get("module"), "IMG_DIR")
+
+    return IMG_DIR
+
 
 def _import_nltk_data_path() -> Any:
     basic_config_load = CONFIG_IMPORTS.get("_basic_config.py")
@@ -477,6 +484,8 @@ def __getattr__(name: str) -> Any:
         return _import_chatchat_root()
     elif name == "DATA_PATH":
         return _import_data_path()
+    elif name == "IMG_DIR":
+        return _import_img_dir()
     elif name == "NLTK_DATA_PATH":
         return _import_nltk_data_path()
     elif name == "LOG_FORMAT":
diff --git a/chatchat-server/chatchat/configs/_basic_config.py b/chatchat-server/chatchat/configs/_basic_config.py
index d172c90b..46eec939 100644
--- a/chatchat-server/chatchat/configs/_basic_config.py
+++ b/chatchat-server/chatchat/configs/_basic_config.py
@@ -19,6 +19,11 @@ DATA_PATH = os.path.join(CHATCHAT_ROOT, "data")
 if not os.path.exists(DATA_PATH):
     os.mkdir(DATA_PATH)
 
+# 项目相关图片
+IMG_DIR = os.path.join(CHATCHAT_ROOT, "img")
+if not os.path.exists(IMG_DIR):
+    os.mkdir(IMG_DIR)
+
 # nltk 模型存储路径
 NLTK_DATA_PATH = os.path.join(DATA_PATH, "nltk_data")
 import nltk
diff --git a/chatchat-server/chatchat/webui.py b/chatchat-server/chatchat/webui.py
index b24f0ed0..4dc8c75a 100644
--- a/chatchat-server/chatchat/webui.py
+++ b/chatchat-server/chatchat/webui.py
@@ -17,7 +17,7 @@ if __name__ == "__main__":
 
     st.set_page_config(
         "Langchain-Chatchat WebUI",
-        get_img_url("chatchat_icon_blue_square_v2.png"),
+        get_img_base64("chatchat_icon_blue_square_v2.png"),
         initial_sidebar_state="expanded",
         menu_items={
             'Get Help': 'https://github.com/chatchat-space/Langchain-Chatchat',
@@ -46,7 +46,7 @@ if __name__ == "__main__":
 
     with st.sidebar:
         st.image(
-            get_img_url('logo-long-chatchat-trans-v2.png'),
+            get_img_base64('logo-long-chatchat-trans-v2.png'),
             use_column_width=True
         )
         st.caption(
diff --git a/chatchat-server/chatchat/webui_pages/dialogue/dialogue.py b/chatchat-server/chatchat/webui_pages/dialogue/dialogue.py
index cf6775d9..d211097e 100644
--- a/chatchat-server/chatchat/webui_pages/dialogue/dialogue.py
+++ b/chatchat-server/chatchat/webui_pages/dialogue/dialogue.py
@@ -19,7 +19,7 @@ from chatchat.webui_pages.utils import *
 from chatchat.webui_pages.dialogue.utils import process_files
 
 chat_box = ChatBox(
-    assistant_avatar=get_img_url("chatchat_icon_blue_square_v2.png")
+    assistant_avatar=get_img_base64("chatchat_icon_blue_square_v2.png")
 )
 
 
diff --git a/chatchat-server/chatchat/webui_pages/utils.py b/chatchat-server/chatchat/webui_pages/utils.py
index 7d8cc5e5..4c988cfe 100644
--- a/chatchat-server/chatchat/webui_pages/utils.py
+++ b/chatchat-server/chatchat/webui_pages/utils.py
@@ -693,12 +693,17 @@ def check_success_msg(data: Union[str, dict, list], key: str = "msg") -> str:
     return ""
 
 
-def get_img_url(file_name: str) -> str:
+def get_img_base64(file_name: str) -> str:
     '''
-    image url used in streamlit.
+    get_img_base64 used in streamlit.
     absolute local path not working on windows.
     '''
-    return f"{api_address()}/img/{file_name}"
+    image = f"{IMG_DIR}/img/{file_name}"
+    # 读取图片
+    with open(image, "rb") as f:
+        buffer = BytesIO(f.read())
+        base_str = base64.b64encode(buffer.getvalue()).decode()
+    return f"data:image/png;base64,{base_str}"
 
 
 if __name__ == "__main__":