mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-07-07 18:01:17 -03:00
refactor(llm): use catalog-based max_tokens, remove JSON retry, reduce Ollama num_ctx
- Parse limit.output from model catalog alongside model IDs for per-model max output token limits - Use catalog lookup in chat_completion_json() to set max_tokens; fall back to 4096 for unknown models (e.g. local Ollama) - Remove the JSON retry (response_format → plain text fallback); keep _try_salvage_json as last-resort for truncated responses - Reduce Ollama num_ctx from 32768 to 8192 (sufficient for metadata enrichment, saves VRAM) - Fix stale test comment referencing removed retry
This commit is contained in:
@@ -27,18 +27,27 @@ _MODEL_CATALOG_URL = "https://models.dev/api.json"
|
|||||||
|
|
||||||
# In-memory cache: maps provider slug -> list of model ID strings.
|
# In-memory cache: maps provider slug -> list of model ID strings.
|
||||||
_catalog_cache: Optional[Dict[str, List[str]]] = None
|
_catalog_cache: Optional[Dict[str, List[str]]] = None
|
||||||
|
|
||||||
|
# Per-model max output token limits parsed from the catalog.
|
||||||
|
# ``{provider_id: {model_id: max_output_tokens}}``.
|
||||||
|
_model_output_limits: Dict[str, Dict[str, int]] = {}
|
||||||
|
|
||||||
_CATALOG_TIMEOUT = aiohttp.ClientTimeout(total=30)
|
_CATALOG_TIMEOUT = aiohttp.ClientTimeout(total=30)
|
||||||
|
|
||||||
|
|
||||||
async def _load_model_catalog() -> Dict[str, List[str]]:
|
async def _load_model_catalog() -> Dict[str, List[str]]:
|
||||||
"""Fetch and parse the model catalog, returning ``{provider_id: [model_id, ...]}``.
|
"""Fetch and parse the model catalog.
|
||||||
|
|
||||||
|
Returns ``{provider_id: [model_id, ...]}`` and also populates
|
||||||
|
:data:`_model_output_limits` with per-model ``limit.output`` values
|
||||||
|
for use by :func:`_get_model_max_output`.
|
||||||
|
|
||||||
The JSON at ``_MODEL_CATALOG_URL`` is a dict keyed by provider slug; each
|
The JSON at ``_MODEL_CATALOG_URL`` is a dict keyed by provider slug; each
|
||||||
value has a ``models`` sub-dict keyed by model ID. Only the model IDs are
|
value has a ``models`` sub-dict keyed by model ID. The result is cached
|
||||||
kept. The result is cached in memory after the first successful fetch.
|
in memory after the first successful fetch.
|
||||||
Subsequent calls return the cached data immediately.
|
Subsequent calls return the cached data immediately.
|
||||||
"""
|
"""
|
||||||
global _catalog_cache
|
global _catalog_cache, _model_output_limits
|
||||||
if _catalog_cache is not None:
|
if _catalog_cache is not None:
|
||||||
return _catalog_cache
|
return _catalog_cache
|
||||||
|
|
||||||
@@ -58,25 +67,52 @@ async def _load_model_catalog() -> Dict[str, List[str]]:
|
|||||||
return _catalog_cache or {}
|
return _catalog_cache or {}
|
||||||
|
|
||||||
result: Dict[str, List[str]] = {}
|
result: Dict[str, List[str]] = {}
|
||||||
|
output_limits: Dict[str, Dict[str, int]] = {}
|
||||||
for provider_id, provider_info in data.items():
|
for provider_id, provider_info in data.items():
|
||||||
if not isinstance(provider_info, dict):
|
if not isinstance(provider_info, dict):
|
||||||
continue
|
continue
|
||||||
models_dict = provider_info.get("models")
|
models_dict = provider_info.get("models")
|
||||||
if not isinstance(models_dict, dict):
|
if not isinstance(models_dict, dict):
|
||||||
continue
|
continue
|
||||||
model_ids = [str(mid) for mid in models_dict.keys() if isinstance(mid, str)]
|
model_ids: List[str] = []
|
||||||
|
provider_limits: Dict[str, int] = {}
|
||||||
|
for mid, model_info in models_dict.items():
|
||||||
|
if not isinstance(mid, str):
|
||||||
|
continue
|
||||||
|
model_ids.append(mid)
|
||||||
|
if isinstance(model_info, dict):
|
||||||
|
limit = model_info.get("limit")
|
||||||
|
if isinstance(limit, dict):
|
||||||
|
output = limit.get("output")
|
||||||
|
if isinstance(output, (int, float)) and output > 0:
|
||||||
|
provider_limits[mid] = int(output)
|
||||||
if model_ids:
|
if model_ids:
|
||||||
result[provider_id] = model_ids
|
result[provider_id] = model_ids
|
||||||
|
if provider_limits:
|
||||||
|
output_limits[provider_id] = provider_limits
|
||||||
|
|
||||||
_catalog_cache = result
|
_catalog_cache = result
|
||||||
|
_model_output_limits = output_limits
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"Loaded model catalog: %d providers, %d total models",
|
"Loaded model catalog: %d providers, %d total models "
|
||||||
|
"(%d providers have output limits)",
|
||||||
len(result),
|
len(result),
|
||||||
sum(len(m) for m in result.values()),
|
sum(len(m) for m in result.values()),
|
||||||
|
len(output_limits),
|
||||||
)
|
)
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def _get_model_max_output(provider: str, model: str) -> Optional[int]:
|
||||||
|
"""Return the model's max output token limit from the catalog, or ``None``.
|
||||||
|
|
||||||
|
Returns ``None`` when the provider or model is not found in the catalog
|
||||||
|
(e.g. local Ollama models, custom models, or user-typed model names).
|
||||||
|
Callers should fall back to a safe default.
|
||||||
|
"""
|
||||||
|
return _model_output_limits.get(provider, {}).get(model)
|
||||||
|
|
||||||
|
|
||||||
# Short timeout for Ollama's local API
|
# Short timeout for Ollama's local API
|
||||||
_OLLAMA_API_TIMEOUT = aiohttp.ClientTimeout(total=8)
|
_OLLAMA_API_TIMEOUT = aiohttp.ClientTimeout(total=8)
|
||||||
|
|
||||||
@@ -364,9 +400,11 @@ class LLMService:
|
|||||||
"think": False,
|
"think": False,
|
||||||
"options": {
|
"options": {
|
||||||
"temperature": temperature,
|
"temperature": temperature,
|
||||||
# Allow up to 32K context so the model has room to think
|
# 8K context is sufficient for metadata enrichment
|
||||||
# AND produce output without hitting the 4K default limit.
|
# (prompt ~2-5K, output ~0.2-1K tokens). The old 32K
|
||||||
"num_ctx": 32768,
|
# value was excessive for this use case and increased
|
||||||
|
# Ollama VRAM usage unnecessarily.
|
||||||
|
"num_ctx": 8192,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
if response_format is not None:
|
if response_format is not None:
|
||||||
@@ -480,11 +518,15 @@ class LLMService:
|
|||||||
temperature: float = 0.3,
|
temperature: float = 0.3,
|
||||||
max_tokens: Optional[int] = None,
|
max_tokens: Optional[int] = None,
|
||||||
) -> Dict[str, Any]:
|
) -> Dict[str, Any]:
|
||||||
"""Call the LLM and return parsed JSON.
|
"""Call the LLM with ``response_format=json_object`` and return parsed JSON.
|
||||||
|
|
||||||
Sends ``response_format: {"type": "json_object"}`` when the provider
|
``max_tokens`` is resolved in this order:
|
||||||
supports it, and parses the response content as JSON. If parsing
|
1. Explicit caller-supplied ``max_tokens``
|
||||||
fails, retries once with a clarifying system message.
|
2. Per-model ``limit.output`` from the model catalog
|
||||||
|
3. A safe default of 4096 (sufficient for metadata enrichment)
|
||||||
|
|
||||||
|
If the response content is empty or not valid JSON, attempts
|
||||||
|
:func:`_try_salvage_json` before raising.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
system_prompt: System-level instructions
|
system_prompt: System-level instructions
|
||||||
@@ -499,7 +541,7 @@ class LLMService:
|
|||||||
Raises:
|
Raises:
|
||||||
LLMNotConfiguredError: Provider not configured
|
LLMNotConfiguredError: Provider not configured
|
||||||
LLMRateLimitError: Rate limited
|
LLMRateLimitError: Rate limited
|
||||||
LLMResponseError: JSON parse failure after retry
|
LLMResponseError: Empty response or JSON parse failure
|
||||||
"""
|
"""
|
||||||
|
|
||||||
messages = [
|
messages = [
|
||||||
@@ -507,10 +549,15 @@ class LLMService:
|
|||||||
{"role": "user", "content": user_prompt},
|
{"role": "user", "content": user_prompt},
|
||||||
]
|
]
|
||||||
|
|
||||||
# First attempt with JSON mode.
|
# Resolve max_tokens: caller override → catalog lookup → safe default
|
||||||
# Use a generous max_tokens so thinking-enabled models (e.g.
|
if max_tokens is None:
|
||||||
# gemma4 via Ollama) have room to reason AND still emit content.
|
cfg = self._get_config()
|
||||||
effective_max = max_tokens or 131072
|
effective_max = _get_model_max_output(cfg["provider"], cfg["model"])
|
||||||
|
else:
|
||||||
|
effective_max = max_tokens
|
||||||
|
if effective_max is None:
|
||||||
|
effective_max = 4096
|
||||||
|
|
||||||
result = await self.chat_completion(
|
result = await self.chat_completion(
|
||||||
messages=messages,
|
messages=messages,
|
||||||
model=model,
|
model=model,
|
||||||
@@ -519,8 +566,15 @@ class LLMService:
|
|||||||
max_tokens=effective_max,
|
max_tokens=effective_max,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
content = result.get("content", "") or ""
|
||||||
|
if not content:
|
||||||
|
raise LLMResponseError(
|
||||||
|
"LLM returned empty content in json_object mode. "
|
||||||
|
f"Raw response: {json.dumps(result)[:500]}"
|
||||||
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
parsed = json.loads(result["content"])
|
parsed = json.loads(content)
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"LLM raw content: %s",
|
"LLM raw content: %s",
|
||||||
json.dumps(parsed, ensure_ascii=False)[:2000],
|
json.dumps(parsed, ensure_ascii=False)[:2000],
|
||||||
@@ -529,52 +583,10 @@ class LLMService:
|
|||||||
except (json.JSONDecodeError, TypeError) as exc:
|
except (json.JSONDecodeError, TypeError) as exc:
|
||||||
logger.info(
|
logger.info(
|
||||||
"LLM raw response (first 800 chars): %s",
|
"LLM raw response (first 800 chars): %s",
|
||||||
(result.get("content") or "")[:800],
|
content[:800],
|
||||||
)
|
|
||||||
logger.warning(
|
|
||||||
"LLM JSON parse failed on first attempt: %s. Retrying.", exc
|
|
||||||
)
|
)
|
||||||
|
|
||||||
# Retry WITHOUT response_format — some providers (Ollama with
|
# Last resort: attempt to salvage partial/truncated JSON
|
||||||
# thinking-enabled models like gemma4) may return empty content
|
|
||||||
# when json_object mode is active. Fall back to a textual
|
|
||||||
# instruction instead.
|
|
||||||
previous_content = result.get("content", "") or ""
|
|
||||||
retry_messages = messages + [
|
|
||||||
{
|
|
||||||
"role": "assistant",
|
|
||||||
"content": previous_content or "(empty response)",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"role": "user",
|
|
||||||
"content": (
|
|
||||||
"The previous response could not be parsed as JSON. "
|
|
||||||
"Please respond with ONLY a valid JSON object, no "
|
|
||||||
"markdown fences or extra text."
|
|
||||||
),
|
|
||||||
},
|
|
||||||
]
|
|
||||||
|
|
||||||
result = await self.chat_completion(
|
|
||||||
messages=retry_messages,
|
|
||||||
model=model,
|
|
||||||
temperature=0.0, # More deterministic for retry
|
|
||||||
max_tokens=effective_max,
|
|
||||||
)
|
|
||||||
|
|
||||||
content = result.get("content", "") or ""
|
|
||||||
if not content:
|
|
||||||
raise LLMResponseError(
|
|
||||||
"LLM response could not be parsed as JSON after retry: "
|
|
||||||
f"Expecting value: line 1 column 1 (char 0)\n"
|
|
||||||
f"Raw content: {content[:500]}"
|
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
|
||||||
return json.loads(content)
|
|
||||||
except (json.JSONDecodeError, TypeError) as parse_err:
|
|
||||||
# Last resort: attempt to salvage partial JSON (closing unclosed
|
|
||||||
# brackets/braces, truncating incomplete strings, etc.)
|
|
||||||
salvaged = _try_salvage_json(content)
|
salvaged = _try_salvage_json(content)
|
||||||
if salvaged is not None:
|
if salvaged is not None:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
@@ -582,10 +594,10 @@ class LLMService:
|
|||||||
len(content),
|
len(content),
|
||||||
)
|
)
|
||||||
return salvaged
|
return salvaged
|
||||||
|
|
||||||
raise LLMResponseError(
|
raise LLMResponseError(
|
||||||
f"LLM response could not be parsed as JSON after retry: {parse_err}\n"
|
f"LLM response could not be parsed as JSON: {content[:200]}"
|
||||||
f"Raw content: {content[:500]}"
|
)
|
||||||
) from parse_err
|
|
||||||
|
|
||||||
|
|
||||||
def _try_salvage_json(raw: str) -> Dict[str, Any] | None:
|
def _try_salvage_json(raw: str) -> Dict[str, Any] | None:
|
||||||
|
|||||||
@@ -219,7 +219,7 @@ class TestLLMServiceChatCompletionJson:
|
|||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_chat_completion_json_raises_on_non_json(self, llm_service):
|
async def test_chat_completion_json_raises_on_non_json(self, llm_service):
|
||||||
# First attempt: non-JSON; second attempt (retry): also non-JSON
|
# Non-JSON content raises LLMResponseError (salvage also fails)
|
||||||
mock_response = MockResponse(
|
mock_response = MockResponse(
|
||||||
200,
|
200,
|
||||||
json_data={
|
json_data={
|
||||||
|
|||||||
Reference in New Issue
Block a user