memory_handler.py

import os
import streamlit as st
from mem0 import MemoryClient
from dotenv import load_dotenv

load_dotenv()


@st.cache_resource
def get_mem0_client():
    """Initializes and caches the Mem0 MemoryClient using Streamlit secrets or environment variables."""
    try:
        api_key = st.secrets["MEM0_API_KEY"]
    except Exception:
        api_key = os.getenv("MEM0_API_KEY", "")

    if not api_key:
        return None
    
    try:
        return MemoryClient(api_key=api_key)
    except Exception:
        return None


class MemoryHandler:
    """Manages fact storage, retrieval, intelligent deduplication, and deletion using standard Mem0 v2 API specifications."""

    def __init__(self):
        self.client = get_mem0_client()

    def is_configured(self) -> bool:
        """Returns True if the Mem0 client is successfully initialized."""
        return self.client is not None

    def add_interaction(self, user_id: str, user_message: str, assistant_message: str, llm_connector=None, existing_memories: list[dict] = None, **kwargs) -> bool:
        """Intelligently extracts core user facts, preferences, updates outdated memories, and stores clean statements in Mem0."""
        if not self.is_configured() or not user_id:
            return False
        
        # Step 1: Use LLM connector to extract durable user facts and preferences
        if llm_connector and hasattr(llm_connector, "extract_user_facts"):
            extracted_facts = llm_connector.extract_user_facts(user_message)
        else:
            extracted_facts = [user_message]

        if not extracted_facts:
            # Skip storing if prompt contains no durable facts or preferences
            return True

        if existing_memories is None:
            existing_memories = self.get_all_memories(user_id)

        # Step 2: Process each extracted fact
        for fact in extracted_facts:
            target_id = None
            if llm_connector and hasattr(llm_connector, "find_target_memory_to_update") and existing_memories:
                target_id = llm_connector.find_target_memory_to_update(fact, existing_memories)

            if target_id:
                # Update existing memory item (e.g. status change or name update)
                self.update_memory(target_id, fact)
            else:
                # Add as a new memory fact (e.g. new preference or distinct topic)
                try:
                    self.client.add(fact, user_id=user_id)
                except Exception as e:
                    st.error(f"Error adding memory fact: {e}")

        return True

    def search_memories(self, user_id: str, query: str) -> list[str]:
        """Retrieves relevant facts from Mem0 for context injection using standard Mem0 v2 filters."""
        if not self.is_configured() or not user_id or not query:
            return []

        try:
            response = self.client.search(query, filters={"user_id": user_id})
            return self._parse_memory_texts(response)
        except Exception as e:
            st.warning(f"Unable to search memories: {e}")
            return []

    def get_all_memories(self, user_id: str) -> list[dict]:
        """Fetches all stored memories for a specific user using standard Mem0 v2 filters."""
        if not self.is_configured() or not user_id:
            return []

        try:
            response = self.client.get_all(filters={"user_id": user_id})
            if isinstance(response, dict) and "results" in response:
                return response["results"]
            elif isinstance(response, list):
                return response
            return []
        except Exception as e:
            st.warning(f"Unable to retrieve memories: {e}")
            return []

    def update_memory(self, memory_id: str, new_text: str) -> bool:
        """Updates an existing memory item by memory_id using Mem0's native update endpoint."""
        if not self.is_configured() or not memory_id or not new_text:
            return False

        try:
            self.client.update(memory_id=memory_id, text=new_text)
            return True
        except Exception as e:
            st.error(f"Failed to update memory item: {e}")
            return False

    def delete_single_memory(self, memory_id: str) -> bool:
        """Deletes a specific memory item by memory_id."""
        if not self.is_configured() or not memory_id:
            return False

        try:
            self.client.delete(memory_id=memory_id)
            return True
        except Exception as e:
            st.error(f"Failed to delete memory item: {e}")
            return False

    def delete_all_memories(self, user_id: str) -> bool:
        """Deletes all stored memories for a specific user ID."""
        if not self.is_configured() or not user_id:
            return False

        try:
            self.client.delete_all(user_id=user_id)
            return True
        except Exception as e:
            st.error(f"Failed to clear memories: {e}")
            return False

    def _parse_memory_texts(self, response) -> list[str]:
        """Helper to extract clean memory text strings from Mem0 response objects."""
        memories = []
        items = response.get("results", []) if isinstance(response, dict) else response
        
        if isinstance(items, list):
            for item in items:
                if isinstance(item, dict):
                    memory_str = item.get("memory") or item.get("text") or item.get("content")
                    if memory_str:
                        memories.append(str(memory_str))
                elif isinstance(item, str):
                    memories.append(item)
        return memories