fix review comments

RasaHQ · Oct 10, 2023 · e278ed8 · e278ed8
1 parent 642aea1
commit e278ed8
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 21 deletions.
diff --git a/data/test_prompt_templates/test_prompt.jinja2 b/data/test_prompt_templates/test_prompt.jinja2
@@ -1 +1 @@
-This is a test prompt.
+This is a test prompt.
diff --git a/rasa/dialogue_understanding/generator/llm_command_generator.py b/rasa/dialogue_understanding/generator/llm_command_generator.py
@@ -81,8 +81,8 @@ def __init__(
     ) -> None:
         self.config = {**self.get_default_config(), **config}
         self.prompt_template = get_prompt_template(
-            DEFAULT_COMMAND_PROMPT_TEMPLATE,
             config.get("prompt"),
+            DEFAULT_COMMAND_PROMPT_TEMPLATE,
         )
         self._model_storage = model_storage
         self._resource = resource

diff --git a/rasa/shared/utils/llm.py b/rasa/shared/utils/llm.py
@@ -212,7 +212,7 @@ def embedder_factory(
 
 
 def get_prompt_template(
-    default_prompt_template: Text, jinja_file_path: Optional[Text]
+    jinja_file_path: Optional[Text], default_prompt_template: Text
 ) -> Text:
     """Returns the prompt template.
 

diff --git a/tests/cdu/generator/test_llm_command_generator.py b/tests/cdu/generator/test_llm_command_generator.py
@@ -22,21 +22,21 @@ def model_storage(tmp_path_factory: TempPathFactory) -> ModelStorage:
     return LocalModelStorage(tmp_path_factory.mktemp(uuid.uuid4().hex))
 
 
-@pytest.mark.parametrize(
-    "config, expected",
-    [
-        (
-            {"prompt": "data/test_prompt_templates/test_prompt.jinja2"},
-            "This is a test prompt.",
-        ),
-        (
-            {},
-            "Your task is to analyze the current conversation",
-        ),
-    ],
-)
-async def test_llm_command_generator_prompt_initialisation(
-    model_storage: ModelStorage, resource: Resource, config: Dict, expected: str
-):
-    generator = LLMCommandGenerator(config, model_storage, resource)
-    assert generator.prompt_template.startswith(expected)
+async def test_llm_command_generator_prompt_init_custom(
+    model_storage: ModelStorage, resource: Resource
+) -> None:
+    generator = LLMCommandGenerator(
+        {"prompt": "data/test_prompt_templates/test_prompt.jinja2"},
+        model_storage,
+        resource,
+    )
+    assert generator.prompt_template.startswith("This is a test prompt.")
+
+
+async def test_llm_command_generator_prompt_init_default(
+    model_storage: ModelStorage, resource: Resource
+) -> None:
+    generator = LLMCommandGenerator({}, model_storage, resource)
+    assert generator.prompt_template.startswith(
+        "Your task is to analyze the current conversation"
+    )