From 2e63d07a80ac02f920c936a9ca566cee3dbc8cd6 Mon Sep 17 00:00:00 2001
From: bparees <bparees@redhat.com>
Date: Fri, 10 Nov 2023 16:03:05 -0500
Subject: [PATCH] Improve yaml gen prompt to ensure only single yaml object
 response

---
 modules/task_processor.py |  2 +-
 modules/yaml_generator.py | 37 ++++++++++++++++++++++++-------------
 2 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/modules/task_processor.py b/modules/task_processor.py
index f39358254..04f200ace 100644
--- a/modules/task_processor.py
+++ b/modules/task_processor.py
@@ -38,7 +38,7 @@ def process_tasks(self, conversation, tasklist, original_query, **kwargs):
         else:
             verbose = False
 
-        settings_string = f"conversation: {conversation}, query: {original_query},model: {model}, verbose: {verbose}"
+        settings_string = f"conversation: {conversation}, tasklist: {tasklist}, query: {original_query},model: {model}, verbose: {verbose}"
         self.logger.info(
             conversation
             + " call settings: "
diff --git a/modules/yaml_generator.py b/modules/yaml_generator.py
index 1d56b1241..2499edd0c 100644
--- a/modules/yaml_generator.py
+++ b/modules/yaml_generator.py
@@ -46,7 +46,17 @@ def generate_yaml(self, conversation, string, **kwargs):
 
         bare_llm = get_watsonx_predictor(model=model)
 
-        prompt_instructions = PromptTemplate.from_template("{string}\n")
+        prompt_instructions = PromptTemplate.from_template(
+            """Instructions:
+- Produce only a yaml response to the user request
+- Do not augment the response with markdown or other formatting beyond standard yaml formatting
+- Only provide a single yaml object containg a single resource type in your response, do not provide multiple yaml objects
+
+User Request: {string}
+"""
+        )
+
+
         llm_chain = LLMChain(llm=bare_llm, verbose=verbose, prompt=prompt_instructions)
 
         task_query = prompt_instructions.format(string=string)
@@ -56,18 +66,19 @@ def generate_yaml(self, conversation, string, **kwargs):
         response = llm_chain(inputs={"string": string})
 
         # https://stackoverflow.com/a/63082323/2328066
-        regex = r"(?:\n+|\A)?(?P<code_block>\`{3,}yaml\n+[\S\s]*\`{3,})"
-
-        match = re.search(regex, response["text"])
-
-        if match:
-            clean_response = match.group("code_block")
-            self.logger.info(conversation + " generated yaml: " + clean_response)
-            return clean_response
-        else:
-            # TODO: need to do the right thing here - raise an exception?
-            self.logger.error(conversation + " could not parse response:\n"+ response["text"])
-            return "some failure"
+        # regex = r"(?:\n+|\A)?(?P<code_block>\`{3,}yaml\n+[\S\s]*\`{3,})"
+
+        # match = re.search(regex, response["text"])
+
+        # if match:
+        #     clean_response = match.group("code_block")
+        #     self.logger.info(conversation + " generated yaml: " + clean_response)
+        #     return clean_response
+        # else:
+        #     # TODO: need to do the right thing here - raise an exception?
+        #     self.logger.error(conversation + " could not parse response:\n"+ response["text"])
+        #     return "some failure"
+        return response["text"]
 
 if __name__ == "__main__":
     """to execute, from the repo root, use python -m modules.yaml_generator"""