khoaliamle/Cooking_Tutor
0
1# memory_updated.py2import re, time, hashlib, asyncio, os3from collections import defaultdict, deque4from typing import List, Dict5import numpy as np6import faiss7from sentence_transformers import SentenceTransformer8from google import genai # must be configured in app.py and imported globally9import logging10from models.summarizer import summarizer11 12_LLM_SMALL = "gemini-2.5-flash-lite-preview-06-17"13# Load embedding model - use standard model that downloads automatically14EMBED = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")15logger = logging.getLogger("rag-agent")16logging.basicConfig(level=logging.INFO, format="%(asctime)s — %(name)s — %(levelname)s — %(message)s", force=True) # Change INFO to DEBUG for full-ctx JSON loader17 18api_key = os.getenv("FlashAPI")19# Initialize client lazily to avoid errors during import20client = None21 22def get_genai_client():23 """Get or create Gemini client"""24 global client25 if client is None and api_key:26 client = genai.Client(api_key=api_key)27 return client28 29class MemoryManager:30 def __init__(self, max_users=1000, history_per_user=20, max_chunks=60):31 # STM: recent conversation summaries (topic + summary), up to 5 entries32 self.stm_summaries = defaultdict(lambda: deque(maxlen=history_per_user)) # deque of {topic,text,vec,timestamp,used}33 # Legacy raw cache (kept for compatibility if needed)34 self.text_cache = defaultdict(lambda: deque(maxlen=history_per_user))35 # LTM: semantic chunk store (approx 3 chunks x 20 rounds)36 self.chunk_index = defaultdict(self._new_index) # user_id -> faiss index37 self.chunk_meta = defaultdict(list) # '' -> list[{text,tag,vec,timestamp,used}]38 self.user_queue = deque(maxlen=max_users) # LRU of users39 self.max_chunks = max_chunks # hard cap per user40 self.chunk_cache = {} # hash(query+resp) -> [chunks]41 42 # ---------- Public API ----------43 def add_exchange(self, user_id: str, query: str, response: str, lang: str = "EN"):44 self._touch_user(user_id)45 # Keep raw record (optional)46 self.text_cache[user_id].append(((query or "").strip(), (response or "").strip()))47 if not response: return []48 # Avoid re-chunking identical response49 cache_key = hashlib.md5((query + response).encode()).hexdigest()50 if cache_key in self.chunk_cache:51 chunks = self.chunk_cache[cache_key]52 else:53 chunks = self.chunk_response(response, lang, question=query)54 self.chunk_cache[cache_key] = chunks55 # Update STM with merging/deduplication56 for chunk in chunks:57 self._upsert_stm(user_id, chunk, lang)58 # Update LTM with merging/deduplication59 self._upsert_ltm(user_id, chunks, lang)60 return chunks61 62 def get_relevant_chunks(self, user_id: str, query: str, top_k: int = 3, min_sim: float = 0.30) -> List[str]:63 """Return texts of chunks whose cosine similarity ≥ min_sim."""64 if self.chunk_index[user_id].ntotal == 0:65 return []66 # Encode chunk67 qvec = self._embed(query)68 sims, idxs = self.chunk_index[user_id].search(np.array([qvec]), k=top_k)69 results = []70 # Append related result with smart-decay to optimize storage and prioritize most-recent chat71 for sim, idx in zip(sims[0], idxs[0]):72 if idx < len(self.chunk_meta[user_id]) and sim >= min_sim:73 chunk = self.chunk_meta[user_id][idx]74 chunk["used"] += 1 # increment usage75 # Decay function76 age_sec = time.time() - chunk["timestamp"]77 decay = 1.0 / (1.0 + age_sec / 300) # 5-min half-life78 score = sim * decay * (1 + 0.1 * chunk["used"])79 # Append chunk with score80 results.append((score, chunk))81 # Sort result on best scored82 results.sort(key=lambda x: x[0], reverse=True)83 # logger.info(f"[Memory] RAG Retrieved Topic: {results}") # Inspect vector data84 return [f"### Topic: {c['tag']}\n{c['text']}" for _, c in results]85 86 def get_recent_chat_history(self, user_id: str, num_turns: int = 5) -> List[Dict]:87 """88 Get the most recent short-term memory summaries.89 Returns: a list of entries containing only the summarized bot context.90 """91 if user_id not in self.stm_summaries:92 return []93 recent = list(self.stm_summaries[user_id])[-num_turns:]94 formatted = []95 for entry in recent:96 formatted.append({97 "user": "",98 "bot": f"Topic: {entry['topic']}\n{entry['text']}",99 "timestamp": entry.get("timestamp", time.time())100 })101 return formatted102 103 def get_context(self, user_id: str, num_turns: int = 5) -> str:104 # Prefer STM summaries105 history = self.get_recent_chat_history(user_id, num_turns=num_turns)106 return "\n".join(h["bot"] for h in history)107 108 def get_contextual_chunks(self, user_id: str, current_query: str, lang: str = "EN") -> str:109 """110 Use NVIDIA Llama to create a summarization of relevant context from both recent history and RAG chunks.111 This ensures conversational continuity while providing a concise summary for the main LLM.112 """113 # Get both types of context114 recent_history = self.get_recent_chat_history(user_id, num_turns=5)115 rag_chunks = self.get_relevant_chunks(user_id, current_query, top_k=3)116 117 logger.info(f"[Contextual] Retrieved {len(recent_history)} recent history items")118 logger.info(f"[Contextual] Retrieved {len(rag_chunks)} RAG chunks")119 120 # Return empty string if no context is found121 if not recent_history and not rag_chunks:122 logger.info(f"[Contextual] No context found, returning empty string")123 return ""124 125 # Prepare context for summarization126 context_parts = []127 # Add recent chat history128 if recent_history:129 history_text = "\n".join([130 f"User: {item['user']}\nBot: {item['bot']}"131 for item in recent_history132 ])133 context_parts.append(f"Recent conversation history:\n{history_text}")134 # Add RAG chunks135 if rag_chunks:136 rag_text = "\n".join(rag_chunks)137 context_parts.append(f"Semantically relevant historical cooking information:\n{rag_text}")138 139 # Combine all context140 full_context = "\n\n".join(context_parts)141 142 # Use summarizer to create concise summary143 try:144 summary = summarizer.summarize_text(full_context, max_length=300)145 logger.info(f"[Contextual] Generated summary using NVIDIA Llama: {len(summary)} characters")146 return summary147 except Exception as e:148 logger.error(f"[Contextual] Summarization failed: {e}")149 return full_context[:500] + "..." if len(full_context) > 500 else full_context150 151 def chunk_response(self, response: str, lang: str, question: str = "") -> List[Dict]:152 """153 Use NVIDIA Llama to chunk and summarize response by cooking topics.154 Returns: [{"tag": ..., "text": ...}, ...]155 """156 if not response: 157 return []158 159 try:160 # Use summarizer to chunk and summarize161 chunks = summarizer.chunk_response(response, max_chunk_size=500)162 163 # Convert to the expected format164 result_chunks = []165 for i, chunk in enumerate(chunks):166 # Extract topic from chunk (first sentence or key cooking terms)167 topic = self._extract_topic_from_chunk(chunk)168 169 result_chunks.append({170 "tag": topic,171 "text": chunk172 })173 174 logger.info(f"[Memory] 📦 NVIDIA Llama summarized {len(result_chunks)} chunks")175 return result_chunks176 177 except Exception as e:178 logger.error(f"[Memory] NVIDIA Llama chunking failed: {e}")179 # Fallback to simple chunking180 return self._fallback_chunking(response)181 182 def _extract_topic_from_chunk(self, chunk: str) -> str:183 """Extract a concise topic from a chunk"""184 # Look for cooking terms or first sentence185 sentences = chunk.split('.')186 if sentences:187 first_sentence = sentences[0].strip()188 if len(first_sentence) > 50:189 first_sentence = first_sentence[:50] + "..."190 return first_sentence191 return "Cooking Information"192 193 def _fallback_chunking(self, response: str) -> List[Dict]:194 """Fallback chunking when NVIDIA Llama fails"""195 # Simple sentence-based chunking196 sentences = re.split(r'[.!?]+', response)197 chunks = []198 current_chunk = ""199 200 for sentence in sentences:201 sentence = sentence.strip()202 if not sentence:203 continue204 205 if len(current_chunk) + len(sentence) > 300:206 if current_chunk:207 chunks.append({208 "tag": "Cooking Information",209 "text": current_chunk.strip()210 })211 current_chunk = sentence212 else:213 current_chunk += sentence + ". "214 215 if current_chunk:216 chunks.append({217 "tag": "Cooking Information", 218 "text": current_chunk.strip()219 })220 221 return chunks222 223 # ---------- Private Methods ----------224 def _touch_user(self, user_id: str):225 """Update LRU queue"""226 if user_id in self.user_queue:227 self.user_queue.remove(user_id)228 self.user_queue.append(user_id)229 230 def _new_index(self):231 """Create new FAISS index"""232 return faiss.IndexFlatIP(384) # 384-dim embeddings233 234 def _upsert_stm(self, user_id: str, chunk: Dict, lang: str):235 """Update short-term memory with merging/deduplication"""236 topic = chunk["tag"]237 text = chunk["text"]238 239 # Check for similar topics in STM240 for entry in self.stm_summaries[user_id]:241 if self._topics_similar(topic, entry["topic"]):242 # Merge with existing entry243 entry["text"] = summarizer.summarize_text(244 f"{entry['text']}\n{text}", 245 max_length=200246 )247 entry["timestamp"] = time.time()248 return249 250 # Add new entry251 self.stm_summaries[user_id].append({252 "topic": topic,253 "text": text,254 "vec": self._embed(f"{topic} {text}"),255 "timestamp": time.time(),256 "used": 0257 })258 259 def _upsert_ltm(self, user_id: str, chunks: List[Dict], lang: str):260 """Update long-term memory with merging/deduplication"""261 for chunk in chunks:262 # Check for similar chunks in LTM263 similar_idx = self._find_similar_chunk(user_id, chunk["text"])264 265 if similar_idx is not None:266 # Merge with existing chunk267 existing = self.chunk_meta[user_id][similar_idx]268 merged_text = summarizer.summarize_text(269 f"{existing['text']}\n{chunk['text']}", 270 max_length=300271 )272 existing["text"] = merged_text273 existing["timestamp"] = time.time()274 else:275 # Add new chunk276 if len(self.chunk_meta[user_id]) >= self.max_chunks:277 # Remove oldest chunk278 self._remove_oldest_chunk(user_id)279 280 vec = self._embed(chunk["text"])281 self.chunk_index[user_id].add(np.array([vec]))282 self.chunk_meta[user_id].append({283 "text": chunk["text"],284 "tag": chunk["tag"],285 "vec": vec,286 "timestamp": time.time(),287 "used": 0288 })289 290 def _topics_similar(self, topic1: str, topic2: str) -> bool:291 """Check if two topics are similar"""292 # Simple similarity check based on common words293 words1 = set(topic1.lower().split())294 words2 = set(topic2.lower().split())295 intersection = words1.intersection(words2)296 return len(intersection) >= 2297 298 def _find_similar_chunk(self, user_id: str, text: str) -> int:299 """Find similar chunk in LTM"""300 if not self.chunk_meta[user_id]:301 return None302 303 text_vec = self._embed(text)304 sims, idxs = self.chunk_index[user_id].search(np.array([text_vec]), k=3)305 306 for sim, idx in zip(sims[0], idxs[0]):307 if sim > 0.8: # High similarity threshold308 return int(idx)309 return None310 311 def _remove_oldest_chunk(self, user_id: str):312 """Remove the oldest chunk from LTM"""313 if not self.chunk_meta[user_id]:314 return315 316 # Find oldest chunk317 oldest_idx = min(range(len(self.chunk_meta[user_id])), 318 key=lambda i: self.chunk_meta[user_id][i]["timestamp"])319 320 # Remove from index and metadata321 self.chunk_meta[user_id].pop(oldest_idx)322 # Note: FAISS doesn't support direct removal, so we rebuild the index323 self._rebuild_index(user_id)324 325 def _rebuild_index(self, user_id: str):326 """Rebuild FAISS index after removal"""327 if not self.chunk_meta[user_id]:328 self.chunk_index[user_id] = self._new_index()329 return330 331 vectors = [chunk["vec"] for chunk in self.chunk_meta[user_id]]332 self.chunk_index[user_id] = self._new_index()333 self.chunk_index[user_id].add(np.array(vectors))334 335 @staticmethod336 def _embed(text: str):337 vec = EMBED.encode(text, convert_to_numpy=True)338 # L2 normalise for cosine on IndexFlatIP339 return vec / (np.linalg.norm(vec) + 1e-9)340