feat(chat.py): add status field to flow_data_store to track build status of flows

The import statement for the BuildStatus enum in the schemas module has been updated to ensure the correct import. Additionally, a new "status" field has been added to the flow_data_store dictionary to track the build status of flows. The status is set to BuildStatus.IN_PROGRESS when a flow is being built, BuildStatus.SUCCESS when the build is successful, and BuildStatus.FAILURE when there is an error during the build process. This allows for better tracking and management of flow builds.
🐛 fix(chat.py): change import statement for BuildStatus in schemas module
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-28 22:29:56 -03:00
commit 48d40bfdd3

View file

@ -6,7 +6,7 @@ from fastapi import (
status,
)
from fastapi.responses import StreamingResponse
from langflow.api.v1.schemas import BuiltResponse, InitResponse, StreamData
from langflow.api.v1.schemas import BuildStatus, BuiltResponse, InitResponse, StreamData
from langflow.chat.manager import ChatManager
from langflow.graph.graph.base import Graph
@ -49,7 +49,10 @@ async def init_build(graph_data: dict):
with chat_manager.in_memory_cache._lock:
chat_manager.in_memory_cache.delete(flow_id)
logger.debug(f"Deleted flow {flow_id} from cache")
flow_data_store[flow_id] = {"graph_data": graph_data, "building": False}
flow_data_store[flow_id] = {
"graph_data": graph_data,
"status": BuildStatus.IN_PROGRESS,
}
return InitResponse(flowId=flow_id)
except Exception as exc:
@ -61,8 +64,9 @@ async def init_build(graph_data: dict):
async def build_status(flow_id: str):
"""Check the flow_id is in the flow_data_store."""
try:
built = flow_id in flow_data_store and not isinstance(
flow_data_store[flow_id], dict
built = (
flow_id in flow_data_store
and flow_data_store[flow_id]["status"] == BuildStatus.SUCCESS
)
return BuiltResponse(
@ -86,7 +90,7 @@ async def stream_build(flow_id: str):
yield str(StreamData(event="error", data={"error": error_message}))
return
if flow_data_store[flow_id].get("building"):
if flow_data_store[flow_id].get("status") == BuildStatus.IN_PROGRESS:
error_message = "Already building"
yield str(StreamData(event="error", data={"error": error_message}))
return
@ -124,6 +128,7 @@ async def stream_build(flow_id: str):
except Exception as exc:
params = str(exc)
valid = False
flow_data_store[flow_id]["status"] = BuildStatus.FAILURE
response = {
"valid": valid,
@ -135,8 +140,10 @@ async def stream_build(flow_id: str):
yield str(StreamData(event="message", data=response))
chat_manager.set_cache(flow_id, graph.build())
flow_data_store[flow_id]["status"] = BuildStatus.SUCCESS
except Exception as exc:
logger.error("Error while building the flow: %s", exc)
flow_data_store[flow_id]["status"] = BuildStatus.FAILURE
yield str(StreamData(event="error", data={"error": str(exc)}))
finally:
yield str(StreamData(event="message", data=final_response))