Upload 7 files

Dockerfile

@@ -0,0 +1,15 @@
+# Usa una imagen base de Python
+FROM python:3.11
+# Establece el directorio de trabajo
+WORKDIR /code
+
+# Copia los archivos necesarios al contenedor
+COPY ./requirements.txt /code/requirements.txt
+RUN pip install --no-cache-dir -r /code/requirements.txt
+
+COPY . .
+
+RUN chmod -R 777 /code
+
+# Comando para ejecutar la aplicación
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,178 @@
+from fastapi import FastAPI, HTTPException, Depends, File, UploadFile
+from typing import List
+from fastapi.responses import FileResponse
+import urllib
+from embeddings import EmbeddingManager
+from model import (
+    AddFilesRequest,
+    CreateVectorStoreRequest,
+    DeleteVectorStoreRequest,
+    DownloadVectorStoreRequest,
+    ListSourcesRequest,
+    SaveTempRequest,
+    SearchSimilarityRequest,
+)
+from vector_db import VectorStoreManager
+import os
+import shutil
+from starlette.responses import RedirectResponse
+
+app = FastAPI()
+
+
+@app.get("/", include_in_schema=False)
+async def redirect_to_docs():
+    return RedirectResponse(url="/docs")
+
+
+# Crear una sola instancia de EmbeddingManager
+embedding_manager = EmbeddingManager()
+embeddings = embedding_manager.get_embeddings
+path_docs = "docs"  # Directorio temporal para almacenar los archivos subidos
+path_db = "database"  # Directorio para almacenar el vectorstore
+
+
+@app.post("/vectorstore/create", tags=["VectorStore"])
+async def create_vectorstore(
+    create_request: CreateVectorStoreRequest = Depends(),  # Usar el modelo como dependencia
+    files: List[UploadFile] = File(...),
+):
+    """Create a vectorstore from the uploaded documents."""
+    try:
+        if os.path.exists(path_docs):
+            shutil.rmtree(path_docs)
+        os.makedirs(path_docs)
+        for file in files:
+            file_path = os.path.join(path_docs, file.filename)
+            with open(file_path, "wb") as f:
+                f.write(await file.read())
+        manager = VectorStoreManager(
+            path=path_docs, name=create_request.name, embeddings=embeddings
+        )
+        if manager.create_vectorstore():
+            shutil.rmtree(path_docs)
+            return {"message": "Vectorstore created successfully."}
+        shutil.rmtree(path_docs)
+        return {"message": "Failed to create vectorstore."}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/vectorstore/search", tags=["Similarity Search"])
+async def search_similarity(search_request: SearchSimilarityRequest = Depends()):
+    """Search for similar documents in the vectorstore."""
+    try:
+        manager = VectorStoreManager(
+            path=path_db,
+            name=search_request.name_database,
+            embeddings=embeddings,
+        )
+        search_request.query = str(urllib.parse.unquote(search_request.query))
+        result = manager.search_similarity(
+            query=search_request.query, fuente=search_request.fuente
+        )
+        return {"results": result}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/vectorstore/sources", tags=["Sources"])
+async def list_sources(list_request: ListSourcesRequest = Depends()):
+    try:
+        manager = VectorStoreManager(
+            path=path_db, name=list_request.nombre_db_vectorial, embeddings=embeddings
+        )
+        sources = manager.list_sources()
+        return {"sources": sources}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/vectorstore/save_temp", tags=["Save Temp"])
+async def save_text_to_file_temp(save_temp: SaveTempRequest = Depends()):
+    """Descripción: Guarda en un archivo temporal el texto de una fuente específica."""
+    try:
+        manager = VectorStoreManager(
+            path=path_db, name=save_temp.nombre_db_vectorial, embeddings=embeddings
+        )
+        saved = manager.save_text_to_file_temp(source=save_temp.fuente)
+        if saved:
+            return {"message": "Text saved to file successfully."}
+        else:
+            return {"message": "No text found to save."}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/vectorstore/add_files", tags=["Add Files"])
+async def add_files_vectorstore(
+    add_files_request: AddFilesRequest = Depends(), files: List[UploadFile] = File(...)
+):
+    try:
+        if os.path.exists(path_docs):
+            shutil.rmtree(path_docs)
+        os.makedirs(path_docs)
+
+        for file in files:
+            file_path = os.path.join(path_docs, file.filename)
+            with open(file_path, "wb") as f:
+                f.write(await file.read())
+        manager = VectorStoreManager(
+            path=path_docs,
+            name=add_files_request.nombre_db_vectorial,
+            embeddings=embeddings,
+        )
+        if manager.add_files_vectorstore():
+            shutil.rmtree(path_docs)
+            return {"message": "Files added to vectorstore successfully."}
+        shutil.rmtree(path_docs)
+        return {"message": "Failed to add files to vectorstore."}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.delete("/vectorstore/delete", tags=["Delete VectorStore"])
+async def delete_vectorstore(delete_request: DeleteVectorStoreRequest = Depends()):
+    """Delete the vectorstore and its data."""
+    try:
+        manager = VectorStoreManager(
+            path=path_db, name=delete_request.nombre_db_vectorial, embeddings=embeddings
+        )
+        if manager.delete_vectorstore():
+            return {"message": "Vectorstore deleted successfully."}
+        return {"message": "Failed to delete vectorstore."}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/vectorstore/download", tags=["Download VectorStore"])
+async def download_vectorstore(
+    download_request: DownloadVectorStoreRequest = Depends(),
+):
+    try:
+        manager = VectorStoreManager(
+            path=path_db,
+            name=download_request.nombre_db_vectorial,
+            embeddings=embeddings,
+        )
+        zip_path = manager.download_vectorstore()
+        return FileResponse(zip_path, filename="vectorstore.zip")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    import os
+
+    try:
+        # crear todas las carpetas necesarias si no existen
+        carpetas = [path_docs, path_db, "temp"]
+        for carpeta in carpetas:
+            if not os.path.exists(carpeta):
+                os.makedirs(carpeta)
+        os.system("uvicorn app:app --port 7860 --host 0.0.0.0")
+    except KeyboardInterrupt:
+        print("Server stopped.")
+    except Exception as e:
+        print(e)
+        print("Failed to start server.")

document_processor.py

@@ -0,0 +1,36 @@
+import os
+from langchain_community.document_loaders import (
+    PyMuPDFLoader,
+    TextLoader,
+    Docx2txtLoader,
+    DirectoryLoader,
+)
+
+class DocumentProcessor:
+    def __init__(self, path: str):
+        self.path = path
+
+    def files_to_texts(self) -> list:
+        loaders_config = {
+            "*.pdf": PyMuPDFLoader,
+            "*.txt": (TextLoader, {"encoding": "utf-8"}),
+            "*.docx": Docx2txtLoader,
+            "*.doc": Docx2txtLoader,
+        }
+
+        loaders = [
+            DirectoryLoader(
+                path=self.path,
+                glob=glob,
+                loader_cls=loader if isinstance(loader, type) else loader[0],
+                loader_kwargs=loader[1] if isinstance(loader, tuple) else None,
+            )
+            for glob, loader in loaders_config.items()
+            if any(fname.endswith(glob[1:]) for fname in os.listdir(self.path))
+        ]
+
+        documents = []
+        for loader in loaders:
+            documents.extend(loader.load())
+
+        return documents

embeddings.py

@@ -0,0 +1,30 @@
+from langchain_huggingface import HuggingFaceEmbeddings
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class EmbeddingManager:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super(EmbeddingManager, cls).__new__(cls, *args, **kwargs)
+            cls._instance.__initialized = False
+        return cls._instance
+
+    def __init__(self):
+        if self.__initialized:
+            return
+        self.__initialized = True
+        self.__embeddings = HuggingFaceEmbeddings(
+            model_name="jinaai/jina-embeddings-v2-base-es",
+            encode_kwargs={"normalize_embeddings": True},
+            model_kwargs={
+                "device": device,
+            },
+        )
+
+    @property
+    def get_embeddings(self):
+        return self.__embeddings

model.py

@@ -0,0 +1,33 @@
+from typing import Optional
+from pydantic import BaseModel
+
+
+class CreateVectorStoreRequest(BaseModel):
+    name: str
+
+
+class SearchSimilarityRequest(BaseModel):
+    name_database: str
+    query: str
+    fuente: Optional[str] = None
+
+
+class ListSourcesRequest(BaseModel):
+    nombre_db_vectorial: str
+
+
+class SaveTempRequest(BaseModel):
+    nombre_db_vectorial: str
+    fuente: str
+
+
+class AddFilesRequest(BaseModel):
+    nombre_db_vectorial: str
+
+
+class DeleteVectorStoreRequest(BaseModel):
+    nombre_db_vectorial: str
+
+
+class DownloadVectorStoreRequest(BaseModel):
+    nombre_db_vectorial: str

requirements.txt

@@ -0,0 +1,18 @@
+requests 
+fastapi 
+langchain 
+langchain-core 
+langchain-openai 
+langchain-community 
+langchain-huggingface
+faiss-cpu 
+duckduckgo-search
+uvicorn
+einops
+python-multipart
+docx2txt
+aiofiles
+pdfplumber 
+python-docx
+sentence-transformers
+pymupdf

vector_db.py

@@ -0,0 +1,160 @@
+import os
+from typing import Optional, List
+import shutil
+from zipfile import ZipFile
+from langchain_community.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.embeddings import Embeddings
+
+from document_processor import DocumentProcessor
+
+
+class VectorStoreManager:
+    def __init__(self, path: str, name: str, embeddings: Embeddings):
+        """
+        Descripción: Clase para gestionar el vectorstore, incluyendo la creación, eliminación y búsqueda de
+        documentos similares.
+
+        Parámetros:
+        - path: str - ruta del directorio que contiene los documentos (usualmente es "database" que es el directorio
+        donde se almacenan las bases de datos).
+        - name: str - nombre del vectorstore. (usualmente, es el nombre de la base de datos que contiene los documentos)
+        - embeddings: Embeddings - modelo de embeddings para el vectorstore.
+
+        """
+        self.path = path
+        self.name = name
+        self.embeddings = embeddings
+        self.vectorstore = None
+
+    def create_vectorstore(self) -> bool:
+        documents = DocumentProcessor(self.path).files_to_texts()
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000, chunk_overlap=200, length_function=len
+        )
+        texts = text_splitter.split_documents(documents)
+        self.vectorstore = FAISS.from_documents(
+            documents=texts, embedding=self.embeddings
+        )
+        base_de_datos_dir = os.path.join("database", self.name)
+        self.vectorstore.save_local(folder_path=base_de_datos_dir)
+        return True
+
+    def delete_vectorstore(self) -> bool:
+        try:
+            shutil.rmtree(f"database/{self.name}")
+        except FileNotFoundError:
+            return False
+        return True
+
+    def search_similarity(self, query: str, fuente: Optional[str] = None) -> str:
+        """
+        Modo de uso:
+        debe ingresar la query y la fuente (opcional) para buscar documentos similares en el vectorstore.
+
+        Nota: debe estar definido el vectorstore para poder realizar la búsqueda.
+
+        Parámetros:
+        query: str - texto de la query.
+        fuente: str - fuente de los documentos a buscar.
+
+        Retorna:
+        str - documentos similares.
+        """
+        if not self.vectorstore:
+            self.vectorstore = self.load_vectorstore()
+
+        if fuente:
+            filtro = {"source": fuente}
+            retriever = self.vectorstore.similarity_search(
+                query=query, k=5, filter=filtro
+            )
+        else:
+            retriever = self.vectorstore.similarity_search(query=query, k=5)
+        busqueda = [
+            {
+                "content": doc.page_content,
+                "title": doc.metadata.get("title", None),
+                "source": doc.metadata.get("source", None),
+            }
+            for doc in retriever
+        ]
+
+        return str(busqueda)
+
+    def list_sources(self) -> List[str]:
+        if not self.vectorstore:
+            self.vectorstore = self.load_vectorstore()
+
+        docstore_dict = self.vectorstore.docstore._dict
+        source_metadata = {}
+        for doc_id, document in docstore_dict.items():
+            source = document.metadata.get("source", None)
+            source_metadata[doc_id] = source
+
+        return list(set(source_metadata.values()))
+
+    def extract_texts_by_source(self, source: str) -> List[str]:
+        if not self.vectorstore:
+            self.vectorstore = self.load_vectorstore()
+
+        docstore_dict = self.vectorstore.docstore._dict
+        texts = []
+        for document in docstore_dict.values():
+            source_doc = document.metadata.get("source", None)
+            if source_doc == source:
+                texts.append(document.page_content)
+        return texts
+
+    def save_text_to_file_temp(self, source: str) -> bool:
+        texts = self.extract_texts_by_source(source)
+        carpeta = "temp"
+        target_source_safe = source.replace("\\", "_").replace("/", "_")
+        file_path = os.path.join(carpeta, target_source_safe + ".txt")
+
+        try:
+            if os.path.exists(carpeta):
+                shutil.rmtree(carpeta)
+            os.makedirs(carpeta)
+
+            with open(file_path, "w", encoding="utf-8") as file:
+                for text in texts:
+                    file.write(text)
+                    file.write("\n")
+            return True
+        except Exception:
+            return False
+
+    def load_vectorstore(self) -> FAISS:
+        return FAISS.load_local(
+            folder_path=os.path.join("database", self.name),
+            embeddings=self.embeddings,
+            allow_dangerous_deserialization=True,
+        )
+
+    def add_files_vectorstore(self) -> Optional[FAISS]:
+        temp_folder = "docs"
+        if not os.path.exists(temp_folder):
+            os.makedirs(temp_folder)
+            return None
+
+        documents = DocumentProcessor(temp_folder).files_to_texts()
+        if not documents:
+            return None
+
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000, chunk_overlap=200, length_function=len
+        )
+        texts = text_splitter.split_documents(documents)
+        self.vectorstore = self.load_vectorstore()
+        self.vectorstore.add_documents(documents=texts)
+        self.vectorstore.save_local(folder_path=os.path.join("database", self.name))
+        return self.vectorstore
+
+    def download_vectorstore(self):
+        # generar un zip de la carpeta del vectorstore, crearlo en la carpeta temp y devolver la ruta
+        with ZipFile("temp/vectorstore.zip", "w") as zip:
+            for root, dirs, files in os.walk(f"database/{self.name}"):
+                for file in files:
+                    zip.write(os.path.join(root, file))
+        return "temp/vectorstore.zip"