Agent 学习：构建丐版 RAG 系统

def chunk_text_fixed_length(text, chunk_size=200, overlap=20):
    """
    固定长度切分，带重叠窗口
    :param text: 输入文本
    :param chunk_size: 每块字符数
    :param overlap: 重叠字符数
    :return: 文本块列表
    """
    chunks = []
    start = 0

    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)

        # 下一块从 end-overlap 开始
        start = end - overlap

        # 防止无限循环
        if start >= len(text):
            break

        if start <= 0:  # 处理极短文本
            break

    return chunks

# 测试
sample_text = """
人工智能（Artificial Intelligence，AI）是计算机科学的一个分支，
它企图了解智能的实质，并生产出一种新的能以人类智能相似的方式做出反应的智能机器。
包括机器人、语言识别、图像识别、自然语言处理和专家系统等。
人工智能的核心问题包括建构能够表现出与人类智能相关的能力，
如推理、学习和解决问题的能力。机器学习是人工智能的一个分支，
它使用统计技术让计算机从数据中"学习"，而不需要明确编程。
深度学习是机器学习的子集，它使用多层神经网络来模拟人类大脑的学习过程。
"""

chunks = chunk_text_fixed_length(sample_text, chunk_size=100, overlap=20)
print(f"原始文本长度: {len(sample_text)}")
print(f"切分成 {len(chunks)} 个片段:")
for i, chunk in enumerate(chunks):
    print(f"\n--- 片段 {i+1} ---")
    print(f"长度: {len(chunk)}")
    print(chunk)

import re

def chunk_text_sentence_aware(text, max_chars=200, max_sentences=3):
    """
    基于句子的智能切分
    :param text: 输入文本
    :param max_chars: 最大字符数
    :param max_sentences: 最大句子数
    :return: 文本块列表
    """
    # 按句子分割（使用更智能的正则）
    sentences = re.split(r'[。！？.!?]\s*', text)
    sentences = [s.strip() for s in sentences if s.strip()]

    chunks = []
    current_chunk = ""

    for sentence in sentences:
        # 如果加上这个句子会超出限制，且当前chunk不为空
        if len(current_chunk + sentence) > max_chars and current_chunk:
            chunks.append(current_chunk)
            current_chunk = sentence
        else:
            if current_chunk:
                current_chunk += "。" + sentence
            else:
                current_chunk = sentence

        # 如果chunk中的句子数已经达到上限
        if current_chunk.count('。') + current_chunk.count('!') + current_chunk.count('?') >= max_sentences:
            chunks.append(current_chunk)
            current_chunk = ""

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

    return chunks

# 测试
chunks = chunk_text_sentence_aware(sample_text, max_chars=150)
print(f"语义切分成 {len(chunks)} 个片段")
for i, chunk in enumerate(chunks):
    print(f"\n--- 片段 {i+1} ---")
    print(chunk)

# experiment_1_embedding.py
from sentence_transformers import SentenceTransformer
import numpy as np

# 加载预训练模型
model = SentenceTransformer('all-MiniLM-L6-v2')

def embed_chunks(chunks):
    """
    将文本块转换为向量
    :param chunks: 文本块列表
    :return: 向量数组 (numpy array)
    """
    embeddings = model.encode(chunks)
    return embeddings

# 实验
chunks = chunk_text_fixed_length(sample_text, chunk_size=100, overlap=20)
embeddings = embed_chunks(chunks)

print(f"文本块数量: {len(chunks)}")
print(f"向量维度: {embeddings.shape}")  # 应该是 (chunk_count, 384)
print(f"向量类型: {type(embeddings)}")

# 查看第一个向量的前5个维度
print(f"\n第一个向量 (前5维): {embeddings[0][:5]}")

import os

def chunk_xyj_fixed_length(text, chunk_size=200, overlap=50):
    """
    固定长度切分，带重叠窗口
    :param text: 输入文本
    :param chunk_size: 每块字符数
    :param overlap: 重叠字符数
    :return: 文本块列表
    """
    chunks = []
    start = 0

    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)

        # 下一块从 end-overlap 开始
        start = end - overlap

        # 防止无限循环
        if start >= len(text):
            break

        if start <= 0:  # 处理极短文本
            break

    return chunks


# 读取《西游记》,文档已经提前将前三回的内容拆出来了
with open('xyj.txt', 'r', encoding='utf-8') as f:
    sample_text = ''.join(f.readlines())


chunks = chunk_text_fixed_length(sample_text, chunk_size=100, overlap=20)
print(f"原始文本长度: {len(sample_text)}")
print(f"切分成 {len(chunks)} 个片段:")
for i, chunk in enumerate(chunks):
    print(f"\n--- 片段 {i+1} ---")
    print(f"长度: {len(chunk)}")
    print(chunk)

# 使用语义切分
import re

def chunk_text_sentence_aware(text, max_chars=200, max_sentences=3):
    """
    基于句子的智能切分
    :param text: 输入文本
    :param max_chars: 最大字符数
    :param max_sentences: 最大句子数
    :return: 文本块列表
    """
    # 按句子分割（使用更智能的正则）
    sentences = re.split(r'[。！？.!?]\s*', text)
    sentences = [s.strip() for s in sentences if s.strip()]

    chunks = []
    current_chunk = ""

    for sentence in sentences:
        # 如果加上这个句子会超出限制，且当前chunk不为空
        if len(current_chunk + sentence) > max_chars and current_chunk:
            chunks.append(current_chunk)
            current_chunk = sentence
        else:
            if current_chunk:
                current_chunk += "。" + sentence
            else:
                current_chunk = sentence

        # 如果chunk中的句子数已经达到上限
        if current_chunk.count('。') + current_chunk.count('!') + current_chunk.count('?') >= max_sentences:
            chunks.append(current_chunk)
            current_chunk = ""

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

    return chunks

# 测试
chunks = chunk_text_sentence_aware(sample_text, max_chars=150)
print(f"语义切分成 {len(chunks)} 个片段")
for i, chunk in enumerate(chunks):
    print(f"\n--- 片段 {i+1} ---")
    print(chunk)

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

class SimpleVectorDB:
    def __init__(self, dimension):
        """
        初始化向量数据库
        :param dimension: 向量维度
        """
        self.dimension = dimension
        # 创建FAISS索引 (L2距离，余弦相似度=1-距离)
        self.index = faiss.IndexFlatL2(dimension)
        # 存储原始文本块
        self.texts = []

    def add_vectors(self, vectors, texts):
        """
        添加向量和文本
        :param vectors: numpy array of vectors
        :param texts: 对应的文本列表
        """
        # 确保向量是float32类型
        vectors = np.array(vectors).astype('float32')
        # 添加到索引
        self.index.add(vectors)
        self.texts.extend(texts)

    def search(self, query_vector, k=3):
        """
        搜索最相似的k个向量
        :param query_vector: 查询向量 (numpy array)
        :param k: 返回top-k结果
        :return: ( distances, indices, texts )
        """
        query_vector = np.array([query_vector]).astype('float32')

        # 搜索
        distances, indices = self.index.search(query_vector, k)

        # 获取对应的文本
        retrieved_texts = [self.texts[idx] for idx in indices[0]]

        return distances[0], indices[0], retrieved_texts

# 实验
model = SentenceTransformer('all-MiniLM-L6-v2')

# 准备数据
documents = [
    "人工智能是计算机科学的重要分支",
    "深度学习使用多层神经网络",
    "机器学习是人工智能的子集",
    "Python是一种编程语言",
    "大数据需要高效的数据处理技术"
]

# 向量化
embeddings = model.encode(documents)

# 构建数据库
vector_db = SimpleVectorDB(dimension=embeddings.shape[1])
vector_db.add_vectors(embeddings, documents)

# 模拟用户查询
query = "什么是深度学习？"
query_embedding = model.encode([query])

# 搜索
distances, indices, retrieved_texts = vector_db.search(query_embedding[0], k=3)

print(f"查询: {query}")
print(f"\nTop-3 检索结果:")
for i, (dist, text) in enumerate(zip(distances, retrieved_texts)):
    # 转换为相似度分数 (0-1, 1为完全相似)
    similarity = 1 / (1 + dist)
    print(f"\n结果 {i+1}:")
    print(f"  相似度分数: {similarity:.4f}")
    print(f"  距离: {dist:.4f}")
    print(f"  文本: {text}")

import time
import matplotlib.pyplot as plt

class PerformanceBenchmark:
    def __init__(self, vector_db):
        self.vector_db = vector_db

    def benchmark_search(self, query_embeddings, k=5):
        """
        基准测试检索性能
        """
        latencies = []

        for query_emb in query_embeddings:
            start_time = time.time()
            distances, indices, texts = self.vector_db.search(query_emb, k=k)
            end_time = time.time()

            latency = (end_time - start_time) * 1000  # 转换为毫秒
            latencies.append(latency)

        avg_latency = np.mean(latencies)
        p95_latency = np.percentile(latencies, 95)
        p99_latency = np.percentile(latencies, 99)

        return {
            'avg_latency_ms': avg_latency,
            'p95_latency_ms': p95_latency,
            'p99_latency_ms': p99_latency,
            'min_latency_ms': min(latencies),
            'max_latency_ms': max(latencies)
        }

# 生成测试数据
test_queries = ["人工智能发展", "机器学习算法", "深度神经网络", "Python编程"] * 10  # 40个查询
test_embeddings = model.encode(test_queries)

# 基准测试
benchmark = PerformanceBenchmark(vector_db)
results = benchmark.benchmark_search(test_embeddings, k=5)

print("=== 检索性能基准测试 ===")
print(f"平均延迟: {results['avg_latency_ms']:.2f} ms")
print(f"P95 延迟: {results['p95_latency_ms']:.2f} ms")
print(f"P99 延迟: {results['p99_latency_ms']:.2f} ms")
print(f"最小延迟: {results['min_latency_ms']:.2f} ms")
print(f"最大延迟: {results['max_latency_ms']:.2f} ms")

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

class SimpleVectorDB:
    def __init__(self, dimension):
        """
        初始化向量数据库
        :param dimension: 向量维度
        """
        self.dimension = dimension
        # 创建FAISS索引 (L2距离，余弦相似度=1-距离)
        self.index = faiss.IndexFlatL2(dimension)
        # 存储原始文本块
        self.texts = []

    def add_vectors(self, vectors, texts):
        """
        添加向量和文本
        :param vectors: numpy array of vectors
        :param texts: 对应的文本列表
        """
        # 确保向量是float32类型
        vectors = np.array(vectors).astype('float32')
        # 添加到索引
        self.index.add(vectors)
        self.texts.extend(texts)

    def search(self, query_vector, k=3):
        """
        搜索最相似的k个向量
        :param query_vector: 查询向量 (numpy array)
        :param k: 返回top-k结果
        :return: ( distances, indices, texts )
        """
        query_vector = np.array([query_vector]).astype('float32')

        # 搜索
        distances, indices = self.index.search(query_vector, k)

        # 获取对应的文本
        retrieved_texts = [self.texts[idx] for idx in indices[0]]

        return distances[0], indices[0], retrieved_texts

# 实验
model = SentenceTransformer('all-MiniLM-L6-v2')

# 准备数据
with open('anthropic.txt', 'r', encoding='utf-8') as f:
    anthropic = ''.join(f.readlines())
documents = anthropic.split('\n')  # 使用 Anthropic 百科文本进行实验 

# 向量化
embeddings = model.encode(documents)

# 构建数据库
vector_db = SimpleVectorDB(dimension=embeddings.shape[1])
vector_db.add_vectors(embeddings, documents)

# 模拟用户查询
query = "Claude是哪个公司开发的？"
query_embedding = model.encode([query])

# 搜索
distances, indices, retrieved_texts = vector_db.search(query_embedding[0], k=3)

print(f"查询: {query}")
print(f"\nTop-3 检索结果:")
for i, (dist, text) in enumerate(zip(distances, retrieved_texts)):
    # 转换为相似度分数 (0-1, 1为完全相似)
    similarity = 1 / (1 + dist)
    print(f"\n结果 {i+1}:")
    print(f"  相似度分数: {similarity:.4f}")
    print(f"  距离: {dist:.4f}")
    print(f"  文本: {text}")
    
print("\n--------------------------------\n")
# 模拟用户查询
query2 = "Anthropic创始人是谁？"
query_embedding2 = model.encode([query2])

# 搜索
distances, indices, retrieved_texts = vector_db.search(query_embedding2[0], k=3)

print(f"查询: {query2}")
print(f"\nTop-3 检索结果:")
for i, (dist, text) in enumerate(zip(distances, retrieved_texts)):
    # 转换为相似度分数 (0-1, 1为完全相似)
    similarity = 1 / (1 + dist)
    print(f"\n结果 {i+1}:")
    print(f"  相似度分数: {similarity:.4f}")
    print(f"  距离: {dist:.4f}")
    print(f"  文本: {text}")

from openai import OpenAI
import os

# 设置API Key（需要您自己的API Key，本例子使用deepseek）
# os.environ['API_KEY'] = 'your-api-key-here'

class SimpleVectorDB:
    def __init__(self, dimension):
        """
        初始化向量数据库
        :param dimension: 向量维度
        """
        self.dimension = dimension
        # 创建FAISS索引 (L2距离，余弦相似度=1-距离)
        self.index = faiss.IndexFlatL2(dimension)
        # 存储原始文本块
        self.texts = []

    def add_vectors(self, vectors, texts):
        """
        添加向量和文本
        :param vectors: numpy array of vectors
        :param texts: 对应的文本列表
        """
        # 确保向量是float32类型
        vectors = np.array(vectors).astype('float32')
        # 添加到索引
        self.index.add(vectors)
        self.texts.extend(texts)

    def search(self, query_vector, k=3):
        """
        搜索最相似的k个向量
        :param query_vector: 查询向量 (numpy array)
        :param k: 返回top-k结果
        :return: ( distances, indices, texts )
        """
        query_vector = np.array([query_vector]).astype('float32')

        # 搜索
        distances, indices = self.index.search(query_vector, k)

        # 获取对应的文本
        retrieved_texts = [self.texts[idx] for idx in indices[0]]

        return distances[0], indices[0], retrieved_texts


class SimpleRAG:
    def __init__(self, vector_db, model_name="deepseek-chat"):
        self.vector_db = vector_db
        self.model_name = model_name
        # 使用本地模型替代（如果无法访问OpenAI API）
        self.use_local = False
        self.local_model = SentenceTransformer('all-MiniLM-L6-v2')

    def build_prompt(self, query, retrieved_texts, max_context_length=2000):
        """
        构建 RAG Prompt
        :param query: 用户问题
        :param retrieved_texts: 检索到的文本列表
        :param max_context_length: 最大上下文长度（字符）
        :return: 完整的Prompt
        """
        # 拼接上下文
        context = ",".join(retrieved_texts)

        # 限制上下文长度
        if len(context) > max_context_length:
            context = context[:max_context_length]

        prompt = f"""你是一个专业的问答助手。请基于提供的上下文信息回答用户问题。
                    上下文信息：{context}
                    用户问题：{query}
                    回答要求：
                    1. 仅基于上下文信息回答，不要添加上下文外的信息
                    2. 如果上下文信息不足以回答问题，请明确说明
                    3. 回答要简洁、准确、有条理
                    4. 如果涉及数字或事实，请确保准确性
                    回答：
                """

        return prompt

    def call_llm(self, prompt):
        """
        调用 LLM 生成回答
        """
        if self.use_local:
            # 使用简单规则引擎模拟LLM回答（仅用于测试）
            return self.simple_rule_based_answer(prompt)
        else:
            # 调用 DeepSeek API
            client = OpenAI(
                api_key=os.getenv('DEEPSEEK_API'),
                base_url="https://api.deepseek.com/v1"  # DeepSeek API 端点
            )

            response = client.chat.completions.create(
                model="deepseek-chat",  # 例如 "deepseek-chat"
                messages=[
                    {"role": "user", "content": prompt}
                ],
                temperature=0.7,
                max_tokens=1000
            )
            return response.choices[0].message.content

    def simple_rule_based_answer(self, prompt):
        """
        简单的基于规则的问答（仅用于测试环境）
        实际生产环境应使用真正的 LLM
        """
        # 提取问题和上下文
        lines = prompt.split('\n')
        question = ""
        context = ""

        in_question = False
        in_context = False

        for line in lines:
            if "用户问题：" in line:
                in_question = True
                in_context = False
                question = line.replace("用户问题：", "").strip()
            elif "上下文信息：" in line:
                in_context = True
                in_question = False
            elif line.startswith("回答："):
                break
            elif in_context:
                context += line + "\n"

        
        for keyword, answer in keywords.items():
            if keyword in question or keyword in context:
                return f"基于上下文信息，{answer}"

        return "抱歉，基于提供的上下文信息，我无法回答这个问题。"

    def query(self, question, k=3):
        """
        完整的RAG查询流程
        """
        # 1. 向量化问题
        query_embedding = self.local_model.encode([question])[0]

        # 2. 检索相关内容
        distances, indices, retrieved_texts = self.vector_db.search(query_embedding, k=k)

        # 3. 构建Prompt
        prompt = self.build_prompt(question, retrieved_texts)

        # 4. 生成回答
        answer = self.call_llm(prompt)

        return {
            'question': question,
            'retrieved_texts': retrieved_texts,
            'distances': distances.tolist(),
            'prompt': prompt,
            'answer': answer
        }

# 简单的关键词匹配回答

documents = [
    "人工智能是计算机科学的重要分支",
    "深度学习使用多层神经网络",
    "机器学习是人工智能的子集",
    "Python是一种编程语言",
    "大数据需要高效的数据处理技术"
]

keywords = {
    "人工智能": "人工智能是计算机科学的重要分支，它企图了解智能的实质。",
    "机器学习": "机器学习是人工智能的一个分支，使用统计技术让计算机从数据中学习。",
    "深度学习": "深度学习是机器学习的子集，使用多层神经网络模拟人脑学习。",
    "Python": "Python是一种高级编程语言，广泛用于数据科学和AI领域。"
}

# 向量化
embeddings = model.encode(documents)

# 构建数据库
vector_db = SimpleVectorDB(dimension=embeddings.shape[1])
vector_db.add_vectors(embeddings, documents)

# 实验
rag = SimpleRAG(vector_db)

# 测试查询
query = "人工智能包括哪些分支？"
result = rag.query(query, k=2)

print(f"=" * 80)
print(f"用户问题: {result['question']}")
print(f"=" * 80)
print(f"\n检索到的文档:")
for i, (text, dist) in enumerate(zip(result['retrieved_texts'], result['distances'])):
    similarity = 1 / (1 + dist)
    print(f"\n文档 {i+1} (相似度: {similarity:.4f}):")
    print(f"  {text}")

print(f"\n" + "=" * 80)
print(f"生成的Prompt:")
print(f"=" * 80)
print(result['prompt'])

print(f"\n" + "=" * 80)
print(f"LLM 回答:")
print(f"=" * 80)
print(result['answer'])

import openai
import os

class AdvancedRAG(SimpleRAG):
    
    def __init__(self, vector_db, model_name="deepseek-chat"):
        self.vector_db = vector_db
        self.model_name = model_name
        # 使用本地模型替代（如果无法访问OpenAI API）
        self.use_local = False
        self.local_model = SentenceTransformer('all-MiniLM-L6-v2')
    
    def build_prompt(self, query, retrieved_texts, metadata_list=None):
        """
        构建更高级的Prompt（包含元数据、结构化输出等）
        """
        if metadata_list is None:
            metadata_list = [{}] * len(retrieved_texts)

        # 构建带编号的上下文
        numbered_context = []
        for i, (text, metadata) in enumerate(zip(retrieved_texts, metadata_list)):
            source_info = ""
            if metadata.get('source'):
                source_info = f" [来源: {metadata['source']}]"
            numbered_context.append(f"[{i+1}] {text}{source_info}")

        context = ",".join(numbered_context)

        prompt = f"""你是一个专业的问答助手，擅长基于提供的文档内容回答问题。
                    ## 任务说明
                    基于以下检索到的文档片段回答用户问题。文档已按相关性排序，编号[1]表示最相关。
                    ## 文档内容
                    {context}
                    ## 用户问题
                    {query}
                    ## 回答要求
                    1. **严格基于文档**：仅使用提供的文档内容，不要添加任何外部知识
                    2. **引用来源**：在回答中标注使用的文档编号，如 [1][2]
                    3. **完整性**：如果单个文档无法完整回答，请综合多个文档的信息
                    4. **诚实性**：如果文档中没有答案，明确说明"文档中没有相关信息"
                    5. **格式**：使用清晰的段落和要点组织答案
                    ## 回答格式
                    **答案：**
                    [你的回答]
                    **参考文档：** [列出使用的文档编号]
                    ---"""

        return prompt

    def query_with_metadata(self, question, k=3):
        """
        带元数据的查询
        """
        query_embedding = self.local_model.encode([question])[0]
        distances, indices, retrieved_texts = self.vector_db.search(query_embedding, k=k)

        # 模拟元数据（实际应用中应从数据库获取）
        metadata_list = [
            {'source': f'doc_{idx}.txt', 'page': 1}
            for idx in indices
        ]

        prompt = self.build_prompt(question, retrieved_texts, metadata_list)
        answer = self.call_llm(prompt)

        return {
            'question': question,
            'retrieved_texts': retrieved_texts,
            'metadata_list': metadata_list,
            'prompt': prompt,
            'answer': answer
        }
        
    def call_llm(self, prompt):
        """
        调用 LLM 生成回答
        """
        if self.use_local:
            # 使用简单规则引擎模拟LLM回答（仅用于测试）
            return self.simple_rule_based_answer(prompt)
        else:
            # 调用 DeepSeek API
            client = OpenAI(
                api_key=os.getenv('DEEPSEEK_API'),
                base_url="https://api.deepseek.com/v1"  # DeepSeek API 端点
            )

            response = client.chat.completions.create(
                model="deepseek-chat",  # 例如 "deepseek-chat"
                messages=[
                    {"role": "user", "content": prompt}
                ],
                temperature=0.7,
                max_tokens=1000
            )
            return response.choices[0].message.content

    def simple_rule_based_answer(self, prompt):
        """
        简单的基于规则的问答（仅用于测试环境）
        实际生产环境应使用真正的 LLM
        """
        # 提取问题和上下文
        lines = prompt.split('\n')
        question = ""
        context = ""

        in_question = False
        in_context = False

        for line in lines:
            if "用户问题：" in line:
                in_question = True
                in_context = False
                question = line.replace("用户问题：", "").strip()
            elif "上下文信息：" in line:
                in_context = True
                in_question = False
            elif line.startswith("回答："):
                break
            elif in_context:
                context += line + "\n"

        
        for keyword, answer in keywords.items():
            if keyword in question or keyword in context:
                return f"基于上下文信息，{answer}"

        return "抱歉，基于提供的上下文信息，我无法回答这个问题。"


# 实验
rag = AdvancedRAG(vector_db)

# 测试查询
query = "人工智能包括哪些分支？"
result = rag.query_with_metadata(query, k=2)

print(f"=" * 80)
print(f"用户问题: {result['question']}")
print(f"=" * 80)
print(f"\n检索到的文档:")
for i, (text, list) in enumerate(zip(result['retrieved_texts'], result['metadata_list'])):
    print(f"  {text}")

print(f"\n" + "=" * 80)
print(f"生成的Prompt:")
print(f"=" * 80)
print(result['prompt'])

print(f"\n" + "=" * 80)
print(f"LLM 回答:")
print(f"=" * 80)
print(result['answer'])

import os
import json
from openai import OpenAI
from sentence_transformers import SentenceTransformer

# 加载预训练模型
model = SentenceTransformer('all-MiniLM-L6-v2')

import faiss
import numpy as np

class SimpleVectorDB:
    def __init__(self, dimension):
        """
        初始化向量数据库
        :param dimension: 向量维度
        """
        self.dimension = dimension
        # 创建FAISS索引 (L2距离，余弦相似度=1-距离)
        self.index = faiss.IndexFlatL2(dimension)
        # 存储原始文本块
        self.texts = []

    def add_vectors(self, vectors, texts):
        """
        添加向量和文本
        :param vectors: numpy array of vectors
        :param texts: 对应的文本列表
        """
        # 确保向量是float32类型
        vectors = np.array(vectors).astype('float32')
        # 添加到索引
        self.index.add(vectors)
        self.texts.extend(texts)

    def search(self, query_vector, k=3):
        """
        搜索最相似的k个向量
        :param query_vector: 查询向量 (numpy array)
        :param k: 返回top-k结果
        :return: ( distances, indices, texts )
        """
        query_vector = np.array([query_vector]).astype('float32')

        # 搜索
        distances, indices = self.index.search(query_vector, k)

        # 获取对应的文本
        retrieved_texts = [self.texts[idx] for idx in indices[0]]

        return distances[0], indices[0], retrieved_texts

def chunk_text_fixed_length(text, chunk_size=300, overlap=50):
        """
        固定长度切分，带重叠窗口
        :param text: 输入文本
        :param chunk_size: 每块字符数
        :param overlap: 重叠字符数
        :return: 文本块列表
        """
        chunks = []
        start = 0

        while start < len(text):
            end = start + chunk_size
            chunk = text[start:end]
            chunks.append(chunk)

            # 下一块从 end-overlap 开始
            start = end - overlap

            # 防止无限循环
            if start >= len(text):
                break

            if start <= 0:  # 处理极短文本
                break

        return chunks

class NovelRAGSystem:
    def __init__(self):
        self.vector_db = None
        self.embeddings = None
        self.chunks = []
        self.metadata = []

    def load_novel_data(self, file_path):
        """
        加载小说文本文件（txt格式）
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            text = f.read()

        # 切分章节（假设用"第X回"分隔）
        import re
        chapters = re.split(r'第\d+回', text)
        chapters = [ch.strip() for ch in chapters if ch.strip()]

        return chapters
    
    def call_llm(self, prompt):
        """
        调用 LLM 生成回答
        """
        # 调用 DeepSeek API
        client = OpenAI(
            api_key=os.getenv('DEEPSEEK_API'),
            base_url="https://api.deepseek.com/v1"  # DeepSeek API 端点
        )

        response = client.chat.completions.create(
            model="deepseek-chat",  # 例如 "deepseek-chat"
            messages=[
                {"role": "user", "content": prompt}
            ],
            temperature=0.7,
            max_tokens=2000
        )
        return response.choices[0].message.content
    
    def preprocess_and_index(self, text_chunks, chunk_size=300, overlap=50):
        """
        预处理和索引文档
        """
        # 1. 切分文本
        all_chunks = []

        for chapter_num, chapter in enumerate(text_chunks):
            chunks = chunk_text_fixed_length(chapter, chunk_size, overlap)

            # 为每个chunk添加元数据
            for chunk_idx, chunk in enumerate(chunks):
                all_chunks.append({
                    'text': chunk,
                    'chapter': chapter_num + 1,
                    'chunk_id': chunk_idx + 1
                })

        self.chunks = all_chunks

        # 2. 向量化
        chunk_texts = [chunk['text'] for chunk in all_chunks]
        self.embeddings = model.encode(chunk_texts)

        # 3. 构建向量数据库
        self.vector_db = SimpleVectorDB(dimension=self.embeddings.shape[1])
        self.vector_db.add_vectors(
            self.embeddings,
            chunk_texts
        )

        print(f"   -索引完成！")
        print(f"   - 总chunk数: {len(all_chunks)}")
        print(f"   - 向量维度: {self.embeddings.shape[1]}")

    def query(self, question, top_k=5):
        """
        查询回答
        """
        if self.vector_db is None:
            raise Exception("请先运行 preprocess_and_index() 构建索引")

        query_embedding = model.encode([question])[0]
        distances, indices, retrieved_texts = self.vector_db.search(query_embedding, top_k)

        # 构建prompt
        prompt = self.build_prompt(question, retrieved_texts)

        # 调用LLM
        answer = self.call_llm(prompt)

        return {
            'question': question,
            'answer': answer,
            'prompt': prompt,
            'retrieved_chunks': [
                {
                    'text': retrieved_texts[i],
                    'distance': distances[i],
                    'metadata': self.chunks[indices[i]]
                }
                for i in range(len(retrieved_texts))
            ]
        }

    def build_prompt(self, question, retrieved_texts):
        """
        构建小说问答的专用Prompt
        """
        context = "\n\n".join([
            f"[片段{i+1}] {text}"
            for i, text in enumerate(retrieved_texts)
        ])

        prompt = f"""你是一个红学专家，对《红楼梦》有深入研究。请基于提供的文本片段回答用户关于红楼梦的问题。
                ## 重要原则
                1. 严格基于提供的文本片段回答
                2. 如果需要引用情节，请说明在哪个片段中找到
                3. 对于人物关系，要准确引用原文描述
                4. 如果片段中没有相关信息，明确说明
                ## 文本片段
                {context}
                ## 问题
                {question}
                ## 回答
                """

        return prompt

    def evaluate(self, test_queries):
        """
        评估系统性能
        """
        results = []

        for query_info in test_queries:
            query = query_info['question']
            expected = query_info.get('expected_answer', '')

            result = self.query(query)
            result['expected'] = expected

            results.append(result)

        return results

question = "《五美吟》的内容是什么？"
result = rag.query(question, top_k=5)

print(f"=" * 80)
print(f"用户问题: {result['question']}")
print(f"=" * 80)
print(f"\n检索到的块:")
print(result['retrieved_chunks'])

print(f"\n" + "=" * 80)
print(f"生成的Prompt:")
print(f"=" * 80)
print(result['prompt'])

print(f"\n" + "=" * 80)
print(f"LLM 回答:")
print(f"=" * 80)
print(result['answer'])