Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -12,17 +12,18 @@ from tooluniverse import ToolUniverse
 from gradio import ChatMessage
 from .toolrag import ToolRAGModel
 import torch
+# near the top of txagent.py
 import logging
-
 logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logging.basicConfig(level=logging.INFO)
 
 from .utils import NoRepeatSentenceProcessor, ReasoningTraceChecker, tool_result_format
 
+
 class TxAgent:
     def __init__(self, model_name,
                  rag_model_name,
-                 tool_files_dict=None,
+                 tool_files_dict=None,  # None leads to the default tool files in ToolUniverse
                  enable_finish=True,
                  enable_rag=True,
                  enable_summary=False,
@@ -46,13 +47,10 @@ class TxAgent:
         self.model = None
         self.rag_model = ToolRAGModel(rag_model_name)
         self.tooluniverse = None
-        self.prompt_multi_step = ("You are a highly skilled medical assistant tasked with analyzing medical records in detail. "
-                                 "Provide comprehensive, step-by-step reasoning to identify oversights, including specific diagnoses, "
-                                 "medication conflicts, incomplete assessments, and abnormal results. For each point, include clinical "
-                                 "rationale, standardized screening tools (e.g., PCL-5, SCID-5-PD), and actionable recommendations for "
-                                 "follow-up, ensuring a thorough and precise response.")
+        # self.tool_desc = None
+        self.prompt_multi_step = "You are a helpful assistant that will solve problems through detailed, step-by-step reasoning and actions based on your reasoning. Typically, your actions will use the provided functions. You have access to the following functions."
         self.self_prompt = "Strictly follow the instruction."
-        self.chat_prompt = "You are a helpful assistant to chat with the user."
+        self.chat_prompt = "You are helpful assistant to chat with the user."
         self.enable_finish = enable_finish
         self.enable_rag = enable_rag
         self.enable_summary = enable_summary
@@ -67,57 +65,42 @@ class TxAgent:
         self.enable_checker = enable_checker
         self.additional_default_tools = additional_default_tools
         self.print_self_values()
-        logger.info("TxAgent initialized with model_name=%s, rag_model_name=%s", model_name, rag_model_name)
 
     def init_model(self):
-        logger.info("Initializing model: %s", self.model_name)
-        try:
-            self.load_models()
-            self.load_tooluniverse()
-            self.load_tool_desc_embedding()
-            logger.info("Model initialization complete")
-        except Exception as e:
-            logger.error("Failed to initialize model: %s", e, exc_info=True)
-            raise
+        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)
+            print(f"{attr}: {value}")
 
     def load_models(self, model_name=None):
         if model_name is not None:
             if model_name == self.model_name:
-                logger.debug("Model %s already loaded", model_name)
                 return f"The model {model_name} is already loaded."
             self.model_name = model_name
 
-        logger.debug("Loading model %s", self.model_name)
         self.model = LLM(model=self.model_name)
         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."
 
     def load_tooluniverse(self):
-        logger.debug("Loading tool universe")
         self.tooluniverse = ToolUniverse(tool_files=self.tool_files_dict)
         self.tooluniverse.load_tools()
         special_tools = self.tooluniverse.prepare_tool_prompts(
             self.tooluniverse.tool_category_dicts["special_tools"])
         self.special_tools_name = [tool['name'] for tool in special_tools]
-        logger.debug("Tool universe loaded with %d special tools", len(self.special_tools_name))
 
     def load_tool_desc_embedding(self):
-        logger.debug("Loading tool description embeddings")
         self.rag_model.load_tool_desc_embedding(self.tooluniverse)
-        logger.debug("Tool description embeddings loaded")
 
     def rag_infer(self, query, top_k=5):
-        logger.debug("Running RAG inference with query: %s", query[:50])
         return self.rag_model.rag_infer(query, top_k)
 
     def initialize_tools_prompt(self, call_agent, call_agent_level, message):
-        logger.debug("Initializing tools prompt, call_agent=%s, level=%d", call_agent, call_agent_level)
         picked_tools_prompt = []
         picked_tools_prompt = self.add_special_tools(
             picked_tools_prompt, call_agent=call_agent)
@@ -129,11 +112,9 @@ class TxAgent:
         if not call_agent:
             picked_tools_prompt += self.tool_RAG(
                 message=message, rag_num=self.init_rag_num)
-        logger.debug("Tools prompt initialized with %d tools", len(picked_tools_prompt))
         return picked_tools_prompt, call_agent_level
 
     def initialize_conversation(self, message, conversation=None, history=None):
-        logger.debug("Initializing conversation with message: %s", message[:50])
         if conversation is None:
             conversation = []
 
@@ -142,7 +123,7 @@ class TxAgent:
         if history is not None:
             if len(history) == 0:
                 conversation = []
-                logger.debug("Cleared conversation")
+                print("clear conversation successfully")
             else:
                 for i in range(len(history)):
                     if history[i]['role'] == 'user':
@@ -156,7 +137,7 @@ class TxAgent:
                             {"role": "assistant", "content": history[i]['content']})
 
         conversation.append({"role": "user", "content": message})
-        logger.debug("Conversation initialized with %d messages", len(conversation))
+
         return conversation
 
     def tool_RAG(self, message=None,
@@ -164,8 +145,7 @@ class TxAgent:
                  existing_tools_prompt=[],
                  rag_num=5,
                  return_call_result=False):
-        logger.debug("Running tool RAG, message=%s, rag_num=%d", message[:50] if message else None, rag_num)
-        extra_factor = 30
+        extra_factor = 30  # Factor to retrieve more than rag_num
         if picked_tool_names is None:
             assert picked_tool_names is not None or message is not None
             picked_tool_names = self.rag_infer(
@@ -183,43 +163,39 @@ class TxAgent:
             picked_tools)
         if return_call_result:
             return picked_tools_prompt, picked_tool_names
-        logger.debug("Tool RAG returned %d tools", len(picked_tools_prompt))
         return picked_tools_prompt
 
     def add_special_tools(self, tools, call_agent=False):
-        logger.debug("Adding special tools, call_agent=%s", call_agent)
         if self.enable_finish:
             tools.append(self.tooluniverse.get_one_tool_by_one_name(
                 'Finish', return_prompt=True))
-            logger.debug("Finish tool added")
+            print("Finish tool is added")
         if call_agent:
             tools.append(self.tooluniverse.get_one_tool_by_one_name(
                 'CallAgent', return_prompt=True))
-            logger.debug("CallAgent tool added")
+            print("CallAgent tool is added")
         else:
             if self.enable_rag:
                 tools.append(self.tooluniverse.get_one_tool_by_one_name(
                     'Tool_RAG', return_prompt=True))
-                logger.debug("Tool_RAG tool added")
+                print("Tool_RAG tool is added")
 
             if self.additional_default_tools is not None:
                 for each_tool_name in self.additional_default_tools:
                     tool_prompt = self.tooluniverse.get_one_tool_by_one_name(
                         each_tool_name, return_prompt=True)
                     if tool_prompt is not None:
-                        logger.debug("%s tool added", each_tool_name)
+                        print(f"{each_tool_name} tool is added")
                         tools.append(tool_prompt)
         return tools
 
     def add_finish_tools(self, tools):
-        logger.debug("Adding finish tools")
         tools.append(self.tooluniverse.get_one_tool_by_one_name(
             'Finish', return_prompt=True))
-        logger.debug("Finish tool added")
+        print("Finish tool is added")
         return tools
 
     def set_system_prompt(self, conversation, sys_prompt):
-        logger.debug("Setting system prompt")
         if len(conversation) == 0:
             conversation.append(
                 {"role": "system", "content": sys_prompt})
@@ -235,7 +211,6 @@ class TxAgent:
                           call_agent_level=None,
                           temperature=None):
 
-        logger.debug("Running function call with input: %s", fcall_str[:50])
         function_call_json, message = self.tooluniverse.extract_function_call_json(
             fcall_str, return_message=return_message, verbose=False)
         call_results = []
@@ -243,7 +218,7 @@ class TxAgent:
         if function_call_json is not None:
             if isinstance(function_call_json, list):
                 for i in range(len(function_call_json)):
-                    logger.debug("Tool Call: %s", function_call_json[i])
+                    print("\033[94mTool Call:\033[0m", function_call_json[i])
                     if function_call_json[i]["name"] == 'Finish':
                         special_tool_call = 'Finish'
                         break
@@ -264,7 +239,7 @@ class TxAgent:
                             )
                             call_result = self.run_multistep_agent(
                                 full_message, temperature=temperature,
-                                max_new_tokens=1024, max_token=8192,
+                                max_new_tokens=1024, max_token=99999,
                                 call_agent=False, call_agent_level=call_agent_level)
                             if call_result is None:
                                 call_result = "⚠️ No content returned from sub-agent."
@@ -278,7 +253,7 @@ class TxAgent:
 
                     call_id = self.tooluniverse.call_id_gen()
                     function_call_json[i]["call_id"] = call_id
-                    logger.debug("Tool Call Result: %s", call_result)
+                    print("\033[94mTool Call Result:\033[0m", call_result)
                     call_results.append({
                         "role": "tool",
                         "content": json.dumps({"tool_name": function_call_json[i]["name"], "content": call_result, "call_id": call_id})
@@ -286,15 +261,16 @@ class TxAgent:
         else:
             call_results.append({
                 "role": "tool",
-                "content": json.dumps({"content": "No valid tool call detected; proceeding with analysis."})
+                "content": json.dumps({"content": "Not a valid function call, please check the function call format."})
             })
 
         revised_messages = [{
             "role": "assistant",
-            "content": message.strip() if message else "Processing...",
-            "tool_calls": json.dumps(function_call_json) if function_call_json else None
+            "content": message.strip(),
+            "tool_calls": json.dumps(function_call_json)
         }] + call_results
-        logger.debug("Function call completed, returning %d messages", len(revised_messages))
+
+        # Yield the final result.
         return revised_messages, existing_tools_prompt, special_tool_call
 
     def run_function_call_stream(self, fcall_str,
@@ -306,104 +282,102 @@ class TxAgent:
                                  temperature=None,
                                  return_gradio_history=True):
 
-        logger.debug("Running function call stream with input: %s", fcall_str[:50])
         function_call_json, message = self.tooluniverse.extract_function_call_json(
             fcall_str, return_message=return_message, verbose=False)
         call_results = []
         special_tool_call = ''
-        gradio_history = []
+        if return_gradio_history:
+            gradio_history = []
+        if function_call_json is not None:
+            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"] == '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 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=1024, max_token=99999,
+                                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: The CallAgent has been disabled. Please proceed with your reasoning process to solve this question."
+                    else:
+                        call_result = self.tooluniverse.run_one_function(
+                            function_call_json[i])
 
-        if function_call_json is None:
-            logger.warning("No valid function call JSON extracted")
+                    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':
+                        if function_call_json[i]["name"] == 'Tool_RAG':
+                            gradio_history.append(ChatMessage(role="assistant", content=str(call_result), metadata={
+                                                  "title": "🧰 "+function_call_json[i]['name'], "log": str(function_call_json[i]['arguments'])}))
+                        else:
+                            gradio_history.append(ChatMessage(role="assistant", content=str(call_result), metadata={
+                                                  "title": "⚒️ "+function_call_json[i]['name'], "log": str(function_call_json[i]['arguments'])}))
+        else:
             call_results.append({
                 "role": "tool",
-                "content": json.dumps({"content": "No tool call detected; continuing analysis."})
+                "content": json.dumps({"content": "Not a valid function call, please check the function call format."})
             })
-            if return_gradio_history:
-                gradio_history.append({"role": "assistant", "content": "No specific tool call identified. Proceeding with medical record analysis."})
-                yield [{"role": "assistant", "content": "Processing..."}], existing_tools_prompt or [], special_tool_call, gradio_history
-            return
-
-        if isinstance(function_call_json, list):
-            for i in range(len(function_call_json)):
-                logger.debug("Processing 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"] == '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 = (existing_tools_prompt or []) + new_tools_prompt
-                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)
-                        sub_result = yield from self.run_gradio_chat(
-                            full_message, history=[], temperature=temperature,
-                            max_new_tokens=1024, max_token=8192,
-                            call_agent=False, call_agent_level=call_agent_level,
-                            conversation=None,
-                            sub_agent_task=sub_agent_task)
-                        call_result = sub_result if isinstance(sub_result, str) else "No content from sub-agent."
-                        if '[FinalAnswer]' in call_result:
-                            call_result = call_result.split('[FinalAnswer]')[-1].strip()
-                    else:
-                        call_result = "CallAgent disabled. Proceeding with reasoning."
-                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({"role": "assistant", "content": str(call_result), "metadata": metadata})
 
         revised_messages = [{
             "role": "assistant",
-            "content": message.strip() if message else "Processing...",
-            "tool_calls": json.dumps(function_call_json) if function_call_json else None
+            "content": message.strip(),
+            "tool_calls": json.dumps(function_call_json)
         }] + call_results
 
         if return_gradio_history:
-            logger.debug("Yielding gradio history with %d entries", len(gradio_history))
-            yield revised_messages, existing_tools_prompt or [], special_tool_call, gradio_history
+            return revised_messages, existing_tools_prompt, special_tool_call, gradio_history
         else:
-            yield revised_messages, existing_tools_prompt or [], special_tool_call
-        logger.debug("Function call stream completed")
+            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):
-        logger.debug("Forcing answer due to unfinished reasoning")
-        if conversation[-1]['role'] == 'assistant':
+        if conversation[-1]['role'] == 'assisant':
             conversation.append(
-                {'role': 'tool', 'content': 'Errors occurred; provide a detailed final answer based on current information.'})
+                {'role': 'tool', 'content': 'Errors happen during the function call, please come up with the final answer with the 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]',
+                                          output_begin_string='Since I cannot continue reasoning, I will provide the final answer based on the current information and general knowledge.\n\n[FinalAnswer]',
                                           skip_special_tokens=True,
                                           max_new_tokens=max_new_tokens, max_token=max_token)
-        logger.debug("Forced finish output: %s", last_outputs_str[:100])
+        print(last_outputs_str)
         return last_outputs_str
 
     def run_multistep_agent(self, message: str,
@@ -413,7 +387,16 @@ class TxAgent:
                             max_round: int = 20,
                             call_agent=False,
                             call_agent_level=0) -> str:
-        logger.info("Starting multistep agent with message: %s", message[:50])
+        """
+        Generate a streaming response using the llama3-8b model.
+        Args:
+            message (str): The input message.
+            temperature (float): The temperature for generating the response.
+            max_new_tokens (int): The maximum number of new tokens to generate.
+        Returns:
+            str: The generated response.
+        """
+        print("\033[1;32;40mstart\033[0m")
         picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(
             call_agent, call_agent_level, message)
         conversation = self.initialize_conversation(message)
@@ -432,7 +415,6 @@ class TxAgent:
         try:
             while next_round and current_round < max_round:
                 current_round += 1
-                logger.debug("Round %d", current_round)
                 if len(outputs) > 0:
                     function_call_messages, picked_tools_prompt, special_tool_call = self.run_function_call(
                         last_outputs, return_message=True,
@@ -450,12 +432,12 @@ class TxAgent:
                                 function_call_messages[0]['content'])
                         content = function_call_messages[0]['content']
                         if content is None:
-                            logger.warning("No content after Finish tool call")
                             return "❌ No content returned after Finish tool call."
-                        logger.debug("Returning final content: %s", content[:50])
                         return content.split('[FinalAnswer]')[-1]
 
                     if (self.enable_summary or token_overflow) and not call_agent:
+                        if token_overflow:
+                            print("token_overflow, using summary")
                         enable_summary = True
                     last_status = self.function_result_summary(
                         conversation, status=last_status, enable_summary=enable_summary)
@@ -466,14 +448,15 @@ class TxAgent:
                             function_call_messages))
                     else:
                         next_round = False
-                        content = ''.join(last_outputs).replace("</s>", "")
-                        logger.debug("Returning content: %s", content[:50])
-                        return content
-
+                        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("Internal error in reasoning: %s", wrong_info)
+                        next_round = False
+                        print(
+                            "Internal error in reasoning: " + wrong_info)
                         break
                 last_outputs = []
                 outputs.append("### TxAgent:\n")
@@ -484,7 +467,7 @@ class TxAgent:
                                                                   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")
+                    print("The number of tokens exceeds the maximum limit.")
                     if self.force_finish:
                         return self.get_answer_based_on_unfinished_reasoning(conversation, temperature, max_new_tokens, max_token)
                     else:
@@ -492,22 +475,21 @@ class TxAgent:
                 else:
                     last_outputs.append(last_outputs_str)
             if max_round == current_round:
-                logger.warning("Max rounds exceeded")
+                print("The number of rounds exceeds the maximum limit!")
             if self.force_finish:
                 return self.get_answer_based_on_unfinished_reasoning(conversation, temperature, max_new_tokens, max_token)
             else:
-                logger.debug("No output due to max rounds")
                 return None
 
         except Exception as e:
-            logger.error("Error in multistep agent: %s", e, exc_info=True)
+            print(f"Error: {e}")
             if self.force_finish:
                 return self.get_answer_based_on_unfinished_reasoning(conversation, temperature, max_new_tokens, max_token)
             else:
                 return None
 
     def build_logits_processor(self, messages, llm):
-        logger.debug("Building logits processor")
+        # Use the tokenizer from the LLM instance.
         tokenizer = llm.get_tokenizer()
         if self.avoid_repeat and len(messages) > 2:
             assistant_messages = []
@@ -519,14 +501,14 @@ class TxAgent:
             forbidden_ids = [tokenizer.encode(
                 msg, add_special_tokens=False) for msg in assistant_messages]
             return [NoRepeatSentenceProcessor(forbidden_ids, 5)]
-        return None
+        else:
+            return None
 
     def llm_infer(self, messages, temperature=0.1, tools=None,
                   output_begin_string=None, max_new_tokens=2048,
-                  max_token=8192, skip_special_tokens=True,
+                  max_token=None, skip_special_tokens=True,
                   model=None, tokenizer=None, terminators=None, seed=None, check_token_status=False):
 
-        logger.debug("Running LLM inference with %d messages", len(messages))
         if model is None:
             model = self.model
 
@@ -534,6 +516,7 @@ class TxAgent:
         sampling_params = SamplingParams(
             temperature=temperature,
             max_tokens=max_new_tokens,
+           
             seed=seed if seed is not None else self.seed,
         )
 
@@ -544,38 +527,24 @@ class TxAgent:
 
         if check_token_status and max_token is not None:
             token_overflow = False
-            input_tokens = self.tokenizer.encode(prompt, return_tensors="pt")[0]
-            num_input_tokens = len(input_tokens)
-            if num_input_tokens > max_token:
-                logger.info("Input tokens: %d, max_token: %d", num_input_tokens, max_token)
-                max_prompt_tokens = max_token - max_new_tokens - 100
-                if max_prompt_tokens > 0:
-                    truncated_input = self.tokenizer.decode(input_tokens[:max_prompt_tokens])
-                    prompt = truncated_input
-                    logger.info("Truncated to %d tokens", len(self.tokenizer.encode(prompt, return_tensors='pt')[0]))
-                    token_overflow = True
-                else:
-                    logger.warning("Cannot truncate effectively")
+            num_input_tokens = len(self.tokenizer.encode(
+                prompt, return_tensors="pt")[0])
+            if max_token is not None:
+                if num_input_tokens > max_token:
                     torch.cuda.empty_cache()
                     gc.collect()
+                    print("Number of input tokens before inference:",
+                          num_input_tokens)
+                    logger.info(
+                        "The number of tokens exceeds the maximum limit!!!!")
+                    token_overflow = True
                     return None, token_overflow
-
         output = model.generate(
             prompt,
             sampling_params=sampling_params,
         )
         output = output[0].outputs[0].text
-        # Deduplicate repetitive output
-        if output:
-            lines = output.split('\n')
-            seen = set()
-            deduped_lines = []
-            for line in lines:
-                if line.strip() and line not in seen:
-                    seen.add(line)
-                    deduped_lines.append(line)
-            output = '\n'.join(deduped_lines)
-        logger.debug("LLM output: %s", output[:50])
+        print("\033[92m" + output + "\033[0m")
         if check_token_status and max_token is not None:
             return output, token_overflow
 
@@ -586,7 +555,7 @@ class TxAgent:
                        max_new_tokens: int,
                        max_token: int) -> str:
 
-        logger.info("Starting self agent with message: %s", message[:50])
+        print("\033[1;32;40mstart self agent\033[0m")
         conversation = []
         conversation = self.set_system_prompt(conversation, self.self_prompt)
         conversation.append({"role": "user", "content": message})
@@ -600,7 +569,7 @@ class TxAgent:
                        max_new_tokens: int,
                        max_token: int) -> str:
 
-        logger.info("Starting chat agent with message: %s", message[:50])
+        print("\033[1;32;40mstart chat agent\033[0m")
         conversation = []
         conversation = self.set_system_prompt(conversation, self.chat_prompt)
         conversation.append({"role": "user", "content": message})
@@ -615,7 +584,7 @@ class TxAgent:
                          max_new_tokens: int,
                          max_token: int) -> str:
 
-        logger.info("Starting format agent")
+        print("\033[1;32;40mstart format agent\033[0m")
         if '[FinalAnswer]' in answer:
             possible_final_answer = answer.split("[FinalAnswer]")[-1]
         elif "\n\n" in answer:
@@ -625,13 +594,12 @@ class TxAgent:
         if len(possible_final_answer) == 1:
             choice = possible_final_answer[0]
             if choice in ['A', 'B', 'C', 'D', 'E']:
-                logger.debug("Returning choice: %s", choice)
                 return choice
         elif len(possible_final_answer) > 1:
             if possible_final_answer[1] == ':':
                 choice = possible_final_answer[0]
                 if choice in ['A', 'B', 'C', 'D', 'E']:
-                    logger.debug("Returning choice: %s", choice)
+                    print("choice", choice)
                     return choice
 
         conversation = []
@@ -649,7 +617,7 @@ class TxAgent:
                           temperature: float,
                           max_new_tokens: int,
                           max_token: int) -> str:
-        logger.info("Running summary agent")
+        print("\033[1;32;40mSummarized Tool Result:\033[0m")
         generate_tool_result_summary_training_prompt = """Thought and function calls: 
 {thought_calls}
 Function calls' responses:
@@ -670,11 +638,20 @@ Generate **one summarized sentence** about "function calls' responses" with nece
 
         if '[' in output:
             output = output.split('[')[0]
-        logger.debug("Summary output: %s", output)
         return output
 
     def function_result_summary(self, input_list, status, enable_summary):
-        logger.debug("Running function result summary, enable_summary=%s", enable_summary)
+        """
+        Processes the input list, extracting information from sequences of 'user', 'tool', 'assistant' roles.
+        Supports 'length' and 'step' modes, and skips the last 'k' groups.
+        Parameters:
+            input_list (list): A list of dictionaries containing role and other information.
+            summary_skip_last_k (int): Number of groups to skip from the end. Defaults to 0.
+            summary_context_length (int): The context length threshold for the 'length' mode.
+            last_processed_index (tuple or int): The last processed index.
+        Returns:
+            list: A list of extracted information from valid sequences.
+        """
         if 'tool_call_step' not in status:
             status['tool_call_step'] = 0
 
@@ -722,14 +699,14 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                         this_thought_calls = None
                     else:
                         if len(function_response) != 0:
-                            logger.debug("Generating internal summary")
+                            print("internal summary")
                             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=1024,
-                                max_token=8192
+                                max_token=99999
                             )
 
                             input_list.insert(
@@ -758,7 +735,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                 function_response=function_response,
                 temperature=0.1,
                 max_new_tokens=1024,
-                max_token=8192
+                max_token=99999
             )
 
             tool_calls = json.loads(input_list[last_call_idx]['tool_calls'])
@@ -769,18 +746,19 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                 last_call_idx+1, {'role': 'tool', 'content': result_summary})
             status['summarized_index'] = last_call_idx + 2
 
-        logger.debug("Function result summary completed")
         return status
 
+    # Following are Gradio related functions
+
+    # General update method that accepts any new arguments through kwargs
     def update_parameters(self, **kwargs):
-        logger.debug("Updating parameters: %s", kwargs)
         for key, value in kwargs.items():
             if hasattr(self, key):
                 setattr(self, key, value)
 
+        # Return the updated attributes
         updated_attributes = {key: value for key,
                               value in kwargs.items() if hasattr(self, key)}
-        logger.debug("Updated attributes: %s", updated_attributes)
         return updated_attributes
 
     def run_gradio_chat(self, message: str,
@@ -806,117 +784,90 @@ Generate **one summarized sentence** about "function calls' responses" with nece
         Returns:
             str: Final assistant message.
         """
-        logger.info("[TxAgent] Chat started with message: %s", message[:100])
-        logger.debug("Initial history: %s", [msg["content"][:50] for msg in history] if history else [])
+        logger.debug(f"[TxAgent] Chat started, message: {message[:100]}...")
+        print("\033[1;32;40m[TxAgent] Chat started\033[0m")
 
-        # Yield initial message to ensure UI updates
-        history.append({"role": "assistant", "content": "Starting analysis..."})
-        yield history
-        logger.debug("Yielded initial history")
+        if not message or len(message.strip()) < 5:
+            yield "Please provide a valid message or upload files to analyze."
+            return "Invalid input."
 
-        try:
-            if not message or len(message.strip()) < 5:
-                logger.warning("Invalid message detected")
-                history.append({"role": "assistant", "content": "Please provide a valid message or upload files to analyze."})
-                yield history
-                return "Invalid input."
-
-            if message.startswith("[\U0001f9f0 Tool_RAG") or message.startswith("⚒️"):
-                logger.debug("Skipping tool-related message")
-                yield history
-                return ""
-
-            outputs = []
-            last_outputs = []
-            picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(
-                call_agent, call_agent_level, message)
-            conversation = self.initialize_conversation(
-                message, conversation=conversation, history=history)
-            history = []  # Reset history to avoid duplication
-            logger.debug("Conversation initialized with %d messages", len(conversation))
-
-            next_round = True
-            function_call_messages = []
-            current_round = 0
-            enable_summary = False
-            last_status = {}
-            token_overflow = False
+        if message.startswith("[\U0001f9f0 Tool_RAG") or message.startswith("⚒️"):
+            return ""
+
+        outputs = []
+        outputs_str = ''
+        last_outputs = []
+
+        picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(
+            call_agent,
+            call_agent_level,
+            message)
+
+        conversation = self.initialize_conversation(
+            message,
+            conversation=conversation,
+            history=history)
+        history = []
+
+        next_round = True
+        function_call_messages = []
+        current_round = 0
+        enable_summary = False
+        last_status = {}
+        token_overflow = False
 
-            if self.enable_checker:
-                checker = ReasoningTraceChecker(
-                    message, conversation, init_index=len(conversation))
+        if self.enable_checker:
+            checker = ReasoningTraceChecker(
+                message, conversation, init_index=len(conversation))
 
+        try:
             while next_round and current_round < max_round:
                 current_round += 1
-                logger.debug("Round %d, conversation length: %d", current_round, len(conversation))
+                logger.debug(f"Round {current_round}, conversation length: {len(conversation)}")
 
                 if last_outputs:
-                    function_call_result = yield from self.run_function_call_stream(
-                        last_outputs[0], return_message=True,
+                    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)
 
-                    if not function_call_result:
-                        logger.warning("Empty result from run_function_call_stream")
-                        history.append({"role": "assistant", "content": "Error: Unable to process tool response. Continuing analysis."})
-                        yield history
-                        last_outputs = []
-                        continue
-
-                    function_call_messages, picked_tools_prompt, special_tool_call, current_gradio_history = function_call_result
-
-                    # Convert history to dicts and deduplicate
-                    unique_history = []
-                    seen_contents = set()
-                    for msg in current_gradio_history:
-                        content = msg["content"] if isinstance(msg, dict) else msg.content
-                        if content not in seen_contents:
-                            unique_history.append({"role": "assistant", "content": content})
-                            seen_contents.add(content)
-                    history.extend(unique_history)
-                    logger.debug("Extended history with %d unique messages", len(unique_history))
+                    history.extend(current_gradio_history)
 
                     if special_tool_call == 'Finish' and function_call_messages:
-                        content = function_call_messages[0]['content']
-                        history.append({"role": "assistant", "content": content})
-                        logger.debug("Yielding final history after Finish: %s", content[:50])
                         yield history
                         next_round = False
                         conversation.extend(function_call_messages)
-                        return content
+                        return function_call_messages[0]['content']
 
                     elif special_tool_call in ['RequireClarification', 'DirectResponse']:
-                        last_msg = history[-1] if history else {"role": "assistant", "content": "Response needed."}
-                        history.append({"role": "assistant", "content": last_msg["content"]})
-                        logger.debug("Yielding history for special tool: %s", last_msg["content"][:50])
+                        last_msg = history[-1] if history else ChatMessage(role="assistant", content="Response needed.")
+                        history.append(ChatMessage(role="assistant", content=last_msg.content))
                         yield history
                         next_round = False
-                        return last_msg["content"]
+                        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)
+                        conversation, status=last_status,
+                        enable_summary=enable_summary)
 
                     if function_call_messages:
                         conversation.extend(function_call_messages)
+                        yield history
                     else:
                         next_round = False
-                        content = ''.join(last_outputs).replace("</s>", "")
-                        history.append({"role": "assistant", "content": content})
-                        conversation.append({"role": "assistant", "content": content})
-                        logger.debug("Yielding history with content: %s", content[:50])
-                        yield history
-                        return content
+                        conversation.append({"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 flagged error: %s", wrong_info)
+                        print("Checker flagged reasoning error: ", wrong_info)
                         break
 
                 last_outputs = []
@@ -926,78 +877,87 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                     tools=picked_tools_prompt,
                     skip_special_tokens=False,
                     max_new_tokens=max_new_tokens,
-                    max_token=8192,
+                    max_token=max_token,
                     seed=seed,
                     check_token_status=True)
 
-                logger.debug("llm_infer output: %s, token_overflow: %s", 
-                            last_outputs_str[:50] if last_outputs_str else None, token_overflow)
+                logger.debug(f"llm_infer output: {last_outputs_str[:100] if last_outputs_str else None}, token_overflow: {token_overflow}")
 
                 if last_outputs_str is None:
-                    logger.warning("llm_infer returned None")
-                    error_msg = "Error: Unable to generate response due to token limit. Please reduce input size."
-                    history.append({"role": "assistant", "content": error_msg})
-                    yield history
-                    return error_msg
+                    logger.warning("llm_infer returned None due to token overflow")
+                    if self.force_finish:
+                        last_outputs_str = self.get_answer_based_on_unfinished_reasoning(
+                            conversation, temperature, max_new_tokens, max_token)
+                        history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
+                        yield history
+                        return last_outputs_str
+                    else:
+                        error_msg = "Token limit exceeded. Please reduce input size or increase max_token."
+                        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 isinstance(msg, dict) and "metadata" in msg and msg["metadata"] is not None:
-                        msg["metadata"]['status'] = 'done'
+                    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({"role": "assistant", "content": final_thought.strip()})
-                    history.append({"role": "assistant", "content": "**🧠 Final Analysis:**\n" + final_answer.strip()})
-                    logger.debug("Yielding final analysis: %s", final_answer[:50])
+                    if len(parts) == 2:
+                        final_thought, final_answer = parts
+                    else:
+                        final_thought, final_answer = 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()))
                     yield history
-                    next_round = False
                 else:
-                    history.append({"role": "assistant", "content": last_thought})
-                    logger.debug("Yielding intermediate history: %s", last_thought[:50])
+                    history.append(ChatMessage(role="assistant", content=last_thought))
                     yield history
 
                 last_outputs.append(last_outputs_str)
 
             if next_round:
-                logger.info("Max rounds reached")
                 if self.force_finish:
                     last_outputs_str = self.get_answer_based_on_unfinished_reasoning(
                         conversation, temperature, max_new_tokens, max_token)
                     if '[FinalAnswer]' in last_outputs_str:
                         parts = last_outputs_str.split('[FinalAnswer]', 1)
-                        final_thought, final_answer = parts if len(parts) == 2 else (last_outputs_str, "")
-                        history.append({"role": "assistant", "content": final_thought.strip()})
-                        history.append({"role": "assistant", "content": "**🧠 Final Analysis:**\n" + final_answer.strip()})
+                        if len(parts) == 2:
+                            final_thought, final_answer = parts
+                        else:
+                            final_thought, final_answer = 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:
-                        history.append({"role": "assistant", "content": last_outputs_str.strip()})
-                    logger.debug("Yielding forced final history")
-                    yield history
+                        history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
+                        yield history
                 else:
-                    error_msg = "The number of reasoning rounds exceeded the limit."
-                    history.append({"role": "assistant", "content": error_msg})
-                    logger.debug("Yielding max rounds error")
-                    yield history
-                    return error_msg
+                    yield "The number of reasoning rounds exceeded the limit."
 
         except Exception as e:
-            logger.error("Exception in run_gradio_chat: %s", e, exc_info=True)
+            logger.error(f"Exception in run_gradio_chat: {e}", exc_info=True)
             error_msg = f"An error occurred: {e}"
-            history.append({"role": "assistant", "content": error_msg})
-            logger.debug("Yielding error history: %s", error_msg)
+            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)
                 if '[FinalAnswer]' in last_outputs_str:
                     parts = last_outputs_str.split('[FinalAnswer]', 1)
-                    final_thought, final_answer = parts if len(parts) == 2 else (last_outputs_str, "")
-                    history.append({"role": "assistant", "content": final_thought.strip()})
-                    history.append({"role": "assistant", "content": "**🧠 Final Analysis:**\n" + final_answer.strip()})
+                    if len(parts) == 2:
+                        final_thought, final_answer = parts
+                    else:
+                        final_thought, final_answer = 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:
-                    history.append({"role": "assistant", "content": last_outputs_str.strip()})
-                logger.debug("Yielding forced final history after error")
-                yield history
+                    history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
+                    yield history
             return error_msg