Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -14,8 +14,9 @@ from .toolrag import ToolRAGModel
 import torch
 import logging
 
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
+# Configure logging with a more specific logger name
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger("TxAgent")
 
 from .utils import NoRepeatSentenceProcessor, ReasoningTraceChecker, tool_result_format
 
@@ -404,16 +405,8 @@ class TxAgent:
         return None
 
     def build_logits_processor(self, messages, llm):
-        # Disabled logits processor due to vLLM V1 limitation
         logger.warning("Logits processor disabled due to vLLM V1 limitation")
         return None
-        # Original code (commented out):
-        # tokenizer = llm.get_tokenizer()
-        # if self.avoid_repeat and len(messages) > 2:
-        #     assistant_messages = [msg['content'] for msg in messages[-3:] if msg['role'] == 'assistant'][:2]
-        #     forbidden_ids = [tokenizer.encode(msg, add_special_tokens=False) for msg in assistant_messages]
-        #     return [NoRepeatSentenceProcessor(forbidden_ids, 5)]
-        # return None
 
     def llm_infer(self, messages, temperature=0.1, tools=None,
                   output_begin_string=None, max_new_tokens=512,
@@ -656,7 +649,10 @@ Summarize the function calls' responses in one sentence with all necessary infor
                         yield history
                         next_round = False
                         conversation.extend(function_call_messages)
-                        return function_call_messages[0]['content']
+                        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.")
@@ -717,6 +713,8 @@ Summarize the function calls' responses in one sentence with all necessary infor
                     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
@@ -734,8 +732,12 @@ Summarize the function calls' responses in one sentence with all necessary infor
                     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:
-                    yield "Reasoning rounds exceeded limit."
+                    error_msg = "Reasoning rounds exceeded limit."
+                    history.append(ChatMessage(role="assistant", content=error_msg))
+                    yield history
+                    return error_msg
 
         except Exception as e:
             logger.error("Exception in run_gradio_chat: %s", e, exc_info=True)
@@ -752,4 +754,5 @@ Summarize the function calls' responses in one sentence with all necessary infor
                 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