app/nodes/rag.py

"""
app/nodes/rag.py
─────────────────
CHECKPOINT 2 — RAG node

Retrieves relevant document chunks and stores them in state so the
LLM node can inject them into its prompt.
"""

from app.state import AgentState
from app.rag.store import retrieve_context


def rag_node(state: AgentState) -> AgentState:
    context = retrieve_context(state["query"])
    print(f"[RAG] Retrieved {len(context.splitlines())} chunk(s).")
    log = state.get("node_log", []) + ["rag"]
    return {**state, "rag_context": context, "node_log": log}