mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/cli: skip blocking GET /models on chat ready
Only fetch remote catalog when interactive model pick needs it; otherwise use config/--model/session and warm Tab cache in the background. Cap list fetch at 5s so a dead API cannot hang startup.
This commit is contained in:
+16
-8
@@ -245,20 +245,26 @@ async def _run_chat( # noqa: C901, PLR0912, PLR0915 — chat orchestration
|
||||
preferred_model=preferred_model,
|
||||
interactive=interactive,
|
||||
)
|
||||
# Build client early so we can always try GET /models for remote-first lists.
|
||||
# Avoid blocking ready on GET /models unless interactive pick needs it.
|
||||
from plyngent.cli.models_source import (
|
||||
client_supports_models,
|
||||
fetch_remote_model_ids,
|
||||
model_choices_for_provider,
|
||||
needs_remote_models_for_selection,
|
||||
)
|
||||
|
||||
client = create_client(provider)
|
||||
remote_ids: list[str] | None = None
|
||||
try:
|
||||
if client_supports_models(client):
|
||||
remote_ids = await fetch_remote_model_ids(client)
|
||||
except RuntimeError, TypeError, OSError, ValueError:
|
||||
remote_ids = None
|
||||
if needs_remote_models_for_selection(
|
||||
provider,
|
||||
preferred_model=preferred_model,
|
||||
interactive=interactive,
|
||||
):
|
||||
client = create_client(provider)
|
||||
try:
|
||||
if client_supports_models(client):
|
||||
remote_ids = await fetch_remote_model_ids(client)
|
||||
except RuntimeError, TypeError, OSError, ValueError, TimeoutError:
|
||||
remote_ids = None
|
||||
choices = model_choices_for_provider(provider, remote_ids=remote_ids)
|
||||
model_id = select_model(
|
||||
provider,
|
||||
@@ -282,9 +288,11 @@ async def _run_chat( # noqa: C901, PLR0912, PLR0915 — chat orchestration
|
||||
interactive_limits=interactive,
|
||||
yolo=yolo,
|
||||
)
|
||||
# Seed remote model cache from startup fetch so Tab/complete stays warm.
|
||||
# Seed cache if we already fetched; else warm in background for Tab.
|
||||
if remote_ids is not None:
|
||||
state.seed_remote_models(remote_ids)
|
||||
elif not oneshot and interactive:
|
||||
state.schedule_remote_models_warm()
|
||||
if not quiet and not oneshot:
|
||||
click.secho(f"workspace: {state.workspace}", fg="bright_black", err=True)
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from typing import TYPE_CHECKING, Literal, Protocol, cast, runtime_checkable
|
||||
|
||||
@@ -10,6 +11,8 @@ if TYPE_CHECKING:
|
||||
|
||||
# Cache remote catalog this long (seconds) unless /models --refresh.
|
||||
DEFAULT_MODELS_CACHE_TTL = 300.0
|
||||
# Bound startup/interactive GET /models so a dead API cannot hang the CLI.
|
||||
DEFAULT_MODELS_FETCH_TIMEOUT = 5.0
|
||||
|
||||
type ModelListPrefer = Literal["remote", "union", "config"]
|
||||
|
||||
@@ -58,21 +61,50 @@ def client_supports_models(client: object) -> bool:
|
||||
return isinstance(client, SupportsModels) or callable(getattr(client, "models", None))
|
||||
|
||||
|
||||
async def fetch_remote_model_ids(client: object) -> list[str]:
|
||||
"""Call ``client.models()``; raise if missing or the call fails."""
|
||||
async def fetch_remote_model_ids(
|
||||
client: object,
|
||||
*,
|
||||
timeout_seconds: float = DEFAULT_MODELS_FETCH_TIMEOUT,
|
||||
) -> list[str]:
|
||||
"""Call ``client.models()``; raise if missing or the call fails.
|
||||
|
||||
*timeout_seconds* bounds the await; use ``0`` or less to wait indefinitely.
|
||||
"""
|
||||
method = getattr(client, "models", None)
|
||||
if not callable(method):
|
||||
msg = "client does not support listing models"
|
||||
raise TypeError(msg)
|
||||
result = method()
|
||||
if inspect.isawaitable(result):
|
||||
result = await result
|
||||
if timeout_seconds > 0:
|
||||
async with asyncio.timeout(timeout_seconds):
|
||||
result = await result
|
||||
else:
|
||||
result = await result
|
||||
if not isinstance(result, list):
|
||||
msg = f"models() returned unexpected type {type(result)!r}"
|
||||
raise TypeError(msg)
|
||||
return [str(item) for item in cast("list[object]", result) if item]
|
||||
|
||||
|
||||
def needs_remote_models_for_selection(
|
||||
provider: Provider,
|
||||
*,
|
||||
preferred_model: str | None,
|
||||
interactive: bool,
|
||||
) -> bool:
|
||||
"""True when interactive model pick needs a remote catalog for a better list.
|
||||
|
||||
Skip network when the model is already known (``--model`` / session) or when
|
||||
config has exactly one model (auto-selected).
|
||||
"""
|
||||
if preferred_model is not None and preferred_model.strip():
|
||||
return False
|
||||
if not interactive:
|
||||
return False
|
||||
return len(config_model_ids(provider)) != 1
|
||||
|
||||
|
||||
def model_choices_for_provider(
|
||||
provider: Provider,
|
||||
*,
|
||||
|
||||
@@ -22,7 +22,7 @@ async def discover_model_ids(provider: Provider) -> list[str]:
|
||||
return []
|
||||
try:
|
||||
return await fetch_remote_model_ids(client)
|
||||
except RuntimeError, TypeError, OSError, ValueError:
|
||||
except RuntimeError, TypeError, OSError, ValueError, TimeoutError:
|
||||
return []
|
||||
|
||||
|
||||
|
||||
@@ -224,6 +224,29 @@ class ReplState:
|
||||
self._remote_models_key = self._models_cache_key()
|
||||
self._remote_models_error = None
|
||||
|
||||
def schedule_remote_models_warm(self) -> None:
|
||||
"""Fire-and-forget ``GET /models`` so Tab complete warms without blocking ready."""
|
||||
import asyncio
|
||||
|
||||
async def _warm() -> None:
|
||||
try:
|
||||
_ = await self.ensure_remote_models(refresh=False)
|
||||
except RuntimeError, TypeError, OSError, ValueError, TimeoutError:
|
||||
return
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return
|
||||
task = loop.create_task(_warm(), name="plyngent-warm-remote-models")
|
||||
# Keep a strong ref until done (same pattern as todo persist tasks).
|
||||
self._todo_persist_tasks.add(task)
|
||||
|
||||
def _done(t: object) -> None:
|
||||
_ = self._todo_persist_tasks.discard(t)
|
||||
|
||||
task.add_done_callback(_done)
|
||||
|
||||
def remember_session_ids(self, ids: Sequence[int]) -> None:
|
||||
"""Cache session ids for Tab completion (``/resume`` / ``/delete``)."""
|
||||
self._session_id_cache = [int(i) for i in ids]
|
||||
@@ -272,7 +295,7 @@ class ReplState:
|
||||
raise TypeError(msg)
|
||||
try:
|
||||
ids = await fetch_remote_model_ids(self.client)
|
||||
except (RuntimeError, TypeError, OSError, ValueError) as exc:
|
||||
except (RuntimeError, TypeError, OSError, ValueError, TimeoutError) as exc:
|
||||
self._remote_models_error = str(exc)
|
||||
raise
|
||||
self._remote_models = list(ids)
|
||||
@@ -286,7 +309,7 @@ class ReplState:
|
||||
remote: list[str] | None
|
||||
try:
|
||||
remote = await self.ensure_remote_models(refresh=refresh)
|
||||
except RuntimeError, TypeError, OSError, ValueError:
|
||||
except RuntimeError, TypeError, OSError, ValueError, TimeoutError:
|
||||
remote = self.cached_remote_models()
|
||||
return model_choices_for_provider(self.provider, remote_ids=remote)
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ from plyngent.cli.models_source import (
|
||||
fetch_remote_model_ids,
|
||||
merge_model_choices,
|
||||
model_choices_for_provider,
|
||||
needs_remote_models_for_selection,
|
||||
)
|
||||
from plyngent.config.models import ModelConfig, OpenAICompatibleProvider
|
||||
|
||||
@@ -55,3 +56,39 @@ async def test_fetch_remote_model_ids() -> None:
|
||||
|
||||
with pytest.raises(TypeError, match="does not support"):
|
||||
_ = await fetch_remote_model_ids(object())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_remote_model_ids_timeout() -> None:
|
||||
import asyncio
|
||||
|
||||
class Slow:
|
||||
async def models(self) -> list[str]:
|
||||
await asyncio.sleep(10)
|
||||
return ["late"]
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
_ = await fetch_remote_model_ids(Slow(), timeout_seconds=0.05)
|
||||
|
||||
|
||||
def test_needs_remote_models_for_selection() -> None:
|
||||
multi = OpenAICompatibleProvider(
|
||||
access_key_or_token="sk",
|
||||
url="https://x/v1",
|
||||
models={"a": ModelConfig(), "b": ModelConfig()},
|
||||
)
|
||||
single = OpenAICompatibleProvider(
|
||||
access_key_or_token="sk",
|
||||
url="https://x/v1",
|
||||
models={"only": ModelConfig()},
|
||||
)
|
||||
empty = OpenAICompatibleProvider(
|
||||
access_key_or_token="sk",
|
||||
url="https://x/v1",
|
||||
models={},
|
||||
)
|
||||
assert needs_remote_models_for_selection(multi, preferred_model="a", interactive=True) is False
|
||||
assert needs_remote_models_for_selection(multi, preferred_model=None, interactive=False) is False
|
||||
assert needs_remote_models_for_selection(single, preferred_model=None, interactive=True) is False
|
||||
assert needs_remote_models_for_selection(multi, preferred_model=None, interactive=True) is True
|
||||
assert needs_remote_models_for_selection(empty, preferred_model=None, interactive=True) is True
|
||||
|
||||
Reference in New Issue
Block a user