import os
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.llms import HuggingFaceHub

# Función para cargar documentos de texto
def load_docs(directory):
    documents = []
    for filename in os.listdir(directory):
        if filename.endswith('.txt'):
            filepath = os.path.join(directory, filename)
            loader = TextLoader(filepath)
            documents.extend(loader.load())
    return documents

# Función para dividir documentos en chunks
def split_docs(documents):
    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    return text_splitter.split_documents(documents)

# Función para crear el índice de vectores
def create_vector_index(chunks):
    embeddings = HuggingFaceEmbeddings()
    vector_store = FAISS.from_documents(chunks, embeddings)
    return vector_store

# Función principal para configurar y ejecutar el asistente
def run_assistant():
    # Cargar documentos
    docs = load_docs('/home/www_ssl/assistant/data/txt')
    
    # Dividir documentos en chunks
    chunks = split_docs(docs)
    
    # Crear índice de vectores
    vector_store = create_vector_index(chunks)
    
    # Configurar el modelo de lenguaje (asegúrate de tener una clave de API de Hugging Face)
    llm = HuggingFaceHub(repo_id="google/flan-t5-xxl", model_kwargs={"temperature":0.5, "max_length":512})
    
    # Configurar el chain de pregunta-respuesta
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever()
    )
    
    # Loop principal de preguntas y respuestas
    while True:
        query = input("Haz una pregunta (o escribe 'salir' para terminar): ")
        if query.lower() == 'salir':
            break
        
        result = qa_chain.run(query)
        print(f"Respuesta: {result}\n")

if __name__ == "__main__":
    run_assistant()