Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -784,6 +784,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
         Returns:
             str: Final assistant message.
         """
+        logger.debug(f"[TxAgent] Chat started, message: {message[:100]}...")
         print("\033[1;32;40m[TxAgent] Chat started\033[0m")
 
         if not message or len(message.strip()) < 5:
@@ -822,6 +823,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
         try:
             while next_round and current_round < max_round:
                 current_round += 1
+                logger.debug(f"Round {current_round}, conversation length: {len(conversation)}")
 
                 if last_outputs:
                     function_call_messages, picked_tools_prompt, special_tool_call, current_gradio_history = yield from self.run_function_call_stream(
@@ -879,6 +881,22 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                     seed=seed,
                     check_token_status=True)
 
+                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 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:
@@ -922,7 +940,10 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                     yield "The number of reasoning rounds exceeded the limit."
 
         except Exception as e:
-            print(f"[TxAgent] Exception occurred: {e}")
+            logger.error(f"Exception in run_gradio_chat: {e}", exc_info=True)
+            error_msg = f"An error occurred: {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)
@@ -936,5 +957,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                     yield history
                     history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
                     yield history
-            else:
-                yield f"An error occurred: {e}"
+                else:
+                    history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
+                    yield history
+            return error_msg