prnvtripathi14 commited on
Commit
315b0b5
·
verified ·
1 Parent(s): 2c9c8ee

update app.py to have more models

Browse files
Files changed (1) hide show
  1. app.py +85 -34
app.py CHANGED
@@ -1,6 +1,6 @@
1
  import logging
2
  import json
3
- from transformers import pipeline
4
  import gradio as gr
5
 
6
  # Configure logging
@@ -8,19 +8,46 @@ logging.basicConfig(level=logging.DEBUG,
8
  format='%(asctime)s - %(levelname)s - %(message)s')
9
  logger = logging.getLogger(__name__)
10
 
11
- # Load the model with explicit logging
12
- try:
13
- model_name = "google/flan-t5-base"
14
- logger.info(f"Attempting to load model: {model_name}")
15
- generator = pipeline("text2text-generation", model=model_name)
16
- logger.info("Model loaded successfully")
17
- except Exception as model_load_error:
18
- logger.error(f"Failed to load model: {model_load_error}")
19
- generator = None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20
 
21
  def generate_test_cases(method, url, headers, payload=""):
22
  try:
23
- # Log input parameters
24
  logger.info(f"Generating test cases for:")
25
  logger.info(f"Method: {method}")
26
  logger.info(f"URL: {url}")
@@ -29,48 +56,72 @@ def generate_test_cases(method, url, headers, payload=""):
29
 
30
  # Validate inputs
31
  if not method or not url:
32
- logger.warning("Method or URL is missing")
33
  return "Error: Method and URL are required"
34
 
35
- # Attempt to parse headers and payload as JSON
36
  try:
37
  headers_dict = json.loads(headers) if headers else {}
38
  payload_dict = json.loads(payload) if payload else {}
39
  except json.JSONDecodeError as json_error:
40
- logger.error(f"JSON parsing error: {json_error}")
41
  return f"JSON Parsing Error: {json_error}"
42
 
43
- # Prepare the prompt
44
  prompt = f"""
45
- Generate comprehensive API test cases:
46
- Method: {method}
47
- URL: {url}
48
- Headers: {headers_dict}
49
- Payload: {payload_dict}
50
- Include:
51
- - Happy path scenarios
52
- - Negative test cases
53
- - Performance test considerations
54
- - Boundary condition tests
55
- - Error handling scenarios
56
- Provide detailed test case descriptions with expected inputs and anticipated outcomes.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
57
  """
58
 
59
  # Check if generator is available
60
  if generator is None:
61
- logger.error("Model generator is not initialized")
62
- return "Error: Model generator not initialized"
63
 
64
  # Generate test cases
65
- logger.info("Attempting to generate test cases")
66
  try:
67
- response = generator(prompt, max_length=500)
68
  generated_text = response[0]['generated_text']
69
 
70
  logger.info("Test cases generated successfully")
71
  logger.debug(f"Generated Text: {generated_text}")
72
 
73
  return generated_text
 
74
  except Exception as generation_error:
75
  logger.error(f"Test case generation error: {generation_error}")
76
  return f"Error generating test cases: {generation_error}"
@@ -79,7 +130,7 @@ def generate_test_cases(method, url, headers, payload=""):
79
  logger.error(f"Unexpected error: {overall_error}")
80
  return f"Unexpected error: {overall_error}"
81
 
82
- # Create the Gradio interface
83
  iface = gr.Interface(
84
  fn=generate_test_cases,
85
  inputs=[
@@ -89,8 +140,8 @@ iface = gr.Interface(
89
  gr.Textbox(label="Payload (JSON format)", placeholder='e.g., {"key": "value"}'),
90
  ],
91
  outputs="text",
92
- title="API Test Case Generator",
93
- description="Generate structured test cases for APIs using Hugging Face models."
94
  )
95
 
96
  # Main execution
 
1
  import logging
2
  import json
3
+ from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
4
  import gradio as gr
5
 
6
  # Configure logging
 
8
  format='%(asctime)s - %(levelname)s - %(message)s')
9
  logger = logging.getLogger(__name__)
10
 
11
+ # Try multiple models in succession
12
+ MODELS_TO_TRY = [
13
+ "google/flan-t5-large", # More capable than base
14
+ "google/flan-t5-xl", # Even more capable
15
+ "facebook/bart-large-cnn", # Alternative model
16
+ "t5-large" # Fallback T5 model
17
+ ]
18
+
19
+ def load_model():
20
+ for model_name in MODELS_TO_TRY:
21
+ try:
22
+ logger.info(f"Attempting to load model: {model_name}")
23
+
24
+ # Load model and tokenizer
25
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
26
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
27
+
28
+ # Create pipeline with specific model and tokenizer
29
+ generator = pipeline(
30
+ "text2text-generation",
31
+ model=model,
32
+ tokenizer=tokenizer,
33
+ max_length=500,
34
+ num_return_sequences=1
35
+ )
36
+
37
+ logger.info(f"Successfully loaded model: {model_name}")
38
+ return generator
39
+ except Exception as model_load_error:
40
+ logger.error(f"Failed to load model {model_name}: {model_load_error}")
41
+
42
+ logger.error("Failed to load any model")
43
+ return None
44
+
45
+ # Load model at startup
46
+ generator = load_model()
47
 
48
  def generate_test_cases(method, url, headers, payload=""):
49
  try:
50
+ # Detailed logging
51
  logger.info(f"Generating test cases for:")
52
  logger.info(f"Method: {method}")
53
  logger.info(f"URL: {url}")
 
56
 
57
  # Validate inputs
58
  if not method or not url:
 
59
  return "Error: Method and URL are required"
60
 
61
+ # Safely parse JSON inputs
62
  try:
63
  headers_dict = json.loads(headers) if headers else {}
64
  payload_dict = json.loads(payload) if payload else {}
65
  except json.JSONDecodeError as json_error:
 
66
  return f"JSON Parsing Error: {json_error}"
67
 
68
+ # Comprehensive prompt for test case generation
69
  prompt = f"""
70
+ Generate detailed API test cases with the following specifications:
71
+
72
+ Test Case Scenario: API Endpoint Testing
73
+
74
+ HTTP Method: {method}
75
+ API Endpoint: {url}
76
+ Request Headers: {json.dumps(headers_dict)}
77
+ Request Payload: {json.dumps(payload_dict)}
78
+
79
+ Test Case Requirements:
80
+ 1. Happy Path Scenarios:
81
+ - Successful request with valid inputs
82
+ - Verify correct response status code
83
+ - Validate response structure and content
84
+
85
+ 2. Negative Test Scenarios:
86
+ - Invalid authentication
87
+ - Malformed request payload
88
+ - Missing required headers
89
+ - Out-of-range parameter values
90
+
91
+ 3. Edge Case Considerations:
92
+ - Maximum/minimum input limits
93
+ - Special character handling
94
+ - Unicode and internationalization testing
95
+
96
+ 4. Performance and Security Checks:
97
+ - Response time validation
98
+ - Payload size limits
99
+ - Basic security vulnerability checks
100
+
101
+ Output Format:
102
+ For each test case, provide:
103
+ - Test Case ID
104
+ - Description
105
+ - Preconditions
106
+ - Input Data
107
+ - Expected Result
108
+ - Actual Result Verification Steps
109
  """
110
 
111
  # Check if generator is available
112
  if generator is None:
113
+ return "Error: No suitable model available for test case generation"
 
114
 
115
  # Generate test cases
 
116
  try:
117
+ response = generator(prompt)
118
  generated_text = response[0]['generated_text']
119
 
120
  logger.info("Test cases generated successfully")
121
  logger.debug(f"Generated Text: {generated_text}")
122
 
123
  return generated_text
124
+
125
  except Exception as generation_error:
126
  logger.error(f"Test case generation error: {generation_error}")
127
  return f"Error generating test cases: {generation_error}"
 
130
  logger.error(f"Unexpected error: {overall_error}")
131
  return f"Unexpected error: {overall_error}"
132
 
133
+ # Gradio Interface
134
  iface = gr.Interface(
135
  fn=generate_test_cases,
136
  inputs=[
 
140
  gr.Textbox(label="Payload (JSON format)", placeholder='e.g., {"key": "value"}'),
141
  ],
142
  outputs="text",
143
+ title="Comprehensive API Test Case Generator",
144
+ description="Advanced test case generation using AI-powered language models"
145
  )
146
 
147
  # Main execution