Repository: TransformerOptimus/SuperAGI
File: superagi/jobs/agent_executor.py (Python)

Code before:
            engine.dispose()

    def __execute_workflow_step(self, agent, agent_config, agent_execution_id, agent_workflow_step, memory,
                                model_api_key, organisation, session):
        logger.info("Executing Workflow step : ",agent_workflow_step.action_type)
        if agent_workflow_step.action_type == "TOOL":
            tool_step_handler = AgentToolStepHandler(session,
                                                     llm=get_model(model=agent_config["model"], api_key=model_api_key,
                                                                   organisation_id=organisation.id)
                                                     , agent_id=agent.id, agent_execution_id=agent_execution_id,
                                                     memory=memory)
            tool_step_handler.execute_step()
        elif agent_workflow_step.action_type == "ITERATION_WORKFLOW":
            iteration_step_handler = AgentIterationStepHandler(session,
                                                               llm=get_model(model=agent_config["model"],
                                                                             api_key=model_api_key,
                                                                             organisation_id=organisation.id)
                                                               , agent_id=agent.id,
                                                               agent_execution_id=agent_execution_id, memory=memory)
            print(get_model(model=agent_config["model"], api_key=model_api_key, organisation_id=organisation.id))
            iteration_step_handler.execute_step()
        elif agent_workflow_step.action_type == "WAIT_STEP":
            (AgentWaitStepHandler(session=session, agent_id=agent.id,
                                  agent_execution_id=agent_execution_id)
             .execute_step())

    @classmethod
    def get_embedding(cls, model_source, model_api_key):
        if "OpenAI" in model_source:
            return OpenAiEmbedding(api_key=model_api_key)
        if "Google" in model_source:
            return GooglePalm(api_key=model_api_key)
        if "Hugging" in model_source:
            return HuggingFace(api_key=model_api_key)
        if "Replicate" in model_source:
            return Replicate(api_key=model_api_key)
        return None

    def _check_for_max_iterations(self, session, organisation_id, agent_config, agent_execution_id):
        db_agent_execution = session.query(AgentExecution).filter(AgentExecution.id == agent_execution_id).first()
        if agent_config["max_iterations"] <= db_agent_execution.num_of_calls:
            db_agent_execution.status = "ITERATION_LIMIT_EXCEEDED"

            EventHandler(session=session).create_event('run_iteration_limit_crossed',
                                                       {'agent_execution_id': db_agent_execution.id,
                                                        'name': db_agent_execution.name,
                                                        'tokens_consumed': db_agent_execution.num_of_tokens,
                                                        "calls": db_agent_execution.num_of_calls},
                                                       db_agent_execution.agent_id, organisation_id)
            session.commit()
            logger.info("ITERATION_LIMIT_CROSSED")

Review comment (suggestion):
The action type 'WAIT_STEP' is hardcoded. Consider using a constant or enum type to avoid hardcoding and improve maintainability.

Code after:
                logger.info("Agent Execution is completed or waiting for permission")
                session.close()
                return
            superagi.worker.execute_agent.apply_async((agent_execution_id, datetime.now()), countdown=2)
            # superagi.worker.execute_agent.delay(agent_execution_id, datetime.now())
        finally:
            session.close()
            engine.dispose()

    def __execute_workflow_step(self, agent, agent_config, agent_execution_id, agent_workflow_step, memory,
                                model_api_key, organisation, session):
        logger.info("Executing Workflow step : ", agent_workflow_step.action_type)
        if agent_workflow_step.action_type == AgentWorkflowStepAction.TOOL.value:
            tool_step_handler = AgentToolStepHandler(session,
                                                     llm=get_model(model=agent_config["model"], api_key=model_api_key,
                                                                   organisation_id=organisation.id)
                                                     , agent_id=agent.id, agent_execution_id=agent_execution_id,
                                                     memory=memory)
            tool_step_handler.execute_step()
        elif agent_workflow_step.action_type == AgentWorkflowStepAction.ITERATION_WORKFLOW.value:
            iteration_step_handler = AgentIterationStepHandler(session,
                                                               llm=get_model(model=agent_config["model"],
                                                                             api_key=model_api_key,
                                                                             organisation_id=organisation.id)
                                                               , agent_id=agent.id,
                                                               agent_execution_id=agent_execution_id, memory=memory)
            print(get_model(model=agent_config["model"], api_key=model_api_key, organisation_id=organisation.id))
            iteration_step_handler.execute_step()
        elif agent_workflow_step.action_type == AgentWorkflowStepAction.WAIT_STEP.value:
            (AgentWaitStepHandler(session=session, agent_id=agent.id,
                                  agent_execution_id=agent_execution_id)
             .execute_step())

    @classmethod
    def get_embedding(cls, model_source, model_api_key):
        if "OpenAI" in model_source:
            return OpenAiEmbedding(api_key=model_api_key)
        if "Google" in model_source:
            return GooglePalm(api_key=model_api_key)
        if "Hugging" in model_source:
            return HuggingFace(api_key=model_api_key)
        if "Replicate" in model_source:
            return Replicate(api_key=model_api_key)
        return None

    def _check_for_max_iterations(self, session, organisation_id, agent_config, agent_execution_id):
        db_agent_execution = session.query(AgentExecution).filter(AgentExecution.id == agent_execution_id).first()
        if agent_config["max_iterations"] <= db_agent_execution.num_of_calls:
            db_agent_execution.status = AgentExecutionStatus.ITERATION_LIMIT_EXCEEDED.value

            EventHandler(session=session).create_event('run_iteration_limit_crossed',