檢索增強生成 RAG 的介紹與實作

什麼是 RAG？

RAG 全名 Retrieval-Augmented Generation 是一種結合「檢索」和「生成」的 AI 技術。簡單來說，它讓大語言模型（LLM）能夠在回答問題時，先從外部知識庫中找到相關資訊，再基於這些資訊生成更準確的答案。

RAG 解決了什麼問題？

傳統的大語言模型有幾個限制：

1. 知識截止點：模型只知道訓練時的資料，無法獲得最新資訊
2. 幻覺問題：可能生成看似合理但實際錯誤的資訊
3. 領域知識不足：對特定領域的專業知識可能不夠深入
4. 無法存取即時資料：無法查詢資料庫或即時資訊

RAG 技術就是為了解決這些問題而生。

RAG 的運作原理

1. 文件處理與儲存

• 將文件切分成較小的片段（chunks）
• 使用嵌入模型（embedding model）將文字轉換成向量
• 將向量儲存在向量資料庫(Vector Database)

2. 檢索與生成

• 將使用者的查詢轉換成向量
• 在向量資料庫中搜尋最相似的文件片段
• 返回最相關的內容作為上下文
• 將檢索到的上下文與使用者問題結合
• 輸入給大語言模型生成回答
• 確保回答基於提供的上下文資訊

實作 RAG 系統

接下來，我們來實作一個簡單的 RAG 系統。這個範例會使用以下技術棧：

• Python 作為主要語言
• uv 作為套件管理工具
• 使用 ChromaDB 向量資料庫
• 使用 Google Gemini 大語言模型

主程式

# main.py
import argparse
from typing import List

import chromadb
from google import genai
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer, CrossEncoder


def parse_arguments():
    """解析命令列參數"""
    parser = argparse.ArgumentParser(description='RAG 問答系統')
    parser.add_argument(
        'query',
        type=str,
        help='要查詢的問題'
    )
    parser.add_argument(
        '--doc-path',
        type=str,
        default='./story.txt',
        help='文檔路徑 (預設: ./story.txt)'
    )
    return parser.parse_args()


def split_into_chunks(doc_file: str) -> List[str]:
    """資料分塊"""
    try:
        with open(doc_file, 'r', encoding='utf-8') as file:
            content = file.read()
        return [chunk.strip() for chunk in content.split("\n\n") if chunk.strip()]
    except FileNotFoundError:
        print(f"錯誤：找不到檔案 {doc_file}")
        return []
    except Exception as e:
        print(f"讀取檔案時發生錯誤：{e}")
        return []


def embed_chunk(chunk: str, embedding_model: SentenceTransformer) -> List[float]:
    """將文本資料轉換為嵌入向量"""
    embedding = embedding_model.encode(chunk, normalize_embeddings=True)
    return embedding.tolist()


def save_embeddings(chunks: List[str], embeddings: List[List[float]], collection) -> None:
    """將嵌入向量保存到資料庫"""
    for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
        collection.add(
            documents=[chunk],
            embeddings=[embedding],
            ids=[str(i)]
        )


def retrieve(query: str, top_k: int, collection, embedding_model: SentenceTransformer) -> List[str]:
    """檢索(搜尋)相關文檔資料"""
    query_embedding = embed_chunk(query, embedding_model)
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k
    )
    return results['documents'][0]


def rerank(query: str, retrieved_chunks: List[str], top_k: int) -> List[str]:
    """重新排序檢索結果"""
    cross_encoder = CrossEncoder('cross-encoder/mmarco-mMiniLMv2-L12-H384-v1')
    pairs = [(query, chunk) for chunk in retrieved_chunks]
    scores = cross_encoder.predict(pairs)

    scored_chunks = list(zip(retrieved_chunks, scores))
    scored_chunks.sort(key=lambda x: x[1], reverse=True)

    return [chunk for chunk, _ in scored_chunks][:top_k]


def generate(query: str, chunks: List[str], google_client) -> str:
    """使用 LLM 生成回答"""
    prompt = f"""你是一位知識助手，請根據使用者的問題和下列相關片段生成準確的回答。

用戶問題: {query}

相關片段:
{"\n\n".join(chunks)}

請基於上述內容作答，不要編造資訊。"""

    try:
        response = google_client.models.generate_content(
            model="gemini-2.5-flash",
            contents=prompt
        )
        return response.text
    except Exception as e:
        return f"生成回答時發生錯誤：{e}"


def main():
    """主要執行函數"""

    # 解析命令列參數
    args = parse_arguments()
    query = args.query
    doc_path = args.doc_path

    # 載入環境變數
    load_dotenv()

    # 初始化 Google Gemini 客戶端
    google_client = genai.Client()


    # 分割文檔
    chunks = split_into_chunks(doc_path)
    if not chunks:
        print("無法載入文檔，程式結束")
        return

    print(f"成功載入 {doc_path}, 內容已經切分成 {len(chunks)} 個分塊 ...")

    # 初始化嵌入模型
    embedding_model = SentenceTransformer("shibing624/text2vec-base-chinese")

    # 生成嵌入向量
    embeddings = [embed_chunk(chunk, embedding_model) for chunk in chunks]

    # 初始化 ChromaDB
    chromadb_client = chromadb.EphemeralClient()
    chromadb_collection = chromadb_client.get_or_create_collection(name="default")

    # 保存嵌入向量
    save_embeddings(chunks, embeddings, chromadb_collection)

    # 查詢和檢索
    retrieved_chunks = retrieve(query, 5, chromadb_collection, embedding_model)

    # 重新排序
    reranked_chunks = rerank(query, retrieved_chunks, 3)

    # 生成回答
    answer = generate(query, reranked_chunks, google_client)
    print(f"\n問題：\n{query}")
    print(f"\n回答：\n{answer}")


if __name__ == "__main__":
    main()

範例資料

以下是 story.txt 中的範例故事，你也可以改成自己的要測試的內容

《失落的圖書館》

在遙遠的沙漠中，有一座被黃沙掩埋的古老城市。傳說中，這裡曾經是知識的聖地，收藏著一座「失落的圖書館」。

第一章：商旅的傳聞
幾個世紀以來，穿越沙漠的商旅偶爾會聽到老人們的低語：有人在風暴中看見半掩的石門；有人說過在夜晚的月光下，隱約能看見石柱的影子。雖然這些故事從未被證實，但它們像火焰一樣，燃起無數冒險家的好奇心。

第二章：年輕學者的啟程
亞倫是一位對古代文明充滿熱情的學者。他在一份破舊的手稿裡找到了一張模糊的地圖，上面標示著一條與傳說相符的路線。懷著期待與不安，他決定踏上尋找「失落圖書館」的旅程。

第三章：穿越沙漠
沙漠之行並不輕鬆。白天酷熱，夜晚寒冷，亞倫的水源逐漸減少。他曾兩度懷疑自己是不是陷入了幻覺：有一次，他看見遠方出現一列石柱；另一次，他在沙丘後聽到彷彿是書頁翻動的聲音。

第四章：石門的出現
終於，在連續七日的沙塵暴過後，風勢稍歇。亞倫在沙丘頂端，看到一道半掩的石門。門上刻著古老的符號，與他手稿裡的標記相符。那一刻，他幾乎無法相信自己的眼睛。

第五章：進入圖書館
石門背後是一條狹長的石階，通往地下深處。當亞倫點燃火把時，眼前出現了一排排高聳的書架，上面堆滿了塵封的卷軸與羊皮紙。空氣中瀰漫著歷史的氣息。

第六章：知識的代價
然而，圖書館並非單純的寶庫。亞倫發現，這裡的書籍帶有一種奇異的力量：有些內容能讓他獲得前所未有的洞察，但也有些文字會讓他陷入混亂與恐懼。他意識到，這不只是尋寶之旅，而是一場對心智與靈魂的考驗。

第七章：抉擇
在圖書館的最深處，有一本書靜靜地放在石桌上，書脊上刻著「始與終」三個字。亞倫知道，打開它可能會改變他的一生。他站在那裡，火光搖曳，呼吸急促，陷入漫長的猶豫。

接著把專案依賴安裝好後，就可以用以下命令進行查詢

uv run python main.py "亞倫是誰？"

[!NOTE] 本文程式已經放到這個 GitHub Repository 具體專案的環境設定請參考專案中的 README.md

RAG 系統的優化方向

雖然我們這個案例很簡單，但完整了呈現的 RAG 的實現邏輯。之後你可以根據需要慢慢優化：

1. 更好的文件切分策略

• 基於語意的切分而非固定長度
• 保持段落和句子的完整性
• 考慮文件的結構（標題、列表等）

2. 進階檢索技術

• 混合檢索（結合關鍵詞和向量搜尋）
• 重排序（re-ranking）機制
• 查詢擴展和改寫

3. 更好的向量化模型

• 使用專門的嵌入模型（如 OpenAI embeddings）
• 多語言支援
• 領域特化的模型

4. 回答品質優化

• 更精細的提示詞工程
• 事實查核機制
• 引用來源的準確性

實際應用場景

RAG 技術可以應用在許多場景中：

1. 企業知識庫：讓員工快速找到內部文件和資訊
2. 客服系統：基於產品文件自動回答客戶問題
3. 學習平台：根據教材內容回答學生問題
4. 法律諮詢：基於法規文件提供法律建議
5. 技術文件助手：幫助開發者快速查找 API 文件

總結

RAG 技術結合了檢索和生成的優勢，讓 AI 系統能夠提供更加客製化、準確的回答。

希望這篇文章對你理解 RAG 技術有所幫助。若您喜歡，歡迎分享 ~