489 lines
18 KiB
Python
489 lines
18 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
ingest.py — RAG Ingestion Pipeline (v2)
|
||
|
||
Ablauf pro Datei:
|
||
1. Extraktion → extractor.py (Text, OCR, DOCX, PDF, Bilder)
|
||
2. Normalisierung → normalizer.py (KI: Rohtext → Template-Struktur)
|
||
3. Chunking → token-basiert mit Overlap
|
||
4. Embedding → multilingual-e5-small (lokal)
|
||
5. Speichern → pgvector via SSH-Tunnel (Hetzner / anythingllm)
|
||
|
||
CLI-Befehle:
|
||
python ingest.py file <pfad> Einzelne Datei
|
||
python ingest.py dir <pfad> Ganzes Verzeichnis
|
||
python ingest.py watch <pfad> Ordner live beobachten
|
||
python ingest.py list DB-Inhalt anzeigen
|
||
python ingest.py delete <titel> Dokument aus DB löschen
|
||
|
||
Flags:
|
||
--force Bestehende Chunks überschreiben
|
||
--no-normalize KI-Normalisierung überspringen
|
||
--dry-run Extrahieren + normalisieren, aber nicht in DB speichern
|
||
--quality-min 0.4 Mindest-Qualitäts-Score (Standard: 0.0 = alles speichern)
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import uuid
|
||
import json
|
||
import argparse
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
import tiktoken
|
||
import psycopg2
|
||
from psycopg2.extras import execute_values
|
||
from sshtunnel import SSHTunnelForwarder
|
||
from dotenv import load_dotenv
|
||
|
||
from extractor import extract, is_supported, SUPPORTED_EXTENSIONS
|
||
from normalizer import normalize
|
||
|
||
load_dotenv(".env.ingest")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Konfiguration (aus .env.ingest)
|
||
# ---------------------------------------------------------------------------
|
||
SSH_HOST = os.getenv("SSH_HOST")
|
||
SSH_USER = os.getenv("SSH_USER", "root")
|
||
SSH_KEY_PATH = os.path.expanduser(os.getenv("SSH_KEY_PATH", "~/.ssh/id_ed25519"))
|
||
SSH_KEY_PASSPHRASE = os.getenv("SSH_KEY_PASSPHRASE")
|
||
DB_HOST = os.getenv("DB_HOST", "127.0.0.1")
|
||
DB_PORT = int(os.getenv("DB_PORT", 5432))
|
||
DB_NAME = os.getenv("DB_NAME", "anythingllm")
|
||
DB_USER = os.getenv("DB_USER", "anythingllm")
|
||
DB_PASSWORD = os.getenv("DB_PASSWORD")
|
||
NAMESPACE = os.getenv("NAMESPACE", "mein-workspace")
|
||
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", 500))
|
||
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", 50))
|
||
EMBED_MODEL = "intfloat/multilingual-e5-small"
|
||
|
||
_embed_model = None # lazy load
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Chunking
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def chunk_text(text: str) -> list[str]:
|
||
"""Teilt Text in überlappende Token-Chunks auf."""
|
||
enc = tiktoken.get_encoding("cl100k_base")
|
||
tokens = enc.encode(text)
|
||
chunks = []
|
||
start = 0
|
||
while start < len(tokens):
|
||
end = min(start + CHUNK_SIZE, len(tokens))
|
||
chunks.append(enc.decode(tokens[start:end]))
|
||
start += CHUNK_SIZE - CHUNK_OVERLAP
|
||
return chunks
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Embedding
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def embed(texts: list[str]) -> list[list[float]]:
|
||
"""Erstellt Embeddings via multilingual-e5-small (lokal, einmaliger Download ~120 MB)."""
|
||
global _embed_model
|
||
if _embed_model is None:
|
||
from sentence_transformers import SentenceTransformer
|
||
print(" ⏳ Lade Embedding-Modell (einmalig)...")
|
||
_embed_model = SentenceTransformer(EMBED_MODEL)
|
||
prefixed = [f"passage: {t}" for t in texts]
|
||
return _embed_model.encode(prefixed).tolist()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Datenbank
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _db_connect(tunnel_port: int):
|
||
return psycopg2.connect(
|
||
host="127.0.0.1",
|
||
port=tunnel_port,
|
||
dbname=DB_NAME,
|
||
user=DB_USER,
|
||
password=DB_PASSWORD,
|
||
)
|
||
|
||
|
||
def _tunnel():
|
||
return SSHTunnelForwarder(
|
||
(SSH_HOST, 22),
|
||
ssh_username=SSH_USER,
|
||
ssh_pkey=SSH_KEY_PATH,
|
||
ssh_private_key_password=SSH_KEY_PASSPHRASE,
|
||
remote_bind_address=("127.0.0.1", DB_PORT),
|
||
)
|
||
|
||
|
||
def _count_existing(cur, title: str) -> int:
|
||
cur.execute(
|
||
"SELECT COUNT(*) FROM anythingllm_vectors "
|
||
"WHERE metadata->>'title' = %s AND namespace = %s",
|
||
(title, NAMESPACE),
|
||
)
|
||
return cur.fetchone()[0]
|
||
|
||
|
||
def _delete_existing(cur, title: str):
|
||
cur.execute(
|
||
"DELETE FROM anythingllm_vectors "
|
||
"WHERE metadata->>'title' = %s AND namespace = %s",
|
||
(title, NAMESPACE),
|
||
)
|
||
|
||
|
||
def _insert_chunks(cur, chunks, embeddings, source_path: Path, extra_meta: dict):
|
||
now = datetime.now().isoformat()
|
||
records = []
|
||
for chunk, embedding in zip(chunks, embeddings):
|
||
metadata = {
|
||
"id": str(uuid.uuid4()),
|
||
"url": f"file://{source_path.resolve()}",
|
||
"text": chunk,
|
||
"title": source_path.name,
|
||
"docSource": "rag-ingestion-v2",
|
||
"published": now,
|
||
"wordCount": len(chunk.split()),
|
||
"chunkSource": str(source_path.resolve()),
|
||
**extra_meta,
|
||
}
|
||
records.append((
|
||
str(uuid.uuid4()),
|
||
json.dumps(metadata),
|
||
NAMESPACE,
|
||
embedding,
|
||
))
|
||
execute_values(
|
||
cur,
|
||
"INSERT INTO anythingllm_vectors (id, metadata, namespace, embedding) VALUES %s",
|
||
records,
|
||
template="(%s::uuid, %s::jsonb, %s, %s::vector)",
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Kern-Pipeline
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def ingest_file(
|
||
file_path: str,
|
||
force: bool = False,
|
||
normalize_doc: bool = True,
|
||
dry_run: bool = False,
|
||
quality_min: float = 0.0,
|
||
) -> bool:
|
||
path = Path(file_path).resolve()
|
||
|
||
if not path.exists():
|
||
print(f"❌ Datei nicht gefunden: {path}")
|
||
return False
|
||
|
||
if not is_supported(path):
|
||
print(f"⏭️ Übersprungen (Format nicht unterstützt): {path.suffix}")
|
||
return False
|
||
|
||
print(f"\n📄 {path.name}")
|
||
|
||
# ── 1. Extraktion ────────────────────────────────────────────────────────
|
||
try:
|
||
raw_text = extract(path)
|
||
print(f" ✅ Text extrahiert ({len(raw_text):,} Zeichen)")
|
||
except Exception as e:
|
||
print(f" ❌ Extraktion fehlgeschlagen: {e}")
|
||
return False
|
||
|
||
if not raw_text.strip():
|
||
print(" ⚠️ Kein Text extrahiert — Datei übersprungen")
|
||
return False
|
||
|
||
# ── 2. KI-Normalisierung ─────────────────────────────────────────────────
|
||
try:
|
||
norm = normalize(
|
||
raw_text,
|
||
filename=path.name,
|
||
skip_normalization=not normalize_doc,
|
||
)
|
||
text_to_index = norm["normalized_text"]
|
||
except Exception as e:
|
||
print(f" ❌ Normalisierung fehlgeschlagen: {e}")
|
||
print(" ↩️ Verwende Rohtext als Fallback")
|
||
norm = {
|
||
"normalized_text": raw_text,
|
||
"doc_type": "unbekannt",
|
||
"template_file": "",
|
||
"quality_score": 0.5,
|
||
"missing_fields": 0,
|
||
"was_normalized": False,
|
||
}
|
||
text_to_index = raw_text
|
||
|
||
# Qualitäts-Filter
|
||
if norm["quality_score"] < quality_min:
|
||
print(
|
||
f" ⏭️ Qualitäts-Score {norm['quality_score']:.0%} "
|
||
f"< Mindest-Score {quality_min:.0%} — übersprungen"
|
||
)
|
||
return False
|
||
|
||
extra_meta = {
|
||
"doc_type": norm["doc_type"],
|
||
"quality_score": norm["quality_score"],
|
||
"missing_fields": norm["missing_fields"],
|
||
"was_normalized": norm["was_normalized"],
|
||
"template_file": norm["template_file"],
|
||
"source_format": path.suffix.lower(),
|
||
}
|
||
|
||
# ── 3. Chunking ──────────────────────────────────────────────────────────
|
||
chunks = chunk_text(text_to_index)
|
||
print(f" ✅ {len(chunks)} Chunks erstellt")
|
||
|
||
# ── 4. Embedding ─────────────────────────────────────────────────────────
|
||
try:
|
||
embeddings = embed(chunks)
|
||
print(f" ✅ {len(embeddings)} Embeddings erstellt")
|
||
except Exception as e:
|
||
print(f" ❌ Embedding fehlgeschlagen: {e}")
|
||
return False
|
||
|
||
# ── Dry-Run Ende ─────────────────────────────────────────────────────────
|
||
if dry_run:
|
||
print(f" 🔍 Dry-Run — Typ: {norm['doc_type']} | "
|
||
f"Score: {norm['quality_score']:.0%} | "
|
||
f"Fehlend: {norm['missing_fields']}")
|
||
if norm["was_normalized"]:
|
||
print("\n" + "─" * 60)
|
||
print(text_to_index[:600] + ("..." if len(text_to_index) > 600 else ""))
|
||
print("─" * 60)
|
||
return True
|
||
|
||
# ── 5. Datenbank ─────────────────────────────────────────────────────────
|
||
print(" ⏳ Verbinde mit Datenbank (SSH-Tunnel)...")
|
||
try:
|
||
with _tunnel() as tunnel:
|
||
conn = _db_connect(tunnel.local_bind_port)
|
||
cur = conn.cursor()
|
||
|
||
existing = _count_existing(cur, path.name)
|
||
if existing > 0:
|
||
if force:
|
||
_delete_existing(cur, path.name)
|
||
print(f" 🗑️ {existing} bestehende Chunks gelöscht")
|
||
else:
|
||
print(
|
||
f" ⚠️ '{path.name}' bereits in DB ({existing} Chunks). "
|
||
"Nutze --force zum Überschreiben."
|
||
)
|
||
cur.close(); conn.close()
|
||
return False
|
||
|
||
_insert_chunks(cur, chunks, embeddings, path, extra_meta)
|
||
conn.commit()
|
||
cur.close(); conn.close()
|
||
print(f" ✅ {len(chunks)} Chunks gespeichert "
|
||
f"(Typ: {norm['doc_type']}, Score: {norm['quality_score']:.0%})")
|
||
except Exception as e:
|
||
print(f" ❌ Datenbankfehler: {e}")
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
def ingest_directory(
|
||
dir_path: str,
|
||
force: bool = False,
|
||
normalize_doc: bool = True,
|
||
dry_run: bool = False,
|
||
quality_min: float = 0.0,
|
||
):
|
||
path = Path(dir_path).resolve()
|
||
if not path.is_dir():
|
||
print(f"❌ Verzeichnis nicht gefunden: {path}")
|
||
return
|
||
|
||
files = [f for f in sorted(path.iterdir()) if f.is_file() and is_supported(f)]
|
||
if not files:
|
||
print(f"❌ Keine unterstützten Dateien in: {path}")
|
||
print(f" Unterstützte Formate: {', '.join(sorted(SUPPORTED_EXTENSIONS))}")
|
||
return
|
||
|
||
print(f"\n📁 {len(files)} Datei(en) in: {path}")
|
||
success = 0
|
||
for f in files:
|
||
if ingest_file(str(f), force=force, normalize_doc=normalize_doc,
|
||
dry_run=dry_run, quality_min=quality_min):
|
||
success += 1
|
||
|
||
print(f"\n{'─' * 50}")
|
||
print(f"✅ Fertig: {success}/{len(files)} Dateien verarbeitet")
|
||
|
||
|
||
def watch_directory(
|
||
dir_path: str,
|
||
force: bool = False,
|
||
normalize_doc: bool = True,
|
||
quality_min: float = 0.0,
|
||
):
|
||
try:
|
||
from watchdog.observers import Observer
|
||
from watchdog.events import FileSystemEventHandler
|
||
except ImportError:
|
||
print("❌ watchdog nicht installiert: pip install watchdog")
|
||
sys.exit(1)
|
||
|
||
class _Handler(FileSystemEventHandler):
|
||
def on_created(self, event):
|
||
if event.is_directory:
|
||
return
|
||
p = Path(event.src_path)
|
||
if is_supported(p):
|
||
print(f"\n🆕 Neue Datei erkannt: {p.name}")
|
||
ingest_file(str(p), force=force, normalize_doc=normalize_doc,
|
||
quality_min=quality_min)
|
||
|
||
observer = Observer()
|
||
observer.schedule(_Handler(), dir_path, recursive=False)
|
||
observer.start()
|
||
print(f"👁️ Beobachte: {dir_path}")
|
||
print(" Drücke Ctrl+C zum Beenden\n")
|
||
try:
|
||
import time
|
||
while True:
|
||
time.sleep(1)
|
||
except KeyboardInterrupt:
|
||
observer.stop()
|
||
observer.join()
|
||
|
||
|
||
def list_documents():
|
||
print(f"\n📋 Dokumente in Namespace '{NAMESPACE}':\n")
|
||
try:
|
||
with _tunnel() as tunnel:
|
||
conn = _db_connect(tunnel.local_bind_port)
|
||
cur = conn.cursor()
|
||
cur.execute("""
|
||
SELECT
|
||
metadata->>'title' AS titel,
|
||
metadata->>'doc_type' AS typ,
|
||
metadata->>'quality_score' AS score,
|
||
metadata->>'source_format' AS format,
|
||
COUNT(*) AS chunks,
|
||
MAX(metadata->>'published') AS datum
|
||
FROM anythingllm_vectors
|
||
WHERE namespace = %s
|
||
GROUP BY titel, typ, score, format
|
||
ORDER BY datum DESC
|
||
""", (NAMESPACE,))
|
||
rows = cur.fetchall()
|
||
cur.close(); conn.close()
|
||
|
||
if not rows:
|
||
print(" Keine Dokumente gefunden.")
|
||
return
|
||
|
||
header = f" {'Titel':<45} {'Typ':<20} {'Score':>6} {'Fmt':>5} {'Chunks':>6} Datum"
|
||
print(header)
|
||
print(" " + "─" * (len(header) - 2))
|
||
for titel, typ, score, fmt, chunks, datum in rows:
|
||
score_str = f"{float(score):.0%}" if score else "–"
|
||
print(
|
||
f" {(titel or '?'):<45} {(typ or '?'):<20} "
|
||
f"{score_str:>6} {(fmt or '?'):>5} {chunks:>6} "
|
||
f"{(datum or '?')[:19]}"
|
||
)
|
||
print()
|
||
except Exception as e:
|
||
print(f"❌ Datenbankfehler: {e}")
|
||
|
||
|
||
def delete_document(title: str):
|
||
print(f"\n🗑️ Lösche '{title}' aus Namespace '{NAMESPACE}'...")
|
||
try:
|
||
with _tunnel() as tunnel:
|
||
conn = _db_connect(tunnel.local_bind_port)
|
||
cur = conn.cursor()
|
||
existing = _count_existing(cur, title)
|
||
if existing == 0:
|
||
print(f" ⚠️ Kein Dokument mit Titel '{title}' gefunden")
|
||
cur.close(); conn.close()
|
||
return
|
||
_delete_existing(cur, title)
|
||
conn.commit()
|
||
cur.close(); conn.close()
|
||
print(f" ✅ {existing} Chunks gelöscht")
|
||
except Exception as e:
|
||
print(f" ❌ Datenbankfehler: {e}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _build_parser() -> argparse.ArgumentParser:
|
||
p = argparse.ArgumentParser(
|
||
description="RAG Ingestion v2 — Dokumente in pgvector importieren",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
Beispiele:
|
||
python ingest.py file ~/Downloads/fremdes_cv.pdf
|
||
python ingest.py file ~/Desktop/scan.png --dry-run
|
||
python ingest.py dir ~/Dokumente/bewerbung/ --force
|
||
python ingest.py watch ~/Desktop/scan-eingang/
|
||
python ingest.py list
|
||
python ingest.py delete "fremdes_cv.pdf"
|
||
"""
|
||
)
|
||
sub = p.add_subparsers(dest="command", required=True)
|
||
|
||
common = argparse.ArgumentParser(add_help=False)
|
||
common.add_argument("--force", action="store_true", help="Bestehende Chunks überschreiben")
|
||
common.add_argument("--no-normalize", action="store_true", help="KI-Normalisierung überspringen")
|
||
common.add_argument("--quality-min", type=float, default=0.0, metavar="0.0-1.0",
|
||
help="Mindest-Qualitäts-Score (Standard: 0.0)")
|
||
common.add_argument("--dry-run", action="store_true",
|
||
help="Nur extrahieren + normalisieren, nicht speichern")
|
||
|
||
pf = sub.add_parser("file", parents=[common], help="Einzelne Datei importieren")
|
||
pf.add_argument("path", help="Pfad zur Datei")
|
||
|
||
pd = sub.add_parser("dir", parents=[common], help="Verzeichnis importieren")
|
||
pd.add_argument("path", help="Pfad zum Verzeichnis")
|
||
|
||
pw = sub.add_parser("watch", parents=[common], help="Ordner live beobachten")
|
||
pw.add_argument("path", help="Pfad zum Verzeichnis")
|
||
|
||
sub.add_parser("list", help="Alle Dokumente in der DB anzeigen")
|
||
|
||
pdel = sub.add_parser("delete", help="Dokument aus DB löschen")
|
||
pdel.add_argument("title", help="Titel (Dateiname) des Dokuments")
|
||
|
||
return p
|
||
|
||
|
||
def main():
|
||
args = _build_parser().parse_args()
|
||
do_normalize = not getattr(args, "no_normalize", False)
|
||
force = getattr(args, "force", False)
|
||
dry_run = getattr(args, "dry_run", False)
|
||
quality_min = getattr(args, "quality_min", 0.0)
|
||
|
||
if args.command == "file":
|
||
ingest_file(args.path, force=force, normalize_doc=do_normalize,
|
||
dry_run=dry_run, quality_min=quality_min)
|
||
elif args.command == "dir":
|
||
ingest_directory(args.path, force=force, normalize_doc=do_normalize,
|
||
dry_run=dry_run, quality_min=quality_min)
|
||
elif args.command == "watch":
|
||
watch_directory(args.path, force=force, normalize_doc=do_normalize,
|
||
quality_min=quality_min)
|
||
elif args.command == "list":
|
||
list_documents()
|
||
elif args.command == "delete":
|
||
delete_document(args.title)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |