Query → [Retrieve] → [Grade] → [Generate/Rewrite/Web Search] → Response
              ↓           ↓
         Documents    Quality Check
                          ↓
                   Route Decision:
                   - Good docs → Generate
                   - Poor docs → Rewrite query
                   - No docs → Web fallback

from langgraph.graph import StateGraph, START, END
from typing import TypedDict, List, Annotated
from langchain_core.documents import Document
import operator

class RAGState(TypedDict):
    """State for agentic RAG workflows."""
    question: str
    documents: Annotated[List[Document], operator.add]
    generation: str
    web_search_needed: bool
    retry_count: int
    relevance_scores: dict[str, float]

def retrieve(state: RAGState) -> dict:
    """Retrieve documents from vector store."""
    question = state["question"]
    documents = retriever.invoke(question)
    return {"documents": documents, "question": question}

from pydantic import BaseModel, Field

class GradeDocuments(BaseModel):
    """Binary score for document relevance."""
    binary_score: str = Field(
        description="Relevance score 'yes' or 'no'"
    )

def grade_documents(state: RAGState) -> dict:
    """Grade documents for relevance - core Self-RAG pattern."""
    question = state["question"]
    documents = state["documents"]

    filtered_docs = []
    relevance_scores = {}

    for doc in documents:
        score = retrieval_grader.invoke({
            "question": question,
            "document": doc.page_content
        })
        doc_id = doc.metadata.get("id", hash(doc.page_content))
        relevance_scores[doc_id] = 1.0 if score.binary_score == "yes" else 0.0

        if score.binary_score == "yes":
            filtered_docs.append(doc)

    # Trigger web search if too many docs filtered out
    web_search_needed = len(filtered_docs) < len(documents) // 2

    return {
        "documents": filtered_docs,
        "web_search_needed": web_search_needed,
        "relevance_scores": relevance_scores
    }

def transform_query(state: RAGState) -> dict:
    """Transform query for better retrieval."""
    question = state["question"]

    better_question = question_rewriter.invoke({
        "question": question,
        "feedback": "Rephrase to improve retrieval. Be specific."
    })

    return {
        "question": better_question,
        "retry_count": state.get("retry_count", 0) + 1
    }

def web_search(state: RAGState) -> dict:
    """Fallback to web search when documents insufficient."""
    question = state["question"]

    web_results = tavily_client.search(
        question,
        max_results=5,
        search_depth="advanced"
    )

    web_docs = [
        Document(
            page_content=r["content"],
            metadata={"source": r["url"], "type": "web"}
        )
        for r in web_results
    ]

    return {"documents": web_docs, "web_search_needed": False}

def generate(state: RAGState) -> dict:
    """Generate answer from documents."""
    question = state["question"]
    documents = state["documents"]

    context = "\n\n".join([
        f"[{i+1}] {doc.page_content}"
        for i, doc in enumerate(documents)
    ])

    generation = rag_chain.invoke({
        "context": context,
        "question": question
    })

    return {"generation": generation}

def route_after_grading(state: RAGState) -> str:
    """Route based on document quality."""
    if state["web_search_needed"]:
        if state.get("retry_count", 0) < 2:
            return "transform_query"  # Try rewriting first
        return "web_search"  # Fallback to web
    return "generate"  # Documents are good

workflow.add_conditional_edges(
    "grade",
    route_after_grading,
    {
        "generate": "generate",
        "transform_query": "transform_query",
        "web_search": "web_search"
    }
)

def build_crag_workflow() -> StateGraph:
    """Build Corrective-RAG workflow with web fallback."""
    workflow = StateGraph(RAGState)

    # Add nodes
    workflow.add_node("retrieve", retrieve)
    workflow.add_node("grade", grade_documents)
    workflow.add_node("generate", generate)
    workflow.add_node("web_search", web_search)
    workflow.add_node("transform_query", transform_query)

    # Define edges
    workflow.add_edge(START, "retrieve")
    workflow.add_edge("retrieve", "grade")

    # Conditional routing based on document quality
    workflow.add_conditional_edges(
        "grade",
        route_after_grading,
        {
            "generate": "generate",
            "transform_query": "transform_query",
            "web_search": "web_search"
        }
    )

    # After query transform, retry retrieval
    workflow.add_edge("transform_query", "retrieve")

    # Web search leads to generation
    workflow.add_edge("web_search", "generate")

    workflow.add_edge("generate", END)

    return workflow.compile()