Update model_service.py

model_service.py

@@ -1,8 +1,11 @@
-# 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
+from langchain.prompts import PromptTemplate
+from config import BASE_MODEL_PATH, GOOGLE_DRIVE_FOLDER_ID
 from drive_service import DriveService  # Import DriveService
 
 logger = logging.getLogger(__name__)
@@ -12,106 +15,139 @@ class ModelService:
         self.loaded_models = {}
         self.drive_service = DriveService()
 
-    
-    @app.post("/load-model/")
-    async def load_model(request: LoadModelRequest):
+    def load_model(self, model_name: str, temperature: float = 0.7):
         """Load a model from Google Drive."""
         try:
-            logger.info(f"Loading model: {request.model_name} with temperature: {request.temperature}")
-    
+            logger.info(f"Loading model: {model_name} with temperature: {temperature}")
+
             # Download model files from Google Drive
             logger.info("Downloading model files from Google Drive...")
-            drive_service.download_model_files_from_subfolder(
+            self.drive_service.download_model_files_from_subfolder(
                 parent_folder_id=GOOGLE_DRIVE_FOLDER_ID,
-                subfolder_name=request.model_name
+                subfolder_name=model_name
             )
-    
+
             # Load the downloaded model
-            model_path = os.path.join(BASE_MODEL_PATH, request.model_name)
+            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 = model_service.configure_chain(request.temperature)
-    
+            chain = self.configure_chain(temperature)
+
             # Store the loaded model in memory
-            model_service.loaded_models[request.model_name] = {
+            self.loaded_models[model_name] = {
                 "vector_store": vector_store,
                 "chain": chain
             }
-    
-            logger.info(f"Model '{request.model_name}' loaded successfully")
+
+            logger.info(f"Model '{model_name}' loaded successfully")
             return {
                 "status": "success",
-                "message": f"Model '{request.model_name}' loaded successfully"
+                "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 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:
+        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."""
+        prompt_template = """
+        You are an AI assistant for SBBU SBA university. Your task is to provide clear, accurate, and helpful responses based on the context provided, as well as to respond to basic greetings and conversational queries. However, if the user makes inappropriate or offensive remarks, you should respond politely and professionally, redirecting the conversation back to helpful topics.
+
+        Instructions:
+        1. **Greeting Responses**: If the user greets you (e.g., "Hello," "Hi," "Hey," "Salam," etc.), respond warmly and politely. Example responses could be:
+           - "Hello! How can I assist you today?"
+           - "Hi there! How can I help you?"
+           - "Salam! What can I do for you today?"
+
+        2. **Casual and Playful Inquiries**: If the user says something playful or informal like "I kiss you" or similar, acknowledge it politely but redirect the conversation back to the main topic. Example:
+           - "Thank you for the kind words! How can I assist you further?"
+           - "I appreciate your enthusiasm! How can I help you today?"
+
+        3. **Inappropriate or Offensive Remarks**: If the user makes inappropriate, disrespectful, or offensive comments, such as offensive language or sexually explicit remarks, respond politely but firmly, maintaining professionalism:
+           - "I strive to maintain a respectful conversation. How can I assist you with your queries?"
+           - "Let's keep the conversation respectful. How can I help you today?"
+           - "I apologize, but I cannot engage in that kind of discussion. Please ask a relevant question related to the university."
+
+        4. **Contextual Responses**:
+           - If the context contains relevant information to the question, provide a clear and direct answer.
+           - If the context only provides partial information, provide a helpful response based on available data and related details.
+           - If the context has no relevant information, respond with: "I apologize, but I don't have specific information about that. Could you please ask something else about the university?"
+
+        5. **Accuracy and Clarity**: Ensure your responses are clear, concise, and accurate. Avoid unnecessary details or over-explanation.
+
+        6. **Clarification**: If the user's question is unclear or lacks sufficient context, ask for clarification. For example:
+           - "Could you please clarify your question?"
+           - "I'm not sure I understand. Can you rephrase your question?"
+
+        Context Information:
+        ---------------------
+        {context}
+
+        Question:
+        {question}
+
+        Response:
+        Provide a friendly, clear, and direct response based on the context. Always aim to be helpful, especially for greetings or casual inquiries, and suggest follow-up questions or clarifications if needed.
+            no preamble
+        """
+
+        try:
+            model = ChatGoogleGenerativeAI(
+                model="gemini-pro",
+                temperature=temperature,
+                google_api_key=os.getenv("GOOGLE_API_KEY")
+            )
+            prompt = PromptTemplate(
+                template=prompt_template,
+                input_variables=["context", "question"]
+            )
+            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")