Update model_service.py

model_service.py

@@ -1,6 +1,8 @@
 # model_service.py
 
 import os
+import logging
+from fastapi import HTTPException
 from langchain_community.vectorstores import FAISS
 from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
 from langchain.chains.question_answering import load_qa_chain
@@ -8,32 +10,37 @@ from langchain.prompts import PromptTemplate
 from config import BASE_MODEL_PATH, GOOGLE_DRIVE_FOLDER_ID
 from drive_service import DriveService
 
+logger = logging.getLogger(__name__)
+
 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:
-            # Initialize Google Drive API
-            drive_service = DriveService()
+            logger.info(f"Loading model: {model_name} with temperature: {temperature}")
 
             # Download model files from Google Drive
-            drive_service.download_model_files_from_subfolder(
+            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")
             )
 
-            # Load the local FAISS index and vector store
+            logger.info("Loading FAISS vector store...")
             vector_store = FAISS.load_local(
                 model_path,
                 embeddings,
@@ -41,6 +48,7 @@ class ModelService:
             )
 
             # Configure the QA chain
+            logger.info("Configuring QA chain...")
             chain = self.configure_chain(temperature)
 
             # Store the loaded model in memory
@@ -49,6 +57,7 @@ class ModelService:
                 "chain": chain
             }
 
+            logger.info(f"Model '{model_name}' loaded successfully")
             return {
                 "status": "success",
                 "message": f"Model '{model_name}' loaded successfully"
@@ -113,33 +122,4 @@ class ModelService:
             return load_qa_chain(model, chain_type="stuff", prompt=prompt)
         except Exception as e:
             logger.error(f"Error configuring chain: {str(e)}")
-            raise HTTPException(status_code=500, detail="Failed to configure model chain")
-
-    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)}")
+            raise HTTPException(status_code=500, detail="Failed to configure model chain")