Add successors to VertexBuildResponse and return only first layer

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-22 18:25:13 -03:00
commit 5e30aac399

View file

@ -85,7 +85,7 @@ async def try_running_celery_task(vertex, user_id):
@router.get("/build/{flow_id}/vertices", response_model=VerticesOrderResponse) @router.get("/build/{flow_id}/vertices", response_model=VerticesOrderResponse)
async def get_vertices( async def get_vertices(
flow_id: str, flow_id: str,
component_id: Optional[str] = None, stop_vertex_id: Optional[str] = None,
chat_service: "ChatService" = Depends(get_chat_service), chat_service: "ChatService" = Depends(get_chat_service),
session=Depends(get_session), session=Depends(get_session),
): ):
@ -96,9 +96,9 @@ async def get_vertices(
if cache := chat_service.get_cache(flow_id): if cache := chat_service.get_cache(flow_id):
graph: Graph = cache.get("result") graph: Graph = cache.get("result")
graph = build_and_cache_graph(flow_id, session, chat_service, graph) graph = build_and_cache_graph(flow_id, session, chat_service, graph)
if component_id: if stop_vertex_id:
try: try:
vertices = graph.sort_vertices(component_id) vertices = graph.sort_vertices(stop_vertex_id)
except Exception as exc: except Exception as exc:
logger.error(exc) logger.error(exc)
vertices = graph.sort_vertices() vertices = graph.sort_vertices()
@ -189,6 +189,7 @@ async def build_vertex(
vertex.add_build_time(timedelta) vertex.add_build_time(timedelta)
return VertexBuildResponse( return VertexBuildResponse(
successors=vertex.successors_ids,
valid=valid, valid=valid,
params=params, params=params,
id=vertex.id, id=vertex.id,