Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -264,7 +264,7 @@ class TxAgent:
                             )
                             call_result = self.run_multistep_agent(
                                 full_message, temperature=temperature,
-                                max_new_tokens=1024, max_token=99999,
+                                max_new_tokens=1024, max_token=8192,
                                 call_agent=False, call_agent_level=call_agent_level)
                             if call_result is None:
                                 call_result = "⚠️ No content returned from sub-agent."
@@ -286,13 +286,13 @@ class TxAgent:
         else:
             call_results.append({
                 "role": "tool",
-                "content": json.dumps({"content": "Not a valid function call, please check the function call format."})
+                "content": json.dumps({"content": "No valid tool call detected; proceeding with analysis."})
             })
 
         revised_messages = [{
             "role": "assistant",
-            "content": message.strip(),
-            "tool_calls": json.dumps(function_call_json)
+            "content": message.strip() if message else "Processing...",
+            "tool_calls": json.dumps(function_call_json) if function_call_json else None
         }] + call_results
         logger.debug("Function call completed, returning %d messages", len(revised_messages))
         return revised_messages, existing_tools_prompt, special_tool_call
@@ -317,11 +317,11 @@ class TxAgent:
             logger.warning("No valid function call JSON extracted")
             call_results.append({
                 "role": "tool",
-                "content": json.dumps({"content": "Invalid function call format."})
+                "content": json.dumps({"content": "No tool call detected; continuing analysis."})
             })
             if return_gradio_history:
-                gradio_history.append(ChatMessage(role="assistant", content="Error: Invalid tool call format."))
-                yield call_results, existing_tools_prompt or [], special_tool_call, gradio_history
+                gradio_history.append({"role": "assistant", "content": "No specific tool call identified. Proceeding with medical record analysis."})
+                yield [{"role": "assistant", "content": "Processing..."}], existing_tools_prompt or [], special_tool_call, gradio_history
             return
 
         if isinstance(function_call_json, list):
@@ -352,9 +352,9 @@ class TxAgent:
                             str(solution_plan)
                         )
                         sub_agent_task = "Sub TxAgent plan: " + str(solution_plan)
-                        sub_result = yield from self.run_gradio_chat(
+                        sub_result = yield from self.run EXEMPLARgradio_chat(
                             full_message, history=[], temperature=temperature,
-                            max_new_tokens=1024, max_token=99999,
+                            max_new_tokens=1024, max_token=8192,
                             call_agent=False, call_agent_level=call_agent_level,
                             conversation=None,
                             sub_agent_task=sub_agent_task)
@@ -375,12 +375,12 @@ class TxAgent:
 
                 if return_gradio_history and function_call_json[i]["name"] != 'Finish':
                     metadata = {"title": f"⚒️ {function_call_json[i]['name']}", "log": str(function_call_json[i]['arguments'])}
-                    gradio_history.append(ChatMessage(role="assistant", content=str(call_result), metadata=metadata))
+                    gradio_history.append({"role": "assistant", "content": str(call_result), "metadata": metadata})
 
         revised_messages = [{
             "role": "assistant",
-            "content": message.strip(),
-            "tool_calls": json.dumps(function_call_json)
+            "content": message.strip() if message else "Processing...",
+            "tool_calls": json.dumps(function_call_json) if function_call_json else None
         }] + call_results
 
         if return_gradio_history:
@@ -523,7 +523,7 @@ class TxAgent:
 
     def llm_infer(self, messages, temperature=0.1, tools=None,
                   output_begin_string=None, max_new_tokens=2048,
-                  max_token=None, skip_special_tokens=True,
+                  max_token=8192, skip_special_tokens=True,
                   model=None, tokenizer=None, terminators=None, seed=None, check_token_status=False):
 
         logger.debug("Running LLM inference with %d messages", len(messages))
@@ -565,6 +565,16 @@ class TxAgent:
             sampling_params=sampling_params,
         )
         output = output[0].outputs[0].text
+        # Deduplicate repetitive output
+        if output:
+            lines = output.split('\n')
+            seen = set()
+            deduped_lines = []
+            for line in lines:
+                if line.strip() and line not in seen:
+                    seen.add(line)
+                    deduped_lines.append(line)
+            output = '\n'.join(deduped_lines)
         logger.debug("LLM output: %s", output[:50])
         if check_token_status and max_token is not None:
             return output, token_overflow
@@ -719,7 +729,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                                 function_response=function_response,
                                 temperature=0.1,
                                 max_new_tokens=1024,
-                                max_token=99999
+                                max_token=8192
                             )
 
                             input_list.insert(
@@ -748,7 +758,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                 function_response=function_response,
                 temperature=0.1,
                 max_new_tokens=1024,
-                max_token=99999
+                max_token=8192
             )
 
             tool_calls = json.loads(input_list[last_call_idx]['tool_calls'])
@@ -842,7 +852,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
 
                 if last_outputs:
                     function_call_result = yield from self.run_function_call_stream(
-                        last_outputs, return_message=True,
+                        last_outputs[0], return_message=True,
                         existing_tools_prompt=picked_tools_prompt,
                         message_for_call_agent=message,
                         call_agent=call_agent,
@@ -851,17 +861,18 @@ 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({"role": "assistant", "content": "Error: Tool call processing failed."})
+                        history.append({"role": "assistant", "content": "Error: Unable to process tool response. Continuing analysis."})
                         yield history
-                        return "Error: Tool call processing failed."
+                        last_outputs = []
+                        continue
 
                     function_call_messages, picked_tools_prompt, special_tool_call, current_gradio_history = function_call_result
 
-                    # Convert ChatMessage to dicts and deduplicate
+                    # Convert history to dicts and deduplicate
                     unique_history = []
                     seen_contents = set()
                     for msg in current_gradio_history:
-                        content = msg.content
+                        content = msg["content"] if isinstance(msg, dict) else msg.content
                         if content not in seen_contents:
                             unique_history.append({"role": "assistant", "content": content})
                             seen_contents.add(content)
@@ -915,7 +926,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                     tools=picked_tools_prompt,
                     skip_special_tokens=False,
                     max_new_tokens=max_new_tokens,
-                    max_token=max_token,
+                    max_token=8192,
                     seed=seed,
                     check_token_status=True)