improve: add lang parameter in mcp chat

Author: ulleo

backend/apps/chat/models/chat_model.py

@@ -293,6 +293,7 @@ class McpQuestion(BaseModel):
     chat_id: int = Body(description='会话ID')
     token: str = Body(description='token')
     stream: Optional[bool] = Body(description='是否流式输出，默认为true开启, 关闭false则返回JSON对象', default=True)
+    lang: Optional[str] = Body(description='语言：zh-CN|en|ko-KR', default='zh-CN')
 
 
 class AxisObj(BaseModel):

backend/apps/chat/task/llm.py

@@ -49,6 +49,7 @@
 from common.core.deps import CurrentAssistant, CurrentUser
 from common.error import SingleMessageError, SQLBotDBError, ParseSQLResultError, SQLBotDBConnectionError
 from common.utils.data_format import DataFormat
+from common.utils.locale import I18n, I18nHelper
 from common.utils.utils import SQLBotLogUtil, extract_nested_json, prepare_for_orjson
 
 warnings.filterwarnings("ignore")
@@ -62,6 +63,8 @@
 
 session_maker = scoped_session(sessionmaker(bind=engine, class_=Session))
 
+i18n = I18n()
+
 
 class LLMService:
     ds: CoreDatasource
@@ -86,6 +89,8 @@ class LLMService:
     chunk_list: List[str] = []
     future: Future
 
+    trans: I18nHelper = None
+
     last_execute_sql_error: str = None
     articles_number: int = 4
 
@@ -125,6 +130,7 @@ def __init__(self, session: Session, current_user: CurrentUser, chat_question: C
         self.change_title = not get_chat_brief_generate(session=session, chat_id=chat_id)
 
         chat_question.lang = get_lang_name(current_user.language)
+        self.trans = i18n(lang=current_user.language)
 
         self.ds = (
             ds if isinstance(ds, AssistantOutDsSchema) else CoreDatasource(**ds.model_dump())) if ds else None
@@ -972,7 +978,7 @@ def run_task(self, in_chat: bool = True, stream: bool = True,
                     {'type': 'question', 'question': self.get_record().question}).decode() + '\n\n'
             else:
                 if stream:
-                    yield '> ID: ' + str(self.get_record().id) + '\n'
+                    yield '> ' + self.trans('i18n_chat.record_id_in_mcp') + str(self.get_record().id) + '\n'
                     yield '> ' + self.get_record().question + '\n\n'
             if not stream:
                 json_result['record_id'] = self.get_record().id
@@ -1173,7 +1179,7 @@ def run_task(self, in_chat: bool = True, stream: bool = True,
                 # generate picture
                 try:
                     if chart.get('type') != 'table':
-                        yield '### generated chart picture\n\n'
+                        # yield '### generated chart picture\n\n'
                         image_url, error = request_picture(self.record.chat_id, self.record.id, chart,
                                                            format_json_data(result))
                         SQLBotLogUtil.info(image_url)
@@ -1185,6 +1191,8 @@ def run_task(self, in_chat: bool = True, stream: bool = True,
                             raise error
                 except Exception as e:
                     if stream:
+                        if chart.get('type') != 'table':
+                            yield 'generate or fetch chart picture error.\n\n'
                         raise e
 
             if not stream:
@@ -1263,7 +1271,7 @@ def run_analysis_or_predict_task(self, action_type: str, in_chat: bool = True, s
                 yield 'data:' + orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n'
             else:
                 if stream:
-                    yield '> ID: ' + str(self.get_record().id) + '\n'
+                    yield '> ' + self.trans('i18n_chat.record_id_in_mcp') + str(self.get_record().id) + '\n'
                     yield '> ' + self.get_record().question + '\n\n'
             if not stream:
                 json_result['record_id'] = self.get_record().id
@@ -1331,7 +1339,7 @@ def run_analysis_or_predict_task(self, action_type: str, in_chat: bool = True, s
                         # generate picture
                         try:
                             if chart.get('type') != 'table':
-                                yield '### generated chart picture\n\n'
+                                # yield '### generated chart picture\n\n'
 
                                 _data = get_chat_chart_data(_session, self.record.id)
                                 _data['data'] = _data.get('data') + predict_data
@@ -1347,6 +1355,8 @@ def run_analysis_or_predict_task(self, action_type: str, in_chat: bool = True, s
                                     raise error
                         except Exception as e:
                             if stream:
+                                if chart.get('type') != 'table':
+                                    yield 'generate or fetch chart picture error.\n\n'
                                 raise e
                 else:
                     if in_chat:

backend/apps/mcp/mcp.py

@@ -91,6 +91,7 @@ async def mcp_question(session: SessionDep, chat: McpQuestion):
     db_user: UserModel = get_db_user(session=session, user_id=token_data.id)
     session_user = UserInfoDTO.model_validate(db_user.model_dump())
     session_user.isAdmin = session_user.id == 1 and session_user.account == 'admin'
+    session_user.language = chat.lang
     if session_user.isAdmin:
         session_user = session_user
     ws_model: UserWsModel = session.exec(

backend/common/utils/locale.py

@@ -1,8 +1,10 @@
-from pathlib import Path
 import json
-from typing import Dict, Optional, Any
+from pathlib import Path
+from typing import Dict, Any
+
 from fastapi import Request
 
+
 class I18n:
     def __init__(self, locale_dir: str = "locales"):
         self.locale_dir = Path(locale_dir)
@@ -18,40 +20,45 @@ def load_translations(self):
             with open(lang_file, 'r', encoding='utf-8') as f:
                 self.translations[lang_file.stem.lower()] = json.load(f)
 
-    def get_language(self, request: Request) -> str:
-        accept_language = request.headers.get('accept-language', 'en')
-        primary_lang = accept_language.split(',')[0].lower()
-        
+    def get_language(self, request: Request = None, lang: str = None) -> str:
+        primary_lang: str | None = None
+        if lang is not None:
+            primary_lang = lang.lower()
+        elif request is not None:
+            accept_language = request.headers.get('accept-language', 'en')
+            primary_lang = accept_language.split(',')[0].lower()
+
         return primary_lang if primary_lang in self.translations else 'zh-cn'
 
-    def __call__(self, request: Request) -> 'I18nHelper':
-        return I18nHelper(self, request)
+    def __call__(self, request: Request = None, lang: str = None) -> 'I18nHelper':
+        return I18nHelper(self, request, lang)
+
 
 class I18nHelper:
-    def __init__(self, i18n: I18n, request: Request):
+    def __init__(self, i18n: I18n, request: Request = None, lang: str = None):
         self.i18n = i18n
         self.request = request
-        self.lang = i18n.get_language(request)
+        self.lang = i18n.get_language(request, lang)
 
     def _get_nested_translation(self, data: Dict[str, Any], key_path: str) -> str:
         keys = key_path.split('.')
         current = data
-        
+
         for key in keys:
             if isinstance(current, dict) and key in current:
                 current = current[key]
             else:
                 return key_path  # 如果找不到，返回原键
-        
+
         return current if isinstance(current, str) else key_path
 
     def __call__(self, arg_key: str, **kwargs) -> str:
         lang_data = self.i18n.translations.get(self.lang, {})
         text = self._get_nested_translation(lang_data, arg_key)
-        
+
         if kwargs:
             try:
                 return text.format(**kwargs)
             except (KeyError, ValueError):
                 return text
-        return text
+        return text

backend/locales/en.json

@@ -40,6 +40,9 @@
   "i18n_embedded": {
     "invalid_origin": "Domain name validation failed【{origin}】"
   },
+  "i18n_chat": {
+    "record_id_in_mcp": "Answer ID: "
+  },
   "i18n_terminology": {
     "terminology_not_exists": "This terminology does not exist",
     "datasource_cannot_be_none": "Datasource cannot be empty",

backend/locales/ko-KR.json

@@ -40,6 +40,9 @@
   "i18n_embedded": {
     "invalid_origin": "도메인 이름 검증 실패 【{origin}】"
   },
+  "i18n_chat": {
+    "record_id_in_mcp": "응답 ID: "
+  },
   "i18n_terminology": {
     "datasource_list_is_not_found": "데이터 소스 목록을 찾을 수 없습니다",
     "datasource_id_not_found": "데이터 소스 ID: {key}를 찾을 수 없습니다",

backend/locales/zh-CN.json

@@ -40,6 +40,9 @@
   "i18n_embedded": {
     "invalid_origin": "域名校验失败【{origin}】"
   },
+  "i18n_chat": {
+    "record_id_in_mcp": "响应ID: "
+  },
   "i18n_terminology": {
     "terminology_not_exists": "该术语不存在",
     "datasource_cannot_be_none": "数据源不能为空",