Python RAG 关键词过滤

作者：追风剑情发布于：2026-5-20 15:05 分类：AI
import requests

# ========== 初始化本地 Ollama 模型 ==========
OLLAMA_GENERATE_URL = "http://localhost:11434/api/generate"
GENERATE_MODEL = "deepseek-r1:7b"  # 使用您本地的模型

def generate_response(query, retrieved_doc):
    """使用本地 Ollama 模型生成响应"""
    prompt = f"""问题: {query}
相关文档内容: {retrieved_doc}

请基于以上文档内容回答问题。"""

    response = requests.post(
        OLLAMA_GENERATE_URL,
        json={
            "model": GENERATE_MODEL,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0.7,
                "top_k": 50,
                "top_p": 0.9
            }
        }
    )
    if response.status_code == 200:
        return response.json()["response"]
    else:
        return f"生成失败: {response.status_code}"

def filter_response(response_text, forbidden_keywords):
    """
    过滤生成的响应，检查是否包含禁止关键词
    
    参数:
        response_text: 生成的原始文本
        forbidden_keywords: 禁止关键词列表
    
    返回:
        如果包含禁止关键词，返回过滤后的提示信息；
        否则返回原始文本
    """
    for keyword in forbidden_keywords:
        if keyword in response_text:
            return f"[已过滤] 生成的响应包含敏感词: '{keyword}'，已阻止输出。"
    return response_text

# ========== 示例1：正常内容（无敏感词）==========
print("=" * 60)
print("示例1：正常查询（无敏感词）")
print("=" * 60)

retrieved_doc = "法国是欧洲的一个国家，首都位于巴黎。巴黎以埃菲尔铁塔和卢浮宫闻名于世。"
query = "法国的首都是哪里？"

# 禁用关键词列表
forbidden_keywords = ["暴力", "色情", "非法", "仇恨", "歧视"]

print(f"查询: {query}")
print(f"检索到的文档: {retrieved_doc}")
print(f"禁止关键词: {forbidden_keywords}")
print()

# 生成响应并进行过滤
print("正在生成响应...")
response = generate_response(query, retrieved_doc)
print(f"生成的原始响应: {response}")

filtered_response = filter_response(response, forbidden_keywords)
print(f"过滤后的响应: {filtered_response}")

# ========== 示例2：包含敏感词的内容 ==========
print("\n" + "=" * 60)
print("示例2：包含敏感词的查询")
print("=" * 60)

retrieved_doc_sensitive = "某组织计划使用暴力和非法手段达成其政治目的。"
query_sensitive = "这个组织的计划是什么？"

print(f"查询: {query_sensitive}")
print(f"检索到的文档: {retrieved_doc_sensitive}")
print(f"禁止关键词: {forbidden_keywords}")
print()

# 生成响应并进行过滤
print("正在生成响应...")
response_sensitive = generate_response(query_sensitive, retrieved_doc_sensitive)
print(f"生成的原始响应: {response_sensitive}")

filtered_response_sensitive = filter_response(response_sensitive, forbidden_keywords)
print(f"过滤后的响应: {filtered_response_sensitive}")

# ========== 示例3：更高级的过滤函数 ==========
print("\n" + "=" * 60)
print("示例3：高级过滤（记录日志 + 关键词高亮）")
print("=" * 60)

def advanced_filter(response_text, forbidden_keywords, log_file="filter_log.txt"):
    """
    高级过滤函数：
    1. 检测敏感词
    2. 记录日志到文件
    3. 标记敏感词位置
    """
    detected_keywords = []
    for keyword in forbidden_keywords:
        if keyword in response_text:
            detected_keywords.append(keyword)
    
    if detected_keywords:
        # 记录日志
        import datetime
        log_entry = f"[{datetime.datetime.now()}] 检测到敏感词: {detected_keywords}\n文本: {response_text}\n"
        with open(log_file, "a", encoding="utf-8") as f:
            f.write(log_entry)
        
        # 标记敏感词位置（用 *** 替换）
        filtered = response_text
        for kw in detected_keywords:
            filtered = filtered.replace(kw, f"[***{kw}***]")
        
        return f"[警告] 检测到敏感词 {detected_keywords}，已记录日志。\n过滤后的文本: {filtered}"
    else:
        return response_text

# 测试高级过滤
test_text = "这篇文章包含暴力和色情内容，属于非法传播。"
print(f"测试文本: {test_text}")
result = advanced_filter(test_text, ["暴力", "色情", "非法"])
print(f"高级过滤结果:\n{result}")
运行测试