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()