From e23f63bd93fd1f935010ad881b50111c9a4c48ab Mon Sep 17 00:00:00 2001 From: Taranum Wasu <81034301+Taranum01@users.noreply.github.com> Date: Wed, 1 Jul 2026 07:00:54 +0530 Subject: [PATCH] fix(agent): prevent empty LLM user message after prompt fitting (#16413) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Treat `max_tokens=0` as unset (`or 8192`) when building model context budgets, fixing agents that silently zeroed prompts when a vLLM model had `max_tokens: 0` in tenant config - Replace trailing same-role canvas history in `LLM._sys_prompt_and_msg` instead of skipping the current user prompt - Add `LLM.fit_messages()` validation after `message_fit_in` on agent paths so empty user content fails fast with a clear error instead of reaching vLLM Fixes #16411 ## Root cause Agent canvas workflow called `message_fit_in` with `int(max_length * 0.97)`. When `max_length` was `0`, both system and user content were trimmed to empty strings. The `[HISTORY STREAMLY]` log showing only `{"role":"user","content":""}` matches this. A secondary bug skipped appending the formatted user prompt when history ended with a `user` role message. ## Test plan - [x] Added `test/unit_test/agent/component/test_llm_prompt.py` for role-replace, validation, and zero-budget fitting - [x] Added `test_message_fit_in_zero_budget_preserves_non_empty_messages` in `test_generator_message_fit_in.py` - [ ] CI unit tests - [ ] Manual: agent canvas `begin → Retrieval → Agent → Message` with vLLM Qwen3; confirm user message reaches LLM Made with [Cursor](https://cursor.com) --------- Co-authored-by: Taranum Wasu Co-authored-by: Cursor --- agent/component/agent_with_tools.py | 45 +++-- agent/component/llm.py | 189 ++++++++---------- api/db/joint_services/tenant_model_service.py | 27 ++- api/db/services/dialog_service.py | 15 +- api/db/services/tenant_llm_service.py | 6 +- rag/prompts/generator.py | 143 ++++++------- .../agent/component/test_llm_prompt.py | 61 ++++++ .../prompts/test_generator_message_fit_in.py | 27 ++- 8 files changed, 282 insertions(+), 231 deletions(-) create mode 100644 test/unit_test/agent/component/test_llm_prompt.py diff --git a/agent/component/agent_with_tools.py b/agent/component/agent_with_tools.py index 6ac4f220db..93b9f2eff2 100644 --- a/agent/component/agent_with_tools.py +++ b/agent/component/agent_with_tools.py @@ -115,12 +115,12 @@ class Agent(LLM, ToolBase): if self.tool_meta: self.chat_mdl.bind_tools(self.toolcall_session, self.tool_meta) - def _fit_messages(self, prompt: str, msg: list[dict]) -> list[dict]: - _, fitted_messages = message_fit_in( - [{"role": "system", "content": prompt}, *msg], - int(self.chat_mdl.max_length * 0.97), - ) - return fitted_messages + def _fit_messages(self, prompt: str, msg: list[dict]) -> tuple[list[dict] | None, str | None]: + msg_fit, fit_error = LLM.fit_messages(prompt, msg, self.chat_mdl.max_length) + if fit_error: + logging.error("Agent prompt fit error: %s", fit_error) + return None, fit_error + return msg_fit, None @staticmethod def _append_system_prompt(msg: list[dict], extra_prompt: str) -> None: @@ -185,7 +185,7 @@ class Agent(LLM, ToolBase): {"role": "system", "content": schema_prompt + "\nIMPORTANT: Output ONLY valid JSON. No markdown, no extra text."}, {"role": "user", "content": text}, ] - _, fmt_msgs = message_fit_in(fmt_msgs, int(self.chat_mdl.max_length * 0.97)) + _, fmt_msgs = message_fit_in(fmt_msgs, LLM.context_fit_budget(self.chat_mdl.max_length)) return await self._generate_async(fmt_msgs) def _invoke(self, **kwargs): @@ -196,11 +196,14 @@ class Agent(LLM, ToolBase): if self.check_if_canceled("Agent processing"): return + user_prompt = kwargs.get("user_prompt") + user_prompt_text = "" if user_prompt is None else str(user_prompt) _logger.debug( - "[Agent] _invoke_async called. Component: %s, Keys in kwargs: %s, user_prompt: %s, tools count: %d", + "[Agent] _invoke_async called. Component: %s, Keys in kwargs: %s, user_prompt_present: %s, user_prompt_length: %d, tools count: %d", self._id, list(kwargs.keys()), - json.dumps(kwargs.get("user_prompt", ""), ensure_ascii=False, default=str)[:300], + bool(user_prompt_text.strip()), + len(user_prompt_text), len(self.tools) if self.tools else 0, ) @@ -215,8 +218,7 @@ class Agent(LLM, ToolBase): else: usr_pmt = str(kwargs["user_prompt"]) self._param.prompts = [{"role": "user", "content": usr_pmt}] - _logger.debug("[Agent] Built user prompt with length=%d, reasoning=%s, context=%s", - len(usr_pmt), bool(kwargs.get("reasoning")), bool(kwargs.get("context"))) + _logger.debug("[Agent] Built user prompt with length=%d, reasoning=%s, context=%s", len(usr_pmt), bool(kwargs.get("reasoning")), bool(kwargs.get("context"))) if not self.tools: if self.check_if_canceled("Agent processing"): @@ -240,7 +242,14 @@ class Agent(LLM, ToolBase): self.set_output("content", partial(self.stream_output_with_tools_async, prompt, deepcopy(msg), user_defined_prompt)) return - msg = self._fit_messages(prompt, msg) + msg, fit_error = self._fit_messages(prompt, msg) + if fit_error: + if self.get_exception_default_value(): + self.set_output("content", self.get_exception_default_value()) + else: + self.set_output("_ERROR", fit_error) + return + self._append_system_prompt(msg, schema_prompt) _logger.debug("[Agent] Calling LLM with %d messages, has_schema=%s", len(msg), bool(schema_prompt)) ans = await self._generate_async(msg) @@ -283,7 +292,17 @@ class Agent(LLM, ToolBase): self.callback("Multi-turn conversation optimization", {}, user_request, elapsed_time=timer() - st) msg = [*msg[:-1], {"role": "user", "content": user_request}] - msg = self._fit_messages(prompt, msg) + msg, fit_error = self._fit_messages(prompt, msg) + if fit_error: + if self.get_exception_default_value(): + fallback = self.get_exception_default_value() + self.set_output("content", fallback) + yield fallback + else: + self.set_output("_ERROR", fit_error) + self.set_output("content", fit_error) + yield fit_error + return need2cite = self._param.cite and self._canvas.get_reference()["chunks"] and self._id.find("-->") < 0 cited = False diff --git a/agent/component/llm.py b/agent/component/llm.py index 59fc60de72..50a95743f0 100644 --- a/agent/component/llm.py +++ b/agent/component/llm.py @@ -23,6 +23,7 @@ from typing import Any, AsyncGenerator import json_repair from functools import partial from common.constants import LLMType +from api.db.services.dialog_service import _stream_with_think_delta from api.db.services.llm_service import LLMBundle from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_model_type_by_name from agent.component.base import ComponentBase, ComponentParamBase @@ -60,6 +61,7 @@ class LLMParam(ComponentParamBase): def gen_conf(self): conf = {} + def get_attr(nm): try: return getattr(self, nm) @@ -87,18 +89,13 @@ class LLM(ComponentBase): model_types = get_model_type_by_name(self._canvas.get_tenant_id(), self._param.llm_id) model_type = "chat" if "chat" in model_types else model_types[0] chat_model_config = get_model_config_from_provider_instance(self._canvas.get_tenant_id(), model_type, self._param.llm_id) - self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), chat_model_config, - max_retries=self._param.max_retries, - retry_interval=self._param.delay_after_error) + self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), chat_model_config, max_retries=self._param.max_retries, retry_interval=self._param.delay_after_error) self.imgs = [] def get_input_form(self) -> dict[str, dict]: res = {} for k, v in self.get_input_elements().items(): - res[k] = { - "type": "line", - "name": v["name"] - } + res[k] = {"type": "line", "name": v["name"]} return res def get_input_elements(self) -> dict[str, Any]: @@ -119,14 +116,41 @@ class LLM(ComponentBase): def _sys_prompt_and_msg(self, msg, args): if isinstance(self._param.prompts, str): self._param.prompts = [{"role": "user", "content": self._param.prompts}] + history_size = len(msg) for p in self._param.prompts: - if msg and msg[-1]["role"] == p["role"]: - continue - p = deepcopy(p) - p["content"] = self.string_format(p["content"], args) - msg.append(p) + formatted = deepcopy(p) + formatted["content"] = self.string_format(formatted["content"], args) + if len(msg) == history_size and msg and msg[-1]["role"] == formatted["role"]: + msg[-1] = formatted + else: + msg.append(formatted) return msg, self.string_format(self._param.sys_prompt, args) + @staticmethod + def effective_context_length(max_length) -> int: + return max_length or 8192 + + @classmethod + def context_fit_budget(cls, max_length) -> int: + return int(cls.effective_context_length(max_length) * 0.97) + + @staticmethod + def validate_fitted_messages(msg_fit: list[dict]) -> str | None: + if len(msg_fit) < 2: + return "**ERROR**: message_fit_in produced insufficient messages for LLM" + last = msg_fit[-1] + if last.get("role") != "user" or not str(last.get("content") or "").strip(): + return "**ERROR**: LLM user message is empty after prompt fitting; check model max_tokens context setting" + return None + + @classmethod + def fit_messages(cls, system_prompt: str, msg: list[dict], max_length) -> tuple[list[dict], str | None]: + _, msg_fit = message_fit_in( + [{"role": "system", "content": system_prompt}, *deepcopy(msg)], + cls.context_fit_budget(max_length), + ) + return msg_fit, cls.validate_fitted_messages(msg_fit) + @staticmethod def _extract_data_images(value) -> list[str]: imgs = [] @@ -255,7 +279,9 @@ class LLM(ComponentBase): text_parts.append(f) logging.info( "[LLM] sys.files split: text_parts=%d image_data_uris=%d (explicit=%s)", - len(text_parts), len(image_data_uris), explicit, + len(text_parts), + len(image_data_uris), + explicit, ) return text_parts, image_data_uris @@ -286,7 +312,9 @@ class LLM(ComponentBase): self.imgs = self._uniq_images(self.imgs + extracted_imgs + sys_file_imgs) logging.debug( "[LLM] imgs rebuilt: total=%d sys_files_added=%d unique_dropped=%d", - len(self.imgs), len(sys_file_imgs), max(0, prev_img_count + len(sys_file_imgs) - len(self.imgs)), + len(self.imgs), + len(sys_file_imgs), + max(0, prev_img_count + len(sys_file_imgs) - len(self.imgs)), ) model_types = get_model_type_by_name(self._canvas.get_tenant_id(), self._param.llm_id) if self.imgs and LLMType.IMAGE2TEXT.value in model_types: @@ -297,9 +325,7 @@ class LLM(ComponentBase): model_type = model_types[0] model_config = get_model_config_from_provider_instance(self._canvas.get_tenant_id(), model_type, self._param.llm_id) if self.imgs: - self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), model_config, max_retries=self._param.max_retries, - retry_interval=self._param.delay_after_error - ) + self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), model_config, max_retries=self._param.max_retries, retry_interval=self._param.delay_after_error) msg, sys_prompt = self._sys_prompt_and_msg(self._canvas.get_history(self._param.message_history_window_size)[:-1], args) @@ -316,7 +342,9 @@ class LLM(ComponentBase): merged_idx = len(msg) - 1 logging.info( "[LLM] sys.files text merged into msg: parts=%d total_chars=%d msg_index=%d action=%s", - len(sys_file_texts), len(joined), merged_idx, + len(sys_file_texts), + len(joined), + merged_idx, "merged_into_existing_user" if merged_idx < len(msg) - 1 or msg[merged_idx].get("content", "") != joined else "appended_new_user", ) @@ -329,11 +357,11 @@ class LLM(ComponentBase): def _extract_prompts(self, sys_prompt): pts = {} for tag in ["TASK_ANALYSIS", "PLAN_GENERATION", "REFLECTION", "CONTEXT_SUMMARY", "CONTEXT_RANKING", "CITATION_GUIDELINES"]: - r = re.search(rf"<{tag}>(.*?)", sys_prompt, flags=re.DOTALL|re.IGNORECASE) + r = re.search(rf"<{tag}>(.*?)", sys_prompt, flags=re.DOTALL | re.IGNORECASE) if not r: continue pts[tag.lower()] = r.group(1) - sys_prompt = re.sub(rf"<{tag}>(.*?)", "", sys_prompt, flags=re.DOTALL|re.IGNORECASE) + sys_prompt = re.sub(rf"<{tag}>(.*?)", "", sys_prompt, flags=re.DOTALL | re.IGNORECASE) return pts, sys_prompt async def _generate_async(self, msg: list[dict], **kwargs) -> str: @@ -342,84 +370,33 @@ class LLM(ComponentBase): return await self.chat_mdl.async_chat(msg[0]["content"], msg[1:], self._param.gen_conf(), images=self.imgs, **kwargs) async def _generate_streamly(self, msg: list[dict], **kwargs) -> AsyncGenerator[str, None]: - async def delta_wrapper(txt_iter): - ans = "" - last_idx = 0 - endswith_think = False - - def delta(txt): - nonlocal ans, last_idx, endswith_think - delta_ans = txt[last_idx:] - ans = txt - - if delta_ans.find("") == 0: - last_idx += len("") - return "" - elif delta_ans.find("") > 0: - delta_ans = txt[last_idx:last_idx + delta_ans.find("")] - last_idx += delta_ans.find("") - return delta_ans - elif delta_ans.endswith(""): - endswith_think = True - elif endswith_think: - endswith_think = False - return "" - - last_idx = len(ans) - if ans.endswith(""): - last_idx -= len("") - return re.sub(r"(|)", "", delta_ans) - - async for t in txt_iter: - yield delta(t) - - if not self.imgs: - async for t in delta_wrapper(self.chat_mdl.async_chat_streamly(msg[0]["content"], msg[1:], self._param.gen_conf(), **kwargs)): - yield t - return - - async for t in delta_wrapper(self.chat_mdl.async_chat_streamly(msg[0]["content"], msg[1:], self._param.gen_conf(), images=self.imgs, **kwargs)): - yield t + stream_kwargs = {"images": self.imgs} if self.imgs else {} + stream_kwargs.update(kwargs) + stream = self.chat_mdl.async_chat_streamly_delta(msg[0]["content"], msg[1:], self._param.gen_conf(), **stream_kwargs) + async for _, value, _ in _stream_with_think_delta(stream, min_tokens=0): + yield value async def _stream_output_async(self, prompt, msg): - _, msg = message_fit_in([{"role": "system", "content": prompt}, *msg], int(self.chat_mdl.max_length * 0.97)) + msg_fit, fit_error = self.fit_messages(prompt, msg, self.chat_mdl.max_length) + if fit_error: + logging.error("LLM streaming prompt fit error: %s", fit_error) + if self.get_exception_default_value(): + fallback = self.get_exception_default_value() + self.set_output("content", fallback) + yield fallback + else: + self.set_output("_ERROR", fit_error) + return + answer = "" - last_idx = 0 - endswith_think = False - - def delta(txt): - nonlocal answer, last_idx, endswith_think - delta_ans = txt[last_idx:] - answer = txt - - if delta_ans.find("") == 0: - last_idx += len("") - return "" - elif delta_ans.find("") > 0: - delta_ans = txt[last_idx:last_idx + delta_ans.find("")] - last_idx += delta_ans.find("") - return delta_ans - elif delta_ans.endswith(""): - endswith_think = True - elif endswith_think: - endswith_think = False - return "" - - last_idx = len(answer) - if answer.endswith(""): - last_idx -= len("") - return re.sub(r"(|)", "", delta_ans) - stream_kwargs = {"images": self.imgs} if self.imgs else {} extra_chat_kwargs = self._get_chat_template_kwargs() stream_kwargs.update(extra_chat_kwargs) - async for ans in self.chat_mdl.async_chat_streamly(msg[0]["content"], msg[1:], self._param.gen_conf(), **stream_kwargs): + stream = self.chat_mdl.async_chat_streamly_delta(msg_fit[0]["content"], msg_fit[1:], self._param.gen_conf(), **stream_kwargs) + async for _, ans, _ in _stream_with_think_delta(stream, min_tokens=0): if self.check_if_canceled("LLM streaming"): return - if isinstance(ans, int): - continue - if ans.find("**ERROR**") >= 0: if self.get_exception_default_value(): self.set_output("content", self.get_exception_default_value()) @@ -428,11 +405,12 @@ class LLM(ComponentBase): self.set_output("_ERROR", ans) return - yield delta(ans) + answer += ans + yield ans self.set_output("content", answer) - @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60))) + @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60))) async def _invoke_async(self, **kwargs): if self.check_if_canceled("LLM processing"): return @@ -457,10 +435,11 @@ class LLM(ComponentBase): if self.check_if_canceled("LLM processing"): return - _, msg_fit = message_fit_in( - [{"role": "system", "content": prompt_with_schema}, *deepcopy(msg)], - int(self.chat_mdl.max_length * 0.97), - ) + msg_fit, fit_error = self.fit_messages(prompt_with_schema, msg, self.chat_mdl.max_length) + if fit_error: + logging.error("LLM structured prompt fit error: %s", fit_error) + self.set_output("_ERROR", fit_error) + return error = "" ans = await self._generate_async(msg_fit, **extra_chat_kwargs) msg_fit.pop(0) @@ -480,9 +459,7 @@ class LLM(ComponentBase): downstreams = self._canvas.get_component(self._id)["downstream"] if self._canvas.get_component(self._id) else [] ex = self.exception_handler() - if any([self._canvas.get_component_obj(cid).component_name.lower() == "message" for cid in downstreams]) and not ( - ex and ex["goto"] - ): + if any([self._canvas.get_component_obj(cid).component_name.lower() == "message" for cid in downstreams]) and not (ex and ex["goto"]): self.set_output("content", partial(self._stream_output_async, prompt, deepcopy(msg))) return @@ -491,9 +468,11 @@ class LLM(ComponentBase): if self.check_if_canceled("LLM processing"): return - _, msg_fit = message_fit_in( - [{"role": "system", "content": prompt}, *deepcopy(msg)], int(self.chat_mdl.max_length * 0.97) - ) + msg_fit, fit_error = self.fit_messages(prompt, msg, self.chat_mdl.max_length) + if fit_error: + logging.error("LLM prompt fit error: %s", fit_error) + error = fit_error + break error = "" ans = await self._generate_async(msg_fit, **extra_chat_kwargs) msg_fit.pop(0) @@ -510,7 +489,7 @@ class LLM(ComponentBase): else: self.set_output("_ERROR", error) - @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60))) + @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60))) def _invoke(self, **kwargs): return asyncio.run(self._invoke_async(**kwargs)) @@ -532,11 +511,11 @@ class LLM(ComponentBase): return {} return {"chat_template_kwargs": chat_template_kwargs} - async def add_memory(self, user:str, assist:str, func_name: str, params: dict, results: str, user_defined_prompt:dict={}): + async def add_memory(self, user: str, assist: str, func_name: str, params: dict, results: str, user_defined_prompt: dict = {}): summ = await tool_call_summary(self.chat_mdl, func_name, params, results, user_defined_prompt) logging.info(f"[MEMORY]: {summ}") self._canvas.add_memory(user, assist, summ) def thoughts(self) -> str: - _, msg,_ = self._prepare_prompt_variables() - return "⌛Give me a moment—starting from: \n\n" + re.sub(r"(User's query:|[\\]+)", '', msg[-1]['content'], flags=re.DOTALL) + "\n\nI’ll figure out our best next move." + _, msg, _ = self._prepare_prompt_variables() + return "⌛Give me a moment—starting from: \n\n" + re.sub(r"(User's query:|[\\]+)", "", msg[-1]["content"], flags=re.DOTALL) + "\n\nI’ll figure out our best next move." diff --git a/api/db/joint_services/tenant_model_service.py b/api/db/joint_services/tenant_model_service.py index fdc3196d59..96f234754e 100644 --- a/api/db/joint_services/tenant_model_service.py +++ b/api/db/joint_services/tenant_model_service.py @@ -32,6 +32,8 @@ def _factory_model_types(llm: dict) -> list[str]: if isinstance(model_type, list): return model_type return [model_type] if model_type else [] + + def _decode_api_key_config(raw_api_key: str) -> tuple[str, bool | None, str | None]: if not raw_api_key: return raw_api_key, None, None @@ -132,7 +134,7 @@ def ensure_paddleocr_from_env(tenant_id: str) -> str | None: ) -def get_tenant_default_model_by_type(tenant_id: str, model_type: str|enum.Enum): +def get_tenant_default_model_by_type(tenant_id: str, model_type: str | enum.Enum): exist, tenant = TenantService.get_by_id(tenant_id) if not exist: raise LookupError("Tenant not found") @@ -142,7 +144,7 @@ def get_tenant_default_model_by_type(tenant_id: str, model_type: str|enum.Enum): case LLMType.EMBEDDING.value: model_name = tenant.embd_id case LLMType.SPEECH2TEXT.value: - model_name = tenant.asr_id + model_name = tenant.asr_id case LLMType.IMAGE2TEXT.value: model_name = tenant.img2txt_id case LLMType.CHAT.value: @@ -185,10 +187,7 @@ def _resolve_instance_for_model(provider_obj, instance_name: str, model_name: st if instance_name != "default": raise LookupError(f"Instance {instance_name} not found for model {model_name}.") - active_instances = [ - inst for inst in TenantModelInstanceService.get_all_by_provider_id(provider_obj.id) - if inst.status == ActiveStatusEnum.ACTIVE.value - ] + active_instances = [inst for inst in TenantModelInstanceService.get_all_by_provider_id(provider_obj.id) if inst.status == ActiveStatusEnum.ACTIVE.value] if len(active_instances) == 1: logger.warning( "Model instance fallback applied for legacy default instance name", @@ -204,16 +203,13 @@ def _resolve_instance_for_model(provider_obj, instance_name: str, model_name: st raise LookupError(f"Instance {instance_name} not found for model {model_name}.") -def get_model_config_from_provider_instance(tenant_id, model_type: str|enum.Enum, model_name: str): +def get_model_config_from_provider_instance(tenant_id, model_type: str | enum.Enum, model_name: str): pure_model_name, instance_name, provider_name = split_model_name(model_name) model_type_val = model_type if isinstance(model_type, str) else model_type.value # Builtin embedding model compose_profiles = os.getenv("COMPOSE_PROFILES", "") is_tei_builtin_embedding = ( - model_type_val == LLMType.EMBEDDING.value - and "tei-" in compose_profiles - and pure_model_name == os.getenv("TEI_MODEL", "") - and (provider_name == "Builtin" or not provider_name) + model_type_val == LLMType.EMBEDDING.value and "tei-" in compose_profiles and pure_model_name == os.getenv("TEI_MODEL", "") and (provider_name == "Builtin" or not provider_name) ) if is_tei_builtin_embedding: # configured local embedding model @@ -249,7 +245,7 @@ def get_model_config_from_provider_instance(tenant_id, model_type: str|enum.Enum "api_base": extra_fields.get("base_url", ""), "model_type": model_obj.model_type, "is_tools": model_extra.get("is_tools", is_tool), - "max_tokens": model_extra.get("max_tokens", 8192), + "max_tokens": model_extra.get("max_tokens") or 8192, } if api_key_payload is not None: model_config["api_key_payload"] = api_key_payload @@ -277,7 +273,7 @@ def get_model_config_from_provider_instance(tenant_id, model_type: str|enum.Enum "api_base": extra_fields.get("base_url", ""), "model_type": model_type_val, "is_tools": llm_info.get("is_tools", is_tool), - "max_tokens": llm_info.get("max_tokens", 8192), + "max_tokens": llm_info.get("max_tokens") or 8192, } if api_key_payload is not None: model_config["api_key_payload"] = api_key_payload @@ -318,7 +314,10 @@ def get_model_type_by_name(tenant_id: str, model_name: str): if not llm_list: raise LookupError(f"Model {pure_model_name} not found for model {model_name}.") types_in_json = _factory_model_types(llm_list[0]) - return list(set(types_in_json + [model_obj.model_type for model_obj in model_objs if model_obj.status != ActiveStatusEnum.UNSUPPORTED.value]) - {model_obj.model_type for model_obj in model_objs if model_obj.status == ActiveStatusEnum.UNSUPPORTED.value}) + return list( + set(types_in_json + [model_obj.model_type for model_obj in model_objs if model_obj.status != ActiveStatusEnum.UNSUPPORTED.value]) + - {model_obj.model_type for model_obj in model_objs if model_obj.status == ActiveStatusEnum.UNSUPPORTED.value} + ) def delete_models_by_instance_ids(instance_ids: list[str]): diff --git a/api/db/services/dialog_service.py b/api/db/services/dialog_service.py index 860a1dd20c..f7c53a9df3 100644 --- a/api/db/services/dialog_service.py +++ b/api/db/services/dialog_service.py @@ -53,6 +53,7 @@ from rag.utils.tts_cache import synthesize_with_cache from common.string_utils import remove_redundant_spaces from common import settings + def _chunk_kb_id_for_doc(row_dict, kb_ids, doc_id): if len(kb_ids or []) == 1: return kb_ids[0] @@ -105,6 +106,7 @@ async def _hydrate_chunk_vectors(retriever, chunks, tenant_ids, kb_ids): if cid and cid in vectors: ck["vector"] = vectors[cid] + def _normalize_internet_flag(value): if isinstance(value, bool): return value @@ -134,7 +136,6 @@ def _enrich_chunks_with_document_metadata(chunks, metadata_fields=None): enrich_chunks_with_document_metadata(chunks, metadata_fields) - class DialogService(CommonService): model = Dialog @@ -563,7 +564,7 @@ async def async_chat(dialog, messages, stream=True, **kwargs): llm_model_config = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.CHAT) factory = llm_model_config.get("llm_factory", "") if llm_model_config else "" - max_tokens = llm_model_config.get("max_tokens", 8192) + max_tokens = llm_model_config.get("max_tokens") or 8192 check_llm_ts = timer() @@ -619,8 +620,7 @@ async def async_chat(dialog, messages, stream=True, **kwargs): if include_reference_metadata and ans.get("reference", {}).get("chunks"): if len(dialog.kb_ids) != 1 and any(not c.get("kb_id") for c in ans["reference"]["chunks"]): logging.warning( - "Skipping some _enrich_chunks_with_document_metadata results because " - "dialog.kb_ids has %d entries and use_sql returned chunks without kb_id.", + "Skipping some _enrich_chunks_with_document_metadata results because dialog.kb_ids has %d entries and use_sql returned chunks without kb_id.", len(dialog.kb_ids), ) _enrich_chunks_with_document_metadata(ans["reference"]["chunks"], metadata_fields) @@ -741,7 +741,9 @@ async def async_chat(dialog, messages, stream=True, **kwargs): kbinfos["doc_aggs"].extend(tav_res["doc_aggs"]) if prompt_config.get("use_kg"): default_chat_model = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.CHAT) - ck = await settings.kg_retriever.retrieval(" ".join(questions), tenant_ids, dialog.kb_ids, embd_mdl, LLMBundle(dialog.tenant_id, default_chat_model, trace_context=trace_context, langfuse_session_id=session_id)) + ck = await settings.kg_retriever.retrieval( + " ".join(questions), tenant_ids, dialog.kb_ids, embd_mdl, LLMBundle(dialog.tenant_id, default_chat_model, trace_context=trace_context, langfuse_session_id=session_id) + ) if ck["content_with_weight"]: kbinfos["chunks"].insert(0, ck) @@ -1667,8 +1669,7 @@ async def async_ask(question, kb_ids, tenant_id, chat_llm_name=None, search_conf except TypeError: full_text_weight = None logger.debug( - "Search async_ask retrieval weight: search_id=%s tenant_id=%s kb_count=%s " - "vector_similarity_weight=%s full_text_weight=%s", + "Search async_ask retrieval weight: search_id=%s tenant_id=%s kb_count=%s vector_similarity_weight=%s full_text_weight=%s", search_id, tenant_id, len(kb_ids), diff --git a/api/db/services/tenant_llm_service.py b/api/db/services/tenant_llm_service.py index b61a50e067..d4b89cc6e6 100644 --- a/api/db/services/tenant_llm_service.py +++ b/api/db/services/tenant_llm_service.py @@ -512,7 +512,7 @@ class LLM4Tenant: self.model_config = model_config self.mdl = TenantLLMService.model_instance(model_config, lang=lang, **kwargs) assert self.mdl, "Can't find model for {}/{}/{}".format(tenant_id, model_config["model_type"], model_config["llm_name"]) - self.max_length = model_config.get("max_tokens", 8192) + self.max_length = model_config.get("max_tokens") or 8192 self.is_tools = model_config.get("is_tools", False) self.verbose_tool_use = kwargs.get("verbose_tool_use") @@ -543,7 +543,7 @@ class LLM4Tenant: if self.langfuse: try: self.langfuse.flush() - if hasattr(self.langfuse, 'shutdown'): + if hasattr(self.langfuse, "shutdown"): self.langfuse.shutdown() except Exception: # Ignore errors during cleanup @@ -552,7 +552,7 @@ class LLM4Tenant: self.langfuse = None # Release underlying model instance if it has a close method - if self.mdl and hasattr(self.mdl, 'close') and callable(getattr(self.mdl, 'close')): + if self.mdl and hasattr(self.mdl, "close") and callable(getattr(self.mdl, "close")): try: self.mdl.close() except Exception: diff --git a/rag/prompts/generator.py b/rag/prompts/generator.py index 267954a93a..49103d56a0 100644 --- a/rag/prompts/generator.py +++ b/rag/prompts/generator.py @@ -67,6 +67,10 @@ def chunks_format(reference): def message_fit_in(msg, max_length=4000): + if max_length <= 0: + logging.debug("message_fit_in normalizing non-positive max_length=%s to 8192", max_length) + max_length = 8192 + def count(): nonlocal msg tks_cnts = [] @@ -158,7 +162,7 @@ def kb_prompt(kbinfos, max_tokens, hash_id=False): for i, ck in enumerate(kbinfos["chunks"][:chunks_num]): cnt = "\nID: {}".format(i if not hash_id else hash_str2int(get_value(ck, "id", "chunk_id"), 500)) cnt += draw_node("Title", get_value(ck, "docnm_kwd", "document_name")) - cnt += draw_node("URL", ck.get('url', '')) + cnt += draw_node("URL", ck.get("url", "")) meta = ck.get("document_metadata") or {} for k, v in meta.items(): cnt += draw_node(k, v) @@ -204,9 +208,7 @@ RANK_MEMORY = load_prompt("rank_memory") META_FILTER = load_prompt("meta_filter") ASK_SUMMARY = load_prompt("ask_summary") -PROMPT_JINJA_ENV = SandboxedEnvironment( - autoescape=False, trim_blocks=True, lstrip_blocks=True -) +PROMPT_JINJA_ENV = SandboxedEnvironment(autoescape=False, trim_blocks=True, lstrip_blocks=True) def citation_prompt(user_defined_prompts: dict = {}) -> str: @@ -290,7 +292,7 @@ async def cross_languages(tenant_id, llm_id, query, languages=[]): from api.db.services.llm_service import LLMBundle from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_tenant_default_model_by_type, get_model_type_by_name - if llm_id and "image2text" in get_model_type_by_name(tenant_id, llm_id) : + if llm_id and "image2text" in get_model_type_by_name(tenant_id, llm_id): chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.IMAGE2TEXT, llm_id) else: if not llm_id: @@ -299,11 +301,9 @@ async def cross_languages(tenant_id, llm_id, query, languages=[]): chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, llm_id) chat_mdl = LLMBundle(tenant_id, chat_model_config) rendered_sys_prompt = PROMPT_JINJA_ENV.from_string(CROSS_LANGUAGES_SYS_PROMPT_TEMPLATE).render() - rendered_user_prompt = PROMPT_JINJA_ENV.from_string(CROSS_LANGUAGES_USER_PROMPT_TEMPLATE).render(query=query, - languages=languages) + rendered_user_prompt = PROMPT_JINJA_ENV.from_string(CROSS_LANGUAGES_USER_PROMPT_TEMPLATE).render(query=query, languages=languages) - ans = await chat_mdl.async_chat(rendered_sys_prompt, [{"role": "user", "content": rendered_user_prompt}], - {"temperature": 0.2}) + ans = await chat_mdl.async_chat(rendered_sys_prompt, [{"role": "user", "content": rendered_user_prompt}], {"temperature": 0.2}) if ans.find("**ERROR**") >= 0: logging.info("[cross_languages] LLM returned error, falling back to original query") return query @@ -380,20 +380,14 @@ def tool_schema(tools_description: list[dict], complete_task=False): "function": { "name": COMPLETE_TASK, "description": "When you have the final answer and are ready to complete the task, call this function with your answer", - "parameters": { - "type": "object", - "properties": { - "answer": {"type": "string", "description": "The final answer to the user's question"}}, - "required": ["answer"] - } - } + "parameters": {"type": "object", "properties": {"answer": {"type": "string", "description": "The final answer to the user's question"}}, "required": ["answer"]}, + }, } for idx, tool in enumerate(tools_description): name = tool["function"]["name"] desc[name] = tool - return "\n\n".join([f"## {i + 1}. {fnm}\n{json.dumps(des, ensure_ascii=False, indent=4)}" for i, (fnm, des) in - enumerate(desc.items())]) + return "\n\n".join([f"## {i + 1}. {fnm}\n{json.dumps(des, ensure_ascii=False, indent=4)}" for i, (fnm, des) in enumerate(desc.items())]) def form_history(history, limit=-6): @@ -408,8 +402,7 @@ def form_history(history, limit=-6): return context -async def analyze_task_async(chat_mdl, prompt, task_name, tools_description: list[dict], - user_defined_prompts: dict = {}): +async def analyze_task_async(chat_mdl, prompt, task_name, tools_description: list[dict], user_defined_prompts: dict = {}): tools_desc = tool_schema(tools_description) context = "" @@ -427,8 +420,7 @@ async def analyze_task_async(chat_mdl, prompt, task_name, tools_description: lis return kwd -async def next_step_async(chat_mdl, history: list, tools_description: list[dict], task_desc, - user_defined_prompts: dict = {}): +async def next_step_async(chat_mdl, history: list, tools_description: list[dict], task_desc, user_defined_prompts: dict = {}): if not tools_description: return "", 0 desc = tool_schema(tools_description) @@ -482,20 +474,16 @@ def structured_output_prompt(schema=None) -> str: async def tool_call_summary(chat_mdl, name: str, params: dict, result: str, user_defined_prompts: dict = {}) -> str: template = PROMPT_JINJA_ENV.from_string(SUMMARY4MEMORY) - system_prompt = template.render(name=name, - params=json.dumps(params, ensure_ascii=False, indent=2), - result=result) + system_prompt = template.render(name=name, params=json.dumps(params, ensure_ascii=False, indent=2), result=result) user_prompt = "→ Summary: " _, msg = message_fit_in(form_message(system_prompt, user_prompt), chat_mdl.max_length) ans = await chat_mdl.async_chat(msg[0]["content"], msg[1:]) return re.sub(r"^.*", "", ans, flags=re.DOTALL) -async def rank_memories_async(chat_mdl, goal: str, sub_goal: str, tool_call_summaries: list[str], - user_defined_prompts: dict = {}): +async def rank_memories_async(chat_mdl, goal: str, sub_goal: str, tool_call_summaries: list[str], user_defined_prompts: dict = {}): template = PROMPT_JINJA_ENV.from_string(RANK_MEMORY) - system_prompt = template.render(goal=goal, sub_goal=sub_goal, - results=[{"i": i, "content": s} for i, s in enumerate(tool_call_summaries)]) + system_prompt = template.render(goal=goal, sub_goal=sub_goal, results=[{"i": i, "content": s} for i, s in enumerate(tool_call_summaries)]) user_prompt = " → rank: " _, msg = message_fit_in(form_message(system_prompt, user_prompt), chat_mdl.max_length) ans = await chat_mdl.async_chat(msg[0]["content"], msg[1:], stop="<|stop|>") @@ -530,10 +518,7 @@ async def gen_meta_filter(chat_mdl, meta_data: dict, query: str, constraints: di meta_data_structure[key] = list(values.keys()) if isinstance(values, dict) else values sys_prompt = PROMPT_JINJA_ENV.from_string(META_FILTER).render( - current_date=datetime.datetime.today().strftime('%Y-%m-%d'), - metadata_keys=json.dumps(meta_data_structure), - user_question=query, - constraints=json.dumps(constraints) if constraints else None + current_date=datetime.datetime.today().strftime("%Y-%m-%d"), metadata_keys=json.dumps(meta_data_structure), user_question=query, constraints=json.dumps(constraints) if constraints else None ) user_prompt = "Generate filters:" ans = await chat_mdl.async_chat(sys_prompt, [{"role": "user", "content": user_prompt}]) @@ -551,6 +536,7 @@ async def gen_meta_filter(chat_mdl, meta_data: dict, query: str, constraints: di async def gen_json(system_prompt: str, user_prompt: str, chat_mdl, gen_conf={}, max_retry=2): from rag.graphrag.utils import get_llm_cache, set_llm_cache + cached = get_llm_cache(chat_mdl.llm_name, system_prompt, user_prompt, gen_conf) if cached: return json_repair.loads(cached) @@ -577,8 +563,7 @@ TOC_DETECTION = load_prompt("toc_detection") async def detect_table_of_contents(page_1024: list[str], chat_mdl): toc_secs = [] for i, sec in enumerate(page_1024[:22]): - ans = await gen_json(PROMPT_JINJA_ENV.from_string(TOC_DETECTION).render(page_txt=sec), "Only JSON please.", - chat_mdl) + ans = await gen_json(PROMPT_JINJA_ENV.from_string(TOC_DETECTION).render(page_txt=sec), "Only JSON please.", chat_mdl) if toc_secs and not ans["exists"]: break toc_secs.append(sec) @@ -593,8 +578,7 @@ async def extract_table_of_contents(toc_pages, chat_mdl): if not toc_pages: return [] - return await gen_json(PROMPT_JINJA_ENV.from_string(TOC_EXTRACTION).render(toc_page="\n".join(toc_pages)), - "Only JSON please.", chat_mdl) + return await gen_json(PROMPT_JINJA_ENV.from_string(TOC_EXTRACTION).render(toc_page="\n".join(toc_pages)), "Only JSON please.", chat_mdl) async def toc_index_extractor(toc: list[dict], content: str, chat_mdl): @@ -619,8 +603,7 @@ async def toc_index_extractor(toc: list[dict], content: str, chat_mdl): If the title of the section are not in the provided pages, do not add the physical_index to it. Directly return the final JSON structure. Do not output anything else.""" - prompt = tob_extractor_prompt + '\nTable of contents:\n' + json.dumps(toc, ensure_ascii=False, - indent=2) + '\nDocument pages:\n' + content + prompt = tob_extractor_prompt + "\nTable of contents:\n" + json.dumps(toc, ensure_ascii=False, indent=2) + "\nDocument pages:\n" + content return await gen_json(prompt, "Only JSON please.", chat_mdl) @@ -700,10 +683,7 @@ async def table_of_contents_index(toc_arr: list[dict], sections: list[str], chat e = toc_arr[e]["indices"][0] for j in range(st_i, min(e + 1, len(sections))): - ans = await gen_json(PROMPT_JINJA_ENV.from_string(TOC_INDEX).render( - structure=it["structure"], - title=it["title"], - text=sections[j]), "Only JSON please.", chat_mdl) + ans = await gen_json(PROMPT_JINJA_ENV.from_string(TOC_INDEX).render(structure=it["structure"], title=it["title"], text=sections[j]), "Only JSON please.", chat_mdl) if ans["exist"] == "yes": it["indices"].append(j) break @@ -725,9 +705,9 @@ async def check_if_toc_transformation_is_complete(content, toc, chat_mdl): }} Directly return the final JSON structure. Do not output anything else.""" - prompt = prompt + '\n Raw Table of contents:\n' + content + '\n Cleaned Table of contents:\n' + toc + prompt = prompt + "\n Raw Table of contents:\n" + content + "\n Cleaned Table of contents:\n" + toc response = await gen_json(prompt, "Only JSON please.", chat_mdl) - return response['completed'] + return response["completed"] async def toc_transformer(toc_pages, chat_mdl): @@ -749,16 +729,14 @@ async def toc_transformer(toc_pages, chat_mdl): Directly return the final JSON structure, do not output anything else. """ toc_content = "\n".join(toc_pages) - prompt = init_prompt + '\n Given table of contents\n:' + toc_content + prompt = init_prompt + "\n Given table of contents\n:" + toc_content def clean_toc(arr): for a in arr: a["title"] = re.sub(r"[.·….]{2,}", "", a["title"]) last_complete = await gen_json(prompt, "Only JSON please.", chat_mdl) - if_complete = await check_if_toc_transformation_is_complete(toc_content, - json.dumps(last_complete, ensure_ascii=False, indent=2), - chat_mdl) + if_complete = await check_if_toc_transformation_is_complete(toc_content, json.dumps(last_complete, ensure_ascii=False, indent=2), chat_mdl) clean_toc(last_complete) if if_complete == "yes": return last_complete @@ -780,9 +758,7 @@ async def toc_transformer(toc_pages, chat_mdl): break clean_toc(new_complete) last_complete.extend(new_complete) - if_complete = await check_if_toc_transformation_is_complete(toc_content, - json.dumps(last_complete, ensure_ascii=False, - indent=2), chat_mdl) + if_complete = await check_if_toc_transformation_is_complete(toc_content, json.dumps(last_complete, ensure_ascii=False, indent=2), chat_mdl) return last_complete @@ -793,12 +769,7 @@ TOC_LEVELS = load_prompt("assign_toc_levels") async def assign_toc_levels(toc_secs, chat_mdl, gen_conf={"temperature": 0.2}): if not toc_secs: return [] - return await gen_json( - PROMPT_JINJA_ENV.from_string(TOC_LEVELS).render(), - str(toc_secs), - chat_mdl, - gen_conf - ) + return await gen_json(PROMPT_JINJA_ENV.from_string(TOC_LEVELS).render(), str(toc_secs), chat_mdl, gen_conf) TOC_FROM_TEXT_SYSTEM = load_prompt("toc_from_text_system") @@ -812,10 +783,9 @@ async def gen_toc_from_text(txt_info: dict, chat_mdl, callback=None): try: ans = await gen_json( PROMPT_JINJA_ENV.from_string(TOC_FROM_TEXT_SYSTEM).render(), - PROMPT_JINJA_ENV.from_string(TOC_FROM_TEXT_USER).render( - text="\n".join([json.dumps(d, ensure_ascii=False) for d in txt_info["chunks"]])), + PROMPT_JINJA_ENV.from_string(TOC_FROM_TEXT_USER).render(text="\n".join([json.dumps(d, ensure_ascii=False) for d in txt_info["chunks"]])), chat_mdl, - gen_conf={"temperature": 0.0, "top_p": 0.9} + gen_conf={"temperature": 0.0, "top_p": 0.9}, ) txt_info["toc"] = ans if ans and not isinstance(ans, str) else [] except Exception as e: @@ -844,9 +814,7 @@ def split_chunks(chunks, max_length: int): async def run_toc_from_text(chunks, chat_mdl, callback=None): - input_budget = int(chat_mdl.max_length * INPUT_UTILIZATION) - num_tokens_from_string( - TOC_FROM_TEXT_USER + TOC_FROM_TEXT_SYSTEM - ) + input_budget = int(chat_mdl.max_length * INPUT_UTILIZATION) - num_tokens_from_string(TOC_FROM_TEXT_USER + TOC_FROM_TEXT_SYSTEM) input_budget = 1024 if input_budget > 1024 else input_budget chunk_sections = split_chunks(chunks, input_budget) @@ -923,26 +891,32 @@ async def run_toc_from_text(chunks, chat_mdl, callback=None): for _, (toc_item, src_item) in enumerate(zip(toc_with_levels, filtered)): if prune and toc_item.get("level", "0") >= max_lvl: continue - merged.append({ - "level": toc_item.get("level", "0"), - "title": toc_item.get("title", ""), - "chunk_id": src_item.get("chunk_id", ""), - }) + merged.append( + { + "level": toc_item.get("level", "0"), + "title": toc_item.get("title", ""), + "chunk_id": src_item.get("chunk_id", ""), + } + ) return merged TOC_RELEVANCE_SYSTEM = load_prompt("toc_relevance_system") TOC_RELEVANCE_USER = load_prompt("toc_relevance_user") + + async def relevant_chunks_with_toc(query: str, toc: list[dict], chat_mdl, topn: int = 6): import numpy as np + try: ans = await gen_json( PROMPT_JINJA_ENV.from_string(TOC_RELEVANCE_SYSTEM).render(), - PROMPT_JINJA_ENV.from_string(TOC_RELEVANCE_USER).render(query=query, toc_json="[\n%s\n]\n" % "\n".join( - [json.dumps({"level": d["level"], "title": d["title"]}, ensure_ascii=False) for d in toc])), + PROMPT_JINJA_ENV.from_string(TOC_RELEVANCE_USER).render( + query=query, toc_json="[\n%s\n]\n" % "\n".join([json.dumps({"level": d["level"], "title": d["title"]}, ensure_ascii=False) for d in toc]) + ), chat_mdl, - gen_conf={"temperature": 0.0, "top_p": 0.9} + gen_conf={"temperature": 0.0, "top_p": 0.9}, ) id2score = {} for ti, sc in zip(toc, ans): @@ -951,7 +925,7 @@ async def relevant_chunks_with_toc(query: str, toc: list[dict], chat_mdl, topn: for id in ti.get("ids", []): if id not in id2score: id2score[id] = [] - id2score[id].append(sc["score"] / 5.) + id2score[id].append(sc["score"] / 5.0) for id in id2score.keys(): id2score[id] = np.mean(id2score[id]) return [(id, sc) for id, sc in list(id2score.items()) if sc >= 0.3][:topn] @@ -961,6 +935,8 @@ async def relevant_chunks_with_toc(query: str, toc: list[dict], chat_mdl, topn: META_DATA = load_prompt("meta_data") + + async def gen_metadata(chat_mdl, schema: dict, content: str): if not schema: return "" @@ -981,30 +957,25 @@ async def gen_metadata(chat_mdl, schema: dict, content: str): SUFFICIENCY_CHECK = load_prompt("sufficiency_check") + + async def sufficiency_check(chat_mdl, question: str, ret_content: str): try: - return await gen_json( - PROMPT_JINJA_ENV.from_string(SUFFICIENCY_CHECK).render(question=question, retrieved_docs=ret_content), - "Output:\n", - chat_mdl - ) + return await gen_json(PROMPT_JINJA_ENV.from_string(SUFFICIENCY_CHECK).render(question=question, retrieved_docs=ret_content), "Output:\n", chat_mdl) except Exception as e: logging.exception(e) return {} MULTI_QUERIES_GEN = load_prompt("multi_queries_gen") -async def multi_queries_gen(chat_mdl, question: str, query:str, missing_infos:list[str], ret_content: str): + + +async def multi_queries_gen(chat_mdl, question: str, query: str, missing_infos: list[str], ret_content: str): try: return await gen_json( - PROMPT_JINJA_ENV.from_string(MULTI_QUERIES_GEN).render( - original_question=question, - original_query=query, - missing_info="\n - ".join(missing_infos), - retrieved_docs=ret_content - ), + PROMPT_JINJA_ENV.from_string(MULTI_QUERIES_GEN).render(original_question=question, original_query=query, missing_info="\n - ".join(missing_infos), retrieved_docs=ret_content), "Output:\n", - chat_mdl + chat_mdl, ) except Exception as e: logging.exception(e) diff --git a/test/unit_test/agent/component/test_llm_prompt.py b/test/unit_test/agent/component/test_llm_prompt.py new file mode 100644 index 0000000000..ebc011b08b --- /dev/null +++ b/test/unit_test/agent/component/test_llm_prompt.py @@ -0,0 +1,61 @@ +import pytest + +from agent.component.llm import LLM, LLMParam + + +def _llm(prompts=None): + cpn = LLM.__new__(LLM) + cpn._param = LLMParam() + cpn._param.prompts = prompts or [{"role": "user", "content": "User query: {sys.query}"}] + return cpn + + +@pytest.mark.p1 +def test_sys_prompt_and_msg_replaces_trailing_user_instead_of_skipping(): + cpn = _llm() + msg, _ = cpn._sys_prompt_and_msg([{"role": "user", "content": ""}], {"sys.query": "test"}) + assert msg == [{"role": "user", "content": "User query: test"}] + + +@pytest.mark.p1 +def test_sys_prompt_and_msg_keeps_consecutive_configured_prompts(): + cpn = _llm( + [ + {"role": "user", "content": "Context: {sys.query}"}, + {"role": "user", "content": "User query: {sys.query}"}, + ] + ) + msg, _ = cpn._sys_prompt_and_msg([], {"sys.query": "test"}) + assert msg == [ + {"role": "user", "content": "Context: test"}, + {"role": "user", "content": "User query: test"}, + ] + + +@pytest.mark.p1 +def test_validate_fitted_messages_requires_trailing_user(): + err = LLM.validate_fitted_messages( + [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "still here"}, + {"role": "assistant", "content": "reply"}, + ] + ) + assert err and "empty" in err.lower() + + +@pytest.mark.p1 +def test_validate_fitted_messages_rejects_empty_user(): + err = LLM.validate_fitted_messages([{"role": "system", "content": "system"}, {"role": "user", "content": ""}]) + assert err and "empty" in err.lower() + + +@pytest.mark.p1 +def test_fit_messages_uses_default_context_when_max_length_zero(): + msg_fit, err = LLM.fit_messages( + "s" * 845, + [{"role": "user", "content": "User query: test"}], + 0, + ) + assert err is None + assert msg_fit[-1]["content"] == "User query: test" diff --git a/test/unit_test/rag/prompts/test_generator_message_fit_in.py b/test/unit_test/rag/prompts/test_generator_message_fit_in.py index 925c203e68..01b6a35c50 100644 --- a/test/unit_test/rag/prompts/test_generator_message_fit_in.py +++ b/test/unit_test/rag/prompts/test_generator_message_fit_in.py @@ -71,9 +71,7 @@ def _load_generator_module(monkeypatch): template_mod.load_prompt = lambda *_args, **_kwargs: "" monkeypatch.setitem(sys.modules, "rag.prompts.template", template_mod) - spec = importlib.util.spec_from_file_location( - "rag.prompts.generator", repo_root / "rag" / "prompts" / "generator.py" - ) + spec = importlib.util.spec_from_file_location("rag.prompts.generator", repo_root / "rag" / "prompts" / "generator.py") module = importlib.util.module_from_spec(spec) monkeypatch.setitem(sys.modules, "rag.prompts.generator", module) spec.loader.exec_module(module) @@ -149,3 +147,26 @@ def test_message_fit_in_clamps_dominant_last_message_to_budget(monkeypatch): assert used_tokens == 8 assert trimmed[0]["content"] == "" assert trimmed[-1]["content"] == "abcdefgh" + + +@pytest.mark.p1 +def test_message_fit_in_zero_budget_preserves_non_empty_messages(monkeypatch): + generator = _load_generator_module(monkeypatch) + monkeypatch.setattr(generator, "num_tokens_from_string", lambda text: len(text)) + monkeypatch.setattr(generator, "encoder", _CharEncoder()) + + system_len = 8100 + user_content = "User query: test" + messages = [ + {"role": "system", "content": "s" * system_len}, + {"role": "user", "content": user_content}, + ] + expected_total = system_len + len(user_content) + + used_tokens, trimmed = generator.message_fit_in(messages, max_length=0) + + assert expected_total > 861 + assert expected_total < 8192 + assert used_tokens == expected_total + assert trimmed[0]["content"] == "s" * system_len + assert trimmed[-1]["content"] == user_content