jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Refactor chat.py to support vertex streaming

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-27 20:40:24 -03:00
commit 34d48992d5
1 changed files with 22 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

28
src/backend/langflow/api/v1/chat.py
View file

@ -1,6 +1,6 @@
import time import time
from typing import Optional
import uuid import uuid
from typing import TYPE_CHECKING, Optional
from fastapi import ( from fastapi import (
APIRouter, APIRouter,
@ -35,6 +35,9 @@ from langflow.services.chat.service import ChatService
from langflow.services.deps import get_chat_service, get_session from langflow.services.deps import get_chat_service, get_session
from langflow.services.monitor.utils import log_vertex_build from langflow.services.monitor.utils import log_vertex_build
if TYPE_CHECKING:
from langflow.graph.vertex.types import ChatVertex
router = APIRouter(tags=["Chat"]) router = APIRouter(tags=["Chat"])
@ -237,27 +240,40 @@ async def build_vertex_stream(
else: else:
graph = cache.get("result") graph = cache.get("result")
vertex = graph.get_vertex(vertex_id) vertex: "ChatVertex" = graph.get_vertex(vertex_id)
if not hasattr(vertex, "stream"):
raise ValueError(f"Vertex {vertex_id} does not support streaming")
if not vertex.pinned or not vertex._built: if not vertex.pinned or not vertex._built:
logger.debug(f"Streaming vertex {vertex_id}")
stream_data = StreamData( stream_data = StreamData(
event="message", event="message",
data={"message": "Building vertex"}, data={"message": f"Streaming vertex {vertex_id}"},
) )
yield str(stream_data) yield str(stream_data)
number_of_chunks = 0
async for chunk in vertex.stream(): async for chunk in vertex.stream():
stream_data = StreamData( stream_data = StreamData(
event="message", event="message",
data={"chunk": chunk}, data={"chunk": chunk},
) )
number_of_chunks += 1
yield str(stream_data) yield str(stream_data)
logger.debug(f"Number of chunks: {number_of_chunks}")
elif vertex.result is not None:
stream_data = StreamData(
event="message",
data={"chunk": vertex._built_result},
)
yield str(stream_data)
else: else:
raise ValueError(f"No result found for vertex {vertex_id}") raise ValueError(f"No result found for vertex {vertex_id}")
except Exception as exc: except Exception as exc:
logger.error(f"Error building vertex: {exc}")
yield str(StreamData(event="error", data={"error": str(exc)})) yield str(StreamData(event="error", data={"error": str(exc)}))
finally:
yield str(StreamData(event="close", data={"message": "Stream closed"})) logger.debug("Closing stream")
yield str(StreamData(event="close", data={"message": "Stream closed"}))
return StreamingResponse(stream_vertex(), media_type="text/event-stream") return StreamingResponse(stream_vertex(), media_type="text/event-stream")
except Exception as exc: except Exception as exc: