Update app.py

app.py

@@ -1,76 +1,43 @@
-import torch
-from transformers import pipeline, set_seed
-import time
+import streamlit as st
+import backend  # Importing backend functions
 
-# Set random seed for reproducibility
-set_seed(42)
+# Title
+st.title("Test Case Generator - QA User Story")
 
-# Initialize model pipeline with better error handling
-print("Initializing distilgpt2 model pipeline...")
+# Input field for the user story
+user_story = st.text_area("Enter the QA User Story", "")
 
-def initialize_model():
-    try:
-        # Set up the model pipeline with distilgpt2
-        model = pipeline(
-            "text-generation",
-            model="distilgpt2",  # Load the distilgpt2 model
-            tokenizer="distilgpt2",  # Ensure tokenizer is loaded as well
-            torch_dtype=torch.float32  # Ensures compatibility with CPU
-        )
-        print("Model pipeline initialized successfully.")
-        return model
-    except Exception as e:
-        print(f"Error during model initialization: {e}")
-        return None
-
-# Retry logic to handle initialization issues
-model = None
-attempts = 0
-max_attempts = 3
-while model is None and attempts < max_attempts:
-    print(f"Attempt {attempts + 1}/{max_attempts} to initialize the model...")
-    model = initialize_model()
-    attempts += 1
-    if model is None:
-        print(f"Retrying after a delay...")
-        time.sleep(10)  # Wait 10 seconds before retrying
-
-if model is None:
-    print("Failed to initialize the model after several attempts. Exiting.")
-    exit(1)  # Exit if model is not initialized successfully
-
-# Generate test cases function
-def generate_test_cases(requirement):
-    prompt = f"Generate detailed test cases in JSON format for the following software requirement: '{requirement}'. Please return only the JSON array with test cases."
-    
-    try:
-        print("Generating test cases...")
-        result = model(prompt, max_length=500, num_return_sequences=1)[0]["generated_text"]
-        print("Test cases generated successfully.")
-        
-        # Clean up output to only return JSON array
-        clean_result = result.strip().split("\n")[0]
-        if clean_result.startswith('[') and clean_result.endswith(']'):
-            return clean_result
-        else:
-            return f"Error: Unexpected output format.\n{result.strip()}"
-    except Exception as e:
-        print(f"Error during generation: {e}")
-        return None
-
-# Example usage
-if __name__ == "__main__":
-    requirements = [
-        "User login functionality with email and password",
-        "Search functionality on an e-commerce website",
-        "Order placement process for an online store",
-        "Admin dashboard to manage users and content"
-    ]
-    
-    for req in requirements:
-        print(f"\nTest cases for requirement: {req}")
-        test_cases = generate_test_cases(req)
+# Button to generate test cases
+if st.button("Generate Test Cases"):
+    if user_story:
+        with st.spinner("Generating test cases..."):
+            test_cases = backend.generate_testcases(user_story)
+        st.subheader("Generated Test Cases")
         if test_cases:
-            print(test_cases)
+            for case in test_cases:
+                st.write(f"- **Test Case**: {case.get('test_case', 'N/A')}")
+            st.session_state.test_cases = test_cases  # Store in session state
         else:
-            print("Failed to generate test cases.")
+            st.error("Failed to generate test cases. Try a clearer input.")
+    else:
+        st.error("Please enter a user story.")
+
+# Export test cases
+st.sidebar.title("Export Test Cases")
+if st.sidebar.button("Export Test Cases as CSV"):
+    if 'test_cases' in st.session_state:
+        test_cases = st.session_state.test_cases
+        csv_content = backend.export_test_cases(test_cases, 'csv')
+        if csv_content:
+            st.sidebar.download_button(
+                label="Download CSV",
+                data=csv_content,
+                file_name="test_cases.csv",
+                mime="text/csv"
+            )
+    else:
+        st.sidebar.error("No test cases available.")
+
+# Footer
+st.markdown("---")
+st.write("Built with ❤️ by QA Team")