""" Recommendation engine using embeddings + LLM for SHL assessments """ import pandas as pd import numpy as np from sentence_transformers import SentenceTransformer from sklearn.metrics.pairwise import cosine_similarity import os import pickle from typing import List, Dict, Tuple import google.generativeai as genai from dotenv import load_dotenv load_dotenv() class RecommendationEngine: def __init__(self, catalog_path: str = 'data/shl_catalogue.csv'): """Initialize the recommendation engine""" print("Loading catalog...") self.df = pd.read_csv(catalog_path) # Load embedding model (lightweight for speed) print("Loading embedding model...") self.model = SentenceTransformer('all-MiniLM-L6-v2') # Configure Gemini api_key = os.getenv('GEMINI_API_KEY') if api_key and api_key != 'your_api_key_here': genai.configure(api_key=api_key) self.use_llm = True else: print("Warning: GEMINI_API_KEY not found. LLM explanations disabled.") self.use_llm = False # Create embeddings cache path in the same directory as catalog catalog_dir = os.path.dirname(os.path.abspath(catalog_path)) self.embeddings_path = os.path.join(catalog_dir, 'embeddings.pkl') # Load or create embeddings self.embeddings = self.load_or_create_embeddings() print(f"Loaded {len(self.df)} assessments") def load_or_create_embeddings(self) -> np.ndarray: """Load embeddings from cache or create new ones""" if os.path.exists(self.embeddings_path): print("Loading cached embeddings...") with open(self.embeddings_path, 'rb') as f: return pickle.load(f) else: print("Creating embeddings...") return self.create_embeddings() def create_embeddings(self) -> np.ndarray: """Create embeddings for all assessments""" # Combine relevant fields for embedding texts = [] for _, row in self.df.iterrows(): text = f"{row['assessment_name']} {row['description']} {row['domain']} {row['test_type']}" texts.append(text) # Generate embeddings embeddings = self.model.encode(texts, show_progress_bar=True) # Cache embeddings with open(self.embeddings_path, 'wb') as f: pickle.dump(embeddings, f) return embeddings def retrieve_candidates(self, query: str, top_k: int = 20) -> List[Tuple[int, float]]: """ Retrieve top-k candidate assessments based on semantic similarity Returns list of (index, score) tuples """ # Encode query query_embedding = self.model.encode([query]) # Calculate similarities similarities = cosine_similarity(query_embedding, self.embeddings)[0] # Get top-k indices top_indices = np.argsort(similarities)[::-1][:top_k] return [(idx, similarities[idx]) for idx in top_indices] def balance_recommendations(self, candidates: List[Dict], query: str) -> List[Dict]: """ Balance recommendations to include diverse test types Focus on balancing Personality (P) and Knowledge/Skills (K) when query mentions both """ query_lower = query.lower() # Check if query needs balanced recommendations needs_personality = any(word in query_lower for word in ['collaborate', 'team', 'behavior', 'personality', 'communication', 'leadership', 'soft skill', 'stakeholder']) needs_technical = any(word in query_lower for word in ['java', 'python', 'sql', 'javascript', 'coding', 'programming', 'technical', 'developer', 'engineer', 'analyst', 'data']) needs_cognitive = any(word in query_lower for word in ['cognitive', 'reasoning', 'analytical', 'thinking', 'aptitude']) if (needs_personality and needs_technical) or (needs_personality and needs_cognitive): # Need balanced mix personality_tests = [c for c in candidates if c['test_type'] == 'P'] technical_tests = [c for c in candidates if c['test_type'] in ['K', 'C']] # Aim for roughly equal split target = len(candidates) // 2 balanced = personality_tests[:target] + technical_tests[:target] # Fill remaining with highest scored remaining_slots = len(candidates) - len(balanced) other_candidates = [c for c in candidates if c not in balanced] balanced.extend(other_candidates[:remaining_slots]) return balanced[:len(candidates)] return candidates def generate_explanation(self, query: str, recommendations: List[Dict]) -> Dict: """Generate LLM explanation for recommendations""" if not self.use_llm: return { 'explanation': 'Recommendations based on semantic similarity to query.', 'best_recommendation': recommendations[0]['assessment_name'] if recommendations else None } try: # Create prompt prompt = f"""You are an HR assessment recommendation expert. Job Query: {query} Top Recommended Assessments: """ for i, rec in enumerate(recommendations[:5], 1): prompt += f"\n{i}. {rec['assessment_name']} (Type: {rec['test_type_label']})" prompt += f"\n Description: {rec['description'][:150]}..." prompt += """ Task: 1. For each assessment, explain in 1-2 sentences why it's relevant for this role 2. Identify which ONE assessment is the best overall fit and explain why 3. If the role requires both technical and behavioral skills, ensure you recommend a balanced mix Format your response as: **Assessment 1: [Name]** [Explanation] **Assessment 2: [Name]** [Explanation] ... **Best Overall: [Name]** [Brief reasoning why this is the best fit] """ # Call Gemini model = genai.GenerativeModel('gemini-2.5-flash') response = model.generate_content(prompt) # Extract best recommendation best_rec = recommendations[0]['assessment_name'] if 'Best Overall:' in response.text: best_line = [line for line in response.text.split('\n') if 'Best Overall:' in line] if best_line: best_rec = best_line[0].split('Best Overall:')[1].strip().split('**')[0].strip() return { 'explanation': response.text, 'best_recommendation': best_rec } except Exception as e: print(f"LLM error: {e}") return { 'explanation': f'Recommendations based on semantic similarity. Error generating detailed explanation: {str(e)}', 'best_recommendation': recommendations[0]['assessment_name'] if recommendations else None } def recommend(self, query: str, top_k: int = 10) -> Dict: """ Main recommendation function Returns top-k recommendations with explanations """ # Retrieve candidates candidates = self.retrieve_candidates(query, top_k=min(top_k * 2, 30)) # Convert to list of dicts recommendations = [] for idx, score in candidates: row = self.df.iloc[idx] test_type_map = { 'P': 'Personality & Behavior', 'K': 'Knowledge & Skills', 'C': 'Cognitive', 'G': 'General' } recommendations.append({ 'assessment_name': row['assessment_name'], 'url': row['url'], 'description': row['description'], 'test_type': row['test_type'], 'test_type_label': test_type_map.get(row['test_type'], 'General'), 'domain': row['domain'], 'similarity_score': float(score) }) # Balance recommendations recommendations = self.balance_recommendations(recommendations, query) # Take top-k after balancing recommendations = recommendations[:top_k] # Generate LLM explanation llm_result = self.generate_explanation(query, recommendations) return { 'query': query, 'recommendations': recommendations, 'explanation': llm_result['explanation'], 'best_recommendation': llm_result['best_recommendation'], 'total_results': len(recommendations) } def main(): """Test the recommendation engine""" engine = RecommendationEngine() # Test query test_query = "I am hiring for Java developers who can also collaborate effectively with my business teams." print(f"\n{'='*80}") print(f"Query: {test_query}") print('='*80) result = engine.recommend(test_query, top_k=5) print(f"\nTop {len(result['recommendations'])} Recommendations:") for i, rec in enumerate(result['recommendations'], 1): print(f"\n{i}. {rec['assessment_name']}") print(f" Type: {rec['test_type_label']} | Domain: {rec['domain']}") print(f" Score: {rec['similarity_score']:.3f}") print(f" URL: {rec['url']}") if result['explanation']: print(f"\n{'='*80}") print("LLM Explanation:") print('='*80) print(result['explanation']) if __name__ == "__main__": main()