cbazza
de9c529333
- .gitignore: env-Dateien explizit ausgeschlossen - ingest.py: load_dotenv auf .env.ingest umgestellt - ingest_anythingllm.py: load_dotenv auf .env.anythingllm umgestellt - README: Dateistruktur, list/delete Befehle dokumentiert
324 lines
10 KiB
Python
324 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
RAG Ingestion Script für apply4jobs.de
|
||
Liest Dokumente (MD, TXT, PDF), chunked sie und speichert Embeddings in pgvector.
|
||
|
||
Späterer Ausbau: Tkinter UI
|
||
"""
|
||
|
||
import os
|
||
import uuid
|
||
import json
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from dotenv import load_dotenv
|
||
from openai import OpenAI
|
||
import psycopg2
|
||
from psycopg2.extras import execute_values
|
||
from sshtunnel import SSHTunnelForwarder
|
||
import tiktoken
|
||
|
||
# PDF Support (optional)
|
||
try:
|
||
from pypdf import PdfReader
|
||
PDF_SUPPORT = True
|
||
except ImportError:
|
||
PDF_SUPPORT = False
|
||
|
||
load_dotenv(".env.ingest")
|
||
|
||
# --- Konfiguration ---
|
||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||
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", None)
|
||
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"
|
||
|
||
client = OpenAI(api_key=OPENAI_API_KEY)
|
||
embed_model = None # lazy load
|
||
|
||
|
||
# --- Text-Extraktion ---
|
||
|
||
def read_file(path: Path) -> str:
|
||
"""Liest MD, TXT oder PDF und gibt den Text zurück."""
|
||
suffix = path.suffix.lower()
|
||
if suffix in (".md", ".txt"):
|
||
return path.read_text(encoding="utf-8")
|
||
elif suffix == ".pdf":
|
||
if not PDF_SUPPORT:
|
||
raise ImportError("pypdf nicht installiert: pip install pypdf")
|
||
reader = PdfReader(str(path))
|
||
return "\n\n".join(page.extract_text() or "" for page in reader.pages)
|
||
else:
|
||
raise ValueError(f"Nicht unterstütztes Dateiformat: {suffix}")
|
||
|
||
|
||
# --- Chunking ---
|
||
|
||
def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
|
||
"""Teilt Text in überlappende Chunks auf (token-basiert)."""
|
||
enc = tiktoken.get_encoding("cl100k_base")
|
||
tokens = enc.encode(text)
|
||
chunks = []
|
||
start = 0
|
||
while start < len(tokens):
|
||
end = min(start + chunk_size, len(tokens))
|
||
chunk_tokens = tokens[start:end]
|
||
chunks.append(enc.decode(chunk_tokens))
|
||
start += chunk_size - overlap
|
||
return chunks
|
||
|
||
|
||
# --- Embeddings ---
|
||
|
||
def get_embeddings(texts: list[str]) -> list[list[float]]:
|
||
"""Erstellt Embeddings via multilingual-e5-small (lokal)."""
|
||
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 get_db_connection(tunnel_port: int):
|
||
return psycopg2.connect(
|
||
host="127.0.0.1",
|
||
port=tunnel_port,
|
||
dbname=DB_NAME,
|
||
user=DB_USER,
|
||
password=DB_PASSWORD
|
||
)
|
||
|
||
|
||
def check_existing(cur, source_title: str) -> int:
|
||
"""Prüft ob Dokument bereits in der DB vorhanden ist."""
|
||
cur.execute(
|
||
"SELECT COUNT(*) FROM anythingllm_vectors WHERE metadata->>'title' = %s AND namespace = %s",
|
||
(source_title, NAMESPACE)
|
||
)
|
||
return cur.fetchone()[0]
|
||
|
||
|
||
def delete_existing(cur, source_title: str):
|
||
"""Löscht alle Chunks eines Dokuments."""
|
||
cur.execute(
|
||
"DELETE FROM anythingllm_vectors WHERE metadata->>'title' = %s AND namespace = %s",
|
||
(source_title, NAMESPACE)
|
||
)
|
||
print(f" 🗑️ Bestehende Chunks für '{source_title}' gelöscht.")
|
||
|
||
|
||
def insert_chunks(cur, chunks: list[str], embeddings: list[list[float]], source_path: Path):
|
||
"""Fügt Chunks mit Embeddings in pgvector ein."""
|
||
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-script",
|
||
"published": now,
|
||
"wordCount": len(chunk.split()),
|
||
"chunkSource": str(source_path.resolve()),
|
||
}
|
||
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)"
|
||
)
|
||
|
||
|
||
# --- Haupt-Ingestion ---
|
||
|
||
def ingest_file(file_path: str, force: bool = False):
|
||
"""
|
||
Verarbeitet eine einzelne Datei:
|
||
1. Text lesen
|
||
2. Chunken
|
||
3. Embeddings erstellen
|
||
4. In pgvector speichern (via SSH Tunnel)
|
||
"""
|
||
path = Path(file_path).resolve()
|
||
if not path.exists():
|
||
print(f"❌ Datei nicht gefunden: {path}")
|
||
return False
|
||
|
||
print(f"\n📄 Verarbeite: {path.name}")
|
||
|
||
# Text lesen
|
||
try:
|
||
text = read_file(path)
|
||
print(f" ✅ Text gelesen ({len(text)} Zeichen)")
|
||
except Exception as e:
|
||
print(f" ❌ Fehler beim Lesen: {e}")
|
||
return False
|
||
|
||
# Chunken
|
||
chunks = chunk_text(text)
|
||
print(f" ✅ {len(chunks)} Chunks erstellt")
|
||
|
||
# Embeddings
|
||
print(f" ⏳ Erstelle Embeddings via OpenAI...")
|
||
try:
|
||
embeddings = get_embeddings(chunks)
|
||
print(f" ✅ {len(embeddings)} Embeddings erstellt")
|
||
except Exception as e:
|
||
print(f" ❌ Fehler bei Embeddings: {e}")
|
||
return False
|
||
|
||
# SSH Tunnel + DB
|
||
print(f" ⏳ Verbinde mit Datenbank via SSH Tunnel...")
|
||
try:
|
||
with 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)
|
||
) as tunnel:
|
||
conn = get_db_connection(tunnel.local_bind_port)
|
||
cur = conn.cursor()
|
||
|
||
# Prüfen ob Dokument bereits existiert
|
||
existing = check_existing(cur, path.name)
|
||
if existing > 0:
|
||
if force:
|
||
delete_existing(cur, path.name)
|
||
else:
|
||
print(f" ⚠️ '{path.name}' bereits in DB ({existing} Chunks). Nutze --force zum Überschreiben.")
|
||
cur.close()
|
||
conn.close()
|
||
return False
|
||
|
||
# Einfügen
|
||
insert_chunks(cur, chunks, embeddings, path)
|
||
conn.commit()
|
||
cur.close()
|
||
conn.close()
|
||
print(f" ✅ {len(chunks)} Chunks in pgvector gespeichert")
|
||
|
||
except Exception as e:
|
||
print(f" ❌ Datenbankfehler: {e}")
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
def ingest_directory(dir_path: str, force: bool = False):
|
||
"""Verarbeitet alle MD, TXT und PDF Dateien in einem Verzeichnis."""
|
||
path = Path(dir_path).resolve()
|
||
if not path.is_dir():
|
||
print(f"❌ Verzeichnis nicht gefunden: {path}")
|
||
return
|
||
|
||
files = list(path.glob("*.md")) + list(path.glob("*.txt")) + list(path.glob("*.pdf"))
|
||
if not files:
|
||
print(f"❌ Keine unterstützten Dateien in: {path}")
|
||
return
|
||
|
||
print(f"\n📁 Verarbeite {len(files)} Dateien aus: {path}")
|
||
success = 0
|
||
for f in files:
|
||
if ingest_file(str(f), force=force):
|
||
success += 1
|
||
|
||
print(f"\n✅ Fertig: {success}/{len(files)} Dateien erfolgreich importiert.")
|
||
|
||
|
||
def list_documents():
|
||
"""Zeigt alle Dokumente in der Datenbank an."""
|
||
print(f"\n📋 Dokumente in Namespace '{NAMESPACE}':")
|
||
try:
|
||
with 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)
|
||
) as tunnel:
|
||
conn = get_db_connection(tunnel.local_bind_port)
|
||
cur = conn.cursor()
|
||
cur.execute("""
|
||
SELECT metadata->>'title', COUNT(*), MAX(metadata->>'published')
|
||
FROM anythingllm_vectors
|
||
WHERE namespace = %s
|
||
GROUP BY metadata->>'title'
|
||
ORDER BY MAX(metadata->>'published') DESC
|
||
""", (NAMESPACE,))
|
||
rows = cur.fetchall()
|
||
cur.close()
|
||
conn.close()
|
||
|
||
if not rows:
|
||
print(" Keine Dokumente gefunden.")
|
||
else:
|
||
print(f" {'Titel':<50} {'Chunks':>6} {'Erstellt'}")
|
||
print(" " + "-" * 75)
|
||
for title, count, published in rows:
|
||
print(f" {(title or 'Unknown'):<50} {count:>6} {published or '-'}")
|
||
print()
|
||
except Exception as e:
|
||
print(f"❌ Fehler: {e}")
|
||
|
||
|
||
# --- CLI ---
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(
|
||
description="RAG Ingestion Script – Dokumente in pgvector importieren"
|
||
)
|
||
subparsers = parser.add_subparsers(dest="command")
|
||
|
||
# ingest file
|
||
p_file = subparsers.add_parser("file", help="Einzelne Datei importieren")
|
||
p_file.add_argument("path", help="Pfad zur Datei (MD, TXT, PDF)")
|
||
p_file.add_argument("--force", action="store_true", help="Bestehende Chunks überschreiben")
|
||
|
||
# ingest directory
|
||
p_dir = subparsers.add_parser("dir", help="Verzeichnis importieren")
|
||
p_dir.add_argument("path", help="Pfad zum Verzeichnis")
|
||
p_dir.add_argument("--force", action="store_true", help="Bestehende Chunks überschreiben")
|
||
|
||
# list
|
||
subparsers.add_parser("list", help="Alle Dokumente in der DB anzeigen")
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.command == "file":
|
||
ingest_file(args.path, force=args.force)
|
||
elif args.command == "dir":
|
||
ingest_directory(args.path, force=args.force)
|
||
elif args.command == "list":
|
||
list_documents()
|
||
else:
|
||
parser.print_help()
|