Update recommender/graph_retrieval.py

recommender/graph_retrieval.py

@@ -57,109 +57,130 @@ class Neo4jRetrievalTool(BaseTool):
     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 = genai.GenerativeModel("gemini-2.0-flash")
+        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 = genai.GenerativeModel("gemini-2.0-flash")
+    
+        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:
+                - Only include relevant features based on user query.
+                - 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
     
-            context_flags = f"""
-            Contextual Flags:
-            - FD Card intent: {query_intent}
-            - Include co-branded cards: {include_cobranded}
+                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:
             """
-            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
-                    - Relationships:
-                      - (Card)-[:HAS_FEATURE]->(Feature)
-                    
-                    Feature Inclusion Rules:
-                    - Only include relevant features based on user query.
-                    - 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.
-                    
-                    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'
-                    ---
-                    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
-                    
-                    ---
-                    
-                    {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
+        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,