分支型 RAG Workflow

You are a reasoning assistant with the ability to perform web searches to help you answer the user’s question
accurately. You have special tools:
To perform a search: write <|begin_search_query|> your query here <|end_search_query|>.
Then, the system will search and analyze relevant web pages, then provide you with helpful information in the
format <|begin_search_result|> ...search results... <|end_search_result|>.
You can repeat the search process multiple times if necessary. The maximum number of search attempts is
limited to {{MAX_SEARCH_LIMIT}}.
Once you have all the information you need, continue your reasoning.
Example:
Question: “...”
Assistant thinking steps:
- I might need to look up details about ...
Assistant:
<|begin_search_query|>...<|end_search_query|>
(System returns processed information from relevant web pages)
Assistant continues reasoning with the new information...
Remember:
- Use <|begin_search_query|> to request a web search and end with <|end_search_query|>.
- When done searching, continue your reasoning.
Please answer the following question. You should think step by step to solve it.\n\n
Provide your final answer in the format \\boxed{YOUR_ANSWER}.\n\n
Question:\n{{question}}\n\n

@app.prompt(output="q_ls, template -> prompt_ls")
def search_o1_init(
    q_ls: List[str],
    template: str | Path,
) -> List[PromptMessage]:
    template: Template = load_prompt_template(template)
    # 目前这个变量写固定了
    MAX_SEARCH_LIMIT = 10
    ret = []
    for q in q_ls:
        p = template.render(question=q, MAX_SEARCH_LIMIT=MAX_SEARCH_LIMIT)
        ret.append(p)
    return ret

@app.prompt(
    output="prompt_ls,extract_query_list,ret_psg,reasoning_indoc_template->prompt_ls"
)
def searcho1_reasoning_indocument(
    prompt_ls: List[PromptMessage],
    extract_query_list: List[str],
    ret_psg: List[str | Any],
    template: str | Path,
) -> List[PromptMessage]:
    template: Template = load_prompt_template(template)
    ret = []
    for prompt, squery, psg in zip(prompt_ls, extract_query_list, ret_psg):
        # passages = [psg[index]["segment"] for index in range(min(5, len(psg)))]
        passages = psg[:3]
        passage_text = "\n".join(passages)
        _pro = prompt.content.text
        p = template.render(
            prev_reasoning=_pro, search_query=squery, document=passage_text
        )
        ret.append(p)
    return ret

template: prompt/qa_boxed.jinja
reasoning_indoc_template: prompt/qa_boxed.jinja

@app.prompt(output="prompt_ls,ans_ls->prompt_ls")
def search_o1_insert(
    prompt_ls: List[PromptMessage],
    ans_ls: List[str],
) -> List[PromptMessage]:
    ret = []
    for prompt, ans in zip(prompt_ls, ans_ls):
        _pro = prompt.content.text
        p = _pro + "<|begin_search_result|>" + ans + "<|end_search_result|>"
        ret.append(p)
    return ret

@app.tool(output="ans_ls->ans_ls")
def search_o1_check(ans_ls: List[str]) -> Dict[str, List[Dict[str, str]]]:
    def get_eos(text):
        # 如果包含特殊终止 token，则直接返回 True
        if "<|im_end|>" in text:
            return True
        elif "<|end_search_query|>" in text:
            return False

    ans_ls = [
        {
            "data": answer,
            "state": "stop" if get_eos(answer) else "retrieve",
        }
        for answer in ans_ls
    ]
    return {"ans_ls": ans_ls}

@app.tool(output="ans_ls->extract_query_list")
def search_r1_query_extract(ans_ls: List[str]) -> Dict[str, List[str]]:

    def get_query(text):
        import re

        # 匹配最后一个 <search> 及其后面的内容
        pattern = re.compile(r"<search>([^<]*)", re.DOTALL)
        matches = pattern.findall(text)

        if matches:
            query = matches[-1].strip()
            if not query.endswith("?"):
                query += "?"
            return query
        else:
            return "There is no query."

    query = [get_query(answer) for answer in ans_ls]

    return {"extract_query_list": query}

# Search-o1 demo

# MCP Servers
servers:
  benchmark: servers/benchmark
  generation: servers/generation
  retriever: servers/retriever
  prompt: servers/prompt
  evaluation: servers/evaluation
  router: servers/router
  custom: servers/custom

pipeline:
# 1. 加载数据集
- benchmark.get_data
# 2. 初始化 (构造 instruction + 问题)
- prompt.search_o1_init
# 3. 首轮初始生成 (得到初始 reasoning / 可能的第一个 <SEARCH_QUERY>)
- generation.generate
# 4. 循环：多轮搜索 + 推理
- loop:
    times: 3
    steps:
    - branch:
        router:
        - router.search_o1_check # 输出 e.g. status=incomplete/complete
        branches:
          retrieve:
          # 4.1 抽取最新需要执行的查询（若没有则在 need_search_check 中应返回 complete）
          - custom.search_o1_query_extract
          # 4.3 执行检索 (将抽取的 query 列表传入)
          - retriever.retriever_deploy_search:
              input:
                query_list: extract_query_list
          # 4.4 reasoning_indocument
          - prompt.searcho1_reasoning_indocument
          - generation.generate
          # 4.5 追加 <BEGIN_SEARCH_RESULT> …）
          - prompt.search_o1_insert
          # 4.6 生成新一轮 reasoning / 可能新的查询
          - generation.generate
          stop: []
# 5. 评估 (使用最终答案 / reasoning)
- evaluation.evaluate:
    input:
      pred_ls: ans_ls

benchmark:
  benchmark:
    key_map:
      gt_ls: golden_answers
      q_ls: question
    limit: 2
    name: asqa
    path: data/sample_asqa_5.jsonl
custom: {}
evaluation:
  metrics:
  - acc
  - f1
  - em
  - coverem
  - stringem
  - rouge-1
  - rouge-2
  - rouge-l
  save_path: output/asqa.json
generation:
  base_url: http://localhost:8000/v1
  model_name: openbmb/MiniCPM4-8B
  sampling_params:
    extra_body:
      chat_template_kwargs:
        enable_thinking: false
      include_stop_str_in_output: true
      top_k: 20
    max_tokens: 2048
    temperature: 0.7
    top_p: 0.8
prompt:
  reasoning_indoc_template: prompt/search_o1_refinement.jinja
  template: prompt/search_o1_reasoning.jinja
retriever:
  query_instruction: 'Query: '
  retriever_url: http://localhost:8080
  top_k: 5