From 3f36713d8b856f8ce360a821b3d0214fbffb9058 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Sun, 6 Oct 2024 19:34:44 +0200 Subject: [PATCH] return a single chat message instead of an empty list --- src/banks/prompt.py | 6 ++++++ tests/test_prompt.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/src/banks/prompt.py b/src/banks/prompt.py index 82c4969..56c513f 100644 --- a/src/banks/prompt.py +++ b/src/banks/prompt.py @@ -134,6 +134,12 @@ def chat_messages(self, data: dict[str, Any] | None = None) -> list[ChatMessage] except ValidationError: # Ignore lines that are not a message pass + + if not messages: + # fallback, if there was no {% chat %} block in the template, + # put the text into a single user message + messages.append(ChatMessage(role="user", content=rendered)) + return messages diff --git a/tests/test_prompt.py b/tests/test_prompt.py index a206ec3..b1db95c 100644 --- a/tests/test_prompt.py +++ b/tests/test_prompt.py @@ -111,3 +111,9 @@ def test_chat_messages_cached(): p = Prompt(p_file.read_text(), render_cache=mock_cache) p.chat_messages() mock_cache.set.assert_called_once() + + +def test_chat_message_no_chat_tag(): + text = "This is raw text" + p = Prompt(text=text) + assert p.chat_messages() == [ChatMessage(role="user", content=text)]