from typing import List
from langchain.prompts import ChatPromptTemplate


FOLLOW_UP_TEMPLATE = """Based on the previous question and answer, generate 2-3 relevant follow-up questions that would help explore the topic further.

Previous Question: {user_input}
Previous Answer: {answer}

Generate short, concise, focused follow-up questions
You don't need a full question as it will be reformulated later as a standalone question with the context. Eg. "Details the first point"
"""

def make_follow_up_node(llm):
    prompt = ChatPromptTemplate.from_template(FOLLOW_UP_TEMPLATE)
    
    def generate_follow_up(state):
        print("---- Generate_follow_up ----")
        if not state.get("answer"):
            return state
            
        response = llm.invoke(prompt.format(
            user_input=state["user_input"],
            answer=state["answer"]
        ))
        
        # Extract questions from response
        follow_ups = [q.strip() for q in response.content.split("\n") if q.strip()]
        state["follow_up_questions"] = follow_ups
        
        return state
        
    return generate_follow_up