dbagent-cli 0.1.0__py3-none-any.whl

dbagent/llm/ollama_provider.py

@@ -0,0 +1,145 @@
+"""
+Local Ollama LLM provider (100% Free, Offline, Zero-cost).
+Auto-installs and configures Ollama on first use — fully independent.
+"""
+
+import json
+import requests
+from typing import Optional, Callable, List
+from dbagent.llm.base import BaseLLMProvider
+
+
+class OllamaProvider(BaseLLMProvider):
+    """Local offline provider using Ollama. Auto-sets up on first use."""
+
+    def __init__(self, base_url: str = "http://localhost:11434", default_model: str = "qwen2.5-coder:1.5b"):
+        self.base_url = base_url.rstrip("/")
+        self.default_model = default_model
+        self._setup_done = False
+
+    @property
+    def name(self) -> str:
+        return "ollama"
+
+    def is_available(self) -> bool:
+        """Check if local Ollama daemon is running."""
+        try:
+            r = requests.get(f"{self.base_url}/api/tags", timeout=0.5)
+            return r.status_code == 200
+        except Exception:
+            return False
+
+    def list_models(self) -> List[str]:
+        """List all installed local models."""
+        try:
+            r = requests.get(f"{self.base_url}/api/tags", timeout=3)
+            if r.status_code == 200:
+                data = r.json()
+                return [m.get("name") for m in data.get("models", [])]
+        except Exception:
+            pass
+        return []
+
+    def _ensure_ready(self, print_fn=None) -> None:
+        """Auto-install, start, and pull model if needed. Called once per session."""
+        if self._setup_done and self.is_available():
+            return
+
+        from dbagent.llm.auto_setup import ensure_ollama_ready, get_best_available_model
+
+        success, msg = ensure_ollama_ready(
+            model=self.default_model,
+            print_fn=print_fn,
+        )
+
+        if not success:
+            raise ConnectionError(msg)
+
+        # Use the best available model
+        best = get_best_available_model()
+        if best:
+            self.default_model = best
+
+        self._setup_done = True
+
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        model: Optional[str] = None,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> str:
+        # Auto-setup on first generate call
+        if not self.is_available():
+            # Build a print function that uses Rich console if available
+            def _print_fn(msg, level="info"):
+                try:
+                    from dbagent.ui.console import print_info, print_success, print_warning, print_error
+                    if level == "success":
+                        print_success(msg)
+                    elif level == "warning":
+                        print_warning(msg)
+                    elif level == "error":
+                        print_error(msg)
+                    else:
+                        print_info(msg)
+                except Exception:
+                    print(f"[{level.upper()}] {msg}")
+
+            self._ensure_ready(print_fn=_print_fn)
+
+        target_model = model or self.default_model
+
+        # Check if requested model exists locally, pick best available
+        available = self.list_models()
+        if available and target_model not in available:
+            # Try partial match
+            for m in available:
+                if target_model.split(":")[0] in m:
+                    target_model = m
+                    break
+            else:
+                target_model = available[0]
+
+        payload = {
+            "model": target_model,
+            "prompt": prompt,
+            "system": system_prompt or "You are an expert database AI engineer.",
+            "stream": stream_callback is not None,
+            "options": {
+                "temperature": 0.1,
+            },
+        }
+
+        try:
+            if stream_callback:
+                response = requests.post(
+                    f"{self.base_url}/api/generate",
+                    json=payload,
+                    stream=True,
+                    timeout=120,
+                )
+                full_text = []
+                for line in response.iter_lines():
+                    if line:
+                        chunk = json.loads(line)
+                        text_part = chunk.get("response", "")
+                        full_text.append(text_part)
+                        stream_callback(text_part)
+                return "".join(full_text)
+            else:
+                response = requests.post(
+                    f"{self.base_url}/api/generate",
+                    json=payload,
+                    timeout=120,
+                )
+                if response.status_code == 200:
+                    return response.json().get("response", "")
+                else:
+                    raise RuntimeError(f"Ollama error ({response.status_code}): {response.text}")
+        except requests.exceptions.ConnectionError:
+            raise ConnectionError(
+                "Could not connect to Ollama. Try running `ollama serve` manually."
+            )
+        except Exception as e:
+            raise RuntimeError(f"Ollama generation failed: {str(e)}")

dbagent/llm/openrouter_provider.py

@@ -0,0 +1,100 @@
+"""
+OpenRouter Provider for free tier models (e.g., Llama 3.3 70B, Qwen 2.5 Coder 32B Free).
+"""
+
+import json
+import os
+import requests
+from typing import Optional, Callable, List
+from dbagent.llm.base import BaseLLMProvider
+
+
+class OpenRouterProvider(BaseLLMProvider):
+    """OpenRouter Cloud Provider with access to free-tier open models."""
+
+    def __init__(self, api_key: Optional[str] = None, default_model: str = "meta-llama/llama-3.3-70b-instruct:free"):
+        self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
+        self.default_model = default_model
+        self.endpoint = "https://openrouter.ai/api/v1/chat/completions"
+
+    @property
+    def name(self) -> str:
+        return "openrouter"
+
+    def is_available(self) -> bool:
+        return bool(self.api_key)
+
+    def list_models(self) -> List[str]:
+        return [
+            "meta-llama/llama-3.3-70b-instruct:free",
+            "qwen/qwen-2.5-coder-32b-instruct:free",
+            "mistralai/mistral-7b-instruct:free",
+            "google/gemini-2.0-flash-exp:free",
+        ]
+
+    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(
+                "OpenRouter API key is not configured. Set OPENROUTER_API_KEY environment variable or run `db-agent config`."
+            )
+
+        target_model = model or self.default_model
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "HTTP-Referer": "https://github.com/db-agent/db-agent",
+            "X-Title": "DB-Agent",
+            "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"OpenRouter 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"OpenRouter API error ({response.status_code}): {response.text}")
+        except Exception as e:
+            raise RuntimeError(f"OpenRouter generation error: {str(e)}")

dbagent/schema/formatter.py

@@ -0,0 +1,123 @@
+"""
+Schema formatters for converting database metadata into rich LLM context.
+"""
+
+from typing import List, Optional
+from dbagent.schema.models import DatabaseSchema, TableModel
+
+
+class SchemaFormatter:
+    """Formats DatabaseSchema into prompts for LLMs."""
+
+    @staticmethod
+    def to_markdown(
+        schema: DatabaseSchema,
+        selected_tables: Optional[List[str]] = None,
+        include_samples: bool = True,
+        max_sample_rows: int = 2,
+    ) -> str:
+        """Render schema into a comprehensive markdown description."""
+        lines = []
+        lines.append(f"# Database Catalog: {schema.database_name}")
+        lines.append(f"- **Dialect**: `{schema.dialect_name}`")
+        if schema.server_version:
+            lines.append(f"- **Server Version**: {schema.server_version}")
+        lines.append(f"- **Total Tables**: {schema.total_tables}")
+        lines.append(f"- **Total Views**: {schema.total_views}")
+        lines.append("")
+
+        tables_to_include = schema.tables
+        if selected_tables is not None:
+            lower_selected = [t.lower() for t in selected_tables]
+            tables_to_include = [t for t in schema.tables if t.name.lower() in lower_selected]
+
+        for table in tables_to_include:
+            table_type = "View" if table.is_view else "Table"
+            lines.append(f"## {table_type}: `{table.name}`")
+            if table.comment:
+                lines.append(f"*{table.comment}*")
+            if table.row_count is not None:
+                lines.append(f"- Approximate Row Count: {table.row_count:,}")
+            lines.append("")
+
+            # Columns table
+            lines.append("| Column | Type | Nullable | PK | Default | Comment |")
+            lines.append("|---|---|:---:|:---:|---|---|")
+            for col in table.columns:
+                pk_mark = "✅ PK" if col.is_primary_key else ""
+                null_mark = "NULL" if col.is_nullable else "NOT NULL"
+                default_val = f"`{col.default_value}`" if col.default_value else ""
+                comment = col.comment or ""
+                lines.append(
+                    f"| `{col.name}` | `{col.data_type}` | {null_mark} | {pk_mark} | {default_val} | {comment} |"
+                )
+            lines.append("")
+
+            # Foreign Keys
+            if table.foreign_keys:
+                lines.append("### Foreign Key Relationships")
+                for fk in table.foreign_keys:
+                    cols = ", ".join(f"`{c}`" for c in fk.constrained_columns)
+                    ref_cols = ", ".join(f"`{c}`" for c in fk.referred_columns)
+                    lines.append(f"- ({cols}) ➡️ `{fk.referred_table}`({ref_cols})")
+                lines.append("")
+
+            # Indexes
+            if table.indexes:
+                lines.append("### Indexes")
+                for idx in table.indexes:
+                    idx_type = "UNIQUE " if idx.is_unique else ""
+                    cols = ", ".join(f"`{c}`" for c in idx.columns)
+                    lines.append(f"- {idx_type}INDEX `{idx.name}` ({cols})")
+                lines.append("")
+
+            # View definition
+            if table.is_view and table.view_definition:
+                lines.append("### View Definition")
+                lines.append("```sql")
+                lines.append(table.view_definition.strip())
+                lines.append("```")
+                lines.append("")
+
+            # Sample Rows
+            if include_samples and table.sample_rows:
+                lines.append("### Sample Data (Shape & Types)")
+                samples = table.sample_rows[:max_sample_rows]
+                if samples:
+                    headers = list(samples[0].keys())
+                    lines.append("| " + " | ".join(headers) + " |")
+                    lines.append("| " + " | ".join(["---"] * len(headers)) + " |")
+                    for row in samples:
+                        row_vals = []
+                        for h in headers:
+                            val = row.get(h)
+                            if val is None:
+                                row_vals.append("`NULL`")
+                            else:
+                                val_str = str(val).replace("\n", " ")
+                                if len(val_str) > 30:
+                                    val_str = val_str[:27] + "..."
+                                row_vals.append(f"`{val_str}`")
+                        lines.append("| " + " | ".join(row_vals) + " |")
+                lines.append("")
+
+            lines.append("---")
+            lines.append("")
+
+        return "\n".join(lines)
+
+    @staticmethod
+    def to_compact(schema: DatabaseSchema) -> str:
+        """Render a compact single-line per table summary (for huge databases)."""
+        lines = [f"Database: {schema.database_name} ({schema.dialect_name})"]
+        for t in schema.tables:
+            cols = []
+            for c in t.columns:
+                pk = " [PK]" if c.is_primary_key else ""
+                cols.append(f"{c.name}: {c.data_type}{pk}")
+            fks = []
+            for fk in t.foreign_keys:
+                fks.append(f"{','.join(fk.constrained_columns)}->{fk.referred_table}({','.join(fk.referred_columns)})")
+            fk_str = f" | FK: {'; '.join(fks)}" if fks else ""
+            lines.append(f"- {t.name} ({', '.join(cols)}){fk_str}")
+        return "\n".join(lines)

dbagent/schema/models.py

@@ -0,0 +1,72 @@
+"""
+Schema catalog models for structured representation of databases.
+"""
+
+from typing import List, Dict, Any, Optional
+from pydantic import BaseModel, Field
+
+
+class ColumnModel(BaseModel):
+    """Represents a database table column."""
+    name: str
+    data_type: str
+    is_nullable: bool = True
+    is_primary_key: bool = False
+    default_value: Optional[str] = None
+    comment: Optional[str] = None
+    is_autoincrement: bool = False
+
+
+class ForeignKeyModel(BaseModel):
+    """Represents a foreign key constraint."""
+    name: Optional[str] = None
+    constrained_columns: List[str]
+    referred_table: str
+    referred_columns: List[str]
+
+
+class IndexModel(BaseModel):
+    """Represents a database index."""
+    name: str
+    columns: List[str]
+    is_unique: bool = False
+
+
+class TableModel(BaseModel):
+    """Represents a database table or view."""
+    name: str
+    schema_name: Optional[str] = None
+    is_view: bool = False
+    columns: List[ColumnModel] = Field(default_factory=list)
+    primary_key: List[str] = Field(default_factory=list)
+    foreign_keys: List[ForeignKeyModel] = Field(default_factory=list)
+    indexes: List[IndexModel] = Field(default_factory=list)
+    row_count: Optional[int] = None
+    sample_rows: List[Dict[str, Any]] = Field(default_factory=list)
+    view_definition: Optional[str] = None
+    comment: Optional[str] = None
+
+
+class DatabaseSchema(BaseModel):
+    """Complete catalog of an introspected database."""
+    dialect_name: str
+    database_name: str
+    server_version: Optional[str] = None
+    tables: List[TableModel] = Field(default_factory=list)
+
+    @property
+    def total_tables(self) -> int:
+        return len([t for t in self.tables if not t.is_view])
+
+    @property
+    def total_views(self) -> int:
+        return len([t for t in self.tables if t.is_view])
+
+    def get_table(self, table_name: str) -> Optional[TableModel]:
+        for t in self.tables:
+            if t.name.lower() == table_name.lower():
+                return t
+        return None
+
+    def get_table_names(self) -> List[str]:
+        return [t.name for t in self.tables]

dbagent/schema/selector.py

@@ -0,0 +1,70 @@
+"""
+Smart schema selector for filtering relevant tables in large databases.
+"""
+
+import re
+from typing import List, Set
+from dbagent.schema.models import DatabaseSchema, TableModel
+
+
+class SchemaSelector:
+    """Selects subset of relevant tables from a large schema based on a user prompt."""
+
+    @classmethod
+    def select_relevant_tables(
+        cls,
+        schema: DatabaseSchema,
+        user_prompt: str,
+        max_tables: int = 15,
+    ) -> List[str]:
+        """
+        Extract relevant table names from user query and include connected foreign key tables.
+        If total tables in database <= max_tables, returns all tables.
+        """
+        all_table_names = [t.name for t in schema.tables]
+        if len(all_table_names) <= max_tables:
+            return all_table_names
+
+        prompt_lower = user_prompt.lower()
+        matched_tables: Set[str] = set()
+
+        # 1. Direct match or word boundary match
+        for table in schema.tables:
+            t_name = table.name.lower()
+            # Check exact word or substring
+            pattern = r"\b" + re.escape(t_name) + r"\b"
+            if re.search(pattern, prompt_lower) or t_name in prompt_lower:
+                matched_tables.add(table.name)
+                continue
+
+            # Check if any column name is specifically mentioned
+            for col in table.columns:
+                c_name = col.name.lower()
+                if len(c_name) > 3 and re.search(r"\b" + re.escape(c_name) + r"\b", prompt_lower):
+                    matched_tables.add(table.name)
+                    break
+
+        # 2. Add foreign-key related tables for matched tables
+        related_tables: Set[str] = set()
+        for t_name in matched_tables:
+            t_model = schema.get_table(t_name)
+            if t_model:
+                for fk in t_model.foreign_keys:
+                    if fk.referred_table in all_table_names:
+                        related_tables.add(fk.referred_table)
+
+        # Also find any other table that references the matched table
+        for table in schema.tables:
+            for fk in table.foreign_keys:
+                if fk.referred_table in matched_tables:
+                    related_tables.add(table.name)
+
+        combined = matched_tables.union(related_tables)
+
+        # If nothing matched (e.g. general prompt like "list everything"), return top tables
+        if not combined:
+            return all_table_names[:max_tables]
+
+        # Convert to list capped at max_tables
+        result = list(combined)[:max_tables]
+        return result

dbagent/ui/console.py

@@ -0,0 +1,77 @@
+"""
+Rich terminal UI utilities and console formatting.
+"""
+
+import sys
+from rich.console import Console
+from rich.panel import Panel
+from rich.syntax import Syntax
+from rich.text import Text
+from rich.table import Table
+from typing import Optional, List, Dict, Any
+
+# Ensure UTF-8 on Windows terminal streams
+if sys.platform == "win32":
+    try:
+        if hasattr(sys.stdout, "reconfigure"):
+            sys.stdout.reconfigure(encoding="utf-8")
+        if hasattr(sys.stderr, "reconfigure"):
+            sys.stderr.reconfigure(encoding="utf-8")
+    except Exception:
+        pass
+
+console = Console()
+
+
+def print_banner() -> None:
+    """Print the DB-Agent terminal banner."""
+    banner_text = Text()
+    banner_text.append("[*] DB-AGENT ", style="bold cyan")
+    banner_text.append("- Universal Database Introspector & AI Script Generator\n", style="bold white")
+    banner_text.append("Supports PostgreSQL | MySQL | SQLite | MSSQL | Oracle | DuckDB | MongoDB", style="dim")
+    panel = Panel(banner_text, border_style="cyan", padding=(0, 2))
+    console.print(panel)
+
+
+def print_success(message: str) -> None:
+    console.print(f"[bold green][OK][/bold green] {message}")
+
+
+def print_error(message: str) -> None:
+    console.print(f"[bold red][ERROR][/bold red] {message}")
+
+
+def print_warning(message: str) -> None:
+    console.print(f"[bold yellow][WARN][/bold yellow] {message}")
+
+
+def print_info(message: str) -> None:
+    console.print(f"[bold blue][INFO][/bold blue] {message}")
+
+
+def print_code(code: str, language: str = "sql", title: Optional[str] = None) -> None:
+    """Print syntax-highlighted code block."""
+    syntax = Syntax(code, language, theme="monokai", line_numbers=True)
+    panel = Panel(syntax, title=f"[bold cyan]{title or language.upper()}[/bold cyan]", border_style="dim")
+    console.print(panel)
+
+
+def print_results_table(columns: List[str], rows: List[Dict[str, Any]], title: str = "Query Results") -> None:
+    """Print query results in a clean Rich table."""
+    if not columns:
+        console.print("[dim]Query returned no rows.[/dim]")
+        return
+
+    table = Table(title=title, show_header=True, header_style="bold magenta", border_style="dim")
+    for col in columns:
+        table.add_column(str(col))
+
+    for row in rows:
+        row_vals = []
+        for col in columns:
+            v = row.get(col)
+            row_vals.append(str(v) if v is not None else "[dim]NULL[/dim]")
+        table.add_row(*row_vals)
+
+    console.print(table)
+    console.print(f"[dim]Total rows: {len(rows)}[/dim]\n")

dbagent/ui/viewer.py

@@ -0,0 +1,93 @@
+"""
+Rich terminal viewer for database schema introspection.
+"""
+
+from rich.table import Table
+from rich.tree import Tree
+from rich.panel import Panel
+from rich.console import Console
+from typing import Optional
+from dbagent.schema.models import DatabaseSchema, TableModel
+from dbagent.ui.console import console
+
+
+class SchemaViewer:
+    """Renders database schema visual representations in the terminal."""
+
+    @staticmethod
+    def display_schema_summary(schema: DatabaseSchema) -> None:
+        """Display an overview table of all discovered tables and views."""
+        table = Table(
+            title=f"[Catalog] Database: [bold cyan]{schema.database_name}[/bold cyan] ({schema.dialect_name})",
+            show_header=True,
+            header_style="bold cyan",
+            border_style="dim",
+        )
+        table.add_column("Type", justify="center", width=8)
+        table.add_column("Table Name", style="bold white")
+        table.add_column("Columns", justify="right")
+        table.add_column("Primary Key", style="yellow")
+        table.add_column("Foreign Keys", style="green")
+        table.add_column("Est. Rows", justify="right", style="cyan")
+
+        for t in schema.tables:
+            type_tag = "[cyan]View[/cyan]" if t.is_view else "[white]Table[/white]"
+            pk_str = ", ".join(t.primary_key) if t.primary_key else "[dim]None[/dim]"
+            fk_count = len(t.foreign_keys)
+            fk_str = f"{fk_count} relations" if fk_count > 0 else "[dim]None[/dim]"
+            rows_str = f"{t.row_count:,}" if t.row_count is not None else "[dim]N/A[/dim]"
+
+            table.add_row(
+                type_tag,
+                t.name,
+                str(len(t.columns)),
+                pk_str,
+                fk_str,
+                rows_str,
+            )
+
+        console.print(table)
+        console.print(
+            f"[dim]Discovered [bold]{schema.total_tables}[/bold] tables and [bold]{schema.total_views}[/bold] views.[/dim]\n"
+        )
+
+    @staticmethod
+    def display_table_detail(table: TableModel) -> None:
+        """Display detailed column attributes for a single table."""
+        tbl = Table(
+            title=f"Table Structure: [bold cyan]{table.name}[/bold cyan]",
+            show_header=True,
+            header_style="bold magenta",
+            border_style="dim",
+        )
+        tbl.add_column("Column", style="bold white")
+        tbl.add_column("Type", style="cyan")
+        tbl.add_column("PK", justify="center", style="yellow")
+        tbl.add_column("Nullable", justify="center")
+        tbl.add_column("Default", style="dim")
+        tbl.add_column("Comment", style="italic dim")
+
+        for col in table.columns:
+            pk_str = "[PK]" if col.is_primary_key else ""
+            null_str = "[green]YES[/green]" if col.is_nullable else "[red]NO[/red]"
+            def_str = col.default_value or ""
+            comm_str = col.comment or ""
+            tbl.add_row(col.name, col.data_type, pk_str, null_str, def_str, comm_str)
+
+        console.print(tbl)
+
+        if table.foreign_keys:
+            console.print("[bold green]Foreign Keys:[/bold green]")
+            for fk in table.foreign_keys:
+                c = ", ".join(fk.constrained_columns)
+                r = ", ".join(fk.referred_columns)
+                console.print(f"  * ({c}) -> [bold]{fk.referred_table}[/bold]({r})")
+            console.print("")
+
+        if table.indexes:
+            console.print("[bold blue]Indexes:[/bold blue]")
+            for idx in table.indexes:
+                u = "[yellow]UNIQUE [/yellow]" if idx.is_unique else ""
+                c = ", ".join(idx.columns)
+                console.print(f"  * {u}{idx.name} on ({c})")
+            console.print("")