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