Refactor vector store loading logic to return the vector store object instead of a string message and enhance error handling during download

app.py

@@ -60,24 +60,14 @@ def load_vector_store():
         print(f"Debug - Download result: success={success}, result_type={type(result)}")
         
         if success:
-            if result is None:
-                print("Debug - Vector store is None despite success=True")
-                return None
-                
             if isinstance(result, str):
-                print(f"Debug - Vector store is a string: {result}")
-                return None
-                
-            # Check if the result has a similarity_search method
-            if hasattr(result, 'similarity_search'):
-                print("Debug - Vector store loaded successfully with similarity_search method")
-                return result
-            else:
-                print(f"Debug - Vector store object does not have similarity_search method: {type(result)}")
+                print(f"Debug - Error message received: {result}")
                 return None
+            return result
         else:
             print(f"Debug - Failed to load vector store: {result}")
             return None
+            
     except Exception as e:
         import traceback
         print(f"Exception loading knowledge base: {str(e)}")

src/knowledge_base/dataset.py

@@ -9,7 +9,8 @@ from typing import Tuple, List, Dict, Any, Optional, Union
 from datetime import datetime
 from huggingface_hub import HfApi, HfFolder
 from langchain_community.vectorstores import FAISS
-from config.settings import VECTOR_STORE_PATH, HF_TOKEN
+from config.settings import VECTOR_STORE_PATH, HF_TOKEN, EMBEDDING_MODEL
+from langchain.embeddings import HuggingFaceEmbeddings
 
 class DatasetManager:
     def __init__(self, dataset_name="Rulga/status-law-knowledge-base", token: Optional[str] = None):
@@ -188,7 +189,17 @@ class DatasetManager:
         try:
             # Check if local files exist and force is False
             if not force and os.path.exists(os.path.join(VECTOR_STORE_PATH, "index.faiss")):
-                return True, "Local files exist"
+                # Instead of returning string, load and return the vector store
+                embeddings = HuggingFaceEmbeddings(
+                    model_name=EMBEDDING_MODEL,
+                    model_kwargs={'device': 'cpu'}
+                )
+                vector_store = FAISS.load_local(
+                    VECTOR_STORE_PATH,
+                    embeddings,
+                    allow_dangerous_deserialization=True
+                )
+                return True, vector_store
             
             # Ensure vector store directory exists
             os.makedirs(VECTOR_STORE_PATH, exist_ok=True)
@@ -209,10 +220,20 @@ class DatasetManager:
                     local_dir=VECTOR_STORE_PATH
                 )
                 
-                return True, "Vector store downloaded successfully"
-            
-            except Exception as download_error:
-                return False, f"Failed to download vector store: {str(download_error)}"
+                # After successful download, load and return the vector store
+                embeddings = HuggingFaceEmbeddings(
+                    model_name=EMBEDDING_MODEL,
+                    model_kwargs={'device': 'cpu'}
+                )
+                vector_store = FAISS.load_local(
+                    VECTOR_STORE_PATH,
+                    embeddings,
+                    allow_dangerous_deserialization=True
+                )
+                return True, vector_store
+                
+            except Exception as e:
+                return False, f"Failed to download vector store: {str(e)}"
             
         except Exception as e:
             return False, f"Error in download_vector_store: {str(e)}"