diff --git a/graphify/__main__.py b/graphify/__main__.py index b4243275a..2eff6ed5f 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -589,7 +589,8 @@ def _run_cli() -> None: print(" --top-k-edges N per-symbol outbound edges in inspector (default 12)") print(" --label NAME project label in header") print(" extract headless full extraction (AST + semantic LLM) for CI/scripts") - print(" --backend B gemini|kimi|claude|openai|deepseek|ollama (default: whichever API key is set)") + print(" --backend B gemini|kimi|claude|openai|deepseek|atlascloud|ollama") + print(" (default: whichever API key is set)") print(" openai also reaches self-hosted OpenAI-compatible servers (llama.cpp,") print(" vLLM, LM Studio): set OPENAI_BASE_URL (e.g. http://localhost:8080/v1)") print(" and OPENAI_MODEL to the model name your server serves") diff --git a/graphify/cli.py b/graphify/cli.py index 5d13a47e9..942199c4f 100644 --- a/graphify/cli.py +++ b/graphify/cli.py @@ -2118,7 +2118,7 @@ def _to_simple(g: "_nx.Graph") -> "_nx.Graph": # has an API key set. if len(sys.argv) < 3: print( - "Usage: graphify extract [--backend gemini|kimi|claude|openai|deepseek|ollama] " + "Usage: graphify extract [--backend gemini|kimi|claude|openai|deepseek|atlascloud|ollama] " "[--model M] [--mode deep] [--out DIR] [--google-workspace] [--no-cluster] " "[--max-workers N] [--token-budget N] [--max-concurrency N] " "[--api-timeout S] [--postgres DSN] [--cargo] [--allow-partial] [--timing]", @@ -2482,7 +2482,8 @@ def _parse_float(name: str, raw: str) -> float: "error: no LLM API key found (" + "; ".join(reasons) + "). " "Set GEMINI_API_KEY or GOOGLE_API_KEY (gemini), MOONSHOT_API_KEY " "(kimi), ANTHROPIC_API_KEY (claude), OPENAI_API_KEY (openai), " - "DEEPSEEK_API_KEY (deepseek), or pass --backend. A code-only " + "DEEPSEEK_API_KEY (deepseek), ATLASCLOUD_API_KEY (atlascloud), " + "or pass --backend. A code-only " "corpus needs no key." + hint, file=sys.stderr, ) diff --git a/graphify/llm.py b/graphify/llm.py index ef094b53f..a885c4043 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -138,6 +138,21 @@ def _get_tokenizer(): "temperature": 0, "max_tokens": 16384, }, + "atlascloud": { + # Atlas Cloud exposes an OpenAI-compatible endpoint. Prefer the Atlas + # API_BASE naming used by their docs, while still accepting BASE_URL for + # consistency with graphify's other OpenAI-compatible backends. + "base_url": os.environ.get( + "ATLASCLOUD_API_BASE", + os.environ.get("ATLASCLOUD_BASE_URL", "https://api.atlascloud.ai/v1"), + ), + "default_model": os.environ.get("ATLASCLOUD_MODEL", "qwen/qwen3.5-flash"), + "env_key": "ATLASCLOUD_API_KEY", + "model_env_key": "GRAPHIFY_ATLASCLOUD_MODEL", + "pricing": {"input": 0.0, "output": 0.0}, + "temperature": 0, + "max_tokens": 16384, + }, "azure": { # Azure OpenAI Service — uses AzureOpenAI SDK client, not the standard # OpenAI client, so it has its own call path (_call_azure). @@ -1559,7 +1574,7 @@ def extract_files_direct( if backend is None: raise ValueError( "No LLM backend configured. Set one of: GEMINI_API_KEY, ANTHROPIC_API_KEY, " - "OPENAI_API_KEY, DEEPSEEK_API_KEY, MOONSHOT_API_KEY, " + "OPENAI_API_KEY, DEEPSEEK_API_KEY, ATLASCLOUD_API_KEY, MOONSHOT_API_KEY, " "AZURE_OPENAI_API_KEY+AZURE_OPENAI_ENDPOINT, OLLAMA_BASE_URL, " "or AWS credentials. Pass backend= explicitly to select a provider." ) @@ -2600,7 +2615,7 @@ def _validate_ollama_base_url(url: str, *, warn: bool = True) -> None: def detect_backend() -> str | None: """Return the name of whichever backend has an API key set, or None. - Priority: gemini → kimi → claude → openai → deepseek → azure → bedrock → ollama (last, opt-in). + Priority: gemini → kimi → claude → openai → deepseek → atlascloud → azure → bedrock → ollama (last, opt-in). Ollama is intentionally checked LAST so a paid API key (Anthropic/OpenAI/etc.) is never silently shadowed by an incidental OLLAMA_BASE_URL in the environment @@ -2608,7 +2623,7 @@ def detect_backend() -> str | None: key now keeps you on the paid backend; remove the paid key (or pass --backend ollama explicitly) to route to the local model. """ - for backend in ("gemini", "kimi", "claude", "openai", "deepseek"): + for backend in ("gemini", "kimi", "claude", "openai", "deepseek", "atlascloud"): if _get_backend_api_key(backend): return backend if _get_backend_api_key("azure") and os.environ.get("AZURE_OPENAI_ENDPOINT"): @@ -2620,7 +2635,10 @@ def detect_backend() -> str | None: _validate_ollama_base_url(ollama_url) return "ollama" for name in BACKENDS: - if name not in ("gemini", "kimi", "claude", "openai", "deepseek", "azure", "bedrock", "ollama", "claude-cli"): + if name not in ( + "gemini", "kimi", "claude", "openai", "deepseek", "atlascloud", + "azure", "bedrock", "ollama", "claude-cli", + ): if _get_backend_api_key(name): return name return None diff --git a/tests/test_extract_cli.py b/tests/test_extract_cli.py index d18bf5ab8..b42012ed7 100644 --- a/tests/test_extract_cli.py +++ b/tests/test_extract_cli.py @@ -465,7 +465,7 @@ def _clear_backend_keys(monkeypatch): """Clear every env var that detect_backend() or _get_backend_api_key() reads.""" for key in ( "GEMINI_API_KEY", "GOOGLE_API_KEY", "OPENAI_API_KEY", - "ANTHROPIC_API_KEY", "DEEPSEEK_API_KEY", "MOONSHOT_API_KEY", + "ANTHROPIC_API_KEY", "DEEPSEEK_API_KEY", "ATLASCLOUD_API_KEY", "MOONSHOT_API_KEY", # bedrock: presence of any of these is treated as a valid credential "AWS_PROFILE", "AWS_REGION", "AWS_DEFAULT_REGION", "AWS_ACCESS_KEY_ID", # ollama: a set OLLAMA_BASE_URL triggers backend detection diff --git a/tests/test_llm_backends.py b/tests/test_llm_backends.py index e6556ba4a..d7d416149 100644 --- a/tests/test_llm_backends.py +++ b/tests/test_llm_backends.py @@ -16,6 +16,7 @@ def _clear_backend_env(monkeypatch): "ANTHROPIC_API_KEY", "OPENAI_API_KEY", "DEEPSEEK_API_KEY", + "ATLASCLOUD_API_KEY", "AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", ): @@ -56,6 +57,14 @@ def test_openai_backend_detected(monkeypatch): assert llm._get_backend_api_key("openai") == "openai-key" +def test_atlascloud_backend_detected(monkeypatch): + _clear_backend_env(monkeypatch) + monkeypatch.setenv("ATLASCLOUD_API_KEY", "atlas-key") + + assert llm.detect_backend() == "atlascloud" + assert llm._get_backend_api_key("atlascloud") == "atlas-key" + + def test_extract_files_direct_routes_gemini_through_openai_compat(tmp_path, monkeypatch): _clear_backend_env(monkeypatch) monkeypatch.setenv("GOOGLE_API_KEY", "google-key") @@ -82,6 +91,26 @@ def test_extract_files_direct_routes_gemini_through_openai_compat(tmp_path, monk assert call.call_args.kwargs["max_completion_tokens"] == 16384 +def test_extract_files_direct_routes_atlascloud_through_openai_compat(tmp_path, monkeypatch): + _clear_backend_env(monkeypatch) + monkeypatch.setenv("ATLASCLOUD_API_KEY", "atlas-key") + source = tmp_path / "note.md" + source.write_text("# Architecture\n\nThe runner emits a snapshot.\n") + result = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 1, "output_tokens": 1} + + with patch("graphify.llm._call_openai_compat", return_value=result) as call: + assert llm.extract_files_direct([source], backend="atlascloud", root=tmp_path) is result + + assert call.call_args.args[:3] == ( + "https://api.atlascloud.ai/v1", + "atlas-key", + "qwen/qwen3.5-flash", + ) + assert call.call_args.kwargs["temperature"] == 0 + assert call.call_args.kwargs["max_completion_tokens"] == 16384 + assert call.call_args.kwargs["backend"] == "atlascloud" + + @pytest.mark.parametrize( "backend, env_key", [ @@ -89,6 +118,7 @@ def test_extract_files_direct_routes_gemini_through_openai_compat(tmp_path, monk ("deepseek", "DEEPSEEK_API_KEY"), ("openai", "OPENAI_API_KEY"), ("kimi", "MOONSHOT_API_KEY"), + ("atlascloud", "ATLASCLOUD_API_KEY"), ], ) def test_openai_compat_backends_resolve_full_output_cap(tmp_path, monkeypatch, backend, env_key): @@ -123,6 +153,20 @@ def test_gemini_model_can_be_overridden_by_env(tmp_path, monkeypatch): assert call.call_args.args[2] == "gemini-3.1-pro-preview" +def test_atlascloud_model_can_be_overridden_by_env(tmp_path, monkeypatch): + _clear_backend_env(monkeypatch) + monkeypatch.setenv("ATLASCLOUD_API_KEY", "atlas-key") + monkeypatch.setenv("GRAPHIFY_ATLASCLOUD_MODEL", "deepseek-ai/deepseek-v4-pro") + source = tmp_path / "note.md" + source.write_text("# Architecture\n") + result = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 1, "output_tokens": 1} + + with patch("graphify.llm._call_openai_compat", return_value=result) as call: + llm.extract_files_direct([source], backend="atlascloud", root=tmp_path) + + assert call.call_args.args[2] == "deepseek-ai/deepseek-v4-pro" + + def test_missing_gemini_key_names_both_supported_env_vars(monkeypatch): _clear_backend_env(monkeypatch) @@ -957,6 +1001,7 @@ def _backend_base_url(backend: str, env_extra: dict) -> str: ("kimi", "KIMI_BASE_URL", "https://proxy.example/kimi/v1"), ("gemini", "GEMINI_BASE_URL", "https://proxy.example/gemini"), ("deepseek", "DEEPSEEK_BASE_URL", "https://proxy.example/deepseek"), + ("atlascloud", "ATLASCLOUD_API_BASE", "https://proxy.example/atlas/v1"), ]) def test_base_url_env_overrides(backend, env_var, override): assert _backend_base_url(backend, {env_var: override}) == override @@ -966,10 +1011,20 @@ def test_base_url_env_overrides(backend, env_var, override): ("kimi", "https://api.moonshot.ai/v1"), ("gemini", "https://generativelanguage.googleapis.com/v1beta/openai/"), ("deepseek", "https://api.deepseek.com"), + ("atlascloud", "https://api.atlascloud.ai/v1"), ]) def test_base_url_defaults_without_env(backend, default): # Ensure the override env vars are unset so the hardcoded default is used. - cleared = {k: "" for k in ("KIMI_BASE_URL", "GEMINI_BASE_URL", "DEEPSEEK_BASE_URL")} + cleared = { + k: "" + for k in ( + "KIMI_BASE_URL", + "GEMINI_BASE_URL", + "DEEPSEEK_BASE_URL", + "ATLASCLOUD_API_BASE", + "ATLASCLOUD_BASE_URL", + ) + } # empty string would be falsy-but-set; delete instead by reconstructing env without them env = {k: v for k, v in os.environ.items() if k not in cleared} out = subprocess.run( diff --git a/tests/test_provider_registry.py b/tests/test_provider_registry.py index 0366c13ff..a6bad68db 100644 --- a/tests/test_provider_registry.py +++ b/tests/test_provider_registry.py @@ -154,7 +154,7 @@ def test_detect_backend_custom_provider_after_builtins(monkeypatch): }) 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"): + "OPENAI_API_KEY", "DEEPSEEK_API_KEY", "ATLASCLOUD_API_KEY", "OLLAMA_BASE_URL"): monkeypatch.delenv(key, raising=False) monkeypatch.delenv("AWS_PROFILE", raising=False) monkeypatch.delenv("AWS_REGION", raising=False)