refactor: Update build_vertex_stream function to use cache for retrieving graph and session data

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-21 11:18:28 -03:00
commit 984dc5fb67
2 changed files with 11 additions and 11 deletions

View file

@ -295,18 +295,12 @@ async def build_vertex_stream(
async def stream_vertex():
try:
if not session_id:
cache = await chat_service.get_cache(flow_id_str)
if not cache:
# If there's no cache
raise ValueError(f"No cache found for {flow_id_str}.")
else:
graph = cache.get("result")
cache = await chat_service.get_cache(flow_id_str)
if not cache:
# If there's no cache
raise ValueError(f"No cache found for {flow_id_str}.")
else:
session_data = await session_service.load_session(session_id, flow_id=flow_id_str)
graph, artifacts = session_data if session_data else (None, None)
if not graph:
raise ValueError(f"No graph found for {flow_id_str}.")
graph = cache.get("result")
vertex: "InterfaceVertex" = graph.get_vertex(vertex_id)
if not hasattr(vertex, "stream"):

View file

@ -290,6 +290,12 @@ class Graph:
raise ValueError(f"Vertex {vertex_id} not found")
vertex.update_raw_params({"session_id": session_id})
# Process the graph
try:
cache_service = get_chat_service()
await cache_service.set_cache(self.flow_id, self)
except Exception as exc:
logger.exception(exc)
try:
start_component_id = next(
(vertex_id for vertex_id in self._is_input_vertices if "chat" in vertex_id.lower()), None