fix main solver pipeline

Author: royzhao

kag/common/conf.py

@@ -208,10 +208,16 @@ def update_conf(self, configs: dict):
 KAG_QA_TASK_CONFIG stores per-task configuration and should be cleaned up after use.
 """
 KAG_QA_TASK_CONFIG  = knext.common.cache.LinkCache(maxsize=100, ttl=300)
-KAG_QA_TASK_CONFIG_LOCK = threading.Lock()
-
 
 class KAGConfigAccessor:
+    @staticmethod
+    def get_default_config():
+        if KAG_CONFIG.global_config.project_id:
+            return KAG_CONFIG
+        for k in KAG_QA_TASK_CONFIG.cache.keys():
+            return KAG_QA_TASK_CONFIG.get(k)
+        return KAG_CONFIG
+
     @staticmethod
     def get_config(task_with_kb_id=None) -> KAGConfigMgr:
         """
@@ -224,9 +230,8 @@ def get_config(task_with_kb_id=None) -> KAGConfigMgr:
         :return: Corresponding configuration object
         """
         if task_with_kb_id is not None:
-            with KAG_QA_TASK_CONFIG_LOCK:
-                return KAG_QA_TASK_CONFIG.get(task_with_kb_id)
-        return KAG_CONFIG
+            return KAG_QA_TASK_CONFIG.get(task_with_kb_id)
+        return KAGConfigAccessor.get_default_config()
 
     @staticmethod
     def set_task_config(task_with_kb_id, config: KAGConfigMgr):
@@ -236,8 +241,7 @@ def set_task_config(task_with_kb_id, config: KAGConfigMgr):
         :param task_with_kb_id: Task ID
         :param config: Configuration object to store
         """
-        with KAG_QA_TASK_CONFIG_LOCK:
-            KAG_QA_TASK_CONFIG.put(task_with_kb_id, config)
+        KAG_QA_TASK_CONFIG.put(task_with_kb_id, config)
 
 
 def init_env(config_file: str = None):

kag/solver/executor/retriever/local_knowledge_base/chunk_retrieved_executor.py

@@ -2,7 +2,7 @@
 import time
 from typing import Any, Optional
 
-from kag.interface import ExecutorABC, ToolABC
+from kag.interface import ExecutorABC, RetrieverABC
 from kag.interface.solver.reporter_abc import ReporterABC
 from kag.solver.executor.retriever.local_knowledge_base.kag_retriever.kag_hybrid_executor import (
     KAGRetrievedResponse,
@@ -16,7 +16,7 @@
 
 @ExecutorABC.register("chunk_retrieved_executor")
 class ChunkRetrievedExecutor(ExecutorABC):
-    def __init__(self, top_k, retriever: ToolABC, **kwargs):
+    def __init__(self, top_k, retriever: RetrieverABC, **kwargs):
         super().__init__(**kwargs)
         self.retriever = retriever
         self.top_k = top_k
@@ -43,23 +43,11 @@ def invoke(self, query: str, task: Any, context: dict, **kwargs):
             "FINISH",
             overwrite=False,
         )
-        retrieved_result = self.retriever.invoke(query=task_query, top_k=self.top_k)
+        retrieved_result = self.retriever.invoke(task, context=context, **kwargs)
 
         # Log the retrieved results
         logger.debug(f"Retrieved results: {retrieved_result}")
 
-        chunk_datas = []
-        for k, v in retrieved_result.items():
-            chunk_datas.append(
-                ChunkData(
-                    content=v["content"],
-                    title=v["name"],
-                    chunk_id=k,
-                    score=v["score"],
-                    properties=v,
-                )
-            )
-        kag_response.chunk_datas = chunk_datas
         self.report_content(
             reporter,
             "reference",
@@ -72,7 +60,7 @@ def invoke(self, query: str, task: Any, context: dict, **kwargs):
             reporter,
             f"{task_query}_begin_kag_retriever",
             f"{task_query}_end_kag_retriever",
-            f"{len(chunk_datas)}",
+            f"{len(retrieved_result.chunks)}",
             "FINISH",
         )
 
@@ -81,8 +69,7 @@ def invoke(self, query: str, task: Any, context: dict, **kwargs):
         logger.info(
             f"Finished retrieval process for query: {task_query}. Duration: {end_time - start_time} bytes"
         )
-        kag_response.summary = "retrieved by local knowledgebase"
-        store_results(task, kag_response)
+        task.update_result(retrieved_result)
 
     def schema(self) -> dict:
         """Function schema definition for OpenAI Function Calling

kag/solver/main_solver.py

@@ -149,32 +149,6 @@ def is_chinese(text):
     chinese_pattern = re.compile(r"[\u4e00-\u9fff]+")
     return bool(chinese_pattern.search(text))
 
-
-async def do_index_pipeline(query, qa_config, reporter):
-    if "chat" not in qa_config or "index_list" not in qa_config["chat"]:
-        raise RuntimeError("chat or index_list not found in qa_config.")
-    index_names = qa_config.get("chat", {}).get("index_list", [])
-    retriever_configs = []
-    for index_name in index_names:
-        try:
-            index_manager = KAGIndexManager.from_config(
-                {
-                    "type": index_name,
-                    "llm_config": qa_config.get("llm", {}),
-                    "vectorize_model_config": qa_config.get("vectorize_model", {}),
-                }
-            )
-            retriever_configs += index_manager.build_retriever_config(
-                qa_config.get("llm", {}), qa_config.get("vectorize_model", {})
-            )
-        except Exception as e:
-            raise RuntimeError(f"not found index {index_name}")
-    qa_config["retrievers"] = retriever_configs
-    pipeline_config = get_pipeline_conf("index_pipeline", qa_config)
-    pipeline = SolverPipelineABC.from_config(pipeline_config)
-    return await pipeline.ainvoke(query, reporter=reporter)
-
-
 async def do_qa_pipeline(
     use_pipeline, query, qa_config, reporter, task_id, kb_project_ids
 ):
@@ -219,10 +193,12 @@ async def do_qa_pipeline(
         custom_pipeline_conf = copy.deepcopy(qa_config.get(use_pipeline, None))
     else:
         custom_pipeline_conf = copy.deepcopy(qa_config.get("solver_pipeline", None))
-
-    self_cognition_conf = get_pipeline_conf("self_cognition_pipeline", qa_config)
-    self_cognition_pipeline = SolverPipelineABC.from_config(self_cognition_conf)
-    self_cognition_res = await self_cognition_pipeline.ainvoke(query, reporter=reporter)
+    if use_pipeline not in ["index_pipeline"]:
+        self_cognition_conf = get_pipeline_conf("self_cognition_pipeline", qa_config)
+        self_cognition_pipeline = SolverPipelineABC.from_config(self_cognition_conf)
+        self_cognition_res = await self_cognition_pipeline.ainvoke(query, reporter=reporter)
+    else:
+        self_cognition_res = False
     if not self_cognition_res:
         if custom_pipeline_conf:
             pipeline_config = custom_pipeline_conf
@@ -262,7 +238,8 @@ async def qa(task_id, query, project_id, host_addr, app_id, params={}):
 
     kb_configs = {}
     kb_project_ids = []
-
+    vectorize_model = {}
+    global_index_set = main_config.get("chat", {}).get("index_list", [])
     if isinstance(main_config.get("kb"), list):
         kbs = main_config["kb"]
         for kb in kbs:
@@ -293,12 +270,24 @@ async def qa(task_id, query, project_id, host_addr, app_id, params={}):
                     kb_conf.update_conf({"llm": main_config["llm"]})
                 if "vectorizer" in kb:
                     kb_conf.update_conf({"vectorize_model": kb["vectorizer"]})
-
+                    vectorize_model = kb["vectorizer"]
+                if "index_list" not in kb and global_index_set:
+                    kb["index_list"] = global_index_set
                 KAGConfigAccessor.set_task_config(kb_task_project_id, kb_conf)
                 kb_configs[kb_project_id] = (kb_task_project_id, kb_conf)
-
             except Exception as e:
                 logger.error(f"KB配置初始化失败: {str(e)}", exc_info=True)
+    if "vectorize_model" not in main_config.keys():
+        main_config["vectorize_model"] = vectorize_model
+
+    if vectorize_model:
+        KAG_CONFIG.update_conf({
+            "vectorize_model": vectorize_model
+        })
+    if main_config["llm"]:
+        KAG_CONFIG.update_conf({
+            "llm": main_config["llm"]
+        })
     reporter_map = {
         "kag_thinker_pipeline": "kag_open_spg_reporter"
     }
@@ -315,17 +304,15 @@ async def qa(task_id, query, project_id, host_addr, app_id, params={}):
 
     try:
         await reporter.start()
-        if use_pipeline == "index_pipeline":
-            answer = await do_index_pipeline(query, main_config, reporter)
-        else:
-            answer = await do_qa_pipeline(
-                use_pipeline,
-                query,
-                main_config,
-                reporter,
-                task_id=task_id,
-                kb_project_ids=kb_project_ids,
-            )
+        answer = await do_qa_pipeline(
+            use_pipeline,
+            query,
+            main_config,
+            reporter,
+            task_id=task_id,
+            kb_project_ids=kb_project_ids,
+        )
+
         if answer:
             reporter.add_report_line("answer", "Final Answer", answer, "FINISH")
 

kag/solver/pipeline/naive_rag_pipeline.py

@@ -71,7 +71,7 @@ async def planning(self, query, context, **kwargs):
         """
         tasks_dep = {}
         tasks_dep[0] = {
-            "executor": "Retriever",
+            "executor": "ChunkRetriever",
             "dependent_task_ids": [],
             "arguments": {"query": query},
         }

kag/solver/pipelineconf/naive_rag.yaml

@@ -9,6 +9,7 @@ chunk_retrieved_executor: &chunk_retrieved_executor_conf
   top_k: 10
   retriever:
     type: vector_chunk_retriever
+    score_threshold: 0.65
     vectorize_model: "{vectorize_model}"
 
 
@@ -17,7 +18,7 @@ solver_pipeline:
   executors:
     - *chunk_retrieved_executor_conf
   generator:
-    type: llm_generator
+    type: llm_index_generator
     llm_client: "{chat_llm}"
     generated_prompt:
       type: default_refer_generator_prompt