rag_agent.ail

# RAG Agent 完整示例
# 展示：with context、parallel、loop、@plan、@extract、@validate、@eval、retry、memory

use tool vector_search(query: str, top_k: int) -> list[Document]
use tool keyword_search(query: str) -> list[Document]
use tool rerank(query: str, docs: list[Document]) -> list[Document]

type Document:
    id:      str
    title:   str
    content: str
    score:   float = 0.0

type QueryInfo:
    intent:     str
    keywords:   list[str]
    multi_step: bool

def rag_answer(user_query: str, max_retrieval_rounds=3) -> (str, list[str]):
    with context() as ctx:

        # Step 1：理解问题
        prompt analyze_query = """
            分析用户问题：{user_query}
            请输出：1.核心意图 2.关键词列表 3.是否需要多步推理
        """
        analysis = @ask(analyze_query)
        info = @extract(analysis, type=QueryInfo)

        # Step 2：并行检索（向量 + 关键词同时进行）
        vec_query = @ask("将意图「{info.intent}」改写为适合向量检索的语句")
        kw_query  = @ask("从关键词「{info.keywords}」中提取最重要的 3 个")

        vec_docs, kw_docs = parallel:
            vec_docs = vector_search(vec_query, top_k=10)
            kw_docs  = keyword_search(kw_query)

        top_docs = rerank(user_query, vec_docs + kw_docs)[:5]

        # Step 3：文档不够时继续补充检索
        prompt find_gap = """
            当前文档还缺少哪些信息才能回答「{user_query}」
        """
        loop max=max_retrieval_rounds until @judge("以下文档能否充分回答「{user_query}」：{top_docs}"):
            gap      = @ask(find_gap)
            new_docs = vector_search(@ask("根据缺口生成补充检索语句：{gap}"), top_k=5)
            top_docs = rerank(user_query, top_docs + new_docs)[:5]

        # Step 4：推理生成（支持单步和多步）
        if info.multi_step:
            steps = @plan("针对「{user_query}」制定推理步骤，参考：{top_docs}")
            result = ""
            for step in steps:
                result = @ask("执行「{step}」，已有结论：{result}，参考：{top_docs}")
            final_context = result
        else:
            final_context = @ask("整理文档中与「{user_query}」相关的核心内容：{top_docs}")

        # Step 5：生成答案，校验质量
        prompt generate_answer = """
            基于以下内容回答用户问题。
            问题：{user_query}
            参考内容：{final_context}
            要求：有理有据，引用来源，不要编造
        """
        retry max=3:
            answer = @ask(generate_answer)
            @validate(answer, "回答必须基于参考内容，不得出现参考内容中没有的事实")

        quality = @eval(answer, type={"relevance": float, "completeness": float})
        if quality.relevance < 0.7 or quality.completeness < 0.7:
            answer = @ask("改进以上回答，使其更相关、更完整，当前评分：{quality}")

        # Step 6：提取引用来源，保存记忆
        citations = @extract(answer, "提取引用来源，与以下文档对照", type=list[str], context=top_docs)
        memory.save(answer, key=f"rag:{user_query[:20]}", tags=["history"])

    return answer, citations