rag_retriever.py

# rag_retriever.py
import faiss
import numpy as np
import json
import os
from zhipuai import ZhipuAI
from typing import List, Dict, Literal
from app_config import load_settings

# --- 1. 配置 ---

# 知识库文件前缀
RULES_KB_NAME = "coc_rules"
MODULE_KB_NAME = "current_module"

# 知识库类型
KnowledgeBaseType = Literal["coc_rules", "current_module"]

# --- 2. RAG 检索器类 ---

# --- 确保这个类定义存在 ---
class RAGRetriever:
    """
    负责加载知识库索引，并提供检索服务
    """
    def __init__(self, api_key: str = None):
        if api_key is None:
            api_key = load_settings().require_embedding_api_key()
        self.client = ZhipuAI(api_key=api_key)
        self.indexes: Dict[KnowledgeBaseType, faiss.Index] = {}
        self.chunk_maps: Dict[KnowledgeBaseType, Dict[int, str]] = {}

        # 加载所有可用的知识库
        self._load_kb(RULES_KB_NAME)
        self._load_kb(MODULE_KB_NAME)

    def _load_kb(self, kb_name: KnowledgeBaseType):
        """加载单个知识库（索引 + 映射文件）"""
        index_file = f"{kb_name}.index"
        map_file = f"{kb_name}_chunk_map.json"

        if os.path.exists(index_file) and os.path.exists(map_file):
            print(f"正在加载知识库: {kb_name}...")
            try:
                # 加载FAISS索引
                self.indexes[kb_name] = faiss.read_index(index_file)

                # 加载JSON块映射
                with open(map_file, "r", encoding="utf-8") as f:
                    # JSON的key是字符串，需要转回为int
                    loaded_map = json.load(f)
                    self.chunk_maps[kb_name] = {int(k): v for k, v in loaded_map.items()}

                print(f"[成功] 加载 {kb_name} (包含 {self.indexes[kb_name].ntotal} 个向量)")

            except Exception as e:
                print(f"[错误] 加载 {kb_name} 失败: {e}")
        else:
            print(f"[警告] 未找到知识库文件: {kb_name} (index或map不存在)，跳过加载。")

    def get_embedding(self, text: str) -> np.ndarray:
        """获取单个文本的嵌入向量"""
        try:
            response = self.client.embeddings.create(
                model="embedding-2",
                input=[text]
            )
            embedding = response.data[0].embedding
            return np.array([embedding], dtype='float32')
        except Exception as e:
            print(f"获取嵌入失败: {e}")
            return None

    def search(self,
               query_text: str,
               kb_type: KnowledgeBaseType = "coc_rules",
               top_k: int = 3) -> List[Dict[str, str]]:
        """
        在指定的知识库中执行语义检索

        :param query_text: 用户的查询
        :param kb_type: 'coc_rules' (规则库) 或 'current_module' (模组库)
        :param top_k: 返回最相似的 K 个结果
        :return: 包含原文和来源的字典列表
        """
        if kb_type not in self.indexes:
            print(f"错误：知识库 '{kb_type}' 未加载。")
            return []

        index = self.indexes[kb_type]
        chunk_map = self.chunk_maps[kb_type]

        # 1. 获取查询的嵌入
        query_embedding = self.get_embedding(query_text)
        if query_embedding is None:
            return []

        # 2. 执行 FAISS 搜索
        distances, indices = index.search(query_embedding, k=top_k)

        # 3. 组装结果
        results = []
        for i in range(len(indices[0])):
            retrieved_index = indices[0][i]
            if retrieved_index == -1:
                continue

            original_text = chunk_map.get(retrieved_index)

            if original_text:
                results.append({
                    "kb_type": kb_type,
                    "index_id": int(retrieved_index),
                    "text_chunk": original_text,
                    "distance": float(distances[0][i])
                })

        return results

# --- 3. 独立测试 (保持不变) ---
if __name__ == "__main__":
    retriever = RAGRetriever()
    print("\n" + "="*50)
    print("\n测试 [规则库] 查询: 'How does Sanity work?'")
    rule_results = retriever.search(query_text="How does Sanity work?", kb_type="coc_rules", top_k=2)
    for i, res in enumerate(rule_results):
        print(f"--- 规则结果 {i+1} (Distance: {res['distance']:.4f}) ---")
        print(res['text_chunk'][:200] + "...")
    print("\n" + "="*50)
    print("\n测试 [模组库] 查询: 'What should I do at the old house?'")
    module_results = retriever.search(query_text="What should I do at the old house?", kb_type="current_module", top_k=2)
    for i, res in enumerate(module_results):
        print(f"--- 模组结果 {i+1} (Distance: {res['distance']:.4f}) ---")
        print(res['text_chunk'][:200] + "...")