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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.11.1
+Version: 0.11.2
 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.2}/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.2"
 __author__ = "Santhosh Gupta"

dbagent_cli-0.11.2/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.2}/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.2}/dbagent/agent/pipeline.py

@@ -160,7 +160,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 +445,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 +482,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:
@@ -547,8 +568,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 +661,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
 

{dbagent_cli-0.11.1 → dbagent_cli-0.11.2}/dbagent_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.11.1
+Version: 0.11.2
 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.2}/dbagent_cli.egg-info/SOURCES.txt

@@ -82,6 +82,7 @@ tests/test_formats_and_exit_codes.py
 tests/test_generator.py
 tests/test_graph.py
 tests/test_knowledge.py
+tests/test_multi_turn_conversation.py
 tests/test_plugins.py
 tests/test_prompt_queue.py
 tests/test_prompt_queue_and_sql.py

{dbagent_cli-0.11.1 → dbagent_cli-0.11.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dbagent-cli"
-version = "0.11.1"
+version = "0.11.2"
 description = "The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)"
 readme = "README.md"
 requires-python = ">=3.9"

dbagent_cli-0.11.2/tests/test_multi_turn_conversation.py

@@ -0,0 +1,102 @@
+"""
+Tests for Multi-Turn Conversational Memory, Query Refinements, and Rollback Generation.
+"""
+
+from typing import List, Optional, Callable
+from dbagent.agent.conversation import ConversationContext
+from dbagent.agent.pipeline import QueryPipeline
+from dbagent.connectors.relational import RelationalConnector
+from dbagent.llm.base import BaseLLMProvider
+from dbagent.sample_db import init_sample_database
+from dbagent.security.modes import OperationMode
+
+
+class MockWriteProvider(BaseLLMProvider):
+    """Custom provider returning specific SQL for multi-turn test."""
+    def __init__(self):
+        self.call_count = 0
+
+    @property
+    def name(self) -> str:
+        return "mock-write"
+
+    def is_available(self) -> bool:
+        return True
+
+    def list_models(self) -> List[str]:
+        return ["mock-v1"]
+
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        model: Optional[str] = None,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> str:
+        self.call_count += 1
+        if self.call_count == 1:
+            return "UPDATE customers SET first_name = 'Alice' WHERE customer_id = 1;"
+        else:
+            return "UPDATE customers SET first_name = 'Original' WHERE customer_id = 1;"
+
+
+def test_conversation_context_tracking_and_followup():
+    """Verify that ConversationContext detects follow-ups and records unexecuted proposals."""
+    ctx = ConversationContext()
+    assert ctx.is_followup_prompt("give the rollback query for the above one") is False  # Empty history
+
+    # Add a proposed write query that was not executed
+    ctx.add_turn(
+        user_prompt="Update customer 1 to inactive",
+        sql="UPDATE customers SET status = 'INACTIVE' WHERE customer_id = 1;",
+        tables=["customers"],
+        was_executed=False,
+        query_type="write",
+    )
+
+    assert len(ctx.turns) == 1
+    assert "customers" in ctx.get_active_tables()
+    assert ctx.get_last_sql() == "UPDATE customers SET status = 'INACTIVE' WHERE customer_id = 1;"
+    assert ctx.is_followup_prompt("okay also give the rollback query for the above one") is True
+    assert ctx.is_followup_prompt("modify that to customer_id = 2") is True
+
+    prompt_context = ctx.format_context_for_prompt()
+    assert "UPDATE customers SET status = 'INACTIVE'" in prompt_context
+    assert "Generated / Proposed" in prompt_context
+
+
+def test_pipeline_multi_turn_followup_table_inheritance():
+    """Verify that QueryPipeline automatically inherits active tables and targets for follow-up prompts."""
+    _, db_url = init_sample_database(force_recreate=False)
+    connector = RelationalConnector(db_url)
+    llm = MockWriteProvider()
+    ctx = ConversationContext()
+
+    # Reject confirmation to simulate user answering [n]
+    def _reject_confirmation(sql, violations):
+        return False
+
+    pipeline = QueryPipeline(
+        connector=connector,
+        llm=llm,
+        mode=OperationMode.CONFIRM,
+        confirm_callback=_reject_confirmation,
+        conversation_context=ctx,
+    )
+
+    # Turn 1: Propose update on customers (user says [n] so it is not executed)
+    res1 = pipeline.run("change customer 1 name to Alice")
+    assert res1.needs_confirmation is True
+    assert res1.was_executed is False
+    assert len(ctx.turns) == 1
+    assert "customers" in ctx.get_active_tables()
+    assert ctx.get_last_sql() == "UPDATE customers SET first_name = 'Alice' WHERE customer_id = 1;"
+
+    # Turn 2: Follow-up prompt with no table names in the text
+    res2 = pipeline.run("give the rollback query for the above one")
+    # Exact tables and schema context should be inherited from previous turn
+    assert "customers" in res2.exact_tables
+    assert res2.schema is not None
+    assert "customers" in [t.name.lower() for t in res2.schema.tables]
+
+    connector.close()

dbagent_cli-0.11.1/dbagent/agent/conversation.py

@@ -1,93 +0,0 @@
-"""
-Multi-Turn Conversational Context Tracker.
-Maintains active tables, filter predicates, entities, and sort orders across chat turns
-so users can refine queries naturally (e.g. 'Which sites?', 'Only Hyderabad', 'Top 10?').
-"""
-
-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
-    is_chat: bool = False
-
-
-class ConversationContext:
-    """
-    Tracks state and multi-turn conversational history across interactive chat turns.
-    """
-
-    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,
-        is_chat: bool = False,
-    ) -> None:
-        """Records a completed turn and updates active context."""
-        turn = ConversationTurn(
-            user_prompt=user_prompt,
-            sql=sql,
-            assistant_response=assistant_response,
-            tables=tables or [],
-            filters=filters or [],
-            is_chat=is_chat,
-        )
-        self.turns.append(turn)
-        if len(self.turns) > self.max_history:
-            self.turns.pop(0)
-
-        if tables:
-            self.active_tables.update(t.lower() for t in tables)
-
-    def get_active_tables(self) -> List[str]:
-        """Returns currently active database table names discussed in this session."""
-        return list(self.active_tables)
-
-    def get_last_turn(self) -> Optional[ConversationTurn]:
-        if self.turns:
-            return self.turns[-1]
-        return None
-
-    def format_context_for_prompt(self, max_turns: int = 5) -> str:
-        """Formats multi-turn conversation memory for inclusion in LLM prompt."""
-        if not self.turns:
-            return ""
-
-        recent = self.turns[-max_turns:]
-        lines = ["=== Previous Chat Session Memory (Use for follow-ups & context) ==="]
-        for i, t in enumerate(recent, 1):
-            lines.append(f"Turn {i}:")
-            lines.append(f"  User: \"{t.user_prompt}\"")
-            if t.sql:
-                lines.append(f"  Assistant (SQL): {t.sql.strip()}")
-            elif t.assistant_response:
-                clean_resp = t.assistant_response.strip().replace("\n", " ")
-                preview = (clean_resp[:120] + "...") if len(clean_resp) > 120 else clean_resp
-                lines.append(f"  Assistant: \"{preview}\"")
-        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()