diff --git a/src/backend/langflow/api/v1/callback.py b/src/backend/langflow/api/v1/callback.py index 90799499f..bda05865b 100644 --- a/src/backend/langflow/api/v1/callback.py +++ b/src/backend/langflow/api/v1/callback.py @@ -3,7 +3,7 @@ from uuid import UUID from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler -from langflow.api.v1.schemas import ChatResponse +from langflow.api.v1.schemas import ChatResponse, PromptResponse from typing import Any, Dict, List, Optional @@ -92,10 +92,7 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler): 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="", + resp = PromptResponse( prompt=text, ) await self.websocket.send_json(resp.dict()) diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py index 9a4df7c4f..37e7d712d 100644 --- a/src/backend/langflow/api/v1/schemas.py +++ b/src/backend/langflow/api/v1/schemas.py @@ -86,7 +86,7 @@ class ChatResponse(ChatMessage): """Chat response schema.""" intermediate_steps: str - prompt: Optional[str] = "" + type: str is_bot: bool = True files: list = [] @@ -98,6 +98,14 @@ class ChatResponse(ChatMessage): return v +class PromptResponse(ChatMessage): + """Prompt response schema.""" + + prompt: str + type: str = "prompt" + is_bot: bool = True + + class FileResponse(ChatMessage): """File response schema."""