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)