llm/wizard_llm.py
Rama main · réplica de lectura del repositorio original.
# SPDX-License-Identifier: BUSL-1.1
# Copyright (c) 2024-2026 Ali Sadeghi Aghili
"""
llm/wizard_llm.py — LLM abstraction layer for the interactive setup wizard.
Supports two providers:
* ``openai`` — the OpenAI-compatible transport (a local ``gpt-oss``/
vLLM/llama.cpp server, LM Studio, or OpenAI's own
hosted API — selected entirely by ``base_url``)
* ``mock`` — returns empty stubs, no network needed (CI / tests)
All calls enforce a 30-second timeout. On JSON parse failure the call is
retried ONCE with an explicit repair prompt before raising ValueError. The
alias call (:meth:`WizardLLM.generate_aliases`) is the exception: a model that
loops on the same phrases until ``max_tokens`` stops it leaves JSON that cannot
be repaired, so that call keeps the complete strings written before the cut,
caps its own completion length, and when nothing is usable asks again with the
original prompt and a stricter instruction rather than sending the garbage back.
Consolidation note
-------------------
This module used to support four providers (``ollama`` / ``openai`` /
``anthropic`` / ``mock``), each built on :mod:`llm.providers` /
``llm.ollama_backend`` — the single place that transport logic lived,
rather than a second, wizard-only copy of it. The Ollama-specific and
Anthropic transports are gone now (see ``llm/providers.py``'s module
docstring): OpenAI-compatible is this project's only protocol, so the
provider surface shrinks to match. What stays here, because it is
genuinely wizard-specific and not something the production engine needs,
is the JSON-parse-then-repair-and-retry-once loop in
:meth:`WizardLLM.generate` — the wizard talks to the model in a very
different shape (one-shot alias/rule extraction) than the engine's
SQL-generation self-correction loop.
Typical usage::
from llm.wizard_llm import WizardLLM
llm = WizardLLM(provider="openai", model="gpt-oss-20b")
result = llm.generate("Return JSON with key x equal to 1", expect_json=True)
# {"x": 1}
# Override endpoint for a local server:
llm = WizardLLM(provider="openai", model="local-model",
base_url="http://localhost:8000/v1")
"""
from __future__ import annotations
import json
import logging
import os
from dataclasses import dataclass, field
from typing import Any, Union
from core.wizard_drafts import ALIAS_LIMIT, tidy_aliases
from llm.base import LLMBackend
from llm.providers import MockBackend, OpenAIBackend, parse_json_response, salvage_string_list
logger = logging.getLogger(__name__)
#: Completion ceiling, in tokens, for one alias call (``max_tokens``), in place
#: of ``LLM_NUM_PREDICT`` -- which is sized for a SQL statement and, when a
#: model loops, is how long the loop runs. The answer is ``{"aliases": [up to
#: 8 short phrases], "description": "one sentence"}``. Worst case, written out
#: in Persian (8 phrases of 3-4 words, a 130-character sentence): 386
#: characters, 671 UTF-8 bytes; the same in English: 232 characters. At 0.5-0.7
#: tokens per Persian character that is about 190-270 tokens, so 320 fits
#: with room to spare. These are estimates from character counts: no
#: tokenizer for the deployed model was available to count tokens. A cap that
#: cuts an answer anyway costs little, because the complete aliases before the
#: cut are kept (:func:`~llm.providers.salvage_string_list`).
ALIAS_MAX_TOKENS = 320
#: The default of ``LLM_NUM_PREDICT``. A larger value is how the setup guide
#: (section 2.9) tells an operator to make room for a reasoning model, whose
#: thinking counts against ``max_tokens``; see :func:`alias_token_cap`.
_SQL_SIZED_BUDGET = 512
#: Stub JSON returned by the ``mock`` provider — no network needed (CI / tests).
_MOCK_STUB = json.dumps(
{"aliases": [], "description": "", "rules": {"rule_text": ""}, "examples": []}
)
def build_backend(model: str, base_url: str | None = None) -> LLMBackend:
"""Construct the SQL-generation backend used by ``app.py``'s REPL.
A thin, non-router wrapper: unlike ``api/runner.py`` and
``session/engine.py`` (which route every call through
:class:`~llm.router.LLMRouter` for fallback chains and remote-provider
governance), the REPL has always talked to one backend directly — see
``generate_sql`` below. ``base_url``/``api_key`` default to
:mod:`config`'s plain ``OPENAI_*`` settings when not given.
Examples
--------
>>> import config as cfg
>>> with cfg.override_settings(openai_api_key="k"):
... backend = build_backend(model="m")
>>> backend.name
'openai:m'
"""
import config as cfg
return OpenAIBackend(
model=model,
api_key=cfg.settings.openai_api_key,
base_url=base_url or cfg.settings.openai_base_url,
)
def generate_sql(question: str, system_prompt: str) -> str:
"""Generate SQL for *question* using the configured OpenAI-compatible backend.
This function only calls the LLM — it does **not** execute the
generated SQL against any database.
The returned SQL has already been passed through
:func:`~security.sql_guard.ensure_top` (capped at
``cfg.settings.default_top_n`` if the model didn't include its own
row-limit clause) — callers such as ``app.py``'s REPL, which execute
this SQL directly rather than through :class:`~llm.sql_agent.SQLAgent`,
would otherwise get no server-side row cap at all.
Raises
------
ValueError("OUT_OF_SCOPE")
Passed through from the model sentinel.
RuntimeError
When the endpoint is unreachable after all retries.
"""
import config as cfg
from retrieval.context_retriever import ContextRetriever
from prompt_engine.builder import PromptBuilder
from security.sql_guard import clean_sql, ensure_top, validate_sql
from llm.source_routing import choose_source
context = ContextRetriever.retrieve(question)
prompt = PromptBuilder.build(
question=question,
system_prompt=system_prompt,
context=context,
# None (nothing changes) unless several data sources are configured.
source=choose_source(question, context),
)
backend = build_backend(model=cfg.settings.openai_model)
raw = backend.generate(prompt)
sql = clean_sql(raw)
validate_sql(sql)
sql = ensure_top(sql, cfg.settings.default_top_n)
return sql
# ---------------------------------------------------------------------------
# The alias call
# ---------------------------------------------------------------------------
def alias_token_cap(configured: int, extra_body: str = "") -> int | None:
"""The ``max_tokens`` for one alias call, or ``None`` to leave ``LLM_NUM_PREDICT`` alone.
A short answer should not be allowed the whole SQL budget: when the model
loops, the loop runs until that budget is spent. But a reasoning model
spends the same ``max_tokens`` on thinking before it writes a character
of the answer, and a small cap would starve it. The setup guide's two
ways to make room for such a model are to raise ``LLM_NUM_PREDICT`` and to
set a reasoning switch in ``LLM_EXTRA_BODY``; either is taken as a sign
that the operator has made room, and the cap then stays out of the way.
Otherwise the cap is :data:`ALIAS_MAX_TOKENS`, and never more than the
configured value.
Parameters
----------
configured:
``LLM_NUM_PREDICT``.
extra_body:
``LLM_EXTRA_BODY``, as the text of the setting.
Returns
-------
int | None
The ceiling for this call; ``None`` when none should be sent.
Raises
------
Nothing. This function never raises.
Examples
--------
>>> alias_token_cap(512)
320
>>> alias_token_cap(256)
256
>>> alias_token_cap(2048) is None # raised for a reasoning model: no cap
True
>>> alias_token_cap(512, '{"think": false}') is None # a reasoning switch is set: no cap
True
"""
if configured > _SQL_SIZED_BUDGET or extra_body.strip():
return None
return min(configured, ALIAS_MAX_TOKENS)
def _strict_instruction(limit: int) -> str:
"""What is added to the original alias prompt when the first answer was unusable."""
return (
"\n\nYour previous answer could not be used: it repeated itself and never finished. "
f"Answer again with at most {limit} aliases, every one different from the others. "
"Do not repeat a phrase and do not reorder its words. "
"Close the JSON right after the list and write nothing after it."
)
class UnusableAnswerError(ValueError):
"""The model gave no usable alias list, even after one stricter retry.
``str(error)`` says what each answer was, with a short excerpt of the
last one; :attr:`reason` is the one-line version for a progress line.
"""
def __init__(self, message: str, reason: str) -> None:
super().__init__(message)
self.reason = reason
@dataclass
class AliasAnswer:
"""What the alias call produced, and how it got there.
Attributes
----------
aliases:
The aliases to use: text, no repeats, at most the limit
(:func:`~core.wizard_drafts.tidy_aliases`).
answered:
The items of the list as the model wrote them, before that cleaning.
description:
The model's one-sentence description; empty when the answer was cut
off before it.
salvaged:
``True`` when the answer was cut off or malformed and the list was
recovered from its complete strings.
retried:
``True`` when the first answer gave nothing and the second one did.
cut_off:
``True`` when the server stopped the answer at the token limit.
repeated:
``True`` when the salvaged list repeats an item (after folding).
"""
aliases: list[str]
answered: list[Any] = field(default_factory=list)
description: str = ""
salvaged: bool = False
retried: bool = False
cut_off: bool = False
repeated: bool = False
@classmethod
def from_parsed(cls, parsed: Any, limit: int = ALIAS_LIMIT) -> "AliasAnswer":
"""The answer for a reply that parsed as JSON: ``{"aliases": [...], "description": "..."}``.
Parameters
----------
parsed:
What :func:`~llm.providers.parse_json_response` returned. A bare
list is read as the aliases themselves.
limit:
The most aliases to keep.
Returns
-------
AliasAnswer
With nothing salvaged; an answer of another shape has no aliases.
Raises
------
Nothing. This function never raises.
Examples
--------
>>> AliasAnswer.from_parsed({"aliases": ["a", "A", "b"], "description": "d"}).aliases
['a', 'b']
>>> AliasAnswer.from_parsed(["x"]).aliases, AliasAnswer.from_parsed({"aliases": "x"}).aliases
(['x'], [])
"""
listed = parsed.get("aliases") if isinstance(parsed, dict) else parsed
answered = list(listed) if isinstance(listed, list) else []
description = parsed.get("description") if isinstance(parsed, dict) else ""
return cls(
aliases=tidy_aliases(answered, limit),
answered=answered,
description=description if isinstance(description, str) else "",
)
def _why_unusable(raw: str, meta: dict[str, Any]) -> str:
"""One short phrase for an answer that gave no alias."""
if meta.get("finish_reason") == "length":
what = "cut off at the token limit"
elif not raw.strip():
return "empty"
else:
what = "not valid JSON"
return f"{what}, {len(raw)} characters"
# ---------------------------------------------------------------------------
# Public WizardLLM class
# ---------------------------------------------------------------------------
class WizardLLM:
"""Unified LLM interface supporting openai and mock.
A thin wrapper around a single :class:`~llm.base.LLMBackend` instance
(built from :mod:`llm.providers` — see the module docstring) that adds
the wizard's own JSON-parse-and-retry contract on top of
``backend.generate()``.
Parameters
----------
provider:
One of ``"openai"``, ``"mock"``.
model:
Model identifier, e.g. ``"gpt-oss-20b"``.
base_url:
Optional override for the provider's default endpoint. Useful for
LM Studio, vLLM, or a custom server.
Raises
------
ValueError
If *provider* is not one of the two supported values, or if a
required API key environment variable is missing.
"""
_SUPPORTED = frozenset({"openai", "mock"})
def __init__(
self,
provider: str,
model: str,
base_url: str | None = None,
) -> None:
provider = provider.strip().lower()
if provider not in self._SUPPORTED:
raise ValueError(
f"Unsupported provider {provider!r}. "
f"Choose one of: {', '.join(sorted(self._SUPPORTED))}."
)
self.provider = provider
self.model = model
self._backend: LLMBackend = self._build_backend(provider, model, base_url)
@staticmethod
def _build_backend(provider: str, model: str, base_url: str | None) -> LLMBackend:
"""Construct the underlying :class:`~llm.base.LLMBackend` for *provider*."""
if provider == "openai":
key = os.getenv("OPENAI_API_KEY", "")
if not key:
raise ValueError(
"OPENAI_API_KEY environment variable is required "
"for provider='openai'."
)
url = base_url or os.getenv("WIZARD_LLM_BASE_URL", "https://api.openai.com/v1")
return OpenAIBackend(model=model, api_key=key, base_url=url)
return MockBackend(response=_MOCK_STUB) # provider == "mock"
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def test_connection(self) -> bool:
"""Return True if the configured backend is reachable, False otherwise."""
return self._backend.test_connection()
def generate(
self,
prompt: str,
expect_json: bool = True,
) -> Union[str, dict, list]:
"""Send *prompt* to the LLM and return the response.
Parameters
----------
prompt:
The full prompt string.
expect_json:
When True (default), parse the response as JSON and return a
``dict`` or ``list``. On the first parse failure the call is
retried ONCE with a repair prompt prepended. Raises
``ValueError`` if the retry also fails.
When False, the raw response string is returned as-is.
Returns
-------
str | dict | list
Raw string when ``expect_json=False``, otherwise parsed JSON.
Raises
------
requests.Timeout
If the backend does not respond within 30 seconds.
requests.HTTPError
On non-2xx HTTP responses.
ValueError
If ``expect_json=True`` and valid JSON cannot be extracted
after the retry.
"""
raw = self._backend.generate(prompt)
if not expect_json:
return raw
# --- first extraction attempt ---
try:
return parse_json_response(raw)
except ValueError:
logger.debug(
"[WizardLLM] First JSON parse failed for provider=%s; retrying.",
self.provider,
)
# --- single retry with explicit repair instruction ---
repair_prompt = (
"The text below must be a valid JSON object or array but failed to parse.\n"
"Return ONLY the corrected JSON — no explanation, no markdown, no code fences.\n\n"
+ raw
)
raw2 = self._backend.generate(repair_prompt)
try:
return parse_json_response(raw2)
except ValueError as exc:
raise ValueError(
f"[WizardLLM] Could not extract valid JSON after retry "
f"(provider={self.provider!r}, model={self.model!r}).\n"
f"Last raw response (first 500 chars): {raw2[:500]!r}"
) from exc
def generate_aliases(self, prompt: str, *, limit: int = ALIAS_LIMIT) -> AliasAnswer:
"""Ask for the aliases of one table; survive an answer that loops and is cut off.
The reply should be ``{"aliases": [...], "description": "..."}``. A
model that repeats the same phrases until ``max_tokens`` stops it
leaves JSON that never closes, and :meth:`generate` would send that
text back to be "corrected", which loops again. Here:
1. The call carries its own, small ``max_tokens`` (:func:`alias_token_cap`).
2. A reply that does not parse is searched for the complete strings of
its ``aliases`` list (:func:`~llm.providers.salvage_string_list`);
those are kept, deduplicated and capped at *limit*.
3. Only when that gives nothing, the model is asked ONCE more with the
original prompt and a stricter instruction. The text of the first
answer is never sent back.
Parameters
----------
prompt:
The alias prompt.
limit:
The most aliases to keep.
Returns
-------
AliasAnswer
The aliases and how they were obtained.
Raises
------
UnusableAnswerError
If neither answer held a complete alias (a ``ValueError``).
requests.Timeout
If the backend does not respond in time.
requests.HTTPError
On a non-2xx response.
RuntimeError
If the endpoint cannot be reached.
Examples
--------
>>> llm = WizardLLM(provider="mock", model="mock")
>>> llm.generate_aliases("anything").aliases
[]
"""
import config as cfg # noqa: PLC0415
cap = alias_token_cap(cfg.settings.llm_num_predict, cfg.settings.llm_extra_body_json or "")
raw, meta = self._ask(prompt, cap)
answer = self._read_aliases(raw, meta, limit)
if answer is not None:
return answer
logger.debug(
"[WizardLLM] No usable alias list (%s); asking once more with a stricter prompt.",
_why_unusable(raw, meta),
)
raw2, meta2 = self._ask(prompt + _strict_instruction(limit), cap)
answer = self._read_aliases(raw2, meta2, limit)
if answer is not None:
answer.retried = True
return answer
first, second = _why_unusable(raw, meta), _why_unusable(raw2, meta2)
hint = ""
if meta.get("reasoning_detected") or (meta.get("reasoning_tokens") or 0) > 0:
hint = " The model spent its tokens on reasoning; see LLM_EXTRA_BODY and LLM_NUM_PREDICT."
raise UnusableAnswerError(
f"[WizardLLM] No usable alias list (provider={self.provider!r}, model={self.model!r}): "
f"first answer {first}; second answer, after a stricter prompt, {second}.{hint} "
f"Start of the last answer: {raw2[:80]!r}",
reason=f"no usable answer: {first}; again {second}",
)
def _ask(self, prompt: str, max_tokens: int | None) -> tuple[str, dict[str, Any]]:
"""One backend call, with a per-call ``max_tokens`` only when there is one."""
if max_tokens is None:
return self._backend.generate_with_meta(prompt)
return self._backend.generate_with_meta(prompt, max_tokens=max_tokens)
@staticmethod
def _read_aliases(raw: str, meta: dict[str, Any], limit: int) -> AliasAnswer | None:
"""The alias list in *raw*: parsed whole, else salvaged; ``None`` when there is none.
A reply that parses is taken as it is (an empty list is an answer). A
reply that does not is searched for the complete strings of its list.
"""
cut_off = meta.get("finish_reason") == "length"
try:
answer = AliasAnswer.from_parsed(parse_json_response(raw), limit)
answer.cut_off = cut_off
return answer
except ValueError:
pass
items = salvage_string_list(raw, "aliases")
if not items:
return None
return AliasAnswer(
aliases=tidy_aliases(items, limit),
answered=items,
salvaged=True,
cut_off=cut_off,
repeated=len(tidy_aliases(items, len(items))) < len(items),
)
# ------------------------------------------------------------------
# Convenience factory
# ------------------------------------------------------------------
@classmethod
def from_config(cls) -> "WizardLLM":
"""Construct a WizardLLM from the project's config / environment.
Reads (in order of precedence):
* ``WIZARD_LLM_PROVIDER`` — default ``openai``
* ``WIZARD_LLM_MODEL`` — default ``gpt-4o-mini`` (or ``OPENAI_MODEL``)
* ``WIZARD_LLM_BASE_URL`` — optional endpoint override
"""
try:
import config as cfg # noqa: PLC0415
default_model = cfg.settings.openai_model
except Exception: # noqa: BLE001
default_model = "gpt-4o-mini"
return cls(
provider=os.getenv("WIZARD_LLM_PROVIDER", "openai"),
model=os.getenv("WIZARD_LLM_MODEL", default_model),
base_url=os.getenv("WIZARD_LLM_BASE_URL") or None,
)