sqlsaber 0.29.1__py3-none-any.whl → 0.30.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/base.py +1 -1
- sqlsaber/agents/pydantic_ai_agent.py +39 -17
- sqlsaber/cli/models.py +1 -1
- sqlsaber/cli/streaming.py +2 -11
- sqlsaber/theme/manager.py +4 -1
- sqlsaber/tools/__init__.py +0 -5
- sqlsaber/tools/base.py +0 -31
- sqlsaber/tools/registry.py +6 -39
- sqlsaber/tools/sql_tools.py +0 -42
- {sqlsaber-0.29.1.dist-info → sqlsaber-0.30.0.dist-info}/METADATA +3 -44
- {sqlsaber-0.29.1.dist-info → sqlsaber-0.30.0.dist-info}/RECORD +14 -19
- {sqlsaber-0.29.1.dist-info → sqlsaber-0.30.0.dist-info}/entry_points.txt +0 -2
- sqlsaber/agents/mcp.py +0 -21
- sqlsaber/mcp/__init__.py +0 -5
- sqlsaber/mcp/mcp.py +0 -129
- sqlsaber/tools/enums.py +0 -19
- sqlsaber/tools/instructions.py +0 -231
- {sqlsaber-0.29.1.dist-info → sqlsaber-0.30.0.dist-info}/WHEEL +0 -0
- {sqlsaber-0.29.1.dist-info → sqlsaber-0.30.0.dist-info}/licenses/LICENSE +0 -0
sqlsaber/agents/base.py
CHANGED
|
@@ -61,7 +61,7 @@ class BaseSQLAgent(ABC):
|
|
|
61
61
|
def _init_tools(self) -> None:
|
|
62
62
|
"""Initialize SQL tools with database connection."""
|
|
63
63
|
# Get all SQL tools and set their database connection
|
|
64
|
-
for tool_name in tool_registry.list_tools(
|
|
64
|
+
for tool_name in tool_registry.list_tools():
|
|
65
65
|
tool = tool_registry.get_tool(tool_name)
|
|
66
66
|
if isinstance(tool, SQLTool):
|
|
67
67
|
tool.set_connection(self.db)
|
|
@@ -24,7 +24,6 @@ from sqlsaber.database import (
|
|
|
24
24
|
SQLiteConnection,
|
|
25
25
|
)
|
|
26
26
|
from sqlsaber.memory.manager import MemoryManager
|
|
27
|
-
from sqlsaber.tools.instructions import InstructionBuilder
|
|
28
27
|
from sqlsaber.tools.registry import tool_registry
|
|
29
28
|
from sqlsaber.tools.sql_tools import SQLTool
|
|
30
29
|
|
|
@@ -43,7 +42,6 @@ class SQLSaberAgent:
|
|
|
43
42
|
self.database_name = database_name
|
|
44
43
|
self.config = Config()
|
|
45
44
|
self.memory_manager = memory_manager or MemoryManager()
|
|
46
|
-
self.instruction_builder = InstructionBuilder(tool_registry)
|
|
47
45
|
self.db_type = self._get_database_type_name()
|
|
48
46
|
|
|
49
47
|
# Thinking configuration (CLI override or config default)
|
|
@@ -61,7 +59,7 @@ class SQLSaberAgent:
|
|
|
61
59
|
|
|
62
60
|
def _configure_sql_tools(self) -> None:
|
|
63
61
|
"""Ensure SQL tools receive the active database connection."""
|
|
64
|
-
for tool_name in tool_registry.list_tools(
|
|
62
|
+
for tool_name in tool_registry.list_tools():
|
|
65
63
|
tool = tool_registry.get_tool(tool_name)
|
|
66
64
|
if isinstance(tool, SQLTool):
|
|
67
65
|
tool.set_connection(self.db_connection)
|
|
@@ -165,30 +163,54 @@ class SQLSaberAgent:
|
|
|
165
163
|
return Agent(model_obj, name="sqlsaber")
|
|
166
164
|
|
|
167
165
|
def _setup_system_prompt(self, agent: Agent) -> None:
|
|
168
|
-
"""
|
|
166
|
+
"""Configure the agent's system prompt using a simple prompt string."""
|
|
169
167
|
if not self.is_oauth:
|
|
170
168
|
|
|
171
169
|
@agent.system_prompt(dynamic=True)
|
|
172
170
|
async def sqlsaber_system_prompt(ctx: RunContext) -> str:
|
|
173
|
-
|
|
174
|
-
db_type=self.db_type
|
|
175
|
-
)
|
|
176
|
-
|
|
177
|
-
# Add memory context if available
|
|
178
|
-
mem = ""
|
|
179
|
-
if self.database_name:
|
|
180
|
-
mem = self.memory_manager.format_memories_for_prompt(
|
|
181
|
-
self.database_name
|
|
182
|
-
)
|
|
183
|
-
|
|
184
|
-
parts = [p for p in (instructions, mem) if p and p.strip()]
|
|
185
|
-
return "\n\n".join(parts) if parts else ""
|
|
171
|
+
return self.system_prompt_text(include_memory=True)
|
|
186
172
|
else:
|
|
187
173
|
|
|
188
174
|
@agent.system_prompt(dynamic=True)
|
|
189
175
|
async def sqlsaber_system_prompt(ctx: RunContext) -> str:
|
|
176
|
+
# OAuth clients (Claude Code) ignore custom system prompts; we inject later.
|
|
190
177
|
return "You are Claude Code, Anthropic's official CLI for Claude."
|
|
191
178
|
|
|
179
|
+
def system_prompt_text(self, include_memory: bool = True) -> str:
|
|
180
|
+
"""Return the original SQLSaber system prompt as a single string."""
|
|
181
|
+
db = self.db_type
|
|
182
|
+
base = (
|
|
183
|
+
f"You are a helpful SQL assistant that helps users query their {db} database.\n\n"
|
|
184
|
+
"Your responsibilities:\n"
|
|
185
|
+
"1. Understand user's natural language requests, think and convert them to SQL\n"
|
|
186
|
+
"2. Use the provided tools efficiently to explore database schema\n"
|
|
187
|
+
"3. Generate appropriate SQL queries\n"
|
|
188
|
+
"4. Execute queries safely - queries that modify the database are not allowed\n"
|
|
189
|
+
"5. Format and explain results clearly\n\n"
|
|
190
|
+
"IMPORTANT - Tool Usage Strategy:\n"
|
|
191
|
+
"1. ALWAYS start with 'list_tables' to see available tables and row counts. Use this first to discover available tables.\n"
|
|
192
|
+
"2. Use 'introspect_schema' with a table_pattern to get details ONLY for relevant tables. Use table patterns like 'sample%' or '%experiment%' to filter related tables.\n"
|
|
193
|
+
"3. Execute SQL queries safely with automatic LIMIT clauses for SELECT statements. Only SELECT queries are permitted for security.\n\n"
|
|
194
|
+
"Tool-Specific Guidelines:\n"
|
|
195
|
+
"- introspect_schema: Use 'introspect_schema' with a table_pattern to get details ONLY for relevant tables. Use table patterns like 'sample%' or '%experiment%' to filter related tables.\n"
|
|
196
|
+
"- execute_sql: Execute SQL queries safely with automatic LIMIT clauses for SELECT statements. Only SELECT queries are permitted for security.\n\n"
|
|
197
|
+
"Guidelines:\n"
|
|
198
|
+
"- Use proper JOIN syntax and avoid cartesian products\n"
|
|
199
|
+
"- Include appropriate WHERE clauses to limit results\n"
|
|
200
|
+
"- Explain what the query does in simple terms\n"
|
|
201
|
+
"- Handle errors gracefully and suggest fixes\n"
|
|
202
|
+
"- Be security conscious - use parameterized queries when needed\n"
|
|
203
|
+
"- Timestamp columns must be converted to text when you write queries\n"
|
|
204
|
+
"- Use table patterns like 'sample%' or '%experiment%' to filter related tables"
|
|
205
|
+
)
|
|
206
|
+
|
|
207
|
+
if include_memory and self.database_name:
|
|
208
|
+
mem = self.memory_manager.format_memories_for_prompt(self.database_name)
|
|
209
|
+
mem = mem.strip()
|
|
210
|
+
if mem:
|
|
211
|
+
return f"{base}\n\n{mem}"
|
|
212
|
+
return base
|
|
213
|
+
|
|
192
214
|
def _register_tools(self, agent: Agent) -> None:
|
|
193
215
|
"""Register all the SQL tools with the agent."""
|
|
194
216
|
|
sqlsaber/cli/models.py
CHANGED
|
@@ -152,7 +152,7 @@ def list():
|
|
|
152
152
|
table.add_column("Name", style="success")
|
|
153
153
|
table.add_column("Description", style="info")
|
|
154
154
|
table.add_column("Context", style="warning", justify="right")
|
|
155
|
-
table.add_column("Current", style="
|
|
155
|
+
table.add_column("Current", style="accent", justify="center")
|
|
156
156
|
|
|
157
157
|
current_model = model_manager.get_current_model()
|
|
158
158
|
|
sqlsaber/cli/streaming.py
CHANGED
|
@@ -144,17 +144,8 @@ class StreamingQueryHandler:
|
|
|
144
144
|
prepared_prompt: str | list[str] = user_query
|
|
145
145
|
no_history = not message_history
|
|
146
146
|
if sqlsaber_agent.is_oauth and no_history:
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
)
|
|
150
|
-
mem = ""
|
|
151
|
-
if sqlsaber_agent.database_name:
|
|
152
|
-
mem = sqlsaber_agent.memory_manager.format_memories_for_prompt(
|
|
153
|
-
sqlsaber_agent.database_name
|
|
154
|
-
)
|
|
155
|
-
parts = [p for p in (instructions, mem) if p and str(p).strip()]
|
|
156
|
-
if parts:
|
|
157
|
-
injected = "\n\n".join(parts)
|
|
147
|
+
injected = sqlsaber_agent.system_prompt_text(include_memory=True)
|
|
148
|
+
if injected and str(injected).strip():
|
|
158
149
|
prepared_prompt = [injected, user_query]
|
|
159
150
|
|
|
160
151
|
# Show a transient status until events start streaming
|
sqlsaber/theme/manager.py
CHANGED
|
@@ -9,7 +9,7 @@ from typing import Dict
|
|
|
9
9
|
from platformdirs import user_config_dir
|
|
10
10
|
from prompt_toolkit.styles import Style as PTStyle
|
|
11
11
|
from prompt_toolkit.styles.pygments import style_from_pygments_cls
|
|
12
|
-
from pygments.styles import get_style_by_name
|
|
12
|
+
from pygments.styles import get_all_styles, get_style_by_name
|
|
13
13
|
from pygments.token import Token
|
|
14
14
|
from pygments.util import ClassNotFound
|
|
15
15
|
from rich.console import Console
|
|
@@ -200,6 +200,9 @@ def get_theme_manager() -> ThemeManager:
|
|
|
200
200
|
user_cfg = _load_user_theme_config()
|
|
201
201
|
env_name = os.getenv("SQLSABER_THEME")
|
|
202
202
|
|
|
203
|
+
if env_name and env_name.lower() not in get_all_styles():
|
|
204
|
+
env_name = None
|
|
205
|
+
|
|
203
206
|
name = (
|
|
204
207
|
env_name or user_cfg.get("theme", {}).get("name") or DEFAULT_THEME_NAME
|
|
205
208
|
).lower()
|
sqlsaber/tools/__init__.py
CHANGED
|
@@ -1,8 +1,6 @@
|
|
|
1
1
|
"""SQLSaber tools module."""
|
|
2
2
|
|
|
3
3
|
from .base import Tool
|
|
4
|
-
from .enums import ToolCategory, WorkflowPosition
|
|
5
|
-
from .instructions import InstructionBuilder
|
|
6
4
|
from .registry import ToolRegistry, register_tool, tool_registry
|
|
7
5
|
|
|
8
6
|
# Import concrete tools to register them
|
|
@@ -10,12 +8,9 @@ from .sql_tools import ExecuteSQLTool, IntrospectSchemaTool, ListTablesTool, SQL
|
|
|
10
8
|
|
|
11
9
|
__all__ = [
|
|
12
10
|
"Tool",
|
|
13
|
-
"ToolCategory",
|
|
14
|
-
"WorkflowPosition",
|
|
15
11
|
"ToolRegistry",
|
|
16
12
|
"tool_registry",
|
|
17
13
|
"register_tool",
|
|
18
|
-
"InstructionBuilder",
|
|
19
14
|
"SQLTool",
|
|
20
15
|
"ListTablesTool",
|
|
21
16
|
"IntrospectSchemaTool",
|
sqlsaber/tools/base.py
CHANGED
|
@@ -3,8 +3,6 @@
|
|
|
3
3
|
from abc import ABC, abstractmethod
|
|
4
4
|
from typing import Any
|
|
5
5
|
|
|
6
|
-
from .enums import ToolCategory, WorkflowPosition
|
|
7
|
-
|
|
8
6
|
|
|
9
7
|
class Tool(ABC):
|
|
10
8
|
"""Abstract base class for all tools."""
|
|
@@ -42,32 +40,3 @@ class Tool(ABC):
|
|
|
42
40
|
JSON string with the tool's output
|
|
43
41
|
"""
|
|
44
42
|
pass
|
|
45
|
-
|
|
46
|
-
@property
|
|
47
|
-
def category(self) -> ToolCategory:
|
|
48
|
-
"""Return the tool category. Override to customize."""
|
|
49
|
-
return ToolCategory.GENERAL
|
|
50
|
-
|
|
51
|
-
def get_usage_instructions(self) -> str | None:
|
|
52
|
-
"""Return tool-specific usage instructions for LLM guidance.
|
|
53
|
-
|
|
54
|
-
Returns:
|
|
55
|
-
Usage instructions string, or None for no specific guidance
|
|
56
|
-
"""
|
|
57
|
-
return None
|
|
58
|
-
|
|
59
|
-
def get_priority(self) -> int:
|
|
60
|
-
"""Return priority for tool ordering in instructions.
|
|
61
|
-
|
|
62
|
-
Returns:
|
|
63
|
-
Priority number (lower = higher priority, default = 100)
|
|
64
|
-
"""
|
|
65
|
-
return 100
|
|
66
|
-
|
|
67
|
-
def get_workflow_position(self) -> WorkflowPosition:
|
|
68
|
-
"""Return the typical workflow position for this tool.
|
|
69
|
-
|
|
70
|
-
Returns:
|
|
71
|
-
WorkflowPosition enum value
|
|
72
|
-
"""
|
|
73
|
-
return WorkflowPosition.OTHER
|
sqlsaber/tools/registry.py
CHANGED
|
@@ -3,7 +3,6 @@
|
|
|
3
3
|
from typing import Type
|
|
4
4
|
|
|
5
5
|
from .base import Tool
|
|
6
|
-
from .enums import ToolCategory
|
|
7
6
|
|
|
8
7
|
|
|
9
8
|
class ToolRegistry:
|
|
@@ -61,45 +60,13 @@ class ToolRegistry:
|
|
|
61
60
|
|
|
62
61
|
return self._instances[name]
|
|
63
62
|
|
|
64
|
-
def list_tools(self
|
|
65
|
-
"""List all registered tool names.
|
|
63
|
+
def list_tools(self) -> list[str]:
|
|
64
|
+
"""List all registered tool names."""
|
|
65
|
+
return list(self._tools.keys())
|
|
66
66
|
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
Returns:
|
|
71
|
-
List of tool names
|
|
72
|
-
"""
|
|
73
|
-
if category is None:
|
|
74
|
-
return list(self._tools.keys())
|
|
75
|
-
|
|
76
|
-
# Convert string to enum
|
|
77
|
-
if isinstance(category, str):
|
|
78
|
-
try:
|
|
79
|
-
category = ToolCategory(category)
|
|
80
|
-
except ValueError:
|
|
81
|
-
# If string doesn't match any enum, return empty list
|
|
82
|
-
return []
|
|
83
|
-
|
|
84
|
-
# Filter by category
|
|
85
|
-
result = []
|
|
86
|
-
for name, tool_class in self._tools.items():
|
|
87
|
-
tool = self.get_tool(name)
|
|
88
|
-
if tool.category == category:
|
|
89
|
-
result.append(name)
|
|
90
|
-
return result
|
|
91
|
-
|
|
92
|
-
def get_all_tools(self, category: str | ToolCategory | None = None) -> list[Tool]:
|
|
93
|
-
"""Get all tool instances.
|
|
94
|
-
|
|
95
|
-
Args:
|
|
96
|
-
category: Optional category to filter by (string or ToolCategory enum)
|
|
97
|
-
|
|
98
|
-
Returns:
|
|
99
|
-
List of tool instances
|
|
100
|
-
"""
|
|
101
|
-
names = self.list_tools(category)
|
|
102
|
-
return [self.get_tool(name) for name in names]
|
|
67
|
+
def get_all_tools(self) -> list[Tool]:
|
|
68
|
+
"""Get all tool instances."""
|
|
69
|
+
return [self.get_tool(name) for name in self.list_tools()]
|
|
103
70
|
|
|
104
71
|
|
|
105
72
|
# Global registry instance
|
sqlsaber/tools/sql_tools.py
CHANGED
|
@@ -7,7 +7,6 @@ from sqlsaber.database import BaseDatabaseConnection
|
|
|
7
7
|
from sqlsaber.database.schema import SchemaManager
|
|
8
8
|
|
|
9
9
|
from .base import Tool
|
|
10
|
-
from .enums import ToolCategory, WorkflowPosition
|
|
11
10
|
from .registry import register_tool
|
|
12
11
|
from .sql_guard import add_limit, validate_read_only
|
|
13
12
|
|
|
@@ -26,11 +25,6 @@ class SQLTool(Tool):
|
|
|
26
25
|
self.db = db_connection
|
|
27
26
|
self.schema_manager = SchemaManager(db_connection)
|
|
28
27
|
|
|
29
|
-
@property
|
|
30
|
-
def category(self) -> ToolCategory:
|
|
31
|
-
"""SQL tools belong to the 'sql' category."""
|
|
32
|
-
return ToolCategory.SQL
|
|
33
|
-
|
|
34
28
|
|
|
35
29
|
@register_tool
|
|
36
30
|
class ListTablesTool(SQLTool):
|
|
@@ -52,18 +46,6 @@ class ListTablesTool(SQLTool):
|
|
|
52
46
|
"required": [],
|
|
53
47
|
}
|
|
54
48
|
|
|
55
|
-
def get_usage_instructions(self) -> str | None:
|
|
56
|
-
"""Return usage instructions for this tool."""
|
|
57
|
-
return "ALWAYS start with 'list_tables' to see available tables and row counts. Use this first to discover available tables."
|
|
58
|
-
|
|
59
|
-
def get_priority(self) -> int:
|
|
60
|
-
"""Return priority for tool ordering."""
|
|
61
|
-
return 10 # High priority - should be used first
|
|
62
|
-
|
|
63
|
-
def get_workflow_position(self) -> WorkflowPosition:
|
|
64
|
-
"""Return workflow position."""
|
|
65
|
-
return WorkflowPosition.DISCOVERY
|
|
66
|
-
|
|
67
49
|
async def execute(self, **kwargs) -> str:
|
|
68
50
|
"""List all tables in the database."""
|
|
69
51
|
if not self.db or not self.schema_manager:
|
|
@@ -101,18 +83,6 @@ class IntrospectSchemaTool(SQLTool):
|
|
|
101
83
|
"required": [],
|
|
102
84
|
}
|
|
103
85
|
|
|
104
|
-
def get_usage_instructions(self) -> str | None:
|
|
105
|
-
"""Return usage instructions for this tool."""
|
|
106
|
-
return "Use 'introspect_schema' with a table_pattern to get details ONLY for relevant tables. Use table patterns like 'sample%' or '%experiment%' to filter related tables."
|
|
107
|
-
|
|
108
|
-
def get_priority(self) -> int:
|
|
109
|
-
"""Return priority for tool ordering."""
|
|
110
|
-
return 20 # Should come after list_tables
|
|
111
|
-
|
|
112
|
-
def get_workflow_position(self) -> WorkflowPosition:
|
|
113
|
-
"""Return workflow position."""
|
|
114
|
-
return WorkflowPosition.ANALYSIS
|
|
115
|
-
|
|
116
86
|
async def execute(self, **kwargs) -> str:
|
|
117
87
|
"""Introspect database schema."""
|
|
118
88
|
if not self.db or not self.schema_manager:
|
|
@@ -184,18 +154,6 @@ class ExecuteSQLTool(SQLTool):
|
|
|
184
154
|
"required": ["query"],
|
|
185
155
|
}
|
|
186
156
|
|
|
187
|
-
def get_usage_instructions(self) -> str | None:
|
|
188
|
-
"""Return usage instructions for this tool."""
|
|
189
|
-
return "Execute SQL queries safely with automatic LIMIT clauses for SELECT statements. Only SELECT queries are permitted for security."
|
|
190
|
-
|
|
191
|
-
def get_priority(self) -> int:
|
|
192
|
-
"""Return priority for tool ordering."""
|
|
193
|
-
return 30 # Should come after schema tools
|
|
194
|
-
|
|
195
|
-
def get_workflow_position(self) -> WorkflowPosition:
|
|
196
|
-
"""Return workflow position."""
|
|
197
|
-
return WorkflowPosition.EXECUTION
|
|
198
|
-
|
|
199
157
|
async def execute(self, **kwargs) -> str:
|
|
200
158
|
"""Execute a SQL query."""
|
|
201
159
|
if not self.db:
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: sqlsaber
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.30.0
|
|
4
4
|
Summary: SQLsaber - Open-source agentic SQL assistant
|
|
5
5
|
License-File: LICENSE
|
|
6
6
|
Requires-Python: >=3.12
|
|
@@ -9,7 +9,6 @@ Requires-Dist: aiosqlite>=0.21.0
|
|
|
9
9
|
Requires-Dist: asyncpg>=0.30.0
|
|
10
10
|
Requires-Dist: cyclopts>=3.22.1
|
|
11
11
|
Requires-Dist: duckdb>=0.9.2
|
|
12
|
-
Requires-Dist: fastmcp>=2.9.0
|
|
13
12
|
Requires-Dist: httpx>=0.28.1
|
|
14
13
|
Requires-Dist: keyring>=25.6.0
|
|
15
14
|
Requires-Dist: platformdirs>=4.0.0
|
|
@@ -44,10 +43,7 @@ Ask your questions in natural language and `sqlsaber` will gather the right cont
|
|
|
44
43
|
- [Resume Past Conversation](#resume-past-conversation)
|
|
45
44
|
- [Database Selection](#database-selection)
|
|
46
45
|
- [Examples](#examples)
|
|
47
|
-
|
|
48
|
-
- [Starting the MCP Server](#starting-the-mcp-server)
|
|
49
|
-
- [Configuring MCP Clients](#configuring-mcp-clients)
|
|
50
|
-
- [Available MCP Tools](#available-mcp-tools)
|
|
46
|
+
|
|
51
47
|
- [How It Works](#how-it-works)
|
|
52
48
|
- [Contributing](#contributing)
|
|
53
49
|
- [License](#license)
|
|
@@ -60,7 +56,7 @@ Ask your questions in natural language and `sqlsaber` will gather the right cont
|
|
|
60
56
|
- Interactive REPL mode
|
|
61
57
|
- Conversation threads (store, display, and resume conversations)
|
|
62
58
|
- Support for PostgreSQL, MySQL, SQLite, DuckDB, and CSVs
|
|
63
|
-
|
|
59
|
+
|
|
64
60
|
- Extended thinking mode for select models (Anthropic, OpenAI, Google, Groq)
|
|
65
61
|
- Beautiful formatted output
|
|
66
62
|
|
|
@@ -220,43 +216,6 @@ saber "show me orders with customer details for this week"
|
|
|
220
216
|
saber "which products had the highest sales growth last quarter?"
|
|
221
217
|
```
|
|
222
218
|
|
|
223
|
-
## MCP Server Integration
|
|
224
|
-
|
|
225
|
-
SQLSaber includes an MCP (Model Context Protocol) server that allows AI agents like Claude Code to directly leverage tools available in SQLSaber.
|
|
226
|
-
|
|
227
|
-
### Starting the MCP Server
|
|
228
|
-
|
|
229
|
-
Run the MCP server using uvx:
|
|
230
|
-
|
|
231
|
-
```bash
|
|
232
|
-
uvx --from sqlsaber saber-mcp
|
|
233
|
-
```
|
|
234
|
-
|
|
235
|
-
### Configuring MCP Clients
|
|
236
|
-
|
|
237
|
-
#### Claude Code
|
|
238
|
-
|
|
239
|
-
Add SQLSaber as an MCP server in Claude Code:
|
|
240
|
-
|
|
241
|
-
```bash
|
|
242
|
-
claude mcp add sqlsaber -- uvx --from sqlsaber saber-mcp
|
|
243
|
-
```
|
|
244
|
-
|
|
245
|
-
#### Other MCP Clients
|
|
246
|
-
|
|
247
|
-
For other MCP clients, configure them to run the command: `uvx --from sqlsaber saber-mcp`
|
|
248
|
-
|
|
249
|
-
### Available MCP Tools
|
|
250
|
-
|
|
251
|
-
Once connected, the MCP client will have access to these tools:
|
|
252
|
-
|
|
253
|
-
- `get_databases()` - Lists all configured databases
|
|
254
|
-
- `list_tables(database)` - Get all tables in a database with row counts
|
|
255
|
-
- `introspect_schema(database, table_pattern?)` - Get detailed schema information
|
|
256
|
-
- `execute_sql(database, query, limit?)` - Execute SQL queries (read-only)
|
|
257
|
-
|
|
258
|
-
The MCP server uses your existing SQLSaber database configurations, so make sure to set up your databases using `saber db add` first.
|
|
259
|
-
|
|
260
219
|
## How It Works
|
|
261
220
|
|
|
262
221
|
SQLsaber uses a multi-step agentic process to gather the right context and execute SQL queries to answer your questions:
|
|
@@ -1,9 +1,8 @@
|
|
|
1
1
|
sqlsaber/__init__.py,sha256=HjS8ULtP4MGpnTL7njVY45NKV9Fi4e_yeYuY-hyXWQc,73
|
|
2
2
|
sqlsaber/__main__.py,sha256=RIHxWeWh2QvLfah-2OkhI5IJxojWfy4fXpMnVEJYvxw,78
|
|
3
3
|
sqlsaber/agents/__init__.py,sha256=qYI6rLY4q5AbF47vXH5RVoM08-yQjymBSaePh4lFIW4,116
|
|
4
|
-
sqlsaber/agents/base.py,sha256=
|
|
5
|
-
sqlsaber/agents/
|
|
6
|
-
sqlsaber/agents/pydantic_ai_agent.py,sha256=wBxKz0pjOkL-HI-TXV6B67bczZNgu7k26Rr3w5usR3o,10064
|
|
4
|
+
sqlsaber/agents/base.py,sha256=T05UsMZPwAlMhsJFpuuVI1RNDhdiwiEsgCWr9MbPoAU,2654
|
|
5
|
+
sqlsaber/agents/pydantic_ai_agent.py,sha256=cHHwQHJf9TqrBhItWJgnScL31lvyvKLcCBTSjRSwWug,12002
|
|
7
6
|
sqlsaber/application/__init__.py,sha256=KY_-d5nEdQyAwNOsK5r-f7Tb69c63XbuEkHPeLpJal8,84
|
|
8
7
|
sqlsaber/application/auth_setup.py,sha256=D94dyU9bOVfnNHLnnFJb5PaeWsKPTL21CiS_DLcY93A,5114
|
|
9
8
|
sqlsaber/application/db_setup.py,sha256=ZSgR9rJJVHttIjsbYQS9GEIyzkM09k5RLrVGdegrfYc,6859
|
|
@@ -17,9 +16,9 @@ sqlsaber/cli/database.py,sha256=Tqy8H5MnjsrmOSPcbA5Qy-u-IOYJCIXRJVhk0veLNDk,1072
|
|
|
17
16
|
sqlsaber/cli/display.py,sha256=WB5JCumhXadziDEX1EZHG3vN1Chol5FNAaTXHieqFK0,17892
|
|
18
17
|
sqlsaber/cli/interactive.py,sha256=PcY6mszImo_3PsqjjWmx_cOfj44OmKvD9ENOvGA-wjU,13715
|
|
19
18
|
sqlsaber/cli/memory.py,sha256=IKq09DUbqpvvtATsyDlpm7rDlGqWEhdUX9wgnR-oiq4,7850
|
|
20
|
-
sqlsaber/cli/models.py,sha256=
|
|
19
|
+
sqlsaber/cli/models.py,sha256=nbn75gCnkRciGt4Q47yxa8wImiZcCkDdQZNVeehDim8,8530
|
|
21
20
|
sqlsaber/cli/onboarding.py,sha256=iBGT-W-OJFRvQoEpuHYyO1c9Mym5c97eIefRvxGHtTg,11265
|
|
22
|
-
sqlsaber/cli/streaming.py,sha256=
|
|
21
|
+
sqlsaber/cli/streaming.py,sha256=eggj25ZlA-xKrAF726S29vfS2MHTFC5wTmgXLbS-RvM,6515
|
|
23
22
|
sqlsaber/cli/theme.py,sha256=hP0kmsMLCtqaT7b5wB1dk1hW1hV94oP4BHdz8S6887A,4243
|
|
24
23
|
sqlsaber/cli/threads.py,sha256=o9q9Hst1Wt7cxSyrpAtwG6pkUct6csgiAmN_0P_WO3k,13637
|
|
25
24
|
sqlsaber/config/__init__.py,sha256=olwC45k8Nc61yK0WmPUk7XHdbsZH9HuUAbwnmKe3IgA,100
|
|
@@ -39,24 +38,20 @@ sqlsaber/database/postgresql.py,sha256=fuf2Wl29NKXvD3mqsR08PDleNQ1PG-fNvWSxT6HDh
|
|
|
39
38
|
sqlsaber/database/resolver.py,sha256=wSCcn__aCqwIfpt_LCjtW2Zgb8RpG5PlmwwZHli1q_U,3628
|
|
40
39
|
sqlsaber/database/schema.py,sha256=CuV0ewoVaERe1gj_fJFJFWAP8aEPgepmn6X6B7bgkfQ,6962
|
|
41
40
|
sqlsaber/database/sqlite.py,sha256=iReEIiSpkhhS1VzITd79ZWqSL3fHMyfe3DRCDpM0DvE,9421
|
|
42
|
-
sqlsaber/mcp/__init__.py,sha256=COdWq7wauPBp5Ew8tfZItFzbcLDSEkHBJSMhxzy8C9c,112
|
|
43
|
-
sqlsaber/mcp/mcp.py,sha256=tpNPHpkaCre1Xjp7c4DHXbTKeuYpDQ8qhmJZvAyr7Vk,3939
|
|
44
41
|
sqlsaber/memory/__init__.py,sha256=GiWkU6f6YYVV0EvvXDmFWe_CxarmDCql05t70MkTEWs,63
|
|
45
42
|
sqlsaber/memory/manager.py,sha256=p3fybMVfH-E4ApT1ZRZUnQIWSk9dkfUPCyfkmA0HALs,2739
|
|
46
43
|
sqlsaber/memory/storage.py,sha256=ne8szLlGj5NELheqLnI7zu21V8YS4rtpYGGC7tOmi-s,5745
|
|
47
44
|
sqlsaber/theme/__init__.py,sha256=qCICX1Cg4B6yCbZ1UrerxglWxcqldRFVSRrSs73na_8,188
|
|
48
|
-
sqlsaber/theme/manager.py,sha256=
|
|
45
|
+
sqlsaber/theme/manager.py,sha256=TPourIKGU-UzHtImgexgtazpuDaFhqUYtVauMblgGAQ,6480
|
|
49
46
|
sqlsaber/threads/__init__.py,sha256=Hh3dIG1tuC8fXprREUpslCIgPYz8_6o7aRLx4yNeO48,139
|
|
50
47
|
sqlsaber/threads/storage.py,sha256=rsUdxT4CR52D7xtGir9UlsFnBMk11jZeflzDrk2q4ME,11183
|
|
51
|
-
sqlsaber/tools/__init__.py,sha256=
|
|
52
|
-
sqlsaber/tools/base.py,sha256=
|
|
53
|
-
sqlsaber/tools/
|
|
54
|
-
sqlsaber/tools/instructions.py,sha256=X-x8maVkkyi16b6Tl0hcAFgjiYceZaSwyWTfmrvx8U8,9024
|
|
55
|
-
sqlsaber/tools/registry.py,sha256=HWOQMsNIdL4XZS6TeNUyrL-5KoSDH6PHsWd3X66o-18,3211
|
|
48
|
+
sqlsaber/tools/__init__.py,sha256=O6eqkMk8mkhYDniQD1eYgAElOjiHz03I2bGARdgkDkk,421
|
|
49
|
+
sqlsaber/tools/base.py,sha256=NKEEooliPKTJj_Pomwte_wW0Xd9Z5kXNfVdCRfTppuw,883
|
|
50
|
+
sqlsaber/tools/registry.py,sha256=XmBzERq0LJXtg3BZ-r8cEyt8J54NUekgUlTJ_EdSYMk,2204
|
|
56
51
|
sqlsaber/tools/sql_guard.py,sha256=dTDwcZP-N4xPGzcr7MQtKUxKrlDzlc1irr9aH5a4wvk,6182
|
|
57
|
-
sqlsaber/tools/sql_tools.py,sha256=
|
|
58
|
-
sqlsaber-0.
|
|
59
|
-
sqlsaber-0.
|
|
60
|
-
sqlsaber-0.
|
|
61
|
-
sqlsaber-0.
|
|
62
|
-
sqlsaber-0.
|
|
52
|
+
sqlsaber/tools/sql_tools.py,sha256=eo-NTxiXGHMopAjujvDDjmv9hf5bQNbiy3nTpxoJ_E8,7369
|
|
53
|
+
sqlsaber-0.30.0.dist-info/METADATA,sha256=8fYZ1qVQDu-oNhxPSTeLDwewRz1OMy3BK5iTF80nHaQ,5823
|
|
54
|
+
sqlsaber-0.30.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
|
55
|
+
sqlsaber-0.30.0.dist-info/entry_points.txt,sha256=tw1mB0fjlkXQiOsC0434X6nE-o1cFCuQwt2ZYHv_WAE,91
|
|
56
|
+
sqlsaber-0.30.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
|
|
57
|
+
sqlsaber-0.30.0.dist-info/RECORD,,
|
sqlsaber/agents/mcp.py
DELETED
|
@@ -1,21 +0,0 @@
|
|
|
1
|
-
"""Generic SQL agent implementation for MCP tools."""
|
|
2
|
-
|
|
3
|
-
from typing import AsyncIterator
|
|
4
|
-
|
|
5
|
-
from sqlsaber.agents.base import BaseSQLAgent
|
|
6
|
-
from sqlsaber.database import BaseDatabaseConnection
|
|
7
|
-
|
|
8
|
-
|
|
9
|
-
class MCPSQLAgent(BaseSQLAgent):
|
|
10
|
-
"""MCP SQL Agent for MCP tool operations without LLM-specific logic."""
|
|
11
|
-
|
|
12
|
-
def __init__(self, db_connection: BaseDatabaseConnection):
|
|
13
|
-
super().__init__(db_connection)
|
|
14
|
-
|
|
15
|
-
async def query_stream(
|
|
16
|
-
self, user_query: str, use_history: bool = True
|
|
17
|
-
) -> AsyncIterator:
|
|
18
|
-
"""Not implemented for generic agent as it's only used for tool operations."""
|
|
19
|
-
raise NotImplementedError(
|
|
20
|
-
"MCPSQLAgent does not support query streaming. Use specific agent implementations for conversation."
|
|
21
|
-
)
|
sqlsaber/mcp/__init__.py
DELETED
sqlsaber/mcp/mcp.py
DELETED
|
@@ -1,129 +0,0 @@
|
|
|
1
|
-
"""FastMCP server implementation for SQLSaber."""
|
|
2
|
-
|
|
3
|
-
import json
|
|
4
|
-
|
|
5
|
-
from fastmcp import FastMCP
|
|
6
|
-
|
|
7
|
-
from sqlsaber.agents.mcp import MCPSQLAgent
|
|
8
|
-
from sqlsaber.config.database import DatabaseConfigManager
|
|
9
|
-
from sqlsaber.database import DatabaseConnection
|
|
10
|
-
from sqlsaber.tools import SQLTool, tool_registry
|
|
11
|
-
from sqlsaber.tools.instructions import InstructionBuilder
|
|
12
|
-
|
|
13
|
-
# Initialize the instruction builder
|
|
14
|
-
instruction_builder = InstructionBuilder(tool_registry)
|
|
15
|
-
|
|
16
|
-
# Generate dynamic instructions
|
|
17
|
-
DYNAMIC_INSTRUCTIONS = instruction_builder.build_mcp_instructions()
|
|
18
|
-
|
|
19
|
-
# Create the FastMCP server instance with dynamic instructions
|
|
20
|
-
mcp = FastMCP(name="SQL Assistant", instructions=DYNAMIC_INSTRUCTIONS)
|
|
21
|
-
|
|
22
|
-
# Initialize the database config manager
|
|
23
|
-
config_manager = DatabaseConfigManager()
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
async def _create_agent_for_database(database_name: str) -> MCPSQLAgent | None:
|
|
27
|
-
"""Create a MCPSQLAgent for the specified database."""
|
|
28
|
-
try:
|
|
29
|
-
# Look up configured database connection
|
|
30
|
-
db_config = config_manager.get_database(database_name)
|
|
31
|
-
if not db_config:
|
|
32
|
-
return None
|
|
33
|
-
connection_string = db_config.to_connection_string()
|
|
34
|
-
|
|
35
|
-
# Create database connection
|
|
36
|
-
db_conn = DatabaseConnection(connection_string)
|
|
37
|
-
|
|
38
|
-
# Create and return the agent
|
|
39
|
-
agent = MCPSQLAgent(db_conn)
|
|
40
|
-
return agent
|
|
41
|
-
|
|
42
|
-
except Exception:
|
|
43
|
-
return None
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
@mcp.tool
|
|
47
|
-
def get_databases() -> dict:
|
|
48
|
-
"""List all configured databases with their types."""
|
|
49
|
-
databases = []
|
|
50
|
-
for db_config in config_manager.list_databases():
|
|
51
|
-
databases.append(
|
|
52
|
-
{
|
|
53
|
-
"name": db_config.name,
|
|
54
|
-
"type": db_config.type,
|
|
55
|
-
"database": db_config.database,
|
|
56
|
-
"host": db_config.host,
|
|
57
|
-
"port": db_config.port,
|
|
58
|
-
"is_default": db_config.name == config_manager.get_default_name(),
|
|
59
|
-
}
|
|
60
|
-
)
|
|
61
|
-
|
|
62
|
-
return {"databases": databases, "count": len(databases)}
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
async def _execute_with_connection(tool_name: str, database: str, **kwargs) -> str:
|
|
66
|
-
"""Execute a SQL tool with database connection management.
|
|
67
|
-
|
|
68
|
-
Args:
|
|
69
|
-
tool_name: Name of the tool to execute
|
|
70
|
-
database: Database name to connect to
|
|
71
|
-
**kwargs: Tool-specific parameters
|
|
72
|
-
|
|
73
|
-
Returns:
|
|
74
|
-
JSON string with the tool's output
|
|
75
|
-
"""
|
|
76
|
-
try:
|
|
77
|
-
agent = await _create_agent_for_database(database)
|
|
78
|
-
if not agent:
|
|
79
|
-
return json.dumps(
|
|
80
|
-
{"error": f"Database '{database}' not found or could not connect"}
|
|
81
|
-
)
|
|
82
|
-
|
|
83
|
-
# Get the tool and set up connection
|
|
84
|
-
tool = tool_registry.get_tool(tool_name)
|
|
85
|
-
if isinstance(tool, SQLTool):
|
|
86
|
-
tool.set_connection(agent.db)
|
|
87
|
-
|
|
88
|
-
# Execute the tool
|
|
89
|
-
result = await tool.execute(**kwargs)
|
|
90
|
-
await agent.db.close()
|
|
91
|
-
return result
|
|
92
|
-
|
|
93
|
-
except Exception as e:
|
|
94
|
-
return json.dumps({"error": f"Error in {tool_name}: {str(e)}"})
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
# SQL Tool Wrappers with explicit signatures
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
@mcp.tool
|
|
101
|
-
async def list_tables(database: str) -> str:
|
|
102
|
-
"""Get a list of all tables in the database with row counts. Use this first to discover available tables."""
|
|
103
|
-
return await _execute_with_connection("list_tables", database)
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
@mcp.tool
|
|
107
|
-
async def introspect_schema(database: str, table_pattern: str = None) -> str:
|
|
108
|
-
"""Introspect database schema to understand table structures."""
|
|
109
|
-
kwargs = {}
|
|
110
|
-
if table_pattern is not None:
|
|
111
|
-
kwargs["table_pattern"] = table_pattern
|
|
112
|
-
return await _execute_with_connection("introspect_schema", database, **kwargs)
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
@mcp.tool
|
|
116
|
-
async def execute_sql(database: str, query: str, limit: int = 100) -> str:
|
|
117
|
-
"""Execute a SQL query against the database."""
|
|
118
|
-
return await _execute_with_connection(
|
|
119
|
-
"execute_sql", database, query=query, limit=limit
|
|
120
|
-
)
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
def main():
|
|
124
|
-
"""Entry point for the MCP server console script."""
|
|
125
|
-
mcp.run()
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
if __name__ == "__main__":
|
|
129
|
-
main()
|
sqlsaber/tools/enums.py
DELETED
|
@@ -1,19 +0,0 @@
|
|
|
1
|
-
"""Enums for tool categories and workflow positions."""
|
|
2
|
-
|
|
3
|
-
from enum import Enum
|
|
4
|
-
|
|
5
|
-
|
|
6
|
-
class ToolCategory(Enum):
|
|
7
|
-
"""Tool categories for organizing and filtering tools."""
|
|
8
|
-
|
|
9
|
-
GENERAL = "general"
|
|
10
|
-
SQL = "sql"
|
|
11
|
-
|
|
12
|
-
|
|
13
|
-
class WorkflowPosition(Enum):
|
|
14
|
-
"""Workflow positions for organizing tools by usage order."""
|
|
15
|
-
|
|
16
|
-
DISCOVERY = "discovery"
|
|
17
|
-
ANALYSIS = "analysis"
|
|
18
|
-
EXECUTION = "execution"
|
|
19
|
-
OTHER = "other"
|
sqlsaber/tools/instructions.py
DELETED
|
@@ -1,231 +0,0 @@
|
|
|
1
|
-
"""Dynamic instruction builder for tools."""
|
|
2
|
-
|
|
3
|
-
from .base import Tool
|
|
4
|
-
from .enums import ToolCategory, WorkflowPosition
|
|
5
|
-
from .registry import ToolRegistry
|
|
6
|
-
|
|
7
|
-
|
|
8
|
-
class InstructionBuilder:
|
|
9
|
-
"""Builds dynamic instructions based on available tools."""
|
|
10
|
-
|
|
11
|
-
def __init__(self, tool_registry: ToolRegistry):
|
|
12
|
-
"""Initialize with a tool registry."""
|
|
13
|
-
self.registry = tool_registry
|
|
14
|
-
|
|
15
|
-
def build_instructions(
|
|
16
|
-
self,
|
|
17
|
-
db_type: str = "database",
|
|
18
|
-
category: str | ToolCategory | None = None,
|
|
19
|
-
include_base_instructions: bool = True,
|
|
20
|
-
) -> str:
|
|
21
|
-
"""Build dynamic instructions from available tools.
|
|
22
|
-
|
|
23
|
-
Args:
|
|
24
|
-
db_type: Type of database (PostgreSQL, MySQL, SQLite, etc.)
|
|
25
|
-
category: Optional category to filter tools by (string or ToolCategory enum)
|
|
26
|
-
include_base_instructions: Whether to include base SQL assistant instructions
|
|
27
|
-
|
|
28
|
-
Returns:
|
|
29
|
-
Complete instruction string for LLM
|
|
30
|
-
"""
|
|
31
|
-
# Get available tools
|
|
32
|
-
tools = self.registry.get_all_tools(category)
|
|
33
|
-
|
|
34
|
-
if not tools:
|
|
35
|
-
return self._get_base_instructions(db_type)
|
|
36
|
-
|
|
37
|
-
# Sort tools by priority and workflow position
|
|
38
|
-
sorted_tools = self._sort_tools_by_workflow(tools)
|
|
39
|
-
|
|
40
|
-
# Build instruction components
|
|
41
|
-
instructions_parts = []
|
|
42
|
-
|
|
43
|
-
if include_base_instructions:
|
|
44
|
-
instructions_parts.append(self._get_base_instructions(db_type))
|
|
45
|
-
|
|
46
|
-
# Add tool-specific workflow guidance
|
|
47
|
-
workflow_instructions = self._build_workflow_instructions(sorted_tools)
|
|
48
|
-
if workflow_instructions:
|
|
49
|
-
instructions_parts.append(workflow_instructions)
|
|
50
|
-
|
|
51
|
-
# Add tool descriptions and guidelines
|
|
52
|
-
tool_guidelines = self._build_tool_guidelines(sorted_tools)
|
|
53
|
-
if tool_guidelines:
|
|
54
|
-
instructions_parts.append(tool_guidelines)
|
|
55
|
-
|
|
56
|
-
# Add general guidelines
|
|
57
|
-
general_guidelines = self._build_general_guidelines(sorted_tools)
|
|
58
|
-
if general_guidelines:
|
|
59
|
-
instructions_parts.append(general_guidelines)
|
|
60
|
-
|
|
61
|
-
return "\n\n".join(instructions_parts)
|
|
62
|
-
|
|
63
|
-
def _get_base_instructions(self, db_type: str) -> str:
|
|
64
|
-
"""Get base SQL assistant instructions."""
|
|
65
|
-
return f"""You are also a helpful SQL assistant that helps users query their {db_type} database.
|
|
66
|
-
|
|
67
|
-
Your responsibilities:
|
|
68
|
-
1. Understand user's natural language requests, think and convert them to SQL
|
|
69
|
-
2. Use the provided tools efficiently to explore database schema
|
|
70
|
-
3. Generate appropriate SQL queries
|
|
71
|
-
4. Execute queries safely - queries that modify the database are not allowed
|
|
72
|
-
5. Format and explain results clearly"""
|
|
73
|
-
|
|
74
|
-
def _sort_tools_by_workflow(self, tools: list[Tool]) -> list[Tool]:
|
|
75
|
-
"""Sort tools by priority and workflow position."""
|
|
76
|
-
# Define workflow position ordering
|
|
77
|
-
position_order = {
|
|
78
|
-
WorkflowPosition.DISCOVERY: 1,
|
|
79
|
-
WorkflowPosition.ANALYSIS: 2,
|
|
80
|
-
WorkflowPosition.EXECUTION: 3,
|
|
81
|
-
WorkflowPosition.OTHER: 4,
|
|
82
|
-
}
|
|
83
|
-
|
|
84
|
-
return sorted(
|
|
85
|
-
tools,
|
|
86
|
-
key=lambda tool: (
|
|
87
|
-
position_order.get(tool.get_workflow_position(), 4),
|
|
88
|
-
tool.get_priority(),
|
|
89
|
-
tool.name,
|
|
90
|
-
),
|
|
91
|
-
)
|
|
92
|
-
|
|
93
|
-
def _build_workflow_instructions(self, sorted_tools: list[Tool]) -> str:
|
|
94
|
-
"""Build workflow-based instructions."""
|
|
95
|
-
# Group tools by workflow position
|
|
96
|
-
workflow_groups = {}
|
|
97
|
-
for tool in sorted_tools:
|
|
98
|
-
position = tool.get_workflow_position()
|
|
99
|
-
if position not in workflow_groups:
|
|
100
|
-
workflow_groups[position] = []
|
|
101
|
-
workflow_groups[position].append(tool)
|
|
102
|
-
|
|
103
|
-
# Build workflow instructions
|
|
104
|
-
instructions = ["IMPORTANT - Tool Usage Strategy:"]
|
|
105
|
-
step = 1
|
|
106
|
-
|
|
107
|
-
# Add discovery tools first
|
|
108
|
-
if WorkflowPosition.DISCOVERY in workflow_groups:
|
|
109
|
-
discovery_tools = workflow_groups[WorkflowPosition.DISCOVERY]
|
|
110
|
-
for tool in discovery_tools:
|
|
111
|
-
usage = tool.get_usage_instructions()
|
|
112
|
-
if usage:
|
|
113
|
-
instructions.append(f"{step}. {usage}")
|
|
114
|
-
else:
|
|
115
|
-
instructions.append(
|
|
116
|
-
f"{step}. Use '{tool.name}' to {tool.description.lower()}"
|
|
117
|
-
)
|
|
118
|
-
step += 1
|
|
119
|
-
|
|
120
|
-
# Add analysis tools
|
|
121
|
-
if WorkflowPosition.ANALYSIS in workflow_groups:
|
|
122
|
-
analysis_tools = workflow_groups[WorkflowPosition.ANALYSIS]
|
|
123
|
-
for tool in analysis_tools:
|
|
124
|
-
usage = tool.get_usage_instructions()
|
|
125
|
-
if usage:
|
|
126
|
-
instructions.append(f"{step}. {usage}")
|
|
127
|
-
else:
|
|
128
|
-
instructions.append(
|
|
129
|
-
f"{step}. Use '{tool.name}' to {tool.description.lower()}"
|
|
130
|
-
)
|
|
131
|
-
step += 1
|
|
132
|
-
|
|
133
|
-
# Add execution tools
|
|
134
|
-
if WorkflowPosition.EXECUTION in workflow_groups:
|
|
135
|
-
execution_tools = workflow_groups[WorkflowPosition.EXECUTION]
|
|
136
|
-
for tool in execution_tools:
|
|
137
|
-
usage = tool.get_usage_instructions()
|
|
138
|
-
if usage:
|
|
139
|
-
instructions.append(f"{step}. {usage}")
|
|
140
|
-
else:
|
|
141
|
-
instructions.append(
|
|
142
|
-
f"{step}. Use '{tool.name}' to {tool.description.lower()}"
|
|
143
|
-
)
|
|
144
|
-
step += 1
|
|
145
|
-
|
|
146
|
-
return "\n".join(instructions) if len(instructions) > 1 else ""
|
|
147
|
-
|
|
148
|
-
def _build_tool_guidelines(self, sorted_tools: list[Tool]) -> str:
|
|
149
|
-
"""Build tool-specific guidelines."""
|
|
150
|
-
guidelines = []
|
|
151
|
-
|
|
152
|
-
for tool in sorted_tools:
|
|
153
|
-
usage = tool.get_usage_instructions()
|
|
154
|
-
if usage and not self._is_usage_in_workflow(usage):
|
|
155
|
-
guidelines.append(f"- {tool.name}: {usage}")
|
|
156
|
-
|
|
157
|
-
if guidelines:
|
|
158
|
-
return "Tool-Specific Guidelines:\n" + "\n".join(guidelines)
|
|
159
|
-
return ""
|
|
160
|
-
|
|
161
|
-
def _build_general_guidelines(self, sorted_tools: list[Tool]) -> str:
|
|
162
|
-
"""Build general usage guidelines."""
|
|
163
|
-
guidelines = [
|
|
164
|
-
"Guidelines:",
|
|
165
|
-
"- Use proper JOIN syntax and avoid cartesian products",
|
|
166
|
-
"- Include appropriate WHERE clauses to limit results",
|
|
167
|
-
"- Explain what the query does in simple terms",
|
|
168
|
-
"- Handle errors gracefully and suggest fixes",
|
|
169
|
-
"- Be security conscious - use parameterized queries when needed",
|
|
170
|
-
]
|
|
171
|
-
|
|
172
|
-
# Add category-specific guidelines
|
|
173
|
-
categories = {tool.category for tool in sorted_tools}
|
|
174
|
-
|
|
175
|
-
if ToolCategory.SQL in categories:
|
|
176
|
-
guidelines.extend(
|
|
177
|
-
[
|
|
178
|
-
"- Timestamp columns must be converted to text when you write queries",
|
|
179
|
-
"- Use table patterns like 'sample%' or '%experiment%' to filter related tables",
|
|
180
|
-
]
|
|
181
|
-
)
|
|
182
|
-
|
|
183
|
-
return "\n".join(guidelines)
|
|
184
|
-
|
|
185
|
-
def _is_usage_in_workflow(self, usage: str) -> bool:
|
|
186
|
-
"""Check if usage instruction is already covered in workflow section."""
|
|
187
|
-
# Simple heuristic - if usage starts with workflow words, it's probably in workflow
|
|
188
|
-
workflow_words = ["always start", "first", "use this", "begin with", "start by"]
|
|
189
|
-
usage_lower = usage.lower()
|
|
190
|
-
return any(word in usage_lower for word in workflow_words)
|
|
191
|
-
|
|
192
|
-
def build_mcp_instructions(self) -> str:
|
|
193
|
-
"""Build instructions specifically for MCP server."""
|
|
194
|
-
instructions = [
|
|
195
|
-
"This server provides helpful resources and tools that will help you address users queries on their database.",
|
|
196
|
-
"",
|
|
197
|
-
]
|
|
198
|
-
|
|
199
|
-
# Add database discovery
|
|
200
|
-
instructions.append("- Get all databases using `get_databases()`")
|
|
201
|
-
|
|
202
|
-
# Add tool-specific instructions
|
|
203
|
-
sql_tools = self.registry.get_all_tools(category=ToolCategory.SQL)
|
|
204
|
-
sorted_tools = self._sort_tools_by_workflow(sql_tools)
|
|
205
|
-
|
|
206
|
-
for tool in sorted_tools:
|
|
207
|
-
instructions.append(f"- Call `{tool.name}()` to {tool.description.lower()}")
|
|
208
|
-
|
|
209
|
-
# Add workflow guidelines
|
|
210
|
-
instructions.extend(["", "Guidelines:"])
|
|
211
|
-
|
|
212
|
-
workflow_instructions = self._build_workflow_instructions(sorted_tools)
|
|
213
|
-
if workflow_instructions:
|
|
214
|
-
# Extract just the numbered steps without the "IMPORTANT" header
|
|
215
|
-
lines = workflow_instructions.split("\n")[1:] # Skip header
|
|
216
|
-
for line in lines:
|
|
217
|
-
if line.strip():
|
|
218
|
-
# Convert numbered steps to bullet points
|
|
219
|
-
if line.strip()[0].isdigit():
|
|
220
|
-
instructions.append(f"- {line.strip()[3:]}") # Remove "X. "
|
|
221
|
-
|
|
222
|
-
# Add general guidelines
|
|
223
|
-
instructions.extend(
|
|
224
|
-
[
|
|
225
|
-
"- Use proper JOIN syntax and avoid cartesian products",
|
|
226
|
-
"- Include appropriate WHERE clauses to limit results",
|
|
227
|
-
"- Handle errors gracefully and suggest fixes",
|
|
228
|
-
]
|
|
229
|
-
)
|
|
230
|
-
|
|
231
|
-
return "\n".join(instructions)
|
|
File without changes
|
|
File without changes
|