diff --git a/atlasbot/engine/answerer.py b/atlasbot/engine/answerer.py index 575b7e3..e5c82c0 100644 --- a/atlasbot/engine/answerer.py +++ b/atlasbot/engine/answerer.py @@ -340,6 +340,28 @@ class AnswerEngine: model=plan.model, tag="evidence_fix", ) + if runbook_paths and resolved_runbook and _needs_runbook_reference(normalized, runbook_paths, reply): + if observer: + observer("runbook_enforce", "enforcing runbook path") + enforce_prompt = prompts.RUNBOOK_ENFORCE_PROMPT.format(path=resolved_runbook) + reply = await call_llm( + prompts.RUNBOOK_ENFORCE_SYSTEM, + enforce_prompt + "\nAnswer: " + reply, + context=context, + model=plan.model, + tag="runbook_enforce", + ) + + if _needs_focus_fix(normalized, reply, classify): + if observer: + observer("focus_fix", "tightening answer") + reply = await call_llm( + prompts.EVIDENCE_FIX_SYSTEM, + prompts.FOCUS_FIX_PROMPT + "\nQuestion: " + normalized + "\nDraft: " + reply, + context=context, + model=plan.model, + tag="focus_fix", + ) if plan.use_critic: if observer: @@ -877,6 +899,18 @@ def _needs_evidence_fix(reply: str, classify: dict[str, Any]) -> bool: return False +def _needs_focus_fix(question: str, reply: str, classify: dict[str, Any]) -> bool: + if not reply: + return False + q_lower = (question or "").lower() + if classify.get("question_type") not in {"metric", "diagnostic"} and not re.search(r"\b(how many|list|count)\b", q_lower): + return False + if reply.count(".") <= 1: + return False + extra_markers = ("for more", "if you need", "additional", "based on") + return any(marker in reply.lower() for marker in extra_markers) + + def _extract_keywords(normalized: str, sub_questions: list[str], keywords: list[Any] | None) -> list[str]: stopwords = { "the", diff --git a/atlasbot/llm/prompts.py b/atlasbot/llm/prompts.py index 2fc9ac3..bbb5dd6 100644 --- a/atlasbot/llm/prompts.py +++ b/atlasbot/llm/prompts.py @@ -101,7 +101,20 @@ EVIDENCE_FIX_PROMPT = ( "rewrite the answer to include those values. " "If data is truly missing, keep the draft concise and honest. " "If AllowedRunbooks are provided, use an exact path from that list when answering " - "documentation or checklist questions and do not invent new paths." + "documentation or checklist questions and do not invent new paths. " + "If ResolvedRunbook is provided, you must include that exact path and must not say it is missing." +) + +RUNBOOK_ENFORCE_SYSTEM = ( + CLUSTER_SYSTEM + + " Ensure the answer includes the required runbook path. " + + "Return a corrected answer only." +) + +RUNBOOK_ENFORCE_PROMPT = ( + "Rewrite the answer so it explicitly cites the required runbook path. " + "If the answer already includes it, keep it. " + "Required path: {path}." ) RUNBOOK_SELECT_SYSTEM = ( @@ -130,6 +143,11 @@ CRITIC_PROMPT = ( "Return JSON with fields: issues (list), missing_data (list), risky_claims (list)." ) +FOCUS_FIX_PROMPT = ( + "Rewrite the answer to be concise and directly answer the question. " + "Remove tangential details and speculative statements." +) + REVISION_SYSTEM = ( CLUSTER_SYSTEM + " Revise the answer based on critique. "