Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import logging | |
| import gc | |
| import psutil | |
| from functools import wraps | |
| import time | |
| import threading | |
| import json | |
| from model.generate import generate_test_cases, get_generator, monitor_memory | |
| # Configure logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' | |
| ) | |
| logger = logging.getLogger(__name__) | |
| # Thread-safe initialization | |
| _init_lock = threading.Lock() | |
| _initialized = False | |
| def init_model(): | |
| """Initialize model on startup""" | |
| try: | |
| # Skip AI model loading in low memory environments | |
| memory_mb = psutil.Process().memory_info().rss / 1024 / 1024 | |
| if memory_mb > 200 or os.environ.get('HUGGINGFACE_SPACE'): | |
| logger.info("⚠️ Skipping AI model loading due to memory constraints") | |
| logger.info("🔧 Using template-based generation mode") | |
| return True | |
| logger.info("🚀 Initializing AI model...") | |
| generator = get_generator() | |
| model_info = generator.get_model_info() | |
| logger.info(f"✅ Model initialized: {model_info['model_name']} | Memory: {model_info['memory_usage']}") | |
| return True | |
| except Exception as e: | |
| logger.error(f"❌ Model initialization failed: {e}") | |
| logger.info("🔧 Falling back to template-based generation") | |
| return False | |
| def check_health(): | |
| """Check system health""" | |
| try: | |
| memory_mb = psutil.Process().memory_info().rss / 1024 / 1024 | |
| return { | |
| "status": "healthy" if memory_mb < 450 else "warning", | |
| "memory_usage": f"{memory_mb:.1f}MB", | |
| "memory_limit": "512MB" | |
| } | |
| except Exception: | |
| return {"status": "unknown", "memory_usage": "unavailable"} | |
| def smart_memory_monitor(func): | |
| """Enhanced memory monitoring with automatic cleanup""" | |
| def wrapper(*args, **kwargs): | |
| start_time = time.time() | |
| try: | |
| initial_memory = psutil.Process().memory_info().rss / 1024 / 1024 | |
| logger.info(f"🔍 {func.__name__} started | Memory: {initial_memory:.1f}MB") | |
| if initial_memory > 400: | |
| logger.warning("⚠️ High memory detected, forcing cleanup...") | |
| gc.collect() | |
| result = func(*args, **kwargs) | |
| return result | |
| except Exception as e: | |
| logger.error(f"❌ Error in {func.__name__}: {str(e)}") | |
| return { | |
| "error": "Internal server error occurred", | |
| "message": "Please try again or contact support" | |
| } | |
| finally: | |
| final_memory = psutil.Process().memory_info().rss / 1024 / 1024 | |
| execution_time = time.time() - start_time | |
| logger.info(f"✅ {func.__name__} completed | Memory: {final_memory:.1f}MB | Time: {execution_time:.2f}s") | |
| if final_memory > 450: | |
| logger.warning("🧹 High memory usage, forcing aggressive cleanup...") | |
| gc.collect() | |
| post_cleanup_memory = psutil.Process().memory_info().rss / 1024 / 1024 | |
| logger.info(f"🧹 Post-cleanup memory: {post_cleanup_memory:.1f}MB") | |
| return wrapper | |
| def ensure_initialized(): | |
| """Ensure model is initialized (thread-safe)""" | |
| global _initialized | |
| if not _initialized: | |
| with _init_lock: | |
| if not _initialized: | |
| logger.info("🚀 Gradio app starting up on Hugging Face Spaces...") | |
| success = init_model() | |
| if success: | |
| logger.info("✅ Startup completed successfully") | |
| else: | |
| logger.warning("⚠️ Model initialization failed, using template mode") | |
| _initialized = True | |
| def read_uploaded_file(file_obj): | |
| """Read and extract text from uploaded file""" | |
| if file_obj is None: | |
| return "" | |
| try: | |
| file_path = file_obj.name | |
| file_extension = os.path.splitext(file_path)[1].lower() | |
| if file_extension in ['.txt', '.md']: | |
| with open(file_path, 'r', encoding='utf-8') as f: | |
| content = f.read() | |
| elif file_extension in ['.doc', '.docx']: | |
| try: | |
| import docx | |
| doc = docx.Document(file_path) | |
| content = '\n'.join([paragraph.text for paragraph in doc.paragraphs]) | |
| except ImportError: | |
| logger.warning("python-docx not available, trying to read as text") | |
| with open(file_path, 'r', encoding='utf-8', errors='ignore') as f: | |
| content = f.read() | |
| elif file_extension == '.pdf': | |
| try: | |
| import PyPDF2 | |
| with open(file_path, 'rb') as f: | |
| reader = PyPDF2.PdfReader(f) | |
| content = '' | |
| for page in reader.pages: | |
| content += page.extract_text() + '\n' | |
| except ImportError: | |
| logger.warning("PyPDF2 not available, cannot read PDF files") | |
| return "❌ PDF support requires PyPDF2. Please install it or use text/Word files." | |
| else: | |
| # Try to read as plain text | |
| with open(file_path, 'r', encoding='utf-8', errors='ignore') as f: | |
| content = f.read() | |
| logger.info(f"📄 File read successfully: {len(content)} characters") | |
| return content | |
| except Exception as e: | |
| logger.error(f"❌ Error reading file: {str(e)}") | |
| return f"❌ Error reading file: {str(e)}" | |
| def combine_inputs(prompt_text, uploaded_file): | |
| """Combine prompt text and uploaded file content""" | |
| file_content = "" | |
| if uploaded_file is not None: | |
| file_content = read_uploaded_file(uploaded_file) | |
| if file_content.startswith("❌"): | |
| return file_content # Return error message | |
| # Combine both inputs | |
| combined_text = "" | |
| if prompt_text and prompt_text.strip(): | |
| combined_text += "PROMPT:\n" + prompt_text.strip() + "\n\n" | |
| if file_content and not file_content.startswith("❌"): | |
| combined_text += "DOCUMENT CONTENT:\n" + file_content.strip() | |
| if not combined_text.strip(): | |
| return "❌ Please provide either text input or upload a document." | |
| return combined_text.strip() | |
| # Initialize on startup | |
| ensure_initialized() | |
| def generate_test_cases_api(prompt_text, uploaded_file): | |
| """Main API function for generating test cases with dual input support""" | |
| # Combine inputs | |
| combined_input = combine_inputs(prompt_text, uploaded_file) | |
| if combined_input.startswith("❌"): | |
| return { | |
| "error": combined_input, | |
| "test_cases": [], | |
| "count": 0 | |
| } | |
| if len(combined_input) > 8000: | |
| logger.warning(f"Input text truncated from {len(combined_input)} to 8000 characters") | |
| combined_input = combined_input[:8000] | |
| try: | |
| logger.info(f"🎯 Generating test cases for combined input ({len(combined_input)} chars)") | |
| test_cases = generate_test_cases(combined_input) | |
| if not test_cases or len(test_cases) == 0: | |
| logger.error("No test cases generated") | |
| return { | |
| "error": "Failed to generate test cases", | |
| "test_cases": [], | |
| "count": 0 | |
| } | |
| try: | |
| generator = get_generator() | |
| model_info = generator.get_model_info() | |
| model_used = model_info.get("model_name", "Unknown Model") | |
| generation_method = model_info.get("status", "unknown") | |
| except Exception: | |
| model_used = "Template-Based Generator" | |
| generation_method = "template_mode" | |
| if model_used == "Template-Based Generator": | |
| model_algorithm = "Enhanced Rule-based Template" | |
| model_reason = "Used enhanced rule-based generation with pattern matching and context analysis." | |
| elif "distilgpt2" in model_used: | |
| model_algorithm = "Transformer-based LM" | |
| model_reason = "Used DistilGPT2 for balanced performance and memory efficiency." | |
| elif "DialoGPT" in model_used: | |
| model_algorithm = "Transformer-based LM" | |
| model_reason = "Used DialoGPT-small as it fits within memory limits and handles conversational input well." | |
| else: | |
| model_algorithm = "Transformer-based LM" | |
| model_reason = "Used available Hugging Face causal LM due to sufficient resources." | |
| logger.info(f"✅ Successfully generated {len(test_cases)} test cases") | |
| return { | |
| "test_cases": test_cases, | |
| "count": len(test_cases), | |
| "model_used": model_used, | |
| "generation_method": generation_method, | |
| "model_algorithm": model_algorithm, | |
| "model_reason": model_reason, | |
| "input_source": "Combined (Prompt + Document)" if (prompt_text and uploaded_file) else | |
| "Document Upload" if uploaded_file else "Text Prompt" | |
| } | |
| except Exception as e: | |
| logger.error(f"❌ Test case generation failed: {str(e)}") | |
| return { | |
| "error": "Failed to generate test cases", | |
| "message": "Please try again with different input", | |
| "test_cases": [], | |
| "count": 0 | |
| } | |
| def format_test_cases_output(result): | |
| """Format the test cases for display""" | |
| if "error" in result: | |
| return f"❌ Error: {result['error']}", "" | |
| test_cases = result.get("test_cases", []) | |
| if not test_cases: | |
| return "No test cases generated", "" | |
| # Format test cases for display | |
| formatted_output = f"✅ Generated {result['count']} Test Cases\n\n" | |
| formatted_output += f"📥 Input Source: {result.get('input_source', 'Unknown')}\n" | |
| formatted_output += f"🤖 Model: {result['model_used']}\n" | |
| formatted_output += f"🔧 Algorithm: {result['model_algorithm']}\n" | |
| formatted_output += f"💡 Reason: {result['model_reason']}\n\n" | |
| formatted_output += "=" * 60 + "\n" | |
| formatted_output += "GENERATED TEST CASES\n" | |
| formatted_output += "=" * 60 + "\n\n" | |
| for i, tc in enumerate(test_cases, 1): | |
| formatted_output += f"🔹 Test Case {i}:\n" | |
| formatted_output += f" ID: {tc.get('id', f'TC_{i:03d}')}\n" | |
| formatted_output += f" Title: {tc.get('title', 'N/A')}\n" | |
| formatted_output += f" Priority: {tc.get('priority', 'Medium')}\n" | |
| formatted_output += f" Category: {tc.get('category', 'Functional')}\n" | |
| formatted_output += f" Description: {tc.get('description', 'N/A')}\n" | |
| # Pre-conditions | |
| preconditions = tc.get('preconditions', []) | |
| if preconditions: | |
| formatted_output += f" Pre-conditions:\n" | |
| for j, precond in enumerate(preconditions, 1): | |
| formatted_output += f" • {precond}\n" | |
| # Test steps | |
| steps = tc.get('steps', []) | |
| if isinstance(steps, list) and steps: | |
| formatted_output += f" Test Steps:\n" | |
| for j, step in enumerate(steps, 1): | |
| formatted_output += f" {j}. {step}\n" | |
| else: | |
| formatted_output += f" Test Steps: {steps if steps else 'N/A'}\n" | |
| formatted_output += f" Expected Result: {tc.get('expected', 'N/A')}\n" | |
| # Post-conditions | |
| postconditions = tc.get('postconditions', []) | |
| if postconditions: | |
| formatted_output += f" Post-conditions:\n" | |
| for postcond in postconditions: | |
| formatted_output += f" • {postcond}\n" | |
| formatted_output += f" Test Data: {tc.get('test_data', 'N/A')}\n" | |
| formatted_output += "\n" + "-" * 40 + "\n\n" | |
| # Return JSON for API access | |
| json_output = json.dumps(result, indent=2) | |
| return formatted_output, json_output | |
| def gradio_generate_test_cases(prompt_text, uploaded_file): | |
| """Gradio interface function""" | |
| result = generate_test_cases_api(prompt_text, uploaded_file) | |
| return format_test_cases_output(result) | |
| def get_system_status(): | |
| """Get system status information""" | |
| health_data = check_health() | |
| try: | |
| generator = get_generator() | |
| model_info = generator.get_model_info() | |
| except Exception: | |
| model_info = { | |
| "model_name": "Enhanced Template-Based Generator", | |
| "status": "template_mode", | |
| "optimization": "memory_safe" | |
| } | |
| status_info = f""" | |
| 🏥 SYSTEM STATUS | |
| ================ | |
| Status: {health_data["status"]} | |
| Memory Usage: {health_data["memory_usage"]} | |
| Memory Limit: 512MB | |
| 🤖 MODEL INFORMATION | |
| =================== | |
| Model Name: {model_info["model_name"]} | |
| Status: {model_info["status"]} | |
| Optimization: {model_info.get("optimization", "standard")} | |
| 🚀 APPLICATION INFO | |
| ================== | |
| Version: 2.0.0-enhanced-input | |
| Environment: Hugging Face Spaces | |
| Backend: Gradio | |
| Features: Text Input + File Upload | |
| Supported Files: .txt, .md, .doc, .docx, .pdf | |
| """ | |
| return status_info | |
| def get_model_info_detailed(): | |
| """Get detailed model information""" | |
| try: | |
| generator = get_generator() | |
| info = generator.get_model_info() | |
| health_data = check_health() | |
| detailed_info = f""" | |
| 🔧 DETAILED MODEL INFORMATION | |
| ============================ | |
| Model Name: {info.get('model_name', 'N/A')} | |
| Status: {info.get('status', 'N/A')} | |
| Memory Usage: {info.get('memory_usage', 'N/A')} | |
| Optimization Level: {info.get('optimization', 'N/A')} | |
| 📊 SYSTEM METRICS | |
| ================ | |
| System Status: {health_data['status']} | |
| Current Memory: {health_data['memory_usage']} | |
| Memory Limit: {health_data.get('memory_limit', 'N/A')} | |
| ⚙️ CONFIGURATION | |
| =============== | |
| Environment: {"Hugging Face Spaces" if os.environ.get('SPACE_ID') else "Local"} | |
| Backend: Gradio | |
| Threading: Enabled | |
| Memory Monitoring: Active | |
| Input Methods: Text + File Upload | |
| File Support: TXT, MD, DOC, DOCX, PDF | |
| """ | |
| return detailed_info | |
| except Exception as e: | |
| return f"❌ Error getting model info: {str(e)}" | |
| # Create Gradio interface | |
| with gr.Blocks(title="AI Test Case Generator - Enhanced", theme=gr.themes.Soft()) as app: | |
| gr.Markdown(""" | |
| # 🧪 AI Test Case Generator - Enhanced Edition | |
| Generate comprehensive test cases from Software Requirements Specification (SRS) documents using AI models. | |
| **New Features:** | |
| - 📝 **Dual Input Support**: Text prompt AND/OR document upload | |
| - 📄 **File Upload**: Support for .txt, .md, .doc, .docx, .pdf files | |
| - 🎯 **Enhanced Test Cases**: More detailed and comprehensive test case generation | |
| - 🔧 **Improved Templates**: Better rule-based fallback with pattern matching | |
| - 📊 **Better Formatting**: Enhanced output with priorities, categories, and conditions | |
| """) | |
| with gr.Tab("🧪 Generate Test Cases"): | |
| gr.Markdown("### Choose your input method: Enter text directly, upload a document, or use both!") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| # Text input | |
| srs_input = gr.Textbox( | |
| label="📝 Text Input (SRS, Requirements, or Prompt)", | |
| placeholder="Enter your requirements, user stories, or specific prompt here...\n\nExample:\n- The system shall provide user authentication with username and password\n- Users should be able to login, logout, and reset passwords\n- The system should validate input and display appropriate error messages\n- Performance requirement: Login should complete within 3 seconds", | |
| lines=8, | |
| max_lines=15 | |
| ) | |
| # File upload | |
| file_upload = gr.File( | |
| label="📄 Upload Document (Optional)", | |
| file_types=[".txt", ".md", ".doc", ".docx", ".pdf"], | |
| type="filepath" | |
| ) | |
| gr.Markdown(""" | |
| **💡 Tips:** | |
| - Use **text input** for quick requirements or specific prompts | |
| - Use **file upload** for complete SRS documents | |
| - Use **both** to combine a specific prompt with a detailed document | |
| - Supported formats: TXT, Markdown, Word (.doc/.docx), PDF | |
| """) | |
| generate_btn = gr.Button("🚀 Generate Test Cases", variant="primary", size="lg") | |
| with gr.Column(scale=3): | |
| output_display = gr.Textbox( | |
| label="📋 Generated Test Cases", | |
| lines=25, | |
| max_lines=35, | |
| interactive=False | |
| ) | |
| with gr.Row(): | |
| json_output = gr.Textbox( | |
| label="📄 JSON Output (for API use)", | |
| lines=12, | |
| max_lines=20, | |
| interactive=False | |
| ) | |
| with gr.Tab("📊 System Status"): | |
| with gr.Column(): | |
| status_display = gr.Textbox( | |
| label="🏥 System Health & Status", | |
| lines=18, | |
| interactive=False | |
| ) | |
| refresh_status_btn = gr.Button("🔄 Refresh Status", variant="secondary") | |
| with gr.Tab("🔧 Model Information"): | |
| with gr.Column(): | |
| model_info_display = gr.Textbox( | |
| label="🤖 Detailed Model Information", | |
| lines=22, | |
| interactive=False | |
| ) | |
| refresh_model_btn = gr.Button("🔄 Refresh Model Info", variant="secondary") | |
| with gr.Tab("📚 API Documentation"): | |
| gr.Markdown(""" | |
| ## 🔌 Enhanced API Endpoints | |
| This Gradio app supports both text input and file upload through API: | |
| ### Generate Test Cases (Text Only) | |
| **Endpoint:** `/api/predict` | |
| **Method:** POST | |
| **Body:** | |
| ```json | |
| { | |
| "data": ["Your SRS text here", null] | |
| } | |
| ``` | |
| ### Generate Test Cases (With File) | |
| **Endpoint:** `/api/predict` | |
| **Method:** POST (multipart/form-data) | |
| - Upload file and include text in the data array | |
| **Response Format:** | |
| ```json | |
| { | |
| "data": [ | |
| "Formatted test cases output", | |
| "JSON output with enhanced test cases" | |
| ] | |
| } | |
| ``` | |
| ### Enhanced Test Case Structure | |
| ```json | |
| { | |
| "test_cases": [ | |
| { | |
| "id": "TC_001", | |
| "title": "Test Case Title", | |
| "priority": "High/Medium/Low", | |
| "category": "Functional/Security/Performance/UI", | |
| "description": "Detailed test description", | |
| "preconditions": ["Pre-condition 1", "Pre-condition 2"], | |
| "steps": ["Step 1", "Step 2", "Step 3"], | |
| "expected": "Expected result", | |
| "postconditions": ["Post-condition 1"], | |
| "test_data": "Required test data" | |
| } | |
| ], | |
| "count": 5, | |
| "model_used": "distilgpt2", | |
| "model_algorithm": "Enhanced Rule-based Template", | |
| "model_reason": "Detailed selection reasoning...", | |
| "input_source": "Combined (Prompt + Document)" | |
| } | |
| ``` | |
| ### Example Usage (Python with File): | |
| ```python | |
| import requests | |
| # Text only | |
| response = requests.post( | |
| "YOUR_SPACE_URL/api/predict", | |
| json={"data": ["User login requirements...", None]} | |
| ) | |
| # With file upload (requires multipart handling) | |
| files = {'file': open('requirements.pdf', 'rb')} | |
| data = {'data': json.dumps(["Additional prompt", "file_placeholder"])} | |
| response = requests.post("YOUR_SPACE_URL/api/predict", files=files, data=data) | |
| ``` | |
| ## 📋 Supported File Formats | |
| - **Text Files**: .txt, .md | |
| - **Word Documents**: .doc, .docx (requires python-docx) | |
| - **PDF Files**: .pdf (requires PyPDF2) | |
| - **Fallback**: Any text-readable format | |
| ## 🎯 Enhanced Features | |
| - **Dual Input**: Combine text prompts with document uploads | |
| - **Better Test Cases**: Includes priorities, categories, pre/post-conditions | |
| - **Smart Parsing**: Automatically detects requirement types and generates appropriate tests | |
| - **Memory Optimized**: Handles large documents efficiently | |
| """) | |
| # Event handlers | |
| generate_btn.click( | |
| fn=gradio_generate_test_cases, | |
| inputs=[srs_input, file_upload], | |
| outputs=[output_display, json_output] | |
| ) | |
| refresh_status_btn.click( | |
| fn=get_system_status, | |
| outputs=[status_display] | |
| ) | |
| refresh_model_btn.click( | |
| fn=get_model_info_detailed, | |
| outputs=[model_info_display] | |
| ) | |
| # Load initial status | |
| app.load( | |
| fn=get_system_status, | |
| outputs=[status_display] | |
| ) | |
| app.load( | |
| fn=get_model_info_detailed, | |
| outputs=[model_info_display] | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| port = int(os.environ.get("PORT", 7860)) | |
| logger.info(f"🚀 Starting Enhanced Gradio app on port {port}") | |
| logger.info(f"🖥️ Environment: {'Hugging Face Spaces' if os.environ.get('SPACE_ID') else 'Local'}") | |
| logger.info("📝 Features: Text Input + File Upload Support") | |
| app.launch( | |
| server_name="0.0.0.0", | |
| server_port=port, | |
| share=False, | |
| show_error=True | |
| ) |