Spaces:
Sleeping
Sleeping
| import logging | |
| import json | |
| from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer | |
| import gradio as gr | |
| # Configure logging | |
| logging.basicConfig(level=logging.DEBUG, | |
| format='%(asctime)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger(__name__) | |
| # Models to try | |
| MODELS_TO_TRY = [ | |
| "google/flan-t5-xxl", # Powerful instruction-following model | |
| "bigscience/T0pp", # Optimized for zero-shot tasks | |
| "t5-large", # General-purpose text generation | |
| "google/flan-t5-large" # Lightweight instruction-tuned model | |
| ] | |
| def load_model(): | |
| """ | |
| Attempt to load a suitable model for text generation. | |
| """ | |
| 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 the text generation pipeline | |
| generator = pipeline( | |
| "text2text-generation", | |
| model=model, | |
| tokenizer=tokenizer, | |
| max_length=512, | |
| num_return_sequences=1 | |
| ) | |
| logger.info(f"Successfully loaded model: {model_name}") | |
| return generator | |
| except Exception as e: | |
| logger.error(f"Failed to load model {model_name}: {e}") | |
| logger.error("All model attempts failed. No model loaded.") | |
| return None | |
| # Load the generator at startup | |
| generator = load_model() | |
| def generate_test_cases(method, url, headers, payload=""): | |
| """ | |
| Generate detailed API test cases using a language model. | |
| """ | |
| 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 | |
| except Exception as e: | |
| logger.error(f"Error during test case generation: {e}") | |
| return f"Error: {e}" | |
| # Gradio Interface | |
| iface = gr.Interface( | |
| fn=generate_test_cases, | |
| 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"}'), | |
| ], | |
| outputs="text", | |
| title="API Test Case Generator", | |
| description="Generate detailed API test cases using AI models." | |
| ) | |
| # Main execution | |
| if __name__ == "__main__": | |
| try: | |
| logger.info("Launching Gradio interface...") | |
| iface.launch() | |
| except Exception as e: | |
| logger.error(f"Failed to launch Gradio interface: {e}") | |