knowledge_base.py

import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
import hashlib
import json
import os
import shutil
from errors import get_logger, GenerAIError, ErrorCode, fmt_exc

log = get_logger("knowledge_base")

DB_PATH = "./database"
COUNTS_FILE = f"{DB_PATH}/query_counts.json"
SEEDED_FLAG = f"{DB_PATH}/.seeded_v2"

EMBED_MODEL = "paraphrase-multilingual-MiniLM-L12-v2"


def _load_counts() -> dict:
    if os.path.exists(COUNTS_FILE):
        try:
            with open(COUNTS_FILE) as f:
                data = json.load(f)
            log.debug("Caricati %d contatori da %s", len(data), COUNTS_FILE)
            return data
        except Exception as e:
            log.warning("[%s] Impossibile leggere %s — %s. Uso contatori vuoti.", ErrorCode.KB_INIT_FAILED.value, COUNTS_FILE, fmt_exc(e))
            return {}
    return {}


def _save_counts(counts: dict):
    os.makedirs(DB_PATH, exist_ok=True)
    try:
        with open(COUNTS_FILE, "w") as f:
            json.dump(counts, f)
    except Exception as e:
        log.error("[%s] Impossibile salvare contatori: %s", ErrorCode.KB_WRITE_FAILED.value, fmt_exc(e))


def _normalize(text: str) -> str:
    return " ".join(text.lower().strip().split()[:12])


class KnowledgeBase:
    ARCHIVE_THRESHOLD = 3

    def __init__(self):
        log.info("Inizializzazione KnowledgeBase in: %s", os.path.abspath(DB_PATH))
        try:
            os.makedirs(DB_PATH, exist_ok=True)
        except Exception as e:
            raise GenerAIError(ErrorCode.KB_INIT_FAILED, f"Impossibile creare la directory DB: {fmt_exc(e)}", cause=e)

        # Wipe old schema if needed
        if not os.path.exists(SEEDED_FLAG):
            old_flag = f"{DB_PATH}/.seeded"
            if os.path.exists(old_flag):
                os.remove(old_flag)
                log.info("Rimosso flag vecchio schema.")
            chroma_dir = os.path.join(DB_PATH, "chroma.sqlite3")
            if os.path.exists(chroma_dir):
                os.remove(chroma_dir)
                log.info("Rimosso database ChromaDB obsoleto.")

        try:
            log.debug("Caricamento modello embedding: %s", EMBED_MODEL)
            ef = SentenceTransformerEmbeddingFunction(model_name=EMBED_MODEL)
            self._client = chromadb.PersistentClient(path=DB_PATH)
            self._col = self._client.get_or_create_collection(
                name="generai",
                embedding_function=ef,
                metadata={"hnsw:space": "cosine"},
            )
            log.info("ChromaDB pronto. Documenti in memoria: %d", self._col.count())
        except Exception as e:
            raise GenerAIError(ErrorCode.KB_INIT_FAILED, f"ChromaDB non avviato: {fmt_exc(e)}", cause=e)

        self._counts = _load_counts()
        self._seed_if_needed()

    def _seed_if_needed(self):
        if os.path.exists(SEEDED_FLAG):
            return
        open(SEEDED_FLAG, "w").close()
        log.info("KB vuota — seeding disabilitato, si parte da zero.")

    # ── Public API ─────────────────────────────────────────────────────────────

    def search(self, query: str, n_results: int = 3) -> list[dict]:
        count = self._col.count()
        if count == 0:
            log.debug("KB vuota, nessuna ricerca eseguita.")
            return []
        try:
            res = self._col.query(
                query_texts=[query],
                n_results=min(n_results, count),
                include=["metadatas", "distances"],
            )
            metas = res["metadatas"][0]
            dists = res["distances"][0]
            log.debug("KB search per %r → %d risultati (best dist=%.3f)", query, len(metas), dists[0] if dists else -1)
            return [
                {"answer": metas[i].get("answer", ""), "metadata": metas[i], "distance": dists[i]}
                for i in range(len(metas))
            ]
        except Exception as e:
            err = GenerAIError(ErrorCode.KB_SEARCH_FAILED, f"Ricerca KB fallita: {fmt_exc(e)}", cause=e)
            err.log(log)
            return []

    def increment_and_should_archive(self, query: str) -> bool:
        key = _normalize(query)
        self._counts[key] = self._counts.get(key, 0) + 1
        _save_counts(self._counts)
        count = self._counts[key]
        log.debug("Query %r vista %d volte (threshold=%d)", key, count, self.ARCHIVE_THRESHOLD)
        return count >= self.ARCHIVE_THRESHOLD

    def get_query_count(self, query: str) -> int:
        return self._counts.get(_normalize(query), 0)

    def add(self, query: str, content: str, source_url: str = "") -> str:
        doc_id = hashlib.md5(_normalize(query).encode()).hexdigest()
        try:
            self._col.upsert(
                documents=[query],
                ids=[doc_id],
                metadatas=[{
                    "query": query[:300],
                    "answer": content[:4000],
                    "source": source_url[:500],
                    "feedback_score": 0,
                }],
            )
            log.info("Documento archiviato (id=%s) per query: %r", doc_id[:8], query[:60])
        except Exception as e:
            err = GenerAIError(ErrorCode.KB_WRITE_FAILED, f"Impossibile salvare in KB: {fmt_exc(e)}", cause=e)
            err.log(log)
        return doc_id

    def reinforce(self, doc_id: str, positive: bool):
        action = "rinforzo positivo" if positive else "eliminazione"
        log.info("Feedback (%s) per doc_id=%s", action, doc_id[:8])
        try:
            result = self._col.get(ids=[doc_id], include=["metadatas"])
            if not result["ids"]:
                log.warning("[%s] doc_id=%s non trovato in KB.", ErrorCode.KB_REINFORCE_FAILED.value, doc_id[:8])
                return
            if positive:
                meta = result["metadatas"][0]
                meta["feedback_score"] = meta.get("feedback_score", 0) + 1
                self._col.update(ids=[doc_id], metadatas=[meta])
                log.debug("feedback_score aggiornato a %d", meta["feedback_score"])
            else:
                self._col.delete(ids=[doc_id])
                key = _normalize(result["metadatas"][0].get("query", ""))
                self._counts.pop(key, None)
                _save_counts(self._counts)
                log.info("Documento eliminato dalla KB.")
        except Exception as e:
            err = GenerAIError(ErrorCode.KB_REINFORCE_FAILED, f"Feedback fallito: {fmt_exc(e)}", cause=e)
            err.log(log)

    def count(self) -> int:
        return self._col.count()