Spaces:
Running
Running
| """ | |
| NeuraPrompt Agent — Memory Module v7.6 (MongoDB Integrated) | |
| Properly connects to your existing MongoDB system from main.py | |
| """ | |
| from typing import List, Dict, Optional | |
| import logging | |
| from datetime import datetime, timezone | |
| log = logging.getLogger("agent.memory.v7.6") | |
| # These will be injected from main.py | |
| long_term_memory_col = None | |
| chat_history_col = None | |
| def set_collections(long_term_col, chat_history_col_ref): | |
| """Call this from main.py to inject MongoDB collections""" | |
| global long_term_memory_col, chat_history_col | |
| long_term_memory_col = long_term_col | |
| chat_history_col = chat_history_col_ref | |
| class MemoryManager: | |
| def __init__(self, user_id: str): | |
| self.user_id = user_id | |
| self.short_term: List[Dict] = [] | |
| def load_long_term_memory(self) -> Dict: | |
| """Load user's long-term memory from MongoDB""" | |
| if not long_term_memory_col: | |
| return {} | |
| try: | |
| doc = long_term_memory_col.find_one({"user_id": self.user_id}) | |
| return doc or {} | |
| except Exception as e: | |
| log.error(f"Failed to load long-term memory: {e}") | |
| return {} | |
| def load_recent_chat(self, limit: int = 8) -> List[Dict]: | |
| """Load recent chat history from MongoDB""" | |
| if not chat_history_col: | |
| return [] | |
| try: | |
| cursor = chat_history_col.find( | |
| {"user_id": self.user_id} | |
| ).sort("timestamp", -1).limit(limit) | |
| messages = list(cursor) | |
| messages.reverse() | |
| return messages | |
| except Exception as e: | |
| log.error(f"Failed to load chat history: {e}") | |
| return [] | |
| def get_context_for_agent(self) -> str: | |
| """Build rich context string for the agent""" | |
| long_term = self.load_long_term_memory() | |
| recent = self.load_recent_chat(6) | |
| context_parts = [] | |
| # Long-term facts | |
| if long_term: | |
| facts = [] | |
| for key, value in long_term.items(): | |
| if key not in ["_id", "user_id", "last_updated"]: | |
| facts.append(f"{key}: {value}") | |
| if facts: | |
| context_parts.append("Known about user: " + ", ".join(facts[:7])) | |
| # Recent conversation | |
| if recent: | |
| chat_lines = [] | |
| for msg in recent[-5:]: | |
| role = msg.get("role", "unknown").capitalize() | |
| content = msg.get("content", "")[:180] | |
| chat_lines.append(f"{role}: {content}") | |
| context_parts.append("Recent conversation:\n" + "\n".join(chat_lines)) | |
| return "\n\n".join(context_parts) if context_parts else "No previous memory available." | |
| def save_important_fact(self, key: str, value: str): | |
| """Save important fact to long-term memory""" | |
| if not long_term_memory_col: | |
| return | |
| try: | |
| long_term_memory_col.update_one( | |
| {"user_id": self.user_id}, | |
| {"$set": { | |
| key: value, | |
| "last_updated": datetime.now(timezone.utc) | |
| }}, | |
| upsert=True | |
| ) | |
| log.info(f"Saved fact '{key}' for user {self.user_id}") | |
| except Exception as e: | |
| log.error(f"Failed to save fact: {e}") | |
| def get_memory_manager(user_id: str) -> MemoryManager: | |
| return MemoryManager(user_id) | |