Repository: Significant-Gravitas/Auto-GPT
File: autogpt_platform/backend/backend/server/v2/library/db.py (Python)

Code before:
    Raises:
        AgentNotFoundError: If the specified agent does not exist.
        DatabaseError: If there's an error during creation or if image generation fails.
    """
    logger.info(
        f"Creating library agent for graph #{graph.id} v{graph.version}; "
        f"user #{user_id}"
    )

    try:
        return await prisma.models.LibraryAgent.prisma().create(
            data=prisma.types.LibraryAgentCreateInput(
                isCreatedByUser=(user_id == graph.user_id),
                useGraphIsActiveVersion=True,
                User={"connect": {"id": user_id}},
                # Creator={"connect": {"id": agent.userId}},
                AgentGraph={
                    "connect": {
                        "graphVersionId": {"id": graph.id, "version": graph.version}
                    }
                },
            ),
            include={"AgentGraph": True},
        )
    except prisma.errors.PrismaError as e:
        logger.error(f"Database error creating agent in library: {e}")
        raise store_exceptions.DatabaseError("Failed to create agent in library") from e


async def update_agent_version_in_library(
    user_id: str,
    agent_graph_id: str,
    agent_graph_version: int,
) -> None:
    """
    Updates the agent version in the library if useGraphIsActiveVersion is True.

    Args:
        user_id: Owner of the LibraryAgent.
        agent_graph_id: The agent graph's ID to update.
        agent_graph_version: The new version of the agent graph.

    Raises:
        DatabaseError: If there's an error with the update.
    """
    logger.debug(
        f"Updating agent version in library for user #{user_id}, "
        f"agent #{agent_graph_id} v{agent_graph_version}"
    )
    try:
        library_agent = await prisma.models.LibraryAgent.prisma().find_first_or_raise(

Review comment (suggestion):
I'd advise changing the return type of this function to `library_model.LibraryAgent` instead of the prisma model

Code after:
    Returns:
        The newly created LibraryAgent record.

    Raises:
        AgentNotFoundError: If the specified agent does not exist.
        DatabaseError: If there's an error during creation or if image generation fails.
    """
    logger.info(
        f"Creating library agent for graph #{graph.id} v{graph.version}; "
        f"user #{user_id}"
    )

    try:
        agent = await prisma.models.LibraryAgent.prisma().create(
            data=prisma.types.LibraryAgentCreateInput(
                isCreatedByUser=(user_id == graph.user_id),
                useGraphIsActiveVersion=True,
                User={"connect": {"id": user_id}},
                # Creator={"connect": {"id": agent.userId}},
                AgentGraph={
                    "connect": {
                        "graphVersionId": {"id": graph.id, "version": graph.version}
                    }
                },
            ),
            include={"AgentGraph": True},
        )
        return library_model.LibraryAgent.from_db(agent)
    except prisma.errors.PrismaError as e:
        logger.error(f"Database error creating agent in library: {e}")
        raise store_exceptions.DatabaseError("Failed to create agent in library") from e


async def update_agent_version_in_library(
    user_id: str,
    agent_graph_id: str,
    agent_graph_version: int,
) -> None:
    """
    Updates the agent version in the library if useGraphIsActiveVersion is True.

    Args:
        user_id: Owner of the LibraryAgent.
        agent_graph_id: The agent graph's ID to update.
        agent_graph_version: The new version of the agent graph.

    Raises:
        DatabaseError: If there's an error with the update.
    """
    logger.debug(
        f"Updating agent version in library for user #{user_id}, "