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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.11.0
+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.0 → 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.0"
+__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.0 → dbagent_cli-0.11.2}/dbagent/agent/generator.py

@@ -59,6 +59,20 @@ Generate ONLY the raw executable SQL query for the given database. No markdown f
    - Check the column data type and sample values shown in the schema.
    - If the column is VARCHAR(1) or CHAR(1) or text, check if values use 'Y'/'N', 'YES'/'NO', '1'/'0', 'ACTIVE'/'INACTIVE', 'TRUE'/'FALSE'.
    - If sample values show 'Y'/'N' (or if it's a single-character VARCHAR/CHAR column), use `active = 'Y'` or `active IN ('Y', '1', 'true', 'TRUE')` for active/enabled status.
+10. For Deletions & Cascading Mapping Cleanup (when deleting a record/user without foreign key constraint violations):
+   - 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.0 → dbagent_cli-0.11.2}/dbagent/agent/pipeline.py

@@ -160,6 +160,12 @@ 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", "change", "update", "rollback", "undo", "revert"
+        ])
+        max_tables_limit = 25 if is_delete_or_mapping else 12
+
         if target_names:
             try:
                 all_table_names = self.connector.get_table_names()
@@ -174,24 +180,29 @@ class QueryPipeline:
                     t_lower = t.lower()
                     if t in target_names:
                         continue
-                    is_map = any(m in t_lower for m in ["_map", "mapping", "_rel", "_junction", "_link", "_access", "_config"])
+                    is_map = any(m in t_lower for m in ["_map", "mapping", "_rel", "_junction", "_link", "_access", "_config", "_preference", "_profile", "_activity", "_log"])
                     matches_stem = any(s in t_lower for s in stems)
-                    if (is_map and matches_stem) or sum(1 for s in stems if s in t_lower) >= 2:
-                        if t not in target_names:
-                            target_names.append(t)
-                    if len(target_names) >= 12:
+                    if is_delete_or_mapping:
+                        if matches_stem or (is_map and matches_stem) or sum(1 for s in stems if s in t_lower) >= 1:
+                            if t not in target_names:
+                                target_names.append(t)
+                    else:
+                        if (is_map and matches_stem) or sum(1 for s in stems if s in t_lower) >= 2:
+                            if t not in target_names:
+                                target_names.append(t)
+                    if len(target_names) >= max_tables_limit:
                         break
             except Exception:
                 pass
 
             targeted = self.connector.inspect_schema(table_names=target_names, include_samples=True, max_samples=2)
             retriever = ContextRetriever(targeted)
-            return retriever.retrieve_context(user_prompt, exact_tables=exact, fuzzy_tables=fuzzy, max_tables=12)
+            return retriever.retrieve_context(user_prompt, exact_tables=exact, fuzzy_tables=fuzzy, max_tables=max_tables_limit)
 
         # 3. Fallback: inspect targeted
-        targeted = self.connector.inspect_targeted(user_prompt, max_tables=12, include_samples=True)
+        targeted = self.connector.inspect_targeted(user_prompt, max_tables=max_tables_limit, include_samples=True)
         retriever = ContextRetriever(targeted)
-        return retriever.retrieve_context(user_prompt, exact_tables=exact, fuzzy_tables=fuzzy, max_tables=12)
+        return retriever.retrieve_context(user_prompt, exact_tables=exact, fuzzy_tables=fuzzy, max_tables=max_tables_limit)
 
     def run(
         self,
@@ -434,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:
@@ -461,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:
@@ -539,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
@@ -612,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.0 → dbagent_cli-0.11.2}/dbagent/execution/prompt_queue.py

@@ -83,15 +83,27 @@ class PromptQueueManager:
         if not raw_clean:
             return raw_clean
 
-        # Do not split direct SQL statements like "SELECT 1; SELECT 2;"
-        if raw_clean.lower().startswith(("select ", "insert ", "update ", "delete ", "create ", "alter ", "drop ", "with ")):
-            return raw_clean
-
         # Do not split special REPL commands
         if raw_clean.startswith(":"):
             return raw_clean
 
-        # Split by ';' or '&&'
+        # Do not split multi-line text with newlines containing SQL statements/keywords
+        if "\n" in raw_clean and any(kw in raw_clean.upper() for kw in [
+            "SELECT", "INSERT", "UPDATE", "DELETE", "FROM", "WHERE", "JOIN", "COMMIT", "BEGIN", "WITH", "CREATE", "ALTER", "DROP"
+        ]):
+            return raw_clean
+
+        # Do not split direct SQL statements or code blocks containing SQL DML/DDL
+        if re.search(r'\b(?:DELETE\s+FROM|INSERT\s+INTO|UPDATE\s+[a-zA-Z0-9_\"`.]+\s+SET|SELECT\s+.*?FROM|CREATE\s+TABLE|ALTER\s+TABLE|DROP\s+TABLE|COMMIT|ROLLBACK|BEGIN)\b', raw_clean, re.IGNORECASE | re.DOTALL):
+            return raw_clean
+
+        # Do not split prompts that ask to review, fix, or correct a code snippet (e.g. "... ; correct this")
+        if any(phrase in raw_clean.lower() for phrase in [
+            "correct this", "fix this", "check this", "review this", "optimize this", "explain this", "what does this do", "run this", "execute this"
+        ]):
+            return raw_clean
+
+        # Split by ';' or '&&' for distinct natural language prompts (e.g. "count users ; show top 5 sites")
         delimiter_pattern = r'\s*(?:;|&&)\s*'
         parts = [p.strip() for p in re.split(delimiter_pattern, raw_clean) if p.strip()]
 

{dbagent_cli-0.11.0 → dbagent_cli-0.11.2}/dbagent/sql/validator.py

@@ -150,17 +150,23 @@ class SQLValidator:
 
         stmt = statements[0]
 
-        for pat in cls.READ_PATTERNS:
-            if re.search(pat, stmt, re.IGNORECASE | re.DOTALL):
-                return "read"
+        # Check for data-modifying CTEs (WITH ... DELETE|INSERT|UPDATE ...)
+        if re.match(r"^\s*WITH\b", stmt, re.IGNORECASE):
+            if re.search(r"\b(?:DELETE\s+FROM|INSERT\s+INTO|UPDATE\s+[a-zA-Z0-9_\"`.]+\s+SET)\b", stmt, re.IGNORECASE):
+                return "write"
+            return "read"
+
+        for pat in cls.WRITE_PATTERNS:
+            if re.search(pat, stmt, re.IGNORECASE):
+                return "write"
 
         for pat in cls.DDL_PATTERNS:
             if re.search(pat, stmt, re.IGNORECASE):
                 return "ddl"
 
-        for pat in cls.WRITE_PATTERNS:
-            if re.search(pat, stmt, re.IGNORECASE):
-                return "write"
+        for pat in cls.READ_PATTERNS:
+            if re.search(pat, stmt, re.IGNORECASE | re.DOTALL):
+                return "read"
 
         return "unknown"
 

{dbagent_cli-0.11.0 → dbagent_cli-0.11.2}/dbagent/ui/console.py

@@ -53,8 +53,8 @@ def print_info(message: str) -> None:
 
 
 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)
+    """Print syntax-highlighted code block without line numbers for easy copy-pasting."""
+    syntax = Syntax(code, language, theme="monokai", line_numbers=False)
     panel = Panel(syntax, title=f"[bold cyan]{title or language.upper()}[/bold cyan]", border_style="dim")
     console.print(panel)
 

{dbagent_cli-0.11.0 → 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.0
+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.0 → dbagent_cli-0.11.2}/dbagent_cli.egg-info/SOURCES.txt

@@ -82,8 +82,10 @@ 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
 tests/test_providers.py
 tests/test_recovery.py
 tests/test_safety.py

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dbagent-cli"
-version = "0.11.0"
+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.2/tests/test_prompt_queue_and_sql.py

@@ -0,0 +1,48 @@
+"""
+Tests for PromptQueueManager SQL script handling and SQLValidator CTE classification.
+"""
+
+from dbagent.execution.prompt_queue import PromptQueueManager
+from dbagent.sql.validator import SQLValidator
+
+
+def test_queue_does_not_split_multiline_sql_with_semicolons():
+    """Verify that multi-line SQL scripts with internal semicolons are not split into queued fragments."""
+    queue_mgr = PromptQueueManager()
+
+    sql_input = """
+    DELETE FROM ivis_customer_user_map WHERE fk_user_id = 5201;
+    DELETE FROM ivis_user_access WHERE fk_user_id = 5201;
+    DELETE FROM user_activity WHERE user_id = 5201;
+    DELETE FROM ivis_user WHERE pk_user_id = 5201;
+    COMMIT; correct this
+    """
+
+    res = queue_mgr.split_and_enqueue_chained(sql_input)
+    assert res == sql_input.strip()
+    assert queue_mgr.has_next() is False
+
+
+def test_classify_query_with_cte_delete():
+    """Verify that data-modifying WITH CTE statements are classified as 'write'."""
+    cte_delete = """
+    WITH target_user AS (
+        SELECT pk_user_id FROM ivis_user WHERE username = 'target_user'
+    ),
+    del_access AS (
+        DELETE FROM ivis_user_access WHERE fk_user_id IN (SELECT pk_user_id FROM target_user)
+    )
+    DELETE FROM ivis_user WHERE pk_user_id IN (SELECT pk_user_id FROM target_user);
+    """
+    assert SQLValidator.classify_query(cte_delete) == "write"
+
+
+def test_classify_query_with_cte_read():
+    """Verify that regular SELECT WITH CTE statements remain classified as 'read'."""
+    cte_read = """
+    WITH active_users AS (
+        SELECT pk_user_id, username FROM ivis_user WHERE active = 'Y'
+    )
+    SELECT * FROM active_users LIMIT 10;
+    """
+    assert SQLValidator.classify_query(cte_read) == "read"

dbagent_cli-0.11.0/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()