diff --git a/src/backend/langflow/api/v1/callback.py b/src/backend/langflow/api/v1/callback.py index 3427cd090..c907bea7b 100644 --- a/src/backend/langflow/api/v1/callback.py +++ b/src/backend/langflow/api/v1/callback.py @@ -10,6 +10,7 @@ from typing import Any, Dict, List, Optional from fastapi import WebSocket +from langflow.utils.util import remove_ansi_escape_codes from langchain.schema import AgentAction, AgentFinish from loguru import logger @@ -85,6 +86,16 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler): # This runs when first sending the prompt # to the LLM, adding it will send the final prompt # to the frontend + if "Prompt after formatting" in text: + text = text.replace("Prompt after formatting:\n", "") + text = remove_ansi_escape_codes(text) + resp = ChatResponse( + message="", + type="stream", + intermediate_steps="", + prompt=text, + ) + await self.websocket.send_json(resp.dict()) async def on_agent_action(self, action: AgentAction, **kwargs: Any): log = f"Thought: {action.log}" diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py index 9c6ac6d60..1ccb79063 100644 --- a/src/backend/langflow/api/v1/schemas.py +++ b/src/backend/langflow/api/v1/schemas.py @@ -85,6 +85,7 @@ class ChatResponse(ChatMessage): """Chat response schema.""" intermediate_steps: str + prompt: Optional[str] = "" type: str is_bot: bool = True files: list = [] diff --git a/src/backend/langflow/services/chat/manager.py b/src/backend/langflow/services/chat/manager.py index 3fe937ce8..b1f030f83 100644 --- a/src/backend/langflow/services/chat/manager.py +++ b/src/backend/langflow/services/chat/manager.py @@ -202,7 +202,9 @@ class ChatService(Service): json_payload = await websocket.receive_json() try: payload = orjson.loads(json_payload) - except Exception: + # except TypeError or JSONDecodeError how? + except Exception as exc: + logger.error(f"Error decoding JSON: {exc}") payload = json_payload if "clear_history" in payload: self.chat_history.history[client_id] = [] @@ -220,7 +222,7 @@ class ChatService(Service): ) except Exception as exc: # Handle any exceptions that might occur - logger.error(f"Error handling websocket: {exc}") + logger.exception(f"Error handling websocket: {exc}") await self.close_connection( client_id=client_id, code=status.WS_1011_INTERNAL_ERROR, diff --git a/src/backend/langflow/utils/util.py b/src/backend/langflow/utils/util.py index 4fd9350d2..bad123480 100644 --- a/src/backend/langflow/utils/util.py +++ b/src/backend/langflow/utils/util.py @@ -11,6 +11,10 @@ from langflow.utils import constants from langchain.schema import Document +def remove_ansi_escape_codes(text): + return re.sub(r"\x1b\[[0-9;]*[a-zA-Z]", "", text) + + def build_template_from_function( name: str, type_to_loader_dict: Dict, add_function: bool = False ): diff --git a/src/frontend/src/controllers/API/index.ts b/src/frontend/src/controllers/API/index.ts index a4af5e6a4..ddc8509b7 100644 --- a/src/frontend/src/controllers/API/index.ts +++ b/src/frontend/src/controllers/API/index.ts @@ -1,4 +1,4 @@ -import axios, { AxiosResponse } from "axios"; +import { AxiosResponse } from "axios"; import { ReactFlowJsonObject } from "reactflow"; import { BASE_URL_API } from "../../constants/constants"; import { api } from "../../controllers/API/api"; @@ -522,4 +522,4 @@ export async function deleteApiKey(api_key: string) { console.log("Error:", error); throw error; } -} \ No newline at end of file +} diff --git a/src/frontend/src/modals/formModal/index.tsx b/src/frontend/src/modals/formModal/index.tsx index 49a03af57..0952cd2f0 100644 --- a/src/frontend/src/modals/formModal/index.tsx +++ b/src/frontend/src/modals/formModal/index.tsx @@ -125,12 +125,13 @@ export default function FormModal({ function updateLastMessage({ str, thought, + prompt, end = false, files, }: { str?: string; thought?: string; - // end param default is false + prompt?: string; end?: boolean; files?: Array; }) { @@ -150,6 +151,9 @@ export default function FormModal({ if (files) { newChat[newChat.length - 1].files = files; } + if (prompt) { + newChat[newChat.length - 2].template = prompt; + } return newChat; }); } @@ -201,7 +205,6 @@ export default function FormModal({ } function handleWsMessage(data: any) { - console.log(data); if (Array.isArray(data) && data.length > 0) { //set chat history setChatHistory((_) => { @@ -267,7 +270,11 @@ export default function FormModal({ isStream = false; } if (data.type === "stream" && isStream) { - updateLastMessage({ str: data.message }); + if (data.prompt) { + updateLastMessage({ prompt: data.prompt }); + } else { + updateLastMessage({ str: data.message }); + } } } diff --git a/src/frontend/src/types/chat/index.ts b/src/frontend/src/types/chat/index.ts index 2e7a65fdb..71a081be7 100644 --- a/src/frontend/src/types/chat/index.ts +++ b/src/frontend/src/types/chat/index.ts @@ -8,5 +8,6 @@ export type ChatMessageType = { isSend: boolean; thought?: string; files?: Array<{ data: string; type: string; data_type: string }>; + prompt?: string; chatKey: string; };