From b65c578c71244bacb83b6a2c8c2e223f3fa8e2ba Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Sun, 11 Oct 2026 21:18:44 +0800 Subject: [PATCH 1/2] fix(providers): keep model metadata lookups offline when requested Honor allow_fetch=False for context-window and output-ceiling metadata, including a loaded LiteLLM whose SDK can query provider endpoints. Keep cached rows and explicitly allowed SDK lookups available. Cover cache misses, aliases, model options, and ACP session creation at EOF. --- raven/providers/rates.py | 37 +++++++++-------- tests/test_acp_server.py | 71 ++++++++++++++++++++++++++++++++ tests/test_provider_rates.py | 80 ++++++++++++++++++++++++++++++++---- tests/test_rpc_model.py | 33 +++++++++++++++ 4 files changed, 197 insertions(+), 24 deletions(-) diff --git a/raven/providers/rates.py b/raven/providers/rates.py index 300187275..7d5145791 100644 --- a/raven/providers/rates.py +++ b/raven/providers/rates.py @@ -218,10 +218,9 @@ def may_prompt(model: str) -> bool: return any((drivers / part / "authenticator.py").exists() for part in model.split("/") if part) -# LiteLLM answers a huggingface id's context window by fetching the repo's -# `config.json` over the network -- `utils.py::_get_max_position_embeddings`, -# on every call, cached nowhere. `allow_fetch=False` does not reach it: that -# flag governs this module's own OpenRouter fetch, not LiteLLM's. +# LiteLLM fetches a huggingface id's config.json for its context window. +# Skip this driver even when SDK lookups are allowed: the routing prefix +# becomes part of the repository name, so the request returns 404. _NETWORK_METADATA_VENDORS = frozenset({"huggingface"}) @@ -794,7 +793,7 @@ def _litellm_ready() -> bool: return not getattr(getattr(module, "__spec__", None), "_initializing", False) -def _try_litellm_max_output(model: str, *, allow_import: bool = True) -> int | None: +def _try_litellm_max_output(model: str, *, allow_import: bool = True, allow_fetch: bool = True) -> int | None: """The model's own output ceiling, from the same metadata as the window. ``max_tokens`` is a legitimate fallback here in a way it is not on the @@ -804,6 +803,9 @@ def _try_litellm_max_output(model: str, *, allow_import: bool = True) -> int | N Both tiers are filtered through ``_trustworthy_ceiling``: the rows that get this wrong answer the table and ``get_model_info`` alike, so guarding one would hand back exactly what the other just rejected. + + ``allow_fetch=False`` reads only the static table: ``get_model_info`` can + query a provider even when LiteLLM itself is already imported. """ if not allow_import and not _litellm_ready(): return None @@ -818,7 +820,7 @@ def _try_litellm_max_output(model: str, *, allow_import: bool = True) -> int | N ceiling = _trustworthy_ceiling(_table_entry(candidate)) if ceiling: return ceiling - if may_prompt(candidate) or may_fetch(candidate): + if not allow_fetch or may_prompt(candidate) or may_fetch(candidate): continue try: info = litellm.get_model_info(candidate) @@ -880,9 +882,9 @@ def declared_max_output_tokens(model: str | None, *, allow_fetch: bool = True) - """ if not model: return None - return _try_litellm_max_output(model, allow_import=allow_fetch) or _try_openrouter_max_output( - model, allow_fetch=allow_fetch - ) + return _try_litellm_max_output( + model, allow_import=allow_fetch, allow_fetch=allow_fetch + ) or _try_openrouter_max_output(model, allow_fetch=allow_fetch) def resolve_max_output_tokens(model: str | None, *, window: int | None = None, allow_fetch: bool = True) -> int: @@ -941,8 +943,8 @@ def resolve_max_output_tokens(model: str | None, *, window: int | None = None, a return max(1, min(ceiling, room, MAX_OUTPUT_TOKENS_PER_ITERATION)) -def _try_litellm_context_window(model: str, *, allow_import: bool = True) -> int | None: - """LiteLLM's static model metadata -- offline, covers most mapped providers. +def _try_litellm_context_window(model: str, *, allow_import: bool = True, allow_fetch: bool = True) -> int | None: + """LiteLLM's model metadata, covering most mapped providers. Falls back to ``max_tokens`` (the output ceiling) for the handful of rows that carry no ``max_input_tokens``. Not because the two mean the same @@ -958,6 +960,9 @@ def _try_litellm_context_window(model: str, *, allow_import: bool = True) -> int only, not to wait out the import it is trying to defer. Once LiteLLM is in hand the check is free and the lookup proceeds exactly as with ``allow_import=True``. + + ``allow_fetch=False`` reads only the static table: ``get_model_info`` can + query a provider even when LiteLLM itself is already imported. """ if not allow_import and not _litellm_ready(): return None @@ -974,7 +979,7 @@ def _try_litellm_context_window(model: str, *, allow_import: bool = True) -> int window = _numeric(_table_entry(candidate), "max_input_tokens", "max_tokens") if window: return int(window) - if may_prompt(candidate) or may_fetch(candidate): + if not allow_fetch or may_prompt(candidate) or may_fetch(candidate): continue try: info = litellm.get_model_info(candidate) @@ -1003,11 +1008,11 @@ def resolve_context_window(model: str, *, allow_fetch: bool = True) -> int | Non ``allow_fetch=False`` means "answer from what is already on hand": it passes through to the OpenRouter tier (see ``_fetch_openrouter_models``) - and also tells the LiteLLM tier not to import LiteLLM on this caller's - behalf (see ``_try_litellm_context_window``) -- a caller cheap enough to - pass this is cheap enough not to pay a fresh import either. + and limits the LiteLLM tier to its already loaded static table. No fresh + import or SDK metadata lookup runs on this caller's behalf (see + ``_try_litellm_context_window``). """ - window = _try_litellm_context_window(model, allow_import=allow_fetch) + window = _try_litellm_context_window(model, allow_import=allow_fetch, allow_fetch=allow_fetch) if window: return window diff --git a/tests/test_acp_server.py b/tests/test_acp_server.py index 68ba83947..03e6d30cb 100644 --- a/tests/test_acp_server.py +++ b/tests/test_acp_server.py @@ -116,6 +116,77 @@ async def test_the_engine_is_built_on_the_acp_channel(self, stub): assert stub.channel == "acp" assert server.ACP_CHANNEL == "acp" + async def test_session_creation_with_uncached_ollama_models_finishes_before_eof_cleanup( + self, stub, tmp_path, monkeypatch + ): + from litellm.llms.ollama.completion.transformation import OllamaConfig + + from raven.providers.litellm_setup import import_litellm + from raven.rpc.methods.model import model_options + + model = "ollama/raven-uncached-model" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "agents": {"defaults": {"model": model, "provider": "ollama_chat"}}, + "providers": {"ollama_chat": {"apiBase": "http://127.0.0.1:11434", "models": [model]}}, + } + ), + encoding="utf-8", + ) + monkeypatch.setattr("raven.home._current_config_path", config_path) + import_litellm() + asked = [] + + def info(self, model, api_base=None): + asked.append(model) + return {} + + monkeypatch.setattr(OllamaConfig, "get_model_info", info) + + async def subscribe(params): + return {"subscription_id": "sub-1"} + + stub.dispatcher.register("turn.subscribe", subscribe) + stub.dispatcher.register("model.options", model_options) + payload = b"".join( + json.dumps(frame).encode() + b"\n" + for frame in ( + {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": 1}}, + { + "jsonrpc": "2.0", + "id": 2, + "method": "session/new", + "params": {"cwd": str(tmp_path), "mcpServers": []}, + }, + { + "jsonrpc": "2.0", + "id": 3, + "method": "session/prompt", + "params": {"sessionId": "acp:missing", "prompt": []}, + }, + ) + ) + out = io.BytesIO() + + await server.serve(_reader(payload), out) + + answers = {frame["id"]: frame for frame in _frames(out) if "id" in frame} + assert sorted(answers) == [1, 2, 3] + assert answers[2]["result"]["sessionId"].startswith("acp:") + options = answers[2]["result"]["configOptions"] + model_option = next(option for option in options if option.get("category") == "model") + assert model_option["currentValue"] == f"ollama_chat/{model}" + assert any( + option["value"] == model_option["currentValue"] + for group in model_option["options"] + for option in group["options"] + ) + assert answers[3]["error"]["code"] == -32002 + assert asked == [] + assert stub.torn_down == 1 + class TestConcurrency: async def test_a_cancel_reaches_a_session_whose_prompt_is_still_running(self, stub, monkeypatch): diff --git a/tests/test_provider_rates.py b/tests/test_provider_rates.py index b9fa4d232..09f9c916b 100644 --- a/tests/test_provider_rates.py +++ b/tests/test_provider_rates.py @@ -733,8 +733,9 @@ def test_resolve_context_window_allow_fetch_false_also_forwards_allow_import_fal LiteLLM tier as allow_import=False, not just the OpenRouter tier.""" seen = {} - def _fake_litellm_tier(model, *, allow_import=True): + def _fake_litellm_tier(model, *, allow_import=True, allow_fetch=True): seen["allow_import"] = allow_import + seen["allow_fetch"] = allow_fetch return None monkeypatch.setattr(rates, "_try_litellm_context_window", _fake_litellm_tier) @@ -743,6 +744,7 @@ def _fake_litellm_tier(model, *, allow_import=True): rates.resolve_context_window("openrouter/deepseek/deepseek-v4-pro", allow_fetch=False) assert seen["allow_import"] is False + assert seen["allow_fetch"] is False # --- Disk persistence of the OpenRouter catalog --- @@ -990,6 +992,69 @@ def _seed_catalog(monkeypatch, table: dict) -> None: monkeypatch.setattr(rates, "_OPENROUTER_CACHE_TIME", 0.0, raising=False) +@pytest.mark.parametrize( + "model", + [ + "ollama/raven-uncached-model", + "ollama_chat/raven-uncached-model", + "openrouter/ollama/raven-uncached-model", + "custom/raven-uncached-model", + ], +) +@pytest.mark.parametrize("lookup", [rates.resolve_context_window, rates.declared_max_output_tokens]) +def test_no_fetch_metadata_does_not_ask_litellm_on_a_table_miss(monkeypatch, model, lookup): + asked = [] + + def info(candidate): + asked.append(candidate) + return {"max_input_tokens": 16384, "max_output_tokens": 4096} + + _patch_table(monkeypatch, {}, info=info) + + assert lookup(model, allow_fetch=False) is None + assert asked == [] + + +@pytest.mark.parametrize( + ("lookup", "expected"), + [(rates.resolve_context_window, 16384), (rates.declared_max_output_tokens, 4096)], +) +def test_no_fetch_metadata_reads_a_loaded_alias_row(monkeypatch, lookup, expected): + model = "ollama/raven-cached-model" + asked = [] + + def info(candidate): + asked.append(candidate) + return {} + + _patch_table( + monkeypatch, + {f"openrouter/{model}": {"max_input_tokens": 16384, "max_output_tokens": 4096}}, + info=info, + ) + + assert lookup(model, allow_fetch=False) == expected + assert asked == [] + + +@pytest.mark.parametrize( + ("lookup", "expected"), + [(rates.resolve_context_window, 16384), (rates.declared_max_output_tokens, 4096)], +) +def test_metadata_lookup_can_ask_litellm_when_fetch_is_allowed(monkeypatch, lookup, expected): + model = "ollama/raven-uncached-model" + asked = [] + + def info(candidate): + asked.append(candidate) + return {"max_input_tokens": 16384, "max_output_tokens": 4096} + + _patch_table(monkeypatch, {}, info=info) + + assert lookup(model, allow_fetch=True) == expected + assert asked == [model] + + def test_a_model_the_catalogue_does_not_know_gets_the_one_fallback(monkeypatch): """One number for every unknown model, claude or not. @@ -1449,15 +1514,14 @@ def _boom(*_a, **_k): assert table == {"m": {"pricing": {}}} -def test_a_huggingface_id_is_never_handed_to_litellm_for_its_window(monkeypatch): +@pytest.mark.parametrize("allow_fetch", [False, True]) +def test_a_huggingface_id_is_never_handed_to_litellm_for_its_window(monkeypatch, allow_fetch): """The lookup that was supposed to read a table went to the network. LiteLLM resolves a huggingface id's window by fetching the repo's - `config.json`, on every call and cached nowhere, and `allow_fetch=False` - does not reach it. `model.options` resolves a window per configured vendor, - so the settings page paid six sequential round trips for answers that were - 404 anyway -- the id still carries its routing prefix, so the address is - `huggingface.co/huggingface/`, which is not a repo. + `config.json`, on every call and cached nowhere. Even where fetching is + allowed, this driver is skipped: the id still carries its routing prefix, + so the address is `huggingface.co/huggingface/`, which is not a repo. """ from raven.providers import rates from raven.providers.wire import metadata_candidates @@ -1477,6 +1541,6 @@ def get_model_info(candidate: str) -> dict: monkeypatch.setattr("raven.providers.litellm_setup.import_litellm", lambda: _Fake) monkeypatch.setattr(rates, "_table_entry", lambda _c: None) - rates.resolve_context_window(model, allow_fetch=False) + rates.resolve_context_window(model, allow_fetch=allow_fetch) assert not [c for c in asked if "huggingface" in c], f"handed over: {asked}" diff --git a/tests/test_rpc_model.py b/tests/test_rpc_model.py index d40ad277b..121cd8d26 100644 --- a/tests/test_rpc_model.py +++ b/tests/test_rpc_model.py @@ -1624,6 +1624,39 @@ async def test_the_context_window_a_picker_shows_is_the_one_a_request_is_sized_w assert shown == resolve_context_window("openai/gpt-4o", allow_fetch=False) +async def test_options_read_cached_windows_without_asking_the_ollama_daemon(fake_home, monkeypatch): + from litellm.llms.ollama.completion.transformation import OllamaConfig + + from raven.providers.litellm_setup import import_litellm + + uncached = "ollama/raven-uncached-model" + cached = "ollama/raven-cached-model" + _write_config( + fake_home, + { + "agents": {"defaults": {"model": uncached, "provider": "ollama_chat"}}, + "providers": {"ollama_chat": {"apiBase": "http://127.0.0.1:11434", "models": [uncached, cached]}}, + }, + ) + litellm = import_litellm() + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, cached: {"max_input_tokens": 16384}}) + asked = [] + + def info(self, model, api_base=None): + asked.append(model) + return {} + + monkeypatch.setattr(OllamaConfig, "get_model_info", info) + + result = await model_options({}) + + entry = _entry(result, "ollama_chat") + assert uncached in entry["models"] + assert entry["model_labels"][cached]["context_window"] == 16384 + assert "context_window" not in entry["model_labels"].get(uncached, {}) + assert asked == [] + + async def test_options_config_reads_do_not_scale_with_the_row_count( fake_home: Path, monkeypatch: pytest.MonkeyPatch, From 4ba33b8fe3b16fb6246d1dd911c66707b4d6c1af Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Mon, 12 Oct 2026 11:38:15 +0000 Subject: [PATCH 2/2] fix(providers): keep offline normalization on cache-only metadata lookups allow_fetch=False skipped get_model_info for every candidate. LiteLLM's table files most direct vendors under the bare model name, so an id such as openai/gpt-4o-mini or vertex_ai/gemini-2.5-pro is found only once get_model_info strips the routing prefix, which it does offline. The cache-only lookup fell back to the default window or to the openrouter/ alias row and its OpenRouter numbers: 203 of 413 measured ids changed from the base. Only the Ollama drivers answer a table miss by asking the serving daemon (POST /api/show), which is the request that delayed session/new. Skip get_model_info for those under allow_fetch=False and keep it for every other driver. huggingface stays skipped under both policies. --- raven/providers/rates.py | 35 +++++++++++++++------ tests/test_provider_rates.py | 60 ++++++++++++++++++++++++++++++++++-- 2 files changed, 84 insertions(+), 11 deletions(-) diff --git a/raven/providers/rates.py b/raven/providers/rates.py index 7d5145791..9ed6d7525 100644 --- a/raven/providers/rates.py +++ b/raven/providers/rates.py @@ -253,6 +253,23 @@ def may_fetch(model: str) -> bool: return any(part in _NETWORK_METADATA_VENDORS for part in model.split("/") if part) +# Ollama's driver answers a table miss by POSTing `{api_base}/api/show` to the +# daemon serving the model. That is a request, not a read: worth making when the +# caller allows a fetch, and what held `session/new` past its EOF grace when not. +_ENDPOINT_METADATA_VENDORS = frozenset({"ollama", "ollama_chat"}) + + +def _may_query_endpoint(model: str) -> bool: + """Would ``get_model_info`` ask this model's own server rather than a table? + + Every other driver answers offline: it strips the routing prefix and reads + the table, which files most direct vendors under the bare name + (``openai/gpt-4o-mini`` is keyed ``gpt-4o-mini``). That normalization is + what a cache-only lookup still needs. Any segment counts, as in ``may_fetch``. + """ + return any(part in _ENDPOINT_METADATA_VENDORS for part in model.split("/") if part) + + def _numeric(entry: dict | None, *fields: str) -> float | None: """First numeric value among ``fields``, or None. @@ -804,8 +821,8 @@ def _try_litellm_max_output(model: str, *, allow_import: bool = True, allow_fetc this wrong answer the table and ``get_model_info`` alike, so guarding one would hand back exactly what the other just rejected. - ``allow_fetch=False`` reads only the static table: ``get_model_info`` can - query a provider even when LiteLLM itself is already imported. + ``allow_fetch=False`` skips ``get_model_info`` only where it would ask the + model's own server (see ``_may_query_endpoint``). """ if not allow_import and not _litellm_ready(): return None @@ -820,7 +837,7 @@ def _try_litellm_max_output(model: str, *, allow_import: bool = True, allow_fetc ceiling = _trustworthy_ceiling(_table_entry(candidate)) if ceiling: return ceiling - if not allow_fetch or may_prompt(candidate) or may_fetch(candidate): + if may_prompt(candidate) or may_fetch(candidate) or (not allow_fetch and _may_query_endpoint(candidate)): continue try: info = litellm.get_model_info(candidate) @@ -961,8 +978,8 @@ def _try_litellm_context_window(model: str, *, allow_import: bool = True, allow_ hand the check is free and the lookup proceeds exactly as with ``allow_import=True``. - ``allow_fetch=False`` reads only the static table: ``get_model_info`` can - query a provider even when LiteLLM itself is already imported. + ``allow_fetch=False`` skips ``get_model_info`` only where it would ask the + model's own server (see ``_may_query_endpoint``). """ if not allow_import and not _litellm_ready(): return None @@ -979,7 +996,7 @@ def _try_litellm_context_window(model: str, *, allow_import: bool = True, allow_ window = _numeric(_table_entry(candidate), "max_input_tokens", "max_tokens") if window: return int(window) - if not allow_fetch or may_prompt(candidate) or may_fetch(candidate): + if may_prompt(candidate) or may_fetch(candidate) or (not allow_fetch and _may_query_endpoint(candidate)): continue try: info = litellm.get_model_info(candidate) @@ -1008,9 +1025,9 @@ def resolve_context_window(model: str, *, allow_fetch: bool = True) -> int | Non ``allow_fetch=False`` means "answer from what is already on hand": it passes through to the OpenRouter tier (see ``_fetch_openrouter_models``) - and limits the LiteLLM tier to its already loaded static table. No fresh - import or SDK metadata lookup runs on this caller's behalf (see - ``_try_litellm_context_window``). + and keeps the LiteLLM tier to lookups that read: no fresh import, and no + metadata request to the model's own server, runs on this caller's behalf + (see ``_try_litellm_context_window``). """ window = _try_litellm_context_window(model, allow_import=allow_fetch, allow_fetch=allow_fetch) if window: diff --git a/tests/test_provider_rates.py b/tests/test_provider_rates.py index 09f9c916b..e64617280 100644 --- a/tests/test_provider_rates.py +++ b/tests/test_provider_rates.py @@ -998,11 +998,10 @@ def _seed_catalog(monkeypatch, table: dict) -> None: "ollama/raven-uncached-model", "ollama_chat/raven-uncached-model", "openrouter/ollama/raven-uncached-model", - "custom/raven-uncached-model", ], ) @pytest.mark.parametrize("lookup", [rates.resolve_context_window, rates.declared_max_output_tokens]) -def test_no_fetch_metadata_does_not_ask_litellm_on_a_table_miss(monkeypatch, model, lookup): +def test_no_fetch_metadata_does_not_ask_the_ollama_driver_on_a_table_miss(monkeypatch, model, lookup): asked = [] def info(candidate): @@ -1015,6 +1014,63 @@ def info(candidate): assert asked == [] +@pytest.mark.parametrize( + ("lookup", "expected"), + [(rates.resolve_context_window, 16384), (rates.declared_max_output_tokens, 4096)], +) +def test_no_fetch_metadata_still_asks_litellm_for_a_driver_that_reads_a_table(monkeypatch, lookup, expected): + model = "custom/raven-uncached-model" + asked = [] + + def info(candidate): + asked.append(candidate) + return {"max_input_tokens": 16384, "max_output_tokens": 4096} + + _patch_table(monkeypatch, {}, info=info) + + assert lookup(model, allow_fetch=False) == expected + assert asked == [model] + + +@pytest.mark.parametrize( + "model", + ["openai/gpt-4o-mini", "openai/gpt-4o", "vertex_ai/gemini-2.5-pro", "anthropic/claude-sonnet-4-5"], +) +def test_no_fetch_metadata_answers_a_direct_vendor_the_way_a_fetch_does(model): + """The table files these under the bare name, so only ``get_model_info`` finds them. + + It finds them offline, by stripping the routing prefix. Skipping it on a + cache-only lookup sized the turn on the fallback window, or on the + ``openrouter/`` alias row, which carries OpenRouter's numbers rather than + the vendor's. + """ + assert rates._table_entry(model) is None + assert rates._try_litellm_context_window(model, allow_fetch=False) == rates._try_litellm_context_window(model) + assert rates._try_litellm_max_output(model, allow_fetch=False) == rates._try_litellm_max_output(model) + assert rates._try_litellm_context_window(model, allow_fetch=False) + assert rates._try_litellm_max_output(model, allow_fetch=False) + + +@pytest.mark.parametrize("model", ["ollama/raven-uncached-model", "ollama_chat/raven-uncached-model"]) +def test_no_fetch_metadata_does_not_reach_the_ollama_daemon(monkeypatch, model): + from litellm.llms.ollama.completion.transformation import OllamaConfig + + asked = [] + + def info(self, model, api_base=None): + asked.append(model) + return {"max_input_tokens": 16384, "max_output_tokens": 4096} + + monkeypatch.setattr(OllamaConfig, "get_model_info", info) + + assert rates.resolve_context_window(model, allow_fetch=False) is None + assert rates.declared_max_output_tokens(model, allow_fetch=False) is None + assert asked == [] + + assert rates.resolve_context_window(model, allow_fetch=True) == 16384 + assert asked + + @pytest.mark.parametrize( ("lookup", "expected"), [(rates.resolve_context_window, 16384), (rates.declared_max_output_tokens, 4096)],