diff --git a/backend/alembic/versions/040_modify_ai_model.py b/backend/alembic/versions/040_modify_ai_model.py new file mode 100644 index 00000000..dc924881 --- /dev/null +++ b/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 + ) diff --git a/backend/apps/ai_model/model_factory.py b/backend/apps/ai_model/model_factory.py index e0186e9f..03479fd8 100644 --- a/backend/apps/ai_model/model_factory.py +++ b/backend/apps/ai_model/model_factory.py @@ -10,6 +10,7 @@ from sqlmodel import Session, select 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 @@ class LLMFactory: 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( diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 5cc0e125..f4a0e640 100644 --- a/backend/apps/chat/api/chat.py +++ b/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() diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 00c9e618..a5c3caa1 100644 --- a/backend/apps/chat/task/llm.py +++ b/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 @@ class LLMService: 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 @@ class LLMService: 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) diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index 19afffe7..15878792 100644 --- a/backend/apps/mcp/mcp.py +++ b/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") diff --git a/backend/apps/system/api/aimodel.py b/backend/apps/system/api/aimodel.py index c33fd4ee..a96354da 100644 --- a/backend/apps/system/api/aimodel.py +++ b/backend/apps/system/api/aimodel.py @@ -9,6 +9,7 @@ from sqlmodel import func, select, update 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 diff --git a/backend/apps/system/crud/aimodel_manage.py b/backend/apps/system/crud/aimodel_manage.py new file mode 100644 index 00000000..77d6cb5e --- /dev/null +++ b/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("✅ 异步加密已有模型的密钥和地址完成") + + + \ No newline at end of file diff --git a/backend/common/utils/crypto.py b/backend/common/utils/crypto.py index ce5d7b95..dcd1a13c 100644 --- a/backend/common/utils/crypto.py +++ b/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) \ No newline at end of file + return await xpack_sqlbot_decrypt(text) + +async def sqlbot_encrypt(text: str) -> str: + return await xpack_sqlbot_encrypt(text) \ No newline at end of file diff --git a/backend/main.py b/backend/main.py index 8ca5a69e..6ec2771f 100644 --- a/backend/main.py +++ b/backend/main.py @@ -10,6 +10,7 @@ from starlette.middleware.cors import CORSMiddleware 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 应用关闭") diff --git a/frontend/src/api/system.ts b/frontend/src/api/system.ts index 9713ec9e..aa7db56d 100644 --- a/frontend/src/api/system.ts +++ b/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}`),