dbagent-cli 0.2.1__tar.gz → 0.2.2__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.2.1 → dbagent_cli-0.2.2}/PKG-INFO +1 -1
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/__init__.py +1 -1
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/agent/generator.py +52 -1
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/agent/pipeline.py +36 -10
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/cli.py +8 -1
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent_cli.egg-info/PKG-INFO +1 -1
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/pyproject.toml +1 -1
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/README.md +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/agent/validator.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/config.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/connectors/base.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/connectors/factory.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/connectors/mongo.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/connectors/relational.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/llm/auto_setup.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/llm/base.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/llm/factory.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/llm/gemini_provider.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/llm/groq_provider.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/llm/mock_provider.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/llm/ollama_provider.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/llm/openrouter_provider.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/schema/formatter.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/schema/models.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/schema/selector.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/ui/console.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent/ui/viewer.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent_cli.egg-info/SOURCES.txt +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent_cli.egg-info/dependency_links.txt +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent_cli.egg-info/entry_points.txt +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent_cli.egg-info/requires.txt +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/dbagent_cli.egg-info/top_level.txt +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/setup.cfg +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/tests/test_cli.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/tests/test_connectors.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/tests/test_generator.py +0 -0
- {dbagent_cli-0.2.1 → dbagent_cli-0.2.2}/tests/test_schema.py +0 -0
|
@@ -63,6 +63,22 @@ Generate ONLY the raw executable SQL query for the given database. No markdown f
|
|
|
63
63
|
"""
|
|
64
64
|
|
|
65
65
|
|
|
66
|
+
# --- Conversational chat prompt ---
|
|
67
|
+
|
|
68
|
+
CHAT_SYSTEM_PROMPT = """You are DB-Agent, a helpful, intelligent, and friendly AI database assistant.
|
|
69
|
+
Target Database Context:
|
|
70
|
+
- **Dialect**: {dialect_name}
|
|
71
|
+
- **Database Name**: {database_name}
|
|
72
|
+
- **Available Tables**: {table_summary}
|
|
73
|
+
|
|
74
|
+
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.
|
|
76
|
+
- 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
|
+
- If the user asks general questions about the schema, tables, or database concepts, provide clear, concise, and helpful answers in formatted markdown.
|
|
78
|
+
- Do NOT generate raw SQL queries or try to execute anything in this conversational mode unless the user explicitly asks for SQL examples.
|
|
79
|
+
"""
|
|
80
|
+
|
|
81
|
+
|
|
66
82
|
# --- Retry prompt when SQL fails ---
|
|
67
83
|
|
|
68
84
|
RETRY_PROMPT_TEMPLATE = """The previous SQL query failed with this error:
|
|
@@ -90,11 +106,22 @@ class ScriptGenerator:
|
|
|
90
106
|
@staticmethod
|
|
91
107
|
def classify_intent(user_prompt: str) -> str:
|
|
92
108
|
"""
|
|
93
|
-
Classify user's natural language intent as 'read', 'write', or 'ddl'.
|
|
109
|
+
Classify user's natural language intent as 'chat', 'read', 'write', or 'ddl'.
|
|
94
110
|
Used before SQL generation to predict the query type.
|
|
95
111
|
"""
|
|
96
112
|
prompt_lower = user_prompt.lower().strip()
|
|
97
113
|
|
|
114
|
+
# Chat / Conversational indicators
|
|
115
|
+
chat_triggers = [
|
|
116
|
+
"hello", "hi", "hey", "howdy", "greetings", "good morning", "good evening",
|
|
117
|
+
"who are you", "what can you do", "what are you", "help me", "how to use",
|
|
118
|
+
"what is this", "tell me about this database", "how does this work",
|
|
119
|
+
"can you help", "thanks", "thank you", "bye", "goodbye",
|
|
120
|
+
]
|
|
121
|
+
for trig in chat_triggers:
|
|
122
|
+
if prompt_lower == trig or prompt_lower.startswith(f"{trig} ") or prompt_lower.startswith(f"{trig},") or prompt_lower.startswith(f"{trig}!"):
|
|
123
|
+
return "chat"
|
|
124
|
+
|
|
98
125
|
# DDL indicators
|
|
99
126
|
ddl_keywords = [
|
|
100
127
|
"create table", "alter table", "drop table", "add column",
|
|
@@ -119,6 +146,30 @@ class ScriptGenerator:
|
|
|
119
146
|
# Default: read
|
|
120
147
|
return "read"
|
|
121
148
|
|
|
149
|
+
def generate_chat_response(
|
|
150
|
+
self,
|
|
151
|
+
schema: DatabaseSchema,
|
|
152
|
+
user_prompt: str,
|
|
153
|
+
model: Optional[str] = None,
|
|
154
|
+
) -> str:
|
|
155
|
+
"""Generate a natural conversational response to the user's message."""
|
|
156
|
+
table_names = [t.name for t in schema.tables]
|
|
157
|
+
table_summary = ", ".join(table_names[:30])
|
|
158
|
+
if len(table_names) > 30:
|
|
159
|
+
table_summary += f" ... (and {len(table_names) - 30} more)"
|
|
160
|
+
|
|
161
|
+
system_prompt = CHAT_SYSTEM_PROMPT.format(
|
|
162
|
+
dialect_name=schema.dialect_name,
|
|
163
|
+
database_name=schema.database_name,
|
|
164
|
+
table_summary=table_summary if table_summary else "No tables found",
|
|
165
|
+
)
|
|
166
|
+
|
|
167
|
+
return self.llm.generate(
|
|
168
|
+
prompt=user_prompt,
|
|
169
|
+
system_prompt=system_prompt,
|
|
170
|
+
model=model,
|
|
171
|
+
)
|
|
172
|
+
|
|
122
173
|
# --- Full generation (for `generate` command) ---
|
|
123
174
|
|
|
124
175
|
def generate_script(
|
|
@@ -28,8 +28,11 @@ class PipelineResult:
|
|
|
28
28
|
was_executed: bool = False
|
|
29
29
|
was_auto_executed: bool = False
|
|
30
30
|
needs_confirmation: bool = False
|
|
31
|
+
# Conversational chat response
|
|
32
|
+
is_chat: bool = False
|
|
33
|
+
message: str = ""
|
|
31
34
|
# Query classification
|
|
32
|
-
query_type: str = "unknown" # read, write, ddl, unknown
|
|
35
|
+
query_type: str = "unknown" # read, write, ddl, chat, unknown
|
|
33
36
|
intent: str = "read" # from natural language
|
|
34
37
|
# Safety
|
|
35
38
|
safety_warnings: List[str] = field(default_factory=list)
|
|
@@ -91,25 +94,48 @@ class QueryPipeline:
|
|
|
91
94
|
Full pipeline: prompt -> SQL -> execute -> results.
|
|
92
95
|
|
|
93
96
|
Steps:
|
|
94
|
-
1. Classify user intent (read/write/ddl)
|
|
95
|
-
2.
|
|
96
|
-
3.
|
|
97
|
-
4.
|
|
98
|
-
5.
|
|
99
|
-
6.
|
|
100
|
-
7.
|
|
101
|
-
8.
|
|
97
|
+
1. Classify user intent (chat/read/write/ddl)
|
|
98
|
+
2. If chat: generate natural language response and return
|
|
99
|
+
3. Resolve tables (exact vs fuzzy match, detect ambiguity)
|
|
100
|
+
4. Fetch targeted schema on-demand
|
|
101
|
+
5. Generate SQL via LLM
|
|
102
|
+
6. Classify generated SQL
|
|
103
|
+
7. Safety check
|
|
104
|
+
8. Auto-execute or confirm
|
|
105
|
+
9. On error: retry with error context (up to MAX_RETRIES)
|
|
102
106
|
"""
|
|
103
107
|
result = PipelineResult()
|
|
104
108
|
result.intent = ScriptGenerator.classify_intent(user_prompt)
|
|
105
109
|
|
|
110
|
+
# --- Fast Path: Conversational Chat ---
|
|
111
|
+
if result.intent == "chat":
|
|
112
|
+
try:
|
|
113
|
+
schema = self.connector.inspect_schema(
|
|
114
|
+
include_samples=False,
|
|
115
|
+
max_samples=0,
|
|
116
|
+
include_row_counts=False,
|
|
117
|
+
)
|
|
118
|
+
result.schema = schema
|
|
119
|
+
result.message = self.generator.generate_chat_response(
|
|
120
|
+
schema=schema,
|
|
121
|
+
user_prompt=user_prompt,
|
|
122
|
+
model=self.model,
|
|
123
|
+
)
|
|
124
|
+
result.is_chat = True
|
|
125
|
+
result.was_executed = True
|
|
126
|
+
result.query_type = "chat"
|
|
127
|
+
return result
|
|
128
|
+
except Exception as e:
|
|
129
|
+
# If conversational response fails, record error or fall through
|
|
130
|
+
result.error = f"Chat generation failed: {str(e)}"
|
|
131
|
+
return result
|
|
132
|
+
|
|
106
133
|
# --- Step 1: Resolve tables ---
|
|
107
134
|
try:
|
|
108
135
|
exact, fuzzy = self.connector.resolve_tables(user_prompt)
|
|
109
136
|
result.exact_tables = exact
|
|
110
137
|
result.fuzzy_tables = fuzzy
|
|
111
138
|
except AttributeError:
|
|
112
|
-
# Connector doesn't support resolve_tables (e.g. old version)
|
|
113
139
|
exact, fuzzy = [], []
|
|
114
140
|
|
|
115
141
|
# Detect ambiguity: user said "users" but we have multiple fuzzy matches and no exact
|
|
@@ -11,6 +11,7 @@ from typing import Optional, List
|
|
|
11
11
|
import typer
|
|
12
12
|
from rich.prompt import Prompt, Confirm
|
|
13
13
|
from rich.table import Table
|
|
14
|
+
from rich.markdown import Markdown
|
|
14
15
|
from prompt_toolkit import PromptSession
|
|
15
16
|
from prompt_toolkit.history import FileHistory
|
|
16
17
|
|
|
@@ -113,7 +114,13 @@ def _choice_callback(prompt_text: str, choices: List[str]) -> Optional[str]:
|
|
|
113
114
|
|
|
114
115
|
|
|
115
116
|
def _display_pipeline_result(result: PipelineResult, show_sql: bool = True) -> None:
|
|
116
|
-
"""Display a pipeline result with SQL
|
|
117
|
+
"""Display a pipeline result with SQL, results table, or conversational chat message."""
|
|
118
|
+
if result.is_chat and result.message:
|
|
119
|
+
console.print("")
|
|
120
|
+
console.print(Markdown(result.message))
|
|
121
|
+
console.print("")
|
|
122
|
+
return
|
|
123
|
+
|
|
117
124
|
if result.error:
|
|
118
125
|
print_error(f"Error: {result.error}")
|
|
119
126
|
if result.sql:
|
|
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
|