workbench_grounded_theory/supervisor.py

# ============================================================================
# supervisor.py — routing node for the Grounded Theory workbench
# ============================================================================
#
# COMPLIANCE
# ----------
# The supervisor's entire job is to let the LLM decide which node runs
# next. This file used to have Python guards that overrode the LLM's
# decision when it picked an out-of-order step. Those guards are gone.
# The LLM decides. If it misroutes, the prompt gets fixed.
#
# There is no MAX_ITERATIONS check. LangGraph's own recursion_limit
# (set in graph.run) is the single source of truth for loop termination.
# ============================================================================

import providers
from .prompts import SUPERVISOR_PROMPT


def supervisor_node(state):
    iteration = state.get("iteration", 0) + 1

    prompt = SUPERVISOR_PROMPT.format(
        user_message=state["user_message"],
        detection_done=state.get("detection_result") is not None,
        refinement_done=state.get("refinement_result") is not None,
        confirmation_done=state.get("confirmation_result") is not None,
    )

    client = providers.get_llm_client(state["llm_provider"], state["llm_key"])
    model_name = providers.get_llm_model(state["llm_provider"])

    resp = client.chat.complete(
        model=model_name,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
        max_tokens=20,
    )
    text = (resp.choices[0].message.content or "").strip().lower()

    # Parse the first matching keyword. If none match, action stays as
    # "respond" — a sensible default that means "we're done".
    keywords = {"detect": "detect", "refine": "refine",
                "confirm": "confirm", "respond": "respond"}
    action = next((v for k, v in keywords.items() if k in text), "respond")

    return {
        "next_action": action,
        "iteration": iteration,
        "steps": [{
            "step": iteration,
            "node": "supervisor",
            "action": f"route -> {action}",
            "detail": text[:60],
        }],
    }


def route_from_supervisor(state):
    """Conditional edge function. Maps action keyword to node name."""
    mapping = {
        "detect": "pattern_detection",
        "refine": "pattern_refinement",
        "confirm": "pattern_confirmation",
    }
    return mapping.get(state["next_action"], "respond")