diff --git a/shell_memory/scripts.py b/shell_memory/scripts.py new file mode 100644 index 0000000..d41ad18 --- /dev/null +++ b/shell_memory/scripts.py @@ -0,0 +1,164 @@ +"""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) \ No newline at end of file