agent/nodes/classify_node.py

"""
classify_node.py — Identifies the DSA topic of the given problem.
Uses structured output for reliability; no JSON parsing errors possible.
"""

from agent.models import AgentState
from agent.llm_factory import get_llm
from agent.prompts import CLASSIFY_PROMPT

_llm = get_llm()


def classify_problem(state: AgentState) -> dict:
    """Classifies the problem into a DSA topic and updates problem_topic in state."""
    chain = CLASSIFY_PROMPT | _llm
    result = chain.invoke({"problem": state["problem"]})
    topic = result.content.strip()
    return {"problem_topic": topic}