perf: Encryption of sensitive data in ai_model table #9

Author: fit2cloud-chenyw

backend/alembic/versions/040_modify_ai_model.py

@@ -0,0 +1,51 @@
+"""040_modify_ai_model
+
+Revision ID: 0fc14c2cfe41
+Revises: 25cbc85766fd
+Create Date: 2025-08-26 23:30:50.192799
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+
+
+# revision identifiers, used by Alembic.
+revision = '0fc14c2cfe41'
+down_revision = '25cbc85766fd'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.alter_column(
+        'ai_model',
+        'api_key',
+        type_=sa.Text(),
+        existing_type=sa.String(length=255),
+        existing_nullable=True
+    )
+    op.alter_column(
+        'ai_model',
+        'api_domain',
+        type_=sa.Text(),
+        existing_type=sa.String(length=255),
+        existing_nullable=False
+    )
+
+
+def downgrade():
+    op.alter_column(
+        'ai_model',
+        'api_key',
+        type_=sa.String(),
+        existing_type=sa.Text(),
+        existing_nullable=True
+    )
+    op.alter_column(
+        'ai_model',
+        'api_domain',
+        type_=sa.String(),
+        existing_type=sa.Text(),
+        existing_nullable=False
+    )

backend/apps/ai_model/model_factory.py

@@ -10,6 +10,7 @@
 from apps.ai_model.openai.llm import BaseChatOpenAI
 from apps.system.models.system_model import AiModelDetail
 from common.core.db import engine
+from common.utils.crypto import sqlbot_decrypt
 from common.utils.utils import prepare_model_arg
 from langchain_community.llms import VLLMOpenAI
 from langchain_openai import AzureChatOpenAI
@@ -137,7 +138,7 @@ def register_llm(cls, model_type: str, llm_class: Type[BaseLLM]):
     return config """
 
 
-def get_default_config() -> LLMConfig:
+async def get_default_config() -> LLMConfig:
     with Session(engine) as session:
         db_model = session.exec(
             select(AiModelDetail).where(AiModelDetail.default_model == True)
@@ -152,6 +153,11 @@ def get_default_config() -> LLMConfig:
                 additional_params = {item["key"]: prepare_model_arg(item.get('val')) for item in config_raw if "key" in item and "val" in item}
             except Exception:
                 pass
+        if not db_model.api_domain.startswith("http"):
+            db_model.api_domain = await sqlbot_decrypt(db_model.api_domain)
+            if db_model.api_key:
+                db_model.api_key = await sqlbot_decrypt(db_model.api_key)
+            
 
         # 构造 LLMConfig
         return LLMConfig(

backend/apps/chat/api/chat.py

@@ -114,7 +114,7 @@ async def recommend_questions(session: SessionDep, current_user: CurrentUser, ch
             )
         request_question = ChatQuestion(chat_id=record.chat_id, question=record.question if record.question else '')
 
-        llm_service = LLMService(current_user, request_question, current_assistant, True)
+        llm_service = await LLMService.create(current_user, request_question, current_assistant, True)
         llm_service.set_record(record)
         llm_service.run_recommend_questions_task_async()
     except Exception as e:
@@ -142,7 +142,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
     """
 
     try:
-        llm_service = LLMService(current_user, request_question, current_assistant)
+        llm_service = await LLMService.create(current_user, request_question, current_assistant)
         llm_service.init_record()
         llm_service.run_task_async()
     except Exception as e:
@@ -189,7 +189,7 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch
     request_question = ChatQuestion(chat_id=record.chat_id, question=record.question)
 
     try:
-        llm_service = LLMService(current_user, request_question, current_assistant)
+        llm_service = await LLMService.create(current_user, request_question, current_assistant)
         llm_service.run_analysis_or_predict_task_async(action_type, record)
     except Exception as e:
         traceback.print_exc()

backend/apps/chat/task/llm.py

@@ -70,7 +70,7 @@ class LLMService:
     future: Future
 
     def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
-                 current_assistant: Optional[CurrentAssistant] = None, no_reasoning: bool = False):
+                 current_assistant: Optional[CurrentAssistant] = None, no_reasoning: bool = False, config: LLMConfig = None):
         self.chunk_list = []
         engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))
         session_maker = sessionmaker(bind=engine)
@@ -110,7 +110,7 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
 
         self.ds = (ds if isinstance(ds, AssistantOutDsSchema) else CoreDatasource(**ds.model_dump())) if ds else None
         self.chat_question = chat_question
-        self.config = get_default_config()
+        self.config = config
         if no_reasoning:
             # only work while using qwen
             if self.config.additional_params:
@@ -125,6 +125,14 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
         llm_instance = LLMFactory.create_llm(self.config)
         self.llm = llm_instance.llm
 
+        self.init_messages()
+    
+    @classmethod
+    async def create(cls, *args, **kwargs):
+        config: LLMConfig = await get_default_config()
+        instance = cls(*args, **kwargs, config=config)
+        return instance
+
     def is_running(self, timeout=0.5):
         try:
             r = concurrent.futures.wait([self.future], timeout)

backend/apps/mcp/mcp.py

@@ -107,7 +107,7 @@ async def mcp_question(session: SessionDep, chat: McpQuestion):
 
     mcp_chat = ChatMcp(token=chat.token, chat_id=chat.chat_id, question=chat.question)
     # ask
-    llm_service = LLMService(session_user, mcp_chat)
+    llm_service = await LLMService.create(session_user, mcp_chat)
     llm_service.init_record()
 
     return StreamingResponse(llm_service.run_task(False), media_type="text/event-stream")

backend/apps/system/api/aimodel.py

@@ -9,6 +9,7 @@
 
 from apps.system.models.system_model import AiModelDetail
 from common.core.deps import SessionDep, Trans
+from common.utils.crypto import sqlbot_decrypt
 from common.utils.time import get_timestamp
 from common.utils.utils import SQLBotLogUtil, prepare_model_arg
 
@@ -102,6 +103,10 @@ async def get_model_by_id(
             config_list = [AiModelConfigItem(**item) for item in raw]
         except Exception:
             pass
+    if db_model.api_key:
+        db_model.api_key = await sqlbot_decrypt(db_model.api_key)
+    if db_model.api_domain:
+        db_model.api_domain = await sqlbot_decrypt(db_model.api_domain)
     data = AiModelDetail.model_validate(db_model).model_dump(exclude_unset=True)
     data.pop("config", None)
     data["config_list"] = config_list

backend/apps/system/crud/aimodel_manage.py

@@ -0,0 +1,26 @@
+
+from apps.system.models.system_model import AiModelDetail
+from common.core.db import engine
+from sqlmodel import Session, select
+from common.utils.crypto import sqlbot_encrypt
+from common.utils.utils import SQLBotLogUtil
+
+async def async_model_info():
+    with Session(engine) as session:
+        model_list = session.exec(select(AiModelDetail)).all()
+        any_model_change = False
+        if model_list:
+            for model in model_list:
+                if model.api_domain.startswith("http"):
+                    if model.api_key:
+                        model.api_key = await sqlbot_encrypt(model.api_key)
+                    if model.api_domain:
+                        model.api_domain = await sqlbot_encrypt(model.api_domain)
+                    session.add(model)
+                    any_model_change = True
+        if any_model_change:
+            session.commit()
+            SQLBotLogUtil.info("✅ 异步加密已有模型的密钥和地址完成")           
+            
+            
+        

backend/common/utils/crypto.py

@@ -1,4 +1,7 @@
-from sqlbot_xpack.core import sqlbot_decrypt as xpack_sqlbot_decrypt
+from sqlbot_xpack.core import sqlbot_decrypt as xpack_sqlbot_decrypt, sqlbot_encrypt as xpack_sqlbot_encrypt
 
 async def sqlbot_decrypt(text: str) -> str:
-    return await xpack_sqlbot_decrypt(text)
+    return await xpack_sqlbot_decrypt(text)
+
+async def sqlbot_encrypt(text: str) -> str:
+    return await xpack_sqlbot_encrypt(text)

backend/main.py

@@ -10,6 +10,7 @@
 
 from alembic import command
 from apps.api import api_router
+from apps.system.crud.aimodel_manage import async_model_info
 from apps.system.crud.assistant import init_dynamic_cors
 from apps.system.middleware.auth import TokenMiddleware
 from apps.terminology.curd.terminology import fill_empty_embeddings
@@ -36,6 +37,7 @@ async def lifespan(app: FastAPI):
     init_embedding_data()
     SQLBotLogUtil.info("✅ SQLBot 初始化完成")
     await sqlbot_xpack.core.clean_xpack_cache()
+    await async_model_info()  # 异步加密已有模型的密钥和地址
     yield
     SQLBotLogUtil.info("SQLBot 应用关闭")
 

frontend/src/api/system.ts

@@ -3,8 +3,26 @@ import { request } from '@/utils/request'
 export const modelApi = {
   queryAll: (keyword?: string) =>
     request.get('/system/aimodel', { params: keyword ? { keyword } : {} }),
-  add: (data: any) => request.post('/system/aimodel', data),
-  edit: (data: any) => request.put('/system/aimodel', data),
+  add: (data: any) => {
+    const param = data
+    if (param.api_key) {
+      param.api_key = LicenseGenerator.sqlbotEncrypt(data.api_key)
+    }
+    if (param.api_domain) {
+      param.api_domain = LicenseGenerator.sqlbotEncrypt(data.api_domain)
+    }
+    return request.post('/system/aimodel', param)
+  },
+  edit: (data: any) => {
+    const param = data
+    if (param.api_key) {
+      param.api_key = LicenseGenerator.sqlbotEncrypt(data.api_key)
+    }
+    if (param.api_domain) {
+      param.api_domain = LicenseGenerator.sqlbotEncrypt(data.api_domain)
+    }
+    return request.put('/system/aimodel', param)
+  },
   delete: (id: number) => request.delete(`/system/aimodel/${id}`),
   query: (id: number) => request.get(`/system/aimodel/${id}`),
   setDefault: (id: number) => request.put(`/system/aimodel/default/${id}`),