Replacing chain with simple LLM

Signed-off-by: Jiri Podivin <[email protected]>

ols/src/query_helpers/question_validator.py

@@ -3,9 +3,6 @@
 import logging
 from typing import Any
 
-from langchain.chains import LLMChain
-from langchain.prompts import PromptTemplate
-
 from ols import config
 from ols.app.metrics import TokenMetricUpdater
 from ols.constants import SUBJECT_REJECTED, GenericLLMParameters
@@ -53,10 +50,9 @@ def validate_question(
         )
         logger.info("%s call settings: %s", conversation_id, settings_string)
 
-        prompt_instructions = PromptTemplate.from_template(
-            prompts.QUESTION_VALIDATOR_PROMPT_TEMPLATE
+        prompt_instructions = prompts.QUESTION_VALIDATOR_PROMPT_TEMPLATE.replace(
+            "{query}", query
         )
-
         bare_llm = self.llm_loader(
             self.provider, self.model, self.generic_llm_params, self.streaming
         )
@@ -70,23 +66,17 @@ def validate_question(
             query, model_config.context_window_size, self.max_tokens_for_response
         )
 
-        llm_chain = LLMChain(
-            llm=bare_llm,
-            prompt=prompt_instructions,
-            verbose=verbose,
-        )
-
         logger.debug("%s validating user query: %s", conversation_id, query)
 
         with TokenMetricUpdater(
             llm=bare_llm,
             provider=provider_config.type,
             model=self.model,
         ) as generic_token_counter:
-            response = llm_chain.invoke(
-                input={"query": query}, config={"callbacks": [generic_token_counter]}
+            response = bare_llm.invoke(
+                input=prompt_instructions, config={"callbacks": [generic_token_counter]}
             )
-        clean_response = str(response["text"]).strip()
+        clean_response = response.strip()
 
         logger.debug(
             "%s query validation response: %s", conversation_id, clean_response

tests/mock_classes/mock_llm_loader.py

@@ -19,6 +19,10 @@ async def astream(self, llm_input, **kwargs):
         # yield input prompt/user query
         yield llm_input[1].content
 
+    def invoke(self, input, config=None, **kwargs):
+        """Transform a single input into an output."""
+        return input
+
 
 def mock_llm_loader(llm=None, expected_params=None):
     """Construct mock for load_llm."""

tests/unit/query_helpers/test_question_validator.py

@@ -1,6 +1,5 @@
 """Unit tests for QuestionValidator class."""
 
-from unittest.mock import patch
 
 import pytest
 
@@ -14,7 +13,6 @@
     QueryHelper,
     QuestionValidator,
 )
-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
 
 
@@ -63,7 +61,6 @@ def test_passing_parameters():
     )
 
 
-@patch("ols.src.query_helpers.question_validator.LLMChain", new=mock_llm_chain(None))
 def test_validate_question_llm_loader():
     """Test that LLM is loaded within validate_question method with proper parameters."""
     # it is needed to initialize configuration in order to be able