File size: 7,439 Bytes
33d4a41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
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()