app.py

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.")