fix: preserve requested sources in quick plans

This commit is contained in:
Hiten Shah
2026-05-17 15:54:45 -07:00
parent 850c7e0185
commit 444e07d141
2 changed files with 50 additions and 2 deletions
+26 -2
View File
@@ -274,7 +274,15 @@ def _sanitize_plan(
freshness_mode=freshness_mode,
cluster_mode=cluster_mode,
raw_topic=topic,
subqueries=_normalize_subquery_weights(_trim_subqueries_for_depth(subqueries, intent, depth, eligible_sources)),
subqueries=_normalize_subquery_weights(
_trim_subqueries_for_depth(
subqueries,
intent,
depth,
eligible_sources,
requested_sources=requested_sources,
)
),
source_weights=source_weights,
notes=[str(note).strip() for note in raw.get("notes") or [] if str(note).strip()],
)
@@ -307,6 +315,7 @@ def _trim_subqueries_for_depth(
intent: str,
depth: str,
available_sources: list[str],
requested_sources: list[str] | None = None,
) -> list[schema.SubQuery]:
# At non-quick depth, expand sources: use capability routing for intents
# that define it, or all available sources otherwise. The LLM planner may
@@ -336,6 +345,15 @@ def _trim_subqueries_for_depth(
for subquery in subqueries:
if depth in {"quick", "default"}:
preferred_sources = ranked_sources[:limit]
if requested_sources:
requested = [
source
for source in requested_sources
if source in available_sources and source in subquery.sources
]
for source in requested:
if source not in preferred_sources:
preferred_sources.append(source)
else:
preferred_sources = [source for source in ranked_sources if source in subquery.sources][:limit]
if len(preferred_sources) < limit:
@@ -428,7 +446,13 @@ def _fallback_plan(
cluster_mode=_default_cluster_mode(intent),
raw_topic=topic,
subqueries=_normalize_subquery_weights(
_trim_subqueries_for_depth(subqueries[:_max_subqueries(intent, topic)], intent, depth, list(source_weights))
_trim_subqueries_for_depth(
subqueries[:_max_subqueries(intent, topic)],
intent,
depth,
list(source_weights),
requested_sources=requested_sources,
)
),
source_weights=_normalize_weights(source_weights),
notes=[note],