knowledge.py

"""Knowledge base module for ORTOS Telegram Bot.
Hybrid search: local bge-m3 + BM25 + FAISS (RRF fusion).
Fallback: HF Inference API ? TF-IDF.
"""

import json, os, requests, logging
import numpy as np
from pydantic import BaseModel
from rank_bm25 import BM25Okapi
import faiss

logger = logging.getLogger(__name__)

class KnowledgeItem(BaseModel):
    id: str
    title: str
    content: str

_index = None
_sections = []
_bm25 = None
_bm25_corpus = []
_local_model = None
_tfidf_backup = None
_items_backup = None
_hf_token = os.getenv("HF_TOKEN")


_morph = None
_morph_loaded = False

def _get_morph():
    global _morph, _morph_loaded
    if _morph_loaded:
        return _morph
    _morph_loaded = True
    try:
        from pymorphy3 import MorphAnalyzer
        _morph = MorphAnalyzer()
        logger.info("pymorphy3 loaded for BM25")
    except Exception as e:
        logger.warning(f"pymorphy3 unavailable: {e}")
        _morph = None
    return _morph


def _stem_word(word: str) -> str:
    if not word.isalpha():
        return word
    m = _get_morph()
    if m is not None:
        try:
            return m.parse(word)[0].normal_form
        except Exception:
            pass
    return word


def _tokenize(text: str) -> list[str]:
    if _get_morph() is None:
        return text.lower().split()
    return [_stem_word(w) for w in text.lower().split()]


def _build_bm25(sections: list[KnowledgeItem]):
    global _bm25, _bm25_corpus
    _bm25_corpus = [_tokenize(s.title + " " + s.content) for s in sections]
    _bm25 = BM25Okapi(_bm25_corpus)


def _load_faiss():
    global _index
    index_path = "embeddings/index.faiss"
    if os.path.exists(index_path):
        _index = faiss.read_index(index_path)
        return True
    return False


def _load_sections(paths):
    global _sections
    if isinstance(paths, str):
        paths = [paths]
    _sections = []
    for path in paths:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
        for name, sec in data["sections"].items():
            _sections.append(KnowledgeItem(
                id=f"{path}:{name}",
                title=sec["title"],
                content=sec["content"],
            ))
    return _sections


def _init_tfidf_backup(paths):
    global _tfidf_backup, _items_backup
    try:
        from knowledge_tfidf_backup import reload_knowledge as tfidf_reload
        _items_backup, _tfidf_backup = tfidf_reload(paths)
    except Exception:
        pass


def _init_local_model():
    global _local_model
    try:
        from sentence_transformers import SentenceTransformer
        logger.info("Loading bge-m3 locally...")
        _local_model = SentenceTransformer("BAAI/bge-m3", trust_remote_code=True)
        logger.info("bge-m3 loaded locally")
        return True
    except Exception as e:
        logger.warning(f"Local model unavailable: {e}")
        return False


def _embed_local(query: str) -> np.ndarray | None:
    if _local_model is None:
        return None
    try:
        emb = _local_model.encode([query], normalize_embeddings=True)
        return emb.astype(np.float32)
    except Exception:
        return None


def _embed_via_hfapi(query: str) -> np.ndarray | None:
    headers = {"Authorization": f"Bearer {_hf_token}"} if _hf_token else {}
    try:
        resp = requests.post(
            "https://api-inference.huggingface.co/models/BAAI/bge-m3",
            headers=headers,
            json={"inputs": query},
            timeout=15,
        )
        if resp.status_code == 200:
            emb = np.array(resp.json(), dtype=np.float32)
            emb = emb / np.linalg.norm(emb)
            return emb.reshape(1, -1).astype(np.float32)
    except Exception:
        pass
    return None


def _embed_query(query: str) -> np.ndarray | None:
    emb = _embed_local(query)
    if emb is not None:
        return emb
    emb = _embed_via_hfapi(query)
    if emb is not None:
        return emb
    return None


def load_knowledge_base(paths: list[str] | str) -> list[KnowledgeItem]:
    return _load_sections(paths)


def search(query: str, top_k: int = 2, items: list[KnowledgeItem] = None,
           tfidf: dict = None) -> list[KnowledgeItem]:
    return search_debug(query, top_k)["items"]


def search_debug(query: str, top_k: int = 2) -> dict:
    """Returns {items, method, details: [{title, bm25_score, embed_rank, rrf_score}]}"""
    n = len(_sections)
    if n == 0:
        return {"items": [], "method": "none", "details": []}

    bm25_scores = None
    if _bm25 is not None:
        bm25_scores = _bm25.get_scores(_tokenize(query))

    embed_ranks = None
    qvec = _embed_query(query)
    if qvec is not None and _index is not None and _index.ntotal > 0:
        embed_scores, embed_indices = _index.search(qvec, min(n, _index.ntotal))
        embed_ranks = {int(idx): rank for rank, idx in enumerate(embed_indices[0])}

    # --- RRF fusion ---
    if bm25_scores is not None and embed_ranks is not None:
        rrf = {}
        for i in range(n):
                bm25_rank = sorted(range(n), key=lambda j: -bm25_scores[j]).index(i)
                rrf_score = 0.0
                rrf_score += 0.3 * (1 / (bm25_rank + 1))
                if i in embed_ranks:
                    rrf_score += 0.7 * (1 / (embed_ranks[i] + 1))
                rrf[i] = rrf_score
        top_indices = sorted(rrf.keys(), key=lambda i: -rrf[i])[:top_k]
        details = [{
            "id": _sections[i].id,
            "title": _sections[i].title,
            "bm25_rank": sorted(range(n), key=lambda j: -bm25_scores[j]).index(i),
            "embed_rank": embed_ranks.get(i, None),
            "rrf_score": round(rrf[i], 4),
        } for i in top_indices]
        return {"items": [_sections[i] for i in top_indices], "method": "hybrid (bge-m3+BM25)", "details": details}

    # --- Embedding only ---
    if embed_ranks is not None:
        top_indices = sorted(embed_ranks.keys(), key=lambda i: embed_ranks[i])[:top_k]
        details = [{"id": _sections[i].id, "title": _sections[i].title, "embed_rank": embed_ranks[i], "rrf_score": None} for i in top_indices]
        return {"items": [_sections[i] for i in top_indices], "method": "bge-m3 only", "details": details}

    # --- BM25 only ---
    if bm25_scores is not None:
        top_indices = sorted(range(n), key=lambda i: -bm25_scores[i])[:top_k]
        details = [{"id": _sections[i].id, "title": _sections[i].title, "bm25_rank": i, "rrf_score": None} for i in top_indices]
        return {"items": [_sections[i] for i in top_indices], "method": "BM25 only", "details": details}

    # --- TF-IDF fallback ---
    if _tfidf_backup is not None and _items_backup:
        from knowledge_tfidf_backup import search as tfidf_search
        items = tfidf_search(query, top_k=top_k, items=_items_backup, tfidf=_tfidf_backup)
        details = [{"id": it.id, "title": it.title, "rrf_score": None} for it in items]
        return {"items": items, "method": "TF-IDF fallback", "details": details}

    return {"items": [], "method": "none", "details": []}


def reload_knowledge(paths: list[str] | str) -> tuple[list[KnowledgeItem], dict]:
    if isinstance(paths, str):
        paths = [paths]
    _load_faiss()
    items = _load_sections(paths)
    _build_bm25(items)
    _init_tfidf_backup(paths)
    _init_local_model()
    return items, {}