agent/nodes/solution_node.py

"""
solution_node.py — Reveals the full optimal solution with explanation.

Improvements over v1:
  - Uses llm.with_structured_output() for guaranteed schema compliance
  - Imports SolutionOutput from models (no local import needed)
"""

from agent.models import AgentState, SolutionOutput
from agent.llm_factory import get_llm
from agent.prompts import SOLUTION_PROMPT

_llm = get_llm()
_structured_llm = _llm.with_structured_output(SolutionOutput, method="function_calling")


def reveal_solution(state: AgentState) -> dict:
    """Provides the full, optimal solution with explanation and complexity analysis."""
    topic = state.get("problem_topic", "DSA")

    try:
        result: SolutionOutput = _structured_llm.invoke(
            SOLUTION_PROMPT.format_messages(
                topic=topic,
                problem=state["problem"],
            )
        )
        return {
            "final_response": {
                "solution":    result.solution_code,
                "explanation": result.explanation,
                "complexity":  result.complexity_analysis,
                "type":        "Solution",
                "score":       0,  # Requested solution — no independent credit
            }
        }
    except Exception as e:
        print(f"[solution_node] Structured output error: {e}")
        return {
            "final_response": {
                "solution":    "# Error generating solution",
                "explanation": "Failed to generate solution. Please try again.",
                "complexity":  "N/A",
                "type":        "Solution",
                "score":       0,
            }
        }