pip install openai chromadb sentence-transformers numpy

# mini_rag.py
import numpy as np
from openai import OpenAI
client = OpenAI()
# ========== 知識(shí)庫(kù) ==========
documents = [
    "公司的年假政策：入職滿(mǎn)1年有5天年假，滿(mǎn)3年有10天年假，滿(mǎn)5年有15天年假。",
    "報(bào)銷(xiāo)流程：填寫(xiě)報(bào)銷(xiāo)單 → 部門(mén)主管審批 → 財(cái)務(wù)審核 → 打款，一般3-5個(gè)工作日完成。",
    "上班時(shí)間為周一至周五 9:00-18:00，午休時(shí)間 12:00-13:30。",
    "遠(yuǎn)程辦公政策：每周可申請(qǐng)最多2天遠(yuǎn)程辦公，需提前一天在OA系統(tǒng)申請(qǐng)。",
    "試用期薪資為正式薪資的80%，試用期一般為3個(gè)月，表現(xiàn)優(yōu)秀可提前轉(zhuǎn)正。",
    "公司提供五險(xiǎn)一金，公積金繳納比例為12%，個(gè)人和公司各承擔(dān)一半。",
    "年終獎(jiǎng)發(fā)放規(guī)則：入職滿(mǎn)一年的員工可獲得1-3個(gè)月薪資的年終獎(jiǎng)。",
    "加班政策：工作日加班按1.5倍計(jì)算，周末加班按2倍計(jì)算，法定節(jié)假日按3倍計(jì)算。",
]
def simple_embed(texts: list[str]) -> list[list[float]]:
    """使用 OpenAI API 生成文本向量"""
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    return [item.embedding for item in response.data]
def cosine_similarity(a, b):
    """計(jì)算余弦相似度"""
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
def rag_query(question: str, top_k: int = 2) -> str:
    """最簡(jiǎn) RAG 查詢(xún)"""
    # 1. 檢索：找最相關(guān)的文檔
    doc_vectors = simple_embed(documents)
    question_vector = simple_embed([question])[0]
    # 計(jì)算相似度并排序
    scores = [cosine_similarity(question_vector, dv) for dv in doc_vectors]
    top_indices = np.argsort(scores)[-top_k:][::-1]
    # 2. 拼接上下文
    context = "\n".join([f"[文檔{i+1}] {documents[idx]}" for i, idx in enumerate(top_indices)])
    # 3. 生成回答
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": f"根據(jù)以下參考文檔回答用戶(hù)問(wèn)題。如果文檔中沒(méi)有相關(guān)信息，請(qǐng)說(shuō)明。\n\n參考文檔：\n{context}"},
            {"role": "user", "content": question}
        ]
    )
    return response.choices[0].message.content
# 測(cè)試
if __name__ == "__main__":
    questions = [
        "我入職半年了，能休多少天年假？",
        "我想報(bào)銷(xiāo)一筆費(fèi)用，流程是什么？",
        "加班工資怎么算？",
        "公司年終獎(jiǎng)怎么發(fā)的？",
    ]
    for q in questions:
        print(f"? 問(wèn)：{q}")
        answer = rag_query(q)
        print(f"?? 答：{answer}")
        print("-" * 60)

# chunker.py
import re
from dataclasses import dataclass
@dataclass
class Chunk:
    """文檔塊"""
    content: str
    metadata: dict  # 存儲(chǔ)來(lái)源、頁(yè)碼等元信息
def split_text(
    text: str,
    chunk_size: int = 300,
    chunk_overlap: int = 50,
    separator: str = "\n"
) -> list[Chunk]:
    """
    智能文本切分
    Args:
        text: 原始文本
        chunk_size: 每個(gè)塊的最大字符數(shù)
        chunk_overlap: 相鄰塊的重疊字符數(shù)
        separator: 切分分隔符
    Returns:
        切分后的文檔塊列表
    """
    # 按分隔符先切分
    segments = text.split(separator)
    segments = [s.strip() for s in segments if s.strip()]
    chunks = []
    current_chunk = ""
    for segment in segments:
        # 如果單個(gè)段落就超長(zhǎng)，按句子再切
        if len(segment) > chunk_size:
            sentences = re.split(r'[。??？；\.\!\?;]', segment)
            sentences = [s.strip() for s in sentences if s.strip()]
        else:
            sentences = [segment]
        for sentence in sentences:
            if len(current_chunk) + len(sentence) + 1 > chunk_size:
                if current_chunk:
                    chunks.append(Chunk(
                        content=current_chunk,
                        metadata={"char_count": len(current_chunk)}
                    ))
                # 保留重疊部分
                overlap_text = current_chunk[-chunk_overlap:] if chunk_overlap > 0 else ""
                current_chunk = overlap_text + sentence
            else:
                current_chunk = current_chunk + separator + sentence if current_chunk else sentence
    if current_chunk:
        chunks.append(Chunk(
            content=current_chunk,
            metadata={"char_count": len(current_chunk)}
        ))
    return chunks

# vector_store.py
import chromadb
from chromadb.config import Settings
class KnowledgeBase:
    """基于 ChromaDB 的知識(shí)庫(kù)"""
    def __init__(self, collection_name: str = "my_knowledge"):
        self.client = chromadb.PersistentClient(path="./chroma_db")
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
        )
        self.embed_fn = None
    def set_embed_function(self, embed_fn):
        """設(shè)置向量化函數(shù)"""
        self.embed_fn = embed_fn
    def add_documents(self, chunks: list):
        """添加文檔塊到知識(shí)庫(kù)"""
        if not self.embed_fn:
            raise ValueError("請(qǐng)先設(shè)置向量化函數(shù)")
        contents = [chunk.content for chunk in chunks]
        embeddings = self.embed_fn(contents)
        ids = [f"doc_{i}" for i in range(self.collection.count(), self.collection.count() + len(chunks))]
        metadatas = [chunk.metadata for chunk in chunks]
        self.collection.add(
            ids=ids,
            documents=contents,
            embeddings=embeddings,
            metadatas=metadatas
        )
        print(f"? 已添加 {len(chunks)} 個(gè)文檔塊，知識(shí)庫(kù)總量：{self.collection.count()}")
    def search(self, query: str, top_k: int = 3) -> list[dict]:
        """檢索最相關(guān)的文檔塊"""
        query_embedding = self.embed_fn([query])[0]
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            include=["documents", "metadatas", "distances"]
        )
        return [
            {
                "content": doc,
                "metadata": meta,
                "score": 1 - dist  # 距離轉(zhuǎn)相似度
            }
            for doc, meta, dist in zip(
                results["documents"][0],
                results["metadatas"][0],
                results["distances"][0]
            )
        ]

# embedder.py
from sentence_transformers import SentenceTransformer
class LocalEmbedder:
    """使用本地模型生成向量，完全免費(fèi)"""
    def __init__(self, model_name: str = "shibing624/text2vec-base-chinese"):
        """
        中文文本向量化模型
        Args:
            model_name: 模型名稱(chēng)，首次使用會(huì)自動(dòng)下載
                       推薦：shibing624/text2vec-base-chinese（中文，1024維）
                       備選：BAAI/bge-small-zh-v1.5（中文，512維，更快）
        """
        print(f"正在加載向量化模型：{model_name}...")
        self.model = SentenceTransformer(model_name)
        print("模型加載完成！")
    def embed(self, texts: list[str]) -> list[list[float]]:
        """批量生成文本向量"""
        embeddings = self.model.encode(texts, show_progress_bar=False)
        return embeddings.tolist()

# rag_system.py
import json
from openai import OpenAI
from chunker import Chunk, split_text
from vector_store import KnowledgeBase
from embedder import LocalEmbedder
client = OpenAI()
class RAGSystem:
    """完整的 RAG 問(wèn)答系統(tǒng)"""
    def __init__(self):
        # 初始化向量化模型
        self.embedder = LocalEmbedder()
        # 初始化知識(shí)庫(kù)
        self.kb = KnowledgeBase()
        self.kb.set_embed_function(self.embedder.embed)
    def ingest(self, text: str, source: str = "unknown"):
        """將文檔導(dǎo)入知識(shí)庫(kù)"""
        chunks = split_text(text, chunk_size=300, chunk_overlap=50)
        # 添加來(lái)源信息
        for chunk in chunks:
            chunk.metadata["source"] = source
        self.kb.add_documents(chunks)
    def query(self, question: str, top_k: int = 3, show_context: bool = False) -> str:
        """
        問(wèn)答主函數(shù)
        Args:
            question: 用戶(hù)問(wèn)題
            top_k: 檢索的文檔塊數(shù)量
            show_context: 是否顯示檢索到的上下文
        Returns:
            模型回答
        """
        # 1. 檢索相關(guān)文檔
        results = self.kb.search(question, top_k=top_k)
        if not results:
            return "抱歉，知識(shí)庫(kù)中暫無(wú)相關(guān)文檔。"
        # 過(guò)濾低相關(guān)度結(jié)果
        results = [r for r in results if r["score"] > 0.3]
        if show_context:
            print("\n?? 檢索到的相關(guān)文檔：")
            for i, r in enumerate(results):
                print(f"  [{i+1}] (相似度: {r['score']:.3f}) {r['content'][:100]}...")
        # 2. 拼接上下文
        context_parts = []
        for i, r in enumerate(results):
            context_parts.append(f"[參考資料{i+1}]（來(lái)源：{r['metadata'].get('source', '未知')}）\n{r['content']}")
        context = "\n\n".join(context_parts)
        # 3. 調(diào)用大模型生成回答
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": f"""你是一個(gè)精準(zhǔn)的知識(shí)庫(kù)問(wèn)答助手。請(qǐng)嚴(yán)格根據(jù)以下參考資料回答用戶(hù)問(wèn)題。
規(guī)則：
1. 只根據(jù)參考資料中的信息回答，不要編造
2. 如果參考資料中沒(méi)有相關(guān)信息，明確說(shuō)明
3. 引用信息時(shí)標(biāo)注出處（如"根據(jù)參考資料1"）
4. 回答要簡(jiǎn)潔準(zhǔn)確
參考資料：
{context}"""
                },
                {"role": "user", "content": question}
            ],
            temperature=0.1  # 低溫度，減少"幻覺(jué)"
        )
        return response.choices[0].message.content
# ========== 使用示例 ==========
if __name__ == "__main__":
    # 初始化 RAG 系統(tǒng)
    rag = RAGSystem()
    # 導(dǎo)入知識(shí)文檔
    company_docs = """
    公司員工手冊(cè) 2026版
    第一章 考勤管理
    上班時(shí)間為周一至周五 9:00-18:00，午休時(shí)間 12:00-13:30。
    遲到15分鐘以?xún)?nèi)扣50元，超過(guò)15分鐘按曠工半天處理。
    每月全勤獎(jiǎng)200元，需當(dāng)月無(wú)遲到早退記錄。
    第二章 休假制度
    年假：入職滿(mǎn)1年5天，滿(mǎn)3年10天，滿(mǎn)5年15天，滿(mǎn)10年20天。
    病假：每年帶薪病假5天，需提供醫(yī)院證明。
    事假：需提前申請(qǐng)，每年事假不超過(guò)10天，超過(guò)部分按曠工處理。
    婚假：法定婚假3天，晚婚增加7天。
    產(chǎn)假：女員工產(chǎn)假158天，男員工陪產(chǎn)假15天。
    第三章 薪酬福利
    薪資結(jié)構(gòu)：基本工資 + 績(jī)效獎(jiǎng)金 + 餐補(bǔ)(500元/月) + 交通補(bǔ)貼(300元/月)。
    五險(xiǎn)一金：公積金繳納比例12%，社保按國(guó)家標(biāo)準(zhǔn)繳納。
    年終獎(jiǎng)：入職滿(mǎn)一年可獲1-3個(gè)月薪資，根據(jù)年度績(jī)效評(píng)定。
    調(diào)薪：每年4月和10月各有一次調(diào)薪窗口。
    第四章 培訓(xùn)發(fā)展
    新員工入職培訓(xùn)為期3天，包含公司文化、制度規(guī)范、安全教育。
    每季度有一次內(nèi)部技術(shù)分享會(huì)。
    每年可申請(qǐng)最高5000元的外部培訓(xùn)費(fèi)用報(bào)銷(xiāo)。
    晉升評(píng)審每半年一次，分別在1月和7月。
    """.strip()
    rag.ingest(company_docs, source="公司員工手冊(cè)2026版")
    # 測(cè)試問(wèn)答
    questions = [
        "我剛?cè)肼?，可以?qǐng)年假嗎？",
        "遲到會(huì)怎樣？",
        "我想?yún)⒓油獠颗嘤?xùn)，公司有補(bǔ)貼嗎？",
        "公司什么時(shí)候調(diào)薪？",
    ]
    for q in questions:
        print(f"\n{'='*60}")
        print(f"? 問(wèn)：{q}")
        print(f"{'='*60}")
        answer = rag.query(q, show_context=True)
        print(f"\n?? 答：{answer}")

# 推薦參數(shù)
chunk_size = 300     # 中文場(chǎng)景 200-500 字符
chunk_overlap = 50   # 重疊 10%-20%
top_k = 3            # 檢索 3-5 個(gè)塊

def hybrid_search(query: str, kb, bm25_results: list, alpha: float = 0.7) -> list:
    """
    混合檢索：結(jié)合向量檢索和關(guān)鍵詞檢索
    Args:
        alpha: 向量檢索的權(quán)重（0-1），1.0表示純向量檢索
    """
    vector_results = kb.search(query, top_k=5)
    # 簡(jiǎn)單的分?jǐn)?shù)融合
    combined = {}
    for r in vector_results:
        key = r["content"]
        combined[key] = combined.get(key, 0) + alpha * r["score"]
    for r in bm25_results:
        key = r["content"]
        combined[key] = combined.get(key, 0) + (1 - alpha) * r["score"]
    # 按綜合分?jǐn)?shù)排序
    return sorted(combined.items(), key=lambda x: -x[1])

def evaluate_retrieval(test_cases: list[dict], rag_system) -> dict:
    """
    評(píng)估檢索質(zhì)量
    test_cases 格式：
    [
        {"question": "年假多少天", "expected_keywords": ["5天", "10天", "15天"]},
        ...
    ]
    """
    results = {"total": len(test_cases), "hit": 0, "details": []}
    for case in test_cases:
        search_results = rag_system.kb.search(case["question"], top_k=3)
        retrieved_text = " ".join([r["content"] for r in search_results])
        # 檢查關(guān)鍵詞是否被檢索到
        hit = any(kw in retrieved_text for kw in case["expected_keywords"])
        if hit:
            results["hit"] += 1
        results["details"].append({
            "question": case["question"],
            "hit": hit,
            "expected": case["expected_keywords"],
            "retrieved": retrieved_text[:200]
        })
    results["recall"] = results["hit"] / results["total"]
    return results
# 使用示例
test_data = [
    {"question": "年假多少天", "expected_keywords": ["5天", "10天", "15天"]},
    {"question": "遲到怎么扣錢(qián)", "expected_keywords": ["15分鐘", "50元"]},
    {"question": "培訓(xùn)報(bào)銷(xiāo)額度", "expected_keywords": ["5000"]},
]
eval_result = evaluate_retrieval(test_data, rag)
print(f"檢索召回率：{eval_result['recall']:.1%}")

# 1. 調(diào)整 chunk_size
chunk_size = 200  # 試試更小的塊
# 2. 增加 top_k
top_k = 5  # 檢索更多候選
# 3. 換一個(gè)更好的向量化模型
# BAAI/bge-large-zh-v1.5（中文最佳）

# 1. 降低 temperature
temperature = 0.0
# 2. 強(qiáng)化系統(tǒng)提示詞
system_prompt = """嚴(yán)格根據(jù)參考資料回答。如果參考資料中沒(méi)有相關(guān)信息，
必須回答"根據(jù)現(xiàn)有資料，我無(wú)法回答這個(gè)問(wèn)題"，絕不可編造。"""
# 3. 添加來(lái)源引用要求
system_prompt += "\n每個(gè)回答必須標(biāo)注參考資料的編號(hào)。"

def multi_turn_rag(messages: list[dict], rag_system) -> str:
    """支持多輪對(duì)話(huà)的 RAG"""
    # 用最后一條用戶(hù)消息檢索
    last_question = messages[-1]["content"]
    # 如果是追問(wèn)，結(jié)合上下文改寫(xiě)問(wèn)題
    if len(messages) > 1:
        rewrite_prompt = f"根據(jù)對(duì)話(huà)歷史，將用戶(hù)的最新提問(wèn)改寫(xiě)為獨(dú)立問(wèn)題：\n"
        for msg in messages:
            rewrite_prompt += f"{msg['role']}: {msg['content']}\n"
        rewrite = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": rewrite_prompt}]
        )
        search_query = rewrite.choices[0].message.content
    else:
        search_query = last_question
    return rag_system.query(search_query)