Repository: Significant-Gravitas/Auto-GPT
File: autogpt_platform/backend/backend/data/graph.py (Python)

Code before:
    for_export: bool = False,
) -> GraphModel | None:
    """
    Retrieves a graph from the DB.
    Defaults to the version with `is_active` if `version` is not passed,
    or the latest version with `is_template` if `template=True`.

    Returns `None` if the record is not found.
    """
    where_clause: AgentGraphWhereInput = {
        "id": graph_id,
    }

    if version is not None:
        where_clause["version"] = version
    else:
        where_clause["isActive"] = True

    graph = await AgentGraph.prisma().find_first(
        where=where_clause,
        include=AGENT_GRAPH_INCLUDE,
        order={"version": "desc"},
    )

    if not graph:
        return None

    if graph.userId == user_id:
        return GraphModel.from_db(graph, for_export) if graph else None

    # If the graph is not owned by the user, we need to check if it's a store listing.
    if not version:
        version = graph.version

    store_listing_where: StoreListingWhereInput = {
        "agentId": graph_id,
        "agentVersion": version,
    }

    store_listing = await StoreListing.prisma().find_first(where=store_listing_where)

    # If it does not belong to the user nor is not a store listing, return None
    if not store_listing:
        return None

    # If it is a store listing, return the graph model
    return GraphModel.from_db(graph, for_export) if graph else None


async def set_graph_active_version(graph_id: str, version: int, user_id: str) -> None:
    # Activate the requested version if it exists and is owned by the user.

Review comment (suggestion):
If we already have this, we can remove `if graph else None` in both returns of this function.

Code after:
    for_export: bool = False,
) -> GraphModel | None:
    """
    Retrieves a graph from the DB.
    Defaults to the version with `is_active` if `version` is not passed,
    or the latest version with `is_template` if `template=True`.

    Returns `None` if the record is not found.
    """
    where_clause: AgentGraphWhereInput = {
        "id": graph_id,
    }

    if version is not None:
        where_clause["version"] = version
    else:
        where_clause["isActive"] = True

    graph = await AgentGraph.prisma().find_first(
        where=where_clause,
        include=AGENT_GRAPH_INCLUDE,
        order={"version": "desc"},
    )

    # The Graph has to be owned by the user or a store listing.
    if (
        graph is None
        or graph.userId != user_id
        and not (
            await StoreListingVersion.prisma().find_first(
                where=prisma.types.StoreListingVersionWhereInput(
                    agentId=graph_id,
                    agentVersion=version or graph.version,
                    isDeleted=False,
                    StoreListing={"is": {"isApproved": True}},
                )
            )
        )
    ):
        return None

    return GraphModel.from_db(graph, for_export)


async def set_graph_active_version(graph_id: str, version: int, user_id: str) -> None:
    # Activate the requested version if it exists and is owned by the user.
    updated_count = await AgentGraph.prisma().update_many(
        data={"isActive": True},
        where={
            "id": graph_id,
            "version": version,