Update model_service.py

model_service.py

@@ -14,85 +14,141 @@ class ModelService:
     def __init__(self):
         self.loaded_models = {}
         self.drive_service = DriveService()
-
+    
     def load_model(self, model_name: str, temperature: float = 0.7):
         """Load a model from Google Drive."""
         try:
             logger.info(f"Loading model: {model_name} with temperature: {temperature}")
-
+    
             # Download model files from Google Drive
             logger.info("Downloading model files from Google Drive...")
             self.drive_service.download_model_files_from_subfolder(
                 parent_folder_id=GOOGLE_DRIVE_FOLDER_ID,
                 subfolder_name=model_name
             )
-
+    
             # Load the downloaded model
-            model_path = os.path.join(BASE_MODEL_PATH, model_name)
+            model_path = os.path.join(BASE_MODEL_PATH, model_name, "faiss_index")  # Add "faiss_index" to the path
             logger.info(f"Model path: {model_path}")
-
+    
+            # Verify the model files exist
+            if not os.path.exists(os.path.join(model_path, "index.faiss")):
+                raise FileNotFoundError(f"FAISS index not found at {model_path}")
+    
             # Initialize embeddings and load vector store
             logger.info("Initializing embeddings...")
             embeddings = GoogleGenerativeAIEmbeddings(
                 model="models/embedding-001",
                 google_api_key=os.getenv("GOOGLE_API_KEY")
             )
-
+    
             logger.info("Loading FAISS vector store...")
             vector_store = FAISS.load_local(
-                model_path,
+                model_path,  # This path should now point to the faiss_index directory
                 embeddings,
                 allow_dangerous_deserialization=True
             )
-
+    
             # Configure the QA chain
             logger.info("Configuring QA chain...")
             chain = self.configure_chain(temperature)
-
+    
             # Store the loaded model in memory
             self.loaded_models[model_name] = {
                 "vector_store": vector_store,
                 "chain": chain
             }
-
+    
             logger.info(f"Model '{model_name}' loaded successfully")
             return {
                 "status": "success",
                 "message": f"Model '{model_name}' loaded successfully"
             }
+        except FileNotFoundError as e:
+            logger.error(f"File not found error: {str(e)}")
+            raise HTTPException(status_code=404, detail=str(e))
         except Exception as e:
             logger.error(f"Error loading model: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}")
-
-    def chat_with_model(self, model_name: str, question: str):
-        """Generate a response using the loaded model."""
-        if model_name not in self.loaded_models:
-            raise HTTPException(
-                status_code=404,
-                detail=f"Model '{model_name}' not loaded. Please load it first."
-            )
-
-        try:
-            model_data = self.loaded_models[model_name]
-            docs = model_data["vector_store"].similarity_search(question)
-            response = model_data["chain"](
-                {
-                    "input_documents": docs,
-                    "question": question
-                },
-                return_only_outputs=True
-            )
-
-            return {
-                "status": "success",
-                "response": response["output_text"]
-            }
-        except Exception as e:
-            logger.error(f"Error generating response: {str(e)}")
-            raise HTTPException(
-                status_code=500,
-                detail=f"Failed to generate response: {str(e)}"
-            )
+    
+    # def load_model(self, model_name: str, temperature: float = 0.7):
+    #     """Load a model from Google Drive."""
+    #     try:
+    #         logger.info(f"Loading model: {model_name} with temperature: {temperature}")
+
+    #         # Download model files from Google Drive
+    #         logger.info("Downloading model files from Google Drive...")
+    #         self.drive_service.download_model_files_from_subfolder(
+    #             parent_folder_id=GOOGLE_DRIVE_FOLDER_ID,
+    #             subfolder_name=model_name
+    #         )
+
+    #         # Load the downloaded model
+    #         model_path = os.path.join(BASE_MODEL_PATH, model_name)
+    #         logger.info(f"Model path: {model_path}")
+
+    #         # Initialize embeddings and load vector store
+    #         logger.info("Initializing embeddings...")
+    #         embeddings = GoogleGenerativeAIEmbeddings(
+    #             model="models/embedding-001",
+    #             google_api_key=os.getenv("GOOGLE_API_KEY")
+    #         )
+
+    #         logger.info("Loading FAISS vector store...")
+    #         vector_store = FAISS.load_local(
+    #             model_path,
+    #             embeddings,
+    #             allow_dangerous_deserialization=True
+    #         )
+
+    #         # Configure the QA chain
+    #         logger.info("Configuring QA chain...")
+    #         chain = self.configure_chain(temperature)
+
+    #         # Store the loaded model in memory
+    #         self.loaded_models[model_name] = {
+    #             "vector_store": vector_store,
+    #             "chain": chain
+    #         }
+
+    #         logger.info(f"Model '{model_name}' loaded successfully")
+    #         return {
+    #             "status": "success",
+    #             "message": f"Model '{model_name}' loaded successfully"
+    #         }
+    #     except Exception as e:
+    #         logger.error(f"Error loading model: {str(e)}")
+    #         raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}")
+
+    # def chat_with_model(self, model_name: str, question: str):
+    #     """Generate a response using the loaded model."""
+    #     if model_name not in self.loaded_models:
+    #         raise HTTPException(
+    #             status_code=404,
+    #             detail=f"Model '{model_name}' not loaded. Please load it first."
+    #         )
+
+    #     try:
+    #         model_data = self.loaded_models[model_name]
+    #         docs = model_data["vector_store"].similarity_search(question)
+    #         response = model_data["chain"](
+    #             {
+    #                 "input_documents": docs,
+    #                 "question": question
+    #             },
+    #             return_only_outputs=True
+    #         )
+
+    #         return {
+    #             "status": "success",
+    #             "response": response["output_text"]
+    #         }
+    #     except Exception as e:
+    #         logger.error(f"Error generating response: {str(e)}")
+    #         raise HTTPException(
+    #             status_code=500,
+    #             detail=f"Failed to generate response: {str(e)}"
+    #         )
 
     def configure_chain(self, temperature: float):
         """Configure the QA chain with the updated prompt template."""