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.
- sqlsaber/agents/pydantic_ai_agent.py +5 -0
- sqlsaber/cli/commands.py +11 -41
- sqlsaber/cli/database.py +3 -1
- sqlsaber/cli/display.py +180 -13
- sqlsaber/cli/interactive.py +34 -27
- sqlsaber/cli/streaming.py +98 -55
- sqlsaber/cli/threads.py +14 -7
- {sqlsaber-0.19.0.dist-info → sqlsaber-0.21.0.dist-info}/METADATA +1 -1
- {sqlsaber-0.19.0.dist-info → sqlsaber-0.21.0.dist-info}/RECORD +12 -12
- {sqlsaber-0.19.0.dist-info → sqlsaber-0.21.0.dist-info}/WHEEL +0 -0
- {sqlsaber-0.19.0.dist-info → sqlsaber-0.21.0.dist-info}/entry_points.txt +0 -0
- {sqlsaber-0.19.0.dist-info → sqlsaber-0.21.0.dist-info}/licenses/LICENSE +0 -0
|
@@ -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
|
|
6
|
-
from rich.
|
|
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:
|
|
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
|
-
|
|
193
|
+
# Normalized leading blank line before tool headers
|
|
194
|
+
self.show_newline()
|
|
38
195
|
if tool_name == "list_tables":
|
|
39
|
-
self.console.print(
|
|
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(
|
|
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("
|
|
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 =
|
|
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 =
|
|
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:
|
sqlsaber/cli/interactive.py
CHANGED
|
@@ -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
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
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
|
-
|
|
219
|
-
|
|
220
|
-
|
|
221
|
-
|
|
222
|
-
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
|
|
228
|
-
|
|
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
|
-
|
|
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
|
-
"""
|
|
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
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
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
|
-
|
|
70
|
-
|
|
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
|
-
|
|
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.
|
|
132
|
-
|
|
133
|
-
|
|
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
|
|
@@ -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=
|
|
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=
|
|
9
|
+
sqlsaber/cli/commands.py,sha256=mjLG9i1bXf0TEroxkIxq5O7Hhjufz3Ad72cyJz7vE1k,8128
|
|
10
10
|
sqlsaber/cli/completers.py,sha256=HsUPjaZweLSeYCWkAcgMl8FylQ1xjWBWYTEL_9F6xfU,6430
|
|
11
|
-
sqlsaber/cli/database.py,sha256=
|
|
12
|
-
sqlsaber/cli/display.py,sha256=
|
|
13
|
-
sqlsaber/cli/interactive.py,sha256=
|
|
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=
|
|
17
|
-
sqlsaber/cli/threads.py,sha256=
|
|
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.
|
|
44
|
-
sqlsaber-0.
|
|
45
|
-
sqlsaber-0.
|
|
46
|
-
sqlsaber-0.
|
|
47
|
-
sqlsaber-0.
|
|
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,,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|