src/app_kit/embedding.py

import json
import math
import numpy as np
from dataclasses import dataclass, field
import os

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

def _get_model():
    global _model
    if _model is None and SentenceTransformer is not None:
        _model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
    return _model

def vectorize(text: str) -> list[float]:
    model = _get_model()
    if model:
        import logging
        logging.getLogger('embedder').info(f"Running inference on sentence-transformers/all-MiniLM-L6-v2 for text length {len(text)}")
        return model.encode([text])[0].tolist()
    return []

def cosine_similarity(left: list[float], right: list[float]) -> float:
    if not left or not right:
        return 0.0
    dot = sum(l * r for l, r in zip(left, right))
    left_norm = math.sqrt(sum(v * v for v in left))
    right_norm = math.sqrt(sum(v * v for v in right))
    if not left_norm or not right_norm:
        return 0.0
    return dot / (left_norm * right_norm)

@dataclass
class SimpleEmbeddingIndex:
    entries: dict[str, list[float]] = field(default_factory=dict)

    def add(self, record_id: str, text: str) -> None:
        self.entries[record_id] = vectorize(text)

    def search(self, query: str, limit: int = 5) -> list[tuple[str, float]]:
        qvec = vectorize(query)
        scored = [(record_id, cosine_similarity(qvec, vec)) for record_id, vec in self.entries.items()]
        return sorted(scored, key=lambda item: item[1], reverse=True)[:limit]

def extract_keywords(text: str, limit: int = 6) -> list[str]:
    # Keeping extract_keywords simple as it's not a model response
    import re
    from collections import Counter
    tokens = [tok.lower() for tok in re.findall(r"[A-Za-z0-9']+", text or '') if len(tok) > 2]
    return [word for word, _ in Counter(tokens).most_common(limit)]