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
This commit is contained in:
Otto
2026-01-31 14:13:12 +00:00
parent 9558d3b67a
commit 1b34c871d7
2 changed files with 35 additions and 29 deletions

View File

@@ -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,

View File

@@ -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,