diff --git a/shellgenius/script_generator.py b/shellgenius/script_generator.py new file mode 100644 index 0000000..1f7a757 --- /dev/null +++ b/shellgenius/script_generator.py @@ -0,0 +1,67 @@ +"""Script generation and review using Ollama.""" + +import logging +from typing import Any + +from shellgenius.ollama_client import OllamaClient + +logger = logging.getLogger(__name__) + + +class ScriptGenerator: + """Generate and review shell scripts using Ollama.""" + + def __init__(self, client: OllamaClient, config: dict[str, Any]): + """Initialize script generator. + + Args: + client: Ollama client instance + config: Configuration dictionary + """ + self.client = client + self.config = config + + def generate(self, prompt: str, shell: str = "bash") -> str: + """Generate a shell script from a natural language prompt. + + Args: + prompt: Description of what the script should do + shell: Target shell type (bash, zsh, sh) + + Returns: + Generated shell script + """ + full_prompt = f"""Generate a {shell} script that does the following: +{prompt} + +Return ONLY the script code, no explanations or markdown formatting. +The script should be production-ready with proper error handling. +""" + response = self.client.generate(full_prompt) + if response.get("success"): + return response["response"]["response"] + return f"# Error: {response.get('error', 'Unknown error')}" + + def review(self, script: str) -> str: + """Review and explain a shell script. + + Args: + script: Shell script content to review + + Returns: + Review and explanation of the script + """ + review_prompt = f"""Review this {self.config.get('default_shell', 'bash')} script and provide: +1. What the script does +2. Any issues or improvements needed +3. Security concerns + +Script: +{script} + +Provide a clear, concise review. +""" + response = self.client.generate(review_prompt) + if response.get("success"): + return response["response"]["response"] + return f"# Error: {response.get('error', 'Unknown error')}"