diff --git a/atlasbot/engine/answerer.py b/atlasbot/engine/answerer.py index ed83f02..61b145b 100644 --- a/atlasbot/engine/answerer.py +++ b/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): diff --git a/atlasbot/llm/prompts.py b/atlasbot/llm/prompts.py index d3accf1..0bacee4 100644 --- a/atlasbot/llm/prompts.py +++ b/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. "