From a3decd2061bf89f43c188c73f7cb7287ed4bef22 Mon Sep 17 00:00:00 2001 From: Pavel Tisnovsky Date: Mon, 27 Jan 2025 11:16:17 +0100 Subject: [PATCH] Merge changes into unit test for DocsSummarizer --- .../query_helpers/test_docs_summarizer.py | 40 ++++++++++++++----- 1 file changed, 31 insertions(+), 9 deletions(-) diff --git a/tests/unit/query_helpers/test_docs_summarizer.py b/tests/unit/query_helpers/test_docs_summarizer.py index 91ec3e8d..c82d6347 100644 --- a/tests/unit/query_helpers/test_docs_summarizer.py +++ b/tests/unit/query_helpers/test_docs_summarizer.py @@ -7,15 +7,25 @@ from langchain_core.messages import AIMessage, HumanMessage from ols import config -from ols.app.models.config import LoggingConfig -from ols.src.query_helpers.docs_summarizer import DocsSummarizer, QueryHelper -from ols.utils import suid -from ols.utils.logging_configurator import configure_logging -from tests import constants -from tests.mock_classes.mock_langchain_interface import mock_langchain_interface -from tests.mock_classes.mock_llama_index import MockLlamaIndex -from tests.mock_classes.mock_llm_chain import mock_llm_chain -from tests.mock_classes.mock_llm_loader import mock_llm_loader + +# needs to be setup there before is_user_authorized is imported +config.ols_config.authentication_config.module = "k8s" + + +from ols.app.models.config import LoggingConfig # noqa:E402 +from ols.src.query_helpers.docs_summarizer import ( # noqa:E402 + DocsSummarizer, + QueryHelper, +) +from ols.utils import suid # noqa:E402 +from ols.utils.logging_configurator import configure_logging # noqa:E402 +from tests import constants # noqa:E402 +from tests.mock_classes.mock_langchain_interface import ( # noqa:E402 + mock_langchain_interface, +) +from tests.mock_classes.mock_llama_index import MockLlamaIndex # noqa:E402 +from tests.mock_classes.mock_llm_chain import mock_llm_chain # noqa:E402 +from tests.mock_classes.mock_llm_loader import mock_llm_loader # noqa:E402 conversation_id = suid.get_suid() @@ -51,6 +61,18 @@ def test_if_system_prompt_was_updated(): assert summarizer.system_prompt == expected_prompt +def test_docs_summarizer_streaming_parameter(): + """Test if optional streaming parameter is stored.""" + summarizer = DocsSummarizer(llm_loader=mock_llm_loader(None)) + assert summarizer.streaming is False + + summarizer = DocsSummarizer(llm_loader=mock_llm_loader(None), streaming=False) + assert summarizer.streaming is False + + summarizer = DocsSummarizer(llm_loader=mock_llm_loader(None), streaming=True) + assert summarizer.streaming is True + + @patch("ols.utils.token_handler.RAG_SIMILARITY_CUTOFF", 0.4) @patch("ols.utils.token_handler.MINIMUM_CONTEXT_TOKEN_LIMIT", 1) @patch("ols.src.query_helpers.docs_summarizer.LLMChain", new=mock_llm_chain(None))