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

Code before:
    """
    Update all LLM models in all AI blocks that don't exist in the enum.
    Note: Only updates top level LlmModel SchemaFields of blocks (won't update nested fields).
    """
    logger.info("Migrating LLM models")
    # Scan all blocks and search for LlmModel fields
    llm_model_fields: dict[str, str] = {}  # {block_id: field_name}

    # Search for all LlmModel fields
    for block_type in get_blocks().values():
        block = block_type()
        from pydantic.fields import FieldInfo

        fields: dict[str, FieldInfo] = block.input_schema.model_fields

        # Collect top-level LlmModel fields
        for field_name, field in fields.items():
            if field.annotation == LlmModel:
                llm_model_fields[block.id] = field_name

    # Update each block
    for id, path in llm_model_fields.items():
        # Convert enum values to a list of strings for the SQL query
        enum_values = [v.value for v in LlmModel.__members__.values()]

        await db.execute_raw(
            f"""
            UPDATE "AgentNode"
            SET "constantInput" = jsonb_set("constantInput", '{{{path}}}', '"{migrate_to.value}"', true)
            WHERE "agentBlockId" = '{id}'
            AND "constantInput" ? '{path}'
            AND NOT ("constantInput"->'{path}')::text IN ({','.join(f"'{value}'" for value in enum_values)})
            """
        )


Review comment (suggestion):
I had to do a raw query because Prisma doesn't handle:
1. Checking if a key exists in jsonb column (and throws when checks value for non-existing key)
2. Updating individual value for a key in jsonb

-> 1. would force me to get all AgentNodes (*a lot* of nodes if we remove popular model in future) and 2. to update entire json for each AgentNode individually.
I think this query is safe, none of the variables come from the user.

Code after:
    """
    # Create a new Graph object, inheriting properties from CreatableGraph
    return GraphModel(
        **creatable_graph.model_dump(exclude={"nodes"}),
        user_id=user_id,
        nodes=[
            NodeModel(
                **creatable_node.model_dump(),
                graph_id=creatable_graph.id,
                graph_version=creatable_graph.version,
            )
            for creatable_node in creatable_graph.nodes
        ],
    )


async def fix_llm_provider_credentials():
    """Fix node credentials with provider `llm`"""
    from backend.integrations.credentials_store import IntegrationCredentialsStore

    from .user import get_user_integrations

    store = IntegrationCredentialsStore()

    broken_nodes = []
    try:
        broken_nodes = await prisma.get_client().query_raw(
            """
            SELECT    graph."userId"       user_id,
                  node.id              node_id,
                  node."constantInput" node_preset_input
        FROM      platform."AgentNode"  node
        LEFT JOIN platform."AgentGraph" graph
        ON        node."agentGraphId" = graph.id
        WHERE     node."constantInput"::jsonb->'credentials'->>'provider' = 'llm'
        ORDER BY  graph."userId";
        """
        )
        logger.info(f"Fixing LLM credential inputs on {len(broken_nodes)} nodes")
    except Exception as e:
        logger.error(f"Error fixing LLM credential inputs: {e}")

    user_id: str = ""
    user_integrations = None
    for node in broken_nodes:
        if node["user_id"] != user_id:
            # Save queries by only fetching once per user
            user_id = node["user_id"]
            user_integrations = await get_user_integrations(user_id)
        elif not user_integrations:
            raise RuntimeError(f"Impossible state while processing node {node}")