From 3ed25b7f6faaf6ea3295e7a756031cbe0c4697c2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pavel=20Ti=C5=A1novsk=C3=BD?= Date: Mon, 27 Jan 2025 15:27:49 +0100 Subject: [PATCH] Test formatted history from Granite (#321) --- tests/unit/utils/test_token_handler.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/unit/utils/test_token_handler.py b/tests/unit/utils/test_token_handler.py index ddcb57b9..32b71e0c 100644 --- a/tests/unit/utils/test_token_handler.py +++ b/tests/unit/utils/test_token_handler.py @@ -7,6 +7,7 @@ from langchain_core.messages import AIMessage, HumanMessage from ols.constants import TOKEN_BUFFER_WEIGHT, ModelFamily +from ols.src.prompts.prompt_generator import restructure_history from ols.utils.token_handler import PromptTooLongError, TokenHandler from tests.mock_classes.mock_retrieved_node import MockRetrievedNode @@ -282,3 +283,12 @@ def test_limit_conversation_history(self): # history should truncate to empty list and flag should be True assert truncated_history == [] assert truncated + + # Test formatted history for granite + model = ModelFamily.GRANITE + truncated_history, truncated = ( + self._token_handler_obj.limit_conversation_history(history, model, 22) + ) + assert len(truncated_history) == 1 + assert truncated_history[-1] == restructure_history(history[-1], model) + assert truncated