| @@ -280,7 +280,7 @@ class BaseAgentRunner(AppRunner): | |||
| def create_agent_thought( | |||
| self, message_id: str, message: str, tool_name: str, tool_input: str, messages_ids: list[str] | |||
| ) -> MessageAgentThought: | |||
| ) -> str: | |||
| """ | |||
| Create agent thought | |||
| """ | |||
| @@ -313,16 +313,15 @@ class BaseAgentRunner(AppRunner): | |||
| db.session.add(thought) | |||
| db.session.commit() | |||
| db.session.refresh(thought) | |||
| db.session.close() | |||
| agent_thought_id = str(thought.id) | |||
| self.agent_thought_count += 1 | |||
| db.session.close() | |||
| return thought | |||
| return agent_thought_id | |||
| def save_agent_thought( | |||
| self, | |||
| agent_thought: MessageAgentThought, | |||
| agent_thought_id: str, | |||
| tool_name: str | None, | |||
| tool_input: Union[str, dict, None], | |||
| thought: str | None, | |||
| @@ -335,12 +334,9 @@ class BaseAgentRunner(AppRunner): | |||
| """ | |||
| Save agent thought | |||
| """ | |||
| updated_agent_thought = ( | |||
| db.session.query(MessageAgentThought).where(MessageAgentThought.id == agent_thought.id).first() | |||
| ) | |||
| if not updated_agent_thought: | |||
| agent_thought = db.session.query(MessageAgentThought).where(MessageAgentThought.id == agent_thought_id).first() | |||
| if not agent_thought: | |||
| raise ValueError("agent thought not found") | |||
| agent_thought = updated_agent_thought | |||
| if thought: | |||
| agent_thought.thought += thought | |||
| @@ -355,7 +351,7 @@ class BaseAgentRunner(AppRunner): | |||
| except Exception: | |||
| tool_input = json.dumps(tool_input) | |||
| updated_agent_thought.tool_input = tool_input | |||
| agent_thought.tool_input = tool_input | |||
| if observation: | |||
| if isinstance(observation, dict): | |||
| @@ -364,27 +360,27 @@ class BaseAgentRunner(AppRunner): | |||
| except Exception: | |||
| observation = json.dumps(observation) | |||
| updated_agent_thought.observation = observation | |||
| agent_thought.observation = observation | |||
| if answer: | |||
| agent_thought.answer = answer | |||
| if messages_ids is not None and len(messages_ids) > 0: | |||
| updated_agent_thought.message_files = json.dumps(messages_ids) | |||
| agent_thought.message_files = json.dumps(messages_ids) | |||
| if llm_usage: | |||
| updated_agent_thought.message_token = llm_usage.prompt_tokens | |||
| updated_agent_thought.message_price_unit = llm_usage.prompt_price_unit | |||
| updated_agent_thought.message_unit_price = llm_usage.prompt_unit_price | |||
| updated_agent_thought.answer_token = llm_usage.completion_tokens | |||
| updated_agent_thought.answer_price_unit = llm_usage.completion_price_unit | |||
| updated_agent_thought.answer_unit_price = llm_usage.completion_unit_price | |||
| updated_agent_thought.tokens = llm_usage.total_tokens | |||
| updated_agent_thought.total_price = llm_usage.total_price | |||
| agent_thought.message_token = llm_usage.prompt_tokens | |||
| agent_thought.message_price_unit = llm_usage.prompt_price_unit | |||
| agent_thought.message_unit_price = llm_usage.prompt_unit_price | |||
| agent_thought.answer_token = llm_usage.completion_tokens | |||
| agent_thought.answer_price_unit = llm_usage.completion_price_unit | |||
| agent_thought.answer_unit_price = llm_usage.completion_unit_price | |||
| agent_thought.tokens = llm_usage.total_tokens | |||
| agent_thought.total_price = llm_usage.total_price | |||
| # check if tool labels is not empty | |||
| labels = updated_agent_thought.tool_labels or {} | |||
| tools = updated_agent_thought.tool.split(";") if updated_agent_thought.tool else [] | |||
| labels = agent_thought.tool_labels or {} | |||
| tools = agent_thought.tool.split(";") if agent_thought.tool else [] | |||
| for tool in tools: | |||
| if not tool: | |||
| continue | |||
| @@ -395,7 +391,7 @@ class BaseAgentRunner(AppRunner): | |||
| else: | |||
| labels[tool] = {"en_US": tool, "zh_Hans": tool} | |||
| updated_agent_thought.tool_labels_str = json.dumps(labels) | |||
| agent_thought.tool_labels_str = json.dumps(labels) | |||
| if tool_invoke_meta is not None: | |||
| if isinstance(tool_invoke_meta, dict): | |||
| @@ -404,7 +400,7 @@ class BaseAgentRunner(AppRunner): | |||
| except Exception: | |||
| tool_invoke_meta = json.dumps(tool_invoke_meta) | |||
| updated_agent_thought.tool_meta_str = tool_invoke_meta | |||
| agent_thought.tool_meta_str = tool_invoke_meta | |||
| db.session.commit() | |||
| db.session.close() | |||
| @@ -97,13 +97,13 @@ class CotAgentRunner(BaseAgentRunner, ABC): | |||
| message_file_ids: list[str] = [] | |||
| agent_thought = self.create_agent_thought( | |||
| agent_thought_id = self.create_agent_thought( | |||
| message_id=message.id, message="", tool_name="", tool_input="", messages_ids=message_file_ids | |||
| ) | |||
| if iteration_step > 1: | |||
| self.queue_manager.publish( | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought.id), PublishFrom.APPLICATION_MANAGER | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER | |||
| ) | |||
| # recalc llm max tokens | |||
| @@ -133,7 +133,7 @@ class CotAgentRunner(BaseAgentRunner, ABC): | |||
| # publish agent thought if it's first iteration | |||
| if iteration_step == 1: | |||
| self.queue_manager.publish( | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought.id), PublishFrom.APPLICATION_MANAGER | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER | |||
| ) | |||
| for chunk in react_chunks: | |||
| @@ -168,7 +168,7 @@ class CotAgentRunner(BaseAgentRunner, ABC): | |||
| usage_dict["usage"] = LLMUsage.empty_usage() | |||
| self.save_agent_thought( | |||
| agent_thought=agent_thought, | |||
| agent_thought_id=agent_thought_id, | |||
| tool_name=(scratchpad.action.action_name if scratchpad.action and not scratchpad.is_final() else ""), | |||
| tool_input={scratchpad.action.action_name: scratchpad.action.action_input} if scratchpad.action else {}, | |||
| tool_invoke_meta={}, | |||
| @@ -181,7 +181,7 @@ class CotAgentRunner(BaseAgentRunner, ABC): | |||
| if not scratchpad.is_final(): | |||
| self.queue_manager.publish( | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought.id), PublishFrom.APPLICATION_MANAGER | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER | |||
| ) | |||
| if not scratchpad.action: | |||
| @@ -212,7 +212,7 @@ class CotAgentRunner(BaseAgentRunner, ABC): | |||
| scratchpad.agent_response = tool_invoke_response | |||
| self.save_agent_thought( | |||
| agent_thought=agent_thought, | |||
| agent_thought_id=agent_thought_id, | |||
| tool_name=scratchpad.action.action_name, | |||
| tool_input={scratchpad.action.action_name: scratchpad.action.action_input}, | |||
| thought=scratchpad.thought or "", | |||
| @@ -224,7 +224,7 @@ class CotAgentRunner(BaseAgentRunner, ABC): | |||
| ) | |||
| self.queue_manager.publish( | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought.id), PublishFrom.APPLICATION_MANAGER | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER | |||
| ) | |||
| # update prompt tool message | |||
| @@ -244,7 +244,7 @@ class CotAgentRunner(BaseAgentRunner, ABC): | |||
| # save agent thought | |||
| self.save_agent_thought( | |||
| agent_thought=agent_thought, | |||
| agent_thought_id=agent_thought_id, | |||
| tool_name="", | |||
| tool_input={}, | |||
| tool_invoke_meta={}, | |||
| @@ -80,7 +80,7 @@ class FunctionCallAgentRunner(BaseAgentRunner): | |||
| prompt_messages_tools = [] | |||
| message_file_ids: list[str] = [] | |||
| agent_thought = self.create_agent_thought( | |||
| agent_thought_id = self.create_agent_thought( | |||
| message_id=message.id, message="", tool_name="", tool_input="", messages_ids=message_file_ids | |||
| ) | |||
| @@ -114,7 +114,7 @@ class FunctionCallAgentRunner(BaseAgentRunner): | |||
| for chunk in chunks: | |||
| if is_first_chunk: | |||
| self.queue_manager.publish( | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought.id), PublishFrom.APPLICATION_MANAGER | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER | |||
| ) | |||
| is_first_chunk = False | |||
| # check if there is any tool call | |||
| @@ -172,7 +172,7 @@ class FunctionCallAgentRunner(BaseAgentRunner): | |||
| result.message.content = "" | |||
| self.queue_manager.publish( | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought.id), PublishFrom.APPLICATION_MANAGER | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER | |||
| ) | |||
| yield LLMResultChunk( | |||
| @@ -205,7 +205,7 @@ class FunctionCallAgentRunner(BaseAgentRunner): | |||
| # save thought | |||
| self.save_agent_thought( | |||
| agent_thought=agent_thought, | |||
| agent_thought_id=agent_thought_id, | |||
| tool_name=tool_call_names, | |||
| tool_input=tool_call_inputs, | |||
| thought=response, | |||
| @@ -216,7 +216,7 @@ class FunctionCallAgentRunner(BaseAgentRunner): | |||
| llm_usage=current_llm_usage, | |||
| ) | |||
| self.queue_manager.publish( | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought.id), PublishFrom.APPLICATION_MANAGER | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER | |||
| ) | |||
| final_answer += response + "\n" | |||
| @@ -276,7 +276,7 @@ class FunctionCallAgentRunner(BaseAgentRunner): | |||
| if len(tool_responses) > 0: | |||
| # save agent thought | |||
| self.save_agent_thought( | |||
| agent_thought=agent_thought, | |||
| agent_thought_id=agent_thought_id, | |||
| tool_name="", | |||
| tool_input="", | |||
| thought="", | |||
| @@ -291,7 +291,7 @@ class FunctionCallAgentRunner(BaseAgentRunner): | |||
| messages_ids=message_file_ids, | |||
| ) | |||
| self.queue_manager.publish( | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought.id), PublishFrom.APPLICATION_MANAGER | |||
| QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER | |||
| ) | |||
| # update prompt tool | |||