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:
Will Miao
2026-07-06 09:13:42 +08:00
parent 07fa454f72
commit 5983eaa1ce
2 changed files with 84 additions and 72 deletions

View File

@@ -27,18 +27,27 @@ _MODEL_CATALOG_URL = "https://models.dev/api.json"
# In-memory cache: maps provider slug -> list of model ID strings.
_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)
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
value has a ``models`` sub-dict keyed by model ID. Only the model IDs are
kept. The result is cached in memory after the first successful fetch.
value has a ``models`` sub-dict keyed by model ID. The result is cached
in memory after the first successful fetch.
Subsequent calls return the cached data immediately.
"""
global _catalog_cache
global _catalog_cache, _model_output_limits
if _catalog_cache is not None:
return _catalog_cache
@@ -58,25 +67,52 @@ async def _load_model_catalog() -> Dict[str, List[str]]:
return _catalog_cache or {}
result: Dict[str, List[str]] = {}
output_limits: Dict[str, Dict[str, int]] = {}
for provider_id, provider_info in data.items():
if not isinstance(provider_info, dict):
continue
models_dict = provider_info.get("models")
if not isinstance(models_dict, dict):
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:
result[provider_id] = model_ids
if provider_limits:
output_limits[provider_id] = provider_limits
_catalog_cache = result
_model_output_limits = output_limits
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),
sum(len(m) for m in result.values()),
len(output_limits),
)
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
_OLLAMA_API_TIMEOUT = aiohttp.ClientTimeout(total=8)
@@ -364,9 +400,11 @@ class LLMService:
"think": False,
"options": {
"temperature": temperature,
# Allow up to 32K context so the model has room to think
# AND produce output without hitting the 4K default limit.
"num_ctx": 32768,
# 8K context is sufficient for metadata enrichment
# (prompt ~2-5K, output ~0.2-1K tokens). The old 32K
# value was excessive for this use case and increased
# Ollama VRAM usage unnecessarily.
"num_ctx": 8192,
},
}
if response_format is not None:
@@ -480,11 +518,15 @@ class LLMService:
temperature: float = 0.3,
max_tokens: Optional[int] = None,
) -> 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
supports it, and parses the response content as JSON. If parsing
fails, retries once with a clarifying system message.
``max_tokens`` is resolved in this order:
1. Explicit caller-supplied ``max_tokens``
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:
system_prompt: System-level instructions
@@ -499,7 +541,7 @@ class LLMService:
Raises:
LLMNotConfiguredError: Provider not configured
LLMRateLimitError: Rate limited
LLMResponseError: JSON parse failure after retry
LLMResponseError: Empty response or JSON parse failure
"""
messages = [
@@ -507,10 +549,15 @@ class LLMService:
{"role": "user", "content": user_prompt},
]
# First attempt with JSON mode.
# Use a generous max_tokens so thinking-enabled models (e.g.
# gemma4 via Ollama) have room to reason AND still emit content.
effective_max = max_tokens or 131072
# Resolve max_tokens: caller override → catalog lookup → safe default
if max_tokens is None:
cfg = self._get_config()
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(
messages=messages,
model=model,
@@ -519,8 +566,15 @@ class LLMService:
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:
parsed = json.loads(result["content"])
parsed = json.loads(content)
logger.debug(
"LLM raw content: %s",
json.dumps(parsed, ensure_ascii=False)[:2000],
@@ -529,64 +583,22 @@ class LLMService:
except (json.JSONDecodeError, TypeError) as exc:
logger.info(
"LLM raw response (first 800 chars): %s",
(result.get("content") or "")[:800],
content[:800],
)
# Last resort: attempt to salvage partial/truncated JSON
salvaged = _try_salvage_json(content)
if salvaged is not None:
logger.warning(
"LLM JSON parse failed on first attempt: %s. Retrying.", exc
"LLM JSON salvaged from partial content (%d chars raw)",
len(content),
)
return salvaged
# Retry WITHOUT response_format — some providers (Ollama with
# 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,
raise LLMResponseError(
f"LLM response could not be parsed as JSON: {content[:200]}"
)
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)
if salvaged is not None:
logger.warning(
"LLM JSON salvaged from partial content (%d chars raw)",
len(content),
)
return salvaged
raise LLMResponseError(
f"LLM response could not be parsed as JSON after retry: {parse_err}\n"
f"Raw content: {content[:500]}"
) from parse_err
def _try_salvage_json(raw: str) -> Dict[str, Any] | None:
"""Attempt to repair and parse a truncated JSON string.

View File

@@ -219,7 +219,7 @@ class TestLLMServiceChatCompletionJson:
@pytest.mark.asyncio
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(
200,
json_data={