RAG：检索增强生成

"""rag_travel.py — 基础 RAG：加载 → 切块 → 内存向量库 → 检索 → 生成"""
from __future__ import annotations

import math
from dataclasses import dataclass, field


# ── 极简运行时 ──────────────────────────────────────────

@dataclass(frozen=True)
class Message:
    role: str
    content: str

class MockLLM:
    def __init__(self, responses: list[str]):
        self._responses = list(responses)
        self._idx = 0
    def complete(self, messages: list[Message]) -> str:
        if self._idx >= len(self._responses):
            raise RuntimeError("MockLLM 回复用完了")
        out = self._responses[self._idx]
        self._idx += 1
        return out


# ── 文档和切块 ──────────────────────────────────────────

@dataclass(frozen=True)
class Document:
    doc_id: str
    text: str

@dataclass(frozen=True)
class Chunk:
    doc_id: str
    chunk_id: str
    text: str


def load_documents() -> list[Document]:
    """模拟加载旅游知识库。"""
    return [
        Document(
            doc_id="tea_museum",
            text=(
                "中国茶叶博物馆位于杭州龙井路 88 号，"
                "开放时间 9:00-17:00（16:30 停止入馆），"
                "每周一闭馆（法定节假日除外），"
                "免费参观，凭身份证入馆。"
                "馆内设有茶史厅、茶萃厅、茶事厅、茶缘厅。"
            ),
        ),
        Document(
            doc_id="west_lake",
            text=(
                "西湖位于杭州市中心，全天开放，免费。"
                "主要景点包括断桥残雪、苏堤春晓、三潭印月。"
                "环湖一圈约 10 公里，步行约 3 小时。"
                "游船票价 55 元/人（三潭印月上岛）。"
            ),
        ),
        Document(
            doc_id="hefang_street",
            text=(
                "河坊街是杭州的历史文化街区，位于吴山脚下。"
                "以小吃和手工艺品闻名，全天开放。"
                "推荐：定胜糕、龙须糖、吴山烤禽。"
                "从西湖步行约 15 分钟可达。"
            ),
        ),
    ]


def chunk_documents(docs: list[Document], max_chars: int = 120) -> list[Chunk]:
    """简单切块：按句号分割，合并不超过 max_chars 的连续句子。"""
    chunks: list[Chunk] = []
    for doc in docs:
        sentences = [s.strip() for s in doc.text.split("。") if s.strip()]
        buf = ""
        chunk_idx = 0
        for sent in sentences:
            candidate = buf + sent + "。" if buf else sent + "。"
            if len(candidate) > max_chars and buf:
                chunks.append(Chunk(
                    doc_id=doc.doc_id,
                    chunk_id=f"{doc.doc_id}_{chunk_idx}",
                    text=buf,
                ))
                chunk_idx += 1
                buf = sent + "。"
            else:
                buf = candidate
        if buf:
            chunks.append(Chunk(
                doc_id=doc.doc_id,
                chunk_id=f"{doc.doc_id}_{chunk_idx}",
                text=buf,
            ))
    return chunks


# ── 内存向量库（用词频代替嵌入，演示用） ──────────────

def _tokenize(text: str) -> list[str]:
    """极简分词：按字符切，小写，只保留字母数字和中文。"""
    tokens: list[str] = []
    buf = ""
    for ch in text:
        if ch.isalnum() or '\u4e00' <= ch <= '\u9fff':
            buf += ch.lower()
        else:
            if buf:
                tokens.append(buf)
                buf = ""
    if buf:
        tokens.append(buf)
    return tokens


@dataclass
class VectorStore:
    """用 TF 评分的内存向量库。生产环境替换为真正的嵌入 + 向量数据库。"""
    chunks: list[Chunk] = field(default_factory=list)

    def add(self, chunks: list[Chunk]) -> None:
        self.chunks.extend(chunks)

    def search(self, query: str, k: int = 3) -> list[tuple[Chunk, float]]:
        query_tokens = _tokenize(query)
        if not query_tokens:
            return []
        scored: list[tuple[Chunk, float]] = []
        for chunk in self.chunks:
            chunk_tokens = _tokenize(chunk.text)
            score = sum(1 for qt in query_tokens for ct in chunk_tokens if qt == ct)
            if score > 0:
                scored.append((chunk, float(score)))
        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:k]


# ── RAG 主流程 ──────────────────────────────────────────

def rag_answer(
    model: MockLLM,
    question: str,
    store: VectorStore,
    k: int = 3,
) -> str:
    """
    基础 RAG：
    1. 用 question 检索 top-k 片段
    2. 把片段塞进 prompt
    3. 模型基于片段回答
    """
    # 检索
    results = store.search(question, k=k)
    print(f"[检索] query='{question}', 命中 {len(results)} 条：")
    context_parts: list[str] = []
    for chunk, score in results:
        print(f"  [{chunk.chunk_id}] score={score:.0f} | {chunk.text[:60]}...")
        context_parts.append(f"[{chunk.chunk_id}] {chunk.text}")

    context = "\n\n".join(context_parts)

    # 生成
    answer = model.complete([
        Message(
            role="system",
            content="根据提供的参考资料回答问题。引用来源用 [chunk_id]。如果资料不足以回答，说明不确定。",
        ),
        Message(
            role="user",
            content=f"问题：{question}\n\n参考资料：\n{context}",
        ),
    ])
    return answer


# ── 运行 ────────────────────────────────────────────────

def main() -> None:
    # 1. 加载文档
    docs = load_documents()
    print(f"[加载] {len(docs)} 份文档")

    # 2. 切块
    chunks = chunk_documents(docs, max_chars=120)
    print(f"[切块] {len(chunks)} 个片段")

    # 3. 存入向量库
    store = VectorStore()
    store.add(chunks)

    # 4. 检索 + 生成
    model = MockLLM([
        (
            "中国茶叶博物馆免费参观，凭身份证入馆，开放时间 9:00-17:00，"
            "每周一闭馆。[tea_museum_0]"
        ),
    ])

    question = "茶叶博物馆要门票吗？几点开门？"
    answer = rag_answer(model, question, store, k=3)
    print(f"\n回答：\n{answer}")


if __name__ == "__main__":
    main()

[加载] 3 份文档
[切块] 5 个片段
[检索] query='茶叶博物馆要门票吗？几点开门？', 命中 3 条：
  [tea_museum_0] score=3 | 中国茶叶博物馆位于杭州龙井路 88 号，开放时间 9:00-17:00（16:30 停止入馆），每周一闭馆...
  [tea_museum_1] score=1 | 馆内设有茶史厅、茶萃厅、茶事厅、茶缘厅。...
  [hefang_street_0] score=1 | 河坊街是杭州的历史文化街区，位于吴山脚下。以小吃和手工艺品闻名，全天开放。...

回答：
中国茶叶博物馆免费参观，凭身份证入馆，开放时间 9:00-17:00，每周一闭馆。[tea_museum_0]

索引构建 = n_docs × 切块 + n_chunks × 嵌入
每次查询 = 1 × 嵌入（query） + top-k 检索 + 1 × 生成