Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -44,7 +44,7 @@ class TxAgent:
         self.rag_model_name = rag_model_name
         self.tool_files_dict = tool_files_dict
         self.model = None
-        self.rag_model = ToolRAGModel(rag_model_name) if enable_rag else None  # MODIFIED: Skip RAG model if disabled
+        self.rag_model = ToolRAGModel(rag_model_name) if enable_rag else None
         self.tooluniverse = None
         self.prompt_multi_step = "You are a helpful assistant that will solve problems through detailed, step-by-step reasoning."
         self.self_prompt = "Strictly follow the instruction."
@@ -66,7 +66,7 @@ class TxAgent:
 
     def init_model(self):
         self.load_models()
-        if self.enable_rag:  # MODIFIED: Only load tools if RAG enabled
+        if self.enable_rag:
             self.load_tooluniverse()
             self.load_tool_desc_embedding()
 
@@ -79,8 +79,7 @@ class TxAgent:
             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, enforce_eager=True)  # MODIFIED: Force no torch.compile
+        self.model = LLM(model=self.model_name, enforce_eager=True)
         self.chat_template = Template(self.model.get_tokenizer().chat_template)
         self.tokenizer = self.model.get_tokenizer()
         return f"Model {model_name} loaded successfully."
@@ -106,7 +105,7 @@ class TxAgent:
 
     def initialize_tools_prompt(self, call_agent, call_agent_level, message):
         picked_tools_prompt = []
-        if not self.enable_rag:  # MODIFIED: No tools if RAG disabled
+        if not self.enable_rag:
             return picked_tools_prompt, call_agent_level
         picked_tools_prompt = self.add_special_tools(
             picked_tools_prompt, call_agent=call_agent)
@@ -164,18 +163,18 @@ class TxAgent:
         return picked_tools_prompt
 
     def add_special_tools(self, tools, call_agent=False):
-        if not self.enable_rag and not self.enable_finish:  # MODIFIED: No tools if RAG disabled
+        if not self.enable_rag and not self.enable_finish:
             return tools
-        if self.enable_finish:
+        if self.enable_finish and self.tooluniverse:  # MODIFIED: Check tooluniverse
             tools.append(self.tooluniverse.get_one_tool_by_one_name('Finish', return_prompt=True))
             logger.info("Finish tool is added")
-        if call_agent:
+        if call_agent and self.tooluniverse:  # MODIFIED: Check tooluniverse
             tools.append(self.tooluniverse.get_one_tool_by_one_name('CallAgent', return_prompt=True))
             logger.info("CallAgent tool is added")
-        elif self.enable_rag:
+        elif self.enable_rag and self.tooluniverse:  # MODIFIED: Check tooluniverse
             tools.append(self.tooluniverse.get_one_tool_by_one_name('Tool_RAG', return_prompt=True))
             logger.info("Tool_RAG tool is added")
-        if self.additional_default_tools is not None:
+        if self.additional_default_tools is not None and self.tooluniverse:  # MODIFIED: Check tooluniverse
             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:
@@ -184,7 +183,7 @@ class TxAgent:
         return tools
 
     def add_finish_tools(self, tools):
-        if not self.enable_finish:
+        if not self.enable_finish or not self.tooluniverse:  # MODIFIED: Check tooluniverse
             return tools
         tools.append(self.tooluniverse.get_one_tool_by_one_name('Finish', return_prompt=True))
         logger.info("Finish tool is added")
@@ -204,7 +203,7 @@ class TxAgent:
                           call_agent=False,
                           call_agent_level=None,
                           temperature=None):
-        if not self.enable_rag:  # MODIFIED: Skip function calls if RAG disabled
+        if not self.enable_rag:
             return [{"role": "assistant", "content": fcall_str.strip()}], existing_tools_prompt or [], ''
         function_call_json, message = self.tooluniverse.extract_function_call_json(fcall_str, return_message=return_message, verbose=False)
         call_results = []
@@ -270,7 +269,7 @@ class TxAgent:
                                 call_agent_level=None,
                                 temperature=None,
                                 return_gradio_history=True):
-        if not self.enable_rag:  # MODIFIED: Skip function calls if RAG disabled
+        if not self.enable_rag:
             gradio_history = [] if return_gradio_history else None
             return [{"role": "assistant", "content": fcall_str.strip()}], existing_tools_prompt or [], '', gradio_history
         function_call_json, message = self.tooluniverse.extract_function_call_json(fcall_str, return_message=return_message, verbose=False)
@@ -452,7 +451,7 @@ class TxAgent:
                     if len(assistant_messages) == 2:
                         break
             forbidden_ids = [tokenizer.encode(msg, add_special_tokens=False) for msg in assistant_messages]
-            return [NoRepeatSentenceProcessor(forbidden_ids, 3)]  # MODIFIED: Stricter penalty
+            return [NoRepeatSentenceProcessor(forbidden_ids, 3)]
         return None
 
     def llm_infer(self, messages, temperature=0.1, tools=None,
@@ -661,7 +660,7 @@ Generate one summarized sentence about "function calls' responses" with necessar
             call_agent, call_agent_level, message)
         conversation = self.initialize_conversation(
             message, conversation=conversation, history=history)
-        history = []
+        history = [] if not history else history  # MODIFIED: Simplify history
         next_round = True
         function_call_messages = []
         current_round = 0
@@ -674,7 +673,7 @@ Generate one summarized sentence about "function calls' responses" with necessar
             while next_round and current_round < max_round:
                 current_round += 1
                 logger.debug(f"Round {current_round}, conversation length: {len(conversation)}")
-                if last_outputs and self.enable_rag:  # MODIFIED: Skip function calls if RAG disabled
+                if last_outputs and self.enable_rag:
                     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,