安全知识库 RAG 系统：架构设计、技术选型与实现

┌──────────────────────────────────────────────────────────────────────┐
│                          用户交互层 (React)                           │
│  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────────────────────┐ │
│  │ 对话界面  │  │ 文档管理  │  │ 反馈面板  │  │ 管理后台（权限/审计）│ │
│  └────┬─────┘  └────┬─────┘  └────┬─────┘  └──────────┬───────────┘ │
│       └──────────────┼─────────────┼───────────────────┘             │
│                      │  SSE / WebSocket                              │
└──────────────────────┼───────────────────────────────────────────────┘
                       │
┌──────────────────────┼───────────────────────────────────────────────┐
│                FastAPI 服务层（API Gateway）                          │
│  ┌──────────┐  ┌─────┴─────┐  ┌──────────┐  ┌───────────────────┐  │
│  │ 路由分发  │  │ 认证鉴权   │  │ 限流熔断  │  │ 审计日志          │  │
│  └────┬─────┘  └───────────┘  └──────────┘  └───────────────────┘  │
│       │                                                              │
│  ┌────┴─────────────────────────────────────────────────────────┐   │
│  │              LangChain RAG Pipeline                          │   │
│  │  ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌───────────┐ │   │
│  │  │ 查询改写    │ → 混合检索   │ → 重排序     │ → LLM生成   │ │   │
│  │  └────────────┘ └─────┬──────┘ └─────┬──────┘ └─────┬─────┘ │   │
│  └───────────────────────┼──────────────┼──────────────┼────────┘   │
└──────────────────────────┼──────────────┼──────────────┼────────────┘
                           │              │              │
┌──────────────────────────┼──────────────┼──────────────┼────────────┐
│                   数据存储层                                                  │
│  ┌──────────────┐  ┌─────┴──────┐  ┌───┴────┐  ┌─────┴──────────┐ │
│  │  PostgreSQL   │  │  PGVector  │  │ Redis  │  │ MinIO / S3     │ │
│  │  (元数据管理)  │  │ (向量存储)  │  │ (缓存) │  │ (文档原始文件) │ │
│  └──────────────┘  └────────────┘  └────────┘  └────────────────┘ │
└─────────────────────────────────────────────────────────────────────┘

from langchain_community.document_loaders import (
    PyPDFLoader,
    UnstructuredMarkdownLoader,
    BSHTMLLoader,
)
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import PGVector
from langchain.chains import RetrievalQA

LOADERS = {
    ".pdf": PyPDFLoader,
    ".md": UnstructuredMarkdownLoader,
    ".html": BSHTMLLoader,
}

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=64,
    separators=["\n\n", "\n", "。", "；", " "],
    length_function=len,
)

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE security_documents (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    source_url TEXT,
    doc_type VARCHAR(50),
    tags TEXT[],
    access_level INTEGER DEFAULT 0,
    created_at TIMESTAMPTZ DEFAULT NOW(),
    updated_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE document_chunks (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    document_id UUID REFERENCES security_documents(id) ON DELETE CASCADE,
    chunk_index INTEGER NOT NULL,
    content TEXT NOT NULL,
    embedding vector(1536),
    token_count INTEGER,
    metadata JSONB
);

CREATE INDEX idx_chunks_embedding ON document_chunks
    USING ivfflat (embedding vector_cosine_ops)
    WITH (lists = 100);

CREATE INDEX idx_chunks_doc_type ON security_documents(doc_type);
CREATE INDEX idx_chunks_access_level ON security_documents(access_level);

from fastapi import FastAPI, Depends
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

app = FastAPI(title="Security Knowledge Base API", version="1.0.0")


class QueryRequest(BaseModel):
    question: str
    session_id: str | None = None
    filters: dict | None = None


class QueryResponse(BaseModel):
    answer: str
    sources: list[dict]
    confidence: float
    session_id: str


@app.post("/api/v1/ask", response_model=QueryResponse)
async def ask_question(
    request: QueryRequest,
    user=Depends(get_current_user),
):
    result = await rag_chain.ainvoke(
        query=request.question,
        filters=request.filters,
        user_access_level=user.access_level,
    )
    return QueryResponse(
        answer=result["answer"],
        sources=result["sources"],
        confidence=result["confidence"],
        session_id=result["session_id"],
    )


@app.post("/api/v1/ask/stream")
async def ask_question_stream(
    request: QueryRequest,
    user=Depends(get_current_user),
):
    async def generate():
        async for event in rag_chain.astream(
            query=request.question,
            filters=request.filters,
            user_access_level=user.access_level,
        ):
            yield f"data: {event.json()}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")

原始文档 ──→ 格式解析 ──→ 内容清洗 ──→ 语义分块 ──→ 向量化 ──→ 索引入库
  │              │              │              │            │           │
  │         PDF/MD/HTML     去噪/去重      递归分割     Embedding   PGVector
  │         飞书/Confluence  元数据提取    上下文保持     批量处理    元数据关联
  │                                                  维度管理

import hashlib
from pathlib import Path
from langchain.schema import Document


class SecurityDocProcessor:
    def __init__(self, embeddings, vector_store, text_splitter):
        self.embeddings = embeddings
        self.vector_store = vector_store
        self.text_splitter = text_splitter

    async def process_document(
        self, file_path: Path, metadata: dict
    ) -> dict:
        loader_class = LOADERS.get(file_path.suffix.lower())
        if not loader_class:
            raise ValueError(f"Unsupported format: {file_path.suffix}")

        loader = loader_class(str(file_path))
        raw_docs = loader.load()

        cleaned_docs = self._clean_documents(raw_docs)
        chunks = self.text_splitter.split_documents(cleaned_docs)

        for i, chunk in enumerate(chunks):
            chunk.metadata.update({
                **metadata,
                "chunk_index": i,
                "total_chunks": len(chunks),
                "source_file": file_path.name,
                "content_hash": hashlib.md5(
                    chunk.page_content.encode()
                ).hexdigest(),
            })

        ids = [str(uuid4()) for _ in chunks]
        await self.vector_store.aadd_documents(chunks, ids=ids)

        return {
            "file": file_path.name,
            "chunks": len(chunks),
            "total_tokens": sum(
                len(c.page_content) for c in chunks
            ),
        }

    def _clean_documents(self, docs: list[Document]) -> list[Document]:
        cleaned = []
        seen_hashes = set()
        for doc in docs:
            content = doc.page_content.strip()
            if not content or len(content) < 20:
                continue
            content_hash = hashlib.md5(content.encode()).hexdigest()
            if content_hash in seen_hashes:
                continue
            seen_hashes.add(content_hash)
            cleaned.append(Document(
                page_content=content,
                metadata=doc.metadata,
            ))
        return cleaned

from langchain.text_splitter import MarkdownHeaderTextSplitter

headers_to_split = [
    ("#", "h1"),
    ("##", "h2"),
    ("###", "h3"),
]

markdown_splitter = MarkdownHeaderTextSplitter(
    headers_to_split_on=headers_to_split,
    strip_headers=False,
)

def split_security_doc(text: str, doc_type: str) -> list[Document]:
    if doc_type in ("playbook", "vulnerability_advisory"):
        md_sections = markdown_splitter.split_text(text)
        final_chunks = []
        for section in md_sections:
            if len(section.page_content) > MAX_CHUNK_SIZE:
                sub_chunks = text_splitter.split_documents([section])
                final_chunks.extend(sub_chunks)
            else:
                final_chunks.append(section)
        return final_chunks
    else:
        return text_splitter.split_text(text)

from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever


def build_hybrid_retriever(vector_store, documents, k=10):
    vector_retriever = vector_store.as_retriever(
        search_type="similarity",
        search_kwargs={"k": k},
    )

    bm25_retriever = BM25Retriever.from_documents(documents)
    bm25_retriever.k = k

    ensemble_retriever = EnsembleRetriever(
        retrievers=[vector_retriever, bm25_retriever],
        weights=[0.6, 0.4],
    )
    return ensemble_retriever

from sentence_transformers import CrossEncoder

reranker = CrossEncoder("BAAI/bge-reranker-v2-m3")


def rerank_candidates(
    query: str, candidates: list[Document], top_k: int = 5
) -> list[Document]:
    pairs = [(query, doc.page_content) for doc in candidates]
    scores = reranker.predict(pairs)

    scored_docs = list(zip(candidates, scores))
    scored_docs.sort(key=lambda x: x[1], reverse=True)

    return [doc for doc, score in scored_docs[:top_k]]

SECURITY_QA_TEMPLATE = """你是一个专业的网络安全知识助手。基于以下检索到的安全知识片段回答用户的问题。

重要规则：
1. 只基于提供的知识片段回答，不要使用你自己的知识
2. 每个关键陈述必须标注来源文档编号 [1], [2] 等
3. 如果知识片段中没有相关信息，明确告知用户"当前知识库中未找到相关资料"
4. 如果答案涉及具体操作步骤，请按顺序列出
5. 如果存在多种方案，请分别说明适用场景
6. 不要给出未经验证的安全建议

检索到的知识片段：
{context}

用户问题：{question}

回答："""

from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough


def build_rag_chain(retriever, llm):
    prompt = ChatPromptTemplate.from_template(SECURITY_QA_TEMPLATE)

    def format_docs(docs):
        formatted = []
        for i, doc in enumerate(docs, 1):
            source = doc.metadata.get("source_file", "unknown")
            title = doc.metadata.get("title", "untitled")
            formatted.append(
                f"[{i}] 来源: {source} | 标题: {title}\n"
                f"{doc.page_content}"
            )
        return "\n\n".join(formatted)

    chain = (
        {
            "context": retriever | format_docs,
            "question": RunnablePassthrough(),
        }
        | prompt
        | llm
        | StrOutputParser()
    )
    return chain

import React, { useState, useCallback } from "react";
import { Input, Button, Spin, Typography } from "antd";
import ReactMarkdown from "react-markdown";

const { Text } = Typography;

interface Source {
  doc_id: string;
  title: string;
  source_file: string;
  score: number;
  content: string;
}

interface ChatMessage {
  role: "user" | "assistant";
  content: string;
  sources?: Source[];
  isStreaming?: boolean;
}

function ChatInterface() {
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [inputValue, setInputValue] = useState("");
  const [isLoading, setIsLoading] = useState(false);

  const handleSend = useCallback(async () => {
    if (!inputValue.trim() || isLoading) return;

    const userMessage: ChatMessage = {
      role: "user",
      content: inputValue,
    };

    const assistantMessage: ChatMessage = {
      role: "assistant",
      content: "",
      isStreaming: true,
    };

    setMessages((prev) => [...prev, userMessage, assistantMessage]);
    setInputValue("");
    setIsLoading(true);

    try {
      const response = await fetch("/api/v1/ask/stream", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ question: inputValue }),
      });

      const reader = response.body?.getReader();
      const decoder = new TextDecoder();

      while (true) {
        const { done, value } = await reader!.read();
        if (done) break;

        const chunk = decoder.decode(value);
        const lines = chunk.split("\n").filter((line) => line.startsWith("data: "));

        for (const line of lines) {
          const data = line.slice(6);
          if (data === "[DONE]") break;

          const event = JSON.parse(data);
          if (event.type === "token") {
            setMessages((prev) => {
              const updated = [...prev];
              const last = updated[updated.length - 1];
              updated[updated.length - 1] = {
                ...last,
                content: last.content + event.content,
              };
              return updated;
            });
          } else if (event.type === "sources") {
            setMessages((prev) => {
              const updated = [...prev];
              updated[updated.length - 1] = {
                ...updated[updated.length - 1],
                sources: event.sources,
                isStreaming: false,
              };
              return updated;
            });
          }
        }
      }
    } finally {
      setIsLoading(false);
    }
  }, [inputValue, isLoading]);

  return (
    <div style={{ display: "flex", height: "100vh" }}>
      <div style={{ flex: 1, display: "flex", flexDirection: "column" }}>
        <div style={{ flex: 1, overflow: "auto", padding: 24 }}>
          {messages.map((msg, idx) => (
            <div key={idx} style={{ marginBottom: 16 }}>
              <Text strong>{msg.role === "user" ? "我" : "安全助手"}</Text>
              <ReactMarkdown>{msg.content}</ReactMarkdown>
              {msg.sources && <SourcePanel sources={msg.sources} />}
            </div>
          ))}
        </div>
        <div style={{ padding: 16, borderTop: "1px solid #f0f0f0" }}>
          <Input.Search
            value={inputValue}
            onChange={(e) => setInputValue(e.target.value)}
            enterButton="发送"
            onSearch={handleSend}
            loading={isLoading}
            placeholder="输入安全相关问题..."
          />
        </div>
      </div>
    </div>
  );
}

function SourcePanel({ sources }: { sources: Source[] }) {
  const [expandedIdx, setExpandedIdx] = useState<number | null>(null);

  return (
    <div style={{ marginTop: 8 }}>
      <Text type="secondary" style={{ fontSize: 12 }}>
        参考来源 ({sources.length})
      </Text>
      {sources.map((source, idx) => (
        <div
          key={source.doc_id}
          style={{
            padding: "4px 8px",
            margin: "4px 0",
            borderRadius: 4,
            background: "#fafafa",
            cursor: "pointer",
          }}
          onClick={() => setExpandedIdx(expandedIdx === idx ? null : idx)}
        >
          <Text style={{ fontSize: 12 }}>
            [{idx + 1}] {source.title}
            <Text type="secondary"> ({source.source_file})</Text>
            <Text type="secondary"> 相似度: {(source.score * 100).toFixed(1)}%</Text>
          </Text>
          {expandedIdx === idx && (
            <div style={{ marginTop: 4, fontSize: 12, color: "#666" }}>
              {source.content}
            </div>
          )}
        </div>
      ))}
      <FeedbackButtons sources={sources} />
    </div>
  );
}

version: "3.9"

services:
  api:
    build:
      context: .
      dockerfile: Dockerfile
    ports:
      - "8000:8000"
    environment:
      - DATABASE_URL=postgresql://admin:secret@db:5432/security_kb
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - REDIS_URL=redis://cache:6379/0
    depends_on:
      db:
        condition: service_healthy
      cache:
        condition: service_started
    volumes:
      - ./uploads:/app/uploads
    command: uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload

  web:
    build:
      context: ./frontend
      dockerfile: Dockerfile
    ports:
      - "3000:3000"
    environment:
      - REACT_APP_API_URL=http://localhost:8000
    depends_on:
      - api

  db:
    image: pgvector/pgvector:pg16
    ports:
      - "5432:5432"
    environment:
      POSTGRES_DB: security_kb
      POSTGRES_USER: admin
      POSTGRES_PASSWORD: secret
    volumes:
      - pgdata:/var/lib/postgresql/data
      - ./sql/init.sql:/docker-entrypoint-initdb.d/init.sql
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U admin -d security_kb"]
      interval: 5s
      timeout: 5s
      retries: 5

  cache:
    image: redis:7-alpine
    ports:
      - "6379:6379"

  worker:
    build:
      context: .
      dockerfile: Dockerfile
    environment:
      - DATABASE_URL=postgresql://admin:secret@db:5432/security_kb
      - OPENAI_API_KEY=${OPENAI_API_KEY}
    depends_on:
      - db
      - cache
    command: celery -A app.tasks worker --loglevel=info --concurrency=4

volumes:
  pgdata:

import hashlib
import json
from redis.asyncio import Redis


class SemanticCache:
    def __init__(self, redis: Redis, ttl: int = 86400):
        self.redis = redis
        self.ttl = ttl

    async def get(self, query: str) -> dict | None:
        key = f"rag_cache:{hashlib.sha256(query.encode()).hexdigest()}"
        cached = await self.redis.get(key)
        return json.loads(cached) if cached else None

    async def set(self, query: str, result: dict):
        key = f"rag_cache:{hashlib.sha256(query.encode()).hexdigest()}"
        await self.redis.setex(key, self.ttl, json.dumps(result))

eval_dataset = [
    {
        "question": "如何检测 Redis 未授权访问漏洞？",
        "ground_truth": "通过执行 redis-cli INFO 命令检查服务信息...",
        "contexts": ["Redis 安全加固指南 v2.1 第三章..."],
        "metadata": {"category": "vulnerability_detection", "difficulty": "easy"},
    },
    {
        "question": "内网横向移动的常见手法有哪些？",
        "ground_truth": "横向移动手法包括：Pass-the-Hash、WMI远程执行...",
        "contexts": ["横向移动检测与防御 Playbook..."],
        "metadata": {"category": "threat_analysis", "difficulty": "hard"},
    },
    {
        "question": "CVE-2024-3094 影响范围和修复方案",
        "ground_truth": "XZ Utils 5.6.0 和 5.6.1 版本存在后门...",
        "contexts": ["CVE-2024-3094 应急通告..."],
        "metadata": {"category": "vulnerability_response", "difficulty": "medium"},
    },
]

from fastapi import APIRouter

router = APIRouter()


@router.post("/api/v1/feedback")
async def submit_feedback(
    feedback: FeedbackRequest,
    user=Depends(get_current_user),
):
    await db.execute(
        """INSERT INTO feedback
           (session_id, message_id, rating, comment, user_id, created_at)
           VALUES ($1, $2, $3, $4, $5, NOW())""",
        feedback.session_id,
        feedback.message_id,
        feedback.rating,
        feedback.comment,
        user.id,
    )

    if feedback.rating == "harmful":
        await alert_service.notify_security_team(
            message_id=feedback.message_id,
            reason=feedback.comment,
        )

    return {"status": "recorded"}

用户提问
    │
    ▼
┌─────────────┐    ┌─────────────┐    ┌──────────────┐
│ 查询改写     │ →  │ 查询分类     │ →  │ 路由分发      │
│ (HyDE/扩展)  │    │ (意图识别)   │    │ (检索策略)    │
└─────────────┘    └─────────────┘    └──────┬───────┘
                                             │
                                    ┌────────┴────────┐
                                    │                 │
                                    ▼                 ▼
                              ┌──────────┐    ┌──────────┐
                              │ 向量检索  │    │ BM25检索  │
                              │ (语义)    │    │ (关键词)  │
                              └────┬─────┘    └────┬─────┘
                                   │               │
                                   └───────┬───────┘
                                           ▼
                                   ┌──────────────┐
                                   │  RRF 分数融合  │
                                   └──────┬───────┘
                                          ▼
                                   ┌──────────────┐
                                   │  BGE 重排序    │
                                   └──────┬───────┘
                                          ▼
                                   ┌──────────────┐
                                   │  上下文构建    │
                                   └──────┬───────┘
                                          ▼
                                   ┌──────────────┐
                                   │  LLM 流式生成  │
                                   └──────┬───────┘
                                          ▼
                                   ┌──────────────┐
                                   │  来源标注输出   │
                                   └──────────────┘

security-kb-rag/
├── backend/
│   ├── app/
│   │   ├── main.py              # FastAPI 入口
│   │   ├── api/                 # 路由定义
│   │   │   ├── ask.py           # 问答接口
│   │   │   ├── documents.py     # 文档管理接口
│   │   │   └── feedback.py      # 反馈接口
│   │   ├── core/
│   │   │   ├── config.py        # 配置管理
│   │   │   ├── security.py      # 认证鉴权
│   │   │   └── dependencies.py  # 依赖注入
│   │   ├── rag/
│   │   │   ├── chain.py         # RAG Chain 编排
│   │   │   ├── retriever.py     # 混合检索器
│   │   │   ├── reranker.py      # 重排序模块
│   │   │   ├── prompt.py        # Prompt 模板
│   │   │   └── cache.py         # 语义缓存
│   │   ├── document/
│   │   │   ├── loader.py        # 多格式文档加载
│   │   │   ├── splitter.py      # 分块策略
│   │   │   └── processor.py     # 文档处理管线
│   │   └── models/
│   │       ├── document.py      # 文档数据模型
│   │       └── feedback.py      # 反馈数据模型
│   ├── sql/
│   │   └── init.sql             # 数据库初始化
│   ├── Dockerfile
│   └── requirements.txt
├── frontend/
│   ├── src/
│   │   ├── components/
│   │   │   ├── ChatInterface.tsx
│   │   │   ├── SourcePanel.tsx
│   │   │   └── FeedbackButtons.tsx
│   │   ├── hooks/
│   │   │   └── useStreamChat.ts
│   │   ├── services/
│   │   │   └── api.ts
│   │   └── App.tsx
│   ├── Dockerfile
│   └── package.json
├── docker-compose.yml
├── docker-compose.prod.yml
└── README.md                     # <!-- [TODO](link-to-repo) -->