Refactor get_vertices function to improve flow ID handling and caching

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-16 18:13:03 -03:00
commit 71f7bf9c58

View file

@ -285,15 +285,19 @@ async def get_vertices(
"""Check the flow_id is in the flow_data_store.""" """Check the flow_id is in the flow_data_store."""
try: try:
# First, we need to check if the flow_id is in the cache # First, we need to check if the flow_id is in the cache
graph = None
if cache := chat_service.get_cache(flow_id): if cache := chat_service.get_cache(flow_id):
graph = cache.get("result") graph: Graph = cache.get("result")
flow: Flow = session.get(Flow, flow_id)
if not flow or not flow.data:
raise ValueError("Invalid flow ID")
other_graph = Graph.from_payload(flow.data)
if graph is None:
graph = other_graph
else: else:
flow: Flow = session.get(Flow, flow_id) graph = graph.update(other_graph)
if not flow or not flow.data: chat_service.set_cache(flow_id, graph)
raise ValueError("Invalid flow ID")
graph = Graph.from_payload(flow.data)
chat_service.set_cache(flow_id, graph)
if component_id: if component_id:
vertices = graph.sort_up_to_vertex(component_id) vertices = graph.sort_up_to_vertex(component_id)
@ -350,12 +354,15 @@ async def build_vertex(
duration=duration, duration=duration,
timedelta=timedelta, timedelta=timedelta,
) )
chat_service.set_cache(flow_id, graph)
except Exception as exc: except Exception as exc:
params = str(exc) params = str(exc)
valid = False valid = False
result_dict = ResultDict(results={}) result_dict = ResultDict(results={})
artifacts = {} artifacts = {}
chat_service.set_cache(flow_id, graph) # If there's an error building the vertex
# we need to clear the cache
chat_service.clear_cache(flow_id)
await log_vertex_build( await log_vertex_build(
flow_id=flow_id, flow_id=flow_id,
vertex_id=vertex_id, vertex_id=vertex_id,