利用RAG低成本构建一个专业AI助手

# 固定长度分块
def chunk_by_size(text, chunk_size=500, overlap=50):
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)
        start = end - overlap  # 重叠部分避免语义割裂
    return chunks

# 按语义单元分块（更智能）
def chunk_by_semantic(text):
    # 按段落、章节、句子等自然边界切分
    paragraphs = text.split('\n\n')
    return [p for p in paragraphs if len(p.strip()) > 50]

from sentence_transformers import SentenceTransformer

# 使用开源的Embedding模型
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

chunks = [
    "糖尿病是一种代谢性疾病，特征是血糖水平持续升高",
    "高血压是指动脉血压持续高于正常范围",
    "Python是一种高级编程语言"
]

embeddings = model.encode(chunks)
print(embeddings.shape)  # (3, 384)

import chromadb

# 初始化Chroma客户端
client = chromadb.Client()
collection = client.create_collection(name="knowledge_base")

# 添加文档
collection.add(
    documents=chunks,
    embeddings=embeddings.tolist(),
    metadatas=[{"source": "medical_handbook.pdf", "page": i} for i in range(len(chunks))],
    ids=[f"doc_{i}" for i in range(len(chunks))]
)

def retrieve(query, top_k=3):
    # 问题向量化
    query_embedding = model.encode([query])[0]

    # 检索最相似的文档
    results = collection.query(
        query_embeddings=[query_embedding.tolist()],
        n_results=top_k
    )

    return results['documents'][0], results['metadatas'][0]

def generate_answer(query, contexts):
    prompt = f"""你是一个专业的医疗助手。请基于以下参考资料回答用户的问题。

参考资料：
{chr(10).join(f"[{i+1}] {ctx}" for i, ctx in enumerate(contexts))}

用户问题：{query}

回答要求：
1. 仅基于参考资料回答，不要编造信息
2. 如果参考资料不足以回答问题，请明确说明
3. 引用参考资料时标注来源编号

回答："""

    response = llm.generate(prompt)
    return response

def hybrid_search(query, top_k=5):
    # 向量检索
    vector_results = vector_search(query, top_k * 2)

    # BM25关键词检索
    bm25_results = bm25_search(query, top_k * 2)

    # 融合排序（RRF算法）
    merged = reciprocal_rank_fusion([vector_results, bm25_results])
    return merged[:top_k]

from sentence_transformers import CrossEncoder

reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

def rerank(query, candidates):
    pairs = [(query, doc) for doc in candidates]
    scores = reranker.predict(pairs)

    # 按分数排序并过滤低分文档
    ranked = [(doc, score) for doc, score in zip(candidates, scores) if score > 0.5]
    ranked.sort(key=lambda x: x[1], reverse=True)
    return [doc for doc, _ in ranked]

def compress_documents(query, docs):
    compressed = []
    for doc in docs:
        prompt = f"提取以下文档中与问题「{query}」相关的核心信息：\n{doc}"
        summary = llm.generate(prompt, max_tokens=200)
        compressed.append(summary)
    return compressed

def rag_fusion(original_query):
    # 1. 生成多个相关查询
    expansion_prompt = f"针对问题「{original_query}」，生成3个不同角度的相关查询："
    related_queries = llm.generate(expansion_prompt).split('\n')

    # 2. 每个查询独立检索
    all_results = []
    for query in related_queries:
        results = retrieve(query, top_k=5)
        all_results.append(results)

    # 3. 融合排序
    final_docs = reciprocal_rank_fusion(all_results)
    return final_docs[:5]

def verify_answer(answer, sources):
    sentences = answer.split('。')
    verified = []

    for sent in sentences:
        prompt = f"句子「{sent}」是否被以下任一资料支持？\n{sources}\n只回答：是/否"
        is_supported = llm.generate(prompt).strip()

        if is_supported == "是":
            verified.append(sent)
        else:
            verified.append(f"[未验证] {sent}")

    return '。'.join(verified)

import chromadb
from sentence_transformers import SentenceTransformer, CrossEncoder
from openai import OpenAI

class RAGSystem:
    def __init__(self):
        self.embedder = SentenceTransformer('moka-ai/m3e-base')
        self.reranker = CrossEncoder('BAAI/bge-reranker-large')
        self.llm = OpenAI(api_key="your-key")

        self.chroma_client = chromadb.Client()
        self.collection = self.chroma_client.get_or_create_collection(
            name="knowledge_base",
            metadata={"hnsw:space": "cosine"}
        )

    def index_documents(self, documents, metadatas):
        """索引文档到向量数据库"""
        chunks = []
        chunk_metas = []

        for doc, meta in zip(documents, metadatas):
            # 智能分块
            doc_chunks = self._chunk_document(doc)
            chunks.extend(doc_chunks)
            chunk_metas.extend([{**meta, 'chunk_id': i} for i in range(len(doc_chunks))])

        # 批量向量化
        embeddings = self.embedder.encode(chunks, show_progress_bar=True)

        # 存储
        self.collection.add(
            documents=chunks,
            embeddings=embeddings.tolist(),
            metadatas=chunk_metas,
            ids=[f"doc_{i}" for i in range(len(chunks))]
        )

        print(f"已索引 {len(chunks)} 个文档块")

    def _chunk_document(self, text, chunk_size=800, overlap=100):
        """智能文档分块"""
        # 优先按段落分
        paragraphs = text.split('\n\n')
        chunks = []
        current_chunk = ""

        for para in paragraphs:
            if len(current_chunk) + len(para) < chunk_size:
                current_chunk += para + '\n\n'
            else:
                if current_chunk:
                    chunks.append(current_chunk.strip())
                current_chunk = para + '\n\n'

        if current_chunk:
            chunks.append(current_chunk.strip())

        return chunks

    def retrieve(self, query, top_k=5):
        """检索相关文档"""
        # 1. 向量检索（召回更多候选）
        query_embedding = self.embedder.encode([query])[0]
        results = self.collection.query(
            query_embeddings=[query_embedding.tolist()],
            n_results=top_k * 3  # 召回3倍候选
        )

        candidates = results['documents'][0]
        metadatas = results['metadatas'][0]

        # 2. Rerank精排
        pairs = [[query, doc] for doc in candidates]
        scores = self.reranker.predict(pairs)

        # 3. 按分数排序并过滤
        ranked = sorted(
            zip(candidates, scores, metadatas),
            key=lambda x: x[1],
            reverse=True
        )

        # 只保留分数>阈值的结果
        filtered = [(doc, meta) for doc, score, meta in ranked if score > 0.3]

        return filtered[:top_k]

    def generate_answer(self, query, contexts):
        """基于检索结果生成答案"""
        context_str = "\n\n".join([
            f"[文档{i+1}] {doc}\n来源：{meta.get('source', '未知')}"
            for i, (doc, meta) in enumerate(contexts)
        ])

        prompt = f"""你是一个专业的知识助手。请基于以下参考资料回答用户的问题。

参考资料：
{context_str}

用户问题：{query}

回答要求：
1. 仅基于参考资料回答，引用时标注文档编号如[1]
2. 如果资料不足，明确说明"根据现有资料无法完整回答"
3. 保持专业、客观、准确

回答："""

        response = self.llm.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3
        )

        return response.choices[0].message.content

    def query(self, question, top_k=3):
        """完整的RAG查询流程"""
        print(f"问题：{question}\n")

        # 1. 检索
        contexts = self.retrieve(question, top_k)
        if not contexts:
            return "抱歉，没有找到相关资料。"

        print(f"检索到 {len(contexts)} 个相关文档片段\n")

        # 2. 生成答案
        answer = self.generate_answer(question, contexts)

        # 3. 返回结果（包含来源）
        sources = [meta.get('source', '未知') for _, meta in contexts]

        return {
            'answer': answer,
            'sources': list(set(sources)),
            'context_count': len(contexts)
        }

# 使用示例
if __name__ == "__main__":
    rag = RAGSystem()

    # 索引文档
    documents = [
        "糖尿病（Diabetes）是一种代谢性疾病，特征是血糖水平持续升高...",
        "高血压的定义是收缩压≥140 mmHg或舒张压≥90 mmHg...",
        # 更多文档...
    ]

    metadatas = [
        {"source": "医学手册.pdf", "category": "内分泌"},
        {"source": "医学手册.pdf", "category": "心血管"},
    ]

    rag.index_documents(documents, metadatas)

    # 查询
    result = rag.query("糖尿病有哪些典型症状？")
    print(f"回答：{result['answer']}\n")
    print(f"来源：{', '.join(result['sources'])}")