Repository: aipotheosis-labs/aci
File: aipolabs/cli/commands/create_agent.py (Python)

Code before:
    """
    Create an agent in db.
    """
    return create_agent_helper(
        project_id,
        name,
        description,
        excluded_apps,
        excluded_functions,
        custom_instructions,
        skip_dry_run,
    )


def create_agent_helper(
    project_id: UUID,
    name: str,
    description: str,
    excluded_apps: list[str],
    excluded_functions: list[str],
    custom_instructions: dict[str, str],
    skip_dry_run: bool,
) -> UUID:
    with utils.create_db_session(config.DB_FULL_URL) as db_session:

        agent = crud.projects.create_agent(
            db_session,
            project_id,
            name,
            description,
            excluded_apps,
            excluded_functions,
            custom_instructions,
        )

        if not skip_dry_run:
            click.echo(create_headline(f"will create new agent {agent.name}"))
            click.echo(agent)
            click.echo(create_headline("provide --skip-dry-run to commit changes"))
            db_session.rollback()
        else:
            click.echo(create_headline(f"committing creation of agent {agent.name}"))
            click.echo(agent)
            db_session.commit()

        return agent.id  # type: ignore


Review comment (refactor):
create_agent_helper is used in code directly in different places, it's a bit weird to have `custom_instructions` as string type.
instead it's better to parse `custom_instructions` before calling create_agent_helper, and let `custom_instructions` in `create_agent_helper` be the type it is supposed to be.

e.g., imagine the situation that someone writing some test cases that calls this method and need to pass a mock dict as string.

Code after:
    excluded_apps: list[str],
    excluded_functions: list[str],
    custom_instructions: str | None,
    skip_dry_run: bool,
) -> UUID:
    """
    Create an agent in db.
    """
    return create_agent_helper(
        project_id,
        name,
        description,
        excluded_apps,
        excluded_functions,
        custom_instructions,
        skip_dry_run,
    )


def create_agent_helper(
    project_id: UUID,
    name: str,
    description: str,
    excluded_apps: list[str],
    excluded_functions: list[str],
    custom_instructions: str | None,
    skip_dry_run: bool,
) -> UUID:
    with utils.create_db_session(config.DB_FULL_URL) as db_session:
        if custom_instructions:
            custom_instructions = json.loads(custom_instructions)

        agent = crud.projects.create_agent(
            db_session,
            project_id,
            name,
            description,
            excluded_apps,
            excluded_functions,
            custom_instructions,  # type: ignore
        )

        if not skip_dry_run:
            click.echo(create_headline(f"will create new agent {agent.name}"))
            click.echo(agent)
            click.echo(create_headline("provide --skip-dry-run to commit changes"))
            db_session.rollback()
        else:
            click.echo(create_headline(f"committing creation of agent {agent.name}"))
            click.echo(agent)
            db_session.commit()