app/graph/workflow.py

import json
from typing import List, Optional, TypedDict
from langchain_core.prompts import PromptTemplate
from langchain_core.documents import Document
from langchain_openai import ChatOpenAI
from langgraph.graph import START, END, StateGraph
import httpx

from app.core.config import settings
from app.core.logging import logger
from app.engine.context_builder import build_context, source_citations
from app.engine.retriever import retrieve_documents
from app.engine.reranker import rerank_documents, evaluate_nli_groundedness


class GraphState(TypedDict):
    question: str
    chat_history: List[dict]
    generation: str
    documents: List[Document]
    sources: Optional[list[dict]]
    run_count: int
    confidence_score: float
    grounded: str
    summary: Optional[str]
    optimistic_route: Optional[bool]
    max_similarity: Optional[float]


_http_client = httpx.AsyncClient(
    http2=True,
    limits=httpx.Limits(max_keepalive_connections=20, max_connections=50),
    timeout=httpx.Timeout(60.0, connect=10.0),
)

llm = ChatOpenAI(
    model=settings.LLM_MODEL,
    temperature=0,
    openai_api_key=settings.OPENROUTER_API_KEY,
    openai_api_base=settings.OPENROUTER_BASE_URL,
    default_headers={"HTTP-Referer": "https://localhost:3000", "X-Title": "Support Docs Copilot"},
    http_async_client=_http_client,
)

llm_slow = ChatOpenAI(
    model=getattr(settings, "SLOW_LLM_MODEL", settings.LLM_MODEL),
    temperature=0,
    openai_api_key=settings.OPENROUTER_API_KEY,
    openai_api_base=settings.OPENROUTER_BASE_URL,
    default_headers={"HTTP-Referer": "https://localhost:3000", "X-Title": "Support Docs Copilot"},
    http_async_client=_http_client,
)


async def retrieve(state: GraphState):
    logger.info("NODE: RETRIEVE DOCS")
    question = state["question"]
    chat_history = state.get("chat_history", [])
    run_count = state.get("run_count", 0)
    documents = await retrieve_documents(question, chat_history)
    max_sim = max([d.metadata.get("similarity_score", 0.0) for d in documents] + [0.0])
    optimistic = max_sim >= 0.82
    if optimistic:
        logger.info(f"OPTIMISTIC ROUTE TRIGGERED: Top similarity score {max_sim:.4f} >= 0.82")
    return {
        "documents": documents,
        "sources": source_citations(documents),
        "question": question,
        "run_count": run_count,
        "optimistic_route": optimistic,
        "max_similarity": max_sim,
    }


async def grade_documents(state: GraphState):
    logger.info("NODE: GRADE DOCUMENT RELEVANCE (VIA HYBRID RERANKER)")
    question = state["question"]
    documents = state.get("documents", [])
    
    reranked_docs = rerank_documents(question, documents, top_k=settings.RERANKER_TOP_N)
    if not reranked_docs:
        return {"documents": []}
        
    filtered_docs = []
    for doc in reranked_docs:
        score = doc.metadata.get("relevance_score", doc.metadata.get("rerank_score", 1.0))
        if score >= settings.MIN_RELEVANCE_SCORE:
            filtered_docs.append(doc)
            
    if not filtered_docs and reranked_docs:
        top_score = reranked_docs[0].metadata.get("rerank_score", 0.0)
        if top_score > 0.0:
            filtered_docs = [reranked_docs[0]]
            
    logger.info(f"Relevance grader filtered {len(reranked_docs)} docs down to {len(filtered_docs)} relevant docs (threshold >= {settings.MIN_RELEVANCE_SCORE}).")
    return {"documents": filtered_docs}


async def decide_to_generate(state: GraphState):
    if not state.get("documents"):
        logger.info("ROUTE: ALL DOCS IRRELEVANT")
        return "end"
    logger.info("ROUTE: RELEVANT DOCS FOUND")
    return "generate"


async def generate(state: GraphState):
    logger.info("NODE: GENERATE ANSWER")
    question = state["question"]
    documents = state["documents"]
    chat_history = state.get("chat_history", [])
    summary = state.get("summary", "")
    run_count = state.get("run_count", 0) + 1
    
    history_lines = [f"{msg['role']}: {msg['content']}" for msg in chat_history[-6:]]
    if summary:
        history_lines.insert(0, summary if summary.startswith("System Summary:") else f"System Summary: {summary}")
    history_str = "\n".join(history_lines)
    context = build_context(documents)
    prompt = PromptTemplate(
        template="""You are a Support Docs Copilot. Use only the retrieved context to answer the question concisely.
        
        CRITICAL INSTRUCTION (Cite-to-Write):
        You must append [doc_id] to the end of every sentence. Do not write a sentence if you cannot cite a source from the retrieved context. If the context does not contain the answer, say "I don't know".
        
        Chat History:
        {chat_history}
        
        Question: {question} 
        Context: {context} 
        Answer:""",
        input_variables=["question", "context", "chat_history"],
    )
    selected_llm = llm_slow if run_count > 1 else llm
    if run_count > 1:
        logger.info(f"Using slow reasoning model ({getattr(settings, 'SLOW_LLM_MODEL', 'default')}) for retry attempt #{run_count}")
    rag_chain = prompt | selected_llm
    generation = await rag_chain.ainvoke({"context": context, "question": question, "chat_history": history_str})
    return {"generation": generation.content, "sources": source_citations(documents), "run_count": run_count}


async def evaluate_answer(state: GraphState):
    logger.info("NODE: EVALUATE ANSWER")
    documents = state["documents"]
    generation = state["generation"]
    
    context = build_context(documents)
    grade, confidence = evaluate_nli_groundedness(context, generation)
        
    return {"grounded": grade, "confidence_score": confidence}


async def check_hallucinations(state: GraphState):
    run_count = state["run_count"]
    
    if run_count >= 3:
        logger.info("ROUTE: MAX RETRIES REACHED")
        return "end"
        
    grade = state.get("grounded", "yes")
        
    if grade.lower() == "yes":
        logger.info("ROUTE: GROUNDED")
        return "end"
    logger.info("ROUTE: HALLUCINATION DETECTED")
    return "regenerate"


async def decide_optimistic_or_grade(state: GraphState):
    if state.get("optimistic_route") and state.get("documents"):
        logger.info(f"OPTIMISTIC STREAMING: High similarity ({state.get('max_similarity', 0.0):.4f} >= 0.82). Skipping LLM grader node!")
        return "generate"
    return "grade_documents"


def compile_workflow():
    workflow = StateGraph(GraphState)
    workflow.add_node("retrieve", retrieve)
    workflow.add_node("grade_documents", grade_documents)
    workflow.add_node("generate", generate)
    workflow.add_node("evaluate_answer", evaluate_answer)
    workflow.add_edge(START, "retrieve")
    workflow.add_conditional_edges("retrieve", decide_optimistic_or_grade, {"generate": "generate", "grade_documents": "grade_documents"})
    workflow.add_conditional_edges("grade_documents", decide_to_generate, {"generate": "generate", "end": END})
    workflow.add_edge("generate", "evaluate_answer")
    workflow.add_conditional_edges("evaluate_answer", check_hallucinations, {"end": END, "regenerate": "generate"})
    return workflow.compile()