""" rag_engine.py - Retrieval-Augmented Generation Engine Handles: 1. Query embedding 2. Message-level retrieval (cosine similarity over all 191K embeddings) 3. Topic-level retrieval (precomputed centroids) 4. Checkpoint retrieval 5. Answer generation via flan-t5-small """ import json import numpy as np from pathlib import Path from sentence_transformers import SentenceTransformer from transformers import AutoModelForSeq2SeqLM, AutoTokenizer # Paths BASE_DIR = Path(__file__).parent.parent EMB_PATH = BASE_DIR / "data" / "processed" / "embeddings.npy" INDEX_PATH = BASE_DIR / "data" / "processed" / "embedding_index.json" JSONL_PATH = BASE_DIR / "data" / "processed" / "processed_messages.jsonl" TOPICS_PATH = BASE_DIR / "data" / "processed" / "topic_segments.json" SUMMARIES_PATH = BASE_DIR / "data" / "processed" / "summaries.json" PERSONA_PATH = BASE_DIR / "data" / "processed" / "persona.json" EMBED_MODEL = "all-MiniLM-L6-v2" QA_MODEL = "google/flan-t5-small" class RAGEngine: """Lazy-loaded RAG engine. Call .initialize() once at startup.""" def __init__(self): self.ready = False self.embedder = None self.qa_model = None self.embeddings = None # shape [N, 384], already L2-normalized self.msg_index = None # list[int]: array_pos -> msg_id self.messages = {} # msg_id -> record dict self.topics = [] # list of topic segment dicts self.topic_centroids = None # shape [T, 384] self.checkpoints = [] # list of checkpoint dicts self.persona = {} # Initialization def initialize(self): """Load all assets into memory. Call once at server startup.""" import time t0 = time.time() print("[RAG] Loading sentence-transformer...") self.embedder = SentenceTransformer(EMBED_MODEL) print("[RAG] Loading QA model (flan-t5-small)...") self.qa_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL) self.qa_model = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL) print("[RAG] Loading embeddings...") raw = np.load(EMB_PATH).astype(np.float32) # Ensure L2-normalized norms = np.linalg.norm(raw, axis=1, keepdims=True) self.embeddings = raw / np.maximum(norms, 1e-10) print("[RAG] Loading message index...") with open(INDEX_PATH, 'r') as f: idx = json.load(f) self.msg_index = [int(idx[str(i)]) for i in range(len(idx))] self.mid_to_index = {mid: i for i, mid in enumerate(self.msg_index)} print("[RAG] Loading messages...") with open(JSONL_PATH, 'r', encoding='utf-8') as f: for line in f: rec = json.loads(line) self.messages[rec["msg_id"]] = rec print("[RAG] Loading topics...") with open(TOPICS_PATH, 'r', encoding='utf-8') as f: self.topics = json.load(f) print("[RAG] Precomputing topic centroids...") self._precompute_centroids() print("[RAG] Loading summaries...") self._load_summaries() print("[RAG] Loading persona...") with open(PERSONA_PATH, 'r', encoding='utf-8') as f: self.persona = json.load(f) elapsed = time.time() - t0 print(f"[RAG] Ready in {elapsed:.1f}s") self.ready = True def _precompute_centroids(self): """Precompute L2-normalized centroid for each topic segment.""" centroids = [] for topic in self.topics: start_id = topic["start_msg_id"] end_id = topic["end_msg_id"] # Fast lookup indices = [self.mid_to_index[mid] for mid in range(start_id, end_id + 1) if mid in self.mid_to_index] if indices: chunk = self.embeddings[indices] centroid = chunk.mean(axis=0) norm = np.linalg.norm(centroid) centroids.append(centroid / max(norm, 1e-10)) else: centroids.append(np.zeros(384, dtype=np.float32)) self.topic_centroids = np.array(centroids, dtype=np.float32) def _load_summaries(self): """Load checkpoint summaries from summaries.json (may not exist yet).""" if SUMMARIES_PATH.exists(): with open(SUMMARIES_PATH, 'r', encoding='utf-8') as f: data = json.load(f) self.checkpoints = data.get("checkpoint_summaries", []) # Also update topic summaries if available topic_sum_map = {ts["topic_id"]: ts["summary"] for ts in data.get("topic_summaries", [])} for t in self.topics: if t["topic_id"] in topic_sum_map: t["summary"] = topic_sum_map[t["topic_id"]] print(f"[RAG] Loaded {len(self.checkpoints)} checkpoint summaries") print("[RAG] Embedding topic summaries...") topic_summaries_text = [t.get("summary", "") for t in self.topics] topic_embs = self.embedder.encode(topic_summaries_text, convert_to_numpy=True) t_norms = np.linalg.norm(topic_embs, axis=1, keepdims=True) self.topic_summary_embeddings = topic_embs / np.maximum(t_norms, 1e-10) print("[RAG] Embedding checkpoint summaries...") ckpt_summaries_text = [ck.get("summary", "") for ck in self.checkpoints] if ckpt_summaries_text: ck_embs = self.embedder.encode(ckpt_summaries_text, convert_to_numpy=True) ck_norms = np.linalg.norm(ck_embs, axis=1, keepdims=True) self.checkpoint_summary_embeddings = ck_embs / np.maximum(ck_norms, 1e-10) else: self.checkpoint_summary_embeddings = np.array([]) else: print("[RAG] summaries.json not found - checkpoint retrieval will be empty") self.checkpoints = [] self.topic_summary_embeddings = None self.checkpoint_summary_embeddings = None # Query embedding def embed_query(self, query: str) -> np.ndarray: """Embed and L2-normalize a query string.""" # Pre-process the query to remove exact persona name bias. # Since retrieval already filters by target_user, asking "what is user 1 job" # biases the embedding heavily towards "Hi, I'm user 1". We rewrite to "you/your". processed = query.lower() processed = processed.replace("user 1's", "your") processed = processed.replace("user 2's", "your") processed = processed.replace("user 1", "you") processed = processed.replace("user 2", "you") vec = self.embedder.encode([processed], convert_to_numpy=True)[0] norm = np.linalg.norm(vec) return vec / max(norm, 1e-10) # Retrieval def retrieve_relevant_messages(self, query_emb: np.ndarray, target_user: str = None, target_topic: str = None, top_k: int = 10) -> list[dict]: """Top-k message retrieval via cosine similarity.""" start_id = 0 end_id = float('inf') # --- BULLETPROOF TOPIC FILTER --- if target_topic: # Extract topic_id if it's sent as a dict, string, or int topic_val = None if isinstance(target_topic, dict): topic_val = target_topic.get("topic_id") else: topic_val = target_topic try: # Strip the 't' prefix added by the frontend BEFORE casting to int clean_val = str(topic_val).replace("t", "").strip() topic_id_int = int(clean_val) for t in self.topics: if t["topic_id"] == topic_id_int: start_id = t.get("start_msg_id", 0) end_id = t.get("end_msg_id", float('inf')) break except (ValueError, TypeError): pass # Ignore malformed topic inputs # --- BULLETPROOF USER FILTER --- normalized_user = None if target_user: user_str = str(target_user).lower().replace("_", "").replace(" ", "") if "1" in user_str: normalized_user = "User 1" elif "2" in user_str: normalized_user = "User 2" else: normalized_user = str(target_user).title() # --- APPLY FILTERS --- if normalized_user or (start_id != 0 or end_id != float('inf')): indices = [] for i, mid in enumerate(self.msg_index): if not (start_id <= mid <= end_id): continue if normalized_user and self.messages.get(mid, {}).get("sender") != normalized_user: continue indices.append(i) if not indices: return [] embs = self.embeddings[indices] sims = embs @ query_emb else: indices = list(range(len(self.msg_index))) sims = self.embeddings @ query_emb k = min(top_k, len(sims)) if k == 0: return [] top_local_idx = np.argpartition(sims, -k)[-k:] top_local_idx = top_local_idx[np.argsort(sims[top_local_idx])[::-1]] results = [] for loc_idx in top_local_idx: orig_idx = indices[loc_idx] mid = self.msg_index[orig_idx] rec = self.messages.get(mid, {}) results.append({ "msg_id": mid, "text": rec.get("message_text", ""), "sender": rec.get("sender", ""), "day": rec.get("day", 0), "conversation_id": rec.get("conversation_id", 0), "similarity_score": float(sims[loc_idx]) }) return results def retrieve_relevant_topics(self, query_emb: np.ndarray, top_k: int = 3) -> list[dict]: """Top-k topic retrieval via summary embedding cosine similarity.""" if not hasattr(self, 'topic_summary_embeddings') or self.topic_summary_embeddings is None or len(self.topic_summary_embeddings) == 0: return [] sims = self.topic_summary_embeddings @ query_emb # shape [T] k = min(top_k, len(sims)) top_indices = np.argpartition(sims, -k)[-k:] top_indices = top_indices[np.argsort(sims[top_indices])[::-1]] results = [] for idx in top_indices: t = self.topics[idx] results.append({ "topic_id": t["topic_id"], "start_msg_id": t["start_msg_id"], "end_msg_id": t["end_msg_id"], "start_day": t["start_day"], "end_day": t["end_day"], "msg_range": f"{t['start_msg_id']}-{t['end_msg_id']}", "num_messages": t["num_messages"], "summary": t.get("summary", ""), "similarity_score": float(sims[idx]) }) return results def retrieve_relevant_checkpoints(self, query_emb: np.ndarray, top_k: int = 2) -> list[dict]: """Top-k checkpoint retrieval via summary embedding cosine similarity.""" if not hasattr(self, 'checkpoint_summary_embeddings') or self.checkpoint_summary_embeddings is None or len(self.checkpoint_summary_embeddings) == 0: return [] sims = self.checkpoint_summary_embeddings @ query_emb k = min(top_k, len(sims)) top_indices = np.argpartition(sims, -k)[-k:] top_indices = top_indices[np.argsort(sims[top_indices])[::-1]] results = [] for idx in top_indices: ck = self.checkpoints[idx] results.append({ "checkpoint_id": ck["checkpoint_id"], "msg_range": ck["msg_range"], "start_day": ck["start_day"], "end_day": ck["end_day"], "summary": ck.get("summary", ""), "similarity_score": float(sims[idx]) }) return results # Answer generation def generate_answer( self, query: str, messages: list[dict], topics: list[dict], checkpoints: list[dict] ) -> str: """Generate a grounded answer using flan-t5-small with strict anti-echo prompt.""" context_str = "Topic Summaries:\n" for t in topics: context_str += f"- {t.get('summary', 'N/A')}\n" context_str += "\nCheckpoint Summaries:\n" for c in checkpoints: context_str += f"- {c.get('summary', 'N/A')}\n" context_str += "\nRelevant Messages:\n" for m in messages: context_str += f"- {m['sender']}: {m['text']}\n" prompt = f"""Based ONLY on the context below, answer the user's question. Do not repeat the context. Synthesize a natural answer. If the context doesn't contain the answer, say 'I don't know based on the provided context.' Context: {context_str} Question: {query} Answer:""" try: inputs = self.qa_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024) outputs = self.qa_model.generate(**inputs, max_new_tokens=150, do_sample=False) ans = self.qa_tokenizer.decode(outputs[0], skip_special_tokens=True).strip() return ans except Exception as e: return f"[Generation error: {e}]" # Main query interface def query(self, user_query: str, target_user: str = None, target_topic: str = None, top_k_msgs: int = 10, top_k_topics: int = 3, top_k_checkpoints: int = 2) -> dict: """Full RAG pipeline: embed → retrieve → generate → return.""" if not self.ready: return {"error": "RAG engine not initialized"} query_emb = self.embed_query(user_query) msgs = self.retrieve_relevant_messages(query_emb, target_user=target_user, target_topic=target_topic, top_k=top_k_msgs) topics = self.retrieve_relevant_topics(query_emb, top_k=top_k_topics) checkpoints = self.retrieve_relevant_checkpoints(query_emb, top_k=top_k_checkpoints) max_sim = msgs[0]["similarity_score"] if msgs else 0 if max_sim < 0.30: answer = "I couldn't find any highly relevant messages to answer this question accurately. Please try rephrasing or ask about a different topic." no_results = True else: if is_persona_query(user_query) and msgs: answer = format_persona_answer(user_query, self.persona) else: answer = self.generate_answer(user_query, msgs, topics, checkpoints) no_results = False return { "query": user_query, "answer": answer, "no_results": no_results, "sources": { "topics_used": [ { "id": t["topic_id"], "range": t["msg_range"], "summary": t["summary"], "score": t["similarity_score"] } for t in topics ], "checkpoints_used": [ { "id": ck["checkpoint_id"], "range": ck["msg_range"], "summary": ck["summary"], "score": ck["similarity_score"] } for ck in checkpoints ], "messages_used": [ { "msg_id": m["msg_id"], "text": m["text"][:200], "sender": m["sender"], "day": m["day"], "score": m["similarity_score"] } for m in msgs[:5] ] } } # Persona query helpers PERSONA_KEYWORDS = { "habit": ["habit", "routine", "sleep", "wake", "morning", "night"], "talk": ["talk", "speak", "communicate", "message", "write", "style"], "person": ["person", "personality", "character", "like", "who", "kind of"], "job": ["job", "work", "career", "profession", "do for"], "location": ["live", "from", "location", "city", "country", "where"], "relationship": ["relationship", "family", "friend", "partner", "boyfriend", "girlfriend"], "hobby": ["hobby", "interest", "enjoy", "fun", "leisure"], } def is_persona_query(query: str) -> bool: """Detect if query is specifically about persona/traits vs conversation content.""" q = query.lower() # Strong persona signals - these clearly ask about WHO the person is strong_triggers = [ "what kind of person", "personality", "habits", "how do they talk", "how do they communicate", "how do they speak", "communication style", "describe their personality", "describe the person", "user 1 like", "user 2 like", "what are they like", "profile", "tell me about their habits", "tell me about their personality", "tell me about their communication", "live", "location", "hobbies", "hobby", "who is user", "who was user", "tell me about user" ] return any(trigger in q for trigger in strong_triggers) def format_persona_answer(query: str, global_persona: dict) -> str: """Format global persona data into a natural-language answer.""" q = query.lower() u1 = global_persona.get("persona_user_1", {}) u2 = global_persona.get("persona_user_2", {}) lines = ["Based on an analysis of all 11,000+ conversation threads:\n"] # Detect what aspect they're asking about if any(w in q for w in ["habit", "sleep", "routine", "late", "early"]): lines.append("**Habits:**") for label, u in [("User 1", u1), ("User 2", u2)]: if not u: continue h = u.get("habits", {}) traits = [] if isinstance(h.get("late_sleeper"), dict) and h["late_sleeper"].get("detected"): traits.append(f"late sleeper ({h['late_sleeper']['evidence_count']} content mentions)") if isinstance(h.get("early_bird"), dict) and h["early_bird"].get("detected"): traits.append(f"early bird ({h['early_bird']['evidence_count']} content mentions)") if h.get("brief_communicator"): traits.append("brief communicator") if h.get("verbose_communicator"): traits.append("verbose communicator") lines.append(f" {label}: {', '.join(traits) if traits else 'no distinct habits detected'}") elif any(w in q for w in ["talk", "speak", "communicate", "style", "message"]): lines.append("**Communication Style:**") for label, u in [("User 1", u1), ("User 2", u2)]: if not u: continue s = u.get("communication_style", {}) lines.append(f" {label}:") lines.append(f" • Avg message length: {s.get('avg_message_length', 0)} chars") lines.append(f" • Uses exclamations: {s.get('exclamation_rate', 0)*100:.0f}% of messages") lines.append(f" • Asks questions: {s.get('question_rate', 0)*100:.0f}% of messages") lines.append(f" • Emoji usage: {s.get('emoji_usage_rate', 0)*100:.1f}%") elif any(w in q for w in ["job", "work", "career", "profession"]): lines.append("**Most commonly mentioned jobs:**") for label, u in [("User 1", u1), ("User 2", u2)]: if not u: continue jobs = u.get("personal_facts", {}).get("job_mentions", {}) filtered = {k: v for k, v in jobs.items() if not any(stop in k for stop in ["glad", "sorry", "sure", "same", "free", "relax"])} top = list(filtered.items())[:5] lines.append(f" {label}: {', '.join(f'{j} ({c})' for j,c in top) if top else 'none mentioned'}") elif any(w in q for w in ["personality", "person", "character", "kind of"]): lines.append("**Personality Traits:**") for label, u in [("User 1", u1), ("User 2", u2)]: if not u: continue t = u.get("personality_traits", {}) detected = [name for name, val in t.items() if isinstance(val, dict) and val.get("detected")] lines.append(f" {label}: {', '.join(detected) if detected else 'neutral/mixed'}") elif any(w in q for w in ["location", "live", "from", "where"]): lines.append("**Commonly mentioned locations:**") for label, u in [("User 1", u1), ("User 2", u2)]: if not u: continue locs = u.get("personal_facts", {}).get("location_mentions", {}) top = list(locs.items())[:5] lines.append(f" {label}: {', '.join(f'{l} ({c})' for l,c in top) if top else 'none mentioned'}") else: # General overview lines.append("**Overview of both participants:**") for label, u in [("User 1", u1), ("User 2", u2)]: if not u: continue s = u.get("communication_style", {}) t = u.get("personality_traits", {}) detected_traits = [name for name, val in t.items() if isinstance(val, dict) and val.get("detected")] jobs = list(u.get("personal_facts", {}).get("job_mentions", {}).keys())[:3] filtered_jobs = [j for j in jobs if not any(stop in j for stop in ["glad", "sorry", "sure", "same", "free", "relax"])] lines.append(f" {label} ({u.get('total_messages_analyzed', 0):,} messages):") lines.append(f" - Traits: {', '.join(detected_traits) or 'neutral'}") lines.append(f" - Avg message: {s.get('avg_message_length', 0)} chars") lines.append(f" - Top jobs mentioned: {', '.join(filtered_jobs) if filtered_jobs else 'none'}") return "\n".join(lines)