diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py index 611407e8d..2b985584e 100644 --- a/src/backend/langflow/api/v1/chat.py +++ b/src/backend/langflow/api/v1/chat.py @@ -6,19 +6,16 @@ from langflow.api.v1.schemas import BuildStatus, BuiltResponse, InitResponse, St from langflow.services import service_manager, ServiceType from langflow.graph.graph.base import Graph from langflow.utils.logger import logger -from cachetools import LRUCache router = APIRouter(tags=["Chat"]) -flow_data_store: LRUCache = LRUCache(maxsize=10) - @router.websocket("/chat/{client_id}") async def chat(client_id: str, websocket: WebSocket): """Websocket endpoint for chat.""" try: chat_manager = service_manager.get(ServiceType.CHAT_MANAGER) - if client_id in chat_manager.in_memory_cache: + if client_id in chat_manager.cache_manager: await chat_manager.handle_websocket(client_id, websocket) else: # We accept the connection but close it immediately @@ -34,23 +31,23 @@ 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): """Initialize the build by storing graph data and returning a unique session ID.""" - + flow_data_store = service_manager.get(ServiceType.CACHE_MANAGER) try: if flow_id is None: raise ValueError("No ID provided") # Check if already building if ( flow_id in flow_data_store - and flow_data_store[flow_id]["status"] == BuildStatus.IN_PROGRESS + and isinstance(flow_data_store[flow_id], dict) + and flow_data_store[flow_id].get("status") == BuildStatus.IN_PROGRESS ): return InitResponse(flowId=flow_id) # Delete from cache if already exists chat_manager = service_manager.get(ServiceType.CHAT_MANAGER) - if flow_id in chat_manager.in_memory_cache: - with chat_manager.in_memory_cache._lock: - chat_manager.in_memory_cache.delete(flow_id) - logger.debug(f"Deleted flow {flow_id} from cache") + if flow_id in chat_manager.cache_manager: + chat_manager.cache_manager.delete(flow_id) + logger.debug(f"Deleted flow {flow_id} from cache") flow_data_store[flow_id] = { "graph_data": graph_data, "status": BuildStatus.STARTED, @@ -65,6 +62,7 @@ async def init_build(graph_data: dict, flow_id: str): @router.get("/build/{flow_id}/status", response_model=BuiltResponse) async def build_status(flow_id: str): """Check the flow_id is in the flow_data_store.""" + flow_data_store = service_manager.get(ServiceType.CACHE_MANAGER) try: built = ( flow_id in flow_data_store @@ -83,6 +81,7 @@ async def build_status(flow_id: str): @router.get("/build/stream/{flow_id}", response_class=StreamingResponse) async def stream_build(flow_id: str): """Stream the build process based on stored flow data.""" + flow_data_store = service_manager.get(ServiceType.CACHE_MANAGER) async def event_stream(flow_id): final_response = {"end_of_stream": True} @@ -163,7 +162,7 @@ async def stream_build(flow_id: str): } yield str(StreamData(event="message", data=input_keys_response)) chat_manager = service_manager.get(ServiceType.CHAT_MANAGER) - chat_manager.set_cache(flow_id, langchain_object) + chat_manager.set_cache(f"{flow_id}_chat", langchain_object) # We need to reset the chat history chat_manager.chat_history.empty_history(flow_id) flow_data_store[flow_id]["status"] = BuildStatus.SUCCESS diff --git a/src/backend/langflow/chat/manager.py b/src/backend/langflow/chat/manager.py new file mode 100644 index 000000000..6c78a2d77 --- /dev/null +++ b/src/backend/langflow/chat/manager.py @@ -0,0 +1,217 @@ +from collections import defaultdict +from fastapi import WebSocket, status +from langflow.api.v1.schemas import ChatMessage, ChatResponse, FileResponse +from langflow.cache import cache_manager +from langflow.cache.manager import Subject +from langflow.chat.utils import process_graph +from langflow.interface.utils import pil_to_base64 +from langflow.utils.logger import logger + + +import asyncio +import json +from typing import Any, Dict, List + +from langflow.cache.flow import InMemoryCache + + +class ChatHistory(Subject): + def __init__(self): + super().__init__() + self.history: Dict[str, List[ChatMessage]] = defaultdict(list) + + def add_message(self, client_id: str, message: ChatMessage): + """Add a message to the chat history.""" + + self.history[client_id].append(message) + + if not isinstance(message, FileResponse): + self.notify() + + def get_history(self, client_id: str, filter_messages=True) -> List[ChatMessage]: + """Get the chat history for a client.""" + if history := self.history.get(client_id, []): + if filter_messages: + return [msg for msg in history if msg.type not in ["start", "stream"]] + return history + else: + return [] + + def empty_history(self, client_id: str): + """Empty the chat history for a client.""" + self.history[client_id] = [] + + +class ChatManager: + def __init__(self): + self.active_connections: Dict[str, WebSocket] = {} + self.chat_history = ChatHistory() + self.cache_manager = cache_manager + self.cache_manager.attach(self.update) + self.in_memory_cache = InMemoryCache() + + def on_chat_history_update(self): + """Send the last chat message to the client.""" + client_id = self.cache_manager.current_client_id + if client_id in self.active_connections: + chat_response = self.chat_history.get_history( + client_id, filter_messages=False + )[-1] + if chat_response.is_bot: + # Process FileResponse + if isinstance(chat_response, FileResponse): + # If data_type is pandas, convert to csv + if chat_response.data_type == "pandas": + chat_response.data = chat_response.data.to_csv() + elif chat_response.data_type == "image": + # Base64 encode the image + chat_response.data = pil_to_base64(chat_response.data) + # get event loop + loop = asyncio.get_event_loop() + + coroutine = self.send_json(client_id, chat_response) + asyncio.run_coroutine_threadsafe(coroutine, loop) + + def update(self): + if self.cache_manager.current_client_id in self.active_connections: + self.last_cached_object_dict = self.cache_manager.get_last() + # Add a new ChatResponse with the data + chat_response = FileResponse( + message=None, + type="file", + data=self.last_cached_object_dict["obj"], + data_type=self.last_cached_object_dict["type"], + ) + + self.chat_history.add_message( + self.cache_manager.current_client_id, chat_response + ) + + async def connect(self, client_id: str, websocket: WebSocket): + await websocket.accept() + self.active_connections[client_id] = websocket + + def disconnect(self, client_id: str): + self.active_connections.pop(client_id, None) + + async def send_message(self, client_id: str, message: str): + websocket = self.active_connections[client_id] + await websocket.send_text(message) + + async def send_json(self, client_id: str, message: ChatMessage): + websocket = self.active_connections[client_id] + await websocket.send_json(message.dict()) + + async def close_connection(self, client_id: str, code: int, reason: str): + if websocket := self.active_connections[client_id]: + try: + await websocket.close(code=code, reason=reason) + self.disconnect(client_id) + except RuntimeError as exc: + # This is to catch the following error: + # Unexpected ASGI message 'websocket.close', after sending 'websocket.close' + if "after sending" in str(exc): + logger.error(f"Error closing connection: {exc}") + + async def process_message( + self, client_id: str, payload: Dict, langchain_object: Any + ): + # Process the graph data and chat message + chat_inputs = payload.pop("inputs", {}) + chat_inputs = ChatMessage(message=chat_inputs) + self.chat_history.add_message(client_id, chat_inputs) + + # graph_data = payload + start_resp = ChatResponse(message=None, type="start", intermediate_steps="") + await self.send_json(client_id, start_resp) + + # is_first_message = len(self.chat_history.get_history(client_id=client_id)) <= 1 + # Generate result and thought + try: + logger.debug("Generating result and thought") + + result, intermediate_steps = await process_graph( + langchain_object=langchain_object, + chat_inputs=chat_inputs, + websocket=self.active_connections[client_id], + ) + except Exception as e: + # Log stack trace + logger.exception(e) + self.chat_history.empty_history(client_id) + raise e + # Send a response back to the frontend, if needed + intermediate_steps = intermediate_steps or "" + history = self.chat_history.get_history(client_id, filter_messages=False) + file_responses = [] + if history: + # Iterate backwards through the history + for msg in reversed(history): + if isinstance(msg, FileResponse): + if msg.data_type == "image": + # Base64 encode the image + if isinstance(msg.data, str): + continue + msg.data = pil_to_base64(msg.data) + file_responses.append(msg) + if msg.type == "start": + break + + response = ChatResponse( + message=result, + intermediate_steps=intermediate_steps.strip(), + type="end", + files=file_responses, + ) + await self.send_json(client_id, response) + self.chat_history.add_message(client_id, response) + + def set_cache(self, client_id: str, langchain_object: Any) -> bool: + """ + Set the cache for a client. + """ + client_id = f"{client_id}_chat" + self.cache_manager.set(client_id, langchain_object) + return client_id in self.cache_manager + + async def handle_websocket(self, client_id: str, websocket: WebSocket): + await self.connect(client_id, websocket) + + try: + chat_history = self.chat_history.get_history(client_id) + # iterate and make BaseModel into dict + chat_history = [chat.dict() for chat in chat_history] + await websocket.send_json(chat_history) + + while True: + json_payload = await websocket.receive_json() + try: + payload = json.loads(json_payload) + except TypeError: + payload = json_payload + if "clear_history" in payload: + self.chat_history.history[client_id] = [] + continue + + with self.cache_manager.set_client_id(client_id): + langchain_object = self.in_memory_cache.get(f"{client_id}_chat") + await self.process_message(client_id, payload, langchain_object) + + except Exception as exc: + # Handle any exceptions that might occur + logger.error(f"Error handling websocket: {exc}") + await self.close_connection( + client_id=client_id, + code=status.WS_1011_INTERNAL_ERROR, + reason=str(exc)[:120], + ) + finally: + try: + await self.close_connection( + client_id=client_id, + code=status.WS_1000_NORMAL_CLOSURE, + reason="Client disconnected", + ) + except Exception as exc: + logger.error(f"Error closing connection: {exc}") + self.disconnect(client_id)