Upload 5 files

app/services/embedding_service.py

@@ -0,0 +1,102 @@
+from sentence_transformers import SentenceTransformer
+import numpy as np
+import os
+from loguru import logger
+
+
+class EmbeddingService:
+    """ Handles text embeddings using sentence transformers
+    Pretty straightforward - just wraps the model"""
+    
+    def __init__(self, model_name=None):
+        try:
+            self.model_name = model_name or os.getenv("EMBED_MODEL", "all-MiniLM-L6-v2")
+            logger.info(f"Loading embedding model: {self.model_name}")
+            self.model = SentenceTransformer(self.model_name)
+            self.dimension = self.model.get_sentence_embedding_dimension()
+            logger.success(f"Model loaded. Embedding dimension: {self.dimension}")
+        except Exception as e:
+            logger.error(f"Failed to load embedding model: {str(e)}")
+            raise
+    
+    def embed_single(self, text):
+        try:
+            if not text or not text.strip():
+                return np.zeros(self.dimension, dtype=np.float32)
+            
+            embedding = self.model.encode(
+                text,
+                normalize_embeddings=True,
+                show_progress_bar=False,
+                convert_to_numpy=True
+            )
+            return embedding.astype(np.float32)
+        except Exception as e:
+            logger.error(f"Single embedding failed: {str(e)}")
+            return np.zeros(self.dimension, dtype=np.float32)
+    
+    def embed_batch(self, texts, batch_size=32):
+        try:
+            if not texts:
+                return np.array([], dtype=np.float32)
+            
+            valid_texts = []
+            valid_indices = []
+            for i, text in enumerate(texts):
+                if text and text.strip():
+                    valid_texts.append(text)
+                    valid_indices.append(i)
+            
+            if not valid_texts:
+                return np.zeros((len(texts), self.dimension), dtype=np.float32)
+            
+            # batch 
+            embeddings = self.model.encode(
+                valid_texts,
+                batch_size=batch_size,
+                normalize_embeddings=True,
+                show_progress_bar=False,
+                convert_to_numpy=True
+            )
+            
+            # put embeddings back in right positions
+            output = np.zeros((len(texts), self.dimension), dtype=np.float32)
+            for i, valid_idx in enumerate(valid_indices):
+                output[valid_idx] = embeddings[i]
+            
+            return output.astype(np.float32)
+        except Exception as e:
+            logger.error(f"Batch embedding failed: {str(e)}")
+            return np.zeros((len(texts), self.dimension), dtype=np.float32)
+    
+    def cosine_similarity(self, vec1, vec2):
+        # similarity between two vectors
+        try:
+            norm1 = np.linalg.norm(vec1)
+            norm2 = np.linalg.norm(vec2)
+            
+            if norm1 == 0 or norm2 == 0:
+                return 0.0
+            
+            similarity = np.dot(vec1, vec2) / (norm1 * norm2)
+            return float(np.clip(similarity, -1.0, 1.0))
+        except Exception as e:
+            logger.error(f"Cosine similarity failed: {str(e)}")
+            return 0.0
+    
+    def batch_cosine_similarity(self, query_vec, corpus_vecs):
+        # compare one query against many vectors
+        # faster than looping through each one
+        try:
+            query_norm = query_vec / (np.linalg.norm(query_vec) + 1e-8)
+            
+            # dot product = cosine sim for normalized vectors
+            similarities = np.dot(corpus_vecs, query_norm)
+            return np.clip(similarities, -1.0, 1.0)
+        except Exception as e:
+            logger.error(f"Batch cosine similarity failed: {str(e)}")
+            return np.zeros(len(corpus_vecs))
+    
+    def get_dimension(self):
+        """Get embedding dimension"""
+        return self.dimension

app/services/key_mapper.py

@@ -0,0 +1,222 @@
+import numpy as np
+from rank_bm25 import BM25Okapi
+import re
+import os
+from loguru import logger
+from app.constants import SAMPLE_STORE_KEYS, build_key_search_text
+from app.models import KeyMapping
+from app.services.embedding_service import EmbeddingService
+
+
+class KeyMapper:
+    """ Hybrid approach - combines semantic search with keyword matching
+    TODO: maybe add cross-encoder reranking later if needed """
+    
+    def __init__(self, embedding_service):
+        try:
+            self.embed_service = embedding_service
+            self.rrf_k = int(os.getenv("RRF_K", "60"))  # k=60 worked best in testing
+            self.threshold = float(os.getenv("SIM_THRESHOLD", "0.7"))
+            
+            logger.info("Initializing KeyMapper...")
+            
+            self.keys = SAMPLE_STORE_KEYS
+            
+            # build text for each key to search against
+            self.key_texts = [build_key_search_text(k) for k in self.keys]
+            logger.debug(f"Built {len(self.key_texts)} key search texts")
+            
+            # precompute embeddings so we dont have to do it every time
+            logger.info("Computing key embeddings...")
+            self.key_embeddings = self.embed_service.embed_batch(self.key_texts)
+            logger.debug(f"Key embeddings shape: {self.key_embeddings.shape}")
+            
+            # setup BM25 for keyword matching
+            logger.info("Building BM25 index...")
+            self.tokenized_keys = [self.tokenize(text) for text in self.key_texts]
+            self.bm25 = BM25Okapi(self.tokenized_keys)
+            logger.success("KeyMapper initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize KeyMapper: {str(e)}")
+            raise
+    
+    def tokenize(self, text):
+        # simple tokenization - just split on word boundaries
+        try:
+            tokens = re.findall(r'\w+', text.lower())
+            return tokens
+        except Exception as e:
+            logger.error(f"Tokenization failed: {str(e)}")
+            return []
+    
+    def extract_key_phrases(self, prompt):
+        # extract different phrase combinations from prompt
+        # helps match to specific parts of the prompt
+        try:
+            phrases = []
+            
+            phrases.append(prompt.strip())
+            
+            tokens = self.tokenize(prompt)
+            
+            # bigrams - pairs of words
+            for i in range(len(tokens) - 1):
+                phrases.append(f"{tokens[i]} {tokens[i+1]}")
+            
+            # trigrams - three word combos
+            for i in range(len(tokens) - 2):
+                phrases.append(f"{tokens[i]} {tokens[i+1]} {tokens[i+2]}")
+            
+            # add longer tokens only (skip short words like 'is', 'or')
+            phrases.extend([t for t in tokens if len(t) > 3])
+            
+            # remove dupes but keep order
+            seen = set()
+            unique = []
+            for p in phrases:
+                if p not in seen:
+                    seen.add(p)
+                    unique.append(p)
+            
+            return unique[:15]  # limit to avoid too many
+        except Exception as e:
+            logger.error(f"Phrase extraction failed: {str(e)}")
+            return [prompt]  # fallback to just the prompt
+    
+    def compute_dense_ranks(self, prompt):
+        # get semantic similarity using embeddings
+        try:
+            prompt_emb = self.embed_service.embed_single(prompt)
+            
+            similarities = self.embed_service.batch_cosine_similarity(
+                prompt_emb, 
+                self.key_embeddings
+            )
+            
+            # sort by similarity
+            ranks = np.argsort(-similarities)
+            
+            # convert to rank positions starting from 1
+            rank_positions = np.zeros(len(self.keys), dtype=int)
+            for pos, idx in enumerate(ranks):
+                rank_positions[idx] = pos + 1
+            
+            return rank_positions, similarities
+        except Exception as e:
+            logger.error(f"Dense ranking failed: {str(e)}")
+            # return default ranks if something breaks
+            default_ranks = np.arange(1, len(self.keys) + 1)
+            default_sims = np.zeros(len(self.keys))
+            return default_ranks, default_sims
+    
+    def compute_sparse_ranks(self, prompt):
+        # keyword-based matching with BM25
+        try:
+            prompt_tokens = self.tokenize(prompt)
+            bm25_scores = self.bm25.get_scores(prompt_tokens)
+            
+            ranks = np.argsort(-bm25_scores)
+            
+            rank_positions = np.zeros(len(self.keys), dtype=int)
+            for pos, idx in enumerate(ranks):
+                rank_positions[idx] = pos + 1
+            
+            return rank_positions, bm25_scores
+        except Exception as e:
+            logger.error(f"Sparse ranking failed: {str(e)}")
+            default_ranks = np.arange(1, len(self.keys) + 1)
+            default_scores = np.zeros(len(self.keys))
+            return default_ranks, default_scores
+    
+    def apply_rrf(self, dense_ranks, sparse_ranks):
+        # reciprocal rank fusion - combines both ranking methods
+        # formula from research paper, works better than weighted average
+        try:
+            rrf_scores = (1.0 / (self.rrf_k + dense_ranks)) + \
+                         (1.0 / (self.rrf_k + sparse_ranks))
+            return rrf_scores
+        except Exception as e:
+            logger.error(f"RRF fusion failed: {str(e)}")
+            # fallback to just dense ranks
+            return 1.0 / (self.rrf_k + dense_ranks)
+    
+    def map_keys(self, prompt, top_k=5):
+        """Map user prompt to actual store keys"""
+        try:
+            logger.info(f"Mapping keys for prompt: {prompt[:50]}...")
+            
+            # get rankings from both methods
+            dense_ranks, dense_sims = self.compute_dense_ranks(prompt)
+            sparse_ranks, sparse_scores = self.compute_sparse_ranks(prompt)
+            
+            # combine them using RRF
+            rrf_scores = self.apply_rrf(dense_ranks, sparse_ranks)
+            
+            # sort by combined score
+            sorted_indices = np.argsort(-rrf_scores)
+            
+            # extract phrases from prompt
+            key_phrases = self.extract_key_phrases(prompt)
+            
+            # build the mappings
+            mappings = []
+            for idx in sorted_indices:
+                # normalize score to 0-1 range
+                max_rrf = 2.0 / (self.rrf_k + 1)
+                normalized_score = float(rrf_scores[idx] / max_rrf)
+                
+                # find which phrase matches this key best
+                key_emb = self.key_embeddings[idx]
+                best_phrase = prompt  # default to full prompt
+                best_phrase_sim = dense_sims[idx]
+                
+                # check each phrase
+                for phrase in key_phrases:
+                    phrase_emb = self.embed_service.embed_single(phrase)
+                    phrase_sim = self.embed_service.cosine_similarity(phrase_emb, key_emb)
+                    if phrase_sim > best_phrase_sim:
+                        best_phrase = phrase
+                        best_phrase_sim = phrase_sim
+                
+                mappings.append(KeyMapping(
+                    user_phrase=best_phrase[:50],
+                    mapped_to=self.keys[idx]['value'],
+                    similarity=float(np.clip(normalized_score, 0.0, 1.0))
+                ))
+                
+                if len(mappings) >= top_k:
+                    break
+            
+            logger.success(f"Mapped {len(mappings)} keys successfully")
+            return mappings
+        
+        except Exception as e:
+            logger.error(f"Key mapping failed: {str(e)}")
+            # return empty list if everything breaks
+            return []
+    
+    def get_top_keys(self, prompt, top_k=5, min_similarity=None):
+        """Get top keys with full metadata"""
+        try:
+            threshold = min_similarity if min_similarity is not None else self.threshold
+            
+            # get more than needed then filter
+            mappings = self.map_keys(prompt, top_k=top_k * 2)
+            
+            filtered = [m for m in mappings if m.similarity >= threshold]
+            
+            # add full key details
+            result = []
+            for mapping in filtered[:top_k]:
+                key_obj = next((k for k in self.keys if k['value'] == mapping.mapped_to), None)
+                if key_obj:
+                    result.append({
+                        **key_obj,
+                        'similarity': mapping.similarity,
+                        'matched_phrase': mapping.user_phrase
+                    })
+            
+            return result
+        except Exception as e:
+            logger.error(f"get_top_keys failed: {str(e)}")
+            return []

app/services/rag_service.py

@@ -0,0 +1,270 @@
+import numpy as np
+import faiss
+import os
+from openai import OpenAI
+from loguru import logger
+from app.services.embedding_service import EmbeddingService
+from app.constants import POLICIES
+
+
+class RAGService:
+    """ Handles policy retrieval with FAISS 
+    CRAG = corrective RAG, basically retries if results are bad """
+    
+    def __init__(self, embedding_service):
+        try:
+            self.embed_service = embedding_service
+            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+            
+            logger.info("Initializing RAG Service...")
+            
+            self.policies = POLICIES.copy()
+            self.build_index()
+            
+            logger.success(f"RAG Service initialized with {len(self.policies)} policies")
+        except Exception as e:
+            logger.error(f"Failed to initialize RAG Service: {str(e)}")
+            raise
+    
+    def build_index(self):
+        """Build FAISS index from policy docs"""
+        try:
+            if not self.policies:
+                logger.warning("No policies to index")
+                self.index = None
+                self.policy_embeddings = None
+                return
+            
+            logger.info(f"Embedding {len(self.policies)} policy documents...")
+            
+            # embed all policies
+            self.policy_embeddings = self.embed_service.embed_batch(self.policies)
+            
+            # FAISS index - using inner product since vectors are normalized
+            dimension = self.embed_service.get_dimension()
+            self.index = faiss.IndexFlatIP(dimension)
+            
+            self.index.add(self.policy_embeddings.astype('float32'))
+            
+            logger.debug(f"FAISS index built with {self.index.ntotal} vectors")
+        except Exception as e:
+            logger.error(f"Index building failed: {str(e)}")
+            self.index = None
+            self.policy_embeddings = None
+    
+    def add_documents(self, new_docs):
+        """Add new docs to index on the fly"""
+        try:
+            if not new_docs:
+                return
+            
+            logger.info(f"Adding {len(new_docs)} temporary documents...")
+            
+            new_embeddings = self.embed_service.embed_batch(new_docs)
+            self.index.add(new_embeddings.astype('float32'))
+            self.policies.extend(new_docs)
+            
+            # stack new embeddings with old ones
+            self.policy_embeddings = np.vstack([self.policy_embeddings, new_embeddings])
+            
+            logger.debug(f"Index now contains {self.index.ntotal} documents")
+        except Exception as e:
+            logger.error(f"Failed to add documents: {str(e)}")
+    
+    def retrieve(self, query, top_k=3):
+        """Basic retrieval from FAISS"""
+        try:
+            if self.index is None or self.index.ntotal == 0:
+                logger.warning("Index is empty, returning no results")
+                return []
+            
+            # embed and search
+            query_emb = self.embed_service.embed_single(query).reshape(1, -1)
+            scores, indices = self.index.search(query_emb.astype('float32'), top_k)
+            
+            results = []
+            for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
+                if idx < len(self.policies):
+                    results.append({
+                        'text': self.policies[idx],
+                        'score': float(score),
+                        'index': int(idx),
+                        'rank': i + 1
+                    })
+            
+            logger.debug(f"Retrieved {len(results)} documents")
+            return results
+        except Exception as e:
+            logger.error(f"Retrieval failed: {str(e)}")
+            return []
+    
+    def judge_relevance(self, query, documents):
+        # use llm to score how relevant each doc is
+        # helps filter out garbage results
+        try:
+            if not documents:
+                return []
+            
+            doc_texts = "\n\n".join([
+                f"DOCUMENT {i+1}:\n{doc['text']}" 
+                for i, doc in enumerate(documents)
+            ])
+            
+            judge_prompt = f"""You are an expert relevance evaluator for a loan application rule generation system.
+
+QUERY: {query}
+
+RETRIEVED DOCUMENTS:
+{doc_texts}
+
+Task: Rate the relevance of each document to the query on a scale of 0.0 to 1.0.
+- 1.0 = Highly relevant, directly helps answer the query
+- 0.5 = Somewhat relevant, provides context
+- 0.0 = Not relevant at all
+
+Respond ONLY with a JSON array of scores, one per document in order.
+Example: [0.9, 0.6, 0.2]
+
+Scores:"""
+
+            response = self.client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[
+                    {"role": "system", "content": "You are a relevance scoring expert. Respond only with a JSON array of numbers."},
+                    {"role": "user", "content": judge_prompt}
+                ],
+                temperature=0.1,
+                max_tokens=100
+            )
+            
+            content = response.choices[0].message.content.strip()
+            import json
+            scores = json.loads(content)
+            
+            # clamp to 0-1 range
+            scores = [max(0.0, min(1.0, float(s))) for s in scores]
+            
+            # pad if llm didnt return enough scores
+            while len(scores) < len(documents):
+                scores.append(0.5)
+            
+            logger.debug(f"LLM judge scores: {scores}")
+            return scores[:len(documents)]
+            
+        except Exception as e:
+            logger.error(f"LLM judge failed: {str(e)}")
+            # fallback - just use retrieval scores
+            return [doc['score'] / (doc['score'] + 1.0) for doc in documents]
+    
+    def refine_query(self, original_query, low_relevance_docs):
+        # if results are bad, ask llm to rewrite the query
+        # usually helps by adding more specific terms
+        try:
+            refine_prompt = f"""Original query: "{original_query}"
+
+The retrieved documents were not very relevant. Suggest a better search query that focuses on key loan application terms like:
+- Bureau score, credit score, CIBIL
+- Business vintage, age
+- Overdue amounts, DPD
+- Income, FOIR
+- GST, banking metrics
+
+Respond with ONLY the improved query, no explanation.
+
+Improved query:"""
+
+            response = self.client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[
+                    {"role": "system", "content": "You are a query refinement expert for loan application rules."},
+                    {"role": "user", "content": refine_prompt}
+                ],
+                temperature=0.3,
+                max_tokens=100
+            )
+            
+            refined = response.choices[0].message.content.strip().strip('"')
+            logger.info(f"Refined query: {refined}")
+            return refined if refined else original_query
+            
+        except Exception as e:
+            logger.error(f"Query refinement failed: {str(e)}")
+            return original_query
+    
+    def retrieve_with_crag(self, query, top_k=2, relevance_threshold=0.7):
+        """
+        CRAG = Corrective RAG
+        retrieves docs, checks if theyre good, retries if not
+        """
+        try:
+            logger.info(f"CRAG: Retrieving for query: '{query[:50]}...'")
+            
+            docs = self.retrieve(query, top_k=top_k)
+            
+            if not docs:
+                logger.warning("No documents retrieved")
+                return [], 0.0
+            
+            # judge how relevant results are
+            relevance_scores = self.judge_relevance(query, docs)
+            
+            for doc, score in zip(docs, relevance_scores):
+                doc['relevance'] = score
+            
+            avg_relevance = np.mean(relevance_scores)
+            logger.debug(f"CRAG: Initial relevance: {avg_relevance:.3f}")
+            
+            # if relevance sucks, refine and try again
+            if avg_relevance < relevance_threshold:
+                logger.info("CRAG: Low relevance detected, refining query...")
+                
+                refined_query = self.refine_query(query, [d['text'] for d in docs])
+                logger.debug(f"CRAG: Refined query: '{refined_query[:50]}...'")
+                
+                # try again with better query
+                refined_docs = self.retrieve(refined_query, top_k=top_k)
+                
+                if refined_docs:
+                    refined_relevance = self.judge_relevance(refined_query, refined_docs)
+                    
+                    for doc, score in zip(refined_docs, refined_relevance):
+                        doc['relevance'] = score
+                    
+                    refined_avg = np.mean(refined_relevance)
+                    logger.debug(f"CRAG: Refined relevance: {refined_avg:.3f}")
+                    
+                    # use refined results only if theyre better
+                    if refined_avg > avg_relevance:
+                        docs = refined_docs
+                        avg_relevance = refined_avg
+                        logger.info("CRAG: Using refined results")
+                    else:
+                        logger.info("CRAG: Keeping original results")
+            
+            # sort by relevance score
+            docs.sort(key=lambda x: x['relevance'], reverse=True)
+            
+            return docs, avg_relevance
+        
+        except Exception as e:
+            logger.error(f"CRAG failed: {str(e)}")
+            return [], 0.0
+    
+    def format_context(self, documents, max_length=500):
+        """Format docs into a string for llm context"""
+        try:
+            if not documents:
+                return "No relevant policies found."
+            
+            context_parts = []
+            for i, doc in enumerate(documents):
+                text = doc['text'][:max_length]
+                relevance = doc.get('relevance', doc.get('score', 0))
+                context_parts.append(
+                    f"Policy {i+1} (relevance: {relevance:.2f}):\n{text}"
+                )
+            
+            return "\n\n".join(context_parts)
+        except Exception as e:
+            logger.error(f"Context formatting failed: {str(e)}")
+            return "Error formatting policy context."

app/services/rule_service.py

@@ -0,0 +1,315 @@
+import json
+import os
+from openai import OpenAI
+from json_logic import jsonLogic
+from loguru import logger
+from app.constants import MOCK_STORE_SAMPLES
+from app.models import KeyMapping
+
+
+class RuleGenerationService:
+    """ Generates json logic rules using gpt-4o-mini
+        Uses self-consistency voting to pick best rule from multiple attempts
+    """
+    
+    def __init__(self):
+        try:
+            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+            self.model = "gpt-4o-mini"
+            logger.success("RuleGenerationService initialized")
+        except Exception as e:
+            logger.error(f"Failed to initialize RuleGenerationService: {str(e)}")
+            raise
+    
+    def build_system_prompt(self, available_keys, policy_context):
+        
+        try:
+            keys_str = json.dumps(available_keys, indent=2)
+            
+            system_prompt = f"""You are an expert JSON Logic rule generator for loan application systems.
+
+AVAILABLE KEYS (use ONLY these in {{"var": "key"}}):
+{keys_str}
+
+POLICY CONTEXT:
+{policy_context}
+
+JSON LOGIC OPERATORS:
+- Logical: "and", "or", "!", "if"
+- Comparison: ">", "<", ">=", "<=", "==", "!="
+- Arrays: "in", "some", "all"
+- Math: "+", "-", "*", "/"
+
+RULES:
+1. Use ONLY the available keys listed above
+2. All keys must be referenced using {{"var": "key.name"}}
+3. Generate valid JSON Logic syntax
+4. Be precise with thresholds from policies
+5. Use "and" for multiple conditions, "or" for alternatives
+
+OUTPUT FORMAT (must be valid JSON):
+{{
+  "json_logic": {{"and": [...]}},
+  "explanation": "Brief 1-2 sentence explanation",
+  "used_keys": ["key1", "key2"],
+  "confidence": 0.0-1.0
+}}
+
+EXAMPLES:
+
+User: "Approve if bureau score > 700"
+Output:
+{{
+  "json_logic": {{">": [{{"var": "bureau.score"}}, 700]}},
+  "explanation": "Approves applications where bureau score exceeds 700.",
+  "used_keys": ["bureau.score"],
+  "confidence": 0.95
+}}
+
+User: "Reject if wilful default OR suit filed"
+Output:
+{{
+  "json_logic": {{"or": [
+    {{"==": [{{"var": "bureau.wilful_default"}}, true]}},
+    {{"==": [{{"var": "bureau.suit_filed"}}, true]}}
+  ]}},
+  "explanation": "Rejects applications with wilful default or suit filed status.",
+  "used_keys": ["bureau.wilful_default", "bureau.suit_filed"],
+  "confidence": 0.92
+}}
+
+User: "Approve if age between 25 and 60"
+Output:
+{{
+  "json_logic": {{"and": [
+    {{">=": [{{"var": "primary_applicant.age"}}, 25]}},
+    {{"<=": [{{"var": "primary_applicant.age"}}, 60]}}
+  ]}},
+  "explanation": "Approves when primary applicant age is between 25 and 60 years inclusive.",
+  "used_keys": ["primary_applicant.age"],
+  "confidence": 0.90
+}}"""
+
+            return system_prompt
+        except Exception as e:
+            logger.error(f"Failed to build system prompt: {str(e)}")
+            return ""
+    
+    def generate_single_rule(self, prompt, system_prompt, temperature=0.2):
+        # calls llm once to get a rule
+        # temperature controls randomness - lower = more deterministic
+        try:
+            response = self.client.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt}
+                ],
+                temperature=temperature,
+                max_tokens=800,
+                response_format={"type": "json_object"}  
+            )
+            
+            content = response.choices[0].message.content
+            result = json.loads(content)
+            
+            # check if all required fields are there
+            if not all(k in result for k in ['json_logic', 'explanation', 'used_keys']):
+                logger.warning("LLM response missing required fields")
+                raise ValueError("Missing required fields in LLM response")
+            
+            # add default confidence if llm forgot to include it
+            if 'confidence' not in result:
+                result['confidence'] = 0.8
+            
+            return result
+            
+        except Exception as e:
+            logger.error(f"Rule generation failed: {str(e)}")
+            return None
+    
+    def validate_rule(self, rule, available_keys):
+        # checks if rule only uses allowed keys
+        # extracts all {"var": "..."} and validates them
+        try:
+            if not rule:
+                return False, "Empty rule"
+            
+            # recursively find all var references
+            def extract_vars(obj):
+                vars_found = []
+                if isinstance(obj, dict):
+                    if "var" in obj:
+                        vars_found.append(obj["var"])
+                    for value in obj.values():
+                        vars_found.extend(extract_vars(value))
+                elif isinstance(obj, list):
+                    for item in obj:
+                        vars_found.extend(extract_vars(item))
+                return vars_found
+            
+            used_vars = extract_vars(rule)
+            
+            # check all vars are in allowed list
+            invalid_vars = [v for v in used_vars if v not in available_keys]
+            if invalid_vars:
+                logger.warning(f"Rule uses invalid keys: {invalid_vars}")
+                return False, f"Invalid keys used: {invalid_vars}"
+            
+            return True, ""
+        
+        except Exception as e:
+            logger.error(f"Rule validation failed: {str(e)}")
+            return False, str(e)
+    
+    def test_rule_on_mocks(self, rule, num_samples=5):
+        # runs the rule against mock data to see if it breaks
+        # doesn't check correctness, just that it executes
+        try:
+            if not rule:
+                return 0.0
+            
+            successes = 0
+            samples = MOCK_STORE_SAMPLES[:num_samples]
+            
+            for sample in samples:
+                try:
+                    # apply json logic rule
+                    result = jsonLogic(rule, sample)
+                    successes += 1
+                except Exception as e:
+                    # rule broke on this sample
+                    logger.debug(f"Rule test failed on sample: {str(e)}")
+                    continue
+            
+            success_rate = successes / len(samples) if samples else 0.0
+            return success_rate
+        
+        except Exception as e:
+            logger.error(f"Mock testing failed: {str(e)}")
+            return 0.0
+    
+    def self_consistency_vote(self, variants):
+        # picks the best rule from multiple variants
+        # scores based on confidence, validation rate, and simplicity
+        try:
+            if not variants:
+                return None
+            
+            if len(variants) == 1:
+                return variants[0]
+            
+            scored_variants = []
+            for variant in variants:
+                score = 0.0
+                
+                # llm's own confidence
+                score += variant.get('confidence', 0.5) * 0.4
+                
+                # how well it ran on mock data
+                validation_rate = variant.get('validation_rate', 0.0)
+                score += validation_rate * 0.4
+                
+                # prefer simpler rules (less json length)
+                rule_str = json.dumps(variant['json_logic'])
+                complexity_penalty = min(len(rule_str) / 500, 0.2)
+                score += (0.2 - complexity_penalty)
+                
+                scored_variants.append((score, variant))
+            
+            # sort by score descending
+            scored_variants.sort(key=lambda x: x[0], reverse=True)
+            
+            scores_str = [f'{s:.3f}' for s, _ in scored_variants]
+            logger.debug(f"Self-consistency scores: {scores_str}")
+            
+            return scored_variants[0][1]  # return best one
+        
+        except Exception as e:
+            logger.error(f"Self-consistency voting failed: {str(e)}")
+            return variants[0] if variants else None
+    
+    def generate_rule(self, prompt, key_mappings, policy_context, num_variants=3):
+        """
+        Main method - generates rule with self-consistency
+        tries multiple times with different temperatures and picks best
+        """
+        try:
+            logger.info(f"Generating {num_variants} rule variants...")
+            
+            # get list of allowed keys
+            available_keys = [m.mapped_to for m in key_mappings]
+            
+            # build the big system prompt
+            system_prompt = self.build_system_prompt(available_keys, policy_context)
+            
+            # generate multiple variants
+            variants = []
+            temperatures = [0.1, 0.3, 0.5][:num_variants]
+            
+            for i, temp in enumerate(temperatures):
+                logger.debug(f"Generating variant {i+1} with temp={temp}...")
+                
+                result = self.generate_single_rule(prompt, system_prompt, temperature=temp)
+                
+                if result:
+                    # validate it uses correct keys
+                    is_valid, error_msg = self.validate_rule(
+                        result['json_logic'], 
+                        available_keys
+                    )
+                    
+                    if not is_valid:
+                        logger.warning(f"Variant {i+1} validation failed: {error_msg}")
+                        continue
+                    
+                    # test on mock data
+                    validation_rate = self.test_rule_on_mocks(result['json_logic'])
+                    result['validation_rate'] = validation_rate
+                    
+                    logger.info(f"Variant {i+1}: conf={result['confidence']:.3f}, val={validation_rate:.3f}")
+                    
+                    variants.append(result)
+            
+            if not variants:
+                logger.error("Failed to generate any valid rule variants")
+                raise ValueError("Failed to generate any valid rule variants")
+            
+            # vote for best rule
+            best_rule = self.self_consistency_vote(variants)
+            
+            logger.success(f"Selected best rule")
+            
+            return best_rule
+        
+        except Exception as e:
+            logger.error(f"Rule generation failed: {str(e)}")
+            raise
+    
+    def calculate_confidence_score(self, rule_result, key_mappings, policy_relevance):
+        """Calculate overall confidence from multiple factors"""
+        try:
+            # how well keys matched (40%)
+            avg_key_sim = sum(m.similarity for m in key_mappings) / len(key_mappings) if key_mappings else 0.0
+            
+            # how relevant policies were (30%)
+            policy_score = policy_relevance
+            
+            # llm confidence + validation rate (30%)
+            llm_confidence = rule_result.get('confidence', 0.8)
+            validation_rate = rule_result.get('validation_rate', 0.8)
+            generation_score = (llm_confidence + validation_rate) / 2
+            
+            # weighted average
+            confidence = (
+                avg_key_sim * 0.4 +
+                policy_score * 0.3 +
+                generation_score * 0.3
+            )
+            
+            # clamp to 0-1
+            return float(min(max(confidence, 0.0), 1.0))
+        
+        except Exception as e:
+            logger.error(f"Confidence calculation failed: {str(e)}")
+            return 0.5  # default fallback