Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -46,9 +46,11 @@ class TxAgent:
         self.model = None
         self.rag_model = ToolRAGModel(rag_model_name)
         self.tooluniverse = None
-        self.prompt_multi_step = ("You are a helpful assistant that solves problems through detailed, step-by-step reasoning "
-                                 "and actions based on your reasoning. Provide comprehensive and clinically precise responses, "
-                                 "including specific diagnoses, tools, and actionable recommendations when analyzing medical data.")
+        self.prompt_multi_step = ("You are a highly skilled medical assistant tasked with analyzing medical records in detail. "
+                                 "Provide comprehensive, step-by-step reasoning to identify oversights, including specific diagnoses, "
+                                 "medication conflicts, incomplete assessments, and abnormal results. For each point, include clinical "
+                                 "rationale, standardized screening tools (e.g., PCL-5, SCID-5-PD), and actionable recommendations for "
+                                 "follow-up, ensuring a thorough and precise response.")
         self.self_prompt = "Strictly follow the instruction."
         self.chat_prompt = "You are a helpful assistant to chat with the user."
         self.enable_finish = enable_finish
@@ -281,76 +283,76 @@ class TxAgent:
                                  temperature=None,
                                  return_gradio_history=True):
 
+        logger.debug(f"Running function call stream with input: {fcall_str[:100]}...")
         function_call_json, message = self.tooluniverse.extract_function_call_json(
             fcall_str, return_message=return_message, verbose=False)
         call_results = []
         special_tool_call = ''
-        if return_gradio_history:
-            gradio_history = []
-        if function_call_json is not None:
-            if isinstance(function_call_json, list):
-                for i in range(len(function_call_json)):
-                    if function_call_json[i]["name"] == 'Finish':
-                        special_tool_call = 'Finish'
-                        break
-                    elif function_call_json[i]["name"] == 'Tool_RAG':
-                        new_tools_prompt, call_result = self.tool_RAG(
-                            message=message,
-                            existing_tools_prompt=existing_tools_prompt,
-                            rag_num=self.step_rag_num,
-                            return_call_result=True)
-                        existing_tools_prompt += new_tools_prompt
-                    elif function_call_json[i]["name"] == 'DirectResponse':
-                        call_result = function_call_json[i]['arguments']['respose']
-                        special_tool_call = 'DirectResponse'
-                    elif function_call_json[i]["name"] == 'RequireClarification':
-                        call_result = function_call_json[i]['arguments']['unclear_question']
-                        special_tool_call = 'RequireClarification'
-                    elif function_call_json[i]["name"] == 'CallAgent':
-                        if call_agent_level < 2 and call_agent:
-                            solution_plan = function_call_json[i]['arguments']['solution']
-                            full_message = (
-                                message_for_call_agent +
-                                "\nYou must follow the following plan to answer the question: " +
-                                str(solution_plan)
-                            )
-                            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=1024, max_token=99999,
-                                call_agent=False, call_agent_level=call_agent_level,
-                                conversation=None,
-                                sub_agent_task=sub_agent_task)
+        gradio_history = []
 
-                            if call_result is not None and isinstance(call_result, str):
-                                call_result = call_result.split('[FinalAnswer]')[-1]
-                            else:
-                                call_result = "⚠️ No content returned from sub-agent."
-                        else:
-                            call_result = "Error: The CallAgent has been disabled. Please proceed with your reasoning process to solve this question."
-                    else:
-                        call_result = self.tooluniverse.run_one_function(
-                            function_call_json[i])
-
-                    call_id = self.tooluniverse.call_id_gen()
-                    function_call_json[i]["call_id"] = call_id
-                    call_results.append({
-                        "role": "tool",
-                        "content": json.dumps({"tool_name": function_call_json[i]["name"], "content": call_result, "call_id": call_id})
-                    })
-                    if return_gradio_history and function_call_json[i]["name"] != 'Finish':
-                        if function_call_json[i]["name"] == 'Tool_RAG':
-                            gradio_history.append(ChatMessage(role="assistant", content=str(call_result), metadata={
-                                                  "title": "🧰 "+function_call_json[i]['name'], "log": str(function_call_json[i]['arguments'])}))
-                        else:
-                            gradio_history.append(ChatMessage(role="assistant", content=str(call_result), metadata={
-                                                  "title": "⚒️ "+function_call_json[i]['name'], "log": str(function_call_json[i]['arguments'])}))
-        else:
+        if function_call_json is None:
+            logger.warning("No valid function call JSON extracted")
             call_results.append({
                 "role": "tool",
-                "content": json.dumps({"content": "Not a valid function call, please check the function call format."})
+                "content": json.dumps({"content": "Invalid function call format."})
             })
+            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
+            return
+
+        if isinstance(function_call_json, list):
+            for i in range(len(function_call_json)):
+                logger.debug(f"Processing tool call: {function_call_json[i]}")
+                if function_call_json[i]["name"] == 'Finish':
+                    special_tool_call = 'Finish'
+                    break
+                elif function_call_json[i]["name"] == 'Tool_RAG':
+                    new_tools_prompt, call_result = self.tool_RAG(
+                        message=message,
+                        existing_tools_prompt=existing_tools_prompt,
+                        rag_num=self.step_rag_num,
+                        return_call_result=True)
+                    existing_tools_prompt = (existing_tools_prompt or []) + new_tools_prompt
+                elif function_call_json[i]["name"] == 'DirectResponse':
+                    call_result = function_call_json[i]['arguments']['respose']
+                    special_tool_call = 'DirectResponse'
+                elif function_call_json[i]["name"] == 'RequireClarification':
+                    call_result = function_call_json[i]['arguments']['unclear_question']
+                    special_tool_call = 'RequireClarification'
+                elif function_call_json[i]["name"] == 'CallAgent':
+                    if call_agent_level < 2 and call_agent:
+                        solution_plan = function_call_json[i]['arguments']['solution']
+                        full_message = (
+                            message_for_call_agent +
+                            "\nYou must follow the following plan to answer the question: " +
+                            str(solution_plan)
+                        )
+                        sub_agent_task = "Sub TxAgent plan: " + str(solution_plan)
+                        sub_result = yield from self.run_gradio_chat(
+                            full_message, history=[], temperature=temperature,
+                            max_new_tokens=1024, max_token=99999,
+                            call_agent=False, call_agent_level=call_agent_level,
+                            conversation=None,
+                            sub_agent_task=sub_agent_task)
+                        call_result = sub_result if isinstance(sub_result, str) else "No content from sub-agent."
+                        if '[FinalAnswer]' in call_result:
+                            call_result = call_result.split('[FinalAnswer]')[-1].strip()
+                    else:
+                        call_result = "CallAgent disabled. Proceeding with reasoning."
+                else:
+                    call_result = self.tooluniverse.run_one_function(function_call_json[i])
+
+                call_id = self.tooluniverse.call_id_gen()
+                function_call_json[i]["call_id"] = call_id
+                call_results.append({
+                    "role": "tool",
+                    "content": json.dumps({"tool_name": function_call_json[i]["name"], "content": call_result, "call_id": call_id})
+                })
+
+                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))
 
         revised_messages = [{
             "role": "assistant",
@@ -359,23 +361,24 @@ class TxAgent:
         }] + call_results
 
         if return_gradio_history:
-            return revised_messages, existing_tools_prompt, special_tool_call, gradio_history
+            logger.debug(f"Yielding gradio history with {len(gradio_history)} entries")
+            yield revised_messages, existing_tools_prompt or [], special_tool_call, gradio_history
         else:
-            return revised_messages, existing_tools_prompt, special_tool_call
+            yield revised_messages, existing_tools_prompt or [], special_tool_call
 
     def get_answer_based_on_unfinished_reasoning(self, conversation, temperature, max_new_tokens, max_token, outputs=None):
         if conversation[-1]['role'] == 'assistant':
             conversation.append(
-                {'role': 'tool', 'content': 'Errors happened during the function call, please come up with the final answer with the current information.'})
+                {'role': 'tool', 'content': 'Errors occurred; provide a detailed final answer based on current information.'})
         finish_tools_prompt = self.add_finish_tools([])
 
         last_outputs_str = self.llm_infer(messages=conversation,
                                           temperature=temperature,
                                           tools=finish_tools_prompt,
-                                          output_begin_string='Since I cannot continue reasoning, I will provide the final answer based on the current information and general knowledge.\n\n[FinalAnswer]',
+                                          output_begin_string='[FinalAnswer]',
                                           skip_special_tokens=True,
                                           max_new_tokens=max_new_tokens, max_token=max_token)
-        print(last_outputs_str)
+        logger.debug(f"Forced finish output: {last_outputs_str[:100]}...")
         return last_outputs_str
 
     def run_multistep_agent(self, message: str,
@@ -782,7 +785,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
             message,
             conversation=conversation,
             history=history)
-        history = []
+        history = []  # Reset history to avoid duplication
 
         next_round = True
         function_call_messages = []
@@ -801,7 +804,7 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                 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(
+                    function_call_result = yield from self.run_function_call_stream(
                         last_outputs, return_message=True,
                         existing_tools_prompt=picked_tools_prompt,
                         message_for_call_agent=message,
@@ -809,9 +812,27 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                         call_agent_level=call_agent_level,
                         temperature=temperature)
 
-                    history.extend(current_gradio_history)
+                    # Ensure function_call_result is valid
+                    if not function_call_result:
+                        logger.warning("Empty result from run_function_call_stream")
+                        error_msg = "Error: Tool call processing failed."
+                        history.append(ChatMessage(role="assistant", content=error_msg))
+                        yield history
+                        return error_msg
+
+                    function_call_messages, picked_tools_prompt, special_tool_call, current_gradio_history = function_call_result
+
+                    # Deduplicate history entries
+                    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)
+                    history.extend(unique_history)
 
                     if special_tool_call == 'Finish' and function_call_messages:
+                        history.append(ChatMessage(role="assistant", content=function_call_messages[0]['content']))
                         yield history
                         next_round = False
                         conversation.extend(function_call_messages)
@@ -833,11 +854,13 @@ Generate **one summarized sentence** about "function calls' responses" with nece
 
                     if function_call_messages:
                         conversation.extend(function_call_messages)
-                        yield history
                     else:
                         next_round = False
-                        conversation.append({"role": "assistant", "content": ''.join(last_outputs)})
-                        return ''.join(last_outputs).replace("</s>", "")
+                        content = ''.join(last_outputs).replace("</s>", "")
+                        history.append(ChatMessage(role="assistant", content=content))
+                        conversation.append({"role": "assistant", "content": content})
+                        yield history
+                        return content
 
                 if self.enable_checker:
                     good_status, wrong_info = checker.check_conversation()
@@ -875,7 +898,6 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                     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()))
-                    yield history
                     history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
                     yield history
                     next_round = False
@@ -894,12 +916,10 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                         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()))
-                        yield history
                         history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
-                        yield history
                     else:
                         history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
-                        yield history
+                    yield history
                 else:
                     error_msg = "The number of reasoning rounds exceeded the limit."
                     history.append(ChatMessage(role="assistant", content=error_msg))
@@ -918,10 +938,8 @@ Generate **one summarized sentence** about "function calls' responses" with nece
                     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()))
-                    yield history
                     history.append(ChatMessage(role="assistant", content="**🧠 Final Analysis:**\n" + final_answer.strip()))
-                    yield history
                 else:
                     history.append(ChatMessage(role="assistant", content=last_outputs_str.strip()))
-                    yield history
+                yield history
             return error_msg