graph.py

import asyncio
import logging
import re
from typing import List

from ddgs import DDGS
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_huggingface import HuggingFaceEmbeddings
from langgraph.graph import END, START, StateGraph

from config import config
from model_router import ModelRouter
from models import (
    BatchedDocEval,
    BatchedSentenceFilter,
    DomainClassification,
    QueryTypeClassification,
    State,
    WebSearchQuery,
)
from policy import Stage
from retrieval import AgenticSearchLoop, HybridRetriever, load_vector_store

logger = logging.getLogger("k_state.graph")

# Vector Infrastructure Initializations
embeddings = HuggingFaceEmbeddings(
    model_name=config.embedding_model,
    model_kwargs={"device": "cpu"},
    encode_kwargs={"normalize_embeddings": True},
)

vector_store, all_chunks, bm25 = load_vector_store(embeddings)
hybrid_retriever = HybridRetriever(vector_store, all_chunks, bm25)
router = ModelRouter()

search_wrapper = DuckDuckGoSearchAPIWrapper(
    max_results=5, region="wt-wt", safesearch="moderate", time="m", backend="api"
)


async def web_search_docs(query: str) -> List[Document]:
    if not query.strip():
        return []
    try:
        raw_results = await asyncio.to_thread(search_wrapper.results, query, num_results=5)
        web_docs = []
        for r in raw_results:
            formatted_content = f"[SOURCE: {r.get('title', 'N/A')}]\n{r.get('link', 'N/A')}\n\n{r.get('snippet', 'No content')}"
            web_docs.append(Document(page_content=formatted_content, metadata={"source": "duckduckgo"}))
        if web_docs: 
            return web_docs
        raise RuntimeError("Wrapper empty")
    except Exception:
        def scrape_fallback():
            with DDGS() as ddgs:
                return list(ddgs.text(query, region="wt-wt", safesearch="moderate", max_results=5))
        try:
            results = await asyncio.to_thread(scrape_fallback)
            return [Document(page_content=f"TITLE: {r.get('title')}\nCONTENT: {r.get('body')}", metadata={"source": "duckduckgo"}) for r in results]
        except Exception:
            return []


agentic_search_loop = AgenticSearchLoop(
    hybrid_retriever=hybrid_retriever,
    web_search_func=web_search_docs,
    llm_router=router,
    max_turns=config.max_agentic_turns,
)

# --- Systematic Prompt Topographies ---
domain_prompt = ChatPromptTemplate.from_messages([
    ("system", "Evaluate if the question maps to standard machine learning/data science literature. Output JSON conforming to schema properties."),
    ("human", "Question: {question}"),
])
query_type_prompt = ChatPromptTemplate.from_messages([
    ("system", "Classify the question as factual, complex, or temporal. Output JSON conforming to schema properties."),
    ("human", "Question: {question}"),
])
doc_eval_prompt = ChatPromptTemplate.from_messages([
    ("system", "Review the following context documents and provide individual relevance scores for the question. Output JSON conforming to schema properties."),
    ("human", "Question: {question}\n\nRetrieved Chunks:\n{chunks}"),
])
rewrite_prompt = ChatPromptTemplate.from_messages([
    ("system", "Rewrite the user question into a concise keyword-based web query. Output JSON conforming to schema properties."),
    ("human", "Question: {question}"),
])
filter_prompt = ChatPromptTemplate.from_messages([
    ("system", "Review the list of sentences extracted from the context. Select and extract ONLY the precise sentences that are relevant to answer the question. Output JSON conforming to schema properties."),
    ("human", "Question: {question}\n\nSentence Pool:\n{sentences}"),
])
answer_prompt = ChatPromptTemplate.from_messages([
    ("system", "Synthesize a concise response using only the refined context provided."),
    ("human", "Question: {question}\n\nRefined context:\n{refined_context}"),
])


def decompose_to_sentences(text: str) -> List[str]:
    text = re.sub(r"\s+", " ", text).strip()
    return [s.strip() for s in re.split(r"(?<=[.!?])\s+", text) if len(s.strip()) >= 25]


def is_future_event(question: str) -> bool:
    q = question.lower()
    if "gpt-5" in q or "gpt5" in q or "gpt-6" in q or "gpt-7" in q: 
        return True
    if any(ind in q for ind in ["upcoming", "future", "next year", "latest"]): 
        return True
    return bool(re.findall(r"20(2[6-9]|3[0-9])", q))


# --- Optimized Asynchronous Node Operations ---

async def check_domain_scope(state: State, config: RunnableConfig) -> State:
    try:
        res: DomainClassification = await router.ainvoke_prompt(
            Stage.INTENT, domain_prompt, {"question": state["question"]}, DomainClassification, config=config
        )
        return {**state, "domain": res.domain, "domain_confidence": res.confidence}
    except Exception:
        return {**state, "domain": "out_of_domain", "domain_confidence": 1.0}


async def retrieve(state: State, config: RunnableConfig) -> State:
    q = state["question"]
    
    # Core Protection: Short-circuit tracking metrics if query points to unavailable horizons
    if is_future_event(q):
        return {
            **state,
            "query_type": "temporal",
            "verdict": "UNSUPPORTED",
            "reason": "Query targets future unreleased software horizons."
        }

    try:
        class_res: QueryTypeClassification = await router.ainvoke_prompt(
            Stage.INTENT, query_type_prompt, {"question": q}, QueryTypeClassification, config=config
        )
        q_type = class_res.type
    except Exception:
        q_type = "factual"

    # Offload index traversal processing safely to background thread pools
    docs = await asyncio.to_thread(hybrid_retriever.hybrid_search, q)

    if q_type == "complex":
        agentic_res = await agentic_search_loop.agentic_search(q, docs, config_run=config)
        return {
            **state, 
            "query_type": q_type, 
            "docs": agentic_res["all_docs"], 
            "agentic_turns": agentic_res["turns"], 
            "agentic_context": agentic_res["context"], 
            "subqueries": agentic_res["subqueries"]
        }

    return {**state, "query_type": q_type, "docs": docs, "agentic_turns": 0}


async def eval_each_doc_node(state: State, config: RunnableConfig) -> State:
    q = state["question"]
    docs_to_eval = state.get("docs", [])
    if not docs_to_eval:
        return {**state, "good_docs": [], "doc_scores": [], "verdict": "INCORRECT", "reason": "No document chunks retrieved."}

    # Compress cross-chunk validations into one clean batched payload
    formatted_chunks = "\n\n".join(f"--- CHUNK INDEX {i} ---\n{d.page_content}" for i, d in enumerate(docs_to_eval))
    try:
        batch_res: BatchedDocEval = await router.ainvoke_prompt(
            Stage.DOC_GRADING, doc_eval_prompt, {"question": q, "chunks": formatted_chunks}, BatchedDocEval, config=config
        )
        scores = [0.0] * len(docs_to_eval)
        good = []
        for item in batch_res.evaluations:
            if 0 <= item.index < len(docs_to_eval):
                scores[item.index] = item.score
                if item.score > state["lower_th"]:
                    good.append(docs_to_eval[item.index])
    except Exception:
        scores = [0.0] * len(docs_to_eval)
        good = []

    verdict = "INCORRECT"
    if any(s > state["upper_th"] for s in scores): 
        verdict = "CORRECT"
    elif any(s > state["lower_th"] for s in scores): 
        verdict = "AMBIGUOUS"

    return {**state, "good_docs": good, "doc_scores": scores, "verdict": verdict}


async def rewrite_query_node(state: State, config: RunnableConfig) -> State:
    try:
        res: WebSearchQuery = await router.ainvoke_prompt(
            Stage.QUERY_REWRITE, rewrite_prompt, {"question": state["question"]}, WebSearchQuery, config=config
        )
        return {**state, "web_query": res.query}
    except Exception:
        return {**state, "web_query": state["question"]}


async def web_search(state: State) -> State:
    q = state.get("web_query") or state["question"]
    web_docs = await web_search_docs(q)
    return {**state, "web_docs": web_docs}


async def refine(state: State, config: RunnableConfig) -> State:
    ctx = state.get("agentic_context", "")
    if not ctx:
        source_docs = state.get("good_docs", []) + state.get("web_docs", [])
        ctx = "\n\n".join(d.page_content for d in source_docs).strip()

    if not ctx:
        return {**state, "refined_context": ""}

    strips = decompose_to_sentences(ctx)
    if not strips:
        return {**state, "refined_context": ctx}

    formatted_sentences = "\n".join(f"- {s}" for s in strips)
    try:
        filter_res: BatchedSentenceFilter = await router.ainvoke_prompt(
            Stage.REFINEMENT, filter_prompt, {"question": state["question"], "sentences": formatted_sentences}, BatchedSentenceFilter, config=config
        )
        refined = "\n".join(filter_res.kept_sentences).strip()
    except Exception:
        refined = "\n".join(strips).strip()

    return {**state, "refined_context": refined or "\n".join(strips).strip()}


async def generate(state: State, config: RunnableConfig) -> State:
    refined = state.get("refined_context", "")
    if not refined:
        return {**state, "answer": "I'm specialized in machine learning and AI, but no matching context was found to support an accurate answer."}
    try:
        out = await router.ainvoke_prompt(Stage.FINAL, answer_prompt, {"question": state["question"], "refined_context": refined}, config=config)
        return {**state, "answer": getattr(out, "content", str(out))}
    except Exception:
        return {**state, "answer": "Generation error fallback channel triggered."}


async def fail_node(state: State) -> State:
    """Unified operational funnel processing all system exclusions predictably."""
    reason = state.get("reason")
    if state.get("domain") == "out_of_domain":
        ans = "I'm specialized in machine learning and AI architecture. This question is outside my target domain scope."
    elif reason and "future" in reason.lower():
        ans = "This query targets future or unreleased AI horizons that are currently unsupported by my indexing baseline."
    else:
        ans = "I am unable to resolve this request with sufficient confidence within my verified parameters."
        
    return {
        **state,
        "verdict": state.get("verdict") or "UNSUPPORTED",
        "answer": ans
    }


# --- Evidence-Driven Corrective Routing Filters ---

def route_by_domain(state: State) -> str:
    """Strict Domain firewall routing gatekeeper."""
    if state.get("domain") == "out_of_domain":
        return "out_of_domain"
    return "in_domain"


def route_by_retrieval_quality(state: State) -> str:
    """Core CRAG Branching Matrix: Evaluates strictly based on evidence quality."""
    if state.get("verdict") == "UNSUPPORTED":
        return "fail"
        
    # If a complex agent turn already compiled an integrated context patch, proceed straight to refinement
    if state.get("agentic_context"):
        return "refine"
        
    return "eval_each_doc"


def route_by_evaluation_verdict(state: State) -> str:
    """Routes execution paths by measuring evidence density scores."""
    verdict = state.get("verdict")
    if verdict == "CORRECT":
        return "refine"
    elif verdict == "AMBIGUOUS":
        # Missing context components for relevant queries triggered web recovery
        return "rewrite_query"
    return "fail"


def build_graph():
    g = StateGraph(State)
    
    # Register Node Elements
    g.add_node("check_domain_scope", check_domain_scope)
    g.add_node("retrieve", retrieve)
    g.add_node("eval_each_doc", eval_each_doc_node)
    g.add_node("rewrite_query", rewrite_query_node)
    g.add_node("web_search", web_search)
    g.add_node("refine", refine)
    g.add_node("generate", generate)
    g.add_node("fail", fail_node)

    # Topological Flows Construction
    g.add_edge(START, "check_domain_scope")
    
    # 1. Domain Guardrail Execution Path
    g.add_conditional_edges(
        "check_domain_scope",
        route_by_domain,
        {
            "out_of_domain": "fail",
            "in_domain": "retrieve"
        }
    )
    
    # 2. Retrieval Transition Mapping
    g.add_conditional_edges(
        "retrieve",
        route_by_retrieval_quality,
        {
            "eval_each_doc": "eval_each_doc",
            "refine": "refine",
            "fail": "fail"
        }
    )
    
    # 3. Post-Evaluation Branching (Corrective RAG Core Decision Matrix)
    g.add_conditional_edges(
        "eval_each_doc",
        route_by_evaluation_verdict,
        {
            "refine": "refine",
            "rewrite_query": "rewrite_query",
            "fail": "fail"
        }
    )
    
    # 4. Data Extraction Extension Pathways
    g.add_edge("rewrite_query", "web_search")
    g.add_edge("web_search", "refine")
    g.add_edge("refine", "generate")
    
    # 5. Safe System Enclosures Funneling through Unified Outputs
    g.add_edge("generate", END)
    g.add_edge("fail", END)
    
    return g.compile()


app = build_graph()

def build_initial_state(question: str) -> State:
    return {
        "question": question, "query_type": "", "query_confidence": 0.0, "domain": "", "domain_confidence": 0.0,
        "retrieval_route": "", "docs": [], "good_docs": [], "verdict": "", "reason": "",
        "upper_th": config.upper_th, "lower_th": config.lower_th, "strips": [], "kept_strips": [],
        "refined_context": "", "web_docs": [], "web_query": "", "answer": "", "agentic_turns": 0,
        "agentic_context": "", "subqueries": [], "total_tokens": 0, "prompt_tokens": 0, "completion_tokens": 0, "doc_scores": []
    }