Pass streaming parameter to LLM loader

Signed-off-by: Pavel Tisnovsky <[email protected]>

ols/app/endpoints/ols.py

@@ -376,6 +376,7 @@ def generate_response(
             provider=llm_request.provider,
             model=llm_request.model,
             system_prompt=llm_request.system_prompt,
+            streaming=streaming,
         )
         history = CacheEntry.cache_entries_to_history(previous_input)
         if streaming:

ols/src/llms/llm_loader.py

@@ -53,14 +53,18 @@ def resolve_provider_config(
 
 
 def load_llm(
-    provider: str, model: str, generic_llm_params: Optional[dict] = None
+    provider: str,
+    model: str,
+    generic_llm_params: Optional[dict] = None,
+    streaming: Optional[bool] = None,
 ) -> LLM:
     """Load LLM according to input provider and model.
 
     Args:
         provider: The provider name.
         model: The model name.
         generic_llm_params: The optional parameters that will be converted into LLM-specific ones.
+        streaming: The optional parameter that enable streaming on LLM side if set to True.
 
     Raises:
         LLMConfigurationError: If the whole provider configuration is missing.

ols/src/query_helpers/docs_summarizer.py

@@ -42,7 +42,7 @@ def _prepare_llm(self) -> None:
             GenericLLMParameters.MAX_TOKENS_FOR_RESPONSE: self.model_config.parameters.max_tokens_for_response  # noqa: E501
         }
         self.bare_llm = self.llm_loader(
-            self.provider, self.model, self.generic_llm_params
+            self.provider, self.model, self.generic_llm_params, self.streaming
         )
 
     def _get_system_prompt(self) -> None:

ols/src/query_helpers/query_helper.py

@@ -23,6 +23,7 @@ def __init__(
         generic_llm_params: Optional[dict] = None,
         llm_loader: Optional[Callable[[str, str, dict], LLM]] = None,
         system_prompt: Optional[str] = None,
+        streaming: Optional[bool] = None,
     ) -> None:
         """Initialize query helper."""
         # NOTE: As signature of this method is evaluated before the config,
@@ -32,6 +33,7 @@ def __init__(
         self.model = model or config.ols_config.default_model
         self.generic_llm_params = generic_llm_params or {}
         self.llm_loader = llm_loader or load_llm
+        self.streaming = streaming or False
 
         self._system_prompt = (
             (config.dev_config.enable_system_prompt_override and system_prompt)