101 lines
3.5 KiB
Python
101 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
from scripts.big_brain_provider import (
|
|
build_generate_payload,
|
|
infer_backend,
|
|
load_big_brain_provider,
|
|
resolve_big_brain_provider,
|
|
resolve_models_url,
|
|
resolve_generate_url,
|
|
)
|
|
|
|
|
|
def test_load_big_brain_provider_from_config(tmp_path: Path) -> None:
|
|
cfg = tmp_path / "config.yaml"
|
|
cfg.write_text(
|
|
yaml.safe_dump(
|
|
{
|
|
"custom_providers": [
|
|
{"name": "Local Ollama", "base_url": "http://localhost:11434/v1", "model": "qwen3:30b"},
|
|
{"name": "Big Brain", "base_url": "https://pod-11434.proxy.runpod.net/v1", "model": "gemma4:latest"},
|
|
]
|
|
}
|
|
)
|
|
)
|
|
|
|
provider = load_big_brain_provider(cfg)
|
|
|
|
assert provider["name"] == "Big Brain"
|
|
assert provider["base_url"] == "https://pod-11434.proxy.runpod.net/v1"
|
|
assert provider["model"] == "gemma4:latest"
|
|
|
|
|
|
def test_infer_backend_distinguishes_openai_compat_from_ollama() -> None:
|
|
assert infer_backend("https://pod-11434.proxy.runpod.net/v1") == "openai"
|
|
assert infer_backend("http://localhost:11434") == "ollama"
|
|
|
|
|
|
def test_resolve_big_brain_provider_prefers_env_overrides(tmp_path: Path, monkeypatch) -> None:
|
|
cfg = tmp_path / "config.yaml"
|
|
cfg.write_text(
|
|
yaml.safe_dump(
|
|
{
|
|
"custom_providers": [
|
|
{"name": "Big Brain", "base_url": "https://old-endpoint/v1", "model": "gemma3:27b"}
|
|
]
|
|
}
|
|
)
|
|
)
|
|
|
|
monkeypatch.setenv("BIG_BRAIN_BASE_URL", "https://vertex-proxy.example/v1")
|
|
monkeypatch.setenv("BIG_BRAIN_MODEL", "gemma4:latest")
|
|
monkeypatch.setenv("BIG_BRAIN_BACKEND", "openai")
|
|
|
|
provider = resolve_big_brain_provider(cfg)
|
|
|
|
assert provider["base_url"] == "https://vertex-proxy.example/v1"
|
|
assert provider["model"] == "gemma4:latest"
|
|
assert provider["backend"] == "openai"
|
|
|
|
|
|
def test_openai_compat_urls_and_payload() -> None:
|
|
provider = {"base_url": "https://pod.proxy.runpod.net/v1", "model": "gemma4:latest", "backend": "openai"}
|
|
|
|
assert resolve_models_url(provider) == "https://pod.proxy.runpod.net/v1/models"
|
|
assert resolve_generate_url(provider) == "https://pod.proxy.runpod.net/v1/chat/completions"
|
|
|
|
payload = build_generate_payload(provider, prompt="Say READY")
|
|
assert payload["model"] == "gemma4:latest"
|
|
assert payload["messages"][0]["content"] == "Say READY"
|
|
assert payload["stream"] is False
|
|
assert payload["max_tokens"] == 32
|
|
|
|
|
|
def test_ollama_urls_and_payload() -> None:
|
|
provider = {"base_url": "http://localhost:11434", "model": "gemma4:latest", "backend": "ollama"}
|
|
|
|
assert resolve_models_url(provider) == "http://localhost:11434/api/tags"
|
|
assert resolve_generate_url(provider) == "http://localhost:11434/api/generate"
|
|
|
|
payload = build_generate_payload(provider, prompt="Say READY")
|
|
assert payload == {"model": "gemma4:latest", "prompt": "Say READY", "stream": False, "options": {"num_predict": 32}}
|
|
|
|
|
|
def test_repo_config_big_brain_is_gemma4_not_hardcoded_dead_pod() -> None:
|
|
config = Path("config.yaml").read_text()
|
|
assert "- name: Big Brain" in config
|
|
assert "model: gemma4:latest" in config
|
|
assert "8lfr3j47a5r3gn-11434.proxy.runpod.net" not in config
|
|
|
|
|
|
def test_big_brain_readme_mentions_runpod_and_vertex() -> None:
|
|
readme = Path("scripts/README_big_brain.md").read_text()
|
|
assert "RunPod" in readme
|
|
assert "Vertex AI" in readme
|
|
assert "gemma4:latest" in readme
|