From a9d6be6537b619b17e6ddc4798e2950b3decfb70 Mon Sep 17 00:00:00 2001 From: Safi Date: Sat, 30 May 2026 13:25:02 +0100 Subject: [PATCH] add custom LLM provider registry via providers.json (closes #1084) Co-Authored-By: Claude Sonnet 4.6 --- graphify/__main__.py | 113 ++++++++++++++++++++++++++++++++ graphify/llm.py | 30 +++++++++ tests/test_provider_registry.py | 95 +++++++++++++++++++++++++++ 3 files changed, 238 insertions(+) create mode 100644 tests/test_provider_registry.py diff --git a/graphify/__main__.py b/graphify/__main__.py index 0eafba73..ef1dc729 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -1778,6 +1778,119 @@ def main() -> None: else: print("Usage: graphify antigravity [install|uninstall]", file=sys.stderr) sys.exit(1) + elif cmd == "provider": + from graphify.llm import _custom_providers_path, BACKENDS + import json as _json + subcmd = sys.argv[2] if len(sys.argv) > 2 else "" + global_path = _custom_providers_path(global_=True) + + if subcmd == "list": + global_path.parent.mkdir(parents=True, exist_ok=True) + existing: dict = {} + if global_path.is_file(): + try: + existing = _json.loads(global_path.read_text(encoding="utf-8")) + except Exception: + pass + if not existing: + print("No custom providers registered.") + else: + for name in existing: + print(f" {name} ({existing[name].get('base_url', '')})") + + elif subcmd == "show": + name = sys.argv[3] if len(sys.argv) > 3 else "" + if not name: + print("Usage: graphify provider show ", file=sys.stderr) + sys.exit(1) + existing = {} + if global_path.is_file(): + try: + existing = _json.loads(global_path.read_text(encoding="utf-8")) + except Exception: + pass + if name not in existing: + print(f"Provider '{name}' not found.", file=sys.stderr) + sys.exit(1) + print(_json.dumps({name: existing[name]}, indent=2)) + + elif subcmd == "add": + args = sys.argv[3:] + name = args[0] if args and not args[0].startswith("-") else "" + if not name: + print("Usage: graphify provider add --base-url URL --default-model MODEL --env-key KEY", file=sys.stderr) + sys.exit(1) + if name in BACKENDS: + print(f"Error: '{name}' is a built-in provider and cannot be overridden.", file=sys.stderr) + sys.exit(1) + base_url = "" + default_model = "" + env_key = "" + pricing_input = 0.0 + pricing_output = 0.0 + i = 1 + while i < len(args): + a = args[i] + if a == "--base-url" and i + 1 < len(args): + base_url = args[i + 1]; i += 2 + elif a.startswith("--base-url="): + base_url = a.split("=", 1)[1]; i += 1 + elif a == "--default-model" and i + 1 < len(args): + default_model = args[i + 1]; i += 2 + elif a.startswith("--default-model="): + default_model = a.split("=", 1)[1]; i += 1 + elif a == "--env-key" and i + 1 < len(args): + env_key = args[i + 1]; i += 2 + elif a.startswith("--env-key="): + env_key = a.split("=", 1)[1]; i += 1 + elif a == "--pricing-input" and i + 1 < len(args): + pricing_input = float(args[i + 1]); i += 2 + elif a == "--pricing-output" and i + 1 < len(args): + pricing_output = float(args[i + 1]); i += 2 + else: + i += 1 + if not base_url or not default_model or not env_key: + print("Error: --base-url, --default-model, and --env-key are required.", file=sys.stderr) + sys.exit(1) + global_path.parent.mkdir(parents=True, exist_ok=True) + existing = {} + if global_path.is_file(): + try: + existing = _json.loads(global_path.read_text(encoding="utf-8")) + except Exception: + pass + existing[name] = { + "base_url": base_url, + "default_model": default_model, + "env_key": env_key, + "pricing": {"input": pricing_input, "output": pricing_output}, + "temperature": 0, + } + global_path.write_text(_json.dumps(existing, indent=2) + "\n", encoding="utf-8") + print(f"Provider '{name}' added. Use with: graphify extract . --backend {name}") + + elif subcmd == "remove": + name = sys.argv[3] if len(sys.argv) > 3 else "" + if not name: + print("Usage: graphify provider remove ", file=sys.stderr) + sys.exit(1) + existing = {} + if global_path.is_file(): + try: + existing = _json.loads(global_path.read_text(encoding="utf-8")) + except Exception: + pass + if name not in existing: + print(f"Provider '{name}' not found.", file=sys.stderr) + sys.exit(1) + del existing[name] + global_path.write_text(_json.dumps(existing, indent=2) + "\n", encoding="utf-8") + print(f"Provider '{name}' removed.") + + else: + print("Usage: graphify provider [add|list|show|remove]", file=sys.stderr) + if subcmd: + sys.exit(1) elif cmd == "prs": from graphify.prs import cmd_prs cmd_prs(sys.argv[2:]) diff --git a/graphify/llm.py b/graphify/llm.py index 8aba58b6..730975ed 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -118,6 +118,32 @@ BACKENDS: dict[str, dict] = { } +def _custom_providers_path(global_: bool = True) -> Path: + if global_: + return Path.home() / ".graphify" / "providers.json" + return Path(".graphify") / "providers.json" + + +def _load_custom_providers() -> dict[str, dict]: + providers: dict[str, dict] = {} + for path in (_custom_providers_path(global_=False), _custom_providers_path(global_=True)): + if path.is_file(): + try: + data = json.loads(path.read_text(encoding="utf-8")) + if isinstance(data, dict): + for name, cfg in data.items(): + if isinstance(name, str) and isinstance(cfg, dict) and name not in BACKENDS: + if "pricing" not in cfg: + cfg = dict(cfg, pricing={"input": 0.0, "output": 0.0}) + providers[name] = cfg + except Exception: + pass + return providers + + +BACKENDS.update(_load_custom_providers()) + + def _resolve_max_tokens(default: int) -> int: """Honour GRAPHIFY_MAX_OUTPUT_TOKENS env var override, else use backend default.""" raw = os.environ.get("GRAPHIFY_MAX_OUTPUT_TOKENS", "").strip() @@ -1221,4 +1247,8 @@ def detect_backend() -> str | None: if ollama_url: _validate_ollama_base_url(ollama_url) return "ollama" + for name in BACKENDS: + if name not in ("gemini", "kimi", "claude", "openai", "deepseek", "bedrock", "ollama", "claude-cli"): + if _get_backend_api_key(name): + return name return None diff --git a/tests/test_provider_registry.py b/tests/test_provider_registry.py new file mode 100644 index 00000000..bbf082ca --- /dev/null +++ b/tests/test_provider_registry.py @@ -0,0 +1,95 @@ +import json +import pytest +from pathlib import Path + + +def test_custom_provider_add_list_show_remove(tmp_path, monkeypatch): + """Full round-trip: add → list → show → remove via providers.json.""" + providers_file = tmp_path / "providers.json" + providers_file.write_text("{}", encoding="utf-8") + + from graphify import llm + monkeypatch.setattr(llm, "_custom_providers_path", lambda global_=True: providers_file if global_ else tmp_path / "local.json") + monkeypatch.setattr(llm, "BACKENDS", {**llm.BACKENDS}) + + providers_file.write_text(json.dumps({ + "nvidia": { + "base_url": "https://integrate.api.nvidia.com/v1", + "default_model": "minimaxai/minimax-m2.7", + "env_key": "NVIDIA_API_KEY", + "pricing": {"input": 0.0, "output": 0.0}, + "temperature": 0, + } + }), encoding="utf-8") + + loaded = llm._load_custom_providers() + assert "nvidia" in loaded + assert loaded["nvidia"]["base_url"] == "https://integrate.api.nvidia.com/v1" + + +def test_custom_provider_pricing_defaults_to_zero(tmp_path): + """Missing pricing field defaults to zero so estimate_cost doesn't blow up.""" + providers_file = tmp_path / "providers.json" + providers_file.write_text(json.dumps({ + "mymodel": { + "base_url": "http://localhost:8080/v1", + "default_model": "llama3", + "env_key": "MY_API_KEY", + } + }), encoding="utf-8") + + from graphify import llm + import importlib + from unittest.mock import patch + + with patch.object(llm, "_custom_providers_path", side_effect=lambda global_=True: providers_file if global_ else tmp_path / "local.json"): + loaded = llm._load_custom_providers() + + assert "mymodel" in loaded + assert loaded["mymodel"]["pricing"] == {"input": 0.0, "output": 0.0} + + +def test_custom_provider_cannot_shadow_builtin(tmp_path): + """Built-in provider names are protected from being overridden.""" + providers_file = tmp_path / "providers.json" + providers_file.write_text(json.dumps({ + "claude": { + "base_url": "http://evil.example.com/v1", + "default_model": "evil-model", + "env_key": "EVIL_KEY", + } + }), encoding="utf-8") + + from graphify import llm + from unittest.mock import patch + + with patch.object(llm, "_custom_providers_path", side_effect=lambda global_=True: providers_file if global_ else tmp_path / "local.json"): + loaded = llm._load_custom_providers() + + assert "claude" not in loaded + + +def test_detect_backend_custom_provider_after_builtins(monkeypatch): + """Custom providers appear after all built-ins in detect_backend() priority.""" + from graphify import llm + + monkeypatch.setattr(llm, "BACKENDS", { + **llm.BACKENDS, + "myprovider": { + "base_url": "http://example.com/v1", + "default_model": "mymodel", + "env_key": "MY_CUSTOM_KEY", + "pricing": {"input": 0.0, "output": 0.0}, + "temperature": 0, + } + }) + monkeypatch.setenv("MY_CUSTOM_KEY", "test-key") + for key in ("GEMINI_API_KEY", "GOOGLE_API_KEY", "MOONSHOT_API_KEY", "ANTHROPIC_API_KEY", + "OPENAI_API_KEY", "DEEPSEEK_API_KEY", "OLLAMA_BASE_URL"): + monkeypatch.delenv(key, raising=False) + monkeypatch.delenv("AWS_PROFILE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False) + + result = llm.detect_backend() + assert result == "myprovider"