Tune probe selection budget and optional sequential species pass

psq_rag/llm/select.py

@@ -1563,10 +1563,27 @@ def _split_probe_tags_by_bundle(
     return out
 
 
+def _species_cap_from_count_tags(selected_tags: Sequence[str]) -> Optional[int]:
+    """Derive species max-picks cap from count-like probe tags."""
+    s = set(selected_tags)
+    if "zero_pictured" in s:
+        return 0
+    if "solo" in s:
+        return 1
+    if "duo" in s:
+        return 2
+    if "trio" in s:
+        return 3
+    if "group" in s:
+        # Conservative finite cap for open-ended "group".
+        return 4
+    return None
+
+
 def llm_infer_probe_tags(
     query_text: str,
     log=None,
-    *,
+    *,
     temperature: float = 0.0,
     max_tokens: int = 512,
     retries: int = 2,
@@ -1615,36 +1632,101 @@ def llm_infer_probe_tags(
     except Exception:
         split_calls = 1
 
-    probe_chunks = _split_probe_tags_by_bundle(probe_tags, split_calls, log=log)
+    try:
+        # Default probe cap is 2.
+        # Rationale (caption-evident n=10 evals, 2026-03-12):
+        # - max1 reduced pollution but hurt recall too much
+        # - max3 increased pollution and reduced overall F1
+        # - max2 provided the best observed precision/recall tradeoff and highest overall F1
+        max_pick_override = int((os.environ.get("PSQ_PROBE_MAX_PICK_OVERRIDE", "2") or "2").strip())
+    except Exception:
+        max_pick_override = 0
+    if max_pick_override < 0:
+        max_pick_override = 0
+
+    sequential_species = (os.environ.get("PSQ_PROBE_SEQUENTIAL_SPECIES", "0") or "0").strip().lower() in {
+        "1", "true", "yes", "on"
+    }
+
+    def _call_chunks(
+        chunks: Sequence[Sequence[str]],
+        *,
+        label: str,
+        max_pick_cap: Optional[int] = None,
+    ) -> List[str]:
+        out_tags: List[str] = []
+        out_seen: Set[str] = set()
+        for chunk_idx, chunk_tags_seq in enumerate(chunks, start=1):
+            chunk_tags = list(chunk_tags_seq)
+            if not chunk_tags:
+                continue
+            if len(chunks) > 1 and log:
+                log(f"Stage3p {label}: call {chunk_idx}/{len(chunks)} with {len(chunk_tags)} probe tags")
+
+            per_call_budget = len(chunk_tags)
+            if max_pick_override > 0:
+                per_call_budget = min(per_call_budget, max_pick_override)
+            if max_pick_cap is not None:
+                per_call_budget = min(per_call_budget, max(0, int(max_pick_cap)))
+            if per_call_budget <= 0:
+                if log:
+                    log(f"Stage3p {label}: skipping call with budget=0")
+                continue
+
+            out = llm_select_indices(
+                query_text=query_text,
+                candidates=chunk_tags,
+                max_pick=per_call_budget,
+                log=log,
+                retries=retries,
+                mode="single_shot",
+                chunk_size=max(1, len(chunk_tags)),
+                per_phrase_k=max(1, per_call_budget),
+                temperature=temperature,
+                max_tokens=max_tokens,
+                return_metadata=False,
+                return_diagnostics=False,
+                min_why=None,
+                candidate_display=candidate_display,
+                user_template=_get_probe_user_template(),
+                model_override=probe_model_override,
+            )
+            for i in out:
+                if 0 <= i < len(chunk_tags):
+                    t = chunk_tags[i]
+                    if t not in out_seen:
+                        out_seen.add(t)
+                        out_tags.append(t)
+        return out_tags
+
     selected: List[str] = []
-    seen: Set[str] = set()
-    for chunk_idx, chunk_tags in enumerate(probe_chunks, start=1):
-        if split_calls > 1 and log:
-            log(f"Stage3p: call {chunk_idx}/{len(probe_chunks)} with {len(chunk_tags)} probe tags")
-        out = llm_select_indices(
-            query_text=query_text,
-            candidates=chunk_tags,
-            max_pick=len(chunk_tags),
-            log=log,
-            retries=retries,
-            mode="single_shot",
-            chunk_size=max(1, len(chunk_tags)),
-            per_phrase_k=max(1, len(chunk_tags)),
-            temperature=temperature,
-            max_tokens=max_tokens,
-            return_metadata=False,
-            return_diagnostics=False,
-            min_why=None,
-            candidate_display=candidate_display,
-            user_template=_get_probe_user_template(),
-            model_override=probe_model_override,
-        )
-        for i in out:
-            if 0 <= i < len(chunk_tags):
-                t = chunk_tags[i]
-                if t not in seen:
-                    seen.add(t)
+
+    if sequential_species:
+        bundle_by_tag = _probe_bundle_map(log=log)
+        species_tags = [t for t in probe_tags if bundle_by_tag.get(t, "") == "species_taxonomy"]
+        non_species_tags = [t for t in probe_tags if t not in set(species_tags)]
+        if log:
+            log(
+                "Stage3p: sequential species mode on "
+                f"(non_species={len(non_species_tags)}, species={len(species_tags)})"
+            )
+
+        non_species_chunks = _split_probe_tags_by_bundle(non_species_tags, split_calls, log=log) if non_species_tags else []
+        selected_non_species = _call_chunks(non_species_chunks, label="non_species")
+        selected.extend(selected_non_species)
+
+        if species_tags:
+            species_cap = _species_cap_from_count_tags(selected_non_species)
+            if log:
+                log(f"Stage3p: species cap from count tags = {species_cap!r}")
+            species_chunks = [species_tags]
+            selected_species = _call_chunks(species_chunks, label="species", max_pick_cap=species_cap)
+            for t in selected_species:
+                if t not in selected:
                     selected.append(t)
+    else:
+        probe_chunks = _split_probe_tags_by_bundle(probe_tags, split_calls, log=log)
+        selected = _call_chunks(probe_chunks, label="all")
 
     selected = _apply_species_anchor_mapping(selected, query_text=query_text, log=log)