app/Nlu.py

from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from app.schemas import ConversationState
from app.config import HUGGINGFACEHUB_API_TOKEN
from app.Llm_utils import invoke_with_retry

class NLUStage:
    """Stage 1: Reads the conversation and extracts structured facts. 
    Never recommends tests or writes user-facing prose.
    """
    def __init__(self):
        # Using Llama-3.1-8B-Instruct
        self.llm = ChatOpenAI(
            model="meta-llama/Llama-3.1-8B-Instruct:novita",
            api_key=HUGGINGFACEHUB_API_TOKEN,
            base_url="https://router.huggingface.co/v1",
            max_tokens=500,
            temperature=0.0,
        )
        self.parser = JsonOutputParser(pydantic_object=ConversationState)
        self.prompt = PromptTemplate(
            template="""You read a conversation between a user and an SHL assessment consultant and extract facts. You do NOT answer the user, recommend tests, or write any reply.

SCOPE CHECK: set "in_scope" to false ONLY if the Latest User Message is completely unrelated to business, hiring, job roles, or SHL assessments. Vague statements like "We need a solution for senior leadership" ARE in-scope.

INTENT — pick exactly one:
- "new_request": describes a hiring need not yet discussed in this conversation.
- "clarifying_answer": answers a question the consultant just asked.
- "comparison_question": asks to compare or explain the difference between named tests.
- "confirmation": says an existing shortlist/answer is right and they're satisfied.
- "addition_removal": asks to add, drop, or swap a specific skill or test.
- "pushback_feedback": says the consultant's last reply was wrong, OR asks regulatory/legal compliance questions about a test's legality.
- "off_topic": completely unrelated to business/hiring.

Extract from the FULL conversation so far:
- role_summary: a short phrase describing who this is for, or null if truly unknown.
- purpose: selection vs. development, or another explicit purpose the user gave, or null if not stated.
- topic_keywords: list of skills, tech stack, domain, or role-type words useful for a catalog search.
- compared_tests: test names the user wants compared, else [].
- ready_to_recommend: true if the user provided BOTH a job role AND at least one specific requirement (skills, traits, industry context, or test type). 
- CONSULTANT BEHAVIOR: Do not withhold recommendations to get a "perfect" request. If the user provides a role and key attributes (e.g., "Plant Operator, safety-focused"), set this to TRUE immediately. You can always refine the shortlist in the next turn if they give more info. Set this to false ONLY if the request is extremely vague (e.g., "I need a test").
EXAMPLES OF SUFFICIENT VS INSUFFICIENT INFO:

Latest User Message: "We need a solution for senior leadership."
{{"in_scope": true, "intent": "new_request", "role_summary": "senior leadership", "purpose": null, "topic_keywords": ["senior leadership"], "compared_tests": [], "ready_to_recommend": false, "missing_info_question": "Happy to help narrow that down. What specific skills or responsibilities are you looking to assess for this leadership pool?"}}

Latest User Message: "Hiring graduate financial analysts — final-year students, no work experience. We need numerical reasoning and a finance knowledge test."
{{"in_scope": true, "intent": "new_request", "role_summary": "graduate financial analysts", "purpose": "selection", "topic_keywords": ["financial analyst", "numerical reasoning", "finance", "graduate"], "compared_tests": [], "ready_to_recommend": true, "missing_info_question": null}}

Latest User Message: "We're hiring plant operators for a chemical facility. Safety is absolute top priority — reliability, procedure compliance, never cutting corners."
{{"in_scope": true, "intent": "new_request", "role_summary": "plant operators, chemical facility", "purpose": "selection", "topic_keywords": ["plant operator", "safety", "reliability", "procedure compliance", "manufacturing"], "compared_tests": [], "ready_to_recommend": true, "missing_info_question": null}}

Conversation so far:
{history}

Latest User Message:
{input}

Respond with ONLY the JSON object matching this schema, nothing else:
{format_instructions}""",
            input_variables=["history", "input"],
            partial_variables={"format_instructions": self.parser.get_format_instructions()},
        )
        self.chain = self.prompt | self.llm | self.parser

    def run(self, history: str, latest_message: str) -> dict:
        try:
            return invoke_with_retry(self.chain, {"history": history, "input": latest_message})
        except Exception as e:
            print(f"[NLU] Fallback triggered: {e}")
            # If the API drops completely, force a bypass so the user isn't trapped
            return {
                "in_scope": True,
                "intent": "new_request",
                "role_summary": "Candidate",
                "purpose": "selection",
                "topic_keywords": [w for w in latest_message.replace("'", "").split() if len(w) > 4],
                "compared_tests": [],
                "ready_to_recommend": True, 
                "missing_info_question": None,
            }