Update rag_pipeline.py

rag_pipeline.py

@@ -4,6 +4,18 @@ 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...")
@@ -13,54 +25,98 @@ model = SentenceTransformer("Sahajtomar/German-semantic")
 url_index = "https://drive.google.com/uc?export=download&id=1QBg4vjitJ2xHEyp3Ae8TWJHwEHjbwgOO"
 url_chunks = "https://drive.google.com/uc?export=download&id=1nsrAm_ozsK4GlmMui9yqZBjmgUfqU2qa"
 
-# === Lokale Dateipfade
 local_index = "faiss_index.index"
 local_chunks = "chunks_mapping.pkl"
 
-# === Datei-Download bei Bedarf
-def download_if_missing(url, local_path):
-    if not os.path.exists(local_path):
-        print(f"⬇️ Lade {local_path} von Google Drive...")
+# === 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(local_path, "wb") as f:
+            with open(path, "wb") as f:
                 f.write(r.content)
-            print(f"✅ Heruntergeladen: {local_path}")
+            print(f"✅ Heruntergeladen: {path}")
         else:
-            raise Exception(f"❌ Download fehlgeschlagen für {local_path}")
+            raise Exception(f"❌ Fehler beim Herunterladen von {path}")
 
 download_if_missing(url_index, local_index)
 download_if_missing(url_chunks, local_chunks)
 
-# === FAISS & Chunks laden
-print("📂 Lade FAISS Index und Text-Chunks...")
-index = faiss.read_index(local_index)
-
+# === 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.")
 
-print(f"✅ Geladene Chunks: {len(token_split_texts)}")
+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.")
 
-# === Embedding nur auf den ersten 10 Chunks testen
-print("⚙️ Starte Embedding-Berechnung auf 10 Chunks...")
-test_chunks = token_split_texts[:10]
-chunk_embeddings = model.encode(test_chunks, convert_to_numpy=True)
-print("✅ Embeddings kodiert")
-
-# === Abruffunktion
+# === Ähnliche Chunks abrufen
 def retrieve(query, k=5):
     query_embedding = model.encode([query], convert_to_numpy=True)
     distances, indices = index.search(query_embedding, k)
-    retrieved_texts = [test_chunks[i] for i in indices[0]]
-    return retrieved_texts
+    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 Builder
+# === Prompt zusammenbauen
 def build_prompt(query, texts):
     context = "\n\n".join(texts)
-    return f"Beantworte die folgende Frage basierend auf dem Kontext:\n\nKontext:\n{context}\n\nFrage:\n{query}"
+    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):
-    retrieved = retrieve(query, k)
-    prompt = build_prompt(query, retrieved)
-    return f"🔍 Kontext gefunden:\n\n{prompt}"
+    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: {str(e)}"    return f"🔍 Kontext gefunden:\n\n{prompt}"