rag_pipeline.py

import os
import pickle
import requests
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from dotenv import load_dotenv
from openai import OpenAI
from groq import Groq

# === API Keys laden ===
load_dotenv()
openai_key = os.getenv("OPENAI_API_KEY")
groq_key = os.getenv("GROQ_API_KEY")

openai_client = OpenAI(api_key=openai_key) if openai_key else None
groq_client = Groq(api_key=groq_key) if groq_key else None

# === Modell laden ===
print("🧠 Lade SentenceTransformer...")
model = SentenceTransformer("Sahajtomar/German-semantic")

# === Google Drive Direktlinks
url_index = "https://drive.google.com/uc?export=download&id=1QBg4vjitJ2xHEyp3Ae8TWJHwEHjbwgOO"
url_chunks = "https://drive.google.com/uc?export=download&id=1nsrAm_ozsK4GlmMui9yqZBjmgUfqU2qa"

local_index = "faiss_index.index"
local_chunks = "chunks_mapping.pkl"

# === Datei-Download bei Bedarf
def download_if_missing(url, path):
    if not os.path.exists(path):
        print(f"⬇️ Lade {path} von Google Drive...")
        r = requests.get(url)
        if r.status_code == 200:
            with open(path, "wb") as f:
                f.write(r.content)
            print(f"✅ Heruntergeladen: {path}")
        else:
            raise Exception(f"❌ Fehler beim Herunterladen von {path}")

download_if_missing(url_index, local_index)
download_if_missing(url_chunks, local_chunks)

# === FAISS laden
print("📂 Lade FAISS & Chunks...")
with open(local_chunks, "rb") as f:
    token_split_texts = pickle.load(f)
print(f"✅ {len(token_split_texts)} Chunks geladen.")

chunk_embeddings = model.encode(token_split_texts, convert_to_numpy=True)
d = chunk_embeddings.shape[1]
index = faiss.IndexFlatL2(d)
index.add(chunk_embeddings)
print(f"✅ FAISS Index mit {index.ntotal} Einträgen.")

# === Ähnliche Chunks abrufen
def retrieve(query, k=5):
    query_embedding = model.encode([query], convert_to_numpy=True)
    distances, indices = index.search(query_embedding, k)
    safe_indices = [i for i in indices[0] if i < len(token_split_texts)]
    return [token_split_texts[i] for i in safe_indices]

# === Prompt zusammenbauen
def build_prompt(query, texts):
    context = "\n\n".join(texts)
    return f"""Beantworte die folgende Frage basierend auf dem Kontext.

Kontext:
{context}

Frage:
{query}
"""

# === Anfrage an OpenAI
def ask_openai(prompt):
    if not openai_client:
        return "❌ Kein OpenAI API Key gefunden"
    res = openai_client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "Du bist ein hilfsbereiter Catan-Regel-Experte."},
            {"role": "user", "content": prompt}
        ]
    )
    return res.choices[0].message.content.strip()

# === Anfrage an Groq
def ask_groq(prompt):
    if not groq_client:
        return "❌ Kein Groq API Key gefunden"
    res = groq_client.chat.completions.create(
        model="llama3-70b-8192",
        messages=[
            {"role": "system", "content": "Du bist ein hilfsbereiter Catan-Regel-Experte."},
            {"role": "user", "content": prompt}
        ]
    )
    return res.choices[0].message.content.strip()

# === Hauptfunktion für Gradio
def run_qa_pipeline(query, k=5):
    try:
        retrieved = retrieve(query, k)
        if not retrieved:
            return "⚠️ Keine relevanten Textstellen gefunden."

        prompt = build_prompt(query, retrieved)
        print("📨 Prompt gesendet...")

        if openai_client:
            answer = ask_openai(prompt)
        elif groq_client:
            answer = ask_groq(prompt)
        else:
            return "⚠️ Kein LLM API-Key vorhanden. Bitte OPENAI_API_KEY oder GROQ_API_KEY hinterlegen."

        return f"📌 Frage: {query}\n\n📖 Antwort:\n{answer}"

    except Exception as e:
        return f"❌ Fehler beim Verarbeiten der Anfrage:\n{str(e)}"