fdaxcnjk

model/generate.py

@@ -5,7 +5,6 @@ import logging
 import psutil
 import re
 import gc
-from typing import List, Dict, Union, Tuple
 
 # Initialize logger
 logger = logging.getLogger(__name__)
@@ -22,76 +21,7 @@ MEMORY_OPTIMIZED_MODELS = [
 # Singleton state
 _generator_instance = None
 
-# Enhanced keyword mapping for more specific test cases
-KEYWORD_TEST_MAPPING = {
-    'login': [
-        {
-            "title": "Valid Credentials Login",
-            "steps": ["Navigate to login page", "Enter valid username", "Enter valid password", "Click login button"],
-            "expected": "User should be redirected to dashboard"
-        },
-        {
-            "title": "Invalid Password Login",
-            "steps": ["Navigate to login page", "Enter valid username", "Enter invalid password", "Click login button"],
-            "expected": "System should display 'Invalid credentials' error"
-        },
-        {
-            "title": "Empty Fields Login",
-            "steps": ["Navigate to login page", "Leave username empty", "Leave password empty", "Click login button"],
-            "expected": "System should display validation errors for both fields"
-        }
-    ],
-    'authentication': [
-        {
-            "title": "Session Timeout Test",
-            "steps": ["Login successfully", "Wait for session timeout period", "Attempt to access protected resource"],
-            "expected": "System should redirect to login page"
-        },
-        {
-            "title": "Concurrent Sessions Test",
-            "steps": ["Login from device A", "Login from device B with same credentials", "Attempt actions on both devices"],
-            "expected": "System should handle concurrent sessions appropriately"
-        }
-    ],
-    'database': [
-        {
-            "title": "Data Integrity Test",
-            "steps": ["Insert test data", "Retrieve same data", "Compare results"],
-            "expected": "Stored data should match retrieved data exactly"
-        },
-        {
-            "title": "Large Data Volume Test",
-            "steps": ["Insert 10,000 records", "Perform search operations", "Measure response times"],
-            "expected": "System should handle large data volumes within acceptable performance thresholds"
-        }
-    ],
-    'api': [
-        {
-            "title": "API Authentication Test",
-            "steps": ["Make API request without authentication", "Make API request with valid credentials", "Make API request with invalid credentials"],
-            "expected": "Only authenticated requests should succeed"
-        },
-        {
-            "title": "API Input Validation Test",
-            "steps": ["Send malformed input to API", "Send extreme values to API", "Send valid input to API"],
-            "expected": "API should properly validate all inputs"
-        }
-    ],
-    'default': [
-        {
-            "title": "Basic Functionality Smoke Test",
-            "steps": ["Access the system", "Perform core operation", "Verify results"],
-            "expected": "System should perform basic functions without errors"
-        },
-        {
-            "title": "Error Handling Test",
-            "steps": ["Force error condition", "Verify system response"],
-            "expected": "System should handle errors gracefully with appropriate messages"
-        }
-    ]
-}
-
-def get_optimal_model_for_memory() -> Union[str, None]:
+def get_optimal_model_for_memory():
     """Select the best model based on available memory."""
     available_memory = psutil.virtual_memory().available / (1024 * 1024)  # MB
     logger.info(f"Available memory: {available_memory:.1f}MB")
@@ -103,7 +33,7 @@ def get_optimal_model_for_memory() -> Union[str, None]:
     else:
         return "distilgpt2"
 
-def load_model_with_memory_optimization(model_name: str) -> Tuple[Union[AutoTokenizer, None], Union[AutoModelForCausalLM, None]]:
+def load_model_with_memory_optimization(model_name):
     """Load model with low memory settings."""
     try:
         logger.info(f"Loading {model_name} with memory optimizations...")
@@ -130,172 +60,116 @@ def load_model_with_memory_optimization(model_name: str) -> Tuple[Union[AutoToke
         logger.error(f"❌ Failed to load model {model_name}: {e}")
         return None, None
 
-def extract_keywords(text: str) -> List[str]:
-    """Extract relevant keywords from text for test case generation."""
+def extract_keywords(text):
     common_keywords = [
         'login', 'authentication', 'user', 'password', 'database', 'data',
         'interface', 'api', 'function', 'feature', 'requirement', 'system',
-        'input', 'output', 'validation', 'error', 'security', 'performance',
-        'storage', 'retrieval', 'search', 'filter', 'export', 'import'
+        'input', 'output', 'validation', 'error', 'security', 'performance'
     ]
     words = re.findall(r'\b\w+\b', text.lower())
     return [word for word in words if word in common_keywords]
 
-def generate_template_based_test_cases(srs_text: str) -> List[Dict[str, Union[str, List[str]]]]:
-    """Generate test cases based on templates matching keywords in requirements."""
+def generate_template_based_test_cases(srs_text):
     keywords = extract_keywords(srs_text)
     test_cases = []
-    case_counter = 1
-
-    # Generate test cases for each matched keyword
-    for keyword, test_templates in KEYWORD_TEST_MAPPING.items():
-        if keyword in keywords:
-            for template in test_templates:
-                test_cases.append({
-                    "id": f"TC_{case_counter:03d}",
-                    "title": template["title"],
-                    "description": f"Test for {keyword} functionality: {template['title']}",
-                    "steps": template["steps"],
-                    "expected": template["expected"]
-                })
-                case_counter += 1
-
-    # Add default test cases if no specific ones were generated
-    if not test_cases:
-        for template in KEYWORD_TEST_MAPPING['default']:
-            test_cases.append({
-                "id": f"TC_{case_counter:03d}",
-                "title": template["title"],
-                "description": template["title"],
-                "steps": template["steps"],
-                "expected": template["expected"]
-            })
-            case_counter += 1
 
-    # Add boundary and edge case tests
-    if any(kw in keywords for kw in ['input', 'validation']):
+    if any(word in keywords for word in ['login', 'authentication', 'user', 'password']):
         test_cases.extend([
             {
-                "id": f"TC_{case_counter:03d}",
-                "title": "Boundary Value Analysis",
-                "description": "Test input boundaries and limits",
-                "steps": ["Enter minimum valid input", "Enter maximum valid input", "Enter just beyond boundaries"],
-                "expected": "System should accept valid inputs and reject invalid ones"
+                "id": "TC_001",
+                "title": "Valid Login Test",
+                "description": "Test login with valid credentials",
+                "steps": ["Enter valid username", "Enter valid password", "Click login"],
+                "expected": "User should be logged in successfully"
             },
             {
-                "id": f"TC_{case_counter+1:03d}",
-                "title": "Data Type Validation",
-                "description": "Test with invalid data types",
-                "steps": ["Enter text in numeric fields", "Enter numbers in text fields", "Enter special characters"],
-                "expected": "System should validate data types properly"
+                "id": "TC_002",
+                "title": "Invalid Login Test",
+                "description": "Test login with invalid credentials",
+                "steps": ["Enter invalid username", "Enter invalid password", "Click login"],
+                "expected": "Error message should be displayed"
             }
         ])
-        case_counter += 2
 
-    # Add security tests if relevant keywords exist
-    if any(kw in keywords for kw in ['security', 'authentication']):
+    if any(word in keywords for word in ['database', 'data', 'store', 'save']):
         test_cases.append({
-            "id": f"TC_{case_counter:03d}",
-            "title": "SQL Injection Test",
-            "description": "Test for SQL injection vulnerability",
-            "steps": ["Enter SQL injection string in input fields", "Submit form"],
-            "expected": "System should sanitize inputs and prevent SQL injection"
+            "id": "TC_003",
+            "title": "Data Storage Test",
+            "description": "Test data storage functionality",
+            "steps": ["Enter data", "Save data", "Verify storage"],
+            "expected": "Data should be stored correctly"
         })
-        case_counter += 1
+
+    if not test_cases:
+        test_cases = [
+            {
+                "id": "TC_001",
+                "title": "Basic Functionality Test",
+                "description": "Test basic system functionality",
+                "steps": ["Access the system", "Perform basic operations", "Verify results"],
+                "expected": "System should work as expected"
+            }
+        ]
 
     return test_cases
 
-def parse_generated_test_cases(generated_text: str) -> List[Dict[str, Union[str, List[str]]]]:
-    """Parse AI-generated text into structured test cases with enhanced parsing."""
-    lines = [line.strip() for line in generated_text.split('\n') if line.strip()]
+def parse_generated_test_cases(generated_text):
+    lines = generated_text.split('\n')
     test_cases = []
     current_case = {}
     case_counter = 1
-    step_pattern = re.compile(r'^\d+\.|step\s?\d+:|steps?:', re.IGNORECASE)
-    expected_pattern = re.compile(r'expected:|result:', re.IGNORECASE)
 
     for line in lines:
-        # Detect test case title
-        if re.match(r'^(test case|tc|\d+)\.?\s', line, re.IGNORECASE):
+        line = line.strip()
+        if line.startswith(('1.', '2.', '3.', 'TC', 'Test')):
             if current_case:
                 test_cases.append(current_case)
-                case_counter += 1
             current_case = {
                 "id": f"TC_{case_counter:03d}",
                 "title": line,
                 "description": line,
-                "steps": [],
-                "expected": "Verify expected behavior"
+                "steps": ["Execute the test"],
+                "expected": "Test should pass"
             }
-        # Detect steps
-        elif step_pattern.match(line):
-            step = re.sub(step_pattern, '', line).strip()
-            if step:
-                if 'steps' not in current_case:
-                    current_case['steps'] = []
-                current_case['steps'].append(step)
-        # Detect expected results
-        elif expected_pattern.match(line):
-            expected = re.sub(expected_pattern, '', line).strip()
-            if expected:
-                current_case['expected'] = expected
+            case_counter += 1
 
     if current_case:
-        # Ensure at least one step exists
-        if not current_case.get('steps'):
-            current_case['steps'] = ["Execute the test according to requirements"]
         test_cases.append(current_case)
 
-    # Fallback if no test cases were parsed
     if not test_cases:
         return [{
             "id": "TC_001",
-            "title": "Comprehensive Functionality Test",
-            "description": "End-to-end test of system functionality",
-            "steps": [
-                "Review all requirements",
-                "Execute core functionality tests",
-                "Verify all expected outcomes"
-            ],
-            "expected": "System meets all specified requirements"
+            "title": "Generated Test Case",
+            "description": "Auto-generated test case based on requirements",
+            "steps": ["Review requirements", "Execute test", "Verify results"],
+            "expected": "Requirements should be met"
         }]
 
     return test_cases
 
-def generate_with_ai_model(srs_text: str, tokenizer: AutoTokenizer, model: AutoModelForCausalLM) -> List[Dict[str, Union[str, List[str]]]]:
-    """Generate test cases using AI model with enhanced prompt engineering."""
-    max_input_length = 512  # Increased from 200 to capture more context
+def generate_with_ai_model(srs_text, tokenizer, model):
+    max_input_length = 200
     if len(srs_text) > max_input_length:
         srs_text = srs_text[:max_input_length]
 
-    prompt = f"""Generate comprehensive test cases for these software requirements. For each test case, include:
-1. Clear title describing the test scenario
-2. Detailed steps to execute the test
-3. Expected results
-
-Requirements:
+    prompt = f"""Generate test cases for this software requirement:
 {srs_text}
 
 Test Cases:
-1. Functional Test - Verify basic functionality:
-   Steps: 1. Access the system
-          2. Perform core operation
-          3. Verify results
-   Expected: System performs as specified in requirements
-2."""
+1."""
 
     try:
         inputs = tokenizer.encode(
             prompt,
             return_tensors="pt",
-            max_length=512,
+            max_length=150,
             truncation=True
         )
 
         with torch.no_grad():
             outputs = model.generate(
                 inputs,
-                max_new_tokens=300,  # Increased from 100 for more detailed output
+                max_new_tokens=100,
                 num_return_sequences=1,
                 temperature=0.7,
                 do_sample=True,
@@ -312,8 +186,7 @@ Test Cases:
         logger.error(f"❌ AI generation failed: {e}")
         raise
 
-def generate_with_fallback(srs_text: str) -> Tuple[List[Dict[str, Union[str, List[str]]]], str, str, str]:
-    """Generate test cases with AI or fallback to templates with enhanced logic."""
+def generate_with_fallback(srs_text):
     model_name = get_optimal_model_for_memory()
 
     if model_name:
@@ -330,12 +203,11 @@ def generate_with_fallback(srs_text: str) -> Tuple[List[Dict[str, Union[str, Lis
     test_cases = generate_template_based_test_cases(srs_text)
     return test_cases, "Template-Based Generator", "rule-based", "Low memory - fallback to rule-based generation"
 
-def generate_test_cases(srs_text: str) -> List[Dict[str, Union[str, List[str]]]]:
-    """Generate test cases from requirements text (primary interface)."""
+# ✅ Function exposed to app.py
+def generate_test_cases(srs_text):
     return generate_with_fallback(srs_text)[0]
 
 def get_generator():
-    """Get singleton generator instance with memory monitoring."""
     global _generator_instance
     if _generator_instance is None:
         class Generator:
@@ -346,7 +218,7 @@ def get_generator():
                 if self.model_name:
                     self.tokenizer, self.model = load_model_with_memory_optimization(self.model_name)
 
-            def get_model_info(self) -> Dict[str, str]:
+            def get_model_info(self):
                 mem = psutil.Process().memory_info().rss / 1024 / 1024
                 return {
                     "model_name": self.model_name if self.model_name else "Template-Based Generator",
@@ -360,15 +232,14 @@ def get_generator():
     return _generator_instance
 
 def monitor_memory():
-    """Monitor and log memory usage with automatic cleanup."""
     mem = psutil.Process().memory_info().rss / 1024 / 1024
     logger.info(f"Memory usage: {mem:.1f}MB")
     if mem > 450:
         gc.collect()
         logger.info("Memory cleanup triggered")
 
-def generate_test_cases_and_info(input_text: str) -> Dict[str, Union[str, List[Dict[str, Union[str, List[str]]]]]]:
-    """Generate test cases with detailed metadata about generation method."""
+# ✅ NEW FUNCTION for enhanced output: test cases + model info + reason
+def generate_test_cases_and_info(input_text):
     test_cases, model_name, algorithm_used, reason = generate_with_fallback(input_text)
     return {
         "model": model_name,
@@ -377,34 +248,15 @@ def generate_test_cases_and_info(input_text: str) -> Dict[str, Union[str, List[D
         "test_cases": test_cases
     }
 
-def get_algorithm_reason(model_name: Union[str, None]) -> str:
-    """Provide detailed explanation for algorithm selection."""
-    reasons = {
-        "microsoft/DialoGPT-small": "Selected due to low memory availability; DialoGPT-small provides conversational understanding in limited memory environments while maintaining reasonable generation quality.",
-        "distilgpt2": "Chosen for optimal balance between performance and memory usage. DistilGPT2 offers 82% of GPT-2's performance with half the memory footprint, ideal for resource-constrained environments.",
-        "gpt2": "Selected when sufficient memory is available for higher quality generation. GPT2 provides more coherent and context-aware outputs than its distilled versions.",
-        None: "Insufficient memory for model loading. Comprehensive template-based generation activated with 25+ predefined test scenarios covering common software testing needs."
-    }
-    return reasons.get(model_name, "Model selected based on best tradeoff between available resources and generation capabilities.")
-
-if __name__ == "__main__":
-    # Example usage
-    sample_requirements = """
-    The system shall provide user authentication via username and password.
-    All user data must be stored securely in the database.
-    The API should validate all inputs before processing.
-    """
-    
-    print("Generating test cases...")
-    result = generate_test_cases_and_info(sample_requirements)
-    print(f"\nModel used: {result['model']}")
-    print(f"Algorithm: {result['algorithm']}")
-    print(f"Reason: {result['reason']}\n")
-    
-    for tc in result["test_cases"]:
-        print(f"Test Case {tc['id']}: {tc['title']}")
-        print(f"Description: {tc['description']}")
-        print("Steps:")
-        for i, step in enumerate(tc['steps'], 1):
-            print(f"  {i}. {step}")
-        print(f"Expected: {tc['expected']}\n")
+# ✅ Explain why each algorithm is selected
+def get_algorithm_reason(model_name):
+    if model_name == "microsoft/DialoGPT-small":
+        return "Selected due to low memory availability; DialoGPT-small provides conversational understanding in limited memory environments."
+    elif model_name == "distilgpt2":
+        return "Selected for its balance between performance and low memory usage. Ideal for small environments needing causal language modeling."
+    elif model_name == "gpt2":
+        return "Chosen for general-purpose text generation with moderate memory headroom."
+    elif model_name is None:
+        return "No model used due to insufficient memory. Rule-based template generation chosen instead."
+    else:
+        return "Model selected based on best tradeoff between memory usage and language generation capability."