dbagent-cli 0.11.1__tar.gz → 0.11.3__tar.gz

{dbagent_cli-0.11.1 → dbagent_cli-0.11.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.11.1
+Version: 0.11.3
 Summary: The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)
 Author: Santhosh Gupta
 License: MIT

{dbagent_cli-0.11.1 → dbagent_cli-0.11.3}/dbagent/__init__.py

@@ -2,5 +2,5 @@
 DB-Agent: Universal Database Introspection and Script Generation AI Agent (CLI).
 """
 
-__version__ = "0.11.1"
+__version__ = "0.11.3"
 __author__ = "Santhosh Gupta"

dbagent_cli-0.11.3/dbagent/agent/conversation.py

@@ -0,0 +1,140 @@
+"""
+Multi-Turn Conversational Context Tracker.
+Maintains active tables, filter predicates, entities, and generated SQL across chat turns
+so users can refine, modify, rollback, or question queries naturally (e.g. 'modify the above', 'give rollback query', 'only active users').
+"""
+
+import re
+from typing import Dict, Any, List, Optional, Set
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ConversationTurn:
+    """A single turn in the conversation."""
+    user_prompt: str
+    sql: str = ""
+    assistant_response: str = ""
+    tables: List[str] = field(default_factory=list)
+    filters: List[str] = field(default_factory=list)
+    entities: Dict[str, str] = field(default_factory=dict)
+    limit: Optional[int] = None
+    was_executed: bool = False
+    is_chat: bool = False
+    query_type: str = "read"
+
+
+class ConversationContext:
+    """
+    Tracks state and multi-turn conversational history across interactive chat turns.
+    """
+
+    FOLLOWUP_INDICATORS = [
+        "above", "previous", "that", "this", "it", "rollback", "undo", "reverse",
+        "modify", "change it", "change that", "update that", "fix that", "correct this",
+        "instead", "also", "same query", "same table", "same as above", "again",
+        "for the above", "for that", "for this", "what about", "how about",
+        "now filter", "now show", "now sort", "add where", "add order", "add limit",
+        "remove limit", "remove where", "only active", "only inactive",
+    ]
+
+    def __init__(self, max_history: int = 15):
+        self.max_history = max_history
+        self.turns: List[ConversationTurn] = []
+        self.active_tables: Set[str] = set()
+        self.active_filters: List[str] = []
+
+    def add_turn(
+        self,
+        user_prompt: str,
+        sql: str = "",
+        assistant_response: str = "",
+        tables: Optional[List[str]] = None,
+        filters: Optional[List[str]] = None,
+        was_executed: bool = False,
+        is_chat: bool = False,
+        query_type: str = "read",
+    ) -> None:
+        """Records a completed or proposed turn and updates active context."""
+        cleaned_tables = [t.lower() for t in (tables or []) if t]
+
+        # If SQL is present, also extract table names from FROM/JOIN clauses
+        if sql:
+            for m in re.findall(r'\b(?:FROM|JOIN|UPDATE|INTO|TABLE)\s+([a-zA-Z0-9_]+)', sql, re.IGNORECASE):
+                t_lower = m.lower()
+                if t_lower not in ("select", "where", "set", "values", "as") and t_lower not in cleaned_tables:
+                    cleaned_tables.append(t_lower)
+
+        turn = ConversationTurn(
+            user_prompt=user_prompt,
+            sql=sql,
+            assistant_response=assistant_response,
+            tables=cleaned_tables,
+            filters=filters or [],
+            was_executed=was_executed,
+            is_chat=is_chat,
+            query_type=query_type,
+        )
+        self.turns.append(turn)
+        if len(self.turns) > self.max_history:
+            self.turns.pop(0)
+
+        if cleaned_tables:
+            self.active_tables.update(cleaned_tables)
+
+    def get_active_tables(self) -> List[str]:
+        """Returns currently active database table names discussed in this session."""
+        if not self.active_tables and self.turns:
+            for t in reversed(self.turns):
+                if t.tables:
+                    return list(t.tables)
+        return list(self.active_tables)
+
+    def get_last_turn(self) -> Optional[ConversationTurn]:
+        """Returns the most recent conversation turn."""
+        if self.turns:
+            return self.turns[-1]
+        return None
+
+    def get_last_sql(self) -> Optional[str]:
+        """Returns the most recent SQL statement generated in the session."""
+        for t in reversed(self.turns):
+            if t.sql and t.sql.strip():
+                return t.sql.strip()
+        return None
+
+    def is_followup_prompt(self, user_prompt: str) -> bool:
+        """Checks if a user prompt is modifying, rolling back, or referencing previous turns."""
+        if not self.turns:
+            return False
+        p_lower = user_prompt.lower()
+        return any(ind in p_lower for ind in self.FOLLOWUP_INDICATORS)
+
+    def format_context_for_prompt(self, max_turns: int = 6) -> str:
+        """Formats multi-turn conversation memory with full SQL for inclusion in LLM prompt."""
+        if not self.turns:
+            return ""
+
+        recent = self.turns[-max_turns:]
+        lines = ["=== Previous Chat Session Memory (Use for follow-ups, modifications, rollbacks & context) ==="]
+        for i, t in enumerate(recent, 1):
+            lines.append(f"Turn {i}:")
+            lines.append(f"  User Request: \"{t.user_prompt}\"")
+            if t.sql:
+                status = "Executed" if t.was_executed else "Generated / Proposed"
+                lines.append(f"  Previous SQL ({status}):\n{t.sql.strip()}")
+            if t.assistant_response:
+                clean_resp = t.assistant_response.strip().replace("\n", " ")
+                preview = (clean_resp[:250] + "...") if len(clean_resp) > 250 else clean_resp
+                lines.append(f"  Assistant Output: \"{preview}\"")
+            if t.tables:
+                lines.append(f"  Tables Referenced: {', '.join(t.tables)}")
+        lines.append("=== End Chat Session Memory ===")
+
+        return "\n".join(lines)
+
+    def clear(self) -> None:
+        """Resets the conversation history."""
+        self.turns.clear()
+        self.active_tables.clear()
+        self.active_filters.clear()

{dbagent_cli-0.11.1 → dbagent_cli-0.11.3}/dbagent/agent/generator.py

@@ -63,6 +63,16 @@ Generate ONLY the raw executable SQL query for the given database. No markdown f
    - For PostgreSQL: Use a single Common Table Expression (CTE) with data-modifying statements deleting from all child, mapping, activity, profile, and preference tables first, and finally deleting from the parent table (e.g. `WITH target AS (SELECT pk_user_id FROM ivis_user WHERE ...), del_act AS (DELETE FROM user_activity WHERE user_id IN (SELECT pk_user_id FROM target)), del_map AS (DELETE FROM ivis_user_access WHERE fk_user_id IN (SELECT pk_user_id FROM target)) DELETE FROM ivis_user WHERE pk_user_id IN (SELECT pk_user_id FROM target);`).
    - For MySQL / SQLite: Generate clean, ordered DELETE statements deleting referencing child records before the parent record.
    - Always delete from all referencing tables present in the schema to avoid foreign key violations.
+11. Multi-Turn Session Memory & Query Refinement:
+   - When the user asks to modify, refine, change, or fix a query (e.g. "modify the above", "change it to X", "add a filter", "order by Y", "show site names instead", "fix the error"):
+   - Locate the previous SQL statement in the Chat Session Memory, apply the requested changes, and output the complete, valid modified SQL query.
+12. Rollback & Reversal Queries:
+   - When the user asks for a "rollback query", "undo query", or "reversal script" for a previous query (e.g. "give the rollback query for the above one", "undo that update", "how to revert this"):
+   - Inspect the previous SQL statement in the Chat Session Memory.
+   - For an UPDATE: Generate the inverse UPDATE statement restoring the original/previous values or setting them back to the previous keys.
+   - For an INSERT: Generate the corresponding DELETE statement.
+   - For a DELETE: If the previous query deleted specific records or used a CTE, generate the corresponding recovery script or reverse CTE.
+   - NEVER output dummy placeholders like `SELECT 1` when a prior query exists to be modified or rolled back.
 """
 
 

{dbagent_cli-0.11.1 → dbagent_cli-0.11.3}/dbagent/agent/pipeline.py

@@ -134,8 +134,13 @@ class QueryPipeline:
         # Database identifier
         self.db_name = "database"
         if hasattr(connector, "engine") and connector.engine:
-            self.db_name = connector.engine.url.database or "database"
-            self.db_id = str(connector.engine.url)
+            try:
+                db_val = getattr(getattr(connector.engine, "url", None), "database", None)
+                if db_val and isinstance(db_val, str):
+                    self.db_name = db_val
+                self.db_id = str(getattr(connector.engine, "url", self.db_name))
+            except Exception:
+                self.db_id = self.db_name
         else:
             self.db_id = self.db_name
         self.cache_mgr = SchemaCacheManager()
@@ -160,7 +165,10 @@ class QueryPipeline:
 
         # 2. Targeted on-demand table introspection with automatic junction expansion
         target_names = list(set(exact + fuzzy))
-        is_delete_or_mapping = any(w in user_prompt.lower() for w in ["delete", "remove", "drop", "cascade", "mapping", "map", "link", "clean", "without any mappings"])
+        is_delete_or_mapping = any(w in user_prompt.lower() for w in [
+            "delete", "remove", "drop", "cascade", "mapping", "map", "link", "clean",
+            "without any mappings", "change", "update", "rollback", "undo", "revert"
+        ])
         max_tables_limit = 25 if is_delete_or_mapping else 12
 
         if target_names:
@@ -442,7 +450,17 @@ class QueryPipeline:
         except AttributeError:
             exact, fuzzy = [], []
 
-        if not exact and len(fuzzy) > 1:
+        # If user prompt is a follow-up or has no table mentions, inherit active tables from previous conversation turns
+        is_followup = self.conversation.is_followup_prompt(user_prompt)
+        if (not exact and not fuzzy) or (not exact and is_followup):
+            active_from_ctx = self.conversation.get_active_tables()
+            if active_from_ctx:
+                exact = active_from_ctx[:10]
+                result.exact_tables = exact
+                # Also prevent false ambiguity prompt if we have strong session context
+                result.ambiguous = False
+
+        if not exact and len(fuzzy) > 1 and not is_followup:
             result.ambiguous = True
             result.ambiguous_choices = fuzzy
             if self.choice_callback:
@@ -469,6 +487,14 @@ class QueryPipeline:
         conv_context = self.conversation.format_context_for_prompt()
 
         augmented_prompt = user_prompt
+        if is_followup and self.conversation.get_last_sql():
+            last_sql = self.conversation.get_last_sql()
+            followup_hint = (
+                f"[Follow-Up Note: The user's query refers directly to the previous query in the session.]\n"
+                f"[Target Previous SQL:\n{last_sql}\n]"
+            )
+            augmented_prompt = f"{followup_hint}\n\n{augmented_prompt}"
+
         if dict_context:
             augmented_prompt = f"{dict_context}\n\n{augmented_prompt}"
         if kb_context:
@@ -492,6 +518,42 @@ class QueryPipeline:
             result.error = "LLM returned empty SQL."
             return result
 
+        # --- 5.5. Schema Ground-Truth Column & Key Pre-Validation & Auto-Correction ---
+        if sub_schema and sub_schema.tables:
+            for _ in range(self.MAX_RETRIES):
+                is_schema_valid, schema_errors = SchemaGroundTruthValidator.validate_sql_against_schema(
+                    result.sql, sub_schema
+                )
+                if is_schema_valid:
+                    break
+
+                # Hallucinated columns / keys detected -> Auto-correct using physical schema definitions
+                result.retries += 1
+                error_feedback = (
+                    "Schema Ground-Truth Column & Key Validation Discrepancy:\n"
+                    + "\n".join(f"- {err}" for err in schema_errors)
+                    + "\nPlease rewrite the query using ONLY the exact physical primary keys, foreign keys, and column names defined in the schema."
+                )
+                try:
+                    corrected_sql = self.generator.generate_retry_sql(
+                        schema=sub_schema,
+                        user_prompt=augmented_prompt,
+                        failed_sql=result.sql,
+                        error_message=error_feedback,
+                        model=self.model,
+                    )
+                    if corrected_sql and corrected_sql.strip():
+                        result.sql = corrected_sql
+                    else:
+                        break
+                except Exception:
+                    break
+
+            # Attach ground truth verification metadata
+            result.ground_truth = SchemaGroundTruthValidator.validate_response(
+                f"```sql\n{result.sql}\n```", sub_schema
+            )
+
         # --- 6. AST & Security Mode Validation ---
         result.query_type = SQLValidator.classify_query(result.sql)
         is_allowed, violations = SQLValidator.validate_for_mode(result.sql, mode=self.mode)
@@ -547,8 +609,28 @@ class QueryPipeline:
                     if self.confirm_callback(result.sql, violations):
                         should_execute = True
                     else:
+                        turn_tables = list(set((exact or []) + (fuzzy or []) + [t.name.lower() for t in (sub_schema.tables if sub_schema else [])]))
+                        self.conversation.add_turn(
+                            user_prompt=user_prompt,
+                            sql=result.sql,
+                            tables=turn_tables,
+                            was_executed=False,
+                            is_chat=False,
+                            query_type=result.query_type,
+                        )
+                        result.execution_time_ms = (time.time() - t_start) * 1000
                         return result
                 else:
+                    turn_tables = list(set((exact or []) + (fuzzy or []) + [t.name.lower() for t in (sub_schema.tables if sub_schema else [])]))
+                    self.conversation.add_turn(
+                        user_prompt=user_prompt,
+                        sql=result.sql,
+                        tables=turn_tables,
+                        was_executed=False,
+                        is_chat=False,
+                        query_type=result.query_type,
+                    )
+                    result.execution_time_ms = (time.time() - t_start) * 1000
                     return result
         elif self.mode == OperationMode.ADMIN:
             should_execute = True
@@ -620,19 +702,16 @@ class QueryPipeline:
             error=result.error,
         )
 
-        if result.success:
-            sql_tables = set()
-            if result.sql:
-                import re as _re
-                for m in _re.findall(r'\b(?:FROM|JOIN)\s+([a-zA-Z0-9_]+)', result.sql, _re.IGNORECASE):
-                    sql_tables.add(m.lower())
-            turn_tables = list(set((exact or []) + (fuzzy or []) + list(sql_tables)))
-            self.conversation.add_turn(
-                user_prompt=user_prompt,
-                sql=result.sql,
-                tables=turn_tables,
-                is_chat=False,
-            )
+        turn_tables = list(set((exact or []) + (fuzzy or []) + [t.name.lower() for t in (sub_schema.tables if sub_schema else [])]))
+        self.conversation.add_turn(
+            user_prompt=user_prompt,
+            sql=result.sql,
+            assistant_response=result.error or result.friendly_error or "",
+            tables=turn_tables,
+            was_executed=result.was_executed,
+            is_chat=False,
+            query_type=result.query_type,
+        )
 
         return result