From a027dfb44984be554d77af067f574bc6e698d487 Mon Sep 17 00:00:00 2001 From: ai-business-hql Date: Thu, 20 Nov 2025 14:18:09 +0800 Subject: [PATCH] feat: mem version1 --- backend/dao/session_message_table.py | 190 +++++++++++++++++++ backend/service/mcp_client.py | 8 +- backend/service/message_memory.py | 261 +++++++++++++++++++++++++++ backend/service/summary_agent.py | 151 ++++++++++++++++ 4 files changed, 609 insertions(+), 1 deletion(-) create mode 100644 backend/dao/session_message_table.py create mode 100644 backend/service/message_memory.py create mode 100644 backend/service/summary_agent.py diff --git a/backend/dao/session_message_table.py b/backend/dao/session_message_table.py new file mode 100644 index 0000000..d30a664 --- /dev/null +++ b/backend/dao/session_message_table.py @@ -0,0 +1,190 @@ +''' +Author: ai-business-hql ai.bussiness.hql@gmail.com +Date: 2025-11-19 +LastEditors: ai-business-hql ai.bussiness.hql@gmail.com +LastEditTime: 2025-11-19 +FilePath: /ComfyUI-Copilot/backend/dao/session_message_table.py +Description: Session message table for conversation memory management with compression +''' + +import os +import json +from typing import Dict, Any, Optional, List +from sqlalchemy import create_engine, Column, Integer, String, DateTime, Text +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import sessionmaker +from datetime import datetime + +# 创建数据库基类 +Base = declarative_base() + +# 定义session_message表模型 +class SessionMessage(Base): + __tablename__ = 'session_message' + + id = Column(Integer, primary_key=True, autoincrement=True) + session_id = Column(String(255), nullable=False, unique=True) # 每个session只有一条记录 + messages = Column(Text, nullable=False) # JSON字符串，存储完整消息列表 + index = Column(Integer, default=0) # 已压缩的消息数量，messages[:index]已被摘要 + summary = Column(Text, nullable=True) # 压缩后的历史摘要 + attributes = Column(Text, nullable=True) # JSON字符串，存储额外属性 + created_at = Column(DateTime, default=datetime.utcnow) + + def to_dict(self): + return { + 'id': self.id, + 'session_id': self.session_id, + 'messages': json.loads(self.messages) if self.messages else [], + 'index': self.index, + 'summary': self.summary, + 'attributes': json.loads(self.attributes) if self.attributes else None, + 'created_at': self.created_at.isoformat() if self.created_at else None + } + +class SessionMessageManager: + """会话消息管理器""" + + def __init__(self, db_path: str = None): + if db_path is None: + # 默认数据库路径 + current_dir = os.path.dirname(os.path.abspath(__file__)) + db_dir = os.path.join(current_dir, '..', 'data') + os.makedirs(db_dir, exist_ok=True) + db_path = os.path.join(db_dir, 'session_message.db') + + self.db_path = db_path + self.engine = create_engine(f'sqlite:///{db_path}', echo=False) + self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine) + + # 创建表 + Base.metadata.create_all(bind=self.engine) + + def get_session(self): + """获取数据库会话""" + return self.SessionLocal() + + def get_session_message(self, session_id: str) -> Optional[Dict[str, Any]]: + """获取指定session的消息记录""" + session = self.get_session() + try: + record = session.query(SessionMessage)\ + .filter(SessionMessage.session_id == session_id)\ + .first() + + if record: + return record.to_dict() + return None + finally: + session.close() + + def save_session_message( + self, + session_id: str, + messages: List[Dict[str, Any]], + index: int = 0, + summary: str = None, + attributes: Optional[Dict[str, Any]] = None + ) -> int: + """保存或更新会话消息记录""" + session = self.get_session() + try: + # 查找是否已存在 + record = session.query(SessionMessage)\ + .filter(SessionMessage.session_id == session_id)\ + .first() + + if record: + # 更新已有记录 + record.messages = json.dumps(messages, ensure_ascii=False) + record.index = index + record.summary = summary + if attributes: + record.attributes = json.dumps(attributes, ensure_ascii=False) + session.commit() + session.refresh(record) + return record.id + else: + # 创建新记录 + new_record = SessionMessage( + session_id=session_id, + messages=json.dumps(messages, ensure_ascii=False), + index=index, + summary=summary, + attributes=json.dumps(attributes, ensure_ascii=False) if attributes else None + ) + session.add(new_record) + session.commit() + session.refresh(new_record) + return new_record.id + except Exception as e: + session.rollback() + raise e + finally: + session.close() + + def update_summary(self, session_id: str, summary: str, index: int) -> bool: + """更新指定session的摘要和索引""" + session = self.get_session() + try: + record = session.query(SessionMessage)\ + .filter(SessionMessage.session_id == session_id)\ + .first() + + if record: + record.summary = summary + record.index = index + session.commit() + return True + return False + except Exception as e: + session.rollback() + raise e + finally: + session.close() + + def delete_session_message(self, session_id: str) -> bool: + """删除指定session的消息记录""" + session = self.get_session() + try: + record = session.query(SessionMessage)\ + .filter(SessionMessage.session_id == session_id)\ + .first() + + if record: + session.delete(record) + session.commit() + return True + return False + except Exception as e: + session.rollback() + raise e + finally: + session.close() + +# 全局会话消息管理器实例 +session_message_manager = SessionMessageManager() + +def get_session_message(session_id: str) -> Optional[Dict[str, Any]]: + """获取会话消息的便捷函数""" + return session_message_manager.get_session_message(session_id) + +def save_session_message( + session_id: str, + messages: List[Dict[str, Any]], + index: int = 0, + summary: str = None, + attributes: Optional[Dict[str, Any]] = None +) -> int: + """保存会话消息的便捷函数""" + return session_message_manager.save_session_message( + session_id, messages, index, summary, attributes + ) + +def update_summary(session_id: str, summary: str, index: int) -> bool: + """更新摘要的便捷函数""" + return session_message_manager.update_summary(session_id, summary, index) + +def delete_session_message(session_id: str) -> bool: + """删除会话消息的便捷函数""" + return session_message_manager.delete_session_message(session_id) + diff --git a/backend/service/mcp_client.py b/backend/service/mcp_client.py index c65c7fe..e4aeac4 100644 --- a/backend/service/mcp_client.py +++ b/backend/service/mcp_client.py @@ -2,7 +2,7 @@ Author: ai-business-hql qingli.hql@alibaba-inc.com Date: 2025-06-16 16:50:17 LastEditors: ai-business-hql ai.bussiness.hql@gmail.com -LastEditTime: 2025-11-18 19:18:30 +LastEditTime: 2025-11-19 20:16:03 FilePath: /comfyui_copilot/backend/service/mcp-client.py Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE ''' @@ -34,6 +34,7 @@ except Exception: from ..agent_factory import create_agent from ..service.workflow_rewrite_agent import create_workflow_rewrite_agent +from ..service.message_memory import message_memory_optimize from ..utils.request_context import get_rewrite_context, get_session_id, get_config from ..utils.logger import log from openai.types.responses import ResponseTextDeltaEvent @@ -117,6 +118,11 @@ async def comfyui_agent_invoke(messages: List[Dict[str, Any]], images: List[Imag if not config: raise ValueError("No config found in request context") + # Optimize messages with memory compression + log.info(f"[MCP] Original messages count: {len(messages)}") + messages = message_memory_optimize(session_id, messages) + log.info(f"[MCP] Optimized messages count: {len(messages)}, messages: {messages}") + # Create MCP server instances mcp_server = MCPServerSse( params= { diff --git a/backend/service/message_memory.py b/backend/service/message_memory.py new file mode 100644 index 0000000..008c0f1 --- /dev/null +++ b/backend/service/message_memory.py @@ -0,0 +1,261 @@ +''' +Author: ai-business-hql ai.bussiness.hql@gmail.com +Date: 2025-11-19 +LastEditors: ai-business-hql ai.bussiness.hql@gmail.com +LastEditTime: 2025-11-19 +FilePath: /ComfyUI-Copilot/backend/service/message_memory.py +Description: Message memory manager with automatic compression +''' + +from typing import List, Dict, Any, Tuple +from ..dao.session_message_table import ( + get_session_message, + save_session_message, + update_summary +) +from .summary_agent import generate_summary +from ..utils.logger import log + + +# 配置参数 +COMPRESSION_THRESHOLD = 8 # 超过10条未压缩消息时触发压缩（4个来回） +KEEP_RECENT_MESSAGES = 4 # 保留最新4条消息不压缩（2个来回） + + +def message_memory_optimize( + session_id: str, + messages: List[Dict[str, Any]] +) -> List[Dict[str, Any]]: + """ + 优化消息内存，自动压缩历史消息 + + 工作流程： + 1. 加载 session 数据（如果不存在则创建） + 2. 更新完整的消息列表 + 3. 检查是否需要压缩（未压缩消息 > 10条） + 4. 如果需要，执行压缩并更新数据库 + 5. 返回优化后的消息列表（summary + 未压缩的消息） + + Args: + session_id: 会话ID + messages: 当前的完整消息列表 + + Returns: + 优化后的消息列表，可以直接传给 LLM + - 如果有压缩：[{"role": "system", "content": "Previous summary: ..."}, ...recent_messages] + - 如果无压缩：原始 messages + """ + try: + log.info(f"[Memory] Starting memory optimization for session: {session_id}") + log.info(f"[Memory] Input messages count: {len(messages)}") + + # 1. 加载或创建 session 记录 + record = get_session_message(session_id) + + if record is None: + # 首次访问，创建新记录 + log.info(f"[Memory] New session detected, creating record") + save_session_message( + session_id=session_id, + messages=messages, + index=0, + summary=None + ) + # 新会话不需要压缩，直接返回原始消息 + return messages + + # 2. 更新消息列表（保存当前完整的消息） + current_index = record['index'] + current_summary = record['summary'] + + log.info(f"[Memory] Existing session - index: {current_index}, summary length: {len(current_summary) if current_summary else 0}") + + # 保存更新后的完整消息列表 + save_session_message( + session_id=session_id, + messages=messages, + index=current_index, # 暂时保持不变 + summary=current_summary + ) + + # 3. 检查是否需要压缩 + uncompressed_count = len(messages) - current_index + log.info(f"[Memory] Uncompressed messages: {uncompressed_count} (threshold: {COMPRESSION_THRESHOLD})") + + if uncompressed_count <= COMPRESSION_THRESHOLD: + # 不需要压缩，构建返回消息 + log.info(f"[Memory] No compression needed") + return _build_optimized_messages(current_summary, messages, current_index) + + # 4. 执行压缩 + log.info(f"[Memory] Compression triggered!") + + # 计算压缩范围 + # 待压缩：messages[current_index : len(messages) - KEEP_RECENT_MESSAGES] + compress_end_index = len(messages) - KEEP_RECENT_MESSAGES + messages_to_compress = messages[current_index:compress_end_index] + + log.info(f"[Memory] Compressing messages from index {current_index} to {compress_end_index}") + log.info(f"[Memory] Messages to compress: {len(messages_to_compress)}") + + # 生成新的摘要（整合历史摘要） + new_summary = generate_summary( + messages=messages_to_compress, + previous_summary=current_summary + ) + + log.info(f"[Memory] Generated new summary: {new_summary[:100]}...") + + # 更新数据库 + update_summary( + session_id=session_id, + summary=new_summary, + index=compress_end_index + ) + + log.info(f"[Memory] Updated index to: {compress_end_index}") + + # 5. 返回优化后的消息 + optimized_messages = _build_optimized_messages(new_summary, messages, compress_end_index) + log.info(f"[Memory] Returning {len(optimized_messages)} optimized messages") + + return optimized_messages + + except Exception as e: + log.error(f"[Memory] Error in message_memory_optimize: {str(e)}") + log.error(f"[Memory] Falling back to original messages") + # 发生错误时返回原始消息，确保系统继续运行 + return messages + + +def _build_optimized_messages( + summary: str, + full_messages: List[Dict[str, Any]], + index: int +) -> List[Dict[str, Any]]: + """ + 构建优化后的消息列表 + + Args: + summary: 压缩后的历史摘要 + full_messages: 完整的消息列表 + index: 已压缩的消息索引 + + Returns: + 优化后的消息列表 + """ + # 获取未压缩的消息 + recent_messages = full_messages[index:] + + if summary and len(summary.strip()) > 0 and len(recent_messages) > 0: + # 有摘要且有未压缩消息，将 summary 作为上下文信息添加到第一条消息前 + # 注意：不能使用 role="system"，因为 create_agent 已经有 instructions 了 + # 将 summary 作为一条独立的 user 消息添加在开头 + summary_context = { + "role": "user", + "content": f"[Context from previous conversation]: {summary}" + } + return [summary_context] + recent_messages + else: + # 无摘要或无消息，直接返回未压缩的消息 + return recent_messages + + +def get_optimized_messages(session_id: str) -> List[Dict[str, Any]]: + """ + 获取优化后的消息列表（不触发新的压缩） + + Args: + session_id: 会话ID + + Returns: + 优化后的消息列表 + """ + try: + record = get_session_message(session_id) + + if record is None: + return [] + + return _build_optimized_messages( + summary=record['summary'], + full_messages=record['messages'], + index=record['index'] + ) + + except Exception as e: + log.error(f"[Memory] Error in get_optimized_messages: {str(e)}") + return [] + + +def get_compression_stats(session_id: str) -> Dict[str, Any]: + """ + 获取压缩统计信息 + + Args: + session_id: 会话ID + + Returns: + 压缩统计信息字典 + """ + try: + record = get_session_message(session_id) + + if record is None: + return { + "exists": False, + "total_messages": 0, + "compressed_messages": 0, + "uncompressed_messages": 0, + "has_summary": False, + "summary_length": 0 + } + + total = len(record['messages']) + compressed = record['index'] + uncompressed = total - compressed + + return { + "exists": True, + "total_messages": total, + "compressed_messages": compressed, + "uncompressed_messages": uncompressed, + "has_summary": bool(record['summary']), + "summary_length": len(record['summary']) if record['summary'] else 0, + "compression_ratio": f"{compressed}/{total}" if total > 0 else "0/0" + } + + except Exception as e: + log.error(f"[Memory] Error in get_compression_stats: {str(e)}") + return {"error": str(e)} + + +# 测试函数 +def test_message_memory(): + """测试消息内存管理""" + test_session_id = "test_session_123" + + # 模拟对话 + messages = [] + for i in range(15): + messages.append({"role": "user", "content": f"User message {i+1}"}) + messages.append({"role": "assistant", "content": f"Assistant response {i+1}"}) + + print(f"Total messages: {len(messages)}") + + # 第一次优化 + optimized = message_memory_optimize(test_session_id, messages) + print(f"Optimized messages count: {len(optimized)}") + + # 获取统计信息 + stats = get_compression_stats(test_session_id) + print(f"Compression stats: {stats}") + + # 打印优化后的消息 + for msg in optimized[:3]: + print(f"{msg['role']}: {msg['content'][:50]}...") + + +if __name__ == "__main__": + test_message_memory() + diff --git a/backend/service/summary_agent.py b/backend/service/summary_agent.py new file mode 100644 index 0000000..4249b1b --- /dev/null +++ b/backend/service/summary_agent.py @@ -0,0 +1,151 @@ +''' +Author: ai-business-hql ai.bussiness.hql@gmail.com +Date: 2025-11-19 +LastEditors: ai-business-hql ai.bussiness.hql@gmail.com +LastEditTime: 2025-11-19 +FilePath: /ComfyUI-Copilot/backend/service/summary_agent.py +Description: Summary agent for compressing conversation history +''' + +import json +from typing import List, Dict, Any +from pydantic import BaseModel +from openai import OpenAI + +from ..utils.key_utils import workflow_config_adapt +from ..utils.globals import WORKFLOW_MODEL_NAME, get_comfyui_copilot_api_key, LLM_DEFAULT_BASE_URL +from ..utils.request_context import get_config +from ..utils.logger import log + + +class SummaryResponse(BaseModel): + """ + 对话历史摘要响应，必须是简洁的英文文本 + """ + summary: str + + +def generate_summary(messages: List[Dict[str, Any]], previous_summary: str = None) -> str: + """ + 生成对话历史摘要 + + Args: + messages: 需要摘要的消息列表，格式为 [{"role": "user/assistant", "content": "..."}] + previous_summary: 之前的摘要（如果有），新摘要会整合历史信息 + + Returns: + 生成的摘要文本，不超过 200 words + """ + try: + # 构建消息内容文本 + messages_text = "\n\n".join([ + f"**{msg['role'].upper()}**: {msg['content']}" + for msg in messages + if isinstance(msg.get('content'), str) # 只处理文本内容 + ]) + + # 构建用户提示 + if previous_summary: + user_prompt = f"""Please generate a concise summary of the following conversation history, integrating it with the previous summary. + +## Previous Summary +{previous_summary} + +## New Conversation Messages +{messages_text} + +## Requirements +- Maximum 200 words +- Focus on key topics, decisions, and important context +- Use clear, objective language +- Integrate previous summary seamlessly with new information +- Avoid redundancy +""" + else: + user_prompt = f"""Please generate a concise summary of the following conversation history. + +## Conversation Messages +{messages_text} + +## Requirements +- Maximum 200 words +- Focus on key topics, decisions, and important context +- Use clear, objective language +""" + + # 系统提示词 + system_prompt = """You are an expert conversation summarizer. Your task is to create concise, informative summaries of conversation histories that preserve the most important information while reducing token usage. + +Focus on: +- Main topics discussed +- Key decisions or conclusions +- Important technical details or requirements +- User's goals and preferences +- Context that would be needed for future interactions + +Keep summaries factual, objective, and efficient.""" + + # 获取配置 + config = get_config() + + # 创建OpenAI客户端 + client = OpenAI( + base_url=config.get("openai_base_url") or LLM_DEFAULT_BASE_URL, + api_key=config.get("openai_api_key") or get_comfyui_copilot_api_key() or "" + ) + + # 调用LLM生成摘要 + completion = client.chat.completions.parse( + model=config.get("model_select"), + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt} + ], + max_tokens=8192, # 限制输出长度 + temperature=0.3, # 较低的温度确保摘要一致性 + response_format=SummaryResponse + ) + + result = completion.choices[0].message.parsed + log.info(f"Generated summary: {result.summary[:100]}...") + + if result.summary: + return result.summary + else: + log.warning("Summary generation returned empty result") + return "" + + except Exception as e: + log.error(f"Failed to generate summary: {str(e)}") + # 返回一个简单的回退摘要 + fallback = f"Conversation with {len(messages)} messages" + if previous_summary: + fallback = f"{previous_summary} | {fallback}" + return fallback + + +def test_summary_agent(): + """测试摘要生成功能""" + test_messages = [ + {"role": "user", "content": "I want to create an image generation workflow using Stable Diffusion"}, + {"role": "assistant", "content": "I can help you with that. Let me search for relevant workflows..."}, + {"role": "user", "content": "I need it to support LoRA models"}, + {"role": "assistant", "content": "I've found some workflows with LoRA support. Here are the options..."}, + ] + + summary = generate_summary(test_messages) + print("Generated Summary:", summary) + + # 测试增量摘要 + new_messages = [ + {"role": "user", "content": "Can you also add upscaling?"}, + {"role": "assistant", "content": "Sure, I'll modify the workflow to include upscaling nodes..."}, + ] + + updated_summary = generate_summary(new_messages, previous_summary=summary) + print("Updated Summary:", updated_summary) + + +if __name__ == "__main__": + test_summary_agent() +