from typing import Dict, List, Optional, Tuple, Sequence, Any import json import re import os from dataclasses import dataclass from logging import getLogger try: from openai import OpenAI except ImportError: OpenAI = None from open_webui.models.chats import Chats from open_webui.config import OPENAI_API_KEYS, OPENAI_API_BASE_URLS log = getLogger(__name__) # --- Constants & Prompts from persona_extractor --- SUMMARY_PROMPT = """你是一名“对话历史整理员”，请在保持事实准确的前提下，概括当前为止的聊天记录。 ## 要求 1. 最终摘要不得超过 1000 字。 2. 聚焦人物状态、事件节点、情绪/意图等关键信息，将片段整合为连贯文字。 3. 输出需包含 who / how / why / what 四个字段，每项不超过 50 字。 4. 禁止臆测或脏话，所有内容都必须能在聊天中找到对应描述。 5. 目标：帮助后续对话快速回忆上下文与人物设定。 已存在的摘要（如无则写“无”）： {existing_summary} 聊天片段： ---CHATS--- {chat_transcript} ---END--- 请严格输出下列 JSON： {{ "summary": "不超过1000字的连贯摘要", "table": {{ "who": "不超过50字", "how": "不超过50字", "why": "不超过50字", "what": "不超过50字" }} }} """ MERGE_ONLY_PROMPT = """你是一名“对话历史整理员”。 请将以下两段对话摘要（A 和 B）合并为一段连贯的、更新后的对话历史摘要。 摘要 A 是较早的时间段，摘要 B 是较新的时间段。 【摘要 A (旧)】 {summary_a} 【摘要 B (新)】 {summary_b} ## 要求 1. 保持时间线的连贯性，将新发生的事自然接续在旧事之后。 2. 最终摘要不得超过 1000 字。 3. 依然提取 who / how / why / what 四个关键要素（基于合并后的全貌）。 4. 禁止臆测，只基于提供的摘要内容。 请严格输出下列 JSON： {{ "summary": "合并后的连贯摘要", "table": {{ "who": "不超过50字", "how": "不超过50字", "why": "不超过50字", "what": "不超过50字" }} }} """ @dataclass class HistorySummary: summary: str table: dict[str, str] @dataclass(slots=True) class ChatMessage: role: str content: str timestamp: Optional[Any] = None def formatted(self) -> str: return f"{self.role}: {self.content}" class HistorySummarizer: def __init__( self, *, client: Optional[Any] = None, model: str = "gpt-4.1-mini", max_output_tokens: int = 800, temperature: float = 0.1, max_messages: int = 120, ) -> None: if client is None: if OpenAI is None: log.warning("OpenAI client not available. Install openai>=1.0.0.") else: try: # 尝试从配置获取 API Key 和 Base URL api_keys = OPENAI_API_KEYS.value if hasattr(OPENAI_API_KEYS, "value") else [] base_urls = OPENAI_API_BASE_URLS.value if hasattr(OPENAI_API_BASE_URLS, "value") else [] api_key = api_keys[0] if api_keys else os.environ.get("OPENAI_API_KEY") base_url = base_urls[0] if base_urls else os.environ.get("OPENAI_API_BASE_URL") if api_key: kwargs = {"api_key": api_key} if base_url: kwargs["base_url"] = base_url client = OpenAI(**kwargs) else: log.warning("No OpenAI API key found.") except Exception as e: log.warning(f"Failed to init OpenAI client: {e}") self._client = client self._model = model self._max_output_tokens = max_output_tokens self._temperature = temperature self._max_messages = max_messages def summarize( self, messages: Sequence[Dict], *, existing_summary: Optional[str] = None, max_tokens: Optional[int] = None, ) -> Optional[HistorySummary]: if not messages and not existing_summary: return None # 转换 dict 消息为 ChatMessage 格式用于 prompt 生成 # 确保消息按时间戳排序，防止乱序导致切片错误 sorted_messages = sorted(messages, key=lambda m: m.get('timestamp', 0) if isinstance(m.get('timestamp'), (int, float)) else 0) # 如果有 existing_summary，我们可以适当减少这里的消息量，或者依然取最近的 # 但为了逻辑简单，我们还是取最近的 max_messages trail = sorted_messages[-self._max_messages :] transcript = "\n".join(f"{m.get('role', 'user')}: {m.get('content', '')}" for m in trail) prompt = SUMMARY_PROMPT.format( existing_summary=existing_summary.strip() if existing_summary else "无", chat_transcript=transcript, ) if not self._client: log.error("No OpenAI client available for summarization.") return None log.info(f"Starting summary generation for {len(messages)} messages...") # Try primary client first try: # 增加 max_tokens 限制，避免摘要过长被截断，同时留给 JSON 结构足够的空间 # 根据经验，1000 字摘要 + JSON 结构大约需要 1500 tokens safe_max_tokens = max(max_tokens or self._max_output_tokens, 2000) response = self._client.chat.completions.create( model=self._model, messages=[{"role": "user", "content": prompt}], max_tokens=safe_max_tokens, temperature=self._temperature, ) # Debug: Print full response to investigate empty content issues log.info(f"Full Summary API Response: {response}") payload = response.choices[0].message.content or "" finish_reason = response.choices[0].finish_reason if finish_reason == "length": log.warning("Summary generation was truncated due to length limit!") log.info(f"Summary generation completed. Payload length: {len(payload)}") log.info(f"Summary Content:\n{payload}") return self._parse_response(payload) except Exception as e: log.warning(f"Summarization failed: {e}") return None def merge_summaries( self, summary_a: str, summary_b: str, *, max_tokens: Optional[int] = None, ) -> Optional[HistorySummary]: if not summary_a and not summary_b: return None prompt = MERGE_ONLY_PROMPT.format( summary_a=summary_a or "无", summary_b=summary_b or "无", ) if not self._client: return None log.info(f"Starting summary merge (A len={len(summary_a)}, B len={len(summary_b)})...") # Try primary client try: response = self._client.chat.completions.create( model=self._model, messages=[{"role": "user", "content": prompt}], max_tokens=max_tokens or self._max_output_tokens, temperature=self._temperature, ) payload = response.choices[0].message.content or "" log.info("Summary merge completed successfully.") return self._parse_response(payload) except Exception as e: log.warning(f"Merge failed: {e}") return None def _parse_response(self, payload: str) -> HistorySummary: data = _safe_json_loads(payload) # 如果解析出的 data 是空或者不是 dict，尝试直接用 payload if not isinstance(data, dict) or (not data and not payload.strip().startswith("{")): summary = payload.strip() table = {} else: summary = str(data.get("summary", "")).strip() table_payload = data.get("table", {}) or {} table = { "who": str(table_payload.get("who", "")).strip(), "how": str(table_payload.get("how", "")).strip(), "why": str(table_payload.get("why", "")).strip(), "what": str(table_payload.get("what", "")).strip(), } if not summary: summary = payload.strip() if len(summary) > 1000: summary = summary[:1000].rstrip() + "..." return HistorySummary(summary=summary, table=table) def _safe_json_loads(raw: str) -> Dict[str, Any]: try: return json.loads(raw) except json.JSONDecodeError: # 简单的正则提取尝试 match = re.search(r'(\{.*\})', raw, re.DOTALL) if match: try: return json.loads(match.group(1)) except json.JSONDecodeError: pass return {} # --- Core Logic Modules --- def build_ordered_messages( messages_map: Optional[Dict], anchor_id: Optional[str] = None ) -> List[Dict]: """ 将消息 map 还原为有序列表 策略： 1. 优先：基于 parentId 链条追溯（从 anchor_id 向上回溯到根消息） 2. 退化：按时间戳排序（无 anchor_id 或追溯失败时） 参数： messages_map: 消息 map，格式 {"msg-id": {"role": "user", "content": "...", "parentId": "...", "timestamp": 123456}} anchor_id: 锚点消息 ID（链尾），从此消息向上追溯 返回： 有序的消息列表，每个消息包含 id 字段 """ if not messages_map: return [] # 补齐消息的 id 字段 def with_id(message_id: str, message: Dict) -> Dict: return {**message, **({"id": message_id} if "id" not in message else {})} # 模式 1：基于 parentId 链条追溯 if anchor_id and anchor_id in messages_map: ordered: List[Dict] = [] current_id: Optional[str] = anchor_id while current_id: current_msg = messages_map.get(current_id) if not current_msg: break ordered.insert(0, with_id(current_id, current_msg)) current_id = current_msg.get("parentId") return ordered # 模式 2：基于时间戳排序 sortable: List[Tuple[int, str, Dict]] = [] for mid, message in messages_map.items(): ts = ( message.get("createdAt") or message.get("created_at") or message.get("timestamp") or 0 ) sortable.append((int(ts), mid, message)) sortable.sort(key=lambda x: x[0]) return [with_id(mid, msg) for _, mid, msg in sortable] def get_recent_messages_by_user_id(user_id: str, chat_id: str, num: int) -> List[Dict]: """ 获取指定用户的最近 N 条消息（优先当前会话，然后按时间顺序） 参数： user_id: 用户 ID chat_id: 当前会话 ID（用于优先提取） num: 需要获取的消息数量（<= 0 时返回全部） 返回： 有序的消息列表（优先当前会话，不足时由全局最近补齐） """ current_chat_messages: List[Dict] = [] other_messages: List[Dict] = [] # 遍历用户的所有聊天 chats = Chats.get_chat_list_by_user_id(user_id, include_archived=True) for chat in chats: messages_map = chat.chat.get("history", {}).get("messages", {}) or {} # 简单判断是否为当前会话 is_current_chat = (str(chat.id) == str(chat_id)) for mid, msg in messages_map.items(): # 跳过空内容 if msg.get("content", "") == "": continue ts = ( msg.get("createdAt") or msg.get("created_at") or msg.get("timestamp") or 0 ) entry = {**msg, "id": mid} entry.setdefault("chat_id", chat.id) entry.setdefault("timestamp", int(ts)) if is_current_chat: current_chat_messages.append(entry) else: other_messages.append(entry) # 分别排序 current_chat_messages.sort(key=lambda m: m.get("timestamp", 0)) other_messages.sort(key=lambda m: m.get("timestamp", 0)) if num <= 0: combined = current_chat_messages + other_messages combined.sort(key=lambda m: m.get("timestamp", 0)) return combined # 策略：优先保留当前会话消息 if len(current_chat_messages) >= num: return current_chat_messages[-num:] # 补充不足的部分 needed = num - len(current_chat_messages) supplement = other_messages[-needed:] if other_messages else [] # 合并并最终按时间排序 final_list = supplement + current_chat_messages final_list.sort(key=lambda m: m.get("timestamp", 0)) return final_list def slice_messages_with_summary( messages_map: Dict, boundary_message_id: Optional[str], anchor_id: Optional[str], pre_boundary: int = 20, ) -> List[Dict]: """ 基于摘要边界裁剪消息列表（返回摘要前 N 条 + 摘要后全部消息） 策略：保留摘要边界前 N 条消息（提供上下文）+ 摘要后全部消息（最新对话） 目的：降低 token 消耗，同时保留足够的上下文信息 参数： messages_map: 消息 map boundary_message_id: 摘要边界消息 ID（None 时返回全量消息） anchor_id: 锚点消息 ID（链尾） pre_boundary: 摘要边界前保留的消息数量（默认 20） 返回： 裁剪后的有序消息列表 示例： 100 条消息，摘要边界在第 50 条，pre_boundary=20 → 返回消息 29-99（共 71 条） """ ordered = build_ordered_messages(messages_map, anchor_id) if boundary_message_id: try: # 查找摘要边界消息的索引 boundary_idx = next( idx for idx, msg in enumerate(ordered) if msg.get("id") == boundary_message_id ) # 计算裁剪起点 start_idx = max(boundary_idx - pre_boundary, 0) ordered = ordered[start_idx:] except StopIteration: # 边界消息不存在，返回全量 pass return ordered def summarize(messages: List[Dict], old_summary: Optional[str] = None, model: Optional[str] = None) -> str: """ 生成对话摘要 参数： messages: 需要摘要的消息列表 old_summary: 旧摘要 model: 指定使用的模型 ID（如果为 None，则使用类内部默认值） 返回： 摘要字符串 """ summarizer = HistorySummarizer(model=model) if model else HistorySummarizer() result = summarizer.summarize(messages, existing_summary=old_summary) return result.summary if result else "" def compute_token_count(messages: List[Dict]) -> int: """ 计算消息的 token 数量（占位实现） 当前算法：4 字符 ≈ 1 token（粗略估算） TODO：接入真实 tokenizer（如 tiktoken for OpenAI models） """ total_chars = 0 for msg in messages: content = msg.get('content') if isinstance(content, str): total_chars += len(content) elif isinstance(content, list): for item in content: if isinstance(item, dict) and 'text' in item: total_chars += len(item['text']) return max(total_chars // 4, 0)