Create backend.py

backend.py

@@ -0,0 +1,66 @@
+import csv
+import io
+from transformers import pipeline
+from bs4 import BeautifulSoup
+
+# Initialize the model (e.g., FLAN-T5)
+model = pipeline(
+    "text-generation",
+    model="google/flan-t5-xxl",  # Replace with your chosen model
+    torch_dtype=torch.float16  # Reduce memory usage
+)
+
+def clean_test_case_output(text):
+    text = text.strip()
+    soup = BeautifulSoup(text, 'html.parser')
+    return soup.get_text(separator="\n").strip()
+
+def generate_testcases(user_story):
+    few_shot_examples = """
+    "Generate as many test cases as possible (minimum 6)"
+    "Understand the user story thoroughly"
+    "If it's a portal-related user story, test in ODAC Portal; else, test in Tech360 iOS App"
+    """
+    
+    prompt = f"{few_shot_examples}\nUser Story: {user_story}"
+    
+    try:
+        result = model(prompt, max_length=4096)[0]["generated_text"]
+        cleaned_text = clean_test_case_output(result)
+        
+        try:
+            test_cases = json.loads(cleaned_text)
+            if isinstance(test_cases, list):
+                return test_cases
+            else:
+                return [{"test_case": cleaned_text}]
+        except json.JSONDecodeError:
+            return [{"test_case": cleaned_text}]
+    
+    except Exception as e:
+        return [{"test_case": f"Error: {str(e)}"}]
+
+def export_test_cases(test_cases, format='csv'):
+    if not test_cases:
+        return "No test cases to export."
+    
+    if format == 'csv':
+        output = io.StringIO()
+        writer = csv.writer(output)
+        writer.writerow(["Test Case"])
+        for case in test_cases:
+            writer.writerow([case.get("test_case", "N/A")])
+        return output.getvalue()
+    else:
+        raise ValueError("Unsupported format. Only CSV is supported.")
+
+def save_test_cases_as_file(test_cases, format='csv'):
+    if format == 'csv':
+        with open('test_cases.csv', 'w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow(["Test Case"])
+            for case in test_cases:
+                writer.writerow([case.get("test_case", "N/A")])
+        return "test_cases.csv saved"
+    else:
+        return f"Unsupported format: {format}"