Skip to content

Commit 38ebfd8

Browse files
committed
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.
1 parent fa29b27 commit 38ebfd8

4 files changed

Lines changed: 197 additions & 24 deletions

File tree

‎raven/providers/rates.py‎

Lines changed: 21 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -218,10 +218,9 @@ def may_prompt(model: str) -> bool:
218218
return any((drivers / part / "authenticator.py").exists() for part in model.split("/") if part)
219219

220220

221-
# LiteLLM answers a huggingface id's context window by fetching the repo's
222-
# `config.json` over the network -- `utils.py::_get_max_position_embeddings`,
223-
# on every call, cached nowhere. `allow_fetch=False` does not reach it: that
224-
# flag governs this module's own OpenRouter fetch, not LiteLLM's.
221+
# LiteLLM fetches a huggingface id's config.json for its context window.
222+
# Skip this driver even when SDK lookups are allowed: the routing prefix
223+
# becomes part of the repository name, so the request returns 404.
225224
_NETWORK_METADATA_VENDORS = frozenset({"huggingface"})
226225

227226

@@ -794,7 +793,7 @@ def _litellm_ready() -> bool:
794793
return not getattr(getattr(module, "__spec__", None), "_initializing", False)
795794

796795

797-
def _try_litellm_max_output(model: str, *, allow_import: bool = True) -> int | None:
796+
def _try_litellm_max_output(model: str, *, allow_import: bool = True, allow_fetch: bool = True) -> int | None:
798797
"""The model's own output ceiling, from the same metadata as the window.
799798
800799
``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
804803
Both tiers are filtered through ``_trustworthy_ceiling``: the rows that get
805804
this wrong answer the table and ``get_model_info`` alike, so guarding one
806805
would hand back exactly what the other just rejected.
806+
807+
``allow_fetch=False`` reads only the static table: ``get_model_info`` can
808+
query a provider even when LiteLLM itself is already imported.
807809
"""
808810
if not allow_import and not _litellm_ready():
809811
return None
@@ -818,7 +820,7 @@ def _try_litellm_max_output(model: str, *, allow_import: bool = True) -> int | N
818820
ceiling = _trustworthy_ceiling(_table_entry(candidate))
819821
if ceiling:
820822
return ceiling
821-
if may_prompt(candidate) or may_fetch(candidate):
823+
if not allow_fetch or may_prompt(candidate) or may_fetch(candidate):
822824
continue
823825
try:
824826
info = litellm.get_model_info(candidate)
@@ -880,9 +882,9 @@ def declared_max_output_tokens(model: str | None, *, allow_fetch: bool = True) -
880882
"""
881883
if not model:
882884
return None
883-
return _try_litellm_max_output(model, allow_import=allow_fetch) or _try_openrouter_max_output(
884-
model, allow_fetch=allow_fetch
885-
)
885+
return _try_litellm_max_output(
886+
model, allow_import=allow_fetch, allow_fetch=allow_fetch
887+
) or _try_openrouter_max_output(model, allow_fetch=allow_fetch)
886888

887889

888890
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
941943
return max(1, min(ceiling, room, MAX_OUTPUT_TOKENS_PER_ITERATION))
942944

943945

944-
def _try_litellm_context_window(model: str, *, allow_import: bool = True) -> int | None:
945-
"""LiteLLM's static model metadata -- offline, covers most mapped providers.
946+
def _try_litellm_context_window(model: str, *, allow_import: bool = True, allow_fetch: bool = True) -> int | None:
947+
"""LiteLLM's model metadata, covering most mapped providers.
946948
947949
Falls back to ``max_tokens`` (the output ceiling) for the handful of rows
948950
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
958960
only, not to wait out the import it is trying to defer. Once LiteLLM is in
959961
hand the check is free and the lookup proceeds exactly as with
960962
``allow_import=True``.
963+
964+
``allow_fetch=False`` reads only the static table: ``get_model_info`` can
965+
query a provider even when LiteLLM itself is already imported.
961966
"""
962967
if not allow_import and not _litellm_ready():
963968
return None
@@ -974,7 +979,7 @@ def _try_litellm_context_window(model: str, *, allow_import: bool = True) -> int
974979
window = _numeric(_table_entry(candidate), "max_input_tokens", "max_tokens")
975980
if window:
976981
return int(window)
977-
if may_prompt(candidate) or may_fetch(candidate):
982+
if not allow_fetch or may_prompt(candidate) or may_fetch(candidate):
978983
continue
979984
try:
980985
info = litellm.get_model_info(candidate)
@@ -1003,11 +1008,11 @@ def resolve_context_window(model: str, *, allow_fetch: bool = True) -> int | Non
10031008
10041009
``allow_fetch=False`` means "answer from what is already on hand": it
10051010
passes through to the OpenRouter tier (see ``_fetch_openrouter_models``)
1006-
and also tells the LiteLLM tier not to import LiteLLM on this caller's
1007-
behalf (see ``_try_litellm_context_window``) -- a caller cheap enough to
1008-
pass this is cheap enough not to pay a fresh import either.
1011+
and limits the LiteLLM tier to its already loaded static table. No fresh
1012+
import or SDK metadata lookup runs on this caller's behalf (see
1013+
``_try_litellm_context_window``).
10091014
"""
1010-
window = _try_litellm_context_window(model, allow_import=allow_fetch)
1015+
window = _try_litellm_context_window(model, allow_import=allow_fetch, allow_fetch=allow_fetch)
10111016
if window:
10121017
return window
10131018

‎tests/test_acp_server.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,77 @@ async def test_the_engine_is_built_on_the_acp_channel(self, stub):
116116
assert stub.channel == "acp"
117117
assert server.ACP_CHANNEL == "acp"
118118

119+
async def test_session_creation_with_uncached_ollama_models_finishes_before_eof_cleanup(
120+
self, stub, tmp_path, monkeypatch
121+
):
122+
from litellm.llms.ollama.completion.transformation import OllamaConfig
123+
124+
from raven.providers.litellm_setup import import_litellm
125+
from raven.rpc.methods.model import model_options
126+
127+
model = "ollama/raven-uncached-model"
128+
config_path = tmp_path / "config.json"
129+
config_path.write_text(
130+
json.dumps(
131+
{
132+
"agents": {"defaults": {"model": model, "provider": "ollama_chat"}},
133+
"providers": {"ollama_chat": {"apiBase": "http://127.0.0.1:11434", "models": [model]}},
134+
}
135+
),
136+
encoding="utf-8",
137+
)
138+
monkeypatch.setattr("raven.home._current_config_path", config_path)
139+
import_litellm()
140+
asked = []
141+
142+
def info(self, model, api_base=None):
143+
asked.append(model)
144+
return {}
145+
146+
monkeypatch.setattr(OllamaConfig, "get_model_info", info)
147+
148+
async def subscribe(params):
149+
return {"subscription_id": "sub-1"}
150+
151+
stub.dispatcher.register("turn.subscribe", subscribe)
152+
stub.dispatcher.register("model.options", model_options)
153+
payload = b"".join(
154+
json.dumps(frame).encode() + b"\n"
155+
for frame in (
156+
{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": 1}},
157+
{
158+
"jsonrpc": "2.0",
159+
"id": 2,
160+
"method": "session/new",
161+
"params": {"cwd": str(tmp_path), "mcpServers": []},
162+
},
163+
{
164+
"jsonrpc": "2.0",
165+
"id": 3,
166+
"method": "session/prompt",
167+
"params": {"sessionId": "acp:missing", "prompt": []},
168+
},
169+
)
170+
)
171+
out = io.BytesIO()
172+
173+
await server.serve(_reader(payload), out)
174+
175+
answers = {frame["id"]: frame for frame in _frames(out) if "id" in frame}
176+
assert sorted(answers) == [1, 2, 3]
177+
assert answers[2]["result"]["sessionId"].startswith("acp:")
178+
options = answers[2]["result"]["configOptions"]
179+
model_option = next(option for option in options if option.get("category") == "model")
180+
assert model_option["currentValue"] == f"ollama_chat/{model}"
181+
assert any(
182+
option["value"] == model_option["currentValue"]
183+
for group in model_option["options"]
184+
for option in group["options"]
185+
)
186+
assert answers[3]["error"]["code"] == -32002
187+
assert asked == []
188+
assert stub.torn_down == 1
189+
119190

120191
class TestConcurrency:
121192
async def test_a_cancel_reaches_a_session_whose_prompt_is_still_running(self, stub, monkeypatch):

‎tests/test_provider_rates.py‎

Lines changed: 72 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -733,8 +733,9 @@ def test_resolve_context_window_allow_fetch_false_also_forwards_allow_import_fal
733733
LiteLLM tier as allow_import=False, not just the OpenRouter tier."""
734734
seen = {}
735735

736-
def _fake_litellm_tier(model, *, allow_import=True):
736+
def _fake_litellm_tier(model, *, allow_import=True, allow_fetch=True):
737737
seen["allow_import"] = allow_import
738+
seen["allow_fetch"] = allow_fetch
738739
return None
739740

740741
monkeypatch.setattr(rates, "_try_litellm_context_window", _fake_litellm_tier)
@@ -743,6 +744,7 @@ def _fake_litellm_tier(model, *, allow_import=True):
743744
rates.resolve_context_window("openrouter/deepseek/deepseek-v4-pro", allow_fetch=False)
744745

745746
assert seen["allow_import"] is False
747+
assert seen["allow_fetch"] is False
746748

747749

748750
# --- Disk persistence of the OpenRouter catalog ---
@@ -990,6 +992,69 @@ def _seed_catalog(monkeypatch, table: dict) -> None:
990992
monkeypatch.setattr(rates, "_OPENROUTER_CACHE_TIME", 0.0, raising=False)
991993

992994

995+
@pytest.mark.parametrize(
996+
"model",
997+
[
998+
"ollama/raven-uncached-model",
999+
"ollama_chat/raven-uncached-model",
1000+
"openrouter/ollama/raven-uncached-model",
1001+
"custom/raven-uncached-model",
1002+
],
1003+
)
1004+
@pytest.mark.parametrize("lookup", [rates.resolve_context_window, rates.declared_max_output_tokens])
1005+
def test_no_fetch_metadata_does_not_ask_litellm_on_a_table_miss(monkeypatch, model, lookup):
1006+
asked = []
1007+
1008+
def info(candidate):
1009+
asked.append(candidate)
1010+
return {"max_input_tokens": 16384, "max_output_tokens": 4096}
1011+
1012+
_patch_table(monkeypatch, {}, info=info)
1013+
1014+
assert lookup(model, allow_fetch=False) is None
1015+
assert asked == []
1016+
1017+
1018+
@pytest.mark.parametrize(
1019+
("lookup", "expected"),
1020+
[(rates.resolve_context_window, 16384), (rates.declared_max_output_tokens, 4096)],
1021+
)
1022+
def test_no_fetch_metadata_reads_a_loaded_alias_row(monkeypatch, lookup, expected):
1023+
model = "ollama/raven-cached-model"
1024+
asked = []
1025+
1026+
def info(candidate):
1027+
asked.append(candidate)
1028+
return {}
1029+
1030+
_patch_table(
1031+
monkeypatch,
1032+
{f"openrouter/{model}": {"max_input_tokens": 16384, "max_output_tokens": 4096}},
1033+
info=info,
1034+
)
1035+
1036+
assert lookup(model, allow_fetch=False) == expected
1037+
assert asked == []
1038+
1039+
1040+
@pytest.mark.parametrize(
1041+
("lookup", "expected"),
1042+
[(rates.resolve_context_window, 16384), (rates.declared_max_output_tokens, 4096)],
1043+
)
1044+
def test_metadata_lookup_can_ask_litellm_when_fetch_is_allowed(monkeypatch, lookup, expected):
1045+
model = "ollama/raven-uncached-model"
1046+
asked = []
1047+
1048+
def info(candidate):
1049+
asked.append(candidate)
1050+
return {"max_input_tokens": 16384, "max_output_tokens": 4096}
1051+
1052+
_patch_table(monkeypatch, {}, info=info)
1053+
1054+
assert lookup(model, allow_fetch=True) == expected
1055+
assert asked == [model]
1056+
1057+
9931058
def test_a_model_the_catalogue_does_not_know_gets_the_one_fallback(monkeypatch):
9941059
"""One number for every unknown model, claude or not.
9951060
@@ -1449,15 +1514,14 @@ def _boom(*_a, **_k):
14491514
assert table == {"m": {"pricing": {}}}
14501515

14511516

1452-
def test_a_huggingface_id_is_never_handed_to_litellm_for_its_window(monkeypatch):
1517+
@pytest.mark.parametrize("allow_fetch", [False, True])
1518+
def test_a_huggingface_id_is_never_handed_to_litellm_for_its_window(monkeypatch, allow_fetch):
14531519
"""The lookup that was supposed to read a table went to the network.
14541520
14551521
LiteLLM resolves a huggingface id's window by fetching the repo's
1456-
`config.json`, on every call and cached nowhere, and `allow_fetch=False`
1457-
does not reach it. `model.options` resolves a window per configured vendor,
1458-
so the settings page paid six sequential round trips for answers that were
1459-
404 anyway -- the id still carries its routing prefix, so the address is
1460-
`huggingface.co/huggingface/<repo>`, which is not a repo.
1522+
`config.json`, on every call and cached nowhere. Even where fetching is
1523+
allowed, this driver is skipped: the id still carries its routing prefix,
1524+
so the address is `huggingface.co/huggingface/<repo>`, which is not a repo.
14611525
"""
14621526
from raven.providers import rates
14631527
from raven.providers.wire import metadata_candidates
@@ -1477,6 +1541,6 @@ def get_model_info(candidate: str) -> dict:
14771541
monkeypatch.setattr("raven.providers.litellm_setup.import_litellm", lambda: _Fake)
14781542
monkeypatch.setattr(rates, "_table_entry", lambda _c: None)
14791543

1480-
rates.resolve_context_window(model, allow_fetch=False)
1544+
rates.resolve_context_window(model, allow_fetch=allow_fetch)
14811545

14821546
assert not [c for c in asked if "huggingface" in c], f"handed over: {asked}"

‎tests/test_rpc_model.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1624,6 +1624,39 @@ async def test_the_context_window_a_picker_shows_is_the_one_a_request_is_sized_w
16241624
assert shown == resolve_context_window("openai/gpt-4o", allow_fetch=False)
16251625

16261626

1627+
async def test_options_read_cached_windows_without_asking_the_ollama_daemon(fake_home, monkeypatch):
1628+
from litellm.llms.ollama.completion.transformation import OllamaConfig
1629+
1630+
from raven.providers.litellm_setup import import_litellm
1631+
1632+
uncached = "ollama/raven-uncached-model"
1633+
cached = "ollama/raven-cached-model"
1634+
_write_config(
1635+
fake_home,
1636+
{
1637+
"agents": {"defaults": {"model": uncached, "provider": "ollama_chat"}},
1638+
"providers": {"ollama_chat": {"apiBase": "http://127.0.0.1:11434", "models": [uncached, cached]}},
1639+
},
1640+
)
1641+
litellm = import_litellm()
1642+
monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, cached: {"max_input_tokens": 16384}})
1643+
asked = []
1644+
1645+
def info(self, model, api_base=None):
1646+
asked.append(model)
1647+
return {}
1648+
1649+
monkeypatch.setattr(OllamaConfig, "get_model_info", info)
1650+
1651+
result = await model_options({})
1652+
1653+
entry = _entry(result, "ollama_chat")
1654+
assert uncached in entry["models"]
1655+
assert entry["model_labels"][cached]["context_window"] == 16384
1656+
assert "context_window" not in entry["model_labels"].get(uncached, {})
1657+
assert asked == []
1658+
1659+
16271660
async def test_options_config_reads_do_not_scale_with_the_row_count(
16281661
fake_home: Path,
16291662
monkeypatch: pytest.MonkeyPatch,

0 commit comments

Comments
 (0)