recommender/graph_retrieval.py

from data import get_model,debug_print,get_driver,eligibility_df
from nodes.intent import CreditCardState
from neo4j.exceptions import ServiceUnavailable, TransientError, SessionExpired
from langchain.tools import BaseTool
from typing import List
import time
import re


#Cypher query generation and graph retrieval

class Neo4jConnectionError(Exception):
    pass
    
class Neo4jRetrievalTool(BaseTool):
    name: str = "neo4j_card_retriever"
    description: str = "Runs Cypher and builds FAISS on filtered cards."
    
    def generate_cypher(self, user_query: str, query_intent: bool, include_cobranded: bool) -> str:
        debug_print("TOOL", f"generate_cypher called with query: '{user_query}'")
        debug_print("TOOL", f"Parameters: query_intent={query_intent}, include_cobranded={include_cobranded}")
        
        model = get_model("gemini-2.0-flash", use_chat=False)
    
        context_flags = f"""
        Contextual Flags:
        - FD Card intent: {query_intent}
        - Include co-branded cards: {include_cobranded}
        """
        prompt = """
                You are an expert Neo4j Cypher query generator.
                
                Given a user’s question, graph schema, and **contextual flags**, generate the correct Cypher query. The query should return only the card names`c.name as name`.
                
                ONLY output the Cypher query. Do NOT explain anything.
                
                ---
                
                Graph Schema:
                - Nodes:
                  - (Card): Properties = name, bank_name, card_type, premium, co_branded
                  - (Feature): Properties = name
                  - (Partner): Properties = name
                - Relationships:
                  - (Card)-[:HAS_FEATURE]->(Feature)
                  - (Card)-[:PARTNER_WITH]->(Partner)
                
                Feature Inclusion Rules:
                - Include only relevant features mentioned or implied in the refined query. If multiple applicable features from the available list match the query, include all of them together in the f.name IN [...] clause.
                - For vague or broad terms such as “vacation”, “travel”, “frequent flyer”, or “general spending”, **analyze the available feature list below carefully** and include **all relevant travel or spending-related features** that accurately represent the intent. Do not limit to only 1-2 features include relevant features covering all aspects
                - Forex markup fee and foreign transaction fee are the same.
                - If FD Card intent is true then include the features if the query contains any and also include “General Cashback” or “General Reward Points”
                - Don’t add “General Cashback” or “General Reward Points” if it is not required.
                - If fuel is mentioned, include both `Fuel Benefits` and `Fuel Surcharge Waiver`.
                - **ALWAYS** match features using: `f.name IN [...]` — even if there is only **one** feature.
                
                Partner Inclusion Rules:
                - If the user query mentions any available partner brand names, include:
                  MATCH (c)-[:PARTNER_WITH]->(p:Partner)
                  AND p.name IN [<matched partner names>]
                - Always use `p.name IN [...]` — even if there is only one partner.
                
                Valid values:
                - card_type: 'FD Card' or 'Regular'
                - premium: true (no concept of false — just include it if applicable)
                - co_branded: true (no concept of false — just include it if applicable)
                
                MANDATORY Condition Rules:
                - If FD Card intent is true → include: `c.card_type = 'FD Card'`
                - Else → include: `c.card_type = 'Regular'`
                - If the query is based on beginners or students or people with no or low credit history then use FD Card.
                - If the query uses words like "premium", "elite", "luxury", "exclusive", "infinia", "black", etc. → include: `AND c.premium = true`
                - If the query includes low spending, without high spending or budget → include: `(c.premium IS NULL OR c.premium = false)`
                - If include co-branded is false → include: `AND (c.co_branded IS NULL OR c.co_branded = false)`
                - Use exact values for `bank_name` as in the database: ["SBI", "HDFC", "Axis", "ICICI", "YES", "HSBC", "IDFC", "American Express", "SMB", "Federal Bank", "AU Bank", "IDBI", "Kotak Mahindra Bank","IndusInd","RBL"]
                
                ---
                
                Available features:
                      'Fuel Surcharge Waiver','Insurance','Shopping Benefits','Airport Lounge Access','Co-Branded',
                      'Daily Spends (Grocery)','Dining Benefits','Domestic Travel Benefits','Entertainment',
                      'General Reward Points','Movie Benefits','Rupay Network Support','Student',
                      'UPI Transaction Support','Welcome Bonus','International Travel Benefits','premium',
                      'Flight Discounts','Hotel Benefits','Travel Benefits','Railway Benefits','Railway Lounge',
                      'Utility','Beginners (Entry Level)','E-commerce Platform Benefits','Air Miles',
                      'Jewellery Spends','Concierge Services','Food Delivery Benefits','Lifestyle & Luxury Perks',
                      'Spa Access Benefits','Golf Access & Perks','Super Premium','Frequent Flyer Benefits',
                      'Health Benefits','Rent Payment Benefits','Education','Lifetime Free','Roadside Assistance',
                      'EMI Conversion Options','No Forex Markup Fee','Secured FD Based','Cashback','Fuel Benefits','Business'
    
                Available partners:
                      "Marriott Bonvoy", "Accor", "Taj", "ITC", "The Postcard", "Indigo", "United Airlines", "Emirates", "Etihad", "Club Vistara", "Air India", "Turkish airlines"
                
                ---
                Few-shot Examples:
                
                User Query: Show premium cards with airport lounge access  
                Cypher:
                MATCH (c:Card)-[:HAS_FEATURE]->(f:Feature)
                WHERE f.name IN ["Airport Lounge Access"]
                  AND c.card_type = 'Regular'
                  AND c.premium = true
                RETURN DISTINCT c.name AS name
                
                User Query: I want FD cards with spa access and golf perks  
                Cypher:
                MATCH (c:Card)-[:HAS_FEATURE]->(f:Feature)
                WHERE f.name IN ["Spa Access Benefits", "Golf Access & Perks"]
                  AND c.card_type = 'FD Card'
                RETURN DISTINCT c.name AS name
                
                User Query: Cards that support UPI but are not co-branded  
                Cypher:
                MATCH (c:Card)-[:HAS_FEATURE]->(f:Feature)
                WHERE f.name IN ["UPI Transaction Support"]
                  AND c.card_type = 'Regular'
                  AND (c.co_branded IS NULL OR c.co_branded = false)
                RETURN DISTINCT c.name AS name
    
                User Query: Cards partnered with Indigo and Vistara that offer flight benefits  
                Cypher:
                MATCH (c:Card)-[:PARTNER_WITH]->(p:Partner)
                MATCH (c)-[:HAS_FEATURE]->(f:Feature)
                WHERE f.name IN ["Flight Discounts"]
                  AND p.name IN ["Indigo", "Vistara"]
                  AND c.card_type = 'Regular'
                RETURN DISTINCT c.name AS name
                
                ---
                
                {context_flags}
                
                User Query: {user_query}  
                Cypher:
            """
    
        cypher_prompt = prompt.format(context_flags=context_flags, user_query=user_query)
        debug_print("TOOL", f"Calling Gemini to generate Cypher query, prompt length: {len(cypher_prompt)}")
        
        cypher_code = model.generate_content(cypher_prompt).text.strip()
        cleaned_cypher = cypher_code.strip("`").replace("cypher", "").strip()
        
        debug_print("TOOL", f"Generated Cypher query: {cleaned_cypher}")
        return cleaned_cypher

    def _run(
        self,
        query_text: str,
        query_intent: bool = False,
        excluded_cards: List[str]=[],
        include_cobranded: bool = True,
        use_eligibility: bool = False,
        age: int = None,
        income: float = None,
        cibil: int = None,
        min_joining_fee: float = None,
        max_joining_fee: float = None,
        min_annual_fee: float = None,
        max_annual_fee: float = None
    ):
        debug_print("TOOL", f"neo4j_card_retriever _run called with query: '{query_text}'")
        debug_print("TOOL", f"Eligibility filter: {use_eligibility}")
        
        cypher = self.generate_cypher(query_text, query_intent, include_cobranded)
        
        debug_print("TOOL", f"Executing Cypher query against Neo4j")
        attempt = 0
        while attempt < 3:
            try:
                with get_driver().session() as session:
                    matched = [rec["name"] for rec in session.run(cypher)]
                
                break  # success: exit retry loop
    
            except (ServiceUnavailable, TransientError, Neo4jConnectionError,SessionExpired, OSError,TimeoutError) as e:
                attempt += 1
                print(f"Neo4j connection error (attempt {attempt}/{3}): {e}")
                if attempt >= 3:
                    raise Neo4jConnectionError("Failed to connect to the Neo4j database after retries.") from e
                time.sleep(2 * attempt)

        
        debug_print("TOOL", f"Neo4j returned {len(matched)} cards")
        print(excluded_cards)
        print("Before filtering out")
        print(matched)
        
        excluded_cards_clean = [normalize_card_name(c) for c in excluded_cards]
        
        matched = [card for card in matched if normalize_card_name(card) not in excluded_cards_clean]
  
        if use_eligibility:
            debug_print("TOOL", f"Applying eligibility filter with: age={age}, income={income}, cibil={cibil}")
            matched = eligibility_filter(
                matched,
                income,
                cibil,
                age,
                min_joining_fee,
                max_joining_fee,
                min_annual_fee,
                max_annual_fee
            )
            debug_print("TOOL", f"After eligibility filter: {len(matched)} cards remain")
        
        return matched

def neo4j_retrieval_node(state: dict):
    debug_print("NODE", f"Entering neo4j_retrieval_node with query: '{state['query']}'")
    debug_print("NODE", f"Query intent: {state['query_intent']}, Include cobranded: {state['include_cobranded']}")

    try:
        tool = Neo4jRetrievalTool()
        cards = tool._run(
            query_text=state["query"],
            query_intent=state["query_intent"],
            include_cobranded=state["include_cobranded"],
            use_eligibility=state["use_eligibility"],
            excluded_cards=state.get("excluded_cards", []),
            age=state["age"],
            income=state["income"],
            cibil=state["cibil"],
            min_joining_fee=state["min_joining_fee"],
            max_joining_fee=state["max_joining_fee"],
            min_annual_fee=state["min_annual_fee"],
            max_annual_fee=state["max_annual_fee"]
        )
        state["cards"] = cards
        state["neo4j_error"] = False
        return state

    except Neo4jConnectionError as e:
        debug_print("ERROR", f"Neo4j connection failed: {e}")
        debug_print("NODE", "Setting neo4j_error to True due to exception")
        state["neo4j_error"] = True
        state["cards"] = []
        debug_print("NODE", f"Returning from neo4j_retrieval_node with state: {state}")
        return state

    except Exception as e:
        debug_print("ERROR", f"Unexpected exception in neo4j_retrieval_node: {type(e)} - {e}")
        state["neo4j_error"] = True
        state["cards"] = []
        return state


def neo4j_error_handler_node(state: CreditCardState):
    message = "Sorry,the graph database is temporarily unavailable. Please try again in a few minutes."
    print("inside neo4j handler")
    state["top_card"] = ""
    state["top_card_description"] = [message]
    state["card_rows"] = []
    state["card_names"] = []
    state["card_lookup"] = {}

    return state

# --- Eligibility Filter ---
def eligibility_filter(cards, user_income, user_cibil, user_age,min_joining_fee, max_joining_fee,
                            min_annual_fee, max_annual_fee): 
    eligible_cards = []
    for card_name in cards:
        eligibility = eligibility_df[eligibility_df["Name"] == card_name]
        if not eligibility.empty:
            min_income = eligibility.iloc[0]["Minimum Income (LPA)"]
            min_cibil = eligibility.iloc[0]["Minimum Credit Score"]
            min_age = eligibility.iloc[0]["Minimum Age"]
            max_age = eligibility.iloc[0]["Maximum Age"]
            joining_fee=eligibility.iloc[0]["Joining fee"]
            annual_fee=eligibility.iloc[0]["Annual fee"]
            if (user_income >= min_income and 
                user_cibil >= min_cibil and 
                min_age <= user_age <= max_age and
                min_joining_fee<=joining_fee<=max_joining_fee and
                min_annual_fee<=annual_fee<=max_annual_fee):
                eligible_cards.append(card_name)
    return eligible_cards

def normalize_card_name(name):
    name = name.lower()
    name = name.replace("+", "plus")
    name = re.sub(r"[^a-z0-9 ]", "", name) 
    name = re.sub(r"\s+", " ", name).strip()
    return name