From 853ce351c9020075fdb99e428b96b28f7cc60802 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 25 Aug 2023 15:24:45 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(chat.py):=20add=20missing=20?= =?UTF-8?q?import=20statement=20for=20get=5Fcurrent=5Factive=5Fuser=20func?= =?UTF-8?q?tion=20=E2=9C=A8=20feat(chat.py):=20add=20current=5Fuser=20depe?= =?UTF-8?q?ndency=20to=20chat=20and=20init=5Fbuild=20endpoints=20to=20enfo?= =?UTF-8?q?rce=20authentication=20and=20authorization=20=F0=9F=94=A7=20ref?= =?UTF-8?q?actor(chat.py):=20pass=20user=5Fid=20to=20vertex.build()=20meth?= =?UTF-8?q?od=20in=20stream=5Fbuild=20endpoint=20to=20associate=20the=20bu?= =?UTF-8?q?ild=20with=20the=20current=20user?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/api/v1/chat.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py index f2b494803..459abcf51 100644 --- a/src/backend/langflow/api/v1/chat.py +++ b/src/backend/langflow/api/v1/chat.py @@ -1,10 +1,18 @@ -from fastapi import APIRouter, HTTPException, WebSocket, WebSocketException, status +from fastapi import ( + APIRouter, + Depends, + HTTPException, + WebSocket, + WebSocketException, + status, +) from fastapi.responses import StreamingResponse from langflow.api.utils import build_input_keys_response from langflow.api.v1.schemas import BuildStatus, BuiltResponse, InitResponse, StreamData from langflow.services import service_manager, ServiceType from langflow.graph.graph.base import Graph +from langflow.services.auth.utils import get_current_active_user from langflow.utils.logger import logger from cachetools import LRUCache @@ -14,7 +22,9 @@ flow_data_store: LRUCache = LRUCache(maxsize=10) @router.websocket("/chat/{client_id}") -async def chat(client_id: str, websocket: WebSocket): +async def chat( + client_id: str, websocket: WebSocket, current_user=Depends(get_current_active_user) +): """Websocket endpoint for chat.""" try: chat_manager = service_manager.get(ServiceType.CHAT_MANAGER) @@ -32,7 +42,9 @@ async def chat(client_id: str, websocket: WebSocket): @router.post("/build/init/{flow_id}", response_model=InitResponse, status_code=201) -async def init_build(graph_data: dict, flow_id: str): +async def init_build( + graph_data: dict, flow_id: str, current_user=Depends(get_current_active_user) +): """Initialize the build by storing graph data and returning a unique session ID.""" try: @@ -54,6 +66,7 @@ async def init_build(graph_data: dict, flow_id: str): flow_data_store[flow_id] = { "graph_data": graph_data, "status": BuildStatus.STARTED, + "user_id": current_user.id, } return InitResponse(flowId=flow_id) @@ -99,6 +112,7 @@ async def stream_build(flow_id: str): return graph_data = flow_data_store[flow_id].get("graph_data") + user_id = flow_data_store[flow_id]["user_id"] if not graph_data: error_message = "No data provided" @@ -119,7 +133,7 @@ async def stream_build(flow_id: str): "log": f"Building node {vertex.vertex_type}", } yield str(StreamData(event="log", data=log_dict)) - vertex.build() + vertex.build(user_id) params = vertex._built_object_repr() valid = True logger.debug(f"Building node {str(vertex.vertex_type)}")