前言

昨晚想通了「開卷考試」的邏輯後，今天一早就來實作向量資料庫。

市面上有很多向量資料庫（Pinecone、Qdrant、Milvus 等），但考慮到本地開發方便、免設定雲端 Key，我最後選了 ChromaDB。它可以直接以 SQLite 檔案的形式存在本地資料夾，輕量又直覺，非常適合拿來做我們的陪伴 Agent 外掛記憶庫！

今天我的目標很明確：

1. 安裝並初始化 ChromaDB
2. 建立「角色設定 (Lore)」與「使用者記憶 (User Memory)」雙軌 Collection
3. 編寫簡單的寫入與檢索測試 script

第一步：安裝環境與套件

首先安裝 ChromaDB 以及 Google genai 相關 SDK：

pip install chromadb google-genai

技術說明：ChromaDB 本身主要負責向量資料庫的儲存與檢索。雖然 ChromaDB 有預設的模型，但為了讓繁體中文語意更精準，本篇我們自訂了 GeminiEmbeddingFunction，讓文字一律透過 Gemini 的 Embedding API 轉為向量後再存入 ChromaDB。

第二步：實作向量資料庫模組（rag_memory.py）

把 ChromaDB 的操作與 Gemini Embedding 封裝成一個獨立的模組 rag_memory.py，負責處理角色設定（Lore）與使用者記憶（User Memory）的雙軌寫入與向量檢索。

來寫 rag_memory.py：

import chromadb
from dotenv import load_dotenv
from google import genai
from google.genai import types

load_dotenv()

EMBEDDING_MODEL = "gemini-embedding-001"

class GeminiEmbeddingFunction:
    """用 Gemini 的 embedding API 取代 ChromaDB 的預設模型。

    預設的 DefaultEmbeddingFunction 底層是 all-MiniLM-L6-v2，
    只用英文語料訓練，看中文等於看亂碼，實測問「甜點」時，
    甜點那筆設定的距離(1.26)反而比不相干的設定(1.01)還遠，排序基本上是隨機的。

    task_type 是關鍵：寫入時用 RETRIEVAL_DOCUMENT、查詢時用 RETRIEVAL_QUERY，
    Gemini 會為兩者產生不對稱的向量，檢索準確度比兩邊都用同一種好。
    """

    def __init__(self, task_type: str):
        self._client = genai.Client()
        self._task_type = task_type

    # ChromaDB 用這個名字辨識 embedding function (存進 collection metadata)
    @staticmethod
    def name() -> str:
        return "gemini"

    def __call__(self, input: list[str]) -> list[list[float]]:
        response = self._client.models.embed_content(
            model=EMBEDDING_MODEL,
            contents=input,
            config=types.EmbedContentConfig(task_type=self._task_type),
        )
        return [e.values for e in response.embeddings]

# 在本地建立永久儲存資料夾 (persist)
chroma_client = chromadb.PersistentClient(path="./chroma_db")

# 寫入與查詢不同的 task_type，所以需要兩個 instance
_doc_ef = GeminiEmbeddingFunction("RETRIEVAL_DOCUMENT")
_query_ef = GeminiEmbeddingFunction("RETRIEVAL_QUERY")

# 建立/取得雙軌 Collection
# 軌道 A：角色設定庫 (Lore)
lore_collection = chroma_client.get_or_create_collection(
    name="character_lore", embedding_function=_doc_ef
)

# 軌道 B：使用者長期對話記憶庫 (User Memory)
user_memory_collection = chroma_client.get_or_create_collection(
    name="user_memory", embedding_function=_doc_ef
)

# --- 軌道 A：角色設定相關 API ---
def add_character_lore(lore_id: str, text: str, category: str = "general"):
    """寫入角色的背景設定、喜好或經典事件"""
    lore_collection.upsert(
        documents=[text],
        metadatas=[{"category": category}],
        ids=[lore_id]
    )

def query_character_lore(query_text: str, n_results: int = 2) -> list[str]:
    """根據使用者的對話，搜尋最相關的角色設定"""
    # 自己算 query 向量再傳 query_embeddings，而不是用 query_texts。
    # 因為 query_texts 會走 collection 綁定的 _doc_ef (RETRIEVAL_DOCUMENT)，
    # 這裡要的是 RETRIEVAL_QUERY
    results = lore_collection.query(
        query_embeddings=_query_ef([query_text]),
        n_results=n_results
    )
    # results['documents'] 回傳的是 list of list
    return results['documents'][0] if results['documents'] else []

# --- 軌道 B：使用者記憶相關 API ---
def add_user_memory(memory_id: str, text: str, user_id: str = "default_user"):
    """記錄使用者的偏好、說過的小事"""
    user_memory_collection.upsert(
        documents=[text],
        metadatas=[{"user_id": user_id}],
        ids=[memory_id]
    )

def query_user_memory(query_text: str, n_results: int = 2) -> list[str]:
    """搜尋使用者之前講過的相關私密小事"""
    results = user_memory_collection.query(
        query_embeddings=_query_ef([query_text]),
        n_results=n_results
    )
    return results['documents'][0] if results['documents'] else []

資料權責劃分：
Firestore 保存的是原始聊天紀錄與狀態資料；ChromaDB 的 User Memory 則是將適合語意檢索的長期資訊獨立保存，方便之後用語意相似度找回。

第三步：灌入測試資料並實測搜尋！

模組寫好後，寫一個測試腳本，順便拿莉莉的設定跟私密小事來當測試資料：

# test_rag.py
import rag_memory

# 1. 灌入莉莉的角色設定 (Lore)
rag_memory.add_character_lore("lore_1", "莉莉非常擅長製作各種甜點與手工烘焙，特別是烤薄餅，最喜歡推薦甜點給別人", "food")
rag_memory.add_character_lore("lore_2", "莉莉極度重視家人之間的羈絆，任何傷害她家人的人都會被她視為敵人。", "family")
rag_memory.add_character_lore("lore_3", "莉莉喜歡帥哥類型的偶像，對於少女漫畫風格的情節容易害羞。", "hobby")

# 2. 灌入關於使用者的長期記憶 (User Memory)
rag_memory.add_user_memory("mem_1", "使用者討厭吃香菜，看到香菜會眉頭深鎖。")
rag_memory.add_user_memory("mem_2", "使用者下周二要參加資工系的期末考，最近壓力很大。")

# 3. 測試語意搜尋 (Similarity Search)
print("--- 測試 1：問 Agent 甜點相關話題 ---")
query = "我最近下午茶想吃點甜的，你有推薦嗎？"
related_lores = rag_memory.query_character_lore(query)
print(f"檢索到的角色設定：{related_lores}\n")

print("--- 測試 2：提及考試壓力 ---")
query_user = "我最近唸書唸得好累喔..."
related_memories = rag_memory.query_user_memory(query_user)
print(f"檢索到的使用者記憶：{related_memories}\n")

魔法輸出：

在終端機執行 python test_rag.py：

(venv) PS D:\Character Interaction and Situations Agent> python test_rag.py
--- 測試 1：問 Agent 甜點相關話題 ---
檢索到的角色設定：['莉莉非常擅長製作各種甜點與手工烘焙，特別是烤薄餅，最喜歡推薦甜點給別人',
                 '莉莉喜歡帥哥類型的偶像，對於少女漫畫風格的情節容易害羞。']

--- 測試 2：提及考試壓力 ---
檢索到的使用者記憶：['使用者下周二要參加資工系的期末考，最近壓力很大。',
                 '使用者討厭吃香菜，看到香菜會眉頭深鎖。']

跑出來的測試結果：

• 測試 1 成功檢索出：「莉莉非常擅長做手工烘焙與甜點，特別是鬆餅...」
• 測試 2 成功檢索出：「使用者下周二要參加資工系的期末考，最近壓力很大。」

完全不用精確比對關鍵字，靠著語意（"唸書累" ⇒ "期末考壓力大"），ChromaDB 就自動抓出最準確的記憶片段了！