第一次學 Kotlin Koog AI 就上手 Day 31：RAG：讓 AI 擁有知識庫的超能力

suspend fun main() {

    // 建立 Embedder 和文件儲存系統
    val embedder = LLMEmbedder(OllamaClient(), OllamaEmbeddingModels.NOMIC_EMBED_TEXT)
    val documentEmbedder = JVMTextDocumentEmbedder(embedder)
    val rankedDocumentStorage = EmbeddingBasedDocumentStorage(documentEmbedder, InMemoryVectorStorage())

    // 儲存 FAQ 文件
    // 建立示範文件內容（實際應用中會從檔案系統讀取）
    val faqData = listOf(
        "./faq/shipping.txt" to "商品包裏出貨時間為 1-3 個工作天，超商取貨需要額外 1-2 天。",
        "./faq/returns.txt" to "商品可在收到後 7 天內退換貨，需保持原包裝完整。",
        "./faq/warranty.txt" to "電子產品提供一年保固，保固期內免費維修或更換。",
        "./faq/payment.txt" to "支援信用卡、ATM轉帳、超商付款等多種付款方式。",
        "./faq/account.txt" to "可透過官網註冊會員帳號，享受更多優惠和服務。"
    )

    // 建立目錄並存儲文件
    faqData.forEach { (filename, content) ->
        val path = Path.of(filename)
        Files.createDirectories(path.parent)
        Files.writeString(path, content)
        rankedDocumentStorage.store(path)
        println("已載入：$filename")
    }

    // 搜尋相關文件
    val query = "我的包裏什麼時候會到？"
    val relevantDocs = rankedDocumentStorage.mostRelevantDocuments(
        query = query,
        count = 2,                    // 取前 2 個最相關的文件
        similarityThreshold = 0.5     // 相似度閥值（調整為更寬鬆的閥值）
    )

    // 顯示搜尋結果
    println("用戶問題：$query")
    println("找到 ${relevantDocs.size} 個相關文件：")

    relevantDocs.forEachIndexed { index, doc ->
        println("${index + 1}. 文件：${doc.fileName}")
        try {
            val content = Files.readString(doc)
            println("內容：$content")
        } catch (e: Exception) {
            println("無法讀取檔案：${e.message}")
        }
        println("-".repeat(50))
    }
}

已載入：./faq/shipping.txt
已載入：./faq/returns.txt
已載入：./faq/warranty.txt
已載入：./faq/payment.txt
已載入：./faq/account.txt
用戶問題：我的包裏什麼時候會到？
找到 2 個相關文件：
1. 文件：account.txt
內容：可透過官網註冊會員帳號，享受更多優惠和服務。
--------------------------------------------------
2. 文件：shipping.txt
內容：商品包裏出貨時間為 1-3 個工作天，超商取貨需要額外 1-2 天。
--------------------------------------------------

// 建立 AI 客服工具集
@LLMDescription("AI 客服工具，能夠搜尋公司知識庫回答客戶問題")
class CustomerServiceTools(
    private val documentStorage: EmbeddingBasedDocumentStorage<Path>
) : ToolSet {

    @Tool
    @LLMDescription("搜尋公司知識庫中與客戶問題相關的文件")
    suspend fun searchKnowledgeBase(
        @LLMDescription("客戶的問題或查詢內容")
        query: String,
        @LLMDescription("需要檢索的文件數量，預設為 3")
        count: Int = 3
    ): String {
        // 搜尋相關文件
        println("AI query document -- $query")
        val relevantDocs = documentStorage.mostRelevantDocuments(
            query = query,
            count = count,
            similarityThreshold = 0.5  // 適中的相似度閥值
        ).toList()

        if (relevantDocs.isEmpty()) {
            return "抱歉，在知識庫中找不到相關資訊。建議聯繫人工客服。"
        }

        // 整合搜尋結果
        val sb = StringBuilder("根據知識庫搜尋，找到以下相關資訊：\n\n")

        relevantDocs.forEachIndexed { index, document ->
            try {
                val content = Files.readString(document)
                sb.append("參考資料 ${index + 1}：${document.fileName}\n")
                sb.append("內容：$content\n\n")
            } catch (e: Exception) {
                println("警告：無法讀取文件 ${document.fileName}，原因：${e.message}")
                sb.append("參考資料 ${index + 1}：${document.fileName} (檔案讀取失敗)\n")
            }
        }

        return sb.toString()
    }
}

suspend fun main() {
    // 初始化 RAG 系統
    val embedder = LLMEmbedder(OllamaClient(), OllamaEmbeddingModels.NOMIC_EMBED_TEXT)
    val documentEmbedder = JVMTextDocumentEmbedder(embedder)
    val documentStorage = EmbeddingBasedDocumentStorage(documentEmbedder, InMemoryVectorStorage())

    // 載入知識庫文件（實際應用中可能從資料庫或 API 載入）
    val knowledgeBaseFiles = listOf(
        "faq/shipping.txt",
        "faq/returns.txt",
        "faq/warranty.txt",
        "faq/payment.txt",
        "faq/product-info.txt"
    )

    knowledgeBaseFiles.forEach { fileName ->
        val path = Path.of("./$fileName")
        if (Files.exists(path)) {
            documentStorage.store(path)
            println("已載入：$fileName")
        }
    }

    // 建立工具註冊器
    val toolRegistry = ToolRegistry {
        tools(CustomerServiceTools(documentStorage))
    }

    // 建立 AI 客服 Agent
    val customerServiceAgent = AIAgent(
        executor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")),
        systemPrompt = """
            你是一位專業的客服代表，負責回答客戶問題。

            工作原則：
            1. 使用 searchKnowledgeBase 工具搜尋相關資訊
            2. 基於搜尋結果提供準確、有用的回答
            3. 如果知識庫中沒有相關資訊，誠實告知並建議其他解決方案
            4. 保持友善、專業的語調
            5. 回答要簡潔明瞭，避免冗長
        """.trimIndent(),
        llmModel = OpenAIModels.Chat.GPT4o,
        toolRegistry = toolRegistry
    )

    // 模擬客戶對話
    val customerQueries = listOf(
        "我的訂單什麼時候會到貨？",
        "如何申請退換貨？",
        "產品保固期是多久？",
        "支援哪些付款方式？"
    )

    customerQueries.forEach { query ->
        println("\n" + "=".repeat(50))
        println("客戶問題：$query")
        println("客服回覆：")

        try {
            val response = customerServiceAgent.run(query)
            println(response)
        } catch (e: Exception) {
            println("處理問題時發生錯誤：${e.message}")
        }
    }
}

已載入：faq/AI/shipping.txt
已載入：faq/AI/returns.txt
已載入：faq/AI/warranty.txt
已載入：faq/AI/payment.txt
已載入：faq/AI/product-info.txt

==================================================
客戶問題：我的訂單什麼時候會到貨？
客服回覆：
AI query document -- 訂單 到貨 時間
您好，訂單一般標準運送時間為3-5個工作天（不含週末及國定假日），急件配送則可於1-2個工作天送達，偏遠地區及離島會稍微延長運送時間。出貨後會提供物流追蹤碼，您可透過官方網站查詢最新的運送狀態。若需要更詳細的訂單狀態，歡迎提供訂單編號，我們可以幫您查詢。

==================================================
客戶問題：如何申請退換貨？
客服回覆：
AI query document -- 退換貨申請
您好，申請退換貨的流程如下：

1. 請於商品到貨後7天猶豫期內，聯繫客服人員申請退換貨，並提供訂單編號及退換貨原因。
2. 客服確認後會提供退貨地址及相關注意事項。
3. 請於3個工作天內將商品寄回指定地址，並務必保持商品及包裝完整（商品本體、配件、贈品及發票等）。
4. 商品寄回時請妥善包裝，建議使用有追蹤號碼的宅配服務。
5. 若是換貨，請先確認換購商品有庫存，並支付相應的價差及運費（瑕疵商品換貨運費由公司負擔）。
6. 收到退貨商品後，會在3-5個工作天內完成退款或換貨處理。

備註：
- 已使用或損毀商品、個人衛生用品、食品類、客製化及特價促銷商品可能不接受退換貨。
- 瑕疵商品可隨時申請退換貨，無須受7天猶豫期限制，請立即拍照存證並聯繫客服。

若您需要進一步協助，歡迎隨時聯繫客服。

==================================================
客戶問題：產品保固期是多久？
客服回覆：
AI query document -- 產品保固期
您好，產品的保固期限一般電子產品為1年，大型家電為2年，保固期從購買日（以發票日期為準）開始計算。特殊商品如手機、平板等採製造商原廠保固，期限依品牌規定。進口商品保固期可能不同，建議參考商品頁面說明。如需更詳細的保固範圍、申請流程等資訊，歡迎隨時聯繫我們！

==================================================
客戶問題：支援哪些付款方式？
客服回覆：
AI query document -- 付款方式
目前沒有找到具體的付款方式詳細資訊。建議您可以直接聯繫我們客服人員，或查看官網的付款方式頁面以獲得最新的相關資訊。若有其他問題，也歡迎隨時詢問！

// 嚴格篩選：高品質，但可能漏掉相關內容
val strictResults = documentStorage.mostRelevantDocuments(
    query = "退貨流程",
    count = 3,
    similarityThreshold = 0.85  // 非常高的閾值
)

// 平衡篩選：大多數情況下的最佳選擇
val balancedResults = documentStorage.mostRelevantDocuments(
    query = "退貨流程",
    count = 3,
    similarityThreshold = 0.7   // 適中閾值
)

// 寬鬆篩選：可能包含噪音，但不會漏掉相關內容
val relaxedResults = documentStorage.mostRelevantDocuments(
    query = "退貨流程",
    count = 5,
    similarityThreshold = 0.5   // 低閾值
)

// 實用技巧：動態調整閾值
fun adaptiveSearch(query: String, storage: EmbeddingBasedDocumentStorage<Path>): List<Path> {
    // 先用高閾值搜尋
    var results = storage.mostRelevantDocuments(query, count = 3, similarityThreshold = 0.8).toList()

    // 如果結果太少，降低閾值
    if (results.size < 2) {
        results = storage.mostRelevantDocuments(query, count = 5, similarityThreshold = 0.6).toList()
    }

    return results
}

// 結構化的文件管理方式
class KnowledgeBaseManager(
    private val documentStorage: EmbeddingBasedDocumentStorage<Path>
) {
    private val categoryFolders = mapOf(
        "shipping" to listOf(
            "shipping-policy.txt" to "運送政策和時間說明",
            "delivery-time.txt" to "各地區配送時間表"
        ),
        "returns" to listOf(
            "return-process.txt" to "退貨流程和注意事項",
            "refund-policy.txt" to "退款政策和處理時間"
        ),
        "products" to listOf(
            "product-specs.txt" to "產品規格和功能說明",
            "warranty-info.txt" to "保固條款和維修服務"
        )
    )

    suspend fun loadKnowledgeBase(basePath: String = "./knowledge-base"): LoadResult {
        var successCount = 0
        var failureCount = 0
        val failedFiles = mutableListOf<String>()

        categoryFolders.forEach { (category, files) ->
            files.forEach { (fileName, description) ->
                val path = Path.of("$basePath/$category/$fileName")
                try {
                    if (Files.exists(path)) {
                        documentStorage.store(path)
                        println("✓ 已載入 $category/$fileName - $description")
                        successCount++
                    } else {
                        println("⚠ 文件不存在：$path")
                        failedFiles.add("$category/$fileName")
                        failureCount++
                    }
                } catch (e: Exception) {
                    println("✗ 載入失敗：$category/$fileName - ${e.message}")
                    failedFiles.add("$category/$fileName")
                    failureCount++
                }
            }
        }

        return LoadResult(successCount, failureCount, failedFiles)
    }

    data class LoadResult(
        val successCount: Int,
        val failureCount: Int,
        val failedFiles: List<String>
    )
}

// 使用示範
suspend fun Main() {
    val embedder = LLMEmbedder(OllamaClient(), OllamaEmbeddingModels.NOMIC_EMBED_TEXT)
    val documentEmbedder = JVMTextDocumentEmbedder(embedder)
    val documentStorage = EmbeddingBasedDocumentStorage(documentEmbedder, InMemoryVectorStorage())

    val manager = KnowledgeBaseManager(documentStorage)
    val result = manager.loadKnowledgeBase()

    println("\n知識庫載入完成：")
    println("成功：${result.successCount} 個文件")
    if (result.failureCount > 0) {
        println("失敗：${result.failureCount} 個文件")
        println("失敗清單：${result.failedFiles.joinToString(", ")}")
    }
}

// 使用持久化儲存替代記憶體儲存
val fileVectorStorage = JVMFileVectorStorage(root = Path.of("./vector-storage"))
val documentStorage = EmbeddingBasedDocumentStorage(documentEmbedder, fileVectorStorage)

// 批量處理大量文件
class BatchDocumentProcessor(
    private val documentStorage: EmbeddingBasedDocumentStorage<Path>,
    private val batchSize: Int = 50
) {
    suspend fun processBatch(documentPaths: List<Path>) {
        documentPaths.chunked(batchSize).forEach { batch ->
            withContext(Dispatchers.IO) {
                batch.forEach { path ->
                    try {
                        documentStorage.store(path)
                        println("✓ 已處理：${path.fileName}")
                    } catch (e: Exception) {
                        println("✗ 處理失敗：${path.fileName} - ${e.message}")
                    }
                }
            }
            println("已完成 ${batch.size} 個文件的處理")
        }
    }
}

// 搜尋結果快取（簡化版本）
class SearchCache {
    private val cache = mutableMapOf<String, Pair<List<Path>, Long>>()
    private val cacheTimeout = 300_000L // 5 分鐘

    fun getCachedResults(query: String): List<Path>? {
        val cached = cache[query] ?: return null
        return if (System.currentTimeMillis() - cached.second < cacheTimeout) {
            cached.first
        } else {
            cache.remove(query)
            null
        }
    }

    fun cacheResults(query: String, results: List<Path>) {
        cache[query] = Pair(results, System.currentTimeMillis())
    }
}

// AI 搜尋策略
class SmartSearchManager(
    private val documentStorage: EmbeddingBasedDocumentStorage<Path>,
    private val cache: SearchCache = SearchCache()
) {
    suspend fun smartSearch(
        query: String,
        maxResults: Int = 5,
        useCache: Boolean = true
    ): List<Path> {
        // 先檢查快取
        if (useCache) {
            cache.getCachedResults(query)?.let { return it }
        }

        // 分階段搜尋策略
        val results = when {
            // 簡單關鍵字：使用寬鬆閾值
            query.split(" ").size <= 2 -> {
                documentStorage.mostRelevantDocuments(
                    query, maxResults, similarityThreshold = 0.6
                ).toList()
            }
            // 複雜問題：使用中等閾值
            query.length > 20 -> {
                documentStorage.mostRelevantDocuments(
                    query, maxResults, similarityThreshold = 0.7
                ).toList()
            }
            // 預設情況
            else -> {
                documentStorage.mostRelevantDocuments(
                    query, maxResults, similarityThreshold = 0.65
                ).toList()
            }
        }

        // 將結果存入快取
        if (useCache && results.isNotEmpty()) {
            cache.cacheResults(query, results)
        }

        return results
    }
}

// 不同分塊策略的比較範例
class RAGPerformanceComparison {

    // 小分塊：精確但可能缺乏上下文
    fun smallChunks(): ChunkStrategy = ChunkStrategy(
        size = 200,     // 200 字元
        overlap = 50,   // 50 字元重疊
        pros = listOf("檢索精確", "回應針對性強"),
        cons = listOf("可能缺乏完整上下文", "需要更多分塊")
    )

    // 中等分塊：平衡效能與品質
    fun mediumChunks(): ChunkStrategy = ChunkStrategy(
        size = 800,     // 800 字元
        overlap = 100,  // 100 字元重疊
        pros = listOf("保持適當上下文", "效能適中"),
        cons = listOf("可能包含無關資訊")
    )

    // 大分塊：上下文豐富但檢索可能不精確
    fun largeChunks(): ChunkStrategy = ChunkStrategy(
        size = 2000,    // 2000 字元
        overlap = 200,  // 200 字元重疊
        pros = listOf("完整上下文", "減少分塊數量"),
        cons = listOf("檢索精確度可能降低", "包含更多噪音")
    )
}

data class ChunkStrategy(
    val size: Int,
    val overlap: Int,
    val pros: List<String>,
    val cons: List<String>
)