src/vector_store.py

import os
from langchain.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings
from dotenv import load_dotenv

load_dotenv()

def get_embedding_model(model_choice):
    if model_choice == "OpenAI (Paid)":
        return OpenAIEmbeddings(api_key=os.getenv("OPENAI_API_KEY"))
    else:
        return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

def build_vector_store(chunks, save_path, model_choice):
    embedding_model = get_embedding_model(model_choice)
    db = FAISS.from_documents(chunks, embedding_model)
    db.save_local(save_path)
    return db

def load_vector_store(load_path, model_choice):
    embedding_model = get_embedding_model(model_choice)
    return FAISS.load_local(load_path, embedding_model)