1  # Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  #
  3  # Licensed under the Apache License, Version 2.0 (the "License");
  4  # you may not use this file except in compliance with the License.
  5  # You may obtain a copy of the License at
  6  #
  7  #     http://www.apache.org/licenses/LICENSE-2.0
  8  #
  9  # Unless required by applicable law or agreed to in writing, software
 10  # distributed under the License is distributed on an "AS IS" BASIS,
 11  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  # See the License for the specific language governing permissions and
 13  # limitations under the License.
 14  #
 15  # Requirement: Any integration or derivative work must explicitly attribute
 16  # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  # documentation or user interface, as detailed in the NOTICE file.
 18  
 19  import sys
 20  import gettext
 21  import os
 22  from typing import Literal, Union
 23  from datetime import datetime
 24  from pydantic import BaseModel
 25  from loguru import logger as base_logger
 26  
 27  class contentSchema(BaseModel):
 28      timestamp: str | None = None
 29  
 30  class newPlanStep(contentSchema):
 31      stepId: str
 32      title: str
 33  
 34  class statusUpdate(contentSchema):
 35      stepId: str
 36      brief: str
 37      description: str
 38      status: Literal["running", "completed", "failed"]
 39  
 40  class toolUsed(contentSchema):
 41      stepId: str
 42      tool_id: str
 43      tool_name: str | None = None
 44      brief: str
 45      status: Literal["todo", "doing", "done"]
 46  
 47  class actionLog(contentSchema):
 48      tool_id: str
 49      tool_name: str
 50      stepId: str
 51      log: str
 52  
 53  class resultUpdate(contentSchema):
 54      msgType: Literal["text", "markdown", "file", "json"]
 55      content: str | dict | list
 56      status: bool | None = None
 57  
 58  class PromptSecurityLog(BaseModel):
 59      type: Literal["error", "newPlanStep", "statusUpdate", "toolUsed", "actionLog", "resultUpdate"]
 60      content: Union[str, newPlanStep, statusUpdate, toolUsed, actionLog, resultUpdate]
 61  
 62  class PromptSecurityLogger:
 63      def __init__(self, base_logger, lang='en_US'):
 64          self._base_logger = base_logger
 65          self._base_logger.remove()
 66          self._base_logger.add(sys.stdout, filter=lambda record: not record["extra"].get("aig_log", False), level="DEBUG")
 67          self._base_logger.add(sys.stdout, filter=lambda record: record["extra"].get("aig_log", False), format="{message}")
 68          self.lang = lang
 69          self._setup_i18n()
 70      
 71      def add(self, *args, **kwargs):
 72          self._base_logger.add(*args, **kwargs)
 73  
 74      def info(self, *args, **kwargs):
 75          self._base_logger.opt(depth=1).info(*args, **kwargs)
 76  
 77      def debug(self, *args, **kwargs):
 78          self._base_logger.opt(depth=1).debug(*args, **kwargs)
 79  
 80      def warning(self, *args, **kwargs):
 81          self._base_logger.opt(depth=1).warning(*args, **kwargs)
 82  
 83      def error(self, *args, **kwargs):
 84          self._base_logger.opt(depth=1).error(*args, **kwargs)
 85      
 86      def exception(self, *args, **kwargs):
 87          self._base_logger.opt(depth=1).exception(*args, **kwargs)
 88  
 89      def disable(self):
 90          self._base_logger.disable("")
 91      
 92      def enable(self):
 93          self._base_logger.enable("")
 94  
 95      def _setup_i18n(self):
 96          localedir = os.path.join(os.path.dirname(__file__), 'locales')
 97          try:
 98              self.trans = gettext.translation('messages', localedir=localedir, languages=[self.lang])
 99              self.lang = self.trans.info().get("language", self.lang)
100              self.trans.install()
101              self._ = self.trans.gettext
102          except FileNotFoundError:
103              gettext.install('messages')
104              self._ = gettext.gettext
105      
106      def set_language(self, lang):
107          """动态切换语言"""
108          self.lang = lang
109          self._setup_i18n()
110  
111      def translated_msg(self, msg, *args, **kwargs):
112          translated_msg = self._(msg)
113          if args or kwargs:
114              translated_msg = translated_msg.format(*args, **kwargs)
115          return translated_msg
116  
117      def _create_log(self, log_type: str, content: Union[str, contentSchema]) -> str:
118          """创建符合PromptSecurityLog格式的日志"""
119          if isinstance(content, contentSchema) and "timestamp" not in content:
120              content.timestamp = datetime.now().isoformat()
121          
122          log_entry = PromptSecurityLog(
123              type=log_type,
124              content=content
125          )
126          return log_entry.model_dump_json(exclude_none=True)
127      
128      def log(self, log_type: str, content: Union[str, contentSchema]):
129          """记录日志"""
130          log_message = self._create_log(log_type, content)
131          self._base_logger.bind(aig_log=True).opt(depth=2).log("INFO", "\n" + log_message)
132      
133      # 为每种日志类型创建便捷方法
134      def critical_issue(self, content: str):
135          self.log("error", content)
136          
137      def new_plan_step(self, content: newPlanStep):
138          self.log("newPlanStep", content)
139      
140      def status_update(self, content: statusUpdate):
141          self.log("statusUpdate", content)
142      
143      def tool_used(self, content: toolUsed):
144          self.log("toolUsed", content)
145      
146      def action_log(self, content: actionLog):
147          self.log("actionLog", content)
148      
149      def result_update(self, content: resultUpdate):
150          self.log("resultUpdate", content)
151  
152  logger = PromptSecurityLogger(base_logger)