基于 PostgreSQL 构建 Agentic RAG 应用

┌──────────────────────────────────────────────────────┐
│                 Agentic RAG 系统                      │
├──────────────────────────────────────────────────────┤
│                                                      │
│  ┌──────────┐    ┌───────────┐    ┌──────────────┐   │
│  │ Planner  │    │ Retriever │    │ Generator    │   │
│  │ Agent    │ →  │ Agent     │ →  │ Agent        │   │
│  │ (规划)   │    │ (检索)    │    │ (生成)       │   │
│  └──────────┘    └───────────┘    └──────────────┘   │
│       ↓               ↓                ↓             │
│  ┌─────────────────────────────────────────────┐     │
│  │        PostgreSQL（统一存储层）              │     │
│  ├─────────────────────────────────────────────┤     │
│  │ pgvector     │ Apache AGE │ 记忆表          │     │
│  │ (向量检索)   │ (知识图谱) │ (Agent 记忆)    │     │
│  │              │            │                 │     │
│  │ tencentdb_ai（大模型调用）                  │     │
│  └─────────────────────────────────────────────┘     │
└──────────────────────────────────────────────────────┘

-- 知识库表
CREATE TABLE rag.documents (
    id BIGSERIAL PRIMARY KEY,
    title TEXT,
    content TEXT NOT NULL,
    chunk_index INT DEFAULT 0,
    embedding vector(1024),
    source TEXT,
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE INDEX idx_docs_embedding ON rag.documents
USING hnsw (embedding vector_cosine_ops);

-- Agent 检索记忆
CREATE TABLE rag.search_memory (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    session_id UUID NOT NULL,
    query TEXT NOT NULL,
    strategy TEXT,
    results_summary TEXT,
    satisfaction FLOAT,  -- Agent 对结果的满意度
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- 查询改写日志
CREATE TABLE rag.query_rewrites (
    id BIGSERIAL PRIMARY KEY,
    original_query TEXT,
    rewritten_query TEXT,
    rewrite_reason TEXT,
    session_id UUID,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Agent 使用大模型改写查询
CREATE OR REPLACE FUNCTION rag.rewrite_query(original TEXT)
RETURNS TEXT AS $$
DECLARE
    rewritten TEXT;
BEGIN
    SELECT tencentdb_ai.chat_completions(
        '<your-llm-model>',
        '你是一个查询优化专家。请将以下用户问题改写为更适合向量检索的查询。只输出改写后的查询，不要解释。' ||
        E'\\n\\n原始问题：' || original
    ) INTO rewritten;

    -- 记录改写
    INSERT INTO rag.query_rewrites (original_query, rewritten_query)
    VALUES (original, rewritten);

    RETURN rewritten;
END;
$$ LANGUAGE plpgsql;

-- 向量检索
CREATE OR REPLACE FUNCTION rag.vector_search(query TEXT, top_k INT DEFAULT 5)
RETURNS TABLE(id BIGINT, content TEXT, score FLOAT) AS $$
BEGIN
    RETURN QUERY
    SELECT d.id, d.content,
           1 - (d.embedding <=> tencentdb_ai.get_embedding('bge-m3', query)) AS score
    FROM rag.documents d
    ORDER BY d.embedding <=> tencentdb_ai.get_embedding('bge-m3', query)
    LIMIT top_k;
END;
$$ LANGUAGE plpgsql;

-- 关键词检索（BM25 风格）
CREATE OR REPLACE FUNCTION rag.keyword_search(keywords TEXT, top_k INT DEFAULT 5)
RETURNS TABLE(id BIGINT, content TEXT, score FLOAT) AS $$
BEGIN
    RETURN QUERY
    SELECT d.id, d.content,
           ts_rank(to_tsvector('chinese', d.content), plainto_tsquery('chinese', keywords))::FLOAT AS score
    FROM rag.documents d
    WHERE to_tsvector('chinese', d.content) @@ plainto_tsquery('chinese', keywords)
    ORDER BY score DESC
    LIMIT top_k;
END;
$$ LANGUAGE plpgsql;

-- 混合检索（RRF 融合）
CREATE OR REPLACE FUNCTION rag.hybrid_search(query TEXT, top_k INT DEFAULT 5)
RETURNS TABLE(id BIGINT, content TEXT, rrf_score FLOAT) AS $$
BEGIN
    RETURN QUERY
    WITH vec AS (
        SELECT id, content, ROW_NUMBER() OVER (ORDER BY score DESC) AS rank
        FROM rag.vector_search(query, 20)
    ),
    kw AS (
        SELECT id, content, ROW_NUMBER() OVER (ORDER BY score DESC) AS rank
        FROM rag.keyword_search(query, 20)
    )
    SELECT COALESCE(v.id, k.id),
           COALESCE(v.content, k.content),
           (COALESCE(1.0/(60+v.rank), 0) + COALESCE(1.0/(60+k.rank), 0))::FLOAT AS rrf_score
    FROM vec v FULL OUTER JOIN kw k ON v.id = k.id
    ORDER BY rrf_score DESC
    LIMIT top_k;
END;
$$ LANGUAGE plpgsql;

-- 完整的 Agentic RAG 函数
CREATE OR REPLACE FUNCTION rag.agentic_answer(question TEXT)
RETURNS TEXT AS $$
DECLARE
    rewritten_query TEXT;
    context TEXT;
    answer TEXT;
BEGIN
    -- Step 1: 查询改写
    rewritten_query := rag.rewrite_query(question);

    -- Step 2: 混合检索
    SELECT string_agg(content, E'\\n---\\n') INTO context
    FROM rag.hybrid_search(rewritten_query, 5);

    -- Step 3: 生成回答
    SELECT tencentdb_ai.chat_completions(
        '<your-llm-model>',
        '基于以下参考资料回答用户问题。如果资料不足，明确说明。' ||
        E'\\n\\n参考资料：\\n' || COALESCE(context, '无相关资料') ||
        E'\\n\\n用户问题：' || question
    ) INTO answer;

    RETURN answer;
END;
$$ LANGUAGE plpgsql;