Update app.py

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 from transformers import pipeline
+import json
 import signal
 
-# Load model
-generator = pipeline("text-generation", model="google/flan-t5-base")
+# Load a small, free, instruction-following model
+generator = pipeline("text2text-generation", model="google/flan-t5-base")
 
 # Function to handle timeouts
 class TimeoutException(Exception):
@@ -17,15 +18,30 @@ signal.signal(signal.SIGALRM, timeout_handler)
 def generate_test_cases(user_story):
     try:
         # Set a 3-minute timeout
-        signal.alarm(120)
-
-        prompt = f"Generate structured test cases in JSON format from this user story: {user_story}"
-        output = generator(prompt, max_length=300, do_sample=True)[0]["generated_text"]
+        signal.alarm(180)
+
+        # Force JSON output using a structured prompt
+        prompt = (
+            f"Convert the following user story into structured JSON test cases.\n"
+            f"User story: {user_story}\n"
+            f"Output format:\n"
+            f'{{"test_cases": [{{"id": 1, "title": "Test case title", "steps": ["Step 1", "Step 2"], "expected_result": "Expected outcome"}}]}}'
+        )
+
+        # Generate output
+        output = generator(prompt, max_length=512, do_sample=False)[0]["generated_text"]
+
+        # Try parsing as JSON to ensure proper formatting
+        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."
 
         # Disable the alarm if processing is successful
         signal.alarm(0)
 
-        return output
+        return formatted_output
 
     except TimeoutException:
         return "Processing timed out. Please try again with a simpler input."