fix(agent): prevent empty LLM user message after prompt fitting (#16413)

## 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 <taranumwasu@Taranums-MacBook-Pro.local>
Co-authored-by: Cursor <cursoragent@cursor.com>

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

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}>(.*?)</{tag}>", sys_prompt, flags=re.DOTALL|re.IGNORECASE)
+            r = re.search(rf"<{tag}>(.*?)</{tag}>", sys_prompt, flags=re.DOTALL | re.IGNORECASE)
             if not r:
                 continue
             pts[tag.lower()] = r.group(1)
-            sys_prompt = re.sub(rf"<{tag}>(.*?)</{tag}>", "", sys_prompt, flags=re.DOTALL|re.IGNORECASE)
+            sys_prompt = re.sub(rf"<{tag}>(.*?)</{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("<think>") == 0:
-                    last_idx += len("<think>")
-                    return "<think>"
-                elif delta_ans.find("<think>") > 0:
-                    delta_ans = txt[last_idx:last_idx + delta_ans.find("<think>")]
-                    last_idx += delta_ans.find("<think>")
-                    return delta_ans
-                elif delta_ans.endswith("</think>"):
-                    endswith_think = True
-                elif endswith_think:
-                    endswith_think = False
-                    return "</think>"
-
-                last_idx = len(ans)
-                if ans.endswith("</think>"):
-                    last_idx -= len("</think>")
-                return re.sub(r"(<think>|</think>)", "", 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("<think>") == 0:
-                last_idx += len("<think>")
-                return "<think>"
-            elif delta_ans.find("<think>") > 0:
-                delta_ans = txt[last_idx:last_idx + delta_ans.find("<think>")]
-                last_idx += delta_ans.find("<think>")
-                return delta_ans
-            elif delta_ans.endswith("</think>"):
-                endswith_think = True
-            elif endswith_think:
-                endswith_think = False
-                return "</think>"
-
-            last_idx = len(answer)
-            if answer.endswith("</think>"):
-                last_idx -= len("</think>")
-            return re.sub(r"(<think>|</think>)", "", 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."

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]):

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),

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:

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"^.*</think>", "", 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)

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"

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