Update main.py

main.py

@@ -6,7 +6,7 @@ import aiofiles
 import uuid
 import shutil
 
-# from dotenv import load_dotenv
+from dotenv import load_dotenv
 
 from langchain_community.document_loaders import TextLoader, Docx2txtLoader, PyPDFLoader
 from langchain.prompts import ChatPromptTemplate, PromptTemplate
@@ -18,11 +18,12 @@ from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from langchain_community.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
 
-# load_dotenv()
+load_dotenv()
 
 app = FastAPI()
 
 origins = ["https://viboognesh-react-chat.static.hf.space"]
+# origins = ["http://localhost:3000"]
 
 app.add_middleware(
     CORSMiddleware,
@@ -34,12 +35,21 @@ app.add_middleware(
 
 
 class ConversationChainManager:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super(ConversationChainManager, cls).__new__(
+                cls, *args, **kwargs
+            )
+        return cls._instance
+
     def __init__(self):
         self.conversation_chain = None
         self.llm_model = ChatOpenAI()
         self.embeddings = OpenAIEmbeddings()
 
-    def create_conversational_chain(self, file_paths: List[str], session_id: str):
+    def create_conversational_chain(self, file_paths: List[str]):
         docs = self.get_docs(file_paths)
         memory = ConversationBufferMemory(
             memory_key="chat_history", return_messages=True
@@ -47,8 +57,6 @@ class ConversationChainManager:
         vectordb = Chroma.from_documents(
             docs,
             self.embeddings,
-            collection_name=session_id,
-            persist_directory="./chroma_db",
         )
         retriever = vectordb.as_retriever()
         self.conversation_chain = ConversationalRetrievalChain.from_llm(
@@ -90,6 +98,7 @@ class ConversationChainManager:
                 loader = PyPDFLoader(file_path)
                 pdf_documents = loader.load_and_split()
                 docs.extend(pdf_documents)
+            os.remove(file_path)
         return docs
 
     @staticmethod
@@ -136,13 +145,17 @@ class ConversationChainManager:
         return ChatPromptTemplate.from_messages(messages)
 
 
+app.state.conversational_chain_manager = ConversationChainManager()
+
+
 @app.post("/upload_files/")
 async def upload_files(
     files: List[UploadFile] = File(...),
-    conversation_chain_manager: ConversationChainManager = Depends(),
+    conversation_chain_manager: ConversationChainManager = Depends(
+        lambda: app.state.conversational_chain_manager
+    ),
 ):
-    session_id = str(uuid.uuid4())
-    session_folder = f"uploads/{session_id}"
+    session_folder = f"uploads"
     os.makedirs(session_folder, exist_ok=True)
     file_paths = []
     for file in files:
@@ -152,15 +165,17 @@ async def upload_files(
             await out_file.write(content)
         file_paths.append(file_path)
 
-    conversation_chain_manager.create_conversational_chain(file_paths, session_id)
-    # shutil.rmtree(session_folder)
+    conversation_chain_manager.create_conversational_chain(file_paths)
     print("conversational_chain_manager created")
     return {"message": "ConversationalRetrievalChain is created. Please ask questions."}
 
 
 @app.get("/predict/")
 async def predict(
-    query: str, conversation_chain_manager: ConversationChainManager = Depends()
+    query: str,
+    conversation_chain_manager: ConversationChainManager = Depends(
+        lambda: app.state.conversational_chain_manager
+    ),
 ):
     if conversation_chain_manager.conversation_chain is None:
         system_prompt = "Answer the question and also ask the user to upload files to ask questions from the files.\n"