feat(solver): support kag thinker

KAG_VERSION

@@ -1 +1 @@
-0.8.0
+0.8.0

kag/common/utils.py

@@ -463,9 +463,17 @@ def resolve_instance(
 
 
 def extract_tag_content(text):
-    # 匹配<tag>和</tag>之间的内容，支持任意标签名
-    matches = re.findall(r"<([^>]+)>(.*?)</\1>", text, flags=re.DOTALL)
-    return [(tag, content.strip()) for tag, content in matches]
+    pattern = r"<(\w+)\b[^>]*>(.*?)</\1>|<(\w+)\b[^>]*>([^<]*)|([^<]+)"
+    results = []
+    for match in re.finditer(pattern, text, re.DOTALL):
+        tag1, content1, tag2, content2, raw_text = match.groups()
+        if tag1:
+            results.append((tag1, content1))  # 保留原始内容（含空格）
+        elif tag2:
+            results.append((tag2, content2))  # 保留原始内容（含空格）
+        elif raw_text:
+            results.append(("", raw_text))  # 保留原始空格
+    return results
 
 
 def extract_specific_tag_content(text, tag):

kag/solver/executor/math/py_based_math_executor.py

@@ -131,17 +131,19 @@ def invoke(self, query: str, task: Task, context: Context, **kwargs):
         )
 
         parent_results = format_task_dep_context(task.parents)
-        parent_results = "\n".join(parent_results)
+        coder_content = context.kwargs.get("planner_thought", "") + "\n\n".join(
+            parent_results
+        )
 
-        parent_results += "\n\n" + contents
+        coder_content += "\n\n" + contents
         tries = self.tries
         error = None
 
         while tries > 0:
             tries -= 1
             rst, error, code = self.run_once(
                 math_query,
-                parent_results,
+                coder_content,
                 error,
                 segment_name=tag_id,
                 tag_name=f"{task_query}_code_generator",

kag/solver/executor/retriever/kag_hybrid_retrieval_executor.py

@@ -42,6 +42,15 @@
 logger = logging.getLogger()
 
 
+def _wrapped_invoke(retriever, task, context, segment_name, kwargs):
+    start_time = time.time()
+    output = retriever.invoke(
+        task, context=context, segment_name=segment_name, **kwargs
+    )
+    elapsed_time = time.time() - start_time
+    return output, elapsed_time
+
+
 @ExecutorABC.register("kag_hybrid_retrieval_executor")
 class KAGHybridRetrievalExecutor(ExecutorABC):
     def __init__(
@@ -76,6 +85,7 @@ def __init__(
         self.context_select_prompt = context_select_prompt or PromptABC.from_config(
             {"type": "context_select_prompt"}
         )
+        self.with_llm_select = kwargs.get("with_llm_select", True)
 
     @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1))
     def context_select_call(self, variables):
@@ -152,22 +162,30 @@ def do_retrieval(
                         "FINISH",
                         component_name=retriever.name,
                     )
-
+                    # Record start time before submitting the task
+                    start_time = time.time()
                     # Prepare function and submit to thread pool
                     func = partial(
-                        retriever.invoke,
+                        _wrapped_invoke,
+                        retriever,
                         task,
-                        context=context,
-                        segment_name=tag_id,
-                        **kwargs,
+                        context,
+                        tag_id,
+                        kwargs.copy(),
                     )
                     future = executor.submit(func)
+                    # Save future, retriever, and start_time together
                     futures.append((future, retriever))
 
                 # Collect results from each future
                 for future, retriever in futures:
                     try:
-                        output = future.result()  # Wait for result
+                        output, elapsed_time = future.result()  # Wait for result
+
+                        # Log the elapsed time for this retriever
+                        logger.info(
+                            f"Retriever {retriever.name} executed in {elapsed_time:.2f} seconds"
+                        )
                         outputs.append(output)
 
                         # Log data report after successful execution
@@ -241,13 +259,18 @@ def do_summary(
         selected_rel = list(set(selected_rel))
         formatted_docs = [str(rel) for rel in selected_rel]
         if retrieved_data.chunks:
-            try:
-                selected_chunks = self.context_select(task_query, retrieved_data.chunks)
-            except Exception as e:
-                logger.warning(
-                    f"select context failed {e}, we use default top 10 to summary",
-                    exc_info=True,
-                )
+            if self.with_llm_select:
+                try:
+                    selected_chunks = self.context_select(
+                        task_query, retrieved_data.chunks
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"select context failed {e}, we use default top 10 to summary",
+                        exc_info=True,
+                    )
+                    selected_chunks = retrieved_data.chunks[:10]
+            else:
                 selected_chunks = retrieved_data.chunks[:10]
             for doc in selected_chunks:
                 formatted_docs.append(f"{doc.content}")
@@ -280,69 +303,82 @@ def invoke(self, query, task, context: Context, **kwargs) -> RetrieverOutput:
         task_query = task.arguments["query"]
 
         tag_id = f"{task_query}_begin_task"
-        self.report_content(reporter, "thinker", tag_id, "", "FINISH", step=task.name)
+        self.report_content(reporter, "thinker", tag_id, "", "INIT", step=task.name)
         try:
-            retrieved_data = self.do_main(task_query, tag_id, task, context, **kwargs)
-        except Exception as e:
-            logger.warning(f"kag hybrid retrieval failed! {e}", exc_info=True)
-            retrieved_data = RetrieverOutput(
-                retriever_method=self.schema().get("name", ""), err_msg=str(e)
-            )
+            try:
+                retrieved_data = self.do_main(
+                    task_query, tag_id, task, context, **kwargs
+                )
+            except Exception as e:
+                logger.warning(f"kag hybrid retrieval failed! {e}", exc_info=True)
+                retrieved_data = RetrieverOutput(
+                    retriever_method=self.schema().get("name", ""), err_msg=str(e)
+                )
 
-        self.report_content(
-            reporter,
-            "reference",
-            f"{task_query}_kag_retriever_result",
-            retrieved_data,
-            "FINISH",
-        )
+            self.report_content(
+                reporter,
+                "reference",
+                f"{task_query}_kag_retriever_result",
+                retrieved_data,
+                "FINISH",
+            )
 
-        retrieved_data.task = task
-        logical_node = task.arguments.get("logic_form_node", None)
-        if (
-            logical_node
-            and isinstance(logical_node, GetSPONode)
-            and retrieved_data.summary
-        ):
-            if isinstance(retrieved_data.summary, str):
-                target_answer = retrieved_data.summary.split("Answer:")[-1].strip()
-                s_entities = context.variables_graph.get_entity_by_alias(
-                    logical_node.s.alias_name
+            retrieved_data.task = task
+            logical_node = task.arguments.get("logic_form_node", None)
+            if (
+                logical_node
+                and isinstance(logical_node, GetSPONode)
+                and retrieved_data.summary
+            ):
+                if isinstance(retrieved_data.summary, str):
+                    target_answer = retrieved_data.summary.split("Answer:")[-1].strip()
+                    s_entities = context.variables_graph.get_entity_by_alias(
+                        logical_node.s.alias_name
+                    )
+                    if (
+                        not s_entities
+                        and not logical_node.s.get_mention_name()
+                        and isinstance(logical_node.s, SPOEntity)
+                    ):
+                        logical_node.s.entity_name = target_answer
+                        context.kwargs[logical_node.s.alias_name] = logical_node.s
+                    o_entities = context.variables_graph.get_entity_by_alias(
+                        logical_node.o.alias_name
+                    )
+                    if (
+                        not o_entities
+                        and not logical_node.o.get_mention_name()
+                        and isinstance(logical_node.o, SPOEntity)
+                    ):
+                        logical_node.o.entity_name = target_answer
+                        context.kwargs[logical_node.o.alias_name] = logical_node.o
+
+                context.variables_graph.add_answered_alias(
+                    logical_node.s.alias_name.alias_name, retrieved_data.summary
                 )
-                if (
-                    not s_entities
-                    and not logical_node.s.get_mention_name()
-                    and isinstance(logical_node.s, SPOEntity)
-                ):
-                    logical_node.s.entity_name = target_answer
-                    context.kwargs[logical_node.s.alias_name] = logical_node.s
-                o_entities = context.variables_graph.get_entity_by_alias(
-                    logical_node.o.alias_name
+                context.variables_graph.add_answered_alias(
+                    logical_node.p.alias_name.alias_name, retrieved_data.summary
                 )
-                if (
-                    not o_entities
-                    and not logical_node.o.get_mention_name()
-                    and isinstance(logical_node.o, SPOEntity)
-                ):
-                    logical_node.o.entity_name = target_answer
-                    context.kwargs[logical_node.o.alias_name] = logical_node.o
-
-            context.variables_graph.add_answered_alias(
-                logical_node.s.alias_name.alias_name, retrieved_data.summary
-            )
-            context.variables_graph.add_answered_alias(
-                logical_node.p.alias_name.alias_name, retrieved_data.summary
+                context.variables_graph.add_answered_alias(
+                    logical_node.o.alias_name.alias_name, retrieved_data.summary
+                )
+
+            task.update_result(retrieved_data)
+            logger.debug(
+                f"kag hybrid retrieval {task_query} cost={time.time() - start_time}"
             )
-            context.variables_graph.add_answered_alias(
-                logical_node.o.alias_name.alias_name, retrieved_data.summary
+            return retrieved_data
+        finally:
+            self.report_content(
+                reporter,
+                "thinker",
+                tag_id,
+                "",
+                "FINISH",
+                step=task.name,
+                overwrite=False,
             )
 
-        task.update_result(retrieved_data)
-        logger.debug(
-            f"kag hybrid retrieval {task_query} cost={time.time() - start_time}"
-        )
-        return retrieved_data
-
     def schema(self) -> dict:
         """Function schema definition for OpenAI Function Calling
 
@@ -403,7 +439,7 @@ def do_data_report(
                 node_type=chunk.properties.get("__labels__"),
             )
             entity_prop = dict(chunk.properties) if chunk.properties else {}
-            entity_prop["content"] = chunk.content
+            entity_prop["content"] = f"{chunk.content[:10]}..."
             entity_prop["score"] = chunk.score
             entity.prop = Prop.from_dict(entity_prop, "Chunk", None)
             chunk_graph.append(entity)

kag/solver/main_solver.py

@@ -140,8 +140,6 @@ def get_pipeline_conf(use_pipeline_name, config):
             raise RuntimeError("mcpServers not found in config.")
         default_solver_pipeline["executors"] = mcp_executors
 
-    # update KAG_CONFIG
-    KAG_CONFIG.update_conf(default_pipeline_conf)
     return default_solver_pipeline
 
 
@@ -167,8 +165,11 @@ async def do_qa_pipeline(
                     f"Knowledge base with id {kb_project_id} not found in qa_config['kb']"
                 )
                 continue
-
-            for index_name in matched_kb.get("index_list", []):
+            index_list = matched_kb.get("index_list", [])
+            if use_pipeline in ["default_pipeline"]:
+                # we only use chunk index
+                index_list = ["chunk_index"]
+            for index_name in index_list:
                 index_manager = KAGIndexManager.from_config(
                     {
                         "type": index_name,
@@ -339,7 +340,7 @@ class SolverMain:
     def invoke(
         self,
         project_id: int,
-        task_id: int,
+        task_id,
         query: str,
         session_id: str = "0",
         is_report=True,

kag/solver/pipelineconf/naive_rag.yaml

@@ -3,20 +3,17 @@ pipeline_name: default_pipeline
 
 #------------kag-solver configuration start----------------#
 
-
-chunk_retrieved_executor: &chunk_retrieved_executor_conf
-  type: chunk_retrieved_executor
-  top_k: 10
-  retriever:
-    type: vector_chunk_retriever
-    score_threshold: 0.65
-    vectorize_model: "{vectorize_model}"
-
+kag_retriever_executor: &kag_retriever_executor_conf
+  type: kag_hybrid_retrieval_executor
+  retrievers: "{retrievers}"
+  merger:
+    type: kag_merger
+  enable_summary:  false
 
 solver_pipeline:
   type: naive_rag_pipeline
   executors:
-    - *chunk_retrieved_executor_conf
+    - *kag_retriever_executor_conf
   generator:
     type: llm_index_generator
     llm_client: "{chat_llm}"

kag/solver/planner/kag_model_planner.py

@@ -186,6 +186,7 @@ async def ainvoke(self, query, **kwargs) -> List[Task]:
             .replace("</answer>", "")
             .strip()
         )
+        context.kwargs["planner_thought"] = logic_form_response
 
         sub_queries, logic_forms = parse_logic_form_with_str(logic_form_str)
         logic_forms = self.logic_node_parser.parse_logic_form_set(