9. 相似度检索原理

# -*- coding: utf-8 -*-
"""
@Created on ： 2026/6/22 10:18
@creator ： er_nao
@File ：day_94.py
@Description ：相似度检索原理
"""
import faiss
import numpy as np
import json
import os
from config import TONGYI_API_KEY
# 复用你之前的文本向量化函数，给用户问题生成向量
import dashscope
from dashscope import TextEmbedding

# ====================== 核心配置区：替换成你自己的项目路径 ======================
# 1. 你生成的faiss索引文件路径
INDEX_PATH = r"F:\\RAG-Learning-Project\\vector-database\\rag_vector_index.faiss"
# 2. 你Day93生成的元数据文件路径
META_DATA_PATH = r"F:\\RAG-Learning-Project\\vector-database\\rag_meta_data.json"
# 3. 你的通义千问API-KEY（和之前文本向量化用的一样）
dashscope.api_key = TONGYI_API_KEY
# 4. 检索配置：返回前3个最相关的结果，相似度阈值0.7
TOP_K = 3
SIMILARITY_THRESHOLD = 0.7
# ==================================================================================

# ====================== 核心函数1：加载faiss向量库和元数据 ======================
def load_vector_lib_and_meta(index_path:str, meta_data_path:str):
    """
        加载你之前生成的向量库和元数据，不用管底层怎么存的，直接调用就行
        返回：加载好的索引对象、元数据列表
    """

    if not os.path.exists(index_path):
        print(f"索引不存在，路径:{index_path}")
        return None, None
    if not os.path.exists(meta_data_path):
        print(f"源数据文件不存在，路径：{meta_data_path}")
        return None, None

    # 加载faiss索引，不用管底层怎么加载的，直接调用就行
    try:
        index = faiss.read_index(index_path)
        print(f"向量库加载成功，索引总向量数：{index.ntotal},维度：{index.d}")
    except Exception as e:
        print(f"索引加载失败，错误信息：{e}")
        return None, None

    # 加载元数据，不用管底层怎么存的，直接调用就行
    try:
        with open(meta_data_path,"r",encoding="utf-8") as f:
            meta_data = json.load(f)
        print(f"源数据加载成功，一共{len(meta_data)}条元数据")
    except Exception as e:
        print(f"元数据加载失败，错误信息：{e}")
        return None, None

    return index, meta_data


# ====================== 核心函数2：给用户问题生成向量（复用之前的函数） ======================
def generate_question_embedding(question:str):
    """
        给用户的提问生成向量，和之前文本块的向量化完全一致，不用管底层怎么生成的
        返回：问题的向量（numpy数组）
    """
    if not question or not question.strip():
        print("问题不能为空")
        return None

    # 调用通义向量模型，生成问题的向量，不用管底层怎么算的
    try:
        response = TextEmbedding.call(
            model= TextEmbedding.Models.text_embedding_v2,
            input= question
        )
        if response.status_code ==200:
            # 提取向量，转成faiss要求的格式
            embedding = np.array(response.output["embeddings"][0]["embedding"], dtype=np.float32).reshape(1,-1)
            print(f"问题向量生成成功，维度：{embedding.shape[1]}")
            return embedding
        else:
            print(f"问题向量生成失败，原因：{response.message}")
            return None
    except Exception as e:
        print(f"问题生成异常，错误信息：{str(e)}")
        return None

# ====================== 核心函数3：相似度检索（核心中的核心，只讲怎么用） ======================
def similarity_search(index:faiss.IndexFlatL2, question_embedding:np.ndarray, top_k: int=3):
    """
        执行相似度检索，不用管底层怎么算的，直接调用就行
        输入：索引对象、问题向量、要返回的结果数量
        输出：按相似度排序的文本块ID列表、对应的相似度数值
    """

    if not index or question_embedding is None:
        print("错误：索引或问题向量为空，无法检索")
        return None,None

    # 执行相似度检索，不用管底层怎么算的，直接调用就行
    # distances：相似度数值，越小说明语义越像
    # indices：对应的文本块ID，按相似度从高到低排序
    distances, indices = index.search(question_embedding, top_k)

    result_indices =  indices[0].tolist()
    result_distances = distances[0].tolist()

    print(f"相似度检索完成，返回前{top_k}个结果")

    for i in range(top_k):
        print(f"  排名{i + 1}：文本块ID={result_indices[i]}，相似度数值={result_distances[i]}")

    return result_indices, result_distances

# ====================== 核心函数4：根据检索ID，获取对应的原文内容 ======================
def get_related_content(indices:list, meta_data:list, similarity_threshold:float=0.7):

    if not indices or not meta_data:
        print("错误：检索ID或元数据为空")
        return []

    related_content_list =[]
    for idx, block_id in enumerate(indices):
        # 从元数据里找到对应的文本内容
        for meta in meta_data:
            if meta["index_id"] == block_id:

                if idx < len(indices) and indices[idx] < len(meta_data):
                    related_content_list.append({
                        "rank": idx + 1,
                        "block_id": block_id,
                        "chunk_content": meta["chunk_content"],
                        "similarity": indices[idx]
                    })
            break

    print(f" 相关内容提取完成，共提取到 {len(related_content_list)} 条有效内容")
    return related_content_list

# ====================== 主函数：一键执行RAG相似度检索全流程 ======================
def main_rag_search_pipeline(user_question:str):
    print("=" * 80)
    print(f" RAG相似度检索全流程开始，用户问题：{user_question}")
    print("=" * 80)

    # 1. 加载向量库和元数据
    index, meta_data = load_vector_lib_and_meta(INDEX_PATH, META_DATA_PATH)
    if not index or not meta_data:
        print(" 全流程终止：向量库或元数据加载失败")
        return

    # 2. 给用户问题生成向量
    question_embedding = generate_question_embedding(user_question)
    if question_embedding is None:
        print(" 全流程终止：问题向量生成失败")
        return

    # 3. 执行相似度检索
    result_indices, result_distances = similarity_search(index, question_embedding, TOP_K)
    if not result_indices:
        print(" 全流程终止：相似度检索失败")
        return

    # 4. 获取对应的原文内容
    related_content = get_related_content(result_indices, meta_data, SIMILARITY_THRESHOLD)
    if not related_content:
        print(" 全流程终止：未找到相关内容")
        return

    # 5. 打印最终结果，你可以直接把这些内容给大模型
    print("\n" + "=" * 80)
    print(" RAG相似度检索全流程完成！最终相关内容：")
    print("=" * 80)
    for content in related_content:
        print(f"\n【排名{content['rank']} | 文本块ID{content['block_id']}】")
        print(f"相关内容：{content['chunk_content']}")

    return related_content

# ====================== 程序入口：修改用户问题，一键运行 ======================
if __name__ == "__main__":
    # 替换成你要问的问题，比如「Day87的学习内容是什么？」「RAG的核心作用是什么？」
    USER_QUESTION = "Day87的学习内容是什么？"
    # 一键执行全流程
    main_rag_search_pipeline(USER_QUESTION)