sqlsaber 0.5.0__tar.gz → 0.6.0__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.

Potentially problematic release.


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

Files changed (50) hide show
  1. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/CHANGELOG.md +24 -0
  2. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/PKG-INFO +1 -1
  3. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/pyproject.toml +1 -1
  4. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/agents/anthropic.py +55 -17
  5. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/agents/base.py +13 -3
  6. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/cli/display.py +11 -3
  7. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/cli/interactive.py +80 -7
  8. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/cli/streaming.py +21 -2
  9. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/uv.lock +1 -1
  10. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/.github/workflows/publish.yml +0 -0
  11. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/.gitignore +0 -0
  12. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/.python-version +0 -0
  13. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/CLAUDE.md +0 -0
  14. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/LICENSE +0 -0
  15. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/README.md +0 -0
  16. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/pytest.ini +0 -0
  17. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/__init__.py +0 -0
  18. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/__main__.py +0 -0
  19. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/agents/__init__.py +0 -0
  20. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/agents/mcp.py +0 -0
  21. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/agents/streaming.py +0 -0
  22. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/cli/__init__.py +0 -0
  23. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/cli/commands.py +0 -0
  24. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/cli/database.py +0 -0
  25. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/cli/memory.py +0 -0
  26. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/cli/models.py +0 -0
  27. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/config/__init__.py +0 -0
  28. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/config/api_keys.py +0 -0
  29. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/config/database.py +0 -0
  30. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/config/settings.py +0 -0
  31. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/database/__init__.py +0 -0
  32. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/database/connection.py +0 -0
  33. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/database/schema.py +0 -0
  34. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/mcp/__init__.py +0 -0
  35. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/mcp/mcp.py +0 -0
  36. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/memory/__init__.py +0 -0
  37. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/memory/manager.py +0 -0
  38. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/memory/storage.py +0 -0
  39. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/models/__init__.py +0 -0
  40. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/models/events.py +0 -0
  41. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/src/sqlsaber/models/types.py +0 -0
  42. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/tests/__init__.py +0 -0
  43. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/tests/conftest.py +0 -0
  44. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/tests/test_cli/__init__.py +0 -0
  45. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/tests/test_cli/test_commands.py +0 -0
  46. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/tests/test_config/__init__.py +0 -0
  47. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/tests/test_config/test_database.py +0 -0
  48. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/tests/test_config/test_settings.py +0 -0
  49. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/tests/test_database/__init__.py +0 -0
  50. {sqlsaber-0.5.0 → sqlsaber-0.6.0}/tests/test_database/test_connection.py +0 -0
@@ -4,6 +4,30 @@ All notable changes to SQLSaber will be documented in this file.
4
4
 
5
5
  ## [Unreleased]
6
6
 
7
+ ## [0.6.0] - 2025-06-30
8
+
9
+ ### Added
10
+
11
+ - Slash command autocomplete in interactive mode
12
+ - Commands now use slash prefix: `/clear`, `/exit`, `/quit`
13
+ - Autocomplete shows when typing `/` at the start of a line
14
+ - Press Tab to select suggestion
15
+ - Query interruption with Ctrl+C in interactive mode
16
+ - Press Ctrl+C during query execution to gracefully cancel ongoing operations
17
+ - Preserves conversation history up to the interruption point
18
+
19
+ ### Changed
20
+
21
+ - Updated table display for better readability: limit to first 15 columns on wide tables
22
+ - Shows warning when columns are truncated
23
+ - Interactive commands now require slash prefix (breaking change)
24
+ - `clear` → `/clear`
25
+ - `exit` → `/exit`
26
+ - `quit` → `/quit`
27
+ - Removed default limit of 100. Now model will decide it.
28
+
29
+ ## [0.5.0] - 2025-06-27
30
+
7
31
  ### Added
8
32
 
9
33
  - Added support for plotting data from query results.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: sqlsaber
3
- Version: 0.5.0
3
+ Version: 0.6.0
4
4
  Summary: SQLSaber - Agentic SQL assistant like Claude Code
5
5
  License-File: LICENSE
6
6
  Requires-Python: >=3.12
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "sqlsaber"
3
- version = "0.5.0"
3
+ version = "0.6.0"
4
4
  description = "SQLSaber - Agentic SQL assistant like Claude Code"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.12"
@@ -1,7 +1,8 @@
1
1
  """Anthropic-specific SQL agent implementation."""
2
2
 
3
+ import asyncio
3
4
  import json
4
- from typing import Any, AsyncIterator, Dict, List, Optional
5
+ from typing import Any, AsyncIterator, Dict, List
5
6
 
6
7
  from anthropic import AsyncAnthropic
7
8
 
@@ -21,7 +22,7 @@ class AnthropicSQLAgent(BaseSQLAgent):
21
22
  """SQL Agent using Anthropic SDK directly."""
22
23
 
23
24
  def __init__(
24
- self, db_connection: BaseDatabaseConnection, database_name: Optional[str] = None
25
+ self, db_connection: BaseDatabaseConnection, database_name: str | None = None
25
26
  ):
26
27
  super().__init__(db_connection)
27
28
 
@@ -164,7 +165,7 @@ Guidelines:
164
165
 
165
166
  return base_prompt
166
167
 
167
- def add_memory(self, content: str) -> Optional[str]:
168
+ def add_memory(self, content: str) -> str | None:
168
169
  """Add a memory for the current database."""
169
170
  if not self.database_name:
170
171
  return None
@@ -174,7 +175,7 @@ Guidelines:
174
175
  self.system_prompt = self._build_system_prompt()
175
176
  return memory.id
176
177
 
177
- async def execute_sql(self, query: str, limit: Optional[int] = 100) -> str:
178
+ async def execute_sql(self, query: str, limit: int | None = None) -> str:
178
179
  """Execute a SQL query against the database with streaming support."""
179
180
  # Call parent implementation for core functionality
180
181
  result = await super().execute_sql(query, limit)
@@ -203,10 +204,18 @@ Guidelines:
203
204
  return await super().process_tool_call(tool_name, tool_input)
204
205
 
205
206
  async def _process_stream_events(
206
- self, stream, content_blocks: List[Dict], tool_use_blocks: List[Dict]
207
+ self,
208
+ stream,
209
+ content_blocks: List[Dict],
210
+ tool_use_blocks: List[Dict],
211
+ cancellation_token: asyncio.Event | None = None,
207
212
  ) -> AsyncIterator[StreamEvent]:
208
213
  """Process stream events and yield appropriate StreamEvents."""
209
214
  async for event in stream:
215
+ # Only check cancellation if token is provided
216
+ if cancellation_token is not None and cancellation_token.is_set():
217
+ return
218
+
210
219
  if event.type == "content_block_start":
211
220
  if hasattr(event.content_block, "type"):
212
221
  if event.content_block.type == "tool_use":
@@ -253,11 +262,17 @@ Guidelines:
253
262
  return "stop"
254
263
 
255
264
  async def _process_tool_results(
256
- self, response: StreamingResponse
265
+ self,
266
+ response: StreamingResponse,
267
+ cancellation_token: asyncio.Event | None = None,
257
268
  ) -> AsyncIterator[StreamEvent]:
258
269
  """Process tool results and yield appropriate events."""
259
270
  tool_results = []
260
271
  for block in response.content:
272
+ # Only check cancellation if token is provided
273
+ if cancellation_token is not None and cancellation_token.is_set():
274
+ return
275
+
261
276
  if block.get("type") == "tool_use":
262
277
  yield StreamEvent(
263
278
  "tool_use",
@@ -304,7 +319,10 @@ Guidelines:
304
319
  yield StreamEvent("tool_result_data", tool_results)
305
320
 
306
321
  async def query_stream(
307
- self, user_query: str, use_history: bool = True
322
+ self,
323
+ user_query: str,
324
+ use_history: bool = True,
325
+ cancellation_token: asyncio.Event | None = None,
308
326
  ) -> AsyncIterator[StreamEvent]:
309
327
  """Process a user query and stream responses."""
310
328
  # Initialize for tracking state
@@ -322,7 +340,11 @@ Guidelines:
322
340
  try:
323
341
  # Create initial stream and get response
324
342
  response = None
325
- async for event in self._create_and_process_stream(messages):
343
+ async for event in self._create_and_process_stream(
344
+ messages, cancellation_token
345
+ ):
346
+ if cancellation_token is not None and cancellation_token.is_set():
347
+ return
326
348
  if event.type == "response_ready":
327
349
  response = event.data
328
350
  else:
@@ -332,14 +354,21 @@ Guidelines:
332
354
 
333
355
  # Process tool calls if needed
334
356
  while response is not None and response.stop_reason == "tool_use":
357
+ # Check for cancellation at the start of tool cycle
358
+ if cancellation_token is not None and cancellation_token.is_set():
359
+ return
360
+
335
361
  # Add assistant's response to conversation
336
362
  collected_content.append(
337
363
  {"role": "assistant", "content": response.content}
338
364
  )
339
365
 
340
- # Process tool results
366
+ # Process tool results - DO NOT check cancellation during tool execution
367
+ # as this would break the tool_use -> tool_result API contract
341
368
  tool_results = []
342
- async for event in self._process_tool_results(response):
369
+ async for event in self._process_tool_results(
370
+ response, None
371
+ ): # Pass None to disable cancellation checks
343
372
  if event.type == "tool_result_data":
344
373
  tool_results = event.data
345
374
  else:
@@ -347,6 +376,12 @@ Guidelines:
347
376
 
348
377
  # Continue conversation with tool results
349
378
  collected_content.append({"role": "user", "content": tool_results})
379
+ if use_history:
380
+ self.conversation_history.extend(collected_content)
381
+
382
+ # Check for cancellation AFTER tool results are complete
383
+ if cancellation_token is not None and cancellation_token.is_set():
384
+ return
350
385
 
351
386
  # Signal that we're processing the tool results
352
387
  yield StreamEvent("processing", "Analyzing results...")
@@ -354,8 +389,10 @@ Guidelines:
354
389
  # Get next response
355
390
  response = None
356
391
  async for event in self._create_and_process_stream(
357
- messages + collected_content
392
+ messages + collected_content, cancellation_token
358
393
  ):
394
+ if cancellation_token is not None and cancellation_token.is_set():
395
+ return
359
396
  if event.type == "response_ready":
360
397
  response = event.data
361
398
  else:
@@ -363,21 +400,19 @@ Guidelines:
363
400
 
364
401
  # Update conversation history if using history
365
402
  if use_history:
366
- self.conversation_history.append(
367
- {"role": "user", "content": user_query}
368
- )
369
- self.conversation_history.extend(collected_content)
370
403
  # Add final assistant response
371
404
  if response is not None:
372
405
  self.conversation_history.append(
373
406
  {"role": "assistant", "content": response.content}
374
407
  )
375
408
 
409
+ except asyncio.CancelledError:
410
+ return
376
411
  except Exception as e:
377
412
  yield StreamEvent("error", str(e))
378
413
 
379
414
  async def _create_and_process_stream(
380
- self, messages: List[Dict]
415
+ self, messages: List[Dict], cancellation_token: asyncio.Event | None = None
381
416
  ) -> AsyncIterator[StreamEvent]:
382
417
  """Create a stream and yield events while building response."""
383
418
  stream = await self.client.messages.create(
@@ -393,8 +428,11 @@ Guidelines:
393
428
  tool_use_blocks = []
394
429
 
395
430
  async for event in self._process_stream_events(
396
- stream, content_blocks, tool_use_blocks
431
+ stream, content_blocks, tool_use_blocks, cancellation_token
397
432
  ):
433
+ # Only check cancellation if token is provided
434
+ if cancellation_token is not None and cancellation_token.is_set():
435
+ return
398
436
  yield event
399
437
 
400
438
  # Finalize tool blocks and create response
@@ -1,5 +1,6 @@
1
1
  """Abstract base class for SQL agents."""
2
2
 
3
+ import asyncio
3
4
  import json
4
5
  from abc import ABC, abstractmethod
5
6
  from typing import Any, AsyncIterator, Dict, List, Optional
@@ -27,9 +28,18 @@ class BaseSQLAgent(ABC):
27
28
 
28
29
  @abstractmethod
29
30
  async def query_stream(
30
- self, user_query: str, use_history: bool = True
31
+ self,
32
+ user_query: str,
33
+ use_history: bool = True,
34
+ cancellation_token: asyncio.Event | None = None,
31
35
  ) -> AsyncIterator[StreamEvent]:
32
- """Process a user query and stream responses."""
36
+ """Process a user query and stream responses.
37
+
38
+ Args:
39
+ user_query: The user's query to process
40
+ use_history: Whether to include conversation history
41
+ cancellation_token: Optional event to signal cancellation
42
+ """
33
43
  pass
34
44
 
35
45
  def clear_history(self):
@@ -86,7 +96,7 @@ class BaseSQLAgent(ABC):
86
96
  except Exception as e:
87
97
  return json.dumps({"error": f"Error listing tables: {str(e)}"})
88
98
 
89
- async def execute_sql(self, query: str, limit: Optional[int] = 100) -> str:
99
+ async def execute_sql(self, query: str, limit: Optional[int] = None) -> str:
90
100
  """Execute a SQL query against the database."""
91
101
  try:
92
102
  # Security check - only allow SELECT queries unless write is enabled
@@ -62,12 +62,20 @@ class DisplayManager:
62
62
  )
63
63
 
64
64
  # Create table with columns from first result
65
- columns = list(results[0].keys())
66
- table = self._create_table(columns)
65
+ all_columns = list(results[0].keys())
66
+ display_columns = all_columns[:15] # Limit to first 15 columns
67
+
68
+ # Show warning if columns were truncated
69
+ if len(all_columns) > 15:
70
+ self.console.print(
71
+ f"[yellow]Note: Showing first 15 of {len(all_columns)} columns[/yellow]"
72
+ )
73
+
74
+ table = self._create_table(display_columns)
67
75
 
68
76
  # Add rows (show first 20 rows)
69
77
  for row in results[:20]:
70
- table.add_row(*[str(row[key]) for key in columns])
78
+ table.add_row(*[str(row[key]) for key in display_columns])
71
79
 
72
80
  self.console.print(table)
73
81
 
@@ -1,6 +1,10 @@
1
1
  """Interactive mode handling for the CLI."""
2
2
 
3
+ import asyncio
4
+ from typing import Optional
5
+
3
6
  import questionary
7
+ from prompt_toolkit.completion import Completer, Completion
4
8
  from rich.console import Console
5
9
  from rich.panel import Panel
6
10
 
@@ -9,6 +13,34 @@ from sqlsaber.cli.display import DisplayManager
9
13
  from sqlsaber.cli.streaming import StreamingQueryHandler
10
14
 
11
15
 
16
+ class SlashCommandCompleter(Completer):
17
+ """Custom completer for slash commands."""
18
+
19
+ def get_completions(self, document, complete_event):
20
+ """Get completions for slash commands."""
21
+ # Only provide completions if the line starts with "/"
22
+ text = document.text
23
+ if text.startswith("/"):
24
+ # Get the partial command after the slash
25
+ partial_cmd = text[1:]
26
+
27
+ # Define available commands with descriptions
28
+ commands = [
29
+ ("clear", "Clear conversation history"),
30
+ ("exit", "Exit the interactive session"),
31
+ ("quit", "Exit the interactive session"),
32
+ ]
33
+
34
+ # Yield completions that match the partial command
35
+ for cmd, description in commands:
36
+ if cmd.startswith(partial_cmd):
37
+ yield Completion(
38
+ cmd,
39
+ start_position=-len(partial_cmd),
40
+ display_meta=description,
41
+ )
42
+
43
+
12
44
  class InteractiveSession:
13
45
  """Manages interactive CLI sessions."""
14
46
 
@@ -17,6 +49,8 @@ class InteractiveSession:
17
49
  self.agent = agent
18
50
  self.display = DisplayManager(console)
19
51
  self.streaming_handler = StreamingQueryHandler(console)
52
+ self.current_task: Optional[asyncio.Task] = None
53
+ self.cancellation_token: Optional[asyncio.Event] = None
20
54
 
21
55
  def show_welcome_message(self):
22
56
  """Display welcome message for interactive mode."""
@@ -28,7 +62,7 @@ class InteractiveSession:
28
62
  Panel.fit(
29
63
  "[bold green]SQLSaber - Use the agent Luke![/bold green]\n\n"
30
64
  "[bold]Your agentic SQL assistant.[/bold]\n\n\n"
31
- "[dim]Use 'clear' to reset conversation, 'exit' or 'quit' to leave.[/dim]\n\n"
65
+ "[dim]Use '/clear' to reset conversation, '/exit' or '/quit' to leave.[/dim]\n\n"
32
66
  "[dim]Start a message with '#' to add something to agent's memory for this database.[/dim]",
33
67
  border_style="green",
34
68
  )
@@ -38,8 +72,31 @@ class InteractiveSession:
38
72
  )
39
73
  self.console.print(
40
74
  "[dim]Press Esc-Enter or Meta-Enter to submit your query.[/dim]\n"
75
+ "[dim]Press Ctrl+C during query execution to interrupt and return to prompt.[/dim]\n"
41
76
  )
42
77
 
78
+ async def _execute_query_with_cancellation(self, user_query: str):
79
+ """Execute a query with cancellation support."""
80
+ # Create cancellation token
81
+ self.cancellation_token = asyncio.Event()
82
+
83
+ # Create the query task
84
+ query_task = asyncio.create_task(
85
+ self.streaming_handler.execute_streaming_query(
86
+ user_query, self.agent, self.cancellation_token
87
+ )
88
+ )
89
+ self.current_task = query_task
90
+
91
+ try:
92
+ # Simply await the query task
93
+ # Ctrl+C will be handled by the KeyboardInterrupt exception in run()
94
+ await query_task
95
+
96
+ finally:
97
+ self.current_task = None
98
+ self.cancellation_token = None
99
+
43
100
  async def run(self):
44
101
  """Run the interactive session loop."""
45
102
  self.show_welcome_message()
@@ -51,12 +108,16 @@ class InteractiveSession:
51
108
  qmark="",
52
109
  multiline=True,
53
110
  instruction="",
111
+ completer=SlashCommandCompleter(),
54
112
  ).ask_async()
55
113
 
56
- if user_query.lower() in ["exit", "quit", "q"]:
114
+ if not user_query:
115
+ continue
116
+
117
+ if user_query in ["/exit", "/quit"]:
57
118
  break
58
119
 
59
- if user_query.lower() == "clear":
120
+ if user_query == "/clear":
60
121
  self.agent.clear_history()
61
122
  self.console.print("[green]Conversation history cleared.[/green]\n")
62
123
  continue
@@ -85,12 +146,24 @@ class InteractiveSession:
85
146
  )
86
147
  continue
87
148
 
88
- await self.streaming_handler.execute_streaming_query(
89
- user_query, self.agent
90
- )
149
+ # Execute query with cancellation support
150
+ await self._execute_query_with_cancellation(user_query)
91
151
  self.display.show_newline() # Empty line for readability
92
152
 
93
153
  except KeyboardInterrupt:
94
- self.console.print("\n[yellow]Use 'exit' or 'quit' to leave.[/yellow]")
154
+ # Handle Ctrl+C - cancel current task if running
155
+ if self.current_task and not self.current_task.done():
156
+ if self.cancellation_token is not None:
157
+ self.cancellation_token.set()
158
+ self.current_task.cancel()
159
+ try:
160
+ await self.current_task
161
+ except asyncio.CancelledError:
162
+ pass
163
+ self.console.print("\n[yellow]Query interrupted[/yellow]")
164
+ else:
165
+ self.console.print(
166
+ "\n[yellow]Use '/exit' or '/quit' to leave.[/yellow]"
167
+ )
95
168
  except Exception as e:
96
169
  self.console.print(f"[bold red]Error:[/bold red] {str(e)}")
@@ -1,5 +1,7 @@
1
1
  """Streaming query handling for the CLI."""
2
2
 
3
+ import asyncio
4
+
3
5
  from rich.console import Console
4
6
 
5
7
  from sqlsaber.agents.base import BaseSQLAgent
@@ -13,7 +15,12 @@ class StreamingQueryHandler:
13
15
  self.console = console
14
16
  self.display = DisplayManager(console)
15
17
 
16
- async def execute_streaming_query(self, user_query: str, agent: BaseSQLAgent):
18
+ async def execute_streaming_query(
19
+ self,
20
+ user_query: str,
21
+ agent: BaseSQLAgent,
22
+ cancellation_token: asyncio.Event | None = None,
23
+ ):
17
24
  """Execute a query with streaming display."""
18
25
 
19
26
  has_content = False
@@ -24,7 +31,12 @@ class StreamingQueryHandler:
24
31
  status.start()
25
32
 
26
33
  try:
27
- async for event in agent.query_stream(user_query):
34
+ async for event in agent.query_stream(
35
+ user_query, cancellation_token=cancellation_token
36
+ ):
37
+ if cancellation_token is not None and cancellation_token.is_set():
38
+ break
39
+
28
40
  if event.type == "tool_use":
29
41
  # Stop any ongoing status, but don't mark has_content yet
30
42
  self._stop_status(status)
@@ -83,6 +95,13 @@ class StreamingQueryHandler:
83
95
  has_content = True
84
96
  self.display.show_error(event.data)
85
97
 
98
+ except asyncio.CancelledError:
99
+ # Handle cancellation gracefully
100
+ self._stop_status(status)
101
+ if explanation_started:
102
+ self.display.show_newline()
103
+ self.console.print("[yellow]Query interrupted[/yellow]")
104
+ return
86
105
  finally:
87
106
  # Make sure status is stopped
88
107
  self._stop_status(status)
@@ -863,7 +863,7 @@ wheels = [
863
863
 
864
864
  [[package]]
865
865
  name = "sqlsaber"
866
- version = "0.5.0"
866
+ version = "0.6.0"
867
867
  source = { editable = "." }
868
868
  dependencies = [
869
869
  { name = "aiomysql" },
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes