Coverage for src/app/dependencies/init_vector_store.py: 100%

1import os 

2from langchain_community.vectorstores import FAISS 

3from langchain_openai import OpenAIEmbeddings 

4 

5VECTOR_STORE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'vector_store')) 

6 

7def store_vector_store(vector_store: FAISS): 

8 """Stores the vector store to disk.""" 

9 

10 vector_store.save_local(VECTOR_STORE_PATH) 

11 

12 print(f"Saved vector store to {VECTOR_STORE_PATH}") 

13 

14def load_vector_store(embedding_model: OpenAIEmbeddings) -> FAISS: 

15 

16 vector_store = FAISS.load_local( 

17 VECTOR_STORE_PATH, 

18 embeddings=embedding_model, 

19 allow_dangerous_deserialization=True # trust the data source 

20 ) 

21 

22 file_size = os.path.getsize(VECTOR_STORE_PATH + "/index.faiss") 

23 print(f"Vector store loaded from {VECTOR_STORE_PATH}") 

24 print(f"Size of the FAISS file: {file_size / (1024 * 1024):.2f} MB") 

25 print(f"Number of vectors: {vector_store.index.ntotal}") 

26 print(f"Dimension of vectors: {vector_store.index.d}") 

27 print(f"Metric type: {vector_store.index.metric_type}") 

28 

29 return vector_store 

30