129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
"""Command generation logic."""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
from .prompts import PromptTemplates
|
|
|
|
|
|
@dataclass
|
|
class GenerationResult:
|
|
"""Result of command generation."""
|
|
|
|
command: str
|
|
explanation: str
|
|
confidence: float
|
|
|
|
|
|
class CommandGenerator:
|
|
"""Generate shell commands from natural language descriptions."""
|
|
|
|
def __init__(self, backend):
|
|
"""Initialize the command generator.
|
|
|
|
Args:
|
|
backend: LLM backend instance for generating commands.
|
|
"""
|
|
self.backend = backend
|
|
self.templates = PromptTemplates()
|
|
|
|
def generate(
|
|
self,
|
|
description: str,
|
|
shell: str = "bash",
|
|
context: Optional[str] = None,
|
|
) -> GenerationResult:
|
|
"""Generate a shell command from natural language.
|
|
|
|
Args:
|
|
description: Natural language description of what to do.
|
|
shell: Target shell (bash or zsh).
|
|
context: Optional context (current directory, files, etc.).
|
|
|
|
Returns:
|
|
GenerationResult with command, explanation, and confidence.
|
|
"""
|
|
prompt = self.templates.generate_prompt(description, shell, context)
|
|
|
|
response = self.backend.generate(prompt)
|
|
|
|
command = self._extract_command(response)
|
|
explanation = self._extract_explanation(response)
|
|
confidence = self._estimate_confidence(response)
|
|
|
|
return GenerationResult(
|
|
command=command.strip(),
|
|
explanation=explanation.strip(),
|
|
confidence=confidence,
|
|
)
|
|
|
|
def _extract_command(self, response: str) -> str:
|
|
"""Extract the shell command from LLM response."""
|
|
lines = response.strip().split("\n")
|
|
command = ""
|
|
|
|
for line in lines:
|
|
line = line.strip()
|
|
if line.startswith("```"):
|
|
continue
|
|
if line.startswith("Command:") or line.startswith("```"):
|
|
if "```" in line:
|
|
parts = line.split("```")
|
|
for part in parts:
|
|
if part.strip() and not part.strip().startswith("Command"):
|
|
command = part.strip()
|
|
break
|
|
else:
|
|
command = line.replace("Command:", "").strip()
|
|
elif line and not line.startswith("#") and not line.startswith("//"):
|
|
if command:
|
|
command += " " + line
|
|
else:
|
|
command = line
|
|
|
|
if not command:
|
|
for line in lines:
|
|
if line.startswith("$") or line.startswith("!"):
|
|
command = line[1:].strip()
|
|
break
|
|
|
|
return command if command else response.strip()
|
|
|
|
def _extract_explanation(self, response: str) -> str:
|
|
"""Extract the explanation from LLM response."""
|
|
lines = response.strip().split("\n")
|
|
explanation = []
|
|
in_explanation = False
|
|
|
|
for line in lines:
|
|
line_lower = line.lower().strip()
|
|
if "explanation" in line_lower or "what this does" in line_lower:
|
|
in_explanation = True
|
|
continue
|
|
if line.startswith("```"):
|
|
in_explanation = False
|
|
continue
|
|
if in_explanation:
|
|
explanation.append(line)
|
|
elif line.startswith("#") or line.startswith("//"):
|
|
explanation.append(line.lstrip("#/").strip())
|
|
|
|
return " ".join(explanation) if explanation else "Generated command"
|
|
|
|
def _estimate_confidence(self, response: str) -> float:
|
|
"""Estimate confidence score based on response characteristics."""
|
|
confidence = 0.5
|
|
|
|
if "```" in response:
|
|
confidence += 0.2
|
|
if len(response) > 50:
|
|
confidence += 0.1
|
|
if "sure" in response.lower() or "certain" in response.lower():
|
|
confidence += 0.1
|
|
if "?" in response:
|
|
confidence -= 0.2
|
|
if "not sure" in response.lower() or "uncertain" in response.lower():
|
|
confidence -= 0.3
|
|
|
|
return max(0.0, min(1.0, confidence))
|