shellmate-cli 3.0.0__py3-none-any.whl

shellmate/providers/openrouter_provider.py

@@ -0,0 +1,149 @@
+import os
+import json
+from typing import Optional
+from openai import AsyncOpenAI
+from .base import AIProvider, ProviderError
+
+
+def _format_context(context: Optional[dict]) -> str:
+    if not context:
+        return ""
+    parts = []
+    watcher = context.get("watcher", {})
+    if watcher:
+        parts.append("System Context:")
+        if watcher.get("os"):
+            parts.append(f"  OS: {watcher['os']}")
+        if watcher.get("shell"):
+            parts.append(f"  Shell: {watcher['shell']}")
+        if watcher.get("current_directory"):
+            parts.append(f"  Directory: {watcher['current_directory']}")
+        git = watcher.get("git_context")
+        if git:
+            git_parts = []
+            if git.get("branch"):
+                git_parts.append(f"branch={git['branch']}")
+            if git.get("staged"):
+                git_parts.append(f"staged={git['staged']}")
+            if git.get("unstaged"):
+                git_parts.append(f"unstaged={git['unstaged']}")
+            if git_parts:
+                parts.append(f"  Git: {', '.join(git_parts)}")
+        stack = watcher.get("tech_stack")
+        if stack:
+            parts.append(f"  Tech stack: {', '.join(stack[:5])}")
+    librarian = context.get("librarian", {})
+    if librarian and librarian.get("matches"):
+        parts.append("\nRelevant project files:")
+        for m in librarian["matches"][:3]:
+            path = m.get("file", m.get("path", "?"))
+            snippet = m.get("content", m.get("snippet", ""))[:200]
+            parts.append(f"  - {path}: {snippet}")
+    if not parts:
+        return ""
+    return "\n" + "\n".join(parts)
+
+
+class OpenRouterProvider(AIProvider):
+    def __init__(self):
+        self.api_key = os.environ.get("OPENROUTER_API_KEY")
+        self.client = None
+        self.base_url = "https://openrouter.ai/api/v1"
+        self.primary_model = os.environ.get("OPENROUTER_PRIMARY_MODEL", "gpt-4o-mini")
+        self.fallback_models = [
+            "meta-llama/llama-3.3-70b-instruct",
+            "qwen/qwen2.5-72b-instruct",
+        ]
+
+    def _ensure_client(self):
+        if not self.api_key:
+            raise ProviderError("OPENROUTER_API_KEY environment variable not set", "openrouter")
+        if self.client is None:
+            self.client = AsyncOpenAI(
+                base_url=self.base_url,
+                api_key=self.api_key,
+            )
+
+    async def analyze_risk(self, command: str, context: Optional[dict] = None) -> dict:
+        raise NotImplementedError("OpenRouter provider only handles remediation")
+
+    async def generate_remediation(self, command: str, risk: dict, context: Optional[dict] = None) -> dict:
+        self._ensure_client()
+
+        risk_level = risk.get("risk", risk.get("severity", "UNKNOWN"))
+        reason = risk.get("reason", risk.get("message", "No reason provided"))
+
+        context_str = _format_context(context)
+        risk_upper = str(risk_level).upper()
+        if risk_upper in ("SAFE", "LOW"):
+            advisor_role = (
+                "You are a helpful shell command advisor. The user's command is low risk. "
+                "They likely need usage help: explain what they might be trying to do, "
+                "what each argument likely means, and suggest a corrected command using context."
+            )
+            alt_hint = (
+                'Start alternatives with "Maybe you\'re trying to..." then parameter notes, '
+                'then "Try: <corrected command>".'
+            )
+        else:
+            advisor_role = (
+                "You are a command safety advisor. A user typed a potentially dangerous command. "
+                "Provide safe alternatives that preserve the original intent."
+            )
+            alt_hint = "Maximum 3 safer alternatives. Never alter the user's original intent."
+
+        prompt = (
+            f"{advisor_role}\n\n"
+            f"Command: {command}\n"
+            f"Risk Level: {risk_level}\n"
+            f"Reason: {reason}"
+            f"{context_str}\n\n"
+            "Respond ONLY with valid JSON in this exact format:\n"
+            '{"explanation": "plain language explanation", "alternatives": ["alt1", "alt2", "alt3"]}\n'
+            f"{alt_hint}"
+        )
+
+        models_to_try = [self.primary_model] + self.fallback_models
+        last_error = None
+
+        for model in models_to_try:
+            try:
+                response = await self.client.chat.completions.create(
+                    model=model,
+                    messages=[{"role": "user", "content": prompt}],
+                    temperature=0.3,
+                    max_tokens=512,
+                )
+
+                content = response.choices[0].message.content.strip()
+                if content.startswith("```"):
+                    content = content.split("\n", 1)[1] if "\n" in content else content
+                    if content.endswith("```"):
+                        content = content[:-3]
+                    content = content.strip()
+
+                result = json.loads(content)
+                result["provider"] = "openrouter"
+                result["model"] = model
+                return result
+
+            except Exception as e:
+                last_error = e
+                continue
+
+        from .base import wrap_provider_exception
+
+        wrapped = wrap_provider_exception(last_error or Exception("unknown"), "openrouter")
+        raise ProviderError(
+            f"All models failed. Last error: {wrapped}",
+            "openrouter",
+            status_code=wrapped.status_code,
+        )
+
+    async def close(self):
+        if self.client is not None:
+            await self.client.aclose()
+            self.client = None
+
+    def is_available(self) -> bool:
+        return bool(self.api_key)

shellmate/providers/router.py

@@ -0,0 +1,272 @@
+import os
+from .base import ProviderError, format_provider_failure, wrap_provider_exception, credential_key_for
+from .groq_provider import GroqProvider
+from .openrouter_provider import OpenRouterProvider
+from .xai_provider import XaiProvider
+from shellmate.config.credentials import load_env
+
+_groq_provider = None
+_openrouter_provider = None
+_xai_provider = None
+
+_RISK_CAPABLE = ("groq", "xai")
+_REMEDIATION_CAPABLE = ("openrouter", "xai")
+
+
+def _inject_env():
+    env = load_env()
+    for key, value in env.items():
+        if value and not os.environ.get(key):
+            os.environ[key] = value
+
+
+def get_groq_provider():
+    global _groq_provider
+    _inject_env()
+    if _groq_provider is None:
+        _groq_provider = GroqProvider()
+    return _groq_provider
+
+
+def get_openrouter_provider():
+    global _openrouter_provider
+    _inject_env()
+    if _openrouter_provider is None:
+        _openrouter_provider = OpenRouterProvider()
+    return _openrouter_provider
+
+
+def get_xai_provider():
+    global _xai_provider
+    _inject_env()
+    if _xai_provider is None:
+        _xai_provider = XaiProvider()
+    return _xai_provider
+
+
+def _provider_instance(name: str):
+    if name == "groq":
+        return get_groq_provider()
+    if name == "openrouter":
+        return get_openrouter_provider()
+    if name == "xai":
+        return get_xai_provider()
+    return None
+
+
+def _role_primary(role: str, default: str) -> str:
+    try:
+        from shellmate.config import get_config
+        cfg = get_config()
+        primary = (cfg.get("providers") or {}).get(role, {}).get("primary") or default
+        return str(primary).lower()
+    except Exception:
+        return default
+
+
+def _apply_role_model(provider, role: str) -> None:
+    try:
+        from shellmate.config import get_config
+        model = (get_config().get("providers") or {}).get(role, {}).get("model")
+        if not model:
+            return
+        if hasattr(provider, "model"):
+            provider.model = model
+        elif hasattr(provider, "primary_model"):
+            provider.primary_model = model
+    except Exception:
+        pass
+
+
+def _risk_provider_order() -> list:
+    primary = _role_primary("risk_analysis", "groq")
+    order = []
+    if primary in _RISK_CAPABLE:
+        order.append(primary)
+    for name in _RISK_CAPABLE:
+        if name not in order:
+            order.append(name)
+    return order
+
+
+def _remediation_provider_order() -> list:
+    primary = _role_primary("remediation", "openrouter")
+    order = []
+    if primary in _REMEDIATION_CAPABLE:
+        order.append(primary)
+    for name in _REMEDIATION_CAPABLE:
+        if name not in order:
+            order.append(name)
+    return order
+
+
+async def check_risk_model_reachable():
+    """Best-effort startup check: is the configured risk model actually reachable?
+
+    Returns a warning dict when a key IS configured but the resolved risk model
+    is not in the account's accessible-model list — the case where AI analysis
+    would silently degrade to a static MEDIUM fallback. Returns None when there
+    is nothing to warn about (no key, model reachable, or the check is
+    inconclusive). Never raises and never blocks daemon readiness.
+    """
+    try:
+        order = _risk_provider_order()
+        name = order[0] if order else "groq"
+        provider = _provider_instance(name)
+        # No key configured -> AI tier is intentionally offline, not a misconfig.
+        if provider is None or not provider.is_available():
+            return None
+        _apply_role_model(provider, "risk_analysis")
+        model = getattr(provider, "model", None)
+        lister = getattr(provider, "list_available_models", None)
+        if not model or lister is None:
+            return None
+        available = await lister()
+        # Could not determine (network/API error) -> stay silent, don't cry wolf.
+        if not available or model in available:
+            return None
+        return {
+            "provider": name,
+            "model": model,
+            "credential": credential_key_for(name),
+            "accessible_models": sorted(available)[:12],
+            "message": (
+                f"Configured {name} risk model '{model}' is NOT accessible with your "
+                f"{credential_key_for(name)}. AI risk analysis will silently fall back to a "
+                f"static MEDIUM rating. Set providers.risk_analysis.model in config.yaml to a "
+                f"model your key can access (see `sm status` for the accessible list)."
+            ),
+        }
+    except Exception:
+        return None
+
+
+async def analyze_risk(command, context=None):
+    last_err = None
+    last_provider = "groq"
+
+    for name in _risk_provider_order():
+        provider = _provider_instance(name)
+        if provider is None or not provider.is_available():
+            continue
+        last_provider = name
+        try:
+            _apply_role_model(provider, "risk_analysis")
+            return await provider.analyze_risk(command, context)
+        except ProviderError as e:
+            last_err = wrap_provider_exception(e, name)
+            continue
+        except NotImplementedError:
+            continue
+
+    err = last_err or ProviderError("No risk analysis provider available", last_provider)
+    return {
+        "risk": "MEDIUM",
+        "confidence": 0.5,
+        "reason": format_provider_failure(last_provider, err, purpose="AI analysis"),
+        "provider": "fallback",
+        "model": "deterministic",
+        "failed_provider": last_provider,
+        "failed_credential": credential_key_for(last_provider),
+    }
+
+
+async def generate_remediation(command, risk, context=None):
+    last_err = None
+    last_provider = "openrouter"
+
+    for name in _remediation_provider_order():
+        provider = _provider_instance(name)
+        if provider is None or not provider.is_available():
+            continue
+        last_provider = name
+        try:
+            _apply_role_model(provider, "remediation")
+            return await provider.generate_remediation(command, risk, context)
+        except ProviderError as e:
+            last_err = wrap_provider_exception(e, name)
+            continue
+        except NotImplementedError:
+            continue
+
+    err = last_err or ProviderError("No remediation provider available", last_provider)
+    groq_guide = await _remediation_from_command_guide(command, risk, context)
+    if groq_guide:
+        groq_guide["openrouter_error"] = format_provider_failure(
+            last_provider, err, purpose="Remediation"
+        )
+        return groq_guide
+    return {
+        "explanation": format_provider_failure(last_provider, err, purpose="Remediation")
+        + " "
+        + risk.get("message", risk.get("reason", "Command may be dangerous.")),
+        "alternatives": ["Review the command carefully before executing."],
+        "provider": "fallback",
+        "model": "deterministic",
+        "failed_provider": last_provider,
+        "failed_credential": credential_key_for(last_provider),
+    }
+
+
+async def _remediation_from_command_guide(command, risk, context=None):
+    groq = get_groq_provider()
+    if not groq.is_available():
+        return None
+    try:
+        from shellmate.providers.groq_provider import _format_context
+        from shellmate.providers.groq_command_guide import generate_command_guide, alternatives_from_guide
+
+        context_str = _format_context(context)
+        reason = risk.get("reason", risk.get("message", ""))
+        guide = await generate_command_guide(
+            command,
+            context=context_str,
+            risk_reason=reason,
+        )
+        alternatives = alternatives_from_guide(guide)
+        if not alternatives:
+            return None
+        return {
+            "explanation": guide.get("when_to_use") or reason or "Usage suggestion based on your command.",
+            "alternatives": alternatives,
+            "provider": "fallback",
+            "model": "groq_command_guide",
+            "command_guide": guide,
+        }
+    except Exception:
+        return None
+
+
+async def close_providers():
+    if _groq_provider is not None:
+        await _groq_provider.close()
+    if _openrouter_provider is not None:
+        await _openrouter_provider.close()
+    if _xai_provider is not None:
+        await _xai_provider.close()
+
+
+def get_provider_status():
+    groq = get_groq_provider()
+    openrouter = get_openrouter_provider()
+    xai = get_xai_provider()
+
+    return {
+        "groq": {
+            "configured": groq.is_available(),
+            "status": "available" if groq.is_available() else "not_configured",
+            "credential": "GROQ_API_KEY",
+        },
+        "openrouter": {
+            "configured": openrouter.is_available(),
+            "status": "available" if openrouter.is_available() else "not_configured",
+            "credential": "OPENROUTER_API_KEY",
+            "fallback_models": openrouter.fallback_models,
+        },
+        "xai": {
+            "configured": xai.is_available(),
+            "status": "available" if xai.is_available() else "not_configured",
+            "credential": "XAI_API_KEY",
+            "model": xai.model,
+        },
+    }

shellmate/providers/validate_keys.py

@@ -0,0 +1,98 @@
+"""Live API-key validation against provider endpoints."""
+
+from __future__ import annotations
+
+import json
+import urllib.error
+import urllib.request
+from typing import Optional, Tuple
+
+from shellmate.config.credentials import CREDENTIAL_KEYS
+
+# credential env var -> (url, extra headers)
+_VALIDATION_ENDPOINTS = {
+    "GROQ_API_KEY": (
+        "https://api.groq.com/openai/v1/models",
+        {},
+    ),
+    "OPENROUTER_API_KEY": (
+        "https://openrouter.ai/api/v1/models",
+        {"HTTP-Referer": "https://github.com/sayan-does/shellmate", "X-Title": "ShellMate"},
+    ),
+    "XAI_API_KEY": (
+        "https://api.x.ai/v1/models",
+        {},
+    ),
+}
+
+DEFAULT_TIMEOUT_SEC = 10
+
+
+def validate_api_key(
+    credential_key: str,
+    api_key: str,
+    *,
+    timeout: float = DEFAULT_TIMEOUT_SEC,
+) -> Tuple[bool, str]:
+    """
+    Ping the provider endpoint with the given API key.
+
+    Returns (ok, message). Context7 and unknown keys without live endpoints
+    return ok=True with a note that they were not live-checked.
+    """
+    if not api_key or not api_key.strip():
+        return False, "API key is empty"
+
+    meta = CREDENTIAL_KEYS.get(credential_key, {})
+    if not meta.get("live_check", True) or credential_key not in _VALIDATION_ENDPOINTS:
+        label = meta.get("label", credential_key)
+        return True, f"{label} saved (not live-checked)"
+
+    url, extra_headers = _VALIDATION_ENDPOINTS[credential_key]
+    return _ping_models_endpoint(url, api_key.strip(), extra_headers, timeout=timeout)
+
+
+def _ping_models_endpoint(
+    url: str,
+    api_key: str,
+    extra_headers: dict,
+    *,
+    timeout: float,
+) -> Tuple[bool, str]:
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Accept": "application/json",
+        **extra_headers,
+    }
+    request = urllib.request.Request(url, headers=headers, method="GET")
+    try:
+        with urllib.request.urlopen(request, timeout=timeout) as response:
+            status = getattr(response, "status", None) or response.getcode()
+            body = response.read(2048)
+            if status and 200 <= int(status) < 300:
+                # Soft sanity check: body should look like JSON list/object
+                if body:
+                    try:
+                        json.loads(body.decode("utf-8", errors="replace"))
+                    except json.JSONDecodeError:
+                        pass
+                return True, "Key is valid"
+            return False, f"Unexpected response status {status}"
+    except urllib.error.HTTPError as e:
+        if e.code in (401, 403):
+            return False, f"Key is invalid or unauthorized ({e.code})"
+        if e.code == 429:
+            return False, "Rate limited (429); try again later"
+        return False, f"Provider returned HTTP {e.code}"
+    except urllib.error.URLError as e:
+        reason = getattr(e, "reason", e)
+        return False, f"Could not reach provider: {reason}"
+    except TimeoutError:
+        return False, "Could not reach provider: timed out"
+    except Exception as e:
+        return False, f"Could not reach provider: {e}"
+
+
+def validate_stored_key(credential_key: str, api_key: Optional[str] = None) -> Tuple[bool, str]:
+    """Validate a credential using an explicit key or the empty-string no-op."""
+    return validate_api_key(credential_key, api_key or "")

shellmate/providers/xai_provider.py

@@ -0,0 +1,142 @@
+import os
+import json
+import asyncio
+from typing import Optional
+from openai import AsyncOpenAI
+from .base import AIProvider, ProviderError
+from .groq_provider import _format_context
+
+
+class XaiProvider(AIProvider):
+    """Grok (xAI) provider — OpenAI-compatible API for risk and remediation."""
+
+    def __init__(self):
+        self.api_key = os.environ.get("XAI_API_KEY")
+        self.client = None
+        self.base_url = "https://api.x.ai/v1"
+        self.model = os.environ.get("XAI_MODEL", "grok-3-mini")
+
+    def _ensure_client(self):
+        if not self.api_key:
+            raise ProviderError("XAI_API_KEY environment variable not set", "xai")
+        if self.client is None:
+            self.client = AsyncOpenAI(
+                base_url=self.base_url,
+                api_key=self.api_key,
+            )
+
+    async def analyze_risk(self, command: str, context: Optional[dict] = None) -> dict:
+        self._ensure_client()
+
+        context_str = _format_context(context)
+        prompt = (
+            "You are a command safety analyzer. Analyze the following shell command and classify its risk level. "
+            "Respond ONLY with valid JSON in this exact format:\n"
+            '{"risk": "SAFE|LOW|MEDIUM|HIGH|CRITICAL", "confidence": 0.0-1.0, "reason": "brief explanation"}\n\n'
+            f"Command: {command}"
+            f"{context_str}"
+        )
+
+        try:
+            response = await self.client.chat.completions.create(
+                model=self.model,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.1,
+                max_tokens=256,
+            )
+
+            content = response.choices[0].message.content.strip()
+            if content.startswith("```"):
+                content = content.split("\n", 1)[1] if "\n" in content else content
+                if content.endswith("```"):
+                    content = content[:-3]
+                content = content.strip()
+
+            result = json.loads(content)
+            result["provider"] = "xai"
+            result["model"] = self.model
+            return result
+
+        except json.JSONDecodeError as e:
+            raise ProviderError(f"Invalid JSON response: {e}", "xai")
+        except ProviderError:
+            raise
+        except Exception as e:
+            from .base import wrap_provider_exception
+            raise wrap_provider_exception(e, "xai")
+
+    async def generate_remediation(self, command: str, risk: dict, context: Optional[dict] = None) -> dict:
+        self._ensure_client()
+
+        risk_level = risk.get("risk", risk.get("severity", "UNKNOWN"))
+        reason = risk.get("reason", risk.get("message", "No reason provided"))
+
+        context_str = _format_context(context)
+        risk_upper = str(risk_level).upper()
+        if risk_upper in ("SAFE", "LOW"):
+            advisor_role = (
+                "You are a helpful shell command advisor. The user's command is low risk. "
+                "They likely need usage help: explain what they might be trying to do, "
+                "what each argument likely means, and suggest a corrected command using context."
+            )
+            alt_hint = (
+                'Start alternatives with "Maybe you\'re trying to..." then parameter notes, '
+                'then "Try: <corrected command>".'
+            )
+        else:
+            advisor_role = (
+                "You are a command safety advisor. A user typed a potentially dangerous command. "
+                "Provide safe alternatives that preserve the original intent."
+            )
+            alt_hint = "Maximum 3 safer alternatives. Never alter the user's original intent."
+
+        prompt = (
+            f"{advisor_role}\n\n"
+            f"Command: {command}\n"
+            f"Risk Level: {risk_level}\n"
+            f"Reason: {reason}"
+            f"{context_str}\n\n"
+            "Respond ONLY with valid JSON in this exact format:\n"
+            '{"explanation": "plain language explanation", "alternatives": ["alt1", "alt2", "alt3"]}\n'
+            f"{alt_hint}"
+        )
+
+        try:
+            response = await self.client.chat.completions.create(
+                model=self.model,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.3,
+                max_tokens=512,
+            )
+
+            content = response.choices[0].message.content.strip()
+            if content.startswith("```"):
+                content = content.split("\n", 1)[1] if "\n" in content else content
+                if content.endswith("```"):
+                    content = content[:-3]
+                content = content.strip()
+
+            result = json.loads(content)
+            result["provider"] = "xai"
+            result["model"] = self.model
+            return result
+
+        except json.JSONDecodeError as e:
+            raise ProviderError(f"Invalid JSON response: {e}", "xai")
+        except ProviderError:
+            raise
+        except Exception as e:
+            from .base import wrap_provider_exception
+            raise wrap_provider_exception(e, "xai")
+
+    async def close(self):
+        if self.client is not None:
+            close = getattr(self.client, "aclose", None) or getattr(self.client, "close", None)
+            if close:
+                result = close()
+                if asyncio.iscoroutine(result):
+                    await result
+            self.client = None
+
+    def is_available(self) -> bool:
+        return bool(self.api_key)