From dc72f2e52fd6e41de17d078d2bdfd77ce81056ea Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Tue, 25 Apr 2023 21:06:29 -0300 Subject: [PATCH 1/3] refactor(test_websocket.py): change sender field to is_bot field in chat history tests --- tests/test_websocket.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_websocket.py b/tests/test_websocket.py index 74e147075..9ec128bcf 100644 --- a/tests/test_websocket.py +++ b/tests/test_websocket.py @@ -28,7 +28,7 @@ def test_chat_history(client: TestClient): # Receive the response from the server response = websocket.receive_json() assert json.loads(response) == { - "sender": "bot", + "is_bot": True, "message": None, "intermediate_steps": "", "type": "start", @@ -40,7 +40,7 @@ def test_chat_history(client: TestClient): # Receive the response from the server response = websocket.receive_json() assert json.loads(response) == { - "sender": "bot", + "is_bot": True, "message": "Hello, I'm a mock response!", "intermediate_steps": "", "type": "end", From ea210af19b461a8a549f5a66661859bf43b2eeac Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Tue, 25 Apr 2023 21:46:02 -0300 Subject: [PATCH 2/3] refactor(chat_manager.py, utils.py): remove unused imports and functions, move pil_to_base64 and try_setting_streaming_options to utils module --- src/backend/langflow/api/chat_manager.py | 41 +++--------------------- src/backend/langflow/interface/utils.py | 34 ++++++++++++++++++++ 2 files changed, 39 insertions(+), 36 deletions(-) diff --git a/src/backend/langflow/api/chat_manager.py b/src/backend/langflow/api/chat_manager.py index b71f652d1..f0847b894 100644 --- a/src/backend/langflow/api/chat_manager.py +++ b/src/backend/langflow/api/chat_manager.py @@ -1,24 +1,17 @@ import asyncio -import base64 -from io import BytesIO from typing import Dict, List from collections import defaultdict from fastapi import WebSocket import json -from langchain.llms import OpenAI, AzureOpenAI -from langchain.chat_models import ChatOpenAI, AzureChatOpenAI from langflow.api.schemas import ChatMessage, ChatResponse, FileResponse -from langflow.cache.manager import AsyncSubject, Subject -from langchain.callbacks.base import AsyncCallbackManager -from langflow.api.callback import StreamingLLMCallbackHandler +from langflow.cache.manager import Subject from langflow.interface.run import ( - async_get_result_and_steps, get_result_and_steps, load_or_build_langchain_object, ) +from langflow.interface.utils import pil_to_base64, try_setting_streaming_options from langflow.utils.logger import logger from langflow.cache import cache_manager -from PIL.Image import Image class ChatHistory(Subject): @@ -149,6 +142,9 @@ class ChatManager: payload = json.loads(json_payload) except TypeError: payload = json_payload + if "clear_history" in payload: + self.chat_history.history[client_id] = [] + with self.cache_manager.set_client_id(client_id): await self.process_message(client_id, payload) except Exception as e: @@ -187,30 +183,3 @@ async def process_graph( # Log stack trace logger.exception(e) raise e - - -def try_setting_streaming_options(langchain_object, websocket): - # If the LLM type is OpenAI or ChatOpenAI, - # set streaming to True - # First we need to find the LLM - llm = None - if hasattr(langchain_object, "llm"): - llm = langchain_object.llm - elif hasattr(langchain_object, "llm_chain") and hasattr( - langchain_object.llm_chain, "llm" - ): - llm = langchain_object.llm_chain.llm - if isinstance(llm, (OpenAI, ChatOpenAI, AzureOpenAI, AzureChatOpenAI)): - llm.streaming = bool(hasattr(llm, "streaming")) - stream_handler = StreamingLLMCallbackHandler(websocket) - stream_manager = AsyncCallbackManager([stream_handler]) - llm.callback_manager = stream_manager - - return langchain_object - - -def pil_to_base64(image: Image) -> str: - buffered = BytesIO() - image.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()) - return img_str.decode("utf-8") diff --git a/src/backend/langflow/interface/utils.py b/src/backend/langflow/interface/utils.py index b3b154790..604df6f28 100644 --- a/src/backend/langflow/interface/utils.py +++ b/src/backend/langflow/interface/utils.py @@ -1,5 +1,12 @@ +import base64 +from io import BytesIO import json import os +from PIL.Image import Image +from langchain.callbacks.base import AsyncCallbackManager +from langchain.chat_models import AzureChatOpenAI, ChatOpenAI +from langchain.llms import AzureOpenAI, OpenAI +from langflow.api.callback import StreamingLLMCallbackHandler import yaml @@ -20,3 +27,30 @@ def load_file_into_dict(file_path: str) -> dict: raise ValueError("Unsupported file type. Please provide a JSON or YAML file.") return data + + +def pil_to_base64(image: Image) -> str: + buffered = BytesIO() + image.save(buffered, format="PNG") + img_str = base64.b64encode(buffered.getvalue()) + return img_str.decode("utf-8") + + +def try_setting_streaming_options(langchain_object, websocket): + # If the LLM type is OpenAI or ChatOpenAI, + # set streaming to True + # First we need to find the LLM + llm = None + if hasattr(langchain_object, "llm"): + llm = langchain_object.llm + elif hasattr(langchain_object, "llm_chain") and hasattr( + langchain_object.llm_chain, "llm" + ): + llm = langchain_object.llm_chain.llm + if isinstance(llm, (OpenAI, ChatOpenAI, AzureOpenAI, AzureChatOpenAI)): + llm.streaming = bool(hasattr(llm, "streaming")) + stream_handler = StreamingLLMCallbackHandler(websocket) + stream_manager = AsyncCallbackManager([stream_handler]) + llm.callback_manager = stream_manager + + return langchain_object From 167d6c0f1e5e1e9e286f6e98d37ad67a7ed23168 Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Tue, 25 Apr 2023 21:47:50 -0300 Subject: [PATCH 3/3] refactor(chat_manager.py): add continue statement to clear_history block to skip message processing --- src/backend/langflow/api/chat_manager.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/backend/langflow/api/chat_manager.py b/src/backend/langflow/api/chat_manager.py index f0847b894..ca31139ab 100644 --- a/src/backend/langflow/api/chat_manager.py +++ b/src/backend/langflow/api/chat_manager.py @@ -144,6 +144,7 @@ class ChatManager: payload = json_payload if "clear_history" in payload: self.chat_history.history[client_id] = [] + continue with self.cache_manager.set_client_id(client_id): await self.process_message(client_id, payload)