app/Planner.py

"""Stage 2 of the pipeline: pure Python decision logic. No LLM call happens here.

This is the core fix for the 'rigid' / hallucinated behavior you were seeing: the
old single prompt asked a 3B model to *decide* — in the same breath as writing
prose — whether it had enough info, whether the user was confirming, whether a
question was a comparison, etc. Small models are unreliable at that kind of
in-context judgment call. Here, the judgment call is a plain if/else over facts
the NLU stage already extracted, so it's 100% consistent and free.
"""
from typing import List, Dict, Any


class Action:
    def __init__(self, kind: str, **kwargs):
        self.kind = kind
        self.__dict__.update(kwargs)


def _to_dict(rec) -> dict:
    if hasattr(rec, "model_dump"):
        return rec.model_dump()
    if hasattr(rec, "dict"):
        return rec.dict()
    return rec


def find_last_recommendations(conversation) -> List[Dict[str, Any]]:
    """Scans the raw conversation (newest first) for the last Agent turn that
    carried a non-empty recommendations list. Requires the frontend to echo
    `recommendations` back on Agent turns (see schemas.Turn / main.py) — this is
    what lets a 'confirmation' turn reuse the REAL prior shortlist instead of the
    model having to recall or reconstruct it from plain text."""
    for turn in reversed(conversation):
        if turn.role.lower() == "agent" and getattr(turn, "recommendations", None):
            return [_to_dict(r) for r in turn.recommendations]
    return []


def _last_agent_message(conversation):
    for turn in reversed(conversation):
        if turn.role.lower() == "agent":
            return turn.content.strip()
    return None


def decide(state: dict, conversation) -> Action:
    if not state.get("in_scope", True):
        return Action("off_topic")

    intent = state.get("intent", "new_request")
    prior_recs = find_last_recommendations(conversation)

    if intent == "off_topic":
        return Action("off_topic")

    if intent == "comparison_question":
        return Action("compare", compared_tests=state.get("compared_tests", []),
                       topic_keywords=state.get("compared_tests", []) or state.get("topic_keywords", []))

    if intent == "confirmation":
        if prior_recs:
            return Action("close", reuse_recommendations=prior_recs)
        # Nothing to confirm yet (e.g. "thanks" mid-clarification) — keep going.
        intent = "new_request"

    if intent == "pushback_feedback":
        return Action(
            "redirect",
            role_summary=state.get("role_summary"),
            purpose=state.get("purpose"),
            topic_keywords=state.get("topic_keywords", []),
            prior_recommendations=prior_recs,
        )

    # --- Smart Update Detection ---
    # If we already have prior recommendations, any new requirement is an update to 
    # the existing list, even if the NLU missed the strict 'addition_removal' intent.
    updating_flag = False
    if intent == "addition_removal":
        updating_flag = True
    elif intent in ["new_request", "clarifying_answer"] and len(prior_recs) > 0:
        updating_flag = True

    # --- THE AMNESIA & LOOP FIX ---
    # If we already have a shortlist from the conversation history, we know we are ready to recommend.
    # This overrides the NLU if it "forgets" the role summary mid-conversation.
    if updating_flag and prior_recs:
        ready = True
    elif intent == "clarifying_answer":
        # If the user just gave us more info, stop asking and just recommend!
        ready = True 
    else:
        ready = bool(state.get("ready_to_recommend")) and bool(state.get("role_summary"))

    candidate_question = (
        state.get("missing_info_question")
        or "Could you tell me more about the role and the purpose of this assessment?"
    )

    # LOOP-BREAKER: if the question we're about to ask is the same one already
    # asked last turn, the NLU stage isn't making progress on it. 
    # Proceed with whatever is known instead of asking a third time.
    if not ready:
        last_agent_msg = _last_agent_message(conversation)
        if last_agent_msg and candidate_question.strip().lower() == last_agent_msg.strip().lower():
            ready = bool(state.get("role_summary"))  # proceed if we at least know who this is for

    if ready:
        return Action("recommend", updating=updating_flag, prior_recommendations=prior_recs,
                       topic_keywords=state.get("topic_keywords", []))

    return Action("ask_question", question=candidate_question)