Update app.py

app.py

@@ -18,14 +18,24 @@ def extract_text_from_block(pdf_path, block_pages):
     return text
 
 def generate_test_cases(text, prompt_template):
-    prompt = prompt_template.format(text=text)
+    # Define a termination phrase that signals completion
+    termination_phrase = "### END OF TEST CASES ###"
+    
+    # Append termination instruction to the prompt template
+    full_prompt_template = prompt_template + (
+        "\n\nPlease ensure that your response ends with '" 
+        + termination_phrase 
+        + "' to indicate that no further test cases are needed."
+    )
+    
+    prompt = full_prompt_template.format(text=text)
     messages = [
         {"role": "system", "content": "You are a helpful assistant that generates test cases based on given text."},
         {"role": "user", "content": prompt}
     ]
     
     full_content = ""
-    max_iterations = 10  
+    max_iterations = 10 
     iteration = 0
 
     while iteration < max_iterations:
@@ -46,6 +56,12 @@ def generate_test_cases(text, prompt_template):
         part = response.choices[0].message['content'].strip()
         full_content += "\n" + part
         
+        # Check if the termination phrase is in the output
+        if termination_phrase in part:
+            # Optionally remove the termination phrase before returning
+            full_content = full_content.replace(termination_phrase, "")
+            break
+
         finish_reason = response.choices[0].get("finish_reason", None)
         # If the response appears truncated, ask the model to continue
         if finish_reason == "length" or part.endswith("..."):
@@ -61,6 +77,7 @@ def generate_test_cases(text, prompt_template):
 
 
 
+
 def process_pdf(pdf_file, selected_block, prompt_template):
     block_config = {
         "5.4.1 OSPI-xSPI-QSPI-SPI Boot": (482, 488),