Update app.py

app.py

@@ -3,7 +3,7 @@ from transformers import pipeline
 import json
 import signal
 
-# Load a small, instruction-following model
+# Load a small, free, instruction-following model
 generator = pipeline("text2text-generation", model="google/flan-t5-large")
 
 # Timeout handling
@@ -19,23 +19,37 @@ def generate_test_cases(user_story):
     try:
         signal.alarm(180)  # Set a 3-minute timeout
 
-        # Improved structured prompt
+        # Structured prompt for better response
         prompt = (
-            f"Given the user story below, generate software test cases in JSON format.\n"
+            f"Generate structured test cases for the following user story:\n"
             f"User Story: {user_story}\n"
-            f"Format:\n"
-            f'{{"test_cases": [{{"id": 1, "title": "Test case title", "steps": ["Step 1", "Step 2"], "expected_result": "Expected outcome"}}]}}\n'
-            f"Ensure the output follows the JSON format strictly."
+            f"Provide output in a clear, structured way with a test case title, steps, and expected result."
         )
 
         output = generator(prompt, max_length=512, do_sample=False)[0]["generated_text"]
 
-        # Attempt to parse JSON
-        try:
-            json_output = json.loads(output)
-            formatted_output = json.dumps(json_output, indent=4)
-        except json.JSONDecodeError:
-            formatted_output = "Error: Model did not return valid JSON. Try re-running."
+        # Simple manual post-processing to force JSON format
+        test_cases = []
+        cases = output.split("\n\n")  # Split into test cases
+
+        for i, case in enumerate(cases, start=1):
+            lines = case.split("\n")
+            if len(lines) >= 3:
+                title = lines[0].strip()
+                steps = [line.strip() for line in lines[1:-1]]
+                expected_result = lines[-1].strip()
+
+                test_cases.append({
+                    "id": i,
+                    "title": title,
+                    "steps": steps,
+                    "expected_result": expected_result
+                })
+
+        if not test_cases:
+            return "Error: Model did not return structured test cases. Try again."
+
+        formatted_output = json.dumps({"test_cases": test_cases}, indent=4)
 
         signal.alarm(0)  # Disable timeout if successful
         return formatted_output