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