diff --git a/.env.example b/.env.example index 76a6814..e53d2f8 100644 --- a/.env.example +++ b/.env.example @@ -16,6 +16,10 @@ LLM_HOST=localhost # when started with OLLAMA_HOST=0.0.0.0:11434. # OLLAMA_BASE_URL=http://host.docker.internal:11434/v1 +# Optional LM Studio URL. In Docker, host LM Studio is reachable here +# when LM Studio is set to serve on all interfaces (0.0.0.0). +# LM_STUDIO_URL=http://host.docker.internal:1234 + # OpenAI API key (only needed if using OpenAI models). # Do not commit real keys. Keep this commented until needed. # OPENAI_API_KEY=your_openai_api_key_here diff --git a/src/model_discovery.py b/src/model_discovery.py index ab3ef13..f109651 100644 --- a/src/model_discovery.py +++ b/src/model_discovery.py @@ -74,15 +74,33 @@ class ModelDiscovery: self.default_host = default_host self.openai_api_key = openai_api_key self.openai_compat_path = "/v1/chat/completions" + # Custom ports from env vars, merged into the scan list by discover_models. + self._extra_ports: set = set() def _get_hosts(self) -> List[str]: """Get all hosts to scan, using env override, Tailscale, or default.""" + self._extra_ports = set() + def _append_host(out: List[str], host: str) -> None: host = (host or "").strip() if not host or host in out: return out.append(host) + def _append_env_hosts(out: List[str]) -> None: + """Add hosts (and any custom ports) from provider-specific env vars.""" + for env_name in ("OLLAMA_BASE_URL", "OLLAMA_URL", "LM_STUDIO_URL"): + raw = os.getenv(env_name, "").strip() + if not raw: + continue + try: + parsed = urlparse(raw if "://" in raw else "http://" + raw) + _append_host(out, parsed.hostname or "") + if parsed.port: + self._extra_ports.add(parsed.port) + except Exception: + pass + # Manual override takes priority extra = os.getenv("LLM_HOSTS", "").strip() if extra: @@ -91,6 +109,7 @@ class ModelDiscovery: if self.default_host not in hosts: hosts.insert(0, self.default_host) _append_host(hosts, "host.docker.internal") + _append_env_hosts(hosts) return hosts # Try Tailscale discovery @@ -100,23 +119,30 @@ class ModelDiscovery: if self.default_host not in ts_hosts: ts_hosts.insert(0, self.default_host) _append_host(ts_hosts, "host.docker.internal") + _append_env_hosts(ts_hosts) return ts_hosts hosts = [self.default_host] # Docker desktop/Linux compose maps this to the host machine. That is # the common "I started Ollama normally on this computer" case. _append_host(hosts, "host.docker.internal") - for env_name in ("OLLAMA_BASE_URL", "OLLAMA_URL"): - raw = os.getenv(env_name, "").strip() - if not raw: - continue - try: - parsed = urlparse(raw if "://" in raw else "http://" + raw) - _append_host(hosts, parsed.hostname or "") - except Exception: - pass + _append_env_hosts(hosts) return hosts + def _fingerprint_provider(self, host: str, port: int) -> Optional[str]: + """Identify the server software via its native API, independent of port.""" + try: + r = httpx.get(f"http://{host}:{port}/api/v1/models", timeout=1.5) + if r.is_success: + models = (r.json() or {}).get("models") + if (isinstance(models, list) and models + and isinstance(models[0], dict) + and "key" in models[0] and "architecture" in models[0]): + return "lmstudio" + except Exception: + pass + return None + def _check_port(self, host: str, port: int) -> Optional[Dict[str, Any]]: """Check a single host:port for models.""" base = f"http://{host}:{port}/v1" @@ -132,7 +158,8 @@ class ModelDiscovery: "port": port, "url": f"http://{host}:{port}{self.openai_compat_path}", "models": ids, - "models_display": [i.lstrip("/") for i in ids] + "models_display": [i.lstrip("/") for i in ids], + "provider": self._fingerprint_provider(host, port), } except Exception: pass @@ -145,9 +172,10 @@ class ModelDiscovery: logger.info(f"Scanning {len(hosts)} hosts for models: {hosts}") - # Build list of (host, port) to check. 8000-8020 catches vLLM, - # llama.cpp, SGLang, and Cookbook serves; 11434 catches Ollama. - ports = list(range(8000, 8021)) + [11434] + # Well-known ports: 8000-8020 (vLLM, llama.cpp, SGLang, Cookbook), + # 1234 (LM Studio), 11434 (Ollama) + ports = list(range(8000, 8021)) + [1234, 11434] + ports += [p for p in sorted(self._extra_ports) if p not in ports] targets = [(h, p) for h in hosts for p in ports] seen_models = set() # dedupe by (port, model_ids) to avoid same machine via different IPs diff --git a/tests/test_lmstudio_discovery.py b/tests/test_lmstudio_discovery.py new file mode 100644 index 0000000..d12eead --- /dev/null +++ b/tests/test_lmstudio_discovery.py @@ -0,0 +1,184 @@ +"""Tests for LM Studio model discovery: port scanning, env host scanning, +and native-API provider fingerprinting.""" +from src.model_discovery import ModelDiscovery + + +class _FakeResponse: + def __init__(self, payload, ok=True): + self._payload = payload + self.is_success = ok + + def json(self): + return self._payload + + +# ════════════════════════════════════════════════════════════ +# ModelDiscovery — ports list includes 1234 +# ════════════════════════════════════════════════════════════ + +class TestModelDiscoveryPorts: + def test_discover_models_scans_port_1234(self, monkeypatch): + """discover_models must include port 1234 among the scan targets.""" + discovery = ModelDiscovery(default_host="localhost") + scanned_ports = [] + + def fake_check_port(host, port): + scanned_ports.append(port) + return None + + monkeypatch.setattr(discovery, "_check_port", fake_check_port) + monkeypatch.setattr( + "src.model_discovery.discover_tailscale_hosts", + lambda: [], + ) + + discovery.discover_models() + assert 1234 in scanned_ports + + def test_discover_models_scans_custom_lm_studio_port(self, monkeypatch): + """A non-default port in LM_STUDIO_URL must be added to the scan targets.""" + monkeypatch.delenv("LLM_HOSTS", raising=False) + monkeypatch.setenv("LM_STUDIO_URL", "http://my-lm-box:5000") + monkeypatch.setattr( + "src.model_discovery.discover_tailscale_hosts", lambda: [], + ) + discovery = ModelDiscovery(default_host="localhost") + scanned = [] + + def fake_check_port(host, port): + scanned.append((host, port)) + return None + + monkeypatch.setattr(discovery, "_check_port", fake_check_port) + discovery.discover_models() + assert ("my-lm-box", 5000) in scanned + + +# ════════════════════════════════════════════════════════════ +# _fingerprint_provider — native API identification +# ════════════════════════════════════════════════════════════ + +class TestFingerprintProvider: + LMSTUDIO_NATIVE = { + "models": [ + {"type": "llm", "key": "qwen3.6-27b", "architecture": "qwen35", + "quantization": {"name": "Q5_K_XL"}, "format": "gguf"}, + ] + } + + def test_lmstudio_native_format_detected(self, monkeypatch): + discovery = ModelDiscovery(default_host="localhost") + monkeypatch.setattr( + "src.model_discovery.httpx.get", + lambda url, timeout=None: _FakeResponse(self.LMSTUDIO_NATIVE), + ) + assert discovery._fingerprint_provider("localhost", 1234) == "lmstudio" + + def test_lmstudio_detected_on_nonstandard_port(self, monkeypatch): + discovery = ModelDiscovery(default_host="localhost") + monkeypatch.setattr( + "src.model_discovery.httpx.get", + lambda url, timeout=None: _FakeResponse(self.LMSTUDIO_NATIVE), + ) + assert discovery._fingerprint_provider("localhost", 8080) == "lmstudio" + + def test_openai_compatible_server_not_lmstudio(self, monkeypatch): + discovery = ModelDiscovery(default_host="localhost") + monkeypatch.setattr( + "src.model_discovery.httpx.get", + lambda url, timeout=None: _FakeResponse({"data": [{"id": "gpt-4o"}]}, ok=False), + ) + assert discovery._fingerprint_provider("localhost", 8000) is None + + def test_ollama_tags_shape_not_lmstudio(self, monkeypatch): + discovery = ModelDiscovery(default_host="localhost") + ollama_shape = {"models": [{"name": "llama3", "modified_at": "x", "size": 1}]} + monkeypatch.setattr( + "src.model_discovery.httpx.get", + lambda url, timeout=None: _FakeResponse(ollama_shape), + ) + assert discovery._fingerprint_provider("localhost", 11434) is None + + def test_unreachable_returns_none(self, monkeypatch): + discovery = ModelDiscovery(default_host="localhost") + def boom(url, timeout=None): + raise OSError("connection refused") + monkeypatch.setattr("src.model_discovery.httpx.get", boom) + assert discovery._fingerprint_provider("localhost", 1234) is None + + def test_check_port_attaches_provider(self, monkeypatch): + discovery = ModelDiscovery(default_host="localhost") + + def fake_get(url, timeout=None): + if url.endswith("/api/v1/models"): + return _FakeResponse(self.LMSTUDIO_NATIVE) + return _FakeResponse({"data": [{"id": "qwen3.6-27b"}]}) + + monkeypatch.setattr("src.model_discovery.httpx.get", fake_get) + result = discovery._check_port("localhost", 1234) + assert result is not None + assert result["provider"] == "lmstudio" + assert result["models"] == ["qwen3.6-27b"] + + +# ════════════════════════════════════════════════════════════ +# _get_hosts — LM_STUDIO_URL env var +# ════════════════════════════════════════════════════════════ + +class TestGetHostsLmStudioUrl: + def test_lm_studio_url_adds_host_default_branch(self, monkeypatch): + """LM_STUDIO_URL hostname must appear in hosts when Tailscale is absent.""" + monkeypatch.delenv("LLM_HOSTS", raising=False) + monkeypatch.setenv("LM_STUDIO_URL", "http://my-lm-box:1234") + monkeypatch.setattr( + "src.model_discovery.discover_tailscale_hosts", + lambda: [], + ) + discovery = ModelDiscovery(default_host="localhost") + hosts = discovery._get_hosts() + assert "my-lm-box" in hosts + + def test_lm_studio_url_adds_host_tailscale_branch(self, monkeypatch): + """LM_STUDIO_URL hostname must also appear when Tailscale hosts are present.""" + monkeypatch.delenv("LLM_HOSTS", raising=False) + monkeypatch.setenv("LM_STUDIO_URL", "http://my-lm-box:1234") + monkeypatch.setattr( + "src.model_discovery.discover_tailscale_hosts", + lambda: ["100.64.0.1"], + ) + discovery = ModelDiscovery(default_host="localhost") + hosts = discovery._get_hosts() + assert "my-lm-box" in hosts + + def test_lm_studio_url_adds_host_llm_hosts_branch(self, monkeypatch): + """LM_STUDIO_URL hostname must also appear when LLM_HOSTS is set.""" + monkeypatch.setenv("LLM_HOSTS", "10.0.0.5") + monkeypatch.setenv("LM_STUDIO_URL", "http://my-lm-box:1234") + discovery = ModelDiscovery(default_host="localhost") + hosts = discovery._get_hosts() + assert "my-lm-box" in hosts + + def test_lm_studio_url_no_duplicate(self, monkeypatch): + """If the hostname is already in the list it should not be added twice.""" + monkeypatch.delenv("LLM_HOSTS", raising=False) + monkeypatch.setenv("LM_STUDIO_URL", "http://localhost:1234") + monkeypatch.setattr( + "src.model_discovery.discover_tailscale_hosts", + lambda: [], + ) + discovery = ModelDiscovery(default_host="localhost") + hosts = discovery._get_hosts() + assert hosts.count("localhost") == 1 + + def test_lm_studio_url_not_set_no_extra_host(self, monkeypatch): + """When LM_STUDIO_URL is absent, no phantom host is added.""" + monkeypatch.delenv("LLM_HOSTS", raising=False) + monkeypatch.delenv("LM_STUDIO_URL", raising=False) + monkeypatch.setattr( + "src.model_discovery.discover_tailscale_hosts", + lambda: [], + ) + discovery = ModelDiscovery(default_host="localhost") + hosts = discovery._get_hosts() + # Only localhost + host.docker.internal expected + assert "my-lm-box" not in hosts