Refactor build_vertex function and improve code readability

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-07 16:55:29 -03:00
commit 34d578b813

View file

@ -93,7 +93,7 @@ async def build_vertex(
current_user=Depends(get_current_active_user), current_user=Depends(get_current_active_user),
): ):
"""Build a vertex instead of the entire graph.""" """Build a vertex instead of the entire graph."""
{"inputs": {"input_value": "some value"}}
start_time = time.perf_counter() start_time = time.perf_counter()
next_vertices_ids = [] next_vertices_ids = []
try: try:
@ -101,8 +101,12 @@ async def build_vertex(
cache = chat_service.get_cache(flow_id) cache = chat_service.get_cache(flow_id)
if not cache: if not cache:
# If there's no cache # If there's no cache
logger.warning(f"No cache found for {flow_id}. Building graph starting at {vertex_id}") logger.warning(
graph = build_and_cache_graph(flow_id=flow_id, session=next(get_session()), chat_service=chat_service) f"No cache found for {flow_id}. Building graph starting at {vertex_id}"
)
graph = build_and_cache_graph(
flow_id=flow_id, session=next(get_session()), chat_service=chat_service
)
else: else:
graph = cache.get("result") graph = cache.get("result")
result_data_response = ResultDataResponse(results={}) result_data_response = ResultDataResponse(results={})
@ -122,7 +126,9 @@ async def build_vertex(
else: else:
raise ValueError(f"No result found for vertex {vertex_id}") raise ValueError(f"No result found for vertex {vertex_id}")
next_vertices_ids = vertex.successors_ids next_vertices_ids = vertex.successors_ids
next_vertices_ids = [v for v in next_vertices_ids if graph.should_run_vertex(v)] next_vertices_ids = [
v for v in next_vertices_ids if graph.should_run_vertex(v)
]
result_data_response = ResultDataResponse(**result_dict.model_dump()) result_data_response = ResultDataResponse(**result_dict.model_dump())
@ -205,7 +211,9 @@ async def build_vertex_stream(
else: else:
graph = cache.get("result") graph = cache.get("result")
else: else:
session_data = await session_service.load_session(session_id, flow_id=flow_id) session_data = await session_service.load_session(
session_id, flow_id=flow_id
)
graph, artifacts = session_data if session_data else (None, None) graph, artifacts = session_data if session_data else (None, None)
if not graph: if not graph:
raise ValueError(f"No graph found for {flow_id}.") raise ValueError(f"No graph found for {flow_id}.")