mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-01-31 09:58:19 -05:00
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:
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user