dbagent-cli 0.1.0__py3-none-any.whl

dbagent/__init__.py

@@ -0,0 +1,5 @@
+"""
+DB-Agent: Universal Database Introspection and Script Generation AI Agent (CLI).
+"""
+
+__version__ = "0.1.0"

dbagent/agent/generator.py

@@ -0,0 +1,264 @@
+"""
+Agent Script Generator Engine.
+Constructs dialect-aware, schema-informed prompts and generates SQL, migrations, ETL, and backend code.
+Includes a fast SQL-only mode for the auto-execute pipeline.
+"""
+
+import re
+from typing import Optional, Callable, Dict, Any, List
+from dbagent.schema.models import DatabaseSchema
+from dbagent.schema.formatter import SchemaFormatter
+from dbagent.schema.selector import SchemaSelector
+from dbagent.llm.base import BaseLLMProvider
+
+
+# --- Full explanatory prompt (used by `generate` command) ---
+
+SYSTEM_PROMPT_TEMPLATE = """You are DB-Agent, an elite Database Engineer, SQL Architect, and Full-Stack Developer.
+Your goal is to inspect the user's database schema and generate accurate, high-performance, dialect-specific scripts, queries, migrations, or backend APIs based on the user's instructions.
+
+### Target Database Metadata:
+- **Dialect**: {dialect_name}
+- **Database Name**: {database_name}
+
+### Database Schema Context:
+{schema_context}
+
+### Core Rules:
+1. Always write queries/scripts precisely tailored for the target SQL dialect (`{dialect_name}`). Use native dialect features (e.g., PostgreSQL JSONB/CTEs/Window functions, MySQL DATE_FORMAT, SQLite strftime, etc.).
+2. Always respect table relationships, primary keys, foreign keys, and column data types defined in the schema above.
+3. When writing SQL queries, format them cleanly with uppercase keywords and descriptive aliases.
+4. When writing Python ETL or API code, use modern best practices (type annotations, SQLAlchemy 2.0 / Pydantic v2 / FastAPI / Pandas).
+5. Always enclose the primary executable script/query in a markdown code block with appropriate language tag (```sql, ```python, ```bash, etc.).
+6. Provide a concise explanation of how the script works and any performance or indexing notes.
+"""
+
+
+# --- Fast SQL-only prompt (used by pipeline `ask` command) ---
+
+FAST_SQL_SYSTEM_PROMPT = """You are DB-Agent, an expert SQL engineer.
+Generate ONLY the raw executable SQL query for the given database. No markdown fences, no explanations, no comments — just the pure SQL statement.
+
+### Target Database:
+- **Dialect**: {dialect_name}
+- **Database**: {database_name}
+
+### Schema:
+{schema_context}
+
+### Rules:
+1. Output ONLY the SQL query — nothing else. No ```sql fences, no explanations.
+2. Use the exact table and column names from the schema above.
+3. Use dialect-specific syntax for `{dialect_name}`.
+4. Write clean, optimized SQL with uppercase keywords.
+5. If the user asks for recent/time-based data, use the dialect's native datetime functions.
+6. Always include appropriate WHERE clauses, JOINs, and ORDER BY as needed.
+7. Default to LIMIT 50 unless the user specifies a count.
+"""
+
+
+# --- Retry prompt when SQL fails ---
+
+RETRY_PROMPT_TEMPLATE = """The previous SQL query failed with this error:
+```
+{error_message}
+```
+
+The failed query was:
+```sql
+{failed_sql}
+```
+
+Generate a corrected SQL query that fixes this error. Output ONLY the raw SQL — no markdown, no explanation.
+"""
+
+
+class ScriptGenerator:
+    """Core AI agent for generating code and scripts from database schema."""
+
+    def __init__(self, llm_provider: BaseLLMProvider):
+        self.llm = llm_provider
+
+    # --- User intent classification (from natural language prompt) ---
+
+    @staticmethod
+    def classify_intent(user_prompt: str) -> str:
+        """
+        Classify user's natural language intent as 'read', 'write', or 'ddl'.
+        Used before SQL generation to predict the query type.
+        """
+        prompt_lower = user_prompt.lower().strip()
+
+        # DDL indicators
+        ddl_keywords = [
+            "create table", "alter table", "drop table", "add column",
+            "remove column", "rename column", "rename table", "modify column",
+            "create index", "drop index", "truncate", "migration",
+        ]
+        for kw in ddl_keywords:
+            if kw in prompt_lower:
+                return "ddl"
+
+        # Write indicators
+        write_keywords = [
+            "insert", "add record", "add row", "add new",
+            "update", "change", "modify", "set ",
+            "delete", "remove", "drop record",
+            "increment", "decrement",
+        ]
+        for kw in write_keywords:
+            if kw in prompt_lower:
+                return "write"
+
+        # Default: read
+        return "read"
+
+    # --- Full generation (for `generate` command) ---
+
+    def generate_script(
+        self,
+        schema: DatabaseSchema,
+        user_prompt: str,
+        script_type: str = "auto",
+        model: Optional[str] = None,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> str:
+        """
+        Generate dialect-specific script based on database schema and user prompt.
+        Returns full markdown response with explanations.
+        """
+        # 1. Filter relevant tables to keep prompt compact and focused
+        relevant_tables = SchemaSelector.select_relevant_tables(schema, user_prompt, max_tables=20)
+        schema_markdown = SchemaFormatter.to_markdown(
+            schema,
+            selected_tables=relevant_tables,
+            include_samples=True,
+            max_sample_rows=2,
+        )
+
+        # 2. Build system prompt
+        system_prompt = SYSTEM_PROMPT_TEMPLATE.format(
+            dialect_name=schema.dialect_name,
+            database_name=schema.database_name,
+            schema_context=schema_markdown,
+        )
+
+        # 3. Augment user prompt with intent specifics if needed
+        augmented_prompt = user_prompt
+        if script_type == "sql":
+            augmented_prompt = f"Write an optimized SQL query for {schema.dialect_name}: {user_prompt}"
+        elif script_type == "migration":
+            augmented_prompt = f"Generate database migration script (DDL / Alembic / Flyway) for {schema.dialect_name}: {user_prompt}"
+        elif script_type == "etl":
+            augmented_prompt = f"Generate a Python ETL / Data pipeline script for {schema.dialect_name}: {user_prompt}"
+        elif script_type == "api":
+            augmented_prompt = f"Generate FastAPI / Backend REST API endpoints for this schema: {user_prompt}"
+        elif script_type == "optimization":
+            augmented_prompt = f"Analyze performance and recommend indexes / query optimizations for: {user_prompt}"
+
+        # 4. Generate response
+        return self.llm.generate(
+            prompt=augmented_prompt,
+            system_prompt=system_prompt,
+            model=model,
+            stream_callback=stream_callback,
+        )
+
+    # --- Fast SQL-only generation (for pipeline `ask` command) ---
+
+    def generate_sql_only(
+        self,
+        schema: DatabaseSchema,
+        user_prompt: str,
+        model: Optional[str] = None,
+    ) -> str:
+        """
+        Generate a raw executable SQL query — no markdown, no explanations.
+        Used by the auto-execute pipeline for speed.
+        """
+        schema_compact = SchemaFormatter.to_compact(schema)
+
+        system_prompt = FAST_SQL_SYSTEM_PROMPT.format(
+            dialect_name=schema.dialect_name,
+            database_name=schema.database_name,
+            schema_context=schema_compact,
+        )
+
+        raw_response = self.llm.generate(
+            prompt=user_prompt,
+            system_prompt=system_prompt,
+            model=model,
+        )
+
+        # Strip any accidental markdown fences the LLM may have added
+        return self._strip_markdown_fences(raw_response)
+
+    def generate_retry_sql(
+        self,
+        schema: DatabaseSchema,
+        user_prompt: str,
+        failed_sql: str,
+        error_message: str,
+        model: Optional[str] = None,
+    ) -> str:
+        """
+        Re-generate SQL after a failed execution, feeding the error back to the LLM.
+        """
+        schema_compact = SchemaFormatter.to_compact(schema)
+
+        system_prompt = FAST_SQL_SYSTEM_PROMPT.format(
+            dialect_name=schema.dialect_name,
+            database_name=schema.database_name,
+            schema_context=schema_compact,
+        )
+
+        retry_prompt = RETRY_PROMPT_TEMPLATE.format(
+            error_message=error_message,
+            failed_sql=failed_sql,
+        )
+
+        full_prompt = f"Original request: {user_prompt}\n\n{retry_prompt}"
+
+        raw_response = self.llm.generate(
+            prompt=full_prompt,
+            system_prompt=system_prompt,
+            model=model,
+        )
+
+        return self._strip_markdown_fences(raw_response)
+
+    # --- Utilities ---
+
+    @staticmethod
+    def _strip_markdown_fences(text: str) -> str:
+        """Strip markdown code fences from LLM output to get raw SQL."""
+        text = text.strip()
+        # Remove ```sql ... ``` or ``` ... ```
+        pattern = r"```(?:\w+)?\s*\n?(.*?)```"
+        match = re.search(pattern, text, re.DOTALL)
+        if match:
+            return match.group(1).strip()
+        # Remove leading/trailing ``` if present
+        if text.startswith("```"):
+            text = re.sub(r"^```\w*\n?", "", text)
+        if text.endswith("```"):
+            text = re.sub(r"\n?```$", "", text)
+        return text.strip()
+
+    @staticmethod
+    def extract_code_block(response_text: str, preferred_lang: Optional[str] = None) -> Optional[str]:
+        """Extract the main code block from LLM markdown response."""
+        # Find all code blocks ```lang ... ```
+        pattern = r"```(?:(\w+)\n)?(.*?)```"
+        matches = re.findall(pattern, response_text, re.DOTALL)
+        if not matches:
+            return None
+
+        if preferred_lang:
+            for lang, code in matches:
+                if lang.lower() == preferred_lang.lower():
+                    return code.strip()
+
+        # Return the longest code block
+        longest = max(matches, key=lambda m: len(m[1]))
+        return longest[1].strip()

dbagent/agent/pipeline.py

@@ -0,0 +1,259 @@
+"""
+QueryPipeline: Zero-friction auto-execute engine.
+
+User Prompt -> Table Resolution -> Schema Fetch -> SQL Generation ->
+Safety Check -> Auto-Execute -> Error Retry -> Results
+
+Works like a native database CLI — ask a question, get results.
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, Any, List, Optional, Callable
+from dbagent.connectors.base import BaseConnector
+from dbagent.agent.generator import ScriptGenerator
+from dbagent.agent.validator import ScriptValidator
+from dbagent.schema.models import DatabaseSchema
+from dbagent.llm.base import BaseLLMProvider
+
+
+@dataclass
+class PipelineResult:
+    """Result of a pipeline execution."""
+    # The generated SQL
+    sql: str = ""
+    # Query results
+    columns: List[str] = field(default_factory=list)
+    rows: List[Dict[str, Any]] = field(default_factory=list)
+    # Status
+    was_executed: bool = False
+    was_auto_executed: bool = False
+    needs_confirmation: bool = False
+    # Query classification
+    query_type: str = "unknown"  # read, write, ddl, unknown
+    intent: str = "read"  # from natural language
+    # Safety
+    safety_warnings: List[str] = field(default_factory=list)
+    # Error handling
+    error: Optional[str] = None
+    retries: int = 0
+    # Table resolution
+    exact_tables: List[str] = field(default_factory=list)
+    fuzzy_tables: List[str] = field(default_factory=list)
+    ambiguous: bool = False
+    ambiguous_choices: List[str] = field(default_factory=list)
+    # Schema used
+    schema: Optional[DatabaseSchema] = None
+
+    @property
+    def success(self) -> bool:
+        return self.error is None and self.was_executed
+
+    @property
+    def has_results(self) -> bool:
+        return len(self.rows) > 0
+
+
+class QueryPipeline:
+    """
+    Zero-friction query pipeline.
+    Generates SQL from natural language, validates, and auto-executes.
+    """
+
+    MAX_RETRIES = 2
+
+    def __init__(
+        self,
+        connector: BaseConnector,
+        llm: BaseLLMProvider,
+        model: Optional[str] = None,
+        auto_execute: bool = True,
+        confirm_callback: Optional[Callable[[str, List[str]], bool]] = None,
+        choice_callback: Optional[Callable[[str, List[str]], Optional[str]]] = None,
+    ):
+        """
+        Args:
+            connector: Database connector
+            llm: LLM provider for SQL generation
+            model: Specific model name
+            auto_execute: Auto-execute safe read queries
+            confirm_callback: Called for write/ddl queries — fn(sql, warnings) -> bool
+            choice_callback: Called for ambiguous tables — fn(prompt, choices) -> selected_choice or None
+        """
+        self.connector = connector
+        self.generator = ScriptGenerator(llm)
+        self.model = model
+        self.auto_execute = auto_execute
+        self.confirm_callback = confirm_callback
+        self.choice_callback = choice_callback
+
+    def run(self, user_prompt: str, force: bool = False) -> PipelineResult:
+        """
+        Full pipeline: prompt -> SQL -> execute -> results.
+
+        Steps:
+        1. Classify user intent (read/write/ddl)
+        2. Resolve tables (exact vs fuzzy match, detect ambiguity)
+        3. Fetch targeted schema on-demand
+        4. Generate SQL via LLM
+        5. Classify generated SQL
+        6. Safety check
+        7. Auto-execute or confirm
+        8. On error: retry with error context (up to MAX_RETRIES)
+        """
+        result = PipelineResult()
+        result.intent = ScriptGenerator.classify_intent(user_prompt)
+
+        # --- Step 1: Resolve tables ---
+        try:
+            exact, fuzzy = self.connector.resolve_tables(user_prompt)
+            result.exact_tables = exact
+            result.fuzzy_tables = fuzzy
+        except AttributeError:
+            # Connector doesn't support resolve_tables (e.g. old version)
+            exact, fuzzy = [], []
+
+        # Detect ambiguity: user said "users" but we have multiple fuzzy matches and no exact
+        if not exact and len(fuzzy) > 1:
+            result.ambiguous = True
+            result.ambiguous_choices = fuzzy
+            # Ask user to choose if callback is available
+            if self.choice_callback:
+                chosen = self.choice_callback(
+                    f"Multiple tables match your query. Which table(s) did you mean?",
+                    fuzzy,
+                )
+                if chosen:
+                    # User chose — treat as exact match
+                    exact = [chosen] if isinstance(chosen, str) else chosen
+                    result.ambiguous = False
+                else:
+                    # User cancelled
+                    result.error = "Ambiguous table reference — please specify the exact table name."
+                    return result
+
+        # --- Step 2: Fetch targeted schema on-demand ---
+        try:
+            schema = self.connector.inspect_targeted(
+                user_prompt=user_prompt,
+                max_tables=10,
+                include_samples=False,  # Skip samples for speed in pipeline
+            )
+            result.schema = schema
+        except Exception as e:
+            result.error = f"Schema introspection failed: {str(e)}"
+            return result
+
+        # --- Step 3: Generate SQL ---
+        try:
+            sql = self.generator.generate_sql_only(
+                schema=schema,
+                user_prompt=user_prompt,
+                model=self.model,
+            )
+            result.sql = sql
+        except Exception as e:
+            result.error = f"SQL generation failed: {str(e)}"
+            return result
+
+        if not result.sql or not result.sql.strip():
+            result.error = "LLM returned empty SQL."
+            return result
+
+        # --- Step 4: Classify & safety check ---
+        result.query_type = ScriptValidator.classify_query(result.sql)
+        is_safe, warnings = ScriptValidator.analyze_safety(result.sql)
+        result.safety_warnings = warnings
+
+        # --- Step 5: Decide execute or confirm ---
+        should_execute = False
+
+        if force:
+            should_execute = True
+        elif self.auto_execute and result.query_type == "read" and is_safe:
+            # Safe read query — auto-execute
+            should_execute = True
+            result.was_auto_executed = True
+        elif result.query_type in ("write", "ddl") or not is_safe:
+            # Needs confirmation
+            result.needs_confirmation = True
+            if self.confirm_callback:
+                confirmed = self.confirm_callback(result.sql, warnings)
+                if confirmed:
+                    should_execute = True
+                else:
+                    # User declined — return SQL without executing
+                    return result
+            else:
+                # No callback — just return the SQL for the caller to handle
+                return result
+        else:
+            # Unknown query type but safe — auto-execute
+            if self.auto_execute:
+                should_execute = True
+                result.was_auto_executed = True
+
+        # --- Step 6: Execute ---
+        if should_execute:
+            self._execute_with_retry(result, user_prompt, schema)
+
+        return result
+
+    def execute_confirmed(self, result: PipelineResult) -> PipelineResult:
+        """
+        Execute a previously generated query that was awaiting confirmation.
+        Called after user confirms a write/ddl operation.
+        """
+        if not result.sql:
+            result.error = "No SQL to execute."
+            return result
+
+        user_prompt = ""  # Not needed for direct execution
+        schema = result.schema
+
+        self._execute_with_retry(result, user_prompt, schema)
+        return result
+
+    def _execute_with_retry(
+        self,
+        result: PipelineResult,
+        user_prompt: str,
+        schema: Optional[DatabaseSchema],
+    ) -> None:
+        """Execute SQL with retry-on-error logic."""
+        current_sql = result.sql
+
+        for attempt in range(self.MAX_RETRIES + 1):
+            cols, rows, error = self.connector.execute_query(current_sql, limit=100)
+
+            if error is None:
+                # Success
+                result.sql = current_sql
+                result.columns = cols
+                result.rows = rows
+                result.was_executed = True
+                result.error = None
+                return
+
+            # Execution failed
+            result.retries = attempt + 1
+
+            if attempt < self.MAX_RETRIES and schema is not None and user_prompt:
+                # Retry: feed error back to LLM for correction
+                try:
+                    current_sql = self.generator.generate_retry_sql(
+                        schema=schema,
+                        user_prompt=user_prompt,
+                        failed_sql=current_sql,
+                        error_message=error,
+                        model=self.model,
+                    )
+                    if current_sql and current_sql.strip():
+                        result.sql = current_sql
+                        continue
+                except Exception:
+                    pass
+
+            # All retries exhausted
+            result.error = error
+            result.was_executed = False
+            return

dbagent/agent/validator.py

@@ -0,0 +1,116 @@
+"""
+SQL and script safety validator with query intent classification.
+"""
+
+import re
+from typing import Tuple, List
+
+
+class ScriptValidator:
+    """Validates SQL and script safety before execution, and classifies query intent."""
+
+    DESTRUCTIVE_KEYWORDS = [
+        r"\bDROP\s+TABLE\b",
+        r"\bDROP\s+DATABASE\b",
+        r"\bTRUNCATE\b",
+        r"\bALTER\s+TABLE\b",
+        r"\bDELETE\s+FROM\b",
+    ]
+
+    # --- Query Intent Classification ---
+
+    READ_PATTERNS = [
+        r"^\s*SELECT\b",
+        r"^\s*SHOW\b",
+        r"^\s*DESCRIBE\b",
+        r"^\s*DESC\b",
+        r"^\s*EXPLAIN\b",
+        r"^\s*PRAGMA\b",
+        r"^\s*WITH\b.*\bSELECT\b",
+    ]
+
+    WRITE_PATTERNS = [
+        r"^\s*INSERT\b",
+        r"^\s*UPDATE\b",
+        r"^\s*DELETE\b",
+        r"^\s*REPLACE\b",
+        r"^\s*MERGE\b",
+        r"^\s*UPSERT\b",
+    ]
+
+    DDL_PATTERNS = [
+        r"^\s*CREATE\b",
+        r"^\s*ALTER\b",
+        r"^\s*DROP\b",
+        r"^\s*TRUNCATE\b",
+        r"^\s*RENAME\b",
+        r"^\s*GRANT\b",
+        r"^\s*REVOKE\b",
+    ]
+
+    @classmethod
+    def classify_query(cls, sql: str) -> str:
+        """
+        Classify a SQL query as 'read', 'write', 'ddl', or 'unknown'.
+        Used by the pipeline to decide auto-execute vs confirmation.
+        """
+        if not sql or not sql.strip():
+            return "unknown"
+
+        # Strip leading comments and whitespace
+        cleaned = re.sub(r"--[^\n]*\n?", "", sql).strip()
+        cleaned = re.sub(r"/\*.*?\*/", "", cleaned, flags=re.DOTALL).strip()
+
+        if not cleaned:
+            return "unknown"
+
+        for pattern in cls.READ_PATTERNS:
+            if re.search(pattern, cleaned, re.IGNORECASE | re.DOTALL):
+                return "read"
+
+        for pattern in cls.DDL_PATTERNS:
+            if re.search(pattern, cleaned, re.IGNORECASE):
+                return "ddl"
+
+        for pattern in cls.WRITE_PATTERNS:
+            if re.search(pattern, cleaned, re.IGNORECASE):
+                return "write"
+
+        return "unknown"
+
+    @classmethod
+    def analyze_safety(cls, query: str) -> Tuple[bool, List[str]]:
+        """
+        Analyze whether a query contains destructive or modifying statements.
+        Returns (is_safe_read_only, list_of_warnings).
+        """
+        warnings = []
+        q_upper = query.upper()
+
+        for pattern in cls.DESTRUCTIVE_KEYWORDS:
+            if re.search(pattern, q_upper):
+                clean_cmd = pattern.replace(r"\b", "").replace(r"\s+", " ")
+                warnings.append(f"Contains potentially destructive command: '{clean_cmd}'")
+
+        # Check for DELETE without WHERE
+        if "DELETE FROM" in q_upper and "WHERE" not in q_upper:
+            warnings.append("DELETE statement without WHERE clause detected (deletes entire table)!")
+
+        # Check for UPDATE without WHERE
+        if "UPDATE " in q_upper and "WHERE" not in q_upper:
+            warnings.append("UPDATE statement without WHERE clause detected (modifies all rows)!")
+
+        is_safe = len(warnings) == 0
+        return is_safe, warnings
+
+    @classmethod
+    def is_auto_executable(cls, sql: str) -> bool:
+        """
+        Quick check: can this query be auto-executed without user confirmation?
+        True only for pure read queries with no safety warnings.
+        """
+        query_type = cls.classify_query(sql)
+        if query_type != "read":
+            return False
+        is_safe, _ = cls.analyze_safety(sql)
+        return is_safe