From 1b34c871d79a5140fbd672c43dae86ee0148a1a4 Mon Sep 17 00:00:00 2001
From: Otto <otto@agpt.co>
Date: Sat, 31 Jan 2026 14:13:12 +0000
Subject: [PATCH] refactor: move batch graph query to data layer

- Added get_graphs_by_ids to backend/data/graph.py
- Uses StoreListingVersion lookup (same pattern as get_graph for public graphs)
- Removed inline Prisma query from feature code
---
 .../chat/tools/agent_generator/core.py        | 30 +---------------
 .../backend/backend/data/graph.py             | 34 +++++++++++++++++++
 2 files changed, 35 insertions(+), 29 deletions(-)

diff --git a/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/core.py b/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/core.py
index b359c48e88..9f04848889 100644
--- a/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/core.py
+++ b/autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/core.py
@@ -5,20 +5,17 @@ import re
 import uuid
 from typing import Any, NotRequired, TypedDict
 
-import prisma.models
-
 from backend.api.features.library import db as library_db
 from backend.api.features.store import db as store_db
 from backend.data.graph import (
     Graph,
-    GraphModel,
     Link,
     Node,
     create_graph,
     get_graph,
     get_graph_all_versions,
+    get_graphs_by_ids,
 )
-from backend.data.includes import AGENT_GRAPH_INCLUDE
 from backend.util.exceptions import DatabaseError, NotFoundError
 
 from .service import (
@@ -267,31 +264,6 @@ async def get_library_agents_for_generation(
         return []
 
 
-async def get_graphs_by_ids(
-    *graph_ids: str,
-) -> dict[str, GraphModel]:
-    """Batch-fetch multiple graphs by their IDs using a single query.
-
-    Args:
-        *graph_ids: Variable number of graph IDs to fetch
-
-    Returns:
-        Dict mapping graph_id to GraphModel for successfully fetched graphs
-    """
-    if not graph_ids:
-        return {}
-
-    try:
-        graphs = await prisma.models.AgentGraph.prisma().find_many(
-            where={"id": {"in": list(graph_ids)}, "isActive": True},
-            include=AGENT_GRAPH_INCLUDE,
-        )
-        return {graph.id: GraphModel.from_db(graph) for graph in graphs}
-    except Exception as e:
-        logger.warning(f"Failed to batch fetch graphs: {e}")
-        return {}
-
-
 async def search_marketplace_agents_for_generation(
     search_query: str,
     max_results: int = 10,
diff --git a/autogpt_platform/backend/backend/data/graph.py b/autogpt_platform/backend/backend/data/graph.py
index c1f38f81d5..f83c6ab2e2 100644
--- a/autogpt_platform/backend/backend/data/graph.py
+++ b/autogpt_platform/backend/backend/data/graph.py
@@ -1028,6 +1028,40 @@ async def get_graph(
     return GraphModel.from_db(graph, for_export)
 
 
+async def get_graphs_by_ids(*graph_ids: str) -> dict[str, GraphModel]:
+    """Batch-fetch multiple public/marketplace graphs by their IDs.
+
+    This fetches graphs that are publicly available via approved store listings.
+    Used for marketplace agent schema lookups where we need multiple graphs
+    in a single query.
+
+    Args:
+        *graph_ids: Variable number of graph IDs to fetch
+
+    Returns:
+        Dict mapping graph_id to GraphModel for successfully fetched graphs
+    """
+    if not graph_ids:
+        return {}
+
+    store_listings = await StoreListingVersion.prisma().find_many(
+        where={
+            "agentGraphId": {"in": list(graph_ids)},
+            "submissionStatus": SubmissionStatus.APPROVED,
+            "isDeleted": False,
+        },
+        include={"AgentGraph": {"include": AGENT_GRAPH_INCLUDE}},
+        distinct=["agentGraphId"],
+        order={"agentGraphVersion": "desc"},
+    )
+
+    return {
+        listing.agentGraphId: GraphModel.from_db(listing.AgentGraph)
+        for listing in store_listings
+        if listing.AgentGraph
+    }
+
+
 async def get_graph_as_admin(
     graph_id: str,
     version: int | None = None,