296ec08dcf
Chat catalogs and the session switch treated image and video generation models as chat. Exclude them by the capability type and name shape the catalog already publishes, reject selecting one as the session model, and do not restore a primary already known to be non-chat. Desktop shows the fallback switch in the transcript.
137 lines
4.5 KiB
Python
137 lines
4.5 KiB
Python
"""Admit only chat models into shared chat catalogs.
|
|
|
|
Generation models are identified by the capability type or name shape the
|
|
catalog already publishes. Missing, malformed, or unknown metadata fails open.
|
|
No single model id is special-cased.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from typing import Any
|
|
|
|
# Structured catalog types that mean "this row is not a chat completion".
|
|
_GENERATION_TYPES = frozenset({
|
|
"image",
|
|
"image_generation",
|
|
"image-generation",
|
|
"image-gen",
|
|
"video",
|
|
"video_generation",
|
|
"video-generation",
|
|
"video-gen",
|
|
})
|
|
|
|
_GENERATION_OUTPUTS = frozenset({"image", "video"})
|
|
_CHAT_OUTPUTS = frozenset({"text", "chat"})
|
|
|
|
# Separator-qualified generation tokens, plus the generation id tokens the
|
|
# DeepInfra chat catalog already excludes. ``vision`` / ``vl`` are not tokens
|
|
# here — those remain chat models.
|
|
_GENERATION_ID_RE = re.compile(
|
|
r"(?i)(?:^|[/_.-])(?:"
|
|
r"image(?:[/_.-]gen(?:eration)?)?|"
|
|
r"video(?:[/_.-]gen(?:eration)?)?|"
|
|
r"text[/_.-]to[/_.-](?:image|video)|"
|
|
r"stable-diffusion|sdxl|flux"
|
|
r")(?=$|[/_.-])"
|
|
)
|
|
|
|
# Ids a live catalog item already proved are generation-only, even when the
|
|
# id itself has no generation token. Process-local: the proof came from this
|
|
# process's catalog parse.
|
|
_seen_generation_ids: set[str] = set()
|
|
|
|
|
|
def note_generation_model(model_id: Any) -> None:
|
|
key = str(model_id or "").strip().lower()
|
|
if key:
|
|
_seen_generation_ids.add(key)
|
|
|
|
|
|
def model_id_is_generation(model_id: Any) -> bool:
|
|
"""Whether *model_id* names an image/video generation surface."""
|
|
return bool(_GENERATION_ID_RE.search(str(model_id or "").strip()))
|
|
|
|
|
|
def _tag_list_is_generation(tags: Any) -> bool:
|
|
if not isinstance(tags, list):
|
|
return False
|
|
return any(str(tag or "").strip().lower() in _GENERATION_TYPES for tag in tags)
|
|
|
|
|
|
def _output_modalities_are_generation(outputs: Any) -> bool:
|
|
if not isinstance(outputs, list) or not outputs:
|
|
return False
|
|
kinds = {str(item or "").strip().lower() for item in outputs if str(item or "").strip()}
|
|
return bool(kinds & _GENERATION_OUTPUTS) and not bool(kinds & _CHAT_OUTPUTS)
|
|
|
|
|
|
def catalog_item_is_generation(item: Any) -> bool:
|
|
"""Fail open unless this catalog row's own type or name proves it is generation-only."""
|
|
if not isinstance(item, dict):
|
|
return False
|
|
if model_id_is_generation(item.get("id")):
|
|
return True
|
|
capabilities = item.get("capabilities")
|
|
if isinstance(capabilities, dict):
|
|
model_type = str(capabilities.get("type") or "").strip().lower()
|
|
if model_type in _GENERATION_TYPES:
|
|
return True
|
|
if _tag_list_is_generation(capabilities.get("tags")):
|
|
return True
|
|
top_type = item.get("type")
|
|
if isinstance(top_type, str) and top_type.strip().lower() in _GENERATION_TYPES:
|
|
return True
|
|
if _tag_list_is_generation(item.get("tags")):
|
|
return True
|
|
architecture = item.get("architecture")
|
|
if isinstance(architecture, dict) and _output_modalities_are_generation(architecture.get("output_modalities")):
|
|
return True
|
|
metadata = item.get("metadata")
|
|
if isinstance(metadata, dict) and _tag_list_is_generation(metadata.get("tags")):
|
|
return True
|
|
return False
|
|
|
|
|
|
def note_catalog_item(item: Any) -> bool:
|
|
"""Remember a generation row and report whether it must leave the chat catalog."""
|
|
if not catalog_item_is_generation(item):
|
|
return False
|
|
if isinstance(item, dict):
|
|
note_generation_model(item.get("id"))
|
|
return True
|
|
|
|
|
|
def is_known_non_chat_model(model_id: Any) -> bool:
|
|
"""True when the id's name shape, or a catalog item already parsed, says it is not chat."""
|
|
mid = str(model_id or "").strip()
|
|
return model_id_is_generation(mid) or mid.lower() in _seen_generation_ids
|
|
|
|
|
|
def without_generation_models(models):
|
|
"""Drop generation ids, preserving a list subclass (curated-fallback marker)."""
|
|
if not models:
|
|
return models
|
|
kept = [model for model in models if not is_known_non_chat_model(model)]
|
|
if len(kept) == len(list(models)):
|
|
return models
|
|
try:
|
|
return type(models)(kept)
|
|
except TypeError:
|
|
return kept
|
|
|
|
|
|
def chat_catalog_ids(items: Any) -> list[str]:
|
|
"""Model ids from a ``/models`` payload, minus generation rows."""
|
|
if not isinstance(items, list):
|
|
return []
|
|
ids: list[str] = []
|
|
for item in items:
|
|
if not isinstance(item, dict) or "id" not in item:
|
|
continue
|
|
if note_catalog_item(item):
|
|
continue
|
|
ids.append(item["id"])
|
|
return ids
|