from sentence_transformers import SentenceTransformer
from chromadb import PersistentClient
from transformers import pipeline

# Conectar con la base vectorial
client = PersistentClient(path="./vector_db")
collection = client.get_collection("codehelper_csharp")

# Cargar modelo de embeddings
model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

# Cargar traductor ingls  espaol
translator = pipeline("translation_en_to_es", model="Helsinki-NLP/opus-mt-en-es")

# Pregunta del usuario
query = input("Escrib tu pregunta sobre C# o .NET: ")

# Generar embedding y buscar
embedding = model.encode(query).tolist()
results = collection.query(query_embeddings=[embedding], n_results=3)

# Mostrar resultados con opcin de ver original
print("\n Resultados ms relevantes (traducidos al espaol):")
for i, doc in enumerate(results["documents"][0]):
    print(f"\n Fragmento {i+1}:")

    # Traduccin al espaol
    translated = translator(doc[:500])[0]["translation_text"]
    print(translated)

    # Preguntar si quiere ver el original
    ver_original = input("Quers ver el texto original en ingls? (s/n): ").strip().lower()
    if ver_original == "s":
        print("\n Original:\n" + doc[:500])