dbagent-cli 0.7.0__tar.gz → 0.7.2__tar.gz

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.7.0
+Version: 0.7.2
 Summary: Enterprise AI Database Copilot, Schema Graph Introspector & Autonomous Recovery Agent (CLI)
 Author: Santhosh Gupta
 Requires-Python: >=3.9

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/dbagent/__init__.py

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

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/dbagent/agent/conversation.py

@@ -12,7 +12,8 @@ from dataclasses import dataclass, field
 class ConversationTurn:
     """A single turn in the conversation."""
     user_prompt: str
-    sql: 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)
@@ -22,10 +23,10 @@ class ConversationTurn:
 
 class ConversationContext:
     """
-    Tracks state across interactive chat turns.
+    Tracks state and multi-turn conversational history across interactive chat turns.
     """
 
-    def __init__(self, max_history: int = 10):
+    def __init__(self, max_history: int = 15):
         self.max_history = max_history
         self.turns: List[ConversationTurn] = []
         self.active_tables: Set[str] = set()
@@ -34,7 +35,8 @@ class ConversationContext:
     def add_turn(
         self,
         user_prompt: str,
-        sql: str,
+        sql: str = "",
+        assistant_response: str = "",
         tables: Optional[List[str]] = None,
         filters: Optional[List[str]] = None,
         is_chat: bool = False,
@@ -43,6 +45,7 @@ class ConversationContext:
         turn = ConversationTurn(
             user_prompt=user_prompt,
             sql=sql,
+            assistant_response=assistant_response,
             tables=tables or [],
             filters=filters or [],
             is_chat=is_chat,
@@ -59,19 +62,23 @@ class ConversationContext:
             return self.turns[-1]
         return None
 
-    def format_context_for_prompt(self) -> str:
-        """Formats recent context for inclusion in LLM prompt."""
+    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_valid = [t for t in self.turns[-3:] if t.sql and not t.is_chat]
-        if not recent_valid:
-            return ""
-
-        lines = ["Previous Conversation Context (Use for follow-up refinements):"]
-        for i, t in enumerate(recent_valid, 1):
-            lines.append(f"  {i}. User: \"{t.user_prompt}\"")
-            lines.append(f"     Last SQL: `{t.sql.strip()}`")
+        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)
 

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/dbagent/agent/generator.py

@@ -169,9 +169,10 @@ class ScriptGenerator:
         self,
         schema: DatabaseSchema,
         user_prompt: str,
+        conversation_history: Optional[str] = None,
         model: Optional[str] = None,
     ) -> str:
-        """Generate a natural conversational response to the user's message."""
+        """Generate a natural conversational response to the user's message with full session memory."""
         table_names = [t.name for t in schema.tables]
         table_summary = ", ".join(table_names[:30])
         if len(table_names) > 30:
@@ -183,8 +184,12 @@ class ScriptGenerator:
             table_summary=table_summary if table_summary else "No tables found",
         )
 
+        full_prompt = user_prompt
+        if conversation_history:
+            full_prompt = f"{conversation_history}\n\nUser: {user_prompt}"
+
         return self.llm.generate(
-            prompt=user_prompt,
+            prompt=full_prompt,
             system_prompt=system_prompt,
             model=model,
         )

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/dbagent/agent/pipeline.py

@@ -53,6 +53,7 @@ class PipelineResult:
     expires_at: Optional[str] = None
     # Conversational chat response
     is_chat: bool = False
+    is_direct_sql: bool = False
     message: str = ""
     # Query classification
     query_type: str = "unknown"  # read, write, ddl, chat, unknown
@@ -183,22 +184,145 @@ class QueryPipeline:
                     schema = self.connector.inspect_schema(table_names=quick_tables, include_samples=False, max_samples=0)
 
                 result.schema = schema
+                conv_history = self.conversation.format_context_for_prompt()
                 result.message = self.generator.generate_chat_response(
                     schema=schema,
                     user_prompt=user_prompt,
+                    conversation_history=conv_history,
                     model=self.model,
                 )
                 result.is_chat = True
                 result.was_executed = True
                 result.query_type = "chat"
                 result.execution_time_ms = (time.time() - t_start) * 1000
-                self.conversation.add_turn(user_prompt=user_prompt, sql="", is_chat=True)
+                self.conversation.add_turn(user_prompt=user_prompt, sql="", assistant_response=result.message, is_chat=True)
                 return result
             except Exception as e:
                 result.error = f"Chat generation failed: {str(e)}"
                 result.execution_time_ms = (time.time() - t_start) * 1000
                 return result
 
+        # --- 1.5. Direct Raw SQL Fast-Path (0ms AI latency, 0 tokens) ---
+        if SQLValidator.is_direct_sql(user_prompt):
+            raw_sql = user_prompt.strip().rstrip(";") + ";"
+            result.sql = raw_sql
+            result.is_direct_sql = True
+            result.query_type = SQLValidator.classify_query(raw_sql)
+            result.intent = result.query_type
+
+            # AST & Security Mode Validation
+            is_allowed, violations = SQLValidator.validate_for_mode(raw_sql, mode=self.mode)
+            result.safety_violations = violations
+            if not is_allowed:
+                result.error = f"Security Policy Violation in {self.mode.value} mode: " + "; ".join(violations)
+                result.execution_time_ms = (time.time() - t_start) * 1000
+                return result
+
+            # Check Result Cache
+            if self.use_cache and not bypass_cache and result.query_type == "read":
+                cached = self.result_cache.get(self.db_name, raw_sql)
+                if cached:
+                    cols, rows, age = cached
+                    result.columns = cols
+                    result.rows = rows
+                    result.was_executed = True
+                    result.is_cached = True
+                    result.cache_age_seconds = age
+                    result.execution_time_ms = (time.time() - t_start) * 1000
+                    result.visualization = ResultVisualizer.analyze_result_shape(cols, rows)
+                    return result
+
+            # Decide Execute or Confirm
+            should_execute = False
+            if force:
+                should_execute = True
+            elif self.mode in (OperationMode.SAFE, OperationMode.READ_ONLY):
+                should_execute = (result.query_type == "read")
+                result.was_auto_executed = should_execute
+            elif self.mode == OperationMode.CONFIRM:
+                if result.query_type == "read" and not violations:
+                    should_execute = True
+                    result.was_auto_executed = True
+                else:
+                    result.needs_confirmation = True
+                    if self.confirm_callback:
+                        if self.confirm_callback(result.sql, violations):
+                            should_execute = True
+                        else:
+                            result.execution_time_ms = (time.time() - t_start) * 1000
+                            return result
+                    else:
+                        result.execution_time_ms = (time.time() - t_start) * 1000
+                        return result
+            elif self.mode == OperationMode.ADMIN:
+                should_execute = True
+
+            if should_execute:
+                # 1. Attempt direct execution without AI (Zero AI Latency)
+                cols, rows, error = self.connector.execute_query(raw_sql, limit=100)
+                if error is None:
+                    result.columns = cols
+                    result.rows = rows
+                    result.was_executed = True
+                    result.error = None
+                    if self.mask_sensitive_data and result.rows:
+                        result.rows = SensitiveDataMasker.mask_row_set(result.columns, result.rows)
+                    if result.rows:
+                        result.visualization = ResultVisualizer.analyze_result_shape(result.columns, result.rows)
+                    if self.use_cache and result.query_type == "read":
+                        self.result_cache.set(self.db_name, raw_sql, result.columns, result.rows)
+                    exec_duration = (time.time() - t_start) * 1000
+                    result.execution_time_ms = exec_duration
+                    self.audit_logger.log_execution(
+                        database=self.db_name,
+                        user=self.user_id,
+                        question=user_prompt,
+                        sql=result.sql,
+                        mode=self.mode.value,
+                        status="SUCCESS",
+                        execution_time_ms=exec_duration,
+                        rows_returned=len(result.rows),
+                        risk_level="LOW",
+                    )
+                    self.conversation.add_turn(user_prompt=user_prompt, sql=result.sql, is_chat=False)
+                    return result
+
+                # 2. Direct query failed with DB error -> Fallback to AI Auto-Repair
+                try:
+                    exact, fuzzy = self.connector.resolve_tables(raw_sql)
+                except Exception:
+                    exact, fuzzy = [], []
+                sub_schema, join_steps = self._get_fast_schema_context(raw_sql, exact, fuzzy)
+                result.schema = sub_schema
+                result.join_steps = join_steps
+                repair_prompt = f"The following direct SQL query failed with a database error. Please correct the query using the database schema:\nFailed SQL: {raw_sql}\nError: {error}"
+                self._execute_with_retry(result, repair_prompt, sub_schema)
+
+                if result.success:
+                    if self.mask_sensitive_data and result.rows:
+                        result.rows = SensitiveDataMasker.mask_row_set(result.columns, result.rows)
+                    if result.rows:
+                        result.visualization = ResultVisualizer.analyze_result_shape(result.columns, result.rows)
+                    if self.use_cache and result.query_type == "read":
+                        self.result_cache.set(self.db_name, result.sql, result.columns, result.rows)
+                    self.conversation.add_turn(user_prompt=user_prompt, sql=result.sql, is_chat=False)
+
+                exec_duration = (time.time() - t_start) * 1000
+                result.execution_time_ms = exec_duration
+                self.audit_logger.log_execution(
+                    database=self.db_name,
+                    user=self.user_id,
+                    question=user_prompt,
+                    sql=result.sql,
+                    mode=self.mode.value,
+                    status="SUCCESS" if result.success else "FAILED",
+                    execution_time_ms=exec_duration,
+                    rows_returned=len(result.rows),
+                    risk_level=result.complexity.risk_level if result.complexity else "LOW",
+                    error=result.error,
+                )
+                return result
+
         # --- 2. Resolve Tables & Detect Ambiguity ---
         try:
             exact, fuzzy = self.connector.resolve_tables(user_prompt)

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/dbagent/sql/validator.py

@@ -52,6 +52,61 @@ class SQLValidator:
         r"^\s*REVOKE\b",
     ]
 
+    DIRECT_SQL_REGEX = re.compile(
+        r"^\s*(?:"
+        r"WITH\b\s+[a-zA-Z0-9_\"`]+\s+AS\b|"
+        r"INSERT\s+(?:INTO\s+)?[a-zA-Z0-9_\"`.]+|"
+        r"UPDATE\s+[a-zA-Z0-9_\"`.]+\s+SET\b|"
+        r"DELETE\s+FROM\s+[a-zA-Z0-9_\"`.]+|"
+        r"CREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|INDEX|VIEW|DATABASE|SCHEMA|UNIQUE\s+INDEX)\b|"
+        r"ALTER\s+(?:TABLE|VIEW)\b|"
+        r"DROP\s+(?:TABLE|VIEW|INDEX|DATABASE|SCHEMA)\b|"
+        r"TRUNCATE\s+(?:TABLE\s+)?[a-zA-Z0-9_\"`.]+|"
+        r"SHOW\s+(?:TABLES|DATABASES|COLUMNS|SCHEMAS|PROCESSLIST|VARIABLES|STATUS|CREATE\s+TABLE)\b|"
+        r"DESCRIBE\s+[a-zA-Z0-9_\"`.]+|"
+        r"DESC\s+[a-zA-Z0-9_\"`.]+|"
+        r"EXPLAIN\s+(?:ANALYZE\s+)?(?:SELECT|INSERT|UPDATE|DELETE|WITH)\b|"
+        r"PRAGMA\s+[a-zA-Z0-9_\"`]+"
+        r")",
+        re.IGNORECASE | re.DOTALL,
+    )
+
+    @classmethod
+    def is_direct_sql(cls, text: str) -> bool:
+        """
+        Determines if the input string is directly a raw SQL query (bypassing AI generation).
+        """
+        if not text or not text.strip():
+            return False
+        clean = re.sub(r"--[^\n]*\n?", "", text).strip()
+        clean = re.sub(r"/\*.*?\*/", "", clean, flags=re.DOTALL).strip()
+        clean_lower = clean.lower()
+
+        # Check explicit non-SELECT DDL/DML/Utility SQL patterns
+        if cls.DIRECT_SQL_REGEX.search(clean):
+            return True
+
+        # Check SELECT statements
+        if clean_lower.startswith("select"):
+            tokens = clean.split()
+            if len(tokens) < 2:
+                return False
+
+            # Natural language conversational words indicate prompt rather than raw SQL
+            nl_words = ["the ", " who ", " which ", " that ", " whose ", " please ", " me ", " their ", " those "]
+            if any(w in clean_lower for w in nl_words):
+                return False
+
+            # Single expression queries: SELECT 1;, SELECT NOW();, SELECT version();
+            if len(tokens) == 2 and (tokens[1].rstrip(";").isdigit() or tokens[1].endswith("()") or tokens[1].startswith("@@")):
+                return True
+
+            # Standard SQL SELECT ... FROM ...
+            if " from " in clean_lower:
+                return True
+
+        return False
+
     @classmethod
     def split_statements(cls, sql: str) -> List[str]:
         """Splits SQL script into individual statements, ignoring semicolons inside quotes/literals."""

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/dbagent/ui/chat_screen.py

@@ -87,12 +87,17 @@ def render_assistant_turn(
     ai_model: Optional[str] = None,
 ) -> None:
     """Render the assistant response card with SQL, data grid, and execution timing."""
-    model_display = f" ({ai_model})" if ai_model else ""
-    timing_str = f"⚡ {result.execution_time_ms / 1000:.1f}s" if result.execution_time_ms > 0 else ""
+    if result.is_direct_sql and result.retries == 0:
+        agent_title = "[bold cyan]🤖 DB-Agent[/bold cyan] [dim](Direct SQL • 0 tokens)[/dim]"
+    else:
+        model_display = f" ({ai_model})" if ai_model else ""
+        agent_title = f"[bold cyan]🤖 DB-Agent{model_display}[/bold cyan]"
+        
+    timing_str = f"⚡ {result.execution_time_ms / 1000:.2f}s" if result.execution_time_ms > 0 else ""
     
     # Conversational Chat Response
     if result.is_chat and result.message:
-        console.print(f"\n[bold cyan]🤖 DB-Agent{model_display}[/bold cyan] [dim]{timing_str}[/dim]")
+        console.print(f"\n{agent_title} [dim]{timing_str}[/dim]")
         console.print(Markdown(result.message))
         console.print("")
         return
@@ -119,8 +124,11 @@ def render_assistant_turn(
 
     # Executed Result Response
     if result.was_executed:
-        console.print(f"\n[bold cyan]🤖 DB-Agent{model_display}[/bold cyan] [dim]{timing_str}[/dim]")
+        console.print(f"\n{agent_title} [dim]{timing_str}[/dim]")
         
+        if result.is_direct_sql and result.retries > 0:
+            console.print(f"[bold green]✓ Direct SQL failed on database and was automatically fixed by AI ({result.retries} retry)[/bold green]")
+
         if show_sql and result.sql:
             sql_clean = result.sql.strip()
             syntax = Syntax(sql_clean, "sql", theme="monokai", line_numbers=False)

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/dbagent_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.7.0
+Version: 0.7.2
 Summary: Enterprise AI Database Copilot, Schema Graph Introspector & Autonomous Recovery Agent (CLI)
 Author: Santhosh Gupta
 Requires-Python: >=3.9

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/dbagent_cli.egg-info/SOURCES.txt

@@ -66,6 +66,7 @@ dbagent_cli.egg-info/top_level.txt
 tests/test_cache.py
 tests/test_cli.py
 tests/test_connectors.py
+tests/test_direct_sql.py
 tests/test_doctor_and_viz.py
 tests/test_generator.py
 tests/test_graph.py

{dbagent_cli-0.7.0 → dbagent_cli-0.7.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dbagent-cli"
-version = "0.7.0"
+version = "0.7.2"
 description = "Enterprise AI Database Copilot, Schema Graph Introspector & Autonomous Recovery Agent (CLI)"
 readme = "README.md"
 requires-python = ">=3.9"

dbagent_cli-0.7.2/tests/test_direct_sql.py

@@ -0,0 +1,83 @@
+"""
+Tests for Direct SQL Fast-Path and Self-Healing AI Auto-Correction.
+"""
+
+from dbagent.sql.validator import SQLValidator
+from dbagent.agent.pipeline import QueryPipeline
+from dbagent.connectors.factory import create_connector
+from dbagent.llm.mock_provider import MockProvider
+
+
+def test_is_direct_sql_detection():
+    # Direct SQL queries
+    assert SQLValidator.is_direct_sql("SELECT * FROM users;") is True
+    assert SQLValidator.is_direct_sql("select pk_user_id, username from ivis_user where active = 'Y'") is True
+    assert SQLValidator.is_direct_sql("SELECT count(*) FROM orders") is True
+    assert SQLValidator.is_direct_sql("SELECT 1;") is True
+    assert SQLValidator.is_direct_sql("WITH cte AS (SELECT 1) SELECT * FROM cte") is True
+    assert SQLValidator.is_direct_sql("INSERT INTO users (id, name) VALUES (1, 'Alice')") is True
+    assert SQLValidator.is_direct_sql("UPDATE users SET active = 'Y' WHERE id = 1") is True
+    assert SQLValidator.is_direct_sql("DELETE FROM users WHERE id = 1") is True
+    assert SQLValidator.is_direct_sql("CREATE TABLE test (id INT);") is True
+    assert SQLValidator.is_direct_sql("SHOW TABLES;") is True
+    assert SQLValidator.is_direct_sql("DESCRIBE users;") is True
+    assert SQLValidator.is_direct_sql("EXPLAIN SELECT * FROM users") is True
+    assert SQLValidator.is_direct_sql("PRAGMA table_info(users)") is True
+
+    # Natural language queries (must NOT be detected as direct SQL)
+    assert SQLValidator.is_direct_sql("list top 10 most recent active users") is False
+    assert SQLValidator.is_direct_sql("show me the active users") is False
+    assert SQLValidator.is_direct_sql("how many users registered today") is False
+    assert SQLValidator.is_direct_sql("select the active users who joined last week") is False
+    assert SQLValidator.is_direct_sql("hello") is False
+
+
+def test_pipeline_direct_sql_fast_path():
+    # Set up in-memory sqlite database
+    connector = create_connector("sqlite:///:memory:")
+    connector.execute_query("CREATE TABLE users (id INT, name VARCHAR, active VARCHAR(1));")
+    connector.execute_query("INSERT INTO users VALUES (1, 'Alice', 'Y'), (2, 'Bob', 'N');")
+
+    mock_llm = MockProvider()
+    pipeline = QueryPipeline(
+        connector=connector,
+        llm=mock_llm,
+    )
+
+    # 1. Direct valid query executes without LLM (0 retries)
+    result = pipeline.run("SELECT * FROM users WHERE active = 'Y';")
+    assert result.was_executed is True
+    assert result.is_direct_sql is True
+    assert result.retries == 0
+    assert len(result.rows) == 1
+    assert result.rows[0]["name"] == "Alice"
+    assert result.error is None
+
+    connector.close()
+
+
+def test_pipeline_direct_sql_auto_repair():
+    connector = create_connector("sqlite:///:memory:")
+    connector.execute_query("CREATE TABLE users (id INT, name VARCHAR, active VARCHAR(1));")
+    connector.execute_query("INSERT INTO users VALUES (1, 'Alice', 'Y'), (2, 'Bob', 'N');")
+
+    class RepairMockLLM(MockProvider):
+        def generate(self, prompt: str, system_prompt: str = None, model: str = None, stream_callback=None) -> str:
+            # Emulate AI fixing the typo 'usrs' -> 'users'
+            return "SELECT * FROM users WHERE active = 'Y';"
+
+    pipeline = QueryPipeline(
+        connector=connector,
+        llm=RepairMockLLM(),
+    )
+
+    # Direct query with typo in table name 'usrs' -> fails initially on DB, then AI auto-repairs to 'users'
+    result = pipeline.run("SELECT * FROM usrs WHERE active = 'Y';")
+    assert result.was_executed is True
+    assert result.is_direct_sql is True
+    assert result.retries == 1
+    assert "users" in result.sql
+    assert len(result.rows) == 1
+    assert result.rows[0]["name"] == "Alice"
+
+    connector.close()