arbitration_system_V20/law_rag/retrieval/retrieval.py

import os
import time
import json
import numpy as np
import faiss
from text2vec import SentenceModel
from openai import OpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI
import config.config
# 向量存储目录结构
VECTOR_STORE_BASE = config.config.VECTOR_STORE_BASE
MODEL_PATH = config.config.MODEL_PATH
# DeepSeek API配置
DEEPSEEK_API_KEY = config.config.DEEPSEEK_API  # 请替换为您的实际API密钥
DEEPSEEK_BASE_URL = "https://api.deepseek.com"

# 加载文本向量化模型
def load_embedding_model(device="cuda"):
    """加载text2vec-base-chinese模型"""
    print(f"正在加载向量化模型: text2vec-base-chinese")
    print(f"使用设备: {device}")
    return SentenceModel(MODEL_PATH, device=device)

# 转译用户问题的链
def build_question_translation_chain():
    """构建用于转译用户问题的LLM链"""
    llm = ChatOpenAI(
        api_key=DEEPSEEK_API_KEY,
        base_url=DEEPSEEK_BASE_URL,
        model="deepseek-chat",
        temperature=0.3,
    )
    
    system_template = """
    你是一个法律问题转译专家。你的任务是将用户的自然语言法律问题转译为更加专业、简洁、直接的形式，以便于向量检索系统能够找到最相关的法律条文。
    
    遵循以下原则：
    1. 提取问题中的关键法律概念和术语
    2. 使用法律专业术语重新表述问题
    3. 去除无关信息，保留核心法律问题
    4. 确保转译后的问题简洁明了，直指法律要点
    
    请直接输出转译后的问题，不要包含额外解释。
    """
    
    prompt = ChatPromptTemplate(
        [
            ("system", system_template),
            ("human", "{question}"),
        ]
    )
    
    return prompt | llm | StrOutputParser()

# 使用DeepSeek API直接调用
def translate_question_with_deepseek(question):
    """使用DeepSeek API直接调用模型进行问题转译"""
    client = OpenAI(api_key=DEEPSEEK_API_KEY, base_url=DEEPSEEK_BASE_URL)
    
    system_template = """
    你是一个法律问题转译专家。你的任务是将用户的自然语言法律问题转译为更加专业、简洁、直接的形式，以便于向量检索系统能够找到最相关的法律条文。
    
    遵循以下原则：
    1. 提取问题中的关键法律概念和术语
    2. 使用法律专业术语重新表述问题
    3. 去除无关信息，保留核心法律问题
    4. 确保转译后的问题简洁明了，直指法律要点
    
    请直接输出转译后的问题，不要包含额外解释。
    """
    
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": system_template},
            {"role": "user", "content": question},
        ],
        temperature=0.3,
        stream=False
    )
    
    return response.choices[0].message.content

# 向量检索相关函数
def load_vector_store(level_dir):
    """加载指定层级的向量存储"""
    store_dir = os.path.join(VECTOR_STORE_BASE, level_dir)
    index_path = os.path.join(store_dir, "index")
    
    if not os.path.exists(f"{index_path}_index"):
        raise FileNotFoundError(f"向量索引文件不存在: {index_path}_index")
    
    index = faiss.read_index(f"{index_path}_index")
    
    # 加载向量、文本和元数据
    vectors = np.load(f"{index_path}_vectors.npy")
    
    # texts 现在包含的是原始完整文本（包含"第xx条  "前缀）
    texts = np.load(f"{index_path}_texts.npy", allow_pickle=True)
    metadata = np.load(f"{index_path}_metadata.npy", allow_pickle=True)
    
    print(f"成功加载{level_dir}向量库: {len(vectors)}条记录")
    
    return {
        "index": index,
        "vectors": vectors,
        "texts": texts,
        "metadata": metadata
    }

def search_in_level(query_embedding, level_dir, top_k=10, store=None):
    """在指定层级中搜索，返回top_k个结果"""
    if store is None:
        store = load_vector_store(level_dir)
    
    # 搜索
    distances, indices = store["index"].search(query_embedding.reshape(1, -1), top_k)
    
    results = []
    for i, idx in enumerate(indices[0]):
        if idx != -1:  # Faiss返回-1表示找不到足够的结果
            results.append({
                "text": store["texts"][idx],
                "metadata": store["metadata"][idx],
                "score": float(1.0 - distances[0][i]),  # 将距离转换为相似度分数
                "index": int(idx),
                "level": level_dir
            })
    
    return results

def rank_final_results(results, query_embedding, top_k=5):
    """对所有结果重新排序，返回top_k个最相关的结果"""
    # 提取所有文本并计算与查询的相似度
    texts = [item["text"] for item in results]
    
    # 加载模型
    embedding_model = load_embedding_model()
    
    # 计算文本向量
    embeddings = embedding_model.encode(texts)
    
    # 计算相似度
    similarities = np.dot(embeddings, query_embedding.T) / (
        np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query_embedding)
    )
    
    # 更新相似度分数
    for i, item in enumerate(results):
        item["score"] = float(similarities[i])
    
    # 先按相似度排序，选出前20个结果
    sorted_by_similarity = sorted(results, key=lambda x: x["score"], reverse=True)[:20]
    
    # 去重 - 基于文本内容
    unique_results = []
    seen_texts = set()
    
    for item in sorted_by_similarity:
        if item["text"] not in seen_texts:
            seen_texts.add(item["text"])
            unique_results.append(item)
            if len(unique_results) >= top_k:
                break
    
    # 对筛选后的top_k个不重复结果按层级优先级排序
    level_priority = {"level_1": 1, "level_2": 2, "level_3": 3}
    final_results = sorted(unique_results, key=lambda x: level_priority.get(x["level"], 999))
    
    return final_results

def hierarchical_search(user_query, l1_results=10, l2_sub_searches=3, l3_sub_searches=1, final_results=5):
    """执行三层次的向量检索"""
    print(f"开始处理用户查询: {user_query}")
    
    # 1. 转译用户查询 - 使用DeepSeek API
    translated_query = translate_question_with_deepseek(user_query)
    print(f"转译后的查询: {translated_query}")
    
    # 2. 加载向量化模型
    embedding_model = load_embedding_model()
    
    # 3. 对转译后的查询进行向量化
    query_embedding = embedding_model.encode(translated_query)
    
    # 4. 加载各层级向量库
    level_1_store = load_vector_store("level_1")
    level_2_store = load_vector_store("level_2")
    level_3_store = load_vector_store("level_3")
    
    # 5. 在level_1中搜索前l1_results个结果
    level_1_results = search_in_level(query_embedding, "level_1", l1_results, level_1_store)
    print(f"一级检索结果: {len(level_1_results)}条")
    
    # 6. 对每个level_1结果在level_2中搜索
    all_results = []
    all_results.extend(level_1_results)
    
    level_2_results = []
    for l1_item in level_1_results:
        # 用level_1中的文本创建查询向量
        l1_text_embedding = embedding_model.encode(l1_item["text"])
        # 在level_2中查找相关内容
        l2_items = search_in_level(l1_text_embedding, "level_2", l2_sub_searches, level_2_store)
        level_2_results.extend(l2_items)
    
    print(f"二级检索结果: {len(level_2_results)}条")
    all_results.extend(level_2_results)
    
    # 7. 对每个level_2结果在level_3中搜索
    # level_3_results = []
    # for l2_item in level_2_results:
    #     # 用level_2中的文本创建查询向量
    #     l2_text_embedding = embedding_model.encode(l2_item["text"])
    #     # 在level_3中查找相关内容
    #     l3_items = search_in_level(l2_text_embedding, "level_3", l3_sub_searches, level_3_store)
    #     level_3_results.extend(l3_items)
    
    # print(f"二级检索结果: {len(level_3_results)}条")
    # all_results.extend(level_3_results)
    
    print(f"所有层级检索结果总量: {len(all_results)}条")
    
    # 8. 对所有结果重新排序，选取最相关的前final_results个
    final_results = rank_final_results(all_results, query_embedding, final_results)
    print(f"最终筛选结果: {len(final_results)}条")
    
    return {
        "original_query": user_query,
        "translated_query": translated_query,
        "results": final_results
    }

def main():
    
    while True:
        user_query = input("\n请输入您的法律问题 (输入'q'退出): ")
        if user_query.lower() == 'q':
            break
        
        start_time = time.time()
        results = hierarchical_search(user_query)
        
        print("\n检索结果:")
        print(f"原始问题: {results['original_query']}")
        print(f"转译问题: {results['translated_query']}")
        print(f"耗时: {time.time() - start_time:.2f}秒")
        
        for i, item in enumerate(results['results']):
            print(f"\n[{i+1}] 相似度: {item['score']:.4f} - 层级: {item['level']}")
            print(f"法律ID: {item['metadata'].get('law_id', '未知')}")
            print(f"内容: {item['text']}")  # 这里输出的是原始完整文本
    

if __name__ == "__main__":
    main()