create_vector_db.py

import os
from dotenv import load_dotenv
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_google_genai import GoogleGenerativeAIEmbeddings # 改用 Google
from langchain_chroma import Chroma

# 載入環境變數
load_dotenv()

DATA_PATH = "knowledge_base"
DB_PATH = "chroma_db"

def create_vector_db():
    print(f"🔄 開始讀取 {DATA_PATH} 中的資料...")
    
    loader = DirectoryLoader(DATA_PATH, glob="*.md", loader_cls=TextLoader,loader_kwargs={"encoding": "utf-8"})
    documents = loader.load()
    print(f"✅ 成功讀取 {len(documents)} 份文件。")

    if not documents:
        print("❌ 錯誤：找不到文件。")
        return

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=100,
        separators=["\n\n", "\n", " ", ""]
    )
    chunks = text_splitter.split_documents(documents)
    print(f"✂️  文件已切分為 {len(chunks)} 個區塊。")

    print("💾 正在建立向量資料庫 (ChromaDB - Google Embeddings)...")
    
    # 修改：使用 Google 的 Embedding 模型
    embedding_function = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
    
    vector_db = Chroma.from_documents(
        documents=chunks,
        embedding=embedding_function,
        persist_directory=DB_PATH
    )
    
    print(f"🎉 向量資料庫已建立並儲存於 '{DB_PATH}'")

if __name__ == "__main__":
    create_vector_db()