Faiss_VectorStore_Class.py

import os
import mimetypes
import faiss
import hashlib
import google.generativeai as genai
from uuid import uuid4
from pdf2image import convert_from_path
from google.cloud import vision
from google.oauth2 import service_account
from moviepy import VideoFileClip
from langchain_core.documents import Document
from langchain_community.vectorstores import FAISS
from langchain_core.document_loaders.blob_loaders import Blob
from langchain_community.docstore.in_memory import InMemoryDocstore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader, Docx2txtLoader
from langchain_community.document_loaders.parsers import OpenAIWhisperParser
from langchain_openai import OpenAIEmbeddings

class VectorStoreManager:
    """
    Clase para gestionar un vector store utilizando FAISS y otras herramientas avanzadas.
    Permite procesar archivos de texto, PDFs, imágenes, audios y videos.
    """
    
    def __init__(self, openai_api_key, gemini_api_key, google_credentials_path, faiss_index_path="faiss_index"):
        """
        Inicializa el gestor con las credenciales necesarias y configura los módulos requeridos.
        
        Args:
            openai_api_key (str): API key de OpenAI.
            gemini_api_key (str): API key de Gemini.
            google_credentials_path (str): Ruta al archivo de credenciales de Google.
            faiss_index_path (str): Ruta donde se guardará el índice FAISS.
        """
        try:
            self.api_key = openai_api_key
            self.GOOGLE_CREDENTIALS_PATH = google_credentials_path
            self.GENAI_API_KEY = gemini_api_key
            genai.configure(api_key=self.GENAI_API_KEY)

            # Configuración de embeddings
            self.embeddings = OpenAIEmbeddings(
                api_key=self.api_key,
                model="text-embedding-3-large",
                max_retries=1,
            )

            # Configuración de FAISS
            self.index = faiss.IndexFlatL2(len(self.embeddings.embed_query("test")))
            self.vector_store = FAISS(
                embedding_function=self.embeddings,
                index=self.index,
                docstore=InMemoryDocstore(),
                index_to_docstore_id={},
            )
            self.faiss_index_path = faiss_index_path

            # Configuración del splitter
            self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)
        except Exception as e:
            raise RuntimeError(f"Error al inicializar VectorStoreManager: {e}")

    def generate_gemini_response(self, text):
        """
        Corrige texto OCR utilizando un modelo generativo (Gemini).
        
        Args:
            text (str): Texto a corregir.
        
        Returns:
            str: Texto corregido.
        """
        try:
            model = genai.GenerativeModel("gemini-1.5-flash")
            prompt = f"""
            Corrección profesional de texto extraído por OCR con los siguientes criterios:
            1. Corregir errores de reconocimiento óptico de caracteres
            2. Restaurar la puntuación y la estructura gramatical correcta
            3. Conservar el formato original del documento
            4. Mantener intacto el contenido informativo del texto original
            5. Asegurar la máxima precisión en la transcripción
            Características de salida:
            - Texto completamente legible
            - Sin caracteres extraviados o mal reconocidos
            - Gramática y ortografía corregidas
            - Formato apropiado para archivo .txt
            Texto OCR:
            {text}
            Texto corregido:
            """
            response = model.generate_content(
                prompt,
                generation_config=genai.GenerationConfig(max_output_tokens=4096, temperature=0.2, top_p=0.75)
            )
            return response.text
        except Exception as e:
            raise RuntimeError(f"Error al generar respuesta de Gemini: {e}")

    def _calculate_file_hash(self, file_path):
        """
        Calcula el hash SHA256 de un archivo.
        
        Args:
            file_path (str): Ruta del archivo.
        
        Returns:
            str: Hash SHA256 del archivo.
        """
        try:
            hasher = hashlib.sha256()
            with open(file_path, 'rb') as f:
                while chunk := f.read(8192):
                    hasher.update(chunk)
            return hasher.hexdigest()
        except Exception as e:
            raise RuntimeError(f"Error al calcular hash del archivo {file_path}: {e}")

    def save_vs_to_local(self):
        """
        Guarda el vector store localmente, incluyendo el índice y el docstore.
        """
        try:
            self.vector_store.save_local(self.faiss_index_path)
            print(f"Vector Store guardado en {self.faiss_index_path}.")
        except Exception as e:
            raise RuntimeError(f"Error al guardar el vector store: {e}")

    def load_vs_from_local(self):
        """
        Carga el vector store desde el disco local.
        """
        try:
            self.vector_store.load_local(self.faiss_index_path, self.embeddings, allow_dangerous_deserialization=True)
            print(f"Vector Store cargado desde {self.faiss_index_path}.")
        except Exception as e:
            raise RuntimeError(f"Error al cargar el vector store: {e}")

    def similarity_search(self, query, k=1):
        """
        Realiza una búsqueda de similitud en el vector store.
        
        Args:
            query (str): Consulta.
            k (int): Número de resultados a retornar.
        
        Returns:
            List[Document]: Resultados de la búsqueda.
        """
        try:
            return self.vector_store.similarity_search(query=query, k=k)
        except Exception as e:
            raise RuntimeError(f"Error en la búsqueda de similitud: {e}")

    def ocr_image(self, image_path):
        """
        Realiza OCR en una imagen utilizando Google Vision API.
        
        Args:
            image_path (str): Ruta de la imagen.
        
        Returns:
            str: Texto extraído y corregido.
        """
        try:
            credentials = service_account.Credentials.from_service_account_file(self.GOOGLE_CREDENTIALS_PATH)
            client = vision.ImageAnnotatorClient(credentials=credentials)

            with open(image_path, 'rb') as image_file:
                content = image_file.read()

            image = vision.Image(content=content)
            response = client.text_detection(image=image)

            if response.error.message:
                raise Exception(f"Error en Vision API: {response.error.message}")
            
            return self.generate_gemini_response(response.full_text_annotation.text)
        except Exception as e:
            raise RuntimeError(f"Error al realizar OCR en la imagen {image_path}: {e}")

    def process_pdf(self, file_path):
        """
        Procesa un PDF dividiendo páginas en imágenes y aplicando OCR.
        
        Args:
            file_path (str): Ruta del archivo PDF.
        
        Returns:
            str: Texto extraído de todas las páginas.
        """
        try:
            pages = convert_from_path(file_path)
            extracted_text = ""
            for page in pages:
                temp_image_path = "temp_page.jpg"
                page.save(temp_image_path, "JPEG")
                extracted_text += self.ocr_image(temp_image_path) + "\n"
                os.remove(temp_image_path)
            return extracted_text
        except Exception as e:
            raise RuntimeError(f"Error al procesar el PDF {file_path}: {e}")

    def process_audio(self, file_path):
        """
        Transcribe un archivo de audio usando OpenAI Whisper.
        
        Args:
            file_path (str): Ruta del archivo de audio.
        
        Returns:
            str: Transcripción del audio.
        """
        try:
            audio_blob = Blob.from_path(file_path)
            parser = OpenAIWhisperParser(api_key=self.api_key)
            documents = parser.parse(audio_blob)
            return "\n".join(doc.page_content for doc in documents)
        except Exception as e:
            raise RuntimeError(f"Error al procesar el audio {file_path}: {e}")

    def process_video(self, file_path):
        """
        Extrae el audio de un video y lo transcribe.
        
        Args:
            file_path (str): Ruta del archivo de video.
        
        Returns:
            str: Transcripción del audio extraído.
        """
        try:
            video = VideoFileClip(file_path)
            temp_audio_path = "temp_audio.wav"
            video.audio.write_audiofile(temp_audio_path)
            transcript = self.process_audio(temp_audio_path)
            os.remove(temp_audio_path)
            return transcript
        except Exception as e:
            raise RuntimeError(f"Error al procesar el video {file_path}: {e}")

    def load_directory(self, directory_path):
        """
        Carga y procesa todos los archivos de un directorio de manera recursiva.
        
        Args:
            directory_path (str): Ruta del directorio.
        
        Raises:
            ValueError: Si la ruta proporcionada no es un directorio válido.
        """
        try:
            if not os.path.isdir(directory_path):
                raise ValueError(f"{directory_path} no es un directorio válido.")

            for root, _, files in os.walk(directory_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    try:
                        print(f"Procesando archivo: {file_path}")
                        self.load_file(file_path)
                    except Exception as e:
                        print(f"Error procesando el archivo {file_path}: {e}")

            print(f"Todos los archivos en el directorio '{directory_path}' han sido procesados.")
        except Exception as e:
            raise RuntimeError(f"Error al procesar el directorio {directory_path}: {e}")

    def load_file(self, file_path):
        """
        Carga y procesa un archivo, dividiéndolo y agregándolo al vector store con metadatos.
        
        Args:
            file_path (str): Ruta del archivo.
        """
        try:
            mime_type, _ = mimetypes.guess_type(file_path)
            if mime_type is None:
                raise ValueError(f"No se pudo determinar el tipo de archivo para: {file_path}")

            file_hash = self._calculate_file_hash(file_path)
            file_name = os.path.basename(file_path)

            # Verificar duplicados
            for doc_id, document in self.vector_store.docstore._dict.items():
                if document.metadata.get("hash") == file_hash:
                    print(f"Archivo '{file_name}' ya procesado. Omitiendo...")
                    return

            # Procesar según el tipo de archivo
            documents = []
            if mime_type.startswith("text"):
                loader = TextLoader(file_path, autodetect_encoding=True)
                raw_documents = loader.load()
                documents = [
                    Document(page_content=doc.page_content, metadata={"source": file_name, "hash": file_hash})
                    for doc in raw_documents
                ]
            elif mime_type in ["application/msword",
                           "application/vnd.openxmlformats-officedocument.wordprocessingml.document"]:
            # Procesar archivos .doc o .docx
                loader = Docx2txtLoader(file_path)
                raw_documents = loader.load()
                documents = [
                    Document(page_content=doc.page_content, metadata={"source": file_name, "hash": file_hash})
                    for doc in raw_documents
                ]
            elif mime_type == "application/pdf":
                text = self.process_pdf(file_path)
                documents = [Document(page_content=text, metadata={"source": file_name, "hash": file_hash})]
            elif mime_type.startswith("image"):
                text = self.ocr_image(file_path)
                documents = [Document(page_content=text, metadata={"source": file_name, "hash": file_hash})]
            elif mime_type.startswith("audio"):
                text = self.process_audio(file_path)
                documents = [Document(page_content=text, metadata={"source": file_name, "hash": file_hash})]
            elif mime_type.startswith("video"):
                text = self.process_video(file_path)
                documents = [Document(page_content=text, metadata={"source": file_name, "hash": file_hash})]
            else:
                raise ValueError(f"Tipo de archivo no soportado: {mime_type}")

            # Dividir y agregar al vector store
            for doc in documents:
                chunks = self.text_splitter.split_documents([doc])
                self.vector_store.add_documents(chunks)
            print(f"Archivo '{file_name}' procesado y agregado al vector store.")
        except Exception as e:
            raise RuntimeError(f"Error al cargar el archivo {file_path}: {e}")