RAG 系统测试用例¶

import pytest
import chromadb
from chromadb.config import Settings

class RAGRetriever:
    """RAG检索器"""

    def __init__(self, collection_name: str = "test_collection"):
        """初始化检索器"""
        # 创建 ChromaDB 客户端实例（内存模式）
        self.client = chromadb.Client(Settings())
        # 获取或创建向量集合，用于存储文档及其向量表示
        self.collection = self.client.get_or_create_collection(
            name=collection_name
        )

    def add_documents(
        self,
        documents: list[str],
        metadatas: list[dict] = None,
    ):
        """添加文档"""
        # 若未提供元数据，则为每个文档创建空字典
        if metadatas is None:
            metadatas = [{} for _ in documents]

        # 将文档、元数据和自动生成的 ID 一起存入向量数据库
        self.collection.add(
            documents=documents,
            metadatas=metadatas,
            ids=[f"doc_{i}" for i in range(len(documents))]  # 自动编号
        )

    def retrieve(
        self,
        query: str,
        n_results: int = 5,
    ) -> list[dict]:
        """
        检索文档

        Args:
            query: 查询文本
            n_results: 返回结果数量

        Returns:
            检索结果列表
        """
        # 调用 ChromaDB 的向量相似度检索
        results = self.collection.query(
            query_texts=[query],      # 查询文本（自动向量化）
            n_results=n_results,      # 返回 Top-N 结果
        )

        # 将检索结果重新组织为结构化字典列表
        retrieved_docs = []
        for i in range(len(results['documents'][0])):
            retrieved_docs.append({
                'content': results['documents'][0][i],    # 文档内容
                'metadata': results['metadatas'][0][i],   # 元数据
                'distance': results['distances'][0][i],   # 与查询的向量距离（越小越相关）
                'id': results['ids'][0][i],               # 文档 ID
            })

        return retrieved_docs

def test_basic_retrieval():
    """测试基础检索功能"""
    retriever = RAGRetriever()

    # 添加测试文档
    documents = [
        "Python是一种高级编程语言，由Guido van Rossum于1991年创建。",
        "JavaScript是一种脚本语言，主要用于网页开发。",
        "Java是一种面向对象的编程语言，由Sun Microsystems开发。",
        "C++是一种通用的编程语言，支持面向对象和泛型编程。",
        "Go是Google开发的开源编程语言，强调简洁和高效。",
    ]

    retriever.add_documents(documents)

    # 测试检索：查询与 Python 相关的文档
    query = "Python是什么？"
    results = retriever.retrieve(query, n_results=3)

    # 验证检索结果数量和相关性
    assert len(results) == 3                      # 应返回 3 个结果
    assert "Python" in results[0]['content']       # 最相关文档应包含 Python
    print(f"✓ 检索成功，找到 {len(results)} 个相关文档")

    # 打印检索结果及其向量距离（距离越小代表语义越相似）
    for i, doc in enumerate(results, 1):  # enumerate同时获取索引和元素
        print(f"  {i}. {doc['content'][:50]}... (距离: {doc['distance']:.4f})")

def test_semantic_retrieval_accuracy():
    """测试语义检索准确性"""
    retriever = RAGRetriever()

    # 添加领域相关文档
    documents = [
        "机器学习是人工智能的一个分支，它使计算机能够从数据中学习。",
        "深度学习是机器学习的一种方法，使用神经网络进行学习。",
        "自然语言处理(NLP)是AI的一个领域，专注于计算机与人类语言之间的交互。",
        "计算机视觉是AI的一个领域，使计算机能够理解和解释视觉信息。",
        "强化学习是机器学习的一种方法，通过与环境交互来学习最优策略。",
    ]

    retriever.add_documents(documents)

    # 测试用例：查询文本与期望命中的关键词对
    test_cases = [
        ("什么是深度学习？", "深度学习"),
        ("AI如何理解语言？", "自然语言处理"),
        ("机器如何看懂图片？", "计算机视觉"),
        ("智能体如何学习？", "强化学习"),
    ]

    correct = 0
    for query, keyword in test_cases:
        results = retriever.retrieve(query, n_results=3)

        # 检查 Top-3 结果中是否包含目标关键词（语义匹配）
        found = any(keyword in doc['content'] for doc in results)  # any()任一为True则返回True

        if found:
            correct += 1
            print(f"✓ 正确: {query}")
        else:
            print(f"✗ 错误: {query}")
            print(f"  关键词 '{keyword}' 未在前3个结果中找到")

    accuracy = correct / len(test_cases)
    print(f"\n检索准确率: {accuracy:.2%}")
    assert accuracy >= 0.75, "检索准确率过低"  # assert断言：条件False时抛出AssertionError

def test_metadata_filtering():
    """测试元数据过滤"""
    retriever = RAGRetriever()

    # 添加带元数据的文档
    documents = [
        "Python是一种高级编程语言",
        "JavaScript是一种脚本语言",
        "Java是一种面向对象的语言",
        "C++支持面向对象和泛型编程",
        "Go强调简洁和高效",
    ]

    metadatas = [
        {"category": "programming", "year": 1991},
        {"category": "web", "year": 1995},
        {"category": "programming", "year": 1995},
        {"category": "programming", "year": 1985},
        {"category": "programming", "year": 2009},
    ]

    retriever.add_documents(documents, metadatas)

    # 测试元数据过滤
    query = "编程语言"

    # 无过滤检索：获取所有文档
    results_all = retriever.retrieve(query, n_results=10)
    print(f"无过滤结果: {len(results_all)} 个文档")

    # 在客户端侧按 category 元数据进行过滤
    results_filtered = [
        doc for doc in results_all
        if doc['metadata'].get('category') == 'programming'
    ]
    print(f"过滤后结果: {len(results_filtered)} 个文档")

    # 验证过滤结果：应得到 4 个 programming 类别的文档
    assert len(results_filtered) == 4  # Python, Java, C++, Go
    assert all(
        doc['metadata'].get('category') == 'programming'
        for doc in results_filtered
    )

    print("✓ 元数据过滤功能正常")

class DocumentSplitter:
    """文档切分器"""

    def __init__(
        self,
        chunk_size: int = 500,
        chunk_overlap: int = 50,
    ):
        """
        初始化切分器

        Args:
            chunk_size: 块大小
            chunk_overlap: 块重叠
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def split(self, text: str) -> list[str]:
        """
        切分文档

        Args:
            text: 文本

        Returns:
            切分后的文本块列表
        """
        chunks = []
        start = 0

        # 滑动窗口切分：每次截取 chunk_size 字符，窗口之间保留 overlap 重叠
        while start < len(text):
            end = start + self.chunk_size
            chunk = text[start:end]    # 截取当前窗口的文本块
            chunks.append(chunk)

            # 下一个窗口起点 = 当前结束位 - 重叠量，保证相邻块有交集
            start = end - self.chunk_overlap

        return chunks

def test_fixed_size_splitting():
    """测试固定大小切分"""
    splitter = DocumentSplitter(
        chunk_size=100,
        chunk_overlap=20,
    )

    # 创建测试文本
    text = "这是一个测试文本。" * 50

    # 切分文档
    chunks = splitter.split(text)

    print(f"原文长度: {len(text)}")
    print(f"切分后块数: {len(chunks)}")

    # 验证切分结果：确保产生了多个块
    assert len(chunks) > 1, "应该切分成多个块"

    # 验证每个块不超过指定大小（最后一个块可能较小，跳过）
    for i, chunk in enumerate(chunks[:-1]):  # 最后一个块可能较小
        assert len(chunk) <= 100, f"块 {i} 大小超过限制"

    # 验证相邻块之间存在重叠（确保上下文连贯性）
    overlap_size = len(chunks[0][-20:])
    assert overlap_size >= 20, "重叠大小不足"

    print("✓ 固定大小切分功能正常")

import re

class SemanticSplitter:
    """语义切分器"""

    def __init__(self, max_chunk_size: int = 500):
        """
        初始化语义切分器

        Args:
            max_chunk_size: 最大块大小
        """
        self.max_chunk_size = max_chunk_size

    def split(self, text: str) -> list[str]:
        """
        基于语义切分文档

        Args:
            text: 文本

        Returns:
            切分后的文本块列表
        """
        # 先按双换行分割成段落（保持语义完整性）
        paragraphs = text.split('\n\n')

        chunks = []
        current_chunk = ""

        # 贪心合并：尽量将相邻段落合并到同一块，直到达到大小上限
        for paragraph in paragraphs:
            # 如果当前块加上新段落超过最大大小，则保存当前块并开始新块
            if len(current_chunk) + len(paragraph) > self.max_chunk_size:
                if current_chunk:
                    chunks.append(current_chunk.strip())  # 链式调用：strip去除空白
                current_chunk = paragraph
            else:
                if current_chunk:
                    current_chunk += "\n\n" + paragraph
                else:
                    current_chunk = paragraph

        # 添加最后一个块
        if current_chunk:
            chunks.append(current_chunk.strip())

        return chunks

def test_semantic_splitting():
    """测试语义切分"""
    splitter = SemanticSplitter(max_chunk_size=200)

    # 创建测试文本（多个段落）
    text = """
    这是第一段。它包含一些关于Python的信息。Python是一种高级编程语言。

    这是第二段。它讨论JavaScript。JavaScript主要用于网页开发。

    这是第三段。它介绍Java。Java是一种面向对象的语言。

    这是第四段。它提到C++。C++支持多种编程范式。
    """

    # 切分文档
    chunks = splitter.split(text)

    print(f"切分后块数: {len(chunks)}")

    for i, chunk in enumerate(chunks, 1):
        print(f"\n块 {i}:")
        print(chunk[:100] + "...")

    # 验证切分结果
    assert len(chunks) > 1, "应该切分成多个块"

    # 验证每个块不超过最大大小
    for chunk in chunks:
        assert len(chunk) <= 200, f"块大小超过限制: {len(chunk)}"

    print("\n✓ 语义切分功能正常")

import openai

class RAGGenerator:
    """RAG生成器"""

    def __init__(self, api_key: str, model: str = "gpt-4o"):
        """初始化生成器"""
        self.client = openai.OpenAI(api_key=api_key)
        self.model = model

    def generate(
        self,
        query: str,
        retrieved_docs: list[dict],
        max_tokens: int = 500,
    ) -> str:
        """
        基于检索结果生成回答

        Args:
            query: 查询
            retrieved_docs: 检索到的文档
            max_tokens: 最大token数

        Returns:
            生成的回答
        """
        # 将检索到的文档拼接为上下文字符串，作为 LLM 的参考信息
        context = "\n\n".join([
            f"文档 {i+1}: {doc['content']}"
            for i, doc in enumerate(retrieved_docs)
        ])

        # 构建 RAG Prompt：将检索文档作为上下文，要求模型基于文档回答
        prompt = f"""基于以下文档回答问题。如果文档中没有相关信息，请说明。

文档:
{context}

问题: {query}

回答:"""

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "user", "content": prompt}
            ],
            temperature=0.7,
            max_tokens=max_tokens,
        )

        return response.choices[0].message.content

def test_rag_generation():
    """测试RAG生成"""
    retriever = RAGRetriever()
    generator = RAGGenerator(api_key="your-api-key")

    # 添加文档
    documents = [
        "Python由Guido van Rossum于1991年创建，是一种高级编程语言。",
        "Python支持多种编程范式，包括面向对象、函数式和过程式编程。",
        "Python有丰富的标准库和第三方库，广泛应用于Web开发、数据科学、AI等领域。",
    ]

    retriever.add_documents(documents)

    # 步骤一：检索相关文档
    query = "Python是什么时候创建的？"
    retrieved_docs = retriever.retrieve(query, n_results=3)

    print(f"检索到的文档: {len(retrieved_docs)} 个")
    for doc in retrieved_docs:
        print(f"  - {doc['content']}")

    # 步骤二：基于检索结果生成回答
    answer = generator.generate(query, retrieved_docs)

    print(f"\n生成的回答: {answer}")

    # 步骤三：验证回答质量（包含检索到的关键信息）
    assert len(answer) > 10, "回答过短"
    assert "1991" in answer or "Guido" in answer, "回答未包含关键信息"

def test_citation_generation():
    """测试引用生成"""
    retriever = RAGRetriever()
    generator = RAGGenerator(api_key="your-api-key")

    # 添加文档
    documents = [
        "Python由Guido van Rossum于1991年创建。",
        "JavaScript由Brendan Eich于1995年创建。",
        "Java由James Gosling于1995年创建。",
    ]

    retriever.add_documents(documents)

    # 检索
    query = "Python和JavaScript是什么时候创建的？"
    retrieved_docs = retriever.retrieve(query, n_results=3)

    # 构建要求引用来源的 Prompt，将检索文档编号嵌入上下文
    prompt = f"""基于以下文档回答问题，并在回答中注明引用的文档编号。

文档:
{chr(10).join([f'文档{i+1}: {doc["content"]}' for i, doc in enumerate(retrieved_docs)])}

问题: {query}

回答:"""

    # 调用 LLM 生成带引用的回答
    response = self.client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "user", "content": prompt}
        ],
        temperature=0.7,
        max_tokens=500,
    )

    answer = response.choices[0].message.content

    print(f"生成的回答: {answer}")

    # 验证引用
    assert "文档1" in answer or "文档2" in answer, "回答未包含引用"

class RAGSystem:
    """RAG系统"""

    def __init__(self, api_key: str):
        """初始化系统"""
        self.retriever = RAGRetriever()
        self.generator = RAGGenerator(api_key)

    def query(self, question: str) -> dict:
        """
        执行查询

        Args:
            question: 问题

        Returns:
            查询结果字典
        """
        # 步骤一：检索相关文档（取 Top-3）
        retrieved_docs = self.retriever.retrieve(question, n_results=3)

        # 步骤二：基于检索结果生成回答
        answer = self.generator.generate(question, retrieved_docs)

        # 返回完整的查询结果（包括问题、回答和检索到的文档）
        return {
            'question': question,
            'answer': answer,
            'retrieved_docs': retrieved_docs,
        }

def test_end_to_end_rag():
    """测试端到端RAG"""
    rag = RAGSystem(api_key="your-api-key")

    # 添加知识库
    documents = [
        "机器学习是AI的一个分支，使计算机能够从数据中学习。",
        "监督学习使用标记数据进行训练，如分类和回归。",
        "无监督学习从未标记数据中发现模式，如聚类。",
        "强化学习通过与环境交互学习最优策略。",
        "深度学习使用神经网络进行学习，是机器学习的一个子集。",
    ]

    rag.retriever.add_documents(documents)

    # 测试查询
    questions = [
        "什么是机器学习？",
        "监督学习和无监督学习有什么区别？",
        "深度学习和机器学习的关系是什么？",
    ]

    # 逐问题执行完整 RAG 流程（检索 + 生成）
    for question in questions:
        result = rag.query(question)

        print(f"\n问题: {result['question']}")
        print(f"回答: {result['answer']}")
        print(f"检索到的文档数: {len(result['retrieved_docs'])}")

        # 验证结果：回答不能太短，且必须检索到文档
        assert len(result['answer']) > 10, "回答过短"
        assert len(result['retrieved_docs']) > 0, "未检索到文档"

    print("\n✓ 端到端RAG流程正常")

def test_multi_turn_dialogue():
    """测试多轮对话"""
    rag = RAGSystem(api_key="your-api-key")

    # 添加知识库
    documents = [
        "Python由Guido van Rossum创建，于1991年首次发布。",
        "Python的设计哲学强调代码的可读性和简洁的语法。",
        "Python支持多种编程范式，包括面向对象、函数式和过程式编程。",
        "Python有丰富的标准库，常被称为"自带电池"的语言。",
    ]

    rag.retriever.add_documents(documents)

    # 多轮对话
    dialogue = [
        "Python是什么时候创建的？",
        "谁创建了Python？",
        "Python有什么特点？",
    ]

    # 模拟多轮对话，逐轮测试 RAG 系统的回答质量
    for i, question in enumerate(dialogue, 1):
        result = rag.query(question)

        print(f"\n第{i}轮")
        print(f"问题: {result['question']}")
        print(f"回答: {result['answer']}")

        # 每轮回答都不能太短，确保回答质量
        assert len(result['answer']) > 10, f"第{i}轮回答过短"

    print("\n✓ 多轮对话功能正常")

import time

def test_retrieval_speed():
    """测试检索速度"""
    retriever = RAGRetriever()

    # 添加大量文档测试检索性能
    num_docs = 1000
    documents = [f"这是第{i}个测试文档。" * 10 for i in range(num_docs)]

    print(f"添加 {num_docs} 个文档...")
    retriever.add_documents(documents)

    # 循环执行多次查询，统计平均耗时
    num_queries = 100
    query = "测试文档"

    print(f"执行 {num_queries} 次查询...")
    start_time = time.time()  # 记录起始时间

    for _ in range(num_queries):
        results = retriever.retrieve(query, n_results=10)

    end_time = time.time()    # 记录结束时间

    # 计算性能指标
    avg_time = (end_time - start_time) / num_queries  # 平均查询时间
    throughput = num_queries / (end_time - start_time) # 吞吐量（每秒查询数）

    print(f"平均查询时间: {avg_time * 1000:.2f} ms")
    print(f"吞吐量: {throughput:.2f} queries/s")

    # 验证性能
    assert avg_time < 0.1, "查询速度过慢"

def test_generation_speed():
    """测试生成速度"""
    generator = RAGGenerator(api_key="your-api-key")

    query = "什么是人工智能？"
    retrieved_docs = [
        {'content': '人工智能是计算机科学的一个分支。'},
        {'content': 'AI使机器能够模拟人类智能。'},
    ]

    # 测试 LLM 生成速度
    num_generations = 10

    print(f"执行 {num_generations} 次生成...")
    start_time = time.time()   # 计时开始

    for _ in range(num_generations):
        answer = generator.generate(query, retrieved_docs)

    end_time = time.time()     # 计时结束

    # 计算生成性能指标
    avg_time = (end_time - start_time) / num_generations    # 平均生成时间
    throughput = num_generations / (end_time - start_time)  # 生成吞吐量

    print(f"平均生成时间: {avg_time:.2f} s")
    print(f"吞吐量: {throughput:.2f} generations/s")

    # 验证性能（生成通常较慢）
    assert avg_time < 10, "生成速度过慢"

from rouge_score import rouge_scorer

def test_rouge_evaluation():
    """测试ROUGE评估"""
    generator = RAGGenerator(api_key="your-api-key")

    # 测试用例
    test_cases = [
        {
            'question': 'Python是什么？',
            'reference': 'Python是一种高级编程语言，由Guido van Rossum于1991年创建。',
            'context': [
                {'content': 'Python由Guido van Rossum于1991年创建，是一种高级编程语言。'},
            ],
        },
    ]

    # 初始化 ROUGE 评分器，支持 ROUGE-1/2/L 三种指标
    scorer = rouge_scorer.RougeScorer(
        ['rouge1', 'rouge2', 'rougeL'],  # 评估单词/双词/最长公共子序列重叠
        use_stemmer=True                  # 启用词干提取，提高容错性
    )

    # 分别存储三种 ROUGE 指标的得分
    rouge1_scores = []
    rouge2_scores = []
    rougeL_scores = []

    for test_case in test_cases:
        # 生成回答
        answer = generator.generate(
            test_case['question'],
            test_case['context'],
        )

        # 计算生成回答与参考答案的 ROUGE 得分
        scores = scorer.score(test_case['reference'], answer)

        # 取 F1 分数（综合考虑准确率和召回率）
        rouge1_scores.append(scores['rouge1'].fmeasure)
        rouge2_scores.append(scores['rouge2'].fmeasure)
        rougeL_scores.append(scores['rougeL'].fmeasure)

        print(f"\n问题: {test_case['question']}")
        print(f"参考答案: {test_case['reference']}")
        print(f"生成回答: {answer}")
        print(f"ROUGE-1: {scores['rouge1'].fmeasure:.4f}")
        print(f"ROUGE-2: {scores['rouge2'].fmeasure:.4f}")
        print(f"ROUGE-L: {scores['rougeL'].fmeasure:.4f}")

    # 计算平均分数
    avg_rouge1 = sum(rouge1_scores) / len(rouge1_scores)
    avg_rouge2 = sum(rouge2_scores) / len(rouge2_scores)
    avg_rougeL = sum(rougeL_scores) / len(rougeL_scores)

    print(f"\n平均ROUGE-1: {avg_rouge1:.4f}")
    print(f"平均ROUGE-2: {avg_rouge2:.4f}")
    print(f"平均ROUGE-L: {avg_rougeL:.4f}")

    # 验证质量
    assert avg_rouge1 >= 0.3, "ROUGE-1分数过低"

def test_relevance_evaluation():
    """测试相关性评估"""
    rag = RAGSystem(api_key="your-api-key")

    # 添加文档
    documents = [
        "机器学习是AI的一个分支，使计算机能够从数据中学习。",
        "深度学习使用神经网络，是机器学习的一个子集。",
        "监督学习使用标记数据进行训练。",
    ]

    rag.retriever.add_documents(documents)

    # 测试查询
    question = "什么是深度学习？"
    result = rag.query(question)

    answer = result['answer']
    retrieved_docs = result['retrieved_docs']

    print(f"问题: {question}")
    print(f"回答: {answer}")
    print(f"检索到的文档: {len(retrieved_docs)} 个")

    # 检查回答是否包含检索文档中的关键信息（评估端到端相关性）
    retrieved_content = " ".join([doc['content'] for doc in retrieved_docs])

    # 定义期望在回答中出现的关键词
    keywords = ["神经网络", "机器学习", "子集"]
    # 统计实际命中的关键词
    found_keywords = [kw for kw in keywords if kw in answer]

    print(f"\n找到的关键词: {found_keywords}")

    # 至少命中 2 个关键词才算相关性达标
    assert len(found_keywords) >= 2, "回答与检索文档相关性不足"