dockerbrain 1.1.0__tar.gz → 1.2.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.
Files changed (43) hide show
  1. {dockerbrain-1.1.0/dockerbrain.egg-info → dockerbrain-1.2.0}/PKG-INFO +3 -2
  2. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/PYPI_README.md +1 -1
  3. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/README.md +12 -3
  4. dockerbrain-1.2.0/core/__init__.py +1 -0
  5. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/cli.py +120 -36
  6. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/llm.py +53 -6
  7. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/monitor/display.py +629 -501
  8. dockerbrain-1.2.0/core/monitor/log_analyzer.py +107 -0
  9. {dockerbrain-1.1.0 → dockerbrain-1.2.0/dockerbrain.egg-info}/PKG-INFO +3 -2
  10. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/dockerbrain.egg-info/SOURCES.txt +1 -0
  11. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/dockerbrain.egg-info/requires.txt +1 -0
  12. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/pyproject.toml +2 -1
  13. dockerbrain-1.2.0/tests/test_llm.py +181 -0
  14. dockerbrain-1.1.0/core/__init__.py +0 -1
  15. dockerbrain-1.1.0/tests/test_llm.py +0 -164
  16. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/LICENSE +0 -0
  17. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/__main__.py +0 -0
  18. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/ai_advisor.py +0 -0
  19. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/dockerizer.py +0 -0
  20. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/fixer/__init__.py +0 -0
  21. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/fixer/container.py +0 -0
  22. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/fixer/dockerfile.py +0 -0
  23. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/monitor/__init__.py +0 -0
  24. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/monitor/collector.py +0 -0
  25. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/monitor/snapshot.py +0 -0
  26. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/optimizer/__init__.py +0 -0
  27. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/optimizer/engine.py +0 -0
  28. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/optimizer/rules.py +0 -0
  29. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/storage.py +0 -0
  30. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/templates.py +0 -0
  31. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/core/utils.py +0 -0
  32. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/dockerbrain.egg-info/dependency_links.txt +0 -0
  33. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/dockerbrain.egg-info/entry_points.txt +0 -0
  34. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/dockerbrain.egg-info/top_level.txt +0 -0
  35. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/setup.cfg +0 -0
  36. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/tests/test_ai_advisor.py +0 -0
  37. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/tests/test_dockerizer.py +0 -0
  38. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/tests/test_fixer.py +0 -0
  39. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/tests/test_monitor.py +0 -0
  40. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/tests/test_optimizer.py +0 -0
  41. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/tests/test_storage.py +0 -0
  42. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/tests/test_templates.py +0 -0
  43. {dockerbrain-1.1.0 → dockerbrain-1.2.0}/tests/test_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dockerbrain
3
- Version: 1.1.0
3
+ Version: 1.2.0
4
4
  Summary: Dockerfile and container monitoring, analysis, and optimization CLI.
5
5
  License: Apache-2.0
6
6
  Project-URL: Homepage, https://github.com/iamPulakesh/DockerBrain
@@ -24,6 +24,7 @@ Classifier: Topic :: System :: Systems Administration
24
24
  Requires-Python: >=3.10
25
25
  Description-Content-Type: text/markdown
26
26
  License-File: LICENSE
27
+ Requires-Dist: anthropic>=0.30
27
28
  Requires-Dist: click>=8.0
28
29
  Requires-Dist: docker>=6.0
29
30
  Requires-Dist: google-genai>=1.0
@@ -37,7 +38,7 @@ Dynamic: license-file
37
38
 
38
39
  # Project Info
39
40
 
40
- [![pipeline](https://github.com/iamPulakesh/DockerBrain/actions/workflows/dockerbrain.yml/badge.svg)](https://github.com/iamPulakesh/DockerBrain/actions/workflows/dockerbrain.yml) [![github project](https://img.shields.io/badge/github-project-blue.svg?logo=github)](https://github.com/iamPulakesh/DockerBrain)
41
+ [![pipeline](https://github.com/iamPulakesh/DockerBrain/actions/workflows/dockerbrain.yml/badge.svg)](https://github.com/iamPulakesh/DockerBrain/actions/workflows/dockerbrain.yml) [![github project](https://img.shields.io/badge/github-project-blue.svg?logo=github)](https://github.com/iamPulakesh/DockerBrain) [![PyPI Downloads](https://static.pepy.tech/personalized-badge/dockerbrain?period=total&units=INTERNATIONAL_SYSTEM&left_color=GRAY&right_color=RED&left_text=Downloads)](https://pepy.tech/projects/dockerbrain)
41
42
 
42
43
  ## Core components of DockerBrain
43
44
 
@@ -1,6 +1,6 @@
1
1
  # Project Info
2
2
 
3
- [![pipeline](https://github.com/iamPulakesh/DockerBrain/actions/workflows/dockerbrain.yml/badge.svg)](https://github.com/iamPulakesh/DockerBrain/actions/workflows/dockerbrain.yml) [![github project](https://img.shields.io/badge/github-project-blue.svg?logo=github)](https://github.com/iamPulakesh/DockerBrain)
3
+ [![pipeline](https://github.com/iamPulakesh/DockerBrain/actions/workflows/dockerbrain.yml/badge.svg)](https://github.com/iamPulakesh/DockerBrain/actions/workflows/dockerbrain.yml) [![github project](https://img.shields.io/badge/github-project-blue.svg?logo=github)](https://github.com/iamPulakesh/DockerBrain) [![PyPI Downloads](https://static.pepy.tech/personalized-badge/dockerbrain?period=total&units=INTERNATIONAL_SYSTEM&left_color=GRAY&right_color=RED&left_text=Downloads)](https://pepy.tech/projects/dockerbrain)
4
4
 
5
5
  ## Core components of DockerBrain
6
6
 
@@ -1,15 +1,22 @@
1
+ <div align="center">
2
+
1
3
  # DockerBrain
2
4
 
3
5
  **AI-powered Docker container monitoring, optimization, and Dockerfile generation CLI.**
4
6
 
7
+ [![PyPI Downloads](https://static.pepy.tech/personalized-badge/dockerbrain?period=total&units=INTERNATIONAL_SYSTEM&left_color=GRAY&right_color=RED&left_text=PyPI+downloads)](https://pepy.tech/projects/dockerbrain)
5
8
  [![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
6
9
  [![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-green.svg)](LICENSE)
7
10
  [![PyPI](https://img.shields.io/pypi/v/dockerbrain)](https://pypi.org/project/dockerbrain/)
8
11
 
9
- DockerBrain monitors your Docker containers in real time, detects resource issues, and uses LLMs to generate actionable optimizations for running containers and Dockerfiles. Supports **Gemini**, **Groq**, and **Ollama**. All configuration lives in a single `.dockerbrainrc` file.
12
+ <p align="center">
13
+ DockerBrain monitors your Docker containers in real time, detects resource issues, and uses LLMs to generate actionable optimizations for running containers and Dockerfiles.
14
+ </p>
10
15
 
11
16
  ---
12
17
 
18
+ </div>
19
+
13
20
  ## Features
14
21
 
15
22
  | Feature | Command | LLM Required |
@@ -68,14 +75,16 @@ This creates `.dockerbrainrc` in `~/.dockerbrain/`. Open it with `dockerb config
68
75
 
69
76
  ```ini
70
77
  [llm]
71
- provider = "groq"
72
- model = "openai/gpt-oss-120b"
78
+ provider = "choose_provider"
79
+ model = "choose_model"
73
80
  api_key = "your_key_here"
74
81
  ```
75
82
 
76
83
  | Provider | API Key | Link |
77
84
  |---|:---:|---|
78
85
  | Gemini | Required | [aistudio.google.com](https://aistudio.google.com/app/apikey) |
86
+ | ChatGPT | Required | [platform.openai.com](https://platform.openai.com/api-keys) |
87
+ | Claude | Required | [console.anthropic.com](https://console.anthropic.com/settings/keys) |
79
88
  | Groq | Required | [console.groq.com/keys](https://console.groq.com/keys) |
80
89
  | Ollama | Not Required | [ollama.com](https://ollama.com) |
81
90
 
@@ -0,0 +1 @@
1
+ __version__ = "1.2.0"
@@ -2,6 +2,7 @@ import platform
2
2
  import click
3
3
  from core import __version__
4
4
 
5
+
5
6
  def print_version(ctx: click.Context, _param: click.Parameter, value: bool) -> None:
6
7
  """Print version with build metadata and exit."""
7
8
  if not value or ctx.resilient_parsing:
@@ -26,6 +27,7 @@ class _HiddenHelpCommand(click.Command):
26
27
 
27
28
  class _HiddenHelpGroup(click.Group):
28
29
  """A Click Group whose subcommands all use HiddenHelpCommand."""
30
+
29
31
  command_class = _HiddenHelpCommand
30
32
 
31
33
 
@@ -36,7 +38,9 @@ class _HiddenHelpGroup(click.Group):
36
38
  is_eager=True,
37
39
  expose_value=False,
38
40
  hidden=True,
39
- callback=lambda ctx, param, value: (click.echo(ctx.get_help(), color=ctx.color) or ctx.exit()) if value else None,
41
+ callback=lambda ctx, param, value: (
42
+ (click.echo(ctx.get_help(), color=ctx.color) or ctx.exit()) if value else None
43
+ ),
40
44
  )
41
45
  @click.option(
42
46
  "--version",
@@ -49,21 +53,60 @@ class _HiddenHelpGroup(click.Group):
49
53
  def cli() -> None:
50
54
  pass
51
55
 
56
+
52
57
  # Commands
53
58
  @cli.command()
54
- @click.option("--interval", "-i", default=1, show_default=True, help="Polling interval in seconds.")
55
- @click.option("--duration", "-d", default=None, type=int, help="Total seconds to monitor (default: unlimited).")
59
+ @click.option(
60
+ "--interval",
61
+ "-i",
62
+ default=1,
63
+ show_default=True,
64
+ help="Polling interval in seconds.",
65
+ )
66
+ @click.option(
67
+ "--duration",
68
+ "-d",
69
+ default=None,
70
+ type=int,
71
+ help="Total seconds to monitor (default: unlimited).",
72
+ )
56
73
  def monitor(interval: int, duration: int | None) -> None:
57
74
  """Display containers' stats."""
58
75
  from core.monitor import run_monitor
76
+
59
77
  run_monitor(interval=interval, duration=duration)
60
78
 
79
+
61
80
  @cli.command("suggest")
62
- @click.option("--container", "-c", default=None, help="Target a specific container (default: all).")
63
- @click.option("--window", "-w", default=30, show_default=True, help="Minutes of metric history to include.")
64
- @click.option("--dockerfile", "-f", default=None, type=click.Path(exists=True), help="Analyze a Dockerfile instead of containers.")
65
- @click.option("--no-rules", is_flag=True, default=False, help="Skip rule-based suggestions, send raw metrics only.")
66
- def ai_suggest(container: str | None, window: int, dockerfile: str | None, no_rules: bool) -> None:
81
+ @click.option(
82
+ "--container",
83
+ "-c",
84
+ default=None,
85
+ help="Target a specific container (default: all).",
86
+ )
87
+ @click.option(
88
+ "--window",
89
+ "-w",
90
+ default=30,
91
+ show_default=True,
92
+ help="Minutes of metric history to include.",
93
+ )
94
+ @click.option(
95
+ "--dockerfile",
96
+ "-f",
97
+ default=None,
98
+ type=click.Path(exists=True),
99
+ help="Analyze a Dockerfile instead of containers.",
100
+ )
101
+ @click.option(
102
+ "--no-rules",
103
+ is_flag=True,
104
+ default=False,
105
+ help="Skip rule-based suggestions, send raw metrics only.",
106
+ )
107
+ def ai_suggest(
108
+ container: str | None, window: int, dockerfile: str | None, no_rules: bool
109
+ ) -> None:
67
110
  """Get optimization suggestions for Dockerfiles and containers.
68
111
 
69
112
  Two modes:
@@ -86,11 +129,21 @@ def ai_suggest(container: str | None, window: int, dockerfile: str | None, no_ru
86
129
  no_rules=no_rules,
87
130
  )
88
131
 
132
+
89
133
  @cli.command()
90
- @click.option("--dockerfile", "-f", default=None, type=click.Path(exists=True),
91
- help="Path to a Dockerfile to auto-fix.")
92
- @click.option("--container", "-c", default=None,
93
- help="Target a specific container (default: all).")
134
+ @click.option(
135
+ "--dockerfile",
136
+ "-f",
137
+ default=None,
138
+ type=click.Path(exists=True),
139
+ help="Path to a Dockerfile to auto-fix.",
140
+ )
141
+ @click.option(
142
+ "--container",
143
+ "-c",
144
+ default=None,
145
+ help="Target a specific container (default: all).",
146
+ )
94
147
  def fix(dockerfile: str | None, container: str | None) -> None:
95
148
  """Automatically diagnose and fix Docker issues.
96
149
 
@@ -109,9 +162,14 @@ def fix(dockerfile: str | None, container: str | None) -> None:
109
162
 
110
163
  run_fix(dockerfile_path=dockerfile, container_name=container)
111
164
 
165
+
112
166
  @cli.command()
113
- @click.option("--force", is_flag=True, default=False,
114
- help="Overwrite existing Dockerfile and .dockerignore forcefully.")
167
+ @click.option(
168
+ "--force",
169
+ is_flag=True,
170
+ default=False,
171
+ help="Overwrite existing Dockerfile and .dockerignore forcefully.",
172
+ )
115
173
  def dockerize(force: bool) -> None:
116
174
  """Generate a Dockerfile and .dockerignore for your project.
117
175
 
@@ -126,18 +184,22 @@ def dockerize(force: bool) -> None:
126
184
 
127
185
  run_dockerize(project_path=".", force=force)
128
186
 
187
+
129
188
  class _NoHelpOptGroup(click.Group):
130
189
  """Group that hides the --help option from the help output."""
190
+
131
191
  def get_help_option(self, ctx: click.Context) -> click.Option | None:
132
192
  opt = super().get_help_option(ctx)
133
193
  if opt:
134
194
  opt.hidden = True
135
195
  return opt
136
196
 
197
+
137
198
  @cli.group(cls=_NoHelpOptGroup)
138
199
  def template() -> None:
139
200
  """Browse and use curated Dockerfile templates."""
140
201
 
202
+
141
203
  @template.command("list")
142
204
  def template_list() -> None:
143
205
  """Show all available Dockerfile templates."""
@@ -163,9 +225,8 @@ def template_list() -> None:
163
225
 
164
226
  console.print()
165
227
  console.print(table)
166
- console.print(
167
- "\n[dim]Use a template:[/] [cyan]dockerb template use <name>[/]\n"
168
- )
228
+ console.print("\n[dim]Use a template:[/] [cyan]dockerb template use <name>[/]\n")
229
+
169
230
 
170
231
  @template.command("use")
171
232
  @click.argument("name")
@@ -191,15 +252,14 @@ def template_use(name: str, force: bool) -> None:
191
252
 
192
253
  if tpl is None:
193
254
  console.print(f"[red]Unknown template:[/] [bold]{name}[/]")
194
- console.print(
195
- f"[dim]Available: {', '.join(get_template_names())}[/]"
196
- )
255
+ console.print(f"[dim]Available: {', '.join(get_template_names())}[/]")
197
256
  return
198
257
 
199
258
  dockerfile_path = Path("Dockerfile")
200
259
 
201
260
  if dockerfile_path.exists() and not force:
202
261
  from rich.prompt import Confirm
262
+
203
263
  if not Confirm.ask(
204
264
  "[yellow]Dockerfile already exists.[/] Overwrite?",
205
265
  default=False,
@@ -208,9 +268,7 @@ def template_use(name: str, force: bool) -> None:
208
268
  return
209
269
 
210
270
  dockerfile_path.write_text(tpl["dockerfile"], encoding="utf-8")
211
- console.print(
212
- f"[green]Created Dockerfile[/] [dim]({tpl['name']})[/]"
213
- )
271
+ console.print(f"[green]Created Dockerfile[/] [dim]({tpl['name']})[/]")
214
272
 
215
273
  dockerignore_path = Path(".dockerignore")
216
274
  if not dockerignore_path.exists():
@@ -232,7 +290,9 @@ README.md
232
290
  console.print("[green]Created .dockerignore[/]")
233
291
 
234
292
  console.print()
235
- console.print(Syntax(tpl["dockerfile"], "dockerfile", theme="monokai", line_numbers=True))
293
+ console.print(
294
+ Syntax(tpl["dockerfile"], "dockerfile", theme="monokai", line_numbers=True)
295
+ )
236
296
 
237
297
 
238
298
  @cli.command()
@@ -252,18 +312,31 @@ def init() -> None:
252
312
  config_path = config_dir / ".dockerbrainrc"
253
313
 
254
314
  if config_path.exists():
255
- console.print(f"[yellow] {config_path} already exists.[/]")
256
- return
315
+ from rich.prompt import Confirm
316
+ if not Confirm.ask(
317
+ f"[yellow]{config_path} already exists.[/] Overwrite?",
318
+ default=False,
319
+ ):
320
+ console.print("[dim]Skipped. No changes made.[/]")
321
+ return
257
322
 
258
- config_content = '''\
323
+ config_content = """\
259
324
  # DockerBrain Configuration
260
325
 
261
326
  # LLM Provider & Models, To switch provider or model, edit below.
262
327
 
263
- # GEMINI
328
+ # GEMINI
264
329
  # provider = "gemini"
265
330
  # model = "gemini-3.1-flash-lite-preview", "gemini-flash-latest"
266
331
 
332
+ # CHATGPT
333
+ # provider = "chatgpt"
334
+ # model = "gpt-5.4-mini", "gpt-5.4-nano"
335
+
336
+ # CLAUDE
337
+ # provider = "claude"
338
+ # model = "claude-sonnet-4-6", "claude-haiku-4-6"
339
+
267
340
  # GROQ (Fast)
268
341
  # provider = "groq"
269
342
  # model = "openai/gpt-oss-120b", "llama-3.3-70b-versatile"
@@ -292,11 +365,12 @@ check_base_image = true # Flag large base images
292
365
  check_apt_recommends = true # Flag apt-get without --no-install-recommends
293
366
  check_cache_busting = true # Flag COPY . . before dependency install
294
367
  check_dockerignore = true # Warn if .dockerignore is missing
295
- '''
368
+ """
296
369
 
297
370
  config_path.write_text(config_content, encoding="utf-8")
298
371
  console.print(f"[green]Created [bold]{config_path}[/bold][/]")
299
372
 
373
+
300
374
  @cli.command()
301
375
  def config() -> None:
302
376
  """Open the .dockerbrainrc config file in your default editor."""
@@ -310,9 +384,7 @@ def config() -> None:
310
384
  config_path = Path.home() / ".dockerbrain" / ".dockerbrainrc"
311
385
 
312
386
  if not config_path.exists():
313
- console.print(
314
- "[yellow]Config not found.[/] Run [cyan]dockerb init[/] first."
315
- )
387
+ console.print("[yellow]Config not found.[/] Run [cyan]dockerb init[/] first.")
316
388
  return
317
389
 
318
390
  if sys.platform == "win32":
@@ -323,6 +395,7 @@ def config() -> None:
323
395
  editor = __import__("os").environ.get("EDITOR", "nano")
324
396
  subprocess.run([editor, str(config_path)])
325
397
 
398
+
326
399
  @cli.command()
327
400
  def env() -> None:
328
401
  """Check your environment and configs."""
@@ -352,7 +425,7 @@ def env() -> None:
352
425
  client = docker.from_env()
353
426
  info = client.version()
354
427
  _ok("Docker daemon", f"running (v{info.get('Version', '?')})")
355
- except Exception :
428
+ except Exception:
356
429
  _fail("Docker daemon", f"not reachable")
357
430
 
358
431
  try:
@@ -362,10 +435,15 @@ def env() -> None:
362
435
 
363
436
  try:
364
437
  from core.llm import load_llm_config
438
+
365
439
  cfg = load_llm_config()
366
440
  _ok("LLM Provider", cfg.provider)
367
441
  _ok("LLM Model", cfg.model)
368
- masked = cfg.api_key[:4] + "…" + cfg.api_key[-4:] if len(cfg.api_key) > 8 else "set ✓"
442
+ masked = (
443
+ cfg.api_key[:4] + "…" + cfg.api_key[-4:]
444
+ if len(cfg.api_key) > 8
445
+ else "set ✓"
446
+ )
369
447
  _ok("API Key", f"{masked}")
370
448
  except SystemExit:
371
449
  _fail("API Key", "not set in ~/.dockerbrain/.dockerbrainrc")
@@ -375,7 +453,9 @@ def env() -> None:
375
453
  size_kb = db_path.stat().st_size / 1024
376
454
  try:
377
455
  conn = sqlite3.connect(str(db_path))
378
- row_count = conn.execute("SELECT COUNT(*) FROM container_metrics").fetchone()[0]
456
+ row_count = conn.execute(
457
+ "SELECT COUNT(*) FROM container_metrics"
458
+ ).fetchone()[0]
379
459
  conn.close()
380
460
  _ok("SQLite DB", f"{db_path} ({size_kb:.0f} KB, {row_count:,} rows)")
381
461
  except Exception:
@@ -383,7 +463,11 @@ def env() -> None:
383
463
  else:
384
464
  _fail("SQLite DB", f"not found at {db_path} — run: dockerb monitor")
385
465
 
386
- status = "[bold green]All checks passed!" if all_ok else "[bold yellow]Some checks failed"
466
+ status = (
467
+ "[bold green]All checks passed!"
468
+ if all_ok
469
+ else "[bold yellow]Some checks failed"
470
+ )
387
471
  console.print(
388
472
  Panel(
389
473
  lines,
@@ -10,21 +10,30 @@ console = Console()
10
10
 
11
11
  _PROVIDER_DEFAULTS: dict[str, dict[str, str]] = {
12
12
  "gemini": {
13
- "model": "gemini-3.1-flash-lite-preview",
14
- "base_url": "", # uses google-genai SDK, not OpenAI
13
+ "model": "gemini-3.1-flash-lite-preview",
14
+ "base_url": "", # uses google-genai SDK
15
+ },
16
+ "chatgpt": {
17
+ "model": "gpt-5.4-mini",
18
+ "base_url": "https://api.openai.com/v1",
19
+ },
20
+ "claude": {
21
+ "model": "claude-sonnet-4-6",
22
+ "base_url": "", # uses anthropic SDK
15
23
  },
16
24
  "groq": {
17
- "model": "llama-3.3-70b-versatile",
25
+ "model": "llama-3.3-70b-versatile",
18
26
  "base_url": "https://api.groq.com/openai/v1",
19
27
  },
20
28
  "ollama": {
21
- "model": "llama3.1",
29
+ "model": "llama3.1",
22
30
  "base_url": "http://localhost:11434/v1",
23
31
  },
24
32
  }
25
33
 
26
34
  _VALID_PROVIDERS = set(_PROVIDER_DEFAULTS.keys())
27
35
 
36
+
28
37
  @dataclass
29
38
  class LLMConfig:
30
39
  """Resolved LLM configuration."""
@@ -104,11 +113,12 @@ def _show_missing_key_error(provider: str) -> None:
104
113
  )
105
114
  )
106
115
 
116
+
107
117
  def _strip_code_fences(text: str) -> str:
108
118
  """Strip markdown code fences that LLMs sometimes wrap responses in."""
109
119
  text = text.strip()
110
120
  if text.startswith("```"):
111
- text = text[text.index("\n") + 1:] if "\n" in text else text
121
+ text = text[text.index("\n") + 1 :] if "\n" in text else text
112
122
  if text.endswith("```"):
113
123
  text = text[:-3].rstrip()
114
124
  return text
@@ -129,6 +139,8 @@ def generate(
129
139
 
130
140
  if config.provider == "gemini":
131
141
  return _generate_gemini(prompt, system_instruction, config)
142
+ elif config.provider == "claude":
143
+ return _generate_anthropic(prompt, system_instruction, config)
132
144
  else:
133
145
  return _generate_openai_compat(prompt, system_instruction, config)
134
146
 
@@ -144,6 +156,8 @@ def generate_stream(
144
156
 
145
157
  if config.provider == "gemini":
146
158
  yield from _stream_gemini(prompt, system_instruction, config)
159
+ elif config.provider == "claude":
160
+ yield from _stream_anthropic(prompt, system_instruction, config)
147
161
  else:
148
162
  yield from _stream_openai_compat(prompt, system_instruction, config)
149
163
 
@@ -163,6 +177,7 @@ def _generate_gemini(prompt: str, system_instruction: str, config: LLMConfig) ->
163
177
  )
164
178
  return _strip_code_fences(response.text)
165
179
 
180
+
166
181
  def _stream_gemini(prompt: str, system_instruction: str, config: LLMConfig):
167
182
  from google import genai
168
183
  from google.genai import types
@@ -180,8 +195,39 @@ def _stream_gemini(prompt: str, system_instruction: str, config: LLMConfig):
180
195
  yield chunk.text
181
196
 
182
197
 
198
+ # Anthropic backend
199
+ def _generate_anthropic(prompt: str, system_instruction: str, config: LLMConfig) -> str:
200
+ import anthropic
201
+
202
+ client = anthropic.Anthropic(api_key=config.api_key)
203
+ response = client.messages.create(
204
+ model=config.model,
205
+ max_tokens=4096,
206
+ system=system_instruction,
207
+ messages=[{"role": "user", "content": prompt}],
208
+ )
209
+ text = response.content[0].text if response.content else ""
210
+ return _strip_code_fences(text)
211
+
212
+
213
+ def _stream_anthropic(prompt: str, system_instruction: str, config: LLMConfig):
214
+ import anthropic
215
+
216
+ client = anthropic.Anthropic(api_key=config.api_key)
217
+ with client.messages.stream(
218
+ model=config.model,
219
+ max_tokens=4096,
220
+ system=system_instruction,
221
+ messages=[{"role": "user", "content": prompt}],
222
+ ) as stream:
223
+ for text in stream.text_stream:
224
+ yield text
225
+
226
+
183
227
  # OpenAI-compatible backend
184
- def _generate_openai_compat(prompt: str, system_instruction: str, config: LLMConfig) -> str:
228
+ def _generate_openai_compat(
229
+ prompt: str, system_instruction: str, config: LLMConfig
230
+ ) -> str:
185
231
  from openai import OpenAI
186
232
 
187
233
  client = OpenAI(api_key=config.api_key, base_url=config.base_url)
@@ -194,6 +240,7 @@ def _generate_openai_compat(prompt: str, system_instruction: str, config: LLMCon
194
240
  )
195
241
  return _strip_code_fences(response.choices[0].message.content or "")
196
242
 
243
+
197
244
  def _stream_openai_compat(prompt: str, system_instruction: str, config: LLMConfig):
198
245
  from openai import OpenAI
199
246