Spaces:
Running
Running
| """ | |
| memory_service.py | |
| ───────────────── | |
| Manages Multi-Tier Memory for Atlas Document Intelligence: | |
| 1. Short-Term Session Memory: Strict per-session sliding context window. | |
| 2. Long-Term Session Memory: Rolling conversation summaries for long chats. | |
| 3. Document Knowledge Dossiers: Synthesized core findings & entities per document. | |
| 4. Global Knowledge Memory: Collective synthesis across all ingested documents. | |
| """ | |
| import json | |
| import time | |
| from typing import Any, Dict, List, Optional | |
| from core.storage import storage | |
| from utils.logger import app_logger | |
| class MemoryService: | |
| def __init__(self, llm_service) -> None: | |
| self.llm = llm_service | |
| # ── 1. Short-Term & Long-Term Session Memory ──────────────────────────── | |
| def get_session_context( | |
| self, chat_id: str, max_recent_turns: int = 6, owner_id: Optional[str] = None | |
| ) -> Dict[str, Any]: | |
| """ | |
| Retrieves strictly isolated session memory: | |
| - Recent turns (short-term memory window) | |
| - Rolling session summary (long-term memory) | |
| """ | |
| full_history = storage.get_chat_history(chat_id, owner_id=owner_id) or [] | |
| summary_record = storage.get_chat_summary(chat_id) | |
| session_summary = summary_record.get("summary") if summary_record else None | |
| recent_turns = full_history[-max_recent_turns:] if len(full_history) > max_recent_turns else full_history | |
| return { | |
| "chat_id": chat_id, | |
| "session_summary": session_summary, | |
| "recent_turns": recent_turns, | |
| "total_messages": len(full_history), | |
| } | |
| async def update_session_summary_if_needed( | |
| self, chat_id: str, owner_id: Optional[str] = None, force: bool = False | |
| ) -> Optional[str]: | |
| """Trigger rolling summarization if conversation length exceeds threshold or if force=True.""" | |
| try: | |
| full_history = storage.get_chat_history(chat_id, owner_id=owner_id) or [] | |
| if not full_history: | |
| return None | |
| if not force and len(full_history) < 8: | |
| return None | |
| last_summary = storage.get_chat_summary(chat_id) | |
| last_count = last_summary.get("turn_count", 0) if last_summary else 0 | |
| # Only re-summarize if force=True or at least 4 new messages occurred | |
| if not force and (len(full_history) - last_count) < 4: | |
| return last_summary.get("summary") if last_summary else None | |
| app_logger.info(f"[MEMORY] Updating rolling summary for chat_id={chat_id}") | |
| summary = await self.llm.summarize_conversation(full_history) | |
| if summary: | |
| storage.set_chat_summary(chat_id, summary, turn_count=len(full_history)) | |
| app_logger.info(f"[MEMORY] Successfully updated summary for chat_id={chat_id}") | |
| return summary | |
| return None | |
| except Exception as e: | |
| app_logger.error(f"[MEMORY] Error updating session summary: {e}") | |
| return None | |
| # ── 2. Document Knowledge Dossier ────────────────────────────────────── | |
| async def generate_document_dossier( | |
| self, doc_id: str, filename: str, tree: list, flat_sections: Optional[list] = None, owner_id: Optional[str] = None | |
| ) -> Optional[Dict[str, Any]]: | |
| """ | |
| Analyzes a newly indexed document to generate an executive dossier: | |
| - Executive summary | |
| - Key entities | |
| - Key findings | |
| - Topics | |
| """ | |
| try: | |
| # Build representative text preview from first few sections or flat sections | |
| sections_to_read = flat_sections[:8] if flat_sections else [] | |
| preview_texts = [] | |
| for s in sections_to_read: | |
| title = s.get("title", "Section") | |
| text = (s.get("text") or "")[:800] | |
| preview_texts.append(f"### {title}\n{text}") | |
| if not preview_texts: | |
| # Fallback to tree node text | |
| for node in tree[:5]: | |
| if isinstance(node, dict): | |
| t = node.get("title", "") | |
| c = (node.get("text") or "")[:800] | |
| preview_texts.append(f"### {t}\n{c}") | |
| combined_preview = "\n\n".join(preview_texts)[:4000] | |
| if not combined_preview.strip(): | |
| combined_preview = f"Document: {filename} with {len(tree)} outline nodes." | |
| prompt = f"""You are the Document Intelligence Synthesizer for Atlas. | |
| Analyze the following document preview and extract a structured knowledge dossier. | |
| DOCUMENT: {filename} | |
| CONTENT PREVIEW: | |
| {combined_preview} | |
| ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ | |
| OUTPUT REQUIREMENTS | |
| ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ | |
| Return ONLY valid JSON matching this schema: | |
| {{ | |
| "executive_summary": "A 2-3 paragraph comprehensive summary of the document purpose, findings, and themes.", | |
| "key_entities": ["Entity 1", "Entity 2", "Entity 3"], | |
| "key_findings": ["Key finding 1", "Key finding 2", "Key finding 3"], | |
| "topics": ["Topic 1", "Topic 2", "Topic 3", "Topic 4"] | |
| }} | |
| """ | |
| raw_result = "" | |
| if not self.llm.hf_exhausted: | |
| try: | |
| resp = await self.llm.client.chat_completion( | |
| messages=[{"role": "user", "content": prompt}], | |
| model=self.llm.chat_model, | |
| temperature=0.2, | |
| max_tokens=800, | |
| ) | |
| raw_result = resp.choices[0].message.content or "" | |
| except Exception as e: | |
| self.llm._is_hf_exhausted(e) | |
| app_logger.warning(f"HF dossier generation failed ({e}). Checking Groq fallback...") | |
| if (not raw_result or self.llm.hf_exhausted) and self.llm.groq_client: | |
| groq_resp = await self.llm.groq_client.chat.completions.create( | |
| model=self.llm.groq_chat_model, | |
| messages=[{"role": "user", "content": prompt}], | |
| temperature=0.2, | |
| max_tokens=800, | |
| ) | |
| raw_result = groq_resp.choices[0].message.content or "" | |
| data = self.llm._extract_json(raw_result) | |
| summary = data.get("executive_summary") or f"Knowledge dossier for {filename}." | |
| entities = data.get("key_entities") or [] | |
| findings = data.get("key_findings") or [] | |
| topics = data.get("topics") or [] | |
| storage.save_knowledge_dossier( | |
| doc_id=doc_id, | |
| filename=filename, | |
| executive_summary=summary, | |
| key_entities=entities, | |
| key_findings=findings, | |
| topics=topics, | |
| owner_id=owner_id, | |
| ) | |
| app_logger.info(f"[MEMORY] Saved knowledge dossier for doc_id={doc_id} ({filename})") | |
| # Asynchronously update global knowledge summary | |
| await self.synthesize_global_knowledge(owner_id=owner_id) | |
| return { | |
| "doc_id": doc_id, | |
| "filename": filename, | |
| "executive_summary": summary, | |
| "key_entities": entities, | |
| "key_findings": findings, | |
| "topics": topics, | |
| } | |
| except Exception as e: | |
| app_logger.error(f"[MEMORY] Failed to generate document dossier for {doc_id}: {e}", exc_info=True) | |
| def get_document_dossier(self, doc_id: str, owner_id: Optional[str] = None) -> Optional[Dict[str, Any]]: | |
| """Retrieve structured knowledge dossier for an individual document.""" | |
| return storage.get_knowledge_dossier(doc_id, owner_id=owner_id) | |
| def list_knowledge_dossiers(self, owner_id: Optional[str] = None) -> List[Dict[str, Any]]: | |
| """Retrieve all knowledge dossiers across indexed documents.""" | |
| return storage.list_knowledge_dossiers(owner_id=owner_id) | |
| # ── 3. Global Knowledge Summary Across All Data ──────────────────────── | |
| async def synthesize_global_knowledge(self, owner_id: Optional[str] = None) -> Dict[str, Any]: | |
| """ | |
| Synthesizes a collective master knowledge summary of all indexed documents in the workspace. | |
| """ | |
| dossiers = storage.list_knowledge_dossiers(owner_id=owner_id) | |
| total_docs = len(dossiers) | |
| if not dossiers: | |
| # Check if any documents exist in doc_meta without dossiers yet | |
| all_docs = storage.list_documents(owner_id=owner_id) | |
| if not all_docs: | |
| default_summary = "No documents currently indexed in the knowledge base." | |
| storage.save_global_knowledge_summary(default_summary, 0, [], owner_id=owner_id) | |
| return {"collective_summary": default_summary, "total_documents": 0, "key_insights": []} | |
| # Build combined summary representation of all dossiers | |
| dossier_summaries = [] | |
| for d in dossiers: | |
| fname = d.get("filename", "Doc") | |
| summ = d.get("executive_summary", "")[:600] | |
| findings = "; ".join(d.get("key_findings", [])[:3]) | |
| topics = ", ".join(d.get("topics", [])[:4]) | |
| dossier_summaries.append( | |
| f"- **{fname}**\n Summary: {summ}\n Topics: {topics}\n Findings: {findings}" | |
| ) | |
| combined_input = "\n\n".join(dossier_summaries) | |
| prompt = f"""You are the Chief Knowledge Architect for Atlas Document Intelligence. | |
| Synthesize the collective knowledge across all {total_docs} documents ingested into the platform. | |
| INDEXED DOCUMENTS DOSSIERS: | |
| {combined_input} | |
| ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ | |
| OUTPUT REQUIREMENTS | |
| ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ | |
| Return ONLY valid JSON with this exact schema: | |
| {{ | |
| "collective_summary": "A cohesive 2-4 paragraph overview of the global knowledge contained in all documents, highlighting overarching themes, areas of synergy, and domain expertise.", | |
| "key_insights": [ | |
| "Cross-cutting insight 1", | |
| "Cross-cutting insight 2", | |
| "Cross-cutting insight 3" | |
| ] | |
| }} | |
| """ | |
| raw_result = "" | |
| if not self.llm.hf_exhausted: | |
| try: | |
| resp = await self.llm.client.chat_completion( | |
| messages=[{"role": "user", "content": prompt}], | |
| model=self.llm.chat_model, | |
| temperature=0.3, | |
| max_tokens=1000, | |
| ) | |
| raw_result = resp.choices[0].message.content or "" | |
| except Exception as e: | |
| self.llm._is_hf_exhausted(e) | |
| app_logger.warning(f"HF global knowledge synthesis failed ({e}). Checking Groq fallback...") | |
| if (not raw_result or self.llm.hf_exhausted) and self.llm.groq_client: | |
| try: | |
| groq_resp = await self.llm.groq_client.chat.completions.create( | |
| model=self.llm.groq_chat_model, | |
| messages=[{"role": "user", "content": prompt}], | |
| temperature=0.3, | |
| max_tokens=1000, | |
| ) | |
| raw_result = groq_resp.choices[0].message.content or "" | |
| except Exception as ge: | |
| app_logger.error(f"Groq global knowledge synthesis failed: {ge}") | |
| data = self.llm._extract_json(raw_result) | |
| collective = data.get("collective_summary") or "Atlas Global Knowledge Base synthesis complete." | |
| insights = data.get("key_insights") or [] | |
| storage.save_global_knowledge_summary(collective, total_docs, insights, owner_id=owner_id) | |
| app_logger.info(f"[MEMORY] Successfully synthesized global knowledge across {total_docs} documents.") | |
| return { | |
| "collective_summary": collective, | |
| "total_documents": total_docs, | |
| "key_insights": insights, | |
| "updated_at": time.time(), | |
| } | |
| def get_global_knowledge(self, owner_id: Optional[str] = None) -> Dict[str, Any]: | |
| """Fetch current global knowledge summary and document dossiers.""" | |
| global_rec = storage.get_global_knowledge_summary(owner_id=owner_id) or {} | |
| dossiers = storage.list_knowledge_dossiers(owner_id=owner_id) | |
| return { | |
| "collective_summary": global_rec.get( | |
| "collective_summary", "No collective summary available yet. Ingest documents to build knowledge." | |
| ), | |
| "total_documents": global_rec.get("total_documents", len(dossiers)), | |
| "key_insights": global_rec.get("key_insights", []), | |
| "updated_at": global_rec.get("updated_at"), | |
| "dossiers": dossiers, | |
| } | |