diff --git a/kag/common/rate_limiter.py b/kag/common/rate_limiter.py index ca93baa3c..258c6ffac 100644 --- a/kag/common/rate_limiter.py +++ b/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() diff --git a/kag/interface/common/llm_client.py b/kag/interface/common/llm_client.py index 4c8e92607..b9738b0e8 100644 --- a/kag/interface/common/llm_client.py +++ b/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 diff --git a/kag/solver/executor/retriever/kag_hybrid_retrieval_executor.py b/kag/solver/executor/retriever/kag_hybrid_retrieval_executor.py index d9e410a55..455f1eeb8 100644 --- a/kag/solver/executor/retriever/kag_hybrid_retrieval_executor.py +++ b/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 diff --git a/kag/solver/planner/kag_model_planner.py b/kag/solver/planner/kag_model_planner.py index df95d7300..62bcd6c20 100644 --- a/kag/solver/planner/kag_model_planner.py +++ b/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) diff --git a/kag/solver/prompt/context_select_prompt.py b/kag/solver/prompt/context_select_prompt.py index cededf7ae..f0083945f 100644 --- a/kag/solver/prompt/context_select_prompt.py +++ b/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)