embed.py

import os
from datetime import datetime
from werkzeug.utils import secure_filename
from langchain_community.document_loaders import UnstructuredPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from get_vector_db import get_vector_db

# Usa la misma ubicación temporal que en app.py
TEMP_FOLDER = os.getenv('TEMP_FOLDER', '/tmp/_temp')

# Función para verificar si el archivo está permitido (solo archivos PDF)
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in {'pdf'}

# Función para guardar el archivo subido en la carpeta temporal
def save_file(file):
    # Guarda el archivo subido con un nombre seguro y devuelve la ruta del archivo
    ct = datetime.now()
    ts = ct.timestamp()
    filename = str(ts) + "_" + secure_filename(file.filename)
    file_path = os.path.join(TEMP_FOLDER, filename)
    file.save(file_path)

    return file_path

# Función para cargar y dividir los datos del archivo PDF
def load_and_split_data(file_path):
    # Carga el archivo PDF y divide los datos en fragmentos
    loader = UnstructuredPDFLoader(file_path=file_path)
    data = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=7500, chunk_overlap=100)
    chunks = text_splitter.split_documents(data)

    return chunks

# Función principal para manejar el proceso de incrustación
def embed(file):
    # Verifica si el archivo es válido, lo guarda, carga y divide los datos, los añade a la base de datos y elimina el archivo temporal
    if file.filename != '' and file and allowed_file(file.filename):
        file_path = save_file(file)
        chunks = load_and_split_data(file_path)
        db = get_vector_db()
        db.add_documents(chunks)
        db.persist()
        os.remove(file_path)

        return True

    return False