memory/vector_db.py

"""
memory/vector_db.py
===================
Módulo de memoria a largo plazo para Aurora Core.

Utiliza ChromaDB en modo persistente (sin servidor externo) para almacenar
recuerdos del usuario como embeddings semánticos y recuperarlos por similitud.

Dependencias:
    pip install chromadb

Modelo de embeddings:
    Por defecto: all-MiniLM-L6-v2 vía ONNX (incluido en chromadb).
    Primera ejecución: el modelo (~45 MB) se descarga automáticamente
    en ~/.cache/chroma/onnx_models/. Las siguientes ejecuciones son
    completamente offline.

Estructura en disco generada:
    .aurora_memory_db/
    └── <colección SQLite + índice HNSW>
"""

from __future__ import annotations

import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import chromadb


# ── Constantes ────────────────────────────────────────────────────────
_DB_PATH        = ".aurora_memory_db"      # directorio de persistencia local
_COLLECTION_NAME = "user_long_term_memory"  # una colección por usuario (prefijada con user_id)


class AuroraMemory:
    """
    Interfaz de memoria a largo plazo para Aurora Core.

    Cada instancia opera sobre una colección ChromaDB aislada por ``user_id``,
    lo que permite múltiples perfiles de usuario en la misma base de datos sin
    colisiones de datos.

    Example::

        mem = AuroraMemory(user_id="gabi_master")
        mem.save_memory("Aurora corre en Python 3.11 con FastAPI")
        contexto = mem.retrieve_relevant_context("¿qué stack tecnológico usamos?")
    """

    def __init__(self, user_id: str = "gabi_master") -> None:
        """
        Inicializa el cliente ChromaDB persistente y obtiene (o crea) la
        colección de memoria para el usuario indicado.

        La colección se nombra ``<user_id>__user_long_term_memory`` para
        garantizar aislamiento entre usuarios dentro de la misma carpeta de DB.

        Args:
            user_id: Identificador único del usuario. El nombre de la colección
                     se derivará de este valor para aislar recuerdos entre perfiles.

        Raises:
            Exception: Si ChromaDB no puede inicializar o acceder al directorio
                       de persistencia (permisos, disco lleno, etc.).
        """
        self.user_id = user_id

        # Asegurar que el directorio de persistencia exista
        Path(_DB_PATH).mkdir(parents=True, exist_ok=True)

        # Cliente persistente — sin servidor, todo local en disco
        self._client = chromadb.PersistentClient(path=_DB_PATH)

        # Nombre de colección con prefijo de usuario para aislamiento
        collection_name = f"{user_id}__{_COLLECTION_NAME}"

        # get_or_create_collection: idempotente — no falla si ya existe
        self._collection = self._client.get_or_create_collection(
            name=collection_name,
            metadata={
                "hnsw:space": "cosine",   # similitud coseno para texto semántico
                "description": f"Memoria a largo plazo de {user_id} — Aurora Core",
            },
        )

    # ── save_memory ───────────────────────────────────────────────────

    def save_memory(self, text: str, metadata: Optional[dict] = None) -> str:
        """
        Guarda un fragmento de texto como recuerdo persistente en la base de datos.

        El texto se convierte automáticamente en embedding semántico por ChromaDB
        (usando all-MiniLM-L6-v2) y se indexa para búsqueda por similitud futura.

        Args:
            text:     Fragmento de texto a recordar.
                      Ejemplos:
                        - "Gabi prefiere trabajar de noche en proyectos de IA"
                        - "El proyecto Aurora usa Hostinger para el deploy"
                        - "A Gabi le gusta el café con mucho azúcar"
            metadata: Diccionario opcional de metadatos adicionales (str/int/float/bool).
                      Se combina con los metadatos automáticos (timestamp, user_id, etc.).
                      Ejemplo: {"categoria": "tecnico", "proyecto": "aurora"}

        Returns:
            El ID único generado para este recuerdo (UUID v4 como string).

        Raises:
            ValueError: Si ``text`` es una cadena vacía.
            Exception:  Si ChromaDB falla al escribir (disco lleno, DB bloqueada, etc.).
        """
        if not text or not text.strip():
            raise ValueError("El texto del recuerdo no puede estar vacío.")

        memory_id = str(uuid.uuid4())

        # Metadatos automáticos siempre presentes
        base_metadata: dict = {
            "user_id":   self.user_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "text_len":  len(text),
        }

        # Fusionar con metadatos adicionales del llamador (tienen prioridad)
        if metadata:
            # ChromaDB solo acepta str/int/float/bool en metadatos — filtrar otros tipos
            safe_extra = {
                k: v for k, v in metadata.items()
                if isinstance(v, (str, int, float, bool))
            }
            base_metadata.update(safe_extra)

        self._collection.add(
            ids=[memory_id],
            documents=[text.strip()],
            metadatas=[base_metadata],
        )

        return memory_id

    # ── retrieve_relevant_context ─────────────────────────────────────

    def retrieve_relevant_context(
        self,
        query: str,
        n_results: int = 3,
    ) -> list[str]:
        """
        Recupera los recuerdos más relevantes para el mensaje actual del usuario.

        Realiza una búsqueda semántica por similitud coseno en el espacio de
        embeddings. El resultado puede inyectarse directamente en el system prompt
        de la IA para darle contexto personalizado.

        Args:
            query:     Texto del mensaje actual del usuario (o una reformulación
                       temática del mismo para guiar la búsqueda).
            n_results: Número máximo de recuerdos a devolver. Se ajusta
                       automáticamente si la DB tiene menos documentos que este
                       valor (evita errores de ChromaDB).

        Returns:
            Lista de textos de los recuerdos más relevantes, ordenados de mayor
            a menor similitud semántica. Lista vacía si no hay recuerdos.

        Example::

            contexto = mem.retrieve_relevant_context("¿en qué proyecto trabajamos?")
            # → ["Aurora usa Hostinger para el deploy", "El backend corre en Python 3.11"]
        """
        if not query or not query.strip():
            return []

        total = self._collection.count()
        if total == 0:
            return []

        # Evitar que n_results > documentos existentes (error de ChromaDB)
        safe_n = min(n_results, total)

        try:
            results = self._collection.query(
                query_texts=[query.strip()],
                n_results=safe_n,
                include=["documents", "distances", "metadatas"],
            )
        except Exception as exc:
            # La DB puede estar temporalmente bloqueada en entornos multi-proceso
            print(f"[AuroraMemory] Error al consultar la memoria: {exc}")
            return []

        # results["documents"] → lista de listas (una por query_text)
        # Tomamos la primera (y única) consulta
        docs: list[str] = results["documents"][0] if results["documents"] else []
        return docs

    # ── borrar_todo ───────────────────────────────────────────────────

    def borrar_todo(self) -> dict:
        """
        Elimina **todos** los recuerdos de este usuario. Operación irreversible.

        La llama ``/delete-account`` al dar de baja una cuenta. Si esto no borra
        de verdad, queda texto literal de conversaciones en disco después de que
        alguien pidió expresamente que lo borren.

        Estrategia: ``delete_collection()`` y no ``collection.delete()``.

        La diferencia importa. ``delete()`` saca los documentos del índice, pero
        los archivos del segmento HNSW y las filas de SQLite siguen ocupando
        disco: los datos quedan recuperables con herramientas forenses básicas.
        ``delete_collection()`` es lo único que elimina la colección completa del
        almacenamiento. Y acá se puede usar porque cada usuario tiene su propia
        colección (``<user_id>__user_long_term_memory``), así que borrarla no
        toca los recuerdos de nadie más.

        Después del borrado la colección se **recrea vacía**, para que la
        instancia siga siendo usable: si otro hilo tuviera una referencia y
        llamara a ``save_memory`` o ``retrieve_relevant_context``, obtendría una
        memoria vacía en vez de una excepción.

        Returns:
            Diccionario con el resultado:

            - ``user_id`` (str): usuario afectado.
            - ``collection_name`` (str): colección eliminada.
            - ``borrados`` (int): recuerdos que había antes del borrado.
            - ``ok`` (bool): si el borrado se completó.
            - ``error`` (str, opcional): motivo si ``ok`` es False.

        Example::

            mem = AuroraMemory(user_id="gabi_master")
            informe = mem.borrar_todo()
            # → {"user_id": "gabi_master", "borrados": 12, "ok": True, ...}
        """
        nombre = self._collection.name

        try:
            previos = self._collection.count()
        except Exception:
            previos = -1

        informe: dict = {
            "user_id":         self.user_id,
            "collection_name": nombre,
            "borrados":        previos,
            "ok":              False,
        }

        try:
            self._client.delete_collection(name=nombre)

            # Recrear vacía con los MISMOS metadatos que en __init__, para que
            # una instancia reutilizada se comporte igual que una recién creada.
            self._collection = self._client.get_or_create_collection(
                name=nombre,
                metadata={
                    "hnsw:space": "cosine",
                    "description": f"Memoria a largo plazo de {self.user_id} — Aurora Core",
                },
            )

            informe["ok"] = True
            print(f"[AuroraMemory] Memoria de '{self.user_id}' borrada "
                  f"({previos} recuerdos eliminados).")

        except Exception as exc:
            # No se relanza: el borrado de cuenta tiene que poder seguir con los
            # demás pasos e informar qué falló, en vez de cortarse por la mitad
            # y dejar al usuario sin saber qué quedó borrado y qué no.
            informe["error"] = f"{type(exc).__name__}: {exc}"
            print(f"[AuroraMemory] FALLÓ el borrado de '{self.user_id}': {exc}")

        return informe

    # ── get_memory_stats ──────────────────────────────────────────────

    def get_memory_stats(self) -> dict:
        """
        Devuelve estadísticas básicas del estado actual de la memoria.

        Returns:
            Diccionario con las siguientes claves:

            - ``total_memories`` (int): Número de recuerdos almacenados.
            - ``user_id`` (str): Identificador del usuario de esta instancia.
            - ``collection_name`` (str): Nombre interno de la colección en ChromaDB.
            - ``db_path`` (str): Ruta absoluta del directorio de persistencia.

        Example::

            stats = mem.get_memory_stats()
            # → {"total_memories": 12, "user_id": "gabi_master", ...}
        """
        try:
            total = self._collection.count()
        except Exception:
            total = -1  # -1 indica que no se pudo acceder a la colección

        return {
            "total_memories":  total,
            "user_id":         self.user_id,
            "collection_name": self._collection.name,
            "db_path":         str(Path(_DB_PATH).resolve()),
        }


# ── Bloque de prueba ──────────────────────────────────────────────────

if __name__ == "__main__":
    print("=" * 58)
    print("  Aurora Core — Test de memoria vectorial a largo plazo")
    print("=" * 58)

    # 1. Inicializar
    print("\n[1] Inicializando AuroraMemory...")
    mem = AuroraMemory(user_id="gabi_master")
    stats_inicial = mem.get_memory_stats()
    print(f"    Recuerdos previos en DB: {stats_inicial['total_memories']}")
    print(f"    Colección: {stats_inicial['collection_name']}")
    print(f"    Ruta DB:   {stats_inicial['db_path']}")

    # 2. Guardar recuerdos de prueba
    print("\n[2] Guardando recuerdos de prueba...")

    id_tecnico = mem.save_memory(
        text="El proyecto Aurora tiene su backend en Python con FastAPI "
             "y está alojado en Hostinger con un VPS Ubuntu 22.04.",
        metadata={"categoria": "tecnico", "proyecto": "aurora"},
    )
    print(f"    ✓ Recuerdo técnico guardado  (id: {id_tecnico[:8]}...)")

    id_personal = mem.save_memory(
        text="Gabi prefiere programar de madrugada, entre las 2 y las 5 AM, "
             "cuando hay silencio total y puede concentrarse mejor en proyectos de IA.",
        metadata={"categoria": "personal", "habito": "horario_trabajo"},
    )
    print(f"    ✓ Recuerdo personal guardado (id: {id_personal[:8]}...)")

    # 3. Búsqueda semántica
    print("\n[3] Búsqueda semántica — query: '¿En qué proyecto estábamos trabajando?'")
    query = "¿En qué proyecto estábamos trabajando?"
    contexto = mem.retrieve_relevant_context(query, n_results=3)

    if contexto:
        print(f"    Recuerdos más relevantes recuperados ({len(contexto)}):")
        for i, doc in enumerate(contexto, 1):
            preview = doc[:90] + "..." if len(doc) > 90 else doc
            print(f"      {i}. {preview}")
    else:
        print("    (sin resultados)")

    # 4. Estadísticas finales
    print("\n[4] Estadísticas finales:")
    stats = mem.get_memory_stats()
    for k, v in stats.items():
        print(f"    {k}: {v}")

    print("\n" + "=" * 58)
    print("  Test completado exitosamente.")
    print("=" * 58)