b5f82fb48c
Some checks are pending
CI / ruff (push) Waiting to run
CI / mypy (push) Waiting to run
CI / pytest (push) Waiting to run
CI / package (push) Waiting to run
CI / Install smoke (${{ matrix.label }}) (linux, ubuntu-latest) (push) Blocked by required conditions
CI / Install smoke (${{ matrix.label }}) (macos-arm64, macos-14) (push) Blocked by required conditions
1150 lines
45 KiB
Python
1150 lines
45 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_GLOBAL_RAG: Optional["RAGRetriever"] = None
|
|
|
|
|
|
def set_global_rag(retriever: Optional["RAGRetriever"]) -> None:
|
|
"""Imposta l'istanza RAG globale condivisa dalla pipeline."""
|
|
global _GLOBAL_RAG
|
|
_GLOBAL_RAG = retriever
|
|
|
|
|
|
def get_global_rag() -> Optional["RAGRetriever"]:
|
|
"""Restituisce l'istanza RAG globale (None se non abilitata)."""
|
|
return _GLOBAL_RAG
|
|
|
|
|
|
@dataclass
|
|
class Chunk:
|
|
"""Un singolo chunk indicizzato."""
|
|
|
|
text: str
|
|
source: str
|
|
chunk_index: int
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class SearchResult:
|
|
"""Risultato di una query RAG."""
|
|
|
|
chunk: Chunk
|
|
score: float
|
|
|
|
|
|
def _split_text_recursive(
|
|
text: str,
|
|
chunk_size: int = 512,
|
|
chunk_overlap: int = 64,
|
|
) -> list[str]:
|
|
"""Semplice splitter ricorsivo: paragrafi → frasi → caratteri.
|
|
|
|
Non dipende da langchain per mantenere zero dipendenze extra oltre a
|
|
sentence-transformers.
|
|
"""
|
|
if chunk_size <= 0:
|
|
raise ValueError("chunk_size must be > 0")
|
|
if chunk_overlap >= chunk_size:
|
|
raise ValueError("chunk_overlap must be < chunk_size")
|
|
if len(text) <= chunk_size:
|
|
return [text.strip()] if text.strip() else []
|
|
|
|
separators = [
|
|
("\n\n", True),
|
|
("\n", True),
|
|
(". ", True),
|
|
("? ", True),
|
|
("! ", True),
|
|
("; ", True),
|
|
(", ", True),
|
|
(" ", False),
|
|
("", False),
|
|
]
|
|
|
|
def _split_with(sep: str, keep_sep: bool) -> list[str]:
|
|
if not sep:
|
|
return list(text)
|
|
if keep_sep:
|
|
parts = re.split(rf"(?<={re.escape(sep)})", text)
|
|
else:
|
|
parts = text.split(sep)
|
|
return [p for p in parts if p]
|
|
|
|
chosen_parts: Optional[list[str]] = None
|
|
for sep, keep in separators:
|
|
parts = _split_with(sep, keep)
|
|
if len(parts) > 1:
|
|
chosen_parts = parts
|
|
break
|
|
if chosen_parts is None:
|
|
chosen_parts = list(text)
|
|
|
|
chunks: list[str] = []
|
|
current = ""
|
|
for part in chosen_parts:
|
|
candidate = (current + part).lstrip()
|
|
if len(candidate) <= chunk_size:
|
|
current = candidate
|
|
else:
|
|
if current:
|
|
chunks.append(current.rstrip())
|
|
# fallback: se il singolo pezzo è più grande di chunk_size, tronchiamo
|
|
if len(part) > chunk_size:
|
|
for i in range(0, len(part), chunk_size - chunk_overlap):
|
|
window = part[i : i + chunk_size].strip()
|
|
if window:
|
|
chunks.append(window)
|
|
current = ""
|
|
else:
|
|
# inizi un nuovo chunk con overlap
|
|
tail = ""
|
|
if chunks and chunk_overlap > 0:
|
|
tail = chunks[-1][-chunk_overlap:]
|
|
current = (tail + part).lstrip()
|
|
if current.strip():
|
|
chunks.append(current.rstrip())
|
|
|
|
# merges finale per non avere chunk troppo piccoli (ultimo < 30% chunk_size)
|
|
merged: list[str] = []
|
|
min_len = max(20, chunk_size // 3)
|
|
for c in chunks:
|
|
if merged and len(c) < min_len:
|
|
merged[-1] = (merged[-1] + " " + c).strip()
|
|
else:
|
|
merged.append(c)
|
|
return merged
|
|
|
|
|
|
def _load_text_files(kb_path: Path) -> list[tuple[str, str]]:
|
|
"""Carica tutti i file .md e .txt dalla cartella (ricorsivo).
|
|
|
|
Restituisce lista di (source_rel_path, full_text).
|
|
"""
|
|
docs: list[tuple[str, str]] = []
|
|
for ext in ("*.md", "*.txt"):
|
|
for fp in sorted(kb_path.rglob(ext)):
|
|
if fp.is_file():
|
|
try:
|
|
text = fp.read_text(encoding="utf-8")
|
|
except UnicodeDecodeError:
|
|
text = fp.read_text(encoding="latin-1", errors="replace")
|
|
rel = str(fp.relative_to(kb_path))
|
|
docs.append((rel, text))
|
|
return docs
|
|
|
|
|
|
def _load_jsonl(kb_path: Path) -> list[tuple[str, str, dict[str, Any]]]:
|
|
"""Carica file *.jsonl (opzionale) già pre-chunkizzati.
|
|
|
|
Riga attesa: {"text": "...", "source": "...", "chunk_index": 0 (opzionale), ...altri metadata}
|
|
"""
|
|
prechunked: list[tuple[str, str, dict[str, Any]]] = []
|
|
for fp in sorted(kb_path.rglob("*.jsonl")):
|
|
if not fp.is_file():
|
|
continue
|
|
rel = str(fp.relative_to(kb_path))
|
|
with fp.open("r", encoding="utf-8") as f:
|
|
for lineno, line in enumerate(f, 1):
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
obj = json.loads(line)
|
|
except json.JSONDecodeError as e:
|
|
logger.warning("RAG: riga %s non valida in %s: %s", lineno, rel, e)
|
|
continue
|
|
text = obj.get("text") or obj.get("content") or ""
|
|
if not text:
|
|
continue
|
|
source = obj.get("source") or f"{rel}#riga{lineno}"
|
|
meta = {k: v for k, v in obj.items() if k not in {"text", "content", "source", "chunk_index"}}
|
|
ci = obj.get("chunk_index", lineno - 1)
|
|
prechunked.append((source, text, {**meta, "chunk_index": ci, "from_jsonl": True}))
|
|
return prechunked
|
|
|
|
|
|
class RAGRetriever:
|
|
"""Retriever RAG locale basato su sentence-transformers + numpy.
|
|
|
|
- Indicizza una cartella di file .md/.txt o un JSONL pre-chunkizzato.
|
|
- Salva/carica l'indice NPZ per evitare re-embedding ad ogni avvio.
|
|
- Ricerca coseno-similarità top-k con soglia minima.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
kb_path: str | Path,
|
|
*,
|
|
embedding_model: str = "paraphrase-multilingual-MiniLM-L12-v2",
|
|
device: str = "auto",
|
|
top_k: int = 3,
|
|
threshold: float = 0.25,
|
|
chunk_size: int = 512,
|
|
chunk_overlap: int = 64,
|
|
embedding_batch_size: int = 32,
|
|
language: str = "es",
|
|
inject_as: str = "system",
|
|
) -> None:
|
|
try:
|
|
from sentence_transformers import SentenceTransformer
|
|
except ImportError as e:
|
|
raise RuntimeError(
|
|
"sentence-transformers non installato. Installa le dipendenze RAG con:\n"
|
|
" uv pip install -e .[rag]\n"
|
|
f"Errore originale: {e}"
|
|
) from e
|
|
|
|
self.kb_path = Path(kb_path).expanduser().resolve()
|
|
self.embedding_model_name = embedding_model
|
|
self.device = device
|
|
self.top_k = max(1, int(top_k))
|
|
self.threshold = float(threshold)
|
|
self.chunk_size = max(64, int(chunk_size))
|
|
self.chunk_overlap = max(0, int(chunk_overlap))
|
|
self.embedding_batch_size = max(1, int(embedding_batch_size))
|
|
self.language = language if language in {"es", "it", "en"} else "es"
|
|
self.inject_as = inject_as if inject_as in {"system", "user"} else "system"
|
|
|
|
self.index_path = self.kb_path / "_index.npz"
|
|
self.meta_path = self.kb_path / "_chunks.jsonl"
|
|
self.dynamic_jsonl_path = self.kb_path / "_dynamic.jsonl"
|
|
|
|
self.chunks: list[Chunk] = []
|
|
self.embeddings: Optional[np.ndarray] = None
|
|
self._lock = threading.RLock()
|
|
# File dinamici sono chunk aggiunti via API durante runtime, senza riavvio (vengono
|
|
# persistere sul disk in kb/_dynamic.jsonl tra un riavvio).
|
|
self._dynamic_sources: set[str] = set()
|
|
self._persist_dirty = False
|
|
|
|
import torch
|
|
|
|
resolved = device
|
|
if resolved == "auto":
|
|
if torch.cuda.is_available():
|
|
resolved = "cuda"
|
|
else:
|
|
resolved = "cpu"
|
|
logger.info(
|
|
"RAG: Inizializzazione modello embedding=%s device=%s su kb_path=%s",
|
|
self.embedding_model_name,
|
|
resolved,
|
|
self.kb_path,
|
|
)
|
|
self._model = SentenceTransformer(self.embedding_model_name, device=resolved)
|
|
|
|
# ── Indexing / persistenza ───────────────────────────────────────────────
|
|
|
|
def _index_exists(self) -> bool:
|
|
return self.index_path.is_file() and self.meta_path.is_file()
|
|
|
|
def _save_index(self) -> None:
|
|
assert self.embeddings is not None
|
|
self.kb_path.mkdir(parents=True, exist_ok=True)
|
|
np.savez(self.index_path, embeddings=self.embeddings.astype(np.float32))
|
|
with self.meta_path.open("w", encoding="utf-8") as f:
|
|
for ch in self.chunks:
|
|
f.write(
|
|
json.dumps(
|
|
{
|
|
"text": ch.text,
|
|
"source": ch.source,
|
|
"chunk_index": ch.chunk_index,
|
|
"metadata": ch.metadata,
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
+ "\n"
|
|
)
|
|
logger.info(
|
|
"RAG: Indice salvato: %d chunk, shape embeddings=%s",
|
|
len(self.chunks),
|
|
self.embeddings.shape,
|
|
)
|
|
|
|
def _load_index(self) -> bool:
|
|
if not self._index_exists():
|
|
return False
|
|
try:
|
|
data = np.load(self.index_path)
|
|
embeddings = data["embeddings"]
|
|
chunks: list[Chunk] = []
|
|
with self.meta_path.open("r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
obj = json.loads(line)
|
|
chunks.append(
|
|
Chunk(
|
|
text=obj["text"],
|
|
source=obj["source"],
|
|
chunk_index=int(obj.get("chunk_index", 0)),
|
|
metadata=obj.get("metadata") or {},
|
|
)
|
|
)
|
|
if len(chunks) != embeddings.shape[0]:
|
|
logger.warning(
|
|
"RAG: mismatch chunks(%d) vs embeddings(%s); ricostruisco l'indice.",
|
|
len(chunks),
|
|
embeddings.shape,
|
|
)
|
|
return False
|
|
self.chunks = chunks
|
|
self.embeddings = embeddings.astype(np.float32)
|
|
logger.info(
|
|
"RAG: Indice caricato da disco: %d chunk (shape=%s).",
|
|
len(self.chunks),
|
|
self.embeddings.shape,
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.warning("RAG: caricamento indice fallito (%s); ricostruisco.", e)
|
|
return False
|
|
|
|
def build_index(self, *, force_rebuild: bool = False) -> None:
|
|
"""Popola chunks+embeddings. Carica da NPZ se possibile, altrimenti indicizza."""
|
|
if not force_rebuild and self._load_index():
|
|
return
|
|
|
|
if not self.kb_path.is_dir():
|
|
logger.info("RAG: kb_path=%s non esiste, lo creo vuoto.", self.kb_path)
|
|
self.kb_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
raw_docs = _load_text_files(self.kb_path)
|
|
prechunked = _load_jsonl(self.kb_path)
|
|
|
|
self.chunks: list[Chunk] = []
|
|
texts_to_embed: list[str] = []
|
|
|
|
for source, text in raw_docs:
|
|
pieces = _split_text_recursive(
|
|
text,
|
|
chunk_size=self.chunk_size,
|
|
chunk_overlap=self.chunk_overlap,
|
|
)
|
|
for i, piece in enumerate(pieces):
|
|
ch = Chunk(text=piece, source=source, chunk_index=i, metadata={"splitted": True})
|
|
self.chunks.append(ch)
|
|
texts_to_embed.append(piece)
|
|
|
|
for source, text, meta in prechunked:
|
|
ci = int(meta.pop("chunk_index", 0))
|
|
ch = Chunk(text=text, source=source, chunk_index=ci, metadata=meta)
|
|
self.chunks.append(ch)
|
|
texts_to_embed.append(text)
|
|
|
|
if not self.chunks:
|
|
logger.warning(
|
|
"RAG: Nessun documento trovato in %s. Indice vuoto (retrieval non inietterà niente).",
|
|
self.kb_path,
|
|
)
|
|
self.embeddings = np.zeros((0, 1), dtype=np.float32)
|
|
else:
|
|
logger.info(
|
|
"RAG: Embedding %d chunk (batch=%d)...",
|
|
len(texts_to_embed),
|
|
self.embedding_batch_size,
|
|
)
|
|
arr = self._model.encode(
|
|
texts_to_embed,
|
|
batch_size=self.embedding_batch_size,
|
|
show_progress_bar=False,
|
|
convert_to_numpy=True,
|
|
normalize_embeddings=True,
|
|
)
|
|
self.embeddings = np.asarray(arr, dtype=np.float32)
|
|
logger.info("RAG: Embedding completato, shape=%s", self.embeddings.shape)
|
|
try:
|
|
self._save_index()
|
|
except Exception as e:
|
|
logger.warning("RAG: salvataggio indice NPZ fallito: %s", e)
|
|
|
|
# Appendi chunk dinamici persistiti (kb/_dynamic.jsonl) dall'avvio precedente
|
|
self._load_dynamic_chunks()
|
|
|
|
# ── Search ───────────────────────────────────────────────────────────────
|
|
|
|
def search(
|
|
self,
|
|
query: str,
|
|
*,
|
|
top_k: Optional[int] = None,
|
|
threshold: Optional[float] = None,
|
|
) -> list[SearchResult]:
|
|
"""Cerca i chunk più rilevanti per *query*.
|
|
|
|
Thread-safe: prende lock in lettura.
|
|
Restituisce lista ordinata per score decrescente, filtrata per soglia.
|
|
Lista vuota se l'indice è vuoto o nessun chunk supera la soglia.
|
|
"""
|
|
with self._lock:
|
|
if self.embeddings is None or self.embeddings.shape[0] == 0:
|
|
return []
|
|
if not query or not query.strip():
|
|
return []
|
|
|
|
k = self.top_k if top_k is None else max(1, int(top_k))
|
|
thr = self.threshold if threshold is None else float(threshold)
|
|
k = min(k, self.embeddings.shape[0])
|
|
|
|
q_vec = self._model.encode(
|
|
[query],
|
|
batch_size=1,
|
|
show_progress_bar=False,
|
|
convert_to_numpy=True,
|
|
normalize_embeddings=True,
|
|
)
|
|
q_vec = np.asarray(q_vec, dtype=np.float32).reshape(1, -1)
|
|
|
|
scores = (self.embeddings @ q_vec.T).reshape(-1)
|
|
if k == self.embeddings.shape[0]:
|
|
top_idx = np.argsort(-scores)
|
|
else:
|
|
top_idx = np.argpartition(-scores, k - 1)[:k]
|
|
top_idx = top_idx[np.argsort(-scores[top_idx])]
|
|
|
|
out: list[SearchResult] = []
|
|
for idx in top_idx:
|
|
sc = float(scores[idx])
|
|
if sc < thr:
|
|
break
|
|
out.append(SearchResult(chunk=self.chunks[int(idx)], score=sc))
|
|
return out
|
|
|
|
# ── Dynamic updates (thread-safe) ────────────────────────────────────────
|
|
|
|
def status(self) -> dict[str, Any]:
|
|
"""Restituisce stato interno (per endpoint HTTP /status)."""
|
|
with self._lock:
|
|
n = self.embeddings.shape[0] if self.embeddings is not None else 0
|
|
dim = int(self.embeddings.shape[1]) if self.embeddings is not None and n > 0 else 0
|
|
return {
|
|
"kb_path": str(self.kb_path),
|
|
"index_path": str(self.index_path),
|
|
"dynamic_path": str(self.dynamic_jsonl_path),
|
|
"embedding_model": self.embedding_model_name,
|
|
"device": self.device,
|
|
"num_chunks": n,
|
|
"embedding_dim": dim,
|
|
"top_k": self.top_k,
|
|
"threshold": self.threshold,
|
|
"language": self.language,
|
|
"inject_as": self.inject_as,
|
|
"dynamic_sources_count": len(self._dynamic_sources),
|
|
"sources": sorted({c.source for c in self.chunks}),
|
|
}
|
|
|
|
def add_chunks(
|
|
self,
|
|
items: list[dict[str, Any]],
|
|
*,
|
|
persist: bool = True,
|
|
dynamic: bool = True,
|
|
) -> dict[str, Any]:
|
|
"""Aggiunge chunk pre-costruiti in modo incrementale.
|
|
|
|
Parameters
|
|
----------
|
|
items : list[dict]
|
|
Lista di dict con chiavi:
|
|
- ``text`` (str, obbligatoria)
|
|
- ``source`` (str, opzionale, default='dynamic#N')
|
|
- ``chunk_index`` (int, opzionale)
|
|
- ``metadata`` (dict, opzionale)
|
|
persist : bool
|
|
Se True, scrive i nuovi chunk anche in kb/_dynamic.jsonl per il prossimo
|
|
avvio.
|
|
dynamic : bool
|
|
Se True, marca questi chunk come 'dinamici' (verranno ignorati da
|
|
remove_by_source solo se match esplicito).
|
|
|
|
Returns
|
|
-------
|
|
dict con:
|
|
``added`` : int — numero chunk aggiunti
|
|
``sources`` : list[str] — sorgenti uniche aggiunte
|
|
``total_after`` : int — dimensione totale dell'indice dopo l'add
|
|
"""
|
|
new_chunks: list[Chunk] = []
|
|
new_texts: list[str] = []
|
|
added_sources: set[str] = set()
|
|
with self._lock:
|
|
existing_count = len(self.chunks)
|
|
next_auto = existing_count + 1000000
|
|
for raw in items:
|
|
text = (raw.get("text") or "").strip()
|
|
if not text:
|
|
continue
|
|
source = (raw.get("source") or "").strip() or f"dynamic#{existing_count + len(new_chunks)}"
|
|
ci = int(raw.get("chunk_index") or (next_auto + len(new_chunks)))
|
|
meta = dict(raw.get("metadata") or {})
|
|
if dynamic:
|
|
meta["_dynamic"] = True
|
|
ch = Chunk(text=text, source=source, chunk_index=ci, metadata=meta)
|
|
new_chunks.append(ch)
|
|
new_texts.append(text)
|
|
added_sources.add(source)
|
|
if dynamic:
|
|
self._dynamic_sources.add(source)
|
|
if not new_chunks:
|
|
return {"added": 0, "sources": [], "total_after": len(self.chunks)}
|
|
|
|
logger.info(
|
|
"RAG add_chunks: aggiungo %d chunk (batch=%d)...",
|
|
len(new_chunks),
|
|
self.embedding_batch_size,
|
|
)
|
|
new_embs = self._model.encode(
|
|
new_texts,
|
|
batch_size=self.embedding_batch_size,
|
|
show_progress_bar=False,
|
|
convert_to_numpy=True,
|
|
normalize_embeddings=True,
|
|
)
|
|
new_embs = np.asarray(new_embs, dtype=np.float32)
|
|
# Concatenazione safe: gestisce caso embeddings vuoto (0, 1)
|
|
if self.embeddings is None or self.embeddings.size == 0 or self.embeddings.shape[1] == 1:
|
|
self.embeddings = new_embs
|
|
else:
|
|
if new_embs.shape[1] != self.embeddings.shape[1]:
|
|
raise ValueError(
|
|
f"Dimensione embedding mismatch: nuovi {new_embs.shape[1]} vs esistenti {self.embeddings.shape[1]}"
|
|
)
|
|
self.embeddings = np.vstack([self.embeddings, new_embs])
|
|
self.chunks.extend(new_chunks)
|
|
total_after = len(self.chunks)
|
|
|
|
# Persisti fuori dal lock (non blocca retrieval)
|
|
if persist:
|
|
self._append_dynamic_chunks_to_disk(new_chunks)
|
|
try:
|
|
with self._lock:
|
|
self._save_index()
|
|
except Exception as exc:
|
|
logger.warning("RAG add_chunks: salvataggio indice fallito: %s", exc)
|
|
|
|
logger.info(
|
|
"RAG add_chunks OK: added=%d total_after=%d sources=%s",
|
|
len(new_chunks),
|
|
total_after,
|
|
sorted(added_sources),
|
|
)
|
|
return {
|
|
"added": len(new_chunks),
|
|
"sources": sorted(added_sources),
|
|
"total_after": total_after,
|
|
}
|
|
|
|
def add_document(
|
|
self,
|
|
text: str,
|
|
*,
|
|
source: str,
|
|
metadata: Optional[dict[str, Any]] = None,
|
|
persist: bool = True,
|
|
dynamic: bool = True,
|
|
) -> dict[str, Any]:
|
|
"""Splitta un documento in chunk e li aggiunge in modo incrementale.
|
|
|
|
Equivalente a:
|
|
1. ``text`` → chunking ricorsivo (stesso algoritmo di build_index)
|
|
2. ``add_chunks()`` di tutti i pezzi
|
|
|
|
Parameters
|
|
----------
|
|
text : str
|
|
Testo completo documento.
|
|
source : str
|
|
Nome identificativo (es. ``crm/cliente_123_nota_20260826``).
|
|
metadata : dict, opzionale
|
|
Mappa passata a tutti i chunk derivati.
|
|
persist / dynamic : vedi ``add_chunks``.
|
|
"""
|
|
text = (text or "").strip()
|
|
if not text:
|
|
n = 0
|
|
with self._lock:
|
|
if self.embeddings is not None:
|
|
n = self.embeddings.shape[0]
|
|
return {"added": 0, "sources": [], "total_after": n}
|
|
pieces = _split_text_recursive(
|
|
text,
|
|
chunk_size=self.chunk_size,
|
|
chunk_overlap=self.chunk_overlap,
|
|
)
|
|
meta = dict(metadata or {})
|
|
items = [
|
|
{"text": p, "source": source, "chunk_index": i, "metadata": {**meta}}
|
|
for i, p in enumerate(pieces)
|
|
]
|
|
return self.add_chunks(items, persist=persist, dynamic=dynamic)
|
|
|
|
def remove_by_source(
|
|
self,
|
|
source_prefix: str,
|
|
*,
|
|
exact: bool = False,
|
|
persist: bool = True,
|
|
) -> dict[str, Any]:
|
|
"""Rimuove chunk selezionati in base al nome sorgente.
|
|
|
|
Parameters
|
|
----------
|
|
source_prefix : str
|
|
Prefisso o nome esterno (vedi ``exact``).
|
|
exact : bool
|
|
- ``False`` (default): rimuove TUTTI i chunk la cui ``source``
|
|
**inizia** con ``source_prefix``.
|
|
- ``True``: rimuove **solo** i chunk la cui ``source`` è
|
|
**esattamente uguale** a ``source_prefix``.
|
|
persist : bool
|
|
Se True, dopo la rimozione riscrive ``_dynamic.jsonl`` e
|
|
``_index.npz``.
|
|
|
|
Casi d'uso:
|
|
- ``remove_by_source('crm/cliente_123')`` → tutte le note di un cliente
|
|
- ``remove_by_source('dynamic#')`` → tutti i chunk temporanei auto-numerati
|
|
- ``remove_by_source('crm/cliente_456_nota_20260826', exact=True)`` → solo quel documento
|
|
"""
|
|
removed_indices: list[int] = []
|
|
removed_sources: set[str] = set()
|
|
with self._lock:
|
|
new_chunks: list[Chunk] = []
|
|
kept_indices: list[int] = []
|
|
for i, c in enumerate(self.chunks):
|
|
match = (c.source == source_prefix) if exact else c.source.startswith(source_prefix)
|
|
if match:
|
|
removed_indices.append(i)
|
|
removed_sources.add(c.source)
|
|
if c.source in self._dynamic_sources:
|
|
self._dynamic_sources.discard(c.source)
|
|
else:
|
|
new_chunks.append(c)
|
|
kept_indices.append(i)
|
|
if not removed_indices:
|
|
return {
|
|
"removed": 0,
|
|
"sources": [],
|
|
"total_after": len(self.chunks),
|
|
}
|
|
if self.embeddings is not None and self.embeddings.shape[0] > 0 and not (
|
|
self.embeddings.shape[0] == 1 and self.embeddings.shape[1] == 1
|
|
):
|
|
kept = np.asarray(kept_indices, dtype=np.int64)
|
|
self.embeddings = self.embeddings.take(kept, axis=0)
|
|
self.chunks = new_chunks
|
|
total_after = len(self.chunks)
|
|
|
|
if persist:
|
|
try:
|
|
with self._lock:
|
|
self._save_index()
|
|
self._rewrite_dynamic_chunks_on_disk()
|
|
except Exception as exc:
|
|
logger.warning("RAG remove: salvataggio dopo rimozione fallito: %s", exc)
|
|
|
|
logger.info(
|
|
"RAG remove_by_source(prefix=%r, exact=%s): rimossi %d chunk, %d sorgenti uniche",
|
|
source_prefix,
|
|
exact,
|
|
len(removed_indices),
|
|
len(removed_sources),
|
|
)
|
|
return {
|
|
"removed": len(removed_indices),
|
|
"sources": sorted(removed_sources),
|
|
"total_after": total_after,
|
|
"mode": "exact" if exact else "prefix",
|
|
}
|
|
|
|
def list_sources(self) -> list[dict[str, Any]]:
|
|
"""Lista di tutte le sorgenti uniche nell'indice con il conteggio chunk.
|
|
|
|
Restituisce una lista di ``[{source, chunk_count, has_dynamic}]``
|
|
ordinata per ``source`` alfabeticamente.
|
|
"""
|
|
with self._lock:
|
|
agg: dict[str, dict[str, Any]] = {}
|
|
for c in self.chunks:
|
|
entry = agg.get(c.source)
|
|
if entry is None:
|
|
entry = {"source": c.source, "chunk_count": 0, "has_dynamic": False}
|
|
agg[c.source] = entry
|
|
entry["chunk_count"] += 1
|
|
if c.metadata.get("_dynamic", False):
|
|
entry["has_dynamic"] = True
|
|
return sorted(agg.values(), key=lambda d: d["source"])
|
|
|
|
def list_chunks(
|
|
self,
|
|
*,
|
|
source_prefix: Optional[str] = None,
|
|
source_exact: Optional[str] = None,
|
|
query: Optional[str] = None,
|
|
min_score: Optional[float] = None,
|
|
offset: int = 0,
|
|
limit: int = 100,
|
|
include_text: bool = True,
|
|
include_embedding: bool = False,
|
|
) -> dict[str, Any]:
|
|
"""Elenca chunk con filtri e paginazione.
|
|
|
|
Parameters
|
|
----------
|
|
source_prefix / source_exact : str, opzionale
|
|
Filtro per nome sorgente (solo uno dei due, ``exact`` vince se entrambi).
|
|
query : str, opzionale
|
|
Se fornito, ordina i risultati per score di rilevanza coseno rispetto a
|
|
questa frase (stesso modello embedding di ``search()``).
|
|
min_score : float, opzionale
|
|
Solo per modalità ``query``: filtra chunk con score >= soglia.
|
|
offset / limit : int
|
|
Paginazione. ``limit`` massimo forzato a 1000.
|
|
include_text : bool
|
|
Includi ``chunk.text`` nella risposta (default True).
|
|
include_embedding : bool
|
|
Includi ``chunk.embedding`` (lista float). Default False (risparmia banda).
|
|
|
|
Returns
|
|
-------
|
|
dict con chiavi ``total`` (conteggio prima della paginazione), ``offset``,
|
|
``limit``, ``items`` (lista chunk dict).
|
|
"""
|
|
with self._lock:
|
|
# 1. Filtro sorgente
|
|
candidates: list[tuple[int, Chunk]] = []
|
|
for i, c in enumerate(self.chunks):
|
|
if source_exact is not None:
|
|
if c.source != source_exact:
|
|
continue
|
|
elif source_prefix is not None:
|
|
if not c.source.startswith(source_prefix):
|
|
continue
|
|
candidates.append((i, c))
|
|
|
|
# 2. Query / ordinamento per rilevanza
|
|
if query and query.strip():
|
|
q_vec = self._model.encode(
|
|
[query],
|
|
batch_size=1,
|
|
show_progress_bar=False,
|
|
convert_to_numpy=True,
|
|
normalize_embeddings=True,
|
|
)
|
|
q_vec = np.asarray(q_vec, dtype=np.float32).reshape(1, -1)
|
|
if self.embeddings is not None and self.embeddings.shape[0] > 0 and self.embeddings.shape[1] > 1:
|
|
all_scores = (self.embeddings @ q_vec.T).reshape(-1)
|
|
scored: list[tuple[int, Chunk, float]] = []
|
|
for i, c in candidates:
|
|
sc = float(all_scores[i])
|
|
if min_score is None or sc >= min_score:
|
|
scored.append((i, c, sc))
|
|
scored.sort(key=lambda t: t[2], reverse=True)
|
|
else:
|
|
scored = [(i, c, 0.0) for i, c in candidates]
|
|
else:
|
|
scored = [(i, c, None) for i, c in candidates]
|
|
# Ordine stabile per source + chunk_index
|
|
scored.sort(key=lambda t: (t[1].source, t[1].chunk_index))
|
|
|
|
total = len(scored)
|
|
offset = max(0, int(offset))
|
|
limit = max(1, min(1000, int(limit)))
|
|
page = scored[offset : offset + limit]
|
|
|
|
items: list[dict[str, Any]] = []
|
|
for i, c, sc in page:
|
|
entry: dict[str, Any] = {
|
|
"index": i,
|
|
"source": c.source,
|
|
"chunk_index": c.chunk_index,
|
|
"score": sc,
|
|
"metadata": {k: v for k, v in c.metadata.items() if k != "_dynamic"} if c.metadata else {},
|
|
"is_dynamic": bool(c.metadata.get("_dynamic", False)),
|
|
}
|
|
if include_text:
|
|
entry["text"] = c.text
|
|
if include_embedding and self.embeddings is not None and i < self.embeddings.shape[0]:
|
|
entry["embedding"] = [float(x) for x in self.embeddings[i].tolist()]
|
|
items.append(entry)
|
|
return {
|
|
"total": total,
|
|
"offset": offset,
|
|
"limit": limit,
|
|
"query": query,
|
|
"source_prefix": source_prefix,
|
|
"source_exact": source_exact,
|
|
"items": items,
|
|
}
|
|
|
|
def update_chunk(
|
|
self,
|
|
*,
|
|
index: Optional[int] = None,
|
|
source: Optional[str] = None,
|
|
chunk_index: Optional[int] = None,
|
|
new_text: Optional[str] = None,
|
|
new_metadata: Optional[dict[str, Any]] = None,
|
|
new_source: Optional[str] = None,
|
|
persist: bool = True,
|
|
) -> dict[str, Any]:
|
|
"""Aggiorna un singolo chunk (testo, metadata, sorgente).
|
|
|
|
Il chunk viene identificato tramite **una sola** di queste strategie
|
|
(la prima che matcha in ordine):
|
|
|
|
1. ``index`` — posizione intera nell'array ``self.chunks`` (solo per
|
|
sistemi che conoscono l'indice da ``list_chunks()``).
|
|
2. ``(source, chunk_index)`` — chiave business: sorgente + indice chunk interno.
|
|
3. ``source`` da sola — solo se esiste UN solo chunk con quella source
|
|
(altrimenti errore ``AMBUGUOUS_SOURCE``).
|
|
|
|
Returns
|
|
-------
|
|
dict con ``updated`` (0/1), ``chunk`` (dati aggiornati senza embedding),
|
|
``error`` (se presente).
|
|
"""
|
|
# Trova indice
|
|
found_idx: Optional[int] = None
|
|
with self._lock:
|
|
if index is not None:
|
|
if 0 <= int(index) < len(self.chunks):
|
|
found_idx = int(index)
|
|
elif source is not None and chunk_index is not None:
|
|
for i, c in enumerate(self.chunks):
|
|
if c.source == source and c.chunk_index == int(chunk_index):
|
|
found_idx = i
|
|
break
|
|
elif source is not None:
|
|
matches = [i for i, c in enumerate(self.chunks) if c.source == source]
|
|
if len(matches) == 1:
|
|
found_idx = matches[0]
|
|
elif len(matches) > 1:
|
|
return {
|
|
"updated": 0,
|
|
"error": "AMBIGUOUS_SOURCE",
|
|
"message": f"La source {source!r} ha {len(matches)} chunk — specifica anche chunk_index o usa index.",
|
|
"matches": [self.chunks[i].chunk_index for i in matches],
|
|
}
|
|
if found_idx is None:
|
|
return {"updated": 0, "error": "CHUNK_NOT_FOUND", "message": "Nessun chunk individuato con i parametri forniti."}
|
|
|
|
old = self.chunks[found_idx]
|
|
new_text_clean = (new_text or "").strip() or None
|
|
|
|
# Calcola nuova source/chunk_index
|
|
resulting_source = new_source if (new_source is not None) else old.source
|
|
resulting_chunk_index = old.chunk_index
|
|
resulting_metadata = dict(new_metadata) if new_metadata is not None else dict(old.metadata)
|
|
# Mantieni flag _dynamic se era dinamico
|
|
if old.metadata.get("_dynamic", False):
|
|
resulting_metadata["_dynamic"] = True
|
|
elif resulting_source in self._dynamic_sources:
|
|
resulting_metadata["_dynamic"] = True
|
|
|
|
# Aggiornamento testo → re-embedding
|
|
if new_text_clean is not None and new_text_clean != old.text:
|
|
vec = self._model.encode(
|
|
[new_text_clean],
|
|
batch_size=1,
|
|
show_progress_bar=False,
|
|
convert_to_numpy=True,
|
|
normalize_embeddings=True,
|
|
)
|
|
vec = np.asarray(vec, dtype=np.float32).reshape(1, -1)
|
|
if self.embeddings is not None and self.embeddings.shape[0] > 0 and self.embeddings.shape[1] == vec.shape[1]:
|
|
self.embeddings[found_idx] = vec[0]
|
|
else:
|
|
# fallback: ricalcola tutto se shape cambiato (caso estremo)
|
|
logger.warning("RAG update_chunk: shape embedding incompatibile, ricostruisco la riga")
|
|
if self.embeddings is None or self.embeddings.size == 0:
|
|
self.embeddings = np.zeros((len(self.chunks), vec.shape[1]), dtype=np.float32)
|
|
else:
|
|
new_mat = np.zeros((len(self.chunks), vec.shape[1]), dtype=np.float32)
|
|
n = min(self.embeddings.shape[0], len(self.chunks))
|
|
m = min(self.embeddings.shape[1], vec.shape[1])
|
|
new_mat[:n, :m] = self.embeddings[:n, :m]
|
|
self.embeddings = new_mat
|
|
self.embeddings[found_idx] = vec[0]
|
|
|
|
self.chunks[found_idx] = Chunk(
|
|
text=new_text_clean,
|
|
source=resulting_source,
|
|
chunk_index=resulting_chunk_index,
|
|
metadata=resulting_metadata,
|
|
)
|
|
else:
|
|
# Solo metadata/source (nessun re-embedding)
|
|
self.chunks[found_idx] = Chunk(
|
|
text=old.text,
|
|
source=resulting_source,
|
|
chunk_index=resulting_chunk_index,
|
|
metadata=resulting_metadata,
|
|
)
|
|
updated_chunk = self.chunks[found_idx]
|
|
if resulting_source in self._dynamic_sources or updated_chunk.metadata.get("_dynamic", False):
|
|
self._dynamic_sources.add(resulting_source)
|
|
if new_source is not None and new_source != old.source and old.source in self._dynamic_sources:
|
|
# vecchia source non viene rimossa da _dynamic_sources se
|
|
# ci sono altri chunk con quella source (verificato dopo persist)
|
|
pass
|
|
|
|
saved = {
|
|
"index": found_idx,
|
|
"source": updated_chunk.source,
|
|
"chunk_index": updated_chunk.chunk_index,
|
|
"text": updated_chunk.text,
|
|
"metadata": {k: v for k, v in updated_chunk.metadata.items() if k != "_dynamic"},
|
|
"is_dynamic": bool(updated_chunk.metadata.get("_dynamic", False)),
|
|
}
|
|
|
|
if persist:
|
|
try:
|
|
with self._lock:
|
|
self._save_index()
|
|
self._rewrite_dynamic_chunks_on_disk()
|
|
except Exception as exc:
|
|
logger.warning("RAG update_chunk: salvataggio fallito: %s", exc)
|
|
|
|
logger.info(
|
|
"RAG update_chunk OK: idx=%d source=%s text_changed=%s",
|
|
found_idx,
|
|
updated_chunk.source,
|
|
new_text_clean is not None,
|
|
)
|
|
return {"updated": 1, "chunk": saved}
|
|
|
|
def upsert_document(
|
|
self,
|
|
text: str,
|
|
*,
|
|
source: str,
|
|
metadata: Optional[dict[str, Any]] = None,
|
|
persist: bool = True,
|
|
dynamic: bool = True,
|
|
) -> dict[str, Any]:
|
|
"""Aggiorna in modo **atomico** un intero documento identificato da ``source``.
|
|
|
|
Operazioni (tutte dentro lo stesso lock):
|
|
1. Rimuovi tutti i chunk esistenti con ``source`` esatto
|
|
2. Splitta ``text`` con lo stesso chunker ricorsivo
|
|
3. Aggiungi i nuovi chunk con embedding e metadata
|
|
"""
|
|
text_clean = (text or "").strip()
|
|
source = (source or "").strip()
|
|
if not source:
|
|
return {"error": "SOURCE_EMPTY", "message": "source è obbligatoria."}
|
|
with self._lock:
|
|
# ── 1. Rimuovi per source exact (SENZA persist, facciamo tutto alla fine)
|
|
removed_count = 0
|
|
kept_indices: list[int] = []
|
|
for i, c in enumerate(self.chunks):
|
|
if c.source == source:
|
|
removed_count += 1
|
|
else:
|
|
kept_indices.append(i)
|
|
if removed_count > 0:
|
|
if self.embeddings is not None and self.embeddings.shape[0] > 0 and not (
|
|
self.embeddings.shape[0] == 1 and self.embeddings.shape[1] == 1
|
|
):
|
|
kept = np.asarray(kept_indices, dtype=np.int64)
|
|
self.embeddings = self.embeddings.take(kept, axis=0)
|
|
self.chunks = [self.chunks[i] for i in kept_indices]
|
|
|
|
if not text_clean:
|
|
# Testo vuoto = delete atomico
|
|
if persist:
|
|
try:
|
|
self._save_index()
|
|
self._rewrite_dynamic_chunks_on_disk()
|
|
except Exception as exc:
|
|
logger.warning("RAG upsert(delete): salvataggio fallito: %s", exc)
|
|
logger.info("RAG upsert_document %r → cancellato (testo vuoto). rimossi=%d", source, removed_count)
|
|
return {"removed": removed_count, "added": 0, "sources": [source], "total_after": len(self.chunks), "mode": "delete"}
|
|
|
|
# ── 2. + 3. Split + add (tutto dentro lo stesso lock così atomico)
|
|
pieces = _split_text_recursive(
|
|
text_clean,
|
|
chunk_size=self.chunk_size,
|
|
chunk_overlap=self.chunk_overlap,
|
|
)
|
|
new_chunks: list[Chunk] = []
|
|
new_texts: list[str] = []
|
|
meta = dict(metadata or {})
|
|
if dynamic:
|
|
meta["_dynamic"] = True
|
|
self._dynamic_sources.add(source)
|
|
for i, p in enumerate(pieces):
|
|
ch = Chunk(text=p, source=source, chunk_index=i, metadata={**meta})
|
|
new_chunks.append(ch)
|
|
new_texts.append(p)
|
|
|
|
logger.info(
|
|
"RAG upsert_document %r: rimossi=%d, aggiungo=%d chunk...",
|
|
source,
|
|
removed_count,
|
|
len(new_chunks),
|
|
)
|
|
new_embs = self._model.encode(
|
|
new_texts,
|
|
batch_size=self.embedding_batch_size,
|
|
show_progress_bar=False,
|
|
convert_to_numpy=True,
|
|
normalize_embeddings=True,
|
|
)
|
|
new_embs = np.asarray(new_embs, dtype=np.float32)
|
|
if self.embeddings is None or self.embeddings.size == 0 or self.embeddings.shape[1] == 1:
|
|
self.embeddings = new_embs
|
|
else:
|
|
if new_embs.shape[1] != self.embeddings.shape[1]:
|
|
raise ValueError(
|
|
f"Dimensione embedding mismatch: nuovi {new_embs.shape[1]} vs esistenti {self.embeddings.shape[1]}"
|
|
)
|
|
self.embeddings = np.vstack([self.embeddings, new_embs])
|
|
self.chunks.extend(new_chunks)
|
|
total_after = len(self.chunks)
|
|
|
|
if persist:
|
|
try:
|
|
with self._lock:
|
|
self._save_index()
|
|
self._rewrite_dynamic_chunks_on_disk()
|
|
except Exception as exc:
|
|
logger.warning("RAG upsert_document: salvataggio fallito: %s", exc)
|
|
|
|
logger.info(
|
|
"RAG upsert_document OK %s: removed=%d added=%d total_after=%d",
|
|
source,
|
|
removed_count,
|
|
len(new_chunks),
|
|
total_after,
|
|
)
|
|
return {
|
|
"removed": removed_count,
|
|
"added": len(new_chunks),
|
|
"sources": [source],
|
|
"total_after": total_after,
|
|
"mode": "upsert",
|
|
}
|
|
|
|
def reload_from_disk(self, *, force_rebuild: bool = True) -> dict[str, Any]:
|
|
"""Rileggere TUTTI i file md/txt/jsonl dalla cartella kb/ e ricostruisce indice.
|
|
|
|
Equivalente a riavviare con ``--rag_force_rebuild`` ma a runtime, senza spegnere.
|
|
"""
|
|
logger.info("RAG reload_from_disk: ricostruzione indice (force_rebuild=%s)...", force_rebuild)
|
|
self.build_index(force_rebuild=True)
|
|
with self._lock:
|
|
n = len(self.chunks)
|
|
dim = int(self.embeddings.shape[1]) if self.embeddings is not None and n > 0 else 0
|
|
logger.info("RAG reload_from_disk OK: %d chunk, dim=%d", n, dim)
|
|
return {"chunks": n, "embedding_dim": dim}
|
|
|
|
# ── persistence internals for dynamic chunks ──────────────────────────
|
|
|
|
def _load_dynamic_chunks(self) -> None:
|
|
"""Avvio: carica kb/_dynamic.jsonl, embedding e appende all'indice esistente."""
|
|
if not self.dynamic_jsonl_path.is_file():
|
|
return
|
|
try:
|
|
items: list[dict[str, Any]] = []
|
|
with self.dynamic_jsonl_path.open("r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
obj = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
items.append(obj)
|
|
if not items:
|
|
return
|
|
res = self.add_chunks(items, persist=False, dynamic=True)
|
|
logger.info(
|
|
"RAG _load_dynamic_chunks: caricati %d chunk da %s (added=%d)",
|
|
len(items),
|
|
self.dynamic_jsonl_path.name,
|
|
res.get("added", 0),
|
|
)
|
|
except Exception as exc:
|
|
logger.warning("RAG _load_dynamic_chunks fallito: %s", exc)
|
|
|
|
def _append_dynamic_chunks_to_disk(self, new_chunks: list[Chunk]) -> None:
|
|
try:
|
|
self.kb_path.mkdir(parents=True, exist_ok=True)
|
|
with self.dynamic_jsonl_path.open("a", encoding="utf-8") as f:
|
|
for ch in new_chunks:
|
|
f.write(
|
|
json.dumps(
|
|
{
|
|
"text": ch.text,
|
|
"source": ch.source,
|
|
"chunk_index": ch.chunk_index,
|
|
"metadata": ch.metadata,
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
+ "\n"
|
|
)
|
|
except Exception as exc:
|
|
logger.warning("RAG append _dynamic.jsonl fallito: %s", exc)
|
|
|
|
def _rewrite_dynamic_chunks_on_disk(self) -> None:
|
|
"""Riscrittura completa di _dynamic.jsonl (chiamato dopo remove)."""
|
|
try:
|
|
with self.dynamic_jsonl_path.open("w", encoding="utf-8") as f:
|
|
for ch in self.chunks:
|
|
if not ch.metadata.get("_dynamic", False):
|
|
continue
|
|
f.write(
|
|
json.dumps(
|
|
{
|
|
"text": ch.text,
|
|
"source": ch.source,
|
|
"chunk_index": ch.chunk_index,
|
|
"metadata": ch.metadata,
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
+ "\n"
|
|
)
|
|
except Exception as exc:
|
|
logger.warning("RAG rewrite _dynamic.jsonl fallito: %s", exc)
|
|
|
|
# ── Helpers per il prompt injection ──────────────────────────────────────
|
|
|
|
@staticmethod
|
|
def format_results(results: list[SearchResult], *, language: str = "es") -> str:
|
|
"""Formatta i risultati retrieval in un blocco testuale per il system prompt.
|
|
|
|
Lingua supportate per l'intestazione: 'es' (spagnolo, default), 'it', 'en'.
|
|
"""
|
|
if not results:
|
|
return ""
|
|
headers = {
|
|
"es": "Fragmentos relevantes recuperados de la base de conocimientos (utilízalos solo si responden directamente a la pregunta del usuario):",
|
|
"it": "Frammenti rilevanti recuperati dalla knowledge base (usali solo se rispondono direttamente alla domanda dell'utente):",
|
|
"en": "Relevant snippets retrieved from the knowledge base (use them only if they directly answer the user's question):",
|
|
}
|
|
hdr = headers.get(language, headers["es"])
|
|
lines = [hdr]
|
|
for i, r in enumerate(results, 1):
|
|
src = r.chunk.source or "sconosciuto"
|
|
lines.append(f"[{i}] (source={src}, score={r.score:.3f})")
|
|
lines.append(r.chunk.text.strip())
|
|
return "\n".join(lines)
|