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

Refactor build_vertex_stream function to handle session_id parameter

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-27 22:37:55 -03:00
commit 13317d61e2
1 changed files with 15 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -32,8 +32,9 @@ from langflow.services.auth.utils import (
get_current_user_for_websocket, get_current_user_for_websocket,
) )
from langflow.services.chat.service import ChatService 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, get_session_service
from langflow.services.monitor.utils import log_vertex_build from langflow.services.monitor.utils import log_vertex_build
from langflow.services.session.service import SessionService
if TYPE_CHECKING: if TYPE_CHECKING:
from langflow.graph.vertex.types import ChatVertex from langflow.graph.vertex.types import ChatVertex
@ -227,19 +228,27 @@ async def build_vertex(
async def build_vertex_stream( async def build_vertex_stream(
flow_id: str, flow_id: str,
vertex_id: str, vertex_id: str,
session_id: Optional[str] = None,
chat_service: "ChatService" = Depends(get_chat_service), chat_service: "ChatService" = Depends(get_chat_service),
session_service: "SessionService" = Depends(get_session_service),
): ):
"""Build a vertex instead of the entire graph.""" """Build a vertex instead of the entire graph."""
try: try:
async def stream_vertex(): async def stream_vertex():
try: try:
cache = chat_service.get_cache(flow_id) if not session_id:
if not cache: cache = chat_service.get_cache(flow_id)
# If there's no cache if not cache:
raise ValueError(f"No cache found for {flow_id}.") # If there's no cache
raise ValueError(f"No cache found for {flow_id}.")
else:
graph = cache.get("result")
else: else:
graph = cache.get("result") session_data = await session_service.load_session(session_id)
graph, artifacts = session_data if session_data else (None, None)
if not graph:
raise ValueError(f"No graph found for {flow_id}.")
vertex: "ChatVertex" = graph.get_vertex(vertex_id) vertex: "ChatVertex" = graph.get_vertex(vertex_id)
if not hasattr(vertex, "stream"): if not hasattr(vertex, "stream"):