fix(context): use Ollama runtime allocation

This commit is contained in:
RaresKeY 2026-07-20 14:09:14 +00:00
parent 77179724ae
commit 8999418d23
2 changed files with 216 additions and 2 deletions

View file

@ -236,6 +236,7 @@ KNOWN_CONTEXT_WINDOWS = {
# Cache
# ---------------------------------------------------------------------------
_context_cache: Dict[Tuple[str, str], Tuple[int, bool]] = {}
_ollama_runtime_endpoints: set[str] = set()
def _get_context_length_cached(endpoint_url: str, model: str) -> Tuple[int, bool]:
@ -243,19 +244,23 @@ def _get_context_length_cached(endpoint_url: str, model: str) -> Tuple[int, bool
bare DEFAULT_CONTEXT fallback (no endpoint report and not in the known table)."""
configured_kind = _configured_endpoint_kind(endpoint_url)
is_local = is_local_endpoint(endpoint_url)
has_dynamic_runtime_context = endpoint_url in _ollama_runtime_endpoints
# Key on (endpoint_url, model): the same model id can be served by two
# different remote endpoints with different real context windows (e.g. a
# capped proxy vs. the full provider), so caching by model id alone would
# serve one endpoint's window for the other (issue #2603).
cache_key = (endpoint_url, model)
if not is_local and cache_key in _context_cache:
if not is_local and not has_dynamic_runtime_context and cache_key in _context_cache:
return _context_cache[cache_key]
ctx, known = _query_context_length(endpoint_url, model)
# Only cache non-default values to allow retry on next request.
# Local endpoints can restart with a different --max-model-len while keeping
# the same model id, so always re-query them instead of serving stale cache.
if not is_local and (ctx != DEFAULT_CONTEXT or configured_kind in ("api", "proxy")):
has_dynamic_runtime_context = endpoint_url in _ollama_runtime_endpoints
if not is_local and not has_dynamic_runtime_context and (
ctx != DEFAULT_CONTEXT or configured_kind in ("api", "proxy")
):
_context_cache[cache_key] = (ctx, known)
logger.info(f"Context length for {model}: {ctx}")
return ctx, known
@ -393,6 +398,107 @@ def _proxy_catalog_context(endpoint_url: str, model: str) -> Optional[int]:
return None
def _ollama_ps_url(endpoint_url: str, configured_kind: Optional[str]) -> Optional[str]:
"""Return a local Ollama ``/api/ps`` URL for a compatible endpoint URL."""
if configured_kind == "proxy":
return None
try:
parsed = urlparse(endpoint_url or "")
host = (parsed.hostname or "").lower()
if parsed.scheme not in ("http", "https") or not host:
return None
named_local_ollama = (
configured_kind in ("api", "local")
and "." not in host
and "ollama" in host
)
network_local = (
host in _LOCAL_HOSTS
or _is_private_ip_literal(host)
or _in_tailscale_range(host)
or named_local_ollama
)
path = (parsed.path or "").rstrip("/")
ollama_looking = (
parsed.port == 11434
or "ollama" in host
or path == "/api"
or "/api/" in path
or (host in _LOCAL_HOSTS and configured_kind == "api")
)
if not network_local or not ollama_looking:
return None
api_path = ""
for suffix in (
"/v1/chat/completions",
"/v1/completions",
"/v1/models",
"/v1",
):
if path.endswith(suffix):
api_path = path[: -len(suffix)] + "/api"
break
if not api_path:
for suffix in ("/api/chat", "/api/generate", "/api/tags", "/api/ps"):
if path.endswith(suffix):
api_path = path[: -len(suffix)] + "/api"
break
if not api_path and (path == "/api" or path.endswith("/api")):
api_path = path
if not api_path and not path:
api_path = "/api"
if not api_path:
return None
return parsed._replace(
path=api_path.rstrip("/") + "/ps",
params="",
query="",
fragment="",
).geturl()
except Exception:
return None
def _probe_ollama_runtime_context(
endpoint_url: str,
model: str,
*,
configured_kind: Optional[str],
) -> Optional[int]:
"""Ask Ollama for the allocation of a currently loaded model."""
ps_url = _ollama_ps_url(endpoint_url, configured_kind)
if not ps_url:
return None
try:
response = httpx.get(ps_url, timeout=REQUEST_TIMEOUT)
if not response.is_success:
return None
payload = response.json()
if not isinstance(payload, dict) or not isinstance(payload.get("models"), list):
return None
# A valid /api/ps response means this endpoint has dynamic runtime
# allocation even when the requested model is not loaded yet. Avoid
# pinning a static-table fallback in the remote/API cache in that case.
_ollama_runtime_endpoints.add(endpoint_url)
from src.model_capability_readers import ollama
runtime = ollama.runtime_context_from_ps_payload(
model,
payload,
base_url=endpoint_url,
)
return runtime.allocated_context_tokens if runtime is not None else None
except Exception as exc:
logger.debug("Ollama runtime context probe failed for %s: %s", model, exc)
return None
def _query_context_length(endpoint_url: str, model: str) -> Tuple[int, bool]:
"""Query the model API for context length. Returns (context_length, known) where
``known`` is False only for the bare DEFAULT_CONTEXT fallback."""
@ -400,6 +506,28 @@ def _query_context_length(endpoint_url: str, model: str) -> Tuple[int, bool]:
api_ctx = None
configured_kind = _configured_endpoint_kind(endpoint_url)
# Manually registered local Ollama endpoints may be stored as
# endpoint_kind="api". Ask the native runtime endpoint before that kind's
# static/catalog short-circuit.
if configured_kind == "api":
runtime_ctx = _probe_ollama_runtime_context(
endpoint_url,
model,
configured_kind=configured_kind,
)
if runtime_ctx:
logger.info(
"Ollama loaded runtime allocation for %s: %s",
model,
runtime_ctx,
)
return runtime_ctx, True
if endpoint_url in _ollama_runtime_endpoints:
# The endpoint answered as Ollama, but this model is not loaded, so
# no effective allocation exists yet. Do not promote a model-table
# maximum into runtime truth or send it back as options.num_ctx.
return DEFAULT_CONTEXT, False
# Large OpenAI-compatible proxies can make /models expensive. If the
# endpoint is explicitly configured as API/proxy, prefer known context
# metadata (or the default) over downloading the full catalog.
@ -432,6 +560,22 @@ def _query_context_length(endpoint_url: str, model: str) -> Tuple[int, bool]:
except Exception:
pass
if configured_kind != "api":
runtime_ctx = _probe_ollama_runtime_context(
endpoint_url,
model,
configured_kind=configured_kind,
)
if runtime_ctx:
logger.info(
"Ollama loaded runtime allocation for %s: %s",
model,
runtime_ctx,
)
return runtime_ctx, True
if endpoint_url in _ollama_runtime_endpoints:
return DEFAULT_CONTEXT, False
# GitHub Copilot's /models requires auth + X-GitHub-Api-Version headers that
# aren't available here; an unauthenticated probe just 400s. All Copilot
# picker models are major API models covered by the known-context table, so

View file

@ -204,6 +204,7 @@ class TestGetContextLength:
def setup_method(self):
model_context._context_cache.clear()
model_context._catalog_ctx_cache.clear()
model_context._ollama_runtime_endpoints.clear()
def test_local_endpoint_requeries_same_model_after_restart(self, monkeypatch):
calls = []
@ -253,6 +254,75 @@ class TestGetContextLength:
)
])
def _api_db(self, monkeypatch, base_url="http://127.0.0.1:11434/v1"):
_install_endpoint_db(monkeypatch, [
types.SimpleNamespace(
base_url=base_url,
endpoint_kind="api",
api_key=None,
is_enabled=True,
)
])
def test_configured_api_ollama_uses_loaded_runtime_allocation(self, monkeypatch):
self._api_db(monkeypatch)
fetches = []
def fake_get(url, *args, **kwargs):
fetches.append(url)
assert url == "http://127.0.0.1:11434/api/ps"
return _FakeResp({
"models": [
{"model": "qwen3:latest", "context_length": 65536},
]
})
monkeypatch.setattr(model_context.httpx, "get", fake_get)
endpoint = "http://127.0.0.1:11434/v1/chat/completions"
assert model_context.get_context_length(endpoint, "qwen3") == 65536
assert fetches == ["http://127.0.0.1:11434/api/ps"]
def test_configured_api_ollama_retries_after_model_load(self, monkeypatch):
self._api_db(monkeypatch)
responses = [
{"models": []},
{"models": [{"model": "qwen3", "context_length": 32768}]},
]
def fake_get(url, *args, **kwargs):
assert url == "http://127.0.0.1:11434/api/ps"
return _FakeResp(responses.pop(0))
monkeypatch.setattr(model_context.httpx, "get", fake_get)
endpoint = "http://127.0.0.1:11434/v1/chat/completions"
assert model_context.get_context_length_known(endpoint, "qwen3") == (
model_context.DEFAULT_CONTEXT,
False,
)
assert model_context.get_context_length_known(endpoint, "qwen3") == (32768, True)
assert responses == []
@pytest.mark.parametrize(
("endpoint", "expected"),
[
("http://localhost:11434", "http://localhost:11434/api/ps"),
("http://ollama:11434/v1", "http://ollama:11434/api/ps"),
("http://localhost:11434/api/chat", "http://localhost:11434/api/ps"),
("http://localhost:11434/v1/chat/completions", "http://localhost:11434/api/ps"),
("http://localhost:11434/ollama/v1", "http://localhost:11434/ollama/api/ps"),
],
)
def test_ollama_ps_url_preserves_supported_prefixes(self, endpoint, expected):
assert model_context._ollama_ps_url(endpoint, "api") == expected
def test_ollama_ps_url_excludes_remote_and_configured_proxy(self):
assert model_context._ollama_ps_url("https://api.openai.com/v1", "api") is None
assert model_context._ollama_ps_url("https://ollama.example.com/v1", "api") is None
assert model_context._ollama_ps_url("http://127.0.0.1:11434/v1", "proxy") is None
assert model_context._ollama_ps_url("http://127.0.0.1:8080/v1", None) is None
def test_configured_proxy_known_model_skips_model_listing(self, monkeypatch):
# A model covered by the known-context table must still resolve without
# touching /models — the cheap path the proxy short-circuit exists for.