app.py

# app.py - v1.5
# Beschreibung: Finale, robuste Version mit umfangreichem Debugging, Assertions
#               und Korrekturen für API-Inkompatibilitäten (Gemma-3 Prozessor, LangChain Embeddings).

import os
import torch
import gradio as gr
import time  # Für Debug-Timings

from typing import List, Tuple, Generator, Dict
from threading import Thread

# ML / Transformers
from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextStreamer

# Dokumentenverarbeitung & RAG
from pypdf import PdfReader
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
# NEUER, KORREKTER IMPORT für Zukunftssicherheit
from langchain_huggingface import HuggingFaceEmbeddings

# --------------------------------------------------------------------
# Konfiguration & Globale States
# --------------------------------------------------------------------
EMBED_MODEL_ID = "google/embeddinggemma-300m"
LLM_MODEL_ID = "google/gemma-3-4b-it"

EMBEDDING_FUNCTION: HuggingFaceEmbeddings = None
LLM_MODEL: Gemma3ForConditionalGeneration = None
LLM_PROCESSOR: AutoProcessor = None
VECTOR_STORE: FAISS = None

# --- DEBUGGING HELPER ---
def print_debug(message: str):
    """Konsistente Debug-Ausgabe mit Zeitstempel."""
    print(f"[DEBUG {time.strftime('%H:%M:%S')}] {message}")

# --------------------------------------------------------------------
# Model Loading
# --------------------------------------------------------------------
def get_device() -> torch.device:
    """Gibt das verfügbare torch-Device zurück (CUDA, wenn verfügbar)."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def get_embedding_function() -> HuggingFaceEmbeddings:
    """Lädt das Embedding-Modell über den korrekten LangChain-Wrapper."""
    global EMBEDDING_FUNCTION
    if EMBEDDING_FUNCTION is None:
        device = get_device()
        print_debug(f"Initialisiere Embedding-Modell '{EMBED_MODEL_ID}' auf Device '{device}'.")
        # Nutzt jetzt die dedizierte, zukunftssichere Klasse.
        EMBEDDING_FUNCTION = HuggingFaceEmbeddings(
            model_name=EMBED_MODEL_ID,
            model_kwargs={'device': device}
        )
        print_debug("Embedding-Modell erfolgreich initialisiert.")
    return EMBEDDING_FUNCTION


def get_llm() -> Tuple[Gemma3ForConditionalGeneration, AutoProcessor]:
    """Lädt und initialisiert das LLM und den zugehörigen Prozessor."""
    global LLM_MODEL, LLM_PROCESSOR
    if LLM_MODEL is None or LLM_PROCESSOR is None:
        device = get_device()
        print_debug(f"Initialisiere LLM '{LLM_MODEL_ID}' auf Device '{device}'.")
        dtype = torch.bfloat16 if "cuda" in device.type else torch.float32

        LLM_MODEL = Gemma3ForConditionalGeneration.from_pretrained(
            LLM_MODEL_ID,
            torch_dtype=dtype,
            device_map="auto",
        ).eval()
        LLM_PROCESSOR = AutoProcessor.from_pretrained(LLM_MODEL_ID)
        print_debug("LLM und Prozessor erfolgreich initialisiert.")

    return LLM_MODEL, LLM_PROCESSOR


# --------------------------------------------------------------------
# Datei-Handling & Chunking
# --------------------------------------------------------------------
def extract_text_from_file(path: str) -> str:
    """Extrahiert Text aus verschiedenen Dateitypen."""
    # ... (Keine Änderungen in dieser Funktion, bleibt wie in v1.4)
    ext = os.path.splitext(path)[1].lower()
    if ext in [".txt", ".md", ".markdown"]:
        with open(path, "r", encoding="utf-8", errors="ignore") as f: return f.read()
    if ext == ".pdf":
        text_parts = []
        try:
            reader = PdfReader(path)
            for page in reader.pages:
                page_text = page.extract_text()
                if page_text: text_parts.append(page_text)
            return "\n\n".join(text_parts)
        except Exception as e:
            print(f"Fehler beim Lesen von PDF {path}: {e}"); return ""
    try:
        with open(path, "r", encoding="utf-8", errors="ignore") as f: return f.read()
    except Exception: return ""

def get_text_splitter() -> RecursiveCharacterTextSplitter:
    """Erstellt einen semantisch bewussten Text-Splitter."""
    return RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200, length_function=len)

# --------------------------------------------------------------------
# Indexing / RAG mit FAISS
# --------------------------------------------------------------------
def index_files(file_paths: List[str], progress=gr.Progress(track_tqdm=True)) -> str:
    """Liest Dateien, erstellt Chunks und baut/aktualisiert einen FAISS-Vektorindex."""
    global VECTOR_STORE
    if not file_paths:
        return "Keine Dateien zum Indexieren ausgewählt."
    print_debug(f"Indexierung gestartet für {len(file_paths)} Datei(en).")

    embedding_function = get_embedding_function()
    text_splitter = get_text_splitter()

    documents: List[Document] = []
    for path in progress.tqdm(file_paths, desc="1/2: Dateien verarbeiten & chunken"):
        if path is None: continue

        text = extract_text_from_file(path)
        if not text.strip():
            print_debug(f"Datei '{os.path.basename(path)}' enthält keinen extrahierbaren Text.")
            continue

        chunks = text_splitter.split_text(text)
        source_name = os.path.basename(path)
        for chunk in chunks:
            doc = Document(page_content=chunk, metadata={"source": source_name})
            documents.append(doc)

    # ASSERT: Sicherstellen, dass wir eine Liste von Document-Objekten haben.
    assert all(isinstance(d, Document) for d in documents), "Alle Elemente in 'documents' müssen vom Typ langchain.Document sein."
    print_debug(f"Erfolgreich {len(documents)} Chunks aus den Dateien erstellt.")

    if not documents:
        return "Kein Text in den Dateien gefunden, der indexiert werden konnte."

    progress(0.5, desc="2/2: Embeddings erstellen & FAISS Index aufbauen...")
    new_store = FAISS.from_documents(documents, embedding_function)
    print_debug("FAISS Index erfolgreich aus Dokumenten erstellt.")

    if VECTOR_STORE is None:
        VECTOR_STORE = new_store
    else:
        VECTOR_STORE.add_documents(documents)
        print_debug("Neuen Index mit bestehendem Index zusammengeführt.")

    # ASSERT: Sicherstellen, dass der Index nun existiert und Elemente enthält.
    assert VECTOR_STORE is not None and VECTOR_STORE.index.ntotal > 0, "VECTOR_STORE wurde nicht korrekt initialisiert."

    final_count = VECTOR_STORE.index.ntotal
    print_debug(f"Indexierung abgeschlossen. Gesamtanzahl der Chunks im Index: {final_count}")
    return f"Index aktualisiert: {final_count} Chunks insgesamt."

def clear_index() -> str:
    """Leert den Vektorindex."""
    global VECTOR_STORE
    VECTOR_STORE = None
    import gc; gc.collect()
    print_debug("Vektor-Index wurde geleert.")
    return "Index geleert."

def retrieve_relevant_chunks(query: str, top_k: int = 5) -> List[Dict]:
    """Sucht die relevantesten Chunks mit FAISS."""
    if VECTOR_STORE is None:
        print_debug("Retrieval versucht, aber Vektor-Index ist leer.")
        return []

    print_debug(f"Suche nach {top_k} relevanten Chunks für die Anfrage: '{query[:80]}...'")
    results_with_scores = VECTOR_STORE.similarity_search_with_score(query, k=top_k)

    formatted_results = []
    for doc, score in results_with_scores:
        formatted_results.append({
            "content": doc.page_content, "source": doc.metadata.get("source", "Unbekannt"), "score": 1 - score
        })

    # ASSERT: Sicherstellen, dass die zurückgegebene Struktur korrekt ist.
    assert isinstance(formatted_results, list), "Retrieval-Ergebnis muss eine Liste sein."
    if formatted_results:
        assert all("content" in r and "source" in r and "score" in r for r in formatted_results), "Jedes Retrieval-Ergebnis muss 'content', 'source' und 'score' enthalten."

    print_debug(f"{len(formatted_results)} Chunks gefunden.")
    return formatted_results

# --------------------------------------------------------------------
# LLM-Generierung mit Streaming
# --------------------------------------------------------------------
def build_rag_prompt(user_question: str, retrieved_chunks: List[Dict]) -> str:
    # ... (Keine Änderungen in dieser Funktion, bleibt wie in v1.4)
    if not retrieved_chunks:
        context_str = "Es wurden keine relevanten Dokumente im Kontext gefunden."
    else:
        context_parts = []
        for i, ch in enumerate(retrieved_chunks, start=1):
            context_parts.append(
                f"Dokument [{i}] (Quelle: {ch['source']}, Relevanz: {ch['score']:.3f}):\n\"{ch['content']}\""
            )
        context_str = "\n\n".join(context_parts)
    prompt = (f"Du bist ein präziser, hilfreicher Assistent. Deine Aufgabe ist es, die folgende Benutzerfrage ausschließlich "
              f"basierend auf den unten stehenden Kontext-Dokumenten zu beantworten. "
              f"Wenn die Antwort nicht in den Dokumenten enthalten ist, gib klar an: 'Die Information ist in den bereitgestellten Dokumenten nicht enthalten.' "
              f"Antworte auf Deutsch und fasse die relevanten Informationen zusammen, anstatt die Dokumente wörtlich zu zitieren.\n\n"
              f"--- Kontext-Dokumente ---\n{context_str}\n\n"
              f"--- Benutzerfrage ---\n{user_question}\n\n"
              f"--- Deine Antwort ---\n")
    return prompt

def answer_with_rag(question: str, history: list) -> Generator[str, None, None]:
    """Führt RAG durch und generiert eine gestreamte Antwort."""
    print_debug("Starte RAG-Antwort-Generierung.")
    model, processor = get_llm()
    streamer = TextStreamer(processor, skip_prompt=True, skip_special_tokens=True)

    retrieved = retrieve_relevant_chunks(question, top_k=5)
    prompt = build_rag_prompt(question, retrieved)
    print_debug(f"Generierter RAG-Prompt (erste 200 Zeichen): '{prompt[:200].replace(chr(10), ' ')}...'")

    # KORREKTUR: Gemma-3 erwartet eine Liste von Inhalts-Dictionaries.
    messages = [
        {"role": "user", "content": [{"type": "text", "text": prompt}]}
    ]

    # ASSERT: Überprüft die korrekte, verschachtelte Struktur vor dem Aufruf.
    print_debug(f"Nachrichten-Struktur wird für Prozessor vorbereitet: {str(messages)[:200]}...")
    assert isinstance(messages, list) and len(messages) > 0, "Messages muss eine nicht-leere Liste sein."
    assert isinstance(messages[0], dict) and "role" in messages[0] and "content" in messages[0], "Nachricht muss ein Dictionary mit 'role' und 'content' sein."
    assert isinstance(messages[0]["content"], list) and len(messages[0]["content"]) > 0, "Content muss eine nicht-leere Liste sein."
    assert isinstance(messages[0]["content"][0], dict) and "type" in messages[0]["content"][0] and "text" in messages[0]["content"][0], "Content-Block muss ein Dictionary mit 'type' und 'text' sein."
    print_debug("ASSERTIONS für Nachrichten-Struktur erfolgreich bestanden.")

    inputs = processor.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
    ).to(model.device)

    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024, do_sample=True, temperature=0.7, top_p=0.9)

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    print_debug("LLM-Generierungs-Thread gestartet.")

    for new_text in streamer:
        yield new_text
    print_debug("LLM-Generierung abgeschlossen.")

# --------------------------------------------------------------------
# Gradio UI
# --------------------------------------------------------------------
def build_demo() -> gr.Blocks:
    with gr.Blocks(title="Gemma 3 RAG v1.5", theme="soft") as demo:
        gr.Markdown(
            """
            # 🔍 Gemma 3 RAG v1.5 – Robust & Debug-Fähig
            **Eine "State of the Art" RAG-Pipeline mit `google/embeddinggemma-300m` und `google/gemma-3-4b-it`**
            1.  Lade deine Dokumente hoch und klicke auf "Index aktualisieren".
            2.  Stelle deine Fragen im Chatfenster. Die Antworten werden live gestreamt.
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                # ... (Keine UI-Änderungen in dieser Spalte)
                gr.Markdown("### 📁 Dokumenten-Management")
                file_uploader = gr.File(label="Dateien hochladen (.pdf, .txt, .md)", file_count="multiple", type="filepath")
                with gr.Row():
                    index_button = gr.Button("🔄 Index aktualisieren", variant="primary")
                    clear_index_button = gr.Button("🧹 Index leeren")
                index_status = gr.Markdown("Index ist leer.")
                index_button.click(fn=index_files, inputs=file_uploader, outputs=index_status)
                clear_index_button.click(fn=clear_index, inputs=None, outputs=index_status)

            with gr.Column(scale=2):
                # ... (Keine UI-Änderungen in dieser Spalte)
                gr.Markdown("### 💬 Chat über deine Dokumente")
                chatbot = gr.Chatbot(label="Gemma-3 Chat", type="messages", show_copy_button=True, height=600)
                with gr.Row():
                    msg_textbox = gr.Textbox(label="Deine Frage", placeholder="Stelle eine Frage zu den Dokumenten...", scale=4, autofocus=True)
                    send_btn = gr.Button("Senden", variant="primary", scale=1)

        def chat_interface(message: str, history: list):
            if not message or not message.strip():
                return history
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": ""})
            for token in answer_with_rag(message, history):
                history[-1]["content"] += token
                yield history

        msg_textbox.submit(fn=chat_interface, inputs=[msg_textbox, chatbot], outputs=chatbot).then(fn=lambda: gr.update(value=""), outputs=msg_textbox)
        send_btn.click(fn=chat_interface, inputs=[msg_textbox, chatbot], outputs=chatbot).then(fn=lambda: gr.update(value=""), outputs=msg_textbox)

    return demo

if __name__ == "__main__":
    # Stelle sicher, dass die Modelle beim Start geladen werden,
    # um eine Verzögerung bei der ersten Anfrage zu vermeiden.
    get_embedding_function()
    get_llm()

    app_demo = build_demo()
    app_demo.launch()