fix(solver): add rate limiter and llm retry

kag/common/rate_limiter.py

@@ -11,6 +11,9 @@
 # or implied.
 
 from aiolimiter import AsyncLimiter
+import threading
+import time
+from collections import deque
 
 
 class RateLimiterManger:
@@ -27,3 +30,71 @@ def get_rate_limiter(
 
 
 RATE_LIMITER_MANGER = RateLimiterManger()
+
+
+class SyncRateLimiter:
+    def __init__(self, max_rate: int, time_period: float = 1.0):
+        """
+        Initialize the rate limiter.
+
+        :param max_rate: The maximum number of requests allowed within the time window.
+        :param time_period: The length of the time window in seconds.
+        """
+        self.max_rate = max_rate
+        self.time_period = time_period
+        self.lock = threading.Lock()
+        self.call_timestamps = deque()
+
+    def acquire(self):
+        """
+        Acquire a permit to proceed. If the rate limit is exceeded, block and wait until a permit becomes available.
+        """
+        with self.lock:
+            now = time.time()
+            # Remove records outside the current time window
+            while (
+                self.call_timestamps
+                and now - self.call_timestamps[0] >= self.time_period
+            ):
+                self.call_timestamps.popleft()
+
+            if len(self.call_timestamps) >= self.max_rate:
+                # Wait until the earliest request exits the window
+                sleep_time = self.call_timestamps[0] + self.time_period - now
+                if sleep_time > 0:
+                    time.sleep(sleep_time)
+                    now = time.time()  # Update current time after waiting
+                    # Clean up expired requests again
+                    while (
+                        self.call_timestamps
+                        and now - self.call_timestamps[0] >= self.time_period
+                    ):
+                        self.call_timestamps.popleft()
+
+            # Append the timestamp of the current request
+            self.call_timestamps.append(now)
+
+
+class SyncRateLimiterManager:
+    def __init__(self):
+        self.limiter_map = {}
+        self.lock = threading.Lock()
+
+    def get_rate_limiter(
+        self, name: str, max_rate: int = 1000, time_period: float = 1.0
+    ):
+        """
+        Get or create a SyncRateLimiter instance by name in a thread-safe manner.
+
+        :param name: Unique identifier for the rate limiter.
+        :param max_rate: Maximum number of requests allowed within the time window.
+        :param time_period: Length of the time window in seconds.
+        :return: A thread-safe SyncRateLimiter instance.
+        """
+        with self.lock:
+            if name not in self.limiter_map:
+                self.limiter_map[name] = SyncRateLimiter(max_rate, time_period)
+        return self.limiter_map[name]
+
+
+SYNC_RATE_LIMITER_MANAGER = SyncRateLimiterManager()

kag/interface/common/llm_client.py

@@ -27,7 +27,7 @@
 from tenacity import retry, stop_after_attempt, wait_exponential
 from kag.interface import PromptABC
 from kag.common.registry import Registrable
-from kag.common.rate_limiter import RATE_LIMITER_MANGER
+from kag.common.rate_limiter import RATE_LIMITER_MANGER, SYNC_RATE_LIMITER_MANAGER
 from kag.common.conf import KAGConstants, KAGConfigAccessor
 
 logger = logging.getLogger(__name__)
@@ -170,6 +170,9 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self.limiter = RATE_LIMITER_MANGER.get_rate_limiter(name, max_rate, time_period)
+        self.sync_limiter = SYNC_RATE_LIMITER_MANAGER.get_rate_limiter(
+            name, max_rate, time_period
+        )
         self.enable_check = kwargs.get("enable_check", True)
         self.max_tokens = kwargs.get("max_tokens", 8192)
         task_id = kwargs.get(KAGConstants.KAG_QA_TASK_CONFIG_KEY, None)
@@ -324,6 +327,7 @@ def invoke(
         if tools:
             with_json_parse = False
         try:
+            self.sync_limiter.acquire()
             response = (
                 self.call_with_json_parse(prompt=prompt, **kwargs)
                 if with_json_parse

kag/solver/executor/retriever/kag_hybrid_retrieval_executor.py

@@ -15,8 +15,11 @@
 from typing import List, Optional
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
+
+from tenacity import stop_after_attempt, retry, wait_exponential
+
 from kag.common.conf import KAGConstants, KAGConfigAccessor
-from kag.common.config import get_default_chat_llm_config, LogicFormConfiguration
+from kag.common.config import get_default_chat_llm_config
 from kag.common.parser.logic_node_parser import GetSPONode
 from kag.interface import (
     ExecutorABC,
@@ -74,6 +77,10 @@ def __init__(
             {"type": "context_select_prompt"}
         )
 
+    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1))
+    def context_select_call(self, variables):
+        return self.context_select_llm.invoke(variables, self.context_select_prompt)
+
     def context_select(self, query: str, sorted_chunks):
         chunks = []
         for idx, item in enumerate(sorted_chunks):
@@ -89,6 +96,24 @@ def context_select(self, query: str, sorted_chunks):
                 indices.append(i)
         return [sorted_chunks[x] for x in indices]
 
+    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1))
+    def summary_answer(
+        self, tag_id, task_query, deps_context, formatted_docs, **kwargs
+    ):
+        return self.llm_module.invoke(
+            {
+                "cur_question": task_query,
+                "questions": "\n\n".join(deps_context),
+                "docs": "\n\n".join(formatted_docs),
+            },
+            self.summary_prompt,
+            with_json_parse=False,
+            with_except=True,
+            segment_name=tag_id,
+            tag_name=f"begin_summary_{task_query}",
+            **kwargs,
+        )
+
     def do_retrieval(
         self, task_query, tag_id, task, context: Context, **kwargs
     ) -> RetrieverOutput:
@@ -216,23 +241,24 @@ def do_summary(
         selected_rel = list(set(selected_rel))
         formatted_docs = [str(rel) for rel in selected_rel]
         if retrieved_data.chunks:
-            selected_chunks = self.context_select(task_query, 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,
+                )
+                selected_chunks = retrieved_data.chunks[:10]
             for doc in selected_chunks:
                 formatted_docs.append(f"{doc.content}")
 
         deps_context = format_task_dep_context(task.parents)
         summary_start_time = time.time()
-        summary_response = self.llm_module.invoke(
-            {
-                "cur_question": task_query,
-                "questions": "\n\n".join(deps_context),
-                "docs": "\n\n".join(formatted_docs),
-            },
-            self.summary_prompt,
-            with_json_parse=False,
-            with_except=True,
-            segment_name=tag_id,
-            tag_name=f"begin_summary_{task_query}",
+        summary_response = self.summary_answer(
+            tag_id=tag_id,
+            task_query=task_query,
+            deps_context=deps_context,
+            formatted_docs=formatted_docs,
             **kwargs,
         )
         retrieved_data.summary = summary_response

kag/solver/planner/kag_model_planner.py

@@ -9,6 +9,7 @@
 # Unless required by applicable law or agreed to in writing, software distributed under the License
 # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 # or implied.
+import json
 from typing import List, Optional
 
 from kag.common.parser.logic_node_parser import (
@@ -20,6 +21,7 @@
     GetNode,
 )
 from kag.interface import PlannerABC, Task, LLMClient, PromptABC, Context
+from kag.interface.solver.planner_abc import format_task_dep_context
 
 
 def _get_dep_task_id(logic_forms):
@@ -77,6 +79,7 @@ def __init__(
         llm: LLMClient,
         system_prompt: PromptABC,
         clarification_prompt: PromptABC,
+        rewrite_prompt: PromptABC,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -85,6 +88,59 @@ def __init__(
         self.logic_node_parser = ParseLogicForm(schema=None, schema_retrieval=None)
         self.system_prompt = system_prompt
         self.clarification_prompt = clarification_prompt
+        self.rewrite_prompt = rewrite_prompt
+
+    def check_require_rewrite(self, task: Task):
+        """Determines if query rewriting is needed based on parameter patterns.
+
+        Args:
+            task (Task): Task to check for rewrite requirements
+
+        Returns:
+            bool: True if query contains dynamic parameter references (e.g., {{1.output}})
+        """
+        return task.arguments.get("is_need_rewrite", False)
+
+    async def query_rewrite(self, task: Task, **kwargs):
+        """Performs asynchronous query rewriting using LLM and context.
+
+        Args:
+            task (Task): Task containing the query to rewrite
+
+        Returns:
+            str: Rewritten query with resolved dynamic references
+        """
+        query = task.arguments["query"]
+        tag_id = f"{query}_begin_task"
+        # print(f"Old query: {query}")
+        deps_context = format_task_dep_context(task.parents)
+        generate_context = {
+            "target question": kwargs.get("query"),
+            "history_qa": deps_context,
+        }
+        new_query = await self.llm.ainvoke(
+            {
+                "input": query,
+                "content": json.dumps(generate_context, indent=2, ensure_ascii=False),
+            },
+            self.rewrite_prompt,
+            segment_name=tag_id,
+            tag_name="Rewrite query",
+            with_json_parse=self.rewrite_prompt.is_json_format(),
+            **kwargs,
+        )
+        logic_form_node = task.arguments.get("logic_form_node", None)
+        if logic_form_node:
+            logic_form_node.sub_query = new_query
+            return {
+                "rewrite_query": new_query,
+                "origin_query": query,
+                "query": new_query,
+                "logic_form_node": logic_form_node,
+            }
+        # print(f"query rewrite context = {context}")
+        # print(f"New query: {new_query}")
+        return {"query": new_query}
 
     async def ainvoke(self, query, **kwargs) -> List[Task]:
         """Asynchronously generates task plan using LLM.
@@ -146,7 +202,7 @@ async def ainvoke(self, query, **kwargs) -> List[Task]:
                 "arguments": {
                     "query": logic_form.sub_query,
                     "logic_form_node": logic_form,
-                    "is_need_rewrite": False,
+                    "is_need_rewrite": True if task_deps else False,
                 },
             }
         return Task.create_tasks_from_dag(tasks_dep)

kag/solver/prompt/context_select_prompt.py

@@ -19,7 +19,10 @@
 @PromptABC.register("context_select_prompt")
 class ContextSelectPrompt(PromptABC):
     template_en = """
-The Question field is the user's question, and Context is a batch of articles retrieved from the knowledge base that may contain the answer to the question. Please carefully analyze the question and each context, and return up to 5 contexts that may contain the answer to the question, and answer the question. If you think all the retrieved contexts are irrelevant to the question, return an empty list and set the answer to UNKNOWN. The return format refers to the example：
+The Question field is the user's question, and Context is a batch of articles retrieved from the knowledge base that may contain the answer to the question. Please carefully analyze the question and each context, and return up to 5 contexts that may contain the answer to the question, and answer the question. If you think all the retrieved contexts are irrelevant to the question, return an empty list and set the answer to UNKNOWN. 
+Note: 
+1. response must be json format, only json format, without explanation
+2. The return format refers to the example：
 
 Question:
 who is the performer of Hello Love?
@@ -68,7 +71,10 @@ class ContextSelectPrompt(PromptABC):
 """
 
     template_zh = """
-问题字段是用户的问题，Context是从知识库中检索到的可能包含问题答案的一批文章。请仔细分析问题和每个上下文，返回最多5个可能包含问题答案的上下文，并回答问题。如果你认为所有检索到的上下文都与问题无关，请返回空列表并将答案设置为UNKNOWN。返回格式参考示例：
+问题字段是用户的问题，Context是从知识库中检索到的可能包含问题答案的一批文章。请仔细分析问题和每个上下文，返回最多5个可能包含问题答案的上下文，并回答问题。如果你认为所有检索到的上下文都与问题无关，请返回空列表并将答案设置为UNKNOWN。
+注意：
+1、输出必须是json格式，只输出json格式内容，不需要带上解释
+2、返回格式参考下面示例：
 
 问题:
 谁是《Hello Love》的演唱者？
@@ -127,13 +133,17 @@ def build_prompt(self, variables: Dict) -> str:
         else:
             return f"{self.template_en}\nQuestion:\n{question}\nContext:\n{context}"
 
-    def parse_response(self, response: str, **kwargs):
+    def parse_response(self, rsp: str, **kwargs):
         if self.language == "zh":
-            if isinstance(response, str) and "答案:" in response:
-                response = response.split("答案:")[1]
+            if isinstance(rsp, str) and "答案:" in rsp:
+                response = rsp.split("答案:")[1]
+            else:
+                response = rsp
         else:
-            if isinstance(response, str) and "Answer:" in response:
-                response = response.split("Answer:")[1]
+            if isinstance(rsp, str) and "Answer:" in rsp:
+                response = rsp.split("Answer:")[1]
+            else:
+                response = rsp
         if isinstance(response, str):
             response = json.loads(response)