dbagent-cli 0.6.1__tar.gz → 0.7.1__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.6.1 → dbagent_cli-0.7.1}/PKG-INFO +1 -1
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/__init__.py +1 -1
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/agent/pipeline.py +122 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/cli.py +124 -45
- dbagent_cli-0.7.1/dbagent/llm/anthropic_provider.py +126 -0
- dbagent_cli-0.7.1/dbagent/llm/deepseek_provider.py +128 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/llm/factory.py +47 -10
- dbagent_cli-0.7.1/dbagent/llm/mistral_provider.py +129 -0
- dbagent_cli-0.7.1/dbagent/llm/openai_provider.py +131 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/sql/validator.py +55 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/ui/chat_screen.py +12 -4
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent_cli.egg-info/PKG-INFO +1 -1
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent_cli.egg-info/SOURCES.txt +6 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/pyproject.toml +1 -1
- dbagent_cli-0.7.1/tests/test_direct_sql.py +83 -0
- dbagent_cli-0.7.1/tests/test_providers.py +69 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/README.md +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/agent/confidence.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/agent/conversation.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/agent/doctor.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/agent/error_classifier.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/agent/generator.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/agent/recommender.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/agent/validator.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/config.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/connectors/base.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/connectors/factory.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/connectors/mongo.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/connectors/relational.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/execution/__init__.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/execution/cache.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/execution/changeset.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/execution/cost.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/execution/recovery.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/execution/sandbox.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/execution/streaming.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/llm/auto_setup.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/llm/base.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/llm/gemini_provider.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/llm/groq_provider.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/llm/mock_provider.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/llm/ollama_provider.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/llm/openrouter_provider.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/schema/cache.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/schema/formatter.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/schema/graph.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/schema/models.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/schema/retriever.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/schema/selector.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/security/__init__.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/security/audit.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/security/environment.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/security/masking.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/security/modes.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/security/rbac.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/semantic/__init__.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/semantic/dictionary.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/semantic/templates.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/sql/__init__.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/sql/optimizer.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/ui/console.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/ui/grid_window.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/ui/viewer.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent/ui/visualizer.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent_cli.egg-info/dependency_links.txt +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent_cli.egg-info/entry_points.txt +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent_cli.egg-info/requires.txt +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/dbagent_cli.egg-info/top_level.txt +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/setup.cfg +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_cache.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_cli.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_connectors.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_doctor_and_viz.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_generator.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_graph.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_recovery.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_safety.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_sandbox.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_schema.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_security.py +0 -0
- {dbagent_cli-0.6.1 → dbagent_cli-0.7.1}/tests/test_semantic.py +0 -0
|
@@ -53,6 +53,7 @@ class PipelineResult:
|
|
|
53
53
|
expires_at: Optional[str] = None
|
|
54
54
|
# Conversational chat response
|
|
55
55
|
is_chat: bool = False
|
|
56
|
+
is_direct_sql: bool = False
|
|
56
57
|
message: str = ""
|
|
57
58
|
# Query classification
|
|
58
59
|
query_type: str = "unknown" # read, write, ddl, chat, unknown
|
|
@@ -199,6 +200,127 @@ class QueryPipeline:
|
|
|
199
200
|
result.execution_time_ms = (time.time() - t_start) * 1000
|
|
200
201
|
return result
|
|
201
202
|
|
|
203
|
+
# --- 1.5. Direct Raw SQL Fast-Path (0ms AI latency, 0 tokens) ---
|
|
204
|
+
if SQLValidator.is_direct_sql(user_prompt):
|
|
205
|
+
raw_sql = user_prompt.strip().rstrip(";") + ";"
|
|
206
|
+
result.sql = raw_sql
|
|
207
|
+
result.is_direct_sql = True
|
|
208
|
+
result.query_type = SQLValidator.classify_query(raw_sql)
|
|
209
|
+
result.intent = result.query_type
|
|
210
|
+
|
|
211
|
+
# AST & Security Mode Validation
|
|
212
|
+
is_allowed, violations = SQLValidator.validate_for_mode(raw_sql, mode=self.mode)
|
|
213
|
+
result.safety_violations = violations
|
|
214
|
+
if not is_allowed:
|
|
215
|
+
result.error = f"Security Policy Violation in {self.mode.value} mode: " + "; ".join(violations)
|
|
216
|
+
result.execution_time_ms = (time.time() - t_start) * 1000
|
|
217
|
+
return result
|
|
218
|
+
|
|
219
|
+
# Check Result Cache
|
|
220
|
+
if self.use_cache and not bypass_cache and result.query_type == "read":
|
|
221
|
+
cached = self.result_cache.get(self.db_name, raw_sql)
|
|
222
|
+
if cached:
|
|
223
|
+
cols, rows, age = cached
|
|
224
|
+
result.columns = cols
|
|
225
|
+
result.rows = rows
|
|
226
|
+
result.was_executed = True
|
|
227
|
+
result.is_cached = True
|
|
228
|
+
result.cache_age_seconds = age
|
|
229
|
+
result.execution_time_ms = (time.time() - t_start) * 1000
|
|
230
|
+
result.visualization = ResultVisualizer.analyze_result_shape(cols, rows)
|
|
231
|
+
return result
|
|
232
|
+
|
|
233
|
+
# Decide Execute or Confirm
|
|
234
|
+
should_execute = False
|
|
235
|
+
if force:
|
|
236
|
+
should_execute = True
|
|
237
|
+
elif self.mode in (OperationMode.SAFE, OperationMode.READ_ONLY):
|
|
238
|
+
should_execute = (result.query_type == "read")
|
|
239
|
+
result.was_auto_executed = should_execute
|
|
240
|
+
elif self.mode == OperationMode.CONFIRM:
|
|
241
|
+
if result.query_type == "read" and not violations:
|
|
242
|
+
should_execute = True
|
|
243
|
+
result.was_auto_executed = True
|
|
244
|
+
else:
|
|
245
|
+
result.needs_confirmation = True
|
|
246
|
+
if self.confirm_callback:
|
|
247
|
+
if self.confirm_callback(result.sql, violations):
|
|
248
|
+
should_execute = True
|
|
249
|
+
else:
|
|
250
|
+
result.execution_time_ms = (time.time() - t_start) * 1000
|
|
251
|
+
return result
|
|
252
|
+
else:
|
|
253
|
+
result.execution_time_ms = (time.time() - t_start) * 1000
|
|
254
|
+
return result
|
|
255
|
+
elif self.mode == OperationMode.ADMIN:
|
|
256
|
+
should_execute = True
|
|
257
|
+
|
|
258
|
+
if should_execute:
|
|
259
|
+
# 1. Attempt direct execution without AI (Zero AI Latency)
|
|
260
|
+
cols, rows, error = self.connector.execute_query(raw_sql, limit=100)
|
|
261
|
+
if error is None:
|
|
262
|
+
result.columns = cols
|
|
263
|
+
result.rows = rows
|
|
264
|
+
result.was_executed = True
|
|
265
|
+
result.error = None
|
|
266
|
+
if self.mask_sensitive_data and result.rows:
|
|
267
|
+
result.rows = SensitiveDataMasker.mask_row_set(result.columns, result.rows)
|
|
268
|
+
if result.rows:
|
|
269
|
+
result.visualization = ResultVisualizer.analyze_result_shape(result.columns, result.rows)
|
|
270
|
+
if self.use_cache and result.query_type == "read":
|
|
271
|
+
self.result_cache.set(self.db_name, raw_sql, result.columns, result.rows)
|
|
272
|
+
exec_duration = (time.time() - t_start) * 1000
|
|
273
|
+
result.execution_time_ms = exec_duration
|
|
274
|
+
self.audit_logger.log_execution(
|
|
275
|
+
database=self.db_name,
|
|
276
|
+
user=self.user_id,
|
|
277
|
+
question=user_prompt,
|
|
278
|
+
sql=result.sql,
|
|
279
|
+
mode=self.mode.value,
|
|
280
|
+
status="SUCCESS",
|
|
281
|
+
execution_time_ms=exec_duration,
|
|
282
|
+
rows_returned=len(result.rows),
|
|
283
|
+
risk_level="LOW",
|
|
284
|
+
)
|
|
285
|
+
self.conversation.add_turn(user_prompt=user_prompt, sql=result.sql, is_chat=False)
|
|
286
|
+
return result
|
|
287
|
+
|
|
288
|
+
# 2. Direct query failed with DB error -> Fallback to AI Auto-Repair
|
|
289
|
+
try:
|
|
290
|
+
exact, fuzzy = self.connector.resolve_tables(raw_sql)
|
|
291
|
+
except Exception:
|
|
292
|
+
exact, fuzzy = [], []
|
|
293
|
+
sub_schema, join_steps = self._get_fast_schema_context(raw_sql, exact, fuzzy)
|
|
294
|
+
result.schema = sub_schema
|
|
295
|
+
result.join_steps = join_steps
|
|
296
|
+
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}"
|
|
297
|
+
self._execute_with_retry(result, repair_prompt, sub_schema)
|
|
298
|
+
|
|
299
|
+
if result.success:
|
|
300
|
+
if self.mask_sensitive_data and result.rows:
|
|
301
|
+
result.rows = SensitiveDataMasker.mask_row_set(result.columns, result.rows)
|
|
302
|
+
if result.rows:
|
|
303
|
+
result.visualization = ResultVisualizer.analyze_result_shape(result.columns, result.rows)
|
|
304
|
+
if self.use_cache and result.query_type == "read":
|
|
305
|
+
self.result_cache.set(self.db_name, result.sql, result.columns, result.rows)
|
|
306
|
+
self.conversation.add_turn(user_prompt=user_prompt, sql=result.sql, is_chat=False)
|
|
307
|
+
|
|
308
|
+
exec_duration = (time.time() - t_start) * 1000
|
|
309
|
+
result.execution_time_ms = exec_duration
|
|
310
|
+
self.audit_logger.log_execution(
|
|
311
|
+
database=self.db_name,
|
|
312
|
+
user=self.user_id,
|
|
313
|
+
question=user_prompt,
|
|
314
|
+
sql=result.sql,
|
|
315
|
+
mode=self.mode.value,
|
|
316
|
+
status="SUCCESS" if result.success else "FAILED",
|
|
317
|
+
execution_time_ms=exec_duration,
|
|
318
|
+
rows_returned=len(result.rows),
|
|
319
|
+
risk_level=result.complexity.risk_level if result.complexity else "LOW",
|
|
320
|
+
error=result.error,
|
|
321
|
+
)
|
|
322
|
+
return result
|
|
323
|
+
|
|
202
324
|
# --- 2. Resolve Tables & Detect Ambiguity ---
|
|
203
325
|
try:
|
|
204
326
|
exact, fuzzy = self.connector.resolve_tables(user_prompt)
|
|
@@ -687,16 +687,18 @@ def chat_command(
|
|
|
687
687
|
if not cmd_arg:
|
|
688
688
|
curr_model = model or getattr(llm, "default_model", "default")
|
|
689
689
|
console.print(f"\n[bold cyan]Active AI Provider:[/bold cyan] [bold white]{llm.name}[/bold white] | Model: [bold green]{curr_model}[/bold green]")
|
|
690
|
-
console.print("[dim]
|
|
690
|
+
console.print("[dim]Supported: [bold]claude, openai, deepseek, mistral, gemini, groq, openrouter, ollama[/bold][/dim]\n")
|
|
691
691
|
continue
|
|
692
|
-
prov_target = cmd_arg.lower()
|
|
693
|
-
if prov_target in ["
|
|
692
|
+
prov_target = cmd_arg.lower().strip()
|
|
693
|
+
if prov_target in ["claude", "anthropic", "openai", "chatgpt", "gpt", "deepseek", "mistral", "gemini", "groq", "openrouter", "ollama", "local", "mock"]:
|
|
694
694
|
llm = get_llm_provider(provider_name=prov_target, config_mgr=config_mgr)
|
|
695
695
|
model = None
|
|
696
696
|
pipeline.generator.llm = llm # Update reused pipeline
|
|
697
697
|
print_success(f"Switched AI provider to [bold cyan]{llm.name}[/bold cyan]")
|
|
698
|
+
if not llm.is_available() and llm.name not in ["ollama", "mock"]:
|
|
699
|
+
print_warning(f"Note: API key for {llm.name} is not set. Run `db-agent config` or set {llm.name.upper()}_API_KEY.")
|
|
698
700
|
else:
|
|
699
|
-
print_error(f"Unknown provider '{cmd_arg}'. Supported:
|
|
701
|
+
print_error(f"Unknown provider '{cmd_arg}'. Supported: claude, openai, deepseek, mistral, gemini, groq, openrouter, ollama")
|
|
700
702
|
continue
|
|
701
703
|
|
|
702
704
|
if cmd_name in [":model", ":models", ":providers", ":ai"]:
|
|
@@ -710,18 +712,41 @@ def chat_command(
|
|
|
710
712
|
if avail_models:
|
|
711
713
|
console.print(f"[bold]Available Models for {llm.name}:[/bold]")
|
|
712
714
|
model_descriptions = {
|
|
713
|
-
|
|
715
|
+
# Gemini
|
|
716
|
+
"gemini-flash-lite-latest": "⚡ Ultra-Fast (~0.8s) • Instant chat & queries",
|
|
714
717
|
"gemini-3.6-flash": "⚡ High Quality (~3.5s) • Strong SQL reasoning & accuracy",
|
|
715
718
|
"gemini-3.7-flash": "🧠 Deep Reasoning (~13s) • Latest flagship model",
|
|
716
719
|
"gemini-3.1-flash-lite": "⚡ Fast Lite (~4.0s) • Low latency",
|
|
717
720
|
"gemini-flash-latest": "⚡ Standard Flash (~7.0s)",
|
|
721
|
+
# Claude
|
|
722
|
+
"claude-3-7-sonnet-20250219": "🧠 Hybrid Reasoning Flagship • Deepest architecture understanding",
|
|
723
|
+
"claude-3-5-sonnet-20241022": "⚡ High Intelligence • Expert SQL engineer & schema specialist",
|
|
724
|
+
"claude-3-5-haiku-20241022": "⚡ Ultra-Fast (~1s) • High speed and cost efficient",
|
|
725
|
+
"claude-3-opus-20240229": "Deep complex schema reasoning",
|
|
726
|
+
# OpenAI
|
|
727
|
+
"gpt-4o": "⚡ Omni Flagship • High speed multimodal & complex SQL reasoning",
|
|
728
|
+
"gpt-4o-mini": "⚡ Ultra-Fast (~1s) • High speed and cost efficient",
|
|
729
|
+
"gpt-4.5-preview": "🧠 Next-Generation Reasoning model",
|
|
730
|
+
"o3-mini": "🧠 Deep STEM & Advanced Coding Reasoning",
|
|
731
|
+
"o1": "🧠 Maximum Reasoning & Math Architecture",
|
|
732
|
+
# DeepSeek
|
|
733
|
+
"deepseek-chat": "⚡ DeepSeek-V3 Ultra-Fast • Strong code & SQL",
|
|
734
|
+
"deepseek-reasoner": "🧠 DeepSeek-R1 Full Chain-of-Thought Reasoning",
|
|
735
|
+
# Mistral
|
|
736
|
+
"codestral-latest": "⚡ Specialized Code & SQL Architect (22B)",
|
|
737
|
+
"mistral-large-latest": "🧠 Flagship Reasoning & Multi-language",
|
|
738
|
+
"mistral-small-latest": "⚡ Fast & lightweight",
|
|
739
|
+
# Groq
|
|
740
|
+
"llama-3.3-70b-versatile": "⚡ Llama 3.3 70B (500+ tokens/sec on Groq LPUs)",
|
|
741
|
+
"llama-3.1-8b-instant": "⚡ Ultra-Fast Llama 3.1 8B (800+ tokens/sec)",
|
|
742
|
+
"qwen-2.5-coder-32b": "⚡ Qwen 2.5 Coder 32B specialized SQL engineer",
|
|
718
743
|
}
|
|
719
744
|
for idx, m_name in enumerate(avail_models, 1):
|
|
720
745
|
is_active = " [bold green]✓ Active[/bold green]" if m_name == curr_model else ""
|
|
721
746
|
desc = model_descriptions.get(m_name, "")
|
|
722
747
|
desc_str = f" - [dim]{desc}[/dim]" if desc else ""
|
|
723
748
|
console.print(f" [bold cyan]{idx}[/bold cyan]. [bold white]{m_name}[/bold white]{desc_str}{is_active}")
|
|
724
|
-
console.print("\n[dim]Quick Switch: [bold]:model
|
|
749
|
+
console.print("\n[dim]Quick Switch: [bold]:model <number|name>[/bold] (e.g. [bold]:model 1[/bold] or [bold]:model 3.7[/bold] or [bold]:model 4o[/bold] or [bold]:model r1[/bold])[/dim]\n")
|
|
725
750
|
continue
|
|
726
751
|
|
|
727
752
|
# Handle numeric shortcut or name shortcut
|
|
@@ -731,13 +756,23 @@ def chat_command(
|
|
|
731
756
|
if 0 <= choice_idx < len(avail_models):
|
|
732
757
|
target_model = avail_models[choice_idx]
|
|
733
758
|
elif target_model in ["3.7", "37"]:
|
|
734
|
-
target_model = "gemini-3.7-flash"
|
|
759
|
+
target_model = "claude-3-7-sonnet-20250219" if llm.name == "claude" else "gemini-3.7-flash"
|
|
760
|
+
elif target_model in ["3.5", "35"]:
|
|
761
|
+
target_model = "claude-3-5-sonnet-20241022" if llm.name == "claude" else "gemini-3.5-flash"
|
|
735
762
|
elif target_model in ["3.6", "36"]:
|
|
736
763
|
target_model = "gemini-3.6-flash"
|
|
737
|
-
elif target_model in ["
|
|
738
|
-
target_model = "
|
|
739
|
-
elif target_model in ["
|
|
740
|
-
target_model = "
|
|
764
|
+
elif target_model in ["4o", "gpt4o"]:
|
|
765
|
+
target_model = "gpt-4o"
|
|
766
|
+
elif target_model in ["mini", "4o-mini"]:
|
|
767
|
+
target_model = "gpt-4o-mini"
|
|
768
|
+
elif target_model in ["r1", "reasoner"]:
|
|
769
|
+
target_model = "deepseek-reasoner"
|
|
770
|
+
elif target_model in ["v3", "chat"]:
|
|
771
|
+
target_model = "deepseek-chat"
|
|
772
|
+
elif target_model in ["codestral"]:
|
|
773
|
+
target_model = "codestral-latest"
|
|
774
|
+
elif target_model in ["lite", "fast", "flash-lite"]:
|
|
775
|
+
target_model = "gemini-flash-lite-latest" if llm.name == "gemini" else "claude-3-5-haiku-20241022"
|
|
741
776
|
elif avail_models:
|
|
742
777
|
matches = [m for m in avail_models if target_model.lower() in m.lower()]
|
|
743
778
|
if matches:
|
|
@@ -1586,7 +1621,7 @@ def run_command(
|
|
|
1586
1621
|
|
|
1587
1622
|
@app.command(name="config")
|
|
1588
1623
|
def config_command():
|
|
1589
|
-
"""Interactive wizard to configure AI providers and
|
|
1624
|
+
"""Interactive wizard to configure AI providers and credentials."""
|
|
1590
1625
|
print_banner()
|
|
1591
1626
|
console.print("[bold cyan]DB-Agent Configuration Wizard[/bold cyan]\n")
|
|
1592
1627
|
|
|
@@ -1594,14 +1629,38 @@ def config_command():
|
|
|
1594
1629
|
console.print(f"Current Default AI Provider: [bold green]{curr_provider}[/bold green]")
|
|
1595
1630
|
provider_choice = Prompt.ask(
|
|
1596
1631
|
"Select Default AI Provider",
|
|
1597
|
-
choices=["
|
|
1632
|
+
choices=["gemini", "claude", "openai", "deepseek", "mistral", "groq", "openrouter", "ollama", "keep"],
|
|
1598
1633
|
default="keep",
|
|
1599
1634
|
)
|
|
1600
1635
|
if provider_choice != "keep":
|
|
1601
1636
|
config_mgr.set_setting("default_provider", provider_choice)
|
|
1602
1637
|
print_success(f"Default provider updated to [bold]{provider_choice}[/bold]")
|
|
1603
1638
|
|
|
1604
|
-
if provider_choice
|
|
1639
|
+
if provider_choice in ["claude", "anthropic"]:
|
|
1640
|
+
console.print("[dim]Get your Anthropic API key at: [bold]https://console.anthropic.com[/bold][/dim]")
|
|
1641
|
+
key = Prompt.ask("Enter ANTHROPIC_API_KEY", default=config_mgr.get_setting("anthropic_api_key", "")).strip()
|
|
1642
|
+
if key:
|
|
1643
|
+
config_mgr.set_setting("anthropic_api_key", key)
|
|
1644
|
+
print_success("Anthropic Claude API key saved.")
|
|
1645
|
+
elif provider_choice in ["openai", "chatgpt"]:
|
|
1646
|
+
console.print("[dim]Get your OpenAI API key at: [bold]https://platform.openai.com/api-keys[/bold][/dim]")
|
|
1647
|
+
key = Prompt.ask("Enter OPENAI_API_KEY", default=config_mgr.get_setting("openai_api_key", "")).strip()
|
|
1648
|
+
if key:
|
|
1649
|
+
config_mgr.set_setting("openai_api_key", key)
|
|
1650
|
+
print_success("OpenAI API key saved.")
|
|
1651
|
+
elif provider_choice == "deepseek":
|
|
1652
|
+
console.print("[dim]Get your DeepSeek API key at: [bold]https://platform.deepseek.com[/bold][/dim]")
|
|
1653
|
+
key = Prompt.ask("Enter DEEPSEEK_API_KEY", default=config_mgr.get_setting("deepseek_api_key", "")).strip()
|
|
1654
|
+
if key:
|
|
1655
|
+
config_mgr.set_setting("deepseek_api_key", key)
|
|
1656
|
+
print_success("DeepSeek API key saved.")
|
|
1657
|
+
elif provider_choice == "mistral":
|
|
1658
|
+
console.print("[dim]Get your Mistral API key at: [bold]https://console.mistral.ai[/bold][/dim]")
|
|
1659
|
+
key = Prompt.ask("Enter MISTRAL_API_KEY", default=config_mgr.get_setting("mistral_api_key", "")).strip()
|
|
1660
|
+
if key:
|
|
1661
|
+
config_mgr.set_setting("mistral_api_key", key)
|
|
1662
|
+
print_success("Mistral API key saved.")
|
|
1663
|
+
elif provider_choice == "gemini":
|
|
1605
1664
|
console.print("[dim]Get your free Gemini API key at: [bold]https://aistudio.google.com/apikey[/bold][/dim]")
|
|
1606
1665
|
key = Prompt.ask("Enter GEMINI_API_KEY", default=config_mgr.get_setting("gemini_api_key", "")).strip()
|
|
1607
1666
|
if key:
|
|
@@ -1652,33 +1711,29 @@ def profiles_command():
|
|
|
1652
1711
|
|
|
1653
1712
|
@app.command(name="models")
|
|
1654
1713
|
def models_command():
|
|
1655
|
-
"""List available models across local Ollama and
|
|
1714
|
+
"""List available models across local Ollama and all major cloud providers."""
|
|
1656
1715
|
print_banner()
|
|
1657
|
-
console.print("[bold cyan]Supported
|
|
1658
|
-
|
|
1659
|
-
|
|
1660
|
-
|
|
1661
|
-
|
|
1662
|
-
|
|
1663
|
-
|
|
1664
|
-
|
|
1665
|
-
|
|
1666
|
-
|
|
1667
|
-
|
|
1668
|
-
|
|
1669
|
-
gemini_status = "[green]Configured[/green]" if gemini.is_available() else "[yellow]Not Configured (Free Key: aistudio.google.com)[/yellow]"
|
|
1670
|
-
console.print(f"\n2. [bold white]Google Gemini (Free Tier)[/bold white] - Status: {gemini_status}")
|
|
1671
|
-
console.print(" Available Models: `gemini-3.6-flash`, `gemini-2.5-flash`, `gemini-2.5-pro`")
|
|
1672
|
-
|
|
1673
|
-
groq = get_llm_provider("groq")
|
|
1674
|
-
groq_status = "[green]Configured[/green]" if groq.is_available() else "[yellow]Not Configured (Free Key: console.groq.com)[/yellow]"
|
|
1675
|
-
console.print(f"\n3. [bold white]Groq (Free Ultra-Fast Tier)[/bold white] - Status: {groq_status}")
|
|
1676
|
-
console.print(" Available Models: `llama-3.3-70b-versatile`, `qwen-2.5-32b`")
|
|
1716
|
+
console.print("[bold cyan]Supported AI Providers & Available Models:[/bold cyan]\n")
|
|
1717
|
+
|
|
1718
|
+
providers_to_show = [
|
|
1719
|
+
("gemini", "Google Gemini", "https://aistudio.google.com/apikey"),
|
|
1720
|
+
("claude", "Anthropic Claude", "https://console.anthropic.com"),
|
|
1721
|
+
("openai", "OpenAI / ChatGPT", "https://platform.openai.com"),
|
|
1722
|
+
("deepseek", "DeepSeek AI", "https://platform.deepseek.com"),
|
|
1723
|
+
("mistral", "Mistral AI", "https://console.mistral.ai"),
|
|
1724
|
+
("groq", "Groq (Ultra-Fast LPUs)", "https://console.groq.com"),
|
|
1725
|
+
("openrouter", "OpenRouter (100+ Models)", "https://openrouter.ai"),
|
|
1726
|
+
("ollama", "Ollama (100% Offline / Local)", "http://localhost:11434"),
|
|
1727
|
+
]
|
|
1677
1728
|
|
|
1678
|
-
|
|
1679
|
-
|
|
1680
|
-
|
|
1681
|
-
|
|
1729
|
+
for idx, (p_slug, p_label, p_url) in enumerate(providers_to_show, 1):
|
|
1730
|
+
prov = get_llm_provider(p_slug, config_mgr=config_mgr)
|
|
1731
|
+
status = "[bold green]✓ Configured & Ready[/bold green]" if prov.is_available() else f"[yellow]Not Configured ({p_url})[/yellow]"
|
|
1732
|
+
console.print(f"{idx}. [bold white]{p_label}[/bold white] - Status: {status}")
|
|
1733
|
+
models = prov.list_models() if hasattr(prov, "list_models") else []
|
|
1734
|
+
if models:
|
|
1735
|
+
console.print(f" Models: [cyan]{', '.join(models[:5])}[/cyan]")
|
|
1736
|
+
console.print("")
|
|
1682
1737
|
|
|
1683
1738
|
|
|
1684
1739
|
# ---- setup ----
|
|
@@ -1699,7 +1754,7 @@ def setup_command(
|
|
|
1699
1754
|
else:
|
|
1700
1755
|
console.print("[bold white]Choose how you would like to setup DB-Agent:[/bold white]")
|
|
1701
1756
|
console.print(" [bold green]1. Local Offline AI[/bold green] (Ollama + Coding Model) - [italic]100% Free, Private, Zero API Keys[/italic]")
|
|
1702
|
-
console.print(" [bold cyan]2.
|
|
1757
|
+
console.print(" [bold cyan]2. Cloud AI[/bold cyan] (Claude / OpenAI / DeepSeek / Mistral / Gemini / Groq) - [italic]High speed & accuracy[/italic]")
|
|
1703
1758
|
console.print(" [bold yellow]3. Normal Database CLI[/bold yellow] - [italic]No AI, use standard connect, scan & SQL tools only[/italic]\n")
|
|
1704
1759
|
|
|
1705
1760
|
choice = Prompt.ask(
|
|
@@ -1720,17 +1775,41 @@ def setup_command(
|
|
|
1720
1775
|
print_warning("Local setup incomplete. You can still use cloud AI or standard DB tools.")
|
|
1721
1776
|
chosen_provider = "cloud"
|
|
1722
1777
|
elif choice == "2":
|
|
1723
|
-
console.print("\n[bold cyan]Setting up
|
|
1724
|
-
prov = Prompt.ask("Choose provider", choices=["gemini", "groq", "openrouter"], default="gemini")
|
|
1778
|
+
console.print("\n[bold cyan]Setting up Cloud AI Provider...[/bold cyan]")
|
|
1779
|
+
prov = Prompt.ask("Choose provider", choices=["gemini", "claude", "openai", "deepseek", "mistral", "groq", "openrouter"], default="gemini")
|
|
1725
1780
|
config_mgr.set_setting("default_provider", prov)
|
|
1726
|
-
if prov
|
|
1727
|
-
console.print("[dim]Get
|
|
1781
|
+
if prov in ["claude", "anthropic"]:
|
|
1782
|
+
console.print("[dim]Get Anthropic key at: [bold]https://console.anthropic.com[/bold][/dim]")
|
|
1783
|
+
k = Prompt.ask("Enter ANTHROPIC_API_KEY", default="")
|
|
1784
|
+
if k:
|
|
1785
|
+
config_mgr.set_setting("anthropic_api_key", k)
|
|
1786
|
+
print_success("Anthropic Claude API key saved.")
|
|
1787
|
+
elif prov in ["openai", "chatgpt"]:
|
|
1788
|
+
console.print("[dim]Get OpenAI key at: [bold]https://platform.openai.com/api-keys[/bold][/dim]")
|
|
1789
|
+
k = Prompt.ask("Enter OPENAI_API_KEY", default="")
|
|
1790
|
+
if k:
|
|
1791
|
+
config_mgr.set_setting("openai_api_key", k)
|
|
1792
|
+
print_success("OpenAI API key saved.")
|
|
1793
|
+
elif prov == "deepseek":
|
|
1794
|
+
console.print("[dim]Get DeepSeek key at: [bold]https://platform.deepseek.com[/bold][/dim]")
|
|
1795
|
+
k = Prompt.ask("Enter DEEPSEEK_API_KEY", default="")
|
|
1796
|
+
if k:
|
|
1797
|
+
config_mgr.set_setting("deepseek_api_key", k)
|
|
1798
|
+
print_success("DeepSeek API key saved.")
|
|
1799
|
+
elif prov == "mistral":
|
|
1800
|
+
console.print("[dim]Get Mistral key at: [bold]https://console.mistral.ai[/bold][/dim]")
|
|
1801
|
+
k = Prompt.ask("Enter MISTRAL_API_KEY", default="")
|
|
1802
|
+
if k:
|
|
1803
|
+
config_mgr.set_setting("mistral_api_key", k)
|
|
1804
|
+
print_success("Mistral API key saved.")
|
|
1805
|
+
elif prov == "gemini":
|
|
1806
|
+
console.print("[dim]Get free Gemini key at: [bold]https://aistudio.google.com/apikey[/bold][/dim]")
|
|
1728
1807
|
k = Prompt.ask("Enter GEMINI_API_KEY", default="")
|
|
1729
1808
|
if k:
|
|
1730
1809
|
config_mgr.set_setting("gemini_api_key", k)
|
|
1731
1810
|
print_success("Gemini API key saved.")
|
|
1732
1811
|
elif prov == "groq":
|
|
1733
|
-
console.print("[dim]Get free key at: [bold]https://console.groq.com[/bold][/dim]")
|
|
1812
|
+
console.print("[dim]Get free Groq key at: [bold]https://console.groq.com[/bold][/dim]")
|
|
1734
1813
|
k = Prompt.ask("Enter GROQ_API_KEY", default="")
|
|
1735
1814
|
if k:
|
|
1736
1815
|
config_mgr.set_setting("groq_api_key", k)
|
|
@@ -0,0 +1,126 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Anthropic Claude LLM Provider.
|
|
3
|
+
Supports claude-3-7-sonnet-20250219, claude-3-5-sonnet-20241022, claude-3-5-haiku-20241022, claude-3-opus-20240229.
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
import json
|
|
7
|
+
import os
|
|
8
|
+
import requests
|
|
9
|
+
from typing import Optional, Callable, List
|
|
10
|
+
from dbagent.llm.base import BaseLLMProvider
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class ClaudeProvider(BaseLLMProvider):
|
|
14
|
+
"""Anthropic Claude provider using direct HTTP REST API."""
|
|
15
|
+
|
|
16
|
+
FALLBACK_MODELS = [
|
|
17
|
+
"claude-3-7-sonnet-20250219",
|
|
18
|
+
"claude-3-5-sonnet-20241022",
|
|
19
|
+
"claude-3-5-haiku-20241022",
|
|
20
|
+
"claude-3-opus-20240229",
|
|
21
|
+
]
|
|
22
|
+
|
|
23
|
+
def __init__(self, api_key: Optional[str] = None, default_model: str = "claude-3-7-sonnet-20250219"):
|
|
24
|
+
if api_key is not None:
|
|
25
|
+
self.api_key = api_key
|
|
26
|
+
else:
|
|
27
|
+
self.api_key = os.getenv("ANTHROPIC_API_KEY") or os.getenv("CLAUDE_API_KEY")
|
|
28
|
+
self.default_model = default_model
|
|
29
|
+
self._session = requests.Session()
|
|
30
|
+
self._session.headers.update({
|
|
31
|
+
"User-Agent": "DB-Agent/0.7.0",
|
|
32
|
+
"anthropic-version": "2023-06-01",
|
|
33
|
+
"content-type": "application/json",
|
|
34
|
+
})
|
|
35
|
+
if self.api_key:
|
|
36
|
+
self._session.headers.update({"x-api-key": self.api_key})
|
|
37
|
+
|
|
38
|
+
@property
|
|
39
|
+
def name(self) -> str:
|
|
40
|
+
return "claude"
|
|
41
|
+
|
|
42
|
+
def is_available(self) -> bool:
|
|
43
|
+
return bool(self.api_key)
|
|
44
|
+
|
|
45
|
+
def list_models(self) -> List[str]:
|
|
46
|
+
return self.FALLBACK_MODELS
|
|
47
|
+
|
|
48
|
+
def generate(
|
|
49
|
+
self,
|
|
50
|
+
prompt: str,
|
|
51
|
+
system_prompt: Optional[str] = None,
|
|
52
|
+
model: Optional[str] = None,
|
|
53
|
+
stream_callback: Optional[Callable[[str], None]] = None,
|
|
54
|
+
) -> str:
|
|
55
|
+
if not self.api_key:
|
|
56
|
+
raise ValueError(
|
|
57
|
+
"Anthropic API key is not configured. Set ANTHROPIC_API_KEY environment variable or run `db-agent config`."
|
|
58
|
+
)
|
|
59
|
+
|
|
60
|
+
requested_model = model or self.default_model
|
|
61
|
+
models_to_try = [requested_model]
|
|
62
|
+
for fb in self.FALLBACK_MODELS:
|
|
63
|
+
if fb != requested_model and fb not in models_to_try:
|
|
64
|
+
models_to_try.append(fb)
|
|
65
|
+
|
|
66
|
+
payload = {
|
|
67
|
+
"model": requested_model,
|
|
68
|
+
"max_tokens": 4096,
|
|
69
|
+
"temperature": 0.1,
|
|
70
|
+
"messages": [{"role": "user", "content": prompt}],
|
|
71
|
+
}
|
|
72
|
+
if system_prompt:
|
|
73
|
+
payload["system"] = system_prompt
|
|
74
|
+
|
|
75
|
+
endpoint = "https://api.anthropic.com/v1/messages"
|
|
76
|
+
last_error = None
|
|
77
|
+
|
|
78
|
+
for target_model in models_to_try:
|
|
79
|
+
payload["model"] = target_model
|
|
80
|
+
try:
|
|
81
|
+
if stream_callback:
|
|
82
|
+
payload["stream"] = True
|
|
83
|
+
response = self._session.post(endpoint, json=payload, stream=True, timeout=15)
|
|
84
|
+
if response.status_code == 200:
|
|
85
|
+
full_text = []
|
|
86
|
+
for line in response.iter_lines():
|
|
87
|
+
if line:
|
|
88
|
+
decoded = line.decode("utf-8")
|
|
89
|
+
if decoded.startswith("data: "):
|
|
90
|
+
raw_json = decoded[6:]
|
|
91
|
+
try:
|
|
92
|
+
chunk = json.loads(raw_json)
|
|
93
|
+
delta = chunk.get("delta", {})
|
|
94
|
+
text_piece = delta.get("text", "")
|
|
95
|
+
if text_piece:
|
|
96
|
+
full_text.append(text_piece)
|
|
97
|
+
stream_callback(text_piece)
|
|
98
|
+
except Exception:
|
|
99
|
+
pass
|
|
100
|
+
return "".join(full_text)
|
|
101
|
+
elif response.status_code in (404, 400, 503, 429):
|
|
102
|
+
last_error = f"Model {target_model} status {response.status_code}"
|
|
103
|
+
continue
|
|
104
|
+
else:
|
|
105
|
+
raise RuntimeError(f"Anthropic API error ({response.status_code}): {response.text}")
|
|
106
|
+
else:
|
|
107
|
+
payload["stream"] = False
|
|
108
|
+
response = self._session.post(endpoint, json=payload, timeout=20)
|
|
109
|
+
if response.status_code == 200:
|
|
110
|
+
data = response.json()
|
|
111
|
+
content_blocks = data.get("content", [])
|
|
112
|
+
text_blocks = [b.get("text", "") for b in content_blocks if b.get("type") == "text"]
|
|
113
|
+
return "".join(text_blocks)
|
|
114
|
+
elif response.status_code in (404, 400, 503, 429):
|
|
115
|
+
last_error = f"Model {target_model} status {response.status_code}"
|
|
116
|
+
continue
|
|
117
|
+
else:
|
|
118
|
+
raise RuntimeError(f"Anthropic API error ({response.status_code}): {response.text}")
|
|
119
|
+
except requests.exceptions.Timeout:
|
|
120
|
+
last_error = f"Model {target_model} timed out"
|
|
121
|
+
continue
|
|
122
|
+
except Exception as e:
|
|
123
|
+
last_error = str(e)
|
|
124
|
+
continue
|
|
125
|
+
|
|
126
|
+
raise RuntimeError(f"Anthropic generation error: {last_error or 'All model attempts failed'}")
|