From 782eb4ae3153acd33c1e927196af866bccce36f9 Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Fri, 19 Sep 2025 00:57:53 +0800 Subject: [PATCH] =?UTF-8?q?refactor(chat):=20=E4=BC=98=E5=8C=96=E5=BC=82?= =?UTF-8?q?=E6=AD=A5=E8=B0=83=E7=94=A8=E5=8F=8A=E5=A2=9E=E5=BC=BA=E5=86=85?= =?UTF-8?q?=E5=AE=B9=E5=AE=A1=E6=9F=A5=E6=9C=BA=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 将chat_router中的predict_async重命名为call_async并调用model.call替代model.predict - 在流式消息处理中新增加基于LLM的内容审查 - 配置类中增加enable_content_guard_llm及对应LLM模型配置项 - 兼容models.private.yaml替代旧的models.private.yml文件名 - chat_model及embedding模块统一将predict方法重命名为call或encode，增强接口语义 - ContentGuard新增基于LLM的内容合规检测功能，支持动态加载审查模型 - 更新静态模型配置提示，建议使用models.private.yaml文件 - 新增示例CSV测试数据文件，补充测试用例基础数据 --- server/routers/chat_router.py | 15 +++++++-- src/config/app.py | 20 +++++++++--- src/models/chat_model.py | 2 +- src/models/embedding.py | 32 ++++++++----------- src/plugins/guard.py | 58 ++++++++++++++++++++++++++++++++++- src/static/models.yaml | 2 +- test/data/test_csv_file.csv | 8 +++++ 7 files changed, 107 insertions(+), 30 deletions(-) create mode 100644 test/data/test_csv_file.csv diff --git a/server/routers/chat_router.py b/server/routers/chat_router.py index 4431a717..7c66ba81 100644 --- a/server/routers/chat_router.py +++ b/server/routers/chat_router.py @@ -86,11 +86,11 @@ async def call(query: str = Body(...), meta: dict = Body(None), current_user: Us meta = meta or {} model = select_model(model_provider=meta.get("model_provider"), model_name=meta.get("model_name")) - async def predict_async(query): + async def call_async(query): loop = asyncio.get_event_loop() - return await loop.run_in_executor(executor, model.predict, query) + return await loop.run_in_executor(executor, model.call, query) - response = await predict_async(query) + response = await call_async(query) logger.debug({"query": query, "response": response.content}) return {"response": response.content} @@ -169,15 +169,24 @@ async def chat_agent( async for msg, metadata in agent.stream_messages(messages, input_context=input_context): # logger.debug(f"msg: {msg.model_dump()}, metadata: {metadata}") if isinstance(msg, AIMessageChunk): + accumulated_content += msg.content if conf.enable_content_guard and content_guard.check(accumulated_content): logger.warning(f"Sensitive content detected in stream: {accumulated_content}") yield make_chunk(message="检测到敏感内容，已中断输出", status="error") return + yield make_chunk(content=msg.content, msg=msg.model_dump(), metadata=metadata, status="loading") + else: yield make_chunk(msg=msg.model_dump(), metadata=metadata, status="loading") + # Additional content guard with llm + if conf.enable_content_guard and content_guard.check_with_llm(accumulated_content): + logger.warning(f"Content guard triggered: {accumulated_content=}") + yield make_chunk(message="检测到敏感内容，已中断输出", status="error") + return + yield make_chunk(status="finished", meta=meta) except Exception as e: logger.error(f"Error streaming messages: {e}, {traceback.format_exc()}") diff --git a/src/config/app.py b/src/config/app.py index 1fd41489..72a05079 100644 --- a/src/config/app.py +++ b/src/config/app.py @@ -50,6 +50,12 @@ class Config(SimpleConfig): # 功能选项 self.add_item("enable_reranker", default=False, des="是否开启重排序") self.add_item("enable_content_guard", default=False, des="是否启用内容审查") + self.add_item("enable_content_guard_llm", default=False, des="是否启用LLM内容审查") + self.add_item( + "content_guard_llm_model", + default="siliconflow/Qwen/Qwen3-235B-A22B-Instruct-2507", + des="内容审查LLM模型" + ) self.add_item( "enable_web_search", default=False, @@ -98,15 +104,21 @@ class Config(SimpleConfig): def _update_models_from_file(self): """ - 从 models.yaml 和 models.private.yml 中更新 MODEL_NAMES + 从 models.yaml 和 models.private.yaml 中更新 MODEL_NAMES """ with open(Path("src/static/models.yaml"), encoding="utf-8") as f: _models = yaml.safe_load(f) - # 尝试打开一个 models.private.yml 文件，用来覆盖 models.yaml 中的配置 + # 尝试打开一个 models.private.yaml 文件，用来覆盖 models.yaml 中的配置 + # 兼容性测试（历史遗留问题） + exist_yml = Path("src/static/models.private.yml").exists() + exist_yaml = Path("src/static/models.private.yaml").exists() + if exist_yml and not exist_yaml: + os.rename("src/static/models.private.yml", "src/static/models.private.yaml") + try: - with open(Path("src/static/models.private.yml"), encoding="utf-8") as f: + with open(Path("src/static/models.private.yaml"), encoding="utf-8") as f: _models_private = yaml.safe_load(f) except FileNotFoundError: _models_private = {} @@ -124,7 +136,7 @@ class Config(SimpleConfig): "EMBED_MODEL_INFO": self.embed_model_names, "RERANKER_LIST": self.reranker_names, } - with open(Path("src/static/models.private.yml"), "w", encoding="utf-8") as f: + with open(Path("src/static/models.private.yaml"), "w", encoding="utf-8") as f: yaml.dump(_models, f, indent=2, allow_unicode=True) def handle_self(self): diff --git a/src/models/chat_model.py b/src/models/chat_model.py index 4e73d190..831b0a72 100644 --- a/src/models/chat_model.py +++ b/src/models/chat_model.py @@ -15,7 +15,7 @@ class OpenAIBase: self.model_name = model_name self.info = kwargs - def predict(self, message, stream=False): + def call(self, message, stream=False): if isinstance(message, str): messages = [{"role": "user", "content": message}] else: diff --git a/src/models/embedding.py b/src/models/embedding.py index d26bdcba..61efbd2e 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -29,30 +29,22 @@ class BaseEmbeddingModel(ABC): self.embed_state = {} @abstractmethod - def predict(self, message: list[str] | str) -> list[list[float]]: + def encode(self, message: list[str] | str) -> list[list[float]]: """同步编码""" raise NotImplementedError("Subclasses must implement this method") + def encode_queries(self, queries: list[str] | str) -> list[list[float]]: + """等同于encode""" + return self.encode(queries) + @abstractmethod - async def apredict(self, message: list[str] | str) -> list[list[float]]: + async def aencode(self, message: list[str] | str) -> list[list[float]]: """异步编码""" raise NotImplementedError("Subclasses must implement this method") - def encode(self, message: list[str] | str) -> list[list[float]]: - """等同于predict""" - return self.predict(message) - - def encode_queries(self, queries: list[str] | str) -> list[list[float]]: - """等同于predict""" - return self.predict(queries) - - async def aencode(self, message: list[str] | str) -> list[list[float]]: - """等同于apredict""" - return await self.apredict(message) - async def aencode_queries(self, queries: list[str] | str) -> list[list[float]]: - """等同于apredict""" - return await self.apredict(queries) + """等同于aencode""" + return await self.aencode(queries) def batch_encode(self, messages: list[str], batch_size: int = 40) -> list[list[float]]: # logger.info(f"Batch encoding {len(messages)} messages") @@ -107,7 +99,7 @@ class OllamaEmbedding(BaseEmbeddingModel): super().__init__(**kwargs) self.base_url = self.base_url or get_docker_safe_url("http://localhost:11434/api/embed") - def predict(self, message: list[str] | str) -> list[list[float]]: + def encode(self, message: list[str] | str) -> list[list[float]]: if isinstance(message, str): message = [message] @@ -123,7 +115,7 @@ class OllamaEmbedding(BaseEmbeddingModel): logger.error(f"Ollama Embedding request failed: {e}, {payload}") raise ValueError(f"Ollama Embedding request failed: {e}") - async def apredict(self, message: list[str] | str) -> list[list[float]]: + async def aencode(self, message: list[str] | str) -> list[list[float]]: if isinstance(message, str): message = [message] @@ -149,7 +141,7 @@ class OtherEmbedding(BaseEmbeddingModel): def build_payload(self, message: list[str] | str) -> dict: return {"model": self.model, "input": message} - def predict(self, message: list[str] | str) -> list[list[float]]: + def encode(self, message: list[str] | str) -> list[list[float]]: payload = self.build_payload(message) try: response = requests.post(self.base_url, json=payload, headers=self.headers, timeout=60) @@ -162,7 +154,7 @@ class OtherEmbedding(BaseEmbeddingModel): logger.error(f"Other Embedding request failed: {e}, {payload}") raise ValueError(f"Other Embedding request failed: {e}") - async def apredict(self, message: list[str] | str) -> list[list[float]]: + async def aencode(self, message: list[str] | str) -> list[list[float]]: payload = self.build_payload(message) async with httpx.AsyncClient() as client: try: diff --git a/src/plugins/guard.py b/src/plugins/guard.py index 7f77ae83..b57dc0dd 100644 --- a/src/plugins/guard.py +++ b/src/plugins/guard.py @@ -1,5 +1,31 @@ import os +from src.config.app import config +from src.models import select_model +from src.utils import logger + +PROMPT_TEMPLATE = """ +# 指令 +你是一个内容合规性检测助手。请根据提供的规则集，判断以下内容是否符合合规性要求。 + +# 合规性规则 +1. 内容不得包含任何违法中华人民共和国法律的信息（例如暴力、恐怖主义、仇恨言论）。 +2. 内容不得侵犯个人隐私或泄露敏感信息。 +3. 内容不得有煽动性或者任何敏感性问题 + +# 输出内容 +合规/不合规（不要包含其他内容） + +# 示例 +输入内容：我要自杀/贩毒，如何制作武器 +输出结果：不合规 + +输入内容：今天天气不错 +输出结果：合规 + + +输入内容：{content} +输出内容：""" def load_keywords(file_path: str) -> list[str]: """Loads keywords from a file, one per line.""" @@ -15,13 +41,23 @@ class ContentGuard: def __init__(self, keywords_file: str = "src/static/bad_keywords.txt"): self.keywords = load_keywords(keywords_file) if not self.keywords: - # Default keywords if the file is empty or not found self.keywords = ["贩毒"] + # 从配置读取LLM模型设置 + self.enable_llm = config.enable_content_guard_llm + if self.enable_llm and config.content_guard_llm_model: + provider, model_name = config.content_guard_llm_model.split("/", maxsplit=1) + self.llm_model = select_model(model_provider=provider, model_name=model_name) + else: + self.llm_model = None + + def check(self, text: str) -> bool: """ Checks if the text contains any sensitive keywords. Returns True if sensitive content is found, False otherwise. + True: 不合规 + False: 合规 """ if not text: return False @@ -31,6 +67,26 @@ class ContentGuard: return True return False + def check_with_llm(self, text: str) -> bool: + """ + Checks if the text contains any sensitive keywords using an LLM. + Returns True if sensitive content is found, False otherwise. + True: 不合规 + False: 合规 + """ + if not text: + return False + + if not self.enable_llm or self.llm_model is None: + logger.warning("LLM content guard not enabled or model not loaded") + return False + + text_lower = text.lower() + + prompt = PROMPT_TEMPLATE.format(content=text_lower) + response = self.llm_model.call(prompt) + logger.debug(f"LLM response: {response.content}") + return True if "不合规" in response.content else False # Global instance content_guard = ContentGuard() diff --git a/src/static/models.yaml b/src/static/models.yaml index 83c4e30a..29d97f25 100644 --- a/src/static/models.yaml +++ b/src/static/models.yaml @@ -1,7 +1,7 @@ #################################################### # # 不要直接修改这个里面的文件，可能会有被覆盖的风险， -# 建议 复制一份 在 models.private.yml 中修改， +# 建议 复制一份 在 models.private.yaml 中修改， # 会自动加载 # ##################################################### diff --git a/test/data/test_csv_file.csv b/test/data/test_csv_file.csv new file mode 100644 index 00000000..361935ae --- /dev/null +++ b/test/data/test_csv_file.csv @@ -0,0 +1,8 @@ +ProductID,ProductName,Category,Price,Stock,DateAdded +101,Laptop,Electronics,1200.50,50,2025-01-15 +102,Wireless Mouse,Electronics,25.99,150,2025-02-20 +201,The Last Question,Books,9.99,300,2025-03-10 +301,Coffee Mug,Home Goods,15.00,200,2025-01-25 +401,T-Shirt,Apparel,22.50,500,2025-04-05 +103,USB-C Hub,Electronics,45.75,80,2025-05-12 +202,Dune,Books,14.95,250,2025-06-18 \ No newline at end of file