fix: reimplemented search-o1 (#158)

Author: mssssss123

examples/search_o1.yaml

@@ -15,6 +15,7 @@ pipeline:
 - benchmark.get_data
 - retriever.retriever_init
 - generation.generation_init
+- custom.search_o1_init_list
 - prompt.search_o1_init
 - generation.generate
 - loop:
@@ -29,8 +30,12 @@ pipeline:
           - retriever.retriever_search:
               input:
                 query_list: extract_query_list
-          - prompt.searcho1_reasoning_indocument
+          - custom.search_o1_reasoning_extract
+          - custom.search_o1_combine_list
+          - prompt.search_o1_reasoning_indocument
           - generation.generate
+          - custom.search_o1_extract_final_information
+          - custom.search_o1_combine_final_information
           - prompt.search_o1_insert
           - generation.generate
           stop: []   

servers/custom/src/custom.py

@@ -119,6 +119,46 @@ def ircot_extract_ans(ans_ls: List[str]) -> Dict[str, List[str]]:
     return {"pred_ls": ret}
 
 
+@app.tool(output="q_ls->total_subq_list,total_reason_list,total_final_info_list")
+def search_o1_init_list(q_ls: List[str]) -> Dict[str, List[Any]]:
+    n = len(q_ls)
+
+    return {
+        "total_subq_list": [["<PAD>"] for _ in range(n)],
+        "total_reason_list": [["<PAD>"] for _ in range(n)],
+        "total_final_info_list": [["<PAD>"] for _ in range(n)],
+    }
+
+@app.tool(
+    output="total_subq_list, extract_query_list, total_reason_list, extract_reason_list"
+           "->total_subq_list, total_reason_list"
+)
+def search_o1_combine_list(
+    total_subq_list: List[List[Any]],
+    extract_query_list: List[str],
+    total_reason_list: List[List[Any]],
+    extract_reason_list: List[str],
+) -> Dict[str, List[Any]]:
+    
+    PAD = "<PAD>"
+
+    for q, bucket in zip(extract_query_list, total_subq_list):
+        if len(bucket) == 1 and bucket[0] == PAD:
+            bucket[0] = q            
+        else:
+            bucket.append(q)
+
+    for c, bucket in zip(extract_reason_list, total_reason_list):
+        if len(bucket) == 1 and bucket[0] == PAD:
+            bucket[0] = c           
+        else:
+            bucket.append(c)
+
+    return {
+        "total_subq_list": total_subq_list,
+        "total_reason_list": total_reason_list,
+    }
+
 @app.tool(output="ans_ls->extract_query_list")
 def search_o1_query_extract(ans_ls: List[str]) -> Dict[str, List[str]]:
     import re
@@ -139,6 +179,58 @@ def get_query(text):
 
     return {"extract_query_list": query}
 
+@app.tool(output="ans_ls->extract_reason_list")
+def search_o1_reasoning_extract(ans_ls: List[str]) -> Dict[str, List[str]]:
+
+    BEGIN = "<|begin_search_query|>"
+
+    def get_content_before(text):
+        if BEGIN not in text:
+            return text.strip()
+        
+
+        return text.split(BEGIN, 1)[0].strip()
+
+    content_list = [get_content_before(answer) for answer in ans_ls]
+
+    return {"extract_reason_list": content_list}
+
+@app.tool(output="ans_ls->extract_final_infor_list")
+def search_o1_extract_final_information(ans_ls: List[str]) -> Dict[str, List[str]]:
+
+    BEGIN = "**Final Information**"
+
+    def get_content_after(text):
+        if BEGIN not in text:
+            return ""
+    
+        return BEGIN + "\n" + text.split(BEGIN, 1)[1].strip()
+
+    content_list = [get_content_after(answer) for answer in ans_ls]
+
+    return {"extract_final_infor_list": content_list}
+
+@app.tool(output="total_final_info_list, extract_final_infor_list->total_final_info_list")
+def search_o1_combine_final_information(
+    total_final_info_list: List[List[str]],
+    extract_final_infor_list: List[str],
+) -> Dict[str, List[Any]]:
+    
+    PAD = "<PAD>"
+
+    for c, bucket in zip(extract_final_infor_list, total_final_info_list):
+        if len(bucket) == 1 and bucket[0] == PAD:
+            bucket[0] = c           
+        else:
+            bucket.append(c)
+
+    app.logger.warning(f"len total_final_info_list: {len(total_final_info_list)}")
+    app.logger.warning(f"total_final_info_list: {total_final_info_list}")
+
+    return {
+        "total_final_info_list": total_final_info_list,
+    }
+
 @app.tool(output="temp_psg,ret_psg->ret_psg")
 def merge_passages(
     temp_psg: List[str | Any],

servers/prompt/src/prompt.py

@@ -288,41 +288,71 @@ def search_o1_init(
         ret.append(p)
     return ret
 
-
 @app.prompt(
-    output="prompt_ls,extract_query_list,ret_psg,searcho1_refine_template->prompt_ls"
+    output="extract_query_list, ret_psg, total_reason_list, searcho1_refine_template -> prompt_ls"
 )
-def searcho1_reasoning_indocument(
-    prompt_ls: List[PromptMessage],
-    extract_query_list: List[str],
-    ret_psg: List[str | Any],
+def search_o1_reasoning_indocument(
+    extract_query_list: List[str], 
+    ret_psg: List[List[str]],       
+    total_reason_list: List[List[str]], 
     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):
-        passage_text = "\n".join(psg)
-        _pro = prompt.content.text
+
+    for squery, psg_list, history_steps in zip(extract_query_list, ret_psg, total_reason_list):
+
+        passage_text = "\n".join(psg_list)
+
+        if len(history_steps) <= 3:
+            selected_history = history_steps[:]  
+        else:
+            selected_history = [history_steps[0]] + history_steps[-3:]
+
+        formatted_history_parts = [
+            f"Step {i}: {reason}"
+            for i, reason in enumerate(selected_history, 1)
+        ]
+        formatted_history_str = "\n\n".join(formatted_history_parts)
+
         p = template.render(
-            prev_reasoning=_pro, search_query=squery, document=passage_text
+            prev_reasoning=formatted_history_str, 
+            search_query=squery, 
+            document=passage_text
         )
         ret.append(p)
-    return ret
 
+    return ret
 
-@app.prompt(output="prompt_ls,ans_ls->prompt_ls")
+@app.prompt(output="q_ls,total_subq_list,total_final_info_list,searcho1_reasoning_template->prompt_ls") 
 def search_o1_insert(
-    prompt_ls: List[PromptMessage],
-    ans_ls: List[str],
+    q_ls: List[str],
+    total_subq_list: List[List[str]], 
+    total_final_info_list: List[List[str]],
+    template: str | Path,
 ) -> List[PromptMessage]:
+    template: Template = load_prompt_template(template)
+    prompt_ls = []
+    for q in q_ls:
+        p = template.render(question=q)
+        prompt_ls.append(p)
+    
     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)
+    for prompt, sub_queries, sub_reasons in zip(prompt_ls, total_subq_list, total_final_info_list):
+        
+        
+        for query, reason in zip(sub_queries, sub_reasons):
+            part = (
+                "<|begin_search_query|>" + str(query) + "<|end_search_query|>" + 
+                '\n' + 
+                "<|begin_search_result|>" + str(reason) + "<|end_search_result|>"
+            )
+            prompt += part
+        
+        ret.append(prompt)
+        
     return ret
 
-
 # prompt for loop and branch demo
 @app.prompt(output="q_ls,ret_psg,gen_subq_template->prompt_ls")
 def gen_subq(

servers/router/src/router.py

@@ -1,5 +1,4 @@
-from typing import List, Dict
-
+from typing import List, Dict, Any
 from ultrarag.server import UltraRAG_MCP_Server
 
 
@@ -105,24 +104,49 @@ def get_eos(text):
     return {"ans_ls": ans_ls}
 
 
-@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):
+@app.tool(
+    output=("ans_ls,q_ls,total_subq_list,total_reason_list,total_final_info_list->ans_ls,q_ls,total_subq_list,total_reason_list,total_final_info_list")
+)
+def search_o1_check(
+    ans_ls: List[str],
+    q_ls: List[str],
+    total_subq_list: List[List[Any]],
+    total_reason_list: List[List[Any]],
+    total_final_info_list: List[List[Any]],
+) -> Dict[str, List[Dict[str, Any]]]:
+
+    def get_eos(text: str) -> bool:
         if "<|im_end|>" in text:
-            return True               
+            return True
         elif "<|end_search_query|>" in text:
-            return False              
+            return False
         else:
-            return True  
+            return True
 
-    ans_ls = [
-        {
-            "data": answer,
-            "state": "stop" if get_eos(answer) else "retrieve",
-        }
-        for answer in ans_ls
-    ]
-    return {"ans_ls": ans_ls}
+    ans_out: List[Dict[str, Any]] = []
+    q_out: List[Dict[str, Any]] = []
+    subq_out: List[Dict[str, Any]] = []
+    reason_out: List[Dict[str, Any]] = []
+    info_out: List[Dict[str, Any]] = []
+
+    for ans, q, subq, reason, info in zip(
+        ans_ls, q_ls, total_subq_list, total_reason_list, total_final_info_list
+    ):
+        state = "stop" if get_eos(ans) else "retrieve"
+
+        ans_out.append({"data": ans, "state": state})
+        q_out.append({"data": q, "state": state})
+        subq_out.append({"data": subq, "state": state})
+        reason_out.append({"data": reason, "state": state})
+        info_out.append({"data": info, "state": state})
+
+    return {
+        "ans_ls": ans_out,
+        "q_ls": q_out,
+        "total_subq_list": subq_out,
+        "total_reason_list": reason_out,
+        "total_final_info_list": info_out,
+    }
 
 
 @app.tool(output="ans_ls->ans_ls")