Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -14,33 +14,11 @@ from .toolrag import ToolRAGModel
 import torch
 import logging
 
-# Configure logging
+# Configure logging with a more specific logger name
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger("TxAgent")
 
 from .utils import NoRepeatSentenceProcessor, ReasoningTraceChecker, tool_result_format
-import re
-
-def extract_function_call_json(text: str, return_message: bool = False):
-    """
-    Simple extraction of function call JSON from model output text.
-    """
-    match = re.search(r"\[TOOL_CALLS\](.*?)$", text, re.DOTALL)
-    if not match:
-        return [], text
-
-    try:
-        json_block = match.group(1).strip()
-        func_calls = json.loads(json_block)
-        if isinstance(func_calls, dict):
-            func_calls = [func_calls]
-        message = text.split("[TOOL_CALLS]")[0].strip()
-        if return_message:
-            return func_calls, message
-        return func_calls
-    except Exception as e:
-        logger.error(f"Error parsing function call JSON: {e}")
-        return [], text
 
 class TxAgent:
     def __init__(self, model_name,
@@ -96,20 +74,19 @@ class TxAgent:
                 return f"The model {model_name} is already loaded."
             self.model_name = model_name
 
-        # OPTIMIZED: maximize batching for big GPU
         self.model = LLM(
             model=self.model_name,
             dtype="float16",
             max_model_len=131072,
-            max_num_batched_tokens=65536,  # Bigger batch size
-            gpu_memory_utilization=0.95,    # Use 95% of GPU
+            max_num_batched_tokens=32768,  # Increased for A100 80GB
+            gpu_memory_utilization=0.9,    # Higher utilization for better performance
             trust_remote_code=True
         )
         self.chat_template = Template(self.model.get_tokenizer().chat_template)
         self.tokenizer = self.model.get_tokenizer()
         logger.info(
             "Model %s loaded with max_model_len=%d, max_num_batched_tokens=%d, gpu_memory_utilization=%.2f",
-            self.model_name, 131072, 65536, 0.95
+            self.model_name, 131072, 32768, 0.9
         )
         return f"Model {model_name} loaded successfully."
 
@@ -129,6 +106,7 @@ class TxAgent:
             self.rag_model.load_tool_desc_embedding(self.tooluniverse)
             self.rag_model.save_embeddings(cache_path)
         logger.debug("Tool description embeddings loaded")
+
     def rag_infer(self, query, top_k=5):
         return self.rag_model.rag_infer(query, top_k)
 
@@ -158,6 +136,29 @@ class TxAgent:
         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
+        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_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))
+        if return_call_result:
+            return picked_tools_prompt, picked_tool_names
+        return picked_tools_prompt
+
     def add_special_tools(self, tools, call_agent=False):
         if self.enable_finish:
             tools.append(self.tooluniverse.get_one_tool_by_one_name('Finish', return_prompt=True))
@@ -167,6 +168,11 @@ class TxAgent:
             logger.debug("CallAgent tool added")
         return tools
 
+    def add_finish_tools(self, tools):
+        tools.append(self.tooluniverse.get_one_tool_by_one_name('Finish', return_prompt=True))
+        logger.debug("Finish tool added")
+        return tools
+
     def set_system_prompt(self, conversation, sys_prompt):
         if not conversation:
             conversation.append({"role": "system", "content": sys_prompt})
@@ -174,12 +180,246 @@ 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
+
+        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=131072,
+                                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})
+                    })
+        else:
+            call_results.append({
+                "role": "tool",
+                "content": json.dumps({"content": "Invalid or no function call detected."})
+            })
+
+        revised_messages = [{
+            "role": "assistant",
+            "content": message.strip(),
+            "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
+
+        call_results = []
+        special_tool_call = ''
+        if return_gradio_history:
+            gradio_history = []
+        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=131072,
+                                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))
+        else:
+            call_results.append({
+                "role": "tool",
+                "content": json.dumps({"content": "Invalid or no function call detected."})
+            })
+
+        revised_messages = [{
+            "role": "assistant",
+            "content": message.strip(),
+            "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
+
+    def get_answer_based_on_unfinished_reasoning(self, conversation, temperature, max_new_tokens, max_token, outputs=None):
+        if conversation[-1]['role'] == 'assistant':
+            conversation.append(
+                {'role': 'tool', 'content': 'Errors occurred during function call; provide final answer with current information.'})
+        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])
+        picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(
+            call_agent, call_agent_level, message)
+        conversation = self.initialize_conversation(message)
+        outputs = []
+        last_outputs = []
+        next_round = True
+        current_round = 0
+        token_overflow = False
+        enable_summary = False
+        last_status = {}
+
+        while next_round and current_round < max_round:
+            current_round += 1
+            if len(outputs) > 0:
+                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)
+
+                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]
+
+                if (self.enable_summary or token_overflow) and not call_agent:
+                    enable_summary = True
+                last_status = self.function_result_summary(
+                    conversation, status=last_status, enable_summary=enable_summary)
+
+                if function_call_messages:
+                    conversation.extend(function_call_messages)
+                    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>", "")
+
+            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=2048,
+                max_token=131072,
+                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:
+            logger.warning("Max rounds exceeded")
+        if self.force_finish:
+            return self.get_answer_based_on_unfinished_reasoning(
+                conversation, temperature, max_new_tokens, max_token)
+        return None
+
     def build_logits_processor(self, messages, llm):
         logger.warning("Logits processor disabled due to vLLM V1 limitation")
         return None
 
-    def llm_infer(self, messages, temperature=0.0, tools=None,
-                  output_begin_string=None, max_new_tokens=4096,
+    def llm_infer(self, messages, temperature=0.1, tools=None,
+                  output_begin_string=None, max_new_tokens=512,
                   max_token=131072, skip_special_tokens=True,
                   model=None, tokenizer=None, terminators=None,
                   seed=None, check_token_status=False):
@@ -188,7 +428,7 @@ class TxAgent:
 
         logits_processor = self.build_logits_processor(messages, model)
         sampling_params = SamplingParams(
-            temperature=temperature,  # Force deterministic
+            temperature=temperature,
             max_tokens=max_new_tokens,
             seed=seed if seed is not None else self.seed,
         )
@@ -211,9 +451,13 @@ class TxAgent:
         output = model.generate(prompt, sampling_params=sampling_params)
         output_text = output[0].outputs[0].text
         output_tokens = len(self.tokenizer.encode(output_text, add_special_tokens=False))
-        logger.debug("Inference output (truncated): %s (output tokens: %d)", output_text[:120], output_tokens)
+        logger.debug("Inference output: %s (output tokens: %d)", output_text[:100], output_tokens)
+        torch.cuda.empty_cache()
+        gc.collect()
+        if check_token_status and max_token is not None:
+            return output_text, token_overflow
+        return output_text
 
-        return output_text if not check_token_status else (output_text, token_overflow)
     def run_self_agent(self, message: str,
                        temperature: float,
                        max_new_tokens: int,
@@ -248,10 +492,21 @@ class TxAgent:
                          max_new_tokens: int,
                          max_token: int):
         logger.info("Starting format agent")
-        possible_final_answer = answer.split("[FinalAnswer]")[-1] if "[FinalAnswer]" in answer else answer.strip()
+        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
+        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: " + possible_final_answer})
+        conversation.append({"role": "user", "content": message +
+                            "\nAgent's answer: " + answer + "\nAnswer (must be a letter):"})
         return self.llm_infer(
             messages=conversation,
             temperature=temperature,
@@ -259,17 +514,108 @@ class TxAgent:
             max_new_tokens=max_new_tokens,
             max_token=max_token)
 
-    def build_chat_messages(self, history, user_message):
-        conversation = [{"role": "system", "content": self.prompt_multi_step}]
-        for item in history:
-            conversation.append(item)
-        conversation.append({"role": "user", "content": user_message})
-        return conversation
+    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' l 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)
+        if '[' in output:
+            output = output.split('[')[0]
+        return output
+
+    def function_result_summary(self, input_list, status, enable_summary):
+        if 'tool_call_step' not in status:
+            status['tool_call_step'] = 0
+        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]:
+                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', [])
+
+        function_response = ''
+        idx = status['summarized_index']
+        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=131072)
+                        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:
+                    function_response += input_list[idx]['content']
+                    del input_list[idx]
+                    idx -= 1
+            else:
+                break
+            idx += 1
+
+        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=131072)
+            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})
+            status['summarized_index'] = last_call_idx + 2
+
+        return status
+
+    def update_parameters(self, **kwargs):
+        updated_attributes = {}
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+                updated_attributes[key] = value
+        logger.info("Updated parameters: %s", updated_attributes)
+        return updated_attributes
 
     def run_gradio_chat(self, message: str,
                         history: list,
                         temperature: float,
-                        max_new_tokens: int = 4096,
+                        max_new_tokens: int = 2048,
                         max_token: int = 131072,
                         call_agent: bool = False,
                         conversation: gr.State = None,
@@ -278,38 +624,71 @@ class TxAgent:
                         call_agent_level: int = 0,
                         sub_agent_task: str = None,
                         uploaded_files: list = None):
-        logger.info("Chat started, message: %s", message[:120])
+        logger.info("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
 
-        picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(call_agent, call_agent_level, message)
-        conversation = self.initialize_conversation(message, conversation, history)
-
-        outputs = []
+        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 = []
-        token_overflow = False
+
         next_round = True
         current_round = 0
+        enable_summary = False
+        last_status = {}
+        token_overflow = False
 
         try:
             while next_round and current_round < max_round:
                 current_round += 1
-                logger.info("Starting round %d/%d", current_round, max_round)
-
+                logger.debug("Starting round %d/%d", current_round, max_round)
                 if last_outputs:
-                    function_call_messages, picked_tools_prompt, special_tool_call = self.tool_function_handling(
-                        last_outputs, picked_tools_prompt, message, call_agent, call_agent_level, temperature)
+                    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)
+                    history.extend(current_gradio_history)
 
-                    conversation.extend(function_call_messages)
-                    outputs.append(tool_result_format(function_call_messages))
                     if special_tool_call == 'Finish':
+                        logger.info("Finish tool called, ending chat")
+                        yield history
+                        next_round = False
+                        conversation.extend(function_call_messages)
+                        content = function_call_messages[0]['content']
+                        if content:
+                            return content
+                        return "No content returned after Finish tool call."
+
+                    elif 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))
+                        logger.info("Special tool %s called, ending chat", special_tool_call)
+                        yield history
+                        next_round = False
+                        return last_msg.content
+
+                    if (self.enable_summary or token_overflow) and not call_agent:
+                        enable_summary = True
+                    last_status = self.function_result_summary(
+                        conversation, status=last_status, enable_summary=enable_summary)
+
+                    if function_call_messages:
+                        conversation.extend(function_call_messages)
+                        yield history
+                    else:
                         next_round = False
-                        final_content = function_call_messages[0]["content"]
-                        yield final_content.split('[FinalAnswer]')[-1] if final_content else "⚠️ No content returned."
-                        return
+                        conversation.append({"role": "assistant", "content": ''.join(last_outputs)})
+                        logger.info("No function call messages, ending chat")
+                        return ''.join(last_outputs).replace("</s>", "")
 
-                logger.info("Running model inference...")
+                last_outputs = []
                 last_outputs_str, token_overflow = self.llm_infer(
                     messages=conversation,
                     temperature=temperature,
@@ -321,103 +700,95 @@ class TxAgent:
                     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)
-                        yield last_outputs_str
+                        history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
+                        yield history
                         return last_outputs_str
-                    yield "⚠️ Token limit exceeded."
-                    return
+                    error_msg = "Token limit exceeded."
+                    history.append(ChatMessage(role="assistant", content=error_msg))
+                    yield history
+                    return error_msg
+
+                last_thought = last_outputs_str.split("[TOOL_CALLS]")[0]
+                for msg in history:
+                    if msg.metadata is not None:
+                        msg.metadata['status'] = 'done'
+
+                if '[FinalAnswer]' in last_thought:
+                    parts = last_thought.split('[FinalAnswer]', 1)
+                    final_thought, final_answer = parts if len(parts) == 2 else (last_thought, "")
+                    history.append(ChatMessage(role="assistant", content=final_thought.strip()))
+                    yield history
+                    history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
+                    logger.info("Final answer provided: %s", final_answer[:100])
+                    yield history
+                    next_round = False  # Ensure we exit after final answer
+                    return final_answer
+                else:
+                    history.append(ChatMessage(role="assistant", content=last_thought))
+                    yield history
 
                 last_outputs.append(last_outputs_str)
 
-                if "[FinalAnswer]" in last_outputs_str:
-                    final_answer = last_outputs_str.split("[FinalAnswer]")[-1].strip()
-                    logger.info("Final answer provided: %s", final_answer[:100])
-                    yield final_answer
-                    return
+            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()))
+                    logger.info("Forced final answer: %s", final_answer[:100])
+                    yield history
+                    return final_answer
+                else:
+                    error_msg = "Reasoning rounds exceeded limit."
+                    history.append(ChatMessage(role="assistant", content=error_msg))
+                    yield history
+                    return error_msg
 
         except Exception as e:
-            logger.error("Exception in run_gradio_chat: %s", e)
-            yield f"⚠️ Error occurred: {e}"
-    def tool_function_handling(self, last_outputs, picked_tools_prompt, message, call_agent, call_agent_level, temperature):
-        try:
-            from .utils import extract_function_call_json
-            function_call_json, _ = extract_function_call_json(last_outputs[-1], return_message=True)
-        except Exception as e:
-            logger.warning("Tool call extraction failed: %s", e)
-            function_call_json = []
-
-        function_call_messages = []
-        special_tool_call = ''
-        if function_call_json:
-            for call in function_call_json:
-                if call['name'] == 'Finish':
-                    special_tool_call = 'Finish'
-                # Can handle other special tools here
-
-        return function_call_messages, picked_tools_prompt, special_tool_call
-
-    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])
-        picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(call_agent, call_agent_level, message)
-        conversation = self.initialize_conversation(message)
-
-        outputs = []
-        last_outputs = []
-        next_round = True
-        current_round = 0
-
-        while next_round and current_round < max_round:
-            current_round += 1
-            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=self.seed,
-                check_token_status=True)
-
-            if last_outputs_str is None:
-                logger.warning("Token limit exceeded inside multistep agent")
-                return "⚠️ Token overflow."
-
-            outputs.append(last_outputs_str)
-
-            if "[FinalAnswer]" in last_outputs_str:
-                logger.info("Multistep Final Answer Provided")
-                return last_outputs_str.split("[FinalAnswer]")[-1]
-
-            last_outputs = [last_outputs_str]
-
-        return "⚠️ Max rounds exceeded."
+            logger.error("Exception in run_gradio_chat: %s", e, exc_info=True)
+            error_msg = f"Error: {e}"
+            history.append(ChatMessage(role="assistant", content=error_msg))
+            yield history
+            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()))
+                logger.info("Forced final answer after error: %s", final_answer[:100])
+                yield history
+                return final_answer
+            return error_msg
 
     def run_gradio_chat_batch(self, messages: List[str],
                              temperature: float,
-                             max_new_tokens: int = 4096,
+                             max_new_tokens: int = 2048,
                              max_token: int = 131072,
                              call_agent: bool = False,
                              conversation: List = None,
                              max_round: int = 5,
                              seed: int = None,
                              call_agent_level: int = 0):
-        """Batch processing optimization."""
+        """Run batch inference for multiple messages."""
         logger.info("Starting batch chat for %d messages", len(messages))
         batch_results = []
-
+        
         for message in messages:
+            # Initialize conversation for each message
             conv = self.initialize_conversation(message, conversation, history=None)
             picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(
                 call_agent, call_agent_level, message)
-
+            
+            # Run single inference for simplicity (extend for multi-round if needed)
             output, token_overflow = self.llm_infer(
                 messages=conv,
                 temperature=temperature,
@@ -428,12 +799,12 @@ class TxAgent:
                 seed=seed,
                 check_token_status=True
             )
-
+            
             if output is None:
                 logger.warning("Token limit exceeded for message: %s", message[:100])
-                batch_results.append("⚠️ Token limit exceeded.")
+                batch_results.append("Token limit exceeded.")
             else:
                 batch_results.append(output)
-
+        
         logger.info("Batch chat completed for %d messages", len(messages))
-        return batch_results
+        return batch_results