JARVIS/nucleo/saber/busca.py

"""Buscar en el saber de COFRE lo que hace falta para contestar una pregunta.

Esto es la mitad "recuperar" de RAG: dada una pregunta, devuelve los fragmentos
de los apuntes que mas se le parecen, para pasarselos al cerebro. El cerebro
construye la respuesta desde ESE texto, no desde lo que recuerde, que es lo que
evita que se invente flags.

## Por que se carga una vez y se guarda

El indice son 2.866 vectores. Cargarlo y encodear la pregunta cuesta, pero solo
la primera vez: el modelo y los vectores se quedan en memoria. Una consulta
despues es un producto escalar contra 2.866 filas, milisegundos.

## Por que no una base de datos vectorial

FAISS, Chroma y compañia son para millones de vectores. Con 2.866, numpy y un
producto escalar es mas rapido de arrancar, no añade una dependencia pesada, y
el indice es un jsonl que se puede abrir y leer con los ojos.
"""
import json
import os

AQUI = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
INDICE = os.path.join(AQUI, "datos", "saber.jsonl")
# El indice portable que viene en el repo. Se usa tal cual mientras no exista el
# personal, para que un clon recien hecho ya sepa de pentesting sin reindexar.
INDICE_PUBLICO = os.path.join(AQUI, "datos", "saber-publico.jsonl")

_modelo = None
_entradas = None
_matriz = None


def _fuente():
    """El indice personal si existe; si no, el publico que trae el repo."""
    if os.path.exists(INDICE):
        return INDICE
    if os.path.exists(INDICE_PUBLICO):
        return INDICE_PUBLICO
    return None


def _carga():
    global _modelo, _entradas, _matriz
    if _entradas is not None:
        return _entradas is not False

    fuente = _fuente()
    if not fuente:
        _entradas = False
        return False

    import numpy as np
    from model2vec import StaticModel

    entradas, vectores = [], []
    with open(fuente, encoding="utf-8") as f:
        for l in f:
            try:
                d = json.loads(l)
            except json.JSONDecodeError:
                continue
            v = d.pop("v", None)
            if v is None:
                continue
            entradas.append(d)
            vectores.append(v)

    _entradas = entradas
    _matriz = np.array(vectores, dtype="float32")
    # normalizar una vez: asi la similitud es un simple producto escalar
    normas = np.linalg.norm(_matriz, axis=1, keepdims=True)
    _matriz = _matriz / np.clip(normas, 1e-9, None)
    _modelo = StaticModel.from_pretrained("minishlab/potion-multilingual-128M")
    return True


def _norm(t: str) -> str:
    import re
    import unicodedata
    t = unicodedata.normalize("NFD", (t or "").lower())
    t = "".join(c for c in t if unicodedata.category(c) != "Mn")
    return re.sub(r"[^a-z0-9 ]", " ", t)


# El texto normalizado de cada entrada y su nombre, cacheados para el bono.
_texto_norm = None
_nombre_norm = None


def _bono_palabras(palabras):
    """Un vector de bonos, uno por entrada, segun cuantas palabras compartan."""
    import numpy as np
    global _texto_norm, _nombre_norm
    if _texto_norm is None:
        _texto_norm = [set(_norm(e["texto"]).split()) for e in _entradas]
        _nombre_norm = [_norm(e.get("herramienta", "")) for e in _entradas]
    bono = np.zeros(len(_entradas), dtype="float32")
    for i, (palabras_e, nombre) in enumerate(zip(_texto_norm, _nombre_norm)):
        comunes = len(palabras & palabras_e)
        b = 0.08 * comunes                       # cada palabra compartida suma
        if nombre and nombre in palabras:        # y nombrar la herramienta, mucho
            b += 0.5
        bono[i] = b
    return bono


def busca(pregunta: str, cuantos: int = 5, perfil: str = None) -> list:
    """Los fragmentos mas parecidos a la pregunta, mejor primero.

    Cada uno lleva su similitud (0-1), de que herramienta y fichero viene, y el
    texto. perfil filtra por carpeta de COFRE si se quiere acotar.
    """
    if not _carga() or not pregunta.strip():
        return []
    import numpy as np

    q = _modelo.encode([pregunta])[0]
    q = q / max(float(np.linalg.norm(q)), 1e-9)
    sim = _matriz @ q

    # Impulso por palabras exactas (busqueda hibrida).
    #
    # Los embeddings estaticos fallan cuando la pregunta va en castellano y el
    # texto en ingles —"cambiar permisos" no se parece a "change file mode
    # bits"—. Pero si la pregunta NOMBRA la herramienta ("como uso chmod") o
    # comparte palabras con el texto, eso es una señal fuerte que la similitud
    # semantica sola no aprovecha. Se suma un bono por cada palabra de la
    # pregunta que aparezca, y uno grande si coincide el nombre de la
    # herramienta: asi "usa nmap para..." lleva nmap al primer puesto.
    palabras = {w for w in _norm(pregunta).split() if len(w) > 3}
    if palabras:
        bono = _bono_palabras(palabras)
        sim = sim + bono

    orden = np.argsort(-sim)[: cuantos * 6]
    salida, vistos = [], set()
    for i in orden:
        e = _entradas[i]
        if perfil and e.get("perfil") != perfil:
            continue
        # sin duplicados: OSCP_APUNTES es copia de los apuntes de primer nivel.
        # Se comparan los primeros 120 caracteres, que basta para reconocerlos.
        firma = e["texto"][:120]
        if firma in vistos:
            continue
        vistos.add(firma)
        salida.append({**e, "sim": round(float(sim[i]), 3)})
        if len(salida) >= cuantos:
            break
    return salida


def contexto(pregunta: str, cuantos: int = 5, minimo: float = 0.35) -> str:
    """Lo mismo, pero ya montado como texto para meterle al cerebro.

    Se corta por debajo de `minimo` de similitud: pasarle al modelo fragmentos
    que no vienen a cuento es lo que le hace mezclar herramientas que no tienen
    nada que ver. Mejor darle poco y bueno que mucho y ruidoso.
    """
    trozos = [t for t in busca(pregunta, cuantos) if t["sim"] >= minimo]
    if not trozos:
        return ""
    lineas = ["De los apuntes de COFRE del usuario:\n"]
    for t in trozos:
        fuente = t.get("fichero") or f"{t['perfil']}/{t['herramienta']}"
        lineas.append(f"--- {fuente} ---\n{t['texto']}\n")
    return "\n".join(lineas)


def resumen(tope_temas: int = 10) -> str:
    """Un resumen de QUE hay indexado, para que el modelo sepa de que sabe.

    Se deriva del indice de VERDAD, no de una lista fija en el prompt. Asi es
    honesto en cualquier maquina: en la del usuario dice sus areas reales; en un
    clon sin apuntes propios dira lo que haya (comandos de Linux) o nada. Lo usa
    el cerebro para el prompt de sistema; evita que presuma de un conocimiento
    que en esa maquina no tiene."""
    if not _carga() or not _entradas:
        return ""
    from collections import Counter
    perfiles, temas = Counter(), Counter()
    n_comandos = 0
    for e in _entradas:
        t = e.get("tipo")
        if t == "comando":
            n_comandos += 1
        elif t in ("apunte", "metodologia", "ficha", "pregunta") and e.get("perfil"):
            perfiles[e["perfil"]] += 1
        if e.get("tema"):
            temas[e["tema"]] += 1
    partes = []
    if perfiles:
        partes.append("apuntes de " + ", ".join(p for p, _ in perfiles.most_common(6)))
    if temas:
        partes.append("metodologia en " + ", ".join(t for t, _ in temas.most_common(tope_temas)))
    if n_comandos:
        partes.append(f"{n_comandos} comandos de Linux")
    if not partes:
        return ""
    return ("En los apuntes que tienes indexados (buscalos con buscar_en_apuntes) "
            "hay: " + "; ".join(partes) + ".")


def disponible() -> bool:
    return _fuente() is not None