dbagent-cli 0.5.1__tar.gz → 0.6.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/PKG-INFO +1 -1
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/__init__.py +1 -1
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/agent/generator.py +15 -5
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/agent/pipeline.py +4 -2
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/cli.py +54 -22
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/connectors/relational.py +3 -3
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/schema/formatter.py +7 -1
- dbagent_cli-0.6.0/dbagent/ui/chat_screen.py +211 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent_cli.egg-info/PKG-INFO +1 -1
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent_cli.egg-info/SOURCES.txt +1 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/pyproject.toml +1 -1
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_cli.py +5 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/README.md +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/agent/confidence.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/agent/conversation.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/agent/doctor.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/agent/error_classifier.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/agent/recommender.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/agent/validator.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/config.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/connectors/base.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/connectors/factory.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/connectors/mongo.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/execution/__init__.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/execution/cache.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/execution/changeset.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/execution/cost.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/execution/recovery.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/execution/sandbox.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/execution/streaming.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/llm/auto_setup.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/llm/base.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/llm/factory.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/llm/gemini_provider.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/llm/groq_provider.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/llm/mock_provider.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/llm/ollama_provider.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/llm/openrouter_provider.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/schema/cache.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/schema/graph.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/schema/models.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/schema/retriever.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/schema/selector.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/security/__init__.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/security/audit.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/security/environment.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/security/masking.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/security/modes.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/security/rbac.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/semantic/__init__.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/semantic/dictionary.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/semantic/templates.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/sql/__init__.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/sql/optimizer.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/sql/validator.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/ui/console.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/ui/grid_window.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/ui/viewer.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent/ui/visualizer.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent_cli.egg-info/dependency_links.txt +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent_cli.egg-info/entry_points.txt +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent_cli.egg-info/requires.txt +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/dbagent_cli.egg-info/top_level.txt +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/setup.cfg +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_cache.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_connectors.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_doctor_and_viz.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_generator.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_graph.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_recovery.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_safety.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_sandbox.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_schema.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_security.py +0 -0
- {dbagent_cli-0.5.1 → dbagent_cli-0.6.0}/tests/test_semantic.py +0 -0
|
@@ -60,6 +60,10 @@ Generate ONLY the raw executable SQL query for the given database. No markdown f
|
|
|
60
60
|
WHERE tc.constraint_type = 'FOREIGN KEY';
|
|
61
61
|
7. Always include appropriate WHERE clauses, JOINs, and ORDER BY as needed.
|
|
62
62
|
8. Default to LIMIT 50 unless the user specifies a count.
|
|
63
|
+
9. For status, active, or flag columns (e.g. `active`, `is_active`, `status`, `enabled`, `deleted`):
|
|
64
|
+
- Check the column data type and sample values shown in the schema.
|
|
65
|
+
- 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'.
|
|
66
|
+
- 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.
|
|
63
67
|
"""
|
|
64
68
|
|
|
65
69
|
|
|
@@ -72,7 +76,8 @@ Target Database Context:
|
|
|
72
76
|
- **Available Tables**: {table_summary}
|
|
73
77
|
|
|
74
78
|
You are chatting with the user in the interactive database shell.
|
|
75
|
-
- If the user says hello, introduce yourself and mention the database you are connected to.
|
|
79
|
+
- If the user says hello or greets you, introduce yourself and mention the database you are connected to.
|
|
80
|
+
- If the user gives feedback or praise (e.g. "good", "great", "thanks", "nice"), acknowledge it warmly and ask what they would like to query next.
|
|
76
81
|
- If the user asks what you can do or asks for help, explain that you can query data in plain English, inspect schemas (:table, :tables), find relationships (:fk), generate migrations/ETL pipelines, or execute SQL queries.
|
|
77
82
|
- If the user asks general questions about the schema, tables, or database concepts, provide clear, concise, and helpful answers in formatted markdown.
|
|
78
83
|
- Do NOT generate raw SQL queries or try to execute anything in this conversational mode unless the user explicitly asks for SQL examples.
|
|
@@ -110,16 +115,21 @@ class ScriptGenerator:
|
|
|
110
115
|
Used before SQL generation to predict the query type.
|
|
111
116
|
"""
|
|
112
117
|
prompt_lower = user_prompt.lower().strip()
|
|
118
|
+
clean_prompt = prompt_lower.rstrip(".!?,:;~ ")
|
|
113
119
|
|
|
114
|
-
# Chat / Conversational indicators
|
|
120
|
+
# Chat / Conversational indicators (greetings, feedback, praise, acknowledgements)
|
|
115
121
|
chat_triggers = [
|
|
116
|
-
"hello", "hi", "hey", "howdy", "greetings", "good morning", "good evening",
|
|
122
|
+
"hello", "hi", "hey", "howdy", "greetings", "good morning", "good evening", "good afternoon",
|
|
117
123
|
"who are you", "what can you do", "what are you", "help me", "how to use",
|
|
118
124
|
"what is this", "tell me about this database", "how does this work",
|
|
119
|
-
"can you help", "thanks", "thank you", "bye", "goodbye",
|
|
125
|
+
"can you help", "thanks", "thank you", "thx", "bye", "goodbye", "see ya",
|
|
126
|
+
"good", "great", "nice", "cool", "awesome", "perfect", "ok", "okay",
|
|
127
|
+
"yes", "no", "sure", "got it", "understood", "nice work", "well done",
|
|
128
|
+
"sounds good", "looks good", "alright", "super", "neat", "excellent",
|
|
129
|
+
"amazing", "wonderful", "cool thanks", "great thanks", "thanks a lot",
|
|
120
130
|
]
|
|
121
131
|
for trig in chat_triggers:
|
|
122
|
-
if
|
|
132
|
+
if clean_prompt == trig or prompt_lower.startswith(f"{trig} ") or prompt_lower.startswith(f"{trig},") or prompt_lower.startswith(f"{trig}!"):
|
|
123
133
|
return "chat"
|
|
124
134
|
|
|
125
135
|
# DDL indicators
|
|
@@ -155,12 +155,12 @@ class QueryPipeline:
|
|
|
155
155
|
# 2. Targeted on-demand table introspection (~50ms instead of 5000ms for full scan)
|
|
156
156
|
target_names = list(set(exact + fuzzy))
|
|
157
157
|
if target_names:
|
|
158
|
-
targeted = self.connector.inspect_schema(table_names=target_names, include_samples=
|
|
158
|
+
targeted = self.connector.inspect_schema(table_names=target_names, include_samples=True, max_samples=2)
|
|
159
159
|
retriever = ContextRetriever(targeted)
|
|
160
160
|
return retriever.retrieve_context(user_prompt, exact_tables=exact, fuzzy_tables=fuzzy, max_tables=8)
|
|
161
161
|
|
|
162
162
|
# 3. Fallback: inspect targeted
|
|
163
|
-
targeted = self.connector.inspect_targeted(user_prompt, max_tables=8)
|
|
163
|
+
targeted = self.connector.inspect_targeted(user_prompt, max_tables=8, include_samples=True)
|
|
164
164
|
retriever = ContextRetriever(targeted)
|
|
165
165
|
return retriever.retrieve_context(user_prompt, exact_tables=exact, fuzzy_tables=fuzzy, max_tables=8)
|
|
166
166
|
|
|
@@ -191,10 +191,12 @@ class QueryPipeline:
|
|
|
191
191
|
result.is_chat = True
|
|
192
192
|
result.was_executed = True
|
|
193
193
|
result.query_type = "chat"
|
|
194
|
+
result.execution_time_ms = (time.time() - t_start) * 1000
|
|
194
195
|
self.conversation.add_turn(user_prompt=user_prompt, sql="", is_chat=True)
|
|
195
196
|
return result
|
|
196
197
|
except Exception as e:
|
|
197
198
|
result.error = f"Chat generation failed: {str(e)}"
|
|
199
|
+
result.execution_time_ms = (time.time() - t_start) * 1000
|
|
198
200
|
return result
|
|
199
201
|
|
|
200
202
|
# --- 2. Resolve Tables & Detect Ambiguity ---
|
|
@@ -14,9 +14,10 @@ import typer
|
|
|
14
14
|
from rich.prompt import Prompt, Confirm
|
|
15
15
|
from rich.table import Table
|
|
16
16
|
from rich.markdown import Markdown
|
|
17
|
-
from
|
|
18
|
-
from prompt_toolkit import PromptSession
|
|
17
|
+
from prompt_toolkit import PromptSession, HTML
|
|
19
18
|
from prompt_toolkit.history import FileHistory
|
|
19
|
+
from prompt_toolkit.completion import WordCompleter
|
|
20
|
+
from prompt_toolkit.styles import Style
|
|
20
21
|
|
|
21
22
|
from dbagent import __version__, __author__
|
|
22
23
|
from dbagent.config import ConfigManager
|
|
@@ -63,6 +64,13 @@ from dbagent.ui.console import (
|
|
|
63
64
|
)
|
|
64
65
|
from dbagent.ui.viewer import SchemaViewer
|
|
65
66
|
from dbagent.ui.grid_window import open_desktop_grid_window, open_browser_grid_window
|
|
67
|
+
from dbagent.ui.chat_screen import (
|
|
68
|
+
clear_screen,
|
|
69
|
+
render_chat_header,
|
|
70
|
+
render_user_turn,
|
|
71
|
+
render_assistant_turn,
|
|
72
|
+
render_session_history,
|
|
73
|
+
)
|
|
66
74
|
|
|
67
75
|
app = typer.Typer(
|
|
68
76
|
name="db-agent",
|
|
@@ -447,27 +455,38 @@ def chat_command(
|
|
|
447
455
|
dialect = connector.engine.dialect.name
|
|
448
456
|
current_mode = OperationMode(mode.upper()) if mode.upper() in OperationMode.__members__ else OperationMode.CONFIRM
|
|
449
457
|
|
|
450
|
-
console.print(f"[bold green]Connected:[/bold green] {db_name} ({dialect}) - {len(table_names)} tables")
|
|
451
|
-
console.print(f"[bold cyan]AI:[/bold cyan] {llm.name} | [bold yellow]Mode:[/bold yellow] {current_mode.value}")
|
|
452
|
-
console.print("[dim]Type your question in plain English or use colon commands (e.g. :help, :tables, :doctor).[/dim]")
|
|
453
|
-
|
|
454
458
|
# Display AI recommendations on connect (instant <0.1ms with zero network overhead)
|
|
459
|
+
recs = []
|
|
455
460
|
try:
|
|
456
461
|
cached_schema = cache_mgr.load_cached_schema(db_url)
|
|
457
462
|
if cached_schema:
|
|
458
463
|
recs = QueryRecommender.generate_recommendations(cached_schema, max_suggestions=3)
|
|
459
464
|
else:
|
|
460
465
|
recs = QueryRecommender.generate_from_table_names(table_names, max_suggestions=3)
|
|
461
|
-
if recs:
|
|
462
|
-
console.print("\n[bold cyan]💡 Try asking:[/bold cyan]")
|
|
463
|
-
for r in recs:
|
|
464
|
-
console.print(f" → [italic]{r}[/italic]")
|
|
465
466
|
except Exception:
|
|
466
467
|
pass
|
|
467
|
-
|
|
468
|
+
|
|
469
|
+
# Clear screen and render Claude-style clean header card
|
|
470
|
+
clear_screen()
|
|
471
|
+
render_chat_header(
|
|
472
|
+
db_name=db_name,
|
|
473
|
+
dialect=dialect,
|
|
474
|
+
table_count=len(table_names),
|
|
475
|
+
ai_provider=llm.name,
|
|
476
|
+
ai_model=model or getattr(llm, "default_model", None),
|
|
477
|
+
mode=current_mode.value,
|
|
478
|
+
recommendations=recs,
|
|
479
|
+
)
|
|
468
480
|
|
|
469
481
|
history_file = config_mgr.config_dir / "chat_history.txt"
|
|
470
|
-
|
|
482
|
+
chat_completer = WordCompleter([
|
|
483
|
+
":help", ":tables", ":schema", ":dt", ":table", ":fk", ":path", ":join",
|
|
484
|
+
":doctor", ":mode", ":changes", ":provider", ":model", ":models",
|
|
485
|
+
":page", ":card", ":wide", ":scroll", ":window", ":grid", ":web",
|
|
486
|
+
":explain", ":history", ":auto", ":noauto", ":run", ":clear", ":export",
|
|
487
|
+
":exit", ":quit", ":q",
|
|
488
|
+
], ignore_case=True)
|
|
489
|
+
session = PromptSession(history=FileHistory(str(history_file)), completer=chat_completer)
|
|
471
490
|
|
|
472
491
|
auto_execute = True
|
|
473
492
|
conversation_ctx = ConversationContext()
|
|
@@ -488,7 +507,7 @@ def chat_command(
|
|
|
488
507
|
|
|
489
508
|
while True:
|
|
490
509
|
try:
|
|
491
|
-
user_input = session.prompt("db-agent> ").strip()
|
|
510
|
+
user_input = session.prompt(HTML("<ansicyan><b>db-agent</b></ansicyan> <ansiblue><b>❯</b></ansiblue> ")).strip()
|
|
492
511
|
if not user_input:
|
|
493
512
|
continue
|
|
494
513
|
|
|
@@ -710,14 +729,20 @@ def chat_command(
|
|
|
710
729
|
continue
|
|
711
730
|
|
|
712
731
|
if cmd_name == ":history":
|
|
713
|
-
|
|
714
|
-
|
|
715
|
-
|
|
716
|
-
|
|
717
|
-
|
|
718
|
-
|
|
719
|
-
|
|
720
|
-
|
|
732
|
+
render_session_history(query_history)
|
|
733
|
+
continue
|
|
734
|
+
|
|
735
|
+
if cmd_name in [":clear", ":cls"]:
|
|
736
|
+
clear_screen()
|
|
737
|
+
render_chat_header(
|
|
738
|
+
db_name=db_name,
|
|
739
|
+
dialect=dialect,
|
|
740
|
+
table_count=len(table_names),
|
|
741
|
+
ai_provider=llm.name,
|
|
742
|
+
ai_model=model or getattr(llm, "default_model", None),
|
|
743
|
+
mode=current_mode.value,
|
|
744
|
+
recommendations=recs,
|
|
745
|
+
)
|
|
721
746
|
continue
|
|
722
747
|
|
|
723
748
|
if cmd_name in [":page", ":cols", ":p"]:
|
|
@@ -817,10 +842,17 @@ def chat_command(
|
|
|
817
842
|
pipeline.mode = current_mode
|
|
818
843
|
pipeline.auto_execute = auto_execute
|
|
819
844
|
|
|
845
|
+
render_user_turn(user_input)
|
|
846
|
+
|
|
820
847
|
with DynamicSpinner():
|
|
821
848
|
result = pipeline.run(user_prompt=user_input)
|
|
822
849
|
|
|
823
|
-
|
|
850
|
+
render_assistant_turn(
|
|
851
|
+
result,
|
|
852
|
+
show_sql=True,
|
|
853
|
+
ai_provider=llm.name,
|
|
854
|
+
ai_model=model or getattr(llm, "default_model", None),
|
|
855
|
+
)
|
|
824
856
|
|
|
825
857
|
last_result = result
|
|
826
858
|
query_history.append(result)
|
|
@@ -246,7 +246,7 @@ class RelationalConnector(BaseConnector):
|
|
|
246
246
|
self,
|
|
247
247
|
user_prompt: str,
|
|
248
248
|
max_tables: int = 4,
|
|
249
|
-
include_samples: bool =
|
|
249
|
+
include_samples: bool = True,
|
|
250
250
|
) -> DatabaseSchema:
|
|
251
251
|
"""
|
|
252
252
|
Ultra-fast on-demand targeted inspection with fuzzy matching (<30ms).
|
|
@@ -260,8 +260,8 @@ class RelationalConnector(BaseConnector):
|
|
|
260
260
|
|
|
261
261
|
return self.inspect_schema(
|
|
262
262
|
table_names=matched[:max_tables],
|
|
263
|
-
include_samples=
|
|
264
|
-
max_samples=0,
|
|
263
|
+
include_samples=include_samples,
|
|
264
|
+
max_samples=2 if include_samples else 0,
|
|
265
265
|
include_views=False,
|
|
266
266
|
include_row_counts=False,
|
|
267
267
|
)
|
|
@@ -119,5 +119,11 @@ class SchemaFormatter:
|
|
|
119
119
|
for fk in t.foreign_keys:
|
|
120
120
|
fks.append(f"{','.join(fk.constrained_columns)}->{fk.referred_table}({','.join(fk.referred_columns)})")
|
|
121
121
|
fk_str = f" | FK: {'; '.join(fks)}" if fks else ""
|
|
122
|
-
|
|
122
|
+
sample_str = ""
|
|
123
|
+
if t.sample_rows:
|
|
124
|
+
s = t.sample_rows[0]
|
|
125
|
+
sample_pairs = [f"{k}={repr(v)}" for k, v in list(s.items())[:8] if v is not None]
|
|
126
|
+
if sample_pairs:
|
|
127
|
+
sample_str = f" | Sample: {{{', '.join(sample_pairs)}}}"
|
|
128
|
+
lines.append(f"- {t.name} ({', '.join(cols)}){fk_str}{sample_str}")
|
|
123
129
|
return "\n".join(lines)
|
|
@@ -0,0 +1,211 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Claude-Style Clean Chat Screen UI Component for DB-Agent.
|
|
3
|
+
Renders clean status headers, message turn cards, visual history tables, and styled prompt layouts.
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
import os
|
|
7
|
+
from typing import List, Optional, Dict, Any
|
|
8
|
+
from rich.console import Console
|
|
9
|
+
from rich.panel import Panel
|
|
10
|
+
from rich.table import Table
|
|
11
|
+
from rich.text import Text
|
|
12
|
+
from rich.markdown import Markdown
|
|
13
|
+
from rich.syntax import Syntax
|
|
14
|
+
from rich.rule import Rule
|
|
15
|
+
|
|
16
|
+
from dbagent import __version__
|
|
17
|
+
from dbagent.agent.pipeline import PipelineResult
|
|
18
|
+
from dbagent.ui.console import print_results_table, print_error, print_warning, print_info
|
|
19
|
+
from dbagent.ui.visualizer import ResultVisualizer
|
|
20
|
+
|
|
21
|
+
console = Console()
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def clear_screen() -> None:
|
|
25
|
+
"""Clear terminal screen cleanly across Windows, macOS, and Linux."""
|
|
26
|
+
os.system("cls" if os.name == "nt" else "clear")
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def render_chat_header(
|
|
30
|
+
db_name: str,
|
|
31
|
+
dialect: str,
|
|
32
|
+
table_count: int,
|
|
33
|
+
ai_provider: str,
|
|
34
|
+
ai_model: Optional[str] = None,
|
|
35
|
+
mode: str = "CONFIRM",
|
|
36
|
+
recommendations: Optional[List[str]] = None,
|
|
37
|
+
) -> None:
|
|
38
|
+
"""Render the top Claude-style status header card."""
|
|
39
|
+
model_str = f" ({ai_model})" if ai_model else ""
|
|
40
|
+
|
|
41
|
+
header_content = Text()
|
|
42
|
+
header_content.append("⚡ DB-AGENT ", style="bold cyan")
|
|
43
|
+
header_content.append(f"v{__version__}\n", style="bold green")
|
|
44
|
+
|
|
45
|
+
header_content.append("🟢 Connected: ", style="bold green")
|
|
46
|
+
header_content.append(f"{db_name} ", style="bold white")
|
|
47
|
+
header_content.append(f"({dialect}) • {table_count} tables\n", style="dim")
|
|
48
|
+
|
|
49
|
+
header_content.append("🤖 AI: ", style="bold cyan")
|
|
50
|
+
header_content.append(f"{ai_provider}{model_str}", style="bold white")
|
|
51
|
+
header_content.append(" | ", style="dim")
|
|
52
|
+
header_content.append("🛡️ Mode: ", style="bold yellow")
|
|
53
|
+
header_content.append(f"{mode}", style="bold yellow")
|
|
54
|
+
header_content.append(" | ", style="dim")
|
|
55
|
+
header_content.append("⌨️ Type ", style="dim")
|
|
56
|
+
header_content.append(":help", style="bold cyan")
|
|
57
|
+
header_content.append(" for commands", style="dim")
|
|
58
|
+
|
|
59
|
+
panel = Panel(
|
|
60
|
+
header_content,
|
|
61
|
+
border_style="cyan",
|
|
62
|
+
padding=(0, 1),
|
|
63
|
+
)
|
|
64
|
+
console.print(panel)
|
|
65
|
+
|
|
66
|
+
if recommendations:
|
|
67
|
+
console.print("[dim italic]💡 Try asking:[/dim italic]")
|
|
68
|
+
for rec in recommendations[:3]:
|
|
69
|
+
console.print(f" [dim]→[/dim] [cyan]{rec}[/cyan]")
|
|
70
|
+
console.print("")
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def render_user_turn(user_prompt: str) -> None:
|
|
74
|
+
"""Render a user turn message bubble/card."""
|
|
75
|
+
user_text = Text()
|
|
76
|
+
user_text.append("💬 You: ", style="bold blue")
|
|
77
|
+
user_text.append(user_prompt, style="bold white")
|
|
78
|
+
console.print("")
|
|
79
|
+
console.print(Rule(style="dim blue"))
|
|
80
|
+
console.print(user_text)
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def render_assistant_turn(
|
|
84
|
+
result: PipelineResult,
|
|
85
|
+
show_sql: bool = True,
|
|
86
|
+
ai_provider: str = "gemini",
|
|
87
|
+
ai_model: Optional[str] = None,
|
|
88
|
+
) -> None:
|
|
89
|
+
"""Render the assistant response card with SQL, data grid, and execution timing."""
|
|
90
|
+
model_display = f" ({ai_model})" if ai_model else ""
|
|
91
|
+
timing_str = f"⚡ {result.execution_time_ms / 1000:.1f}s" if result.execution_time_ms > 0 else ""
|
|
92
|
+
|
|
93
|
+
# Conversational Chat Response
|
|
94
|
+
if result.is_chat and result.message:
|
|
95
|
+
console.print(f"\n[bold cyan]🤖 DB-Agent{model_display}[/bold cyan] [dim]{timing_str}[/dim]")
|
|
96
|
+
console.print(Markdown(result.message))
|
|
97
|
+
console.print("")
|
|
98
|
+
return
|
|
99
|
+
|
|
100
|
+
# Error Response
|
|
101
|
+
if result.error:
|
|
102
|
+
console.print(f"\n[bold red]🤖 DB-Agent (Error)[/bold red] [dim]{timing_str}[/dim]")
|
|
103
|
+
if result.friendly_error:
|
|
104
|
+
print_error(result.friendly_error)
|
|
105
|
+
else:
|
|
106
|
+
print_error(f"Error: {result.error}")
|
|
107
|
+
|
|
108
|
+
if result.error_suggestions:
|
|
109
|
+
console.print(f"[dim]Did you mean: {', '.join(result.error_suggestions)}[/dim]")
|
|
110
|
+
|
|
111
|
+
if result.sql:
|
|
112
|
+
console.print("[dim]Generated SQL:[/dim]")
|
|
113
|
+
syntax = Syntax(result.sql.strip(), "sql", theme="monokai", line_numbers=False)
|
|
114
|
+
console.print(Panel(syntax, border_style="red", padding=(0, 1)))
|
|
115
|
+
if result.retries > 0:
|
|
116
|
+
console.print(f"[dim](Retried {result.retries} time(s))[/dim]")
|
|
117
|
+
console.print("")
|
|
118
|
+
return
|
|
119
|
+
|
|
120
|
+
# Executed Result Response
|
|
121
|
+
if result.was_executed:
|
|
122
|
+
console.print(f"\n[bold cyan]🤖 DB-Agent{model_display}[/bold cyan] [dim]{timing_str}[/dim]")
|
|
123
|
+
|
|
124
|
+
if show_sql and result.sql:
|
|
125
|
+
sql_clean = result.sql.strip()
|
|
126
|
+
syntax = Syntax(sql_clean, "sql", theme="monokai", line_numbers=False)
|
|
127
|
+
console.print(Panel(syntax, title="[dim]SQL Statement[/dim]", border_style="dim", padding=(0, 1)))
|
|
128
|
+
if result.is_cached:
|
|
129
|
+
console.print(f"[dim italic]⚡ Cached result ({round(result.cache_age_seconds, 1)}s old)[/dim italic]")
|
|
130
|
+
|
|
131
|
+
if result.has_results:
|
|
132
|
+
print_results_table(result.columns, result.rows, title="Results")
|
|
133
|
+
|
|
134
|
+
# ASCII visualization if detected
|
|
135
|
+
if result.visualization and result.visualization.chart_type == "BAR":
|
|
136
|
+
console.print("")
|
|
137
|
+
bar_chart = ResultVisualizer.render_ascii_bar_chart(
|
|
138
|
+
result.visualization.x_column,
|
|
139
|
+
result.visualization.y_column,
|
|
140
|
+
result.rows,
|
|
141
|
+
)
|
|
142
|
+
console.print(bar_chart)
|
|
143
|
+
console.print("")
|
|
144
|
+
else:
|
|
145
|
+
print_info("Query executed successfully. No rows returned.")
|
|
146
|
+
|
|
147
|
+
if result.confidence:
|
|
148
|
+
console.print(result.confidence.format_summary())
|
|
149
|
+
|
|
150
|
+
if result.retries > 0:
|
|
151
|
+
console.print(f"[dim](Auto-corrected after {result.retries} retry(ies))[/dim]")
|
|
152
|
+
|
|
153
|
+
# Recovery details for writes
|
|
154
|
+
if result.change_id:
|
|
155
|
+
console.print(f"\n[bold green]✓ Change committed successfully[/bold green]")
|
|
156
|
+
console.print(f" * Change ID: [bold white]{result.change_id}[/bold white]")
|
|
157
|
+
if result.recovery_status == "AVAILABLE":
|
|
158
|
+
console.print(f" * Recovery: [bold cyan]{result.recovery_message}[/bold cyan]")
|
|
159
|
+
console.print(f" * Rollback command: [bold]db-agent rollback {result.change_id}[/bold]\n")
|
|
160
|
+
else:
|
|
161
|
+
console.print(f" * Recovery: [yellow]{result.recovery_message or result.recovery_status}[/yellow]\n")
|
|
162
|
+
console.print("")
|
|
163
|
+
elif result.needs_confirmation and not result.was_executed:
|
|
164
|
+
if result.sql:
|
|
165
|
+
syntax = Syntax(result.sql.strip(), "sql", theme="monokai", line_numbers=False)
|
|
166
|
+
console.print(Panel(syntax, title=f"Generated {result.query_type.upper()} Query", border_style="yellow"))
|
|
167
|
+
console.print("[dim]This query was not executed (requires confirmation).[/dim]\n")
|
|
168
|
+
elif result.sql:
|
|
169
|
+
syntax = Syntax(result.sql.strip(), "sql", theme="monokai", line_numbers=False)
|
|
170
|
+
console.print(Panel(syntax, title="Generated SQL", border_style="dim"))
|
|
171
|
+
console.print("")
|
|
172
|
+
|
|
173
|
+
|
|
174
|
+
def render_session_history(query_history: List[PipelineResult]) -> None:
|
|
175
|
+
"""Render a clean visual session history table."""
|
|
176
|
+
if not query_history:
|
|
177
|
+
console.print("[dim]No query history in this session yet.[/dim]\n")
|
|
178
|
+
return
|
|
179
|
+
|
|
180
|
+
table = Table(title="📜 Session Query History", border_style="cyan", show_lines=True)
|
|
181
|
+
table.add_column("#", style="dim", width=4, justify="right")
|
|
182
|
+
table.add_column("Status", width=8, justify="center")
|
|
183
|
+
table.add_column("Type", width=8, justify="center")
|
|
184
|
+
table.add_column("Query / Message Preview", style="white", min_width=30)
|
|
185
|
+
table.add_column("Rows", style="cyan", width=6, justify="right")
|
|
186
|
+
table.add_column("Time", style="green", width=8, justify="right")
|
|
187
|
+
|
|
188
|
+
for i, r in enumerate(query_history, 1):
|
|
189
|
+
if r.is_chat:
|
|
190
|
+
status = "[bold blue]CHAT[/bold blue]"
|
|
191
|
+
q_type = "chat"
|
|
192
|
+
preview = (r.message or "")[:60] + "..." if len(r.message or "") > 60 else (r.message or "")
|
|
193
|
+
rows = "-"
|
|
194
|
+
elif r.success:
|
|
195
|
+
status = "[bold green]OK[/bold green]"
|
|
196
|
+
q_type = r.query_type
|
|
197
|
+
sql_str = r.sql or ""
|
|
198
|
+
preview = (sql_str[:60] + "...") if len(sql_str) > 60 else sql_str
|
|
199
|
+
rows = str(len(r.rows)) if r.has_results else "0"
|
|
200
|
+
else:
|
|
201
|
+
status = "[bold red]FAIL[/bold red]"
|
|
202
|
+
q_type = r.query_type or "error"
|
|
203
|
+
preview = (r.error or r.sql or "")[:60]
|
|
204
|
+
rows = "0"
|
|
205
|
+
|
|
206
|
+
time_str = f"{r.execution_time_ms / 1000:.1f}s" if r.execution_time_ms > 0 else "-"
|
|
207
|
+
table.add_row(str(i), status, q_type, preview, rows, time_str)
|
|
208
|
+
|
|
209
|
+
console.print("")
|
|
210
|
+
console.print(table)
|
|
211
|
+
console.print("")
|
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "dbagent-cli"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.6.0"
|
|
8
8
|
description = "Enterprise AI Database Copilot, Schema Graph Introspector & Autonomous Recovery Agent (CLI)"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.9"
|
|
@@ -417,6 +417,11 @@ def test_classify_intent():
|
|
|
417
417
|
assert ScriptGenerator.classify_intent("create table products") == "ddl"
|
|
418
418
|
assert ScriptGenerator.classify_intent("drop table temp_data") == "ddl"
|
|
419
419
|
assert ScriptGenerator.classify_intent("add column age to users") == "ddl"
|
|
420
|
+
assert ScriptGenerator.classify_intent("good") == "chat"
|
|
421
|
+
assert ScriptGenerator.classify_intent("nice") == "chat"
|
|
422
|
+
assert ScriptGenerator.classify_intent("great") == "chat"
|
|
423
|
+
assert ScriptGenerator.classify_intent("thanks") == "chat"
|
|
424
|
+
assert ScriptGenerator.classify_intent("hello") == "chat"
|
|
420
425
|
|
|
421
426
|
|
|
422
427
|
def test_strip_markdown_fences():
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|