src/agents/final_answer.py

"""
final_answer.py -- Final Answer Agent
========================================
Responsibility: turn verified reasoning notes into the user-facing answer,
with explicit sections for findings, evidence, sources, confidence, and
(for research-oriented questions) research gaps.

Primary path: a real local LLM (flan-t5) is prompted with ONLY the
retrieved+verified evidence text and asked to write the explanation --
never asked to use outside knowledge. If unavailable, a deterministic
extractive fallback assembles the answer directly from the Reasoning
Agent's grounded notes (still not fabricated, just less fluent).

Confidence is computed from measurable signals: sub-question coverage,
proportion of SUPPORTED (vs weak) evidence, and presence of unresolved
contradictions -- not asserted by the LLM.
"""
from __future__ import annotations

from typing import List

from src.agents.evidence import _domain_floor_for_backend
from src.agents.state import ResearchState, TraceEvent
from src.rag.embeddings import get_embedding_backend
from src.utils.llm import get_llm


def _compute_confidence(state: ResearchState) -> tuple[str, float]:
    evidence = state.get("evidence", [])
    if not evidence:
        return "Low", 0.1
    supported = sum(1 for e in evidence if e["relevance_label"] == "SUPPORTED")
    weak = sum(1 for e in evidence if e["relevance_label"] == "WEAKLY_SUPPORTED")
    total_used = supported + weak
    support_ratio = supported / max(1, total_used)

    sub_qs = state.get("sub_questions", [])
    answered = len({e["sub_question"] for e in evidence if e["relevance_label"] in ("SUPPORTED", "WEAKLY_SUPPORTED")})
    coverage = answered / max(1, len(sub_qs))

    contradiction_penalty = 0.15 * len(state.get("contradictions", []))
    verification_bonus = 0.1 if state.get("verification_passed") else -0.1

    score = max(0.0, min(1.0, 0.5 * support_ratio + 0.4 * coverage + verification_bonus - contradiction_penalty))
    if score >= 0.7:
        label = "High"
    elif score >= 0.4:
        label = "Moderate"
    else:
        label = "Low"
    return label, round(score, 2)


def _research_gaps(state: ResearchState) -> List[str]:
    gaps = []
    for c in state.get("contradictions", []):
        gaps.append(
            f"Conflicting reported results between '{c['doc_a']}' and '{c['doc_b']}' are unresolved in "
            f"the literature retrieved -- a controlled replication would help ({c['reason']})."
        )
    for claim in state.get("unsupported_claims", []):
        gaps.append(claim.replace("No sufficiently relevant evidence found for:", "Under-explored in the retrieved corpus:"))
    if not gaps:
        gaps.append("No explicit research gaps were surfaced by the retrieved evidence for this query.")
    return gaps


def _extractive_answer(state: ResearchState) -> str:
    notes = state.get("reasoning_notes", [])
    if not notes:
        return "Insufficient evidence to provide a reliable answer."
    return "\n\n".join(f"- {n}" for n in notes)


def finalize(state: ResearchState) -> ResearchState:
    trace: List[TraceEvent] = list(state.get("trace", []))
    evidence = state.get("evidence", [])
    reasoning_notes = state.get("reasoning_notes", [])

    best_score = max((e["relevance_score"] for e in evidence), default=0.0)
    domain_floor = _domain_floor_for_backend(get_embedding_backend().backend_name)
    on_topic = best_score >= domain_floor

    if not evidence or not any(e["relevance_label"] in ("SUPPORTED", "WEAKLY_SUPPORTED") for e in evidence) or not on_topic:
        state["final_answer"] = (
            "Insufficient evidence to provide a reliable answer. The retrieved corpus does not appear "
            "to contain material relevant to this question (best passage relevance "
            f"{best_score:.2f} vs. required {domain_floor:.2f})."
            if evidence else "Insufficient evidence to provide a reliable answer."
        )
        state["key_findings"] = []
        state["research_gaps"] = ["No usable, on-topic evidence was retrieved for this question."]
        state["confidence"], state["confidence_score"] = "Low", 0.0
        trace.append({"agent": "Final Answer", "message": "No usable on-topic evidence -- returning an honest insufficiency notice."})
        state["trace"] = trace
        return state

    llm = get_llm()
    used_llm = False
    answer_text = None
    if llm.available:
        context = "\n".join(f"- {n}" for n in reasoning_notes)[:2500]
        prompt = (
            "Using ONLY the notes below (do not add outside knowledge), write a clear, "
            "well-organized answer to the research question. Be concise.\n\n"
            f"Question: {state['question']}\n\nNotes:\n{context}\n\nAnswer:"
        )
        generated = llm.generate(prompt, max_new_tokens=250)
        if generated and len(generated.split()) >= 5:
            answer_text = generated
            used_llm = True
    if answer_text is None:
        answer_text = _extractive_answer(state)

    key_findings = [n.split(": ", 1)[0] + ": " + n.split(": ", 1)[1].split(". ")[0] for n in reasoning_notes if ": " in n][:6]

    state["final_answer"] = answer_text
    state["key_findings"] = key_findings
    state["research_gaps"] = _research_gaps(state)
    state["confidence"], state["confidence_score"] = _compute_confidence(state)

    method = "a local instruction-tuned LLM grounded strictly in retrieved evidence" if used_llm else "deterministic extractive synthesis of verified evidence"
    trace.append({
        "agent": "Final Answer",
        "message": f"Generated final answer using {method}. Confidence: {state['confidence']} ({state['confidence_score']}).",
    })
    state["trace"] = trace
    return state