semantic_memory.py

from __future__ import annotations

import os
from difflib import SequenceMatcher
from typing import Any

# ZeroGPU must be imported before anything that imports torch.
try:
    import spaces  # type: ignore
    _SPACES_AVAILABLE = True
except ImportError:
    _SPACES_AVAILABLE = False

    class _Spaces:
        @staticmethod
        def GPU(function=None, *args, **kwargs):
            if function is not None and callable(function):
                return function

            def decorator(fn):
                return fn

            return decorator

    spaces = _Spaces()  # type: ignore

try:
    import torch
except Exception:
    torch = None

import numpy as np

from memory_engine import (
    build_scene,
    memory_context_score,
    normalize_text,
    _context_allows_memory,
)

try:
    from sentence_transformers import SentenceTransformer
except ImportError:
    SentenceTransformer = None

try:
    from sklearn.feature_extraction.text import TfidfVectorizer
    from sklearn.metrics.pairwise import cosine_similarity
except ImportError:
    TfidfVectorizer = None
    cosine_similarity = None

EMBEDDING_MODEL_NAME = os.getenv(
    "LEXI_EMBEDDING_MODEL",
    "sentence-transformers/all-MiniLM-L6-v2",
)
USE_ZEROGPU = os.getenv("LEXI_USE_ZEROGPU", "1") != "0"

_MINILM_ERROR = ""

if torch is not None and (_SPACES_AVAILABLE or torch.cuda.is_available()):
    DEVICE = "cuda" if USE_ZEROGPU else "cpu"
else:
    DEVICE = "cpu"

_EMBEDDER = None
if SentenceTransformer is not None:
    try:
        _EMBEDDER = SentenceTransformer(
            EMBEDDING_MODEL_NAME,
            device=DEVICE,
        )
    except Exception as exc:
        _MINILM_ERROR = f"MiniLM initialization failed: {exc}"
        _EMBEDDER = None
else:
    _MINILM_ERROR = "sentence-transformers is not installed."


@spaces.GPU(duration=12)
def _encode_gpu(texts: list[str]) -> np.ndarray:
    if _EMBEDDER is None:
        raise RuntimeError(_MINILM_ERROR or "MiniLM model is unavailable.")

    return np.asarray(
        _EMBEDDER.encode(
            texts,
            normalize_embeddings=True,
            convert_to_numpy=True,
            show_progress_bar=False,
            batch_size=min(32, max(1, len(texts))),
        ),
        dtype=np.float32,
    )


_PROFILE_CACHE: dict[tuple[str, str], np.ndarray] = {}
_CACHE_LIMIT = 512


def _encode(texts: list[str]) -> np.ndarray:
    if not texts:
        return np.empty((0, 384), dtype=np.float32)
    if _EMBEDDER is None:
        raise RuntimeError(_MINILM_ERROR or "MiniLM model is unavailable.")
    return _encode_gpu(texts)


def _candidate_retrieval_pool(
    memories: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    pool = []
    for memory in memories or []:
        status = str(memory.get("status", "candidate") or "candidate")
        if status in {"active", "candidate", "ambiguous"}:
            confidence = float(memory.get("confidence", 0.0) or 0.0)
            positives = int(memory.get("positive_evidence_count", 0) or 0)
            if status == "active" or (confidence >= 0.60 and positives >= 1):
                pool.append(memory)
    return pool


def retrieval_memories(memories: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return _candidate_retrieval_pool(memories)


def _memory_profile(memory: dict[str, Any]) -> str:
    positives = memory.get("positive_contexts", [])[-8:]
    negatives = memory.get("negative_contexts", [])[-8:]
    contexts = " ".join(item.get("context", "") for item in positives)
    negative_contexts = " ".join(item.get("context", "") for item in negatives)
    tags = " ".join(
        tag
        for item in positives
        for tag in item.get("tags", [])
    )
    return " ".join(
        part
        for part in [
            memory.get("source", ""),
            memory.get("target", ""),
            memory.get("kind", ""),
            "preferred contexts",
            contexts,
            "avoid contexts",
            negative_contexts,
            tags,
        ]
        if part
    )[:3500]


def _lexical_score(query: str, memory: dict[str, Any]) -> float:
    query_norm = normalize_text(query)
    best = 0.0
    for term in [memory.get("source", ""), *memory.get("aliases", [])]:
        term_norm = normalize_text(term)
        if not term_norm:
            continue
        if term_norm in query_norm:
            return 1.0
        for token in query_norm.split():
            if len(token) >= 3 and len(term_norm) >= 3:
                best = max(
                    best,
                    SequenceMatcher(
                        None,
                        token,
                        term_norm,
                        autojunk=False,
                    ).ratio(),
                )
    return min(best, 0.94)


def _term_present(query: str, memory: dict[str, Any]) -> bool:
    query_norm = normalize_text(query)
    target_norm = normalize_text(memory.get("target", ""))
    for term in [memory.get("source", ""), *memory.get("aliases", [])]:
        term_norm = normalize_text(term)
        if not term_norm:
            continue
        if term_norm == target_norm and memory.get("kind") != "casing":
            continue
        if term_norm in query_norm:
            return True
    return False


def _context_features(
    query: str,
    memory: dict[str, Any],
    note_title: str,
    scene_text: str,
) -> dict[str, Any]:
    scene = build_scene(note_title, scene_text or "", query)
    ctx = memory_context_score(memory, scene)
    positive = float(ctx["positive"])
    negative = float(ctx["negative"])
    context_fit = max(0.0, positive - 0.65 * negative)
    blocked = (
        not _context_allows_memory(memory, scene, positive, negative)
        or ctx["exception"] >= 0.76
        or (negative >= 0.58 and negative > positive)
    )
    return {
        "scene": scene,
        "positive": positive,
        "negative": negative,
        "exception": float(ctx["exception"]),
        "context_fit": context_fit,
        "blocked": blocked,
    }


def _quick_results(
    query: str,
    memories: list[dict],
    note_title: str,
    scene_text: str,
    top_k: int,
) -> list[dict[str, Any]]:
    results = []
    for memory in memories:
        features = _context_features(query, memory, note_title, scene_text)
        lexical = _lexical_score(query, memory)
        context_fit = features["context_fit"]
        evidence = min(
            1.0,
            float(memory.get("positive_evidence_count", 0) or 0) / 3.0,
        )
        score = (
            0.47 * lexical
            + 0.28 * context_fit
            + 0.15 * float(memory.get("confidence", 0.0) or 0.0)
            + 0.10 * evidence
        )
        if features["blocked"]:
            score *= 0.20

        results.append({
            "memory_id": memory["id"],
            "source": memory["source"],
            "target": memory["target"],
            "kind": memory["kind"],
            "status": memory.get("status", "candidate"),
            "confidence": round(float(memory.get("confidence", 0.0) or 0.0), 3),
            "evidence_count": memory.get("evidence_count", 0),
            "semantic_score": None,
            "lexical_score": round(lexical, 4),
            "context_score": round(context_fit, 4),
            "positive_context_score": round(features["positive"], 4),
            "negative_context_score": round(features["negative"], 4),
            "retrieval_score": round(score, 4),
            "retrieval_backend": "fast lexical + context",
            "blocked_by_context": bool(features["blocked"]),
            "term_present": _term_present(query, memory),
            "reason": "Fast lexical/context retrieval.",
        })

    results.sort(key=lambda x: x["retrieval_score"], reverse=True)
    return results[:top_k]


def quick_retrieve_memories(
    query: str,
    memories: list[dict],
    top_k: int = 4,
    fuzzy_threshold: float = 0.96,
    note_title: str = "",
    scene_text: str = "",
):
    pool = retrieval_memories(memories)
    if not pool or not query.strip():
        return [], True

    results = _quick_results(
        query,
        pool,
        note_title,
        scene_text,
        max(top_k, 8),
    )
    strong = [
        item
        for item in results
        if item["term_present"]
        and item.get("status") == "active"
        and int(item.get("evidence_count", 0) or 0) >= 2
        and not item["blocked_by_context"]
        and item["lexical_score"] >= fuzzy_threshold
        and (
            item["context_score"] >= 0.45
            or item["negative_context_score"] < 0.50
        )
    ]

    unresolved_candidate = any(
        item.get("status") in {"candidate", "ambiguous"}
        and not item.get("blocked_by_context")
        and (
            item.get("term_present")
            or item.get("lexical_score", 0.0) >= 0.45
            or item.get("context_score", 0.0) >= 0.45
        )
        for item in results
    )

    if strong and not unresolved_candidate:
        return strong[:top_k], True

    return results[:top_k], False


def _statistical_batch_scores(
    query: str,
    memories: list[dict[str, Any]],
) -> tuple[np.ndarray, str]:
    if TfidfVectorizer is None or cosine_similarity is None:
        return np.zeros(len(memories), dtype=np.float32), "ML unavailable"

    texts = [_memory_profile(memory) for memory in memories]
    corpus = [query, *texts]

    word_vectorizer = TfidfVectorizer(
        lowercase=True,
        ngram_range=(1, 2),
        min_df=1,
        sublinear_tf=True,
    )
    char_vectorizer = TfidfVectorizer(
        lowercase=True,
        analyzer="char_wb",
        ngram_range=(2, 5),
        min_df=1,
        sublinear_tf=True,
    )

    word_matrix = word_vectorizer.fit_transform(corpus)
    char_matrix = char_vectorizer.fit_transform(corpus)

    word_scores = cosine_similarity(word_matrix[0:1], word_matrix[1:]).ravel()
    char_scores = cosine_similarity(char_matrix[0:1], char_matrix[1:]).ravel()

    scores = 0.40 * word_scores + 0.60 * char_scores
    return np.asarray(scores, dtype=np.float32), "ML fallback: word+char TF-IDF"


def _semantic_batch_scores(
    query: str,
    memories: list[dict[str, Any]],
) -> tuple[np.ndarray, str]:
    profiles = [_memory_profile(memory) for memory in memories]
    missing = []
    missing_keys = []
    vectors = [None] * len(memories)

    for index, (memory, profile) in enumerate(zip(memories, profiles)):
        key = (
            memory["id"],
            f'{memory.get("updated_at", "")}::{profile}',
        )
        cached = _PROFILE_CACHE.get(key)
        if cached is None:
            missing.append(profile)
            missing_keys.append((index, key))
        else:
            vectors[index] = cached

    encoded = _encode([query, *missing])
    query_vector = encoded[0]

    for offset, (index, key) in enumerate(missing_keys, start=1):
        vectors[index] = encoded[offset]
        _PROFILE_CACHE[key] = encoded[offset]

    while len(_PROFILE_CACHE) > _CACHE_LIMIT:
        _PROFILE_CACHE.pop(next(iter(_PROFILE_CACHE)))

    scores = np.asarray(
        [float(np.dot(vector, query_vector)) for vector in vectors],
        dtype=np.float32,
    )
    return scores, "MiniLM via ZeroGPU" if DEVICE == "cuda" else "MiniLM on CPU"


def retrieve_memories(
    query: str,
    memories: list[dict],
    note_title: str = "",
    scene_text: str = "",
    top_k: int = 4,
) -> list[dict]:
    pool = retrieval_memories(memories)
    if not pool or not query.strip():
        return []

    quick, strong = quick_retrieve_memories(
        query,
        pool,
        top_k=max(top_k, 8),
        note_title=note_title,
        scene_text=scene_text,
    )

    if strong:
        return quick[:top_k]

    scene = build_scene(note_title, scene_text or "", query)

    ml_error = ""
    try:
        semantic_scores, backend = _semantic_batch_scores(
            scene["text"][:3000],
            pool,
        )
    except Exception as exc:
        ml_error = str(exc)
        semantic_scores, backend = _statistical_batch_scores(
            scene["text"][:3000],
            pool,
        )
        if backend == "ML unavailable":
            return quick[:top_k]

    results = []
    for index, memory in enumerate(pool):
        features = _context_features(query, memory, note_title, scene_text)
        semantic = max(0.0, min(1.0, float(semantic_scores[index])))
        lexical = _lexical_score(query, memory)
        context_fit = features["context_fit"]
        evidence = min(
            1.0,
            float(memory.get("positive_evidence_count", 0) or 0) / 3.0,
        )

        score = (
            0.42 * semantic
            + 0.23 * lexical
            + 0.25 * context_fit
            + 0.10 * float(memory.get("confidence", 0.0) or 0.0)
        )
        if features["blocked"]:
            score *= 0.20

        results.append({
            "memory_id": memory["id"],
            "source": memory["source"],
            "target": memory["target"],
            "kind": memory["kind"],
            "status": memory.get("status", "candidate"),
            "confidence": round(float(memory.get("confidence", 0.0) or 0.0), 3),
            "evidence_count": memory.get("evidence_count", 0),
            "semantic_score": round(semantic, 4),
            "lexical_score": round(lexical, 4),
            "context_score": round(context_fit, 4),
            "positive_context_score": round(features["positive"], 4),
            "negative_context_score": round(features["negative"], 4),
            "retrieval_score": round(score, 4),
            "retrieval_backend": backend,
            "blocked_by_context": bool(features["blocked"]),
            "term_present": _term_present(query, memory),
            "reason": (
                "MiniLM semantic retrieval validated the scene."
                if backend.startswith("MiniLM")
                else "Statistical ML fallback handled semantic/phonetic retrieval."
            ),
            "ml_error": ml_error if ml_error else None,
        })

    results.sort(key=lambda x: x["retrieval_score"], reverse=True)
    return results[:top_k]



def ml_probe() -> dict[str, Any]:
    started = __import__("time").perf_counter()
    result = {
        "ok": False,
        "backend": backend_status(),
        "model": EMBEDDING_MODEL_NAME,
        "device": DEVICE,
        "minilm_loaded": _EMBEDDER is not None,
        "error": None,
    }
    try:
        if _EMBEDDER is None:
            result["backend"] = "ML fallback: word+char TF-IDF"
            vector_scores, backend = _statistical_batch_scores(
                "Lexi ML health check",
                [normalize_memory_for_probe()],
            )
            result["backend"] = backend
            result["ok"] = len(vector_scores) == 1
        else:
            vectors = _encode(["Lexi ML health check"])
            result["backend"] = "MiniLM via ZeroGPU" if DEVICE == "cuda" else "MiniLM on CPU"
            result["ok"] = bool(vectors.shape[0] == 1)
    except Exception as exc:
        result["error"] = str(exc)
    result["latency_ms"] = round((__import__("time").perf_counter() - started) * 1000, 1)
    return result


def normalize_memory_for_probe() -> dict[str, Any]:
    return {
        "id": "ml_probe",
        "source": "lexi",
        "target": "Lexi",
        "kind": "product_service",
        "confidence": 0.9,
        "positive_evidence_count": 1,
        "evidence_count": 1,
        "positive_contexts": [{"context": "Lexi personal writing assistant", "tags": ["software"]}],
        "negative_contexts": [],
        "exceptions": [],
    }

def backend_status() -> str:
    if _EMBEDDER is not None:
        if DEVICE == "cuda":
            return "MiniLM semantic retrieval via ZeroGPU (loaded)"
        return "MiniLM semantic retrieval on CPU (loaded)"
    if _MINILM_ERROR:
        return f"ML fallback enabled; {_MINILM_ERROR}"
    return "ML fallback enabled"


def ml_diagnostics() -> dict[str, Any]:
    return {
        "minilm_loaded": _EMBEDDER is not None,
        "model": EMBEDDING_MODEL_NAME,
        "device": DEVICE,
        "zero_gpu_requested": USE_ZEROGPU,
        "spaces_available": _SPACES_AVAILABLE,
        "minilm_error": _MINILM_ERROR,
        "fallback_available": TfidfVectorizer is not None,
        "backend": backend_status(),
    }