update

app.py

@@ -1,110 +1,68 @@
 import logging
-import json
-from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import pipeline
 import gradio as gr
 
-# Configure logging
-logging.basicConfig(level=logging.DEBUG,
-                    format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
+# Set up logging
+logging.basicConfig(
+    filename="app.log",
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
 
-# Model to use
-MODEL_NAME = "google/flan-t5-base"
+# Load the generative AI model
+logging.info("Loading the Hugging Face model...")
+try:
+    model = pipeline("text2text-generation", model="google/flan-t5-large")  # Replace with your chosen model
+    logging.info("Model loaded successfully.")
+except Exception as e:
+    logging.error(f"Error loading the model: {e}")
+    raise
 
-def load_model():
-    """
-    Load the selected model and tokenizer using PyTorch.
-    """
+# Function to generate test cases
+def generate_test_cases(api_info):
+    logging.info(f"Generating test cases for API info: {api_info}")
     try:
-        logger.info(f"Loading model: {MODEL_NAME} with PyTorch backend")
-        
-        # Load the model and tokenizer
-        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, framework="pt")
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        
-        # Create the text generation pipeline
-        generator = pipeline(
-            "text2text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            framework="pt",  # Specify PyTorch framework
-            max_length=512,
-            num_return_sequences=1
+        prompt = (
+            f"Generate API test cases for the following API:\n\n{api_info}\n\n"
+            f"Test cases should include:\n- Happy path\n- Negative tests\n- Edge cases"
         )
-        logger.info(f"Successfully loaded model: {MODEL_NAME}")
-        return generator
+        result = model(prompt, max_length=512, num_return_sequences=1)
+        logging.info(f"Test cases generated successfully.")
+        return result[0]['generated_text']
     except Exception as e:
-        logger.error(f"Failed to load model {MODEL_NAME}: {e}")
-        return None
-
-# Load the generator at startup
-generator = load_model()
+        logging.error(f"Error generating test cases: {e}")
+        return "An error occurred while generating test cases."
 
-def generate_test_cases(method, url, headers, payload=""):
-    """
-    Generate detailed API test cases using a language model.
-    """
+# Process input and generate output
+def process_input(url, method, headers, payload):
     try:
-        # Input validation and logging
-        logger.info(f"Received inputs: Method={method}, URL={url}, Headers={headers}, Payload={payload}")
-        
-        if not method or not url:
-            return "Error: HTTP Method and API URL are required inputs."
-        
-        # Parse headers and payload as JSON
-        try:
-            headers_dict = json.loads(headers) if headers.strip() else {}
-            payload_dict = json.loads(payload) if payload.strip() else {}
-        except json.JSONDecodeError as e:
-            return f"JSON Parsing Error: {e}"
-
-        # Prompt for the model
-        prompt = f"""
-        Generate comprehensive API test cases for the following:
-        
-        HTTP Method: {method}
-        API URL: {url}
-        Headers: {json.dumps(headers_dict, indent=2)}
-        Payload: {json.dumps(payload_dict, indent=2)}
-        
-        Requirements:
-        - Include Happy Path, Negative, and Edge Cases.
-        - Provide validation steps and expected results.
-        """
-
-        # Ensure model is loaded
-        if generator is None:
-            return "Error: No model is available for test case generation."
-
-        # Generate test cases
-        response = generator(prompt, max_length=500, num_return_sequences=1)
-        generated_text = response[0]['generated_text']
-        
-        logger.info("Successfully generated test cases.")
-        return generated_text
-
+        logging.info("Received user input.")
+        api_info = f"URL: {url}\nMethod: {method}\nHeaders: {headers}\nPayload: {payload}"
+        logging.debug(f"Formatted API info: {api_info}")
+        test_cases = generate_test_cases(api_info)
+        return test_cases
     except Exception as e:
-        logger.error(f"Error during test case generation: {e}")
-        return f"Error: {e}"
+        logging.error(f"Error processing input: {e}")
+        return "An error occurred. Please check the input format and try again."
 
-# Gradio Interface
-iface = gr.Interface(
-    fn=generate_test_cases,
+# Define Gradio interface
+interface = gr.Interface(
+    fn=process_input,
     inputs=[
-        gr.Textbox(label="HTTP Method (GET, POST, etc.)", placeholder="e.g., GET"),
-        gr.Textbox(label="API URL", placeholder="e.g., https://api.example.com/endpoint"),
-        gr.Textbox(label="Headers (JSON format)", placeholder='e.g., {"Content-Type": "application/json"}'),
-        gr.Textbox(label="Payload (JSON format)", placeholder='e.g., {"key": "value"}'),
+        gr.Textbox(label="API URL"),
+        gr.Textbox(label="HTTP Method"),
+        gr.Textbox(label="Headers (JSON format)"),
+        gr.Textbox(label="Payload (JSON format)"),
     ],
     outputs="text",
-    title="API Test Case Generator",
-    description="Generate detailed API test cases using AI models."
+    title="API Test Case Generator"
 )
 
-# Main execution
+# Launch Gradio app
 if __name__ == "__main__":
     try:
-        logger.info("Launching Gradio interface...")
-        iface.launch()
+        logging.info("Starting the Gradio app...")
+        interface.launch()
+        logging.info("Gradio app launched successfully.")
     except Exception as e:
-        logger.error(f"Failed to launch Gradio interface: {e}")
+        logging.error(f"Error launching the Gradio app: {e}")