Refactor build_vertex function to include timing information

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-20 16:58:56 -03:00
commit 4abff7f80c

View file

@ -20,7 +20,7 @@ from langflow.api.v1.schemas import (
BuildStatus, BuildStatus,
BuiltResponse, BuiltResponse,
InitResponse, InitResponse,
ResultDict, ResultData,
StreamData, StreamData,
VertexBuildResponse, VertexBuildResponse,
VerticesOrderResponse, VerticesOrderResponse,
@ -322,6 +322,7 @@ async def build_vertex(
inputs: dict = Body(None), inputs: dict = Body(None),
): ):
"""Build a vertex instead of the entire graph.""" """Build a vertex instead of the entire graph."""
start_time = time.perf_counter()
try: try:
cache = chat_service.get_cache(flow_id) cache = chat_service.get_cache(flow_id)
if not cache: if not cache:
@ -336,7 +337,6 @@ async def build_vertex(
graph = cache.get("result") graph = cache.get("result")
result_dict = {} result_dict = {}
duration = "" duration = ""
start_time = time.perf_counter()
if tweaks: if tweaks:
graph = process_tweaks_on_graph(graph, tweaks) graph = process_tweaks_on_graph(graph, tweaks)
if not (vertex := graph.get_vertex(vertex_id)): if not (vertex := graph.get_vertex(vertex_id)):
@ -351,13 +351,9 @@ async def build_vertex(
# to the frontend # to the frontend
vertex.set_artifacts() vertex.set_artifacts()
artifacts = vertex.artifacts artifacts = vertex.artifacts
timedelta = time.perf_counter() - start_time result_dict = ResultData(
duration = format_elapsed_time(timedelta)
result_dict = ResultDict(
results=result_dict, results=result_dict,
artifacts=artifacts, artifacts=artifacts,
duration=duration,
timedelta=timedelta,
) )
vertex.set_result(result_dict) vertex.set_result(result_dict)
elif vertex.result is not None: elif vertex.result is not None:
@ -370,7 +366,7 @@ async def build_vertex(
except Exception as exc: except Exception as exc:
params = str(exc) params = str(exc)
valid = False valid = False
result_dict = ResultDict(results={}) result_dict = ResultData(results={})
artifacts = {} artifacts = {}
# If there's an error building the vertex # If there's an error building the vertex
# we need to clear the cache # we need to clear the cache
@ -383,6 +379,12 @@ async def build_vertex(
data=result_dict, data=result_dict,
artifacts=artifacts, artifacts=artifacts,
) )
timedelta = time.perf_counter() - start_time
duration = format_elapsed_time(timedelta)
result_dict.duration = duration
result_dict.timedelta = timedelta
return VertexBuildResponse( return VertexBuildResponse(
valid=valid, valid=valid,
params=params, params=params,