Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -1,4 +1,3 @@
-import gradio as gr
 import os
 import sys
 import json
@@ -6,13 +5,13 @@ import gc
 import numpy as np
 from vllm import LLM, SamplingParams
 from jinja2 import Template
-from typing import List
+from typing import List, Dict, Optional, Union, Generator
 import types
 from tooluniverse import ToolUniverse
-from gradio import ChatMessage
 from .toolrag import ToolRAGModel
 import torch
 import logging
+from datetime import datetime
 
 # Configure logging with a more specific logger name
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -21,28 +20,50 @@ logger = logging.getLogger("TxAgent")
 from .utils import NoRepeatSentenceProcessor, ReasoningTraceChecker, tool_result_format
 
 class TxAgent:
-    def __init__(self, model_name,
-                 rag_model_name,
-                 tool_files_dict=None,
-                 enable_finish=True,
-                 enable_rag=False,
-                 enable_summary=False,
-                 init_rag_num=0,
-                 step_rag_num=0,
-                 summary_mode='step',
-                 summary_skip_last_k=0,
-                 summary_context_length=None,
-                 force_finish=True,
-                 avoid_repeat=True,
-                 seed=None,
-                 enable_checker=False,
-                 enable_chat=False,
-                 additional_default_tools=None):
+    def __init__(self, 
+                 model_name: str,
+                 rag_model_name: str,
+                 tool_files_dict: Optional[Dict] = None,
+                 enable_finish: bool = True,
+                 enable_rag: bool = False,
+                 enable_summary: bool = False,
+                 init_rag_num: int = 0,
+                 step_rag_num: int = 0,
+                 summary_mode: str = 'step',
+                 summary_skip_last_k: int = 0,
+                 summary_context_length: Optional[int] = None,
+                 force_finish: bool = True,
+                 avoid_repeat: bool = True,
+                 seed: Optional[int] = None,
+                 enable_checker: bool = False,
+                 enable_chat: bool = False,
+                 additional_default_tools: Optional[List] = None):
+        """
+        Initialize the TxAgent with specified configuration.
+        
+        Args:
+            model_name: Name of the main LLM model
+            rag_model_name: Name of the RAG model
+            tool_files_dict: Dictionary of tool files
+            enable_finish: Whether to enable the Finish tool
+            enable_rag: Whether to enable RAG functionality
+            enable_summary: Whether to enable summarization
+            init_rag_num: Initial number of RAG tools to retrieve
+            step_rag_num: Number of RAG tools to retrieve per step
+            summary_mode: Mode for summarization ('step' or 'length')
+            summary_skip_last_k: Number of last steps to skip in summarization
+            summary_context_length: Context length threshold for summarization
+            force_finish: Whether to force finish when max rounds reached
+            avoid_repeat: Whether to avoid repeating similar responses
+            seed: Random seed for reproducibility
+            enable_checker: Whether to enable reasoning trace checker
+            enable_chat: Whether to enable chat mode
+            additional_default_tools: Additional tools to include by default
+        """
         self.model_name = model_name
         self.tokenizer = None
-        self.terminators = None
         self.rag_model_name = rag_model_name
-        self.tool_files_dict = tool_files_dict
+        self.tool_files_dict = tool_files_dict or {}
         self.model = None
         self.rag_model = ToolRAGModel(rag_model_name)
         self.tooluniverse = None
@@ -61,106 +82,166 @@ class TxAgent:
         self.avoid_repeat = avoid_repeat
         self.seed = seed
         self.enable_checker = enable_checker
-        self.additional_default_tools = additional_default_tools
+        self.additional_default_tools = additional_default_tools or []
         logger.info("TxAgent initialized with model: %s, RAG: %s", model_name, rag_model_name)
 
-    def init_model(self):
+    def init_model(self) -> None:
+        """Initialize both the main model and tool universe."""
         self.load_models()
         self.load_tooluniverse()
+        logger.info("Model and tools initialized successfully")
 
-    def load_models(self, model_name=None):
+    def load_models(self, model_name: Optional[str] = None) -> str:
+        """
+        Load the specified model or the default model if none specified.
+        
+        Args:
+            model_name: Name of the model to load
+            
+        Returns:
+            Status message indicating success or failure
+        """
         if model_name is not None:
             if model_name == self.model_name:
                 return f"The model {model_name} is already loaded."
             self.model_name = model_name
 
-        self.model = LLM(
-            model=self.model_name,
-            dtype="float16",
-            max_model_len=131072,
-            max_num_batched_tokens=65536,  # Increased for A100 80GB
-            max_num_seqs=512,
-            gpu_memory_utilization=0.95,    # 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, 32768, 0.9
-        )
-        return f"Model {model_name} loaded successfully."
+        try:
+            self.model = LLM(
+                model=self.model_name,
+                dtype="float16",
+                max_model_len=131072,
+                max_num_batched_tokens=65536,
+                max_num_seqs=512,
+                gpu_memory_utilization=0.95,
+                trust_remote_code=True,
+            )
+            self.tokenizer = self.model.get_tokenizer()
+            self.chat_template = Template(self.tokenizer.chat_template)
+            logger.info(
+                "Model %s loaded with max_model_len=%d, max_num_batched_tokens=%d",
+                self.model_name, 131072, 65536
+            )
+            return f"Model {model_name} loaded successfully."
+        except Exception as e:
+            logger.error("Failed to load model: %s", str(e))
+            raise RuntimeError(f"Failed to load model: {str(e)}")
 
-    def load_tooluniverse(self):
-        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("ToolUniverse loaded with %d special tools", len(self.special_tools_name))
+    def load_tooluniverse(self) -> None:
+        """Load and initialize the tool universe with specified tools."""
+        try:
+            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.info("ToolUniverse loaded with %d special tools", len(self.special_tools_name))
+        except Exception as e:
+            logger.error("Failed to load tools: %s", str(e))
+            raise RuntimeError(f"Failed to load tools: {str(e)}")
 
-    def load_tool_desc_embedding(self):
+    def load_tool_desc_embedding(self) -> None:
+        """Load tool description embeddings from cache or generate new ones."""
         cache_path = os.path.join(os.path.dirname(self.tool_files_dict["new_tool"]), "tool_embeddings.pkl")
-        if os.path.exists(cache_path):
-            self.rag_model.load_cached_embeddings(cache_path)
-        else:
-            self.rag_model.load_tool_desc_embedding(self.tooluniverse)
-            self.rag_model.save_embeddings(cache_path)
-        logger.debug("Tool description embeddings loaded")
+        try:
+            if os.path.exists(cache_path):
+                self.rag_model.load_cached_embeddings(cache_path)
+            else:
+                self.rag_model.load_tool_desc_embedding(self.tooluniverse)
+                self.rag_model.save_embeddings(cache_path)
+            logger.info("Tool description embeddings loaded successfully")
+        except Exception as e:
+            logger.error("Failed to load tool embeddings: %s", str(e))
+            raise RuntimeError(f"Failed to load tool embeddings: {str(e)}")
 
-    def rag_infer(self, query, top_k=5):
+    def rag_infer(self, query: str, top_k: int = 5) -> List[str]:
+        """
+        Perform RAG inference to retrieve relevant tools.
+        
+        Args:
+            query: The query to search for
+            top_k: Number of top results to return
+            
+        Returns:
+            List of relevant tool names
+        """
+        if not self.enable_rag:
+            return []
         return self.rag_model.rag_infer(query, top_k)
 
-    def initialize_tools_prompt(self, call_agent, call_agent_level, message):
-        picked_tools_prompt = []
-        picked_tools_prompt = self.add_special_tools(
-            picked_tools_prompt, call_agent=call_agent)
-        if call_agent:
-            call_agent_level += 1
-            if call_agent_level >= 2:
-                call_agent = False
-        return picked_tools_prompt, call_agent_level
-
-    def initialize_conversation(self, message, conversation=None, history=None):
+    def initialize_conversation(self, 
+                              message: str, 
+                              conversation: Optional[List[Dict]] = None,
+                              history: Optional[List[Dict]] = None) -> List[Dict]:
+        """
+        Initialize or extend a conversation with the given message and history.
+        
+        Args:
+            message: The new message to add
+            conversation: Existing conversation to extend
+            history: Chat history to incorporate
+            
+        Returns:
+            Updated conversation list
+        """
         if conversation is None:
             conversation = []
 
-        conversation = self.set_system_prompt(
-            conversation, self.prompt_multi_step)
+        conversation = self.set_system_prompt(conversation, self.prompt_multi_step)
         if history:
-            for i in range(len(history)):
-                if history[i]['role'] == 'user':
-                    conversation.append({"role": "user", "content": history[i]['content']})
-                elif history[i]['role'] == 'assistant':
-                    conversation.append({"role": "assistant", "content": history[i]['content']})
+            for msg in history:
+                if msg['role'] == 'user':
+                    conversation.append({"role": "user", "content": msg['content']})
+                elif msg['role'] == 'assistant':
+                    conversation.append({"role": "assistant", "content": msg['content']})
         conversation.append({"role": "user", "content": message})
         logger.debug("Conversation initialized with %d messages", len(conversation))
         return conversation
 
-    def tool_RAG(self, message=None,
-                 picked_tool_names=None,
-                 existing_tools_prompt=[],
-                 rag_num=0,
-                 return_call_result=False):
+    def tool_RAG(self, 
+                 message: Optional[str] = None,
+                 picked_tool_names: Optional[List[str]] = None,
+                 existing_tools_prompt: List = [],
+                 rag_num: int = 0,
+                 return_call_result: bool = False) -> Union[List, Tuple[List, List]]:
+        """
+        Retrieve relevant tools using RAG.
+        
+        Args:
+            message: The query message for RAG
+            picked_tool_names: Pre-selected tool names
+            existing_tools_prompt: Existing tools to include
+            rag_num: Number of tools to retrieve
+            return_call_result: Whether to return tool names
+            
+        Returns:
+            List of tool prompts or tuple with tool names if return_call_result is True
+        """
         if not self.enable_rag:
-            return []
+            return [] if not return_call_result else ([], [])
+
         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]
+            if message is None:
+                raise ValueError("Either message or picked_tool_names must be provided")
+            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):
+    def add_special_tools(self, tools: List, call_agent: bool = False) -> List:
+        """Add special tools (Finish and optionally CallAgent) to the tools list."""
         if self.enable_finish:
             tools.append(self.tooluniverse.get_one_tool_by_one_name('Finish', return_prompt=True))
             logger.debug("Finish tool added")
@@ -169,25 +250,37 @@ 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):
+    def set_system_prompt(self, conversation: List[Dict], sys_prompt: str) -> List[Dict]:
+        """Set or update the system prompt in the conversation."""
         if not conversation:
             conversation.append({"role": "system", "content": sys_prompt})
         else:
             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):
+    def run_function_call(self, 
+                         fcall_str: str,
+                         return_message: bool = False,
+                         existing_tools_prompt: Optional[List] = None,
+                         message_for_call_agent: Optional[str] = None,
+                         call_agent: bool = False,
+                         call_agent_level: Optional[int] = None,
+                         temperature: Optional[float] = None) -> Tuple[List[Dict], List, str]:
+        """
+        Execute function calls from the model's output.
+        
+        Args:
+            fcall_str: The function call string from the model
+            return_message: Whether to return the message part
+            existing_tools_prompt: Existing tools to consider
+            message_for_call_agent: Original message for CallAgent
+            call_agent: Whether CallAgent is enabled
+            call_agent_level: Current CallAgent level
+            temperature: Temperature for sub-agent calls
+            
+        Returns:
+            Tuple of (revised_messages, tools_prompt, special_tool_call)
+        """
         try:
             function_call_json, message = self.tooluniverse.extract_function_call_json(
                 fcall_str, return_message=return_message, verbose=False)
@@ -206,13 +299,12 @@ class TxAgent:
                         special_tool_call = 'Finish'
                         break
                     elif function_call_json[i]["name"] == 'CallAgent':
-                        if call_agent_level < 2 and call_agent:
+                        if call_agent_level is not None and call_agent_level < 2 and call_agent:
                             solution_plan = function_call_json[i]['arguments']['solution']
                             full_message = (
-                                message_for_call_agent +
+                                (message_for_call_agent or "") +
                                 "\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,
@@ -230,7 +322,11 @@ class TxAgent:
                     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})
+                        "content": json.dumps({
+                            "tool_name": function_call_json[i]["name"],
+                            "content": call_result,
+                            "call_id": call_id
+                        })
                     })
         else:
             call_results.append({
@@ -243,109 +339,99 @@ class TxAgent:
             "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
+        return revised_messages, existing_tools_prompt or [], special_tool_call
+
+    def llm_infer(self, 
+                  messages: List[Dict], 
+                  temperature: float = 0.1, 
+                  tools: Optional[List] = None,
+                  output_begin_string: Optional[str] = None, 
+                  max_new_tokens: int = 512,
+                  max_token: int = 131072, 
+                  skip_special_tokens: bool = True,
+                  model: Optional[LLM] = None, 
+                  check_token_status: bool = False) -> Union[str, Tuple[str, bool]]:
+        """
+        Perform inference using the LLM.
+        
+        Args:
+            messages: Conversation history
+            temperature: Sampling temperature
+            tools: List of tools to include
+            output_begin_string: Prefix for output
+            max_new_tokens: Maximum new tokens to generate
+            max_token: Maximum total tokens allowed
+            skip_special_tokens: Whether to skip special tokens
+            model: Optional custom model to use
+            check_token_status: Whether to check token limits
+            
+        Returns:
+            Generated text or tuple with text and overflow flag if check_token_status
+        """
+        model = model or self.model
+        tokenizer = self.tokenizer
 
-        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."})
-            })
+        sampling_params = SamplingParams(
+            temperature=temperature,
+            max_tokens=max_new_tokens,
+            seed=self.seed,
+        )
 
-        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
+        prompt = self.chat_template.render(
+            messages=messages, tools=tools, add_generation_prompt=True)
+        if output_begin_string is not None:
+            prompt += output_begin_string
 
-    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
+        token_overflow = False
+        if check_token_status and max_token is not None:
+            num_input_tokens = len(tokenizer.encode(prompt, add_special_tokens=False))
+            logger.info("Input prompt tokens: %d, max_token: %d", num_input_tokens, max_token)
+            if num_input_tokens > max_token:
+                torch.cuda.empty_cache()
+                gc.collect()
+                logger.warning("Token overflow: %d > %d", num_input_tokens, max_token)
+                return (None, True) if check_token_status else None
 
-    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):
+        try:
+            output = model.generate(prompt, sampling_params=sampling_params)
+            output_text = output[0].outputs[0].text
+            output_tokens = len(tokenizer.encode(output_text, add_special_tokens=False))
+            logger.debug("Inference output: %s (output tokens: %d)", output_text[:100], output_tokens)
+            
+            if skip_special_tokens:
+                output_text = output_text.replace("</s>", "").strip()
+                
+            torch.cuda.empty_cache()
+            gc.collect()
+            
+            return (output_text, token_overflow) if check_token_status else output_text
+        except Exception as e:
+            logger.error("Inference failed: %s", str(e))
+            raise RuntimeError(f"Inference failed: {str(e)}")
+
+    def run_multistep_agent(self, 
+                           message: str,
+                           temperature: float,
+                           max_new_tokens: int,
+                           max_token: int,
+                           max_round: int = 5,
+                           call_agent: bool = False,
+                           call_agent_level: int = 0) -> Optional[str]:
+        """
+        Run multi-step reasoning with the agent.
+        
+        Args:
+            message: Input message
+            temperature: Sampling temperature
+            max_new_tokens: Max new tokens per step
+            max_token: Max total tokens
+            max_round: Maximum reasoning rounds
+            call_agent: Whether to enable CallAgent
+            call_agent_level: Current CallAgent level
+            
+        Returns:
+            Final answer or None if failed
+        """
         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)
@@ -415,196 +501,109 @@ class TxAgent:
                 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.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):
-        if model is None:
-            model = self.model
-
-        logits_processor = self.build_logits_processor(messages, model)
-        sampling_params = SamplingParams(
-            temperature=temperature,
-            max_tokens=max_new_tokens,
-            seed=seed if seed is not None else self.seed,
-        )
-
-        prompt = self.chat_template.render(
-            messages=messages, tools=tools, add_generation_prompt=True)
-        if output_begin_string is not None:
-            prompt += output_begin_string
-
-        if check_token_status and max_token is not None:
-            token_overflow = False
-            num_input_tokens = len(self.tokenizer.encode(prompt, add_special_tokens=False))
-            logger.info("Input prompt tokens: %d, max_token: %d", num_input_tokens, max_token)
-            if num_input_tokens > max_token:
-                torch.cuda.empty_cache()
-                gc.collect()
-                logger.warning("Token overflow: %d > %d", num_input_tokens, max_token)
-                return None, True
+    def analyze_document(self, 
+                        file_path: str,
+                        temperature: float = 0.1,
+                        max_new_tokens: int = 2048,
+                        max_token: int = 131072) -> Dict[str, Union[str, List]]:
+        """
+        Analyze a document and return structured results.
+        
+        Args:
+            file_path: Path to the document
+            temperature: Sampling temperature
+            max_new_tokens: Max new tokens per step
+            max_token: Max total tokens
+            
+        Returns:
+            Dictionary with analysis results
+        """
+        logger.info("Starting document analysis for: %s", file_path)
+        start_time = time.time()
+        
+        try:
+            extracted_text = self.extract_text(file_path)
+            if not extracted_text:
+                raise ValueError("Could not extract text from document")
 
-        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: %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
-
-    def run_self_agent(self, message: str,
-                       temperature: float,
-                       max_new_tokens: int,
-                       max_token: int):
-        logger.info("Starting self agent")
-        conversation = self.set_system_prompt([], self.self_prompt)
-        conversation.append({"role": "user", "content": message})
-        return self.llm_infer(
-            messages=conversation,
-            temperature=temperature,
-            tools=None,
-            max_new_tokens=max_new_tokens,
-            max_token=max_token)
+            chunks = self.split_text(extracted_text)
+            batches = self.batch_chunks(chunks, batch_size=1)
+            batch_results = []
+            
+            for batch in batches:
+                prompt = "\n\n".join(self.build_prompt(chunk) for chunk in batch)
+                response = self.run_multistep_agent(
+                    prompt,
+                    temperature=temperature,
+                    max_new_tokens=max_new_tokens,
+                    max_token=max_token,
+                    call_agent=False
+                )
+                batch_results.append(self.clean_response(response or "No response"))
 
-    def run_chat_agent(self, message: str,
-                       temperature: float,
-                       max_new_tokens: int,
-                       max_token: int):
-        logger.info("Starting chat agent")
-        conversation = self.set_system_prompt([], self.chat_prompt)
-        conversation.append({"role": "user", "content": message})
-        return self.llm_infer(
-            messages=conversation,
-            temperature=temperature,
-            tools=None,
-            max_new_tokens=max_new_tokens,
-            max_token=max_token)
+            combined = "\n\n".join([res for res in batch_results if not res.startswith("❌")])
+            if not combined:
+                raise ValueError("No valid batch responses generated")
 
-    def run_format_agent(self, message: str,
-                         answer: str,
-                         temperature: float,
-                         max_new_tokens: int,
-                         max_token: int):
-        logger.info("Starting format agent")
-        if '[FinalAnswer]' in answer:
-            possible_final_answer = answer.split("[FinalAnswer]")[-1]
-        elif "\n\n" in answer:
-            possible_final_answer = answer.split("\n\n")[-1]
-        else:
-            possible_final_answer = answer.strip()
-        if len(possible_final_answer) == 1 and possible_final_answer in ['A', 'B', 'C', 'D', 'E']:
-            return possible_final_answer
-        elif len(possible_final_answer) > 1 and possible_final_answer[1] == ':' and possible_final_answer[0] in ['A', 'B', 'C', 'D', 'E']:
-            return possible_final_answer[0]
-
-        conversation = self.set_system_prompt(
-            [], "Transform the agent's answer to a single letter: 'A', 'B', 'C', 'D'.")
-        conversation.append({"role": "user", "content": message +
-                            "\nAgent's answer: " + answer + "\nAnswer (must be a letter):"})
-        return self.llm_infer(
+            final_summary = self.generate_final_summary(self, combined)
+            
+            return {
+                "status": "success",
+                "summary": final_summary,
+                "batch_results": batch_results,
+                "processing_time": time.time() - start_time
+            }
+            
+        except Exception as e:
+            logger.error("Document analysis failed: %s", str(e))
+            return {
+                "status": "error",
+                "message": str(e),
+                "processing_time": time.time() - start_time
+            }
+
+    def get_answer_based_on_unfinished_reasoning(self, 
+                                                conversation: List[Dict],
+                                                temperature: float,
+                                                max_new_tokens: int,
+                                                max_token: int) -> str:
+        """
+        Generate a final answer when reasoning is incomplete.
+        
+        Args:
+            conversation: Current conversation history
+            temperature: Sampling temperature
+            max_new_tokens: Max new tokens
+            max_token: Max total tokens
+            
+        Returns:
+            Final answer string
+        """
+        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=None,
+            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_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):
+    def update_parameters(self, **kwargs) -> Dict:
+        """
+        Update agent parameters dynamically.
+        
+        Args:
+            kwargs: Parameter names and values to update
+            
+        Returns:
+            Dictionary of updated parameters
+        """
         updated_attributes = {}
         for key, value in kwargs.items():
             if hasattr(self, key):
@@ -613,199 +612,18 @@ Summarize the function calls' l responses in one sentence with all necessary inf
         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 = 2048,
-                        max_token: int = 131072,
-                        call_agent: bool = False,
-                        conversation: gr.State = None,
-                        max_round: int = 5,
-                        seed: int = None,
-                        call_agent_level: int = 0,
-                        sub_agent_task: str = None,
-                        uploaded_files: list = None):
-        logger.info("Chat started, message: %s", message[:100])
-        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)
-        history = []
-        last_outputs = []
-
-        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.debug("Starting round %d/%d", current_round, max_round)
-                if last_outputs:
-                    function_call_messages, picked_tools_prompt, special_tool_call, current_gradio_history = yield from self.run_function_call_stream(
-                        last_outputs, return_message=True,
-                        existing_tools_prompt=picked_tools_prompt,
-                        message_for_call_agent=message,
-                        call_agent=call_agent,
-                        call_agent_level=call_agent_level,
-                        temperature=temperature)
-                    history.extend(current_gradio_history)
-
-                    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
-                        conversation.append({"role": "assistant", "content": ''.join(last_outputs)})
-                        logger.info("No function call messages, ending chat")
-                        return ''.join(last_outputs).replace("</s>", "")
-
-                last_outputs = []
-                last_outputs_str, token_overflow = self.llm_infer(
-                    messages=conversation,
-                    temperature=temperature,
-                    tools=picked_tools_prompt,
-                    skip_special_tokens=False,
-                    max_new_tokens=max_new_tokens,
-                    max_token=max_token,
-                    seed=seed,
-                    check_token_status=True)
-
-                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)
-                        history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
-                        yield history
-                        return last_outputs_str
-                    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 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
+    def cleanup(self) -> None:
+        """Clean up resources and clear memory."""
+        if hasattr(self, 'model'):
+            del self.model
+        if hasattr(self, 'rag_model'):
+            del self.rag_model
+        if hasattr(self, 'tooluniverse'):
+            del self.tooluniverse
+        torch.cuda.empty_cache()
+        gc.collect()
+        logger.info("TxAgent resources cleaned up")
 
-        except Exception as e:
-            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 = 2048,
-                             max_token: int = 131072,
-                             call_agent: bool = False,
-                             conversation: List = None,
-                             max_round: int = 5,
-                             seed: int = None,
-                             call_agent_level: int = 0):
-        """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,
-                tools=picked_tools_prompt,
-                max_new_tokens=max_new_tokens,
-                max_token=max_token,
-                skip_special_tokens=False,
-                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.")
-            else:
-                batch_results.append(output)
-        
-        logger.info("Batch chat completed for %d messages", len(messages))
-        return batch_results
+    def __del__(self):
+        """Destructor to ensure proper cleanup."""
+        self.cleanup()