summarizer.py

"""
模块 4: 大模型总结 (Summarizer)

将搜索结果 + 网页内容交给大模型进行智能总结。
支持多种总结模式: 简明摘要 / 结构化 / 引用来源。
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Optional

import httpx

from config import ModuleConfig
from crawler import SearchResponse


@dataclass
class SummaryResult:
    """总结结果"""
    answer: str                    # 主要回答
    key_points: list[str]          # 关键要点
    sources: list[dict]            # 引用的来源 [{title, url, snippet}]
    confidence: str = "medium"     # high / medium / low
    raw_model_output: str = ""     # 原始输出 (调试用)


SUMMARIZE_SYSTEM_PROMPT = """你是一个信息总结助手。你的任务是根据提供的搜索结果，回答用户的问题。

要求:
1. 基于提供的搜索结果作答，不要编造信息
2. 如果搜索结果信息不足或矛盾，诚实说明
3. 回答要结构清晰，先给核心答案，再展开细节
4. 引用具体来源时标注编号 [1] [2] 等
5. 对时效性信息，注明时间

最后请严格按 JSON 格式输出:
{
  "answer": "核心回答 (markdown 格式)",
  "key_points": ["要点1", "要点2", "要点3"],
  "confidence": "high/medium/low",
  "source_refs": [1, 2, 3]
}"""


def _build_summarize_prompt(
    query: str,
    search_response: SearchResponse,
) -> str:
    """构建总结 prompt"""
    context = search_response.context_text
    
    # 如果内容太长, 截断
    max_context_chars = 8000
    if len(context) > max_context_chars:
        context = context[:max_context_chars] + "\n\n...[更多内容已截断]"
    
    prompt = f"""搜索结果如下:

{context}

---
用户问题: {query}

请基于以上搜索结果作答 (JSON 格式):"""
    
    return prompt


async def summarize(
    query: str,
    search_response: SearchResponse,
    config: ModuleConfig,
) -> SummaryResult:
    """
    对搜索结果进行 LLM 总结。
    
    Args:
        query: 用户原始问题
        search_response: 搜索 + 抓取结果
        config: 模块配置
    """
    if search_response.error:
        return SummaryResult(
            answer=f"搜索过程出错: {search_response.error}",
            key_points=[],
            sources=[],
            confidence="low",
        )
    
    if not search_response.results:
        return SummaryResult(
            answer="未找到相关搜索结果，请尝试更换搜索词。",
            key_points=[],
            sources=[],
            confidence="low",
        )
    
    prompt = _build_summarize_prompt(query, search_response)
    
    if config.verbose:
        prompt_len = len(prompt)
        print(f"[Summarizer] Prompt 长度: {prompt_len} chars, "
              f"{len(search_response.results)} 条搜索结果")
    
    async with httpx.AsyncClient(timeout=60) as client:
        resp = await client.post(
            f"{config.llm.api_base}/chat/completions",
            headers={
                "Authorization": f"Bearer {config.llm.api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": config.llm.summary_model,
                "messages": [
                    {"role": "system", "content": SUMMARIZE_SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                "max_tokens": config.llm.max_tokens,
                "temperature": config.llm.temperature,
            },
        )
        resp.raise_for_status()
        data = resp.json()
    
    raw = data["choices"][0]["message"]["content"].strip()
    
    # 提取 JSON
    json_match = re.search(r"\{.*\}", raw, re.DOTALL)
    json_str = json_match.group(0) if json_match else raw
    
    try:
        result = json.loads(json_str)
    except json.JSONDecodeError:
        # fallback: 直接用原始输出
        return SummaryResult(
            answer=raw,
            key_points=[],
            sources=_extract_sources(search_response),
            confidence="medium",
            raw_model_output=raw,
        )
    
    # 构建来源引用
    source_refs = result.get("source_refs", [])
    sources = []
    for ref in source_refs:
        idx = ref - 1
        if 0 <= idx < len(search_response.results):
            r = search_response.results[idx]
            sources.append({
                "title": r.title,
                "url": r.url,
                "snippet": r.snippet[:200],
            })
    
    return SummaryResult(
        answer=result.get("answer", raw),
        key_points=result.get("key_points", []),
        sources=sources,
        confidence=result.get("confidence", "medium"),
        raw_model_output=raw,
    )


def _extract_sources(response: SearchResponse) -> list[dict]:
    """从搜索结果提取来源信息"""
    return [
        {"title": r.title, "url": r.url, "snippet": r.snippet[:200]}
        for r in response.results[:5]
    ]


# ── 格式化输出 ──────────────────────────────────────────────────

def format_summary(result: SummaryResult, query: str = "") -> str:
    """将总结结果格式化为美观的文本输出"""
    
    lines = []
    
    if query:
        lines.append(f"## 🔍 {query}")
    
    lines.append("")
    lines.append(result.answer)
    
    if result.key_points:
        lines.append("")
        lines.append("### 📌 关键要点")
        for i, point in enumerate(result.key_points, 1):
            lines.append(f"{i}. {point}")
    
    if result.sources:
        lines.append("")
        lines.append("### 📚 参考来源")
        for i, src in enumerate(result.sources, 1):
            lines.append(f"{i}. [{src['title']}]({src['url']})")
            lines.append(f"   _{src['snippet'][:150]}_")
    
    if result.confidence:
        emoji_map = {"high": "🟢", "medium": "🟡", "low": "🔴"}
        emoji = emoji_map.get(result.confidence, "⚪")
        lines.append(f"\n> 置信度: {emoji} {result.confidence}")
    
    return "\n".join(lines)