dbagent-cli 0.1.0__py3-none-any.whl

dbagent/llm/auto_setup.py

@@ -0,0 +1,269 @@
+"""
+Ollama Auto-Setup: Automatically install, start, and configure Ollama
+so DB-Agent works 100% independently with zero manual setup.
+"""
+
+import os
+import sys
+import time
+import subprocess
+import shutil
+import requests
+from pathlib import Path
+from typing import Optional, Tuple
+
+# Small, fast model that's good at SQL generation (~1.6 GB)
+DEFAULT_MODEL = "qwen2.5-coder:1.5b"
+OLLAMA_BASE_URL = "http://localhost:11434"
+OLLAMA_WINDOWS_INSTALLER_URL = "https://ollama.com/download/OllamaSetup.exe"
+
+
+def find_ollama_binary() -> Optional[str]:
+    """Find ollama executable on PATH or in standard Windows installation directories."""
+    which_path = shutil.which("ollama")
+    if which_path:
+        return which_path
+
+    # Common Windows install locations
+    local_app_data = os.environ.get("LOCALAPPDATA", "")
+    user_profile = os.environ.get("USERPROFILE", "")
+    candidates = [
+        Path(local_app_data) / "Programs" / "Ollama" / "ollama.exe" if local_app_data else None,
+        Path(user_profile) / "AppData" / "Local" / "Programs" / "Ollama" / "ollama.exe" if user_profile else None,
+        Path("C:/Program Files/Ollama/ollama.exe"),
+        Path("C:/Program Files (x86)/Ollama/ollama.exe"),
+    ]
+    for p in candidates:
+        if p and p.exists():
+            # Add to PATH for current session
+            os.environ["PATH"] = str(p.parent) + os.pathsep + os.environ.get("PATH", "")
+            return str(p)
+
+    return None
+
+
+def is_ollama_installed() -> bool:
+    """Check if ollama binary is available."""
+    return find_ollama_binary() is not None
+
+
+def is_ollama_running() -> bool:
+    """Check if Ollama server is responding."""
+    try:
+        r = requests.get(f"{OLLAMA_BASE_URL}/api/tags", timeout=2)
+        return r.status_code == 200
+    except Exception:
+        return False
+
+
+def get_installed_models() -> list:
+    """Get list of locally installed Ollama models."""
+    try:
+        r = requests.get(f"{OLLAMA_BASE_URL}/api/tags", timeout=4)
+        if r.status_code == 200:
+            data = r.json()
+            return [m["name"] for m in data.get("models", [])]
+    except Exception:
+        pass
+    return []
+
+
+def start_ollama_server() -> bool:
+    """Start ollama serve in the background."""
+    if is_ollama_running():
+        return True
+
+    exe = find_ollama_binary() or "ollama"
+    try:
+        if sys.platform == "win32":
+            subprocess.Popen(
+                [exe, "serve"],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                creationflags=subprocess.CREATE_NO_WINDOW | subprocess.DETACHED_PROCESS,
+            )
+        else:
+            subprocess.Popen(
+                [exe, "serve"],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                start_new_session=True,
+            )
+
+        # Wait for server to become ready
+        for _ in range(15):
+            time.sleep(1)
+            if is_ollama_running():
+                return True
+        return False
+    except Exception:
+        return False
+
+
+def pull_model(model_name: str, progress_callback=None) -> bool:
+    """Pull (download) an Ollama model with progress reporting."""
+    try:
+        resp = requests.post(
+            f"{OLLAMA_BASE_URL}/api/pull",
+            json={"name": model_name},
+            stream=True,
+            timeout=900,
+        )
+        if resp.status_code != 200:
+            return False
+
+        for line in resp.iter_lines():
+            if line and progress_callback:
+                import json
+                try:
+                    data = json.loads(line)
+                    status = data.get("status", "")
+                    total = data.get("total", 0)
+                    completed = data.get("completed", 0)
+                    progress_callback(status, completed, total)
+                except Exception:
+                    pass
+        return True
+    except Exception:
+        return False
+
+
+def download_ollama_installer(dest_dir: Optional[Path] = None, progress_callback=None) -> Optional[Path]:
+    """Download the Ollama Windows installer."""
+    if dest_dir is None:
+        dest_dir = Path.home() / ".dbagent" / "installers"
+    dest_dir.mkdir(parents=True, exist_ok=True)
+    installer_path = dest_dir / "OllamaSetup.exe"
+
+    if installer_path.exists() and installer_path.stat().st_size > 10_000_000:
+        return installer_path
+
+    try:
+        resp = requests.get(OLLAMA_WINDOWS_INSTALLER_URL, stream=True, timeout=180)
+        if resp.status_code != 200:
+            return None
+
+        total = int(resp.headers.get("content-length", 0))
+        downloaded = 0
+
+        with open(installer_path, "wb") as f:
+            for chunk in resp.iter_content(chunk_size=1024 * 128):
+                f.write(chunk)
+                downloaded += len(chunk)
+                if progress_callback and total > 0:
+                    pct = int(downloaded / total * 100)
+                    progress_callback(f"Downloading installer: {pct}%", downloaded, total)
+
+        return installer_path
+    except Exception:
+        if installer_path.exists():
+            try:
+                installer_path.unlink()
+            except Exception:
+                pass
+        return None
+
+
+def install_ollama_windows(installer_path: Path) -> bool:
+    """Run the Ollama installer on Windows."""
+    try:
+        # Run silent installer
+        subprocess.run(
+            [str(installer_path), "/VERYSILENT", "/NORESTART"],
+            timeout=180,
+            capture_output=True,
+        )
+        # Give Windows a moment to complete writing files
+        for _ in range(10):
+            time.sleep(2)
+            if is_ollama_installed():
+                return True
+
+        return is_ollama_installed()
+    except Exception:
+        return False
+
+
+def ensure_ollama_ready(
+    model: str = DEFAULT_MODEL,
+    print_fn=None,
+) -> Tuple[bool, str]:
+    """
+    Full auto-setup pipeline: install -> start -> pull model.
+    Returns (success, status_message).
+    """
+    def log(msg: str, level: str = "info"):
+        if print_fn:
+            print_fn(msg, level)
+
+    # --- Step 1: Check if Ollama is installed ---
+    if not is_ollama_installed():
+        if sys.platform != "win32":
+            return False, (
+                "Ollama is not installed. Install it from https://ollama.com/download "
+                "or run: curl -fsSL https://ollama.com/install.sh | sh"
+            )
+
+        log("[*] Ollama not found on system. Auto-downloading standalone setup...", "info")
+
+        installer = download_ollama_installer(
+            progress_callback=lambda status, c, t: log(f"  {status}", "progress")
+        )
+        if not installer:
+            return False, "Failed to download Ollama installer. Please install manually from https://ollama.com/download"
+
+        log("[*] Installing Ollama standalone AI engine (one-time setup)...", "info")
+        if not install_ollama_windows(installer):
+            return False, "Ollama installation failed. Please install from https://ollama.com/download"
+
+        log("[OK] Ollama installed successfully!", "success")
+
+    # --- Step 2: Start Ollama server if not running ---
+    if not is_ollama_running():
+        log("[*] Starting local Ollama AI server...", "info")
+        if not start_ollama_server():
+            return False, "Failed to start Ollama server. Try running `ollama serve` manually."
+        log("[OK] Local Ollama server is running and ready!", "success")
+
+    # --- Step 3: Check if model is available, pull if not ---
+    installed_models = get_installed_models()
+
+    if installed_models:
+        for m in installed_models:
+            if model.split(":")[0] in m:
+                return True, f"Ollama ready with local model '{m}'."
+        return True, f"Ollama ready with local model '{installed_models[0]}'."
+
+    # No models installed — pull one
+    log(f"[*] Downloading standalone coding model '{model}' (~1.2GB, one-time)...", "info")
+
+    last_pct = [-1]
+    def progress_cb(status, completed, total):
+        if total > 0:
+            pct = int(completed / total * 100)
+            if pct != last_pct[0] and pct % 10 == 0:
+                last_pct[0] = pct
+                log(f"  Downloading model: {pct}%", "progress")
+        elif status:
+            log(f"  {status}", "progress")
+
+    if not pull_model(model, progress_callback=progress_cb):
+        return False, f"Failed to download model '{model}'. Check your internet connection."
+
+    log(f"[OK] Model '{model}' downloaded and ready!", "success")
+    return True, f"Ollama ready with model '{model}'."
+
+
+def get_best_available_model() -> Optional[str]:
+    """Get the best available model that's already pulled."""
+    models = get_installed_models()
+    if not models:
+        return None
+
+    preferred_order = ["qwen2.5-coder", "codellama", "deepseek-coder", "qwen", "llama", "mistral"]
+    for pref in preferred_order:
+        for m in models:
+            if pref in m:
+                return m
+
+    return models[0]

dbagent/llm/base.py

@@ -0,0 +1,37 @@
+"""
+Base interface for LLM providers.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional, Callable, List, Dict, Any, Generator
+
+
+class BaseLLMProvider(ABC):
+    """Abstract base class for all AI/LLM providers."""
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Name of the provider (e.g. 'ollama', 'gemini', 'groq')."""
+        pass
+
+    @abstractmethod
+    def is_available(self) -> bool:
+        """Check if the provider is currently reachable and configured."""
+        pass
+
+    @abstractmethod
+    def list_models(self) -> List[str]:
+        """List available models for this provider."""
+        pass
+
+    @abstractmethod
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        model: Optional[str] = None,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> str:
+        """Generate response from LLM."""
+        pass

dbagent/llm/factory.py

@@ -0,0 +1,73 @@
+"""
+LLM Provider Factory.
+Auto-detects available providers: Ollama (local) -> Gemini Free -> Groq Free -> OpenRouter Free -> Offline Fallback.
+"""
+
+from typing import Optional
+from dbagent.config import ConfigManager
+from dbagent.llm.base import BaseLLMProvider
+from dbagent.llm.ollama_provider import OllamaProvider
+from dbagent.llm.gemini_provider import GeminiProvider
+from dbagent.llm.groq_provider import GroqProvider
+from dbagent.llm.openrouter_provider import OpenRouterProvider
+from dbagent.llm.mock_provider import MockProvider
+
+
+def get_llm_provider(
+    provider_name: Optional[str] = None,
+    model: Optional[str] = None,
+    config_mgr: Optional[ConfigManager] = None,
+) -> BaseLLMProvider:
+    """Instantiate and return the requested or best available LLM provider."""
+    cfg = config_mgr or ConfigManager()
+
+    # 1. Explicit provider requested
+    if provider_name:
+        p = provider_name.lower()
+        if p == "ollama":
+            base_url = cfg.get_setting("ollama_base_url", "http://localhost:11434")
+            return OllamaProvider(base_url=base_url, default_model=model or "qwen2.5-coder")
+        elif p == "gemini":
+            key = cfg.get_setting("gemini_api_key")
+            return GeminiProvider(api_key=key, default_model=model or "gemini-3.6-flash")
+        elif p == "groq":
+            key = cfg.get_setting("groq_api_key")
+            return GroqProvider(api_key=key, default_model=model or "llama-3.3-70b-versatile")
+        elif p == "openrouter":
+            key = cfg.get_setting("openrouter_api_key")
+            return OpenRouterProvider(api_key=key, default_model=model or "meta-llama/llama-3.3-70b-instruct:free")
+        elif p in ["mock", "offline"]:
+            return MockProvider()
+
+    # 2. Check saved default provider
+    saved_default = cfg.get_setting("default_provider")
+    if saved_default and saved_default != "auto":
+        try:
+            return get_llm_provider(provider_name=saved_default, model=model, config_mgr=cfg)
+        except Exception:
+            pass
+
+    # 3. Auto-detection priority:
+    # A. Local Ollama if running
+    ollama_url = cfg.get_setting("ollama_base_url", "http://localhost:11434")
+    ollama = OllamaProvider(base_url=ollama_url)
+    if ollama.is_available():
+        return ollama
+
+    # B. Gemini if API key is present
+    gemini_key = cfg.get_setting("gemini_api_key")
+    if gemini_key:
+        return GeminiProvider(api_key=gemini_key)
+
+    # C. Groq if API key is present
+    groq_key = cfg.get_setting("groq_api_key")
+    if groq_key:
+        return GroqProvider(api_key=groq_key)
+
+    # D. OpenRouter if API key is present
+    openrouter_key = cfg.get_setting("openrouter_api_key")
+    if openrouter_key:
+        return OpenRouterProvider(api_key=openrouter_key)
+
+    # E. Fallback to Ollama or Mock
+    return ollama

dbagent/llm/gemini_provider.py

@@ -0,0 +1,110 @@
+"""
+Google Gemini Free Tier LLM Provider.
+Uses Gemini 2.0 Flash / Gemini 1.5 Flash via direct HTTP API.
+"""
+
+import json
+import os
+import requests
+from typing import Optional, Callable, List
+from dbagent.llm.base import BaseLLMProvider
+
+
+class GeminiProvider(BaseLLMProvider):
+    """Google Gemini LLM provider (Free Tier available)."""
+
+    def __init__(self, api_key: Optional[str] = None, default_model: str = "gemini-3.6-flash"):
+        self.api_key = api_key or os.getenv("GEMINI_API_KEY")
+        self.default_model = default_model
+
+    @property
+    def name(self) -> str:
+        return "gemini"
+
+    def is_available(self) -> bool:
+        return bool(self.api_key)
+
+    def list_models(self) -> List[str]:
+        return [
+            "gemini-3.6-flash",
+            "gemini-2.5-flash",
+            "gemini-2.5-pro",
+        ]
+
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        model: Optional[str] = None,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> str:
+        if not self.api_key:
+            raise ValueError(
+                "Gemini API key is not configured. Set GEMINI_API_KEY environment variable or run `db-agent config`."
+            )
+
+        target_model = model or self.default_model
+        # Strip model prefixes if provided
+        if target_model.startswith("models/"):
+            target_model = target_model[7:]
+
+        endpoint = f"https://generativelanguage.googleapis.com/v1beta/models/{target_model}:generateContent?key={self.api_key}"
+        if stream_callback:
+            endpoint = f"https://generativelanguage.googleapis.com/v1beta/models/{target_model}:streamGenerateContent?alt=sse&key={self.api_key}"
+
+        contents = []
+        if system_prompt:
+            contents.append({
+                "role": "user",
+                "parts": [{"text": f"System Instructions: {system_prompt}\n\nTask: {prompt}"}],
+            })
+        else:
+            contents.append({
+                "role": "user",
+                "parts": [{"text": prompt}],
+            })
+
+        payload = {
+            "contents": contents,
+            "generationConfig": {
+                "temperature": 0.1,
+            },
+        }
+
+        try:
+            if stream_callback:
+                response = requests.post(endpoint, json=payload, stream=True, timeout=90)
+                if response.status_code != 200:
+                    raise RuntimeError(f"Gemini API error ({response.status_code}): {response.text}")
+
+                full_text = []
+                for line in response.iter_lines():
+                    if line:
+                        decoded = line.decode("utf-8")
+                        if decoded.startswith("data: "):
+                            raw_json = decoded[6:]
+                            try:
+                                chunk = json.loads(raw_json)
+                                candidates = chunk.get("candidates", [])
+                                if candidates:
+                                    parts = candidates[0].get("content", {}).get("parts", [])
+                                    for p in parts:
+                                        t = p.get("text", "")
+                                        full_text.append(t)
+                                        stream_callback(t)
+                            except Exception:
+                                pass
+                return "".join(full_text)
+            else:
+                response = requests.post(endpoint, json=payload, timeout=90)
+                if response.status_code == 200:
+                    data = response.json()
+                    candidates = data.get("candidates", [])
+                    if candidates:
+                        parts = candidates[0].get("content", {}).get("parts", [])
+                        return "".join(p.get("text", "") for p in parts)
+                    return ""
+                else:
+                    raise RuntimeError(f"Gemini API error ({response.status_code}): {response.text}")
+        except Exception as e:
+            raise RuntimeError(f"Gemini generation error: {str(e)}")

dbagent/llm/groq_provider.py

@@ -0,0 +1,99 @@
+"""
+Groq Free Tier LLM Provider.
+Provides ultra-fast inference with Llama 3.3, Qwen 2.5, DeepSeek via OpenAI-compatible API.
+"""
+
+import json
+import os
+import requests
+from typing import Optional, Callable, List
+from dbagent.llm.base import BaseLLMProvider
+
+
+class GroqProvider(BaseLLMProvider):
+    """Groq Cloud Provider (Free Rate-Limited Tier)."""
+
+    def __init__(self, api_key: Optional[str] = None, default_model: str = "llama-3.3-70b-versatile"):
+        self.api_key = api_key or os.getenv("GROQ_API_KEY")
+        self.default_model = default_model
+        self.endpoint = "https://api.groq.com/openai/v1/chat/completions"
+
+    @property
+    def name(self) -> str:
+        return "groq"
+
+    def is_available(self) -> bool:
+        return bool(self.api_key)
+
+    def list_models(self) -> List[str]:
+        return [
+            "llama-3.3-70b-versatile",
+            "llama-3.1-8b-instant",
+            "qwen-2.5-32b",
+            "deepseek-r1-distill-llama-70b",
+        ]
+
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        model: Optional[str] = None,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> str:
+        if not self.api_key:
+            raise ValueError(
+                "Groq API key is not configured. Set GROQ_API_KEY environment variable or run `db-agent config`."
+            )
+
+        target_model = model or self.default_model
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+        messages = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+
+        payload = {
+            "model": target_model,
+            "messages": messages,
+            "temperature": 0.1,
+            "stream": stream_callback is not None,
+        }
+
+        try:
+            if stream_callback:
+                response = requests.post(self.endpoint, headers=headers, json=payload, stream=True, timeout=60)
+                if response.status_code != 200:
+                    raise RuntimeError(f"Groq API error ({response.status_code}): {response.text}")
+
+                full_text = []
+                for line in response.iter_lines():
+                    if line:
+                        decoded = line.decode("utf-8")
+                        if decoded.startswith("data: "):
+                            raw_json = decoded[6:].strip()
+                            if raw_json == "[DONE]":
+                                break
+                            try:
+                                chunk = json.loads(raw_json)
+                                delta = chunk.get("choices", [{}])[0].get("delta", {})
+                                text_part = delta.get("content", "")
+                                if text_part:
+                                    full_text.append(text_part)
+                                    stream_callback(text_part)
+                            except Exception:
+                                pass
+                return "".join(full_text)
+            else:
+                response = requests.post(self.endpoint, headers=headers, json=payload, timeout=60)
+                if response.status_code == 200:
+                    data = response.json()
+                    return data.get("choices", [{}])[0].get("message", {}).get("content", "")
+                else:
+                    raise RuntimeError(f"Groq API error ({response.status_code}): {response.text}")
+        except Exception as e:
+            raise RuntimeError(f"Groq generation error: {str(e)}")

dbagent/llm/mock_provider.py

@@ -0,0 +1,42 @@
+"""
+Offline Rule-Based Mock Provider for zero-dependency testing and fallback.
+"""
+
+from typing import Optional, Callable, List
+from dbagent.llm.base import BaseLLMProvider
+
+
+class MockProvider(BaseLLMProvider):
+    """Fallback provider when no AI keys or local Ollama are available."""
+
+    @property
+    def name(self) -> str:
+        return "offline-mock"
+
+    def is_available(self) -> bool:
+        return True
+
+    def list_models(self) -> List[str]:
+        return ["offline-rule-based"]
+
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        model: Optional[str] = None,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> str:
+        response = (
+            "```sql\n"
+            "-- [Offline Rule-Based Generator: Connect Ollama or configure Gemini/Groq for advanced AI]\n"
+            "SELECT *\n"
+            "FROM (\n"
+            "    -- Query generated based on schema introspection\n"
+            "    SELECT 1 AS sample_result\n"
+            ");\n"
+            "```\n\n"
+            "💡 *Tip: Run `db-agent config` to configure Ollama (local offline) or free Gemini/Groq API keys for full AI capability.*"
+        )
+        if stream_callback:
+            stream_callback(response)
+        return response