atlasbot: validate metric key coverage

atlasbot/engine/answerer.py

@@ -1313,12 +1313,59 @@ async def _select_metric_chunks(
         selected = _merge_metric_keys(selected, candidate_keys, max_keys)
     if not selected and candidate_keys:
         selected = candidate_keys[:max_keys]
+    if available_keys:
+        missing = await _validate_metric_keys(
+            call_llm,
+            question,
+            sub_questions,
+            selected,
+            available_keys,
+            plan,
+        )
+        if missing:
+            selected = _merge_metric_keys(selected, missing, max_keys)
     if not selected:
         return [], []
     ids = _chunk_ids_for_keys(chunks, selected)
     return selected, ids
 
 
+async def _validate_metric_keys(
+    call_llm: Callable[..., Awaitable[str]],
+    question: str,
+    sub_questions: list[str],
+    selected: list[str],
+    available: list[str],
+    plan: ModePlan,
+) -> list[str]:
+    if not available:
+        return []
+    cap = max(12, plan.max_subquestions * 4)
+    available_list = available[:cap]
+    prompt = prompts.METRIC_KEYS_VALIDATE_PROMPT.format(
+        question=question,
+        sub_questions=json.dumps(sub_questions),
+        selected=json.dumps(selected),
+        available="\n".join(available_list),
+    )
+    raw = await call_llm(
+        prompts.METRIC_KEYS_VALIDATE_SYSTEM,
+        prompt,
+        model=plan.fast_model,
+        tag="metric_keys_validate",
+    )
+    parsed = _parse_json_block(raw, fallback={})
+    items = parsed.get("missing") if isinstance(parsed, dict) else []
+    if not isinstance(items, list):
+        return []
+    available_set = set(available_list)
+    out: list[str] = []
+    for item in items:
+        if isinstance(item, str) and item in available_set and item not in out:
+            out.append(item)
+    return out
+
+
 def _metric_ctx_values(ctx: dict[str, Any]) -> tuple[list[str], str, list[str], list[str], set[str]]:
     summary_lines = ctx.get("summary_lines") if isinstance(ctx, dict) else None
     if not isinstance(summary_lines, list):

atlasbot/llm/prompts.py

@@ -83,6 +83,22 @@ METRIC_KEYS_PROMPT = (
     "Limit to at most {max_keys} keys."
 )
 
+METRIC_KEYS_VALIDATE_SYSTEM = (
+    CLUSTER_SYSTEM
+    + " Verify whether selected metric keys cover the question. "
+    + "Only reference keys from the provided list. "
+    + "Return JSON with field: missing (list)."
+)
+
+METRIC_KEYS_VALIDATE_PROMPT = (
+    "Question: {question}\n"
+    "SubQuestions: {sub_questions}\n"
+    "SelectedKeys: {selected}\n\n"
+    "AvailableKeys:\n{available}\n\n"
+    "Return JSON with field: missing (list). "
+    "List any keys from AvailableKeys that are needed but missing."
+)
+
 TOOL_SYSTEM = (
     CLUSTER_SYSTEM
     + " Suggest a safe, read-only command that could refine the answer. "