|
29 | 29 | from rich import console, table, print as pp |
30 | 30 | from click import Choice |
31 | 31 | from dremioai.config import settings |
| 32 | +from enum import StrEnum, auto |
| 33 | +from json import load, dump as jdump |
| 34 | +from shutil import which |
32 | 35 | import asyncio |
33 | | -from yaml import dump |
| 36 | +from yaml import dump, add_representer |
34 | 37 |
|
35 | 38 |
|
36 | 39 | def init( |
@@ -146,23 +149,158 @@ def main( |
146 | 149 | ) |
147 | 150 |
|
148 | 151 |
|
| 152 | +class ConfigTypes(StrEnum): |
| 153 | + dremioai = auto() |
| 154 | + claude = auto() |
| 155 | + |
| 156 | + |
149 | 157 | @tc.command("list", help="Show default configuration, if it exists") |
150 | 158 | def show_default_config( |
151 | 159 | show_filename: Annotated[ |
152 | 160 | bool, Option(help="Show the filename for default config file") |
153 | 161 | ] = False, |
| 162 | + type: Annotated[ |
| 163 | + Optional[ConfigTypes], |
| 164 | + Option(help="The type of configuration to show", show_default=True), |
| 165 | + ] = ConfigTypes.dremioai, |
154 | 166 | ): |
155 | | - dc = settings.default_config() |
156 | | - pp(f"Default config file: {dc!s} (exists = {dc.exists()!s})") |
157 | | - if not show_filename: |
158 | | - settings.configure(dc) |
159 | | - pp( |
160 | | - dump( |
161 | | - settings.instance().model_dump( |
162 | | - exclude_none=True, mode="json", exclude_unset=True |
| 167 | + |
| 168 | + match type: |
| 169 | + case ConfigTypes.dremioai: |
| 170 | + dc = settings.default_config() |
| 171 | + pp(f"Default config file: {dc!s} (exists = {dc.exists()!s})") |
| 172 | + if not show_filename: |
| 173 | + settings.configure(dc) |
| 174 | + pp( |
| 175 | + dump( |
| 176 | + settings.instance().model_dump( |
| 177 | + exclude_none=True, mode="json", exclude_unset=True |
| 178 | + ) |
| 179 | + ) |
163 | 180 | ) |
| 181 | + case ConfigTypes.claude: |
| 182 | + cc = ( |
| 183 | + Path.home() |
| 184 | + / "Library" |
| 185 | + / "Application Support" |
| 186 | + / "Claude" |
| 187 | + / "claude_desktop_config.json" |
164 | 188 | ) |
| 189 | + pp(f"Default config file: '{cc!s}' (exists = {cc.exists()!s})") |
| 190 | + if not show_filename: |
| 191 | + pp(load(cc.open())) |
| 192 | + |
| 193 | + |
| 194 | +cc = Typer( |
| 195 | + context_settings=dict(help_option_names=["-h", "--help"]), |
| 196 | + name="create", |
| 197 | + help="Create DremioAI or LLM configuration files", |
| 198 | +) |
| 199 | +tc.add_typer(cc) |
| 200 | + |
| 201 | + |
| 202 | +@cc.command("claude", help="Create a default configuration file for Claude") |
| 203 | +def create_default_config( |
| 204 | + dry_run: Annotated[ |
| 205 | + bool, Option(help="Dry run, do not overwrite the config file. Just print it") |
| 206 | + ] = False, |
| 207 | +): |
| 208 | + if (uv := which("uv")) is not None: |
| 209 | + uv = Path(uv).resolve() |
| 210 | + dir = str(Path(os.getcwd()).resolve()) |
| 211 | + dmcp = { |
| 212 | + "Dremio": { |
| 213 | + "command": str(uv), |
| 214 | + "args": ["run", "--directory", dir, "dremio-mcp-server", "run"], |
| 215 | + } |
| 216 | + } |
| 217 | + cc = ( |
| 218 | + Path.home() |
| 219 | + / "Library" |
| 220 | + / "Application Support" |
| 221 | + / "Claude" |
| 222 | + / "claude_desktop_config.json" |
165 | 223 | ) |
| 224 | + c = load(cc.open()) if cc.exists() else {"mcpServers": {}} |
| 225 | + c["mcpServers"].update(dmcp) |
| 226 | + if dry_run: |
| 227 | + pp(c) |
| 228 | + else: |
| 229 | + if not cc.exists(): |
| 230 | + cc.parent.mkdir(parents=True, exist_ok=True) |
| 231 | + with cc.open("w") as f: |
| 232 | + jdump(c, f) |
| 233 | + pp(f"Created default config file: {cc!s}") |
| 234 | + else: |
| 235 | + raise FileNotFoundError("uv command not found. Please install uv") |
| 236 | + |
| 237 | + |
| 238 | +@cc.command("dremioai", help="Create a default configuration file") |
| 239 | +def create_default_config( |
| 240 | + uri: Annotated[ |
| 241 | + str, |
| 242 | + Option( |
| 243 | + help=f"The Dremio URL or shorthand for Dremio Cloud regions ({ ','.join(settings.DremioCloudUri)})" |
| 244 | + ), |
| 245 | + ], |
| 246 | + pat: Annotated[ |
| 247 | + str, |
| 248 | + Option( |
| 249 | + help="The Dremio PAT. If it starts with @ then treat the rest is treated as a filename" |
| 250 | + ), |
| 251 | + ], |
| 252 | + project_id: Annotated[ |
| 253 | + Optional[str], |
| 254 | + Option(help="The Dremio project id, only if connecting to Dremio Cloud"), |
| 255 | + ] = None, |
| 256 | + mode: Annotated[ |
| 257 | + Optional[List[str]], |
| 258 | + Option("-m", "--mode", help="MCP server mode", click_type=Choice(_mode())), |
| 259 | + ] = [tools.ToolType.FOR_DATA_PATTERNS.name], |
| 260 | + enable_experimental: Annotated[ |
| 261 | + bool, Option(help="Enable experimental features") |
| 262 | + ] = False, |
| 263 | + dry_run: Annotated[ |
| 264 | + bool, Option(help="Dry run, do not overwrite the config file. Just print it") |
| 265 | + ] = False, |
| 266 | +): |
| 267 | + mode = "|".join([tools.ToolType[m.upper()].name for m in mode]) |
| 268 | + dremio = settings.Dremio.model_validate( |
| 269 | + { |
| 270 | + "uri": uri, |
| 271 | + "pat": pat, |
| 272 | + "project_id": project_id, |
| 273 | + "enable_experimental": enable_experimental, |
| 274 | + } |
| 275 | + ) |
| 276 | + ts = settings.Tools.model_validate({"server_mode": tools.ToolType.FOR_SELF}) |
| 277 | + settings.configure(settings.default_config(), force=True) |
| 278 | + settings.instance().dremio = dremio |
| 279 | + settings.instance().tools = ts |
| 280 | + d = settings.instance().model_dump( |
| 281 | + exclude_none=True, mode="json", exclude_unset=True |
| 282 | + ) |
| 283 | + |
| 284 | + add_representer( |
| 285 | + str, |
| 286 | + lambda dumper, data: dumper.represent_scalar( |
| 287 | + "tag:yaml.org,2002:str", data, style=('"' if "@" in data else None) |
| 288 | + ), |
| 289 | + ) |
| 290 | + if pat.startswith("@") and d["dremio"]["pat"] != pat: |
| 291 | + d["dremio"]["pat"] = pat |
| 292 | + d["tools"]["server_mode"] = mode |
| 293 | + |
| 294 | + if dry_run: |
| 295 | + pp(dump(d)) |
| 296 | + return |
| 297 | + |
| 298 | + dc = settings.default_config() |
| 299 | + if not dc.exists(): |
| 300 | + dc.parent.mkdir(parents=True, exist_ok=True) |
| 301 | + with dc.open("w") as f: |
| 302 | + dump(d, f) |
| 303 | + pp(f"Created default config file: {dc!s}") |
166 | 304 |
|
167 | 305 |
|
168 | 306 | # -------------------------------------------------------------------------------- |
|
0 commit comments