update app.py to have more models

app.py

@@ -1,6 +1,6 @@
 import logging
 import json
-from transformers import pipeline
+from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
 import gradio as gr
 
 # Configure logging
@@ -8,19 +8,46 @@ logging.basicConfig(level=logging.DEBUG,
                     format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Load the model with explicit logging
-try:
-    model_name = "google/flan-t5-base"
-    logger.info(f"Attempting to load model: {model_name}")
-    generator = pipeline("text2text-generation", model=model_name)
-    logger.info("Model loaded successfully")
-except Exception as model_load_error:
-    logger.error(f"Failed to load model: {model_load_error}")
-    generator = None
+# Try multiple models in succession
+MODELS_TO_TRY = [
+    "google/flan-t5-large",  # More capable than base
+    "google/flan-t5-xl",     # Even more capable
+    "facebook/bart-large-cnn",  # Alternative model
+    "t5-large"               # Fallback T5 model
+]
+
+def load_model():
+    for model_name in MODELS_TO_TRY:
+        try:
+            logger.info(f"Attempting to load model: {model_name}")
+            
+            # Load model and tokenizer
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            
+            # Create pipeline with specific model and tokenizer
+            generator = pipeline(
+                "text2text-generation", 
+                model=model, 
+                tokenizer=tokenizer,
+                max_length=500,
+                num_return_sequences=1
+            )
+            
+            logger.info(f"Successfully loaded model: {model_name}")
+            return generator
+        except Exception as model_load_error:
+            logger.error(f"Failed to load model {model_name}: {model_load_error}")
+    
+    logger.error("Failed to load any model")
+    return None
+
+# Load model at startup
+generator = load_model()
 
 def generate_test_cases(method, url, headers, payload=""):
     try:
-        # Log input parameters
+        # Detailed logging
         logger.info(f"Generating test cases for:")
         logger.info(f"Method: {method}")
         logger.info(f"URL: {url}")
@@ -29,48 +56,72 @@ def generate_test_cases(method, url, headers, payload=""):
 
         # Validate inputs
         if not method or not url:
-            logger.warning("Method or URL is missing")
             return "Error: Method and URL are required"
 
-        # Attempt to parse headers and payload as JSON
+        # Safely parse JSON inputs
         try:
             headers_dict = json.loads(headers) if headers else {}
             payload_dict = json.loads(payload) if payload else {}
         except json.JSONDecodeError as json_error:
-            logger.error(f"JSON parsing error: {json_error}")
             return f"JSON Parsing Error: {json_error}"
 
-        # Prepare the prompt
+        # Comprehensive prompt for test case generation
         prompt = f"""
-        Generate comprehensive API test cases:
-        Method: {method}
-        URL: {url}
-        Headers: {headers_dict}
-        Payload: {payload_dict}
-        Include:
-        - Happy path scenarios
-        - Negative test cases
-        - Performance test considerations
-        - Boundary condition tests
-        - Error handling scenarios
-        Provide detailed test case descriptions with expected inputs and anticipated outcomes.
+        Generate detailed API test cases with the following specifications:
+        
+        Test Case Scenario: API Endpoint Testing
+        
+        HTTP Method: {method}
+        API Endpoint: {url}
+        Request Headers: {json.dumps(headers_dict)}
+        Request Payload: {json.dumps(payload_dict)}
+        
+        Test Case Requirements:
+        1. Happy Path Scenarios:
+           - Successful request with valid inputs
+           - Verify correct response status code
+           - Validate response structure and content
+        
+        2. Negative Test Scenarios:
+           - Invalid authentication
+           - Malformed request payload
+           - Missing required headers
+           - Out-of-range parameter values
+        
+        3. Edge Case Considerations:
+           - Maximum/minimum input limits
+           - Special character handling
+           - Unicode and internationalization testing
+        
+        4. Performance and Security Checks:
+           - Response time validation
+           - Payload size limits
+           - Basic security vulnerability checks
+        
+        Output Format:
+        For each test case, provide:
+        - Test Case ID
+        - Description
+        - Preconditions
+        - Input Data
+        - Expected Result
+        - Actual Result Verification Steps
         """
 
         # Check if generator is available
         if generator is None:
-            logger.error("Model generator is not initialized")
-            return "Error: Model generator not initialized"
+            return "Error: No suitable model available for test case generation"
 
         # Generate test cases
-        logger.info("Attempting to generate test cases")
         try:
-            response = generator(prompt, max_length=500)
+            response = generator(prompt)
             generated_text = response[0]['generated_text']
             
             logger.info("Test cases generated successfully")
             logger.debug(f"Generated Text: {generated_text}")
             
             return generated_text
+        
         except Exception as generation_error:
             logger.error(f"Test case generation error: {generation_error}")
             return f"Error generating test cases: {generation_error}"
@@ -79,7 +130,7 @@ def generate_test_cases(method, url, headers, payload=""):
         logger.error(f"Unexpected error: {overall_error}")
         return f"Unexpected error: {overall_error}"
 
-# Create the Gradio interface
+# Gradio Interface
 iface = gr.Interface(
     fn=generate_test_cases,
     inputs=[
@@ -89,8 +140,8 @@ iface = gr.Interface(
         gr.Textbox(label="Payload (JSON format)", placeholder='e.g., {"key": "value"}'),
     ],
     outputs="text",
-    title="API Test Case Generator",
-    description="Generate structured test cases for APIs using Hugging Face models."
+    title="Comprehensive API Test Case Generator",
+    description="Advanced test case generation using AI-powered language models"
 )
 
 # Main execution