import gradio as gr import os import sys import json import gc import numpy as np from vllm import LLM, SamplingParams from jinja2 import Template from typing import List import types from tooluniverse import ToolUniverse from gradio import ChatMessage from .toolrag import ToolRAGModel import torch # near the top of txagent.py import logging logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) from .utils import NoRepeatSentenceProcessor, ReasoningTraceChecker, tool_result_format class TxAgent: def __init__(self, model_name, rag_model_name, tool_files_dict=None, # None leads to the default tool files in ToolUniverse enable_finish=True, enable_rag=True, enable_summary=False, init_rag_num=0, step_rag_num=10, summary_mode='step', summary_skip_last_k=0, summary_context_length=None, force_finish=True, avoid_repeat=True, seed=None, enable_checker=False, enable_chat=False, additional_default_tools=None, ): self.model_name = model_name self.tokenizer = None self.terminators = None self.rag_model_name = rag_model_name self.tool_files_dict = tool_files_dict self.model = None self.rag_model = ToolRAGModel(rag_model_name) self.tooluniverse = None # self.tool_desc = None self.prompt_multi_step = "You are a helpful assistant that will solve problems through detailed, step-by-step reasoning and actions based on your reasoning. Typically, your actions will use the provided functions. You have access to the following functions." self.self_prompt = "Strictly follow the instruction." self.chat_prompt = "You are helpful assistant to chat with the user." self.enable_finish = enable_finish self.enable_rag = enable_rag self.enable_summary = enable_summary self.summary_mode = summary_mode self.summary_skip_last_k = summary_skip_last_k self.summary_context_length = summary_context_length self.init_rag_num = init_rag_num self.step_rag_num = step_rag_num self.force_finish = force_finish self.avoid_repeat = avoid_repeat self.seed = seed self.enable_checker = enable_checker self.additional_default_tools = additional_default_tools self.print_self_values() def init_model(self): self.load_models() self.load_tooluniverse() self.load_tool_desc_embedding() def print_self_values(self): for attr, value in self.__dict__.items(): print(f"{attr}: {value}") def load_models(self, model_name=None): if model_name is not None: if model_name == self.model_name: return f"The model {model_name} is already loaded." self.model_name = model_name self.model = LLM(model=self.model_name) self.chat_template = Template(self.model.get_tokenizer().chat_template) self.tokenizer = self.model.get_tokenizer() return f"Model {model_name} loaded successfully." def load_tooluniverse(self): self.tooluniverse = ToolUniverse(tool_files=self.tool_files_dict) self.tooluniverse.load_tools() special_tools = self.tooluniverse.prepare_tool_prompts( self.tooluniverse.tool_category_dicts["special_tools"]) self.special_tools_name = [tool['name'] for tool in special_tools] def load_tool_desc_embedding(self): self.rag_model.load_tool_desc_embedding(self.tooluniverse) def rag_infer(self, query, top_k=5): return self.rag_model.rag_infer(query, top_k) def initialize_tools_prompt(self, call_agent, call_agent_level, message): picked_tools_prompt = [] picked_tools_prompt = self.add_special_tools( picked_tools_prompt, call_agent=call_agent) if call_agent: call_agent_level += 1 if call_agent_level >= 2: call_agent = False if not call_agent: picked_tools_prompt += self.tool_RAG( message=message, rag_num=self.init_rag_num) return picked_tools_prompt, call_agent_level def initialize_conversation(self, message, conversation=None, history=None): if conversation is None: conversation = [] conversation = self.set_system_prompt( conversation, self.prompt_multi_step) if history is not None: if len(history) == 0: conversation = [] print("clear conversation successfully") else: for i in range(len(history)): if history[i]['role'] == 'user': if i-1 >= 0 and history[i-1]['role'] == 'assistant': conversation.append( {"role": "assistant", "content": history[i-1]['content']}) conversation.append( {"role": "user", "content": history[i]['content']}) if i == len(history)-1 and history[i]['role'] == 'assistant': conversation.append( {"role": "assistant", "content": history[i]['content']}) conversation.append({"role": "user", "content": message}) return conversation def tool_RAG(self, message=None, picked_tool_names=None, existing_tools_prompt=[], rag_num=5, return_call_result=False): extra_factor = 30 # Factor to retrieve more than rag_num if picked_tool_names is None: assert picked_tool_names is not None or message is not None picked_tool_names = self.rag_infer( message, top_k=rag_num*extra_factor) picked_tool_names_no_special = [] for tool in picked_tool_names: if tool not in self.special_tools_name: picked_tool_names_no_special.append(tool) picked_tool_names_no_special = picked_tool_names_no_special[:rag_num] picked_tool_names = picked_tool_names_no_special[:rag_num] picked_tools = self.tooluniverse.get_tool_by_name(picked_tool_names) picked_tools_prompt = self.tooluniverse.prepare_tool_prompts( picked_tools) if return_call_result: return picked_tools_prompt, picked_tool_names return picked_tools_prompt def add_special_tools(self, tools, call_agent=False): if self.enable_finish: tools.append(self.tooluniverse.get_one_tool_by_one_name( 'Finish', return_prompt=True)) print("Finish tool is added") if call_agent: tools.append(self.tooluniverse.get_one_tool_by_one_name( 'CallAgent', return_prompt=True)) print("CallAgent tool is added") else: if self.enable_rag: tools.append(self.tooluniverse.get_one_tool_by_one_name( 'Tool_RAG', return_prompt=True)) print("Tool_RAG tool is added") if self.additional_default_tools is not None: for each_tool_name in self.additional_default_tools: tool_prompt = self.tooluniverse.get_one_tool_by_one_name( each_tool_name, return_prompt=True) if tool_prompt is not None: print(f"{each_tool_name} tool is added") tools.append(tool_prompt) return tools def add_finish_tools(self, tools): tools.append(self.tooluniverse.get_one_tool_by_one_name( 'Finish', return_prompt=True)) print("Finish tool is added") return tools def set_system_prompt(self, conversation, sys_prompt): if len(conversation) == 0: conversation.append( {"role": "system", "content": sys_prompt}) else: conversation[0] = {"role": "system", "content": sys_prompt} return conversation def run_function_call(self, fcall_str, return_message=False, existing_tools_prompt=None, message_for_call_agent=None, call_agent=False, call_agent_level=None, temperature=None): function_call_json, message = self.tooluniverse.extract_function_call_json( fcall_str, return_message=return_message, verbose=False) call_results = [] special_tool_call = '' if function_call_json is not None: if isinstance(function_call_json, list): for i in range(len(function_call_json)): print("\033[94mTool Call:\033[0m", 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 += new_tools_prompt 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) ) call_result = self.run_multistep_agent( full_message, temperature=temperature, max_new_tokens=1024, max_token=99999, call_agent=False, call_agent_level=call_agent_level) call_result = call_result.split( '[FinalAnswer]')[-1].strip() 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 print("\033[94mTool Call Result:\033[0m", call_result) call_results.append({ "role": "tool", "content": json.dumps({"content": call_result, "call_id": call_id}) }) else: call_results.append({ "role": "tool", "content": json.dumps({"content": "Not a valid function call, please check the function call format."}) }) revised_messages = [{ "role": "assistant", "content": message.strip(), "tool_calls": json.dumps(function_call_json) }] + call_results # Yield the final result. return revised_messages, existing_tools_prompt, special_tool_call def run_function_call_stream(self, fcall_str, return_message=False, existing_tools_prompt=None, message_for_call_agent=None, call_agent=False, call_agent_level=None, temperature=None, return_gradio_history=True): 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) # When streaming, yield responses as they arrive. 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) call_result = call_result.split( '[FinalAnswer]')[-1] 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({"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: call_results.append({ "role": "tool", "content": json.dumps({"content": "Not a valid function call, please check the function call format."}) }) revised_messages = [{ "role": "assistant", "content": message.strip(), "tool_calls": json.dumps(function_call_json) }] + call_results # Yield the final result. if return_gradio_history: return revised_messages, existing_tools_prompt, special_tool_call, gradio_history else: return revised_messages, existing_tools_prompt, special_tool_call def get_answer_based_on_unfinished_reasoning(self, conversation, temperature, max_new_tokens, max_token, outputs=None): if conversation[-1]['role'] == 'assisant': conversation.append( {'role': 'tool', 'content': 'Errors happen during the function call, please come up with the final answer with the 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]', skip_special_tokens=True, max_new_tokens=max_new_tokens, max_token=max_token) print(last_outputs_str) return last_outputs_str def run_multistep_agent(self, message: str, temperature: float, max_new_tokens: int, max_token: int, max_round: int = 20, call_agent=False, call_agent_level=0) -> str: """ Generate a streaming response using the llama3-8b model. Args: message (str): The input message. temperature (float): The temperature for generating the response. max_new_tokens (int): The maximum number of new tokens to generate. Returns: str: The generated response. """ print("\033[1;32;40mstart\033[0m") picked_tools_prompt, call_agent_level = self.initialize_tools_prompt( call_agent, call_agent_level, message) conversation = self.initialize_conversation(message) outputs = [] last_outputs = [] next_round = True function_call_messages = [] current_round = 0 token_overflow = False enable_summary = False last_status = {} if self.enable_checker: checker = ReasoningTraceChecker(message, conversation) try: while next_round and current_round < max_round: current_round += 1 if len(outputs) > 0: function_call_messages, picked_tools_prompt, special_tool_call = self.run_function_call( last_outputs, return_message=True, existing_tools_prompt=picked_tools_prompt, message_for_call_agent=message, call_agent=call_agent, call_agent_level=call_agent_level, temperature=temperature) if special_tool_call == 'Finish': next_round = False conversation.extend(function_call_messages) if isinstance(function_call_messages[0]['content'], types.GeneratorType): function_call_messages[0]['content'] = next( function_call_messages[0]['content']) return function_call_messages[0]['content'].split('[FinalAnswer]')[-1] if (self.enable_summary or token_overflow) and not call_agent: if token_overflow: print("token_overflow, using summary") enable_summary = True last_status = self.function_result_summary( conversation, status=last_status, enable_summary=enable_summary) if function_call_messages is not None: conversation.extend(function_call_messages) outputs.append(tool_result_format( function_call_messages)) else: next_round = False conversation.extend( [{"role": "assistant", "content": ''.join(last_outputs)}]) return ''.join(last_outputs).replace("", "") if self.enable_checker: good_status, wrong_info = checker.check_conversation() if not good_status: next_round = False print( "Internal error in reasoning: " + wrong_info) break last_outputs = [] outputs.append("### TxAgent:\n") last_outputs_str, token_overflow = self.llm_infer(messages=conversation, temperature=temperature, tools=picked_tools_prompt, skip_special_tokens=False, max_new_tokens=max_new_tokens, max_token=max_token, check_token_status=True) if last_outputs_str is None: next_round = False print( "The number of tokens exceeds the maximum limit.") else: last_outputs.append(last_outputs_str) if max_round == current_round: print("The number of rounds exceeds the maximum limit!") if self.force_finish: return self.get_answer_based_on_unfinished_reasoning(conversation, temperature, max_new_tokens, max_token) else: return None except Exception as e: print(f"Error: {e}") if self.force_finish: return self.get_answer_based_on_unfinished_reasoning(conversation, temperature, max_new_tokens, max_token) else: return None def build_logits_processor(self, messages, llm): # Use the tokenizer from the LLM instance. tokenizer = llm.get_tokenizer() if self.avoid_repeat and len(messages) > 2: assistant_messages = [] for i in range(1, len(messages) + 1): if messages[-i]['role'] == 'assistant': assistant_messages.append(messages[-i]['content']) if len(assistant_messages) == 2: break forbidden_ids = [tokenizer.encode( msg, add_special_tokens=False) for msg in assistant_messages] return [NoRepeatSentenceProcessor(forbidden_ids, 5)] else: return None 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, model=None, tokenizer=None, terminators=None, seed=None, check_token_status=False): if model is None: model = self.model logits_processor = self.build_logits_processor(messages, model) sampling_params = SamplingParams( temperature=temperature, max_tokens=max_new_tokens, seed=seed if seed is not None else self.seed, ) prompt = self.chat_template.render( messages=messages, tools=tools, add_generation_prompt=True) if output_begin_string is not None: prompt += output_begin_string if check_token_status and max_token is not None: token_overflow = False num_input_tokens = len(self.tokenizer.encode( prompt, return_tensors="pt")[0]) if max_token is not None: if num_input_tokens > max_token: torch.cuda.empty_cache() gc.collect() print("Number of input tokens before inference:", num_input_tokens) logger.info( "The number of tokens exceeds the maximum limit!!!!") token_overflow = True return None, token_overflow output = model.generate( prompt, sampling_params=sampling_params, ) output = output[0].outputs[0].text print("\033[92m" + output + "\033[0m") if check_token_status and max_token is not None: return output, token_overflow return output def run_self_agent(self, message: str, temperature: float, max_new_tokens: int, max_token: int) -> str: print("\033[1;32;40mstart self agent\033[0m") conversation = [] conversation = self.set_system_prompt(conversation, self.self_prompt) conversation.append({"role": "user", "content": message}) return self.llm_infer(messages=conversation, temperature=temperature, tools=None, max_new_tokens=max_new_tokens, max_token=max_token) def run_chat_agent(self, message: str, temperature: float, max_new_tokens: int, max_token: int) -> str: print("\033[1;32;40mstart chat agent\033[0m") conversation = [] conversation = self.set_system_prompt(conversation, self.chat_prompt) conversation.append({"role": "user", "content": message}) return self.llm_infer(messages=conversation, temperature=temperature, tools=None, max_new_tokens=max_new_tokens, max_token=max_token) def run_format_agent(self, message: str, answer: str, temperature: float, max_new_tokens: int, max_token: int) -> str: print("\033[1;32;40mstart format agent\033[0m") if '[FinalAnswer]' in answer: possible_final_answer = answer.split("[FinalAnswer]")[-1] elif "\n\n" in answer: possible_final_answer = answer.split("\n\n")[-1] else: possible_final_answer = answer.strip() if len(possible_final_answer) == 1: choice = possible_final_answer[0] if choice in ['A', 'B', 'C', 'D', 'E']: return choice elif len(possible_final_answer) > 1: if possible_final_answer[1] == ':': choice = possible_final_answer[0] if choice in ['A', 'B', 'C', 'D', 'E']: print("choice", choice) return choice conversation = [] format_prompt = f"You are helpful assistant to transform the answer of agent to the final answer of 'A', 'B', 'C', 'D'." conversation = self.set_system_prompt(conversation, format_prompt) conversation.append({"role": "user", "content": message + "\nThe final answer of agent:" + answer + "\n The answer is (must be a letter):"}) return self.llm_infer(messages=conversation, temperature=temperature, tools=None, max_new_tokens=max_new_tokens, max_token=max_token) def run_summary_agent(self, thought_calls: str, function_response: str, temperature: float, max_new_tokens: int, max_token: int) -> str: print("\033[1;32;40mSummarized Tool Result:\033[0m") generate_tool_result_summary_training_prompt = """Thought and function calls: {thought_calls} Function calls' responses: \"\"\" {function_response} \"\"\" Based on the Thought and function calls, and the function calls' responses, you need to generate a summary of the function calls' responses that fulfills the requirements of the thought. The summary MUST BE ONE sentence and include all necessary information. Directly respond with the summarized sentence of the function calls' responses only. Generate **one summarized sentence** about "function calls' responses" with necessary information, and respond with a string: """.format(thought_calls=thought_calls, function_response=function_response) conversation = [] conversation.append( {"role": "user", "content": generate_tool_result_summary_training_prompt}) output = self.llm_infer(messages=conversation, temperature=temperature, tools=None, max_new_tokens=max_new_tokens, max_token=max_token) if '[' in output: output = output.split('[')[0] return output def function_result_summary(self, input_list, status, enable_summary): """ Processes the input list, extracting information from sequences of 'user', 'tool', 'assistant' roles. Supports 'length' and 'step' modes, and skips the last 'k' groups. Parameters: input_list (list): A list of dictionaries containing role and other information. summary_skip_last_k (int): Number of groups to skip from the end. Defaults to 0. summary_context_length (int): The context length threshold for the 'length' mode. last_processed_index (tuple or int): The last processed index. Returns: list: A list of extracted information from valid sequences. """ if 'tool_call_step' not in status: status['tool_call_step'] = 0 for idx in range(len(input_list)): pos_id = len(input_list)-idx-1 if input_list[pos_id]['role'] == 'assistant': if 'tool_calls' in input_list[pos_id]: if 'Tool_RAG' in str(input_list[pos_id]['tool_calls']): status['tool_call_step'] += 1 break if 'step' in status: status['step'] += 1 else: status['step'] = 0 if not enable_summary: return status if 'summarized_index' not in status: status['summarized_index'] = 0 if 'summarized_step' not in status: status['summarized_step'] = 0 if 'previous_length' not in status: status['previous_length'] = 0 if 'history' not in status: status['history'] = [] function_response = '' idx = 0 current_summarized_index = status['summarized_index'] status['history'].append(self.summary_mode == 'step' and status['summarized_step'] < status['step']-status['tool_call_step']-self.summary_skip_last_k) idx = current_summarized_index while idx < len(input_list): if (self.summary_mode == 'step' and status['summarized_step'] < status['step']-status['tool_call_step']-self.summary_skip_last_k) or (self.summary_mode == 'length' and status['previous_length'] > self.summary_context_length): if input_list[idx]['role'] == 'assistant': if 'Tool_RAG' in str(input_list[idx]['tool_calls']): this_thought_calls = None else: if len(function_response) != 0: print("internal summary") status['summarized_step'] += 1 result_summary = self.run_summary_agent( thought_calls=this_thought_calls, function_response=function_response, temperature=0.1, max_new_tokens=1024, max_token=99999 ) input_list.insert( last_call_idx+1, {'role': 'tool', 'content': result_summary}) status['summarized_index'] = last_call_idx + 2 idx += 1 last_call_idx = idx this_thought_calls = input_list[idx]['content'] + \ input_list[idx]['tool_calls'] function_response = '' elif input_list[idx]['role'] == 'tool' and this_thought_calls is not None: function_response += input_list[idx]['content'] del input_list[idx] idx -= 1 else: break idx += 1 if len(function_response) != 0: status['summarized_step'] += 1 result_summary = self.run_summary_agent( thought_calls=this_thought_calls, function_response=function_response, temperature=0.1, max_new_tokens=1024, max_token=99999 ) tool_calls = json.loads(input_list[last_call_idx]['tool_calls']) for tool_call in tool_calls: del tool_call['call_id'] input_list[last_call_idx]['tool_calls'] = json.dumps(tool_calls) input_list.insert( last_call_idx+1, {'role': 'tool', 'content': result_summary}) status['summarized_index'] = last_call_idx + 2 return status # Following are Gradio related functions # General update method that accepts any new arguments through kwargs def update_parameters(self, **kwargs): for key, value in kwargs.items(): if hasattr(self, key): setattr(self, key, value) # Return the updated attributes updated_attributes = {key: value for key, value in kwargs.items() if hasattr(self, key)} return updated_attributes def run_gradio_chat(self, message: str, history: list, temperature: float, max_new_tokens: int, max_token: int, call_agent: bool, conversation: gr.State, max_round: int = 20, seed: int = None, call_agent_level: int = 0, sub_agent_task: str = None) -> str: """ Generate a streaming response using the llama3-8b model. Args: message (str): The input message. history (list): The conversation history used by ChatInterface. temperature (float): The temperature for generating the response. max_new_tokens (int): The maximum number of new tokens to generate. Returns: str: The generated response. """ print("\033[1;32;40mstart\033[0m") print("len(message)", len(message)) if len(message) <= 10: yield "Hi, I am TxAgent, an assistant for answering biomedical questions. Please provide a valid message with a string longer than 10 characters." return "Please provide a valid message." outputs = [] outputs_str = '' last_outputs = [] picked_tools_prompt, call_agent_level = self.initialize_tools_prompt( call_agent, call_agent_level, message) conversation = self.initialize_conversation( message, conversation=conversation, history=history) history = [] next_round = True function_call_messages = [] current_round = 0 enable_summary = False last_status = {} # for summary token_overflow = False if self.enable_checker: checker = ReasoningTraceChecker( message, conversation, init_index=len(conversation)) try: while next_round and current_round < max_round: current_round += 1 if len(last_outputs) > 0: function_call_messages, picked_tools_prompt, special_tool_call, current_gradio_history = yield from self.run_function_call_stream( last_outputs, return_message=True, existing_tools_prompt=picked_tools_prompt, message_for_call_agent=message, call_agent=call_agent, call_agent_level=call_agent_level, temperature=temperature) history.extend(current_gradio_history) if special_tool_call == 'Finish': yield history next_round = False conversation.extend(function_call_messages) return function_call_messages[0]['content'] elif special_tool_call == 'RequireClarification' or special_tool_call == 'DirectResponse': history.append( ChatMessage(role="assistant", content=history[-1].content)) yield history next_round = False return history[-1].content if (self.enable_summary or token_overflow) and not call_agent: if token_overflow: print("token_overflow, using summary") enable_summary = True last_status = self.function_result_summary( conversation, status=last_status, enable_summary=enable_summary) if function_call_messages is not None: conversation.extend(function_call_messages) formated_md_function_call_messages = tool_result_format( function_call_messages) yield history else: next_round = False conversation.extend( [{"role": "assistant", "content": ''.join(last_outputs)}]) return ''.join(last_outputs).replace("", "") if self.enable_checker: good_status, wrong_info = checker.check_conversation() if not good_status: next_round = False print("Internal error in reasoning: " + wrong_info) break last_outputs = [] last_outputs_str, token_overflow = self.llm_infer( messages=conversation, temperature=temperature, tools=picked_tools_prompt, skip_special_tokens=False, max_new_tokens=max_new_tokens, max_token=max_token, seed=seed, check_token_status=True) last_thought = last_outputs_str.split("[TOOL_CALLS]")[0] for each in history: if each.metadata is not None: each.metadata['status'] = 'done' if '[FinalAnswer]' in last_thought: final_thought, final_answer = last_thought.split( '[FinalAnswer]') history.append( ChatMessage(role="assistant", content=final_thought.strip()) ) yield history history.append( ChatMessage( role="assistant", content="**Answer**:\n"+final_answer.strip()) ) yield history else: history.append(ChatMessage( role="assistant", content=last_thought)) yield history last_outputs.append(last_outputs_str) if next_round: if self.force_finish: last_outputs_str = self.get_answer_based_on_unfinished_reasoning( conversation, temperature, max_new_tokens, max_token) for each in history: if each.metadata is not None: each.metadata['status'] = 'done' if '[FinalAnswer]' in last_thought: final_thought, final_answer = last_thought.split( '[FinalAnswer]') history.append( ChatMessage(role="assistant", content=final_thought.strip()) ) yield history history.append( ChatMessage( role="assistant", content="**Answer**:\n"+final_answer.strip()) ) yield history else: yield "The number of rounds exceeds the maximum limit!" except Exception as e: print(f"Error: {e}") if self.force_finish: last_outputs_str = self.get_answer_based_on_unfinished_reasoning( conversation, temperature, max_new_tokens, max_token) for each in history: if each.metadata is not None: each.metadata['status'] = 'done' if '[FinalAnswer]' in last_thought or '"name": "Finish",' in last_outputs_str: if '[FinalAnswer]' in last_thought: final_thought, final_answer = last_thought.split('[FinalAnswer]', 1) else: final_thought = "" final_answer = last_thought history.append( ChatMessage(role="assistant", content=final_thought.strip()) ) yield history history.append( ChatMessage( role="assistant", content="**Answer**:\n" + final_answer.strip()) ) yield history else: return None