Create chatbot_service.py

chatbot_service.py

@@ -0,0 +1,76 @@
+# app/services/chatbot_service.py
+
+import os
+import pickle
+from typing import List
+from fastapi import HTTPException, UploadFile, File
+from PyPDF2 import PdfReader
+from langchain_community.vectorstores import FAISS
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from app.services.drive_service import DriveService
+from app.config import CUSTOM_CHATBOTS_DIR
+
+class ChatbotService:
+    def __init__(self):
+        self.drive_service = DriveService()
+
+    def create_chatbot(self, folder_name: str, pdf_files: List[UploadFile]):
+        """Create a new chatbot using only PDF files."""
+        try:
+            # Validate input
+            if not pdf_files:
+                raise HTTPException(status_code=400, detail="At least one PDF file is required.")
+
+            # Create folder for the chatbot
+            folder_path = os.path.join(CUSTOM_CHATBOTS_DIR, folder_name)
+            os.makedirs(folder_path, exist_ok=True)
+
+            # Process PDFs
+            raw_text = ""
+            for pdf_file in pdf_files:
+                pdf_reader = PdfReader(pdf_file.file)
+                for page in pdf_reader.pages:
+                    raw_text += page.extract_text() or ""
+
+            # Split text into chunks
+            text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
+            text_chunks = text_splitter.split_text(raw_text)
+
+            # Create vector store
+            embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
+            vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
+
+            # Save the vector store and documents
+            vector_store.save_local(os.path.join(folder_path, "faiss_index"))
+            with open(os.path.join(folder_path, "documents.pkl"), 'wb') as f:
+                pickle.dump(text_chunks, f)
+
+            # Upload files to Google Drive
+            # Step 1: Create or find the folder with the provided `folder_name`
+            chatbot_folder_id = self.drive_service.find_folder_by_name(folder_name, GOOGLE_DRIVE_FOLDER_ID)
+            if not chatbot_folder_id:
+                chatbot_folder_id = self.drive_service.create_folder_in_google_drive(folder_name, GOOGLE_DRIVE_FOLDER_ID)
+
+            # Step 2: Create or find the `faiss_index` subfolder inside the `folder_name` folder
+            faiss_index_folder_id = self.drive_service.find_folder_by_name("faiss_index", chatbot_folder_id)
+            if not faiss_index_folder_id:
+                faiss_index_folder_id = self.drive_service.create_folder_in_google_drive("faiss_index", chatbot_folder_id)
+
+            # Step 3: Upload index.faiss and index.pkl to the `faiss_index` subfolder
+            index_pkl_path = os.path.join(folder_path, "faiss_index", "index.pkl")
+            os.rename(os.path.join(folder_path, "documents.pkl"), index_pkl_path)
+
+            faiss_index_path = os.path.join(folder_path, "faiss_index", "index.faiss")
+            self.drive_service.upload_file_to_google_drive(faiss_index_path, faiss_index_folder_id)
+            self.drive_service.upload_file_to_google_drive(index_pkl_path, faiss_index_folder_id)
+
+            return {
+                "status": "success",
+                "message": f"Chatbot '{folder_name}' created and files uploaded to Google Drive successfully!"
+            }
+        except HTTPException:
+            raise
+        except Exception as e:
+            logger.error(f"Error creating chatbot: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Failed to create chatbot: {str(e)}")