mirror of
https://github.com/safishamsi/graphify.git
synced 2026-08-26 16:26:42 +00:00
add custom LLM provider registry via providers.json (closes #1084)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
006e1594eb
commit
a9d6be6537
@@ -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 <name>", 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 <name> --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 <name>", 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:])
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user