Files
ai.bewerbung---vectorize-do…/ingest_job_matching.py
T

342 lines
11 KiB
Python
Raw Normal View History

#!/usr/bin/env python3
"""
RAG Ingestion Script für Job-Matching Workflow
Liest MD/TXT/PDF Dokumente, chunked sie und speichert Embeddings
in der lokalen pgvector Instanz auf der NAS (Port 5433).
Verwendung:
python ingest_job_matching.py file <pfad> # Einzelne Datei
python ingest_job_matching.py dir <pfad> # Verzeichnis
python ingest_job_matching.py list # Alle Dokumente anzeigen
python ingest_job_matching.py clear <filename> # Dokument löschen
"""
import os
import uuid
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
import tiktoken
# PDF Support (optional)
try:
from pypdf import PdfReader
PDF_SUPPORT = True
except ImportError:
PDF_SUPPORT = False
# Lade job_matching spezifische .env Datei
env_path = Path(__file__).parent / ".env.job_matching"
load_dotenv(dotenv_path=env_path)
# --- Konfiguration ---
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
DB_HOST = os.getenv("JM_DB_HOST", "192.168.178.128")
DB_PORT = int(os.getenv("JM_DB_PORT", 5433))
DB_NAME = os.getenv("JM_DB_NAME", "job_matching")
DB_USER = os.getenv("JM_DB_USER", "jobmatch")
DB_PASSWORD = os.getenv("JM_DB_PASSWORD")
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", 500))
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", 50))
EMBED_MODEL = "text-embedding-3-small"
EMBED_DIM = 1536
client = OpenAI(api_key=OPENAI_API_KEY)
# --- Datenbank ---
def get_db_connection():
"""Direkte Verbindung zur pgvector Instanz auf der NAS."""
return psycopg2.connect(
host=DB_HOST,
port=DB_PORT,
dbname=DB_NAME,
user=DB_USER,
password=DB_PASSWORD
)
def ensure_table(conn):
"""Stellt sicher dass die profile_documents Tabelle existiert."""
with conn.cursor() as cur:
cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
cur.execute(f"""
CREATE TABLE IF NOT EXISTS profile_documents (
id SERIAL PRIMARY KEY,
filename TEXT NOT NULL,
chunk_index INTEGER NOT NULL,
content TEXT NOT NULL,
embedding vector({EMBED_DIM}),
created_at TIMESTAMP DEFAULT NOW(),
UNIQUE(filename, chunk_index)
);
""")
cur.execute("""
CREATE INDEX IF NOT EXISTS profile_documents_embedding_idx
ON profile_documents
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 10);
""")
conn.commit()
print(" ✅ Tabelle profile_documents bereit")
# --- 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 OpenAI text-embedding-3-small."""
response = client.embeddings.create(
model=EMBED_MODEL,
input=texts
)
return [item.embedding for item in response.data]
# --- Ingestion ---
def check_existing(conn, filename: str) -> int:
"""Prüft ob Dokument bereits in der DB vorhanden ist."""
with conn.cursor() as cur:
cur.execute(
"SELECT COUNT(*) FROM profile_documents WHERE filename = %s",
(filename,)
)
return cur.fetchone()[0]
def delete_existing(conn, filename: str):
"""Löscht alle Chunks eines Dokuments."""
with conn.cursor() as cur:
cur.execute(
"DELETE FROM profile_documents WHERE filename = %s",
(filename,)
)
conn.commit()
print(f" 🗑️ Bestehende Chunks für '{filename}' gelöscht.")
def insert_chunks(conn, chunks: list[str], embeddings: list[list[float]], filename: str):
"""Fügt Chunks mit Embeddings in pgvector ein."""
records = [
(filename, i, chunk, embedding)
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings))
]
with conn.cursor() as cur:
execute_values(
cur,
"""
INSERT INTO profile_documents (filename, chunk_index, content, embedding)
VALUES %s
ON CONFLICT (filename, chunk_index) DO UPDATE
SET content = EXCLUDED.content,
embedding = EXCLUDED.embedding,
created_at = NOW()
""",
records,
template="(%s, %s, %s, %s::vector)"
)
conn.commit()
def ingest_file(file_path: str, force: bool = False) -> bool:
"""Verarbeitet eine einzelne Datei."""
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 ({EMBED_MODEL})...")
try:
embeddings = []
batch_size = 100
for i in range(0, len(chunks), batch_size):
batch = chunks[i:i + batch_size]
embeddings.extend(get_embeddings(batch))
print(f" ✅ {len(embeddings)} Embeddings erstellt")
except Exception as e:
print(f" ❌ Fehler bei Embeddings: {e}")
return False
# Datenbank
print(f" ⏳ Verbinde mit pgvector auf {DB_HOST}:{DB_PORT}...")
try:
conn = get_db_connection()
ensure_table(conn)
existing = check_existing(conn, path.name)
if existing > 0:
if force:
delete_existing(conn, path.name)
else:
print(f" ⚠️ '{path.name}' bereits in DB ({existing} Chunks). Nutze --force zum Überschreiben.")
conn.close()
return False
insert_chunks(conn, chunks, embeddings, path.name)
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"))
)
files = [f for f in files if not f.name.startswith(".")]
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 sorted(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 profile_documents ({DB_HOST}:{DB_PORT}/{DB_NAME}):")
try:
conn = get_db_connection()
with conn.cursor() as cur:
cur.execute("""
SELECT filename, COUNT(*) as chunks, MAX(created_at) as imported_at
FROM profile_documents
GROUP BY filename
ORDER BY MAX(created_at) DESC
""")
rows = cur.fetchall()
conn.close()
if not rows:
print(" Keine Dokumente gefunden.")
else:
print(f"\n {'Dateiname':<55} {'Chunks':>6} {'Importiert'}")
print(" " + "-" * 80)
for filename, chunks, imported_at in rows:
ts = imported_at.strftime("%d.%m.%Y %H:%M") if imported_at else "-"
print(f" {filename:<55} {chunks:>6} {ts}")
print()
except Exception as e:
print(f"❌ Fehler: {e}")
def clear_document(filename: str):
"""Löscht ein Dokument aus der Datenbank."""
try:
conn = get_db_connection()
existing = check_existing(conn, filename)
if existing == 0:
print(f"⚠️ '{filename}' nicht in der Datenbank gefunden.")
else:
delete_existing(conn, filename)
print(f"✅ '{filename}' gelöscht ({existing} Chunks).")
conn.close()
except Exception as e:
print(f"❌ Fehler: {e}")
# --- CLI ---
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(
description="Job-Matching RAG Ingestion Profil-Dokumente in pgvector importieren"
)
subparsers = parser.add_subparsers(dest="command")
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")
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")
subparsers.add_parser("list", help="Alle Dokumente in der DB anzeigen")
p_clear = subparsers.add_parser("clear", help="Dokument aus DB löschen")
p_clear.add_argument("filename", help="Dateiname (z.B. 'Lebenslauf - 2026 - Sebastian Fröhlich.md')")
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()
elif args.command == "clear":
clear_document(args.filename)
else:
parser.print_help()