AI知识库搭建全流程实战：从向量数据库选型到RAG系统部署

# 创建Python虚拟环境（推荐使用 conda）
conda create -n rag python=3.11 -y
conda activate rag

# 安装核心依赖
# chromadb: 向量数据库
# langchain: RAG框架
# langchain-openai/langchain-anthropic: 大模型API
# sentence-transformers: Embedding模型
pip install chromadb langchain langchain-community langchain-openai \
    sentence-transformers pypdf python-docx tqdm

import chromadb

# 创建一个持久化的ChromaDB客户端（数据保存在本地磁盘）
client = chromadb.PersistentClient(path="./chroma_data")

# 创建或获取一个集合（类似关系数据库的表）
collection = client.get_or_create_collection(name="my_knowledge_base")

# 插入一条向量数据测试
collection.add(
    ids=["doc_001"],  # 文档ID，唯一
    documents=["这是一个测试文档，用于验证ChromaDB是否正常工作。"],  # 原始文本
    metadatas=[{"source": "test", "page": 1}]  # 元数据，用于过滤
)

# 查询最相似的文档
results = collection.query(
    query_texts=["验证系统是否工作"],
    n_results=1
)

print(results)
# 输出类似：{'ids': [['doc_001']], 'documents': [['这是一个测试文档...']], 'distances': [[0.32]]}
# distances是余弦距离，越小越相似，0表示完全匹配

# 方式一：Docker部署Qdrant（推荐用于开发测试）
docker pull qdrant/qdrant:latest
docker run -d \
  --name qdrant \
  -p 6333:6333 \      # REST API 端口
  -p 6334:6334 \      # gRPC 端口（用于高性能查询）
  -v $(pwd)/qdrant_storage:/qdrant/storage \
  qdrant/qdrant:latest

# 验证Qdrant是否启动成功
curl http://localhost:6333/collections 2>/dev/null | python3 -m json.tool
# 输出：{"result": {"collections": []}, "status": "ok"}

# 安装Qdrant Python客户端
pip install qdrant-client

# 连接Qdrant
from qdrant_client import QdrantClient

client = QdrantClient(host="localhost", port=6333)

# 创建集合，指定向量维度（text-embedding-ada-002 输出1536维）
client.create_collection(
    collection_name="my_knowledge_base",
    vectors_config={
        "size": 1536,       # Embedding维度
        "distance": "Cosine"  # 余弦距离，1表示完全相反，0表示完全相同
    }
)
print("集合创建成功！")

from sentence_transformers import SentenceTransformer
import torch

# 选择设备：优先使用GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"使用设备：{device}")

# 加载BGE-large中文模型（约1.1GB，第一次加载会下载）
# 模型下载地址：https://huggingface.co/BAAI/bge-large-zh-v1.5
model_name = "BAAI/bge-large-zh-v1.5"
model = SentenceTransformer(model_name, device=device)

# 测试Embedding效果
test_texts = [
    "人工智能是什么？",
    "机器学习和深度学习的区别是什么？",
    "Python编程基础教程"
]

# encode方法支持批量处理，性能比逐条处理高很多
embeddings = model.encode(test_texts, normalize_embeddings=True)
# normalize_embeddings=True 表示返回单位向量，余弦相似度计算简化为点积

print(f"Embedding矩阵形状：{embeddings.shape}")
# 输出：Embedding矩阵形状：(3, 1024)
# 3条文本，每条1024维向量

# 计算两条文本的相似度（用余弦相似度）
from numpy import dot
from numpy.linalg import norm

def cosine_sim(a, b):
    return dot(a, b) / (norm(a) * norm(b))

# 比较第一和第二条文本的相似度
sim = cosine_sim(embeddings[0], embeddings[1])
print(f"文本1与文本2的余弦相似度：{sim:.4f}")
# 预期输出：约0.75-0.85（两条都是AI相关，相似度高）

import os
from openai import OpenAI

# 设置API Key（建议用环境变量，不写在代码里）
os.environ["OPENAI_API_KEY"] = "sk-xxxxx"  # 替换成你的真实Key

client = OpenAI()

def get_embedding(text: str, model: str = "text-embedding-3-small") -> list:
    """获取文本的Embedding向量"""
    response = client.embeddings.create(
        model=model,
        input=text
    )
    # 返回向量（OpenAI返回格式：data[0].embedding）
    return response.data[0].embedding

# 测试
emb = get_embedding("人工智能技术应用")
print(f"向量维度：{len(emb)}")  # text-embedding-3-small 输出1536维
print(f"向量前5位：{emb[:5]}")  # 浮点数数组

# ============================================
# 策略1：固定大小分块（最简单，但效果一般）
# ============================================
def fixed_chunk(text: str, chunk_size: int = 500, overlap: int = 50) -> list:
    """
    固定大小分块
    
    参数:
        text: 输入文本
        chunk_size: 每块token数（中文约1字符≈1token，英文约4字符≈1token）
        overlap: 块之间的重叠token数，保持上下文连贯
    
    返回:
        块列表
    """
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)
        # 滑动窗口：起点移动 chunk_size - overlap
        start = start + chunk_size - overlap
    
    return chunks

# 示例
sample_text = "人工智能是计算机科学的一个分支，它企图了解智能的实质，并生产出一种新的能以人类智能相似的方式做出反应的智能机器。该领域的研究包括机器人、语言识别、图像识别、自然语言处理和专家系统等。"
chunks = fixed_chunk(sample_text, chunk_size=50, overlap=10)
print(f"分块数量：{len(chunks)}")
for i, c in enumerate(chunks):
    print(f"块{i+1}：{c}")
    print("---")

# ============================================
# 策略2：语义分块（按段落/句子边界，更智能）
# ============================================
import re

def semantic_chunk(text: str, min_chars: int = 100, max_chars: int = 1000) -> list:
    """
    语义分块：按段落边界分块，同时限制单块长度
    
    参数:
        text: 输入文本
        min_chars: 单块最小字符数（太小则合并到下一块）
        max_chars: 单块最大字符数（超长则强制截断）
    
    返回:
        块列表（每块附带元数据：段落索引）
    """
    # 按换行符分割段落（支持 \n \n 或空行）
    paragraphs = re.split(r'\n\s*\n', text)
    
    chunks = []
    current_chunk = ""
    current_para_idx = 0
    
    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
            
        # 如果当前块加上新段落会超过最大值，先保存当前块
        if len(current_chunk) + len(para) > max_chars and current_chunk:
            # 只保存有内容的块
            if len(current_chunk) >= min_chars:
                chunks.append({
                    "content": current_chunk.strip(),
                    "para_index": current_para_idx,
                    "char_count": len(current_chunk)
                })
            current_chunk = ""
            current_para_idx += 1
        
        # 如果单个段落就超长，强制按句子分割
        if len(para) > max_chars:
            sentences = re.split(r'([。!?；])', para)
            for i in range(0, len(sentences)-1, 2):
                sentence = sentences[i] + (sentences[i+1] if i+1 < len(sentences) else "")
                if len(current_chunk) + len(sentence) > max_chars:
                    if len(current_chunk) >= min_chars:
                        chunks.append({
                            "content": current_chunk.strip(),
                            "para_index": current_para_idx
                        })
                    current_chunk = sentence
                    current_para_idx += 1
                else:
                    current_chunk += sentence + "。"
        else:
            current_chunk += para + "\n\n"
    
    # 不要忘记最后一个块
    if len(current_chunk) >= min_chars:
        chunks.append({
            "content": current_chunk.strip(),
            "para_index": current_para_idx
        })
    
    return chunks

# 使用示例
with open("sample_article.txt", "r", encoding="utf-8") as f:
    article_text = f.read()

chunks = semantic_chunk(article_text)
print(f"语义分块结果：共 {len(chunks)} 个块")
for i, chunk in enumerate(chunks[:3]):  # 只打印前3个
    print(f"块{i+1}（{chunk['char_count']}字符）：{chunk['content'][:80]}...")
    print("---")

# ============================================
# 策略3：基于LangChain的RecursiveCharacterTextSplitter
# ============================================
from langchain.text_splitter import RecursiveCharacterTextSplitter

def langchain_chunk(documents: list, chunk_size: int = 500, chunk_overlap: int = 50):
    """
    LangChain官方推荐的分块器
    原理：依次尝试按段落分割、按句子分割、按单词分割，保证每块尽可能完整
    
    参数:
        documents: LangChain Document对象列表
        chunk_size: 每块字符数
        chunk_overlap: 重叠字符数
    
    返回:
        分割后的Document列表
    """
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,           # 每块字符数
        chunk_overlap=chunk_overlap,       # 块之间重叠字符数
        length_function=len,              # 计算文本长度的函数
        separators=["\n\n", "\n", "。", "！", "？", " ", ""]  
        # 优先按段落分割，然后按句子，按中文标点符号，最后按空格
    )
    
    # split_documents接受Document对象列表，保留metadata
    chunks = text_splitter.split_documents(documents)
    
    print(f"分块完成：原始 {len(documents)} 个文档 → {len(chunks)} 个块")
    
    # 打印每个块的信息
    for i, chunk in enumerate(chunks[:5]):
        print(f"块{i+1} | 来源：{chunk.metadata.get('source','unknown')} | 字符数：{len(chunk.page_content)}")
        print(f"  内容预览：{chunk.page_content[:60]}...")
        print()
    
    return chunks

# 使用示例：读取PDF后分块
from langchain_community.document_loaders import PyPDFLoader

loader = PyPDFLoader("sample.pdf")
pages = loader.load_and_split()  # 按页分割
# pages 是 Document 对象列表，每页一个Document

chunks = langchain_chunk(pages, chunk_size=500, chunk_overlap=50)

# ============================================
# 完整RAG管道：PDF知识库问答系统
# ============================================

import os
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA

# ===================== 第1步：加载PDF文档 =====================
print("=" * 50)
print("第1步：加载PDF文档")
print("=" * 50)

loader = PyPDFLoader("knowledge_base/faq.pdf")
documents = loader.load()
print(f"加载文档数：{len(documents)}")
print(f"文档总页数：{len(documents)}页")
# 查看前两页的内容
for i, doc in enumerate(documents[:2]):
    print(f"\n--- 第{i+1}页预览 ---")
    print(doc.page_content[:200] + "...")

# ===================== 第2步：文档分块 =====================
print("\n" + "=" * 50)
print("第2步：文档分块")
print("=" * 50)

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,       # 每块500字符
    chunk_overlap=100,    # 相邻块重叠100字符，保持上下文
    length_function=len,  # 用len函数计算字符数
    separators=["\n\n", "\n", "。", "！", "？"]  # 按优先级尝试分割
)

chunks = text_splitter.split_documents(documents)
print(f"分块结果：{len(documents)} 个文档 → {len(chunks)} 个块")

# 统计每块长度
chunk_lengths = [len(chunk.page_content) for chunk in chunks]
print(f"块长度统计：平均{sum(chunk_lengths)/len(chunk_lengths):.0f}字符")
print(f"  最短：{min(chunk_lengths)}字符，最长：{max(chunk_lengths)}字符")

# ===================== 第3步：生成Embedding并存储到向量数据库 =====================
print("\n" + "=" * 50)
print("第3步：Embedding + 向量存储")
print("=" * 50)

# 初始化Embedding模型（使用BGE-large中文模型）
# 如果是第一次使用，模型会自动从HuggingFace下载（约1.1GB）
embeddings = HuggingFaceBgeEmbeddings(
    model_name="BAAI/bge-large-zh-v1.5",
    model_kwargs={'device': 'cpu'},  # 有GPU可以改成 'cuda'
    encode_kwargs={'normalize_embeddings': True}  # 单位向量，余弦相似度=点积
)

# 将文档块存储到ChromaDB
# persist_directory：持久化存储路径，重启后数据不丢失
vectorstore = Chroma.from_documents(
    documents=chunks,           # 文档块列表
    embedding=embeddings,       # Embedding模型
    persist_directory="./chroma_knowledge_base"  # 数据存储路径
)

print(f"向量数据库创建成功！")
print(f"向量维度：{vectorstore._collection.metadata()['hnsw:dim']}")
print(f"向量总数：{vectorstore._collection.count()}")

# ===================== 第4步：检索器配置 =====================
print("\n" + "=" * 50)
print("第4步：配置检索器")
print("=" * 50)

# 创建检索器
# k=5 表示返回最相关的5个块
retriever = vectorstore.as_retriever(
    search_type="similarity",  # 按相似度检索（还有mmr最大边际相关性）
    search_kwargs={
        "k": 5,                # 返回5个最相关的块
        "score_threshold": 0.7  # 只返回相似度>0.7的块（可选，0-1，越高越严格）
    }
)

# 测试检索效果
test_query = "如何申请产品试用？"
retrieved_docs = retriever.invoke(test_query)

print(f"查询：「{test_query}」")
print(f"检索到 {len(retrieved_docs)} 个相关块")
print()
for i, doc in enumerate(retrieved_docs):
    print(f"--- 结果{i+1}（来源：{doc.metadata.get('source','unknown')}, 页：{doc.metadata.get('page','?')}）---")
    print(doc.page_content[:150] + "...")
    print()

# ===================== 第5步：构建QA Chain（问答链） =====================
print("\n" + "=" * 50)
print("第5步：构建问答链")
print("=" * 50)

# 初始化大模型（使用GPT-4o-mini，性价比高）
llm = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0.3,  # 温度0.3：既有创造性又不会胡编
    api_key=os.getenv("OPENAI_API_KEY")
)

# 创建RetrievalQA链
# retriever=retriever 让模型只从知识库中找答案
# return_source_documents=True 让模型回答时附带引用来源
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",  # stuff模式：将所有检索到的块拼成Prompt
    retriever=retriever,
    return_source_documents=True,  # 返回使用的文档块（可用于溯源）
    chain_type_kwargs={
        "verbose": True  # 打印中间步骤日志
    }
)

print("QA链创建成功！")

# ===================== 第6步：执行问答 =====================
print("\n" + "=" * 50)
print("第6步：测试问答")
print("=" * 50)

# 问一个知识库相关的问题
query = "你们的产品支持哪些支付方式？"

print(f"问题：{query}")
print("-" * 50)

result = qa_chain.invoke({"query": query})

print("\n📝 AI回答：")
print(result["result"])

print("\n📄 引用来源：")
for i, doc in enumerate(result["source_documents"]):
    source = doc.metadata.get('source', 'unknown')
    page = doc.metadata.get('page', '?')
    print(f"  [{i+1}] {source} 第{page}页")
    print(f"      内容：{doc.page_content[:80]}...")

# ============================================
# 一个完整可运行的RAG问答系统（完整版）
# 文件：rag_qa_system.py
# 使用方法：python rag_qa_system.py
# ============================================

import os
from typing import List
from langchain_community.document_loaders import PyPDFLoader, TextLoader, WebLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.schema import Document

class RAGQASystem:
    """RAG问答系统封装类"""
    
    def __init__(self, 
                 embedding_model: str = "BAAI/bge-large-zh-v1.5",
                 llm_model: str = "gpt-4o-mini",
                 vector_store_path: str = "./chroma_db"):
        self.embedding_model = embedding_model
        self.llm_model = llm_model
        self.vector_store_path = vector_store_path
        self.vectorstore = None
        self.qa_chain = None
        self._init_components()
    
    def _init_components(self):
        """初始化各组件"""
        print("🔧 初始化组件...")
        
        # 初始化Embedding模型
        self.embeddings = HuggingFaceBgeEmbeddings(
            model_name=self.embedding_model,
            model_kwargs={'device': 'cpu'},
            encode_kwargs={'normalize_embeddings': True}
        )
        print("✅ Embedding模型加载完成")
        
        # 初始化LLM
        self.llm = ChatOpenAI(
            model=self.llm_model,
            temperature=0.3
        )
        print("✅ 大模型加载完成")
    
    def load_documents(self, file_paths: List[str], file_type: str = "pdf"):
        """加载文档"""
        print(f"\n📂 加载 {len(file_paths)} 个文档...")
        
        all_docs = []
        for path in file_paths:
            if file_type == "pdf":
                loader = PyPDFLoader(path)
            elif file_type == "txt":
                loader = TextLoader(path, encoding="utf-8")
            else:
                raise ValueError(f"不支持的文件类型：{file_type}")
            
            docs = loader.load()
            print(f"  - {path}: {len(docs)} 页/节")
            all_docs.extend(docs)
        
        print(f"✅ 共加载 {len(all_docs)} 个文档")
        return all_docs
    
    def split_documents(self, documents: List[Document], 
                        chunk_size: int = 500, 
                        chunk_overlap: int = 100):
        """分块文档"""
        print(f"\n✂️ 分块处理（chunk_size={chunk_size}, overlap={chunk_overlap}）...")
        
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=["\n\n", "\n", "。", "！", "？"]
        )
        
        chunks = text_splitter.split_documents(documents)
        print(f"✅ 分块完成：{len(chunks)} 个块")
        return chunks
    
    def build_vectorstore(self, chunks: List[Document]):
        """构建向量数据库"""
        print(f"\n🗄️ 构建向量数据库...")
        
        self.vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=self.embeddings,
            persist_directory=self.vector_store_path
        )
        
        print(f"✅ 向量数据库构建完成，共 {self.vectorstore._collection.count()} 条向量")
    
    def build_qa_chain(self, search_k: int = 5):
        """构建问答链"""
        print(f"\n🔗 构建问答链...")
        
        retriever = self.vectorstore.as_retriever(
            search_kwargs={"k": search_k}
        )
        
        self.qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=retriever,
            return_source_documents=True
        )
        
        print("✅ 问答链构建完成")
    
    def query(self, question: str) -> dict:
        """问答"""
        if not self.qa_chain:
            raise RuntimeError("QA链未构建，请先调用 build_qa_chain()")
        
        result = self.qa_chain.invoke({"query": question})
        return {
            "answer": result["result"],
            "sources": [
                {
                    "content": doc.page_content[:100],
                    "source": doc.metadata.get("source", "unknown")
                }
                for doc in result["source_documents"]
            ]
        }
    
    @classmethod
    def from_existing_vectorstore(cls, vector_store_path: str, llm_model: str = "gpt-4o-mini"):
        """从已有向量数据库加载（不需要重新构建）"""
        instance = cls(vector_store_path=vector_store_path, llm_model=llm_model)
        
        # 加载已有向量数据库
        instance.vectorstore = Chroma(
            persist_directory=vector_store_path,
            embedding_function=instance.embeddings
        )
        
        # 重建QA链
        instance.build_qa_chain()
        
        return instance


# ===================== 使用示例 =====================
if __name__ == "__main__":
    # 初始化系统（首次使用需要构建向量库）
    rag = RAGQASystem()
    
    # 第1步：加载文档
    docs = rag.load_documents(["docs/faq.pdf", "docs/manual.pdf"], file_type="pdf")
    
    # 第2步：分块
    chunks = rag.split_documents(docs)
    
    # 第3步：构建向量数据库
    rag.build_vectorstore(chunks)
    
    # 第4步：构建问答链
    rag.build_qa_chain()
    
    # 第5步：提问
    while True:
        question = input("\n❓ 请输入问题（输入q退出）：")
        if question.lower() == 'q':
            break
        
        result = rag.query(question)
        print(f"\n📝 回答：{result['answer']}")
        print(f"📄 来源：{len(result['sources'])}个")
        for s in result['sources']:
            print(f"  - {s['source']}: {s['content']}...")

# ============================================
# 混合检索：向量检索 + 关键词检索
# ============================================
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain_community.vectorstores import Chroma

# 假设你已经有了 chunks（文档块列表）
# 和 vectorstore（Chroma向量数据库）

# 1. 创建向量检索器
vector_retriever = vectorstore.as_retriever(
    search_kwargs={"k": 5}
)

# 2. 创建BM25关键词检索器
# BM25是一种经典的文本检索算法，不依赖向量
bm25_retriever = BM25Retriever.from_texts(
    texts=[chunk.page_content for chunk in chunks],
    metadatas=[chunk.metadata for chunk in chunks],
    k=5  # 返回5个结果
)

# 3. 组合成混合检索器
# weights=[0.5, 0.5] 表示向量检索和BM25各占50%权重
# 可以根据实际效果调整，比如[0.7, 0.3]表示更信任向量检索
ensemble_retriever = EnsembleRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    weights=[0.5, 0.5]
)

# 测试混合检索效果
query = "产品报价是多少？"
results = ensemble_retriever.invoke(query)
print(f"混合检索返回 {len(results)} 个结果")
for i, r in enumerate(results):
    print(f"  [{i+1}] {r.page_content[:60]}...")

# ============================================
# Query改写：让AI帮你优化问题
# ============================================
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)

# Query改写模板
rewrite_prompt = PromptTemplate(
    template="""你是一个专业的搜索顾问。用户可能不太擅长提问，请将用户的问题改写得更清晰、更适合检索知识库。

原始问题：{original_query}

请输出3个改写后的问题（每个一行，直接输出问题，不需要解释）：
1.""",
    input_variables=["original_query"]
)

def rewrite_query(query: str) -> list:
    """将原始问题改写成3个更精准的版本"""
    response = llm.invoke(rewrite_prompt.format(original_query=query))
    # 解析返回的3个问题
    lines = [l.strip() for l in response.content.split('\n') if l.strip()]
    # 过滤非问题的行
    questions = [l for l in lines if l and l[0].isdigit()]
    return questions

# 示例
original = "你们那个软件怎么卖"
rewritten = rewrite_query(original)
print("原始问题：", original)
print("改写后：")
for q in rewritten:
    print(f"  - {q}")
# 输出类似：
#   - 软件的定价方案是什么？
#   - 软件的价格和收费模式
#   - 产品购买渠道和费用标准

# ============================================
# Context压缩：只保留与问题最相关的部分
# ============================================
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_community.vectorstores import Chroma
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini")

# 创建Context压缩器
# 从每个检索到的文档中提取与问题最相关的句子
compressor = LLMChainExtractor.from_llm(llm)

# 创建自压缩向量检索器
# 工作原理：先检索N个块，然后对每个块压缩，只保留与问题相关的内容
self_retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 10}  # 先检索10个块
)

from langchain.retrievers import ContextualCompressionRetriever

compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=self_retriever
)

# 测试压缩效果
query = "产品的退款政策是什么？"
compressed_docs = compression_retriever.invoke(query)

print(f"原始检索10个块 → 压缩后{len(compressed_docs)}个块")
print("压缩后内容示例：")
for doc in compressed_docs[:2]:
    print(f"  - {doc.page_content[:100]}...")

# ============================================
# 使用Cohere进行重排序（效果最好）
# ============================================
import os

# 需要安装 cohere 包
# pip install cohere
import cohere

co = cohere.Client(os.getenv("COHERE_API_KEY"))

def rerank_documents(query: str, documents: list, top_n: int = 3):
    """
    使用Cohere Rerank API对文档重新排序
    
    参数:
        query: 用户问题
        documents: 检索到的文档列表（字符串列表）
        top_n: 返回前N个最相关的文档
    
    返回:
        重排序后的文档列表（按相关性从高到低）
    """
    results = co.rerank(
        query=query,
        documents=documents,
        top_n=top_n,
        model="rerank-multilingual-v2.0"  # 支持多语言
    )
    
    # results 是 RerankResult 对象列表
    reranked = []
    for r in results:
        reranked.append({
            "index": r.index,
            "score": r.relevance_score,
            "document": documents[r.index]
        })
    
    return reranked

# 使用示例
vector_results = vectorstore.similarity_search("产品的支持服务有哪些", k=10)
doc_texts = [doc.page_content for doc in vector_results]

reranked = rerank_documents(
    query="产品的支持服务有哪些",
    documents=doc_texts,
    top_n=3
)

print("重排序后的Top3结果：")
for i, r in enumerate(reranked):
    print(f"  [{i+1}] 相关性分数={r['score']:.4f}")
    print(f"      内容：{r['document'][:60]}...")

# ============================================
# docker-compose.yml - RAG系统生产部署
# ============================================
version: '3.8'

services:
  # FastAPI后端服务
  rag-api:
    build:
      context: ./rag_api
      dockerfile: Dockerfile
    container_name: rag-api
    ports:
      - "8000:8000"
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - COHERE_API_KEY=${COHERE_API_KEY}
      - CHROMA_PATH=/app/chroma_db
      - LOG_LEVEL=INFO
    volumes:
      # 挂载向量数据库（持久化存储）
      - ./chroma_data:/app/chroma_db
      # 挂载知识库文档（可以随时更新）
      - ./documents:/app/documents:ro
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
      interval: 30s
      timeout: 10s
      retries: 3
    restart: unless-stopped

  # Nginx反向代理（可选，生产环境建议加）
  nginx:
    image: nginx:alpine
    container_name: rag-nginx
    ports:
      - "80:80"
      - "443:443"
    volumes:
      - ./nginx.conf:/etc/nginx/nginx.conf:ro
      - ./ssl:/etc/nginx/ssl:ro
    depends_on:
      - rag-api
    restart: unless-stopped

# ============================================
# FastAPI RAG服务入口
# 文件：rag_api/main.py
# ============================================
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import os
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="RAG问答API", version="1.0.0")

# 延迟导入，避免启动时就加载大模型
rag_system = None

class QuestionRequest(BaseModel):
    """问答请求"""
    question: str
    top_k: Optional[int] = 5
    use_rerank: Optional[bool] = True

class SourceDocument(BaseModel):
    """来源文档"""
    content: str
    source: str
    score: Optional[float] = None

class AnswerResponse(BaseModel):
    """回答响应"""
    answer: str
    sources: List[SourceDocument]
    latency_ms: float

@app.on_event("startup")
async def startup_event():
    """启动时初始化RAG系统"""
    global rag_system
    
    logger.info("🚀 启动RAG系统...")
    
    # 检查环境变量
    if not os.getenv("OPENAI_API_KEY"):
        raise RuntimeError("OPENAI_API_KEY 环境变量未设置")
    
    # 导入RAG系统
    from rag_qa import RAGQASystem
    
    # 尝试加载已有的向量数据库，如果没有则构建
    chroma_path = os.getenv("CHROMA_PATH", "./chroma_db")
    if os.path.exists(chroma_path):
        logger.info(f"📂 从已有向量数据库加载：{chroma_path}")
        rag_system = RAGQASystem.from_existing_vectorstore(chroma_path)
    else:
        logger.info("🆕 构建新的向量数据库...")
        rag_system = RAGQASystem()
        docs = rag_system.load_documents(
            file_paths=["./documents/faq.pdf"],
            file_type="pdf"
        )
        chunks = rag_system.split_documents(docs)
        rag_system.build_vectorstore(chunks)
        rag_system.build_qa_chain()
    
    logger.info("✅ RAG系统就绪！")

@app.get("/health")
async def health_check():
    """健康检查接口"""
    return {"status": "healthy", "rag_ready": rag_system is not None}

@app.post("/api/ask", response_model=AnswerResponse)
async def ask_question(request: QuestionRequest):
    """问答接口"""
    import time
    
    if rag_system is None:
        raise HTTPException(status_code=503, detail="RAG系统未初始化")
    
    start_time = time.time()
    
    try:
        result = rag_system.query(request.question)
        
        latency_ms = (time.time() - start_time) * 1000
        
        return AnswerResponse(
            answer=result["answer"],
            sources=[
                SourceDocument(
                    content=s["content"],
                    source=s["source"],
                    score=s.get("score")
                )
                for s in result["sources"]
            ],
            latency_ms=round(latency_ms, 2)
        )
    except Exception as e:
        logger.error(f"问答请求失败：{str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/documents")
async def list_documents():
    """列出已索引的文档"""
    if rag_system is None:
        raise HTTPException(status_code=503, detail="RAG系统未初始化")
    
    # 从向量数据库获取统计信息
    count = rag_system.vectorstore._collection.count()
    return {
        "total_vectors": count,
        "status": "indexed"
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

# nginx.conf - 生产环境Nginx配置
worker_processes auto;
error_log /var/log/nginx/error.log warn;

events {
    worker_connections 1024;
}

http {
    include /etc/nginx/mime.types;
    default_type application/octet-stream;
    
    # 日志格式
    log_format main '$remote_addr - $remote_user [$time_local] "$request" '
                    '$status $body_bytes_sent "$http_referer" '
                    '"$http_user_agent" "$http_x_forwarded_for" '
                    'rt=$request_time';
    
    access_log /var/log/nginx/access.log main;
    
    # 性能优化
    sendfile on;
    tcp_nopush on;
    tcp_nodelay on;
    keepalive_timeout 65;
    types_hash_max_size 2048;
    
    # Gzip压缩（减少传输量，API响应通常为JSON）
    gzip on;
    gzip_vary on;
    gzip_min_length 1000;
    gzip_types application/json text/plain text/css;
    gzip_proxied any;
    
    upstream rag_api {
        server rag-api:8000;
        keepalive 32;  # 保持连接，避免频繁建连
    }
    
    server {
        listen 80;
        server_name your-domain.com;  # 替换成你的域名
        
        # 安全头
        add_header X-Frame-Options "SAMEORIGIN" always;
        add_header X-Content-Type-Options "nosniff" always;
        add_header X-XSS-Protection "1; mode=block" always;
        
        # 限流（保护后端API）
        limit_req_zone $binary_remote_addr zone=api_limit:10m rate=10r/s;
        
        location / {
            limit_req zone=api_limit burst=20 nodelay;
            
            proxy_pass http://rag_api;
            proxy_http_version 1.1;
            proxy_set_header Host $host;
            proxy_set_header X-Real-IP $remote_addr;
            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
            proxy_set_header X-Forwarded-Proto $scheme;
            
            # 超时配置（AI生成可能需要较长时间）
            proxy_read_timeout 120s;
            proxy_connect_timeout 10s;
            proxy_send_timeout 120s;
            
            # 禁用缓存
            proxy_cache off;
        }
        
        # 监控接口
        location /health {
            proxy_pass http://rag_api;
            proxy_http_version 1.1;
            access_log off;
        }
    }
}

#!/bin/bash
# deploy_rag.sh - 一键部署RAG系统

set -e  # 遇到错误立即退出

echo "=========================================="
echo "🚀 RAG系统部署脚本"
echo "=========================================="

# 检查环境变量
if [ -z "$OPENAI_API_KEY" ]; then
    echo "❌ 错误：请设置 OPENAI_API_KEY 环境变量"
    echo "   export OPENAI_API_KEY=sk-xxxxx"
    exit 1
fi

# 创建必要的目录
mkdir -p documents chroma_data ssl logs

echo "📂 目录结构："
ls -la

# 构建Docker镜像
echo ""
echo "🔨 构建Docker镜像..."
docker build -t rag-api:latest ./rag_api

# 启动服务
echo ""
echo "🐳 启动Docker服务..."
docker-compose up -d

# 等待服务就绪
echo ""
echo "⏳ 等待服务启动..."
sleep 10

# 检查健康状态
echo ""
echo "🏥 检查服务健康状态..."
curl -s http://localhost:8000/health | python3 -m json.tool

# 查看运行中的容器
echo ""
echo "📋 运行中的容器："
docker ps --format "table {{.Names}}\t{{.Status}}\t{{.Ports}}"

echo ""
echo "=========================================="
echo "✅ 部署完成！"
echo "   API地址：http://localhost:8000"
echo "   文档： curl -X POST http://localhost:8000/api/ask "
echo "          -H 'Content-Type: application/json' "
echo "          -d '{\"question\":\"你的问题\"}'"
echo "=========================================="