wehub-resource-sync
219 lines
13 KiB
Python
219 lines
13 KiB
Python
from datetime import datetime
|
|
import json
|
|
from sqlalchemy import asc
|
|
from sqlalchemy.sql.operators import and_
|
|
import logging
|
|
import superagi
|
|
from superagi.agent.agent_message_builder import AgentLlmMessageBuilder
|
|
from superagi.agent.agent_prompt_builder import AgentPromptBuilder
|
|
from superagi.agent.output_handler import ToolOutputHandler, get_output_handler
|
|
from superagi.agent.task_queue import TaskQueue
|
|
from superagi.agent.tool_builder import ToolBuilder
|
|
from superagi.apm.event_handler import EventHandler
|
|
from superagi.config.config import get_config
|
|
from superagi.helper.error_handler import ErrorHandler
|
|
from superagi.helper.token_counter import TokenCounter
|
|
from superagi.lib.logger import logger
|
|
from superagi.models.agent import Agent
|
|
from superagi.models.agent_config import AgentConfiguration
|
|
from superagi.models.agent_execution import AgentExecution
|
|
from superagi.models.agent_execution_config import AgentExecutionConfiguration
|
|
from superagi.models.agent_execution_feed import AgentExecutionFeed
|
|
from superagi.models.agent_execution_permission import AgentExecutionPermission
|
|
from superagi.models.organisation import Organisation
|
|
from superagi.models.tool import Tool
|
|
from superagi.models.workflows.agent_workflow import AgentWorkflow
|
|
from superagi.models.workflows.agent_workflow_step import AgentWorkflowStep
|
|
from superagi.models.workflows.iteration_workflow import IterationWorkflow
|
|
from superagi.models.workflows.iteration_workflow_step import IterationWorkflowStep
|
|
from superagi.resource_manager.resource_summary import ResourceSummarizer
|
|
from superagi.tools.resource.query_resource import QueryResourceTool
|
|
from superagi.tools.thinking.tools import ThinkingTool
|
|
from superagi.apm.call_log_helper import CallLogHelper
|
|
|
|
|
|
class AgentIterationStepHandler:
|
|
""" Handles iteration workflow steps in the agent workflow."""
|
|
def __init__(self, session, llm, agent_id: int, agent_execution_id: int, memory=None):
|
|
self.session = session
|
|
self.llm = llm
|
|
self.agent_execution_id = agent_execution_id
|
|
self.agent_id = agent_id
|
|
self.memory = memory
|
|
self.organisation = Agent.find_org_by_agent_id(self.session, agent_id=self.agent_id)
|
|
self.task_queue = TaskQueue(str(self.agent_execution_id))
|
|
|
|
def execute_step(self):
|
|
agent_config = Agent.fetch_configuration(self.session, self.agent_id)
|
|
execution = AgentExecution.get_agent_execution_from_id(self.session, self.agent_execution_id)
|
|
iteration_workflow_step = IterationWorkflowStep.find_by_id(self.session, execution.iteration_workflow_step_id)
|
|
agent_execution_config = AgentExecutionConfiguration.fetch_configuration(self.session, self.agent_execution_id)
|
|
if not self._handle_wait_for_permission(execution, agent_config, agent_execution_config,
|
|
iteration_workflow_step):
|
|
return
|
|
|
|
workflow_step = AgentWorkflowStep.find_by_id(self.session, execution.current_agent_step_id)
|
|
organisation = Agent.find_org_by_agent_id(self.session, agent_id=self.agent_id)
|
|
iteration_workflow = IterationWorkflow.find_by_id(self.session, workflow_step.action_reference_id)
|
|
agent_feeds = AgentExecutionFeed.fetch_agent_execution_feeds(self.session, self.agent_execution_id)
|
|
if not agent_feeds:
|
|
self.task_queue.clear_tasks()
|
|
|
|
agent_tools = self._build_tools(agent_config, agent_execution_config)
|
|
prompt = self._build_agent_prompt(iteration_workflow=iteration_workflow,
|
|
agent_config=agent_config,
|
|
agent_execution_config=agent_execution_config,
|
|
prompt=iteration_workflow_step.prompt,
|
|
agent_tools=agent_tools)
|
|
|
|
messages = AgentLlmMessageBuilder(self.session, self.llm, self.llm.get_model(), self.agent_id, self.agent_execution_id) \
|
|
.build_agent_messages(prompt, agent_feeds, history_enabled=iteration_workflow_step.history_enabled,
|
|
completion_prompt=iteration_workflow_step.completion_prompt)
|
|
|
|
logger.debug("Prompt messages:", messages)
|
|
current_tokens = TokenCounter.count_message_tokens(messages = messages, model = self.llm.get_model())
|
|
response = self.llm.chat_completion(messages, TokenCounter(session=self.session, organisation_id=organisation.id).token_limit(self.llm.get_model()) - current_tokens)
|
|
|
|
if 'error' in response and response['message'] is not None:
|
|
ErrorHandler.handle_openai_errors(self.session, self.agent_id, self.agent_execution_id, response['message'])
|
|
|
|
if 'content' not in response or response['content'] is None:
|
|
raise RuntimeError(f"Failed to get response from llm")
|
|
|
|
total_tokens = current_tokens + TokenCounter.count_message_tokens(response['content'], self.llm.get_model())
|
|
AgentExecution.update_tokens(self.session, self.agent_execution_id, total_tokens)
|
|
try:
|
|
content = json.loads(response['content'])
|
|
tool = content.get('tool', {})
|
|
tool_name = tool.get('name', '') if tool else ''
|
|
except json.JSONDecodeError:
|
|
print("Decoding JSON has failed")
|
|
tool_name = ''
|
|
|
|
CallLogHelper(session=self.session, organisation_id=organisation.id).create_call_log(execution.name,
|
|
agent_config['agent_id'], total_tokens, tool_name, agent_config['model'])
|
|
|
|
assistant_reply = response['content']
|
|
output_handler = get_output_handler(iteration_workflow_step.output_type,
|
|
agent_execution_id=self.agent_execution_id,
|
|
agent_config=agent_config,memory=self.memory, agent_tools=agent_tools)
|
|
response = output_handler.handle(self.session, assistant_reply)
|
|
if response.status == "COMPLETE":
|
|
execution.status = "COMPLETED"
|
|
self.session.commit()
|
|
|
|
self._update_agent_execution_next_step(execution, iteration_workflow_step.next_step_id, "COMPLETE")
|
|
EventHandler(session=self.session).create_event('run_completed',
|
|
{'agent_execution_id': execution.id,
|
|
'name': execution.name,
|
|
'tokens_consumed': execution.num_of_tokens,
|
|
"calls": execution.num_of_calls},
|
|
execution.agent_id, organisation.id)
|
|
elif response.status == "WAITING_FOR_PERMISSION":
|
|
execution.status = "WAITING_FOR_PERMISSION"
|
|
execution.permission_id = response.permission_id
|
|
self.session.commit()
|
|
else:
|
|
# moving to next step of iteration or workflow
|
|
self._update_agent_execution_next_step(execution, iteration_workflow_step.next_step_id)
|
|
logger.info(f"Starting next job for agent execution id: {self.agent_execution_id}")
|
|
|
|
self.session.flush()
|
|
|
|
def _update_agent_execution_next_step(self, execution, next_step_id, step_response: str = "default"):
|
|
if next_step_id == -1:
|
|
next_step = AgentWorkflowStep.fetch_next_step(self.session, execution.current_agent_step_id, step_response)
|
|
if str(next_step) == "COMPLETE":
|
|
execution.current_agent_step_id = -1
|
|
execution.status = "COMPLETED"
|
|
else:
|
|
AgentExecution.assign_next_step_id(self.session, self.agent_execution_id, next_step.id)
|
|
else:
|
|
execution.iteration_workflow_step_id = next_step_id
|
|
self.session.commit()
|
|
|
|
def _build_agent_prompt(self, iteration_workflow: IterationWorkflow, agent_config: dict,
|
|
agent_execution_config: dict,
|
|
prompt: str, agent_tools: list):
|
|
max_token_limit = int(get_config("MAX_TOOL_TOKEN_LIMIT", 600))
|
|
prompt = AgentPromptBuilder.replace_main_variables(prompt, agent_execution_config["goal"],
|
|
agent_execution_config["instruction"],
|
|
agent_config["constraints"], agent_tools,
|
|
(not iteration_workflow.has_task_queue))
|
|
if iteration_workflow.has_task_queue:
|
|
response = self.task_queue.get_last_task_details()
|
|
last_task, last_task_result = (response["task"], response["response"]) if response is not None else ("", "")
|
|
current_task = self.task_queue.get_first_task() or ""
|
|
token_limit = TokenCounter(session=self.session, organisation_id=self.organisation.id).token_limit() - max_token_limit
|
|
prompt = AgentPromptBuilder.replace_task_based_variables(prompt, current_task, last_task, last_task_result,
|
|
self.task_queue.get_tasks(),
|
|
self.task_queue.get_completed_tasks(), token_limit)
|
|
return prompt
|
|
|
|
def _build_tools(self, agent_config: dict, agent_execution_config: dict):
|
|
agent_tools = [ThinkingTool()]
|
|
|
|
config_data = AgentConfiguration.get_model_api_key(self.session, self.agent_id, agent_config["model"])
|
|
model_api_key = config_data['api_key']
|
|
tool_builder = ToolBuilder(self.session, self.agent_id, self.agent_execution_id)
|
|
resource_summary = ResourceSummarizer(session=self.session, agent_id=self.agent_id, model=agent_config['model']).fetch_or_create_agent_resource_summary(default_summary=agent_config.get("resource_summary"))
|
|
if resource_summary is not None:
|
|
agent_tools.append(QueryResourceTool())
|
|
user_tools = self.session.query(Tool).filter(
|
|
and_(Tool.id.in_(agent_execution_config["tools"]), Tool.file_name is not None)).all()
|
|
for tool in user_tools:
|
|
agent_tools.append(tool_builder.build_tool(tool))
|
|
agent_tools = [tool_builder.set_default_params_tool(tool, agent_config, agent_execution_config,
|
|
model_api_key, resource_summary,self.memory) for tool in agent_tools]
|
|
return agent_tools
|
|
|
|
def _handle_wait_for_permission(self, agent_execution, agent_config: dict, agent_execution_config: dict,
|
|
iteration_workflow_step: IterationWorkflowStep):
|
|
"""
|
|
Handles the wait for permission when the agent execution is waiting for permission.
|
|
|
|
Args:
|
|
agent_execution (AgentExecution): The agent execution.
|
|
agent_config (dict): The agent configuration.
|
|
agent_execution_config (dict): The agent execution configuration.
|
|
iteration_workflow_step (IterationWorkflowStep): The iteration workflow step.
|
|
|
|
Raises:
|
|
Returns permission success or failure
|
|
"""
|
|
if agent_execution.status != "WAITING_FOR_PERMISSION":
|
|
return True
|
|
agent_execution_permission = self.session.query(AgentExecutionPermission).filter(
|
|
AgentExecutionPermission.id == agent_execution.permission_id).first()
|
|
if agent_execution_permission.status == "PENDING":
|
|
logger.error("handle_wait_for_permission: Permission is still pending")
|
|
return False
|
|
if agent_execution_permission.status == "APPROVED":
|
|
agent_tools = self._build_tools(agent_config, agent_execution_config)
|
|
tool_output_handler = ToolOutputHandler(self.agent_execution_id, agent_config, agent_tools,self.memory)
|
|
tool_result = tool_output_handler.handle_tool_response(self.session,
|
|
agent_execution_permission.assistant_reply)
|
|
result = tool_result.result
|
|
else:
|
|
result = f"User denied the permission to run the tool {agent_execution_permission.tool_name}" \
|
|
f"{' and has given the following feedback : ' + agent_execution_permission.user_feedback if agent_execution_permission.user_feedback else ''}"
|
|
|
|
agent_execution_feed = AgentExecutionFeed(agent_execution_id=agent_execution_permission.agent_execution_id,
|
|
agent_id=agent_execution_permission.agent_id,
|
|
feed=agent_execution_permission.assistant_reply,
|
|
role="assistant",
|
|
feed_group_id=agent_execution.current_feed_group_id)
|
|
self.session.add(agent_execution_feed)
|
|
agent_execution_feed1 = AgentExecutionFeed(agent_execution_id=agent_execution_permission.agent_execution_id,
|
|
agent_id=agent_execution_permission.agent_id,
|
|
feed=result, role="user",
|
|
feed_group_id=agent_execution.current_feed_group_id)
|
|
self.session.add(agent_execution_feed1)
|
|
agent_execution.status = "RUNNING"
|
|
execution = AgentExecution.find_by_id(self.session, agent_execution_permission.agent_execution_id)
|
|
self._update_agent_execution_next_step(execution, iteration_workflow_step.next_step_id)
|
|
self.session.commit()
|
|
|
|
|
|
return True
|