llm/source_routing.py
Rama main · réplica de lectura del repositorio original.
# SPDX-License-Identifier: BUSL-1.1
# Copyright (c) 2024-2026 Ali Sadeghi Aghili
"""Per-request data-source routing shared by both engines.
With several data sources configured, each question is first routed to one
source (:mod:`retrieval.source_selector`) and its prompt is built for that
source alone (:func:`llm.router.build_prompt_segments` with ``source=``).
This module is the glue the two engines -- ``api/runner.py`` /
:class:`~llm.sql_agent.SQLAgent` for ``/query`` and ``session/engine.py``
for conversations -- share, so they cannot drift apart:
* :class:`SourceRouting` holds one request's state: the selection, the
source currently in use, the prompt built for it, and whether the one
permitted retry has been used.
* :func:`generate_with_source_fallback` makes the model call and, when the
model answers ``OUT_OF_SCOPE``, retries once with the next candidate
source. **At most one extra attempt per request**: the second
``OUT_OF_SCOPE`` is raised as it always was, and a deployment with one
source never retries (``plan`` returns ``None`` for it, and nothing here
runs).
* :meth:`SourceRouting.audit` is the ``datasource_selection`` block of the
audit record.
Why the retry exists: source selection is a heuristic (keywords, retrieval
evidence), and a model shown the wrong source's tables can only say it has
nothing to answer with. The retry costs one more model call, only for that
case, and only when another source exists.
The statement the model then writes is still checked by the SQL guard and
routed by the tables it reads (:mod:`database.routing`); the selection
decides what the model *sees*, never where a statement runs.
"""
from __future__ import annotations
import logging
from typing import Any, Callable
from llm.router import PromptSegments, RouteResult, _is_out_of_scope
from retrieval.source_selector import (
REASON_SESSION,
SourceSelection,
select_source_for_question,
)
logger = logging.getLogger(__name__)
__all__ = [
"SourceRouting",
"choose_source",
"continuity_audit",
"generate_with_source_fallback",
"is_out_of_scope_failure",
]
def is_out_of_scope_failure(exc: BaseException) -> bool:
"""True when *exc*, or the exception it wraps, is the model's ``OUT_OF_SCOPE``.
The router re-raises a decline unwrapped, but a caller further out may
see it as the ``__cause__`` of a chain-exhausted ``RuntimeError``; both
shapes count.
Examples
--------
>>> is_out_of_scope_failure(ValueError("OUT_OF_SCOPE"))
True
>>> wrapper = RuntimeError("Every backend in the chain failed")
>>> wrapper.__cause__ = ValueError("OUT_OF_SCOPE")
>>> is_out_of_scope_failure(wrapper)
True
>>> is_out_of_scope_failure(ValueError("something else"))
False
"""
return _is_out_of_scope(exc.__cause__ or exc)
class SourceRouting:
"""One request's data-source routing: selection, prompt, and retry state.
Build one with :meth:`plan`; construct directly only in tests.
Parameters
----------
selection:
What :func:`retrieval.source_selector.select_source` decided.
build:
``build(source) -> PromptSegments``: assembles the request's prompt
for that source (the caller's closure over its question, system
prompt, retrieval context and session context).
Attributes
----------
selection:
The selection as made.
current:
The source the prompt in :attr:`segments` is for.
fallback_from:
The source that answered ``OUT_OF_SCOPE`` when the retry happened,
else ``None``.
segments:
The base prompt for :attr:`current` (the caller appends any
correction text to its question segment, never to the prefix).
Examples
--------
>>> sel = SourceSelection("sales", "keyword", ("sales", "inventory"))
>>> built = []
>>> def build(source):
... built.append(source)
... return PromptSegments(static_prefix=f"prefix for {source}", question="q")
>>> routing = SourceRouting(sel, build)
>>> routing.current, routing.segments.static_prefix
('sales', 'prefix for sales')
>>> routing.fall_back().static_prefix
'prefix for inventory'
>>> routing.audit()
{'chosen': 'inventory', 'reason': 'keyword', 'candidates': ['sales', 'inventory'], 'fallback_from': 'sales'}
>>> routing.can_fall_back()
False
>>> built
['sales', 'inventory']
"""
def __init__(
self, selection: SourceSelection, build: Callable[[str], PromptSegments],
) -> None:
self.selection = selection
self._build = build
self.current = selection.chosen
self.fallback_from: str | None = None
self.segments = build(selection.chosen)
@classmethod
def plan(
cls,
question: str,
context: object,
*,
build: Callable[[str], PromptSegments],
previous_source: str | None = None,
) -> "SourceRouting | None":
"""Select a source for *question* and build its prompt.
Parameters
----------
question:
The natural-language question.
context:
The retrieval context over the whole schema.
build:
See the class docstring.
previous_source:
The previous turn's source, for a follow-up turn.
Returns
-------
SourceRouting | None
``None`` when fewer than two data sources are configured: the
caller then builds its prompt exactly as it always did.
"""
selection = select_source_for_question(
question, context, previous_source=previous_source,
)
if selection is None:
return None
logger.debug(
"Data source %r selected (%s); candidates %s",
selection.chosen, selection.reason, selection.candidates,
)
return cls(selection, build)
def can_fall_back(self) -> bool:
"""Whether an ``OUT_OF_SCOPE`` may still be retried on another source."""
return self.fallback_from is None and self.selection.after(self.current) is not None
def fall_back(self) -> PromptSegments:
"""Switch to the next candidate source and return its prompt.
Raises
------
RuntimeError
If no retry is left (:meth:`can_fall_back` is false).
"""
following = self.selection.after(self.current) if self.fallback_from is None else None
if following is None:
raise RuntimeError("no data source left to fall back to")
logger.info(
"Model answered OUT_OF_SCOPE for data source %r; retrying once with %r",
self.current, following,
)
self.fallback_from = self.current
self.current = following
self.segments = self._build(following)
return self.segments
def audit(self) -> dict[str, Any]:
"""The ``datasource_selection`` block for the audit record."""
return self.selection.audit(chosen=self.current, fallback_from=self.fallback_from)
def generate_with_source_fallback(
generate: Callable[[PromptSegments], RouteResult],
routing: SourceRouting | None,
segments: PromptSegments,
) -> tuple[RouteResult, PromptSegments | None]:
"""Call *generate* with *segments*; on ``OUT_OF_SCOPE`` retry once on the next source.
Parameters
----------
generate:
The model call: ``generate(segments) -> RouteResult``. May raise
anything the router raises.
routing:
The request's :class:`SourceRouting`, or ``None`` for a
single-source deployment, in which case this is just
``generate(segments)``.
segments:
The prompt for this attempt (the routing's base prompt, plus any
correction text the caller appended to its question segment).
Returns
-------
tuple[RouteResult, PromptSegments | None]
The model's result, and ``None`` -- or, when the retry was used and
answered, the new source's base prompt. A caller that keeps
correction state must restart it from that prompt: the corrections
it accumulated describe tables of the other source.
Raises
------
Exception
Whatever *generate* raises, unchanged. In particular the second
``OUT_OF_SCOPE`` (the retry's) propagates, so a request costs at
most one extra model call and still ends in ``OUT_OF_SCOPE`` when
no source can answer.
Examples
--------
>>> sel = SourceSelection("sales", "default", ("sales", "inventory"))
>>> routing = SourceRouting(sel, lambda s: PromptSegments(static_prefix=s, question="q"))
>>> calls = []
>>> def generate(segments):
... calls.append(segments.static_prefix)
... if segments.static_prefix == "sales":
... raise ValueError("OUT_OF_SCOPE")
... return RouteResult(text="SELECT 1", structured=None, meta={}, provider="p", fallback_used=False)
>>> result, replaced = generate_with_source_fallback(generate, routing, routing.segments)
>>> result.text, replaced.static_prefix, calls
('SELECT 1', 'inventory', ['sales', 'inventory'])
With no routing there is no retry:
>>> generate_with_source_fallback(generate, None, PromptSegments(static_prefix="sales", question="q"))
Traceback (most recent call last):
...
ValueError: OUT_OF_SCOPE
"""
try:
return generate(segments), None
except Exception as exc: # noqa: BLE001 - only OUT_OF_SCOPE is handled here
if routing is None or not is_out_of_scope_failure(exc) or not routing.can_fall_back():
raise
# The retry. Its own failure -- a second OUT_OF_SCOPE included -- is not
# caught: one extra attempt, never more.
replacement = routing.fall_back()
return generate(replacement), replacement
def choose_source(
question: str, context: object, *, previous_source: str | None = None,
) -> str | None:
"""The source to build *question*'s prompt for, or ``None`` for a single source.
For the callers that make one model call and keep no retry state (the
evaluation harness, the REPL's :mod:`llm.wizard_llm`): they build their
prompt with ``source=choose_source(...)``.
Examples
--------
>>> import config as cfg
>>> from core.models import RetrievalContext
>>> with cfg.override_settings(project_config_dir="/nonexistent"):
... choose_source("anything", RetrievalContext()) is None
True
"""
selection = select_source_for_question(
question, context, previous_source=previous_source,
)
return None if selection is None else selection.chosen
def continuity_audit(previous_source: str | None) -> dict[str, Any] | None:
"""The ``datasource_selection`` of a turn that inherits the previous turn's source.
A refinement composed over the previous turn's SQL (``session.composer``)
reads whatever that statement read, so it stays on the previous source
without any selection.
Parameters
----------
previous_source:
The previous turn's source, or ``None`` when unknown.
Returns
-------
dict | None
``{"chosen": previous_source, "reason": "session", "candidates":
[previous_source], "fallback_from": None}``; ``None`` when
*previous_source* is ``None`` or fewer than two sources are
configured (the single-source audit record has no selection).
Examples
--------
>>> continuity_audit(None) is None
True
>>> from unittest.mock import patch
>>> with patch("database.datasources.datasource_names", return_value=("sales", "inventory")):
... continuity_audit("inventory")
{'chosen': 'inventory', 'reason': 'session', 'candidates': ['inventory'], 'fallback_from': None}
"""
if previous_source is None:
return None
import database.datasources as datasources
if len(datasources.datasource_names()) < 2:
return None
return SourceSelection(
previous_source, REASON_SESSION, (previous_source,),
).audit()