Update back.py

back.py

@@ -1,3 +1,4 @@
+# back.py
 import os
 import logging
 from typing import List, Dict, Any, Optional, Union
@@ -9,6 +10,7 @@ from langchain_core.embeddings import Embeddings
 import google.generativeai as genai
 from datetime import datetime
 import json
+import pickle
 
 @dataclass
 class UserInfo:
@@ -26,10 +28,11 @@ class ChatConfig:
     """Configuration for the chatbot"""
     embedding_model_name: str = 'all-MiniLM-L6-v2'
     device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
-    max_history: int = 3
+    max_history: int = 6
     gemini_api_key: str = os.getenv("GEMINI_API")  # Replace with your API key
     log_file: str = "chat_history.txt"
     user_data_file: str = "user_data.json"
+    database_file: str = "faiss_db.pkl" # Added database file path
 
 # In the UserManager class, modify these methods:
 class UserManager:
@@ -178,7 +181,7 @@ class GeminiRAG:
             "max_output_tokens": 8192,
         }
         self.model = genai.GenerativeModel(
-            model_name="gemini-1.5-flash",
+            model_name="gemini-2.0-flash-exp",
             generation_config=self.generation_config,
             safety_settings={'HATE': 'BLOCK_NONE','HARASSMENT': 'BLOCK_NONE','SEXUAL' : 'BLOCK_NONE','DANGEROUS' : 'BLOCK_NONE'}
         )
@@ -253,7 +256,23 @@ class ProductDatabase:
             device=config.device
         )
         self.vectorstore = None
+        self.config = config
+        self.load_database()
         
+    def load_database(self):
+        """Loads the FAISS database from file"""
+        try:
+            if os.path.exists(self.config.database_file):
+                with open(self.config.database_file, "rb") as f:
+                    self.vectorstore = pickle.load(f)
+                print("Database loaded successfully from file.")
+            else:
+                print("Database file not found. Please run setup.py to create it.")
+        except Exception as e:
+            logging.error(f"Error loading database: {str(e)}")
+            print(f"Error loading database: {str(e)}")
+            self.vectorstore = None
+    
     def process_markdown(self, markdown_content: str):
         """Process markdown content and create vector store"""
         try:
@@ -278,16 +297,19 @@ class ProductDatabase:
             texts = [doc["content"] for doc in documents]
             metadatas = [{"section": doc["section"]} for doc in documents]
             
-            self.vectorstore = FAISS.from_texts(
-                texts=texts,
-                embedding=self.embeddings,
-                metadatas=metadatas
-            )
+            if self.vectorstore is None:
+                self.vectorstore = FAISS.from_texts(
+                    texts=texts,
+                    embedding=self.embeddings,
+                    metadatas=metadatas
+                )
+            else:
+                self.vectorstore.add_texts(texts=texts, metadatas=metadatas, embedding=self.embeddings)
             
         except Exception as e:
             raise Exception(f"Error processing markdown content: {str(e)}")
         
-    def search(self, query: str, k: int = 3) -> List[Dict[str, Any]]:
+    def search(self, query: str, k: int = 15) -> List[Dict[str, Any]]:
         """Search for relevant documents"""
         if not self.vectorstore:
             raise ValueError("Database not initialized. Please process documents first.")