dbagent-cli 0.3.3__tar.gz → 0.3.5__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.
Files changed (74) hide show
  1. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/PKG-INFO +2 -2
  2. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/__init__.py +1 -1
  3. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/connectors/relational.py +36 -60
  4. dbagent_cli-0.3.5/dbagent/llm/gemini_provider.py +131 -0
  5. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/ui/console.py +25 -10
  6. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent_cli.egg-info/PKG-INFO +2 -2
  7. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/pyproject.toml +2 -2
  8. dbagent_cli-0.3.3/dbagent/llm/gemini_provider.py +0 -110
  9. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/README.md +0 -0
  10. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/agent/confidence.py +0 -0
  11. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/agent/conversation.py +0 -0
  12. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/agent/doctor.py +0 -0
  13. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/agent/error_classifier.py +0 -0
  14. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/agent/generator.py +0 -0
  15. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/agent/pipeline.py +0 -0
  16. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/agent/recommender.py +0 -0
  17. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/agent/validator.py +0 -0
  18. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/cli.py +0 -0
  19. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/config.py +0 -0
  20. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/connectors/base.py +0 -0
  21. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/connectors/factory.py +0 -0
  22. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/connectors/mongo.py +0 -0
  23. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/execution/__init__.py +0 -0
  24. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/execution/cache.py +0 -0
  25. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/execution/changeset.py +0 -0
  26. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/execution/cost.py +0 -0
  27. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/execution/recovery.py +0 -0
  28. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/execution/sandbox.py +0 -0
  29. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/execution/streaming.py +0 -0
  30. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/llm/auto_setup.py +0 -0
  31. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/llm/base.py +0 -0
  32. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/llm/factory.py +0 -0
  33. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/llm/groq_provider.py +0 -0
  34. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/llm/mock_provider.py +0 -0
  35. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/llm/ollama_provider.py +0 -0
  36. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/llm/openrouter_provider.py +0 -0
  37. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/schema/cache.py +0 -0
  38. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/schema/formatter.py +0 -0
  39. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/schema/graph.py +0 -0
  40. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/schema/models.py +0 -0
  41. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/schema/retriever.py +0 -0
  42. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/schema/selector.py +0 -0
  43. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/security/__init__.py +0 -0
  44. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/security/audit.py +0 -0
  45. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/security/environment.py +0 -0
  46. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/security/masking.py +0 -0
  47. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/security/modes.py +0 -0
  48. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/security/rbac.py +0 -0
  49. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/semantic/__init__.py +0 -0
  50. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/semantic/dictionary.py +0 -0
  51. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/semantic/templates.py +0 -0
  52. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/sql/__init__.py +0 -0
  53. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/sql/optimizer.py +0 -0
  54. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/sql/validator.py +0 -0
  55. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/ui/viewer.py +0 -0
  56. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent/ui/visualizer.py +0 -0
  57. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent_cli.egg-info/SOURCES.txt +0 -0
  58. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent_cli.egg-info/dependency_links.txt +0 -0
  59. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent_cli.egg-info/entry_points.txt +0 -0
  60. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent_cli.egg-info/requires.txt +0 -0
  61. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/dbagent_cli.egg-info/top_level.txt +0 -0
  62. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/setup.cfg +0 -0
  63. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_cache.py +0 -0
  64. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_cli.py +0 -0
  65. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_connectors.py +0 -0
  66. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_doctor_and_viz.py +0 -0
  67. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_generator.py +0 -0
  68. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_graph.py +0 -0
  69. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_recovery.py +0 -0
  70. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_safety.py +0 -0
  71. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_sandbox.py +0 -0
  72. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_schema.py +0 -0
  73. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_security.py +0 -0
  74. {dbagent_cli-0.3.3 → dbagent_cli-0.3.5}/tests/test_semantic.py +0 -0
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dbagent-cli
3
- Version: 0.3.3
4
- Summary: Enterprise AI Database Copilot with Instant Schema Caching & Zero-Latency Chat (CLI)
3
+ Version: 0.3.5
4
+ Summary: Production-Grade AI Database Copilot with Ultra-Fast Inference & Clean UI (CLI)
5
5
  Author: Santhosh Gupta
6
6
  Requires-Python: >=3.9
7
7
  Description-Content-Type: text/markdown
@@ -2,5 +2,5 @@
2
2
  DB-Agent: Universal Database Introspection and Script Generation AI Agent (CLI).
3
3
  """
4
4
 
5
- __version__ = "0.3.3"
5
+ __version__ = "0.3.5"
6
6
  __author__ = "Santhosh Gupta"
@@ -65,45 +65,41 @@ def _fuzzy_match_tables(table_names: List[str], user_prompt: str) -> Tuple[List[
65
65
 
66
66
  for t_name in table_names:
67
67
  t_lower = t_name.lower()
68
+ t_singular = t_lower.rstrip("s")
69
+ t_no_ies = re.sub(r"ies$", "y", t_lower)
68
70
 
69
- # 1. Exact word boundary match
71
+ # 1. Exact word boundary or exact singular/plural match
72
+ is_exact = False
70
73
  if re.search(r"\b" + re.escape(t_lower) + r"\b", prompt_lower):
74
+ is_exact = True
75
+ elif t_lower in prompt_words or t_singular in prompt_words or t_no_ies in prompt_words:
76
+ is_exact = True
77
+ else:
78
+ for w in prompt_words:
79
+ w_sing = w.rstrip("s")
80
+ if w == t_lower or w == t_singular or w_sing == t_lower or w_sing == t_singular:
81
+ is_exact = True
82
+ break
83
+
84
+ if is_exact:
71
85
  exact.append(t_name)
72
86
  continue
73
87
 
74
- # 2. Substring match (table name mentioned in prompt)
75
- if t_lower in prompt_lower:
88
+ # 2. Substring match (full table name appears in prompt)
89
+ if t_lower in prompt_lower or t_singular in prompt_lower:
76
90
  exact.append(t_name)
77
91
  continue
78
92
 
79
- # 3. Singular/plural normalization
80
- # "users" matches "user", "categories" matches "category"
81
- t_singular = t_lower.rstrip("s")
82
- t_no_ies = re.sub(r"ies$", "y", t_lower)
93
+ # 3. Partial prefix/suffix/component match (fuzzy)
83
94
  for word in prompt_words:
84
- w_singular = word.rstrip("s")
85
- w_no_ies = re.sub(r"ies$", "y", word)
86
- # Check various forms
87
- if (word == t_lower or
88
- word == t_singular or
89
- w_singular == t_lower or
90
- w_singular == t_singular or
91
- word == t_no_ies or
92
- w_no_ies == t_lower):
93
- fuzzy.append(t_name)
94
- break
95
- else:
96
- # 4. Partial prefix/suffix match (e.g. "user" matches "user_logins", "app_users")
97
- for word in prompt_words:
98
- if len(word) >= 4:
99
- if t_lower.startswith(word) or t_lower.endswith(word):
100
- fuzzy.append(t_name)
101
- break
102
- # Check if word is a component of underscore-separated name
103
- t_parts = t_lower.split("_")
104
- if word in t_parts or word.rstrip("s") in t_parts:
105
- fuzzy.append(t_name)
106
- break
95
+ if len(word) >= 4:
96
+ t_parts = t_lower.split("_")
97
+ if word in t_parts or word.rstrip("s") in t_parts:
98
+ fuzzy.append(t_name)
99
+ break
100
+ if t_lower.startswith(word) or t_lower.endswith(word):
101
+ fuzzy.append(t_name)
102
+ break
107
103
 
108
104
  return exact, fuzzy
109
105
 
@@ -249,44 +245,24 @@ class RelationalConnector(BaseConnector):
249
245
  def inspect_targeted(
250
246
  self,
251
247
  user_prompt: str,
252
- max_tables: int = 8,
253
- include_samples: bool = True,
248
+ max_tables: int = 4,
249
+ include_samples: bool = False,
254
250
  ) -> DatabaseSchema:
255
251
  """
256
- On-demand targeted inspection with fuzzy matching.
257
- Returns (schema, exact_matches, fuzzy_matches) for conflict resolution.
252
+ Ultra-fast on-demand targeted inspection with fuzzy matching (<30ms).
258
253
  """
259
254
  all_tables = self.get_table_names()
260
255
  exact, fuzzy = _fuzzy_match_tables(all_tables, user_prompt)
261
256
 
262
- # Combine exact + fuzzy, but prefer exact
263
- matched = exact[:]
264
- for f in fuzzy:
265
- if f not in matched:
266
- matched.append(f)
267
-
268
- if matched:
269
- # Also pull in FK-related tables for completeness
270
- fk_related = self._get_fk_related_tables(matched[:max_tables])
271
- all_to_inspect = matched[:max_tables]
272
- for r in fk_related:
273
- if r not in all_to_inspect and len(all_to_inspect) < max_tables:
274
- all_to_inspect.append(r)
275
-
276
- return self.inspect_schema(
277
- table_names=all_to_inspect,
278
- include_samples=include_samples,
279
- max_samples=2,
280
- include_views=False,
281
- include_row_counts=False,
282
- )
257
+ matched = exact[:] if exact else fuzzy[:max_tables]
258
+ if not matched:
259
+ matched = all_tables[:max_tables]
283
260
 
284
- # Fallback: inspect top max_tables
285
261
  return self.inspect_schema(
286
- table_names=all_tables[:max_tables],
287
- include_samples=include_samples,
288
- max_samples=2,
289
- include_views=True,
262
+ table_names=matched[:max_tables],
263
+ include_samples=False,
264
+ max_samples=0,
265
+ include_views=False,
290
266
  include_row_counts=False,
291
267
  )
292
268
 
@@ -0,0 +1,131 @@
1
+ """
2
+ Google Gemini Free Tier LLM Provider.
3
+ Uses Gemini 2.5 Flash / Gemini 2.0 Flash / Gemini 1.5 Flash via direct HTTP API.
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 GeminiProvider(BaseLLMProvider):
14
+ """Google Gemini LLM provider (Free Tier available)."""
15
+
16
+ FALLBACK_MODELS = [
17
+ "gemini-2.5-flash",
18
+ "gemini-2.0-flash",
19
+ "gemini-1.5-flash",
20
+ "gemini-1.5-pro",
21
+ ]
22
+
23
+ def __init__(self, api_key: Optional[str] = None, default_model: str = "gemini-2.5-flash"):
24
+ self.api_key = api_key or os.getenv("GEMINI_API_KEY")
25
+ self.default_model = default_model
26
+
27
+ @property
28
+ def name(self) -> str:
29
+ return "gemini"
30
+
31
+ def is_available(self) -> bool:
32
+ return bool(self.api_key)
33
+
34
+ def list_models(self) -> List[str]:
35
+ return self.FALLBACK_MODELS
36
+
37
+ def generate(
38
+ self,
39
+ prompt: str,
40
+ system_prompt: Optional[str] = None,
41
+ model: Optional[str] = None,
42
+ stream_callback: Optional[Callable[[str], None]] = None,
43
+ ) -> str:
44
+ if not self.api_key:
45
+ raise ValueError(
46
+ "Gemini API key is not configured. Set GEMINI_API_KEY environment variable or run `db-agent config`."
47
+ )
48
+
49
+ requested_model = model or self.default_model
50
+ if requested_model.startswith("models/"):
51
+ requested_model = requested_model[7:]
52
+
53
+ models_to_try = [requested_model]
54
+ for fb in self.FALLBACK_MODELS:
55
+ if fb != requested_model and fb not in models_to_try:
56
+ models_to_try.append(fb)
57
+
58
+ contents = []
59
+ if system_prompt:
60
+ contents.append({
61
+ "role": "user",
62
+ "parts": [{"text": f"System Instructions: {system_prompt}\n\nTask: {prompt}"}],
63
+ })
64
+ else:
65
+ contents.append({
66
+ "role": "user",
67
+ "parts": [{"text": prompt}],
68
+ })
69
+
70
+ payload = {
71
+ "contents": contents,
72
+ "generationConfig": {
73
+ "temperature": 0.1,
74
+ },
75
+ }
76
+
77
+ last_error = None
78
+ for target_model in models_to_try:
79
+ endpoint = f"https://generativelanguage.googleapis.com/v1beta/models/{target_model}:generateContent?key={self.api_key}"
80
+ if stream_callback:
81
+ endpoint = f"https://generativelanguage.googleapis.com/v1beta/models/{target_model}:streamGenerateContent?alt=sse&key={self.api_key}"
82
+
83
+ try:
84
+ if stream_callback:
85
+ response = requests.post(endpoint, json=payload, stream=True, timeout=15)
86
+ if response.status_code == 200:
87
+ full_text = []
88
+ for line in response.iter_lines():
89
+ if line:
90
+ decoded = line.decode("utf-8")
91
+ if decoded.startswith("data: "):
92
+ raw_json = decoded[6:]
93
+ try:
94
+ chunk = json.loads(raw_json)
95
+ candidates = chunk.get("candidates", [])
96
+ if candidates:
97
+ parts = candidates[0].get("content", {}).get("parts", [])
98
+ for p in parts:
99
+ t = p.get("text", "")
100
+ full_text.append(t)
101
+ stream_callback(t)
102
+ except Exception:
103
+ pass
104
+ return "".join(full_text)
105
+ elif response.status_code in (404, 400):
106
+ last_error = f"Model {target_model} not available ({response.status_code})"
107
+ continue
108
+ else:
109
+ raise RuntimeError(f"Gemini API error ({response.status_code}): {response.text}")
110
+ else:
111
+ response = requests.post(endpoint, json=payload, timeout=15)
112
+ if response.status_code == 200:
113
+ data = response.json()
114
+ candidates = data.get("candidates", [])
115
+ if candidates:
116
+ parts = candidates[0].get("content", {}).get("parts", [])
117
+ return "".join(p.get("text", "") for p in parts)
118
+ return ""
119
+ elif response.status_code in (404, 400):
120
+ last_error = f"Model {target_model} not available ({response.status_code})"
121
+ continue
122
+ else:
123
+ raise RuntimeError(f"Gemini API error ({response.status_code}): {response.text}")
124
+ except requests.exceptions.Timeout:
125
+ last_error = f"Model {target_model} request timed out"
126
+ continue
127
+ except Exception as e:
128
+ last_error = str(e)
129
+ continue
130
+
131
+ raise RuntimeError(f"Gemini generation error: {last_error or 'All model attempts failed'}")
@@ -56,22 +56,37 @@ def print_code(code: str, language: str = "sql", title: Optional[str] = None) ->
56
56
  console.print(panel)
57
57
 
58
58
 
59
- def print_results_table(columns: List[str], rows: List[Dict[str, Any]], title: str = "Query Results") -> None:
60
- """Print query results in a clean Rich table."""
61
- if not columns:
62
- console.print("[dim]Query returned no rows.[/dim]")
59
+ def print_results_table(columns: List[str], rows: List[Dict[str, Any]], title: str = "Results") -> None:
60
+ """Print query results in a clean, readable Rich table with wide-table protection."""
61
+ if not columns or not rows:
62
+ console.print("[dim]Query returned 0 rows.[/dim]\n")
63
63
  return
64
64
 
65
- table = Table(title=title, show_header=True, header_style="bold magenta", border_style="dim")
66
- for col in columns:
67
- table.add_column(str(col))
65
+ MAX_DISPLAY_COLS = 8
66
+ is_truncated = len(columns) > MAX_DISPLAY_COLS
67
+ display_cols = columns[:MAX_DISPLAY_COLS] if is_truncated else columns
68
+
69
+ table = Table(title=f"[bold cyan]{title}[/bold cyan]", show_header=True, header_style="bold magenta", border_style="dim")
70
+
71
+ for col in display_cols:
72
+ table.add_column(str(col), overflow="ellipsis", max_width=35)
68
73
 
69
74
  for row in rows:
70
75
  row_vals = []
71
- for col in columns:
76
+ for col in display_cols:
72
77
  v = row.get(col)
73
- row_vals.append(str(v) if v is not None else "[dim]NULL[/dim]")
78
+ val_str = str(v) if v is not None else "[dim]NULL[/dim]"
79
+ # Clean up excessively long raw json/strings
80
+ if len(val_str) > 40:
81
+ val_str = val_str[:37] + "..."
82
+ row_vals.append(val_str)
74
83
  table.add_row(*row_vals)
75
84
 
85
+ console.print("")
76
86
  console.print(table)
77
- console.print(f"[dim]Total rows: {len(rows)}[/dim]\n")
87
+ if is_truncated:
88
+ remaining = len(columns) - MAX_DISPLAY_COLS
89
+ console.print(f"[dim](Showing {MAX_DISPLAY_COLS} of {len(columns)} columns. Hidden: {', '.join(columns[MAX_DISPLAY_COLS:MAX_DISPLAY_COLS+5])}{'...' if remaining > 5 else ''})[/dim]")
90
+ console.print(f"[dim]Tip: Ask for specific fields (e.g., 'list id, username, email, status from ivis_user') for targeted view.[/dim]")
91
+ console.print(f"[bold green]Total rows:[/bold green] {len(rows)}\n")
92
+
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dbagent-cli
3
- Version: 0.3.3
4
- Summary: Enterprise AI Database Copilot with Instant Schema Caching & Zero-Latency Chat (CLI)
3
+ Version: 0.3.5
4
+ Summary: Production-Grade AI Database Copilot with Ultra-Fast Inference & Clean UI (CLI)
5
5
  Author: Santhosh Gupta
6
6
  Requires-Python: >=3.9
7
7
  Description-Content-Type: text/markdown
@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "dbagent-cli"
7
- version = "0.3.3"
8
- description = "Enterprise AI Database Copilot with Instant Schema Caching & Zero-Latency Chat (CLI)"
7
+ version = "0.3.5"
8
+ description = "Production-Grade AI Database Copilot with Ultra-Fast Inference & Clean UI (CLI)"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.9"
11
11
  authors = [
@@ -1,110 +0,0 @@
1
- """
2
- Google Gemini Free Tier LLM Provider.
3
- Uses Gemini 2.0 Flash / Gemini 1.5 Flash via direct HTTP API.
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 GeminiProvider(BaseLLMProvider):
14
- """Google Gemini LLM provider (Free Tier available)."""
15
-
16
- def __init__(self, api_key: Optional[str] = None, default_model: str = "gemini-3.6-flash"):
17
- self.api_key = api_key or os.getenv("GEMINI_API_KEY")
18
- self.default_model = default_model
19
-
20
- @property
21
- def name(self) -> str:
22
- return "gemini"
23
-
24
- def is_available(self) -> bool:
25
- return bool(self.api_key)
26
-
27
- def list_models(self) -> List[str]:
28
- return [
29
- "gemini-3.6-flash",
30
- "gemini-2.5-flash",
31
- "gemini-2.5-pro",
32
- ]
33
-
34
- def generate(
35
- self,
36
- prompt: str,
37
- system_prompt: Optional[str] = None,
38
- model: Optional[str] = None,
39
- stream_callback: Optional[Callable[[str], None]] = None,
40
- ) -> str:
41
- if not self.api_key:
42
- raise ValueError(
43
- "Gemini API key is not configured. Set GEMINI_API_KEY environment variable or run `db-agent config`."
44
- )
45
-
46
- target_model = model or self.default_model
47
- # Strip model prefixes if provided
48
- if target_model.startswith("models/"):
49
- target_model = target_model[7:]
50
-
51
- endpoint = f"https://generativelanguage.googleapis.com/v1beta/models/{target_model}:generateContent?key={self.api_key}"
52
- if stream_callback:
53
- endpoint = f"https://generativelanguage.googleapis.com/v1beta/models/{target_model}:streamGenerateContent?alt=sse&key={self.api_key}"
54
-
55
- contents = []
56
- if system_prompt:
57
- contents.append({
58
- "role": "user",
59
- "parts": [{"text": f"System Instructions: {system_prompt}\n\nTask: {prompt}"}],
60
- })
61
- else:
62
- contents.append({
63
- "role": "user",
64
- "parts": [{"text": prompt}],
65
- })
66
-
67
- payload = {
68
- "contents": contents,
69
- "generationConfig": {
70
- "temperature": 0.1,
71
- },
72
- }
73
-
74
- try:
75
- if stream_callback:
76
- response = requests.post(endpoint, json=payload, stream=True, timeout=90)
77
- if response.status_code != 200:
78
- raise RuntimeError(f"Gemini API error ({response.status_code}): {response.text}")
79
-
80
- full_text = []
81
- for line in response.iter_lines():
82
- if line:
83
- decoded = line.decode("utf-8")
84
- if decoded.startswith("data: "):
85
- raw_json = decoded[6:]
86
- try:
87
- chunk = json.loads(raw_json)
88
- candidates = chunk.get("candidates", [])
89
- if candidates:
90
- parts = candidates[0].get("content", {}).get("parts", [])
91
- for p in parts:
92
- t = p.get("text", "")
93
- full_text.append(t)
94
- stream_callback(t)
95
- except Exception:
96
- pass
97
- return "".join(full_text)
98
- else:
99
- response = requests.post(endpoint, json=payload, timeout=90)
100
- if response.status_code == 200:
101
- data = response.json()
102
- candidates = data.get("candidates", [])
103
- if candidates:
104
- parts = candidates[0].get("content", {}).get("parts", [])
105
- return "".join(p.get("text", "") for p in parts)
106
- return ""
107
- else:
108
- raise RuntimeError(f"Gemini API error ({response.status_code}): {response.text}")
109
- except Exception as e:
110
- raise RuntimeError(f"Gemini generation error: {str(e)}")
File without changes
File without changes
File without changes