Upload knowledge_base.py with huggingface_hub

knowledge_base.py

@@ -0,0 +1,183 @@
+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
+        log.info("Seeding grammatica italiana...")
+        try:
+            from seed_italian import GRAMMAR_SEED
+            for item in GRAMMAR_SEED:
+                doc_id = hashlib.md5(item["q"].encode()).hexdigest()
+                self._col.upsert(
+                    documents=[item["q"]],
+                    ids=[doc_id],
+                    metadatas=[{
+                        "query": item["q"],
+                        "answer": item["a"],
+                        "source": "grammatica_italiana",
+                        "feedback_score": 5,
+                    }],
+                )
+            open(SEEDED_FLAG, "w").close()
+            log.info("Grammatica italiana caricata: %d regole.", len(GRAMMAR_SEED))
+        except Exception as e:
+            raise GenerAIError(ErrorCode.KB_SEED_FAILED, f"Seeding fallito: {fmt_exc(e)}", cause=e)
+
+    # ── 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()