pip install openai chromadb sentence-transformers numpy

# mini_rag.py
import numpy as np
from openai import OpenAI
client = OpenAI()
# ========== 知識庫 ==========
documents = [
    "公司的年假政策：入職滿1年有5天年假，滿3年有10天年假，滿5年有15天年假。",
    "報銷流程：填寫報銷單 → 部門主管審批 → 財務(wù)審核 → 打款，一般3-5個工作日完成。",
    "上班時間為周一至周五 9:00-18:00，午休時間 12:00-13:30。",
    "遠程辦公政策：每周可申請最多2天遠程辦公，需提前一天在OA系統(tǒng)申請。",
    "試用期薪資為正式薪資的80%，試用期一般為3個月，表現(xiàn)優(yōu)秀可提前轉(zhuǎn)正。",
    "公司提供五險一金，公積金繳納比例為12%，個人和公司各承擔一半。",
    "年終獎發(fā)放規(guī)則：入職滿一年的員工可獲得1-3個月薪資的年終獎。",
    "加班政策：工作日加班按1.5倍計算，周末加班按2倍計算，法定節(jié)假日按3倍計算。",
]
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):
    """計算余弦相似度"""
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
def rag_query(question: str, top_k: int = 2) -> str:
    """最簡 RAG 查詢"""
    # 1. 檢索：找最相關(guān)的文檔
    doc_vectors = simple_embed(documents)
    question_vector = simple_embed([question])[0]
    # 計算相似度并排序
    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ù)以下參考文檔回答用戶問題。如果文檔中沒有相關(guān)信息，請說明。\n\n參考文檔：\n{context}"},
            {"role": "user", "content": question}
        ]
    )
    return response.choices[0].message.content
# 測試
if __name__ == "__main__":
    questions = [
        "我入職半年了，能休多少天年假？",
        "我想報銷一筆費用，流程是什么？",
        "加班工資怎么算？",
        "公司年終獎怎么發(fā)的？",
    ]
    for q in questions:
        print(f"? 問：{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  # 存儲來源、頁碼等元信息
def split_text(
    text: str,
    chunk_size: int = 300,
    chunk_overlap: int = 50,
    separator: str = "\n"
) -> list[Chunk]:
    """
    智能文本切分
    Args:
        text: 原始文本
        chunk_size: 每個塊的最大字符數(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:
        # 如果單個段落就超長，按句子再切
        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 的知識庫"""
    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):
        """添加文檔塊到知識庫"""
        if not self.embed_fn:
            raise ValueError("請先設(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)} 個文檔塊，知識庫總量：{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:
    """使用本地模型生成向量，完全免費"""
    def __init__(self, model_name: str = "shibing624/text2vec-base-chinese"):
        """
        中文文本向量化模型
        Args:
            model_name: 模型名稱，首次使用會自動下載
                       推薦：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 問答系統(tǒng)"""
    def __init__(self):
        # 初始化向量化模型
        self.embedder = LocalEmbedder()
        # 初始化知識庫
        self.kb = KnowledgeBase()
        self.kb.set_embed_function(self.embedder.embed)
    def ingest(self, text: str, source: str = "unknown"):
        """將文檔導入知識庫"""
        chunks = split_text(text, chunk_size=300, chunk_overlap=50)
        # 添加來源信息
        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:
        """
        問答主函數(shù)
        Args:
            question: 用戶問題
            top_k: 檢索的文檔塊數(shù)量
            show_context: 是否顯示檢索到的上下文
        Returns:
            模型回答
        """
        # 1. 檢索相關(guān)文檔
        results = self.kb.search(question, top_k=top_k)
        if not results:
            return "抱歉，知識庫中暫無相關(guān)文檔。"
        # 過濾低相關(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}]（來源：{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"""你是一個精準的知識庫問答助手。請嚴格根據(jù)以下參考資料回答用戶問題。
規(guī)則：
1. 只根據(jù)參考資料中的信息回答，不要編造
2. 如果參考資料中沒有相關(guān)信息，明確說明
3. 引用信息時標注出處（如"根據(jù)參考資料1"）
4. 回答要簡潔準確
參考資料：
{context}"""
                },
                {"role": "user", "content": question}
            ],
            temperature=0.1  # 低溫度，減少"幻覺"
        )
        return response.choices[0].message.content
# ========== 使用示例 ==========
if __name__ == "__main__":
    # 初始化 RAG 系統(tǒng)
    rag = RAGSystem()
    # 導入知識文檔
    company_docs = """
    公司員工手冊 2026版
    第一章 考勤管理
    上班時間為周一至周五 9:00-18:00，午休時間 12:00-13:30。
    遲到15分鐘以內(nèi)扣50元，超過15分鐘按曠工半天處理。
    每月全勤獎200元，需當月無遲到早退記錄。
    第二章 休假制度
    年假：入職滿1年5天，滿3年10天，滿5年15天，滿10年20天。
    病假：每年帶薪病假5天，需提供醫(yī)院證明。
    事假：需提前申請，每年事假不超過10天，超過部分按曠工處理。
    婚假：法定婚假3天，晚婚增加7天。
    產(chǎn)假：女員工產(chǎn)假158天，男員工陪產(chǎn)假15天。
    第三章 薪酬福利
    薪資結(jié)構(gòu)：基本工資 + 績效獎金 + 餐補(500元/月) + 交通補貼(300元/月)。
    五險一金：公積金繳納比例12%，社保按國家標準繳納。
    年終獎：入職滿一年可獲1-3個月薪資，根據(jù)年度績效評定。
    調(diào)薪：每年4月和10月各有一次調(diào)薪窗口。
    第四章 培訓發(fā)展
    新員工入職培訓為期3天，包含公司文化、制度規(guī)范、安全教育。
    每季度有一次內(nèi)部技術(shù)分享會。
    每年可申請最高5000元的外部培訓費用報銷。
    晉升評審每半年一次，分別在1月和7月。
    """.strip()
    rag.ingest(company_docs, source="公司員工手冊2026版")
    # 測試問答
    questions = [
        "我剛?cè)肼?，可以請年假嗎?,
        "遲到會怎樣？",
        "我想?yún)⒓油獠颗嘤?，公司有補貼嗎？",
        "公司什么時候調(diào)薪？",
    ]
    for q in questions:
        print(f"\n{'='*60}")
        print(f"? 問：{q}")
        print(f"{'='*60}")
        answer = rag.query(q, show_context=True)
        print(f"\n?? 答：{answer}")

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

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)
    # 簡單的分數(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"]
    # 按綜合分數(shù)排序
    return sorted(combined.items(), key=lambda x: -x[1])

def evaluate_retrieval(test_cases: list[dict], rag_system) -> dict:
    """
    評估檢索質(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": "遲到怎么扣錢", "expected_keywords": ["15分鐘", "50元"]},
    {"question": "培訓報銷額度", "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. 換一個更好的向量化模型
# BAAI/bge-large-zh-v1.5（中文最佳）

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

def multi_turn_rag(messages: list[dict], rag_system) -> str:
    """支持多輪對話的 RAG"""
    # 用最后一條用戶消息檢索
    last_question = messages[-1]["content"]
    # 如果是追問，結(jié)合上下文改寫問題
    if len(messages) > 1:
        rewrite_prompt = f"根據(jù)對話歷史，將用戶的最新提問改寫為獨立問題：\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)