Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -73,7 +73,7 @@ class TxAgent:
                 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=1024, gpu_memory_utilization=0.8)
+        self.model = LLM(model=self.model_name, dtype="float16", max_model_len=2048, gpu_memory_utilization=0.8)
         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)
@@ -176,8 +176,14 @@ class TxAgent:
                           call_agent=False,
                           call_agent_level=None,
                           temperature=None):
-        function_call_json, message = self.tooluniverse.extract_function_call_json(
-            fcall_str, return_message=return_message, verbose=False)
+        try:
+            function_call_json, message = self.tooluniverse.extract_function_call_json(
+                fcall_str, return_message=return_message, verbose=False)
+        except Exception as e:
+            logger.error("Tool call parsing failed: %s", e)
+            function_call_json = []
+            message = fcall_str
+
         call_results = []
         special_tool_call = ''
         if function_call_json:
@@ -197,7 +203,7 @@ class TxAgent:
                             )
                             call_result = self.run_multistep_agent(
                                 full_message, temperature=temperature,
-                                max_new_tokens=512, max_token=1024,
+                                max_new_tokens=512, max_token=2048,
                                 call_agent=False, call_agent_level=call_agent_level)
                             if call_result is None:
                                 call_result = "⚠️ No content returned from sub-agent."
@@ -217,7 +223,7 @@ class TxAgent:
         else:
             call_results.append({
                 "role": "tool",
-                "content": json.dumps({"content": "Invalid function call format."})
+                "content": json.dumps({"content": "Invalid or no function call detected."})
             })
 
         revised_messages = [{
@@ -235,8 +241,14 @@ class TxAgent:
                                  call_agent_level=None,
                                  temperature=None,
                                  return_gradio_history=True):
-        function_call_json, message = self.tooluniverse.extract_function_call_json(
-            fcall_str, return_message=return_message, verbose=False)
+        try:
+            function_call_json, message = self.tooluniverse.extract_function_call_json(
+                fcall_str, return_message=return_message, verbose=False)
+        except Exception as e:
+            logger.error("Tool call parsing failed: %s", e)
+            function_call_json = []
+            message = fcall_str
+
         call_results = []
         special_tool_call = ''
         if return_gradio_history:
@@ -264,7 +276,7 @@ class TxAgent:
                             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=1024,
+                                max_new_tokens=512, max_token=2048,
                                 call_agent=False, call_agent_level=call_agent_level,
                                 conversation=None, sub_agent_task=sub_agent_task)
                             if call_result is not None and isinstance(call_result, str):
@@ -287,7 +299,7 @@ class TxAgent:
         else:
             call_results.append({
                 "role": "tool",
-                "content": json.dumps({"content": "Invalid function call format."})
+                "content": json.dumps({"content": "Invalid or no function call detected."})
             })
 
         revised_messages = [{
@@ -300,6 +312,13 @@ class TxAgent:
         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):
+        # Truncate conversation to fit within max_token
+        tokenized = self.tokenizer.encode(json.dumps(conversation), add_special_tokens=False)
+        if len(tokenized) > max_token - 100:
+            logger.warning("Truncating conversation to fit max_token=%d", max_token)
+            while len(tokenized) > max_token - 100 and len(conversation) > 1:
+                conversation.pop(1)  # Keep system prompt and latest message
+                tokenized = self.tokenizer.encode(json.dumps(conversation), add_special_tokens=False)
         if conversation[-1]['role'] == 'assistant':
             conversation.append(
                 {'role': 'tool', 'content': 'Errors occurred during function call; provide final answer with current information.'})
@@ -401,7 +420,7 @@ class TxAgent:
 
     def llm_infer(self, messages, temperature=0.1, tools=None,
                   output_begin_string=None, max_new_tokens=512,
-                  max_token=1024, skip_special_tokens=True,
+                  max_token=2048, skip_special_tokens=True,
                   model=None, tokenizer=None, terminators=None,
                   seed=None, check_token_status=False):
         if model is None:
@@ -550,7 +569,7 @@ Summarize the function calls' responses in one sentence with all necessary infor
                             function_response=function_response,
                             temperature=0.1,
                             max_new_tokens=512,
-                            max_token=1024)
+                            max_token=2048)
                         input_list.insert(last_call_idx + 1, {'role': 'tool', 'content': result_summary})
                         status['summarized_index'] = last_call_idx + 2
                         idx += 1
@@ -572,7 +591,7 @@ Summarize the function calls' responses in one sentence with all necessary infor
                 function_response=function_response,
                 temperature=0.1,
                 max_new_tokens=512,
-                max_token=1024)
+                max_token=2048)
             tool_calls = json.loads(input_list[last_call_idx]['tool_calls'])
             for tool_call in tool_calls:
                 del tool_call['call_id']