mr

Build error

App Files Files Community

JPBianchi commited on Jun 24, 2024

Commit

747f388

1 Parent(s): 24f6e72

fixed call to main_reflex with uvicorn

Browse files

Files changed (2) hide show

Dockerfile +1 -1
app/main.py +93 -53

Dockerfile CHANGED Viewed

@@ -48,4 +48,4 @@ RUN guardrails hub install hub://guardrails/detect_pii
 RUN guardrails hub install hub://guardrails/qa_relevance_llm_eval
-CMD ["uvicorn", "app.main_reflex:app", "--host", "0.0.0.0", "--port", "7860"]


48	RUN guardrails hub install hub://guardrails/qa_relevance_llm_eval
49
50
51	+ CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app/main.py CHANGED Viewed

@@ -1,4 +1,3 @@
 import os, random, logging, pickle, shutil
 from dotenv import load_dotenv, find_dotenv
 from typing import Optional
@@ -7,15 +6,21 @@ from pydantic import BaseModel, Field
 from fastapi import FastAPI, HTTPException, File, UploadFile, status
 from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
-from engine.processing import process_pdf, index_data, empty_collection, vector_search
-from rag.rag import rag_it
-from engine.logger import logger
-from settings import datadir
-os.makedirs(datadir, exist_ok=True)
 app = FastAPI()
@@ -37,74 +42,83 @@ try:
 except Exception as e:
     pass
-@app.get("/", response_class=HTMLResponse)
 def read_root():
-    logger.info("Title displayed on home page")
     return """
     <html>
         <body>
-            <h1>Welcome to FinExpert, a RAG system designed by JP Bianchi!</h1>
         </body>
     </html>
     """
-@app.get("/ping/")
 def ping():
     """ Testing """
-    logger.info("Someone is pinging the server")
     return {"answer": str(int(random.random() * 100))}
-@app.delete("/erase_data/")
 def erase_data():
-    """ Erase all files in the data directory, but not the vector store """
     if len(os.listdir(datadir)) == 0:
-        logger.info("No data to erase")
         return {"message": "No data to erase"}
-    shutil.rmtree(datadir, ignore_errors=True)
-    os.mkdir(datadir)
-    logger.warning("All data has been erased")
     return {"message": "All data has been erased"}
-@app.delete("/empty_collection/")
 def delete_vectors():
     """ Empty the collection in the vector store """
     try:
         status = empty_collection()
-        return {f"""message": "Collection{'' if status else ' NOT'} erased!"""}
     except Exception as e:
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
-@app.get("/list_files/")
 def list_files():
     """ List all files in the data directory """
     files = os.listdir(datadir)
-    logger.info(f"Files in data directory: {files}")
     return {"files": files}
-@app.post("/upload/")
 # @limiter.limit("5/minute") see 'slowapi' for rate limiting
 async def upload_file(file: UploadFile = File(...)):
     """  Uploads a file in data directory, for later indexing """
     try:
         filepath = os.path.join(datadir, file.filename)
-        logger.info(f"Fiename detected: {file.filename}")
         if os.path.exists(filepath):
-            logger.warning(f"File {file.filename} already exists: no processing done")
             return {"message": f"File {file.filename} already exists: no processing done"}
         else:
-            logger.info(f"Receiving file: {file.filename}")
             contents = await file.read()
-            logger.info(f"File reception complete!")
     except Exception as e:
-        logger.error(f"Error during file upload: {str(e)}")
         return {"message": f"Error during file upload:  {str(e)}"}
     if file.filename.endswith('.pdf'):
@@ -112,9 +126,14 @@ async def upload_file(file: UploadFile = File(...)):
         # let's save the file in /data even if it's temp storage on HF
         with open(filepath, 'wb') as f:
             f.write(contents)
         try:
-            logger.info(f"Starting to process {file.filename}")
             new_content = process_pdf(filepath)
             success = {"message": f"Successfully uploaded {file.filename}"}
             success.update(new_content)
@@ -122,15 +141,35 @@ async def upload_file(file: UploadFile = File(...)):
         except Exception as e:
             return {"message": f"Failed to extract text from PDF: {str(e)}"}
     else:
-        return {"message": "Only PDF files are accepted"}
-@app.post("/create_index/")
 async def create_index():
     """ Create an index for the uploaded files """
-    logger.info("Creating index for uploaded files")
     try:
         msg = index_data()
         return {"message": msg}
@@ -141,50 +180,51 @@ async def create_index():
 class Question(BaseModel):
     question: str
-@app.post("/ask/")
 async def hybrid_search(question: Question):
-    logger.info(f"Processing question: {question.question}")
     try:
         search_results = vector_search(question.question)
-        logger.info(f"Answer: {search_results}")
         return {"answer": search_results}
     except Exception as e:
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
-@app.post("/ragit/")
 async def ragit(question: Question):
-    logger.info(f"Processing question: {question.question}")
     try:
-        search_results = vector_search(question.question)
-        logger.info(f"Search results generated: {search_results}")
         answer = rag_it(question.question, search_results)
-        logger.info(f"Answer: {answer}")
         return {"answer": answer}
     except Exception as e:
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
 if __name__ == '__main__':
     import uvicorn
     from os import getenv
-    port = int(getenv("PORT", 80))
     print(f"Starting server on port {port}")
     reload = True if environment == "dev" else False
     uvicorn.run("main:app", host="0.0.0.0", port=port, reload=reload)
 # Examples:
-# curl -X POST "http://localhost:80/upload" -F "file=@test.pdf"
-# curl -X DELETE "http://localhost:80/erase_data/"
-# curl -X GET "http://localhost:80/list_files/"
-# hf space is at https://jpbianchi-finrag.hf.space/
-# code given by https://jpbianchi-finrag.hf.space/docs
 # Space must be public
-# curl -X POST "https://jpbianchi-finrag.hf.space/upload/" -F "file=@test.pdf"
-# curl -X POST http://localhost:80/ask/ -H "Content-Type: application/json" -d '{"question": "what is Amazon loss"}'
-# curl -X POST http://localhost:80/ragit/ -H "Content-Type: application/json" -d '{"question": "Does ATT have postpaid phone customers?"}'

 import os, random, logging, pickle, shutil
 from dotenv import load_dotenv, find_dotenv
 from typing import Optional
 from fastapi import FastAPI, HTTPException, File, UploadFile, status
 from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
+from app.engine.processing import ( # << creates the collection already
+    process_pdf,
+    process_txt,
+    index_data,
+    empty_collection,
+    vector_search,
+    vector_search_raw,
+)
+from app.rag.rag import rag_it
+from app.engine.logger import logger
+from app.settings import datadir, datadir2
+EXTENSIONS = ["pdf", "txt"]
 app = FastAPI()
 except Exception as e:
     pass
 def read_root():
+    logger("Title displayed on home page")
     return """
     <html>
         <body>
+            <h1>Welcome to MultiRAG, a RAG system designed by JP Bianchi!</h1>
         </body>
     </html>
     """
+# already provided by Reflex
+# @app.get("/ping/")
 def ping():
     """ Testing """
+    logger("Someone is pinging the server")
     return {"answer": str(int(random.random() * 100))}
+# @app.delete("/erase_data/")
 def erase_data():
+    """ Erase all files in the data directory at the first level only,
+        (in case we would like to use it for something else)
+        but not the vector store or the parquet file.
+        We can do it since the embeddings are in the parquet file already.
+    """
     if len(os.listdir(datadir)) == 0:
+        logger("No data to erase")
         return {"message": "No data to erase"}
+    # if we try to rmtree datadir, it looks like /data can't be deleted on HF
+    for f in os.listdir(datadir):
+        if f == '.DS_Store' or f.split('.')[-1].lower() in EXTENSIONS:
+            print(f"Removing {f}")
+            os.remove(os.path.join(datadir, f))
+            # we don't remove the parquet file, create_index does that
+    logger("All data has been erased")
     return {"message": "All data has been erased"}
+# @app.delete("/empty_collection/")
 def delete_vectors():
     """ Empty the collection in the vector store """
     try:
         status = empty_collection()
+        return {"message": f"Collection{'' if status else ' NOT'} erased!"}
     except Exception as e:
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
+# @app.get("/list_files/")
 def list_files():
     """ List all files in the data directory """
+    print("Listing files")
     files = os.listdir(datadir)
+    logger(f"Files in data directory: {files}")
     return {"files": files}
+# @app.post("/upload/")
 # @limiter.limit("5/minute") see 'slowapi' for rate limiting
 async def upload_file(file: UploadFile = File(...)):
     """  Uploads a file in data directory, for later indexing """
     try:
         filepath = os.path.join(datadir, file.filename)
+        logger(f"Fiename detected: {file.filename}")
         if os.path.exists(filepath):
+            logger(f"File {file.filename} already exists: no processing done")
             return {"message": f"File {file.filename} already exists: no processing done"}
         else:
+            logger(f"Receiving file: {file.filename}")
             contents = await file.read()
+            logger(f"File reception complete!")
     except Exception as e:
+        logger(f"Error during file upload: {str(e)}")
         return {"message": f"Error during file upload:  {str(e)}"}
     if file.filename.endswith('.pdf'):
         # let's save the file in /data even if it's temp storage on HF
         with open(filepath, 'wb') as f:
             f.write(contents)
+        # save it also in assets/data because data can be cleared
+        filepath2 = os.path.join(datadir2, file.filename)
+        with open(filepath2, 'wb') as f:
+            f.write(contents)
         try:
+            logger(f"Starting to process {file.filename}")
             new_content = process_pdf(filepath)
             success = {"message": f"Successfully uploaded {file.filename}"}
             success.update(new_content)
         except Exception as e:
             return {"message": f"Failed to extract text from PDF: {str(e)}"}
+    elif file.filename.endswith('.txt'):
+        with open(filepath, 'wb') as f:
+            f.write(contents)
+        filepath2 = os.path.join(datadir2, file.filename)
+        with open(filepath2, 'wb') as f:
+            f.write(contents)
+        try:
+            logger(f"Reading {file.filename}")
+            new_content = process_txt(filepath)
+            success = {"message": f"Successfully uploaded {file.filename}"}
+            success.update(new_content)
+            return success
+        except Exception as e:
+            return {"message": f"Failed to extract text from TXT: {str(e)}"}
     else:
+        return {"message": "Only PDF & txt files are accepted"}
+# @app.post("/create_index/")
 async def create_index():
     """ Create an index for the uploaded files """
+    logger("Creating index for uploaded files")
     try:
         msg = index_data()
         return {"message": msg}
 class Question(BaseModel):
     question: str
+# @app.post("/ask/")
 async def hybrid_search(question: Question):
+    logger(f"Processing question: {question.question}")
     try:
         search_results = vector_search(question.question)
+        logger(f"Answer: {search_results}")
         return {"answer": search_results}
     except Exception as e:
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
+# @app.post("/ragit/")
 async def ragit(question: Question):
+    logger(f"Processing question: {question.question}")
     try:
+        search_results = vector_search_raw(question.question)
+        logger(f"Search results generated: {search_results}")
         answer = rag_it(question.question, search_results)
+        logger(f"Answer: {answer}")
         return {"answer": answer}
     except Exception as e:
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
 if __name__ == '__main__':
     import uvicorn
     from os import getenv
+    port = int(getenv("PORT", 8000))
     print(f"Starting server on port {port}")
     reload = True if environment == "dev" else False
     uvicorn.run("main:app", host="0.0.0.0", port=port, reload=reload)
 # Examples:
+# curl -X POST "http://localhost:8001/upload" -F "file=@test.pdf"
+# curl -X DELETE "http://localhost:8001/erase_data/"
+# curl -X GET "http://localhost:8001/list_files/"
+# hf space is at https://jpbianchi-mr.hf.space/
+# code given by https://jpbianchi-mr.hf.space/docs
 # Space must be public
+# curl -X POST "https://jpbianchi-mr.hf.space/upload/" -F "file=@test.pdf"
+# curl -X POST http://localhost:8000/ask/ -H "Content-Type: application/json" -d '{"question": "what is Amazon loss"}'
+# curl -X POST http://localhost:8000/ragit/ -H "Content-Type: application/json" -d '{"question": "Does ATT have postpaid phone customers?"}'
+# see more in notebook upload_index.ipynb