++
This commit is contained in:
jack
@@ -0,0 +1,408 @@
|
||||
import os
|
||||
import json
|
||||
import shutil
|
||||
import chromadb
|
||||
from fastembed import TextEmbedding
|
||||
from watchdog.observers import Observer
|
||||
from watchdog.events import FileSystemEventHandler
|
||||
import threading
|
||||
|
||||
|
||||
PREFERRED_MODELS = [
|
||||
"jinaai/jina-embeddings-v2-base-zh", # 中英混合友好，~0.64GB
|
||||
"sentence-transformers/paraphrase-multilingual-mpnet-base-v2", # 多语 ~50 语种
|
||||
"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", # 多语轻量版
|
||||
"intfloat/multilingual-e5-large", # 多语更强，体积约 2.2GB
|
||||
]
|
||||
|
||||
# Final chosen model will be detected at runtime from supported list
|
||||
MODEL_NAME = None
|
||||
COLLECTION_NAME = "brain_kb_v5"
|
||||
BATCH_SIZE = 128 # batch upserts to avoid huge single writes
|
||||
|
||||
# Optional imports for different file types
|
||||
try:
|
||||
from pypdf import PdfReader
|
||||
except ImportError:
|
||||
PdfReader = None
|
||||
|
||||
try:
|
||||
from docx import Document
|
||||
except ImportError:
|
||||
Document = None
|
||||
|
||||
class KnowledgeBase:
|
||||
def __init__(self, kb_path="knowledge", db_path="vector_db"):
|
||||
self.kb_path = os.path.abspath(kb_path)
|
||||
self.db_path = os.path.abspath(db_path)
|
||||
self.meta_path = os.path.join(self.db_path, "_meta.json")
|
||||
self.manifest_path = os.path.join(self.db_path, "_manifest.json")
|
||||
self._collection_reset_guard = False
|
||||
self._query_reset_guard = False
|
||||
self._sync_lock = threading.Lock()
|
||||
|
||||
if not os.path.exists(self.kb_path):
|
||||
os.makedirs(self.kb_path)
|
||||
|
||||
# Initialize Embedding Model (BAAI/bge-small-zh-v1.5 is ~100MB)
|
||||
# This will load from cache if already downloaded
|
||||
# Pick the first available model from the preferred list
|
||||
_supported_raw = TextEmbedding.list_supported_models()
|
||||
supported = set()
|
||||
for item in _supported_raw:
|
||||
if isinstance(item, dict) and "model" in item:
|
||||
supported.add(item["model"])
|
||||
elif isinstance(item, str):
|
||||
supported.add(item)
|
||||
chosen = None
|
||||
for name in PREFERRED_MODELS:
|
||||
if name in supported:
|
||||
chosen = name
|
||||
break
|
||||
if not chosen:
|
||||
raise RuntimeError(
|
||||
"No preferred embedding models are supported by fastembed. "
|
||||
"Please check available models via TextEmbedding.list_supported_models()."
|
||||
)
|
||||
|
||||
print(f"Loading Knowledge Base Embedding Model: {chosen} (may take some time on first run)...")
|
||||
try:
|
||||
self.model = TextEmbedding(model_name=chosen)
|
||||
print("Embedding Model loaded successfully.")
|
||||
except Exception as e:
|
||||
print(f"Error loading embedding model: {e}")
|
||||
raise
|
||||
|
||||
# Store chosen model name for reference
|
||||
global MODEL_NAME
|
||||
MODEL_NAME = chosen
|
||||
|
||||
# Cache embedding dimension (detects library/model changes that corrupt existing indexes)
|
||||
self.embed_dim = self._get_embedding_dim()
|
||||
self.chroma_version = getattr(chromadb, "__version__", "unknown")
|
||||
|
||||
# If the stored index was built with a different model/dimension/chromadb version, wipe it
|
||||
self._maybe_reset_for_incompatibility(chosen, self.embed_dim, self.chroma_version)
|
||||
|
||||
# Initialize Vector DB
|
||||
self._init_collection()
|
||||
self._healthcheck()
|
||||
|
||||
# Initial sync
|
||||
self.sync_knowledge()
|
||||
|
||||
# Start Watcher
|
||||
self.start_watcher()
|
||||
|
||||
def _init_collection(self, recreate: bool = False):
|
||||
"""(Re)initialize Chroma client/collection. If recreate=True, wipe on-disk index."""
|
||||
if recreate and os.path.exists(self.db_path):
|
||||
shutil.rmtree(self.db_path, ignore_errors=True)
|
||||
try:
|
||||
self.client = chromadb.PersistentClient(path=self.db_path)
|
||||
self.collection = self.client.get_or_create_collection(
|
||||
name=COLLECTION_NAME,
|
||||
metadata={"hnsw:space": "cosine"}
|
||||
)
|
||||
except Exception as exc:
|
||||
# If collection load itself fails, wipe and retry once to clear corrupted segments
|
||||
if not recreate:
|
||||
shutil.rmtree(self.db_path, ignore_errors=True)
|
||||
return self._init_collection(recreate=True)
|
||||
raise
|
||||
|
||||
# Persist metadata about the embedding model used to build this index
|
||||
try:
|
||||
os.makedirs(self.db_path, exist_ok=True)
|
||||
with open(self.meta_path, "w", encoding="utf-8") as f:
|
||||
json.dump({
|
||||
"model": MODEL_NAME,
|
||||
"embed_dim": self.embed_dim,
|
||||
"chroma_version": self.chroma_version,
|
||||
}, f)
|
||||
except Exception:
|
||||
pass # Metadata failure should not block runtime
|
||||
|
||||
def _healthcheck(self):
|
||||
"""Validate index readability right after startup; rebuild if corrupted."""
|
||||
try:
|
||||
_ = self.collection.count()
|
||||
except Exception as e:
|
||||
msg = str(e).lower()
|
||||
if any(x in msg for x in ["hnsw", "segment", "compaction", "backfill"]):
|
||||
print("Detected index corruption on startup. Rebuilding vector_db...")
|
||||
shutil.rmtree(self.db_path, ignore_errors=True)
|
||||
self._init_collection(recreate=True)
|
||||
self.sync_knowledge(allow_reset=False)
|
||||
else:
|
||||
print(f"Index healthcheck encountered an unexpected error: {e}")
|
||||
|
||||
def _maybe_reset_for_incompatibility(self, chosen_model: str, embed_dim: int, chroma_version: str):
|
||||
"""If existing index meta differs (model/dimension/chromadb), wipe it."""
|
||||
if not os.path.exists(self.db_path):
|
||||
return
|
||||
try:
|
||||
with open(self.meta_path, "r", encoding="utf-8") as f:
|
||||
meta = json.load(f)
|
||||
prev_model = meta.get("model")
|
||||
prev_dim = meta.get("embed_dim")
|
||||
prev_chroma = meta.get("chroma_version")
|
||||
if prev_model != chosen_model or prev_dim != embed_dim or prev_chroma != chroma_version:
|
||||
shutil.rmtree(self.db_path, ignore_errors=True)
|
||||
except Exception:
|
||||
# If meta cannot be read, assume stale/corrupted and rebuild
|
||||
shutil.rmtree(self.db_path, ignore_errors=True)
|
||||
|
||||
def _get_embedding_dim(self) -> int:
|
||||
for vec in self.model.embed(["dimension_probe"]):
|
||||
try:
|
||||
return len(vec)
|
||||
except Exception:
|
||||
return len(list(vec))
|
||||
raise RuntimeError("Failed to determine embedding dimension")
|
||||
|
||||
def sync_knowledge(self, allow_reset: bool = True):
|
||||
"""Scans the knowledge folder and updates the vector database."""
|
||||
if not self._sync_lock.acquire(blocking=False):
|
||||
print("Sync already running, skip this trigger.")
|
||||
return
|
||||
|
||||
print("Syncing knowledge base...")
|
||||
manifest = self._load_manifest()
|
||||
updated_manifest = {}
|
||||
supported_extensions = (".txt", ".md", ".pdf", ".docx", ".json")
|
||||
current_files = []
|
||||
try:
|
||||
for filename in os.listdir(self.kb_path):
|
||||
file_path = os.path.join(self.kb_path, filename)
|
||||
if os.path.isfile(file_path) and filename.lower().endswith(supported_extensions):
|
||||
current_files.append(filename)
|
||||
mtime = os.path.getmtime(file_path)
|
||||
size = os.path.getsize(file_path)
|
||||
prev_meta = manifest.get(filename)
|
||||
# Skip unchanged files
|
||||
if prev_meta and prev_meta.get("mtime") == mtime and prev_meta.get("size") == size:
|
||||
updated_manifest[filename] = prev_meta
|
||||
continue
|
||||
try:
|
||||
content = self._extract_text(file_path)
|
||||
if content:
|
||||
# Sliding window chunking on original text
|
||||
chunk_size = 800
|
||||
overlap = 80
|
||||
original_chunks = []
|
||||
for i in range(0, len(content), chunk_size - overlap):
|
||||
chunk = content[i:i + chunk_size].strip()
|
||||
if chunk:
|
||||
original_chunks.append(chunk)
|
||||
|
||||
if original_chunks:
|
||||
# Normalize for embedding generation only (not for storage)
|
||||
normalized_chunks = [c.lower().replace('_', ' ') for c in original_chunks]
|
||||
|
||||
ids = [f"{filename}_{i}" for i in range(len(original_chunks))]
|
||||
metadatas = [{"source": filename, "chunk": i} for i in range(len(original_chunks))]
|
||||
|
||||
# Compute embeddings from normalized text
|
||||
embeddings = []
|
||||
for v in self.model.embed(normalized_chunks):
|
||||
try:
|
||||
embeddings.append(v.tolist())
|
||||
except Exception:
|
||||
embeddings.append(list(v))
|
||||
|
||||
# Store ORIGINAL text (not normalized) so users see the real content
|
||||
for start in range(0, len(original_chunks), BATCH_SIZE):
|
||||
end = start + BATCH_SIZE
|
||||
self.collection.upsert(
|
||||
documents=original_chunks[start:end],
|
||||
ids=ids[start:end],
|
||||
metadatas=metadatas[start:end],
|
||||
embeddings=embeddings[start:end]
|
||||
)
|
||||
print(f" ✓ Indexed {filename}: {len(original_chunks)} chunks (batched)")
|
||||
updated_manifest[filename] = {"mtime": mtime, "size": size}
|
||||
except Exception as e:
|
||||
err_msg = str(e)
|
||||
print(f"Error processing {filename}: {err_msg}")
|
||||
# Auto-recover if HNSW/compaction/index errors occur
|
||||
if allow_reset and any(x in err_msg.lower() for x in ["hnsw", "compaction", "segment reader"]):
|
||||
if not self._collection_reset_guard:
|
||||
print("Detected index corruption. Rebuilding vector_db and retrying sync once...")
|
||||
self._collection_reset_guard = True
|
||||
self._init_collection(recreate=True)
|
||||
return self.sync_knowledge(allow_reset=False)
|
||||
# Remove deleted files from the index
|
||||
deleted_files = set(manifest.keys()) - set(current_files)
|
||||
for filename in deleted_files:
|
||||
try:
|
||||
self.collection.delete(where={"source": filename})
|
||||
print(f" ✓ Removed deleted file from index: {filename}")
|
||||
except Exception as e:
|
||||
print(f" ! Failed to remove {filename}: {e}")
|
||||
# Persist manifest
|
||||
self._save_manifest(updated_manifest)
|
||||
print("Knowledge base sync complete.")
|
||||
finally:
|
||||
self._sync_lock.release()
|
||||
|
||||
def _extract_text(self, file_path):
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
if ext == ".txt":
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
return f.read()
|
||||
elif ext == ".md":
|
||||
# Treat Markdown as plain text for retrieval
|
||||
with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
|
||||
return f.read()
|
||||
elif ext == ".pdf":
|
||||
if PdfReader:
|
||||
reader = PdfReader(file_path)
|
||||
text = ""
|
||||
for page in reader.pages:
|
||||
text += page.extract_text() + "\n"
|
||||
return text
|
||||
else:
|
||||
print("pypdf not installed, skipping PDF.")
|
||||
elif ext == ".docx":
|
||||
if Document:
|
||||
doc = Document(file_path)
|
||||
return "\n".join([para.text for para in doc.paragraphs])
|
||||
else:
|
||||
print("python-docx not installed, skipping Word.")
|
||||
elif ext == ".json":
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
return json.dumps(data, ensure_ascii=False, indent=2)
|
||||
return None
|
||||
|
||||
def query(self, text, top_k=5, distance_threshold=0.8, allow_reset: bool = True):
|
||||
"""Retrieves relevant snippets from the knowledge base.
|
||||
|
||||
Uses cosine distance (lower is better). A result is treated as a hit only
|
||||
when best_distance <= distance_threshold.
|
||||
Returns:
|
||||
dict: {"hit": bool, "context": str, "hits": [{source, chunk, distance, text}, ...]}
|
||||
"""
|
||||
try:
|
||||
# Normalize query same as indexed content
|
||||
normalized_text = text.lower().replace('_', ' ')
|
||||
|
||||
q_vec = None
|
||||
for v in self.model.embed([normalized_text]):
|
||||
try:
|
||||
q_vec = v.tolist()
|
||||
except Exception:
|
||||
q_vec = list(v)
|
||||
break
|
||||
if q_vec is None:
|
||||
return {"hit": False, "context": "", "hits": []}
|
||||
|
||||
results = self.collection.query(
|
||||
query_embeddings=[q_vec],
|
||||
n_results=top_k,
|
||||
include=["documents", "metadatas", "distances"]
|
||||
)
|
||||
|
||||
docs = (results or {}).get("documents") or []
|
||||
metas = (results or {}).get("metadatas") or []
|
||||
dists = (results or {}).get("distances") or []
|
||||
|
||||
if not docs or not docs[0]:
|
||||
print("[KB Query] No results returned from collection")
|
||||
return {"hit": False, "context": "", "hits": []}
|
||||
|
||||
docs0 = docs[0]
|
||||
metas0 = metas[0] if metas and metas[0] else [{} for _ in docs0]
|
||||
dists0 = dists[0] if dists and dists[0] else [None for _ in docs0]
|
||||
|
||||
hits = []
|
||||
for doc_text, meta, dist in zip(docs0, metas0, dists0):
|
||||
hits.append({
|
||||
"source": (meta or {}).get("source", ""),
|
||||
"chunk": (meta or {}).get("chunk", None),
|
||||
"distance": dist,
|
||||
"text": doc_text,
|
||||
})
|
||||
|
||||
best = hits[0].get("distance")
|
||||
is_hit = (best is not None) and (best <= distance_threshold)
|
||||
|
||||
# Debug log
|
||||
best_str = f"{best:.4f}" if best is not None else "N/A"
|
||||
print(f"[KB Query] '{text[:50]}...' -> best_dist={best_str}, threshold={distance_threshold}, hit={is_hit}")
|
||||
if hits:
|
||||
top3_dists = [f"{h['distance']:.4f}" if h['distance'] is not None else "N/A" for h in hits[:3]]
|
||||
print(f"[KB Query] Top 3 distances: {top3_dists}")
|
||||
|
||||
context = "\n---\n".join([h["text"] for h in hits]) if is_hit else ""
|
||||
return {"hit": is_hit, "context": context, "hits": hits}
|
||||
except Exception as e:
|
||||
err_msg = str(e)
|
||||
print(f"Query error: {err_msg}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
# Auto-recover if HNSW/compaction/backfill errors surface during query
|
||||
if allow_reset and any(x in err_msg.lower() for x in ["hnsw", "compaction", "segment reader", "backfill"]):
|
||||
if not self._query_reset_guard:
|
||||
print("Detected index corruption during query. Rebuilding vector_db and retrying once...")
|
||||
self._query_reset_guard = True
|
||||
try:
|
||||
self._init_collection(recreate=True)
|
||||
self.sync_knowledge(allow_reset=False)
|
||||
# Retry query once with guard disabled to avoid loops
|
||||
self._query_reset_guard = False
|
||||
return self.query(text, top_k=top_k, distance_threshold=distance_threshold, allow_reset=False)
|
||||
except Exception as inner_e:
|
||||
print(f"Auto-rebuild after query failure also failed: {inner_e}")
|
||||
self._query_reset_guard = False
|
||||
return {"hit": False, "context": "", "hits": []}
|
||||
|
||||
def start_watcher(self):
|
||||
event_handler = KBHandler(self)
|
||||
self.observer = Observer()
|
||||
self.observer.schedule(event_handler, self.kb_path, recursive=False)
|
||||
self.observer.start()
|
||||
|
||||
def _load_manifest(self):
|
||||
if not os.path.exists(self.manifest_path):
|
||||
return {}
|
||||
try:
|
||||
with open(self.manifest_path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
def _save_manifest(self, data):
|
||||
try:
|
||||
os.makedirs(self.db_path, exist_ok=True)
|
||||
with open(self.manifest_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
except Exception as e:
|
||||
print(f" ! Failed to save manifest: {e}")
|
||||
|
||||
class KBHandler(FileSystemEventHandler):
|
||||
def __init__(self, kb_instance):
|
||||
self.kb = kb_instance
|
||||
self.supported_extensions = (".txt", ".md", ".pdf", ".docx", ".json")
|
||||
self._debounce_timer = None
|
||||
|
||||
def _trigger_sync(self):
|
||||
def run():
|
||||
self.kb.sync_knowledge()
|
||||
if self._debounce_timer and self._debounce_timer.is_alive():
|
||||
return
|
||||
self._debounce_timer = threading.Timer(0.5, run)
|
||||
self._debounce_timer.start()
|
||||
|
||||
def on_modified(self, event):
|
||||
if not event.is_directory and event.src_path.lower().endswith(self.supported_extensions):
|
||||
print(f"File modified: {event.src_path}. Re-syncing...")
|
||||
self._trigger_sync()
|
||||
|
||||
def on_created(self, event):
|
||||
if not event.is_directory and event.src_path.lower().endswith(self.supported_extensions):
|
||||
print(f"File created: {event.src_path}. Syncing...")
|
||||
self._trigger_sync()
|
||||
|
||||
Reference in New Issue
Block a user