graph.py

import os
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langchain_core.messages import  AIMessage
from typing import Literal
from nodes.intent import CreditCardState,intent_classifier_node,general_info_handler_node,oos_handler_node
from recommender.vectror_db import query_refiner_node
from recommender.graph_retrieval import neo4j_error_handler_node, neo4j_retrieval_node
from recommender.vectordb_retrieval import ranked_card_retrieval_node
from nodes.agent import agent_node,TOOLS
from nodes.format import format_output_node
from data import debug_print
from nodes.chat import chat_agent_node,chat_tool_node
from nodes.compare import compare_node_fn

# --- Graph Definition ---
graph = StateGraph(CreditCardState)
graph.add_node("intent_classifier", intent_classifier_node)
graph.add_node("general_info_handler", general_info_handler_node)
graph.add_node("oos_handler", oos_handler_node)
graph.add_node("query_refiner", query_refiner_node)
graph.add_node("neo4j_retriever", neo4j_retrieval_node)
graph.add_node("neo4j_error_handler", neo4j_error_handler_node)
graph.add_node("ranked_card_retrieval", ranked_card_retrieval_node)
graph.add_node("agent", agent_node)  # Using our new React agent implementation
graph.add_node("tools", ToolNode(TOOLS))
graph.add_node("format_output", format_output_node)

graph.set_entry_point("intent_classifier")

def route_after_intent_classification(state: CreditCardState):
    intent = state["intent"]
    debug_print("ROUTE", f"Intent classification routing with intent: '{intent}'")

    if intent == "credit-card-recommendation":
        return "query_refiner"
    elif intent == "general-credit-related":
        return "general_info_handler"
    else: # This will catch 'out-of-scope' and any other unexpected values
        return "oos_handler"
    
def route_after_format_output(state: CreditCardState):
    if state.get("trigger_compare", False):
        return "compare_node"
    elif state.get("trigger_chat", False):
        return "chat_node"
    else:
        return END

graph.add_conditional_edges(
    "intent_classifier",
    route_after_intent_classification,
    {
        "query_refiner": "query_refiner",
        "general_info_handler": "general_info_handler",
        "oos_handler": "oos_handler",
    },
)

graph.add_edge("general_info_handler", END)
graph.add_edge("oos_handler", END)
graph.add_edge("query_refiner", "neo4j_retriever")

def route_after_neo4j_retriever(state: CreditCardState):
    debug_print("ROUTE", f"neo4j_error: {state.get('neo4j_error')}")
    if state.get("neo4j_error", False):
        return "neo4j_error_handler"
    else:
        return "ranked_card_retrieval"


graph.add_conditional_edges(
    "neo4j_retriever",
    route_after_neo4j_retriever,
    {
        "neo4j_error_handler": "neo4j_error_handler",
        "ranked_card_retrieval": "ranked_card_retrieval",
    },
)

graph.add_edge("neo4j_error_handler", END)
graph.add_edge("ranked_card_retrieval", "agent")


def route_agent_output(state: CreditCardState) -> Literal["format_output", "tools"]:
    """Determine the next node based on the model's output.

    This function checks if the model's last message contains tool calls.

    Args:
        state (CreditCardState): The current state of the conversation.

    Returns:
        str: The name of the next node to call ("format_output" or "tools").
    """
    last_message = state["messages"][-1]
    if not isinstance(last_message, AIMessage):
        raise ValueError(
            f"Expected AIMessage in output edges, but got {type(last_message).__name__}"
        )
    # If there is no tool call, then we move to format_output
    if not last_message.tool_calls:
        return "format_output"
    # Otherwise we execute the requested actions
    return "tools"

graph.add_conditional_edges(
    "agent",
    route_agent_output
)

graph.add_edge("tools", "agent")

graph.add_edge("format_output",END)

app = graph.compile()

# --- Pipeline Function ---
def run_langgraph_pipeline(
    query: str,
    preferences: str,
    query_intent: bool,
    include_cobranded: bool,
    use_eligibility: bool = False,
    age=None,
    income=None,
    cibil=None,
    min_joining_fee=None,
    max_joining_fee=None,
    min_annual_fee=None,
    max_annual_fee=None
):
    debug_print("PIPELINE", f"Starting pipeline with query: '{query}'")
    debug_print("PIPELINE", f"Preferences: '{preferences}'")
    debug_print("PIPELINE", f"Query intent: {query_intent}, Include cobranded: {include_cobranded}")
    debug_print("PIPELINE", f"Eligibility: {use_eligibility}, Age: {age}, Income: {income}, CIBIL: {cibil}")
    debug_print("PIPELINE", f"Join fee: {min_joining_fee}-{max_joining_fee}, Annual fee: {min_annual_fee}-{max_annual_fee}")
    
    inputs = {
        "query": query,
        "preferences": preferences,
        "query_intent": query_intent,
        "include_cobranded": include_cobranded,
        "use_eligibility": use_eligibility,
        "age": age,
        "income": income,
        "cibil": cibil,
        "min_joining_fee": min_joining_fee,
        "max_joining_fee": max_joining_fee,
        "min_annual_fee": min_annual_fee,
        "max_annual_fee": max_annual_fee,
        "agent_outcome": None,
        "messages": [],
        "trigger_chat": False,
        "trigger_compare": False,
        "selected_cards": [],
        "user_message": "",
    }
    
    debug_print("PIPELINE", f"Invoking LangGraph app")
    result = app.invoke(inputs)
    debug_print("PIPELINE", f"LangGraph execution complete")
    
    # Ensure card_file is a valid file path or None
    card_file = result.get("card_file", None)
    if card_file and os.path.isfile(card_file):
        debug_print("PIPELINE", f"Valid card file found: {card_file}")
    else:
        debug_print("PIPELINE", f"Invalid or missing card file: {card_file}")
        card_file = None
    
    debug_print("PIPELINE", f"Pipeline complete, returning results")
    return (
        result.get("top_card_html", ""),
        result.get("card_rows", []),
        card_file,  # Use the validated file path
        result.get("card_names", []),
        result.get("card_lookup", {}),
    )


def passthrough_node(state: CreditCardState) -> CreditCardState:
    return state

def utility_router(state: CreditCardState):
    if state.get("trigger_compare", False):
        return "compare_node"
    elif state.get("trigger_chat", False):
        return "chat_agent" # Start with the agent
    else:
        raise ValueError("No trigger flag set for utility graph.")

utility_graph = StateGraph(CreditCardState)

# Correct: Add a valid node that returns the full state
utility_graph.add_node("router", passthrough_node)

# Add your real action nodes
utility_graph.add_node("compare_node", compare_node_fn)
utility_graph.add_node("chat_agent", chat_agent_node)
utility_graph.add_node("chat_tools", chat_tool_node)

# Set router as entry point
utility_graph.set_entry_point("router")

# Add conditional edges using utility_router (not as a node)
utility_graph.add_conditional_edges(
    "router",
    utility_router,
    {
        "compare_node": "compare_node",
        "chat_agent": "chat_agent",
    },
)

def route_chat_agent_output(state: CreditCardState) -> Literal["chat_tools", "__end__"]:
    last_message = state["messages"][-1]
    if not isinstance(last_message, AIMessage):
        # This can happen if the tool call fails with an error
        return "__end__"
    if not last_message.tool_calls:
        return "__end__"
    return "chat_tools"

utility_graph.add_conditional_edges(
    "chat_agent",
    route_chat_agent_output,
    {
        "chat_tools": "chat_tools",
        "__end__": "__end__"
    }
)

utility_graph.add_edge("chat_tools", "chat_agent")
utility_graph.add_edge("compare_node", END)

utility_app = utility_graph.compile()