core/vector_store.py

from typing import Optional
from pathlib import Path
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings

from . import utils
from .config import logger, VECTOR_STORE_DIR, EMBEDDING_MODEL

# Global instance
_vector_store_instance = None

def get_vector_store() -> Optional[FAISS]:
    """Get the current vector store instance"""
    global _vector_store_instance
    return _vector_store_instance

def create_vector_store():
    """Create a new vector store instance"""
    global _vector_store_instance
    
    # Try to load existing vector store, create if not found
    try:
        company_documents = utils.create_company_documents()
        company_chunks = utils.split_documents(company_documents)
        _vector_store_instance = utils.create_company_vector_store(company_chunks)
        logger.info("Vector store created successfully")
        return _vector_store_instance
    except Exception as e:
        logger.error(f"Error creating vector store: {str(e)}")
        raise


def load_vector_store() -> Optional[FAISS]:
    """Load existing vector store with proper error handling"""
    global _vector_store_instance
    try:
        if Path(VECTOR_STORE_DIR).exists():
            _vector_store_instance = FAISS.load_local(
                str(VECTOR_STORE_DIR), 
                EMBEDDING_MODEL, 
                allow_dangerous_deserialization=True
            )
            logger.info("Successfully loaded existing vector store")
            return _vector_store_instance
        else:
            logger.info("No existing vector store found")
            logger.info("Creating new vector store...")
            return create_vector_store()
    except Exception as e:
        logger.error(f"Failed to load vector store: {e}")
        return None