feat: add history to prompt (#480)

Author: mhordynski

packages/ragbits-core/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 - Make the score in VectorStoreResult consistent (always bigger is better)
 - Add router option to LiteLLMEmbedder (#440)
+- New methods in Prompt class for appending conversation history (#480)
 - Fix: make unflatten_dict symmetric to flatten_dict (#461)
 - Cost and capabilities config for custom litellm models (#481)
 

packages/ragbits-core/src/ragbits/core/prompt/prompt.py

@@ -142,6 +142,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # Additional few shot examples that can be added dynamically using methods
         # (in opposite to the static `few_shots` attribute which is defined in the class)
         self._instance_few_shots: list[FewShotExample[InputT, OutputT]] = []
+
+        # Additional conversation history that can be added dynamically using methods
+        self._conversation_history: list[dict[str, Any]] = []
         super().__init__()
 
     @property
@@ -166,6 +169,7 @@ def chat(self) -> ChatFormat:
             ),
             *self.list_few_shots(),
             {"role": "user", "content": user_content},
+            *self._conversation_history,
         ]
         return chat
 
@@ -216,6 +220,57 @@ def list_few_shots(self) -> ChatFormat:
             result.append({"role": "assistant", "content": assistant_content})
         return result
 
+    def add_user_message(self, message: str | dict[str, Any] | InputT) -> "Prompt[InputT, OutputT]":
+        """
+        Add a user message to the conversation history.
+
+        Args:
+            message (str | dict[str, Any] | InputT): The user message content. Can be:
+                - A string: Used directly as content
+                - A dictionary: With format {"type": "text", "text": "message"} or image content
+                - An InputT model: Will be rendered using the user prompt template
+
+        Returns:
+            Prompt[InputT, OutputT]: The current prompt instance to allow chaining.
+        """
+        content: str | list[dict[str, Any]] | dict[str, Any] | InputT
+
+        if isinstance(message, BaseModel):
+            # Type checking to ensure we're passing InputT to the methods
+            input_model: InputT = cast(InputT, message)
+
+            # Render the message using the template if it's an input model
+            rendered_text = self._render_template(self.user_prompt_template, input_model)
+            images_in_input = self._get_images_from_input_data(input_model)
+
+            if images_in_input:
+                content = [{"type": "text", "text": rendered_text}] + [
+                    self._create_message_with_image(image) for image in images_in_input
+                ]
+            else:
+                content = rendered_text
+        else:
+            # Use the message directly if it's a string or dict
+            content = message
+
+        self._conversation_history.append({"role": "user", "content": content})
+        return self
+
+    def add_assistant_message(self, message: str | OutputT) -> "Prompt[InputT, OutputT]":
+        """
+        Add an assistant message to the conversation history.
+
+        Args:
+            message (str): The assistant message content.
+
+        Returns:
+            Prompt[InputT, OutputT]: The current prompt instance to allow chaining.
+        """
+        if isinstance(message, BaseModel):
+            message = message.model_dump_json()
+        self._conversation_history.append({"role": "assistant", "content": str(message)})
+        return self
+
     def list_images(self) -> list[str]:
         """
         Returns the images in form of URLs or base64 encoded strings.

packages/ragbits-core/tests/unit/prompts/test_prompt.py

@@ -537,3 +537,99 @@ def sync_parser(response: str) -> str:
     test_prompt.response_parser = sync_parser
     resp_sync = await test_prompt.parse_response(resp)
     assert resp_sync == "hello human"
+
+
+def test_add_user_message_with_string():
+    """Test adding a user message with a string content."""
+
+    class TestPrompt(Prompt):
+        user_prompt = "Hello"
+
+    prompt = TestPrompt()
+    prompt.add_user_message("Additional message")
+
+    assert prompt.chat == [{"role": "user", "content": "Hello"}, {"role": "user", "content": "Additional message"}]
+
+
+def test_add_user_message_with_input_model():
+    """Test adding a user message with an input model."""
+
+    class TestPrompt(Prompt[_PromptInput, str]):
+        user_prompt = "Hello {{ name }}"
+
+    prompt = TestPrompt(_PromptInput(name="Alice", age=30, theme="rock"))
+    prompt.add_user_message(_PromptInput(name="Bob", age=25, theme="jazz"))
+
+    assert prompt.chat == [{"role": "user", "content": "Hello Alice"}, {"role": "user", "content": "Hello Bob"}]
+
+
+def test_add_user_message_with_image():
+    """Test adding a user message with an image."""
+
+    class ImagePrompt(Prompt):
+        user_prompt = "What is on this image?"
+        image_input_fields = ["image"]
+
+    prompt = ImagePrompt(_ImagePromptInput(image=_get_image_bytes()))
+    prompt.add_user_message(_ImagePromptInput(image=_get_image_bytes()))
+
+    assert len(prompt.chat) == 2
+    assert prompt.chat[0]["role"] == "user"
+    assert prompt.chat[1]["role"] == "user"
+    assert len(prompt.chat[0]["content"]) == 2  # text + image
+    assert len(prompt.chat[1]["content"]) == 2  # text + image
+
+
+def test_add_assistant_message():
+    """Test adding an assistant message."""
+
+    class TestPrompt(Prompt[_PromptInput, _PromptOutput]):
+        user_prompt = "Hello {{ name }}"
+
+    prompt = TestPrompt(_PromptInput(name="Alice", age=30, theme="rock"))
+    prompt.add_assistant_message("Assistant response")
+
+    assert prompt.chat == [
+        {"role": "user", "content": "Hello Alice"},
+        {"role": "assistant", "content": "Assistant response"},
+    ]
+
+
+def test_add_assistant_message_with_model():
+    """Test adding an assistant message with a model output."""
+
+    class TestPrompt(Prompt[_PromptInput, _PromptOutput]):
+        user_prompt = "Hello {{ name }}"
+
+    prompt = TestPrompt(_PromptInput(name="Alice", age=30, theme="rock"))
+    output = _PromptOutput(song_title="Test Song", song_lyrics="Test Lyrics")
+    prompt.add_assistant_message(output)
+
+    assert prompt.chat == [
+        {"role": "user", "content": "Hello Alice"},
+        {"role": "assistant", "content": output.model_dump_json()},
+    ]
+
+
+def test_conversation_history():
+    """Test building a complete conversation history with multiple messages."""
+
+    class TestPrompt(Prompt[_PromptInput, _PromptOutput]):
+        user_prompt = "Hello {{ name }}"
+
+    prompt = TestPrompt(_PromptInput(name="Alice", age=30, theme="rock"))
+    prompt.add_user_message("How are you?")
+    prompt.add_assistant_message("I'm doing well!")
+    prompt.add_user_message(_PromptInput(name="Bob", age=25, theme="jazz"))
+    prompt.add_assistant_message(_PromptOutput(song_title="Jazz Song", song_lyrics="Jazz lyrics"))
+
+    assert prompt.chat == [
+        {"role": "user", "content": "Hello Alice"},
+        {"role": "user", "content": "How are you?"},
+        {"role": "assistant", "content": "I'm doing well!"},
+        {"role": "user", "content": "Hello Bob"},
+        {
+            "role": "assistant",
+            "content": _PromptOutput(song_title="Jazz Song", song_lyrics="Jazz lyrics").model_dump_json(),
+        },
+    ]