Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -24,17 +24,17 @@ class TxAgent:
                  rag_model_name,
                  tool_files_dict=None,
                  enable_finish=True,
-                 enable_rag=False,
+                 enable_rag=True,
                  enable_summary=False,
-                 init_rag_num=0,
-                 step_rag_num=0,
+                 init_rag_num=2,  # Reduced for faster initial tool selection
+                 step_rag_num=4,  # Reduced for fewer RAG calls
                  summary_mode='step',
                  summary_skip_last_k=0,
                  summary_context_length=None,
                  force_finish=True,
                  avoid_repeat=True,
                  seed=None,
-                 enable_checker=False,
+                 enable_checker=False,  # Disabled by default for speed
                  enable_chat=False,
                  additional_default_tools=None):
         self.model_name = model_name
@@ -45,9 +45,9 @@ class TxAgent:
         self.model = None
         self.rag_model = ToolRAGModel(rag_model_name)
         self.tooluniverse = None
-        self.prompt_multi_step = "You are a helpful assistant that solves problems through step-by-step reasoning."
-        self.self_prompt = "Strictly follow the instruction."
-        self.chat_prompt = "You are a helpful assistant for user chat."
+        self.prompt_multi_step = "You are a medical assistant solving clinical oversight issues step-by-step using provided tools."
+        self.self_prompt = "Follow instructions precisely."
+        self.chat_prompt = "You are a helpful assistant for clinical queries."
         self.enable_finish = enable_finish
         self.enable_rag = enable_rag
         self.enable_summary = enable_summary
@@ -61,23 +61,28 @@ class TxAgent:
         self.seed = seed
         self.enable_checker = enable_checker
         self.additional_default_tools = additional_default_tools
-        logger.info("TxAgent initialized with model: %s, RAG: %s", model_name, rag_model_name)
+        logger.debug("TxAgent initialized with parameters: %s", self.__dict__)
 
     def init_model(self):
         self.load_models()
         self.load_tooluniverse()
+        self.load_tool_desc_embedding()
+
+    def print_self_values(self):
+        for attr, value in self.__dict__.items():
+            logger.debug("%s: %s", attr, value)
 
     def load_models(self, model_name=None):
-        if model_name is not None:
-            if model_name == self.model_name:
-                return f"The model {model_name} is already loaded."
+        if model_name is not None and model_name == self.model_name:
+            return f"The model {model_name} is already loaded."
+        if model_name:
             self.model_name = model_name
 
-        self.model = LLM(model=self.model_name, dtype="float16", max_model_len=2048, gpu_memory_utilization=0.8)
+        self.model = LLM(model=self.model_name, dtype="float16")  # Enable FP16
         self.chat_template = Template(self.model.get_tokenizer().chat_template)
         self.tokenizer = self.model.get_tokenizer()
         logger.info("Model %s loaded successfully", self.model_name)
-        return f"Model {model_name} loaded successfully."
+        return f"Model {self.model_name} loaded successfully."
 
     def load_tooluniverse(self):
         self.tooluniverse = ToolUniverse(tool_files=self.tool_files_dict)
@@ -88,12 +93,7 @@ class TxAgent:
         logger.debug("ToolUniverse loaded with %d special tools", len(self.special_tools_name))
 
     def load_tool_desc_embedding(self):
-        cache_path = os.path.join(os.path.dirname(self.tool_files_dict["new_tool"]), "tool_embeddings.pkl")
-        if os.path.exists(cache_path):
-            self.rag_model.load_cached_embeddings(cache_path)
-        else:
-            self.rag_model.load_tool_desc_embedding(self.tooluniverse)
-            self.rag_model.save_embeddings(cache_path)
+        self.rag_model.load_tool_desc_embedding(self.tooluniverse)
         logger.debug("Tool description embeddings loaded")
 
     def rag_infer(self, query, top_k=5):
@@ -107,43 +107,39 @@ class TxAgent:
             call_agent_level += 1
             if call_agent_level >= 2:
                 call_agent = False
+
+        if not call_agent and self.enable_rag:
+            picked_tools_prompt += self.tool_RAG(
+                message=message, rag_num=self.init_rag_num)
         return picked_tools_prompt, call_agent_level
 
     def initialize_conversation(self, message, conversation=None, history=None):
         if conversation is None:
             conversation = []
 
-        conversation = self.set_system_prompt(
-            conversation, self.prompt_multi_step)
+        conversation = self.set_system_prompt(conversation, self.prompt_multi_step)
         if history:
-            for i in range(len(history)):
-                if history[i]['role'] == 'user':
-                    conversation.append({"role": "user", "content": history[i]['content']})
-                elif history[i]['role'] == 'assistant':
-                    conversation.append({"role": "assistant", "content": history[i]['content']})
+            conversation.extend(
+                {"role": h['role'], "content": h['content']}
+                for h in history if h['role'] in ['user', 'assistant']
+            )
         conversation.append({"role": "user", "content": message})
         logger.debug("Conversation initialized with %d messages", len(conversation))
         return conversation
 
-    def tool_RAG(self, message=None,
-                 picked_tool_names=None,
-                 existing_tools_prompt=[],
-                 rag_num=0,
-                 return_call_result=False):
-        if not self.enable_rag:
-            return []
-        extra_factor = 10
+    def tool_RAG(self, message=None, picked_tool_names=None,
+                 existing_tools_prompt=None, rag_num=4, return_call_result=False):
+        extra_factor = 10  # Reduced from 30 for efficiency
         if picked_tool_names is None:
-            assert picked_tool_names is not None or message is not None
-            picked_tool_names = self.rag_infer(
-                message, top_k=rag_num * extra_factor)
-
-        picked_tool_names_no_special = [tool for tool in picked_tool_names if tool not in self.special_tools_name]
-        picked_tool_names = picked_tool_names_no_special[:rag_num]
+            picked_tool_names = self.rag_infer(message, top_k=rag_num * extra_factor)
 
+        picked_tool_names = [
+            tool for tool in picked_tool_names
+            if tool not in self.special_tools_name
+        ][:rag_num]
         picked_tools = self.tooluniverse.get_tool_by_name(picked_tool_names)
         picked_tools_prompt = self.tooluniverse.prepare_tool_prompts(picked_tools)
-        logger.debug("Retrieved %d tools via RAG", len(picked_tools_prompt))
+        logger.debug("RAG selected %d tools: %s", len(picked_tool_names), picked_tool_names)
         if return_call_result:
             return picked_tools_prompt, picked_tool_names
         return picked_tools_prompt
@@ -155,6 +151,15 @@ class TxAgent:
         if call_agent:
             tools.append(self.tooluniverse.get_one_tool_by_one_name('CallAgent', return_prompt=True))
             logger.debug("CallAgent tool added")
+        elif self.enable_rag:
+            tools.append(self.tooluniverse.get_one_tool_by_one_name('Tool_RAG', return_prompt=True))
+            logger.debug("Tool_RAG tool added")
+        if self.additional_default_tools:
+            for tool_name in self.additional_default_tools:
+                tool_prompt = self.tooluniverse.get_one_tool_by_one_name(tool_name, return_prompt=True)
+                if tool_prompt:
+                    tools.append(tool_prompt)
+                    logger.debug("%s tool added", tool_name)
         return tools
 
     def add_finish_tools(self, tools):
@@ -169,179 +174,135 @@ class TxAgent:
             conversation[0] = {"role": "system", "content": sys_prompt}
         return conversation
 
-    def run_function_call(self, fcall_str,
-                          return_message=False,
-                          existing_tools_prompt=None,
-                          message_for_call_agent=None,
-                          call_agent=False,
-                          call_agent_level=None,
-                          temperature=None):
-        try:
-            function_call_json, message = self.tooluniverse.extract_function_call_json(
-                fcall_str, return_message=return_message, verbose=False)
-        except Exception as e:
-            logger.error("Tool call parsing failed: %s", e)
-            function_call_json = []
-            message = fcall_str
-
+    def run_function_call(self, fcall_str, return_message=False,
+                          existing_tools_prompt=None, message_for_call_agent=None,
+                          call_agent=False, call_agent_level=None, temperature=None):
+        function_call_json, message = self.tooluniverse.extract_function_call_json(
+            fcall_str, return_message=return_message, verbose=False)
         call_results = []
         special_tool_call = ''
         if function_call_json:
-            if isinstance(function_call_json, list):
-                for i in range(len(function_call_json)):
-                    logger.info("Tool Call: %s", function_call_json[i])
-                    if function_call_json[i]["name"] == 'Finish':
-                        special_tool_call = 'Finish'
-                        break
-                    elif function_call_json[i]["name"] == 'CallAgent':
-                        if call_agent_level < 2 and call_agent:
-                            solution_plan = function_call_json[i]['arguments']['solution']
-                            full_message = (
-                                message_for_call_agent +
-                                "\nYou must follow the following plan to answer the question: " +
-                                str(solution_plan)
-                            )
-                            call_result = self.run_multistep_agent(
-                                full_message, temperature=temperature,
-                                max_new_tokens=512, max_token=2048,
-                                call_agent=False, call_agent_level=call_agent_level)
-                            if call_result is None:
-                                call_result = "⚠️ No content returned from sub-agent."
-                            else:
-                                call_result = call_result.split('[FinalAnswer]')[-1].strip()
-                        else:
-                            call_result = "Error: CallAgent disabled."
-                    else:
-                        call_result = self.tooluniverse.run_one_function(function_call_json[i])
-                    call_id = self.tooluniverse.call_id_gen()
-                    function_call_json[i]["call_id"] = call_id
-                    logger.info("Tool Call Result: %s", call_result)
-                    call_results.append({
-                        "role": "tool",
-                        "content": json.dumps({"tool_name": function_call_json[i]["name"], "content": call_result, "call_id": call_id})
-                    })
+            for func in function_call_json if isinstance(function_call_json, list) else [function_call_json]:
+                logger.debug("Tool Call: %s", func)
+                if func["name"] == 'Finish':
+                    special_tool_call = 'Finish'
+                    break
+                elif func["name"] == 'Tool_RAG':
+                    new_tools_prompt, call_result = self.tool_RAG(
+                        message=message, existing_tools_prompt=existing_tools_prompt,
+                        rag_num=self.step_rag_num, return_call_result=True)
+                    existing_tools_prompt += new_tools_prompt
+                elif func["name"] == 'CallAgent' and call_agent and call_agent_level < 2:
+                    solution_plan = func['arguments']['solution']
+                    full_message = (
+                        message_for_call_agent + "\nFollow this plan: " + str(solution_plan)
+                    )
+                    call_result = self.run_multistep_agent(
+                        full_message, temperature=temperature, max_new_tokens=512,
+                        max_token=2048, call_agent=False, call_agent_level=call_agent_level)
+                    call_result = call_result.split('[FinalAnswer]')[-1].strip() if call_result else "⚠️ No content from sub-agent."
+                else:
+                    call_result = self.tooluniverse.run_one_function(func)
+
+                call_id = self.tooluniverse.call_id_gen()
+                func["call_id"] = call_id
+                logger.debug("Tool Call Result: %s", call_result)
+                call_results.append({
+                    "role": "tool",
+                    "content": json.dumps({"tool_name": func["name"], "content": call_result, "call_id": call_id})
+                })
         else:
             call_results.append({
                 "role": "tool",
-                "content": json.dumps({"content": "Invalid or no function call detected."})
+                "content": json.dumps({"content": "Invalid function call format."})
             })
 
         revised_messages = [{
             "role": "assistant",
-            "content": message.strip(),
+            "content": message.strip() if message else "",
             "tool_calls": json.dumps(function_call_json)
         }] + call_results
         return revised_messages, existing_tools_prompt, special_tool_call
 
-    def run_function_call_stream(self, fcall_str,
-                                 return_message=False,
-                                 existing_tools_prompt=None,
-                                 message_for_call_agent=None,
-                                 call_agent=False,
-                                 call_agent_level=None,
-                                 temperature=None,
-                                 return_gradio_history=True):
-        try:
-            function_call_json, message = self.tooluniverse.extract_function_call_json(
-                fcall_str, return_message=return_message, verbose=False)
-        except Exception as e:
-            logger.error("Tool call parsing failed: %s", e)
-            function_call_json = []
-            message = fcall_str
-
+    def run_function_call_stream(self, fcall_str, return_message=False,
+                                existing_tools_prompt=None, message_for_call_agent=None,
+                                call_agent=False, call_agent_level=None, temperature=None,
+                                return_gradio_history=True):
+        function_call_json, message = self.tooluniverse.extract_function_call_json(
+            fcall_str, return_message=return_message, verbose=False)
         call_results = []
         special_tool_call = ''
-        if return_gradio_history:
-            gradio_history = []
+        gradio_history = [] if return_gradio_history else None
         if function_call_json:
-            if isinstance(function_call_json, list):
-                for i in range(len(function_call_json)):
-                    if function_call_json[i]["name"] == 'Finish':
-                        special_tool_call = 'Finish'
-                        break
-                    elif function_call_json[i]["name"] == 'DirectResponse':
-                        call_result = function_call_json[i]['arguments']['respose']
-                        special_tool_call = 'DirectResponse'
-                    elif function_call_json[i]["name"] == 'RequireClarification':
-                        call_result = function_call_json[i]['arguments']['unclear_question']
-                        special_tool_call = 'RequireClarification'
-                    elif function_call_json[i]["name"] == 'CallAgent':
-                        if call_agent_level < 2 and call_agent:
-                            solution_plan = function_call_json[i]['arguments']['solution']
-                            full_message = (
-                                message_for_call_agent +
-                                "\nYou must follow the following plan to answer the question: " +
-                                str(solution_plan)
-                            )
-                            sub_agent_task = "Sub TxAgent plan: " + str(solution_plan)
-                            call_result = yield from self.run_gradio_chat(
-                                full_message, history=[], temperature=temperature,
-                                max_new_tokens=512, max_token=2048,
-                                call_agent=False, call_agent_level=call_agent_level,
-                                conversation=None, sub_agent_task=sub_agent_task)
-                            if call_result is not None and isinstance(call_result, str):
-                                call_result = call_result.split('[FinalAnswer]')[-1]
-                            else:
-                                call_result = "⚠️ No content returned from sub-agent."
-                        else:
-                            call_result = "Error: CallAgent disabled."
-                    else:
-                        call_result = self.tooluniverse.run_one_function(function_call_json[i])
-                    call_id = self.tooluniverse.call_id_gen()
-                    function_call_json[i]["call_id"] = call_id
-                    call_results.append({
-                        "role": "tool",
-                        "content": json.dumps({"tool_name": function_call_json[i]["name"], "content": call_result, "call_id": call_id})
-                    })
-                    if return_gradio_history and function_call_json[i]["name"] != 'Finish':
-                        metadata = {"title": f"🧰 {function_call_json[i]['name']}", "log": str(function_call_json[i]['arguments'])}
-                        gradio_history.append(ChatMessage(role="assistant", content=str(call_result), metadata=metadata))
+            for func in function_call_json if isinstance(function_call_json, list) else [function_call_json]:
+                if func["name"] == 'Finish':
+                    special_tool_call = 'Finish'
+                    break
+                elif func["name"] == 'Tool_RAG':
+                    new_tools_prompt, call_result = self.tool_RAG(
+                        message=message, existing_tools_prompt=existing_tools_prompt,
+                        rag_num=self.step_rag_num, return_call_result=True)
+                    existing_tools_prompt += new_tools_prompt
+                elif func["name"] == 'DirectResponse':
+                    call_result = func['arguments']['response']
+                    special_tool_call = 'DirectResponse'
+                elif func["name"] == 'RequireClarification':
+                    call_result = func['arguments']['unclear_question']
+                    special_tool_call = 'RequireClarification'
+                elif func["name"] == 'CallAgent' and call_agent and call_agent_level < 2:
+                    solution_plan = func['arguments']['solution']
+                    full_message = (
+                        message_for_call_agent + "\nFollow this plan: " + str(solution_plan)
+                    )
+                    sub_agent_task = "Sub TxAgent plan: " + str(solution_plan)
+                    call_result = yield from self.run_gradio_chat(
+                        full_message, history=[], temperature=temperature,
+                        max_new_tokens=512, max_token=2048, call_agent=False,
+                        call_agent_level=call_agent_level, conversation=None,
+                        sub_agent_task=sub_agent_task)
+                    call_result = call_result.split('[FinalAnswer]')[-1] if call_result else "⚠️ No content from sub-agent."
+                else:
+                    call_result = self.tooluniverse.run_one_function(func)
+
+                call_id = self.tooluniverse.call_id_gen()
+                func["call_id"] = call_id
+                call_results.append({
+                    "role": "tool",
+                    "content": json.dumps({"tool_name": func["name"], "content": call_result, "call_id": call_id})
+                })
+                if return_gradio_history and func["name"] != 'Finish':
+                    title = f"{'🧰' if func['name'] == 'Tool_RAG' else '⚒️'} {func['name']}"
+                    gradio_history.append(ChatMessage(
+                        role="assistant", content=str(call_result),
+                        metadata={"title": title, "log": str(func['arguments'])}
+                    ))
         else:
             call_results.append({
                 "role": "tool",
-                "content": json.dumps({"content": "Invalid or no function call detected."})
+                "content": json.dumps({"content": "Invalid function call format."})
             })
 
         revised_messages = [{
             "role": "assistant",
-            "content": message.strip(),
+            "content": message.strip() if message else "",
             "tool_calls": json.dumps(function_call_json)
         }] + call_results
-        if return_gradio_history:
-            return revised_messages, existing_tools_prompt, special_tool_call, gradio_history
-        return revised_messages, existing_tools_prompt, special_tool_call
+        return revised_messages, existing_tools_prompt, special_tool_call, gradio_history
 
-    def get_answer_based_on_unfinished_reasoning(self, conversation, temperature, max_new_tokens, max_token, outputs=None):
-        # Truncate conversation to fit within max_token
-        tokenized = self.tokenizer.encode(json.dumps(conversation), add_special_tokens=False)
-        if len(tokenized) > max_token - 100:
-            logger.warning("Truncating conversation to fit max_token=%d", max_token)
-            while len(tokenized) > max_token - 100 and len(conversation) > 1:
-                conversation.pop(1)  # Keep system prompt and latest message
-                tokenized = self.tokenizer.encode(json.dumps(conversation), add_special_tokens=False)
+    def get_answer_based_on_unfinished_reasoning(self, conversation, temperature, max_new_tokens, max_token):
         if conversation[-1]['role'] == 'assistant':
             conversation.append(
-                {'role': 'tool', 'content': 'Errors occurred during function call; provide final answer with current information.'})
+                {'role': 'tool', 'content': 'Errors occurred; provide final answer with current info.'})
         finish_tools_prompt = self.add_finish_tools([])
-        last_outputs_str = self.llm_infer(
-            messages=conversation,
-            temperature=temperature,
-            tools=finish_tools_prompt,
-            output_begin_string='[FinalAnswer]',
-            skip_special_tokens=True,
-            max_new_tokens=max_new_tokens,
-            max_token=max_token)
-        logger.info("Unfinished reasoning answer: %s", last_outputs_str[:100])
-        return last_outputs_str
-
-    def run_multistep_agent(self, message: str,
-                            temperature: float,
-                            max_new_tokens: int,
-                            max_token: int,
-                            max_round: int = 5,
-                            call_agent=False,
-                            call_agent_level=0):
-        logger.info("Starting multistep agent for message: %s", message[:100])
+        output = self.llm_infer(
+            messages=conversation, temperature=temperature, tools=finish_tools_prompt,
+            output_begin_string='[FinalAnswer]', max_new_tokens=max_new_tokens, max_token=max_token)
+        logger.debug("Unfinished reasoning output: %s", output)
+        return output
+
+    def run_multistep_agent(self, message: str, temperature: float, max_new_tokens: int,
+                            max_token: int, max_round: int = 10, call_agent=False, call_agent_level=0):
+        logger.debug("Starting multistep agent for message: %s", message[:100])
         picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(
             call_agent, call_agent_level, message)
         conversation = self.initialize_conversation(message)
@@ -353,24 +314,22 @@ class TxAgent:
         enable_summary = False
         last_status = {}
 
+        if self.enable_checker:
+            checker = ReasoningTraceChecker(message, conversation)
+
         while next_round and current_round < max_round:
             current_round += 1
-            if len(outputs) > 0:
+            if last_outputs:
                 function_call_messages, picked_tools_prompt, special_tool_call = self.run_function_call(
-                    last_outputs, return_message=True,
-                    existing_tools_prompt=picked_tools_prompt,
-                    message_for_call_agent=message,
-                    call_agent=call_agent,
-                    call_agent_level=call_agent_level,
-                    temperature=temperature)
+                    last_outputs, return_message=True, existing_tools_prompt=picked_tools_prompt,
+                    message_for_call_agent=message, call_agent=call_agent,
+                    call_agent_level=call_agent_level, temperature=temperature)
 
                 if special_tool_call == 'Finish':
                     next_round = False
                     conversation.extend(function_call_messages)
                     content = function_call_messages[0]['content']
-                    if content is None:
-                        return "❌ No content returned after Finish tool call."
-                    return content.split('[FinalAnswer]')[-1]
+                    return content.split('[FinalAnswer]')[-1] if content else "❌ No content after Finish."
 
                 if (self.enable_summary or token_overflow) and not call_agent:
                     enable_summary = True
@@ -382,28 +341,26 @@ class TxAgent:
                     outputs.append(tool_result_format(function_call_messages))
                 else:
                     next_round = False
-                    conversation.extend([{"role": "assistant", "content": ''.join(last_outputs)}])
                     return ''.join(last_outputs).replace("</s>", "")
 
+            if self.enable_checker:
+                good_status, wrong_info = checker.check_conversation()
+                if not good_status:
+                    logger.warning("Checker error: %s", wrong_info)
+                    break
+
             last_outputs = []
-            outputs.append("### TxAgent:\n")
             last_outputs_str, token_overflow = self.llm_infer(
-                messages=conversation,
-                temperature=temperature,
-                tools=picked_tools_prompt,
-                skip_special_tokens=False,
-                max_new_tokens=max_new_tokens,
-                max_token=max_token,
-                check_token_status=True)
+                messages=conversation, temperature=temperature, tools=picked_tools_prompt,
+                max_new_tokens=max_new_tokens, max_token=max_token, check_token_status=True)
             if last_outputs_str is None:
-                logger.warning("Token limit exceeded")
                 if self.force_finish:
                     return self.get_answer_based_on_unfinished_reasoning(
                         conversation, temperature, max_new_tokens, max_token)
                 return "❌ Token limit exceeded."
             last_outputs.append(last_outputs_str)
 
-        if max_round == current_round:
+        if current_round >= max_round:
             logger.warning("Max rounds exceeded")
         if self.force_finish:
             return self.get_answer_based_on_unfinished_reasoning(
@@ -413,16 +370,16 @@ class TxAgent:
     def build_logits_processor(self, messages, llm):
         tokenizer = llm.get_tokenizer()
         if self.avoid_repeat and len(messages) > 2:
-            assistant_messages = [msg['content'] for msg in messages[-3:] if msg['role'] == 'assistant'][:2]
+            assistant_messages = [
+                m['content'] for m in messages[-3:] if m['role'] == 'assistant'
+            ][:2]
             forbidden_ids = [tokenizer.encode(msg, add_special_tokens=False) for msg in assistant_messages]
             return [NoRepeatSentenceProcessor(forbidden_ids, 5)]
         return None
 
-    def llm_infer(self, messages, temperature=0.1, tools=None,
-                  output_begin_string=None, max_new_tokens=512,
-                  max_token=2048, skip_special_tokens=True,
-                  model=None, tokenizer=None, terminators=None,
-                  seed=None, check_token_status=False):
+    def llm_infer(self, messages, temperature=0.1, tools=None, output_begin_string=None,
+                  max_new_tokens=512, max_token=2048, skip_special_tokens=True,
+                  model=None, tokenizer=None, terminators=None, seed=None, check_token_status=False):
         if model is None:
             model = self.model
 
@@ -431,108 +388,73 @@ class TxAgent:
             temperature=temperature,
             max_tokens=max_new_tokens,
             seed=seed if seed is not None else self.seed,
+            logits_processors=logits_processor
         )
 
-        prompt = self.chat_template.render(
-            messages=messages, tools=tools, add_generation_prompt=True)
-        if output_begin_string is not None:
+        prompt = self.chat_template.render(messages=messages, tools=tools, add_generation_prompt=True)
+        if output_begin_string:
             prompt += output_begin_string
 
-        if check_token_status and max_token is not None:
-            token_overflow = False
+        if check_token_status and max_token:
             num_input_tokens = len(self.tokenizer.encode(prompt, return_tensors="pt")[0])
             if num_input_tokens > max_token:
                 torch.cuda.empty_cache()
                 gc.collect()
                 logger.info("Token overflow: %d > %d", num_input_tokens, max_token)
                 return None, True
+            logger.debug("Input tokens: %d", num_input_tokens)
 
         output = model.generate(prompt, sampling_params=sampling_params)
         output = output[0].outputs[0].text
         logger.debug("Inference output: %s", output[:100])
-        torch.cuda.empty_cache()
-        gc.collect()
-        if check_token_status and max_token is not None:
-            return output, token_overflow
+        torch.cuda.empty_cache()  # Clear CUDA cache
+        if check_token_status:
+            return output, False
         return output
 
-    def run_self_agent(self, message: str,
-                       temperature: float,
-                       max_new_tokens: int,
-                       max_token: int):
-        logger.info("Starting self agent")
+    def run_self_agent(self, message: str, temperature: float, max_new_tokens: int, max_token: int):
+        logger.debug("Starting self agent")
         conversation = self.set_system_prompt([], self.self_prompt)
         conversation.append({"role": "user", "content": message})
-        return self.llm_infer(
-            messages=conversation,
-            temperature=temperature,
-            tools=None,
-            max_new_tokens=max_new_tokens,
-            max_token=max_token)
-
-    def run_chat_agent(self, message: str,
-                       temperature: float,
-                       max_new_tokens: int,
-                       max_token: int):
-        logger.info("Starting chat agent")
+        return self.llm_infer(messages=conversation, temperature=temperature,
+                              max_new_tokens=max_new_tokens, max_token=max_token)
+
+    def run_chat_agent(self, message: str, temperature: float, max_new_tokens: int, max_token: int):
+        logger.debug("Starting chat agent")
         conversation = self.set_system_prompt([], self.chat_prompt)
         conversation.append({"role": "user", "content": message})
-        return self.llm_infer(
-            messages=conversation,
-            temperature=temperature,
-            tools=None,
-            max_new_tokens=max_new_tokens,
-            max_token=max_token)
-
-    def run_format_agent(self, message: str,
-                         answer: str,
-                         temperature: float,
-                         max_new_tokens: int,
-                         max_token: int):
-        logger.info("Starting format agent")
+        return self.llm_infer(messages=conversation, temperature=temperature,
+                              max_new_tokens=max_new_tokens, max_token=max_token)
+
+    def run_format_agent(self, message: str, answer: str, temperature: float, max_new_tokens: int, max_token: int):
+        logger.debug("Starting format agent")
         if '[FinalAnswer]' in answer:
             possible_final_answer = answer.split("[FinalAnswer]")[-1]
         elif "\n\n" in answer:
             possible_final_answer = answer.split("\n\n")[-1]
         else:
             possible_final_answer = answer.strip()
-        if len(possible_final_answer) == 1 and possible_final_answer in ['A', 'B', 'C', 'D', 'E']:
-            return possible_final_answer
+
+        if len(possible_final_answer) >= 1 and possible_final_answer[0] in ['A', 'B', 'C', 'D', 'E']:
+            return possible_final_answer[0]
         elif len(possible_final_answer) > 1 and possible_final_answer[1] == ':' and possible_final_answer[0] in ['A', 'B', 'C', 'D', 'E']:
             return possible_final_answer[0]
 
         conversation = self.set_system_prompt(
-            [], "Transform the agent's answer to a single letter: 'A', 'B', 'C', 'D'.")
-        conversation.append({"role": "user", "content": message +
-                            "\nAgent's answer: " + answer + "\nAnswer (must be a letter):"})
-        return self.llm_infer(
-            messages=conversation,
-            temperature=temperature,
-            tools=None,
-            max_new_tokens=max_new_tokens,
-            max_token=max_token)
-
-    def run_summary_agent(self, thought_calls: str,
-                          function_response: str,
-                          temperature: float,
-                          max_new_tokens: int,
-                          max_token: int):
-        logger.info("Summarizing tool result")
-        prompt = f"""Thought and function calls: 
-{thought_calls}
-Function calls' responses:
-\"\"\"
-{function_response}
-\"\"\"
-Summarize the function calls' responses in one sentence with all necessary information.
-"""
+            [], "Transform the answer to a single letter: 'A', 'B', 'C', 'D', or 'E'.")
+        conversation.append({"role": "user", "content": f"Original: {message}\nAnswer: {answer}\nFinal answer (letter):"})
+        return self.llm_infer(messages=conversation, temperature=temperature,
+                              max_new_tokens=max_new_tokens, max_token=max_token)
+
+    def run_summary_agent(self, thought_calls: str, function_response: str,
+                         temperature: float, max_new_tokens: int, max_token: int):
+        logger.debug("Starting summary agent")
+        prompt = f"""Thought and function calls: {thought_calls}
+Function responses: {function_response}
+Summarize the function responses in one sentence with all necessary information."""
         conversation = [{"role": "user", "content": prompt}]
-        output = self.llm_infer(
-            messages=conversation,
-            temperature=temperature,
-            tools=None,
-            max_new_tokens=max_new_tokens,
-            max_token=max_token)
+        output = self.llm_infer(messages=conversation, temperature=temperature,
+                                max_new_tokens=max_new_tokens, max_token=max_token)
         if '[' in output:
             output = output.split('[')[0]
         return output
@@ -540,43 +462,55 @@ Summarize the function calls' responses in one sentence with all necessary infor
     def function_result_summary(self, input_list, status, enable_summary):
         if 'tool_call_step' not in status:
             status['tool_call_step'] = 0
+        if 'step' not in status:
+            status['step'] = 0
+        status['step'] += 1
+
         for idx in range(len(input_list)):
             pos_id = len(input_list) - idx - 1
             if input_list[pos_id]['role'] == 'assistant' and 'tool_calls' in input_list[pos_id]:
+                if 'Tool_RAG' in str(input_list[pos_id]['tool_calls']):
+                    status['tool_call_step'] += 1
                 break
 
-        status['step'] = status.get('step', 0) + 1
         if not enable_summary:
             return status
 
-        status['summarized_index'] = status.get('summarized_index', 0)
-        status['summarized_step'] = status.get('summarized_step', 0)
-        status['previous_length'] = status.get('previous_length', 0)
-        status['history'] = status.get('history', [])
+        if 'summarized_index' not in status:
+            status['summarized_index'] = 0
+        if 'summarized_step' not in status:
+            status['summarized_step'] = 0
+        if 'previous_length' not in status:
+            status['previous_length'] = 0
+        if 'history' not in status:
+            status['history'] = []
+
+        status['history'].append(
+            self.summary_mode == 'step' and status['summarized_step'] < status['step'] - status['tool_call_step'] - self.summary_skip_last_k)
 
-        function_response = ''
         idx = status['summarized_index']
+        function_response = ''
         this_thought_calls = None
-
         while idx < len(input_list):
             if (self.summary_mode == 'step' and status['summarized_step'] < status['step'] - status['tool_call_step'] - self.summary_skip_last_k) or \
                (self.summary_mode == 'length' and status['previous_length'] > self.summary_context_length):
                 if input_list[idx]['role'] == 'assistant':
-                    if function_response:
-                        status['summarized_step'] += 1
-                        result_summary = self.run_summary_agent(
-                            thought_calls=this_thought_calls,
-                            function_response=function_response,
-                            temperature=0.1,
-                            max_new_tokens=512,
-                            max_token=2048)
-                        input_list.insert(last_call_idx + 1, {'role': 'tool', 'content': result_summary})
-                        status['summarized_index'] = last_call_idx + 2
-                        idx += 1
-                    last_call_idx = idx
-                    this_thought_calls = input_list[idx]['content'] + input_list[idx]['tool_calls']
-                    function_response = ''
-                elif input_list[idx]['role'] == 'tool' and this_thought_calls is not None:
+                    if 'Tool_RAG' in str(input_list[idx]['tool_calls']):
+                        this_thought_calls = None
+                    else:
+                        if function_response:
+                            status['summarized_step'] += 1
+                            result_summary = self.run_summary_agent(
+                                thought_calls=this_thought_calls, function_response=function_response,
+                                temperature=0.1, max_new_tokens=512, max_token=2048)
+                            input_list.insert(
+                                last_call_idx + 1, {'role': 'tool', 'content': result_summary})
+                            status['summarized_index'] = last_call_idx + 2
+                            idx += 1
+                        last_call_idx = idx
+                        this_thought_calls = input_list[idx]['content'] + input_list[idx]['tool_calls']
+                        function_response = ''
+                elif input_list[idx]['role'] == 'tool' and this_thought_calls:
                     function_response += input_list[idx]['content']
                     del input_list[idx]
                     idx -= 1
@@ -587,16 +521,14 @@ Summarize the function calls' responses in one sentence with all necessary infor
         if function_response:
             status['summarized_step'] += 1
             result_summary = self.run_summary_agent(
-                thought_calls=this_thought_calls,
-                function_response=function_response,
-                temperature=0.1,
-                max_new_tokens=512,
-                max_token=2048)
+                thought_calls=this_thought_calls, function_response=function_response,
+                temperature=0.1, max_new_tokens=512, max_token=2048)
             tool_calls = json.loads(input_list[last_call_idx]['tool_calls'])
             for tool_call in tool_calls:
                 del tool_call['call_id']
             input_list[last_call_idx]['tool_calls'] = json.dumps(tool_calls)
-            input_list.insert(last_call_idx + 1, {'role': 'tool', 'content': result_summary})
+            input_list.insert(
+                last_call_idx + 1, {'role': 'tool', 'content': result_summary})
             status['summarized_index'] = last_call_idx + 2
 
         return status
@@ -607,32 +539,27 @@ Summarize the function calls' responses in one sentence with all necessary infor
             if hasattr(self, key):
                 setattr(self, key, value)
                 updated_attributes[key] = value
-        logger.info("Updated parameters: %s", updated_attributes)
+        logger.debug("Updated parameters: %s", updated_attributes)
         return updated_attributes
 
-    def run_gradio_chat(self, message: str,
-                        history: list,
-                        temperature: float,
-                        max_new_tokens: int,
-                        max_token: int,
-                        call_agent: bool,
-                        conversation: gr.State,
-                        max_round: int = 5,
-                        seed: int = None,
-                        call_agent_level: int = 0,
-                        sub_agent_task: str = None,
+    def run_gradio_chat(self, message: str, history: list, temperature: float,
+                        max_new_tokens: int, max_token: int, call_agent: bool,
+                        conversation: gr.State, max_round: int = 10, seed: int = None,
+                        call_agent_level: int = 0, sub_agent_task: str = None,
                         uploaded_files: list = None):
-        logger.info("Chat started, message: %s", message[:100])
+        logger.debug("Chat started, message: %s", message[:100])
         if not message or len(message.strip()) < 5:
             yield "Please provide a valid message or upload files to analyze."
             return
 
+        if message.startswith("[\U0001f9f0 Tool_RAG") or message.startswith("⚒️"):
+            return
+
         picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(
             call_agent, call_agent_level, message)
         conversation = self.initialize_conversation(
             message, conversation, history)
         history = []
-        last_outputs = []  # Initialize last_outputs to avoid UnboundLocalError
 
         next_round = True
         current_round = 0
@@ -640,17 +567,18 @@ Summarize the function calls' responses in one sentence with all necessary infor
         last_status = {}
         token_overflow = False
 
+        if self.enable_checker:
+            checker = ReasoningTraceChecker(message, conversation, init_index=len(conversation))
+
         try:
             while next_round and current_round < max_round:
                 current_round += 1
-                if last_outputs:  # Process previous outputs if any
+                last_outputs = []
+                if last_outputs:
                     function_call_messages, picked_tools_prompt, special_tool_call, current_gradio_history = yield from self.run_function_call_stream(
-                        last_outputs, return_message=True,
-                        existing_tools_prompt=picked_tools_prompt,
-                        message_for_call_agent=message,
-                        call_agent=call_agent,
-                        call_agent_level=call_agent_level,
-                        temperature=temperature)
+                        last_outputs, return_message=True, existing_tools_prompt=picked_tools_prompt,
+                        message_for_call_agent=message, call_agent=call_agent,
+                        call_agent_level=call_agent_level, temperature=temperature)
                     history.extend(current_gradio_history)
 
                     if special_tool_call == 'Finish':
@@ -659,7 +587,7 @@ Summarize the function calls' responses in one sentence with all necessary infor
                         conversation.extend(function_call_messages)
                         return function_call_messages[0]['content']
 
-                    elif special_tool_call in ['RequireClarification', 'DirectResponse']:
+                    if special_tool_call in ['RequireClarification', 'DirectResponse']:
                         last_msg = history[-1] if history else ChatMessage(role="assistant", content="Response needed.")
                         history.append(ChatMessage(role="assistant", content=last_msg.content))
                         yield history
@@ -676,22 +604,19 @@ Summarize the function calls' responses in one sentence with all necessary infor
                         yield history
                     else:
                         next_round = False
-                        conversation.append({"role": "assistant", "content": ''.join(last_outputs)})
                         return ''.join(last_outputs).replace("</s>", "")
 
-                last_outputs = []
+                if self.enable_checker:
+                    good_status, wrong_info = checker.check_conversation()
+                    if not good_status:
+                        logger.warning("Checker error: %s", wrong_info)
+                        break
+
                 last_outputs_str, token_overflow = self.llm_infer(
-                    messages=conversation,
-                    temperature=temperature,
-                    tools=picked_tools_prompt,
-                    skip_special_tokens=False,
-                    max_new_tokens=max_new_tokens,
-                    max_token=max_token,
-                    seed=seed,
-                    check_token_status=True)
+                    messages=conversation, temperature=temperature, tools=picked_tools_prompt,
+                    max_new_tokens=max_new_tokens, max_token=max_token, seed=seed, check_token_status=True)
 
                 if last_outputs_str is None:
-                    logger.warning("Token limit exceeded")
                     if self.force_finish:
                         last_outputs_str = self.get_answer_based_on_unfinished_reasoning(
                             conversation, temperature, max_new_tokens, max_token)
@@ -705,7 +630,7 @@ Summarize the function calls' responses in one sentence with all necessary infor
 
                 last_thought = last_outputs_str.split("[TOOL_CALLS]")[0]
                 for msg in history:
-                    if msg.metadata is not None:
+                    if msg.metadata:
                         msg.metadata['status'] = 'done'
 
                 if '[FinalAnswer]' in last_thought:
@@ -721,18 +646,15 @@ Summarize the function calls' responses in one sentence with all necessary infor
 
                 last_outputs.append(last_outputs_str)
 
-            if next_round:
-                if self.force_finish:
-                    last_outputs_str = self.get_answer_based_on_unfinished_reasoning(
-                        conversation, temperature, max_new_tokens, max_token)
-                    parts = last_outputs_str.split('[FinalAnswer]', 1)
-                    final_thought, final_answer = parts if len(parts) == 2 else (last_outputs_str, "")
-                    history.append(ChatMessage(role="assistant", content=final_thought.strip()))
-                    yield history
-                    history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
-                    yield history
-                else:
-                    yield "Reasoning rounds exceeded limit."
+            if next_round and self.force_finish:
+                last_outputs_str = self.get_answer_based_on_unfinished_reasoning(
+                    conversation, temperature, max_new_tokens, max_token)
+                parts = last_outputs_str.split('[FinalAnswer]', 1)
+                final_thought, final_answer = parts if len(parts) == 2 else (last_outputs_str, "")
+                history.append(ChatMessage(role="assistant", content=final_thought.strip()))
+                yield history
+                history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
+                yield history
 
         except Exception as e:
             logger.error("Exception in run_gradio_chat: %s", e, exc_info=True)