Add config and ollama client modules

shellgenius/config.py

@@ -0,0 +1,142 @@
+"""Configuration loader for ShellGenius."""
+
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import yaml
+
+
+class Config:
+    """Configuration management for ShellGenius."""
+
+    def __init__(self, config_path: Optional[str] = None):
+        """Initialize configuration.
+
+        Args:
+            config_path: Path to config.yaml file
+        """
+        self.config_path = config_path or os.environ.get(
+            "SHELLGENIUS_CONFIG", "config.yaml"
+        )
+        self.config: Dict[str, Any] = self._load_config()
+
+    def _load_config(self) -> Dict[str, Any]:
+        """Load configuration from YAML file.
+
+        Returns:
+            Dictionary containing configuration
+        """
+        default_config = {
+            "ollama": {
+                "host": "localhost:11434",
+                "model": "codellama",
+                "timeout": 120,
+            },
+            "safety": {
+                "level": "moderate",
+                "warn_on_destructive": True,
+                "allow_dangerous": False,
+                "blocked_patterns": [
+                    "rm -rf /",
+                    ":(){:|:&};:",
+                    "chmod 777",
+                    "sudo su",
+                    "dd if=/dev/zero",
+                ],
+            },
+            "ui": {
+                "theme": "default",
+                "show_syntax_highlighting": True,
+                "auto_suggest": True,
+                "max_history": 100,
+            },
+            "shell": {
+                "default": "bash",
+                "prefer_shebang": True,
+            },
+            "history": {
+                "enabled": True,
+                "storage_path": "~/.config/shellgenius/history.yaml",
+                "similarity_threshold": 0.7,
+            },
+        }
+
+        if self.config_path and Path(self.config_path).exists():
+            try:
+                with open(self.config_path, "r") as f:
+                    user_config = yaml.safe_load(f) or {}
+                return self._merge_configs(default_config, user_config)
+            except Exception:
+                return default_config
+        return default_config
+
+    def _merge_configs(
+        self, default: Dict[str, Any], user: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Merge user config with defaults.
+
+        Args:
+            default: Default configuration
+            user: User-provided configuration
+
+        Returns:
+            Merged configuration
+        """
+        result = default.copy()
+        for key, value in user.items():
+            if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+                result[key] = self._merge_configs(result[key], value)
+            else:
+                result[key] = value
+        return result
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get configuration value by key.
+
+        Args:
+            key: Dot-separated key path (e.g., "ollama.host")
+            default: Default value if key not found
+
+        Returns:
+            Configuration value
+        """
+        keys = key.split(".")
+        value = self.config
+        for k in keys:
+            if isinstance(value, dict) and k in value:
+                value = value[k]
+            else:
+                return default
+        return value
+
+    @property
+    def ollama_host(self) -> str:
+        """Get Ollama host."""
+        return os.environ.get("OLLAMA_HOST", self.get("ollama.host", "localhost:11434"))
+
+    @property
+    def ollama_model(self) -> str:
+        """Get Ollama model."""
+        return os.environ.get("OLLAMA_MODEL", self.get("ollama.model", "codellama"))
+
+    @property
+    def safety_level(self) -> str:
+        """Get safety level."""
+        return os.environ.get("SHELLGENIUS_SAFETY", self.get("safety.level", "moderate"))
+
+    def reload(self) -> None:
+        """Reload configuration from file."""
+        self.config = self._load_config()
+
+
+def get_config(config_path: Optional[str] = None) -> Config:
+    """Get configuration instance.
+
+    Args:
+        config_path: Optional path to config file
+
+    Returns:
+        Config instance
+    """
+    return Config(config_path)