sqlsaber 0.19.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

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.

Potentially problematic release.


This version of sqlsaber might be problematic. Click here for more details.

@@ -8,6 +8,7 @@ import httpx
8
8
  from pydantic_ai import Agent, RunContext
9
9
  from pydantic_ai.models.anthropic import AnthropicModel
10
10
  from pydantic_ai.models.google import GoogleModel
11
+ from pydantic_ai.models.openai import OpenAIResponsesModel
11
12
  from pydantic_ai.providers.anthropic import AnthropicProvider
12
13
  from pydantic_ai.providers.google import GoogleProvider
13
14
 
@@ -79,6 +80,10 @@ def build_sqlsaber_agent(
79
80
  provider_obj = AnthropicProvider(api_key="placeholder", http_client=http_client)
80
81
  model_obj = AnthropicModel(model_name_only, provider=provider_obj)
81
82
  agent = Agent(model_obj, name="sqlsaber")
83
+ elif provider == "openai":
84
+ # Use OpenAI Responses Model for structured output capabilities
85
+ model_obj = OpenAIResponsesModel(model_name_only)
86
+ agent = Agent(model_obj, name="sqlsaber")
82
87
  else:
83
88
  agent = Agent(cfg.model_name, name="sqlsaber")
84
89
 
sqlsaber/cli/commands.py CHANGED
@@ -7,6 +7,12 @@ from typing import Annotated
7
7
  import cyclopts
8
8
  from rich.console import Console
9
9
 
10
+ from sqlsaber.cli.auth import create_auth_app
11
+ from sqlsaber.cli.database import create_db_app
12
+ from sqlsaber.cli.memory import create_memory_app
13
+ from sqlsaber.cli.models import create_models_app
14
+ from sqlsaber.cli.threads import create_threads_app
15
+
10
16
  # Lazy imports - only import what's needed for CLI parsing
11
17
  from sqlsaber.config.database import DatabaseConfigManager
12
18
 
@@ -24,6 +30,11 @@ app = cyclopts.App(
24
30
  help="SQLsaber - Open-source agentic SQL assistant for your database",
25
31
  )
26
32
 
33
+ app.command(create_auth_app(), name="auth")
34
+ app.command(create_db_app(), name="db")
35
+ app.command(create_memory_app(), name="memory")
36
+ app.command(create_models_app(), name="models")
37
+ app.command(create_threads_app(), name="threads")
27
38
 
28
39
  console = Console()
29
40
  config_manager = DatabaseConfigManager()
@@ -195,47 +206,6 @@ def query(
195
206
  sys.exit(e.exit_code)
196
207
 
197
208
 
198
- # Use lazy imports for fast CLI startup time
199
- @app.command(name="auth")
200
- def auth(*args, **kwargs):
201
- """Manage authentication configuration."""
202
- from sqlsaber.cli.auth import create_auth_app
203
-
204
- return create_auth_app()(*args, **kwargs)
205
-
206
-
207
- @app.command(name="db")
208
- def db(*args, **kwargs):
209
- """Manage database connections."""
210
- from sqlsaber.cli.database import create_db_app
211
-
212
- return create_db_app()(*args, **kwargs)
213
-
214
-
215
- @app.command(name="memory")
216
- def memory(*args, **kwargs):
217
- """Manage database-specific memories."""
218
- from sqlsaber.cli.memory import create_memory_app
219
-
220
- return create_memory_app()(*args, **kwargs)
221
-
222
-
223
- @app.command(name="models")
224
- def models(*args, **kwargs):
225
- """Select and manage models."""
226
- from sqlsaber.cli.models import create_models_app
227
-
228
- return create_models_app()(*args, **kwargs)
229
-
230
-
231
- @app.command(name="threads")
232
- def threads(*args, **kwargs):
233
- """Manage SQLsaber threads."""
234
- from sqlsaber.cli.threads import create_threads_app
235
-
236
- return create_threads_app()(*args, **kwargs)
237
-
238
-
239
209
  def main():
240
210
  """Entry point for the CLI application."""
241
211
  app()
sqlsaber/cli/database.py CHANGED
@@ -12,7 +12,6 @@ from rich.console import Console
12
12
  from rich.table import Table
13
13
 
14
14
  from sqlsaber.config.database import DatabaseConfig, DatabaseConfigManager
15
- from sqlsaber.database.connection import DatabaseConnection
16
15
 
17
16
  # Global instances for CLI commands
18
17
  console = Console()
@@ -343,6 +342,9 @@ def test(
343
342
  """Test a database connection."""
344
343
 
345
344
  async def test_connection():
345
+ # Lazy import to keep CLI startup fast
346
+ from sqlsaber.database.connection import DatabaseConnection
347
+
346
348
  if name:
347
349
  db_config = config_manager.get_database(name)
348
350
  if not db_config:
sqlsaber/cli/display.py CHANGED
@@ -1,12 +1,167 @@
1
- """Display utilities for the CLI interface."""
1
+ """Display utilities for the CLI interface.
2
+
3
+ All rendering occurs on the event loop thread.
4
+ Streaming segments use Live Markdown; transient status and SQL blocks are also
5
+ rendered with Live.
6
+ """
2
7
 
3
8
  import json
9
+ from typing import Sequence, Type
4
10
 
5
- from rich.console import Console
6
- from rich.markdown import Markdown
11
+ from pydantic_ai.messages import ModelResponsePart, TextPart
12
+ from rich.columns import Columns
13
+ from rich.console import Console, ConsoleOptions, RenderResult
14
+ from rich.live import Live
15
+ from rich.markdown import CodeBlock, Markdown
7
16
  from rich.panel import Panel
17
+ from rich.spinner import Spinner
8
18
  from rich.syntax import Syntax
9
19
  from rich.table import Table
20
+ from rich.text import Text
21
+
22
+
23
+ class _SimpleCodeBlock(CodeBlock):
24
+ def __rich_console__(
25
+ self, console: Console, options: ConsoleOptions
26
+ ) -> RenderResult:
27
+ code = str(self.text).rstrip()
28
+ yield Syntax(
29
+ code,
30
+ self.lexer_name,
31
+ theme=self.theme,
32
+ background_color="default",
33
+ word_wrap=True,
34
+ )
35
+
36
+
37
+ class LiveMarkdownRenderer:
38
+ """Handles Live markdown rendering with segment separation.
39
+
40
+ Supports different segment kinds: 'assistant', 'thinking', 'sql'.
41
+ Adds visible paragraph breaks between segments and renders code fences
42
+ with nicer formatting.
43
+ """
44
+
45
+ _patched_fences = False
46
+
47
+ def __init__(self, console: Console):
48
+ self.console = console
49
+ self._live: Live | None = None
50
+ self._status_live: Live | None = None
51
+ self._buffer: str = ""
52
+ self._current_kind: Type[ModelResponsePart] | None = None
53
+
54
+ def prepare_code_blocks(self) -> None:
55
+ """Patch rich Markdown fence rendering once for nicer code blocks."""
56
+ if LiveMarkdownRenderer._patched_fences:
57
+ return
58
+ # Guard with class check to avoid re-patching if already applied
59
+ if Markdown.elements.get("fence") is not _SimpleCodeBlock:
60
+ Markdown.elements["fence"] = _SimpleCodeBlock
61
+ LiveMarkdownRenderer._patched_fences = True
62
+
63
+ def ensure_segment(self, kind: Type[ModelResponsePart]) -> None:
64
+ """
65
+ Ensure a markdown Live segment is active for the given kind.
66
+
67
+ When switching kinds, end the previous segment and add a paragraph break.
68
+ """
69
+ # If a transient status is showing, clear it first (no paragraph break)
70
+ if self._status_live is not None:
71
+ self.end_status()
72
+ if self._live is not None and self._current_kind == kind:
73
+ return
74
+ if self._live is not None:
75
+ self.end()
76
+ self.paragraph_break()
77
+
78
+ self._start()
79
+ self._current_kind = kind
80
+
81
+ def append(self, text: str | None) -> None:
82
+ """Append text to the current markdown segment and refresh."""
83
+ if not text:
84
+ return
85
+ if self._live is None:
86
+ # default to assistant if no segment was ensured
87
+ self.ensure_segment(TextPart)
88
+
89
+ self._buffer += text
90
+ self._live.update(Markdown(self._buffer))
91
+
92
+ def end(self) -> None:
93
+ """Finalize and stop the current Live segment, if any."""
94
+ if self._live is None:
95
+ return
96
+ if self._buffer:
97
+ self._live.update(Markdown(self._buffer))
98
+ self._live.stop()
99
+ self._live = None
100
+ self._buffer = ""
101
+ self._current_kind = None
102
+
103
+ def end_if_active(self) -> None:
104
+ self.end()
105
+
106
+ def paragraph_break(self) -> None:
107
+ self.console.print()
108
+
109
+ def start_sql_block(self, sql: str) -> None:
110
+ """Render a SQL block using a transient Live markdown segment."""
111
+ if not sql or not isinstance(sql, str) or not sql.strip():
112
+ return
113
+ # Separate from surrounding content
114
+ self.end_if_active()
115
+ self.paragraph_break()
116
+ self._buffer = f"```sql\n{sql}\n```"
117
+ # Use context manager to auto-stop and persist final render
118
+ with Live(
119
+ Markdown(self._buffer),
120
+ console=self.console,
121
+ vertical_overflow="visible",
122
+ refresh_per_second=12,
123
+ ):
124
+ pass
125
+
126
+ def start_status(self, message: str = "Crunching data...") -> None:
127
+ """Show a transient status line with a spinner until streaming starts."""
128
+ if self._status_live is not None:
129
+ # Update existing status text
130
+ self._status_live.update(self._status_renderable(message))
131
+ return
132
+ live = Live(
133
+ self._status_renderable(message),
134
+ console=self.console,
135
+ transient=True, # disappear when stopped
136
+ refresh_per_second=12,
137
+ )
138
+ self._status_live = live
139
+ live.start()
140
+
141
+ def end_status(self) -> None:
142
+ live = self._status_live
143
+ if live is None:
144
+ return
145
+ live.stop()
146
+ self._status_live = None
147
+
148
+ def _status_renderable(self, message: str):
149
+ spinner = Spinner("dots", style="yellow")
150
+ text = Text(f" {message}", style="yellow")
151
+ return Columns([spinner, text], expand=False)
152
+
153
+ def _start(self, initial_markdown: str = "") -> None:
154
+ if self._live is not None:
155
+ self.end()
156
+ self._buffer = initial_markdown or ""
157
+ live = Live(
158
+ Markdown(self._buffer),
159
+ console=self.console,
160
+ vertical_overflow="visible",
161
+ refresh_per_second=12,
162
+ )
163
+ self._live = live
164
+ live.start()
10
165
 
11
166
 
12
167
  class DisplayManager:
@@ -14,10 +169,11 @@ class DisplayManager:
14
169
 
15
170
  def __init__(self, console: Console):
16
171
  self.console = console
172
+ self.live = LiveMarkdownRenderer(console)
17
173
 
18
174
  def _create_table(
19
175
  self,
20
- columns: list,
176
+ columns: Sequence[str | dict[str, str]],
21
177
  header_style: str = "bold blue",
22
178
  title: str | None = None,
23
179
  ) -> Table:
@@ -34,17 +190,24 @@ class DisplayManager:
34
190
 
35
191
  def show_tool_executing(self, tool_name: str, tool_input: dict):
36
192
  """Display tool execution details."""
37
- self.console.print(f"\n[yellow]🔧 Using tool: {tool_name}[/yellow]")
193
+ # Normalized leading blank line before tool headers
194
+ self.show_newline()
38
195
  if tool_name == "list_tables":
39
- self.console.print("[dim] → Discovering available tables[/dim]")
196
+ self.console.print(
197
+ "[dim bold]:gear: Discovering available tables[/dim bold]"
198
+ )
40
199
  elif tool_name == "introspect_schema":
41
200
  pattern = tool_input.get("table_pattern", "all tables")
42
- self.console.print(f"[dim] → Examining schema for: {pattern}[/dim]")
201
+ self.console.print(
202
+ f"[dim bold]:gear: Examining schema for: {pattern}[/dim bold]"
203
+ )
43
204
  elif tool_name == "execute_sql":
205
+ # For streaming, we render SQL via LiveMarkdownRenderer; keep Syntax
206
+ # rendering for threads show/resume. Controlled by include_sql flag.
44
207
  query = tool_input.get("query", "")
45
- self.console.print("\n[bold green]Executing SQL:[/bold green]")
208
+ self.console.print("[dim bold]:gear: Executing SQL:[/dim bold]")
46
209
  self.show_newline()
47
- syntax = Syntax(query, "sql")
210
+ syntax = Syntax(query, "sql", background_color="default", word_wrap=True)
48
211
  self.console.print(syntax)
49
212
 
50
213
  def show_text_stream(self, text: str):
@@ -99,10 +262,12 @@ class DisplayManager:
99
262
  """Display a newline for spacing."""
100
263
  self.console.print()
101
264
 
102
- def show_table_list(self, tables_data: str):
265
+ def show_table_list(self, tables_data: str | dict):
103
266
  """Display the results from list_tables tool."""
104
267
  try:
105
- data = json.loads(tables_data)
268
+ data = (
269
+ json.loads(tables_data) if isinstance(tables_data, str) else tables_data
270
+ )
106
271
 
107
272
  # Handle error case
108
273
  if "error" in data:
@@ -143,10 +308,12 @@ class DisplayManager:
143
308
  except Exception as e:
144
309
  self.show_error(f"Error displaying table list: {str(e)}")
145
310
 
146
- def show_schema_info(self, schema_data: str):
311
+ def show_schema_info(self, schema_data: str | dict):
147
312
  """Display the results from introspect_schema tool."""
148
313
  try:
149
- data = json.loads(schema_data)
314
+ data = (
315
+ json.loads(schema_data) if isinstance(schema_data, str) else schema_data
316
+ )
150
317
 
151
318
  # Handle error case
152
319
  if "error" in data:
@@ -2,13 +2,16 @@
2
2
 
3
3
  import asyncio
4
4
  from pathlib import Path
5
+ from textwrap import dedent
5
6
 
6
7
  import platformdirs
7
8
  from prompt_toolkit import PromptSession
8
9
  from prompt_toolkit.history import FileHistory
10
+ from prompt_toolkit.patch_stdout import patch_stdout
9
11
  from prompt_toolkit.styles import Style
10
12
  from pydantic_ai import Agent
11
13
  from rich.console import Console
14
+ from rich.markdown import Markdown
12
15
  from rich.panel import Panel
13
16
 
14
17
  from sqlsaber.cli.completers import (
@@ -101,14 +104,14 @@ class InteractiveSession:
101
104
  )
102
105
  )
103
106
  self.console.print(
104
- "\n",
105
- "[dim] > Use '/clear' to reset conversation",
106
- "[dim] > Use 'Ctrl+D', '/exit' or '/quit' to leave[/dim]",
107
- "[dim] > Use 'Ctrl+C' to interrupt and return to prompt\n\n",
108
- "[dim] > Start message with '#' to add something to agent's memory for this database",
109
- "[dim] > Type '@' to get table name completions",
110
- "[dim] > Press 'Esc-Enter' or 'Meta-Enter' to submit your question",
111
- sep="\n",
107
+ Markdown(
108
+ dedent("""
109
+ - Use `/` for slash commands
110
+ - Type `@` to get table name completions
111
+ - Start message with `#` to add something to agent's memory
112
+ - Use `Ctrl+C` to interrupt and `Ctrl+D` to exit
113
+ """)
114
+ )
112
115
  )
113
116
 
114
117
  self.console.print(
@@ -118,6 +121,15 @@ class InteractiveSession:
118
121
  if self._thread_id:
119
122
  self.console.print(f"[dim]Resuming thread:[/dim] {self._thread_id}\n")
120
123
 
124
+ async def _end_thread_and_display_resume_hint(self):
125
+ """End thread and display command to resume thread"""
126
+ # Print resume hint if there is an active thread
127
+ if self._thread_id:
128
+ await self._threads.end_thread(self._thread_id)
129
+ self.console.print(
130
+ f"[dim]You can continue this thread using:[/dim] saber threads resume {self._thread_id}"
131
+ )
132
+
121
133
  async def _update_table_cache(self):
122
134
  """Update the table completer cache with fresh data."""
123
135
  try:
@@ -215,33 +227,27 @@ class InteractiveSession:
215
227
 
216
228
  while True:
217
229
  try:
218
- # with patch_stdout():
219
- user_query = await session.prompt_async(
220
- "",
221
- multiline=True,
222
- mouse_support=True,
223
- completer=CompositeCompleter(
224
- SlashCommandCompleter(), self.table_completer
225
- ),
226
- show_frame=True,
227
- bottom_toolbar=bottom_toolbar,
228
- style=style,
229
- )
230
+ with patch_stdout():
231
+ user_query = await session.prompt_async(
232
+ "",
233
+ multiline=True,
234
+ completer=CompositeCompleter(
235
+ SlashCommandCompleter(), self.table_completer
236
+ ),
237
+ show_frame=True,
238
+ bottom_toolbar=bottom_toolbar,
239
+ style=style,
240
+ )
230
241
 
231
242
  if not user_query:
232
243
  continue
233
244
 
234
245
  if (
235
- user_query in ["/exit", "/quit"]
246
+ user_query in ["/exit", "/quit", "exit", "quit"]
236
247
  or user_query.startswith("/exit")
237
248
  or user_query.startswith("/quit")
238
249
  ):
239
- # Print resume hint if there is an active thread
240
- if self._thread_id:
241
- await self._threads.end_thread(self._thread_id)
242
- self.console.print(
243
- f"[dim]You can continue this thread using:[/dim] saber threads resume {self._thread_id}"
244
- )
250
+ await self._end_thread_and_display_resume_hint()
245
251
  break
246
252
 
247
253
  if user_query == "/clear":
@@ -313,6 +319,7 @@ class InteractiveSession:
313
319
  )
314
320
  except EOFError:
315
321
  # Exit when Ctrl+D is pressed
322
+ await self._end_thread_and_display_resume_hint()
316
323
  break
317
324
  except Exception as e:
318
325
  self.console.print(f"[bold red]Error:[/bold red] {str(e)}")
sqlsaber/cli/streaming.py CHANGED
@@ -1,7 +1,13 @@
1
- """Streaming query handling for the CLI (pydantic-ai based)."""
1
+ """Streaming query handling for the CLI (pydantic-ai based).
2
+
3
+ This module uses DisplayManager's LiveMarkdownRenderer to stream Markdown
4
+ incrementally as the agent outputs tokens. Tool calls and results are
5
+ rendered via DisplayManager helpers.
6
+ """
2
7
 
3
8
  import asyncio
4
9
  import json
10
+ from functools import singledispatchmethod
5
11
  from typing import AsyncIterable
6
12
 
7
13
  from pydantic_ai import Agent, RunContext
@@ -22,56 +28,98 @@ from sqlsaber.cli.display import DisplayManager
22
28
 
23
29
 
24
30
  class StreamingQueryHandler:
25
- """Handles streaming query execution and display using pydantic-ai events."""
31
+ """
32
+ Handles streaming query execution and display using pydantic-ai events.
33
+
34
+ Uses DisplayManager.live to render Markdown incrementally as text streams in.
35
+ """
26
36
 
27
37
  def __init__(self, console: Console):
28
38
  self.console = console
29
39
  self.display = DisplayManager(console)
30
40
 
31
- self.status = self.console.status(
32
- "[yellow]Crunching data...[/yellow]", spinner="bouncingBall"
33
- )
34
-
35
41
  async def _event_stream_handler(
36
42
  self, ctx: RunContext, event_stream: AsyncIterable[AgentStreamEvent]
37
43
  ) -> None:
44
+ """
45
+ Handle pydantic-ai streaming events and update Live Markdown via DisplayManager.
46
+ """
47
+
38
48
  async for event in event_stream:
39
- if isinstance(event, PartStartEvent):
40
- if isinstance(event.part, (TextPart, ThinkingPart)):
41
- self.status.stop()
42
- self.display.show_text_stream(event.part.content)
43
-
44
- elif isinstance(event, PartDeltaEvent):
45
- if isinstance(event.delta, (TextPartDelta, ThinkingPartDelta)):
46
- delta = event.delta.content_delta or ""
47
- if delta:
48
- self.status.stop()
49
- self.display.show_text_stream(delta)
50
-
51
- elif isinstance(event, FunctionToolCallEvent):
52
- # Show tool execution start
53
- self.status.stop()
54
- args = event.part.args_as_dict()
55
- self.display.show_newline()
56
- self.display.show_tool_executing(event.part.tool_name, args)
57
-
58
- elif isinstance(event, FunctionToolResultEvent):
59
- self.status.stop()
60
- # Route tool result to appropriate display
61
- tool_name = event.result.tool_name
62
- content = event.result.content
63
- if tool_name == "list_tables":
64
- self.display.show_table_list(content)
65
- elif tool_name == "introspect_schema":
66
- self.display.show_schema_info(content)
67
- elif tool_name == "execute_sql":
49
+ await self.on_event(event, ctx)
50
+
51
+ # --- Event routing via singledispatchmethod ---------------------------------------
52
+ @singledispatchmethod
53
+ async def on_event(
54
+ self, event: AgentStreamEvent, ctx: RunContext
55
+ ) -> None: # default
56
+ return
57
+
58
+ @on_event.register
59
+ async def _(self, event: PartStartEvent, ctx: RunContext) -> None:
60
+ if isinstance(event.part, TextPart):
61
+ self.display.live.ensure_segment(TextPart)
62
+ self.display.live.append(event.part.content)
63
+ elif isinstance(event.part, ThinkingPart):
64
+ self.display.live.ensure_segment(ThinkingPart)
65
+ self.display.live.append(event.part.content)
66
+
67
+ @on_event.register
68
+ async def _(self, event: PartDeltaEvent, ctx: RunContext) -> None:
69
+ d = event.delta
70
+ if isinstance(d, TextPartDelta):
71
+ delta = d.content_delta or ""
72
+ if delta:
73
+ self.display.live.ensure_segment(TextPart)
74
+ self.display.live.append(delta)
75
+ elif isinstance(d, ThinkingPartDelta):
76
+ delta = d.content_delta or ""
77
+ if delta:
78
+ self.display.live.ensure_segment(ThinkingPart)
79
+ self.display.live.append(delta)
80
+
81
+ @on_event.register
82
+ async def _(self, event: FunctionToolCallEvent, ctx: RunContext) -> None:
83
+ # Clear any status/markdown Live so tool output sits between
84
+ self.display.live.end_status()
85
+ self.display.live.end_if_active()
86
+ args = event.part.args_as_dict()
87
+
88
+ # Special handling: display SQL via Live as markdown code block
89
+ if event.part.tool_name == "execute_sql":
90
+ query = args.get("query") or ""
91
+ if isinstance(query, str) and query.strip():
92
+ self.display.live.start_sql_block(query)
93
+ else:
94
+ self.display.show_tool_executing(event.part.tool_name, args)
95
+
96
+ @on_event.register
97
+ async def _(self, event: FunctionToolResultEvent, ctx: RunContext) -> None:
98
+ # Route tool result to appropriate display
99
+ tool_name = event.result.tool_name
100
+ content = event.result.content
101
+ if tool_name == "list_tables":
102
+ self.display.show_table_list(content)
103
+ elif tool_name == "introspect_schema":
104
+ self.display.show_schema_info(content)
105
+ elif tool_name == "execute_sql":
106
+ data = {}
107
+ if isinstance(content, str):
108
+ try:
109
+ data = json.loads(content)
110
+ except (json.JSONDecodeError, TypeError) as exc:
68
111
  try:
69
- data = json.loads(content)
70
- if data.get("success") and data.get("results"):
71
- self.display.show_query_results(data["results"]) # type: ignore[arg-type]
72
- except json.JSONDecodeError:
73
- # If not JSON, ignore here
112
+ self.console.log(f"Malformed execute_sql result: {exc}")
113
+ except Exception:
74
114
  pass
115
+ elif isinstance(content, dict):
116
+ data = content
117
+ if isinstance(data, dict) and data.get("success") and data.get("results"):
118
+ self.display.show_query_results(data["results"]) # type: ignore[arg-type]
119
+ # Add a blank line after tool output to separate from next segment
120
+ self.display.show_newline()
121
+ # Show status while agent sends a follow-up request to the model
122
+ self.display.live.start_status("Crunching data...")
75
123
 
76
124
  async def execute_streaming_query(
77
125
  self,
@@ -80,7 +128,8 @@ class StreamingQueryHandler:
80
128
  cancellation_token: asyncio.Event | None = None,
81
129
  message_history: list | None = None,
82
130
  ):
83
- self.status.start()
131
+ # Prepare nicer code block rendering for Markdown
132
+ self.display.live.prepare_code_blocks()
84
133
  try:
85
134
  # If Anthropic OAuth, inject SQLsaber instructions before the first user prompt
86
135
  prepared_prompt: str | list[str] = user_query
@@ -104,30 +153,24 @@ class StreamingQueryHandler:
104
153
  injected = "\n\n".join(parts)
105
154
  prepared_prompt = [injected, user_query]
106
155
 
156
+ # Show a transient status until events start streaming
157
+ self.display.live.start_status("Crunching data...")
158
+
107
159
  # Run the agent with our event stream handler
108
160
  run = await agent.run(
109
161
  prepared_prompt,
110
162
  message_history=message_history,
111
163
  event_stream_handler=self._event_stream_handler,
112
164
  )
113
- # After the run completes, show the assistant's final text as markdown if available
114
- try:
115
- output = run.output
116
- if isinstance(output, str) and output.strip():
117
- self.display.show_newline()
118
- self.display.show_markdown_response(
119
- [{"type": "text", "text": output}]
120
- )
121
- except Exception as e:
122
- self.display.show_error(str(e))
123
- self.display.show_newline()
124
165
  return run
125
166
  except asyncio.CancelledError:
167
+ # Show interruption message outside of Live
126
168
  self.display.show_newline()
127
169
  self.console.print("[yellow]Query interrupted[/yellow]")
128
170
  return None
129
171
  finally:
172
+ # End any active status and live markdown segments
130
173
  try:
131
- self.status.stop()
132
- except Exception:
133
- pass
174
+ self.display.live.end_status()
175
+ finally:
176
+ self.display.live.end_if_active()
sqlsaber/cli/threads.py CHANGED
@@ -12,17 +12,10 @@ from rich.markdown import Markdown
12
12
  from rich.panel import Panel
13
13
  from rich.table import Table
14
14
 
15
- from sqlsaber.agents import build_sqlsaber_agent
16
- from sqlsaber.cli.display import DisplayManager
17
- from sqlsaber.cli.interactive import InteractiveSession
18
- from sqlsaber.config.database import DatabaseConfigManager
19
- from sqlsaber.database.connection import DatabaseConnection
20
- from sqlsaber.database.resolver import DatabaseResolutionError, resolve_database
21
15
  from sqlsaber.threads import ThreadStorage
22
16
 
23
17
  # Globals consistent with other CLI modules
24
18
  console = Console()
25
- config_manager = DatabaseConfigManager()
26
19
 
27
20
 
28
21
  threads_app = cyclopts.App(
@@ -41,6 +34,9 @@ def _render_transcript(
41
34
  console: Console, all_msgs: list[ModelMessage], last_n: int | None = None
42
35
  ) -> None:
43
36
  """Render conversation turns from ModelMessage[] using DisplayManager."""
37
+ # Lazy import to avoid pulling UI helpers at startup
38
+ from sqlsaber.cli.display import DisplayManager
39
+
44
40
  dm = DisplayManager(console)
45
41
 
46
42
  # Locate indices of user prompts
@@ -237,6 +233,16 @@ def resume(
237
233
  store = ThreadStorage()
238
234
 
239
235
  async def _run() -> None:
236
+ # Lazy imports to avoid heavy modules at CLI startup
237
+ from sqlsaber.agents import build_sqlsaber_agent
238
+ from sqlsaber.cli.interactive import InteractiveSession
239
+ from sqlsaber.config.database import DatabaseConfigManager
240
+ from sqlsaber.database.connection import DatabaseConnection
241
+ from sqlsaber.database.resolver import (
242
+ DatabaseResolutionError,
243
+ resolve_database,
244
+ )
245
+
240
246
  thread = await store.get_thread(thread_id)
241
247
  if not thread:
242
248
  console.print(f"[red]Thread not found:[/red] {thread_id}")
@@ -248,6 +254,7 @@ def resume(
248
254
  )
249
255
  return
250
256
  try:
257
+ config_manager = DatabaseConfigManager()
251
258
  resolved = resolve_database(db_selector, config_manager)
252
259
  connection_string = resolved.connection_string
253
260
  db_name = resolved.name
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: sqlsaber
3
- Version: 0.19.0
3
+ Version: 0.21.0
4
4
  Summary: SQLsaber - Open-source agentic SQL assistant
5
5
  License-File: LICENSE
6
6
  Requires-Python: >=3.12
@@ -3,18 +3,18 @@ sqlsaber/__main__.py,sha256=RIHxWeWh2QvLfah-2OkhI5IJxojWfy4fXpMnVEJYvxw,78
3
3
  sqlsaber/agents/__init__.py,sha256=i_MI2eWMQaVzGikKU71FPCmSQxNDKq36Imq1PrYoIPU,130
4
4
  sqlsaber/agents/base.py,sha256=7zOZTHKxUuU0uMc-NTaCkkBfDnU3jtwbT8_eP1ZtJ2k,2615
5
5
  sqlsaber/agents/mcp.py,sha256=GcJTx7YDYH6aaxIADEIxSgcWAdWakUx395JIzVnf17U,768
6
- sqlsaber/agents/pydantic_ai_agent.py,sha256=dGdsgyxCZvfK-v-MH8KimKOr-xb2aSfSWY8CMcOUCT8,6795
6
+ sqlsaber/agents/pydantic_ai_agent.py,sha256=6RvG2O7G8P6NN9QaRXUodg5Q26QJ4ShGWoTGYbVQ5K4,7065
7
7
  sqlsaber/cli/__init__.py,sha256=qVSLVJLLJYzoC6aj6y9MFrzZvAwc4_OgxU9DlkQnZ4M,86
8
8
  sqlsaber/cli/auth.py,sha256=jTsRgbmlGPlASSuIKmdjjwfqtKvjfKd_cTYxX0-QqaQ,7400
9
- sqlsaber/cli/commands.py,sha256=CmCqDC6KiE8JD6Vkpsry4lBQInCiS8TBeKKx3gdxZcM,8689
9
+ sqlsaber/cli/commands.py,sha256=mjLG9i1bXf0TEroxkIxq5O7Hhjufz3Ad72cyJz7vE1k,8128
10
10
  sqlsaber/cli/completers.py,sha256=HsUPjaZweLSeYCWkAcgMl8FylQ1xjWBWYTEL_9F6xfU,6430
11
- sqlsaber/cli/database.py,sha256=atwg3l8acQ3YTDuhq7vNrBN6tpOv0syz6V62KTF-Bh8,12910
12
- sqlsaber/cli/display.py,sha256=wa7BjTBwXwqLT145Q1AEL0C28pQJTrvDN10mnFMjqsg,8554
13
- sqlsaber/cli/interactive.py,sha256=uoJdLWPoaDBkfdFN-59u6zduU8XGY98L16zpNsGu7nE,12964
11
+ sqlsaber/cli/database.py,sha256=JKtHSN-BFzBa14REf0phFVQB7d67m1M5FFaD8N6DdrY,12966
12
+ sqlsaber/cli/display.py,sha256=bul9Yzw8KFYkof-kDzeajpx2TtG9CjTaUiwWaTv95dQ,14293
13
+ sqlsaber/cli/interactive.py,sha256=7uM4LoXbhPJr8o5yNjICSzL0uxZkp1psWrVq4G9V0OI,13118
14
14
  sqlsaber/cli/memory.py,sha256=OufHFJFwV0_GGn7LvKRTJikkWhV1IwNIUDOxFPHXOaQ,7794
15
15
  sqlsaber/cli/models.py,sha256=ZewtwGQwhd9b-yxBAPKePolvI1qQG-EkmeWAGMqtWNQ,8986
16
- sqlsaber/cli/streaming.py,sha256=WNqBYYbWtL5CNQkRg5YWhYpWKI8qz7JmqneB2DXTOHY,5259
17
- sqlsaber/cli/threads.py,sha256=xti7_kvh3loQfLb7_GC8wSULJ4Oj56jXY8GQp69CQCI,11111
16
+ sqlsaber/cli/streaming.py,sha256=BeG7H38-I1n8b9R8XSBV-IqkxDRZhsWFW6sdvtbVi3o,6879
17
+ sqlsaber/cli/threads.py,sha256=XUnLcCUe2wa_85IKdKmryqfiHTQu_IylET2Qo8oy1nk,11324
18
18
  sqlsaber/config/__init__.py,sha256=olwC45k8Nc61yK0WmPUk7XHdbsZH9HuUAbwnmKe3IgA,100
19
19
  sqlsaber/config/api_keys.py,sha256=RqWQCko1tY7sES7YOlexgBH5Hd5ne_kGXHdBDNqcV2U,3649
20
20
  sqlsaber/config/auth.py,sha256=b5qB2h1doXyO9Bn8z0CcL8LAR2jF431gGXBGKLgTmtQ,2756
@@ -40,8 +40,8 @@ sqlsaber/tools/enums.py,sha256=CH32mL-0k9ZA18911xLpNtsgpV6tB85TktMj6uqGz54,411
40
40
  sqlsaber/tools/instructions.py,sha256=X-x8maVkkyi16b6Tl0hcAFgjiYceZaSwyWTfmrvx8U8,9024
41
41
  sqlsaber/tools/registry.py,sha256=HWOQMsNIdL4XZS6TeNUyrL-5KoSDH6PHsWd3X66o-18,3211
42
42
  sqlsaber/tools/sql_tools.py,sha256=hM6tKqW5MDhFUt6MesoqhTUqIpq_5baIIDoN1MjDCXY,9647
43
- sqlsaber-0.19.0.dist-info/METADATA,sha256=oH4qQnKuu5n7ucEEtSccDyvI1obYPPeZV3xNAn-8Mzc,6178
44
- sqlsaber-0.19.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
45
- sqlsaber-0.19.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
46
- sqlsaber-0.19.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
47
- sqlsaber-0.19.0.dist-info/RECORD,,
43
+ sqlsaber-0.21.0.dist-info/METADATA,sha256=7NgOXIfUrri2-SHemXf6TVRLBqj2KuzweGafg-NPZuQ,6178
44
+ sqlsaber-0.21.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
45
+ sqlsaber-0.21.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
46
+ sqlsaber-0.21.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
47
+ sqlsaber-0.21.0.dist-info/RECORD,,