Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -797,17 +797,17 @@ Generate **one summarized sentence** about "function calls' responses" with nece
             str: Final assistant message.
         """
         logger.info("[TxAgent] Chat started with message: %s", message[:100])
-        logger.debug("Initial history: %s", [msg.content[:50] for msg in history] if history else [])
+        logger.debug("Initial history: %s", [msg["content"][:50] for msg in history] if history else [])
 
         # Yield initial message to ensure UI updates
-        history.append(ChatMessage(role="assistant", content="Starting analysis..."))
+        history.append({"role": "assistant", "content": "Starting analysis..."})
         yield history
         logger.debug("Yielded initial history")
 
         try:
             if not message or len(message.strip()) < 5:
                 logger.warning("Invalid message detected")
-                history.append(ChatMessage(role="assistant", content="Please provide a valid message or upload files to analyze."))
+                history.append({"role": "assistant", "content": "Please provide a valid message or upload files to analyze."})
                 yield history
                 return "Invalid input."
 
@@ -851,37 +851,39 @@ Generate **one summarized sentence** about "function calls' responses" with nece
 
                     if not function_call_result:
                         logger.warning("Empty result from run_function_call_stream")
-                        history.append(ChatMessage(role="assistant", content="Error: Tool call processing failed."))
+                        history.append({"role": "assistant", "content": "Error: Tool call processing failed."})
                         yield history
                         return "Error: Tool call processing failed."
 
                     function_call_messages, picked_tools_prompt, special_tool_call, current_gradio_history = function_call_result
 
-                    # Deduplicate history entries
+                    # Convert ChatMessage to dicts and deduplicate
                     unique_history = []
                     seen_contents = set()
                     for msg in current_gradio_history:
-                        if msg.content not in seen_contents:
-                            unique_history.append(msg)
-                            seen_contents.add(msg.content)
+                        content = msg.content
+                        if content not in seen_contents:
+                            unique_history.append({"role": "assistant", "content": content})
+                            seen_contents.add(content)
                     history.extend(unique_history)
                     logger.debug("Extended history with %d unique messages", len(unique_history))
 
                     if special_tool_call == 'Finish' and function_call_messages:
-                        history.append(ChatMessage(role="assistant", content=function_call_messages[0]['content']))
-                        logger.debug("Yielding final history after Finish: %s", function_call_messages[0]['content'][:50])
+                        content = function_call_messages[0]['content']
+                        history.append({"role": "assistant", "content": content})
+                        logger.debug("Yielding final history after Finish: %s", content[:50])
                         yield history
                         next_round = False
                         conversation.extend(function_call_messages)
-                        return function_call_messages[0]['content']
+                        return content
 
                     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.debug("Yielding history for special tool: %s", last_msg.content[:50])
+                        last_msg = history[-1] if history else {"role": "assistant", "content": "Response needed."}
+                        history.append({"role": "assistant", "content": last_msg["content"]})
+                        logger.debug("Yielding history for special tool: %s", last_msg["content"][:50])
                         yield history
                         next_round = False
-                        return last_msg.content
+                        return last_msg["content"]
 
                     if (self.enable_summary or token_overflow) and not call_agent:
                         enable_summary = True
@@ -894,7 +896,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                     else:
                         next_round = False
                         content = ''.join(last_outputs).replace("</s>", "")
-                        history.append(ChatMessage(role="assistant", content=content))
+                        history.append({"role": "assistant", "content": content})
                         conversation.append({"role": "assistant", "content": content})
                         logger.debug("Yielding history with content: %s", content[:50])
                         yield history
@@ -923,26 +925,26 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                 if last_outputs_str is None:
                     logger.warning("llm_infer returned None")
                     error_msg = "Error: Unable to generate response due to token limit. Please reduce input size."
-                    history.append(ChatMessage(role="assistant", content=error_msg))
+                    history.append({"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 isinstance(msg, dict) and "metadata" in msg and 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()))
-                    history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
+                    history.append({"role": "assistant", "content": final_thought.strip()})
+                    history.append({"role": "assistant", "content": "**🧠 Final Analysis:**\n" + final_answer.strip()})
                     logger.debug("Yielding final analysis: %s", final_answer[:50])
                     yield history
                     next_round = False
                 else:
-                    history.append(ChatMessage(role="assistant", content=last_thought))
+                    history.append({"role": "assistant", "content": last_thought})
                     logger.debug("Yielding intermediate history: %s", last_thought[:50])
                     yield history
 
@@ -956,15 +958,15 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                     if '[FinalAnswer]' in last_outputs_str:
                         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()))
-                        history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
+                        history.append({"role": "assistant", "content": final_thought.strip()})
+                        history.append({"role": "assistant", "content": "**🧠 Final Analysis:**\n" + final_answer.strip()})
                     else:
-                        history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
+                        history.append({"role": "assistant", "content": last_outputs_str.strip()})
                     logger.debug("Yielding forced final history")
                     yield history
                 else:
                     error_msg = "The number of reasoning rounds exceeded the limit."
-                    history.append(ChatMessage(role="assistant", content=error_msg))
+                    history.append({"role": "assistant", "content": error_msg})
                     logger.debug("Yielding max rounds error")
                     yield history
                     return error_msg
@@ -972,7 +974,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
         except Exception as e:
             logger.error("Exception in run_gradio_chat: %s", e, exc_info=True)
             error_msg = f"An error occurred: {e}"
-            history.append(ChatMessage(role="assistant", content=error_msg))
+            history.append({"role": "assistant", "content": error_msg})
             logger.debug("Yielding error history: %s", error_msg)
             yield history
             if self.force_finish:
@@ -981,10 +983,10 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                 if '[FinalAnswer]' in last_outputs_str:
                     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()))
-                    history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
+                    history.append({"role": "assistant", "content": final_thought.strip()})
+                    history.append({"role": "assistant", "content": "**🧠 Final Analysis:**\n" + final_answer.strip()})
                 else:
-                    history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
+                    history.append({"role": "assistant", "content": last_outputs_str.strip()})
                 logger.debug("Yielding forced final history after error")
                 yield history
             return error_msg