feat: 更新聊天代理和任务管理器,支持通过线程ID管理消息,优化日志记录和敏感内容检查

This commit is contained in:
Wenjie Zhang 2025-10-12 13:09:29 +08:00
parent 54bb8e1e4f
commit 2ac690eb81
8 changed files with 132 additions and 86 deletions

View File

@ -2,15 +2,21 @@
## 内容安全
系统内置内容审查机制,保障服务内容的合规性。目前配置了关键词过滤以及 LLM 对内容进行审查。管理员可在 `设置``基本设置` 页面中进行配置并选择安全模型。
系统内置内容审查机制(默认是关闭状态),保障服务内容的合规性。目前配置了关键词过滤以及 LLM 对内容进行审查。管理员可在 `设置``基本设置` 页面中进行配置并选择安全模型。
敏感词词库位于 `src/config/static/bad_keywords.txt` 文件,每行一个关键词,实时生效,无需重启服务。
检测流程为,接收到用户输入之后,就对用户的输入进行检测是否合规,同时在流式传输的过程中进行实时检测(仅关键词)。当流式输出结束之后,则开始检测整个内容。
**注意**,使用 LLM 检测虽然可以大大缓解提示词注入带来的问题,但也会在用户交互上带来延迟影响,需要考虑是否启用。
对于关键词检测,敏感词词库位于 `src/config/static/bad_keywords.txt` 文件,每行一个关键词,实时生效,无需重启服务。
对于 LLM 检测Prompt 可以看到 `src/plugins/guard.py`
<<< @/../src/plugins/guard.py#guard_prompt
## 网页搜索
系统内置了基于 Tavily 的联网搜索能力,配置完成后,大模型会自动在需要时调用 `enable_web_search` 对应的工具,为回答提供实时网页信息。
### 配置步骤
1. 前往 [Tavily 官网](https://app.tavily.com/) 注册并在控制台创建 API Key。
2. 在项目根目录的 `.env`(或 `docker-compose.yml` 中的对应环境变量段)写入:

View File

@ -9,8 +9,9 @@
- [x] 部分 doc 格式的文件支持有问题
- [x] 当出现不支持的文件类型的时候,前端没有限制
- [ ] 当消息生成的时候有报错的时候,前端无显示
- [ ] 另外一个智能体的历史对话无法显示
- [ ] 调用统计的统计结果有问题Token 计算方法可能也不对)
- [x] 另外一个智能体的历史对话无法显示
- [ ] 调用统计的统计结果疑似有问题Token 计算方法可能也不对)
- [ ] 【重要】传输给 agent 的上下文消息有问题,需要基于新的 conv 构建阶段算法
---
@ -27,13 +28,13 @@
📝 **Base**
- [ ] 优化全局配置的管理模型,以及配置文件的管理,子配置的管理
- [ ] 新建 tasker 模块用来管理所有的后台任务UI 上使用侧边栏管理。
- [x] 新建 tasker 模块用来管理所有的后台任务UI 上使用侧边栏管理。
- [ ] 新增 files 模块,用来管理文件上传,下载等
## 未来可能会支持
下面的功能**可能**会放在后续版本实现,暂时未定
- [ ] 添加测试脚本,覆盖最常见的功能
- [ ] 优化对文档信息的检索展示(检索结果页、详情页)
- [x] 添加测试脚本覆盖最常见的功能已覆盖API
- [x] 优化对文档信息的检索展示(检索结果页、详情页)
- [ ] 集成 LangFuse (观望)添加用户日志与用户反馈模块,可以在 AgentView 中查看信息

View File

@ -20,8 +20,7 @@ cd Yuxi-Know
::: warning 版本说明
- `0.2.2`: 当前稳定版本(推荐)
- `stable`: 旧版本稳定分支(与现版本不兼容)
- `main`: 最新开发版本(可能不稳定)
- `main`: 最新开发版本(不稳定,新特性可能会导致新 bug
:::
#### 2. 配置环境变量

View File

@ -142,7 +142,7 @@ async def chat_agent(
async def save_messages_from_langgraph_state(
agent_instance,
conversation,
thread_id,
conv_mgr,
config_dict,
):
@ -162,11 +162,11 @@ async def chat_agent(
logger.debug(f"Retrieved {len(messages)} messages from LangGraph state")
# 获取已保存的消息数量,避免重复保存
existing_messages = conv_mgr.get_messages(conversation.id)
existing_messages = conv_mgr.get_messages_by_thread_id(thread_id)
existing_count = len(existing_messages)
# 只保存新增的消息
new_messages = messages[existing_count:]
new_messages = messages
for msg in new_messages:
msg_dict = msg.model_dump() if hasattr(msg, "model_dump") else {}
@ -190,8 +190,8 @@ async def chat_agent(
msg_dict["response_metadata"]["model_name"] = model_name[: len(model_name) // repeat_count]
# 保存 AI 消息
ai_msg = conv_mgr.add_message(
conversation_id=conversation.id,
ai_msg = conv_mgr.add_message_by_thread_id(
thread_id=thread_id,
role="assistant",
content=content,
message_type="text",
@ -236,6 +236,10 @@ async def chat_agent(
else:
logger.warning(f"Tool call {tool_call_id} not found for update")
else:
logger.warning(f"Unknown message type: {msg_type}, skipping")
continue
logger.debug(f"Processed message type={msg_type}")
logger.info(f"Saved {len(new_messages)} new messages from LangGraph state")
@ -249,7 +253,7 @@ async def chat_agent(
yield make_chunk(status="init", meta=meta, msg=HumanMessage(content=query).model_dump())
# Input guard
if conf.enable_content_guard and content_guard.check(query):
if conf.enable_content_guard and await content_guard.check(query):
yield make_chunk(status="error", message="输入内容包含敏感词", meta=meta)
return
@ -265,91 +269,74 @@ async def chat_agent(
# 构造运行时配置如果没有thread_id则生成一个
user_id = str(current_user.id)
thread_id = config.get("thread_id")
input_context = {"user_id": user_id, "thread_id": thread_id}
if not thread_id:
thread_id = str(uuid.uuid4())
logger.warning(f"No thread_id provided, generated new thread_id: {thread_id}")
# Initialize conversation manager
conv_manager = ConversationManager(db)
# Get or create conversation
conversation = None
if thread_id:
conversation = conv_manager.get_conversation_by_thread_id(thread_id)
if not conversation:
try:
# Auto-create conversation for existing thread
conversation = conv_manager.create_conversation(
user_id=user_id,
agent_id=agent_id,
title=(query[:50] + "..." if len(query) > 50 else query) if query else "新的对话",
thread_id=thread_id,
)
logger.info(f"Auto-created conversation for thread_id {thread_id}")
except Exception as e:
logger.error(f"Failed to auto-create conversation: {e}")
conversation = None
# Save user message
if conversation:
try:
conv_manager.add_message(
conversation_id=conversation.id,
role="user",
content=query,
message_type="text",
extra_metadata={"raw_message": HumanMessage(content=query).model_dump()},
)
except Exception as e:
logger.error(f"Error saving user message: {e}")
try:
conv_manager.add_message_by_thread_id(
thread_id=thread_id,
role="user",
content=query,
message_type="text",
extra_metadata={"raw_message": HumanMessage(content=query).model_dump()},
)
except Exception as e:
logger.error(f"Error saving user message: {e}")
try:
# Stream messages (only for display, don't save yet)
full_ai_content = ""
async for msg, metadata in agent.stream_messages(messages, input_context=input_context):
if isinstance(msg, AIMessageChunk):
# Content guard
if conf.enable_content_guard and content_guard.check(msg.content):
full_ai_content += msg.content
if conf.enable_content_guard and await content_guard.check_with_keywords(full_ai_content[-20:]):
logger.warning("Sensitive content detected in stream")
yield make_chunk(message="检测到敏感内容,已中断输出", status="error")
return
yield make_chunk(content=msg.content, msg=msg.model_dump(), metadata=metadata, status="loading")
elif isinstance(msg, ToolMessage):
yield make_chunk(msg=msg.model_dump(), metadata=metadata, status="loading")
else:
yield make_chunk(msg=msg.model_dump(), metadata=metadata, status="loading")
if conf.enable_content_guard and await content_guard.check(full_ai_content):
logger.warning("Sensitive content detected in final message")
yield make_chunk(message="检测到敏感内容,已中断输出", status="error")
return
yield make_chunk(status="finished", meta=meta)
# After streaming finished, save all messages from LangGraph state
if conversation:
langgraph_config = {"configurable": {"thread_id": thread_id, "user_id": user_id}}
await save_messages_from_langgraph_state(
agent_instance=agent,
conversation=conversation,
conv_mgr=conv_manager,
config_dict=langgraph_config,
)
langgraph_config = {"configurable": input_context}
await save_messages_from_langgraph_state(
agent_instance=agent,
thread_id=thread_id,
conv_mgr=conv_manager,
config_dict=langgraph_config,
)
except (asyncio.CancelledError, ConnectionError) as e:
# 客户端主动中断连接,尝试保存已生成的部分内容
logger.info(f"Client disconnected for thread {thread_id}: {e}")
try:
if conversation:
langgraph_config = {"configurable": {"thread_id": thread_id, "user_id": user_id}}
await save_messages_from_langgraph_state(
agent_instance=agent,
conversation=conversation,
conv_mgr=conv_manager,
config_dict=langgraph_config,
)
except Exception as save_error:
logger.error(f"Error saving partial messages after disconnect: {save_error}")
logger.warning(f"Client disconnected for thread {thread_id}: {e}")
langgraph_config = {"configurable": input_context}
await save_messages_from_langgraph_state(
agent_instance=agent,
thread_id=thread_id,
conv_mgr=conv_manager,
config_dict=langgraph_config,
)
# 通知前端中断(可能发送不到,但用于一致性)
try:
yield make_chunk(status="interrupted", message="对话已中断", meta=meta)
except Exception:
pass
return
yield make_chunk(status="interrupted", message="对话已中断", meta=meta)
except Exception as e:
logger.error(f"Error streaming messages: {e}, {traceback.format_exc()}")
yield make_chunk(message=f"Error streaming messages: {e}", status="error")

View File

@ -105,7 +105,7 @@ class Tasker:
worker = asyncio.create_task(self._worker_loop(), name="tasker-worker")
self._workers.append(worker)
self._started = True
logger.info("Tasker started with %s workers", self.worker_count)
logger.info("Tasker started with {} workers", self.worker_count)
async def shutdown(self) -> None:
async with self._lock:
@ -133,7 +133,7 @@ class Tasker:
self._tasks[task_id] = task
await self._persist_state()
await self._queue.put((task_id, coroutine))
logger.info("Enqueued task %s (%s)", task_id, name)
logger.info("Enqueued task {} ({})", task_id, name)
return task
async def list_tasks(self, status: Optional[str] = None) -> List[Dict[str, Any]]:
@ -159,7 +159,7 @@ class Tasker:
task.cancel_requested = True
task.updated_at = _utc_timestamp()
await self._persist_state()
logger.info("Cancellation requested for task %s", task_id)
logger.info("Cancellation requested for task {}", task_id)
return True
async def _worker_loop(self) -> None:
@ -191,7 +191,7 @@ class Tasker:
except asyncio.CancelledError:
await self._mark_cancelled(task_id, "任务被取消")
except Exception as exc: # noqa: BLE001
logger.error("Task %s failed: %s", task_id, exc, exc_info=True)
logger.exception("Task {} failed: {}", task_id, exc)
await self._update_task(
task_id,
status="failed",
@ -205,7 +205,7 @@ class Tasker:
except asyncio.CancelledError:
break
except Exception as exc: # noqa: BLE001
logger.error("Tasker worker error: %s", exc, exc_info=True)
logger.exception("Tasker worker error: {}", exc)
async def _get_task_instance(self, task_id: str) -> Optional[Task]:
async with self._lock:
@ -277,9 +277,9 @@ class Tasker:
task.message = "服务重启时任务未继续执行"
task.updated_at = _utc_timestamp()
self._tasks[task.id] = task
logger.info("Loaded %s task records from storage", len(tasks))
logger.info("Loaded {} task records from storage", len(tasks))
except Exception as exc: # noqa: BLE001
logger.error("Failed to load task state: %s", exc, exc_info=True)
logger.exception("Failed to load task state: {}", exc)
async def _persist_state(self) -> None:
tasks = [task.to_dict() for task in self._tasks.values()]

View File

@ -54,7 +54,7 @@ async def get_mcp_client(
logger.info(f"Initialized MCP client with servers: {list(configs.keys())}")
return client
except Exception as e:
logger.error("Failed to initialize MCP client: %s", e)
logger.error("Failed to initialize MCP client: {}", e)
return None

View File

@ -4,6 +4,7 @@ from src.config.app import config
from src.models import select_model
from src.utils import logger
# region guard_prompt
PROMPT_TEMPLATE = """
# 指令
你是一个内容合规性检测助手请根据提供的规则集判断以下内容是否符合合规性要求
@ -26,6 +27,7 @@ PROMPT_TEMPLATE = """
输入内容{content}
输出内容"""
# endregion guard_prompt
def load_keywords(file_path: str) -> list[str]:
@ -52,22 +54,38 @@ class ContentGuard:
else:
self.llm_model = None
def check(self, text: str) -> bool:
async 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 keywords_result := await self.check_with_keywords(text):
return keywords_result
if self.llm_model:
return await self.check_with_llm(text)
return False
async def check_with_keywords(self, text: str) -> bool:
"""
Checks if the text contains any sensitive keywords from the predefined list.
Returns True if sensitive content is found, False otherwise.
True: 不合规
False: 合规
"""
if not text:
return False
text_lower = text.lower()
for keyword in self.keywords:
if keyword in text_lower:
logger.debug(f"Keyword match found: {keyword}")
return True
return False
def check_with_llm(self, text: str) -> bool:
async 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.

View File

@ -116,6 +116,40 @@ class ConversationManager:
logger.debug(f"Added {role} message to conversation {conversation_id}")
return message
def add_message_by_thread_id(
self,
thread_id: str,
role: str,
content: str,
message_type: str = "text",
extra_metadata: dict | None = None,
) -> Message | None:
"""
Add a message to a conversation by thread ID
Args:
thread_id: Thread ID
role: Message role (user/assistant/system/tool)
content: Message content
message_type: Message type (text/tool_call/tool_result)
extra_metadata: Additional metadata (complete message dump)
Returns:
Created Message object or None if conversation not found
"""
conversation = self.get_conversation_by_thread_id(thread_id)
if not conversation:
logger.warning(f"Conversation not found for thread_id: {thread_id}")
return None
return self.add_message(
conversation_id=conversation.id,
role=role,
content=content,
message_type=message_type,
extra_metadata=extra_metadata,
)
def add_tool_call(
self,
message_id: int,
@ -196,6 +230,7 @@ class ConversationManager:
"""
conversation = self.get_conversation_by_thread_id(thread_id)
if not conversation:
logger.warning(f"Conversation not found for thread_id: {thread_id}")
return []
return self.get_messages(conversation.id, limit, offset)