RAG 优化 | 文档版本管理

"""任务：实现文档版本控制场景：文档会更新，需要追踪历史"""from dataclasses import dataclassfrom datetime import datetimefrom typing import List, Optional, Dict@dataclassclass DocumentVersion:    """文档版本"""    version: int    content: str    created_at: datetime    created_by: str    changelog: strclass VersionedDocument:    """带版本的文档"""    def __init__(self, doc_id: str):        self.doc_id = doc_id        self.versions: List[DocumentVersion] = []   # 所有版本列表        self.current_version = 0                    # 当前是第几版    def add_version(        self,        content: str,        created_by: str,        changelog: str    ):        """添加新版本"""        version = DocumentVersion(            version=self.current_version + 1,            content=content,            created_at=datetime.now(),            created_by=created_by,            changelog=changelog        )        self.versions.append(version)        self.current_version += 1    def get_version(self, version: int) -> Optional[DocumentVersion]:        for v in self.versions:            if v.version == version:                return v        return None    def get_current(self) -> DocumentVersion:        """获取当前版本"""        return self.versions[-1] if self.versions else None    def rollback(self, version: int):        """回滚到指定版本，回滚不是真的删掉历史，而是创建一个新版本，内容是旧版本的内容。"""        # 1. 找到目标版本        target = self.get_version(version)        if not target:            raise ValueError(f"版本 {version} 不存在")        # 2. 添加新版本（内容是目标版本的内容）        self.add_version(            content=target.content,         # 用旧版本的内容            created_by="system",            # 系统自动操作            changelog=f"回滚到版本 {version}"  # 说明是回滚        )class VersionControlledRAG:    """带版本控制的RAG"""    def __init__(self, vector_store, embedding_service):        self.vector_store = vector_store        self.embedding_service = embedding_service        self.documents: Dict[str, VersionedDocument] = {}    # doc_id → 版本文档    def update_document(        self,        doc_id: str,        new_content: str,        user: str,        changelog: str    ):        """        更新文档        1. 创建新版本        2. 重新生成embedding        3. 更新向量数据库        """        # 1. 如果文档第一次创建，先初始化        if doc_id not in self.documents:            self.documents[doc_id] = VersionedDocument(doc_id)        # 2. 添加新版本        versioned_doc = self.documents[doc_id]        versioned_doc.add_version(new_content, user, changelog)        # 重新生成embedding并更新向量库        embedding = self.embedding_service.embed([new_content])[0]        self.vector_store.delete([doc_id])    # 删除旧的        self.vector_store.add([doc_id], [embedding], [new_content])    # 添加新的# 模拟服务 class MockEmbedding:    def embed(self, texts):        return [[0.1]*5 for _ in texts]class MockVectorStore:    def __init__(self):        self.docs = {}    def delete(self, ids):        for i in ids:            self.docs.pop(i, None)    def add(self, ids, embeddings, texts):        for i, e, t in zip(ids, embeddings, texts):            self.docs[i] = (e, t)# 测试 rag = VersionControlledRAG(MockVectorStore(), MockEmbedding())# 第一次创建rag.update_document("doc1", "RAG 是检索增强生成", "张三", "初版")# 修改rag.update_document("doc1", "RAG = 检索增强生成，用于AI问答", "李四", "优化描述")# 回滚doc = rag.documents["doc1"]doc.rollback(1)print("\n当前版本：", doc.get_current())

# 1. 初始化rag = VersionControlledRAG(vector_store, embedding_service)# 2. 第一次创建文档rag.update_document(    doc_id="doc_001",    new_content="RAG是检索增强生成",    user="张三",    changelog="创建文档")# 版本历史: [v1: "RAG是检索增强生成"]# 3. 修改文档rag.update_document(    doc_id="doc_001",    new_content="RAG = 检索 + 生成",    user="张三",    changelog="简化描述")# 版本历史: [v1: "RAG是检索增强生成", v2: "RAG = 检索 + 生成"]# 4. 再次修改rag.update_document(    doc_id="doc_001",    new_content="RAG是检索增强生成技术，结合了检索和LLM",    user="李四",    changelog="补充详细说明")# 版本历史: [v1, v2, v3]# 5. 回滚到 v1versioned_doc = rag.documents["doc_001"]versioned_doc.rollback(1)# 版本历史: [v1, v2, v3, v4]  (v4内容和v1相同)

# 不用 dataclass 的写法class DocumentVersion:    def __init__(self, version, content, created_at, created_by, changelog):        self.version = version        self.content = content        self.created_at = created_at        self.created_by = created_by        self.changelog = changelog# 用 dataclass 的写法（自动生成 __init__）@dataclassclass DocumentVersion:    version: int    content: str    created_at: datetime    created_by: str    changelog: str# 同样能用，代码更简洁

# 不正确的回滚（删除后面的版本）v1, v2, v3 → 删除v2,v3 → 只有v1 ❌# 正确的回滚（追加新版本）v1, v2, v3 → 追加v4(内容是v1的) → v1,v2,v3,v4 ✅

# 文档内容变了，向量也要重新生成self.vector_store.delete([doc_id])      # 删旧的self.vector_store.add([doc_id], [embedding], [new_content])  # 加新的

@dataclassclass DocumentVersion:    version: int    content: str    created_at: datetime    created_by: str    changelog: str

class VersionedDocument:    def __init__(self, doc_id):        self.doc_id = doc_id        self.versions = []        self.current_version = 0    def add_version(self, content, created_by, changelog):        self.current_version +=1        self.versions.append(DocumentVersion(            self.current_version, content, datetime.now(), created_by, changelog        ))    def rollback(self, version):        target = self.get_version(version)        self.add_version(target.content, "system", f"回滚到 {version}")