# query_files.py (modifiziert für Ollama)

from sentence_transformers import SentenceTransformer
import chromadb
import ollama # Neu importieren

# --- Gleicher Setup-Code wie zuvor ---
INDEX_DIR = './index/'
model = SentenceTransformer('all-MiniLM-L6-v2')
client = chromadb.PersistentClient(path=INDEX_DIR)
collection = client.get_or_create_collection(name="robot_docs")

def query_docs(question):
    """Frage stellen und relevante Chunks abrufen."""
    # Embedding der Frage
    query_embedding = model.encode(question).tolist()
    # Suche Top-5 relevante Chunks
    results = collection.query(query_embeddings=[query_embedding], n_results=5)
    
    # Wir geben jetzt nur die Text-Chunks zurück
    return results['documents'][0]

def generate_answer_with_ollama(question, context_chunks):
    """Generiert eine Antwort mit Llama 3.1 basierend auf dem Kontext."""
    
    # Die gefundenen Chunks zu einem einzigen Text zusammenfügen
    context = "\n\n---\n\n".join(context_chunks)
    
    # Ein "Prompt Template" erstellen. Dies ist die Anweisung an das LLM.
    prompt = f"""
    Du bist ein hilfreicher Assistent. Beantworte die folgende Frage ausschließlich basierend auf dem bereitgestellten Kontext.
    Wenn die Antwort im Kontext nicht zu finden ist, antworte mit "Ich konnte die Antwort in den Dokumenten nicht finden.".
    Antworte auf Deutsch.

    Kontext:
    {context}

    Frage: {question}

    Antwort:
    """
    
    print("\n--- Sende Anfrage an Llama 3.1... ---")
    
    # Anfrage an Ollama senden
    try:
        response = ollama.chat(
            model='llama3.1',
            messages=[{'role': 'user', 'content': prompt}]
        )
        return response['message']['content']
    except Exception as e:
        return f"Fehler bei der Kommunikation mit Ollama: {e}"


if __name__ == "__main__":
    # 1. Frage vom Benutzer erhalten
    question = input("Stelle deine Frage (z.B. 'Fehler bei Robotergruppe A'): ")
    
    # 2. Relevante Dokumente aus der Datenbank abrufen (Retrieval)
    retrieved_chunks = query_docs(question)
    
    # Optional: Gefundene Chunks anzeigen
    print("\n--- Relevante Informationen gefunden ---")
    for i, chunk in enumerate(retrieved_chunks):
        print(f"Chunk {i+1}: {chunk[:150]}...")
    
    # 3. Antwort mit Ollama und den Chunks generieren (Generation)
    final_answer = generate_answer_with_ollama(question, retrieved_chunks)
    
    # 4. Finale Antwort ausgeben
    print("\n================ ANTWORT VON LLAMA 3.1 ================")
    print(final_answer)
    print("=======================================================")