"""Natural language to shell script generation.""" from typing import List, Optional, Dict from dataclasses import dataclass from .models import ScriptTemplate from .database import Database @dataclass class GeneratedScript: """Result of script generation.""" script: str confidence: float matched_keywords: List[str] template_used: Optional[ScriptTemplate] class ScriptGenerator: """Generates shell scripts from natural language descriptions.""" def __init__(self, db: Database): self.db = db self._ensure_default_templates() def _ensure_default_templates(self): existing = self.db.get_script_templates() if not existing: default_templates = [ ScriptTemplate( keywords=["deploy", "app", "application", "production"], template="#!/bin/bash\n\n# Deploy application\necho \"Deploying application...\"\nsource venv/bin/activate\npip install -r requirements.txt\npython manage.py migrate\npython manage.py collectstatic --noinput\nsudo systemctl restart gunicorn\necho \"Deployment complete!\"", description="Deploy a Python web application to production", ), ScriptTemplate( keywords=["git", "commit", "message", "push"], template="#!/bin/bash\n\n# Commit and push changes\necho \"Committing and pushing changes...\"\ngit add .\necho \"Enter commit message:\"\nread -r msg\ngit commit -m \"$msg\"\ngit push origin main\necho \"Changes committed and pushed!\"", description="Commit changes to git with message and push", ), ScriptTemplate( keywords=["docker", "build", "image", "container"], template="#!/bin/bash\n\n# Build and run Docker container\necho \"Building Docker image...\"\ndocker build -t ${IMAGE_NAME:-myapp} .\necho \"Running container...\"\ndocker run -d -p ${PORT:-8000}:8000 --name ${CONTAINER_NAME:-myapp-run} ${IMAGE_NAME:-myapp}\necho \"Container started!\"", description="Build and run a Docker container", ), ScriptTemplate( keywords=["backup", "database", "db", "sql"], template="#!/bin/bash\n\n# Backup database\necho \"Creating database backup...\"\ndatetime=$(date +%Y%m%d_%H%M%S)\npg_dump -U ${DB_USER:-postgres} ${DB_NAME:-mydb} > backup_${DB_NAME:-mydb}_$datetime.sql\ngzip backup_${DB_NAME:-mydb}_$datetime.sql\necho \"Backup created: backup_${DB_NAME:-mydb}_$datetime.sql.gz\"", description="Create a compressed database backup", ), ScriptTemplate( keywords=["test", "pytest", "coverage"], template="#!/bin/bash\n\n# Run tests with coverage\necho \"Running tests...\"\npytest --cov=./ --cov-report=term-missing --cov-report=html\n\necho \"Opening coverage report...\"\nopen htmlcov/index.html 2>/dev/null || xdg-open htmlcov/index.html 2>/dev/null || echo \"Coverage report generated in htmlcov/\"", description="Run pytest with coverage reporting", ), ] for template in default_templates: self.db.add_script_template(template) def generate(self, description: str, custom_vars: Optional[Dict[str, str]] = None) -> GeneratedScript: description_lower = description.lower() words = description_lower.split() templates = self.db.get_script_templates() best_match = None best_score = 0 matched_keywords = [] for template in templates: score = 0 template_matched = [] for keyword in template.keywords: if keyword in description_lower or keyword in words: score += 1 template_matched.append(keyword) if score > best_score: best_score = score best_match = template matched_keywords = template_matched if best_match and best_score > 0: script = best_match.template if custom_vars: for key, value in custom_vars.items(): script = script.replace(f"${{{key}}}", value) script = script.replace(f"${key}", value) confidence = best_score / max(len(best_match.keywords), 1) return GeneratedScript( script=script, confidence=confidence, matched_keywords=matched_keywords, template_used=best_match, ) simple_script = self._generate_simple_script(description) return GeneratedScript( script=simple_script, confidence=0.3, matched_keywords=[], template_used=None, ) def _generate_simple_script(self, description: str) -> str: words = description.lower().split() script_lines = ["#!/bin/bash", "", f"# Generated from: {description}", ""] if "docker" in words or "container" in words: script_lines.extend([ "# Docker operations", "docker ps", "docker images", "echo 'Docker commands ready'", ]) elif "git" in words or "commit" in words: script_lines.extend([ "# Git operations", "git status", "git add .", "echo 'Enter commit message:'", "read -r msg", "git commit -m \"$msg\"", "git push", ]) elif "test" in words: script_lines.extend([ "# Test operations", "pytest -v", "echo 'Tests completed'", ]) elif "backup" in words: script_lines.extend([ "# Backup operations", "echo 'Creating backup...'", f"BACKUP_FILE=backup_$(date +%Y%m%d_%H%M%S).tar.gz", "tar -czf $BACKUP_FILE .", "echo 'Backup created: $BACKUP_FILE'", ]) else: script_lines.extend([ "# Script generated from your description", f"# Description: {description}", "", "echo 'Running: {description}'", "", "# TODO: Customize this script based on your needs", "# Add your commands here", ]) return "\n".join(script_lines) def list_templates(self) -> List[ScriptTemplate]: return self.db.get_script_templates() def add_template(self, keywords: List[str], template: str, description: str = "") -> Optional[int]: script_template = ScriptTemplate( keywords=keywords, template=template, description=description, ) return self.db.add_script_template(script_template)