pgnode 0.1.0__py3-none-any.whl

app/__init__.py

@@ -0,0 +1 @@
+"""Application package."""

app/agent/__init__.py

@@ -0,0 +1 @@
+"""Agent orchestration: NL to validated SQL execution."""

app/agent/agent.py

@@ -0,0 +1,216 @@
+"""MVP agent: schema -> LLM SQL -> validate -> execute."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from app.agent.intent_agent import run_intent_agent
+from app.agent.sql_metadata import extract_sql_metadata
+from app.agent.sql_quality import detect_quality_issue
+from app.agent.sql_quoter import quote_schema_identifiers
+from app.agent.validator import validate_sql
+from app.db.query_executor import run_query
+from app.db.schema_loader import get_full_schema
+from app.llm.prompt_builder import strip_sql_fences
+from app.llm.sql_generator import generate_sql, repair_sql
+
+_MAX_SQL_RETRIES = 1
+
+
+@dataclass
+class AgentResult:
+    """Outcome of a single agent turn."""
+
+    user_prompt: str
+    sql: str | None
+    validation_error: str | None
+    execution_result: list[dict[str, Any]] | str | None
+    llm_error: str | None = None
+    operation: str | None = None
+    table: str | None = None
+    changed_columns: list[str] | None = None
+    where_clause: str | None = None
+    out_of_scope: bool = False
+
+
+def run_agent(
+    user_prompt: str,
+    *,
+    dry_run: bool = False,
+    model: str | None = None,
+    conversation_context: str | None = None,
+    forced_sql: str | None = None,
+) -> AgentResult:
+    """
+    Load schema, generate SQL via local LLM, validate, optionally execute.
+
+    If dry_run is True, the query executor returns a dry-run message without executing.
+    """
+    try:
+        schema = get_full_schema()
+    except RuntimeError as exc:
+        return AgentResult(
+            user_prompt=user_prompt,
+            sql=None,
+            validation_error=str(exc),
+            execution_result=None,
+            llm_error=None,
+        )
+
+    try:
+        intent = run_intent_agent(
+            user_prompt,
+            schema,
+            model=model,
+            conversation_context=conversation_context,
+        )
+    except RuntimeError as exc:
+        return AgentResult(
+            user_prompt=user_prompt,
+            sql=None,
+            validation_error=None,
+            execution_result=None,
+            llm_error=str(exc),
+        )
+
+    if intent.intent == "out_of_scope":
+        return AgentResult(
+            user_prompt=user_prompt,
+            sql=None,
+            validation_error=None,
+            execution_result=intent.response,
+            out_of_scope=True,
+        )
+
+    if forced_sql is not None:
+        sql = quote_schema_identifiers(strip_sql_fences(forced_sql).strip(), schema)
+    else:
+        try:
+            raw_sql = generate_sql(
+                user_prompt,
+                schema,
+                model=model,
+                conversation_context=conversation_context,
+            )
+        except RuntimeError as exc:
+            return AgentResult(
+                user_prompt=user_prompt,
+                sql=None,
+                validation_error=None,
+                execution_result=None,
+                llm_error=str(exc),
+            )
+
+        sql = quote_schema_identifiers(strip_sql_fences(raw_sql).strip(), schema)
+    sql = _maybe_repair_sql_quality(
+        user_prompt=user_prompt,
+        sql=sql,
+        schema=schema,
+        model=model,
+        conversation_context=conversation_context,
+    )
+    metadata = extract_sql_metadata(sql)
+
+    validation_error = validate_sql(sql)
+    if validation_error is not None:
+        return AgentResult(
+            user_prompt=user_prompt,
+            sql=sql,
+            validation_error=validation_error,
+            execution_result=None,
+            operation=metadata.operation,
+            table=metadata.table,
+            changed_columns=metadata.changed_columns,
+            where_clause=metadata.where_clause,
+        )
+
+    execution_result = run_query(sql, dry_run=dry_run)
+    attempts = 0
+    while (
+        not dry_run
+        and isinstance(execution_result, str)
+        and execution_result.startswith("Query execution failed:")
+        and attempts < _MAX_SQL_RETRIES
+    ):
+        attempts += 1
+        try:
+            repaired_sql = repair_sql(
+                user_request=user_prompt,
+                schema=schema,
+                previous_sql=sql,
+                db_error=execution_result,
+                model=model,
+                conversation_context=conversation_context,
+            ).strip()
+        except RuntimeError as exc:
+            return AgentResult(
+                user_prompt=user_prompt,
+                sql=sql,
+                validation_error=None,
+                execution_result=execution_result,
+                llm_error=f"Repair attempt failed: {exc}",
+                operation=metadata.operation,
+                table=metadata.table,
+                changed_columns=metadata.changed_columns,
+                where_clause=metadata.where_clause,
+            )
+
+        repair_validation_error = validate_sql(repaired_sql)
+        if repair_validation_error is not None:
+            return AgentResult(
+                user_prompt=user_prompt,
+                sql=repaired_sql,
+                validation_error=repair_validation_error,
+                execution_result=execution_result,
+                operation=metadata.operation,
+                table=metadata.table,
+                changed_columns=metadata.changed_columns,
+                where_clause=metadata.where_clause,
+            )
+
+        sql = quote_schema_identifiers(repaired_sql, schema)
+        metadata = extract_sql_metadata(sql)
+        execution_result = run_query(sql, dry_run=False)
+
+    return AgentResult(
+        user_prompt=user_prompt,
+        sql=sql,
+        validation_error=None,
+        execution_result=execution_result,
+        operation=metadata.operation,
+        table=metadata.table,
+        changed_columns=metadata.changed_columns,
+        where_clause=metadata.where_clause,
+    )
+
+
+def _maybe_repair_sql_quality(
+    *,
+    user_prompt: str,
+    sql: str,
+    schema: dict[str, list[dict[str, str]]],
+    model: str | None,
+    conversation_context: str | None,
+) -> str:
+    """Run one quality-focused repair when generated SQL violates browse/lookup heuristics."""
+    quality_issue = detect_quality_issue(user_prompt, sql)
+    if quality_issue is None:
+        return sql
+
+    try:
+        repaired_sql = repair_sql(
+            user_request=user_prompt,
+            schema=schema,
+            previous_sql=sql,
+            db_error=f"Quality issue: {quality_issue}",
+            model=model,
+            conversation_context=conversation_context,
+        ).strip()
+    except RuntimeError:
+        return sql
+
+    if validate_sql(repaired_sql) is not None:
+        return sql
+
+    return quote_schema_identifiers(repaired_sql, schema)

app/agent/followup.py

@@ -0,0 +1,127 @@
+"""Follow-up prompt grounding after writes and other contextual turns."""
+
+from __future__ import annotations
+
+import re
+
+from app.agent.memory import ConversationMemory, MemoryTurn
+from app.agent.sql_metadata import extract_sql_metadata
+
+_SHOW_VERBS = re.compile(
+    r"\b(show|display|list|see|view|fetch|get)\b",
+    re.IGNORECASE,
+)
+_AFTER_WRITE_HINTS = re.compile(
+    r"\b(updated|update|modified|changes?|changed|result|after)\b",
+    re.IGNORECASE,
+)
+_ANAPHORA = re.compile(
+    r"\b(that|it|same|previous|those|the change|the update)\b",
+    re.IGNORECASE,
+)
+_QUOTED_IDENT = r'"([^"]+)"'
+_BARE_IDENT = r"([a-zA-Z_][\w]*)"
+
+
+def is_show_after_write_prompt(text: str) -> bool:
+    """True when the user likely wants to read back data after a write."""
+    lowered = text.lower()
+    if not _SHOW_VERBS.search(lowered):
+        return False
+    if _AFTER_WRITE_HINTS.search(lowered):
+        return True
+    return "table" in lowered or "rows" in lowered or "record" in lowered
+
+
+def try_build_show_after_write_sql(
+    user_prompt: str,
+    memory: ConversationMemory | None,
+) -> str | None:
+    """
+    Build a SELECT for the table/row affected by the most recent write.
+
+    Returns None when this follow-up pattern does not apply.
+    """
+    if memory is None or not is_show_after_write_prompt(user_prompt):
+        return None
+
+    last_write = _find_last_write_turn(memory)
+    if last_write is None or not last_write.sql:
+        return None
+
+    metadata = extract_sql_metadata(last_write.sql)
+    table = metadata.table or _table_from_sql(last_write.sql)
+    if not table:
+        return None
+
+    table_ref = _quote_table_ref(table)
+    if metadata.where_clause:
+        return f"SELECT * FROM {table_ref} WHERE {metadata.where_clause} LIMIT 10;"
+    return f"SELECT * FROM {table_ref} LIMIT 10;"
+
+
+def ground_followup_prompt(text: str, memory: ConversationMemory | None) -> str:
+    """Enrich vague follow-ups using recent conversation turns."""
+    if memory is None:
+        return text
+
+    if is_show_after_write_prompt(text):
+        last_write = _find_last_write_turn(memory)
+        if last_write is not None:
+            metadata = extract_sql_metadata(last_write.sql or "")
+            table = metadata.table or _table_from_sql(last_write.sql or "") or "unknown"
+            where = metadata.where_clause or ""
+            cols = ", ".join(last_write.changed_columns) if last_write.changed_columns else "*"
+            return (
+                f"{text}\n\n"
+                "Context from previous turn:\n"
+                f"- Previous operation: {last_write.operation}\n"
+                f"- Table: {table}\n"
+                f"- Changed columns: {cols}\n"
+                f"- Previous filter: {where or '<none>'}\n"
+                "The user wants to VIEW current rows after that write.\n"
+                "MANDATORY: Generate a SELECT query only (NOT UPDATE/INSERT/DELETE).\n"
+                f"Prefer: SELECT * FROM {_quote_table_ref(table)}"
+                + (f" WHERE {where}" if where else "")
+                + " LIMIT 10;"
+            )
+
+    if not _ANAPHORA.search(text):
+        return text
+
+    for turn in reversed(memory.recent(10)):
+        if turn.operation in {"UPDATE", "INSERT"} and (turn.table or turn.sql):
+            metadata = extract_sql_metadata(turn.sql or "")
+            table = metadata.table or turn.table or "unknown"
+            cols = ", ".join(turn.changed_columns) if turn.changed_columns else "*"
+            where = metadata.where_clause or turn.where_clause or ""
+            return (
+                f"{text}\n"
+                "Context from previous turn:\n"
+                f"- Last {turn.operation} table: {table}\n"
+                f"- Changed columns: {cols}\n"
+                f"- Filter: {where or '<none>'}\n"
+                "If user asks to show the change, produce SELECT for that table/row."
+            )
+
+    return text
+
+
+def _find_last_write_turn(memory: ConversationMemory) -> MemoryTurn | None:
+    for turn in reversed(memory.recent(10)):
+        if turn.operation in {"UPDATE", "INSERT"}:
+            return turn
+    return None
+
+
+def _table_from_sql(sql: str) -> str | None:
+    metadata = extract_sql_metadata(sql)
+    return metadata.table
+
+
+def _quote_table_ref(table: str) -> str:
+    if table.startswith('"') and table.endswith('"'):
+        return table
+    if table != table.lower():
+        return f'"{table}"'
+    return table

app/agent/intent_agent.py

@@ -0,0 +1,21 @@
+"""Intent Agent wrapper: classify user input before SQL generation."""
+
+from __future__ import annotations
+
+from app.llm.intent_classifier import IntentClassification, classify_intent
+
+
+def run_intent_agent(
+    user_prompt: str,
+    schema: dict[str, list[dict[str, str]]],
+    *,
+    model: str | None = None,
+    conversation_context: str | None = None,
+) -> IntentClassification:
+    """Classify whether the user message should reach the SQL generator agent."""
+    return classify_intent(
+        user_prompt,
+        schema,
+        model=model,
+        conversation_context=conversation_context,
+    )

app/agent/memory.py

@@ -0,0 +1,61 @@
+"""Ephemeral in-process conversation memory for interactive CLI sessions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class MemoryTurn:
+    """Single turn retained in in-memory conversation context."""
+
+    user_prompt: str
+    sql: str | None
+    validation_error: str | None
+    execution_summary: str | None
+    operation: str | None
+    table: str | None
+    changed_columns: list[str]
+    where_clause: str | None
+
+
+class ConversationMemory:
+    """Simple rolling memory for one running conversation."""
+
+    def __init__(self) -> None:
+        self._turns: list[MemoryTurn] = []
+
+    def append(self, turn: MemoryTurn) -> None:
+        self._turns.append(turn)
+
+    def clear(self) -> None:
+        self._turns.clear()
+
+    def recent(self, limit: int) -> list[MemoryTurn]:
+        if limit <= 0:
+            return []
+        return self._turns[-limit:]
+
+    def format_for_prompt(self, limit: int) -> str:
+        """
+        Format recent turns as compact plain text context for the SQL prompt.
+        """
+        rows: list[str] = []
+        for idx, turn in enumerate(self.recent(limit), start=1):
+            rows.append(f"Turn {idx}:")
+            rows.append(f"- User: {turn.user_prompt}")
+            if turn.sql:
+                rows.append(f"- SQL: {turn.sql}")
+            if turn.operation:
+                rows.append(f"- Operation: {turn.operation}")
+            if turn.table:
+                rows.append(f"- Table: {turn.table}")
+            if turn.changed_columns:
+                rows.append(f"- Changed columns: {', '.join(turn.changed_columns)}")
+            if turn.where_clause:
+                rows.append(f"- Where: {turn.where_clause}")
+            if turn.validation_error:
+                rows.append(f"- Validation: {turn.validation_error}")
+            if turn.execution_summary:
+                rows.append(f"- Result: {turn.execution_summary}")
+        return "\n".join(rows)

app/agent/sql_metadata.py

@@ -0,0 +1,101 @@
+"""Lightweight SQL metadata extraction for conversational follow-ups."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+_TABLE_PATTERN = r'(?:"([^"]+)"|([a-zA-Z_][\w]*))'
+
+
+@dataclass
+class SqlMetadata:
+    """Best-effort structured details from generated SQL."""
+
+    operation: str | None
+    table: str | None
+    changed_columns: list[str]
+    where_clause: str | None
+
+
+def first_keyword(sql: str) -> str:
+    """Return the first SQL keyword (uppercased)."""
+    cleaned = sql.strip()
+    if not cleaned:
+        return ""
+    return cleaned.split(maxsplit=1)[0].upper()
+
+
+def is_write_operation(sql: str) -> bool:
+    """True when SQL starts with INSERT or UPDATE."""
+    return first_keyword(sql) in {"INSERT", "UPDATE"}
+
+
+def extract_sql_metadata(sql: str) -> SqlMetadata:
+    """
+    Best-effort extraction for SELECT/INSERT/UPDATE.
+
+    This is intentionally simple regex parsing for MVP behavior.
+    """
+    operation = first_keyword(sql) or None
+    table: str | None = None
+    changed_columns: list[str] = []
+    where_clause: str | None = None
+    cleaned = sql.strip().rstrip(";")
+
+    if operation == "UPDATE":
+        match = re.search(
+            rf"^\s*UPDATE\s+{_TABLE_PATTERN}\s+SET\s+(.*)$",
+            cleaned,
+            re.IGNORECASE,
+        )
+        if match:
+            table = match.group(1) or match.group(2)
+            remainder = match.group(3)
+            where_split = re.split(r"\bWHERE\b", remainder, maxsplit=1, flags=re.IGNORECASE)
+            set_expr = where_split[0]
+            if len(where_split) > 1:
+                where_clause = where_split[1].strip()
+            changed_columns = _extract_set_columns(set_expr)
+    elif operation == "INSERT":
+        match = re.search(
+            rf"^\s*INSERT\s+INTO\s+{_TABLE_PATTERN}",
+            cleaned,
+            re.IGNORECASE,
+        )
+        if match:
+            table = match.group(1) or match.group(2)
+    elif operation == "SELECT":
+        match = re.search(
+            rf"\bFROM\s+{_TABLE_PATTERN}",
+            cleaned,
+            re.IGNORECASE,
+        )
+        if match:
+            table = match.group(1) or match.group(2)
+        where_match = re.search(r"\bWHERE\b\s+(.*)$", cleaned, re.IGNORECASE)
+        if where_match:
+            where_clause = where_match.group(1).strip()
+
+    return SqlMetadata(
+        operation=operation,
+        table=table,
+        changed_columns=changed_columns,
+        where_clause=where_clause,
+    )
+
+
+def _extract_set_columns(set_expr: str) -> list[str]:
+    columns: list[str] = []
+    for part in set_expr.split(","):
+        left = part.split("=", maxsplit=1)[0].strip()
+        if not left:
+            continue
+        quoted = re.match(r'"([^"]+)"', left)
+        if quoted:
+            columns.append(quoted.group(1))
+            continue
+        bare = re.match(r"([a-zA-Z_][\w]*)", left)
+        if bare:
+            columns.append(bare.group(1))
+    return columns

app/agent/sql_quality.py

@@ -0,0 +1,117 @@
+"""Heuristic SQL quality checks before execution."""
+
+from __future__ import annotations
+
+import re
+
+_BROAD_VERBS = re.compile(
+    r"\b(show|list|display|get|fetch|give|view|browse|see)\b",
+    re.IGNORECASE,
+)
+_SPECIFIC_FIELD_HINTS = re.compile(
+    r"\b(only|just|columns?|fields?|select)\s+[\w,\s]+|"
+    r"\b(name|price|email|count|sum|avg|min|max|total)\b",
+    re.IGNORECASE,
+)
+_EXPLICIT_COLUMN_THRESHOLD = 8
+_EXACT_LABEL_MATCH = re.compile(
+    r"""=\s*'[^']+'""",
+    re.IGNORECASE,
+)
+_JOIN_PATTERN = re.compile(r"\bJOIN\b", re.IGNORECASE)
+_LABEL_COLUMN_NAMES = frozenset({"name", "title", "label", "slug", "code"})
+
+
+def detect_quality_issue(user_prompt: str, sql: str) -> str | None:
+    """
+    Return a human-readable quality issue, or None if SQL looks acceptable.
+    """
+    cleaned = sql.strip()
+    if not cleaned.upper().startswith("SELECT"):
+        return None
+
+    if is_broad_list_request(user_prompt) and _over_selects_for_broad_request(cleaned):
+        return (
+            "Query lists too many explicit columns for a broad show/list request. "
+            "Use SELECT * FROM the target table with a reasonable LIMIT instead."
+        )
+
+    if _uses_exact_label_match_on_join(cleaned):
+        return (
+            "Query uses an exact string match on a related entity label. "
+            "Do not assume user text matches stored casing/spacing. "
+            "Resolve the related row with case-insensitive LOWER(...) LIKE LOWER(...) "
+            "on likely label columns, then filter the target table by foreign key "
+            "using a subquery or IN (...)."
+        )
+
+    return None
+
+
+def is_broad_list_request(user_prompt: str) -> bool:
+    """True when the user asks to show/list rows without naming specific output fields."""
+    text = user_prompt.strip()
+    if not text:
+        return False
+    if not _BROAD_VERBS.search(text):
+        return False
+    return not _SPECIFIC_FIELD_HINTS.search(text)
+
+
+def count_explicit_select_columns(sql: str) -> int:
+    """Count columns in the SELECT list; returns 0 for SELECT *."""
+    match = re.search(
+        r"^\s*SELECT\s+(DISTINCT\s+)?(.*?)\s+FROM\b",
+        sql,
+        re.IGNORECASE | re.DOTALL,
+    )
+    if not match:
+        return 0
+
+    select_list = match.group(2).strip()
+    if not select_list or select_list == "*":
+        return 0
+
+    # Rough split on commas not inside parentheses
+    depth = 0
+    parts: list[str] = []
+    current: list[str] = []
+    for char in select_list:
+        if char == "(":
+            depth += 1
+        elif char == ")":
+            depth = max(0, depth - 1)
+        elif char == "," and depth == 0:
+            parts.append("".join(current).strip())
+            current = []
+            continue
+        current.append(char)
+    if current:
+        parts.append("".join(current).strip())
+    return len([part for part in parts if part])
+
+
+def _over_selects_for_broad_request(sql: str) -> bool:
+    if re.search(r"^\s*SELECT\s+\*\s+FROM\b", sql, re.IGNORECASE):
+        return False
+    return count_explicit_select_columns(sql) >= _EXPLICIT_COLUMN_THRESHOLD
+
+
+def _uses_exact_label_match_on_join(sql: str) -> bool:
+    if not _JOIN_PATTERN.search(sql):
+        return False
+    if not _EXACT_LABEL_MATCH.search(sql):
+        return False
+
+    where_match = re.search(
+        r"\bWHERE\b(.+?)(?:\bORDER\b|\bGROUP\b|\bLIMIT\b|$)", sql, re.IGNORECASE | re.DOTALL
+    )
+    if not where_match:
+        return False
+
+    where_clause = where_match.group(1)
+    for column in _LABEL_COLUMN_NAMES:
+        pattern = rf'["\w.]*{column}["\w.]*\s*=\s*\'[^\']+\''
+        if re.search(pattern, where_clause, re.IGNORECASE):
+            return True
+    return False